diff --git a/apps/code/src/main/services/auth/port-adapters.ts b/apps/code/src/main/services/auth/port-adapters.ts index c354a9e5d5..784d379df1 100644 --- a/apps/code/src/main/services/auth/port-adapters.ts +++ b/apps/code/src/main/services/auth/port-adapters.ts @@ -37,12 +37,12 @@ import { @injectable() export class TokenCipherPortAdapter implements IAuthTokenCipher { - encrypt(plaintext: string): string { - return encrypt(plaintext); + encrypt(plaintext: string): Promise { + return Promise.resolve(encrypt(plaintext)); } - decrypt(encrypted: string): string | null { - return decrypt(encrypted); + decrypt(encrypted: string): Promise { + return Promise.resolve(decrypt(encrypted)); } } diff --git a/packages/agent/src/acp-extensions.ts b/packages/agent/src/acp-extensions.ts index cfa5e7b5be..26783a49f3 100644 --- a/packages/agent/src/acp-extensions.ts +++ b/packages/agent/src/acp-extensions.ts @@ -86,8 +86,22 @@ export const POSTHOG_NOTIFICATIONS = { /** RTK output-compression token savings tallied at the end of a run */ RTK_SAVINGS: "_posthog/rtk_savings", + + /** Latest native Codex goal state, persisted so cold cloud resumes can restore it. */ + CODEX_GOAL: "_posthog/codex_goal", } as const; +export type NativeGoalState = { + objective: string; + status: + | "active" + | "paused" + | "blocked" + | "usageLimited" + | "budgetLimited" + | "complete"; +}; + /** * Custom request methods for PostHog-specific operations that need a response * (request/response, not fire-and-forget). Used with 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 7781fb3234..e91cf2b80c 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 @@ -14,6 +14,9 @@ import { sandboxPolicyFor } from "./session-config"; // Required-field invariants the native codex app-server enforces on each request. const REQUIRED_FIELDS: Record = { + "thread/goal/clear": ["threadId"], + "thread/goal/get": ["threadId"], + "thread/goal/set": ["threadId"], "turn/interrupt": ["threadId", "turnId"], "turn/steer": ["threadId", "input", "expectedTurnId"], }; @@ -43,7 +46,12 @@ function makeStubRpc(responses: Record) { message: `Invalid request: missing field \`${missing}\``, }; } - return (responses[method] ?? {}) as T; + const response = responses[method]; + return ( + typeof response === "function" + ? await response(params) + : (response ?? {}) + ) as T; }, notify() {}, async close() {}, @@ -295,6 +303,274 @@ describe("CodexAppServerAgent", () => { ).toHaveLength(1); }); + it.each([ + { + label: "reads an empty goal", + prompt: "/goal", + method: "thread/goal/get", + response: { goal: null }, + expectedParams: { threadId: "thr_1" }, + expectedText: "No goal set. Usage: `/goal `", + expectedGoal: undefined, + }, + { + label: "reads an active goal", + prompt: "/goal", + method: "thread/goal/get", + response: { goal: { objective: "Ship the fix", status: "active" } }, + expectedParams: { threadId: "thr_1" }, + expectedText: "Goal active: Ship the fix", + expectedGoal: undefined, + }, + { + label: "sets a goal", + prompt: "/goal Ship the fix", + method: "thread/goal/set", + response: { goal: { objective: "Ship the fix", status: "active" } }, + expectedParams: { threadId: "thr_1", objective: "Ship the fix" }, + expectedText: "Goal set: Ship the fix", + expectedGoal: { objective: "Ship the fix", status: "active" }, + }, + { + label: "clears a goal", + prompt: "/goal clear", + method: "thread/goal/clear", + response: { cleared: true }, + expectedParams: { threadId: "thr_1" }, + expectedText: "Goal cleared.", + expectedGoal: null, + }, + { + label: "pauses a goal", + prompt: "/goal pause", + method: "thread/goal/set", + response: { goal: { objective: "Ship the fix", status: "paused" } }, + expectedParams: { threadId: "thr_1", status: "paused" }, + expectedText: "Goal paused: Ship the fix", + expectedGoal: { objective: "Ship the fix", status: "paused" }, + }, + { + label: "resumes a goal", + prompt: "/goal resume", + method: "thread/goal/set", + response: { goal: { objective: "Ship the fix", status: "active" } }, + expectedParams: { threadId: "thr_1", status: "active" }, + expectedText: "Goal resumed: Ship the fix", + expectedGoal: { objective: "Ship the fix", status: "active" }, + }, + ])("$label without starting a model turn", async (testCase) => { + const stub = makeStubRpc({ + "thread/start": { thread: { id: "thr_1" } }, + [testCase.method]: testCase.response, + }); + const { client, sessionUpdates, extNotifications } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/bundle/codex" }, + rpcFactory: stub.factory, + }); + + await agent.newSession({ cwd: "/repo" } as unknown as NewSessionRequest); + const result = await agent.prompt({ + sessionId: "thr_1", + prompt: [{ type: "text", text: testCase.prompt }], + } as unknown as PromptRequest); + + expect(result.stopReason).toBe("end_turn"); + expect(stub.requests).toContainEqual({ + method: testCase.method, + params: testCase.expectedParams, + }); + expect( + stub.requests.some((request) => request.method === "turn/start"), + ).toBe(false); + expect(sessionUpdates).toContainEqual({ + sessionId: "thr_1", + update: { + sessionUpdate: "agent_message_chunk", + content: { type: "text", text: testCase.expectedText }, + }, + }); + if (testCase.expectedGoal !== undefined) { + expect(extNotifications).toContainEqual({ + method: "_posthog/codex_goal", + params: { goal: testCase.expectedGoal }, + }); + } + }); + + it("handles a goal command wrapped in hidden cold-resume context", async () => { + const stub = makeStubRpc({ + "thread/start": { thread: { id: "thr_1" } }, + "thread/goal/get": { + goal: { objective: "Ship the fix", status: "paused" }, + }, + }); + const { client, sessionUpdates } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/bundle/codex" }, + rpcFactory: stub.factory, + }); + await agent.newSession({ cwd: "/repo" } as unknown as NewSessionRequest); + + await agent.prompt({ + sessionId: "thr_1", + prompt: [ + { + type: "text", + text: "Previous conversation context", + _meta: { ui: { hidden: true } }, + }, + { type: "text", text: "/goal" }, + { + type: "text", + text: "Respond to the user above", + _meta: { ui: { hidden: true } }, + }, + ], + } as unknown as PromptRequest); + + expect(stub.requests).toContainEqual({ + method: "thread/goal/get", + params: { threadId: "thr_1" }, + }); + expect( + stub.requests.some((request) => request.method === "turn/start"), + ).toBe(false); + expect(sessionUpdates).toContainEqual({ + sessionId: "thr_1", + update: { + sessionUpdate: "user_message_chunk", + content: { type: "text", text: "/goal" }, + }, + }); + }); + + it("restores a persisted goal when starting a replacement thread", async () => { + const restoredGoal = { objective: "Ship the fix", status: "paused" }; + const stub = makeStubRpc({ + "thread/start": { thread: { id: "thr_1" } }, + "thread/goal/set": { goal: restoredGoal }, + }); + const { client, extNotifications } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/bundle/codex" }, + rpcFactory: stub.factory, + }); + + await agent.newSession({ + cwd: "/repo", + _meta: { nativeGoal: restoredGoal }, + } as unknown as NewSessionRequest); + + expect(stub.requests).toContainEqual({ + method: "thread/goal/set", + params: { threadId: "thr_1", ...restoredGoal }, + }); + expect(extNotifications).toContainEqual({ + method: "_posthog/codex_goal", + params: { goal: restoredGoal }, + }); + }); + + it("interrupts a native goal turn that was already queued when paused", async () => { + const stub = makeStubRpc({ + "thread/start": { thread: { id: "thr_1" } }, + "thread/goal/set": { + goal: { objective: "Ship the fix", status: "paused" }, + }, + }); + const { client } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/bundle/codex" }, + rpcFactory: stub.factory, + }); + await agent.newSession({ cwd: "/repo" } as unknown as NewSessionRequest); + + await agent.prompt({ + sessionId: "thr_1", + prompt: [{ type: "text", text: "/goal pause" }], + } as unknown as PromptRequest); + stub.emit("turn/started", { turn: { id: "goal_tick_1" } }); + await Promise.resolve(); + + expect(stub.requests).toContainEqual({ + method: "turn/interrupt", + params: { threadId: "thr_1", turnId: "goal_tick_1" }, + }); + }); + + it("interrupts a native goal turn that started before it was paused", async () => { + const stub = makeStubRpc({ + "thread/start": { thread: { id: "thr_1" } }, + "thread/goal/set": { + goal: { objective: "Ship the fix", status: "paused" }, + }, + }); + const { client } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/bundle/codex" }, + rpcFactory: stub.factory, + }); + await agent.newSession({ cwd: "/repo" } as unknown as NewSessionRequest); + stub.emit("turn/started", { turn: { id: "goal_tick_1" } }); + + await agent.prompt({ + sessionId: "thr_1", + prompt: [{ type: "text", text: "/goal pause" }], + } as unknown as PromptRequest); + + expect(stub.requests).toContainEqual({ + method: "turn/interrupt", + params: { threadId: "thr_1", turnId: "goal_tick_1" }, + }); + }); + + it("retries queued goal cancellation after an interrupt failure", async () => { + let interruptAttempts = 0; + const stub = makeStubRpc({ + "thread/start": { thread: { id: "thr_1" } }, + "thread/goal/set": { + goal: { objective: "Ship the fix", status: "paused" }, + }, + "turn/interrupt": () => { + interruptAttempts++; + if (interruptAttempts === 1) { + throw new Error("interrupt failed"); + } + return {}; + }, + }); + const { client } = makeFakeClient(); + const agent = new CodexAppServerAgent(client, { + processOptions: { binaryPath: "/bundle/codex" }, + rpcFactory: stub.factory, + }); + await agent.newSession({ cwd: "/repo" } as unknown as NewSessionRequest); + + await agent.prompt({ + sessionId: "thr_1", + prompt: [{ type: "text", text: "/goal pause" }], + } as unknown as PromptRequest); + stub.emit("turn/started", { turn: { id: "goal_tick_1" } }); + await Promise.resolve(); + await Promise.resolve(); + stub.emit("turn/started", { turn: { id: "goal_tick_2" } }); + await Promise.resolve(); + + expect( + stub.requests.filter((request) => request.method === "turn/interrupt"), + ).toEqual([ + { + method: "turn/interrupt", + params: { threadId: "thr_1", turnId: "goal_tick_1" }, + }, + { + method: "turn/interrupt", + params: { threadId: "thr_1", turnId: "goal_tick_2" }, + }, + ]); + }); + it("includes buffered command output when completion omits aggregatedOutput", async () => { const stub = makeStubRpc({ initialize: {}, @@ -1713,6 +1989,7 @@ describe("CodexAppServerAgent", () => { { skills: [ { name: "deploy", description: "Deploy", enabled: true }, + { name: "goal", description: "Duplicate", enabled: true }, { name: "danger", description: "Disabled", enabled: false }, ], }, @@ -1731,7 +2008,14 @@ describe("CodexAppServerAgent", () => { (u: any) => u.update?.sessionUpdate === "available_commands_update", ) as any )?.update?.availableCommands; - expect(cmds.map((c: { name: string }) => c.name)).toEqual(["deploy"]); + expect(cmds).toEqual([ + { + name: "goal", + description: "Set or view the goal for a long-running task", + input: { hint: "[|clear|pause|resume]" }, + }, + { name: "deploy", description: "Deploy" }, + ]); }); it("emits _posthog/sdk_session when a taskRunId is present", async () => { 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 406ec589a3..7ed57f3339 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 @@ -20,7 +20,10 @@ import type { StopReason, } from "@agentclientprotocol/sdk"; import { mcpToolKey, posthogToolMeta } from "@posthog/shared"; -import { POSTHOG_NOTIFICATIONS } from "../../acp-extensions"; +import { + type NativeGoalState, + POSTHOG_NOTIFICATIONS, +} from "../../acp-extensions"; import { DEFAULT_CODEX_MODEL } from "../../gateway-models"; import type { ProcessSpawnedCallback } from "../../types"; import { ALLOW_BYPASS } from "../../utils/common"; @@ -87,6 +90,7 @@ type AppServerSessionMeta = { channelMode?: boolean; spokenNarration?: boolean; baseBranch?: string; + nativeGoal?: NativeGoalState; }; /** The subset of codex's `Thread` the adapter reads: id + persisted `turns` for history replay. */ @@ -95,6 +99,64 @@ type AppServerThread = { turns?: Array<{ items?: Parameters[1][] }>; }; +type ThreadGoal = { + objective: string; + status: NativeGoalState["status"]; +}; + +type GoalCommand = + | { kind: "get" } + | { kind: "clear" } + | { kind: "pause" } + | { kind: "resume" } + | { kind: "set"; objective: string }; + +type CodexSkill = { + name?: string; + description?: string; + enabled?: boolean; +}; + +const GOAL_COMMAND = { + name: "goal", + description: "Set or view the goal for a long-running task", + input: { hint: "[|clear|pause|resume]" }, +}; + +function isHiddenPromptBlock(block: PromptRequest["prompt"][number]): boolean { + const meta = block._meta as { ui?: { hidden?: boolean } } | undefined; + return meta?.ui?.hidden === true; +} + +function visiblePromptBlocks( + prompt: PromptRequest["prompt"], +): PromptRequest["prompt"] { + return prompt.filter((block) => !isHiddenPromptBlock(block)); +} + +function parseGoalCommand(prompt: PromptRequest["prompt"]): GoalCommand | null { + const visible = visiblePromptBlocks(prompt); + if (visible.some((block) => block.type !== "text")) return null; + const text = visible + .map((block) => (block.type === "text" ? block.text : "")) + .join("\n") + .trim(); + const match = text.match(/^\/goal(?:\s+([\s\S]*))?$/); + if (!match) return null; + const argument = match[1]?.trim(); + if (!argument) return { kind: "get" }; + switch (argument.toLowerCase()) { + case "clear": + return { kind: "clear" }; + case "pause": + return { kind: "pause" }; + case "resume": + return { kind: "resume" }; + default: + return { kind: "set", objective: argument }; + } +} + // The native app-server owns its config; BaseAcpAgent only calls dispose() on this. class NoopSettingsManager implements BaseSettingsManager { constructor(private cwd: string) {} @@ -158,6 +220,10 @@ export class CodexAppServerAgent extends BaseAcpAgent { private readonly mcp = new McpManager(); private readonly turns = new TurnController(); private readonly usage = new UsageTracker(); + /** Pause/clear can race a goal continuation already queued by app-server. */ + private cancelNextGoalTurn = false; + /** Native goal ticks start outside prompt(), so TurnController does not own them. */ + private nativeGoalTurnId?: string; constructor( client: AgentSideConnection, @@ -363,6 +429,8 @@ export class CodexAppServerAgent extends BaseAcpAgent { additionalDirectories?: string[]; }, ): Promise<{ threadId: string; thread: AppServerThread | undefined }> { + this.cancelNextGoalTurn = false; + this.nativeGoalTurnId = undefined; this.jsonSchema = params.meta?.jsonSchema ?? undefined; this.taskRunId = params.meta?.taskRunId; this.environment = params.meta?.environment; @@ -418,6 +486,9 @@ export class CodexAppServerAgent extends BaseAcpAgent { } this.threadId = threadId; this.sessionId = threadId; + if (method === APP_SERVER_METHODS.THREAD_START && params.meta?.nativeGoal) { + await this.restoreGoal(params.meta.nativeGoal); + } await this.loadModelConfig(); this.emitConfigOptions(); await this.emitAvailableCommands(); @@ -519,17 +590,29 @@ export class CodexAppServerAgent extends BaseAcpAgent { /** skills/list → available_commands_update so the slash-command menu fills. */ private async emitAvailableCommands(): Promise { if (!this.sessionId) return; - let commands: Array<{ name: string; description: string }> = []; + let commands: Array<{ + name: string; + description: string; + input?: { hint: string }; + }> = [GOAL_COMMAND]; try { - const res = await this.rpc.request<{ data?: Array<{ skills?: any[] }> }>( - APP_SERVER_METHODS.SKILLS_LIST, - {}, - ); - commands = (res?.data ?? []) + const res = await this.rpc.request<{ + data?: Array<{ skills?: CodexSkill[] }>; + }>(APP_SERVER_METHODS.SKILLS_LIST, {}); + const skills = (res?.data ?? []) .flatMap((entry) => entry?.skills ?? []) // Drop explicitly-disabled skills; lenient `!== false` so a malformed payload still shows. - .filter((s) => s?.name && s?.enabled !== false) - .map((s: any) => ({ name: s.name, description: s.description ?? "" })); + .filter( + (skill): skill is CodexSkill & { name: string } => + !!skill.name && + skill.name !== GOAL_COMMAND.name && + skill.enabled !== false, + ) + .map((skill) => ({ + name: skill.name, + description: skill.description ?? "", + })); + commands = [GOAL_COMMAND, ...skills]; } catch (err) { this.logger.warn("skills/list failed", { error: String(err) }); } @@ -548,6 +631,13 @@ export class CodexAppServerAgent extends BaseAcpAgent { if (!this.threadId) { throw new Error("prompt() called before newSession()"); } + const goalCommand = parseGoalCommand(params.prompt); + if (goalCommand) { + this.broadcastUserInput(visiblePromptBlocks(params.prompt)); + await this.handleGoalCommand(goalCommand); + return { stopReason: "end_turn" }; + } + this.cancelNextGoalTurn = false; // Reopen the notification gate (a prior interrupt may have left session.cancelled set). this.session.cancelled = false; // A new prompt while the plan handoff awaits approval implicitly declines it: @@ -643,6 +733,117 @@ export class CodexAppServerAgent extends BaseAcpAgent { return { stopReason: await this.maybeOfferPlanImplementation(stopReason) }; } + private async handleGoalCommand(command: GoalCommand): Promise { + if (!this.threadId) return; + if (command.kind === "clear") { + this.cancelNextGoalTurn = true; + const result = await this.rpc.request<{ cleared?: boolean }>( + APP_SERVER_METHODS.THREAD_GOAL_CLEAR, + { threadId: this.threadId }, + ); + await this.emitGoalState(null); + await this.cancelRunningGoalTurn(); + this.broadcastAgentText( + result.cleared ? "Goal cleared." : "No goal was set.", + ); + return; + } + + if (command.kind === "get") { + const result = await this.rpc.request<{ goal?: ThreadGoal | null }>( + APP_SERVER_METHODS.THREAD_GOAL_GET, + { threadId: this.threadId }, + ); + this.broadcastAgentText( + result.goal + ? `Goal ${result.goal.status}: ${result.goal.objective}` + : "No goal set. Usage: `/goal `", + ); + return; + } + + if (command.kind === "pause") { + this.cancelNextGoalTurn = true; + } + const params = + command.kind === "set" + ? { threadId: this.threadId, objective: command.objective } + : { + threadId: this.threadId, + status: command.kind === "pause" ? "paused" : "active", + }; + const result = await this.rpc.request<{ goal: ThreadGoal }>( + APP_SERVER_METHODS.THREAD_GOAL_SET, + params, + ); + await this.emitGoalState(result.goal); + if (command.kind === "pause") { + await this.cancelRunningGoalTurn(); + } + const prefix = + command.kind === "set" + ? "Goal set" + : command.kind === "pause" + ? "Goal paused" + : "Goal resumed"; + this.broadcastAgentText(`${prefix}: ${result.goal.objective}`); + } + + private async restoreGoal(goal: NativeGoalState): Promise { + if (!this.threadId) return; + const result = await this.rpc.request<{ goal: ThreadGoal }>( + APP_SERVER_METHODS.THREAD_GOAL_SET, + { + threadId: this.threadId, + objective: goal.objective, + status: goal.status, + }, + ); + await this.emitGoalState(result.goal); + } + + private async emitGoalState(goal: NativeGoalState | null): Promise { + await this.client + .extNotification(POSTHOG_NOTIFICATIONS.CODEX_GOAL, { goal }) + .catch((error) => + this.logger.warn("Failed to persist Codex goal state", error), + ); + } + + private async cancelRunningGoalTurn(): Promise { + if (this.turns.isRunning) { + this.cancelNextGoalTurn = false; + await this.interrupt(); + return; + } + if (this.nativeGoalTurnId) { + await this.interruptNativeGoalTurn(this.nativeGoalTurnId); + } + } + + private interruptQueuedGoalTurn(turnId: string | undefined): void { + if (!this.cancelNextGoalTurn || !this.threadId || !turnId) return; + void this.interruptNativeGoalTurn(turnId); + } + + private async interruptNativeGoalTurn(turnId: string): Promise { + if (!this.threadId) return; + await this.rpc + .request(APP_SERVER_METHODS.TURN_INTERRUPT, { + threadId: this.threadId, + turnId, + }) + .then(() => { + this.cancelNextGoalTurn = false; + if (this.nativeGoalTurnId === turnId) { + this.nativeGoalTurnId = undefined; + } + }) + .catch((error) => + this.logger.warn("Native goal turn interrupt failed", error), + ); + } + /** Start one codex turn and await its completion. */ private async runTurn(input: CodexUserInput[]): Promise { this.lastAgentMessage = ""; @@ -922,6 +1123,7 @@ export class CodexAppServerAgent extends BaseAcpAgent { async closeSession(): Promise { this.commandOutputs.clear(); + this.nativeGoalTurnId = undefined; this.session.abortController.abort(); this.session.cancelled = true; this.planHandoffCancel?.(); @@ -971,7 +1173,12 @@ export class CodexAppServerAgent extends BaseAcpAgent { if (method === APP_SERVER_NOTIFICATIONS.TURN_STARTED) { // Capture the active turn id (steer precondition / interrupt target). - this.turns.onStarted((params as { turn?: { id?: string } })?.turn?.id); + const turnId = (params as { turn?: { id?: string } })?.turn?.id; + if (!this.turns.isPending && turnId) { + this.nativeGoalTurnId = turnId; + } + this.turns.onStarted(turnId); + this.interruptQueuedGoalTurn(turnId); } // codex auto-compaction surfaces as a contextCompaction item: item/started → in progress, @@ -1014,6 +1221,9 @@ export class CodexAppServerAgent extends BaseAcpAgent { this.commandOutputs.clear(); const turn = (params as { turn?: { id?: string; status?: string } }) ?.turn; + if (turn?.id === this.nativeGoalTurnId) { + this.nativeGoalTurnId = undefined; + } // Drop the late completion of an already-interrupted turn (else it cancels the follow-up). if (this.turns.shouldDropCompletion(turn?.id)) return; void this.finalizeTurn(mapTurnStopReason(turn?.status)); diff --git a/packages/agent/src/adapters/codex-app-server/protocol.ts b/packages/agent/src/adapters/codex-app-server/protocol.ts index e50304b0d6..37c61a75e9 100644 --- a/packages/agent/src/adapters/codex-app-server/protocol.ts +++ b/packages/agent/src/adapters/codex-app-server/protocol.ts @@ -15,6 +15,9 @@ export const APP_SERVER_METHODS = { TURN_INTERRUPT: "turn/interrupt", MODEL_LIST: "model/list", SKILLS_LIST: "skills/list", + THREAD_GOAL_SET: "thread/goal/set", + THREAD_GOAL_GET: "thread/goal/get", + THREAD_GOAL_CLEAR: "thread/goal/clear", THREAD_LIST: "thread/list", } as const; diff --git a/packages/agent/src/resume.ts b/packages/agent/src/resume.ts index fe4c600f29..fad960cf62 100644 --- a/packages/agent/src/resume.ts +++ b/packages/agent/src/resume.ts @@ -16,6 +16,7 @@ */ import type { ContentBlock } from "@agentclientprotocol/sdk"; +import type { NativeGoalState } from "./acp-extensions"; import { selectRecentTurns } from "./adapters/claude/session/jsonl-hydration"; import type { PostHogAPIClient } from "./posthog-api"; import { ResumeSaga } from "./sagas/resume-saga"; @@ -29,6 +30,7 @@ export interface ResumeState { lastDevice?: DeviceInfo; logEntryCount: number; sessionId: string | null; + nativeGoal?: NativeGoalState | null; } export interface ConversationTurn { @@ -95,6 +97,7 @@ export async function resumeFromLog( lastDevice: result.data.lastDevice, logEntryCount: result.data.logEntryCount, sessionId: result.data.sessionId, + nativeGoal: result.data.nativeGoal, }; } diff --git a/packages/agent/src/sagas/resume-saga.test.ts b/packages/agent/src/sagas/resume-saga.test.ts index a26e6c7824..a3919c6de7 100644 --- a/packages/agent/src/sagas/resume-saga.test.ts +++ b/packages/agent/src/sagas/resume-saga.test.ts @@ -617,6 +617,69 @@ describe("ResumeSaga", () => { }); }); + describe("Codex goal state", () => { + it.each([ + { + name: "restores the latest persisted goal", + entries: [ + createNotification(POSTHOG_NOTIFICATIONS.CODEX_GOAL, { + goal: { objective: "Old goal", status: "active" }, + }), + createNotification(`_${POSTHOG_NOTIFICATIONS.CODEX_GOAL}`, { + goal: { objective: "Ship the fix", status: "paused" }, + }), + ], + expected: { objective: "Ship the fix", status: "paused" }, + }, + { + name: "preserves an explicit goal clear", + entries: [ + createNotification(POSTHOG_NOTIFICATIONS.CODEX_GOAL, { + goal: { objective: "Ship the fix", status: "active" }, + }), + createNotification(POSTHOG_NOTIFICATIONS.CODEX_GOAL, { goal: null }), + ], + expected: null, + }, + { + name: "skips malformed newer goal notifications", + entries: [ + createNotification(POSTHOG_NOTIFICATIONS.CODEX_GOAL, { + goal: { objective: "Ship the fix", status: "paused" }, + }), + createNotification(POSTHOG_NOTIFICATIONS.CODEX_GOAL, { + goal: { objective: "Invalid goal", status: "unknown" }, + }), + createNotification(POSTHOG_NOTIFICATIONS.CODEX_GOAL, {}), + ], + expected: { objective: "Ship the fix", status: "paused" }, + }, + { + name: "leaves goal state undefined when no notification exists", + entries: [createUserMessage("Hello")], + expected: undefined, + }, + ])("$name", async ({ entries, expected }) => { + (mockApiClient.getTaskRun as ReturnType).mockResolvedValue( + createTaskRun(), + ); + ( + mockApiClient.fetchTaskRunLogs as ReturnType + ).mockResolvedValue(entries); + + const result = await new ResumeSaga(mockLogger).run({ + taskId: "task-1", + runId: "run-1", + repositoryPath: repo.path, + apiClient: mockApiClient, + }); + + expect(result.success).toBe(true); + if (!result.success) return; + expect(result.data.nativeGoal).toEqual(expected); + }); + }); + describe("log entry count", () => { it("reports correct log entry count", async () => { (mockApiClient.getTaskRun as ReturnType).mockResolvedValue( diff --git a/packages/agent/src/sagas/resume-saga.ts b/packages/agent/src/sagas/resume-saga.ts index 09ad7e45e7..d18309675f 100644 --- a/packages/agent/src/sagas/resume-saga.ts +++ b/packages/agent/src/sagas/resume-saga.ts @@ -1,6 +1,6 @@ import type { ContentBlock } from "@agentclientprotocol/sdk"; import { Saga } from "@posthog/shared"; -import { POSTHOG_NOTIFICATIONS } from "../acp-extensions"; +import { type NativeGoalState, POSTHOG_NOTIFICATIONS } from "../acp-extensions"; import type { PostHogAPIClient } from "../posthog-api"; import type { DeviceInfo, @@ -37,6 +37,7 @@ export interface ResumeOutput { lastDevice?: DeviceInfo; logEntryCount: number; sessionId: string | null; + nativeGoal?: NativeGoalState | null; } export class ResumeSaga extends Saga { @@ -91,6 +92,9 @@ export class ResumeSaga extends Saga { const sessionId = await this.readOnlyStep("find_session_id", () => Promise.resolve(this.findSessionId(entries)), ); + const nativeGoal = await this.readOnlyStep("find_native_goal", () => + Promise.resolve(this.findNativeGoal(entries)), + ); this.log.info("Resume state rebuilt", { turns: conversation.length, @@ -106,6 +110,7 @@ export class ResumeSaga extends Saga { lastDevice, logEntryCount: entries.length, sessionId, + nativeGoal, }; } @@ -116,9 +121,47 @@ export class ResumeSaga extends Saga { interrupted: false, logEntryCount: 0, sessionId: null, + nativeGoal: undefined, }; } + private findNativeGoal( + entries: StoredNotification[], + ): NativeGoalState | null | undefined { + const statuses = new Set([ + "active", + "paused", + "blocked", + "usageLimited", + "budgetLimited", + "complete", + ]); + const methods = new Set([ + POSTHOG_NOTIFICATIONS.CODEX_GOAL, + `_${POSTHOG_NOTIFICATIONS.CODEX_GOAL}`, + ]); + for (let index = entries.length - 1; index >= 0; index--) { + const notification = entries[index].notification; + if (!methods.has(notification?.method ?? "")) continue; + const goal = (notification?.params as { goal?: unknown } | undefined) + ?.goal; + if (goal === null) return null; + if (!goal || typeof goal !== "object") continue; + const value = goal as Record; + if ( + typeof value.objective === "string" && + typeof value.status === "string" && + statuses.has(value.status as NativeGoalState["status"]) + ) { + return { + objective: value.objective, + status: value.status as NativeGoalState["status"], + }; + } + } + return undefined; + } + private findSessionId(entries: StoredNotification[]): string | null { const runStarted = POSTHOG_NOTIFICATIONS.RUN_STARTED; for (let i = entries.length - 1; i >= 0; i--) { diff --git a/packages/agent/src/server/agent-server.test.ts b/packages/agent/src/server/agent-server.test.ts index bc43a7858b..7edaffac73 100644 --- a/packages/agent/src/server/agent-server.test.ts +++ b/packages/agent/src/server/agent-server.test.ts @@ -242,6 +242,10 @@ interface TestableServer { buildClaudeCodeSessionMeta( runtimeAdapter: Adapter, ): { claudeCode: { options: Record } } | undefined; + resumeState: ResumeState | null; + getNativeGoalForFreshSession( + runtimeAdapter: Adapter, + ): ResumeState["nativeGoal"]; } interface NativeResumeTestServer { @@ -436,6 +440,37 @@ describe("AgentServer HTTP Mode", () => { ); }; + it("replays ACP notifications emitted before cloud session assignment", () => { + const testServer = createServer() as unknown as { + session: { sseController: null } | null; + pendingEvents: Record[]; + preSessionEvents: Record[]; + handleAcpTransportMessage(message: unknown): void; + flushPreSessionEvents(): void; + }; + const message = { + method: "session/update", + params: { + update: { + sessionUpdate: "available_commands_update", + availableCommands: [{ name: "goal" }], + }, + }, + }; + + testServer.handleAcpTransportMessage(message); + expect(testServer.preSessionEvents).toHaveLength(1); + + testServer.session = { sseController: null }; + testServer.flushPreSessionEvents(); + + expect(testServer.preSessionEvents).toHaveLength(0); + expect(testServer.pendingEvents).toContainEqual( + expect.objectContaining({ notification: message }), + ); + testServer.session = null; + }); + describe("GET /health", () => { it("returns ok status with active session", async () => { await createServer().start(); @@ -2345,6 +2380,22 @@ describe("AgentServer HTTP Mode", () => { }); describe("native resume", () => { + it("restores persisted Codex goals only for fresh Codex sessions", () => { + const s = createServer() as unknown as TestableServer; + const goal = { objective: "Ship the fix", status: "paused" as const }; + s.resumeState = { + conversation: [], + latestGitCheckpoint: null, + interrupted: false, + logEntryCount: 1, + sessionId: "prior-session", + nativeGoal: goal, + }; + + expect(s.getNativeGoalForFreshSession("codex")).toEqual(goal); + expect(s.getNativeGoalForFreshSession("claude")).toBeUndefined(); + }); + it.each([ { retryOutcome: "succeeds", retryFails: false }, { retryOutcome: "fails", retryFails: true }, diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index 5210b7540f..01e69aa342 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -368,6 +368,8 @@ export class AgentServer { // causing a second session to be created and duplicate Slack messages to be sent. private initializationPromise: Promise | null = null; private pendingEvents: Record[] = []; + /** ACP notifications emitted by newSession/resumeSession before this.session is assigned. */ + private preSessionEvents: Record[] = []; private deliveredMessageIds = new Set(); private pendingCompactContinuationMessageIds = new Set(); private pendingPermissions = new Map< @@ -800,6 +802,13 @@ export class AgentServer { return { sessionId: priorSessionId, warm }; } + private getNativeGoalForFreshSession( + runtimeAdapter: Adapter, + ): ResumeState["nativeGoal"] { + if (runtimeAdapter !== "codex") return undefined; + return this.resumeState?.nativeGoal; + } + async stop(): Promise { this.logger.debug("Stopping agent server..."); @@ -1220,6 +1229,7 @@ export class AgentServer { this.resumeState = null; this.nativeResume = null; + this.preSessionEvents = []; this.prewarmedRun = false; this.warmAutoPublishResolved = false; @@ -1427,6 +1437,9 @@ export class AgentServer { sessionCwd, initialPermissionMode, ); + let effectiveSessionMeta: typeof sessionMeta & { + nativeGoal?: NonNullable; + } = sessionMeta; let acpSessionId: string | null = null; if (nativeResume) { @@ -1435,7 +1448,7 @@ export class AgentServer { sessionId: nativeResume.sessionId, cwd: sessionCwd, mcpServers: this.config.mcpServers ?? [], - _meta: { ...sessionMeta, sessionId: nativeResume.sessionId }, + _meta: { ...effectiveSessionMeta, sessionId: nativeResume.sessionId }, }); acpSessionId = nativeResume.sessionId; this.nativeResume = nativeResume; @@ -1454,10 +1467,15 @@ export class AgentServer { } } if (!acpSessionId) { + const restoredNativeGoal = + this.getNativeGoalForFreshSession(runtimeAdapter); + effectiveSessionMeta = restoredNativeGoal + ? { ...sessionMeta, nativeGoal: restoredNativeGoal } + : sessionMeta; const sessionResponse = await clientConnection.newSession({ cwd: sessionCwd, mcpServers: this.config.mcpServers ?? [], - _meta: sessionMeta, + _meta: effectiveSessionMeta, }); acpSessionId = sessionResponse.sessionId; this.logger.debug("ACP session created", { @@ -1480,8 +1498,9 @@ export class AgentServer { permissionMode: initialPermissionMode, hasDesktopConnected: sseController !== null, pendingHandoffGitState: undefined, - sessionMeta, + sessionMeta: effectiveSessionMeta, }; + this.flushPreSessionEvents(); this.logger = new Logger({ debug: true, @@ -3977,6 +3996,7 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} } this.pendingEvents = []; + this.preSessionEvents = []; this.lastReportedBranch = null; // Run usage is per run: a later session on this instance (e.g. a resume // with a different run_id) must not inherit the previous run's totals. @@ -4098,11 +4118,16 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} } this.adapterEmittedTurnComplete = true; } - this.broadcastEvent({ + const event = { type: "notification", timestamp: new Date().toISOString(), notification: message, - }); + }; + if (!this.session) { + this.preSessionEvents.push(event); + return; + } + this.broadcastEvent(event); } private broadcastTurnComplete(stopReason: string): void { @@ -4145,6 +4170,15 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions} } } + private flushPreSessionEvents(): void { + if (!this.session || this.preSessionEvents.length === 0) return; + const events = this.preSessionEvents; + this.preSessionEvents = []; + for (const event of events) { + this.broadcastEvent(event); + } + } + private replayPendingEvents(): void { if (!this.session?.sseController || this.pendingEvents.length === 0) return; const events = this.pendingEvents; diff --git a/packages/core/src/auth/auth.test.ts b/packages/core/src/auth/auth.test.ts index 00bd2ee104..38be8ae05e 100644 --- a/packages/core/src/auth/auth.test.ts +++ b/packages/core/src/auth/auth.test.ts @@ -64,8 +64,8 @@ function createPreferencePort(): IAuthPreferenceStore { } const identityCipher: IAuthTokenCipher = { - encrypt: (plaintext) => plaintext, - decrypt: (encrypted) => encrypted, + encrypt: (plaintext) => Promise.resolve(plaintext), + decrypt: (encrypted) => Promise.resolve(encrypted), }; const mockLogger: RootLogger = { @@ -413,6 +413,52 @@ describe("AuthService", () => { } }); + it("dedupes concurrent refreshes into a single token request", async () => { + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ + accessToken: "initial-access-token", + refreshToken: "initial-refresh-token", + }), + ); + stubAuthFetch(); + + // Keep the refresh in flight so both callers overlap; if the refresh + // isn't deduped, the rotating refresh token is spent twice. + let resolveRefresh!: (value: unknown) => void; + oauthFlow.refreshToken.mockReturnValue( + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + + await service.initialize(); + await service.login("us"); + + // Two forced refreshes fired in the same tick. The dedup guard must close + // synchronously — resolving the stored session now awaits decryption, so + // if that await sat between the guard and the refreshPromise assignment, + // both callers would slip through and fire their own request. + oauthFlow.refreshToken.mockClear(); + const both = Promise.all([ + service.refreshAccessToken(), + service.refreshAccessToken(), + ]); + + await new Promise((r) => setTimeout(r, 0)); + expect(oauthFlow.refreshToken).toHaveBeenCalledTimes(1); + + resolveRefresh( + mockTokenResponse({ + accessToken: "refreshed-access-token", + refreshToken: "refreshed-refresh-token", + }), + ); + const [a, b] = await both; + expect(a.accessToken).toBe("refreshed-access-token"); + expect(b.accessToken).toBe("refreshed-access-token"); + expect(oauthFlow.refreshToken).toHaveBeenCalledTimes(1); + }); + it("forces a token refresh when explicitly requested", async () => { oauthFlow.startFlow.mockResolvedValue( mockTokenResponse({ @@ -500,6 +546,141 @@ describe("AuthService", () => { }); }); + it("keeps the prior selection committed when persisting a new selection fails", async () => { + const orgs = { + "org-1": { + name: "Org 1", + projects: [ + { id: 42, name: "Project 42" }, + { id: 84, name: "Project 84" }, + ], + }, + }; + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ + accessToken: "initial-access-token", + refreshToken: "initial-refresh-token", + }), + ); + stubAuthFetch({ orgs }); + + // Encrypts fine during login, then rejects so the selectProject persist + // fails mid-flow (mirrors the browser cipher's Web Crypto rejecting). + let failEncrypt = false; + const flakyCipher: IAuthTokenCipher = { + encrypt: (plaintext) => + failEncrypt + ? Promise.reject(new Error("encryption unavailable")) + : Promise.resolve(plaintext), + decrypt: (encrypted) => Promise.resolve(encrypted), + }; + service = new AuthService( + preferencePort, + sessionPort, + oauthFlow as unknown as IAuthOAuthFlowService, + connectivity, + flakyCipher, + mockPowerManager as unknown as IPowerManager, + mockLogger, + null, + ); + + // Initialize once while the session store is empty so login/selectProject + // don't later trigger the stored-session restore path. + service.init(); + await service.initialize(); + + await service.login("us"); + const priorProjectId = service.getState().currentProjectId; + expect(priorProjectId).toBe(42); + + failEncrypt = true; + await expect(service.selectProject(84)).rejects.toThrow( + "encryption unavailable", + ); + + // A failed persist must not commit the new project anywhere: published + // state, the stored session, and the saved preference all stay on the + // prior selection (the preference is the tell — the buggy order saved it + // to 84 before the encrypt rejected). + expect(service.getState().currentProjectId).toBe(priorProjectId); + expect(sessionPort.getCurrent()?.selectedProjectId).toBe(priorProjectId); + expect(preferencePort.get("user-1", "us")?.lastSelectedProjectId).toBe( + priorProjectId, + ); + }); + + it("keeps overlapping selections consistent so the latest one wins", async () => { + const orgs = { + "org-1": { + name: "Org 1", + projects: [ + { id: 42, name: "Project 42" }, + { id: 84, name: "Project 84" }, + { id: 99, name: "Project 99" }, + ], + }, + }; + oauthFlow.startFlow.mockResolvedValue( + mockTokenResponse({ + accessToken: "initial-access-token", + refreshToken: "initial-refresh-token", + }), + ); + stubAuthFetch({ orgs }); + + // Encryption is gated so overlapping selection commits stay in flight and + // can be resolved out of order below. + let deferEncrypt = false; + const pending: Array<() => void> = []; + const gatedCipher: IAuthTokenCipher = { + encrypt: (plaintext) => + deferEncrypt + ? new Promise((resolve) => { + pending.push(() => resolve(plaintext)); + }) + : Promise.resolve(plaintext), + decrypt: (encrypted) => Promise.resolve(encrypted), + }; + service = new AuthService( + preferencePort, + sessionPort, + oauthFlow as unknown as IAuthOAuthFlowService, + connectivity, + gatedCipher, + mockPowerManager as unknown as IPowerManager, + mockLogger, + null, + ); + service.init(); + await service.initialize(); + await service.login("us"); + expect(service.getState().currentProjectId).toBe(42); + + deferEncrypt = true; + // Two overlapping selections: 84 first, then 99 (the latest intent). + const first = service.selectProject(84); + const second = service.selectProject(99); + + // Drain pending encryptions newest-first each round to surface any + // out-of-order completion, until both commits settle. Serialized commits + // only expose one pending encryption at a time; unserialized ones would + // let the stale 84 commit land last. + const flush = () => new Promise((r) => setTimeout(r, 0)); + await flush(); + while (pending.length > 0) { + for (const resolve of pending.splice(0).reverse()) resolve(); + await flush(); + } + await Promise.all([first, second]); + + // The latest selection wins everywhere — no stale overwrite and no split + // between the in-memory session (getState), stored session, and preference. + expect(service.getState().currentProjectId).toBe(99); + expect(sessionPort.getCurrent()?.selectedProjectId).toBe(99); + expect(preferencePort.get("user-1", "us")?.lastSelectedProjectId).toBe(99); + }); + it("restores the selected project after app restart while logged out", async () => { const orgs = { "org-1": { diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts index 18e3882e88..9f9452ec04 100644 --- a/packages/core/src/auth/auth.ts +++ b/packages/core/src/auth/auth.ts @@ -86,6 +86,9 @@ export class AuthService extends TypedEventEmitter { private session: InMemorySession | null = null; private initializePromise: Promise | null = null; private refreshPromise: Promise | null = null; + // Serializes session-state commits so overlapping selections can't + // interleave across async encryption (see commitSessionState). + private commitChain: Promise = Promise.resolve(); constructor( @inject(AUTH_PREFERENCE_STORE) private readonly authPreference: IAuthPreferenceStore, @@ -254,7 +257,7 @@ export class AuthService extends TypedEventEmitter { ? await this.applyOrgChange(session, newOrgId) : session.orgProjectsMap; - this.commitSessionState(session, { + await this.commitSessionState(session, { orgProjectsMap, currentOrgId: newOrgId, currentProjectId: projectId, @@ -277,7 +280,7 @@ export class AuthService extends TypedEventEmitter { orgId, ); - this.commitSessionState(session, { + await this.commitSessionState(session, { orgProjectsMap, currentOrgId: orgId, currentProjectId, @@ -333,8 +336,26 @@ export class AuthService extends TypedEventEmitter { currentOrgId: string | null; currentProjectId: number | null; }, - ): void { - this.session = { + ): Promise { + // Serialize commits onto a chain so overlapping selections can't + // interleave across async encryption and clobber a newer one. The chain + // swallows rejections so one failure doesn't wedge later commits; the + // returned promise still rejects for the caller. + const run = this.commitChain.then(() => + this.applyCommittedSession(prevSession, next), + ); + this.commitChain = run.catch(() => {}); + return run; + } + private async applyCommittedSession( + prevSession: InMemorySession, + next: { + orgProjectsMap: OrgProjectsMap; + currentOrgId: string | null; + currentProjectId: number | null; + }, + ): Promise { + const nextSession: InMemorySession = { ...prevSession, orgProjectsMap: next.orgProjectsMap, currentOrgId: next.currentOrgId, @@ -342,13 +363,18 @@ export class AuthService extends TypedEventEmitter { orgProjectsIncomplete: false, }; - this.persistProjectPreference(this.session); - this.persistSession({ - refreshToken: this.session.refreshToken, - cloudRegion: this.session.cloudRegion, + // Persist the durable session first — the only step that can fail (async + // encryption may reject). Mutate this.session, the preference, and + // published state only after it resolves, so a rejection leaves every + // layer on the prior session. + await this.persistSession({ + refreshToken: nextSession.refreshToken, + cloudRegion: nextSession.cloudRegion, selectedProjectId: next.currentProjectId, }); + this.session = nextSession; + this.persistProjectPreference(nextSession); this.updateState({ orgProjectsMap: next.orgProjectsMap, currentOrgId: next.currentOrgId, @@ -413,7 +439,7 @@ export class AuthService extends TypedEventEmitter { return; } - const storedSession = this.resolveStoredSession(); + const storedSession = await this.resolveStoredSession(); if (!storedSession) { this.logger.warn("Stored auth session could not be decrypted"); this.authSession.clearCurrent(); @@ -500,9 +526,12 @@ export class AuthService extends TypedEventEmitter { return this.refreshPromise; } - const sessionInput = this.getSessionInputForRefresh(); - + // Assign refreshPromise synchronously — no await before this — so + // concurrent callers dedupe onto one refresh. Resolving the stored session + // (now async) must happen INSIDE refreshAndSync, else two callers both + // refresh and burn the rotating token twice. const refreshAndSync = async (): Promise => { + const sessionInput = await this.getSessionInputForRefresh(); let session: InMemorySession; try { session = await this.refreshSession(sessionInput); @@ -532,7 +561,7 @@ export class AuthService extends TypedEventEmitter { return this.refreshPromise; } - private getSessionInputForRefresh(): StoredSessionInput { + private async getSessionInputForRefresh(): Promise { if (this.session) { return { refreshToken: this.session.refreshToken, @@ -541,7 +570,7 @@ export class AuthService extends TypedEventEmitter { }; } - const storedSession = this.resolveStoredSession(); + const storedSession = await this.resolveStoredSession(); if (!storedSession) { throw new NotAuthenticatedError(); } @@ -794,7 +823,7 @@ export class AuthService extends TypedEventEmitter { session: InMemorySession, ): Promise { this.persistProjectPreference(session); - this.persistSession({ + await this.persistSession({ refreshToken: session.refreshToken, cloudRegion: session.cloudRegion, selectedProjectId: session.currentProjectId, @@ -816,15 +845,15 @@ export class AuthService extends TypedEventEmitter { void this.refreshOrgProjects(); } } - private persistSession(input: { + private async persistSession(input: { refreshToken: string; cloudRegion: CloudRegion; selectedProjectId: number | null; - }): void { + }): Promise { const priorSelected = this.authSession.getCurrent()?.selectedProjectId ?? null; this.authSession.saveCurrent({ - refreshTokenEncrypted: this.cipher.encrypt(input.refreshToken), + refreshTokenEncrypted: await this.cipher.encrypt(input.refreshToken), cloudRegion: input.cloudRegion, selectedProjectId: input.selectedProjectId ?? priorSelected, scopeVersion: OAUTH_SCOPE_VERSION, @@ -1051,11 +1080,13 @@ export class AuthService extends TypedEventEmitter { private handleResume = (): void => { this.attemptSessionRecovery(); }; - private resolveStoredSession(): StoredSessionInput | null { + private async resolveStoredSession(): Promise { const stored = this.authSession.getCurrent(); if (!stored) return null; - const refreshToken = this.cipher.decrypt(stored.refreshTokenEncrypted); + const refreshToken = await this.cipher.decrypt( + stored.refreshTokenEncrypted, + ); if (!refreshToken) return null; return { @@ -1077,14 +1108,10 @@ export class AuthService extends TypedEventEmitter { if (!stored) return; if (stored.scopeVersion < OAUTH_SCOPE_VERSION) return; - if (!this.resolveStoredSession()) return; - - // Route through ensureValidSession so a refresh already in flight (e.g. the - // background bootstrap restore past its deadline) is shared instead of - // kicking a second concurrent token refresh that would burn the same - // rotating refresh token twice. - this.recoveryPromise = this.ensureValidSession() - .then(() => undefined) + // Claim the recovery slot synchronously so concurrent triggers don't both + // kick a token refresh; decryptability is now async (Web Crypto), so it's + // validated inside recoverSession. + this.recoveryPromise = this.recoverSession() .catch((error) => { this.logger.warn("Session recovery failed", { error }); }) @@ -1092,6 +1119,16 @@ export class AuthService extends TypedEventEmitter { this.recoveryPromise = null; }); } + private async recoverSession(): Promise { + // Bail before touching the network if the stored token can't be decrypted. + if (!(await this.resolveStoredSession())) return; + + // Route through ensureValidSession so a refresh already in flight (e.g. the + // background bootstrap restore past its deadline) is shared instead of + // kicking a second concurrent token refresh that would burn the same + // rotating refresh token twice. + await this.ensureValidSession(); + } private refreshOrgProjects(): Promise { if (this.orgProjectsRefreshPromise) { @@ -1157,7 +1194,7 @@ export class AuthService extends TypedEventEmitter { null, lastSelectedOrgId: lastPrefs?.lastSelectedOrgId ?? null, }); - this.commitSessionState(session, { + await this.commitSessionState(session, { orgProjectsMap: map, currentOrgId: session.currentOrgId, currentProjectId, diff --git a/packages/core/src/auth/identifiers.ts b/packages/core/src/auth/identifiers.ts index 300efbd74a..b0cb002595 100644 --- a/packages/core/src/auth/identifiers.ts +++ b/packages/core/src/auth/identifiers.ts @@ -89,11 +89,12 @@ export const AUTH_OAUTH_FLOW_SERVICE = Symbol.for( /** * Machine-bound symmetric cipher for the refresh token at rest. Desktop adapter - * wraps the existing encryption util (node:crypto + machine id). + * wraps the existing encryption util (node:crypto + machine id); the web adapter + * uses a non-extractable Web Crypto key (async), so the contract is async. */ export interface IAuthTokenCipher { - encrypt(plaintext: string): string; - decrypt(encrypted: string): string | null; + encrypt(plaintext: string): Promise; + decrypt(encrypted: string): Promise; } export const AUTH_TOKEN_CIPHER = Symbol.for("posthog.core.auth.tokenCipher");