From cad0f41b37f54bcd9fd3faf0056bebffd731b6d5 Mon Sep 17 00:00:00 2001 From: Cody Date: Thu, 2 Jul 2026 19:39:12 +0000 Subject: [PATCH] Fix custom image terminal session expiry Separate short-lived terminal connect token expiry from the live browser SSH session deadline so established setup sessions are not closed after the five-minute token TTL. Add focused coverage for live sessions surviving token expiry and closing at setup expiry. Co-Authored-By: Paperclip --- ...vironment-custom-image-terminal-ws.test.ts | 49 +++++++++++++++++++ .../environment-custom-image-terminal-ws.ts | 6 +-- server/src/routes/environments.ts | 5 +- ...ent-custom-image-terminal-sessions.test.ts | 9 ++-- ...ironment-custom-image-terminal-sessions.ts | 20 +++++--- ui/src/pages/CompanyEnvironments.tsx | 29 ++++++++++- 6 files changed, 101 insertions(+), 17 deletions(-) diff --git a/server/src/__tests__/environment-custom-image-terminal-ws.test.ts b/server/src/__tests__/environment-custom-image-terminal-ws.test.ts index b0119bb67d2..03aa9329988 100644 --- a/server/src/__tests__/environment-custom-image-terminal-ws.test.ts +++ b/server/src/__tests__/environment-custom-image-terminal-ws.test.ts @@ -108,6 +108,10 @@ async function waitForAssertion(assertion: () => void) { throw lastError; } +async function waitForDuration(ms: number) { + await new Promise((resolve) => setTimeout(resolve, ms)); +} + function deferred() { let resolve!: (value: T) => void; let reject!: (err: Error) => void; @@ -370,6 +374,51 @@ describe("custom image terminal websocket bridge", () => { }); }); + it("keeps established terminal sessions alive past connect-token expiry and closes them at setup expiry", async () => { + const setupExpiresAt = new Date(Date.now() + 2500); + customImages.getSessionById.mockResolvedValue(createSession({ expiresAt: setupExpiresAt })); + customImages.refreshSetupSession.mockResolvedValue({ + session: createSession({ expiresAt: setupExpiresAt }), + connectionPayload: { + type: "ssh", + command: "ssh fresh-token@fresh.example.test -p 2200", + expiresAt: futureDate(30).toISOString(), + }, + }); + const { port } = await startHarness(); + const minted = sessionStore.create({ + setupSessionId: "session-1", + companyId: "company-1", + environmentId: "env-1", + provider: "daytona", + ssh: { username: "old-token", host: "old.example.test", port: 22 }, + setupExpiresAt, + now: new Date(Date.now() - 5 * 60 * 1000 + 750), + }); + expect(minted.session.connectExpiresAt.getTime()).toBeLessThan(minted.session.sessionExpiresAt.getTime()); + + const ws = new WebSocket(terminalUrl(port, { + terminalSessionId: minted.session.id, + token: minted.token, + })); + let closed = false; + ws.on("close", () => { + closed = true; + }); + const readyPromise = waitForJsonMessage(ws, (frame) => frame.type === "ready"); + const closePromise = waitForClose(ws); + + await waitForOpen(ws); + await readyPromise; + await waitForDuration(Math.max(0, minted.session.connectExpiresAt.getTime() - Date.now()) + 400); + expect(closed).toBe(false); + expect(fakeShell.closeCalls).toBe(0); + + await closePromise; + expect(closed).toBe(true); + expect(fakeShell.closeCalls).toBeGreaterThan(0); + }); + it("applies the latest resize sent while the SSH shell is still opening", async () => { const pendingShell = deferred(); connector.connect.mockReturnValueOnce(pendingShell.promise); diff --git a/server/src/realtime/environment-custom-image-terminal-ws.ts b/server/src/realtime/environment-custom-image-terminal-ws.ts index 84934e76cbc..403fc6e3f17 100644 --- a/server/src/realtime/environment-custom-image-terminal-ws.ts +++ b/server/src/realtime/environment-custom-image-terminal-ws.ts @@ -300,7 +300,7 @@ async function validateTerminalUpgrade(input: { input.sessionStore.delete(terminalSession.id); throw conflict(`Cannot open terminal for setup status "${refreshed.session.status}".`); } - requireFutureSetupExpiry(refreshed.session, input.now); + const sessionExpiresAt = requireFutureSetupExpiry(refreshed.session, input.now); const payloadValidation = validateCustomImageSetupSshPayload(refreshed.connectionPayload, input.now); if (!payloadValidation.ok) { @@ -308,7 +308,7 @@ async function validateTerminalUpgrade(input: { throw terminalPayloadValidationError(payloadValidation); } - return { ...terminalSession, ssh: payloadValidation.ssh }; + return { ...terminalSession, ssh: payloadValidation.ssh, sessionExpiresAt }; } class Ssh2Shell implements EnvironmentCustomImageSshShell { @@ -490,7 +490,7 @@ export function setupEnvironmentCustomImageTerminalWebSocketServer( }, }); - const expiresInMs = Math.max(0, context.terminalSession.expiresAt.getTime() - Date.now()); + const expiresInMs = Math.max(0, context.terminalSession.sessionExpiresAt.getTime() - Date.now()); expiryTimer = setTimeout(() => { closeTerminal("expired", 1008, "expired"); }, expiresInMs); diff --git a/server/src/routes/environments.ts b/server/src/routes/environments.ts index 3f66da5c602..6c3f05e3c84 100644 --- a/server/src/routes/environments.ts +++ b/server/src/routes/environments.ts @@ -466,14 +466,15 @@ export function environmentRoutes( session: setupSessionActivityDetails(refreshed.session), terminalSession: { connectionType: "ssh", - expiresAt: minted.session.expiresAt.toISOString(), + connectExpiresAt: minted.session.connectExpiresAt.toISOString(), + sessionExpiresAt: minted.session.sessionExpiresAt.toISOString(), }, }, }); res.status(201).json({ id: minted.session.id, token: minted.token, - expiresAt: minted.session.expiresAt.toISOString(), + expiresAt: minted.session.connectExpiresAt.toISOString(), setupSessionId: minted.session.setupSessionId, environmentId: minted.session.environmentId, connectionType: "ssh", diff --git a/server/src/services/environment-custom-image-terminal-sessions.test.ts b/server/src/services/environment-custom-image-terminal-sessions.test.ts index 23c1034dcb6..66a3d14efe4 100644 --- a/server/src/services/environment-custom-image-terminal-sessions.test.ts +++ b/server/src/services/environment-custom-image-terminal-sessions.test.ts @@ -108,7 +108,7 @@ describe("validateCustomImageSetupSshPayload", () => { }); describe("EnvironmentCustomImageTerminalSessionStore", () => { - it("mints opaque tokens and caps expiry by setup and provider payload expiry", () => { + it("mints opaque connect tokens and tracks live session expiry separately", () => { const store = new EnvironmentCustomImageTerminalSessionStore(); const now = new Date("2026-06-25T20:00:00.000Z"); const minted = store.create({ @@ -118,13 +118,14 @@ describe("EnvironmentCustomImageTerminalSessionStore", () => { provider: "daytona", ssh: { username: "ssh-token-secret", host: "203.0.113.10", port: 2222 }, setupExpiresAt: new Date("2026-06-25T20:30:00.000Z"), - connectionExpiresAt: new Date("2026-06-25T20:02:00.000Z"), + connectionExpiresAt: new Date("2026-06-25T20:10:00.000Z"), now, }); expect(minted.token).toHaveLength(43); expect(minted.session.id).toMatch(/^[0-9a-f-]{36}$/); - expect(minted.session.expiresAt.toISOString()).toBe("2026-06-25T20:02:00.000Z"); + expect(minted.session.connectExpiresAt.toISOString()).toBe("2026-06-25T20:05:00.000Z"); + expect(minted.session.sessionExpiresAt.toISOString()).toBe("2026-06-25T20:30:00.000Z"); expect(store.get({ id: minted.session.id, token: minted.token, @@ -140,7 +141,7 @@ describe("EnvironmentCustomImageTerminalSessionStore", () => { expect(store.get({ id: minted.session.id, token: minted.token, - }, new Date("2026-06-25T20:02:00.000Z"))).toBeNull(); + }, new Date("2026-06-25T20:05:00.000Z"))).toBeNull(); }); it("deletes all tokens for a setup session", () => { diff --git a/server/src/services/environment-custom-image-terminal-sessions.ts b/server/src/services/environment-custom-image-terminal-sessions.ts index 1a77a3866f1..b5b8e850e0d 100644 --- a/server/src/services/environment-custom-image-terminal-sessions.ts +++ b/server/src/services/environment-custom-image-terminal-sessions.ts @@ -38,7 +38,8 @@ export interface EnvironmentCustomImageTerminalSessionRecord { connectionType: "ssh"; ssh: ParsedCustomImageSetupSshCommand; createdAt: Date; - expiresAt: Date; + connectExpiresAt: Date; + sessionExpiresAt: Date; } export interface MintedEnvironmentCustomImageTerminalSession { @@ -192,12 +193,16 @@ export class EnvironmentCustomImageTerminalSessionStore { const now = input.now ?? new Date(); this.cleanupExpired(now); + const setupExpiresAt = toValidFutureDate(input.setupExpiresAt, now); + if (!setupExpiresAt) { + throw new Error("Terminal sessions require a future setup session expiry."); + } const candidateExpirations = [ new Date(now.getTime() + DEFAULT_TERMINAL_SESSION_TOKEN_TTL_MS), - toValidFutureDate(input.setupExpiresAt, now), + setupExpiresAt, toValidFutureDate(input.connectionExpiresAt, now), ].filter((date): date is Date => date !== null); - const expiresAt = minDate(candidateExpirations); + const connectExpiresAt = minDate(candidateExpirations); const token = randomBytes(TERMINAL_SESSION_TOKEN_BYTES).toString("base64url"); const id = randomUUID(); const session: EnvironmentCustomImageTerminalSessionRecord = { @@ -209,7 +214,8 @@ export class EnvironmentCustomImageTerminalSessionStore { connectionType: "ssh", ssh: input.ssh, createdAt: now, - expiresAt, + connectExpiresAt, + sessionExpiresAt: setupExpiresAt, }; this.sessionsById.set(id, { tokenHash: hashTerminalSessionToken(token), @@ -223,7 +229,7 @@ export class EnvironmentCustomImageTerminalSessionStore { const stored = this.sessionsById.get(input.id) ?? null; if (!stored) return null; if (stored.tokenHash !== hashTerminalSessionToken(input.token)) return null; - if (stored.session.expiresAt.getTime() <= now.getTime()) { + if (stored.session.connectExpiresAt.getTime() <= now.getTime()) { this.sessionsById.delete(input.id); return null; } @@ -234,7 +240,7 @@ export class EnvironmentCustomImageTerminalSessionStore { if (!id) return null; const stored = this.sessionsById.get(id) ?? null; if (!stored) return null; - if (stored.session.expiresAt.getTime() <= now.getTime()) { + if (stored.session.connectExpiresAt.getTime() <= now.getTime()) { this.sessionsById.delete(id); return null; } @@ -260,7 +266,7 @@ export class EnvironmentCustomImageTerminalSessionStore { cleanupExpired(now = new Date()): number { let removed = 0; for (const [id, stored] of this.sessionsById) { - if (stored.session.expiresAt.getTime() <= now.getTime()) { + if (stored.session.connectExpiresAt.getTime() <= now.getTime()) { this.sessionsById.delete(id); removed += 1; } diff --git a/ui/src/pages/CompanyEnvironments.tsx b/ui/src/pages/CompanyEnvironments.tsx index f10c28f7bde..f813c633913 100644 --- a/ui/src/pages/CompanyEnvironments.tsx +++ b/ui/src/pages/CompanyEnvironments.tsx @@ -288,6 +288,33 @@ function customImageTerminalStatusCopy(state: CustomImageTerminalConnectionState } } +function customImageTerminalCloseReasonCopy(reason: unknown) { + if ( + reason !== "expired" + && reason !== "ssh_closed" + && reason !== "server_shutdown" + && reason !== "setup_finished" + && reason !== "setup_cancelled" + ) { + return typeof reason === "string" && reason.trim() ? "Terminal closed." : null; + } + + switch (reason) { + case "expired": + return "Setup session expired."; + case "ssh_closed": + return "SSH session closed."; + case "server_shutdown": + return "Terminal server shut down."; + case "setup_finished": + return "Setup session finished."; + case "setup_cancelled": + return "Setup session cancelled."; + default: + return null; + } +} + function EnvironmentCustomImageBrowserTerminal({ autoConnect = false, sessionId, @@ -516,7 +543,7 @@ function EnvironmentCustomImageBrowserTerminal({ if (frame.type === "closed") { setConnectionState("closed"); - setErrorMessage(typeof frame.reason === "string" ? frame.reason : null); + setErrorMessage(customImageTerminalCloseReasonCopy(frame.reason)); } };