diff --git a/packages/agent/src/adapters/claude/session/jsonl-hydration.test.ts b/packages/agent/src/adapters/claude/session/jsonl-hydration.test.ts index 2bb4f3d38e..9e094a5c8a 100644 --- a/packages/agent/src/adapters/claude/session/jsonl-hydration.test.ts +++ b/packages/agent/src/adapters/claude/session/jsonl-hydration.test.ts @@ -1595,6 +1595,37 @@ describe("hydrateSessionJsonl", () => { }; } + it("returns the selected conversation when hydrating from logs", async () => { + const posthogAPI = { + getTaskRun: vi.fn().mockResolvedValue({ log_url: "https://logs.test" }), + fetchTaskRunLogs: vi.fn().mockResolvedValue([ + entry("user_message", { + content: { type: "text", text: "previous request" }, + }), + ]), + } as unknown as PostHogAPIClient; + const log = { info: vi.fn(), warn: vi.fn() }; + + const result = await hydrateSessionJsonl({ + sessionId, + cwd, + taskId: "t1", + runId: "r1", + posthogAPI, + log, + }); + + expect(result).toEqual({ + hasSession: true, + conversation: [ + { + role: "user", + content: [{ type: "text", text: "previous request" }], + }, + ], + }); + }); + it("sanitizes an existing file and skips S3 hydration", async () => { const file = await writeSessionFile(); const { posthogAPI, log } = makeDeps(); @@ -1608,7 +1639,7 @@ describe("hydrateSessionJsonl", () => { log, }); - expect(result).toBe(true); + expect(result).toEqual({ hasSession: true }); expect( (posthogAPI as unknown as { getTaskRun: ReturnType }) .getTaskRun, @@ -1635,7 +1666,7 @@ describe("hydrateSessionJsonl", () => { log, }); - expect(result).toBe(true); + expect(result).toEqual({ hasSession: true }); expect(log.warn).toHaveBeenCalledWith( "Failed to sanitize existing session JSONL", expect.anything(), diff --git a/packages/agent/src/adapters/claude/session/jsonl-hydration.ts b/packages/agent/src/adapters/claude/session/jsonl-hydration.ts index 65cf55fe65..6f9d5f4920 100644 --- a/packages/agent/src/adapters/claude/session/jsonl-hydration.ts +++ b/packages/agent/src/adapters/claude/session/jsonl-hydration.ts @@ -10,12 +10,17 @@ import { isEmptyContentBlock } from "../../../utils/acp-content"; import { neutralizeUnprocessableImages } from "../image-sanitization"; import { supports1MContext } from "./models"; -interface ConversationTurn { +export interface ConversationTurn { role: "user" | "assistant"; content: ContentBlock[]; toolCalls?: ToolCallInfo[]; } +export interface HydrateSessionJsonlResult { + hasSession: boolean; + conversation?: ConversationTurn[]; +} + interface ToolCallInfo { toolCallId: string; toolName: string; @@ -681,7 +686,7 @@ export async function hydrateSessionJsonl(params: { permissionMode?: string; posthogAPI: PostHogAPIClient; log: HydrationLog; -}): Promise { +}): Promise { const { posthogAPI, log } = params; try { @@ -702,7 +707,7 @@ export async function hydrateSessionJsonl(params: { error: err instanceof Error ? err.message : String(err), }); } - return true; + return { hasSession: true }; } catch { // File doesn't exist, proceed with hydration } @@ -710,13 +715,13 @@ export async function hydrateSessionJsonl(params: { const taskRun = await posthogAPI.getTaskRun(params.taskId, params.runId); if (!taskRun.log_url) { log.info("No log URL, skipping JSONL hydration"); - return false; + return { hasSession: false }; } const entries = await posthogAPI.fetchTaskRunLogs(taskRun); if (entries.length === 0) { log.info("No S3 log entries, skipping JSONL hydration"); - return false; + return { hasSession: false }; } const entryCounts: Record = {}; @@ -742,7 +747,7 @@ export async function hydrateSessionJsonl(params: { if (allTurns.length === 0) { log.info("No conversation to hydrate, skipping JSONL hydration"); - return false; + return { hasSession: false }; } const maxTokens = supports1MContext(params.model ?? "") @@ -774,12 +779,12 @@ export async function hydrateSessionJsonl(params: { turns: conversation.length, lines: jsonlLines.length, }); - return true; + return { hasSession: true, conversation }; } catch (err) { log.warn("Failed to hydrate session JSONL, continuing", { sessionId: params.sessionId, error: err instanceof Error ? err.message : String(err), }); - return false; + return { hasSession: false }; } } diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 3a1853f7b9..75dbb4c5a5 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -841,7 +841,7 @@ export class AgentServer { } try { - const hasSession = await hydrateSessionJsonl({ + const { hasSession } = await hydrateSessionJsonl({ sessionId: priorSessionId, cwd, taskId: payload.task_id, diff --git a/packages/workspace-server/src/services/agent/agent.test.ts b/packages/workspace-server/src/services/agent/agent.test.ts index 14204e4065..1109fbf868 100644 --- a/packages/workspace-server/src/services/agent/agent.test.ts +++ b/packages/workspace-server/src/services/agent/agent.test.ts @@ -15,6 +15,12 @@ const mockNewSession = vi.hoisted(() => configOptions: [], }), ); +const mockResumeSession = vi.hoisted(() => + vi.fn().mockResolvedValue({ configOptions: [] }), +); +const mockPrompt = vi.hoisted(() => + vi.fn().mockResolvedValue({ stopReason: "end_turn" }), +); const mockAcpClient = vi.hoisted(() => ({ current: undefined as @@ -40,7 +46,8 @@ const mockClientSideConnection = vi.hoisted(() => this.initialize = vi.fn().mockResolvedValue({}); this.newSession = mockNewSession; this.loadSession = vi.fn().mockResolvedValue({ configOptions: [] }); - this.resumeSession = vi.fn().mockResolvedValue({ configOptions: [] }); + this.resumeSession = mockResumeSession; + this.prompt = mockPrompt; this.setSessionConfigOption = vi.fn( async ({ value }: { value: string }) => ({ configOptions: [ @@ -70,6 +77,12 @@ const mockAgentRun = vi.hoisted(() => ), ); +const mockResumeFromLog = vi.hoisted(() => vi.fn()); +const mockFormatConversationForResume = vi.hoisted(() => vi.fn()); +const mockHydrateSessionJsonl = vi.hoisted(() => + vi.fn().mockResolvedValue({ hasSession: false }), +); + const mockAgentConstructor = vi.hoisted(() => vi.fn().mockImplementation(function (this: Record) { this.run = mockAgentRun; @@ -103,6 +116,11 @@ vi.mock("@posthog/agent/posthog-api", () => ({ getLlmGatewayUrl: vi.fn(() => "https://gateway.example.com"), })); +vi.mock("@posthog/agent/resume", () => ({ + resumeFromLog: mockResumeFromLog, + formatConversationForResume: mockFormatConversationForResume, +})); + vi.mock("@posthog/agent/gateway-models", () => ({ DEFAULT_GATEWAY_MODEL: "claude-opus-4-8", DEFAULT_CODEX_MODEL: "gpt-5.5", @@ -113,7 +131,7 @@ vi.mock("@posthog/agent/gateway-models", () => ({ })); vi.mock("@posthog/agent/adapters/claude/session/jsonl-hydration", () => ({ - hydrateSessionJsonl: vi.fn().mockResolvedValue(undefined), + hydrateSessionJsonl: mockHydrateSessionJsonl, })); vi.mock("node:fs", async (importOriginal) => { @@ -334,6 +352,148 @@ describe("AgentService", () => { }); }); + describe("reconnect", () => { + it("preserves conversation context when native reconnect fails", async () => { + const apiClient = {}; + mockAgentConstructor.mockImplementationOnce(function ( + this: Record, + ) { + this.run = mockAgentRun; + this.cleanup = vi.fn().mockResolvedValue(undefined); + this.getPosthogAPI = vi.fn(() => apiClient); + this.flushAllLogs = vi.fn().mockResolvedValue(undefined); + }); + mockResumeFromLog.mockResolvedValue({ conversation: [{ role: "user" }] }); + mockFormatConversationForResume.mockReturnValue("User: previous request"); + mockResumeSession.mockRejectedValueOnce(new Error("not found")); + await service.reconnectSession({ + ...baseSessionParams, + adapter: "codex", + sessionId: "old-session", + }); + + await service.prompt("run-1", [{ type: "text", text: "next request" }]); + await service.prompt("run-1", [{ type: "text", text: "later request" }]); + + expect(mockPrompt.mock.calls[0][0].prompt).toEqual([ + expect.objectContaining({ + type: "text", + text: expect.stringContaining("previous request"), + }), + { type: "text", text: "next request" }, + ]); + expect(mockPrompt.mock.calls[1][0].prompt).toEqual([ + { type: "text", text: "later request" }, + ]); + }); + + it("preserves conversation context when reconnect has no session ID", async () => { + mockAgentConstructor.mockImplementationOnce(function ( + this: Record, + ) { + this.run = mockAgentRun; + this.cleanup = vi.fn().mockResolvedValue(undefined); + this.getPosthogAPI = vi.fn(() => ({})); + this.flushAllLogs = vi.fn().mockResolvedValue(undefined); + }); + mockResumeFromLog.mockResolvedValue({ conversation: [{ role: "user" }] }); + mockFormatConversationForResume.mockReturnValue("User: previous request"); + + await service.reconnectSession({ + ...baseSessionParams, + adapter: "codex", + }); + await service.prompt("run-1", [{ type: "text", text: "next request" }]); + + expect(mockPrompt.mock.calls[0][0].prompt[0].text).toContain( + "previous request", + ); + }); + + it("reuses hydrated conversation when Claude resume fails", async () => { + mockAgentConstructor.mockImplementationOnce(function ( + this: Record, + ) { + this.run = mockAgentRun; + this.cleanup = vi.fn().mockResolvedValue(undefined); + this.getPosthogAPI = vi.fn(() => ({})); + this.flushAllLogs = vi.fn().mockResolvedValue(undefined); + }); + mockHydrateSessionJsonl.mockResolvedValueOnce({ + hasSession: true, + conversation: [{ role: "user", content: [] }], + }); + mockFormatConversationForResume.mockReturnValue("User: hydrated request"); + mockResumeSession.mockRejectedValueOnce(new Error("not found")); + + await service.reconnectSession({ + ...baseSessionParams, + adapter: "claude", + sessionId: "old-session", + }); + await service.prompt("run-1", [{ type: "text", text: "next request" }]); + + expect(mockResumeFromLog).not.toHaveBeenCalled(); + expect(mockPrompt.mock.calls[0][0].prompt[0].text).toContain( + "hydrated request", + ); + }); + + it("does not resend hydrated conversation after native resume succeeds", async () => { + mockAgentConstructor.mockImplementationOnce(function ( + this: Record, + ) { + this.run = mockAgentRun; + this.cleanup = vi.fn().mockResolvedValue(undefined); + this.getPosthogAPI = vi.fn(() => ({})); + this.flushAllLogs = vi.fn().mockResolvedValue(undefined); + }); + mockHydrateSessionJsonl.mockResolvedValueOnce({ + hasSession: true, + conversation: [{ role: "user", content: [] }], + }); + mockFormatConversationForResume.mockReturnValue("User: hydrated request"); + + await service.reconnectSession({ + ...baseSessionParams, + adapter: "claude", + sessionId: "old-session", + }); + await service.prompt("run-1", [{ type: "text", text: "next request" }]); + + expect(mockPrompt.mock.calls[0][0].prompt).toEqual([ + { type: "text", text: "next request" }, + ]); + }); + + it("retries recovered context after prompt failure", async () => { + mockAgentConstructor.mockImplementationOnce(function ( + this: Record, + ) { + this.run = mockAgentRun; + this.cleanup = vi.fn().mockResolvedValue(undefined); + this.getPosthogAPI = vi.fn(() => ({})); + this.flushAllLogs = vi.fn().mockResolvedValue(undefined); + }); + mockResumeFromLog.mockResolvedValue({ conversation: [{ role: "user" }] }); + mockFormatConversationForResume.mockReturnValue("User: previous request"); + mockPrompt.mockRejectedValueOnce(new Error("connection lost")); + + await service.reconnectSession({ + ...baseSessionParams, + adapter: "codex", + }); + await expect( + service.prompt("run-1", [{ type: "text", text: "first attempt" }]), + ).rejects.toThrow("connection lost"); + await service.prompt("run-1", [{ type: "text", text: "retry" }]); + + expect(mockPrompt.mock.calls[1][0].prompt[0].text).toContain( + "previous request", + ); + }); + }); + describe("MCP servers", () => { it("marks desktop sessions as local even though they have a taskRunId", async () => { await service.startSession({ diff --git a/packages/workspace-server/src/services/agent/agent.ts b/packages/workspace-server/src/services/agent/agent.ts index 9fa5c7db8a..e9fc690702 100644 --- a/packages/workspace-server/src/services/agent/agent.ts +++ b/packages/workspace-server/src/services/agent/agent.ts @@ -45,6 +45,10 @@ import { wasCreatedByLogin, wasCreatedRecently, } from "@posthog/agent/pr-url-detector"; +import { + formatConversationForResume, + resumeFromLog, +} from "@posthog/agent/resume"; import type * as AgentTypes from "@posthog/agent/types"; import { execGh } from "@posthog/git/gh"; import { getCurrentBranch } from "@posthog/git/queries"; @@ -845,6 +849,8 @@ If a repository IS genuinely required, attach one in this priority order: debug: isDevBuild(), onLog: this.onAgentLog, }); + let fallbackResumeContext: string | undefined; + let hydratedResumeContext: string | undefined; try { const systemPrompt = this.buildSystemPrompt( @@ -1009,6 +1015,13 @@ If a repository IS genuinely required, attach one in this priority order: let configOptions: SessionConfigOption[] | undefined; let agentSessionId: string | undefined; + if (isReconnect && !config.sessionId) { + fallbackResumeContext = await this.loadFallbackResumeContext( + agent, + config, + ); + } + // Imported Claude Code CLI session: the transcript JSONL was copied // into CLAUDE_CONFIG_DIR at import time, so load it directly and let // the adapter replay its history to the client. On failure, fall @@ -1065,7 +1078,7 @@ If a repository IS genuinely required, attach one in this priority order: if (adapter !== "codex") { const posthogAPI = agent.getPosthogAPI(); if (posthogAPI) { - const hasSession = await hydrateSessionJsonl({ + const hydration = await hydrateSessionJsonl({ sessionId: existingSessionId, cwd: repoPath, taskId, @@ -1074,11 +1087,19 @@ If a repository IS genuinely required, attach one in this priority order: posthogAPI, log: this.log, }); - if (!hasSession) { + if (hydration.conversation) { + hydratedResumeContext = this.formatFallbackResumeContext( + hydration.conversation, + ); + } + if (!hydration.hasSession) { this.log.info( "No session JSONL to resume, creating new session instead", { taskId, taskRunId }, ); + fallbackResumeContext ??= + hydratedResumeContext ?? + (await this.loadFallbackResumeContext(agent, config)); config.sessionId = undefined; } } @@ -1169,6 +1190,7 @@ If a repository IS genuinely required, attach one in this priority order: toolInstallations, evaluatedPrUrls: new Set(), prAttachChain: Promise.resolve(), + pendingContext: fallbackResumeContext, }; this.sessions.set(taskRunId, session); @@ -1179,6 +1201,16 @@ If a repository IS genuinely required, attach one in this priority order: } return session; } catch (err) { + if ( + fallbackResumeContext === undefined && + isReconnect && + !isRetry && + !isAuthError(err) + ) { + fallbackResumeContext = + hydratedResumeContext ?? + (await this.loadFallbackResumeContext(agent, config)); + } try { await agent.cleanup(); } catch { @@ -1231,13 +1263,48 @@ If a repository IS genuinely required, attach one in this priority order: sessionId: config.sessionId, }); config.sessionId = undefined; - return this.getOrCreateSession(config, false, false); + const session = await this.getOrCreateSession(config, false, false); + session.pendingContext = fallbackResumeContext; + return session; } if (isReconnect) return null; throw err; } } + private async loadFallbackResumeContext( + agent: Agent, + config: SessionConfig, + ): Promise { + const apiClient = agent.getPosthogAPI(); + if (!apiClient) return undefined; + + try { + const state = await resumeFromLog({ + taskId: config.taskId, + runId: config.taskRunId, + repositoryPath: config.repoPath, + apiClient, + }); + return this.formatFallbackResumeContext(state.conversation); + } catch (err) { + this.log.warn("Failed to restore conversation for fallback session", { + taskId: config.taskId, + taskRunId: config.taskRunId, + error: err instanceof Error ? err.message : String(err), + }); + return undefined; + } + } + + private formatFallbackResumeContext( + conversation: Parameters[0], + ): string | undefined { + const history = formatConversationForResume(conversation); + if (!history) return undefined; + return `You are resuming a previous conversation after the native session could not be restored. Here is the conversation history from the previous session:\n\n${history}\n\nContinue from where you left off when responding to the user's next message.`; + } + private async filterReachableMcpServers< T extends { name: string; @@ -1341,12 +1408,13 @@ If a repository IS genuinely required, attach one in this priority order: // Prepend pending context if present let finalPrompt = prompt; - if (session.pendingContext) { + const pendingContext = session.pendingContext; + if (pendingContext) { this.log.info("Prepending context to prompt", { sessionId }); finalPrompt = [ { type: "text", - text: `_${session.pendingContext}_\n\n`, + text: `_${pendingContext}_\n\n`, _meta: { ui: { hidden: true } }, }, ...prompt, @@ -1360,14 +1428,21 @@ If a repository IS genuinely required, attach one in this priority order: this.sleepService.acquire(sessionId); try { - const result = await session.clientSideConnection.prompt({ - sessionId: getAgentSessionId(session), - prompt: finalPrompt, - }); - return { - stopReason: result.stopReason, - _meta: result._meta as PromptOutput["_meta"], - }; + try { + const result = await session.clientSideConnection.prompt({ + sessionId: getAgentSessionId(session), + prompt: finalPrompt, + }); + return { + stopReason: result.stopReason, + _meta: result._meta as PromptOutput["_meta"], + }; + } catch (err) { + if (pendingContext && session.pendingContext === undefined) { + session.pendingContext = pendingContext; + } + throw err; + } } finally { session.promptPending = false; session.lastActivityAt = Date.now();