diff --git a/CHANGELOG.md b/CHANGELOG.md index fb759651..b5d357c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed +- CrewAI listeners preserve lifecycle callback order and per-crew context so repeated invocations retain + all spans and their parent relationships. Standalone calls and overlapping crews retain telemetry and token + usage, child callback errors do not end healthy crews, and spans use original event timestamps. +- CrewAI OTLP completion flushes run outside the shared event loop when explicitly enabled. Legacy ingestion-hook + flushes wait for active runs to finish and remain synchronous to preserve their whole-trace batches. - Restore completed SDK operations to their caller context when framework callbacks finish in a copied async task, preserving caller span recording and preventing stale parents and missing telemetry on subsequent LangGraph invocations. - LangChain callbacks record the terminal response of a LangGraph root instead of accumulated checkpoint history, diff --git a/src/splunk_ao/handlers/crewai/handler.py b/src/splunk_ao/handlers/crewai/handler.py index 6adc870c..2374e48a 100644 --- a/src/splunk_ao/handlers/crewai/handler.py +++ b/src/splunk_ao/handlers/crewai/handler.py @@ -1,7 +1,11 @@ import hashlib import json import logging -from datetime import datetime +from collections.abc import Callable +from contextvars import Context, copy_context +from dataclasses import dataclass, field +from datetime import UTC, datetime +from functools import wraps from typing import Any from uuid import UUID @@ -10,7 +14,7 @@ from splunk_ao.handlers.base_handler import SplunkAOBaseHandler from splunk_ao.logger import SplunkAOLogger from splunk_ao.schema.handlers import NodeType -from splunk_ao.utils.serialization import serialize_to_str +from splunk_ao.utils.serialization import convert_time_delta_to_ns, serialize_to_str _logger = logging.getLogger(__name__) @@ -53,6 +57,19 @@ def _resolve_crewai_imports() -> None: litellm = None +@dataclass +class _CrewRun: + """Callback state retained until one crew or standalone operation ends.""" + + handler: SplunkAOBaseHandler + context: Context + tasks: set[str] = field(default_factory=set) + agent_run_ids: dict[tuple[str, str], UUID] = field(default_factory=dict) + active_tool_run_id: UUID | None = None + warned: bool = False + source: Any = None + + class CrewAIEventListener: """ CrewAI event listener for logging traces to the Splunk AO platform. @@ -78,6 +95,9 @@ def __init__( integration="crewai", ) self._active_tool_run_id: UUID | None = None + self._agent_run_ids: dict[tuple[str, str], UUID] = {} + self._crew_runs: dict[UUID, _CrewRun] = {} + self._standalone_runs: dict[UUID, _CrewRun] = {} if CREWAI_AVAILABLE: # Register with the event bus directly instead of inheriting from @@ -97,6 +117,135 @@ def __init__( litellm.success_callback = [] litellm.success_callback.append(self.lite_llm_usage_callback) + def _on(self, event_bus: Any, event_type: type) -> Callable: + """Use CrewAI's event loop instead of its unordered synchronous worker pool.""" + # Preserve registration for older event buses without an async event API. + if not callable(getattr(event_bus, "aemit", None)): + return event_bus.on(event_type) + + def register(callback: Callable[[Any, Any], None]) -> Callable: + @wraps(callback) + async def receive(source: Any, event: Any) -> None: + # Lifecycle mutations cannot yield; export waiting happens afterward. + flush_logger = self._dispatch_event(callback, source, event) + if flush_logger is not None: + if flush_logger._ingestion_hook: + # Legacy hooks flush the whole mutable trace list. Wait for + # active runs and do not yield while that list is flushed. + if not self._crew_runs and not self._standalone_runs: + flush_logger.flush() + else: + await flush_logger.async_flush() + + return event_bus.on(event_type)(receive) + + return register + + def _dispatch_event(self, callback: Callable[[Any, Any], None], source: Any, event: Any) -> SplunkAOLogger | None: + """Keep each crew's callbacks in one context, independent of delivery tasks.""" + name = callback.__name__ + run_id = self._generate_run_id(source, event) + runs = self._crew_runs + standalone_failure = False + crew_id: UUID | None + if name.startswith("on_crew_kickoff_"): + crew_id = self._to_uuid(run_id) + else: + task = getattr(event, "task", None) + agent = getattr(event, "agent", None) or getattr(task, "agent", None) + crew = getattr(agent, "crew", None) + crew_id = self._to_uuid(getattr(crew, "id", None)) + if crew_id is None: + task_id = str(getattr(event, "task_id", None) or getattr(task, "id", "")) + agent_id = str(getattr(event, "agent_id", None) or getattr(agent, "id", "")) + matches = [ + (registry, owner) + for registry in (self._crew_runs, self._standalone_runs) + for owner, run in registry.items() + if (task_id and task_id in run.tasks) + or (agent_id and any(key[0] == agent_id for key in run.agent_run_ids)) + ] + if len(matches) == 1: + runs, crew_id = matches[0] + elif not matches: + # An unlinked call is independent, even while a crew is active. + runs, crew_id = self._standalone_runs, self._to_uuid(run_id) + if crew_id not in runs and name == "on_llm_call_failed": + # Older failures have no call ID or messages; use only an unambiguous source. + owners = [ + owner + for owner, run in runs.items() + if run.source is source + and run.handler._root_node is not None + and run.handler._root_node.node_type == NodeType.LLM.value + ] + if len(owners) == 1: + crew_id, standalone_failure = owners[0], True + + if crew_id is None: + _logger.debug("Ignoring CrewAI callback without an identifiable crew: %s", name) + return None + + is_root = ( + standalone_failure + or name.startswith("on_crew_kickoff_") + or (runs is self._standalone_runs and crew_id == self._to_uuid(run_id)) + ) + terminal = is_root and not name.endswith("_started") + if is_root and name.endswith("_started") and crew_id not in runs: + handler = self._handler + if self._crew_runs or self._standalone_runs: + handler = SplunkAOBaseHandler( + splunk_ao_logger=handler._splunk_ao_logger, + start_new_trace=handler._start_new_trace, + flush_on_chain_end=handler._flush_on_chain_end, + integration="crewai", + ) + runs[crew_id] = _CrewRun(handler, copy_context(), source=source) + + run = runs.get(crew_id) + if run is None: + _logger.debug("Ignoring CrewAI callback without an active crew: %s", name) + return None + + previous_handler, previous_tool, previous_agents = self._handler, self._active_tool_run_id, self._agent_run_ids + self._handler, self._active_tool_run_id, self._agent_run_ids = ( + run.handler, + run.active_tool_run_id, + run.agent_run_ids, + ) + if name == "on_task_started": + run.tasks.add(str(getattr(getattr(event, "task", source), "id", ""))) + flush = run.handler._flush_on_chain_end + run.handler._flush_on_chain_end = False + try: + run.context.run(callback, source, event) + except Exception as exc: + if is_root: + # A root callback failure must not retain an unfinished invocation. + try: + run.context.run(run.handler._conclude_owned_state_on_failure) + finally: + run.context.run(run.handler._reset_handler_state) + runs.pop(crew_id, None) + raise + if not run.warned: + _logger.warning( + "CrewAI telemetry callback %s failed (%s); continuing the crew", name, type(exc).__name__ + ) + run.warned = True + finally: + run.handler._flush_on_chain_end = flush + run.active_tool_run_id = self._active_tool_run_id + self._handler, self._active_tool_run_id, self._agent_run_ids = ( + previous_handler, + previous_tool, + previous_agents, + ) + if terminal: + runs.pop(crew_id, None) + return run.handler._splunk_ao_logger if terminal and flush else None + def setup_listeners(self, crewai_event_bus: Any) -> None: """Setup event listeners for CrewAI events.""" if not CREWAI_AVAILABLE: @@ -179,108 +328,108 @@ def setup_listeners(self, crewai_event_bus: Any) -> None: MemoryRetrievalFailedEvent = None # Crew event handlers - @crewai_event_bus.on(CrewKickoffStartedEvent) + @self._on(crewai_event_bus, CrewKickoffStartedEvent) def on_crew_kickoff_started(source: Any, event: CrewKickoffStartedEvent) -> None: self._handle_crew_kickoff_started(source, event) - @crewai_event_bus.on(CrewKickoffCompletedEvent) + @self._on(crewai_event_bus, CrewKickoffCompletedEvent) def on_crew_kickoff_completed(source: Any, event: CrewKickoffCompletedEvent) -> None: self._handle_crew_kickoff_completed(source, event) - @crewai_event_bus.on(CrewKickoffFailedEvent) + @self._on(crewai_event_bus, CrewKickoffFailedEvent) def on_crew_kickoff_failed(source: Any, event: CrewKickoffFailedEvent) -> None: self._handle_crew_kickoff_failed(source, event) # Agent event handlers - @crewai_event_bus.on(AgentExecutionStartedEvent) + @self._on(crewai_event_bus, AgentExecutionStartedEvent) def on_agent_execution_started(source: Any, event: AgentExecutionStartedEvent) -> None: self._handle_agent_execution_started(source, event) - @crewai_event_bus.on(AgentExecutionCompletedEvent) + @self._on(crewai_event_bus, AgentExecutionCompletedEvent) def on_agent_execution_completed(source: Any, event: AgentExecutionCompletedEvent) -> None: self._handle_agent_execution_completed(source, event) - @crewai_event_bus.on(AgentExecutionErrorEvent) + @self._on(crewai_event_bus, AgentExecutionErrorEvent) def on_agent_execution_error(source: Any, event: AgentExecutionErrorEvent) -> None: self._handle_agent_execution_error(source, event) # Task event handlers - @crewai_event_bus.on(TaskStartedEvent) + @self._on(crewai_event_bus, TaskStartedEvent) def on_task_started(source: Any, event: TaskStartedEvent) -> None: self._handle_task_started(source, event) - @crewai_event_bus.on(TaskCompletedEvent) + @self._on(crewai_event_bus, TaskCompletedEvent) def on_task_completed(source: Any, event: TaskCompletedEvent) -> None: self._handle_task_completed(source, event) - @crewai_event_bus.on(TaskFailedEvent) + @self._on(crewai_event_bus, TaskFailedEvent) def on_task_failed(source: Any, event: TaskFailedEvent) -> None: self._handle_task_failed(source, event) # Tool event handlers - @crewai_event_bus.on(ToolUsageStartedEvent) + @self._on(crewai_event_bus, ToolUsageStartedEvent) def on_tool_usage_started(source: Any, event: ToolUsageStartedEvent) -> None: self._handle_tool_usage_started(source, event) - @crewai_event_bus.on(ToolUsageFinishedEvent) + @self._on(crewai_event_bus, ToolUsageFinishedEvent) def on_tool_usage_finished(source: Any, event: ToolUsageFinishedEvent) -> None: self._handle_tool_usage_finished(source, event) - @crewai_event_bus.on(ToolUsageErrorEvent) + @self._on(crewai_event_bus, ToolUsageErrorEvent) def on_tool_usage_error(source: Any, event: ToolUsageErrorEvent) -> None: self._handle_tool_usage_error(source, event) # LLM event handlers - @crewai_event_bus.on(LLMCallStartedEvent) + @self._on(crewai_event_bus, LLMCallStartedEvent) def on_llm_call_started(source: Any, event: LLMCallStartedEvent) -> None: self._handle_llm_call_started(source, event) - @crewai_event_bus.on(LLMCallCompletedEvent) + @self._on(crewai_event_bus, LLMCallCompletedEvent) def on_llm_call_completed(source: Any, event: LLMCallCompletedEvent) -> None: self._handle_llm_call_completed(source, event) - @crewai_event_bus.on(LLMCallFailedEvent) + @self._on(crewai_event_bus, LLMCallFailedEvent) def on_llm_call_failed(source: Any, event: LLMCallFailedEvent) -> None: self._handle_llm_call_failed(source, event) # Memory event handlers (only available in CrewAI >= 0.177.0) if CREWAI_EVENTS_MODULE_AVAILABLE: - @crewai_event_bus.on(MemoryQueryStartedEvent) + @self._on(crewai_event_bus, MemoryQueryStartedEvent) def on_memory_query_started(source: Any, event: MemoryQueryStartedEvent) -> None: self._handle_memory_query_started(source, event) - @crewai_event_bus.on(MemoryQueryCompletedEvent) + @self._on(crewai_event_bus, MemoryQueryCompletedEvent) def on_memory_query_completed(source: Any, event: MemoryQueryCompletedEvent) -> None: self._handle_memory_query_completed(source, event) - @crewai_event_bus.on(MemoryQueryFailedEvent) + @self._on(crewai_event_bus, MemoryQueryFailedEvent) def on_memory_query_failed(source: Any, event: MemoryQueryFailedEvent) -> None: self._handle_memory_query_failed(source, event) - @crewai_event_bus.on(MemorySaveStartedEvent) + @self._on(crewai_event_bus, MemorySaveStartedEvent) def on_memory_save_started(source: Any, event: MemorySaveStartedEvent) -> None: self._handle_memory_save_started(source, event) - @crewai_event_bus.on(MemorySaveCompletedEvent) + @self._on(crewai_event_bus, MemorySaveCompletedEvent) def on_memory_save_completed(source: Any, event: MemorySaveCompletedEvent) -> None: self._handle_memory_save_completed(source, event) - @crewai_event_bus.on(MemorySaveFailedEvent) + @self._on(crewai_event_bus, MemorySaveFailedEvent) def on_memory_save_failed(source: Any, event: MemorySaveFailedEvent) -> None: self._handle_memory_save_failed(source, event) - @crewai_event_bus.on(MemoryRetrievalStartedEvent) + @self._on(crewai_event_bus, MemoryRetrievalStartedEvent) def on_memory_retrieval_started(source: Any, event: MemoryRetrievalStartedEvent) -> None: self._handle_memory_retrieval_started(source, event) - @crewai_event_bus.on(MemoryRetrievalCompletedEvent) + @self._on(crewai_event_bus, MemoryRetrievalCompletedEvent) def on_memory_retrieval_completed(source: Any, event: MemoryRetrievalCompletedEvent) -> None: self._handle_memory_retrieval_completed(source, event) if MemoryRetrievalFailedEvent is not None: - @crewai_event_bus.on(MemoryRetrievalFailedEvent) + @self._on(crewai_event_bus, MemoryRetrievalFailedEvent) def on_memory_retrieval_failed(source: Any, event: MemoryRetrievalFailedEvent) -> None: self._handle_memory_retrieval_failed(source, event) @@ -355,6 +504,30 @@ def _generate_run_id(self, source: Any, event: Any) -> UUID: f"{getattr(event, 'crew_name', '')}_{getattr(event, 'agent', '')}_{getattr(event, 'task', '')}" ) + @staticmethod + def _event_timestamp(event: Any) -> datetime | None: + timestamp = getattr(event, "timestamp", None) + if not isinstance(timestamp, datetime): + return None + return timestamp.replace(tzinfo=UTC) if timestamp.tzinfo is None else timestamp.astimezone(UTC) + + def _start_node(self, event: Any, **kwargs: Any) -> None: + timestamp = self._event_timestamp(event) + if timestamp is not None: + existing = self._handler.get_node(kwargs["run_id"]) + kwargs["created_at"] = existing.span_params.get("created_at", timestamp) if existing else timestamp + self._handler.start_node(**kwargs) + + def _end_node(self, event: Any, *, run_id: UUID, **kwargs: Any) -> None: + timestamp = self._event_timestamp(event) + node = self._handler.get_node(run_id) + created_at = node.span_params.get("created_at") if node else None + if timestamp is not None and isinstance(created_at, datetime): + if created_at.tzinfo is None: + created_at = created_at.replace(tzinfo=UTC) + kwargs["duration_ns"] = max(0, convert_time_delta_to_ns(timestamp - created_at)) + self._handler.end_node(run_id=run_id, **kwargs) + def _extract_metadata(self, event: Any) -> dict: """Extract metadata from event for span attributes.""" metadata = {} @@ -380,7 +553,8 @@ def _handle_crew_kickoff_started(self, source: Any, event: Any) -> None: input = getattr(event, "inputs", {}) - self._handler.start_node( + self._start_node( + event, node_type=NodeType.CHAIN.value, parent_run_id=None, # Root node run_id=run_id, @@ -415,7 +589,8 @@ def _handle_crew_kickoff_completed(self, source: Any, event: Any) -> None: run_id = self._generate_run_id(source, event) output = getattr(event, "output", {}) self._update_crew_input(str(run_id)) - self._handler.end_node(run_id=run_id, output=serialize_to_str(getattr(output, "raw", output))) + self._end_node(event, run_id=run_id, output=serialize_to_str(getattr(output, "raw", output))) + self._agent_run_ids.clear() def _handle_crew_kickoff_failed(self, source: Any, event: Any) -> None: """Handle crew execution failure.""" @@ -423,13 +598,17 @@ def _handle_crew_kickoff_failed(self, source: Any, event: Any) -> None: metadata = self._extract_metadata(event) metadata["error"] = event.error - self._handler.end_node(run_id=run_id, output=f"Error: {event.error}", metadata=metadata) + self._end_node(event, run_id=run_id, output=f"Error: {event.error}", metadata=metadata) + self._agent_run_ids.clear() # Agent event handlers def _handle_agent_execution_started(self, source: Any, event: Any) -> None: """Handle agent execution start.""" run_id = self._generate_run_id(source, event) event_task = getattr(event, "task", None) + agent_id = str(getattr(getattr(event, "agent", None), "id", "")) + task_id = str(getattr(event_task, "id", "")) + self._agent_run_ids[agent_id, task_id] = run_id parent_run_id = self._to_uuid(getattr(event_task, "id", None)) if event_task else None # When a delegation tool creates a temporary task, that task has no node. @@ -449,7 +628,8 @@ def _handle_agent_execution_started(self, source: Any, event: Any) -> None: if hasattr(event, "tools") and event.tools: metadata["available_tools"] = [str(getattr(tool, "name", tool)) for tool in event.tools] - self._handler.start_node( + self._start_node( + event, node_type=NodeType.AGENT.value, parent_run_id=parent_run_id, run_id=run_id, @@ -461,7 +641,7 @@ def _handle_agent_execution_started(self, source: Any, event: Any) -> None: def _handle_agent_execution_completed(self, source: Any, event: Any) -> None: """Handle agent execution completion.""" run_id = self._generate_run_id(source, event) - self._handler.end_node(run_id=run_id, output=serialize_to_str(getattr(event, "output", ""))) + self._end_node(event, run_id=run_id, output=serialize_to_str(getattr(event, "output", ""))) def _handle_agent_execution_error(self, source: Any, event: Any) -> None: """Handle agent execution error.""" @@ -469,8 +649,8 @@ def _handle_agent_execution_error(self, source: Any, event: Any) -> None: metadata = self._extract_metadata(event) metadata["error"] = getattr(event, "error", "Unknown error") - self._handler.end_node( - run_id=run_id, output=f"Error: {getattr(event, 'error', 'Unknown error')}", metadata=metadata + self._end_node( + event, run_id=run_id, output=f"Error: {getattr(event, 'error', 'Unknown error')}", metadata=metadata ) # Task event handlers @@ -498,7 +678,8 @@ def _handle_task_started(self, source: Any, event: Any) -> None: if not input: input = task.description if task and hasattr(task, "description") else "" - self._handler.start_node( + self._start_node( + event, node_type=NodeType.CHAIN.value, parent_run_id=parent_run_id, run_id=run_id, @@ -514,7 +695,7 @@ def _handle_task_completed(self, source: Any, event: Any) -> None: if hasattr(event, "output") and event.output: output = event.output.raw if hasattr(event.output, "raw") else str(event.output) - self._handler.end_node(run_id=run_id, output=serialize_to_str(output)) + self._end_node(event, run_id=run_id, output=serialize_to_str(output)) def _handle_task_failed(self, source: Any, event: Any) -> None: """Handle task failure.""" @@ -522,7 +703,7 @@ def _handle_task_failed(self, source: Any, event: Any) -> None: metadata = self._extract_metadata(event) metadata["error"] = event.error - self._handler.end_node(run_id=run_id, output=f"Error: {event.error}", metadata=metadata) + self._end_node(event, run_id=run_id, output=f"Error: {event.error}", metadata=metadata) # Tool event handlers def _handle_tool_usage_started(self, source: Any, event: Any) -> None: @@ -547,7 +728,8 @@ def _handle_tool_usage_started(self, source: Any, event: Any) -> None: if input: metadata["tool_args"] = str(input) - self._handler.start_node( + self._start_node( + event, node_type=NodeType.TOOL.value, parent_run_id=parent_run_id, run_id=run_id, @@ -560,7 +742,7 @@ def _handle_tool_usage_started(self, source: Any, event: Any) -> None: def _handle_tool_usage_finished(self, source: Any, event: Any) -> None: """Handle tool usage completion.""" run_id = self._generate_run_id(source, event) - self._handler.end_node(run_id=run_id, output=serialize_to_str(getattr(event, "output", ""))) + self._end_node(event, run_id=run_id, output=serialize_to_str(getattr(event, "output", ""))) self._active_tool_run_id = None def _handle_tool_usage_error(self, source: Any, event: Any) -> None: @@ -569,8 +751,8 @@ def _handle_tool_usage_error(self, source: Any, event: Any) -> None: metadata = self._extract_metadata(event) metadata["error"] = getattr(event, "error", "Unknown error") - self._handler.end_node( - run_id=run_id, output=f"Error: {getattr(event, 'error', 'Unknown error')}", metadata=metadata + self._end_node( + event, run_id=run_id, output=f"Error: {getattr(event, 'error', 'Unknown error')}", metadata=metadata ) self._active_tool_run_id = None @@ -588,7 +770,9 @@ def _to_uuid(self, id: str | None | UUID) -> UUID | None: def _handle_llm_call_started(self, source: Any, event: Any) -> None: """Handle LLM call start.""" run_id = self._generate_run_id(source, event) - parent_run_id = self._to_uuid(getattr(event, "agent_id", None)) + parent_run_id = self._agent_run_ids.get( + (str(getattr(event, "agent_id", "")), str(getattr(event, "task_id", ""))) + ) or self._to_uuid(getattr(event, "agent_id", None)) llm_name = getattr(event, "model", None) or getattr(source, "model", "Unknown Model") metadata = self._extract_metadata(event) @@ -596,7 +780,8 @@ def _handle_llm_call_started(self, source: Any, event: Any) -> None: if hasattr(source, "temperature"): metadata["temperature"] = str(source.temperature) - self._handler.start_node( + self._start_node( + event, node_type=NodeType.LLM.value, parent_run_id=parent_run_id, run_id=run_id, @@ -624,16 +809,19 @@ def _handle_llm_call_completed(self, source: Any, event: Any) -> None: token_kwargs["num_output_tokens"] = _get("completion_tokens", 0) token_kwargs["total_tokens"] = _get("total_tokens", 0) - self._handler.end_node(run_id=run_id, output=serialize_to_str(response or ""), **token_kwargs) + self._end_node(event, run_id=run_id, output=serialize_to_str(response or ""), **token_kwargs) def _handle_llm_call_failed(self, source: Any, event: Any) -> None: """Handle LLM call failure.""" run_id = self._generate_run_id(source, event) + root = self._handler._root_node + if self._handler.get_node(run_id) is None and root is not None and root.node_type == NodeType.LLM.value: + run_id = root.run_id metadata = self._extract_metadata(event) metadata["error"] = getattr(event, "error", "Unknown error") - self._handler.end_node( - run_id=run_id, output=f"Error: {getattr(event, 'error', 'Unknown error')}", metadata=metadata + self._end_node( + event, run_id=run_id, output=f"Error: {getattr(event, 'error', 'Unknown error')}", metadata=metadata ) # Memory event handlers @@ -649,7 +837,8 @@ def _handle_memory_query_started(self, source: Any, event: Any) -> None: task_run_id = self._to_uuid(task_id) parent_node = self._handler.get_node(task_run_id) if task_run_id else None if parent_node: - self._handler.start_node( + self._start_node( + event, node_type=NodeType.AGENT.value, parent_run_id=task_run_id, run_id=parent_run_id, @@ -665,7 +854,8 @@ def _handle_memory_query_started(self, source: Any, event: Any) -> None: if hasattr(event, "agent_role") and event.agent_role: metadata["agent_role"] = event.agent_role - self._handler.start_node( + self._start_node( + event, node_type=NodeType.RETRIEVER.value, parent_run_id=parent_run_id, run_id=run_id, @@ -683,7 +873,7 @@ def _handle_memory_query_completed(self, source: Any, event: Any) -> None: metadata["results_count"] = len(event.results) if hasattr(event.results, "__len__") else 0 output = serialize_to_str(event.results) if event.results else "No results found" - self._handler.end_node(run_id=run_id, output=output, metadata=metadata) + self._end_node(event, run_id=run_id, output=output, metadata=metadata) def _handle_memory_query_failed(self, source: Any, event: Any) -> None: """Handle memory query failure.""" @@ -692,7 +882,7 @@ def _handle_memory_query_failed(self, source: Any, event: Any) -> None: metadata = self._extract_metadata(event) metadata["error"] = event.error - self._handler.end_node(run_id=run_id, output=f"Memory query failed: {event.error}", metadata=metadata) + self._end_node(event, run_id=run_id, output=f"Memory query failed: {event.error}", metadata=metadata) def _handle_memory_save_started(self, source: Any, event: Any) -> None: """Handle memory save start.""" @@ -707,7 +897,8 @@ def _handle_memory_save_started(self, source: Any, event: Any) -> None: input_value = event.value if event.value else "Memory content" - self._handler.start_node( + self._start_node( + event, node_type=NodeType.CHAIN.value, parent_run_id=parent_run_id, run_id=run_id, @@ -723,7 +914,7 @@ def _handle_memory_save_completed(self, source: Any, event: Any) -> None: metadata = self._extract_metadata(event) metadata["save_time_ms"] = event.save_time_ms - self._handler.end_node(run_id=run_id, output=f"Memory saved successfully: {event.value}", metadata=metadata) + self._end_node(event, run_id=run_id, output=f"Memory saved successfully: {event.value}", metadata=metadata) def _handle_memory_save_failed(self, source: Any, event: Any) -> None: """Handle memory save failure.""" @@ -732,7 +923,7 @@ def _handle_memory_save_failed(self, source: Any, event: Any) -> None: metadata = self._extract_metadata(event) metadata["error"] = event.error - self._handler.end_node(run_id=run_id, output=f"Memory save failed: {event.error}", metadata=metadata) + self._end_node(event, run_id=run_id, output=f"Memory save failed: {event.error}", metadata=metadata) def _handle_memory_retrieval_started(self, source: Any, event: Any) -> None: """Handle memory retrieval start.""" @@ -743,7 +934,8 @@ def _handle_memory_retrieval_started(self, source: Any, event: Any) -> None: if hasattr(event, "task_id") and event.task_id: metadata["task_id"] = event.task_id - self._handler.start_node( + self._start_node( + event, node_type=NodeType.CHAIN.value, parent_run_id=parent_run_id, run_id=run_id, @@ -759,7 +951,7 @@ def _handle_memory_retrieval_completed(self, source: Any, event: Any) -> None: metadata = self._extract_metadata(event) metadata["retrieval_time_ms"] = event.retrieval_time_ms - self._handler.end_node(run_id=run_id, output=serialize_to_str(event.memory_content), metadata=metadata) + self._end_node(event, run_id=run_id, output=serialize_to_str(event.memory_content), metadata=metadata) def _handle_memory_retrieval_failed(self, source: Any, event: Any) -> None: """Handle memory retrieval failure.""" @@ -768,7 +960,7 @@ def _handle_memory_retrieval_failed(self, source: Any, event: Any) -> None: metadata = self._extract_metadata(event) metadata["error"] = getattr(event, "error", "Unknown error") - self._handler.end_node(run_id=run_id, output=f"Memory retrieval failed: {metadata['error']}", metadata=metadata) + self._end_node(event, run_id=run_id, output=f"Memory retrieval failed: {metadata['error']}", metadata=metadata) def lite_llm_usage_callback( self, @@ -780,11 +972,17 @@ def lite_llm_usage_callback( node_id = self._generate_run_id(kwargs, kwargs) node = self._handler.get_node(node_id) + if node is None: + for run in (*tuple(self._crew_runs.values()), *tuple(self._standalone_runs.values())): + node = run.handler.get_node(node_id) + if node is not None: + break if not node: _logger.debug(f"No node exists for run_id {node_id}") return usage = completion_response.model_extra["usage"] node.span_params["usage"] = usage.model_dump() if hasattr(usage, "model_dump") else usage - node.span_params["num_input_tokens"] = getattr(usage, "prompt_tokens", 0) - node.span_params["num_output_tokens"] = getattr(usage, "completion_tokens", 0) - node.span_params["total_tokens"] = getattr(usage, "total_tokens", 0) + get_usage = usage.get if isinstance(usage, dict) else lambda key, default=0: getattr(usage, key, default) + node.span_params["num_input_tokens"] = get_usage("prompt_tokens", 0) + node.span_params["num_output_tokens"] = get_usage("completion_tokens", 0) + node.span_params["total_tokens"] = get_usage("total_tokens", 0) diff --git a/tests/test_crewai_handler.py b/tests/test_crewai_handler.py index 96d36e10..b4b8895f 100644 --- a/tests/test_crewai_handler.py +++ b/tests/test_crewai_handler.py @@ -1,10 +1,16 @@ +import asyncio +import inspect import sys +import threading import uuid -from datetime import datetime +from contextvars import copy_context +from datetime import UTC, datetime, timedelta +from types import ModuleType from typing import Any from unittest.mock import Mock, patch import pytest +from opentelemetry import trace from opentelemetry.sdk.trace import ReadableSpan # Skip all tests in this module on Python 3.14+ (crewai doesn't support it yet) @@ -94,6 +100,780 @@ def shutdown(self) -> None: return None +@pytest.fixture +def crewai_bus(): + # Given: synchronous registrations are reversed so reverting async registration fails the regression. + # Async listeners run without yielding, as in CrewAI's dedicated event loop. + modules = {} + events = {} + for group, names in { + "crew_events": ["CrewKickoffStartedEvent", "CrewKickoffCompletedEvent", "CrewKickoffFailedEvent"], + "task_events": ["TaskStartedEvent", "TaskCompletedEvent", "TaskFailedEvent"], + "agent_events": ["AgentExecutionStartedEvent", "AgentExecutionCompletedEvent", "AgentExecutionErrorEvent"], + "llm_events": ["LLMCallStartedEvent", "LLMCallCompletedEvent", "LLMCallFailedEvent"], + "tool_usage_events": ["ToolUsageStartedEvent", "ToolUsageFinishedEvent", "ToolUsageErrorEvent"], + }.items(): + module = ModuleType(f"crewai.utilities.events.{group}") + for name in names: + events[name] = type(name, (MockEvent,), {}) + setattr(module, name, events[name]) + modules[module.__name__] = module + + callbacks = {} + bus = Mock(spec=["on", "aemit"]) + + def on(event_type): + def register(callback): + callbacks[event_type] = callback + return callback + + return register + + bus.on.side_effect = on + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + with ( + patch("splunk_ao.handlers.crewai.handler._crewai_imports_resolved", True), + patch("splunk_ao.handlers.crewai.handler.CREWAI_AVAILABLE", False), + patch("splunk_ao.handlers.crewai.handler.LITE_LLM_AVAILABLE", False), + ): + listener = CrewAIEventListener(splunk_ao_logger=logger) + with ( + patch.dict(sys.modules, modules), + patch("splunk_ao.handlers.crewai.handler.CREWAI_AVAILABLE", True), + patch("splunk_ao.handlers.crewai.handler.CREWAI_EVENTS_MODULE_AVAILABLE", False), + ): + listener.setup_listeners(bus) + + async def deliver(batch): + pending = [] + for source, event in batch: + callback = callbacks[type(event)] + if inspect.iscoroutinefunction(callback): + # Each event is delivered in a distinct task/copied context. + pending.append(asyncio.create_task(callback(source, event))) + else: + pending.append((callback, source, event)) + if pending and isinstance(pending[0], tuple): + for callback, source, event in reversed(pending): + copy_context().run(callback, source, event) + else: + await asyncio.gather(*pending) + + try: + yield listener, logger, sink, events, deliver + finally: + logger.terminate() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("legacy_agent_source", [False, True]) +def test_crewai_event_bus_preserves_callback_order_and_context(crewai_bus, asynchronous, legacy_agent_source) -> None: + listener, logger, sink, events, deliver = crewai_bus + crew = MockCrew() + agent = MockAgent(crew=crew) + task = MockTask(agent=agent) + agent_source = MockSource() if legacy_agent_source else agent + + async def run(turns): + for turn in turns: + call_id = f"call-{turn}" + await deliver( + [ + (crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={"turn": turn})), + (task, events["TaskStartedEvent"](task=task)), + ( + agent_source, + events["AgentExecutionStartedEvent"](agent=agent, task=task, task_prompt="Question"), + ), + ( + MockSource(), + events["LLMCallStartedEvent"]( + call_id=call_id, + model="local-model", + agent_id=str(agent.id), + task_id=str(task.id), + messages=[], + ), + ), + ( + MockSource(), + events["LLMCallCompletedEvent"]( + call_id=call_id, agent_id=str(agent.id), task_id=str(task.id), response=f"Reply {turn}" + ), + ), + ( + agent_source, + events["AgentExecutionCompletedEvent"](agent=agent, task=task, output=f"Reply {turn}"), + ), + (task, events["TaskCompletedEvent"](task=task, output=MockOutput(f"Reply {turn}"))), + (crew, events["CrewKickoffCompletedEvent"](output=MockOutput(f"Reply {turn}"))), + ] + ) + # Then: one complete tree, clean state, and no context leaked into the caller. + spans = sink.spans[turn * 4 :] + assert len(spans) == 4 + by_operation = {span.attributes["gen_ai.operation.name"]: span for span in spans} + llm = by_operation["chat"] + agent_span = by_operation["invoke_agent"] + task_span, root = [span for span in spans if span.attributes["gen_ai.operation.name"] == "invoke_workflow"] + assert llm.parent == agent_span.context + assert agent_span.parent == task_span.context + assert task_span.parent == root.context + assert logger.current_parent() is None + assert not trace.get_current_span().get_span_context().is_valid + assert not listener._handler.get_nodes() + assert not listener._handler._active_steps + assert not listener._crew_runs + + try: + # When: the same listener/crew is reused from sync or async application code. + if asynchronous: + asyncio.run(run(range(3))) + else: + for turn in range(3): + asyncio.run(run([turn])) + assert len({span.context.trace_id for span in sink.spans}) == 3 + finally: + logger.terminate() + + +def test_crewai_event_bus_isolates_overlapping_crews(crewai_bus) -> None: + # Given: two crews sharing one listener and logger, with interleaved callbacks. + listener, logger, sink, events, deliver = crewai_bus + crews = [MockCrew(), MockCrew()] + agents = [MockAgent(crew=crew) for crew in crews] + tasks = [MockTask(agent=agent) for agent in agents] + starts, ends = [], [] + for index, (crew, agent, task) in enumerate(zip(crews, agents, tasks, strict=True)): + starts.extend( + [ + (crew, events["CrewKickoffStartedEvent"](crew_name=f"Crew {index}", inputs={})), + (task, events["TaskStartedEvent"](task=task)), + (agent, events["AgentExecutionStartedEvent"](agent=agent, task=task, task_prompt="Question")), + ( + MockSource(), + events["LLMCallStartedEvent"]( + call_id=str(index), + agent_id=str(agent.id), + task_id=str(task.id), + model="local-model", + messages=[], + ), + ), + ] + ) + ends.extend( + [ + ( + MockSource(), + events["LLMCallCompletedEvent"]( + call_id=str(index), agent_id=str(agent.id), task_id=str(task.id), response="Reply" + ), + ), + (agent, events["AgentExecutionCompletedEvent"](agent=agent, task=task, output="Reply")), + (task, events["TaskCompletedEvent"](task=task, output=MockOutput("Reply"))), + (crew, events["CrewKickoffCompletedEvent"](output=MockOutput("Reply"))), + ] + ) + + # When: both trees open before either completes. + asyncio.run(deliver(starts + ends)) + + # Then: each tree has its own trace and complete, consistent hierarchy. + assert len(sink.spans) == 8 + roots = [span for span in sink.spans if span.parent is None] + assert len(roots) == 2 + for root in roots: + tree = [span for span in sink.spans if span.context.trace_id == root.context.trace_id] + assert len(tree) == 4 + assert all(span.parent in [other.context for other in tree] for span in tree if span is not root) + assert not listener._crew_runs + assert not listener._handler._active_steps + assert logger.current_parent() is None + assert not trace.get_current_span().get_span_context().is_valid + + +def test_crewai_event_bus_failure_releases_unfinished_children(crewai_bus) -> None: + # Given: a crew fails while a child is unfinished. + listener, logger, sink, events, deliver = crewai_bus + crew = MockCrew() + task = MockTask(agent=MockAgent(crew=crew)) + asyncio.run( + deliver( + [ + (crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={})), + (task, events["TaskStartedEvent"](task=task)), + (crew, events["CrewKickoffFailedEvent"](error="synthetic failure")), + ] + ) + ) + + # Then: only the failed root is exported and no invocation/context remains. + assert len(sink.spans) == 1 + assert not listener._crew_runs + assert not listener._handler.get_nodes() + assert not listener._handler._active_steps + assert logger.current_parent() is None + assert not trace.get_current_span().get_span_context().is_valid + + +def test_crewai_event_bus_preserves_caller_owned_context(crewai_bus) -> None: + # Given: the application has an active operation when CrewAI emits its events. + listener, logger, sink, events, deliver = crewai_bus + listener._handler._start_new_trace = False + logger.start_trace(input="Request") + caller = logger.add_workflow_span(input="Request", name="Caller") + caller_context = trace.get_current_span().get_span_context() + crew = MockCrew() + task = MockTask(agent=MockAgent(crew=crew)) + try: + # When: callbacks finish in independent delivery tasks. + asyncio.run( + deliver( + [ + (crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={})), + (task, events["TaskStartedEvent"](task=task)), + (task, events["TaskCompletedEvent"](task=task, output=MockOutput("Reply"))), + (crew, events["CrewKickoffCompletedEvent"](output=MockOutput("Reply"))), + ] + ) + ) + + # Then: the subtree inherits the caller, and callback cleanup leaves it open. + assert len(sink.spans) == 2 + assert sink.spans[-1].parent == caller_context + assert logger.current_parent() is caller + assert trace.get_current_span().get_span_context() == caller_context + assert not listener._crew_runs + finally: + logger.conclude(output="Reply") + logger.conclude(output="Reply") + + +def test_crewai_event_bus_inherits_native_otel_parent(crewai_bus) -> None: + # Given: standalone handler ownership beneath a caller-owned native OTel span. + listener, logger, sink, events, deliver = crewai_bus + caller_context = trace.SpanContext(trace_id=123, span_id=456, is_remote=True, trace_flags=trace.TraceFlags(1)) + crew = MockCrew() + with trace.use_span(trace.NonRecordingSpan(caller_context)): + # When: the handler opens and concludes its own SDK trace. + asyncio.run( + deliver( + [ + (crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={})), + (crew, events["CrewKickoffCompletedEvent"](output=MockOutput("Reply"))), + ] + ) + ) + + # Then: it inherits W3C parentage and leaves the caller context active. + assert len(sink.spans) == 1 + assert sink.spans[0].parent == caller_context + assert trace.get_current_span().get_span_context() == caller_context + assert logger.current_parent() is None + assert not listener._crew_runs + + +def test_crewai_event_bus_callback_exception_releases_invocation(crewai_bus) -> None: + # Given: a malformed callback fails after the crew root was activated. + listener, logger, _sink, events, deliver = crewai_bus + crew = MockCrew() + asyncio.run(deliver([(crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={}))])) + + # When: processing the completion fails before it can conclude the root. + with patch.object(listener, "_handle_crew_kickoff_completed", side_effect=ValueError("synthetic event failure")): + with pytest.raises(ValueError, match="synthetic event failure"): + asyncio.run(deliver([(crew, events["CrewKickoffCompletedEvent"](output=MockOutput()))])) + + # Then: no pending callback state or caller context remains. + assert not listener._crew_runs + assert not listener._handler.get_nodes() + assert not listener._handler._active_steps + assert logger.current_parent() is None + assert not trace.get_current_span().get_span_context().is_valid + + +def test_crewai_event_bus_without_async_api_retains_registration(crewai_callback) -> None: + bus = Mock(spec=["on"]) + callback = Mock() + + crewai_callback._on(bus, MockEvent)(callback) + + bus.on.assert_called_once_with(MockEvent) + bus.on.return_value.assert_called_once_with(callback) + + +@pytest.mark.parametrize("phase", ["start", "end"]) +def test_crewai_review_child_callback_failure_keeps_crew(crewai_bus, caplog, enable_galileo_logging, phase) -> None: + listener, logger, sink, events, deliver = crewai_bus + crew = MockCrew() + agent = MockAgent(crew=crew) + task = MockTask(agent=agent) + source = MockSource() + + def llm_event(kind, call_id): + return events[kind]( + call_id=call_id, + task_id=str(task.id), + agent_id=str(agent.id), + model="local-model", + messages=[], + response="Reply", + ) + + async def run(): + await deliver( + [ + (crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={})), + (task, events["TaskStartedEvent"](task=task)), + (agent, events["AgentExecutionStartedEvent"](agent=agent, task=task)), + ] + ) + if phase == "end": + await deliver([(source, llm_event("LLMCallStartedEvent", "bad"))]) + method = "_handle_llm_call_started" if phase == "start" else "_handle_llm_call_completed" + kind = "LLMCallStartedEvent" if phase == "start" else "LLMCallCompletedEvent" + with patch.object(listener, method, side_effect=ValueError("private event payload")): + await deliver([(source, llm_event(kind, "bad")), (source, llm_event(kind, "bad"))]) + assert crew.id in listener._crew_runs + await deliver( + [ + (source, llm_event("LLMCallStartedEvent", "good")), + (source, llm_event("LLMCallCompletedEvent", "good")), + (agent, events["AgentExecutionCompletedEvent"](agent=agent, task=task, output="Reply")), + (task, events["TaskCompletedEvent"](task=task, output=MockOutput("Reply"))), + (crew, events["CrewKickoffCompletedEvent"](output=MockOutput("Reply"))), + ] + ) + + asyncio.run(run()) + assert len(sink.spans) == 4 + assert len({span.context.trace_id for span in sink.spans}) == 1 + assert sink.spans[-1].status.status_code != trace.StatusCode.ERROR + warnings = [record for record in caplog.records if record.name == "splunk_ao.handlers.crewai.handler"] + assert len(warnings) == 1 + assert "private event payload" not in caplog.text + assert not listener._crew_runs + assert not listener._handler._active_steps + assert not logger._otel_ids + + +@pytest.mark.parametrize("crew_active", [False, True]) +def test_crewai_review_standalone_llms_keep_separate_traces(crewai_bus, crew_active) -> None: + listener, logger, sink, events, deliver = crewai_bus + crew = MockCrew() + source = MockSource() + + async def run(): + if crew_active: + await deliver([(crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={}))]) + for index in range(2): + await deliver( + [ + ( + source, + events["LLMCallStartedEvent"](call_id=f"standalone-{index}", model="local-model", messages=[]), + ) + ] + ) + assert logger.current_parent() is None + assert not trace.get_current_span().get_span_context().is_valid + for index in range(2): + await deliver( + [(source, events["LLMCallCompletedEvent"](call_id=f"standalone-{index}", response=f"Reply {index}"))] + ) + if crew_active: + await deliver([(crew, events["CrewKickoffCompletedEvent"](output=MockOutput("Crew reply")))]) + + asyncio.run(run()) + assert len(sink.spans) == 2 + int(crew_active) + assert all(span.parent is None for span in sink.spans) + assert len({span.context.trace_id for span in sink.spans}) == len(sink.spans) + assert sum(span.attributes["gen_ai.operation.name"] == "chat" for span in sink.spans) == 2 + assert not listener._crew_runs + assert not listener._standalone_runs + assert not logger._otel_ids + + +@pytest.mark.parametrize("string_id", [False, True]) +def test_crewai_review_standalone_task_ids_can_be_reused(crewai_bus, string_id) -> None: + listener, logger, sink, events, deliver = crewai_bus + task_id = uuid.uuid4() + task = MockTask(task_id=str(task_id) if string_id else task_id) + + async def run(): + for _ in range(2): + await deliver( + [ + (task, events["TaskStartedEvent"](task=task)), + (task, events["TaskCompletedEvent"](task=task, output=MockOutput("Reply"))), + ] + ) + assert not listener._standalone_runs + assert logger.current_parent() is None + + asyncio.run(run()) + assert len(sink.spans) == 2 + assert all(span.parent is None for span in sink.spans) + assert len({span.context.trace_id for span in sink.spans}) == 2 + assert not logger._otel_ids + + +@pytest.mark.parametrize("call_id", [None, "standalone-failure"]) +def test_crewai_review_failed_standalone_llm_releases_state(crewai_bus, call_id) -> None: + listener, logger, sink, events, deliver = crewai_bus + source = MockSource(model="local-model") + correlation = {"call_id": call_id} if call_id else {} + # CrewAI 1.6 failure events omit messages and have no call_id. + asyncio.run( + deliver( + [ + ( + source, + events["LLMCallStartedEvent"]( + messages=[{"role": "user", "content": "Question"}], model="local-model", **correlation + ), + ), + (source, events["LLMCallFailedEvent"](error="synthetic failure", **correlation)), + ] + ) + ) + assert len(sink.spans) == 1 + assert sink.spans[0].attributes["gen_ai.operation.name"] == "chat" + assert "synthetic failure" in sink.spans[0].attributes["gen_ai.output.messages"] + assert not listener._standalone_runs + assert not logger._otel_ids + assert logger.current_parent() is None + + +def test_crewai_review_usage_reaches_overlapping_crews(crewai_bus) -> None: + listener, _logger, sink, events, deliver = crewai_bus + calls = [] + starts, ends = [], [] + for index in range(2): + crew = MockCrew() + agent = MockAgent(crew=crew) + task = MockTask(agent=agent) + source = MockSource() + messages = [{"role": "user", "content": f"Question {index}"}] + calls.append({"messages": messages}) + starts.extend( + [ + (crew, events["CrewKickoffStartedEvent"](crew_name=f"Crew {index}", inputs={})), + (task, events["TaskStartedEvent"](task=task)), + (agent, events["AgentExecutionStartedEvent"](agent=agent, task=task)), + ( + source, + events["LLMCallStartedEvent"]( + messages=messages, agent_id=str(agent.id), task_id=str(task.id), model="local-model" + ), + ), + ] + ) + ends.extend( + [ + ( + source, + events["LLMCallCompletedEvent"]( + messages=messages, agent_id=str(agent.id), task_id=str(task.id), response="Reply" + ), + ), + (agent, events["AgentExecutionCompletedEvent"](agent=agent, task=task, output="Reply")), + (task, events["TaskCompletedEvent"](task=task, output=MockOutput("Reply"))), + (crew, events["CrewKickoffCompletedEvent"](output=MockOutput("Reply"))), + ] + ) + + async def run(): + await deliver(starts) + for index, kwargs in enumerate(calls): + usage = Mock(prompt_tokens=10 + index, completion_tokens=5, total_tokens=15 + index) + usage.model_dump.return_value = { + "prompt_tokens": 10 + index, + "completion_tokens": 5, + "total_tokens": 15 + index, + } + listener.lite_llm_usage_callback( + kwargs, Mock(model_extra={"usage": usage}), datetime.now(UTC), datetime.now(UTC) + ) + await deliver(ends) + + asyncio.run(run()) + chats = [span for span in sink.spans if span.attributes["gen_ai.operation.name"] == "chat"] + assert len(sink.spans) == 8 + assert len(chats) == 2 + assert [span.attributes.get("gen_ai.usage.input_tokens") for span in chats] == [10, 11] + assert [span.attributes.get("gen_ai.usage.output_tokens") for span in chats] == [5, 5] + assert not listener._crew_runs + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_crewai_hook_flush_preserves_other_runs(crewai_bus, monkeypatch, overlap) -> None: + # Given: a legacy hook shares one mutable trace list across crew invocations. + listener, logger, _sink, events, deliver = crewai_bus + listener._handler._flush_on_chain_end = True + batches = [] + + def hook(request): + batches.append([(item.name, item.output) for item in request.traces]) + + monkeypatch.setattr(logger, "_ingestion_hook", hook) + crews = [MockCrew(name="First"), MockCrew(name="Second")] + starts = [(crew, events["CrewKickoffStartedEvent"](crew_name=crew.name, inputs={})) for crew in crews] + ends = [(crew, events["CrewKickoffCompletedEvent"](output=MockOutput(f"Reply {crew.name}"))) for crew in crews] + + async def run(): + # When: another run is open, or its start is already queued behind the first completion. + if overlap: + await deliver(starts) + await deliver([ends[0]]) + assert not batches + assert len(logger.traces) == 2 + await deliver([ends[1]]) + else: + await deliver([starts[0], ends[0], starts[1], ends[1]]) + + asyncio.run(run()) + + # Then: only completed traces reach the hook, exactly once, with neither run lost. + assert [item for batch in batches for item in batch] == [(crew.name, f"Reply {crew.name}") for crew in crews] + assert len(batches) == (1 if overlap else 2) + assert not logger.traces + assert not listener._crew_runs + assert not listener._standalone_runs + + +def test_crewai_review_flush_runs_off_event_loop(crewai_bus, monkeypatch) -> None: + listener, _logger, sink, events, deliver = crewai_bus + listener._handler._flush_on_chain_end = True + crew = MockCrew() + flush_threads = [] + loop_thread = threading.get_ident() + + def flush(): + flush_threads.append(threading.get_ident()) + return True + + monkeypatch.setattr(sink, "force_flush", flush) + asyncio.run( + deliver( + [ + (crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={})), + (crew, events["CrewKickoffCompletedEvent"](output=MockOutput("Reply"))), + ] + ) + ) + assert len(sink.spans) == 1 + assert len(flush_threads) == 1 + assert flush_threads[0] != loop_thread + assert listener._handler._flush_on_chain_end is True + + +@pytest.mark.parametrize("naive", [False, True]) +def test_crewai_review_uses_original_event_times(crewai_bus, naive) -> None: + listener, _logger, sink, events, deliver = crewai_bus + crew = MockCrew() + agent = MockAgent(crew=crew) + agent_source = MockSource() + task = MockTask(agent=agent) + base = datetime(2020, 1, 2, tzinfo=UTC) + event_base = base.replace(tzinfo=None) if naive else base + + def event(kind, seconds, **kwargs): + return events[kind](timestamp=event_base + timedelta(seconds=seconds), **kwargs) + + asyncio.run( + deliver( + [ + (crew, event("CrewKickoffStartedEvent", 0, crew_name="Crew", inputs={})), + (task, event("TaskStartedEvent", 1, task=task)), + (agent_source, event("AgentExecutionStartedEvent", 2, agent=agent, task=task)), + ( + MockSource(), + event( + "ToolUsageStartedEvent", + 3, + tool_name="lookup", + tool_args={}, + agent_id=str(agent.id), + task_id=str(task.id), + ), + ), + ( + MockSource(), + event( + "ToolUsageFinishedEvent", + 3.5, + tool_name="lookup", + tool_args={}, + agent_id=str(agent.id), + task_id=str(task.id), + output="result", + ), + ), + ( + MockSource(), + event( + "LLMCallStartedEvent", + 4, + call_id="timed", + model="local-model", + agent_id=str(agent.id), + task_id=str(task.id), + messages=[], + ), + ), + ( + MockSource(), + event( + "LLMCallCompletedEvent", + 5, + call_id="timed", + agent_id=str(agent.id), + task_id=str(task.id), + response="Reply", + ), + ), + (agent_source, event("AgentExecutionCompletedEvent", 6, agent=agent, task=task, output="Reply")), + (task, event("TaskCompletedEvent", 7, task=task, output=MockOutput("Reply"))), + (crew, event("CrewKickoffCompletedEvent", 8, output=MockOutput("Reply"))), + ] + ) + ) + expected = { + "invoke_workflow Crew": (0, 8), + "invoke_workflow Test task": (1, 7), + "invoke_agent Test Agent": (2, 6), + "execute_tool lookup": (3, 3.5), + "chat local-model": (4, 5), + } + assert len(sink.spans) == len(expected) + for span in sink.spans: + start, end = expected[span.name] + assert span.start_time == int((base + timedelta(seconds=start)).timestamp() * 1e9) + assert span.end_time - span.start_time == int((end - start) * 1e9) + assert not listener._crew_runs + + +def test_crewai_review_duplicate_events_emit_once(crewai_bus) -> None: + listener, logger, sink, events, deliver = crewai_bus + crew = MockCrew() + task = MockTask(agent=MockAgent(crew=crew)) + crew_start = (crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={})) + start = (task, events["TaskStartedEvent"](task=task)) + end = (task, events["TaskCompletedEvent"](task=task, output=MockOutput("Reply"))) + crew_end = (crew, events["CrewKickoffCompletedEvent"](output=MockOutput("Reply"))) + asyncio.run(deliver([crew_start, crew_start, start, start, end, end, crew_end, crew_end])) + assert len(sink.spans) == 2 + assert sink.spans[0].parent == sink.spans[1].context + assert not listener._crew_runs + assert not logger._otel_ids + + +@pytest.mark.parametrize("caller_owned", [False, True]) +def test_crewai_review_inherits_context_across_loop_thread(crewai_bus, caller_owned) -> None: + listener, logger, sink, events, deliver = crewai_bus + loop = asyncio.new_event_loop() + ready = threading.Event() + + def serve(): + asyncio.set_event_loop(loop) + loop.call_soon(ready.set) + loop.run_forever() + + worker = threading.Thread(target=serve, daemon=True) + worker.start() + crew = MockCrew() + remote = trace.SpanContext(123, 456, is_remote=True, trace_flags=trace.TraceFlags(1)) + try: + assert ready.wait(timeout=5) + with trace.use_span(trace.NonRecordingSpan(remote)): + if caller_owned: + listener._handler._start_new_trace = False + logger.start_trace(input="caller") + logger.add_workflow_span(input="caller", name="caller") + caller = logger.current_parent() + expected = trace.get_current_span().get_span_context() + for _ in range(2): + future = asyncio.run_coroutine_threadsafe( + deliver( + [ + (crew, events["CrewKickoffStartedEvent"](crew_name="Crew", inputs={})), + (crew, events["CrewKickoffCompletedEvent"](output=MockOutput("Reply"))), + ] + ), + loop, + ) + future.result(timeout=5) + assert logger.current_parent() is caller + assert trace.get_current_span().get_span_context() == expected + assert len(sink.spans) == 2 + assert all(span.parent == expected for span in sink.spans) + assert not listener._crew_runs + if caller_owned: + logger.conclude(output="Reply", conclude_all=True) + finally: + loop.call_soon_threadsafe(loop.stop) + worker.join(timeout=5) + assert not worker.is_alive() + loop.close() + + +def test_crewai_review_two_agent_turns_keep_all_twenty_one_spans(crewai_bus) -> None: + listener, logger, sink, events, deliver = crewai_bus + crew = MockCrew() + agents = [MockAgent(crew=crew, role=role) for role in ("Researcher", "Reviewer")] + tasks = [MockTask(agent=agent) for agent in agents] + + async def run(): + for turn in range(3): + batch = [(crew, events["CrewKickoffStartedEvent"](crew_name="LocalCrew", inputs={"turn": turn}))] + for index, (agent, task) in enumerate(zip(agents, tasks, strict=True)): + source = MockSource() + call = {"call_id": f"{turn}-{index}", "agent_id": str(agent.id), "task_id": str(task.id)} + batch.extend( + [ + (task, events["TaskStartedEvent"](task=task)), + (source, events["AgentExecutionStartedEvent"](agent=agent, task=task)), + (source, events["LLMCallStartedEvent"](**call, model="local-model", messages=[])), + (source, events["LLMCallCompletedEvent"](**call, response=f"Reply {turn}")), + ( + source, + events["AgentExecutionCompletedEvent"](agent=agent, task=task, output=f"Reply {turn}"), + ), + (task, events["TaskCompletedEvent"](task=task, output=MockOutput(f"Reply {turn}"))), + ] + ) + batch.append((crew, events["CrewKickoffCompletedEvent"](output=MockOutput(f"Reply {turn}")))) + before = len(sink.spans) + await deliver(batch) + spans = sink.spans[before:] + assert len(spans) == 7 + roots = [span for span in spans if span.parent is None] + assert len(roots) == 1 + by_id = {span.context.span_id: span for span in spans} + for span in spans: + operation = span.attributes["gen_ai.operation.name"] + if span is roots[0]: + assert operation == "invoke_workflow" + elif operation == "chat": + assert by_id[span.parent.span_id].attributes["gen_ai.operation.name"] == "invoke_agent" + elif operation == "invoke_agent": + assert by_id[span.parent.span_id].attributes["gen_ai.operation.name"] == "invoke_workflow" + else: + assert span.parent == roots[0].context + assert not listener._crew_runs + assert not listener._standalone_runs + assert not logger._otel_ids + assert logger.current_parent() is None + assert not trace.get_current_span().get_span_context().is_valid + + asyncio.run(run()) + assert len(sink.spans) == 21 + assert len({span.context.trace_id for span in sink.spans}) == 3 + + def test_crewai_adapter_enqueues_task_at_callback_end_without_flush() -> None: sink = RecordingSink() logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink)