fix(agents): emit terminal stream END only after finalize_turn

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:

packages/tracecat-ee/tracecat_ee/agent/workflows/durable.py

@@ -492,9 +492,6 @@ def _preserved_agents_binding(
 
 FINALIZE_TURN_PATCH = "durable-agent-finalize-turn-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"
 
 
 @workflow.defn
@@ -517,7 +514,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
@@ -922,8 +918,10 @@ async def run(self, args: AgentWorkflowArgs) -> AgentOutput:
             # 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.
+            terminal_stream_id = self.active_stream_id
             if workflow.patched(FINALIZE_TURN_PATCH):
                 await self._finalize_turn()
+            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)."""
@@ -1039,6 +1037,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 captured terminal stream after durable state is publishable."""
+        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,8 +1270,6 @@ 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)
-
         # Prepare executor input
         executor_input = AgentExecutorInput(
             session_id=self.session_id,
@@ -1265,7 +1285,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 +1395,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 +1533,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,
                 )