diff --git a/packages/tracecat-ee/tracecat_ee/agent/activities.py b/packages/tracecat-ee/tracecat_ee/agent/activities.py index 167972bb6..901a7db98 100644 --- a/packages/tracecat-ee/tracecat_ee/agent/activities.py +++ b/packages/tracecat-ee/tracecat_ee/agent/activities.py @@ -167,11 +167,11 @@ class EmitSessionCancelledInputs(BaseModel): curr_run_id: uuid.UUID | None = None active_stream_id: uuid.UUID | None = None emit_stream: bool = True - """Whether to also push the cancelled/done frames onto the live stream. + """Whether to also push the cancelled frame onto the live stream. - False when the executor loopback already emitted them (a second done - marker would race the client's stream teardown); the activity then only - persists the timeline marker row. + False when the executor loopback already emitted the notice; the activity + then only persists the timeline marker row. The workflow owns the terminal + END after finalizing the turn. """ @@ -533,7 +533,8 @@ async def emit_session_error(self, args: EmitSessionErrorInputs) -> None: signal the inbox reads) and, for pre-stream failures, also pushes the error onto the SSE stream since those happen before the loopback is wired up. The runtime path streams inline already and passes - ``should_stream=False`` to persist-only. + ``should_stream=False`` to persist-only. The workflow owns the terminal + END after finalizing the turn. Best-effort: a persistence failure must not mask the agent's real error or abort propagation, so it is logged and swallowed. @@ -565,7 +566,6 @@ async def emit_session_error(self, args: EmitSessionErrorInputs) -> None: stream = await self._open_session_stream(args) await stream.error(args.message) - await stream.done() @staticmethod async def _open_session_stream(args: _SessionStreamInputs) -> AgentStream: @@ -590,8 +590,9 @@ async def emit_session_cancelled(self, args: EmitSessionCancelledInputs) -> None Every cancelled turn persists a marker row so the "stopped by user" divider survives DB reloads. Stream emission is conditional: approval -wait cancels happen outside a running executor activity and must push - the cancelled/done frames here, while executor cancels already emitted - them from the loopback (``emit_stream=False``). + the cancelled frame here, while executor cancels already emitted it + from the loopback (``emit_stream=False``). The workflow owns the + terminal END after finalizing the turn. """ # Local import: tracecat.agent.session.service imports tracecat_ee # modules, so a top-level import here would create a cycle. @@ -616,7 +617,6 @@ async def emit_session_cancelled(self, args: EmitSessionCancelledInputs) -> None tool_call_ids=args.interrupted_tool_call_ids, ) ) - await stream.done() @activity.defn async def execute_remote_mcp_tool(self, args: ExecuteRemoteMCPToolArgs) -> str: diff --git a/packages/tracecat-ee/tracecat_ee/agent/workflows/durable.py b/packages/tracecat-ee/tracecat_ee/agent/workflows/durable.py index c733a9372..f779d257f 100644 --- a/packages/tracecat-ee/tracecat_ee/agent/workflows/durable.py +++ b/packages/tracecat-ee/tracecat_ee/agent/workflows/durable.py @@ -64,6 +64,7 @@ from tracecat.agent.session.activities import ( CreateSessionInput, FinalizeTurnInput, + FinalizeTurnResult, LoadSessionInput, LoadSessionMessagesInput, LoadSessionResult, @@ -491,9 +492,8 @@ def _preserved_agents_binding( FINALIZE_TURN_PATCH = "durable-agent-finalize-turn-v1" +FINALIZE_TURN_WITH_END_PATCH = "durable-agent-finalize-turn-with-end-v1" REMINT_SCOPE_TOKENS_PATCH = "durable-agent-remint-scope-tokens-v1" -# Gates the approval-stream lifecycle as one capability: persist approvals before -# closing the pause stream, rotate continuations, and best-effort stream closure. APPROVAL_STREAM_V2_PATCH = "durable-agent-approval-stream-v2" @@ -517,7 +517,6 @@ def __init__(self, args: AgentWorkflowArgs): self.organization_id = args.role.organization_id self.session_id = args.agent_args.session_id self.active_stream_id = args.agent_args.active_stream_id - self._approval_stream_v2 = False self.harness_type = args.harness_type or "claude_code" self.approvals = ApprovalManager(role=self.role) self.max_requests = args.agent_args.max_requests @@ -920,13 +919,28 @@ async def run(self, args: AgentWorkflowArgs) -> AgentOutput: # Terminal boundary only: approval-pause awaits inside the executor # loop and never reaches here. Clear the active-turn pointers so the # mid-turn DB filter releases the final rows and reconnect -> 204. - # Patch-gated: finalize_turn_activity is a new command, so old - # histories recorded before this change must not replay it. - if workflow.patched(FINALIZE_TURN_PATCH): - await self._finalize_turn() + # The v2 patch folds Redis END into finalize_turn_activity. The + # legacy branch preserves histories that recorded a standalone END, + # with or without pointer cleanup, while the activity-result fallback + # covers a v2 workflow whose task is picked up by a pre-v2 worker. + terminal_stream_id = self.active_stream_id + if workflow.patched(FINALIZE_TURN_WITH_END_PATCH): + await self._finalize_turn( + terminal_stream_id, + emit_terminal_done=True, + ) + else: + if workflow.patched(FINALIZE_TURN_PATCH): + await self._finalize_turn(None, emit_terminal_done=False) + await self._emit_terminal_done(terminal_stream_id) - async def _finalize_turn(self) -> None: - """Clear active-turn pointers at terminal (compare-and-clear by run_id).""" + async def _finalize_turn( + self, + active_stream_id: uuid.UUID | None, + *, + emit_terminal_done: bool, + ) -> None: + """Finalize terminal state and bridge workers without combined support.""" # Use the workflow-id token (same as the persisted curr_run_id), not # args.agent_args.curr_run_id, which is None for DSL/workflow callers and # would skip cleanup. workflow.info() is replay-safe. @@ -934,25 +948,36 @@ async def _finalize_turn(self) -> None: workflow.info().workflow_id ).session_id try: - await workflow.execute_activity( + result: FinalizeTurnResult | None = await workflow.execute_activity( finalize_turn_activity, FinalizeTurnInput( role=self.role, session_id=self.session_id, run_id=run_id, + active_stream_id=active_stream_id, + emit_terminal_done=emit_terminal_done, ), - start_to_close_timeout=timedelta(seconds=10), - # Idempotent (compare-and-clear by run_id); retry so a transient - # failure doesn't leave curr_run_id set and hide the final row. + # Preserve the two former 10-second operation budgets when DB + # cleanup and Redis completion run inside one activity. + start_to_close_timeout=timedelta( + seconds=20 if emit_terminal_done else 10 + ), + # Retry-safe: pointer cleanup is compare-and-clear by run_id, + # and clients tolerate an ambiguous duplicate END. retry_policy=RETRY_POLICIES["activity:fail_slow"], ) except ActivityError as exc: logger.warning( - "Failed to finalize agent turn pointers", + "Failed to finalize agent turn", session_id=str(self.session_id), run_id=str(run_id), + active_stream_id=str(active_stream_id), error=str(exc), ) + return + + if emit_terminal_done and (result is None or not result.terminal_done_emitted): + await self._emit_terminal_done(active_stream_id) async def _finalize_session_error( self, message: str, *, should_stream: bool @@ -1039,6 +1064,30 @@ async def _emit_session_cancelled( error=str(emit_error), ) + async def _emit_terminal_done(self, active_stream_id: uuid.UUID | None) -> None: + """Close the stream for legacy histories or a legacy-worker fallback.""" + try: + await workflow.execute_activity_method( + AgentActivities.emit_session_done, + EmitSessionDoneInputs( + role=self.role, + session_id=self.session_id, + workspace_id=self.workspace_id, + active_stream_id=active_stream_id, + ), + start_to_close_timeout=timedelta(seconds=10), + # Retry-safe: an ambiguous failure may append a duplicate END, + # which clients already tolerate. + retry_policy=RETRY_POLICIES["activity:fail_slow"], + ) + except ActivityError as emit_error: + logger.warning( + "Failed to emit terminal agent stream done", + session_id=self.session_id, + active_stream_id=str(active_stream_id), + error=str(emit_error), + ) + async def _emit_approval_pause_done(self) -> None: """Close the approval-pause stream after approval rows are durable.""" try: @@ -1248,7 +1297,12 @@ async def _run_with_agent_executor( model_settings=cfg.model_settings, routes=compiled_run.llm_routes, ) - self._approval_stream_v2 = workflow.patched(APPROVAL_STREAM_V2_PATCH) + # Replay bridge for histories that recorded the v2 marker. Keep this + # until those histories have drained. This does not make pre-v2 + # histories that already advanced past an approval pause compatible + # with the now-unconditional emit_session_done command; verify that no + # such executions remain RUNNING before rollout. + workflow.deprecate_patch(APPROVAL_STREAM_V2_PATCH) # Prepare executor input executor_input = AgentExecutorInput( @@ -1265,7 +1319,6 @@ async def _run_with_agent_executor( subagents=compiled_run.sandbox_subagents, sdk_session_id=load_result.sdk_session_id, sdk_session_data=load_result.sdk_session_data, - defer_done_on_approval=self._approval_stream_v2, is_fork=load_result.is_fork, ) @@ -1376,8 +1429,7 @@ async def _run_with_agent_executor( tool_call_parts, request_metadata=request_metadata, ) - if self._approval_stream_v2: - await self._emit_approval_pause_done() + await self._emit_approval_pause_done() # Wait for either approval decisions or a user cancellation. await workflow.wait_condition( lambda: self.approvals.is_ready() or self._cancel_requested @@ -1515,7 +1567,6 @@ async def _run_with_agent_executor( subagents=compiled_run.sandbox_subagents, sdk_session_id=reload_result.sdk_session_id, sdk_session_data=reload_result.sdk_session_data, - defer_done_on_approval=self._approval_stream_v2, is_fork=reload_result.is_fork, is_approval_continuation=True, ) diff --git a/tests/temporal/test_durable_agent_workflow.py b/tests/temporal/test_durable_agent_workflow.py index 9087f9001..e121a572d 100644 --- a/tests/temporal/test_durable_agent_workflow.py +++ b/tests/temporal/test_durable_agent_workflow.py @@ -10,7 +10,7 @@ import asyncio import os import uuid -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from datetime import timedelta from typing import Any @@ -21,6 +21,7 @@ from pydantic_ai.tools import ToolApproved, ToolDenied from temporalio import activity +from temporalio import workflow as temporal_workflow from temporalio.api.enums.v1 import EventType from temporalio.client import ( Client, @@ -31,10 +32,12 @@ from temporalio.exceptions import ApplicationError from temporalio.worker import Replayer, UnsandboxedWorkflowRunner, Worker from tracecat_ee.agent.activities import ( + AgentActivities, ApplyApprovalResultsActivityInputs, BuildAgentToolDefsArgs, BuildAgentToolDefsResult, BuildToolDefsResult, + EmitSessionCancelledInputs, EmitSessionDoneInputs, EmitSessionErrorInputs, ExecuteRemoteMCPToolArgs, @@ -47,9 +50,11 @@ ) from tracecat_ee.agent.types import AgentWorkflowID from tracecat_ee.agent.workflows.durable import ( + APPROVAL_STREAM_V2_PATCH, AgentWorkflowArgs, DurableAgentWorkflow, WorkflowApprovalSubmission, + WorkflowCancelRequest, ) from tracecat import config @@ -70,6 +75,8 @@ from tracecat.agent.session.activities import ( CreateSessionInput, CreateSessionResult, + FinalizeTurnInput, + FinalizeTurnResult, LoadSessionInput, LoadSessionMessagesInput, LoadSessionMessagesResult, @@ -94,6 +101,7 @@ from tracecat.agent.workflow_config import agent_config_to_payload from tracecat.auth.types import Role from tracecat.authz.scopes import SERVICE_PRINCIPAL_SCOPES +from tracecat.chat.enums import MessageKind from tracecat.chat.schemas import ChatMessage from tracecat.db.models import AgentSessionHistory, User from tracecat.dsl.common import RETRY_POLICIES @@ -113,6 +121,24 @@ def enable_agent_approvals_entitlement(monkeypatch): ) +@pytest.fixture(autouse=True) +async def close_shared_redis_client(): + """Close the shared Redis client while this test's event loop is alive. + + Tests that register the real ``create_session_activity`` initialize the + module-level ``RedisClient`` (stream-cursor reset) on the current test's + event loop. anyio gives every test a fresh loop, so without closing here + the pooled connection outlives its loop and the next caller to touch it — + e.g. the concurrency-limits fixture's defensive ``close()`` — raises + "Event loop is closed". + """ + yield + from tracecat.redis.client import get_redis_client + + client = await get_redis_client() + await client.close() + + # ============================================================================= # Mock Activity Factories # ============================================================================= @@ -284,7 +310,7 @@ def create_mock_emit_session_done_activity( call_order: list[str] | None = None, done_event: asyncio.Event | None = None, ) -> Callable[..., Any]: - """Create a mock emit_session_done activity for approval-pause closure.""" + """Create a mock emit_session_done activity for stream closure.""" @activity.defn(name="emit_session_done") async def mock_emit_session_done(input: EmitSessionDoneInputs) -> None: @@ -298,6 +324,33 @@ async def mock_emit_session_done(input: EmitSessionDoneInputs) -> None: return mock_emit_session_done +def create_mock_finalize_turn_activity() -> Callable[..., Any]: + """Create a no-op finalize activity for workflow tests without a DB session.""" + + @activity.defn(name="finalize_turn_activity") + async def mock_finalize_turn_activity(input: FinalizeTurnInput) -> None: + del input + + return mock_finalize_turn_activity + + +def create_legacy_finalize_turn_activity() -> Callable[..., Any]: + """Create a DB-only finalize activity that emulates a pre-combined worker.""" + + @activity.defn(name="finalize_turn_activity") + async def legacy_finalize_turn_activity(input: FinalizeTurnInput) -> None: + await finalize_turn_activity( + input.model_copy( + update={ + "active_stream_id": None, + "emit_terminal_done": False, + } + ) + ) + + return legacy_finalize_turn_activity + + def create_activities_with_mock_executor( response_callback: Callable[[int, AgentExecutorInput], AgentExecutorResult], tool_exec_callback: Callable[[RunActionInput], InlineObject[dict[str, str]]] @@ -305,6 +358,8 @@ def create_activities_with_mock_executor( tool_definitions: dict[str, MCPToolDefinition] | None = None, message_load_inputs: list[LoadSessionMessagesInput] | None = None, session_messages: list[ChatMessage] | None = None, + done_inputs: list[EmitSessionDoneInputs] | None = None, + done_call_order: list[str] | None = None, ) -> Sequence[Callable[..., Any]]: """Create a full activity list with mocked agent executor. @@ -327,7 +382,11 @@ def create_activities_with_mock_executor( create_mock_run_agent_activity(response_callback), create_mock_execute_action_activity(tool_exec_callback), create_mock_reconcile_tool_results_activity(), - create_mock_emit_session_done_activity(), + create_mock_finalize_turn_activity(), + create_mock_emit_session_done_activity( + captured_inputs=done_inputs, + call_order=done_call_order, + ), *ApprovalManager.get_activities(), ] return activities @@ -346,6 +405,29 @@ async def replay_durable_agent_workflow_history( assert replay_result.replay_failure is None +async def recorded_patch_ids( + temporal_client: Client, + history: WorkflowHistory, +) -> set[str]: + """Decode patch IDs recorded in a workflow history.""" + patch_ids: set[str] = set() + for event in history.events: + if not event.HasField("marker_recorded_event_attributes"): + continue + attributes = event.marker_recorded_event_attributes + if attributes.marker_name != "core_patch": + continue + if "patch-data" not in attributes.details: + continue + patch_data = await temporal_client.data_converter.decode( + attributes.details["patch-data"].payloads + ) + for data in patch_data: + if isinstance(data, Mapping) and isinstance(data.get("id"), str): + patch_ids.add(data["id"]) + return patch_ids + + async def fetch_history_after_completed_workflow_task( handle: WorkflowHandle[Any, Any], ) -> WorkflowHistory: @@ -539,6 +621,8 @@ async def mock_emit_session_error(args: EmitSessionErrorInputs) -> None: create_mock_load_session_activity(), mock_build_tool_definitions, mock_emit_session_error, + create_mock_finalize_turn_activity(), + create_mock_emit_session_done_activity(), ] async with agent_worker_factory( @@ -572,17 +656,19 @@ async def test_agent_workflow_persists_runtime_terminal_error_without_streaming( """A runtime failure that already streamed its error persists last_error only. ``terminal_stream_error_emitted=True`` means the loopback already pushed the - terminal error onto the SSE stream, so ``emit_session_error`` runs with - ``should_stream=False`` to record the durable last_error signal without - re-emitting a duplicate terminal marker. + error event onto the SSE stream, so ``emit_session_error`` runs with + ``should_stream=False`` to record the durable last_error signal. The workflow + emits END separately after finalization. """ queue = f"test-agent-queue-{mock_session_id}" emitted_errors: list[EmitSessionErrorInputs] = [] + emitted_done: list[EmitSessionDoneInputs] = [] + call_order: list[str] = [] def runtime_failure( call_count: int, input: AgentExecutorInput ) -> AgentExecutorResult: - del call_count, input + del call_count return AgentExecutorResult( success=False, error="runtime exploded", @@ -592,9 +678,14 @@ def runtime_failure( @activity.defn(name="emit_session_error") async def mock_emit_session_error(args: EmitSessionErrorInputs) -> None: emitted_errors.append(args) + call_order.append("persist_error") activities = [ - *create_activities_with_mock_executor(runtime_failure), + *create_activities_with_mock_executor( + runtime_failure, + done_inputs=emitted_done, + done_call_order=call_order, + ), mock_emit_session_error, ] @@ -613,12 +704,14 @@ async def mock_emit_session_error(args: EmitSessionErrorInputs) -> None: assert isinstance(exc_info.value.cause, ApplicationError) assert "Agent execution failed: runtime exploded" in str(exc_info.value.cause) - # Persist-only: exactly one emit, carrying last_error, with streaming off so - # the already-emitted terminal marker is not duplicated. + # Persist-only: exactly one emit, carrying last_error, with error streaming + # off so the already-emitted inline error is not duplicated. assert len(emitted_errors) == 1 assert emitted_errors[0].session_id == mock_session_id assert emitted_errors[0].message == "Agent execution failed: runtime exploded" assert emitted_errors[0].should_stream is False + assert len(emitted_done) == 1 + assert call_order == ["persist_error", "emit_session_done"] @pytest.mark.anyio @@ -631,11 +724,12 @@ async def test_agent_workflow_streams_executor_pre_stream_failure( ) -> None: queue = f"test-agent-queue-{mock_session_id}" emitted_errors: list[EmitSessionErrorInputs] = [] + emitted_done: list[EmitSessionDoneInputs] = [] def setup_failure( call_count: int, input: AgentExecutorInput ) -> AgentExecutorResult: - del call_count, input + del call_count return AgentExecutorResult( success=False, error="executor setup failed", @@ -647,7 +741,10 @@ async def mock_emit_session_error(args: EmitSessionErrorInputs) -> None: emitted_errors.append(args) activities = [ - *create_activities_with_mock_executor(setup_failure), + *create_activities_with_mock_executor( + setup_failure, + done_inputs=emitted_done, + ), mock_emit_session_error, ] @@ -669,6 +766,8 @@ async def mock_emit_session_error(args: EmitSessionErrorInputs) -> None: assert len(emitted_errors) == 1 assert emitted_errors[0].session_id == mock_session_id assert emitted_errors[0].message == "Agent execution failed: executor setup failed" + assert emitted_errors[0].should_stream is True + assert len(emitted_done) == 1 @pytest.mark.anyio @@ -721,6 +820,474 @@ def mock_executor( assert [input.session_id for input in message_load_inputs] == [mock_session_id] +@pytest.mark.anyio +@pytest.mark.integration +async def test_agent_workflow_replays_approval_stream_v2_patch_history( + svc_role: Role, + temporal_client: Client, + agent_worker_factory, + mock_session_id: uuid.UUID, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The deprecated bridge must replay suspended histories with the v2 marker.""" + queue = f"test-agent-queue-{mock_session_id}" + approval_done_emitted = asyncio.Event() + + def approval_executor( + call_count: int, input: AgentExecutorInput + ) -> AgentExecutorResult: + del input + assert call_count == 0 + return AgentExecutorResult( + success=True, + approval_requested=True, + approval_items=[ + ToolCallContent( + id="call-1", + name="core__http_request", + input={"url": "https://example.com", "method": "GET"}, + ) + ], + ) + + @activity.defn(name="record_approval_requests") + async def mock_record_approval_requests( + input: PersistApprovalsActivityInputs, + ) -> None: + assert [item.tool_call_id for item in input.approvals] == ["call-1"] + + workflow_args = AgentWorkflowArgs( + role=svc_role, + agent_args=RunAgentArgs( + session_id=mock_session_id, + user_prompt="Request an approved action", + config=AgentConfig( + model_name="claude-3-5-sonnet-20241022", + model_provider="anthropic", + actions=["core.http_request"], + tool_approvals={"core.http_request": True}, + ), + ), + entity_type=AgentSessionEntity.WORKFLOW, + entity_id=uuid.uuid4(), + ) + activities = [ + create_mock_create_session_activity(), + create_mock_load_session_activity(), + create_mock_load_session_messages_activity(), + create_mock_build_tool_definitions_activity(), + create_mock_run_agent_activity(approval_executor), + mock_record_approval_requests, + create_mock_emit_session_done_activity(done_event=approval_done_emitted), + ] + + # Simulate the prior worker version at the exact production call site: + # patched() records the v2 marker, while deprecate_patch() preserves replay + # compatibility without writing that marker into new histories. + with monkeypatch.context() as legacy_patch: + legacy_patch.setattr( + temporal_workflow, + "deprecate_patch", + temporal_workflow.patched, + ) + async with agent_worker_factory( + temporal_client, + task_queue=queue, + custom_activities=activities, + ): + wf_handle = await temporal_client.start_workflow( + DurableAgentWorkflow.run, + workflow_args, + id=AgentWorkflowID(mock_session_id), + task_queue=queue, + retry_policy=RETRY_POLICIES["workflow:fail_fast"], + execution_timeout=timedelta(seconds=30), + ) + await asyncio.wait_for(approval_done_emitted.wait(), timeout=10) + marked_history = await fetch_history_after_completed_workflow_task( + wf_handle + ) + + await wf_handle.terminate(reason="Replay regression history captured") + assert APPROVAL_STREAM_V2_PATCH in await recorded_patch_ids( + temporal_client, + marked_history, + ) + await replay_durable_agent_workflow_history(temporal_client, marked_history) + + +@pytest.mark.anyio +@pytest.mark.integration +async def test_terminal_end_observes_publishable_final_history( + svc_role: Role, + temporal_client: Client, + agent_worker_factory, + mock_session_id: uuid.UUID, + test_user: User, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """END is emitted only after final rows are visible to canonical reads.""" + del test_user + queue = f"test-agent-queue-{mock_session_id}" + active_stream_id = uuid.uuid4() + call_order: list[str] = [] + + @activity.defn(name="run_agent_activity") + async def mock_run_agent_activity( + input: AgentExecutorInput, + ) -> AgentExecutorResult: + assert input.active_stream_id == active_stream_id + assert input.curr_run_id == mock_session_id + async with AgentSessionService.with_session(role=input.role) as service: + service.session.add( + AgentSessionHistory( + session_id=input.session_id, + workspace_id=input.workspace_id, + kind=MessageKind.CHAT_MESSAGE.value, + curr_run_id=input.curr_run_id, + content={ + "uuid": str(uuid.uuid4()), + "sessionId": "sdk-session", + "type": "assistant", + "message": { + "role": "assistant", + "content": [{"type": "text", "text": "Final answer"}], + }, + }, + ) + ) + await service.session.commit() + call_order.append("persist_final_history") + return AgentExecutorResult(success=True, output={"status": "done"}) + + @activity.defn(name="load_session_messages_activity") + async def observed_load_session_messages_activity( + input: LoadSessionMessagesInput, + ) -> LoadSessionMessagesResult: + result = await load_session_messages_activity(input) + call_order.append("load_terminal_history") + return result + + class ObservedTerminalStream: + async def done(self) -> None: + async with AgentSessionService.with_session(role=svc_role) as service: + session = await service.get_session(mock_session_id) + assert session is not None + assert session.curr_run_id is None + messages = await service.list_messages( + mock_session_id, + include_active=False, + ) + assert len(messages) == 1 + assert messages[0].kind is MessageKind.CHAT_MESSAGE + assert messages[0].message is not None + call_order.append("emit_session_done") + + class InitialSessionStream: + async def clear_buffer(self) -> None: + return None + + async def observed_stream_new( + *, + workspace_id: uuid.UUID, + session_id: uuid.UUID, + stream_id: uuid.UUID | None = None, + ) -> ObservedTerminalStream | InitialSessionStream: + assert workspace_id == svc_role.workspace_id + assert session_id == mock_session_id + if stream_id is None: + return InitialSessionStream() + assert stream_id == active_stream_id + return ObservedTerminalStream() + + monkeypatch.setattr( + "tracecat.agent.session.service.AgentStream.new", + observed_stream_new, + ) + + @activity.defn(name="finalize_turn_activity") + async def observed_finalize_turn_activity( + input: FinalizeTurnInput, + ) -> FinalizeTurnResult | None: + result = await finalize_turn_activity(input) + call_order.append("finalize_turn") + return result + + workflow_args = AgentWorkflowArgs( + role=svc_role, + agent_args=RunAgentArgs( + session_id=mock_session_id, + user_prompt="Return a final answer", + config=AgentConfig( + model_name="claude-3-5-sonnet-20241022", + model_provider="anthropic", + actions=[], + ), + active_stream_id=active_stream_id, + ), + entity_type=AgentSessionEntity.WORKFLOW, + entity_id=uuid.uuid4(), + ) + activities = [ + create_session_activity, + load_session_activity, + observed_load_session_messages_activity, + observed_finalize_turn_activity, + create_mock_build_tool_definitions_activity({}), + mock_run_agent_activity, + ] + + async with agent_worker_factory( + temporal_client, task_queue=queue, custom_activities=activities + ): + handle = await temporal_client.start_workflow( + DurableAgentWorkflow.run, + workflow_args, + id=AgentWorkflowID(mock_session_id), + task_queue=queue, + retry_policy=RETRY_POLICIES["workflow:fail_fast"], + execution_timeout=timedelta(seconds=30), + ) + result = await handle.result() + completed_history = await handle.fetch_history() + await replay_durable_agent_workflow_history( + temporal_client, + completed_history, + ) + + assert result.output == {"status": "done"} + assert call_order == [ + "persist_final_history", + "load_terminal_history", + "emit_session_done", + "finalize_turn", + ] + + +@pytest.mark.anyio +@pytest.mark.integration +async def test_executor_cancellation_persists_marker_before_terminal_end( + svc_role: Role, + temporal_client: Client, + agent_worker_factory, + mock_session_id: uuid.UUID, + test_user: User, +) -> None: + del test_user + queue = f"test-agent-queue-{mock_session_id}" + active_stream_id = uuid.uuid4() + terminal_done_observed = asyncio.Event() + + def cancelled_executor( + call_count: int, input: AgentExecutorInput + ) -> AgentExecutorResult: + del call_count + return AgentExecutorResult( + success=True, + cancelled=True, + cancelled_reason="user_cancel", + interrupted_tool_call_ids=["tool-call-1"], + ) + + @activity.defn(name="emit_session_done") + async def observe_terminal_done(input: EmitSessionDoneInputs) -> None: + assert input.active_stream_id == active_stream_id + async with AgentSessionService.with_session(role=input.role) as service: + session = await service.get_session(input.session_id) + assert session is not None + assert session.curr_run_id is None + messages = await service.list_messages( + input.session_id, + include_active=False, + ) + cancelled_messages = [ + message for message in messages if message.kind is MessageKind.CANCELLED + ] + assert len(cancelled_messages) == 1 + assert cancelled_messages[0].cancelled == { + "reason": "user_cancel", + "tool_call_ids": ["tool-call-1"], + } + terminal_done_observed.set() + + agent_activities = AgentActivities() + + @activity.defn(name="emit_session_cancelled") + async def persist_executor_cancellation( + input: EmitSessionCancelledInputs, + ) -> None: + assert input.emit_stream is False + await agent_activities.emit_session_cancelled(input) + + workflow_args = AgentWorkflowArgs( + role=svc_role, + agent_args=RunAgentArgs( + session_id=mock_session_id, + user_prompt="Start then cancel", + config=AgentConfig( + model_name="claude-3-5-sonnet-20241022", + model_provider="anthropic", + actions=[], + ), + active_stream_id=active_stream_id, + ), + entity_type=AgentSessionEntity.WORKFLOW, + entity_id=uuid.uuid4(), + ) + activities = [ + create_session_activity, + load_session_activity, + load_session_messages_activity, + create_legacy_finalize_turn_activity(), + create_mock_build_tool_definitions_activity({}), + create_mock_run_agent_activity(cancelled_executor), + persist_executor_cancellation, + observe_terminal_done, + ] + + async with agent_worker_factory( + temporal_client, task_queue=queue, custom_activities=activities + ): + result = await temporal_client.execute_workflow( + DurableAgentWorkflow.run, + workflow_args, + id=AgentWorkflowID(mock_session_id), + task_queue=queue, + retry_policy=RETRY_POLICIES["workflow:fail_fast"], + execution_timeout=timedelta(seconds=30), + ) + + assert result.output is None + assert terminal_done_observed.is_set() + + +@pytest.mark.anyio +@pytest.mark.integration +async def test_approval_wait_cancellation_defers_end_until_marker_and_finalize( + svc_role: Role, + temporal_client: Client, + agent_worker_factory, + mock_session_id: uuid.UUID, + test_user: User, +) -> None: + del test_user + queue = f"test-agent-queue-{mock_session_id}" + active_stream_id = uuid.uuid4() + approval_pause_done = asyncio.Event() + done_session_run_ids: list[uuid.UUID | None] = [] + cancelled_inputs: list[EmitSessionCancelledInputs] = [] + + def approval_executor( + call_count: int, input: AgentExecutorInput + ) -> AgentExecutorResult: + assert call_count == 0 + return AgentExecutorResult( + success=True, + approval_requested=True, + approval_items=[ + ToolCallContent( + id="tool-call-1", + name="core__http_request", + input={"url": "https://example.com", "method": "GET"}, + ) + ], + ) + + @activity.defn(name="record_approval_requests") + async def mock_record_approval_requests( + input: PersistApprovalsActivityInputs, + ) -> None: + assert [item.tool_call_id for item in input.approvals] == ["tool-call-1"] + + @activity.defn(name="apply_approval_decisions") + async def mock_apply_approval_decisions( + input: ApplyApprovalResultsActivityInputs, + ) -> None: + assert [item.tool_call_id for item in input.decisions] == ["tool-call-1"] + assert input.decisions[0].approved is False + + @activity.defn(name="emit_session_cancelled") + async def persist_cancelled_notice(input: EmitSessionCancelledInputs) -> None: + cancelled_inputs.append(input) + assert input.emit_stream is True + async with AgentSessionService.with_session(role=input.role) as service: + await service.append_cancelled_marker( + input.session_id, + reason=input.reason, + interrupted_tool_call_ids=input.interrupted_tool_call_ids, + curr_run_id=input.curr_run_id, + ) + + @activity.defn(name="emit_session_done") + async def observe_pause_or_terminal_done(input: EmitSessionDoneInputs) -> None: + assert input.active_stream_id == active_stream_id + async with AgentSessionService.with_session(role=input.role) as service: + session = await service.get_session(input.session_id) + assert session is not None + done_session_run_ids.append(session.curr_run_id) + if session.curr_run_id is None: + messages = await service.list_messages( + input.session_id, + include_active=False, + ) + assert any( + message.kind is MessageKind.CANCELLED for message in messages + ) + else: + approval_pause_done.set() + + workflow_args = AgentWorkflowArgs( + role=svc_role, + agent_args=RunAgentArgs( + session_id=mock_session_id, + user_prompt="Request an approved action", + config=AgentConfig( + model_name="claude-3-5-sonnet-20241022", + model_provider="anthropic", + actions=["core.http_request"], + tool_approvals={"core.http_request": True}, + ), + active_stream_id=active_stream_id, + ), + entity_type=AgentSessionEntity.WORKFLOW, + entity_id=uuid.uuid4(), + ) + activities = [ + create_session_activity, + load_session_activity, + load_session_messages_activity, + create_legacy_finalize_turn_activity(), + create_mock_build_tool_definitions_activity(), + create_mock_run_agent_activity(approval_executor), + mock_record_approval_requests, + mock_apply_approval_decisions, + persist_cancelled_notice, + observe_pause_or_terminal_done, + ] + + async with agent_worker_factory( + temporal_client, task_queue=queue, custom_activities=activities + ): + handle = await temporal_client.start_workflow( + DurableAgentWorkflow.run, + workflow_args, + id=AgentWorkflowID(mock_session_id), + task_queue=queue, + retry_policy=RETRY_POLICIES["workflow:fail_fast"], + execution_timeout=timedelta(seconds=30), + ) + await asyncio.wait_for(approval_pause_done.wait(), timeout=10) + await handle.execute_update( + DurableAgentWorkflow.request_cancel, + WorkflowCancelRequest(reason="user_cancel"), + ) + result = await handle.result() + + assert result.output is None + assert len(cancelled_inputs) == 1 + assert done_session_run_ids == [mock_session_id, None] + + @pytest.mark.anyio @pytest.mark.integration async def test_agent_workflow_preserves_stored_subagent_binding_on_resume( @@ -832,6 +1399,8 @@ async def mock_run_agent_activity( mock_run_agent_activity, create_mock_execute_action_activity(), create_mock_reconcile_tool_results_activity(), + create_mock_finalize_turn_activity(), + create_mock_emit_session_done_activity(), *ApprovalManager.get_activities(), ] @@ -991,11 +1560,14 @@ async def test_agent_workflow_routes_approved_tools_to_executor_and_reconciles_h approval_request_recorded = asyncio.Event() approval_done_emitted = asyncio.Event() approval_pause_call_order: list[str] = [] + emitted_done_inputs: list[EmitSessionDoneInputs] = [] resumed_after_approval = asyncio.Event() agent_executor_task_queues: list[str] = [] executor_task_queues: list[str] = [] captured_run_inputs: list[RunActionInput] = [] captured_executor_roles: list[Role] = [] + initial_stream_id = uuid.uuid4() + rotated_stream_id = uuid.uuid4() class _FakeStream: async def append(self, event: Any) -> None: @@ -1058,6 +1630,7 @@ async def mock_run_agent_activity( activity.heartbeat("Mock agent running") if run_agent_call_count == 0: + assert input.active_stream_id == initial_stream_id assistant_uuid = str(uuid.uuid4()) async with AgentSessionService.with_session(role=input.role) as service: session = await service.get_session(input.session_id) @@ -1141,6 +1714,7 @@ async def mock_run_agent_activity( assert input.is_approval_continuation is True assert input.sdk_session_id == "sdk-session" assert input.sdk_session_data is None + assert input.active_stream_id == rotated_stream_id resumed_after_approval.set() run_agent_call_count += 1 @@ -1176,6 +1750,7 @@ async def mock_execute_action_activity( ) mock_emit_session_done = create_mock_emit_session_done_activity( + captured_inputs=emitted_done_inputs, call_order=approval_pause_call_order, done_event=approval_done_emitted, ) @@ -1193,6 +1768,7 @@ async def mock_execute_action_activity( user_prompt="Make a test HTTP request", config=agent_config_with_approvals, curr_run_id=mock_session_id, + active_stream_id=initial_stream_id, ), entity_type=AgentSessionEntity.WORKFLOW, entity_id=uuid.uuid4(), @@ -1206,7 +1782,7 @@ async def mock_execute_action_activity( load_session_activity, load_session_messages_activity, reconcile_tool_results_activity, - finalize_turn_activity, + create_legacy_finalize_turn_activity(), mock_build_tool_definitions, mock_record_approval_requests, mock_apply_approval_decisions, @@ -1262,6 +1838,7 @@ async def mock_execute_action_activity( WorkflowApprovalSubmission( approvals={"call_123": True}, approved_by=svc_role.user_id, + new_stream_id=rotated_stream_id, ), ) @@ -1280,6 +1857,10 @@ async def mock_execute_action_activity( ] assert executor_task_queues == [executor_queue] assert resumed_after_approval.is_set() + assert [input.active_stream_id for input in emitted_done_inputs] == [ + initial_stream_id, + rotated_stream_id, + ] assert len(captured_run_inputs) == 1 assert len(captured_executor_roles) == 1 assert captured_run_inputs[0].task.action == "core__http_request" diff --git a/tests/unit/test_agent_activities.py b/tests/unit/test_agent_activities.py index 841b91d5b..886f63db2 100644 --- a/tests/unit/test_agent_activities.py +++ b/tests/unit/test_agent_activities.py @@ -25,6 +25,7 @@ BuildAgentToolDefsArgs, BuildToolDefsArgs, EmitSessionDoneInputs, + EmitSessionErrorInputs, ) from tracecat.agent.common.protocol import RuntimeInitPayload @@ -1257,6 +1258,48 @@ async def test_emit_session_done_pushes_done_to_active_stream( ) stream.done.assert_awaited_once() + @pytest.mark.anyio + async def test_emit_session_error_leaves_done_to_workflow( + self, + mock_role: Role, + mock_session_id: uuid.UUID, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + workspace_id = mock_role.workspace_id or uuid.uuid4() + active_stream_id = uuid.uuid4() + agent_session = SimpleNamespace(last_error=None) + service = SimpleNamespace( + get_session=AsyncMock(return_value=agent_session), + session=SimpleNamespace(add=MagicMock(), commit=AsyncMock()), + ) + service_context = AsyncMock() + service_context.__aenter__.return_value = service + with_session = MagicMock(return_value=service_context) + monkeypatch.setattr( + "tracecat.agent.session.service.AgentSessionService.with_session", + with_session, + ) + stream = SimpleNamespace(error=AsyncMock(), done=AsyncMock()) + stream_new = AsyncMock(return_value=stream) + monkeypatch.setattr( + "tracecat_ee.agent.activities.AgentStream.new", + stream_new, + ) + + await AgentActivities().emit_session_error( + EmitSessionErrorInputs( + role=mock_role, + session_id=mock_session_id, + workspace_id=workspace_id, + active_stream_id=active_stream_id, + message="runtime failed", + ) + ) + + assert agent_session.last_error == "runtime failed" + stream.error.assert_awaited_once_with("runtime failed") + stream.done.assert_not_awaited() + @pytest.mark.anyio async def test_returns_approval_requested_on_approval_interrupt( self, mock_executor_input: AgentExecutorInput diff --git a/tests/unit/test_agent_executor_loopback.py b/tests/unit/test_agent_executor_loopback.py index 7f338cd28..edddecbe6 100644 --- a/tests/unit/test_agent_executor_loopback.py +++ b/tests/unit/test_agent_executor_loopback.py @@ -164,7 +164,7 @@ async def test_emit_terminal_error_uses_redis_when_external_lookup_errors( stream_id=loopback_input.active_stream_id, ) fake_stream.error.assert_awaited_once_with("runtime exited before connect") - fake_stream.done.assert_awaited_once() + fake_stream.done.assert_not_awaited() @pytest.mark.anyio @@ -195,7 +195,7 @@ async def test_emit_terminal_error_emits_failed_compaction_when_pending( assert failed_event.type == StreamEventType.COMPACTION assert failed_event.metadata == {"phase": "failed"} fake_stream.error.assert_awaited_once_with("runtime exited before connect") - fake_stream.done.assert_awaited_once() + fake_stream.done.assert_not_awaited() @pytest.mark.anyio @@ -215,12 +215,11 @@ async def test_prepare_initializes_stream_sink_once( initialize_stream_sink.assert_awaited_once() -def _make_handler(*, defer_done_on_approval: bool = False) -> LoopbackHandler: +def _make_handler() -> LoopbackHandler: return LoopbackHandler( input=LoopbackInput( session_id=UUID("00000000-0000-0000-0000-000000000001"), workspace_id=UUID("00000000-0000-0000-0000-000000000002"), - defer_done_on_approval=defer_done_on_approval, ) ) @@ -477,37 +476,81 @@ def __init__(self, session: object, role: Role) -> None: @pytest.mark.anyio -@pytest.mark.parametrize( - ("defer_done", "approval_requested", "expect_done"), - [ - (False, True, True), - (True, True, False), - (True, False, True), - ], - ids=["legacy-approval", "deferred-approval", "deferred-normal-turn"], -) -async def test_emit_stream_done_policy( - defer_done: bool, - approval_requested: bool, - expect_done: bool, -) -> None: - handler = _make_handler(defer_done_on_approval=defer_done) +async def test_close_external_stream_leaves_redis_open() -> None: + handler = _make_handler() + stream = _FakeStream() + handler._stream_sink = stream + + await handler._close_external_stream() + + stream.done.assert_not_awaited() + + +@pytest.mark.anyio +async def test_terminal_success_closes_only_external_sink() -> None: + handler = _make_handler() + redis_stream = _FakeStream() + external_stream = _FakeExternalSink() + handler._stream_sink = FanoutStreamSink( + sinks=( + AgentStreamSink(stream=cast(AgentStream, redis_stream)), + external_stream, + ) + ) + + await handler.send_result(output={"status": "done"}) + await handler.send_done() + await handler.send_done() + + assert handler.build_result().success is True + redis_stream.done.assert_not_awaited() + external_stream.done.assert_awaited_once() + assert handler._external_stream_done_emitted is True + + +@pytest.mark.anyio +async def test_terminal_error_streams_error_and_closes_only_external_sink() -> None: + handler = _make_handler() + redis_stream = _FakeStream() + external_stream = _FakeExternalSink() + handler._stream_sink = FanoutStreamSink( + sinks=( + AgentStreamSink(stream=cast(AgentStream, redis_stream)), + external_stream, + ) + ) + + await handler.send_error("runtime failed") + + redis_stream.error.assert_awaited_once_with("runtime failed") + external_stream.error.assert_awaited_once_with("runtime failed") + redis_stream.done.assert_not_awaited() + external_stream.done.assert_awaited_once() + assert handler._external_stream_done_emitted is True + + +@pytest.mark.anyio +async def test_terminal_error_leaves_redis_open_for_workflow() -> None: + handler = _make_handler() stream = _FakeStream() + event_order: list[str] = [] + + async def record_error(error: str) -> None: + assert error == "runtime failed" + event_order.append("error") + + stream.error.side_effect = record_error handler._stream_sink = stream - handler._result.approval_requested = approval_requested - await handler._emit_stream_done() + await handler.send_error("runtime failed") - if expect_done: - stream.done.assert_awaited_once() - else: - stream.done.assert_not_awaited() - assert handler._stream_done_emitted is expect_done + assert event_order == ["error"] + stream.done.assert_not_awaited() @pytest.mark.anyio -async def test_emit_stream_done_closes_external_sink_on_approval_pause() -> None: - handler = _make_handler(defer_done_on_approval=True) +async def test_close_external_stream_is_deduplicated_on_approval_pause() -> None: + handler = _make_handler() redis_stream = _FakeStream() external_stream = _FakeExternalSink() handler._stream_sink = FanoutStreamSink( @@ -518,12 +561,11 @@ async def test_emit_stream_done_closes_external_sink_on_approval_pause() -> None ) handler._result.approval_requested = True - await handler._emit_stream_done() - await handler._emit_stream_done() + await handler._close_external_stream() + await handler._close_external_stream() redis_stream.done.assert_not_awaited() external_stream.done.assert_awaited_once() - assert handler._stream_done_emitted is False assert handler._external_stream_done_emitted is True @@ -559,7 +601,7 @@ async def test_process_runtime_events_emits_failed_compaction_on_runtime_error() {"phase": "failed"}, ] stream.error.assert_awaited_once_with("request_timeout: LLM gateway timed out") - stream.done.assert_awaited_once() + stream.done.assert_not_awaited() @pytest.mark.anyio @@ -592,7 +634,7 @@ async def test_process_runtime_events_emits_failed_compaction_on_done_without_bo {"phase": "failed"}, ] stream.error.assert_not_awaited() - stream.done.assert_awaited_once() + stream.done.assert_not_awaited() @pytest.mark.anyio @@ -606,7 +648,7 @@ async def test_process_runtime_events_fails_when_done_arrives_without_result() - assert handler._result.error == "Runtime completed without final result" stream.error.assert_awaited_once_with("Runtime completed without final result") - stream.done.assert_awaited_once() + stream.done.assert_not_awaited() @pytest.mark.anyio @@ -650,7 +692,7 @@ async def test_send_done_preserves_existing_error_state() -> None: assert handler._result.success is False assert handler._result.error == "runtime failed" stream.error.assert_not_awaited() - stream.done.assert_awaited_once() + stream.done.assert_not_awaited() @pytest.mark.anyio diff --git a/tests/unit/test_agent_sandbox_litellm.py b/tests/unit/test_agent_sandbox_litellm.py index 2067bcaf7..1afa198f3 100644 --- a/tests/unit/test_agent_sandbox_litellm.py +++ b/tests/unit/test_agent_sandbox_litellm.py @@ -854,7 +854,8 @@ async def fake_persist_session_line( assert any(request.get("model") == "customer-alias" for request in proxy.requests) assert stream_sink.errors == [] - assert stream_sink.done_count == 1 + # Terminal END is emitted by the workflow, not the executor loopback. + assert stream_sink.done_count == 0 async def _run_mcp_compression_initialize_case( @@ -1070,7 +1071,8 @@ async def fast_send_control_request( request["accept_encoding"] == "identity" for request in mcp_proxy.requests ) assert stream_sink.errors == [] - assert stream_sink.done_count == 1 + # Terminal END is emitted by the workflow, not the executor loopback. + assert stream_sink.done_count == 0 else: assert result.success is False assert result.error == "Unexpected error: Control request timeout: initialize" @@ -1532,7 +1534,8 @@ async def approval_script( StreamEventType.APPROVAL_REQUEST ] assert stream_sink.errors == [] - assert stream_sink.done_count == 1 + # Approval-pause END is emitted by the workflow, not the executor loopback. + assert stream_sink.done_count == 0 assert persisted_session_lines == [] assert fake_proxy.started is True @@ -1618,7 +1621,8 @@ async def resume_script( assert [event.type for event in stream_sink.events] == [StreamEventType.TEXT_DELTA] assert stream_sink.errors == [] - assert stream_sink.done_count == 1 + # Terminal END is emitted by the workflow, not the executor loopback. + assert stream_sink.done_count == 0 assert persisted_session_lines == [ ("child-sdk-session", session_line, False), ] diff --git a/tests/unit/test_agent_session_activities.py b/tests/unit/test_agent_session_activities.py index 0a752f73f..c74983543 100644 --- a/tests/unit/test_agent_session_activities.py +++ b/tests/unit/test_agent_session_activities.py @@ -6,14 +6,69 @@ import pytest from tracecat.agent.session.activities import ( + FinalizeTurnInput, PendingToolResult, ReconcileToolResultsInput, + finalize_turn_activity, reconcile_tool_results_activity, ) from tracecat.auth.types import Role from tracecat.storage.object import ExternalObject, ObjectRef +@pytest.mark.anyio +@pytest.mark.parametrize( + "emit_terminal_done", + [False, True], + ids=["legacy-workflow", "combined-workflow"], +) +async def test_finalize_turn_activity_bridges_workflow_versions( + monkeypatch: pytest.MonkeyPatch, + *, + emit_terminal_done: bool, +) -> None: + service = AsyncMock() + ctx = AsyncMock() + ctx.__aenter__.return_value = service + monkeypatch.setattr( + "tracecat.agent.session.activities.AgentSessionService.with_session", + MagicMock(return_value=ctx), + ) + + role = Role( + type="service", + workspace_id=uuid.uuid4(), + organization_id=uuid.uuid4(), + service_id="tracecat-service", + ) + session_id = uuid.uuid4() + run_id = uuid.uuid4() + active_stream_id = uuid.uuid4() + input = FinalizeTurnInput( + role=role, + session_id=session_id, + run_id=run_id, + active_stream_id=active_stream_id, + emit_terminal_done=emit_terminal_done, + ) + + result = await finalize_turn_activity(input) + + if emit_terminal_done: + assert result is not None + assert result.terminal_done_emitted is True + service.finalize_turn.assert_awaited_once_with( + session_id, + run_id, + active_stream_id=active_stream_id, + ) + service.clear_turn_pointers.assert_not_awaited() + else: + assert result is None + service.clear_turn_pointers.assert_awaited_once_with(session_id, run_id) + service.finalize_turn.assert_not_awaited() + + @pytest.mark.anyio async def test_reconcile_tool_results_raises_on_materialization_failure( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_agent_session_finalize.py b/tests/unit/test_agent_session_finalize.py new file mode 100644 index 000000000..1b427ee5b --- /dev/null +++ b/tests/unit/test_agent_session_finalize.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import uuid +from types import SimpleNamespace +from typing import Any, cast +from unittest.mock import AsyncMock + +import pytest + +from tracecat.agent.session.service import AgentSessionService +from tracecat.auth.types import Role + + +def _build_service( + call_order: list[str], +) -> tuple[AgentSessionService, SimpleNamespace, Role]: + role = Role( + type="service", + service_id="tracecat-api", + workspace_id=uuid.uuid4(), + organization_id=uuid.uuid4(), + scopes=frozenset({"agent:execute"}), + ) + + async def execute(_: object) -> None: + call_order.append("execute") + + async def commit() -> None: + call_order.append("commit") + + session = SimpleNamespace( + execute=AsyncMock(side_effect=execute), + commit=AsyncMock(side_effect=commit), + ) + return AgentSessionService(cast(Any, session), role), session, role + + +@pytest.mark.anyio +async def test_finalize_turn_commits_pointers_before_emitting_done( + monkeypatch: pytest.MonkeyPatch, +) -> None: + call_order: list[str] = [] + service, _, role = _build_service(call_order) + session_id = uuid.uuid4() + run_id = uuid.uuid4() + active_stream_id = uuid.uuid4() + expected_session_id = session_id + + async def done() -> None: + call_order.append("done") + + stream = SimpleNamespace(done=done) + + async def open_stream( + *, + workspace_id: uuid.UUID, + session_id: uuid.UUID, + stream_id: uuid.UUID | None, + ) -> SimpleNamespace: + assert workspace_id == role.workspace_id + assert session_id == expected_session_id + assert stream_id == active_stream_id + call_order.append("open_stream") + return stream + + monkeypatch.setattr( + "tracecat.agent.session.service.AgentStream.new", + open_stream, + ) + + await service.finalize_turn( + session_id, + run_id, + active_stream_id=active_stream_id, + ) + + assert call_order == ["execute", "commit", "open_stream", "done"] + + +@pytest.mark.anyio +async def test_finalize_turn_propagates_stream_failure_after_commit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + call_order: list[str] = [] + service, session, _ = _build_service(call_order) + + async def done() -> None: + call_order.append("done") + raise RuntimeError("redis unavailable") + + async def open_stream(**_: object) -> SimpleNamespace: + call_order.append("open_stream") + return SimpleNamespace(done=done) + + monkeypatch.setattr( + "tracecat.agent.session.service.AgentStream.new", + open_stream, + ) + + with pytest.raises(RuntimeError, match="redis unavailable"): + await service.finalize_turn( + uuid.uuid4(), + uuid.uuid4(), + active_stream_id=uuid.uuid4(), + ) + + session.commit.assert_awaited_once() + assert call_order == ["execute", "commit", "open_stream", "done"] diff --git a/tests/unit/test_durable_agent_workflow_search_attributes.py b/tests/unit/test_durable_agent_workflow_search_attributes.py index f92903f0f..40bd2bd70 100644 --- a/tests/unit/test_durable_agent_workflow_search_attributes.py +++ b/tests/unit/test_durable_agent_workflow_search_attributes.py @@ -16,6 +16,7 @@ BUILD_AGENT_TOOL_DEFINITIONS_PATCH, EMIT_PRE_STREAM_SESSION_ERRORS_PATCH, FINALIZE_TURN_PATCH, + FINALIZE_TURN_WITH_END_PATCH, LOAD_TERMINAL_MESSAGE_HISTORY_PATCH, PERSIST_SESSION_ERROR_PATCH, UPSERT_TRACECAT_SEARCH_ATTRIBUTES_PATCH, @@ -31,6 +32,7 @@ from tracecat.agent.executor.schemas import ApprovedToolCall from tracecat.agent.preset.activities import ResolveAgentPresetConfigActivityInput from tracecat.agent.schemas import AgentOutput, RunAgentArgs +from tracecat.agent.session.activities import FinalizeTurnInput, FinalizeTurnResult from tracecat.agent.session.types import AgentSessionEntity from tracecat.agent.types import AgentConfig from tracecat.agent.workflow_config import agent_config_to_payload @@ -88,7 +90,7 @@ def test_agent_workflow_args_ignores_legacy_workspace_credentials() -> None: assert not hasattr(workflow_args.agent_args, "use_workspace_credentials") -def test_legacy_workflow_rotates_stream_from_new_approval_update() -> None: +def test_workflow_rotates_stream_from_new_approval_update() -> None: role = Role( type="user", service_id="tracecat-api", @@ -101,7 +103,6 @@ def test_legacy_workflow_rotates_stream_from_new_approval_update() -> None: previous_stream_id = uuid.uuid4() new_stream_id = uuid.uuid4() workflow_instance.active_stream_id = previous_stream_id - workflow_instance._approval_stream_v2 = False cast(Any, workflow_instance.set_approvals)( WorkflowApprovalSubmission( @@ -260,18 +261,232 @@ async def test_run_skips_search_attribute_upsert_without_patch_marker() -> None: "_run_with_agent_executor", AsyncMock(return_value=expected_output), ) as run_mock, + patch.object( + workflow_instance, + "_finalize_turn", + AsyncMock(), + ) as finalize_turn_mock, + patch.object( + workflow_instance, + "_emit_terminal_done", + AsyncMock(), + ) as emit_terminal_done_mock, ): result = await workflow_instance.run(workflow_args) - # run() gates the search-attribute upsert (entry) and finalize_turn (finally) - # behind patch markers; legacy histories see False for both. + # Search-attribute upsert and both finalization history shapes retain their + # independent patch gates. assert patched_mock.call_args_list == [ ((UPSERT_TRACECAT_SEARCH_ATTRIBUTES_PATCH,),), + ((FINALIZE_TURN_WITH_END_PATCH,),), ((FINALIZE_TURN_PATCH,),), ] upsert_mock.assert_not_called() run_mock.assert_awaited_once_with(workflow_args, cfg) + finalize_turn_mock.assert_not_awaited() + emit_terminal_done_mock.assert_awaited_once_with(None) + assert result == expected_output + + +@pytest.mark.anyio +async def test_run_preserves_v1_finalize_then_done_history_shape() -> None: + role = Role( + type="user", + service_id="tracecat-api", + workspace_id=uuid.uuid4(), + organization_id=uuid.uuid4(), + user_id=uuid.uuid4(), + scopes=frozenset({"agent:execute", "secret:read"}), + ) + workflow_args = _build_workflow_args(role) + workflow_instance = DurableAgentWorkflow(workflow_args) + cfg = cast(Any, workflow_args.agent_args.config) + active_stream_id = uuid.uuid4() + expected_output = AgentOutput( + output="ok", + duration=0.1, + session_id=workflow_args.agent_args.session_id, + ) + call_order: list[str] = [] + + async def terminal_run(*_: object) -> AgentOutput: + workflow_instance.active_stream_id = active_stream_id + return expected_output + + async def finalize(*_: object, **__: object) -> None: + call_order.append("finalize") + + async def emit_done(*_: object) -> None: + call_order.append("done") + + with ( + patch( + "tracecat_ee.agent.workflows.durable.workflow.patched", + side_effect=lambda patch_id: patch_id == FINALIZE_TURN_PATCH, + ), + patch( + "tracecat_ee.agent.workflows.durable.workflow.unsafe.is_replaying", + return_value=False, + ), + patch.object(workflow_instance, "_build_config", AsyncMock(return_value=cfg)), + patch.object( + workflow_instance, + "_run_with_agent_executor", + AsyncMock(side_effect=terminal_run), + ), + patch.object( + workflow_instance, + "_finalize_turn", + AsyncMock(side_effect=finalize), + ) as finalize_turn_mock, + patch.object( + workflow_instance, + "_emit_terminal_done", + AsyncMock(side_effect=emit_done), + ) as emit_terminal_done_mock, + ): + result = await workflow_instance.run(workflow_args) + assert result == expected_output + finalize_turn_mock.assert_awaited_once_with(None, emit_terminal_done=False) + emit_terminal_done_mock.assert_awaited_once_with(active_stream_id) + assert call_order == ["finalize", "done"] + + +@pytest.mark.anyio +async def test_run_does_not_emit_done_after_combined_finalize_failure() -> None: + """A failed DB/Redis finalization must not claim the stream is complete.""" + role = Role( + type="user", + service_id="tracecat-api", + workspace_id=uuid.uuid4(), + organization_id=uuid.uuid4(), + user_id=uuid.uuid4(), + scopes=frozenset({"agent:execute", "secret:read"}), + ) + workflow_args = _build_workflow_args(role) + workflow_instance = DurableAgentWorkflow(workflow_args) + cfg = cast(Any, workflow_args.agent_args.config) + rotated_stream_id = uuid.uuid4() + expected_output = AgentOutput( + output="ok", + duration=0.1, + session_id=workflow_args.agent_args.session_id, + ) + finalize_error = ActivityError( + "finalize failed", + scheduled_event_id=1, + started_event_id=2, + identity="worker", + activity_type="finalize_turn_activity", + activity_id="activity-id", + retry_state=None, + ) + + async def terminal_run(*args: object) -> AgentOutput: + del args + workflow_instance.active_stream_id = rotated_stream_id + return expected_output + + with ( + patch( + "tracecat_ee.agent.workflows.durable.workflow.patched", + side_effect=lambda patch_id: patch_id == FINALIZE_TURN_WITH_END_PATCH, + ), + patch( + "tracecat_ee.agent.workflows.durable.workflow.unsafe.is_replaying", + return_value=False, + ), + patch( + "tracecat_ee.agent.workflows.durable.workflow.info", + return_value=SimpleNamespace( + workflow_id=f"agent/{workflow_args.agent_args.session_id}" + ), + ), + patch.object(workflow_instance, "_build_config", AsyncMock(return_value=cfg)), + patch.object( + workflow_instance, + "_run_with_agent_executor", + AsyncMock(side_effect=terminal_run), + ), + patch( + "tracecat_ee.agent.workflows.durable.workflow.execute_activity", + AsyncMock(side_effect=finalize_error), + ) as execute_activity_mock, + patch( + "tracecat_ee.agent.workflows.durable.workflow.execute_activity_method", + AsyncMock(), + ) as execute_activity_method_mock, + ): + result = await workflow_instance.run(workflow_args) + + assert result == expected_output + execute_activity_mock.assert_awaited_once() + execute_activity_method_mock.assert_not_awaited() + + +@pytest.mark.anyio +@pytest.mark.parametrize( + ("activity_result", "expects_fallback"), + [ + pytest.param(None, True, id="legacy-worker"), + pytest.param( + FinalizeTurnResult(terminal_done_emitted=True), + False, + id="combined-worker", + ), + ], +) +async def test_finalize_turn_falls_back_only_for_legacy_worker_result( + activity_result: FinalizeTurnResult | None, + *, + expects_fallback: bool, +) -> None: + role = Role( + type="user", + service_id="tracecat-api", + workspace_id=uuid.uuid4(), + organization_id=uuid.uuid4(), + user_id=uuid.uuid4(), + scopes=frozenset({"agent:execute", "secret:read"}), + ) + workflow_args = _build_workflow_args(role) + workflow_instance = DurableAgentWorkflow(workflow_args) + active_stream_id = uuid.uuid4() + + with ( + patch( + "tracecat_ee.agent.workflows.durable.workflow.info", + return_value=SimpleNamespace( + workflow_id=f"agent/{workflow_args.agent_args.session_id}" + ), + ), + patch( + "tracecat_ee.agent.workflows.durable.workflow.execute_activity", + AsyncMock(return_value=activity_result), + ) as execute_activity_mock, + patch.object( + workflow_instance, + "_emit_terminal_done", + AsyncMock(), + ) as emit_terminal_done_mock, + ): + await workflow_instance._finalize_turn( + active_stream_id, + emit_terminal_done=True, + ) + + execute_activity_mock.assert_awaited_once() + activity_call = execute_activity_mock.await_args + assert activity_call is not None + finalize_input = activity_call.args[1] + assert isinstance(finalize_input, FinalizeTurnInput) + assert finalize_input.active_stream_id == active_stream_id + assert finalize_input.emit_terminal_done is True + if expects_fallback: + emit_terminal_done_mock.assert_awaited_once_with(active_stream_id) + else: + emit_terminal_done_mock.assert_not_awaited() @pytest.mark.anyio @@ -301,7 +516,7 @@ async def test_run_skips_activity_error_emission_without_patch_marker() -> None: with ( patch( "tracecat_ee.agent.workflows.durable.workflow.patched", - side_effect=[False, False, False, False], + side_effect=[False, False, False, False, False], ) as patched_mock, patch( "tracecat_ee.agent.workflows.durable.workflow.unsafe.is_replaying", @@ -319,6 +534,11 @@ async def test_run_skips_activity_error_emission_without_patch_marker() -> None: "tracecat_ee.agent.workflows.durable.workflow.execute_activity_method", AsyncMock(), ) as execute_activity_mock, + patch.object( + workflow_instance, + "_emit_terminal_done", + AsyncMock(), + ) as emit_terminal_done_mock, ): with pytest.raises(ActivityError): await workflow_instance.run(workflow_args) @@ -330,9 +550,11 @@ async def test_run_skips_activity_error_emission_without_patch_marker() -> None: ((UPSERT_TRACECAT_SEARCH_ATTRIBUTES_PATCH,),), ((EMIT_PRE_STREAM_SESSION_ERRORS_PATCH,),), ((PERSIST_SESSION_ERROR_PATCH,),), + ((FINALIZE_TURN_WITH_END_PATCH,),), ((FINALIZE_TURN_PATCH,),), ] execute_activity_mock.assert_not_awaited() + emit_terminal_done_mock.assert_awaited_once_with(None) @pytest.mark.anyio diff --git a/tracecat/agent/executor/activity.py b/tracecat/agent/executor/activity.py index 977aa9da9..bc2432542 100644 --- a/tracecat/agent/executor/activity.py +++ b/tracecat/agent/executor/activity.py @@ -138,9 +138,6 @@ class AgentExecutorInput(BaseModel): sdk_session_data: str | None = Field(default=None, deprecated=True) # True when resuming after an approval decision. is_approval_continuation: bool = False - # True when the durable workflow will emit stream.done() after approval - # rows are persisted, instead of the executor loopback doing it immediately. - defer_done_on_approval: bool = False # True when forking from parent session (SDK should use fork_session=True) is_fork: bool = False @@ -452,7 +449,6 @@ async def run(self) -> AgentExecutorResult: workspace_id=self.input.workspace_id, active_stream_id=self.input.active_stream_id, curr_run_id=self.input.curr_run_id, - defer_done_on_approval=self.input.defer_done_on_approval, ) handler = LoopbackHandler(input=loopback_input) diff --git a/tracecat/agent/executor/loopback.py b/tracecat/agent/executor/loopback.py index cb6737fe5..5bf145ed2 100644 --- a/tracecat/agent/executor/loopback.py +++ b/tracecat/agent/executor/loopback.py @@ -98,7 +98,6 @@ class LoopbackInput: workspace_id: uuid.UUID active_stream_id: uuid.UUID | None = None curr_run_id: uuid.UUID | None = None - defer_done_on_approval: bool = False @dataclass(kw_only=True, slots=True) @@ -282,7 +281,6 @@ def __init__(self, input: LoopbackInput) -> None: self._stream_sink: LoopbackEventSink | None = None self._result = LoopbackResult(success=False) self._sdk_session_id: str | None = None # Track SDK session ID for this run - self._stream_done_emitted: bool = False # Dedupe flag for stream.done() self._external_stream_done_emitted: bool = False self._interrupt_notice_emitted: bool = False # Dedupe for cancelled event # Track which session lines have been persisted to avoid duplicates @@ -368,40 +366,21 @@ async def prepare(self) -> LoopbackEventSink: self._stream_sink = await self._initialize_stream_sink() return self._stream_sink - async def _emit_stream_done(self) -> None: - """Emit stream.done() exactly once. - - This helper ensures the stream end marker is emitted exactly once, - even if multiple code paths could trigger it (e.g., error + finally). - """ - if self._should_defer_done_for_approval(): - if ( - isinstance(self._stream_sink, FanoutStreamSink) - and not self._external_stream_done_emitted - ): - self._external_stream_done_emitted = True - try: - await self._stream_sink.done_external() - except Exception as e: - logger.warning( - "Failed to emit external stream done", - error=str(e), - ) + async def _close_external_stream(self) -> None: + """Close external sinks exactly once while leaving Redis to the workflow.""" + if ( + not isinstance(self._stream_sink, FanoutStreamSink) + or self._external_stream_done_emitted + ): return - if self._stream_sink and not self._stream_done_emitted: - self._stream_done_emitted = True - try: - await self._stream_sink.done() - except Exception as e: - logger.warning("Failed to emit stream done", error=str(e)) - - def _should_defer_done_for_approval(self) -> bool: - return ( - self.input.defer_done_on_approval - and self._result.approval_requested - and self._result.error is None - and not self._result.cancelled - ) + self._external_stream_done_emitted = True + try: + await self._stream_sink.done_external() + except Exception as e: + logger.warning( + "Failed to emit external stream done", + error=str(e), + ) async def _emit_terminal_stream_error( self, @@ -410,7 +389,7 @@ async def _emit_terminal_stream_error( ) -> None: await self._emit_failed_compaction_if_pending() await stream_sink.error(error) - await self._emit_stream_done() + await self._close_external_stream() self._result.terminal_stream_error_emitted = True def mark_cancelled(self, reason: str) -> None: @@ -556,8 +535,9 @@ async def handle_connection( except TimeoutError: logger.warning("Timeout emitting stream error") finally: - # ALWAYS emit done on any exit path to prevent SSE consumers from hanging - await self._emit_stream_done() + # External channels finish with the runtime. The durable workflow + # owns Redis completion after approval persistence or finalization. + await self._close_external_stream() writer.close() await writer.wait_closed() @@ -881,7 +861,7 @@ async def _handle_done(self) -> bool: ) await self._emit_failed_compaction_if_pending() if self._result.error is not None: - await self._emit_stream_done() + await self._close_external_stream() return True if validation_error := self._validate_runtime_completion(): await self._emit_terminal_stream_error(stream_sink, validation_error) @@ -889,7 +869,7 @@ async def _handle_done(self) -> bool: return True self._result.success = True await self._emit_interrupt_notice_if_cancelled(stream_sink) - await self._emit_stream_done() + await self._close_external_stream() return True async def send_done(self) -> None: diff --git a/tracecat/agent/session/activities.py b/tracecat/agent/session/activities.py index 46e9eabc1..cf053410e 100644 --- a/tracecat/agent/session/activities.py +++ b/tracecat/agent/session/activities.py @@ -451,28 +451,53 @@ class FinalizeTurnInput(BaseModel): role: Role session_id: uuid.UUID run_id: uuid.UUID + # Defaults preserve old workflow inputs while workers roll. New workflows + # set emit_terminal_done=True and pass their captured per-turn stream ID. + active_stream_id: uuid.UUID | None = None + emit_terminal_done: bool = False -@activity.defn -async def finalize_turn_activity(input: FinalizeTurnInput) -> None: - """Clear active-turn pointers at terminal (compare-and-clear by run_id). +class FinalizeTurnResult(BaseModel): + """Result describing whether this worker emitted the terminal stream marker.""" + + terminal_done_emitted: bool + - Idempotent and replay-safe: nulls curr_run_id/active_stream_id only while the - session still points at this run, so a stale terminal never clears a newer - live turn. +@activity.defn +async def finalize_turn_activity( + input: FinalizeTurnInput, +) -> FinalizeTurnResult | None: + """Finalize a turn while bridging old workflow and worker versions. + + New workflows ask this activity to commit pointer cleanup and then append + ``END`` as one retryable operation. Old workflow inputs omit that capability, + so new workers retain the previous database-only behavior. The optional + return type lets new workflows recognize an old worker's legacy ``None`` + result and use their temporary stream-emission fallback. """ ctx_role.set(input.role) try: async with AgentSessionService.with_session(role=input.role) as service: - await service.finalize_turn(input.session_id, input.run_id) + if input.emit_terminal_done: + await service.finalize_turn( + input.session_id, + input.run_id, + active_stream_id=input.active_stream_id, + ) + else: + await service.clear_turn_pointers(input.session_id, input.run_id) except Exception as e: logger.warning( - "Failed to finalize agent turn pointers", + "Failed to finalize agent turn", session_id=str(input.session_id), run_id=str(input.run_id), error=str(e), ) raise + if not input.emit_terminal_done: + # Old workflow code expects this activity to have no result payload. + return None + return FinalizeTurnResult(terminal_done_emitted=True) def get_session_activities() -> list: diff --git a/tracecat/agent/session/service.py b/tracecat/agent/session/service.py index f21665cb3..566edd3d2 100644 --- a/tracecat/agent/session/service.py +++ b/tracecat/agent/session/service.py @@ -1012,12 +1012,12 @@ async def update_last_stream_id( await self.session.refresh(agent_session) return agent_session - async def finalize_turn( + async def clear_turn_pointers( self, session_id: uuid.UUID, run_id: uuid.UUID, ) -> None: - """Clear the active-turn pointers at terminal (compare-and-clear). + """Clear the active-turn pointers with a compare-and-clear update. Nulls ``curr_run_id`` and ``active_stream_id`` only while the session still points at ``run_id``. The compare guards against a newer turn that @@ -1038,6 +1038,33 @@ async def finalize_turn( await self.session.execute(stmt) await self.session.commit() + async def finalize_turn( + self, + session_id: uuid.UUID, + run_id: uuid.UUID, + *, + active_stream_id: uuid.UUID | None, + ) -> None: + """Publish terminal DB state, then close the captured Redis stream. + + When this method returns, this run no longer owns the session's active + pointers and its captured stream contains ``END``. Both operations are + retry-safe: the database update is compare-and-clear by ``run_id``, and + clients tolerate duplicate terminal stream markers. + + Args: + session_id: Agent session being finalized. + run_id: Run that is relinquishing ownership of the active pointers. + active_stream_id: Stream captured by that run before pointer cleanup. + """ + await self.clear_turn_pointers(session_id, run_id) + stream = await AgentStream.new( + workspace_id=self.workspace_id, + session_id=session_id, + stream_id=active_stream_id, + ) + await stream.done() + async def append_cancelled_marker( self, session_id: uuid.UUID,