diff --git a/packages/plugins/sandbox-providers/daytona/src/plugin.ts b/packages/plugins/sandbox-providers/daytona/src/plugin.ts index 9e38e0b381..52819929b5 100644 --- a/packages/plugins/sandbox-providers/daytona/src/plugin.ts +++ b/packages/plugins/sandbox-providers/daytona/src/plugin.ts @@ -2870,15 +2870,16 @@ const plugin = definePlugin({ daytonaLoginPtyByRoute.set(params.hostRouteId, entry); daytonaLoginPtyBySession.set(workerSessionId, entry); // Register the output listener before the first input, so no early output - // chunk is lost. The client stamps the worker session id, so the host binds - // the output to the open route. + // chunk is lost. The client stamps the host route identifier and the worker + // session identifier, so the host can hold more than one concurrent login + // pseudo-terminal on this worker and binds the output to its own route. session.onData((chunk) => { - pluginContext?.loginPty.output(workerSessionId, chunk); + pluginContext?.loginPty.output(params.hostRouteId, workerSessionId, chunk); }); // Forward the child exit one time. The host resolves the login run on it. void session.wait().then( - (result) => pluginContext?.loginPty.exit(workerSessionId, result.exitCode), - () => pluginContext?.loginPty.exit(workerSessionId, null), + (result) => pluginContext?.loginPty.exit(params.hostRouteId, workerSessionId, result.exitCode), + () => pluginContext?.loginPty.exit(params.hostRouteId, workerSessionId, null), ); return { workerSessionId }; }, diff --git a/packages/plugins/sdk/src/protocol.ts b/packages/plugins/sdk/src/protocol.ts index 81ba4e6161..071c1ac36a 100644 --- a/packages/plugins/sdk/src/protocol.ts +++ b/packages/plugins/sdk/src/protocol.ts @@ -1086,6 +1086,12 @@ export interface PluginLoginPtyCloseResult { /** The worker→host pseudo-terminal output notification parameters. Modeled on `execute.log`. */ export interface PluginLoginPtyOutputParams { + /** + * The host route identifier the open request carried. The worker echoes it, + * so the host can hold more than one concurrent login pseudo-terminal per + * worker and route each chunk to its own route. + */ + hostRouteId: string; /** The worker session identifier that the open reply returned. */ workerSessionId: string; /** The raw terminal output bytes. */ @@ -1094,6 +1100,12 @@ export interface PluginLoginPtyOutputParams { /** The worker→host pseudo-terminal exit notification parameters. */ export interface PluginLoginPtyExitParams { + /** + * The host route identifier the open request carried. The worker echoes it, + * so the host can hold more than one concurrent login pseudo-terminal per + * worker and resolve the exit against its own route. + */ + hostRouteId: string; /** The worker session identifier that the open reply returned. */ workerSessionId: string; /** The child exit code, or null when the child ended with no code. */ diff --git a/packages/plugins/sdk/src/testing.ts b/packages/plugins/sdk/src/testing.ts index 6cb96790da..8f7d04bfdb 100644 --- a/packages/plugins/sdk/src/testing.ts +++ b/packages/plugins/sdk/src/testing.ts @@ -2479,10 +2479,10 @@ export function createTestHarness(options: TestHarnessOptions): TestHarness { }, }, loginPty: { - output(_workerSessionId: string, _chunk: string) { + output(_hostRouteId: string, _workerSessionId: string, _chunk: string) { // No-op in test harness — the host login route is not wired here. }, - exit(_workerSessionId: string, _exitCode: number | null) { + exit(_hostRouteId: string, _workerSessionId: string, _exitCode: number | null) { // No-op in test harness — the host login route is not wired here. }, }, diff --git a/packages/plugins/sdk/src/types.ts b/packages/plugins/sdk/src/types.ts index a3525bc422..56a561a298 100644 --- a/packages/plugins/sdk/src/types.ts +++ b/packages/plugins/sdk/src/types.ts @@ -2025,28 +2025,32 @@ export interface PluginExecutionClient { * from a sandbox provider worker to the host. * * The worker opener registers the output listener on the session and forwards - * each raw chunk through `output(workerSessionId, chunk)`. It forwards the child - * exit through `exit(workerSessionId, exitCode)`. Each call carries the worker - * session identifier the open reply returned, so the host binds the output to - * the open route by that identifier while the route is open. The host drops a - * chunk or an exit that carries an unknown or a mismatched identifier, and it - * never logs the raw bytes. The default is a no-op that never throws. + * each raw chunk through `output(hostRouteId, workerSessionId, chunk)`. It + * forwards the child exit through `exit(hostRouteId, workerSessionId, + * exitCode)`. Each call carries the host route identifier the open request + * carried and the worker session identifier the open reply returned, so the + * host can hold more than one concurrent login pseudo-terminal per worker and + * bind each chunk to its own route. The host drops a chunk or an exit that + * carries an unknown, a stale, or a mismatched identifier, and it never logs + * the raw bytes. The default is a no-op that never throws. */ export interface PluginLoginPtyClient { /** * Deliver one raw output chunk of a live login pseudo-terminal. * + * @param hostRouteId - The host route identifier the open request carried. The worker echoes it, so the host routes the chunk to its own route. * @param workerSessionId - The worker session identifier the open reply returned. * @param chunk - The raw terminal output text. */ - output(workerSessionId: string, chunk: string): void; + output(hostRouteId: string, workerSessionId: string, chunk: string): void; /** * Deliver the child exit of a live login pseudo-terminal. * + * @param hostRouteId - The host route identifier the open request carried. The worker echoes it, so the host resolves the exit against its own route. * @param workerSessionId - The worker session identifier the open reply returned. * @param exitCode - The child exit code, or null when the child ended with no code. */ - exit(workerSessionId: string, exitCode: number | null): void; + exit(hostRouteId: string, workerSessionId: string, exitCode: number | null): void; } /** diff --git a/packages/plugins/sdk/src/worker-rpc-host.ts b/packages/plugins/sdk/src/worker-rpc-host.ts index 73fd52c44a..5e56ef200c 100644 --- a/packages/plugins/sdk/src/worker-rpc-host.ts +++ b/packages/plugins/sdk/src/worker-rpc-host.ts @@ -1381,23 +1381,28 @@ export function startWorkerRpcHost(options: WorkerRpcHostOptions): WorkerRpcHost }, loginPty: { - output(workerSessionId: string, chunk: string): void { + output(hostRouteId: string, workerSessionId: string, chunk: string): void { // Forward one raw output chunk of a live login pseudo-terminal. The - // notification carries the worker session identifier, so the host binds - // the chunk to the open route by that identifier while the route is - // open. The host drops an unknown or a mismatched identifier and never - // logs the raw bytes. This notification carries no invocation id, - // because it fires after the open reply returns. + // notification echoes the host route identifier and the worker session + // identifier, so the host can hold more than one concurrent login + // pseudo-terminal per worker and binds the chunk to its own route + // while that route is open. The host drops an unknown, a stale, or a + // mismatched identifier and never logs the raw bytes. This + // notification carries no invocation id, because it fires after the + // open reply returns. + if (typeof hostRouteId !== "string" || hostRouteId.length === 0) return; if (typeof workerSessionId !== "string" || workerSessionId.length === 0) return; if (typeof chunk !== "string" || chunk.length === 0) return; - notifyHost(LOGIN_PTY_OUTPUT_NOTIFICATION, { workerSessionId, chunk }); + notifyHost(LOGIN_PTY_OUTPUT_NOTIFICATION, { hostRouteId, workerSessionId, chunk }); }, - exit(workerSessionId: string, exitCode: number | null): void { + exit(hostRouteId: string, workerSessionId: string, exitCode: number | null): void { // Forward the child exit of a live login pseudo-terminal. The host - // resolves the open route's wait promise by the worker session - // identifier while the route is open. + // resolves its own route's wait promise by the host route identifier + // and the bound worker session identifier while that route is open. + if (typeof hostRouteId !== "string" || hostRouteId.length === 0) return; if (typeof workerSessionId !== "string" || workerSessionId.length === 0) return; notifyHost(LOGIN_PTY_EXIT_NOTIFICATION, { + hostRouteId, workerSessionId, exitCode: typeof exitCode === "number" ? exitCode : null, }); diff --git a/packages/plugins/sdk/tests/worker-rpc-host.test.ts b/packages/plugins/sdk/tests/worker-rpc-host.test.ts index 5a32f858b0..6652995631 100644 --- a/packages/plugins/sdk/tests/worker-rpc-host.test.ts +++ b/packages/plugins/sdk/tests/worker-rpc-host.test.ts @@ -778,8 +778,8 @@ describe("worker setup-token pseudo-terminal dispatch", () => { // worker session id. The test drives them through the captured emitters. const controllablePlugin = definePlugin({ async setup(ctx) { - emitOutput = (chunk: string) => ctx.loginPty.output("ws-1", chunk); - resolveWait = (value) => ctx.loginPty.exit("ws-1", value.exitCode); + emitOutput = (chunk: string) => ctx.loginPty.output("route-1", "ws-1", chunk); + resolveWait = (value) => ctx.loginPty.exit("route-1", "ws-1", value.exitCode); }, async onLoginPtyOpen(params) { // The open carries the host route id, the closed command key, and the @@ -896,13 +896,13 @@ describe("worker setup-token pseudo-terminal dispatch", () => { (note) => note.method === "loginPty.output", ); expect(outputNotes.map((note) => note.params)).toEqual([ - { workerSessionId: "ws-1", chunk: "prompt output" }, + { hostRouteId: "route-1", workerSessionId: "ws-1", chunk: "prompt output" }, ]); const exitNotes = notifications.filter( (note) => note.method === "loginPty.exit", ); expect(exitNotes.map((note) => note.params)).toEqual([ - { workerSessionId: "ws-1", exitCode: 0 }, + { hostRouteId: "route-1", workerSessionId: "ws-1", exitCode: 0 }, ]); } finally { worker.stop(); diff --git a/server/src/__tests__/agent-device-login-routes.test.ts b/server/src/__tests__/agent-device-login-routes.test.ts index 2fd3657310..f2ce4267b5 100644 --- a/server/src/__tests__/agent-device-login-routes.test.ts +++ b/server/src/__tests__/agent-device-login-routes.test.ts @@ -294,6 +294,19 @@ function createMemoryStore(): AdapterAuthSessionStore & { rows: Map { expect(harness.acquisitions).toHaveLength(0); }); - it("delivers the one-time prompt to the owner on the first read only", async () => { + it("delivers the prompt to the owner on every read while the session is active, and clears it on a terminal transition", async () => { const app = await createApp(); const start = await request(app) @@ -656,17 +669,30 @@ describe("adapter device-login routes", () => { expect(start.status, JSON.stringify(start.body)).toBe(201); const sessionId = start.body.sessionId as string; - // The first authorized owner read receives the one-time prompt. + // The first authorized owner read receives the prompt. The response + // repeats the live prompt on every read, so it carries the same private + // no-store policy as the `.../login-sessions/active` route (see the tests + // below). const first = await request(app).get(`${loginPath(COMPANY_1)}/${sessionId}`); expect(first.status, JSON.stringify(first.body)).toBe(200); expect(first.body.prompt).toEqual({ url: DEVICE_LOGIN_URL, code: PROMPT_CODE }); + expect(first.headers["cache-control"]).toBe("no-store, private"); - // A second authorized owner read no longer carries the prompt. The status - // stays available, so the owner still tracks the session. + // A second authorized owner read still carries the prompt while the + // session is active. The status stays available too, so the owner still + // tracks the session across a page reload. const second = await request(app).get(`${loginPath(COMPANY_1)}/${sessionId}`); expect(second.status, JSON.stringify(second.body)).toBe(200); - expect(second.body.prompt).toBeNull(); + expect(second.body.prompt).toEqual({ url: DEVICE_LOGIN_URL, code: PROMPT_CODE }); expect(second.body.status).toBe(first.body.status); + + // Once the login reaches a terminal state, the prompt is gone. + harness.releaseGate(); + await vi.waitFor(async () => { + const afterTerminal = await request(app).get(`${loginPath(COMPANY_1)}/${sessionId}`); + expect(["authenticated", "failed"]).toContain(afterTerminal.body.status); + expect(afterTerminal.body.prompt).toBeNull(); + }); }); it("starts a Grok session, delivers the Grok prompt once, and a codex_local read finds no row", async () => { @@ -739,6 +765,52 @@ describe("adapter device-login routes", () => { expect(status.body.prompt).toBeUndefined(); }); + it("returns the caller's active login session with no session id in the URL", async () => { + const app = await createApp(); + + const start = await request(app) + .post(loginPath(COMPANY_1)) + .send({ environmentId: SANDBOX_ENV_1 }); + expect(start.status, JSON.stringify(start.body)).toBe(201); + + const active = await request(app).get(`${loginPath(COMPANY_1)}/active`); + expect(active.status, JSON.stringify(active.body)).toBe(200); + expect(active.body.sessionId).toBe(start.body.sessionId); + expect(active.body.prompt).toEqual({ url: DEVICE_LOGIN_URL, code: PROMPT_CODE }); + expect(active.headers["cache-control"]).toBe("no-store, private"); + }); + + it("returns the identical 404 on the active route for no active session, another owner, another company, and another adapter", async () => { + const app = await createApp(); + + // No active session exists yet. + const none = await request(app).get(`${loginPath(COMPANY_1)}/active`); + expect(none.status, JSON.stringify(none.body)).toBe(404); + expect(none.body).toEqual({ error: "Adapter login session not found" }); + + const start = await request(app) + .post(loginPath(COMPANY_1)) + .send({ environmentId: SANDBOX_ENV_1 }); + expect(start.status, JSON.stringify(start.body)).toBe(201); + + // A different board user in the same company holds no active session. + currentActor = boardActor(OWNER_B); + const otherOwner = await request(app).get(`${loginPath(COMPANY_1)}/active`); + expect(otherOwner.status, JSON.stringify(otherOwner.body)).toBe(404); + expect(otherOwner.body).toEqual(none.body); + currentActor = boardActor(OWNER_A); + + // The same owner under a different company holds no active session there. + const otherCompany = await request(app).get(`${loginPath(COMPANY_2)}/active`); + expect(otherCompany.status, JSON.stringify(otherCompany.body)).toBe(404); + expect(otherCompany.body).toEqual(none.body); + + // The owner's active session belongs to `codex_local`, not `grok_local`. + const otherAdapter = await request(app).get(`${loginPath(COMPANY_1, "grok_local")}/active`); + expect(otherAdapter.status, JSON.stringify(otherAdapter.body)).toBe(404); + expect(otherAdapter.body).toEqual(none.body); + }); + it("durably cancels a login for the owner and releases the company slot", async () => { const app = await createApp(); diff --git a/server/src/__tests__/device-login-service.test.ts b/server/src/__tests__/device-login-service.test.ts index ad7c479dc2..13566ca385 100644 --- a/server/src/__tests__/device-login-service.test.ts +++ b/server/src/__tests__/device-login-service.test.ts @@ -259,6 +259,19 @@ function createMemoryStore(): AdapterAuthSessionStore & { } return null; }, + async getActiveByOwner(companyId, startedByUserId, adapterType) { + for (const row of rows.values()) { + if ( + row.companyId === companyId && + row.startedByUserId === startedByUserId && + row.adapterType === adapterType && + isActive(row.status) + ) { + return { ...row }; + } + } + return null; + }, async withCompanyAdapterPromotionLock(_companyId, _startedByUserId, _adapterType, fn) { // The in-memory store runs on a single event loop, so it needs no real // lock. The pass-through keeps the store contract satisfied. @@ -305,8 +318,20 @@ describe("device login service", () => { const activity: LoginSessionActivityEvent[] = []; const promoted: Buffer[] = []; const promotionContexts: { sessionId: string; companyId: string }[] = []; + // Gate the login after the prompt surfaces, so the row stays in + // `waiting_for_user` long enough for the test to read the prompt (twice) + // before it releases the gate and lets the login race to `authenticated`. + let releaseGate!: () => void; + const gate = new Promise((resolve) => { + releaseGate = resolve; + }); + const execGated: ExecBehavior = async ({ onStdout }) => { + onStdout(PROMPT_OUTPUT); + await gate; + return { exitCode: 0 }; + }; const { runtime, deleteCalls } = createFakeRuntime({ - exec: execSuccess, + exec: execGated, authBytes: Buffer.from('{"token":"secret"}'), }); const companyId = randomUUID(); @@ -344,6 +369,14 @@ describe("device login service", () => { const other = await service.readOwnerSession(session.sessionId, companyId, OWNER_B); expect(other?.prompt).toBeNull(); + // The prompt survives a repeated read while the session is still active. The + // read never consumes it, so a page reload does not lose the code. + const repeatWhileActive = await service.readOwnerSession(session.sessionId, companyId, OWNER_A); + expect(repeatWhileActive?.prompt).toEqual({ url: DEVICE_LOGIN_URL, code: PROMPT_CODE }); + + // Release the gate, so the login proceeds to the promotion and the + // authenticated terminal. + releaseGate(); const outcome = await completed; expect(outcome.status).toBe("authenticated"); expect(outcome.cleanupPending).toBe(false); @@ -357,9 +390,11 @@ describe("device login service", () => { expect(promotionContexts).toEqual([{ sessionId: internalRow?.id, companyId }]); expect(internalRow?.id).not.toBe(session.sessionId); - // The prompt is one-time: a second owner read returns null. - const secondRead = await service.readOwnerSession(session.sessionId, companyId, OWNER_A); - expect(secondRead?.prompt).toBeNull(); + // The session is now terminal (authenticated). A read after a terminal + // transition returns a null prompt, and it carries `Cache-Control: + // no-store, private` at the route layer (see agent-device-login-routes.test.ts). + const afterTerminal = await service.readOwnerSession(session.sessionId, companyId, OWNER_A); + expect(afterTerminal?.prompt).toBeNull(); const row = await store.getByPublicId(session.sessionId, companyId); expect(row?.status).toBe("authenticated"); @@ -465,11 +500,19 @@ describe("device login service", () => { // the trusted adapter type, so a `grok_local` session runs the Grok parser, // not the Codex parser. const store = createMemoryStore(); - const execGrokSuccess: ExecBehavior = async ({ onStdout }) => { + // Gate the login after the prompt surfaces, so the row stays in + // `waiting_for_user` long enough for the test to read the prompt before the + // login races to its own terminal. + let releaseGate!: () => void; + const gate = new Promise((resolve) => { + releaseGate = resolve; + }); + const execGrokGatedSuccess: ExecBehavior = async ({ onStdout }) => { onStdout(GROK_PROMPT_OUTPUT); + await gate; return { exitCode: 0 }; }; - const { runtime } = createFakeRuntime({ exec: execGrokSuccess, authBytes: Buffer.from("{}") }); + const { runtime } = createFakeRuntime({ exec: execGrokGatedSuccess, authBytes: Buffer.from("{}") }); const companyId = randomUUID(); const service = makeService({ store, runtime }); const { session, completed } = await service.start({ @@ -485,6 +528,7 @@ describe("device login service", () => { const owner = await service.readOwnerSession(session.sessionId, companyId, OWNER_A); expect(owner?.prompt).toEqual({ url: GROK_DEVICE_LOGIN_URL, code: GROK_CODE }); + releaseGate(); await completed; }); @@ -605,9 +649,17 @@ describe("device login service", () => { const store = createMemoryStore(); const activity: LoginSessionActivityEvent[] = []; + // Gate the login after the prompt surfaces, so the row stays in + // `waiting_for_user` long enough for the test to read the prompt before the + // login races to its own terminal. + let releaseGate!: () => void; + const gate = new Promise((resolve) => { + releaseGate = resolve; + }); const { runtime } = createFakeRuntime({ exec: async ({ onStdout }) => { onStdout(GROK_PROMPT_OUTPUT); + await gate; return { exitCode: 0 }; }, authBytes: grokAuthBytes, @@ -630,6 +682,7 @@ describe("device login service", () => { await new Promise((resolve) => setImmediate(resolve)); const owner = await service.readOwnerSession(session.sessionId, companyId, OWNER_A); + releaseGate(); const outcome = await completed; expect(outcome.status).toBe("authenticated"); @@ -811,6 +864,36 @@ describe("device login service", () => { expect(row?.status).toBe("cancelled"); }); + it("clears the prompt on a cancellation, so the owner read returns a null prompt with Cache-Control set at the route layer", async () => { + const store = createMemoryStore(); + const { runtime } = createFakeRuntime({ exec: execHang }); + const service = makeService({ store, runtime }); + const controller = new AbortController(); + const companyId = randomUUID(); + const { session, completed } = await service.start({ + companyId, + environmentId: randomUUID(), + adapterType: ADAPTER_TYPE, + startedByUserId: OWNER_A, + signal: controller.signal, + }); + + await waitForStatus(store, session.sessionId, companyId, "waiting_for_user"); + await new Promise((resolve) => setImmediate(resolve)); + + // The prompt is present while the session is active. + const whileActive = await service.readOwnerSession(session.sessionId, companyId, OWNER_A); + expect(whileActive?.prompt).toEqual({ url: DEVICE_LOGIN_URL, code: PROMPT_CODE }); + + controller.abort(); + await completed; + + // The session is now terminal (cancelled). The prompt is gone. + const afterCancel = await service.readOwnerSession(session.sessionId, companyId, OWNER_A); + expect(afterCancel?.status).toBe("cancelled"); + expect(afterCancel?.prompt).toBeNull(); + }); + it("releases the lease when a transition fails after acquisition", async () => { const store = createMemoryStore(); // A rejecting `recordLeaseAcquired` forces a transition failure right after @@ -1091,6 +1174,79 @@ describe("device login service", () => { }); }); + describe("active session lookup", () => { + it("finds the caller's active session for a company and adapter, with no session id", async () => { + const store = createMemoryStore(); + const { runtime } = createFakeRuntime({ exec: execHang }); + const service = makeService({ store, runtime }); + const controller = new AbortController(); + const companyId = randomUUID(); + const { session, completed } = await service.start({ + companyId, + environmentId: randomUUID(), + adapterType: ADAPTER_TYPE, + startedByUserId: OWNER_A, + signal: controller.signal, + }); + await waitForStatus(store, session.sessionId, companyId, "waiting_for_user"); + await new Promise((resolve) => setImmediate(resolve)); + + const active = await service.readActiveOwnerSession(companyId, ADAPTER_TYPE, OWNER_A); + expect(active?.sessionId).toBe(session.sessionId); + expect(active?.prompt).toEqual({ url: DEVICE_LOGIN_URL, code: PROMPT_CODE }); + + controller.abort(); + await completed; + }); + + it("returns null for another owner, another company, and another adapter", async () => { + const store = createMemoryStore(); + const { runtime } = createFakeRuntime({ exec: execHang }); + const service = makeService({ store, runtime }); + const controller = new AbortController(); + const companyId = randomUUID(); + await service.start({ + companyId, + environmentId: randomUUID(), + adapterType: ADAPTER_TYPE, + startedByUserId: OWNER_A, + signal: controller.signal, + }); + + expect(await service.readActiveOwnerSession(companyId, ADAPTER_TYPE, OWNER_B)).toBeNull(); + expect( + await service.readActiveOwnerSession(randomUUID(), ADAPTER_TYPE, OWNER_A), + ).toBeNull(); + expect(await service.readActiveOwnerSession(companyId, "grok_local", OWNER_A)).toBeNull(); + + controller.abort(); + }); + + it("returns null once the caller's session has no active row", async () => { + const store = createMemoryStore(); + const { runtime } = createFakeRuntime({ exec: execHang }); + const service = makeService({ store, runtime }); + const controller = new AbortController(); + const companyId = randomUUID(); + const { completed } = await service.start({ + companyId, + environmentId: randomUUID(), + adapterType: ADAPTER_TYPE, + startedByUserId: OWNER_A, + signal: controller.signal, + }); + + // No active session for this owner before any start ever ran. + expect(await service.readActiveOwnerSession(companyId, ADAPTER_TYPE, OWNER_A)).not.toBeNull(); + + controller.abort(); + await completed; + + // The session is now terminal, so it no longer counts as active. + expect(await service.readActiveOwnerSession(companyId, ADAPTER_TYPE, OWNER_A)).toBeNull(); + }); + }); + describe("five-minute host timeout", () => { it("holds the session active until exactly five minutes, then times out and deletes", async () => { vi.useFakeTimers(); diff --git a/server/src/__tests__/fixtures/plugin-worker-login-pty.cjs b/server/src/__tests__/fixtures/plugin-worker-login-pty.cjs index 27d9f1c1e5..b143a0a29b 100644 --- a/server/src/__tests__/fixtures/plugin-worker-login-pty.cjs +++ b/server/src/__tests__/fixtures/plugin-worker-login-pty.cjs @@ -9,11 +9,20 @@ // - `mode`: "normal" | "malformed-open" | "no-open-reply" | "duplicate-open-reply" | // "exit-before-open-reply" // - `workerSessionId`: the worker session id the open reply returns (default "ws-1") -// - `outputs`: an array of `{ chunk, sid? }`. The fixture emits each as an output -// notification after the open reply. `sid` defaults to the real worker session -// id; a test sets a wrong `sid` to prove the host drops a mismatched -// notification. +// - `outputs`: an array of `{ chunk, sid?, crossRoute?, omitHostRouteId? }`. The +// fixture emits each as an output notification after the open reply. `sid` +// defaults to the real worker session id; a test sets a wrong `sid` to prove +// the host drops a mismatched notification. `crossRoute: true` sends this +// route's own worker session id under a DIFFERENT, already-open route's host +// route identifier (any other entry currently registered), so a test proves +// the host drops a swapped `(hostRouteId, workerSessionId)` pair instead of +// misdelivering it. `omitHostRouteId: true` sends the notification with no +// `hostRouteId` field at all, so a test proves the host warns about a plugin +// build old enough to omit the field, instead of silently dropping it. // - `exitCode`: when set, the fixture emits an exit notification after the outputs. +// - `omitHostRouteIdOnExit`: when true, the main `exitCode` exit notification +// carries no `hostRouteId` field, so a test proves the host still resolves +// a legacy worker's exit by the worker session id. // - `extraExits`: an array of `{ exitCode, sid? }`. The fixture emits each as a // further exit notification, after the main `exitCode` exit. `sid` defaults // to the real worker session id; a test sets a wrong `sid` to script a worker @@ -43,30 +52,55 @@ function send(message) { process.stdout.write(`${JSON.stringify(message)}\n`); } -// Serialize one array of `{ chunk, sid? }` entries as newline-delimited output -// notification lines. A test sets a wrong `sid` on one entry to force a -// mismatch. -function outputLines(entries, workerSessionId) { +// Find some other host route identifier the fixture already registered, for +// a `crossRoute` output entry. Returns null when no other route is open, +// which a well-formed test never triggers. +function pickOtherHostRouteId(ownHostRouteId) { + for (const key of routes.keys()) { + if (key !== ownHostRouteId) return key; + } + return null; +} + +// Serialize one array of `{ chunk, sid?, crossRoute?, omitHostRouteId? }` +// entries as newline-delimited output notification lines. A test sets a +// wrong `sid` on one entry to force a mismatch. `crossRoute: true` stamps +// some OTHER already-open route's host route identifier on this route's own +// worker session id, so a test proves the host drops a swapped pair. Every +// other line echoes this route's own host route identifier, so the host can +// route each chunk to its own route when the worker holds more than one. +// `omitHostRouteId: true` sends the notification with no `hostRouteId` field +// at all, so a test proves the host warns instead of silently dropping it. +function outputLines(entries, hostRouteId, workerSessionId) { let lines = ""; for (const entry of entries) { + const effectiveHostRouteId = entry.crossRoute + ? pickOtherHostRouteId(hostRouteId) ?? hostRouteId + : entry.hostRouteId ?? hostRouteId; + const params = { + workerSessionId: entry.sid ?? workerSessionId, + chunk: entry.chunk, + }; + if (!entry.omitHostRouteId) params.hostRouteId = effectiveHostRouteId; lines += `${JSON.stringify({ jsonrpc: "2.0", method: "loginPty.output", - params: { - workerSessionId: entry.sid ?? workerSessionId, - chunk: entry.chunk, - }, + params, })}\n`; } return lines; } -// Serialize one exit notification line for the given worker session id. -function exitLine(workerSessionId, exitCode) { +// Serialize one exit notification line for the given host route and worker +// session id. `omit: true` sends the notification with no `hostRouteId` +// field at all, matching a plugin build old enough to predate it. +function exitLine(hostRouteId, workerSessionId, exitCode, omit) { + const params = { workerSessionId, exitCode }; + if (!omit) params.hostRouteId = hostRouteId; return `${JSON.stringify({ jsonrpc: "2.0", method: "loginPty.exit", - params: { workerSessionId, exitCode }, + params, })}\n`; } @@ -75,14 +109,18 @@ function exitLine(workerSessionId, exitCode) { // fixed composition otherwise emits, in order: the pre-exit outputs, then // the main exit, then every extra exit, then the post-exit outputs. The // batch mode writes these together with the open reply in one stdout write. -function scriptedOutputLines(directive, workerSessionId) { +function scriptedOutputLines(directive, hostRouteId, workerSessionId) { if (Array.isArray(directive.sequence)) { let lines = ""; for (const entry of directive.sequence) { if (entry.type === "output") { - lines += outputLines([{ chunk: entry.chunk, sid: entry.sid }], workerSessionId); + lines += outputLines( + [{ chunk: entry.chunk, sid: entry.sid, hostRouteId: entry.hostRouteId, crossRoute: entry.crossRoute }], + hostRouteId, + workerSessionId, + ); } else if (entry.type === "exit") { - lines += exitLine(entry.sid ?? workerSessionId, entry.exitCode); + lines += exitLine(entry.hostRouteId ?? hostRouteId, entry.sid ?? workerSessionId, entry.exitCode); } } return lines; @@ -92,14 +130,14 @@ function scriptedOutputLines(directive, workerSessionId) { ? directive.outputsAfterExit : []; const extraExits = Array.isArray(directive.extraExits) ? directive.extraExits : []; - let lines = outputLines(outputs, workerSessionId); + let lines = outputLines(outputs, hostRouteId, workerSessionId); if (typeof directive.exitCode === "number") { - lines += exitLine(workerSessionId, directive.exitCode); + lines += exitLine(hostRouteId, workerSessionId, directive.exitCode, directive.omitHostRouteIdOnExit); } for (const exit of extraExits) { - lines += exitLine(exit.sid ?? workerSessionId, exit.exitCode); + lines += exitLine(exit.hostRouteId ?? hostRouteId, exit.sid ?? workerSessionId, exit.exitCode); } - lines += outputLines(outputsAfterExit, workerSessionId); + lines += outputLines(outputsAfterExit, hostRouteId, workerSessionId); return lines; } @@ -156,7 +194,7 @@ rl.on("line", (line) => { // Emit the scripted pre-bind outputs, then exit with no open reply. The // route never binds, so a test proves the worker-exit path clears the // pre-bind queue. - process.stdout.write(scriptedOutputLines(directive, workerSessionId)); + process.stdout.write(scriptedOutputLines(directive, params.hostRouteId, workerSessionId)); process.exit(1); return; } @@ -176,7 +214,9 @@ rl.on("line", (line) => { // notification arrives before the route binds — even a malformed reply // that never binds. The host must queue and, on a bind, replay it; on a // malformed reply, it must clear the queue instead. - process.stdout.write(openReplyLine + scriptedOutputLines(directive, workerSessionId)); + process.stdout.write( + openReplyLine + scriptedOutputLines(directive, params.hostRouteId, workerSessionId), + ); return; } @@ -193,7 +233,7 @@ rl.on("line", (line) => { // Emit the scripted output and the exit after the open reply, so the host // binds the route first. setImmediate(() => { - process.stdout.write(scriptedOutputLines(directive, workerSessionId)); + process.stdout.write(scriptedOutputLines(directive, params.hostRouteId, workerSessionId)); }); return; } @@ -201,12 +241,16 @@ rl.on("line", (line) => { if (method === "loginPtyInput") { // Echo the input back as one output notification for the bound session, so a // test proves the input reaches the worker and the output routes back. - for (const entry of routes.values()) { + for (const [hostRouteId, entry] of routes.entries()) { if (entry.workerSessionId === params.workerSessionId) { send({ jsonrpc: "2.0", method: "loginPty.output", - params: { workerSessionId: entry.workerSessionId, chunk: `echo:${params.data}` }, + params: { + hostRouteId, + workerSessionId: entry.workerSessionId, + chunk: `echo:${params.data}`, + }, }); } } diff --git a/server/src/__tests__/plugin-worker-manager.test.ts b/server/src/__tests__/plugin-worker-manager.test.ts index 56c268dc7a..03035ff2dc 100644 --- a/server/src/__tests__/plugin-worker-manager.test.ts +++ b/server/src/__tests__/plugin-worker-manager.test.ts @@ -29,6 +29,7 @@ vi.mock("../middleware/logger.js", () => { import { logger } from "../middleware/logger.js"; import { appendStderrExcerpt, + createDuplexRouteSlotController, createPluginWorkerHandle, formatWorkerFailureMessage, resolveRpcCallTimeoutMs, @@ -1169,24 +1170,21 @@ describe("plugin worker manager setup-token pty route gate", () => { } }); - it("permits one active credential pseudo-terminal per worker", async () => { + it("permits two concurrent credential pseudo-terminals on one worker", async () => { const handle = makeLoginPtyHandle(); try { await handle.start(); + // Neither open waits on the other. Both return a live route on the same + // worker, so a second owner is never blocked by a first owner's session. const first = await handle.openLoginPtySession( - ptyOpenInput({ mode: "normal" }), + ptyOpenInput({ mode: "normal", workerSessionId: "ws-A" }), ); - // A second open while the first route is not closed rejects with one fixed - // non-secret error before it reaches the worker. - await expect( - handle.openLoginPtySession(ptyOpenInput({ mode: "normal" })), - ).rejects.toThrow("LOGIN_PTY_ROUTE_BUSY"); - await first.close(); - // After the first route closes and the worker acknowledges the close, a new - // open is admitted. const second = await handle.openLoginPtySession( - ptyOpenInput({ mode: "normal" }), + ptyOpenInput({ mode: "normal", workerSessionId: "ws-B" }), ); + expect(first).toBeDefined(); + expect(second).toBeDefined(); + await first.close(); await second.close(); } finally { await handle.stop().catch(() => undefined); @@ -1365,7 +1363,526 @@ describe("plugin worker manager setup-token pty route gate", () => { }); // --------------------------------------------------------------------------- -// Login pseudo-terminal pre-bind queue (GH #12122 / PAP-5132) +// Login pseudo-terminal concurrency +// --------------------------------------------------------------------------- +// One shared plugin worker now holds more than one live login pseudo-terminal +// route at once, so a second owner is never blocked by a first owner's +// session. The host resolves an output or an exit notification by the host +// route identifier first, then checks the bound worker session identifier, +// so each route receives only its own data even while another route on the +// same worker is live. + +describe("plugin worker manager login pseudo-terminal concurrency", () => { + it("delivers output to each concurrent route only, with no cross-talk", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + const routeA = await handle.openLoginPtySession( + ptyOpenInput({ + workerSessionId: "ws-A", + outputs: [{ chunk: "a-1" }], + exitCode: 0, + }), + ); + const routeB = await handle.openLoginPtySession( + ptyOpenInput({ + workerSessionId: "ws-B", + outputs: [{ chunk: "b-1" }], + exitCode: 0, + }), + ); + const aChunks: string[] = []; + const bChunks: string[] = []; + routeA.onData((chunk) => aChunks.push(chunk)); + routeB.onData((chunk) => bChunks.push(chunk)); + await expect(routeA.wait()).resolves.toEqual({ exitCode: 0 }); + await expect(routeB.wait()).resolves.toEqual({ exitCode: 0 }); + // The host resolves each notification by its own host route identifier, so + // neither route ever sees the other's chunk. + expect(aChunks).toEqual(["a-1"]); + expect(bChunks).toEqual(["b-1"]); + await routeA.close(); + await routeB.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("settles the exit of one route only, while a second concurrent route stays live", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + const routeA = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A" }), + ); + const routeB = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-B", exitCode: 0 }), + ); + await expect(routeB.wait()).resolves.toEqual({ exitCode: 0 }); + // Route A never received an exit, so its wait is still pending. Race it + // against a short delay to prove it has not settled. + const stillPending = await Promise.race([ + routeA.wait().then(() => "settled"), + new Promise((resolve) => setTimeout(() => resolve("pending"), 50)), + ]); + expect(stillPending).toBe("pending"); + await routeA.close(); + await routeB.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("drops a swapped (hostRouteId, workerSessionId) notification and delivers it to neither route", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + const routeA = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A" }), + ); + // Route B emits one output that carries route A's host route identifier + // but route B's own worker session identifier — a swapped pair. The host + // resolves route A by the host route identifier, then finds route B's + // worker session identifier does not match route A's bound identifier, + // so it drops the notification. It never reaches route A, and route B + // never sent it under its own host route identifier, so it never + // reaches route B either. + const routeB = await handle.openLoginPtySession( + ptyOpenInput({ + workerSessionId: "ws-B", + outputs: [{ chunk: "swapped", crossRoute: true }], + }), + ); + const aChunks: string[] = []; + const bChunks: string[] = []; + routeA.onData((chunk) => aChunks.push(chunk)); + routeB.onData((chunk) => bChunks.push(chunk)); + // Give the swapped notification time to arrive and be dropped. + await new Promise((resolve) => setTimeout(resolve, 60)); + expect(aChunks).toEqual([]); + expect(bChunks).toEqual([]); + await routeA.close(); + await routeB.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("drops an output notification for an unknown host route identifier", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + const session = await handle.openLoginPtySession( + ptyOpenInput({ + workerSessionId: "ws-A", + sequence: [ + { type: "output", chunk: "ghost", hostRouteId: "totally-unknown-route" }, + { type: "output", chunk: "good" }, + ], + }), + ); + const chunks: string[] = []; + session.onData((chunk) => chunks.push(chunk)); + await vi.waitFor(() => expect(chunks).toContain("good")); + // The unknown host route identifier resolves to no route on this worker, + // so the host drops it before it can reach any listener. + expect(chunks).toEqual(["good"]); + await session.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("replays a pre-bind output record to its own route while a second route is already open", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + const routeA = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A" }), + ); + // Route B's fixture writes its open reply and its output notification in + // one stdout write, so the host reads both before route B's bind, while + // route A is already open. The bind replays the held record into route + // B's own queue; it never reaches route A. + const routeB = await handle.openLoginPtySession( + ptyOpenInput({ + batchWithOpenReply: true, + workerSessionId: "ws-B", + outputs: [{ chunk: "b-batched" }], + }), + ); + const aChunks: string[] = []; + const bChunks: string[] = []; + routeA.onData((chunk) => aChunks.push(chunk)); + routeB.onData((chunk) => bChunks.push(chunk)); + expect(aChunks).toEqual([]); + expect(bChunks).toEqual(["b-batched"]); + await routeA.close(); + await routeB.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("closes one route without affecting a second concurrent route on the same worker", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + const routeA = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A" }), + ); + const routeB = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-B" }), + ); + await routeA.close(); + // Route B is still live: writing to it still round-trips through the + // worker and back to the listener. + const bChunks: string[] = []; + routeB.onData((chunk) => bChunks.push(chunk)); + routeB.write("still-here"); + await vi.waitFor(() => expect(bChunks).toContain("echo:still-here")); + await routeB.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("settles both concurrent routes exactly once when the worker exits", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + const routeA = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A" }), + ); + const routeB = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-B" }), + ); + const waitA = routeA.wait(); + const waitB = routeB.wait(); + await handle.stop(); + // The worker exit closes every route it still holds and resolves each + // login wait with the fixed non-secret exit, in one sweep. + await expect(waitA).resolves.toEqual({ exitCode: null }); + await expect(waitB).resolves.toEqual({ exitCode: null }); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("retires the worker on an unconfirmed close, settling every concurrent route and logging no raw output", async () => { + const handle = makeLoginPtyHandle({ + loginPtyLimits: { closeTimeoutMs: 200 }, + }); + const secretMarker = "super-secret-login-code-must-never-reach-a-log-line"; + vi.mocked(logger.warn).mockClear(); + vi.mocked(logger.error).mockClear(); + vi.mocked(logger.info).mockClear(); + vi.mocked(logger.debug).mockClear(); + try { + await handle.start(); + const exited = new Promise((resolve) => { + handle.on("exit", () => resolve()); + }); + const routeA = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A", closeMode: "bad-ack" }), + ); + const routeB = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-B", outputs: [{ chunk: secretMarker }] }), + ); + const waitB = routeB.wait(); + await routeA.close(); + // Route A's close acknowledgement carried a mismatched host route id, so + // the host fails closed and retires the whole worker before any reuse. + // That retirement settles route B too, even though B's own close was + // never called. + await exited; + await expect(waitB).resolves.toEqual({ exitCode: null }); + await expect( + handle.openLoginPtySession(ptyOpenInput({ mode: "normal" })), + ).rejects.toThrow(); + const loggedText = [ + ...vi.mocked(logger.warn).mock.calls, + ...vi.mocked(logger.error).mock.calls, + ...vi.mocked(logger.info).mock.calls, + ...vi.mocked(logger.debug).mock.calls, + ] + .flat() + .map((arg) => JSON.stringify(arg)) + .join("\n"); + expect(loggedText).not.toContain(secretMarker); + } finally { + await handle.stop().catch(() => undefined); + } + }); +}); + +// --------------------------------------------------------------------------- +// The missing-hostRouteId diagnostic. `hostRouteId` is the routing key the +// current protocol tags every login pseudo-terminal notification with, but a +// plugin build old enough to predate it still tags the same notification with +// the worker session identifier, the sole routing key the previous protocol +// used. The host warns once about the old build, then still resolves the +// route by that identifier, so a legacy worker keeps its login output and its +// exit notification instead of losing them. +// --------------------------------------------------------------------------- + +describe("plugin worker manager login pseudo-terminal missing hostRouteId diagnostic", () => { + it("delivers output from a legacy worker with no hostRouteId, resolved by the worker session id, and warns once", async () => { + const handle = makeLoginPtyHandle(); + vi.mocked(logger.warn).mockClear(); + try { + await handle.start(); + const chunks: string[] = []; + const route = await handle.openLoginPtySession( + ptyOpenInput({ + workerSessionId: "ws-A", + outputs: [{ chunk: "legacy-output", omitHostRouteId: true }], + }), + ); + route.onData((chunk) => chunks.push(chunk)); + await vi.waitFor(() => expect(chunks).toContain("legacy-output")); + + const warnCalls = vi.mocked(logger.warn).mock.calls.flat().map((arg) => JSON.stringify(arg)); + const matches = warnCalls.filter((call) => call.includes("hostRouteId")); + expect(matches).toHaveLength(1); + expect(matches[0]).toContain("plugin build is too old"); + // No identifier and no raw output in the warning. + expect(matches[0]).not.toContain("legacy-output"); + expect(matches[0]).not.toContain("ws-A"); + + await route.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("settles the login wait on a legacy worker's exit notification with no hostRouteId", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + const route = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A", exitCode: 0, omitHostRouteIdOnExit: true }), + ); + await expect(route.wait()).resolves.toEqual({ exitCode: 0 }); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("repeats the warning no more than one time for one worker", async () => { + const handle = makeLoginPtyHandle(); + vi.mocked(logger.warn).mockClear(); + try { + await handle.start(); + const route = await handle.openLoginPtySession( + ptyOpenInput({ + workerSessionId: "ws-A", + outputs: [ + { chunk: "one", omitHostRouteId: true }, + { chunk: "two", omitHostRouteId: true }, + { chunk: "three", omitHostRouteId: true }, + ], + }), + ); + // Give every notification time to arrive. + await new Promise((resolve) => setTimeout(resolve, 60)); + + const warnCalls = vi.mocked(logger.warn).mock.calls.flat().map((arg) => JSON.stringify(arg)); + const matches = warnCalls.filter((call) => call.includes("hostRouteId")); + expect(matches).toHaveLength(1); + + await route.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("drops a message with no hostRouteId and no worker session id", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + const chunks: string[] = []; + const route = await handle.openLoginPtySession( + ptyOpenInput({ + workerSessionId: "ws-A", + outputs: [ + { chunk: "unroutable", omitHostRouteId: true, sid: "" }, + { chunk: "good" }, + ], + }), + ); + route.onData((chunk) => chunks.push(chunk)); + await vi.waitFor(() => expect(chunks).toContain("good")); + expect(chunks).toEqual(["good"]); + await route.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("drops a no-hostRouteId message naming a concurrent route's session instead of cross-delivering it", async () => { + const handle = makeLoginPtyHandle(); + try { + await handle.start(); + // Two live routes share this worker. A worker session id alone cannot + // prove which route a hostRouteId-less message belongs to once a + // second route is live, so the fallback must refuse to guess. + const first = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A" }), + ); + const second = await handle.openLoginPtySession( + ptyOpenInput({ + workerSessionId: "ws-B", + outputs: [{ chunk: "cross-route", sid: "ws-A", omitHostRouteId: true }], + }), + ); + const firstChunks: string[] = []; + const secondChunks: string[] = []; + first.onData((chunk) => firstChunks.push(chunk)); + second.onData((chunk) => secondChunks.push(chunk)); + // Give the ambiguous notification time to arrive and be dropped. + await new Promise((resolve) => setTimeout(resolve, 60)); + expect(firstChunks).toEqual([]); + expect(secondChunks).toEqual([]); + + await first.close(); + await second.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); +}); + +// --------------------------------------------------------------------------- +// The process-wide login pseudo-terminal route ceiling. A host process can +// run only so many concurrent login pseudo-terminal routes before the +// underlying resources are exhausted, so the ceiling protects the host +// process itself, not one user or one worker. +// --------------------------------------------------------------------------- +// Every login pseudo-terminal route reserves one slot from the same +// process-wide aggregate route-slot controller the duplex channel route +// uses (`createDuplexRouteSlotController`), so the two route types share one +// ceiling. This is a host-process safety ceiling across every worker in the +// process, not a per-user or a per-worker quota. + +describe("plugin worker manager login pseudo-terminal route ceiling", () => { + it("rejects the second open with the fixed capacity error before the worker call, when the process-wide ceiling is full", async () => { + // A process-scoped ceiling of one slot. The manager injects one shared + // controller into every worker; the test injects a small one directly — + // the same controller shape the duplex channel route shares. + const handle = makeLoginPtyHandle({ + duplexRouteSlots: createDuplexRouteSlotController(1), + loginPtyLimits: { openTimeoutMs: 300 }, + }); + try { + await handle.start(); + const first = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A" }), + ); + const startedAt = Date.now(); + // The refused open scripts `no-open-reply`: if it wrongly reached the + // worker, the promise would only settle after the 300ms open timeout, + // and with a different, timeout-shaped error — not the capacity error. + // A fast rejection with the exact capacity error proves the worker + // never received this open's request. + await expect( + handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-B", mode: "no-open-reply" }), + ), + ).rejects.toThrow("LOGIN_PTY_ROUTES_AT_CAPACITY"); + expect(Date.now() - startedAt).toBeLessThan(150); + await first.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("admits a later open after the first route closes and releases its slot", async () => { + const handle = makeLoginPtyHandle({ + duplexRouteSlots: createDuplexRouteSlotController(1), + }); + try { + await handle.start(); + const first = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A" }), + ); + await expect( + handle.openLoginPtySession(ptyOpenInput({ workerSessionId: "ws-B" })), + ).rejects.toThrow("LOGIN_PTY_ROUTES_AT_CAPACITY"); + // Closing the first route releases its slot, so a later open is admitted. + await first.close(); + const second = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-C" }), + ); + await second.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); + + it("admits a later open on a different worker after a worker exit releases the shared slot", async () => { + // One shared controller instance, the same shape the manager injects into + // every worker handle, so both handles below draw from the SAME ceiling. + const sharedSlots = createDuplexRouteSlotController(1); + const handle = makeLoginPtyHandle({ duplexRouteSlots: sharedSlots }); + const secondHandle = makeLoginPtyHandle({ duplexRouteSlots: sharedSlots }); + try { + await handle.start(); + const first = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A" }), + ); + const waitResult = first.wait(); + await handle.stop(); + // The worker exit settles the route and releases its slot. + await expect(waitResult).resolves.toEqual({ exitCode: null }); + // Before the release, a second worker's open against the same shared + // ceiling would have rejected with the capacity error. After the + // release, the shared ceiling of one admits a fresh open on a + // different worker. + await secondHandle.start(); + const second = await secondHandle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-B" }), + ); + await second.close(); + } finally { + await handle.stop().catch(() => undefined); + await secondHandle.stop().catch(() => undefined); + } + }); + + it("releases exactly one slot when a route's terminal exit is followed by an explicit close", async () => { + // A double release (once on the terminal exit, once on the later close) + // would let a third open through a ceiling of one before its true + // capacity. This test proves the ceiling still holds after both events. + const handle = makeLoginPtyHandle({ + duplexRouteSlots: createDuplexRouteSlotController(1), + }); + try { + await handle.start(); + const first = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-A", exitCode: 0 }), + ); + await expect(first.wait()).resolves.toEqual({ exitCode: 0 }); + // The terminal exit already released the one slot. A second open now + // succeeds, consuming that same slot. + const second = await handle.openLoginPtySession( + ptyOpenInput({ workerSessionId: "ws-B" }), + ); + // The explicit close on the first (already-exited) route must not + // release a second slot it no longer holds. + await first.close(); + await expect( + handle.openLoginPtySession(ptyOpenInput({ workerSessionId: "ws-C" })), + ).rejects.toThrow("LOGIN_PTY_ROUTES_AT_CAPACITY"); + await second.close(); + } finally { + await handle.stop().catch(() => undefined); + } + }); +}); + +// --------------------------------------------------------------------------- +// Login pseudo-terminal pre-bind queue // --------------------------------------------------------------------------- // The host reads the worker pipe and `readline` dispatches every line of one // chunk synchronously. The route only becomes `open` inside the `await` diff --git a/server/src/routes/agents.ts b/server/src/routes/agents.ts index 5ee35842b9..2455b1ca1c 100644 --- a/server/src/routes/agents.ts +++ b/server/src/routes/agents.ts @@ -145,6 +145,7 @@ import { SetupTokenSessionError, assessConfidentialStartup, evaluateConfidentialTransport, + isTerminalSessionState, SETUP_TOKEN_START_FAILED, SETUP_TOKEN_SESSION_NOT_FOUND, SETUP_TOKEN_PROVIDER_UNSUPPORTED, @@ -571,6 +572,28 @@ export function agentRoutes( row.boundAt = Date.now(); return { ...row }; }, + async cancelDurable(identity, cancellableStates): Promise { + const row = setupTokenCleanupRows.get(identity.sessionId); + if (!row || !scopeMatchesRow(row, identity) || !cancellableStates.includes(row.state)) { + return null; + } + row.state = "cancelled"; + return { ...row }; + }, + async findActiveDurable(key, now): Promise { + for (const row of setupTokenCleanupRows.values()) { + if ( + row.companyId === key.companyId && + row.ownerUserId === key.ownerUserId && + row.adapterType === key.adapterType && + !isTerminalSessionState(row.state) && + row.deadline > now + ) { + return { ...row }; + } + } + return null; + }, }; const deferredSetupTokenLoginFactory: SetupTokenLoginProcessFactory = () => { @@ -3328,8 +3351,37 @@ export function agentRoutes( }, ); + // Read the caller's active login session for one company and adapter, with no + // session id. The browser rediscovers its own session after a reload with no + // local state. A non-owner, a foreign company, an unknown adapter, and no + // active session all receive the same 404. + // + // This route registers before the `:sessionId` route below, so Express + // never matches the literal `active` segment as a session id. + router.get( + "/companies/:companyId/adapters/:type/login-sessions/active", + async (req, res) => { + const companyId = req.params.companyId as string; + const type = req.params.type as string; + const ownerUserId = await assertCanManageAdapterLogin(req, companyId); + assertDeviceLoginAdapter(type); + res.setHeader("Cache-Control", "no-store, private"); + + const owner = await adapterLoginService.readActiveOwnerSession(companyId, type, ownerUserId); + if (!owner) { + res.status(404).json({ error: "Adapter login session not found" }); + return; + } + res.json(owner); + }, + ); + // Read a login session. The owner receives the status and the one-time prompt. - // A non-owner or a cross-company caller receives a 404. + // A non-owner or a cross-company caller receives a 404. The owner response + // repeats the live prompt on every read while the session holds an active + // public status, so this sets the same private no-store policy as the + // active-session route above: a shared or a browser cache never stores the + // authenticated response between one poll and the next. router.get( "/companies/:companyId/adapters/:type/login-sessions/:sessionId", async (req, res) => { @@ -3338,6 +3390,7 @@ export function agentRoutes( const sessionId = req.params.sessionId as string; const ownerUserId = await assertCanManageAdapterLogin(req, companyId); assertDeviceLoginAdapter(type); + res.setHeader("Cache-Control", "no-store, private"); const owner = await readOwnerLoginSession(companyId, type, sessionId, ownerUserId); if (!owner) { @@ -5662,6 +5715,41 @@ export function agentRoutes( } }); + // Read the caller's active Claude setup-token login session, with no session + // id. The browser rediscovers its own session after a reload with no local + // state. The response carries the panel mode and the one-time prompt, the + // same owner response shape the start route returns. A caller with no active + // session receives the same fixed not-found error as a foreign session. + // + // This route registers before the `:sessionId` route below, so Express never + // matches the literal `active` segment as a session id. + router.get("/companies/:companyId/setup-token-login-sessions/active", async (req, res) => { + const companyId = req.params.companyId as string; + const ownerUserId = resolveCompanySessionOwner(req, companyId, res); + if (ownerUserId === null) return; + res.setHeader("Cache-Control", "no-store, private"); + const descriptor = await setupTokenLoginService.findActiveByScope( + companySetupTokenKey(companyId, ownerUserId), + ); + if (!descriptor) { + res.status(404).json({ error: SETUP_TOKEN_SESSION_NOT_FOUND }); + return; + } + // Read the panel mode from the adapter capability, the same way the start + // route does. The full login URL rides in this response, guarded by the + // same transport advisory the prompt route attaches. + const panelMode = + getRegistryLoginCapability(SETUP_TOKEN_ADAPTER_TYPE)?.panelMode ?? "submitted_browser_code"; + const body: ClaudeSetupTokenSessionOwnerResponse = { + ...toClaudePublicResponse(descriptor), + panelMode, + prompt: descriptor.loginUrl + ? { authorizationUrl: descriptor.loginUrl, transportAdvisory: assessSetupTokenTransport(req) } + : null, + }; + res.json(body); + }); + router.get("/companies/:companyId/setup-token-login-sessions/:sessionId", async (req, res) => { const companyId = req.params.companyId as string; const ownerUserId = resolveCompanySessionOwner(req, companyId, res); @@ -5779,11 +5867,14 @@ export function agentRoutes( res.setHeader("Cache-Control", "no-store"); const sessionId = req.params.sessionId as string; try { - const scope = setupTokenLoginService.resolveCompanyScope( + // `cancelByScope` tries the live in-memory session first, then falls back + // to a durable-only cancel when no live session matches — for example, + // after a restart drops the in-memory session but the durable row still + // holds the company slot. + await setupTokenLoginService.cancelByScope( sessionId, companySetupTokenKey(companyId, ownerUserId), ); - await setupTokenLoginService.cancel(sessionId, scope); res.status(200).json({}); } catch (err) { // Cancel is idempotent. The service removes a session when it reaches a diff --git a/server/src/routes/openapi.ts b/server/src/routes/openapi.ts index 90bb0ba9a2..dda1c36904 100644 --- a/server/src/routes/openapi.ts +++ b/server/src/routes/openapi.ts @@ -2265,6 +2265,17 @@ registry.registerPath({ }, }); +registry.registerPath({ + method: "get", + path: "/api/companies/{companyId}/adapters/{type}/login-sessions/active", + tags: ["adapters"], + summary: "Read the caller's active adapter device login session", + request: { + params: z.object({ companyId: z.string(), type: z.string() }), + }, + responses: { 200: r.ok(), 401: r.unauthorized, 403: r.forbidden, 404: r.notFound }, +}); + registry.registerPath({ method: "get", path: "/api/companies/{companyId}/adapters/{type}/login-sessions/{sessionId}", @@ -4957,6 +4968,20 @@ registry.registerPath({ }, }); +registry.registerPath({ + method: "get", + path: "/api/companies/{companyId}/setup-token-login-sessions/active", + tags: ["companies"], + summary: "Read the caller's active Claude setup-token login session", + request: { params: z.object({ companyId: z.string() }) }, + responses: { + 200: r.ok(claudeSetupTokenSessionOwnerResponseSchema), + 401: r.unauthorized, + 403: r.forbidden, + 404: r.notFound, + }, +}); + registry.registerPath({ method: "get", path: "/api/companies/{companyId}/setup-token-login-sessions/{sessionId}", diff --git a/server/src/routes/setup-token-route.test.ts b/server/src/routes/setup-token-route.test.ts index 09c583021f..aaa448f902 100644 --- a/server/src/routes/setup-token-route.test.ts +++ b/server/src/routes/setup-token-route.test.ts @@ -24,6 +24,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { SETUP_TOKEN_SESSION_NOT_FOUND, SETUP_TOKEN_PROVIDER_UNSUPPORTED, + isTerminalSessionState, type SetupTokenCleanupRecord, type SetupTokenCleanupStore, type SetupTokenLeaseManager, @@ -289,6 +290,26 @@ function buildTransport(opts: { onSubmit?: "complete" | "throw" | "pending" } = row.boundAt = Date.now(); return { ...row }; }, + async cancelDurable(identity, cancellableStates) { + const row = rows.get(identity.sessionId); + if (!row || !cancellableStates.includes(row.state)) return null; + row.state = "cancelled"; + return { ...row }; + }, + async findActiveDurable(key, now) { + for (const row of rows.values()) { + if ( + row.companyId === key.companyId && + row.ownerUserId === key.ownerUserId && + row.adapterType === key.adapterType && + !isTerminalSessionState(row.state) && + row.deadline > now + ) { + return { ...row }; + } + } + return null; + }, }; const leases: SetupTokenLeaseManager = { async acquire() { @@ -639,6 +660,50 @@ describe("company-and-environment setup-token route — object-level authorizati expect(res.body.error).toBe(SETUP_TOKEN_SESSION_NOT_FOUND); }); + it("returns the caller's active session with no session id in the URL", async () => { + const transport = buildTransport({ onSubmit: "pending" }); + const { app } = await createApp({ transport }); + + const startRes = await startCompanySession(app); + const sessionId = startRes.body.sessionId as string; + + const active = await request(app).get(`${COMPANY_BASE}/active`).send(); + expect(active.status, JSON.stringify(active.body)).toBe(200); + expect(active.body.sessionId).toBe(sessionId); + expect(active.body.status).toBe("waiting_for_user"); + expect(active.body.panelMode).toBe("submitted_browser_code"); + // The full login URL rides in this owner response, the same way it rides in + // the guarded `.../prompt` response — this is not a public, secret-free + // surface, so it is not checked against `expectNoSecret`. + expect(active.body.prompt).toEqual({ authorizationUrl: FULL_LOGIN_URL, transportAdvisory: null }); + expect(active.headers["cache-control"]).toBe("no-store, private"); + }); + + it("returns the identical not-found on the active route for no active session, a cross-owner caller, and a cross-company caller", async () => { + const transport = buildTransport({ onSubmit: "pending" }); + const { app } = await createApp({ transport }); + + // No session has started yet. + const none = await request(app).get(`${COMPANY_BASE}/active`).send(); + expect(none.status).toBe(404); + expect(none.body.error).toBe(SETUP_TOKEN_SESSION_NOT_FOUND); + + await startCompanySession(app); + + // A different board user in the same company holds no active session. + useOwner(OTHER_USER_ID); + const otherOwner = await request(app).get(`${COMPANY_BASE}/active`).send(); + expect(otherOwner.status).toBe(404); + expect(otherOwner.body.error).toBe(SETUP_TOKEN_SESSION_NOT_FOUND); + useOwner(); + + // The same owner under a different company holds no active session there. + const crossCompanyBase = `/api/companies/${OTHER_COMPANY_ID}/setup-token-login-sessions`; + const otherCompany = await request(app).get(`${crossCompanyBase}/active`).send(); + expect(otherCompany.status).toBe(404); + expect(otherCompany.body.error).toBe(SETUP_TOKEN_SESSION_NOT_FOUND); + }); + it("returns the fixed not-found for a non-member on every action across a company boundary", async () => { const transport = buildTransport({ onSubmit: "complete" }); const { app } = await createApp({ transport }); diff --git a/server/src/services/device-login-service.ts b/server/src/services/device-login-service.ts index 4e722319ea..5e7453038c 100644 --- a/server/src/services/device-login-service.ts +++ b/server/src/services/device-login-service.ts @@ -342,6 +342,17 @@ export interface AdapterAuthSessionStore { companyId: string, adapterType?: AgentAdapterType, ): Promise; + /** + * Read the active session for one company, owner, and adapter, with no + * session id. The predicate matches the same three columns as the active-slot + * partial unique index, so at most one row can match. It returns null when the + * owner holds no active session for the adapter. + */ + getActiveByOwner( + companyId: string, + startedByUserId: string, + adapterType: AgentAdapterType, + ): Promise; /** Run `fn` while the process holds the promotion critical-section lock for the * company, owner, and adapter slot. The reaper reclaims a stale `promoting` * row inside this lock, so a reclaim never interleaves with a live credential @@ -585,6 +596,24 @@ export function createDbAdapterAuthSessionStore( const row = rows[0]; return row ? toRow(row) : null; }, + async getActiveByOwner(companyId, startedByUserId, adapterType) { + // The predicate matches the same three columns as the active-slot partial + // unique index, plus the active-status set, so at most one row can match. + const rows = await db + .select() + .from(adapterAuthSessions) + .where( + and( + eq(adapterAuthSessions.companyId, companyId), + eq(adapterAuthSessions.startedByUserId, startedByUserId), + eq(adapterAuthSessions.adapterType, adapterType), + inArray(adapterAuthSessions.status, [...ADAPTER_AUTH_ACTIVE_STATUSES]), + ), + ) + .limit(1); + const row = rows[0]; + return row ? toRow(row) : null; + }, async listExpiredActiveSessions(nowAt) { // The partial index on the active statuses and the index on `expiresAt` // both support this scan. The scan is bounded by the active-status set, so @@ -1159,6 +1188,9 @@ export function createDeviceLoginService(deps: DeviceLoginServiceDeps) { activity: (phase: LoginSessionActivityPhase) => void; }): Promise { const { sessionId, lease, activity } = ctx; + // The reaper already reclaimed the row, so this process no longer owns the + // slot. Delete a prompt this process may still hold for it. + promptsBySession.delete(sessionId); const observation = await observeSandboxDelete(() => lease.deleteSandbox()); if (observation.confirmed) { activity("sandbox_deleted"); @@ -1193,6 +1225,11 @@ export function createDeviceLoginService(deps: DeviceLoginServiceDeps) { const { sessionId, lease, terminal, reason, expectedStatuses, conditionalTransition, activity } = ctx; + // Delete the in-memory prompt on every terminal transition this function + // handles: an authenticated success, a failure, a timeout, and a + // cancellation. A terminal response then always carries a null prompt. + promptsBySession.delete(sessionId); + // The cleanup-state handoff. The service owns and observes the provider // delete on every terminal path. The reaper path shares the same delete // observation and the same durable-write choice. @@ -1233,6 +1270,30 @@ export function createDeviceLoginService(deps: DeviceLoginServiceDeps) { }; } + // Build the owner response for a row. The prompt survives every read while + // the session holds an active public status: the read never deletes it. A + // terminal transition deletes the prompt at its own write site instead (see + // `terminate`, `abandonAfterLostClaim`, and `cancelOwnerSession`), so a + // terminal response always carries a null prompt. + function buildOwnerResponse( + row: AdapterAuthSessionRow, + requestingUserId: string, + ): AdapterAuthSessionOwnerResponse { + const isOwner = row.startedByUserId === requestingUserId; + const status = resolvePublicStatus(row); + // The prompt map keys on the internal id, so read it by `row.id`, not the + // public id. Only the owner principal ever reads the prompt. + const prompt = isOwner ? promptsBySession.get(row.id) ?? null : null; + return { + sessionId: row.publicSessionId, + environmentId: row.environmentId, + status, + expiresAt: row.expiresAt?.toISOString() ?? null, + failure: buildFailure(row, status), + prompt: prompt ? { url: prompt.url, code: prompt.code } : null, + }; + } + async function readOwnerSession( publicSessionId: string, companyId: string, @@ -1242,26 +1303,21 @@ export function createDeviceLoginService(deps: DeviceLoginServiceDeps) { // foreign-company caller reads nothing, and the internal id never matches. const row = await store.getByPublicId(publicSessionId, companyId); if (!row) return null; - const isOwner = row.startedByUserId === requestingUserId; - const status = resolvePublicStatus(row); - // Deliver the one-time prompt to the owner principal exactly once. The read - // and the delete run with no await between them, so the first authorized - // owner read consumes the prompt and every later read returns null. This - // keeps the short-lived device code out of a repeated response. The prompt - // map keys on the internal id, so read it by `row.id`, not the public id. - let prompt: DeviceLoginPrompt | null = null; - if (isOwner) { - prompt = promptsBySession.get(row.id) ?? null; - if (prompt) promptsBySession.delete(row.id); - } - return { - sessionId: row.publicSessionId, - environmentId: row.environmentId, - status, - expiresAt: row.expiresAt?.toISOString() ?? null, - failure: buildFailure(row, status), - prompt: prompt ? { url: prompt.url, code: prompt.code } : null, - }; + return buildOwnerResponse(row, requestingUserId); + } + + // Read the caller's active session for one company and adapter, with no + // session id. The browser rediscovers its own session after a reload with no + // local state. The store predicate matches only an active row for this exact + // owner, so this never surfaces a foreign owner's session. + async function readActiveOwnerSession( + companyId: string, + adapterType: AgentAdapterType, + requestingUserId: string, + ): Promise { + const row = await store.getActiveByOwner(companyId, requestingUserId, adapterType); + if (!row) return null; + return buildOwnerResponse(row, requestingUserId); } // Cancel a login session for its owner. The write is durable, so a cancel @@ -1296,10 +1352,14 @@ export function createDeviceLoginService(deps: DeviceLoginServiceDeps) { finishedAt: now(), promotionExpiresAt: null, }); + // Delete the in-memory prompt on every cancel, whether or not the write + // above won. A lost write means the row already left the cancellable + // states, so no fresh prompt exists for it either. + promptsBySession.delete(row.id); return readOwnerSession(publicSessionId, companyId, requestingUserId); } - return { start, readOwnerSession, cancelOwnerSession }; + return { start, readOwnerSession, readActiveOwnerSession, cancelOwnerSession }; } export type DeviceLoginService = ReturnType; diff --git a/server/src/services/plugin-worker-manager.ts b/server/src/services/plugin-worker-manager.ts index 459ceb585e..bdcf10cb8b 100644 --- a/server/src/services/plugin-worker-manager.ts +++ b/server/src/services/plugin-worker-manager.ts @@ -163,10 +163,18 @@ const LOGIN_PTY_CLOSE_TIMEOUT_MS = 10_000; * select an arbitrary command in the sandbox. */ const LOGIN_PTY_COMMAND_NOT_ALLOWED = "LOGIN_PTY_COMMAND_NOT_ALLOWED"; -/** The fixed non-secret error a rejected second credential open returns. */ -const LOGIN_PTY_ROUTE_BUSY = "LOGIN_PTY_ROUTE_BUSY"; /** The fixed non-secret error a failed open returns. */ const LOGIN_PTY_OPEN_FAILED = "LOGIN_PTY_OPEN_FAILED"; +/** + * The fixed non-secret error a login pseudo-terminal open returns when the + * process-wide aggregate route-slot ceiling is full. This is a distinct + * condition from the removed per-worker single-route gate: it means the whole + * process is at capacity, across every worker, not that one worker already + * holds a terminal. It is also distinct from {@link DUPLEX_CHANNEL_ROUTE_BUSY}: + * `packages/adapter-utils/src/execution-target.ts` matches that exact text as + * a marker for the duplex path, and a login failure must never enter it. + */ +const LOGIN_PTY_ROUTES_AT_CAPACITY = "LOGIN_PTY_ROUTES_AT_CAPACITY"; // Bounds and timeouts for the generic duplex channel route. The route mirrors the // login pseudo-terminal route, but it carries no command allowlist and adds seven @@ -1317,18 +1325,26 @@ export function createPluginWorkerHandle( // ----------------------------------------------------------------------- // Host-owned login pseudo-terminal route gate // ----------------------------------------------------------------------- - // The manager owns one live login pseudo-terminal route per worker. It mints a - // host-owned opaque route identifier, carries it in the open call, and keys the - // close on it, so it closes a worker-created terminal even when the open reply - // was lost and no worker session identifier arrived. It binds the worker - // session identifier one time while the route is `opening`, for output only. It - // never trusts a worker-supplied identifier as proof of origin: it delivers - // output only while the route is `open` and the notification carries the exact - // bound identifier and valid bounded bytes, and it never logs the raw bytes. It - // terminalizes the route exactly once on every open failure path, closes the - // terminal by the host route identifier, and admits a new open only after it - // verifies a close acknowledgement bound to that identifier; it retires the - // worker on an unconfirmed close. + // The manager owns every live login pseudo-terminal route on a worker. One + // worker admits more than one concurrent route, so more than one owner can + // hold an active credential login on the same shared worker at once. The + // manager mints a host-owned opaque route identifier for each open, carries + // it in the open call, and keys the close on it, so it closes a + // worker-created terminal even when the open reply was lost and no worker + // session identifier arrived. It binds the worker session identifier one + // time while a route is `opening`, for output only. It never trusts a + // worker-supplied identifier as proof of origin: for each output or exit + // notification, it resolves the route by the host route identifier first, + // then delivers only while that route is `open` and the notification + // carries the exact bound worker session identifier and valid bounded + // bytes; it drops an unknown, a stale, a duplicate, a malformed, or a + // mismatched notification, and it never logs the raw bytes. It terminalizes + // one route exactly once on every open failure path, closes the terminal by + // its host route identifier, and admits a new open on that same identifier + // only after it verifies a close acknowledgement bound to it. On an + // unconfirmed close it retires the whole worker, which settles and clears + // every route the worker still holds — a possibly live terminal must never + // reach reuse or a wrong delivery. // A single-consumer route state. The login pseudo-terminal route and the // generic duplex channel route share it. @@ -1399,9 +1415,47 @@ export function createPluginWorkerHandle( // record. preBindChars: number; } - // At most one active credential pseudo-terminal per worker. A non-null route - // blocks a second open until the manager confirms the first route's close. - let loginPtyRoute: LoginPtyRoute | null = null; + // The live login pseudo-terminal routes on this worker, keyed by the host + // route id. A route lives here from reservation — before the open call — + // until it terminalizes. An output or an exit notification resolves its + // route through this map first, so more than one route can be `reserved`, + // `opening`, or `open` on one worker at once. + const loginPtyRoutesByHostRouteId = new Map(); + // True once this worker has logged the missing-`hostRouteId` warning below. + // A plugin build old enough to omit the field would otherwise log one line + // for every dropped message, and a login pseudo-terminal can carry a high + // volume of output notifications. One warning per worker is enough for an + // operator to see the cause. + let loggedMissingLoginPtyHostRouteId = false; + // The bound routes on this worker, keyed by the worker session id. A route + // enters this map once its worker session id binds and leaves it when the + // route terminalizes. The host checks this map at bind time, so one worker + // session id can never bind to two live routes at once. + const loginPtyRoutesByWorkerSessionId = new Map(); + + // The routes that currently hold one process-wide aggregate route slot. The + // host releases a slot one time per route, so a double terminalize, or a + // terminal exit followed by a later close, never releases two slots. + const loginPtyRouteSlotHolders = new Set(); + + // Try to reserve one aggregate route slot for a route. Return true when the + // route holds a slot after the call. When no controller is present, the + // route always holds a slot. This calls the SAME shared controller instance + // the duplex channel route uses (`duplexRouteSlots`), so the two route types + // share one process-wide ceiling. + function acquireLoginPtyRouteSlot(route: LoginPtyRoute): boolean { + if (!duplexRouteSlots) return true; + if (!duplexRouteSlots.tryAcquire()) return false; + loginPtyRouteSlotHolders.add(route); + return true; + } + + // Release the aggregate route slot a route holds, one time. A route that + // never held a slot, or already released it, releases nothing. + function releaseLoginPtyRouteSlot(route: LoginPtyRoute): void { + if (!loginPtyRouteSlotHolders.delete(route)) return; + duplexRouteSlots?.release(); + } // Close the worker terminal by the host route identifier and verify the bound // acknowledgement. Return true only when the worker returns an acknowledgement @@ -1421,9 +1475,13 @@ export function createPluginWorkerHandle( } } - // Terminalize the route exactly once. Resolve the login wait, close the worker - // terminal by the host route identifier, and free the per-worker slot only - // after the close resolves. Retire the worker when the close is unconfirmed. + // Terminalize one route exactly once. Remove it from both maps at once, + // before the worker close call, so no later notification for its host route + // identifier or its worker session identifier can still resolve to it. + // Resolve the login wait, close the worker terminal by the host route + // identifier, and retire the whole worker when the close is unconfirmed — + // every other route this worker still holds settles through the worker-exit + // path that follows. async function terminalizeLoginPtyRoute(route: LoginPtyRoute): Promise { if (route.terminalized) return; route.terminalized = true; @@ -1433,14 +1491,19 @@ export function createPluginWorkerHandle( // A terminalized route never replays a queued pre-bind record. route.preBind = []; route.preBindChars = 0; + loginPtyRoutesByHostRouteId.delete(route.hostRouteId); + if (route.workerSessionId !== null) { + loginPtyRoutesByWorkerSessionId.delete(route.workerSessionId); + } + releaseLoginPtyRouteSlot(route); // A terminalized route reports a null exit code, which the runner treats as a // failure. settleRouteWait(route, { exitCode: null }); const confirmed = await closeLoginPtyTerminal(route.hostRouteId); - if (loginPtyRoute === route) loginPtyRoute = null; if (!confirmed) { // The worker did not acknowledge the close, so the host cannot prove the - // terminal is gone. Fail closed: retire the worker before any reuse. + // terminal is gone. Fail closed: retire the worker before any reuse. This + // also settles and clears every other route the worker still holds. log.error( { pluginId }, "login pseudo-terminal close not acknowledged; retiring worker", @@ -1449,20 +1512,61 @@ export function createPluginWorkerHandle( } } + // Resolve one login pseudo-terminal notification to its route. Return null + // for an unknown or a stale (already terminalized) identifier, so the caller + // drops the notification. + // + // The current protocol tags every notification with the host route + // identifier, so a current build resolves through + // `loginPtyRoutesByHostRouteId` directly, even while more than one route is + // open on this worker. + // + // A plugin build old enough to predate the host route identifier tags the + // notification with only the worker session identifier, the sole routing + // key the previous protocol used. The worker itself controls that + // identifier, so it alone cannot prove which route a message belongs to + // once this worker holds two or more concurrent routes: a message that + // omits `hostRouteId` and names a different, live route's session + // identifier would otherwise deliver into, or end, that other route. A + // build old enough to omit `hostRouteId` never ran a concurrent route, so + // the fallback only trusts the worker session identifier while this worker + // holds exactly one route. The host warns once for this worker — never + // once for each dropped message, since a live pseudo-terminal can send + // many — so an operator can see the build is too old, then still delivers + // the route's output and exit notifications instead of losing them. + function resolveLoginPtyRoute(params: Record): LoginPtyRoute | null { + const hostRouteId = readNonEmptyString(params.hostRouteId); + if (hostRouteId) { + return loginPtyRoutesByHostRouteId.get(hostRouteId) ?? null; + } + if (!loggedMissingLoginPtyHostRouteId) { + loggedMissingLoginPtyHostRouteId = true; + log.warn( + "login pseudo-terminal message has no hostRouteId; the plugin build is too old", + ); + } + if (loginPtyRoutesByHostRouteId.size !== 1) return null; + const workerSessionId = readNonEmptyString(params.workerSessionId); + if (!workerSessionId) return null; + return loginPtyRoutesByWorkerSessionId.get(workerSessionId) ?? null; + } + // Route one login pseudo-terminal output notification to the per-session - // listener. Deliver only while the route is `open` and the notification carries - // the exact bound worker session identifier and valid bounded bytes. Queue the - // notification while the route is still `opening`. Drop an unknown, late, - // malformed, or mismatched notification. Never log the raw bytes. + // listener. Resolve the route by the host route identifier first. Deliver + // only while that route is `open` and the notification carries the exact + // bound worker session identifier and valid bounded bytes. Queue the + // notification while the route is still `opening`. Drop an unknown, a stale, + // a duplicate, a malformed, or a mismatched notification. Never log the raw + // bytes. function routeLoginPtyOutput(notification: JsonRpcNotification): void { - const route = loginPtyRoute; + const params = isRecord(notification.params) ? notification.params : {}; + const route = resolveLoginPtyRoute(params); if (!route) return; if (route.state === "opening") { queuePreBindLoginPtyOutput(route, notification); return; } if (route.state !== "open") return; - const params = isRecord(notification.params) ? notification.params : {}; const workerSessionId = readNonEmptyString(params.workerSessionId); if (!workerSessionId || workerSessionId !== route.workerSessionId) return; const chunk = params.chunk; @@ -1483,24 +1587,28 @@ export function createPluginWorkerHandle( else route.buffered.push(chunk); } - // Route one login pseudo-terminal exit notification to the login wait. Resolve - // only while the route is `open` and the notification carries the exact bound - // worker session identifier. Queue the notification while the route is still - // `opening`. A resolved exit moves the state off `open`, so a later record — - // live or replayed — finds a closed route and drops there. + // Route one login pseudo-terminal exit notification to the login wait. + // Resolve the route by the host route identifier first. Settle the wait + // only while that route is `open` and the notification carries the exact + // bound worker session identifier. Queue the notification while the route + // is still `opening`. A resolved exit moves the state off `open`, so a + // later record — live or replayed — finds a closed route and drops there. function routeLoginPtyExit(notification: JsonRpcNotification): void { - const route = loginPtyRoute; + const params = isRecord(notification.params) ? notification.params : {}; + const route = resolveLoginPtyRoute(params); if (!route) return; if (route.state === "opening") { queuePreBindLoginPtyExit(route, notification); return; } if (route.state !== "open") return; - const params = isRecord(notification.params) ? notification.params : {}; const workerSessionId = readNonEmptyString(params.workerSessionId); if (!workerSessionId || workerSessionId !== route.workerSessionId) return; const exitCode = typeof params.exitCode === "number" ? params.exitCode : null; route.state = "closed"; + // A terminal exit is its own slot-release path, distinct from the later + // explicit close, so a route that already exited returns its slot at once. + releaseLoginPtyRouteSlot(route); settleRouteWait(route, { exitCode }); } @@ -1583,38 +1691,51 @@ export function createPluginWorkerHandle( routeLoginPtyOutput({ jsonrpc: "2.0", method: LOGIN_PTY_OUTPUT_NOTIFICATION, - params: { workerSessionId: record.workerSessionId, chunk: record.chunk }, + params: { + hostRouteId: route.hostRouteId, + workerSessionId: record.workerSessionId, + chunk: record.chunk, + }, }); } else { routeLoginPtyExit({ jsonrpc: "2.0", method: LOGIN_PTY_EXIT_NOTIFICATION, - params: { workerSessionId: record.workerSessionId, exitCode: record.exitCode }, + params: { + hostRouteId: route.hostRouteId, + workerSessionId: record.workerSessionId, + exitCode: record.exitCode, + }, }); } } } - // Close the one route on a worker exit. The worker is gone, so the manager - // resolves the login wait with the fixed non-secret exit and clears the route - // one time. The pending pseudo-terminal calls reject through `rejectAllPending`. + // Close every route on a worker exit. The worker is gone, so the manager + // resolves each login wait with the fixed non-secret exit and clears both + // maps one time. The pending pseudo-terminal calls reject through + // `rejectAllPending`. function closeLoginPtyRouteOnWorkerExit(): void { - const route = loginPtyRoute; - if (!route) return; - loginPtyRoute = null; - route.terminalized = true; - route.state = "closed"; - route.listener = null; - route.buffered = []; - route.preBind = []; - route.preBindChars = 0; - settleRouteWait(route, { exitCode: null }); + const routes = [...loginPtyRoutesByHostRouteId.values()]; + loginPtyRoutesByHostRouteId.clear(); + loginPtyRoutesByWorkerSessionId.clear(); + for (const route of routes) { + route.terminalized = true; + route.state = "closed"; + route.listener = null; + route.buffered = []; + route.preBind = []; + route.preBindChars = 0; + releaseLoginPtyRouteSlot(route); + settleRouteWait(route, { exitCode: null }); + } } - // Open one live login pseudo-terminal route. Reserve the route - // before the open call, bind the worker session identifier one time on the - // first successful open reply, and return a session the login transport drives. - // Terminalize the route on every open failure path. + // Open one live login pseudo-terminal route. Reserve one process-wide + // aggregate route slot before the open call, bind the worker session + // identifier one time on the first successful open reply, and return a + // session the login transport drives. Terminalize the route on every open + // failure path. One worker can hold more than one concurrent route. async function openLoginPtySession( input: LoginPtyOpenInput, ): Promise { @@ -1629,11 +1750,6 @@ export function createPluginWorkerHandle( // The binding validated it at the service boundary; this is the last host gate // before the worker call, so a malformed home fails closed here too. validateLoginSessionHome(input.sessionHome); - if (loginPtyRoute) { - // A route for this worker is not yet closed and confirmed. Reject the - // second open with one fixed non-secret error before it reaches the worker. - throw new Error(LOGIN_PTY_ROUTE_BUSY); - } const hostRouteId = randomUUID(); let settleWait: (value: { exitCode: number | null }) => void = () => {}; const waitPromise = new Promise<{ exitCode: number | null }>((resolve) => { @@ -1651,7 +1767,18 @@ export function createPluginWorkerHandle( preBind: [], preBindChars: 0, }; - loginPtyRoute = route; + // Reserve one process-wide aggregate route slot before any work. When the + // ceiling is full, reject with the fixed capacity error and open nothing, + // so an active login route never downgrades and the ceiling never + // overcommits. This never reveals the live count, the ceiling, or any + // other tenant. + if (!acquireLoginPtyRouteSlot(route)) { + throw new Error(LOGIN_PTY_ROUTES_AT_CAPACITY); + } + // Reserve the route by its host route identifier before the open call, so + // a notification that echoes this identifier can queue against it even + // before the worker replies. + loginPtyRoutesByHostRouteId.set(hostRouteId, route); route.state = "opening"; let openResult: HostToWorkerMethods["loginPtyOpen"][1]; @@ -1683,9 +1810,16 @@ export function createPluginWorkerHandle( await terminalizeLoginPtyRoute(route); throw new Error(LOGIN_PTY_OPEN_FAILED); } + if (loginPtyRoutesByWorkerSessionId.has(workerSessionId)) { + // A live route already owns this worker session identifier. Fail closed + // instead of binding a second route to it. + await terminalizeLoginPtyRoute(route); + throw new Error(LOGIN_PTY_OPEN_FAILED); + } // Bind the worker session identifier one time and move the route to `open`. route.workerSessionId = workerSessionId; route.state = "open"; + loginPtyRoutesByWorkerSessionId.set(workerSessionId, route); // Replay every record the route queued before the bind, in arrival order. replayPreBindLoginPtyRecords(route); diff --git a/server/src/services/setup-token-session.test.ts b/server/src/services/setup-token-session.test.ts index 5be8887aa7..9b85419fb1 100644 --- a/server/src/services/setup-token-session.test.ts +++ b/server/src/services/setup-token-session.test.ts @@ -23,6 +23,7 @@ import { SETUP_TOKEN_CAP_EXCEEDED, SETUP_TOKEN_TOKEN_UNAVAILABLE, SETUP_TOKEN_STORAGE_FAILED, + SETUP_TOKEN_CANCELLABLE_STATES, type SetupTokenCleanupIdentity, type SetupTokenCleanupRecord, type SetupTokenCleanupStore, @@ -36,6 +37,7 @@ import { type SetupTokenRateLimiter, type SetupTokenSecretWriter, type SetupTokenSessionScope, + type SetupTokenSessionState, } from "./setup-token-session.js"; import { redactSensitive } from "../middleware/redact-sensitive.js"; import { sanitizeRecord } from "../redaction.js"; @@ -217,6 +219,34 @@ class FakeStore implements SetupTokenCleanupStore { row.boundAt = Date.now(); return { ...row }; } + async cancelDurable( + identity: SetupTokenCleanupIdentity, + cancellableStates: readonly SetupTokenSessionState[], + ): Promise { + const row = this.rows.get(identity.sessionId); + if (!row || !identityMatchesRow(row, identity) || !cancellableStates.includes(row.state)) { + return null; + } + row.state = "cancelled"; + return { ...row }; + } + async findActiveDurable( + key: Pick, + now: number, + ): Promise { + for (const row of this.rows.values()) { + if ( + row.companyId === key.companyId && + row.ownerUserId === key.ownerUserId && + row.adapterType === key.adapterType && + !isTerminalSessionState(row.state) && + row.deadline > now + ) { + return { ...row }; + } + } + return null; + } } function allowAllRateLimiter(): SetupTokenRateLimiter { @@ -654,6 +684,141 @@ describe("SetupTokenSessionService durable reaper", () => { }); }); +describe("SetupTokenSessionService.cancelByScope", () => { + it("cancels the live in-memory session the normal way when one exists", async () => { + const { service, store } = buildService(); + const { sessionId } = await service.start(OWNER_SCOPE); + const result = await service.cancelByScope(sessionId, OWNER_SCOPE); + expect(result.state).toBe("cancelled"); + expect(store.rows.has(sessionId)).toBe(false); + }); + + it("releases the durable row when no live session exists (a restart dropped it)", async () => { + const store = new FakeStore(); + // Simulate a durable row a prior process created. No `service.start()` ever + // ran for it in THIS service instance, so `sessions` holds nothing for it — + // the shape a restart leaves behind. + await store.record({ + sessionId: "durable-only-1", + companyId: OWNER_SCOPE.companyId, + ownerUserId: OWNER_SCOPE.ownerUserId, + adapterType: OWNER_SCOPE.adapterType, + environmentId: OWNER_SCOPE.environmentId, + leaseId: "lease-durable-only-1", + deadline: Date.now() + 60_000, + state: "awaiting_code", + boundAt: null, + }); + const { service } = buildService({ store }); + + const result = await service.cancelByScope("durable-only-1", OWNER_SCOPE); + expect(result.state).toBe("cancelled"); + expect(store.rows.get("durable-only-1")?.state).toBe("cancelled"); + }); + + it("cancels only the caller's row: a foreign scope leaves a same-id row untouched and throws not-found", async () => { + const store = new FakeStore(); + await store.record({ + sessionId: "durable-only-2", + companyId: OWNER_SCOPE.companyId, + ownerUserId: OWNER_SCOPE.ownerUserId, + adapterType: OWNER_SCOPE.adapterType, + environmentId: OWNER_SCOPE.environmentId, + leaseId: "lease-durable-only-2", + deadline: Date.now() + 60_000, + state: "awaiting_code", + boundAt: null, + }); + const { service } = buildService({ store }); + + const foreignKey = { ...OWNER_SCOPE, ownerUserId: "user-2" }; + await expect(service.cancelByScope("durable-only-2", foreignKey)).rejects.toThrow( + SETUP_TOKEN_SESSION_NOT_FOUND, + ); + // The foreign-scope attempt never touched the real owner's row. + expect(store.rows.get("durable-only-2")?.state).toBe("awaiting_code"); + }); + + it("does not interrupt a session in the submitting (credential-write) state", async () => { + const store = new FakeStore(); + await store.record({ + sessionId: "durable-only-3", + companyId: OWNER_SCOPE.companyId, + ownerUserId: OWNER_SCOPE.ownerUserId, + adapterType: OWNER_SCOPE.adapterType, + environmentId: OWNER_SCOPE.environmentId, + leaseId: "lease-durable-only-3", + deadline: Date.now() + 60_000, + state: "submitting", + boundAt: null, + }); + const { service } = buildService({ store }); + + await expect(service.cancelByScope("durable-only-3", OWNER_SCOPE)).rejects.toThrow( + SETUP_TOKEN_SESSION_NOT_FOUND, + ); + // The row stays in `submitting`: the durable-only cancel never interrupts + // the credential write in progress. + expect(store.rows.get("durable-only-3")?.state).toBe("submitting"); + }); + + it("throws the fixed not-found error when nothing matches at all", async () => { + const { service } = buildService(); + await expect(service.cancelByScope("never-existed", OWNER_SCOPE)).rejects.toThrow( + SETUP_TOKEN_SESSION_NOT_FOUND, + ); + }); +}); + +describe("SetupTokenSessionService.findActiveByScope", () => { + it("finds the caller's live active session with no session id", async () => { + const { service } = buildService(); + const { sessionId } = await service.start(OWNER_SCOPE); + const found = await service.findActiveByScope(OWNER_SCOPE); + expect(found?.sessionId).toBe(sessionId); + }); + + it("returns null for another owner and for a terminal session", async () => { + const { service } = buildService(); + const { sessionId } = await service.start(OWNER_SCOPE); + expect( + await service.findActiveByScope({ ...OWNER_SCOPE, ownerUserId: "user-2" }), + ).toBeNull(); + + await service.cancel(sessionId, OWNER_SCOPE); + expect(await service.findActiveByScope(OWNER_SCOPE)).toBeNull(); + }); + + it("finds the durable active row after a restart drops the in-memory session", async () => { + const store = new FakeStore(); + const { service: before } = buildService({ store }); + const { sessionId } = await before.start(OWNER_SCOPE); + + // Simulate a restart: a fresh service shares the durable store but starts + // with an empty in-memory session map. + const { service: after } = buildService({ store }); + const found = await after.findActiveByScope(OWNER_SCOPE); + expect(found?.sessionId).toBe(sessionId); + // The full login URL lives only in memory (SR-5), so a restart-recovered + // descriptor never carries one. + expect(found?.loginUrl).toBeNull(); + }); + + it("does not find a foreign-scope or an expired durable row after a restart", async () => { + const store = new FakeStore(); + const { service: before } = buildService({ store, ttlMs: 1_000 }); + await before.start(OWNER_SCOPE); + + const { service: after } = buildService({ store }); + expect( + await after.findActiveByScope({ ...OWNER_SCOPE, ownerUserId: "user-2" }), + ).toBeNull(); + + const { service: expiredAfter } = buildService({ store, now: () => Date.now() + 60_000 }); + expect(await expiredAfter.findActiveByScope(OWNER_SCOPE)).toBeNull(); + }); +}); + describe("SetupTokenSessionService.completeSession", () => { it("returns the non-secret storedSessionId after a successful secret write, and no token", async () => { const { service, processes, leases, store, secretWrites, logs } = buildService(); diff --git a/server/src/services/setup-token-session.ts b/server/src/services/setup-token-session.ts index 813286a168..0a922f5a03 100644 --- a/server/src/services/setup-token-session.ts +++ b/server/src/services/setup-token-session.ts @@ -32,7 +32,7 @@ // testable. import { randomBytes } from "node:crypto"; -import { and, eq, gt, inArray, isNotNull, isNull, lte, or, sql } from "drizzle-orm"; +import { and, eq, gt, inArray, isNotNull, isNull, lte, notInArray, or, sql } from "drizzle-orm"; import type { Db } from "@paperclipai/db"; import { adapterAuthSessions } from "@paperclipai/db"; import type { AgentAdapterType } from "@paperclipai/shared"; @@ -85,6 +85,18 @@ export function isTerminalSessionState(state: SetupTokenSessionState): boolean { return SETUP_TOKEN_TERMINAL_STATES.includes(state); } +/** + * The cancellable pre-promotion states. A durable-only cancel (no live + * in-memory session) may transition a row only from one of these. `submitting` + * is the credential-write phase — the setup-token analogue of the device-login + * `promoting` claim — so it is excluded: a durable cancel never interrupts a + * write in progress. `stored` and every terminal state are excluded too. + */ +export const SETUP_TOKEN_CANCELLABLE_STATES: readonly SetupTokenSessionState[] = [ + "starting", + "awaiting_code", +]; + /** * The immutable owner scope of a session. The service builds the scope once at * start and never changes it. The session identity is the company, the owner, @@ -264,6 +276,31 @@ export interface SetupTokenCleanupStore { * step. The agent-service transaction calls this method. */ consumeStoredClaim(identity: SetupTokenCleanupIdentity): Promise; + /** + * Cancels the exact durable row with one conditional write. The predicate + * matches the full owner scope, the session id, and one of + * `cancellableStates`. It returns the updated record only on a successful + * transition. It returns null for a missing row, a foreign-scope row, and a + * row outside the cancellable states — a caller cannot tell these apart. + */ + cancelDurable( + identity: SetupTokenCleanupIdentity, + cancellableStates: readonly SetupTokenSessionState[], + ): Promise; + /** + * Returns the caller's active durable row for a scope, with no session id. A + * restarted server process holds no in-memory session, so this is the + * fallback source of truth for session discovery: the row survives the + * restart even though the live process and the in-memory session do not. It + * returns a row only when the company, the owner, and the adapter match, the + * state is not terminal, and the deadline is not yet past. It returns null + * for a missing row, a foreign-scope row, a terminal row, and an expired + * row. + */ + findActiveDurable( + key: Pick, + now: number, + ): Promise; } /** The counts one reaper sweep produced over the durable cleanup store. */ @@ -1124,6 +1161,78 @@ export class SetupTokenSessionService { return { state: session.state }; } + /** + * Finds the caller's active session for a scope, with no session id. The + * browser rediscovers its own session after a reload with no local state. It + * matches the company, the owner, and the adapter, and it returns only a + * non-terminal session, so a caller with no active login and a caller with a + * foreign scope both find nothing. + * + * It checks the in-memory session first. A server restart drops every + * in-memory session, so it then falls back to the durable active row. The + * durable-only descriptor carries no login URL, because the full URL lives + * only in memory (SR-5). This fallback keeps the caller's start route from + * retrying a start that the durable active-row uniqueness constraint would + * reject. + */ + async findActiveByScope( + key: Pick, + ): Promise { + for (const session of this.sessions.values()) { + if (isTerminalSessionState(session.state)) continue; + if ( + session.scope.companyId === key.companyId && + session.scope.ownerUserId === key.ownerUserId && + session.scope.adapterType === key.adapterType + ) { + return this.describeOwned(session.id, session.scope); + } + } + const durable = await this.store.findActiveDurable(key, this.now()); + if (!durable) return null; + return { + sessionId: durable.sessionId, + state: durable.state, + environmentId: durable.environmentId, + deadline: durable.deadline, + loginUrl: null, + }; + } + + /** + * Cancels a session by its full scope, with a durable fallback when no live + * session exists. It first tries the live in-memory session that matches the + * scope and the session id, and cancels it the normal way. When no live + * session matches — for example, a restart dropped the in-memory session — + * it falls back to a durable-only cancel: a conditional transition of the + * exact row from a cancellable pre-promotion state to `cancelled`. This + * releases the company slot even though this process holds no live process + * to stop, so it aborts no local process. It throws the fixed not-found + * error for a missing row, a foreign owner, a foreign company, a foreign + * adapter, and a row outside the cancellable states — the caller cannot tell + * these apart. + */ + async cancelByScope( + sessionId: string, + key: Pick, + ): Promise<{ state: SetupTokenSessionState }> { + const session = this.sessions.get(sessionId); + if ( + session && + session.scope.companyId === key.companyId && + session.scope.ownerUserId === key.ownerUserId && + session.scope.adapterType === key.adapterType + ) { + return this.cancel(sessionId, session.scope); + } + const identity: SetupTokenCleanupIdentity = { sessionId, ...key }; + const cancelled = await this.store.cancelDurable(identity, SETUP_TOKEN_CANCELLABLE_STATES); + if (!cancelled) { + throw new SetupTokenSessionError(404, SETUP_TOKEN_SESSION_NOT_FOUND); + } + return { state: "cancelled" }; + } + /** * Expires a session on a timeout. It stops the direct child before it releases * the lease. The harness can call it, and the deadline timer calls the same @@ -1500,5 +1609,43 @@ export function createDbSetupTokenCleanupStore(db: Db): SetupTokenCleanupStore { const row = changed[0]; return row ? toCleanupRecord(row) : null; }, + + async cancelDurable(identity, cancellableStates): Promise { + // One conditional write. The predicate carries the company, the owner, the + // adapter, the session id, and one of `cancellableStates`. It never + // interrupts a `submitting` write, a `stored` claim, or a terminal row. + const changed = await db + .update(adapterAuthSessions) + .set({ + status: "cancelled", + finishedAt: sql`clock_timestamp()`, + updatedAt: sql`clock_timestamp()`, + }) + .where(and(scopeMatch(identity), inArray(adapterAuthSessions.status, [...cancellableStates]))) + .returning(); + const row = changed[0]; + return row ? toCleanupRecord(row) : null; + }, + + async findActiveDurable(key, now): Promise { + // The active slot cap is one row per company, owner, and adapter, so at + // most one row can match. The scan filters by the setup-token adapter, so + // it never reads a Codex device-login row on the shared table. + const rows = await db + .select() + .from(adapterAuthSessions) + .where( + and( + eq(adapterAuthSessions.companyId, key.companyId), + eq(adapterAuthSessions.startedByUserId, key.ownerUserId), + eq(adapterAuthSessions.adapterType, key.adapterType as AgentAdapterType), + notInArray(adapterAuthSessions.status, [...SETUP_TOKEN_TERMINAL_STATES]), + gt(adapterAuthSessions.expiresAt, new Date(now)), + ), + ) + .limit(1); + const row = rows[0]; + return row ? toCleanupRecord(row) : null; + }, }; } diff --git a/server/src/services/setup-token-transport-binding.test.ts b/server/src/services/setup-token-transport-binding.test.ts index a9f3d1f34e..e78b3dc79d 100644 --- a/server/src/services/setup-token-transport-binding.test.ts +++ b/server/src/services/setup-token-transport-binding.test.ts @@ -124,6 +124,23 @@ function createRecordingStore() { async consumeStoredClaim() { return null; }, + async cancelDurable() { + return null; + }, + async findActiveDurable(key, now) { + for (const row of rows.values()) { + if ( + row.companyId === key.companyId && + row.ownerUserId === key.ownerUserId && + row.adapterType === key.adapterType && + !isTerminalSessionState(row.state) && + row.deadline > now + ) { + return { ...row }; + } + } + return null; + }, }; return { diff --git a/tests/e2e/onboarding.spec.ts b/tests/e2e/onboarding.spec.ts index 2c6824387f..b425051524 100644 --- a/tests/e2e/onboarding.spec.ts +++ b/tests/e2e/onboarding.spec.ts @@ -305,6 +305,213 @@ test.describe("Onboarding wizard", () => { expect(pageErrors, pageErrors.join("\n")).toHaveLength(0); }); + test("a reload during a login resumes the same session", async ({ page }) => { + // The last piece of the resume feature: a customer who reloads the page + // mid-login must see the same login still running, not the tile row + // again and not a second session started behind their back. + // + // A bare `/onboarding` reload is not the right way to exercise this: that + // route always forces its own fixed entry step (see + // `resolveRouteOnboardingOptions` — a bare `/onboarding` hit is a request + // to start a new company, by design, and always wins over a saved draft's + // step). A returning visit to an existing company's own onboarding path + // (`/{prefix}/onboarding`) is the real "reopen onboarding" entry point, so + // this test reloads through that path instead, landing on the agent step + // with the draft's agent name already restored, then advances once to the + // connect step — where the resumed login must already be running. + const pageErrors: string[] = []; + page.on("pageerror", (err) => pageErrors.push(err.message)); + + const flagRes = await page.request.patch("/api/instance/settings/experimental", { + data: { enableConferenceRoomChat: true }, + }); + expect(flagRes.ok()).toBe(true); + + const FAKE_SANDBOX_ENVIRONMENT_ID = "e2e-fake-sandbox-environment-reload"; + const FAKE_SANDBOX_PROVIDER = "e2e-fake-provider-reload"; + + await page.route("**/environments", async (route) => { + const response = await route.fetch(); + const environments = await response.json(); + environments.push({ + id: FAKE_SANDBOX_ENVIRONMENT_ID, + name: "E2E fake sandbox", + description: null, + driver: "sandbox", + status: "active", + config: { provider: FAKE_SANDBOX_PROVIDER }, + envVars: {}, + metadata: {}, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }); + await route.fulfill({ response, json: environments }); + }); + + await page.route("**/environments/capabilities", async (route) => { + const response = await route.fetch(); + const capabilities = await response.json(); + capabilities.sandboxProviders[FAKE_SANDBOX_PROVIDER] = { + status: "supported", + supportsSavedProbe: true, + supportsUnsavedProbe: true, + supportsRunExecution: true, + supportsReusableLeases: false, + supportsInteractiveSetup: false, + interactiveSetupConnectionTypes: [], + supportsTemplateCapture: false, + supportsTemplateDelete: false, + supportsLoginPty: true, + source: "plugin", + }; + await route.fulfill({ response, json: capabilities }); + }); + + await page.route("**/instance/settings", async (route) => { + const response = await route.fetch(); + const settings = await response.json(); + settings.defaultEnvironmentId = FAKE_SANDBOX_ENVIRONMENT_ID; + await route.fulfill({ response, json: settings }); + }); + + await page.route("**/adapters/*/auth-signal*", (route) => + route.fulfill({ + contentType: "application/json", + body: JSON.stringify({ status: "absent" }), + }), + ); + + const SESSION_ID = "e2e-reload-setup-token-session"; + const AUTHORIZATION_URL = "https://claude.ai/oauth/authorize?code=true&client=e2e-reload"; + let startCalls = 0; + // Flips once the login actually starts, so the owner-scoped resume read + // below answers "no active session" until then — matching the real + // route, and keeping the pre-Connect part of this test the same as the + // ordinary sign-in test above. + let sessionStarted = false; + await page.route("**/setup-token-login-sessions", (route) => { + if (route.request().method() === "POST") { + startCalls += 1; + sessionStarted = true; + } + return route.fulfill({ + contentType: "application/json", + body: JSON.stringify({ + sessionId: SESSION_ID, + status: "pending", + expiresAt: new Date(Date.now() + 600_000).toISOString(), + }), + }); + }); + await page.route(`**/setup-token-login-sessions/${SESSION_ID}/prompt`, (route) => + route.fulfill({ + contentType: "application/json", + body: JSON.stringify({ + authorizationUrl: AUTHORIZATION_URL, + transportAdvisory: null, + }), + }), + ); + await page.route(`**/setup-token-login-sessions/${SESSION_ID}`, (route) => { + if (route.request().method() !== "GET") return route.continue(); + return route.fulfill({ + contentType: "application/json", + body: JSON.stringify({ + sessionId: SESSION_ID, + status: "pending", + expiresAt: new Date(Date.now() + 600_000).toISOString(), + }), + }); + }); + // The owner-scoped resume read. Both the wizard's own step-restore effect + // and the panel's own resume-on-mount read use this, with no session id + // in the URL — the caller rediscovers its own session. + await page.route("**/setup-token-login-sessions/active", (route) => { + if (!sessionStarted) { + return route.fulfill({ + status: 404, + contentType: "application/json", + body: JSON.stringify({ error: "Setup-token login session not found." }), + }); + } + return route.fulfill({ + contentType: "application/json", + body: JSON.stringify({ + sessionId: SESSION_ID, + status: "pending", + expiresAt: new Date(Date.now() + 600_000).toISOString(), + panelMode: "submitted_browser_code", + prompt: { authorizationUrl: AUTHORIZATION_URL, transportAdvisory: null }, + }), + }); + }); + + await page.goto("/onboarding"); + + const startBtn = page.getByRole("button", { + name: /Start Onboarding|New Organization|Add Agent/, + }); + if (await startBtn.count()) { + await startBtn.first().click(); + } + const createCard = page.getByRole("button", { name: /Build a new organization/ }); + if (await createCard.count()) { + await createCard.first().click(); + } + + await expect( + page.getByRole("heading", { name: "What is the name of your organization?" }), + ).toBeVisible({ timeout: 15_000 }); + await page.getByPlaceholder("e.g. Northwind Labs").fill(`${COMPANY_NAME}-reload`); + await page.getByRole("button", { name: /^Continue/ }).click(); + + await page.waitForSelector("#onboarding-agent-name", { timeout: 30_000 }); + await page.locator("#onboarding-agent-name").fill("Ada"); + await page.getByRole("button", { name: "Next" }).click(); + + const source = page.getByRole("radio").first(); + await source.waitFor({ timeout: 30_000 }); + + // Answering the row is what starts the sign-in now — there is no separate + // Connect press between choosing a source and being signed in. + await source.click(); + + const cardInstruction = page.getByText("then come back and enter authorization code"); + const authorizationLink = page.getByRole("link", { name: /^Sign in to / }); + + await expect(cardInstruction).toBeVisible({ timeout: 30_000 }); + await expect(authorizationLink).toBeVisible({ timeout: 15_000 }); + await expect(authorizationLink).toHaveAttribute("href", /claude\.ai\/oauth\/authorize/); + expect(startCalls).toBe(1); + + const companiesRes = await page.request.get("/api/companies"); + expect(companiesRes.ok()).toBe(true); + const companies = await companiesRes.json(); + const company = companies.find( + (c: { name: string }) => c.name === `${COMPANY_NAME}-reload`, + ); + expect(company, "the created company should exist").toBeTruthy(); + + // Reload through the company's own onboarding path — the real "reopen + // onboarding" entry point for a company that already exists. + await page.goto(`/${company.issuePrefix}/onboarding`); + + // This entry point always re-enters on the agent step, with the agent + // name restored from the draft. One more press reaches the connect step. + await page.waitForSelector("#onboarding-agent-name", { timeout: 30_000 }); + await expect(page.locator("#onboarding-agent-name")).toHaveValue("Ada"); + await page.getByRole("button", { name: "Next" }).click(); + + // The resumed login shows with no new source pick and no fresh sign-in + // press: the panel and the step both discover it from the caller's active + // session. + await expect(cardInstruction).toBeVisible({ timeout: 15_000 }); + await expect(authorizationLink).toBeVisible({ timeout: 15_000 }); + expect(startCalls, "the reload must resume the session, not start a new one").toBe(1); + + expect(pageErrors, pageErrors.join("\n")).toHaveLength(0); + }); + test("connect step blocks the hire when the environment probe fails and no sign-in is needed", async ({ page, }) => { diff --git a/ui/src/api/agents.ts b/ui/src/api/agents.ts index 530f3006d2..5f7cc46cf9 100644 --- a/ui/src/api/agents.ts +++ b/ui/src/api/agents.ts @@ -252,9 +252,21 @@ export const agentsApi = { `/companies/${encodeURIComponent(companyId)}/adapters/${encodeURIComponent(type)}/login-sessions`, data, ), + // The owner response repeats the live prompt on every read while the + // session holds an active public status, so this polling call carries the + // same no-store request option as the active-session read below. getAdapterAuthLoginStatus: (companyId: string, type: string, sessionId: string) => api.get( `/companies/${encodeURIComponent(companyId)}/adapters/${encodeURIComponent(type)}/login-sessions/${encodeURIComponent(sessionId)}`, + { cache: "no-store" }, + ), + // Reads the caller's active login session for one company and adapter, with + // no session id, so the browser rediscovers its own session after a reload + // with no local state. A 404 means no active session for the caller. + getActiveAdapterAuthLoginSession: (companyId: string, type: string) => + api.get( + `/companies/${encodeURIComponent(companyId)}/adapters/${encodeURIComponent(type)}/login-sessions/active`, + { cache: "no-store" }, ), cancelAdapterAuthLogin: (companyId: string, type: string, sessionId: string) => api.post( @@ -287,6 +299,14 @@ export const agentsApi = { api.get( `/companies/${encodeURIComponent(companyId)}/setup-token-login-sessions/${encodeURIComponent(sessionId)}`, ), + // Reads the caller's active Claude setup-token login session, with no + // session id, so the browser rediscovers its own session after a reload with + // no local state. A 404 means no active session for the caller. + getActiveClaudeSetupTokenLoginSession: (companyId: string) => + api.get( + `/companies/${encodeURIComponent(companyId)}/setup-token-login-sessions/active`, + { cache: "no-store" }, + ), getClaudeSetupTokenLoginPrompt: (companyId: string, sessionId: string) => api.get( `/companies/${encodeURIComponent(companyId)}/setup-token-login-sessions/${encodeURIComponent(sessionId)}/prompt`, diff --git a/ui/src/api/client.ts b/ui/src/api/client.ts index 3966cb0d65..88bb9a9cd1 100644 --- a/ui/src/api/client.ts +++ b/ui/src/api/client.ts @@ -20,6 +20,9 @@ export interface RequestOptions { signal?: AbortSignal; /** Extra request headers (e.g. the async-import opt-in). Mutations only. */ headers?: Record; + /** The `fetch` cache mode. Use `"no-store"` for a response that must never + * come from the browser's HTTP cache. */ + cache?: RequestCache; } function abortError(): DOMException { @@ -88,7 +91,11 @@ function coalescedGet(path: string, options?: RequestOptions): Promise { let entry = inflightGets.get(path); if (!entry) { const controller = new AbortController(); - const promise = request(path, { method: "GET", signal: controller.signal }); + const promise = request(path, { + method: "GET", + signal: controller.signal, + ...(options?.cache ? { cache: options.cache } : {}), + }); const created: InflightGet = { promise, controller, refs: new Set() }; // Clear the shared entry once settled so later calls issue a fresh request. promise.then( diff --git a/ui/src/components/AgentConfigForm.render.test.tsx b/ui/src/components/AgentConfigForm.render.test.tsx index fdc6c1697d..f9fa9dcea5 100644 --- a/ui/src/components/AgentConfigForm.render.test.tsx +++ b/ui/src/components/AgentConfigForm.render.test.tsx @@ -21,9 +21,11 @@ const mockAgentsApi = vi.hoisted(() => ({ testEnvironment: vi.fn(), startAdapterAuthLogin: vi.fn(), getAdapterAuthLoginStatus: vi.fn(), + getActiveAdapterAuthLoginSession: vi.fn(), cancelAdapterAuthLogin: vi.fn(), startClaudeSetupTokenLogin: vi.fn(), getClaudeSetupTokenLoginStatus: vi.fn(), + getActiveClaudeSetupTokenLoginSession: vi.fn(), getClaudeSetupTokenLoginPrompt: vi.fn(), submitClaudeSetupTokenBrowserCode: vi.fn(), completeClaudeSetupTokenLogin: vi.fn(), @@ -31,6 +33,14 @@ const mockAgentsApi = vi.hoisted(() => ({ getClaudeOAuthTokenStatus: vi.fn(), })); +// The default resume read for a test that does not exercise resume: no active +// session for the caller. +function noActiveSession() { + return Promise.reject( + new ApiError("Adapter login session not found", 404, { error: "Adapter login session not found" }), + ); +} + const mockClipboard = vi.hoisted(() => ({ copyTextToClipboard: vi.fn(), })); @@ -645,6 +655,10 @@ describe("AgentConfigForm environment selector", () => { mockEnvironmentsApi.capabilities.mockResolvedValue(SANDBOX_CAPABILITIES); mockSecretsApi.list.mockResolvedValue([]); mockSecretsApi.listProposals.mockResolvedValue([]); + // Default: the caller has no active session. A resume test overrides this + // with a resolved session body. + mockAgentsApi.getActiveAdapterAuthLoginSession.mockImplementation(noActiveSession); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockImplementation(noActiveSession); mockAgentsApi.startAdapterAuthLogin.mockResolvedValue({ sessionId: "session-1", environmentId: "sandbox-1", @@ -1603,17 +1617,14 @@ describe("AgentConfigForm environment selector", () => { expect(result.container.textContent).not.toContain("WXYZ-1234"); }); - it("releases an active login session when the panel unmounts", async () => { - // The server holds a one-per-owner reservation until the session reaches a - // terminal state, so a panel that disappears mid-login would leave the owner - // unable to start another until it expires. - // - // Reachable in the settings form only by navigating away, which is why this - // went unnoticed. The connect step unmounts the panel routinely — Cancel - // closes the canvas, switching source remounts it under a new key, closing - // the wizard drops it — so the cleanup is what keeps an immediate retry - // possible. Deliberately not pushed to `roots`: this test does the unmount - // itself, and that unmount is the thing under test. + it("does not cancel an active login session when the panel unmounts", async () => { + // The owner-scoped active-session read and the manual Cancel button now + // take over the purpose the unmount cancel used to serve. The connect step + // unmounts the panel routinely — Cancel closes the canvas, switching source + // remounts it under a new key, closing the wizard drops it — and each of + // those unmounts must leave the session reachable by a later mount's resume + // read, not release it. Deliberately not pushed to `roots`: this test does + // the unmount itself, and that unmount is the thing under test. mockAgentsApi.testEnvironment.mockResolvedValue(AUTH_MISSING_RESULT); const result = await renderCodexSandbox(); @@ -1626,11 +1637,163 @@ describe("AgentConfigForm environment selector", () => { result.root.unmount(); }); + expect(mockAgentsApi.cancelAdapterAuthLogin).not.toHaveBeenCalled(); + }); + + it("shows a reachable Cancel control in the onboarding chrome and cancels the session", async () => { + mockAgentsApi.testEnvironment.mockResolvedValue(AUTH_MISSING_RESULT); + const onCancel = vi.fn(); + + const container = document.createElement("div"); + document.body.appendChild(container); + const root = createRoot(container); + roots.push(root); + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + + await act(async () => { + root.render( + + + + + + + , + ); + }); + await flushUntil(() => Boolean(findButton(container, "Cancel"))); + + await clickByText(container, "Cancel"); + expect(mockAgentsApi.cancelAdapterAuthLogin).toHaveBeenCalledWith( "company-1", "codex_local", "session-1", ); + expect(onCancel).toHaveBeenCalledTimes(1); + }); + + it("resumes an active login session on mount, adopting its session id and prompt", async () => { + // A page reload loses every piece of local state, so the panel must read + // the caller's active session and adopt it instead of starting a new one. + mockAgentsApi.testEnvironment.mockResolvedValue(AUTH_MISSING_RESULT); + mockAgentsApi.getActiveAdapterAuthLoginSession.mockResolvedValue({ + sessionId: "resumed-session-1", + environmentId: "sandbox-1", + status: "waiting_for_user", + expiresAt: null, + failure: null, + prompt: { url: "https://auth.example.test/resumed", code: "RESUME-1" }, + }); + // The prompt survives every owner read while the session stays active, so + // the status poll for the resumed session agrees with the resumed read. + mockAgentsApi.getAdapterAuthLoginStatus.mockResolvedValue({ + sessionId: "resumed-session-1", + environmentId: "sandbox-1", + status: "waiting_for_user", + expiresAt: null, + failure: null, + prompt: { url: "https://auth.example.test/resumed", code: "RESUME-1" }, + }); + const result = await renderCodexSandbox(); + roots.push(result.root); + + await runTest(result.container); + await flushUntil(() => (result.container.textContent ?? "").includes("RESUME-1")); + + // The prompt from the resumed read shows without a fresh start. + expect(result.container.textContent).toContain("RESUME-1"); + expect(result.container.textContent).toContain("https://auth.example.test/resumed"); + expect(mockAgentsApi.startAdapterAuthLogin).not.toHaveBeenCalled(); + // The panel polls the resumed session id, not a freshly started one. + expect(mockAgentsApi.getAdapterAuthLoginStatus).toHaveBeenCalledWith( + "company-1", + "codex_local", + "resumed-session-1", + ); + expect(findButton(result.container, "Cancel")).toBeTruthy(); + }); + + it("starts a new session when the active-session read finds none", async () => { + // The default mock already answers with no active session (a 404). This + // pins the fallback: the panel still waits for the caller's press. + mockAgentsApi.testEnvironment.mockResolvedValue(AUTH_MISSING_RESULT); + const result = await renderCodexSandbox(); + roots.push(result.root); + + await runTest(result.container); + await flushUntil(() => mockAgentsApi.getActiveAdapterAuthLoginSession.mock.calls.length > 0); + await flushReact(); + + expect(mockAgentsApi.startAdapterAuthLogin).not.toHaveBeenCalled(); + expect(findButton(result.container, "Sign in")?.disabled).toBe(false); + + await startLogin(result.container); + + expect(mockAgentsApi.startAdapterAuthLogin).toHaveBeenCalledWith("company-1", "codex_local", { + environmentId: "sandbox-1", + }); + }); + + it("releases a resumed session after an unrecoverable resume error, waiting for the cancel response", async () => { + // The active-session read finds a session, but by the time the status poll + // reaches the server the session is already gone (a race between the two + // reads). The panel cannot resume it, so it releases the reservation + // explicitly instead of trusting the 404 alone, and it waits for that + // release before it returns to its start state. + mockAgentsApi.testEnvironment.mockResolvedValue(AUTH_MISSING_RESULT); + mockAgentsApi.getActiveAdapterAuthLoginSession.mockResolvedValue({ + sessionId: "resumed-session-1", + environmentId: "sandbox-1", + status: "waiting_for_user", + expiresAt: null, + failure: null, + prompt: null, + }); + mockAgentsApi.getAdapterAuthLoginStatus.mockRejectedValue( + new ApiError("Adapter login session not found", 404, { + error: "Adapter login session not found", + }), + ); + let resolveCancel!: () => void; + mockAgentsApi.cancelAdapterAuthLogin.mockReturnValue( + new Promise((resolve) => { + resolveCancel = () => resolve({ + sessionId: "resumed-session-1", + environmentId: "sandbox-1", + status: "cancelled", + expiresAt: null, + failure: null, + prompt: null, + }); + }), + ); + const result = await renderCodexSandbox(); + roots.push(result.root); + + await runTest(result.container); + await flushUntil(() => + mockAgentsApi.cancelAdapterAuthLogin.mock.calls.some((call) => call[2] === "resumed-session-1"), + ); + + // The cancel call fired, but the panel still shows the resumed login as + // active because it is waiting for the cancel response. + expect(findButton(result.container, "Cancel")).toBeTruthy(); + + resolveCancel(); + await flushUntil(() => findButton(result.container, "Sign in")?.disabled === false); + + expect(findButton(result.container, "Cancel")).toBeFalsy(); + expect(findButton(result.container, "Sign in")?.disabled).toBe(false); }); it("announces the login state through a polite live region", async () => { @@ -2028,6 +2191,46 @@ describe("AgentConfigForm environment selector", () => { expect(onStored).toHaveBeenCalledWith("stored-session-1"); }); + it("shows a reachable Cancel control in the onboarding chrome and cancels the session", async () => { + const onCancel = vi.fn(); + + const container = document.createElement("div"); + document.body.appendChild(container); + const root = createRoot(container); + roots.push(root); + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false }, mutations: { retry: false } }, + }); + + await act(async () => { + root.render( + + + + + + + , + ); + }); + await flushUntil(() => Boolean(findButton(container, "Cancel"))); + + await clickByText(container, "Cancel"); + + expect(mockAgentsApi.cancelClaudeSetupTokenLogin).toHaveBeenCalledWith( + "company-1", + "claude-session-1", + ); + expect(onCancel).toHaveBeenCalledTimes(1); + }); + it("offers an apply-existing affordance when the status route reports a stored value", async () => { mockAgentsApi.getClaudeOAuthTokenStatus.mockResolvedValue({ secretId: "secret-1", @@ -2265,7 +2468,10 @@ describe("AgentConfigForm environment selector", () => { expect(result.container.textContent).not.toContain("Could not cancel the login."); }); - it("cancels the active server session when the panel unmounts", async () => { + it("does not cancel an active login session when the panel unmounts", async () => { + // The owner-scoped active-session read and the manual Cancel button now + // take over the purpose the unmount cancel used to serve, so an unmount + // must leave the session reachable by a later mount's resume read. mockAgentsApi.testEnvironment.mockResolvedValue(CLAUDE_AUTH_MISSING_RESULT); const result = await renderClaudeSandbox(); @@ -2277,16 +2483,12 @@ describe("AgentConfigForm environment selector", () => { // The login is active before the unmount. expect(findButton(result.container, "Cancel")).toBeTruthy(); + mockAgentsApi.cancelClaudeSetupTokenLogin.mockClear(); await act(async () => { result.root.unmount(); }); - // The unmount released the active server session, so the abandoned session - // does not hold the per-owner reservation until the server deadline. - expect(mockAgentsApi.cancelClaudeSetupTokenLogin).toHaveBeenCalledWith( - "company-1", - "claude-session-1", - ); + expect(mockAgentsApi.cancelClaudeSetupTokenLogin).not.toHaveBeenCalled(); }); it("does not cancel on unmount when no login is active", async () => { @@ -2305,6 +2507,134 @@ describe("AgentConfigForm environment selector", () => { expect(mockAgentsApi.cancelClaudeSetupTokenLogin).not.toHaveBeenCalled(); }); + it("resumes an active Claude login session on mount, adopting its session id and authorization URL", async () => { + mockAgentsApi.testEnvironment.mockResolvedValue(CLAUDE_AUTH_MISSING_RESULT); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockResolvedValue({ + sessionId: "resumed-claude-session-1", + environmentId: "sandbox-1", + status: "waiting_for_user", + expiresAt: null, + failure: null, + panelMode: "submitted_browser_code", + prompt: { authorizationUrl: "https://claude.example.test/resumed" }, + }); + const result = await renderClaudeSandbox(); + roots.push(result.root); + + await runTest(result.container); + await flushUntil(() => + (result.container.textContent ?? "").includes("https://claude.example.test/resumed"), + ); + + expect(result.container.textContent).toContain("https://claude.example.test/resumed"); + expect(mockAgentsApi.startClaudeSetupTokenLogin).not.toHaveBeenCalled(); + expect(mockAgentsApi.getClaudeSetupTokenLoginStatus).toHaveBeenCalledWith( + "company-1", + "resumed-claude-session-1", + ); + expect(findButton(result.container, "Cancel")).toBeTruthy(); + }); + + it("starts a new Claude login when the active-session read finds none", async () => { + // The default mock already answers with no active session (a 404). + mockAgentsApi.testEnvironment.mockResolvedValue(CLAUDE_AUTH_MISSING_RESULT); + const result = await renderClaudeSandbox(); + roots.push(result.root); + + await runTest(result.container); + await flushUntil(() => mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mock.calls.length > 0); + await flushReact(); + + expect(mockAgentsApi.startClaudeSetupTokenLogin).not.toHaveBeenCalled(); + expect(findButton(result.container, "Sign in")?.disabled).toBe(false); + + await startLogin(result.container); + + expect(mockAgentsApi.startClaudeSetupTokenLogin).toHaveBeenCalled(); + }); + + it("cancels a resumed Claude login session with the manual Cancel button", async () => { + mockAgentsApi.testEnvironment.mockResolvedValue(CLAUDE_AUTH_MISSING_RESULT); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockResolvedValue({ + sessionId: "resumed-claude-session-1", + environmentId: "sandbox-1", + status: "waiting_for_user", + expiresAt: null, + failure: null, + panelMode: "submitted_browser_code", + prompt: { authorizationUrl: "https://claude.example.test/resumed" }, + }); + const result = await renderClaudeSandbox(); + roots.push(result.root); + + await runTest(result.container); + await flushUntil(() => + (result.container.textContent ?? "").includes("https://claude.example.test/resumed"), + ); + + await clickByText(result.container, "Cancel"); + await flushReact(); + + expect(mockAgentsApi.cancelClaudeSetupTokenLogin).toHaveBeenCalledWith( + "company-1", + "resumed-claude-session-1", + ); + expect(findButton(result.container, "Sign in")?.disabled).toBe(false); + expect(findButton(result.container, "Cancel")).toBeFalsy(); + }); + + it("releases a resumed Claude login after an unrecoverable resume error, waiting for the cancel response", async () => { + // The active-session read finds a session, but the status poll for that + // resumed session finds it already gone (a race between the two reads). + // The panel cannot resume it, so it releases the reservation explicitly + // and waits for that release before it returns to its start state. + mockAgentsApi.testEnvironment.mockResolvedValue(CLAUDE_AUTH_MISSING_RESULT); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockResolvedValue({ + sessionId: "resumed-claude-session-1", + environmentId: "sandbox-1", + status: "waiting_for_user", + expiresAt: null, + failure: null, + panelMode: "submitted_browser_code", + prompt: null, + }); + mockAgentsApi.getClaudeSetupTokenLoginStatus.mockRejectedValue( + new ApiError("Setup-token login session not found.", 404, { + error: "Setup-token login session not found.", + }), + ); + mockAgentsApi.getClaudeSetupTokenLoginPrompt.mockRejectedValue( + new ApiError("Setup-token login session not found.", 404, { + error: "Setup-token login session not found.", + }), + ); + let resolveCancel!: () => void; + mockAgentsApi.cancelClaudeSetupTokenLogin.mockReturnValue( + new Promise((resolve) => { + resolveCancel = () => resolve(undefined); + }), + ); + const result = await renderClaudeSandbox(); + roots.push(result.root); + + await runTest(result.container); + await flushUntil(() => + mockAgentsApi.cancelClaudeSetupTokenLogin.mock.calls.some( + (call) => call[1] === "resumed-claude-session-1", + ), + ); + + // The cancel call fired, but the panel still shows the resumed login as + // active because it is waiting for the cancel response. + expect(findButton(result.container, "Sign in")?.disabled).toBe(true); + + resolveCancel(); + await flushUntil(() => findButton(result.container, "Sign in")?.disabled === false); + + expect(findButton(result.container, "Cancel")).toBeFalsy(); + expect(findButton(result.container, "Sign in")?.disabled).toBe(false); + }); + it("stops both polls and shows the timed-out state at the server deadline", async () => { vi.useFakeTimers(); // Pin the clock to a known base. The panel arms the timed-out timer from the @@ -2539,6 +2869,8 @@ describe("AgentConfigForm create-mode Claude OAuth binding", () => { mockEnvironmentsApi.capabilities.mockResolvedValue(SANDBOX_CAPABILITIES); mockSecretsApi.list.mockResolvedValue([]); mockSecretsApi.listProposals.mockResolvedValue([]); + mockAgentsApi.getActiveAdapterAuthLoginSession.mockImplementation(noActiveSession); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockImplementation(noActiveSession); mockAgentsApi.testEnvironment.mockResolvedValue(CLAUDE_AUTH_MISSING_RESULT); mockAgentsApi.startClaudeSetupTokenLogin.mockResolvedValue({ sessionId: "claude-session-1", @@ -2737,6 +3069,8 @@ describe("AgentConfigForm edit-mode Claude OAuth binding", () => { mockEnvironmentsApi.capabilities.mockResolvedValue(SANDBOX_CAPABILITIES); mockSecretsApi.list.mockResolvedValue([]); mockSecretsApi.listProposals.mockResolvedValue([]); + mockAgentsApi.getActiveAdapterAuthLoginSession.mockImplementation(noActiveSession); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockImplementation(noActiveSession); mockAgentsApi.testEnvironment.mockResolvedValue(CLAUDE_AUTH_MISSING_RESULT); mockAgentsApi.startClaudeSetupTokenLogin.mockResolvedValue({ sessionId: "claude-session-1", diff --git a/ui/src/components/AgentConfigForm.tsx b/ui/src/components/AgentConfigForm.tsx index e13eefd2ba..2497cfdac0 100644 --- a/ui/src/components/AgentConfigForm.tsx +++ b/ui/src/components/AgentConfigForm.tsx @@ -2020,10 +2020,10 @@ export type AdapterLoginDescriptor = { // // They are props on the existing panels rather than a second implementation // because the part onboarding needs unchanged is the whole of it: the session -// start, the two polls, the server deadline, the one-shot completion read, the -// unmount release. A copy drawn to the new design would have had to reproduce -// all of that correctly, and the first thing to rot would have been the -// timeout and cleanup paths, which are the ones nobody exercises by hand. +// start, the two polls, the server deadline, the one-shot completion read. A +// copy drawn to the new design would have had to reproduce all of that +// correctly, and the first thing to rot would have been the timeout and +// cleanup paths, which are the ones nobody exercises by hand. export type AdapterLoginPanelProps = AdapterLoginDescriptor & { onStored?: (storedSessionId: string) => void; onApplyStored?: () => void; @@ -2047,8 +2047,9 @@ export type AdapterLoginPanelProps = AdapterLoginDescriptor & { * its own button is what sends the customer there, and a prompt arriving is * what moves the step from waiting to ready. Everything else it needs the * panel already does — the paste submits itself, success is reported through - * `onConnected`, and an unmount releases the session — so this stays a single - * value rather than a whole session handed upward. + * `onConnected`, and the customer's own Cancel press is reported through + * `onCancel` — so this stays a single value rather than a whole session + * handed upward. */ onPromptReady?: (authorizationUrl: string | null) => void; }; @@ -2102,9 +2103,18 @@ function DisplayedCodeLoginPanel({ // URL. const [latchedPrompt, setLatchedPrompt] = useState(null); + // True for the session currently held in `sessionId` when it came from the + // owner-scoped resume read rather than a fresh `startLogin`. It marks the + // one case that needs the extra release-on-error path below: a session this + // browser instance did not just start, so a broken poll cannot fall back to + // the ordinary "let the user press Sign in again" recovery — the owner has + // no local memory of ever starting it. + const resumedRef = useRef(false); + const startLogin = useMutation({ mutationFn: () => agentsApi.startAdapterAuthLogin(companyId, adapterType, { environmentId }), onSuccess: (session) => { + resumedRef.current = false; setStartError(null); setLatchedPrompt(null); setSessionId(session.sessionId); @@ -2114,24 +2124,54 @@ function DisplayedCodeLoginPanel({ }, }); + // Reset local state, so the panel returns to its idle start state and the + // Sign in button is available again. + const clearActiveSession = useCallback(() => { + resumedRef.current = false; + setSessionId(null); + setLatchedPrompt(null); + setStartError(null); + }, []); + const cancelLogin = useMutation({ mutationFn: () => agentsApi.cancelAdapterAuthLogin(companyId, adapterType, sessionId!), - onSuccess: () => { - // Reset local state, so the panel returns to its idle start state and the - // Log in button is available again. - setSessionId(null); - setLatchedPrompt(null); - setStartError(null); - }, + onSuccess: clearActiveSession, onError: (error) => { setStartError(error instanceof Error ? error.message : "Could not cancel the login."); }, }); + // Read the caller's active session on mount, with no session id, so the + // browser rediscovers its own session after a reload with no local state. A + // 404 means no active session for the caller. + const activeSessionQuery = useQuery({ + queryKey: ["adapter-login-active-session", companyId, adapterType], + queryFn: async () => { + try { + return await agentsApi.getActiveAdapterAuthLoginSession(companyId, adapterType); + } catch (error) { + if (error instanceof ApiError && error.status === 404) return null; + throw error; + } + }, + retry: false, + }); + + // While the panel releases a resumed session it cannot recover (see below), + // it keeps showing the login as active rather than dropping back to idle, so + // it does not clear local state before the release finishes. + const [releasingResumedSession, setReleasingResumedSession] = useState(false); + const statusQuery = useQuery({ queryKey: ["adapter-login-status", companyId, adapterType, sessionId], queryFn: () => agentsApi.getAdapterAuthLoginStatus(companyId, adapterType, sessionId!), - enabled: Boolean(sessionId), + enabled: Boolean(sessionId) && !releasingResumedSession, + // A status 404 is unrecoverable: the server removed the row, so a retry + // cannot bring it back. Stop at once and fail loudly. + retry: (failureCount, error) => { + if (error instanceof ApiError && error.status === 404) return false; + return failureCount < 3; + }, refetchInterval: (query) => { const status = query.state.data?.status; return status && ADAPTER_LOGIN_TERMINAL_STATUSES.has(status) @@ -2154,54 +2194,79 @@ function DisplayedCodeLoginPanel({ const isActive = Boolean(sessionId) && !isTerminal; const startDisabled = startLogin.isPending || isActive; - // Release the server session at once, without a change to the panel state, the - // way the submitted-browser-code panel does. The server holds a per-owner - // reservation until the session reaches a terminal state, so an abandoned - // session locks the owner out until the server deadline. Fire-and-forget: a - // 404 means the server already removed a terminal session, and a cleanup path - // cannot surface any other error either, so it drops them all. The manual - // Cancel button keeps using the `cancelLogin` mutation, because that path also - // returns the panel to its idle start state. - const releaseServerSession = useCallback( - (id: string) => { - void agentsApi.cancelAdapterAuthLogin(companyId, adapterType, id).catch(() => { - // Drop the error, as above. - }); - }, - [companyId, adapterType], - ); - - // Hold the active session id for the unmount cleanup. Onboarding removes this - // panel as soon as Cancel is pressed — `handleCancel` fires the request and - // calls `onCancel` without waiting for it — so the panel can be gone before - // the cancel resolves. Without this, a failed cancel, or any other unmount - // (navigating away, the step advancing), would leave the reservation held - // until the server deadline and an immediate retry unable to start. The ref is - // null once the session leaves the active state, so the cleanup never cancels - // a session the server already removed. - const activeSessionRef = useRef(null); - activeSessionRef.current = isActive ? sessionId : null; - + // Adopt the caller's active session once, on mount. This is what makes a + // page reload keep the session: with no local state at all, the panel would + // otherwise show its idle start state even though the server still holds an + // active login for this owner. + const resumeAttemptedRef = useRef(false); useEffect(() => { - return () => { - const id = activeSessionRef.current; - if (id) releaseServerSession(id); - }; - }, [releaseServerSession]); + if (resumeAttemptedRef.current || !activeSessionQuery.isFetched) return; + resumeAttemptedRef.current = true; + const active = activeSessionQuery.data; + if (!active) return; + resumedRef.current = true; + setStartError(null); + setLatchedPrompt(active.prompt ?? null); + setSessionId(active.sessionId); + }, [activeSessionQuery.isFetched, activeSessionQuery.data]); - // Start once, on mount, when the caller has already taken the press. The ref - // is the guard rather than the mutation's own pending flag: `startLogin` - // settles, and without a latch a re-render after it settles would read "not - // pending, no session yet" during the gap before the session id lands and - // start a second login the server would count against the per-owner cap. + // A resumed session's status poll found the session already gone: the read + // that discovered it and the poll that tried to use it raced, and the + // session lost. The panel cannot resume it, and there is no unmount cleanup + // left to fall back on, so it releases the reservation itself and waits for + // that release before it returns to the idle start state. + useEffect(() => { + const error = statusQuery.error; + if (!(error instanceof ApiError && error.status === 404)) return; + if (!resumedRef.current || releasingResumedSession) return; + setReleasingResumedSession(true); + const id = sessionId; + void (async () => { + if (id) { + await agentsApi.cancelAdapterAuthLogin(companyId, adapterType, id).catch(() => { + // The session is already gone either way; nothing more to do. + }); + } + setReleasingResumedSession(false); + clearActiveSession(); + })(); + }, [statusQuery.error, releasingResumedSession, sessionId, companyId, adapterType, clearActiveSession]); + + // Start once, on mount, when the caller has already taken the press, and + // only once the resume read has answered: a resumed session takes over + // instead of a fresh start. The ref is the guard rather than the mutation's + // own pending flag: `startLogin` settles, and without a latch a re-render + // after it settles would read "not pending, no session yet" during the gap + // before the session id lands and start a second login the server would + // count against the per-owner cap. const autoStartedRef = useRef(false); const startLoginRef = useRef(startLogin.mutate); startLoginRef.current = startLogin.mutate; useEffect(() => { if (!autoStart || autoStartedRef.current) return; + // A failed lookup is not proof that no session exists: only a successful + // lookup is. Show the failure to the user instead of starting a second + // login the server would reject against the per-owner cap. + if (activeSessionQuery.isError) { + autoStartedRef.current = true; + setStartError( + activeSessionQuery.error instanceof Error + ? activeSessionQuery.error.message + : "Could not check for an active login.", + ); + return; + } + if (!activeSessionQuery.isSuccess) return; autoStartedRef.current = true; + if (activeSessionQuery.data) return; startLoginRef.current(); - }, [autoStart]); + }, [ + autoStart, + activeSessionQuery.isSuccess, + activeSessionQuery.isError, + activeSessionQuery.data, + activeSessionQuery.error, + ]); // Report success upward once. `authenticated` is this panel's terminal // success: unlike the Claude login there is no completion read after it, so @@ -2236,6 +2301,7 @@ function DisplayedCodeLoginPanel({ return ( {/* The same destination as the step's own button. Two ways to one @@ -2465,6 +2531,13 @@ function SubmittedBrowserCodeLoginPanel({ // apply-existing path binds the fixed reference with no new login round trip, // so the panel shows the applied confirmation and hides the apply affordance. const [appliedStored, setAppliedStored] = useState(false); + // True for the session currently held in `sessionId` when it came from the + // owner-scoped resume read rather than a fresh `startLogin`. It marks the + // one case that needs the extra release-on-error path below: a session this + // browser instance did not just start, so a broken poll cannot fall back to + // the ordinary "let the user press Sign in again" recovery — the owner has + // no local memory of ever starting it. + const resumedRef = useRef(false); const resetLocalState = () => { setStartError(null); @@ -2520,6 +2593,7 @@ function SubmittedBrowserCodeLoginPanel({ : {}), }), onSuccess: (session) => { + resumedRef.current = false; resetLocalState(); setSessionId(session.sessionId); }, @@ -2531,6 +2605,7 @@ function SubmittedBrowserCodeLoginPanel({ const clearActiveSession = () => { // Return the panel to its idle start state. The Log in button is available // again, and both polls stop because the session id is null. + resumedRef.current = false; setSessionId(null); resetLocalState(); }; @@ -2555,7 +2630,7 @@ function SubmittedBrowserCodeLoginPanel({ }); // Release the server session at once, without a change to the panel state. The - // client-cutoff timer and the unmount path both use this. The server holds a + // client-cutoff timer uses this. The server holds a // per-owner reservation until the session reaches a terminal state, so an // abandoned session locks the owner out until the server deadline. A best- // effort cancel frees that reservation now, so the same owner can start a new @@ -2575,11 +2650,33 @@ function SubmittedBrowserCodeLoginPanel({ [companyId], ); + // Read the caller's active Claude setup-token session on mount, with no + // session id, so the browser rediscovers its own session after a reload + // with no local state. A 404 means no active session for the caller. + const activeSessionQuery = useQuery({ + queryKey: ["claude-setup-token-active-session", companyId], + queryFn: async () => { + try { + return await agentsApi.getActiveClaudeSetupTokenLoginSession(companyId); + } catch (error) { + if (error instanceof ApiError && error.status === 404) return null; + throw error; + } + }, + retry: false, + }); + + // While the panel releases a resumed session it cannot recover (see below), + // it keeps showing the login as active rather than dropping back to idle, so + // it does not clear local state before the release finishes. + const [releasingResumedSession, setReleasingResumedSession] = useState(false); + // Both polls run only while a session is active and the client cap has not // passed. The timeout stops the polls, so the panel never polls forever. A // status 404 also stops the polls: the server cleaned up the session, so the // panel enters a terminal failure state instead. - const pollingEnabled = Boolean(sessionId) && !timedOut && !statusGone; + const pollingEnabled = + Boolean(sessionId) && !timedOut && !statusGone && !releasingResumedSession; const statusQuery = useQuery({ queryKey: ["claude-setup-token-status", companyId, sessionId], @@ -2606,12 +2703,51 @@ function SubmittedBrowserCodeLoginPanel({ // race against the next poll. React Query keeps the last successful data on // error, so without this branch the panel would hold stale data and show // nothing. Enter the terminal failure state, which stops both polls. + // + // A resumed session takes a different path: the read that discovered it and + // the poll that tried to use it raced, and the session lost. There is no + // unmount cleanup left to fall back on, so the panel releases the + // reservation itself and waits for that release before it returns to the + // idle start state, instead of trusting the 404 alone. useEffect(() => { const error = statusQuery.error; - if (error instanceof ApiError && error.status === 404) { - setStatusGone(true); + if (!(error instanceof ApiError && error.status === 404)) return; + if (resumedRef.current) { + if (releasingResumedSession) return; + setReleasingResumedSession(true); + const id = sessionId; + void (async () => { + if (id) { + await agentsApi.cancelClaudeSetupTokenLogin(companyId, id).catch(() => { + // The session is already gone either way; nothing more to do. + }); + } + setReleasingResumedSession(false); + clearActiveSession(); + })(); + return; } - }, [statusQuery.error]); + setStatusGone(true); + }, [statusQuery.error, releasingResumedSession, sessionId, companyId]); + + // Adopt the caller's active session once, on mount. This is what makes a + // page reload keep the session: with no local state at all, the panel would + // otherwise show its idle start state even though the server still holds an + // active login for this owner. + const resumeAttemptedRef = useRef(false); + useEffect(() => { + if (resumeAttemptedRef.current || !activeSessionQuery.isFetched) return; + resumeAttemptedRef.current = true; + const active = activeSessionQuery.data; + if (!active) return; + resumedRef.current = true; + resetLocalState(); + setSessionId(active.sessionId); + if (active.prompt) { + setAuthorizationUrl(active.prompt.authorizationUrl); + if (active.prompt.transportAdvisory) setTransportInsecure(true); + } + }, [activeSessionQuery.isFetched, activeSessionQuery.data]); // Poll the guarded prompt route until it returns the authorization URL. The // route returns 404 until the URL is ready, so the panel treats a 404 as @@ -2686,21 +2822,6 @@ function SubmittedBrowserCodeLoginPanel({ const isActive = Boolean(sessionId) && !isStored && !isFailure && !timedOut; const startDisabled = startLogin.isPending || isActive; - // Hold the active session id for the unmount cleanup. The panel updates it on - // every render. When the panel unmounts, or the parent removes it as the login - // closes, with an active, non-terminal session, the cleanup releases that - // session on the server. The ref is null once the session leaves the active - // state, so the cleanup never cancels a session the server already removed. - const activeSessionRef = useRef(null); - activeSessionRef.current = isActive ? sessionId : null; - - useEffect(() => { - return () => { - const id = activeSessionRef.current; - if (id) releaseServerSession(id); - }; - }, [releaseServerSession]); - // Cap the active login at the server deadline. The timer arms when the login // becomes active and clears when the login leaves the active state (a terminal // status, a stored success, or a new login). It re-arms when `expiresAt` @@ -2746,17 +2867,39 @@ function SubmittedBrowserCodeLoginPanel({ setBrowserCode(""); }; - // Start once, on mount, when the caller has already taken the press. Latched - // for the same reason as the displayed-code panel: a second start would burn - // an owner reservation, and here it would also rotate the stored token twice. + // Start once, on mount, when the caller has already taken the press, and + // only once the resume read has answered: a resumed session takes over + // instead of a fresh start. Latched for the same reason as the + // displayed-code panel: a second start would burn an owner reservation, and + // here it would also rotate the stored token twice. const autoStartedRef = useRef(false); const startLoginRef = useRef(startLogin.mutate); startLoginRef.current = startLogin.mutate; useEffect(() => { if (!autoStart || autoStartedRef.current) return; + // A failed lookup is not proof that no session exists: only a successful + // lookup is. Show the failure to the user instead of starting a second + // login the server would reject against the per-owner cap. + if (activeSessionQuery.isError) { + autoStartedRef.current = true; + setStartError( + activeSessionQuery.error instanceof Error + ? activeSessionQuery.error.message + : "Could not check for an active login.", + ); + return; + } + if (!activeSessionQuery.isSuccess) return; autoStartedRef.current = true; + if (activeSessionQuery.data) return; startLoginRef.current(); - }, [autoStart]); + }, [ + autoStart, + activeSessionQuery.isSuccess, + activeSessionQuery.isError, + activeSessionQuery.data, + activeSessionQuery.error, + ]); /** * Submit the pasted code without a press. @@ -2819,6 +2962,7 @@ function SubmittedBrowserCodeLoginPanel({ return ( ({ prompt: { url: "https://auth.openai.com/codex/device", code: "Q2RJ-E1YIF" }, })), cancelAdapterAuthLogin: vi.fn(async () => ({})), + // No default implementation: the top-level `beforeEach` sets the "no active + // session" 404 rejection, matching the real route. + getActiveAdapterAuthLoginSession: vi.fn(), + getActiveClaudeSetupTokenLoginSession: vi.fn(), })); // The adapter registry mock below always returns this function, so a test // can shape the built adapter config (e.g. a configured ANTHROPIC_API_KEY) @@ -326,6 +330,16 @@ describe("OnboardingWizard restore-gate (stale localStorage across accounts)", ( mockAgentsApi.getClaudeOAuthTokenStatus.mockRejectedValue( new ApiError("Not found", 404, null), ); + // Default: no active login session for the caller. A resume test + // overrides this with a resolved session body. + mockAgentsApi.getActiveAdapterAuthLoginSession.mockReset(); + mockAgentsApi.getActiveAdapterAuthLoginSession.mockRejectedValue( + new ApiError("Not found", 404, null), + ); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockReset(); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockRejectedValue( + new ApiError("Not found", 404, null), + ); // Reset to each mock's original default. `mockResolvedValue` / // `mockReturnValue` overrides a mock's implementation permanently — it // is not undone by `afterEach`'s `vi.clearAllMocks()`, which only clears @@ -2418,6 +2432,42 @@ describe("OnboardingWizard restore-gate (stale localStorage across accounts)", ( await act(async () => root.unmount()); }); + it("shows the failure instead of starting a codex_local login when the active-session read fails", async () => { + // A transient active-session lookup failure is not proof that no + // session exists — only a successful lookup that resolves to null is. + // The step must not start a second login the server would reject + // against the per-owner cap; it shows the failure instead. + mockAgentsApi.getAdapterAuthSignal.mockResolvedValue({ status: "unknown" }); + mockAgentsApi.getActiveAdapterAuthLoginSession.mockReset(); + mockAgentsApi.getActiveAdapterAuthLoginSession.mockRejectedValue( + new ApiError("Service unavailable", 503, null), + ); + const { root } = await openStep4({ adapterType: "codex_local" }); + + await pickSource(/OpenAI/); + + expect(mockAgentsApi.startAdapterAuthLogin).not.toHaveBeenCalled(); + expect(document.body.textContent).toContain("Service unavailable"); + + await act(async () => root.unmount()); + }); + + it("shows the failure instead of starting a claude_local login when the active-session read fails", async () => { + mockAgentsApi.getAdapterAuthSignal.mockResolvedValue({ status: "absent" }); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockReset(); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockRejectedValue( + new ApiError("Service unavailable", 503, null), + ); + const { root } = await openStep4({ adapterType: "claude_local" }); + + await pickSource(/Claude/); + + expect(mockAgentsApi.startClaudeSetupTokenLogin).not.toHaveBeenCalled(); + expect(document.body.textContent).toContain("Service unavailable"); + + await act(async () => root.unmount()); + }); + it("hires on Connect, with no sign-in, when the signal reports a ready credential", async () => { mockAgentsApi.getAdapterAuthSignal.mockResolvedValue({ status: "present" }); const { root } = await openStep4({ adapterType: "claude_local" }); @@ -2494,5 +2544,38 @@ describe("OnboardingWizard restore-gate (stale localStorage across accounts)", ( expect(mockAgentsApi.getAdapterAuthSignal).not.toHaveBeenCalled(); await act(async () => root.unmount()); }); + + it("restores the started-login state after a reload, resuming the active session with no press", async () => { + // A reload loses `connectPhase` — the draft deliberately does not carry + // it, because a login is a live server session with a deadline, not + // wizard state to replay blindly. The step must instead re-derive it + // from the caller's active session, so a customer who reloads mid-login + // sees their sign-in still running rather than the tile row again. + mockAgentsApi.getAdapterAuthSignal.mockResolvedValue({ status: "absent" }); + mockAgentsApi.getActiveClaudeSetupTokenLoginSession.mockResolvedValue({ + sessionId: "claude-session-1", + environmentId: "env-sandbox-1", + status: "waiting_for_user", + expiresAt: null, + failure: null, + panelMode: "submitted_browser_code", + prompt: { authorizationUrl: "https://claude.ai/oauth/authorize?code=true" }, + }); + const { root } = await openStep4({ adapterType: "claude_local" }); + // A reload's resume now runs one layer deeper than a fresh press: the + // step's own active-session read has to land before it opens the + // sequence, and only then does the mounted panel run its own resume + // read. Each is a further round trip `flushReact` has to catch up to. + for (let i = 0; i < 10; i += 1) await flushReact(); + + // No press: the resumed session is discovered on load and the card + // shows the login already running. + expect(mockAgentsApi.startClaudeSetupTokenLogin).not.toHaveBeenCalled(); + expect(document.body.textContent).toContain( + "Sign in to Claude then come back and enter authorization code", + ); + + await act(async () => root.unmount()); + }); }); }); diff --git a/ui/src/components/OnboardingWizard.tsx b/ui/src/components/OnboardingWizard.tsx index cf997c1f9f..91ac467f53 100644 --- a/ui/src/components/OnboardingWizard.tsx +++ b/ui/src/components/OnboardingWizard.tsx @@ -1106,6 +1106,46 @@ function OnboardingWizardInner({ const authSignalStatus = authSignalQuery.data?.status ?? null; const showAdapterLoginPanel = canShowAdapterLogin && (authSignalStatus === "absent" || authSignalStatus === "unknown"); + /** + * Restores the connect sequence after a reload. + * + * The panel resumes an active session on its own mount, but this step only + * mounts the panel once the sequence has moved off `idle` — and a reload + * starts the sequence at `idle` again, deliberately: `connectPhase` is not + * in the draft. Without this read, a reload during a login would leave the + * panel unmounted and the resumed session unreachable from this step. A 404 + * means no active session for the caller. + */ + const activeLoginSessionQuery = useQuery({ + queryKey: createdCompanyId + ? queryKeys.agents.activeLoginSession(createdCompanyId, adapterType) + : ["agents", "none", "active-login-session", adapterType], + queryFn: async () => { + try { + return adapterCaps.login?.panelMode === "submitted_browser_code" + ? await agentsApi.getActiveClaudeSetupTokenLoginSession(createdCompanyId!) + : await agentsApi.getActiveAdapterAuthLoginSession(createdCompanyId!, adapterType); + } catch (error) { + if (error instanceof ApiError && error.status === 404) return null; + throw error; + } + }, + enabled: + Boolean(createdCompanyId) && effectiveOnboardingOpen && step === 4 && canShowAdapterLogin, + }); + useEffect(() => { + if (!activeLoginSessionQuery.data) return; + // Re-derive the row's answer along with the sequence: a resumed session + // implies a source was already picked, and the row stays a question + // otherwise (see `sourcePicked` above). + setSourcePicked(true); + // Skip straight past the collapsing beat — that animation is for a press + // landing on a step already on screen, not for a reload that should show + // the running login at once. The panel's own mount reports the resumed + // prompt through `onPromptReady`, below, which is what moves this beat + // from `loading` to `ready`, exactly as a fresh press would. + setConnectPhase((phase) => (phase === "idle" ? "loading" : phase)); + }, [activeLoginSessionQuery.data]); /** * The signal is being fetched and has not answered yet. * @@ -1353,20 +1393,18 @@ function OnboardingWizardInner({ /** * Back, on the connect step, unwinds the sign-in before it leaves the step. * - * With no Cancel on the card this is the only way out, and what it undoes - * depends on how far in you are. Unmounting the panel is what releases the - * server session — see the release-on-unmount effect in `AdapterLoginPanel` - * — so the card leaving is the cancel, not a separate call. + * This hides the card; it does not cancel the login. Unmounting the panel no + * longer releases the server session — the session stays reachable for a + * later resume, the same read that restores it after a reload — so backing + * out and returning shows the sign-in still running, not a fresh one. The + * card's own Cancel button is the only explicit release; `onCancel` below + * puts the step back here when it fires. */ function unwindConnectStep() { setConnectAuthUrl(null); // Where the reverse starts depends on how far the sequence got. Backing out - // during the collapse has no card to close and no room to give back, and - // entering `unwindCard` regardless mounted the panel — which starts a - // server login on mount — purely so the unmount could cancel it. Should - // that cancel fail, the reservation is held to the server deadline and an - // immediate retry cannot start. With no card open, the row is the whole of - // the unwind. + // during the collapse has no card to close and no room to give back. + // With no card open, the row is the whole of the unwind. setConnectPhase(connectCardLive ? "unwindCard" : "unwindRow"); } @@ -1502,7 +1540,9 @@ function OnboardingWizardInner({ * Nothing else needs to reset it. A source can only change by being picked, * and picking sets the phase itself; the credential mode can only change * before the sequence starts, because its control is inert once the row has - * collapsed. + * collapsed. Switching either no longer needs its own reset: the panel keeps + * its server session reachable for a later resume instead of releasing it on + * the remount, so there is nothing left here for that change to undo. */ useEffect(() => { if (step === 4) return; @@ -3131,9 +3171,11 @@ function OnboardingWizardInner({ in the connect step's chrome. It owns the session; the step owns the sequence around it. - Unmounting it is the cancel: the panel releases its - server session on unmount, so Back closing the card is - what frees the owner's reservation. + Unmounting it is no longer the cancel: the session + stays reachable for a later resume, so Back and a + source switch only hide the card. `onCancel` fires + from the card's own Cancel button, the one explicit + release, and puts the step back where Back would. No "Use saved login" control: the hire step already applies a stored login on its own. */ @@ -3144,6 +3186,7 @@ function OnboardingWizardInner({ environmentId={resolvedLoginEnvironmentId} chrome="onboarding" autoStart + onCancel={unwindConnectStep} onPromptReady={(url) => { setConnectAuthUrl(url); // The prompt arriving is what ends the waiting beat. diff --git a/ui/src/lib/queryKeys.ts b/ui/src/lib/queryKeys.ts index 3ce435a550..407c00ef80 100644 --- a/ui/src/lib/queryKeys.ts +++ b/ui/src/lib/queryKeys.ts @@ -224,6 +224,8 @@ export const queryKeys = { ["agents", companyId, "detect-model", adapterType] as const, authSignal: (companyId: string, adapterType: string, environmentId?: string | null) => ["agents", companyId, "auth-signal", adapterType, environmentId ?? null] as const, + activeLoginSession: (companyId: string, adapterType: string) => + ["agents", companyId, "active-login-session", adapterType] as const, }, builtInAgents: { list: (companyId: string) => ["built-in-agents", companyId] as const,