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

packages/tracecat-ee/tracecat_ee/agent/activities.py

@@ -136,6 +136,8 @@ class EmitSessionErrorInputs(BaseModel):
     # error path has already streamed the error inline via the loopback, so it
     # persists-only; pre-stream failures stream too.
     should_stream: bool = True
+    # The workflow emits END after finalizing the turn when this is true.
+    defer_done: bool = False
 
 
 class EmitSessionDoneInputs(BaseModel):
@@ -167,12 +169,14 @@ 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. END is emitted here only when
+    ``defer_done`` is false.
     """
+    # The workflow emits END after finalizing the turn when this is true.
+    defer_done: bool = False
 
 
 class _SessionStreamInputs(Protocol):
@@ -533,7 +537,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. A patch-gated workflow can set
+        ``defer_done`` so this activity leaves END to the post-finalize boundary.
 
         Best-effort: a persistence failure must not mask the agent's real error
         or abort propagation, so it is logged and swallowed.
@@ -565,7 +570,8 @@ async def emit_session_error(self, args: EmitSessionErrorInputs) -> None:
 
         stream = await self._open_session_stream(args)
         await stream.error(args.message)
-        await stream.done()
+        if not args.defer_done:
+            await stream.done()
 
     @staticmethod
     async def _open_session_stream(args: _SessionStreamInputs) -> AgentStream:
@@ -590,8 +596,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``). ``defer_done`` leaves END to
+        the workflow's post-finalize boundary.
         """
         # Local import: tracecat.agent.session.service imports tracecat_ee
         # modules, so a top-level import here would create a cycle.
@@ -616,7 +623,8 @@ async def emit_session_cancelled(self, args: EmitSessionCancelledInputs) -> None
                 tool_call_ids=args.interrupted_tool_call_ids,
             )
         )
-        await stream.done()
+        if not args.defer_done:
+            await stream.done()
 
     @activity.defn
     async def execute_remote_mcp_tool(self, args: ExecuteRemoteMCPToolArgs) -> str:

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

@@ -491,6 +491,7 @@ def _preserved_agents_binding(
 
 
 FINALIZE_TURN_PATCH = "durable-agent-finalize-turn-v1"
+TERMINAL_END_AFTER_FINALIZE_PATCH = "durable-agent-terminal-end-after-finalize-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.
@@ -518,6 +519,7 @@ def __init__(self, args: AgentWorkflowArgs):
         self.session_id = args.agent_args.session_id
         self.active_stream_id = args.agent_args.active_stream_id
         self._approval_stream_v2 = False
+        self._terminal_end_after_finalize = False
         self.harness_type = args.harness_type or "claude_code"
         self.approvals = ApprovalManager(role=self.role)
         self.max_requests = args.agent_args.max_requests
@@ -885,6 +887,9 @@ async def run(self, args: AgentWorkflowArgs) -> AgentOutput:
         """Run the agent until completion. The agent will call tools until it needs human approval."""
         if workflow.patched(UPSERT_TRACECAT_SEARCH_ATTRIBUTES_PATCH):
             self._upsert_tracecat_search_attributes()
+        self._terminal_end_after_finalize = workflow.patched(
+            TERMINAL_END_AFTER_FINALIZE_PATCH
+        )
         logger.debug(
             "DurableAgentWorkflow run", args=args, harness_type=self.harness_type
         )
@@ -922,8 +927,11 @@ 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()
+            if self._terminal_end_after_finalize:
+                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)."""
@@ -981,6 +989,7 @@ async def _finalize_session_error(
                     # None falls back to the per-session key for non-chat turns.
                     active_stream_id=self.active_stream_id,
                     should_stream=should_stream,
+                    defer_done=self._terminal_end_after_finalize,
                 ),
                 start_to_close_timeout=timedelta(seconds=10),
                 retry_policy=RETRY_POLICIES["activity:fail_fast"],
@@ -1021,6 +1030,7 @@ async def _emit_session_cancelled(
                     # None falls back to the per-session key for non-chat turns.
                     active_stream_id=self.active_stream_id,
                     emit_stream=emit_stream,
+                    defer_done=self._terminal_end_after_finalize,
                     interrupted_tool_call_ids=interrupted_tool_call_ids,
                     # Pin the marker to this run explicitly: the session row's
                     # curr_run_id may already point at a newer turn by the time
@@ -1039,6 +1049,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),
+                # Idempotent: a retry may append a duplicate END after an
+                # ambiguous failure, 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:
@@ -1266,6 +1300,7 @@ async def _run_with_agent_executor(
             sdk_session_id=load_result.sdk_session_id,
             sdk_session_data=load_result.sdk_session_data,
             defer_done_on_approval=self._approval_stream_v2,
+            defer_done_on_terminal=self._terminal_end_after_finalize,
             is_fork=load_result.is_fork,
         )
 
@@ -1516,6 +1551,7 @@ async def _run_with_agent_executor(
                     sdk_session_id=reload_result.sdk_session_id,
                     sdk_session_data=reload_result.sdk_session_data,
                     defer_done_on_approval=self._approval_stream_v2,
+                    defer_done_on_terminal=self._terminal_end_after_finalize,
                     is_fork=reload_result.is_fork,
                     is_approval_continuation=True,
                 )