diff --git a/lib/crewai/src/crewai/crew.py b/lib/crewai/src/crewai/crew.py index fda456be4..e643e562e 100644 --- a/lib/crewai/src/crewai/crew.py +++ b/lib/crewai/src/crewai/crew.py @@ -1048,8 +1048,9 @@ class Crew(FlowTrackable, BaseModel): ) raise finally: - if self._memory is not None and hasattr(self._memory, "drain_writes"): - self._memory.drain_writes() + # Safety net for the exception path; the success path already + # drained in _create_crew_output before emitting completion. + self._drain_memory_writes() clear_files(self.id) detach(token) crewai_event_bus._exit_runtime_scope(runtime_scope) @@ -1841,6 +1842,37 @@ class Crew(FlowTrackable, BaseModel): output=output.raw, ) + def _drain_memory_writes(self) -> None: + """Block until all pending background memory saves have completed. + + Covers both the crew memory and per-agent memories — agents save + through ``agent.memory`` when set (see + ``BaseAgentExecutor._save_to_memory``), so draining only + ``self._memory`` can miss in-flight saves. Scope/slice views are + unwrapped to their backing ``Memory`` so each pool is drained once. + + Must run before ``CrewKickoffCompletedEvent`` is emitted: listeners + (e.g. telemetry sessions) tear down on that event, and any + ``MemorySaveCompletedEvent``/``MemorySaveFailedEvent`` emitted after + teardown is lost, leaving the save span orphaned. + """ + seen: set[int] = set() + candidates = [ + self._memory, + self.memory, + *(getattr(agent, "memory", None) for agent in self.agents), + ] + for mem in candidates: + if mem is None or isinstance(mem, bool): + continue + backing = getattr(mem, "_memory", None) or mem + if id(backing) in seen: + continue + seen.add(id(backing)) + drain = getattr(backing, "drain_writes", None) + if callable(drain): + drain() + def _create_crew_output(self, task_outputs: list[TaskOutput]) -> CrewOutput: if not task_outputs: raise ValueError("No task outputs available to create crew output.") @@ -1853,6 +1885,10 @@ class Crew(FlowTrackable, BaseModel): final_string_output = final_task_output.raw self._finish_execution(final_string_output) self.token_usage = self.calculate_usage_metrics() + # Ensure background memory saves finish (and emit their + # completed/failed events) before the kickoff-completed event below + # triggers listener teardown/finalization. + self._drain_memory_writes() crewai_event_bus.flush() crewai_event_bus.emit( self, diff --git a/lib/crewai/src/crewai/flow/runtime/__init__.py b/lib/crewai/src/crewai/flow/runtime/__init__.py index 531aa9dbc..9aa88ee11 100644 --- a/lib/crewai/src/crewai/flow/runtime/__init__.py +++ b/lib/crewai/src/crewai/flow/runtime/__init__.py @@ -956,6 +956,22 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): return self.memory.remember_many(content, **kwargs) return self.memory.remember(content, **kwargs) + def _drain_memory_writes(self) -> None: + """Block until pending background memory saves for this flow finish. + + Must run before ``FlowFinishedEvent`` is emitted: listeners (e.g. + telemetry sessions) tear down on that event, and any + ``MemorySaveCompletedEvent``/``MemorySaveFailedEvent`` emitted after + teardown is lost, leaving the save span orphaned. + """ + mem = self.memory + if mem is None: + return + backing = getattr(mem, "_memory", None) or mem + drain = getattr(backing, "drain_writes", None) + if callable(drain): + drain() + def extract_memories(self, content: str) -> list[str]: """Extract discrete memories from content. Delegates to this flow's memory. @@ -1474,6 +1490,10 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): not self.suppress_flow_events and not self._should_defer_trace_finalization() ): + # Background memory saves must finish (and emit their + # completed/failed events) before flow-finished triggers + # listener teardown/finalization. + self._drain_memory_writes() future = crewai_event_bus.emit( self, FlowFinishedEvent( @@ -2285,6 +2305,10 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): # flag is read from either the instance attribute or an extension # definition. if not self._should_defer_trace_finalization(): + # Background memory saves must finish (and emit their + # completed/failed events) before flow-finished triggers + # listener teardown/finalization. + self._drain_memory_writes() future = crewai_event_bus.emit( self, FlowFinishedEvent( @@ -2317,9 +2341,9 @@ class Flow(BaseModel, Generic[T], metaclass=FlowMeta): return final_output finally: - # Ensure all background memory saves complete before returning - if self.memory is not None and hasattr(self.memory, "drain_writes"): - self.memory.drain_writes() + # Safety net for the exception path; the success path already + # drained before emitting FlowFinishedEvent. + self._drain_memory_writes() # Drain pending LLMCallCompletedEvent handlers before # detaching so `flow.usage_metrics` reflects every call # emitted during this kickoff — mirrors `Crew.kickoff()`, diff --git a/lib/crewai/tests/test_crew.py b/lib/crewai/tests/test_crew.py index 557f826d5..90d648711 100644 --- a/lib/crewai/tests/test_crew.py +++ b/lib/crewai/tests/test_crew.py @@ -4598,6 +4598,98 @@ def test_reset_memory_uses_full_unified_memory_reset(researcher): reset.assert_not_called() +def test_kickoff_drains_pending_memory_saves_before_completion_event(researcher): + """Background memory saves must finish (and emit their completion events) + before CrewKickoffCompletedEvent, otherwise listeners that tear down on + kickoff-completed (e.g. telemetry sessions) see the save span as orphaned.""" + import time + + from crewai.events.types.crew_events import CrewKickoffCompletedEvent + + order: list[str] = [] + + def slow_save(): + time.sleep(0.3) + order.append("save-done") + + crew = Crew( + agents=[researcher], + process=Process.sequential, + tasks=[ + Task(description="Task 1", expected_output="output", agent=researcher), + ], + memory=True, + task_callback=lambda _output: crew._memory._submit_save(slow_save), + ) + + completed = threading.Event() + + with crewai_event_bus.scoped_handlers(): + + @crewai_event_bus.on(CrewKickoffCompletedEvent) + def on_completed(_source, _event): + order.append("kickoff-completed") + completed.set() + + with patch.object(Agent, "execute_task", return_value="ok"): + crew.kickoff() + + assert completed.wait(timeout=5) + + assert order.index("save-done") < order.index("kickoff-completed") + + +def test_kickoff_drains_agent_memory_saves_before_completion_event(tmp_path): + """Agents save through their own ``agent.memory`` when set; those pools + must also be drained before CrewKickoffCompletedEvent.""" + import time + + from crewai.events.types.crew_events import CrewKickoffCompletedEvent + + agent_memory = Memory(storage=str(tmp_path / "agent-mem")) + agent_with_memory = Agent( + role="Researcher", + goal="Research things", + backstory="Experienced researcher", + memory=agent_memory, + ) + + order: list[str] = [] + + def slow_save(): + time.sleep(0.3) + order.append("save-done") + + crew = Crew( + agents=[agent_with_memory], + process=Process.sequential, + tasks=[ + Task( + description="Task 1", + expected_output="output", + agent=agent_with_memory, + ), + ], + task_callback=lambda _output: agent_memory._submit_save(slow_save), + ) + + completed = threading.Event() + + with crewai_event_bus.scoped_handlers(): + + @crewai_event_bus.on(CrewKickoffCompletedEvent) + def on_completed(_source, _event): + order.append("kickoff-completed") + completed.set() + + with patch.object(Agent, "execute_task", return_value="ok"): + crew.kickoff() + + assert completed.wait(timeout=5) + + assert order.index("save-done") < order.index("kickoff-completed") + + def test_reset_knowledge_with_only_crew_knowledge(researcher, writer): mock_ks = MagicMock(spec=Knowledge) diff --git a/lib/crewai/tests/test_flow.py b/lib/crewai/tests/test_flow.py index 06a420c4d..5e31f08ce 100644 --- a/lib/crewai/tests/test_flow.py +++ b/lib/crewai/tests/test_flow.py @@ -2353,3 +2353,41 @@ def test_locked_dict_proxy_ior(): def test_locked_dict_proxy_reversed(): flow = _make_dict_flow() assert list(reversed(flow.state.data)) == ["c", "b", "a"] + + +def test_flow_drains_pending_memory_saves_before_finished_event(tmp_path): + """Background memory saves must finish (and emit their completion events) + before FlowFinishedEvent, otherwise listeners that tear down on + flow-finished (e.g. telemetry sessions) see the save span as orphaned.""" + import time + + from crewai.memory.unified_memory import Memory + + order: list[str] = [] + + def slow_save(): + time.sleep(0.3) + order.append("save-done") + + class MemoryFlow(Flow): + @start() + def step_1(self): + self.memory._submit_save(slow_save) + return "done" + + flow = MemoryFlow(memory=Memory(storage=str(tmp_path / "flow-mem"))) + + finished = threading.Event() + + with crewai_event_bus.scoped_handlers(): + + @crewai_event_bus.on(FlowFinishedEvent) + def on_finished(_source, _event): + order.append("flow-finished") + finished.set() + + flow.kickoff() + + assert finished.wait(timeout=5) + + assert order.index("save-done") < order.index("flow-finished")