From eba3248d86491b06eec043988b7d133331b0dcac Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Fri, 17 Jul 2026 13:39:22 +0100 Subject: [PATCH 1/5] feat(cloud): add native steering to agent adapters --- .../claude/claude-agent.streamed-text.test.ts | 58 ++++++++ .../agent/src/adapters/claude/claude-agent.ts | 31 ++++- packages/agent/src/adapters/claude/types.ts | 1 + .../app-server-client.test.ts | 9 +- .../codex-app-server/app-server-client.ts | 19 ++- .../codex-app-server-agent.test.ts | 126 +++++++++++++++++- .../codex-app-server-agent.ts | 52 +++++--- .../codex-app-server/turn-controller.ts | 5 - 8 files changed, 273 insertions(+), 28 deletions(-) diff --git a/packages/agent/src/adapters/claude/claude-agent.streamed-text.test.ts b/packages/agent/src/adapters/claude/claude-agent.streamed-text.test.ts index 96775fa040..4bd3d8e13f 100644 --- a/packages/agent/src/adapters/claude/claude-agent.streamed-text.test.ts +++ b/packages/agent/src/adapters/claude/claude-agent.streamed-text.test.ts @@ -243,6 +243,64 @@ describe("ClaudeAcpAgent.prompt — streamed assistant text wiring", () => { ]); }); + it("keeps the original turn open until a pending steer is consumed", async () => { + const { agent, client } = makeAgent(); + const sessionId = "s-steer-ordering"; + const { query, input } = installFakeSession(agent, sessionId); + + const promptPromise = agent.prompt({ + sessionId, + prompt: [{ type: "text", text: "use orange" }], + }); + let promptSettled = false; + void promptPromise.then(() => { + promptSettled = true; + }); + await tick(); + await echoUserMessage(query, input); + + const steerResult = await agent.prompt({ + sessionId, + prompt: [{ type: "text", text: "use green instead" }], + _meta: { steer: true }, + }); + expect(steerResult._meta).toEqual({ steer: true }); + + await send(query, assistantMessage(sessionId, "msg_orange", "ORANGE")); + await send(query, resultSuccess(sessionId)); + expect(promptSettled).toBe(false); + + await echoUserMessage(query, input); + await send(query, assistantMessage(sessionId, "msg_green", "GREEN")); + await send(query, resultSuccess(sessionId)); + + await expect(promptPromise).resolves.toMatchObject({ + stopReason: "end_turn", + }); + expect(messageChunkTexts(client.sessionUpdate.mock.calls)).toEqual([ + "ORANGE", + "GREEN", + ]); + }); + + it("declines an explicit steer after the active turn has ended", async () => { + const { agent } = makeAgent(); + const sessionId = "s-expired-steer"; + installFakeSession(agent, sessionId); + + await expect( + agent.prompt({ + sessionId, + prompt: [{ type: "text", text: "too late" }], + _meta: { steer: true }, + }), + ).resolves.toMatchObject({ _meta: { steer: false } }); + + const session = (agent as unknown as { session: { turnQueue: unknown[] } }) + .session; + expect(session.turnQueue).toHaveLength(0); + }); + it("reconnects a disconnected signed-commit server before the turn", async () => { const { agent } = makeAgent(); const sessionId = "s-heal"; diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index 402d8dbdd4..983567d8a5 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -482,12 +482,20 @@ export class ClaudeAcpAgent extends BaseAcpAgent { const hasInFlightTurns = this.session.activeTurn !== null || this.session.turnQueue.length > 0; - if (hasInFlightTurns && isSteerMeta(params._meta)) { + const isSteer = isSteerMeta(params._meta); + if (hasInFlightTurns && isSteer) { // Fold into the running turn (promptToClaude tagged it priority:"next"); // the benign end_turn is ignored by clients, which key off _meta.steer. + const owner = + this.session.activeTurn ?? + this.session.turnQueue.find((turn) => !turn.settled); + owner?.pendingSteerUuids.add(promptUuid); this.session.input.push(userMessage); await this.broadcastUserMessage(params); - return { stopReason: "end_turn" }; + return { stopReason: "end_turn", _meta: { steer: true } }; + } + if (isSteer) { + return { stopReason: "end_turn", _meta: { steer: false } }; } if (!hasInFlightTurns && !isLocalOnlyCommand) { @@ -507,6 +515,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent { const turn: Turn = { promptUuid, + pendingSteerUuids: new Set(), isLocalOnlyCommand, commandName: commandMatch?.[1], broadcast: () => this.broadcastUserMessage(params), @@ -1059,6 +1068,21 @@ export class ClaudeAcpAgent extends BaseAcpAgent { }, ); + if ( + !isTaskNotification && + session.activeTurn && + session.activeTurn.pendingSteerUuids.size > 0 + ) { + this.logger.debug( + "Deferring turn completion until pending steers are consumed", + { + sessionId, + pendingSteers: session.activeTurn.pendingSteerUuids.size, + }, + ); + break; + } + if ( (message as { stop_reason?: string }).stop_reason === "refusal" ) { @@ -1198,6 +1222,9 @@ export class ClaudeAcpAgent extends BaseAcpAgent { // active one first), then drops from the feed. Runs before the // cancelled guard so a turn enqueued after a cancel still starts. if (message.type === "user" && "uuid" in message && message.uuid) { + if (session.activeTurn?.pendingSteerUuids.delete(message.uuid)) { + break; + } const queued = session.turnQueue.find( (t) => t.promptUuid === message.uuid && !t.settled, ); diff --git a/packages/agent/src/adapters/claude/types.ts b/packages/agent/src/adapters/claude/types.ts index 80078dfa8f..37d46ac0d6 100644 --- a/packages/agent/src/adapters/claude/types.ts +++ b/packages/agent/src/adapters/claude/types.ts @@ -42,6 +42,7 @@ export type BackgroundTerminal = /** One in-flight `prompt()` call, settled by the session's consumer. */ export type Turn = { promptUuid: string; + pendingSteerUuids: Set; isLocalOnlyCommand: boolean; commandName?: string; /** Invoked once at activation, matching the pre-consumer broadcast timing. */ diff --git a/packages/agent/src/adapters/codex-app-server/app-server-client.test.ts b/packages/agent/src/adapters/codex-app-server/app-server-client.test.ts index 9501f019ce..a46cc96091 100644 --- a/packages/agent/src/adapters/codex-app-server/app-server-client.test.ts +++ b/packages/agent/src/adapters/codex-app-server/app-server-client.test.ts @@ -4,7 +4,7 @@ import { createBidirectionalStreams, type StreamPair, } from "../../utils/streams"; -import { AppServerClient } from "./app-server-client"; +import { AppServerClient, AppServerRequestError } from "./app-server-client"; interface RpcMessage { id?: number | string; @@ -86,7 +86,12 @@ describe("AppServerClient", () => { error: { code: -32001, message: "Server overloaded; retry later." }, }); - await expect(pending).rejects.toThrow("Server overloaded; retry later."); + const error = await pending.catch((requestError: unknown) => requestError); + expect(error).toBeInstanceOf(AppServerRequestError); + expect(error).toMatchObject({ + code: -32001, + message: "Server overloaded; retry later.", + }); await client.close(); }); diff --git a/packages/agent/src/adapters/codex-app-server/app-server-client.ts b/packages/agent/src/adapters/codex-app-server/app-server-client.ts index ff32050f5c..7d911173e2 100644 --- a/packages/agent/src/adapters/codex-app-server/app-server-client.ts +++ b/packages/agent/src/adapters/codex-app-server/app-server-client.ts @@ -24,6 +24,17 @@ export interface AppServerRpc { close(): Promise; } +export class AppServerRequestError extends Error { + constructor( + readonly code: number, + message: string, + readonly data?: unknown, + ) { + super(message); + this.name = "AppServerRequestError"; + } +} + /** * Bidirectional newline-delimited JSON-RPC client for the native Codex `app-server` subprocess. * Transport-agnostic via a {@link StreamPair} so tests can drive it over in-memory streams. @@ -173,7 +184,13 @@ export class AppServerClient implements AppServerRpc { } this.pending.delete(message.id); if (message.error) { - call.reject(new Error(message.error.message)); + call.reject( + new AppServerRequestError( + message.error.code, + message.error.message, + message.error.data, + ), + ); } else { call.resolve(message.result); } diff --git a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts index 9761794839..7062d98f70 100644 --- a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts +++ b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.test.ts @@ -10,6 +10,7 @@ import type { AppServerClientHandlers, AppServerRpc, } from "./app-server-client"; +import { AppServerRequestError } from "./app-server-client"; import { CodexAppServerAgent } from "./codex-app-server-agent"; import { sandboxPolicyFor } from "./session-config"; @@ -48,6 +49,9 @@ function makeStubRpc(responses: Record) { }; } const response = responses[method]; + if (response instanceof Error) { + throw response; + } return ( typeof response === "function" ? await response(params) @@ -2133,7 +2137,12 @@ describe("CodexAppServerAgent", () => { prompt: [{ type: "text", text: "more context" }], } as unknown as PromptRequest); - // The single turn/completed resolves both the original and the folded prompt. + await expect(second).resolves.toMatchObject({ + stopReason: "end_turn", + _meta: { steer: true }, + }); + + // The original prompt remains the sole owner of turn completion. stub.emit("thread/tokenUsage/updated", { tokenUsage: { last: { @@ -2145,12 +2154,11 @@ describe("CodexAppServerAgent", () => { }, }); stub.emit("turn/completed", { turn: { status: "completed" } }); - const [firstResult, secondResult] = await Promise.all([first, second]); + const firstResult = await first; expect(firstResult).toMatchObject({ stopReason: "end_turn", usage: { totalTokens: 45 }, }); - expect(secondResult).toEqual({ stopReason: "end_turn" }); const steer = stub.requests.find((r) => r.method === "turn/steer"); expect(steer?.params).toMatchObject({ @@ -2164,6 +2172,118 @@ describe("CodexAppServerAgent", () => { ); }); + it("rejects a failed turn/steer without echoing or acknowledging it", async () => { + const stub = makeStubRpc({ + "thread/start": { thread: { id: "t" } }, + "turn/start": { turn: { id: "turn_1" } }, + "turn/steer": new Error("steer transport failed"), + }); + const { client, sessionUpdates } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/x/codex" }, + rpcFactory: stub.factory, + }); + + await agent.newSession({ cwd: "/r" } as unknown as NewSessionRequest); + const first = agent.prompt({ + sessionId: "t", + prompt: [{ type: "text", text: "one" }], + } as unknown as PromptRequest); + stub.emit("turn/started", { threadId: "t", turn: { id: "turn_1" } }); + + await expect( + agent.prompt({ + sessionId: "t", + prompt: [{ type: "text", text: "lost steer" }], + _meta: { steer: true }, + } as unknown as PromptRequest), + ).rejects.toThrow("steer transport failed"); + expect(sessionUpdates).not.toContainEqual( + expect.objectContaining({ + update: expect.objectContaining({ + sessionUpdate: "user_message_chunk", + content: { type: "text", text: "lost steer" }, + }), + }), + ); + + stub.emit("turn/completed", { turn: { status: "completed" } }); + await first; + }); + + it("declines a stale turn/steer so the caller can queue it normally", async () => { + const stub = makeStubRpc({ + "thread/start": { thread: { id: "t" } }, + "turn/start": { turn: { id: "turn_1" } }, + "turn/steer": new AppServerRequestError( + -32600, + "expected active turn id `turn_1` but found `turn_2`", + ), + }); + const { client, sessionUpdates } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/x/codex" }, + rpcFactory: stub.factory, + }); + + await agent.newSession({ cwd: "/r" } as unknown as NewSessionRequest); + const first = agent.prompt({ + sessionId: "t", + prompt: [{ type: "text", text: "one" }], + } as unknown as PromptRequest); + stub.emit("turn/started", { threadId: "t", turn: { id: "turn_1" } }); + + await expect( + agent.prompt({ + sessionId: "t", + prompt: [{ type: "text", text: "queue me" }], + _meta: { steer: true }, + } as unknown as PromptRequest), + ).resolves.toMatchObject({ _meta: { steer: false } }); + expect(sessionUpdates).not.toContainEqual( + expect.objectContaining({ + update: expect.objectContaining({ + sessionUpdate: "user_message_chunk", + content: { type: "text", text: "queue me" }, + }), + }), + ); + + stub.emit("turn/completed", { turn: { status: "completed" } }); + await first; + }); + + it("declines an explicit steer after the active turn has ended", async () => { + const stub = makeStubRpc({ + "thread/start": { thread: { id: "t" } }, + }); + const { client, sessionUpdates } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/x/codex" }, + rpcFactory: stub.factory, + }); + + await agent.newSession({ cwd: "/r" } as unknown as NewSessionRequest); + await expect( + agent.prompt({ + sessionId: "t", + prompt: [{ type: "text", text: "too late" }], + _meta: { steer: true }, + } as unknown as PromptRequest), + ).resolves.toMatchObject({ _meta: { steer: false } }); + expect( + stub.requests.filter((request) => request.method === "turn/start"), + ).toHaveLength(0); + expect(sessionUpdates).not.toContainEqual( + expect.objectContaining({ + update: expect.objectContaining({ + sessionUpdate: "user_message_chunk", + content: { type: "text", text: "too late" }, + }), + }), + ); + }); + it("refreshes the live turnId from each turn/steer response", async () => { const stub = makeStubRpc({ "thread/start": { thread: { id: "t" } }, diff --git a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts index 823a23c6e3..5029f2127a 100644 --- a/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts +++ b/packages/agent/src/adapters/codex-app-server/codex-app-server-agent.ts @@ -50,6 +50,7 @@ import { resolveSpokenNarration } from "../session-meta"; import { AppServerClient, type AppServerClientHandlers, + AppServerRequestError, type AppServerRpc, } from "./app-server-client"; import { handleServerRequest } from "./approvals"; @@ -88,6 +89,16 @@ import { parseStructuredOutput } from "./structured-output"; import { TurnController } from "./turn-controller"; import { UsageTracker } from "./usage-tracker"; +function isStaleTurnSteerError(error: unknown): boolean { + if (!(error instanceof AppServerRequestError) || error.code !== -32600) { + return false; + } + return ( + error.message === "no active turn to steer" || + /^expected active turn id `.*` but found `.*`$/.test(error.message) + ); +} + type AppServerSessionMeta = { // The host sends either a plain string or the Claude-style `{ append }` form. systemPrompt?: string | { append?: string }; @@ -753,32 +764,43 @@ export class CodexAppServerAgent extends BaseAcpAgent { if (dropped > 0) { this.logger.warn("Dropped non-text/non-image prompt blocks", { dropped }); } - // Echo the user prompt (codex emits none), for fresh turns and steering alike. - this.broadcastUserInput(params.prompt); - if (this.turns.isRunning) { // A turn is already running: fold the message in via turn/steer (precondition: the // active turnId). Refresh from the response's rotated turnId so a later steer/interrupt // still targets the live turn (no turn/started is re-emitted for a steer). - const steerRes = await this.rpc - .request<{ turnId?: string }>(APP_SERVER_METHODS.TURN_STEER, { - threadId: this.threadId, - input, - expectedTurnId: this.turns.activeTurnId, - }) - .catch((err) => { - this.logger.warn("turn/steer failed", err); - return undefined; - }); + let steerRes: { turnId?: string }; + try { + steerRes = await this.rpc.request<{ turnId?: string }>( + APP_SERVER_METHODS.TURN_STEER, + { + threadId: this.threadId, + input, + expectedTurnId: this.turns.activeTurnId, + }, + ); + } catch (error) { + if ( + (params._meta as { steer?: unknown } | undefined)?.steer === true && + isStaleTurnSteerError(error) + ) { + return { stopReason: "end_turn", _meta: { steer: false } }; + } + throw error; + } this.turns.onSteered(steerRes?.turnId); - const response = await this.turns.awaitCompletion(); - return { stopReason: response.stopReason }; + this.broadcastUserInput(params.prompt); + return { stopReason: "end_turn", _meta: { steer: true } }; + } + if ((params._meta as { steer?: unknown } | undefined)?.steer === true) { + return { stopReason: "end_turn", _meta: { steer: false } }; } if (this.turns.isPending) { // A turn is pending but has no turnId yet, so we can't steer; fail fast. throw new Error("prompt() called while a turn is already in progress"); } + // Codex does not echo user input, so emit it only once delivery can proceed. + this.broadcastUserInput(params.prompt); const response = await this.runTurn(input); return this.maybeOfferPlanImplementation(response); } diff --git a/packages/agent/src/adapters/codex-app-server/turn-controller.ts b/packages/agent/src/adapters/codex-app-server/turn-controller.ts index e6f6483a9a..90e7bcb21b 100644 --- a/packages/agent/src/adapters/codex-app-server/turn-controller.ts +++ b/packages/agent/src/adapters/codex-app-server/turn-controller.ts @@ -48,11 +48,6 @@ export class TurnController { if (typeof id === "string") this.turnId = id; } - /** Await the in-flight turn's completion (the steer path reuses the original). */ - awaitCompletion(): Promise { - return this.completion ?? Promise.resolve({ stopReason: "end_turn" }); - } - /** Atomically claim the pending turn (clears the slot + turnId synchronously), or undefined if already claimed. */ claim(): PendingTurn | undefined { const pending = this.pending; From a04d2b18844afebd8b3a9e431771f325495f4418 Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Fri, 17 Jul 2026 13:39:31 +0100 Subject: [PATCH 2/5] feat(cloud): negotiate idempotent steer delivery --- .../agent/src/server/agent-server.test.ts | 258 +++++++++++ packages/agent/src/server/agent-server.ts | 430 +++++++++++------- packages/agent/src/server/schemas.ts | 2 + .../api-client/src/posthog-client.test.ts | 13 +- packages/api-client/src/posthog-client.ts | 22 +- packages/shared/src/sessions.test.ts | 9 +- packages/shared/src/sessions.ts | 7 +- 7 files changed, 569 insertions(+), 172 deletions(-) diff --git a/packages/agent/src/server/agent-server.test.ts b/packages/agent/src/server/agent-server.test.ts index 087f2b55b7..bce20573d8 100644 --- a/packages/agent/src/server/agent-server.test.ts +++ b/packages/agent/src/server/agent-server.test.ts @@ -2267,6 +2267,106 @@ describe("AgentServer HTTP Mode", () => { expect(prompt).toHaveBeenCalledTimes(4); }, 20000); + it("steers an active turn without emitting a separate turn completion", async () => { + const s = createServer(); + await s.start(); + const prompt = vi.fn(async () => ({ + stopReason: "end_turn", + _meta: { steer: true }, + })); + const broadcastTurnComplete = vi.fn(); + const resetTurnMessages = vi.fn(); + const serverInternals = s as unknown as { + activeOwnedTurnCount: number; + broadcastTurnComplete: typeof broadcastTurnComplete; + session: { + clientConnection: { prompt: typeof prompt }; + logWriter: { resetTurnMessages: typeof resetTurnMessages }; + }; + }; + serverInternals.activeOwnedTurnCount = 1; + serverInternals.broadcastTurnComplete = broadcastTurnComplete; + serverInternals.session.clientConnection.prompt = prompt; + serverInternals.session.logWriter.resetTurnMessages = resetTurnMessages; + + const response = await fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${createToken()}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: "steer-1", + method: "user_message", + params: { content: "change direction", steer: true }, + }), + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + result: { stopReason: "steered", steered: true }, + }); + expect(prompt).toHaveBeenCalledWith( + expect.objectContaining({ + _meta: expect.objectContaining({ steer: true }), + }), + ); + expect(broadcastTurnComplete).not.toHaveBeenCalled(); + expect(resetTurnMessages).not.toHaveBeenCalled(); + }, 20000); + + it("declines steering without blocking on a fallback normal turn", async () => { + const s = createServer(); + await s.start(); + const prompt = vi.fn(); + const broadcastTurnComplete = vi.fn(); + const resetTurnMessages = vi.fn(); + const serverInternals = s as unknown as { + activeOwnedTurnCount: number; + broadcastTurnComplete: typeof broadcastTurnComplete; + session: { + clientConnection: { prompt: typeof prompt }; + logWriter: { resetTurnMessages: typeof resetTurnMessages }; + }; + }; + serverInternals.activeOwnedTurnCount = 1; + prompt.mockImplementationOnce(async () => { + serverInternals.activeOwnedTurnCount = 0; + return { stopReason: "end_turn", _meta: { steer: false } }; + }); + serverInternals.broadcastTurnComplete = broadcastTurnComplete; + serverInternals.session.clientConnection.prompt = prompt; + serverInternals.session.logWriter.resetTurnMessages = resetTurnMessages; + + const response = await fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${createToken()}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: "steer-race", + method: "user_message", + params: { content: "continue normally", steer: true }, + }), + }); + + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ + result: { stopReason: "steer_declined", steered: false }, + }); + expect(prompt).toHaveBeenCalledTimes(1); + expect(prompt.mock.calls[0]?.[0]).toEqual( + expect.objectContaining({ + _meta: expect.objectContaining({ steer: true }), + }), + ); + expect(resetTurnMessages).not.toHaveBeenCalled(); + expect(broadcastTurnComplete).not.toHaveBeenCalled(); + }, 20000); + it("redelivers a messageId whose first delivery failed before producing a turn", async () => { const s = createServer(); await s.start(); @@ -2307,6 +2407,163 @@ describe("AgentServer HTTP Mode", () => { expect(prompt).toHaveBeenCalledTimes(2); }, 20000); + it("keeps a recoverable delivery committed across an ambiguous retry", async () => { + const s = createServer(); + await s.start(); + const prompt = vi + .fn() + .mockRejectedValue(new Error("API Error: The operation timed out.")); + const serverInternals = s as unknown as { + session: { clientConnection: { prompt: typeof prompt } }; + }; + serverInternals.session.clientConnection.prompt = prompt; + + const token = createToken(); + const send = async (requestId: string) => + fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: requestId, + method: "user_message", + params: { + content: "do the thing", + messageId: "m-recoverable", + }, + }), + }); + + const first = await send("first-attempt"); + await expect(first.json()).resolves.toMatchObject({ + result: { stopReason: "error_recoverable" }, + }); + expect(prompt).toHaveBeenCalledTimes(1); + + const retry = await send("ambiguous-retry"); + await expect(retry.json()).resolves.toMatchObject({ + result: { stopReason: "duplicate_delivery", duplicate: true }, + }); + expect(prompt).toHaveBeenCalledTimes(1); + }, 20000); + + it("shares a failed in-flight messageId outcome with concurrent retries", async () => { + const s = createServer(); + await s.start(); + let rejectFirstDelivery!: (error: Error) => void; + const prompt = vi + .fn() + .mockImplementationOnce( + () => + new Promise((_resolve, reject) => { + rejectFirstDelivery = reject; + }), + ) + .mockResolvedValueOnce({ stopReason: "end_turn" }); + const serverInternals = s as unknown as { + logger: { info: (...args: unknown[]) => void }; + session: { clientConnection: { prompt: typeof prompt } }; + }; + serverInternals.session.clientConnection.prompt = prompt; + const infoLog = vi.spyOn(serverInternals.logger, "info"); + + const token = createToken(); + const send = async (requestId: string) => + fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: requestId, + method: "user_message", + params: { content: "do the thing", messageId: "m-concurrent" }, + }), + }); + + const firstResponse = send("first-attempt"); + await vi.waitFor(() => expect(prompt).toHaveBeenCalledTimes(1)); + + let retrySettled = false; + const retryResponse = send("concurrent-retry").finally(() => { + retrySettled = true; + }); + await vi.waitFor(() => { + expect(infoLog).toHaveBeenCalledWith( + "Awaiting in-flight user_message delivery", + { messageId: "m-concurrent" }, + ); + expect(prompt).toHaveBeenCalledTimes(1); + expect(retrySettled).toBe(false); + }); + + rejectFirstDelivery(new Error("sdk connection lost")); + const [first, retry] = await Promise.all([firstResponse, retryResponse]); + await expect(first.json()).resolves.toMatchObject({ + error: { message: "sdk connection lost" }, + }); + await expect(retry.json()).resolves.toMatchObject({ + error: { message: "sdk connection lost" }, + }); + expect(prompt).toHaveBeenCalledTimes(1); + }, 20000); + + it("keeps an accepted messageId committed when teardown clears the active session", async () => { + const s = createServer(); + await s.start(); + let finishPrompt!: (result: { stopReason: "end_turn" }) => void; + const prompt = vi.fn( + () => + new Promise<{ stopReason: "end_turn" }>((resolve) => { + finishPrompt = resolve; + }), + ); + const serverInternals = s as unknown as { + session: { clientConnection: { prompt: typeof prompt } } | null; + }; + const acceptedSession = serverInternals.session; + if (!acceptedSession) throw new Error("expected active test session"); + acceptedSession.clientConnection.prompt = prompt; + + const token = createToken(); + const send = async (requestId: string) => + fetch(`http://localhost:${port}/command`, { + method: "POST", + headers: { + Authorization: `Bearer ${token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + jsonrpc: "2.0", + id: requestId, + method: "user_message", + params: { content: "do the thing", messageId: "m-teardown" }, + }), + }); + + const firstResponse = send("first-attempt"); + await vi.waitFor(() => expect(prompt).toHaveBeenCalledTimes(1)); + + serverInternals.session = null; + finishPrompt({ stopReason: "end_turn" }); + const first = await firstResponse; + await expect(first.json()).resolves.toMatchObject({ + result: { stopReason: "end_turn" }, + }); + + serverInternals.session = acceptedSession; + const retry = await send("retry"); + await expect(retry.json()).resolves.toMatchObject({ + result: { stopReason: "duplicate_delivery", duplicate: true }, + }); + expect(prompt).toHaveBeenCalledTimes(1); + }, 20000); + // Shared plumbing for the relay-echo tests: install a controllable // prompt, stub the log writer so relayAgentResponse has an answer to // relay, and spy on the relay_message client call. @@ -2447,6 +2704,7 @@ describe("AgentServer HTTP Mode", () => { expect(runStarted?.notification?.params).toMatchObject({ runId: "test-run-id", taskId: "test-task-id", + steering: "native", }); // Agent reports its semver so clients can gate UI features // against agent capabilities (e.g. `>=0.40.1`). The exact value diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 208c643c4b..c25336e948 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -298,6 +298,15 @@ function isManualCompactPrompt(prompt: ContentBlock[]): boolean { return /^\/compact(?:\s|$)/.test(promptBlocksToText(prompt).trimStart()); } +function extractSteeringCapability(result: unknown): string | undefined { + const steering = ( + result as { + agentCapabilities?: { _meta?: { posthog?: { steering?: unknown } } }; + } + )?.agentCapabilities?._meta?.posthog?.steering; + return typeof steering === "string" ? steering : undefined; +} + interface LocalSkillPromptContext { /** Set when the message is a bare `/skill` invocation the adapter should strip. */ skillName?: string; @@ -377,6 +386,8 @@ export class AgentServer { private preSessionEvents: Record[] = []; private deliveredMessageIds = new Set(); private pendingCompactContinuationMessageIds = new Set(); + private inFlightMessageDeliveries = new Map>(); + private activeOwnedTurnCount = 0; private pendingPermissions = new Map< string, { @@ -937,49 +948,51 @@ export class AgentServer { switch (method) { case POSTHOG_NOTIFICATIONS.USER_MESSAGE: case "user_message": { - this.logger.debug("Received user_message command", { - hasContent: - typeof params.content === "string" - ? params.content.trim().length > 0 - : Array.isArray(params.content) && params.content.length > 0, - artifactCount: Array.isArray(params.artifacts) - ? params.artifacts.length - : 0, - }); - const builtPrompt = await this.buildPromptFromContentAndArtifacts({ - content: params.content as string | ContentBlock[] | undefined, - artifacts: Array.isArray(params.artifacts) - ? (params.artifacts as TaskRunArtifact[]) - : [], - taskId: this.session.payload.task_id, - runId: this.session.payload.run_id, - }); - const prompt = builtPrompt.prompt; - if (prompt.length === 0) { - throw new Error("User message cannot be empty"); - } - + const commandSession = this.session; const messageId = typeof params.messageId === "string" && params.messageId ? params.messageId : undefined; + const inFlightDelivery = messageId + ? this.inFlightMessageDeliveries.get(messageId) + : undefined; + if (inFlightDelivery) { + this.logger.info("Awaiting in-flight user_message delivery", { + messageId, + }); + return await inFlightDelivery; + } + let retryCompactContinuation = false; - if (messageId) { - if (this.deliveredMessageIds.has(messageId)) { - if (this.pendingCompactContinuationMessageIds.has(messageId)) { - retryCompactContinuation = true; - this.logger.info("Retrying pending compact continuation", { - messageId, - }); - } else { - this.logger.info("Duplicate user_message delivery ignored", { - messageId, - }); - return { stopReason: "duplicate_delivery", duplicate: true }; - } + if (messageId && this.deliveredMessageIds.has(messageId)) { + if (this.pendingCompactContinuationMessageIds.has(messageId)) { + retryCompactContinuation = true; + this.logger.info("Retrying pending compact continuation", { + messageId, + }); } else { - this.deliveredMessageIds.add(messageId); + this.logger.info("Duplicate user_message delivery ignored", { + messageId, + }); + return { stopReason: "duplicate_delivery", duplicate: true }; } + } + + let resolveDelivery: (result: unknown) => void = () => {}; + let rejectDelivery: (error: unknown) => void = () => {}; + const deliveryOutcome = new Promise((resolve, reject) => { + resolveDelivery = resolve; + rejectDelivery = reject; + }); + void deliveryOutcome.catch(() => {}); + if (messageId) { + this.inFlightMessageDeliveries.set(messageId, deliveryOutcome); + } + let deliveryCommitted = retryCompactContinuation; + const commitDelivery = (): void => { + deliveryCommitted = true; + if (!messageId) return; + this.deliveredMessageIds.add(messageId); if (this.deliveredMessageIds.size > 500) { const oldest = this.deliveredMessageIds.values().next().value; if (oldest !== undefined) { @@ -987,137 +1000,213 @@ export class AgentServer { this.pendingCompactContinuationMessageIds.delete(oldest); } } - } - this.logger.debug("Built user_message prompt", { - blockTypes: prompt.map((block) => block.type), - }); - const promptPreview = promptBlocksToText(prompt); + }; - this.logger.debug( - `Processing user message (detectedPrUrl=${this.detectedPrUrl ?? "none"}): ${promptPreview.substring(0, 100)}...`, - ); + try { + this.logger.debug("Received user_message command", { + hasContent: + typeof params.content === "string" + ? params.content.trim().length > 0 + : Array.isArray(params.content) && params.content.length > 0, + artifactCount: Array.isArray(params.artifacts) + ? params.artifacts.length + : 0, + }); + const builtPrompt = await this.buildPromptFromContentAndArtifacts({ + content: params.content as string | ContentBlock[] | undefined, + artifacts: Array.isArray(params.artifacts) + ? (params.artifacts as TaskRunArtifact[]) + : [], + taskId: commandSession.payload.task_id, + runId: commandSession.payload.run_id, + }); + const prompt = builtPrompt.prompt; + if (prompt.length === 0) { + throw new Error("User message cannot be empty"); + } - this.session.logWriter.resetTurnMessages(this.session.payload.run_id); + this.logger.debug("Built user_message prompt", { + blockTypes: prompt.map((block) => block.type), + }); + const promptPreview = promptBlocksToText(prompt); - // Resolve before buildDetectedPrContext so a warm auto-publish upgrade - // also flips the detected-PR context to its push variant. - const autoPublishUpgrade = await this.resolveWarmAutoPublishUpgrade(); - const hostContext = [ - ...(autoPublishUpgrade ? [autoPublishUpgrade] : []), - ...(this.detectedPrUrl - ? [this.buildDetectedPrContext(this.detectedPrUrl)] - : []), - ]; - const promptMeta: Record = { - ...(builtPrompt.meta ?? {}), - ...(hostContext.length > 0 - ? { prContext: hostContext.join("\n\n") } - : {}), - }; + this.logger.debug( + `Processing user message (detectedPrUrl=${this.detectedPrUrl ?? "none"}): ${promptPreview.substring(0, 100)}...`, + ); - const manualCompactPrompt = isManualCompactPrompt(prompt); - const acpSessionId = this.session.acpSessionId; - const continueAfterCompaction = (): Promise => - this.promptWithUpstreamRetry({ - sessionId: acpSessionId, - prompt: [ - hiddenTextBlock( - "Compaction is complete. Continue working on the task from the compacted context, following the user's instructions from the /compact command.", - ), - ], - }); + // Resolve before buildDetectedPrContext so a warm auto-publish upgrade + // also flips the detected-PR context to its push variant. + const autoPublishUpgrade = await this.resolveWarmAutoPublishUpgrade(); + const hostContext = [ + ...(autoPublishUpgrade ? [autoPublishUpgrade] : []), + ...(this.detectedPrUrl + ? [this.buildDetectedPrContext(this.detectedPrUrl)] + : []), + ]; + const promptMeta: Record = { + ...(builtPrompt.meta ?? {}), + ...(hostContext.length > 0 + ? { prContext: hostContext.join("\n\n") } + : {}), + }; - let compactCommandCompleted = retryCompactContinuation; - let result: PromptResponse; - this.suppressAdapterTurnComplete = - manualCompactPrompt || retryCompactContinuation; - try { - if (retryCompactContinuation) { - result = await continueAfterCompaction(); - if (messageId) { - this.pendingCompactContinuationMessageIds.delete(messageId); + if (params.steer === true) { + if (this.activeOwnedTurnCount > 0) { + const result = await commandSession.clientConnection.prompt({ + sessionId: commandSession.acpSessionId, + prompt, + _meta: { ...promptMeta, steer: true }, + }); + const accepted = + (result._meta as { steer?: unknown } | undefined)?.steer === + true; + if (accepted) { + commitDelivery(); + const outcome = { stopReason: "steered", steered: true }; + resolveDelivery(outcome); + return outcome; + } } - } else { - result = await this.session.clientConnection.prompt({ - sessionId: this.session.acpSessionId, - prompt, - ...(Object.keys(promptMeta).length > 0 - ? { _meta: promptMeta } - : {}), + const outcome = { + stopReason: "steer_declined", + steered: false, + }; + resolveDelivery(outcome); + return outcome; + } + + commandSession.logWriter.resetTurnMessages( + commandSession.payload.run_id, + ); + + const manualCompactPrompt = isManualCompactPrompt(prompt); + const acpSessionId = commandSession.acpSessionId; + const continueAfterCompaction = (): Promise => + this.promptWithUpstreamRetry({ + sessionId: acpSessionId, + prompt: [ + hiddenTextBlock( + "Compaction is complete. Continue working on the task from the compacted context, following the user's instructions from the /compact command.", + ), + ], }); - if (result.stopReason === "end_turn" && manualCompactPrompt) { - compactCommandCompleted = true; - if (messageId) { - this.pendingCompactContinuationMessageIds.add(messageId); - } - // `/compact` is an SDK-local command, so without a follow-up the - // cloud run reports completion before the model resumes the task. - this.recordTurnUsage(result.usage); - result = await continueAfterCompaction(); + let result: PromptResponse; + this.suppressAdapterTurnComplete = + manualCompactPrompt || retryCompactContinuation; + try { + if (retryCompactContinuation) { + result = await this.runOwnedTurn(continueAfterCompaction); if (messageId) { this.pendingCompactContinuationMessageIds.delete(messageId); } + } else { + result = await this.runOwnedTurn(() => { + const promptResult = commandSession.clientConnection.prompt({ + sessionId: commandSession.acpSessionId, + prompt, + ...(Object.keys(promptMeta).length > 0 + ? { _meta: promptMeta } + : {}), + }); + if (!promptResult) { + throw new Error("Agent connection did not accept the prompt"); + } + return promptResult; + }); + + if (result.stopReason === "end_turn" && manualCompactPrompt) { + commitDelivery(); + if (messageId) { + this.pendingCompactContinuationMessageIds.add(messageId); + } + // `/compact` is an SDK-local command, so without a follow-up the + // cloud run reports completion before the model resumes the task. + this.recordTurnUsage(result.usage); + result = await this.runOwnedTurn(continueAfterCompaction); + if (messageId) { + this.pendingCompactContinuationMessageIds.delete(messageId); + } + } } + } catch (error) { + await commandSession.logWriter.flushAll(); + const { recoverable } = await this.handleTurnFailure( + commandSession.payload, + "followup", + error, + ); + if (!recoverable) { + throw error; + } + commitDelivery(); + const outcome = { stopReason: "error_recoverable" }; + resolveDelivery(outcome); + return outcome; + } finally { + this.suppressAdapterTurnComplete = false; } - } catch (error) { - if (messageId && !compactCommandCompleted) { - this.deliveredMessageIds.delete(messageId); - } - await this.session.logWriter.flushAll(); - const { recoverable } = await this.handleTurnFailure( - this.session.payload, - "followup", - error, - ); - if (!recoverable) { - throw error; - } - return { stopReason: "error_recoverable" }; - } finally { - this.suppressAdapterTurnComplete = false; - } + commitDelivery(); - this.logger.debug("User message completed", { - stopReason: result.stopReason, - }); + this.logger.debug("User message completed", { + stopReason: result.stopReason, + }); - if (result.stopReason === "end_turn") { - void this.syncCloudBranchMetadata(this.session.payload); - } + if (result.stopReason === "end_turn") { + void this.syncCloudBranchMetadata(commandSession.payload); + } - this.recordTurnUsage(result.usage); - this.broadcastTurnComplete(result.stopReason); + this.recordTurnUsage(result.usage); + this.broadcastTurnComplete(result.stopReason); - if (result.stopReason === "end_turn") { - // Relay the response to Slack. For follow-ups this is the primary - // delivery path — the HTTP caller only handles reactions. Echo the - // initiating message's id so the backend can attribute the answer. - this.relayAgentResponse(this.session.payload, messageId).catch( - (err) => - this.logger.debug("Failed to relay follow-up response", err), - ); - } + if (result.stopReason === "end_turn") { + // Relay the response to Slack. For follow-ups this is the primary + // delivery path — the HTTP caller only handles reactions. Echo the + // initiating message's id so the backend can attribute the answer. + this.relayAgentResponse(commandSession.payload, messageId).catch( + (err) => + this.logger.debug("Failed to relay follow-up response", err), + ); + } - // Flush logs and include the assistant's response text so callers - // (e.g. Slack follow-up forwarding) can extract it without racing - // against async log persistence to object storage. - let assistantMessage: string | undefined; - try { - await this.session.logWriter.flush(this.session.payload.run_id, { - coalesce: true, - }); - assistantMessage = this.session.logWriter.getFullAgentResponse( - this.session.payload.run_id, - ); - } catch { - this.logger.debug("Failed to extract assistant message from logs"); - } + // Flush logs and include the assistant's response text so callers + // (e.g. Slack follow-up forwarding) can extract it without racing + // against async log persistence to object storage. + let assistantMessage: string | undefined; + try { + await commandSession.logWriter.flush( + commandSession.payload.run_id, + { + coalesce: true, + }, + ); + assistantMessage = commandSession.logWriter.getFullAgentResponse( + commandSession.payload.run_id, + ); + } catch { + this.logger.debug("Failed to extract assistant message from logs"); + } - return { - stopReason: result.stopReason, - ...(assistantMessage && { assistant_message: assistantMessage }), - }; + const outcome = { + stopReason: result.stopReason, + ...(assistantMessage && { assistant_message: assistantMessage }), + }; + resolveDelivery(outcome); + return outcome; + } catch (error) { + if (messageId && !deliveryCommitted) { + this.deliveredMessageIds.delete(messageId); + } + rejectDelivery(error); + throw error; + } finally { + if ( + messageId && + this.inFlightMessageDeliveries.get(messageId) === deliveryOutcome + ) { + this.inFlightMessageDeliveries.delete(messageId); + } + } } case POSTHOG_NOTIFICATIONS.CANCEL: @@ -1467,10 +1556,11 @@ export class AgentServer { clientStream, ); - await clientConnection.initialize({ + const initializeResult = await clientConnection.initialize({ protocolVersion: PROTOCOL_VERSION, clientCapabilities: {}, }); + const steering = extractSteeringCapability(initializeResult); const runState = preTaskRun?.state as Record | undefined; // Preserve native Codex modes for cloud runs so they behave the same as @@ -1624,6 +1714,7 @@ export class AgentServer { runId: payload.run_id, taskId: payload.task_id, agentVersion: this.config.version ?? packageJson.version, + ...(steering ? { steering } : {}), }, }; this.broadcastEvent({ @@ -1686,6 +1777,15 @@ export class AgentServer { return { classification: classifyAgentError(message), message }; } + private async runOwnedTurn(operation: () => Promise): Promise { + this.activeOwnedTurnCount += 1; + try { + return await operation(); + } finally { + this.activeOwnedTurnCount -= 1; + } + } + /** * Send an initial/resume turn prompt, absorbing transient upstream * failures with a bounded number of retries. These turns run unattended @@ -1900,12 +2000,18 @@ export class AgentServer { }); this.session.logWriter.resetTurnMessages(payload.run_id); + const acpSessionId = this.session.acpSessionId; + if (!acpSessionId) { + throw new Error("Agent session is missing its ACP session ID"); + } - const result = await this.promptWithUpstreamRetry({ - sessionId: this.session.acpSessionId, - prompt: initialPrompt, - ...(initialPromptMeta ? { _meta: initialPromptMeta } : {}), - }); + const result = await this.runOwnedTurn(() => + this.promptWithUpstreamRetry({ + sessionId: acpSessionId, + prompt: initialPrompt, + ...(initialPromptMeta ? { _meta: initialPromptMeta } : {}), + }), + ); this.logger.debug("Initial task message completed", { stopReason: result.stopReason, @@ -2112,12 +2218,18 @@ export class AgentServer { const builtPrompt = await buildPrompt(); this.session.logWriter.resetTurnMessages(payload.run_id); + const acpSessionId = this.session.acpSessionId; + if (!acpSessionId) { + throw new Error("Agent session is missing its ACP session ID"); + } - const result = await this.promptWithUpstreamRetry({ - sessionId: this.session.acpSessionId, - prompt: builtPrompt.prompt, - ...(builtPrompt.meta ? { _meta: builtPrompt.meta } : {}), - }); + const result = await this.runOwnedTurn(() => + this.promptWithUpstreamRetry({ + sessionId: acpSessionId, + prompt: builtPrompt.prompt, + ...(builtPrompt.meta ? { _meta: builtPrompt.meta } : {}), + }), + ); this.logger.debug(`${logLabel} completed`, { stopReason: result.stopReason, diff --git a/packages/agent/src/server/schemas.ts b/packages/agent/src/server/schemas.ts index fd3f2789fe..faf79588fd 100644 --- a/packages/agent/src/server/schemas.ts +++ b/packages/agent/src/server/schemas.ts @@ -65,6 +65,8 @@ export const userMessageParamsSchema = z ]) .optional(), artifacts: z.array(z.record(z.string(), z.unknown())).optional(), + messageId: z.string().min(1).optional(), + steer: z.boolean().optional(), }) .refine( (params) => { diff --git a/packages/api-client/src/posthog-client.test.ts b/packages/api-client/src/posthog-client.test.ts index 3e6509ac08..472a8eeb74 100644 --- a/packages/api-client/src/posthog-client.test.ts +++ b/packages/api-client/src/posthog-client.test.ts @@ -1443,7 +1443,7 @@ describe("PostHogAPIClient", () => { }); }); - it("returns the entries collected so far when a later page fails", async () => { + it("marks entries collected before a failed page as incomplete", async () => { const fetch = vi .fn() .mockResolvedValueOnce(page(makeEntries(50, "a"), true)) @@ -1455,11 +1455,14 @@ describe("PostHogAPIClient", () => { }); const client = makeClient(fetch); - const result = await client.getTaskRunSessionLogs("task-1", "run-1", { - limit: 100000, - }); + const result = await client.getTaskRunSessionLogsResult( + "task-1", + "run-1", + { limit: 100000 }, + ); - expect(result).toHaveLength(50); + expect(result).toEqual({ entries: expect.any(Array), complete: false }); + expect(result.entries).toHaveLength(50); expect(fetch).toHaveBeenCalledTimes(2); }); diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index 5478769308..29c687a4fb 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -131,6 +131,11 @@ export const CLOUD_USAGE_LIMIT_ERROR_MESSAGE = "Cloud usage limit reached"; export const SESSION_LOGS_MAX_PAGE_SIZE = 5000; +export interface TaskRunSessionLogsResult { + entries: StoredLogEntry[]; + complete: boolean; +} + /** Thrown when the backend rejects a cloud run with a 429 usage-limit error. */ export class CloudUsageLimitError extends Error { limitType: UsageLimitType; @@ -3157,6 +3162,15 @@ export class PostHogAPIClient { runId: string, options?: { limit?: number; after?: string }, ): Promise { + return (await this.getTaskRunSessionLogsResult(taskId, runId, options)) + .entries; + } + + async getTaskRunSessionLogsResult( + taskId: string, + runId: string, + options?: { limit?: number; after?: string }, + ): Promise { const maxEntries = options?.limit ?? SESSION_LOGS_MAX_PAGE_SIZE; const entries: StoredLogEntry[] = []; try { @@ -3187,21 +3201,21 @@ export class PostHogAPIClient { log.warn( `Failed to fetch session logs page at offset ${offset}: ${response.status} ${response.statusText}`, ); - break; + return { entries, complete: false }; } const page = (await response.json()) as StoredLogEntry[]; entries.push(...page); const hasMore = response.headers.get("X-Has-More") === "true"; if (!hasMore || page.length === 0) { - break; + return { entries, complete: true }; } offset += page.length; } - return entries; + return { entries, complete: false }; } catch (err) { log.warn("Failed to fetch task run session logs", err); - return entries; + return { entries, complete: false }; } } diff --git a/packages/shared/src/sessions.test.ts b/packages/shared/src/sessions.test.ts index 85d736fbc3..8706578afe 100644 --- a/packages/shared/src/sessions.test.ts +++ b/packages/shared/src/sessions.test.ts @@ -95,15 +95,20 @@ describe("sessionSupportsNativeSteer", () => { { isCloud: false, steering: "interrupt-resend", adapter: "claude" }, false, ], - // Cloud runs queue/resend; they never steer locally regardless of capability. + // Cloud runs steer only when the sandbox explicitly advertises support. [ "cloud claude native", { isCloud: true, steering: "native", adapter: "claude" }, - false, + true, ], [ "cloud codex native", { isCloud: true, steering: "native", adapter: "codex" }, + true, + ], + [ + "cloud without capability", + { isCloud: true, steering: undefined, adapter: "claude" }, false, ], ])("%s", (_label, session, expected) => { diff --git a/packages/shared/src/sessions.ts b/packages/shared/src/sessions.ts index 493bb5b5ac..fd4fca99c3 100644 --- a/packages/shared/src/sessions.ts +++ b/packages/shared/src/sessions.ts @@ -62,6 +62,9 @@ export interface AgentSession { promptStartedAt: number | null; currentPromptId?: number | null; logUrl?: string; + /** Full cloud transcript entry count across the resume chain. */ + cloudTranscriptEntryCount?: number; + /** Leaf-run cursor used to reconcile live cloud log updates. */ processedLineCount?: number; framework?: "claude"; adapter?: Adapter; @@ -222,7 +225,7 @@ export function resolveBypassRevertMode( * Whether a mid-turn message can be folded into the running turn (steered) * rather than interrupt-and-resent. Decided by the adapter's negotiated * `steering` capability: "native" folds (claude, codex app-server); - * "interrupt-resend" (legacy) does not. Cloud runs never steer locally. + * "interrupt-resend" (legacy) does not. * * Fallback: if `steering` is unset (a start path that predates capability * plumbing), Claude is still treated as native — it has always steered — so the @@ -231,7 +234,7 @@ export function resolveBypassRevertMode( export function sessionSupportsNativeSteer( session: Pick, ): boolean { - if (session.isCloud) return false; if (session.steering === "native") return true; + if (session.isCloud) return false; return session.steering == null && session.adapter === "claude"; } From 4b049ccf27e34a598086d1dd7bc230d33465a352 Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Fri, 17 Jul 2026 13:40:22 +0100 Subject: [PATCH 3/5] fix(cloud): reconcile resumed events by leaf identity --- packages/core/src/sessions/sessionEvents.ts | 88 +- packages/core/src/sessions/sessionService.ts | 875 ++++++++++++++++++ .../sessions/sessionServiceHydration.test.ts | 197 ++++ 3 files changed, 1144 insertions(+), 16 deletions(-) create mode 100644 packages/core/src/sessions/sessionServiceHydration.test.ts diff --git a/packages/core/src/sessions/sessionEvents.ts b/packages/core/src/sessions/sessionEvents.ts index 7370397428..f3f09f6cd8 100644 --- a/packages/core/src/sessions/sessionEvents.ts +++ b/packages/core/src/sessions/sessionEvents.ts @@ -23,20 +23,60 @@ import { skillTagsToSlashCommands } from "../message-editor/skillTags"; import { isNotification, POSTHOG_NOTIFICATIONS } from "./acpNotifications"; import { extractPromptDisplayContent } from "./promptContent"; +export interface StoredLogEventPosition { + taskRunId: string; + entryIndex: number; +} + +export interface StoredLogEventPositionOptions { + taskRunId: string; + startEntryIndex: number; + firstPositionedEntryIndex?: number; +} + +// Ordinals are local reconciliation provenance, so keep them out of the ACP +// event shape that crosses host and renderer boundaries. +const storedLogEventPositions = new WeakMap< + AcpMessage, + StoredLogEventPosition +>(); + +export function getStoredLogEventPosition( + event: AcpMessage, +): StoredLogEventPosition | undefined { + return storedLogEventPositions.get(event); +} + +function recordStoredLogEventPosition( + event: AcpMessage, + position: StoredLogEventPosition | undefined, +): AcpMessage { + if (position) storedLogEventPositions.set(event, position); + return event; +} + /** * Convert a stored log entry to an ACP message. */ -function storedEntryToAcpMessage(entry: StoredLogEntry): AcpMessage { +function storedEntryToAcpMessage( + entry: StoredLogEntry, + position?: StoredLogEventPosition, +): AcpMessage { const ts = entry.timestamp ? new Date(entry.timestamp).getTime() : Date.now(); const promoted = promoteImportedUserPrompt(entry, ts); // Freeze at creation: events assigned via setSession bypass the store's // per-append freeze, so this keeps them read-only once stored. - if (promoted) return Object.freeze(promoted); - return Object.freeze({ - type: "acp_message", - ts, - message: (entry.notification ?? {}) as JsonRpcMessage, - }); + if (promoted) { + return recordStoredLogEventPosition(Object.freeze(promoted), position); + } + return recordStoredLogEventPosition( + Object.freeze({ + type: "acp_message", + ts, + message: (entry.notification ?? {}) as JsonRpcMessage, + }), + position, + ); } /** @@ -192,13 +232,16 @@ function withToolCallUpdate( update: Record, ): AcpMessage { const msg = event.message as { params?: SessionNotification }; - return Object.freeze({ - ...event, - message: { - ...msg, - params: { ...msg.params, update }, - } as JsonRpcMessage, - }); + return recordStoredLogEventPosition( + Object.freeze({ + ...event, + message: { + ...msg, + params: { ...msg.params, update }, + } as JsonRpcMessage, + }), + getStoredLogEventPosition(event), + ); } /** @@ -262,6 +305,7 @@ export function collapseSupersededToolCallUpdates( export function convertStoredEntriesToEvents( entries: StoredLogEntry[], taskDescription?: string, + positionOptions?: StoredLogEventPositionOptions, ): AcpMessage[] { const events: AcpMessage[] = []; @@ -272,8 +316,20 @@ export function convertStoredEntriesToEvents( events.push(createUserMessageEvent(taskDescription, startTs)); } - for (const entry of entries) { - events.push(storedEntryToAcpMessage(entry)); + const firstPositionedEntryIndex = + positionOptions?.firstPositionedEntryIndex ?? 0; + for (let entryIndex = 0; entryIndex < entries.length; entryIndex += 1) { + const position = + positionOptions && entryIndex >= firstPositionedEntryIndex + ? { + taskRunId: positionOptions.taskRunId, + entryIndex: + positionOptions.startEntryIndex + + entryIndex - + firstPositionedEntryIndex, + } + : undefined; + events.push(storedEntryToAcpMessage(entries[entryIndex], position)); } return collapseSupersededToolCallUpdates(events); diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index 3846d2164f..a3ce8c66c0 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -84,6 +84,7 @@ import { convertStoredEntriesToEvents, createUserShellExecuteEvent, extractPromptText, + getStoredLogEventPosition, getUserShellExecutesSinceLastPrompt, hasSessionPromptEvent, isTurnCompleteEvent, @@ -528,6 +529,880 @@ function entriesScopedToTaskRun( }); } +function appendHydrationHash(hash: number, value: string): number { + let nextHash = hash; + for (let index = 0; index < value.length; index += 1) { + nextHash ^= value.charCodeAt(index); + nextHash = Math.imul(nextHash, 16_777_619); + } + return nextHash >>> 0; +} + +function hashHydrationValue(value: unknown, hash = 2_166_136_261): number { + if (value === null) return appendHydrationHash(hash, "null"); + if (Array.isArray(value)) { + let nextHash = appendHydrationHash(hash, "["); + for (const item of value) { + nextHash = hashHydrationValue(item, nextHash); + nextHash = appendHydrationHash(nextHash, ","); + } + return appendHydrationHash(nextHash, "]"); + } + switch (typeof value) { + case "boolean": + return appendHydrationHash(hash, value ? "true" : "false"); + case "number": + return appendHydrationHash(hash, `number:${value}`); + case "string": + return appendHydrationHash(hash, `string:${value}`); + case "undefined": + return appendHydrationHash(hash, "undefined"); + case "object": { + let nextHash = appendHydrationHash(hash, "{"); + const record = value as Record; + for (const key of Object.keys(record).sort()) { + nextHash = appendHydrationHash(nextHash, key); + nextHash = hashHydrationValue(record[key], nextHash); + } + return appendHydrationHash(nextHash, "}"); + } + default: + return appendHydrationHash(hash, typeof value); + } +} + +function hydrationValuesEqual(left: unknown, right: unknown): boolean { + if (left === right) return true; + if (left === null || right === null || typeof left !== typeof right) { + return false; + } + if (Array.isArray(left) || Array.isArray(right)) { + if (!Array.isArray(left) || !Array.isArray(right)) return false; + return ( + left.length === right.length && + left.every((value, index) => hydrationValuesEqual(value, right[index])) + ); + } + if (typeof left !== "object" || typeof right !== "object") return false; + const leftRecord = left as Record; + const rightRecord = right as Record; + const leftKeys = Object.keys(leftRecord); + const rightKeys = Object.keys(rightRecord); + return ( + leftKeys.length === rightKeys.length && + leftKeys.every( + (key) => + Object.hasOwn(rightRecord, key) && + hydrationValuesEqual(leftRecord[key], rightRecord[key]), + ) + ); +} + +function cloudHydrationMessageHash(event: AcpMessage): number { + return hashHydrationValue(event.message); +} + +function cloudHydrationMessagesEqual( + left: AcpMessage, + right: AcpMessage, +): boolean { + const leftPosition = getStoredLogEventPosition(left); + const rightPosition = getStoredLogEventPosition(right); + if (leftPosition && rightPosition) { + return ( + leftPosition.taskRunId === rightPosition.taskRunId && + leftPosition.entryIndex === rightPosition.entryIndex + ); + } + return hydrationValuesEqual(left.message, right.message); +} + +function cloudHydrationPositionsEqual( + left: AcpMessage, + right: AcpMessage, +): boolean { + const leftPosition = getStoredLogEventPosition(left); + const rightPosition = getStoredLogEventPosition(right); + return ( + leftPosition !== undefined && + rightPosition !== undefined && + leftPosition.taskRunId === rightPosition.taskRunId && + leftPosition.entryIndex === rightPosition.entryIndex + ); +} + +interface HydrationTurn { + events: AcpMessage[]; + eventHashes: number[]; + promptEvent?: AcpMessage; + promptHash?: number; + taskRunId?: string; +} + +interface HydrationPromptPositions { + all: number[]; + unscoped: number[]; + byTaskRunId: Map; +} + +function sessionEventTaskRunMarker(event: AcpMessage): string | undefined { + if (!isJsonRpcNotification(event.message)) return undefined; + const params = (event.message.params ?? {}) as { + runId?: unknown; + taskRunId?: unknown; + }; + if ( + isNotification(event.message.method, POSTHOG_NOTIFICATIONS.SDK_SESSION) && + typeof params.taskRunId === "string" + ) { + return params.taskRunId; + } + if ( + isNotification(event.message.method, POSTHOG_NOTIFICATIONS.RUN_STARTED) && + typeof params.runId === "string" + ) { + return params.runId; + } + return undefined; +} + +function splitHydrationTurns(events: AcpMessage[]): HydrationTurn[] { + const turns: HydrationTurn[] = []; + let taskRunId: string | undefined; + let currentEvents: AcpMessage[] = []; + let currentPromptEvent: AcpMessage | undefined; + const finishCurrent = (): void => { + if (currentEvents.length === 0) return; + turns.push({ + events: currentEvents, + eventHashes: currentEvents.map(cloudHydrationMessageHash), + promptEvent: currentPromptEvent, + promptHash: currentPromptEvent + ? cloudHydrationMessageHash(currentPromptEvent) + : undefined, + taskRunId, + }); + }; + + for (const event of events) { + const marker = sessionEventTaskRunMarker(event); + if (marker) { + finishCurrent(); + taskRunId = marker; + currentEvents = [event]; + currentPromptEvent = undefined; + continue; + } + + if (isSessionPromptEvent(event)) { + finishCurrent(); + currentEvents = [event]; + currentPromptEvent = event; + continue; + } + currentEvents.push(event); + } + finishCurrent(); + return turns; +} + +function hydrationTurnScopesMatch( + liveTurn: HydrationTurn, + hydratedTurn: HydrationTurn, +): boolean { + return ( + liveTurn.taskRunId === undefined || + hydratedTurn.taskRunId === undefined || + liveTurn.taskRunId === hydratedTurn.taskRunId + ); +} + +function indexHydratedPromptTurns( + hydratedTurns: HydrationTurn[], +): Map { + const positionsByPrompt = new Map(); + for (let index = 0; index < hydratedTurns.length; index += 1) { + const turn = hydratedTurns[index]; + if (turn.promptHash === undefined) continue; + let positions = positionsByPrompt.get(turn.promptHash); + if (!positions) { + positions = { + all: [], + unscoped: [], + byTaskRunId: new Map(), + }; + positionsByPrompt.set(turn.promptHash, positions); + } + positions.all.push(index); + if (turn.taskRunId === undefined) { + positions.unscoped.push(index); + continue; + } + const scopedPositions = positions.byTaskRunId.get(turn.taskRunId) ?? []; + scopedPositions.push(index); + positions.byTaskRunId.set(turn.taskRunId, scopedPositions); + } + return positionsByPrompt; +} + +function latestPositionIndexAtOrBefore( + positions: number[] | undefined, + maximum: number, +): number { + if (!positions || positions.length === 0) return -1; + let low = 0; + let high = positions.length - 1; + let match = -1; + while (low <= high) { + const middle = Math.floor((low + high) / 2); + if (positions[middle] <= maximum) { + match = middle; + low = middle + 1; + } else { + high = middle - 1; + } + } + return match; +} + +function findPromptHydrationTurn( + liveTurn: HydrationTurn, + hydratedTurns: HydrationTurn[], + positionsByPrompt: Map, + maximum: number, +): number { + const livePrompt = liveTurn.promptEvent; + if (livePrompt === undefined || liveTurn.promptHash === undefined) { + return -1; + } + const positions = positionsByPrompt.get(liveTurn.promptHash); + if (!positions) return -1; + const matchingPosition = ( + candidatePositions: number[] | undefined, + ): number => { + if (!candidatePositions) return -1; + let candidateIndex = latestPositionIndexAtOrBefore( + candidatePositions, + maximum, + ); + while (candidateIndex >= 0) { + const position = candidatePositions[candidateIndex]; + const hydratedPrompt = hydratedTurns[position].promptEvent; + if ( + hydratedPrompt && + cloudHydrationMessagesEqual(livePrompt, hydratedPrompt) + ) { + return position; + } + candidateIndex -= 1; + } + return -1; + }; + if (liveTurn.taskRunId === undefined) { + return matchingPosition(positions.all); + } + return Math.max( + matchingPosition(positions.byTaskRunId.get(liveTurn.taskRunId)), + matchingPosition(positions.unscoped), + ); +} + +interface PromptlessHydrationMatch { + hydratedTurnIndex: number; + liveMessageIndexOffset: number; +} + +interface HydrationEventOverlap { + hydratedEventIndex: number; + liveEventIndex: number; +} + +interface HydrationEventPosition { + turnIndex: number; + eventIndex: number; +} + +type HydrationEventIndex = Map; + +function indexHydratedTurnEvents( + hydratedTurns: HydrationTurn[], +): HydrationEventIndex { + const positionsByHash: HydrationEventIndex = new Map(); + for (let turnIndex = 0; turnIndex < hydratedTurns.length; turnIndex += 1) { + const turn = hydratedTurns[turnIndex]; + for (let eventIndex = 0; eventIndex < turn.events.length; eventIndex += 1) { + if (!isStrongPromptlessOverlapEvent(turn.events[eventIndex])) continue; + const hash = turn.eventHashes[eventIndex]; + const positions = positionsByHash.get(hash) ?? []; + positions.push({ turnIndex, eventIndex }); + positionsByHash.set(hash, positions); + } + } + return positionsByHash; +} + +function isStrongPromptlessOverlapEvent(event: AcpMessage): boolean { + if (!isJsonRpcNotification(event.message)) return false; + if (event.message.method !== "session/update") return false; + const update = ( + event.message.params as { update?: { sessionUpdate?: unknown } } | undefined + )?.update; + return ( + typeof update?.sessionUpdate === "string" && + agentMessageUpdateKind(event) !== "ignored" + ); +} + +function findHydrationEventOverlap( + liveTurn: HydrationTurn, + hydratedTurn: HydrationTurn, + allowWeakOverlap: boolean, +): HydrationEventOverlap | undefined { + if (!hydrationTurnScopesMatch(liveTurn, hydratedTurn)) return undefined; + const hydratedPositions = new Map(); + for ( + let hydratedIndex = 0; + hydratedIndex < hydratedTurn.events.length; + hydratedIndex += 1 + ) { + const event = hydratedTurn.events[hydratedIndex]; + if (!allowWeakOverlap && !isStrongPromptlessOverlapEvent(event)) { + continue; + } + const hash = hydratedTurn.eventHashes[hydratedIndex]; + const positions = hydratedPositions.get(hash) ?? []; + positions.push(hydratedIndex); + hydratedPositions.set(hash, positions); + } + const findMatch = ( + kind: "stable" | "boundary" | "any", + ): HydrationEventOverlap | undefined => { + for ( + let liveIndex = liveTurn.events.length - 1; + liveIndex >= 0; + liveIndex -= 1 + ) { + const liveEvent = liveTurn.events[liveIndex]; + if (!allowWeakOverlap && !isStrongPromptlessOverlapEvent(liveEvent)) { + continue; + } + if ( + kind === "boundary" && + (!isStrongPromptlessOverlapEvent(liveEvent) || + agentMessageUpdateKind(liveEvent) != null) + ) { + continue; + } + const positions = hydratedPositions.get(liveTurn.eventHashes[liveIndex]); + if (!positions) continue; + for ( + let positionIndex = positions.length - 1; + positionIndex >= 0; + positionIndex -= 1 + ) { + const hydratedEventIndex = positions[positionIndex]; + const hydratedEvent = hydratedTurn.events[hydratedEventIndex]; + if ( + kind === "stable" && + !cloudHydrationPositionsEqual(liveEvent, hydratedEvent) + ) { + continue; + } + if (cloudHydrationMessagesEqual(liveEvent, hydratedEvent)) { + return { hydratedEventIndex, liveEventIndex: liveIndex }; + } + } + } + return undefined; + }; + return findMatch("stable") ?? findMatch("boundary") ?? findMatch("any"); +} + +function agentMessageIndexBeforeEvent( + events: AcpMessage[], + eventIndex: number, +): number { + const position: AgentMessagePosition = { + messageIndex: 0, + chunkRunActive: false, + }; + for (let index = 0; index < eventIndex; index += 1) { + const updateKind = agentMessageUpdateKind(events[index]); + if (updateKind === "ignored") continue; + if (updateKind === "chunk") { + position.chunkRunActive = true; + } else if (updateKind === "final") { + position.messageIndex += 1; + position.chunkRunActive = false; + } else { + finishAgentMessageChunkRun(position); + } + } + return position.messageIndex; +} + +function indexAgentMessagePositions( + events: AcpMessage[], + startingIndex: number, +): WeakMap { + const positions = new WeakMap(); + const position: AgentMessagePosition = { + messageIndex: startingIndex, + chunkRunActive: false, + }; + for (const event of events) { + const updateKind = agentMessageUpdateKind(event); + if (updateKind === "chunk") { + positions.set(event, position.messageIndex); + position.chunkRunActive = true; + } else if (updateKind === "final") { + positions.set(event, position.messageIndex); + position.messageIndex += 1; + position.chunkRunActive = false; + } else if (updateKind !== "ignored") { + finishAgentMessageChunkRun(position); + } + } + return positions; +} + +function promptlessTailStrictlyPredatesPrompt( + liveTurn: HydrationTurn, + hydratedTurn: HydrationTurn, +): boolean { + const promptTimestamp = hydratedTurn.promptEvent?.ts; + return ( + promptTimestamp === undefined || + liveTurn.events.every((event) => event.ts < promptTimestamp) + ); +} + +function hasLaterUnmatchedAssistantBoundary( + liveTurn: HydrationTurn, + hydratedTurn: HydrationTurn, + overlap: HydrationEventOverlap, +): boolean { + if (agentMessageUpdateKind(liveTurn.events[overlap.liveEventIndex]) == null) { + return false; + } + return hydratedTurn.events + .slice(overlap.hydratedEventIndex + 1) + .some( + (event) => + isStrongPromptlessOverlapEvent(event) && + agentMessageUpdateKind(event) == null, + ); +} + +interface IndexedHydrationEventOverlap extends HydrationEventOverlap { + hydratedTurnIndex: number; +} + +function latestEventPositionAtOrBeforeTurn( + positions: HydrationEventPosition[], + maximumTurnIndex: number, +): number { + let low = 0; + let high = positions.length - 1; + let match = -1; + while (low <= high) { + const middle = Math.floor((low + high) / 2); + if (positions[middle].turnIndex <= maximumTurnIndex) { + match = middle; + low = middle + 1; + } else { + high = middle - 1; + } + } + return match; +} + +function findIndexedPromptlessOverlap( + liveTurn: HydrationTurn, + hydratedTurns: HydrationTurn[], + hydratedEventIndex: HydrationEventIndex, + minimumTurnIndex: number, + maximumTurnIndex: number, +): IndexedHydrationEventOverlap | undefined { + let latestMatch: IndexedHydrationEventOverlap | undefined; + for ( + let liveEventIndex = liveTurn.events.length - 1; + liveEventIndex >= 0; + liveEventIndex -= 1 + ) { + const liveEvent = liveTurn.events[liveEventIndex]; + if (!isStrongPromptlessOverlapEvent(liveEvent)) continue; + const positions = hydratedEventIndex.get( + liveTurn.eventHashes[liveEventIndex], + ); + if (!positions) continue; + for ( + let positionIndex = latestEventPositionAtOrBeforeTurn( + positions, + maximumTurnIndex, + ); + positionIndex >= 0; + positionIndex -= 1 + ) { + const position = positions[positionIndex]; + if (position.turnIndex < minimumTurnIndex) break; + if (latestMatch && position.turnIndex < latestMatch.hydratedTurnIndex) { + break; + } + const hydratedTurn = hydratedTurns[position.turnIndex]; + const hydratedEvent = hydratedTurn.events[position.eventIndex]; + if ( + !isStrongPromptlessOverlapEvent(hydratedEvent) || + !hydrationTurnScopesMatch(liveTurn, hydratedTurn) || + !cloudHydrationMessagesEqual(liveEvent, hydratedEvent) + ) { + continue; + } + latestMatch = { + hydratedTurnIndex: position.turnIndex, + hydratedEventIndex: position.eventIndex, + liveEventIndex, + }; + break; + } + } + return latestMatch; +} + +function findPromptlessHydrationTurn( + liveTurn: HydrationTurn, + hydratedTurns: HydrationTurn[], + hydratedEventIndex: HydrationEventIndex, + maximum: number, +): PromptlessHydrationMatch | undefined { + if (maximum < 0) return undefined; + const leafTaskRunId = hydratedTurns[maximum]?.taskRunId; + let minimum = 0; + if (leafTaskRunId !== undefined) { + minimum = maximum; + while ( + minimum > 0 && + hydratedTurns[minimum - 1].taskRunId === leafTaskRunId + ) { + minimum -= 1; + } + } + const newestHydratedTurn = hydratedTurns[maximum]; + const newestOverlap = findHydrationEventOverlap( + liveTurn, + newestHydratedTurn, + true, + ); + if (newestOverlap) { + if ( + !cloudHydrationPositionsEqual( + liveTurn.events[newestOverlap.liveEventIndex], + newestHydratedTurn.events[newestOverlap.hydratedEventIndex], + ) && + hasLaterUnmatchedAssistantBoundary( + liveTurn, + newestHydratedTurn, + newestOverlap, + ) + ) { + return undefined; + } + return { + hydratedTurnIndex: maximum, + liveMessageIndexOffset: + agentMessageIndexBeforeEvent( + newestHydratedTurn.events, + newestOverlap.hydratedEventIndex, + ) - + agentMessageIndexBeforeEvent( + liveTurn.events, + newestOverlap.liveEventIndex, + ), + }; + } + if (maximum === minimum) { + return undefined; + } + const olderOverlap = findIndexedPromptlessOverlap( + liveTurn, + hydratedTurns, + hydratedEventIndex, + minimum, + maximum - 1, + ); + if (!olderOverlap) return undefined; + const hydratedTurn = hydratedTurns[olderOverlap.hydratedTurnIndex]; + const hasStableOverlap = cloudHydrationPositionsEqual( + liveTurn.events[olderOverlap.liveEventIndex], + hydratedTurn.events[olderOverlap.hydratedEventIndex], + ); + if ( + !hasStableOverlap && + !promptlessTailStrictlyPredatesPrompt(liveTurn, newestHydratedTurn) + ) { + return undefined; + } + return { + hydratedTurnIndex: olderOverlap.hydratedTurnIndex, + liveMessageIndexOffset: + agentMessageIndexBeforeEvent( + hydratedTurn.events, + olderOverlap.hydratedEventIndex, + ) - + agentMessageIndexBeforeEvent( + liveTurn.events, + olderOverlap.liveEventIndex, + ), + }; +} + +function discardExactHydratedEvents( + liveTurn: Pick, + hydratedTurn: Pick, + liveMessagePositions: WeakMap, + hydratedMessagePositions: WeakMap, +): AcpMessage[] { + const keep = new Array(liveTurn.events.length).fill(true); + const hydratedPositions = new Map(); + for (let index = 0; index < hydratedTurn.eventHashes.length; index += 1) { + const eventHash = hydratedTurn.eventHashes[index]; + const positions = hydratedPositions.get(eventHash) ?? []; + positions.push(index); + hydratedPositions.set(eventHash, positions); + } + let hydratedIndex = hydratedTurn.eventHashes.length - 1; + for ( + let liveIndex = liveTurn.eventHashes.length - 1; + liveIndex >= 0; + liveIndex -= 1 + ) { + const positions = hydratedPositions.get(liveTurn.eventHashes[liveIndex]); + if (!positions) continue; + let positionIndex = latestPositionIndexAtOrBefore(positions, hydratedIndex); + while (positionIndex >= 0) { + const matchedIndex = positions[positionIndex]; + const liveMessagePosition = liveMessagePositions.get( + liveTurn.events[liveIndex], + ); + const hydratedMessagePosition = hydratedMessagePositions.get( + hydratedTurn.events[matchedIndex], + ); + if ( + (liveMessagePosition !== undefined || + hydratedMessagePosition !== undefined) && + liveMessagePosition !== hydratedMessagePosition + ) { + positionIndex -= 1; + continue; + } + if ( + cloudHydrationMessagesEqual( + liveTurn.events[liveIndex], + hydratedTurn.events[matchedIndex], + ) + ) { + keep[liveIndex] = false; + hydratedIndex = matchedIndex - 1; + break; + } + positionIndex -= 1; + } + } + return liveTurn.events.filter((_event, index) => keep[index]); +} + +function agentMessageUpdateKind( + event: AcpMessage, +): "final" | "chunk" | "ignored" | undefined { + const message = event.message; + if (!isJsonRpcNotification(message) || message.method !== "session/update") { + return undefined; + } + const update = ( + message.params as + | { + update?: { + sessionUpdate?: string; + content?: unknown; + }; + } + | undefined + )?.update; + if (update?.sessionUpdate === "agent_message") return "final"; + if (update?.sessionUpdate === "agent_message_chunk") return "chunk"; + if (update?.sessionUpdate === "agent_thought_chunk") { + const content = update.content as + | { type?: string; text?: string; thinking?: string } + | null + | undefined; + if ( + (content?.type === "text" && !content.text) || + (content?.type === "thinking" && !content.thinking) + ) { + return "ignored"; + } + } + return undefined; +} + +interface AgentMessagePosition { + messageIndex: number; + chunkRunActive: boolean; +} + +function isSessionPromptEvent(event: AcpMessage): boolean { + return ( + isJsonRpcRequest(event.message) && event.message.method === "session/prompt" + ); +} + +function finishAgentMessageChunkRun(position: AgentMessagePosition): void { + if (!position.chunkRunActive) return; + position.messageIndex += 1; + position.chunkRunActive = false; +} + +function discardChunksSupersededByHydratedMessages( + liveTurn: HydrationTurn, + hydratedTurn: HydrationTurn, + liveMessageIndexOffset: number, +): Pick { + const hydratedMessagePositions = new Set(); + const hydratedPosition: AgentMessagePosition = { + messageIndex: 0, + chunkRunActive: false, + }; + for (const event of hydratedTurn.events) { + if (isSessionPromptEvent(event)) continue; + const updateKind = agentMessageUpdateKind(event); + if (updateKind === "ignored") continue; + if (updateKind === "chunk") { + hydratedPosition.chunkRunActive = true; + continue; + } + if (updateKind === "final") { + hydratedMessagePositions.add(hydratedPosition.messageIndex); + hydratedPosition.messageIndex += 1; + hydratedPosition.chunkRunActive = false; + continue; + } + finishAgentMessageChunkRun(hydratedPosition); + } + + // SessionLogWriter treats consecutive chunks as one assistant message. A + // direct agent_message replaces that buffered message, while a later chunk + // run after a tool/update boundary is a separate message. Match those + // turn-local message positions instead of timestamps so direct finals and + // same-millisecond events reconcile correctly. + const livePosition: AgentMessagePosition = { + messageIndex: Math.max(0, liveMessageIndexOffset), + chunkRunActive: false, + }; + let discardChunkRun = false; + const events: AcpMessage[] = []; + const eventHashes: number[] = []; + for ( + let eventIndex = 0; + eventIndex < liveTurn.events.length; + eventIndex += 1 + ) { + const event = liveTurn.events[eventIndex]; + let keep = true; + if (isSessionPromptEvent(event)) { + discardChunkRun = false; + } else { + const updateKind = agentMessageUpdateKind(event); + if (updateKind === "chunk") { + if (!livePosition.chunkRunActive) { + livePosition.chunkRunActive = true; + discardChunkRun = hydratedMessagePositions.has( + livePosition.messageIndex, + ); + } + keep = !discardChunkRun; + } else if (updateKind === "final") { + livePosition.messageIndex += 1; + livePosition.chunkRunActive = false; + discardChunkRun = false; + } else if (updateKind !== "ignored") { + finishAgentMessageChunkRun(livePosition); + discardChunkRun = false; + } + } + if (keep) { + events.push(event); + eventHashes.push(liveTurn.eventHashes[eventIndex]); + } + } + return { events, eventHashes }; +} + +export function reconcileLiveEventsWithHydratedEvents( + liveEvents: AcpMessage[], + hydratedEvents: AcpMessage[], +): AcpMessage[] { + const liveTurns = splitHydrationTurns(liveEvents); + const hydratedTurns = splitHydrationTurns(hydratedEvents); + const promptPositions = indexHydratedPromptTurns(hydratedTurns); + const hydratedEventIndex = indexHydratedTurnEvents(hydratedTurns); + const reconciledTurns = new Array(liveTurns.length); + let hydratedTurnIndex = hydratedTurns.length - 1; + + for ( + let liveTurnIndex = liveTurns.length - 1; + liveTurnIndex >= 0; + liveTurnIndex -= 1 + ) { + const liveTurn = liveTurns[liveTurnIndex]; + let liveMessageIndexOffset = 0; + let matchedHydratedTurnIndex = findPromptHydrationTurn( + liveTurn, + hydratedTurns, + promptPositions, + hydratedTurnIndex, + ); + if (liveTurn.promptEvent === undefined) { + const promptlessMatch = findPromptlessHydrationTurn( + liveTurn, + hydratedTurns, + hydratedEventIndex, + hydratedTurnIndex, + ); + if (promptlessMatch) { + matchedHydratedTurnIndex = promptlessMatch.hydratedTurnIndex; + liveMessageIndexOffset = promptlessMatch.liveMessageIndexOffset; + } + } + if (matchedHydratedTurnIndex === -1) { + reconciledTurns[liveTurnIndex] = liveTurn.events; + continue; + } + + const hydratedTurn = hydratedTurns[matchedHydratedTurnIndex]; + const liveMessagePositions = indexAgentMessagePositions( + liveTurn.events, + Math.max(0, liveMessageIndexOffset), + ); + const hydratedMessagePositions = indexAgentMessagePositions( + hydratedTurn.events, + 0, + ); + reconciledTurns[liveTurnIndex] = discardExactHydratedEvents( + discardChunksSupersededByHydratedMessages( + liveTurn, + hydratedTurn, + liveMessageIndexOffset, + ), + hydratedTurn, + liveMessagePositions, + hydratedMessagePositions, + ); + hydratedTurnIndex = matchedHydratedTurnIndex - 1; + } + + return reconciledTurns.flat(); +} + export function derivePendingPermissionRequests( entries: StoredLogEntry[], options?: { taskRunId?: string }, diff --git a/packages/core/src/sessions/sessionServiceHydration.test.ts b/packages/core/src/sessions/sessionServiceHydration.test.ts new file mode 100644 index 0000000000..27dfda38ce --- /dev/null +++ b/packages/core/src/sessions/sessionServiceHydration.test.ts @@ -0,0 +1,197 @@ +import type { AcpMessage, StoredLogEntry } from "@posthog/shared"; +import { describe, expect, it } from "vitest"; +import { convertStoredEntriesToEvents } from "./sessionEvents"; +import { reconcileLiveEventsWithHydratedEvents } from "./sessionService"; + +function prompt(id: number, text: string, ts: number): AcpMessage { + return { + type: "acp_message", + ts, + message: { + jsonrpc: "2.0", + id, + method: "session/prompt", + params: { prompt: [{ type: "text", text }] }, + }, + }; +} + +function agentMessage(text: string, ts: number): AcpMessage { + return { + type: "acp_message", + ts, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + update: { + sessionUpdate: "agent_message", + content: { type: "text", text }, + }, + }, + }, + }; +} + +function agentMessageChunk(text: string, ts: number): AcpMessage { + return { + type: "acp_message", + ts, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text }, + }, + }, + }, + }; +} + +function toolCall(toolCallId: string, ts: number): AcpMessage { + return { + type: "acp_message", + ts, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + update: { + sessionUpdate: "tool_call", + toolCallId, + title: "Run command", + kind: "execute", + status: "completed", + }, + }, + }, + }; +} + +function turnComplete(ts: number): AcpMessage { + return { + type: "acp_message", + ts, + message: { + jsonrpc: "2.0", + method: "_posthog/turn_complete", + params: { stopReason: "end_turn" }, + }, + }; +} + +function storedEntry(event: AcpMessage): StoredLogEntry { + return { + type: "notification", + timestamp: new Date(event.ts).toISOString(), + notification: event.message, + }; +} + +describe("resume hydration reconciliation", () => { + it("discards a stale completed tail from an earlier leaf turn", () => { + const firstResponse = agentMessage("first response", 20); + const firstCompletion = turnComplete(30); + const hydratedEvents = [ + prompt(1, "first request", 10), + firstResponse, + firstCompletion, + prompt(2, "second request", 40), + ]; + + expect( + reconcileLiveEventsWithHydratedEvents( + [ + { ...firstResponse, ts: 21 }, + { ...firstCompletion, ts: 31 }, + ], + hydratedEvents, + ), + ).toEqual([]); + }); + + it("preserves a new assistant response after an overlapping tool boundary", () => { + const boundary = toolCall("tool-1", 30); + const nextResponse = agentMessageChunk("after tool", 40); + const completion = turnComplete(50); + const hydratedEvents = [ + prompt(1, "run a command", 10), + agentMessage("before tool", 20), + boundary, + ]; + + expect( + reconcileLiveEventsWithHydratedEvents( + [{ ...boundary, ts: 31 }, nextResponse, completion], + hydratedEvents, + ), + ).toEqual([nextResponse, completion]); + }); + + it("preserves an identical response belonging to the unmatched current prompt", () => { + const currentResponse = agentMessage("Done", 50); + const currentCompletion = turnComplete(60); + const hydratedEvents = [ + prompt(1, "first request", 10), + agentMessage("Done", 20), + turnComplete(30), + prompt(2, "second request", 40), + ]; + + expect( + reconcileLiveEventsWithHydratedEvents( + [currentResponse, currentCompletion], + hydratedEvents, + ), + ).toEqual([currentResponse, currentCompletion]); + }); + + it("preserves a repeated response when a partial live tail omits the tool boundary", () => { + const boundary = toolCall("tool-1", 30); + const repeatedResponse = agentMessage("Done", 40); + const completion = turnComplete(50); + const hydratedEvents = [ + prompt(1, "run a command", 10), + agentMessage("Done", 20), + boundary, + ]; + + expect( + reconcileLiveEventsWithHydratedEvents( + [repeatedResponse, completion], + hydratedEvents, + ), + ).toEqual([repeatedResponse, completion]); + expect( + reconcileLiveEventsWithHydratedEvents( + [{ ...boundary, ts: 31 }, repeatedResponse, completion], + hydratedEvents, + ), + ).toEqual([repeatedResponse, completion]); + }); + + it("discards a positioned stale tail across a same-millisecond prompt boundary", () => { + const firstPrompt = prompt(1, "first request", 10); + const firstResponse = agentMessage("first response", 20); + const firstCompletion = turnComplete(40); + const secondPrompt = prompt(2, "second request", 40); + const hydratedEvents = convertStoredEntriesToEvents( + [firstPrompt, firstResponse, firstCompletion, secondPrompt].map( + storedEntry, + ), + undefined, + { taskRunId: "run-1", startEntryIndex: 0 }, + ); + const staleLiveTail = convertStoredEntriesToEvents( + [firstResponse, firstCompletion].map(storedEntry), + undefined, + { taskRunId: "run-1", startEntryIndex: 1 }, + ); + + expect( + reconcileLiveEventsWithHydratedEvents(staleLiveTail, hydratedEvents), + ).toEqual([]); + }); +}); From 4bd7ca292eba48016ffedbd4a0f4a6f69cd39a99 Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Tue, 21 Jul 2026 09:47:23 +0100 Subject: [PATCH 4/5] fix(cloud): hydrate resumed transcripts atomically (#3540) --- packages/core/src/sessions/sessionService.ts | 538 ++++++-- .../sessions/sessionServiceHost.test.ts | 1078 +++++++++++++++-- 2 files changed, 1393 insertions(+), 223 deletions(-) diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index a3ce8c66c0..fc34739647 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -152,6 +152,26 @@ type TrpcSubscription = { ) => { unsubscribe: () => void }; }; +interface CloudHydrationResult { + historyEntryCount: number; + liveStreamLineCount: number; +} + +interface CloudTaskWatcher { + runId: string; + apiHost: string; + teamId: number; + startToken: number; + resumeFromEntryCount?: number; + resumeHistoryCountOffset?: number; + resumeHydrationToken: number; + bufferResumeUpdates: boolean; + bufferedResumeUpdates: CloudTaskUpdatePayload[]; + processCloudUpdate: (update: CloudTaskUpdatePayload) => void; + subscription: { unsubscribe: () => void }; + onStatusChange?: () => void; +} + export interface SessionTrpc { agent: { start: TrpcMutation; @@ -529,6 +549,32 @@ function entriesScopedToTaskRun( }); } +function suffixPrefixOverlap(left: string[], right: string[]): number { + if (left.length === 0 || right.length === 0) return 0; + + const separator = Symbol("resume-chain-separator"); + const patternAndTail: (string | symbol)[] = [ + ...right, + separator, + ...left.slice(-right.length), + ]; + const prefixLengths = new Array(patternAndTail.length).fill(0); + for (let index = 1; index < patternAndTail.length; index += 1) { + let prefixLength = prefixLengths[index - 1]; + while ( + prefixLength > 0 && + patternAndTail[index] !== patternAndTail[prefixLength] + ) { + prefixLength = prefixLengths[prefixLength - 1]; + } + if (patternAndTail[index] === patternAndTail[prefixLength]) { + prefixLength += 1; + } + prefixLengths[index] = prefixLength; + } + return prefixLengths[prefixLengths.length - 1]; +} + function appendHydrationHash(hash: number, value: string): number { let nextHash = hash; for (let index = 0; index < value.length; index += 1) { @@ -1528,17 +1574,7 @@ export class SessionService { } >(); /** Active cloud task watchers, keyed by taskId */ - private cloudTaskWatchers = new Map< - string, - { - runId: string; - apiHost: string; - teamId: number; - startToken: number; - subscription: { unsubscribe: () => void }; - onStatusChange?: () => void; - } - >(); + private cloudTaskWatchers = new Map(); private cloudLogGapReconciler: CloudLogGapReconciler; /** Maps toolCallId → cloud requestId for routing permission responses */ private cloudPermissionRequestIds = new Map(); @@ -1554,6 +1590,10 @@ export class SessionService { { startedAtTs: number; agentTextChunks: number; agentOutputEvents: number } >(); private pendingPermissionHydratedRuns = new Set(); + private cloudHydrationPromises = new Map< + string, + Promise + >(); private idleKilledSubscription: { unsubscribe: () => void } | null = null; /** * Cached preview-config-options responses keyed by `${apiHost}::${adapter}`. @@ -4324,8 +4364,10 @@ export class SessionService { newSession.optimisticItems = ( this.getSessionByRunId(session.taskRunId)?.optimisticItems ?? [] ).filter((item) => item.type === "user_message" && item.pinToTop === false); - const resumeFromEntryCount = session.processedLineCount ?? 0; - newSession.processedLineCount = resumeFromEntryCount; + const resumeFromEntryCount = + session.cloudTranscriptEntryCount ?? session.processedLineCount ?? 0; + newSession.cloudTranscriptEntryCount = resumeFromEntryCount; + newSession.processedLineCount = 0; this.d.store.setSession(newSession); // Start the watcher immediately so we don't miss status updates. @@ -5465,6 +5507,19 @@ export class SessionService { initialReasoningEffort, ); } + if ( + typeof runState?.resume_from_run_id === "string" && + !this.pendingPermissionHydratedRuns.has(taskRunId) + ) { + void this.hydrateResumeCloudTaskSessionFromLogs( + taskId, + taskRunId, + logUrl, + taskDescription, + runStatus, + runState, + ); + } return () => {}; } @@ -5523,12 +5578,16 @@ export class SessionService { !this.pendingPermissionHydratedRuns.has(taskRunId) && (isTerminalStatus(existing.cloudStatus) || (runStatus !== undefined && isTerminalStatus(runStatus))); + const shouldHydrateResumeChain = + Boolean(runState?.resume_from_run_id) && + !this.pendingPermissionHydratedRuns.has(taskRunId); const shouldHydrateSession = !existing || existing.taskRunId !== taskRunId || shouldResetExistingSession || existing.events.length === 0 || - shouldHydratePersistedPermissions; + shouldHydratePersistedPermissions || + shouldHydrateResumeChain; if ( !existing || @@ -5589,47 +5648,97 @@ export class SessionService { initialReasoningEffort, ); - if (shouldHydrateSession) { - this.hydrateCloudTaskSessionFromLogs( - taskId, - taskRunId, - logUrl, - taskDescription, - runStatus, - runState, - ); - } + const processCloudUpdate = (update: CloudTaskUpdatePayload): void => { + if (update.kind === "logs" || update.kind === "snapshot") { + this.d.store.updateSession(taskRunId, { + cloudTranscriptEntryCount: update.totalEntryCount, + }); + } + const watcher = this.cloudTaskWatchers.get(taskId); + const resumeHistoryCountOffset = + watcher?.runId === runId ? (watcher.resumeHistoryCountOffset ?? 0) : 0; + const normalizedUpdate: CloudTaskUpdatePayload = + resumeHistoryCountOffset > 0 && + (update.kind === "logs" || update.kind === "snapshot") + ? { + ...update, + totalEntryCount: Math.max( + 0, + update.totalEntryCount - resumeHistoryCountOffset, + ), + } + : update; + this.handleCloudTaskUpdate(taskRunId, normalizedUpdate); + if ( + (update.kind === "status" || + update.kind === "snapshot" || + update.kind === "error") && + watcher?.onStatusChange + ) { + watcher.onStatusChange(); + } + }; + + const watcher: CloudTaskWatcher = { + runId, + apiHost, + teamId, + startToken, + resumeFromEntryCount, + resumeHistoryCountOffset: shouldHydrateResumeChain + ? resumeFromEntryCount + : 0, + resumeHydrationToken: 0, + bufferResumeUpdates: false, + bufferedResumeUpdates: [], + processCloudUpdate, + subscription: { unsubscribe: () => undefined }, + onStatusChange, + }; + this.cloudTaskWatchers.set(taskId, watcher); // Subscribe before starting the main-process watcher so the first replayed // SSE/log burst cannot race ahead of the renderer subscription. - const subscription = this.d.trpc.cloudTask.onUpdate.subscribe( + watcher.subscription = this.d.trpc.cloudTask.onUpdate.subscribe( { taskId, runId }, { onData: (update: CloudTaskUpdatePayload) => { - this.handleCloudTaskUpdate(taskRunId, update); - const watcher = this.cloudTaskWatchers.get(taskId); - if ( - (update.kind === "status" || - update.kind === "snapshot" || - update.kind === "error") && - watcher?.onStatusChange - ) { - watcher.onStatusChange(); + const activeWatcher = this.cloudTaskWatchers.get(taskId); + if (!activeWatcher || activeWatcher.runId !== runId) { + return; } + if (activeWatcher.bufferResumeUpdates) { + activeWatcher.bufferedResumeUpdates.push(update); + return; + } + activeWatcher.processCloudUpdate(update); }, onError: (err: unknown) => this.d.log.error("Cloud task subscription error", { taskId, err }), }, ); - this.cloudTaskWatchers.set(taskId, { - runId, - apiHost, - teamId, - startToken, - subscription, - onStatusChange, - }); + if (shouldHydrateSession) { + if (shouldHydrateResumeChain) { + void this.hydrateResumeCloudTaskSessionFromLogs( + taskId, + taskRunId, + logUrl, + taskDescription, + runStatus, + runState, + ); + } else { + void this.hydrateCloudTaskSessionFromLogs( + taskId, + taskRunId, + logUrl, + taskDescription, + runStatus, + runState, + ); + } + } // Start main-process watcher after the subscription is attached. void (async () => { @@ -5683,108 +5792,277 @@ export class SessionService { taskDescription?: string, runStatus?: TaskRunStatus, runState?: Record, + ): Promise { + const existing = this.cloudHydrationPromises.get(taskRunId); + if (existing) { + return existing; + } + const hydration = this.performCloudTaskSessionHydration( + taskId, + taskRunId, + logUrl, + taskDescription, + runStatus, + runState, + ).catch((err: unknown) => { + this.d.log.warn("Failed to hydrate cloud task session from logs", { + taskId, + taskRunId, + err, + }); + return undefined; + }); + this.cloudHydrationPromises.set(taskRunId, hydration); + void hydration.finally(() => { + if (this.cloudHydrationPromises.get(taskRunId) === hydration) { + this.cloudHydrationPromises.delete(taskRunId); + } + }); + return hydration; + } + + private async hydrateResumeCloudTaskSessionFromLogs( + taskId: string, + taskRunId: string, + logUrl?: string, + taskDescription?: string, + runStatus?: TaskRunStatus, + runState?: Record, + ): Promise { + const watcher = this.cloudTaskWatchers.get(taskId); + if (!watcher || watcher.runId !== taskRunId) return; + const hydrationToken = ++watcher.resumeHydrationToken; + watcher.bufferResumeUpdates = true; + + const result = await this.hydrateCloudTaskSessionFromLogs( + taskId, + taskRunId, + logUrl, + taskDescription, + runStatus, + runState, + ); + const activeWatcher = this.cloudTaskWatchers.get(taskId); + if ( + !activeWatcher || + activeWatcher.runId !== taskRunId || + activeWatcher.resumeHydrationToken !== hydrationToken + ) { + return; + } + + this.applyResumeHydrationOffset(taskId, taskRunId, result); + activeWatcher.bufferResumeUpdates = false; + const bufferedUpdates = activeWatcher.bufferedResumeUpdates.splice(0); + for (const update of bufferedUpdates) { + activeWatcher.processCloudUpdate(update); + } + } + + private applyResumeHydrationOffset( + taskId: string, + taskRunId: string, + result: CloudHydrationResult | undefined, ): void { - void (async () => { - let rawEntries: StoredLogEntry[]; - let totalLineCount: number; - const isResumeRun = Boolean(runState?.resume_from_run_id); - if (isTerminalStatus(runStatus) || isResumeRun) { - // Resume chains need the full history even while the leaf run is still - // active; otherwise a renderer restart hydrates only the final run. - // Non-resume in-progress runs keep using the single-run log so hydrate - // cannot race the live stream and double the active turn. - const authStatus = await this.getAuthCredentialsStatus(); - if (authStatus.kind !== "ready") { - return; - } - try { - rawEntries = await authStatus.auth.client.getTaskRunSessionLogs( + if (!result) return; + const watcher = this.cloudTaskWatchers.get(taskId); + if (!watcher || watcher.runId !== taskRunId) return; + watcher.resumeHistoryCountOffset = Math.max( + 0, + result.historyEntryCount - result.liveStreamLineCount, + ); + } + + private async performCloudTaskSessionHydration( + taskId: string, + taskRunId: string, + logUrl?: string, + taskDescription?: string, + runStatus?: TaskRunStatus, + runState?: Record, + ): Promise { + let rawEntries: StoredLogEntry[]; + let liveStreamLineCount: number; + let resumeLeafEntryStartIndex: number | undefined; + const resumeFromRunId = + typeof runState?.resume_from_run_id === "string" + ? runState.resume_from_run_id + : undefined; + const isResumeRun = Boolean(resumeFromRunId); + if (isTerminalStatus(runStatus) || isResumeRun) { + // Resume chains need the full history even while the leaf run is still + // active; otherwise a renderer restart hydrates only the final run. + // Non-resume in-progress runs keep using the single-run log so hydrate + // cannot race the live stream and double the active turn. + const authStatus = await this.getAuthCredentialsStatus(); + if (authStatus.kind !== "ready") { + return; + } + if (resumeFromRunId) { + const [ancestorResult, currentRunResult] = await Promise.all([ + authStatus.auth.client.getTaskRunSessionLogsResult( + taskId, + resumeFromRunId, + { limit: 100000 }, + ), + authStatus.auth.client.getTaskRunSessionLogsResult( taskId, taskRunId, { limit: 100000 }, - ); - } catch (err) { - this.d.log.warn("Failed to fetch session-log chain for hydrate", { + ), + ]); + if (!ancestorResult.complete || !currentRunResult.complete) { + this.d.log.warn("Resume session log hydration was incomplete", { taskId, taskRunId, - err, + resumeFromRunId, + ancestorComplete: ancestorResult.complete, + currentRunComplete: currentRunResult.complete, }); return; } - totalLineCount = rawEntries.length; + const ancestorEntries: StoredLogEntry[] = ancestorResult.entries; + const currentRunEntries: StoredLogEntry[] = currentRunResult.entries; + + const ancestorKeys = ancestorEntries.map((entry) => + JSON.stringify(entry), + ); + const currentKeys = currentRunEntries.map((entry) => + JSON.stringify(entry), + ); + const overlap = suffixPrefixOverlap(ancestorKeys, currentKeys); + const persistedLeafEntries = currentRunEntries.slice(overlap); + const leafLogs = await this.fetchSessionLogs(logUrl, taskRunId); + const leafKeys = new Set( + persistedLeafEntries.map((entry) => JSON.stringify(entry)), + ); + rawEntries = [ + ...ancestorEntries, + ...persistedLeafEntries, + ...leafLogs.rawEntries.filter( + (entry) => !leafKeys.has(JSON.stringify(entry)), + ), + ]; + resumeLeafEntryStartIndex = ancestorEntries.length; + liveStreamLineCount = Math.max( + leafLogs.totalLineCount, + persistedLeafEntries.length, + ); } else { - const parsed = await this.fetchSessionLogs(logUrl, taskRunId); - rawEntries = parsed.rawEntries; - totalLineCount = parsed.totalLineCount; + const result = await authStatus.auth.client.getTaskRunSessionLogsResult( + taskId, + taskRunId, + { limit: 100000 }, + ); + if (!result.complete) { + this.d.log.warn("Session log hydration was incomplete", { + taskId, + taskRunId, + }); + return; + } + rawEntries = result.entries; + liveStreamLineCount = rawEntries.length; } + } else { + const parsed = await this.fetchSessionLogs(logUrl, taskRunId); + rawEntries = parsed.rawEntries; + liveStreamLineCount = parsed.totalLineCount; + } - const session = this.d.store.getSessionByTaskId(taskId); - if (!session || session.taskRunId !== taskRunId) { - return; - } + const session = this.d.store.getSessionByTaskId(taskId); + if (!session || session.taskRunId !== taskRunId) { + return; + } - const events = convertStoredEntriesToEvents(rawEntries); - const hasUserPrompt = events.some( - (e: AcpMessage) => - isJsonRpcRequest(e.message) && e.message.method === "session/prompt", + let events = convertStoredEntriesToEvents(rawEntries, undefined, { + taskRunId, + startEntryIndex: 0, + firstPositionedEntryIndex: resumeLeafEntryStartIndex, + }); + if (isResumeRun && session.events.length > 0) { + const inheritedEvents = reconcileLiveEventsWithHydratedEvents( + session.events, + events, ); - - // Seed the optimistic user-message bubble whenever the agent has - // not yet recorded an initial `session/prompt` request — covers the - // brand-new task case as well as "agent has emitted lifecycle - // notifications but hasn't received its first prompt yet". Prefer the - // stashed initial prompt (which carries the channel CONTEXT.md block, so - // its chip renders right away) over the bare task description. - const seedContent = - this.initialCloudOptimisticPrompt.get(taskId) ?? taskDescription; - if (!hasUserPrompt && seedContent?.trim()) { - this.d.store.appendOptimisticItem(taskRunId, { - type: "user_message", - content: seedContent, - timestamp: Date.now(), - }); - } - if (hasUserPrompt) { - // The real prompt has landed; the stash is no longer needed. - this.initialCloudOptimisticPrompt.delete(taskId); - this.d.store.clearTailOptimisticItems(taskRunId); + events = [...events, ...inheritedEvents]; + const watcher = this.cloudTaskWatchers.get(taskId); + const hasLeafLocalWatcherCursor = + watcher?.runId === taskRunId && + watcher.resumeHistoryCountOffset !== undefined; + if (hasLeafLocalWatcherCursor) { + liveStreamLineCount = Math.max( + liveStreamLineCount, + session.processedLineCount ?? 0, + ); } + } + const hasUserPrompt = events.some( + (e: AcpMessage) => + isJsonRpcRequest(e.message) && e.message.method === "session/prompt", + ); - if (rawEntries.length === 0) { - this.pendingPermissionHydratedRuns.add(taskRunId); - return; - } + // Seed the optimistic user-message bubble whenever the agent has + // not yet recorded an initial `session/prompt` request — covers the + // brand-new task case as well as "agent has emitted lifecycle + // notifications but hasn't received its first prompt yet". Prefer the + // stashed initial prompt (which carries the channel CONTEXT.md block, so + // its chip renders right away) over the bare task description. + const seedContent = + this.initialCloudOptimisticPrompt.get(taskId) ?? taskDescription; + if (!hasUserPrompt && seedContent?.trim()) { + this.d.store.appendOptimisticItem(taskRunId, { + type: "user_message", + content: seedContent, + timestamp: Date.now(), + }); + } + if (hasUserPrompt) { + // The real prompt has landed; the stash is no longer needed. + this.initialCloudOptimisticPrompt.delete(taskId); + this.d.store.clearTailOptimisticItems(taskRunId); + } - // If live updates already populated a processed count, don't overwrite - // that newer state with the persisted baseline fetched during startup. - if ( - session.processedLineCount !== undefined && - session.processedLineCount > 0 - ) { - this.surfacePersistedPendingPermissions(taskRunId, rawEntries); - this.pendingPermissionHydratedRuns.add(taskRunId); - return; - } + if (rawEntries.length === 0) { + this.pendingPermissionHydratedRuns.add(taskRunId); + return { + historyEntryCount: 0, + liveStreamLineCount, + }; + } - this.d.store.updateSession(taskRunId, { - events, - isCloud: true, - logUrl: logUrl ?? session.logUrl, - processedLineCount: totalLineCount, - }); + // If live updates already populated a processed count, don't overwrite + // that newer state with the persisted baseline fetched during startup. + if ( + session.processedLineCount !== undefined && + session.processedLineCount > 0 && + !isResumeRun + ) { this.surfacePersistedPendingPermissions(taskRunId, rawEntries); this.pendingPermissionHydratedRuns.add(taskRunId); - // Without this the "Galumphing…" indicator stays hidden when the hydrated - // baseline already contains an in-flight session/prompt — the live delta - // path otherwise sees delta <= 0 and never re-evaluates the tail. - this.updatePromptStateFromEvents(taskRunId, events); - })().catch((err: unknown) => { - this.d.log.warn("Failed to hydrate cloud task session from logs", { - taskId, - taskRunId, - err, - }); + return { + historyEntryCount: rawEntries.length, + liveStreamLineCount: session.processedLineCount, + }; + } + + this.d.store.updateSession(taskRunId, { + events, + isCloud: true, + logUrl: logUrl ?? session.logUrl, + cloudTranscriptEntryCount: rawEntries.length, + processedLineCount: liveStreamLineCount, }); + this.surfacePersistedPendingPermissions(taskRunId, rawEntries); + this.pendingPermissionHydratedRuns.add(taskRunId); + // Without this the "Galumphing…" indicator stays hidden when the hydrated + // baseline already contains an in-flight session/prompt — the live delta + // path otherwise sees delta <= 0 and never re-evaluates the tail. + this.updatePromptStateFromEvents(taskRunId, events); + return { + historyEntryCount: rawEntries.length, + liveStreamLineCount, + }; } private isCurrentCloudTaskWatcher( @@ -6683,7 +6961,14 @@ export class SessionService { // Already caught up — skip duplicate entries } else if (plan.kind === "append-tail") { const entriesToAppend = update.newEntries.slice(-plan.tailCount); - const newEvents = convertStoredEntriesToEvents(entriesToAppend); + const newEvents = convertStoredEntriesToEvents( + entriesToAppend, + undefined, + { + taskRunId, + startEntryIndex: expectedCount - entriesToAppend.length, + }, + ); if (hasSessionPromptEvent(newEvents)) { this.d.store.clearTailOptimisticItems(taskRunId); } @@ -6962,7 +7247,10 @@ export class SessionService { logUrl: string | undefined, processedLineCount: number, ): void { - const events = convertStoredEntriesToEvents(rawEntries); + const events = convertStoredEntriesToEvents(rawEntries, undefined, { + taskRunId, + startEntryIndex: 0, + }); if (hasSessionPromptEvent(events)) { this.d.store.clearTailOptimisticItems(taskRunId); } diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index 0c84840a5e..791269d1ac 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.test.ts @@ -135,6 +135,7 @@ const mockAuthenticatedClient = vi.hoisted(() => ({ finalizeTaskStagedArtifactUploads: vi.fn(), startGithubUserIntegrationConnect: vi.fn(), getTaskRunSessionLogs: vi.fn(), + getTaskRunSessionLogsResult: vi.fn(), })); type MockAuthenticatedClient = typeof mockAuthenticatedClient; @@ -346,7 +347,13 @@ vi.mock("@posthog/shared", async (importOriginal) => ({ ), })); const mockConvertStoredEntriesToEvents = vi.hoisted(() => - vi.fn<(entries: unknown[]) => unknown[]>(() => []), + vi.fn< + ( + entries: unknown[], + taskDescription?: string, + positionOptions?: unknown, + ) => unknown[] + >(() => []), ); vi.mock("@posthog/core/sessions/sessionEvents", async () => { @@ -376,6 +383,7 @@ vi.mock("@posthog/core/sessions/sessionEvents", async () => { message: {}, })), extractPromptText: vi.fn((p) => (typeof p === "string" ? p : "text")), + getStoredLogEventPosition: actual.getStoredLogEventPosition, getUserShellExecutesSinceLastPrompt: vi.fn(() => []), hasSessionPromptEvent: actual.hasSessionPromptEvent, isAbsoluteFolderPath: actual.isAbsoluteFolderPath, @@ -449,6 +457,10 @@ describe("SessionService", () => { mockAuthenticatedClient.getTaskRunSessionLogs.mockResolvedValue([]); mockSessionConfigStore.getPersistedConfigOptions.mockReturnValue(undefined); mockAdapterFns.getAdapter.mockReturnValue(undefined); + mockAuthenticatedClient.getTaskRunSessionLogsResult.mockResolvedValue({ + entries: [], + complete: true, + }); mockSessionStoreSetters.getSessionByTaskId.mockReturnValue(undefined); mockSessionStoreSetters.getSessions.mockReturnValue({}); mockAuth.fetchAuthState.mockResolvedValue({ @@ -3463,60 +3475,63 @@ describe("SessionService", () => { mockSessionStoreSetters.getSessions.mockReturnValue({ "run-123": completedSession, }); - mockAuthenticatedClient.getTaskRunSessionLogs.mockResolvedValue([ - { - type: "notification", - notification: { - method: "_posthog/sdk_session", - params: { - taskRunId: "run-123", - sessionId: "acp-session-1", - adapter: "claude", + mockAuthenticatedClient.getTaskRunSessionLogsResult.mockResolvedValue({ + complete: true, + entries: [ + { + type: "notification", + notification: { + method: "_posthog/sdk_session", + params: { + taskRunId: "run-123", + sessionId: "acp-session-1", + adapter: "claude", + }, }, }, - }, - { - type: "notification", - notification: { - method: "_posthog/run_started", - params: { - sessionId: "acp-session-1", - runId: "run-123", - taskId: "task-123", + { + type: "notification", + notification: { + method: "_posthog/run_started", + params: { + sessionId: "acp-session-1", + runId: "run-123", + taskId: "task-123", + }, }, }, - }, - { - type: "notification", - notification: { - method: "_posthog/permission_request", - params: { - requestId: "request-1", - toolCall: { - toolCallId: "tool-1", - title: "What animal do you prefer?", - kind: "other", - _meta: { - codeToolKind: "question", - questions: [ - { - question: "What animal do you prefer?", - options: [ - { label: "cats", description: "Cats" }, - { label: "dogs", description: "Dogs" }, - ], - }, - ], + { + type: "notification", + notification: { + method: "_posthog/permission_request", + params: { + requestId: "request-1", + toolCall: { + toolCallId: "tool-1", + title: "What animal do you prefer?", + kind: "other", + _meta: { + codeToolKind: "question", + questions: [ + { + question: "What animal do you prefer?", + options: [ + { label: "cats", description: "Cats" }, + { label: "dogs", description: "Dogs" }, + ], + }, + ], + }, }, + options: [ + { optionId: "option_0", name: "cats", kind: "allow_once" }, + { optionId: "option_1", name: "dogs", kind: "allow_once" }, + ], }, - options: [ - { optionId: "option_0", name: "cats", kind: "allow_once" }, - { optionId: "option_1", name: "dogs", kind: "allow_once" }, - ], }, }, - }, - ]); + ], + }); service.watchCloudTask( "task-123", @@ -3626,23 +3641,258 @@ describe("SessionService", () => { ).not.toHaveBeenCalled(); }); - it("hydrates an in-progress resumed run from the full session-log chain", async () => { + it.each([ + { name: "leaf-only response", responseShape: "leaf" }, + { name: "full-chain response", responseShape: "full" }, + { name: "overlapping chain window", responseShape: "overlap" }, + ])( + "hydrates an in-progress resumed run from a $name", + async ({ responseShape }) => { + const service = getSessionService(); + const priorPrompt = { + type: "acp_message" as const, + ts: 1700000000, + message: { + jsonrpc: "2.0" as const, + id: 1, + method: "session/prompt", + params: { prompt: [{ type: "text", text: "first request" }] }, + }, + }; + const resumePrompt = { + type: "acp_message" as const, + ts: 1700000060, + message: { + jsonrpc: "2.0" as const, + id: 2, + method: "session/prompt", + params: { prompt: [{ type: "text", text: "continue" }] }, + }, + }; + const resumeCompletion = { + type: "acp_message" as const, + ts: 1700000120, + message: { + jsonrpc: "2.0" as const, + method: "_posthog/turn_complete", + params: { sessionId: "session-1", stopReason: "end_turn" }, + }, + }; + const resumedSession = createMockSession({ + taskRunId: "run-456", + taskId: "task-123", + status: "disconnected", + isCloud: true, + events: [resumePrompt], + processedLineCount: 1, + optimisticItems: [ + { + id: "optimistic-follow-up", + type: "user_message", + content: "continue", + timestamp: 1700000001, + pinToTop: false, + }, + ], + }); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( + resumedSession, + ); + mockSessionStoreSetters.getSessions.mockReturnValue({ + "run-456": resumedSession, + }); + const parentEntries = [ + { timestamp: "2024-01-01T00:00:00Z", notification: {} }, + { timestamp: "2024-01-01T00:00:30Z", notification: {} }, + ]; + const leafEntries = [ + { timestamp: "2024-01-01T00:01:00Z", notification: {} }, + ]; + mockAuthenticatedClient.getTaskRunSessionLogsResult + .mockResolvedValueOnce({ entries: parentEntries, complete: true }) + .mockResolvedValueOnce({ + entries: + responseShape === "full" + ? [...parentEntries, ...leafEntries] + : responseShape === "overlap" + ? [parentEntries[1], ...leafEntries] + : leafEntries, + complete: true, + }); + mockTrpcLogs.readLocalLogs.query.mockResolvedValue( + JSON.stringify(leafEntries[0]), + ); + mockTrpcLogs.fetchS3Logs.query.mockResolvedValue(""); + mockConvertStoredEntriesToEvents.mockReturnValueOnce([ + priorPrompt, + resumePrompt, + resumeCompletion, + ]); + + service.watchCloudTask( + "task-123", + "run-456", + "https://api.anthropic.com", + 123, + undefined, + "https://logs.example.com/run-456", + undefined, + "claude", + undefined, + "first request", + undefined, + "in_progress", + undefined, + { resume_from_run_id: "run-123" }, + ); + + const subscribeOptions = mockTrpcCloudTask.onUpdate.subscribe.mock + .calls[0][1] as { onData: (update: unknown) => void }; + subscribeOptions.onData({ + kind: "snapshot", + taskId: "task-123", + runId: "run-456", + totalEntryCount: 3, + newEntries: [...parentEntries, ...leafEntries], + status: "in_progress", + }); + expect(mockSessionStoreSetters.appendEvents).not.toHaveBeenCalled(); + + expect(mockTrpcCloudTask.watch.mutate).toHaveBeenCalledWith({ + taskId: "task-123", + runId: "run-456", + apiHost: "https://api.anthropic.com", + teamId: 123, + resumeFromEntryCount: undefined, + }); + await vi.waitFor(() => { + expect( + mockAuthenticatedClient.getTaskRunSessionLogsResult, + ).toHaveBeenCalledWith("task-123", "run-123", { limit: 100000 }); + }); + expect( + mockAuthenticatedClient.getTaskRunSessionLogsResult, + ).toHaveBeenCalledWith("task-123", "run-456", { limit: 100000 }); + expect(mockConvertStoredEntriesToEvents).toHaveBeenCalledWith( + [...parentEntries, ...leafEntries], + undefined, + { + taskRunId: "run-456", + startEntryIndex: 0, + firstPositionedEntryIndex: parentEntries.length, + }, + ); + expect(mockSessionStoreSetters.updateSession).toHaveBeenCalledWith( + "run-456", + expect.objectContaining({ + events: [priorPrompt, resumePrompt, resumeCompletion], + processedLineCount: 1, + }), + ); + expect(mockSessionStoreSetters.updateSession).toHaveBeenCalledWith( + "run-456", + expect.objectContaining({ + isPromptPending: false, + promptStartedAt: null, + currentPromptId: null, + }), + ); + expect( + mockSessionStoreSetters.clearTailOptimisticItems, + ).toHaveBeenCalledWith("run-456"); + expect( + mockSessionStoreSetters.appendOptimisticItem, + ).not.toHaveBeenCalled(); + expect(mockSessionStoreSetters.appendEvents).not.toHaveBeenCalled(); + }, + ); + + it("reconciles repeated prompt occurrences and promptless live tails", async () => { const service = getSessionService(); + const ancestorPrompt = { + type: "acp_message" as const, + ts: 1700000010, + message: { + jsonrpc: "2.0" as const, + id: 1, + method: "session/prompt", + params: { prompt: [{ type: "text", text: "repeat request" }] }, + }, + }; + const currentLivePrompt = { + type: "acp_message" as const, + ts: 1700000040, + message: ancestorPrompt.message, + }; + const persistedCurrentPrompt = { + ...currentLivePrompt, + ts: 1700000041, + }; + const persistedAncestorMessage = { + type: "acp_message" as const, + ts: 1700000020, + message: { + jsonrpc: "2.0" as const, + method: "session/update", + params: { + update: { + sessionUpdate: "agent_message", + content: { type: "text", text: "ancestor complete" }, + }, + }, + }, + }; + const currentLiveChunk = { + type: "acp_message" as const, + ts: 1700000050, + message: { + jsonrpc: "2.0" as const, + method: "session/update", + params: { + sessionId: "current-session", + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text: "current partial" }, + }, + }, + }, + }; + const ancestorCompletion = { + type: "acp_message" as const, + ts: 1700000030, + message: { + jsonrpc: "2.0" as const, + method: "_posthog/turn_complete", + params: { stopReason: "end_turn" }, + }, + }; + const currentCompletion = { + ...ancestorCompletion, + ts: 1700000060, + }; + const promptlessLiveOnlyEvent = { + type: "acp_message" as const, + ts: 1700000035, + message: { + jsonrpc: "2.0" as const, + method: "_posthog/usage_update", + params: { used: 42 }, + }, + }; const resumedSession = createMockSession({ taskRunId: "run-456", taskId: "task-123", - status: "disconnected", + status: "connected", isCloud: true, - events: [], - optimisticItems: [ - { - id: "optimistic-follow-up", - type: "user_message", - content: "continue", - timestamp: 1700000001, - pinToTop: false, - }, + events: [ + promptlessLiveOnlyEvent, + persistedAncestorMessage, + ancestorCompletion, + currentLivePrompt, + currentLiveChunk, + currentCompletion, ], + processedLineCount: 1, }); mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( resumedSession, @@ -3650,43 +3900,119 @@ describe("SessionService", () => { mockSessionStoreSetters.getSessions.mockReturnValue({ "run-456": resumedSession, }); - const chainedEntries = [ - { timestamp: "2024-01-01T00:00:00Z", notification: {} }, - { timestamp: "2024-01-01T00:01:00Z", notification: {} }, - ]; - mockAuthenticatedClient.getTaskRunSessionLogs.mockResolvedValue( - chainedEntries, - ); - mockTrpcLogs.readLocalLogs.query.mockResolvedValue( - "leaf log should not be used", - ); - mockTrpcLogs.fetchS3Logs.query.mockResolvedValue( - "leaf s3 log should not be used", + const parentEntry = { + timestamp: "2024-01-01T00:01:00Z", + notification: {}, + }; + mockAuthenticatedClient.getTaskRunSessionLogsResult + .mockResolvedValueOnce({ entries: [parentEntry], complete: true }) + .mockResolvedValueOnce({ entries: [parentEntry], complete: true }); + mockTrpcLogs.readLocalLogs.query.mockResolvedValue(""); + mockTrpcLogs.fetchS3Logs.query.mockResolvedValue(""); + mockConvertStoredEntriesToEvents.mockReturnValueOnce([ + ancestorPrompt, + persistedAncestorMessage, + ancestorCompletion, + persistedCurrentPrompt, + ]); + + service.watchCloudTask( + "task-123", + "run-456", + "https://api.anthropic.com", + 123, + undefined, + "https://logs.example.com/run-456", + undefined, + "claude", + undefined, + undefined, + undefined, + "in_progress", + undefined, + { resume_from_run_id: "run-123" }, ); - const priorPrompt = { + await vi.waitFor(() => { + expect(mockSessionStoreSetters.updateSession).toHaveBeenCalledWith( + "run-456", + expect.objectContaining({ + events: [ + ancestorPrompt, + persistedAncestorMessage, + ancestorCompletion, + persistedCurrentPrompt, + promptlessLiveOnlyEvent, + currentLiveChunk, + currentCompletion, + ], + }), + ); + }); + }); + + it("preserves a promptless current completion that only matches an ancestor turn", async () => { + const service = getSessionService(); + const ancestorPrompt = { type: "acp_message" as const, - ts: 1700000000, + ts: 1700000010, message: { jsonrpc: "2.0" as const, id: 1, method: "session/prompt", - params: { prompt: [{ type: "text", text: "first request" }] }, + params: { prompt: [{ type: "text", text: "ancestor request" }] }, }, }; - const resumePrompt = { + const currentPrompt = { type: "acp_message" as const, - ts: 1700000060, + ts: 1700000040, message: { jsonrpc: "2.0" as const, id: 2, method: "session/prompt", - params: { prompt: [{ type: "text", text: "continue" }] }, + params: { prompt: [{ type: "text", text: "current request" }] }, }, }; + const ancestorCompletion = { + type: "acp_message" as const, + ts: 1700000030, + message: { + jsonrpc: "2.0" as const, + method: "_posthog/turn_complete", + params: { stopReason: "end_turn" }, + }, + }; + const currentCompletion = { + ...ancestorCompletion, + ts: 1700000060, + }; + const resumedSession = createMockSession({ + taskRunId: "run-456", + taskId: "task-123", + status: "connected", + isCloud: true, + events: [currentCompletion], + processedLineCount: 1, + }); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( + resumedSession, + ); + mockSessionStoreSetters.getSessions.mockReturnValue({ + "run-456": resumedSession, + }); + const parentEntry = { + timestamp: "2024-01-01T00:01:00Z", + notification: {}, + }; + mockAuthenticatedClient.getTaskRunSessionLogsResult + .mockResolvedValueOnce({ entries: [parentEntry], complete: true }) + .mockResolvedValueOnce({ entries: [parentEntry], complete: true }); + mockTrpcLogs.readLocalLogs.query.mockResolvedValue(""); + mockTrpcLogs.fetchS3Logs.query.mockResolvedValue(""); mockConvertStoredEntriesToEvents.mockReturnValueOnce([ - priorPrompt, - resumePrompt, + ancestorPrompt, + ancestorCompletion, + currentPrompt, ]); service.watchCloudTask( @@ -3699,7 +4025,7 @@ describe("SessionService", () => { undefined, "claude", undefined, - "first request", + undefined, undefined, "in_progress", undefined, @@ -3707,26 +4033,582 @@ describe("SessionService", () => { ); await vi.waitFor(() => { - expect( - mockAuthenticatedClient.getTaskRunSessionLogs, - ).toHaveBeenCalledWith("task-123", "run-456", { limit: 100000 }); + expect(mockSessionStoreSetters.updateSession).toHaveBeenCalledWith( + "run-456", + expect.objectContaining({ + events: [ + ancestorPrompt, + ancestorCompletion, + currentPrompt, + currentCompletion, + ], + }), + ); }); - expect(mockConvertStoredEntriesToEvents).toHaveBeenCalledWith( - chainedEntries, - ); - expect(mockSessionStoreSetters.updateSession).toHaveBeenCalledWith( - "run-456", - expect.objectContaining({ - events: [priorPrompt, resumePrompt], - processedLineCount: chainedEntries.length, - }), - ); - expect( - mockSessionStoreSetters.clearTailOptimisticItems, - ).toHaveBeenCalledWith("run-456"); - expect( - mockSessionStoreSetters.appendOptimisticItem, - ).not.toHaveBeenCalled(); + }); + + it("keeps immediate-resume watcher counts leaf-local while flushing buffered updates", async () => { + const service = getSessionService(); + const ancestorEvent = { + type: "acp_message" as const, + ts: 1700000000, + message: { + jsonrpc: "2.0" as const, + id: 1, + method: "session/prompt", + params: { prompt: [{ type: "text", text: "first request" }] }, + }, + }; + const leafEvent = { + type: "acp_message" as const, + ts: 1700000060, + message: { + jsonrpc: "2.0" as const, + id: 2, + method: "session/prompt", + params: { prompt: [{ type: "text", text: "continue" }] }, + }, + }; + const liveEvent = { + type: "acp_message" as const, + ts: 1700000120, + message: { + jsonrpc: "2.0" as const, + method: "session/update", + params: { update: { sessionUpdate: "agent_message_chunk" } }, + }, + }; + const resumedSession = createMockSession({ + taskRunId: "run-456", + taskId: "task-123", + status: "connected", + isCloud: true, + events: [ancestorEvent], + cloudTranscriptEntryCount: 3, + processedLineCount: 0, + }); + mockSessionStoreSetters.getSessionByTaskId.mockImplementation( + () => resumedSession, + ); + mockSessionStoreSetters.getSessions.mockImplementation(() => ({ + "run-456": resumedSession, + })); + mockSessionStoreSetters.updateSession.mockImplementation( + (_runId, updates) => Object.assign(resumedSession, updates), + ); + mockSessionStoreSetters.appendEvents.mockImplementation( + (_runId, events, processedLineCount) => { + resumedSession.events.push(...events); + if (processedLineCount !== undefined) { + resumedSession.processedLineCount = processedLineCount; + } + }, + ); + + const ancestorEntries = [ + { timestamp: "2024-01-01T00:00:00Z", notification: {} }, + { timestamp: "2024-01-01T00:00:01Z", notification: {} }, + { timestamp: "2024-01-01T00:00:02Z", notification: {} }, + ]; + const leafEntry = { + timestamp: "2024-01-01T00:01:00Z", + notification: {}, + }; + const liveEntry = { + timestamp: "2024-01-01T00:02:00Z", + notification: { method: "session/update" }, + }; + let resolveAncestor!: (result: { + entries: typeof ancestorEntries; + complete: boolean; + }) => void; + mockAuthenticatedClient.getTaskRunSessionLogsResult + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveAncestor = resolve; + }), + ) + .mockResolvedValueOnce({ + entries: [...ancestorEntries, leafEntry], + complete: true, + }); + mockTrpcLogs.readLocalLogs.query.mockResolvedValue( + JSON.stringify(leafEntry), + ); + mockTrpcLogs.fetchS3Logs.query.mockResolvedValue(""); + mockConvertStoredEntriesToEvents + .mockReturnValueOnce([ancestorEvent, leafEvent]) + .mockReturnValueOnce([liveEvent]); + + service.watchCloudTask( + "task-123", + "run-456", + "https://api.anthropic.com", + 123, + undefined, + "https://logs.example.com/run-456", + undefined, + "claude", + undefined, + "first request", + 3, + "in_progress", + undefined, + { resume_from_run_id: "run-123" }, + ); + + await vi.waitFor(() => { + expect( + mockAuthenticatedClient.getTaskRunSessionLogsResult, + ).toHaveBeenCalledTimes(2); + }); + + const subscribeOptions = mockTrpcCloudTask.onUpdate.subscribe.mock + .calls[0][1] as { onData: (update: unknown) => void }; + subscribeOptions.onData({ + kind: "logs", + taskId: "task-123", + runId: "run-456", + totalEntryCount: 5, + newEntries: [leafEntry, liveEntry], + }); + expect(mockSessionStoreSetters.appendEvents).not.toHaveBeenCalled(); + + resolveAncestor({ entries: ancestorEntries, complete: true }); + await vi.waitFor(() => { + expect(mockSessionStoreSetters.appendEvents).toHaveBeenCalledWith( + "run-456", + [liveEvent], + 2, + ); + }); + expect(mockSessionStoreSetters.updateSession).toHaveBeenCalledWith( + "run-456", + expect.objectContaining({ + events: expect.arrayContaining([ancestorEvent, leafEvent]), + processedLineCount: 1, + }), + ); + expect(resumedSession.events).toEqual([ + ancestorEvent, + leafEvent, + liveEvent, + ]); + expect(resumedSession.processedLineCount).toBe(2); + expect(mockTrpcCloudTask.watch.mutate).toHaveBeenCalledWith({ + taskId: "task-123", + runId: "run-456", + apiHost: "https://api.anthropic.com", + teamId: 123, + resumeFromEntryCount: 3, + }); + }); + + it("uses the full A→B transcript count when B resumes into C", async () => { + const service = getSessionService(); + const aEvent = { + type: "acp_message" as const, + ts: 1700000000, + message: { + jsonrpc: "2.0" as const, + id: 1, + method: "session/prompt", + params: { prompt: [{ type: "text", text: "start A" }] }, + }, + }; + const bEvent = { + type: "acp_message" as const, + ts: 1700000060, + message: { + jsonrpc: "2.0" as const, + id: 2, + method: "session/prompt", + params: { prompt: [{ type: "text", text: "resume B" }] }, + }, + }; + const cEvent = { + type: "acp_message" as const, + ts: 1700000120, + message: { + jsonrpc: "2.0" as const, + id: 3, + method: "session/prompt", + params: { prompt: [{ type: "text", text: "resume C" }] }, + }, + }; + const liveCEvent = { + type: "acp_message" as const, + ts: 1700000180, + message: { + jsonrpc: "2.0" as const, + method: "session/update", + params: { update: { sessionUpdate: "agent_message_chunk" } }, + }, + }; + let activeSession = createMockSession({ + taskRunId: "run-b", + taskId: "task-123", + status: "connected", + isCloud: true, + cloudStatus: "completed", + cloudBranch: "feature/resume-chain", + events: [aEvent, bEvent], + cloudTranscriptEntryCount: 7, + processedLineCount: 2, + }); + mockSessionStoreSetters.getSessionByTaskId.mockImplementation( + () => activeSession, + ); + mockSessionStoreSetters.getSessions.mockImplementation(() => ({ + [activeSession.taskRunId]: activeSession, + })); + mockSessionStoreSetters.setSession.mockImplementation((session) => { + activeSession = session; + }); + mockSessionStoreSetters.updateSession.mockImplementation( + (_runId, updates) => Object.assign(activeSession, updates), + ); + mockSessionStoreSetters.appendEvents.mockImplementation( + (_runId, events, processedLineCount) => { + activeSession.events.push(...events); + if (processedLineCount !== undefined) { + activeSession.processedLineCount = processedLineCount; + } + }, + ); + + mockAuthenticatedClient.getTaskRun.mockResolvedValue({ + id: "run-b", + task: "task-123", + team: 123, + branch: "feature/resume-chain", + runtime_adapter: "claude", + model: "claude-sonnet-4-20250514", + reasoning_effort: null, + environment: "cloud", + status: "completed", + log_url: "https://example.com/logs/run-b", + error_message: null, + output: {}, + state: { resume_from_run_id: "run-a" }, + created_at: "2026-04-14T00:00:00Z", + updated_at: "2026-04-14T00:05:00Z", + completed_at: "2026-04-14T00:05:00Z", + }); + mockAuthenticatedClient.getTask.mockResolvedValue(createMockTask()); + mockAuthenticatedClient.runTaskInCloud.mockResolvedValue( + createMockTask({ + latest_run: { + id: "run-c", + task: "task-123", + team: 123, + branch: "feature/resume-chain", + runtime_adapter: "claude", + model: "claude-sonnet-4-20250514", + reasoning_effort: null, + environment: "cloud", + status: "queued", + log_url: "https://example.com/logs/run-c", + error_message: null, + output: {}, + state: { resume_from_run_id: "run-b" }, + created_at: "2026-04-14T00:06:00Z", + updated_at: "2026-04-14T00:06:00Z", + completed_at: null, + }, + }), + ); + + const inheritedEntries = Array.from({ length: 7 }, (_, index) => ({ + timestamp: `2024-01-01T00:00:0${index}Z`, + notification: {}, + })); + const cEntry = { + timestamp: "2024-01-01T00:02:00Z", + notification: {}, + }; + const liveCEntry = { + timestamp: "2024-01-01T00:03:00Z", + notification: { method: "session/update" }, + }; + let resolveInherited!: (result: { + entries: typeof inheritedEntries; + complete: boolean; + }) => void; + mockAuthenticatedClient.getTaskRunSessionLogsResult + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveInherited = resolve; + }), + ) + .mockResolvedValueOnce({ + entries: [...inheritedEntries, cEntry], + complete: true, + }); + mockTrpcLogs.readLocalLogs.query.mockResolvedValue( + JSON.stringify(cEntry), + ); + mockTrpcLogs.fetchS3Logs.query.mockResolvedValue(""); + mockConvertStoredEntriesToEvents + .mockReturnValueOnce([aEvent, bEvent, cEvent]) + .mockReturnValueOnce([liveCEvent]); + + const result = await service.sendPrompt("task-123", "resume C"); + expect(result.stopReason).toBe("queued"); + expect(activeSession).toEqual( + expect.objectContaining({ + taskRunId: "run-c", + cloudTranscriptEntryCount: 7, + processedLineCount: 0, + }), + ); + await vi.waitFor(() => { + expect( + mockAuthenticatedClient.getTaskRunSessionLogsResult, + ).toHaveBeenCalledTimes(2); + }); + + const subscribeOptions = mockTrpcCloudTask.onUpdate.subscribe.mock + .calls[0][1] as { onData: (update: unknown) => void }; + subscribeOptions.onData({ + kind: "logs", + taskId: "task-123", + runId: "run-c", + totalEntryCount: 9, + newEntries: [cEntry, liveCEntry], + }); + expect(mockSessionStoreSetters.appendEvents).not.toHaveBeenCalled(); + + resolveInherited({ entries: inheritedEntries, complete: true }); + await vi.waitFor(() => { + expect(mockSessionStoreSetters.appendEvents).toHaveBeenCalledWith( + "run-c", + [liveCEvent], + 2, + ); + }); + expect(activeSession.events).toEqual([ + aEvent, + bEvent, + cEvent, + liveCEvent, + ]); + expect(activeSession.processedLineCount).toBe(2); + expect(activeSession.cloudTranscriptEntryCount).toBe(9); + expect(mockTrpcCloudTask.watch.mutate).toHaveBeenCalledWith({ + taskId: "task-123", + runId: "run-c", + apiHost: "https://api.anthropic.com", + teamId: 123, + resumeFromEntryCount: 7, + }); + }); + + it("switches a cold-reload watcher to leaf-local counts after hydration recovers", async () => { + const service = getSessionService(); + const resumePrompt = { + type: "acp_message" as const, + ts: 1700000060, + message: { + jsonrpc: "2.0" as const, + id: 2, + method: "session/prompt", + params: { prompt: [{ type: "text", text: "continue" }] }, + }, + }; + const resumedSession = createMockSession({ + taskRunId: "run-456", + taskId: "task-123", + status: "connected", + isCloud: true, + events: [], + processedLineCount: 0, + }); + mockSessionStoreSetters.getSessionByTaskId.mockImplementation( + () => resumedSession, + ); + mockSessionStoreSetters.getSessions.mockImplementation(() => ({ + "run-456": resumedSession, + })); + mockSessionStoreSetters.updateSession.mockImplementation( + (_runId, updates) => + Object.assign(resumedSession, { + ...updates, + ...(updates.events ? { events: [...updates.events] } : {}), + }), + ); + mockSessionStoreSetters.appendEvents.mockImplementation( + (_runId, events, processedLineCount) => { + resumedSession.events.push(...events); + if (processedLineCount !== undefined) { + resumedSession.processedLineCount = processedLineCount; + } + }, + ); + const parentEntries = Array.from({ length: 7 }, (_, index) => ({ + timestamp: `2024-01-01T00:00:0${index}Z`, + notification: {}, + })); + const leafEntry = { + timestamp: "2024-01-01T00:01:00Z", + notification: {}, + }; + const liveEntry = { + timestamp: "2024-01-01T00:02:00Z", + notification: { method: "session/update" }, + }; + const liveEvent = { + type: "acp_message" as const, + ts: 1700000120, + message: { + jsonrpc: "2.0" as const, + method: "session/update", + params: { update: { sessionUpdate: "agent_message_chunk" } }, + }, + }; + let resolveAncestor!: (result: { + entries: typeof parentEntries; + complete: boolean; + }) => void; + mockAuthenticatedClient.getTaskRunSessionLogsResult + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveAncestor = resolve; + }), + ) + .mockResolvedValueOnce({ + entries: [...parentEntries, leafEntry], + complete: true, + }); + mockConvertStoredEntriesToEvents.mockImplementation((entries) => + entries.some( + (entry) => + (entry as { timestamp?: string }).timestamp === liveEntry.timestamp, + ) + ? [liveEvent] + : [resumePrompt], + ); + + const watch = (): void => { + service.watchCloudTask( + "task-123", + "run-456", + "https://api.anthropic.com", + 123, + undefined, + "https://logs.example.com/run-456", + undefined, + "claude", + undefined, + "first request", + undefined, + "in_progress", + undefined, + { resume_from_run_id: "run-123" }, + ); + }; + + watch(); + await vi.waitFor(() => { + expect( + mockAuthenticatedClient.getTaskRunSessionLogsResult, + ).toHaveBeenCalledTimes(2); + }); + expect(mockTrpcCloudTask.watch.mutate).toHaveBeenCalledWith({ + taskId: "task-123", + runId: "run-456", + apiHost: "https://api.anthropic.com", + teamId: 123, + resumeFromEntryCount: undefined, + }); + watch(); + expect( + mockAuthenticatedClient.getTaskRunSessionLogsResult, + ).toHaveBeenCalledTimes(2); + + const subscribeOptions = mockTrpcCloudTask.onUpdate.subscribe.mock + .calls[0][1] as { onData: (update: unknown) => void }; + subscribeOptions.onData({ + kind: "logs", + taskId: "task-123", + runId: "run-456", + totalEntryCount: 8, + newEntries: [...parentEntries, leafEntry], + }); + expect(mockSessionStoreSetters.appendEvents).not.toHaveBeenCalled(); + + resolveAncestor({ entries: parentEntries, complete: false }); + await vi.waitFor(() => { + expect(mockSessionStoreSetters.appendEvents).toHaveBeenCalledWith( + "run-456", + [resumePrompt], + 8, + ); + }); + expect(resumedSession.processedLineCount).toBe(8); + + let resolveRetryAncestor!: (result: { + entries: typeof parentEntries; + complete: boolean; + }) => void; + mockAuthenticatedClient.getTaskRunSessionLogsResult + .mockImplementationOnce( + () => + new Promise((resolve) => { + resolveRetryAncestor = resolve; + }), + ) + .mockResolvedValueOnce({ + entries: [...parentEntries, leafEntry], + complete: true, + }); + mockTrpcLogs.readLocalLogs.query.mockResolvedValue( + JSON.stringify(leafEntry), + ); + mockTrpcLogs.fetchS3Logs.query.mockResolvedValue(""); + + watch(); + await vi.waitFor(() => { + expect( + mockAuthenticatedClient.getTaskRunSessionLogsResult, + ).toHaveBeenCalledTimes(4); + }); + const appendCountBeforeRetryUpdate = + mockSessionStoreSetters.appendEvents.mock.calls.length; + subscribeOptions.onData({ + kind: "logs", + taskId: "task-123", + runId: "run-456", + totalEntryCount: 9, + newEntries: [liveEntry], + }); + expect(mockSessionStoreSetters.appendEvents).toHaveBeenCalledTimes( + appendCountBeforeRetryUpdate, + ); + + resolveRetryAncestor({ entries: parentEntries, complete: true }); + await vi.waitFor(() => { + expect(mockSessionStoreSetters.updateSession).toHaveBeenCalledWith( + "run-456", + expect.objectContaining({ + events: [resumePrompt], + processedLineCount: 1, + }), + ); + }); + await vi.waitFor(() => { + expect(mockSessionStoreSetters.appendEvents).toHaveBeenLastCalledWith( + "run-456", + [liveEvent], + 2, + ); + }); + expect(resumedSession.events).toEqual([resumePrompt, liveEvent]); + expect(resumedSession.processedLineCount).toBe(2); + expect(resumedSession.cloudTranscriptEntryCount).toBe(9); }); it("ignores stale async starts when the same watcher is replaced", async () => { From 2b5dc8c1e2220b96eaac6b17fc71cf215502f000 Mon Sep 17 00:00:00 2001 From: Alessandro Pogliaghi Date: Tue, 21 Jul 2026 09:50:18 +0100 Subject: [PATCH 5/5] feat(cloud): expose steering controls for cloud sessions (#3541) --- packages/core/src/sessions/sessionService.ts | 81 ++++++++++++------- .../sessions/components/SessionView.tsx | 6 +- .../sessions/hooks/useMessagingMode.ts | 4 +- .../sessions/sessionServiceHost.test.ts | 46 +++++++++-- 4 files changed, 96 insertions(+), 41 deletions(-) diff --git a/packages/core/src/sessions/sessionService.ts b/packages/core/src/sessions/sessionService.ts index fc34739647..2879033ee8 100644 --- a/packages/core/src/sessions/sessionService.ts +++ b/packages/core/src/sessions/sessionService.ts @@ -2904,7 +2904,11 @@ export class SessionService { isNotification(msg.method, POSTHOG_NOTIFICATIONS.RUN_STARTED) ) { const session = this.d.store.getSessions()[taskRunId]; - const params = (msg as { params?: { agentVersion?: unknown } }).params; + const params = ( + msg as { + params?: { agentVersion?: unknown; steering?: unknown }; + } + ).params; const agentVersion = typeof params?.agentVersion === "string" ? params.agentVersion @@ -2913,6 +2917,12 @@ export class SessionService { if (agentVersion && session?.agentVersion !== agentVersion) { updates.agentVersion = agentVersion; } + if ( + typeof params?.steering === "string" && + session?.steering !== params.steering + ) { + updates.steering = params.steering; + } if (session?.isCloud && session.status !== "connected") { updates.status = "connected"; } @@ -3387,22 +3397,27 @@ export class SessionService { // Steer: the user sent a message mid-turn and asked to fold it into the // running turn rather than queue it. Adapters that negotiated // `steering: "native"` (Claude, codex) inject at the next tool boundary; - // unknown adapters cancel and resend. Cloud has no real mid-turn steer - // (the backend only delivers messages between turns), so it falls through - // to the queue; compaction too. - if ( - options?.steer && - !session.isCloud && - session.isPromptPending && - !session.isCompacting - ) { + // unknown local adapters cancel and resend. Cloud sessions only enter this + // path after the sandbox advertises native steering; compaction still queues. + if (options?.steer && session.isPromptPending && !session.isCompacting) { if (sessionSupportsNativeSteer(session)) { - return this.sendSteerPrompt(session, prompt); + if (session.isCloud) { + if (session.status === "connected") { + return this.sendCloudPrompt(session, prompt, { + skipQueueGuard: true, + steer: true, + }); + } + } else { + return this.sendSteerPrompt(session, prompt); + } } - await this.cancelPrompt(taskId); - const refreshed = this.d.store.getSessionByTaskId(taskId); - if (refreshed) { - session = refreshed; + if (!session.isCloud) { + await this.cancelPrompt(taskId); + const refreshed = this.d.store.getSessionByTaskId(taskId); + if (refreshed) { + session = refreshed; + } } } @@ -3934,7 +3949,7 @@ export class SessionService { private async sendCloudPrompt( session: AgentSession, prompt: string | ContentBlock[], - options?: { skipQueueGuard?: boolean }, + options?: { skipQueueGuard?: boolean; steer?: boolean }, ): Promise<{ stopReason: string }> { const normalizedPrompt = await this.resolveCloudPrompt(prompt); const transport = this.d.h.getCloudPromptTransport(normalizedPrompt); @@ -4066,19 +4081,24 @@ export class SessionService { if (artifactIds.length > 0) { params.artifact_ids = artifactIds; } + if (options?.steer) { + params.steer = true; + } const currentSessionBeforeSend = this.getSessionByRunId(session.taskRunId) ?? session; const idleEvidenceBeforeSend = this.cloudRunIdleTracker.capture( currentSessionBeforeSend, ); - this.d.store.updateSession(session.taskRunId, { - isPromptPending: true, - promptStartedAt: Date.now(), - pausedDurationMs: 0, - agentIdleForRunId: undefined, - }); - this.cloudRunIdleTracker.markBusy(currentSessionBeforeSend); + if (!options?.steer) { + this.d.store.updateSession(session.taskRunId, { + isPromptPending: true, + promptStartedAt: Date.now(), + pausedDurationMs: 0, + agentIdleForRunId: undefined, + }); + this.cloudRunIdleTracker.markBusy(currentSessionBeforeSend); + } this.d.store.appendOptimisticItem(session.taskRunId, { type: "user_message", content: transport.promptText, @@ -4091,6 +4111,7 @@ export class SessionService { is_initial: session.events.length === 0, execution_type: "cloud", prompt_length_chars: transport.promptText.length, + ...(options?.steer ? { is_steer: true } : {}), }); try { @@ -4108,7 +4129,7 @@ export class SessionService { } const commandResult = result.result as - | { queued?: boolean; stopReason?: string } + | { queued?: boolean; steered?: boolean; stopReason?: string } | undefined; const stopReason = commandResult?.queued ? "queued" @@ -4116,15 +4137,17 @@ export class SessionService { return { stopReason }; } catch (error) { - this.d.store.updateSession(session.taskRunId, { - isPromptPending: false, - promptStartedAt: null, - }); + if (!options?.steer) { + this.d.store.updateSession(session.taskRunId, { + isPromptPending: false, + promptStartedAt: null, + }); + } this.d.store.clearTailOptimisticItems(session.taskRunId); const currentSessionAfterFailure = this.getSessionByRunId( session.taskRunId, ); - if (currentSessionAfterFailure) { + if (currentSessionAfterFailure && !options?.steer) { const restoreResult = this.cloudRunIdleTracker.restoreAfterFailedSend( idleEvidenceBeforeSend, currentSessionAfterFailure, diff --git a/packages/ui/src/features/sessions/components/SessionView.tsx b/packages/ui/src/features/sessions/components/SessionView.tsx index cb9566e436..d9236f3a1d 100644 --- a/packages/ui/src/features/sessions/components/SessionView.tsx +++ b/packages/ui/src/features/sessions/components/SessionView.tsx @@ -704,13 +704,11 @@ export function SessionView({ ) : null } messagingModeToggle={ - taskId && !isCloudRun ? ( + taskId ? ( ) : undefined } - onToggleMessagingMode={ - isCloudRun ? undefined : toggleMessagingMode - } + onToggleMessagingMode={toggleMessagingMode} onPromptRecall={handlePromptRecall} onBeforeSubmit={handleBeforeSubmit} onSubmit={handleSubmit} diff --git a/packages/ui/src/features/sessions/hooks/useMessagingMode.ts b/packages/ui/src/features/sessions/hooks/useMessagingMode.ts index f4a93dd63e..480bbda90f 100644 --- a/packages/ui/src/features/sessions/hooks/useMessagingMode.ts +++ b/packages/ui/src/features/sessions/hooks/useMessagingMode.ts @@ -19,8 +19,8 @@ export function useMessagingMode(taskId: string | undefined): MessagingMode { * Whether the task's session steers natively (folds a mid-turn message into the * running turn) versus falling back to interrupt-and-resend. Driven by the * adapter's negotiated `steering` capability — same decision as the host's - * sendPrompt gate — so Claude and codex steer, while cloud - * resend. Drives the steer label/tooltip, not whether steer is allowed. + * sendPrompt gate, including capability-advertising cloud sandboxes. Drives + * the steer label/tooltip, not whether steer is allowed. */ export function useSupportsNativeSteer(taskId: string | undefined): boolean { return useSessionStore((s) => { diff --git a/packages/ui/src/features/sessions/sessionServiceHost.test.ts b/packages/ui/src/features/sessions/sessionServiceHost.test.ts index 791269d1ac..28a710fc2c 100644 --- a/packages/ui/src/features/sessions/sessionServiceHost.test.ts +++ b/packages/ui/src/features/sessions/sessionServiceHost.test.ts @@ -3226,7 +3226,7 @@ describe("SessionService", () => { ); }); - it("captures agentVersion from run_started params onto the session", async () => { + it("captures agent capabilities from run_started params onto the session", async () => { const service = getSessionService(); const hydratedSession = createMockSession({ taskRunId: "run-123", @@ -3256,6 +3256,7 @@ describe("SessionService", () => { runId: "run-123", taskId: "task-123", agentVersion: "0.42.3", + steering: "native", }, }, }; @@ -3275,6 +3276,7 @@ describe("SessionService", () => { "run-123", expect.objectContaining({ agentVersion: "0.42.3", + steering: "native", status: "connected", }), ); @@ -5472,11 +5474,7 @@ describe("SessionService", () => { expect(mockTrpcCloudTask.sendCommand.mutate).not.toHaveBeenCalled(); }); - it("queues a cloud steer instead of interrupting the running turn", async () => { - // Regression: cloud has no native mid-turn steer, so steering used to - // fall back to cancel-then-resend — which surfaced as a jarring user - // interruption. Cloud steer must now queue like a normal message and - // never cancel the running turn. + it("sends a native cloud steer immediately", async () => { const service = getSessionService(); mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( createMockSession({ @@ -5484,6 +5482,42 @@ describe("SessionService", () => { cloudStatus: "in_progress", status: "connected", isPromptPending: true, + steering: "native", + }), + ); + mockTrpcCloudTask.sendCommand.mutate.mockResolvedValue({ + success: true, + result: { stopReason: "steered", steered: true }, + }); + + const prompt: ContentBlock[] = [{ type: "text", text: "steer me" }]; + const result = await service.sendPrompt("task-123", prompt, { + steer: true, + }); + + expect(result.stopReason).toBe("steered"); + expect(mockSessionStoreSetters.enqueueMessage).not.toHaveBeenCalled(); + expect(mockTrpcCloudTask.sendCommand.mutate).toHaveBeenCalledWith( + expect.objectContaining({ + method: "user_message", + params: { content: "steer me", steer: true }, + }), + ); + expect(mockSessionStoreSetters.updateSession).not.toHaveBeenCalledWith( + "run-123", + expect.objectContaining({ isPromptPending: false }), + ); + }); + + it("queues a cloud steer when the sandbox lacks the capability", async () => { + const service = getSessionService(); + mockSessionStoreSetters.getSessionByTaskId.mockReturnValue( + createMockSession({ + isCloud: true, + cloudStatus: "in_progress", + status: "connected", + isPromptPending: true, + steering: undefined, }), );