diff --git a/packages/cli/src/__tests__/pi-tui-mcp-status.test.ts b/packages/cli/src/__tests__/pi-tui-mcp-status.test.ts index cdb4cb6dae..9d3a4c6afe 100644 --- a/packages/cli/src/__tests__/pi-tui-mcp-status.test.ts +++ b/packages/cli/src/__tests__/pi-tui-mcp-status.test.ts @@ -422,6 +422,39 @@ describe('MCP management overlay', () => { }); } + test('Esc aborts the active MCP action and ignores its late result', async () => { + let actionSignal: AbortSignal | undefined; + let resolveAction!: (result: TuiMcpActionResult) => void; + const actionResult = new Promise((resolve) => { + resolveAction = resolve; + }); + const mcp = surface(listSnapshot()); + mcp.execute = async (_action, options?: { signal?: AbortSignal }) => { + actionSignal = options?.signal; + return actionResult; + }; + const overlay = new McpManagementOverlay({ + locale: 'en', + surface: mcp, + viewportRows: () => 8, + onClose: () => undefined, + onChange: () => undefined, + }); + + overlay.handleInput('t'); + await new Promise((resolve) => setImmediate(resolve)); + assert.equal(actionSignal?.aborted, false); + + overlay.handleInput('\u001b'); + assert.equal(actionSignal?.aborted, true); + resolveAction({ status: 'tested', test: testResult(true), effect: 'published' }); + await new Promise((resolve) => setImmediate(resolve)); + + const text = overlay.render(100).map(stripAnsi).join('\n'); + assert.doesNotMatch(text, /MCP server responded successfully/u); + assert.match(text, /filesystem/u); + }); + const resultCopy = TUI_COPY_RESOURCES['mcp-status'].en.editor.results; const RESULT_CASES: readonly [TuiMcpActionResult, keyof typeof resultCopy][] = [ [{ status: 'conflict', reason: 'exists' }, 'exists'], @@ -430,9 +463,11 @@ describe('MCP management overlay', () => { [{ status: 'conflict', reason: 'stale_import' }, 'stale_import'], [{ status: 'conflict', reason: 'missing' }, 'missing'], [{ status: 'failed', reason: 'closed' }, 'closed'], + [{ status: 'failed', reason: 'cancelled' }, 'cancelled'], [{ status: 'failed', reason: 'invalid-config' }, 'invalid-config'], [{ status: 'failed', reason: 'credential-cleanup-failed' }, 'credential-cleanup-failed'], [{ status: 'failed', reason: 'persist-failed' }, 'persist-failed'], + [{ status: 'failed', reason: 'rollback-failed' }, 'rollback-failed'], [{ status: 'failed', reason: 'manager-failed' }, 'manager-failed'], [{ status: 'applied', effect: 'published' }, 'published'], [{ status: 'applied', effect: 'pending_host' }, 'pending_host'], diff --git a/packages/cli/src/__tests__/tui-mcp-control.test.ts b/packages/cli/src/__tests__/tui-mcp-control.test.ts index 32d893e483..249f2dbfd3 100644 --- a/packages/cli/src/__tests__/tui-mcp-control.test.ts +++ b/packages/cli/src/__tests__/tui-mcp-control.test.ts @@ -17,19 +17,21 @@ * under the License. */ -import { deferred } from '@maka/core/test-only/async-primitives'; +import { deferred, waitFor as pollFor } from '@maka/core/test-only/async-primitives'; import assert from 'node:assert/strict'; -import { mkdtemp, rm } from 'node:fs/promises'; +import { mkdtemp, readFile, rm } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import { join } from 'node:path'; +import { fileURLToPath } from 'node:url'; import test from 'node:test'; import type { McpConfigFile, McpServerStatus, McpToolSnapshot } from '@maka/core/mcp'; -import type { McpClientManager } from '@maka/mcp'; +import { createCredentialMcpOAuthStorage, McpClientManager } from '@maka/mcp'; import type { ClientCapabilityProvider, RuntimeHostConnectionAvailability, } from '@maka/runtime-host/client'; import { createMcpConfigStore } from '@maka/storage/mcp-config-store'; +import { createFileCredentialStore } from '@maka/storage/credential-store'; import { createTuiMcpController, type TuiMcpPublicationAvailability } from '../tui-mcp-control.js'; import { waitFor } from './tui-terminal-mock.js'; @@ -82,8 +84,14 @@ test('TUI MCP serializes remote provider credential changes through its publicat let removed = 0; let closed = 0; const connection = { - replaceClientCapabilities: async () => ({ registrationId: 'registration', revision: 1 }), - unregisterClientCapabilities: async () => ({ registrationId: 'registration', revision: 1 }), + replaceClientCapabilities: async () => ({ + registrationId: 'registration', + revision: 1, + }), + unregisterClientCapabilities: async () => ({ + registrationId: 'registration', + revision: 1, + }), subscribeConnectionAvailability: (next: (value: TuiMcpPublicationAvailability) => void) => { listener = next; next(availability); @@ -93,7 +101,11 @@ test('TUI MCP serializes remote provider credential changes through its publicat }, setCredential: async (credential: string) => { credentials.push(credential); - availability = { kind: 'connected', hostEpoch: 'host-1', connectionId: 'provider-1' }; + availability = { + kind: 'connected', + hostEpoch: 'host-1', + connectionId: 'provider-1', + }; listener?.(availability); }, removeCredential: async () => { @@ -143,6 +155,79 @@ test('TUI MCP serializes remote provider credential changes through its publicat assert.equal(closed, 1); }); +test('TUI MCP forwards cancellation into a publication credential change', async () => { + let credentialSignal: AbortSignal | undefined; + const connection = { + ...connectionHarness().connection, + setCredential: async (_credential: string, options?: { readonly signal?: AbortSignal }) => { + credentialSignal = options?.signal; + await new Promise((_resolve, reject) => { + options?.signal?.addEventListener( + 'abort', + () => reject(options.signal?.reason ?? new Error('cancelled')), + { once: true }, + ); + }); + }, + removeCredential: async () => undefined, + }; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection }, + { + configStore: configStoreHarness(async () => emptyConfig()), + manager: managerHarness(0, []).manager, + createProvider: () => undefined, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + 'remote MCP controller initialization before credential cancellation', + ); + const abort = new AbortController(); + const setting = controller.execute( + { kind: 'set_publication_credential', credential: 'provider-secret' }, + { signal: abort.signal }, + ); + await waitFor(() => credentialSignal !== undefined, 'credential target to receive its signal'); + abort.abort(new Error('cancel credential change')); + + assert.deepEqual(await setting, { status: 'failed', reason: 'cancelled' }); + assert.equal(credentialSignal?.aborted, true); + await controller.close(); +}); + +test('TUI MCP does not claim a credential cancellation before its target settles', async () => { + const credentialWrite = deferred(); + const connection = { + ...connectionHarness().connection, + setCredential: async () => credentialWrite.promise, + removeCredential: async () => undefined, + }; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection }, + { + configStore: configStoreHarness(async () => emptyConfig()), + manager: managerHarness(0, []).manager, + createProvider: () => undefined, + actionTimeoutMs: 50, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + 'remote MCP controller initialization before uncooperative credential cancellation', + ); + + assert.deepEqual( + await controller.execute({ + kind: 'set_publication_credential', + credential: 'provider-secret', + }), + { status: 'failed', reason: 'rollback-failed' }, + ); + credentialWrite.resolve(); + await controller.close(); +}); + test('TUI MCP publication coalesces a discovery change behind the in-flight revision', async () => { const manager = managerHarness(1, [connectedStatus('local', 1)]); const connection = connectionHarness(); @@ -183,6 +268,47 @@ test('TUI MCP publication coalesces a discovery change behind the in-flight revi await controller.close(); }); +test('TUI MCP unregisters a capability replacement that becomes stale after Host commit', async () => { + const manager = managerHarness(0, []); + const connection = connectionHarness(); + const publication = deferred(); + connection.replace = async () => { + connection.replacements.push('replace'); + await publication.promise; + }; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connection.connection }, + { + configStore: configStoreHarness(async () => emptyConfig()), + manager: manager.manager, + createProvider: (current) => + current.toolSnapshot().tools.length === 0 ? undefined : provider('provider'), + }, + ); + + await waitFor( + () => controller.snapshot().publication === 'not_published', + "initial empty TUI MCP publication to reach 'not_published'", + ); + manager.changeRevision(1, 1); + await waitFor( + () => connection.replacements.length === 1, + 'capability replacement before its snapshot becomes stale', + ); + manager.changeRevision(2, 0); + publication.resolve(); + + await waitFor( + () => connection.unregisters === 1, + 'stale capability replacement to be removed from the Host', + ); + await waitFor( + () => controller.snapshot().publication === 'not_published', + "TUI MCP publication to return to 'not_published'", + ); + await controller.close(); +}); + test('TUI MCP invalidates a lost generation and republishes on its replacement', async () => { const manager = managerHarness(1, [connectedStatus('local', 1)]); const connection = connectionHarness(); @@ -201,7 +327,11 @@ test('TUI MCP invalidates a lost generation and republishes on its replacement', ); connection.emit({ kind: 'unavailable' }); assert.equal(controller.snapshot().publication, 'host_unavailable'); - connection.emit({ kind: 'connected', hostEpoch: 'host-2', connectionId: 'connection-2' }); + connection.emit({ + kind: 'connected', + hostEpoch: 'host-2', + connectionId: 'connection-2', + }); await waitFor( () => connection.replacements.length === 2, 'second capability publication after generation replacement', @@ -342,7 +472,11 @@ test('TUI MCP retires endpoint credentials before persistence and aborts on clea const connection = connectionHarness(); const controller = createTuiMcpController( { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store.store, manager: manager.manager, createProvider: () => undefined }, + { + configStore: store.store, + manager: manager.manager, + createProvider: () => undefined, + }, ); await waitFor( () => controller.snapshot().initialization === 'ready', @@ -359,7 +493,10 @@ test('TUI MCP retires endpoint credentials before persistence and aborts on clea config: { url: 'https://new.example/mcp' }, }); - assert.deepEqual(result, { status: 'failed', reason: 'credential-cleanup-failed' }); + assert.deepEqual(result, { + status: 'failed', + reason: 'credential-cleanup-failed', + }); assert.deepEqual(order, ['transform', 'forget:docs']); const stored = (await store.store.get()).mcpServers.docs; assert.ok(stored && 'url' in stored); @@ -377,7 +514,11 @@ test('TUI MCP rejects an invalid endpoint before retiring the previous credentia const connection = connectionHarness(); const controller = createTuiMcpController( { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store.store, manager: manager.manager, createProvider: () => undefined }, + { + configStore: store.store, + manager: manager.manager, + createProvider: () => undefined, + }, ); await waitFor( () => controller.snapshot().initialization === 'ready', @@ -410,7 +551,11 @@ test('TUI MCP edit rejects a stale revision without touching credentials or disk const connection = connectionHarness(); const controller = createTuiMcpController( { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store.store, manager: manager.manager, createProvider: () => undefined }, + { + configStore: store.store, + manager: manager.manager, + createProvider: () => undefined, + }, ); await waitFor( () => controller.snapshot().initialization === 'ready', @@ -418,7 +563,10 @@ test('TUI MCP edit rejects a stale revision without touching credentials or disk ); const edit = controller.configForEdit('docs'); assert.ok(edit); - store.replace({ version: 3, mcpServers: { docs: { url: 'https://other.example/mcp' } } }); + store.replace({ + version: 3, + mcpServers: { docs: { url: 'https://other.example/mcp' } }, + }); order.length = 0; const result = await controller.execute({ @@ -443,7 +591,11 @@ test('TUI MCP import preserves unrelated external edits and rejects changed prev const connection = connectionHarness(); const controller = createTuiMcpController( { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store.store, manager: manager.manager, createProvider: () => undefined }, + { + configStore: store.store, + manager: manager.manager, + createProvider: () => undefined, + }, ); await waitFor( () => controller.snapshot().initialization === 'ready', @@ -458,7 +610,10 @@ test('TUI MCP import preserves unrelated external edits and rejects changed prev }); assert.deepEqual( - await controller.execute({ kind: 'commit_import', previewId: preview.preview.previewId }), + await controller.execute({ + kind: 'commit_import', + previewId: preview.preview.previewId, + }), { status: 'applied', effect: 'published', @@ -479,7 +634,10 @@ test('TUI MCP import preserves unrelated external edits and rejects changed prev }, }); assert.deepEqual( - await controller.execute({ kind: 'commit_import', previewId: stale.preview.previewId }), + await controller.execute({ + kind: 'commit_import', + previewId: stale.preview.previewId, + }), { status: 'conflict', reason: 'stale_import', @@ -495,7 +653,11 @@ test('TUI MCP keeps a durable mutation visible when manager synchronization fail const connection = connectionHarness(); const controller = createTuiMcpController( { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store.store, manager: manager.manager, createProvider: () => undefined }, + { + configStore: store.store, + manager: manager.manager, + createProvider: () => undefined, + }, ); await waitFor( () => controller.snapshot().initialization === 'ready', @@ -532,7 +694,11 @@ test('TUI MCP reports a committed action as pending while the Host is unavailabl connection.emit({ kind: 'unavailable' }); const controller = createTuiMcpController( { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store.store, manager: manager.manager, createProvider: () => undefined }, + { + configStore: store.store, + manager: manager.manager, + createProvider: () => undefined, + }, ); await waitFor( () => controller.snapshot().initialization === 'ready', @@ -540,7 +706,11 @@ test('TUI MCP reports a committed action as pending while the Host is unavailabl ); assert.deepEqual( - await controller.execute({ kind: 'add', serverId: 'local', config: { command: 'server' } }), + await controller.execute({ + kind: 'add', + serverId: 'local', + config: { command: 'server' }, + }), { status: 'applied', effect: 'pending_host' }, ); await controller.close(); @@ -570,7 +740,11 @@ test('TUI MCP close fences an admitted mutation before persistence', async () => const connection = connectionHarness(); const controller = createTuiMcpController( { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store, manager: manager.manager, createProvider: () => undefined }, + { + configStore: store, + manager: manager.manager, + createProvider: () => undefined, + }, ); await waitFor( () => controller.snapshot().initialization === 'ready', @@ -591,129 +765,1136 @@ test('TUI MCP close fences an admitted mutation before persistence', async () => assert.equal(writes, 0); }); -test('TUI MCP waits for manager synchronization before publishing an action snapshot', async () => { - const store = mutableConfigStore(emptyConfig(), []); - const actionSync = deferred(); +test('TUI MCP aborts a queued mutation before it reaches persistence', async () => { + const firstSync = deferred(); let syncCount = 0; - let listener: (() => void) | undefined; - let revision = 0; - const manager = { - sync: async () => { - syncCount += 1; - revision += 1; - listener?.(); - if (syncCount === 2) await actionSync.promise; - }, - statuses: () => [], - toolSnapshot: () => ({ revision, tools: [{}] }) as unknown as McpToolSnapshot, - callTool: async () => ({ content: [] }), - test: async () => ({ ok: true, status: connectedStatus('local', 1), latencyMs: 1 }), - reconnect: async () => connectedStatus('local', 1), - forgetServerCredentials: async () => undefined, - onChange: (next: () => void) => { - listener = next; - return () => { - listener = undefined; - }; + let transforms = 0; + const manager = managementManager([]); + manager.manager.sync = async () => { + syncCount += 1; + if (syncCount === 2) await firstSync.promise; + }; + const store = { + get: async () => emptyConfig(), + transform: async ( + apply: (current: McpConfigFile) => McpConfigFile | Promise, + ) => { + transforms += 1; + return apply(emptyConfig()); }, - close: async () => undefined, - } as unknown as ReturnType['manager']; - const connection = connectionHarness(); + }; const controller = createTuiMcpController( - { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store.store, manager, createProvider: () => provider('provider') }, + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: store, + manager: manager.manager, + createProvider: () => undefined, + }, ); await waitFor( - () => controller.snapshot().publication === 'published', - "initial TUI MCP publication to reach 'published' before action sync", + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before queued cancellation", ); - connection.replacements.length = 0; - const executing = controller.execute({ + const first = controller.execute({ kind: 'add', - serverId: 'local', - config: { command: 'server' }, + serverId: 'first', + config: { command: 'first' }, }); - await waitFor( - () => syncCount === 2, - 'manager sync to reach count 2 before publishing action snapshot', + await waitFor(() => syncCount === 2, 'first action to own the manager sync'); + const abort = new AbortController(); + const second = controller.execute( + { kind: 'add', serverId: 'second', config: { command: 'second' } }, + { signal: abort.signal }, ); - assert.equal(connection.replacements.length, 0); - assert.equal(controller.snapshot().configuration, 'synchronizing'); - actionSync.resolve(); + abort.abort(new Error('cancel queued action')); + firstSync.resolve(); - assert.deepEqual(await executing, { status: 'applied', effect: 'published' }); - assert.equal(connection.replacements.length, 1); + assert.deepEqual(await first, { status: 'applied', effect: 'published' }); + assert.deepEqual(await second, { status: 'failed', reason: 'cancelled' }); + assert.equal(transforms, 1); await controller.close(); }); -test('TUI MCP rebases an action over an unrelated concurrent config edit', async () => { - let config: McpConfigFile = { - version: 3, - mcpServers: { existing: { command: 'before' } }, +test('TUI MCP action deadline starts when a queued action begins running', async () => { + const firstSync = deferred(); + let syncCount = 0; + const manager = managementManager([]); + manager.manager.sync = async () => { + syncCount += 1; + if (syncCount === 2) await firstSync.promise; }; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: mutableConfigStore(emptyConfig(), []).store, + manager: manager.manager, + createProvider: () => undefined, + actionTimeoutMs: 50, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before queued action deadline", + ); + const first = controller.execute({ + kind: 'add', + serverId: 'first', + config: { command: 'first' }, + }); + await waitFor(() => syncCount === 2, 'first action to own the manager sync'); + const second = controller.execute({ + kind: 'add', + serverId: 'second', + config: { command: 'second' }, + }); + await new Promise((resolve) => setTimeout(resolve, 60)); + firstSync.resolve(); + + assert.deepEqual(await first, { status: 'failed', reason: 'cancelled' }); + assert.deepEqual(await second, { status: 'applied', effect: 'published' }); + await controller.close(); +}); + +test('TUI MCP action deadline bounds a config transaction waiting before its callback', async () => { + let transforms = 0; + const blocked = deferred(); const store = { - get: async () => structuredClone(config), - transform: async ( - apply: (current: McpConfigFile) => McpConfigFile | Promise, - ) => { - config = await apply({ - version: 3, - mcpServers: { existing: { command: 'concurrent' } }, - }); - return structuredClone(config); + get: async () => emptyConfig(), + transform: async () => { + transforms += 1; + return blocked.promise; }, }; - const manager = managementManager([]); - const connection = connectionHarness(); const controller = createTuiMcpController( - { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store, manager: manager.manager, createProvider: () => undefined }, + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: store, + manager: managementManager([]).manager, + createProvider: () => undefined, + actionTimeoutMs: 50, + }, ); await waitFor( () => controller.snapshot().initialization === 'ready', - "TUI MCP initialization to reach 'ready' before rebased concurrent edit", + "TUI MCP initialization to reach 'ready' before blocked config transaction", ); + const startedAt = Date.now(); assert.deepEqual( - await controller.execute({ kind: 'add', serverId: 'local', config: { command: 'server' } }), - { status: 'applied', effect: 'published' }, + await controller.execute({ + kind: 'add', + serverId: 'blocked', + config: { command: 'server' }, + }), + { status: 'failed', reason: 'rollback-failed' }, ); - const existing = config.mcpServers.existing; - assert.ok(existing && 'command' in existing); - assert.equal(existing.command, 'concurrent'); - assert.ok(config.mcpServers.local); + assert.ok(Date.now() - startedAt < 500); + assert.equal(transforms, 1); + assert.equal(controller.snapshot().configuration, 'out_of_sync'); + blocked.resolve(emptyConfig()); await controller.close(); }); -test('independent TUI controllers preserve concurrent additions in one workspace', async (t) => { - const root = await mkdtemp(join(tmpdir(), 'maka-tui-mcp-concurrent-')); - t.after(() => rm(root, { recursive: true, force: true })); - await createMcpConfigStore(root).transform((current) => current); - const left = createTuiMcpController( - { workspaceRoot: root, connection: connectionHarness().connection }, - { - configStore: createMcpConfigStore(root), - manager: managementManager([]).manager, - createProvider: () => undefined, +test('TUI MCP compensates a config transaction that commits after its cleanup deadline', async () => { + let config = emptyConfig(); + let transforms = 0; + const lateWrite = deferred(); + const store = { + get: async () => structuredClone(config), + transform: async ( + apply: (current: McpConfigFile) => McpConfigFile | Promise, + ) => { + transforms += 1; + const next = structuredClone(await apply(structuredClone(config))); + if (transforms === 1) await lateWrite.promise; + config = next; + return structuredClone(config); }, - ); - const right = createTuiMcpController( - { workspaceRoot: root, connection: connectionHarness().connection }, + }; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, { - configStore: createMcpConfigStore(root), + configStore: store, manager: managementManager([]).manager, createProvider: () => undefined, + actionTimeoutMs: 50, }, ); await waitFor( - () => left.snapshot().initialization === 'ready' && right.snapshot().initialization === 'ready', - "both TUI controllers to reach 'ready' before concurrent additions", + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before late config commit", + ); + + assert.deepEqual( + await controller.execute({ + kind: 'add', + serverId: 'late', + config: { command: 'server' }, + }), + { status: 'failed', reason: 'rollback-failed' }, + ); + lateWrite.resolve(); + await waitFor( + () => transforms === 2 && config.mcpServers.late === undefined, + 'late config commit to be conditionally compensated', + ); + assert.equal(controller.snapshot().configuration, 'ready'); + await controller.close(); +}); + +test('TUI MCP cancellation rolls back a committed add and reconciles the manager', async () => { + const store = mutableConfigStore(emptyConfig(), []); + const actionSync = deferred(); + let syncCount = 0; + let disconnects = 0; + let disconnectSignal: AbortSignal | undefined; + const syncSnapshots: McpConfigFile[] = []; + const manager = managementManager([], { + sync: async (config, options) => { + syncCount += 1; + syncSnapshots.push(structuredClone(config)); + if (syncCount !== 2) return; + await new Promise((resolve, reject) => { + const onAbort = () => reject(options?.signal?.reason ?? new Error('cancelled')); + options?.signal?.addEventListener('abort', onAbort, { once: true }); + void actionSync.promise.then(resolve, reject); + }); + }, + disconnect: async (_serverId, _remove, options) => { + disconnects += 1; + disconnectSignal = options?.signal; + }, + }).manager; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { configStore: store.store, manager, createProvider: () => undefined }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before committed cancellation", + ); + const abort = new AbortController(); + const adding = controller.execute( + { kind: 'add', serverId: 'late', config: { command: 'server' } }, + { signal: abort.signal }, + ); + await waitFor(() => syncCount === 2, 'cancelled action to reach manager synchronization'); + + abort.abort(new Error('cancel committed add')); + assert.deepEqual(await adding, { status: 'failed', reason: 'cancelled' }); + assert.equal((await store.store.get()).mcpServers.late, undefined); + assert.equal(syncSnapshots.length, 3); + assert.equal(disconnects, 1); + const committed = syncSnapshots[1]?.mcpServers.late; + assert.ok(committed && 'command' in committed); + assert.equal(committed.command, 'server'); + assert.equal(syncSnapshots[2]?.mcpServers.late, undefined); + assert.ok(disconnectSignal); + assert.equal(disconnectSignal.aborted, false); + assert.equal(controller.snapshot().servers.length, 0); + assert.equal(controller.snapshot().configuration, 'ready'); + actionSync.resolve(); + await controller.close(); +}); + +test('TUI MCP forwards its action deadline into credential retirement before persistence', async () => { + const initial = { + version: 3, + mcpServers: { + docs: { url: 'https://old.example/mcp', oauth: { clientId: 'client' } }, + }, + } satisfies McpConfigFile; + const store = mutableConfigStore(initial, []); + let cleanupSignal: AbortSignal | undefined; + const manager = managementManager([], { + forgetServerCredentials: async (_serverId, _config, options) => { + cleanupSignal = options?.signal; + if (!cleanupSignal) throw new Error('missing credential cleanup signal'); + await new Promise((_resolve, reject) => { + cleanupSignal?.addEventListener( + 'abort', + () => reject(cleanupSignal?.reason ?? new Error('deadline')), + { once: true }, + ); + }); + }, + }).manager; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: store.store, + manager, + createProvider: () => undefined, + actionTimeoutMs: 100, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before credential cleanup deadline", + ); + const edit = controller.configForEdit('docs'); + assert.ok(edit); + + assert.deepEqual( + await controller.execute({ + kind: 'edit', + serverId: 'docs', + expectedRevision: edit.revision, + config: { url: 'https://new.example/mcp', oauth: { clientId: 'client' } }, + }), + { status: 'failed', reason: 'cancelled' }, + ); + assert.ok(cleanupSignal); + assert.equal(cleanupSignal.aborted, true); + const current = await store.store.get(); + const docs = current.mcpServers.docs; + assert.ok(docs && 'url' in docs); + assert.equal(docs.url, 'https://old.example/mcp'); + await controller.close(); +}); + +test('TUI MCP commits an endpoint edit once credential retirement has started', async () => { + const initial = { + version: 3, + mcpServers: { + docs: { url: 'https://old.example/mcp', oauth: { clientId: 'client' } }, + }, + } satisfies McpConfigFile; + const store = mutableConfigStore(initial, []); + const retirement = deferred(); + let cleanupSignal: AbortSignal | undefined; + const manager = managementManager([], { + forgetServerCredentials: async (_serverId, _config, options) => { + cleanupSignal = options?.signal; + options?.onCommitStarted?.(); + await retirement.promise; + }, + }).manager; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: store.store, + manager, + createProvider: () => undefined, + actionTimeoutMs: 1_000, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + 'TUI MCP initialization before credential retirement settlement test', + ); + const edit = controller.configForEdit('docs'); + assert.ok(edit); + const abort = new AbortController(); + const editing = controller.execute( + { + kind: 'edit', + serverId: 'docs', + expectedRevision: edit.revision, + config: { url: 'https://new.example/mcp', oauth: { clientId: 'client' } }, + }, + { signal: abort.signal }, + ); + await waitFor(() => cleanupSignal !== undefined, 'credential retirement to start'); + abort.abort(new Error('cancel credential retirement')); + assert.equal(cleanupSignal?.aborted, true); + + assert.equal(await settlesWithin(editing, 10), false); + retirement.resolve(); + assert.deepEqual(await editing, { status: 'applied', effect: 'sync_failed' }); + const current = await store.store.get(); + const docs = current.mcpServers.docs; + assert.ok(docs && 'url' in docs); + assert.equal(docs.url, 'https://new.example/mcp'); + await controller.close(); +}); + +test('TUI MCP never late-rolls back an endpoint after credential retirement starts', async () => { + const initial = { + version: 3, + mcpServers: { + docs: { url: 'https://old.example/mcp', oauth: { clientId: 'client' } }, + }, + } satisfies McpConfigFile; + const store = mutableConfigStore(initial, []); + const retirement = deferred(); + let commitStarted = false; + const manager = managementManager([], { + forgetServerCredentials: async (_serverId, _config, options) => { + commitStarted = true; + options?.onCommitStarted?.(); + // Model a storage backend whose tombstone write is already in flight, + // but whose promise settles after both the action deadline and cleanup + // reserve have elapsed. + await retirement.promise; + }, + }).manager; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: store.store, + manager, + createProvider: () => undefined, + actionTimeoutMs: 50, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + 'TUI MCP initialization before late credential retirement', + ); + const edit = controller.configForEdit('docs'); + assert.ok(edit); + + const editing = controller.execute({ + kind: 'edit', + serverId: 'docs', + expectedRevision: edit.revision, + config: { url: 'https://new.example/mcp', oauth: { clientId: 'client' } }, + }); + + await waitFor(() => commitStarted, 'credential retirement commit to start'); + assert.equal(await settlesWithin(editing, 100), false); + const beforeSettlement = await store.store.get(); + const oldDocs = beforeSettlement.mcpServers.docs; + assert.ok(oldDocs && 'url' in oldDocs); + assert.equal(oldDocs.url, 'https://old.example/mcp'); + + retirement.resolve(); + assert.deepEqual(await editing, { status: 'applied', effect: 'sync_failed' }); + const current = await store.store.get(); + const docs = current.mcpServers.docs; + assert.ok(docs && 'url' in docs); + assert.equal(docs.url, 'https://new.example/mcp'); + await controller.close(); +}); + +test('TUI MCP keeps the new endpoint when synchronization fails after credential retirement', async () => { + const initial = { + version: 3, + mcpServers: { + docs: { url: 'https://old.example/mcp', oauth: { clientId: 'client' } }, + }, + } satisfies McpConfigFile; + const store = mutableConfigStore(initial, []); + const syncStarted = deferred(); + const releaseSync = deferred(); + const manager = managementManager([], { + forgetServerCredentials: async (_serverId, _config, options) => { + options?.onCommitStarted?.(); + }, + sync: async (config) => { + const docs = config.mcpServers.docs; + if (!docs || !('url' in docs) || docs.url !== 'https://new.example/mcp') return; + syncStarted.resolve(); + await releaseSync.promise; + throw new Error('new endpoint synchronization failed'); + }, + }).manager; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { configStore: store.store, manager, createProvider: () => undefined }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + 'TUI MCP initialization before post-retirement sync failure', + ); + const edit = controller.configForEdit('docs'); + assert.ok(edit); + const abort = new AbortController(); + const editing = controller.execute( + { + kind: 'edit', + serverId: 'docs', + expectedRevision: edit.revision, + config: { url: 'https://new.example/mcp', oauth: { clientId: 'client' } }, + }, + { signal: abort.signal }, + ); + + await syncStarted.promise; + abort.abort(new Error('cancel after irreversible credential retirement')); + releaseSync.resolve(); + + assert.deepEqual(await editing, { status: 'applied', effect: 'sync_failed' }); + const current = await store.store.get(); + const docs = current.mcpServers.docs; + assert.ok(docs && 'url' in docs); + assert.equal(docs.url, 'https://new.example/mcp'); + assert.equal(controller.snapshot().configuration, 'out_of_sync'); + await controller.close(); +}); + +test('TUI MCP cancels post-retirement synchronization without rolling back the committed endpoint', async () => { + const initial = { + version: 3, + mcpServers: { + docs: { url: 'https://old.example/mcp', oauth: { clientId: 'client' } }, + }, + } satisfies McpConfigFile; + const store = mutableConfigStore(initial, []); + const syncStarted = deferred(); + let syncSignal: AbortSignal | undefined; + const manager = managementManager([], { + forgetServerCredentials: async (_serverId, _config, options) => { + options?.onCommitStarted?.(); + }, + sync: async (config, options) => { + const docs = config.mcpServers.docs; + if (!docs || !('url' in docs) || docs.url !== 'https://new.example/mcp') return; + syncSignal = options?.signal; + syncStarted.resolve(); + await new Promise((_resolve, reject) => { + options?.signal?.addEventListener( + 'abort', + () => reject(options.signal?.reason ?? new Error('cancelled')), + { once: true }, + ); + }); + }, + }).manager; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: store.store, + manager, + createProvider: () => undefined, + actionTimeoutMs: 1_000, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + 'TUI MCP initialization before post-retirement synchronization cancellation', + ); + const edit = controller.configForEdit('docs'); + assert.ok(edit); + const abort = new AbortController(); + const editing = controller.execute( + { + kind: 'edit', + serverId: 'docs', + expectedRevision: edit.revision, + config: { url: 'https://new.example/mcp', oauth: { clientId: 'client' } }, + }, + { signal: abort.signal }, + ); + + await syncStarted.promise; + abort.abort(new Error('cancel post-retirement synchronization')); + + assert.equal(await settlesWithin(editing, 100), true); + assert.deepEqual(await editing, { status: 'applied', effect: 'sync_failed' }); + assert.equal(syncSignal?.aborted, true); + assert.equal(controller.snapshot().configuration, 'out_of_sync'); + assert.equal(controller.snapshot().servers[0]?.synchronized, false); + const current = await store.store.get(); + const docs = current.mcpServers.docs; + assert.ok(docs && 'url' in docs); + assert.equal(docs.url, 'https://new.example/mcp'); + await controller.close(); +}); + +test('TUI MCP bounds a real post-retirement stdio connection and keeps the committed endpoint', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'maka-tui-mcp-post-retirement-cancel-')); + const eventLog = join(root, 'events.jsonl'); + const store = createMcpConfigStore(root); + await store.upsert('docs', { + url: 'https://old.example/mcp', + enabled: false, + oauth: { clientId: 'client' }, + }); + const fixturePath = fileURLToPath( + new URL(import.meta.resolve('@maka/mcp/test-only/stdio-server')), + ); + const controller = createTuiMcpController( + { workspaceRoot: root, connection: connectionHarness().connection }, + { + manager: new McpClientManager({ + oauthStorage: createCredentialMcpOAuthStorage(createFileCredentialStore(root)), + timeouts: { stdioConnectMs: 5_000, listToolsMs: 5_000, callToolMs: 5_000 }, + }), + createProvider: () => undefined, + // Leave enough room for a loaded CI runner to spawn the real child and + // reach tools/list; this test triggers cancellation explicitly below. + actionTimeoutMs: 5_000, + }, + ); + let childPid: number | undefined; + t.after(async () => { + if (childPid && processExists(childPid)) process.kill(childPid, 'SIGKILL'); + await controller.close().catch(() => undefined); + await rm(root, { recursive: true, force: true }); + }); + await waitFor( + () => controller.snapshot().initialization === 'ready', + 'TUI MCP initialization before real post-retirement cancellation', + ); + const edit = controller.configForEdit('docs'); + assert.ok(edit); + const abort = new AbortController(); + const editing = controller.execute( + { + kind: 'edit', + serverId: 'docs', + expectedRevision: edit.revision, + config: { + command: process.execPath, + args: [fixturePath, '--slow-tool-list'], + env: { MAKA_MCP_STDIO_EVENT_LOG: eventLog }, + protocol: 'legacy', + }, + }, + { signal: abort.signal }, + ); + await pollFor( + async () => { + const events = await readFixtureEvents(eventLog); + const start = events.find((event) => event.event === 'start'); + if (start) childPid = start.pid; + return events.some((event) => event.event === 'tools-list'); + }, + { timeoutMs: 4_000, pollMs: 5 }, + ); + + abort.abort(new Error('cancel real post-retirement connection')); + + assert.equal(await settlesWithin(editing, 750), true); + assert.deepEqual(await editing, { status: 'applied', effect: 'sync_failed' }); + const current = await store.get(); + const docs = current.mcpServers.docs; + assert.ok(docs && 'command' in docs); + assert.equal(docs.command, process.execPath); + assert.equal(controller.snapshot().configuration, 'out_of_sync'); + assert.equal(controller.snapshot().servers[0]?.synchronized, false); + assert.ok(childPid); + const cancelledPid = childPid; + assert.equal(processExists(cancelledPid), true); + await pollFor(() => !processExists(cancelledPid), { timeoutMs: 3_000, pollMs: 5 }); +}); + +test('TUI MCP marks configuration out of sync when persistence fails after credential retirement', async () => { + const initial = { + version: 3, + mcpServers: { + docs: { url: 'https://old.example/mcp', oauth: { clientId: 'client' } }, + }, + } satisfies McpConfigFile; + let credentialRetired = false; + const store = { + get: async () => structuredClone(initial), + transform: async ( + apply: (current: McpConfigFile) => McpConfigFile | Promise, + ) => { + await apply(structuredClone(initial)); + const error = new Error('config write failed'); + Object.assign(error, { code: 'ENOSPC' }); + throw error; + }, + }; + const manager = managementManager([], { + forgetServerCredentials: async (_serverId, _config, options) => { + options?.onCommitStarted?.(); + credentialRetired = true; + }, + }).manager; + (manager.statuses() as McpServerStatus[]).push({ + serverId: 'docs', + state: 'connected', + transport: 'streamable-http', + negotiatedProtocol: { era: 'modern', revision: '2025-11-25' }, + toolCount: 0, + tools: [], + updatedAt: 1, + }); + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { configStore: store, manager, createProvider: () => undefined }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + 'TUI MCP initialization before post-retirement persistence failure', + ); + assert.equal(controller.snapshot().configuration, 'ready'); + assert.equal(controller.snapshot().servers[0]?.synchronized, true); + const edit = controller.configForEdit('docs'); + assert.ok(edit); + + assert.deepEqual( + await controller.execute({ + kind: 'edit', + serverId: 'docs', + expectedRevision: edit.revision, + config: { url: 'https://new.example/mcp', oauth: { clientId: 'client' } }, + }), + { status: 'failed', reason: 'persist-failed' }, + ); + assert.equal(credentialRetired, true); + const persisted = await store.get(); + const docs = persisted.mcpServers.docs; + assert.ok(docs && 'url' in docs); + assert.equal(docs.url, 'https://old.example/mcp'); + assert.equal(controller.snapshot().configuration, 'out_of_sync'); + assert.equal(controller.snapshot().servers[0]?.synchronized, false); + await controller.close(); +}); + +test('TUI MCP bounds an uncooperative credential retirement after cancellation', async () => { + const initial = { + version: 3, + mcpServers: { + docs: { url: 'https://old.example/mcp', oauth: { clientId: 'client' } }, + }, + } satisfies McpConfigFile; + const store = mutableConfigStore(initial, []); + let cleanupSignal: AbortSignal | undefined; + const manager = managementManager([], { + forgetServerCredentials: async (_serverId, _config, options) => { + cleanupSignal = options?.signal; + await new Promise(() => undefined); + }, + }).manager; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: store.store, + manager, + createProvider: () => undefined, + actionTimeoutMs: 50, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + 'TUI MCP initialization before uncooperative credential retirement', + ); + const edit = controller.configForEdit('docs'); + assert.ok(edit); + + const editing = controller.execute({ + kind: 'edit', + serverId: 'docs', + expectedRevision: edit.revision, + config: { url: 'https://new.example/mcp', oauth: { clientId: 'client' } }, + }); + + assert.equal(await settlesWithin(editing, 500), true); + assert.deepEqual(await editing, { + status: 'failed', + reason: 'rollback-failed', + }); + assert.equal(cleanupSignal?.aborted, true); + const current = await store.store.get(); + const docs = current.mcpServers.docs; + assert.ok(docs && 'url' in docs); + assert.equal(docs.url, 'https://old.example/mcp'); + await controller.close(); +}); + +test('TUI MCP reports a failed cancellation rollback instead of claiming cancellation completed', async () => { + let config = emptyConfig(); + let transforms = 0; + const store = { + get: async () => structuredClone(config), + transform: async ( + apply: (current: McpConfigFile) => McpConfigFile | Promise, + ) => { + transforms += 1; + if (transforms === 2) throw new Error('rollback write failed'); + config = structuredClone(await apply(structuredClone(config))); + return structuredClone(config); + }, + }; + const actionSync = deferred(); + let syncCount = 0; + const manager = managementManager([], { + sync: async (_config, options) => { + syncCount += 1; + if (syncCount !== 2) return; + await new Promise((resolve, reject) => { + options?.signal?.addEventListener( + 'abort', + () => reject(options.signal?.reason ?? new Error('cancelled')), + { once: true }, + ); + void actionSync.promise.then(resolve, reject); + }); + }, + }).manager; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { configStore: store, manager, createProvider: () => undefined }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before rollback failure", + ); + const abort = new AbortController(); + const adding = controller.execute( + { kind: 'add', serverId: 'late', config: { command: 'server' } }, + { signal: abort.signal }, + ); + await waitFor(() => syncCount === 2, 'cancelled action to reach manager synchronization'); + + abort.abort(new Error('cancel committed add')); + assert.deepEqual(await adding, { + status: 'failed', + reason: 'rollback-failed', + }); + assert.ok(config.mcpServers.late); + assert.equal(controller.snapshot().configuration, 'out_of_sync'); + actionSync.resolve(); + await controller.close(); +}); + +test('TUI MCP rollback preserves a newer concurrent edit to the same server', async () => { + const store = mutableConfigStore(emptyConfig(), []); + const actionSync = deferred(); + let syncCount = 0; + const manager = managementManager([], { + sync: async (_config, options) => { + syncCount += 1; + if (syncCount !== 2) return; + await new Promise((resolve, reject) => { + options?.signal?.addEventListener( + 'abort', + () => reject(options.signal?.reason ?? new Error('cancelled')), + { once: true }, + ); + void actionSync.promise.then(resolve, reject); + }); + }, + }).manager; + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { configStore: store.store, manager, createProvider: () => undefined }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before concurrent rollback", + ); + const abort = new AbortController(); + const adding = controller.execute( + { kind: 'add', serverId: 'docs', config: { command: 'old-command' } }, + { signal: abort.signal }, + ); + await waitFor(() => syncCount === 2, 'cancelled action to reach manager synchronization'); + store.replace({ + version: 3, + mcpServers: { docs: { command: 'newer-command' } }, + }); + + abort.abort(new Error('cancel superseded add')); + assert.deepEqual(await adding, { status: 'failed', reason: 'cancelled' }); + const docs = (await store.store.get()).mcpServers.docs; + assert.ok(docs && 'command' in docs); + assert.equal(docs.command, 'newer-command'); + actionSync.resolve(); + await controller.close(); +}); + +test('TUI MCP forwards action cancellation to manager test without accepting its late result', async () => { + const manager = cancellableTestManager({ lateSuccessfulTest: true }); + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: configStoreHarness(async () => docsConfig()), + manager: manager.manager, + createProvider: () => undefined, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before test cancellation", + ); + const abort = new AbortController(); + const testing = controller.execute({ kind: 'test', serverId: 'docs' }, { signal: abort.signal }); + await waitFor(() => manager.testSignal !== undefined, 'manager test to receive its signal'); + + abort.abort(new Error('cancel test')); + assert.deepEqual(await testing, { status: 'failed', reason: 'cancelled' }); + assert.equal(manager.testSignal?.aborted, true); + assert.equal(controller.snapshot().servers[0]?.state, 'disconnected'); + await controller.close(); +}); + +test('TUI MCP cancellation removes the previously published tools after test cleanup', async () => { + const manager = cancellableTestManager({ + status: connectedStatus('docs', 1), + revision: 1, + }); + const connection = connectionHarness(); + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connection.connection }, + { + configStore: configStoreHarness(async () => docsConfig()), + manager: manager.manager, + createProvider: (current) => + current.toolSnapshot().tools.length === 0 ? undefined : provider('provider'), + }, + ); + await waitFor( + () => controller.snapshot().publication === 'published' && connection.replacements.length === 1, + 'initial MCP tools to be published before test cancellation', + ); + const abort = new AbortController(); + const testing = controller.execute({ kind: 'test', serverId: 'docs' }, { signal: abort.signal }); + await new Promise((resolve) => setImmediate(resolve)); + + abort.abort(new Error('cancel test')); + assert.deepEqual(await testing, { status: 'failed', reason: 'cancelled' }); + assert.equal(connection.unregisters, 1); + assert.equal(controller.snapshot().publication, 'not_published'); + await controller.close(); + assert.equal(connection.unregisters, 1); +}); + +test('TUI MCP returns a typed failure when cancelled publication cleanup times out', async () => { + const manager = cancellableTestManager({ + status: connectedStatus('docs', 1), + revision: 1, + notifyOnDisconnect: false, + }); + const connection = connectionHarness(); + const allowUnregister = deferred(); + connection.connection.unregisterClientCapabilities = async () => + allowUnregister.promise.then(() => ({ + registrationId: 'registration', + revision: 1, + })); + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connection.connection }, + { + configStore: configStoreHarness(async () => docsConfig()), + manager: manager.manager, + createProvider: (current) => + current.toolSnapshot().tools.length === 0 ? undefined : provider('provider'), + actionTimeoutMs: 50, + }, + ); + await waitFor( + () => controller.snapshot().publication === 'published', + 'initial publication before bounded cancellation cleanup', + ); + + const result = await controller.execute({ kind: 'test', serverId: 'docs' }); + + assert.deepEqual(result, { status: 'failed', reason: 'manager-failed' }); + assert.equal(controller.snapshot().servers[0]?.state, 'disconnected'); + allowUnregister.resolve(); + await controller.close(); +}); + +test('TUI MCP action deadline aborts a hung test and restores disconnected state', async () => { + const manager = cancellableTestManager(); + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: configStoreHarness(async () => docsConfig()), + manager: manager.manager, + createProvider: () => undefined, + actionTimeoutMs: 10, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before action deadline", + ); + + assert.deepEqual(await controller.execute({ kind: 'test', serverId: 'docs' }), { + status: 'failed', + reason: 'cancelled', + }); + assert.equal(manager.disconnects, 1); + assert.equal(controller.snapshot().servers[0]?.state, 'disconnected'); + await controller.close(); +}); + +test('TUI MCP starts a fresh bounded cleanup window after an early cancellation', async () => { + const manager = cancellableTestManager({ + disconnect: async (signal) => { + await new Promise((resolve, reject) => { + signal?.addEventListener( + 'abort', + () => reject(signal.reason ?? new Error('cleanup deadline')), + { once: true }, + ); + setTimeout(resolve, 30); + }); + }, + }); + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connectionHarness().connection }, + { + configStore: configStoreHarness(async () => docsConfig()), + manager: manager.manager, + createProvider: () => undefined, + actionTimeoutMs: 500, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before early cancellation cleanup", + ); + const abort = new AbortController(); + const testing = controller.execute({ kind: 'test', serverId: 'docs' }, { signal: abort.signal }); + await waitFor( + () => manager.testSignal !== undefined, + 'manager test to begin before early cancellation', + ); + abort.abort(new Error('cancel test immediately')); + + assert.deepEqual(await testing, { status: 'failed', reason: 'cancelled' }); + assert.ok(manager.disconnectSignal); + assert.equal(manager.disconnectSignal.aborted, false); + await controller.close(); +}); + +test('TUI MCP waits for manager synchronization before publishing an action snapshot', async () => { + const store = mutableConfigStore(emptyConfig(), []); + const actionSync = deferred(); + let syncCount = 0; + let listener: (() => void) | undefined; + let revision = 0; + const manager = { + sync: async () => { + syncCount += 1; + revision += 1; + listener?.(); + if (syncCount === 2) await actionSync.promise; + }, + statuses: () => [], + toolSnapshot: () => ({ revision, tools: [{}] }) as unknown as McpToolSnapshot, + callTool: async () => ({ content: [] }), + test: async () => ({ + ok: true, + status: connectedStatus('local', 1), + latencyMs: 1, + }), + reconnect: async () => connectedStatus('local', 1), + forgetServerCredentials: async () => undefined, + onChange: (next: () => void) => { + listener = next; + return () => { + listener = undefined; + }; + }, + close: async () => undefined, + } as unknown as ReturnType['manager']; + const connection = connectionHarness(); + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connection.connection }, + { + configStore: store.store, + manager, + createProvider: () => provider('provider'), + }, + ); + await waitFor( + () => controller.snapshot().publication === 'published', + "initial TUI MCP publication to reach 'published' before action sync", + ); + connection.replacements.length = 0; + + const executing = controller.execute({ + kind: 'add', + serverId: 'local', + config: { command: 'server' }, + }); + await waitFor( + () => syncCount === 2, + 'manager sync to reach count 2 before publishing action snapshot', + ); + assert.equal(connection.replacements.length, 0); + assert.equal(controller.snapshot().configuration, 'synchronizing'); + actionSync.resolve(); + + assert.deepEqual(await executing, { status: 'applied', effect: 'published' }); + assert.equal(connection.replacements.length, 1); + await controller.close(); +}); + +test('TUI MCP rebases an action over an unrelated concurrent config edit', async () => { + let config: McpConfigFile = { + version: 3, + mcpServers: { existing: { command: 'before' } }, + }; + const store = { + get: async () => structuredClone(config), + transform: async ( + apply: (current: McpConfigFile) => McpConfigFile | Promise, + ) => { + config = await apply({ + version: 3, + mcpServers: { existing: { command: 'concurrent' } }, + }); + return structuredClone(config); + }, + }; + const manager = managementManager([]); + const connection = connectionHarness(); + const controller = createTuiMcpController( + { workspaceRoot: '/unused', connection: connection.connection }, + { + configStore: store, + manager: manager.manager, + createProvider: () => undefined, + }, + ); + await waitFor( + () => controller.snapshot().initialization === 'ready', + "TUI MCP initialization to reach 'ready' before rebased concurrent edit", + ); + + assert.deepEqual( + await controller.execute({ + kind: 'add', + serverId: 'local', + config: { command: 'server' }, + }), + { status: 'applied', effect: 'published' }, + ); + const existing = config.mcpServers.existing; + assert.ok(existing && 'command' in existing); + assert.equal(existing.command, 'concurrent'); + assert.ok(config.mcpServers.local); + await controller.close(); +}); + +test('independent TUI controllers preserve concurrent additions in one workspace', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'maka-tui-mcp-concurrent-')); + t.after(() => rm(root, { recursive: true, force: true })); + await createMcpConfigStore(root).transform((current) => current); + const left = createTuiMcpController( + { workspaceRoot: root, connection: connectionHarness().connection }, + { + configStore: createMcpConfigStore(root), + manager: managementManager([]).manager, + createProvider: () => undefined, + }, + ); + const right = createTuiMcpController( + { workspaceRoot: root, connection: connectionHarness().connection }, + { + configStore: createMcpConfigStore(root), + manager: managementManager([]).manager, + createProvider: () => undefined, + }, + ); + await waitFor( + () => left.snapshot().initialization === 'ready' && right.snapshot().initialization === 'ready', + "both TUI controllers to reach 'ready' before concurrent additions", ); const [leftResult, rightResult] = await Promise.all([ - left.execute({ kind: 'add', serverId: 'left', config: { command: 'left-server' } }), - right.execute({ kind: 'add', serverId: 'right', config: { command: 'right-server' } }), + left.execute({ + kind: 'add', + serverId: 'left', + config: { command: 'left-server' }, + }), + right.execute({ + kind: 'add', + serverId: 'right', + config: { command: 'right-server' }, + }), ]); assert.equal(leftResult.status, 'applied'); @@ -727,7 +1908,9 @@ test('independent TUI controllers preserve concurrent additions in one workspace test('same-server credential retirement stays inside the shared config transaction', async (t) => { const root = await mkdtemp(join(tmpdir(), 'maka-tui-mcp-retirement-')); t.after(() => rm(root, { recursive: true, force: true })); - await createMcpConfigStore(root).upsert('docs', { url: 'https://old.example/mcp' }); + await createMcpConfigStore(root).upsert('docs', { + url: 'https://old.example/mcp', + }); const retirement = deferred(); const leftOrder: string[] = []; const rightOrder: string[] = []; @@ -735,7 +1918,9 @@ test('same-server credential retirement stays inside the shared config transacti { workspaceRoot: root, connection: connectionHarness().connection }, { configStore: createMcpConfigStore(root), - manager: managementManager(leftOrder, { credentialWait: retirement.promise }).manager, + manager: managementManager(leftOrder, { + credentialWait: retirement.promise, + }).manager, createProvider: () => undefined, }, ); @@ -800,7 +1985,11 @@ test('TUI MCP manages enabled state, tests, reconnects, and removes through one const connection = connectionHarness(); const controller = createTuiMcpController( { workspaceRoot: '/unused', connection: connection.connection }, - { configStore: store.store, manager: manager.manager, createProvider: () => undefined }, + { + configStore: store.store, + manager: manager.manager, + createProvider: () => undefined, + }, ); await waitFor( () => controller.snapshot().initialization === 'ready', @@ -809,7 +1998,11 @@ test('TUI MCP manages enabled state, tests, reconnects, and removes through one order.length = 0; assert.deepEqual( - await controller.execute({ kind: 'set_enabled', serverId: 'docs', enabled: true }), + await controller.execute({ + kind: 'set_enabled', + serverId: 'docs', + enabled: true, + }), { status: 'applied', effect: 'published' }, ); assert.equal((await store.store.get()).mcpServers.docs?.enabled, true); @@ -831,6 +2024,102 @@ test('TUI MCP manages enabled state, tests, reconnects, and removes through one await controller.close(); }); +function docsConfig(): McpConfigFile { + return { version: 3, mcpServers: { docs: { command: 'server' } } }; +} + +function cancellableTestManager( + options: { + readonly status?: McpServerStatus; + readonly revision?: number; + readonly lateSuccessfulTest?: boolean; + readonly notifyOnDisconnect?: boolean; + readonly disconnect?: (signal?: AbortSignal) => Promise; + } = {}, +) { + let listener: (() => void) | undefined; + let revision = options.revision ?? 0; + let status: McpServerStatus = options.status ?? { + serverId: 'docs', + state: 'disconnected', + toolCount: 0, + tools: [], + updatedAt: 0, + }; + let testSignal: AbortSignal | undefined; + let disconnectSignal: AbortSignal | undefined; + let disconnects = 0; + const manager = managementManager([], { + test: async (_serverId, testOptions) => { + testSignal = testOptions?.signal; + if (options.lateSuccessfulTest) { + await new Promise((resolve) => { + testSignal?.addEventListener('abort', () => resolve(), { + once: true, + }); + }); + status = connectedStatus('docs', 1); + listener?.(); + return { ok: true, status, latencyMs: 1 }; + } + return new Promise((_resolve, reject) => { + testSignal?.addEventListener( + 'abort', + () => reject(testSignal?.reason ?? new Error('cancelled')), + { once: true }, + ); + }); + }, + disconnect: async (_serverId, _remove, disconnectOptions) => { + disconnects += 1; + disconnectSignal = disconnectOptions?.signal; + revision += 1; + status = { ...status, state: 'disconnected', toolCount: 0, updatedAt: 2 }; + await options.disconnect?.(disconnectSignal); + if (options.notifyOnDisconnect !== false) listener?.(); + }, + }).manager; + return { + manager: { + ...manager, + statuses: () => [status], + toolSnapshot: () => + ({ + revision, + tools: new Array(status.toolCount).fill({}), + }) as unknown as McpToolSnapshot, + reconnect: async () => status, + onChange: (next: () => void) => { + listener = next; + return () => { + if (listener === next) listener = undefined; + }; + }, + } as unknown as Pick< + McpClientManager, + | 'sync' + | 'statuses' + | 'toolSnapshot' + | 'callTool' + | 'test' + | 'reconnect' + | 'disconnect' + | 'forgetServerCredentials' + | 'onChange' + | 'close' + >, + get testSignal() { + return testSignal; + }, + get disconnectSignal() { + return disconnectSignal; + }, + get disconnects() { + return disconnects; + }, + }; +} + function mutableConfigStore(initial: McpConfigFile, order: string[]) { let config = structuredClone(initial); const store = { @@ -856,14 +2145,38 @@ function mutableConfigStore(initial: McpConfigFile, order: string[]) { function managementManager( order: string[], - options: { readonly credentialFailure?: boolean; readonly credentialWait?: Promise } = {}, + options: { + readonly credentialFailure?: boolean; + readonly credentialWait?: Promise; + readonly sync?: ( + config: McpConfigFile, + options?: { readonly signal?: AbortSignal }, + ) => Promise; + readonly test?: ( + serverId: string, + options?: { readonly signal?: AbortSignal }, + ) => ReturnType; + readonly disconnect?: ( + serverId: string, + remove: boolean, + options?: { readonly signal?: AbortSignal }, + ) => Promise; + readonly forgetServerCredentials?: ( + serverId: string, + config: unknown, + options?: { + readonly signal?: AbortSignal; + readonly onCommitStarted?: () => void; + }, + ) => Promise; + } = {}, ) { let listener: (() => void) | undefined; let syncFailure = false; let revision = 0; const statuses: McpServerStatus[] = []; const manager = { - sync: async () => { + sync: async (config: McpConfigFile, syncOptions?: { readonly signal?: AbortSignal }) => { order.push('sync'); if (syncFailure) { syncFailure = false; @@ -871,21 +2184,41 @@ function managementManager( } revision += 1; listener?.(); + await options.sync?.(config, syncOptions); }, statuses: () => statuses, toolSnapshot: () => ({ revision, tools: [] }) as McpToolSnapshot, callTool: async () => ({ content: [] }), - test: async (serverId: string) => { + test: async (serverId: string, testOptions?: { readonly signal?: AbortSignal }) => { order.push(`test:${serverId}`); + if (options.test) return options.test(serverId, testOptions); return { ok: true, status: connectedStatus(serverId, 0), latencyMs: 1 }; }, reconnect: async (serverId: string) => { order.push(`reconnect:${serverId}`); return connectedStatus(serverId, 0); }, - forgetServerCredentials: async (serverId: string) => { + disconnect: async ( + serverId: string, + remove = false, + disconnectOptions?: { readonly signal?: AbortSignal }, + ) => { + await options.disconnect?.(serverId, remove, disconnectOptions); + }, + forgetServerCredentials: async ( + serverId: string, + config: unknown, + credentialOptions?: { + readonly signal?: AbortSignal; + readonly onCommitStarted?: () => void; + }, + ) => { order.push(`forget:${serverId}`); if (options.credentialFailure) throw new Error('credential cleanup failed'); + if (options.forgetServerCredentials) { + await options.forgetServerCredentials(serverId, config, credentialOptions); + return; + } await options.credentialWait; }, onChange: (next: () => void) => { @@ -903,6 +2236,7 @@ function managementManager( | 'callTool' | 'test' | 'reconnect' + | 'disconnect' | 'forgetServerCredentials' | 'onChange' | 'close' @@ -929,7 +2263,11 @@ function managerHarness(revision: number, statuses: McpServerStatus[]) { tools: new Array(toolCount).fill({}), }) as McpToolSnapshot, callTool: async () => ({ content: [] }), - test: async () => ({ ok: true, status: connectedStatus('local', 1), latencyMs: 1 }), + test: async () => ({ + ok: true, + status: connectedStatus('local', 1), + latencyMs: 1, + }), reconnect: async () => connectedStatus('local', 1), forgetServerCredentials: async () => undefined, onChange: (next: () => void) => { @@ -1034,6 +2372,30 @@ function emptyConfig(): McpConfigFile { return { version: 3, mcpServers: {} }; } +async function readFixtureEvents( + path: string, +): Promise> { + try { + return (await readFile(path, 'utf8')) + .trim() + .split('\n') + .filter(Boolean) + .map((line) => JSON.parse(line) as { event: string; pid: number }); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') return []; + throw error; + } +} + +function processExists(pid: number): boolean { + try { + process.kill(pid, 0); + return true; + } catch (error) { + return (error as NodeJS.ErrnoException).code === 'EPERM'; + } +} + function configStoreHarness(get: () => Promise) { return { get, @@ -1048,3 +2410,20 @@ function deferredValue() { }); return { promise, resolve }; } + +async function settlesWithin(promise: Promise, timeoutMs: number): Promise { + let timer: NodeJS.Timeout | undefined; + try { + return await Promise.race([ + promise.then( + () => true, + () => true, + ), + new Promise((resolve) => { + timer = setTimeout(() => resolve(false), timeoutMs); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } +} diff --git a/packages/cli/src/__tests__/tui-mcp-remote-publication.test.ts b/packages/cli/src/__tests__/tui-mcp-remote-publication.test.ts index 77b8338f5d..face50ae3a 100644 --- a/packages/cli/src/__tests__/tui-mcp-remote-publication.test.ts +++ b/packages/cli/src/__tests__/tui-mcp-remote-publication.test.ts @@ -102,7 +102,10 @@ test('remote TUI publication activates, rotates, and removes one profile-bound c }, ); let latest = await availability(target); - assert.deepEqual(latest(), { kind: 'unavailable', reason: 'credential_required' }); + assert.deepEqual(latest(), { + kind: 'unavailable', + reason: 'credential_required', + }); await target.setCredential?.('provider-secret-a'); await waitFor(() => latest().kind === 'connected', 'provider companion to connect'); @@ -129,13 +132,287 @@ test('remote TUI publication activates, rotates, and removes one profile-bound c assert.equal(identityPaths[0], identityPaths[1]); await target.removeCredential?.(); - assert.deepEqual(latest(), { kind: 'unavailable', reason: 'credential_required' }); + assert.deepEqual(latest(), { + kind: 'unavailable', + reason: 'credential_required', + }); assert.equal(connections[1]?.unregisters, 0); assert.equal(connections[1]?.closes, 1); assert.equal(credentials.values.has('office\0incarnation-a\0terminal-client'), false); await target.closePublication?.(); }); +test('remote TUI publication aborts and rolls back a cancelled credential write', async () => { + const credentials = credentialHarness(); + const writeStarted = deferred(); + const allowWrite = deferred(); + let connects = 0; + const target = createRemoteTuiMcpPublicationTarget( + { + clientDataRoot: '/client-data', + profile: PROFILE, + profileIncarnationId: PROFILE_INCARNATION_ID, + ownerClientInstanceId: 'terminal-client', + }, + { + ...profileDeps(), + credentials: { + ...credentials.store, + compareAndSet: async (...args) => { + writeStarted.resolve(); + await allowWrite.promise; + assert.ok(credentials.store.compareAndSet); + return credentials.store.compareAndSet(...args); + }, + }, + connectProfile: async () => { + connects += 1; + return connectionHarness('unexpected').connection; + }, + }, + ); + const latest = await availability(target); + assert.deepEqual(latest(), { + kind: 'unavailable', + reason: 'credential_required', + }); + const abort = new AbortController(); + const setting = target.setCredential?.('provider-secret', { + signal: abort.signal, + }); + assert.ok(setting); + await writeStarted.promise; + abort.abort(new Error('cancel credential write')); + allowWrite.resolve(); + + await assert.rejects(setting, /cancel credential write/u); + assert.equal(credentials.values.size, 0); + assert.equal(connects, 0); + assert.deepEqual(latest(), { + kind: 'unavailable', + reason: 'credential_required', + }); + await target.closePublication?.(); +}); + +test('remote TUI publication connects a newer credential after cancellation before disconnect', async () => { + const credentials = credentialHarness('provider-secret-a'); + const writeStarted = deferred(); + const allowWrite = deferred(); + const connected: string[] = []; + const connections: ConnectionHarness[] = []; + let firstMutation = true; + const target = createRemoteTuiMcpPublicationTarget( + { + clientDataRoot: '/client-data', + profile: PROFILE, + profileIncarnationId: PROFILE_INCARNATION_ID, + ownerClientInstanceId: 'terminal-client', + }, + { + ...profileDeps(), + credentials: { + ...credentials.store, + compareAndSet: async (...args) => { + assert.ok(credentials.store.compareAndSet); + if (!firstMutation) return credentials.store.compareAndSet(...args); + firstMutation = false; + writeStarted.resolve(); + await allowWrite.promise; + const result = await credentials.store.compareAndSet(...args); + credentials.setConcurrent('newer-secret'); + return result; + }, + }, + loadClientInstanceId: async () => 'provider-client', + connectProfile: async (input) => { + connected.push(input.credential ?? ''); + const connection = connectionHarness(`connection-${connections.length + 1}`); + connections.push(connection); + return connection.connection; + }, + }, + ); + const latest = await availability(target); + await waitFor(() => latest().kind === 'connected', 'initial provider companion to connect'); + const abort = new AbortController(); + const setting = target.setCredential?.('provider-secret', { + signal: abort.signal, + }); + assert.ok(setting); + await writeStarted.promise; + abort.abort(new Error('cancel superseded credential write')); + allowWrite.resolve(); + + await assert.rejects(setting, /cancel superseded credential write/u); + assert.equal(credentials.values.get('office\0incarnation-a\0terminal-client'), 'newer-secret'); + await waitFor( + () => latest().kind === 'connected' && connected.at(-1) === 'newer-secret', + 'concurrent winner credential to connect', + ); + assert.deepEqual(connected, ['provider-secret-a', 'newer-secret']); + assert.equal(connections[0]?.closes, 1); + await target.closePublication?.(); +}); + +test('remote TUI publication connects the concurrent winner after cancelled rotation', async () => { + const credentials = credentialHarness('provider-secret-a'); + const rotatedConnectStarted = deferred(); + const allowRotatedConnect = deferred(); + const connected: string[] = []; + const target = createRemoteTuiMcpPublicationTarget( + { + clientDataRoot: '/client-data', + profile: PROFILE, + profileIncarnationId: PROFILE_INCARNATION_ID, + ownerClientInstanceId: 'terminal-client', + }, + { + ...profileDeps(), + credentials: credentials.store, + loadClientInstanceId: async () => 'provider-client', + connectProfile: async (input) => { + connected.push(input.credential ?? ''); + if (input.credential === 'provider-secret-b') { + rotatedConnectStarted.resolve(); + await allowRotatedConnect.promise; + throw input.signal?.reason ?? new Error('rotated connection cancelled'); + } + return connectionHarness(`connection-${connected.length}`).connection; + }, + }, + ); + const latest = await availability(target); + await waitFor(() => latest().kind === 'connected', 'initial provider companion to connect'); + const abort = new AbortController(); + const setting = target.setCredential?.('provider-secret-b', { + signal: abort.signal, + }); + assert.ok(setting); + await Promise.race([rotatedConnectStarted.promise, setting]); + credentials.setConcurrent('provider-secret-c'); + abort.abort(new Error('cancel credential rotation')); + allowRotatedConnect.resolve(); + + await assert.rejects(setting, /cancel credential rotation/u); + assert.equal( + credentials.values.get('office\0incarnation-a\0terminal-client'), + 'provider-secret-c', + ); + await waitFor( + () => latest().kind === 'connected' && connected.at(-1) === 'provider-secret-c', + 'concurrent winner credential to connect', + ); + assert.deepEqual(connected, ['provider-secret-a', 'provider-secret-b', 'provider-secret-c']); + await target.closePublication?.(); +}); + +test('remote TUI publication does not restore through a same-value ABA rotation', async () => { + const credentials = credentialHarness('provider-secret-a'); + const rotatedConnectStarted = deferred(); + const allowRotatedConnect = deferred(); + const connected: string[] = []; + const target = createRemoteTuiMcpPublicationTarget( + { + clientDataRoot: '/client-data', + profile: PROFILE, + profileIncarnationId: PROFILE_INCARNATION_ID, + ownerClientInstanceId: 'terminal-client', + }, + { + ...profileDeps(), + credentials: credentials.store, + loadClientInstanceId: async () => 'provider-client', + connectProfile: async (input) => { + connected.push(input.credential ?? ''); + if (input.credential === 'provider-secret-b' && connected.length === 2) { + rotatedConnectStarted.resolve(); + await allowRotatedConnect.promise; + throw input.signal?.reason ?? new Error('rotated connection cancelled'); + } + return connectionHarness(`connection-${connected.length}`).connection; + }, + }, + ); + const latest = await availability(target); + await waitFor(() => latest().kind === 'connected', 'initial provider companion to connect'); + const abort = new AbortController(); + const setting = target.setCredential?.('provider-secret-b', { + signal: abort.signal, + }); + assert.ok(setting); + await Promise.race([rotatedConnectStarted.promise, setting]); + credentials.setConcurrent('provider-secret-c'); + credentials.setConcurrent('provider-secret-b'); + abort.abort(new Error('cancel ABA credential rotation')); + allowRotatedConnect.resolve(); + + await assert.rejects(setting, /cancel ABA credential rotation/u); + assert.equal( + credentials.values.get('office\0incarnation-a\0terminal-client'), + 'provider-secret-b', + ); + await waitFor( + () => latest().kind === 'connected' && connected.length === 3, + 'same-value concurrent winner credential to connect', + ); + assert.deepEqual(connected, ['provider-secret-a', 'provider-secret-b', 'provider-secret-b']); + await target.closePublication?.(); +}); + +test('remote TUI publication restores a removed credential when cancellation lands late', async () => { + const credentials = credentialHarness('provider-secret'); + const deleting = deferred(); + const allowDelete = deferred(); + const connections: ConnectionHarness[] = []; + const target = createRemoteTuiMcpPublicationTarget( + { + clientDataRoot: '/client-data', + profile: PROFILE, + profileIncarnationId: PROFILE_INCARNATION_ID, + ownerClientInstanceId: 'terminal-client', + }, + { + ...profileDeps(), + credentials: { + ...credentials.store, + compareAndSet: async (...args) => { + deleting.resolve(); + await allowDelete.promise; + assert.ok(credentials.store.compareAndSet); + return credentials.store.compareAndSet(...args); + }, + }, + loadClientInstanceId: async () => 'provider-client', + connectProfile: async () => { + const connection = connectionHarness(`connection-${connections.length + 1}`); + connections.push(connection); + return connection.connection; + }, + }, + ); + const latest = await availability(target); + await waitFor( + () => latest().kind === 'connected', + 'provider companion to connect before removal', + ); + const abort = new AbortController(); + const removing = target.removeCredential?.({ signal: abort.signal }); + assert.ok(removing); + await deleting.promise; + abort.abort(new Error('cancel credential removal')); + allowDelete.resolve(); + + await assert.rejects(removing, /cancel credential removal/u); + assert.equal(credentials.values.get('office\0incarnation-a\0terminal-client'), 'provider-secret'); + await waitFor( + () => latest().kind === 'connected', + 'provider companion to reconnect after rollback', + ); + assert.equal(connections.length, 2); + await target.closePublication?.(); +}); + test('remote TUI publication fails closed while the same provider lifetime is active', async () => { const credentials = credentialHarness('provider-secret'); const profiles = profileHarness(); @@ -500,7 +777,10 @@ test('remote TUI publication cannot write into a recreated profile incarnation', }, ); const latest = await availability(target); - assert.deepEqual(latest(), { kind: 'unavailable', reason: 'credential_required' }); + assert.deepEqual(latest(), { + kind: 'unavailable', + reason: 'credential_required', + }); const heldMutation = profiles.holdNextMutation(); const setCredential = target.setCredential?.('replacement-secret'); @@ -513,7 +793,10 @@ test('remote TUI publication cannot write into a recreated profile incarnation', await assert.rejects(setCredential, /profile is no longer current/u); assert.equal(credentials.values.has('office\0incarnation-a\0terminal-client'), false); assert.equal(credentials.values.has('office\0incarnation-b\0terminal-client'), false); - assert.deepEqual(latest(), { kind: 'unavailable', reason: 'target_mismatch' }); + assert.deepEqual(latest(), { + kind: 'unavailable', + reason: 'target_mismatch', + }); await target.closePublication?.(); }); @@ -552,7 +835,10 @@ test('remote TUI publication cannot register capabilities after profile recreati error instanceof RuntimeHostProfileConnectionError && error.reason === 'target_mismatch', ); assert.equal(connection.replacements, 0); - assert.deepEqual(latest(), { kind: 'unavailable', reason: 'target_mismatch' }); + assert.deepEqual(latest(), { + kind: 'unavailable', + reason: 'target_mismatch', + }); await target.closePublication?.(); }); @@ -670,7 +956,9 @@ test('remote TUI publication closes its direct peer after a permanent reconnect loadClientInstanceId: async () => 'provider-client', connectProfile: async () => initial.connection, createPeerClient: () => - ({ close: async () => void (peerCloses += 1) }) as RuntimeHostPeerClient, + ({ + close: async () => void (peerCloses += 1), + }) as RuntimeHostPeerClient, createReconnectingConnection: async (input) => { fatal = input.onFatalError; return reconnectingConnection(initial.connection); @@ -685,7 +973,10 @@ test('remote TUI publication closes its direct peer after a permanent reconnect fatal?.(new RuntimeHostProfileConnectionError('credential_rejected', 'revoked')); await waitFor(() => peerCloses === 1, 'direct peer endpoint to close'); - assert.deepEqual(latest(), { kind: 'unavailable', reason: 'credential_rejected' }); + assert.deepEqual(latest(), { + kind: 'unavailable', + reason: 'credential_rejected', + }); await target.closePublication?.(); assert.equal(peerCloses, 1); }); @@ -732,31 +1023,79 @@ function publicationLeaseHarness() { function credentialHarness(initial?: string) { const values = new Map(); - if (initial) values.set('office\0incarnation-a\0terminal-client', initial); + const revisions = new Map(); + const initialKey = 'office\0incarnation-a\0terminal-client'; + if (initial) { + values.set(initialKey, initial); + revisions.set(initialKey, 1); + } const key = (target: RuntimeHostRemoteProfileIncarnation, ownerClientInstanceId: string) => `${target.profile.id}\0${target.profileIncarnationId}\0${ownerClientInstanceId}`; + const mutate = (entryKey: string, credential: string | null) => { + if (credential === null) values.delete(entryKey); + else values.set(entryKey, credential); + const revision = (revisions.get(entryKey) ?? 0) + 1; + revisions.set(entryKey, revision); + return String(revision); + }; const store: RuntimeHostCapabilityProviderCredentialStore = { get: async (target, ownerClientInstanceId) => values.get(key(target, ownerClientInstanceId)) ?? null, + read: async (target, ownerClientInstanceId) => { + const entryKey = key(target, ownerClientInstanceId); + return { + credential: values.get(entryKey) ?? null, + revision: revisions.has(entryKey) ? String(revisions.get(entryKey)) : null, + }; + }, set: async (target, ownerClientInstanceId, credential) => { - values.set(key(target, ownerClientInstanceId), credential); + mutate(key(target, ownerClientInstanceId), credential); }, delete: async (target, ownerClientInstanceId) => { - values.delete(key(target, ownerClientInstanceId)); + mutate(key(target, ownerClientInstanceId), null); + }, + compareAndSet: async (target, ownerClientInstanceId, expectedRevision, credential) => { + const entryKey = key(target, ownerClientInstanceId); + const currentRevision = revisions.has(entryKey) ? String(revisions.get(entryKey)) : null; + if (currentRevision !== expectedRevision) { + return { + committed: false, + current: { + credential: values.get(entryKey) ?? null, + revision: currentRevision, + }, + }; + } + return { committed: true, revision: mutate(entryKey, credential) }; + }, + }; + return { + store, + values, + setConcurrent(credential: string | null) { + mutate(initialKey, credential); }, }; - return { store, values }; } function profileDeps(profile: RemoteRuntimeHostProfile = PROFILE) { const profiles = profileHarness(profile); - return { profiles: profiles.catalog, subscribeProfileChanges: profiles.subscribe }; + return { + profiles: profiles.catalog, + subscribeProfileChanges: profiles.subscribe, + }; } function profileHarness(initial: RemoteRuntimeHostProfile = PROFILE) { let current: - | { readonly profile: RemoteRuntimeHostProfile; readonly profileIncarnationId: string } - | undefined = { profile: initial, profileIncarnationId: PROFILE_INCARNATION_ID }; + | { + readonly profile: RemoteRuntimeHostProfile; + readonly profileIncarnationId: string; + } + | undefined = { + profile: initial, + profileIncarnationId: PROFILE_INCARNATION_ID, + }; const listeners = new Set<(error?: Error) => void>(); let heldValidation: | { @@ -875,11 +1214,17 @@ function connectionHarness( closed, replaceClientCapabilities: async () => { harness.replacements += 1; - return { registrationId: 'registration-a', revision: harness.replacements }; + return { + registrationId: 'registration-a', + revision: harness.replacements, + }; }, unregisterClientCapabilities: async () => { harness.unregisters += 1; - return { registrationId: 'registration-a', revision: harness.unregisters }; + return { + registrationId: 'registration-a', + revision: harness.unregisters, + }; }, subscribeConfigurationChanges: () => () => undefined, subscribeConnectionCatalogChanges: () => () => undefined, diff --git a/packages/cli/src/pi-tui-mcp-status.ts b/packages/cli/src/pi-tui-mcp-status.ts index ff61bcae81..6e7339ec9c 100644 --- a/packages/cli/src/pi-tui-mcp-status.ts +++ b/packages/cli/src/pi-tui-mcp-status.ts @@ -167,6 +167,7 @@ export class McpManagementOverlay implements Component { private editor: OverlayTextInput | undefined; private closed = false; private actionAttempt = 0; + private actionAbort: AbortController | undefined; constructor( private readonly input: { @@ -197,10 +198,9 @@ export class McpManagementOverlay implements Component { } if (this.phase.kind === 'busy') { if (matchesKey(data, Key.escape)) { - this.actionAttempt += 1; + this.cancelActiveAction(); this.backToList(); } else if (matchesKey(data, 'q')) { - this.actionAttempt += 1; this.close(); } return; @@ -457,14 +457,17 @@ export class McpManagementOverlay implements Component { } this.clearEditor(); const attempt = ++this.actionAttempt; + const abort = new AbortController(); + this.actionAbort = abort; this.phase = { kind: 'busy', label: actionLabel(action, this.input.locale) }; this.input.onChange(); let result: TuiMcpActionResult; try { - result = await management.execute(action); + result = await management.execute(action, { signal: abort.signal }); } catch { result = { status: 'failed', reason: 'manager-failed' }; } + if (this.actionAbort === abort) this.actionAbort = undefined; if (this.closed || attempt !== this.actionAttempt) return; this.phase = { kind: 'list' }; this.notice = actionNotice(result, this.input.locale); @@ -585,6 +588,7 @@ export class McpManagementOverlay implements Component { private close(): void { if (this.closed) return; this.closed = true; + this.cancelActiveAction(); if (this.phase.kind === 'confirm_import') { this.management()?.discardImportPreview(this.phase.preview.previewId); } @@ -593,6 +597,13 @@ export class McpManagementOverlay implements Component { this.input.onClose(); } + private cancelActiveAction(): void { + this.actionAttempt += 1; + const abort = this.actionAbort; + this.actionAbort = undefined; + abort?.abort(new Error('MCP action cancelled')); + } + private keepSelectionVisible(): void { const rows = this.serverRows[this.selected]; if (!rows || this.bodyRows <= 0) return; diff --git a/packages/cli/src/tui-copy-catalog.ts b/packages/cli/src/tui-copy-catalog.ts index 44f8efa249..2ffd471cd9 100644 --- a/packages/cli/src/tui-copy-catalog.ts +++ b/packages/cli/src/tui-copy-catalog.ts @@ -149,12 +149,15 @@ export const TUI_COPY_RESOURCES = { stale_import: 'An imported entry changed; preview the import again.', missing: 'That server no longer exists.', closed: 'The MCP controller is closed.', + cancelled: 'The MCP action was cancelled.', 'invalid-config': 'The server configuration is invalid.', 'credential-cleanup-failed': 'Stored credentials could not be removed; the configuration was not changed.', 'publication-credential-failed': 'The provider credential could not be stored or applied.', 'persist-failed': 'The configuration could not be saved.', + 'rollback-failed': + 'The MCP action was cancelled, but cleanup failed. The configuration may be out of sync.', 'manager-failed': 'The MCP connection action failed.', turn_active: 'MCP cannot be changed while a turn or another control action is running.', invalid: 'Check the value and try again.', @@ -250,10 +253,12 @@ export const TUI_COPY_RESOURCES = { stale_import: '导入项已变化，请重新预览。', missing: '该服务器已不存在。', closed: 'MCP 控制器已关闭。', + cancelled: 'MCP 操作已取消。', 'invalid-config': '服务器配置无效。', 'credential-cleanup-failed': '无法删除旧凭据，配置未修改。', 'publication-credential-failed': '无法保存或应用 Provider 凭据。', 'persist-failed': '无法保存配置。', + 'rollback-failed': 'MCP 操作已取消，但清理失败，配置可能不同步。', 'manager-failed': 'MCP 连接操作失败。', turn_active: 'Turn 或其他控制操作运行期间不能修改 MCP。', invalid: '请检查输入后重试。', @@ -350,10 +355,12 @@ export const TUI_COPY_RESOURCES = { stale_import: '匯入項目已變更，請重新預覽。', missing: '該伺服器已不存在。', closed: 'MCP 控制器已關閉。', + cancelled: 'MCP 操作已取消。', 'invalid-config': '伺服器設定無效。', 'credential-cleanup-failed': '無法移除已儲存的認證資料；設定未變更。', 'publication-credential-failed': '無法儲存或套用 Provider 認證資料。', 'persist-failed': '無法儲存設定。', + 'rollback-failed': 'MCP 操作已取消，但清理失敗，設定可能不同步。', 'manager-failed': 'MCP 連線操作失敗。', turn_active: 'Turn 或其他控制操作執行期間無法修改 MCP。', invalid: '請檢查輸入後重試。', @@ -453,7 +460,7 @@ export const TUI_COPY_RESOURCES = { en: { modelPickerTitle: 'Select Model', modelSwitchCacheWarning: - '⚠ Switching models may rebuild the prompt cache; the next request may be slower or cost more.', + '\u26a0 Switching models may rebuild the prompt cache; the next request may be slower or cost more.', modelSearchHint: 'Search models / providers / connections · ↑↓ select · Enter confirm · Esc cancel', searchLabel: 'Search', @@ -553,7 +560,8 @@ export const TUI_COPY_RESOURCES = { }, 'zh-CN': { modelPickerTitle: '选择模型', - modelSwitchCacheWarning: '⚠ 切换模型可能需要重建提示缓存；下一次请求可能更慢或成本更高。', + modelSwitchCacheWarning: + '\u26a0 切换模型可能需要重建提示缓存；下一次请求可能更慢或成本更高。', modelSearchHint: '搜索模型 / 服务商 / 连接 · ↑↓ 选择 · Enter 确认 · Esc 取消', searchLabel: '搜索', noMatchingModels: '没有匹配的模型', @@ -642,7 +650,8 @@ export const TUI_COPY_RESOURCES = { }, 'zh-TW': { modelPickerTitle: '選擇模型', - modelSwitchCacheWarning: '⚠ 切換模型可能需要重建提示快取；下一次請求可能較慢或成本較高。', + modelSwitchCacheWarning: + '\u26a0 切換模型可能需要重建提示快取；下一次請求可能較慢或成本較高。', modelSearchHint: '搜尋模型 / 服務商 / 連線 · ↑↓ 選擇 · Enter 確認 · Esc 取消', searchLabel: '搜尋', noMatchingModels: '沒有符合的模型', diff --git a/packages/cli/src/tui-mcp-control.ts b/packages/cli/src/tui-mcp-control.ts index e0df998355..9a23f4f8e5 100644 --- a/packages/cli/src/tui-mcp-control.ts +++ b/packages/cli/src/tui-mcp-control.ts @@ -47,6 +47,8 @@ import type { import { createMcpCapabilityProvider } from './mcp-capability-provider.js'; const RUNTIME_HOST_CREDENTIAL_ENV = 'MAKA_RUNTIME_HOST_ACCESS_CREDENTIAL'; +const MCP_ACTION_TIMEOUT_MS = 90_000; +const MCP_ACTION_CLEANUP_RESERVE_MS = 5_000; export type TuiMcpPublicationState = | 'waiting' @@ -113,7 +115,11 @@ export type TuiMcpImportPreviewResult = }; export type TuiMcpAction = - | { readonly kind: 'add'; readonly serverId: string; readonly config: McpServerConfig } + | { + readonly kind: 'add'; + readonly serverId: string; + readonly config: McpServerConfig; + } | { readonly kind: 'edit'; readonly serverId: string; @@ -121,7 +127,11 @@ export type TuiMcpAction = readonly expectedRevision: string; } | { readonly kind: 'commit_import'; readonly previewId: string } - | { readonly kind: 'set_enabled'; readonly serverId: string; readonly enabled: boolean } + | { + readonly kind: 'set_enabled'; + readonly serverId: string; + readonly enabled: boolean; + } | { readonly kind: 'remove'; readonly serverId: string } | { readonly kind: 'test'; readonly serverId: string } | { readonly kind: 'reconnect'; readonly serverId: string } @@ -136,7 +146,11 @@ export type TuiMcpActionEffect = export type TuiMcpActionResult = | { readonly status: 'applied'; readonly effect: TuiMcpActionEffect } - | { readonly status: 'tested'; readonly test: McpTestResult; readonly effect: TuiMcpActionEffect } + | { + readonly status: 'tested'; + readonly test: McpTestResult; + readonly effect: TuiMcpActionEffect; + } | { readonly status: 'conflict'; readonly reason: 'exists' | 'stale_config' | 'stale_edit' | 'stale_import' | 'missing'; @@ -145,10 +159,12 @@ export type TuiMcpActionResult = readonly status: 'failed'; readonly reason: | 'closed' + | 'cancelled' | 'invalid-config' | 'credential-cleanup-failed' | 'publication-credential-failed' | 'persist-failed' + | 'rollback-failed' | 'manager-failed'; }; @@ -156,7 +172,10 @@ export interface TuiMcpManagement extends TuiMcpSurface { configForEdit(serverId: string): TuiMcpEditConfig | undefined; previewImport(source: string): TuiMcpImportPreviewResult; discardImportPreview(previewId: string): void; - execute(action: TuiMcpAction): Promise; + execute( + action: TuiMcpAction, + options?: { readonly signal?: AbortSignal }, + ): Promise; } export interface TuiMcpController extends TuiMcpManagement { @@ -165,16 +184,15 @@ export interface TuiMcpController extends TuiMcpManagement { type TuiMcpManager = Pick< McpClientManager, - | 'sync' - | 'statuses' - | 'toolSnapshot' - | 'callTool' - | 'onChange' - | 'test' - | 'reconnect' - | 'forgetServerCredentials' - | 'close' ->; + 'sync' | 'statuses' | 'toolSnapshot' | 'callTool' | 'onChange' | 'test' | 'reconnect' | 'close' +> & + Pick & { + disconnect?( + serverId: string, + remove?: boolean, + options?: { signal?: AbortSignal }, + ): Promise; + }; export type TuiMcpPublicationUnavailableReason = | 'host_unavailable' @@ -198,8 +216,8 @@ export interface TuiMcpPublicationTarget subscribeConnectionAvailability( listener: (availability: TuiMcpPublicationAvailability) => void, ): () => void; - setCredential?(credential: string): Promise; - removeCredential?(): Promise; + setCredential?(credential: string, options?: { readonly signal?: AbortSignal }): Promise; + removeCredential?(options?: { readonly signal?: AbortSignal }): Promise; closePublication?(): Promise; } @@ -207,6 +225,7 @@ interface TuiMcpControllerDeps { readonly configStore: Pick; readonly manager: TuiMcpManager; readonly createProvider: (manager: TuiMcpManager) => ClientCapabilityProvider | undefined; + readonly actionTimeoutMs: number; } export function createTuiMcpController( @@ -227,6 +246,7 @@ export function createTuiMcpController( configStore: overrides.configStore ?? createMcpConfigStore(input.workspaceRoot), manager, createProvider: overrides.createProvider ?? createMcpCapabilityProvider, + actionTimeoutMs: overrides.actionTimeoutMs ?? MCP_ACTION_TIMEOUT_MS, }); } @@ -248,6 +268,7 @@ class TuiMcpControllerImpl implements TuiMcpController { } | undefined; #actionLane: Promise = Promise.resolve(); + readonly #lifetimeAbort = new AbortController(); #publicationSuppressed = false; #publicationRequested = false; #publicationTask: Promise | undefined; @@ -318,7 +339,10 @@ class TuiMcpControllerImpl implements TuiMcpController { configForEdit(serverId: string): TuiMcpEditConfig | undefined { const config = this.#config?.mcpServers[serverId]; if (!config) return undefined; - return { config: structuredClone(config), revision: configRevision(config) }; + return { + config: structuredClone(config), + revision: configRevision(config), + }; } previewImport(source: string): TuiMcpImportPreviewResult { @@ -359,14 +383,30 @@ class TuiMcpControllerImpl implements TuiMcpController { if (this.#preparedImport?.previewId === previewId) this.#preparedImport = undefined; } - execute(action: TuiMcpAction): Promise { + execute( + action: TuiMcpAction, + options: { readonly signal?: AbortSignal } = {}, + ): Promise { if (this.#closed) return Promise.resolve({ status: 'failed', reason: 'closed' }); - return this.#serializeAction(() => this.#executeAction(action)); + return this.#serializeAction(async () => { + const cleanupReserveMs = Math.min( + MCP_ACTION_CLEANUP_RESERVE_MS, + Math.max(1, Math.floor(this.#deps.actionTimeoutMs / 10)), + ); + const operationDeadline = AbortSignal.timeout( + Math.max(1, this.#deps.actionTimeoutMs - cleanupReserveMs), + ); + const signals = [this.#lifetimeAbort.signal, operationDeadline]; + if (options.signal) signals.push(options.signal); + const signal = AbortSignal.any(signals); + return this.#executeAction(action, signal, cleanupReserveMs); + }); } async close(): Promise { if (this.#closed) return; this.#closed = true; + this.#lifetimeAbort.abort(new Error('MCP controller closed')); this.#disposeManagerChange(); this.#disposeConnectionAvailability(); this.#listeners.clear(); @@ -396,7 +436,10 @@ class TuiMcpControllerImpl implements TuiMcpController { this.#requestPublication(); } catch { if (this.#closed) return; - this.#updateSnapshot({ initialization: 'error', publication: 'not_published' }); + this.#updateSnapshot({ + initialization: 'error', + publication: 'not_published', + }); } } @@ -409,16 +452,38 @@ class TuiMcpControllerImpl implements TuiMcpController { return run; } - async #executeAction(action: TuiMcpAction): Promise { + async #executeAction( + action: TuiMcpAction, + signal?: AbortSignal, + cleanupTimeoutMs?: number, + ): Promise { if (this.#closed) return { status: 'failed', reason: 'closed' }; + if (signal?.aborted) return { status: 'failed', reason: 'cancelled' }; if (action.kind === 'set_publication_credential') { if (!this.#connection.setCredential) { return { status: 'failed', reason: 'publication-credential-failed' }; } + const operation = this.#connection.setCredential(action.credential, { + signal, + }); try { - await this.#connection.setCredential(action.credential); - return { status: 'applied', effect: await this.#settlePublication() }; - } catch { + await waitForAbort(operation, signal); + throwIfAborted(signal); + const effect = await this.#settlePublication(signal); + throwIfAborted(signal); + return { status: 'applied', effect }; + } catch (error) { + if (signal?.aborted) { + return { + status: 'failed', + reason: await this.#settleCancelledCredentialOperation( + operation, + signal, + cleanupTimeoutMs, + ), + }; + } + if (error === signal?.reason) return { status: 'failed', reason: 'cancelled' }; return { status: 'failed', reason: 'publication-credential-failed' }; } } @@ -426,31 +491,78 @@ class TuiMcpControllerImpl implements TuiMcpController { if (!this.#connection.removeCredential) { return { status: 'failed', reason: 'publication-credential-failed' }; } + const operation = this.#connection.removeCredential({ signal }); try { - await this.#connection.removeCredential(); + await waitForAbort(operation, signal); + throwIfAborted(signal); return { status: 'applied', effect: 'pending_host' }; - } catch { + } catch (error) { + if (signal?.aborted) { + return { + status: 'failed', + reason: await this.#settleCancelledCredentialOperation( + operation, + signal, + cleanupTimeoutMs, + ), + }; + } + if (error === signal?.reason) return { status: 'failed', reason: 'cancelled' }; return { status: 'failed', reason: 'publication-credential-failed' }; } } if (action.kind === 'test') { try { - const test = await this.#deps.manager.test(action.serverId); - return { status: 'tested', test, effect: await this.#settlePublication() }; - } catch { + const test = await this.#deps.manager.test(action.serverId, { signal }); + throwIfAborted(signal); + const effect = await this.#settlePublication(signal); + throwIfAborted(signal); + return { status: 'tested', test, effect }; + } catch (error) { + if (signal?.aborted) { + const cleaned = await this.#settleCancelledConnection( + action.serverId, + cleanupSignal(cleanupTimeoutMs), + ); + this.#refreshManagerSnapshot(); + const publicationCleaned = cleaned + ? await this.#settleCancelledPublication(cleanupSignal(cleanupTimeoutMs)) + : false; + return { + status: 'failed', + reason: cleaned && publicationCleaned ? 'cancelled' : 'manager-failed', + }; + } return { status: 'failed', reason: 'manager-failed' }; } } if (action.kind === 'reconnect') { try { - await this.#deps.manager.reconnect(action.serverId); - return { status: 'applied', effect: await this.#settlePublication() }; - } catch { + await this.#deps.manager.reconnect(action.serverId, { signal }); + throwIfAborted(signal); + const effect = await this.#settlePublication(signal); + throwIfAborted(signal); + return { status: 'applied', effect }; + } catch (error) { + if (signal?.aborted) { + const cleaned = await this.#settleCancelledConnection( + action.serverId, + cleanupSignal(cleanupTimeoutMs), + ); + this.#refreshManagerSnapshot(); + const publicationCleaned = cleaned + ? await this.#settleCancelledPublication(cleanupSignal(cleanupTimeoutMs)) + : false; + return { + status: 'failed', + reason: cleaned && publicationCleaned ? 'cancelled' : 'manager-failed', + }; + } this.#refreshManagerSnapshot(); return { status: 'failed', reason: 'manager-failed' }; } } - const result = await this.#commitMutation(action); + const result = await this.#commitMutation(action, signal, cleanupTimeoutMs); if (action.kind === 'commit_import') this.discardImportPreview(action.previewId); return result; } @@ -461,62 +573,174 @@ class TuiMcpControllerImpl implements TuiMcpController { | { kind: 'test' | 'reconnect' } | { kind: 'set_publication_credential' | 'remove_publication_credential' } >, + signal?: AbortSignal, + cleanupTimeoutMs?: number, ): Promise { + let previous: McpConfigFile | undefined; + let changedIds: string[] = []; let committed: McpConfigFile; - try { - committed = await this.#deps.configStore.transform(async (current) => { - if (this.#closed) { - throw new TuiMcpMutationError({ status: 'failed', reason: 'closed' }); + let credentialRetirementStarted = false; + const transaction = this.#deps.configStore.transform(async (current) => { + if (this.#closed) { + throw new TuiMcpMutationError({ status: 'failed', reason: 'closed' }); + } + if (signal?.aborted) { + throw new TuiMcpMutationError({ + status: 'failed', + reason: 'cancelled', + }); + } + const prepared = this.#prepareMutation(current, action); + if ('status' in prepared) throw new TuiMcpMutationError(prepared); + const { next } = prepared; + previous = cloneConfig(current); + changedIds = changedServerIds(current, next); + try { + assertMcpEndpointPolicyOnChanges(current, next); + } catch { + throw new TuiMcpMutationError({ + status: 'failed', + reason: 'invalid-config', + }); + } + try { + for (const [serverId, previous] of Object.entries(current.mcpServers)) { + if (!mcpConfigChangeRetiresCredentials(previous, next.mcpServers[serverId])) continue; + // The manager owns the erase fence. Once credential storage has + // entered its write phase, let that operation settle before this + // transaction reports cancellation; an outer abort race here could + // otherwise leave a tombstone landing after rollback has begun. + await this.#deps.manager.forgetServerCredentials(serverId, previous, { + signal: credentialRetirementStarted ? undefined : signal, + onCommitStarted: () => { + credentialRetirementStarted = true; + }, + }); } - const prepared = this.#prepareMutation(current, action); - if ('status' in prepared) throw new TuiMcpMutationError(prepared); - const { next } = prepared; - try { - assertMcpEndpointPolicyOnChanges(current, next); - } catch { - throw new TuiMcpMutationError({ status: 'failed', reason: 'invalid-config' }); + } catch (error) { + if (error instanceof TuiMcpMutationError) throw error; + if (credentialRetirementStarted) { + // The erase may already be durable even when its promise rejects + // (or another endpoint in the same mutation may already have been + // retired). Keeping the old config is no longer a safe rollback. + // Commit the new config and let manager.sync retry retirement before + // it adopts or connects the new endpoint. + return next; } - try { - for (const [serverId, previous] of Object.entries(current.mcpServers)) { - if (!mcpConfigChangeRetiresCredentials(previous, next.mcpServers[serverId])) continue; - await this.#deps.manager.forgetServerCredentials(serverId, previous); - if (this.#closed) { - throw new TuiMcpMutationError({ status: 'failed', reason: 'closed' }); - } - } - } catch (error) { - if (error instanceof TuiMcpMutationError) throw error; + if (signal?.aborted && !credentialRetirementStarted) { throw new TuiMcpMutationError({ status: 'failed', - reason: 'credential-cleanup-failed', + reason: 'cancelled', }); } - return next; - }); + throw new TuiMcpMutationError({ + status: 'failed', + reason: 'credential-cleanup-failed', + }); + } + return next; + }); + try { + committed = await waitForAbort(transaction, signal); } catch (error) { if (error instanceof TuiMcpMutationError) return error.result; - return { status: 'failed', reason: 'persist-failed' }; + if (this.#closed || signal?.aborted) { + const cleanup = cleanupSignal(cleanupTimeoutMs); + try { + committed = await waitForAbort(transaction, cleanup); + } catch (settlementError) { + if (settlementError instanceof TuiMcpMutationError) return settlementError.result; + if (credentialRetirementStarted) { + // The tombstone write is now irreversible. It may already be + // durable even though the credential operation has not settled, + // so a cleanup deadline cannot hand this transaction to the late + // rollback path. Wait for the matching config mutation instead. + try { + committed = await transaction; + } catch (transactionError) { + if (transactionError instanceof TuiMcpMutationError) { + return transactionError.result; + } + this.#publicationSuppressed = false; + this.#updateSnapshot({ configuration: 'out_of_sync' }); + return { status: 'failed', reason: 'persist-failed' }; + } + } else { + this.#scheduleLateMutationRollback( + transaction, + () => previous, + () => changedIds, + ); + this.#publicationSuppressed = false; + this.#updateSnapshot({ configuration: 'out_of_sync' }); + return { status: 'failed', reason: 'rollback-failed' }; + } + } + if (!credentialRetirementStarted) { + const rolledBack = await this.#rollbackCancelledMutation( + previous, + committed, + changedIds, + cleanup, + ); + if (!rolledBack) return { status: 'failed', reason: 'rollback-failed' }; + return { + status: 'failed', + reason: this.#closed ? 'closed' : 'cancelled', + }; + } + } else { + if (credentialRetirementStarted) { + this.#publicationSuppressed = false; + this.#updateSnapshot({ configuration: 'out_of_sync' }); + this.#refreshManagerSnapshot(); + } + return { status: 'failed', reason: 'persist-failed' }; + } } - if (this.#closed) return { status: 'failed', reason: 'closed' }; this.#preparedImport = undefined; this.#config = cloneConfig(committed); this.#updateSnapshot({ configuration: 'synchronizing' }); this.#refreshManagerSnapshot(); this.#publicationSuppressed = true; try { - await this.#deps.manager.sync(committed); + throwIfAborted(signal); + await waitForAbort(this.#deps.manager.sync(committed, { signal }), signal); + throwIfAborted(signal); + this.#publicationSuppressed = false; + if (this.#closed) throw new Error('MCP controller closed'); + this.#updateSnapshot({ configuration: 'ready' }); + this.#refreshManagerSnapshot(); + const effect = await this.#settlePublication(signal); + throwIfAborted(signal); + return { status: 'applied', effect }; } catch { + if (!credentialRetirementStarted && (this.#closed || signal?.aborted)) { + const rolledBack = await this.#rollbackCancelledMutation( + previous, + committed, + changedIds, + cleanupSignal(cleanupTimeoutMs), + ); + if (!rolledBack) return { status: 'failed', reason: 'rollback-failed' }; + return { + status: 'failed', + reason: this.#closed ? 'closed' : 'cancelled', + }; + } this.#publicationSuppressed = false; this.#updateSnapshot({ configuration: 'out_of_sync' }); this.#refreshManagerSnapshot(); - await this.#settlePublication(); + if (credentialRetirementStarted && (this.#closed || signal?.aborted)) { + const cleanup = cleanupSignal(cleanupTimeoutMs); + await this.#settleCancelledConnections(changedIds, cleanup); + this.#refreshManagerSnapshot(); + await this.#settleCancelledPublication(cleanup); + } else { + await this.#settlePublication(signal); + } return { status: 'applied', effect: 'sync_failed' }; } - this.#publicationSuppressed = false; - if (this.#closed) return { status: 'failed', reason: 'closed' }; - this.#updateSnapshot({ configuration: 'ready' }); - this.#refreshManagerSnapshot(); - return { status: 'applied', effect: await this.#settlePublication() }; } #prepareMutation( @@ -563,17 +787,22 @@ class TuiMcpControllerImpl implements TuiMcpController { } try { return { - next: normalizeMcpConfig({ version: MCP_CONFIG_VERSION, mcpServers: servers }), + next: normalizeMcpConfig({ + version: MCP_CONFIG_VERSION, + mcpServers: servers, + }), }; } catch { return { status: 'failed', reason: 'invalid-config' }; } } - async #settlePublication(): Promise { + async #settlePublication(signal?: AbortSignal): Promise { + throwIfAborted(signal); this.#requestPublication(); while (!this.#closed && (this.#publicationTask || this.#publicationRequested)) { - await this.#publicationTask?.catch(() => undefined); + await waitForAbort(this.#publicationTask ?? Promise.resolve(), signal); + throwIfAborted(signal); } if ( this.#snapshot.publication === 'error' || @@ -622,6 +851,125 @@ class TuiMcpControllerImpl implements TuiMcpController { this.#notify(); } + async #rollbackCancelledMutation( + previous: McpConfigFile | undefined, + committed: McpConfigFile, + serverIds: readonly string[], + signal?: AbortSignal, + ): Promise { + if (!previous) { + this.#publicationSuppressed = false; + this.#updateSnapshot({ configuration: 'out_of_sync' }); + return false; + } + try { + const restored = await waitForAbort( + this.#deps.configStore.transform(async (current) => { + throwIfAborted(signal); + const servers = { ...current.mcpServers }; + for (const serverId of serverIds) { + if ( + configRevision(servers[serverId]) !== configRevision(committed.mcpServers[serverId]) + ) { + continue; + } + const currentEntry = servers[serverId]; + const previousEntry = previous.mcpServers[serverId]; + if (currentEntry && mcpConfigChangeRetiresCredentials(currentEntry, previousEntry)) { + await this.#deps.manager.forgetServerCredentials(serverId, currentEntry, { signal }); + throwIfAborted(signal); + } + if (previousEntry) servers[serverId] = previousEntry; + else delete servers[serverId]; + } + throwIfAborted(signal); + return normalizeMcpConfig({ + version: MCP_CONFIG_VERSION, + mcpServers: servers, + }); + }), + signal, + ); + this.#config = cloneConfig(restored); + if (this.#closed) return true; + for (const serverId of serverIds) { + if (!this.#deps.manager.disconnect) continue; + await waitForAbort(this.#deps.manager.disconnect(serverId, false, { signal }), signal); + } + await waitForAbort(this.#deps.manager.sync(restored, { signal }), signal); + this.#publicationSuppressed = false; + this.#updateSnapshot({ configuration: 'ready' }); + this.#refreshManagerSnapshot(); + await this.#settlePublication(signal); + return true; + } catch { + this.#publicationSuppressed = false; + this.#updateSnapshot({ configuration: 'out_of_sync' }); + this.#refreshManagerSnapshot(); + return false; + } + } + + #scheduleLateMutationRollback( + transaction: Promise, + previous: () => McpConfigFile | undefined, + serverIds: () => readonly string[], + ): void { + void transaction + .then((committed) => this.#rollbackCancelledMutation(previous(), committed, serverIds())) + .catch(() => undefined); + } + + async #settleCancelledConnection(serverId: string, signal?: AbortSignal): Promise { + if (!this.#deps.manager.disconnect) return true; + try { + await waitForAbort(this.#deps.manager.disconnect(serverId, false, { signal }), signal); + return true; + } catch { + return false; + } + } + + async #settleCancelledConnections( + serverIds: readonly string[], + signal?: AbortSignal, + ): Promise { + const disconnect = this.#deps.manager.disconnect; + if (!disconnect) return true; + try { + await Promise.all( + serverIds.map((serverId) => + waitForAbort(disconnect.call(this.#deps.manager, serverId, false, { signal }), signal), + ), + ); + return true; + } catch { + return false; + } + } + + async #settleCancelledPublication(signal?: AbortSignal): Promise { + try { + await this.#settlePublication(signal); + return true; + } catch { + return false; + } + } + + async #settleCancelledCredentialOperation( + operation: Promise, + signal: AbortSignal, + cleanupTimeoutMs?: number, + ): Promise<'cancelled' | 'rollback-failed'> { + try { + await waitForAbort(operation, cleanupSignal(cleanupTimeoutMs)); + return 'rollback-failed'; + } catch (error) { + return error === signal.reason ? 'cancelled' : 'rollback-failed'; + } + } + #requestPublication(): void { if (this.#closed) { this.#publicationRequested = false; @@ -646,7 +994,9 @@ class TuiMcpControllerImpl implements TuiMcpController { async #publishCurrentSnapshot(): Promise { const availability = this.#availability; if (availability.kind !== 'connected') { - this.#updateSnapshot({ publication: availability.reason ?? 'host_unavailable' }); + this.#updateSnapshot({ + publication: availability.reason ?? 'host_unavailable', + }); return; } const identity = connectionIdentity(availability); @@ -663,8 +1013,19 @@ class TuiMcpControllerImpl implements TuiMcpController { provider = this.#deps.createProvider(this.#deps.manager); if (provider) { await this.#connection.replaceClientCapabilities(provider); - } else if (this.#published?.identity === identity && this.#published.registered) { - await this.#connection.unregisterClientCapabilities(); + // The Host owns this registration as soon as replace resolves. Record + // that fact before checking whether the source snapshot is still + // current, so a coalesced empty snapshot can reliably unregister a + // replacement that became stale while the request was in flight. + this.#published = { identity, revision, registered: true }; + } else { + if (this.#published?.identity === identity && this.#published.registered) { + await this.#connection.unregisterClientCapabilities(); + } + // Even when no registration existed, remember that this revision's + // canonical Host state is empty. Otherwise settlePublication can spin: + // there is no mutation to perform, but the revision never converges. + this.#published = { identity, revision, registered: false }; } } catch { await closeProvider(provider); @@ -679,8 +1040,14 @@ class TuiMcpControllerImpl implements TuiMcpController { this.#requestPublication(); return; } - this.#published = { identity, revision, registered: provider !== undefined }; - this.#updateSnapshot({ publication: provider ? 'published' : 'not_published' }); + this.#published = { + identity, + revision, + registered: provider !== undefined, + }; + this.#updateSnapshot({ + publication: provider ? 'published' : 'not_published', + }); } #isCurrent(identity: string, revision: number): boolean { @@ -760,6 +1127,48 @@ function cloneConfig(config: McpConfigFile): McpConfigFile { return structuredClone(config); } +function changedServerIds(before: McpConfigFile, after: McpConfigFile): string[] { + return [...new Set([...Object.keys(before.mcpServers), ...Object.keys(after.mcpServers)])].filter( + (serverId) => + configRevision(before.mcpServers[serverId]) !== configRevision(after.mcpServers[serverId]), + ); +} + +function cleanupSignal(timeoutMs?: number): AbortSignal | undefined { + return timeoutMs === undefined ? undefined : AbortSignal.timeout(timeoutMs); +} + +function throwIfAborted(signal?: AbortSignal): void { + if (!signal?.aborted) return; + throw signal.reason instanceof Error + ? signal.reason + : new Error(String(signal.reason ?? 'MCP action cancelled')); +} + +function waitForAbort(promise: Promise, signal?: AbortSignal): Promise { + if (!signal) return promise; + try { + throwIfAborted(signal); + } catch (error) { + return Promise.reject(error); + } + let rejectAbort!: (reason: Error) => void; + const aborted = new Promise((_resolve, reject) => { + rejectAbort = reject; + }); + const onAbort = () => { + try { + throwIfAborted(signal); + } catch (error) { + rejectAbort(error instanceof Error ? error : new Error(String(error))); + } + }; + signal.addEventListener('abort', onAbort, { once: true }); + return Promise.race([promise, aborted]).finally(() => { + signal.removeEventListener('abort', onAbort); + }); +} + function configRevision(config: McpConfigFile | McpServerConfig | undefined): string { if (!config) return 'missing'; return createHash('sha256').update(JSON.stringify(config)).digest('hex'); diff --git a/packages/cli/src/tui-mcp-remote-publication.ts b/packages/cli/src/tui-mcp-remote-publication.ts index b0dec0b88e..760977019d 100644 --- a/packages/cli/src/tui-mcp-remote-publication.ts +++ b/packages/cli/src/tui-mcp-remote-publication.ts @@ -33,6 +33,8 @@ import { RuntimeHostRemoteCompatibilityError, runtimeHostProfileTargetFingerprint, type RemoteRuntimeHostProfile, + type RuntimeHostCapabilityProviderCredentialMutationResult, + type RuntimeHostCapabilityProviderCredentialSnapshot, type RuntimeHostCapabilityProviderCredentialStore, type RuntimeHostConnection, type RuntimeHostPeerClient, @@ -214,52 +216,117 @@ class RemoteTuiMcpPublicationTarget implements TuiMcpPublicationTarget { return () => this.#listeners.delete(listener); } - setCredential(credential: string): Promise { + setCredential( + credential: string, + options: { readonly signal?: AbortSignal } = {}, + ): Promise { this.#cancelConnect(); return this.#serialize(async () => { if (this.#closed) throw new Error('Remote MCP publication is closed'); - const committed = await this.#deps.profiles.mutateRemoteProfileIfCurrent( - this.#profileTarget(), - (profile) => - this.#deps.credentials.set( - { profile, profileIncarnationId: this.#input.profileIncarnationId }, - this.#input.ownerClientInstanceId, - credential, - ), - ); - if (!committed) { - await this.#retire('target_mismatch'); - throw new RuntimeHostProfileConnectionError( - 'target_mismatch', - 'Remote MCP publication profile is no longer current', + let target: RuntimeHostRemoteProfileIncarnation | undefined; + let previous: RuntimeHostCapabilityProviderCredentialSnapshot | undefined; + let written: RuntimeHostCapabilityProviderCredentialSnapshot | undefined; + let disconnected = false; + try { + throwIfAborted(options.signal); + const committed = await this.#deps.profiles.mutateRemoteProfileIfCurrent( + this.#profileTarget(), + async (profile) => { + throwIfAborted(options.signal); + target = { + profile, + profileIncarnationId: this.#input.profileIncarnationId, + }; + previous = await this.#readCredential(target); + throwIfAborted(options.signal); + written = await this.#writeCredential(target, previous, credential); + throwIfAborted(options.signal); + }, ); + if (!committed) { + await this.#retire('target_mismatch'); + throw new RuntimeHostProfileConnectionError( + 'target_mismatch', + 'Remote MCP publication profile is no longer current', + ); + } + throwIfAborted(options.signal); + await this.#disconnect(); + disconnected = true; + throwIfAborted(options.signal); + await this.#connect(credential, options.signal); + throwIfAborted(options.signal); + } catch (error) { + if (options.signal?.aborted && target && previous && written) { + this.#cancelConnect(); + if (disconnected) await this.#disconnect(); + const restoration = await this.#restoreCredential(target, previous, written); + if (!restoration.restored && !disconnected) { + await this.#disconnect(); + disconnected = true; + } + if (restoration.current.credential === null) { + this.#setUnavailable('credential_required'); + } else if (disconnected) { + await this.#connect(restoration.current.credential); + } + } + throw error; } - await this.#disconnect(); - await this.#connect(credential); }); } - removeCredential(): Promise { + removeCredential(options: { readonly signal?: AbortSignal } = {}): Promise { this.#cancelConnect(); return this.#serialize(async () => { if (this.#closed) throw new Error('Remote MCP publication is closed'); - await this.#disconnect(); - const committed = await this.#deps.profiles.mutateRemoteProfileIfCurrent( - this.#profileTarget(), - (profile) => - this.#deps.credentials.delete( - { profile, profileIncarnationId: this.#input.profileIncarnationId }, - this.#input.ownerClientInstanceId, - ), - ); - if (!committed) { - await this.#retire('target_mismatch'); - throw new RuntimeHostProfileConnectionError( - 'target_mismatch', - 'Remote MCP publication profile is no longer current', + let target: RuntimeHostRemoteProfileIncarnation | undefined; + let previous: RuntimeHostCapabilityProviderCredentialSnapshot | undefined; + let deleted: RuntimeHostCapabilityProviderCredentialSnapshot | undefined; + let disconnected = false; + try { + throwIfAborted(options.signal); + await this.#disconnect(); + disconnected = true; + throwIfAborted(options.signal); + const committed = await this.#deps.profiles.mutateRemoteProfileIfCurrent( + this.#profileTarget(), + async (profile) => { + throwIfAborted(options.signal); + target = { + profile, + profileIncarnationId: this.#input.profileIncarnationId, + }; + previous = await this.#readCredential(target); + throwIfAborted(options.signal); + deleted = await this.#deleteCredential(target, previous); + throwIfAborted(options.signal); + }, ); + if (!committed) { + await this.#retire('target_mismatch'); + throw new RuntimeHostProfileConnectionError( + 'target_mismatch', + 'Remote MCP publication profile is no longer current', + ); + } + throwIfAborted(options.signal); + this.#setUnavailable('credential_required'); + } catch (error) { + if (options.signal?.aborted && disconnected) { + const rollbackTarget = target ?? this.#profileTarget(); + const prior = + previous === undefined ? await this.#readCredential(rollbackTarget) : previous; + const restoration = + deleted !== undefined + ? await this.#restoreCredential(rollbackTarget, prior, deleted) + : { restored: true, current: prior }; + if (restoration.current.credential !== null) { + await this.#connect(restoration.current.credential); + } else this.#setUnavailable('credential_required'); + } + throw error; } - this.#setUnavailable('credential_required'); }); } @@ -277,12 +344,15 @@ class RemoteTuiMcpPublicationTarget implements TuiMcpPublicationTarget { return pending; } - async #connect(credential: string): Promise { + async #connect(credential: string, signal?: AbortSignal): Promise { const generation = ++this.#generation; const abort = new AbortController(); this.#connectAbort = abort; + const forwardAbort = () => abort.abort(abortReason(signal)); + signal?.addEventListener('abort', forwardAbort, { once: true }); this.#setUnavailable('host_unavailable'); try { + throwIfAborted(signal); const clientInstanceId = await this.#deps.loadClientInstanceId( providerIdentityPath(this.#input), ); @@ -343,10 +413,112 @@ class RemoteTuiMcpPublicationTarget implements TuiMcpPublicationTarget { } await this.#closePeer(this.#peerClient); } finally { + signal?.removeEventListener('abort', forwardAbort); if (this.#connectAbort === abort) this.#connectAbort = undefined; } } + async #readCredential( + target: RuntimeHostRemoteProfileIncarnation, + ): Promise { + if (this.#deps.credentials.read) { + return this.#deps.credentials.read(target, this.#input.ownerClientInstanceId); + } + return { + credential: await this.#deps.credentials.get(target, this.#input.ownerClientInstanceId), + revision: null, + }; + } + + async #writeCredential( + target: RuntimeHostRemoteProfileIncarnation, + previous: RuntimeHostCapabilityProviderCredentialSnapshot, + credential: string, + ): Promise { + if (this.#deps.credentials.compareAndSet && this.#deps.credentials.read) { + const result = await this.#deps.credentials.compareAndSet( + target, + this.#input.ownerClientInstanceId, + previous.revision, + credential, + ); + if (!result.committed) { + throw new Error('Runtime Host capability-provider credential changed during update'); + } + return { credential, revision: result.revision }; + } + await this.#deps.credentials.set(target, this.#input.ownerClientInstanceId, credential); + return { credential, revision: null }; + } + + async #deleteCredential( + target: RuntimeHostRemoteProfileIncarnation, + previous: RuntimeHostCapabilityProviderCredentialSnapshot, + ): Promise { + if (previous.credential === null) return previous; + if (this.#deps.credentials.compareAndSet && this.#deps.credentials.read) { + const result = await this.#deps.credentials.compareAndSet( + target, + this.#input.ownerClientInstanceId, + previous.revision, + null, + ); + if (!result.committed) { + throw new Error('Runtime Host capability-provider credential changed during deletion'); + } + return { credential: null, revision: result.revision }; + } + await this.#deps.credentials.delete(target, this.#input.ownerClientInstanceId); + return { credential: null, revision: null }; + } + + async #restoreCredential( + target: RuntimeHostRemoteProfileIncarnation, + previous: RuntimeHostCapabilityProviderCredentialSnapshot, + written: RuntimeHostCapabilityProviderCredentialSnapshot, + ): Promise<{ + readonly restored: boolean; + readonly current: RuntimeHostCapabilityProviderCredentialSnapshot; + }> { + if ( + this.#deps.credentials.compareAndSet && + this.#deps.credentials.read && + written.revision !== null + ) { + const result: RuntimeHostCapabilityProviderCredentialMutationResult = + await this.#deps.credentials.compareAndSet( + target, + this.#input.ownerClientInstanceId, + written.revision, + previous.credential, + ); + if (result.committed) { + return { + restored: true, + current: { credential: previous.credential, revision: result.revision }, + }; + } + return { restored: false, current: result.current }; + } + const current = await this.#deps.credentials.get(target, this.#input.ownerClientInstanceId); + if (current !== written.credential) { + return { + restored: false, + current: { credential: current, revision: null }, + }; + } + if (previous.credential === null) { + await this.#deps.credentials.delete(target, this.#input.ownerClientInstanceId); + } else { + await this.#deps.credentials.set( + target, + this.#input.ownerClientInstanceId, + previous.credential, + ); + } + return { restored: true, current: previous }; + } + async #disconnect(): Promise { this.#cancelConnect(); this.#generation += 1; @@ -536,3 +708,13 @@ function classifyUnavailable(error: unknown): TuiMcpPublicationUnavailableReason if (error instanceof RuntimeHostRemoteCompatibilityError) return 'target_mismatch'; return 'host_unavailable'; } + +function abortReason(signal?: AbortSignal): Error { + return signal?.reason instanceof Error + ? signal.reason + : new Error(String(signal?.reason ?? 'Remote MCP publication cancelled')); +} + +function throwIfAborted(signal?: AbortSignal): void { + if (signal?.aborted) throw abortReason(signal); +} diff --git a/packages/mcp/src/__fixtures__/stdio-fixture-events.ts b/packages/mcp/src/__fixtures__/stdio-fixture-events.ts index de46e80d6a..ce058a09c6 100644 --- a/packages/mcp/src/__fixtures__/stdio-fixture-events.ts +++ b/packages/mcp/src/__fixtures__/stdio-fixture-events.ts @@ -43,9 +43,10 @@ export function installStdioFixtureEvents( fixtureEnv: process.env.MAKA_MCP_STDIO_FIXTURE_VALUE ?? null, }); process.stderr.write(`stdio fixture ${fixture} pid=${process.pid}\n`); + if (process.argv.includes('--hold-stdin-open')) setInterval(() => undefined, 30_000); process.once('SIGTERM', () => { record('signal', { signal: 'SIGTERM' }); - process.exit(0); + if (!process.argv.includes('--ignore-sigterm')) process.exit(0); }); process.once('exit', (code) => record('exit', { code })); return record; diff --git a/packages/mcp/src/__fixtures__/stdio-server.ts b/packages/mcp/src/__fixtures__/stdio-server.ts index 62a1af6608..97c1bb9126 100644 --- a/packages/mcp/src/__fixtures__/stdio-server.ts +++ b/packages/mcp/src/__fixtures__/stdio-server.ts @@ -22,7 +22,7 @@ import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js' import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js'; import { installStdioFixtureEvents } from './stdio-fixture-events.js'; -installStdioFixtureEvents('legacy'); +const recordFixtureEvent = installStdioFixtureEvents('legacy'); if (process.argv.includes('--crash')) { process.stderr.write('fixture startup failed: deliberate diagnostic\n'); @@ -57,6 +57,10 @@ const server = new Server( ); server.setRequestHandler(ListToolsRequestSchema, async ({ params }) => { + if (process.argv.includes('--slow-tool-list')) { + recordFixtureEvent('tools-list'); + await new Promise((resolve) => setTimeout(resolve, 30_000)); + } if (process.argv.includes('--schema-annotations')) { return { tools: [ diff --git a/packages/mcp/src/__tests__/credential-coordinator.test.ts b/packages/mcp/src/__tests__/credential-coordinator.test.ts index ce6dc6c530..65dea58fe4 100644 --- a/packages/mcp/src/__tests__/credential-coordinator.test.ts +++ b/packages/mcp/src/__tests__/credential-coordinator.test.ts @@ -165,4 +165,47 @@ describe('McpCredentialCoordinator', () => { await assert.rejects(coordinator.erase('remote', { signal: aborted.signal }), /abandoned/u); assert.equal(writes.length, 0); }); + + test('an erase ignores cancellation after its irreversible commit boundary', async () => { + let releaseWrite!: () => void; + const writeGate = new Promise((resolve) => { + releaseWrite = resolve; + }); + let writeStarted!: () => void; + const started = new Promise((resolve) => { + writeStarted = resolve; + }); + let stored: McpOAuthRecord | undefined = { + version: 1, + generation: 0, + tokens: { access_token: 'old-token', token_type: 'Bearer' }, + }; + const storage: McpOAuthStorage = { + get: async () => stored, + set: async (_id, record) => { + writeStarted(); + await writeGate; + stored = record; + }, + delete: async () => {}, + }; + const coordinator = new McpCredentialCoordinator(storage); + const round = new AbortController(); + let boundaryCalls = 0; + + const erasing = coordinator.erase('remote', { + signal: round.signal, + onCommitStarted: () => { + boundaryCalls += 1; + round.abort(new Error('too late to cancel erase')); + }, + }); + await started; + assert.equal(round.signal.aborted, true); + assert.equal(boundaryCalls, 1); + + releaseWrite(); + await erasing; + assert.deepEqual(stored, { version: 2, generation: 1 }); + }); }); diff --git a/packages/mcp/src/__tests__/manager.test.ts b/packages/mcp/src/__tests__/manager.test.ts index 68c79926b0..cf67c13fb8 100644 --- a/packages/mcp/src/__tests__/manager.test.ts +++ b/packages/mcp/src/__tests__/manager.test.ts @@ -18,7 +18,10 @@ */ import assert from 'node:assert/strict'; +import { mkdtemp, readFile, rm } from 'node:fs/promises'; import { createServer, type IncomingMessage, type ServerResponse } from 'node:http'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; import { fileURLToPath } from 'node:url'; import { afterEach, describe, test } from 'node:test'; import { NodeStreamableHTTPServerTransport } from '@modelcontextprotocol/node'; @@ -30,7 +33,13 @@ import { type McpProtocolPreference, type McpToolBinding, } from '@maka/core/mcp'; -import { buildStdioEnvironment, McpClientManager, McpToolCallError } from '../index.js'; +import { deferred } from '@maka/core/test-only/async-primitives'; +import { + buildStdioEnvironment, + McpClientManager, + McpToolCallError, + type McpOAuthStorage, +} from '../index.js'; import { waitFor as pollFor } from '@maka/core/test-only/async-primitives'; const fixturePath = fileURLToPath(new URL('../__fixtures__/stdio-server.js', import.meta.url)); @@ -72,6 +81,88 @@ describe('McpClientManager E2E', { concurrency: false }, () => { assertLegacyHandshake(fixture); }); + test('caller abort stops waiting on the initial OAuth credential read', async () => { + const readStarted = deferred(); + const releaseRead = deferred(); + let reads = 0; + const storage: McpOAuthStorage = { + get: async () => { + reads += 1; + if (reads === 1) { + readStarted.resolve(); + await releaseRead.promise; + } + return undefined; + }, + set: async () => undefined, + delete: async () => undefined, + }; + const manager = createManager({ oauthStorage: storage }); + const abort = new AbortController(); + const syncing = manager.sync( + { + version: MCP_CONFIG_VERSION, + mcpServers: { + remote: { + url: 'https://oauth-read.example/mcp', + transport: 'streamable-http', + }, + }, + }, + { signal: abort.signal }, + ); + await readStarted.promise; + + abort.abort(new Error('cancel initial OAuth credential read')); + const rejected = await rejectsWithin(syncing, 100, /cancel initial OAuth credential read/u); + releaseRead.resolve(); + await syncing.catch(() => undefined); + + assert.equal(rejected, true); + assert.equal(manager.status('remote')?.state, 'disconnected'); + }); + + test('caller abort stops waiting on the OAuth authorization-owner read', async () => { + const readStarted = deferred(); + const releaseRead = deferred(); + let reads = 0; + const storage: McpOAuthStorage = { + get: async () => { + reads += 1; + if (reads === 2) { + readStarted.resolve(); + await releaseRead.promise; + } + return undefined; + }, + set: async () => undefined, + delete: async () => undefined, + }; + const manager = createManager({ oauthStorage: storage }); + const abort = new AbortController(); + const syncing = manager.sync( + { + version: MCP_CONFIG_VERSION, + mcpServers: { + remote: { + url: 'https://oauth-owner.example/mcp', + transport: 'streamable-http', + }, + }, + }, + { signal: abort.signal }, + ); + await readStarted.promise; + + abort.abort(new Error('cancel OAuth authorization-owner read')); + const rejected = await rejectsWithin(syncing, 100, /cancel OAuth authorization-owner read/u); + releaseRead.resolve(); + await syncing.catch(() => undefined); + + assert.equal(rejected, true); + assert.equal(manager.status('remote')?.state, 'disconnected'); + }); + test('auto probes before negotiating a legacy Streamable HTTP server', async () => { const fixture = await createRemoteFixture('streamable-http'); const manager = createManager(); @@ -106,7 +197,10 @@ describe('McpClientManager E2E', { concurrency: false }, () => { test('strips configured headers when a redirect leaves the endpoint origin', async () => { // Undici forwards custom headers (X-API-Key) across cross-origin // redirects; the manager's scoped fetch must not. - const crossOriginSeen: Array<{ authorization?: string; apiKey?: string }> = []; + const crossOriginSeen: Array<{ + authorization?: string; + apiKey?: string; + }> = []; const target = createServer((req, res) => { crossOriginSeen.push({ ...(typeof req.headers.authorization === 'string' @@ -126,7 +220,9 @@ describe('McpClientManager E2E', { concurrency: false }, () => { if (!targetAddress || typeof targetAddress === 'string') throw new Error('no target port'); const redirector = createServer((req, res) => { res - .writeHead(307, { location: `http://127.0.0.1:${targetAddress.port}${req.url ?? '/'}` }) + .writeHead(307, { + location: `http://127.0.0.1:${targetAddress.port}${req.url ?? '/'}`, + }) .end(); }); await new Promise((resolve, reject) => { @@ -145,7 +241,10 @@ describe('McpClientManager E2E', { concurrency: false }, () => { remote: { url: `http://127.0.0.1:${redirectorAddress.port}/mcp`, transport: 'streamable-http', - headers: { Authorization: 'Bearer remote-test', 'X-API-Key': 'key-123456' }, + headers: { + Authorization: 'Bearer remote-test', + 'X-API-Key': 'key-123456', + }, }, }, }); @@ -320,7 +419,9 @@ describe('McpClientManager E2E', { concurrency: false }, () => { await waitFor(() => manager.status('remote')?.error?.includes('duplicate tool') === true); assert.equal(manager.toolSnapshot(), replacement); assert.deepEqual( - await manager.callTool(replacement.tools[0]!.binding, { value: 'retained' }), + await manager.callTool(replacement.tools[0]!.binding, { + value: 'retained', + }), { content: [{ type: 'text', text: 'retained' }], structuredContent: undefined, @@ -329,7 +430,9 @@ describe('McpClientManager E2E', { concurrency: false }, () => { }); test('refreshes for legacy list-changed notifications without an advertised flag', async () => { - const fixture = await createRemoteFixture('sse', { advertiseToolListChanges: false }); + const fixture = await createRemoteFixture('sse', { + advertiseToolListChanges: false, + }); const manager = createManager(); await manager.sync(remoteConfig(`${fixture.url}/sse`, 'auto', 'legacy')); @@ -792,7 +895,10 @@ describe('McpClientManager E2E', { concurrency: false }, () => { await fixture.notifyToolListChanged(); await gate.started; - const removal = manager.sync({ version: MCP_CONFIG_VERSION, mcpServers: {} }); + const removal = manager.sync({ + version: MCP_CONFIG_VERSION, + mcpServers: {}, + }); let removedBeforeRelease: boolean; try { removedBeforeRelease = await settlesWithin(removal, 1_000); @@ -1207,6 +1313,248 @@ describe('McpClientManager E2E', { concurrency: false }, () => { assert.deepEqual(manager.toolSnapshot().tools, []); }); + test('caller abort cancels sync and reaps its in-flight stdio child', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'maka-mcp-sync-abort-')); + const eventLog = join(root, 'events.jsonl'); + t.after(() => rm(root, { recursive: true, force: true })); + const manager = createManager(); + const abort = new AbortController(); + const sync = manager.sync( + { + version: MCP_CONFIG_VERSION, + mcpServers: { + fixture: { + command: process.execPath, + args: [fixturePath, '--slow-start', '--ignore-sigterm'], + env: { MAKA_MCP_STDIO_EVENT_LOG: eventLog }, + }, + }, + }, + { signal: abort.signal }, + ); + await pollFor( + async () => { + try { + return (await readFile(eventLog, 'utf8')).includes('"event":"start"'); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') return false; + throw error; + } + }, + { timeoutMs: 1_000, pollMs: 5 }, + ); + const events = await readFile(eventLog, 'utf8'); + const start = events + .split('\n') + .filter(Boolean) + .map((line) => JSON.parse(line) as { event: string; pid: number }) + .find((event) => event.event === 'start'); + assert.ok(start); + abort.abort(new Error('cancelled by caller')); + await assert.rejects(sync, /cancelled by caller/u); + assert.equal(processExists(start.pid), false); + await pollFor(() => manager.status('fixture')?.state === 'disconnected', { + timeoutMs: 2_000, + pollMs: 5, + }); + assert.equal(manager.status('fixture')?.state, 'disconnected'); + assert.deepEqual(manager.toolSnapshot().tools, []); + }); + + test('a repeated disconnect joins the original stdio teardown before reconnecting', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'maka-mcp-repeat-disconnect-')); + const eventLog = join(root, 'events.jsonl'); + t.after(() => rm(root, { recursive: true, force: true })); + const manager = createManager(); + const config = fixtureConfig(['--ignore-sigterm', '--hold-stdin-open']); + config.mcpServers.fixture = { + ...config.mcpServers.fixture, + env: { MAKA_MCP_STDIO_EVENT_LOG: eventLog }, + protocol: 'legacy', + }; + await manager.sync(config); + const events = await readFile(eventLog, 'utf8'); + const start = events + .split('\n') + .filter(Boolean) + .map((line) => JSON.parse(line) as { event: string; pid: number }) + .find((event) => event.event === 'start'); + assert.ok(start); + + const abort = new AbortController(); + const firstDisconnect = manager.disconnect('fixture', false, { signal: abort.signal }); + await pollFor(() => manager.status('fixture')?.state === 'disconnected', { + timeoutMs: 1_000, + pollMs: 5, + }); + assert.equal(processExists(start.pid), true); + abort.abort(new Error('cancel first disconnect wait')); + await assert.rejects(firstDisconnect, /cancel first disconnect wait/u); + const secondDisconnect = manager.disconnect('fixture'); + + assert.equal(await settlesWithin(secondDisconnect, 50), false); + process.kill(start.pid, 'SIGKILL'); + await secondDisconnect; + assert.equal(processExists(start.pid), false); + + await manager.reconnect('fixture'); + const starts = (await readFile(eventLog, 'utf8')) + .split('\n') + .filter(Boolean) + .map((line) => JSON.parse(line) as { event: string; pid: number }) + .filter((event) => event.event === 'start'); + assert.equal(starts.length, 2); + assert.notEqual(starts[1]?.pid, start.pid); + const replacement = starts[1]; + assert.ok(replacement); + process.kill(replacement.pid, 'SIGKILL'); + await pollFor(() => !processExists(replacement.pid), { + timeoutMs: 1_000, + pollMs: 5, + }); + }); + + test('a joining caller aborts only its wait for a shared in-flight connect', async (t) => { + const root = await mkdtemp(join(tmpdir(), 'maka-mcp-joined-connect-abort-')); + const eventLog = join(root, 'events.jsonl'); + t.after(() => rm(root, { recursive: true, force: true })); + const manager = createManager(); + const initialSync = manager.sync({ + version: MCP_CONFIG_VERSION, + mcpServers: { + fixture: { + command: process.execPath, + args: [fixturePath, '--slow-start'], + env: { MAKA_MCP_STDIO_EVENT_LOG: eventLog }, + }, + }, + }); + await pollFor( + async () => { + try { + return (await readFile(eventLog, 'utf8')).includes('"event":"start"'); + } catch (error) { + if ((error as NodeJS.ErrnoException).code === 'ENOENT') return false; + throw error; + } + }, + { timeoutMs: 1_000, pollMs: 5 }, + ); + const abort = new AbortController(); + const joined = manager.connect('fixture', { signal: abort.signal }); + + abort.abort(new Error('joining caller cancelled')); + await assert.rejects(joined, /joining caller cancelled/u); + assert.equal(await settlesWithin(initialSync, 50), false); + const eventsBeforeOwnerCancellation = await readFile(eventLog, 'utf8'); + const start = eventsBeforeOwnerCancellation + .split('\n') + .filter(Boolean) + .map((line) => JSON.parse(line) as { event: string; pid: number }) + .find((event) => event.event === 'start'); + assert.ok(start); + assert.equal(processExists(start.pid), true); + + assert.equal(manager.cancelConnect('fixture'), true); + await initialSync; + await pollFor(() => manager.status('fixture')?.state === 'disconnected', { + timeoutMs: 2_000, + pollMs: 5, + }); + await pollFor(() => !processExists(start.pid), { + timeoutMs: 5_000, + pollMs: 5, + }); + assert.deepEqual(manager.toolSnapshot().tools, []); + }); + + test('caller abort stops waiting on a credential read before sync mutates connections', async () => { + let releaseGet!: () => void; + const getGate = new Promise((resolve) => { + releaseGet = resolve; + }); + const storage: McpOAuthStorage = { + get: async () => { + await getGate; + return undefined; + }, + set: async () => undefined, + delete: async () => undefined, + }; + const manager = createManager({ oauthStorage: storage }); + await manager.sync({ + version: MCP_CONFIG_VERSION, + mcpServers: { fixture: { command: process.execPath, enabled: false } }, + }); + const abort = new AbortController(); + const removing = manager.sync( + { version: MCP_CONFIG_VERSION, mcpServers: {} }, + { signal: abort.signal }, + ); + await new Promise((resolve) => setImmediate(resolve)); + abort.abort(new Error('credential read cancelled')); + + assert.equal(await rejectsWithin(removing, 100, /credential read cancelled/u), true); + assert.equal(manager.status('fixture')?.state, 'disabled'); + releaseGet(); + }); + + test('sync waits for every started removal before it settles after a sibling aborts', async () => { + const fastRead = deferred(); + const slowWrite = deferred(); + let fastReadStarted = false; + let slowWriteStarted = false; + let removalActive = false; + let slowWrites = 0; + const storage: McpOAuthStorage = { + get: async (serverId) => { + if (serverId === 'fast' && removalActive) { + fastReadStarted = true; + await fastRead.promise; + } + return undefined; + }, + set: async (serverId) => { + if (serverId === 'slow' && removalActive) { + slowWrites += 1; + slowWriteStarted = true; + await slowWrite.promise; + } + }, + delete: async () => undefined, + }; + const manager = createManager({ oauthStorage: storage }); + await manager.sync({ + version: MCP_CONFIG_VERSION, + mcpServers: { + fast: { command: process.execPath, enabled: false }, + slow: { url: 'https://slow.example/mcp', enabled: false }, + }, + }); + const abort = new AbortController(); + removalActive = true; + const removing = manager.sync( + { version: MCP_CONFIG_VERSION, mcpServers: {} }, + { signal: abort.signal }, + ); + const observedRemoval = removing.then( + () => ({ status: 'fulfilled' as const }), + (error: unknown) => ({ status: 'rejected' as const, error }), + ); + await waitFor(() => fastReadStarted && slowWriteStarted); + abort.abort(new Error('cancel parallel removals')); + + assert.equal(await settlesWithin(observedRemoval, 50), false); + assert.equal(slowWrites, 1); + slowWrite.resolve(); + fastRead.resolve(); + const result = await observedRemoval; + assert.equal(result.status, 'rejected'); + assert.match(String('error' in result ? result.error : ''), /cancel parallel removals/u); + assert.equal(manager.status('fast')?.state, 'disabled'); + assert.equal(manager.status('slow'), undefined); + }); + test('cancels installation after remote tool discovery starts', async () => { const fixture = await createRemoteFixture('streamable-http'); const manager = createManager(); @@ -1377,6 +1725,15 @@ function fixtureConfig(extraArgs: string[] = []): McpConfigFile { }; } +function processExists(pid: number): boolean { + try { + process.kill(pid, 0); + return true; + } catch (error) { + return (error as NodeJS.ErrnoException).code === 'EPERM'; + } +} + function remoteConfig( url: string, transport: 'auto' | 'streamable-http' = 'streamable-http', @@ -1396,7 +1753,11 @@ function remoteConfig( } async function waitFor(predicate: () => boolean, timeoutMs = 1_000): Promise { - await pollFor(predicate, { timeoutMs, pollMs: 5, message: 'condition was not reached' }); + await pollFor(predicate, { + timeoutMs, + pollMs: 5, + message: 'condition was not reached', + }); } async function settlesWithin(promise: Promise, timeoutMs: number): Promise { @@ -1413,6 +1774,27 @@ async function settlesWithin(promise: Promise, timeoutMs: number): Prom } } +async function rejectsWithin( + promise: Promise, + timeoutMs: number, + pattern: RegExp, +): Promise { + let timer: NodeJS.Timeout | undefined; + try { + return await Promise.race([ + promise.then( + () => false, + (error: unknown) => pattern.test(String(error)), + ), + new Promise((resolve) => { + timer = setTimeout(() => resolve(false), timeoutMs); + }), + ]); + } finally { + if (timer) clearTimeout(timer); + } +} + interface RemoteRequest { method: string; path: string; @@ -1646,7 +2028,9 @@ function createProtocolServer(options: { { name: 'maka-remote-fixture', version: '1.0.0' }, { capabilities: options.advertiseTools - ? { tools: options.advertiseToolListChanges === false ? {} : { listChanged: true } } + ? { + tools: options.advertiseToolListChanges === false ? {} : { listChanged: true }, + } : {}, }, ); @@ -1682,7 +2066,10 @@ function createProtocolServer(options: { { name: 'invalid-schema', inputSchema: { type: 'object' as const }, - outputSchema: { type: 'string' as const, pattern: '[' }, + outputSchema: { + type: 'string' as const, + pattern: '[', + }, }, ] : [remoteToolDefinition('echo'), remoteToolDefinition('invalid-output')], @@ -1702,7 +2089,9 @@ function createProtocolServer(options: { const args = params.arguments ?? {}; if (params.name === 'invalid-output') { if (args.mode === 'missing') { - return { content: [{ type: 'text', text: 'missing structured output' }] }; + return { + content: [{ type: 'text', text: 'missing structured output' }], + }; } if (args.mode === 'is-error') { return { @@ -1719,7 +2108,10 @@ function createProtocolServer(options: { if (args.mode === 'too-many-error-blocks') { return { isError: true, - content: Array.from({ length: 101 }, () => ({ type: 'text' as const, text: 'x' })), + content: Array.from({ length: 101 }, () => ({ + type: 'text' as const, + text: 'x', + })), }; } return { diff --git a/packages/mcp/src/__tests__/oauth.test.ts b/packages/mcp/src/__tests__/oauth.test.ts index c25b265b65..cc9715c0fe 100644 --- a/packages/mcp/src/__tests__/oauth.test.ts +++ b/packages/mcp/src/__tests__/oauth.test.ts @@ -29,6 +29,7 @@ import { ListToolsRequestSchema, } from '@modelcontextprotocol/sdk/types.js'; import { MCP_CONFIG_VERSION, type McpConfigFile } from '@maka/core/mcp'; +import { deferred } from '@maka/core/test-only/async-primitives'; import { createMemoryMcpOAuthStorage, McpClientManager, @@ -67,7 +68,9 @@ describe('McpClientManager OAuth E2E', () => { assert.ok((await storage.get('remote'))?.discovery?.authorizationServerUrl); const redirectUrl = 'http://127.0.0.1:39999/callback'; - const start = await manager.startAuthorization('remote', redirectUrl, { state: 'maka-state' }); + const start = await manager.startAuthorization('remote', redirectUrl, { + state: 'maka-state', + }); assert.equal(start.status, 'redirect'); if (start.status !== 'redirect') return; const authorizationUrl = new URL(start.authorizationUrl); @@ -104,11 +107,16 @@ describe('McpClientManager OAuth E2E', () => { const code = location.searchParams.get('code'); assert.ok(code); - const status = await manager.finishAuthorization('remote', { code, state: 'maka-state' }); + const status = await manager.finishAuthorization('remote', { + code, + state: 'maka-state', + }); assert.equal(status.state, 'connected'); assert.equal(status.authenticated, true); assert.deepEqual( - await manager.callTool(bindingFor(manager, 'remote', 'echo'), { value: 'authorized' }), + await manager.callTool(bindingFor(manager, 'remote', 'echo'), { + value: 'authorized', + }), { content: [{ type: 'text', text: 'authorized' }], structuredContent: undefined, @@ -180,7 +188,9 @@ describe('McpClientManager OAuth E2E', () => { fixture.rotateAccessToken(); await assert.rejects( - manager.callTool(bindingFor(manager, 'remote', 'echo'), { value: 'revoked' }), + manager.callTool(bindingFor(manager, 'remote', 'echo'), { + value: 'revoked', + }), ); assert.equal(manager.status('remote')?.state, 'needs-auth'); }); @@ -363,7 +373,11 @@ describe('McpClientManager OAuth E2E', () => { const storage = createMemoryMcpOAuthStorage(); await storage.set('remote', { serverUrl: fixture.mcpUrl, - tokens: { access_token: fixture.accessToken, token_type: 'Bearer', id_token: idToken }, + tokens: { + access_token: fixture.accessToken, + token_type: 'Bearer', + id_token: idToken, + }, }); const manager = new McpClientManager({ oauthStorage: storage }); managers.push(manager); @@ -378,7 +392,9 @@ describe('McpClientManager OAuth E2E', () => { const fixture = await createOAuthFixture({ mcpFailureBody: () => 'upstream rejected credential k7#', }); - const manager = new McpClientManager({ oauthStorage: createMemoryMcpOAuthStorage() }); + const manager = new McpClientManager({ + oauthStorage: createMemoryMcpOAuthStorage(), + }); managers.push(manager); await manager.sync({ version: MCP_CONFIG_VERSION, @@ -438,12 +454,50 @@ describe('McpClientManager OAuth E2E', () => { assert.equal(manager.status('remote')?.state, 'needs-auth'); }); + test('clearAuthorization aborts an in-flight reconnect instead of reporting success', async () => { + const reconnectStarted = deferred(); + let holdReconnect = false; + const fixture = await createOAuthFixture({ + holdMcpRequest: async () => { + if (!holdReconnect) return false; + reconnectStarted.resolve(); + return true; + }, + }); + const storage = createMemoryMcpOAuthStorage(); + await storage.set('remote', { + serverUrl: fixture.mcpUrl, + tokens: { access_token: fixture.accessToken, token_type: 'Bearer' }, + }); + const manager = new McpClientManager({ + oauthStorage: storage, + timeouts: { remoteConnectMs: 300 }, + }); + managers.push(manager); + await manager.sync(config(fixture.mcpUrl)); + assert.equal(manager.status('remote')?.state, 'connected'); + + holdReconnect = true; + const abort = new AbortController(); + const clearing = manager.clearAuthorization('remote', { + signal: abort.signal, + }); + await reconnectStarted.promise; + abort.abort(new Error('cancelled authorization reconnect')); + + await assert.rejects(clearing, /cancelled authorization reconnect/u); + assert.equal((await storage.get('remote'))?.tokens, undefined); + assert.equal(manager.status('remote')?.state, 'disconnected'); + }); + test('the probe speaks the current protocol version to strict POST-only servers', async () => { const fixture = await createOAuthFixture({ challengeOnPostOnly: true, requireProtocolVersion: LATEST_PROTOCOL_VERSION, }); - const manager = new McpClientManager({ oauthStorage: createMemoryMcpOAuthStorage() }); + const manager = new McpClientManager({ + oauthStorage: createMemoryMcpOAuthStorage(), + }); managers.push(manager); await manager.sync(config(fixture.mcpUrl)); assert.equal(manager.status('remote')?.state, 'needs-auth'); @@ -458,7 +512,9 @@ describe('McpClientManager OAuth E2E', () => { test('a bare 401 on GET does not stop the probe from asking via POST', async () => { const fixture = await createOAuthFixture({ bareChallengeOnGet: true }); - const manager = new McpClientManager({ oauthStorage: createMemoryMcpOAuthStorage() }); + const manager = new McpClientManager({ + oauthStorage: createMemoryMcpOAuthStorage(), + }); managers.push(manager); await manager.sync(config(fixture.mcpUrl)); assert.equal(manager.status('remote')?.state, 'needs-auth'); @@ -504,7 +560,9 @@ describe('McpClientManager OAuth E2E', () => { }); test('a token endpoint reflecting the PKCE verifier does not leak it', async () => { - const fixture = await createOAuthFixture({ reflectVerifierInTokenError: true }); + const fixture = await createOAuthFixture({ + reflectVerifierInTokenError: true, + }); const storage = createMemoryMcpOAuthStorage(); const manager = new McpClientManager({ oauthStorage: storage }); managers.push(manager); @@ -625,7 +683,9 @@ describe('McpClientManager OAuth E2E', () => { test('a token endpoint reflecting the authorization code does not leak it', async () => { const fixture = await createOAuthFixture({ reflectCodeInTokenError: true }); - const manager = new McpClientManager({ oauthStorage: createMemoryMcpOAuthStorage() }); + const manager = new McpClientManager({ + oauthStorage: createMemoryMcpOAuthStorage(), + }); managers.push(manager); await manager.sync(config(fixture.mcpUrl)); @@ -639,7 +699,10 @@ describe('McpClientManager OAuth E2E', () => { assert.ok(code); await assert.rejects( - manager.finishAuthorization('remote', { code, state: 'code-reflect-state' }), + manager.finishAuthorization('remote', { + code, + state: 'code-reflect-state', + }), (error: unknown) => { assert.ok(error instanceof Error); assert.ok(!error.message.includes(code)); @@ -700,7 +763,11 @@ describe('McpClientManager OAuth E2E', () => { assert.ok(iss); // The genuine issuer passes... - const status = await manager.finishAuthorization('remote', { code, iss, state: 'iss-state' }); + const status = await manager.finishAuthorization('remote', { + code, + iss, + state: 'iss-state', + }); assert.equal(status.state, 'connected'); }); @@ -746,7 +813,10 @@ describe('McpClientManager OAuth E2E', () => { const consent = await fetch(start.authorizationUrl, { redirect: 'manual' }); const code = new URL(consent.headers.get('location') ?? '').searchParams.get('code'); assert.ok(code); - await manager.finishAuthorization('remote', { code, state: 'version-state' }); + await manager.finishAuthorization('remote', { + code, + state: 'version-state', + }); const record = await storage.get('remote'); assert.ok(record?.tokens); @@ -788,7 +858,10 @@ describe('McpClientManager OAuth E2E', () => { // The exchange must refuse to overwrite the externally changed record — // and the conflict proves CAS survives the manager's storage wrapper. await assert.rejects( - manager.finishAuthorization('remote', { code, state: 'cas-conflict-state' }), + manager.finishAuthorization('remote', { + code, + state: 'cas-conflict-state', + }), /outside/u, ); assert.ok(tokenWriteConflicts >= 1); @@ -871,7 +944,10 @@ describe('McpClientManager OAuth E2E', () => { manager.sync({ version: MCP_CONFIG_VERSION, mcpServers: { - remote: { url: 'https://changed.example/mcp', transport: 'streamable-http' }, + remote: { + url: 'https://changed.example/mcp', + transport: 'streamable-http', + }, }, }), /credential store unavailable/u, @@ -1034,7 +1110,9 @@ describe('McpClientManager OAuth E2E', () => { }); assert.equal(first.status, 'redirect'); if (first.status !== 'redirect') return; - const firstConsent = await fetch(first.authorizationUrl, { redirect: 'manual' }); + const firstConsent = await fetch(first.authorizationUrl, { + redirect: 'manual', + }); const firstCode = new URL(firstConsent.headers.get('location') ?? '').searchParams.get('code'); assert.ok(firstCode); @@ -1047,12 +1125,17 @@ describe('McpClientManager OAuth E2E', () => { assert.equal(second.status, 'redirect'); if (second.status !== 'redirect') return; await assert.rejects( - manager.finishAuthorization('remote', { code: firstCode, state: 'round-one' }), + manager.finishAuthorization('remote', { + code: firstCode, + state: 'round-one', + }), /superseded/u, ); // Round 2 completes normally. - const secondConsent = await fetch(second.authorizationUrl, { redirect: 'manual' }); + const secondConsent = await fetch(second.authorizationUrl, { + redirect: 'manual', + }); const secondCode = new URL(secondConsent.headers.get('location') ?? '').searchParams.get( 'code', ); @@ -1068,7 +1151,10 @@ describe('McpClientManager OAuth E2E', () => { const storage = createMemoryMcpOAuthStorage(); await storage.set('remote', { serverUrl: 'https://mcp.example/mcp', - clientInformation: { client_id: 'as-a-client', client_secret: 'as-a-secret' }, + clientInformation: { + client_id: 'as-a-client', + client_secret: 'as-a-secret', + }, tokens: { access_token: 'as-a-token', token_type: 'Bearer' }, discovery: { authorizationServerUrl: 'https://as-a.example' } as never, }); @@ -1097,7 +1183,10 @@ describe('McpClientManager OAuth E2E', () => { const storage = createMemoryMcpOAuthStorage(); await storage.set('remote', { serverUrl: 'https://old.example/mcp', - clientInformation: { client_id: 'old-client', client_secret: 'old-secret' }, + clientInformation: { + client_id: 'old-client', + client_secret: 'old-secret', + }, tokens: { access_token: 'old-token', token_type: 'Bearer' }, discovery: { authorizationServerUrl: 'https://as.example' } as never, generation: 3, @@ -1138,7 +1227,11 @@ describe('McpClientManager OAuth E2E', () => { const config: McpConfigFile = { version: MCP_CONFIG_VERSION, mcpServers: { - remote: { url: 'https://mcp.example/mcp', transport: 'streamable-http', enabled: false }, + remote: { + url: 'https://mcp.example/mcp', + transport: 'streamable-http', + enabled: false, + }, }, }; await managerB.sync(config); @@ -1192,8 +1285,16 @@ describe('McpClientManager OAuth E2E', () => { // lands, round B persists a fresh verifier/state. A's abandon must // become a no-op instead of CAS-deleting B's pending round. const rounds: McpOAuthRecord[] = [ - { version: 3, codeVerifier: 'round-a-verifier', pendingState: 'round-a-state' }, - { version: 4, codeVerifier: 'round-b-verifier', pendingState: 'round-b-state' }, + { + version: 3, + codeVerifier: 'round-a-verifier', + pendingState: 'round-a-state', + }, + { + version: 4, + codeVerifier: 'round-b-verifier', + pendingState: 'round-b-state', + }, ]; let stored = rounds[0] as McpOAuthRecord; let reads = 0; @@ -1270,7 +1371,10 @@ describe('McpClientManager OAuth E2E', () => { manager.sync({ version: MCP_CONFIG_VERSION, mcpServers: { - remote: { url: 'https://changed.example/mcp', transport: 'streamable-http' }, + remote: { + url: 'https://changed.example/mcp', + transport: 'streamable-http', + }, }, }), /credential store unavailable/u, @@ -1278,7 +1382,9 @@ describe('McpClientManager OAuth E2E', () => { assert.match(manager.status('remote')?.error ?? '', /could not be removed/u); // Not just connect(): the interactive paths fail closed too. await assert.rejects( - manager.startAuthorization('remote', 'http://127.0.0.1:39998/callback', { state: 's' }), + manager.startAuthorization('remote', 'http://127.0.0.1:39998/callback', { + state: 's', + }), /blocked/u, ); await assert.rejects( @@ -1396,6 +1502,8 @@ async function createOAuthFixture( /** The token endpoint reflects the authorization code it received into * error_description. */ reflectCodeInTokenError?: boolean; + /** Holds the MCP endpoint open until the client aborts the request. */ + holdMcpRequest?: () => Promise; } = {}, ): Promise { let accessToken = `token-${randomUUID()}`; @@ -1414,6 +1522,7 @@ async function createOAuthFixture( const authorization = req.headers.authorization; if (typeof authorization === 'string') lastAuthorization = authorization; mcpRequests.push(typeof authorization === 'string' ? { authorization } : {}); + if (options.holdMcpRequest && (await options.holdMcpRequest())) return; if (options.bareChallengeOnGet && req.method === 'GET') { res.writeHead(401, { 'www-authenticate': 'Bearer realm="mcp"' }).end(); return; @@ -1444,7 +1553,9 @@ async function createOAuthFixture( .end(JSON.stringify({ error: 'unauthorized' })); return; } - const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined }); + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + }); const server = createProtocolServer(); await server.connect(transport); res.once('close', () => { @@ -1475,7 +1586,9 @@ async function createOAuthFixture( res.writeHead(405).end(); return; } - const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined }); + const transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: undefined, + }); const server = createProtocolServer( options.reflectAuthInProtocol ? () => lastAuthorization : undefined, ); @@ -1488,7 +1601,10 @@ async function createOAuthFixture( return; } if (url.pathname === '/.well-known/oauth-protected-resource' && req.method === 'GET') { - json(res, { resource: `${origin}/mcp`, authorization_servers: [origin] }); + json(res, { + resource: `${origin}/mcp`, + authorization_servers: [origin], + }); return; } if (url.pathname === '/.well-known/oauth-authorization-server' && req.method === 'GET') { @@ -1647,7 +1763,10 @@ function createProtocolServer(reflect?: () => string): McpServer { // A server echoing the credential it was just sent — into the tool // metadata the client persists. description: reflect ? `Echo text (${reflect()})` : 'Echo text', - inputSchema: { type: 'object', properties: { value: { type: 'string' } } }, + inputSchema: { + type: 'object', + properties: { value: { type: 'string' } }, + }, }, ], })); diff --git a/packages/mcp/src/credential-coordinator.ts b/packages/mcp/src/credential-coordinator.ts index 9a5bd875a0..8451c61443 100644 --- a/packages/mcp/src/credential-coordinator.ts +++ b/packages/mcp/src/credential-coordinator.ts @@ -96,7 +96,13 @@ export class McpCredentialCoordinator { * The optional signal fences an ABANDONED erase: a logout whose round * timed out must not resume later, adopt whatever record a newer login * just stored as its basis, and tombstone the fresh tokens. */ - async erase(serverId: string, options: { signal?: AbortSignal } = {}): Promise { + async erase( + serverId: string, + options: { + signal?: AbortSignal; + onCommitStarted?: () => void; + } = {}, + ): Promise { this.epochs.set(serverId, this.epoch(serverId) + 1); await this.run(serverId, async () => { this.assertNotAbandoned(serverId, options.signal); @@ -112,6 +118,12 @@ export class McpCredentialCoordinator { generation: (basis?.generation ?? 0) + 1, version: (basis?.version ?? 0) + 1, }; + // From this point the storage implementation owns an in-flight write. + // The caller may stop passing cancellation into later work, but it must + // not report cancellation or compensate related state until this commit + // settles: an atomic backend can durably land the tombstone before its + // promise becomes observable as fulfilled. + options.onCommitStarted?.(); await this.commit(serverId, basis, tombstone); }); } diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 76bdd21acc..def5a30ac8 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -232,6 +232,8 @@ interface Connection { subscription?: McpSubscription; subscriptionDiagnostic?: string; closing: boolean; + teardown?: Promise; + removeAfterTeardown?: boolean; } interface ToolBindingTarget { @@ -308,11 +310,21 @@ export class McpClientManager { private harvestingStorage(storage: McpOAuthStorage): McpOAuthStorage { const harvest = ( serverId: string, - record: { tokens?: unknown; clientInformation?: unknown; codeVerifier?: unknown } | undefined, + record: + | { + tokens?: unknown; + clientInformation?: unknown; + codeVerifier?: unknown; + } + | undefined, ) => { if (!record) return; const tokens = record.tokens as - | { access_token?: unknown; refresh_token?: unknown; id_token?: unknown } + | { + access_token?: unknown; + refresh_token?: unknown; + id_token?: unknown; + } | undefined; const client = record.clientInformation as { client_secret?: unknown } | undefined; const verifier = record.codeVerifier; @@ -390,48 +402,60 @@ export class McpClientManager { return () => this.listeners.delete(listener); } - sync(config: McpConfigFile): Promise { + sync(config: McpConfigFile, options: { signal?: AbortSignal } = {}): Promise { if (this.closed) return Promise.reject(new Error('MCP client manager is closed')); const snapshot = structuredClone(config); - const operation = this.syncQueue.catch(() => {}).then(() => this.syncNow(snapshot)); + const operation = this.syncQueue + .catch(() => {}) + .then(() => { + throwIfAborted(options.signal); + return this.syncNow(snapshot, options.signal); + }); this.syncQueue = operation; return operation; } - private async syncNow(config: McpConfigFile): Promise { + private async syncNow(config: McpConfigFile, signal?: AbortSignal): Promise { + throwIfAborted(signal); const desired = new Set(Object.keys(config.mcpServers)); // A failed erase must not abandon the rest of the reconciliation: the // config file is already written, so stopping here would leave every // OTHER added/changed server diverged until the next sync. The blocked // server stays blocked; the failures reject the sync at the end. const removalFailures: unknown[] = []; - await Promise.all( - [...this.connections.keys()] - .filter((serverId) => !desired.has(serverId)) - .map(async (serverId) => { - const entry = this.connections.get(serverId); - // Credentials first, connection second: a removed server's stored - // OAuth tokens are a hazard — a same-id server added back later - // must not inherit them. Erasing is the authoritative transition; - // only after it succeeds may connection ownership be released. On - // failure the entry stays, blocked — the next sync retries. - try { - await this.forgetAuthorization(serverId, entry?.credentialCleanupOwed ?? entry?.config); - } catch (error) { - if (entry) { - await this.blockForCredentialCleanup( - serverId, - entry, - entry.credentialCleanupOwed ?? entry.config, - error, - ); - } - removalFailures.push(error); - return; + const removals = [...this.connections.keys()] + .filter((serverId) => !desired.has(serverId)) + .map(async (serverId) => { + const entry = this.connections.get(serverId); + // Credentials first, connection second: a removed server's stored + // OAuth tokens are a hazard — a same-id server added back later + // must not inherit them. Erasing is the authoritative transition; + // only after it succeeds may connection ownership be released. On + // failure the entry stays, blocked — the next sync retries. + try { + await this.forgetAuthorization(serverId, entry?.credentialCleanupOwed ?? entry?.config, { + signal, + }); + } catch (error) { + if (signal?.aborted) throw error; + if (entry) { + await this.blockForCredentialCleanup( + serverId, + entry, + entry.credentialCleanupOwed ?? entry.config, + error, + ); } - await this.disconnect(serverId, true); - }), + removalFailures.push(error); + return; + } + await this.disconnect(serverId, true, { signal }); + }); + const removalResults = await Promise.allSettled(removals); + const removalRejection = removalResults.find( + (result): result is PromiseRejectedResult => result.status === 'rejected', ); + if (removalRejection) throw removalRejection.reason; const connectIds: string[] = []; for (const [serverId, serverConfig] of Object.entries(config.mcpServers)) { const fingerprint = stableConfigFingerprint(serverConfig); @@ -451,8 +475,9 @@ export class McpClientManager { : undefined; if (owed) { try { - await this.forgetAuthorization(serverId, owed); + await this.forgetAuthorization(serverId, owed, { signal }); } catch (error) { + if (signal?.aborted) throw error; await this.blockForCredentialCleanup(serverId, current, owed, error); // Same contract as the removal loop: the config is already // written to the NEW endpoint while the old one stays blocked — @@ -463,7 +488,7 @@ export class McpClientManager { } current.credentialCleanupOwed = undefined; } - await this.disconnect(serverId, true); + await this.disconnect(serverId, true, { signal }); } else if (current?.credentialCleanupOwed) { // Same fingerprint again: the config reverted to (or never left) // the endpoint the credentials belong to — nothing is owed. @@ -484,7 +509,10 @@ export class McpClientManager { } if (serverConfig.enabled !== false) connectIds.push(serverId); } - await Promise.all(connectIds.map((serverId) => this.connect(serverId).catch(() => {}))); + await Promise.all( + connectIds.map((serverId) => this.connect(serverId, { signal }).catch(() => {})), + ); + throwIfAborted(signal); if (removalFailures.length > 0) throw removalFailures[0]; } @@ -528,8 +556,12 @@ export class McpClientManager { return this.callableSnapshot; } - async connect(serverId: string): Promise { + async connect( + serverId: string, + options: { signal?: AbortSignal } = {}, + ): Promise { if (this.closed) throw new Error('MCP client manager is closed'); + throwIfAborted(options.signal); const entry = this.requireConnection(serverId); if (entry.closing) throw new Error(`MCP server "${serverId}" is closing`); if (entry.credentialCleanupOwed) { @@ -547,10 +579,19 @@ export class McpClientManager { // state mid-round and rotate the record under the round's fence. return cloneStatus(entry.status); } - if (entry.connectPromise) return entry.connectPromise; + if (entry.connectPromise) { + // A joining caller observes the manager-owned single flight; it does not + // acquire cancellation ownership over the caller that created it. The + // join may stop waiting, while explicit lifecycle owners can still use + // cancelConnect()/disconnect() to terminate the shared transport. + return waitForAbort(entry.connectPromise, options.signal); + } const controller = new AbortController(); + const forwardAbort = () => controller.abort(abortReason(options.signal)); + options.signal?.addEventListener('abort', forwardAbort, { once: true }); entry.connectController = controller; const promise = this.connectEntry(serverId, entry, controller.signal).finally(() => { + options.signal?.removeEventListener('abort', forwardAbort); if (entry.connectPromise === promise) entry.connectPromise = undefined; if (entry.connectController === controller) entry.connectController = undefined; }); @@ -567,15 +608,30 @@ export class McpClientManager { return true; } - async reconnect(serverId: string): Promise { - await this.disconnect(serverId, false); - return this.connect(serverId); + async reconnect( + serverId: string, + options: { signal?: AbortSignal } = {}, + ): Promise { + throwIfAborted(options.signal); + await this.disconnect(serverId, false, options); + throwIfAborted(options.signal); + return this.connect(serverId, options); } - async disconnect(serverId: string, remove = false): Promise { + async disconnect( + serverId: string, + remove = false, + options: { signal?: AbortSignal } = {}, + ): Promise { const entry = this.connections.get(serverId); if (!entry) return; + if (entry.teardown) { + if (remove) entry.removeAfterTeardown = true; + await waitForAbort(entry.teardown, options.signal); + return; + } entry.closing = true; + entry.removeAfterTeardown = remove; entry.connectController?.abort(new Error(`MCP connection closed: ${serverId}`)); const connectPromise = entry.connectPromise; const client = entry.client; @@ -591,17 +647,28 @@ export class McpClientManager { entry.enforceMcpHeaders = false; entry.refreshDiagnostic = undefined; entry.subscriptionDiagnostic = undefined; - await safeClose(client, transport, subscription); - await connectPromise?.catch(() => {}); - if (remove) { - this.connections.delete(serverId); - return; + if (!remove) { + this.update(entry, { + ...this.makeStatus(serverId, entry.config.enabled === false ? 'disabled' : 'disconnected'), + stderrTail: entry.status.stderrTail, + }); } - entry.closing = false; - this.update(entry, { - ...this.makeStatus(serverId, entry.config.enabled === false ? 'disabled' : 'disconnected'), - stderrTail: entry.status.stderrTail, + const cleanup = Promise.all([ + safeClose(client, transport, subscription), + connectPromise?.catch(() => {}), + ]).then(() => undefined); + const teardown = cleanup.then(() => { + if (entry.teardown !== teardown) return; + entry.teardown = undefined; + if (entry.removeAfterTeardown) { + if (this.connections.get(serverId) === entry) this.connections.delete(serverId); + } else if (this.connections.get(serverId) === entry) { + entry.closing = false; + } + entry.removeAfterTeardown = undefined; }); + entry.teardown = teardown; + await waitForAbort(teardown, options.signal); } async close(): Promise { @@ -882,20 +949,25 @@ export class McpClientManager { }; } - async test(serverId: string): Promise { + async test(serverId: string, options: { signal?: AbortSignal } = {}): Promise { const started = this.now(); + throwIfAborted(options.signal); const current = this.requireConnection(serverId); if (current.config.enabled === false) { return { ok: false, - status: { ...cloneStatus(current.status), error: 'MCP server is disabled' }, + status: { + ...cloneStatus(current.status), + error: 'MCP server is disabled', + }, latencyMs: this.now() - started, }; } try { - const status = await this.reconnect(serverId); + const status = await this.reconnect(serverId, options); return { ok: true, status, latencyMs: this.now() - started }; - } catch { + } catch (error) { + if (options.signal?.aborted) throw error; return { ok: false, status: this.status(serverId) ?? this.makeStatus(serverId, 'error'), @@ -1157,7 +1229,9 @@ export class McpClientManager { // A plain static-bearer config with no OAuth involvement keeps its // header untouched. The config store rejects the explicit conflict // (oauth block + Authorization header) outright. - const record = this.coordinator ? await this.coordinator.read(serverId) : undefined; + const record = this.coordinator + ? await waitForAbort(this.coordinator.read(serverId), signal) + : undefined; // The raw record only counts when it is BOUND to this endpoint: after an // offline mcp.json repoint the stale record's tokens will be dropped by // the provider, so they must not strip a configured header either. @@ -1260,7 +1334,10 @@ export class McpClientManager { /** Flow view with generation/version pinned before any remote await. */ private beginFlow(serverId: string, signal?: AbortSignal): Promise { - return this.requireCoordinator().beginFlow(serverId, signal ? { signal } : {}); + return waitForAbort( + this.requireCoordinator().beginFlow(serverId, signal ? { signal } : {}), + signal, + ); } private requireCoordinator(): McpCredentialCoordinator { @@ -1552,7 +1629,10 @@ export class McpClientManager { ) { return undefined; } - return { redirectUrl: record.pendingRedirectUrl, state: record.pendingState }; + return { + redirectUrl: record.pendingRedirectUrl, + state: record.pendingState, + }; } /** Abandons a persisted-but-dead interactive round: clears the verifier @@ -1601,8 +1681,12 @@ export class McpClientManager { async forgetServerCredentials( serverId: string, previousConfig = this.connections.get(serverId)?.config, + options: { + signal?: AbortSignal; + onCommitStarted?: () => void; + } = {}, ): Promise { - await this.forgetAuthorization(serverId, previousConfig); + await this.forgetAuthorization(serverId, previousConfig, options); } /** Drops any stored OAuth record for a server that is being removed or @@ -1618,10 +1702,24 @@ export class McpClientManager { private async forgetAuthorization( serverId: string, config?: McpServerConfig, - options: { signal?: AbortSignal } = {}, + options: { + signal?: AbortSignal; + onCommitStarted?: () => void; + } = {}, ): Promise { if (!this.coordinator) return; - if (config && isMcpStdioConfig(config) && !(await this.coordinator.read(serverId))) return; + throwIfAborted(options.signal); + if ( + config && + isMcpStdioConfig(config) && + !(await waitForAbort(this.coordinator.read(serverId), options.signal)) + ) { + return; + } + throwIfAborted(options.signal); + // erase() owns the write fence: once its storage commit starts it must + // settle before sync can report cancellation, otherwise a tombstone can + // land after the caller has already begun rollback/reconciliation. await this.coordinator.erase(serverId, options); } @@ -1633,7 +1731,12 @@ export class McpClientManager { const { config } = this.requireRemoteEntry(serverId); this.interactiveRounds.delete(serverId); await this.forgetAuthorization(serverId, config, options); - await this.reconnect(serverId).catch(() => {}); + try { + await this.reconnect(serverId, options); + } catch (error) { + if (options.signal?.aborted) throw abortReason(options.signal); + } + throwIfAborted(options.signal); const status = this.status(serverId); if (!status) throw new Error(`Unknown MCP server: ${serverId}`); return status; @@ -1781,7 +1884,10 @@ export class McpClientManager { state: ToolRefreshState, ): Promise { let latestSnapshot: - | { entries: Map; descriptors: McpToolDescriptor[] } + | { + entries: Map; + descriptors: McpToolDescriptor[]; + } | undefined; const finish = ( snapshot: NonNullable, @@ -2073,7 +2179,10 @@ export class McpClientManager { } const tools = entries.map( ({ descriptor, binding }) => - deepFreeze({ descriptor: cloneTool(descriptor), binding }) as McpBoundTool, + deepFreeze({ + descriptor: cloneTool(descriptor), + binding, + }) as McpBoundTool, ); return Object.freeze({ revision: this.callableSnapshot.revision + 1, @@ -2534,7 +2643,9 @@ function enrichStdioError( secrets: SecretInventory = EMPTY_INVENTORY, ): Error { const suffix = stderrTail?.length ? `\nstderr:\n${stderrTail.join('\n')}` : ''; - return new Error(`${errorMessage(error, secrets)}${suffix}`, { cause: error }); + return new Error(`${errorMessage(error, secrets)}${suffix}`, { + cause: error, + }); } async function safeClose( @@ -2550,17 +2661,48 @@ async function safeClose( ]); } +function abortReason(signal?: AbortSignal): Error { + return signal?.reason instanceof Error + ? signal.reason + : new Error(String(signal?.reason ?? 'MCP operation aborted')); +} + +function throwIfAborted(signal?: AbortSignal): void { + if (signal?.aborted) throw abortReason(signal); +} + +function waitForAbort(promise: Promise, signal?: AbortSignal): Promise { + if (!signal) return promise; + if (signal.aborted) return Promise.reject(abortReason(signal)); + return new Promise((resolve, reject) => { + let settled = false; + const settle = (callback: () => void) => { + if (settled) return; + settled = true; + signal.removeEventListener('abort', onAbort); + callback(); + }; + const onAbort = () => settle(() => reject(abortReason(signal))); + signal.addEventListener('abort', onAbort, { once: true }); + void promise.then( + (value) => settle(() => resolve(value)), + (error: unknown) => settle(() => reject(error)), + ); + }); +} + async function connectCandidate( client: Client, transport: Transport, timeout: number, signal: AbortSignal, ): Promise { + let abortClose: Promise | undefined; const closeOnAbort = () => { // SDK v2's server/discover probes currently use their timeout but not the // Client.connect signal. Closing the candidate transport aborts either an // HTTP probe or the disposable stdio sibling before a late session starts. - void transport.close().catch(() => {}); + abortClose ??= closeAbortedTransport(transport); }; if (signal.aborted) { await transport.close().catch(() => {}); @@ -2573,9 +2715,25 @@ async function connectCandidate( await client.connect(transport, { timeout, signal }); } finally { signal.removeEventListener('abort', closeOnAbort); + await abortClose; } } +function closeAbortedTransport(transport: Transport): Promise { + if (transport instanceof StdioClientTransport) { + // @modelcontextprotocol/client is pinned to 2.0.0. In that release the + // public close() clears its child handle before the shutdown finishes, so + // a second close cannot join the in-flight reap. The SDK's own disposable + // stdio-probe path uses _dispose(), which waits for the child `exit` event; + // use that same reaper here so cancellation cannot settle while the old + // server process is still alive. The real-child regression test must stay + // green before changing the SDK version or this private compatibility shim. + const disposable = transport as unknown as { _dispose?(): Promise }; + if (disposable._dispose) return disposable._dispose().catch(() => {}); + } + return transport.close().catch(() => {}); +} + function stableConfigFingerprint(config: McpServerConfig): string { return JSON.stringify(sortValue(config)); } @@ -2673,7 +2831,10 @@ async function probeAuthChallenge( fetchImpl: typeof fetch, ): Promise<{ scope?: string; resourceMetadataUrl?: URL } | undefined> { const attempts: RequestInit[] = [ - { method: 'GET', headers: { accept: 'text/event-stream, application/json' } }, + { + method: 'GET', + headers: { accept: 'text/event-stream, application/json' }, + }, { method: 'POST', headers: { @@ -2793,7 +2954,11 @@ function scopedFetch( redirect: 'manual', // The round's deadline aborts in-flight requests too, not only the // caller's await: a hung endpoint must not keep the flow alive. - ...(signal ? { signal: init?.signal ? AbortSignal.any([init.signal, signal]) : signal } : {}), + ...(signal + ? { + signal: init?.signal ? AbortSignal.any([init.signal, signal]) : signal, + } + : {}), }; }; return (async (input: RequestInfo | URL, init?: RequestInit) => { diff --git a/packages/runtime-host/src/__tests__/host-profile.test.ts b/packages/runtime-host/src/__tests__/host-profile.test.ts index 45cc7b654a..01ab48ca17 100644 --- a/packages/runtime-host/src/__tests__/host-profile.test.ts +++ b/packages/runtime-host/src/__tests__/host-profile.test.ts @@ -172,7 +172,10 @@ describe('Runtime Host profiles', () => { id: 'backup', name: 'Backup', kind: 'remote', - transport: { kind: 'tls', url: 'wss://backup.example.com/runtime-host' }, + transport: { + kind: 'tls', + url: 'wss://backup.example.com/runtime-host', + }, rootId: ROOT_B, }, 'loopback-token', @@ -203,7 +206,10 @@ describe('Runtime Host profiles', () => { id: 'backup', name: 'Backup', kind: 'remote', - transport: { kind: 'tls', url: 'wss://backup.example.com/runtime-host' }, + transport: { + kind: 'tls', + url: 'wss://backup.example.com/runtime-host', + }, rootId: ROOT_B, }, ], @@ -219,7 +225,10 @@ describe('Runtime Host profiles', () => { id: 'backup', name: 'Backup', kind: 'remote', - transport: { kind: 'tls', url: 'wss://backup.example.com/runtime-host' }, + transport: { + kind: 'tls', + url: 'wss://backup.example.com/runtime-host', + }, rootId: ROOT_B, }, ], @@ -408,7 +417,10 @@ describe('Runtime Host profiles', () => { id: 'office', name: 'Office', kind: 'remote', - transport: { kind: 'tls', url: 'wss://runtime.example.com/runtime-host' }, + transport: { + kind: 'tls', + url: 'wss://runtime.example.com/runtime-host', + }, rootId: ROOT_A, }, ], @@ -428,7 +440,10 @@ describe('Runtime Host profiles', () => { profiles: [ { ...valid.profiles[0], - transport: { kind: 'tls', url: 'ws://runtime.example.com/runtime-host' }, + transport: { + kind: 'tls', + url: 'ws://runtime.example.com/runtime-host', + }, }, ], }), @@ -535,12 +550,18 @@ describe('Runtime Host profiles', () => { ); const targetA = remoteProfile('office', 'wss://a.example.com', ROOT_A); const targetB = remoteProfile('office', 'wss://b.example.com', ROOT_B); - const incarnationA = { profile: targetA, profileIncarnationId: 'incarnation-a' }; + const incarnationA = { + profile: targetA, + profileIncarnationId: 'incarnation-a', + }; const recreatedIncarnationA = { profile: targetA, profileIncarnationId: 'incarnation-a-recreated', }; - const incarnationB = { profile: targetB, profileIncarnationId: 'incarnation-b' }; + const incarnationB = { + profile: targetB, + profileIncarnationId: 'incarnation-b', + }; await assert.rejects( () => credentials.set(incarnationA, 'owner-a', 'not a token'), @@ -560,6 +581,115 @@ describe('Runtime Host profiles', () => { assert.equal(await credentials.get(incarnationA, 'owner-b'), 'provider-b'); }); + test('conditionally restores capability-provider credentials without overwriting a newer value', async () => { + const path = await profilePath(); + const credentials = createRuntimeHostCapabilityProviderCredentialStore( + createFileCredentialStore(join(dirname(path), 'credentials')), + ); + assert.ok(credentials.read); + assert.ok(credentials.compareAndSet); + const profile = remoteProfile('office', 'wss://a.example.com', ROOT_A); + const target = { profile, profileIncarnationId: 'incarnation-a' }; + + await credentials.set(target, 'owner-a', 'cancelled-value'); + const written = await credentials.read(target, 'owner-a'); + assert.equal(written.credential, 'cancelled-value'); + assert.equal( + (await credentials.compareAndSet(target, 'owner-a', written.revision, null)).committed, + true, + ); + assert.equal(await credentials.get(target, 'owner-a'), null); + + await credentials.set(target, 'owner-a', 'cancelled-value'); + const stale = await credentials.read(target, 'owner-a'); + await credentials.set(target, 'owner-a', 'newer-value'); + assert.equal( + (await credentials.compareAndSet(target, 'owner-a', stale.revision, 'old-value')).committed, + false, + ); + assert.equal(await credentials.get(target, 'owner-a'), 'newer-value'); + }); + + test('capability-provider revision CAS rejects an ABA cycle at the adapter boundary', async () => { + const path = await profilePath(); + const credentials = createRuntimeHostCapabilityProviderCredentialStore( + createFileCredentialStore(join(dirname(path), 'credentials')), + ); + assert.ok(credentials.read); + assert.ok(credentials.compareAndSet); + const profile = remoteProfile('office', 'wss://a.example.com', ROOT_A); + const target = { profile, profileIncarnationId: 'incarnation-a' }; + + await credentials.set(target, 'owner-a', 'cancelled-value'); + const stale = await credentials.read(target, 'owner-a'); + await credentials.set(target, 'owner-a', 'newer-value'); + await credentials.set(target, 'owner-a', 'cancelled-value'); + + const result = await credentials.compareAndSet(target, 'owner-a', stale.revision, 'old-value'); + assert.equal(result.committed, false); + if (result.committed) return; + assert.equal(result.current.credential, 'cancelled-value'); + assert.notEqual(result.current.revision, stale.revision); + assert.equal(await credentials.get(target, 'owner-a'), 'cancelled-value'); + }); + + test('capability-provider reads retain the raw revision for an owner handoff', async () => { + const path = await profilePath(); + const credentials = createRuntimeHostCapabilityProviderCredentialStore( + createFileCredentialStore(join(dirname(path), 'credentials')), + ); + assert.ok(credentials.read); + assert.ok(credentials.compareAndSet); + const profile = remoteProfile('office', 'wss://a.example.com', ROOT_A); + const target = { profile, profileIncarnationId: 'incarnation-a' }; + + await credentials.set(target, 'owner-a', 'provider-a'); + const handoffBasis = await credentials.read(target, 'owner-b'); + assert.equal(handoffBasis.credential, null); + assert.ok(handoffBasis.revision); + + const result = await credentials.compareAndSet( + target, + 'owner-b', + handoffBasis.revision, + 'provider-b', + ); + assert.equal(result.committed, true); + assert.equal(await credentials.get(target, 'owner-a'), null); + assert.equal(await credentials.get(target, 'owner-b'), 'provider-b'); + }); + + test('capability-provider deletion rejects when its credential changed after the read', async () => { + const path = await profilePath(); + const raw = createFileCredentialStore(join(dirname(path), 'credentials')); + const concurrent = createRuntimeHostCapabilityProviderCredentialStore(raw); + const profile = remoteProfile('office', 'wss://a.example.com', ROOT_A); + const target = { profile, profileIncarnationId: 'incarnation-a' }; + await concurrent.set(target, 'owner-a', 'old-value'); + let raced = false; + const credentials = createRuntimeHostCapabilityProviderCredentialStore({ + getSecret: raw.getSecret.bind(raw), + getSecretSnapshot: raw.getSecretSnapshot?.bind(raw), + setSecret: raw.setSecret.bind(raw), + deleteSecret: raw.deleteSecret.bind(raw), + compareAndSetSecret: raw.compareAndSetSecret?.bind(raw), + compareAndSetSecretRevision: async (...args) => { + if (!raced) { + raced = true; + await concurrent.set(target, 'owner-a', 'newer-value'); + } + assert.ok(raw.compareAndSetSecretRevision); + return raw.compareAndSetSecretRevision(...args); + }, + }); + + await assert.rejects( + credentials.delete(target, 'owner-a'), + /credential changed during deletion/u, + ); + assert.equal(await concurrent.get(target, 'owner-a'), 'newer-value'); + }); + test('removing a profile retires its terminal and provider credentials together', async () => { const path = await profilePath(); const credentialStore = createFileCredentialStore(join(dirname(path), 'credentials')); @@ -572,7 +702,10 @@ describe('Runtime Host profiles', () => { await catalog.save(profile, 'terminal-token'); const target = await catalog.resolve(profile.id); assert.ok(target.profileIncarnationId); - const incarnation = { profile, profileIncarnationId: target.profileIncarnationId }; + const incarnation = { + profile, + profileIncarnationId: target.profileIncarnationId, + }; await providers.set(incarnation, 'owner-a', 'provider-token'); await catalog.remove(profile.id); @@ -601,7 +734,10 @@ describe('Runtime Host profiles', () => { await removingCatalog.save(profile, 'terminal-token'); const resolved = await removingCatalog.resolve(profile.id); assert.ok(resolved.profileIncarnationId); - const incarnation = { profile, profileIncarnationId: resolved.profileIncarnationId }; + const incarnation = { + profile, + profileIncarnationId: resolved.profileIncarnationId, + }; const removal = removingCatalog.remove(profile.id); await removalStarted.promise; @@ -609,7 +745,10 @@ describe('Runtime Host profiles', () => { const mutation = mutatingCatalog.mutateRemoteProfileIfCurrent(incarnation, async (current) => { mutationRan = true; await providers.set( - { profile: current, profileIncarnationId: incarnation.profileIncarnationId }, + { + profile: current, + profileIncarnationId: incarnation.profileIncarnationId, + }, 'owner-a', 'provider-token', ); @@ -642,7 +781,10 @@ describe('Runtime Host profiles', () => { await removingCatalog.save(profile, 'terminal-token'); const resolved = await removingCatalog.resolve(profile.id); assert.ok(resolved.profileIncarnationId); - const incarnation = { profile, profileIncarnationId: resolved.profileIncarnationId }; + const incarnation = { + profile, + profileIncarnationId: resolved.profileIncarnationId, + }; const removal = removingCatalog.remove(profile.id); await removalStarted.promise; @@ -958,7 +1100,10 @@ describe('Runtime Host profiles', () => { events.push('tunnel'); assert.equal(input.remotePort, 43_210); assert.equal(input.websocketPath, '/runtime-host/activated'); - return { url: 'ws://127.0.0.1:43211/runtime-host/activated', resource }; + return { + url: 'ws://127.0.0.1:43211/runtime-host/activated', + resource, + }; }, connect: async () => { events.push('connect'); @@ -986,7 +1131,10 @@ describe('Runtime Host profiles', () => { clientInstanceId: 'client-1', }, { - connect: async () => ({ kind: 'unavailable', reason: 'root_mismatch' }), + connect: async () => ({ + kind: 'unavailable', + reason: 'root_mismatch', + }), }, ), (error: unknown) => { @@ -1173,7 +1321,12 @@ describe('Runtime Host profiles', () => { credential: 'revoked-token', clientInstanceId: 'client-1', }, - { connect: async () => ({ kind: 'unavailable', reason: 'authentication_failed' }) }, + { + connect: async () => ({ + kind: 'unavailable', + reason: 'authentication_failed', + }), + }, ), (error: unknown) => { assert.ok(error instanceof RuntimeHostPermanentReconnectError); @@ -1283,7 +1436,13 @@ async function profilePath(): Promise { } function remoteProfile(id: string, url: string, rootId: string): RemoteRuntimeHostProfile { - return { id, name: id, kind: 'remote', transport: { kind: 'tls', url }, rootId }; + return { + id, + name: id, + kind: 'remote', + transport: { kind: 'tls', url }, + rootId, + }; } function directPeerProfile( @@ -1297,7 +1456,10 @@ function directPeerProfile( name: 'Peer', kind: 'remote', rootId: ROOT_A, - transport: { kind: 'libp2p-direct', reachability: reachability(peerId, routeHints) }, + transport: { + kind: 'libp2p-direct', + reachability: reachability(peerId, routeHints), + }, }; } diff --git a/packages/runtime-host/src/client/host-profile.ts b/packages/runtime-host/src/client/host-profile.ts index e7cf14c65a..4c7c1f1945 100644 --- a/packages/runtime-host/src/client/host-profile.ts +++ b/packages/runtime-host/src/client/host-profile.ts @@ -286,8 +286,30 @@ export interface RuntimeHostCapabilityProviderCredentialStore { credential: string, ): Promise; delete(target: RuntimeHostRemoteProfileIncarnation, ownerClientInstanceId: string): Promise; + read?( + target: RuntimeHostRemoteProfileIncarnation, + ownerClientInstanceId: string, + ): Promise; + compareAndSet?( + target: RuntimeHostRemoteProfileIncarnation, + ownerClientInstanceId: string, + expectedRevision: string | null, + credential: string | null, + ): Promise; } +export interface RuntimeHostCapabilityProviderCredentialSnapshot { + readonly credential: string | null; + readonly revision: string | null; +} + +export type RuntimeHostCapabilityProviderCredentialMutationResult = + | { readonly committed: true; readonly revision: string } + | { + readonly committed: false; + readonly current: RuntimeHostCapabilityProviderCredentialSnapshot; + }; + export type RuntimeHostProfileConnectionFailureReason = | 'credential_required' | 'credential_rejected' @@ -364,35 +386,108 @@ export function createRuntimeHostProfileCredentialStore( } export function createRuntimeHostCapabilityProviderCredentialStore( - credentials: Pick, + credentials: Pick< + CredentialStore, + | 'getSecret' + | 'getSecretSnapshot' + | 'setSecret' + | 'deleteSecret' + | 'compareAndSetSecret' + | 'compareAndSetSecretRevision' + >, ): RuntimeHostCapabilityProviderCredentialStore { - return { - get: async (target, ownerClientInstanceId) => { - const stored = await credentials.getSecret( - profileCredentialSlot(target.profile), - 'runtime_host_capability_provider', - ); - if (stored === null) return null; - const decoded = decodeCapabilityProviderCredential(stored); - return decoded.ownerClientInstanceId === requireClientInstanceId(ownerClientInstanceId) && + const getSecretSnapshot = credentials.getSecretSnapshot?.bind(credentials); + const compareAndSetSecretRevision = credentials.compareAndSetSecretRevision?.bind(credentials); + const locator = ( + target: RuntimeHostRemoteProfileIncarnation, + ownerClientInstanceId: string, + credential: string, + ) => ({ + slot: profileCredentialSlot(target.profile), + encoded: encodeCapabilityProviderCredential(target, ownerClientInstanceId, credential), + }); + const read = async ( + target: RuntimeHostRemoteProfileIncarnation, + ownerClientInstanceId: string, + ): Promise => { + const stored = getSecretSnapshot + ? await getSecretSnapshot( + profileCredentialSlot(target.profile), + 'runtime_host_capability_provider', + ) + : { + value: await credentials.getSecret( + profileCredentialSlot(target.profile), + 'runtime_host_capability_provider', + ), + revision: null, + }; + if (stored.value === null) return { credential: null, revision: stored.revision }; + const decoded = decodeCapabilityProviderCredential(stored.value); + return { + credential: + decoded.ownerClientInstanceId === requireClientInstanceId(ownerClientInstanceId) && decoded.profileIncarnationId === requireProfileIncarnationId(target.profileIncarnationId) - ? decoded.credential - : null; - }, + ? decoded.credential + : null, + revision: stored.revision, + }; + }; + return { + get: async (target, ownerClientInstanceId) => + (await read(target, ownerClientInstanceId)).credential, set: async (target, ownerClientInstanceId, credential) => { - await credentials.setSecret( - profileCredentialSlot(target.profile), - 'runtime_host_capability_provider', - JSON.stringify({ - schemaVersion: 1, - profileIncarnationId: requireProfileIncarnationId(target.profileIncarnationId), - ownerClientInstanceId: requireClientInstanceId(ownerClientInstanceId), - credential: requireRuntimeHostAccessCredential(credential), - }), - ); + const value = locator(target, ownerClientInstanceId, credential); + await credentials.setSecret(value.slot, 'runtime_host_capability_provider', value.encoded); }, delete: (target, ownerClientInstanceId) => deleteCapabilityProviderCredential(credentials, target, ownerClientInstanceId), + ...(getSecretSnapshot && compareAndSetSecretRevision + ? { + read, + compareAndSet: async ( + target: RuntimeHostRemoteProfileIncarnation, + ownerClientInstanceId: string, + expectedRevision: string | null, + credential: string | null, + ) => { + const slot = profileCredentialSlot(target.profile); + const result = await compareAndSetSecretRevision( + slot, + 'runtime_host_capability_provider', + expectedRevision, + credential === null + ? null + : encodeCapabilityProviderCredential(target, ownerClientInstanceId, credential), + ); + if (result.committed) return result; + const current = result.current.value; + if (current === null) { + return { + committed: false as const, + current: { + credential: null, + revision: result.current.revision, + }, + }; + } + const decoded = decodeCapabilityProviderCredential(current); + return { + committed: false as const, + current: { + credential: + decoded.ownerClientInstanceId === + requireClientInstanceId(ownerClientInstanceId) && + decoded.profileIncarnationId === + requireProfileIncarnationId(target.profileIncarnationId) + ? decoded.credential + : null, + revision: result.current.revision, + }, + }; + }, + } + : {}), }; } @@ -989,7 +1084,9 @@ class FileRuntimeHostProfileCatalog implements RuntimeHostProfileCatalog { const value: unknown = JSON.parse(new TextDecoder('utf-8', { fatal: true }).decode(bytes)); return decodeRuntimeHostProfileDocument(value); } catch (error) { - throw new Error('Runtime Host profile document is invalid', { cause: error }); + throw new Error('Runtime Host profile document is invalid', { + cause: error, + }); } } @@ -1138,7 +1235,10 @@ class FileRuntimeHostProfileCatalog implements RuntimeHostProfileCatalog { ) { return { removed: false, document: current }; } - return { removed: true, document: await this.#removeProfile(current, profile) }; + return { + removed: true, + document: await this.#removeProfile(current, profile), + }; }); } @@ -1411,7 +1511,9 @@ export function decodeRuntimeHostRemoteTransport(value: unknown): RuntimeHostRem if (new URL(rawUrl).protocol !== 'ws:') { throw new Error('Runtime Host plaintext URL must use ws'); } - const url = normalizeRemoteRuntimeHostUrl(rawUrl, { allowInsecureRemote: true }); + const url = normalizeRemoteRuntimeHostUrl(rawUrl, { + allowInsecureRemote: true, + }); return Object.freeze({ kind: 'plaintext', url: url.toString(), @@ -1497,11 +1599,39 @@ function profileCredentialSlot(profile: RemoteRuntimeHostProfile): string { } async function deleteCapabilityProviderCredential( - credentials: Pick, + credentials: Pick< + CredentialStore, + | 'getSecret' + | 'getSecretSnapshot' + | 'deleteSecret' + | 'compareAndSetSecret' + | 'compareAndSetSecretRevision' + >, target: RuntimeHostRemoteProfileIncarnation, ownerClientInstanceId: string, ): Promise { const slot = profileCredentialSlot(target.profile); + if (credentials.getSecretSnapshot && credentials.compareAndSetSecretRevision) { + const stored = await credentials.getSecretSnapshot(slot, 'runtime_host_capability_provider'); + if (stored.value === null) return; + const decoded = decodeCapabilityProviderCredential(stored.value); + if ( + decoded.ownerClientInstanceId !== requireClientInstanceId(ownerClientInstanceId) || + decoded.profileIncarnationId !== requireProfileIncarnationId(target.profileIncarnationId) + ) { + return; + } + const result = await credentials.compareAndSetSecretRevision( + slot, + 'runtime_host_capability_provider', + stored.revision, + null, + ); + if (!result.committed) { + throw new Error('Runtime Host capability-provider credential changed during deletion'); + } + return; + } const stored = await credentials.getSecret(slot, 'runtime_host_capability_provider'); if (stored === null) return; const decoded = decodeCapabilityProviderCredential(stored); @@ -1511,6 +1641,18 @@ async function deleteCapabilityProviderCredential( ) { return; } + if (credentials.compareAndSetSecret) { + const result = await credentials.compareAndSetSecret( + slot, + 'runtime_host_capability_provider', + stored, + null, + ); + if (!result.committed) { + throw new Error('Runtime Host capability-provider credential changed during deletion'); + } + return; + } await credentials.deleteSecret(slot, 'runtime_host_capability_provider'); } @@ -1536,10 +1678,25 @@ function decodeCapabilityProviderCredential(value: string): { profileIncarnationId: requireProfileIncarnationId(record.profileIncarnationId), }; } catch (error) { - throw new Error('Runtime Host capability-provider credential is invalid', { cause: error }); + throw new Error('Runtime Host capability-provider credential is invalid', { + cause: error, + }); } } +function encodeCapabilityProviderCredential( + target: RuntimeHostRemoteProfileIncarnation, + ownerClientInstanceId: string, + credential: string, +): string { + return JSON.stringify({ + schemaVersion: 1, + profileIncarnationId: requireProfileIncarnationId(target.profileIncarnationId), + ownerClientInstanceId: requireClientInstanceId(ownerClientInstanceId), + credential: requireRuntimeHostAccessCredential(credential), + }); +} + function encodeProfileCredential(credential: RuntimeHostProfileCredential): string { return `${PROFILE_CREDENTIAL_RECORD_PREFIX}${JSON.stringify({ schemaVersion: 1, @@ -1574,7 +1731,9 @@ function decodeProfileCredential( profileIncarnationId: requireProfileIncarnationId(record.profileIncarnationId), }; } catch (error) { - throw new Error('Runtime Host profile credential is invalid', { cause: error }); + throw new Error('Runtime Host profile credential is invalid', { + cause: error, + }); } } @@ -1777,7 +1936,10 @@ function requireExactRecord( } function emptyProfileDocument(): RuntimeHostProfileDocument { - return Object.freeze({ schemaVersion: PROFILE_SCHEMA_VERSION, profiles: Object.freeze([]) }); + return Object.freeze({ + schemaVersion: PROFILE_SCHEMA_VERSION, + profiles: Object.freeze([]), + }); } async function writeProfileDocument( diff --git a/packages/runtime-host/src/client/index.ts b/packages/runtime-host/src/client/index.ts index c5e727f15d..cbf34f0e6c 100644 --- a/packages/runtime-host/src/client/index.ts +++ b/packages/runtime-host/src/client/index.ts @@ -83,6 +83,8 @@ export { type RuntimeHostProfileCatalog, type RuntimeHostConnectionPhase, type RuntimeHostRemoteProfileIncarnation, + type RuntimeHostCapabilityProviderCredentialMutationResult, + type RuntimeHostCapabilityProviderCredentialSnapshot, type RuntimeHostCapabilityProviderCredentialStore, RuntimeHostProfileConnectionError, type RuntimeHostProfileConnectionFailureReason, @@ -166,7 +168,10 @@ export { createRuntimeHostCandidateLaunchBarrier, type RuntimeHostCandidateLaunchBarrier, } from './candidate-launch-barrier.js'; -export { runHostedExecution, type RunHostedExecutionInput } from './hosted-execution.js'; +export { + runHostedExecution, + type RunHostedExecutionInput, +} from './hosted-execution.js'; export { type ClientCapabilityProvider } from './client-capability.js'; export { readRuntimeHostAgentGraphEpochs, @@ -187,7 +192,10 @@ export { type IssueRuntimeHostOwnerConnectionCodeInput, type RuntimeHostOwnerConnectionCode, } from './owner-connection-code.js'; -export { ensureRuntimeHostPeerIdentity, RuntimeHostPeerError } from '../transport/peer-native.js'; +export { + ensureRuntimeHostPeerIdentity, + RuntimeHostPeerError, +} from '../transport/peer-native.js'; export { createRuntimeHostPeerClient, createRuntimeHostPeerClientFromEnvironment, diff --git a/packages/storage/src/__tests__/credential-store.test.ts b/packages/storage/src/__tests__/credential-store.test.ts index 9f6e195a3b..c39aa1b985 100644 --- a/packages/storage/src/__tests__/credential-store.test.ts +++ b/packages/storage/src/__tests__/credential-store.test.ts @@ -236,7 +236,7 @@ describe('FileCredentialStore compareAndSetSecret', () => { slug: string, kind: CredentialKind, expected: string | null, - value: string, + value: string | null, ): Promise { assert.ok(store.compareAndSetSecret, 'file store exposes the CAS capability'); return store.compareAndSetSecret(slug, kind, expected, value); @@ -312,6 +312,26 @@ describe('FileCredentialStore compareAndSetSecret', () => { }); }); + test('compare-and-delete removes only the credential version it observed', async () => { + await withTempDir(async (dir) => { + const a = createFileCredentialStore(dir); + const b = createFileCredentialStore(dir); + await a.setSecret('acct', 'oauth_token', 'tok-basis'); + + assert.deepEqual(await cas(a, 'acct', 'oauth_token', 'tok-basis', null), { + committed: true, + }); + assert.equal(await a.getSecret('acct', 'oauth_token'), null); + + await a.setSecret('acct', 'oauth_token', 'tok-newer'); + assert.deepEqual(await cas(b, 'acct', 'oauth_token', 'tok-basis', null), { + committed: false, + current: 'tok-newer', + }); + assert.equal(await b.getSecret('acct', 'oauth_token'), 'tok-newer'); + }); + }); + test('concurrent CAS from the same basis: exactly one wins, the loser returns the winner value', async () => { await withTempDir(async (dir) => { const a = createFileCredentialStore(dir); @@ -333,4 +353,62 @@ describe('FileCredentialStore compareAndSetSecret', () => { assert.equal(await reader.getSecret('acct', 'oauth_token'), winnerValue); }); }); + + test('revision CAS rejects an ABA cycle even when the secret value matches again', async () => { + await withTempDir(async (dir) => { + const stale = createFileCredentialStore(dir); + const concurrent = createFileCredentialStore(dir); + await stale.setSecret('acct', 'oauth_token', 'tok-A'); + assert.ok(stale.getSecretSnapshot); + assert.ok(stale.compareAndSetSecretRevision); + const basis = await stale.getSecretSnapshot('acct', 'oauth_token'); + + await concurrent.setSecret('acct', 'oauth_token', 'tok-B'); + await concurrent.setSecret('acct', 'oauth_token', 'tok-A'); + + const result = await stale.compareAndSetSecretRevision( + 'acct', + 'oauth_token', + basis.revision, + 'stale-rollback', + ); + assert.equal(result.committed, false); + assert.equal(result.current.value, 'tok-A'); + assert.notEqual(result.current.revision, basis.revision); + assert.equal(await stale.getSecret('acct', 'oauth_token'), 'tok-A'); + }); + }); + + test('revision CAS rejects a stale restore after an absent-state ABA cycle', async () => { + await withTempDir(async (dir) => { + const stale = createFileCredentialStore(dir); + const concurrent = createFileCredentialStore(dir); + await stale.setSecret('acct', 'oauth_token', 'tok-old'); + assert.ok(stale.getSecretSnapshot); + assert.ok(stale.compareAndSetSecretRevision); + const beforeDelete = await stale.getSecretSnapshot('acct', 'oauth_token'); + const deleted = await stale.compareAndSetSecretRevision( + 'acct', + 'oauth_token', + beforeDelete.revision, + null, + ); + assert.equal(deleted.committed, true); + if (!deleted.committed) return; + + await concurrent.setSecret('acct', 'oauth_token', 'tok-newer'); + await concurrent.deleteSecret('acct', 'oauth_token'); + + const result = await stale.compareAndSetSecretRevision( + 'acct', + 'oauth_token', + deleted.revision, + 'tok-old', + ); + assert.equal(result.committed, false); + assert.equal(result.current.value, null); + assert.notEqual(result.current.revision, deleted.revision); + assert.equal(await stale.getSecret('acct', 'oauth_token'), null); + }); + }); }); diff --git a/packages/storage/src/credential-store.ts b/packages/storage/src/credential-store.ts index 325227e8d5..8a47d12022 100644 --- a/packages/storage/src/credential-store.ts +++ b/packages/storage/src/credential-store.ts @@ -71,8 +71,20 @@ export const CREDENTIAL_SCHEMA_VERSION = 1; interface CredentialFile { version: number; values: Record; + /** Opaque per-entry generations. Entries remain after deletion so absence + * has an identity too and a stale writer cannot pass an ABA cycle. */ + revisions: Record; } +export interface CredentialSecretSnapshot { + readonly value: string | null; + readonly revision: string | null; +} + +export type CredentialRevisionCasResult = + | { readonly committed: true; readonly revision: string } + | { readonly committed: false; readonly current: CredentialSecretSnapshot }; + /** * Outcome of a compare-and-set write. * @@ -93,14 +105,17 @@ export type CredentialCasResult = export interface CredentialStore { getSecret(slug: string, kind: CredentialKind): Promise; + /** Read the value together with an opaque generation suitable for ABA-safe CAS. */ + getSecretSnapshot?(slug: string, kind: CredentialKind): Promise; setSecret(slug: string, kind: CredentialKind, value: string): Promise; /** Delete one kind, or — with no kind — every kind for the slug (e.g. a * connection being removed). */ deleteSecret(slug: string, kind?: CredentialKind): Promise; /** - * Optional compare-and-set write. Persist `value` for `(slug, kind)` only - * while the stored entry still equals `expected` — the basis the caller read - * before deciding to write. `expected: null` asserts the entry is absent. + * Optional compare-and-set mutation. Persist `value`, or delete the entry + * when `value` is null, only while the stored entry still equals `expected` + * — the basis the caller read before deciding to mutate. `expected: null` + * asserts the entry is absent. * * The basis check and the write run together under the same cross-process * lock as `setSecret`, so no concurrent writer can slip in between them; the @@ -116,8 +131,15 @@ export interface CredentialStore { slug: string, kind: CredentialKind, expected: string | null, - value: string, + value: string | null, ): Promise; + /** Compare by opaque generation rather than secret contents. */ + compareAndSetSecretRevision?( + slug: string, + kind: CredentialKind, + expectedRevision: string | null, + value: string | null, + ): Promise; } export function createFileCredentialStore(workspaceRoot: string): CredentialStore { @@ -131,19 +153,25 @@ class FileCredentialStore implements CredentialStore { return this.get(slug, toStoredKind(kind)); } + async getSecretSnapshot(slug: string, kind: CredentialKind): Promise { + const key = this.key(slug, toStoredKind(kind)); + const file = await this.readUnlocked(); + return snapshot(file, key); + } + setSecret(slug: string, kind: CredentialKind, value: string): Promise { return this.set(slug, toStoredKind(kind), value); } async deleteSecret(slug: string, kind?: CredentialKind): Promise { - await this.mutate((values) => { + await this.mutate((file) => { if (kind) { - delete values[this.key(slug, toStoredKind(kind))]; + mutateEntry(file, this.key(slug, toStoredKind(kind)), null); return; } // No kind: clear every kind for the slug in one read-modify-write. for (const storedKind of STORED_CREDENTIAL_KINDS) { - delete values[this.key(slug, storedKind)]; + mutateEntry(file, this.key(slug, storedKind), null); } }); } @@ -154,9 +182,7 @@ class FileCredentialStore implements CredentialStore { } private set(slug: string, kind: StoredCredentialKind, value: string): Promise { - return this.mutate((values) => { - values[this.key(slug, kind)] = value; - }); + return this.mutate((file) => mutateEntry(file, this.key(slug, kind), value)); } /** @@ -170,7 +196,7 @@ class FileCredentialStore implements CredentialStore { slug: string, kind: CredentialKind, expected: string | null, - value: string, + value: string | null, ): Promise { const key = this.key(slug, toStoredKind(kind)); return withCredentialFileLock(this.path, async () => { @@ -180,21 +206,38 @@ class FileCredentialStore implements CredentialStore { if (current !== expected) { return { committed: false, current }; } - file.values[key] = value; + mutateEntry(file, key, value); await this.write(file); return { committed: true }; }); } + compareAndSetSecretRevision( + slug: string, + kind: CredentialKind, + expectedRevision: string | null, + value: string | null, + ): Promise { + const key = this.key(slug, toStoredKind(kind)); + return withCredentialFileLock(this.path, async () => { + const file = await this.readUnlocked(); + const current = snapshot(file, key); + if (current.revision !== expectedRevision) return { committed: false, current }; + const revision = mutateEntry(file, key, value); + await this.write(file); + return { committed: true, revision }; + }); + } + /** * Read-modify-write the whole file under the cross-process lockfile. The lock * serializes concurrent calls on this instance and a second store instance / * process alike, so one mechanism covers both — no separate in-instance queue. */ - private mutate(apply: (values: Record) => void): Promise { + private mutate(apply: (file: CredentialFile) => void): Promise { return withCredentialFileLock(this.path, async () => { const file = await this.readUnlocked(); - apply(file.values); + apply(file); await this.write(file); }); } @@ -209,7 +252,11 @@ class FileCredentialStore implements CredentialStore { raw = await readFile(this.path, 'utf8'); } catch (error) { if ((error as { code?: string }).code === 'ENOENT') { - return { version: CREDENTIAL_SCHEMA_VERSION, values: {} }; + return { + version: CREDENTIAL_SCHEMA_VERSION, + values: {}, + revisions: {}, + }; } throw error; } @@ -233,7 +280,20 @@ class FileCredentialStore implements CredentialStore { throw new Error(`Corrupt credentials.json: value for "${k}" is not a string.`); } } - return { version: CREDENTIAL_SCHEMA_VERSION, values: values as Record }; + const revisions = parsed.revisions ?? {}; + if (revisions === null || typeof revisions !== 'object' || Array.isArray(revisions)) { + throw new Error('Corrupt credentials.json: `revisions` is not an object.'); + } + for (const [k, v] of Object.entries(revisions)) { + if (typeof v !== 'string' || v.length === 0) { + throw new Error(`Corrupt credentials.json: revision for "${k}" is invalid.`); + } + } + return { + version: CREDENTIAL_SCHEMA_VERSION, + values: values as Record, + revisions: revisions as Record, + }; } private write(file: CredentialFile): Promise { @@ -241,6 +301,21 @@ class FileCredentialStore implements CredentialStore { } } +function snapshot(file: CredentialFile, key: string): CredentialSecretSnapshot { + return { + value: file.values[key] ?? null, + revision: file.revisions[key] ?? null, + }; +} + +function mutateEntry(file: CredentialFile, key: string, value: string | null): string { + const revision = randomUUID(); + if (value === null) delete file.values[key]; + else file.values[key] = value; + file.revisions[key] = revision; + return revision; +} + /** * Create (or harden) the directory that holds a secret file: 0700, and * re-chmod a pre-existing looser dir so neither the secret nor the lock can sit