From 8f4ce0c0bbbad7cc2de68445e01d2f0b02fc05f1 Mon Sep 17 00:00:00 2001 From: pradystar Date: Thu, 22 Oct 2026 17:40:57 -0700 Subject: [PATCH 1/3] fix(crewai): prevent span loss from out-of-order callbacks --- CHANGELOG.md | 5 + src/splunk_ao/handlers/crewai/handler.py | 158 ++++++++++-- tests/test_crewai_handler.py | 309 +++++++++++++++++++++++ 3 files changed, 447 insertions(+), 25 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fec4614c..ad555970 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- CrewAI listeners preserve lifecycle callback order and per-crew context so repeated invocations retain + all spans and their parent relationships. + ## [0.6.0] ### Added diff --git a/src/splunk_ao/handlers/crewai/handler.py b/src/splunk_ao/handlers/crewai/handler.py index 6adc870c..2d73cb31 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 collections.abc import Callable +from contextvars import Context, copy_context +from dataclasses import dataclass, field from datetime import datetime +from functools import wraps from typing import Any from uuid import UUID @@ -53,6 +57,17 @@ def _resolve_crewai_imports() -> None: litellm = None +@dataclass +class _CrewRun: + """Callback state retained only until one crew invocation 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 + + class CrewAIEventListener: """ CrewAI event listener for logging traces to the Splunk AO platform. @@ -78,6 +93,8 @@ 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] = {} if CREWAI_AVAILABLE: # Register with the event bus directly instead of inheriting from @@ -97,6 +114,90 @@ 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: + # Deliberately no await: one lifecycle mutation must finish before + # the next event's callback can run on the bus's event loop. + self._dispatch_event(callback, source, event) + + return event_bus.on(event_type)(receive) + + return register + + def _dispatch_event(self, callback: Callable[[Any, Any], None], source: Any, event: Any) -> None: + """Keep each crew's callbacks in one context, independent of delivery tasks.""" + name = callback.__name__ + crew_id: UUID | None + if name.startswith("on_crew_kickoff_"): + crew_id = self._generate_run_id(source, event) + 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", "")) + matches = [run_id for run_id, run in self._crew_runs.items() if task_id in run.tasks] + if len(matches) == 1: + crew_id = matches[0] + elif len(self._crew_runs) == 1: + crew_id = next(iter(self._crew_runs)) + + if crew_id is None: + _logger.debug("Ignoring CrewAI callback without an identifiable crew: %s", name) + return + + if name == "on_crew_kickoff_started" and crew_id not in self._crew_runs: + handler = self._handler + if self._crew_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", + ) + self._crew_runs[crew_id] = _CrewRun(handler, copy_context()) + + run = self._crew_runs.get(crew_id) + if run is None: + _logger.debug("Ignoring CrewAI callback without an active crew: %s", name) + return + + 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", ""))) + try: + run.context.run(callback, source, event) + except Exception: + # A malformed event must not retain an unfinished invocation forever. + try: + run.context.run(run.handler._conclude_owned_state_on_failure) + finally: + run.context.run(run.handler._reset_handler_state) + self._crew_runs.pop(crew_id, None) + raise + finally: + 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 name in ("on_crew_kickoff_completed", "on_crew_kickoff_failed"): + self._crew_runs.pop(crew_id, None) + def setup_listeners(self, crewai_event_bus: Any) -> None: """Setup event listeners for CrewAI events.""" if not CREWAI_AVAILABLE: @@ -179,108 +280,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) @@ -416,6 +517,7 @@ def _handle_crew_kickoff_completed(self, source: Any, event: Any) -> None: 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._agent_run_ids.clear() def _handle_crew_kickoff_failed(self, source: Any, event: Any) -> None: """Handle crew execution failure.""" @@ -424,12 +526,16 @@ def _handle_crew_kickoff_failed(self, source: Any, event: Any) -> None: metadata["error"] = event.error self._handler.end_node(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. @@ -588,7 +694,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) diff --git a/tests/test_crewai_handler.py b/tests/test_crewai_handler.py index 96d36e10..4daf0938 100644 --- a/tests/test_crewai_handler.py +++ b/tests/test_crewai_handler.py @@ -1,10 +1,15 @@ +import asyncio +import inspect import sys import uuid +from contextvars import copy_context from datetime import datetime +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 +99,310 @@ def shutdown(self) -> None: return None +@pytest.fixture +def crewai_bus(): + # Given: a bus that can deliver synchronous callbacks in reverse worker order. + # 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() + + 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) + + 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) From 254fb2d59808ed2fa78ab2603d8a84792b71c90c Mon Sep 17 00:00:00 2001 From: pradystar Date: Sat, 24 Oct 2026 10:26:03 -0700 Subject: [PATCH 2/3] fix(crewai): preserve standalone calls, usage, timestamps, and callback isolation --- CHANGELOG.md | 4 +- src/splunk_ao/handlers/crewai/handler.py | 202 ++++++++--- tests/test_crewai_handler.py | 441 ++++++++++++++++++++++- 3 files changed, 584 insertions(+), 63 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ad555970..57a6c0d9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,9 @@ 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. + 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 completion flushes run outside the shared event loop when explicitly enabled. ## [0.6.0] diff --git a/src/splunk_ao/handlers/crewai/handler.py b/src/splunk_ao/handlers/crewai/handler.py index 2d73cb31..43a206d7 100644 --- a/src/splunk_ao/handlers/crewai/handler.py +++ b/src/splunk_ao/handlers/crewai/handler.py @@ -4,7 +4,7 @@ from collections.abc import Callable from contextvars import Context, copy_context from dataclasses import dataclass, field -from datetime import datetime +from datetime import UTC, datetime from functools import wraps from typing import Any from uuid import UUID @@ -14,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__) @@ -59,13 +59,15 @@ def _resolve_crewai_imports() -> None: @dataclass class _CrewRun: - """Callback state retained only until one crew invocation ends.""" + """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: @@ -95,6 +97,7 @@ def __init__( 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 @@ -123,20 +126,24 @@ def _on(self, event_bus: Any, event_type: type) -> Callable: def register(callback: Callable[[Any, Any], None]) -> Callable: @wraps(callback) async def receive(source: Any, event: Any) -> None: - # Deliberately no await: one lifecycle mutation must finish before - # the next event's callback can run on the bus's event loop. - self._dispatch_event(callback, source, event) + # Lifecycle mutations cannot yield; export waiting happens afterward. + flush_logger = self._dispatch_event(callback, source, event) + if flush_logger is not None: + 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) -> None: + 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._generate_run_id(source, event) + crew_id = self._to_uuid(run_id) else: task = getattr(event, "task", None) agent = getattr(event, "agent", None) or getattr(task, "agent", None) @@ -144,31 +151,56 @@ def _dispatch_event(self, callback: Callable[[Any, Any], None], source: Any, eve 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", "")) - matches = [run_id for run_id, run in self._crew_runs.items() if task_id in run.tasks] + 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: - crew_id = matches[0] - elif len(self._crew_runs) == 1: - crew_id = next(iter(self._crew_runs)) + 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 + return None - if name == "on_crew_kickoff_started" and crew_id not in self._crew_runs: + 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: + 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", ) - self._crew_runs[crew_id] = _CrewRun(handler, copy_context()) + runs[crew_id] = _CrewRun(handler, copy_context(), source=source) - run = self._crew_runs.get(crew_id) + run = runs.get(crew_id) if run is None: _logger.debug("Ignoring CrewAI callback without an active crew: %s", name) - return + 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 = ( @@ -178,25 +210,35 @@ def _dispatch_event(self, callback: Callable[[Any, Any], None], source: Any, eve ) 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: - # A malformed event must not retain an unfinished invocation forever. - try: - run.context.run(run.handler._conclude_owned_state_on_failure) - finally: - run.context.run(run.handler._reset_handler_state) - self._crew_runs.pop(crew_id, None) - raise + 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 name in ("on_crew_kickoff_completed", "on_crew_kickoff_failed"): - self._crew_runs.pop(crew_id, None) + 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.""" @@ -456,6 +498,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 = {} @@ -481,7 +547,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, @@ -516,7 +583,7 @@ 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: @@ -525,7 +592,7 @@ 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 @@ -555,7 +622,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, @@ -567,7 +635,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.""" @@ -575,8 +643,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 @@ -604,7 +672,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, @@ -620,7 +689,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.""" @@ -628,7 +697,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: @@ -653,7 +722,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, @@ -666,7 +736,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: @@ -675,8 +745,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 @@ -704,7 +774,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, @@ -732,16 +803,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 @@ -757,7 +831,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, @@ -773,7 +848,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, @@ -791,7 +867,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.""" @@ -800,7 +876,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.""" @@ -815,7 +891,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, @@ -831,7 +908,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.""" @@ -840,7 +917,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.""" @@ -851,7 +928,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, @@ -867,7 +945,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.""" @@ -876,7 +954,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, @@ -888,11 +966,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 4daf0938..125e413f 100644 --- a/tests/test_crewai_handler.py +++ b/tests/test_crewai_handler.py @@ -1,9 +1,10 @@ import asyncio import inspect import sys +import threading import uuid from contextvars import copy_context -from datetime import datetime +from datetime import UTC, datetime, timedelta from types import ModuleType from typing import Any from unittest.mock import Mock, patch @@ -101,7 +102,7 @@ def shutdown(self) -> None: @pytest.fixture def crewai_bus(): - # Given: a bus that can deliver synchronous callbacks in reverse worker order. + # 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 = {} @@ -119,7 +120,7 @@ def crewai_bus(): modules[module.__name__] = module callbacks = {} - bus = Mock() + bus = Mock(spec=["on", "aemit"]) def on(event_type): def register(callback): @@ -403,6 +404,440 @@ def test_crewai_event_bus_without_async_api_retains_registration(crewai_callback 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 + + +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) From 590e9362dbde897398abf60645b17e0e52559980 Mon Sep 17 00:00:00 2001 From: pradystar Date: Mon, 26 Oct 2026 09:41:14 -0700 Subject: [PATCH 3/3] "fix(crewai): guard legacy hook flushes across concurrent runs" --- CHANGELOG.md | 3 +- src/splunk_ao/handlers/crewai/handler.py | 8 +++++- tests/test_crewai_handler.py | 36 ++++++++++++++++++++++++ 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bcc0b379..b5d357c0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,7 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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 completion flushes run outside the shared event loop when explicitly enabled. +- 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 43a206d7..2374e48a 100644 --- a/src/splunk_ao/handlers/crewai/handler.py +++ b/src/splunk_ao/handlers/crewai/handler.py @@ -129,7 +129,13 @@ 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: - await flush_logger.async_flush() + 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) diff --git a/tests/test_crewai_handler.py b/tests/test_crewai_handler.py index 125e413f..b4b8895f 100644 --- a/tests/test_crewai_handler.py +++ b/tests/test_crewai_handler.py @@ -608,6 +608,42 @@ async def run(): 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