Allow native chat replies to finish after semantic completion

packages/paperclip-runner/src/native-session-runtime.test.ts

@@ -4375,111 +4375,118 @@ describe("executeNativeSession recovery", () => {
   });
 
   it("retains provider output emitted after a durable semantic result", async () => {
-    const cancel = vi.fn(() => ({ cleanup: Promise.resolve() }));
-    const close = vi.fn(async () => undefined);
-    const events: PrpEvent[] = [];
-    const session: NativeSession = {
-      identity: () => identity,
-      async capabilities() {
-        return {
-          resume: true,
-          typedEvents: true,
-          steering: false,
-          interruption: true,
-          structuredResult: true,
-        };
-      },
-      async *events() {
-        yield runnerEvent(1, "run.result.proposed", result);
-        yield runnerEvent(2, "item.completed", {
-          item: { type: "assistant_message", text: "Final response." },
-        });
-        yield runnerEvent(3, "turn.completed");
-      },
-      async startTurn() {
-        return { turnId: "turn-recovery" };
-      },
-      cancel,
-      async result() {
-        return null;
-      },
-      async snapshot() {
-        return {
-          backendKind: "mock",
-          sessionId: identity.sessionId,
-          identity,
-          providerSessionId: "provider-recovery",
-          cursor: "3",
-          activeTurnId: null,
-          pendingRuntimeRequests: [],
-          lineage: [],
-        };
-      },
-      close,
-    };
-    const backend: NativeSessionBackend = {
-      async descriptor() {
-        return {
-          kind: "mock",
-          name: "semantic-result-final-response-backend",
-          version: "1",
-          capabilities: await session.capabilities(),
-        };
-      },
-      async openSession() {
-        return session;
-      },
-    };
-    const port: ControlPlanePort = {
-      async openRun() {},
-      async checkpointSession() {},
-      async appendEvent(event) {
-        events.push(structuredClone(event as PrpEvent));
-        const sourceEvents = events.filter(
-          (candidate) => candidate.sourceInstanceId === event.sourceInstanceId,
-        );
-        return {
-          cursor: events.length,
-          highestContiguousSourceSeq: highestContiguous(sourceEvents),
-          disposition: "committed",
-        };
-      },
-      async replayEvents(replay) {
-        const sourceEvents = events.filter(
-          (event) => event.sourceInstanceId === replay.sourceInstanceId,
-        );
-        return {
-          events: structuredClone(
-            sourceEvents.filter(
-              (event) => event.sourceSeq > replay.afterSourceSeq,
+    vi.useFakeTimers();
+    try {
+      const cancel = vi.fn(() => ({ cleanup: Promise.resolve() }));
+      const close = vi.fn(async () => undefined);
+      const events: PrpEvent[] = [];
+      const session: NativeSession = {
+        identity: () => identity,
+        async capabilities() {
+          return {
+            resume: true,
+            typedEvents: true,
+            steering: false,
+            interruption: true,
+            structuredResult: true,
+          };
+        },
+        async *events() {
+          yield runnerEvent(1, "run.result.proposed", result);
+          // A streamed reply can finish beyond the ordinary five-second grace.
+          await new Promise<void>((resolve) => setTimeout(resolve, 6_000));
+          yield runnerEvent(2, "item.completed", {
+            item: { type: "assistant_message", text: "Final response." },
+          });
+          yield runnerEvent(3, "turn.completed");
+        },
+        async startTurn() {
+          return { turnId: "turn-recovery" };
+        },
+        cancel,
+        async result() {
+          return null;
+        },
+        async snapshot() {
+          return {
+            backendKind: "mock",
+            sessionId: identity.sessionId,
+            identity,
+            providerSessionId: "provider-recovery",
+            cursor: "3",
+            activeTurnId: null,
+            pendingRuntimeRequests: [],
+            lineage: [],
+          };
+        },
+        close,
+      };
+      const backend: NativeSessionBackend = {
+        async descriptor() {
+          return {
+            kind: "mock",
+            name: "semantic-result-final-response-backend",
+            version: "1",
+            capabilities: await session.capabilities(),
+          };
+        },
+        async openSession() {
+          return session;
+        },
+      };
+      const port: ControlPlanePort = {
+        async openRun() {},
+        async checkpointSession() {},
+        async appendEvent(event) {
+          events.push(structuredClone(event as PrpEvent));
+          const sourceEvents = events.filter(
+            (candidate) => candidate.sourceInstanceId === event.sourceInstanceId,
+          );
+          return {
+            cursor: events.length,
+            highestContiguousSourceSeq: highestContiguous(sourceEvents),
+            disposition: "committed",
+          };
+        },
+        async replayEvents(replay) {
+          const sourceEvents = events.filter(
+            (event) => event.sourceInstanceId === replay.sourceInstanceId,
+          );
+          return {
+            events: structuredClone(
+              sourceEvents.filter(
+                (event) => event.sourceSeq > replay.afterSourceSeq,
+              ),
             ),
-          ),
-          highestContiguousSourceSeq: highestContiguous(sourceEvents),
-        };
-      },
-      async completeRun() {},
-    };
+            highestContiguousSourceSeq: highestContiguous(sourceEvents),
+          };
+        },
+        async completeRun() {},
+      };
 
-    await expect(
-      executeNativeSession({
+      const completed = executeNativeSession({
         input,
         backend,
         controlPlane: port,
         runnerInstanceId: "runner-recovery",
         controlPlaneInstanceId: "control-recovery",
-        semanticResultTerminalGraceMs: 50,
-      }),
-    ).resolves.toMatchObject({ result, terminal });
+        semanticResultTerminalGraceMs: 30_000,
+      });
+      await vi.advanceTimersByTimeAsync(6_000);
+      await expect(completed).resolves.toMatchObject({ result, terminal });
 
-    expect(cancel).not.toHaveBeenCalled();
-    expect(close).toHaveBeenCalledOnce();
-    expect(events.map((event) => event.eventType)).toEqual([
-      "run.result.proposed",
-      "item.completed",
-      "turn.completed",
-      "run.result.accepted",
-      "run.terminal",
-    ]);
+      expect(cancel).not.toHaveBeenCalled();
+      expect(close).toHaveBeenCalledOnce();
+      expect(events.map((event) => event.eventType)).toEqual([
+        "run.result.proposed",
+        "item.completed",
+        "turn.completed",
+        "run.result.accepted",
+        "run.terminal",
+      ]);
+    } finally {
+      vi.useRealTimers();
+    }
   });
 
   it("rejects a mismatched checkpoint before it mutates control-plane state", async () => {

server/src/services/heartbeat.ts

@@ -20988,6 +20988,7 @@ export function heartbeatService(
                     executePaperclipNativeSession({
                       db,
                       execution: nativeExecution,
+                      conversationMode: isConversation(issueContext),
                       runnerInstanceId: nativeRunnerInstanceId,
                       leaseOwner: runOptions.nativeLeaseOwner,
                       restartRecovery: runOptions.nativeRestartRecovery,

server/src/services/native-runtime/native-session-executor.test.ts

@@ -3135,6 +3135,19 @@ describe("native session same-turn steering", () => {
 });
 
 describe("native warm session supervision", () => {
+  it.each([true, false])("preserves chat reply grace for per-turn providers: chat=%s", async (conversationMode) => {
+    state.execute.mockReset().mockImplementationOnce(async (options) => {
+      expect(options.semanticResultTerminalGraceMs).toBe(conversationMode ? 30_000 : undefined);
+      return {
+        result: { summary: "Reply completed" },
+        terminal: { runTerminalState: "succeeded" },
+        turnId: "turn-grace", normalizedSessionId: execution.session.normalizedSessionId,
+        providerSessionId: "provider-grace", driverKind: "test", driverVersion: "1",
+        nativeEventCount: 1, highestContiguousSourceSeq: 1,
+      };
+    });
+    await executePaperclipNativeSession({ db: leaseDb(), execution, runnerInstanceId: "runner", conversationMode });
+  });
   it("persists agent-created goal continuity before a per-turn runner settles", async () => {
     const goalCheckpoint = {
       identity: { runId: execution.binding.runId, sessionId: "session" },

server/src/services/native-runtime/native-session-executor.ts

@@ -195,12 +195,11 @@ const TERMINAL_HEARTBEAT_RUN_STATUSES = new Set([
 const NATIVE_SESSION_EXECUTION_LEASE_TTL_MS = 20 * 60_000;
 const NATIVE_SESSION_EXECUTION_LEASE_RENEW_INTERVAL_MS = 5 * 60_000;
 const NATIVE_SESSION_CANCELLATION_CLEANUP_GRACE_MS = 2_000;
-// A reusable provider must publish its terminal suffix before the next run can
-// rotate PRP authority. Remote Codex can take more than the ordinary five-second
-// result grace to flush its final answer over Daytona, so retain the bounded
-// turn long enough to reach a naturally quiescent, reusable state. This adds no
-// delay when the provider terminates normally.
-const NATIVE_WARM_SEMANTIC_RESULT_TERMINAL_GRACE_MS = 30_000;
+// Chat replies and reusable providers need their final output before the next
+// run rotates PRP authority. Streaming answers can take longer than the ordinary
+// five-second result grace, so retain the bounded turn until it is quiescent.
+// This adds no delay when the provider terminates normally.
+const NATIVE_RESPONSE_SEMANTIC_RESULT_TERMINAL_GRACE_MS = 30_000;
 const NATIVE_RUNTIME_REQUEST_RESOLUTION_CACHE_MAX = 256;
 type NativeRuntimeRequestResolution = {
   runId: string;
@@ -4042,6 +4041,8 @@ export async function executePaperclipNativeSession(input: {
   db: Db;
   execution: NativeExecutionInput;
   runnerInstanceId: string;
+  /** Trusted task identity from the heartbeat orchestration. */
+  conversationMode?: boolean;
   leaseOwner?: string;
   restartRecovery?: NativeRestartRecoveryClaim;
   onSpawn?: (meta: {
@@ -4958,9 +4959,9 @@ async function executePaperclipNativeSessionWithinScope(
             sessionGoalControl: input.sessionGoalControl,
             resumeSessionGoalHeartbeat: input.resumeSessionGoalHeartbeat,
             semanticResultTerminalGraceMs:
-              warmSessionId === null
-                ? undefined
-                : NATIVE_WARM_SEMANTIC_RESULT_TERMINAL_GRACE_MS,
+              warmSessionId !== null || input.conversationMode === true
+                ? NATIVE_RESPONSE_SEMANTIC_RESULT_TERMINAL_GRACE_MS
+                : undefined,
             // Every durable runner must finish its bounded suspension before
             // the next run verifies and rotates the saved authority.
             requireSessionCloseBeforeReturn: runnerdBackend !== null,