diff --git a/packages/electron-app/electron/main/client-state-lifecycle.test.ts b/packages/electron-app/electron/main/client-state-lifecycle.test.ts index 9f65b211..89b0bd9a 100644 --- a/packages/electron-app/electron/main/client-state-lifecycle.test.ts +++ b/packages/electron-app/electron/main/client-state-lifecycle.test.ts @@ -72,13 +72,16 @@ test("late old-window detach preserves replacement tracker during shutdown", asy assert.deepEqual(h.calls, ["hide", "renderer", "replacement-native", "stop", "release"]) }) -test("Windows session end vetoes termination until cleanup exits explicitly", async () => { +test("Windows session end starts cleanup without vetoing termination", async () => { const h = harness() let prevented = false h.windows.get("query-session-end")?.({ preventDefault: () => { prevented = true } }) + await (h.lifecycle as any).sessionEndPreparation + assert.equal(prevented, false) + assert.deepEqual(h.calls, ["renderer", "native"]) + assert.equal(h.exits(), 0) h.windows.get("session-end")?.() await (h.lifecycle as any).sessionEnd; await tick() - assert.equal(prevented, true) assert.deepEqual(h.calls, ["renderer", "native", "stop", "release"]) assert.equal(h.exits(), 1) }) @@ -88,8 +91,9 @@ test("session end force-exits after the bounded window when an ordinary shutdown let prevented = false h.appEvents.get("before-quit")?.({ preventDefault: () => {} }) h.windows.get("query-session-end")?.({ preventDefault: () => { prevented = true } }) + h.windows.get("session-end")?.() await delay(25) - assert.equal(prevented, true) + assert.equal(prevented, false) assert.deepEqual(h.calls, ["hide", "renderer", "release"]) assert.equal(h.exits(), 1) }) @@ -120,6 +124,7 @@ test("Windows session-end rejection fails open at the bounded deadline", async ( const h = harness({ stop: async () => { throw new Error("unconfirmed") }, sessionEndCleanupTimeoutMs: 10 }) h.appEvents.get("before-quit")?.({ preventDefault: () => {} }) h.windows.get("query-session-end")?.({ preventDefault: () => {} }) + h.windows.get("session-end")?.() await delay(25) assert.equal(h.exits(), 1) assert.deepEqual(h.calls, ["hide", "renderer", "native", "stop", "release"]) @@ -133,6 +138,7 @@ test("Windows fail-open bounds a hanging primary release before app.exit", async sessionEndReleaseTimeoutMs: 10, }) h.windows.get("query-session-end")?.({ preventDefault: () => {} }) + h.windows.get("session-end")?.() await delay(25) assert.deepEqual(h.calls, ["renderer", "release"]) diff --git a/packages/electron-app/electron/main/client-state-lifecycle.ts b/packages/electron-app/electron/main/client-state-lifecycle.ts index c0e7405c..023436e1 100644 --- a/packages/electron-app/electron/main/client-state-lifecycle.ts +++ b/packages/electron-app/electron/main/client-state-lifecycle.ts @@ -21,6 +21,8 @@ interface ClientStateLifecycleDependencies { export class ClientStateLifecycle { private shutdown: Promise | null = null private sessionEnd: Promise | null = null + private sessionEndPreparation: Promise | null = null + private sessionEndPreparationPending = false private exitAllowed = false private trackedMainWindow: BrowserWindow | null = null private windowStateTracker: WindowStateTracker | null = null @@ -62,10 +64,8 @@ export class ClientStateLifecycle { }) if (this.dependencies.isWindows ?? process.platform === "win32") { - window.on("query-session-end", (event) => { - if (this.exitAllowed) return - event.preventDefault() - this.promoteToSessionEnd(window) + window.on("query-session-end", () => { + this.prepareSessionEnd(window) }) window.on("session-end", () => this.promoteToSessionEnd(window)) } @@ -96,11 +96,10 @@ export class ClientStateLifecycle { await this.runStage("native main-window close flush", () => this.flushNative()) } - private startShutdown(window: BrowserWindow | null): Promise { + private startShutdown(window: BrowserWindow | null, preparedFlush?: Promise): Promise { if (this.shutdown) return this.shutdown const stages = (async () => { - await this.runStage("renderer shutdown flush", () => this.flushRenderer(window)) - await this.runStage("native shutdown flush", () => this.flushNative()) + await (preparedFlush ?? this.flushForShutdown(window)) await this.dependencies.cliManager.shutdown() await this.releasePrimary() })() @@ -111,6 +110,21 @@ export class ClientStateLifecycle { return this.shutdown } + private async flushForShutdown(window: BrowserWindow | null): Promise { + await this.runStage("renderer shutdown flush", () => this.flushRenderer(window)) + await this.runStage("native shutdown flush", () => this.flushNative()) + } + + private prepareSessionEnd(window: BrowserWindow): void { + if (this.exitAllowed || this.sessionEnd || this.shutdown || this.sessionEndPreparationPending) return + this.sessionEndPreparationPending = true + const preparation = this.flushForShutdown(window) + this.sessionEndPreparation = preparation + void preparation.finally(() => { + if (this.sessionEndPreparation === preparation) this.sessionEndPreparationPending = false + }) + } + private hideWindows(): void { for (const window of this.dependencies.getAllWindows()) { if (!window.isDestroyed()) { @@ -131,7 +145,7 @@ export class ClientStateLifecycle { private promoteToSessionEnd(window: BrowserWindow): void { if (this.exitAllowed || this.sessionEnd) return - const cleanup = this.startShutdown(window) + const cleanup = this.shutdown ?? this.startShutdown(window, this.sessionEndPreparation ?? this.flushForShutdown(window)) this.sessionEnd = new Promise((resolve) => { const timeoutMs = this.dependencies.sessionEndCleanupTimeoutMs ?? 5_000 const releaseTimeoutMs = Math.min(timeoutMs, this.dependencies.sessionEndReleaseTimeoutMs ?? 250) diff --git a/packages/electron-app/electron/main/client-state-navigation.test.ts b/packages/electron-app/electron/main/client-state-navigation.test.ts index 8c40b0a9..a8d3f289 100644 --- a/packages/electron-app/electron/main/client-state-navigation.test.ts +++ b/packages/electron-app/electron/main/client-state-navigation.test.ts @@ -87,3 +87,20 @@ test("queued navigation preserves order and distinct generations", async () => { await Promise.all([first, second]) assert.deepEqual(calls, ["start-1", "end-1", "run-2"]) }) + +test("queued navigation exposes whether work was invalidated before it mutates navigation state", async () => { + const calls: string[] = [] + let release!: () => void + const gate = new Promise((resolve) => { release = resolve }) + const navigation = controller(window(), { isPrimary: true }) + const first = navigation.navigate(async (_window, generation) => { + await gate + if (navigation.isCurrent(generation)) calls.push("stale") + }) + const second = navigation.navigate((_window, generation) => { + if (navigation.isCurrent(generation)) calls.push("current") + }) + release() + await Promise.all([first, second]) + assert.deepEqual(calls, ["current"]) +}) diff --git a/packages/electron-app/electron/main/client-state-navigation.ts b/packages/electron-app/electron/main/client-state-navigation.ts index feed061d..a4aa6da3 100644 --- a/packages/electron-app/electron/main/client-state-navigation.ts +++ b/packages/electron-app/electron/main/client-state-navigation.ts @@ -24,6 +24,10 @@ export class ClientStateNavigationController { return request } + isCurrent(generation: number): boolean { + return generation === this.generation + } + private async performNavigation( operation: (window: BrowserWindow, generation: number) => void | Promise, generation: number, diff --git a/packages/electron-app/electron/main/client-state-partitions.ts b/packages/electron-app/electron/main/client-state-partitions.ts index 14d03a8a..df4a482c 100644 --- a/packages/electron-app/electron/main/client-state-partitions.ts +++ b/packages/electron-app/electron/main/client-state-partitions.ts @@ -6,8 +6,8 @@ import { hasErrorCode } from "./client-state-process" export const CLIENT_STATE_PARTITION_PROTOCOL_VERSION = 1 export const CLIENT_STATE_PARTITION_ENVELOPE_VERSION = 2 export const MAX_CLIENT_STATE_ROOT_BYTES = 1024 * 1024 +export const MAX_CLIENT_STATE_PARTITION_COMMIT_BYTES = 256 * 1024 * 1024 const MAX_PARTITION_BYTES = 1024 * 1024 -const MAX_COMMIT_BYTES = 8 * 1024 * 1024 const MAX_PARTITION_KEYS = 4096 const PARTITION_KEY = /^[0-9a-f]{64}$/ const PARTITION_DIRECTORY = "partitions" @@ -87,7 +87,7 @@ export function validateClientStatePartitionCommit(value: unknown): ValidatedCli const size = Buffer.byteLength(content, "utf8") if (size > MAX_PARTITION_BYTES) throw new RangeError("Client state partition exceeds the 1 MiB limit") commitBytes += size - if (commitBytes > MAX_COMMIT_BYTES) throw new RangeError("Client state partition commit exceeds the 8 MiB limit") + if (commitBytes > MAX_CLIENT_STATE_PARTITION_COMMIT_BYTES) throw new RangeError("Client state partition commit exceeds the 256 MiB limit") if (digest(content) !== key) throw new TypeError("Client state partition digest mismatch") } return { snapshot: normalizedSnapshot, partitions, partitionKeys: rootPartitionKeys } diff --git a/packages/electron-app/electron/main/client-state.test.ts b/packages/electron-app/electron/main/client-state.test.ts index 4f662494..0796e49c 100644 --- a/packages/electron-app/electron/main/client-state.test.ts +++ b/packages/electron-app/electron/main/client-state.test.ts @@ -7,6 +7,7 @@ import { join } from "node:path" import test from "node:test" import { ClientStateManager, type ClientStateWriter } from "./client-state" import { deterministicLegacyWindowId, parseClientState } from "./client-state-envelope" +import { MAX_CLIENT_STATE_PARTITION_COMMIT_BYTES } from "./client-state-partitions" test("legacy migration UUIDs are deterministic from exact envelope bytes", () => { const vectors = [ @@ -537,6 +538,7 @@ test("invalid V3 remains byte-frozen until explicit clear", async (t) => { }) test("partition commits validate protocol and hashes", async (t) => { + assert.equal(MAX_CLIENT_STATE_PARTITION_COMMIT_BYTES, 256 * 1024 * 1024) const manager = harness(t, { version: 1, restoreEnabled: true }).create() const content = "partition" const key = partitionKey(content) @@ -561,16 +563,16 @@ test("partition commits validate protocol and hashes", async (t) => { const partitions: Record = {} const partitionKeys: string[] = [] - for (let index = 0; index < 8; index++) { + for (let index = 0; index < 9; index++) { const value = `${index}${"x".repeat(1024 * 1024 - 1)}` const valueKey = partitionKey(value) partitions[valueKey] = value partitionKeys.push(valueKey) } partitionKeys.sort() - assert.throws(() => manager.commitClientStatePartitions({ + assert.equal(await manager.commitClientStatePartitions({ protocolVersion: 1, snapshot: partitionRoot(partitionKeys), partitions, partitionKeys, - }), /exceeds the 8 MiB limit/) + }), true) }) test("partition commit/read preserves the old root on failure and clear sweeps", async (t) => { diff --git a/packages/electron-app/electron/main/ipc.ts b/packages/electron-app/electron/main/ipc.ts index 45035d10..cf939b85 100644 --- a/packages/electron-app/electron/main/ipc.ts +++ b/packages/electron-app/electron/main/ipc.ts @@ -4,6 +4,7 @@ import { requestMicrophoneAccess } from "./permissions" import type { CliProcessManager } from "./process-manager" import { openWorkspaceTarget, type WorkspaceEditor, type WorkspaceOpenTarget } from "./workspace-open" import { setWorkspaceMenuEnabled } from "./menu" +import { requireHttpUrl } from "./navigation-security" interface LocalSender { id: string @@ -171,6 +172,8 @@ export function setupCliIPC(cliManager: CliProcessManager, dependencies: CliIPCD || typeof payload.skipTlsVerify !== "boolean") { throw new Error("Invalid remote window request") } + requireHttpUrl(payload.baseUrl, "baseUrl") + if (payload.entryUrl !== undefined) requireHttpUrl(payload.entryUrl, "entryUrl") await dependencies.openRemoteWindow(payload) return { ok: true } }) diff --git a/packages/electron-app/electron/main/main.ts b/packages/electron-app/electron/main/main.ts index 4dcf513f..4c0a0b29 100644 --- a/packages/electron-app/electron/main/main.ts +++ b/packages/electron-app/electron/main/main.ts @@ -3,7 +3,7 @@ import http from "node:http" import https from "node:https" import { existsSync, mkdirSync, rmSync } from "node:fs" import { dirname, join } from "node:path" -import { fileURLToPath } from "node:url" +import { fileURLToPath, pathToFileURL } from "node:url" import { ClientStateManager } from "./client-state" import { setupClientStateIPC } from "./client-state-ipc" import { ClientStateNavigationController } from "./client-state-navigation" @@ -12,9 +12,10 @@ import { LocalWindowRegistry, type LocalWindowRecord } from "./local-window-regi import { clearWorkspaceMenuWindow, createApplicationMenu, setWorkspaceMenuEnabled } from "./menu" import { resolveFocusedLocalTarget, resolveWindowTarget } from "./menu-target" import { MultiwindowLifecycle } from "./multiwindow-lifecycle" +import { decideNavigation, requireHttpUrl } from "./navigation-security" import { configureMediaPermissionHandlers, isAllowedRendererOrigin } from "./permissions" import { CliProcessManager } from "./process-manager" -import { RemoteWindowRegistry } from "./remote-window-registry" +import { navigateReusedRemoteWindow, RemoteWindowRegistry } from "./remote-window-registry" import { resolveConfiguredRendererOrigins } from "./renderer-origin" import { allocateLocalWindowIdentity, BackendBootstrapCoordinator, createLaunchIntentQueue, isRemoteCertificateAllowed, parseLaunchIntent, resolveRemoteSessionPartition, resolveStorageScope, startPrimaryInstance, type LaunchIntent } from "./startup" import { clampWindowBounds, DEFAULT_WINDOW_HEIGHT, DEFAULT_WINDOW_WIDTH, installWindowZoomInput, restoreWindowState, WindowStateTracker } from "./window-state" @@ -95,13 +96,21 @@ function runPrimary(firstIntent: LaunchIntent) { const candidates = [join(app.getAppPath(), "dist/renderer/loading.html"), join(process.resourcesPath, "dist/renderer/loading.html"), join(mainDirname, "../dist/renderer/loading.html")] return { file: candidates.find(existsSync) ?? candidates[0] } } + const getLoadingUrl = () => { + const target = loadingTarget() + return target.url ?? pathToFileURL(target.file!).toString() + } const loadLoading = async (record: LocalWindowRecord, force = false) => { if (record.window.isDestroyed() || (record.loading && !force)) return record.loading = true - record.backendUrl = null - remoteOrigins.delete(record.window.id) const target = loadingTarget() - await record.navigation.navigate((window) => target.url ? window.loadURL(target.url) : window.loadFile(target.file!)).catch((error) => { + await record.navigation.navigate(async (window, generation) => { + if (!record.navigation.isCurrent(generation)) return + await (target.url ? window.loadURL(target.url) : window.loadFile(target.file!)) + if (!record.navigation.isCurrent(generation)) return + record.backendUrl = null + remoteOrigins.delete(record.window.id) + }).catch((error) => { if (!isIgnorableNavigationError(error)) console.error("[cli] failed to load loading screen", error) }) } @@ -109,14 +118,21 @@ function runPrimary(firstIntent: LaunchIntent) { if (record.window.isDestroyed() || (!record.loading && record.backendUrl === url)) return let origin: string try { origin = new URL(url).origin } catch { return } - const previous = remoteOrigins.get(record.window.id) - remoteOrigins.set(record.window.id, new Set([...(previous ?? []), origin])) - await record.navigation.navigate((window) => window.loadURL(url)).then(() => { + await record.navigation.navigate(async (window, generation) => { + if (!record.navigation.isCurrent(generation)) return + const previous = remoteOrigins.get(record.window.id) + remoteOrigins.set(record.window.id, new Set([...(previous ?? []), origin])) + try { await window.loadURL(url) } catch (error) { + if (record.navigation.isCurrent(generation)) { + if (previous) remoteOrigins.set(record.window.id, previous); else remoteOrigins.delete(record.window.id) + } + throw error + } + if (!record.navigation.isCurrent(generation)) return record.loading = false record.backendUrl = url remoteOrigins.set(record.window.id, new Set([origin])) - }, (error) => { - if (previous) remoteOrigins.set(record.window.id, previous); else remoteOrigins.delete(record.window.id) + }).catch((error) => { if (!isIgnorableNavigationError(error)) console.error("[cli] failed to load backend", error) }) } @@ -154,7 +170,7 @@ function runPrimary(firstIntent: LaunchIntent) { bindClientState(window) lifecycle.attach(record) installWindowZoomInput(window, (level) => tracker ? tracker.setZoomLevel(level) : window.webContents.setZoomLevel(level)) - setupNavigationGuards(window, navigation, getAllowedOrigins) + setupNavigationGuards(window, navigation, getAllowedOrigins, getLoadingUrl) window.webContents.on("did-start-navigation", (_event, _url, _isInPlace, isMainFrame) => { if (isMainFrame) setWorkspaceMenuEnabled(window, false) }) @@ -228,7 +244,7 @@ function runPrimary(firstIntent: LaunchIntent) { bootstrap.reset() backendUrl = null backendTargetUrl = null - for (const record of registry.all()) void loadLoading(record) + for (const record of registry.all()) void loadLoading(record, true) } }) cli.on("error", (error) => registry.fanout("cli:error", { message: error.message })) @@ -267,17 +283,14 @@ function runPrimary(firstIntent: LaunchIntent) { return candidates.find(existsSync) ?? candidates[0] } async function openRemoteWindow(payload: { id: string; name: string; baseUrl: string; entryUrl?: string; proxySessionId?: string; skipTlsVerify: boolean }) { - const base = new URL(payload.baseUrl) - const target = new URL(payload.entryUrl ?? payload.baseUrl) + const base = requireHttpUrl(payload.baseUrl, "baseUrl") + const target = requireHttpUrl(payload.entryUrl ?? payload.baseUrl, "entryUrl") const title = `${payload.name} - ${payload.baseUrl}` const existing = remoteWindows.reuse(payload.id, payload.proxySessionId) if (existing) { const allowedOrigins = new Set([base.origin, target.origin]) - remoteOrigins.set(existing.id, allowedOrigins) - if (payload.skipTlsVerify) insecureOrigins.set(existing.webContents.id, allowedOrigins) - else insecureOrigins.delete(existing.webContents.id) existing.setTitle(title) - await existing.loadURL(target.toString()) + await navigateReusedRemoteWindow(existing, target, allowedOrigins, remoteOrigins, insecureOrigins, payload.skipTlsVerify) return } const remoteSession = session.fromPartition(resolveRemoteSessionPartition(payload.id, payload.proxySessionId)) @@ -294,26 +307,40 @@ function runPrimary(firstIntent: LaunchIntent) { .flatMap((candidate) => [...(remoteOrigins.get(candidate.id) ?? [])]), remoteSession) window.setTitle(title) window.webContents.on("page-title-updated", (event) => { event.preventDefault(); window.setTitle(title) }) - setupNavigationGuards(window, undefined, getAllowedOrigins) + setupNavigationGuards(window, undefined, getAllowedOrigins, getLoadingUrl) + lifecycle.attachSessionEnd(window) window.on("closed", () => { remoteOrigins.delete(window.id); insecureOrigins.delete(window.webContents.id) }) try { await window.loadURL(target.toString()) } catch (error) { - const message = error instanceof Error ? error.message : String(error) - await window.loadURL(`data:text/html;charset=utf-8,${encodeURIComponent(`

${escapeHtml(payload.name)}

${escapeHtml(message)}

`)}`) + console.warn("[electron] failed to load remote window; showing loading screen", error) + remoteOrigins.delete(window.id) + insecureOrigins.delete(window.webContents.id) + const loading = loadingTarget() + await (loading.url ? window.loadURL(loading.url) : window.loadFile(loading.file!)) } } } -function setupNavigationGuards(window: BrowserWindow, navigation: ClientStateNavigationController | undefined, allowedOrigins: (window: BrowserWindow) => string[]) { +function setupNavigationGuards( + window: BrowserWindow, + navigation: ClientStateNavigationController | undefined, + allowedOrigins: (window: BrowserWindow) => string[], + loadingUrl: () => string, +) { const external = (url: string) => shell.openExternal(url).catch((error) => console.error("[cli] failed to open external URL", url, error)) - const shouldOpenExternally = (url: string) => { - try { const parsed = new URL(url); return !["http:", "https:", "file:"].includes(parsed.protocol) || (parsed.protocol !== "file:" && !allowedOrigins(window).includes(parsed.origin)) } catch { return false } - } - window.webContents.setWindowOpenHandler(({ url }) => shouldOpenExternally(url) ? (external(url), { action: "deny" }) : { action: "allow" }) + const decide = (url: string) => decideNavigation(url, allowedOrigins(window), loadingUrl()) + window.webContents.setWindowOpenHandler(({ url }) => { + if (decide(url) === "external") void external(url) + return { action: "deny" } + }) window.webContents.on("will-navigate", (event, url) => { - if (shouldOpenExternally(url)) { event.preventDefault(); void external(url) } + const decision = decide(url) + if (decision !== "allow") { event.preventDefault(); if (decision === "external") void external(url) } else if (navigation) { event.preventDefault(); void navigation.navigate((target) => target.loadURL(url)) } }) - window.webContents.on("will-redirect", (event, url) => { if (shouldOpenExternally(url)) { event.preventDefault(); void external(url) } }) + window.webContents.on("will-redirect", (event, url) => { + const decision = decide(url) + if (decision !== "allow") { event.preventDefault(); if (decision === "external") void external(url) } + }) } function isIgnorableNavigationError(error: unknown): boolean { @@ -321,8 +348,6 @@ function isIgnorableNavigationError(error: unknown): boolean { return text.includes("ERR_ABORTED") || text.includes("ERR_FAILED") } -function escapeHtml(value: string): string { return value.replace(/[&<>"]/g, (char) => ({ "&": "&", "<": "<", ">": ">", '"': """ })[char]!) } - async function exchangeBootstrapToken(baseUrl: string, token: string, cli: CliProcessManager): Promise { const target = new URL("/api/auth/token", baseUrl) const body = JSON.stringify({ token }) diff --git a/packages/electron-app/electron/main/multiwindow-lifecycle.test.ts b/packages/electron-app/electron/main/multiwindow-lifecycle.test.ts index e871ad0d..50b7aa56 100644 --- a/packages/electron-app/electron/main/multiwindow-lifecycle.test.ts +++ b/packages/electron-app/electron/main/multiwindow-lifecycle.test.ts @@ -93,22 +93,49 @@ test("final close retains its record and shutdown stops/releases once", async () assert.equal(calls.filter((call) => call === "release").length, 1) }) -test("Windows session end exits even when CLI shutdown rejects after the query was vetoed", async () => { +test("Windows query preflush leaves the app alive until session end is confirmed", async () => { const calls: string[] = [] const first = windowRecord("one", calls) const lifecycle = new MultiwindowLifecycle({ app: { on: () => {}, quit: () => {}, exit: () => calls.push("exit") } as never, - clientStateManager: { isPrimary: true, flush: async () => {}, drainAndReleasePrimary: async () => calls.push("release") } as never, - cliManager: { shutdown: async () => { throw new Error("CLI failed") } } as never, + clientStateManager: { isPrimary: true, flush: async () => calls.push("aggregate"), drainAndReleasePrimary: async () => calls.push("release") } as never, + cliManager: { shutdown: async () => calls.push("stop") } as never, getLocalWindows: () => [first], getAllWindows: () => [first.window], removeWindowState: async () => {}, getAllowedRendererOrigins: () => ["http://localhost"], isTrustedRendererOrigin: () => true, isWindows: true, sessionEndCleanupTimeoutMs: 20, }) lifecycle.attach(first) let vetoed = false first.events.get("query-session-end")?.({ preventDefault: () => { vetoed = true } }) - await tick(); await tick() - assert.equal(vetoed, true) - assert.deepEqual(calls.filter((call) => call === "exit"), ["exit"]) + await (lifecycle as any).sessionEndPreparation + assert.equal(vetoed, false) + assert.deepEqual(calls, ["renderer:one", "native:one"]) + + first.events.get("session-end")?.() + await (lifecycle as any).sessionEnd + assert.deepEqual(calls, ["renderer:one", "native:one", "aggregate", "stop", "release", "exit"]) +}) + +test("remote windows receive session-end cleanup without local close semantics", async () => { + const calls: string[] = [] + const events = new Map() + const remote = { on: (name: string, handler: Function) => events.set(name, handler), isDestroyed: () => false } + const lifecycle = new MultiwindowLifecycle({ + app: { on: () => {}, exit: () => calls.push("exit") } as never, + clientStateManager: { isPrimary: true, flush: async () => calls.push("aggregate"), drainAndReleasePrimary: async () => calls.push("release") } as never, + cliManager: { shutdown: async () => calls.push("stop") } as never, + getLocalWindows: () => [], getAllWindows: () => [remote as never], removeWindowState: async () => calls.push("remove"), + getAllowedRendererOrigins: () => [], isTrustedRendererOrigin: () => false, isWindows: true, + }) + + lifecycle.attachSessionEnd(remote as never) + let vetoed = false + events.get("query-session-end")?.({ preventDefault: () => { vetoed = true } }) + await (lifecycle as any).sessionEndPreparation + assert.equal(vetoed, false) + assert.deepEqual(calls, []) + events.get("session-end")?.() + await (lifecycle as any).sessionEnd + assert.deepEqual(calls, ["aggregate", "stop", "release", "exit"]) }) test("normal quit reports a failed CLI shutdown without allowing exit", async () => { diff --git a/packages/electron-app/electron/main/multiwindow-lifecycle.ts b/packages/electron-app/electron/main/multiwindow-lifecycle.ts index 976def93..9ed13661 100644 --- a/packages/electron-app/electron/main/multiwindow-lifecycle.ts +++ b/packages/electron-app/electron/main/multiwindow-lifecycle.ts @@ -27,8 +27,12 @@ interface Dependencies { export class MultiwindowLifecycle { private shutdown: Promise | null = null + private sessionEnd: Promise | null = null + private sessionEndPreparation: Promise | null = null + private sessionEndPreparationPending = false private release: Promise | null = null private exitAllowed = false + private readonly sessionEndWindows = new WeakSet() constructor(private readonly dependencies: Dependencies) {} @@ -57,14 +61,14 @@ export class MultiwindowLifecycle { }) }) - if (this.dependencies.isWindows ?? process.platform === "win32") { - record.window.on("query-session-end", (event) => { - if (this.exitAllowed) return - event.preventDefault() - this.startSessionEnd() - }) - record.window.on("session-end", () => this.startSessionEnd()) - } + this.attachSessionEnd(record.window) + } + + attachSessionEnd(window: BrowserWindow): void { + if (!(this.dependencies.isWindows ?? process.platform === "win32") || this.sessionEndWindows.has(window)) return + this.sessionEndWindows.add(window) + window.on("query-session-end", () => this.prepareSessionEnd()) + window.on("session-end", () => this.startSessionEnd()) } registerAppEvents(): void { @@ -77,10 +81,10 @@ export class MultiwindowLifecycle { this.dependencies.app.on("window-all-closed", () => this.dependencies.app.quit()) } - private startShutdown(): Promise { + private startShutdown(preparedFlush?: Promise): Promise { if (this.shutdown) return this.shutdown this.shutdown = (async () => { - await Promise.all(this.dependencies.getLocalWindows().map((record) => this.flushWindow(record))) + await (preparedFlush ?? this.flushLocalWindows()) await this.run("aggregate state flush", () => this.dependencies.clientStateManager.flush()) await this.dependencies.cliManager.shutdown() await this.releasePrimary() @@ -88,11 +92,26 @@ export class MultiwindowLifecycle { return this.shutdown } + private flushLocalWindows(): Promise { + return Promise.all(this.dependencies.getLocalWindows().map((record) => this.flushWindow(record))).then(() => undefined) + } + + private prepareSessionEnd(): void { + if (this.exitAllowed || this.sessionEnd || this.shutdown || this.sessionEndPreparationPending) return + this.sessionEndPreparationPending = true + const preparation = this.flushLocalWindows() + this.sessionEndPreparation = preparation + void preparation.finally(() => { + if (this.sessionEndPreparation === preparation) this.sessionEndPreparationPending = false + }) + } + private startSessionEnd(): void { - if (this.exitAllowed) return + if (this.exitAllowed || this.sessionEnd) return const timeoutMs = this.dependencies.sessionEndCleanupTimeoutMs ?? 5_000 - void Promise.race([ - this.startShutdown(), + const cleanup = this.shutdown ?? this.startShutdown(this.sessionEndPreparation ?? this.flushLocalWindows()) + this.sessionEnd = Promise.race([ + cleanup, new Promise((resolve) => setTimeout(resolve, timeoutMs)), ]).catch((error) => { console.warn("[client-state] OS session-end shutdown failed; exiting at the fail-open boundary", error) diff --git a/packages/electron-app/electron/main/navigation-security.test.ts b/packages/electron-app/electron/main/navigation-security.test.ts new file mode 100644 index 00000000..f74217f9 --- /dev/null +++ b/packages/electron-app/electron/main/navigation-security.test.ts @@ -0,0 +1,22 @@ +import assert from "node:assert/strict" +import test from "node:test" +import { decideNavigation, requireHttpUrl } from "./navigation-security" + +test("remote window URLs require HTTP or HTTPS", () => { + assert.equal(requireHttpUrl("http://localhost:3000/app", "baseUrl").protocol, "http:") + assert.equal(requireHttpUrl("https://example.com/app", "entryUrl").protocol, "https:") + for (const url of ["file:///tmp/index.html", "data:text/html,hello", "javascript:alert(1)"]) { + assert.throws(() => requireHttpUrl(url, "baseUrl"), /must use HTTP or HTTPS/) + } +}) + +test("navigation allows registered origins and only the exact loading file", () => { + const loading = "file:///opt/codenomad/loading.html" + const origins = ["https://renderer.example"] + assert.equal(decideNavigation(loading, origins, loading), "allow") + assert.equal(decideNavigation("file:///opt/codenomad/index.html", origins, loading), "deny") + assert.equal(decideNavigation("https://renderer.example/workspace", origins, loading), "allow") + assert.equal(decideNavigation("https://outside.example/", origins, loading), "external") + assert.equal(decideNavigation("not a URL", origins, loading), "deny") + assert.equal(decideNavigation("http://localhost:5173/loading.html", [], "http://localhost:5173/loading.html"), "allow") +}) diff --git a/packages/electron-app/electron/main/navigation-security.ts b/packages/electron-app/electron/main/navigation-security.ts new file mode 100644 index 00000000..df567f73 --- /dev/null +++ b/packages/electron-app/electron/main/navigation-security.ts @@ -0,0 +1,25 @@ +export type NavigationDecision = "allow" | "external" | "deny" + +export function requireHttpUrl(value: string, name: string): URL { + const url = new URL(value) + if (url.protocol !== "http:" && url.protocol !== "https:") throw new Error(`${name} must use HTTP or HTTPS`) + return url +} + +export function decideNavigation( + value: string, + allowedOrigins: readonly string[], + loadingUrl: string, +): NavigationDecision { + let url: URL + try { url = new URL(value) } catch { return "deny" } + + try { + if (url.toString() === new URL(loadingUrl).toString()) return "allow" + } catch {} + + if (url.protocol === "http:" || url.protocol === "https:") { + return allowedOrigins.includes(url.origin) ? "allow" : "external" + } + return url.protocol === "file:" || url.protocol === "data:" || url.protocol === "javascript:" ? "deny" : "external" +} diff --git a/packages/electron-app/electron/main/process-exit.test.ts b/packages/electron-app/electron/main/process-exit.test.ts new file mode 100644 index 00000000..3adffe15 --- /dev/null +++ b/packages/electron-app/electron/main/process-exit.test.ts @@ -0,0 +1,23 @@ +import assert from "node:assert/strict" +import test from "node:test" +import { resolveManagedProcessExit, shouldReportManagedProcessError } from "./process-exit" + +test("requested and invalidated process exits do not become failures", () => { + assert.deepEqual(resolveManagedProcessExit(undefined, 0, null, true, true), { state: "stopped" }) + assert.equal(resolveManagedProcessExit(undefined, 1, null, false, false), null) +}) + +test("unexpected exits report their code or signal", () => { + assert.deepEqual(resolveManagedProcessExit(undefined, 0, null, false, true), { + state: "error", + error: "CLI exited unexpectedly (code 0)", + }) + assert.match(resolveManagedProcessExit(undefined, null, "SIGTERM", false, true)?.error ?? "", /signal SIGTERM/) + assert.equal(resolveManagedProcessExit("startup failed", 1, null, false, true)?.error, "startup failed") +}) + +test("child errors are reported only for the current process outside a requested stop", () => { + assert.equal(shouldReportManagedProcessError(false, true), true) + assert.equal(shouldReportManagedProcessError(true, true), false) + assert.equal(shouldReportManagedProcessError(false, false), false) +}) diff --git a/packages/electron-app/electron/main/process-exit.ts b/packages/electron-app/electron/main/process-exit.ts new file mode 100644 index 00000000..6b3f4db1 --- /dev/null +++ b/packages/electron-app/electron/main/process-exit.ts @@ -0,0 +1,21 @@ +export interface ManagedProcessExit { + state: "error" | "stopped" + error?: string +} + +export function shouldReportManagedProcessError(requestedStop: boolean, currentGeneration: boolean): boolean { + return currentGeneration && !requestedStop +} + +export function resolveManagedProcessExit( + currentError: string | undefined, + code: number | null, + signal: NodeJS.Signals | null, + requestedStop: boolean, + currentGeneration: boolean, +): ManagedProcessExit | null { + if (!currentGeneration) return null + if (requestedStop) return { state: "stopped" } + const details = [code === null ? null : `code ${code}`, signal ? `signal ${signal}` : null].filter(Boolean).join(", ") + return { state: "error", error: currentError ?? `CLI exited unexpectedly (${details || "unknown status"})` } +} diff --git a/packages/electron-app/electron/main/process-manager.ts b/packages/electron-app/electron/main/process-manager.ts index 00bd7035..0d79c161 100644 --- a/packages/electron-app/electron/main/process-manager.ts +++ b/packages/electron-app/electron/main/process-manager.ts @@ -18,6 +18,7 @@ import { stopManagedChild, } from "./process-stop" import { SerializedLifecycle } from "./serialized-lifecycle" +import { resolveManagedProcessExit, shouldReportManagedProcessError } from "./process-exit" import { buildUserShellCommand, getUserShellEnv, supportsUserShell } from "./user-shell" const nodeRequire = createRequire(import.meta.url) @@ -143,6 +144,7 @@ export class CliProcessManager extends EventEmitter { private bootstrapToken: string | null = null private authCookieName = `${SESSION_COOKIE_NAME_PREFIX}_${process.pid}_${Date.now()}` private requestedStop = false + private cancelPendingStart?: (error: Error) => void private shutdownStatus: "complete" | "incomplete" | null = null private lifecycle = new SerializedLifecycle() @@ -163,7 +165,10 @@ export class CliProcessManager extends EventEmitter { } shutdown(): Promise { - return this.lifecycle.stop(() => this.stopNow()) + return this.lifecycle.stop(() => this.stopNow(), () => { + this.requestedStop = true + this.cancelPendingStart?.(new Error("CLI startup interrupted by shutdown")) + }) } private async startNow(options: StartOptions): Promise { @@ -171,6 +176,7 @@ export class CliProcessManager extends EventEmitter { if (this.child) { await this.stopNow() if (this.child) throw new Error("CLI process did not exit before restart") + if (this.lifecycle.stopped) throw new Error("CLI startup interrupted by shutdown") } this.stdoutBuffer = "" @@ -185,7 +191,8 @@ export class CliProcessManager extends EventEmitter { const listeningMode = this.resolveListeningMode() const host = resolveHostForMode(listeningMode) const args = this.buildCliArgs(options, host) - const cliEntry = await this.resolveCliEntry(options) + const cliEntry = await this.awaitStartupStep(this.resolveCliEntry(options)) + if (this.lifecycle.stopped) throw new Error("CLI startup interrupted by shutdown") console.info( `[cli] launching CodeNomad CLI (${options.dev ? "dev" : "prod"}) using ${cliEntry.runner} at ${cliEntry.entry} (host=${host})`, @@ -229,17 +236,19 @@ export class CliProcessManager extends EventEmitter { }) child.on("error", (error) => { + if (!shouldReportManagedProcessError(this.requestedStop, this.child === child)) return console.error("[cli] failed to start CLI:", error) this.updateStatus({ state: "error", error: error.message }) this.emit("error", error) }) child.on("exit", (code, signal) => { - if (this.child !== child) return - const failed = this.status.state !== "ready" - const error = failed ? this.status.error ?? `CLI exited with code ${code ?? 0}${signal ? ` (${signal})` : ""}` : undefined + const exit = resolveManagedProcessExit(this.status.error, code, signal, this.requestedStop, this.child === child) + if (!exit) return + const failed = exit.state === "error" + const error = exit.error console.info(`[cli] exit (code=${code}, signal=${signal || ""})${error ? ` error=${error}` : ""}`) - this.updateStatus({ state: failed ? "error" : "stopped", error }) + this.updateStatus({ state: exit.state, error }) if (failed && error) { this.emit("error", new Error(error)) } @@ -251,18 +260,22 @@ export class CliProcessManager extends EventEmitter { return new Promise((resolve, reject) => { const timeout = setTimeout(() => { this.handleTimeout() - reject(new Error("CLI startup timeout")) + finish(reject, new Error("CLI startup timeout")) }, 60000) - this.once("ready", (status) => { + const finish = (settle: (value: T) => void, value: T) => { clearTimeout(timeout) - resolve(status) - }) - - this.once("error", (error) => { - clearTimeout(timeout) - reject(error) - }) + this.off("ready", onReady) + this.off("error", onError) + if (this.cancelPendingStart === cancel) this.cancelPendingStart = undefined + settle(value) + } + const onReady = (status: CliStatus) => finish(resolve, status) + const onError = (error: Error) => finish(reject, error) + const cancel = (error: Error) => finish(reject, error) + this.cancelPendingStart = cancel + this.once("ready", onReady) + this.once("error", onError) }) } @@ -354,6 +367,22 @@ export class CliProcessManager extends EventEmitter { return this.authCookieName } + private awaitStartupStep(operation: Promise): Promise { + return new Promise((resolve, reject) => { + const finish = () => { + if (this.cancelPendingStart !== cancel) return false + this.cancelPendingStart = undefined + return true + } + const cancel = (error: Error) => { if (finish()) reject(error) } + this.cancelPendingStart = cancel + operation.then( + (value) => { if (finish()) resolve(value) }, + (error) => { if (finish()) reject(error) }, + ) + }) + } + private resolveListeningMode(): ListeningMode { return readListeningModeFromConfig() } diff --git a/packages/electron-app/electron/main/process-stop.test.ts b/packages/electron-app/electron/main/process-stop.test.ts index 898a79c3..500fc9eb 100644 --- a/packages/electron-app/electron/main/process-stop.test.ts +++ b/packages/electron-app/electron/main/process-stop.test.ts @@ -514,7 +514,7 @@ test("signal dispatch is not confirmation while the captured identity remains", assert.equal(await forceCapturedProcessTree(tree, () => "owned", undefined, kill), false) }) -test("incomplete shutdown status remains terminal", async () => { +test("process manager keeps incomplete shutdown terminal and interrupts a pending startup", async () => { const hooks = registerHooks({ resolve(specifier, context, nextResolve) { if (specifier === "electron") { @@ -536,6 +536,14 @@ test("incomplete shutdown status remains terminal", async () => { assert.equal((manager as any).shutdownStatus, "incomplete") assert.equal(enforcements, 1) + + const pending = new CliProcessManager() + ;(pending as any).resolveCliEntry = () => new Promise(() => {}) + const startup = pending.start({ dev: false }) + await new Promise((resolve) => setImmediate(resolve)) + const shutdown = pending.shutdown() + await assert.rejects(startup, /startup interrupted by shutdown/) + await shutdown } finally { hooks.deregister() } diff --git a/packages/electron-app/electron/main/remote-window-registry.test.ts b/packages/electron-app/electron/main/remote-window-registry.test.ts index bda77c30..e9c451bf 100644 --- a/packages/electron-app/electron/main/remote-window-registry.test.ts +++ b/packages/electron-app/electron/main/remote-window-registry.test.ts @@ -1,6 +1,7 @@ import assert from "node:assert/strict" import test from "node:test" -import { RemoteWindowRegistry } from "./remote-window-registry" +import type { BrowserWindow } from "electron" +import { navigateReusedRemoteWindow, RemoteWindowRegistry } from "./remote-window-registry" function window() { const events = new Map void>() @@ -16,7 +17,7 @@ function window() { focus: () => calls.push("focus"), close: () => calls.push("close"), on: (name: string, callback: () => void) => events.set(name, callback), - } as never, + } as unknown as BrowserWindow, } } @@ -45,3 +46,35 @@ test("proxy replacement and close clean exactly their corresponding sessions", ( second.events.get("closed")?.() assert.deepEqual(cleaned, ["proxy-one", "proxy-two"]) }) + +test("reused remote navigation trusts old and next redirect origins until success", async () => { + const remote = window() + const trusted = new Map([[1, new Set(["https://old.example"])]]) + const insecure = new Map([[2, new Set(["https://old.example"])]]) + Object.assign(remote.value, { id: 1, webContents: { id: 2 } }) + remote.value.loadURL = async () => { + assert.deepEqual([...trusted.get(1)!], ["https://old.example", "https://new.example", "https://redirect.example"]) + } + + const next = new Set(["https://new.example", "https://redirect.example"]) + await navigateReusedRemoteWindow(remote.value, new URL("https://new.example/app"), next, trusted, insecure, false) + assert.deepEqual([...trusted.get(1)!], [...next]) + assert.equal(insecure.has(2), false) +}) + +test("failed reused remote navigation restores trusted and insecure origins", async () => { + const remote = window() + const trusted = new Map([[1, new Set(["https://old.example"])]]) + const insecure = new Map([[2, new Set(["https://old.example"])]]) + Object.assign(remote.value, { id: 1, webContents: { id: 2 } }) + remote.value.loadURL = async () => { + assert.deepEqual([...trusted.get(1)!], ["https://old.example", "https://new.example"]) + assert.deepEqual([...insecure.get(2)!], ["https://old.example", "https://new.example"]) + throw new Error("failed") + } + + const next = new Set(["https://new.example"]) + await assert.rejects(navigateReusedRemoteWindow(remote.value, new URL("https://new.example/app"), next, trusted, insecure, true), /failed/) + assert.deepEqual([...trusted.get(1)!], ["https://old.example"]) + assert.deepEqual([...insecure.get(2)!], ["https://old.example"]) +}) diff --git a/packages/electron-app/electron/main/remote-window-registry.ts b/packages/electron-app/electron/main/remote-window-registry.ts index 67026cc6..41c3f4b5 100644 --- a/packages/electron-app/electron/main/remote-window-registry.ts +++ b/packages/electron-app/electron/main/remote-window-registry.ts @@ -35,3 +35,30 @@ export class RemoteWindowRegistry { }) } } + +export async function navigateReusedRemoteWindow( + window: BrowserWindow, + target: URL, + nextOrigins: ReadonlySet, + trustedOrigins: Map>, + insecureOrigins: Map>, + skipTlsVerify: boolean, +): Promise { + const committedOrigins = new Set(nextOrigins) + const previousTrusted = trustedOrigins.get(window.id) + const previousInsecure = insecureOrigins.get(window.webContents.id) + trustedOrigins.set(window.id, new Set([...(previousTrusted ?? []), ...committedOrigins])) + if (skipTlsVerify) insecureOrigins.set(window.webContents.id, new Set([...(previousInsecure ?? []), ...committedOrigins])) + + try { await window.loadURL(target.toString()) } catch (error) { + if (previousTrusted) trustedOrigins.set(window.id, previousTrusted) + else trustedOrigins.delete(window.id) + if (previousInsecure) insecureOrigins.set(window.webContents.id, previousInsecure) + else insecureOrigins.delete(window.webContents.id) + throw error + } + + trustedOrigins.set(window.id, committedOrigins) + if (skipTlsVerify) insecureOrigins.set(window.webContents.id, committedOrigins) + else insecureOrigins.delete(window.webContents.id) +} diff --git a/packages/electron-app/electron/main/serialized-lifecycle.test.ts b/packages/electron-app/electron/main/serialized-lifecycle.test.ts index 24e7422f..aacd9744 100644 --- a/packages/electron-app/electron/main/serialized-lifecycle.test.ts +++ b/packages/electron-app/electron/main/serialized-lifecycle.test.ts @@ -36,3 +36,16 @@ test("failed shutdown reopens the lifecycle before queued retries run", async () assert.equal(await retry, "restarted") assert.equal(lifecycle.stopped, false) }) + +test("shutdown can interrupt pending work before entering the serialized stop", async () => { + const lifecycle = new SerializedLifecycle() + let interrupt!: () => void + const startup = lifecycle.enqueue(() => new Promise((_resolve, reject) => { + interrupt = () => reject(new Error("startup interrupted")) + })) + await new Promise((resolve) => setImmediate(resolve)) + const shutdown = lifecycle.stop(async () => {}, () => interrupt()) + + await assert.rejects(startup, /startup interrupted/) + await shutdown +}) diff --git a/packages/electron-app/electron/main/serialized-lifecycle.ts b/packages/electron-app/electron/main/serialized-lifecycle.ts index 1f72f5a7..3f1ec89a 100644 --- a/packages/electron-app/electron/main/serialized-lifecycle.ts +++ b/packages/electron-app/electron/main/serialized-lifecycle.ts @@ -8,8 +8,9 @@ export class SerializedLifecycle { return queued } - stop(operation: () => Promise): Promise { + stop(operation: () => Promise, interrupt?: () => void): Promise { this.stopped = true + interrupt?.() return this.enqueue(async () => { try { return await operation() diff --git a/packages/electron-app/electron/preload/index.cjs b/packages/electron-app/electron/preload/index.cjs index 0f0e318a..c21c5fac 100644 --- a/packages/electron-app/electron/preload/index.cjs +++ b/packages/electron-app/electron/preload/index.cjs @@ -22,12 +22,14 @@ const windowId = resolveWindowId() const localElectronAPI = { onCliStatus: (callback) => { - ipcRenderer.on("cli:status", (_, data) => callback(data)) - return () => ipcRenderer.removeAllListeners("cli:status") + const handler = (_, data) => callback(data) + ipcRenderer.on("cli:status", handler) + return () => ipcRenderer.removeListener("cli:status", handler) }, onCliError: (callback) => { - ipcRenderer.on("cli:error", (_, data) => callback(data)) - return () => ipcRenderer.removeAllListeners("cli:error") + const handler = (_, data) => callback(data) + ipcRenderer.on("cli:error", handler) + return () => ipcRenderer.removeListener("cli:error", handler) }, getCliStatus: () => ipcRenderer.invoke("cli:getStatus"), restartCli: () => ipcRenderer.invoke("cli:restart"), diff --git a/packages/electron-app/electron/preload/index.test.ts b/packages/electron-app/electron/preload/index.test.ts new file mode 100644 index 00000000..2d0e5698 --- /dev/null +++ b/packages/electron-app/electron/preload/index.test.ts @@ -0,0 +1,35 @@ +import assert from "node:assert/strict" +import { readFileSync } from "node:fs" +import test from "node:test" +import vm from "node:vm" + +test("CLI event disposers remove only their own wrapper listeners", () => { + const listeners = new Map>() + let api: Record | undefined + const ipcRenderer = { + on(channel: string, listener: Function) { + const channelListeners = listeners.get(channel) ?? new Set() + channelListeners.add(listener) + listeners.set(channel, channelListeners) + }, + removeListener(channel: string, listener: Function) { listeners.get(channel)?.delete(listener) }, + invoke() {}, + } + vm.runInNewContext(readFileSync(new URL("./index.cjs", import.meta.url), "utf8"), { + require: () => ({ + contextBridge: { exposeInMainWorld(name: string, value: Record) { if (name === "electronAPI") api = value } }, + ipcRenderer, + webUtils: {}, + }), + process: { argv: [] }, + }) + + for (const [subscribe, channel] of [["onCliStatus", "cli:status"], ["onCliError", "cli:error"]] as const) { + const calls: string[] = [] + const disposeFirst = api![subscribe]((value: string) => calls.push(`first:${value}`)) + api![subscribe]((value: string) => calls.push(`second:${value}`)) + disposeFirst() + for (const listener of listeners.get(channel) ?? []) listener({}, "event") + assert.deepEqual(calls, ["second:event"]) + } +}) diff --git a/packages/electron-app/package.json b/packages/electron-app/package.json index 7a9fcd69..ffb93e67 100644 --- a/packages/electron-app/package.json +++ b/packages/electron-app/package.json @@ -24,7 +24,7 @@ "prebuild": "npm run prepare:resources", "build": "electron-vite build", "typecheck": "tsc --noEmit -p tsconfig.json", - "test:native": "node --import tsx --test electron/main/client-state-cross-host.test.ts electron/main/client-state-process.test.ts electron/main/client-state.test.ts electron/main/client-state-ipc.test.ts electron/main/client-state-navigation.test.ts electron/main/client-state-lifecycle.test.ts electron/main/local-window-registry.test.ts electron/main/menu-target.test.ts electron/main/multiwindow-lifecycle.test.ts electron/main/process-stop.test.ts electron/main/remote-window-registry.test.ts electron/main/renderer-client-state-flush.test.ts electron/main/renderer-origin.test.ts electron/main/serialized-lifecycle.test.ts electron/main/startup.test.ts electron/main/window-state.test.ts electron/main/workspace-open.test.ts", + "test:native": "node --import tsx --test electron/main/client-state-cross-host.test.ts electron/main/client-state-process.test.ts electron/main/client-state.test.ts electron/main/client-state-ipc.test.ts electron/main/client-state-navigation.test.ts electron/main/client-state-lifecycle.test.ts electron/main/local-window-registry.test.ts electron/main/menu-target.test.ts electron/main/multiwindow-lifecycle.test.ts electron/main/navigation-security.test.ts electron/main/process-exit.test.ts electron/main/process-stop.test.ts electron/main/remote-window-registry.test.ts electron/main/renderer-client-state-flush.test.ts electron/main/renderer-origin.test.ts electron/main/serialized-lifecycle.test.ts electron/main/startup.test.ts electron/main/window-state.test.ts electron/main/workspace-open.test.ts electron/preload/index.test.ts", "preview": "electron-vite preview", "build:binaries": "node scripts/build.js", "build:mac": "node scripts/build.js mac", diff --git a/packages/server/src/server/__tests__/instance-proxy.test.ts b/packages/server/src/server/__tests__/instance-proxy.test.ts index b6b95d40..e082e05b 100644 --- a/packages/server/src/server/__tests__/instance-proxy.test.ts +++ b/packages/server/src/server/__tests__/instance-proxy.test.ts @@ -25,6 +25,7 @@ async function harness( pathMappings: Record = {}, ptyDirectories: Record = {}, shellDirectories: Record = {}, + directoryMappings: Record = {}, ) { const upstream = Fastify() apps.push(upstream) @@ -38,7 +39,7 @@ async function harness( const address = upstream.server.address() assert.ok(address && typeof address === "object") - const owned = new Set([workspacePath, serviceDirectory, "/repo", "/repo/worktree"]) + const owned = new Set([workspacePath, serviceDirectory, "/repo", "/repo/worktree", ...Object.keys(directoryMappings)]) const sessionGets: string[] = [] const pathOwnershipChecks: string[] = [] const servicePathCalls: string[] = [] @@ -90,7 +91,9 @@ async function harness( getSharedServiceEndpoint: async () => ({ url: `http://127.0.0.1:${address.port}` }), getInstanceAuthorizationHeader: () => "Basic internal-secret", getServiceDirectory: () => serviceDirectory, - getServiceDirectoryForPath: async (_id, directory) => directory === workspacePath ? serviceDirectory : owned.has(directory) ? directory : undefined, + getServiceDirectoryForPath: async (_id, directory) => directory === workspacePath + ? serviceDirectory + : owned.has(directory) ? directoryMappings[directory] ?? directory : undefined, getServicePathForPath: async (_id, candidate) => { assert.ok(pathOwnershipChecks.includes(candidate), "prompt path must be ownership-checked before translation") servicePathCalls.push(candidate) @@ -296,6 +299,74 @@ describe("instance proxy location enforcement", () => { assert.deepEqual(sessionGets, ["global"]) }) + it("forwards a validated global Form root location instead of browser routing headers", async () => { + const { app } = await harness("/repo/worktree", {}, {}, "/repo", "/srv/repo") + const response = await app.inject({ + method: "POST", + url: "/workspaces/workspace/instance/api/session/global/form/form-1/reply", + headers: { + "x-opencode-directory": encodeURIComponent("/repo"), + "x-opencode-workspace": "untrusted-workspace", + }, + payload: { answers: {} }, + }) + + assert.equal(response.statusCode, 200) + const upstream = JSON.parse(response.body) + assert.equal(upstream.headers["x-opencode-directory"], encodeURIComponent("/srv/repo")) + assert.equal(upstream.headers["x-opencode-workspace"], undefined) + }) + + it("translates and forwards a validated global Form worktree location", async () => { + const { app } = await harness( + "/repo/worktree", {}, {}, "/repo", "/srv/repo", {}, {}, {}, + { "/repo/worktree": "/srv/worktree" }, + ) + const response = await app.inject({ + method: "POST", + url: "/workspaces/workspace/instance/api/session/global/form/form-1/cancel", + headers: { "x-opencode-directory": encodeURIComponent("/repo/worktree") }, + payload: {}, + }) + + assert.equal(response.statusCode, 200) + assert.equal(JSON.parse(response.body).headers["x-opencode-directory"], encodeURIComponent("/srv/worktree")) + }) + + it("rejects a foreign global Form location before proxying", async () => { + const { app, requestCount } = await harness() + const response = await app.inject({ + method: "POST", + url: "/workspaces/workspace/instance/api/session/global/form/form-1/reply", + headers: { "x-opencode-directory": encodeURIComponent("/other") }, + payload: { answers: {} }, + }) + + assert.equal(response.statusCode, 403) + assert.equal((await app.inject({ + method: "POST", + url: "/workspaces/workspace/instance/api/session/global/form/form-1/reply", + headers: { "x-opencode-directory": "%ZZ" }, + payload: { answers: {} }, + })).statusCode, 400) + assert.equal(requestCount(), 0) + }) + + it("decodes, translates, and re-encodes Unicode global Form locations", async () => { + const directory = "/工作/100% ready" + const serviceDirectory = "/服务/工作 100%" + const { app } = await harness(directory, {}, {}, directory, serviceDirectory) + const response = await app.inject({ + method: "POST", + url: "/workspaces/workspace/instance/api/session/global/form/form-1/reply", + headers: { "x-opencode-directory": encodeURIComponent(directory) }, + payload: { answers: {} }, + }) + + assert.equal(response.statusCode, 200) + assert.equal(JSON.parse(response.body).headers["x-opencode-directory"], encodeURIComponent(serviceDirectory)) + }) + it("rejects deletion through a double-encoded alias of a foreign session", async () => { const { app, sessionGets, requestCount } = await harness("/repo/worktree", {}, { "foreign%25session": "/other", diff --git a/packages/server/src/server/http-server.ts b/packages/server/src/server/http-server.ts index 6248037e..784226d1 100644 --- a/packages/server/src/server/http-server.ts +++ b/packages/server/src/server/http-server.ts @@ -620,6 +620,26 @@ async function proxyWorkspaceRequest(args: { return } const serviceDirectory = workspaceManager.getServiceDirectory?.(workspaceId) ?? workspace.path + let globalFormDirectory: string | undefined + if (isGlobalFormAction(pathname, request.method)) { + const header = request.headers["x-opencode-directory"] + if (Array.isArray(header)) { + reply.code(400).send({ error: "Invalid Form location" }) + return + } + if (header !== undefined) { + try { + globalFormDirectory = decodeURIComponent(header) + } catch { + reply.code(400).send({ error: "Invalid Form location" }) + return + } + if (!globalFormDirectory.trim()) { + reply.code(400).send({ error: "Invalid Form location" }) + return + } + } + } const imported = prepareSessionImport( pathname, request.method, @@ -627,6 +647,7 @@ async function proxyWorkspaceRequest(args: { serviceDirectory, ) const requestLocations = readRequestDirectories(targetUrl, imported.body) + if (globalFormDirectory) requestLocations.directories.push(globalFormDirectory) requestLocations.directories.push(...imported.directories) requestLocations.invalid ||= imported.invalid readNativeCwd(targetUrl, imported.body, requestLocations) @@ -741,6 +762,9 @@ async function proxyWorkspaceRequest(args: { ...(body !== request.body ? { body } : {}), rewriteRequestHeaders: (_originalRequest, headers) => { const outgoingHeaders = sanitizeInstanceProxyRequestHeaders(headers, instanceAuthHeader) + if (globalFormDirectory) { + outgoingHeaders["x-opencode-directory"] = encodeURIComponent(translatedDirectories.get(globalFormDirectory)!) + } if (logger.isLevelEnabled("trace")) { logger.trace( diff --git a/packages/server/src/server/routes/events.test.ts b/packages/server/src/server/routes/events.test.ts index ce3ee0ec..0fd44f65 100644 --- a/packages/server/src/server/routes/events.test.ts +++ b/packages/server/src/server/routes/events.test.ts @@ -1,5 +1,6 @@ import assert from "node:assert/strict" import { describe, it } from "node:test" +import cors from "@fastify/cors" import Fastify from "fastify" import type { EventBus } from "../../events/bus" import type { Logger } from "../../logger" @@ -7,8 +8,9 @@ import { registerEventRoutes } from "./events" const logger = { debug() {}, trace() {}, isLevelEnabled() { return false } } as unknown as Logger -function harness(options: { limit?: number; timeout?: number } = {}) { +function harness(options: { limit?: number; timeout?: number; corsOrigin?: false | string } = {}) { const app = Fastify() + if (options.corsOrigin !== undefined) app.register(cors, { origin: options.corsOrigin, credentials: true }) let listener: ((event: any) => void) | undefined let closeClient: (() => void) | undefined let raw: NodeJS.EventEmitter | undefined @@ -67,6 +69,43 @@ async function waitFor(check: () => boolean): Promise { } describe("SSE backpressure", () => { + it("does not override the central CORS policy", async () => { + const test = harness({ corsOrigin: false }) + try { + const response = test.app.inject({ + method: "GET", + url: "/api/events?clientId=client&connectionId=connection", + headers: { origin: "https://untrusted.example" }, + }) + await waitFor(test.ready) + test.close() + const result = await response + assert.equal(result.headers["access-control-allow-origin"], undefined) + assert.equal(result.headers["access-control-allow-credentials"], undefined) + } finally { + await test.app.close() + } + }) + + it("preserves headers from an allowed central CORS policy", async () => { + const origin = "https://trusted.example" + const test = harness({ corsOrigin: origin }) + try { + const response = test.app.inject({ + method: "GET", + url: "/api/events?clientId=client&connectionId=connection", + headers: { origin }, + }) + await waitFor(test.ready) + test.close() + const result = await response + assert.equal(result.headers["access-control-allow-origin"], origin) + assert.equal(result.headers["access-control-allow-credentials"], "true") + } finally { + await test.app.close() + } + }) + it("queues after write(false), flushes on drain, and remains connected", async () => { const test = harness() try { diff --git a/packages/server/src/server/routes/events.ts b/packages/server/src/server/routes/events.ts index 1309c352..1f0ef63b 100644 --- a/packages/server/src/server/routes/events.ts +++ b/packages/server/src/server/routes/events.ts @@ -31,9 +31,11 @@ export function registerEventRoutes(app: FastifyInstance, deps: RouteDeps) { const connection = ConnectionQuerySchema.parse(request.query ?? {}) deps.logger.debug({ clientId }, "SSE client connected") - const origin = request.headers.origin ?? "*" - reply.raw.setHeader("Access-Control-Allow-Origin", origin) - reply.raw.setHeader("Access-Control-Allow-Credentials", "true") + for (const [name, value] of Object.entries(reply.getHeaders())) { + if ((name === "vary" || name.startsWith("access-control-")) && value !== undefined) { + reply.raw.setHeader(name, value) + } + } reply.raw.setHeader("Content-Type", "text/event-stream") reply.raw.setHeader("Cache-Control", "no-cache") reply.raw.setHeader("Connection", "keep-alive") diff --git a/packages/server/src/workspaces/instance-events.test.ts b/packages/server/src/workspaces/instance-events.test.ts index 8ef612e1..1c81df8d 100644 --- a/packages/server/src/workspaces/instance-events.test.ts +++ b/packages/server/src/workspaces/instance-events.test.ts @@ -32,6 +32,10 @@ function waitFor(check: () => boolean): Promise { }) } +function serverConnected(): OpenCodeEvent { + return { type: "server.connected", data: {} } as OpenCodeEvent +} + function locationlessManager( events: OpenCodeEvent[], sessionLocations: Record, @@ -53,6 +57,7 @@ function locationlessManager( } }, }), subscribeToSharedService: async (signal?: AbortSignal) => (async function* () { + yield serverConnected() yield* events await new Promise((resolve) => signal?.addEventListener("abort", () => resolve(), { once: true })) })(), @@ -68,7 +73,7 @@ describe("InstanceEventBridge", () => { ownsDirectory: async () => true, subscribeToSharedService: async (signal?: AbortSignal) => (async function* () { await gate.promise - yield { type: "permission.asked", location: { directory: "/repo-a" }, data: { id: "p1" } } as OpenCodeEvent + yield serverConnected() await new Promise((resolve) => signal?.addEventListener("abort", () => resolve(), { once: true })) })(), } as unknown as WorkspaceManager @@ -87,6 +92,41 @@ describe("InstanceEventBridge", () => { } }) + it("rejects a stream whose first event is not server.connected and reconnects", async () => { + let subscriptions = 0 + const manager = { + list: () => [{ id: "a", path: "/repo-a" }], + ownsDirectory: async () => true, + subscribeToSharedService: async (signal?: AbortSignal) => { + subscriptions += 1 + return (async function* () { + if (subscriptions === 1) { + yield { type: "permission.asked", location: { directory: "/repo-a" }, data: { id: "p1" } } as OpenCodeEvent + return + } + yield serverConnected() + await new Promise((resolve) => signal?.addEventListener("abort", () => resolve(), { once: true })) + })() + }, + } as unknown as WorkspaceManager + const bus = new EventBus() + const statuses: Array<{ status: string; reason?: string }> = [] + const received: OpenCodeEvent[] = [] + bus.on("instance.eventStatus", (event) => statuses.push(event)) + bus.on("instance.event", (event) => received.push(event.event)) + const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) + try { + bus.publish({ type: "workspace.started", workspace: manager.list()[0] as any }) + await waitFor(() => statuses.some((event) => event.status === "error")) + assert.match(statuses.find((event) => event.status === "error")?.reason ?? "", /expected server\.connected/) + assert.deepEqual(received, []) + await waitFor(() => statuses.some((event) => event.status === "connected")) + assert.equal(subscriptions, 2) + } finally { + bridge.shutdown() + } + }) + it("clears routing caches before reconnecting", async () => { let subscriptions = 0 let ownershipChecks = 0 @@ -98,6 +138,7 @@ describe("InstanceEventBridge", () => { subscriptions += 1 const current = subscriptions return (async function* () { + yield serverConnected() yield event if (current > 1) await new Promise((resolve) => signal?.addEventListener("abort", () => resolve(), { once: true })) })() @@ -105,7 +146,9 @@ describe("InstanceEventBridge", () => { } as unknown as WorkspaceManager const bus = new EventBus() const received: unknown[] = [] - bus.on("instance.event", (value) => received.push(value)) + bus.on("instance.event", (value) => { + if (value.event.type !== "server.connected") received.push(value) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { bus.publish({ type: "workspace.started", workspace: manager.list()[0] as any }) @@ -117,6 +160,7 @@ describe("InstanceEventBridge", () => { }) it("routes root and owned worktree events to the logical workspace and caches ownership", async () => { const events = [ + { id: "0", created: 0, type: "server.connected", data: {} }, { id: "1", created: 1, type: "permission.asked", location: { directory: "/repo-a" }, data: { id: "p1" } }, { id: "2", @@ -133,7 +177,6 @@ describe("InstanceEventBridge", () => { }, }, { id: "3", created: 3, type: "permission.asked", location: { directory: "/other" }, data: { id: "p2" } }, - { id: "4", created: 4, type: "server.connected", data: {} }, { id: "5", created: 5, @@ -173,16 +216,16 @@ describe("InstanceEventBridge", () => { try { bus.publish({ type: "workspace.started", workspace: manager.list()[0] as any }) await waitFor(() => received.length === 6) - assert.equal(received[0].instanceId, "a") - assert.deepEqual(received[0].event.location, { directory: "/repo-a" }) - assert.deepEqual(received[0].event.data, { id: "p1" }) - assert.equal(received[0].event.properties, undefined) - assert.equal(received[1].event.data.sessionID, "session-1") - assert.equal(received[1].event.properties, undefined) - assert.deepEqual(received.slice(2, 4).map((event) => [event.instanceId, event.event.type]), [ + assert.deepEqual(received.slice(0, 2).map((event) => [event.instanceId, event.event.type]), [ ["a", "server.connected"], ["b", "server.connected"], ]) + assert.equal(received[2].instanceId, "a") + assert.deepEqual(received[2].event.location, { directory: "/repo-a" }) + assert.deepEqual(received[2].event.data, { id: "p1" }) + assert.equal(received[2].event.properties, undefined) + assert.equal(received[3].event.data.sessionID, "session-1") + assert.equal(received[3].event.properties, undefined) assert.equal(received[4].instanceId, "a") assert.deepEqual(received[4].event.data, { sessionID: "session-2", @@ -203,13 +246,16 @@ describe("InstanceEventBridge", () => { list: () => [{ id: "first", path: "/repo" }, { id: "second", path: "/repo" }], ownsDirectory: async (_workspaceId: string, directory: string) => directory === "/repo", subscribeToSharedService: async (signal?: AbortSignal) => (async function* () { + yield serverConnected() yield { id: "1", created: 1, type: "permission.asked", location: { directory: "/repo" }, data: { id: "p1" } } as OpenCodeEvent await new Promise((resolve) => signal?.addEventListener("abort", () => resolve(), { once: true })) })(), } as unknown as WorkspaceManager const bus = new EventBus() const received: string[] = [] - bus.on("instance.event", (event) => received.push(event.instanceId)) + bus.on("instance.event", (event) => { + if (event.event.type !== "server.connected") received.push(event.instanceId) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { @@ -231,7 +277,9 @@ describe("InstanceEventBridge", () => { const { manager, sessionGets } = locationlessManager(events, { known: "/repo-a" }) const bus = new EventBus() const received: any[] = [] - bus.on("instance.event", (event) => received.push(event)) + bus.on("instance.event", (event) => { + if (event.event.type !== "server.connected") received.push(event) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { @@ -249,7 +297,9 @@ describe("InstanceEventBridge", () => { const { manager, sessionGets } = locationlessManager(events, { unknown: new Error("not found") }) const bus = new EventBus() const received: any[] = [] - bus.on("instance.event", (event) => received.push(event)) + bus.on("instance.event", (event) => { + if (event.event.type !== "server.connected") received.push(event) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { @@ -271,7 +321,9 @@ describe("InstanceEventBridge", () => { const { manager } = locationlessManager(events, {}, workspaces) const bus = new EventBus() const received: any[] = [] - bus.on("instance.event", (event) => received.push(event)) + bus.on("instance.event", (event) => { + if (event.event.type !== "server.connected") received.push(event) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { @@ -294,7 +346,9 @@ describe("InstanceEventBridge", () => { const { manager } = locationlessManager(events, {}, workspaces) const bus = new EventBus() const received: any[] = [] - bus.on("instance.event", (event) => received.push(event)) + bus.on("instance.event", (event) => { + if (event.event.type !== "server.connected") received.push(event) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { @@ -313,7 +367,9 @@ describe("InstanceEventBridge", () => { const { manager, sessionGets } = locationlessManager(events, { deleted: new Error("not found") }, workspaces) const bus = new EventBus() const received: any[] = [] - bus.on("instance.event", (event) => received.push(event)) + bus.on("instance.event", (event) => { + if (event.event.type !== "server.connected") received.push(event) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { @@ -336,7 +392,9 @@ describe("InstanceEventBridge", () => { const { manager } = locationlessManager(events, { foreign: "/repo-b" }, workspaces) const bus = new EventBus() const received: any[] = [] - bus.on("instance.event", (event) => received.push(event)) + bus.on("instance.event", (event) => { + if (event.event.type !== "server.connected") received.push(event) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { @@ -357,7 +415,9 @@ describe("InstanceEventBridge", () => { const { manager, sessionGets } = locationlessManager(events, { owned: "/repo-b" }, workspaces) const bus = new EventBus() const received: any[] = [] - bus.on("instance.event", (event) => received.push(event)) + bus.on("instance.event", (event) => { + if (event.event.type !== "server.connected") received.push(event) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { bus.publish({ type: "workspace.started", workspace: manager.list()[0] as any }) @@ -388,7 +448,9 @@ describe("InstanceEventBridge", () => { const { manager, sessionGets } = locationlessManager(events, {}, workspaces) const bus = new EventBus() const received: any[] = [] - bus.on("instance.event", (event) => received.push(event)) + bus.on("instance.event", (event) => { + if (event.event.type !== "server.connected") received.push(event) + }) const bridge = new InstanceEventBridge({ workspaceManager: manager, eventBus: bus, logger }) try { bus.publish({ type: "workspace.started", workspace: manager.list()[0] as any }) diff --git a/packages/server/src/workspaces/instance-events.ts b/packages/server/src/workspaces/instance-events.ts index c5078ed6..5b9bc709 100644 --- a/packages/server/src/workspaces/instance-events.ts +++ b/packages/server/src/workspaces/instance-events.ts @@ -80,6 +80,9 @@ export class InstanceEventBridge { for await (const event of events) { if (this.controller.signal.aborted) return if (!confirmed) { + if (event.type !== "server.connected") { + throw new Error(`Shared OpenCode event stream started with ${event.type}, expected server.connected`) + } confirmed = true this.updateStatus("connected") } diff --git a/packages/server/src/workspaces/manager.test.ts b/packages/server/src/workspaces/manager.test.ts index 502f5f01..bb6fba85 100644 --- a/packages/server/src/workspaces/manager.test.ts +++ b/packages/server/src/workspaces/manager.test.ts @@ -23,9 +23,12 @@ function deferred() { class ControlledSharedService { readonly validationStarted = deferred() validationGate?: ReturnType> + ignoreValidationAbort = false afterValidation?: () => void validationCalls: Array<{ location: LocationRef; options?: OpenCodeSharedServiceOptions }> = [] shutdownCalls = 0 + shutdownGate?: ReturnType> + shutdownTimeouts: number[] = [] evictionCalls: Array<{ location: LocationRef; options?: OpenCodeSharedServiceOptions; signal?: AbortSignal }> = [] async endpoint() { @@ -44,14 +47,17 @@ class ControlledSharedService { this.validationCalls.push({ location, options }) this.validationStarted.resolve() if (this.validationGate) { - await Promise.race([ - this.validationGate.promise, - new Promise((_resolve, reject) => { - const cancel = () => reject(requestOptions?.signal?.reason) - requestOptions?.signal?.addEventListener("abort", cancel, { once: true }) - if (requestOptions?.signal?.aborted) cancel() - }), - ]) + if (this.ignoreValidationAbort) await this.validationGate.promise + else { + await Promise.race([ + this.validationGate.promise, + new Promise((_resolve, reject) => { + const cancel = () => reject(requestOptions?.signal?.reason) + requestOptions?.signal?.addEventListener("abort", cancel, { once: true }) + if (requestOptions?.signal?.aborted) cancel() + }), + ]) + } } this.afterValidation?.() return { @@ -69,7 +75,11 @@ class ControlledSharedService { return { async *[Symbol.asyncIterator]() {} } } - async shutdown() { this.shutdownCalls += 1 } + async shutdown(options?: { timeoutMs?: number }) { + this.shutdownCalls += 1 + if (options?.timeoutMs !== undefined) this.shutdownTimeouts.push(options.timeoutMs) + await this.shutdownGate?.promise + } } function createHarness(service = new ControlledSharedService(), overrides: Record = {}) { @@ -85,7 +95,7 @@ function createHarness(service = new ControlledSharedService(), overrides: Recor sharedService: service, ...overrides, }) - return { manager, service, stopped } + return { manager, service, stopped, eventBus } } describe("workspace manager shared service lifecycle", () => { @@ -253,14 +263,103 @@ describe("workspace manager shared service lifecycle", () => { assert.equal(harness.manager.list().length, 1) }) - it("shuts down only the local adapter after deleting workspaces", async () => { + it("shuts down local workspaces without evicting their OpenCode locations", async () => { const harness = createHarness() await harness.manager.create(process.cwd()) await harness.manager.shutdown() assert.deepEqual(harness.manager.list(), []) + assert.equal(harness.service.evictionCalls.length, 0) + assert.equal(harness.service.shutdownCalls, 1) + assert.deepEqual(harness.stopped.length, 1) + }) + + it("cancels and removes an in-flight creation without evicting its location", async () => { + const harness = createHarness() + harness.service.validationGate = deferred() + const creation = harness.manager.create(process.cwd()) + await harness.service.validationStarted.promise + + await harness.manager.shutdown() + await assert.rejects(creation, WorkspaceLaunchCancelledError) + + assert.deepEqual(harness.manager.list(), []) + assert.equal((harness.manager as any).workspaces.size, 0) + assert.equal((harness.manager as any).pendingWorkspaceCreations.size, 0) + assert.equal(harness.service.evictionCalls.length, 0) assert.equal(harness.service.shutdownCalls, 1) }) + it("waits for the underlying launch to settle before completing shutdown", async () => { + const harness = createHarness() + harness.service.validationGate = deferred() + harness.service.ignoreValidationAbort = true + const lifecycleEvents: string[] = [] + harness.eventBus.on("workspace.created", () => lifecycleEvents.push("created")) + harness.eventBus.on("workspace.started", () => lifecycleEvents.push("started")) + const creation = harness.manager.create(process.cwd()) + const creationFailure = assert.rejects(creation, WorkspaceLaunchCancelledError) + await harness.service.validationStarted.promise + const record = [...(harness.manager as any).workspaces.values()][0] + let shutdownSettled = false + + const shutdown = harness.manager.shutdown().then(() => { shutdownSettled = true }) + await new Promise((resolve) => setImmediate(resolve)) + assert.equal(shutdownSettled, false) + harness.service.validationGate.resolve() + await shutdown + await creationFailure + + assert.deepEqual(record.location, { directory: process.cwd() }) + assert.equal(record[Object.getOwnPropertySymbols(record)[0]].locationOwned, false) + assert.deepEqual(lifecycleEvents, []) + assert.equal(harness.service.evictionCalls.length, 0) + assert.equal((harness.manager as any).workspaces.size, 0) + }) + + it("prevents a launch completing after bounded shutdown from mutating its removed record", async () => { + const service = new ControlledSharedService() + service.validationGate = deferred() + service.ignoreValidationAbort = true + const validationFinished = deferred() + service.afterValidation = validationFinished.resolve + const harness = createHarness(service, { shutdownTimeoutMs: 20 }) + const lifecycleEvents: string[] = [] + harness.eventBus.on("workspace.created", () => lifecycleEvents.push("created")) + harness.eventBus.on("workspace.started", () => lifecycleEvents.push("started")) + const creation = harness.manager.create(process.cwd()) + const creationFailure = assert.rejects(creation, WorkspaceLaunchCancelledError) + await service.validationStarted.promise + const record = [...(harness.manager as any).workspaces.values()][0] + + await assert.rejects(harness.manager.shutdown(), /Failed to stop 1 workspace during shutdown/) + assert.equal((harness.manager as any).workspaces.size, 0) + service.validationGate.resolve() + await validationFinished.promise + await new Promise((resolve) => setImmediate(resolve)) + await creationFailure + + assert.deepEqual(record.location, { directory: process.cwd() }) + assert.equal(record[Object.getOwnPropertySymbols(record)[0]].locationOwned, false) + assert.deepEqual(lifecycleEvents, []) + assert.equal(service.evictionCalls.length, 0) + }) + + it("bounds local adapter shutdown after removing workspace records", async () => { + const service = new ControlledSharedService() + service.shutdownGate = deferred() + const { manager } = createHarness(service, { shutdownTimeoutMs: 20 }) + await manager.create(process.cwd()) + const startedAt = Date.now() + + await assert.rejects(manager.shutdown(), /Failed to stop 1 workspace during shutdown/) + + assert.ok(Date.now() - startedAt < 1000) + assert.deepEqual(manager.list(), []) + assert.equal(service.evictionCalls.length, 0) + assert.equal(service.shutdownCalls, 1) + assert.ok((service.shutdownTimeouts[0] ?? 0) > 0) + }) + }) diff --git a/packages/server/src/workspaces/manager.ts b/packages/server/src/workspaces/manager.ts index 0ec55a77..e646f955 100644 --- a/packages/server/src/workspaces/manager.ts +++ b/packages/server/src/workspaces/manager.ts @@ -430,13 +430,15 @@ export class WorkspaceManager { } private startCreation(record: WorkspaceRecord, options: WorkspaceCreateOptions, launchDeadlineAt: number, launchTimeoutMs: number): Promise { - const creation = this.createWithDeadline(record, options, launchDeadlineAt, launchTimeoutMs) + const launch = this.createResolvedWorkspace(record, Math.max(1, launchDeadlineAt - Date.now())) + const creation = this.createWithDeadline(record, options, launchDeadlineAt, launchTimeoutMs, launch) record[WORKSPACE_STATE].creation = creation - record[WORKSPACE_STATE].settlement = creation.then(() => undefined, () => undefined) + record[WORKSPACE_STATE].settlement = launch.then(() => undefined, () => undefined) return creation } private async createWithDeadline(record: WorkspaceRecord, options: WorkspaceCreateOptions, - launchDeadlineAt: number, launchTimeoutMs: number): Promise { + launchDeadlineAt: number, launchTimeoutMs: number, + launch: Promise): Promise { const timeoutMs = Math.max(1, launchDeadlineAt - Date.now()) const state = record[WORKSPACE_STATE] let timeout: ManagerTimeout | null = (this.options.setTimeout ?? setTimeout)(() => { @@ -449,7 +451,7 @@ export class WorkspaceManager { const deadline = new Promise((_resolve, reject) => { state.abortController.signal.addEventListener("abort", () => reject(state.abortController.signal.reason), { once: true }) }) - return await Promise.race([this.createResolvedWorkspace(record, timeoutMs), deadline]) + return await Promise.race([launch, deadline]) } finally { if (timeout) (this.options.clearTimeout ?? clearTimeout)(timeout) } @@ -511,6 +513,7 @@ export class WorkspaceManager { serviceOptions, ), ]) + if (this.shuttingDown) this.throwIfCancelled(record) this.serviceAuthorization = headers?.authorization record.location = { directory: location.directory, workspaceID: location.workspaceID } state.locationOwned = true @@ -526,7 +529,7 @@ export class WorkspaceManager { return { workspace: record, created: true } } catch (error) { const launchFailure = state.abortController.signal.aborted ? state.abortController.signal.reason : error - if (state.locationOwned) { + if (state.locationOwned && !this.shuttingDown) { await this.evictRecordLocation(record, timeoutMs).catch((evictionError) => { this.options.logger.warn( { workspaceId: id, err: evictionError }, @@ -640,24 +643,32 @@ export class WorkspaceManager { this.options.logger.info("Shutting down all workspaces") const shutdownTimeoutMs = Math.max(1, this.options.shutdownTimeoutMs ?? 10000) const deadlineAt = Date.now() + shutdownTimeoutMs - const stopTasks = Array.from(this.workspaces.keys(), (id) => this.delete(id)) - const results = stopTasks.length - ? await this.withTimeout(Promise.allSettled(stopTasks), shutdownTimeoutMs, "shutdown") - : [] - const stopFailures = results.flatMap((result) => result.status === "rejected" ? [result.reason] : []) - if (this.workspaces.size === 0) { - this.pendingWorkspaceCreations.clear() - this.cancelledCreationRequests.clear() - const remaining = deadlineAt - Date.now() - if (remaining <= 0) stopFailures.push(new WorkspaceCleanupTimeoutError("shared service shutdown", shutdownTimeoutMs)) - else await this.withTimeout( - this.sharedService.shutdown({ timeoutMs: remaining }), - remaining, - "shared service shutdown", - ).catch((error) => stopFailures.push(error)) - } else if (!stopFailures.length) stopFailures.push( - new Error(`Workspace cleanup remains incomplete for: ${Array.from(this.workspaces.keys()).join(", ")}`), - ) + const records = Array.from(this.workspaces.entries()) + for (const [id, record] of records) { + const state = record[WORKSPACE_STATE] + if (!state.abortController.signal.aborted) state.abortController.abort(new WorkspaceLaunchCancelledError(id)) + } + const settlements = records.map(([, record]) => { + const state = record[WORKSPACE_STATE] + return state.deletePromise ?? state.settlement ?? Promise.resolve() + }) + const stopFailures: unknown[] = [] + if (settlements.length) { + await this.withTimeout(Promise.allSettled(settlements), shutdownTimeoutMs, "shutdown") + .then((results) => { + stopFailures.push(...results.flatMap((result) => result.status === "rejected" ? [result.reason] : [])) + }) + .catch((error) => stopFailures.push(error)) + } + for (const [id, record] of records) this.removeRecord(id, record, true, "stopped") + this.pendingWorkspaceCreations.clear() + this.cancelledCreationRequests.clear() + const remaining = Math.max(1, deadlineAt - Date.now()) + await this.withTimeout( + this.sharedService.shutdown({ timeoutMs: remaining }), + remaining, + "shared service shutdown", + ).catch((error) => stopFailures.push(error)) if (stopFailures.length) throw new WorkspaceShutdownError(stopFailures) } @@ -793,11 +804,16 @@ export class WorkspaceManager { return environment } - private removeRecord(id: string, record: WorkspaceRecord, publishStopped: boolean): void { + private removeRecord( + id: string, + record: WorkspaceRecord, + publishStopped: boolean, + reason: "deleted" | "stopped" = "deleted", + ): void { if (this.workspaces.get(id) !== record) return this.workspaces.delete(id) clearWorkspaceSearchCache(record.path) - if (publishStopped) this.publishStopped(record, "deleted") + if (publishStopped) this.publishStopped(record, reason) } private publishStopped(record: WorkspaceRecord, reason: "deleted" | "stopped" = "stopped"): void { diff --git a/packages/tauri-app/src-tauri/src/cli_manager.rs b/packages/tauri-app/src-tauri/src/cli_manager.rs index b18087f6..0d4dccf2 100644 --- a/packages/tauri-app/src-tauri/src/cli_manager.rs +++ b/packages/tauri-app/src-tauri/src/cli_manager.rs @@ -22,7 +22,7 @@ use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; use std::thread; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; -use tauri::{webview::cookie::Cookie, AppHandle, Emitter, Manager, Url}; +use tauri::{webview::cookie::Cookie, AppHandle, Manager, Url}; #[cfg(windows)] use std::os::windows::io::AsRawHandle; @@ -754,6 +754,17 @@ impl Default for CliStatus { } } +fn cli_exit_error(status: &CliStatus, exit: &std::process::ExitStatus) -> String { + if status.state == CliState::Ready { + format!("CLI exited unexpectedly after readiness: {exit}") + } else { + status + .error + .clone() + .unwrap_or_else(|| format!("CLI exited early: {exit}")) + } +} + #[derive(Debug, Clone)] pub struct CliProcessManager { status: Arc>, @@ -1186,18 +1197,14 @@ impl CliProcessManager { manager.with_current_generation(generation, || { *manager.local_access.lock() = None; let mut status = manager.status.lock(); - if status.state != CliState::Ready { - status.state = CliState::Error; - if status.error.is_none() { - status.error = Some(format!("CLI exited early: {code}")); - } - let _ = app.emit( - "cli:error", - json!({"message": status.error.clone().unwrap_or_default()}), - ); - } else { - status.state = CliState::Stopped; - } + let message = cli_exit_error(&status, &code); + status.state = CliState::Error; + status.error = Some(message.clone()); + crate::local_windows::emit_all( + &app, + "cli:error", + json!({"message": message}), + ); Self::emit_status(&app, &status); }); return; @@ -1785,6 +1792,64 @@ mod tests { assert_eq!(manager.status().state, CliState::Ready); } + #[test] + fn generation_invalidated_exit_cannot_replace_requested_stop_status() { + let manager = CliProcessManager::new(); + let generation = manager.advance_generation(); + manager.status.lock().state = CliState::Ready; + + manager.advance_generation(); + manager.reset_stopped_status(); + + assert!(manager + .with_current_generation(generation, || { + manager.status.lock().state = CliState::Error; + }) + .is_none()); + assert_eq!(manager.status().state, CliState::Stopped); + } + + #[test] + fn unexpected_ready_exit_is_an_error_with_the_platform_status() { + let status = if cfg!(windows) { + Command::new("cmd.exe") + .args(["/C", "exit", "23"]) + .status() + .unwrap() + } else { + Command::new("sh").args(["-c", "exit 23"]).status().unwrap() + }; + + let message = cli_exit_error( + &CliStatus { + state: CliState::Ready, + ..CliStatus::default() + }, + &status, + ); + assert!( + message.contains("unexpectedly after readiness"), + "{message}" + ); + assert!(message.contains("23"), "{message}"); + } + + #[cfg(unix)] + #[test] + fn unexpected_ready_exit_preserves_the_signal() { + use std::os::unix::process::ExitStatusExt; + + let status = std::process::ExitStatus::from_raw(9); + let message = cli_exit_error( + &CliStatus { + state: CliState::Ready, + ..CliStatus::default() + }, + &status, + ); + assert!(message.contains("signal: 9"), "{message}"); + } + #[test] fn local_cli_access_requires_readiness_and_clears_on_stop() { let manager = CliProcessManager::new(); diff --git a/packages/tauri-app/src-tauri/src/client_state.rs b/packages/tauri-app/src-tauri/src/client_state.rs index 048db00b..499ff282 100644 --- a/packages/tauri-app/src-tauri/src/client_state.rs +++ b/packages/tauri-app/src-tauri/src/client_state.rs @@ -1056,12 +1056,11 @@ pub fn release(app: &AppHandle) { } } -pub fn flush_and_release_without_window_capture(app: &AppHandle) { +pub fn flush_without_window_capture(app: &AppHandle) { if let Some(state) = app.try_state::() { if let Err(err) = state.flush() { eprintln!("[client-state] failed to flush state: {err}"); } - state.release_locks(); } } diff --git a/packages/tauri-app/src-tauri/src/client_state/partitions.rs b/packages/tauri-app/src-tauri/src/client_state/partitions.rs index 9791e5e7..37796d2b 100644 --- a/packages/tauri-app/src-tauri/src/client_state/partitions.rs +++ b/packages/tauri-app/src-tauri/src/client_state/partitions.rs @@ -11,7 +11,7 @@ use std::path::{Path, PathBuf}; pub(super) const PROTOCOL_VERSION: u64 = 1; pub(super) const MAX_ROOT_BYTES: usize = 1024 * 1024; const MAX_PARTITION_BYTES: usize = 1024 * 1024; -const MAX_COMMIT_BYTES: usize = 8 * 1024 * 1024; +pub(super) const MAX_COMMIT_BYTES: usize = 256 * 1024 * 1024; const MAX_PARTITION_KEYS: usize = 4096; const PARTITION_DIRECTORY: &str = "partitions"; @@ -122,10 +122,10 @@ impl PartitionCommit { return Err("Client state partition exceeds the 1 MiB limit".to_string()); } commit_size = commit_size.checked_add(content.len()).ok_or_else(|| { - "Client state partition commit exceeds the 8 MiB limit".to_string() + "Client state partition commit exceeds the 256 MiB limit".to_string() })?; if commit_size > MAX_COMMIT_BYTES { - return Err("Client state partition commit exceeds the 8 MiB limit".to_string()); + return Err("Client state partition commit exceeds the 256 MiB limit".to_string()); } if hex_digest(content.as_bytes()) != *key { return Err("Client state partition digest mismatch".to_string()); diff --git a/packages/tauri-app/src-tauri/src/client_state/tests.rs b/packages/tauri-app/src-tauri/src/client_state/tests.rs index 5d039d7d..a5d41079 100644 --- a/packages/tauri-app/src-tauri/src/client_state/tests.rs +++ b/packages/tauri-app/src-tauri/src/client_state/tests.rs @@ -1,4 +1,5 @@ use super::commands::is_allowed_client_state_origin; +use super::partitions::MAX_COMMIT_BYTES; use super::process::{PRIMARY_LOCK_FILENAME, RUNNING_MARKER_PREFIX, RUNNING_MARKER_SUFFIX}; use super::window::{ clamp_window_bounds, normalize_native_zoom_level, DisplayArea, NativeWindowState, WindowBounds, @@ -1160,6 +1161,7 @@ fn invalid_v3_is_frozen_until_explicit_clear() { #[test] fn partition_protocol_and_hashes_are_validated() { + assert_eq!(MAX_COMMIT_BYTES, 256 * 1024 * 1024); let directory = tempfile::tempdir().unwrap(); let state = ClientState::initialize_at(directory.path()).unwrap(); enable_restore(&state); @@ -1180,14 +1182,14 @@ fn partition_protocol_and_hashes_are_validated() { let mut partitions = serde_json::Map::new(); let mut partition_keys = Vec::new(); - for index in 0..8 { + for index in 0..9 { let content = format!("{index}{}", "x".repeat(1024 * 1024 - 1)); let key = partition_key(&content); partitions.insert(key.clone(), Value::String(content)); partition_keys.push(key); } partition_keys.sort(); - let error = state + assert!(state .commit_partitions_guarded( partition_commit(json!({ "protocolVersion": 1, "snapshot": partition_root(&partition_keys, json!({})), @@ -1195,8 +1197,7 @@ fn partition_protocol_and_hashes_are_validated() { })), || true, ) - .unwrap_err(); - assert!(error.contains("8 MiB")); + .unwrap()); } #[test] diff --git a/packages/tauri-app/src-tauri/src/main.rs b/packages/tauri-app/src-tauri/src/main.rs index 46a58bba..18319d16 100644 --- a/packages/tauri-app/src-tauri/src/main.rs +++ b/packages/tauri-app/src-tauri/src/main.rs @@ -214,6 +214,14 @@ struct RemoteWindowPayload { skip_tls_verify: bool, } +fn require_http_url(value: &str, name: &str) -> Result { + let url = Url::parse(value).map_err(|error| error.to_string())?; + if !matches!(url.scheme(), "http" | "https") { + return Err(format!("{name} must use HTTP or HTTPS")); + } + Ok(url) +} + fn schedule_remote_proxy_session_cleanup(app: AppHandle, session_id: String) { tauri::async_runtime::spawn(async move { if let Err(err) = cleanup_remote_proxy_session(&app, &session_id).await { @@ -407,11 +415,12 @@ async fn open_remote_window_impl( app: AppHandle, payload: RemoteWindowPayload, ) -> Result<(), String> { + require_http_url(&payload.base_url, "baseUrl")?; let entry_url = payload .entry_url .as_deref() .unwrap_or(payload.base_url.as_str()); - let parsed = Url::parse(entry_url).map_err(|err| err.to_string())?; + let parsed = require_http_url(entry_url, "entryUrl")?; let label = format!("remote-{}", payload.id); let title = format!("{} - {}", payload.name, payload.base_url); let requested_profile = RemoteProfileIdentity::new(payload.proxy_session_id.as_deref()); @@ -681,7 +690,8 @@ async fn open_remote_window( .entry_url .as_deref() .unwrap_or(payload.base_url.as_str()); - let parsed = Url::parse(entry_url).map_err(|err| err.to_string())?; + require_http_url(&payload.base_url, "baseUrl")?; + let parsed = require_http_url(entry_url, "entryUrl")?; if payload.proxy_session_id.is_some() && parsed.scheme() == "https" { let local_cert = cert_manager::ensure_local_cert().map_err(|err| { format!( @@ -1557,7 +1567,7 @@ fn build_about_metadata(version: &str, include_update_link: bool) -> AboutMetada #[cfg(test)] mod menu_tests { use super::{ - build_about_metadata, is_allowed_local_origin, run_update_with_fallback, + build_about_metadata, is_allowed_local_origin, require_http_url, run_update_with_fallback, should_allow_registered_origin, should_recreate_remote_window, RemoteProfileIdentity, WakeLockState, RELEASES_URL, REMOTE_WINDOW_CONTEXT_SCRIPT, }; @@ -1671,6 +1681,31 @@ mod menu_tests { )); } + #[test] + fn remote_window_urls_require_http_or_https() { + assert_eq!( + require_http_url("http://localhost:3000/app", "baseUrl") + .unwrap() + .scheme(), + "http" + ); + assert_eq!( + require_http_url("https://example.com/app", "entryUrl") + .unwrap() + .scheme(), + "https" + ); + for value in [ + "file:///tmp/app", + "data:text/html,hi", + "javascript:alert(1)", + ] { + assert!(require_http_url(value, "baseUrl") + .unwrap_err() + .contains("must use HTTP or HTTPS")); + } + } + #[test] fn remote_window_reuse_requires_exact_profile_identity() { let direct = RemoteProfileIdentity::Direct; diff --git a/packages/tauri-app/src-tauri/src/shutdown.rs b/packages/tauri-app/src-tauri/src/shutdown.rs index 21b1078d..cf69b86d 100644 --- a/packages/tauri-app/src-tauri/src/shutdown.rs +++ b/packages/tauri-app/src-tauri/src/shutdown.rs @@ -1,9 +1,7 @@ use crate::{client_state, local_windows::LocalWindows, AppState}; use std::collections::{HashMap, HashSet}; -#[cfg(windows)] -use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Mutex; -use std::time::Duration; +use std::time::{Duration, Instant}; #[cfg(windows)] use tauri::WebviewWindow; use tauri::{AppHandle, Emitter, Manager}; @@ -21,22 +19,38 @@ struct PendingClose { persisted: bool, } +#[cfg(windows)] +struct WindowsSessionEndPreparation { + generation: u64, + deadline: Instant, + requests: Option>, +} + #[derive(Default)] struct ShutdownState { next_generation: u64, local_closes: HashMap, close_allowed: HashSet, global_pending: HashMap, + global_requests: HashMap, shutdown_started: bool, cleanup_started: bool, exit_allowed: bool, + #[cfg(windows)] + windows_session_end_generation: Option, + #[cfg(windows)] + windows_session_end_deadline: Option, + #[cfg(windows)] + windows_session_end_owns_shutdown: bool, + #[cfg(windows)] + windows_renderer_deadline: Option, + #[cfg(windows)] + windows_native_flush_complete: bool, } #[derive(Default)] pub(crate) struct ShutdownCoordinator { state: Mutex, - #[cfg(windows)] - windows_session_end_started: AtomicBool, } impl ShutdownCoordinator { @@ -85,11 +99,13 @@ impl ShutdownCoordinator { } state.shutdown_started = true; state.local_closes.clear(); + state.global_requests.clear(); let mut requests = Vec::new(); for label in labels { state.next_generation += 1; let generation = state.next_generation; state.global_pending.insert(label.clone(), generation); + state.global_requests.insert(label.clone(), generation); requests.push((label, generation)); } Some(requests) @@ -112,17 +128,31 @@ impl ShutdownCoordinator { { return false; } + #[cfg(windows)] + if state.windows_session_end_owns_shutdown { + return false; + } state.global_pending.clear(); state.cleanup_started = true; true } - fn cleanup_failed(&self) { + fn cleanup_failed(&self) -> Vec<(String, u64)> { let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + let cancellations = state.global_requests.drain().collect(); state.cleanup_started = false; state.shutdown_started = false; state.global_pending.clear(); state.close_allowed.clear(); + #[cfg(windows)] + { + state.windows_session_end_generation = None; + state.windows_session_end_deadline = None; + state.windows_session_end_owns_shutdown = false; + state.windows_renderer_deadline = None; + state.windows_native_flush_complete = false; + } + cancellations } fn commit_local_close(&self, label: &str) -> bool { @@ -178,10 +208,143 @@ impl ShutdownCoordinator { } #[cfg(windows)] - fn begin_windows_session_end(&self) -> bool { - !self - .windows_session_end_started - .swap(true, Ordering::SeqCst) + fn begin_windows_session_end( + &self, + labels: impl IntoIterator, + ) -> WindowsSessionEndPreparation { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if let Some(generation) = state.windows_session_end_generation { + return WindowsSessionEndPreparation { + generation, + deadline: state + .windows_session_end_deadline + .unwrap_or_else(Instant::now), + requests: None, + }; + } + state.next_generation += 1; + let session_generation = state.next_generation; + let session_deadline = Instant::now() + WINDOWS_SESSION_END_TIMEOUT; + state.windows_session_end_generation = Some(session_generation); + state.windows_session_end_deadline = Some(session_deadline); + if state.shutdown_started { + return WindowsSessionEndPreparation { + generation: session_generation, + deadline: session_deadline, + requests: Some(Vec::new()), + }; + } + state.windows_session_end_owns_shutdown = true; + state.windows_renderer_deadline = Some(Instant::now() + RENDERER_FLUSH_TIMEOUT); + state.windows_native_flush_complete = false; + state.shutdown_started = true; + state.local_closes.clear(); + state.global_requests.clear(); + let mut requests = Vec::new(); + for label in labels { + state.next_generation += 1; + let generation = state.next_generation; + state.global_pending.insert(label.clone(), generation); + state.global_requests.insert(label.clone(), generation); + requests.push((label, generation)); + } + WindowsSessionEndPreparation { + generation: session_generation, + deadline: session_deadline, + requests: Some(requests), + } + } + + #[cfg(windows)] + fn windows_session_end_owns_shutdown(&self, generation: u64) -> bool { + self.state + .lock() + .map(|state| { + state.windows_session_end_generation == Some(generation) + && state.windows_session_end_owns_shutdown + }) + .unwrap_or(false) + } + + #[cfg(windows)] + fn windows_renderer_wait(&self, generation: u64) -> Option<(bool, Instant)> { + self.state.lock().ok().and_then(|state| { + (state.windows_session_end_generation == Some(generation) + && state.windows_session_end_owns_shutdown) + .then(|| { + state + .windows_renderer_deadline + .map(|deadline| (!state.global_pending.is_empty(), deadline)) + }) + .flatten() + }) + } + + #[cfg(windows)] + fn begin_windows_cleanup(&self, generation: u64) -> bool { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.windows_session_end_generation != Some(generation) + || !state.windows_session_end_owns_shutdown + || state.cleanup_started + { + return false; + } + state.cleanup_started = true; + state.global_pending.clear(); + true + } + + #[cfg(windows)] + fn complete_windows_native_flush(&self, generation: u64) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.windows_session_end_generation == Some(generation) + && state.windows_session_end_owns_shutdown + { + state.windows_native_flush_complete = true; + } + } + + #[cfg(windows)] + fn windows_native_flush_complete(&self, generation: u64) -> bool { + self.state + .lock() + .map(|state| { + state.windows_session_end_generation == Some(generation) + && state.windows_native_flush_complete + }) + .unwrap_or(false) + } + + #[cfg(windows)] + fn windows_session_end_remaining(&self, generation: u64, now: Instant) -> Option { + self.state.lock().ok().and_then(|state| { + (state.windows_session_end_generation == Some(generation)) + .then(|| { + state + .windows_session_end_deadline + .map(|deadline| deadline.saturating_duration_since(now)) + }) + .flatten() + }) + } + + #[cfg(windows)] + fn cancel_windows_session_end(&self) -> Vec<(String, u64)> { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.windows_session_end_generation.is_none() { + return Vec::new(); + } + state.windows_session_end_generation = None; + state.windows_session_end_deadline = None; + state.windows_renderer_deadline = None; + state.windows_native_flush_complete = false; + if !state.windows_session_end_owns_shutdown || state.cleanup_started { + return Vec::new(); + } + state.windows_session_end_owns_shutdown = false; + state.shutdown_started = false; + state.global_pending.clear(); + state.global_requests.drain().collect() } } @@ -196,9 +359,18 @@ fn emit_flush(app: &AppHandle, label: &str, generation: u64) -> bool { }) } -fn emit_flush_cancelled(app: &AppHandle, label: &str) { +fn emit_flush_cancelled(app: &AppHandle, label: &str, generation: u64) { if let Some(window) = app.get_webview_window(label) { - let _ = window.emit(FLUSH_CANCELLED_EVENT, ()); + let _ = window.emit( + FLUSH_CANCELLED_EVENT, + client_state::RendererFlushRequest { generation }, + ); + } +} + +fn emit_flush_cancellations(app: &AppHandle, cancellations: Vec<(String, u64)>) { + for (label, generation) in cancellations { + emit_flush_cancelled(app, &label, generation); } } @@ -244,7 +416,7 @@ fn finish_local_close(app: AppHandle, label: String, window_id: String, generati eprintln!("[client-state] failed to remove closed window: {error}"); app.state::() .rollback_local_close(&label); - emit_flush_cancelled(&app, &label); + emit_flush_cancelled(&app, &label, generation); return; } } @@ -258,7 +430,7 @@ fn finish_local_close(app: AppHandle, label: String, window_id: String, generati close_app .state::() .rollback_local_close(&close_label); - emit_flush_cancelled(&close_app, &close_label); + emit_flush_cancelled(&close_app, &close_label, generation); } } }) @@ -266,7 +438,7 @@ fn finish_local_close(app: AppHandle, label: String, window_id: String, generati { app.state::() .rollback_local_close(&label); - emit_flush_cancelled(&app, &label); + emit_flush_cancelled(&app, &label, generation); } }); } @@ -335,8 +507,8 @@ fn start_cleanup(app: AppHandle, deadline_reached: bool) { }; if let Err(error) = result { eprintln!("[tauri] shutdown cleanup remains unconfirmed: {error}"); - app.state::().cleanup_failed(); - crate::local_windows::emit_all(&app, FLUSH_CANCELLED_EVENT, ()); + let cancellations = app.state::().cleanup_failed(); + emit_flush_cancellations(&app, cancellations); return; } client_state::release(&app); @@ -378,27 +550,104 @@ pub(crate) fn exit_allowed(app: &AppHandle) -> bool { } #[cfg(windows)] -pub(crate) fn request_windows_session_end(app: AppHandle) { +fn prepare_windows_session_end(app: &AppHandle) -> (u64, Instant) { + let labels = app + .state::() + .records() + .into_iter() + .map(|record| record.label) + .collect::>(); + let preparation = app + .state::() + .begin_windows_session_end(labels); + let generation = preparation.generation; + let deadline = preparation.deadline; + let Some(requests) = preparation.requests else { + return (generation, deadline); + }; + for (label, flush_generation) in requests { + if !emit_flush(app, &label, flush_generation) { + app.state::() + .acknowledge_global(&label, flush_generation); + } + } if !app .state::() - .begin_windows_session_end() + .windows_session_end_owns_shutdown(generation) { - return; + return (generation, deadline); } - if app.state::().shutdown_started() { - let deadline = std::time::Instant::now() + WINDOWS_SESSION_END_TIMEOUT; - while !exit_allowed(&app) && std::time::Instant::now() < deadline { + let flush_app = app.clone(); + std::thread::spawn(move || { + while flush_app + .state::() + .windows_renderer_wait(generation) + .is_some_and(|(pending, deadline)| pending && Instant::now() < deadline) + { + std::thread::sleep(Duration::from_millis(10)); + } + if !flush_app + .state::() + .windows_session_end_owns_shutdown(generation) + { + return; + } + client_state::flush_without_window_capture(&flush_app); + flush_app + .state::() + .complete_windows_native_flush(generation); + }); + (generation, deadline) +} + +#[cfg(windows)] +fn cancel_windows_session_end(app: &AppHandle) { + let cancellations = app + .state::() + .cancel_windows_session_end(); + emit_flush_cancellations(app, cancellations); +} + +#[cfg(windows)] +pub(crate) fn request_windows_session_end(app: AppHandle) { + let (generation, session_deadline) = prepare_windows_session_end(&app); + if !app + .state::() + .windows_session_end_owns_shutdown(generation) + { + while !exit_allowed(&app) && Instant::now() < session_deadline { std::thread::sleep(Duration::from_millis(10)); } return; } - let _ = app + + while !app .state::() - .begin_shutdown(std::iter::empty()); + .windows_native_flush_complete(generation) + && Instant::now() < session_deadline + { + std::thread::sleep(Duration::from_millis(10)); + } + if !app + .state::() + .windows_native_flush_complete(generation) + { + eprintln!( + "[tauri] Windows session-end state flush exceeded {:?}", + WINDOWS_SESSION_END_TIMEOUT + ); + return; + } + if !app + .state::() + .begin_windows_cleanup(generation) + { + return; + } + let (finished_tx, finished_rx) = std::sync::mpsc::sync_channel(1); let cleanup_app = app.clone(); std::thread::spawn(move || { - client_state::flush_and_release_without_window_capture(&cleanup_app); let result = { cleanup_app .try_state::() @@ -409,13 +658,15 @@ pub(crate) fn request_windows_session_end(app: AppHandle) { }) .unwrap_or(Ok(())) }; + client_state::release(&cleanup_app); cleanup_app.state::().allow_exit(); let _ = finished_tx.send(result); }); - if finished_rx - .recv_timeout(WINDOWS_SESSION_END_TIMEOUT) - .is_err() - { + let remaining = app + .state::() + .windows_session_end_remaining(generation, Instant::now()) + .unwrap_or_default(); + if finished_rx.recv_timeout(remaining).is_err() { eprintln!( "[tauri] Windows session-end cleanup exceeded {:?}", WINDOWS_SESSION_END_TIMEOUT @@ -450,11 +701,19 @@ unsafe extern "system" fn windows_session_end_proc( return result; } if message == WM_QUERYENDSESSION { + // Windows requires this message to return promptly. Start the bounded state-flush + // worker now, but leave CLI cleanup and lock release for WM_ENDSESSION. + let context = &*(reference_data as *const WindowsSessionEndContext); + prepare_windows_session_end(&context.app); return 1; } - if message == WM_ENDSESSION && wparam != 0 { + if message == WM_ENDSESSION { let context = &*(reference_data as *const WindowsSessionEndContext); - request_windows_session_end(context.app.clone()); + if wparam != 0 { + request_windows_session_end(context.app.clone()); + } else { + cancel_windows_session_end(&context.app); + } return 0; } DefSubclassProc(hwnd, message, wparam, lparam) diff --git a/packages/tauri-app/src-tauri/src/shutdown_tests.rs b/packages/tauri-app/src-tauri/src/shutdown_tests.rs index 68c18eff..df26b9f0 100644 --- a/packages/tauri-app/src-tauri/src/shutdown_tests.rs +++ b/packages/tauri-app/src-tauri/src/shutdown_tests.rs @@ -45,8 +45,8 @@ fn final_shutdown_and_cleanup_start_once() { fn failed_shutdown_paths_emit_the_renderer_resume_event() { let source = include_str!("shutdown.rs"); assert_eq!(FLUSH_CANCELLED_EVENT, "client-state:flush-cancelled"); - assert!(source.matches("emit_flush_cancelled(&app, &label)").count() >= 2); - assert!(source.contains("emit_all(&app, FLUSH_CANCELLED_EVENT")); + assert!(source.contains("RendererFlushRequest { generation }")); + assert!(source.contains("emit_flush_cancellations(&app, cancellations)")); } #[test] @@ -76,3 +76,106 @@ fn bounded_retry_stops_after_success() { .unwrap(); assert_eq!(calls, 2); } + +#[test] +fn failed_global_shutdown_cancels_the_exact_renderer_generations() { + let coordinator = ShutdownCoordinator::default(); + let mut requests = coordinator + .begin_shutdown(["local-a".to_string(), "local-b".to_string()]) + .unwrap(); + for (label, generation) in &requests { + coordinator.acknowledge_global(label, *generation); + } + assert!(coordinator.begin_cleanup(false)); + + let mut cancellations = coordinator.cleanup_failed(); + requests.sort(); + cancellations.sort(); + assert_eq!(cancellations, requests); +} + +#[cfg(windows)] +#[test] +fn windows_session_end_defers_native_cleanup_until_renderer_flush_finishes() { + let coordinator = ShutdownCoordinator::default(); + let preparation = coordinator.begin_windows_session_end(["local-a".to_string()]); + let generation = preparation.generation; + let requests = preparation.requests.unwrap(); + + assert!(!coordinator.begin_cleanup(false)); + assert!(coordinator.windows_renderer_wait(generation).unwrap().0); + assert!(coordinator.acknowledge_global(&requests[0].0, requests[0].1)); + assert!(!coordinator.windows_renderer_wait(generation).unwrap().0); + assert!(!coordinator.windows_native_flush_complete(generation)); + coordinator.complete_windows_native_flush(generation); + assert!(coordinator.windows_native_flush_complete(generation)); + assert!(coordinator.begin_windows_cleanup(generation)); + assert!(!coordinator.begin_windows_cleanup(generation)); +} + +#[cfg(windows)] +#[test] +fn cancelled_windows_session_end_reopens_renderer_persistence() { + let coordinator = ShutdownCoordinator::default(); + coordinator.begin_windows_session_end(["local-a".to_string()]); + + assert_eq!(coordinator.cancel_windows_session_end().len(), 1); + assert!(!coordinator.shutdown_started()); + assert!(coordinator.navigation_allowed()); + assert!(coordinator.begin_shutdown(std::iter::empty()).is_some()); +} + +#[cfg(windows)] +#[test] +fn multi_window_session_end_prepare_reuses_generation_deadline_and_requests() { + let coordinator = ShutdownCoordinator::default(); + let first = + coordinator.begin_windows_session_end(["local-a".to_string(), "local-b".to_string()]); + let second = coordinator.begin_windows_session_end(["local-c".to_string()]); + + assert_eq!(first.requests.as_ref().unwrap().len(), 2); + assert_eq!(second.generation, first.generation); + assert_eq!(second.deadline, first.deadline); + assert!(second.requests.is_none()); +} + +#[cfg(windows)] +#[test] +fn repeated_session_end_calls_share_timeout_and_cleanup_claim() { + let coordinator = ShutdownCoordinator::default(); + let preparation = coordinator.begin_windows_session_end(std::iter::empty()); + let generation = preparation.generation; + + assert_eq!( + coordinator.windows_session_end_remaining(generation, preparation.deadline), + Some(Duration::ZERO) + ); + coordinator.complete_windows_native_flush(generation); + assert!(coordinator.begin_windows_cleanup(generation)); + assert!(!coordinator.begin_windows_cleanup(generation)); + assert_eq!( + coordinator.windows_session_end_remaining( + generation, + preparation.deadline + Duration::from_secs(1) + ), + Some(Duration::ZERO) + ); +} + +#[cfg(windows)] +#[test] +fn cancelled_session_worker_cannot_complete_a_new_session_flush() { + let coordinator = ShutdownCoordinator::default(); + let first = coordinator + .begin_windows_session_end(["local-a".to_string()]) + .generation; + coordinator.cancel_windows_session_end(); + let second = coordinator + .begin_windows_session_end(["local-a".to_string()]) + .generation; + + coordinator.complete_windows_native_flush(first); + assert!(!coordinator.windows_native_flush_complete(second)); + coordinator.complete_windows_native_flush(second); + assert!(coordinator.windows_native_flush_complete(second)); +} diff --git a/packages/ui/src/components/message-history-pagination.test.ts b/packages/ui/src/components/message-history-pagination.test.ts new file mode 100644 index 00000000..62a517eb --- /dev/null +++ b/packages/ui/src/components/message-history-pagination.test.ts @@ -0,0 +1,28 @@ +import assert from "node:assert/strict" +import { describe, it } from "node:test" + +import { MESSAGE_HISTORY_TOP_THRESHOLD_PX, shouldLoadOlderMessages } from "./message-history-pagination.ts" + +describe("message history pagination", () => { + const ready = { + active: true, + failed: false, + hasMore: true, + loading: false, + messageCount: 2, + scrollTop: MESSAGE_HISTORY_TOP_THRESHOLD_PX, + } + + it("loads at the top threshold", () => { + assert.equal(shouldLoadOlderMessages(ready), true) + assert.equal(shouldLoadOlderMessages({ ...ready, scrollTop: MESSAGE_HISTORY_TOP_THRESHOLD_PX + 1 }), false) + }) + + it("guards inactive, exhausted, concurrent, failed, and empty loads", () => { + assert.equal(shouldLoadOlderMessages({ ...ready, active: false }), false) + assert.equal(shouldLoadOlderMessages({ ...ready, hasMore: false }), false) + assert.equal(shouldLoadOlderMessages({ ...ready, loading: true }), false) + assert.equal(shouldLoadOlderMessages({ ...ready, failed: true }), false) + assert.equal(shouldLoadOlderMessages({ ...ready, messageCount: 0 }), false) + }) +}) diff --git a/packages/ui/src/components/message-history-pagination.ts b/packages/ui/src/components/message-history-pagination.ts new file mode 100644 index 00000000..6f1015f6 --- /dev/null +++ b/packages/ui/src/components/message-history-pagination.ts @@ -0,0 +1,17 @@ +export const MESSAGE_HISTORY_TOP_THRESHOLD_PX = 320 + +export function shouldLoadOlderMessages(options: { + active: boolean + failed: boolean + hasMore: boolean + loading: boolean + messageCount: number + scrollTop: number +}): boolean { + return options.active + && !options.failed + && options.hasMore + && !options.loading + && options.messageCount > 0 + && options.scrollTop <= MESSAGE_HISTORY_TOP_THRESHOLD_PX +} diff --git a/packages/ui/src/components/message-section.tsx b/packages/ui/src/components/message-section.tsx index f9bf20b7..143ea9a5 100644 --- a/packages/ui/src/components/message-section.tsx +++ b/packages/ui/src/components/message-section.tsx @@ -21,6 +21,8 @@ import { getMessageSelectionActionPosition } from "../lib/message-selection-posi import { buildSessionSearchMatches } from "../lib/session-search" import type { SessionSearchMatch } from "../lib/session-search" import { resolveThinkingExpansionDefault, resolveToolVisibility } from "./tool-call/tool-registry" +import { MESSAGE_HISTORY_TOP_THRESHOLD_PX, shouldLoadOlderMessages } from "./message-history-pagination" +import { getLogger } from "../lib/logger" const MESSAGE_SCROLL_CACHE_SCOPE = "message-stream" const QUOTE_SELECTION_MAX_LENGTH = 2000 @@ -28,6 +30,7 @@ const STREAMING_TEXT_HOLD_TOP_THRESHOLD_PX = 8 const SEARCH_DEBOUNCE_MS = 250 const SEARCH_MIN_CHARS = 3 const OPEN_SESSION_SEARCH_EVENT = "codenomad:open-session-search" +const log = getLogger("session") export interface MessageSectionProps { instanceId: string @@ -43,6 +46,8 @@ export interface MessageSectionProps { forceCompactStatusLayout?: boolean onQuoteSelection?: (text: string, mode: "quote" | "code") => void onReloadMessages?: () => void + hasMoreMessages?: boolean + onLoadMoreMessages?: () => Promise isActive?: boolean sessionStreamingActive?: boolean explicitBottomPinIntent?: VirtualExplicitBottomPinIntent | null @@ -242,6 +247,8 @@ export default function MessageSection(props: MessageSectionProps) { let restoringScrollSnapshot = false let restoredWithoutSnapshot = false let scrollRestoreGeneration = 0 + let loadingOlderMessages = false + let olderMessageLoadFailed = false function getLastGoodScrollSnapshot(sessionId: string) { return lastGoodScrollSnapshots.get(sessionId) ?? store().getScrollSnapshot(sessionId, MESSAGE_SCROLL_CACHE_SCOPE) @@ -258,6 +265,8 @@ export default function MessageSection(props: MessageSectionProps) { scrollRestoreGeneration += 1 restoringScrollSnapshot = false restoredWithoutSnapshot = false + loadingOlderMessages = false + olderMessageLoadFailed = false setDidRestoreScroll(false) const snapshot = store().getScrollSnapshot(props.sessionId, MESSAGE_SCROLL_CACHE_SCOPE) if (snapshot) setLastGoodScrollSnapshot(props.sessionId, snapshot) @@ -587,6 +596,52 @@ export default function MessageSection(props: MessageSectionProps) { listApi()?.notifyContentRendered() } + async function maybeLoadOlderMessages() { + const api = listApi() + const snapshot = api?.captureScrollSnapshot() + if (!api || !snapshot || !props.onLoadMoreMessages) return + if (!shouldLoadOlderMessages({ + active: isActive(), + failed: olderMessageLoadFailed, + hasMore: Boolean(props.hasMoreMessages), + loading: Boolean(props.loading) || loadingOlderMessages, + messageCount: visibleMessageIds().length, + scrollTop: snapshot.scrollTop, + })) return + + const sessionId = props.sessionId + const firstMessageId = visibleMessageIds()[0] + const anchorSnapshot = snapshot.atBottom && firstMessageId + ? { ...snapshot, atBottom: false, anchorKey: firstMessageId, anchorOffset: 0, followModeType: "escaped" as const } + : snapshot + loadingOlderMessages = true + try { + await props.onLoadMoreMessages() + if (props.sessionId !== sessionId || listApi() !== api) return + await new Promise((resolve) => api.restoreScrollSnapshot(anchorSnapshot, { + behavior: "auto", + fallback: resolve, + onApplied: resolve, + onCancelled: resolve, + })) + } catch (error) { + olderMessageLoadFailed = true + log.error("Failed to load older messages", { instanceId: props.instanceId, sessionId, error }) + } finally { + loadingOlderMessages = false + } + + if (!olderMessageLoadFailed) void maybeLoadOlderMessages() + } + + createEffect(() => { + if (!didRestoreScroll()) return + props.loading + props.hasMoreMessages + visibleMessageIds().length + void maybeLoadOlderMessages() + }) + createEffect(() => { if (!props.onQuoteSelection) { clearQuoteSelection() @@ -758,6 +813,9 @@ export default function MessageSection(props: MessageSectionProps) { onScroll={() => { clearQuoteSelection() persistMessageScrollSnapshot() + const scrollTop = listApi()?.captureScrollSnapshot()?.scrollTop + if (typeof scrollTop === "number" && scrollTop > MESSAGE_HISTORY_TOP_THRESHOLD_PX) olderMessageLoadFailed = false + void maybeLoadOlderMessages() }} onMouseUp={() => handleStreamMouseUp()} onActiveKeyChange={(messageId) => { diff --git a/packages/ui/src/components/session/session-view.tsx b/packages/ui/src/components/session/session-view.tsx index ee6a4292..bddc25ec 100644 --- a/packages/ui/src/components/session/session-view.tsx +++ b/packages/ui/src/components/session/session-view.tsx @@ -8,7 +8,7 @@ import PromptInput from "../prompt-input" import PromptAttachmentsBar from "../prompt-input/PromptAttachmentsBar" import { getAttachments, removeAttachment } from "../../stores/attachments" import { instances, waitForInstanceWorkspaceMetadataHydration } from "../../stores/instances" -import { loadMessages, sendMessage, forkSession, renameSession, isSessionMessagesLoading, getSessionMessagesLoadError, markSessionIdleSeen, ensureSessionAncestorsExpanded, setActiveSessionFromList, runShellCommand, abortSession } from "../../stores/sessions" +import { hasMoreMessages, loadMessages, loadMoreMessages, sendMessage, forkSession, renameSession, isSessionMessagesLoading, getSessionMessagesLoadError, markSessionIdleSeen, ensureSessionAncestorsExpanded, setActiveSessionFromList, runShellCommand, abortSession } from "../../stores/sessions" import { canMarkSessionIdleSeen } from "./session-idle-attention" import { clearSessionIdleFade, IDLE_STATUS_VISIBILITY_MS, getSessionStatus, isSessionBusy as getSessionBusyStatus, markSessionIdleFadeStarted } from "../../stores/session-status" import { showAlertDialog } from "../../stores/alerts" @@ -526,6 +526,8 @@ export const SessionView: Component = (props) => { loading={messagesLoading()} loadError={messagesLoadError()} onReloadMessages={handleReloadMessages} + hasMoreMessages={hasMoreMessages(props.instanceId, activeSession.id)} + onLoadMoreMessages={() => loadMoreMessages(props.instanceId, activeSession.id)} sessionStreamingActive={sessionStreamingActive()} explicitBottomPinIntent={activeSubmitBottomPinIntent()} onExplicitBottomPinCancelled={() => setSubmitBottomPinIntent(null)} diff --git a/packages/ui/src/lib/hooks/use-app-session-capture.test.ts b/packages/ui/src/lib/hooks/use-app-session-capture.test.ts index 090a5bee..eaafcb59 100644 --- a/packages/ui/src/lib/hooks/use-app-session-capture.test.ts +++ b/packages/ui/src/lib/hooks/use-app-session-capture.test.ts @@ -10,16 +10,17 @@ it("preserves settled tabs during native shutdown", () => { }) it("makes native shutdown terminal for reactive captures", () => { - assert.match(capture, /if \(nativeShutdown\) nativeShutdownStarted = true/) - assert.match(capture, /if \(!enabled\(\) \|\| disposed \|\| nativeShutdownStarted\) return/) + assert.match(capture, /if \(nativeShutdown\) nativeShutdownGeneration = nativeShutdownGenerationRequest/) + assert.match(capture, /if \(!enabled\(\) \|\| disposed \|\| nativeShutdownGeneration !== null\) return/) }) it("keeps navigation flushes nonterminal", () => { - assert.match(capture, /flush\(nativeShutdown\)/) + assert.match(capture, /flush\(nativeShutdown \? payload\.generation : undefined\)/) assert.match(capture, /"client-state:flush-requested",[\s\S]*?, true\)/) assert.match(capture, /"client-state:navigation-flush-requested",[\s\S]*?, false\)/) }) it("resumes capture only after native shutdown cancellation", () => { - assert.match(capture, /listen\("client-state:flush-cancelled"[\s\S]*nativeShutdownStarted = false[\s\S]*schedule\(\)/) + assert.match(capture, /listen<\{ generation: number \}>\("client-state:flush-cancelled"/) + assert.match(capture, /if \(nativeShutdownGeneration !== payload\.generation\) return[\s\S]*nativeShutdownGeneration = null[\s\S]*schedule\(\)/) }) diff --git a/packages/ui/src/lib/hooks/use-app-session-capture.ts b/packages/ui/src/lib/hooks/use-app-session-capture.ts index a6a3290e..7c7e5bd4 100644 --- a/packages/ui/src/lib/hooks/use-app-session-capture.ts +++ b/packages/ui/src/lib/hooks/use-app-session-capture.ts @@ -125,7 +125,6 @@ export function useAppSessionCapture() { const instanceLifecycleTokens = new Map() let nextInstanceLifecycleToken = 0 let disposed = false - let nativeShutdownStarted = false const hydrationController = new AbortController() let timer: ReturnType | null = null let preservation: RestorableSessionPreservation | null = null @@ -151,21 +150,23 @@ export function useAppSessionCapture() { currentTabIds: captured.tabIds, currentTabAuthorities: captured.authorities, }) } + let nativeShutdownGeneration: number | null = null const capture = () => { timer = null - if (enabled() && !disposed && !nativeShutdownStarted) { + if (enabled() && !disposed && nativeShutdownGeneration === null) { const state = mergedState() if (state.tabs.length > 0) nativeFallbackState = state updateRestorableSession(state) } } const schedule = () => { - if (!enabled() || disposed || nativeShutdownStarted) return + if (!enabled() || disposed || nativeShutdownGeneration !== null) return if (timer) clearTimeout(timer) timer = setTimeout(capture, CAPTURE_DEBOUNCE_MS) } - const flush = async (nativeShutdown = false) => { - if (nativeShutdown) nativeShutdownStarted = true + const flush = async (nativeShutdownGenerationRequest?: number) => { + const nativeShutdown = nativeShutdownGenerationRequest !== undefined + if (nativeShutdown) nativeShutdownGeneration = nativeShutdownGenerationRequest if (timer) clearTimeout(timer) timer = null if (enabled()) { @@ -181,8 +182,8 @@ export function useAppSessionCapture() { } const nativeUnlisteners: Array<() => void> = [] let nativeDisposed = false - const register = (event: string, acknowledge: (payload: T) => void | Promise, nativeShutdown: boolean) => listen(event, ({ payload }) => { - void flush(nativeShutdown).then(() => acknowledge(payload)).catch((error) => log.error(`Failed to handle ${event}`, error)) + const register = (event: string, acknowledge: (payload: T) => void | Promise, nativeShutdown: boolean) => listen(event, ({ payload }) => { + void flush(nativeShutdown ? payload.generation : undefined).then(() => acknowledge(payload)).catch((error) => log.error(`Failed to handle ${event}`, error)) }).then((unlisten) => { if (nativeDisposed) unlisten() else nativeUnlisteners.push(unlisten) @@ -193,8 +194,9 @@ export function useAppSessionCapture() { ({ generation }) => acknowledgeNativeClientStateRendererFlush(generation), true), register<{ generation: number }>("client-state:navigation-flush-requested", ({ generation }) => acknowledgeNativeClientStateNavigationFlush(generation), false), - listen("client-state:flush-cancelled", () => { - nativeShutdownStarted = false + listen<{ generation: number }>("client-state:flush-cancelled", ({ payload }) => { + if (nativeShutdownGeneration !== payload.generation) return + nativeShutdownGeneration = null schedule() }).then((unlisten) => { if (nativeDisposed) unlisten() diff --git a/packages/ui/src/stores/app-session-workspace-hydration.test.ts b/packages/ui/src/stores/app-session-workspace-hydration.test.ts new file mode 100644 index 00000000..d3772371 --- /dev/null +++ b/packages/ui/src/stores/app-session-workspace-hydration.test.ts @@ -0,0 +1,73 @@ +import assert from "node:assert/strict" +import { it } from "node:test" + +import { sdkManager } from "../lib/sdk-manager.ts" +import { addInstance, removeInstance } from "./instances.ts" +import { hydrateRestoredWorkspaceState } from "./app-session-workspace-hydration.ts" +import { messageStoreBus } from "./message-v2/bus.ts" +import { + clearInstanceDraftPrompts, + clearInstanceSessionSelection, + getSessionDraftPrompt, + setSessions, +} from "./session-state.ts" + +function deferred() { + let resolve!: (value: T) => void + const promise = new Promise((done) => { resolve = done }) + return { promise, resolve } +} + +function apiSession(id: string) { + return { + id, title: id, projectID: "project", location: { directory: "/work" }, cost: 0, + tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }, + time: { created: 1, updated: 1 }, + } +} + +it("restores the selected draft before inactive session hydration settles", async () => { + const instanceId = "selected-draft-first" + const inactive = deferred() + const controller = new AbortController() + const signals: AbortSignal[] = [] + const client = { session: { get: (input: { sessionID: string }, options?: { signal?: AbortSignal }) => { + if (options?.signal) signals.push(options.signal) + return input.sessionID === "active" ? Promise.resolve(apiSession("active")) : inactive.promise + } } } as any + ;(sdkManager as any).clients.set(`${instanceId}:/workspaces/${instanceId}/instance`, client) + addInstance({ id: instanceId, folder: "/work", port: 0, pid: 0, proxyPath: "", status: "ready", client }) + + try { + let settled = false + const hydration = hydrateRestoredWorkspaceState(instanceId, { + kind: "workspace", + folder: "/work", + activeParentSessionId: "active", + activeSessionId: "active", + drafts: { active: "active draft", inactive: "inactive draft" }, + attachments: {}, + scrollSnapshots: {}, + unseenIdleSince: {}, + generationRecovery: {}, + }, controller.signal, () => true).then((value) => { settled = true; return value }) + + await new Promise((resolve) => setImmediate(resolve)) + assert.equal(settled, false) + assert.equal(getSessionDraftPrompt(instanceId, "active"), "active draft") + assert.equal(getSessionDraftPrompt(instanceId, "inactive"), "") + + inactive.resolve(apiSession("inactive")) + assert.deepEqual(await hydration, new Set()) + assert.equal(getSessionDraftPrompt(instanceId, "inactive"), "inactive draft") + assert.equal(signals.length, 2) + assert.equal(signals.every((signal) => signal === controller.signal), true) + } finally { + if (messageStoreBus.getInstance(instanceId)) messageStoreBus.unregisterInstance(instanceId) + clearInstanceDraftPrompts(instanceId) + clearInstanceSessionSelection(instanceId) + setSessions((previous) => { const next = new Map(previous); next.delete(instanceId); return next }) + removeInstance(instanceId, { authoritative: false }) + sdkManager.destroyClientsForInstance(instanceId) + } +}) diff --git a/packages/ui/src/stores/app-session-workspace-hydration.ts b/packages/ui/src/stores/app-session-workspace-hydration.ts index 486c629c..780051be 100644 --- a/packages/ui/src/stores/app-session-workspace-hydration.ts +++ b/packages/ui/src/stores/app-session-workspace-hydration.ts @@ -42,6 +42,12 @@ export async function hydrateRestoredWorkspaceState( if (!hasAuthoritativeSessionSelection(instanceId)) { hydrateActiveSessionSelection(instanceId, selection?.parentSessionId ?? null, selection?.activeSessionId ?? null) } + hydrateWorkspacePromptState( + instanceId, + snapshot, + new Set(sessions.map(({ id }) => id)), + NO_SESSION_DRAFT_SESSION_ID, + ) await hydrateRestoredSessionChain(instanceId, getRestoredSessionIds([ Object.keys(snapshot.drafts), diff --git a/packages/ui/src/stores/client-state-attachments-codec.ts b/packages/ui/src/stores/client-state-attachments-codec.ts index 1253ddda..a8e9052b 100644 --- a/packages/ui/src/stores/client-state-attachments-codec.ts +++ b/packages/ui/src/stores/client-state-attachments-codec.ts @@ -1,5 +1,4 @@ import type { Attachment, AttachmentSource } from "../types/attachment" -import { removeAttachmentPromptTokens } from "../lib/attachment-mentions" export type RestorableAttachmentSource = | { type: "file"; path: string; mime: string; data?: string } @@ -21,37 +20,19 @@ export interface RestorableAttachment { mediaType: string source: RestorableAttachmentSource } -export interface AttachmentCodecBudget { - attachmentsRemaining: number - metadataCharactersRemaining: number - fileDataCharactersRemaining: number -} - -const MAX_SESSIONS = 24 -const MAX_PER_SESSION = 8 -const MAX_ATTACHMENTS = 64 -const MAX_METADATA = 24 * 1024 -const MAX_FILE_BYTES = 64 * 1024 -const MAX_FILE_CHARACTERS = 96 * 1024 const MAX_ID = 512 const MAX_DISPLAY = 1024 const MAX_PATH = 4096 const MAX_MIME = 256 -const MAX_TEXT = 24 * 1024 const isRecord = (value: unknown): value is Record => typeof value === "object" && value !== null && !Array.isArray(value) const isSafeKey = (value: string) => value !== "__proto__" && value !== "constructor" && value !== "prototype" -function takeString(value: unknown, max: number, budget: AttachmentCodecBudget, allowEmpty = false) { - if ( - typeof value !== "string" || - value.length > max || - value.length > budget.metadataCharactersRemaining - ) return +function takeString(value: unknown, max: number, allowEmpty = false) { + if (typeof value !== "string" || value.length > max) return if (!allowEmpty && value.trim().length === 0) return - budget.metadataCharactersRemaining -= value.length return value } @@ -79,18 +60,10 @@ const dataUrlPayload = (value: unknown): string | undefined => typeof value === ? value.match(/^data:[^;,]+;base64,([A-Za-z0-9+/]*={0,2})$/)?.[1] : undefined -function takeFileData(value: unknown, url: unknown, budget: AttachmentCodecBudget): string | undefined { +function takeFileData(value: unknown, url: unknown): string | undefined { const raw = dataUrlPayload(url) ?? value - if (raw instanceof Uint8Array && raw.byteLength > MAX_FILE_BYTES) return const data = raw instanceof Uint8Array ? bytesToBase64(raw) : raw - if ( - !validBase64(data) || - base64ToBytes(data).byteLength > MAX_FILE_BYTES || - data.length > MAX_FILE_CHARACTERS || - data.length > budget.fileDataCharactersRemaining - ) return - budget.fileDataCharactersRemaining -= data.length - return data + return validBase64(data) ? data : undefined } function takePosition(value: unknown): Position | undefined { @@ -102,77 +75,59 @@ function takePosition(value: unknown): Position | undefined { function normalizeSource( value: unknown, url: unknown, - budget: AttachmentCodecBudget, ): RestorableAttachmentSource | undefined { if (!isRecord(value)) return if (value.type === "text") { - const text = takeString(value.value, MAX_TEXT, budget, true) - return text === undefined ? undefined : { type: "text", value: text } + return typeof value.value === "string" ? { type: "text", value: value.value } : undefined } if (value.type === "agent") { - const name = takeString(value.name, MAX_DISPLAY, budget) + const name = takeString(value.name, MAX_DISPLAY) return name === undefined ? undefined : { type: "agent", name } } - const path = takeString(value.path, MAX_PATH, budget) + const path = takeString(value.path, MAX_PATH) if (value.type === "file") { - const mime = takeString(value.mime, MAX_MIME, budget) + const mime = takeString(value.mime, MAX_MIME) if (path === undefined || mime === undefined) return const rawData = dataUrlPayload(url) ?? value.data - const data = takeFileData(value.data, url, budget) - const pathBacked = typeof url === "string" && url.length > 0 && !url.startsWith("data:") - const validPathData = rawData instanceof Uint8Array || validBase64(rawData) - if (rawData !== undefined && data === undefined && (!pathBacked || !validPathData)) return + const data = takeFileData(value.data, url) + if (rawData !== undefined && data === undefined) return return data === undefined ? { type: "file", path, mime } : { type: "file", path, mime, data } } if (value.type !== "symbol" || !isRecord(value.range)) return - const name = takeString(value.name, MAX_DISPLAY, budget) + const name = takeString(value.name, MAX_DISPLAY) const start = takePosition(value.range.start) const end = takePosition(value.range.end) if (path === undefined || name === undefined || !Number.isSafeInteger(value.kind) || !start || !end) return return { type: "symbol", path, name, kind: Number(value.kind), range: { start, end } } } -function normalizeAttachment(value: unknown, budget: AttachmentCodecBudget): RestorableAttachment | undefined { - if (!isRecord(value) || budget.attachmentsRemaining <= 0) return +function normalizeAttachment(value: unknown): RestorableAttachment | undefined { + if (!isRecord(value)) return if (!["file", "text", "symbol", "agent"].includes(String(value.type))) return - const next = { ...budget } - const id = takeString(value.id, MAX_ID, next) - const display = takeString(value.display, MAX_DISPLAY, next) - const filename = takeString(value.filename, MAX_PATH, next) - const mediaType = takeString(value.mediaType, MAX_MIME, next) - const source = normalizeSource(value.source, value.url, next) + const id = takeString(value.id, MAX_ID) + const display = takeString(value.display, MAX_DISPLAY) + const filename = takeString(value.filename, MAX_PATH) + const mediaType = takeString(value.mediaType, MAX_MIME) + const source = normalizeSource(value.source, value.url) const rawUrl = typeof value.url === "string" && !value.url.startsWith("data:") ? value.url : "" - const url = takeString(rawUrl, MAX_PATH, next, true) + const url = takeString(rawUrl, MAX_PATH, true) if ( !id || !display || !filename || !mediaType || !source || source.type !== value.type || url === undefined ) return - next.attachmentsRemaining -= 1 - Object.assign(budget, next) return { id, type: value.type as Attachment["type"], display, url, filename, mediaType, source } } -export function createAttachmentCodecBudget(): AttachmentCodecBudget { - return { - attachmentsRemaining: MAX_ATTACHMENTS, - metadataCharactersRemaining: MAX_METADATA, - fileDataCharactersRemaining: MAX_FILE_CHARACTERS, - } -} - export function normalizeRestorableAttachmentRecord( value: unknown, drafts: Record, - budget: AttachmentCodecBudget, prioritySessionIds: readonly string[] = [], ): { attachments: Record; drafts: Record } | null { if (!isRecord(value)) return null const attachments: Record = Object.create(null) - const nextDrafts = { ...drafts } - let sessions = 0 const priority = [...new Set(prioritySessionIds)] const prioritySet = new Set(priority) const entries = [ @@ -186,16 +141,14 @@ export function normalizeRestorableAttachmentRecord( !isSafeKey(sessionId) || !sessionId || sessionId.length > MAX_ID || !Array.isArray(rawAttachments) ) continue - const persist = sessions++ < MAX_SESSIONS const normalized: RestorableAttachment[] = [] for (const raw of rawAttachments) { - const attachment = persist && normalized.length < MAX_PER_SESSION ? normalizeAttachment(raw, budget) : undefined + const attachment = normalizeAttachment(raw) if (attachment) normalized.push(attachment) - else if (nextDrafts[sessionId]) nextDrafts[sessionId] = removeAttachmentPromptTokens(nextDrafts[sessionId], raw) } if (normalized.length) attachments[sessionId] = normalized } - return { attachments, drafts: nextDrafts } + return { attachments, drafts: { ...drafts } } } export function serializeDraftAttachments( @@ -203,8 +156,12 @@ export function serializeDraftAttachments( attachments: Record, prioritySessionIds: readonly string[] = [], ) { - return normalizeRestorableAttachmentRecord(attachments, drafts, createAttachmentCodecBudget(), prioritySessionIds) - ?? { drafts: { ...drafts }, attachments: {} } + const normalized = normalizeRestorableAttachmentRecord(attachments, drafts, prioritySessionIds) + if (!normalized || Object.entries(attachments).some(([sessionId, values]) => + values.length !== (normalized.attachments[sessionId]?.length ?? 0))) { + throw new Error("Draft attachments could not be persisted safely") + } + return normalized } export function hydrateRestorableAttachment(value: RestorableAttachment): Attachment | null { diff --git a/packages/ui/src/stores/client-state-codec.test.ts b/packages/ui/src/stores/client-state-codec.test.ts index 7925ab30..9a19e071 100644 --- a/packages/ui/src/stores/client-state-codec.test.ts +++ b/packages/ui/src/stores/client-state-codec.test.ts @@ -144,10 +144,9 @@ describe("client state codec", () => { assert.ok(Buffer.byteLength(JSON.stringify(roundTripped), "utf8") < 1024 * 1024) }) - it("drops unsupported attachments but retains oversized path-backed files and exact mentions", () => { + it("drops unsupported attachments without altering drafts and retains large path-backed data", () => { const mention = "@./reports/exact report.txt" const oversizedData = new Uint8Array(64 * 1024 + 1) - oversizedData.subarray = () => { throw new Error("oversized data must not be Base64 encoded") } const large = file("./reports/exact report.txt", { id: "large", display: "@exact report.txt", url: "file:///work/reports/exact%20report.txt", filename: "exact report.txt", source: { type: "file", path: "./reports/exact report.txt", mime: "text/plain", data: oversizedData }, @@ -157,9 +156,10 @@ describe("client state codec", () => { drafts: { session1: `Review ${mention}; remove @unsupported` }, attachments: { session1: [large, unsupported] }, }) assert.equal(tab.attachments.session1?.length, 1) - assert.deepEqual(tab.attachments.session1?.[0]?.source, - { type: "file", path: "./reports/exact report.txt", mime: "text/plain" }) - assert.equal(tab.drafts.session1, `Review ${mention}; remove `) + const source = tab.attachments.session1?.[0]?.source + assert.equal(source?.type, "file") + assert.equal(source?.type === "file" ? source.data?.length : 0, 87384) + assert.equal(tab.drafts.session1, `Review ${mention}; remove @unsupported`) }) it("reserves structural and active-session identity before optional strings", () => { @@ -305,7 +305,7 @@ describe("client state codec", () => { assert.deepEqual(normalized, { tabs: [{ kind: "sidecar", sidecarId: "valid-late-tab" }], activeTabIndex: 0 }) }) - it("removes every picker mention for attachments beyond the per-session limit", () => { + it("retains attachments and picker mentions beyond the former per-session limit", () => { const cases = [ ["relative file", "./dir/f8", "f8", "@f8", ["@./dir/f8", "@f8"], ["@./dir/f80", "@f80"], "text/plain"], @@ -326,12 +326,12 @@ describe("client state codec", () => { attachments: { session: [...keep, file(path, { id: "drop", filename, display, mediaType: mime, source: { type: "file", path, mime } })] }, }) - assert.equal(tab.attachments.session?.length, 8, label) - assert.equal(tab.drafts.session, `remove then keep ${collisions.join(" and ")} and @other`, label) + assert.equal(tab.attachments.session?.length, 9, label) + assert.equal(tab.drafts.session, `remove ${tokens.join(" then ")} keep ${collisions.join(" and ")} and @other`, label) } }) - it("removes loose dropped placeholders without touching ordinary bracket text", () => { + it("keeps large attachment placeholders and never edits drafts for an unsupported attachment", () => { const input = [ "[Image #9]", "[ Image # 9 ]", "[iMaGe # 9]", "Image #9", "[Image #90]", "[Image #9 notes]", @@ -345,12 +345,11 @@ describe("client state codec", () => { attachment({ type: "archive" }, { id: "ordinary", type: "archive", display: "[ordinary bracket text]" }), ] const tab = normalizeWorkspace({ drafts: { session1: input }, attachments: { session1: dropped } }) - assert.deepEqual({ ...tab.attachments }, {}) - assert.equal(tab.drafts.session1, - "|||Image #9|[Image #90]|[Image #9 notes]||||pasted #4|[pasted #40]|[pasted notes]|[ordinary bracket text]") + assert.deepEqual(tab.attachments.session1?.map(({ id }) => id), ["image", "paste"]) + assert.equal(tab.drafts.session1, input) }) - it("keeps a maximal normalized attachment snapshot below the native 1 MiB limit", () => { + it("retains attachment data beyond the former global budget", () => { const normalized = session(...Array.from({ length: 32 }, (_, index) => workspace({ folder: `/work/${index}`, drafts: { session: `[Image #${index + 1}]` }, attachments: { session: [file(`file-${index}.bin`, { id: `file-${index}`, display: `[Image #${index + 1}]`, @@ -358,6 +357,6 @@ describe("client state codec", () => { mime: "application/octet-stream", data: new Uint8Array(64 * 1024) } })] }, }))) assert.ok(normalized) - assert.ok(Buffer.byteLength(JSON.stringify(normalized), "utf8") < 1024 * 1024) + assert.equal(normalized.tabs.filter((tab) => tab.kind === "workspace" && tab.attachments.session?.length).length, 32) }) }) diff --git a/packages/ui/src/stores/client-state-codec.ts b/packages/ui/src/stores/client-state-codec.ts index a89a5818..9e784e07 100644 --- a/packages/ui/src/stores/client-state-codec.ts +++ b/packages/ui/src/stores/client-state-codec.ts @@ -1,6 +1,5 @@ import type { ScrollSnapshot } from "./message-v2/types" -import { createAttachmentCodecBudget, normalizeRestorableAttachmentRecord, - type AttachmentCodecBudget, type RestorableAttachment } from "./client-state-attachments-codec" +import { normalizeRestorableAttachmentRecord, type RestorableAttachment } from "./client-state-attachments-codec" import type { PersistedGenerationRecovery } from "./session-generation-recovery" export interface RestorableWorkspaceTabState { @@ -26,10 +25,10 @@ const MAX_LAYOUT_VALUE = 4096, MAX_DRAFT = 32 * 1024, MAX_ANCHOR_KEY = 1024 const MAX_STRINGS = 96 * 1024, MAX_SCROLLS = 256 const NO_SESSION_DRAFT_SESSION_ID = "__no_session_draft__" -interface StringBudget { remaining: number; scrollSnapshotsRemaining: number; attachments: AttachmentCodecBudget } +interface StringBudget { remaining: number; scrollSnapshotsRemaining: number } function createBudget(): StringBudget { - return { remaining: MAX_STRINGS, scrollSnapshotsRemaining: MAX_SCROLLS, attachments: createAttachmentCodecBudget() } + return { remaining: MAX_STRINGS, scrollSnapshotsRemaining: MAX_SCROLLS } } const isRecord = (value: unknown): value is Record => @@ -151,7 +150,7 @@ function normalizeWorkspaceTab( const remainingAttachments = Object.fromEntries(Object.entries(value.attachments ?? {}) .filter(([id]) => !identity.prioritySessionIds.includes(id))) const attachmentResult = normalizeRestorableAttachmentRecord( - remainingAttachments, drafts, budget.attachments, + remainingAttachments, drafts, ) if (!attachmentResult) return null @@ -227,7 +226,6 @@ function reservePriorityDrafts(identity: WorkspaceIdentity, budget: StringBudget const result = normalizeRestorableAttachmentRecord( priorityAttachments, identity.priorityDrafts, - budget.attachments, identity.prioritySessionIds, ) if (!result) return diff --git a/packages/ui/src/stores/client-state-partitions.test.ts b/packages/ui/src/stores/client-state-partitions.test.ts index e1f3e1e3..6c0d5dd8 100644 --- a/packages/ui/src/stores/client-state-partitions.test.ts +++ b/packages/ui/src/stores/client-state-partitions.test.ts @@ -49,6 +49,39 @@ const sidecarSnapshot: ClientSnapshotV1 = { const loader = (encoded: Awaited>) => async (key: string) => encoded.partitions[key] ?? null +async function blobReferenceSnapshot(dataPartitions: string[], blobs: Record) { + const document = canonicalJson({ + format: 2, + draft: "keep only if the attachment is complete", + attachments: [{ + id: "blob", type: "file", display: "[Image #1]", url: "", filename: "blob.bin", + mediaType: "application/octet-stream", + source: { type: "file", path: "blob.bin", mime: "application/octet-stream", dataPartitions }, + }], + }) + const documentPartition = await sha256(document) + const leafKeys = [...new Set([documentPartition, ...Object.keys(blobs)])].sort() + const workspaceDocument = canonicalJson({ + format: 2, + activeSessionId: "selected", + sessions: { selected: { documentPartition, partitionKeys: leafKeys } }, + }) + const workspacePartition = await sha256(workspaceDocument) + const manifest = canonicalJson({ + format: 2, + session: { tabs: [{ kind: "workspace", folder: "/work", workspacePartition }], activeTabIndex: 0 }, + }) + const sessionPartition = await sha256(manifest) + const partitions = { ...blobs, [documentPartition]: document, [workspacePartition]: workspaceDocument, + [sessionPartition]: manifest } + const partitionKeys = Object.keys(partitions).sort() + return { + root: { version: 2 as const, revision: 1, savedAt: 2, layout: {}, sessionPartition, partitionKeys }, + partitions, + partitionKeys, + } +} + it("canonicalizes recursively without reordering arrays and hashes the exact UTF-8 text", async () => { const value = { z: [{ y: 2, x: 1 }, "first"], a: { d: 4, c: 3 } } const reordered = { a: { c: 3, d: 4 }, z: [{ x: 1, y: 2 }, "first"] } @@ -102,20 +135,115 @@ it("uses the native graph cap without truncating the encoded key list", async () assert.equal(canCommitClientSnapshotV2({ ...encoded, partitionKeys: Array(4097).fill("key") }), false) }) -it("rejects a missing or corrupt inactive session document", async () => { - const encoded = await encodeClientSnapshotV2(graphSnapshot()) +it("drops only a missing or corrupt inactive session leaf", async () => { + const snapshot = graphSnapshot() + const inactive = snapshot.session!.tabs[1] as RestorableWorkspaceTabState + inactive.drafts.stale = "optional stale draft" + inactive.attachments.stale = [attachment("stale-attachment")] + const encoded = await encodeClientSnapshotV2(snapshot) const manifest = JSON.parse(encoded.partitions[encoded.root.sessionPartition]!) const inactiveWorkspace = JSON.parse(encoded.partitions[manifest.session.tabs[1].workspacePartition]!) - const inactiveDocument = inactiveWorkspace.sessions["session-1"] + const inactiveDocument = inactiveWorkspace.sessions.stale.documentPartition for (const failure of ["missing", "corrupt"] as const) { const load = async (key: string) => key === inactiveDocument ? failure === "missing" ? null : `${encoded.partitions[key]} ` : encoded.partitions[key] ?? null - assert.equal(await decodeClientSnapshotV2(encoded.root, 1, load), null, failure) + const decoded = await decodeClientSnapshotV2(encoded.root, 1, load) + const tab = decoded?.session?.tabs[1] + assert.equal(decoded?.session?.tabs.length, 2, failure) + assert.equal(tab?.kind === "workspace" ? tab.drafts["session-1"] : undefined, "second draft", failure) + assert.equal(tab?.kind === "workspace" ? tab.drafts.stale : undefined, undefined, failure) + assert.equal(tab?.kind === "workspace" ? tab.attachments.stale : undefined, undefined, failure) + assert.equal(tab?.kind === "workspace" ? tab.activeSessionId : undefined, "session-1", failure) } }) +it("clears selection and all optional state when the selected leaf is corrupt", async () => { + const encoded = await encodeClientSnapshotV2(graphSnapshot()) + const manifest = JSON.parse(encoded.partitions[encoded.root.sessionPartition]!) + const workspaceDocument = JSON.parse(encoded.partitions[manifest.session.tabs[0].workspacePartition]!) + const selectedDocument = workspaceDocument.sessions["session-1"].documentPartition + const decoded = await decodeClientSnapshotV2(encoded.root, 1, async (key) => + key === selectedDocument ? null : encoded.partitions[key] ?? null) + const tab = decoded?.session?.tabs[0] + + assert.equal(tab?.kind, "workspace") + if (tab?.kind !== "workspace") return + assert.equal(tab.activeSessionId, undefined) + assert.equal(tab.activeParentSessionId, "parent") + assert.deepEqual(tab.expandedSessionIds, ["parent"]) + assert.equal(tab.drafts["session-1"], undefined) + assert.equal(tab.attachments["session-1"], undefined) + assert.equal(tab.scrollSnapshots["session-1"], undefined) +}) + +it("chunks and round trips attachments larger than the former 8 MiB commit limit", async () => { + const data = new Uint8Array(9 * 1024 * 1024 + 17) + const tab = workspace(0, "Review [Image #1]") + tab.attachments["session-1"] = [{ + ...attachment("large-image"), type: "file", display: "[Image #1]", filename: "large.bin", + mediaType: "application/octet-stream", + source: { type: "file", path: "large.bin", mime: "application/octet-stream", + data: Buffer.from(data).toString("base64") }, + }] + const snapshot = { ...graphSnapshot(), session: { tabs: [tab], activeTabIndex: 0 } } + const encoded = await encodeClientSnapshotV2(snapshot) + const decoded = await decodeClientSnapshotV2(encoded.root, 1, loader(encoded)) + const restored = decoded?.session?.tabs[0] + const source = restored?.kind === "workspace" ? restored.attachments["session-1"]?.[0]?.source : undefined + + assert.ok(Object.values(encoded.partitions).every((value) => Buffer.byteLength(value, "utf8") < 1024 * 1024)) + assert.ok(Object.values(encoded.partitions).reduce((total, value) => total + Buffer.byteLength(value, "utf8"), 0) > 8 * 1024 * 1024) + assert.ok(encoded.partitionKeys.length > 4, "attachment was split into content partitions") + assert.equal(source?.type, "file") + assert.deepEqual(source?.type === "file" ? Buffer.from(source.data ?? "", "base64") : null, Buffer.from(data)) +}) + +it("rejects duplicate and excessive blob chunk references without loading amplification chunks", async () => { + const first = canonicalJson({ format: 2, index: 0, data: "YQ==" }) + const firstKey = await sha256(first) + const duplicate = await blobReferenceSnapshot([firstKey, firstKey], { [firstKey]: first }) + const duplicateDecoded = await decodeClientSnapshotV2(duplicate.root, 1, async (key) => duplicate.partitions[key] ?? null) + const duplicateTab = duplicateDecoded?.session?.tabs[0] + assert.equal(duplicateTab?.kind === "workspace" ? duplicateTab.activeSessionId : "missing", undefined) + assert.deepEqual(duplicateTab?.kind === "workspace" ? { ...duplicateTab.drafts } : null, {}) + + const blobs: Record = {} + const keys: string[] = [] + for (let index = 0; index < 257; index += 1) { + const chunk = canonicalJson({ format: 2, index, data: "" }) + const key = await sha256(chunk) + blobs[key] = chunk + keys.push(key) + } + const excessive = await blobReferenceSnapshot(keys, blobs) + const blobKeys = new Set(keys) + let blobLoads = 0 + const excessiveDecoded = await decodeClientSnapshotV2(excessive.root, 1, async (key) => { + if (blobKeys.has(key)) blobLoads += 1 + return excessive.partitions[key] ?? null + }) + const excessiveTab = excessiveDecoded?.session?.tabs[0] + assert.equal(excessiveTab?.kind === "workspace" ? excessiveTab.activeSessionId : "missing", undefined) + assert.equal(blobLoads, 0) +}) + +it("round trips many attachment sessions without count-based truncation", async () => { + const tab = workspace(0, "draft") + tab.attachments = Object.fromEntries(Array.from({ length: 40 }, (_, sessionIndex) => [ + `many-${sessionIndex}`, + Array.from({ length: 12 }, (_, attachmentIndex) => attachment(`item-${sessionIndex}-${attachmentIndex}`)), + ])) + const snapshot = { ...graphSnapshot(), session: { tabs: [tab], activeTabIndex: 0 } } + const encoded = await encodeClientSnapshotV2(snapshot) + const decoded = await decodeClientSnapshotV2(encoded.root, 1, loader(encoded)) + const restored = decoded?.session?.tabs[0] + + assert.equal(restored?.kind === "workspace" ? Object.keys(restored.attachments).length : 0, 40) + assert.ok(restored?.kind === "workspace" && Object.values(restored.attachments).every((values) => values.length === 12)) +}) + it("rejects missing, reordered, and disconnected root graph keys", async () => { const encoded = await encodeClientSnapshotV2(graphSnapshot()) const removable = encoded.root.partitionKeys.find((key) => key !== encoded.root.sessionPartition)! diff --git a/packages/ui/src/stores/client-state-partitions.ts b/packages/ui/src/stores/client-state-partitions.ts index 8ca64413..908cac25 100644 --- a/packages/ui/src/stores/client-state-partitions.ts +++ b/packages/ui/src/stores/client-state-partitions.ts @@ -1,10 +1,14 @@ import { decodeClientSnapshot, normalizeRestorableSession } from "./client-state-codec" import type { ClientSnapshotV1, RestorableSessionState, RestorableWorkspaceTabState } from "./client-state-codec" +import type { RestorableAttachment } from "./client-state-attachments-codec" const MAX_PARTITION_BYTES = 1024 * 1024 const MAX_ROOT_BYTES = 1024 * 1024 const MAX_NATIVE_PARTITIONS = 4096 const MAX_SESSION_ID = 512 +const BLOB_CHUNK_CHARACTERS = 768 * 1024 +const MAX_BLOB_CHUNKS = 256 +const MAX_BLOB_BASE64_CHARACTERS = BLOB_CHUNK_CHARACTERS * MAX_BLOB_CHUNKS const PARTITION_KEY = /^[0-9a-f]{64}$/ const ROOT_KEYS = ["layout", "partitionKeys", "revision", "savedAt", "sessionPartition", "version"] const UNSAFE_KEYS = new Set(["__proto__", "constructor", "prototype"]) @@ -26,9 +30,9 @@ export interface EncodedClientSnapshotV2 { } interface SessionDocument { - format: 1 + format: 2 draft?: string - attachments?: RestorableWorkspaceTabState["attachments"][string] + attachments?: unknown[] scrollSnapshot?: RestorableWorkspaceTabState["scrollSnapshots"][string] unseenIdleSince?: number generationRecovery?: RestorableWorkspaceTabState["generationRecovery"][string] @@ -41,6 +45,8 @@ const isSafePersistedSessionId = (value: unknown): value is string => typeof value === "string" && value.length > 0 && value.length <= MAX_SESSION_ID && value.trim().length > 0 && !UNSAFE_KEYS.has(value) +type AddPartition = (value: unknown) => Promise + function isPartitionKeyArray(value: unknown): value is string[] { return Array.isArray(value) && value.every((key, index) => typeof key === "string" && PARTITION_KEY.test(key) && (index === 0 || value[index - 1]! < key)) @@ -76,10 +82,52 @@ export async function sha256(value: string): Promise { return Array.from(new Uint8Array(digest), (byte) => byte.toString(16).padStart(2, "0")).join("") } +function bytesToBase64(bytes: Uint8Array): string { + let binary = "" + for (let index = 0; index < bytes.length; index += 0x8000) { + binary += String.fromCharCode(...bytes.subarray(index, index + 0x8000)) + } + return btoa(binary) +} + +const base64ToBytes = (value: string) => Uint8Array.from(atob(value), (character) => character.charCodeAt(0)) +const validBase64 = (value: unknown): value is string => typeof value === "string" + && value.length % 4 === 0 && /^[A-Za-z0-9+/]*={0,2}$/.test(value) + +async function addBlob(data: string, addPartition: AddPartition): Promise { + if (!validBase64(data) || data.length > MAX_BLOB_BASE64_CHARACTERS) { + throw new Error("Client state attachment exceeds the 192 MiB encoded blob limit") + } + const keys: string[] = [] + for (let offset = 0, index = 0; offset < data.length || offset === 0; offset += BLOB_CHUNK_CHARACTERS, index += 1) { + keys.push(await addPartition({ format: 2, index, data: data.slice(offset, offset + BLOB_CHUNK_CHARACTERS) })) + } + return keys +} + +async function encodeAttachments(attachments: RestorableAttachment[], addPartition: AddPartition): Promise { + return Promise.all(attachments.map(async (attachment) => { + let source: Record + if (attachment.source.type === "file") { + source = { type: "file", path: attachment.source.path, mime: attachment.source.mime } + if (attachment.source.data !== undefined) source.dataPartitions = await addBlob(attachment.source.data, addPartition) + } else if (attachment.source.type === "text") { + source = { + type: "text", + valuePartitions: await addBlob(bytesToBase64(encoder.encode(attachment.source.value)), addPartition), + } + } else source = attachment.source + return { ...attachment, source } + })) +} + async function encodeSessionGraph(session: RestorableSessionState | null) { const generated: Record = Object.create(null) const addPartition = async (value: unknown) => { const partition = canonicalJson(value) + if (encoder.encode(partition).byteLength > MAX_PARTITION_BYTES) { + throw new Error("Client state partition exceeds the native size limit") + } const key = await sha256(partition) generated[key] = partition return key @@ -96,18 +144,29 @@ async function encodeSessionGraph(session: RestorableSessionState | null) { ...Object.keys(tab.unseenIdleSince), ...Object.keys(tab.generationRecovery), ]) - const sessions: Record = Object.create(null) + const sessions: Record = Object.create(null) for (const sessionId of [...sessionIds].sort()) { - const document: SessionDocument = { format: 1 } + const document: SessionDocument = { format: 2 } if (hasOwn(tab.drafts, sessionId)) document.draft = tab.drafts[sessionId] - if (hasOwn(tab.attachments, sessionId)) document.attachments = tab.attachments[sessionId] + if (hasOwn(tab.attachments, sessionId)) { + document.attachments = await encodeAttachments(tab.attachments[sessionId]!, addPartition) + } if (hasOwn(tab.scrollSnapshots, sessionId)) document.scrollSnapshot = tab.scrollSnapshots[sessionId] if (hasOwn(tab.unseenIdleSince, sessionId)) document.unseenIdleSince = tab.unseenIdleSince[sessionId] if (hasOwn(tab.generationRecovery, sessionId)) document.generationRecovery = tab.generationRecovery[sessionId] - sessions[sessionId] = await addPartition(document) + const documentPartition = await addPartition(document) + const attachmentPartitions = (document.attachments ?? []).flatMap((attachment) => { + if (!isRecord(attachment) || !isRecord(attachment.source)) return [] + const keys = attachment.source.dataPartitions ?? attachment.source.valuePartitions + return Array.isArray(keys) ? keys.filter((key): key is string => typeof key === "string") : [] + }) + sessions[sessionId] = { + documentPartition, + partitionKeys: [...new Set([documentPartition, ...attachmentPartitions])].sort(), + } } - const workspace: Record = { format: 1, sessions } + const workspace: Record = { format: 2, sessions } if (tab.activeParentSessionId !== undefined) workspace.activeParentSessionId = tab.activeParentSessionId if (tab.activeSessionId !== undefined) workspace.activeSessionId = tab.activeSessionId if (tab.expandedSessionIds !== undefined) workspace.expandedSessionIds = tab.expandedSessionIds @@ -214,7 +273,74 @@ async function decodeGraph( return parsed && isRecord(parsed.value) && canonicalJson(parsed.value) === partition ? parsed.value : null } + const sameKeys = (left: Iterable, right: readonly string[]) => { + const sorted = [...new Set(left)].sort() + return sorted.length === right.length && sorted.every((key, index) => key === right[index]) + } + const loadBlob = async (keys: unknown, allowed: ReadonlySet, referenced: Set) => { + if (!Array.isArray(keys) || !keys.length || keys.length > MAX_BLOB_CHUNKS + || !keys.every((key) => typeof key === "string" && PARTITION_KEY.test(key)) + || new Set(keys).size !== keys.length) return null + const chunks: string[] = [] + let length = 0 + for (const [index, key] of (keys as string[]).entries()) { + if (!allowed.has(key)) return null + referenced.add(key) + const chunk = await loadCanonical(key) + if (!chunk || !validBase64(chunk.data)) return null + if (chunk.format === 1) { + if (!hasExactKeys(chunk, ["format", "data"])) return null + } else if (chunk.format === 2) { + if (!hasExactKeys(chunk, ["format", "index", "data"]) || chunk.index !== index) return null + } else return null + if (chunk.data.length > MAX_BLOB_BASE64_CHARACTERS - length) return null + length += chunk.data.length + chunks.push(chunk.data) + } + const data = chunks.join("") + return validBase64(data) ? data : null + } + const decodeAttachments = async ( + value: unknown, + allowed: ReadonlySet, + referenced: Set, + ): Promise => { + if (!Array.isArray(value)) return null + const result: unknown[] = [] + for (const attachment of value) { + if (!isRecord(attachment) || !hasExactKeys(attachment, + ["id", "type", "display", "url", "filename", "mediaType", "source"]) || !isRecord(attachment.source)) return null + const source = attachment.source + if (source.type === "file") { + if (!hasExactKeys(source, ["type", "path", "mime"], ["dataPartitions"])) return null + if (hasOwn(source, "dataPartitions")) { + const data = await loadBlob(source.dataPartitions, allowed, referenced) + if (data === null) return null + result.push({ ...attachment, source: { type: source.type, path: source.path, mime: source.mime, data } }) + } else result.push(attachment) + continue + } + if (source.type === "text") { + if (!hasExactKeys(source, ["type", "valuePartitions"])) return null + const data = await loadBlob(source.valuePartitions, allowed, referenced) + if (data === null) return null + try { + const value = new TextDecoder("utf-8", { fatal: true }).decode(base64ToBytes(data)) + result.push({ ...attachment, source: { type: source.type, value } }) + } catch { + return null + } + continue + } + result.push(attachment) + } + return result + } + const tabs: unknown[] = [] + const declaredGraph = new Set([rootKey]) + let degraded = false + let legacyGraph = false for (const shell of manifest.session.tabs) { if (!isRecord(shell)) return if (shell.kind === "sidecar") { @@ -224,11 +350,14 @@ async function decodeGraph( } if (shell.kind !== "workspace" || !hasExactKeys(shell, ["kind", "folder", "workspacePartition"], ["occurrence", "projectName", "binaryPath"])) return + if (typeof shell.workspacePartition !== "string" || !PARTITION_KEY.test(shell.workspacePartition)) return + declaredGraph.add(shell.workspacePartition) const workspace = await loadCanonical(shell.workspacePartition) - if (!workspace || workspace.format !== 1 + if (!workspace || (workspace.format !== 1 && workspace.format !== 2) || !hasExactKeys(workspace, ["format", "sessions"], ["activeParentSessionId", "activeSessionId", "expandedSessionIds"]) || !isRecord(workspace.sessions)) return + legacyGraph ||= workspace.format === 1 if (!Object.keys(workspace.sessions).every(isSafePersistedSessionId) || (hasOwn(workspace, "activeParentSessionId") && !isSafePersistedSessionId(workspace.activeParentSessionId)) || (hasOwn(workspace, "activeSessionId") && !isSafePersistedSessionId(workspace.activeSessionId)) @@ -250,26 +379,87 @@ async function decodeGraph( for (const key of ["activeParentSessionId", "activeSessionId", "expandedSessionIds"] as const) { if (hasOwn(workspace, key)) tab[key] = workspace[key] } - for (const [sessionId, documentKey] of Object.entries(workspace.sessions)) { + const entries: Array<{ sessionId: string; documentKey: string; partitionKeys: string[] }> = [] + for (const [sessionId, rawEntry] of Object.entries(workspace.sessions)) { + if (workspace.format === 1) { + if (typeof rawEntry !== "string" || !PARTITION_KEY.test(rawEntry)) return + entries.push({ sessionId, documentKey: rawEntry, partitionKeys: [rawEntry] }) + declaredGraph.add(rawEntry) + continue + } + if (!isRecord(rawEntry) || !hasExactKeys(rawEntry, ["documentPartition", "partitionKeys"]) + || typeof rawEntry.documentPartition !== "string" || !PARTITION_KEY.test(rawEntry.documentPartition) + || !isPartitionKeyArray(rawEntry.partitionKeys) + || !rawEntry.partitionKeys.includes(rawEntry.documentPartition)) return + entries.push({ sessionId, documentKey: rawEntry.documentPartition, partitionKeys: rawEntry.partitionKeys }) + rawEntry.partitionKeys.forEach((key) => declaredGraph.add(key)) + } + const activeId = typeof workspace.activeSessionId === "string" ? workspace.activeSessionId : undefined + entries.sort((left, right) => left.sessionId === activeId ? -1 : right.sessionId === activeId ? 1 + : left.sessionId.localeCompare(right.sessionId)) + const dropped = new Set() + for (const { sessionId, documentKey, partitionKeys } of entries) { const document = await loadCanonical(documentKey) - if (!document || document.format !== 1 - || !hasExactKeys(document, ["format"], + const allowed = new Set(partitionKeys) + const referenced = new Set([documentKey]) + const validDocument = document && (document.format === 1 || document.format === 2) + && hasExactKeys(document, ["format"], ["draft", "attachments", "scrollSnapshot", "unseenIdleSince", "generationRecovery"]) - || Object.keys(document).length === 1) return + && Object.keys(document).length > 1 + if (!validDocument) { + degraded = true + dropped.add(sessionId) + continue + } + const decodedAttachments = !hasOwn(document, "attachments") ? undefined : document.format === 1 + ? document.attachments + : await decodeAttachments(document.attachments, allowed, referenced) + if ((hasOwn(document, "attachments") && decodedAttachments === null) + || !sameKeys(referenced, partitionKeys)) { + degraded = true + dropped.add(sessionId) + continue + } + const leaf: RestorableWorkspaceTabState = { + kind: "workspace", folder: "/", + drafts: Object.create(null), attachments: Object.create(null), scrollSnapshots: Object.create(null), + unseenIdleSince: Object.create(null), generationRecovery: Object.create(null), + } const fields = [ ["draft", "drafts"], - ["attachments", "attachments"], ["scrollSnapshot", "scrollSnapshots"], ["unseenIdleSince", "unseenIdleSince"], ["generationRecovery", "generationRecovery"], ] as const for (const [documentField, stateField] of fields) { - if (hasOwn(document, documentField)) (tab[stateField] as Record)[sessionId] = document[documentField] + if (hasOwn(document, documentField)) (leaf[stateField] as Record)[sessionId] = document[documentField] } + if (decodedAttachments !== undefined) leaf.attachments[sessionId] = decodedAttachments as RestorableAttachment[] + const normalizedLeaf = normalizeRestorableSession({ tabs: [leaf], activeTabIndex: 0 })?.tabs[0] + if (!normalizedLeaf || !canonicalEquals(leaf, normalizedLeaf)) { + degraded = true + dropped.add(sessionId) + continue + } + for (const [, stateField] of fields) { + if (hasOwn(leaf[stateField], sessionId)) { + (tab[stateField] as Record)[sessionId] = leaf[stateField][sessionId] + } + } + if (hasOwn(leaf.attachments, sessionId)) { + (tab.attachments as Record)[sessionId] = leaf.attachments[sessionId] + } + } + if (dropped.has(String(tab.activeSessionId))) delete tab.activeSessionId + if (dropped.has(String(tab.activeParentSessionId))) delete tab.activeParentSessionId + if (Array.isArray(tab.expandedSessionIds)) { + tab.expandedSessionIds = tab.expandedSessionIds.filter((sessionId) => !dropped.has(String(sessionId))) } tabs.push(tab) } + if (!sameKeys(declaredGraph, persistedKeys)) return + const reconstructed: unknown = { tabs, activeTabIndex: manifest.session.activeTabIndex, @@ -277,6 +467,7 @@ async function decodeGraph( } const normalized = normalizeRestorableSession(reconstructed) if (!normalized || !canonicalEquals(reconstructed, normalized)) return + if (degraded || legacyGraph) return normalized const reencoded = await encodeSessionGraph(normalized) return graphMatches(reencoded, rootKey, persistedKeys, loaded) ? normalized : undefined } diff --git a/packages/ui/src/stores/client-state.test.ts b/packages/ui/src/stores/client-state.test.ts index 2b8084a6..5638af72 100644 --- a/packages/ui/src/stores/client-state.test.ts +++ b/packages/ui/src/stores/client-state.test.ts @@ -266,6 +266,35 @@ describe("secondary hosts", () => { }) describe("partitioned client state", () => { + it("restores a degraded graph and permits a repairing write", async () => { + const persisted = { + version: 1 as const, revision: 3, savedAt: 4, layout: { [layoutKey]: "390" }, + session: { activeTabIndex: 0, tabs: [{ + kind: "workspace" as const, folder: "/work", activeSessionId: "active", + drafts: { active: "keep", stale: "drop" }, attachments: {}, scrollSnapshots: {}, + unseenIdleSince: {}, generationRecovery: {}, + }] }, + } + const encoded = await encodeClientSnapshotV2(persisted) + const manifest = JSON.parse(encoded.partitions[encoded.root.sessionPartition]!) + const workspace = JSON.parse(encoded.partitions[manifest.session.tabs[0].workspacePartition]!) + const staleDocument = workspace.sessions.stale.documentPartition + const commits: any[] = [] + const state = await boot({ + loadClientState: async () => loadResult(encoded.root, true, 1), + loadClientStatePartition: async (_token, key) => key === staleDocument ? null : encoded.partitions[key] ?? null, + commitClientStatePartitions: async (_token, value) => { commits.push(value); return true }, + }) + + const restored = state.loadedRestorableSession()?.tabs[0] + assert.equal(restored?.kind === "workspace" ? restored.drafts.active : undefined, "keep") + assert.equal(restored?.kind === "workspace" ? restored.drafts.stale : undefined, undefined) + assert.equal(state.readClientLayoutValue(layoutKey), "390") + state.updateRestorableSession(session("repaired")) + await state.flushClientState() + assert.equal(commits.length, 1) + }) + it("commits one atomic graph without a monolithic save or automatic load rewrite", async () => { const encoded = await encodeClientSnapshotV2(snapshot("restored", { [layoutKey]: "380" })) const commits: any[] = []; let monolithicSaves = 0 @@ -308,6 +337,33 @@ describe("partitioned client state", () => { assert.equal(saved[0].session.tabs[0].sidecarId, "fallback") }) + it("keeps an oversized V1 draft attachment dirty instead of truncating or overwriting it", async () => { + let nativeSaves = 0 + const state = await boot({ + loadClientState: async () => loadResult(null), + saveClientState: async () => { nativeSaves += 1; return true }, + }) + const data = Buffer.alloc(1024 * 1024, 7).toString("base64") + state.updateRestorableSession({ tabs: [{ + kind: "workspace", folder: "/work", activeSessionId: "active", + drafts: { active: "Review [Image #1]" }, + attachments: { active: [{ + id: "large", type: "file", display: "[Image #1]", url: "", filename: "large.bin", + mediaType: "application/octet-stream", + source: { type: "file", path: "large.bin", mime: "application/octet-stream", data }, + }] }, + scrollSnapshots: {}, unseenIdleSince: {}, generationRecovery: {}, + }], activeTabIndex: 0 }) + + await assert.rejects(state.flushClientState(), /V1 1 MiB limit/) + await assert.rejects(state.flushClientState(), /V1 1 MiB limit/) + const restored = state.loadedRestorableSession()?.tabs[0] + const source = restored?.kind === "workspace" ? restored.attachments.active?.[0]?.source : undefined + assert.equal(nativeSaves, 0) + assert.equal(restored?.kind === "workspace" ? restored.drafts.active : undefined, "Review [Image #1]") + assert.equal(source?.type === "file" ? source.data : undefined, data) + }) + it("migrates a loaded V1 snapshot on the next real partition-capable save", async () => { const commits: any[] = []; let saves = 0 const state = await boot({ diff --git a/packages/ui/src/stores/client-state.ts b/packages/ui/src/stores/client-state.ts index 06c7f17e..05014bae 100644 --- a/packages/ui/src/stores/client-state.ts +++ b/packages/ui/src/stores/client-state.ts @@ -9,6 +9,7 @@ export type { ClientSnapshotV1, RestorableSessionState, RestorableSidecarTabStat export type { ClientSnapshotV2 } const SAVE_DEBOUNCE_MS = 250 const FLUSH_MAX_ATTEMPTS = 3 +const MAX_V1_SNAPSHOT_BYTES = 1024 * 1024 const MAX_LAYOUT_ENTRIES = 64 const MAX_LAYOUT_KEY_LENGTH = 256 const MAX_LAYOUT_VALUE_LENGTH = 4096 @@ -98,8 +99,13 @@ function enqueuePendingSave(): Promise { const saveAttempt = writeQueue.then(async () => { try { const partitioned = partitionProtocolVersion === 1 ? await encodeClientSnapshotV2(normalizedSnapshot) : null - // ponytail: bounded V1 avoids truncation at the native cap; add a denser graph if V1 stops fitting. - const accepted = partitioned && canCommitClientSnapshotV2(partitioned) + if (partitioned && !canCommitClientSnapshotV2(partitioned)) { + throw new Error("Client snapshot exceeds the native partition count limit") + } + if (!partitioned && new TextEncoder().encode(JSON.stringify(normalizedSnapshot)).byteLength > MAX_V1_SNAPSHOT_BYTES) { + throw new Error("Client snapshot exceeds the V1 1 MiB limit and partition persistence is unavailable") + } + const accepted = partitioned ? await commitNativeClientStatePartitions({ protocolVersion: 1, snapshot: partitioned.root, diff --git a/packages/ui/src/stores/connection-resync-gate.test.ts b/packages/ui/src/stores/connection-resync-gate.test.ts index aaa5fef1..715967ce 100644 --- a/packages/ui/src/stores/connection-resync-gate.test.ts +++ b/packages/ui/src/stores/connection-resync-gate.test.ts @@ -20,4 +20,9 @@ describe("connection resync gate", () => { gate.clear("instance") assert.equal(gate.observe("instance", "connected"), false) }) + + it("does not claim browser transport reconnect authority", () => { + const gate = new ConnectionResyncGate() + assert.equal("observeTransport" in gate, false) + }) }) diff --git a/packages/ui/src/stores/forms.test.ts b/packages/ui/src/stores/forms.test.ts index db0c29c6..c4fb1336 100644 --- a/packages/ui/src/stores/forms.test.ts +++ b/packages/ui/src/stores/forms.test.ts @@ -1,6 +1,18 @@ import assert from "node:assert/strict" import { describe, it } from "node:test" -import { activeInterruption, addPendingForm, removePendingForm } from "./instances.ts" +import { + activeInterruption, + addInstance, + addPendingForm, + removeInstance, + removePendingForm, + sendFormCancel, + sendFormReply, + syncPendingRequests, +} from "./instances.ts" +import { formRequestOptions, getFormQueue } from "./forms.ts" +import { getRootClient } from "./opencode-client.ts" +import { sdkManager } from "../lib/sdk-manager.ts" import { sessions, setSessions } from "./session-state.ts" const form = { @@ -32,4 +44,101 @@ describe("form interruption lifecycle", () => { setSessions((previous) => { const next = new Map(previous); next.delete(instanceId); return next }) } }) + + it("keeps global form locations for response routing while session forms stay unchanged", () => { + const instanceId = "global-form-location" + const globalForm = { + ...form, + id: "global-form", + sessionID: "global", + location: { directory: "/worktree", workspaceID: "workspace-1" }, + } + + try { + addPendingForm(instanceId, globalForm) + assert.deepEqual(getFormQueue(instanceId)[0]?.location, globalForm.location) + assert.deepEqual(formRequestOptions(globalForm), { + headers: { + "x-opencode-directory": "%2Fworktree", + "x-opencode-workspace": "workspace-1", + }, + }) + assert.equal(formRequestOptions(form), undefined) + } finally { + removePendingForm(instanceId, globalForm.id) + } + }) + + it("percent-encodes Unicode and percent signs in global form directories", () => { + assert.deepEqual(formRequestOptions({ + ...form, + sessionID: "global", + location: { directory: "/工作/100% ready" }, + }), { + headers: { + "x-opencode-directory": "%2F%E5%B7%A5%E4%BD%9C%2F100%25%20ready", + }, + }) + }) + + it("sends global replies and cancellations with their location request options", async () => { + const instanceId = "global-form-response-location" + const globalForm = { + ...form, + id: "global-response-form", + sessionID: "global", + location: { directory: "/worktree", workspaceID: "workspace-1" }, + } + const calls: unknown[][] = [] + const client = getRootClient(instanceId) + ;(client.form as any).reply = async (...args: unknown[]) => { calls.push(args) } + ;(client.form as any).cancel = async (...args: unknown[]) => { calls.push(args) } + + try { + addPendingForm(instanceId, globalForm) + await sendFormReply(instanceId, globalForm.id, { channel: "stable" }) + addPendingForm(instanceId, globalForm) + await sendFormCancel(instanceId, globalForm.id) + + assert.deepEqual(calls, [ + [ + { sessionID: "global", formID: globalForm.id, answer: { channel: "stable" } }, + formRequestOptions(globalForm), + ], + [ + { sessionID: "global", formID: globalForm.id }, + formRequestOptions(globalForm), + ], + ]) + } finally { + removePendingForm(instanceId, globalForm.id) + sdkManager.destroyClientsForInstance(instanceId) + } + }) + + it("attaches list response locations only to global forms", async () => { + const instanceId = "global-form-list-location" + const location = { directory: "/worktree", workspaceID: "workspace-1" } + const client = { + permission: { request: { list: async () => ({ location, data: [] }) } }, + form: { request: { list: async () => ({ + location, + data: [ + { ...form, id: "global-list-form", sessionID: "global" }, + { ...form, id: "session-list-form" }, + ], + }) } }, + } + addInstance({ id: instanceId, folder: "/worktree", status: "ready", client } as any) + + try { + await syncPendingRequests(instanceId) + assert.deepEqual(getFormQueue(instanceId).map((entry) => [entry.id, entry.location]), [ + ["global-list-form", location], + ["session-list-form", undefined], + ]) + } finally { + removeInstance(instanceId) + } + }) }) diff --git a/packages/ui/src/stores/forms.ts b/packages/ui/src/stores/forms.ts index 07ae1ca9..82da49a9 100644 --- a/packages/ui/src/stores/forms.ts +++ b/packages/ui/src/stores/forms.ts @@ -1,13 +1,14 @@ import { createSignal } from "solid-js" import type { FormAnswer, FormInfo } from "@opencode-ai/client" +import type { FormWithLocation } from "@opencode-ai/client/solid" -const [formQueues, setFormQueues] = createSignal>(new Map()) +const [formQueues, setFormQueues] = createSignal>(new Map()) -export function getFormQueue(instanceId: string): FormInfo[] { +export function getFormQueue(instanceId: string): FormWithLocation[] { return formQueues().get(instanceId) ?? [] } -export function addFormToQueue(instanceId: string, form: FormInfo): void { +export function addFormToQueue(instanceId: string, form: FormWithLocation): void { setFormQueues((previous) => { const next = new Map(previous) const queue = next.get(instanceId) ?? [] @@ -30,7 +31,7 @@ export function removeFormFromQueue(instanceId: string, formId: string): void { }) } -export function replaceFormQueue(instanceId: string, forms: readonly FormInfo[]): void { +export function replaceFormQueue(instanceId: string, forms: readonly FormWithLocation[]): void { setFormQueues((previous) => { const next = new Map(previous) if (forms.length) next.set(instanceId, [...forms]) @@ -43,4 +44,14 @@ export function clearFormQueue(instanceId: string): void { replaceFormQueue(instanceId, []) } -export type { FormAnswer, FormInfo } +export function formRequestOptions(form: FormWithLocation) { + if (!form.location) return undefined + return { + headers: { + "x-opencode-directory": encodeURIComponent(form.location.directory), + ...(form.location.workspaceID ? { "x-opencode-workspace": form.location.workspaceID } : {}), + }, + } +} + +export type { FormAnswer, FormInfo, FormWithLocation } diff --git a/packages/ui/src/stores/instances.ts b/packages/ui/src/stores/instances.ts index bf6d8a6d..de9d6ad6 100644 --- a/packages/ui/src/stores/instances.ts +++ b/packages/ui/src/stores/instances.ts @@ -71,9 +71,10 @@ import { addFormToQueue, clearFormQueue as clearStoredFormQueue, getFormQueue, + formRequestOptions, removeFormFromQueue, type FormAnswer, - type FormInfo, + type FormWithLocation, } from "./forms" import { invalidateFilesystemCaches } from "../lib/filesystem-events" import { detachInstanceTabMembership, requestInstanceTabClose } from "./app-tab-membership" @@ -181,7 +182,7 @@ class InterruptionRegistry { } const permissionRegistry = new InterruptionRegistry() -const formRegistry = new InterruptionRegistry() +const formRegistry = new InterruptionRegistry() type InterruptionKind = "permission" | "form" @@ -257,7 +258,7 @@ const connectionResyncs = new TrailingResyncCoordinator( const instance = instances().get(instanceId) if (!instance?.client || instance.status !== "ready") return await Promise.all([ - fetchSessions(instanceId, { reset: false }), + fetchSessions(instanceId, { reset: true }), syncPendingRequests(instanceId), refreshVolatileInstanceState(instanceId), ]) @@ -315,7 +316,6 @@ function refreshVolatileInstanceState( serverEvents.on("instance.eventStatus", (event) => { if (event.type !== "instance.eventStatus") return - if (event.status === "connecting") destroyOpenCodeData(event.instanceId) const shouldResync = connectionResyncGate.observe(event.instanceId, event.status, event.reason) if (event.status !== "connected") return if (disconnectedInstance()?.id === event.instanceId) { @@ -415,6 +415,7 @@ function attachClient(descriptor: WorkspaceDescriptor) { if (instance.client) { sdkManager.destroyClientsForInstance(descriptor.id) + destroyOpenCodeData(descriptor.id) } const client = sdkManager.createClient(descriptor.id, nextProxyPath) @@ -434,9 +435,6 @@ function attachClient(descriptor: WorkspaceDescriptor) { workspaceMetadataHydration: sessionHydration.workspaceMetadata, }) initialHydrations.set(descriptor.id, hydration) - if (sseManager.getStatuses().get(descriptor.id) === "connected") { - resyncConnectedInstance(descriptor.id) - } void hydration.catch((error) => { log.error("Failed to hydrate instance data", error) }) @@ -554,10 +552,12 @@ async function syncPendingForms( const mutationEpoch = pendingFormMutationEpochs.get(instanceId) ?? 0 try { - const remote: FormInfo[] = [] + const remote: FormWithLocation[] = [] for (const location of buildV2RequestLocations(instance.folder, getWorktrees(instanceId))) { const response = await instance.client.form.request.list({ location }) - remote.push(...response.data) + remote.push(...response.data.map((form) => form.sessionID === "global" + ? { ...form, location: response.location } + : form)) } if (!isCurrent() || (pendingFormMutationEpochs.get(instanceId) ?? 0) !== mutationEpoch) { if (propagateErrors) throw pendingRequestSyncSuperseded @@ -1627,17 +1627,20 @@ async function sendFormReply(instanceId: string, formId: string, answer: FormAns const form = getFormQueue(instanceId).find((item) => item.id === formId) if (!form) throw new Error(`Form request not found: ${formId}`) bumpEpoch(pendingFormMutationEpochs, instanceId) - await getRootClient(instanceId).form.reply({ sessionID: form.sessionID, formID: form.id, answer }) + await getRootClient(instanceId).form.reply( + { sessionID: form.sessionID, formID: form.id, answer }, + formRequestOptions(form), + ) removePendingForm(instanceId, form.id) } -let pendingFormAddedHandler: ((instanceId: string, form: FormInfo) => void) | undefined +let pendingFormAddedHandler: ((instanceId: string, form: FormWithLocation) => void) | undefined -function setPendingFormAddedHandler(handler: (instanceId: string, form: FormInfo) => void): void { +function setPendingFormAddedHandler(handler: (instanceId: string, form: FormWithLocation) => void): void { pendingFormAddedHandler = handler } -function addPendingForm(instanceId: string, form: FormInfo): FormInfo | undefined { +function addPendingForm(instanceId: string, form: FormWithLocation): FormWithLocation | undefined { bumpEpoch(pendingFormMutationEpochs, instanceId) const previous = getFormQueue(instanceId).find((item) => item.id === form.id) addFormToQueue(instanceId, form) @@ -1665,7 +1668,7 @@ function removePendingForm(instanceId: string, formId: string): void { recomputeActiveInterruption(instanceId) } -function replacePendingForms(instanceId: string, forms: readonly FormInfo[]): void { +function replacePendingForms(instanceId: string, forms: readonly FormWithLocation[]): void { const ids = new Set(forms.map((form) => form.id)) for (const form of getFormQueue(instanceId)) { if (!ids.has(form.id)) removePendingForm(instanceId, form.id) @@ -1686,7 +1689,10 @@ async function sendFormCancel(instanceId: string, formId: string): Promise const form = getFormQueue(instanceId).find((item) => item.id === formId) if (!form) throw new Error(`Form request not found: ${formId}`) bumpEpoch(pendingFormMutationEpochs, instanceId) - await getRootClient(instanceId).form.cancel({ sessionID: form.sessionID, formID: form.id }) + await getRootClient(instanceId).form.cancel( + { sessionID: form.sessionID, formID: form.id }, + formRequestOptions(form), + ) removePendingForm(instanceId, form.id) } @@ -1720,7 +1726,7 @@ function handleInstanceInvalidation(instanceId: string, event: Parameters { it("uses createData to reduce messages, permissions, and forms", () => { @@ -96,6 +98,32 @@ describe("OpenCode data projection", () => { } }) + it("preserves the data controller across server.connected", () => { + const instanceId = "opencode-data-server-connected" + const client = getRootClient(instanceId) + ;(client.session as any).active = async () => ({}) + ;(client.location as any).get = async () => ({ directory: "/work" }) + ;(client.vcs as any).get = async () => ({ location: { directory: "/work" }, data: { branch: "main" } }) + ;(client.project as any).list = async () => [] + try { + const before = applyOpenCodeDataEvent(instanceId, "/work", { + id: "permission", + type: "permission.asked", + created: 1, + data: { id: "permission", sessionID: "session", action: "read", resources: ["*"] }, + } as any) + const after = applyOpenCodeDataEvent(instanceId, "/work", { + id: "connected", type: "server.connected", created: 2, data: {}, + } as any) + + assert.strictEqual(after, before) + assert.equal(after.session.permission.list("session")?.[0]?.id, "permission") + } finally { + destroyOpenCodeData(instanceId) + sdkManager.destroyClientsForInstance(instanceId) + } + }) + it("replaces the optimistic prompt part with its native projection", () => { const instanceId = "opencode-data-optimistic-prompt" const sessionId = "session" diff --git a/packages/ui/src/stores/opencode-data.ts b/packages/ui/src/stores/opencode-data.ts index a6735f44..0d680778 100644 --- a/packages/ui/src/stores/opencode-data.ts +++ b/packages/ui/src/stores/opencode-data.ts @@ -46,7 +46,6 @@ function ensureData(instanceId: string, directory: string) { } export function applyOpenCodeDataEvent(instanceId: string, directory: string, event: OpenCodeEvent): Data { - if (event.type === "server.connected") destroyOpenCodeData(instanceId) const entry = ensureData(instanceId, directory) entry.emit(event) return entry.data diff --git a/packages/ui/src/stores/session-api.ts b/packages/ui/src/stores/session-api.ts index c7089851..602fd3dd 100644 --- a/packages/ui/src/stores/session-api.ts +++ b/packages/ui/src/stores/session-api.ts @@ -82,9 +82,14 @@ const catalogRefreshes = new Map } const agentRequestIds = new Map() const providerRequestIds = new Map() const sessionPageRequests = new Map>() -const MAX_DESCENDANT_SESSION_REQUESTS = 1_000_000 +const messageNextCursors = new Map() +const messagePageRequests = new Map>() let nextSessionListRequestId = 0 +function messagePageKey(instanceId: string, sessionId: string): string { + return `${instanceId}\0${sessionId}` +} + function catalogLocationKey(location: LocationRef): string { return `${location.directory}\0${location.workspaceID ?? ""}` } @@ -138,6 +143,13 @@ function clearSessionCatalogState(instanceId: string): void { catalogRefreshes.delete(instanceId) agentRequestIds.delete(instanceId) providerRequestIds.delete(instanceId) + const prefix = `${instanceId}\0` + for (const key of messageNextCursors.keys()) { + if (key.startsWith(prefix)) messageNextCursors.delete(key) + } + for (const key of messagePageRequests.keys()) { + if (key.startsWith(prefix)) messagePageRequests.delete(key) + } } type V2SessionListOptions = { @@ -177,45 +189,27 @@ function hasMissingParentChain(session: SDKSession, loaded: Map { +async function fetchV2Sessions( + instanceId: string, + options: V2SessionListOptions, + signal?: AbortSignal, +): Promise { const client = getRootClient(instanceId) const project = options.project ?? getInstanceMetadata(instanceId)?.project?.id const listOptions = { ...options, project, order: options.order ?? "desc" as const } if (project) delete listOptions.directory - const response = await client.session.list(buildProjectSessionListOptions(listOptions)) - const sessionsById = new Map(response.data.map((session) => [session.id, session])) if (options.parentID === null && !options.search) { - const pending = response.data.map((session) => session.id) - const visited = new Set() - let requests = 0 - while (pending.length) { - const parentID = pending.shift()! - if (visited.has(parentID)) continue - visited.add(parentID) - let cursor: string | undefined - const seenCursors = new Set() - do { - if (++requests > MAX_DESCENDANT_SESSION_REQUESTS) throw new Error("Descendant session traversal limit exceeded") - const children = await client.session.list(buildProjectSessionListOptions({ - project: project ?? sessionsById.get(parentID)?.projectID, - parentID, - order: "asc", - cursor, - })) - for (const child of children.data) { - sessionsById.set(child.id, child) - if (!visited.has(child.id)) pending.push(child.id) - } - cursor = children.cursor?.next ?? undefined - if (cursor && seenCursors.has(cursor)) throw new Error(`Repeated child session cursor: ${cursor}`) - if (cursor) seenCursors.add(cursor) - } while (cursor) - } + delete listOptions.parentID } + const response = await client.session.list( + buildProjectSessionListOptions(listOptions), + signal ? { signal } : undefined, + ) + return { - data: Array.from(sessionsById.values()), + data: response.data, complete: !response.cursor?.next, nextCursor: response.cursor?.next ?? undefined, } @@ -267,7 +261,7 @@ async function hydrateRestoredSessionChain( if (!session) { try { signal?.throwIfAborted() - const apiSession = await client.session.get({ sessionID: sessionId }) + const apiSession = await client.session.get({ sessionID: sessionId }, signal ? { signal } : undefined) signal?.throwIfAborted() setSessions((prev) => { if (getAuthoritativelyDeletedSessionIdsForInstance(instanceId).has(sessionId) || signal?.aborted) return prev @@ -289,38 +283,22 @@ async function hydrateRestoredSessionChain( } } -async function ensureV2ParentChainsLoaded(instanceId: string, apiSessions: SDKSession[], directory?: string): Promise { +async function ensureV2ParentChainsLoaded(instanceId: string, apiSessions: SDKSession[], signal?: AbortSignal): Promise { const currentSessions = sessions().get(instanceId) ?? new Map() const loaded = new Map(currentSessions) for (const session of apiSessions) loaded.set(session.id, session) - if (!apiSessions.some((session) => hasMissingParentChain(session, loaded))) return - - const page = await fetchV2Sessions(instanceId, { directory }) - const items = getV2SessionItems(page) - if (items.length === 0) return - - setSessions((prev) => { - const next = new Map(prev) - const instanceSessions = new Map(next.get(instanceId) ?? new Map()) - const deletedSessionIds = getAuthoritativelyDeletedSessionIdsForInstance(instanceId) - - for (const apiSession of items) { - if (deletedSessionIds.has(apiSession.id)) continue - const existingSession = instanceSessions.get(apiSession.id) - instanceSessions.set(apiSession.id, toClientSessionV2(instanceId, apiSession, existingSession)) - loaded.set(apiSession.id, apiSession) - } - - next.set(instanceId, instanceSessions) - return next - }) + const missingChains = apiSessions + .filter((session) => hasMissingParentChain(session, loaded)) + .map((session) => session.parentID) + if (missingChains.length > 0) await hydrateRestoredSessionChain(instanceId, missingChains, signal) } async function fetchSessions(instanceId: string, options?: { reset?: boolean strictStatus?: boolean registerInvalidation?: (invalidate: () => void) => void + signal?: AbortSignal }): Promise { const instance = instances().get(instanceId) if (!instance || !instance.client) { @@ -345,8 +323,8 @@ async function fetchSessions(instanceId: string, options?: { log.info("session.list", { instanceId, limit: PROJECT_SESSION_LIST_LIMIT, directory: sessionListOptions.directory }) const [response, activeSessions] = await Promise.all([ - fetchV2Sessions(instanceId, sessionListOptions), - getRootClient(instanceId).session.active().catch((error) => { + fetchV2Sessions(instanceId, sessionListOptions, options?.signal), + getRootClient(instanceId).session.active(options?.signal ? { signal: options.signal } : undefined).catch((error) => { log.warn("Failed to refresh active sessions", { instanceId, error }) return null }), @@ -381,6 +359,8 @@ async function fetchSessions(instanceId: string, options?: { next.set(instanceId, instanceSessions) return next }) + await ensureV2ParentChainsLoaded(instanceId, apiSessions, options?.signal) + if (!isLatestSessionListRequest(instanceId, requestId)) return if (response.complete) { const fetchedIds = new Set(apiSessions.map((session) => session.id)) @@ -482,7 +462,13 @@ async function loadNextSessionPage(instanceId: string): Promise { next.set(instanceId, current) return next }) + await ensureV2ParentChainsLoaded(instanceId, response.data) + if (getSessionNextCursor(instanceId) !== cursor) return const roots = response.data.filter((item) => !item.parentID).map((item) => item.id) + for (const item of response.data) { + const root = getSessionRoot(instanceId, item.id) + if (root && !roots.includes(root.id)) roots.push(root.id) + } setSessionPage(instanceId, roots, Boolean(response.nextCursor), false, response.nextCursor) } @@ -526,7 +512,7 @@ async function searchSessions(instanceId: string, query: string): Promise next.set(instanceId, instanceSessions) return next }) - await ensureV2ParentChainsLoaded(instanceId, searchResults, instance.folder) + await ensureV2ParentChainsLoaded(instanceId, searchResults) if (!isLatestSessionSearch(instanceId, trimmedQuery, requestId)) return @@ -809,6 +795,9 @@ function removeSessionRuntimeState(instanceId: string, sessionId: string, author removeSessionListId(instanceId, sessionId) // Drop normalized message state and caches for this session. + const pageKey = messagePageKey(instanceId, sessionId) + messageNextCursors.delete(pageKey) + messagePageRequests.delete(pageKey) messageStoreBus.getOrCreate(instanceId).clearSession(sessionId) clearCacheForSession(instanceId, sessionId) @@ -946,6 +935,7 @@ async function loadMessages( force?: boolean skipChildren?: boolean registerInvalidation?: (invalidate: () => void) => void + signal?: AbortSignal }, ): Promise { const force = options?.force ?? false @@ -997,20 +987,12 @@ async function loadMessages( try { log.info(`[HTTP] GET /session.${"messages"} for instance ${instanceId}`, { sessionId }) - const apiMessages: SessionMessagesResponse["data"] = [] - let cursor: string | undefined - const seenCursors = new Set() - do { - const response: SessionMessagesResponse = await client.message.list({ - sessionID: sessionId, - limit: 200, - ...(cursor ? { cursor } : { order: "asc" }), - }) - apiMessages.push(...response.data) - cursor = response.cursor?.next ?? undefined - if (cursor && seenCursors.has(cursor)) throw new Error(`Repeated message cursor: ${cursor}`) - if (cursor) seenCursors.add(cursor) - } while (cursor) + const response: SessionMessagesResponse = await client.message.list({ + sessionID: sessionId, + limit: 200, + order: "desc", + }, options?.signal ? { signal: options.signal } : undefined) + const apiMessages = [...response.data].reverse() if (!instances().has(instanceId) || !isCurrentMessageLoad(instanceId, sessionId, loadEpoch) @@ -1045,6 +1027,9 @@ async function loadMessages( next.set(instanceId, loadedSet) return next }) + const nextCursor = response.cursor?.next ?? undefined + if (nextCursor) messageNextCursors.set(messagePageKey(instanceId, sessionId), nextCursor) + else messageNextCursors.delete(messagePageKey(instanceId, sessionId)) } } else { const seenMessageIds = new Set() @@ -1118,6 +1103,9 @@ async function loadMessages( next.set(instanceId, loadedSet) return next }) + const nextCursor = response.cursor?.next ?? undefined + if (nextCursor) messageNextCursors.set(messagePageKey(instanceId, sessionId), nextCursor) + else messageNextCursors.delete(messagePageKey(instanceId, sessionId)) reconcilePendingPermissionsV2(instanceId, sessionId) } } @@ -1148,6 +1136,7 @@ async function loadMessages( force: true, skipChildren, registerInvalidation: options?.registerInvalidation, + signal: options?.signal, }) } @@ -1170,6 +1159,89 @@ async function loadMessages( } } +async function loadMoreMessages(instanceId: string, sessionId: string, signal?: AbortSignal): Promise { + const key = messagePageKey(instanceId, sessionId) + const pending = messagePageRequests.get(key) + if (pending) return pending + const request = loadNextMessagePage(instanceId, sessionId, signal).finally(() => messagePageRequests.delete(key)) + messagePageRequests.set(key, request) + return request +} + +function hasMoreMessages(instanceId: string, sessionId: string): boolean { + return messageNextCursors.has(messagePageKey(instanceId, sessionId)) +} + +async function loadNextMessagePage(instanceId: string, sessionId: string, signal?: AbortSignal): Promise { + const key = messagePageKey(instanceId, sessionId) + const cursor = messageNextCursors.get(key) + if (!cursor) return + const instance = instances().get(instanceId) + const session = sessions().get(instanceId)?.get(sessionId) + if (!instance?.client) throw new Error("Instance not ready") + if (!session) throw new Error("Session not found") + + const loadEpoch = advanceMessageLoadEpoch(instanceId, sessionId) + + const response = await getRootClient(instanceId).message.list({ + sessionID: sessionId, + limit: 200, + order: "desc", + cursor, + }, signal ? { signal } : undefined) + if (!instances().has(instanceId) + || !isCurrentMessageLoad(instanceId, sessionId, loadEpoch) + || !sessions().get(instanceId)?.has(sessionId) + || messageNextCursors.get(key) !== cursor) return + + const store = messageStoreBus.getOrCreate(instanceId) + const existingIds = store.getSessionMessageIds(sessionId) + const existing = new Set(existingIds) + const olderIds: string[] = [] + for (const apiMessage of [...response.data].reverse()) { + const normalized = normalizeSessionMessage(sessionId, apiMessage) + if (existing.has(normalized.message.id)) continue + existing.add(normalized.message.id) + olderIds.push(normalized.message.id) + store.upsertMessage({ + id: normalized.message.id, + sessionId, + role: normalized.message.type, + status: normalized.message.status, + createdAt: normalized.message.timestamp, + updatedAt: normalized.message.timestamp, + parts: normalized.message.parts, + isEphemeral: normalized.message.status === "sending" + || (normalized.message.type === "assistant" && normalized.message.status === "streaming"), + }) + store.setMessageInfo(normalized.info.id, normalized.info) + } + if (olderIds.length > 0) { + store.addOrUpdateSession({ + id: sessionId, + title: session.title, + parentId: session.parentId, + revert: session.revert, + messageIds: [...olderIds, ...existingIds], + }) + store.rebuildUsage(sessionId, store.getSessionMessageIds(sessionId) + .map((id) => store.getMessageInfo(id)) + .filter((info): info is NonNullable => Boolean(info))) + } + const nextCursor = response.cursor?.next ?? undefined + if (nextCursor) messageNextCursors.set(key, nextCursor) + else messageNextCursors.delete(key) + setMessagesLoaded((prev) => { + const next = new Map(prev) + const loadedSet = next.get(instanceId) || new Set() + loadedSet.add(sessionId) + next.set(instanceId, loadedSet) + return next + }) + reconcilePendingPermissionsV2(instanceId, sessionId) + updateSessionInfo(instanceId, sessionId) +} + export { createSession, deleteSession, @@ -1185,6 +1257,8 @@ export { searchSessions, forkSession, loadMessages, + loadMoreMessages, + hasMoreMessages, clearSessionListRequestState, clearSessionCatalogState, } diff --git a/packages/ui/src/stores/session-pagination.test.ts b/packages/ui/src/stores/session-pagination.test.ts index 2e061ca5..01df78a6 100644 --- a/packages/ui/src/stores/session-pagination.test.ts +++ b/packages/ui/src/stores/session-pagination.test.ts @@ -22,9 +22,10 @@ describe("project session list loading", () => { }) it("passes native cursors through unchanged", () => { - assert.deepEqual(buildProjectSessionListOptions({ directory: "/tmp/project", cursor: "next-page" }), { + const cursor = `next-page-${"x".repeat(4096)}` + assert.deepEqual(buildProjectSessionListOptions({ directory: "/tmp/project", cursor }), { directory: "/tmp/project", - cursor: "next-page", + cursor, limit: PROJECT_SESSION_LIST_LIMIT, }) }) @@ -35,4 +36,11 @@ describe("project session list loading", () => { assert.deepEqual(next, { ids: ["root-1", "root-2"], hasMore: false, nextCursor: undefined }) }) + it("replaces stale pages when reconnect refreshes the newest roots", () => { + const stale = applySessionPage(getDefaultSessionPaginationState(), ["old-2", "old-1"], true, true, "old-page-2") + const refreshed = applySessionPage(stale, ["new", "old-2"], true, true, "new-page-2") + + assert.deepEqual(refreshed, { ids: ["new", "old-2"], hasMore: true, nextCursor: "new-page-2" }) + }) + }) diff --git a/packages/ui/src/stores/session-request-authority.test.ts b/packages/ui/src/stores/session-request-authority.test.ts index 4359eda0..acf6d3c9 100644 --- a/packages/ui/src/stores/session-request-authority.test.ts +++ b/packages/ui/src/stores/session-request-authority.test.ts @@ -5,7 +5,7 @@ import { sdkManager } from "../lib/sdk-manager.ts" import type { Session } from "../types/session.ts" import { addInstance, removeInstance } from "./instances.ts" import { messageStoreBus } from "./message-v2/bus.ts" -import { fetchSessions, loadMessages, loadMoreSessions, removeSessionRuntimeState, searchSessions } from "./session-api.ts" +import { fetchSessions, hasMoreMessages, loadMessages, loadMoreMessages, loadMoreSessions, removeSessionRuntimeState, searchSessions } from "./session-api.ts" import { setInstanceMetadata } from "./instance-metadata.ts" import { clearInstanceDeletedSessionAuthority, @@ -68,7 +68,8 @@ describe("session request authority", () => { const search = deferred() const parents = deferred() let calls = 0 - ;(client.session as any).list = () => (++calls === 1 ? search.promise : parents.promise) + ;(client.session as any).list = () => { calls += 1; return search.promise } + ;(client.session as any).get = () => parents.promise try { const request = searchSessions(instanceId, "child") @@ -76,12 +77,13 @@ describe("session request authority", () => { await new Promise((resolve) => setImmediate(resolve)) removeSessionRuntimeState(instanceId, "child") removeSessionRuntimeState(instanceId, "parent") - parents.resolve({ data: [apiSession("parent")] }) + parents.resolve(apiSession("parent")) await request assert.equal(sessions().get(instanceId)?.has("child") ?? false, false) assert.equal(sessions().get(instanceId)?.has("parent") ?? false, false) assert.deepEqual(getSessionSearchResultIds(instanceId), []) + assert.equal(calls, 1) } finally { cleanup() } @@ -163,54 +165,80 @@ describe("session request authority", () => { } }) - it("does not replace messages when a later page fails", async () => { + it("shows the latest message page and preserves it when cursor load-more fails", async () => { const instanceId = "partial-message-pages", sessionId = "session" const { client, cleanup } = setup(instanceId) let failSecondPage = false + let pendingSecondPage: ReturnType> | undefined + const requests: any[] = [] ;(client as any).message = { list: async (input: any) => { + requests.push(input) if (input.cursor && failSecondPage) throw new Error("cursor failed") + if (input.cursor && pendingSecondPage) return pendingSecondPage.promise return input.cursor - ? { data: [apiMessage("old-2")], cursor: {} } - : { data: [apiMessage("old-1")], cursor: { next: "page-2" } } + ? { data: [apiMessage("old-2"), apiMessage("old-1")], cursor: {} } + : { data: [apiMessage("new-2"), apiMessage("new-1")], cursor: { next: "page-2" } } } } setSessions((previous) => new Map(previous).set(instanceId, new Map([[sessionId, session(instanceId, sessionId)]]))) try { await loadMessages(instanceId, sessionId) + assert.deepEqual(messageStoreBus.getOrCreate(instanceId).getSessionMessageIds(sessionId), ["new-1", "new-2"]) + assert.deepEqual(requests, [{ sessionID: sessionId, limit: 200, order: "desc" }]) + assert.equal(hasMoreMessages(instanceId, sessionId), true) failSecondPage = true - await assert.rejects(loadMessages(instanceId, sessionId, { force: true }), /cursor failed/) - assert.deepEqual(messageStoreBus.getOrCreate(instanceId).getSessionMessageIds(sessionId), ["old-1", "old-2"]) + await assert.rejects(loadMoreMessages(instanceId, sessionId), /cursor failed/) + assert.deepEqual(messageStoreBus.getOrCreate(instanceId).getSessionMessageIds(sessionId), ["new-1", "new-2"]) assert.equal(messagesLoaded().get(instanceId)?.has(sessionId), true) + + failSecondPage = false + pendingSecondPage = deferred() + const firstLoadMore = loadMoreMessages(instanceId, sessionId) + const concurrentLoadMore = loadMoreMessages(instanceId, sessionId) + await new Promise((resolve) => setImmediate(resolve)) + assert.equal(loading().loadingMessages.get(instanceId)?.has(sessionId) ?? false, false) + assert.deepEqual(messageStoreBus.getOrCreate(instanceId).getSessionMessageIds(sessionId), ["new-1", "new-2"]) + assert.equal(requests.filter((request: any) => request.cursor === "page-2").length, 2) + pendingSecondPage.resolve({ data: [apiMessage("old-2"), apiMessage("old-1")], cursor: {} }) + await Promise.all([firstLoadMore, concurrentLoadMore]) + assert.deepEqual(messageStoreBus.getOrCreate(instanceId).getSessionMessageIds(sessionId), ["old-1", "old-2", "new-1", "new-2"]) + assert.equal((requests.at(-1) as any)?.cursor, "page-2") + assert.equal(hasMoreMessages(instanceId, sessionId), false) } finally { cleanup() } }) - it("loads every descendant depth by project and applies active state to later pages", async () => { + it("loads paginated project sessions without per-parent requests", async () => { const instanceId = "project-descendants" const { client, cleanup } = setup(instanceId) const requests: any[] = [] - let active: Record = {} + const active: Record = { later: {} } setInstanceMetadata(instanceId, { project: { id: "project", directory: "/work", canonical: "/work" } as any }) ;(client.session as any).active = async () => active ;(client.session as any).list = async (input: any) => { requests.push(input) - if (input.cursor === "root-page-2") return { data: [apiSession("later")], cursor: {} } - if (input.parentID === "root" && !input.cursor) return { data: [{ ...apiSession("child", "root"), subpath: "other-worktree" }], cursor: { next: "child-page-2" } } - if (input.parentID === "child") return { data: [apiSession("grandchild", "child")], cursor: {} } - if (input.parentID) return { data: [], cursor: {} } - return { data: [apiSession("root")], cursor: { next: "root-page-2" } } + if (input.cursor === "page-2") { + return { data: [apiSession("later"), apiSession("grandchild", "child")], cursor: {} } + } + return { + data: [apiSession("root"), { ...apiSession("child", "root"), subpath: "other-worktree" }], + cursor: { next: "page-2" }, + } } try { await fetchSessions(instanceId) - assert.equal(sessions().get(instanceId)?.has("grandchild"), true) + assert.equal(sessions().get(instanceId)?.has("grandchild"), false) assert.equal(requests[0].project, "project") assert.equal("directory" in requests[0], false) - assert.equal(requests.find((request) => request.parentID === "child")?.subpath, undefined) + assert.equal(requests.every((request) => !("parentID" in request)), true) + assert.equal(requests.length, 1) - active = { later: {} } await loadMoreSessions(instanceId) + assert.equal(requests.length, 2) + assert.equal(requests[1].cursor, "page-2") + assert.equal(sessions().get(instanceId)?.has("grandchild"), true) assert.equal(sessions().get(instanceId)?.get("later")?.status, "working") assert.equal(sessions().get(instanceId)?.get("later")?.runtimeStatusKnown, true) } finally { diff --git a/packages/ui/src/stores/sessions.ts b/packages/ui/src/stores/sessions.ts index dda6d4c5..9458b849 100644 --- a/packages/ui/src/stores/sessions.ts +++ b/packages/ui/src/stores/sessions.ts @@ -80,6 +80,8 @@ import { refreshSessionCatalog, fetchSessions, hydrateRestoredSessionChain, + hasMoreMessages, + loadMoreMessages, loadMoreSessions, searchSessions, forkSession, @@ -146,6 +148,8 @@ export { refreshSessionCatalog, fetchSessions, hydrateRestoredSessionChain, + hasMoreMessages, + loadMoreMessages, loadMoreSessions, searchSessions, forkSession,