diff --git a/docs/design/2026-07-11-daemon-workspace-runtime-removal.md b/docs/design/2026-07-11-daemon-workspace-runtime-removal.md new file mode 100644 index 00000000000..db1076199b7 --- /dev/null +++ b/docs/design/2026-07-11-daemon-workspace-runtime-removal.md @@ -0,0 +1,45 @@ +# Daemon Workspace Runtime Removal + +## Context + +Runtime workspace registration and persistent registration are already available, but forgetting a persistent registration does not unload the live bridge, ACP mount, session admission state, or memory lane. This design adds synchronous hot removal for secondary runtimes while preserving the existing registration-forget API. + +## Scope and invariants + +- Only dynamically registered and persistence-restored secondary runtimes are removable. The primary and every `--workspace` runtime are static. +- `DELETE /workspaces/:workspace` removes the runtime and all known persistent aliases. It never removes workspace files, settings, transcripts, archives, or other project data. +- Non-force removal is observational: if the frozen runtime has activity, every gate is rolled back and the request returns `409 workspace_busy`. Force removal terminates that activity. +- Persistence is committed before destructive cleanup. A store failure restores the active runtime. Cleanup failures after the store commit cannot roll the operation back and use synchronous bridge kill as a fallback. +- A removed cwd remains reserved until cleanup completes, then may be registered again with a fresh bridge, ACP dispatcher, connection registry, and memory lane. + +## Protocol + +Production daemons advertise `workspace_runtime_removal` when the removal controller is installed. Capability workspace rows add optional `removable`; old clients and daemons remain compatible. + +`DELETE /workspaces/:workspace` uses the existing workspace-id-or-canonical-cwd selector and accepts an optional JSON body containing a boolean `force`. Success returns the removed identity, whether force was requested, whether any persistent alias was removed, and the final post-drain activity snapshot. A non-force request that is already observably busy may return an earlier pre-drain snapshot without briefly gating the runtime. Existing `DELETE /workspace-registrations/:id` remains forget-only. + +## Lifecycle + +The registry tracks active, draining, and removed runtimes. Public resolution sees only active runtimes; management resolution retains draining runtimes for conflict reporting and cwd reservation. + +Removal first takes a fast activity snapshot. It then synchronously marks the registry draining, closes per-workspace session admission, and drains the ACP mount and memory lane. The final snapshot reads pending session reservations before live bridge counts so a reservation-to-session transition cannot appear idle. A busy non-force request reverses the gates. Otherwise all known registration IDs are deleted atomically, queued memory work is failed, the sub-session launcher and bridge are stopped, the ACP mount is disposed, ownership indexes are cleared, and the registry entry is completed. + +Runtime cleanup is memoized by runtime identity, not cwd, so a later runtime registered at the same path cannot reuse an old cleanup promise. Daemon shutdown seals management operations, waits for them to converge, stops launchers, and then uses the same bridge teardown path for the remaining managed runtimes. + +## Persistence identity + +Restoration records the ID of each raw stored path before canonicalization. Multiple raw aliases that resolve to one runtime are retained as one ID set, including aliases shadowed by an explicit startup workspace. Removal deletes that set plus the canonical registration ID under one store lock without changing the schema. + +## UI + +The Web Shell exposes removal only when both the feature tag and `removable: true` are present. The action remains available for untrusted workspaces. The first confirmation performs a non-force request; `workspace_busy` renders the activity counts and offers force removal. Force is disabled when the current session belongs to the target workspace. Success reconciles capabilities and session lists and falls back to the primary workspace when necessary. + +## Failure and compatibility analysis + +Client disconnects and SDK timeouts do not cancel server-side cleanup. Concurrent add, persistence promotion, and remove operations are serialized per canonical cwd. Shutdown rejects new management operations with `daemon_shutting_down` and waits for already-started work. Old clients ignore the optional capability field and feature; old daemons continue to produce a normal `DaemonHttpError` for the missing route. + +The workspace-scoped channel worker group supplies activity and teardown through a thin adapter. Draining blocks reload and webhook routing for the target workspace; committed removal stops and unregisters only that worker so daemon status and pidfile metadata converge without affecting other workspaces. + +## Verification + +Unit coverage targets registry state transitions and owner cleanup, admission drain rollback, alias batch deletion, busy/force/store-failure route behavior, bridge shutdown reason idempotence, memory-lane cancellation, SDK request encoding, and Web Shell feature and force guards. The E2E plan lives at `.qwen/e2e-tests/workspace-runtime-removal.md`. diff --git a/docs/developers/qwen-serve-protocol.md b/docs/developers/qwen-serve-protocol.md index 7ea1bdee85f..0664e21d190 100644 --- a/docs/developers/qwen-serve-protocol.md +++ b/docs/developers/qwen-serve-protocol.md @@ -198,6 +198,8 @@ registry. Clients **must** gate UI off `features`, not off `mode` (per design `persistent_workspace_registration` advertises durable registration for workspaces added at runtime. `POST /workspaces` accepts `{ "cwd": "/absolute/path", "persist": true }`; success includes `persisted: true`. Registrations are scoped to the daemon's canonical primary workspace under the user's Qwen home and are restored on the next daemon start. Omitting `persist` preserves process-local registration. `GET /workspace-registrations` lists the stored desired set, and `DELETE /workspace-registrations/:id` forgets an entry for the next restart without hot-removing an active runtime. +`workspace_runtime_removal` advertises synchronous hot removal through `DELETE /workspaces/:workspace`. Capability workspace entries add optional `removable`; only rows with `removable: true` may be removed. Removal also forgets every persistent registration alias for the runtime, but never deletes files, settings, transcripts, or archives. + `session_load` and `session_resume` advertise the explicit-restore routes (`POST /session/:id/load` and `POST /session/:id/resume`). Older daemons return `404` for these paths, so SDK clients should pre-flight `caps.features` before calling. `unstable_session_resume` is still advertised as a deprecated alias for compatibility with SDKs that shipped while the underlying ACP method was named `connection.unstable_resumeSession`; new clients should gate on `session_resume`. `session_transcript` advertises `GET /session/:id/transcript`, a read-only paged replay view over the persisted active-session JSONL. It is separate from `/load`: it does not attach a client, seed the live EventBus, create a live session, or change the live replay window. Clients should use it when they need the complete on-disk transcript for a long session, and continue using `/load` only for bounded live replay during cold UI restore. @@ -581,7 +583,7 @@ Stable contract: when `v` increments the frame layout has changed in a backwards > **`workspaceCwd`** is the canonical absolute path for the daemon's primary workspace. Use it to omit `cwd` on `POST /session` (the route falls back to this primary path) and to keep old single-workspace clients compatible. Additive to v=1: pre-§02 v=1 daemons omit the field — clients that target older builds should null-check before consuming it. -> **`workspaces[]`** is present only when `features` contains `multi_workspace_sessions`. Each entry is `{ id, cwd, primary, trusted }`. The first/primary workspace remains mirrored by `workspaceCwd`; new clients choose a non-primary runtime by passing that entry's `cwd` to `POST /session`. Untrusted workspaces are advertised for diagnostics but reject fresh session creation with `403 untrusted_workspace` until trust changes. +> **`workspaces[]`** is present only when `features` contains `multi_workspace_sessions`. Each entry is `{ id, cwd, primary, trusted, removable? }`. The first/primary workspace remains mirrored by `workspaceCwd`; new clients choose a non-primary runtime by passing that entry's `cwd` to `POST /session`. Untrusted workspaces are advertised for diagnostics but reject fresh session creation with `403 untrusted_workspace` until trust changes. `removable` is present on daemons that support runtime removal and is true only for process-dynamic or persistence-restored secondary runtimes. The workspace feature tags and `workspaces[]` are dynamic. Clients that add a workspace must fetch `/capabilities` again after the mutation completes; the daemon does not broadcast capability changes to clients that cached an earlier response. Forgetting persistence does not unload an active runtime, so that runtime remains advertised until restart. @@ -607,9 +609,36 @@ A newly created runtime returns `201`; promoting an already-active secondary wor Errors include `400 invalid_path` / `invalid_persist_flag` / `invalid_persist_target`, `409 workspace_exists` / `workspace_nested` / `workspace_limit_reached`, `500 workspace_registration_store_error` / `runtime_creation_failed`, and `501 persistence_not_available` / `not_implemented`. +### `DELETE /workspaces/:workspace` + +Remove one removable secondary runtime. The selector follows the plural workspace routing rules and accepts either a workspace ID or a URL-encoded absolute cwd. The optional JSON body is `{ "force": boolean }`; omitting it requests non-force removal. + +Non-force removal returns `409 workspace_busy` with an `activity` snapshot when the frozen runtime has sessions, prompts, pending starts, ACP connections, memory tasks, or workspace channel workers. Sending `{ "force": true }` terminates those resources. A successful response is: + +```json +{ + "removed": true, + "workspaceId": "stable-workspace-id", + "workspaceCwd": "/canonical/path/to/secondary-workspace", + "forced": true, + "persistedRegistrationRemoved": true, + "activity": { + "sessions": 2, + "activePrompts": 1, + "pendingSessionStarts": 0, + "acpConnections": 1, + "memoryTasks": 0, + "channelWorkers": 0 + } +} +``` + +An immediately busy non-force request returns a fast pre-drain activity snapshot. Once drain starts, the busy or success response contains the final snapshot taken after admission and ACP drain gates close and before cleanup begins. Errors include `400 invalid_force_flag` / `workspace_mismatch`, `409 workspace_busy` / `primary_workspace_removal_forbidden` / `static_workspace_removal_forbidden` / `workspace_removal_in_progress` / `workspace_registration_in_progress`, `500 workspace_persist_failed` / `workspace_runtime_removal_failed`, `501 workspace_runtime_removal_unsupported`, and `503 daemon_shutting_down`. + ### `GET /workspace-registrations` List the persisted desired workspace set for this primary workspace. Entries remain visible with `active: false` when a stored directory could not be restored during the current start. +An entry remains `active: true` while its runtime is draining because the runtime still owns live resources until removal completes. ```json { diff --git a/integration-tests/cli/qwen-serve-routes.test.ts b/integration-tests/cli/qwen-serve-routes.test.ts index e582e7769f2..6f0be47757f 100644 --- a/integration-tests/cli/qwen-serve-routes.test.ts +++ b/integration-tests/cli/qwen-serve-routes.test.ts @@ -147,6 +147,7 @@ beforeAll(async () => { 'QWEN_SERVE_WRITER_IDLE_TIMEOUT_MS', 'QWEN_SERVE_RATE_LIMIT', 'QWEN_SERVE_NO_MCP_POOL', + 'QWEN_SERVE_NO_PERSISTENT_REGISTRATION', 'QWEN_SERVE_CLIENT_MCP_OVER_WS', 'QWEN_SERVE_CDP_TUNNEL_OVER_WS', ].includes(k), @@ -373,6 +374,8 @@ describe('qwen serve — capabilities envelope', () => { 'session_branch', 'workspace_reload', 'channel_control', + 'persistent_workspace_registration', + 'workspace_runtime_removal', 'workspace_qualified_rest_core', 'workspace_persisted_transcript', 'voice_transcribe', diff --git a/packages/acp-bridge/src/bridge.test.ts b/packages/acp-bridge/src/bridge.test.ts index d58e9a0b624..c1fac1ffda5 100644 --- a/packages/acp-bridge/src/bridge.test.ts +++ b/packages/acp-bridge/src/bridge.test.ts @@ -164,6 +164,51 @@ describe('createAcpSessionBridge', () => { await bridge.shutdown(); }); + it('memoizes shutdown and keeps the first workspace-removal reason', async () => { + const lifecycle: Array<{ type: string; reason?: string }> = []; + const bridge = makeBridge({ + channelFactory: async () => makeChannel().channel, + sessionLifecycle: (event) => lifecycle.push(event), + }); + const session = await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + const iterator = bridge + .subscribeEvents(session.sessionId) + [Symbol.asyncIterator](); + const terminalEvent = iterator.next(); + + const first = bridge.shutdown({ reason: 'workspace_removed' }); + const second = bridge.shutdown({ reason: 'daemon_shutdown' }); + + expect(second).toBe(first); + await first; + await expect(terminalEvent).resolves.toMatchObject({ + value: { + type: 'session_died', + data: { reason: 'workspace_removed' }, + }, + }); + expect(lifecycle.at(-1)).toMatchObject({ + type: 'removed', + reason: 'workspace_removed', + }); + }); + + it('publishes the shutdown promise before lifecycle callbacks can re-enter', async () => { + let reentered: Promise | undefined; + const bridge = makeBridge({ + channelFactory: async () => makeChannel().channel, + sessionLifecycle: (event) => { + if (event.type === 'removed') reentered = bridge.shutdown(); + }, + }); + await bridge.spawnOrAttach({ workspaceCwd: WS_A }); + + const first = bridge.shutdown({ reason: 'workspace_removed' }); + + expect(reentered).toBe(first); + await first; + }); + it('accepts a valid BridgeOptions.eventRingSize at construction time', () => { // Smoke: positive finite integers are accepted; the underlying // EventBus ring-size threading is exercised end-to-end in diff --git a/packages/acp-bridge/src/bridge.ts b/packages/acp-bridge/src/bridge.ts index a637950a625..f750c291a57 100644 --- a/packages/acp-bridge/src/bridge.ts +++ b/packages/acp-bridge/src/bridge.ts @@ -1607,6 +1607,7 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { // (b) `server.close` rejecting new connections, during which a // late-arriving `POST /session` slips a fresh child past cleanup. let shuttingDown = false; + let shutdownPromise: Promise | undefined; // Tee writeServeDebugLine through the optional onDiagnosticLine callback. // The module-level writeServeDebugLine is left intact for other entry points; @@ -6880,100 +6881,111 @@ export function createAcpSessionBridge(opts: BridgeOptions): AcpSessionBridge { } }, - async shutdown() { - // Set BEFORE the snapshot so any racing `spawnOrAttach` triggered - // by an in-flight HTTP connection after `runQwenServe.close()` - // entered the bridge.shutdown() phase fails fast instead of - // spawning a child this teardown won't see. - shuttingDown = true; - cancelIdleTimer(); - stopSessionReaper(); - const entries = Array.from(byId.values()); - // Snapshot every alive channel (typically 1; up to 2 during a - // `killSession`-then-`spawnOrAttach` overlap) — entries are - // intentionally NOT removed from `aliveChannels` here; their - // `channel.exited` handlers clear them once the OS has reaped - // each child. That preserves the BkUyD invariant: a - // double-Ctrl+C arriving mid-SIGTERM-grace can still find every - // alive channel via `killAllSync`. Marking each `isDying` makes - // them invisible to any racing `ensureChannel` call — but - // `shuttingDown` already blocks new `spawnOrAttach` upstream, - // so this is mostly belt-and-suspenders (a direct internal - // `ensureChannel` past the gate would still see the dying - // state and not attach). - const channels = Array.from(aliveChannels); - for (const ci of channels) ci.isDying = true; - // Drain mediator pending state before clearing byId so awaiting - // `requestPermission` callers unwind. Each `forgetSession` - // settles all matching pending as session_closed; the bridge's - // per-entry index gets cleared alongside. - for (const e of entries) { - permissionMediator.forgetSession(e.sessionId); - e.pendingPermissionIds.clear(); - e.pendingInteractions.clear(); - } - defaultEntry = undefined; - byId.clear(); - // Publish a terminal `session_died` BEFORE closing each bus so SSE - // subscribers can distinguish "daemon shut down" from a transient - // network error and don't sit indefinitely retrying. The - // channel.exited handler also publishes this on a child crash, - // but at shutdown time the entry has already been removed from - // `byId` (above), so the handler's `byId.get(...)` is undefined - // and the automatic publish wouldn't fire. - for (const e of entries) { - telemetry.metrics?.sessionLifecycle('die'); - emitSessionLifecycle({ - type: 'removed', - sessionId: e.sessionId, - workspaceCwd: e.workspaceCwd, - reason: 'daemon_shutdown', - }); - try { - e.events.publish({ - type: 'session_died', - data: { sessionId: e.sessionId, reason: 'daemon_shutdown' }, + shutdown(options) { + if (shutdownPromise) return shutdownPromise; + const shutdownReason = options?.reason ?? 'daemon_shutdown'; + let resolveShutdown: (() => void) | undefined; + let rejectShutdown: ((reason?: unknown) => void) | undefined; + shutdownPromise = new Promise((resolve, reject) => { + resolveShutdown = resolve; + rejectShutdown = reject; + }); + void (async () => { + // Set BEFORE the snapshot so any racing `spawnOrAttach` triggered + // by an in-flight HTTP connection after `runQwenServe.close()` + // entered the bridge.shutdown() phase fails fast instead of + // spawning a child this teardown won't see. + shuttingDown = true; + cancelIdleTimer(); + stopSessionReaper(); + const entries = Array.from(byId.values()); + // Snapshot every alive channel (typically 1; up to 2 during a + // `killSession`-then-`spawnOrAttach` overlap) — entries are + // intentionally NOT removed from `aliveChannels` here; their + // `channel.exited` handlers clear them once the OS has reaped + // each child. That preserves the BkUyD invariant: a + // double-Ctrl+C arriving mid-SIGTERM-grace can still find every + // alive channel via `killAllSync`. Marking each `isDying` makes + // them invisible to any racing `ensureChannel` call — but + // `shuttingDown` already blocks new `spawnOrAttach` upstream, + // so this is mostly belt-and-suspenders (a direct internal + // `ensureChannel` past the gate would still see the dying + // state and not attach). + const channels = Array.from(aliveChannels); + for (const ci of channels) ci.isDying = true; + // Drain mediator pending state before clearing byId so awaiting + // `requestPermission` callers unwind. Each `forgetSession` + // settles all matching pending as session_closed; the bridge's + // per-entry index gets cleared alongside. + for (const e of entries) { + permissionMediator.forgetSession(e.sessionId); + e.pendingPermissionIds.clear(); + e.pendingInteractions.clear(); + } + defaultEntry = undefined; + byId.clear(); + // Publish a terminal `session_died` BEFORE closing each bus so SSE + // subscribers can distinguish "daemon shut down" from a transient + // network error and don't sit indefinitely retrying. The + // channel.exited handler also publishes this on a child crash, + // but at shutdown time the entry has already been removed from + // `byId` (above), so the handler's `byId.get(...)` is undefined + // and the automatic publish wouldn't fire. + for (const e of entries) { + telemetry.metrics?.sessionLifecycle('die'); + emitSessionLifecycle({ + type: 'removed', + sessionId: e.sessionId, + workspaceCwd: e.workspaceCwd, + reason: shutdownReason, }); - } catch { - /* bus already closed */ + try { + e.events.publish({ + type: 'session_died', + data: { sessionId: e.sessionId, reason: shutdownReason }, + }); + } catch { + /* bus already closed */ + } + e.events.close(); } - e.events.close(); - } - // Wait for in-flight channel + session spawns. The snapshot - // above only sees what's already registered; a doSpawn past - // `newSession()` but pre-`byId.set` is missed, as is an - // `ensureChannel` past `channelFactory()` but pre-`channelInfo - // = info`. The late-shutdown re-checks at doSpawn/ensureChannel - // catch both — but without these awaits, `bridge.shutdown()` - // would resolve before they finish, and the orphan stderr - // error from a half-built child would fire AFTER the daemon - // claimed graceful shutdown (log-confusing). - const inFlightSessionAwaits = Array.from(inFlightSpawns.values()).map( - (p): Promise => - p.then( - () => undefined, - () => undefined, - ), - ); - const inFlightRestoreAwaits = Array.from(inFlightRestores.values()).map( - (restore): Promise => - restore.promise.then( - () => undefined, - () => undefined, - ), - ); - const inFlightChannelAwait: Promise = inFlightChannelSpawn - ? inFlightChannelSpawn.then( - () => undefined, - () => undefined, - ) - : Promise.resolve(); - await Promise.all([ - ...channels.map((ci) => ci.channel.kill().catch(() => {})), - ...inFlightSessionAwaits, - ...inFlightRestoreAwaits, - inFlightChannelAwait, - ]); + // Wait for in-flight channel + session spawns. The snapshot + // above only sees what's already registered; a doSpawn past + // `newSession()` but pre-`byId.set` is missed, as is an + // `ensureChannel` past `channelFactory()` but pre-`channelInfo + // = info`. The late-shutdown re-checks at doSpawn/ensureChannel + // catch both — but without these awaits, `bridge.shutdown()` + // would resolve before they finish, and the orphan stderr + // error from a half-built child would fire AFTER the daemon + // claimed graceful shutdown (log-confusing). + const inFlightSessionAwaits = Array.from(inFlightSpawns.values()).map( + (p): Promise => + p.then( + () => undefined, + () => undefined, + ), + ); + const inFlightRestoreAwaits = Array.from(inFlightRestores.values()).map( + (restore): Promise => + restore.promise.then( + () => undefined, + () => undefined, + ), + ); + const inFlightChannelAwait: Promise = inFlightChannelSpawn + ? inFlightChannelSpawn.then( + () => undefined, + () => undefined, + ) + : Promise.resolve(); + await Promise.all([ + ...channels.map((ci) => ci.channel.kill().catch(() => {})), + ...inFlightSessionAwaits, + ...inFlightRestoreAwaits, + inFlightChannelAwait, + ]); + })().then(resolveShutdown, rejectShutdown); + return shutdownPromise; }, async preheat() { diff --git a/packages/acp-bridge/src/bridgeErrors.ts b/packages/acp-bridge/src/bridgeErrors.ts index 638fb57b8ed..6bb2f770e0b 100644 --- a/packages/acp-bridge/src/bridgeErrors.ts +++ b/packages/acp-bridge/src/bridgeErrors.ts @@ -531,6 +531,16 @@ export class SessionBusyError extends Error { } } +export class WorkspaceDrainingError extends Error { + readonly code = 'workspace_draining'; + readonly workspaceCwd: string; + constructor(workspaceCwd: string) { + super(`Workspace ${JSON.stringify(workspaceCwd)} is being removed`); + this.name = 'WorkspaceDrainingError'; + this.workspaceCwd = workspaceCwd; + } +} + export class InvalidRewindTargetError extends Error { readonly sessionId: string; constructor(sessionId: string, message?: string) { diff --git a/packages/acp-bridge/src/bridgeTypes.ts b/packages/acp-bridge/src/bridgeTypes.ts index 13b3187a036..08b681a2c5c 100644 --- a/packages/acp-bridge/src/bridgeTypes.ts +++ b/packages/acp-bridge/src/bridgeTypes.ts @@ -1252,8 +1252,8 @@ export interface AcpSessionBridge { */ killAllSync(): void; - /** Close all live child processes; called on daemon shutdown. */ - shutdown(): Promise; + /** Close all live child processes; called on daemon/workspace shutdown. */ + shutdown(options?: BridgeShutdownOptions): Promise; /** * Eagerly spawn the ACP child so the first session doesn't pay @@ -1263,6 +1263,10 @@ export interface AcpSessionBridge { preheat(): Promise; } +export interface BridgeShutdownOptions { + reason?: 'daemon_shutdown' | 'workspace_removed'; +} + export interface ShellCommandResult { exitCode: number | null; output: string; diff --git a/packages/cli/src/serve/acp-http/index.ts b/packages/cli/src/serve/acp-http/index.ts index 9fb4fd78942..9b1654ef81d 100644 --- a/packages/cli/src/serve/acp-http/index.ts +++ b/packages/cli/src/serve/acp-http/index.ts @@ -24,8 +24,8 @@ import type { WorkspaceRuntime, } from '../workspace-registry.js'; import { - resolveRegisteredWorkspaceRuntimeByPathSelector, - resolveWorkspaceRuntimeFromParam, + resolveManagedWorkspaceRuntimeFromParam, + resolveManagedWorkspaceRuntimeByPathSelector, } from '../workspace-route-runtime.js'; import { ConnectionRegistry, @@ -39,9 +39,12 @@ import { SessionArchiveCoordinator } from '../server/session-archive.js'; import { RPC, error as rpcError, + isNotification, isRequest, + isResponse, logSafe, parseInbound, + type JsonRpcInbound, } from './json-rpc.js'; import { parseLastEventId } from '../sse-last-event-id.js'; import { @@ -70,6 +73,33 @@ export const ACP_SESSION_HEADER = 'acp-session-id'; /** Pathname of the Plan C CDP-tunnel endpoint (issue #5626). */ const CDP_PATH = '/cdp'; +function isActiveDrainCorrelation( + registry: ConnectionRegistry, + conn: AcpConnection, + message: JsonRpcInbound, +): boolean { + if (isResponse(message) && typeof message.id === 'string') { + const pending = registry.findPendingClientRequest(message.id); + return ( + pending !== undefined && + (pending.conn === conn || conn.ownsSession(pending.req.sessionId)) + ); + } + if ( + !(isRequest(message) || isNotification(message)) || + message.method !== 'session/cancel' + ) { + return false; + } + const params = + message.params !== null && typeof message.params === 'object' + ? (message.params as { sessionId?: unknown }) + : undefined; + return ( + typeof params?.sessionId === 'string' && conn.ownsSession(params.sessionId) + ); +} + /** Prefix/suffix of the Phase 4 workspace-qualified ACP WS path. */ const PLURAL_ACP_WS_PREFIX = '/workspaces/'; const PLURAL_ACP_WS_SUFFIX = '/acp'; @@ -437,6 +467,10 @@ interface RuntimeAcpMount { readonly rateLimitScope: string; readonly registry: ConnectionRegistry; readonly dispatcher: AcpDispatcher; + readonly workspaceRememberLane: WorkspaceRememberTaskLane; + readonly webSockets: Set; + readonly pendingWebSockets: Set; + draining: boolean; readonly ensureChromeDevToolsMcpRegistered: ( localPort: number | undefined, originatorClientId: string, @@ -492,6 +526,15 @@ export interface AcpHttpHandle { * connections, not just the primary's. */ getSnapshot(): AcpHttpSnapshot; + beginWorkspaceDrain(workspaceId: string): void; + cancelWorkspaceDrain(workspaceId: string): void; + getWorkspaceActivity(workspaceId: string): { + acpConnections: number; + memoryTasks: number; + }; + /** Commit memory teardown while sockets remain open for terminal events. */ + commitWorkspaceRemoval(workspaceId: string): void; + disposeWorkspace(workspaceId: string): void; /** Attach HTTP server post-listen to enable WebSocket upgrade. */ attachServer(server: import('node:http').Server): void; } @@ -529,6 +572,20 @@ export function mountAcpHttp( }); return true; }; + const rejectIfUnavailable = ( + mount: RuntimeAcpMount, + res: Response, + ): boolean => { + if (rejectIfDisposed(res)) return true; + if (!mount.draining) return false; + res.set('Retry-After', '5'); + res.status(503).json({ + error: 'Workspace runtime is being removed', + code: 'workspace_draining', + workspaceCwd: mount.workspaceCwd, + }); + return true; + }; // When a session/connection tears down with a permission still pending, // cancel it on the bridge so the agent's prompt isn't left blocked. const registry = new ConnectionRegistry( @@ -703,6 +760,10 @@ export function mountAcpHttp( rateLimitScope: opts.workspaceRegistry?.primary.workspaceId ?? 'primary', registry, dispatcher, + workspaceRememberLane: opts.workspaceRememberLane, + webSockets: new Set(), + pendingWebSockets: new Set(), + draining: false, ensureChromeDevToolsMcpRegistered, removeChromeDevToolsMcpIfUnused, clientMcpProviderFactory: opts.clientMcpProviderFactory, @@ -715,7 +776,7 @@ export function mountAcpHttp( req: Request, res: Response, ): void => { - if (rejectIfDisposed(res)) return; + if (rejectIfUnavailable(mount, res)) return; const connectionId = headerOf(req, ACP_CONNECTION_HEADER); if (!connectionId) { res.status(400).json({ error: 'Missing Acp-Connection-Id' }); @@ -765,6 +826,18 @@ export function mountAcpHttp( return; } const message = parsed.message; + const cleanupMessage = + (isRequest(message) || isNotification(message)) && + message.method === 'session/cancel'; + if (mount.draining && !isResponse(message) && !cleanupMessage) { + res.set('Retry-After', '5'); + res.status(503).json({ + error: 'Workspace runtime is being removed', + code: 'workspace_draining', + workspaceCwd: mount.workspaceCwd, + }); + return; + } // `initialize` mints a connection and replies inline (200 + JSON). if (isRequest(message) && message.method === 'initialize') { @@ -835,6 +908,17 @@ export function mountAcpHttp( ); return; } + if (mount.draining) { + if (!isActiveDrainCorrelation(mount.registry, conn, message)) { + res.set('Retry-After', '5'); + res.status(503).json({ + error: 'Workspace runtime is being removed', + code: 'workspace_draining', + workspaceCwd: mount.workspaceCwd, + }); + return; + } + } // Rate limit ACP HTTP POST (mirrors the WS checkRate path). if (opts.checkRate && isRequest(message)) { @@ -893,7 +977,7 @@ export function mountAcpHttp( req: Request, res: Response, ): void => { - if (rejectIfDisposed(res)) return; + if (rejectIfUnavailable(mount, res)) return; // RFD: Accept MUST include text/event-stream; otherwise 406. const accept = req.headers['accept'] ?? ''; if (!accept.includes('text/event-stream')) { @@ -1103,11 +1187,15 @@ export function mountAcpHttp( }, opts.maxConnections, ); + const workspaceRememberLane = new WorkspaceRememberTaskLane( + rt.bridge, + rt.workspaceCwd, + ); const secondaryDispatcher = new AcpDispatcher( rt.bridge, rt.workspaceCwd, rt.workspaceService, - new WorkspaceRememberTaskLane(rt.bridge), + workspaceRememberLane, rt.routeFileSystemFactory, // Phase 4: secondary mounts share the daemon-global device-flow registry // (single instance per daemon; OAuth credentials are global state). The @@ -1127,6 +1215,10 @@ export function mountAcpHttp( rateLimitScope: rt.workspaceId, registry: secondaryRegistry, dispatcher: secondaryDispatcher, + workspaceRememberLane, + webSockets: new Set(), + pendingWebSockets: new Set(), + draining: false, ensureChromeDevToolsMcpRegistered: () => {}, removeChromeDevToolsMcpIfUnused: () => {}, // Reverse client-MCP over WS is per-runtime: a connection on this @@ -1146,6 +1238,7 @@ export function mountAcpHttp( }; const secondaryMounts = new Map(); + const drainingWorkspaceIds = new Set(); const getOrCreateSecondaryMount = ( rt: WorkspaceRuntime, ): RuntimeAcpMount | undefined => { @@ -1153,6 +1246,10 @@ export function mountAcpHttp( const existing = secondaryMounts.get(rt.workspaceId); if (existing) return existing; const mount = createSecondaryAcpMount(rt); + if (drainingWorkspaceIds.has(rt.workspaceId)) { + mount.draining = true; + mount.workspaceRememberLane.beginDrain(); + } secondaryMounts.set(rt.workspaceId, mount); return mount; }; @@ -1182,7 +1279,11 @@ export function mountAcpHttp( }); return null; } - const rt = resolveWorkspaceRuntimeFromParam(workspaceRegistry, req, res); + const rt = resolveManagedWorkspaceRuntimeFromParam( + workspaceRegistry, + req, + res, + ); if (!rt) return null; if (!rt.primary && !rt.trusted) { res.status(403).json({ @@ -1423,11 +1524,8 @@ export function mountAcpHttp( } const wsRegistry = opts.workspaceRegistry; const rt = wsRegistry - ? (wsRegistry.getByWorkspaceId(selector) ?? - resolveRegisteredWorkspaceRuntimeByPathSelector( - wsRegistry, - selector, - )) + ? (wsRegistry.getManagedByWorkspaceId(selector) ?? + resolveManagedWorkspaceRuntimeByPathSelector(wsRegistry, selector)) : undefined; if (!rt) { logReject(`workspace-mismatch ${logSafe(selector)}`); @@ -1453,6 +1551,15 @@ export function mountAcpHttp( activeMount = resolvedMount; } + if (activeMount.draining) { + logReject(`workspace-draining ${activeMount.routeLabel}`); + socket.write( + 'HTTP/1.1 503 Service Unavailable\r\nRetry-After: 5\r\n\r\n', + ); + socket.destroy(); + return; + } + wss!.handleUpgrade(req, socket, head, (ws: WebSocket) => { if (disposed) { ws.close(1012, 'Server shutting down'); @@ -1464,6 +1571,8 @@ export function mountAcpHttp( extraRoute.onConnection(ws, req); return; } + activeMount.webSockets.add(ws); + activeMount.pendingWebSockets.add(ws); let initialized = false; const initTimer = setTimeout(() => { if (!initialized) { @@ -1504,6 +1613,8 @@ export function mountAcpHttp( // the socket goes away. WsStream's onClose handles ACP teardown. ws.on('close', () => { clearTimeout(initTimer); + activeMount.webSockets.delete(ws); + activeMount.pendingWebSockets.delete(ws); if (clientMcp) { void clientMcp.dispose('WS closed').catch(() => {}); clientMcp = undefined; @@ -1562,6 +1673,59 @@ export function mountAcpHttp( return; } + const frameType = + parsed !== null && typeof parsed === 'object' + ? (parsed as { type?: unknown }).type + : undefined; + if ( + activeMount.draining && + frameType === 'mcp_message' && + clientMcp !== undefined && + parsed !== null && + typeof parsed === 'object' + ) { + // Only replies to requests pending on this connection resolve here. + const result = await clientMcp.handleFrame( + parsed as Record, + ); + if (result.kind === 'message_resolved') return; + } + if ( + activeMount.draining && + !( + connRef !== undefined && + (isResponse(parsed) || + isRequest(parsed) || + isNotification(parsed)) && + isActiveDrainCorrelation(activeMount.registry, connRef, parsed) + ) + ) { + if ( + opts.cdpTunnelOverWs === true && + cdpEndpoint !== undefined && + parsed !== null && + typeof parsed === 'object' && + isCdpInboundFrameType(frameType) && + cdpEndpoint.routeInbound(parsed as Record) + ) { + return; + } + ws.send( + JSON.stringify( + rpcError( + isRequest(parsed) ? parsed.id : null, + RPC.INTERNAL_ERROR, + 'Workspace runtime is being removed', + { + code: 'workspace_draining', + workspaceCwd: activeMount.workspaceCwd, + }, + ), + ), + ); + return; + } + // ── Client-hosted MCP frames (issue #5626) ─────────────────── // These are NOT JSON-RPC envelopes — they carry a `type` // discriminator (`mcp_register` / `mcp_message` / `mcp_unregister`) @@ -1821,6 +1985,7 @@ export function mountAcpHttp( ); initialized = true; + activeMount.pendingWebSockets.delete(ws); clearTimeout(initTimer); connRef = conn; writeStderrLine( @@ -1980,11 +2145,14 @@ export function mountAcpHttp( upgradeServer = undefined; } registry.dispose(); + opts.workspaceRememberLane.dispose(); // Phase 4: dispose every non-primary runtime's ACP connection registry too, // so their sweep timers + live connections are torn down on shutdown. for (const mount of secondaryMounts.values()) { + mount.workspaceRememberLane.dispose(); mount.registry.dispose(); } + drainingWorkspaceIds.clear(); if (wss) { for (const client of wss.clients) { client.close(1012, 'Server shutting down'); @@ -1994,6 +2162,48 @@ export function mountAcpHttp( } }, registry, + beginWorkspaceDrain: (workspaceId) => { + drainingWorkspaceIds.add(workspaceId); + const mount = secondaryMounts.get(workspaceId); + if (!mount) return; + mount.draining = true; + mount.workspaceRememberLane.beginDrain(); + }, + cancelWorkspaceDrain: (workspaceId) => { + drainingWorkspaceIds.delete(workspaceId); + const mount = secondaryMounts.get(workspaceId); + if (!mount) return; + mount.draining = false; + mount.workspaceRememberLane.cancelDrain(); + }, + getWorkspaceActivity: (workspaceId) => { + const mount = secondaryMounts.get(workspaceId); + return { + acpConnections: + (mount?.registry.size ?? 0) + (mount?.pendingWebSockets.size ?? 0), + memoryTasks: mount?.workspaceRememberLane.pendingCount() ?? 0, + }; + }, + commitWorkspaceRemoval: (workspaceId) => { + secondaryMounts.get(workspaceId)?.workspaceRememberLane.dispose(); + }, + disposeWorkspace: (workspaceId) => { + drainingWorkspaceIds.delete(workspaceId); + const mount = secondaryMounts.get(workspaceId); + if (!mount) return; + secondaryMounts.delete(workspaceId); + try { + mount.workspaceRememberLane.dispose(); + } finally { + try { + for (const ws of mount.webSockets) { + ws.close(1012, 'Workspace removed'); + } + } finally { + mount.registry.dispose(); + } + } + }, getSnapshot: () => { const perMount = [ { diff --git a/packages/cli/src/serve/acp-http/workspace-qualified-acp.test.ts b/packages/cli/src/serve/acp-http/workspace-qualified-acp.test.ts index 34fe274dad4..c4c09ab716f 100644 --- a/packages/cli/src/serve/acp-http/workspace-qualified-acp.test.ts +++ b/packages/cli/src/serve/acp-http/workspace-qualified-acp.test.ts @@ -122,13 +122,22 @@ describe('workspace-qualified ACP (/workspaces/:workspace/acp)', () => { let checkRate: ReturnType; let primaryBridge: HttpAcpBridge; let secondaryBridge: HttpAcpBridge; + let workspaceRegistry: ReturnType; + let secondaryRuntime: WorkspaceRuntime; beforeEach(async () => { primaryBridge = makeBridge(); secondaryBridge = makeBridge(); const untrustedBridge = makeBridge(); - const registry = createWorkspaceRegistry([ + secondaryRuntime = makeRuntime({ + id: 'secondary-id', + cwd: '/ws-b', + primary: false, + trusted: true, + bridge: secondaryBridge, + }); + workspaceRegistry = createWorkspaceRegistry([ makeRuntime({ id: 'primary-id', cwd: '/ws', @@ -136,13 +145,7 @@ describe('workspace-qualified ACP (/workspaces/:workspace/acp)', () => { trusted: true, bridge: primaryBridge, }), - makeRuntime({ - id: 'secondary-id', - cwd: '/ws-b', - primary: false, - trusted: true, - bridge: secondaryBridge, - }), + secondaryRuntime, makeRuntime({ id: 'untrusted-id', cwd: '/ws-c', @@ -167,7 +170,7 @@ describe('workspace-qualified ACP (/workspaces/:workspace/acp)', () => { boundWorkspace: '/ws', workspace: {} as unknown as DaemonWorkspaceService, enabled: true, - workspaceRegistry: registry, + workspaceRegistry, deviceFlowRegistry, cdpTunnelOverWs: true, cdpTunnelRegistry: cdpRegistry, @@ -258,6 +261,9 @@ describe('workspace-qualified ACP (/workspaces/:workspace/acp)', () => { _meta?: { qwen?: { workspaceCwd?: string } }; }; }; + error?: { + data?: { code?: string; workspaceCwd?: string }; + }; }> { return new Promise((resolve, reject) => { const ws = new WebSocket(`ws://127.0.0.1:${port}${pathname}`, { @@ -531,6 +537,21 @@ describe('workspace-qualified ACP (/workspaces/:workspace/acp)', () => { bridge: makeBridge(), }), ); + singleHandle.beginWorkspaceDrain('dynamic-id'); + const drainingDynamic = await fetch( + `http://127.0.0.1:${singlePort}/workspaces/dynamic-id/acp`, + { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: INITIALIZE, + }, + ); + expect(drainingDynamic.status).toBe(503); + await expect(drainingDynamic.json()).resolves.toMatchObject({ + code: 'workspace_draining', + }); + + singleHandle.cancelWorkspaceDrain('dynamic-id'); const dynamic = await fetch( `http://127.0.0.1:${singlePort}/workspaces/dynamic-id/acp`, { @@ -623,6 +644,102 @@ describe('workspace-qualified ACP (/workspaces/:workspace/acp)', () => { } }); + it('counts and closes an uninitialized workspace WebSocket on removal', async () => { + const ws = new WebSocket( + `ws://127.0.0.1:${port}/workspaces/secondary-id/acp`, + { handshakeTimeout: 2000 }, + ); + await new Promise((resolve, reject) => { + ws.once('open', resolve); + ws.once('error', reject); + }); + expect(handle!.getWorkspaceActivity('secondary-id').acpConnections).toBe(1); + + const closed = new Promise((resolve, reject) => { + const timer = setTimeout( + () => + reject(new Error('Workspace WebSocket stayed open after removal')), + 2000, + ); + ws.once('close', () => { + clearTimeout(timer); + resolve(); + }); + }); + handle!.beginWorkspaceDrain('secondary-id'); + handle!.commitWorkspaceRemoval('secondary-id'); + handle!.disposeWorkspace('secondary-id'); + await closed; + + expect(handle!.getWorkspaceActivity('secondary-id').acpConnections).toBe(0); + await expect(initializeWs('/acp')).resolves.toMatchObject({ + result: { + agentCapabilities: { _meta: { qwen: { workspaceCwd: '/ws' } } }, + }, + }); + }); + + it('disposes only the target live WebSocket and allows a fresh mount', async () => { + const connect = (pathname: string) => + new Promise((resolve, reject) => { + const ws = new WebSocket(`ws://127.0.0.1:${port}${pathname}`, { + handshakeTimeout: 2000, + }); + ws.on('open', () => ws.send(INITIALIZE)); + ws.on('message', (data: WebSocket.RawData) => { + const message = JSON.parse(data.toString()) as { id?: number }; + if (message.id === 1) resolve(ws); + }); + ws.on('error', reject); + }); + const primaryWs = await connect('/acp'); + const secondaryWs = await connect('/workspaces/secondary-id/acp'); + const secondaryClosed = new Promise((resolve) => { + secondaryWs.once('close', () => resolve()); + }); + + expect(workspaceRegistry.beginDrain(secondaryRuntime)).toBe(true); + handle!.beginWorkspaceDrain('secondary-id'); + handle!.commitWorkspaceRemoval('secondary-id'); + handle!.disposeWorkspace('secondary-id'); + workspaceRegistry.completeDrain(secondaryRuntime); + await secondaryClosed; + + const primaryReply = new Promise>( + (resolve, reject) => { + primaryWs.once('message', (data: WebSocket.RawData) => { + try { + resolve(JSON.parse(data.toString()) as Record); + } catch (err) { + reject(err as Error); + } + }); + }, + ); + primaryWs.send( + JSON.stringify({ jsonrpc: '2.0', id: 2, method: 'unknown/mutation' }), + ); + expect(await primaryReply).toMatchObject({ id: 2 }); + primaryWs.close(); + + workspaceRegistry.add( + makeRuntime({ + id: 'secondary-id', + cwd: '/ws-b', + primary: false, + trusted: true, + bridge: makeBridge(), + }), + ); + await expect( + initializeWs('/workspaces/secondary-id/acp'), + ).resolves.toMatchObject({ + result: { + agentCapabilities: { _meta: { qwen: { workspaceCwd: '/ws-b' } } }, + }, + }); + }); + it('rejects an upgrade whose listener starts after disposal', async () => { server.prependOnceListener('upgrade', () => handle!.dispose()); @@ -677,6 +794,192 @@ describe('workspace-qualified ACP (/workspaces/:workspace/acp)', () => { ]); }); + it('drains, rolls back, disposes, and recreates a secondary mount', async () => { + const initialized = await postInitialize('/workspaces/secondary-id/acp'); + expect(initialized.status).toBe(200); + expect(handle!.getWorkspaceActivity('secondary-id')).toEqual({ + acpConnections: 1, + memoryTasks: 0, + }); + + handle!.beginWorkspaceDrain('secondary-id'); + const draining = await postInitialize('/workspaces/secondary-id/acp'); + expect(draining.status).toBe(503); + expect(draining.headers.get('retry-after')).toBe('5'); + await expect(draining.json()).resolves.toMatchObject({ + code: 'workspace_draining', + }); + + handle!.cancelWorkspaceDrain('secondary-id'); + expect((await postInitialize('/workspaces/secondary-id/acp')).status).toBe( + 200, + ); + + expect(workspaceRegistry.beginDrain(secondaryRuntime)).toBe(true); + handle!.beginWorkspaceDrain('secondary-id'); + const registryDraining = await postInitialize( + '/workspaces/secondary-id/acp', + ); + expect(registryDraining.status).toBe(503); + await expect(registryDraining.json()).resolves.toMatchObject({ + code: 'workspace_draining', + }); + handle!.commitWorkspaceRemoval('secondary-id'); + handle!.disposeWorkspace('secondary-id'); + workspaceRegistry.completeDrain(secondaryRuntime); + expect(handle!.getWorkspaceActivity('secondary-id')).toEqual({ + acpConnections: 0, + memoryTasks: 0, + }); + expect( + handle! + .getSnapshot() + .mounts.some((mount) => mount.workspaceId === 'secondary-id'), + ).toBe(false); + + const replacementBridge = makeBridge(); + workspaceRegistry.add( + makeRuntime({ + id: 'secondary-id', + cwd: '/ws-b', + primary: false, + trusted: true, + bridge: replacementBridge, + }), + ); + expect((await postInitialize('/workspaces/secondary-id/acp')).status).toBe( + 200, + ); + expect(handle!.getWorkspaceActivity('secondary-id').acpConnections).toBe(1); + }); + + it('returns a structured workspace_draining error on an existing WebSocket', async () => { + const reply = await new Promise>( + (resolve, reject) => { + const ws = new WebSocket( + `ws://127.0.0.1:${port}/workspaces/secondary-id/acp`, + { handshakeTimeout: 2000 }, + ); + ws.on('open', () => ws.send(INITIALIZE)); + ws.on('message', (data: WebSocket.RawData) => { + const message = JSON.parse(data.toString()) as Record< + string, + unknown + >; + if (message['id'] === 1) { + handle!.beginWorkspaceDrain('secondary-id'); + ws.send( + JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'unknown/mutation', + }), + ); + return; + } + if (message['id'] === 2) { + ws.close(); + resolve(message); + } + }); + ws.on('error', reject); + }, + ); + + expect(reply).toMatchObject({ + error: { + data: { + code: 'workspace_draining', + workspaceCwd: '/ws-b', + }, + }, + }); + }); + + it('rejects a new WebSocket upgrade while its workspace is draining', async () => { + expect(workspaceRegistry.beginDrain(secondaryRuntime)).toBe(true); + handle!.beginWorkspaceDrain('secondary-id'); + + const response = await new Promise<{ + status: number | undefined; + retryAfter: string | string[] | undefined; + }>((resolve, reject) => { + const ws = new WebSocket( + `ws://127.0.0.1:${port}/workspaces/secondary-id/acp`, + { handshakeTimeout: 2000 }, + ); + ws.on('unexpected-response', (_req, res) => { + resolve({ + status: res.statusCode, + retryAfter: res.headers['retry-after'], + }); + ws.terminate(); + }); + ws.on('open', () => { + ws.close(); + reject(new Error('draining workspace WS upgrade should not open')); + }); + ws.on('error', reject); + }); + + expect(response).toEqual({ status: 503, retryAfter: '5' }); + }); + + it('rejects unowned and spoofed correlation frames during drain', async () => { + const replies = await new Promise>>( + (resolve, reject) => { + const received: Array> = []; + const ws = new WebSocket( + `ws://127.0.0.1:${port}/workspaces/secondary-id/acp`, + { handshakeTimeout: 2000 }, + ); + ws.on('open', () => ws.send(INITIALIZE)); + ws.on('message', (data: WebSocket.RawData) => { + const message = JSON.parse(data.toString()) as Record< + string, + unknown + >; + if (message['id'] === 1) { + handle!.beginWorkspaceDrain('secondary-id'); + ws.send(JSON.stringify({ jsonrpc: '2.0', id: 99, result: {} })); + ws.send( + JSON.stringify({ type: 'cdp_result', requestId: 'unknown' }), + ); + ws.send( + JSON.stringify({ + type: 'mcp_message', + id: 'unknown', + server: 'missing', + payload: {}, + }), + ); + ws.send( + JSON.stringify({ + jsonrpc: '2.0', + id: 2, + method: 'unknown/mutation', + }), + ); + return; + } + received.push(message); + if (message['id'] === 2) { + ws.close(); + resolve(received); + } + }); + ws.on('error', reject); + }, + ); + + expect(replies).toHaveLength(4); + for (const reply of replies) { + expect(reply).toMatchObject({ + error: { data: { code: 'workspace_draining' } }, + }); + } + }); + it('rejects a raw WS upgrade whose selector is a dot-segment (%2e%2e)', async () => { // `ws` normalizes the client URL (/workspaces/%2e%2e/acp -> /acp), so the // real attack surface — a raw, non-normalized request-target — must be diff --git a/packages/cli/src/serve/acp-session-bridge.ts b/packages/cli/src/serve/acp-session-bridge.ts index 98346bd40e2..5c4d865de26 100644 --- a/packages/cli/src/serve/acp-session-bridge.ts +++ b/packages/cli/src/serve/acp-session-bridge.ts @@ -90,6 +90,7 @@ export type { BridgeDaemonStatusLimits, BridgeDaemonSessionDiagnostic, BridgeDaemonStatusSnapshot, + BridgeShutdownOptions, AcpSessionBridge, HttpAcpBridge, } from '@qwen-code/acp-bridge/bridgeTypes'; @@ -116,6 +117,7 @@ export { McpServerNotFoundError, McpServerRestartFailedError, SessionBusyError, + WorkspaceDrainingError, InvalidRewindTargetError, TotalSessionLimitExceededError, NOT_CURRENTLY_GENERATING_CANCEL_MESSAGE, diff --git a/packages/cli/src/serve/capabilities.ts b/packages/cli/src/serve/capabilities.ts index c45956818c5..324f9f6f778 100644 --- a/packages/cli/src/serve/capabilities.ts +++ b/packages/cli/src/serve/capabilities.ts @@ -274,6 +274,7 @@ export const SERVE_CAPABILITY_REGISTRY = { // only when one daemon hosts more than one registered workspace runtime. multi_workspace_sessions: { since: 'v1' }, persistent_workspace_registration: { since: 'v1' }, + workspace_runtime_removal: { since: 'v1' }, // Workspace-qualified core REST routes under `/workspaces/:workspace/...`. // Covers core file/status/permissions/trust/lifecycle/MCP/tool, memory, // workspace agent CRUD, and persisted session organization surfaces. @@ -369,6 +370,7 @@ export interface AdvertiseFeatureToggles { voiceWsAvailable?: boolean; multiWorkspaceSessionsEnabled?: boolean; persistentWorkspaceRegistrationAvailable?: boolean; + workspaceRuntimeRemovalAvailable?: boolean; /** * Whether the HTTP ACP surface is enabled (default on; opts out via * QWEN_SERVE_ACP_HTTP=0). Workspace-qualified ACP is only advertised when on. @@ -454,6 +456,10 @@ export const CONDITIONAL_SERVE_FEATURES: ReadonlyMap< 'persistent_workspace_registration', (toggles) => toggles.persistentWorkspaceRegistrationAvailable === true, ], + [ + 'workspace_runtime_removal', + (toggles) => toggles.workspaceRuntimeRemovalAvailable === true, + ], [ 'workspace_qualified_acp', // The plural routes are pre-mounted for workspaces registered after app diff --git a/packages/cli/src/serve/channel-worker-group.test.ts b/packages/cli/src/serve/channel-worker-group.test.ts index 9914889e47e..8f4b7facad2 100644 --- a/packages/cli/src/serve/channel-worker-group.test.ts +++ b/packages/cli/src/serve/channel-worker-group.test.ts @@ -13,6 +13,7 @@ import type { CreateChannelWorkerSupervisorOptions, } from './channel-worker-supervisor.js'; import { ChannelWorkerStopError } from './channel-worker-supervisor.js'; +import type { ChannelWorkspaceGroup } from './channel-workspace-grouping.js'; import type { WorkspaceRegistry, WorkspaceRuntime, @@ -45,16 +46,24 @@ function fakeRegistry(runtimes: WorkspaceRuntime[]): WorkspaceRegistry { return { primary: runtimes.find((runtime) => runtime.primary)!, list: () => runtimes, + listManaged: () => runtimes, add: vi.fn(), getByWorkspaceCwd: (cwd) => runtimes.find((runtime) => runtime.workspaceCwd === cwd), getByWorkspaceId: (id) => runtimes.find((runtime) => runtime.workspaceId === id), + getManagedByWorkspaceCwd: (cwd) => + runtimes.find((runtime) => runtime.workspaceCwd === cwd), + getManagedByWorkspaceId: (id) => + runtimes.find((runtime) => runtime.workspaceId === id), resolveWorkspaceCwd: (cwd) => cwd === undefined ? runtimes.find((runtime) => runtime.primary) : runtimes.find((runtime) => runtime.workspaceCwd === cwd), resolveLiveSessionOwner: () => ({ kind: 'not_found' }), + beginDrain: vi.fn(() => true), + cancelDrain: vi.fn(), + completeDrain: vi.fn(), }; } @@ -145,6 +154,118 @@ describe('createChannelWorkerGroup', () => { ); }); + it('drains and removes only the target workspace worker', async () => { + const registry = fakeRegistry([ + fakeRuntime(PRIMARY, true), + fakeRuntime(SECONDARY, false), + ]); + const { createSupervisor, recorded } = makeCreateSupervisor(() => + snapshot({}), + ); + const groups: ChannelWorkspaceGroup[] = [ + { workspaceCwd: PRIMARY, selection: { mode: 'names', names: ['a'] } }, + { workspaceCwd: SECONDARY, selection: { mode: 'names', names: ['b'] } }, + ]; + const group = createChannelWorkerGroup({ + groups, + registry, + createSupervisor, + shared, + }); + + expect(group.workspaceActivity(SECONDARY)).toBe(1); + group.beginWorkspaceDrain(SECONDARY); + await expect(group.enqueueWebhookTask(webhookTask)).rejects.toMatchObject({ + code: 'channel_worker_unavailable', + }); + await expect(group.reconcile(groups, { force: true })).rejects.toThrow( + 'cannot change while a workspace is draining', + ); + expect(recorded).toHaveLength(2); + expect(recorded[1]!.supervisor.stop).not.toHaveBeenCalled(); + expect(recorded[1]!.supervisor.start).not.toHaveBeenCalled(); + + group.cancelWorkspaceDrain(SECONDARY); + recorded[1]!.supervisor.enqueueWebhookTask.mockResolvedValueOnce({ + accepted: true, + }); + await expect(group.enqueueWebhookTask(webhookTask)).resolves.toEqual({ + accepted: true, + }); + + group.beginWorkspaceDrain(SECONDARY); + const firstRemoval = group.removeWorkspace(SECONDARY); + const secondRemoval = group.removeWorkspace(SECONDARY); + expect(firstRemoval).toBe(secondRemoval); + await firstRemoval; + + expect(recorded[1]!.supervisor.stop).toHaveBeenCalledOnce(); + expect(recorded[0]!.supervisor.stop).not.toHaveBeenCalled(); + expect(group.workspaceActivity(SECONDARY)).toBe(0); + expect(group.snapshots()).toEqual([ + expect.objectContaining({ workspaceCwd: PRIMARY }), + ]); + await expect(group.enqueueWebhookTask(webhookTask)).rejects.toMatchObject({ + code: 'channel_worker_unavailable', + }); + + await group.stop(); + expect(recorded[0]!.supervisor.stop).toHaveBeenCalledOnce(); + expect(recorded[1]!.supervisor.stop).toHaveBeenCalledOnce(); + }); + + it('counts a scheduled restart as workspace activity', () => { + const registry = fakeRegistry([fakeRuntime(SECONDARY, false)]); + const { createSupervisor } = makeCreateSupervisor(() => + snapshot({ + state: 'failed', + nextRestartAt: new Date(Date.now() + 1_000).toISOString(), + }), + ); + const group = createChannelWorkerGroup({ + groups: [ + { + workspaceCwd: SECONDARY, + selection: { mode: 'names', names: ['telegram'] }, + }, + ], + registry, + createSupervisor, + shared, + }); + + expect(group.workspaceActivity(SECONDARY)).toBe(1); + }); + + it('falls back to synchronous kill when target worker stop fails', async () => { + const registry = fakeRegistry([ + fakeRuntime(PRIMARY, true), + fakeRuntime(SECONDARY, false), + ]); + const { createSupervisor, recorded } = makeCreateSupervisor(() => + snapshot({}), + ); + const group = createChannelWorkerGroup({ + groups: [ + { workspaceCwd: PRIMARY, selection: { mode: 'names', names: ['a'] } }, + { workspaceCwd: SECONDARY, selection: { mode: 'names', names: ['b'] } }, + ], + registry, + createSupervisor, + shared, + }); + recorded[1]!.supervisor.stop.mockRejectedValueOnce( + new Error('stop failed'), + ); + + await expect(group.removeWorkspace(SECONDARY)).resolves.toBeUndefined(); + + expect(recorded[1]!.supervisor.killAllSync).toHaveBeenCalledOnce(); + expect(group.snapshots()).toEqual([ + expect.objectContaining({ workspaceCwd: PRIMARY }), + ]); + }); + it('routes --channel all webhook tasks to the primary supervisor', async () => { const registry = fakeRegistry([ fakeRuntime(PRIMARY, true), @@ -300,6 +421,164 @@ describe('createChannelWorkerGroup', () => { expect(onStateChange).toHaveBeenCalledTimes(2); }); + it('stops supervisors again after the group is restarted', async () => { + const registry = fakeRegistry([fakeRuntime(PRIMARY, true)]); + const { createSupervisor, recorded } = makeCreateSupervisor(() => + snapshot({}), + ); + const group = createChannelWorkerGroup({ + groups: [{ workspaceCwd: PRIMARY, selection: { mode: 'all' } }], + registry, + createSupervisor, + shared, + }); + + await group.start(); + await group.stop(); + await group.start(); + await group.stop(); + + expect(recorded[0]!.supervisor.start).toHaveBeenCalledTimes(2); + expect(recorded[0]!.supervisor.stop).toHaveBeenCalledTimes(2); + }); + + it('creates a fresh worker and restores webhook routing after re-add', async () => { + const runtimes = [ + fakeRuntime(PRIMARY, true), + fakeRuntime(SECONDARY, false, { VERSION: 'old' }), + ]; + const registry = fakeRegistry(runtimes); + const { createSupervisor, recorded } = makeCreateSupervisor(() => + snapshot({}), + ); + const group = createChannelWorkerGroup({ + groups: [ + { workspaceCwd: PRIMARY, selection: { mode: 'names', names: ['a'] } }, + ], + registry, + createSupervisor, + shared, + }); + await group.start(); + await group.reconcile([ + { workspaceCwd: PRIMARY, selection: { mode: 'names', names: ['a'] } }, + { workspaceCwd: SECONDARY, selection: { mode: 'names', names: ['b'] } }, + ]); + await group.removeWorkspace(SECONDARY); + runtimes.splice(1, 1, fakeRuntime(SECONDARY, false, { VERSION: 'new' })); + + await group.restoreWorkspace(SECONDARY); + recorded[2]!.supervisor.enqueueWebhookTask.mockResolvedValueOnce({ + accepted: true, + }); + + expect(recorded).toHaveLength(3); + expect(recorded[2]!.opts.workerBaseEnv).toEqual({ VERSION: 'new' }); + expect(recorded[2]!.supervisor.start).toHaveBeenCalledOnce(); + await expect(group.enqueueWebhookTask(webhookTask)).resolves.toEqual({ + accepted: true, + }); + await group.removeWorkspace(SECONDARY); + expect(recorded[2]!.supervisor.stop).toHaveBeenCalledOnce(); + }); + + it('removes a restored worker from routing when startup fails', async () => { + const runtimes = [ + fakeRuntime(PRIMARY, true), + fakeRuntime(SECONDARY, false), + ]; + const registry = fakeRegistry(runtimes); + const { createSupervisor, recorded } = makeCreateSupervisor(() => + snapshot({}), + ); + const createSupervisorWithRestoreFailure = ( + opts: CreateChannelWorkerSupervisorOptions, + ) => { + const supervisor = createSupervisor(opts); + if (recorded.length === 3) { + supervisor.start.mockRejectedValueOnce(new Error('restore failed')); + } + return supervisor; + }; + const group = createChannelWorkerGroup({ + groups: [ + { workspaceCwd: PRIMARY, selection: { mode: 'names', names: ['a'] } }, + { workspaceCwd: SECONDARY, selection: { mode: 'names', names: ['b'] } }, + ], + registry, + createSupervisor: createSupervisorWithRestoreFailure, + shared, + }); + await group.start(); + await group.removeWorkspace(SECONDARY); + + await expect(group.restoreWorkspace(SECONDARY)).rejects.toThrow( + 'restore failed', + ); + + expect(group.workspaceActivity(SECONDARY)).toBe(0); + await expect(group.enqueueWebhookTask(webhookTask)).rejects.toMatchObject({ + code: 'channel_worker_unavailable', + }); + + await group.restoreWorkspace(SECONDARY); + expect(recorded).toHaveLength(4); + expect(recorded[3]!.supervisor.start).toHaveBeenCalledOnce(); + }); + + it('stops a restored worker whose start finishes after the group stops', async () => { + const runtimes = [ + fakeRuntime(PRIMARY, true), + fakeRuntime(SECONDARY, false), + ]; + const registry = fakeRegistry(runtimes); + const { createSupervisor, recorded } = makeCreateSupervisor(() => + snapshot({}), + ); + let releaseStart!: () => void; + const createSupervisorWithPendingRestore = ( + opts: CreateChannelWorkerSupervisorOptions, + ) => { + const supervisor = createSupervisor(opts); + if (recorded.length === 3) { + supervisor.start.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseStart = resolve; + }), + ); + } + return supervisor; + }; + const group = createChannelWorkerGroup({ + groups: [ + { workspaceCwd: PRIMARY, selection: { mode: 'names', names: ['a'] } }, + { workspaceCwd: SECONDARY, selection: { mode: 'names', names: ['b'] } }, + ], + registry, + createSupervisor: createSupervisorWithPendingRestore, + shared, + }); + await group.start(); + await group.removeWorkspace(SECONDARY); + + const restoring = group.restoreWorkspace(SECONDARY); + await vi.waitFor(() => { + expect(recorded[2]!.supervisor.start).toHaveBeenCalledOnce(); + }); + await group.stop(); + releaseStart(); + await restoring; + + expect(recorded[2]!.supervisor.stop).toHaveBeenCalledTimes(2); + expect(group.workspaceActivity(SECONDARY)).toBe(0); + expect(group.snapshots()).not.toEqual( + expect.arrayContaining([ + expect.objectContaining({ workspaceCwd: SECONDARY }), + ]), + ); + }); + it('does not start later supervisors when the first start fails', async () => { const registry = fakeRegistry([ fakeRuntime(PRIMARY, true), @@ -694,6 +973,53 @@ describe('createChannelWorkerGroup', () => { }); }); + it('does not commit a worker drained while reconcile is starting it', async () => { + const registry = fakeRegistry([ + fakeRuntime(PRIMARY, true), + fakeRuntime(SECONDARY, false), + ]); + let releaseStart!: () => void; + const test = makeCreateSupervisor(() => snapshot({})); + const createSupervisor = (opts: CreateChannelWorkerSupervisorOptions) => { + const supervisor = test.createSupervisor(opts); + if (opts.workspace === SECONDARY) { + supervisor.start.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseStart = resolve; + }), + ); + } + return supervisor; + }; + const group = createChannelWorkerGroup({ + groups: [ + { workspaceCwd: PRIMARY, selection: { mode: 'names', names: ['a'] } }, + ], + registry, + createSupervisor, + shared, + }); + await group.start(); + + const reconciling = group.reconcile([ + { workspaceCwd: PRIMARY, selection: { mode: 'names', names: ['a'] } }, + { workspaceCwd: SECONDARY, selection: { mode: 'names', names: ['b'] } }, + ]); + await vi.waitFor(() => expect(releaseStart).toBeTypeOf('function')); + group.beginWorkspaceDrain(SECONDARY); + releaseStart(); + await expect(reconciling).rejects.toThrow( + 'Workspace drained during channel worker reconcile.', + ); + + expect(test.recorded[1]!.supervisor.stop).toHaveBeenCalledOnce(); + expect(group.workspaceActivity(SECONDARY)).toBe(0); + await expect(group.enqueueWebhookTask(webhookTask)).rejects.toMatchObject({ + code: 'channel_worker_unavailable', + }); + }); + it('stops failed new workers and restores stopped old workers', async () => { const registry = fakeRegistry([fakeRuntime(PRIMARY, true)]); const recorded: RecordedSupervisor[] = []; diff --git a/packages/cli/src/serve/channel-worker-group.ts b/packages/cli/src/serve/channel-worker-group.ts index 45e5dc3a9f7..ecbb0aaaa1b 100644 --- a/packages/cli/src/serve/channel-worker-group.ts +++ b/packages/cli/src/serve/channel-worker-group.ts @@ -64,6 +64,11 @@ export interface ChannelWorkerGroup { snapshots(): ChannelWorkerGroupSnapshot[]; /** Primary workspace snapshot, backing the legacy single-worker fields. */ primarySnapshot(): ChannelWorkerSnapshot; + beginWorkspaceDrain(workspaceCwd: string): void; + cancelWorkspaceDrain(workspaceCwd: string): void; + workspaceActivity(workspaceCwd: string): number; + removeWorkspace(workspaceCwd: string): Promise; + restoreWorkspace(workspaceCwd: string): Promise; enqueueWebhookTask: ChannelWorkerSupervisor['enqueueWebhookTask']; } @@ -102,6 +107,7 @@ interface ChannelWorkerGroupEntry { selection: ChannelWorkspaceGroup['selection']; generation: number; supervisor: ChannelWorkerSupervisor; + stopPromise?: Promise; } function selectionsEqual( @@ -125,10 +131,16 @@ export function createChannelWorkerGroup( ): ChannelWorkerGroup { let generation = 0; let entries = new Map(); + const groupsByWorkspace = new Map( + opts.groups.map((group) => [group.workspaceCwd, group]), + ); const pendingEntries = new Set(); const pendingGenerations = new Map(); + const drainingWorkspaces = new Set(); + const removalPromises = new Map>(); let reconciling: Promise | undefined; let stopping = false; + let groupStarted = false; const withMeta = ( entry: ChannelWorkerGroupEntry, @@ -139,7 +151,6 @@ export function createChannelWorkerGroup( workspaceCwd: entry.workspaceCwd, primary: entry.primary, }); - const createEntry = ( group: ChannelWorkspaceGroup, ): ChannelWorkerGroupEntry => { @@ -260,12 +271,22 @@ export function createChannelWorkerGroup( withMeta(entry, entry.supervisor.snapshot()), ); - const stopEntry = async (entry: ChannelWorkerGroupEntry): Promise => { - try { - await entry.supervisor.stop(); - } finally { - opts.onStateChange?.(); + const stopEntry = (entry: ChannelWorkerGroupEntry): Promise => { + if (!entry.stopPromise) { + const stopPromise = entry.supervisor + .stop() + .finally(() => opts.onStateChange?.()); + entry.stopPromise = stopPromise; + void stopPromise.then( + () => { + if (entry.stopPromise === stopPromise) entry.stopPromise = undefined; + }, + () => { + if (entry.stopPromise === stopPromise) entry.stopPromise = undefined; + }, + ); } + return entry.stopPromise; }; const stopEntriesBestEffort = async ( @@ -338,6 +359,16 @@ export function createChannelWorkerGroup( return undefined; }; + const detachEntry = (entry: ChannelWorkerGroupEntry): void => { + if (entries.get(entry.workspaceCwd)?.generation === entry.generation) { + entries.delete(entry.workspaceCwd); + } + pendingEntries.delete(entry); + if (pendingGenerations.get(entry.workspaceCwd) === entry.generation) { + pendingGenerations.delete(entry.workspaceCwd); + } + }; + const group: ChannelWorkerGroup = { async start() { // Start sequentially so a failing initial launch can roll back every @@ -346,6 +377,9 @@ export function createChannelWorkerGroup( const started: ChannelWorkerGroupEntry[] = []; try { for (const entry of entries.values()) { + if (drainingWorkspaces.has(entry.workspaceCwd)) { + throw new Error('Workspace drained during channel worker startup.'); + } if (stopping) { throw new Error('Channel worker group stopped during startup.'); } @@ -354,14 +388,23 @@ export function createChannelWorkerGroup( if (stopping) { throw new Error('Channel worker group stopped during startup.'); } + if ( + entries.get(entry.workspaceCwd)?.generation !== entry.generation || + drainingWorkspaces.has(entry.workspaceCwd) + ) { + throw new Error('Workspace drained during channel worker startup.'); + } } + groupStarted = true; } catch (error) { + groupStarted = false; await stopEntriesBestEffort(started); throw error; } }, async stop() { stopping = true; + groupStarted = false; await reconciling?.catch(() => {}); await stopAllEntries([...entries.values()]); }, @@ -374,6 +417,14 @@ export function createChannelWorkerGroup( ), ); } + if (drainingWorkspaces.size > 0) { + return Promise.reject( + new ChannelWorkerReconcileError( + 'Channel worker configuration cannot change while a workspace is draining.', + { rolledBack: true }, + ), + ); + } if (reconciling) return reconciling; reconciling = (async () => { const targets = new Map( @@ -426,6 +477,11 @@ export function createChannelWorkerGroup( const startedNew: ChannelWorkerGroupEntry[] = []; try { for (const entry of newEntries) { + if (drainingWorkspaces.has(entry.workspaceCwd)) { + throw new Error( + 'Workspace drained during channel worker reconcile.', + ); + } if (stopping) { throw new Error('Channel worker group stopped during reconcile.'); } @@ -434,6 +490,16 @@ export function createChannelWorkerGroup( if (stopping) { throw new Error('Channel worker group stopped during reconcile.'); } + if (drainingWorkspaces.has(entry.workspaceCwd)) { + throw new Error( + 'Workspace drained during channel worker reconcile.', + ); + } + } + if (drainingWorkspaces.size > 0) { + throw new Error( + 'Workspace drained during channel worker reconcile.', + ); } } catch (error) { reconcileOptions?.onRollingBack?.(); @@ -461,6 +527,17 @@ export function createChannelWorkerGroup( committed.set(entry.workspaceCwd, entry); } entries = committed; + const targetWorkspaceCwds = new Set( + targetGroups.map((target) => target.workspaceCwd), + ); + for (const workspaceCwd of groupsByWorkspace.keys()) { + if (!targetWorkspaceCwds.has(workspaceCwd)) { + groupsByWorkspace.delete(workspaceCwd); + } + } + for (const target of targetGroups) { + groupsByWorkspace.set(target.workspaceCwd, target); + } return { changed: true, workers: entrySnapshots() }; })().finally(() => { pendingEntries.clear(); @@ -479,6 +556,7 @@ export function createChannelWorkerGroup( }, killAllSync() { stopping = true; + groupStarted = false; for (const entry of entries.values()) { entry.supervisor.killAllSync(); } @@ -491,9 +569,101 @@ export function createChannelWorkerGroup( const primary = [...entries.values()].find((entry) => entry.primary); return primary?.supervisor.snapshot() ?? { ...DISABLED_SNAPSHOT }; }, + beginWorkspaceDrain(workspaceCwd) { + drainingWorkspaces.add(workspaceCwd); + }, + cancelWorkspaceDrain(workspaceCwd) { + drainingWorkspaces.delete(workspaceCwd); + }, + workspaceActivity(workspaceCwd) { + if (pendingGenerations.has(workspaceCwd)) return 1; + const entry = entries.get(workspaceCwd); + if (!entry) return 0; + const snapshot = entry.supervisor.snapshot(); + return snapshot.state === 'starting' || + snapshot.state === 'running' || + snapshot.nextRestartAt !== undefined + ? 1 + : 0; + }, + removeWorkspace(workspaceCwd) { + const existing = removalPromises.get(workspaceCwd); + if (existing) return existing; + drainingWorkspaces.add(workspaceCwd); + const removal = (async () => { + try { + await reconciling?.catch(() => {}); + const entry = entries.get(workspaceCwd); + if (!entry) return; + let killError: unknown; + try { + await stopEntry(entry); + } catch { + try { + entry.supervisor.killAllSync(); + } catch (err) { + killError = err; + } + } finally { + detachEntry(entry); + } + if (killError) throw killError; + } finally { + drainingWorkspaces.delete(workspaceCwd); + } + })(); + removalPromises.set(workspaceCwd, removal); + void removal.then( + () => { + if (removalPromises.get(workspaceCwd) === removal) { + removalPromises.delete(workspaceCwd); + } + }, + () => { + if (removalPromises.get(workspaceCwd) === removal) { + removalPromises.delete(workspaceCwd); + } + }, + ); + return removal; + }, + async restoreWorkspace(workspaceCwd) { + if (entries.has(workspaceCwd)) return; + const target = groupsByWorkspace.get(workspaceCwd); + if (!target) return; + const entry = createEntry(target); + entries.set(workspaceCwd, entry); + if (!groupStarted || stopping) { + detachEntry(entry); + return; + } + try { + await entry.supervisor.start(); + if (stopping || !groupStarted) { + await stopEntry(entry).catch(() => { + try { + entry.supervisor.killAllSync(); + } catch { + // Best-effort cleanup after the group stopped concurrently. + } + }); + detachEntry(entry); + } + } catch (err) { + await stopEntry(entry).catch(() => { + try { + entry.supervisor.killAllSync(); + } catch { + // Preserve the start failure that caused the rollback. + } + }); + detachEntry(entry); + throw err; + } + }, async enqueueWebhookTask(task) { const entry = routeEntry(task.channelName); - if (!entry) { + if (!entry || drainingWorkspaces.has(entry.workspaceCwd)) { throw new ChannelWebhookEnqueueError( 'channel_worker_unavailable', `No channel worker owns channel "${task.channelName}".`, diff --git a/packages/cli/src/serve/channel-worker-manager.test.ts b/packages/cli/src/serve/channel-worker-manager.test.ts index 2d6d8dff9e4..70fd2f56e21 100644 --- a/packages/cli/src/serve/channel-worker-manager.test.ts +++ b/packages/cli/src/serve/channel-worker-manager.test.ts @@ -54,6 +54,11 @@ function fakeGroup( killAllSync: vi.fn(), snapshots: vi.fn(() => snapshots), primarySnapshot: vi.fn(() => snapshots[0]!), + beginWorkspaceDrain: vi.fn(), + cancelWorkspaceDrain: vi.fn(), + workspaceActivity: vi.fn(() => 0), + removeWorkspace: vi.fn(async () => {}), + restoreWorkspace: vi.fn(async () => {}), enqueueWebhookTask: vi.fn(async () => ({ accepted: true as const })), ...overrides, }; @@ -138,6 +143,88 @@ describe('createChannelWorkerManager', () => { }); }); + it('refreshes workspace topology without forcing unchanged workers', async () => { + const test = setup(); + const selection: ServeChannelSelection = { + mode: 'names', + names: ['telegram'], + }; + await test.manager.setSelection(selection); + + await test.manager.refreshWorkspaces(); + + expect(test.resolveGroups).toHaveBeenLastCalledWith(selection, 'reload'); + expect(test.group.reconcile).toHaveBeenCalledWith( + workspaceGroups(selection), + ); + expect(test.onCommittedSelection).toHaveBeenCalledTimes(2); + }); + + it('restores idle and classifies workspace topology reconcile failures', async () => { + const test = setup(); + const selection: ServeChannelSelection = { + mode: 'names', + names: ['telegram'], + }; + await test.manager.setSelection(selection); + vi.mocked(test.group.reconcile).mockRejectedValueOnce( + new ChannelWorkerReconcileError('secondary failed', { + rolledBack: true, + }), + ); + + await expect(test.manager.refreshWorkspaces()).rejects.toMatchObject({ + code: 'channel_worker_start_failed', + rolledBack: true, + }); + expect(test.manager.state()).toMatchObject({ + transition: 'idle', + selection, + }); + }); + + it('restores idle when workspace topology resolution fails', async () => { + const test = setup(); + const selection: ServeChannelSelection = { + mode: 'names', + names: ['telegram'], + }; + await test.manager.setSelection(selection); + test.resolveGroups.mockRejectedValueOnce(new Error('settings invalid')); + + await expect(test.manager.refreshWorkspaces()).rejects.toThrow( + 'settings invalid', + ); + expect(test.manager.state()).toMatchObject({ + transition: 'idle', + selection, + }); + }); + + it('does not reconcile after forced shutdown interrupts workspace refresh', async () => { + const test = setup(); + const selection: ServeChannelSelection = { + mode: 'names', + names: ['telegram'], + }; + await test.manager.setSelection(selection); + let releaseGroups!: () => void; + test.resolveGroups.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseGroups = () => resolve(workspaceGroups(selection)); + }), + ); + + const refreshing = test.manager.refreshWorkspaces(); + await vi.waitFor(() => expect(test.resolveGroups).toHaveBeenCalledTimes(2)); + test.manager.killAllSync(); + releaseGroups(); + + await expect(refreshing).rejects.toMatchObject({ code: 'daemon_draining' }); + expect(test.group.reconcile).not.toHaveBeenCalled(); + }); + it('starts the initial selection through the boot-time path', async () => { const test = setup(); const selection: ServeChannelSelection = { @@ -153,6 +240,21 @@ describe('createChannelWorkerManager', () => { expect(test.manager.state()).toMatchObject({ enabled: true, selection }); }); + it('applies an existing workspace drain before a newly created group starts', async () => { + const test = setup(); + + test.manager.beginWorkspaceDrain(PRIMARY); + await test.manager.setSelection({ + mode: 'names', + names: ['telegram'], + }); + + expect(test.group.beginWorkspaceDrain).toHaveBeenCalledWith(PRIMARY); + expect( + vi.mocked(test.group.beginWorkspaceDrain).mock.invocationCallOrder[0]!, + ).toBeLessThan(vi.mocked(test.group.start).mock.invocationCallOrder[0]!); + }); + it('keeps the boot-time lease reserved when group construction fails', async () => { const test = setup(); test.createGroup.mockImplementationOnce(() => { diff --git a/packages/cli/src/serve/channel-worker-manager.ts b/packages/cli/src/serve/channel-worker-manager.ts index 836109e62ef..5c64d4a4eb0 100644 --- a/packages/cli/src/serve/channel-worker-manager.ts +++ b/packages/cli/src/serve/channel-worker-manager.ts @@ -96,6 +96,12 @@ export interface ChannelWorkerManager { enqueueWebhookTask( task: ChannelWebhookTask, ): ReturnType; + beginWorkspaceDrain(workspaceCwd: string): void; + cancelWorkspaceDrain(workspaceCwd: string): void; + workspaceActivity(workspaceCwd: string): number; + removeWorkspace(workspaceCwd: string): Promise; + restoreWorkspace(workspaceCwd: string): Promise; + refreshWorkspaces(): Promise; workerChanged(): void; shutdown(): Promise; killAllSync(): void; @@ -150,6 +156,7 @@ export function createChannelWorkerManager( let draining = false; let hardKilled = false; let lane: Promise = Promise.resolve(); + const workspaceDrains = new Set(); const snapshot = (): ChannelWorkerControlState => ({ enabled: @@ -292,6 +299,9 @@ export function createChannelWorkerManager( ); } group = candidate; + for (const workspaceCwd of workspaceDrains) { + candidate.beginWorkspaceDrain(workspaceCwd); + } notify(); try { await candidate.start(); @@ -440,6 +450,54 @@ export function createChannelWorkerManager( } return group.enqueueWebhookTask(task); }, + beginWorkspaceDrain(workspaceCwd) { + workspaceDrains.add(workspaceCwd); + group?.beginWorkspaceDrain(workspaceCwd); + }, + cancelWorkspaceDrain(workspaceCwd) { + workspaceDrains.delete(workspaceCwd); + group?.cancelWorkspaceDrain(workspaceCwd); + }, + workspaceActivity(workspaceCwd) { + return group?.workspaceActivity(workspaceCwd) ?? 0; + }, + removeWorkspace(workspaceCwd) { + return enqueue(async () => { + try { + await group?.removeWorkspace(workspaceCwd); + notify(); + } finally { + workspaceDrains.delete(workspaceCwd); + } + }); + }, + restoreWorkspace(workspaceCwd) { + return enqueue(async () => { + await group?.restoreWorkspace(workspaceCwd); + notify(); + }); + }, + refreshWorkspaces() { + return enqueue(async () => { + if (!group || !committedSelection) return; + setTransition('reconciling', committedSelection); + let targetGroups: readonly ChannelWorkspaceGroup[]; + try { + targetGroups = await opts.resolveGroups(committedSelection, 'reload'); + } catch (error) { + setTransition('idle'); + throw error; + } + if (hardKilled) throw drainingError(); + try { + await group.reconcile(targetGroups); + } catch (error) { + setTransition('idle'); + throw classifyFailure(error, 'channel_worker_start_failed'); + } + commit(committedSelection, targetGroups); + }); + }, workerChanged: notify, shutdown() { draining = true; diff --git a/packages/cli/src/serve/routes/capabilities.ts b/packages/cli/src/serve/routes/capabilities.ts index a74c4a3c00d..ea37786e0b1 100644 --- a/packages/cli/src/serve/routes/capabilities.ts +++ b/packages/cli/src/serve/routes/capabilities.ts @@ -39,6 +39,8 @@ export function registerCapabilitiesRoutes( app.get('/capabilities', (_req, res) => { const runtimes = deps.workspaceRegistry.list(); const multiWorkspace = runtimes.length > 1; + const features = deps.currentServeFeatures(); + const runtimeRemoval = features.includes('workspace_runtime_removal'); const envelope: CapabilitiesEnvelope = { v: CAPABILITIES_SCHEMA_VERSION, protocolVersions: getServeProtocolVersions(), @@ -46,7 +48,7 @@ export function registerCapabilitiesRoutes( ? { qwenCodeVersion: deps.qwenCodeVersion } : {}), mode: deps.mode, - features: deps.currentServeFeatures(), + features, modelServices: [], // Surface the primary workspace so clients can omit `cwd` on // `POST /session`; multi-workspace clients use `workspaces[]`. @@ -81,6 +83,9 @@ export function registerCapabilitiesRoutes( cwd: runtime.workspaceCwd, primary: runtime.primary, trusted: runtime.trusted, + ...(runtimeRemoval + ? { removable: runtime.removable === true } + : {}), })), } : {}), diff --git a/packages/cli/src/serve/routes/workspace-git.test.ts b/packages/cli/src/serve/routes/workspace-git.test.ts index 3b51e8bb330..beb232b4f98 100644 --- a/packages/cli/src/serve/routes/workspace-git.test.ts +++ b/packages/cli/src/serve/routes/workspace-git.test.ts @@ -10,9 +10,10 @@ import { describe, expect, it, vi } from 'vitest'; import type { AcpSessionBridge } from '../acp-session-bridge.js'; import { sendBridgeError } from '../server/error-response.js'; import type { WorkspaceGitState } from '../workspace-git-state.js'; -import type { - WorkspaceRegistry, - WorkspaceRuntime, +import { + createWorkspaceRegistry, + type WorkspaceRegistry, + type WorkspaceRuntime, } from '../workspace-registry.js'; import { registerWorkspaceGitRoutes, @@ -34,19 +35,7 @@ function runtime( } function registry(runtimes: WorkspaceRuntime[]): WorkspaceRegistry { - return { - primary: runtimes[0]!, - list: () => runtimes, - getByWorkspaceCwd: (cwd) => - runtimes.find((item) => item.workspaceCwd === cwd), - getByWorkspaceId: (id) => runtimes.find((item) => item.workspaceId === id), - resolveWorkspaceCwd: (cwd) => - cwd === undefined - ? runtimes[0] - : runtimes.find((item) => item.workspaceCwd === cwd), - resolveLiveSessionOwner: () => ({ kind: 'not_found' }), - add: () => {}, - }; + return createWorkspaceRegistry(runtimes); } describe('workspace Git routes', () => { diff --git a/packages/cli/src/serve/routes/workspace-management.test.ts b/packages/cli/src/serve/routes/workspace-management.test.ts index ef9342ee808..f6324847e79 100644 --- a/packages/cli/src/serve/routes/workspace-management.test.ts +++ b/packages/cli/src/serve/routes/workspace-management.test.ts @@ -10,6 +10,7 @@ import request from 'supertest'; import { registerWorkspaceManagementRoutes, type WorkspaceManagementRouteDeps, + type WorkspaceRuntimeRemovalController, } from './workspace-management.js'; import type { WorkspaceRegistry, @@ -17,11 +18,14 @@ import type { } from '../workspace-registry.js'; import { tmpdir } from 'node:os'; import { realpathSync } from 'node:fs'; +import { mkdtemp, rm, symlink } from 'node:fs/promises'; +import { join } from 'node:path'; import { workspaceRegistrationId, WorkspaceRegistrationStoreLimitError, type WorkspaceRegistrationStore, } from '../workspace-registration-store.js'; +import { writeStderrLine } from '../../utils/stdioHelpers.js'; vi.mock('../../utils/stdioHelpers.js', () => ({ writeStderrLine: vi.fn(), @@ -35,6 +39,7 @@ function createMockRegistry( ): WorkspaceRegistry { const byCwd = new Map(runtimes.map((r) => [r.workspaceCwd, r])); const byId = new Map(runtimes.map((r) => [r.workspaceId, r])); + const draining = new Set(); const add = vi.fn((runtime: WorkspaceRuntime) => { runtimes.push(runtime); byCwd.set(runtime.workspaceCwd, runtime); @@ -42,24 +47,60 @@ function createMockRegistry( }); return { primary: runtimes[0]!, - list: () => Object.freeze([...runtimes]) as readonly WorkspaceRuntime[], - getByWorkspaceCwd: (cwd: string) => byCwd.get(cwd), - getByWorkspaceId: (id: string) => byId.get(id), + list: () => + Object.freeze( + runtimes.filter((runtime) => !draining.has(runtime)), + ) as readonly WorkspaceRuntime[], + listManaged: () => + Object.freeze([...runtimes]) as readonly WorkspaceRuntime[], + getByWorkspaceCwd: (cwd: string) => { + const runtime = byCwd.get(cwd); + return runtime && !draining.has(runtime) ? runtime : undefined; + }, + getManagedByWorkspaceCwd: (cwd: string) => byCwd.get(cwd), + getByWorkspaceId: (id: string) => { + const runtime = byId.get(id); + return runtime && !draining.has(runtime) ? runtime : undefined; + }, + getManagedByWorkspaceId: (id: string) => byId.get(id), resolveWorkspaceCwd: () => undefined, resolveLiveSessionOwner: () => ({ kind: 'not_found' }), + beginDrain: vi.fn((runtime: WorkspaceRuntime) => { + if (draining.has(runtime)) return false; + draining.add(runtime); + return true; + }), + cancelDrain: vi.fn((runtime: WorkspaceRuntime) => draining.delete(runtime)), + completeDrain: vi.fn((runtime: WorkspaceRuntime) => { + draining.delete(runtime); + const index = runtimes.indexOf(runtime); + if (index < 0) return false; + runtimes.splice(index, 1); + byCwd.delete(runtime.workspaceCwd); + byId.delete(runtime.workspaceId); + return true; + }), add, } as unknown as WorkspaceRegistry; } -function makeRuntime(cwd: string): WorkspaceRuntime { +function makeRuntime( + cwd: string, + overrides: Partial = {}, +): WorkspaceRuntime { return { workspaceId: `id-${cwd}`, workspaceCwd: cwd, primary: false, trusted: true, + removable: true, bridge: { + sessionCount: 0, + activePromptCount: 0, shutdown: vi.fn().mockResolvedValue(undefined), + killAllSync: vi.fn(), }, + ...overrides, } as unknown as WorkspaceRuntime; } @@ -75,8 +116,20 @@ function createApp(overrides?: Partial) { .mockImplementation((cwd: string) => Promise.resolve(makeRuntime(cwd))), ...overrides, }; - registerWorkspaceManagementRoutes(app, deps); - return { app, deps }; + const handle = registerWorkspaceManagementRoutes(app, deps); + return { app, deps, handle }; +} + +function createRemovalController( + pendingSessionStarts = 0, +): WorkspaceRuntimeRemovalController { + return { + beginDrain: vi.fn(), + cancelDrain: vi.fn(), + completeDrain: vi.fn(), + getActivity: vi.fn(() => ({ pendingSessionStarts, channelWorkers: 0 })), + disposeRuntime: vi.fn().mockResolvedValue(undefined), + }; } describe('POST /workspaces', () => { @@ -166,6 +219,68 @@ describe('POST /workspaces', () => { expect(res.body).not.toHaveProperty('persisted'); }); + it('does not double-count a runtime while its addition hook is pending', async () => { + const firstDir = await mkdtemp(join(REAL_DIR, 'qws-capacity-a-')); + const secondDir = await mkdtemp(join(REAL_DIR, 'qws-capacity-b-')); + try { + const registry = createMockRegistry( + Array.from({ length: 23 }, (_, index) => + makeRuntime(`/registered-${index}`), + ), + ); + let releaseAddition!: () => void; + const additionPending = new Promise((resolve) => { + releaseAddition = resolve; + }); + const runtimeRemoval = createRemovalController(); + runtimeRemoval.runtimeAdded = vi + .fn() + .mockReturnValueOnce(additionPending) + .mockResolvedValue(undefined); + const { app } = createApp({ + workspaceRegistry: registry, + runtimeRemoval, + }); + + const first = request(app).post('/workspaces').send({ cwd: firstDir }); + const firstResult = first.then((response) => response); + await vi.waitFor(() => { + expect(runtimeRemoval.runtimeAdded).toHaveBeenCalledOnce(); + }); + const second = await request(app) + .post('/workspaces') + .send({ cwd: secondDir }); + releaseAddition(); + + expect((await firstResult).status).toBe(201); + expect(second.status).toBe(201); + expect(registry.listManaged()).toHaveLength(25); + } finally { + await Promise.all([ + rm(firstDir, { recursive: true, force: true }), + rm(secondDir, { recursive: true, force: true }), + ]); + } + }); + + it('keeps a registered runtime when an optional adapter fails to attach', async () => { + const registry = createMockRegistry([makeRuntime('/some-other-dir')]); + const runtimeRemoval = createRemovalController(); + runtimeRemoval.runtimeAdded = vi + .fn() + .mockRejectedValue(new Error('worker unavailable')); + const { app } = createApp({ + workspaceRegistry: registry, + runtimeRemoval, + }); + + const res = await request(app).post('/workspaces').send({ cwd: REAL_DIR }); + + expect(res.status).toBe(201); + expect(registry.getByWorkspaceCwd(REAL_DIR)).toBeDefined(); + expect(runtimeRemoval.disposeRuntime).not.toHaveBeenCalled(); + }); + it('does not echo resolved paths in 409 error messages', async () => { const { app } = createApp(); const res = await request(app).post('/workspaces').send({ cwd: REAL_DIR }); @@ -203,7 +318,30 @@ describe('POST /workspaces', () => { .send({ cwd: REAL_DIR, persist: true }); expect(res.status).toBe(200); expect(res.body.persisted).toBe(true); - expect(add).toHaveBeenCalledWith(REAL_DIR); + expect(add).not.toHaveBeenCalled(); + }); + + it('does not duplicate a persisted alias when promoting its runtime', async () => { + const alias = '/raw/workspace-alias'; + const add = vi.fn(); + const runtime = makeRuntime(REAL_DIR, { + registrationIds: [workspaceRegistrationId(alias)], + }); + const { app } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + workspaceRegistrationStore: { + add, + read: vi.fn().mockResolvedValue({ workspaces: [alias] }), + } as unknown as WorkspaceRegistrationStore, + }); + + const res = await request(app) + .post('/workspaces') + .send({ cwd: REAL_DIR, persist: true }); + + expect(res.status).toBe(200); + expect(res.body.persisted).toBe(true); + expect(add).not.toHaveBeenCalled(); }); it('promotes an existing workspace without a dynamic runtime factory', async () => { @@ -373,6 +511,517 @@ describe('POST /workspaces', () => { }); }); +describe('DELETE /workspaces/:workspace', () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it('validates force and protects primary and static runtimes', async () => { + const primary = makeRuntime(REAL_DIR, { + primary: true, + removable: false, + }); + const runtimeRemoval = createRemovalController(); + const primaryApp = createApp({ + workspaceRegistry: createMockRegistry([primary]), + runtimeRemoval, + }).app; + + const invalid = await request(primaryApp) + .delete(`/workspaces/${encodeURIComponent(primary.workspaceId)}`) + .send({ force: 'yes' }); + expect(invalid.status).toBe(400); + expect(invalid.body.code).toBe('invalid_force_flag'); + + const forbidden = await request(primaryApp).delete( + `/workspaces/${encodeURIComponent(primary.workspaceId)}`, + ); + expect(forbidden.status).toBe(409); + expect(forbidden.body.code).toBe('primary_workspace_removal_forbidden'); + + const staticRuntime = makeRuntime(REAL_DIR, { removable: false }); + const staticApp = createApp({ + workspaceRegistry: createMockRegistry([staticRuntime]), + runtimeRemoval, + }).app; + const staticResult = await request(staticApp).delete( + `/workspaces/${encodeURIComponent(staticRuntime.workspaceId)}`, + ); + expect(staticResult.status).toBe(409); + expect(staticResult.body.code).toBe('static_workspace_removal_forbidden'); + }); + + it('returns 501 when runtime removal is unavailable', async () => { + const runtime = makeRuntime(REAL_DIR); + const { app } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval: undefined, + }); + + const res = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + + expect(res.status).toBe(501); + expect(res.body.code).toBe('workspace_runtime_removal_unsupported'); + }); + + it('does not expose workspace counts for an unknown removal selector', async () => { + const { app } = createApp({ + runtimeRemoval: createRemovalController(), + }); + + const res = await request(app).delete('/workspaces/unknown-workspace'); + + expect(res.status).toBe(400); + expect(res.body.code).toBe('workspace_mismatch'); + expect(res.body).not.toHaveProperty('workspaceCount'); + }); + + it('returns the fast busy snapshot without disturbing runtime gates', async () => { + const runtime = makeRuntime(REAL_DIR); + Object.assign(runtime.bridge, { sessionCount: 1, activePromptCount: 1 }); + const runtimeRemoval = createRemovalController(); + const { app, deps } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval, + }); + + const res = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + + expect(res.status).toBe(409); + expect(res.body).toMatchObject({ + code: 'workspace_busy', + activity: { sessions: 1, activePrompts: 1 }, + }); + expect(runtimeRemoval.beginDrain).not.toHaveBeenCalled(); + expect(deps.workspaceRegistry.beginDrain).not.toHaveBeenCalled(); + }); + + it('rolls every gate back when the final frozen snapshot becomes busy', async () => { + const runtime = makeRuntime(REAL_DIR); + const runtimeRemoval = createRemovalController(); + vi.mocked(runtimeRemoval.getActivity) + .mockReturnValueOnce({ pendingSessionStarts: 0, channelWorkers: 0 }) + .mockReturnValueOnce({ pendingSessionStarts: 1, channelWorkers: 0 }); + const acpHandle = { + beginWorkspaceDrain: vi.fn(), + cancelWorkspaceDrain: vi.fn(), + getWorkspaceActivity: vi.fn(() => ({ + acpConnections: 0, + memoryTasks: 0, + })), + }; + const { app, deps } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval, + getAcpHandle: () => acpHandle as never, + }); + + const res = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + + expect(res.status).toBe(409); + expect(res.body.activity.pendingSessionStarts).toBe(1); + expect(acpHandle.beginWorkspaceDrain).toHaveBeenCalledWith( + runtime.workspaceId, + ); + expect(acpHandle.cancelWorkspaceDrain).toHaveBeenCalledWith( + runtime.workspaceId, + ); + expect(runtimeRemoval.cancelDrain).toHaveBeenCalledWith(runtime); + expect(deps.workspaceRegistry.cancelDrain).toHaveBeenCalledWith(runtime); + expect(deps.workspaceRegistry.getByWorkspaceId(runtime.workspaceId)).toBe( + runtime, + ); + }); + + it('continues rolling gates back when cancel hooks throw', async () => { + const runtime = makeRuntime(REAL_DIR); + const runtimeRemoval = createRemovalController(); + vi.mocked(runtimeRemoval.getActivity) + .mockReturnValueOnce({ pendingSessionStarts: 0, channelWorkers: 0 }) + .mockReturnValueOnce({ pendingSessionStarts: 1, channelWorkers: 0 }); + vi.mocked(runtimeRemoval.cancelDrain).mockImplementation(() => { + throw new Error('controller rollback failed'); + }); + const acpHandle = { + beginWorkspaceDrain: vi.fn(), + cancelWorkspaceDrain: vi.fn(() => { + throw new Error('ACP rollback failed'); + }), + getWorkspaceActivity: vi.fn(() => ({ + acpConnections: 0, + memoryTasks: 0, + })), + }; + const { app, deps } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval, + getAcpHandle: () => acpHandle as never, + }); + + const res = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + + expect(res.status).toBe(409); + expect(acpHandle.cancelWorkspaceDrain).toHaveBeenCalledWith( + runtime.workspaceId, + ); + expect(runtimeRemoval.cancelDrain).toHaveBeenCalledWith(runtime); + expect(deps.workspaceRegistry.cancelDrain).toHaveBeenCalledWith(runtime); + expect(deps.workspaceRegistry.getByWorkspaceId(runtime.workspaceId)).toBe( + runtime, + ); + }); + + it('rolls drain back when persistent identity removal fails', async () => { + const runtime = makeRuntime(REAL_DIR); + const runtimeRemoval = createRemovalController(); + const { app, deps } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval, + workspaceRegistrationStore: { + removeByIds: vi.fn().mockRejectedValue(new Error('disk full')), + } as unknown as WorkspaceRegistrationStore, + }); + + const res = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + + expect(res.status).toBe(500); + expect(res.body.code).toBe('workspace_persist_failed'); + expect(runtimeRemoval.cancelDrain).toHaveBeenCalledWith(runtime); + expect(runtimeRemoval.disposeRuntime).not.toHaveBeenCalled(); + expect(deps.workspaceRegistry.getByWorkspaceId(runtime.workspaceId)).toBe( + runtime, + ); + }); + + it('force-removes activity, aliases, runtime resources, and registry state', async () => { + const runtime = makeRuntime(REAL_DIR, { + registrationIds: ['raw-alias-a', 'raw-alias-b'], + }); + Object.assign(runtime.bridge, { sessionCount: 2, activePromptCount: 1 }); + const runtimeRemoval = createRemovalController(1); + const removeByIds = vi.fn().mockResolvedValue(2); + const acpHandle = { + beginWorkspaceDrain: vi.fn(), + cancelWorkspaceDrain: vi.fn(), + getWorkspaceActivity: vi.fn(() => ({ + acpConnections: 1, + memoryTasks: 1, + })), + commitWorkspaceRemoval: vi.fn(), + disposeWorkspace: vi.fn(), + }; + const { app, deps } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval, + getAcpHandle: () => acpHandle as never, + workspaceRegistrationStore: { + removeByIds, + } as unknown as WorkspaceRegistrationStore, + }); + + const res = await request(app) + .delete(`/workspaces/${encodeURIComponent(runtime.workspaceId)}`) + .send({ force: true }); + + expect(res.status).toBe(200); + expect(res.body).toMatchObject({ + removed: true, + workspaceId: runtime.workspaceId, + forced: true, + persistedRegistrationRemoved: true, + activity: { + sessions: 2, + activePrompts: 1, + pendingSessionStarts: 1, + acpConnections: 1, + memoryTasks: 1, + }, + }); + expect(removeByIds).toHaveBeenCalledWith( + expect.arrayContaining([ + 'raw-alias-a', + 'raw-alias-b', + workspaceRegistrationId(runtime.workspaceCwd), + ]), + ); + expect(runtimeRemoval.disposeRuntime).toHaveBeenCalledWith( + runtime, + 'workspace_removed', + ); + expect(runtimeRemoval.completeDrain).toHaveBeenCalledWith(runtime); + expect(acpHandle.commitWorkspaceRemoval).toHaveBeenCalledWith( + runtime.workspaceId, + ); + expect(acpHandle.disposeWorkspace).toHaveBeenCalledWith( + runtime.workspaceId, + ); + expect( + deps.workspaceRegistry.getManagedByWorkspaceId(runtime.workspaceId), + ).toBeUndefined(); + }); + + it('does not reactivate a runtime when cleanup fails after persistence commits', async () => { + const runtime = makeRuntime(REAL_DIR); + Object.assign(runtime.bridge, { sessionCount: 1 }); + const runtimeRemoval = createRemovalController(); + vi.mocked(runtimeRemoval.disposeRuntime).mockRejectedValueOnce( + new Error('bridge cleanup failed'), + ); + vi.mocked(writeStderrLine).mockImplementationOnce(() => { + throw new Error('stderr closed'); + }); + const acpHandle = { + beginWorkspaceDrain: vi.fn(), + cancelWorkspaceDrain: vi.fn(), + getWorkspaceActivity: vi.fn(() => ({ + acpConnections: 0, + memoryTasks: 0, + })), + commitWorkspaceRemoval: vi.fn(() => { + throw new Error('commit cleanup failed'); + }), + disposeWorkspace: vi.fn(() => { + throw new Error('mount cleanup failed'); + }), + }; + const { app, deps } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval, + getAcpHandle: () => acpHandle as never, + workspaceRegistrationStore: { + removeByIds: vi.fn().mockResolvedValue(1), + } as unknown as WorkspaceRegistrationStore, + }); + + const res = await request(app) + .delete(`/workspaces/${encodeURIComponent(runtime.workspaceId)}`) + .send({ force: true }); + + expect(res.status).toBe(200); + expect(res.body.forced).toBe(true); + expect(runtimeRemoval.disposeRuntime).toHaveBeenCalledWith( + runtime, + 'workspace_removed', + ); + expect(runtimeRemoval.cancelDrain).not.toHaveBeenCalled(); + expect(runtimeRemoval.completeDrain).toHaveBeenCalledWith(runtime); + expect(acpHandle.cancelWorkspaceDrain).not.toHaveBeenCalled(); + expect(runtime.bridge.killAllSync).toHaveBeenCalledOnce(); + expect(deps.workspaceRegistry.cancelDrain).not.toHaveBeenCalled(); + expect( + deps.workspaceRegistry.getManagedByWorkspaceId(runtime.workspaceId), + ).toBeUndefined(); + }); + + it('accepts a URL-encoded absolute cwd selector', async () => { + const runtime = makeRuntime(REAL_DIR); + const { app, deps } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval: createRemovalController(), + }); + + const res = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceCwd)}`, + ); + + expect(res.status).toBe(200); + expect(res.body.workspaceCwd).toBe(runtime.workspaceCwd); + expect( + deps.workspaceRegistry.getManagedByWorkspaceCwd(runtime.workspaceCwd), + ).toBeUndefined(); + }); + + it('canonicalizes a symlink cwd selector before removal', async () => { + const selectorRoot = await mkdtemp(join(REAL_DIR, 'qws-selector-')); + const selector = join(selectorRoot, 'workspace-alias'); + await symlink( + REAL_DIR, + selector, + process.platform === 'win32' ? 'junction' : 'dir', + ); + const runtime = makeRuntime(REAL_DIR); + const { app, deps } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval: createRemovalController(), + }); + + try { + const res = await request(app).delete( + `/workspaces/${encodeURIComponent(selector)}`, + ); + + expect(res.status).toBe(200); + expect(res.body.workspaceCwd).toBe(runtime.workspaceCwd); + expect( + deps.workspaceRegistry.getManagedByWorkspaceCwd(runtime.workspaceCwd), + ).toBeUndefined(); + } finally { + await rm(selectorRoot, { recursive: true, force: true }); + } + }); + + it('reserves the cwd against concurrent remove and add operations', async () => { + const runtime = makeRuntime(REAL_DIR); + let finishCleanup!: () => void; + const cleanup = new Promise((resolve) => { + finishCleanup = resolve; + }); + const runtimeRemoval = createRemovalController(); + vi.mocked(runtimeRemoval.disposeRuntime).mockReturnValue(cleanup); + const { app } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval, + }); + + const firstRemoval = request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + const firstResult = firstRemoval.then((res) => res); + await vi.waitFor(() => { + expect(runtimeRemoval.disposeRuntime).toHaveBeenCalledWith( + runtime, + 'workspace_removed', + ); + }); + + const duplicateRemoval = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + expect(duplicateRemoval.status).toBe(409); + expect(duplicateRemoval.body.code).toBe('workspace_removal_in_progress'); + + const concurrentAdd = await request(app) + .post('/workspaces') + .send({ cwd: runtime.workspaceCwd }); + expect(concurrentAdd.status).toBe(409); + expect(concurrentAdd.body.code).toBe('workspace_removal_in_progress'); + + finishCleanup(); + expect((await firstResult).status).toBe(200); + const replacement = await request(app) + .post('/workspaces') + .send({ cwd: runtime.workspaceCwd }); + expect(replacement.status).toBe(201); + }); + + it('waits for an in-flight removal after sealing and rejects new work', async () => { + const runtime = makeRuntime(REAL_DIR); + let finishCleanup!: () => void; + const cleanup = new Promise((resolve) => { + finishCleanup = resolve; + }); + const runtimeRemoval = createRemovalController(); + vi.mocked(runtimeRemoval.disposeRuntime).mockReturnValue(cleanup); + const { app, handle } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval, + }); + const removal = request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + const removalResult = removal.then((res) => res); + await vi.waitFor(() => { + expect(runtimeRemoval.disposeRuntime).toHaveBeenCalled(); + }); + + let sealed = false; + const seal = handle.sealAndWait().then(() => { + sealed = true; + }); + await Promise.resolve(); + expect(sealed).toBe(false); + const rejected = await request(app) + .post('/workspaces') + .send({ cwd: runtime.workspaceCwd }); + expect(rejected.status).toBe(503); + + finishCleanup(); + expect((await removalResult).status).toBe(200); + await seal; + expect(sealed).toBe(true); + }); + + it('finishes sealing when an in-flight removal fails', async () => { + const runtime = makeRuntime(REAL_DIR); + let failPersistence!: (error: Error) => void; + const persistence = new Promise((_resolve, reject) => { + failPersistence = reject; + }); + const removeByIds = vi.fn().mockReturnValue(persistence); + const { app, handle } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval: createRemovalController(), + workspaceRegistrationStore: { + removeByIds, + } as unknown as WorkspaceRegistrationStore, + }); + const removal = request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + const removalResult = removal.then((res) => res); + await vi.waitFor(() => expect(removeByIds).toHaveBeenCalledOnce()); + + let sealed = false; + const seal = handle.sealAndWait().then(() => { + sealed = true; + }); + await Promise.resolve(); + expect(sealed).toBe(false); + + failPersistence(new Error('disk full')); + const result = await removalResult; + expect(result.status).toBe(500); + expect(result.body.code).toBe('workspace_persist_failed'); + await seal; + expect(sealed).toBe(true); + }); + + it('returns a coded error when removal fails before persistence commits', async () => { + const runtime = makeRuntime(REAL_DIR); + const registry = createMockRegistry([runtime]); + vi.mocked(registry.beginDrain).mockImplementationOnce(() => { + throw new Error('drain failed'); + }); + const { app } = createApp({ + workspaceRegistry: registry, + runtimeRemoval: createRemovalController(), + }); + + const res = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + + expect(res.status).toBe(500); + expect(res.body.code).toBe('workspace_runtime_removal_failed'); + }); + + it('rejects removal after workspace management is sealed', async () => { + const runtime = makeRuntime(REAL_DIR); + const { app, handle } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval: createRemovalController(), + }); + await handle.sealAndWait(); + + const res = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + + expect(res.status).toBe(503); + expect(res.body.code).toBe('daemon_shutting_down'); + }); +}); + describe('persistent workspace registrations', () => { it('returns 501 for registration management without a store', async () => { const { app } = createApp(); @@ -389,12 +1038,17 @@ describe('persistent workspace registrations', () => { }); it('lists desired registrations and whether they are active', async () => { + const alias = '/raw/symlink-alias'; + const aliasId = workspaceRegistrationId(alias); const read = vi.fn().mockResolvedValue({ schemaVersion: 1, primaryWorkspace: '/primary', - workspaces: [REAL_DIR, '/currently-unavailable'], + workspaces: [REAL_DIR, alias, '/currently-unavailable'], }); const { app } = createApp({ + workspaceRegistry: createMockRegistry([ + makeRuntime(REAL_DIR, { registrationIds: [aliasId] }), + ]), workspaceRegistrationStore: { read, } as unknown as WorkspaceRegistrationStore, @@ -410,6 +1064,12 @@ describe('persistent workspace registrations', () => { active: true, persisted: true, }), + expect.objectContaining({ + id: aliasId, + cwd: alias, + active: true, + persisted: true, + }), expect.objectContaining({ cwd: '/currently-unavailable', active: false, @@ -419,7 +1079,8 @@ describe('persistent workspace registrations', () => { }); it('forgets persistence without unloading an active runtime', async () => { - const active = makeRuntime(REAL_DIR); + const aliasId = workspaceRegistrationId('/raw/symlink-alias'); + const active = makeRuntime(REAL_DIR, { registrationIds: [aliasId] }); const removeById = vi.fn().mockResolvedValue(true); const { app } = createApp({ workspaceRegistry: createMockRegistry([active]), @@ -440,11 +1101,68 @@ describe('persistent workspace registrations', () => { active: true, restartRequired: true, }); + + const aliasResult = await request(app).delete( + `/workspace-registrations/${aliasId}`, + ); + expect(aliasResult.status).toBe(200); + expect(aliasResult.body).toMatchObject({ + removed: true, + active: true, + restartRequired: true, + }); + }); + + it('does not require restart when forgetting an alias of a static runtime', async () => { + const aliasId = workspaceRegistrationId('/raw/static-alias'); + const active = makeRuntime(REAL_DIR, { + removable: false, + registrationIds: [aliasId], + }); + const { app } = createApp({ + workspaceRegistry: createMockRegistry([active]), + workspaceRegistrationStore: { + removeById: vi.fn().mockResolvedValue(true), + } as unknown as WorkspaceRegistrationStore, + }); + + const res = await request(app).delete( + `/workspace-registrations/${aliasId}`, + ); + + expect(res.status).toBe(200); + expect(res.body).toMatchObject({ + active: true, + restartRequired: false, + }); + }); + + it('treats a draining runtime registration as active', async () => { + const active = makeRuntime(REAL_DIR); + const registry = createMockRegistry([active]); + expect(registry.beginDrain(active)).toBe(true); + const { app } = createApp({ + workspaceRegistry: registry, + workspaceRegistrationStore: { + removeById: vi.fn().mockResolvedValue(true), + } as unknown as WorkspaceRegistrationStore, + }); + + const res = await request(app).delete( + `/workspace-registrations/${workspaceRegistrationId(REAL_DIR)}`, + ); + + expect(res.status).toBe(200); + expect(res.body).toMatchObject({ + active: true, + restartRequired: true, + }); }); it('returns 404 when a registration does not exist', async () => { const { app } = createApp({ workspaceRegistrationStore: { + read: vi.fn().mockResolvedValue({ workspaces: [] }), removeById: vi.fn().mockResolvedValue(false), } as unknown as WorkspaceRegistrationStore, }); @@ -471,6 +1189,7 @@ describe('persistent workspace registrations', () => { it('returns a store error when a registration cannot be forgotten', async () => { const { app } = createApp({ workspaceRegistrationStore: { + read: vi.fn().mockResolvedValue({ workspaces: [] }), removeById: vi.fn().mockRejectedValue(new Error('write failed')), } as unknown as WorkspaceRegistrationStore, }); @@ -480,4 +1199,185 @@ describe('persistent workspace registrations', () => { expect(res.status).toBe(500); expect(res.body.code).toBe('workspace_registration_store_error'); }); + + it('serializes forget and runtime removal for the same workspace', async () => { + const registrationId = 'runtime-alias'; + const runtime = makeRuntime(REAL_DIR, { + registrationIds: [registrationId], + }); + let finishForget!: () => void; + const forgetting = new Promise((resolve) => { + finishForget = () => resolve(true); + }); + let finishRemoval!: () => void; + const removing = new Promise((resolve) => { + finishRemoval = () => resolve(1); + }); + const removeById = vi.fn().mockReturnValue(forgetting); + const removeByIds = vi.fn().mockReturnValue(removing); + const { app } = createApp({ + workspaceRegistry: createMockRegistry([runtime]), + runtimeRemoval: createRemovalController(), + workspaceRegistrationStore: { + read: vi.fn().mockResolvedValue({ workspaces: [] }), + removeById, + removeByIds, + } as unknown as WorkspaceRegistrationStore, + }); + + const pendingForget = request(app).delete( + `/workspace-registrations/${registrationId}`, + ); + const forgetResult = pendingForget.then((res) => res); + await vi.waitFor(() => expect(removeById).toHaveBeenCalledOnce()); + const removalWhileForgetting = await request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + expect(removalWhileForgetting.status).toBe(409); + expect(removalWhileForgetting.body.code).toBe( + 'workspace_registration_in_progress', + ); + finishForget(); + expect((await forgetResult).status).toBe(200); + + const pendingRemoval = request(app).delete( + `/workspaces/${encodeURIComponent(runtime.workspaceId)}`, + ); + const removalResult = pendingRemoval.then((res) => res); + await vi.waitFor(() => expect(removeByIds).toHaveBeenCalledOnce()); + const forgetWhileRemoving = await request(app).delete( + `/workspace-registrations/${registrationId}`, + ); + expect(forgetWhileRemoving.status).toBe(409); + expect(forgetWhileRemoving.body.code).toBe('workspace_removal_in_progress'); + finishRemoval(); + expect((await removalResult).status).toBe(200); + }); + + it('serializes an inactive forget with adding the same workspace', async () => { + const registrationId = workspaceRegistrationId(REAL_DIR); + let finishForget!: () => void; + const forgetting = new Promise((resolve) => { + finishForget = () => resolve(true); + }); + const removeById = vi.fn().mockReturnValue(forgetting); + const { app, deps } = createApp({ + workspaceRegistry: createMockRegistry([makeRuntime('/some-other-dir')]), + workspaceRegistrationStore: { + read: vi.fn().mockResolvedValue({ workspaces: [REAL_DIR] }), + removeById, + } as unknown as WorkspaceRegistrationStore, + }); + + const pendingForget = request(app).delete( + `/workspace-registrations/${registrationId}`, + ); + const forgetResult = pendingForget.then((res) => res); + await vi.waitFor(() => expect(removeById).toHaveBeenCalledOnce()); + const addWhileForgetting = await request(app) + .post('/workspaces') + .send({ cwd: REAL_DIR }); + + expect(addWhileForgetting.status).toBe(409); + expect(addWhileForgetting.body.code).toBe('workspace_exists'); + expect(deps.createWorkspaceRuntime).not.toHaveBeenCalled(); + finishForget(); + expect((await forgetResult).status).toBe(200); + }); + + it('waits for an in-flight forget after sealing and rejects another one', async () => { + let finishForget!: () => void; + const forgetting = new Promise((resolve) => { + finishForget = () => resolve(true); + }); + const removeById = vi.fn().mockReturnValueOnce(forgetting); + const { app, handle } = createApp({ + workspaceRegistrationStore: { + read: vi.fn().mockResolvedValue({ workspaces: [] }), + removeById, + } as unknown as WorkspaceRegistrationStore, + }); + const first = request(app).delete('/workspace-registrations/first'); + const firstResult = first.then((res) => res); + await vi.waitFor(() => expect(removeById).toHaveBeenCalledOnce()); + + let sealed = false; + const seal = handle.sealAndWait().then(() => { + sealed = true; + }); + await Promise.resolve(); + expect(sealed).toBe(false); + const second = await request(app).delete('/workspace-registrations/second'); + expect(second.status).toBe(503); + expect(second.body.code).toBe('daemon_shutting_down'); + + finishForget(); + expect((await firstResult).status).toBe(200); + await seal; + expect(sealed).toBe(true); + }); + + it('waits for an inactive forget that is still reading its registration', async () => { + let finishRead!: () => void; + const reading = new Promise<{ workspaces: string[] }>((resolve) => { + finishRead = () => resolve({ workspaces: [] }); + }); + const read = vi.fn().mockReturnValue(reading); + const removeById = vi.fn().mockResolvedValue(true); + const { app, handle } = createApp({ + workspaceRegistrationStore: { + read, + removeById, + } as unknown as WorkspaceRegistrationStore, + }); + const pendingForget = request(app).delete( + '/workspace-registrations/inactive', + ); + const forgetResult = pendingForget.then((res) => res); + await vi.waitFor(() => expect(read).toHaveBeenCalledOnce()); + + let sealed = false; + const seal = handle.sealAndWait().then(() => { + sealed = true; + }); + await Promise.resolve(); + expect(sealed).toBe(false); + expect(removeById).not.toHaveBeenCalled(); + + finishRead(); + expect((await forgetResult).status).toBe(200); + await seal; + expect(sealed).toBe(true); + expect(removeById).toHaveBeenCalledOnce(); + }); + + it('finishes a sealed forget operation when registration reading fails', async () => { + let failRead!: () => void; + const reading = new Promise((_resolve, reject) => { + failRead = () => reject(new Error('read failed')); + }); + const read = vi.fn().mockReturnValue(reading); + const { app, handle } = createApp({ + workspaceRegistrationStore: { + read, + } as unknown as WorkspaceRegistrationStore, + }); + const pendingForget = request(app).delete( + '/workspace-registrations/inactive', + ); + const forgetResult = pendingForget.then((res) => res); + await vi.waitFor(() => expect(read).toHaveBeenCalledOnce()); + + let sealed = false; + const seal = handle.sealAndWait().then(() => { + sealed = true; + }); + await Promise.resolve(); + expect(sealed).toBe(false); + + failRead(); + expect((await forgetResult).status).toBe(500); + await seal; + expect(sealed).toBe(true); + }); }); diff --git a/packages/cli/src/serve/routes/workspace-management.ts b/packages/cli/src/serve/routes/workspace-management.ts index 8f6f5eb564b..813d5516585 100644 --- a/packages/cli/src/serve/routes/workspace-management.ts +++ b/packages/cli/src/serve/routes/workspace-management.ts @@ -16,8 +16,14 @@ import type { WorkspaceRegistry, WorkspaceRuntime, } from '../workspace-registry.js'; +import type { AcpHttpHandle } from '../acp-http/index.js'; +import { + isPortableAbsolutePath, + resolveManagedWorkspaceRuntimeByPathSelector, +} from '../workspace-route-runtime.js'; import { workspaceRegistrationId, + WorkspaceRegistrationStoreCommittedError, WorkspaceRegistrationStoreLimitError, type WorkspaceRegistrationStore, } from '../workspace-registration-store.js'; @@ -33,23 +39,76 @@ export interface WorkspaceManagementRouteDeps { safeBody: (req: Request) => Record; createWorkspaceRuntime?: (cwd: string) => Promise; workspaceRegistrationStore?: WorkspaceRegistrationStore; + getAcpHandle?: () => AcpHttpHandle | undefined; + runtimeRemoval?: WorkspaceRuntimeRemovalController; +} + +export interface WorkspaceRemovalActivity { + sessions: number; + activePrompts: number; + pendingSessionStarts: number; + acpConnections: number; + memoryTasks: number; + channelWorkers: number; +} + +export interface WorkspaceRuntimeRemovalController { + runtimeAdded?(runtime: WorkspaceRuntime): Promise; + beginDrain(runtime: WorkspaceRuntime): void; + cancelDrain(runtime: WorkspaceRuntime): void; + completeDrain(runtime: WorkspaceRuntime): void; + getActivity(runtime: WorkspaceRuntime): { + pendingSessionStarts: number; + channelWorkers: number; + }; + disposeRuntime( + runtime: WorkspaceRuntime, + reason?: 'daemon_shutdown' | 'workspace_removed', + ): Promise; +} + +export interface WorkspaceManagementHandle { + sealAndWait(): Promise; } export function registerWorkspaceManagementRoutes( app: Application, deps: WorkspaceManagementRouteDeps, -): void { +): WorkspaceManagementHandle { const { workspaceRegistry, mutate, safeBody, createWorkspaceRuntime, workspaceRegistrationStore, + getAcpHandle, + runtimeRemoval, } = deps; - // Canonical cwds with a registration in flight, so two concurrent POSTs for - // the same directory can't both pass the duplicate check and both build - // runtime infrastructure before one add() throws. - const inFlight = new Set(); + // Serialize runtime addition, persistence promotion/forget, and removal by + // canonical cwd so conflicting management mutations cannot cross their + // validation and persistence commit points concurrently. + const inFlight = new Map< + string, + 'addition' | 'promotion' | 'removal' | 'forget' + >(); + let sealed = false; + let activeOperations = 0; + const idleWaiters = new Set<() => void>(); + const operationStarted = (): void => { + activeOperations++; + }; + const operationFinished = (): void => { + activeOperations--; + if (activeOperations !== 0) return; + for (const resolveIdle of idleWaiters) resolveIdle(); + idleWaiters.clear(); + }; + const sendSealed = (res: Response): void => { + res.status(503).json({ + error: 'Daemon is shutting down', + code: 'daemon_shutting_down', + }); + }; app.post( '/workspaces', @@ -121,6 +180,11 @@ export function registerWorkspaceManagementRoutes( return; } + if (sealed) { + sendSealed(res); + return; + } + try { const s = await stat(canonical); if (!s.isDirectory()) { @@ -138,12 +202,28 @@ export function registerWorkspaceManagementRoutes( return; } + // `stat` yields. Shutdown may seal management after the earlier fast + // check, so re-check immediately before claiming the cwd operation. + if (sealed) { + sendSealed(res); + return; + } + // The duplicate / in-flight / nesting checks and `inFlight.add` below run // synchronously (no `await` between them), so concurrent POSTs for the // same canonical cwd can't race past registration. Error messages stay // generic and never echo a resolved path (which could reveal symlink // targets or another workspace's location). - const existingRuntime = workspaceRegistry.getByWorkspaceCwd(canonical); + const activeOperation = inFlight.get(canonical); + if (activeOperation === 'removal') { + res.status(409).json({ + error: 'Workspace removal is in progress', + code: 'workspace_removal_in_progress', + }); + return; + } + const existingRuntime = + workspaceRegistry.getManagedByWorkspaceCwd(canonical); if (existingRuntime?.primary && persist) { res.status(400).json({ error: 'Primary workspace cannot be persisted', @@ -152,9 +232,20 @@ export function registerWorkspaceManagementRoutes( return; } if (existingRuntime && persist && !existingRuntime.primary) { + if (activeOperation) { + res.status(409).json({ + error: 'Workspace registration is in progress', + code: 'workspace_registration_in_progress', + }); + return; + } const nested = [ - ...workspaceRegistry.list().map((runtime) => runtime.workspaceCwd), - ...inFlight, + ...workspaceRegistry + .listManaged() + .map((runtime) => runtime.workspaceCwd), + ...[...inFlight].flatMap(([cwd, operation]) => + operation === 'addition' || operation === 'promotion' ? [cwd] : [], + ), ].some( (boundCwd) => boundCwd !== canonical && @@ -168,12 +259,18 @@ export function registerWorkspaceManagementRoutes( }); return; } + inFlight.set(canonical, 'promotion'); + operationStarted(); try { const snapshot = await workspaceRegistrationStore!.read(); - const alreadyPersisted = snapshot.workspaces.some((stored) => - process.platform === 'win32' - ? stored.toLowerCase() === canonical.toLowerCase() - : stored === canonical, + const alreadyPersisted = snapshot.workspaces.some( + (stored) => + existingRuntime.registrationIds?.includes( + workspaceRegistrationId(stored), + ) === true || + (process.platform === 'win32' + ? stored.toLowerCase() === canonical.toLowerCase() + : stored === canonical), ); if ( !alreadyPersisted && @@ -185,7 +282,20 @@ export function registerWorkspaceManagementRoutes( }); return; } - await workspaceRegistrationStore!.add(canonical); + if (!alreadyPersisted) { + try { + await workspaceRegistrationStore!.add(canonical); + } catch (err) { + if (!(err instanceof WorkspaceRegistrationStoreCommittedError)) { + throw err; + } + try { + writeStderrLine(`qwen serve: ${err.message}`); + } catch { + // The registration is committed; diagnostics are best-effort. + } + } + } res.status(200).json({ id: existingRuntime.workspaceId, cwd: existingRuntime.workspaceCwd, @@ -210,10 +320,13 @@ export function registerWorkspaceManagementRoutes( error: 'Failed to persist workspace registration', code: 'workspace_registration_store_error', }); + } finally { + inFlight.delete(canonical); + operationFinished(); } return; } - if (existingRuntime || inFlight.has(canonical)) { + if (existingRuntime || activeOperation) { res.status(409).json({ error: 'Workspace already registered', code: 'workspace_exists', @@ -232,8 +345,10 @@ export function registerWorkspaceManagementRoutes( // so two concurrent POSTs for parent/child paths (e.g. /project and // /project/sub) can't both pass while neither is in the registry yet. const boundCwds = [ - ...workspaceRegistry.list().map((r) => r.workspaceCwd), - ...inFlight, + ...workspaceRegistry.listManaged().map((r) => r.workspaceCwd), + ...[...inFlight].flatMap(([cwd, operation]) => + operation === 'addition' ? [cwd] : [], + ), ]; for (const existing of boundCwds) { if ( @@ -249,10 +364,13 @@ export function registerWorkspaceManagementRoutes( } } - if ( - workspaceRegistry.list().length + inFlight.size >= - MAX_REGISTERED_WORKSPACES - ) { + const projectedWorkspaceCwds = new Set( + workspaceRegistry.listManaged().map((runtime) => runtime.workspaceCwd), + ); + for (const [cwd, operation] of inFlight) { + if (operation === 'addition') projectedWorkspaceCwds.add(cwd); + } + if (projectedWorkspaceCwds.size >= MAX_REGISTERED_WORKSPACES) { res.status(409).json({ error: 'Workspace registration limit reached', code: 'workspace_limit_reached', @@ -260,7 +378,8 @@ export function registerWorkspaceManagementRoutes( return; } - inFlight.add(canonical); + inFlight.set(canonical, 'addition'); + operationStarted(); let persistenceFailed = false; try { const runtime = await createWorkspaceRuntime(canonical); @@ -268,14 +387,41 @@ export function registerWorkspaceManagementRoutes( try { if (persist) { try { - persistedRecordAdded = - await workspaceRegistrationStore!.add(canonical); + try { + persistedRecordAdded = + await workspaceRegistrationStore!.add(canonical); + } catch (err) { + if ( + !(err instanceof WorkspaceRegistrationStoreCommittedError) + ) { + throw err; + } + persistedRecordAdded = true; + try { + writeStderrLine(`qwen serve: ${err.message}`); + } catch { + // The registration is committed; diagnostics are best-effort. + } + } } catch (err) { persistenceFailed = true; throw err; } } workspaceRegistry.add(runtime); + try { + await runtimeRemoval?.runtimeAdded?.(runtime); + } catch (err) { + try { + writeStderrLine( + `qwen serve: workspace runtime adapter notification failed after registry add: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + } catch { + // The runtime is registered; diagnostics are best-effort. + } + } } catch (err) { if (persistedRecordAdded) { try { @@ -292,7 +438,19 @@ export function registerWorkspaceManagementRoutes( ); } } - await runtime.bridge.shutdown().catch(() => undefined); + if (runtimeRemoval) { + await runtimeRemoval + .disposeRuntime(runtime, 'workspace_removed') + .catch(() => { + try { + runtime.bridge.killAllSync(); + } catch { + // Preserve the original registration failure. + } + }); + } else { + await runtime.bridge.shutdown().catch(() => undefined); + } throw err; } res.status(201).json({ @@ -323,10 +481,322 @@ export function registerWorkspaceManagementRoutes( } } finally { inFlight.delete(canonical); + operationFinished(); + } + }, + ); + + const workspaceActivity = ( + runtime: WorkspaceRuntime, + ): WorkspaceRemovalActivity => { + const controllerActivity = runtimeRemoval?.getActivity(runtime) ?? { + pendingSessionStarts: 0, + channelWorkers: 0, + }; + const acpActivity = getAcpHandle?.()?.getWorkspaceActivity( + runtime.workspaceId, + ) ?? { acpConnections: 0, memoryTasks: 0 }; + return { + pendingSessionStarts: controllerActivity.pendingSessionStarts, + sessions: runtime.bridge.sessionCount, + activePrompts: runtime.bridge.activePromptCount, + acpConnections: acpActivity.acpConnections, + memoryTasks: acpActivity.memoryTasks, + channelWorkers: controllerActivity.channelWorkers, + }; + }; + const isBusy = (activity: WorkspaceRemovalActivity): boolean => + Object.values(activity).some((count) => count > 0); + const resolveManagedRuntime = ( + req: Request, + res: Response, + ): WorkspaceRuntime | undefined => { + const selector = String(req.params['workspace'] ?? ''); + const byId = workspaceRegistry.getManagedByWorkspaceId(selector); + if (byId) return byId; + if (!isPortableAbsolutePath(selector)) { + res.status(400).json({ + error: '`workspace` must decode to a workspace id or absolute path', + code: 'workspace_mismatch', + }); + return undefined; + } + const runtime = resolveManagedWorkspaceRuntimeByPathSelector( + workspaceRegistry, + selector, + ); + if (runtime) return runtime; + res.status(400).json({ + error: + 'Workspace mismatch: the requested workspace is not registered with this daemon.', + code: 'workspace_mismatch', + }); + return undefined; + }; + + app.delete( + '/workspaces/:workspace', + mutate({ strict: true }), + async (req: Request, res: Response) => { + const body = safeBody(req); + const force = body['force']; + if (force !== undefined && typeof force !== 'boolean') { + res.status(400).json({ + error: '`force` must be a boolean when provided', + code: 'invalid_force_flag', + }); + return; + } + if (sealed) { + sendSealed(res); + return; + } + const runtime = resolveManagedRuntime(req, res); + if (!runtime) return; + if (runtime.primary) { + res.status(409).json({ + error: 'The primary workspace cannot be removed at runtime', + code: 'primary_workspace_removal_forbidden', + }); + return; + } + if (runtime.removable !== true) { + res.status(409).json({ + error: 'Startup workspaces cannot be removed at runtime', + code: 'static_workspace_removal_forbidden', + }); + return; + } + if (!runtimeRemoval) { + res.status(501).json({ + error: 'Workspace runtime removal is not available', + code: 'workspace_runtime_removal_unsupported', + }); + return; + } + + const operation = inFlight.get(runtime.workspaceCwd); + if (operation) { + res.status(409).json({ + error: + operation === 'removal' + ? 'Workspace removal is in progress' + : 'Workspace registration is in progress', + code: + operation === 'removal' + ? 'workspace_removal_in_progress' + : 'workspace_registration_in_progress', + }); + return; + } + + const initialActivity = workspaceActivity(runtime); + if (force !== true && isBusy(initialActivity)) { + res.status(409).json({ + error: 'Workspace has active runtime resources', + code: 'workspace_busy', + activity: initialActivity, + }); + return; + } + + inFlight.set(runtime.workspaceCwd, 'removal'); + operationStarted(); + let registryDraining = false; + let controllerDraining = false; + let acpDraining = false; + const rollbackDrain = (): void => { + if (acpDraining) { + try { + getAcpHandle?.()?.cancelWorkspaceDrain(runtime.workspaceId); + } catch { + // Continue rolling back the remaining gates. + } + acpDraining = false; + } + if (controllerDraining) { + try { + runtimeRemoval.cancelDrain(runtime); + } catch { + // Continue rolling back the remaining gates. + } + controllerDraining = false; + } + if (registryDraining) { + try { + workspaceRegistry.cancelDrain(runtime); + } catch { + // Every rollback gate has now been attempted. + } + registryDraining = false; + } + }; + const convergeCommittedRemoval = async (): Promise => { + const logCleanupFailure = (message: string): void => { + try { + writeStderrLine(message); + } catch { + // Cleanup must continue after the persistence commit point. + } + }; + try { + getAcpHandle?.()?.commitWorkspaceRemoval(runtime.workspaceId); + } catch (err) { + logCleanupFailure( + `qwen serve: failed to commit workspace ACP removal: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + } + await runtimeRemoval + .disposeRuntime(runtime, 'workspace_removed') + .catch((err) => { + logCleanupFailure( + `qwen serve: workspace runtime cleanup failed: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + try { + runtime.bridge.killAllSync(); + } catch { + // Logical removal must still converge after persistence commits. + } + }); + try { + getAcpHandle?.()?.disposeWorkspace(runtime.workspaceId); + } catch (err) { + logCleanupFailure( + `qwen serve: failed to dispose workspace ACP mount: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + } + try { + runtimeRemoval.completeDrain(runtime); + } catch (err) { + logCleanupFailure( + `qwen serve: failed to complete workspace admission drain: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + } + try { + workspaceRegistry.completeDrain(runtime); + } catch (err) { + logCleanupFailure( + `qwen serve: failed to complete workspace registry drain: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + } + registryDraining = false; + controllerDraining = false; + acpDraining = false; + }; + + try { + registryDraining = workspaceRegistry.beginDrain(runtime); + if (!registryDraining) { + res.status(409).json({ + error: 'Workspace removal is in progress', + code: 'workspace_removal_in_progress', + }); + return; + } + runtimeRemoval.beginDrain(runtime); + controllerDraining = true; + getAcpHandle?.()?.beginWorkspaceDrain(runtime.workspaceId); + acpDraining = true; + + const activity = workspaceActivity(runtime); + if (force !== true && isBusy(activity)) { + rollbackDrain(); + res.status(409).json({ + error: 'Workspace has active runtime resources', + code: 'workspace_busy', + activity, + }); + return; + } + + let persistedRegistrationRemoved = false; + if (workspaceRegistrationStore) { + try { + const registrationIds = new Set([ + ...(runtime.registrationIds ?? []), + workspaceRegistrationId(runtime.workspaceCwd), + ]); + try { + persistedRegistrationRemoved = + (await workspaceRegistrationStore.removeByIds([ + ...registrationIds, + ])) > 0; + } catch (err) { + if (!(err instanceof WorkspaceRegistrationStoreCommittedError)) { + throw err; + } + persistedRegistrationRemoved = true; + try { + writeStderrLine(`qwen serve: ${err.message}`); + } catch { + // Persistence committed; diagnostics are best-effort. + } + } + } catch (err) { + rollbackDrain(); + writeStderrLine( + `qwen serve: failed to remove workspace persistence: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + res.status(500).json({ + error: 'Failed to persist workspace removal', + code: 'workspace_persist_failed', + }); + return; + } + } + + // Persistence is the commit point. Every cleanup step after it is + // best-effort and logical removal must never roll back to active. + await convergeCommittedRemoval(); + + res.status(200).json({ + removed: true, + workspaceId: runtime.workspaceId, + workspaceCwd: runtime.workspaceCwd, + forced: force === true, + persistedRegistrationRemoved, + activity, + }); + } catch (err) { + rollbackDrain(); + writeStderrLine( + `qwen serve: DELETE /workspaces/:workspace failed: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + if (!res.headersSent) { + res.status(500).json({ + error: 'Failed to remove workspace runtime', + code: 'workspace_runtime_removal_failed', + }); + } + } finally { + inFlight.delete(runtime.workspaceCwd); + operationFinished(); } }, ); + const registrationIsActive = (registrationId: string): boolean => + workspaceRegistry.listManaged().some((runtime) => { + if (workspaceRegistrationId(runtime.workspaceCwd) === registrationId) { + return true; + } + return runtime.registrationIds?.includes(registrationId) === true; + }); + app.get('/workspace-registrations', async (_req, res) => { if (!workspaceRegistrationStore) { res.status(501).json({ @@ -345,7 +815,9 @@ export function registerWorkspaceManagementRoutes( return { id: workspaceRegistrationId(cwd), cwd, - active: runtime !== undefined, + active: + runtime !== undefined || + registrationIsActive(workspaceRegistrationId(cwd)), persisted: true, }; }), @@ -374,16 +846,89 @@ export function registerWorkspaceManagementRoutes( }); return; } + if (sealed) { + sendSealed(res); + return; + } + operationStarted(); + let operationCwd: string | undefined; + let ownsInFlight = false; try { const registrationId = String(req.params['id']); - const active = workspaceRegistry - .list() - .some( - (runtime) => - workspaceRegistrationId(runtime.workspaceCwd) === registrationId, + let runtime = workspaceRegistry + .listManaged() + .find( + (candidate) => + workspaceRegistrationId(candidate.workspaceCwd) === + registrationId || + candidate.registrationIds?.includes(registrationId) === true, ); - const removed = - await workspaceRegistrationStore.removeById(registrationId); + operationCwd = runtime?.workspaceCwd; + if (!operationCwd) { + let storedCwd: string | undefined; + try { + const snapshot = await workspaceRegistrationStore.read(); + storedCwd = snapshot.workspaces.find( + (workspace) => + workspaceRegistrationId(workspace) === registrationId, + ); + } catch (err) { + writeStderrLine( + `qwen serve: failed to read workspace registration before forget: ${ + err instanceof Error ? err.message : String(err) + }`, + ); + res.status(500).json({ + error: 'Failed to read workspace registration', + code: 'workspace_registration_store_error', + }); + return; + } + if (storedCwd) { + try { + operationCwd = realpathSync.native(resolve(storedCwd)); + } catch { + operationCwd = resolve(storedCwd); + } + } + } + const operation = operationCwd ? inFlight.get(operationCwd) : undefined; + if (operation) { + res.status(409).json({ + error: + operation === 'removal' + ? 'Workspace removal is in progress' + : 'Workspace registration is in progress', + code: + operation === 'removal' + ? 'workspace_removal_in_progress' + : 'workspace_registration_in_progress', + }); + return; + } + if (operationCwd) { + inFlight.set(operationCwd, 'forget'); + ownsInFlight = true; + } + runtime = + (operationCwd + ? workspaceRegistry.getManagedByWorkspaceCwd(operationCwd) + : undefined) ?? runtime; + const active = registrationIsActive(registrationId); + let removed: boolean; + try { + removed = await workspaceRegistrationStore.removeById(registrationId); + } catch (err) { + if (!(err instanceof WorkspaceRegistrationStoreCommittedError)) { + throw err; + } + removed = true; + try { + writeStderrLine(`qwen serve: ${err.message}`); + } catch { + // The forget committed; diagnostics are best-effort. + } + } if (!removed) { res.status(404).json({ error: 'Workspace registration not found', @@ -394,7 +939,7 @@ export function registerWorkspaceManagementRoutes( res.json({ removed: true, active, - restartRequired: active, + restartRequired: active && runtime?.removable === true, }); } catch (err) { writeStderrLine( @@ -406,7 +951,18 @@ export function registerWorkspaceManagementRoutes( error: 'Failed to forget workspace registration', code: 'workspace_registration_store_error', }); + } finally { + if (ownsInFlight && operationCwd) inFlight.delete(operationCwd); + operationFinished(); } }, ); + + return { + async sealAndWait() { + sealed = true; + if (activeOperations === 0) return; + await new Promise((resolveIdle) => idleWaiters.add(resolveIdle)); + }, + }; } diff --git a/packages/cli/src/serve/run-qwen-serve.test.ts b/packages/cli/src/serve/run-qwen-serve.test.ts index 54543043de1..b17a90701f2 100644 --- a/packages/cli/src/serve/run-qwen-serve.test.ts +++ b/packages/cli/src/serve/run-qwen-serve.test.ts @@ -38,6 +38,7 @@ import * as serverModule from './server.js'; import * as settingsRuntime from '../config/settings.js'; import * as environmentRuntime from '../config/environment.js'; import * as trustedFoldersRuntime from '../config/trustedFolders.js'; +import * as workspaceServiceRuntime from './workspace-service/index.js'; import type { ChannelWorkerSnapshot, CreateChannelWorkerSupervisorOptions, @@ -48,7 +49,10 @@ import type { } from '../commands/channel/pidfile.js'; import { LARGE_PIPE_FRAME_THRESHOLD_BYTES } from './large-pipe-frame-observer.js'; import type { ChannelWebhookEnqueueError } from './channel-webhook-ipc.js'; -import type { WorkspaceRegistrationStore } from './workspace-registration-store.js'; +import { + workspaceRegistrationId, + type WorkspaceRegistrationStore, +} from './workspace-registration-store.js'; const BASE_BRIDGE_SNAPSHOT: BridgeDaemonStatusSnapshot = { limits: { @@ -77,6 +81,8 @@ function makeRuntimeBridge(): HttpAcpBridge { getEventRing: vi.fn().mockReturnValue({ getAll: () => [] }), resume: vi.fn(), preheat: vi.fn().mockResolvedValue(undefined), + sessionCount: 0, + activePromptCount: 0, getDaemonStatusSnapshot: vi.fn().mockReturnValue(BASE_BRIDGE_SNAPSHOT), isChannelLive: vi.fn().mockReturnValue(true), } as unknown as HttpAcpBridge; @@ -429,6 +435,9 @@ describe('runQwenServe telemetry validation', () => { enabled: false, sensitiveSpanAttributeMaxLength: 1024 * 1024, }); + const createBridge = vi + .spyOn(acpBridge, 'createAcpSessionBridge') + .mockImplementation(() => makeRuntimeBridge()); const handle = await runQwenServe( { @@ -450,25 +459,269 @@ describe('runQwenServe telemetry validation', () => { const body = (await res.json()) as { workspaceCwd: string; features: string[]; - workspaces: Array<{ cwd: string; primary: boolean }>; + workspaces: Array<{ + cwd: string; + primary: boolean; + removable?: boolean; + }>; limits: { maxTotalSessions: number | null }; }; expect(body.workspaceCwd).toBe(canonicalizeWorkspace(primary)); expect(body.features).toContain('multi_workspace_sessions'); + expect(body.features).toContain('workspace_runtime_removal'); expect(body.limits.maxTotalSessions).toBe(2); expect(body.workspaces).toEqual([ expect.objectContaining({ cwd: canonicalizeWorkspace(primary), primary: true, + removable: false, }), expect.objectContaining({ cwd: canonicalizeWorkspace(secondary), primary: false, + removable: false, }), ]); } finally { await handle.close(); } + expect(createBridge).toHaveBeenCalledTimes(2); + for (const result of createBridge.mock.results) { + expect(result.value.shutdown).toHaveBeenCalledWith({ + reason: 'daemon_shutdown', + }); + } + }); + + it('adds, advertises, and hot-removes a dynamic workspace runtime', async () => { + tmpDir = fs.realpathSync( + fs.mkdtempSync(path.join(os.tmpdir(), 'qws-hot-remove-')), + ); + const primary = path.join(tmpDir, 'primary'); + const secondary = path.join(tmpDir, 'secondary'); + fs.mkdirSync(primary); + fs.mkdirSync(secondary); + vi.spyOn(qwenCore, 'resolveTelemetrySettings').mockResolvedValue({ + enabled: false, + sensitiveSpanAttributeMaxLength: 1024 * 1024, + }); + vi.spyOn(settingsRuntime, 'loadSettings').mockReturnValue({ + merged: {}, + } as ReturnType); + vi.spyOn(trustedFoldersRuntime, 'getWorkspaceTrustStatus').mockReturnValue({ + effective: { state: 'trusted' }, + } as ReturnType); + const createBridge = vi + .spyOn(acpBridge, 'createAcpSessionBridge') + .mockImplementation(() => makeRuntimeBridge()); + const removeByIds = vi.fn().mockResolvedValue(1); + const store = { + read: vi.fn().mockResolvedValue({ + schemaVersion: 1, + primaryWorkspace: canonicalizeWorkspace(primary), + workspaces: [], + }), + add: vi.fn().mockResolvedValue(true), + removeByIds, + } as unknown as WorkspaceRegistrationStore; + const handle = await runQwenServe( + { + port: 0, + hostname: '127.0.0.1', + mode: 'http-bridge', + workspace: primary, + token: 'hot-remove-token', + serveWebShell: false, + }, + { + preheatBridge: false, + workspaceRegistrationStore: store, + daemonLogBaseDir: path.join(tmpDir, 'debug'), + }, + ); + const headers = { + Authorization: 'Bearer hot-remove-token', + 'Content-Type': 'application/json', + }; + + try { + const added = await fetch(`${handle.url}/workspaces`, { + method: 'POST', + headers, + body: JSON.stringify({ cwd: secondary, persist: true }), + }); + expect(added.status).toBe(201); + + const before = (await ( + await fetch(`${handle.url}/capabilities`, { headers }) + ).json()) as { + features: string[]; + workspaces: Array<{ + id: string; + cwd: string; + removable?: boolean; + }>; + }; + expect(before.features).toContain('workspace_runtime_removal'); + const removable = before.workspaces.find( + (workspace) => workspace.cwd === canonicalizeWorkspace(secondary), + ); + expect(removable).toMatchObject({ removable: true }); + + const removed = await fetch( + `${handle.url}/workspaces/${encodeURIComponent(removable!.id)}`, + { + method: 'DELETE', + headers, + body: JSON.stringify({ force: true }), + }, + ); + expect(removed.status).toBe(200); + await expect(removed.json()).resolves.toMatchObject({ + removed: true, + workspaceId: removable!.id, + persistedRegistrationRemoved: true, + }); + expect(removeByIds).toHaveBeenCalledWith( + expect.arrayContaining([ + workspaceRegistrationId(canonicalizeWorkspace(secondary)), + ]), + ); + const dynamicBridge = createBridge.mock.results[1]?.value; + expect(dynamicBridge?.shutdown).toHaveBeenCalledWith({ + reason: 'workspace_removed', + }); + + const afterResponse = await fetch(`${handle.url}/capabilities`, { + headers, + }); + expect(afterResponse.status).toBe(200); + const after = (await afterResponse.json()) as { + workspaces?: Array<{ id: string }>; + }; + expect( + (after.workspaces ?? []).some( + (workspace) => workspace.id === removable!.id, + ), + ).toBe(false); + + const readded = await fetch(`${handle.url}/workspaces`, { + method: 'POST', + headers, + body: JSON.stringify({ cwd: secondary, persist: true }), + }); + expect(readded.status).toBe(201); + expect(createBridge).toHaveBeenCalledTimes(3); + let releaseRemoval!: (count: number) => void; + removeByIds.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseRemoval = resolve; + }), + ); + const pendingRemoval = fetch( + `${handle.url}/workspaces/${encodeURIComponent(removable!.id)}`, + { + method: 'DELETE', + headers, + body: JSON.stringify({ force: true }), + }, + ); + await vi.waitFor(() => expect(removeByIds).toHaveBeenCalledTimes(2)); + let closeSettled = false; + const closing = handle.close().then(() => { + closeSettled = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(closeSettled).toBe(false); + + releaseRemoval(1); + expect((await pendingRemoval).status).toBe(200); + await closing; + expect(closeSettled).toBe(true); + } finally { + await handle.close(); + } + }); + + it('kills a half-built dynamic bridge when async construction cleanup fails', async () => { + tmpDir = fs.realpathSync( + fs.mkdtempSync(path.join(os.tmpdir(), 'qws-runtime-failure-')), + ); + const primary = path.join(tmpDir, 'primary'); + const secondary = path.join(tmpDir, 'secondary'); + fs.mkdirSync(primary); + fs.mkdirSync(secondary); + vi.spyOn(qwenCore, 'resolveTelemetrySettings').mockResolvedValue({ + enabled: false, + sensitiveSpanAttributeMaxLength: 1024 * 1024, + }); + vi.spyOn(settingsRuntime, 'loadSettings').mockReturnValue({ + merged: {}, + } as ReturnType); + vi.spyOn(trustedFoldersRuntime, 'getWorkspaceTrustStatus').mockReturnValue({ + effective: { state: 'trusted' }, + } as ReturnType); + + const primaryBridge = makeRuntimeBridge(); + const failedBridge = makeRuntimeBridge(); + vi.mocked(failedBridge.shutdown).mockRejectedValue( + new Error('async cleanup failed'), + ); + vi.spyOn(acpBridge, 'createAcpSessionBridge') + .mockReturnValueOnce( + primaryBridge as ReturnType, + ) + .mockReturnValueOnce( + failedBridge as ReturnType, + ); + const originalCreateWorkspaceService = + workspaceServiceRuntime.createDaemonWorkspaceService; + const createWorkspaceService = vi.spyOn( + workspaceServiceRuntime, + 'createDaemonWorkspaceService', + ); + createWorkspaceService + .mockImplementationOnce(originalCreateWorkspaceService) + .mockImplementationOnce(() => { + throw new Error('workspace service construction failed'); + }); + + const handle = await runQwenServe( + { + port: 0, + hostname: '127.0.0.1', + mode: 'http-bridge', + workspace: primary, + token: 'runtime-failure-token', + serveWebShell: false, + }, + { + preheatBridge: false, + daemonLogBaseDir: path.join(tmpDir, 'debug'), + }, + ); + + try { + const response = await fetch(`${handle.url}/workspaces`, { + method: 'POST', + headers: { + Authorization: 'Bearer runtime-failure-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ cwd: secondary }), + }); + + expect(response.status).toBe(500); + await expect(response.json()).resolves.toMatchObject({ + code: 'runtime_creation_failed', + }); + expect(failedBridge.shutdown).toHaveBeenCalledWith(); + expect(failedBridge.killAllSync).toHaveBeenCalledOnce(); + expect(primaryBridge.shutdown).not.toHaveBeenCalled(); + } finally { + await handle.close(); + } }); it('uses the daemon-wide policy and limits when constructing workspace bridges', async () => { @@ -1944,6 +2197,15 @@ describe('runQwenServe runtime startup failures', () => { fs.mkdirSync(explicitSecondary); fs.mkdirSync(restoredSecondary); fs.mkdirSync(nestedSecondary); + const restoredSecondaryAlias = path.join( + tmpDir, + 'restored-secondary-alias', + ); + fs.symlinkSync( + restoredSecondary, + restoredSecondaryAlias, + process.platform === 'win32' ? 'junction' : 'dir', + ); const canonicalPrimary = canonicalizeWorkspace(primary); const canonicalExplicitSecondary = canonicalizeWorkspace(explicitSecondary); const canonicalRestoredSecondary = canonicalizeWorkspace(restoredSecondary); @@ -1965,6 +2227,8 @@ describe('runQwenServe runtime startup failures', () => { .spyOn(acpBridge, 'createAcpSessionBridge') .mockImplementation(() => makeRuntimeBridge()); let restoredCwds: string[] = []; + let restoredRemovable: Array = []; + let restoredRegistrationIds: Array = []; let advertisedMaxTotalSessions: number | undefined; vi.spyOn(serverModule, 'createServeApp').mockImplementation( (opts, _getPort, deps) => { @@ -1972,6 +2236,13 @@ describe('runQwenServe runtime startup failures', () => { deps?.workspaceRegistry ?.list() .map((runtime) => runtime.workspaceCwd) ?? []; + restoredRemovable = + deps?.workspaceRegistry?.list().map((runtime) => runtime.removable) ?? + []; + restoredRegistrationIds = + deps?.workspaceRegistry + ?.list() + .map((runtime) => runtime.registrationIds) ?? []; advertisedMaxTotalSessions = opts.maxTotalSessions; return express(); }, @@ -1984,6 +2255,7 @@ describe('runQwenServe runtime startup failures', () => { missingPersistedWorkspace, canonicalExplicitSecondary, nestedSecondary, + restoredSecondaryAlias, canonicalRestoredSecondary, ], }), @@ -2013,6 +2285,15 @@ describe('runQwenServe runtime startup failures', () => { canonicalExplicitSecondary, canonicalRestoredSecondary, ]); + expect(restoredRemovable).toEqual([false, false, true]); + expect(restoredRegistrationIds).toEqual([ + [], + [workspaceRegistrationId(canonicalExplicitSecondary)], + [ + workspaceRegistrationId(restoredSecondaryAlias), + workspaceRegistrationId(canonicalRestoredSecondary), + ], + ]); expect(createBridge).toHaveBeenCalledTimes(3); expect(advertisedMaxTotalSessions).toBe(3); expect( @@ -3996,6 +4277,8 @@ describe('runQwenServe runtime startup failures', () => { 'daemon_status', 'workspace_settings', 'workspace_reload', + 'persistent_workspace_registration', + 'workspace_runtime_removal', ]), modelServices: [], workspaceCwd: boundWorkspace, @@ -5248,7 +5531,113 @@ describe('runQwenServe channel worker supervisor', () => { expect(supervisorFactory).not.toHaveBeenCalled(); }); - it('orchestrates and persists distinct workers for multiple workspaces', async () => { + it('records a secondary-only worker added to a primary-only daemon', async () => { + tmpDir = fs.realpathSync( + fs.mkdtempSync(path.join(os.tmpdir(), 'qws-dynamic-worker-pidfile-')), + ); + const primary = path.join(tmpDir, 'primary'); + const secondary = path.join(tmpDir, 'secondary'); + fs.mkdirSync(primary); + fs.mkdirSync(secondary); + const primaryCwd = canonicalizeWorkspace(primary); + const secondaryCwd = canonicalizeWorkspace(secondary); + vi.spyOn(qwenCore, 'resolveTelemetrySettings').mockResolvedValue({ + enabled: false, + sensitiveSpanAttributeMaxLength: 1024 * 1024, + }); + vi.spyOn(settingsRuntime, 'loadSettings').mockImplementation( + (workspace) => + ({ + merged: { + channels: + canonicalizeWorkspace(String(workspace)) === secondaryCwd + ? { feishu: { type: 'feishu' } } + : { telegram: { type: 'telegram' } }, + }, + }) as unknown as ReturnType, + ); + vi.spyOn(trustedFoldersRuntime, 'getWorkspaceTrustStatus').mockReturnValue({ + effective: { state: 'trusted' }, + } as ReturnType); + vi.spyOn(acpBridge, 'createAcpSessionBridge').mockImplementation(() => + makeFakeBridge(), + ); + const worker = makeWorker({ + enabled: true, + state: 'running', + pid: 5678, + channels: ['feishu'], + }); + const workerFactory = makeReadyWorkerFactory(worker); + const pidfile = makePidfileDeps(); + const store = { + read: vi.fn().mockResolvedValue({ + schemaVersion: 1, + primaryWorkspace: primaryCwd, + workspaces: [], + }), + add: vi.fn().mockResolvedValue(true), + } as unknown as WorkspaceRegistrationStore; + const handle = await runQwenServe( + { + port: 0, + hostname: '127.0.0.1', + mode: 'http-bridge', + workspace: primary, + token: 'dynamic-worker-token', + serveWebShell: false, + }, + { + preheatBridge: false, + daemonLogBaseDir: path.join(tmpDir, 'debug'), + channelWorkerSupervisorFactory: workerFactory, + channelServicePidfile: pidfile, + workspaceRegistrationStore: store, + }, + ); + const headers = { + Authorization: 'Bearer dynamic-worker-token', + 'Content-Type': 'application/json', + }; + + try { + await handle.runtimeReady; + const added = await fetch(`${handle.url}/workspaces`, { + method: 'POST', + headers, + body: JSON.stringify({ cwd: secondary }), + }); + expect(added.status).toBe(201); + + const enabled = await fetch(`${handle.url}/workspace/channel`, { + method: 'PUT', + headers, + body: JSON.stringify({ + selection: { mode: 'names', names: ['feishu'] }, + }), + }); + expect(enabled.status).toBe(201); + expect(workerFactory).toHaveBeenCalledOnce(); + expect(workerFactory).toHaveBeenCalledWith( + expect.objectContaining({ workspace: secondaryCwd }), + ); + expect(pidfile.writeServeServiceInfo).toHaveBeenLastCalledWith({ + channels: ['feishu'], + servePid: process.pid, + workers: [ + expect.objectContaining({ + workspaceCwd: secondaryCwd, + channels: ['feishu'], + workerPid: 5678, + }), + ], + }); + } finally { + await handle.close(); + } + }); + + it('orchestrates, persists, and hot-removes distinct workspace workers', async () => { const previousSharedSecret = process.env['QWEN_SHARED_WEBHOOK_SECRET']; process.env['QWEN_SHARED_WEBHOOK_SECRET'] = 'primary-secret'; tmpDir = fs.realpathSync( @@ -5317,15 +5706,16 @@ describe('runQwenServe channel worker supervisor', () => { vi.spyOn(trustedFoldersRuntime, 'getWorkspaceTrustStatus').mockReturnValue({ effective: { state: 'trusted' }, } as ReturnType); - vi.spyOn(acpBridge, 'createAcpSessionBridge').mockImplementation(() => - makeFakeBridge(), - ); + const createBridge = vi + .spyOn(acpBridge, 'createAcpSessionBridge') + .mockImplementation(() => makeFakeBridge()); const snapshots = new Map(); const workerOptions = new Map< string, CreateChannelWorkerSupervisorOptions >(); + const workerSupervisors = new Map>(); const webhookEnqueues = new Map>(); const supervisorFactory = vi.fn( (options: CreateChannelWorkerSupervisorOptions) => { @@ -5343,10 +5733,13 @@ describe('runQwenServe channel worker supervisor', () => { accepted: true as const, })); webhookEnqueues.set(options.workspace, enqueueWebhookTask); - return { + const supervisor = { start: vi.fn(async () => { const capabilitiesResponse = await fetch( `${options.daemonUrl}/capabilities`, + { + headers: { Authorization: 'Bearer worker-remove-token' }, + }, ); expect(capabilitiesResponse.status).toBe(200); expect(await capabilitiesResponse.json()).toMatchObject({ @@ -5369,16 +5762,31 @@ describe('runQwenServe channel worker supervisor', () => { snapshot: vi.fn(() => snapshots.get(options.workspace)!), enqueueWebhookTask, }; + workerSupervisors.set( + options.workspace, + supervisor as ReturnType, + ); + return supervisor; }, ); const pidfile = makePidfileDeps(); + const removeByIds = vi.fn().mockResolvedValue(1); + const workspaceRegistrationStore = { + read: vi.fn().mockResolvedValue({ + schemaVersion: 1, + primaryWorkspace: primaryCwd, + workspaces: [secondaryCwd], + }), + removeByIds, + } as unknown as WorkspaceRegistrationStore; const handle = await runQwenServe( { port: 0, hostname: '127.0.0.1', mode: 'http-bridge', - workspace: [primary, secondary], + workspace: primary, + token: 'worker-remove-token', serveWebShell: false, channelSelection: { mode: 'names', @@ -5391,6 +5799,7 @@ describe('runQwenServe channel worker supervisor', () => { daemonLogBaseDir: path.join(tmpDir, 'debug'), channelWorkerSupervisorFactory: supervisorFactory, channelServicePidfile: pidfile, + workspaceRegistrationStore, deferRuntimeUntilFirstHealth: true, runtimeStartupTimeoutMs: 0, }, @@ -5409,13 +5818,13 @@ describe('runQwenServe channel worker supervisor', () => { }, ); expect(crossWorkspaceSecretResponse.status).toBe(401); - expect(supervisorFactory).not.toHaveBeenCalled(); const webhookResponse = await fetch( `${handle.url}/channels/feishu/webhooks/github-ci`, { method: 'POST', headers: { + Authorization: 'Bearer worker-remove-token', 'content-type': 'application/json', 'x-qwen-webhook-secret': 'secondary-secret', }, @@ -5465,6 +5874,78 @@ describe('runQwenServe channel worker supervisor', () => { ], }); + const capabilities = (await ( + await fetch(`${handle.url}/capabilities`, { + headers: { Authorization: 'Bearer worker-remove-token' }, + }) + ).json()) as { + workspaces: Array<{ id: string; cwd: string; removable?: boolean }>; + }; + const secondaryRuntime = capabilities.workspaces.find( + (workspace) => workspace.cwd === secondaryCwd, + ); + expect(secondaryRuntime).toMatchObject({ removable: true }); + const removalUrl = `${handle.url}/workspaces/${encodeURIComponent( + secondaryRuntime!.id, + )}`; + const busyRemoval = await fetch(removalUrl, { + method: 'DELETE', + headers: { Authorization: 'Bearer worker-remove-token' }, + }); + expect(busyRemoval.status).toBe(409); + await expect(busyRemoval.json()).resolves.toMatchObject({ + code: 'workspace_busy', + activity: { channelWorkers: 1 }, + }); + expect(workerSupervisors.get(secondaryCwd)!.stop).not.toHaveBeenCalled(); + + const forcedRemoval = await fetch(removalUrl, { + method: 'DELETE', + headers: { + Authorization: 'Bearer worker-remove-token', + 'content-type': 'application/json', + }, + body: JSON.stringify({ force: true }), + }); + expect(forcedRemoval.status).toBe(200); + await expect(forcedRemoval.json()).resolves.toMatchObject({ + removed: true, + activity: { channelWorkers: 1 }, + }); + expect(removeByIds).toHaveBeenCalledWith([ + workspaceRegistrationId(secondaryCwd), + ]); + const removedSupervisor = workerSupervisors.get(secondaryCwd)!; + const removedWorkerOptions = workerOptions.get(secondaryCwd)!; + expect(removedSupervisor.stop).toHaveBeenCalledOnce(); + expect(workerSupervisors.get(primaryCwd)!.stop).not.toHaveBeenCalled(); + expect(pidfile.writeServeServiceInfo).toHaveBeenLastCalledWith({ + channels: ['telegram'], + servePid: process.pid, + workerPid: 1234, + workers: [ + expect.objectContaining({ + workspaceCwd: primaryCwd, + channels: ['telegram'], + workerPid: 1234, + }), + ], + }); + + const removedWebhook = await fetch( + `${handle.url}/channels/feishu/webhooks/github-ci`, + { + method: 'POST', + headers: { + Authorization: 'Bearer worker-remove-token', + 'content-type': 'application/json', + 'x-qwen-webhook-secret': 'secondary-secret', + }, + body: JSON.stringify({ eventType: 'check_failed' }), + }, + ); + expect(removedWebhook.status).not.toBe(202); + const failedSecondary: ChannelWorkerSnapshot = { enabled: true, state: 'failed', @@ -5472,16 +5953,16 @@ describe('runQwenServe channel worker supervisor', () => { error: 'worker stopped', }; snapshots.set(secondaryCwd, failedSecondary); - workerOptions.get(secondaryCwd)!.onExit?.(failedSecondary); + removedWorkerOptions.onExit?.(failedSecondary); expect(pidfile.writeServeServiceInfo).toHaveBeenLastCalledWith( expect.objectContaining({ workerPid: 1234, - workers: expect.arrayContaining([ + workers: [ expect.objectContaining({ - workspaceCwd: secondaryCwd, - channels: ['feishu'], + workspaceCwd: primaryCwd, + channels: ['telegram'], }), - ]), + ], }), ); expect( @@ -5490,7 +5971,77 @@ describe('runQwenServe channel worker supervisor', () => { .workers?.find( (worker: ServiceInfoWorker) => worker.workspaceCwd === secondaryCwd, ), - ).not.toHaveProperty('workerPid'); + ).toBeUndefined(); + + snapshots.set(secondaryCwd, { + enabled: true, + state: 'running', + pid: 6789, + channels: ['feishu'], + }); + const readded = await fetch(`${handle.url}/workspaces`, { + method: 'POST', + headers: { + Authorization: 'Bearer worker-remove-token', + 'content-type': 'application/json', + }, + body: JSON.stringify({ cwd: secondaryCwd }), + }); + expect(readded.status).toBe(201); + expect(supervisorFactory).toHaveBeenCalledTimes(3); + expect(createBridge).toHaveBeenCalledTimes(3); + const replacementSupervisor = workerSupervisors.get(secondaryCwd)!; + expect(replacementSupervisor).not.toBe(removedSupervisor); + expect(replacementSupervisor.start).toHaveBeenCalledOnce(); + + const readdedWebhook = await fetch( + `${handle.url}/channels/feishu/webhooks/github-ci`, + { + method: 'POST', + headers: { + Authorization: 'Bearer worker-remove-token', + 'content-type': 'application/json', + 'x-qwen-webhook-secret': 'secondary-secret', + }, + body: JSON.stringify({ + eventType: 'check_failed', + targetRef: 'default', + title: 'CI failed again', + }), + }, + ); + expect(readdedWebhook.status).toBe(202); + expect(webhookEnqueues.get(secondaryCwd)).toHaveBeenCalledWith( + expect.objectContaining({ + channelName: 'feishu', + title: 'CI failed again', + }), + ); + + snapshots.set(secondaryCwd, failedSecondary); + workerOptions.get(secondaryCwd)!.onExit?.(failedSecondary); + const failedWorkerPidfile = + pidfile.writeServeServiceInfo.mock.calls.at(-1)?.[0]; + expect( + failedWorkerPidfile?.workers?.find( + (worker: ServiceInfoWorker) => worker.workspaceCwd === secondaryCwd, + ), + ).toMatchObject({ + workspaceCwd: secondaryCwd, + channels: ['feishu'], + }); + expect( + failedWorkerPidfile?.workers?.find( + (worker: ServiceInfoWorker) => worker.workspaceCwd === secondaryCwd, + )?.workerPid, + ).toBeUndefined(); + + const removeReplacement = await fetch(removalUrl, { + method: 'DELETE', + headers: { Authorization: 'Bearer worker-remove-token' }, + }); + expect(removeReplacement.status).toBe(200); + expect(replacementSupervisor.stop).toHaveBeenCalledOnce(); } finally { await handle.close(); if (previousSharedSecret === undefined) { diff --git a/packages/cli/src/serve/run-qwen-serve.ts b/packages/cli/src/serve/run-qwen-serve.ts index b100cdf12a3..e06b454bf0e 100644 --- a/packages/cli/src/serve/run-qwen-serve.ts +++ b/packages/cli/src/serve/run-qwen-serve.ts @@ -88,11 +88,15 @@ import type { WorkspaceRegistry, WorkspaceRuntime, } from './workspace-registry.js'; -import type { WorkspaceRegistrationStore } from './workspace-registration-store.js'; +import { + workspaceRegistrationId, + type WorkspaceRegistrationStore, +} from './workspace-registration-store.js'; import type { PermissionPolicy } from '@qwen-code/acp-bridge'; import { getCliVersion } from '../utils/version.js'; import { getRateLimiter } from './rate-limit.js'; import type { AcpHttpHandle } from './acp-http/index.js'; +import type { WorkspaceRuntimeRemovalController } from './routes/workspace-management.js'; import { allowOriginMode, listenerMaxConnections, @@ -957,6 +961,8 @@ function currentServeFeaturesForRunQwenServe( reloadAvailable: true, channelReloadAvailable: opts.channelSelection !== undefined, channelControlAvailable: true, + persistentWorkspaceRegistrationAvailable: true, + workspaceRuntimeRemovalAvailable: true, // Advertise the same WS feature flags as the runtime path (serve-features.ts) // so the bootstrap `/capabilities` window doesn't briefly under-report them. clientMcpOverWsEnabled: opts.clientMcpOverWs === true, @@ -1978,6 +1984,8 @@ export async function runQwenServe( const workspaceInputs = rawWorkspaces.map((workspace) => ({ raw: workspace, cwd: validateAndCanonicalizeWorkspace(workspace), + removable: false, + registrationIds: [] as string[], })); const boundWorkspace = workspaceInputs[0]!.cwd; @@ -2041,6 +2049,7 @@ export async function runQwenServe( try { const stored = await workspaceRegistrationStore.read(); for (const storedWorkspace of stored.workspaces) { + const registrationId = workspaceRegistrationId(storedWorkspace); let cwd: string; try { cwd = validateAndCanonicalizeWorkspace(storedWorkspace); @@ -2052,7 +2061,11 @@ export async function runQwenServe( ); continue; } - if (workspaceInputs.some((workspace) => workspace.cwd === cwd)) { + const existingInput = workspaceInputs.find( + (workspace) => workspace.cwd === cwd, + ); + if (existingInput) { + existingInput.registrationIds.push(registrationId); continue; } const nested = workspaceInputs.some( @@ -2076,7 +2089,12 @@ export async function runQwenServe( ); continue; } - workspaceInputs.push({ raw: storedWorkspace, cwd }); + workspaceInputs.push({ + raw: storedWorkspace, + cwd, + removable: true, + registrationIds: [registrationId], + }); } } catch (err) { writeStderrLine( @@ -2493,6 +2511,7 @@ export async function runQwenServe( // a single entry) keeps concurrent per-worker updates from losing each other. const isLiveWorker = (snapshot: ChannelWorkerGroupSnapshot): boolean => snapshot.state === 'running' || snapshot.state === 'starting'; + let channelWorkerPidfileUsesWorkers = workspaceInputs.length > 1; const writeChannelWorkerPidfile = (): void => { if (runtimeStartupError !== undefined) return; if (!channelPidfileReserved || !channelServicePidfile) return; @@ -2513,7 +2532,11 @@ export async function runQwenServe( const primary = snapshots.find((snapshot) => snapshot.primary); // Only surface the per-workspace worker list in multi-workspace mode; a // single-workspace daemon keeps the byte-identical channels/workerPid shape. - const includeWorkers = workspaceInputs.length > 1 && workers.length > 0; + if (workers.length > 1 || snapshots.some((snapshot) => !snapshot.primary)) { + channelWorkerPidfileUsesWorkers = true; + } + const includeWorkers = + channelWorkerPidfileUsesWorkers && workers.length > 0; try { channelServicePidfile.writeServeServiceInfo({ channels, @@ -2545,7 +2568,9 @@ export async function runQwenServe( | WorkspaceRegistry | undefined; const bridges = [ - ...(registry ? registry.list().map((runtime) => runtime.bridge) : []), + ...(registry + ? registry.listManaged().map((runtime) => runtime.bridge) + : []), ...(bridgeRef ? [bridgeRef] : []), ...internalRuntimeBridgesForCleanup, ]; @@ -3129,6 +3154,8 @@ export async function runQwenServe( workspaceCwd: boundWorkspace, primary: true, trusted: trustedWorkspace, + removable: false, + registrationIds: workspaceInputs[0]?.registrationIds ?? [], env: primaryRuntimeEnv, bridge, workspaceService, @@ -3200,6 +3227,15 @@ export async function runQwenServe( // (primary + secondaries). Called during shutdown so no new sub-sessions // are admitted while bridges are being torn down. const subSessionStoppers: Array<() => void> = []; + const subSessionStoppersByWorkspace = new Map void>(); + const runtimeCleanupPromises = new WeakMap< + WorkspaceRuntime, + Promise + >(); + const removeArrayValue = (values: T[], value: T): void => { + const index = values.indexOf(value); + if (index >= 0) values.splice(index, 1); + }; for (const workspaceInput of workspaceInputs.slice(1)) { let secondarySettings: @@ -3326,6 +3362,10 @@ export async function runQwenServe( runtimeBridges.push(secondaryBridge); internalRuntimeBridgesForCleanup.push(secondaryBridge); subSessionStoppers.push(secondarySubSessionLauncher.stop); + subSessionStoppersByWorkspace.set( + workspaceInput.cwd, + secondarySubSessionLauncher.stop, + ); const secondaryWorkspaceService = runtime.createDaemonWorkspaceService({ boundWorkspace: workspaceInput.cwd, contextFilename: contextFilenameForInit ?? 'QWEN.md', @@ -3404,6 +3444,8 @@ export async function runQwenServe( workspaceCwd: workspaceInput.cwd, primary: false, trusted: secondaryTrusted, + removable: workspaceInput.removable, + registrationIds: workspaceInput.registrationIds, env: secondaryEnv.metadata, bridge: secondaryBridge, workspaceService: secondaryWorkspaceService, @@ -3629,143 +3671,170 @@ export async function runQwenServe( getBridge: () => wsBridgeRef, boundWorkspace: cwd, }); - const wsBridge = runtime.createAcpSessionBridge({ - clientMcpSender: wsClientMcpRegistry.lookup, - onCreateSubSession: wsSubSessionLauncher.launch, - maxSessions: opts.maxSessions, - freshSessionAdmission: totalSessionAdmission.admit, - sessionLifecycle: sessionOwnerIndex.handleBridgeSessionLifecycle, - ...(opts.maxPendingPromptsPerSession !== undefined - ? { maxPendingPromptsPerSession: opts.maxPendingPromptsPerSession } - : {}), - ...(opts.eventRingSize !== undefined - ? { eventRingSize: opts.eventRingSize } - : {}), - ...(opts.compactedReplayMaxBytes !== undefined - ? { compactedReplayMaxBytes: opts.compactedReplayMaxBytes } - : {}), - ...(opts.channelIdleTimeoutMs !== undefined - ? { channelIdleTimeoutMs: opts.channelIdleTimeoutMs } - : {}), - ...(opts.sessionReapIntervalMs !== undefined - ? { sessionReapIntervalMs: opts.sessionReapIntervalMs } - : {}), - ...(opts.sessionIdleTimeoutMs !== undefined - ? { sessionIdleTimeoutMs: opts.sessionIdleTimeoutMs } - : {}), - ...(opts.permissionResponseTimeoutMs !== undefined - ? { permissionResponseTimeoutMs: opts.permissionResponseTimeoutMs } - : {}), - boundWorkspace: cwd, - sessionShellCommandEnabled, - childEnvOverrides, - channelFactory: wsChannelFactory, - onDiagnosticLine: diagnosticSink, - telemetry: createRuntimeBridgeTelemetry(wsHash), - ...(permissionPolicy !== undefined ? { permissionPolicy } : {}), - ...(permissionConsensusQuorum !== undefined - ? { permissionConsensusQuorum } - : {}), - permissionAudit: permissionAuditPublisher, - statusProvider: runtime.createDaemonStatusProvider({ - env: wsEnv.effectiveEnv, - }), - fileSystem: createBridgeFileSystemAdapter(wsFsFactory), - persistApprovalMode: (workspace, mode) => - withSettingsLock(workspace, async () => { - const fresh = settingsRuntime.settings.loadSettings(workspace); - fresh.setValue(WORKSPACE_SETTING_SCOPE, 'tools.approvalMode', mode); + let wsBridge: ReturnType; + try { + wsBridge = runtime.createAcpSessionBridge({ + clientMcpSender: wsClientMcpRegistry.lookup, + onCreateSubSession: wsSubSessionLauncher.launch, + maxSessions: opts.maxSessions, + freshSessionAdmission: totalSessionAdmission.admit, + sessionLifecycle: sessionOwnerIndex.handleBridgeSessionLifecycle, + ...(opts.maxPendingPromptsPerSession !== undefined + ? { maxPendingPromptsPerSession: opts.maxPendingPromptsPerSession } + : {}), + ...(opts.eventRingSize !== undefined + ? { eventRingSize: opts.eventRingSize } + : {}), + ...(opts.compactedReplayMaxBytes !== undefined + ? { compactedReplayMaxBytes: opts.compactedReplayMaxBytes } + : {}), + ...(opts.channelIdleTimeoutMs !== undefined + ? { channelIdleTimeoutMs: opts.channelIdleTimeoutMs } + : {}), + ...(opts.sessionReapIntervalMs !== undefined + ? { sessionReapIntervalMs: opts.sessionReapIntervalMs } + : {}), + ...(opts.sessionIdleTimeoutMs !== undefined + ? { sessionIdleTimeoutMs: opts.sessionIdleTimeoutMs } + : {}), + ...(opts.permissionResponseTimeoutMs !== undefined + ? { permissionResponseTimeoutMs: opts.permissionResponseTimeoutMs } + : {}), + boundWorkspace: cwd, + sessionShellCommandEnabled, + childEnvOverrides, + channelFactory: wsChannelFactory, + onDiagnosticLine: diagnosticSink, + telemetry: createRuntimeBridgeTelemetry(wsHash), + ...(permissionPolicy !== undefined ? { permissionPolicy } : {}), + ...(permissionConsensusQuorum !== undefined + ? { permissionConsensusQuorum } + : {}), + permissionAudit: permissionAuditPublisher, + statusProvider: runtime.createDaemonStatusProvider({ + env: wsEnv.effectiveEnv, }), - }); + fileSystem: createBridgeFileSystemAdapter(wsFsFactory), + persistApprovalMode: (workspace, mode) => + withSettingsLock(workspace, async () => { + const fresh = settingsRuntime.settings.loadSettings(workspace); + fresh.setValue( + WORKSPACE_SETTING_SCOPE, + 'tools.approvalMode', + mode, + ); + }), + }); + } catch (err) { + wsSubSessionLauncher.stop(); + throw err; + } wsBridgeRef = wsBridge; - const wsService = runtime.createDaemonWorkspaceService({ - boundWorkspace: cwd, - contextFilename: contextFilenameForInit ?? 'QWEN.md', - statusProvider: runtime.createDaemonStatusProvider({ - env: wsEnv.effectiveEnv, - }), - workspaceProvidersStatusProvider: - runtime.createWorkspaceProvidersStatusProvider({ + let wsService: ReturnType; + try { + wsService = runtime.createDaemonWorkspaceService({ + boundWorkspace: cwd, + contextFilename: contextFilenameForInit ?? 'QWEN.md', + statusProvider: runtime.createDaemonStatusProvider({ env: wsEnv.effectiveEnv, }), - workspaceSkillsStatusProvider: - runtime.createWorkspaceSkillsStatusProvider(), - isChannelLive: () => wsBridge.isChannelLive(), - preheatAcpChild: () => wsBridge.preheat(), - persistDisabledTools: persistDisabledToolsFn, - persistSetting: persistSettingFn, - persistSettings: persistSettingsFn, - reloadDaemonEnv: (workspace) => - withSettingsLock(workspace, async () => { - const fresh = settingsRuntime.settings.loadSettings(workspace, { - skipLoadEnvironment: true, - }); - const result = settingsRuntime.settings.reloadEnvironment( - fresh.merged, - workspace, - ); - // Mirror the startup secondary-workspace path: rebuild the runtime - // env snapshot and update the metadata so `.env` changes actually - // propagate to child processes spawned by this workspace's bridge. - try { - const refreshedRuntimeEnv = - settingsRuntime.environment.buildRuntimeEnvironment( - fresh.merged, - workspace, - daemonRuntimeBaseEnv, - ); - logRuntimeEnvFileReadFailures(workspace, refreshedRuntimeEnv); - wsEnv.replace(refreshedRuntimeEnv.effectiveEnv); - wsEnv.metadata.envFileReadFailed = - refreshedRuntimeEnv.envFileReadFailed; - wsEnv.metadata.envFileReadFailures.splice( - 0, - wsEnv.metadata.envFileReadFailures.length, - ...refreshedRuntimeEnv.envFileReadFailures, - ); - wsEnv.metadata.overlayKeys.splice( - 0, - wsEnv.metadata.overlayKeys.length, - ...refreshedRuntimeEnv.overlayKeys, - ); - wsEnv.metadata.envFilePaths.splice( - 0, - wsEnv.metadata.envFilePaths.length, - ...refreshedRuntimeEnv.envFilePaths, - ); - delete wsEnv.metadata.fallbackReason; - } catch (err) { - wsEnv.metadata.fallbackReason = - err instanceof Error ? err.message : String(err); - daemonLog.warn( - 'failed to rebuild dynamic runtime env snapshot after daemon env reload; preserving previous runtime env', - { - workspace, - error: wsEnv.metadata.fallbackReason, - }, + workspaceProvidersStatusProvider: + runtime.createWorkspaceProvidersStatusProvider({ + env: wsEnv.effectiveEnv, + }), + workspaceSkillsStatusProvider: + runtime.createWorkspaceSkillsStatusProvider(), + isChannelLive: () => wsBridge.isChannelLive(), + preheatAcpChild: () => wsBridge.preheat(), + persistDisabledTools: persistDisabledToolsFn, + persistSetting: persistSettingFn, + persistSettings: persistSettingsFn, + reloadDaemonEnv: (workspace) => + withSettingsLock(workspace, async () => { + const fresh = settingsRuntime.settings.loadSettings(workspace, { + skipLoadEnvironment: true, + }); + const result = settingsRuntime.settings.reloadEnvironment( + fresh.merged, + workspace, ); - } - return result; - }), - queryWorkspaceStatus: (method, idle) => - wsBridge.queryWorkspaceStatus(method, idle), - invokeWorkspaceCommand: (method, params, invokeOpts) => - wsBridge.invokeWorkspaceCommand(method, params, invokeOpts), - refreshExtensionsForAllSessions: () => - wsBridge.refreshExtensionsForAllSessions(), - publishWorkspaceEvent: (event) => wsBridge.publishWorkspaceEvent(event), - }); + // Mirror the startup secondary-workspace path: rebuild the runtime + // env snapshot and update the metadata so `.env` changes actually + // propagate to child processes spawned by this workspace's bridge. + try { + const refreshedRuntimeEnv = + settingsRuntime.environment.buildRuntimeEnvironment( + fresh.merged, + workspace, + daemonRuntimeBaseEnv, + ); + logRuntimeEnvFileReadFailures(workspace, refreshedRuntimeEnv); + wsEnv.replace(refreshedRuntimeEnv.effectiveEnv); + wsEnv.metadata.envFileReadFailed = + refreshedRuntimeEnv.envFileReadFailed; + wsEnv.metadata.envFileReadFailures.splice( + 0, + wsEnv.metadata.envFileReadFailures.length, + ...refreshedRuntimeEnv.envFileReadFailures, + ); + wsEnv.metadata.overlayKeys.splice( + 0, + wsEnv.metadata.overlayKeys.length, + ...refreshedRuntimeEnv.overlayKeys, + ); + wsEnv.metadata.envFilePaths.splice( + 0, + wsEnv.metadata.envFilePaths.length, + ...refreshedRuntimeEnv.envFilePaths, + ); + delete wsEnv.metadata.fallbackReason; + } catch (err) { + wsEnv.metadata.fallbackReason = + err instanceof Error ? err.message : String(err); + daemonLog.warn( + 'failed to rebuild dynamic runtime env snapshot after daemon env reload; preserving previous runtime env', + { + workspace, + error: wsEnv.metadata.fallbackReason, + }, + ); + } + return result; + }), + queryWorkspaceStatus: (method, idle) => + wsBridge.queryWorkspaceStatus(method, idle), + invokeWorkspaceCommand: (method, params, invokeOpts) => + wsBridge.invokeWorkspaceCommand(method, params, invokeOpts), + refreshExtensionsForAllSessions: () => + wsBridge.refreshExtensionsForAllSessions(), + publishWorkspaceEvent: (event) => + wsBridge.publishWorkspaceEvent(event), + }); + } catch (err) { + wsSubSessionLauncher.stop(); + await wsBridge.shutdown().catch(() => { + try { + wsBridge.killAllSync(); + } catch { + // Preserve the workspace-service construction error. + } + }); + throw err; + } // Register shared-array cleanup only after the runtime is fully built, so // a throw during createDaemonWorkspaceService (or any later step) can't // leave an orphaned bridge/channel in the shutdown arrays. runtimeBridges.push(wsBridge); internalRuntimeBridgesForCleanup.push(wsBridge); subSessionStoppers.push(wsSubSessionLauncher.stop); + subSessionStoppersByWorkspace.set(cwd, wsSubSessionLauncher.stop); return { workspaceId: wsHash, workspaceCwd: cwd, primary: false, trusted, + removable: true, + registrationIds: [], env: wsEnv.metadata, bridge: wsBridge, workspaceService: wsService, @@ -3774,10 +3843,143 @@ export async function runQwenServe( }; }; + const workspaceRuntimeRemoval = { + async runtimeAdded(runtimeAdded: WorkspaceRuntime): Promise { + const app = runtimeApp ?? runtimeAppForCleanup; + const startScheduledTaskKeepaliveForWorkspace = app?.locals?.[ + 'startScheduledTaskKeepaliveForWorkspace' + ] as ((runtime: WorkspaceRuntime) => void) | undefined; + startScheduledTaskKeepaliveForWorkspace?.(runtimeAdded); + if (!channelWorkerManager) return; + try { + await channelWorkerManager.restoreWorkspace( + runtimeAdded.workspaceCwd, + ); + await channelWorkerManager.refreshWorkspaces(); + } catch (err) { + daemonLog.error( + 'workspace channel worker startup error', + err instanceof Error ? err : null, + ); + } finally { + writeChannelWorkerPidfile(); + } + }, + beginDrain(runtimeToDrain: WorkspaceRuntime): void { + totalSessionAdmission.beginWorkspaceDrain(runtimeToDrain.workspaceCwd); + channelWorkerManager?.beginWorkspaceDrain(runtimeToDrain.workspaceCwd); + }, + cancelDrain(runtimeToDrain: WorkspaceRuntime): void { + channelWorkerManager?.cancelWorkspaceDrain(runtimeToDrain.workspaceCwd); + totalSessionAdmission.cancelWorkspaceDrain(runtimeToDrain.workspaceCwd); + }, + completeDrain(runtimeToDrain: WorkspaceRuntime): void { + totalSessionAdmission.completeWorkspaceDrain( + runtimeToDrain.workspaceCwd, + ); + }, + getActivity(runtimeToDrain: WorkspaceRuntime) { + return { + pendingSessionStarts: totalSessionAdmission.snapshotForWorkspace( + runtimeToDrain.workspaceCwd, + ).inFlight, + channelWorkers: + channelWorkerManager?.workspaceActivity( + runtimeToDrain.workspaceCwd, + ) ?? 0, + }; + }, + disposeRuntime( + runtimeToDrain: WorkspaceRuntime, + reason: 'daemon_shutdown' | 'workspace_removed' = 'workspace_removed', + ): Promise { + const existing = runtimeCleanupPromises.get(runtimeToDrain); + if (existing) return existing; + const cleanup = (async () => { + const stopSubSessions = subSessionStoppersByWorkspace.get( + runtimeToDrain.workspaceCwd, + ); + try { + stopSubSessions?.(); + } catch { + // Continue to bridge teardown. + } + if (reason === 'workspace_removed' && channelWorkerManager) { + await channelWorkerManager + .removeWorkspace(runtimeToDrain.workspaceCwd) + .catch((err) => { + daemonLog.error( + 'workspace channel worker cleanup error', + err instanceof Error ? err : null, + ); + }); + try { + await channelWorkerManager.refreshWorkspaces(); + } catch (err) { + channelWorkspaceGroups = (channelWorkspaceGroups ?? []).filter( + (group) => group.workspaceCwd !== runtimeToDrain.workspaceCwd, + ); + channelWebhookConfigVersion += 1; + refreshChannelWebhookConfigs?.(); + daemonLog.error( + 'workspace channel worker topology refresh error', + err instanceof Error ? err : null, + ); + } + writeChannelWorkerPidfile(); + } + if (reason === 'workspace_removed') { + const app = runtimeApp ?? runtimeAppForCleanup; + const stopWorkspaceGitStateForWorkspace = app?.locals?.[ + 'stopWorkspaceGitStateForWorkspace' + ] as ((workspaceCwd: string) => void) | undefined; + const stopScheduledTaskKeepaliveForWorkspace = app?.locals?.[ + 'stopScheduledTaskKeepaliveForWorkspace' + ] as ((workspaceCwd: string) => void) | undefined; + stopWorkspaceGitStateForWorkspace?.(runtimeToDrain.workspaceCwd); + stopScheduledTaskKeepaliveForWorkspace?.( + runtimeToDrain.workspaceCwd, + ); + } + let bridgeStopped = false; + try { + if (!shutdownBridges.has(runtimeToDrain.bridge)) { + await runtimeToDrain.bridge.shutdown({ reason }); + } + bridgeStopped = true; + } finally { + if (bridgeStopped || reason === 'workspace_removed') { + subSessionStoppersByWorkspace.delete(runtimeToDrain.workspaceCwd); + if (stopSubSessions) { + removeArrayValue(subSessionStoppers, stopSubSessions); + } + removeArrayValue(runtimeBridges, runtimeToDrain.bridge); + removeArrayValue( + internalRuntimeBridgesForCleanup, + runtimeToDrain.bridge, + ); + shutdownBridges.add(runtimeToDrain.bridge); + } + } + })(); + runtimeCleanupPromises.set(runtimeToDrain, cleanup); + void cleanup.catch(() => { + if ( + reason === 'daemon_shutdown' && + runtimeCleanupPromises.get(runtimeToDrain) === cleanup + ) { + runtimeCleanupPromises.delete(runtimeToDrain); + } + }); + return cleanup; + }, + }; + const app = runtime.createServeApp(opts, () => actualPort, { workspaceRegistry, createWorkspaceRuntime: createDynamicWorkspaceRuntime, workspaceRegistrationStore, + workspaceRuntimeRemoval, bridge, webShellDir, boundWorkspace, @@ -3898,6 +4100,7 @@ export async function runQwenServe( app.locals as { subSessionStoppers?: Array<() => void> } ).subSessionStoppers = subSessionStoppers; subSessionStoppers.push(subSessionLauncher.stop); + subSessionStoppersByWorkspace.set(boundWorkspace, subSessionLauncher.stop); return { app, bridge }; }; @@ -4669,25 +4872,15 @@ export async function runQwenServe( closePromise = new Promise((res, rej) => { shuttingDown = true; channelControlDraining = true; - // Stop the scheduled-task keepalive timer before tearing down the - // bridge it heartbeats. It's already unref()'d so it can't hold the - // process open, but stopping it here keeps it from firing against a - // disposed bridge (matters for embedders that don't process.exit). - ( - app.locals as { stopScheduledTaskKeepalive?: () => void } - ).stopScheduledTaskKeepalive?.(); - ( - app.locals as { stopWorkspaceGitState?: () => void } - ).stopWorkspaceGitState?.(); - // Same rationale for the create_sub_session launchers: stop accepting - // new sub-session spawns before the bridges are torn down. Calls - // every workspace's launcher stop (primary + secondaries). - const stoppers = ( - app.locals as { subSessionStoppers?: Array<() => void> } - ).subSessionStoppers; - if (stoppers) { - for (const stop of stoppers) stop(); - } + const initiallyMountedApp = runtimeApp ?? runtimeAppForCleanup; + const initiallyMountedManagement = initiallyMountedApp?.locals?.[ + 'workspaceManagementHandle' + ] as { sealAndWait?: () => Promise } | undefined; + // Calling an async function runs through its first await + // synchronously. Seal an already-mounted runtime before close() + // yields so no management request can enter the shutdown window. + const initialManagementWait = + initiallyMountedManagement?.sealAndWait?.(); clearRuntimeStartAfterHealthTimer(); clearRuntimeStartFallbackTimer(); cancelDeferredRuntimeStartup(); @@ -4797,6 +4990,29 @@ export async function runQwenServe( daemonLog, ); const appForCleanup = runtimeApp ?? runtimeAppForCleanup; + const workspaceManagementHandle = appForCleanup?.locals?.[ + 'workspaceManagementHandle' + ] as { sealAndWait?: () => Promise } | undefined; + await initialManagementWait; + if (workspaceManagementHandle !== initiallyMountedManagement) { + await workspaceManagementHandle?.sealAndWait?.(); + } + // Stop the scheduled-task keepalive only after workspace + // management has sealed and every accepted operation settled. + const stopScheduledTaskKeepalive = appForCleanup?.locals?.[ + 'stopScheduledTaskKeepalive' + ] as (() => void) | undefined; + stopScheduledTaskKeepalive?.(); + const stopWorkspaceGitState = appForCleanup?.locals?.[ + 'stopWorkspaceGitState' + ] as (() => void) | undefined; + stopWorkspaceGitState?.(); + const stoppers = appForCleanup?.locals?.[ + 'subSessionStoppers' + ] as Array<() => void> | undefined; + if (stoppers) { + for (const stop of stoppers) stop(); + } // Dispose the device-flow registry FIRST so any // in-flight IdP poll is cancelled and timers are cleared // before the bridge tear-down (which would otherwise race @@ -4853,7 +5069,36 @@ export async function runQwenServe( } else { removeCurrentServePidfile(); } + const runtimeRemoval = appForCleanup?.locals?.[ + 'workspaceRuntimeRemoval' + ] as WorkspaceRuntimeRemovalController | undefined; + const workspaceRegistry = appForCleanup?.locals?.[ + 'workspaceRegistry' + ] as WorkspaceRegistry | undefined; + const managedRuntimeBridges = new Set(); + if (runtimeRemoval && workspaceRegistry) { + const managedRuntimes = workspaceRegistry.listManaged(); + for (const workspaceRuntime of managedRuntimes) { + managedRuntimeBridges.add(workspaceRuntime.bridge); + await runtimeRemoval + .disposeRuntime(workspaceRuntime, 'daemon_shutdown') + .catch((err) => { + daemonLog.error( + 'workspace runtime shutdown error', + err instanceof Error ? err : null, + ); + bridgeShutdownError = + err instanceof Error ? err : new Error(String(err)); + try { + workspaceRuntime.bridge.killAllSync(); + } catch { + // Continue shutting down the remaining runtimes. + } + }); + } + } for (const bridgeForShutdown of getRuntimeBridgesForCleanup()) { + if (managedRuntimeBridges.has(bridgeForShutdown)) continue; if (shutdownBridges.has(bridgeForShutdown)) continue; shutdownBridges.add(bridgeForShutdown); await bridgeForShutdown.shutdown().catch((err) => { diff --git a/packages/cli/src/serve/server.test.ts b/packages/cli/src/serve/server.test.ts index 755f61f7648..91baefef8b2 100644 --- a/packages/cli/src/serve/server.test.ts +++ b/packages/cli/src/serve/server.test.ts @@ -80,6 +80,7 @@ import { SessionLimitExceededError, SessionNotFoundError, TotalSessionLimitExceededError, + WorkspaceDrainingError, WorkspaceMismatchError, type BridgeHeartbeatResult, type BridgeHeartbeatState, @@ -395,6 +396,7 @@ const EXPECTED_REGISTERED_FEATURES = [ 'channel_control', 'multi_workspace_sessions', 'persistent_workspace_registration', + 'workspace_runtime_removal', 'workspace_qualified_rest_core', 'workspace_persisted_transcript', 'workspace_qualified_acp', @@ -2256,6 +2258,24 @@ describe('createServeApp', () => { ); continue; } + if (feature === 'workspace_runtime_removal') { + expect(predicate({ workspaceRuntimeRemovalAvailable: true })).toBe( + true, + ); + expect(predicate({ workspaceRuntimeRemovalAvailable: false })).toBe( + false, + ); + expect(predicate({})).toBe(false); + expect( + getAdvertisedServeFeatures(undefined, { + workspaceRuntimeRemovalAvailable: true, + }), + ).toContain(feature); + expect(getAdvertisedServeFeatures(undefined, {})).not.toContain( + feature, + ); + continue; + } if (feature === 'workspace_qualified_acp') { // Advertised only when BOTH multi-workspace sessions and the HTTP ACP // surface are enabled. @@ -19468,6 +19488,27 @@ describe('T2.9 serve-side errorKind taxonomy (issue #4514)', () => { }); describe('sendBridgeError daemonLog routing', () => { + it('maps workspace drain admission failures to 503', async () => { + const bridge = fakeBridge({ + spawnImpl: async () => { + throw new WorkspaceDrainingError('/work/a'); + }, + }); + const app = createServeApp(baseOpts, undefined, { bridge }); + + const res = await request(app) + .post('/session') + .set('Host', `127.0.0.1:${baseOpts.port}`) + .send({ cwd: '/work/a' }); + + expect(res.status).toBe(503); + expect(res.headers['retry-after']).toBe('5'); + expect(res.body).toMatchObject({ + code: 'workspace_draining', + workspaceCwd: '/work/a', + }); + }); + it('routes 5xx errors through daemonLog when provided', async () => { const tmp = await fsp.mkdtemp(path.join(os.tmpdir(), 'daemon-log-')); const stderrLines: string[] = []; diff --git a/packages/cli/src/serve/server.ts b/packages/cli/src/serve/server.ts index 8bcffaf1ba4..958f033d5c6 100644 --- a/packages/cli/src/serve/server.ts +++ b/packages/cli/src/serve/server.ts @@ -176,7 +176,11 @@ import { registerWorkspaceLifecycleRoutes, registerWorkspaceQualifiedLifecycleRoutes, } from './routes/workspace-lifecycle.js'; -import { registerWorkspaceManagementRoutes } from './routes/workspace-management.js'; +import { + registerWorkspaceManagementRoutes, + type WorkspaceManagementHandle, + type WorkspaceRuntimeRemovalController, +} from './routes/workspace-management.js'; import type { WorkspaceRegistrationStore } from './workspace-registration-store.js'; import { registerWorkspaceGitRoutes, @@ -452,6 +456,7 @@ export interface ServeAppDeps { workspaceRegistry?: WorkspaceRegistry; createWorkspaceRuntime?: (cwd: string) => Promise; workspaceRegistrationStore?: WorkspaceRegistrationStore; + workspaceRuntimeRemoval?: WorkspaceRuntimeRemovalController; primaryWorkspaceTrusted?: boolean; primaryRuntimeEnv?: WorkspaceRuntimeEnvMetadata; voiceTranscriber?: WorkspaceVoiceRouteDeps['transcribe']; @@ -682,6 +687,8 @@ export function createServeApp( multiWorkspaceSessionsEnabled: () => workspaceRegistry.list().length > 1, persistentWorkspaceRegistrationAvailable: deps.workspaceRegistrationStore !== undefined, + workspaceRuntimeRemovalAvailable: + deps.workspaceRuntimeRemoval !== undefined, ...(primaryEffectiveEnv ? { env: primaryEffectiveEnv } : {}), }); const statusProvider = @@ -850,6 +857,12 @@ export function createServeApp( const workspaceGitState = new WorkspaceGitState(); (app.locals as { stopWorkspaceGitState?: () => void }).stopWorkspaceGitState = () => workspaceGitState.dispose(); + ( + app.locals as { + stopWorkspaceGitStateForWorkspace?: (workspaceCwd: string) => void; + } + ).stopWorkspaceGitStateForWorkspace = (workspaceCwd) => + workspaceGitState.disposeWorkspace(workspaceCwd); const workspaceQualifiedAcpEnabled = resolveAcpHttpEnabled(); // Order matters: rejection guards (CORS / Host allowlist / bearer auth) @@ -1025,7 +1038,10 @@ export function createServeApp( const buildWorkspaceCtx = createBuildWorkspaceCtx(primaryBoundWorkspace); const acpHandleRef: { current?: AcpHttpHandle } = {}; - const workspaceRememberLane = new WorkspaceRememberTaskLane(primaryBridge); + const workspaceRememberLane = new WorkspaceRememberTaskLane( + primaryBridge, + primaryBoundWorkspace, + ); // Plan C CDP tunnel (issue #5626): process-scoped registry pairing the // extension `/acp` connection with the `/cdp` puppeteer endpoint. Inert until @@ -1195,13 +1211,23 @@ export function createServeApp( }); // Dynamic workspace registration. - registerWorkspaceManagementRoutes(app, { + const workspaceManagementHandle = registerWorkspaceManagementRoutes(app, { workspaceRegistry, mutate, safeBody, createWorkspaceRuntime: deps.createWorkspaceRuntime, workspaceRegistrationStore: deps.workspaceRegistrationStore, + getAcpHandle: () => acpHandleRef.current, + runtimeRemoval: deps.workspaceRuntimeRemoval, }); + ( + app.locals as { workspaceManagementHandle?: WorkspaceManagementHandle } + ).workspaceManagementHandle = workspaceManagementHandle; + ( + app.locals as { + workspaceRuntimeRemoval?: WorkspaceRuntimeRemovalController; + } + ).workspaceRuntimeRemoval = deps.workspaceRuntimeRemoval; const broadcastSettingsChanged = ( key: string, @@ -1487,21 +1513,21 @@ export function createServeApp( // its own cron file + bridge, so a bound task created through the // workspace-qualified route fires (and survives a restart) exactly like a // primary-workspace one — otherwise a secondary workspace's tasks would be - // written to disk but silently never revived. The registry is fully - // populated (primary + every `--workspace`) before createServeApp runs. - // A workspace ADDED at runtime (Phase 4 add-workspace) isn't looped here, so - // its bound-task sessions aren't kept resident for the rest of this process; - // once persisted it registers as a boot workspace and is covered on the next - // restart, which is when its rehydration matters most anyway. - const keepaliveStops = workspaceRegistry.list().map((runtime) => { + // written to disk but silently never revived. + const keepaliveStops = new Map void>(); + const startKeepaliveForWorkspace = (runtime: WorkspaceRuntime) => { + if (keepaliveStops.has(runtime.workspaceCwd)) return; const keepalive = startScheduledTaskKeepalive({ bridge: runtime.bridge, boundWorkspace: runtime.workspaceCwd, intervalMs: keepaliveIntervalMs, }); rehydrateWorkspace(runtime.bridge, runtime.workspaceCwd); - return keepalive.stop; - }); + keepaliveStops.set(runtime.workspaceCwd, keepalive.stop); + }; + for (const runtime of workspaceRegistry.list()) { + startKeepaliveForWorkspace(runtime); + } // Park a combined stop fn on `app.locals` (same pattern as `fsFactory` / // `boundWorkspace` / `acpHandle` above) so the shutdown sequence in @@ -1510,8 +1536,24 @@ export function createServeApp( ( app.locals as { stopScheduledTaskKeepalive?: () => void } ).stopScheduledTaskKeepalive = () => { - for (const stop of keepaliveStops) stop(); + for (const stop of keepaliveStops.values()) stop(); + keepaliveStops.clear(); + }; + ( + app.locals as { + stopScheduledTaskKeepaliveForWorkspace?: (workspaceCwd: string) => void; + } + ).stopScheduledTaskKeepaliveForWorkspace = (workspaceCwd) => { + keepaliveStops.get(workspaceCwd)?.(); + keepaliveStops.delete(workspaceCwd); }; + ( + app.locals as { + startScheduledTaskKeepaliveForWorkspace?: ( + runtime: WorkspaceRuntime, + ) => void; + } + ).startScheduledTaskKeepaliveForWorkspace = startKeepaliveForWorkspace; } registerPermissionRoutes(app, { diff --git a/packages/cli/src/serve/server/error-response.ts b/packages/cli/src/serve/server/error-response.ts index 04c336d23ac..38b59ed0b67 100644 --- a/packages/cli/src/serve/server/error-response.ts +++ b/packages/cli/src/serve/server/error-response.ts @@ -45,6 +45,7 @@ import { WorkspaceInitRaceError, WorkspaceInitSymlinkError, WorkspaceMismatchError, + WorkspaceDrainingError, TotalSessionLimitExceededError, } from '../acp-session-bridge.js'; import type { DaemonLogger } from '../daemon-logger.js'; @@ -188,6 +189,15 @@ export function sendBridgeError( }); return; } + if (err instanceof WorkspaceDrainingError) { + res.set('Retry-After', '5'); + res.status(503).json({ + error: err.message, + code: 'workspace_draining', + workspaceCwd: err.workspaceCwd, + }); + return; + } if (err instanceof WorkspaceInitConflictError) { // The target file already exists with non- // whitespace content and the caller did not pass `force: true`. diff --git a/packages/cli/src/serve/server/serve-features.ts b/packages/cli/src/serve/server/serve-features.ts index 4285c57a4ea..6c9e5d515ce 100644 --- a/packages/cli/src/serve/server/serve-features.ts +++ b/packages/cli/src/serve/server/serve-features.ts @@ -49,6 +49,7 @@ interface CreateServeFeaturesDeps { sessionShellCommandEnabled: boolean; multiWorkspaceSessionsEnabled: () => boolean; persistentWorkspaceRegistrationAvailable: boolean; + workspaceRuntimeRemovalAvailable?: boolean; env?: Readonly>; } @@ -72,6 +73,7 @@ export function createServeFeatures( sessionShellCommandEnabled, multiWorkspaceSessionsEnabled, persistentWorkspaceRegistrationAvailable, + workspaceRuntimeRemovalAvailable, } = deps; const env = deps.env ?? process.env; let cachedVoiceTranscriptionAvailable: boolean | undefined; @@ -108,6 +110,7 @@ export function createServeFeatures( channelControlAvailable, multiWorkspaceSessionsEnabled: multiWorkspaceSessionsEnabled(), persistentWorkspaceRegistrationAvailable, + workspaceRuntimeRemovalAvailable, acpHttpEnabled: resolveAcpHttpEnabled(), clientMcpOverWsEnabled: opts.clientMcpOverWs === true, cdpTunnelOverWsEnabled: opts.cdpTunnelOverWs === true, diff --git a/packages/cli/src/serve/total-session-admission.test.ts b/packages/cli/src/serve/total-session-admission.test.ts index 660f287a5b7..ef914f22bdf 100644 --- a/packages/cli/src/serve/total-session-admission.test.ts +++ b/packages/cli/src/serve/total-session-admission.test.ts @@ -5,7 +5,10 @@ */ import { describe, expect, it } from 'vitest'; -import { TotalSessionLimitExceededError } from './acp-session-bridge.js'; +import { + TotalSessionLimitExceededError, + WorkspaceDrainingError, +} from './acp-session-bridge.js'; import { createTotalSessionAdmissionController } from './total-session-admission.js'; describe('createTotalSessionAdmissionController', () => { @@ -86,4 +89,61 @@ describe('createTotalSessionAdmissionController', () => { if (!nextReservation) throw new Error('expected reservation'); nextReservation.release(); }); + + it('tracks per-workspace reservations and supports drain rollback', () => { + const admission = createTotalSessionAdmissionController({ + getBridges: () => [], + }); + const reservation = admission.admit({ + operation: 'resume', + workspaceCwd: '/work/a', + }); + expect(admission.snapshotForWorkspace('/work/a')).toEqual({ + liveCount: 0, + inFlight: 1, + }); + + admission.beginWorkspaceDrain('/work/a'); + expect(() => + admission.admit({ operation: 'spawn', workspaceCwd: '/work/a' }), + ).toThrow(WorkspaceDrainingError); + admission.cancelWorkspaceDrain('/work/a'); + const afterRollback = admission.admit({ + operation: 'branch', + workspaceCwd: '/work/a', + }); + afterRollback?.release(); + reservation?.release(); + admission.completeWorkspaceDrain('/work/a'); + expect(admission.snapshotForWorkspace('/work/a').inFlight).toBe(0); + const afterCompletion = admission.admit({ + operation: 'spawn', + workspaceCwd: '/work/a', + }); + afterCompletion?.release(); + }); + + it('does not let an old reservation erase a replacement runtime count', () => { + const admission = createTotalSessionAdmissionController({ + getBridges: () => [], + }); + const oldReservation = admission.admit({ + operation: 'resume', + workspaceCwd: '/work/a', + }); + admission.beginWorkspaceDrain('/work/a'); + admission.completeWorkspaceDrain('/work/a'); + const replacementReservation = admission.admit({ + operation: 'spawn', + workspaceCwd: '/work/a', + }); + if (!oldReservation || !replacementReservation) { + throw new Error('expected reservations'); + } + + oldReservation.release(); + expect(admission.snapshotForWorkspace('/work/a').inFlight).toBe(1); + replacementReservation.release(); + expect(admission.snapshotForWorkspace('/work/a').inFlight).toBe(0); + }); }); diff --git a/packages/cli/src/serve/total-session-admission.ts b/packages/cli/src/serve/total-session-admission.ts index 814a73601ab..6f478c976e0 100644 --- a/packages/cli/src/serve/total-session-admission.ts +++ b/packages/cli/src/serve/total-session-admission.ts @@ -6,6 +6,7 @@ import { TotalSessionLimitExceededError, + WorkspaceDrainingError, type BridgeFreshSessionAdmission, type BridgeFreshSessionAdmissionContext, type BridgeFreshSessionReservation, @@ -28,6 +29,12 @@ export interface TotalSessionAdmissionSnapshot { export interface TotalSessionAdmissionController { readonly admit: BridgeFreshSessionAdmission; readonly snapshot: () => TotalSessionAdmissionSnapshot; + readonly snapshotForWorkspace: ( + workspaceCwd: string, + ) => TotalSessionAdmissionSnapshot; + readonly beginWorkspaceDrain: (workspaceCwd: string) => void; + readonly cancelWorkspaceDrain: (workspaceCwd: string) => void; + readonly completeWorkspaceDrain: (workspaceCwd: string) => void; } export function createTotalSessionAdmissionController({ @@ -35,6 +42,8 @@ export function createTotalSessionAdmissionController({ getBridges, }: TotalSessionAdmissionOptions): TotalSessionAdmissionController { let inFlight = 0; + const inFlightByWorkspace = new Map(); + const drainingWorkspaces = new Set(); const limit = maxTotalSessions === undefined || maxTotalSessions === 0 || @@ -46,6 +55,9 @@ export function createTotalSessionAdmissionController({ admit( context: BridgeFreshSessionAdmissionContext, ): BridgeFreshSessionReservation { + if (drainingWorkspaces.has(context.workspaceCwd)) { + throw new WorkspaceDrainingError(context.workspaceCwd); + } if (limit !== Number.POSITIVE_INFINITY) { if (getLiveCount(getBridges()) + inFlight >= limit) { throw Object.assign(new TotalSessionLimitExceededError(limit), { @@ -60,18 +72,47 @@ export function createTotalSessionAdmissionController({ } inFlight++; + inFlightByWorkspace.set( + context.workspaceCwd, + (inFlightByWorkspace.get(context.workspaceCwd) ?? 0) + 1, + ); let released = false; return { release() { if (released) return; released = true; inFlight--; + const workspaceInFlight = + (inFlightByWorkspace.get(context.workspaceCwd) ?? 1) - 1; + if (workspaceInFlight <= 0) { + inFlightByWorkspace.delete(context.workspaceCwd); + } else { + inFlightByWorkspace.set(context.workspaceCwd, workspaceInFlight); + } }, }; }, snapshot() { return { liveCount: getLiveCount(getBridges()), inFlight }; }, + snapshotForWorkspace(workspaceCwd) { + return { + // Per-workspace live sessions remain bridge-owned; removal reads + // `runtime.bridge.sessionCount`. This controller only owns admission + // reservations, so the aggregate snapshot shape uses zero here. + liveCount: 0, + inFlight: inFlightByWorkspace.get(workspaceCwd) ?? 0, + }; + }, + beginWorkspaceDrain(workspaceCwd) { + drainingWorkspaces.add(workspaceCwd); + }, + cancelWorkspaceDrain(workspaceCwd) { + drainingWorkspaces.delete(workspaceCwd); + }, + completeWorkspaceDrain(workspaceCwd) { + drainingWorkspaces.delete(workspaceCwd); + }, }; } diff --git a/packages/cli/src/serve/types.ts b/packages/cli/src/serve/types.ts index d13af838133..0102e365cb7 100644 --- a/packages/cli/src/serve/types.ts +++ b/packages/cli/src/serve/types.ts @@ -314,6 +314,7 @@ export interface CapabilitiesEnvelope { cwd: string; primary: boolean; trusted: boolean; + removable?: boolean; }>; /** * Transport families this daemon supports. Always includes `'rest'`; diff --git a/packages/cli/src/serve/workspace-git-state.test.ts b/packages/cli/src/serve/workspace-git-state.test.ts index af92cac59c1..a0fcaf7aebe 100644 --- a/packages/cli/src/serve/workspace-git-state.test.ts +++ b/packages/cli/src/serve/workspace-git-state.test.ts @@ -104,4 +104,55 @@ describe('WorkspaceGitState', () => { expect(resolveBranchNameMock).toHaveBeenCalledTimes(2); expect(watchRepoBranchMock).toHaveBeenCalledOnce(); }); + + it('disposes only the removed workspace watcher', async () => { + const firstDispose = vi.fn(); + const secondDispose = vi.fn(); + resolveBranchNameMock.mockResolvedValue('main'); + watchRepoBranchMock + .mockResolvedValueOnce(firstDispose) + .mockResolvedValueOnce(secondDispose); + const state = new WorkspaceGitState(); + const bridge = { + publishWorkspaceEvent: vi.fn(), + } as unknown as AcpSessionBridge; + await state.getStatus('/first', bridge); + await state.getStatus('/second', bridge); + + state.disposeWorkspace('/first'); + await vi.waitFor(() => expect(firstDispose).toHaveBeenCalledOnce()); + expect(secondDispose).not.toHaveBeenCalled(); + + state.dispose(); + await vi.waitFor(() => expect(secondDispose).toHaveBeenCalledOnce()); + }); + + it('keeps a replacement entry when the disposed creation later fails', async () => { + let rejectFirst!: (error: Error) => void; + const firstBranch = new Promise((_resolve, reject) => { + rejectFirst = reject; + }); + resolveBranchNameMock + .mockReturnValueOnce(firstBranch) + .mockResolvedValue('main'); + watchRepoBranchMock.mockResolvedValue(vi.fn()); + const state = new WorkspaceGitState(); + const bridge = { + publishWorkspaceEvent: vi.fn(), + } as unknown as AcpSessionBridge; + + const first = state.getStatus('/same', bridge); + state.disposeWorkspace('/same'); + await expect(state.getStatus('/same', bridge)).resolves.toMatchObject({ + branch: 'main', + }); + rejectFirst(new Error('old watcher failed')); + await expect(first).rejects.toThrow('old watcher failed'); + + await expect(state.getStatus('/same', bridge)).resolves.toMatchObject({ + branch: 'main', + }); + expect(resolveBranchNameMock).toHaveBeenCalledTimes(2); + state.dispose(); + }); }); diff --git a/packages/cli/src/serve/workspace-git-state.ts b/packages/cli/src/serve/workspace-git-state.ts index 47f252138e4..1d1e4e35b7d 100644 --- a/packages/cli/src/serve/workspace-git-state.ts +++ b/packages/cli/src/serve/workspace-git-state.ts @@ -36,6 +36,13 @@ export class WorkspaceGitState { this.entries.clear(); } + disposeWorkspace(workspaceCwd: string): void { + const pending = this.entries.get(workspaceCwd); + if (!pending) return; + this.entries.delete(workspaceCwd); + void pending.then((entry) => entry.dispose()).catch(() => {}); + } + private getOrCreateEntry( workspaceCwd: string, bridge: AcpSessionBridge, @@ -44,7 +51,9 @@ export class WorkspaceGitState { if (existing) return existing; const pending = this.createEntry(workspaceCwd, bridge).catch((error) => { - this.entries.delete(workspaceCwd); + if (this.entries.get(workspaceCwd) === pending) { + this.entries.delete(workspaceCwd); + } throw error; }); this.entries.set(workspaceCwd, pending); diff --git a/packages/cli/src/serve/workspace-registration-store.test.ts b/packages/cli/src/serve/workspace-registration-store.test.ts index 509a1007f3d..6f2f1fe99cd 100644 --- a/packages/cli/src/serve/workspace-registration-store.test.ts +++ b/packages/cli/src/serve/workspace-registration-store.test.ts @@ -7,7 +7,7 @@ import * as fs from 'node:fs/promises'; import * as path from 'node:path'; import { tmpdir } from 'node:os'; -import { afterEach, describe, expect, it } from 'vitest'; +import { afterEach, describe, expect, it, vi } from 'vitest'; import { WorkspaceRegistrationStore, WorkspaceRegistrationStoreError, @@ -88,6 +88,26 @@ describe('WorkspaceRegistrationStore', () => { }, ); + it('removes multiple raw and canonical registration identities atomically', async () => { + const home = await tempHome(); + const store = new WorkspaceRegistrationStore('/work/primary', home); + await store.add('/work/raw-alias-a'); + await store.add('/work/raw-alias-b'); + await store.add('/work/other'); + + await expect( + store.removeByIds([ + workspaceRegistrationId('/work/raw-alias-a'), + workspaceRegistrationId('/work/raw-alias-b'), + 'missing', + ]), + ).resolves.toBe(2); + await expect(store.read()).resolves.toMatchObject({ + workspaces: ['/work/other'], + }); + await expect(store.removeByIds([])).resolves.toBe(0); + }); + it('rejects invalid primary and primary-as-secondary inputs', async () => { const home = await tempHome(); expect(() => new WorkspaceRegistrationStore('relative', home)).toThrow( @@ -221,4 +241,90 @@ describe('WorkspaceRegistrationStore', () => { await expect(store.add('/work/secondary')).resolves.toBe(true); await expect(fs.stat(lockPath)).rejects.toMatchObject({ code: 'ENOENT' }); }); + + it('reports lock release failure after a committed write', async () => { + const home = await tempHome(); + vi.resetModules(); + vi.doMock('proper-lockfile', () => ({ + default: { + lock: vi.fn(async () => async () => { + throw new Error('release failed'); + }), + }, + })); + try { + const storeModule = await import('./workspace-registration-store.js'); + const store = new storeModule.WorkspaceRegistrationStore( + '/work/primary', + home, + ); + await fs.mkdir(path.dirname(store.filePath), { recursive: true }); + await fs.writeFile( + store.filePath, + JSON.stringify({ + schemaVersion: 1, + primaryWorkspace: '/work/primary', + workspaces: ['/work/secondary'], + }), + ); + + await expect( + store.removeByIds([ + storeModule.workspaceRegistrationId('/work/secondary'), + ]), + ).rejects.toBeInstanceOf( + storeModule.WorkspaceRegistrationStoreCommittedError, + ); + expect( + JSON.parse(await fs.readFile(store.filePath, 'utf8')), + ).toMatchObject({ workspaces: [] }); + } finally { + vi.doUnmock('proper-lockfile'); + vi.resetModules(); + } + }); + + it('preserves the write failure when lock release also fails', async () => { + const home = await tempHome(); + const writeError = new Error('write failed'); + const releaseError = new Error('release failed'); + vi.resetModules(); + vi.doMock('proper-lockfile', () => ({ + default: { + lock: vi.fn(async () => async () => { + throw releaseError; + }), + }, + })); + vi.doMock('@qwen-code/qwen-code-core', () => ({ + atomicWriteFile: vi.fn().mockRejectedValue(writeError), + })); + try { + const storeModule = await import('./workspace-registration-store.js'); + const store = new storeModule.WorkspaceRegistrationStore( + '/work/primary', + home, + ); + await fs.mkdir(path.dirname(store.filePath), { recursive: true }); + await fs.writeFile( + store.filePath, + JSON.stringify({ + schemaVersion: 1, + primaryWorkspace: '/work/primary', + workspaces: ['/work/secondary'], + }), + ); + + await expect( + store.removeByIds([ + storeModule.workspaceRegistrationId('/work/secondary'), + ]), + ).rejects.toBe(writeError); + expect(writeError.cause).toBe(releaseError); + } finally { + vi.doUnmock('proper-lockfile'); + vi.doUnmock('@qwen-code/qwen-code-core'); + vi.resetModules(); + } + }); }); diff --git a/packages/cli/src/serve/workspace-registration-store.ts b/packages/cli/src/serve/workspace-registration-store.ts index c4b11566f0b..73d404fcc8e 100644 --- a/packages/cli/src/serve/workspace-registration-store.ts +++ b/packages/cli/src/serve/workspace-registration-store.ts @@ -45,6 +45,8 @@ export class WorkspaceRegistrationStoreError extends Error { export class WorkspaceRegistrationStoreLimitError extends WorkspaceRegistrationStoreError {} +export class WorkspaceRegistrationStoreCommittedError extends WorkspaceRegistrationStoreError {} + function normalizedScopePath(primaryWorkspace: string): string { return os.platform() === 'win32' ? primaryWorkspace.toLowerCase() @@ -340,14 +342,24 @@ export class WorkspaceRegistrationStore { } async removeById(id: string): Promise { - return this.update((snapshot) => { - const index = snapshot.workspaces.findIndex( - (workspace) => workspaceRegistrationId(workspace) === id, - ); - if (index < 0) return false; - snapshot.workspaces.splice(index, 1); + return (await this.removeByIds([id])) > 0; + } + + async removeByIds(ids: readonly string[]): Promise { + const requested = new Set(ids); + if (requested.size === 0) return 0; + let removed = 0; + await this.update((snapshot) => { + const retained = snapshot.workspaces.filter((workspace) => { + if (!requested.has(workspaceRegistrationId(workspace))) return true; + removed++; + return false; + }); + if (removed === 0) return false; + snapshot.workspaces.splice(0, snapshot.workspaces.length, ...retained); return true; }); + return removed; } private async update( @@ -356,21 +368,47 @@ export class WorkspaceRegistrationStore { const { atomicWriteFile } = await import('@qwen-code/qwen-code-core'); return withInProcessLock(this.filePath, async () => { const lock = await acquireFileLock(this.filePath); + let committed = false; + let changed = false; + let workError: unknown; try { const snapshot = await this.read(); lock.assertOwned(); - const changed = mutate(snapshot); - if (!changed) return false; - lock.assertOwned(); - await atomicWriteFile( - this.filePath, - `${JSON.stringify(snapshot, null, 2)}\n`, - { mode: 0o600, forceMode: true, noFollow: true }, - ); - return true; - } finally { + changed = mutate(snapshot); + if (changed) { + lock.assertOwned(); + await atomicWriteFile( + this.filePath, + `${JSON.stringify(snapshot, null, 2)}\n`, + { mode: 0o600, forceMode: true, noFollow: true }, + ); + committed = true; + } + } catch (err) { + workError = err; + } + let releaseError: unknown; + try { await lock.release(); + } catch (err) { + releaseError = committed + ? new WorkspaceRegistrationStoreCommittedError( + `Workspace registration update committed but lock release failed: ${ + err instanceof Error ? err.message : String(err) + }`, + ) + : err; + } + if ( + workError instanceof Error && + releaseError !== undefined && + workError.cause === undefined + ) { + workError.cause = releaseError; } + if (workError !== undefined) throw workError; + if (releaseError !== undefined) throw releaseError; + return changed; }); } } diff --git a/packages/cli/src/serve/workspace-registry.test.ts b/packages/cli/src/serve/workspace-registry.test.ts index 0b054f6aea8..4108dd01072 100644 --- a/packages/cli/src/serve/workspace-registry.test.ts +++ b/packages/cli/src/serve/workspace-registry.test.ts @@ -304,4 +304,93 @@ describe('createWorkspaceRegistry', () => { expect(primarySummary).toHaveBeenCalledTimes(2); expect(secondarySummary).toHaveBeenCalledTimes(2); }); + + it('hides draining runtimes, rolls back, and releases cwd ownership on completion', () => { + const primary = makeRuntime('/work/primary', { + workspaceId: 'ws-primary', + primary: true, + }); + const secondary = makeRuntime('/work/secondary', { + workspaceId: 'ws-secondary', + removable: true, + }); + const sessionOwnerIndex = createWorkspaceSessionOwnerIndex(); + sessionOwnerIndex.register('session-secondary', secondary.workspaceCwd); + const registry = createWorkspaceRegistry([primary, secondary], { + sessionOwnerIndex, + }); + + expect(registry.beginDrain(primary)).toBe(false); + expect(registry.beginDrain(secondary)).toBe(true); + expect(registry.list()).toEqual([primary]); + expect(registry.listManaged()).toEqual([primary, secondary]); + expect(registry.getByWorkspaceId(secondary.workspaceId)).toBeUndefined(); + expect(registry.getManagedByWorkspaceId(secondary.workspaceId)).toBe( + secondary, + ); + + registry.cancelDrain(secondary); + expect(registry.list()).toEqual([primary, secondary]); + expect(registry.beginDrain(secondary)).toBe(true); + registry.completeDrain(secondary); + expect(registry.listManaged()).toEqual([primary]); + expect(sessionOwnerIndex.getWorkspaceCwds('session-secondary')).toEqual([]); + + const replacement = makeRuntime('/work/secondary', { + workspaceId: 'ws-secondary', + removable: true, + }); + registry.add(replacement); + expect(registry.getByWorkspaceCwd('/work/secondary')).toBe(replacement); + }); + + it('excludes draining owners from indexed and fallback session resolution', () => { + const primary = makeRuntime('/work/primary', { + workspaceId: 'ws-primary', + primary: true, + bridge: bridgeWithSummary((sessionId: string) => { + throw new SessionNotFoundError(sessionId); + }), + }); + const secondary = makeRuntime('/work/secondary', { + workspaceId: 'ws-secondary', + removable: true, + bridge: bridgeWithSummary((sessionId: string) => ({ + sessionId, + workspaceCwd: '/work/secondary', + })), + }); + const sessionOwnerIndex = createWorkspaceSessionOwnerIndex(); + sessionOwnerIndex.register('indexed', secondary.workspaceCwd); + const registry = createWorkspaceRegistry([primary, secondary], { + sessionOwnerIndex, + }); + + expect(registry.beginDrain(secondary)).toBe(true); + expect(registry.resolveLiveSessionOwner('indexed')).toEqual({ + kind: 'not_found', + }); + expect(registry.resolveLiveSessionOwner('fallback')).toEqual({ + kind: 'not_found', + }); + expect(sessionOwnerIndex.getWorkspaceCwds('indexed')).toEqual([ + secondary.workspaceCwd, + ]); + + registry.cancelDrain(secondary); + expect(registry.resolveLiveSessionOwner('indexed')).toEqual({ + kind: 'found', + runtime: secondary, + }); + expect(registry.resolveLiveSessionOwner('fallback')).toEqual({ + kind: 'found', + runtime: secondary, + }); + expect(sessionOwnerIndex.getWorkspaceCwds('indexed')).toEqual([ + secondary.workspaceCwd, + ]); + expect(sessionOwnerIndex.getWorkspaceCwds('fallback')).toEqual([ + secondary.workspaceCwd, + ]); + }); }); diff --git a/packages/cli/src/serve/workspace-registry.ts b/packages/cli/src/serve/workspace-registry.ts index 0dd4162cb2b..98b60fdd764 100644 --- a/packages/cli/src/serve/workspace-registry.ts +++ b/packages/cli/src/serve/workspace-registry.ts @@ -30,6 +30,10 @@ export interface WorkspaceRuntime { readonly workspaceCwd: string; readonly primary: boolean; readonly trusted: boolean; + /** Whether this runtime may be removed without restarting the daemon. */ + readonly removable?: boolean; + /** Persistent registration ids that restore this runtime on daemon startup. */ + readonly registrationIds?: readonly string[]; readonly env: WorkspaceRuntimeEnvMetadata; readonly bridge: AcpSessionBridge; readonly workspaceService: DaemonWorkspaceService; @@ -61,6 +65,7 @@ export interface WorkspaceSessionOwnerIndex { register(sessionId: string, workspaceCwd: string): void; remove(sessionId: string, workspaceCwd?: string): void; getWorkspaceCwds(sessionId: string): readonly string[]; + removeWorkspace(workspaceCwd: string): void; handleBridgeSessionLifecycle(event: WorkspaceSessionLifecycleEvent): void; } @@ -74,8 +79,16 @@ export interface WorkspaceRegistry { ): WorkspaceRuntime | undefined; resolveLiveSessionOwner(sessionId: string): WorkspaceSessionOwnerResolution; add(runtime: WorkspaceRuntime): void; + listManaged(): readonly WorkspaceRuntime[]; + getManagedByWorkspaceCwd(workspaceCwd: string): WorkspaceRuntime | undefined; + getManagedByWorkspaceId(workspaceId: string): WorkspaceRuntime | undefined; + beginDrain(runtime: WorkspaceRuntime): boolean; + cancelDrain(runtime: WorkspaceRuntime): void; + completeDrain(runtime: WorkspaceRuntime): void; } +type WorkspaceRuntimeState = 'active' | 'draining' | 'removed'; + export interface WorkspaceRegistryOptions { readonly sessionOwnerIndex?: WorkspaceSessionOwnerIndex; } @@ -109,6 +122,12 @@ export function createWorkspaceSessionOwnerIndex(): WorkspaceSessionOwnerIndex { register, remove, getWorkspaceCwds: (sessionId) => [...(bySessionId.get(sessionId) ?? [])], + removeWorkspace: (workspaceCwd) => { + for (const [sessionId, owners] of bySessionId) { + owners.delete(workspaceCwd); + if (owners.size === 0) bySessionId.delete(sessionId); + } + }, handleBridgeSessionLifecycle: (event) => { if (event.type === 'registered') { register(event.sessionId, event.workspaceCwd); @@ -159,6 +178,9 @@ export function createWorkspaceRegistry( } const runtimes: WorkspaceRuntime[] = [...inputRuntimes]; + const states = new WeakMap( + inputRuntimes.map((runtime) => [runtime, 'active']), + ); const primary = primaryRuntimes[0]!; const sessionOwnerIndex = options.sessionOwnerIndex; const scanLiveOwners = ( @@ -166,6 +188,7 @@ export function createWorkspaceRegistry( ): WorkspaceSessionOwnerResolution => { const matches: WorkspaceRuntime[] = []; for (const runtime of runtimes) { + if (states.get(runtime) !== 'active') continue; try { runtime.bridge.getSessionSummary(sessionId); matches.push(runtime); @@ -188,11 +211,31 @@ export function createWorkspaceRegistry( primary, // Return a frozen snapshot: `runtimes` is mutable internally (see `add`), // but callers must not be able to push/splice into the registry's state. - list: () => Object.freeze([...runtimes]) as readonly WorkspaceRuntime[], - getByWorkspaceCwd: (workspaceCwd) => byCwd.get(workspaceCwd), - getByWorkspaceId: (workspaceId) => byId.get(workspaceId), + list: () => + Object.freeze( + runtimes.filter((runtime) => states.get(runtime) === 'active'), + ) as readonly WorkspaceRuntime[], + listManaged: () => + Object.freeze([...runtimes]) as readonly WorkspaceRuntime[], + getByWorkspaceCwd: (workspaceCwd) => { + const runtime = byCwd.get(workspaceCwd); + return runtime && states.get(runtime) === 'active' ? runtime : undefined; + }, + getByWorkspaceId: (workspaceId) => { + const runtime = byId.get(workspaceId); + return runtime && states.get(runtime) === 'active' ? runtime : undefined; + }, + getManagedByWorkspaceCwd: (workspaceCwd) => byCwd.get(workspaceCwd), + getManagedByWorkspaceId: (workspaceId) => byId.get(workspaceId), resolveWorkspaceCwd: (workspaceCwd) => - workspaceCwd === undefined ? primary : byCwd.get(workspaceCwd), + workspaceCwd === undefined + ? primary + : (() => { + const runtime = byCwd.get(workspaceCwd); + return runtime && states.get(runtime) === 'active' + ? runtime + : undefined; + })(), add: (runtime) => { if (byCwd.has(runtime.workspaceCwd)) { throw new Error( @@ -207,6 +250,24 @@ export function createWorkspaceRegistry( byCwd.set(runtime.workspaceCwd, runtime); byId.set(runtime.workspaceId, runtime); runtimes.push(runtime); + states.set(runtime, 'active'); + }, + beginDrain: (runtime) => { + if (runtime.primary || states.get(runtime) !== 'active') return false; + states.set(runtime, 'draining'); + return true; + }, + cancelDrain: (runtime) => { + if (states.get(runtime) === 'draining') states.set(runtime, 'active'); + }, + completeDrain: (runtime) => { + if (runtime.primary || states.get(runtime) !== 'draining') return; + states.set(runtime, 'removed'); + byCwd.delete(runtime.workspaceCwd); + byId.delete(runtime.workspaceId); + const index = runtimes.indexOf(runtime); + if (index >= 0) runtimes.splice(index, 1); + sessionOwnerIndex?.removeWorkspace(runtime.workspaceCwd); }, resolveLiveSessionOwner: (sessionId) => { const indexedCwds = sessionOwnerIndex?.getWorkspaceCwds(sessionId) ?? []; @@ -214,10 +275,11 @@ export function createWorkspaceRegistry( const matches: WorkspaceRuntime[] = []; for (const workspaceCwd of indexedCwds) { const runtime = byCwd.get(workspaceCwd); - if (!runtime) { + if (!runtime || states.get(runtime) === 'removed') { sessionOwnerIndex?.remove(sessionId, workspaceCwd); continue; } + if (states.get(runtime) !== 'active') continue; try { runtime.bridge.getSessionSummary(sessionId); matches.push(runtime); diff --git a/packages/cli/src/serve/workspace-remember.test.ts b/packages/cli/src/serve/workspace-remember.test.ts index 641334e240e..22905672c36 100644 --- a/packages/cli/src/serve/workspace-remember.test.ts +++ b/packages/cli/src/serve/workspace-remember.test.ts @@ -16,6 +16,7 @@ import type { BridgeWorkspaceMemoryRememberRequest, BridgeWorkspaceMemoryRememberResult, } from './acp-session-bridge.js'; +import { WorkspaceDrainingError } from './acp-session-bridge.js'; import type { BridgeEvent } from '@qwen-code/acp-bridge/eventBus'; import { mountWorkspaceMemoryRememberRoutes, @@ -649,6 +650,76 @@ describe('workspace memory remember routes', () => { expect(lane.get(first.taskId)).toBeUndefined(); }); + it('rolls back the enqueue gate and fails queued tasks on disposal', async () => { + const first = deferred(); + const bridge = buildBridgeStub({ + rememberImpl: vi.fn(async () => first.promise), + }); + const lane = new WorkspaceRememberTaskLane(bridge, '/work/remove-me'); + const running = lane.enqueue({ + content: 'running', + contextMode: 'workspace', + }); + const queued = lane.enqueue({ + content: 'queued', + contextMode: 'workspace', + }); + await waitFor(() => lane.get(running.taskId)?.status === 'running'); + + lane.beginDrain(); + expect(() => + lane.enqueue({ content: 'blocked', contextMode: 'workspace' }), + ).toThrow(WorkspaceDrainingError); + lane.cancelDrain(); + const queuedAfterRollback = lane.enqueue({ + content: 'queued after rollback', + contextMode: 'workspace', + }); + + lane.dispose(); + expect(lane.get(queued.taskId)).toMatchObject({ + status: 'failed', + error: { code: 'workspace_removed' }, + }); + expect(lane.get(queuedAfterRollback.taskId)).toMatchObject({ + status: 'failed', + error: { code: 'workspace_removed' }, + }); + expect(lane.pendingCount()).toBe(1); + + first.reject(new Error('bridge closed')); + await waitFor(() => lane.get(running.taskId)?.status === 'failed'); + expect(lane.get(running.taskId)).toMatchObject({ + error: { code: 'workspace_removed' }, + }); + expect(bridge.events).toEqual([]); + expect(bridge.rememberCalls.map((call) => call.content)).toEqual([ + 'running', + ]); + }); + + it('fails a successful bridge result that settles after disposal', async () => { + const first = deferred(); + const bridge = buildBridgeStub({ + rememberImpl: vi.fn(async () => first.promise), + }); + const lane = new WorkspaceRememberTaskLane(bridge, '/work/remove-me'); + const running = lane.enqueue({ + content: 'running', + contextMode: 'workspace', + }); + await waitFor(() => lane.get(running.taskId)?.status === 'running'); + + lane.dispose(); + first.resolve({ filesTouched: [], touchedScopes: [] }); + + await waitFor(() => lane.get(running.taskId)?.status === 'failed'); + expect(lane.get(running.taskId)).toMatchObject({ + error: { code: 'workspace_removed' }, + }); + expect(bridge.events).toEqual([]); + }); + it('runs hidden remember tasks serially within the remember lane', async () => { const first = deferred(); const second = deferred(); @@ -937,6 +1008,24 @@ describe('workspace memory remember routes', () => { ); }); + it('maps a draining workspace to a stable 503 response', async () => { + const bridge = buildBridgeStub({}); + const lane = new WorkspaceRememberTaskLane(bridge, '/work/draining'); + lane.beginDrain(); + const app = buildApp(bridge, undefined, lane); + + await request(app) + .post('/workspace/memory/remember') + .send({ content: 'remember me' }) + .expect(503) + .expect((res) => { + expect(res.body).toEqual({ + error: 'Workspace runtime is being removed.', + code: 'workspace_draining', + }); + }); + }); + it('records bridge failures with stable public error codes', async () => { const bridge = buildBridgeStub({ rememberImpl: vi diff --git a/packages/cli/src/serve/workspace-remember.ts b/packages/cli/src/serve/workspace-remember.ts index 7ab7422dbc2..cc3a20c6f91 100644 --- a/packages/cli/src/serve/workspace-remember.ts +++ b/packages/cli/src/serve/workspace-remember.ts @@ -15,6 +15,7 @@ import type { BridgeWorkspaceMemoryRememberContextMode, BridgeWorkspaceMemoryRememberResult, } from './acp-session-bridge.js'; +import { WorkspaceDrainingError } from './acp-session-bridge.js'; import { createWorkspaceMemoryExtractionErrorLogger, shouldSuppressRememberErrorDetails, @@ -169,6 +170,9 @@ export function publicErrorMessage( ? 'Workspace memory remember queue is full.' : 'Workspace memory task queue is full.'; } + if (code === 'workspace_draining') { + return 'Workspace runtime is being removed.'; + } if ( code === 'remember_timeout' || code === 'forget_timeout' || @@ -180,6 +184,7 @@ export function publicErrorMessage( } export function publicErrorStatus(code: string): number { + if (code === 'workspace_draining') return 503; if (code === 'remember_queue_full') return 429; if (code === 'managed_memory_unavailable') return 409; return 500; @@ -216,8 +221,53 @@ export class WorkspaceRememberTaskLane { ); private readonly tasks = new Map(); private tail: Promise = Promise.resolve(); + private draining = false; + private disposed = false; + + constructor( + private readonly bridge: AcpSessionBridge, + private readonly workspaceCwd = 'workspace', + ) {} + + beginDrain(): void { + this.draining = true; + } + + cancelDrain(): void { + if (!this.disposed) this.draining = false; + } - constructor(private readonly bridge: AcpSessionBridge) {} + pendingCount(): number { + return this.pendingCounts().total; + } + + dispose(): void { + if (this.disposed) return; + this.disposed = true; + this.draining = true; + for (const task of this.tasks.values()) { + if (task.status !== 'queued') continue; + task.status = 'failed'; + task.updatedAt = nowIso(); + task.error = createTaskError( + 'workspace_removed', + task.kind, + 'Workspace runtime was removed before the task started.', + ); + } + } + + private failRunningTaskAfterRemoval(task: WorkspaceMemoryTaskRecord): void { + if (!this.disposed || task.status === 'completed') return; + task.status = 'failed'; + delete task.result; + task.updatedAt = nowIso(); + task.error = createTaskError( + 'workspace_removed', + task.kind, + 'Workspace runtime was removed while the task was running.', + ); + } private pendingCounts(): { total: number; nonRemember: number } { let total = 0; @@ -260,6 +310,9 @@ export class WorkspaceRememberTaskLane { } private assertCapacity(kind: WorkspaceMemoryTaskRecord['kind']): void { + if (this.draining || this.disposed) { + throw new WorkspaceDrainingError(this.workspaceCwd); + } const pending = this.pendingCounts(); if (pending.total >= WorkspaceRememberTaskLane.MAX_PENDING) { throw Object.assign(new Error('Workspace memory task queue is full'), { @@ -288,7 +341,11 @@ export class WorkspaceRememberTaskLane { this.tasks.set(task.taskId, task); this.evictTerminalTasks(); - this.tail = this.tail.then(run, run); + const runIfQueued = async () => { + if (task.status !== 'queued') return; + await run(); + }; + this.tail = this.tail.then(runIfQueued, runIfQueued); void this.tail.catch((err: unknown) => { debugLogger.error('Unhandled task lane error:', err); }); @@ -345,16 +402,18 @@ export class WorkspaceRememberTaskLane { content: params.content, contextMode: params.contextMode, }); - task.status = 'completed'; - task.result = { - summary: - result.filesTouched.length > 0 - ? 'Memory update completed.' - : 'No memory files updated.', - filesTouched: result.filesTouched, - touchedScopes: result.touchedScopes, - }; - task.updatedAt = nowIso(); + if (!this.disposed) { + task.status = 'completed'; + task.result = { + summary: + result.filesTouched.length > 0 + ? 'Memory update completed.' + : 'No memory files updated.', + filesTouched: result.filesTouched, + touchedScopes: result.touchedScopes, + }; + task.updatedAt = nowIso(); + } } catch (err) { const code = workspaceMemoryFailureCode( err, @@ -375,6 +434,7 @@ export class WorkspaceRememberTaskLane { task.error = createTaskError(code, task.kind, diagnostics.details); task.updatedAt = nowIso(); } + this.failRunningTaskAfterRemoval(task); try { if (task.status === 'completed' && task.result) { this.publishManagedMemoryChanged({ @@ -417,16 +477,18 @@ export class WorkspaceRememberTaskLane { const result = await this.bridge.runWorkspaceMemoryForget({ query: params.query, }); - task.status = 'completed'; - task.result = { - summary: - result.summary ?? - formatWorkspaceMemoryForgetSummary(result.removedEntries.length), - removedEntries: result.removedEntries, - touchedTopics: result.touchedTopics, - touchedScopes: result.touchedScopes, - }; - task.updatedAt = nowIso(); + if (!this.disposed) { + task.status = 'completed'; + task.result = { + summary: + result.summary ?? + formatWorkspaceMemoryForgetSummary(result.removedEntries.length), + removedEntries: result.removedEntries, + touchedTopics: result.touchedTopics, + touchedScopes: result.touchedScopes, + }; + task.updatedAt = nowIso(); + } } catch (err) { const code = workspaceMemoryFailureCode( err, @@ -447,6 +509,7 @@ export class WorkspaceRememberTaskLane { task.error = createTaskError(code, task.kind, diagnostics.details); task.updatedAt = nowIso(); } + this.failRunningTaskAfterRemoval(task); try { if (task.status === 'completed' && task.result) { this.publishManagedMemoryChanged({ @@ -486,15 +549,17 @@ export class WorkspaceRememberTaskLane { task.updatedAt = nowIso(); try { const result = await this.bridge.runWorkspaceMemoryDream(); - task.status = 'completed'; - task.result = { - summary: - result.summary ?? - formatWorkspaceMemoryDreamSummary(result.touchedTopics.length), - touchedTopics: result.touchedTopics, - dedupedEntries: result.dedupedEntries, - }; - task.updatedAt = nowIso(); + if (!this.disposed) { + task.status = 'completed'; + task.result = { + summary: + result.summary ?? + formatWorkspaceMemoryDreamSummary(result.touchedTopics.length), + touchedTopics: result.touchedTopics, + dedupedEntries: result.dedupedEntries, + }; + task.updatedAt = nowIso(); + } } catch (err) { const code = workspaceMemoryFailureCode( err, @@ -515,6 +580,7 @@ export class WorkspaceRememberTaskLane { task.error = createTaskError(code, task.kind, diagnostics.details); task.updatedAt = nowIso(); } + this.failRunningTaskAfterRemoval(task); try { if (task.status === 'completed' && task.result) { this.publishManagedMemoryChanged({ diff --git a/packages/cli/src/serve/workspace-route-runtime.ts b/packages/cli/src/serve/workspace-route-runtime.ts index 93cf93cffc7..7558729014a 100644 --- a/packages/cli/src/serve/workspace-route-runtime.ts +++ b/packages/cli/src/serve/workspace-route-runtime.ts @@ -69,6 +69,39 @@ export function resolveRegisteredWorkspaceRuntimeByPathSelector( ); } +export function resolveManagedWorkspaceRuntimeByPathSelector( + registry: WorkspaceRegistry, + selector: string, +): WorkspaceRuntime | undefined { + const exact = registry.getManagedByWorkspaceCwd(selector); + if (exact) return exact; + + if (path.isAbsolute(selector) && !isUncPath(selector)) { + try { + const canonicalSelector = canonicalizeWorkspace(selector); + const canonicalMatch = + registry.getManagedByWorkspaceCwd(canonicalSelector); + if (canonicalMatch) return canonicalMatch; + for (const runtime of registry.listManaged()) { + if (canonicalizeWorkspace(runtime.workspaceCwd) === canonicalSelector) { + return runtime; + } + } + } catch { + // Fall through to lexical matching for unavailable paths. + } + } + + const normalizedSelector = normalizePortableAbsolutePath(selector); + return registry + .listManaged() + .find( + (runtime) => + normalizePortableAbsolutePath(runtime.workspaceCwd) === + normalizedSelector, + ); +} + export function resolveWorkspaceRuntimeFromParam( registry: WorkspaceRegistry, req: Request, @@ -98,6 +131,39 @@ export function resolveWorkspaceRuntimeFromParam( return runtime; } +export function resolveManagedWorkspaceRuntimeFromParam( + registry: WorkspaceRegistry, + req: Request, + res: Response, + paramName = 'workspace', +): WorkspaceRuntime | null { + const selector = req.params[paramName] ?? ''; + const byId = registry.getManagedByWorkspaceId(selector); + if (byId) return byId; + + if (!isPortableAbsolutePath(selector)) { + res.status(400).json({ + error: `\`:${paramName}\` must decode to a workspace id or absolute path`, + code: 'workspace_mismatch', + }); + return null; + } + + const runtime = resolveManagedWorkspaceRuntimeByPathSelector( + registry, + selector, + ); + if (!runtime) { + res.status(400).json({ + error: + 'Workspace mismatch: the requested workspace is not registered with this daemon.', + code: 'workspace_mismatch', + }); + return null; + } + return runtime; +} + export function requireTrustedWorkspaceRuntime( runtime: WorkspaceRuntime, res: Response, diff --git a/packages/sdk-typescript/src/daemon/DaemonClient.ts b/packages/sdk-typescript/src/daemon/DaemonClient.ts index 617d7902483..b880a559d07 100644 --- a/packages/sdk-typescript/src/daemon/DaemonClient.ts +++ b/packages/sdk-typescript/src/daemon/DaemonClient.ts @@ -86,6 +86,7 @@ import type { DaemonWorkspaceMemoryForgetTask, DaemonWorkspaceMemoryRememberOptions, DaemonWorkspaceMemoryRememberTask, + DaemonWorkspaceRemovalResult, HeartbeatResult, PermissionResponse, PromptContentBlock, @@ -3602,6 +3603,31 @@ export class WorkspaceDaemonClient { return this.get('/memory', 'GET /workspaces/:workspace/memory'); } + remove(options?: { + force?: boolean; + timeoutMs?: number; + }): Promise { + const body = + options?.force === undefined + ? undefined + : { + force: options.force, + }; + return this.client.workspaceJsonRequest( + this.workspaceSelector, + '', + 'DELETE /workspaces/:workspace', + { + method: 'DELETE', + ...(body ? { body } : {}), + ...(options?.timeoutMs !== undefined + ? { timeoutMs: options.timeoutMs } + : {}), + mode: 'rest', + }, + ); + } + writeWorkspaceMemory( req: Omit & { scope?: 'workspace' }, clientId?: string, diff --git a/packages/sdk-typescript/src/daemon/index.ts b/packages/sdk-typescript/src/daemon/index.ts index 94f37c9dd9d..bed3c975f1b 100644 --- a/packages/sdk-typescript/src/daemon/index.ts +++ b/packages/sdk-typescript/src/daemon/index.ts @@ -380,6 +380,8 @@ export type { DaemonWorkspaceTrustStatus, DaemonWorkspaceCapability, DaemonWorkspaceGitStatus, + DaemonWorkspaceRemovalActivity, + DaemonWorkspaceRemovalResult, DaemonAvailableCommand, DaemonArchiveSessionsResult, DaemonCapabilities, diff --git a/packages/sdk-typescript/src/daemon/types.ts b/packages/sdk-typescript/src/daemon/types.ts index 2d01001824b..eb13bedae12 100644 --- a/packages/sdk-typescript/src/daemon/types.ts +++ b/packages/sdk-typescript/src/daemon/types.ts @@ -31,6 +31,26 @@ export interface DaemonWorkspaceCapability { cwd: string; primary: boolean; trusted: boolean; + /** Whether this runtime can be removed without restarting the daemon. */ + removable?: boolean; +} + +export interface DaemonWorkspaceRemovalActivity { + sessions: number; + activePrompts: number; + pendingSessionStarts: number; + acpConnections: number; + memoryTasks: number; + channelWorkers: number; +} + +export interface DaemonWorkspaceRemovalResult { + removed: true; + workspaceId: string; + workspaceCwd: string; + forced: boolean; + persistedRegistrationRemoved: boolean; + activity: DaemonWorkspaceRemovalActivity; } /** Current Git branch metadata returned from a workspace Git status route. */ diff --git a/packages/sdk-typescript/src/index.ts b/packages/sdk-typescript/src/index.ts index 29b2ebdaabb..0d914621bd1 100644 --- a/packages/sdk-typescript/src/index.ts +++ b/packages/sdk-typescript/src/index.ts @@ -72,6 +72,8 @@ export { type DaemonWorkspaceInitializedEvent, type DaemonWorkspaceCapability, type DaemonWorkspaceGitStatus, + type DaemonWorkspaceRemovalActivity, + type DaemonWorkspaceRemovalResult, type DaemonAvailableCommand, type DaemonCapabilities, type DaemonEnvCell, diff --git a/packages/sdk-typescript/test/unit/DaemonClient.test.ts b/packages/sdk-typescript/test/unit/DaemonClient.test.ts index 22804a4d4f4..9af102fad10 100644 --- a/packages/sdk-typescript/test/unit/DaemonClient.test.ts +++ b/packages/sdk-typescript/test/unit/DaemonClient.test.ts @@ -4925,6 +4925,58 @@ describe('DaemonClient', () => { 'http://daemon/workspaces/%2Ftmp%2Fwork%20space/session/session-1/organization', ); }); + + it('removes a workspace by id or cwd with an optional force body', async () => { + const result = { + removed: true as const, + workspaceId: 'workspace/id', + workspaceCwd: '/tmp/work space', + forced: true, + persistedRegistrationRemoved: true, + activity: { + sessions: 1, + activePrompts: 0, + pendingSessionStarts: 0, + acpConnections: 0, + memoryTasks: 0, + channelWorkers: 0, + }, + }; + const { fetch, calls } = recordingFetch(() => jsonResponse(200, result)); + const transportFetch = vi.fn(async () => + jsonResponse(404, { error: 'transport route not mapped' }), + ); + const transport: DaemonTransport = { + type: 'acp-http', + supportsReplay: true, + connected: true, + fetch: transportFetch, + async *subscribeEvents() {}, + dispose() {}, + }; + const client = new DaemonClient({ + baseUrl: 'http://daemon', + fetch, + transport, + }); + + await expect( + client.workspaceById('workspace/id').remove({ force: true }), + ).resolves.toEqual(result); + await expect( + client.workspaceByCwd('/tmp/work space').remove(), + ).resolves.toEqual(result); + + expect(calls.map((call) => [call.method, call.url, call.body])).toEqual([ + [ + 'DELETE', + 'http://daemon/workspaces/workspace%2Fid', + JSON.stringify({ force: true }), + ], + ['DELETE', 'http://daemon/workspaces/%2Ftmp%2Fwork%20space', null], + ]); + expect(transportFetch).not.toHaveBeenCalled(); + }); }); describe('addRuntimeMcpServer (T2.8 #4514)', () => { diff --git a/packages/web-shell/client/components/sidebar/WebShellSidebar.module.css b/packages/web-shell/client/components/sidebar/WebShellSidebar.module.css index 94186086844..9f6e2b58048 100644 --- a/packages/web-shell/client/components/sidebar/WebShellSidebar.module.css +++ b/packages/web-shell/client/components/sidebar/WebShellSidebar.module.css @@ -1143,6 +1143,14 @@ line-height: 22px; } +.workspaceRemovalActivityList { + margin: 0; + padding-left: 18px; + color: var(--muted-foreground); + font-size: 14px; + line-height: 22px; +} + .confirmActions { display: flex; justify-content: flex-end; diff --git a/packages/web-shell/client/components/sidebar/WebShellSidebar.tsx b/packages/web-shell/client/components/sidebar/WebShellSidebar.tsx index d8709b5e90b..b8d2fecd257 100644 --- a/packages/web-shell/client/components/sidebar/WebShellSidebar.tsx +++ b/packages/web-shell/client/components/sidebar/WebShellSidebar.tsx @@ -17,6 +17,7 @@ import { useWorkspace, useWorkspaceActions, } from '@qwen-code/webui/daemon-react-sdk'; +import { DaemonHttpError } from '@qwen-code/sdk/daemon'; import type { DaemonSessionGroup, DaemonSessionGroupColor, @@ -24,6 +25,7 @@ import type { DaemonSessionGroupPresetColor, DaemonSessionSummary, DaemonWorkspaceCapability, + DaemonWorkspaceRemovalActivity, } from '@qwen-code/sdk/daemon'; import { ActivityIcon, @@ -509,6 +511,18 @@ export function WebShellSidebar({ const [projectsExpanded, setProjectsExpanded] = useState(true); const [searchOpen, setSearchOpen] = useState(false); const [showAddWorkspaceDialog, setShowAddWorkspaceDialog] = useState(false); + const [workspaceRemovalCandidate, setWorkspaceRemovalCandidate] = + useState(null); + const [workspaceRemovalActivity, setWorkspaceRemovalActivity] = + useState(null); + const [workspaceRemovalSubmitting, setWorkspaceRemovalSubmitting] = + useState(false); + const workspaceRemovalMountedRef = useRef(false); + const workspaceRemovalDismissedRef = useRef(false); + const [ + workspaceRemovalRemoteInProgress, + setWorkspaceRemovalRemoteInProgress, + ] = useState(false); const [workspaceSessionsReloadToken, setWorkspaceSessionsReloadToken] = useState(0); const [autoExpandWorkspace, setAutoExpandWorkspace] = useState<{ @@ -521,6 +535,14 @@ export function WebShellSidebar({ const bumpWorkspaceReload = useCallback(() => { setWorkspaceSessionsReloadToken((v) => v + 1); }, []); + + useEffect(() => { + workspaceRemovalMountedRef.current = true; + return () => { + workspaceRemovalMountedRef.current = false; + workspaceRemovalDismissedRef.current = true; + }; + }, []); const [searchQuery, setSearchQuery] = useState(''); const [isResizing, setIsResizing] = useState(false); const [completedUnreadIds, setCompletedUnreadIds] = useState>( @@ -535,6 +557,9 @@ export function WebShellSidebar({ null, ); const currentSessionId = connection.sessionId; + const workspaceRemovalEnabled = Boolean( + connection.capabilities?.features?.includes('workspace_runtime_removal'), + ); const canExportSessions = connection.capabilities?.features?.includes('session_export') ?? false; const projectName = @@ -976,6 +1001,181 @@ export function WebShellSidebar({ [t, workspaceActions, workspace], ); + const reconcileRemovedWorkspace = useCallback( + async (removed: DaemonWorkspaceCapability) => { + if (!workspaceRemovalMountedRef.current) return; + if (selectedWorkspaceCwd === removed.cwd) { + onSelectWorkspace?.(undefined); + } + setWorkspaceSessionsReloadToken((token) => token + 1); + try { + await workspace.refreshCapabilities?.(); + } catch { + // The mutation already converged; a later refresh will reconcile. + } + if (!workspaceRemovalMountedRef.current) return; + setWorkspaceRemovalCandidate(null); + setWorkspaceRemovalActivity(null); + setWorkspaceRemovalRemoteInProgress(false); + void reload().catch(() => undefined); + void reloadArchived().catch(() => undefined); + }, + [ + onSelectWorkspace, + reload, + reloadArchived, + selectedWorkspaceCwd, + workspace, + ], + ); + + const requestWorkspaceRemoval = useCallback( + (candidate: DaemonWorkspaceCapability) => { + if (workspaceRemovalSubmitting) return; + workspaceRemovalDismissedRef.current = false; + setWorkspaceRemovalActivity(null); + setWorkspaceRemovalRemoteInProgress(false); + setWorkspaceRemovalCandidate(candidate); + }, + [workspaceRemovalSubmitting], + ); + + const confirmWorkspaceRemoval = useCallback(async () => { + const candidate = workspaceRemovalCandidate; + if (!candidate || workspaceRemovalSubmitting) return; + const force = workspaceRemovalActivity !== null; + if ( + force && + connection.sessionId && + connection.workspaceCwd === candidate.cwd + ) { + return; + } + setWorkspaceRemovalSubmitting(true); + try { + await workspaceActions.removeWorkspace(candidate.id, { force }); + await reconcileRemovedWorkspace(candidate); + } catch (error) { + if (!workspaceRemovalMountedRef.current) return; + if (error instanceof DaemonHttpError) { + const body = error.body as + | { + code?: unknown; + activity?: DaemonWorkspaceRemovalActivity; + } + | undefined; + if ( + error.status === 409 && + body?.code === 'workspace_busy' && + body.activity + ) { + setWorkspaceRemovalActivity(body.activity); + return; + } + if (error.status === 400 && body?.code === 'workspace_mismatch') { + await reconcileRemovedWorkspace(candidate); + return; + } + if ( + error.status === 409 && + (body?.code === 'workspace_removal_in_progress' || + body?.code === 'workspace_registration_in_progress') + ) { + setWorkspaceRemovalRemoteInProgress(true); + let lastError: unknown = error; + let exhaustedTransientRetries = true; + for (let attempt = 0; attempt < 20; attempt++) { + if ( + !workspaceRemovalMountedRef.current || + workspaceRemovalDismissedRef.current + ) { + return; + } + await new Promise((resolve) => window.setTimeout(resolve, 250)); + if ( + !workspaceRemovalMountedRef.current || + workspaceRemovalDismissedRef.current + ) { + return; + } + try { + await workspaceActions.removeWorkspace(candidate.id, { force }); + await reconcileRemovedWorkspace(candidate); + return; + } catch (retryError) { + if (!workspaceRemovalMountedRef.current) return; + lastError = retryError; + if (retryError instanceof DaemonHttpError) { + const retryBody = retryError.body as + | { + code?: unknown; + activity?: DaemonWorkspaceRemovalActivity; + } + | undefined; + if ( + retryError.status === 400 && + retryBody?.code === 'workspace_mismatch' + ) { + await reconcileRemovedWorkspace(candidate); + return; + } + if ( + retryError.status === 409 && + retryBody?.code === 'workspace_busy' && + retryBody.activity + ) { + setWorkspaceRemovalRemoteInProgress(false); + setWorkspaceRemovalActivity(retryBody.activity); + return; + } + if ( + retryError.status === 409 && + (retryBody?.code === 'workspace_removal_in_progress' || + retryBody?.code === 'workspace_registration_in_progress') + ) { + continue; + } + } + exhaustedTransientRetries = false; + break; + } + } + if ( + !workspaceRemovalMountedRef.current || + workspaceRemovalDismissedRef.current + ) { + return; + } + setWorkspaceRemovalRemoteInProgress(false); + onError( + exhaustedTransientRetries + ? new Error( + 'Workspace removal remained in progress after retries.', + ) + : lastError, + t('sidebar.removeWorkspaceError'), + ); + return; + } + } + onError(error, t('sidebar.removeWorkspaceError')); + } finally { + if (workspaceRemovalMountedRef.current) { + setWorkspaceRemovalSubmitting(false); + } + } + }, [ + connection.sessionId, + connection.workspaceCwd, + onError, + reconcileRemovedWorkspace, + t, + workspaceActions, + workspaceRemovalActivity, + workspaceRemovalCandidate, + workspaceRemovalSubmitting, + ]); + const handleNewSession = useCallback( (workspaceCwd?: string) => { if (creatingSessionRef.current) return; @@ -2495,6 +2695,115 @@ export function WebShellSidebar({ )} + {workspaceRemovalCandidate && ( + { + if ( + !workspaceRemovalSubmitting || + workspaceRemovalRemoteInProgress + ) { + workspaceRemovalDismissedRef.current = true; + setWorkspaceRemovalCandidate(null); + setWorkspaceRemovalActivity(null); + setWorkspaceRemovalRemoteInProgress(false); + } + }} + > +
+

+ {workspaceRemovalActivity + ? t('sidebar.removeWorkspaceBusy', { + name: workspaceRemovalCandidate.cwd, + }) + : t('sidebar.removeWorkspaceConfirm', { + name: workspaceRemovalCandidate.cwd, + })} +

+ {workspaceRemovalActivity && ( +
    +
  • + {t('sidebar.removeWorkspaceSessions', { + count: workspaceRemovalActivity.sessions, + })} +
    • +
  • + {t('sidebar.removeWorkspacePrompts', { + count: workspaceRemovalActivity.activePrompts, + })} +
    • +
  • + {t('sidebar.removeWorkspaceStarts', { + count: workspaceRemovalActivity.pendingSessionStarts, + })} +
    • +
  • + {t('sidebar.removeWorkspaceConnections', { + count: workspaceRemovalActivity.acpConnections, + })} +
    • +
  • + {t('sidebar.removeWorkspaceMemoryTasks', { + count: workspaceRemovalActivity.memoryTasks, + })} +
    • +
  • + {t('sidebar.removeWorkspaceWorkers', { + count: workspaceRemovalActivity.channelWorkers, + })} +
    • +
+ )} + {workspaceRemovalActivity && + connection.sessionId && + connection.workspaceCwd === workspaceRemovalCandidate.cwd && ( +

+ {t('sidebar.removeWorkspaceCurrentSession')} +

+ )} + {workspaceRemovalRemoteInProgress && ( +

+ {t('sidebar.removeWorkspaceInProgress')} +

+ )} +
+ + +
+
+ + )} {groupEditor && ( ( -
- - -
- )} + {ws.trusted && ( + <> + + + + )} + {canRemove && ( + + + + + + + requestWorkspaceRemoval(ws) + } + > + + {t('sidebar.removeWorkspace')} + + + + )} + + ); + }} /> {ws.primary && (projectExpanded || searchQuery.trim()) ? ( diff --git a/packages/web-shell/client/components/sidebar/WebShellSidebar.workspace-removal.test.tsx b/packages/web-shell/client/components/sidebar/WebShellSidebar.workspace-removal.test.tsx new file mode 100644 index 00000000000..6647faeb41c --- /dev/null +++ b/packages/web-shell/client/components/sidebar/WebShellSidebar.workspace-removal.test.tsx @@ -0,0 +1,278 @@ +// @vitest-environment jsdom +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { act } from 'react'; +import { createRoot, type Root } from 'react-dom/client'; +import { + DaemonHttpError, + type DaemonWorkspaceCapability, +} from '@qwen-code/sdk/daemon'; + +const { connection, workspace, workspaceActions, active, archived } = + vi.hoisted(() => { + const makeSessions = () => ({ + sessions: [], + loading: false, + error: null, + reload: vi.fn().mockResolvedValue(undefined), + deleteSession: vi.fn().mockResolvedValue(true), + archiveSession: vi.fn().mockResolvedValue(true), + unarchiveSession: vi.fn().mockResolvedValue(true), + exportSession: vi.fn(), + }); + return { + connection: { + status: 'connected', + sessionId: null as string | null, + workspaceCwd: '/tmp/project', + capabilities: undefined as + | { + qwenCodeVersion: string; + features: string[]; + workspaces: DaemonWorkspaceCapability[]; + } + | undefined, + }, + workspace: { + capabilities: undefined as + | { + qwenCodeVersion: string; + features: string[]; + workspaces: DaemonWorkspaceCapability[]; + } + | undefined, + client: { + workspaceByCwd: vi.fn(() => ({ + listWorkspaceSessions: vi.fn().mockResolvedValue([]), + listSessionGroups: vi.fn().mockResolvedValue({ groups: [] }), + })), + }, + refreshCapabilities: vi.fn(), + }, + workspaceActions: { + addWorkspace: vi.fn(), + removeWorkspace: vi.fn(), + listSessionGroups: vi.fn(), + }, + active: makeSessions(), + archived: makeSessions(), + }; + }); + +vi.mock('@qwen-code/webui/daemon-react-sdk', () => ({ + useConnection: () => connection, + useActions: () => ({ renameSession: vi.fn() }), + useWorkspace: () => workspace, + useWorkspaceActions: () => workspaceActions, + useSessions: (options?: { archiveState?: string }) => + options?.archiveState === 'archived' ? archived : active, +})); + +const { I18nProvider } = await import('../../i18n'); +const { WebShellSidebar } = await import('./WebShellSidebar'); + +globalThis.IS_REACT_ACT_ENVIRONMENT = true; +if (!globalThis.PointerEvent) { + globalThis.PointerEvent = MouseEvent as typeof PointerEvent; +} +if (!Element.prototype.hasPointerCapture) { + Element.prototype.hasPointerCapture = () => false; +} +if (!Element.prototype.setPointerCapture) { + Element.prototype.setPointerCapture = () => {}; +} +if (!Element.prototype.releasePointerCapture) { + Element.prototype.releasePointerCapture = () => {}; +} +if (!Element.prototype.scrollIntoView) { + Element.prototype.scrollIntoView = () => {}; +} + +const capabilities = { + qwenCodeVersion: '1.2.3', + features: ['multi_workspace_sessions', 'workspace_runtime_removal'], + workspaces: [ + { + id: 'primary', + cwd: '/tmp/project', + primary: true, + trusted: true, + removable: false, + }, + { + id: 'secondary', + cwd: '/tmp/other', + primary: false, + trusted: true, + removable: true, + }, + { + id: 'untrusted', + cwd: '/tmp/danger', + primary: false, + trusted: false, + removable: true, + }, + ], +} satisfies NonNullable; + +let root: Root; +let container: HTMLDivElement; + +function renderSidebar( + overrides: { + selectedWorkspaceCwd?: string; + onSelectWorkspace?: (cwd: string | undefined) => void; + onError?: (error: unknown, message: string) => void; + } = {}, +) { + act(() => { + root.render( + + {}} + onOpenSettings={() => {}} + onOpenDaemonStatus={() => {}} + onOpenScheduledTasks={() => {}} + onOpenSessions={() => {}} + onOpenSplitView={() => {}} + onNewSession={() => false} + onLoadSession={() => {}} + onError={overrides.onError ?? (() => {})} + selectedWorkspaceCwd={overrides.selectedWorkspaceCwd} + onSelectWorkspace={overrides.onSelectWorkspace} + /> + , + ); + }); +} + +function workspaceAction(cwd: string): HTMLButtonElement | undefined { + return Array.from( + container.querySelectorAll( + 'button[aria-label="Workspace actions"]', + ), + ).find((button) => + button.parentElement?.parentElement?.textContent?.includes( + cwd.split('/').at(-1)!, + ), + ); +} + +function click(element: HTMLElement): void { + element.dispatchEvent( + new PointerEvent('pointerdown', { bubbles: true, button: 0 }), + ); + element.dispatchEvent(new PointerEvent('pointerup', { bubbles: true })); + element.dispatchEvent(new MouseEvent('click', { bubbles: true })); +} + +function openRemoval(cwd: string): void { + const trigger = workspaceAction(cwd); + expect(trigger).toBeDefined(); + act(() => click(trigger!)); + const item = document.body.querySelector( + `[aria-label="Remove workspace: ${cwd}"]`, + ); + expect(item).not.toBeNull(); + act(() => click(item!)); +} + +function dialogButton(label: string): HTMLButtonElement { + const button = Array.from( + document.body.querySelectorAll('button'), + ).find((candidate) => candidate.textContent === label); + expect(button).toBeDefined(); + return button!; +} + +beforeEach(() => { + container = document.createElement('div'); + document.body.appendChild(container); + root = createRoot(container); + connection.sessionId = null; + connection.workspaceCwd = '/tmp/project'; + connection.capabilities = capabilities; + workspace.capabilities = capabilities; + workspace.refreshCapabilities.mockReset(); + workspace.refreshCapabilities.mockResolvedValue(capabilities); + workspaceActions.removeWorkspace.mockReset(); + workspaceActions.removeWorkspace.mockResolvedValue({ removed: true }); + active.reload.mockClear(); + archived.reload.mockClear(); +}); + +afterEach(() => { + act(() => root.unmount()); + container.remove(); + vi.restoreAllMocks(); +}); + +describe('WebShellSidebar workspace removal', () => { + it('hides removal when the daemon does not publish the feature', () => { + connection.capabilities = { + ...capabilities, + features: ['multi_workspace_sessions'], + }; + renderSidebar(); + + expect(workspaceAction('/tmp/other')).toBeUndefined(); + expect(workspaceAction('/tmp/danger')).toBeUndefined(); + }); + + it('exposes removal for an untrusted removable workspace', () => { + renderSidebar(); + + expect(workspaceAction('/tmp/danger')).toBeDefined(); + expect(workspaceAction('/tmp/project')).toBeUndefined(); + }); + + it('removes the selected workspace and falls back to primary', async () => { + const onSelectWorkspace = vi.fn(); + renderSidebar({ + selectedWorkspaceCwd: '/tmp/danger', + onSelectWorkspace, + }); + openRemoval('/tmp/danger'); + + await act(async () => click(dialogButton('Remove workspace'))); + + expect(workspaceActions.removeWorkspace).toHaveBeenCalledWith('untrusted', { + force: false, + }); + expect(onSelectWorkspace).toHaveBeenCalledWith(undefined); + expect(workspace.refreshCapabilities).toHaveBeenCalled(); + }); + + it('shows activity and blocks force for the current session workspace', async () => { + connection.sessionId = 'active-session'; + connection.workspaceCwd = '/tmp/other'; + workspaceActions.removeWorkspace.mockRejectedValueOnce( + new DaemonHttpError( + 409, + { + code: 'workspace_busy', + activity: { + sessions: 1, + activePrompts: 1, + pendingSessionStarts: 0, + acpConnections: 1, + memoryTasks: 0, + channelWorkers: 0, + }, + }, + 'busy', + ), + ); + renderSidebar(); + openRemoval('/tmp/other'); + + await act(async () => click(dialogButton('Remove workspace'))); + + expect(document.body.textContent).toContain('Sessions: 1'); + expect(document.body.textContent).toContain( + 'Switch to another workspace or close the current session', + ); + expect(dialogButton('Force remove').disabled).toBe(true); + }); +}); diff --git a/packages/web-shell/client/components/sidebar/WorkspaceSection.tsx b/packages/web-shell/client/components/sidebar/WorkspaceSection.tsx index 1bcb3bbc3e2..bbf837a1cb0 100644 --- a/packages/web-shell/client/components/sidebar/WorkspaceSection.tsx +++ b/packages/web-shell/client/components/sidebar/WorkspaceSection.tsx @@ -270,7 +270,7 @@ export function WorkspaceSection({ )} {readOnly && {readOnlyLabel}} - {!readOnly && !disabled && headerActions?.(actionsVisible)} + {headerActions?.(actionsVisible)} {renderSessions && (expanded || Boolean(searchQuery.trim())) && diff --git a/packages/web-shell/client/i18n.tsx b/packages/web-shell/client/i18n.tsx index e2457a6961f..ecb675f52a8 100644 --- a/packages/web-shell/client/i18n.tsx +++ b/packages/web-shell/client/i18n.tsx @@ -827,6 +827,27 @@ const EN: Messages = { 'sidebar.addWorkspacePersistHint': 'Persist this workspace registration in the daemon configuration.', 'sidebar.addWorkspaceAdding': 'Adding…', + 'sidebar.removeWorkspace': 'Remove workspace', + 'sidebar.workspaceActions': 'Workspace actions', + 'sidebar.forceRemoveWorkspace': 'Force remove', + 'sidebar.removeWorkspaceTitle': 'Remove Workspace', + 'sidebar.removeWorkspaceConfirm': (v) => + `Remove the runtime and persistent registration for “${v?.name ?? ''}”? Files, settings, and session history will not be deleted.`, + 'sidebar.removeWorkspaceBusy': (v) => + `“${v?.name ?? ''}” still has active runtime resources. Force removal will terminate them.`, + 'sidebar.removeWorkspaceCurrentSession': + 'Switch to another workspace or close the current session before forcing removal.', + 'sidebar.removeWorkspaceInProgress': + 'Another client is already removing this workspace. Refresh after that operation finishes.', + 'sidebar.removeWorkspaceError': 'Failed to remove workspace', + 'sidebar.removeWorkspaceSessions': (v) => `Sessions: ${v?.count ?? 0}`, + 'sidebar.removeWorkspacePrompts': (v) => `Active prompts: ${v?.count ?? 0}`, + 'sidebar.removeWorkspaceStarts': (v) => + `Pending session starts: ${v?.count ?? 0}`, + 'sidebar.removeWorkspaceConnections': (v) => + `ACP connections: ${v?.count ?? 0}`, + 'sidebar.removeWorkspaceMemoryTasks': (v) => `Memory tasks: ${v?.count ?? 0}`, + 'sidebar.removeWorkspaceWorkers': (v) => `Channel workers: ${v?.count ?? 0}`, 'sidebar.noSessions': 'No sessions.', 'sidebar.projectFallback': 'Project', 'sidebar.sessionsOverview': 'Session Overview', @@ -2667,6 +2688,25 @@ const ZH: Messages = { 'sidebar.addWorkspacePersist': '服务重启后保留', 'sidebar.addWorkspacePersistHint': '将此工作区注册持久化到守护进程配置中。', 'sidebar.addWorkspaceAdding': '添加中…', + 'sidebar.removeWorkspace': '移除工作区', + 'sidebar.workspaceActions': '工作区操作', + 'sidebar.forceRemoveWorkspace': '强制移除', + 'sidebar.removeWorkspaceTitle': '移除工作区', + 'sidebar.removeWorkspaceConfirm': (v) => + `确定移除“${v?.name ?? ''}”的运行时和持久化注册吗？文件、设置和会话历史不会被删除。`, + 'sidebar.removeWorkspaceBusy': (v) => + `“${v?.name ?? ''}”仍有活动运行时资源。强制移除会终止这些资源。`, + 'sidebar.removeWorkspaceCurrentSession': + '请先切换到其他工作区或关闭当前会话，再执行强制移除。', + 'sidebar.removeWorkspaceInProgress': + '另一个客户端正在移除此工作区，请在操作完成后刷新。', + 'sidebar.removeWorkspaceError': '移除工作区失败', + 'sidebar.removeWorkspaceSessions': (v) => `会话：${v?.count ?? 0}`, + 'sidebar.removeWorkspacePrompts': (v) => `活动提示：${v?.count ?? 0}`, + 'sidebar.removeWorkspaceStarts': (v) => `待启动会话：${v?.count ?? 0}`, + 'sidebar.removeWorkspaceConnections': (v) => `ACP 连接：${v?.count ?? 0}`, + 'sidebar.removeWorkspaceMemoryTasks': (v) => `Memory 任务：${v?.count ?? 0}`, + 'sidebar.removeWorkspaceWorkers': (v) => `Channel worker：${v?.count ?? 0}`, 'sidebar.noSessions': '暂无会话', 'sidebar.projectFallback': '项目', 'sidebar.sessionsOverview': '会话总览', diff --git a/packages/webui/src/daemon/session/DaemonSessionProvider.test.tsx b/packages/webui/src/daemon/session/DaemonSessionProvider.test.tsx index 32a83caea2c..b61d7f1ee55 100644 --- a/packages/webui/src/daemon/session/DaemonSessionProvider.test.tsx +++ b/packages/webui/src/daemon/session/DaemonSessionProvider.test.tsx @@ -33,7 +33,10 @@ import { type DaemonSessionNotice, type DaemonWorkspaceEventSignals, } from './DaemonSessionProvider.js'; -import { DaemonWorkspaceProvider } from '../workspace/DaemonWorkspaceProvider.js'; +import { + DaemonWorkspaceProvider, + useOptionalDaemonWorkspace, +} from '../workspace/DaemonWorkspaceProvider.js'; import { clearSidechannelMidTurnInjected, getSidechannelMidTurnInjected, @@ -879,6 +882,46 @@ describe('DaemonSessionProvider', () => { expect(sdkMocks.capabilities).toHaveBeenCalledTimes(1); }); + it('updates session connection capabilities after a workspace refresh', async () => { + sdkMocks.sessions.push(createMockSession()); + let connection: DaemonConnectionState | undefined; + let refreshCapabilities: (() => Promise) | undefined; + + function Harness() { + connection = useDaemonConnection(); + refreshCapabilities = useOptionalDaemonWorkspace()?.refreshCapabilities; + return null; + } + + container = document.createElement('div'); + document.body.appendChild(container); + root = createRoot(container); + act(() => { + root?.render( + + + + + , + ); + }); + await act(async () => { + await flushPromises(); + }); + sdkMocks.capabilities.mockResolvedValueOnce({ + workspaceCwd: '/mock-workspace', + features: ['workspace_runtime_removal'], + }); + + await act(async () => { + await refreshCapabilities?.(); + }); + + expect(connection?.capabilities?.features).toContain( + 'workspace_runtime_removal', + ); + }); + it('uses session context models over workspace provider defaults', async () => { sdkMocks.workspaceProviders.mockResolvedValue({ v: 1, diff --git a/packages/webui/src/daemon/session/DaemonSessionProvider.tsx b/packages/webui/src/daemon/session/DaemonSessionProvider.tsx index c15245748f1..feab6fba80d 100644 --- a/packages/webui/src/daemon/session/DaemonSessionProvider.tsx +++ b/packages/webui/src/daemon/session/DaemonSessionProvider.tsx @@ -301,6 +301,14 @@ export function DaemonSessionProvider(props: DaemonSessionProviderProps) { }); const connectionRef = useRef(connection); connectionRef.current = connection; + useEffect(() => { + if (!workspace?.capabilities) return; + setConnection((current) => + current.capabilities === workspace.capabilities + ? current + : { ...current, capabilities: workspace.capabilities }, + ); + }, [workspace?.capabilities]); const noticeIdRef = useRef(0); const [notices, setNotices] = useState([]); const addNotice = useCallback((input) => { diff --git a/packages/webui/src/daemon/workspace/actions.test.ts b/packages/webui/src/daemon/workspace/actions.test.ts new file mode 100644 index 00000000000..37424815606 --- /dev/null +++ b/packages/webui/src/daemon/workspace/actions.test.ts @@ -0,0 +1,112 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { createDaemonWorkspaceActions } from './actions.js'; + +describe('workspace actions', () => { + afterEach(() => { + vi.useRealTimers(); + }); + + it('applies the action timeout to workspace removal', async () => { + vi.useFakeTimers(); + const remove = vi.fn(() => new Promise(() => {})); + const actions = createDaemonWorkspaceActions({ + getClient: () => ({ workspaceById: () => ({ remove }) }) as never, + getWorkspaceCwd: () => '/ws', + baseUrl: '', + }); + + const result = actions + .removeWorkspace('secondary', { force: true, timeoutMs: 10 }) + .then( + () => undefined, + (error: unknown) => error, + ); + await vi.advanceTimersByTimeAsync(10); + + const error = await result; + expect(error).toBeInstanceOf(Error); + expect(error).toMatchObject({ + message: 'Remove workspace timed out after 10ms', + }); + expect(remove).toHaveBeenCalledWith({ force: true, timeoutMs: 10 }); + }); + + it('forwards successful workspace removal results', async () => { + const removal = { + removed: true as const, + workspaceId: 'secondary', + workspaceCwd: '/ws/secondary', + forced: false, + persistedRegistrationRemoved: true, + activity: { + sessions: 0, + activePrompts: 0, + pendingSessionStarts: 0, + acpConnections: 0, + memoryTasks: 0, + channelWorkers: 0, + }, + }; + const remove = vi.fn().mockResolvedValue(removal); + const workspaceById = vi.fn(() => ({ remove })); + const actions = createDaemonWorkspaceActions({ + getClient: () => ({ workspaceById }) as never, + getWorkspaceCwd: () => '/ws', + baseUrl: '', + }); + + await expect( + actions.removeWorkspace('secondary', { force: false }), + ).resolves.toEqual(removal); + expect(workspaceById).toHaveBeenCalledWith('secondary'); + expect(remove).toHaveBeenCalledWith({ force: false }); + }); + + it('rejects workspace removal without a connected client', async () => { + const actions = createDaemonWorkspaceActions({ + getClient: () => undefined, + getWorkspaceCwd: () => '/ws', + baseUrl: '', + }); + + await expect(actions.removeWorkspace('secondary')).rejects.toThrow( + 'Remove workspace failed: DaemonClient is not connected', + ); + }); + + it('preserves zero as the disabled timeout sentinel', async () => { + vi.useFakeTimers(); + const removal = { + removed: true as const, + workspaceId: 'secondary', + workspaceCwd: '/ws/secondary', + forced: false, + persistedRegistrationRemoved: false, + activity: { + sessions: 0, + activePrompts: 0, + pendingSessionStarts: 0, + acpConnections: 0, + memoryTasks: 0, + channelWorkers: 0, + }, + }; + const remove = vi.fn().mockResolvedValue(removal); + const actions = createDaemonWorkspaceActions({ + getClient: () => ({ workspaceById: () => ({ remove }) }) as never, + getWorkspaceCwd: () => '/ws', + baseUrl: '', + }); + + await expect( + actions.removeWorkspace('secondary', { timeoutMs: 0 }), + ).resolves.toEqual(removal); + expect(remove).toHaveBeenCalledWith({ timeoutMs: 0 }); + }); +}); diff --git a/packages/webui/src/daemon/workspace/actions.ts b/packages/webui/src/daemon/workspace/actions.ts index f0bfe0104cf..aa272ae44e8 100644 --- a/packages/webui/src/daemon/workspace/actions.ts +++ b/packages/webui/src/daemon/workspace/actions.ts @@ -715,6 +715,17 @@ export function createDaemonWorkspaceActions({ 'Add workspace timed out', ); }, + + async removeWorkspace(workspaceId, options) { + const client = requireClient(getClient, 'Remove workspace failed'); + const removal = client.workspaceById(workspaceId).remove(options); + if (options?.timeoutMs === 0) return removal; + return withActionTimeout( + removal, + 'Remove workspace timed out', + options?.timeoutMs, + ); + }, }; } diff --git a/packages/webui/src/daemon/workspace/types.ts b/packages/webui/src/daemon/workspace/types.ts index 0b6f5cab32f..dda851286be 100644 --- a/packages/webui/src/daemon/workspace/types.ts +++ b/packages/webui/src/daemon/workspace/types.ts @@ -44,6 +44,7 @@ import type { DaemonWorkspaceMcpToolsStatus, DaemonWorkspaceMcpResourcesStatus, DaemonWorkspaceMemoryStatus, + DaemonWorkspaceRemovalResult, DaemonWorkspacePreflightStatus, DaemonWorkspaceProvidersStatus, DaemonWorkspaceSkillsStatus, @@ -447,4 +448,8 @@ export interface DaemonWorkspaceActions { cwd: string, options?: { persist?: boolean }, ): Promise; + removeWorkspace( + workspaceId: string, + options?: { force?: boolean; timeoutMs?: number }, + ): Promise; }