diff --git a/.changeset/respect-mcp-management-readiness.md b/.changeset/respect-mcp-management-readiness.md new file mode 100644 index 0000000000..98032fb55d --- /dev/null +++ b/.changeset/respect-mcp-management-readiness.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": patch +--- + +Respect workspace trust and configuration readiness when managing MCP servers. diff --git a/.changeset/sdk-mcp-management-cwd.md b/.changeset/sdk-mcp-management-cwd.md new file mode 100644 index 0000000000..ea6c807e58 --- /dev/null +++ b/.changeset/sdk-mcp-management-cwd.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code-sdk": patch +--- + +Add optional `cwd` and `verify` parameters to the global MCP management and authorization methods. diff --git a/packages/agent-core-v2/AGENTS.md b/packages/agent-core-v2/AGENTS.md index e9cd9b8092..875dda1a08 100644 --- a/packages/agent-core-v2/AGENTS.md +++ b/packages/agent-core-v2/AGENTS.md @@ -63,6 +63,12 @@ Business code must not `import 'node:fs'`, write SQL, hand-roll append-logs / at One accepted exception: `features/tower/protocol` manages the `.tower/` directory inside the *user's* repository (worktree slots, comms files, activity log) — workspace content, not engine state — and is a verbatim port of the v1 protocol whose semantics (atomic tmp+rename, real `git` CLI for worktrees/merges) are the feature. It keeps direct `node:fs` / `node:child_process` access; do not "modernize" it onto the Stores above without a dedicated migration. +## MCP management plane + +The App scope owns the process-wide MCP management surface, split across three domains: `mcpConfig` (`src/app/mcpConfig/` — the `[mcp]` config section, the layered mcp.json loader with per-entry origins, `IMcpConfigStore` as the single write point for the user-level `mcp.json`, `IMcpOAuthStore` credential persistence, and `IMcpOAuthService`, the process-wide OAuth orchestrator with credential events, single-flight refresh, and proactive refresh timers), `mcpRegistry` (`IMcpRegistryService` — the unified read view over the file layers and plugin manifests with `source`/`origin`/`mutable`; runtime-name collisions keep both entries, and runtime resolution ranks the file entry above plugin entries — a file entry wins by presence, even disabled, while a disabled plugin descriptor is treated as absent; a project layer joins the view only when the queried cwd itself is trusted, matching what the workspace runtime would load), and `mcpManagement` (`IMcpManagementService` — guarded CRUD (project-layer entries reject as read-only; plugin entries never block, so a user-level write may shadow a plugin), connection-test probes, the locator-addressed inspection/auth-status surface, and locator-addressed OAuth begin/complete/cancel/reset with ambiguity rejection plus an idle timeout that cancels abandoned flows). The engine services and the edge exposure (kap-server routes, klient facade) are ungated. On the Workspace side, `workspaceMcpConfig` merges the same sources (same file-over-plugin precedence), watches the files and plugin reloads, follows the store's `onDidWrite` for immediate management-plane reloads, and publishes fingerprint diffs that `workspaceMcp` applies to the handler-shared `McpConnectionManager` (tombstones for removals); `workspaceMcp` also subscribes the OAuth service's credential events to reconnect affected entries. Session overlays (`session/mcp`) keep caller-injected ephemeral servers session-local. + +Name-collision precedence is a deliberate, documented divergence from v1: v1 ranks an enabled plugin above the file layers (#2858), v2 keeps its historical file-over-plugin order, and each engine's management guards follow its own winner — the same `mcp.json` plus plugin set can therefore resolve a collision differently per engine. Do not "re-align" one side without an explicit product decision. + ## Session index `ISessionIndex` (`src/app/sessionIndex/`, App scope) serves session list/resume reads over two paths: the authoritative directory scan (`sessionIndexSource`, always correct, linear) and the minidb-backed derived read model (`IQueryStore` at `/cache/query-store`, keyset-paged, `O(log N + limit)`), gated by the `persistence_minidb_readmodel` flag (default ON; roll back via `KIMI_CODE_EXPERIMENTAL_PERSISTENCE_MINIDB_READMODEL=false` or the `[experimental]` config section). The read model has an explicit lifecycle — `uninitialized → preparing → ready/degraded` via `prepare()`/`status()`; reads while preparing answer from the authoritative store immediately and fold the `ISessionIndexMirror` queue in for read-your-writes; the first list shares one single-flight authoritative scan with the initial projection. The query-store is structural-only — text-index definitions are rejected at definition level, so session operations never touch the global full-text index (`/search-index`, owned by kap-server's search surface). diff --git a/packages/agent-core-v2/src/_base/event.ts b/packages/agent-core-v2/src/_base/event.ts index a80483c2ca..1c53bb85ae 100644 --- a/packages/agent-core-v2/src/_base/event.ts +++ b/packages/agent-core-v2/src/_base/event.ts @@ -112,6 +112,24 @@ export type IWaitUntilData = Omit; export class AsyncEmitter extends Emitter { private _asyncDeliveryQueue?: LinkedList<[(event: T) => void, IWaitUntilData]>; + async fireAsyncConcurrent(data: IWaitUntilData, signal: AbortSignal): Promise { + if (this.isDisposed || this._listeners === undefined || signal.aborted) { + return; + } + const snapshot = Array.from(this._listeners); + await Promise.all( + snapshot.map((entry) => + this.deliverAsync( + (event) => { + entry.listener.call(entry.thisArg, event); + }, + data, + signal, + ), + ), + ); + } + async fireAsync(data: IWaitUntilData, signal: AbortSignal): Promise { if (this.isDisposed || this._listeners === undefined) { return; @@ -129,32 +147,37 @@ export class AsyncEmitter extends Emitter { while (this._asyncDeliveryQueue.size > 0 && !signal.aborted) { const [deliver, eventData] = this._asyncDeliveryQueue.shift()!; - const thenables: Promise[] = []; - - const event = { - ...eventData, - signal, - waitUntil: (p: Promise): void => { - if (Object.isFrozen(thenables)) { - throw new Error('waitUntil can NOT be called asynchronously'); - } - thenables.push(p); - }, - } as T; - - try { - deliver(event); - } catch (error) { - onUnexpectedError(error); - continue; - } + await this.deliverAsync(deliver, eventData, signal); + } + } - void Object.freeze(thenables); - const settled = await Promise.allSettled(thenables); - for (const result of settled) { - if (result.status === 'rejected') { - onUnexpectedError(result.reason); + private async deliverAsync( + deliver: (event: T) => void, + data: IWaitUntilData, + signal: AbortSignal, + ): Promise { + const thenables: Promise[] = []; + const event = { + ...data, + signal, + waitUntil: (p: Promise): void => { + if (Object.isFrozen(thenables)) { + throw new Error('waitUntil can NOT be called asynchronously'); } + thenables.push(p); + }, + } as T; + try { + deliver(event); + } catch (error) { + onUnexpectedError(error); + return; + } + void Object.freeze(thenables); + const settled = await Promise.allSettled(thenables); + for (const result of settled) { + if (result.status === 'rejected') { + onUnexpectedError(result.reason); } } } diff --git a/packages/agent-core-v2/src/_base/utils/paths.ts b/packages/agent-core-v2/src/_base/utils/paths.ts index 757fcff708..bf035216f0 100644 --- a/packages/agent-core-v2/src/_base/utils/paths.ts +++ b/packages/agent-core-v2/src/_base/utils/paths.ts @@ -1,9 +1,34 @@ import nodePath from 'node:path'; +import { isAbsolute, normalize, resolve } from 'pathe'; + +import { workspaceRootKey } from './workdir-slug'; + function normalizeSlashes(p: string): string { return p.replaceAll('\\', '/'); } +export function isWindowsAbsolutePath(value: string): boolean { + return /^[A-Za-z]:[\\/]/.test(value) || /^[\\/]{2}[^\\/]+[\\/][^\\/]+/.test(value); +} + +export function resolvePath(base: string, value: string): string { + if (isWindowsAbsolutePath(base)) { + return nodePath.win32.resolve(base, value).replaceAll('\\', '/'); + } + if (isWindowsAbsolutePath(value)) { + return nodePath.win32.resolve(value).replaceAll('\\', '/'); + } + return isAbsolute(value) ? normalize(value) : resolve(base, value); +} + +export function canonicalWorkspaceRoot(cwd: string): string { + const resolved = isWindowsAbsolutePath(cwd) + ? nodePath.win32.resolve(cwd).replaceAll('\\', '/') + : resolve(cwd); + return workspaceRootKey(resolved) || resolved; +} + export interface UpwardRootPathApi { resolve(dir: string): string; dirname(dir: string): string; diff --git a/packages/agent-core-v2/src/app/mcpConfig/configLoader.ts b/packages/agent-core-v2/src/app/mcpConfig/configLoader.ts new file mode 100644 index 0000000000..574f20ece6 --- /dev/null +++ b/packages/agent-core-v2/src/app/mcpConfig/configLoader.ts @@ -0,0 +1,194 @@ +import { dirname, join, normalize } from 'pathe'; + +import { resolveKimiHome } from '#/app/bootstrap/bootstrap'; +import { findGitWorkTree } from '#/app/git/workTree'; +import { resolvePath } from '#/_base/utils/paths'; +import { ErrorCodes, Error2 } from '#/errors'; +import { McpServerConfigSchema, type McpServerConfig } from '#/mcpCore/config-schema'; +import type { IHostFileSystem } from '#/os/interface/hostFileSystem'; +import { OsFsErrors, HostFsError } from '#/os/interface/hostFsErrors'; + +export interface McpJsonPaths { + readonly user: string; + readonly projectRoot: string; + readonly project: string; +} + +export interface ResolveMcpJsonPathsInput { + readonly fs: IHostFileSystem; + readonly cwd: string; + readonly homeDir?: string; +} + +export async function resolveMcpJsonPaths(input: ResolveMcpJsonPathsInput): Promise { + const start = normalize(input.cwd); + const projectRoot = (await findGitWorkTree(input.fs, start))?.root ?? start; + + return { + user: join(resolveKimiHome(input.homeDir), 'mcp.json'), + projectRoot: join(projectRoot, '.mcp.json'), + project: join(input.cwd, '.kimi-code', 'mcp.json'), + }; +} + +export interface LoadMcpServersInput { + readonly fs: IHostFileSystem; + readonly cwd: string; + readonly homeDir?: string; + readonly includeProject?: boolean; +} + +export interface LoadMcpServersDetailedResult { + /** Later layers override earlier ones with the same key. */ + readonly servers: Record; + /** The file each effective entry was last defined in. */ + readonly origins: Record; +} + +export async function loadMcpServers( + input: LoadMcpServersInput, +): Promise> { + return (await loadMcpServersDetailed(input)).servers; +} + +/** + * {@link loadMcpServers} plus the defining-file origin of every effective + * entry, for management surfaces that show where a server came from. + */ +export async function loadMcpServersDetailed( + input: LoadMcpServersInput, +): Promise { + const paths = await resolveMcpJsonPaths(input); + if (input.includeProject === false) { + const user = await readMcpJson(input.fs, paths.user); + return { servers: user, origins: mapValuesToPath(user, paths.user) }; + } + const layers: readonly [path: string, servers: Record][] = + await Promise.all([ + readMcpJson(input.fs, paths.user), + readMcpJson(input.fs, paths.projectRoot, { stdioCwdBase: dirname(paths.projectRoot) }), + readMcpJson(input.fs, paths.project), + ]).then(([user, projectRoot, project]) => [ + [paths.user, user], + [paths.projectRoot, projectRoot], + [paths.project, project], + ]); + const servers: Record = Object.create(null); + const origins: Record = Object.create(null); + for (const [path, layer] of layers) { + for (const [name, config] of Object.entries(layer)) { + servers[name] = config; + origins[name] = path; + } + } + return { servers, origins }; +} + +interface ReadMcpJsonOptions { + readonly stdioCwdBase?: string; +} + +async function readMcpJson( + fs: IHostFileSystem, + filePath: string, + options: ReadMcpJsonOptions = {}, +): Promise> { + let text: string; + try { + text = await fs.readText(filePath); + } catch (error: unknown) { + if (isFileNotFound(error)) return {}; + throw new Error2( + ErrorCodes.CONFIG_INVALID, + `Failed to read ${filePath}: ${describeError(error)}`, + { + cause: error, + }, + ); + } + + if (text.trim().length === 0) return {}; + + let data: unknown; + try { + data = JSON.parse(text); + } catch (error: unknown) { + throw new Error2( + ErrorCodes.CONFIG_INVALID, + `Invalid JSON in ${filePath}: ${describeError(error)}`, + { + cause: error, + }, + ); + } + + try { + return normalizeMcpServers(parseMcpJsonServers(data), options); + } catch (error: unknown) { + throw new Error2( + ErrorCodes.CONFIG_INVALID, + `Invalid MCP server config in ${filePath}: ${describeError(error)}`, + { + cause: error, + }, + ); + } +} + +function parseMcpJsonServers(data: unknown): Record { + if (!isRecord(data)) { + throw new Error('expected a JSON object'); + } + if (!('mcpServers' in data)) return {}; + const raw = data['mcpServers']; + if (!isRecord(raw)) { + throw new Error('"mcpServers" must be an object'); + } + return Object.fromEntries( + Object.entries(raw).map(([name, value]) => [name, McpServerConfigSchema.parse(value)]), + ); +} + +function normalizeMcpServers( + servers: Record, + options: ReadMcpJsonOptions, +): Record { + const stdioCwdBase = options.stdioCwdBase; + if (stdioCwdBase === undefined) return servers; + + return Object.fromEntries( + Object.entries(servers).map(([name, config]) => [ + name, + normalizeStdioCwd(config, stdioCwdBase), + ]), + ); +} + +function normalizeStdioCwd(config: McpServerConfig, cwdBase: string): McpServerConfig { + if (config.transport !== 'stdio') return config; + const cwd = config.cwd === undefined ? cwdBase : resolvePath(cwdBase, config.cwd); + return { ...config, cwd }; +} + +function mapValuesToPath( + servers: Record, + path: string, +): Record { + const origins: Record = Object.create(null); + for (const name of Object.keys(servers)) { + origins[name] = path; + } + return origins; +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +function isFileNotFound(error: unknown): boolean { + return error instanceof HostFsError && error.code === OsFsErrors.codes.OS_FS_NOT_FOUND; +} + +function describeError(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} diff --git a/packages/agent-core-v2/src/app/mcpConfig/configStore.ts b/packages/agent-core-v2/src/app/mcpConfig/configStore.ts new file mode 100644 index 0000000000..eaa150fafb --- /dev/null +++ b/packages/agent-core-v2/src/app/mcpConfig/configStore.ts @@ -0,0 +1,225 @@ +import { join } from 'pathe'; + +import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; +import { Disposable } from '#/_base/di/lifecycle'; +import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; +import { AsyncEmitter, type Event, type IWaitUntil } from '#/_base/event'; +import { IBootstrapService } from '#/app/bootstrap/bootstrap'; +import { LifecycleScope } from '#/app/scopes'; +import { ErrorCodes, Error2 } from '#/errors'; +import { McpServerConfigSchema, type McpServerConfig } from '#/mcpCore/config-schema'; +import { IFileSystemStorageService } from '#/persistence/interface/storage'; + +export type GlobalMcpServerConfig = McpServerConfig & { readonly name: string }; + +export type McpConfigWriteEvent = IWaitUntil; + +export interface IMcpConfigStore { + readonly _serviceBrand: undefined; + readonly path: string; + readonly onDidWrite: Event; + list(): Promise; + get(name: string): Promise; + add(server: GlobalMcpServerConfig): Promise; + update(server: GlobalMcpServerConfig): Promise; + remove(name: string): Promise; +} + +export const IMcpConfigStore: ServiceIdentifier = + createDecorator('mcpConfigStore'); + +interface McpConfigFile { + readonly raw: Record; + readonly rawServers: Record; + readonly servers: readonly GlobalMcpServerConfig[]; +} + +const CONFIG_SCOPE = ''; +const MCP_CONFIG_KEY = 'mcp.json'; + +const textEncoder = new TextEncoder(); +const textDecoder = new TextDecoder('utf-8', { ignoreBOM: true }); + +export class McpConfigStore extends Disposable implements IMcpConfigStore { + declare readonly _serviceBrand: undefined; + + readonly path: string; + + private readonly writeEmitter = this._register(new AsyncEmitter()); + readonly onDidWrite: Event = this.writeEmitter.event; + private mutationTail: Promise = Promise.resolve(); + private writePending = false; + + constructor( + @IFileSystemStorageService private readonly storage: IFileSystemStorageService, + @IBootstrapService bootstrap: IBootstrapService, + ) { + super(); + this.path = join(bootstrap.homeDir, MCP_CONFIG_KEY); + } + + async list(): Promise { + return (await this.read()).servers; + } + + async get(name: string): Promise { + const normalizedName = normalizeServerName(name); + const server = (await this.read()).servers.find((entry) => entry.name === normalizedName); + if (server !== undefined) return server; + throw serverNotFound(normalizedName); + } + + add(server: GlobalMcpServerConfig): Promise { + return this.mutate(async () => { + const normalized = parseServerInput(server); + const file = await this.read(); + if (Object.hasOwn(file.rawServers, normalized.name)) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + `MCP server "${normalized.name}" already exists`, + ); + } + await this.write(file, { + ...file.rawServers, + [normalized.name]: persistedEntry(normalized), + }); + return this.list(); + }); + } + + update(server: GlobalMcpServerConfig): Promise { + return this.mutate(async () => { + const normalized = parseServerInput(server); + const file = await this.read(); + if (!Object.hasOwn(file.rawServers, normalized.name)) { + throw serverNotFound(normalized.name); + } + await this.write(file, { + ...file.rawServers, + [normalized.name]: persistedEntry(normalized), + }); + return this.list(); + }); + } + + remove(name: string): Promise { + return this.mutate(async () => { + const normalizedName = normalizeServerName(name); + const file = await this.read(); + if (!Object.hasOwn(file.rawServers, normalizedName)) return file.servers; + const nextServers = Object.fromEntries( + Object.entries(file.rawServers).filter(([entryName]) => entryName !== normalizedName), + ); + await this.write(file, nextServers); + return this.list(); + }); + } + + private mutate(work: () => Promise): Promise { + const tail = this.mutationTail.catch(() => undefined).then(work); + this.mutationTail = tail.then( + () => undefined, + () => undefined, + ); + return tail.then(async (result) => { + if (!this.writePending) return result; + this.writePending = false; + await this.writeEmitter.fireAsyncConcurrent({}, NO_ABORT); + return result; + }); + } + + private async read(): Promise { + let bytes: Uint8Array | undefined; + try { + bytes = await this.storage.read(CONFIG_SCOPE, MCP_CONFIG_KEY); + } catch (error: unknown) { + throw configError(`Failed to read ${this.path}: ${describeError(error)}`, error); + } + if (bytes === undefined) { + return { raw: {}, rawServers: {}, servers: [] }; + } + + const text = textDecoder.decode(bytes); + if (text.trim().length === 0) { + return { raw: {}, rawServers: {}, servers: [] }; + } + + let parsed: unknown; + try { + parsed = JSON.parse(text) as unknown; + } catch (error: unknown) { + throw configError(`Invalid JSON in ${this.path}: ${describeError(error)}`, error); + } + if (!isRecord(parsed)) { + throw configError(`Invalid MCP config in ${this.path}: expected a JSON object`); + } + const rawServersValue = parsed['mcpServers']; + if (rawServersValue !== undefined && !isRecord(rawServersValue)) { + throw configError(`Invalid MCP config in ${this.path}: "mcpServers" must be an object`); + } + const rawServers = rawServersValue ?? {}; + const servers = Object.entries(rawServers).map(([name, value]) => parseServer(name, value)); + return { raw: parsed, rawServers, servers }; + } + + private async write(file: McpConfigFile, rawServers: Record): Promise { + const text = `${JSON.stringify({ ...file.raw, mcpServers: rawServers }, null, 2)}\n`; + await this.storage.write(CONFIG_SCOPE, MCP_CONFIG_KEY, textEncoder.encode(text), { + atomic: true, + }); + this.writePending = true; + } +} + +const NO_ABORT = new AbortController().signal; + +function parseServerInput(server: GlobalMcpServerConfig): GlobalMcpServerConfig { + return parseServer(normalizeServerName(server.name), server); +} + +function parseServer(name: string, value: unknown): GlobalMcpServerConfig { + const result = McpServerConfigSchema.safeParse(value); + if (!result.success) { + throw configError( + `Invalid MCP server "${name}" in global config: ${result.error.message}`, + result.error, + ); + } + return { name, ...result.data }; +} + +function persistedEntry(server: GlobalMcpServerConfig): McpServerConfig { + const { name: _name, ...entry } = server; + return entry; +} + +export function normalizeServerName(name: string): string { + const normalized = name.trim(); + if (normalized.length > 0) return normalized; + throw new Error2(ErrorCodes.REQUEST_INVALID, 'MCP server name cannot be empty'); +} + +function serverNotFound(name: string): Error2 { + return new Error2(ErrorCodes.MCP_SERVER_NOT_FOUND, `MCP server "${name}" was not found`); +} + +function configError(message: string, cause?: unknown): Error2 { + return new Error2(ErrorCodes.CONFIG_INVALID, message, { cause }); +} + +function describeError(error: unknown): string { + return error instanceof Error ? error.message : String(error); +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value); +} + +registerScopedService( + LifecycleScope.App, + IMcpConfigStore, + McpConfigStore, + ScopeActivation.OnDemand, + 'mcpConfig', +); diff --git a/packages/agent-core-v2/src/app/mcpConfig/oauthService.ts b/packages/agent-core-v2/src/app/mcpConfig/oauthService.ts new file mode 100644 index 0000000000..732f8fd0e5 --- /dev/null +++ b/packages/agent-core-v2/src/app/mcpConfig/oauthService.ts @@ -0,0 +1,43 @@ +import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; +import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; +import { ILogService } from '#/_base/log/log'; +import { IAgentIdentity } from '#/app/agentIdentity/agentIdentity'; +import { LifecycleScope } from '#/app/scopes'; +import { McpOAuthService } from '#/mcpCore/oauth/service'; + +import { IMcpOAuthStore } from './oauthStore'; + +export const IMcpOAuthService: ServiceIdentifier = + createDecorator('mcpOAuthService'); + +export class AppMcpOAuthService extends McpOAuthService { + constructor( + @IMcpOAuthStore store: IMcpOAuthStore, + @IAgentIdentity identity: IAgentIdentity, + @ILogService log: ILogService, + ) { + super({ + store, + resolveClientName: () => identity.current().slug, + log, + }); + void identity + .resolved() + .then(() => { + const sweep = this.sweepProactiveRefresh(); + this.trackBackgroundTask(sweep); + return sweep; + }) + .catch((error: unknown) => { + log.warn(`mcp oauth proactive-refresh sweep failed: ${String(error)}`); + }); + } +} + +registerScopedService( + LifecycleScope.App, + IMcpOAuthService, + AppMcpOAuthService, + ScopeActivation.OnDemand, + 'mcpConfig', +); diff --git a/packages/agent-core-v2/src/app/mcpConfig/oauthStore.ts b/packages/agent-core-v2/src/app/mcpConfig/oauthStore.ts index 036e3c2f99..aa44809342 100644 --- a/packages/agent-core-v2/src/app/mcpConfig/oauthStore.ts +++ b/packages/agent-core-v2/src/app/mcpConfig/oauthStore.ts @@ -29,6 +29,9 @@ export function createMcpOAuthStore(docs: IAtomicDocumentStore): McpOAuthStore { remove(key) { return docs.delete(CREDENTIALS_SCOPE, key); }, + list(prefix) { + return docs.list(CREDENTIALS_SCOPE, prefix); + }, }; } @@ -52,6 +55,10 @@ export class McpOAuthStoreAdapter implements IMcpOAuthStore { remove(key: string): Promise { return this.delegate.remove(key); } + + list(prefix?: string): Promise { + return this.delegate.list(prefix); + } } registerScopedService( diff --git a/packages/agent-core-v2/src/app/mcpManagement/mcpManagement.ts b/packages/agent-core-v2/src/app/mcpManagement/mcpManagement.ts new file mode 100644 index 0000000000..dbded9e44a --- /dev/null +++ b/packages/agent-core-v2/src/app/mcpManagement/mcpManagement.ts @@ -0,0 +1,175 @@ +import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; + +import type { McpServerConfig } from '#/mcpCore/config-schema'; +import type { McpServerConfigView } from '#/mcpCore/configView'; +import type { + McpRegistryPluginOrigin, + McpRegistryQuery, + McpServerSource, +} from '#/app/mcpRegistry/mcpRegistry'; + +export type GlobalMcpServerConfig = McpServerConfig & { readonly name: string }; + +export interface McpManagedServer { + readonly name: string; + /** + * Mutable (user-level) entries carry the full config so edit UIs can + * prefill values; read-only entries are redacted to sorted key lists + * (`envKeys` / `headerKeys`) and never disclose secret values. + */ + readonly config: McpServerConfig | McpServerConfigView; + readonly source: McpServerSource; + readonly origin: string; + readonly mutable: boolean; + readonly plugin?: McpRegistryPluginOrigin; +} + +export interface McpServerTestTarget { + /** Registry-resolved by name when `server` is omitted. */ + readonly name?: string; + /** Inline config probes as-is — nothing has to be saved first. */ + readonly server?: GlobalMcpServerConfig; + /** Project layers join the resolution; also the stdio working directory. */ + readonly cwd?: string; +} + +export interface McpServerTestResult { + readonly success: boolean; + readonly output: string; +} + +/** + * Stable address of one catalog entry: a global (file-layer) server by name, + * or a plugin server by plugin id + manifest-local server name. + */ +export type McpServerLocator = + | { readonly source: 'global'; readonly name: string } + | { readonly source: 'plugin'; readonly pluginId: string; readonly serverName: string }; + +/** Locator-addressed catalog entry with the redacted config view. */ +export interface McpServerDescriptor { + /** `global:` / `plugin::`, URL-encoded. */ + readonly serverId: string; + readonly locator: McpServerLocator; + readonly runtimeName: string; + /** Canonical credential URL for remote servers; undefined for stdio. */ + readonly canonicalUrl?: string; + readonly origin: McpServerSource; + readonly config: McpServerConfigView; + readonly enabled: boolean; + readonly editable: boolean; +} + +export type McpServerAuthState = + | 'not-applicable' + | 'bearer-token' + | 'oauth-required' + | 'oauth-authorized' + | 'oauth-expired' + | 'unavailable'; + +export interface McpServerInspection extends McpServerDescriptor { + readonly authStatus: McpServerAuthState; + readonly checkedAt?: number; + readonly error?: string; +} + +export interface McpServerAuthStatus { + readonly name: string; + readonly authStatus: McpServerAuthState; +} + +export type McpServerAuthBeginResult = + | { + readonly status: 'authorization-required'; + readonly flowId: string; + readonly authorizationUrl: string; + } + | { readonly status: 'already-authorized' }; + +export interface McpServerAuthFlowHandle { + readonly flowId: string; + readonly timeoutMs?: number; +} + +export interface McpAuthStatusQuery extends McpRegistryQuery { + /** + * Omitted preserves implicit OAuth detection, `false` stays offline, and + * `true` verifies every OAuth candidate through a real connection. + */ + readonly verify?: boolean; +} + +export interface IMcpManagementService { + readonly _serviceBrand: undefined; + + listServers(query?: McpRegistryQuery): Promise; + + getServer(name: string, query?: McpRegistryQuery): Promise; + + /** Writes the user-level file; rejects read-only collisions. Returns the refreshed list. */ + addServer( + server: GlobalMcpServerConfig, + query?: McpRegistryQuery, + ): Promise; + + /** Updates an existing user-level entry; rejects read-only collisions. Returns the refreshed list. */ + updateServer( + server: GlobalMcpServerConfig, + query?: McpRegistryQuery, + ): Promise; + + /** Removes a user-level entry; rejects read-only collisions. Returns the refreshed list. */ + removeServer(name: string, query?: McpRegistryQuery): Promise; + + testServer(target: McpServerTestTarget): Promise; + + /** + * Legacy auth-status surface: per-server OAuth state over the registry + * catalog. Omitted preserves the legacy implicit-OAuth probe for unpinned + * servers without stored credentials; `verify: false` is fully offline; + * `verify: true` probes every candidate. Probes may refresh or invalidate + * stored credentials and broadcast the events. + */ + listAuthStatuses(query?: McpAuthStatusQuery): Promise; + + /** + * The locator-addressed catalog plus a batched real-connection probe of + * every OAuth candidate; a probe that hits an expired grant may refresh or + * invalidate stored credentials and broadcast the events. A runtime name + * shared by enabled entries cannot be probed (or credentialed) + * unambiguously and reports `unavailable`. + */ + inspectServers( + targets?: readonly McpServerLocator[], + query?: McpRegistryQuery, + ): Promise; + + /** + * Resolve a legacy name-only auth target: exactly one enabled entry may + * own the runtime name — under a collision the caller cannot tell which + * credential the flow acts on, so it rejects instead of guessing. + */ + resolveServerByName(name: string, query?: McpRegistryQuery): Promise; + + /** Begin an interactive OAuth flow for a remote server. */ + beginServerAuth( + locator: McpServerLocator, + query?: McpRegistryQuery, + ): Promise; + + /** Await the browser callback and finish the code exchange. Unknown flow → request.invalid. */ + completeServerAuth( + handle: McpServerAuthFlowHandle, + options?: { readonly signal?: AbortSignal }, + ): Promise; + + /** Tear down a flow without finishing it; unknown flows are ignored. */ + cancelServerAuth(handle: Pick): Promise; + + /** Clear stored credentials; the invalidation event reaches live sessions. */ + resetServerAuth(locator: McpServerLocator, query?: McpRegistryQuery): Promise; +} + +export const IMcpManagementService: ServiceIdentifier = + createDecorator('mcpManagementService'); diff --git a/packages/agent-core-v2/src/app/mcpManagement/mcpManagementService.ts b/packages/agent-core-v2/src/app/mcpManagement/mcpManagementService.ts new file mode 100644 index 0000000000..0f30c9cfa6 --- /dev/null +++ b/packages/agent-core-v2/src/app/mcpManagement/mcpManagementService.ts @@ -0,0 +1,653 @@ +import { randomUUID } from 'node:crypto'; + +import { normalize } from 'pathe'; + +import { Disposable } from '#/_base/di/lifecycle'; +import { LifecycleScope } from '#/app/scopes'; +import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; +import { ILogService } from '#/_base/log/log'; + +import { ErrorCodes, Error2 } from '#/errors'; +import { McpConnectionManager } from '#/mcpCore/connection-manager'; +import { McpServerConfigSchema, type McpServerConfig } from '#/mcpCore/config-schema'; +import { toMcpServerConfigView } from '#/mcpCore/configView'; +import { + AlreadyAuthorizedError, + type BeginAuthorizationResult, + type McpOAuthService, + type McpOAuthTokenState, +} from '#/mcpCore/oauth/service'; +import { canonicalMcpOAuthResource } from '#/mcpCore/oauth/store'; +import { IHostEnvironment } from '#/os/interface/hostEnvironment'; +import { IHostProcessService } from '#/os/interface/hostProcess'; +import { LocalRuntime } from '#/runtime/localRuntime'; +import { RuntimeRegistry } from '#/runtime/runtimeRegistry'; +import { IAgentIdentity } from '#/app/agentIdentity/agentIdentity'; +import { IConfigService } from '#/app/config/config'; +import { MCP_SECTION, type McpSection } from '#/app/mcpConfig/configSection'; +import { IMcpConfigStore, normalizeServerName } from '#/app/mcpConfig/configStore'; +import { IMcpOAuthService } from '#/app/mcpConfig/oauthService'; +import { + IMcpRegistryService, + type McpRegistryEntry, + type McpRegistryQuery, +} from '#/app/mcpRegistry/mcpRegistry'; +import { + IRuntimeResolver, + IWorkspaceInstanceManager, +} from '#/workspace/workspaceInstance/workspaceInstanceManager'; + +import { + IMcpManagementService, + type GlobalMcpServerConfig, + type McpAuthStatusQuery, + type McpManagedServer, + type McpServerAuthBeginResult, + type McpServerAuthFlowHandle, + type McpServerAuthState, + type McpServerAuthStatus, + type McpServerDescriptor, + type McpServerInspection, + type McpServerLocator, + type McpServerTestResult, + type McpServerTestTarget, +} from './mcpManagement'; + +const DEFAULT_AUTH_TIMEOUT_MS = 15 * 60_000; +const AUTH_FLOW_IDLE_TIMEOUT_MS = 15 * 60_000; +const MAX_AUTH_TIMEOUT_MS = 2 ** 31 - 1; + +export class McpManagementService extends Disposable implements IMcpManagementService { + declare readonly _serviceBrand: undefined; + + private readonly authFlows = new Map< + string, + { flow: BeginAuthorizationResult; idleTimer: NodeJS.Timeout } + >(); + + constructor( + @IMcpRegistryService private readonly registry: IMcpRegistryService, + @IMcpConfigStore private readonly store: IMcpConfigStore, + @IMcpOAuthService private readonly oauth: McpOAuthService, + @IConfigService private readonly config: IConfigService, + @IAgentIdentity private readonly identity: IAgentIdentity, + @IRuntimeResolver private readonly runtimeResolver: IRuntimeResolver, + @IWorkspaceInstanceManager private readonly workspaceInstances: IWorkspaceInstanceManager, + @IHostEnvironment private readonly hostEnvironment: IHostEnvironment, + @IHostProcessService private readonly hostProcess: IHostProcessService, + @ILogService private readonly log: ILogService, + ) { + super(); + } + + async listServers(query: McpRegistryQuery = {}): Promise { + return (await this.registry.list(query)).map(toManagedServer); + } + + async getServer(name: string, query: McpRegistryQuery = {}): Promise { + return toManagedServer(await this.registry.get(name, query)); + } + + async addServer( + server: GlobalMcpServerConfig, + query: McpRegistryQuery = {}, + ): Promise { + const name = normalizeServerName(server.name); + await this.guardMutation(name, query); + await this.store.add({ ...server, name }); + return this.listServers(query); + } + + async updateServer( + server: GlobalMcpServerConfig, + query: McpRegistryQuery = {}, + ): Promise { + const name = normalizeServerName(server.name); + await this.guardMutation(name, query); + await this.store.update({ ...server, name }); + return this.listServers(query); + } + + async removeServer( + name: string, + query: McpRegistryQuery = {}, + ): Promise { + const normalized = normalizeServerName(name); + await this.guardMutation(normalized, query); + await this.store.remove(normalized); + return this.listServers(query); + } + + async testServer(target: McpServerTestTarget): Promise { + await this.waitForReadiness(); + const resolved = await this.resolveTestTarget(target); + return this.withProbe(resolved, target.cwd, (manager) => + standaloneTestResult(resolved.name, manager), + ); + } + + private async guardMutation(name: string, query: McpRegistryQuery): Promise { + const matches = (await this.registry.list(query)).filter((entry) => entry.name === name); + for (const entry of matches) { + if (entry.source === 'global' && !entry.mutable) throwReadOnlyMcpServer(entry); + } + } + + private async resolveTestTarget(target: McpServerTestTarget): Promise { + const { name, server, cwd } = target; + if (server !== undefined) { + if (name !== undefined && name !== server.name) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + 'Pass either an MCP server name or an inline server config, not both', + ); + } + const parsed = McpServerConfigSchema.safeParse(server); + if (!parsed.success) { + throw new Error2( + ErrorCodes.CONFIG_INVALID, + `Invalid MCP server "${server.name}": ${parsed.error.message}`, + ); + } + return { name: server.name, ...parsed.data }; + } + if (name === undefined) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + 'Pass an MCP server name or an inline server config', + ); + } + const matches = (await this.registry.list({ cwd })).filter((entry) => entry.name === name); + if (matches.length === 0) { + throw new Error2(ErrorCodes.MCP_SERVER_NOT_FOUND, `MCP server "${name}" was not found`); + } + const enabled = matches.filter((entry) => entry.config.enabled !== false); + if (enabled.length > 1) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + `MCP runtime name "${name}" is shared by multiple enabled servers`, + ); + } + const entry = enabled[0] ?? matches[0]!; + return { name: entry.name, ...entry.config }; + } + + private async withProbe( + server: GlobalMcpServerConfig, + cwd: string | undefined, + inspect: (manager: McpConnectionManager) => T, + ): Promise { + await this.waitForReadiness(); + const section = this.config.get(MCP_SECTION); + let workspaceId: string | undefined; + let stdioCwd = cwd; + let runtimeResolver = this.runtimeResolver; + let transientRuntimes: RuntimeRegistry | undefined; + if (server.transport === 'stdio') { + stdioCwd = normalize(cwd ?? process.cwd()); + const workspace = this.workspaceInstances.findContaining(stdioCwd); + if (workspace !== undefined) { + workspaceId = workspace.id; + } else { + const runtimeId = server.runtime_id; + if (runtimeId !== undefined && runtimeId !== 'local') { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + `Cannot probe MCP server "${server.name}" with runtime_id "${runtimeId}": no materialized workspace contains ${stdioCwd}, and an out-of-workspace probe only supports the local runtime`, + ); + } + await this.hostEnvironment.ready; + workspaceId = `mcp-probe-${randomUUID()}`; + transientRuntimes = new RuntimeRegistry(workspaceId); + transientRuntimes.register( + new LocalRuntime( + workspaceId, + this.hostEnvironment, + undefined, + this.hostProcess, + undefined, + undefined, + ), + ); + runtimeResolver = { + _serviceBrand: undefined, + inspect: (binding) => transientRuntimes!.inspect(binding), + acquire: (binding, required) => transientRuntimes!.acquire(binding, required), + }; + } + } + const manager = new McpConnectionManager({ + log: this.log, + stdioCwd, + runtimeResolver, + workspaceId, + runtimeId: workspaceId === undefined ? undefined : 'local', + oauthService: this.oauth, + resolveClientName: () => this.identity.current().slug, + resolveDefaultTimeouts: () => ({ + startupTimeoutMs: section?.startupTimeoutMs, + toolTimeoutMs: section?.toolTimeoutMs, + }), + }); + try { + await manager.connectAll({ [server.name]: mcpConfigWithoutName(server) }); + return inspect(manager); + } finally { + try { + await manager.shutdown(); + } finally { + await transientRuntimes?.dispose(); + } + } + } + + async listAuthStatuses(query: McpAuthStatusQuery = {}): Promise { + await this.waitForReadiness(); + const entries = await this.registry.list({ cwd: query.cwd }); + return Promise.all( + entries.map(async (entry) => ({ + name: entry.name, + authStatus: await this.serverAuthState(entry, query.cwd, query.verify), + })), + ); + } + + async inspectServers( + targets?: readonly McpServerLocator[], + query: McpRegistryQuery = {}, + ): Promise { + await this.waitForReadiness(); + const catalog = await this.serverDescriptors(query); + const descriptors = selectServerDescriptors(catalog, targets); + const inspections = await this.inspectServerDescriptors(descriptors, catalog); + return inspections.map((inspection) => ({ + ...inspection, + config: toMcpServerConfigView(inspection.config), + })); + } + + async resolveServerByName(name: string, query: McpRegistryQuery = {}): Promise { + await this.registry.get(name, query); + const catalog = await this.serverDescriptors(query); + const matches = catalog.filter((candidate) => candidate.runtimeName === name); + const descriptor = matches.find((candidate) => candidate.enabled) ?? matches[0]!; + this.requireUnambiguousRuntimeName(catalog, descriptor); + return descriptor.locator; + } + + async beginServerAuth( + locator: McpServerLocator, + query: McpRegistryQuery = {}, + ): Promise { + await this.waitForReadiness(); + const server = await this.resolveServer(locator, query); + const config = requireOAuthMcpConfig(server.runtimeName, server.config); + try { + const flow = await this.oauth.beginAuthorization(server.runtimeName, config.url); + const flowId = randomUUID(); + const idleTimer = setTimeout(() => { + const expired = this.authFlows.get(flowId); + this.authFlows.delete(flowId); + void expired?.flow.cancel(); + }, AUTH_FLOW_IDLE_TIMEOUT_MS); + idleTimer.unref(); + this.authFlows.set(flowId, { flow, idleTimer }); + return { + status: 'authorization-required', + flowId, + authorizationUrl: flow.authorizationUrl.toString(), + }; + } catch (error) { + if (error instanceof AlreadyAuthorizedError) { + return { status: 'already-authorized' }; + } + throw error; + } + } + + async completeServerAuth( + handle: McpServerAuthFlowHandle, + options?: { readonly signal?: AbortSignal }, + ): Promise { + if ( + handle.timeoutMs !== undefined && + (!Number.isInteger(handle.timeoutMs) || + handle.timeoutMs < 1 || + handle.timeoutMs > MAX_AUTH_TIMEOUT_MS) + ) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + `MCP OAuth timeoutMs must be an integer between 1 and ${MAX_AUTH_TIMEOUT_MS}`, + ); + } + const active = this.authFlows.get(handle.flowId); + if (active === undefined) { + throw new Error2(ErrorCodes.REQUEST_INVALID, `Unknown MCP OAuth flow: ${handle.flowId}`); + } + clearTimeout(active.idleTimer); + try { + await active.flow.complete({ + signal: options?.signal, + timeoutMs: handle.timeoutMs ?? DEFAULT_AUTH_TIMEOUT_MS, + }); + } finally { + this.authFlows.delete(handle.flowId); + } + } + + async cancelServerAuth(handle: Pick): Promise { + const active = this.authFlows.get(handle.flowId); + if (active === undefined) return; + clearTimeout(active.idleTimer); + this.authFlows.delete(handle.flowId); + await active.flow.cancel(); + } + + override dispose(): void { + for (const active of this.authFlows.values()) { + clearTimeout(active.idleTimer); + void active.flow.cancel(); + } + this.authFlows.clear(); + super.dispose(); + } + + async resetServerAuth(locator: McpServerLocator, query: McpRegistryQuery = {}): Promise { + await this.waitForReadiness(); + const server = await this.resolveServer(locator, query); + const config = requireRemoteMcpConfig(server.runtimeName, server.config); + await this.oauth.invalidate(server.runtimeName, config.url); + } + + private async serverDescriptors( + query: McpRegistryQuery = {}, + ): Promise { + return (await this.registry.list(query)).map((entry) => serverDescriptor(entry)); + } + + private async resolveServer( + locator: McpServerLocator, + query: McpRegistryQuery, + ): Promise { + const catalog = await this.serverDescriptors(query); + const server = selectServerDescriptors(catalog, [locator])[0]!; + this.requireUnambiguousRuntimeName(catalog, server); + return server; + } + + private requireUnambiguousRuntimeName( + catalog: readonly McpServerRuntimeDescriptor[], + server: McpServerRuntimeDescriptor, + ): void { + const conflict = catalog.find( + (candidate) => + candidate.serverId !== server.serverId && + candidate.enabled && + candidate.runtimeName === server.runtimeName, + ); + if (conflict !== undefined) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + `MCP runtime name "${server.runtimeName}" is shared by multiple enabled servers; use the locator-addressed RPC instead`, + ); + } + } + + private async serverAuthState( + entry: McpRegistryEntry, + cwd: string | undefined, + verify: boolean | undefined, + ): Promise { + const server = entry.config; + if (server.enabled === false) return 'not-applicable'; + if (server.transport === 'stdio') return 'not-applicable'; + if (server.bearerTokenEnvVar !== undefined) return 'bearer-token'; + if (server.headers !== undefined && server.auth !== 'oauth') return 'not-applicable'; + if (server.transport !== 'http' && server.auth !== 'oauth') return 'not-applicable'; + const tokens = await this.oauth.tokenState(entry.name, server.url); + const offline = (): McpServerAuthState => { + if (tokens.hasTokens) { + return !tokens.expired || tokens.hasRefreshToken ? 'oauth-authorized' : 'oauth-expired'; + } + return server.auth === 'oauth' ? 'oauth-required' : 'not-applicable'; + }; + + const probe = async (): Promise => + this.withProbe({ name: entry.name, ...server }, cwd, (manager) => { + const status = manager.get(entry.name)?.status; + if (status === 'connected') return tokens.hasTokens ? 'oauth-authorized' : 'not-applicable'; + if (status === 'needs-auth') return tokens.hasTokens ? 'oauth-expired' : 'oauth-required'; + return offline(); + }); + + if (verify === true) return probe(); + if (verify === false || tokens.hasTokens || server.auth === 'oauth') return offline(); + return probe(); + } + + private async inspectServerDescriptors( + descriptors: readonly McpServerRuntimeDescriptor[], + catalog: readonly McpServerRuntimeDescriptor[], + ): Promise { + const runtimeNameCounts = new Map(); + for (const server of new Map(catalog.map((item) => [item.serverId, item])).values()) { + if (!server.enabled) continue; + runtimeNameCounts.set(server.runtimeName, (runtimeNameCounts.get(server.runtimeName) ?? 0) + 1); + } + const credentialStates = new Map(); + const probeConfigs = Object.create(null) as Record; + for (const server of descriptors) { + if (configuredMcpAuthState(server) !== undefined) continue; + if (runtimeNameCounts.get(server.runtimeName) !== 1) continue; + const config = requireRemoteMcpConfig(server.runtimeName, server.config); + credentialStates.set( + server.serverId, + await this.oauth.tokenState(server.runtimeName, config.url), + ); + probeConfigs[server.runtimeName] = server.config; + } + let manager: McpConnectionManager | undefined; + try { + if (Object.keys(probeConfigs).length > 0) { + const section = this.config.get(MCP_SECTION); + manager = new McpConnectionManager({ + log: this.log, + oauthService: this.oauth, + resolveClientName: () => this.identity.current().slug, + resolveDefaultTimeouts: () => ({ + startupTimeoutMs: section?.startupTimeoutMs, + toolTimeoutMs: section?.toolTimeoutMs, + }), + }); + await manager.connectAll(probeConfigs); + } + const checkedAt = Date.now(); + return descriptors.map((server) => { + const configured = configuredMcpAuthState(server); + if (configured !== undefined) return { ...server, authStatus: configured }; + if (runtimeNameCounts.get(server.runtimeName) !== 1) { + return { + ...server, + authStatus: 'unavailable' as const, + checkedAt, + error: `MCP runtime name "${server.runtimeName}" is not unique`, + }; + } + const tokens = credentialStates.get(server.serverId); + const entry = manager?.get(server.runtimeName); + if (entry?.status === 'connected') { + return { + ...server, + authStatus: tokens?.hasTokens === true ? 'oauth-authorized' : 'not-applicable', + checkedAt, + }; + } + if (entry?.status === 'needs-auth') { + return { + ...server, + authStatus: tokens?.hasTokens === true ? 'oauth-expired' : 'oauth-required', + checkedAt, + }; + } + return { + ...server, + authStatus: 'unavailable' as const, + checkedAt, + error: entry?.error ?? `MCP server finished with status ${entry?.status ?? 'unknown'}`, + }; + }); + } finally { + await manager?.shutdown(); + } + } + + private async waitForReadiness(): Promise { + await this.config.ready; + await this.identity.resolved(); + } +} + +function throwReadOnlyMcpServer(entry: McpRegistryEntry): void { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + `MCP server "${entry.name}" is read-only: it is defined in ${entry.origin} — edit that file instead`, + ); +} + +function toManagedServer(entry: McpRegistryEntry): McpManagedServer { + return { + name: entry.name, + config: entry.mutable ? entry.config : toMcpServerConfigView(entry.config), + source: entry.source, + origin: entry.origin, + mutable: entry.mutable, + plugin: entry.plugin, + }; +} + +function mcpConfigWithoutName(server: GlobalMcpServerConfig): McpServerConfig { + const { name: _name, ...config } = server; + return config; +} + +type McpRemoteServerConfig = Exclude; + +function requireRemoteMcpConfig(name: string, config: McpServerConfig): McpRemoteServerConfig { + if (config.transport !== 'stdio') return config; + throw new Error2( + ErrorCodes.REQUEST_INVALID, + `MCP server "${name}" does not use a remote transport`, + ); +} + +function requireOAuthMcpConfig(name: string, input: McpServerConfig): McpRemoteServerConfig { + const config = requireRemoteMcpConfig(name, input); + if (config.bearerTokenEnvVar !== undefined) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + `MCP server "${name}" uses a static bearer token`, + ); + } + if (config.headers !== undefined && config.auth !== 'oauth') { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + `MCP server "${name}" uses static headers and is not marked for OAuth`, + ); + } + return config; +} + +/** Stable wire id of a locator: `global:` / `plugin::`. */ +export function mcpServerId(locator: McpServerLocator): string { + if (locator.source === 'global') return `global:${encodeURIComponent(locator.name)}`; + return `plugin:${encodeURIComponent(locator.pluginId)}:${encodeURIComponent(locator.serverName)}`; +} + +export function describeMcpServerLocator(locator: McpServerLocator): string { + if (locator.source === 'global') return locator.name; + return `${locator.pluginId}/${locator.serverName}`; +} + +type McpServerRuntimeDescriptor = Omit & { + readonly config: McpServerConfig; +}; + +type McpServerRuntimeInspection = McpServerRuntimeDescriptor & + Pick; + +function serverDescriptor(entry: McpRegistryEntry): McpServerRuntimeDescriptor { + const locator: McpServerLocator = + entry.source === 'plugin' && entry.plugin !== undefined + ? { source: 'plugin', pluginId: entry.plugin.id, serverName: entry.plugin.name } + : { source: 'global', name: entry.name }; + return { + serverId: mcpServerId(locator), + locator, + runtimeName: entry.name, + canonicalUrl: + entry.config.transport === 'stdio' + ? undefined + : canonicalMcpOAuthResource(entry.config.url), + origin: entry.source, + config: entry.config, + enabled: entry.config.enabled !== false, + editable: entry.mutable, + }; +} + +function selectServerDescriptors( + catalog: readonly McpServerRuntimeDescriptor[], + targets?: readonly McpServerLocator[], +): readonly McpServerRuntimeDescriptor[] { + const effectiveTargets = targets === null ? undefined : targets; + if (effectiveTargets === undefined) return catalog; + const byId = new Map(catalog.map((server) => [server.serverId, server])); + return effectiveTargets.map((target) => { + const server = byId.get(mcpServerId(target)); + if (server !== undefined) return server; + throw new Error2( + ErrorCodes.MCP_SERVER_NOT_FOUND, + `MCP server "${describeMcpServerLocator(target)}" was not found`, + ); + }); +} + +function configuredMcpAuthState( + server: McpServerRuntimeDescriptor, +): McpServerAuthState | undefined { + if (!server.enabled || server.config.enabled === false) return 'not-applicable'; + if (server.config.transport === 'stdio') return 'not-applicable'; + if (server.config.bearerTokenEnvVar !== undefined) return 'bearer-token'; + if (server.config.headers !== undefined && server.config.auth !== 'oauth') { + return 'not-applicable'; + } + return undefined; +} + +function standaloneTestResult( + name: string, + manager: McpConnectionManager, +): McpServerTestResult { + const entry = manager.get(name); + if (entry?.status !== 'connected') { + return { + success: false, + output: entry?.error ?? `MCP server "${name}" finished with status ${entry?.status ?? 'unknown'}`, + }; + } + const tools = manager.resolved(name)?.rawTools ?? []; + const lines = [ + `Connected to MCP server "${name}".`, + `Available tools: ${tools.length}`, + ...tools.map((tool) => `- ${tool.name}${tool.description ? `: ${tool.description}` : ''}`), + ]; + return { success: true, output: lines.join('\n') }; +} + +registerScopedService( + LifecycleScope.App, + IMcpManagementService, + McpManagementService, + ScopeActivation.OnDemand, + 'mcpManagement', +); diff --git a/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistry.ts b/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistry.ts new file mode 100644 index 0000000000..1c63e66f07 --- /dev/null +++ b/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistry.ts @@ -0,0 +1,55 @@ +import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; + +import type { McpServerConfig } from '#/mcpCore/config-schema'; + +export type McpServerSource = 'global' | 'plugin' | 'caller'; + +export interface McpRegistryPluginOrigin { + readonly id: string; + /** Manifest-local server name (without the `plugin-:` runtime prefix). */ + readonly name: string; +} + +export interface McpRegistryEntry { + /** Runtime name — for plugin entries the renamed `plugin-:` form. */ + readonly name: string; + /** Final effective config after source-specific transforms. */ + readonly config: McpServerConfig; + readonly source: McpServerSource; + /** global: the defining file path; plugin: the plugin id; caller: `'caller'`. */ + readonly origin: string; + /** True only for user-level global entries — the management API writes there. */ + readonly mutable: boolean; + readonly plugin?: McpRegistryPluginOrigin; +} + +export interface McpRegistryQuery { + /** + * When set, the project-root and project-local layers join the global + * source. Session-scoped resolutions pass the session workDir; the + * process-global management plane usually omits it. + */ + readonly cwd?: string; +} + +export interface IMcpRegistryService { + readonly _serviceBrand: undefined; + + list(query?: McpRegistryQuery): Promise; + + /** First match wins on a runtime-name collision (globals list first). */ + get(name: string, query?: McpRegistryQuery): Promise; + + /** + * Session-runtime resolution for one server name — the entry a live + * session should actually run, as opposed to the management view which + * lists every collision side by side. Returns `undefined` when no source + * currently defines the name. + */ + resolveRuntimeTarget(name: string, query?: McpRegistryQuery): Promise; +} + +export const IMcpRegistryService: ServiceIdentifier = + createDecorator('mcpRegistryService'); + +export { mcpServerConfigsEqual } from '#/mcpCore/connection-manager'; diff --git a/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistryService.ts b/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistryService.ts new file mode 100644 index 0000000000..2398c1904f --- /dev/null +++ b/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistryService.ts @@ -0,0 +1,116 @@ +import { LifecycleScope } from '#/app/scopes'; +import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; +import { canonicalWorkspaceRoot } from '#/_base/utils/paths'; + +import { ErrorCodes, Error2 } from '#/errors'; +import { IBootstrapService } from '#/app/bootstrap/bootstrap'; +import { loadMcpServersDetailed } from '#/app/mcpConfig/configLoader'; +import { IMcpConfigStore } from '#/app/mcpConfig/configStore'; +import { IPluginService } from '#/app/plugin/plugin'; +import { IHostFileSystem } from '#/os/interface/hostFileSystem'; +import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; +import { readWorkspaceTrust } from '#/workspace/workspaceTrust/trustRecord'; + +import { + IMcpRegistryService, + type McpRegistryEntry, + type McpRegistryQuery, +} from './mcpRegistry'; + +export class McpRegistryService implements IMcpRegistryService { + declare readonly _serviceBrand: undefined; + + constructor( + @IMcpConfigStore private readonly store: IMcpConfigStore, + @IPluginService private readonly plugins: IPluginService, + @IHostFileSystem private readonly fs: IHostFileSystem, + @IBootstrapService private readonly bootstrap: IBootstrapService, + @IAtomicDocumentStore private readonly docs: IAtomicDocumentStore, + ) {} + + async list(query: McpRegistryQuery = {}): Promise { + const out: McpRegistryEntry[] = []; + + if (query.cwd === undefined) { + const userEntries = await this.store.list(); + for (const server of userEntries) { + const { name, ...config } = server; + out.push({ + name, + config, + source: 'global', + origin: this.store.path, + mutable: true, + }); + } + } else { + const cwd = canonicalWorkspaceRoot(query.cwd); + if (!(await readWorkspaceTrust(this.docs, cwd))) { + const userEntries = await this.store.list(); + for (const server of userEntries) { + const { name, ...config } = server; + out.push({ + name, + config, + source: 'global', + origin: this.store.path, + mutable: true, + }); + } + } else { + const detailed = await loadMcpServersDetailed({ + fs: this.fs, + cwd, + homeDir: this.bootstrap.homeDir, + }); + for (const [name, config] of Object.entries(detailed.servers)) { + const origin = detailed.origins[name] ?? this.store.path; + out.push({ + name, + config, + source: 'global', + origin, + mutable: origin === this.store.path, + }); + } + } + } + + for (const entry of await this.plugins.mcpServerEntries()) { + out.push({ + name: entry.name, + config: entry.config, + source: 'plugin', + origin: entry.pluginId, + mutable: false, + plugin: { id: entry.pluginId, name: entry.serverName }, + }); + } + + return out; + } + + async get(name: string, query: McpRegistryQuery = {}): Promise { + const entry = (await this.list(query)).find((candidate) => candidate.name === name); + if (entry !== undefined) return entry; + throw new Error2(ErrorCodes.MCP_SERVER_NOT_FOUND, `MCP server "${name}" was not found`); + } + + async resolveRuntimeTarget( + name: string, + query: McpRegistryQuery = {}, + ): Promise { + const matches = (await this.list(query)).filter((entry) => entry.name === name); + const file = matches.find((entry) => entry.source === 'global'); + if (file !== undefined) return file; + return matches.find((entry) => entry.source === 'plugin' && entry.config.enabled !== false); + } +} + +registerScopedService( + LifecycleScope.App, + IMcpRegistryService, + McpRegistryService, + ScopeActivation.OnDemand, + 'mcpRegistry', +); diff --git a/packages/agent-core-v2/src/app/plugin/manager.ts b/packages/agent-core-v2/src/app/plugin/manager.ts index d041757b51..eee1ace2fb 100644 --- a/packages/agent-core-v2/src/app/plugin/manager.ts +++ b/packages/agent-core-v2/src/app/plugin/manager.ts @@ -2,20 +2,20 @@ import { cp, mkdir, mkdtemp, realpath, rename, rm, stat } from 'node:fs/promises import { tmpdir } from 'node:os'; import path from 'node:path'; -import { BugIndicatingError, Error2, ErrorCodes, PluginErrors } from '#/errors'; import type { HookDef } from '#/features/externalHooks/internal/types'; -import type { McpServerConfig } from '#/mcpCore/config-schema'; -import type { PluginAgentRoot } from './types'; import { discoverFileSkills } from '#/app/skillCatalog/fileSkillDiscovery'; import type { SkillDiscoveryResult } from '#/app/skillCatalog/skillDiscovery'; import type { SkillRoot } from '#/app/skillCatalog/types'; +import { BugIndicatingError, Error2, ErrorCodes, PluginErrors } from '#/errors'; +import type { McpServerConfig } from '#/mcpCore/config-schema'; import { downloadZip, extractZip } from './archive'; import { loadPluginCommand } from './commands'; import { resolveGithubCommitSha, resolveGithubSource } from './github-resolver'; -import { resolveInstallSource } from './source'; import { parseManifest, type ParsedManifestResult } from './manifest'; +import { resolveInstallSource } from './source'; import { readInstalled, writeInstalled, type InstalledRecord } from './store'; +import type { PluginAgentRoot } from './types'; import { normalizePluginId, type EnabledPluginSessionStart, @@ -24,6 +24,7 @@ import { type PluginCommandDef, type PluginGithubMetadata, type PluginInfo, + type PluginMcpServerEntry, type PluginMcpServerInfo, type PluginRecord, type PluginSource, @@ -44,9 +45,7 @@ interface ManagedPluginCopy { export class PluginManager { private readonly kimiHomeDir: string; - private readonly discoverSkills: ( - roots: readonly SkillRoot[], - ) => Promise; + private readonly discoverSkills: (roots: readonly SkillRoot[]) => Promise; private records = new Map(); constructor(options: PluginManagerOptions) { @@ -117,7 +116,8 @@ export class PluginManager { const parsed = await parseManifest(sourceRoot); if (parsed.manifest === undefined) { - const msg = parsed.diagnostics.find((d) => d.severity === 'error')?.message ?? 'no manifest'; + const msg = + parsed.diagnostics.find((d) => d.severity === 'error')?.message ?? 'no manifest'; throw new Error2( ErrorCodes.PLUGIN_LOAD_FAILED, sourceType === 'local-path' @@ -369,6 +369,28 @@ export class PluginManager { return out; } + mcpServerEntries(): readonly PluginMcpServerEntry[] { + const out: PluginMcpServerEntry[] = []; + for (const record of this.records.values()) { + if (record.state !== 'ok' || record.manifest === undefined) continue; + for (const [name, config] of Object.entries(record.manifest.mcpServers ?? {})) { + const enabled = record.enabled && isMcpServerEnabled(record, name, config); + const effective = withPluginMcpRuntime( + withMcpServerEnabled(config, enabled), + record.root, + this.kimiHomeDir, + ); + out.push({ + name: pluginMcpRuntimeName(record.id, name), + config: effective, + pluginId: record.id, + serverName: name, + }); + } + } + return out; + } + summaries(): readonly PluginSummary[] { return this.list().map((record) => recordToSummary(record)); } @@ -583,11 +605,7 @@ async function recordFrom(input: { originalSource: input.originalSource, capabilities: input.capabilities, github: input.github, - skillCount: await countDiscoveredPluginSkills( - input.id, - parsed.manifest, - input.discoverSkills, - ), + skillCount: await countDiscoveredPluginSkills(input.id, parsed.manifest, input.discoverSkills), manifest: parsed.manifest, manifestKind: parsed.manifestKind, manifestPath: parsed.manifestPath, diff --git a/packages/agent-core-v2/src/app/plugin/plugin.ts b/packages/agent-core-v2/src/app/plugin/plugin.ts index fe2754b181..2185469576 100644 --- a/packages/agent-core-v2/src/app/plugin/plugin.ts +++ b/packages/agent-core-v2/src/app/plugin/plugin.ts @@ -1,8 +1,8 @@ import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; import type { Event } from '#/_base/event'; import type { HookDef } from '#/features/externalHooks/internal/types'; -import type { McpServerConfig } from '#/mcpCore/config-schema'; import type { SkillRoot } from '#/app/skillCatalog/types'; +import type { McpServerConfig } from '#/mcpCore/config-schema'; import type { EnabledPluginSessionStart, @@ -10,7 +10,9 @@ import type { PluginAgentRoot, PluginCommandDef, PluginInfo, + PluginMcpServerEntry, PluginMutationSummary, + PluginReloadEvent, PluginSummary, PluginUpdateStatus, ReloadSummary, @@ -56,9 +58,10 @@ export interface IPluginService { enabledSessionStarts(): Promise; enabledSystemPrompts(): Promise; enabledMcpServers(): Promise>; + mcpServerEntries(): Promise; enabledHooks(): Promise; hasLoadedSnapshot(): boolean; - readonly onDidReload: Event; + readonly onDidReload: Event; readonly onDidMutate: Event; } diff --git a/packages/agent-core-v2/src/app/plugin/pluginService.ts b/packages/agent-core-v2/src/app/plugin/pluginService.ts index 9de12b153d..6fa47c80fb 100644 --- a/packages/agent-core-v2/src/app/plugin/pluginService.ts +++ b/packages/agent-core-v2/src/app/plugin/pluginService.ts @@ -1,16 +1,16 @@ import { KIMI_CODE_PROVIDER_NAME } from '@moonshot-ai/kimi-code-oauth'; +import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; import { Service } from '#/_base/di/service'; -import { Emitter, type Event } from '#/_base/event'; +import { AsyncEmitter, Emitter, type Event } from '#/_base/event'; +import type { HookDef } from '#/features/externalHooks/internal/types'; +import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { LifecycleScope } from '#/app/scopes'; -import { ScopeActivation, registerScopedService } from '#/_base/di/scope'; +import { ISkillDiscovery } from '#/app/skillCatalog/skillDiscovery'; +import type { SkillRoot } from '#/app/skillCatalog/types'; import { BugIndicatingError, Error2, PluginErrors } from '#/errors'; -import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { IProviderService } from '#/kosong/provider/provider'; -import { ISkillDiscovery } from '#/app/skillCatalog/skillDiscovery'; -import type { HookDef } from '#/features/externalHooks/internal/types'; import type { McpServerConfig } from '#/mcpCore/config-schema'; -import type { SkillRoot } from '#/app/skillCatalog/types'; import { PluginManager } from './manager'; import { @@ -27,8 +27,10 @@ import type { PluginCommandDef, PluginInfo, PluginAgentRoot, + PluginMcpServerEntry, PluginMutation, PluginMutationSummary, + PluginReloadEvent, PluginSummary, PluginUpdateStatus, ReloadSummary, @@ -37,6 +39,17 @@ import type { const KIMI_CODE_BASE_URL_ENV = 'KIMI_CODE_BASE_URL'; const KIMI_CODE_OAUTH_HOST_ENV = 'KIMI_CODE_OAUTH_HOST'; const KIMI_OAUTH_HOST_ENV = 'KIMI_OAUTH_HOST'; +const NO_ABORT = new AbortController().signal; + +interface PluginReloadNotification { + readonly summary: ReloadSummary; + readonly delivery: Promise; +} + +interface PluginMutationOutcome { + readonly result: T; + readonly notification: PluginReloadNotification; +} export class PluginService extends Service implements IPluginService { declare readonly _serviceBrand: undefined; @@ -49,10 +62,10 @@ export class PluginService extends Service implements IPluginService { private snapshotLoaded = false; private loadError: Error | undefined; private mutationQueue: Promise = Promise.resolve(); - private readonly onDidReloadEmitter = this._register(new Emitter()); + private readonly onDidReloadEmitter = this._register(new AsyncEmitter()); private readonly onDidMutateEmitter = this._register(new Emitter()); - readonly onDidReload: Event = this.onDidReloadEmitter.event; + readonly onDidReload: Event = this.onDidReloadEmitter.event; readonly onDidMutate: Event = this.onDidMutateEmitter.event; constructor( @@ -76,52 +89,64 @@ export class PluginService extends Service implements IPluginService { } installPlugin(input: InstallPluginInput): Promise { - return this.runSerializedOperation(async () => { + return this.runNotifiedMutation(async () => { const record = await this.manager.install(input.source); const info = this.manager.info(record.id); if (info === undefined) throw new BugIndicatingError(`Plugin "${record.id}" missing right after install`); - await this.reloadAndNotify({ mutation: { kind: 'install', id: record.id } }); - return info; + const notification = await this.reloadAndNotify({ + mutation: { kind: 'install', id: record.id }, + }); + return { result: info, notification }; }); } setPluginEnabled(input: SetPluginEnabledInput): Promise { - return this.runSerializedOperation(async () => { + return this.runNotifiedMutation(async () => { await this.manager.setEnabled(input.id, input.enabled); - await this.reloadAndNotify({ + const notification = await this.reloadAndNotify({ mutation: { kind: input.enabled ? 'enable' : 'disable', id: input.id }, }); + return { result: undefined, notification }; }); } setPluginMcpServerEnabled(input: SetPluginMcpServerEnabledInput): Promise { - return this.runSerializedOperation(async () => { + return this.runNotifiedMutation(async () => { await this.manager.setMcpServerEnabled(input.id, input.server, input.enabled); - await this.reloadAndNotify({ mutation: { kind: 'mcp-server', id: input.id } }); + const notification = await this.reloadAndNotify({ + mutation: { kind: 'mcp-server', id: input.id }, + }); + return { result: undefined, notification }; }); } removePlugin(input: RemovePluginInput): Promise { - return this.runSerializedOperation(async () => { + return this.runNotifiedMutation(async () => { await this.manager.remove(input.id); - await this.reloadAndNotify({ mutation: { kind: 'remove', id: input.id } }); + const notification = await this.reloadAndNotify({ + mutation: { kind: 'remove', id: input.id }, + }); + return { result: undefined, notification }; }); } reloadPlugins(): Promise { - const reload = this.enqueueMutation(async () => { - try { - return await this.reloadAndNotify(); - } catch (error) { - this.loadError = error instanceof Error ? error : new Error(String(error)); - throw new Error2( - PluginErrors.codes.PLUGIN_LOAD_FAILED, - `Failed to reload plugins: ${this.loadError.message}`, - { cause: this.loadError, details: { kimiHomeDir: this.homeDir } }, - ); - } - }); + const reload = this.awaitReloadDelivery( + this.enqueueMutation(async () => { + try { + const notification = await this.reloadAndNotify(); + return { result: notification.summary, notification }; + } catch (error) { + this.loadError = error instanceof Error ? error : new Error(String(error)); + throw new Error2( + PluginErrors.codes.PLUGIN_LOAD_FAILED, + `Failed to reload plugins: ${this.loadError.message}`, + { cause: this.loadError, details: { kimiHomeDir: this.homeDir } }, + ); + } + }), + ); this.initialLoadPromise ??= reload.then( () => undefined, () => undefined, @@ -131,14 +156,24 @@ export class PluginService extends Service implements IPluginService { private async reloadAndNotify(options?: { readonly mutation: PluginMutation; - }): Promise { + }): Promise { const summary = await this.manager.reload(); this.snapshotLoaded = true; this.loadError = undefined; - this.onDidReloadEmitter.fire(summary); + const delivery = this.onDidReloadEmitter.fireAsyncConcurrent(summary, NO_ABORT); if (options?.mutation !== undefined) this.onDidMutateEmitter.fire({ ...summary, mutation: options.mutation }); - return summary; + return { summary, delivery }; + } + + private runNotifiedMutation(operation: () => Promise>): Promise { + return this.awaitReloadDelivery(this.runSerializedOperation(operation)); + } + + private async awaitReloadDelivery(operation: Promise>): Promise { + const { result, notification } = await operation; + await notification.delivery; + return result; } getPluginInfo(input: GetPluginInfoInput): Promise { @@ -190,6 +225,17 @@ export class PluginService extends Service implements IPluginService { }); } + mcpServerEntries(): Promise { + return this.runManagementRead(async () => { + const entries = this.manager.mcpServerEntries(); + if (!entries.some((entry) => entry.config.transport === 'stdio')) { + return entries; + } + const managedEnv = await this.managedKimiCodeEnvForPlugins(); + return withManagedKimiPluginEnvOnEntries(entries, managedEnv); + }); + } + enabledHooks(): Promise { return this.runConsumptionRead([], async () => this.manager.enabledHooks()); } @@ -265,8 +311,7 @@ export class PluginService extends Service implements IPluginService { const envBaseUrl = this.envBaseUrl; const envOAuthHost = this.envOAuthHost; const hasEnvOverride = envBaseUrl !== undefined || envOAuthHost !== undefined; - const baseUrl = - envBaseUrl !== undefined ? envBaseUrl.replace(/\/+$/, '') : provider?.baseUrl; + const baseUrl = envBaseUrl !== undefined ? envBaseUrl.replace(/\/+$/, '') : provider?.baseUrl; const oauthHost = hasEnvOverride ? envOAuthHost : provider?.oauth?.oauthHost; const env: Record = {}; if (baseUrl !== undefined) env[KIMI_CODE_BASE_URL_ENV] = baseUrl; @@ -283,13 +328,23 @@ function withManagedKimiPluginEnv( const out: Record = {}; for (const [name, server] of Object.entries(pluginServers)) { out[name] = - server.transport === 'stdio' - ? { ...server, env: { ...server.env, ...managedEnv } } - : server; + server.transport === 'stdio' ? { ...server, env: { ...server.env, ...managedEnv } } : server; } return out; } +function withManagedKimiPluginEnvOnEntries( + entries: readonly PluginMcpServerEntry[], + managedEnv: Record, +): readonly PluginMcpServerEntry[] { + if (Object.keys(managedEnv).length === 0) return entries; + return entries.map((entry) => + entry.config.transport === 'stdio' + ? { ...entry, config: { ...entry.config, env: { ...entry.config.env, ...managedEnv } } } + : entry, + ); +} + registerScopedService( LifecycleScope.App, IPluginService, diff --git a/packages/agent-core-v2/src/app/plugin/types.ts b/packages/agent-core-v2/src/app/plugin/types.ts index 8ae8afbe21..3a7375fed1 100644 --- a/packages/agent-core-v2/src/app/plugin/types.ts +++ b/packages/agent-core-v2/src/app/plugin/types.ts @@ -1,3 +1,4 @@ +import type { IWaitUntil } from '#/_base/event'; import type { HookDefConfig } from '#/features/externalHooks/configSection'; import type { McpServerConfig } from '#/mcpCore/config-schema'; @@ -71,6 +72,13 @@ export interface PluginMcpServerInfo { readonly headerKeys?: readonly string[]; } +export interface PluginMcpServerEntry { + readonly name: string; + readonly config: McpServerConfig; + readonly pluginId: string; + readonly serverName: string; +} + export interface PluginCommandDef { readonly pluginId: string; readonly name: string; @@ -165,6 +173,8 @@ export interface ReloadSummary { readonly errors: ReadonlyArray<{ readonly id: string; readonly message: string }>; } +export type PluginReloadEvent = ReloadSummary & IWaitUntil; + export interface PluginMutation { readonly kind: 'install' | 'enable' | 'disable' | 'remove' | 'mcp-server'; readonly id: string; diff --git a/packages/agent-core-v2/src/index.ts b/packages/agent-core-v2/src/index.ts index 31db38ea9b..f5d6e0484f 100644 --- a/packages/agent-core-v2/src/index.ts +++ b/packages/agent-core-v2/src/index.ts @@ -441,6 +441,14 @@ export { type McpSection, } from '#/app/mcpConfig/configSection'; export * from '#/app/mcpConfig/oauthStore'; +export { IMcpConfigStore } from '#/app/mcpConfig/configStore'; +import '#/app/mcpConfig/configStore'; +export { IMcpOAuthService } from '#/app/mcpConfig/oauthService'; +import '#/app/mcpConfig/oauthService'; +export * from '#/app/mcpRegistry/mcpRegistry'; +import '#/app/mcpRegistry/mcpRegistryService'; +export * from '#/app/mcpManagement/mcpManagement'; +import '#/app/mcpManagement/mcpManagementService'; export * from '#/workspace/workspaceMcpConfig/workspaceMcpConfig'; export * from '#/workspace/workspaceMcpConfig/workspaceMcpConfigService'; export * from '#/workspace/workspaceMcp/workspaceMcp'; diff --git a/packages/agent-core-v2/src/mcpCore/client-http.ts b/packages/agent-core-v2/src/mcpCore/client-http.ts index ba08624797..cba0f91622 100644 --- a/packages/agent-core-v2/src/mcpCore/client-http.ts +++ b/packages/agent-core-v2/src/mcpCore/client-http.ts @@ -15,6 +15,7 @@ import { type UnexpectedCloseReason, } from './client-shared'; import { buildMcpRemoteHeaders } from './client-remote'; +import { createMcpOAuthFetch } from './oauth/provider'; import type { MCPClient, MCPToolDefinition, MCPToolResult } from './types'; export interface HttpMcpClientOptions { @@ -47,7 +48,7 @@ export class HttpMcpClient implements MCPClient { this.transport = new StreamableHTTPClientTransport(new URL(config.url), { requestInit: headers !== undefined ? { headers } : undefined, - fetch: options.fetch, + fetch: createMcpOAuthFetch(options.oauthProvider, options.fetch), authProvider: options.oauthProvider, }); this.client = new Client({ diff --git a/packages/agent-core-v2/src/mcpCore/client-sse.ts b/packages/agent-core-v2/src/mcpCore/client-sse.ts index 250c365e76..d490a09bf7 100644 --- a/packages/agent-core-v2/src/mcpCore/client-sse.ts +++ b/packages/agent-core-v2/src/mcpCore/client-sse.ts @@ -15,6 +15,7 @@ import { type UnexpectedCloseReason, } from './client-shared'; import { buildMcpRemoteHeaders } from './client-remote'; +import { createMcpOAuthFetch } from './oauth/provider'; import type { MCPClient, MCPToolDefinition, MCPToolResult } from './types'; export interface SseMcpClientOptions { @@ -47,7 +48,7 @@ export class SseMcpClient implements MCPClient { this.transport = new SSEClientTransport(new URL(config.url), { requestInit: headers !== undefined ? { headers } : undefined, - fetch: options.fetch, + fetch: createMcpOAuthFetch(options.oauthProvider, options.fetch), authProvider: options.oauthProvider, }); this.client = new Client({ diff --git a/packages/agent-core-v2/src/mcpCore/configView.ts b/packages/agent-core-v2/src/mcpCore/configView.ts new file mode 100644 index 0000000000..0c35e5f5ee --- /dev/null +++ b/packages/agent-core-v2/src/mcpCore/configView.ts @@ -0,0 +1,19 @@ +import type { McpServerConfig } from './config-schema'; + +export type McpServerConfigView = + | (Omit, 'env'> & { + readonly envKeys?: readonly string[]; + }) + | (Omit, 'headers'> & { + readonly headerKeys?: readonly string[]; + }); + +/** Project a full effective config into its wire-facing view. */ +export function toMcpServerConfigView(config: McpServerConfig): McpServerConfigView { + if (config.transport === 'stdio') { + const { env, ...safe } = config; + return env === undefined ? safe : { ...safe, envKeys: Object.keys(env).toSorted() }; + } + const { headers, ...safe } = config; + return headers === undefined ? safe : { ...safe, headerKeys: Object.keys(headers).toSorted() }; +} diff --git a/packages/agent-core-v2/src/mcpCore/connection-manager.ts b/packages/agent-core-v2/src/mcpCore/connection-manager.ts index 6ae2c2606a..98581ee11c 100644 --- a/packages/agent-core-v2/src/mcpCore/connection-manager.ts +++ b/packages/agent-core-v2/src/mcpCore/connection-manager.ts @@ -297,6 +297,17 @@ export class McpConnectionManager implements McpConnectionView { return work; } + /** + * {@link reconnectAndJoin} queued behind any in-flight reconnect: a + * credential that lands while a reconnect is already running triggers one + * more pass instead of being absorbed by the stale run. + */ + async reconnectAfterCurrent(name: string): Promise { + const existing = this.inFlightReconnects.get(name); + if (existing !== undefined) await existing.catch(() => undefined); + await this.reconnectAndJoin(name); + } + async shutdown(): Promise { const entries = Array.from(this.entries.values()); this.entries.clear(); @@ -561,7 +572,12 @@ function stderrTail(client: RuntimeMcpClient | undefined): string | undefined { return snapshot.trimEnd(); } -function mcpServerConfigsEqual(a: McpServerConfig, b: McpServerConfig): boolean { +/** + * Structural equality for effective configs, backing the idempotent-connect + * guard (config reconcilers and explicit callers may issue the same upsert) + * and the management plane's change detection. + */ +export function mcpServerConfigsEqual(a: McpServerConfig, b: McpServerConfig): boolean { return stableConfigJson(a) === stableConfigJson(b); } diff --git a/packages/agent-core-v2/src/mcpCore/oauth/callback-server.ts b/packages/agent-core-v2/src/mcpCore/oauth/callback-server.ts index f190c4dd76..cb4141f867 100644 --- a/packages/agent-core-v2/src/mcpCore/oauth/callback-server.ts +++ b/packages/agent-core-v2/src/mcpCore/oauth/callback-server.ts @@ -8,10 +8,24 @@ export interface CallbackResult { export interface CallbackServer { readonly redirectUri: string; + /** + * Resolves with the OAuth callback payload, or rejects when: + * - `signal` aborts → AbortError + * - `timeoutMs` elapses → Error('OAuth callback timed out') + * - the user's authorization server returns an error → Error('OAuth error: ') + * - `close()` is called → OAuthCallbackClosedError + */ waitForCode(opts: { signal?: AbortSignal; timeoutMs?: number }): Promise; close(): Promise; } +export class OAuthCallbackClosedError extends Error { + constructor() { + super('OAuth callback listener closed'); + this.name = 'OAuthCallbackClosedError'; + } +} + const SUCCESS_HTML = 'Authorized' + '' + @@ -29,12 +43,29 @@ const ERROR_HTML = export async function startCallbackServer(): Promise { let resolveCode: ((value: CallbackResult) => void) | undefined; let rejectCode: ((reason: Error) => void) | undefined; - let settled = false; + let cleanupWait: (() => void) | undefined; + let outcome: + | { readonly status: 'pending' } + | { readonly status: 'resolved'; readonly value: CallbackResult } + | { readonly status: 'rejected'; readonly reason: Error } = { status: 'pending' }; - const settle = (fn: () => void) => { - if (settled) return; - settled = true; - fn(); + const settle = ( + next: + | { readonly status: 'resolved'; readonly value: CallbackResult } + | { readonly status: 'rejected'; readonly reason: Error }, + ) => { + if (outcome.status !== 'pending') return; + outcome = next; + cleanupWait?.(); + cleanupWait = undefined; + if (next.status === 'resolved') { + resolveCode?.(next.value); + } else { + rejectCode?.(next.reason); + } + resolveCode = undefined; + rejectCode = undefined; + void closeServer(); }; const server: Server = createServer((req, res) => { @@ -61,26 +92,26 @@ export async function startCallbackServer(): Promise { if (errorParam !== null) { const description = url.searchParams.get('error_description') ?? ''; res.writeHead(400, { 'content-type': 'text/html; charset=utf-8' }).end(ERROR_HTML); - settle(() => { - rejectCode?.( - new Error(`OAuth error: ${errorParam}${description ? ` — ${description}` : ''}`), - ); + settle({ + status: 'rejected', + reason: new Error( + `OAuth error: ${errorParam}${description ? ` — ${description}` : ''}`, + ), }); return; } const code = url.searchParams.get('code'); if (code === null || code.length === 0) { res.writeHead(400, { 'content-type': 'text/html; charset=utf-8' }).end(ERROR_HTML); - settle(() => { - rejectCode?.(new Error('OAuth callback missing authorization code')); + settle({ + status: 'rejected', + reason: new Error('OAuth callback missing authorization code'), }); return; } const state = url.searchParams.get('state') ?? undefined; res.writeHead(200, { 'content-type': 'text/html; charset=utf-8' }).end(SUCCESS_HTML); - settle(() => { - resolveCode?.({ code, state }); - }); + settle({ status: 'resolved', value: { code, state } }); } await new Promise((resolve, reject) => { @@ -93,44 +124,49 @@ export async function startCallbackServer(): Promise { const port = (server.address() as AddressInfo).port; const redirectUri = `http://127.0.0.1:${port}/callback`; - let closed = false; - const close = async () => { - if (closed) return; - closed = true; - await new Promise((resolve) => { + let closeServerPromise: Promise | undefined; + const closeServer = (): Promise => { + closeServerPromise ??= new Promise((resolve) => { server.close(() => { resolve(); }); }); + return closeServerPromise; + }; + const close = async () => { + settle({ status: 'rejected', reason: new OAuthCallbackClosedError() }); + await closeServer(); }; const waitForCode: CallbackServer['waitForCode'] = ({ signal, timeoutMs } = {}) => { return new Promise((resolve, reject) => { + if (outcome.status === 'resolved') { + resolve(outcome.value); + return; + } + if (outcome.status === 'rejected') { + reject(outcome.reason); + return; + } + let timer: NodeJS.Timeout | undefined; const onAbort = () => { - settle(() => - rejectCode?.( + settle({ + status: 'rejected', + reason: signal?.reason instanceof Error ? signal.reason : new Error('OAuth flow aborted'), - ), - ); + }); }; const cleanup = () => { if (timer !== undefined) clearTimeout(timer); signal?.removeEventListener('abort', onAbort); }; - resolveCode = (value) => { - cleanup(); - void close(); - resolve(value); - }; - rejectCode = (reason) => { - cleanup(); - void close(); - reject(reason); - }; + cleanupWait = cleanup; + resolveCode = resolve; + rejectCode = reject; if (timeoutMs !== undefined) { timer = setTimeout(() => { - settle(() => rejectCode?.(new Error('OAuth callback timed out'))); + settle({ status: 'rejected', reason: new Error('OAuth callback timed out') }); }, timeoutMs); } if (signal !== undefined) { diff --git a/packages/agent-core-v2/src/mcpCore/oauth/provider.ts b/packages/agent-core-v2/src/mcpCore/oauth/provider.ts index f1b9c831a0..bafa3151b7 100644 --- a/packages/agent-core-v2/src/mcpCore/oauth/provider.ts +++ b/packages/agent-core-v2/src/mcpCore/oauth/provider.ts @@ -1,17 +1,19 @@ import { randomBytes } from 'node:crypto'; -import { BugIndicatingError } from '#/errors'; - import type { OAuthClientProvider, OAuthDiscoveryState, } from '@modelcontextprotocol/sdk/client/auth.js'; -import type { - OAuthClientInformationFull, - OAuthClientInformationMixed, - OAuthClientMetadata, - OAuthTokens, +import { + OAuthTokensSchema, + type OAuthClientInformationFull, + type OAuthClientInformationMixed, + type OAuthClientMetadata, + type OAuthTokens, } from '@modelcontextprotocol/sdk/shared/auth.js'; +import { OAuthTokenTransaction } from '@moonshot-ai/kimi-code-oauth'; + +import { BugIndicatingError } from '#/errors'; import { KIMI_MCP_CLIENT_NAME } from '../client-shared'; import { canonicalMcpOAuthResource, mcpOAuthStoreKey, type McpOAuthStore } from './store'; @@ -19,53 +21,108 @@ import { canonicalMcpOAuthResource, mcpOAuthStoreKey, type McpOAuthStore } from const TOKENS_SUFFIX = '-tokens.json'; const CLIENT_SUFFIX = '-client.json'; const DISCOVERY_SUFFIX = '-discovery.json'; +/** Sidecar `-meta.json` suffix; the service scans these on startup. */ +export const META_SUFFIX = '-meta.json'; const PASSIVE_REDIRECT_URI = 'http://127.0.0.1:3118/callback'; +export interface StoredMcpOAuthTokens extends OAuthTokens { + readonly obtained_at?: number; +} + +/** Sidecar `-meta.json` record mapping a store key back to its server. */ +export interface McpOAuthStoreMeta { + readonly serverName: string; + readonly serverUrl: string; +} + export interface McpOAuthProviderOptions { readonly serverName: string; readonly serverUrl: string | URL; readonly store: McpOAuthStore; readonly clientLabel?: string; readonly clientName?: string; -} - -export interface StoredMcpOAuthTokens extends OAuthTokens { - readonly obtained_at?: number; + readonly now?: () => number; + /** Called after tokens are persisted (login, exchange, or refresh). */ + readonly onTokensSaved?: (tokens: StoredMcpOAuthTokens) => void; + /** Called after any credential invalidation, including SDK-driven ones. */ + readonly onCredentialsInvalidated?: ( + scope: 'all' | 'client' | 'tokens' | 'verifier' | 'discovery', + ) => void; + /** Receives every in-flight token-grant promise so shutdown can drain it. */ + readonly track?: (operation: Promise) => void; } export class McpOAuthClientProvider implements OAuthClientProvider { readonly storeKey: string; readonly serverUrl: string; readonly ready: Promise; + private readonly serverName: string; private readonly store: McpOAuthStore; private readonly clientLabel: string; + private readonly onTokensSaved: McpOAuthProviderOptions['onTokensSaved']; + private readonly onCredentialsInvalidated: McpOAuthProviderOptions['onCredentialsInvalidated']; + private readonly now: () => number; private _redirectUrl: URL | undefined; private _codeVerifier: string | undefined; private _state: string | undefined; private _lastAuthorizationUrl: URL | undefined; + private readonly tokenTransaction: OAuthTokenTransaction; private clientCache: OAuthClientInformationMixed | undefined; - private tokensCache: OAuthTokens | undefined; private discoveryCache: OAuthDiscoveryState | undefined; constructor(options: McpOAuthProviderOptions) { this.serverUrl = canonicalMcpOAuthResource(options.serverUrl); this.storeKey = mcpOAuthStoreKey(options.serverName, this.serverUrl); + this.serverName = options.serverName; this.store = options.store; this.clientLabel = options.clientLabel ?? `${options.clientName ?? KIMI_MCP_CLIENT_NAME} (${options.serverName})`; + this.onTokensSaved = options.onTokensSaved; + this.onCredentialsInvalidated = options.onCredentialsInvalidated; + this.now = options.now ?? Date.now; + const tokensFile = `${this.storeKey}${TOKENS_SUFFIX}`; + const metaFile = `${this.storeKey}${META_SUFFIX}`; + this.tokenTransaction = new OAuthTokenTransaction({ + key: this.storeKey, + read: async () => this.store.read(tokensFile), + write: async (tokens) => { + const incoming = tokens as StoredMcpOAuthTokens; + await this.store.write(tokensFile, { + ...incoming, + obtained_at: incoming.obtained_at ?? this.now(), + }); + }, + remove: async () => { + await this.store.remove(tokensFile); + }, + parse: (value) => OAuthTokensSchema.safeParse(value).data, + normalize: (tokens) => OAuthTokensSchema.safeParse(tokens).data ?? tokens, + track: options.track, + afterCommit: async (tokens) => { + if (tokens === undefined) { + await this.store.remove(metaFile); + return; + } + const meta: McpOAuthStoreMeta = { serverName: this.serverName, serverUrl: this.serverUrl }; + await this.store.write(metaFile, meta); + const stamped: StoredMcpOAuthTokens = { + ...tokens, + obtained_at: (tokens as StoredMcpOAuthTokens).obtained_at ?? this.now(), + }; + this.onTokensSaved?.(stamped); + }, + }); this.ready = this.load(); } private async load(): Promise { - const [client, tokens, discovery] = await Promise.all([ + const [client, discovery] = await Promise.all([ this.store.read(`${this.storeKey}${CLIENT_SUFFIX}`), - this.store.read(`${this.storeKey}${TOKENS_SUFFIX}`), this.store.read(`${this.storeKey}${DISCOVERY_SUFFIX}`), ]); this.clientCache = client; - this.tokensCache = tokens; this.discoveryCache = discovery; } @@ -115,19 +172,25 @@ export class McpOAuthClientProvider implements OAuthClientProvider { } async saveClientInformation(info: OAuthClientInformationMixed): Promise { - this.clientCache = info; await this.store.write(`${this.storeKey}${CLIENT_SUFFIX}`, info); + this.clientCache = info; } async tokens(): Promise { - await this.ready; - return this.tokensCache; + return this.store.read(`${this.storeKey}${TOKENS_SUFFIX}`); } async saveTokens(tokens: OAuthTokens): Promise { - const stamped: StoredMcpOAuthTokens = { ...tokens, obtained_at: Date.now() }; - this.tokensCache = stamped; - await this.store.write(`${this.storeKey}${TOKENS_SUFFIX}`, stamped); + await this.tokenTransaction.save(tokens); + } + + /** + * Wrap the fetch used by the SDK's OAuth flow. Refresh-token grants for the + * same MCP identity are serialized, re-read from durable storage inside the + * lock, and committed before the lock is released. + */ + createOAuthFetch(fetchFn: typeof fetch = globalThis.fetch): typeof fetch { + return this.tokenTransaction.createFetch(fetchFn); } redirectToAuthorization(url: URL): void { @@ -146,8 +209,8 @@ export class McpOAuthClientProvider implements OAuthClientProvider { } async saveDiscoveryState(state: OAuthDiscoveryState): Promise { - this.discoveryCache = state; await this.store.write(`${this.storeKey}${DISCOVERY_SUFFIX}`, state); + this.discoveryCache = state; } async discoveryState(): Promise { @@ -162,32 +225,51 @@ export class McpOAuthClientProvider implements OAuthClientProvider { const uris = info.redirect_uris; if (!Array.isArray(uris) || uris.length === 0) return false; if (uris.includes(redirectUri)) return false; - await this.invalidateCredentials('client'); + await this.clearCredentials('client'); return true; } async invalidateCredentials( scope: 'all' | 'client' | 'tokens' | 'verifier' | 'discovery', + ): Promise { + if (scope !== 'tokens' && scope !== 'all') { + await this.clearCredentials(scope); + return; + } + const tokensInvalidated = await this.tokenTransaction.invalidateFromSdk(scope); + if (!tokensInvalidated) return; + if (scope === 'all') { + await this.clearCredentials('client'); + await this.clearCredentials('discovery'); + this._codeVerifier = undefined; + } + this.onCredentialsInvalidated?.(scope); + } + + /** Explicit user-driven reset; unlike the SDK invalidation hook, never preserves tokens. */ + async clearCredentials( + scope: 'all' | 'client' | 'tokens' | 'verifier' | 'discovery', ): Promise { if (scope === 'verifier') { this._codeVerifier = undefined; + this.onCredentialsInvalidated?.(scope); return; } if (scope === 'tokens' || scope === 'all') { - this.tokensCache = undefined; - await this.store.remove(`${this.storeKey}${TOKENS_SUFFIX}`); + await this.tokenTransaction.clear(); } if (scope === 'client' || scope === 'all') { - this.clientCache = undefined; await this.store.remove(`${this.storeKey}${CLIENT_SUFFIX}`); + this.clientCache = undefined; } if (scope === 'discovery' || scope === 'all') { - this.discoveryCache = undefined; await this.store.remove(`${this.storeKey}${DISCOVERY_SUFFIX}`); + this.discoveryCache = undefined; } if (scope === 'all') { this._codeVerifier = undefined; } + this.onCredentialsInvalidated?.(scope); } private effectiveRedirectUri(): string { @@ -204,3 +286,14 @@ function registeredRedirectUri(info: OAuthClientInformationMixed | undefined): s const [redirectUri] = info.redirect_uris; return redirectUri; } + +/** + * Route a transport's fetch through the provider's token transaction when one + * is attached, so refresh grants racing on the same credential serialize. + */ +export function createMcpOAuthFetch( + provider: OAuthClientProvider | undefined, + fetchFn: typeof fetch | undefined, +): typeof fetch | undefined { + return provider instanceof McpOAuthClientProvider ? provider.createOAuthFetch(fetchFn) : fetchFn; +} diff --git a/packages/agent-core-v2/src/mcpCore/oauth/service.ts b/packages/agent-core-v2/src/mcpCore/oauth/service.ts index db666275c9..8d5cd7670d 100644 --- a/packages/agent-core-v2/src/mcpCore/oauth/service.ts +++ b/packages/agent-core-v2/src/mcpCore/oauth/service.ts @@ -1,15 +1,44 @@ import { auth, type OAuthClientProvider } from '@modelcontextprotocol/sdk/client/auth.js'; +import type { ILogger as Logger } from '#/_base/log/log'; import { ErrorCodes, Error2, isError2 } from '#/errors'; import { startCallbackServer, type CallbackServer } from './callback-server'; -import { McpOAuthClientProvider } from './provider'; -import { mcpOAuthStoreKey, type McpOAuthStore } from './store'; +import { + META_SUFFIX, + McpOAuthClientProvider, + type McpOAuthStoreMeta, + type StoredMcpOAuthTokens, +} from './provider'; +import { canonicalMcpOAuthResource, mcpOAuthStoreKey, type McpOAuthStore } from './store'; + +const defaultLog: Logger = { + error: () => {}, + warn: () => {}, + info: () => {}, + debug: () => {}, + child: () => defaultLog, +}; export interface McpOAuthServiceOptions { readonly store: McpOAuthStore; readonly clientLabel?: string; readonly resolveClientName?: () => string | undefined; + readonly log?: Logger; + readonly scheduler?: McpOAuthScheduler; + /** Per-request bound for OAuth-flow HTTP (discovery, registration, grants). */ + readonly authRequestTimeoutMs?: number; + /** Upper bound for awaiting in-flight flows and refreshes during shutdown. */ + readonly shutdownDrainTimeoutMs?: number; +} + +export interface McpOAuthScheduledTask { + cancel(): void; +} + +export interface McpOAuthScheduler { + now(): number; + schedule(delayMs: number, task: () => void | Promise): McpOAuthScheduledTask; } export interface BeginAuthorizationOptions { @@ -18,56 +47,357 @@ export interface BeginAuthorizationOptions { export interface BeginAuthorizationResult { readonly authorizationUrl: URL; + /** + * Awaits the OAuth callback, validates `state`, exchanges the code for + * tokens, and persists them via the provider. Resolves on success; + * rejects on abort, timeout, or auth-server error. + * + * Handles sharing one underlying flow (concurrent `beginAuthorization` + * calls for the same credential) run the wait and the exchange exactly + * once: the first `complete()` call's `signal`/`timeoutMs` apply and the + * rest await the same outcome. + */ complete(opts?: { signal?: AbortSignal; timeoutMs?: number }): Promise; + /** + * Detaches this caller without finishing the flow. The callback listener + * stays active while another handle is attached and closes when the final + * handle detaches. Safe to call repeatedly; called automatically by + * `complete()`. + */ cancel(): Promise; } +interface SharedAuthorizationFlow { + readonly attach: () => BeginAuthorizationResult; + readonly cancelUnderlying: () => Promise; +} + +interface ActiveAuthorization { + readonly started: Promise; + readonly controller: AbortController; + readonly serverRef: { current: CallbackServer | undefined }; +} + +export type McpOAuthEvent = + | { + readonly type: 'tokens-saved'; + readonly serverName: string; + readonly serverUrl: string; + } + | { + readonly type: 'tokens-invalidated'; + readonly serverName: string; + readonly serverUrl: string; + readonly scope: 'all' | 'client' | 'tokens' | 'verifier' | 'discovery'; + } + | { + readonly type: 'refresh-failed'; + readonly serverName: string; + readonly serverUrl: string; + readonly error: string; + }; + +export type McpOAuthEventListener = (event: McpOAuthEvent) => void; + +/** Offline credential snapshot for one server/resource identity. */ +export interface McpOAuthTokenState { + readonly hasTokens: boolean; + readonly hasRefreshToken: boolean; + /** Absolute expiry in epoch ms, when the stored grant carries enough data. */ + readonly expiresAt?: number; + readonly expired: boolean; +} + +const REFRESH_AHEAD_MS = 120_000; +const MAX_TIMER_DELAY_MS = 0x7fffffff; +const DEFAULT_AUTH_REQUEST_TIMEOUT_MS = 30_000; +const DEFAULT_SHUTDOWN_DRAIN_TIMEOUT_MS = 30_000; + +const defaultScheduler: McpOAuthScheduler = { + now: () => Date.now(), + schedule: (delayMs, task) => { + const timer = setTimeout(() => void task(), delayMs); + timer.unref(); + return { cancel: () => clearTimeout(timer) }; + }, +}; + export class McpOAuthService { private readonly store: McpOAuthStore; private readonly clientLabel: string | undefined; private readonly resolveClientName: (() => string | undefined) | undefined; + private readonly log: Logger; + private readonly scheduler: McpOAuthScheduler; + private readonly authRequestTimeoutMs: number; + private readonly shutdownDrainTimeoutMs: number; private readonly providers = new Map(); + private readonly listeners = new Set(); + private readonly refreshes = new Map>(); + private readonly refreshTimers = new Map(); + private readonly activeAuthorizations = new Map(); + private readonly backgroundTasks = new Set>(); + private shuttingDown = false; + private shutdownPromise: Promise | undefined; constructor(options: McpOAuthServiceOptions) { this.store = options.store; this.clientLabel = options.clientLabel; this.resolveClientName = options.resolveClientName; + this.log = options.log ?? defaultLog; + this.scheduler = options.scheduler ?? defaultScheduler; + this.authRequestTimeoutMs = options.authRequestTimeoutMs ?? DEFAULT_AUTH_REQUEST_TIMEOUT_MS; + this.shutdownDrainTimeoutMs = options.shutdownDrainTimeoutMs ?? DEFAULT_SHUTDOWN_DRAIN_TIMEOUT_MS; + } + + dispose(): Promise { + return this.shutdown(); } + /** Returns the cached provider for `serverName` + `serverUrl`, constructing it on first use. */ getProvider(serverName: string, serverUrl: string | URL): McpOAuthClientProvider { const storeKey = mcpOAuthStoreKey(serverName, serverUrl); let provider = this.providers.get(storeKey); if (provider === undefined) { - provider = new McpOAuthClientProvider({ - serverName, - serverUrl, - store: this.store, - clientLabel: this.clientLabel, - clientName: this.resolveClientName?.(), - }); + provider = this.createProvider(serverName, serverUrl); this.providers.set(provider.storeKey, provider); } return provider; } + /** True once the provider has persisted tokens for this server/resource identity. */ async hasTokens(serverName: string, serverUrl: string | URL): Promise { return (await this.getProvider(serverName, serverUrl).tokens()) !== undefined; } + /** + * Offline view of the stored grant. `expired` is only computable when the + * tokens were written with an `obtained_at` stamp and carry `expires_in`; + * older or foreign writes without both are treated as non-expiring. + */ + async tokenState(serverName: string, serverUrl: string | URL): Promise { + const tokens = (await this.getProvider(serverName, serverUrl).tokens()) as + | StoredMcpOAuthTokens + | undefined; + if (tokens === undefined) { + return { hasTokens: false, hasRefreshToken: false, expired: false }; + } + const expiresAt = + typeof tokens.obtained_at === 'number' && typeof tokens.expires_in === 'number' + ? tokens.obtained_at + tokens.expires_in * 1000 + : undefined; + return { + hasTokens: true, + hasRefreshToken: typeof tokens.refresh_token === 'string' && tokens.refresh_token.length > 0, + expiresAt, + expired: expiresAt !== undefined && this.scheduler.now() >= expiresAt, + }; + } + + onEvent(listener: McpOAuthEventListener): () => void { + this.listeners.add(listener); + return () => { + this.listeners.delete(listener); + }; + } + + protected trackBackgroundTask(task: Promise): void { + this.backgroundTasks.add(task); + void task.then( + () => this.backgroundTasks.delete(task), + () => this.backgroundTasks.delete(task), + ); + } + + /** + * Single-flight token refresh per credential: concurrent callers share one + * in-flight SDK `auth()` run, so two sessions expiring together cannot race + * a rotating refresh token. Resolves when the grant is usable again; + * rejects when the refresh token was rejected (or never existed) and an + * interactive login is required. + */ + async refresh(serverName: string, serverUrl: string | URL): Promise { + const storeKey = mcpOAuthStoreKey(serverName, serverUrl); + const existing = this.refreshes.get(storeKey); + if (existing !== undefined) return existing; + if (this.shuttingDown) { + throw new Error2(ErrorCodes.MCP_OAUTH_FAILED, 'MCP OAuth service is shutting down'); + } + const task = this.refreshNow(serverName, serverUrl).finally(() => { + this.refreshes.delete(storeKey); + }); + this.refreshes.set(storeKey, task); + return task; + } + + /** + * Arm the proactive refresh timer for every stored credential that carries + * enough data to expire. Called once at engine start; subsequent token + * writes re-arm through the provider save hook. A malformed meta sidecar + * (or any per-credential failure) is skipped with a warning rather than + * aborting the whole sweep. + */ + async sweepProactiveRefresh(): Promise { + if (this.shuttingDown) return; + const keys = await this.store.list(); + for (const key of keys) { + if (this.shuttingDown) return; + if (!key.endsWith(META_SUFFIX)) continue; + const meta = await readStoreMeta(this.store, key, this.log); + if (meta === undefined) continue; + try { + const state = await this.tokenState(meta.serverName, meta.serverUrl); + if (!state.hasTokens || !state.hasRefreshToken || state.expiresAt === undefined) continue; + this.scheduleRefresh(meta.serverName, meta.serverUrl, state.expiresAt); + } catch (error) { + this.log.warn('skipping MCP OAuth credential during proactive-refresh sweep', { + file: key, + error: error instanceof Error ? error : String(error), + }); + } + } + } + + /** Clear every pending proactive-refresh timer (engine shutdown, tests). */ + stopProactiveRefresh(): void { + for (const timer of this.refreshTimers.values()) timer.cancel(); + this.refreshTimers.clear(); + } + + /** + * Release everything the service owns: pending proactive-refresh timers, + * in-flight refreshes and interactive flows (closing their callback + * listeners), event listeners, and cached providers. Idempotent. + */ + shutdown(): Promise { + if (this.shutdownPromise !== undefined) return this.shutdownPromise; + this.shuttingDown = true; + this.stopProactiveRefresh(); + const authorizations = [...this.activeAuthorizations.values()]; + const refreshes = [...this.refreshes.values()]; + this.activeAuthorizations.clear(); + const deadline = this.drainDeadline(); + this.shutdownPromise = (async () => { + try { + await Promise.race([ + Promise.all([ + Promise.all( + authorizations.map(async (active) => { + active.controller.abort(); + await active.serverRef.current?.close().catch(() => undefined); + const flow = await active.started.catch(() => undefined); + await flow?.cancelUnderlying(); + }), + ), + Promise.allSettled(refreshes), + this.drainBackgroundTasks(), + ]), + deadline.promise, + ]); + } finally { + deadline.cancel(); + this.listeners.clear(); + this.providers.clear(); + } + })(); + return this.shutdownPromise; + } + + private async drainBackgroundTasks(): Promise { + for (;;) { + await Promise.allSettled(this.backgroundTasks); + await new Promise((resolve) => { + setImmediate(resolve); + }); + if (this.backgroundTasks.size === 0) return; + } + } + + private drainDeadline(): { readonly promise: Promise; readonly cancel: () => void } { + let task: McpOAuthScheduledTask | undefined; + const promise = new Promise((resolve) => { + task = this.scheduler.schedule(this.shutdownDrainTimeoutMs, () => { + this.log.warn('mcp oauth shutdown drain timed out; continuing teardown'); + resolve(); + }); + }); + return { promise, cancel: () => task?.cancel() }; + } + + private authFetch( + provider: McpOAuthClientProvider, + signals: readonly AbortSignal[] = [], + ): typeof fetch { + const fetchFn = provider.createOAuthFetch(); + const timeoutMs = this.authRequestTimeoutMs; + return (async (input: Parameters[0], init?: Parameters[1]) => { + const combined: AbortSignal[] = [AbortSignal.timeout(timeoutMs), ...signals]; + if (init?.signal !== undefined && init.signal !== null) combined.push(init.signal); + return fetchFn(input, { ...init, signal: AbortSignal.any(combined) }); + }) as typeof fetch; + } + + /** + * Drive the SDK `auth()` orchestrator far enough to surface an + * authorization URL. The caller is responsible for displaying the URL + * (typically via the synthetic authenticate tool) and then awaiting + * `complete()` to finish the code exchange. + * + * Interactive flows are serialized per credential: while one flow for a + * store key is in flight, further calls join it — same URL, shared + * `complete()`, and a `cancel()` that only detaches the caller — instead + * of resetting the shared provider's PKCE/state mid-flow. + */ async beginAuthorization( serverName: string, serverUrl: string | URL, options: BeginAuthorizationOptions = {}, ): Promise { - const provider = options.clientLabel === undefined - ? this.getProvider(serverName, serverUrl) - : new McpOAuthClientProvider({ - serverName, - serverUrl, - store: this.store, - clientLabel: options.clientLabel, - clientName: this.resolveClientName?.(), - }); + if (this.shuttingDown) { + throw new Error2(ErrorCodes.MCP_OAUTH_FAILED, 'MCP OAuth service is shutting down'); + } + const storeKey = mcpOAuthStoreKey(serverName, serverUrl); + await this.refreshes.get(storeKey)?.catch(() => undefined); + if (this.shuttingDown) { + throw new Error2(ErrorCodes.MCP_OAUTH_FAILED, 'MCP OAuth service is shutting down'); + } + const inFlight = this.activeAuthorizations.get(storeKey); + if (inFlight !== undefined) { + const flow = await inFlight.started; + return flow.attach(); + } + + const controller = new AbortController(); + const serverRef: { current: CallbackServer | undefined } = { current: undefined }; + const started = this.startAuthorizationFlow( + serverName, + serverUrl, + options, + controller.signal, + serverRef, + ); + this.activeAuthorizations.set(storeKey, { started, controller, serverRef }); + let flow: SharedAuthorizationFlow; + try { + flow = await started; + } catch (error) { + this.activeAuthorizations.delete(storeKey); + throw error; + } + return flow.attach(); + } + + private async startAuthorizationFlow( + serverName: string, + serverUrl: string | URL, + options: BeginAuthorizationOptions, + signal: AbortSignal, + serverRef: { current: CallbackServer | undefined }, + ): Promise { + const storeKey = mcpOAuthStoreKey(serverName, serverUrl); + const provider = + options.clientLabel === undefined + ? this.getProvider(serverName, serverUrl) + : this.createProvider(serverName, serverUrl, options.clientLabel); if (options.clientLabel !== undefined) { this.providers.set(provider.storeKey, provider); } @@ -80,24 +410,48 @@ export class McpOAuthService { } catch (error) { throw wrapAuthError('failed to start OAuth callback listener', error); } - - provider.setRedirectUrl(new URL(callbackServer.redirectUri)); - await provider.ready; - await provider.invalidateStaleRegistration(callbackServer.redirectUri); + serverRef.current = callbackServer; let authorizationUrl: URL | undefined; try { - const result = await auth(provider as OAuthClientProvider, { serverUrl }); - if (result !== 'REDIRECT') { - await callbackServer.close(); - throw new AlreadyAuthorizedError(serverName); - } - authorizationUrl = provider.takeAuthorizationUrl(); - if (authorizationUrl === undefined) { - throw new Error2( - ErrorCodes.MCP_OAUTH_FAILED, - 'OAuth provider did not capture an authorization URL', - ); + provider.setRedirectUrl(new URL(callbackServer.redirectUri)); + await provider.ready; + await provider.invalidateStaleRegistration(callbackServer.redirectUri); + let tokensSaved = false; + const unsubscribeTokensSaved = this.onEvent((event) => { + if ( + event.type === 'tokens-saved' && + event.serverName === serverName && + event.serverUrl === canonicalMcpOAuthResource(serverUrl) + ) { + tokensSaved = true; + } + }); + try { + const result = await auth(provider as OAuthClientProvider, { + serverUrl, + fetchFn: this.authFetch(provider, [signal]), + }); + if (result !== 'REDIRECT') { + await callbackServer.close(); + if (!tokensSaved) { + this.emit({ + type: 'tokens-saved', + serverName, + serverUrl: canonicalMcpOAuthResource(serverUrl), + }); + } + throw new AlreadyAuthorizedError(serverName); + } + authorizationUrl = provider.takeAuthorizationUrl(); + if (authorizationUrl === undefined) { + throw new Error2( + ErrorCodes.MCP_OAUTH_FAILED, + 'OAuth provider did not capture an authorization URL', + ); + } + } finally { + unsubscribeTokensSaved(); } } catch (error) { await callbackServer.close().catch(() => undefined); @@ -107,61 +461,235 @@ export class McpOAuthService { } let settled = false; - const cancel = async (): Promise => { + let completion: Promise | undefined; + let attachedHandles = 0; + const settle = async (): Promise => { if (settled) return; settled = true; - await callbackServer.close().catch(() => undefined); + this.activeAuthorizations.delete(storeKey); provider.resetFlow(); + await callbackServer.close().catch(() => undefined); }; - const complete: BeginAuthorizationResult['complete'] = async (opts = {}) => { + const startCompletion: BeginAuthorizationResult['complete'] = (opts = {}) => { + if (completion !== undefined) return completion; if (settled) { - throw new Error2(ErrorCodes.MCP_OAUTH_FAILED, 'OAuth flow already completed or cancelled'); + return Promise.reject( + new Error2(ErrorCodes.MCP_OAUTH_FAILED, 'OAuth flow already completed or cancelled'), + ); } - try { - const { code, state } = await callbackServer.waitForCode({ - signal: opts.signal, - timeoutMs: opts.timeoutMs, - }); - const expectedState = provider.expectedState(); - if (expectedState !== undefined && state !== expectedState) { - throw new Error2( - ErrorCodes.MCP_OAUTH_FAILED, - 'OAuth state mismatch — possible CSRF; refusing token exchange', - ); - } - const finalResult = await auth(provider as OAuthClientProvider, { - serverUrl, - authorizationCode: code, - }); - if (finalResult !== 'AUTHORIZED') { - throw new Error2( - ErrorCodes.MCP_OAUTH_FAILED, - `OAuth code exchange returned "${finalResult}" instead of AUTHORIZED`, - { details: { result: finalResult } }, - ); + completion = (async () => { + try { + const { code, state } = await callbackServer.waitForCode({ + signal: opts.signal, + timeoutMs: opts.timeoutMs, + }); + const expectedState = provider.expectedState(); + if (expectedState !== undefined && state !== expectedState) { + throw new Error2( + ErrorCodes.MCP_OAUTH_FAILED, + 'OAuth state mismatch — possible CSRF; refusing token exchange', + ); + } + const finalResult = await auth(provider as OAuthClientProvider, { + serverUrl, + authorizationCode: code, + fetchFn: this.authFetch( + provider, + opts.signal === undefined ? [signal] : [signal, opts.signal], + ), + }); + if (finalResult !== 'AUTHORIZED') { + throw new Error2( + ErrorCodes.MCP_OAUTH_FAILED, + `OAuth code exchange returned "${finalResult}" instead of AUTHORIZED`, + { details: { result: finalResult } }, + ); + } + } catch (error) { + await settle(); + throw wrapAuthError(`OAuth flow for "${serverName}" failed`, error); } - } catch (error) { - await cancel(); - throw wrapAuthError(`OAuth flow for "${serverName}" failed`, error); - } - settled = true; - await callbackServer.close().catch(() => undefined); - provider.resetFlow(); + await settle(); + })(); + this.trackBackgroundTask(completion); + return completion; }; - return { authorizationUrl, complete, cancel }; + const attach = (): BeginAuthorizationResult => { + attachedHandles += 1; + let detached = false; + const detach = async (): Promise => { + if (detached) return; + detached = true; + attachedHandles -= 1; + if (attachedHandles === 0) await settle(); + }; + return { + authorizationUrl, + complete: async (opts = {}) => { + if (detached) { + throw new Error2( + ErrorCodes.MCP_OAUTH_FAILED, + 'OAuth flow already completed or cancelled', + ); + } + try { + await startCompletion(opts); + } finally { + await detach(); + } + }, + cancel: detach, + }; + }; + + return { + attach, + cancelUnderlying: settle, + }; } + /** + * Clear stored credentials for a server. Use `'all'` after the user + * explicitly signs out; use `'tokens'` to force a re-auth while keeping + * the registered DCR client. + */ invalidate( serverName: string, serverUrl: string | URL, scope: 'all' | 'client' | 'tokens' | 'discovery' = 'all', ): Promise { - return this.getProvider(serverName, serverUrl).invalidateCredentials(scope); + return this.getProvider(serverName, serverUrl).clearCredentials(scope); + } + + /** + * Drop the cached provider for a credential. After an invalidation this + * guarantees the next `beginAuthorization` starts from a clean in-memory + * flow state (files are always re-read, so this is defensive). + */ + forgetProvider(serverName: string, serverUrl: string | URL): void { + this.providers.delete(mcpOAuthStoreKey(serverName, serverUrl)); + } + + private createProvider( + serverName: string, + serverUrl: string | URL, + clientLabel?: string, + ): McpOAuthClientProvider { + const canonicalUrl = canonicalMcpOAuthResource(serverUrl); + return new McpOAuthClientProvider({ + serverName, + serverUrl, + store: this.store, + clientLabel: clientLabel ?? this.clientLabel, + clientName: this.resolveClientName?.(), + now: () => this.scheduler.now(), + track: (task) => { + this.trackBackgroundTask(task); + }, + onTokensSaved: (tokens) => { + this.emit({ type: 'tokens-saved', serverName, serverUrl: canonicalUrl }); + if ( + typeof tokens.obtained_at === 'number' && + typeof tokens.expires_in === 'number' && + typeof tokens.refresh_token === 'string' && + tokens.refresh_token.length > 0 + ) { + this.scheduleRefresh( + serverName, + canonicalUrl, + tokens.obtained_at + tokens.expires_in * 1000, + ); + } + }, + onCredentialsInvalidated: (scope) => { + if (scope === 'tokens' || scope === 'all') { + this.cancelScheduledRefresh(serverName, canonicalUrl); + } + this.emit({ type: 'tokens-invalidated', serverName, serverUrl: canonicalUrl, scope }); + }, + }); + } + + private async refreshNow(serverName: string, serverUrl: string | URL): Promise { + if (this.activeAuthorizations.has(mcpOAuthStoreKey(serverName, serverUrl))) return; + const state = await this.tokenState(serverName, serverUrl); + if (this.activeAuthorizations.has(mcpOAuthStoreKey(serverName, serverUrl))) return; + if (!state.hasTokens || !state.hasRefreshToken) { + throw new Error2( + ErrorCodes.MCP_OAUTH_FAILED, + `MCP server "${serverName}" has no refreshable OAuth grant`, + ); + } + const provider = this.getProvider(serverName, serverUrl); + provider.resetFlow(); + try { + const result = await auth(provider as OAuthClientProvider, { + serverUrl, + fetchFn: this.authFetch(provider), + }); + if (result !== 'AUTHORIZED') { + throw new Error2( + ErrorCodes.MCP_OAUTH_FAILED, + 'the stored OAuth grant requires an interactive login', + ); + } + } finally { + provider.resetFlow(); + } + } + + private scheduleRefresh(serverName: string, serverUrl: string | URL, expiresAt: number): void { + if (this.shuttingDown) return; + const canonicalUrl = canonicalMcpOAuthResource(serverUrl); + const storeKey = mcpOAuthStoreKey(serverName, canonicalUrl); + this.cancelScheduledRefresh(serverName, canonicalUrl); + const now = this.scheduler.now(); + if (expiresAt <= now) return; + const lifetimeMs = expiresAt - now; + const refreshAheadMs = Math.min(REFRESH_AHEAD_MS, lifetimeMs / 2); + const delay = lifetimeMs - refreshAheadMs; + let timer: McpOAuthScheduledTask; + if (delay > MAX_TIMER_DELAY_MS) { + timer = this.scheduler.schedule(MAX_TIMER_DELAY_MS, () => { + this.refreshTimers.delete(storeKey); + this.scheduleRefresh(serverName, canonicalUrl, expiresAt); + }); + } else { + timer = this.scheduler.schedule(delay, async () => { + this.refreshTimers.delete(storeKey); + await this.refresh(serverName, canonicalUrl).catch((error: unknown) => { + this.emit({ + type: 'refresh-failed', + serverName, + serverUrl: canonicalUrl, + error: error instanceof Error ? error.message : String(error), + }); + }); + }); + } + this.refreshTimers.set(storeKey, timer); + } + + private cancelScheduledRefresh(serverName: string, serverUrl: string | URL): void { + const storeKey = mcpOAuthStoreKey(serverName, serverUrl); + const timer = this.refreshTimers.get(storeKey); + timer?.cancel(); + this.refreshTimers.delete(storeKey); + } + + private emit(event: McpOAuthEvent): void { + for (const listener of this.listeners) { + try { + listener(event); + } catch { + } + } } } +/** Thrown by `beginAuthorization` when stored tokens already satisfy the server. */ export class AlreadyAuthorizedError extends Error2 { constructor(serverName: string) { super( @@ -172,6 +700,29 @@ export class AlreadyAuthorizedError extends Error2 { } } +async function readStoreMeta( + store: McpOAuthStore, + key: string, + log: Logger, +): Promise { + const raw: unknown = await store.read(key); + if (raw === undefined) return undefined; + if (typeof raw !== 'object' || raw === null) { + log.warn('ignoring malformed MCP OAuth meta file', { file: key }); + return undefined; + } + const { serverName, serverUrl } = raw as Record; + if (typeof serverName !== 'string' || serverName.length === 0 || typeof serverUrl !== 'string') { + log.warn('ignoring malformed MCP OAuth meta file', { file: key }); + return undefined; + } + if (URL.parse(serverUrl) === null) { + log.warn('ignoring MCP OAuth meta file with unparseable serverUrl', { file: key, serverUrl }); + return undefined; + } + return { serverName, serverUrl }; +} + function wrapAuthError(prefix: string, error: unknown): Error2 { if (isError2(error)) { return error; diff --git a/packages/agent-core-v2/src/mcpCore/oauth/store.ts b/packages/agent-core-v2/src/mcpCore/oauth/store.ts index 67d2ea9669..a858debe85 100644 --- a/packages/agent-core-v2/src/mcpCore/oauth/store.ts +++ b/packages/agent-core-v2/src/mcpCore/oauth/store.ts @@ -5,7 +5,9 @@ import { basename } from 'pathe'; import { ErrorCodes, Error2 } from '#/errors'; export function sanitizeStoreKey(name: string): string { - const safe = basename(name).replaceAll(/[^a-zA-Z0-9_-]/g, '_').replaceAll(/_+/g, '_'); + const safe = basename(name) + .replaceAll(/[^a-zA-Z0-9_-]/g, '_') + .replaceAll(/_+/g, '_'); if (safe.length === 0 || safe.startsWith('.')) { throw new Error2(ErrorCodes.CONFIG_INVALID, `Invalid MCP OAuth store key: "${name}"`); } @@ -34,4 +36,5 @@ export interface McpOAuthStore { read(key: string): Promise; write(key: string, data: unknown): Promise; remove(key: string): Promise; + list(prefix?: string): Promise; } diff --git a/packages/agent-core-v2/src/program/program.ts b/packages/agent-core-v2/src/program/program.ts index 961fa450de..e565993d23 100644 --- a/packages/agent-core-v2/src/program/program.ts +++ b/packages/agent-core-v2/src/program/program.ts @@ -290,8 +290,8 @@ export class Program { const watch = own(new WorkspaceFsWatchService(this.context, dirs, runtime.watch!, runtime.fs!)); const instructions = own(new WorkspaceInstructionsService(this.context, runtime.fs!, runtime.environment, this.dependencies.bootstrap, runtime.watch!, this.dependencies.log, state)); const trust = own(new WorkspaceTrustService(this.context, this.dependencies.docs, state)); - const mcpConfig = own(new WorkspaceMcpConfigService(this.context, this.dependencies.bootstrap, this.dependencies.plugins, this.dependencies.log, this.dependencies.config, runtime.watch!, runtime.fs!, trust)); - const mcp = own(new WorkspaceMcpService(this.context, this.resolver, mcpConfig, this.dependencies.oauthStore, this.dependencies.log, this.dependencies.telemetry, this.dependencies.identity, this.dependencies.sessionManager)); + const mcpConfig = own(new WorkspaceMcpConfigService(this.context, this.dependencies.bootstrap, this.dependencies.plugins, this.dependencies.log, this.dependencies.config, runtime.watch!, runtime.fs!, trust, this.dependencies.configStore)); + const mcp = own(new WorkspaceMcpService(this.context, this.resolver, mcpConfig, this.dependencies.oauth, this.dependencies.log, this.dependencies.telemetry, this.dependencies.identity, this.dependencies.sessionManager)); const userAgentProfiles = own(new UserAgentProfileLoaderService(this.dependencies.bootstrap, runtime.fs!, this.dependencies.log, this.dependencies.builtinAgentProfiles, this.context, this.dependencies.agentProfiles)); const pluginAgentProfiles = own(new PluginAgentProfileLoaderService(this.dependencies.plugins, runtime.fs!, this.dependencies.log, userAgentProfiles, this.context, this.dependencies.agentProfiles)); const explicitAgentProfiles = own(new ExplicitAgentProfileLoaderService(this.context, this.dependencies.bootstrap, runtime.fs!, this.dependencies.log, userAgentProfiles, this.dependencies.agentProfiles)); diff --git a/packages/agent-core-v2/src/program/programDependencies.ts b/packages/agent-core-v2/src/program/programDependencies.ts index edd33be7a0..1df99afab9 100644 --- a/packages/agent-core-v2/src/program/programDependencies.ts +++ b/packages/agent-core-v2/src/program/programDependencies.ts @@ -6,7 +6,8 @@ import type { IAgentProfileRegistry } from '#/app/agentProfileCatalog/agentProfi import type { IBootstrapService } from '#/app/bootstrap/bootstrap'; import type { IConfigService } from '#/app/config/config'; import type { IGitService } from '#/app/git/git'; -import type { IMcpOAuthStore } from '#/app/mcpConfig/oauthStore'; +import type { McpOAuthService } from '#/mcpCore/oauth/service'; +import type { IMcpConfigStore } from '#/app/mcpConfig/configStore'; import type { IPluginService } from '#/app/plugin/plugin'; import type { ISessionManager } from '#/app/sessionManager/sessionManager'; import type { IBuiltinSkillSource } from '#/app/skillCatalog/builtinSkillSource'; @@ -48,7 +49,8 @@ export interface ProgramDependencies { readonly git: LiveRef; readonly identity: IAgentIdentity; readonly log: ILogService; - readonly oauthStore: IMcpOAuthStore; + readonly oauth: McpOAuthService; + readonly configStore: IMcpConfigStore; readonly plugins: IPluginService; readonly sessionManager: LiveRef; readonly agentProfiles: IAgentProfileRegistry; diff --git a/packages/agent-core-v2/src/runtime/localRuntime.ts b/packages/agent-core-v2/src/runtime/localRuntime.ts index 932d935d48..c86104361e 100644 --- a/packages/agent-core-v2/src/runtime/localRuntime.ts +++ b/packages/agent-core-v2/src/runtime/localRuntime.ts @@ -8,7 +8,7 @@ import { IHostFsWatchService } from '#/os/interface/hostFsWatch'; import { IHostProcessService } from '#/os/interface/hostProcess'; import { IHostTerminalService } from '#/os/interface/terminal'; -import type { Runtime, RuntimePath, RuntimeStatus } from './runtime'; +import type { Runtime, RuntimeCapability, RuntimePath, RuntimeStatus } from './runtime'; import type { RuntimeProviderAttachment, RuntimeProviderContext, RuntimeProviderFactory } from './runtimeProvider'; import type { RuntimeProviderHost } from './runtimeUnitHost'; @@ -16,7 +16,7 @@ let nextGeneration = 1; export class LocalRuntime implements Runtime { readonly identity; - readonly capabilities = new Set(['fs', 'process', 'watch', 'terminal'] as const); + readonly capabilities: ReadonlySet; readonly environment; readonly path: RuntimePath; readonly workspace: Runtime['workspace']; @@ -31,12 +31,18 @@ export class LocalRuntime implements Runtime { constructor( workspaceId: string, environment: IHostEnvironment, - fs: IHostFileSystem, - process: IHostProcessService, - watch: IHostFsWatchService, - terminal: IHostTerminalService, + fs: IHostFileSystem | undefined, + process: IHostProcessService | undefined, + watch: IHostFsWatchService | undefined, + terminal: IHostTerminalService | undefined, ) { this.identity = { workspaceId, runtimeId: 'local', generation: `local-${nextGeneration++}` }; + const capabilities = new Set(); + if (fs !== undefined) capabilities.add('fs'); + if (process !== undefined) capabilities.add('process'); + if (watch !== undefined) capabilities.add('watch'); + if (terminal !== undefined) capabilities.add('terminal'); + this.capabilities = capabilities; this.environment = { osKind: environment.osKind, osArch: environment.osArch, diff --git a/packages/agent-core-v2/src/workspace/workspaceInstance/workspaceInstanceManager.ts b/packages/agent-core-v2/src/workspace/workspaceInstance/workspaceInstanceManager.ts index 9395a99811..55c263b7d1 100644 --- a/packages/agent-core-v2/src/workspace/workspaceInstance/workspaceInstanceManager.ts +++ b/packages/agent-core-v2/src/workspace/workspaceInstance/workspaceInstanceManager.ts @@ -22,6 +22,7 @@ export interface IWorkspaceInstanceManager { getOrCreate(ref: WorkspaceInstanceRef): Promise; get(workspaceId: string): WorkspaceInstance | undefined; findByRoot(root: string): WorkspaceInstance | undefined; + findContaining(cwd: string): WorkspaceInstance | undefined; list(): readonly WorkspaceInstance[]; snapshot(): WorkspaceInstancesSnapshot; close(workspaceId: string): Promise; diff --git a/packages/agent-core-v2/src/workspace/workspaceInstance/workspaceInstanceManagerService.ts b/packages/agent-core-v2/src/workspace/workspaceInstance/workspaceInstanceManagerService.ts index ba21976227..21a9c1cd07 100644 --- a/packages/agent-core-v2/src/workspace/workspaceInstance/workspaceInstanceManagerService.ts +++ b/packages/agent-core-v2/src/workspace/workspaceInstance/workspaceInstanceManagerService.ts @@ -11,7 +11,9 @@ import { ICronTaskPersistence } from '#/app/cron/cronTaskPersistence'; import { IEventService } from '#/app/event/event'; import { IFlagService } from '#/app/flag/flag'; import { IGitService } from '#/app/git/git'; -import { IMcpOAuthStore } from '#/app/mcpConfig/oauthStore'; +import { IMcpOAuthService } from '#/app/mcpConfig/oauthService'; +import type { McpOAuthService } from '#/mcpCore/oauth/service'; +import { IMcpConfigStore } from '#/app/mcpConfig/configStore'; import { IPluginService } from '#/app/plugin/plugin'; import { ISessionIndex, ISessionIndexMirror } from '#/app/sessionIndex/sessionIndex'; import { ISessionManager } from '#/app/sessionManager/sessionManager'; @@ -28,6 +30,7 @@ import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStor import { Error2, ErrorCodes } from '#/errors'; import { IHostEnvironment } from '#/os/interface/hostEnvironment'; import { LocalRuntimeProviderFactory } from '#/runtime/localRuntime'; +import { canonicalWorkspaceRoot } from '#/_base/utils/paths'; import type { Runtime, RuntimeBinding, RuntimeCapability, RuntimeLease } from '#/runtime/runtime'; import { RuntimeError, RuntimeRegistry } from '#/runtime/runtimeRegistry'; import type { RuntimeProviderFactory } from '#/runtime/runtimeProvider'; @@ -64,7 +67,8 @@ export class WorkspaceInstanceManager implements IWorkspaceInstanceManager { @ILogService private readonly log: ILogService, @IModelCatalog private readonly modelCatalog: IModelCatalog, @IModelService private readonly models: IModelService, - @IMcpOAuthStore private readonly oauthStore: IMcpOAuthStore, + @IMcpOAuthService private readonly oauth: McpOAuthService, + @IMcpConfigStore private readonly configStore: IMcpConfigStore, @IPluginService private readonly plugins: IPluginService, @IProviderService private readonly modelProviders: IProviderService, @ref(ISessionManager) private readonly sessionManager: LiveRef, @@ -88,6 +92,20 @@ export class WorkspaceInstanceManager implements IWorkspaceInstanceManager { return [...this.instances.values()].find((instance) => instance.root.replace(/[\\/]$/, '') === normalized); } + findContaining(cwd: string): WorkspaceInstance | undefined { + const probe = canonicalWorkspaceRoot(cwd); + let best: { readonly instance: WorkspaceInstance; readonly rootLength: number } | undefined; + for (const instance of this.instances.values()) { + const root = canonicalWorkspaceRoot(instance.root); + const prefix = root.endsWith('/') ? root : `${root}/`; + if (probe !== root && !probe.startsWith(prefix)) continue; + if (best === undefined || root.length > best.rootLength) { + best = { instance, rootLength: root.length }; + } + } + return best?.instance; + } + list(): readonly WorkspaceInstance[] { return [...this.instances.values()]; } @@ -191,7 +209,8 @@ export class WorkspaceInstanceManager implements IWorkspaceInstanceManager { git: this.git, identity: this.identity, log: this.log, - oauthStore: this.oauthStore, + oauth: this.oauth, + configStore: this.configStore, plugins: this.plugins, sessionManager: this.sessionManager, agentProfiles: this.agentProfiles, diff --git a/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts b/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts index 951088adbf..1103949a42 100644 --- a/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts +++ b/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts @@ -1,14 +1,18 @@ -import { Disposable } from '#/_base/di/lifecycle'; import { ref, type LiveRef } from '#/_base/di/instantiation'; +import { Disposable } from '#/_base/di/lifecycle'; import { ILogService } from '#/_base/log/log'; - -import { McpConnectionManager, type McpConnectionView } from '#/mcpCore/connection-manager'; -import type { McpServerConfig } from '#/mcpCore/config-schema'; -import { McpOAuthService } from '#/mcpCore/oauth/service'; import { IAgentIdentity } from '#/app/agentIdentity/agentIdentity'; -import { IMcpOAuthStore } from '#/app/mcpConfig/oauthStore'; -import { ITelemetryService } from '#/app/telemetry/telemetry'; +import { IMcpOAuthService } from '#/app/mcpConfig/oauthService'; import { ISessionManager } from '#/app/sessionManager/sessionManager'; +import { ITelemetryService } from '#/app/telemetry/telemetry'; +import type { McpServerConfig } from '#/mcpCore/config-schema'; +import { + McpConnectionManager, + type McpConnectionView, + type McpServerEntry, +} from '#/mcpCore/connection-manager'; +import type { McpOAuthEvent, McpOAuthService } from '#/mcpCore/oauth/service'; +import { canonicalMcpOAuthResource } from '#/mcpCore/oauth/store'; import { ISessionEphemeralMcpServers } from '#/session/mcp/ephemeralMcpServers'; import { MergedMcpConnectionView } from '#/session/mcp/mergedConnectionView'; import { ISessionMcpHandle } from '#/session/mcp/sessionMcpHandle'; @@ -43,7 +47,7 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ @IWorkspaceContext workspace: IWorkspaceContext, @IRuntimeResolver private readonly runtimeResolver: IRuntimeResolver, @IWorkspaceMcpConfigService private readonly mcpConfig: IWorkspaceMcpConfigService, - @IMcpOAuthStore oauthStore: IMcpOAuthStore, + @IMcpOAuthService oauthService: McpOAuthService, @ILogService private readonly log: ILogService, @ITelemetryService private readonly telemetry: ITelemetryService, @IAgentIdentity private readonly identity: IAgentIdentity, @@ -53,10 +57,7 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ this.sessionLifecycle = sessionLifecycle; this.stdioCwd = workspace.cwd; this.workspaceId = workspace.workspaceId; - this.oauthService = new McpOAuthService({ - store: oauthStore, - resolveClientName: this.resolveClientName, - }); + this.oauthService = oauthService; this.manager = new McpConnectionManager({ log: this.log, oauthService: this.oauthService, @@ -70,9 +71,10 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ this._register({ dispose: () => void this.manager.shutdown() }); this._register( this.mcpConfig.onDidChange((change) => { - this.scheduleApply(change); + change.waitUntil(this.scheduleApply(change)); }), ); + this._register({ dispose: this.oauthEventSubscription(this.manager) }); this.attachSessionLifecycle(); this._register(sessionLifecycle.onDidChange(() => this.attachSessionLifecycle())); this.ready = this.initialize().catch((error: unknown) => { @@ -134,6 +136,7 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ .catch((error: unknown) => { this.log.error('session mcp overlay initial load failed', { error }); }); + const unsubscribeOAuth = this.oauthEventSubscription(sessionManager); const view = new MergedMcpConnectionView( this.manager, sessionManager, @@ -147,10 +150,70 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ connectionManager: view, isBaselineServer: this.sessionBaseline(this.manager, this.ready, Object.keys(servers)), }, - shutdown: () => sessionManager.shutdown(), + shutdown: () => { + unsubscribeOAuth(); + return sessionManager.shutdown(); + }, }; } + private oauthEventSubscription(manager: McpConnectionManager): () => void { + return this.oauthService.onEvent((event) => { + void this.handleMcpOAuthEvent(manager, event).catch((error: unknown) => { + this.log.warn(`mcp oauth event handling failed: ${String(error)}`); + }); + }); + } + + private async handleMcpOAuthEvent( + manager: McpConnectionManager, + event: McpOAuthEvent, + ): Promise { + if (event.type === 'tokens-invalidated' && event.scope !== 'tokens' && event.scope !== 'all') { + return; + } + const entry = manager.get(event.serverName); + if (entry === undefined) return; + const serverUrl = manager.getRemoteServerUrl(event.serverName); + if (serverUrl === undefined || canonicalMcpOAuthResource(serverUrl) !== event.serverUrl) return; + if (event.type === 'tokens-invalidated') { + this.oauthService.forgetProvider(event.serverName, event.serverUrl); + } + if (entry.status === 'disabled' || entry.status === 'removed') return; + if (entry.status === 'pending') { + await new Promise((resolve, reject) => { + let unsubscribe = (): void => {}; + let settled = false; + const reconnect = (next: McpServerEntry | undefined): void => { + if (settled) return; + if (next !== undefined && (next.name !== event.serverName || next.status === 'pending')) { + return; + } + settled = true; + unsubscribe(); + if (next === undefined || next.status === 'disabled' || next.status === 'removed') { + resolve(); + return; + } + void manager.reconnectAfterCurrent(event.serverName).then(resolve, reject); + }; + unsubscribe = manager.onStatusChange(reconnect); + if (settled) unsubscribe(); + else reconnect(manager.get(event.serverName)); + }); + return; + } + if ( + event.type === 'tokens-saved' && + entry.status !== 'needs-auth' && + entry.status !== 'failed' + ) { + return; + } + if (event.type === 'refresh-failed' && entry.status !== 'connected') return; + await manager.reconnectAndJoin(event.serverName); + } + private sessionBaseline( view: McpConnectionView, ready: Promise, @@ -202,8 +265,8 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ this.trackMcpInitialLoad(); } - private scheduleApply(change: McpServersChange): void { - void this.ready + private scheduleApply(change: McpServersChange): Promise { + return this.ready .then(() => this.mutate(() => this.apply(change))) .catch((error) => { this.log.warn(`mcp server change apply failed: ${String(error)}`); @@ -241,4 +304,3 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ } } } - diff --git a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/internal/config-loader.ts b/packages/agent-core-v2/src/workspace/workspaceMcpConfig/internal/config-loader.ts deleted file mode 100644 index fa6302ac91..0000000000 --- a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/internal/config-loader.ts +++ /dev/null @@ -1,127 +0,0 @@ -import { dirname, isAbsolute, join, normalize, resolve } from 'pathe'; - -import { findGitWorkTree } from '#/app/git/workTree'; -import { resolveKimiHome } from '#/app/bootstrap/bootstrap'; -import { OsFsErrors, HostFsError } from '#/os/interface/hostFsErrors'; -import type { IHostFileSystem } from '#/os/interface/hostFileSystem'; -import { McpServerConfigSchema, type McpServerConfig } from '#/mcpCore/config-schema'; -import { ErrorCodes, Error2 } from '#/errors'; -import { z } from 'zod'; - -const McpJsonFileSchema = z.object({ - mcpServers: z.record(z.string(), McpServerConfigSchema).default({}), -}); - -export interface McpJsonPaths { - readonly user: string; - readonly projectRoot: string; - readonly project: string; -} - -export interface ResolveMcpJsonPathsInput { - readonly fs: IHostFileSystem; - readonly cwd: string; - readonly homeDir?: string; -} - -export async function resolveMcpJsonPaths(input: ResolveMcpJsonPathsInput): Promise { - const start = normalize(input.cwd); - const projectRoot = (await findGitWorkTree(input.fs, start))?.root ?? start; - - return { - user: join(resolveKimiHome(input.homeDir), 'mcp.json'), - projectRoot: join(projectRoot, '.mcp.json'), - project: join(input.cwd, '.kimi-code', 'mcp.json'), - }; -} - -export interface LoadMcpServersInput { - readonly fs: IHostFileSystem; - readonly cwd: string; - readonly homeDir?: string; - readonly includeProject?: boolean; -} - -export async function loadMcpServers( - input: LoadMcpServersInput, -): Promise> { - const paths = await resolveMcpJsonPaths(input); - if (input.includeProject === false) { - return readMcpJson(input.fs, paths.user); - } - const [user, projectRoot, project] = await Promise.all([ - readMcpJson(input.fs, paths.user), - readMcpJson(input.fs, paths.projectRoot, { stdioCwdBase: dirname(paths.projectRoot) }), - readMcpJson(input.fs, paths.project), - ]); - return { ...user, ...projectRoot, ...project }; -} - -interface ReadMcpJsonOptions { - readonly stdioCwdBase?: string; -} - -async function readMcpJson( - fs: IHostFileSystem, - filePath: string, - options: ReadMcpJsonOptions = {}, -): Promise> { - let text: string; - try { - text = await fs.readText(filePath); - } catch (error: unknown) { - if (isFileNotFound(error)) return {}; - throw new Error2(ErrorCodes.CONFIG_INVALID, `Failed to read ${filePath}: ${describeError(error)}`, { - cause: error, - }); - } - - if (text.trim().length === 0) return {}; - - let data: unknown; - try { - data = JSON.parse(text); - } catch (error: unknown) { - throw new Error2(ErrorCodes.CONFIG_INVALID, `Invalid JSON in ${filePath}: ${describeError(error)}`, { - cause: error, - }); - } - - try { - return normalizeMcpServers(McpJsonFileSchema.parse(data).mcpServers, options); - } catch (error: unknown) { - throw new Error2(ErrorCodes.CONFIG_INVALID, `Invalid MCP server config in ${filePath}: ${describeError(error)}`, { - cause: error, - }); - } -} - -function normalizeMcpServers( - servers: Record, - options: ReadMcpJsonOptions, -): Record { - const stdioCwdBase = options.stdioCwdBase; - if (stdioCwdBase === undefined) return servers; - - return Object.fromEntries( - Object.entries(servers).map(([name, config]) => [name, normalizeStdioCwd(config, stdioCwdBase)]), - ); -} - -function normalizeStdioCwd(config: McpServerConfig, cwdBase: string): McpServerConfig { - if (config.transport !== 'stdio') return config; - const cwd = config.cwd === undefined ? cwdBase : resolvePath(cwdBase, config.cwd); - return { ...config, cwd }; -} - -function resolvePath(base: string, value: string): string { - return isAbsolute(value) ? normalize(value) : resolve(base, value); -} - -function isFileNotFound(error: unknown): boolean { - return error instanceof HostFsError && error.code === OsFsErrors.codes.OS_FS_NOT_FOUND; -} - -function describeError(error: unknown): string { - return error instanceof Error ? error.message : String(error); -} diff --git a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfig.ts b/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfig.ts index 6c713f955b..783a4c7f3d 100644 --- a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfig.ts +++ b/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfig.ts @@ -1,6 +1,5 @@ import { createDecorator, type ServiceIdentifier } from '#/_base/di/instantiation'; -import type { Event } from '#/_base/event'; - +import type { Event, IWaitUntil } from '#/_base/event'; import type { McpServerConfig } from '#/mcpCore/config-schema'; export interface McpServersChange { @@ -8,6 +7,8 @@ export interface McpServersChange { readonly remove: readonly string[]; } +export type McpServersChangeEvent = McpServersChange & IWaitUntil; + export interface McpTunables { readonly startupTimeoutMs?: number; readonly toolTimeoutMs?: number; @@ -22,7 +23,7 @@ export interface IWorkspaceMcpConfigService { tunables(): McpTunables; - readonly onDidChange: Event; + readonly onDidChange: Event; } export const IWorkspaceMcpConfigService: ServiceIdentifier = diff --git a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfigService.ts b/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfigService.ts index 2d86ab8982..706f80e05f 100644 --- a/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfigService.ts +++ b/packages/agent-core-v2/src/workspace/workspaceMcpConfig/workspaceMcpConfigService.ts @@ -1,24 +1,25 @@ +import { dirname } from 'pathe'; + import { Disposable } from '#/_base/di/lifecycle'; -import { Emitter } from '#/_base/event'; +import { AsyncEmitter } from '#/_base/event'; import { ILogService } from '#/_base/log/log'; -import { TimeoutTimer } from '#/_base/utils/timer'; import { subtreeWatchFilter } from '#/_base/utils/paths'; -import { dirname } from 'pathe'; - -import type { McpServerConfig } from '#/mcpCore/config-schema'; -import { MCP_SECTION, type McpSection } from '#/app/mcpConfig/configSection'; +import { TimeoutTimer } from '#/_base/utils/timer'; import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { IConfigService } from '#/app/config/config'; +import { loadMcpServers, resolveMcpJsonPaths } from '#/app/mcpConfig/configLoader'; +import { MCP_SECTION, type McpSection } from '#/app/mcpConfig/configSection'; +import { IMcpConfigStore } from '#/app/mcpConfig/configStore'; import { IPluginService } from '#/app/plugin/plugin'; +import type { McpServerConfig } from '#/mcpCore/config-schema'; import { IHostFileSystem } from '#/os/interface/hostFileSystem'; import { IHostFsWatchService } from '#/os/interface/hostFsWatch'; import { IWorkspaceContext } from '#/workspace/workspaceContext/workspaceContext'; import { IWorkspaceTrust } from '#/workspace/workspaceTrust/workspaceTrust'; -import { loadMcpServers, resolveMcpJsonPaths } from './internal/config-loader'; import { IWorkspaceMcpConfigService, - type McpServersChange, + type McpServersChangeEvent, type McpTunables, } from './workspaceMcpConfig'; @@ -33,7 +34,7 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM private pluginServers = new Map(); private current: Readonly> = {}; private readonly watchDebounce = this._register(new TimeoutTimer()); - private readonly changeEmitter = this._register(new Emitter()); + private readonly changeEmitter = this._register(new AsyncEmitter()); readonly onDidChange = this.changeEmitter.event; constructor( @@ -45,16 +46,19 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM @IHostFsWatchService private readonly fsWatch: IHostFsWatchService, @IHostFileSystem private readonly fs: IHostFileSystem, @IWorkspaceTrust private readonly trust: IWorkspaceTrust, + @IMcpConfigStore mcpConfigStore: IMcpConfigStore, ) { super(); this.ready = this.initialize().catch((error: unknown) => { this.log.error('mcp config initial load failed', { error }); }); this._register( - this.plugins.onDidReload(() => { - void this.reloadPluginServers().catch((error) => { - this.log.warn(`mcp plugin reload failed: ${String(error)}`); - }); + this.plugins.onDidReload((event) => { + event.waitUntil( + this.reloadPluginServers().catch((error) => { + this.log.warn(`mcp plugin reload failed: ${String(error)}`); + }), + ); }), ); this._register( @@ -64,6 +68,15 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM }); }), ); + this._register( + mcpConfigStore.onDidWrite((event) => { + event.waitUntil( + this.reloadFileServers().catch((error) => { + this.log.warn(`mcp config reload after management write failed: ${String(error)}`); + }), + ); + }), + ); void this.watchConfigFiles(); } @@ -155,7 +168,7 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM includeProject: this.trust.isTrusted(), }); this.fileServers = new Map(Object.entries(fresh)); - this.publishIfChanged(); + await this.publishIfChanged(); }); } @@ -164,13 +177,13 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM await this.mutate(async () => { const fresh = await this.plugins.enabledMcpServers(); this.pluginServers = new Map(Object.entries(fresh)); - this.publishIfChanged(); + await this.publishIfChanged(); }); } - private publishIfChanged(): void { + private async publishIfChanged(): Promise { const next = this.merged(); - const upsert: Record = {}; + const upsert: Record = Object.create(null); const remove: string[] = []; for (const [name, config] of Object.entries(next)) { const previous = this.current[name]; @@ -183,10 +196,12 @@ export class WorkspaceMcpConfigService extends Disposable implements IWorkspaceM } this.current = next; if (Object.keys(upsert).length === 0 && remove.length === 0) return; - this.changeEmitter.fire({ upsert, remove }); + await this.changeEmitter.fireAsync({ upsert, remove }, NO_ABORT); } } +const NO_ABORT = new AbortController().signal; + function fingerprintConfig(config: McpServerConfig): string { return JSON.stringify(sortKeysDeep(config)); } @@ -202,4 +217,3 @@ function sortKeysDeep(value: unknown): unknown { } return value; } - diff --git a/packages/agent-core-v2/src/workspace/workspaceTrust/trustRecord.ts b/packages/agent-core-v2/src/workspace/workspaceTrust/trustRecord.ts new file mode 100644 index 0000000000..be6622f95b --- /dev/null +++ b/packages/agent-core-v2/src/workspace/workspaceTrust/trustRecord.ts @@ -0,0 +1,56 @@ +import { encodeWorkDirKey } from '#/_base/utils/workdir-slug'; +import { canonicalWorkspaceRoot } from '#/_base/utils/paths'; +import type { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; + +const TRUST_SCOPE = 'workspace-trust'; + +interface TrustRecord { + readonly root: string; + readonly trustedAt: number; +} + +export async function readWorkspaceTrust( + docs: IAtomicDocumentStore, + root: string, +): Promise { + try { + const canonicalKey = trustKey(root); + if ((await docs.get(TRUST_SCOPE, canonicalKey)) !== undefined) return true; + + const legacyKey = encodeWorkDirKey(root); + if (legacyKey === canonicalKey) return false; + const legacy = await docs.get(TRUST_SCOPE, legacyKey); + if (legacy === undefined) return false; + try { + await docs.set(TRUST_SCOPE, canonicalKey, legacy); + await docs.delete(TRUST_SCOPE, legacyKey); + } catch {} + return true; + } catch { + return false; + } +} + +export function writeWorkspaceTrust( + docs: IAtomicDocumentStore, + root: string, + trustedAt: number, +): Promise { + return docs.set(TRUST_SCOPE, trustKey(root), { root, trustedAt }); +} + +export function deleteWorkspaceTrust( + docs: IAtomicDocumentStore, + root: string, +): Promise { + const canonicalKey = trustKey(root); + const legacyKey = encodeWorkDirKey(root); + return (async () => { + await docs.delete(TRUST_SCOPE, canonicalKey); + if (legacyKey !== canonicalKey) await docs.delete(TRUST_SCOPE, legacyKey); + })(); +} + +function trustKey(root: string): string { + return encodeWorkDirKey(canonicalWorkspaceRoot(root)); +} diff --git a/packages/agent-core-v2/src/workspace/workspaceTrust/workspaceTrustService.ts b/packages/agent-core-v2/src/workspace/workspaceTrust/workspaceTrustService.ts index 6ad2bd936b..ed489d0c77 100644 --- a/packages/agent-core-v2/src/workspace/workspaceTrust/workspaceTrustService.ts +++ b/packages/agent-core-v2/src/workspace/workspaceTrust/workspaceTrustService.ts @@ -1,19 +1,12 @@ import { Disposable } from '#/_base/di/lifecycle'; import { Emitter } from '#/_base/event'; import { defineState } from '#/state/state'; -import { encodeWorkDirKey } from '#/_base/utils/workdir-slug'; import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; import { IWorkspaceStateService } from '#/workspace/state/workspaceState'; import { IWorkspaceContext } from '#/workspace/workspaceContext/workspaceContext'; import { IWorkspaceTrust, type WorkspaceTrustChange } from './workspaceTrust'; - -const TRUST_SCOPE = 'workspace-trust'; - -interface TrustRecord { - readonly root: string; - readonly trustedAt: number; -} +import { deleteWorkspaceTrust, readWorkspaceTrust, writeWorkspaceTrust } from './trustRecord'; export const workspaceTrustTrustedKey = defineState( 'workspaceTrust.trusted', @@ -25,7 +18,6 @@ export class WorkspaceTrustService extends Disposable implements IWorkspaceTrust readonly ready: Promise; private readonly root: string; - private readonly storeKey: string; private readonly changeEmitter = this._register(new Emitter()); readonly onDidChange = this.changeEmitter.event; @@ -37,7 +29,6 @@ export class WorkspaceTrustService extends Disposable implements IWorkspaceTrust super(); this.states.contributeState(workspaceTrustTrustedKey); this.root = workspace.cwd; - this.storeKey = encodeWorkDirKey(workspace.cwd); this.ready = this.initialize(); } @@ -60,27 +51,19 @@ export class WorkspaceTrustService extends Disposable implements IWorkspaceTrust async trust(): Promise { if (this.trusted) return; - await this.docs.set(TRUST_SCOPE, this.storeKey, { - root: this.root, - trustedAt: Date.now(), - }); + await writeWorkspaceTrust(this.docs, this.root, Date.now()); this.trusted = true; this.changeEmitter.fire({ trusted: true }); } async untrust(): Promise { if (!this.trusted) return; - await this.docs.delete(TRUST_SCOPE, this.storeKey); + await deleteWorkspaceTrust(this.docs, this.root); this.trusted = false; this.changeEmitter.fire({ trusted: false }); } private async initialize(): Promise { - try { - this.trusted = (await this.docs.get(TRUST_SCOPE, this.storeKey)) !== undefined; - } catch { - this.trusted = false; - } + this.trusted = await readWorkspaceTrust(this.docs, this.root); } } - diff --git a/packages/agent-core-v2/test/_base/utils/paths.test.ts b/packages/agent-core-v2/test/_base/utils/paths.test.ts index 7f4ce3547e..1859ff0a38 100644 --- a/packages/agent-core-v2/test/_base/utils/paths.test.ts +++ b/packages/agent-core-v2/test/_base/utils/paths.test.ts @@ -4,7 +4,7 @@ import nodePath, { win32 } from 'node:path'; import { afterEach, beforeEach, describe, expect, it } from 'vitest'; -import { findUpwardRoot, subtreeWatchFilter } from '#/_base/utils/paths'; +import { canonicalWorkspaceRoot, findUpwardRoot, resolvePath, subtreeWatchFilter } from '#/_base/utils/paths'; describe('subtree watch filtering', () => { const root = '/repo'; @@ -168,3 +168,49 @@ describe('findUpwardRoot', () => { expect(found).toBe('E:/repo'); }); }); + +describe('resolvePath', () => { + it('resolves drive-letter absolute values without joining the base', () => { + expect(resolvePath('/repo', 'C:/tools')).toBe('C:/tools'); + expect(resolvePath('/repo', 'C:\\tools\\bin')).toBe('C:/tools/bin'); + }); + + it('resolves values against a Windows base with win32 semantics', () => { + expect(resolvePath('C:/repo', 'tools/mcp')).toBe('C:/repo/tools/mcp'); + expect(resolvePath('C:\\repo', '.\\tools')).toBe('C:/repo/tools'); + expect(resolvePath('C:/repo', 'D:/elsewhere')).toBe('D:/elsewhere'); + }); + + it('keeps UNC bases and values intact', () => { + expect(resolvePath('//server/share/repo', 'tools')).toBe('//server/share/repo/tools'); + expect(resolvePath('/repo', '//server/share/tools')).toBe('//server/share/tools'); + expect(resolvePath('\\\\server\\share\\repo', 'tools')).toBe('//server/share/repo/tools'); + }); + + it('keeps POSIX resolution identical to plain absolute/normalize semantics', () => { + expect(resolvePath('/repo', 'tools/../mcp')).toBe('/repo/mcp'); + expect(resolvePath('/repo', '/abs/path')).toBe('/abs/path'); + }); +}); + +describe('canonicalWorkspaceRoot', () => { + it('case-folds drive-letter spellings and strips trailing separators', () => { + expect(canonicalWorkspaceRoot('C:\\Users\\Foo\\Repo')).toBe('c:/users/foo/repo'); + expect(canonicalWorkspaceRoot('C:/Users/Foo/Repo/')).toBe('c:/users/foo/repo'); + }); + + it('keeps the UNC share slash and case-folds', () => { + expect(canonicalWorkspaceRoot('//server/share/repo')).toBe('//server/share/repo'); + expect(canonicalWorkspaceRoot('\\\\SERVER\\SHARE\\REPO')).toBe('//server/share/repo'); + }); + + it('resolves dot segments in Windows spellings', () => { + expect(canonicalWorkspaceRoot('C:/Users/Foo/../Foo/Repo')).toBe('c:/users/foo/repo'); + }); + + it('keeps POSIX roots untouched apart from trailing-slash and dot-segment cleanup', () => { + expect(canonicalWorkspaceRoot('/Repo/Sub')).toBe('/Repo/Sub'); + expect(canonicalWorkspaceRoot('/Repo/Sub/')).toBe('/Repo/Sub'); + expect(canonicalWorkspaceRoot('/Repo/Sub/../Other')).toBe('/Repo/Other'); + }); +}); diff --git a/packages/agent-core-v2/test/agent/plugin/agentPlugin.test.ts b/packages/agent-core-v2/test/agent/plugin/agentPlugin.test.ts index 7fb6552417..ea78621292 100644 --- a/packages/agent-core-v2/test/agent/plugin/agentPlugin.test.ts +++ b/packages/agent-core-v2/test/agent/plugin/agentPlugin.test.ts @@ -1,7 +1,7 @@ import { afterEach, describe, expect, it } from 'vitest'; import { SyncDescriptor } from '#/_base/di/descriptors'; -import { Emitter } from '#/_base/event'; +import { AsyncEmitter, Emitter } from '#/_base/event'; import { IAgentPluginService } from '#/agent/plugin/agentPlugin'; import { AgentPluginService } from '#/agent/plugin/agentPluginService'; import { USER_PROMPT_ORIGIN } from '#/agent/contextMemory/types'; @@ -12,7 +12,7 @@ import { IPluginService } from '#/app/plugin/plugin'; import type { EnabledPluginSessionStart, PluginMutationSummary, - ReloadSummary, + PluginReloadEvent, } from '#/app/plugin/types'; import { InMemorySkillCatalog } from '#/app/skillCatalog/registry'; import { summarizeSkill } from '#/app/skillCatalog/types'; @@ -363,7 +363,7 @@ describe('AgentPluginService plugin-change reminder', () => { }); it('does not append the plugin_change reminder on an explicit reload', async () => { - const reloadEmitter = new Emitter(); + const reloadEmitter = new AsyncEmitter(); ctx = createTestAgent( { autoConfigure: true }, appService(IPluginService, stubPluginService({ sessionStarts: [], reloadEmitter })), @@ -372,7 +372,10 @@ describe('AgentPluginService plugin-change reminder', () => { ); ctx.get(IAgentPluginService); - reloadEmitter.fire({ added: [], removed: [], errors: [] }); + await reloadEmitter.fireAsyncConcurrent( + { added: [], removed: [], errors: [] }, + new AbortController().signal, + ); expect(findPluginChangeMessages(ctx)).toHaveLength(0); reloadEmitter.dispose(); diff --git a/packages/agent-core-v2/test/agent/pluginCommand/pluginCommand.test.ts b/packages/agent-core-v2/test/agent/pluginCommand/pluginCommand.test.ts index 89f53563b0..284482e949 100644 --- a/packages/agent-core-v2/test/agent/pluginCommand/pluginCommand.test.ts +++ b/packages/agent-core-v2/test/agent/pluginCommand/pluginCommand.test.ts @@ -42,6 +42,7 @@ function pluginServiceStub(commands: readonly PluginCommandDef[]): IPluginServic enabledSessionStarts: async () => [], enabledSystemPrompts: async () => [], enabledMcpServers: async () => ({}), + mcpServerEntries: async () => [], enabledHooks: async () => [], hasLoadedSnapshot: () => true, }; diff --git a/packages/agent-core-v2/test/workspace/workspaceMcpConfig/config-loader.test.ts b/packages/agent-core-v2/test/app/mcpConfig/configLoader.test.ts similarity index 90% rename from packages/agent-core-v2/test/workspace/workspaceMcpConfig/config-loader.test.ts rename to packages/agent-core-v2/test/app/mcpConfig/configLoader.test.ts index e813d97ac7..70d6070594 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcpConfig/config-loader.test.ts +++ b/packages/agent-core-v2/test/app/mcpConfig/configLoader.test.ts @@ -5,7 +5,7 @@ import { join } from 'pathe'; import { afterEach, describe, expect, it } from 'vitest'; import { ErrorCodes, Error2 } from '#/errors'; -import { loadMcpServers, resolveMcpJsonPaths } from '#/workspace/workspaceMcpConfig/internal/config-loader'; +import { loadMcpServers, resolveMcpJsonPaths } from '#/app/mcpConfig/configLoader'; import { HostFileSystem } from '#/os/backends/node-local/hostFsService'; const fs = new HostFileSystem(); @@ -60,6 +60,16 @@ describe('loadMcpServers', () => { expect(servers).toEqual({}); }); + it('rejects a null mcpServers field', async () => { + const home = makeTempDir(); + const cwd = makeTempDir(); + await writeJson(join(home, 'mcp.json'), { mcpServers: null }); + + await expect(loadMcpServers({ fs, cwd, homeDir: home })).rejects.toMatchObject({ + code: ErrorCodes.CONFIG_INVALID, + }); + }); + it('merges project-local mcp.json with user-global, project overriding on conflict', async () => { const home = makeTempDir(); const cwd = makeTempDir(); @@ -217,6 +227,28 @@ describe('loadMcpServers', () => { }); }); + it('keeps Windows drive-letter and UNC stdio cwd values resolved on any host', async () => { + const home = makeTempDir(); + const repoRoot = makeTempDir(); + const cwd = join(repoRoot, 'packages', 'agent-core'); + await mkdir(join(repoRoot, '.git'), { recursive: true }); + await mkdir(cwd, { recursive: true }); + + await writeJson(join(repoRoot, '.mcp.json'), { + mcpServers: { + drive: { command: 'node', cwd: 'C:/tools' }, + driveBackslash: { command: 'node', cwd: 'C:\\tools\\bin' }, + unc: { command: 'node', cwd: '//server/share/tools' }, + }, + }); + + const servers = await loadMcpServers({ fs, cwd, homeDir: home }); + + expect(servers['drive']).toMatchObject({ cwd: 'C:/tools' }); + expect(servers['driveBackslash']).toMatchObject({ cwd: 'C:/tools/bin' }); + expect(servers['unc']).toMatchObject({ cwd: '//server/share/tools' }); + }); + it('throws Error2(config.invalid) on invalid JSON', async () => { const home = makeTempDir(); const cwd = makeTempDir(); diff --git a/packages/agent-core-v2/test/app/mcpConfig/configStore.test.ts b/packages/agent-core-v2/test/app/mcpConfig/configStore.test.ts new file mode 100644 index 0000000000..d7c8f2a16e --- /dev/null +++ b/packages/agent-core-v2/test/app/mcpConfig/configStore.test.ts @@ -0,0 +1,383 @@ +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; + +import { DisposableStore } from '#/_base/di/lifecycle'; +import { createServices } from '#/_base/di/test'; +import { IBootstrapService } from '#/app/bootstrap/bootstrap'; +import { + IMcpConfigStore, + McpConfigStore, + type GlobalMcpServerConfig, +} from '#/app/mcpConfig/configStore'; +import { ErrorCodes } from '#/errors'; +import { InMemoryStorageService } from '#/persistence/backends/memory/inMemoryStorageService'; +import { IFileSystemStorageService } from '#/persistence/interface/storage'; + +const CONFIG_SCOPE = ''; +const CONFIG_KEY = 'mcp.json'; + +const textEncoder = new TextEncoder(); +const textDecoder = new TextDecoder(); + +function stdioServer(name: string, command = 'npx'): GlobalMcpServerConfig { + return { name, transport: 'stdio', command }; +} + +describe('McpConfigStore', () => { + let disposables: DisposableStore; + let storage: InMemoryStorageService; + let store: IMcpConfigStore; + + beforeEach(() => { + disposables = new DisposableStore(); + storage = new InMemoryStorageService(); + const ix = createServices(disposables, { + additionalServices: (reg) => { + reg.defineInstance(IFileSystemStorageService, storage); + reg.definePartialInstance(IBootstrapService, { homeDir: '/kimi-test-home' }); + reg.define(IMcpConfigStore, McpConfigStore); + }, + }); + store = ix.get(IMcpConfigStore); + }); + + afterEach(() => { + disposables.dispose(); + }); + + async function seedRaw(text: string): Promise { + await storage.write(CONFIG_SCOPE, CONFIG_KEY, textEncoder.encode(text)); + } + + async function seedJson(value: unknown): Promise { + await seedRaw(JSON.stringify(value)); + } + + async function readRaw(): Promise { + const bytes = await storage.read(CONFIG_SCOPE, CONFIG_KEY); + return bytes === undefined ? undefined : textDecoder.decode(bytes); + } + + describe('CRUD', () => { + it('round-trips add → get → update → remove against an empty catalog', async () => { + await expect(store.list()).resolves.toEqual([]); + + const added = await store.add(stdioServer('alpha')); + expect(added).toEqual([{ name: 'alpha', transport: 'stdio', command: 'npx' }]); + await expect(store.get('alpha')).resolves.toEqual({ + name: 'alpha', + transport: 'stdio', + command: 'npx', + }); + + const updated = await store.update(stdioServer('alpha', 'node')); + expect(updated).toEqual([{ name: 'alpha', transport: 'stdio', command: 'node' }]); + + const remaining = await store.remove('alpha'); + expect(remaining).toEqual([]); + await expect(store.list()).resolves.toEqual([]); + }); + + it('returns the full catalog from add, update, and remove', async () => { + await store.add(stdioServer('alpha')); + const added = await store.add(stdioServer('beta')); + expect(added.map((server) => server.name)).toEqual(['alpha', 'beta']); + const remaining = await store.remove('alpha'); + expect(remaining.map((server) => server.name)).toEqual(['beta']); + }); + + it('treats a missing file as an empty catalog', async () => { + await expect(store.list()).resolves.toEqual([]); + }); + + it('treats a whitespace-only file as an empty catalog', async () => { + await seedRaw(' \n'); + await expect(store.list()).resolves.toEqual([]); + }); + }); + + describe('byte format', () => { + it('writes two-space-indented JSON with a trailing newline', async () => { + await store.add(stdioServer('alpha')); + + const expected = `${JSON.stringify( + { mcpServers: { alpha: { transport: 'stdio', command: 'npx' } } }, + null, + 2, + )}\n`; + expect(await readRaw()).toBe(expected); + }); + + it('preserves unknown top-level keys and mcpServers ordering on write', async () => { + await seedRaw('{\n "mcpServers": {},\n "futureSetting": { "a": 1 }\n}\n'); + + await store.add(stdioServer('alpha')); + + const expected = `${JSON.stringify( + { + mcpServers: { alpha: { transport: 'stdio', command: 'npx' } }, + futureSetting: { a: 1 }, + }, + null, + 2, + )}\n`; + expect(await readRaw()).toBe(expected); + }); + + it('round-trips an entry byte-identically through update', async () => { + await store.add(stdioServer('alpha')); + const before = await readRaw(); + + await store.update(stdioServer('alpha')); + + expect(await readRaw()).toBe(before); + }); + }); + + describe('name normalization', () => { + it('trims surrounding whitespace from server names', async () => { + const added = await store.add(stdioServer(' alpha ')); + expect(added.map((server) => server.name)).toEqual(['alpha']); + await expect(store.get(' alpha ')).resolves.toMatchObject({ name: 'alpha' }); + expect(JSON.parse((await readRaw())!)).toMatchObject({ mcpServers: { alpha: {} } }); + }); + + it('rejects empty names across all operations', async () => { + await expect(store.add(stdioServer(' '))).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP server name cannot be empty', + }); + await expect(store.update(stdioServer(''))).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + }); + await expect(store.get(' ')).rejects.toMatchObject({ code: ErrorCodes.REQUEST_INVALID }); + await expect(store.remove(' ')).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + }); + }); + }); + + describe('guards', () => { + it('rejects add with an existing name', async () => { + await store.add(stdioServer('alpha')); + await expect(store.add(stdioServer('alpha'))).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP server "alpha" already exists', + }); + }); + + it('rejects update for an unknown server', async () => { + await expect(store.update(stdioServer('ghost'))).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + message: 'MCP server "ghost" was not found', + }); + }); + + it('rejects get for an unknown server', async () => { + await expect(store.get('ghost')).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + message: 'MCP server "ghost" was not found', + }); + }); + + it('treats remove of an unknown server as a no-op returning the current catalog', async () => { + await store.add(stdioServer('alpha')); + const before = await readRaw(); + + let fired = 0; + store.onDidWrite(() => fired++); + const remaining = await store.remove('ghost'); + + expect(remaining.map((server) => server.name)).toEqual(['alpha']); + expect(await readRaw()).toBe(before); + expect(fired).toBe(0); + }); + }); + + describe('read validation', () => { + it('rejects invalid JSON with config.invalid', async () => { + await seedRaw('{not json}'); + await expect(store.list()).rejects.toMatchObject({ code: ErrorCodes.CONFIG_INVALID }); + await expect(store.list()).rejects.toThrow(/^Invalid JSON in /); + }); + + it('rejects a BOM-prefixed file as malformed JSON', async () => { + await seedRaw('\uFEFF{"mcpServers":{}}'); + await expect(store.list()).rejects.toMatchObject({ code: ErrorCodes.CONFIG_INVALID }); + await expect(store.list()).rejects.toThrow(/^Invalid JSON in /); + }); + + it('rejects a non-object top level', async () => { + await seedRaw('["alpha"]'); + await expect(store.list()).rejects.toMatchObject({ + code: ErrorCodes.CONFIG_INVALID, + message: `Invalid MCP config in ${store.path}: expected a JSON object`, + }); + }); + + it('rejects a non-object "mcpServers" value', async () => { + await seedJson({ mcpServers: 'nope' }); + await expect(store.list()).rejects.toMatchObject({ + code: ErrorCodes.CONFIG_INVALID, + message: `Invalid MCP config in ${store.path}: "mcpServers" must be an object`, + }); + }); + + it('rejects an invalid server entry with the v1 message shape', async () => { + await seedJson({ mcpServers: { bad: { transport: 'websocket' } } }); + await expect(store.list()).rejects.toMatchObject({ code: ErrorCodes.CONFIG_INVALID }); + await expect(store.list()).rejects.toThrow(/^Invalid MCP server "bad" in global config: /); + }); + + it('rejects an invalid add payload before touching the file', async () => { + const invalid = { name: 'bad', transport: 'stdio' } as unknown as GlobalMcpServerConfig; + await expect(store.add(invalid)).rejects.toMatchObject({ + code: ErrorCodes.CONFIG_INVALID, + }); + await expect(store.add(invalid)).rejects.toThrow( + /^Invalid MCP server "bad" in global config: /, + ); + expect(await readRaw()).toBeUndefined(); + }); + }); + + describe('__proto__ safety', () => { + it('adds, reads back, and removes a server literally named __proto__', async () => { + const added = await store.add(stdioServer('__proto__')); + expect(added).toEqual([{ name: '__proto__', transport: 'stdio', command: 'npx' }]); + + await expect(store.get('__proto__')).resolves.toMatchObject({ name: '__proto__' }); + + const persisted = JSON.parse((await readRaw())!) as Record; + const rawServers = persisted['mcpServers'] as Record; + expect(Object.hasOwn(rawServers, '__proto__')).toBe(true); + expect(rawServers['__proto__']).toEqual({ transport: 'stdio', command: 'npx' }); + + await expect(store.remove('__proto__')).resolves.toEqual([]); + }); + + it('reads a file declaring a __proto__ server', async () => { + await seedRaw('{"mcpServers":{"__proto__":{"transport":"stdio","command":"npx"}}}'); + await expect(store.list()).resolves.toEqual([ + { name: '__proto__', transport: 'stdio', command: 'npx' }, + ]); + }); + }); + + describe('onDidWrite', () => { + it('fires once after each successful add, update, and remove', async () => { + let fired = 0; + store.onDidWrite(() => fired++); + + await store.add(stdioServer('alpha')); + expect(fired).toBe(1); + + await store.update(stdioServer('alpha', 'node')); + expect(fired).toBe(2); + + await store.remove('alpha'); + expect(fired).toBe(3); + }); + + it('never fires on reads, failed mutations, or no-op removes', async () => { + await store.add(stdioServer('alpha')); + let fired = 0; + store.onDidWrite(() => fired++); + + await store.list(); + await store.get('alpha'); + expect(fired).toBe(0); + + await expect(store.add(stdioServer(' '))).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + }); + await expect(store.get('ghost')).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + }); + await store.remove('ghost'); + expect(fired).toBe(0); + }); + + it('waits for asynchronous listeners before resolving a mutation', async () => { + let resolveStarted!: () => void; + const started = new Promise((resolve) => { + resolveStarted = resolve; + }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + store.onDidWrite((event) => { + resolveStarted(); + event.waitUntil(gate); + }); + + let completed = false; + const mutation = store.add(stdioServer('alpha')).then(() => { + completed = true; + }); + await started; + await Promise.resolve(); + expect(completed).toBe(false); + + release(); + await mutation; + expect(completed).toBe(true); + }); + + it('starts asynchronous listeners concurrently before waiting for completion', async () => { + let resolveStarted!: () => void; + const started = new Promise((resolve) => { + resolveStarted = resolve; + }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + let secondStarted = false; + store.onDidWrite((event) => { + resolveStarted(); + event.waitUntil(gate); + }); + store.onDidWrite(() => { + secondStarted = true; + }); + + const mutation = store.add(stdioServer('alpha')); + await started; + await Promise.resolve(); + const secondStartedBeforeRelease = secondStarted; + release(); + await mutation; + + expect(secondStartedBeforeRelease).toBe(true); + }); + + it('serializes concurrent mutations so no entry is lost', async () => { + await Promise.all([store.add(stdioServer('alpha')), store.add(stdioServer('beta'))]); + + await expect(store.list()).resolves.toEqual([ + { name: 'alpha', transport: 'stdio', command: 'npx' }, + { name: 'beta', transport: 'stdio', command: 'npx' }, + ]); + expect(JSON.parse((await readRaw())!)).toMatchObject({ + mcpServers: { alpha: {}, beta: {} }, + }); + }); + + it('lets a write listener mutate the store without wedging the mutation queue', async () => { + let reentered = false; + store.onDidWrite((event) => { + if (reentered) return; + reentered = true; + event.waitUntil(store.add(stdioServer('beta')).then(() => undefined)); + }); + + await store.add(stdioServer('alpha')); + + await expect(store.list()).resolves.toEqual([ + { name: 'alpha', transport: 'stdio', command: 'npx' }, + { name: 'beta', transport: 'stdio', command: 'npx' }, + ]); + }); + }); +}); diff --git a/packages/agent-core-v2/test/app/mcpConfig/oauthService.test.ts b/packages/agent-core-v2/test/app/mcpConfig/oauthService.test.ts new file mode 100644 index 0000000000..67d6dce60f --- /dev/null +++ b/packages/agent-core-v2/test/app/mcpConfig/oauthService.test.ts @@ -0,0 +1,134 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { DisposableStore } from '#/_base/di/lifecycle'; +import { createServices } from '#/_base/di/test'; +import { ILogService } from '#/_base/log/log'; +import { IAgentIdentity } from '#/app/agentIdentity/agentIdentity'; +import { + AppMcpOAuthService, + IMcpOAuthService, +} from '#/app/mcpConfig/oauthService'; +import { IMcpOAuthStore } from '#/app/mcpConfig/oauthStore'; + +import { stubLog } from '../../_base/log/stubs'; +import { deferredAgentIdentityStub } from '../agentIdentity/stubs'; +import { createMemoryMcpOAuthStore } from '../../mcpCore/stubs'; + +describe('App MCP OAuth bootstrap', () => { + let disposables: DisposableStore; + + beforeEach(() => { + disposables = new DisposableStore(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + disposables.dispose(); + }); + + it('starts the proactive refresh sweep only after identity resolution', async () => { + const memory = createMemoryMcpOAuthStore(); + let signalList: () => void = () => undefined; + const listed = new Promise((resolve) => { + signalList = resolve; + }); + const list = vi.fn(async (prefix?: string) => { + signalList(); + return memory.list(prefix); + }); + const identity = deferredAgentIdentityStub({ slug: 'test-agent' }); + const ix = createServices(disposables, { + strict: true, + additionalServices: (reg) => { + reg.defineInstance(IMcpOAuthStore, { + _serviceBrand: undefined, + ...memory, + list, + }); + reg.defineInstance(ILogService, stubLog()); + reg.defineInstance(IAgentIdentity, identity.identity); + reg.define(IMcpOAuthService, AppMcpOAuthService); + }, + }); + ix.get(IMcpOAuthService); + + await Promise.resolve(); + expect(list).not.toHaveBeenCalled(); + + identity.freeze(); + await listed; + expect(list).toHaveBeenCalledTimes(1); + }); + + it('does not start the proactive refresh sweep after shutdown before identity resolution', async () => { + const memory = createMemoryMcpOAuthStore(); + const list = vi.fn(memory.list); + const identity = deferredAgentIdentityStub({ slug: 'test-agent' }); + const ix = createServices(disposables, { + strict: true, + additionalServices: (reg) => { + reg.defineInstance(IMcpOAuthStore, { + _serviceBrand: undefined, + ...memory, + list, + }); + reg.defineInstance(ILogService, stubLog()); + reg.defineInstance(IAgentIdentity, identity.identity); + reg.define(IMcpOAuthService, AppMcpOAuthService); + }, + }); + const service = ix.get(IMcpOAuthService); + + await service.shutdown(); + identity.freeze(); + await Promise.resolve(); + await Promise.resolve(); + + expect(list).not.toHaveBeenCalled(); + }); + + it('stops a proactive refresh sweep that is still listing credentials during shutdown', async () => { + const memory = createMemoryMcpOAuthStore(); + let releaseList: () => void = () => undefined; + const listed = new Promise((resolve) => { + releaseList = resolve; + }); + let signalList: () => void = () => undefined; + const listStarted = new Promise((resolve) => { + signalList = resolve; + }); + const list = vi.fn(async () => { + signalList(); + await listed; + return ['credential-meta.json']; + }); + const read = vi.fn(); + const identity = deferredAgentIdentityStub({ slug: 'test-agent' }); + const ix = createServices(disposables, { + strict: true, + additionalServices: (reg) => { + reg.defineInstance(IMcpOAuthStore, { + _serviceBrand: undefined, + ...memory, + list, + read: async (key: string) => { + read(key); + return memory.read(key); + }, + }); + reg.defineInstance(ILogService, stubLog()); + reg.defineInstance(IAgentIdentity, identity.identity); + reg.define(IMcpOAuthService, AppMcpOAuthService); + }, + }); + const service = ix.get(IMcpOAuthService); + identity.freeze(); + await listStarted; + + const shutdown = service.shutdown(); + releaseList(); + await shutdown; + + expect(read).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts b/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts new file mode 100644 index 0000000000..57776721e7 --- /dev/null +++ b/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts @@ -0,0 +1,1691 @@ +import { mkdtempSync } from 'node:fs'; +import { mkdir, rm, writeFile } from 'node:fs/promises'; +import { createServer as createHttpServer, type Server as HttpServer } from 'node:http'; +import type { AddressInfo as HttpAddress } from 'node:net'; +import { tmpdir } from 'node:os'; +import { join } from 'pathe'; + +import type { OAuthClientInformationFull } from '@modelcontextprotocol/sdk/shared/auth.js'; +import { afterEach, beforeEach, describe, expect, it, vi, type Mock } from 'vitest'; + +import { DisposableStore } from '#/_base/di/lifecycle'; +import { createServices } from '#/_base/di/test'; +import { ILogService } from '#/_base/log/log'; +import { IBootstrapService } from '#/app/bootstrap/bootstrap'; +import { + IAgentIdentity, + type AgentIdentitySnapshot, +} from '#/app/agentIdentity/agentIdentity'; +import { IConfigService } from '#/app/config/config'; +import { IMcpConfigStore, McpConfigStore } from '#/app/mcpConfig/configStore'; +import { IMcpOAuthService } from '#/app/mcpConfig/oauthService'; +import { + IMcpManagementService, + type GlobalMcpServerConfig, + type McpServerLocator, +} from '#/app/mcpManagement/mcpManagement'; +import { McpManagementService } from '#/app/mcpManagement/mcpManagementService'; +import { IMcpRegistryService } from '#/app/mcpRegistry/mcpRegistry'; +import { McpRegistryService } from '#/app/mcpRegistry/mcpRegistryService'; +import { IPluginService } from '#/app/plugin/plugin'; +import type { PluginMcpServerEntry } from '#/app/plugin/types'; +import { ErrorCodes, Error2 } from '#/errors'; +import { McpOAuthService, type McpOAuthEvent } from '#/mcpCore/oauth/service'; +import { HostFileSystem } from '#/os/backends/node-local/hostFsService'; +import { HostProcessService } from '#/os/backends/node-local/hostProcessService'; +import { IHostEnvironment } from '#/os/interface/hostEnvironment'; +import { IHostFileSystem } from '#/os/interface/hostFileSystem'; +import { IHostProcessService } from '#/os/interface/hostProcess'; +import { InMemoryStorageService } from '#/persistence/backends/memory/inMemoryStorageService'; +import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; +import { IFileSystemStorageService } from '#/persistence/interface/storage'; +import { FakeRuntime } from '#/runtime/fakeRuntime'; +import type { WorkspaceInstance } from '#/workspace/workspaceInstance/workspaceInstance'; +import { + IRuntimeResolver, + IWorkspaceInstanceManager, +} from '#/workspace/workspaceInstance/workspaceInstanceManager'; + +import { stubLog } from '../../_base/log/stubs'; +import { + createMemoryMcpOAuthStore, + startInProcessHttpMcpServer, + stdioFixture, +} from '../../mcpCore/stubs'; +import { stubAgentIdentity } from '../agentIdentity/stubs'; + +function stdioServer(name: string, command = 'npx'): GlobalMcpServerConfig { + return { name, transport: 'stdio', command }; +} + +const CONFIG_SCOPE = ''; +const CONFIG_KEY = 'mcp.json'; + +const textEncoder = new TextEncoder(); + +describe('McpManagementService', () => { + let home: string; + let disposables: DisposableStore; + let tempDirs: string[]; + let httpServers: Array<{ close: () => Promise }>; + let storage: InMemoryStorageService; + let store: IMcpConfigStore; + let pluginEntries: PluginMcpServerEntry[]; + let pluginError: Error | undefined; + let oauth: McpOAuthService; + let configReady: Promise; + let identityReady: Promise; + let identitySnapshot: AgentIdentitySnapshot; + let trusted: boolean; + let getOrCreate: Mock; + let findContaining: Mock; + let management: IMcpManagementService; + + beforeEach(() => { + home = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-home-')); + vi.stubEnv('KIMI_CODE_HOME', home); + disposables = new DisposableStore(); + tempDirs = [home]; + httpServers = []; + storage = new InMemoryStorageService(); + pluginEntries = []; + pluginError = undefined; + oauth = new McpOAuthService({ store: createMemoryMcpOAuthStore() }); + configReady = Promise.resolve(); + identitySnapshot = stubAgentIdentity({ slug: 'test-agent' }).current(); + identityReady = Promise.resolve(identitySnapshot); + trusted = true; + getOrCreate = vi.fn(async () => + ({ id: 'test-workspace' }) as unknown as WorkspaceInstance, + ); + findContaining = vi.fn(() => undefined); + const hostProcess = new HostProcessService(); + const runtime = Object.assign( + new FakeRuntime( + { workspaceId: 'test-workspace', runtimeId: 'local', generation: 'test-generation' }, + { capabilities: ['process'] }, + ), + { process: hostProcess }, + ); + const ix = createServices(disposables, { + additionalServices: (reg) => { + reg.defineInstance(IFileSystemStorageService, storage); + reg.definePartialInstance(IBootstrapService, { homeDir: home }); + reg.define(IMcpConfigStore, McpConfigStore); + reg.definePartialInstance(IPluginService, { + mcpServerEntries: async () => { + if (pluginError !== undefined) throw pluginError; + return pluginEntries; + }, + }); + reg.defineInstance(IHostFileSystem, new HostFileSystem()); + reg.defineInstance(IHostEnvironment, { + _serviceBrand: undefined, + osKind: 'Linux', + osArch: 'x64', + osVersion: 'test', + shellName: 'bash', + shellPath: '/bin/bash', + pathClass: 'posix', + homeDir: home, + ready: Promise.resolve(), + }); + reg.defineInstance(IHostProcessService, hostProcess); + reg.definePartialInstance(IAtomicDocumentStore, { + get: async () => (trusted ? ({} as T) : undefined), + }); + reg.define(IMcpRegistryService, McpRegistryService); + reg.defineInstance(IMcpOAuthService, oauth); + reg.definePartialInstance(IConfigService, { + get ready() { + return configReady; + }, + get: ((_domain: string): T => undefined as T) as IConfigService['get'], + }); + reg.defineInstance(IAgentIdentity, { + _serviceBrand: undefined, + resolved: () => identityReady, + current: () => identitySnapshot, + }); + reg.defineInstance(IRuntimeResolver, { + _serviceBrand: undefined, + inspect: () => runtime, + acquire: () => ({ runtime, track: (resource) => resource, dispose: () => {} }), + }); + reg.definePartialInstance(IWorkspaceInstanceManager, { findContaining, getOrCreate }); + reg.defineInstance(ILogService, stubLog()); + reg.define(IMcpManagementService, McpManagementService); + }, + }); + store = ix.get(IMcpConfigStore); + management = ix.get(IMcpManagementService); + }); + + afterEach(async () => { + disposables.dispose(); + await oauth.dispose(); + vi.unstubAllEnvs(); + await Promise.all(httpServers.map((server) => server.close())); + await Promise.all(tempDirs.map((dir) => rm(dir, { recursive: true, force: true }))); + }); + + async function startHttpServer(): Promise<{ url: string }> { + const server = await startInProcessHttpMcpServer(); + httpServers.push(server); + return server; + } + + async function startCountingServer(): Promise<{ + url: string; + requestCount: () => number; + }> { + let requests = 0; + const httpServer: HttpServer = createHttpServer((_req, res) => { + requests += 1; + res.writeHead(404).end(); + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + httpServers.push({ + close: () => + new Promise((resolve, reject) => { + httpServer.close((err) => (err === undefined || err === null ? resolve() : reject(err))); + }), + }); + const port = (httpServer.address() as HttpAddress).port; + return { + url: `http://127.0.0.1:${port}/mcp`, + requestCount: () => requests, + }; + } + + async function startGatedServer(): Promise<{ origin: string; url: string }> { + const httpServer: HttpServer = createHttpServer((req, res) => { + if (req.method === 'POST' && req.url === '/token') { + res.writeHead(400, { 'content-type': 'application/json' }); + res.end(JSON.stringify({ error: 'invalid_grant' })); + return; + } + res.writeHead(401, { + 'content-type': 'application/json', + 'www-authenticate': 'Bearer realm="mcp"', + }); + res.end(JSON.stringify({ error: 'unauthorized' })); + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + httpServers.push({ + close: () => + new Promise((resolve, reject) => { + httpServer.close((err) => (err === undefined || err === null ? resolve() : reject(err))); + }), + }); + const port = (httpServer.address() as HttpAddress).port; + return { origin: `http://127.0.0.1:${port}`, url: `http://127.0.0.1:${port}/mcp` }; + } + + async function startInteractiveAuthServer(): Promise<{ origin: string }> { + const httpServer: HttpServer = createHttpServer((req, res) => { + if (req.method !== 'POST' || (req.url !== '/register' && req.url !== '/token')) { + res.writeHead(404).end(); + return; + } + let body = ''; + req.on('data', (chunk: Buffer) => { + body += chunk.toString('utf-8'); + }); + req.on('end', () => { + if (req.url === '/register') { + const metadata = JSON.parse(body) as Record; + res.writeHead(201, { 'content-type': 'application/json' }); + res.end(JSON.stringify({ ...metadata, client_id: 'test-client' })); + return; + } + res.writeHead(200, { 'content-type': 'application/json' }); + res.end( + JSON.stringify({ access_token: 'fresh-token', token_type: 'Bearer', expires_in: 3600 }), + ); + }); + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + httpServers.push({ + close: () => + new Promise((resolve, reject) => { + httpServer.close((err) => (err === undefined || err === null ? resolve() : reject(err))); + }), + }); + const port = (httpServer.address() as HttpAddress).port; + return { origin: `http://127.0.0.1:${port}` }; + } + + async function seedDiscovery(name: string, url: string, authServerOrigin: string): Promise { + const provider = oauth.getProvider(name, url); + await provider.ready; + await provider.saveDiscoveryState({ + authorizationServerUrl: authServerOrigin, + authorizationServerMetadata: { + issuer: authServerOrigin, + authorization_endpoint: `${authServerOrigin}/authorize`, + token_endpoint: `${authServerOrigin}/token`, + registration_endpoint: `${authServerOrigin}/register`, + response_types_supported: ['code'], + grant_types_supported: ['authorization_code', 'refresh_token'], + token_endpoint_auth_methods_supported: ['none'], + }, + }); + } + + async function seedClient(name: string, url: string): Promise { + const provider = oauth.getProvider(name, url); + await provider.ready; + await provider.saveClientInformation({ + client_id: 'cached-client', + redirect_uris: ['http://127.0.0.1:45678/callback'], + token_endpoint_auth_method: 'none', + grant_types: ['authorization_code', 'refresh_token'], + response_types: ['code'], + } satisfies OAuthClientInformationFull); + } + + async function seedTokens( + name: string, + url: string, + tokens: { access_token: string; refresh_token?: string; expires_in?: number }, + ): Promise { + const provider = oauth.getProvider(name, url); + await provider.ready; + await provider.saveTokens({ token_type: 'Bearer', ...tokens }); + } + + async function deliverAuthCallback(authorizationUrl: string): Promise { + const url = new URL(authorizationUrl); + const redirectUri = url.searchParams.get('redirect_uri'); + const state = url.searchParams.get('state'); + expect(redirectUri).toBeTruthy(); + const callbackUrl = new URL(redirectUri!); + callbackUrl.searchParams.set('code', 'test-auth-code'); + if (state !== null) callbackUrl.searchParams.set('state', state); + const response = await fetch(callbackUrl); + expect(response.status).toBe(200); + await response.text(); + } + + describe('CRUD', () => { + it('round-trips add → get → update → remove through the real store and registry', async () => { + await expect(management.listServers()).resolves.toEqual([]); + + const added = await management.addServer({ + name: 'alpha', + transport: 'stdio', + command: 'npx', + env: { TOKEN: 'abc' }, + }); + expect(added).toEqual([ + { + name: 'alpha', + config: { transport: 'stdio', command: 'npx', env: { TOKEN: 'abc' } }, + source: 'global', + origin: join(home, 'mcp.json'), + mutable: true, + plugin: undefined, + }, + ]); + await expect(management.getServer('alpha')).resolves.toMatchObject({ + name: 'alpha', + mutable: true, + config: { command: 'npx' }, + }); + + const updated = await management.updateServer(stdioServer('alpha', 'node')); + expect(updated).toHaveLength(1); + expect(updated[0]?.config).toEqual({ transport: 'stdio', command: 'node' }); + await expect(store.get('alpha')).resolves.toMatchObject({ command: 'node' }); + + const remaining = await management.removeServer('alpha'); + expect(remaining).toEqual([]); + await expect(store.list()).resolves.toEqual([]); + }); + + it('waits for live config reconciliation listeners before returning', async () => { + let resolveStarted!: () => void; + const started = new Promise((resolve) => { + resolveStarted = resolve; + }); + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + store.onDidWrite((event) => { + resolveStarted(); + event.waitUntil(gate); + }); + + let completed = false; + const mutation = management.addServer(stdioServer('alpha')).then(() => { + completed = true; + }); + await started; + await Promise.resolve(); + expect(completed).toBe(false); + + release(); + await mutation; + expect(completed).toBe(true); + }); + + it('normalizes server names so the guard, the persisted key, and the list agree', async () => { + const added = await management.addServer(stdioServer(' alpha ')); + + expect(added.map((entry) => entry.name)).toEqual(['alpha']); + await expect(store.get('alpha')).resolves.toMatchObject({ name: 'alpha' }); + + const remaining = await management.removeServer(' alpha '); + expect(remaining).toEqual([]); + }); + + it('keeps the store duplicate error when re-adding a user-level name', async () => { + await management.addServer(stdioServer('alpha')); + + await expect(management.addServer(stdioServer('alpha'))).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP server "alpha" already exists', + }); + }); + + it('keeps the store not-found error when updating an unknown server', async () => { + await expect(management.updateServer(stdioServer('ghost'))).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + message: 'MCP server "ghost" was not found', + }); + }); + + it('treats removing an unknown server as a no-op returning the current catalog', async () => { + await management.addServer(stdioServer('alpha')); + + const remaining = await management.removeServer('ghost'); + expect(remaining.map((entry) => entry.name)).toEqual(['alpha']); + await expect(store.list()).resolves.toHaveLength(1); + }); + }); + + describe('read-only guards', () => { + it.each([ + ['add', (cwd: string) => management.addServer(stdioServer('local'), { cwd })], + ['update', (cwd: string) => management.updateServer(stdioServer('local'), { cwd })], + ['remove', (cwd: string) => management.removeServer('local', { cwd })], + ])('rejects %s when a trusted project-layer entry is read-only', async (_operation, mutate) => { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-read-only-')); + tempDirs.push(project); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ mcpServers: { local: { command: process.execPath } } }), + 'utf8', + ); + + await expect(mutate(project)).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: `MCP server "local" is read-only: it is defined in ${join(project, '.kimi-code', 'mcp.json')} — edit that file instead`, + }); + await expect(store.list()).resolves.toEqual([]); + }); + + it('lets a file entry shadow an enabled plugin entry', async () => { + pluginEntries = [ + { + name: 'plugin-demo:docs', + config: { transport: 'http', url: 'https://example.com/mcp' }, + pluginId: 'demo', + serverName: 'docs', + }, + ]; + const server: GlobalMcpServerConfig = { + name: 'plugin-demo:docs', + transport: 'http', + url: 'https://example.com/v2', + }; + + const added = await management.addServer(server); + const matches = added.filter((entry) => entry.name === 'plugin-demo:docs'); + expect(matches).toHaveLength(2); + expect(matches[0]).toMatchObject({ source: 'global', mutable: true }); + expect(matches[1]).toMatchObject({ source: 'plugin', mutable: false }); + + const remaining = await management.removeServer('plugin-demo:docs'); + expect(remaining.filter((entry) => entry.name === 'plugin-demo:docs')).toEqual([ + expect.objectContaining({ source: 'plugin', mutable: false }), + ]); + await expect(store.list()).resolves.toEqual([]); + }); + + it('rejects update against an enabled plugin entry that has no file entry yet', async () => { + pluginEntries = [ + { + name: 'plugin-demo:docs', + config: { transport: 'http', url: 'https://example.com/mcp' }, + pluginId: 'demo', + serverName: 'docs', + }, + ]; + + await expect( + management.updateServer({ + name: 'plugin-demo:docs', + transport: 'http', + url: 'https://example.com/v2', + }), + ).rejects.toMatchObject({ code: ErrorCodes.MCP_SERVER_NOT_FOUND }); + }); + + it('lets a mutable global entry be maintained past an enabled plugin collision', async () => { + await store.add(stdioServer('plugin-demo:docs', 'global-version')); + pluginEntries = [ + { + name: 'plugin-demo:docs', + config: { transport: 'http', url: 'https://example.com/mcp' }, + pluginId: 'demo', + serverName: 'docs', + }, + ]; + + await expect( + management.addServer(stdioServer('plugin-demo:docs', 'add-version')), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP server "plugin-demo:docs" already exists', + }); + + await management.updateServer(stdioServer('plugin-demo:docs', 'update-version')); + await expect(store.get('plugin-demo:docs')).resolves.toMatchObject({ + command: 'update-version', + }); + + const remaining = await management.removeServer('plugin-demo:docs'); + expect(remaining.filter((entry) => entry.name === 'plugin-demo:docs')).toEqual([ + expect.objectContaining({ source: 'plugin', mutable: false }), + ]); + await expect(store.list()).resolves.toEqual([]); + }); + + it('never blocks mutations on a disabled plugin descriptor', async () => { + pluginEntries = [ + { + name: 'plugin-demo:docs', + config: { transport: 'http', url: 'https://example.com/mcp', enabled: false }, + pluginId: 'demo', + serverName: 'docs', + }, + ]; + + const added = await management.addServer({ + name: 'plugin-demo:docs', + transport: 'http', + url: 'https://example.com/user', + }); + const matches = added.filter((entry) => entry.name === 'plugin-demo:docs'); + expect(matches).toHaveLength(2); + expect(matches[0]).toMatchObject({ source: 'global', mutable: true }); + expect(matches[1]).toMatchObject({ source: 'plugin', mutable: false }); + + await management.updateServer({ + name: 'plugin-demo:docs', + transport: 'http', + url: 'https://example.com/user-v2', + }); + await expect(store.get('plugin-demo:docs')).resolves.toMatchObject({ + url: 'https://example.com/user-v2', + }); + + const remaining = await management.removeServer('plugin-demo:docs'); + expect(remaining.filter((entry) => entry.name === 'plugin-demo:docs')).toHaveLength(1); + expect(remaining[0]).toMatchObject({ source: 'plugin' }); + }); + }); + + describe('mutation guard under a degraded registry', () => { + async function readStoreBytes(): Promise { + return storage.read(CONFIG_SCOPE, CONFIG_KEY); + } + + it('aborts add/update/remove without writing when plugin entries fail to load', async () => { + await management.addServer(stdioServer('alpha')); + const before = await readStoreBytes(); + pluginError = new Error2(ErrorCodes.PLUGIN_LOAD_FAILED, 'plugin state corrupt'); + + await expect(management.addServer(stdioServer('beta'))).rejects.toMatchObject({ + code: ErrorCodes.PLUGIN_LOAD_FAILED, + }); + await expect(management.updateServer(stdioServer('alpha', 'node'))).rejects.toMatchObject({ + code: ErrorCodes.PLUGIN_LOAD_FAILED, + }); + await expect(management.removeServer('alpha')).rejects.toMatchObject({ + code: ErrorCodes.PLUGIN_LOAD_FAILED, + }); + expect(await readStoreBytes()).toEqual(before); + }); + + it('aborts add/update/remove without writing when the user mcp.json is corrupt', async () => { + const corrupt = textEncoder.encode('{not json'); + await storage.write(CONFIG_SCOPE, CONFIG_KEY, corrupt); + + await expect(management.addServer(stdioServer('alpha'))).rejects.toMatchObject({ + code: ErrorCodes.CONFIG_INVALID, + }); + await expect(management.updateServer(stdioServer('alpha'))).rejects.toMatchObject({ + code: ErrorCodes.CONFIG_INVALID, + }); + await expect(management.removeServer('alpha')).rejects.toMatchObject({ + code: ErrorCodes.CONFIG_INVALID, + }); + expect(await readStoreBytes()).toEqual(corrupt); + }); + }); + + describe('redaction', () => { + it('redacts secret values of read-only entries while mutable entries keep full values', async () => { + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { + transport: 'stdio', + command: 'api-mcp', + env: { Z_KEY: 'z-value', A_TOKEN: 'a-value' }, + }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + await management.addServer({ + name: 'alpha', + transport: 'http', + url: 'https://example.com/mcp', + headers: { Authorization: 'Bearer secret' }, + }); + + const list = await management.listServers(); + + const plugin = list.find((entry) => entry.name === 'plugin-demo:api'); + expect(plugin).toMatchObject({ mutable: false, source: 'plugin' }); + expect(plugin?.config).toMatchObject({ envKeys: ['A_TOKEN', 'Z_KEY'] }); + expect(plugin?.config).not.toHaveProperty('env'); + expect(JSON.stringify(plugin?.config)).not.toContain('a-value'); + + const mutable = list.find((entry) => entry.name === 'alpha'); + expect(mutable).toMatchObject({ mutable: true, source: 'global' }); + expect(mutable?.config).toMatchObject({ headers: { Authorization: 'Bearer secret' } }); + + const got = await management.getServer('plugin-demo:api'); + expect(got.config).not.toHaveProperty('env'); + expect(got.config).toMatchObject({ envKeys: ['A_TOKEN', 'Z_KEY'] }); + }); + + it('lists project-layer entries as read-only redacted views when a cwd is given', async () => { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-proj-')); + tempDirs.push(project); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ + mcpServers: { + local: { + transport: 'http', + url: 'https://example.com/local', + headers: { 'X-Key': 'secret' }, + }, + }, + }), + 'utf8', + ); + + const list = await management.listServers({ cwd: project }); + + const local = list.find((entry) => entry.name === 'local'); + expect(local).toMatchObject({ + source: 'global', + mutable: false, + origin: join(project, '.kimi-code', 'mcp.json'), + }); + expect(local?.config).toMatchObject({ headerKeys: ['X-Key'] }); + expect(local?.config).not.toHaveProperty('headers'); + + const got = await management.getServer('local', { cwd: project }); + expect(got.mutable).toBe(false); + expect(got.config).not.toHaveProperty('headers'); + }); + + it('hides project-layer entries when the workspace is untrusted', async () => { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-untrusted-')); + tempDirs.push(project); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ mcpServers: { local: { command: process.execPath } } }), + 'utf8', + ); + await store.add(stdioServer('user', process.execPath)); + trusted = false; + + const list = await management.listServers({ cwd: project }); + + expect(list.map((entry) => entry.name)).toEqual(['user']); + await expect(management.getServer('local', { cwd: project })).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + }); + expect(getOrCreate).not.toHaveBeenCalled(); + }); + }); + + describe('testServer', () => { + it('probes an inline unsaved http config without touching the store', async () => { + const server = await startHttpServer(); + + const result = await management.testServer({ + server: { name: 'unsaved-probe', transport: 'http', url: server.url }, + }); + + expect(result.success).toBe(true); + expect(result.output).toContain('Connected to MCP server "unsaved-probe".'); + expect(result.output).toContain('Available tools: 1'); + expect(result.output).toContain('- echo: Echoes text'); + expect(getOrCreate).not.toHaveBeenCalled(); + await expect(store.list()).resolves.toEqual([]); + }, 20000); + + it('reports a clean failure for an unreachable inline http server', async () => { + const result = await management.testServer({ + server: { + name: 'down', + transport: 'http', + url: 'http://127.0.0.1:1/mcp', + startupTimeoutMs: 5_000, + }, + }); + + expect(result.success).toBe(false); + expect(result.output.length).toBeGreaterThan(0); + await expect(store.list()).resolves.toEqual([]); + }, 20000); + + it('probes an inline stdio config without retaining the probe cwd workspace', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-cwd-')); + tempDirs.push(cwd); + + const result = await management.testServer({ + server: { + name: 'stdio-probe', + transport: 'stdio', + command: process.execPath, + args: [stdioFixture], + }, + cwd, + }); + + expect(result.success).toBe(true); + expect(result.output).toContain('Available tools: 4'); + expect(result.output).toContain('- echo: Echoes input text'); + expect(findContaining).toHaveBeenCalledWith(cwd); + expect(getOrCreate).not.toHaveBeenCalled(); + }, 20000); + + it('probes a nested cwd against the containing workspace runtimes', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-nested-')); + tempDirs.push(cwd); + findContaining.mockReturnValue({ id: 'test-workspace' } as unknown as WorkspaceInstance); + + const result = await management.testServer({ + server: { + name: 'stdio-probe', + transport: 'stdio', + command: process.execPath, + args: [stdioFixture], + }, + cwd, + }); + + expect(result.success).toBe(true); + expect(result.output).toContain('Available tools: 4'); + expect(findContaining).toHaveBeenCalledWith(cwd); + expect(getOrCreate).not.toHaveBeenCalled(); + }, 20000); + + it('rejects a non-local runtime_id probe when no workspace contains the cwd', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-remote-miss-')); + tempDirs.push(cwd); + + await expect( + management.testServer({ + server: { + name: 'stdio-probe', + transport: 'stdio', + command: process.execPath, + args: [stdioFixture], + runtime_id: 'remote', + }, + cwd, + }), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: expect.stringContaining('runtime_id "remote"'), + }); + expect(findContaining).toHaveBeenCalledWith(cwd); + expect(getOrCreate).not.toHaveBeenCalled(); + }); + + it('probes a non-local runtime_id through the containing workspace', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-remote-hit-')); + tempDirs.push(cwd); + findContaining.mockReturnValue({ id: 'test-workspace' } as unknown as WorkspaceInstance); + + const result = await management.testServer({ + server: { + name: 'stdio-probe', + transport: 'stdio', + command: process.execPath, + args: [stdioFixture], + runtime_id: 'remote', + }, + cwd, + }); + + expect(result.success).toBe(true); + expect(findContaining).toHaveBeenCalledWith(cwd); + expect(getOrCreate).not.toHaveBeenCalled(); + }, 20000); + + it('keeps the transient local probe for an explicit local runtime_id', async () => { + const cwd = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-local-explicit-')); + tempDirs.push(cwd); + + const result = await management.testServer({ + server: { + name: 'stdio-probe', + transport: 'stdio', + command: process.execPath, + args: [stdioFixture], + runtime_id: 'local', + }, + cwd, + }); + + expect(result.success).toBe(true); + expect(findContaining).toHaveBeenCalledWith(cwd); + expect(getOrCreate).not.toHaveBeenCalled(); + }, 20000); + + it('rejects an inline probe whose name disagrees with the server config', async () => { + await expect( + management.testServer({ name: 'other', server: stdioServer('inline') }), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'Pass either an MCP server name or an inline server config, not both', + }); + await expect(management.testServer({})).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'Pass an MCP server name or an inline server config', + }); + }); + + it('rejects a name-only probe for an unknown server', async () => { + await expect(management.testServer({ name: 'ghost' })).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + message: 'MCP server "ghost" was not found', + }); + }); + + it('does not execute a project server while the workspace is untrusted', async () => { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-untrusted-probe-')); + tempDirs.push(project); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ + mcpServers: { + local: { command: process.execPath, args: [stdioFixture] }, + }, + }), + 'utf8', + ); + trusted = false; + + await expect(management.testServer({ name: 'local', cwd: project })).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + }); + }); + + it('waits for config and identity readiness before starting a probe', async () => { + let releaseConfig: () => void = () => undefined; + configReady = new Promise((resolve) => { + releaseConfig = resolve; + }); + let releaseIdentity: () => void = () => undefined; + identityReady = new Promise((resolve) => { + releaseIdentity = () => resolve(identitySnapshot); + }); + const cwd = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-ready-')); + tempDirs.push(cwd); + + const probe = management.testServer({ + server: { + name: 'stdio-probe', + transport: 'stdio', + command: process.execPath, + args: [stdioFixture], + }, + cwd, + }); + await Promise.resolve(); + expect(findContaining).not.toHaveBeenCalled(); + + releaseConfig(); + await Promise.resolve(); + expect(findContaining).not.toHaveBeenCalled(); + + releaseIdentity(); + await expect(probe).resolves.toMatchObject({ success: true }); + expect(findContaining).toHaveBeenCalledWith(cwd); + expect(getOrCreate).not.toHaveBeenCalled(); + }, 20000); + + it('rejects a name-only probe under an enabled runtime-name collision', async () => { + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { transport: 'http', url: 'https://example.com/plugin' }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + await store.add({ + name: 'plugin-demo:api', + transport: 'http', + url: 'https://example.com/user', + }); + + await expect(management.testServer({ name: 'plugin-demo:api' })).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP runtime name "plugin-demo:api" is shared by multiple enabled servers', + }); + }); + + it('probes the sole enabled entry when the name collides with a disabled shadow', async () => { + const server = await startHttpServer(); + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { transport: 'http', url: server.url }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + await store.add({ + name: 'plugin-demo:api', + transport: 'http', + url: 'http://127.0.0.1:1/unreachable', + enabled: false, + }); + + const result = await management.testServer({ name: 'plugin-demo:api' }); + + expect(result.success).toBe(true); + expect(result.output).toContain('echo'); + }, 20000); + + it('probes a plugin server by name', async () => { + const server = await startHttpServer(); + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { transport: 'http', url: server.url }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + + const result = await management.testServer({ name: 'plugin-demo:api' }); + + expect(result.success).toBe(true); + expect(result.output).toContain('Connected to MCP server "plugin-demo:api".'); + expect(result.output).toContain('echo'); + }, 20000); + }); + + describe('listAuthStatuses', () => { + it('classifies stored grants offline without probing', async () => { + await management.addServer({ + name: 'stale', + transport: 'http', + url: 'https://stale.example.test/mcp', + auth: 'oauth', + }); + await management.addServer({ + name: 'refreshable', + transport: 'http', + url: 'https://refresh.example.test/mcp', + auth: 'oauth', + }); + await management.addServer({ + name: 'fresh', + transport: 'http', + url: 'https://fresh.example.test/mcp', + auth: 'oauth', + }); + await management.addServer({ + name: 'bearer', + transport: 'http', + url: 'https://bearer.example.test/mcp', + bearerTokenEnvVar: 'API_TOKEN', + }); + await management.addServer(stdioServer('local-tool')); + await seedTokens('stale', 'https://stale.example.test/mcp', { + access_token: 'dead', + expires_in: -60, + }); + await seedTokens('refreshable', 'https://refresh.example.test/mcp', { + access_token: 'old', + refresh_token: 'still-good', + expires_in: -60, + }); + await seedTokens('fresh', 'https://fresh.example.test/mcp', { + access_token: 'good', + expires_in: 3600, + }); + + await expect(management.listAuthStatuses()).resolves.toEqual([ + { name: 'stale', authStatus: 'oauth-expired' }, + { name: 'refreshable', authStatus: 'oauth-authorized' }, + { name: 'fresh', authStatus: 'oauth-authorized' }, + { name: 'bearer', authStatus: 'bearer-token' }, + { name: 'local-tool', authStatus: 'not-applicable' }, + ]); + }); + + it('short-circuits disabled servers even under online verification', async () => { + await management.addServer({ + name: 'off', + transport: 'http', + url: 'https://disabled.example.test/mcp', + auth: 'oauth', + enabled: false, + }); + + await expect(management.listAuthStatuses({ verify: true })).resolves.toEqual([ + { name: 'off', authStatus: 'not-applicable' }, + ]); + }); + + it('classifies unpinned servers without a stored grant offline when verify is false', async () => { + const server = await startCountingServer(); + await management.addServer({ name: 'plain', transport: 'http', url: server.url }); + await management.addServer({ + name: 'challenged', + transport: 'http', + url: 'https://challenged.example.test/mcp', + auth: 'oauth', + }); + + await expect(management.listAuthStatuses({ verify: false })).resolves.toEqual([ + { name: 'plain', authStatus: 'not-applicable' }, + { name: 'challenged', authStatus: 'oauth-required' }, + ]); + expect(server.requestCount()).toBe(0); + }, 20000); + + it('detects an implicit OAuth challenge when verify is omitted', async () => { + const gated = await startGatedServer(); + await management.addServer({ name: 'detected', transport: 'http', url: gated.url }); + + await expect(management.listAuthStatuses()).resolves.toEqual([ + { name: 'detected', authStatus: 'oauth-required' }, + ]); + }, 20000); + + it('verify settles a stored-but-rejected grant as oauth-expired through a real probe', async () => { + const gated = await startGatedServer(); + await management.addServer({ + name: 'stale', + transport: 'http', + url: gated.url, + auth: 'oauth', + }); + await seedDiscovery('stale', gated.url, gated.origin); + await seedClient('stale', gated.url); + await seedTokens('stale', gated.url, { + access_token: 'wrong', + refresh_token: 'dead-refresh', + }); + + await expect(management.listAuthStatuses({ verify: true })).resolves.toEqual([ + { name: 'stale', authStatus: 'oauth-expired' }, + ]); + }, 20000); + }); + + describe('inspectServers', () => { + it('includes trusted project-layer entries when cwd is provided', async () => { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-inspect-project-')); + tempDirs.push(project); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ + mcpServers: { + local: { + transport: 'http', + url: 'https://project.example.test/mcp', + headers: { 'X-Key': 'secret' }, + }, + }, + }), + 'utf8', + ); + + const inspections = await management.inspectServers(undefined, { cwd: project }); + + expect(inspections).toEqual([ + expect.objectContaining({ + serverId: 'global:local', + runtimeName: 'local', + canonicalUrl: 'https://project.example.test/mcp', + editable: false, + authStatus: 'not-applicable', + }), + ]); + }); + + it('lists the locator-addressed catalog with offline classifications and redacted configs', async () => { + const plain = await startHttpServer(); + await management.addServer({ name: 'plain', transport: 'http', url: plain.url }); + await management.addServer(stdioServer('local-tool')); + await management.addServer({ + name: 'bearer', + transport: 'http', + url: 'https://bearer.example.test/mcp', + bearerTokenEnvVar: 'API_TOKEN', + }); + await management.addServer({ + name: 'off', + transport: 'http', + url: 'https://off.example.test/mcp', + enabled: false, + }); + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { transport: 'http', url: plain.url, headers: { 'X-Key': 'secret' } }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + + const inspections = await management.inspectServers(); + const byId = new Map(inspections.map((server) => [server.serverId, server])); + expect([...byId.keys()].toSorted()).toEqual([ + 'global:bearer', + 'global:local-tool', + 'global:off', + 'global:plain', + 'plugin:demo:api', + ]); + + expect(byId.get('global:plain')).toMatchObject({ + locator: { source: 'global', name: 'plain' }, + runtimeName: 'plain', + canonicalUrl: plain.url, + origin: 'global', + enabled: true, + editable: true, + authStatus: 'not-applicable', + }); + expect(byId.get('global:local-tool')).toMatchObject({ + canonicalUrl: undefined, + authStatus: 'not-applicable', + }); + expect(byId.get('global:bearer')).toMatchObject({ authStatus: 'bearer-token' }); + expect(byId.get('global:off')).toMatchObject({ + enabled: false, + authStatus: 'not-applicable', + }); + const plugin = byId.get('plugin:demo:api'); + expect(plugin).toMatchObject({ + locator: { source: 'plugin', pluginId: 'demo', serverName: 'api' }, + runtimeName: 'plugin-demo:api', + canonicalUrl: plain.url, + origin: 'plugin', + enabled: true, + editable: false, + authStatus: 'not-applicable', + }); + expect(plugin?.config).toMatchObject({ headerKeys: ['X-Key'] }); + expect(plugin?.config).not.toHaveProperty('headers'); + expect(JSON.stringify(plugin?.config)).not.toContain('secret'); + }, 20000); + + it('marks a runtime-name collision as unavailable instead of probing it', async () => { + const plain = await startHttpServer(); + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { transport: 'http', url: plain.url }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + await store.add({ name: 'plugin-demo:api', transport: 'http', url: plain.url }); + + const targeted = await management.inspectServers([ + { source: 'plugin', pluginId: 'demo', serverName: 'api' }, + ]); + expect(targeted).toHaveLength(1); + expect(targeted[0]).toMatchObject({ + runtimeName: 'plugin-demo:api', + authStatus: 'unavailable', + error: 'MCP runtime name "plugin-demo:api" is not unique', + }); + + const all = await management.inspectServers(); + expect(all.filter((server) => server.runtimeName === 'plugin-demo:api')).toHaveLength(2); + }, 20000); + + it('settles needs-auth probes by their stored grant: expired with one, required without', async () => { + const gated = await startGatedServer(); + await management.addServer({ + name: 'stale', + transport: 'http', + url: gated.url, + auth: 'oauth', + }); + await management.addServer({ + name: 'challenged', + transport: 'http', + url: `${gated.origin}/other`, + auth: 'oauth', + }); + await seedDiscovery('stale', gated.url, gated.origin); + await seedClient('stale', gated.url); + await seedTokens('stale', gated.url, { + access_token: 'wrong', + refresh_token: 'dead-refresh', + }); + await seedDiscovery('challenged', `${gated.origin}/other`, gated.origin); + await seedClient('challenged', `${gated.origin}/other`); + + const inspections = await management.inspectServers(); + const byName = new Map(inspections.map((server) => [server.runtimeName, server])); + + expect(byName.get('stale')).toMatchObject({ authStatus: 'oauth-expired' }); + expect(byName.get('challenged')).toMatchObject({ authStatus: 'oauth-required' }); + }, 20000); + + it('rejects unknown locators with the shared not-found error', async () => { + await expect( + management.inspectServers([{ source: 'global', name: 'missing' }]), + ).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + message: 'MCP server "missing" was not found', + }); + await expect( + management.inspectServers([{ source: 'plugin', pluginId: 'demo', serverName: 'ghost' }]), + ).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + message: 'MCP server "demo/ghost" was not found', + }); + }); + }); + + describe('resolveServerByName', () => { + it('resolves a project-layer-only name when cwd is provided', async () => { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-resolve-project-')); + tempDirs.push(project); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ mcpServers: { local: { command: process.execPath } } }), + 'utf8', + ); + + await expect(management.resolveServerByName('local', { cwd: project })).resolves.toEqual({ + source: 'global', + name: 'local', + }); + }); + + it('resolves a unique global name to its locator', async () => { + await management.addServer(stdioServer('alpha')); + + await expect(management.resolveServerByName('alpha')).resolves.toEqual({ + source: 'global', + name: 'alpha', + }); + }); + + it('resolves the sole enabled owner past a disabled shadow', async () => { + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { transport: 'http', url: 'https://example.com/mcp' }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + await store.add({ + name: 'plugin-demo:api', + transport: 'http', + url: 'https://example.com/user', + enabled: false, + }); + + await expect(management.resolveServerByName('plugin-demo:api')).resolves.toEqual({ + source: 'plugin', + pluginId: 'demo', + serverName: 'api', + }); + }); + + it('rejects an unknown name with the shared not-found error', async () => { + await expect(management.resolveServerByName('ghost')).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + message: 'MCP server "ghost" was not found', + }); + }); + + it('rejects a name shared by enabled entries, pointing at the locator RPC', async () => { + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { transport: 'http', url: 'https://example.com/mcp' }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + await store.add({ + name: 'plugin-demo:api', + transport: 'http', + url: 'https://example.com/user', + }); + + await expect(management.resolveServerByName('plugin-demo:api')).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: + 'MCP runtime name "plugin-demo:api" is shared by multiple enabled servers; use the locator-addressed RPC instead', + }); + }); + }); + + describe('OAuth operations', () => { + it('begins authorization against the project-layer URL when cwd is provided', async () => { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-begin-project-')); + tempDirs.push(project); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ + mcpServers: { + oauthable: { + transport: 'http', + url: 'https://project.example.test/mcp', + auth: 'oauth', + }, + }, + }), + 'utf8', + ); + const cancel = vi.fn(async () => undefined); + const begin = vi.spyOn(oauth, 'beginAuthorization').mockResolvedValue({ + authorizationUrl: new URL('https://project.example.test/authorize'), + complete: vi.fn(async () => undefined), + cancel, + }); + + const result = await management.beginServerAuth( + { source: 'global', name: 'oauthable' }, + { cwd: project }, + ); + + expect(begin).toHaveBeenCalledWith('oauthable', 'https://project.example.test/mcp'); + if (result.status === 'authorization-required') { + await management.cancelServerAuth({ flowId: result.flowId }); + } + }); + + it('resets credentials for the project-layer URL when cwd is provided', async () => { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-management-reset-project-')); + tempDirs.push(project); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ + mcpServers: { + oauthable: { + transport: 'http', + url: 'https://project.example.test/mcp', + auth: 'oauth', + }, + }, + }), + 'utf8', + ); + const invalidate = vi.spyOn(oauth, 'invalidate').mockResolvedValue(undefined); + + await management.resetServerAuth( + { source: 'global', name: 'oauthable' }, + { cwd: project }, + ); + + expect(invalidate).toHaveBeenCalledWith( + 'oauthable', + 'https://project.example.test/mcp', + ); + }); + + it('rejects begin for entries that cannot run an OAuth flow', async () => { + await management.addServer(stdioServer('local-tool')); + await management.addServer({ + name: 'bearer', + transport: 'http', + url: 'https://bearer.example.test/mcp', + bearerTokenEnvVar: 'API_TOKEN', + }); + await management.addServer({ + name: 'static-headers', + transport: 'http', + url: 'https://static.example.test/mcp', + headers: { 'X-Key': 'v' }, + }); + + await expect( + management.beginServerAuth({ source: 'global', name: 'local-tool' }), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP server "local-tool" does not use a remote transport', + }); + await expect( + management.beginServerAuth({ source: 'global', name: 'bearer' }), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP server "bearer" uses a static bearer token', + }); + await expect( + management.beginServerAuth({ source: 'global', name: 'static-headers' }), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP server "static-headers" uses static headers and is not marked for OAuth', + }); + await expect( + management.beginServerAuth({ source: 'global', name: 'missing' }), + ).rejects.toMatchObject({ code: ErrorCodes.MCP_SERVER_NOT_FOUND }); + }); + + it('refuses credential operations under an enabled runtime-name collision', async () => { + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { transport: 'http', url: 'https://example.com/mcp', auth: 'oauth' }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + await store.add({ + name: 'plugin-demo:api', + transport: 'http', + url: 'https://example.com/user', + auth: 'oauth', + }); + + const ambiguous = { + code: ErrorCodes.REQUEST_INVALID, + message: + 'MCP runtime name "plugin-demo:api" is shared by multiple enabled servers; use the locator-addressed RPC instead', + }; + await expect( + management.beginServerAuth({ source: 'plugin', pluginId: 'demo', serverName: 'api' }), + ).rejects.toMatchObject(ambiguous); + await expect( + management.resetServerAuth({ source: 'global', name: 'plugin-demo:api' }), + ).rejects.toMatchObject(ambiguous); + }); + + it('returns already-authorized when a valid grant is stored', async () => { + const authServer = await startInteractiveAuthServer(); + const mcpUrl = `${authServer.origin}/mcp`; + await management.addServer({ + name: 'oauthable', + transport: 'http', + url: mcpUrl, + auth: 'oauth', + }); + await seedDiscovery('oauthable', mcpUrl, authServer.origin); + await seedTokens('oauthable', mcpUrl, { + access_token: 'stale-access', + refresh_token: 'good-refresh', + }); + + await expect( + management.beginServerAuth({ source: 'global', name: 'oauthable' }), + ).resolves.toEqual({ status: 'already-authorized' }); + }, 20000); + + it('drives a full browser flow: begin → callback → complete → tokens persisted', async () => { + const authServer = await startInteractiveAuthServer(); + const mcpUrl = `${authServer.origin}/mcp`; + await management.addServer({ + name: 'oauthable', + transport: 'http', + url: mcpUrl, + auth: 'oauth', + }); + await seedDiscovery('oauthable', mcpUrl, authServer.origin); + const events: McpOAuthEvent[] = []; + oauth.onEvent((event) => events.push(event)); + + const begun = await management.beginServerAuth({ source: 'global', name: 'oauthable' }); + if (begun.status !== 'authorization-required') { + throw new Error(`expected authorization-required, got ${begun.status}`); + } + expect(begun.authorizationUrl).toContain(`${authServer.origin}/authorize`); + + const completing = management.completeServerAuth({ + flowId: begun.flowId, + timeoutMs: 10_000, + }); + await deliverAuthCallback(begun.authorizationUrl); + await completing; + + expect((await oauth.tokenState('oauthable', mcpUrl)).hasTokens).toBe(true); + expect(events).toContainEqual({ + type: 'tokens-saved', + serverName: 'oauthable', + serverUrl: mcpUrl, + }); + }, 20000); + + it('cancel tears down an active flow, so a later complete rejects as unknown', async () => { + const authServer = await startInteractiveAuthServer(); + const mcpUrl = `${authServer.origin}/mcp`; + await management.addServer({ + name: 'oauthable', + transport: 'http', + url: mcpUrl, + auth: 'oauth', + }); + await seedDiscovery('oauthable', mcpUrl, authServer.origin); + + const begun = await management.beginServerAuth({ source: 'global', name: 'oauthable' }); + if (begun.status !== 'authorization-required') { + throw new Error(`expected authorization-required, got ${begun.status}`); + } + + await management.cancelServerAuth({ flowId: begun.flowId }); + + await expect( + management.completeServerAuth({ flowId: begun.flowId, timeoutMs: 1000 }), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: `Unknown MCP OAuth flow: ${begun.flowId}`, + }); + }, 20000); + + it('keeps a joined flow usable when an earlier flow handle is cancelled', async () => { + const authServer = await startInteractiveAuthServer(); + const mcpUrl = `${authServer.origin}/mcp`; + await management.addServer({ + name: 'oauthable', + transport: 'http', + url: mcpUrl, + auth: 'oauth', + }); + await seedDiscovery('oauthable', mcpUrl, authServer.origin); + + const first = await management.beginServerAuth({ source: 'global', name: 'oauthable' }); + const second = await management.beginServerAuth({ source: 'global', name: 'oauthable' }); + if ( + first.status !== 'authorization-required' || + second.status !== 'authorization-required' + ) { + throw new Error('expected both flows to require authorization'); + } + expect(second.authorizationUrl).toBe(first.authorizationUrl); + + await management.cancelServerAuth({ flowId: first.flowId }); + const completing = management.completeServerAuth({ + flowId: second.flowId, + timeoutMs: 10_000, + }); + await deliverAuthCallback(second.authorizationUrl); + await completing; + + expect((await oauth.tokenState('oauthable', mcpUrl)).hasTokens).toBe(true); + }, 20000); + + it('expires an idle flow: the flow is cancelled and a later complete rejects as unknown', async () => { + await management.addServer({ + name: 'oauthable', + transport: 'http', + url: 'https://oauthable.example.test/mcp', + auth: 'oauth', + }); + const cancel = vi.fn(async () => undefined); + const beginSpy = vi.spyOn(oauth, 'beginAuthorization').mockResolvedValue({ + authorizationUrl: new URL('https://oauthable.example.test/authorize'), + complete: vi.fn(async () => undefined), + cancel, + }); + vi.useFakeTimers(); + try { + const begun = await management.beginServerAuth({ source: 'global', name: 'oauthable' }); + if (begun.status !== 'authorization-required') { + throw new Error(`expected authorization-required, got ${begun.status}`); + } + + await vi.advanceTimersByTimeAsync(15 * 60_000); + + expect(cancel).toHaveBeenCalledTimes(1); + await expect( + management.completeServerAuth({ flowId: begun.flowId, timeoutMs: 1000 }), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: `Unknown MCP OAuth flow: ${begun.flowId}`, + }); + } finally { + vi.useRealTimers(); + beginSpy.mockRestore(); + } + }); + + it('complete rejects on timeout when the browser callback never arrives', async () => { + const authServer = await startInteractiveAuthServer(); + const mcpUrl = `${authServer.origin}/mcp`; + await management.addServer({ + name: 'oauthable', + transport: 'http', + url: mcpUrl, + auth: 'oauth', + }); + await seedDiscovery('oauthable', mcpUrl, authServer.origin); + + const begun = await management.beginServerAuth({ source: 'global', name: 'oauthable' }); + if (begun.status !== 'authorization-required') { + throw new Error(`expected authorization-required, got ${begun.status}`); + } + + await expect( + management.completeServerAuth({ flowId: begun.flowId, timeoutMs: 200 }), + ).rejects.toThrow(/OAuth callback timed out/); + }, 20000); + + it('complete rejects an unknown flow while cancel ignores it', async () => { + await expect(management.completeServerAuth({ flowId: 'unknown-flow' })).rejects.toMatchObject( + { + code: ErrorCodes.REQUEST_INVALID, + message: 'Unknown MCP OAuth flow: unknown-flow', + }, + ); + await expect(management.cancelServerAuth({ flowId: 'unknown-flow' })).resolves.toBeUndefined(); + }); + + it('complete rejects a timeoutMs outside the setTimeout range', async () => { + await expect( + management.completeServerAuth({ flowId: 'unknown-flow', timeoutMs: 2 ** 31 }), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP OAuth timeoutMs must be an integer between 1 and 2147483647', + }); + await expect( + management.completeServerAuth({ flowId: 'unknown-flow', timeoutMs: 0 }), + ).rejects.toMatchObject({ code: ErrorCodes.REQUEST_INVALID }); + await expect( + management.completeServerAuth({ flowId: 'unknown-flow', timeoutMs: 1.5 }), + ).rejects.toMatchObject({ code: ErrorCodes.REQUEST_INVALID }); + }); + + it('reset invalidates stored credentials and broadcasts the event', async () => { + await management.addServer({ + name: 'oauthable', + transport: 'http', + url: 'https://oauth.example.test/mcp', + auth: 'oauth', + }); + await seedTokens('oauthable', 'https://oauth.example.test/mcp', { + access_token: 'good', + expires_in: 3600, + }); + const events: McpOAuthEvent[] = []; + oauth.onEvent((event) => events.push(event)); + + await management.resetServerAuth({ source: 'global', name: 'oauthable' }); + + expect((await oauth.tokenState('oauthable', 'https://oauth.example.test/mcp')).hasTokens).toBe( + false, + ); + expect(events).toContainEqual({ + type: 'tokens-invalidated', + serverName: 'oauthable', + serverUrl: 'https://oauth.example.test/mcp', + scope: 'all', + }); + }); + + it('resets a plugin server by locator', async () => { + pluginEntries = [ + { + name: 'plugin-demo:api', + config: { transport: 'http', url: 'https://example.com/mcp', auth: 'oauth' }, + pluginId: 'demo', + serverName: 'api', + }, + ]; + + await expect( + management.resetServerAuth({ source: 'plugin', pluginId: 'demo', serverName: 'api' }), + ).resolves.toBeUndefined(); + }); + + it('rejects reset for a stdio locator', async () => { + await management.addServer(stdioServer('local-tool')); + + await expect( + management.resetServerAuth({ source: 'global', name: 'local-tool' }), + ).rejects.toMatchObject({ + code: ErrorCodes.REQUEST_INVALID, + message: 'MCP server "local-tool" does not use a remote transport', + }); + }); + }); +}); diff --git a/packages/agent-core-v2/test/app/mcpRegistry/mcpRegistry.test.ts b/packages/agent-core-v2/test/app/mcpRegistry/mcpRegistry.test.ts new file mode 100644 index 0000000000..f9c626006a --- /dev/null +++ b/packages/agent-core-v2/test/app/mcpRegistry/mcpRegistry.test.ts @@ -0,0 +1,428 @@ +import { mkdtempSync } from 'node:fs'; +import { mkdir, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join, relative } from 'pathe'; + +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { DisposableStore } from '#/_base/di/lifecycle'; +import { encodeWorkDirKey } from '#/_base/utils/workdir-slug'; +import { createServices } from '#/_base/di/test'; +import { IBootstrapService } from '#/app/bootstrap/bootstrap'; +import { + IMcpConfigStore, + McpConfigStore, + type GlobalMcpServerConfig, +} from '#/app/mcpConfig/configStore'; +import { IPluginService } from '#/app/plugin/plugin'; +import type { PluginMcpServerEntry } from '#/app/plugin/types'; +import { IMcpRegistryService, mcpServerConfigsEqual } from '#/app/mcpRegistry/mcpRegistry'; +import { McpRegistryService } from '#/app/mcpRegistry/mcpRegistryService'; +import { ErrorCodes } from '#/errors'; +import { HostFileSystem } from '#/os/backends/node-local/hostFsService'; +import { IHostFileSystem } from '#/os/interface/hostFileSystem'; +import { InMemoryStorageService } from '#/persistence/backends/memory/inMemoryStorageService'; +import { IAtomicDocumentStore } from '#/persistence/interface/atomicDocumentStore'; +import { IFileSystemStorageService } from '#/persistence/interface/storage'; + +function stdioServer(name: string, command = 'npx'): GlobalMcpServerConfig { + return { name, transport: 'stdio', command }; +} + +function pluginEntry( + pluginId: string, + serverName: string, + config: PluginMcpServerEntry['config'], +): PluginMcpServerEntry { + return { name: `plugin-${pluginId}:${serverName}`, config, pluginId, serverName }; +} + +async function writeJson(file: string, value: unknown): Promise { + await mkdir(join(file, '..'), { recursive: true }); + await writeFile(file, JSON.stringify(value), 'utf8'); +} + +describe('McpRegistryService', () => { + let home: string; + let disposables: DisposableStore; + let tempDirs: string[]; + let store: IMcpConfigStore; + let pluginEntries: PluginMcpServerEntry[]; + let pluginError: Error | undefined; + let trusted: boolean; + let trustedKey: string | undefined; + let registry: IMcpRegistryService; + + beforeEach(() => { + home = mkdtempSync(join(tmpdir(), 'kimi-mcp-registry-home-')); + vi.stubEnv('KIMI_CODE_HOME', home); + disposables = new DisposableStore(); + tempDirs = [home]; + pluginEntries = []; + pluginError = undefined; + trusted = true; + trustedKey = undefined; + const ix = createServices(disposables, { + additionalServices: (reg) => { + reg.defineInstance(IFileSystemStorageService, new InMemoryStorageService()); + reg.definePartialInstance(IBootstrapService, { homeDir: home }); + reg.define(IMcpConfigStore, McpConfigStore); + reg.definePartialInstance(IPluginService, { + mcpServerEntries: async () => { + if (pluginError !== undefined) throw pluginError; + return pluginEntries; + }, + }); + reg.defineInstance(IHostFileSystem, new HostFileSystem()); + reg.definePartialInstance(IAtomicDocumentStore, { + get: async (_scope: string, key: string) => { + if (!trusted || (trustedKey !== undefined && key !== encodeWorkDirKey(trustedKey))) { + return undefined; + } + return {} as T; + }, + }); + reg.define(IMcpRegistryService, McpRegistryService); + }, + }); + store = ix.get(IMcpConfigStore); + registry = ix.get(IMcpRegistryService); + }); + + afterEach(async () => { + disposables.dispose(); + vi.unstubAllEnvs(); + await Promise.all(tempDirs.map((dir) => rm(dir, { recursive: true, force: true }))); + }); + + async function makeProject(): Promise<{ project: string; sub: string }> { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-registry-proj-')); + tempDirs.push(project); + await mkdir(join(project, '.git'), { recursive: true }); + const sub = join(project, 'pkg'); + await mkdir(sub, { recursive: true }); + return { project, sub }; + } + + describe('list', () => { + it('lists user-level entries with the store path as origin when no cwd is given', async () => { + await store.add({ + name: 'fs', + transport: 'stdio', + command: 'fs-mcp', + args: ['--readonly'], + }); + await store.add({ name: 'docs', transport: 'http', url: 'https://example.com/mcp' }); + + const entries = await registry.list(); + + expect(entries).toHaveLength(2); + expect(entries[0]).toEqual({ + name: 'fs', + config: { transport: 'stdio', command: 'fs-mcp', args: ['--readonly'] }, + source: 'global', + origin: join(home, 'mcp.json'), + mutable: true, + plugin: undefined, + }); + expect(entries[1]).toMatchObject({ + name: 'docs', + source: 'global', + origin: join(home, 'mcp.json'), + mutable: true, + }); + }); + + it('merges the three file layers with origin and mutability tracking when a cwd is given', async () => { + await writeJson(join(home, 'mcp.json'), { + mcpServers: { + shared: { command: 'user-version' }, + userOnly: { command: 'user-only' }, + }, + }); + const { project, sub } = await makeProject(); + await writeJson(join(project, '.mcp.json'), { + mcpServers: { + shared: { command: 'repo-version', cwd: './bin' }, + repoOnly: { command: 'repo-only' }, + }, + }); + await writeJson(join(sub, '.kimi-code', 'mcp.json'), { + mcpServers: { localOnly: { command: 'local-only' } }, + }); + + const entries = await registry.list({ cwd: sub }); + const byName = new Map(entries.map((entry) => [entry.name, entry])); + expect([...byName.keys()].toSorted()).toEqual([ + 'localOnly', + 'repoOnly', + 'shared', + 'userOnly', + ]); + + expect(byName.get('shared')).toMatchObject({ + source: 'global', + mutable: false, + origin: join(project, '.mcp.json'), + }); + expect(byName.get('shared')?.config).toEqual({ + transport: 'stdio', + command: 'repo-version', + cwd: join(project, 'bin'), + }); + expect(byName.get('userOnly')).toMatchObject({ + mutable: true, + origin: join(home, 'mcp.json'), + }); + expect(byName.get('repoOnly')).toMatchObject({ + mutable: false, + origin: join(project, '.mcp.json'), + }); + expect(byName.get('localOnly')).toMatchObject({ + mutable: false, + origin: join(sub, '.kimi-code', 'mcp.json'), + }); + }); + + it('loads only user and plugin entries when the workspace is untrusted', async () => { + await store.add(stdioServer('userOnly', 'user-only')); + const { project, sub } = await makeProject(); + await writeJson(join(project, '.mcp.json'), { + mcpServers: { repoOnly: { command: 'repo-only' } }, + }); + await writeJson(join(sub, '.kimi-code', 'mcp.json'), { + mcpServers: { localOnly: { command: 'local-only' } }, + }); + pluginEntries = [ + pluginEntry('demo', 'api', { transport: 'stdio', command: 'plugin-only' }), + ]; + trusted = false; + + const entries = await registry.list({ cwd: sub }); + + expect(entries.map((entry) => entry.name).toSorted()).toEqual([ + 'plugin-demo:api', + 'userOnly', + ]); + }); + + it('checks trust at the queried cwd rather than the canonical git root', async () => { + const { project, sub } = await makeProject(); + await writeJson(join(project, '.mcp.json'), { + mcpServers: { projectOnly: { command: 'project-only' } }, + }); + trustedKey = project; + + const entries = await registry.list({ cwd: sub }); + + expect(entries.map((entry) => entry.name)).not.toContain('projectOnly'); + }); + + it('lists project layers when the queried subdirectory cwd itself is trusted', async () => { + const { project, sub } = await makeProject(); + await writeJson(join(project, '.mcp.json'), { + mcpServers: { projectOnly: { command: 'project-only' } }, + }); + await writeJson(join(sub, '.kimi-code', 'mcp.json'), { + mcpServers: { localOnly: { command: 'local-only' } }, + }); + trustedKey = sub; + + const entries = await registry.list({ cwd: sub }); + + expect(entries.map((entry) => entry.name)).toEqual( + expect.arrayContaining(['projectOnly', 'localOnly']), + ); + }); + + it('resolves a relative non-git cwd before checking workspace trust', async () => { + const project = mkdtempSync(join(tmpdir(), 'kimi-mcp-registry-non-git-')); + tempDirs.push(project); + await writeJson(join(project, '.mcp.json'), { + mcpServers: { relativeOnly: { command: 'relative-only' } }, + }); + trustedKey = project; + + const entries = await registry.list({ cwd: relative(process.cwd(), project) }); + + expect(entries.map((entry) => entry.name)).toContain('relativeOnly'); + }); + + it('exposes plugin servers as read-only entries with their effective config', async () => { + pluginEntries = [ + pluginEntry('demo', 'finance', { + transport: 'stdio', + command: 'finance-mcp', + enabled: true, + }), + pluginEntry('demo', 'docs', { transport: 'http', url: 'https://example.com/mcp' }), + ]; + + const entries = await registry.list(); + + expect(entries).toHaveLength(2); + expect(entries[0]).toEqual({ + name: 'plugin-demo:finance', + config: { transport: 'stdio', command: 'finance-mcp', enabled: true }, + source: 'plugin', + origin: 'demo', + mutable: false, + plugin: { id: 'demo', name: 'finance' }, + }); + expect(entries[1]).toMatchObject({ + name: 'plugin-demo:docs', + source: 'plugin', + origin: 'demo', + mutable: false, + plugin: { id: 'demo', name: 'docs' }, + }); + }); + + it('keeps both sides of a runtime-name collision instead of hiding one', async () => { + await store.add(stdioServer('plugin-demo:api', 'user-version')); + pluginEntries = [ + pluginEntry('demo', 'api', { transport: 'http', url: 'https://example.com/mcp' }), + ]; + + const matches = (await registry.list()).filter((entry) => entry.name === 'plugin-demo:api'); + + expect(matches).toHaveLength(2); + expect(matches[0]).toMatchObject({ source: 'global', mutable: true }); + expect(matches[1]).toMatchObject({ source: 'plugin', mutable: false, origin: 'demo' }); + }); + + it('propagates a plugin listing failure instead of reading as not configured', async () => { + await store.add(stdioServer('fs')); + pluginError = new Error('plugin state corrupt'); + + await expect(registry.list()).rejects.toThrow('plugin state corrupt'); + await expect(registry.get('fs')).rejects.toThrow('plugin state corrupt'); + await expect(registry.resolveRuntimeTarget('fs')).rejects.toThrow('plugin state corrupt'); + }); + }); + + describe('get', () => { + it('returns the first match on a runtime-name collision (globals list first)', async () => { + await store.add(stdioServer('plugin-demo:api', 'user-version')); + pluginEntries = [ + pluginEntry('demo', 'api', { transport: 'http', url: 'https://example.com/mcp' }), + ]; + + const entry = await registry.get('plugin-demo:api'); + + expect(entry).toMatchObject({ + source: 'global', + mutable: true, + config: { command: 'user-version' }, + }); + }); + + it('rejects unknown names with the shared not-found error', async () => { + await expect(registry.get('missing')).rejects.toMatchObject({ + code: ErrorCodes.MCP_SERVER_NOT_FOUND, + message: 'MCP server "missing" was not found', + }); + }); + }); + + describe('resolveRuntimeTarget', () => { + it('prefers the file entry over an enabled plugin entry', async () => { + await store.add(stdioServer('plugin-demo:api', 'user-version')); + pluginEntries = [ + pluginEntry('demo', 'api', { transport: 'http', url: 'https://example.com/mcp' }), + ]; + + await expect(registry.resolveRuntimeTarget('plugin-demo:api')).resolves.toMatchObject({ + source: 'global', + config: { command: 'user-version' }, + }); + }); + + it('lets a file entry win by presence even when the file entry is disabled', async () => { + await store.add({ ...stdioServer('plugin-demo:api', 'user-version'), enabled: false }); + pluginEntries = [ + pluginEntry('demo', 'api', { transport: 'http', url: 'https://example.com/mcp' }), + ]; + + await expect(registry.resolveRuntimeTarget('plugin-demo:api')).resolves.toMatchObject({ + source: 'global', + config: { command: 'user-version', enabled: false }, + }); + }); + + it('treats a disabled plugin descriptor as absent and falls back to the file entry', async () => { + await store.add(stdioServer('plugin-demo:api', 'user-version')); + pluginEntries = [ + pluginEntry('demo', 'api', { + transport: 'http', + url: 'https://example.com/mcp', + enabled: false, + }), + ]; + + await expect(registry.resolveRuntimeTarget('plugin-demo:api')).resolves.toMatchObject({ + source: 'global', + config: { command: 'user-version' }, + }); + + await store.remove('plugin-demo:api'); + await expect(registry.resolveRuntimeTarget('plugin-demo:api')).resolves.toBeUndefined(); + }); + + it('never picks a disabled plugin descriptor when it is the only entry', async () => { + pluginEntries = [ + pluginEntry('demo', 'api', { + transport: 'http', + url: 'https://example.com/mcp', + enabled: false, + }), + ]; + await expect(registry.resolveRuntimeTarget('plugin-demo:api')).resolves.toBeUndefined(); + }); + + it('resolves an enabled plugin entry', async () => { + pluginEntries = [ + pluginEntry('demo', 'api', { transport: 'http', url: 'https://example.com/mcp' }), + ]; + await expect(registry.resolveRuntimeTarget('plugin-demo:api')).resolves.toMatchObject({ + source: 'plugin', + }); + }); + + it('returns undefined for a name no source defines', async () => { + await expect(registry.resolveRuntimeTarget('ghost')).resolves.toBeUndefined(); + }); + }); +}); + +describe('mcpServerConfigsEqual', () => { + it('ignores key order and undefined fields', () => { + expect( + mcpServerConfigsEqual( + { transport: 'stdio', command: 'a', args: ['x'], enabled: true }, + { command: 'a', transport: 'stdio', args: ['x'], enabled: true, cwd: undefined }, + ), + ).toBe(true); + expect( + mcpServerConfigsEqual( + { transport: 'stdio', command: 'a', env: { A: '1', B: '2' } }, + { transport: 'stdio', command: 'a', env: { B: '2', A: '1' } }, + ), + ).toBe(true); + }); + + it('distinguishes structural differences', () => { + expect( + mcpServerConfigsEqual( + { transport: 'stdio', command: 'a' }, + { transport: 'stdio', command: 'a', args: [] }, + ), + ).toBe(false); + expect( + mcpServerConfigsEqual( + { transport: 'stdio', command: 'a' }, + { transport: 'http', url: 'https://example.com/mcp' }, + ), + ).toBe(false); + }); +}); diff --git a/packages/agent-core-v2/test/app/plugin/manager-consumption.test.ts b/packages/agent-core-v2/test/app/plugin/manager-consumption.test.ts index bf848e156d..c7de88222d 100644 --- a/packages/agent-core-v2/test/app/plugin/manager-consumption.test.ts +++ b/packages/agent-core-v2/test/app/plugin/manager-consumption.test.ts @@ -1,6 +1,6 @@ import { execFileSync } from 'node:child_process'; -import { createServer } from 'node:http'; import { mkdir, mkdtemp, readdir, readFile, realpath, rm, stat, writeFile } from 'node:fs/promises'; +import { createServer } from 'node:http'; import { tmpdir } from 'node:os'; import path from 'node:path'; @@ -97,7 +97,10 @@ async function makePlugin( } async function zipDir(sourceRoot: string): Promise { - const zipPath = path.join(tmpdir(), `plugin-${Date.now()}-${Math.random().toString(36).slice(2)}.zip`); + const zipPath = path.join( + tmpdir(), + `plugin-${Date.now()}-${Math.random().toString(36).slice(2)}.zip`, + ); execFileSync('zip', ['-qr', zipPath, '.'], { cwd: sourceRoot }); const buffer = await readFile(zipPath); await rm(zipPath, { force: true }); @@ -127,8 +130,7 @@ function mockGithubFetch(options: MockGithubFetchOptions): void { vi.stubGlobal( 'fetch', vi.fn(async (input: Parameters[0], init?: RequestInit) => { - const url = - typeof input === 'string' ? input : input instanceof URL ? input.href : input.url; + const url = typeof input === 'string' ? input : input instanceof URL ? input.href : input.url; if (/^https:\/\/github\.com\/[^/]+\/[^/]+\/releases\/latest$/.test(url)) { options.onReleaseLookup?.(); if (options.releaseTag === undefined) { @@ -579,6 +581,92 @@ describe('PluginManager consumption plane', () => { expect(manager.enabledMcpServers()).toEqual({}); }); + it('mcpServerEntries() lists disabled plugins and disabled servers with provenance', async () => { + const home = await makeKimiHome(); + const demo = await makePlugin('demo', { + mcpServers: { + finance: { command: 'finance-mcp' }, + docs: { url: 'https://example.com/mcp' }, + }, + }); + const other = await makePlugin('other', { + mcpServers: { data: { command: 'data-mcp' } }, + }); + const manager = new PluginManager({ kimiHomeDir: home }); + await manager.load(); + await manager.install(demo); + await manager.install(other); + await manager.setMcpServerEnabled('demo', 'finance', false); + await manager.setEnabled('other', false); + + const entries = manager.mcpServerEntries(); + expect(entries).toHaveLength(3); + const finance = entries.find((entry) => entry.name === 'plugin-demo:finance'); + expect(finance).toEqual(expect.objectContaining({ pluginId: 'demo', serverName: 'finance' })); + expect(finance?.config.enabled).toBe(false); + const docs = entries.find((entry) => entry.name === 'plugin-demo:docs'); + expect(docs).toEqual(expect.objectContaining({ pluginId: 'demo', serverName: 'docs' })); + expect(docs?.config.enabled).toBe(true); + const data = entries.find((entry) => entry.name === 'plugin-other:data'); + expect(data).toEqual(expect.objectContaining({ pluginId: 'other', serverName: 'data' })); + expect(data?.config.enabled).toBe(false); + }); + + it('mcpServerEntries() applies the stdio runtime transforms to every entry', async () => { + const home = await makeKimiHome(); + const root = await makePlugin('demo', { + mcpServers: { + finance: { command: 'finance-mcp', env: { CUSTOM: '1' } }, + docs: { url: 'https://example.com/mcp' }, + }, + }); + const manager = new PluginManager({ kimiHomeDir: home }); + await manager.load(); + await manager.install(root); + const managedRoot = await managedPluginRoot(manager, 'demo'); + + const entries = manager.mcpServerEntries(); + const finance = entries.find((entry) => entry.name === 'plugin-demo:finance'); + expect(finance?.config).toEqual( + expect.objectContaining({ + command: 'finance-mcp', + cwd: managedRoot, + env: expect.objectContaining({ + CUSTOM: '1', + KIMI_CODE_HOME: home, + KIMI_PLUGIN_ROOT: managedRoot, + }), + }), + ); + const docs = entries.find((entry) => entry.name === 'plugin-demo:docs'); + expect(docs?.config).toEqual( + expect.objectContaining({ + transport: 'http', + url: 'https://example.com/mcp', + enabled: true, + }), + ); + expect(JSON.stringify(docs?.config)).not.toContain('KIMI_PLUGIN_ROOT'); + }); + + it('mcpServerEntries() skips plugins in error state', async () => { + const home = await makeKimiHome(); + const root = await makePlugin('demo', { + mcpServers: { finance: { command: 'finance-mcp' } }, + }); + const manager = new PluginManager({ kimiHomeDir: home }); + await manager.load(); + await manager.install(root); + await writeFile( + path.join(await managedPluginRoot(manager, 'demo'), 'kimi.plugin.json'), + '{ not json', + 'utf8', + ); + await manager.reload(); + expect(manager.get('demo')?.state).toBe('error'); + expect(manager.mcpServerEntries()).toEqual([]); + }); + it('setMcpServerEnabled() rejects unknown MCP servers', async () => { const home = await makeKimiHome(); const root = await makePlugin('demo'); @@ -753,10 +841,9 @@ describe('PluginManager consumption plane', () => { const root = await makePlugin('rando', { version: '1.0.0' }); const manager = new PluginManager({ kimiHomeDir: home }); await manager.load(); - const record = await (manager.install as (source: string, options?: unknown) => Promise)( - root, - { marketplace: { id: 'rando', tier: 'official' } }, - ); + const record = await ( + manager.install as (source: string, options?: unknown) => Promise + )(root, { marketplace: { id: 'rando', tier: 'official' } }); expect((record as { marketplace?: unknown }).marketplace).toBeUndefined(); }); diff --git a/packages/agent-core-v2/test/app/plugin/pluginService.test.ts b/packages/agent-core-v2/test/app/plugin/pluginService.test.ts index 6f3d4d9af5..85fd0569dc 100644 --- a/packages/agent-core-v2/test/app/plugin/pluginService.test.ts +++ b/packages/agent-core-v2/test/app/plugin/pluginService.test.ts @@ -4,7 +4,7 @@ import path from 'node:path'; import { KIMI_CODE_PROVIDER_NAME } from '@moonshot-ai/kimi-code-oauth'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; -import { LifecycleScope } from '#/app/scopes'; + import { ScopeActivation, _clearScopedRegistryForTests, @@ -14,11 +14,12 @@ import { createScopedTestHost, stubPair, type ScopedTestHost } from '#/_base/di/ import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { IPluginService } from '#/app/plugin/plugin'; import { PluginService } from '#/app/plugin/pluginService'; -import { IProviderService, type ProviderConfig } from '#/kosong/provider/provider'; -import { ISkillDiscovery } from '#/app/skillCatalog/skillDiscovery'; import * as pluginStore from '#/app/plugin/store'; import type { InstalledFile } from '#/app/plugin/store'; import type { PluginMutationSummary, ReloadSummary } from '#/app/plugin/types'; +import { LifecycleScope } from '#/app/scopes'; +import { ISkillDiscovery } from '#/app/skillCatalog/skillDiscovery'; +import { IProviderService, type ProviderConfig } from '#/kosong/provider/provider'; import { stubBootstrap } from '../bootstrap/stubs'; import { stubProviderService } from '../provider/stubs'; @@ -116,10 +117,7 @@ function deferred(): { return { promise, resolve }; } -async function makePluginDir( - name: string, - manifest: Record, -): Promise { +async function makePluginDir(name: string, manifest: Record): Promise { const root = await mkdtemp(path.join(tmpdir(), `plugin-${name}-`)); await writeFile( path.join(root, 'kimi.plugin.json'), @@ -182,6 +180,8 @@ describe('PluginService (plugin boundary)', () => { try { const svc = host.app.accessor.get(IPluginService); await expect(svc.enabledMcpServers()).resolves.toEqual({}); + const failure = await svc.mcpServerEntries().catch((error: unknown) => error); + expect(failure).toMatchObject({ code: 'plugin.load_failed' }); } finally { host.dispose(); } @@ -233,7 +233,7 @@ describe('PluginService (plugin boundary)', () => { createdDirs.push(pluginRoot); await writeInstalledFile(home, JSON.stringify(installedFile('recovery-demo', pluginRoot))); const reloads: ReloadSummary[] = []; - svc.onDidReload((summary) => reloads.push(summary)); + svc.onDidReload(({ added, removed, errors }) => reloads.push({ added, removed, errors })); await expect(svc.reloadPlugins()).resolves.toEqual({ added: ['recovery-demo'], @@ -302,6 +302,73 @@ describe('PluginService (plugin boundary)', () => { } }); + it('resolves a mutation only after reload listeners settle their waitUntil work', async () => { + const home = await makeHome(); + const pluginRoot = await makePluginDir('barrier-demo', {}); + createdDirs.push(pluginRoot); + await writeInstalledFile(home, JSON.stringify(installedFile('barrier-demo', pluginRoot))); + const host = makeHost(home); + try { + const svc = host.app.accessor.get(IPluginService); + await expect(svc.listPlugins()).resolves.toHaveLength(1); + + const listenerCalled = deferred(); + const reconcileGate = deferred(); + let reconciled = false; + svc.onDidReload((event) => { + listenerCalled.resolve(undefined); + event.waitUntil( + reconcileGate.promise.then(() => { + reconciled = true; + }), + ); + }); + + const mutation = svc.setPluginEnabled({ id: 'barrier-demo', enabled: false }); + let mutationSettled = false; + void mutation.then(() => { + mutationSettled = true; + }); + await listenerCalled.promise; + for (let i = 0; i < 5; i++) await Promise.resolve(); + expect(mutationSettled).toBe(false); + + reconcileGate.resolve(undefined); + await mutation; + expect(reconciled).toBe(true); + await expect(svc.listPlugins()).resolves.toEqual([ + expect.objectContaining({ id: 'barrier-demo', enabled: false }), + ]); + } finally { + host.dispose(); + } + }); + + it('does not reject a mutation when a reload listener waitUntil promise rejects', async () => { + const home = await makeHome(); + const pluginRoot = await makePluginDir('tolerant-demo', {}); + createdDirs.push(pluginRoot); + await writeInstalledFile(home, JSON.stringify(installedFile('tolerant-demo', pluginRoot))); + const host = makeHost(home); + try { + const svc = host.app.accessor.get(IPluginService); + await expect(svc.listPlugins()).resolves.toHaveLength(1); + vi.spyOn(console, 'error').mockImplementation(() => {}); + svc.onDidReload((event) => { + event.waitUntil(Promise.reject(new Error('workspace reconcile failed'))); + }); + + await expect( + svc.setPluginEnabled({ id: 'tolerant-demo', enabled: false }), + ).resolves.toBeUndefined(); + await expect(svc.listPlugins()).resolves.toEqual([ + expect.objectContaining({ id: 'tolerant-demo', enabled: false }), + ]); + } finally { + host.dispose(); + } + }); + it('serves enabled plugin system-prompt sections on the consumption plane', async () => { const home = await makeHome(); const pluginRoot = await makePluginDir('prompt-demo', { systemPrompt: 'Always cite sources.' }); @@ -352,10 +419,7 @@ describe('PluginService (plugin boundary)', () => { const host = makeHost(home); try { const svc = host.app.accessor.get(IPluginService); - const [plugins, roots] = await Promise.all([ - svc.listPlugins(), - svc.pluginSkillRoots(), - ]); + const [plugins, roots] = await Promise.all([svc.listPlugins(), svc.pluginSkillRoots()]); expect(plugins).toEqual([expect.objectContaining({ id: 'snapshot-demo' })]); expect(roots).toEqual([ @@ -422,9 +486,9 @@ describe('PluginService (plugin boundary)', () => { await expect(svc.getPluginInfo({ id: 'demo' })).resolves.toEqual( expect.objectContaining({ root: previous.root, version: '1.0.0' }), ); - await expect(readFile(path.join(previous.root, 'kimi.plugin.json'), 'utf8')).resolves.toContain( - '"version":"1.0.0"', - ); + await expect( + readFile(path.join(previous.root, 'kimi.plugin.json'), 'utf8'), + ).resolves.toContain('"version":"1.0.0"'); await expect(readdir(path.join(home, 'plugins', 'managed'))).resolves.toEqual(['demo']); } finally { host.dispose(); @@ -491,7 +555,7 @@ describe('PluginService (plugin boundary)', () => { try { const svc = host.app.accessor.get(IPluginService); const reloads: ReloadSummary[] = []; - svc.onDidReload((summary) => reloads.push(summary)); + svc.onDidReload(({ added, removed, errors }) => reloads.push({ added, removed, errors })); const firstList = svc.listPlugins(); await firstReadStarted.promise; @@ -597,6 +661,55 @@ describe('PluginService (plugin boundary)', () => { } }); + it('merges the managed Kimi endpoint env into stdio MCP server entries with provenance', async () => { + const home = await makeHome(); + await writeValidInstalledFile(home); + const host = makeHost( + home, + stubProviderService({ + [KIMI_CODE_PROVIDER_NAME]: { + baseUrl: 'https://api.example.test/', + oauth: { storage: 'file', key: 'kimi', oauthHost: 'https://auth.example.test' }, + }, + }), + ); + try { + const svc = host.app.accessor.get(IPluginService); + const pluginRoot = await makePluginDir('demo', { + mcpServers: { + finance: { command: 'finance-mcp', env: { CUSTOM: '1' } }, + docs: { url: 'https://example.test/mcp' }, + }, + }); + createdDirs.push(pluginRoot); + await svc.installPlugin({ source: pluginRoot }); + await svc.setPluginMcpServerEnabled({ id: 'demo', server: 'finance', enabled: false }); + + const entries = await svc.mcpServerEntries(); + const managedRoot = await realpath(path.join(home, 'plugins', 'managed', 'demo')); + const finance = entries.find((entry) => entry.name === 'plugin-demo:finance'); + expect(finance).toEqual(expect.objectContaining({ pluginId: 'demo', serverName: 'finance' })); + expect(finance?.config).toEqual( + expect.objectContaining({ + enabled: false, + env: expect.objectContaining({ + KIMI_CODE_BASE_URL: 'https://api.example.test/', + KIMI_CODE_OAUTH_HOST: 'https://auth.example.test', + CUSTOM: '1', + KIMI_CODE_HOME: home, + KIMI_PLUGIN_ROOT: managedRoot, + }), + }), + ); + const docs = entries.find((entry) => entry.name === 'plugin-demo:docs'); + expect(docs).toEqual(expect.objectContaining({ pluginId: 'demo', serverName: 'docs' })); + expect(docs?.config.enabled).toBe(true); + expect(JSON.stringify(docs?.config)).not.toContain('KIMI_CODE_BASE_URL'); + } finally { + host.dispose(); + } + }); + it('waits for provider config before injecting persisted managed endpoints', async () => { const home = await makeHome(); await writeValidInstalledFile(home); diff --git a/packages/agent-core-v2/test/app/plugin/stubs.ts b/packages/agent-core-v2/test/app/plugin/stubs.ts index 01b77d9b4f..17a9a3d51a 100644 --- a/packages/agent-core-v2/test/app/plugin/stubs.ts +++ b/packages/agent-core-v2/test/app/plugin/stubs.ts @@ -3,12 +3,13 @@ import type { IPluginService } from '#/app/plugin/plugin'; import type { EnabledPluginSessionStart, PluginMutationSummary, + PluginReloadEvent, ReloadSummary, } from '#/app/plugin/types'; interface StubPluginServiceOptions { readonly sessionStarts: readonly EnabledPluginSessionStart[]; - readonly reloadEmitter?: Emitter; + readonly reloadEmitter?: Emitter; readonly mutateEmitter?: Emitter; } @@ -33,6 +34,7 @@ export function stubPluginService(options: StubPluginServiceOptions): IPluginSer enabledSessionStarts: async () => options.sessionStarts, enabledSystemPrompts: async () => [], enabledMcpServers: async () => ({}), + mcpServerEntries: async () => [], enabledHooks: async () => [], hasLoadedSnapshot: () => true, }; diff --git a/packages/agent-core-v2/test/mcpCore/oauth/callback-server.test.ts b/packages/agent-core-v2/test/mcpCore/oauth/callback-server.test.ts new file mode 100644 index 0000000000..5ad70c11f1 --- /dev/null +++ b/packages/agent-core-v2/test/mcpCore/oauth/callback-server.test.ts @@ -0,0 +1,134 @@ +import { createServer as createHttpServer, type Server as HttpServer } from 'node:http'; +import type { AddressInfo as HttpAddress } from 'node:net'; + +import { afterEach, describe, expect, it } from 'vitest'; + +import { + type CallbackServer, + OAuthCallbackClosedError, + startCallbackServer, +} from '#/mcpCore/oauth/callback-server'; +import { McpOAuthService } from '#/mcpCore/oauth/service'; + +import { createMemoryMcpOAuthStore } from '../stubs'; + +const cleanups: Array<() => Promise | void> = []; +afterEach(async () => { + while (cleanups.length > 0) { + await cleanups.pop()?.(); + } +}); + +function trackServer(server: CallbackServer): void { + cleanups.push(() => server.close()); +} + +async function startRegistrationServer(): Promise<{ readonly url: string }> { + const httpServer: HttpServer = createHttpServer((req, res) => { + if (req.method !== 'POST' || req.url !== '/register') { + res.writeHead(404).end(); + return; + } + let body = ''; + req.on('data', (chunk: Buffer) => { + body += chunk.toString('utf-8'); + }); + req.on('end', () => { + const metadata = JSON.parse(body) as Record; + res.writeHead(201, { 'content-type': 'application/json' }); + res.end(JSON.stringify({ ...metadata, client_id: 'test-client' })); + }); + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + cleanups.push( + () => + new Promise((resolve, reject) => { + httpServer.close((err) => { + if (err) { + reject(err); + return; + } + resolve(); + }); + }), + ); + const port = (httpServer.address() as HttpAddress).port; + return { url: `http://127.0.0.1:${port}` }; +} + +describe('startCallbackServer', () => { + it('resolves a late waitForCode with a callback that arrived before it', async () => { + const server = await startCallbackServer(); + trackServer(server); + + const response = await fetch(`${server.redirectUri}?code=early-code&state=early-state`); + expect(response.status).toBe(200); + + await expect(server.waitForCode({ timeoutMs: 10_000 })).resolves.toEqual({ + code: 'early-code', + state: 'early-state', + }); + }); + + it('delivers the callback payload to a pending wait', async () => { + const server = await startCallbackServer(); + trackServer(server); + const pending = server.waitForCode({ timeoutMs: 10_000 }); + + await fetch(`${server.redirectUri}?code=code-1&state=state-1`); + + await expect(pending).resolves.toEqual({ code: 'code-1', state: 'state-1' }); + }); + + it('rejects a pending wait with a closed error when explicitly closed', async () => { + const server = await startCallbackServer(); + trackServer(server); + const pending = server.waitForCode({ timeoutMs: 10_000 }); + const rejection = expect(pending).rejects.toBeInstanceOf(OAuthCallbackClosedError); + + await server.close(); + + await rejection; + }); + + it('rejects a late waitForCode after close', async () => { + const server = await startCallbackServer(); + trackServer(server); + + await server.close(); + + await expect(server.waitForCode({ timeoutMs: 10_000 })).rejects.toBeInstanceOf( + OAuthCallbackClosedError, + ); + }); +}); + +describe('McpOAuthService cancellation', () => { + it('rejects an in-flight completion when the authorization flow is cancelled', async () => { + const service = new McpOAuthService({ store: createMemoryMcpOAuthStore() }); + cleanups.push(() => service.dispose()); + const registrationServer = await startRegistrationServer(); + const provider = service.getProvider('example', 'https://mcp.example.test/rpc'); + await provider.ready; + await provider.saveDiscoveryState({ + authorizationServerUrl: registrationServer.url, + authorizationServerMetadata: { + issuer: registrationServer.url, + authorization_endpoint: `${registrationServer.url}/authorize`, + token_endpoint: `${registrationServer.url}/token`, + registration_endpoint: `${registrationServer.url}/register`, + response_types_supported: ['code'], + grant_types_supported: ['authorization_code'], + token_endpoint_auth_methods_supported: ['none'], + }, + }); + + const flow = await service.beginAuthorization('example', 'https://mcp.example.test/rpc'); + const completion = flow.complete({ timeoutMs: 10_000 }); + const rejection = expect(completion).rejects.toThrow('OAuth callback listener closed'); + + await flow.cancel(); + + await rejection; + }, 15000); +}); diff --git a/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts b/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts new file mode 100644 index 0000000000..313f93327a --- /dev/null +++ b/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts @@ -0,0 +1,1402 @@ +import { createServer as createHttpServer, type Server as HttpServer } from 'node:http'; +import type { AddressInfo as HttpAddress } from 'node:net'; + +import type { OAuthClientInformationFull } from '@modelcontextprotocol/sdk/shared/auth.js'; +import { afterEach, describe, expect, it, vi } from 'vitest'; + +import type { ILogger as Logger } from '#/_base/log/log'; +import * as callbackServerModule from '#/mcpCore/oauth/callback-server'; +import { + META_SUFFIX, + type McpOAuthClientProvider, + type McpOAuthStoreMeta, +} from '#/mcpCore/oauth/provider'; +import { + AlreadyAuthorizedError, + McpOAuthService, + type BeginAuthorizationResult, + type McpOAuthEvent, +} from '#/mcpCore/oauth/service'; +import { mcpOAuthStoreKey, type McpOAuthStore } from '#/mcpCore/oauth/store'; + +import { createMemoryMcpOAuthStore, ManualMcpOAuthScheduler } from '../stubs'; + +const SERVER_NAME = 'notion'; +const SERVER_URL = 'https://mcp.example.test/mcp'; + +interface Fixture { + readonly service: McpOAuthService; + readonly store: McpOAuthStore; + readonly events: McpOAuthEvent[]; + readonly scheduler: ManualMcpOAuthScheduler; +} + +function makeFixture( + store: McpOAuthStore = createMemoryMcpOAuthStore(), + options: { + readonly authRequestTimeoutMs?: number; + readonly shutdownDrainTimeoutMs?: number; + readonly log?: Logger; + } = {}, +): Fixture { + const events: McpOAuthEvent[] = []; + const scheduler = new ManualMcpOAuthScheduler(1_000_000); + const service = new McpOAuthService({ store, scheduler, ...options }); + service.onEvent((event) => events.push(event)); + return { service, store, events, scheduler }; +} + +const cleanups: Array<() => Promise | void> = []; +afterEach(async () => { + while (cleanups.length > 0) { + await cleanups.pop()?.(); + } +}); + +async function listMetaKeys(store: McpOAuthStore): Promise { + return (await store.list()).filter((key) => key.endsWith(META_SUFFIX)); +} + +async function readyProvider(fixture: Fixture): Promise { + const provider = fixture.service.getProvider(SERVER_NAME, SERVER_URL); + await provider.ready; + return provider; +} + +interface FakeAuthServer { + readonly url: string; + readonly counts: { register: number; exchange: number; refresh: number }; +} + +async function startFakeAuthServer( + options: { readonly rejectRefreshToken?: boolean; readonly refreshExpiresIn?: number } = {}, +): Promise { + const counts = { register: 0, exchange: 0, refresh: 0 }; + const httpServer: HttpServer = createHttpServer((req, res) => { + if (req.method !== 'POST' || (req.url !== '/token' && req.url !== '/register')) { + res.writeHead(404).end(); + return; + } + let body = ''; + req.on('data', (chunk: Buffer) => { + body += chunk.toString('utf-8'); + }); + req.on('end', () => { + if (req.url === '/register') { + counts.register += 1; + const metadata = JSON.parse(body) as Record; + res.writeHead(201, { 'content-type': 'application/json' }); + res.end(JSON.stringify({ ...metadata, client_id: `test-client-${counts.register}` })); + return; + } + const grantType = new URLSearchParams(body).get('grant_type'); + if (grantType === 'authorization_code') counts.exchange += 1; + if (grantType === 'refresh_token') { + counts.refresh += 1; + if (options.rejectRefreshToken === true) { + res.writeHead(400, { 'content-type': 'application/json' }); + res.end(JSON.stringify({ error: 'invalid_grant' })); + return; + } + } + res.writeHead(200, { 'content-type': 'application/json' }); + res.end( + JSON.stringify({ + access_token: 'fresh-token', + token_type: 'Bearer', + expires_in: options.refreshExpiresIn ?? 3600, + }), + ); + }); + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + cleanups.push( + () => + new Promise((resolve, reject) => { + httpServer.close((err) => { + if (err) { + reject(err); + return; + } + resolve(); + }); + }), + ); + const port = (httpServer.address() as HttpAddress).port; + return { url: `http://127.0.0.1:${port}`, counts }; +} + +async function startHangingServer(): Promise<{ readonly url: string; readonly counts: { requests: number } }> { + const counts = { requests: 0 }; + const httpServer: HttpServer = createHttpServer(() => { + counts.requests += 1; + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + cleanups.push( + () => + new Promise((resolve, reject) => { + httpServer.close((err) => { + if (err) { + reject(err); + return; + } + resolve(); + }); + }), + ); + const port = (httpServer.address() as HttpAddress).port; + return { url: `http://127.0.0.1:${port}`, counts }; +} + +interface GatedExchangeAuthServer { + readonly url: string; + readonly counts: { register: number; exchange: number }; + readonly exchangeStarted: Promise; + readonly releaseExchange: () => void; +} + +async function startGatedExchangeAuthServer(): Promise { + const counts = { register: 0, exchange: 0 }; + let signalExchangeStarted: () => void = () => undefined; + const exchangeStarted = new Promise((resolve) => { + signalExchangeStarted = resolve; + }); + let releaseExchange: () => void = () => undefined; + const exchangeReleased = new Promise((resolve) => { + releaseExchange = resolve; + }); + const httpServer: HttpServer = createHttpServer((req, res) => { + if (req.method !== 'POST' || (req.url !== '/token' && req.url !== '/register')) { + res.writeHead(404).end(); + return; + } + let body = ''; + req.on('data', (chunk: Buffer) => { + body += chunk.toString('utf-8'); + }); + req.on('end', () => { + if (req.url === '/register') { + counts.register += 1; + const metadata = JSON.parse(body) as Record; + res.writeHead(201, { 'content-type': 'application/json' }); + res.end(JSON.stringify({ ...metadata, client_id: `test-client-${counts.register}` })); + return; + } + if (new URLSearchParams(body).get('grant_type') !== 'authorization_code') { + res.writeHead(400, { 'content-type': 'application/json' }); + res.end(JSON.stringify({ error: 'unsupported_grant_type' })); + return; + } + counts.exchange += 1; + signalExchangeStarted(); + res.on('error', () => undefined); + void exchangeReleased.then(() => { + res.writeHead(200, { 'content-type': 'application/json' }); + res.end( + JSON.stringify({ access_token: 'fresh-token', token_type: 'Bearer', expires_in: 3600 }), + ); + }); + }); + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + cleanups.push( + () => + new Promise((resolve, reject) => { + httpServer.close((err) => { + if (err) { + reject(err); + return; + } + resolve(); + }); + }), + ); + const port = (httpServer.address() as HttpAddress).port; + return { url: `http://127.0.0.1:${port}`, counts, exchangeStarted, releaseExchange }; +} + +function authServerState(authServerUrl: string) { + return { + discovery: { + authorizationServerUrl: authServerUrl, + authorizationServerMetadata: { + issuer: authServerUrl, + authorization_endpoint: `${authServerUrl}/authorize`, + token_endpoint: `${authServerUrl}/token`, + registration_endpoint: `${authServerUrl}/register`, + response_types_supported: ['code'], + grant_types_supported: ['authorization_code', 'refresh_token'], + token_endpoint_auth_methods_supported: ['none'], + }, + }, + client: { + client_id: 'cached-client', + redirect_uris: ['http://127.0.0.1:45678/callback'], + token_endpoint_auth_method: 'none', + grant_types: ['authorization_code', 'refresh_token'], + response_types: ['code'], + } satisfies OAuthClientInformationFull, + }; +} + +async function blockedRefreshFixture(options?: { + readonly shutdownDrainTimeoutMs?: number; +}): Promise<{ + readonly fixture: Fixture; + readonly authServer: FakeAuthServer; + readonly writeStarted: Promise; + readonly releaseWrite: () => void; + readonly armMetaGate: () => void; + readonly metaWriteStarted: Promise; + readonly releaseMetaWrite: () => void; +}> { + const memory = createMemoryMcpOAuthStore(); + let signalWriteStarted: () => void = () => undefined; + const writeStarted = new Promise((resolve) => { + signalWriteStarted = resolve; + }); + let releaseWrite: () => void = () => undefined; + const writeReleased = new Promise((resolve) => { + releaseWrite = resolve; + }); + let gateMeta = false; + let signalMetaStarted: () => void = () => undefined; + const metaWriteStarted = new Promise((resolve) => { + signalMetaStarted = resolve; + }); + let releaseMetaWrite: () => void = () => undefined; + const metaReleased = new Promise((resolve) => { + releaseMetaWrite = resolve; + }); + const store: McpOAuthStore = { + ...memory, + async write(key: string, value: unknown): Promise { + const accessToken = + typeof value === 'object' && value !== null + ? (value as { readonly access_token?: unknown }).access_token + : undefined; + if (accessToken === 'fresh-token') { + signalWriteStarted(); + await writeReleased; + } + if (gateMeta && key.endsWith(META_SUFFIX)) { + signalMetaStarted(); + await metaReleased; + } + await memory.write(key, value); + }, + }; + const fixture = makeFixture(store, options ?? {}); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer({ refreshExpiresIn: 60 }); + const provider = await readyProvider(fixture); + const state = authServerState(authServer.url); + await provider.saveDiscoveryState(state.discovery); + await provider.saveClientInformation(state.client); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + }); + return { + fixture, + authServer, + writeStarted, + releaseWrite, + armMetaGate: () => { + gateMeta = true; + }, + metaWriteStarted, + releaseMetaWrite, + }; +} + +async function metaGatedGrantFixture(): Promise<{ + readonly fixture: Fixture; + readonly provider: McpOAuthClientProvider; + readonly tokenUrl: string; + readonly armMetaGate: () => void; + readonly metaWriteStarted: Promise; + readonly releaseMetaWrite: () => void; +}> { + const memory = createMemoryMcpOAuthStore(); + let gateMeta = false; + let signalMetaStarted: () => void = () => undefined; + const metaWriteStarted = new Promise((resolve) => { + signalMetaStarted = resolve; + }); + let releaseMetaWrite: () => void = () => undefined; + const metaReleased = new Promise((resolve) => { + releaseMetaWrite = resolve; + }); + const store: McpOAuthStore = { + ...memory, + async write(key: string, value: unknown): Promise { + if (gateMeta && key.endsWith(META_SUFFIX)) { + signalMetaStarted(); + await metaReleased; + } + await memory.write(key, value); + }, + }; + const fixture = makeFixture(store); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + const provider = await readyProvider(fixture); + const state = authServerState(authServer.url); + await provider.saveDiscoveryState(state.discovery); + await provider.saveClientInformation(state.client); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + }); + return { + fixture, + provider, + tokenUrl: `${authServer.url}/token`, + armMetaGate: () => { + gateMeta = true; + }, + metaWriteStarted, + releaseMetaWrite, + }; +} + +async function deliverCallback(flow: BeginAuthorizationResult): Promise { + const redirectUri = flow.authorizationUrl.searchParams.get('redirect_uri'); + const state = flow.authorizationUrl.searchParams.get('state'); + expect(redirectUri).toBeTruthy(); + const callbackUrl = new URL(redirectUri!); + callbackUrl.searchParams.set('code', 'test-auth-code'); + if (state !== null) callbackUrl.searchParams.set('state', state); + const response = await fetch(callbackUrl); + expect(response.status).toBe(200); + await response.text(); +} + +describe('McpOAuthService credential bookkeeping', () => { + it('stamps token writes with obtained_at and a name/url meta record', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + await fixture.service + .getProvider(SERVER_NAME, SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer', expires_in: 3600 }); + + const state = await fixture.service.tokenState(SERVER_NAME, SERVER_URL); + expect(state.hasTokens).toBe(true); + expect(state.expired).toBe(false); + expect(state.expiresAt).toBe(4_600_000); + + const metaFiles = await listMetaKeys(fixture.store); + expect(metaFiles).toHaveLength(1); + expect(await fixture.store.read(metaFiles[0]!)).toEqual({ + serverName: SERVER_NAME, + serverUrl: SERVER_URL, + }); + + expect(fixture.events).toEqual([ + { type: 'tokens-saved', serverName: SERVER_NAME, serverUrl: SERVER_URL }, + ]); + }); + + it('treats tokens without expiry data as non-expiring', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + expect(await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).toEqual({ + hasTokens: false, + hasRefreshToken: false, + expired: false, + }); + + await fixture.service + .getProvider(SERVER_NAME, SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer' }); + expect(await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).toMatchObject({ + hasTokens: true, + expired: false, + expiresAt: undefined, + }); + }); + + it('treats a grant saved with a negative expires_in as expired', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + await fixture.service.getProvider(SERVER_NAME, SERVER_URL).saveTokens({ + access_token: 'a', + token_type: 'Bearer', + expires_in: -60, + }); + expect(await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).toMatchObject({ + hasTokens: true, + hasRefreshToken: false, + expired: true, + }); + }); + + it('emits tokens-invalidated and drops the meta record when credentials are cleared', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + await fixture.service + .getProvider(SERVER_NAME, SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer' }); + expect(await listMetaKeys(fixture.store)).toHaveLength(1); + + await fixture.service.invalidate(SERVER_NAME, SERVER_URL, 'tokens'); + expect((await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).hasTokens).toBe(false); + expect(await listMetaKeys(fixture.store)).toHaveLength(0); + expect(fixture.events).toContainEqual({ + type: 'tokens-invalidated', + serverName: SERVER_NAME, + serverUrl: SERVER_URL, + scope: 'tokens', + }); + }); + + it('keeps the meta sidecar consistent with the final tokens when save and clear race', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const provider = await readyProvider(fixture); + + await Promise.all([ + provider.saveTokens({ access_token: 'a', token_type: 'Bearer' }), + provider.clearCredentials('tokens'), + ]); + expect(await provider.tokens()).toBeUndefined(); + expect(await listMetaKeys(fixture.store)).toHaveLength(0); + + await Promise.all([ + provider.clearCredentials('tokens'), + provider.saveTokens({ access_token: 'b', token_type: 'Bearer' }), + ]); + expect(await provider.tokens()).toMatchObject({ access_token: 'b' }); + expect(await listMetaKeys(fixture.store)).toHaveLength(1); + }); +}); + +describe('McpOAuthService single-flight refresh', () => { + it('shares one in-flight refresh across concurrent callers', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + let tokenRequests = 0; + const httpServer: HttpServer = createHttpServer((req, res) => { + if (req.url === '/token' && req.method === 'POST') { + tokenRequests += 1; + res.writeHead(200, { 'content-type': 'application/json' }); + res.end( + JSON.stringify({ access_token: 'fresh-token', token_type: 'Bearer', expires_in: 3600 }), + ); + return; + } + res.writeHead(404).end(); + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + cleanups.push( + () => + new Promise((resolve, reject) => { + httpServer.close((err) => { + if (err) { + reject(err); + return; + } + resolve(); + }); + }), + ); + const port = (httpServer.address() as HttpAddress).port; + const authServerUrl = `http://127.0.0.1:${port}`; + + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState({ + authorizationServerUrl: authServerUrl, + authorizationServerMetadata: { + issuer: authServerUrl, + authorization_endpoint: `${authServerUrl}/authorize`, + token_endpoint: `${authServerUrl}/token`, + response_types_supported: ['code'], + grant_types_supported: ['authorization_code', 'refresh_token'], + token_endpoint_auth_methods_supported: ['none'], + }, + }); + await provider.saveClientInformation({ + client_id: 'cached-client', + redirect_uris: ['http://127.0.0.1:45678/callback'], + token_endpoint_auth_method: 'none', + grant_types: ['authorization_code', 'refresh_token'], + response_types: ['code'], + } satisfies OAuthClientInformationFull); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + }); + + await Promise.all([ + fixture.service.refresh(SERVER_NAME, SERVER_URL), + fixture.service.refresh(SERVER_NAME, SERVER_URL), + ]); + expect(tokenRequests).toBe(1); + expect(await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).toMatchObject({ + hasTokens: true, + expired: false, + }); + expect(fixture.events.filter((event) => event.type === 'tokens-saved')).toHaveLength(2); + }, 15000); + + it('bounds a hung authorization-server request by the configured request timeout', async () => { + const hanging = await startHangingServer(); + const fixture = makeFixture(createMemoryMcpOAuthStore(), { authRequestTimeoutMs: 50 }); + cleanups.push(() => fixture.service.dispose()); + const provider = fixture.service.getProvider(SERVER_NAME, hanging.url); + await provider.ready; + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + }); + + await expect(fixture.service.refresh(SERVER_NAME, hanging.url)).rejects.toThrow(); + expect(hanging.counts.requests).toBeGreaterThan(0); + }); + + it('rejects when no refresh token is stored', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + await expect(fixture.service.refresh(SERVER_NAME, SERVER_URL)).rejects.toThrow( + /no refreshable OAuth grant/, + ); + }); + + it('routes the token request through the credential-serialized fetch', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + await provider.saveClientInformation({ + client_id: 'cached-client', + redirect_uris: ['http://127.0.0.1:45678/callback'], + token_endpoint_auth_method: 'none', + grant_types: ['authorization_code', 'refresh_token'], + response_types: ['code'], + } satisfies OAuthClientInformationFull); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + }); + + const fetchSpy = vi.spyOn(provider, 'createOAuthFetch'); + await fixture.service.refresh(SERVER_NAME, SERVER_URL); + expect(fetchSpy).toHaveBeenCalled(); + expect(authServer.counts.refresh).toBe(1); + }, 15000); + + it('emits tokens-invalidated when the SDK invalidates a rejected refresh grant', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer({ rejectRefreshToken: true }); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + await provider.saveClientInformation(authServerState(authServer.url).client); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + }); + + await expect(fixture.service.refresh(SERVER_NAME, SERVER_URL)).rejects.toThrow( + /requires an interactive login/, + ); + expect((await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).hasTokens).toBe(false); + expect(fixture.events).toContainEqual({ + type: 'tokens-invalidated', + serverName: SERVER_NAME, + serverUrl: SERVER_URL, + scope: 'tokens', + }); + }, 15000); + + it('does not resurrect tokens cleared between a grant fetch and the SDK save', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + const grant = { + access_token: 'rotated-access', + refresh_token: 'rotated-refresh', + token_type: 'Bearer', + expires_in: 3600, + }; + const httpServer: HttpServer = createHttpServer((req, res) => { + if (req.url === '/token' && req.method === 'POST') { + res.writeHead(200, { 'content-type': 'application/json' }); + res.end(JSON.stringify(grant)); + return; + } + res.writeHead(404).end(); + }); + await new Promise((resolve) => httpServer.listen(0, '127.0.0.1', resolve)); + cleanups.push( + () => + new Promise((resolve, reject) => { + httpServer.close((err) => { + if (err) { + reject(err); + return; + } + resolve(); + }); + }), + ); + const authServerUrl = `http://127.0.0.1:${(httpServer.address() as HttpAddress).port}`; + + const provider = await readyProvider(fixture); + const state = authServerState(authServerUrl); + await provider.saveDiscoveryState(state.discovery); + await provider.saveClientInformation(state.client); + await provider.saveTokens({ + access_token: 'seed-access', + refresh_token: 'seed-refresh', + token_type: 'Bearer', + }); + + const res = await provider.createOAuthFetch()(`${authServerUrl}/token`, { + method: 'POST', + body: new URLSearchParams({ grant_type: 'refresh_token', refresh_token: 'seed-refresh' }), + }); + const granted = (await res.json()) as Parameters[0]; + + await provider.clearCredentials('all'); + expect(await provider.tokens()).toBeUndefined(); + + await provider.saveTokens(granted); + expect(await provider.tokens()).toBeUndefined(); + }, 15000); +}); + +describe('McpOAuthService interactive flow serialization', () => { + it('closes the callback listener when stale registration cleanup fails', async () => { + const memory = createMemoryMcpOAuthStore(); + const store: McpOAuthStore = { + ...memory, + async remove(key: string): Promise { + if (key.endsWith('-client.json')) throw new Error('disk full'); + await memory.remove(key); + }, + }; + const fixture = makeFixture(store); + cleanups.push(() => fixture.service.dispose()); + const provider = await readyProvider(fixture); + await provider.saveClientInformation({ + client_id: 'cached-client', + redirect_uris: ['http://127.0.0.1:45678/callback'], + token_endpoint_auth_method: 'none', + grant_types: ['authorization_code', 'refresh_token'], + response_types: ['code'], + } satisfies OAuthClientInformationFull); + + const callbackServer: callbackServerModule.CallbackServer = { + redirectUri: 'http://127.0.0.1:45679/callback', + waitForCode: vi.fn(), + close: vi.fn(async () => undefined), + }; + const startSpy = vi + .spyOn(callbackServerModule, 'startCallbackServer') + .mockResolvedValue(callbackServer); + cleanups.push(() => startSpy.mockRestore()); + + await expect( + fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL), + ).rejects.toThrow(/failed to start OAuth flow/); + expect(callbackServer.close).toHaveBeenCalledOnce(); + }); + + it('joins a concurrent flow for the same credential instead of resetting PKCE state', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + + const first = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + const second = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL, { + clientLabel: 'other-client', + }); + expect(second.authorizationUrl.toString()).toBe(first.authorizationUrl.toString()); + + const firstComplete = first.complete({ timeoutMs: 10_000 }); + await deliverCallback(first); + await firstComplete; + await second.complete(); + expect(authServer.counts.exchange).toBe(1); + expect((await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).hasTokens).toBe(true); + }, 15000); + + it('skips a refresh that fires while an interactive flow owns the credential', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer({ rejectRefreshToken: true }); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + expires_in: 3600, + }); + + const flow = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + const complete = flow.complete({ timeoutMs: 10_000 }); + await expect(fixture.service.refresh(SERVER_NAME, SERVER_URL)).resolves.toBeUndefined(); + await deliverCallback(flow); + await complete; + expect(authServer.counts.exchange).toBe(1); + expect((await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).hasTokens).toBe(true); + }, 15000); + + it('serializes an interactive flow behind a refresh whose token read is in flight', async () => { + const memory = createMemoryMcpOAuthStore(); + let releaseTokensRead: () => void = () => undefined; + const tokensReadGate = new Promise((resolve) => { + releaseTokensRead = resolve; + }); + let signalReadHeld: () => void = () => undefined; + const tokensReadHeld = new Promise((resolve) => { + signalReadHeld = resolve; + }); + let gateArmed = false; + const store: McpOAuthStore = { + ...memory, + async read(key: string): Promise { + if (gateArmed && key.endsWith('-tokens.json')) { + gateArmed = false; + signalReadHeld(); + await tokensReadGate; + } + return memory.read(key); + }, + }; + const fixture = makeFixture(store); + cleanups.push(() => fixture.service.dispose()); + cleanups.push(() => releaseTokensRead()); + const authServer = await startFakeAuthServer({ rejectRefreshToken: true }); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + await provider.saveClientInformation(authServerState(authServer.url).client); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + expires_in: 3600, + }); + + gateArmed = true; + const refresh = fixture.service.refresh(SERVER_NAME, SERVER_URL); + await tokensReadHeld; + + let began = false; + const begin = fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL).then((flow) => { + began = true; + return flow; + }); + await Promise.resolve(); + expect(began).toBe(false); + + releaseTokensRead(); + await expect(refresh).rejects.toThrow(/requires an interactive login/); + expect(authServer.counts.refresh).toBe(1); + + const flow = await begin; + const complete = flow.complete({ timeoutMs: 10_000 }); + await deliverCallback(flow); + await complete; + expect(authServer.counts.exchange).toBe(1); + expect((await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).hasTokens).toBe(true); + }, 15000); + + it('waits for an in-flight refresh before starting an interactive flow', async () => { + const { fixture, authServer, writeStarted, releaseWrite } = await blockedRefreshFixture(); + cleanups.push(() => { + releaseWrite(); + }); + const provider = await readyProvider(fixture); + const resetFlow = vi.spyOn(provider, 'resetFlow'); + + const refresh = fixture.service.refresh(SERVER_NAME, SERVER_URL); + await writeStarted; + const resetCountDuringRefresh = resetFlow.mock.calls.length; + + const begin = fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + await Promise.resolve(); + expect(resetFlow).toHaveBeenCalledTimes(resetCountDuringRefresh); + + releaseWrite(); + await refresh; + await expect(begin).rejects.toBeInstanceOf(AlreadyAuthorizedError); + expect(resetFlow.mock.calls.length).toBeGreaterThan(resetCountDuringRefresh); + expect(authServer.counts.exchange).toBe(0); + }, 15000); + + it('keeps the shared flow active when a joined handle cancels, so the first handle completes', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + + const first = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + const second = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + + await second.cancel(); + await expect(second.complete()).rejects.toThrow(/already completed or cancelled/); + + const firstComplete = first.complete({ timeoutMs: 10_000 }); + await deliverCallback(first); + await firstComplete; + expect(authServer.counts.exchange).toBe(1); + }, 15000); + + it('keeps the shared flow active when the first handle cancels, so a joined handle completes', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + + const first = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + const second = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + await first.cancel(); + + const secondComplete = second.complete({ timeoutMs: 10_000 }); + await deliverCallback(second); + await secondComplete; + expect(authServer.counts.exchange).toBe(1); + expect((await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).hasTokens).toBe(true); + }, 15000); + + it('closes the shared flow when its final handle cancels, so the next begin starts fresh', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + + const first = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + const second = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + await first.cancel(); + await second.cancel(); + + const third = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + expect(third.authorizationUrl.toString()).not.toBe(first.authorizationUrl.toString()); + await third.cancel(); + }, 15000); + + it('leaves no shared flow behind when begin reports already-authorized', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + }); + const tokensSavedBefore = fixture.events.filter((event) => event.type === 'tokens-saved').length; + + await expect( + fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL), + ).rejects.toBeInstanceOf(AlreadyAuthorizedError); + await expect( + fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL), + ).rejects.toBeInstanceOf(AlreadyAuthorizedError); + expect(authServer.counts.refresh).toBe(2); + expect(fixture.events.filter((event) => event.type === 'tokens-saved')).toHaveLength( + tokensSavedBefore + 2, + ); + }, 15000); + + it('bounds a hanging protected-resource-metadata request during completion', async () => { + const hanging = await startHangingServer(); + const authServer = await startFakeAuthServer(); + const fixture = makeFixture(createMemoryMcpOAuthStore(), { authRequestTimeoutMs: 50 }); + cleanups.push(() => fixture.service.dispose()); + const provider = fixture.service.getProvider(SERVER_NAME, hanging.url); + await provider.ready; + const state = authServerState(authServer.url); + await provider.saveDiscoveryState(state.discovery); + await provider.saveClientInformation(state.client); + + const flow = await fixture.service.beginAuthorization(SERVER_NAME, hanging.url); + const complete = flow.complete({ timeoutMs: 10_000 }); + await deliverCallback(flow); + await complete; + + expect(hanging.counts.requests).toBeGreaterThan(0); + expect(authServer.counts.exchange).toBe(1); + expect((await fixture.service.tokenState(SERVER_NAME, hanging.url)).hasTokens).toBe(true); + }, 15000); + + it('rejects completion without writing tokens when the caller aborts after the callback', async () => { + const gated = await startGatedExchangeAuthServer(); + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + cleanups.push(() => gated.releaseExchange()); + const provider = fixture.service.getProvider(SERVER_NAME, gated.url); + await provider.ready; + const state = authServerState(gated.url); + await provider.saveDiscoveryState(state.discovery); + await provider.saveClientInformation(state.client); + + const flow = await fixture.service.beginAuthorization(SERVER_NAME, gated.url); + const controller = new AbortController(); + const complete = flow.complete({ signal: controller.signal, timeoutMs: 10_000 }); + await deliverCallback(flow); + await gated.exchangeStarted; + + controller.abort(); + await expect(complete).rejects.toThrow(/OAuth flow for "notion" failed/); + expect(gated.counts.exchange).toBe(1); + expect((await fixture.service.tokenState(SERVER_NAME, gated.url)).hasTokens).toBe(false); + }, 15000); +}); + +describe('McpOAuthService sweepProactiveRefresh resilience', () => { + it('skips malformed meta sidecars and still schedules the valid credential', async () => { + const memory = createMemoryMcpOAuthStore(); + const store: McpOAuthStore = { + ...memory, + async read(key: string): Promise { + if (key === 'corrupt-meta.json') return undefined; + return memory.read(key); + }, + }; + const fixture = makeFixture(store); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + + const state = authServerState(authServer.url); + const storeKey = mcpOAuthStoreKey(SERVER_NAME, SERVER_URL); + await fixture.store.write(`${storeKey}-discovery.json`, state.discovery); + await fixture.store.write(`${storeKey}-client.json`, state.client); + await fixture.store.write(`${storeKey}-tokens.json`, { + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + expires_in: 60, + obtained_at: fixture.scheduler.now(), + }); + await fixture.store.write(`${storeKey}-meta.json`, { + serverName: SERVER_NAME, + serverUrl: SERVER_URL, + } satisfies McpOAuthStoreMeta); + + await fixture.store.write('broken-empty-meta.json', {}); + await fixture.store.write('broken-types-meta.json', { serverName: 1, serverUrl: 42 }); + await fixture.store.write('broken-url-meta.json', { serverName: 'x', serverUrl: 'not a url' }); + await fixture.store.write('corrupt-meta.json', '{not json'); + + await expect(fixture.service.sweepProactiveRefresh()).resolves.toBeUndefined(); + await fixture.scheduler.advanceBy(30_000); + expect(authServer.counts.refresh).toBe(1); + }, 15000); +}); + +describe('McpOAuthService proactive refresh scheduling', () => { + it('delays a 60-second grant refresh until its midpoint', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + + const provider = await readyProvider(fixture); + const state = authServerState(authServer.url); + await provider.saveDiscoveryState(state.discovery); + await provider.saveClientInformation(state.client); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + expires_in: 60, + }); + + await fixture.scheduler.advanceBy(29_999); + expect(authServer.counts.refresh).toBe(0); + + await fixture.scheduler.advanceBy(1); + expect(authServer.counts.refresh).toBe(1); + expect(fixture.events.filter((event) => event.type === 'tokens-saved')).toHaveLength(2); + }, 15000); + + it('waits another midpoint after refreshing into another 60-second grant', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer({ refreshExpiresIn: 60 }); + + const provider = await readyProvider(fixture); + const state = authServerState(authServer.url); + await provider.saveDiscoveryState(state.discovery); + await provider.saveClientInformation(state.client); + await provider.saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + expires_in: 60, + }); + + await fixture.scheduler.advanceBy(30_000); + expect(authServer.counts.refresh).toBe(1); + expect((await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).hasRefreshToken).toBe(true); + + await fixture.scheduler.advanceBy(0); + expect(authServer.counts.refresh).toBe(1); + + await fixture.scheduler.advanceBy(30_000); + expect(authServer.counts.refresh).toBe(2); + }, 15000); + + it('re-arms scheduling for expiries beyond the setTimeout limit', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const maxTimerDelayMs = 0x7fffffff; + const refreshSpy = vi + .spyOn(fixture.service, 'refresh') + .mockRejectedValue(new Error('refresh unavailable in test')); + + await fixture.service.getProvider(SERVER_NAME, SERVER_URL).saveTokens({ + access_token: 'a', + refresh_token: 'r', + token_type: 'Bearer', + expires_in: Math.ceil(maxTimerDelayMs / 1000) + 600, + }); + const expiresAt = (await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).expiresAt!; + + await fixture.scheduler.advanceBy(maxTimerDelayMs); + expect(refreshSpy).not.toHaveBeenCalled(); + + await fixture.scheduler.advanceBy(expiresAt - fixture.scheduler.now() - 120_000); + expect(refreshSpy).toHaveBeenCalledTimes(1); + expect(fixture.events).toContainEqual({ + type: 'refresh-failed', + serverName: SERVER_NAME, + serverUrl: SERVER_URL, + error: 'refresh unavailable in test', + }); + }); + + it('does not proactively refresh an already-expired grant', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const refreshSpy = vi.spyOn(fixture.service, 'refresh'); + + await fixture.service.getProvider(SERVER_NAME, SERVER_URL).saveTokens({ + access_token: 'a', + refresh_token: 'r', + token_type: 'Bearer', + expires_in: -60, + }); + + await fixture.scheduler.advanceBy(10_000); + expect(refreshSpy).not.toHaveBeenCalled(); + }); + + it('does not schedule an expiring grant without a refresh token', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const refreshSpy = vi.spyOn(fixture.service, 'refresh'); + + await fixture.service.getProvider(SERVER_NAME, SERVER_URL).saveTokens({ + access_token: 'a', + token_type: 'Bearer', + expires_in: 60, + }); + + await fixture.scheduler.advanceBy(60_000); + expect(refreshSpy).not.toHaveBeenCalled(); + expect(fixture.events.some((event) => event.type === 'refresh-failed')).toBe(false); + }); +}); + +describe('McpOAuthService shutdown', () => { + it('keeps shutdown pending while a token refresh is in flight', async () => { + const { fixture, writeStarted, releaseWrite } = await blockedRefreshFixture(); + const refresh = fixture.service.refresh(SERVER_NAME, SERVER_URL); + await writeStarted; + + const shutdown = fixture.service.shutdown(); + let shutdownSettled = false; + void shutdown.then(() => { + shutdownSettled = true; + }); + await Promise.resolve(); + const pendingWhileRefreshInFlight = !shutdownSettled; + + releaseWrite(); + await Promise.all([refresh, shutdown]); + expect(pendingWhileRefreshInFlight).toBe(true); + }, 15000); + + it('caps the shutdown drain when an in-flight refresh outlives the drain timeout', async () => { + const { fixture, writeStarted, releaseWrite } = await blockedRefreshFixture({ + shutdownDrainTimeoutMs: 50, + }); + const refresh = fixture.service.refresh(SERVER_NAME, SERVER_URL); + await writeStarted; + + const shutdown = fixture.service.shutdown(); + await fixture.scheduler.advanceBy(60); + const settledBeforeRelease = await Promise.race([ + shutdown.then(() => true), + new Promise((resolve) => setTimeout(() => resolve(false), 1_000)), + ]); + releaseWrite(); + await Promise.all([refresh, shutdown]); + + expect(settledBeforeRelease).toBe(true); + }, 15000); + + it('prevents a completing refresh from scheduling work after shutdown', async () => { + const { fixture, authServer, writeStarted, releaseWrite } = await blockedRefreshFixture(); + const refresh = fixture.service.refresh(SERVER_NAME, SERVER_URL); + await writeStarted; + + const shutdown = fixture.service.shutdown(); + releaseWrite(); + await Promise.all([refresh, shutdown]); + await fixture.scheduler.advanceBy(30_000); + + expect(authServer.counts.refresh).toBe(1); + }, 15000); + + it('cancels active flows on shutdown', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const authServer = await startFakeAuthServer(); + const provider = await readyProvider(fixture); + await provider.saveDiscoveryState(authServerState(authServer.url).discovery); + + const flow = await fixture.service.beginAuthorization(SERVER_NAME, SERVER_URL); + + await fixture.service.shutdown(); + + await expect(flow.complete()).rejects.toThrow(/already completed or cancelled/); + }, 15000); + + it('clears event listeners and cached providers on shutdown', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + const providerBefore = fixture.service.getProvider(SERVER_NAME, SERVER_URL); + + await fixture.service.shutdown(); + + const eventCount = fixture.events.length; + await fixture.service + .getProvider(SERVER_NAME, SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer', expires_in: 3600 }); + expect(fixture.events).toHaveLength(eventCount); + + expect(fixture.service.getProvider(SERVER_NAME, SERVER_URL)).not.toBe(providerBefore); + }); + + it('is idempotent across repeated shutdown calls', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + await fixture.service.shutdown(); + await expect(fixture.service.shutdown()).resolves.toBeUndefined(); + }); + + it('clears pending proactive-refresh timers', async () => { + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + await fixture.service.getProvider(SERVER_NAME, SERVER_URL).saveTokens({ + access_token: 'a', + refresh_token: 'r', + token_type: 'Bearer', + expires_in: 3600, + }); + const refreshSpy = vi.spyOn(fixture.service, 'refresh'); + + await fixture.service.shutdown(); + await fixture.scheduler.advanceBy(3600_000); + expect(refreshSpy).not.toHaveBeenCalled(); + }); + + it('waits for an in-flight transport-driven grant before completing shutdown', async () => { + const { + fixture, + authServer, + writeStarted, + releaseWrite, + armMetaGate, + metaWriteStarted, + releaseMetaWrite, + } = await blockedRefreshFixture(); + cleanups.push(() => { + releaseWrite(); + releaseMetaWrite(); + }); + const provider = fixture.service.getProvider(SERVER_NAME, SERVER_URL); + + const grant = provider.createOAuthFetch()(`${authServer.url}/token`, { + method: 'POST', + body: new URLSearchParams({ + grant_type: 'refresh_token', + refresh_token: 'stale-refresh-token', + }), + }); + await writeStarted; + armMetaGate(); + const save = provider.saveTokens({ + access_token: 'fresh-token', + token_type: 'Bearer', + expires_in: 60, + }); + + const shutdown = fixture.service.shutdown(); + let shutdownSettled = false; + void shutdown.then(() => { + shutdownSettled = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + const pendingWhileGrantInFlight = !shutdownSettled; + + releaseWrite(); + const response = await grant; + await metaWriteStarted; + await new Promise((resolve) => setImmediate(resolve)); + const pendingWhileMetaWriteInFlight = !shutdownSettled; + + releaseMetaWrite(); + await save; + await shutdown; + + expect(response.status).toBe(200); + expect(pendingWhileGrantInFlight).toBe(true); + expect(pendingWhileMetaWriteInFlight).toBe(true); + }, 15000); + + it('drains an SDK save continuation that starts after shutdown began', async () => { + const { + fixture, + provider, + tokenUrl, + armMetaGate, + metaWriteStarted, + releaseMetaWrite, + } = await metaGatedGrantFixture(); + cleanups.push(() => { + releaseMetaWrite(); + }); + + const grant = provider.createOAuthFetch()(tokenUrl, { + method: 'POST', + body: new URLSearchParams({ + grant_type: 'refresh_token', + refresh_token: 'stale-refresh-token', + }), + }); + const response = await grant; + armMetaGate(); + + const shutdown = fixture.service.shutdown(); + let shutdownSettled = false; + void shutdown.then(() => { + shutdownSettled = true; + }); + const save = (async () => + provider.saveTokens((await response.json()) as Parameters[0]))(); + + await metaWriteStarted; + await new Promise((resolve) => setImmediate(resolve)); + const pendingWhileMetaWriteInFlight = !shutdownSettled; + + releaseMetaWrite(); + await save; + await shutdown; + + expect(response.status).toBe(200); + expect(pendingWhileMetaWriteInFlight).toBe(true); + expect(await listMetaKeys(fixture.store)).toHaveLength(1); + }, 15000); + + it('aborts a hung begin on shutdown and closes the callback listener', async () => { + const hanging = await startHangingServer(); + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + + const callbackServer: callbackServerModule.CallbackServer = { + redirectUri: 'http://127.0.0.1:45679/callback', + waitForCode: vi.fn(), + close: vi.fn(async () => undefined), + }; + const startSpy = vi + .spyOn(callbackServerModule, 'startCallbackServer') + .mockResolvedValue(callbackServer); + cleanups.push(() => startSpy.mockRestore()); + + const begin = fixture.service.beginAuthorization(SERVER_NAME, hanging.url); + await vi.waitFor(() => { + expect(hanging.counts.requests).toBeGreaterThan(0); + }); + + const beginRejected = expect(begin).rejects.toThrow(/failed to start OAuth flow/); + await fixture.service.shutdown(); + + await beginRejected; + expect(callbackServer.close).toHaveBeenCalled(); + }, 15000); + + it('writes no tokens when shutdown aborts an in-flight code exchange', async () => { + const gated = await startGatedExchangeAuthServer(); + const fixture = makeFixture(); + cleanups.push(() => fixture.service.dispose()); + cleanups.push(() => gated.releaseExchange()); + const provider = fixture.service.getProvider(SERVER_NAME, gated.url); + await provider.ready; + const state = authServerState(gated.url); + await provider.saveDiscoveryState(state.discovery); + await provider.saveClientInformation(state.client); + + const flow = await fixture.service.beginAuthorization(SERVER_NAME, gated.url); + const complete = flow.complete({ timeoutMs: 10_000 }); + await deliverCallback(flow); + await gated.exchangeStarted; + + const shutdown = fixture.service.shutdown(); + await expect(complete).rejects.toThrow(/OAuth flow for "notion" failed/); + await shutdown; + + expect((await fixture.service.tokenState(SERVER_NAME, gated.url)).hasTokens).toBe(false); + expect(fixture.events.filter((event) => event.type === 'tokens-saved')).toHaveLength(0); + }, 15000); + + it('does not log a drain timeout after a clean shutdown settles', async () => { + const warnings: string[] = []; + const log: Logger = { + error: () => {}, + warn: (message) => { + warnings.push(message); + }, + info: () => {}, + debug: () => {}, + child: () => log, + }; + const fixture = makeFixture(createMemoryMcpOAuthStore(), { + shutdownDrainTimeoutMs: 50, + log, + }); + cleanups.push(() => fixture.service.dispose()); + + await fixture.service.shutdown(); + await fixture.scheduler.advanceBy(1_000); + + expect(warnings).toHaveLength(0); + }); +}); diff --git a/packages/agent-core-v2/test/mcpCore/stubs.ts b/packages/agent-core-v2/test/mcpCore/stubs.ts index 2e11f25c74..957e9d9d87 100644 --- a/packages/agent-core-v2/test/mcpCore/stubs.ts +++ b/packages/agent-core-v2/test/mcpCore/stubs.ts @@ -5,10 +5,14 @@ import type { AddressInfo } from 'node:net'; import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js'; import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; -import type { Tool as KosongTool } from '#/kosong/contract/tool'; import { z } from 'zod'; +import type { Tool as KosongTool } from '#/kosong/contract/tool'; import type { McpOAuthStore } from '#/mcpCore/oauth/store'; +import type { + McpOAuthScheduledTask, + McpOAuthScheduler, +} from '#/mcpCore/oauth/service'; import type { MCPClient, MCPToolDefinition } from '#/mcpCore/types'; import type { ExecutableTool, @@ -20,7 +24,8 @@ import type { export const fixturesDir = new URL('./fixtures/', import.meta.url).pathname; export const stdioFixture = new URL('./fixtures/mock-stdio-server.mjs', import.meta.url).pathname; export const cwdStdioFixture = new URL('./fixtures/cwd-stdio-server.mjs', import.meta.url).pathname; -export const slowStdioFixture = new URL('./fixtures/slow-stdio-server.mjs', import.meta.url).pathname; +export const slowStdioFixture = new URL('./fixtures/slow-stdio-server.mjs', import.meta.url) + .pathname; export const slowToolStdioFixture = new URL( './fixtures/slow-tool-stdio-server.mjs', import.meta.url, @@ -50,9 +55,48 @@ export function createMemoryMcpOAuthStore(): McpOAuthStore { async remove(key: string): Promise { data.delete(key); }, + async list(prefix?: string): Promise { + const keys = [...data.keys()]; + return prefix === undefined ? keys : keys.filter((key) => key.startsWith(prefix)); + }, }; } +export class ManualMcpOAuthScheduler implements McpOAuthScheduler { + private current: number; + private sequence = 0; + private readonly tasks = new Map< + number, + { readonly due: number; readonly task: () => void | Promise } + >(); + + constructor(now = Date.now()) { + this.current = now; + } + + now(): number { + return this.current; + } + + schedule(delayMs: number, task: () => void | Promise): McpOAuthScheduledTask { + const id = this.sequence++; + this.tasks.set(id, { due: this.current + delayMs, task }); + return { cancel: () => this.tasks.delete(id) }; + } + + async advanceBy(deltaMs: number): Promise { + this.current += deltaMs; + while (true) { + const next = [...this.tasks] + .filter(([, task]) => task.due <= this.current) + .toSorted((left, right) => left[1].due - right[1].due || left[0] - right[0])[0]; + if (next === undefined) return; + this.tasks.delete(next[0]); + await next[1].task(); + } + } +} + export function fakeMcpClient( tools: readonly MCPToolDefinition[] = [ { @@ -217,6 +261,8 @@ export function closeServer(server: Server): Promise { }); } -function isPromiseLike(value: ToolExecution | Promise): value is Promise { +function isPromiseLike( + value: ToolExecution | Promise, +): value is Promise { return typeof (value as Promise).then === 'function'; } diff --git a/packages/agent-core-v2/test/workspace/workspaceAgentProfileLoader/agentProfileLoader.test.ts b/packages/agent-core-v2/test/workspace/workspaceAgentProfileLoader/agentProfileLoader.test.ts index d837980b48..acff537790 100644 --- a/packages/agent-core-v2/test/workspace/workspaceAgentProfileLoader/agentProfileLoader.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceAgentProfileLoader/agentProfileLoader.test.ts @@ -4,7 +4,7 @@ import { tmpdir } from 'node:os'; import { join } from 'pathe'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -import { Emitter, Event } from '#/_base/event'; +import { AsyncEmitter, Emitter, Event } from '#/_base/event'; import { SyncDescriptor } from '#/_base/di/descriptors'; import type { ServiceIdentifier } from '#/_base/di/instantiation'; import { InstantiationService } from '#/_base/di/instantiationService'; @@ -12,7 +12,7 @@ import { ServiceCollection } from '#/_base/di/serviceCollection'; import { ILogService } from '#/_base/log/log'; import { EXTRA_AGENT_DIRS_SECTION } from '#/workspace/workspaceAgentProfileLoader/configSection'; import { UserAgentProfileLoaderService } from '#/workspace/workspaceAgentProfileLoader/userAgentProfileLoaderService'; -import type { PluginAgentRoot, ReloadSummary } from '#/app/plugin/types'; +import type { PluginAgentRoot, PluginReloadEvent } from '#/app/plugin/types'; import { DEFAULT_AGENT_PROFILE_NAME, normalizeAgentProfile, @@ -172,7 +172,7 @@ function logStub(warnings?: string[]): ILogService { function pluginStub( agentRoots: readonly PluginAgentRoot[] = [], - reloadEmitter?: Emitter, + reloadEmitter?: Emitter, ): IPluginService { return { _serviceBrand: undefined, @@ -194,6 +194,7 @@ function pluginStub( enabledSessionStarts: async () => [], enabledSystemPrompts: async () => [], enabledMcpServers: async () => ({}), + mcpServerEntries: async () => [], enabledHooks: async () => [], hasLoadedSnapshot: () => true, }; @@ -229,7 +230,7 @@ interface StackOptions { readonly extraAgentDirs?: readonly string[]; readonly explicitFiles?: readonly string[]; readonly pluginAgentRoots?: readonly PluginAgentRoot[]; - readonly pluginReloadEmitter?: Emitter; + readonly pluginReloadEmitter?: Emitter; readonly hostFs?: HostFileSystem; readonly fsWatch?: IHostFsWatchService; } @@ -425,7 +426,7 @@ describe('agent profile loaders + session catalog', () => { await withFixture(async (fixture) => { const pluginAgentsDir = join(fixture.extraDir, 'plugin-agents'); await mkdir(pluginAgentsDir, { recursive: true }); - const reloadEmitter = new Emitter(); + const reloadEmitter = new AsyncEmitter(); await withStack( fixture, { @@ -438,7 +439,10 @@ describe('agent profile loaders + session catalog', () => { await writeAgent(pluginAgentsDir, 'late.md', agentMd('late', 'late plugin agent')); const changed = waitForEvent(stack.catalog.onDidChange); - reloadEmitter.fire({ added: [], removed: [], errors: [] }); + await reloadEmitter.fireAsyncConcurrent( + { added: [], removed: [], errors: [] }, + new AbortController().signal, + ); await changed; expect(stack.catalog.get('late')?.description).toBe('late plugin agent'); diff --git a/packages/agent-core-v2/test/workspace/workspaceInstance/workspaceInstanceManager.test.ts b/packages/agent-core-v2/test/workspace/workspaceInstance/workspaceInstanceManager.test.ts index 700c1f00da..b86fc70d89 100644 --- a/packages/agent-core-v2/test/workspace/workspaceInstance/workspaceInstanceManager.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceInstance/workspaceInstanceManager.test.ts @@ -139,10 +139,10 @@ function manager( { scope: () => 'sessions' }, workspaces, { ready }, - ...Array.from({ length: 22 }, () => undefined), + ...Array.from({ length: 23 }, () => undefined), new TestRuntimeUnitHostFactory(), ]; - args[20] = { entries: () => [] }; + args[21] = { entries: () => [] }; const value = Reflect.construct(WorkspaceInstanceManager, args) as WorkspaceInstanceManager; const providers = (value as unknown as { providers: Map }).providers; providers.clear(); @@ -230,4 +230,51 @@ describe('WorkspaceInstanceManager', () => { expect(instance.snapshot().lifecycle).toBe('active'); await value.dispose(); }); + + describe('findContaining', () => { + function rootedWorkspace(id: string, root: string): Workspace { + return { id, root, name: id, createdAt: 0, lastOpenedAt: 0 }; + } + + it('matches exact and nested cwds, preferring the longest containing root', async () => { + const value = manager([ + rootedWorkspace('repo', '/repo'), + rootedWorkspace('sub', '/repo/sub'), + ]); + await value.getOrCreate({ workspaceId: 'repo' }); + await value.getOrCreate({ workspaceId: 'sub' }); + + expect(value.findContaining('/repo')?.id).toBe('repo'); + expect(value.findContaining('/repo/sub')?.id).toBe('sub'); + expect(value.findContaining('/repo/sub/deep/pkg')?.id).toBe('sub'); + expect(value.findContaining('/repo/other')?.id).toBe('repo'); + expect(value.findContaining('/repo-other')).toBeUndefined(); + expect(value.findContaining('/outside')).toBeUndefined(); + await value.dispose(); + }); + + it('matches across Windows spelling variants', async () => { + const value = manager([rootedWorkspace('win', 'C:\\Users\\Foo\\Repo')]); + await value.getOrCreate({ workspaceId: 'win' }); + + expect(value.findContaining('c:/users/foo/repo')?.id).toBe('win'); + expect(value.findContaining('C:/Users/Foo/Repo/sub')?.id).toBe('win'); + expect(value.findContaining('D:/elsewhere')).toBeUndefined(); + await value.dispose(); + }); + + it('matches any absolute cwd against a workspace rooted at /', async () => { + const value = manager([ + rootedWorkspace('root', '/'), + rootedWorkspace('repo', '/repo'), + ]); + await value.getOrCreate({ workspaceId: 'root' }); + await value.getOrCreate({ workspaceId: 'repo' }); + + expect(value.findContaining('/')?.id).toBe('root'); + expect(value.findContaining('/elsewhere')?.id).toBe('root'); + expect(value.findContaining('/repo/sub')?.id).toBe('repo'); + await value.dispose(); + }); + }); }); diff --git a/packages/agent-core-v2/test/workspace/workspaceMcp/initialization.test.ts b/packages/agent-core-v2/test/workspace/workspaceMcp/initialization.test.ts index 6aa0329168..08f342d8ca 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcp/initialization.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceMcp/initialization.test.ts @@ -12,11 +12,13 @@ import { Event } from '#/_base/event'; import { ILogService } from '#/_base/log/log'; import { McpConnectionManager } from '#/mcpCore/connection-manager'; import { MCP_SECTION, type McpSection } from '#/app/mcpConfig/configSection'; -import { IMcpOAuthStore } from '#/app/mcpConfig/oauthStore'; +import { IMcpOAuthService } from '#/app/mcpConfig/oauthService'; +import { IMcpConfigStore, type McpConfigWriteEvent } from '#/app/mcpConfig/configStore'; +import { McpOAuthService } from '#/mcpCore/oauth/service'; import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { IConfigService } from '#/app/config/config'; import { IPluginService } from '#/app/plugin/plugin'; -import type { ReloadSummary } from '#/app/plugin/types'; +import type { PluginReloadEvent } from '#/app/plugin/types'; import { ITelemetryService, noopTelemetryService } from '#/app/telemetry/telemetry'; import { HostFileSystem } from '#/os/backends/node-local/hostFsService'; import { HostProcessService } from '#/os/backends/node-local/hostProcessService'; @@ -73,9 +75,15 @@ describe('Workspace MCP initialization', () => { reg.definePartialInstance(IWorkspaceContext, { cwd, workspaceId: 'test-workspace' }); reg.definePartialInstance(IPluginService, { enabledMcpServers: async () => ({}), - onDidReload: Event.None as Event, + onDidReload: Event.None as Event, + }); + reg.definePartialInstance( + IMcpOAuthService, + new McpOAuthService({ store: createMemoryMcpOAuthStore() }), + ); + reg.definePartialInstance(IMcpConfigStore, { + onDidWrite: Event.None as Event, }); - reg.definePartialInstance(IMcpOAuthStore, createMemoryMcpOAuthStore()); reg.defineInstance(ILogService, stubLog()); reg.defineInstance(ITelemetryService, noopTelemetryService); const runtime = Object.assign( diff --git a/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts b/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts index 02b3687ecb..0332dc1625 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts @@ -1,43 +1,65 @@ import { mkdtempSync } from 'node:fs'; import { rm } from 'node:fs/promises'; +import { createServer as createHttpServer, type Server as HttpServer } from 'node:http'; +import type { AddressInfo as HttpAddress } from 'node:net'; import { tmpdir } from 'node:os'; -import { join } from 'pathe'; +import { join } from 'pathe'; import { afterEach, beforeEach, describe, expect, it, vi, type Mock } from 'vitest'; +import type { OAuthClientInformationFull } from '@modelcontextprotocol/sdk/shared/auth.js'; + +import type { ServiceIdentifier } from '#/_base/di/instantiation'; import { DisposableStore } from '#/_base/di/lifecycle'; import { createServices } from '#/_base/di/test'; -import type { ServiceIdentifier } from '#/_base/di/instantiation'; -import { Emitter } from '#/_base/event'; +import { AsyncEmitter, Emitter } from '#/_base/event'; import { ILogService } from '#/_base/log/log'; -import { McpConnectionManager } from '#/mcpCore/connection-manager'; -import { HostProcessService } from '#/os/backends/node-local/hostProcessService'; +import { IMcpOAuthService } from '#/app/mcpConfig/oauthService'; +import { ISessionManager } from '#/app/sessionManager/sessionManager'; +import { ITelemetryService, noopTelemetryService } from '#/app/telemetry/telemetry'; import type { McpServerConfig } from '#/mcpCore/config-schema'; +import { + McpConnectionManager, + type McpServerEntry, + type McpServerStatus, +} from '#/mcpCore/connection-manager'; +import { McpOAuthService, type McpOAuthEvent } from '#/mcpCore/oauth/service'; +import { HostProcessService } from '#/os/backends/node-local/hostProcessService'; +import { FakeRuntime } from '#/runtime/fakeRuntime'; import { ISessionEphemeralMcpServers } from '#/session/mcp/ephemeralMcpServers'; import { MergedMcpConnectionView } from '#/session/mcp/mergedConnectionView'; import { ISessionMcpHandle } from '#/session/mcp/sessionMcpHandle'; import { ISessionContext, makeSessionContext } from '#/session/sessionContext/sessionContext'; -import { IMcpOAuthStore } from '#/app/mcpConfig/oauthStore'; -import { ITelemetryService, noopTelemetryService } from '#/app/telemetry/telemetry'; -import { ISessionManager } from '#/app/sessionManager/sessionManager'; +import type { SessionWillCreateEvent } from '#/workspace/sessionLifecycle/sessionLifecycle'; import { IWorkspaceContext } from '#/workspace/workspaceContext/workspaceContext'; import { IRuntimeResolver } from '#/workspace/workspaceInstance/workspaceInstanceManager'; -import { FakeRuntime } from '#/runtime/fakeRuntime'; -import type { SessionWillCreateEvent } from '#/workspace/sessionLifecycle/sessionLifecycle'; +import { + IWorkspaceMcpService, + type ISessionMcpOverlay, +} from '#/workspace/workspaceMcp/workspaceMcp'; +import { WorkspaceMcpService } from '#/workspace/workspaceMcp/workspaceMcpService'; import { IWorkspaceMcpConfigService, - type McpServersChange, + type McpServersChangeEvent, type McpTunables, } from '#/workspace/workspaceMcpConfig/workspaceMcpConfig'; -import { IWorkspaceMcpService, type ISessionMcpOverlay } from '#/workspace/workspaceMcp/workspaceMcp'; -import { WorkspaceMcpService } from '#/workspace/workspaceMcp/workspaceMcpService'; import { stubLog } from '../../_base/log/stubs'; -import { createMemoryMcpOAuthStore, startInProcessHttpMcpServer, stdioFixture } from '../../mcpCore/stubs'; import { registerAgentIdentityStub } from '../../app/agentIdentity/stubs'; +import { + createMemoryMcpOAuthStore, + ManualMcpOAuthScheduler, + startInProcessHttpMcpServer, + stdioFixture, +} from '../../mcpCore/stubs'; function stdioServer(): McpServerConfig { - return { transport: 'stdio', command: process.execPath, args: [stdioFixture], runtime_id: 'local' }; + return { + transport: 'stdio', + command: process.execPath, + args: [stdioFixture], + runtime_id: 'local', + }; } describe('WorkspaceMcpService', () => { @@ -46,8 +68,10 @@ describe('WorkspaceMcpService', () => { let current: Record; let tunablesValue: McpTunables; let tunablesFn: Mock<() => McpTunables>; - let configChanges: Emitter; + let configChanges: AsyncEmitter; let assemblyEvents: Emitter; + let oauthService: McpOAuthService; + let oauthScheduler: ManualMcpOAuthScheduler; let manager: InstanceType | undefined; beforeEach(() => { @@ -56,14 +80,20 @@ describe('WorkspaceMcpService', () => { current = {}; tunablesValue = {}; tunablesFn = vi.fn(() => tunablesValue); - configChanges = new Emitter(); + configChanges = disposables.add(new AsyncEmitter()); assemblyEvents = disposables.add(new Emitter()); + oauthScheduler = new ManualMcpOAuthScheduler(); + oauthService = new McpOAuthService({ + store: createMemoryMcpOAuthStore(), + scheduler: oauthScheduler, + }); manager = undefined; }); afterEach(async () => { vi.restoreAllMocks(); await manager?.shutdown(); + await oauthService.dispose(); disposables.dispose(); await rm(cwd, { recursive: true, force: true }); }); @@ -84,14 +114,21 @@ describe('WorkspaceMcpService', () => { additionalServices: (reg) => { reg.definePartialInstance(IWorkspaceContext, { cwd, workspaceId: 'test-workspace' }); reg.defineInstance(IWorkspaceMcpConfigService, mcpConfigStub()); - reg.definePartialInstance(IMcpOAuthStore, createMemoryMcpOAuthStore()); + reg.definePartialInstance(IMcpOAuthService, oauthService); reg.defineInstance(ILogService, stubLog()); reg.defineInstance(ITelemetryService, noopTelemetryService); const runtime = Object.assign( - new FakeRuntime({ workspaceId: 'test-workspace', runtimeId: 'local', generation: 'test-generation' }, { capabilities: ['process'] }), + new FakeRuntime( + { workspaceId: 'test-workspace', runtimeId: 'local', generation: 'test-generation' }, + { capabilities: ['process'] }, + ), { process: new HostProcessService() }, ); - reg.defineInstance(IRuntimeResolver, { _serviceBrand: undefined, inspect: () => runtime, acquire: () => ({ runtime, track: (resource) => resource, dispose: () => {} }) }); + reg.defineInstance(IRuntimeResolver, { + _serviceBrand: undefined, + inspect: () => runtime, + acquire: () => ({ runtime, track: (resource) => resource, dispose: () => {} }), + }); reg.definePartialInstance(ISessionManager, { onWillCreateSession: assemblyEvents.event, }); @@ -135,23 +172,26 @@ describe('WorkspaceMcpService', () => { await service.ready; expect(manager.get('alpha')?.status).toBe('connected'); - configChanges.fire({ upsert: { beta: stdioServer() }, remove: ['alpha'] }); - - await vi.waitFor( - () => { - expect(manager?.get('alpha')?.status).toBe('removed'); - expect(manager?.get('beta')?.status).toBe('connected'); - }, - { timeout: 10000, interval: 50 }, + await configChanges.fireAsync( + { upsert: { beta: stdioServer() }, remove: ['alpha'] }, + new AbortController().signal, ); + + expect(manager.get('alpha')?.status).toBe('removed'); + expect(manager.get('beta')?.status).toBe('connected'); }, 20000); it('queues change events until the initial connect settles', async () => { current = { alpha: stdioServer() }; let settleConnectAll: () => void = () => undefined; + let signalConnectAllStarted: () => void = () => undefined; + const connectAllStarted = new Promise((resolve) => { + signalConnectAllStarted = resolve; + }); vi.spyOn(McpConnectionManager.prototype, 'connectAll').mockImplementation( () => new Promise((resolve) => { + signalConnectAllStarted(); settleConnectAll = resolve; }), ); @@ -165,8 +205,11 @@ describe('WorkspaceMcpService', () => { const service = createService(); manager = service.connectionManager(); - configChanges.fire({ upsert: { beta: stdioServer() }, remove: ['alpha'] }); - await new Promise((resolvePromise) => setTimeout(resolvePromise, 300)); + void configChanges.fireAsync( + { upsert: { beta: stdioServer() }, remove: ['alpha'] }, + new AbortController().signal, + ); + await connectAllStarted; expect(connect).not.toHaveBeenCalled(); expect(markRemoved).not.toHaveBeenCalled(); @@ -342,7 +385,7 @@ describe('WorkspaceMcpService', () => { const disposers: Array<() => void> = []; const event: SessionWillCreateEvent = { sessionId: 's1', - readSeed: (id: ServiceIdentifier): T => seeds.get(id) as T, + readSeed: (id: ServiceIdentifier): T => seeds.get(id) as T, contributeSeed: (id, value) => { contributed.set(id, value); }, @@ -410,6 +453,304 @@ describe('WorkspaceMcpService', () => { expect(disposers).toHaveLength(0); }); }); + + describe('credential events', () => { + const SERVER_URL = 'https://mcp.example.test/mcp'; + let httpServers: HttpServer[] = []; + + afterEach(async () => { + const servers = httpServers; + httpServers = []; + await Promise.all( + servers.map( + (server) => + new Promise((resolve, reject) => { + server.close((err) => { + if (err) { + reject(err); + return; + } + resolve(); + }); + }), + ), + ); + }); + + function mockManagerEntry( + status: McpServerStatus, + url: string = SERVER_URL, + ): Mock<(name: string) => Promise> { + vi.spyOn(McpConnectionManager.prototype, 'get').mockReturnValue({ + name: 'notion', + transport: 'http', + status, + toolCount: 0, + }); + vi.spyOn(McpConnectionManager.prototype, 'getRemoteServerUrl').mockReturnValue(url); + return vi + .spyOn(McpConnectionManager.prototype, 'reconnectAndJoin') + .mockResolvedValue(undefined); + } + + async function startRefreshFailingServer(): Promise<{ + origin: string; + counts: { refresh: number }; + }> { + const counts = { refresh: 0 }; + const server: HttpServer = createHttpServer((req, res) => { + if (req.method === 'POST' && req.url === '/token') { + counts.refresh += 1; + res.writeHead(500, { 'content-type': 'text/plain' }); + res.end('broken'); + return; + } + res.writeHead(404).end(); + }); + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)); + httpServers.push(server); + const port = (server.address() as HttpAddress).port; + return { origin: `http://127.0.0.1:${port}`, counts }; + } + + async function seedOAuthServerState(authServerOrigin: string): Promise { + const provider = oauthService.getProvider('notion', SERVER_URL); + await provider.ready; + await provider.saveDiscoveryState({ + authorizationServerUrl: authServerOrigin, + authorizationServerMetadata: { + issuer: authServerOrigin, + authorization_endpoint: `${authServerOrigin}/authorize`, + token_endpoint: `${authServerOrigin}/token`, + response_types_supported: ['code'], + grant_types_supported: ['authorization_code', 'refresh_token'], + token_endpoint_auth_methods_supported: ['none'], + }, + }); + await provider.saveClientInformation({ + client_id: 'cached-client', + redirect_uris: ['http://127.0.0.1:45678/callback'], + token_endpoint_auth_method: 'none', + grant_types: ['authorization_code', 'refresh_token'], + response_types: ['code'], + } satisfies OAuthClientInformationFull); + } + + it('reconnects a needs-auth entry when tokens are saved', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + const reconnectAndJoin = mockManagerEntry('needs-auth'); + + await oauthService + .getProvider('notion', SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer' }); + + await vi.waitFor(() => { + expect(reconnectAndJoin).toHaveBeenCalledWith('notion'); + }); + }); + + it('leaves a connected entry alone when tokens are saved', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + const reconnectAndJoin = mockManagerEntry('connected'); + + await oauthService + .getProvider('notion', SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer' }); + + expect(reconnectAndJoin).not.toHaveBeenCalled(); + }); + + it('forgets the provider and reconnects a connected entry on token invalidation', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + const reconnectAndJoin = mockManagerEntry('connected'); + const forgetProvider = vi.spyOn(oauthService, 'forgetProvider'); + + await oauthService.invalidate('notion', SERVER_URL, 'all'); + + await vi.waitFor(() => { + expect(reconnectAndJoin).toHaveBeenCalledWith('notion'); + }); + expect(forgetProvider).toHaveBeenCalledWith('notion', SERVER_URL); + }); + + it('ignores a credential event for a different server URL', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + const reconnectAndJoin = mockManagerEntry('needs-auth', 'https://other.example.test/mcp'); + const forgetProvider = vi.spyOn(oauthService, 'forgetProvider'); + + await oauthService + .getProvider('notion', SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer' }); + + expect(reconnectAndJoin).not.toHaveBeenCalled(); + expect(forgetProvider).not.toHaveBeenCalled(); + }); + + it('defers the reconnect of a pending entry until its initial connect settles', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + mockManagerEntry('pending'); + const reconnectAfterCurrent = vi + .spyOn(McpConnectionManager.prototype, 'reconnectAfterCurrent') + .mockResolvedValue(undefined); + let notifyStatus: ((entry: McpServerEntry) => void) | undefined; + vi.spyOn(McpConnectionManager.prototype, 'onStatusChange').mockImplementation( + (listener) => { + notifyStatus = listener; + return () => undefined; + }, + ); + + await oauthService + .getProvider('notion', SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer' }); + expect(reconnectAfterCurrent).not.toHaveBeenCalled(); + + notifyStatus?.({ name: 'notion', transport: 'http', status: 'connected', toolCount: 0 }); + await vi.waitFor(() => { + expect(reconnectAfterCurrent).toHaveBeenCalledWith('notion'); + }); + }); + + it('reconnects when a pending entry settles before the status listener is attached', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + vi.spyOn(McpConnectionManager.prototype, 'get') + .mockReturnValueOnce({ + name: 'notion', + transport: 'http', + status: 'pending', + toolCount: 0, + }) + .mockReturnValue({ + name: 'notion', + transport: 'http', + status: 'needs-auth', + toolCount: 0, + }); + vi.spyOn(McpConnectionManager.prototype, 'getRemoteServerUrl').mockReturnValue(SERVER_URL); + vi.spyOn(McpConnectionManager.prototype, 'onStatusChange').mockReturnValue(() => undefined); + const reconnectAfterCurrent = vi + .spyOn(McpConnectionManager.prototype, 'reconnectAfterCurrent') + .mockResolvedValue(undefined); + + await oauthService + .getProvider('notion', SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer' }); + + await vi.waitFor(() => { + expect(reconnectAfterCurrent).toHaveBeenCalledWith('notion'); + }); + }); + + it('ignores a client-scope invalidation as flow-local churn', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + const reconnectAndJoin = mockManagerEntry('connected'); + const forgetProvider = vi.spyOn(oauthService, 'forgetProvider'); + + const provider = oauthService.getProvider('notion', SERVER_URL); + await provider.ready; + await provider.clearCredentials('client'); + + expect(reconnectAndJoin).not.toHaveBeenCalled(); + expect(forgetProvider).not.toHaveBeenCalled(); + }); + + it('ignores a discovery-scope invalidation as flow-local churn', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + const reconnectAndJoin = mockManagerEntry('connected'); + const forgetProvider = vi.spyOn(oauthService, 'forgetProvider'); + + const provider = oauthService.getProvider('notion', SERVER_URL); + await provider.ready; + await provider.clearCredentials('discovery'); + + expect(reconnectAndJoin).not.toHaveBeenCalled(); + expect(forgetProvider).not.toHaveBeenCalled(); + }); + + it('reconnects a connected entry when a proactive refresh fails', async () => { + const authServer = await startRefreshFailingServer(); + await seedOAuthServerState(authServer.origin); + const service = createService(); + manager = service.connectionManager(); + await service.ready; + + await oauthService.getProvider('notion', SERVER_URL).saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + expires_in: 60, + }); + const reconnectAndJoin = mockManagerEntry('connected'); + + await oauthScheduler.advanceBy(30_000); + expect(reconnectAndJoin).toHaveBeenCalledWith('notion'); + expect(authServer.counts.refresh).toBe(1); + }); + + it('ignores a failed proactive refresh for a needs-auth entry', async () => { + const authServer = await startRefreshFailingServer(); + await seedOAuthServerState(authServer.origin); + const events: McpOAuthEvent[] = []; + oauthService.onEvent((event) => events.push(event)); + const service = createService(); + manager = service.connectionManager(); + await service.ready; + + await oauthService.getProvider('notion', SERVER_URL).saveTokens({ + access_token: 'stale-access-token', + refresh_token: 'stale-refresh-token', + token_type: 'Bearer', + expires_in: 60, + }); + const reconnectAndJoin = mockManagerEntry('needs-auth'); + + await oauthScheduler.advanceBy(30_000); + expect(events.some((event) => event.type === 'refresh-failed')).toBe(true); + expect(reconnectAndJoin).not.toHaveBeenCalled(); + }); + + it('does not reconnect a disabled entry when tokens are saved', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + const reconnectAndJoin = mockManagerEntry('disabled'); + + await oauthService + .getProvider('notion', SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer' }); + + expect(reconnectAndJoin).not.toHaveBeenCalled(); + }); + + it('does not reconnect a removed entry when tokens are saved', async () => { + const service = createService(); + manager = service.connectionManager(); + await service.ready; + const reconnectAndJoin = mockManagerEntry('removed'); + + await oauthService + .getProvider('notion', SERVER_URL) + .saveTokens({ access_token: 'a', token_type: 'Bearer' }); + + expect(reconnectAndJoin).not.toHaveBeenCalled(); + }); + }); }); describe('MergedMcpConnectionView', () => { @@ -441,11 +782,12 @@ describe('MergedMcpConnectionView', () => { await overlay.connect('eph', disabledStdio('eph-cmd')); const view = new MergedMcpConnectionView(base, overlay, new Set(['shared', 'eph'])); - expect(view.list().map((entry) => entry.name).toSorted()).toEqual([ - 'base-only', - 'eph', - 'shared', - ]); + expect( + view + .list() + .map((entry) => entry.name) + .toSorted(), + ).toEqual(['base-only', 'eph', 'shared']); expect(view.get('shared')?.transport).toBe('http'); expect(view.get('base-only')?.transport).toBe('stdio'); diff --git a/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts b/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts index 4d708fd2c0..d8b509188f 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts @@ -1,20 +1,24 @@ import { mkdtempSync } from 'node:fs'; import { mkdir, rm, writeFile } from 'node:fs/promises'; import { tmpdir } from 'node:os'; -import { join } from 'pathe'; +import { join } from 'pathe'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { DisposableStore } from '#/_base/di/lifecycle'; import { createServices } from '#/_base/di/test'; -import { Emitter } from '#/_base/event'; +import { AsyncEmitter, Emitter } from '#/_base/event'; import { ILogService } from '#/_base/log/log'; -import type { McpServerConfig } from '#/mcpCore/config-schema'; -import { MCP_SECTION, type McpSection } from '#/app/mcpConfig/configSection'; import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { IConfigService } from '#/app/config/config'; +import { MCP_SECTION, type McpSection } from '#/app/mcpConfig/configSection'; +import { + IMcpConfigStore, + type McpConfigWriteEvent, +} from '#/app/mcpConfig/configStore'; import { IPluginService } from '#/app/plugin/plugin'; -import type { ReloadSummary } from '#/app/plugin/types'; +import type { PluginReloadEvent } from '#/app/plugin/types'; +import type { McpServerConfig } from '#/mcpCore/config-schema'; import { HostFileSystem } from '#/os/backends/node-local/hostFsService'; import { IHostFileSystem } from '#/os/interface/hostFileSystem'; import { @@ -23,15 +27,15 @@ import { type IHostFsWatchHandle, } from '#/os/interface/hostFsWatch'; import { IWorkspaceContext } from '#/workspace/workspaceContext/workspaceContext'; -import { - IWorkspaceTrust, - type WorkspaceTrustChange, -} from '#/workspace/workspaceTrust/workspaceTrust'; import { IWorkspaceMcpConfigService, type McpServersChange, } from '#/workspace/workspaceMcpConfig/workspaceMcpConfig'; import { WorkspaceMcpConfigService } from '#/workspace/workspaceMcpConfig/workspaceMcpConfigService'; +import { + IWorkspaceTrust, + type WorkspaceTrustChange, +} from '#/workspace/workspaceTrust/workspaceTrust'; import { stubLog } from '../../_base/log/stubs'; @@ -45,7 +49,8 @@ describe('WorkspaceMcpConfigService', () => { let disposables: DisposableStore; let watchFires: Map>; let pluginServers: Record; - let pluginReloads: Emitter; + let pluginReloads: AsyncEmitter; + let storeWrites: AsyncEmitter; let trusted: boolean; let trustFlips: Emitter; let changes: McpServersChange[]; @@ -56,7 +61,8 @@ describe('WorkspaceMcpConfigService', () => { disposables = new DisposableStore(); watchFires = new Map(); pluginServers = {}; - pluginReloads = new Emitter(); + pluginReloads = disposables.add(new AsyncEmitter()); + storeWrites = disposables.add(new AsyncEmitter()); trusted = true; trustFlips = new Emitter(); changes = []; @@ -98,8 +104,8 @@ describe('WorkspaceMcpConfigService', () => { reg.defineInstance(ILogService, stubLog()); reg.definePartialInstance(IConfigService, { ready: Promise.resolve(), - get: ((domain: string): T => - (domain === MCP_SECTION ? mcpSection : undefined) as T), + get: (domain: string): T => + (domain === MCP_SECTION ? mcpSection : undefined) as T, }); reg.defineInstance(IHostFsWatchService, fsWatchStub()); reg.defineInstance(IHostFileSystem, new HostFileSystem()); @@ -108,11 +114,12 @@ describe('WorkspaceMcpConfigService', () => { isTrusted: () => trusted, onDidChange: trustFlips.event, }); + reg.definePartialInstance(IMcpConfigStore, { onDidWrite: storeWrites.event }); reg.define(IWorkspaceMcpConfigService, WorkspaceMcpConfigService); }, }); const service = ix.get(IWorkspaceMcpConfigService); - service.onDidChange((change) => changes.push(change)); + service.onDidChange(({ upsert, remove }) => changes.push({ upsert, remove })); return service; } @@ -125,7 +132,10 @@ describe('WorkspaceMcpConfigService', () => { } it('merges file and plugin servers in the initial resolve (file wins name collisions)', async () => { - await writeProjectConfig({ shared: stdioConfig('file-version'), fileOnly: stdioConfig('file') }); + await writeProjectConfig({ + shared: stdioConfig('file-version'), + fileOnly: stdioConfig('file'), + }); pluginServers = { shared: stdioConfig('plugin-version'), pluginOnly: stdioConfig('plugin') }; const service = createService(); @@ -209,33 +219,44 @@ describe('WorkspaceMcpConfigService', () => { await vi.waitFor( () => { - expect(changes).toEqual([ - { upsert: { beta: stdioConfig('beta') }, remove: ['alpha'] }, - ]); + expect(changes).toEqual([{ upsert: { beta: stdioConfig('beta') }, remove: ['alpha'] }]); }, { timeout: 10000, interval: 50 }, ); expect(service.servers()).toEqual({ beta: stdioConfig('beta') }); }, 20000); - it('falls back to the same-named plugin server when a file server vanishes', async () => { - const file = await writeProjectConfig({ shared: stdioConfig('file-version') }); + it('revives the same-named plugin entry when the winning file entry vanishes', async () => { + await writeProjectConfig({ shared: stdioConfig('file-version') }); pluginServers = { shared: stdioConfig('plugin-version') }; const service = createService(); await service.ready; expect(service.servers()).toEqual({ shared: stdioConfig('file-version') }); await writeProjectConfig({}); - watchFires.get(cwd)?.fire({ path: file, action: 'modified', kind: 'file' }); + await storeWrites.fireAsync({}, new AbortController().signal); + pluginServers = { + shared: stdioConfig('plugin-version'), + pluginOnly: stdioConfig('plugin'), + }; + await pluginReloads.fireAsyncConcurrent( + { added: [], removed: [], errors: [] }, + new AbortController().signal, + ); await vi.waitFor( () => { expect(changes).toEqual([ { upsert: { shared: stdioConfig('plugin-version') }, remove: [] }, + { upsert: { pluginOnly: stdioConfig('plugin') }, remove: [] }, ]); }, { timeout: 10000, interval: 50 }, ); + expect(service.servers()).toEqual({ + shared: stdioConfig('plugin-version'), + pluginOnly: stdioConfig('plugin'), + }); }, 20000); it('publishes a plugin server that appears on plugin reload', async () => { @@ -243,7 +264,10 @@ describe('WorkspaceMcpConfigService', () => { await service.ready; pluginServers = { gamma: stdioConfig('gamma') }; - pluginReloads.fire({ added: [], removed: [], errors: [] }); + await pluginReloads.fireAsyncConcurrent( + { added: [], removed: [], errors: [] }, + new AbortController().signal, + ); await vi.waitFor( () => { @@ -254,13 +278,30 @@ describe('WorkspaceMcpConfigService', () => { expect(service.servers()).toEqual({ gamma: stdioConfig('gamma') }); }, 20000); + it('settles the plugin reload event only after the workspace reconcile is published', async () => { + const service = createService(); + await service.ready; + + pluginServers = { gamma: stdioConfig('gamma') }; + await pluginReloads.fireAsyncConcurrent( + { added: [], removed: [], errors: [] }, + new AbortController().signal, + ); + + expect(changes).toEqual([{ upsert: { gamma: stdioConfig('gamma') }, remove: [] }]); + expect(service.servers()).toEqual({ gamma: stdioConfig('gamma') }); + }); + it('removes a plugin server that vanishes on plugin reload', async () => { pluginServers = { alpha: stdioConfig('alpha') }; const service = createService(); await service.ready; pluginServers = {}; - pluginReloads.fire({ added: [], removed: [], errors: [] }); + await pluginReloads.fireAsyncConcurrent( + { added: [], removed: [], errors: [] }, + new AbortController().signal, + ); await vi.waitFor( () => { @@ -270,7 +311,7 @@ describe('WorkspaceMcpConfigService', () => { ); }, 20000); - it('stays silent when a vanished plugin server leaves the same-named file entry in place', async () => { + it('keeps the winning file server when the same-named plugin entry vanishes', async () => { await writeProjectConfig({ shared: stdioConfig('file-version') }); pluginServers = { shared: stdioConfig('plugin-version') }; const service = createService(); @@ -278,10 +319,33 @@ describe('WorkspaceMcpConfigService', () => { expect(service.servers()).toEqual({ shared: stdioConfig('file-version') }); pluginServers = {}; - pluginReloads.fire({ added: [], removed: [], errors: [] }); + await pluginReloads.fireAsyncConcurrent( + { added: [], removed: [], errors: [] }, + new AbortController().signal, + ); - await new Promise((resolvePromise) => setTimeout(resolvePromise, 500)); expect(changes).toEqual([]); expect(service.servers()).toEqual({ shared: stdioConfig('file-version') }); + }); + + it('reloads immediately on a management-plane write, without the watch debounce', async () => { + const service = createService(); + await service.ready; + expect(service.servers()).toEqual({}); + + await writeFile( + join(homeDir, 'mcp.json'), + JSON.stringify({ mcpServers: { added: stdioConfig('added') } }), + 'utf8', + ); + await storeWrites.fireAsync({}, new AbortController().signal); + + await vi.waitFor( + () => { + expect(changes).toEqual([{ upsert: { added: stdioConfig('added') }, remove: [] }]); + }, + { timeout: 10000, interval: 20 }, + ); + expect(service.servers()).toEqual({ added: stdioConfig('added') }); }, 20000); }); diff --git a/packages/agent-core-v2/test/workspace/workspaceSkillCatalog/skillCatalog.test.ts b/packages/agent-core-v2/test/workspace/workspaceSkillCatalog/skillCatalog.test.ts index 1204517b73..c4044849e5 100644 --- a/packages/agent-core-v2/test/workspace/workspaceSkillCatalog/skillCatalog.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceSkillCatalog/skillCatalog.test.ts @@ -11,11 +11,11 @@ import { _clearScopedRegistryForTests, registerScopedService, } from '#/_base/di/scope'; -import { Emitter, Event } from '#/_base/event'; +import { AsyncEmitter, Emitter, Event } from '#/_base/event'; import { IBootstrapService } from '#/app/bootstrap/bootstrap'; import { IPluginService } from '#/app/plugin/plugin'; import { PluginService } from '#/app/plugin/pluginService'; -import type { ReloadSummary } from '#/app/plugin/types'; +import type { PluginReloadEvent } from '#/app/plugin/types'; import { IProviderService } from '#/kosong/provider/provider'; import { IHostFsWatchService, @@ -109,7 +109,7 @@ function configStub(): IConfigService & { function pluginStub( skillRoots: readonly SkillRoot[] = [], - reloadEmitter?: Emitter, + reloadEmitter?: Emitter, ): IPluginService { return { _serviceBrand: undefined, @@ -131,6 +131,7 @@ function pluginStub( enabledSessionStarts: async () => [], enabledSystemPrompts: async () => [], enabledMcpServers: async () => ({}), + mcpServerEntries: async () => [], enabledHooks: async () => [], hasLoadedSnapshot: () => true, }; @@ -168,7 +169,7 @@ function makeHost( ws: IWorkspaceContext, pluginRoots: readonly SkillRoot[] = [], explicitDirs?: readonly string[], - pluginReloadEmitter?: Emitter, + pluginReloadEmitter?: Emitter, ) { const config = configStub(); const host = createScopedTestHost([ @@ -585,7 +586,7 @@ describe('WorkspaceSkillCatalogService', () => { store.setPluginSkills([ stubSkill('demo-skill', { source: 'extra', plugin: { id: 'demo' } }), ]); - const reloadEmitter = new Emitter(); + const reloadEmitter = new AsyncEmitter(); const pluginRoot: SkillRoot = { path: '/plugins/demo/skills', source: 'extra', @@ -605,7 +606,10 @@ describe('WorkspaceSkillCatalogService', () => { resolve(sourceId); }); }); - reloadEmitter.fire({ added: [], removed: [], errors: [] }); + await reloadEmitter.fireAsyncConcurrent( + { added: [], removed: [], errors: [] }, + new AbortController().signal, + ); await expect(refreshed).resolves.toBe('plugin'); } finally { @@ -695,7 +699,7 @@ describe('WorkspaceSkillCatalogService', () => { }); it('binds thisArg when forwarding plugin reloads through the plugin skill source', async () => { - const reloadEmitter = new Emitter(); + const reloadEmitter = new AsyncEmitter(); const pluginService = pluginStub([], reloadEmitter); const ws = workspaceContextStub('/work'); const host = createScopedTestHost([ @@ -722,7 +726,10 @@ describe('WorkspaceSkillCatalogService', () => { receiver, ); - reloadEmitter.fire({ added: [], removed: [], errors: [] }); + await reloadEmitter.fireAsyncConcurrent( + { added: [], removed: [], errors: [] }, + new AbortController().signal, + ); expect(seen).toEqual([receiver]); subscription?.dispose(); diff --git a/packages/agent-core-v2/test/workspace/workspaceTrust/workspaceTrust.test.ts b/packages/agent-core-v2/test/workspace/workspaceTrust/workspaceTrust.test.ts index c0a24b4dea..7604e1550b 100644 --- a/packages/agent-core-v2/test/workspace/workspaceTrust/workspaceTrust.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceTrust/workspaceTrust.test.ts @@ -5,6 +5,7 @@ import { join } from 'pathe'; import { afterEach, beforeEach, describe, expect, it } from 'vitest'; +import { encodeWorkDirKey } from '#/_base/utils/workdir-slug'; import { DisposableStore } from '#/_base/di/lifecycle'; import { createServices } from '#/_base/di/test'; import { JsonAtomicDocumentStore } from '#/persistence/backends/node-fs/atomicDocumentStore'; @@ -20,6 +21,11 @@ import { WorkspaceTrustService, workspaceTrustTrustedKey, } from '#/workspace/workspaceTrust/workspaceTrustService'; +import { + deleteWorkspaceTrust, + readWorkspaceTrust, + writeWorkspaceTrust, +} from '#/workspace/workspaceTrust/trustRecord'; import { registerStateServices } from '../../state/stubs'; @@ -111,6 +117,41 @@ describe('WorkspaceTrustService', () => { expect(second.isTrusted()).toBe(true); }); + it('migrates a legacy Windows trust marker to the canonical key', async () => { + const docs = new JsonAtomicDocumentStore(new FileStorageService(homeDir)); + const root = 'C:\\Users\\Foo\\Repo'; + const record = { root, trustedAt: 1 }; + await docs.set('workspace-trust', encodeWorkDirKey(root), record); + + expect(await readWorkspaceTrust(docs, root)).toBe(true); + await expect( + docs.get('workspace-trust', encodeWorkDirKey('c:/users/foo/repo')), + ).resolves.toEqual(record); + }); + + it('shares one trust key across UNC and drive-letter spelling variants', async () => { + const docs = new JsonAtomicDocumentStore(new FileStorageService(homeDir)); + await writeWorkspaceTrust(docs, '//server/share/repo', 1); + expect(await readWorkspaceTrust(docs, '\\\\SERVER\\SHARE\\REPO')).toBe(true); + expect(await readWorkspaceTrust(docs, '//Server/Share/Repo')).toBe(true); + + await writeWorkspaceTrust(docs, 'C:\\Users\\Foo\\Repo', 2); + expect(await readWorkspaceTrust(docs, 'c:/users/foo/repo')).toBe(true); + }); + + it('deletes both canonical and legacy trust markers', async () => { + const docs = new JsonAtomicDocumentStore(new FileStorageService(homeDir)); + const root = 'C:\\Users\\Foo\\Repo'; + const legacyKey = encodeWorkDirKey(root); + await docs.set('workspace-trust', legacyKey, { root, trustedAt: 1 }); + await writeWorkspaceTrust(docs, root, 2); + + await deleteWorkspaceTrust(docs, root); + + await expect(docs.get('workspace-trust', legacyKey)).resolves.toBeUndefined(); + await expect(readWorkspaceTrust(docs, root)).resolves.toBe(false); + }); + it('tracks different roots independently', async () => { const other = mkdtempSync(join(tmpdir(), 'kimi-workspace-trust-other-')); try { diff --git a/packages/agent-core/src/errors/codes.ts b/packages/agent-core/src/errors/codes.ts index 3032f3b0ac..46e569db7a 100644 --- a/packages/agent-core/src/errors/codes.ts +++ b/packages/agent-core/src/errors/codes.ts @@ -67,6 +67,7 @@ export const ErrorCodes = { MCP_SERVER_DISABLED: 'mcp.server_disabled', MCP_STARTUP_FAILED: 'mcp.startup_failed', MCP_TOOL_NAME_COLLISION: 'mcp.tool_name_collision', + MCP_OAUTH_FAILED: 'mcp.oauth_failed', PLUGIN_NOT_FOUND: 'plugin.not_found', PLUGIN_LOAD_FAILED: 'plugin.load_failed', @@ -403,6 +404,12 @@ export const KIMI_ERROR_INFO = { public: true, action: 'Rename one of the colliding MCP tools or servers so their qualified names are unique.', }, + 'mcp.oauth_failed': { + title: 'MCP OAuth authorization failed', + retryable: true, + public: true, + action: 'Begin the authorization flow again; inspect the error details if it keeps failing.', + }, 'plugin.not_found': { title: 'Plugin not found', @@ -456,3 +463,12 @@ export const KIMI_ERROR_INFO = { action: 'Inspect logs or report the issue with diagnostics.', }, } as const satisfies Record; + +/** + * Runtime membership check against the registry. Errors arriving from another + * process or engine generation may carry codes this build does not declare; + * branch on this before indexing `KIMI_ERROR_INFO` with an untrusted code. + */ +export function isKimiErrorCode(code: unknown): code is KimiErrorCode { + return typeof code === 'string' && Object.hasOwn(KIMI_ERROR_INFO, code); +} diff --git a/packages/agent-core/src/errors/index.ts b/packages/agent-core/src/errors/index.ts index 66379cbe48..ea018e1759 100644 --- a/packages/agent-core/src/errors/index.ts +++ b/packages/agent-core/src/errors/index.ts @@ -1,5 +1,6 @@ export { ErrorCodes, + isKimiErrorCode, KIMI_ERROR_INFO, type KimiErrorCode, type KimiErrorInfo, diff --git a/packages/agent-core/src/rpc/core-api.ts b/packages/agent-core/src/rpc/core-api.ts index 1ae6239c7a..787eef6b35 100644 --- a/packages/agent-core/src/rpc/core-api.ts +++ b/packages/agent-core/src/rpc/core-api.ts @@ -402,10 +402,12 @@ export interface GetGlobalMcpServerPayload { export interface PutGlobalMcpServerPayload { readonly server: GlobalMcpServerConfig; + readonly cwd?: string; } export interface GlobalMcpServerNamePayload { readonly name: string; + readonly cwd?: string; } /** @@ -424,10 +426,12 @@ export type McpServerLocator = export interface McpServerLocatorPayload { readonly locator: McpServerLocator; + readonly cwd?: string; } export interface InspectAppMcpServersPayload { readonly targets?: readonly McpServerLocator[]; + readonly cwd?: string; } export type GlobalMcpServerAuthState = @@ -447,9 +451,10 @@ export interface GlobalMcpServerAuthStatus { export interface ListGlobalMcpServerAuthStatusesPayload { readonly cwd?: string; /** - * Verify online: run a real connection probe for OAuth-capable servers so - * an expired/revoked grant surfaces as `oauth-expired` instead of the - * offline `oauth-authorized` guess. + * Omitted preserves implicit OAuth detection for unpinned servers without + * stored credentials. `false` stays fully offline; `true` verifies every + * OAuth-capable server so an expired/revoked grant surfaces as + * `oauth-expired` instead of the offline `oauth-authorized` guess. */ readonly verify?: boolean; } diff --git a/packages/agent-core/src/rpc/core-impl.ts b/packages/agent-core/src/rpc/core-impl.ts index 17b261aa30..34c505f700 100644 --- a/packages/agent-core/src/rpc/core-impl.ts +++ b/packages/agent-core/src/rpc/core-impl.ts @@ -855,24 +855,23 @@ export class KimiCore implements PromisableMethods { ): Promise { await this.awaitMcpRegistryReady(); const entries = await this.mcpRegistry.list({ cwd: input?.cwd }); - const verify = input?.verify === true; return Promise.all( entries.map(async (entry) => ({ name: entry.name, - authStatus: await this.mcpServerAuthState(entry, input?.cwd, verify), + authStatus: await this.mcpServerAuthState(entry, input?.cwd, input?.verify), })), ); } async addGlobalMcpServer( - { server }: PutGlobalMcpServerPayload, + { server, cwd }: PutGlobalMcpServerPayload, ): Promise { await this.awaitMcpRegistryReady(); // Normalize once: the store trims names, so the read-only guard, the // persisted key, and live-session reconciliation must all agree (a padded // name would otherwise persist trimmed but reconcile the raw name). const name = normalizeServerName(server.name); - const existing = await this.mcpRegistry.get(name).catch(() => undefined); + const existing = await this.mcpRegistry.get(name, { cwd }).catch(() => undefined); if (existing !== undefined && !(existing.source === 'global' && existing.mutable)) { // A same-named plugin / project-layer entry already exists; writing a // user-level shadow would silently change precedence, so reject. A @@ -882,15 +881,15 @@ export class KimiCore implements PromisableMethods { } await this.globalMcpConfig.add({ ...server, name }); await this.reconcileMcpServerInSessions([name], 'global-add'); - return this.listGlobalMcpServers({}); + return this.listGlobalMcpServers({ cwd }); } async updateGlobalMcpServer( - { server }: PutGlobalMcpServerPayload, + { server, cwd }: PutGlobalMcpServerPayload, ): Promise { await this.awaitMcpRegistryReady(); const name = normalizeServerName(server.name); - const existing = await this.mcpRegistry.get(name).catch(() => undefined); + const existing = await this.mcpRegistry.get(name, { cwd }).catch(() => undefined); if (existing === undefined) { // Preserve the store's not-found error (and its config validation). await this.globalMcpConfig.update({ ...server, name }); @@ -899,19 +898,19 @@ export class KimiCore implements PromisableMethods { await this.globalMcpConfig.update({ ...server, name }); await this.reconcileMcpServerInSessions([name], 'global-update'); } - return this.listGlobalMcpServers({}); + return this.listGlobalMcpServers({ cwd }); } async removeGlobalMcpServer( - { name }: GlobalMcpServerNamePayload, + { name, cwd }: GlobalMcpServerNamePayload, ): Promise { await this.awaitMcpRegistryReady(); const normalized = normalizeServerName(name); - const existing = await this.mcpRegistry.get(normalized).catch(() => undefined); + const existing = await this.mcpRegistry.get(normalized, { cwd }).catch(() => undefined); if (existing !== undefined) this.throwReadOnlyMcpServer(existing); await this.globalMcpConfig.remove(normalized); await this.reconcileMcpServerInSessions([normalized], 'global-remove'); - return this.listGlobalMcpServers({}); + return this.listGlobalMcpServers({ cwd }); } private throwReadOnlyMcpServer(entry: McpRegistryEntry): void { @@ -1016,15 +1015,16 @@ export class KimiCore implements PromisableMethods { } async beginGlobalMcpServerAuth( - { name }: GlobalMcpServerNamePayload, + { name, cwd }: GlobalMcpServerNamePayload, ): Promise { - return this.beginAppMcpServerAuth(await this.resolveLegacyNamedAppMcpServer(name)); + return this.beginAppMcpServerAuth(await this.resolveLegacyNamedAppMcpServer(name, cwd)); } async beginMcpServerAuth({ locator, + cwd, }: McpServerLocatorPayload): Promise { - return this.beginAppMcpServerAuth(await this.resolveAppMcpServer(locator)); + return this.beginAppMcpServerAuth(await this.resolveAppMcpServer(locator, cwd)); } private async beginAppMcpServerAuth( @@ -1086,14 +1086,14 @@ export class KimiCore implements PromisableMethods { await active.flow.cancel(); } - async resetGlobalMcpServerAuth({ name }: GlobalMcpServerNamePayload): Promise { + async resetGlobalMcpServerAuth({ name, cwd }: GlobalMcpServerNamePayload): Promise { // The legacy name-based surface resolves through the registry too, so a // plugin runtime name works here as well. - await this.appMcpServerDescriptorReset(await this.resolveLegacyNamedAppMcpServer(name)); + await this.appMcpServerDescriptorReset(await this.resolveLegacyNamedAppMcpServer(name, cwd)); } - async resetMcpServerAuth({ locator }: McpServerLocatorPayload): Promise { - await this.appMcpServerDescriptorReset(await this.resolveAppMcpServer(locator)); + async resetMcpServerAuth({ locator, cwd }: McpServerLocatorPayload): Promise { + await this.appMcpServerDescriptorReset(await this.resolveAppMcpServer(locator, cwd)); } private async appMcpServerDescriptorReset( @@ -1107,17 +1107,20 @@ export class KimiCore implements PromisableMethods { async inspectAppMcpServers({ targets, + cwd, }: InspectAppMcpServersPayload): Promise { - const catalog = await this.appMcpServerDescriptors(); + const catalog = await this.appMcpServerDescriptors(cwd); const descriptors = selectAppMcpServerDescriptors(catalog, targets); const inspections = await this.inspectAppMcpServerDescriptors(descriptors, catalog); return inspections.map(sanitizeAppMcpServerInspection); } /** The registry catalog in the locator-addressed shape, with full configs. */ - private async appMcpServerDescriptors(): Promise { + private async appMcpServerDescriptors( + cwd?: string, + ): Promise { await this.awaitMcpRegistryReady(); - return (await this.mcpRegistry.list()).map((entry) => this.appMcpServerDescriptor(entry)); + return (await this.mcpRegistry.list({ cwd })).map((entry) => this.appMcpServerDescriptor(entry)); } private appMcpServerDescriptor(entry: McpRegistryEntry): AppMcpServerRuntimeDescriptor { @@ -1142,8 +1145,9 @@ export class KimiCore implements PromisableMethods { private async resolveAppMcpServer( locator: McpServerLocator, + cwd?: string, ): Promise { - const catalog = await this.appMcpServerDescriptors(); + const catalog = await this.appMcpServerDescriptors(cwd); const server = selectAppMcpServerDescriptors(catalog, [locator])[0]!; this.requireUnambiguousRuntimeName(catalog, server); return server; @@ -1157,11 +1161,12 @@ export class KimiCore implements PromisableMethods { */ private async resolveLegacyNamedAppMcpServer( name: string, + cwd?: string, ): Promise { await this.awaitMcpRegistryReady(); // get() first, preserving its not-found error for unknown names. - await this.mcpRegistry.get(name); - const catalog = await this.appMcpServerDescriptors(); + await this.mcpRegistry.get(name, { cwd }); + const catalog = await this.appMcpServerDescriptors(cwd); const matches = catalog.filter((candidate) => candidate.runtimeName === name); // The sole enabled owner wins over disabled shadows (matching the runtime // and the connection-test path); ambiguity is then judged among the @@ -1357,7 +1362,7 @@ export class KimiCore implements PromisableMethods { private async mcpServerAuthState( entry: McpRegistryEntry, cwd: string | undefined, - verify: boolean, + verify: boolean | undefined, ): Promise { const server = entry.config; // A disabled server never participates in OAuth; keep the historical @@ -1389,11 +1394,12 @@ export class KimiCore implements PromisableMethods { return offline(); }); - if (verify) { + if (verify === true) { // Online verification: a real connection probe settles states the // offline view cannot distinguish (revoked grant, dead refresh token). return probe(); } + if (verify === false) return offline(); if (tokens.hasTokens) return offline(); if (server.auth === 'oauth') return 'oauth-required'; // Unpinned auth with no stored grant: probe once to detect whether the diff --git a/packages/agent-core/test/errors/serialize.test.ts b/packages/agent-core/test/errors/serialize.test.ts index db095b0613..3b74df1a3f 100644 --- a/packages/agent-core/test/errors/serialize.test.ts +++ b/packages/agent-core/test/errors/serialize.test.ts @@ -1,6 +1,8 @@ import { APIProviderQuotaExhaustedError, APIStatusError } from '@moonshot-ai/kosong'; import { describe, expect, it } from 'vitest'; +import { KimiError } from '#/errors/classes'; +import { ErrorCodes, isKimiErrorCode, KIMI_ERROR_INFO } from '#/errors/codes'; import { toKimiErrorPayload } from '#/errors/serialize'; const NGINX_413_HTML = @@ -66,3 +68,31 @@ describe('toKimiErrorPayload — quota-exhausted 429', () => { expect(payload.details).toMatchObject({ statusCode: 429, requestId: 'req-quota' }); }); }); + +describe('toKimiErrorPayload — mcp.oauth_failed registry entry', () => { + it('serializes a KimiError stamped with the engine OAuth failure code', () => { + // The v2 engine raises `mcp.oauth_failed` from its MCP OAuth service; an + // unregistered code would throw on the KIMI_ERROR_INFO lookup here. + const payload = toKimiErrorPayload( + new KimiError(ErrorCodes.MCP_OAUTH_FAILED, 'OAuth flow timed out'), + ); + expect(payload).toMatchObject({ + code: 'mcp.oauth_failed', + message: 'OAuth flow timed out', + retryable: KIMI_ERROR_INFO['mcp.oauth_failed'].retryable, + }); + }); +}); + +describe('isKimiErrorCode', () => { + it('accepts registered codes and rejects unknown or inherited property names', () => { + expect(isKimiErrorCode('mcp.oauth_failed')).toBe(true); + expect(isKimiErrorCode('mcp.future_code')).toBe(false); + // `in` would walk the prototype chain and admit these. + expect(isKimiErrorCode('constructor')).toBe(false); + expect(isKimiErrorCode('toString')).toBe(false); + expect(isKimiErrorCode('__proto__')).toBe(false); + expect(isKimiErrorCode(undefined)).toBe(false); + expect(isKimiErrorCode(42)).toBe(false); + }); +}); diff --git a/packages/kap-server/AGENTS.md b/packages/kap-server/AGENTS.md index d9e14c5d0a..990c0e027d 100644 --- a/packages/kap-server/AGENTS.md +++ b/packages/kap-server/AGENTS.md @@ -15,6 +15,8 @@ No comments — no file headers, no section banners, no statement-level narratio `GET /api/v2/sessions` (`src/routes/v2/sessions.ts`, mounted by `src/routes/registerApiV2Routes.ts`) is the first endpoint of the v2 API. The v2 surface shares v1's wire conventions: every response is wrapped in the `{ code, msg, data, request_id }` envelope with the business outcome in `code` (`40001` invalid query params with `details`, `40922` page_token mismatch), and the HTTP status only reports server-/transport-level outcomes (401 from the global auth hook, `50001` via the catch-all error hook). Pagination is an opaque `page_token` (base64url JSON: version + sha256 query-condition fingerprint + keyset position) — any condition flip mid-pagination fails 40922. Response domains are grouped (`workspace` / `meta` / `activity` always; `git` opt-in via `include=git`, deduped per unique cwd with a 60s TTL cache over `IGitService`, all git/gh failures degrading to cached null fields). Sorts/filters are applied at the edge over the index's canonical `updatedAt desc, id desc` drain, so all three sort orders share one comparator + cursor encoding; `activity.status` maps the core `ISessionActivityView` facts (pending interaction > active turn > failed last turn > idle; cold sessions are always `idle`). +`/api/v2/mcp/*` (`src/routes/v2/mcp.ts`) exposes the agent-core-v2 `mcpManagement` plane: CRUD on the user-level `mcp.json` (`GET/POST/PUT/DELETE /mcp/servers[/{name}]`; `PUT` takes a name-less config body, the path owns the identity), a connection-test probe and the locator-addressed inspection catalog (`POST /mcp/servers:test` / `:inspect`, declared with the doubled-colon static-segment convention), the auth-status surface (`GET /mcp/auth-statuses?verify=`), and the locator-addressed OAuth flow operations (`POST /mcp/auth:begin|complete|cancel|reset`). engine `Error2`s map `mcp.server_not_found` → `40408`, `request.invalid` / `config.invalid` → `40001`, and `mcp.oauth_failed` → `40929`. `auth:complete` is a long poll (up to the flow's `timeoutMs`): the server disables Node's default `requestTimeout` (`start.ts`), and the handler aborts the engine wait when the client connection closes early. The klient facade mirrors the same surface as `global.mcp.*` with identical wire codes, plus the name-only `global.mcp.resolveByName` helper (REST clients compose locators from the `GET /mcp/servers` catalog instead). + ## Transcript surface Implements the op-batch sequencing contract: diff --git a/packages/kap-server/src/protocol/error-codes.ts b/packages/kap-server/src/protocol/error-codes.ts index b154f1cb7d..8238385a86 100644 --- a/packages/kap-server/src/protocol/error-codes.ts +++ b/packages/kap-server/src/protocol/error-codes.ts @@ -61,6 +61,7 @@ export const ErrorCode = { CAPABILITY_UNSUPPORTED: 40925, RUNTIME_UNAVAILABLE: 40926, PROMPT_ID_CONFLICT: 40927, + MCP_OAUTH_FAILED: 40929, APPROVAL_EXPIRED: 41001, QUESTION_EXPIRED: 41002, diff --git a/packages/kap-server/src/routes/registerApiV2Routes.ts b/packages/kap-server/src/routes/registerApiV2Routes.ts index 9095c52cfa..149d9c3b57 100644 --- a/packages/kap-server/src/routes/registerApiV2Routes.ts +++ b/packages/kap-server/src/routes/registerApiV2Routes.ts @@ -1,5 +1,6 @@ import type { Scope } from '@moonshot-ai/agent-core-v2'; +import { registerV2McpRoutes } from './v2/mcp'; import { registerV2SessionsRoutes } from './v2/sessions'; interface ApiV2AppHost { @@ -13,6 +14,7 @@ export async function registerApiV2Routes(app: ApiV2AppHost, core: Scope): Promi await app.register( async (apiV2) => { registerV2SessionsRoutes(apiV2 as Parameters[0], core); + registerV2McpRoutes(apiV2 as Parameters[0], core); }, { prefix: '/api/v2' }, ); diff --git a/packages/kap-server/src/routes/v2/mcp.ts b/packages/kap-server/src/routes/v2/mcp.ts new file mode 100644 index 0000000000..9976e774cb --- /dev/null +++ b/packages/kap-server/src/routes/v2/mcp.ts @@ -0,0 +1,548 @@ +import type { ServerResponse } from 'node:http'; + +import { + ErrorCodes, + IMcpManagementService, + isError2, + type Scope, +} from '@moonshot-ai/agent-core-v2'; +import { + McpServerHttpConfigSchema, + McpServerSseConfigSchema, + McpServerStdioConfigSchema, +} from '@moonshot-ai/agent-core-v2/mcpCore/config-schema'; +import { z } from 'zod'; + +import { defineRoute } from '../../middleware/defineRoute'; +import { errEnvelope, okEnvelope } from '../../protocol/envelope'; +import { ErrorCode } from '../../protocol/error-codes'; + +interface V2McpRouteHost { + get( + path: string, + options: { preHandler: unknown[]; schema?: Record }, + handler: ( + req: { id: string; query: unknown; params: unknown }, + reply: { send(payload: unknown): unknown }, + ) => Promise | void, + ): unknown; + post( + path: string, + options: { preHandler: unknown[]; schema?: Record }, + handler: ( + req: { id: string; body: unknown; query: unknown; params: unknown }, + reply: { send(payload: unknown): unknown }, + ) => Promise | void, + ): unknown; + put( + path: string, + options: { preHandler: unknown[]; schema?: Record }, + handler: ( + req: { id: string; body: unknown; query: unknown; params: unknown }, + reply: { send(payload: unknown): unknown }, + ) => Promise | void, + ): unknown; + delete( + path: string, + options: { preHandler: unknown[]; schema?: Record }, + handler: ( + req: { id: string; query: unknown; params: unknown }, + reply: { send(payload: unknown): unknown }, + ) => Promise | void, + ): unknown; +} + +const serverNameSchema = z.string().min(1); + +const serverNameParamSchema = z.object({ name: serverNameSchema }); + +const serverScopedQuerySchema = z.object({ cwd: z.string().min(1).optional() }); + +const authStatusesQuerySchema = z.object({ + cwd: z.string().min(1).optional(), + verify: z.enum(['true', 'false']).optional(), +}); + +const globalMcpServerConfigSchema = z.discriminatedUnion('transport', [ + McpServerStdioConfigSchema.extend({ name: serverNameSchema }), + McpServerHttpConfigSchema.extend({ name: serverNameSchema }), + McpServerSseConfigSchema.extend({ name: serverNameSchema }), +]); + +const mcpServerConfigBodySchema = z.discriminatedUnion('transport', [ + McpServerStdioConfigSchema, + McpServerHttpConfigSchema, + McpServerSseConfigSchema, +]); + +const testServerBodySchema = z.object({ + name: serverNameSchema.optional(), + server: globalMcpServerConfigSchema.optional(), + cwd: z.string().min(1).optional(), +}); + +const mcpServerLocatorSchema = z.discriminatedUnion('source', [ + z.object({ source: z.literal('global'), name: serverNameSchema }), + z.object({ + source: z.literal('plugin'), + pluginId: z.string().min(1), + serverName: z.string().min(1), + }), +]); + +const inspectServersBodySchema = z.object({ + targets: z.array(mcpServerLocatorSchema).optional(), + cwd: z.string().min(1).optional(), +}); + +const authCompleteBodySchema = z.object({ + flowId: z.string().min(1), + timeoutMs: z + .number() + .int() + .min(1) + .max(2 ** 31 - 1) + .optional(), +}); + +const authCancelBodySchema = z.object({ flowId: z.string().min(1) }); + +const mcpServerSourceSchema = z.enum(['global', 'plugin', 'caller']); + +const mcpServerAuthStateSchema = z.enum([ + 'not-applicable', + 'bearer-token', + 'oauth-required', + 'oauth-authorized', + 'oauth-expired', + 'unavailable', +]); + +const mcpServerConfigDataSchema = z.union([ + McpServerStdioConfigSchema.extend({ envKeys: z.array(z.string()).optional() }), + McpServerHttpConfigSchema.extend({ headerKeys: z.array(z.string()).optional() }), + McpServerSseConfigSchema.extend({ headerKeys: z.array(z.string()).optional() }), +]); + +const mcpManagedServerSchema = z.object({ + name: z.string(), + config: mcpServerConfigDataSchema, + source: mcpServerSourceSchema, + origin: z.string(), + mutable: z.boolean(), + plugin: z.object({ id: z.string(), name: z.string() }).optional(), +}); + +const mcpServerTestResultSchema = z.object({ + success: z.boolean(), + output: z.string(), +}); + +const mcpServerAuthStatusSchema = z.object({ + name: z.string(), + authStatus: mcpServerAuthStateSchema, +}); + +const mcpServerInspectionSchema = z.object({ + serverId: z.string(), + locator: mcpServerLocatorSchema, + runtimeName: z.string(), + canonicalUrl: z.string().optional(), + origin: mcpServerSourceSchema, + config: mcpServerConfigDataSchema, + enabled: z.boolean(), + editable: z.boolean(), + authStatus: mcpServerAuthStateSchema, + checkedAt: z.number().optional(), + error: z.string().optional(), +}); + +const mcpServerAuthBeginResultSchema = z.union([ + z.object({ + status: z.literal('authorization-required'), + flowId: z.string(), + authorizationUrl: z.string(), + }), + z.object({ status: z.literal('already-authorized') }), +]); + +const detailsSchema = z.array(z.object({ path: z.string(), message: z.string() })); + +const baseErrorSchemas = { + [ErrorCode.VALIDATION_FAILED]: { detailsSchema }, +}; + +const namedServerErrorSchemas = { + ...baseErrorSchemas, + [ErrorCode.MCP_SERVER_NOT_FOUND]: {}, +}; + +const oauthErrorSchemas = { + ...baseErrorSchemas, + [ErrorCode.MCP_OAUTH_FAILED]: {}, +}; + +const namedServerOAuthErrorSchemas = { + ...namedServerErrorSchemas, + [ErrorCode.MCP_OAUTH_FAILED]: {}, +}; + +function sendMappedError( + reply: { send(payload: unknown): unknown }, + requestId: string, + err: unknown, +): void { + if (isError2(err)) { + switch (err.code) { + case ErrorCodes.MCP_SERVER_NOT_FOUND: + reply.send(errEnvelope(ErrorCode.MCP_SERVER_NOT_FOUND, err.message, requestId, err.stack)); + return; + case ErrorCodes.REQUEST_INVALID: + case ErrorCodes.CONFIG_INVALID: + reply.send(errEnvelope(ErrorCode.VALIDATION_FAILED, err.message, requestId, err.stack)); + return; + case ErrorCodes.MCP_OAUTH_FAILED: + reply.send(errEnvelope(ErrorCode.MCP_OAUTH_FAILED, err.message, requestId, err.stack)); + return; + } + } + throw err; +} + +export function registerV2McpRoutes(app: V2McpRouteHost, core: Scope): void { + const management = (): IMcpManagementService => core.accessor.get(IMcpManagementService); + + const listServersRoute = defineRoute( + { + method: 'GET', + path: '/mcp/servers', + querystring: serverScopedQuerySchema, + success: { data: z.array(mcpManagedServerSchema) }, + errors: baseErrorSchemas, + description: + 'List every MCP server the management plane knows about (user-level file, plugin manifests; project layers join when `cwd` is given). Read-only entries carry redacted configs.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + const servers = await management().listServers({ cwd: req.query.cwd }); + reply.send(okEnvelope(servers, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.get( + listServersRoute.path, + (listServersRoute.options), + listServersRoute.handler as Parameters[2], + ); + + const getServerRoute = defineRoute( + { + method: 'GET', + path: '/mcp/servers/{name}', + params: serverNameParamSchema, + querystring: serverScopedQuerySchema, + success: { data: mcpManagedServerSchema }, + errors: namedServerErrorSchemas, + description: 'Get one MCP server by runtime name (`40408` when unknown).', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + const server = await management().getServer(req.params.name, { cwd: req.query.cwd }); + reply.send(okEnvelope(server, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.get( + getServerRoute.path, + (getServerRoute.options), + getServerRoute.handler as Parameters[2], + ); + + const addServerRoute = defineRoute( + { + method: 'POST', + path: '/mcp/servers', + querystring: serverScopedQuerySchema, + body: globalMcpServerConfigSchema, + success: { data: z.array(mcpManagedServerSchema) }, + errors: baseErrorSchemas, + description: + 'Add a server to the user-level `mcp.json`; a same-named read-only entry (plugin / project layer) is rejected. Returns the refreshed list.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + const servers = await management().addServer(req.body, { cwd: req.query.cwd }); + reply.send(okEnvelope(servers, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.post( + addServerRoute.path, + (addServerRoute.options), + addServerRoute.handler as Parameters[2], + ); + + const updateServerRoute = defineRoute( + { + method: 'PUT', + path: '/mcp/servers/{name}', + params: serverNameParamSchema, + querystring: serverScopedQuerySchema, + body: mcpServerConfigBodySchema, + success: { data: z.array(mcpManagedServerSchema) }, + errors: namedServerErrorSchemas, + description: + 'Replace the user-level entry named in the path (the body carries no `name`); read-only entries reject the write. Returns the refreshed list.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + const servers = await management().updateServer( + { ...req.body, name: req.params.name }, + { cwd: req.query.cwd }, + ); + reply.send(okEnvelope(servers, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.put( + updateServerRoute.path, + (updateServerRoute.options), + updateServerRoute.handler as Parameters[2], + ); + + const removeServerRoute = defineRoute( + { + method: 'DELETE', + path: '/mcp/servers/{name}', + params: serverNameParamSchema, + querystring: serverScopedQuerySchema, + success: { data: z.array(mcpManagedServerSchema) }, + errors: namedServerErrorSchemas, + description: + 'Remove a user-level entry; read-only entries reject the delete. Returns the refreshed list.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + const servers = await management().removeServer(req.params.name, { cwd: req.query.cwd }); + reply.send(okEnvelope(servers, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.delete( + removeServerRoute.path, + (removeServerRoute.options), + removeServerRoute.handler as Parameters[2], + ); + + const testServerRoute = defineRoute( + { + method: 'POST', + path: '/mcp/servers::test', + body: testServerBodySchema, + success: { data: mcpServerTestResultSchema }, + errors: namedServerErrorSchemas, + description: + 'Probe a real connection to one server: pass `name` to test a registry entry (plugin and project layers included) or an inline `server` config to probe it as-is. Never persists anything.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + const result = await management().testServer(req.body); + reply.send(okEnvelope(result, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.post( + testServerRoute.path, + (testServerRoute.options), + testServerRoute.handler as Parameters[2], + ); + + const inspectServersRoute = defineRoute( + { + method: 'POST', + path: '/mcp/servers::inspect', + body: inspectServersBodySchema, + success: { data: z.array(mcpServerInspectionSchema) }, + errors: namedServerErrorSchemas, + description: + 'The locator-addressed catalog (redacted configs) plus a batched real-connection probe of every OAuth candidate. `targets` narrows the catalog; omitted inspects all. `cwd` includes trusted project layers.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + const inspections = await management().inspectServers(req.body.targets, { + cwd: req.body.cwd, + }); + reply.send(okEnvelope(inspections, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.post( + inspectServersRoute.path, + (inspectServersRoute.options), + inspectServersRoute.handler as Parameters[2], + ); + + const authStatusesRoute = defineRoute( + { + method: 'GET', + path: '/mcp/auth-statuses', + querystring: authStatusesQuerySchema, + success: { data: z.array(mcpServerAuthStatusSchema) }, + errors: baseErrorSchemas, + description: + 'Per-server OAuth state over the registry catalog. Omitted `verify` preserves implicit OAuth detection; `verify=false` is fully offline; `verify=true` verifies every candidate. Probes may refresh or invalidate credentials.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + const statuses = await management().listAuthStatuses({ + cwd: req.query.cwd, + verify: req.query.verify === undefined ? undefined : req.query.verify === 'true', + }); + reply.send(okEnvelope(statuses, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.get( + authStatusesRoute.path, + (authStatusesRoute.options), + authStatusesRoute.handler as Parameters[2], + ); + + const authBeginRoute = defineRoute( + { + method: 'POST', + path: '/mcp/auth::begin', + body: mcpServerLocatorSchema, + querystring: serverScopedQuerySchema, + success: { data: mcpServerAuthBeginResultSchema }, + errors: namedServerOAuthErrorSchemas, + description: + 'Begin an interactive OAuth flow for a remote server. Answers `authorization-required` with the flow handle + URL, or `already-authorized` when a grant exists.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + const result = await management().beginServerAuth(req.body, { cwd: req.query.cwd }); + reply.send(okEnvelope(result, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.post( + authBeginRoute.path, + (authBeginRoute.options), + authBeginRoute.handler as Parameters[2], + ); + + const authCompleteRoute = defineRoute( + { + method: 'POST', + path: '/mcp/auth::complete', + body: authCompleteBodySchema, + success: { data: z.null() }, + errors: oauthErrorSchemas, + description: + 'Await the browser callback of a begun flow and finish the code exchange (`40001` for an unknown `flowId`).', + tags: ['v2-mcp'], + }, + async (req, reply) => { + const { raw } = reply as unknown as { raw: ServerResponse }; + const disconnect = new AbortController(); + const onClose = (): void => { + if (raw.writableFinished) return; + disconnect.abort(); + }; + raw.once('close', onClose); + try { + await management().completeServerAuth(req.body, { signal: disconnect.signal }); + reply.send(okEnvelope(null, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } finally { + raw.off('close', onClose); + } + }, + ); + app.post( + authCompleteRoute.path, + (authCompleteRoute.options), + authCompleteRoute.handler as Parameters[2], + ); + + const authCancelRoute = defineRoute( + { + method: 'POST', + path: '/mcp/auth::cancel', + body: authCancelBodySchema, + success: { data: z.null() }, + errors: oauthErrorSchemas, + description: 'Tear down a begun OAuth flow without finishing it; unknown flows are ignored.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + await management().cancelServerAuth(req.body); + reply.send(okEnvelope(null, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.post( + authCancelRoute.path, + (authCancelRoute.options), + authCancelRoute.handler as Parameters[2], + ); + + const authResetRoute = defineRoute( + { + method: 'POST', + path: '/mcp/auth::reset', + body: mcpServerLocatorSchema, + querystring: serverScopedQuerySchema, + success: { data: z.null() }, + errors: namedServerOAuthErrorSchemas, + description: + 'Clear the stored credentials of one server; the invalidation event reaches live sessions.', + tags: ['v2-mcp'], + }, + async (req, reply) => { + try { + await management().resetServerAuth(req.body, { cwd: req.query.cwd }); + reply.send(okEnvelope(null, req.id)); + } catch (err) { + sendMappedError(reply, req.id, err); + } + }, + ); + app.post( + authResetRoute.path, + (authResetRoute.options), + authResetRoute.handler as Parameters[2], + ); +} diff --git a/packages/kap-server/src/start.ts b/packages/kap-server/src/start.ts index 2e6c533e11..63f7e39d6e 100644 --- a/packages/kap-server/src/start.ts +++ b/packages/kap-server/src/start.ts @@ -7,6 +7,7 @@ import { CapabilityChanged, IConfigService, IEventService, + IMcpOAuthService, IProviderDiscoveryService, ISessionIndex, ISessionIndexMirror, @@ -301,6 +302,7 @@ export async function startServer(opts: ServerStartOptions): Promise resolveRequestId(req.headers), }) as unknown as FastifyInstance; + app.server.requestTimeout = 0; registerRequestLogging(app); app.setValidatorCompiler(() => () => true); app.setSerializerCompiler(() => (data) => JSON.stringify(data)); @@ -346,6 +348,7 @@ export async function startServer(opts: ServerStartOptions): Promise matches the documented v2 route table and meta e "DELETE", "/api/v1/workspaces/{workspace_id}", ], + [ + "DELETE", + "/api/v2/mcp/servers/{name}", + ], [ "GET", "/api/v1/auth", @@ -280,6 +284,18 @@ exports[`API surface snapshot > matches the documented v2 route table and meta e "GET", "/api/v1/workspaces/{workspace_id}/trust", ], + [ + "GET", + "/api/v2/mcp/auth-statuses", + ], + [ + "GET", + "/api/v2/mcp/servers", + ], + [ + "GET", + "/api/v2/mcp/servers/{name}", + ], [ "GET", "/api/v2/sessions", @@ -468,6 +484,34 @@ exports[`API surface snapshot > matches the documented v2 route table and meta e "POST", "/api/v1/workspaces/{workspace_id}/untrust", ], + [ + "POST", + "/api/v2/mcp/auth:begin", + ], + [ + "POST", + "/api/v2/mcp/auth:cancel", + ], + [ + "POST", + "/api/v2/mcp/auth:complete", + ], + [ + "POST", + "/api/v2/mcp/auth:reset", + ], + [ + "POST", + "/api/v2/mcp/servers", + ], + [ + "POST", + "/api/v2/mcp/servers:inspect", + ], + [ + "POST", + "/api/v2/mcp/servers:test", + ], [ "POST", "/api/v2/sessions:archive", @@ -480,6 +524,10 @@ exports[`API surface snapshot > matches the documented v2 route table and meta e "PUT", "/api/v1/providers/{provider_id}", ], + [ + "PUT", + "/api/v2/mcp/servers/{name}", + ], ], } `; diff --git a/packages/kap-server/test/helpers/auth.ts b/packages/kap-server/test/helpers/auth.ts index fa1fd62e55..8b7d0b9aa5 100644 --- a/packages/kap-server/test/helpers/auth.ts +++ b/packages/kap-server/test/helpers/auth.ts @@ -6,6 +6,7 @@ interface FetchOptions { readonly method?: string; readonly headers?: HeaderMap; readonly body?: string; + readonly signal?: AbortSignal; } export function bearerToken(server: RunningServer): string { diff --git a/packages/kap-server/test/openapi.test.ts b/packages/kap-server/test/openapi.test.ts index 0702236042..20c9a13388 100644 --- a/packages/kap-server/test/openapi.test.ts +++ b/packages/kap-server/test/openapi.test.ts @@ -116,6 +116,25 @@ describe('server-v2 OpenAPI', () => { const schema = asRecord(json['schema']); expect(Array.isArray(schema['oneOf'])).toBe(true); }); + + it('documents MCP OAuth failures for auth completion', async () => { + const doc = await fetchOpenApi(); + const authCompleteOp = operation(doc, '/api/v2/mcp/auth:complete', 'post'); + const responses = asRecord(authCompleteOp['responses']); + const response = asRecord(responses['200']); + const content = asRecord(response['content']); + const schema = asRecord(asRecord(content['application/json'])['schema']); + const variants = schema['oneOf']; + + expect(Array.isArray(variants)).toBe(true); + expect( + (variants as unknown[]).some((variant) => { + const properties = asRecord(asRecord(variant)['properties']); + const values = asRecord(properties['code'])['enum']; + return Array.isArray(values) && values.includes(40929); + }), + ).toBe(true); + }); }); function asRecord(value: unknown): Record { diff --git a/packages/kap-server/test/v2Mcp.test.ts b/packages/kap-server/test/v2Mcp.test.ts new file mode 100644 index 0000000000..afd989437e --- /dev/null +++ b/packages/kap-server/test/v2Mcp.test.ts @@ -0,0 +1,512 @@ +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; + +import { + Error2, + ErrorCodes, + IMcpManagementService, + type GlobalMcpServerConfig, + type McpManagedServer, + type McpServerInspection, + type McpServerLocator, + type McpServerTestTarget, +} from '@moonshot-ai/agent-core-v2'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +import { type RunningServer, startServer } from '../src/start'; +import { authedFetch } from './helpers/auth'; +import { TEST_HOST_IDENTITY } from './helpers/hostIdentity'; + +interface EnvelopeWire { + code: number; + msg: string; + data: T | null; + request_id: string; + details?: { path: string; message: string }[]; +} + +const STDIO_A: GlobalMcpServerConfig = { + name: 'a', + transport: 'stdio', + command: 'run-a', + args: ['--verbose'], + env: { TOKEN: 'secret' }, +}; + +interface McpStub { + readonly service: IMcpManagementService; + readonly calls: string[]; + readonly state: { + lastUpdate?: GlobalMcpServerConfig; + lastTestTarget?: McpServerTestTarget; + lastResetLocator?: McpServerLocator; + lastInspectCwd?: string; + lastBeginCwd?: string; + lastResetCwd?: string; + verifySeen?: boolean; + mutationCwds: Array; + }; +} + +function makeMcpStub(): McpStub { + const servers = new Map(); + const calls: string[] = []; + const state: McpStub['state'] = { mutationCwds: [] }; + const list = (): McpManagedServer[] => + [...servers.values()].map((server) => { + const { name, ...config } = server; + return { + name, + config, + source: 'global', + origin: '/home/user/.kimi-code/mcp.json', + mutable: true, + }; + }); + const service: IMcpManagementService = { + _serviceBrand: undefined, + listServers: async () => { + calls.push('listServers'); + return list(); + }, + getServer: async (name) => { + calls.push(`getServer:${name}`); + const server = servers.get(name); + if (server === undefined) { + throw new Error2(ErrorCodes.MCP_SERVER_NOT_FOUND, `MCP server "${name}" was not found`); + } + return list().find((entry) => entry.name === name)!; + }, + addServer: async (server, query) => { + calls.push(`addServer:${server.name}`); + state.mutationCwds.push(query?.cwd); + servers.set(server.name, server); + return list(); + }, + updateServer: async (server, query) => { + calls.push(`updateServer:${server.name}`); + state.mutationCwds.push(query?.cwd); + state.lastUpdate = server; + if (!servers.has(server.name)) { + throw new Error2( + ErrorCodes.MCP_SERVER_NOT_FOUND, + `MCP server "${server.name}" was not found`, + ); + } + servers.set(server.name, server); + return list(); + }, + removeServer: async (name, query) => { + calls.push(`removeServer:${name}`); + state.mutationCwds.push(query?.cwd); + servers.delete(name); + return list(); + }, + testServer: async (target) => { + calls.push('testServer'); + state.lastTestTarget = target; + if (target.name === undefined && target.server === undefined) { + throw new Error2( + ErrorCodes.REQUEST_INVALID, + 'Pass an MCP server name or an inline server config', + ); + } + return { success: true, output: 'probe ok' }; + }, + listAuthStatuses: async (query) => { + calls.push('listAuthStatuses'); + state.verifySeen = query?.verify; + return [...servers.keys()].map((name) => ({ + name, + authStatus: 'not-applicable' as const, + })); + }, + inspectServers: async (targets, query) => { + calls.push('inspectServers'); + state.lastInspectCwd = query?.cwd; + const selected = [...servers.values()].filter( + (server) => + targets === undefined || + targets.some((target) => target.source === 'global' && target.name === server.name), + ); + return selected.map((server): McpServerInspection => { + const { name, ...config } = server; + return { + serverId: `global:${name}`, + locator: { source: 'global', name }, + runtimeName: name, + origin: 'global', + config, + enabled: true, + editable: true, + authStatus: 'not-applicable', + checkedAt: 1000, + }; + }); + }, + resolveServerByName: async (name) => ({ source: 'global', name }), + beginServerAuth: async (_locator, query) => { + state.lastBeginCwd = query?.cwd; + return { + status: 'authorization-required', + flowId: 'flow-1', + authorizationUrl: 'https://example.com/oauth/authorize?client=x', + }; + }, + completeServerAuth: async (handle) => { + if (handle.flowId !== 'flow-1') { + throw new Error2(ErrorCodes.REQUEST_INVALID, `Unknown MCP OAuth flow: ${handle.flowId}`); + } + }, + cancelServerAuth: async () => {}, + resetServerAuth: async (locator, query) => { + state.lastResetLocator = locator; + state.lastResetCwd = query?.cwd; + }, + }; + return { service, calls, state }; +} + +describe('server /api/v2/mcp', () => { + let server: RunningServer | undefined; + let home: string | undefined; + let base: string; + + afterEach(async () => { + vi.unstubAllEnvs(); + if (server !== undefined) { + await server.close(); + server = undefined; + } + if (home !== undefined) { + await new Promise((resolve) => setTimeout(resolve, 25)); + await rm(home, { recursive: true, force: true, maxRetries: 5, retryDelay: 50 } as never); + home = undefined; + } + }); + + async function boot(stub: McpStub): Promise { + home = await mkdtemp(join(tmpdir(), 'kimi-server-v2-mcp-')); + server = await startServer({ + hostIdentity: TEST_HOST_IDENTITY, + host: '127.0.0.1', + port: 0, + homeDir: home, + logLevel: 'silent', + seeds: [[IMcpManagementService, stub.service]], + }); + base = `http://127.0.0.1:${server.port}`; + } + + async function call( + method: string, + path: string, + body?: unknown, + ): Promise<{ status: number; body: EnvelopeWire }> { + const res = await authedFetch(server as RunningServer, base, path, { + method, + headers: body === undefined ? undefined : { 'content-type': 'application/json' }, + body: body === undefined ? undefined : JSON.stringify(body), + }); + return { status: res.status, body: (await res.json()) as EnvelopeWire }; + } + + describe('routes', () => { + it('round-trips a server through add/get/update/remove', async () => { + const stub = makeMcpStub(); + await boot(stub); + + const added = await call( + 'POST', + '/api/v2/mcp/servers?cwd=%2Fworkspace%2Fproject', + STDIO_A, + ); + expect(added.status).toBe(200); + expect(added.body.code).toBe(0); + expect(added.body.data).toEqual([ + { + name: 'a', + config: { + transport: 'stdio', + command: 'run-a', + args: ['--verbose'], + env: { TOKEN: 'secret' }, + }, + source: 'global', + origin: '/home/user/.kimi-code/mcp.json', + mutable: true, + }, + ]); + + const got = await call('GET', '/api/v2/mcp/servers/a'); + expect(got.body.code).toBe(0); + expect(got.body.data).toMatchObject({ name: 'a', config: { command: 'run-a' } }); + + const updated = await call( + 'PUT', + '/api/v2/mcp/servers/a?cwd=%2Fworkspace%2Fproject', + { + transport: 'stdio', + command: 'run-b', + }, + ); + expect(updated.body.code).toBe(0); + expect(updated.body.data).toHaveLength(1); + expect(stub.state.lastUpdate).toEqual({ transport: 'stdio', command: 'run-b', name: 'a' }); + + const removed = await call( + 'DELETE', + '/api/v2/mcp/servers/a?cwd=%2Fworkspace%2Fproject', + ); + expect(removed.body.code).toBe(0); + expect(removed.body.data).toEqual([]); + expect(stub.calls).toContain('removeServer:a'); + expect(stub.state.mutationCwds).toEqual([ + '/workspace/project', + '/workspace/project', + '/workspace/project', + ]); + }); + + it('maps an unknown server name to 40408', async () => { + const stub = makeMcpStub(); + await boot(stub); + const got = await call('GET', '/api/v2/mcp/servers/nope'); + expect(got.status).toBe(200); + expect(got.body.code).toBe(40408); + expect(got.body.data).toBeNull(); + expect(got.body.msg).toContain('nope'); + + const updated = await call('PUT', '/api/v2/mcp/servers/nope', { + transport: 'stdio', + command: 'run-x', + }); + expect(updated.body.code).toBe(40408); + }); + + it('rejects malformed bodies with 40001 + details from the zod preHandler', async () => { + const stub = makeMcpStub(); + await boot(stub); + + const badAdd = await call('POST', '/api/v2/mcp/servers', { name: 'a', command: 'run-a' }); + expect(badAdd.body.code).toBe(40001); + expect(Array.isArray(badAdd.body.details)).toBe(true); + + const badBegin = await call('POST', '/api/v2/mcp/auth:begin', { source: 'global' }); + expect(badBegin.body.code).toBe(40001); + + expect(stub.calls).toEqual([]); + }); + + it('maps the engine request.invalid rejection to 40001', async () => { + const stub = makeMcpStub(); + await boot(stub); + const res = await call('POST', '/api/v2/mcp/servers:test', {}); + expect(res.body.code).toBe(40001); + expect(res.body.data).toBeNull(); + expect(stub.calls).toEqual(['testServer']); + }); + + it('maps the engine config.invalid rejection to 40001', async () => { + const stub = makeMcpStub(); + stub.service.addServer = async () => { + throw new Error2( + ErrorCodes.CONFIG_INVALID, + 'Invalid JSON in /home/user/.kimi-code/mcp.json: Unexpected token', + ); + }; + await boot(stub); + + const res = await call('POST', '/api/v2/mcp/servers', STDIO_A); + expect(res.status).toBe(200); + expect(res.body.code).toBe(40001); + expect(res.body.data).toBeNull(); + }); + + it('maps the engine mcp.oauth_failed rejection to 40929', async () => { + const stub = makeMcpStub(); + stub.service.completeServerAuth = async () => { + throw new Error2( + ErrorCodes.MCP_OAUTH_FAILED, + 'OAuth flow for "a" failed: OAuth callback timed out', + ); + }; + await boot(stub); + + const res = await call('POST', '/api/v2/mcp/auth:complete', { flowId: 'flow-1' }); + expect(res.status).toBe(200); + expect(res.body.code).toBe(40929); + expect(res.body.data).toBeNull(); + }); + + it('maps a delete rejected with mcp.server_not_found to 40408', async () => { + const stub = makeMcpStub(); + stub.service.removeServer = async (name) => { + throw new Error2(ErrorCodes.MCP_SERVER_NOT_FOUND, `MCP server "${name}" was not found`); + }; + await boot(stub); + + const res = await call('DELETE', '/api/v2/mcp/servers/nope'); + expect(res.status).toBe(200); + expect(res.body.code).toBe(40408); + expect(res.body.data).toBeNull(); + expect(res.body.msg).toContain('nope'); + }); + + it('tests an inline server config without saving it', async () => { + const stub = makeMcpStub(); + await boot(stub); + const inline = { name: 'inline', transport: 'http', url: 'https://example.com/mcp' }; + const res = await call('POST', '/api/v2/mcp/servers:test', { server: inline }); + expect(res.body).toMatchObject({ code: 0, data: { success: true, output: 'probe ok' } }); + expect(stub.state.lastTestTarget).toEqual({ server: inline }); + }); + + it('inspects the catalog narrowed by locator targets', async () => { + const stub = makeMcpStub(); + await boot(stub); + await call('POST', '/api/v2/mcp/servers', STDIO_A); + + const res = await call('POST', '/api/v2/mcp/servers:inspect', { + targets: [{ source: 'global', name: 'a' }], + }); + expect(res.body.code).toBe(0); + expect(res.body.data).toEqual([ + { + serverId: 'global:a', + locator: { source: 'global', name: 'a' }, + runtimeName: 'a', + origin: 'global', + config: { + transport: 'stdio', + command: 'run-a', + args: ['--verbose'], + env: { TOKEN: 'secret' }, + }, + enabled: true, + editable: true, + authStatus: 'not-applicable', + checkedAt: 1000, + }, + ]); + }); + + it('forwards cwd through locator-addressed inspection and OAuth operations', async () => { + const stub = makeMcpStub(); + await boot(stub); + + await call('POST', '/api/v2/mcp/servers:inspect', { + targets: [], + cwd: '/workspace/project', + }); + await call('POST', '/api/v2/mcp/auth:begin?cwd=%2Fworkspace%2Fproject', { + source: 'global', + name: 'a', + }); + await call('POST', '/api/v2/mcp/auth:reset?cwd=%2Fworkspace%2Fproject', { + source: 'global', + name: 'a', + }); + + expect(stub.state).toMatchObject({ + lastInspectCwd: '/workspace/project', + lastBeginCwd: '/workspace/project', + lastResetCwd: '/workspace/project', + }); + }); + + it('maps ?verify= onto the boolean auth-status query flag', async () => { + const stub = makeMcpStub(); + await boot(stub); + await call('POST', '/api/v2/mcp/servers', STDIO_A); + + const verified = await call('GET', '/api/v2/mcp/auth-statuses?verify=true'); + expect(verified.body).toMatchObject({ + code: 0, + data: [{ name: 'a', authStatus: 'not-applicable' }], + }); + expect(stub.state.verifySeen).toBe(true); + + const offline = await call('GET', '/api/v2/mcp/auth-statuses'); + expect(offline.body.code).toBe(0); + expect(stub.state.verifySeen).toBeUndefined(); + + const bogus = await call('GET', '/api/v2/mcp/auth-statuses?verify=yes'); + expect(bogus.body.code).toBe(40001); + }); + + it('drives the locator-addressed OAuth flow operations', async () => { + const stub = makeMcpStub(); + await boot(stub); + + const begin = await call('POST', '/api/v2/mcp/auth:begin', { source: 'global', name: 'a' }); + expect(begin.body).toMatchObject({ + code: 0, + data: { + status: 'authorization-required', + flowId: 'flow-1', + authorizationUrl: 'https://example.com/oauth/authorize?client=x', + }, + }); + + const complete = await call('POST', '/api/v2/mcp/auth:complete', { flowId: 'flow-1' }); + expect(complete.body).toMatchObject({ code: 0, data: null }); + + const unknownFlow = await call('POST', '/api/v2/mcp/auth:complete', { flowId: 'nope' }); + expect(unknownFlow.body.code).toBe(40001); + + const cancel = await call('POST', '/api/v2/mcp/auth:cancel', { flowId: 'flow-1' }); + expect(cancel.body).toMatchObject({ code: 0, data: null }); + + const reset = await call('POST', '/api/v2/mcp/auth:reset', { + source: 'plugin', + pluginId: 'p', + serverName: 's', + }); + expect(reset.body).toMatchObject({ code: 0, data: null }); + expect(stub.state.lastResetLocator).toEqual({ source: 'plugin', pluginId: 'p', serverName: 's' }); + }); + + it('rejects an overflowing auth:complete timeoutMs with 40001', async () => { + const stub = makeMcpStub(); + await boot(stub); + + const res = await call('POST', '/api/v2/mcp/auth:complete', { + flowId: 'flow-1', + timeoutMs: 2 ** 31, + }); + + expect(res.body.code).toBe(40001); + expect(stub.calls).toEqual([]); + }); + + it('aborts the engine wait when the client disconnects mid-complete', async () => { + const stub = makeMcpStub(); + let seenSignal: AbortSignal | undefined; + let reached = false; + stub.service.completeServerAuth = async (_handle, options) => { + seenSignal = options?.signal; + reached = true; + await new Promise((resolve) => { + options?.signal?.addEventListener('abort', () => resolve(), { once: true }); + }); + }; + await boot(stub); + + const controller = new AbortController(); + const pending = authedFetch(server as RunningServer, base, '/api/v2/mcp/auth:complete', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ flowId: 'flow-1' }), + signal: controller.signal, + }); + await vi.waitFor(() => expect(reached).toBe(true)); + + controller.abort(); + + await expect(pending).rejects.toThrow(); + await vi.waitFor(() => expect(seenSignal?.aborted).toBe(true)); + }); + }); +}); diff --git a/packages/klient/src/contract/global/mcpManagement.ts b/packages/klient/src/contract/global/mcpManagement.ts new file mode 100644 index 0000000000..0279df3e2e --- /dev/null +++ b/packages/klient/src/contract/global/mcpManagement.ts @@ -0,0 +1,189 @@ +/** + * `mcpManagementService` — the unified MCP management plane. Mirrors + * `agent-core-v2/app/mcpManagement/mcpManagement.ts`; `McpServerSource` / + * `McpRegistryPluginOrigin` / `McpRegistryQuery` mirror + * `agent-core-v2/app/mcpRegistry/mcpRegistry.ts`, and the redacted config + * shape mirrors `agent-core-v2/mcpCore/configView.ts`. + */ + +import { z } from 'zod'; + +import { noResult } from '../helpers.js'; +import { + mcpServerHttpConfigSchema, + mcpServerSseConfigSchema, + mcpServerStdioConfigSchema, +} from '../mcp.js'; +import type { ServiceContract } from '../types.js'; + +export const mcpServerSourceSchema = z.enum(['global', 'plugin', 'caller']); + +export const mcpRegistryPluginOriginSchema = z.object({ + id: z.string(), + /** Manifest-local server name (without the runtime prefix). */ + name: z.string(), +}); + +export const mcpRegistryQuerySchema = z.object({ + cwd: z.string().min(1).optional(), +}); + +export const mcpAuthStatusQuerySchema = z.object({ + cwd: z.string().min(1).optional(), + verify: z.boolean().optional(), +}); + +/** `GlobalMcpServerConfig` — a named full config (add/update, inline test target). */ +export const globalMcpServerConfigSchema = z.discriminatedUnion('transport', [ + mcpServerStdioConfigSchema.extend({ name: z.string().min(1) }), + mcpServerHttpConfigSchema.extend({ name: z.string().min(1) }), + mcpServerSseConfigSchema.extend({ name: z.string().min(1) }), +]); + +/** + * The wire config of a managed/inspected server: mutable entries carry the + * full config (edit UIs prefill from it); read-only entries are redacted — + * `env` / `headers` values are replaced by the sorted key lists `envKeys` / + * `headerKeys`. One schema covers both shapes, mirroring the engine's + * `McpServerConfig | McpServerConfigView` union. + */ +export const mcpServerConfigDataSchema = z.discriminatedUnion('transport', [ + mcpServerStdioConfigSchema.extend({ envKeys: z.array(z.string()).optional() }), + mcpServerHttpConfigSchema.extend({ headerKeys: z.array(z.string()).optional() }), + mcpServerSseConfigSchema.extend({ headerKeys: z.array(z.string()).optional() }), +]); + +export const mcpManagedServerSchema = z.object({ + name: z.string(), + config: mcpServerConfigDataSchema, + source: mcpServerSourceSchema, + origin: z.string(), + mutable: z.boolean(), + plugin: mcpRegistryPluginOriginSchema.optional(), +}); + +export const mcpServerTestTargetSchema = z.object({ + name: z.string().min(1).optional(), + server: globalMcpServerConfigSchema.optional(), + cwd: z.string().min(1).optional(), +}); + +export const mcpServerTestResultSchema = z.object({ + success: z.boolean(), + output: z.string(), +}); + +export const mcpServerLocatorSchema = z.discriminatedUnion('source', [ + z.object({ source: z.literal('global'), name: z.string().min(1) }), + z.object({ + source: z.literal('plugin'), + pluginId: z.string().min(1), + serverName: z.string().min(1), + }), +]); + +export const mcpServerAuthStateSchema = z.enum([ + 'not-applicable', + 'bearer-token', + 'oauth-required', + 'oauth-authorized', + 'oauth-expired', + 'unavailable', +]); + +export const mcpServerDescriptorSchema = z.object({ + /** `global:` / `plugin::`, URL-encoded. */ + serverId: z.string(), + locator: mcpServerLocatorSchema, + runtimeName: z.string(), + canonicalUrl: z.string().optional(), + origin: mcpServerSourceSchema, + config: mcpServerConfigDataSchema, + enabled: z.boolean(), + editable: z.boolean(), +}); + +export const mcpServerInspectionSchema = mcpServerDescriptorSchema.extend({ + authStatus: mcpServerAuthStateSchema, + checkedAt: z.number().optional(), + error: z.string().optional(), +}); + +export const mcpServerAuthStatusSchema = z.object({ + name: z.string(), + authStatus: mcpServerAuthStateSchema, +}); + +export const mcpServerAuthBeginResultSchema = z.discriminatedUnion('status', [ + z.object({ + status: z.literal('authorization-required'), + flowId: z.string(), + authorizationUrl: z.string(), + }), + z.object({ status: z.literal('already-authorized') }), +]); + +export const mcpServerAuthFlowHandleSchema = z.object({ + flowId: z.string().min(1), + // Node overflows setTimeout delays above 2^31-1 into ~1ms; the REST schema + // and the engine reject the same range. + timeoutMs: z.number().int().min(1).max(2 ** 31 - 1).optional(), +}); + +export const mcpManagementContract = { + listServers: { + input: z.tuple([mcpRegistryQuerySchema.optional()]), + output: z.array(mcpManagedServerSchema), + }, + getServer: { + input: z.tuple([z.string().min(1), mcpRegistryQuerySchema.optional()]), + output: mcpManagedServerSchema, + }, + addServer: { + input: z.tuple([globalMcpServerConfigSchema, mcpRegistryQuerySchema.optional()]), + output: z.array(mcpManagedServerSchema), + }, + updateServer: { + input: z.tuple([globalMcpServerConfigSchema, mcpRegistryQuerySchema.optional()]), + output: z.array(mcpManagedServerSchema), + }, + removeServer: { + input: z.tuple([z.string().min(1), mcpRegistryQuerySchema.optional()]), + output: z.array(mcpManagedServerSchema), + }, + testServer: { + input: z.tuple([mcpServerTestTargetSchema]), + output: mcpServerTestResultSchema, + }, + listAuthStatuses: { + input: z.tuple([mcpAuthStatusQuerySchema.optional()]), + output: z.array(mcpServerAuthStatusSchema), + }, + inspectServers: { + input: z.tuple([ + z.array(mcpServerLocatorSchema).optional(), + mcpRegistryQuerySchema.optional(), + ]), + output: z.array(mcpServerInspectionSchema), + }, + resolveServerByName: { + input: z.tuple([z.string().min(1), mcpRegistryQuerySchema.optional()]), + output: mcpServerLocatorSchema, + }, + beginServerAuth: { + input: z.tuple([mcpServerLocatorSchema, mcpRegistryQuerySchema.optional()]), + output: mcpServerAuthBeginResultSchema, + }, + completeServerAuth: { + input: z.tuple([mcpServerAuthFlowHandleSchema]), + output: noResult, + }, + cancelServerAuth: { + input: z.tuple([z.object({ flowId: z.string().min(1) })]), + output: noResult, + }, + resetServerAuth: { + input: z.tuple([mcpServerLocatorSchema, mcpRegistryQuerySchema.optional()]), + output: noResult, + }, +} satisfies ServiceContract; diff --git a/packages/klient/src/contract/index.ts b/packages/klient/src/contract/index.ts index 4155f4b48a..d16cfaf228 100644 --- a/packages/klient/src/contract/index.ts +++ b/packages/klient/src/contract/index.ts @@ -35,6 +35,7 @@ import { filesContract } from './global/files.js'; import { flagsContract } from './global/flags.js'; import { hostFsContract } from './global/hostFs.js'; import { modelsContract } from './global/models.js'; +import { mcpManagementContract } from './global/mcpManagement.js'; import { pluginsContract } from './global/plugins.js'; import { providersContract } from './global/providers.js'; import { sessionsContract } from './global/sessions.js'; @@ -64,6 +65,7 @@ export const globalContract: KlientContract = { hostFolderBrowser: hostFsContract, bootstrapService: envContract, fileService: filesContract, + mcpManagementService: mcpManagementContract, sessionManager: sessionManagerContract, // session scope sessionMetadata: sessionMetadataContract, diff --git a/packages/klient/src/contract/mcp.ts b/packages/klient/src/contract/mcp.ts index 13960df889..ed8827a7d2 100644 --- a/packages/klient/src/contract/mcp.ts +++ b/packages/klient/src/contract/mcp.ts @@ -1,6 +1,8 @@ /** * Shared MCP server wire schema for session creation and plugin manifests. - * Mirrors `agent-core-v2/mcpCore/config-schema.ts`. + * Mirrors `agent-core-v2/mcpCore/config-schema.ts`. Unlike the config files, + * the wire requires the explicit `transport` discriminant (the engine's + * command/url inference preprocess is a file-format convenience). */ import { z } from 'zod'; @@ -17,31 +19,39 @@ const mcpServerCommonFields = { disabledTools: z.array(z.string()).optional(), } as const; +export const mcpServerStdioConfigSchema = z.object({ + transport: z.literal('stdio'), + runtime_id: z.string().min(1).optional(), + command: z.string().min(1), + args: z.array(z.string()).optional(), + env: stringRecordSchema.optional(), + cwd: z.string().optional(), + executor: z.enum(['local', 'kaos']).optional(), + ...mcpServerCommonFields, +}); + +export const mcpServerHttpConfigSchema = z.object({ + transport: z.literal('http'), + url: z.string().url(), + headers: stringRecordSchema.optional(), + auth: z.literal('oauth').optional(), + bearerTokenEnvVar: z.string().min(1).optional(), + ...mcpServerCommonFields, +}); + +export const mcpServerSseConfigSchema = z.object({ + transport: z.literal('sse'), + url: z.string().url(), + headers: stringRecordSchema.optional(), + auth: z.literal('oauth').optional(), + bearerTokenEnvVar: z.string().min(1).optional(), + ...mcpServerCommonFields, +}); + export const mcpServerConfigSchema = z.discriminatedUnion('transport', [ - z.object({ - transport: z.literal('stdio'), - runtime_id: z.string().min(1).optional(), - command: z.string().min(1), - args: z.array(z.string()).optional(), - env: stringRecordSchema.optional(), - cwd: z.string().optional(), - executor: z.enum(['local', 'kaos']).optional(), - ...mcpServerCommonFields, - }), - z.object({ - transport: z.literal('http'), - url: z.string().url(), - headers: stringRecordSchema.optional(), - bearerTokenEnvVar: z.string().min(1).optional(), - ...mcpServerCommonFields, - }), - z.object({ - transport: z.literal('sse'), - url: z.string().url(), - headers: stringRecordSchema.optional(), - bearerTokenEnvVar: z.string().min(1).optional(), - ...mcpServerCommonFields, - }), + mcpServerStdioConfigSchema, + mcpServerHttpConfigSchema, + mcpServerSseConfigSchema, ]); export type McpServerConfig = z.infer; diff --git a/packages/klient/src/core/channel.ts b/packages/klient/src/core/channel.ts index 181290167b..c674ee6b69 100644 --- a/packages/klient/src/core/channel.ts +++ b/packages/klient/src/core/channel.ts @@ -12,6 +12,16 @@ export interface IDisposable { dispose(): void; } +/** Optional per-call knobs a transport may honor. */ +export interface CallOptions { + /** + * Per-call deadline (ms). A transport with a default call timeout (ipc) + * takes it as an override — long-poll calls pass a deadline covering the + * engine-side wait; transports without a timeout (memory) ignore it. + */ + readonly timeoutMs?: number; +} + /** Scope coordinates of a call/subscription. Empty object = core (app) scope. */ export interface ScopeRef { readonly workspaceId?: string; @@ -34,7 +44,13 @@ export type EventSourceRef = export interface KlientChannel { /** Invoke `service.method(...args)` in the given scope; resolves with the raw wire result. */ - call(scope: ScopeRef, service: string, method: string, args: unknown[]): Promise; + call( + scope: ScopeRef, + service: string, + method: string, + args: unknown[], + options?: CallOptions, + ): Promise; /** * Invoke `service.method(...args)` in the given scope and return a streaming * result. The callee must return an `AsyncIterable`; each yielded chunk is diff --git a/packages/klient/src/core/facade/global.ts b/packages/klient/src/core/facade/global.ts index c185812f24..247fd476bd 100644 --- a/packages/klient/src/core/facade/global.ts +++ b/packages/klient/src/core/facade/global.ts @@ -37,6 +37,17 @@ import type { IModelCatalog } from '@moonshot-ai/agent-core-v2/kosong/model/cata import type { IProviderDiscoveryService } from '@moonshot-ai/agent-core-v2/app/kosongConfig/discovery'; import type { McpServerConfig } from '../../contract/mcp.js'; +import type { CallOptions } from '../channel.js'; +import type { + GlobalMcpServerConfig, + McpManagedServer, + McpServerAuthBeginResult, + McpServerAuthStatus, + McpServerInspection, + McpServerLocator, + McpServerTestResult, + McpServerTestTarget, +} from '@moonshot-ai/agent-core-v2/app/mcpManagement/mcpManagement'; import type { AnonymousProviderInput, GenerateEvent, GenerateInput, GenerateParams, ProviderInput } from './kosong-types.js'; import type { PluginCommandDef, @@ -48,7 +59,12 @@ import type { import type { CapabilityStatus } from '@moonshot-ai/agent-core-v2/app/capability/types'; /** Low-level caller the klient factory builds: routes + validates one service call. */ -export type Caller = (service: string, method: string, args: unknown[]) => Promise; +export type Caller = ( + service: string, + method: string, + args: unknown[], + options?: CallOptions, +) => Promise; /** Scoped variant — the factory's real signature; global methods bind the core scope. */ export type ScopedCaller = ( @@ -56,6 +72,7 @@ export type ScopedCaller = ( service: string, method: string, args: unknown[], + options?: CallOptions, ) => Promise; /** Streaming variant of `ScopedCaller` — returns a validated `AsyncIterable`. */ @@ -222,6 +239,50 @@ export interface GlobalHostFsFacade { home(): Promise; } +/** + * The unified MCP management plane (engine `IMcpManagementService`, App + * scope): CRUD on the user-level `mcp.json`, a connection test probe, the + * locator-addressed inspection catalog, the auth-status surface, and the + * locator-addressed OAuth flow operations. + */ +export interface GlobalMcpFacade { + list(input?: { cwd?: string }): Promise; + get(input: { name: string; cwd?: string }): Promise; + /** Add a user-level entry; a same-named read-only entry rejects. Returns the refreshed list. */ + add(input: { + server: GlobalMcpServerConfig; + cwd?: string; + }): Promise; + /** Replace a user-level entry; read-only entries reject. Returns the refreshed list. */ + update(input: { + server: GlobalMcpServerConfig; + cwd?: string; + }): Promise; + /** Remove a user-level entry; read-only entries reject. Returns the refreshed list. */ + remove(input: { name: string; cwd?: string }): Promise; + /** Probe a real connection: a registry `name`, or an inline `server` config as-is. */ + test(input: McpServerTestTarget): Promise; + /** The locator-addressed catalog plus a batched real-connection probe of OAuth candidates. */ + inspect(input?: { + targets?: readonly McpServerLocator[]; + cwd?: string; + }): Promise; + /** Per-server OAuth state; omitted `verify` detects implicit OAuth, `false` stays offline. */ + authStatuses(input?: { + cwd?: string; + verify?: boolean; + }): Promise; + /** Resolve a legacy name-only auth target to its unambiguous locator. */ + resolveByName(input: { name: string; cwd?: string }): Promise; + beginAuth(input: { + locator: McpServerLocator; + cwd?: string; + }): Promise; + completeAuth(input: { flowId: string; timeoutMs?: number }): Promise; + cancelAuth(input: { flowId: string }): Promise; + resetAuth(input: { locator: McpServerLocator; cwd?: string }): Promise; +} + /** One downloaded upload: its metadata plus the buffered bytes. */ export interface FileDownload { readonly meta: FileMeta; @@ -273,6 +334,7 @@ export interface GlobalFacade { readonly capabilities: GlobalCapabilitiesFacade; readonly hostFs: GlobalHostFsFacade; readonly files: GlobalFilesFacade; + readonly mcp: GlobalMcpFacade; env(): Promise; } @@ -296,8 +358,18 @@ const ENV_SCALAR_PROPERTIES = [ 'logsDir', ] as const; +// The IPC transport enforces a per-call deadline (default 30s) that would +// truncate the completeAuth long poll: the engine waits up to +// `DEFAULT_AUTH_TIMEOUT_MS` for the browser callback when the caller omits +// `timeoutMs` (agent-core-v2 `mcpManagementService.ts`), and the +// authorization-code exchange afterwards is itself bounded at 30s per grant +// request (agent-core-v2 `mcpCore/oauth/service.ts`). The per-call deadline +// below covers both, so IPC behaves like the timeout-free memory transport. +const DEFAULT_AUTH_TIMEOUT_MS = 15 * 60_000; +const AUTH_COMPLETION_MARGIN_MS = 30_000; + export function createGlobalFacade(scoped: ScopedCaller, scopedStream: ScopedStreamCaller): GlobalFacade { - const call: Caller = (service, method, args) => scoped({}, service, method, args); + const call: Caller = (service, method, args, options) => scoped({}, service, method, args, options); const streamCall = (service: string, method: string, args: unknown[]) => scopedStream({}, service, method, args); // The bootstrap snapshot is frozen at process start, so the aggregated @@ -508,6 +580,75 @@ export function createGlobalFacade(scoped: ScopedCaller, scopedStream: ScopedStr delete: (fileId) => call('fileService', 'delete', [fileId]) as Promise, }, + mcp: { + list: (input) => + call('mcpManagementService', 'listServers', [ + input === undefined ? undefined : { cwd: input.cwd }, + ]) as Promise, + get: ({ name, cwd }) => + call('mcpManagementService', 'getServer', [ + name, + cwd === undefined ? undefined : { cwd }, + ]) as Promise, + add: ({ server, cwd }) => + call('mcpManagementService', 'addServer', [ + server, + cwd === undefined ? undefined : { cwd }, + ]) as Promise< + readonly McpManagedServer[] + >, + update: ({ server, cwd }) => + call('mcpManagementService', 'updateServer', [ + server, + cwd === undefined ? undefined : { cwd }, + ]) as Promise< + readonly McpManagedServer[] + >, + remove: ({ name, cwd }) => + call('mcpManagementService', 'removeServer', [ + name, + cwd === undefined ? undefined : { cwd }, + ]) as Promise< + readonly McpManagedServer[] + >, + test: (target) => + call('mcpManagementService', 'testServer', [target]) as Promise, + inspect: (input) => + call('mcpManagementService', 'inspectServers', [ + input?.targets, + input === undefined ? undefined : { cwd: input.cwd }, + ]) as Promise< + readonly McpServerInspection[] + >, + authStatuses: (input) => + call('mcpManagementService', 'listAuthStatuses', [ + input === undefined ? undefined : { cwd: input.cwd, verify: input.verify }, + ]) as Promise, + resolveByName: ({ name, cwd }) => + call('mcpManagementService', 'resolveServerByName', [name, { cwd }]) as Promise< + McpServerLocator + >, + beginAuth: ({ locator, cwd }) => + call('mcpManagementService', 'beginServerAuth', [ + locator, + { cwd }, + ]) as Promise, + completeAuth: ({ flowId, timeoutMs }) => + call('mcpManagementService', 'completeServerAuth', [{ flowId, timeoutMs }], { + // Clamp to Node's 32-bit timer ceiling: `timeoutMs` may legally be + // the contract max (2**31 - 1), and adding the margin would + // overflow setTimeout into a ~1ms deadline. + timeoutMs: Math.min( + (timeoutMs ?? DEFAULT_AUTH_TIMEOUT_MS) + AUTH_COMPLETION_MARGIN_MS, + 2 ** 31 - 1, + ), + }) as Promise, + cancelAuth: ({ flowId }) => + call('mcpManagementService', 'cancelServerAuth', [{ flowId }]) as Promise, + resetAuth: ({ locator, cwd }) => + call('mcpManagementService', 'resetServerAuth', [locator, { cwd }]) as Promise, + }, + env, }; } diff --git a/packages/klient/src/core/klient.ts b/packages/klient/src/core/klient.ts index 6eafac7d54..b0312c2c9b 100644 --- a/packages/klient/src/core/klient.ts +++ b/packages/klient/src/core/klient.ts @@ -48,7 +48,7 @@ export function createKlientFromChannel( ): Klient { const validate = options.validate ?? true; - const call: ScopedCaller = async (scope, service, method, args) => { + const call: ScopedCaller = async (scope, service, method, args, options) => { const procedure = globalContract[service]?.[method]; if (procedure === undefined) { // A facade method without a contract entry is a klient bug, not a wire error. @@ -59,7 +59,7 @@ export function createKlientFromChannel( } const name = `${service}.${method}`; const wireArgs = validate ? parseInput(name, procedure, args) : args; - const data = await channel.call(scope, service, method, wireArgs); + const data = await channel.call(scope, service, method, wireArgs, options); return validate ? parseOutput(name, procedure, data) : data; }; diff --git a/packages/klient/src/index.ts b/packages/klient/src/index.ts index 07bfff732d..b59fda687d 100644 --- a/packages/klient/src/index.ts +++ b/packages/klient/src/index.ts @@ -7,6 +7,7 @@ */ export type { + CallOptions, EventSourceRef, IDisposable, KlientChannel, @@ -34,6 +35,7 @@ export type { GlobalFlagsFacade, GlobalHostFsFacade, GlobalKosongFacade, + GlobalMcpFacade, GlobalPluginsFacade, GlobalSessionsFacade, GlobalWorkspacesFacade, @@ -141,5 +143,16 @@ export type { InteractionKind, } from '@moonshot-ai/agent-core-v2/session/interaction/interaction'; export type { SkillSummary } from '@moonshot-ai/agent-core-v2/app/skillCatalog/types'; +export type { + GlobalMcpServerConfig, + McpManagedServer, + McpServerAuthBeginResult, + McpServerAuthState, + McpServerAuthStatus, + McpServerInspection, + McpServerLocator, + McpServerTestResult, + McpServerTestTarget, +} from '@moonshot-ai/agent-core-v2/app/mcpManagement/mcpManagement'; export type { ContentPart } from '@moonshot-ai/agent-core-v2/kosong/contract/message'; export type { PermissionMode } from '@moonshot-ai/agent-core-v2/agent/permissionPolicy/types'; diff --git a/packages/klient/src/transports/ipc/channel.ts b/packages/klient/src/transports/ipc/channel.ts index 8c19937676..0136078dd2 100644 --- a/packages/klient/src/transports/ipc/channel.ts +++ b/packages/klient/src/transports/ipc/channel.ts @@ -10,6 +10,7 @@ import { createConnection, type Socket } from 'node:net'; import type { + CallOptions, EventSourceRef, IDisposable, KlientChannel, @@ -100,17 +101,24 @@ export class IpcChannel implements KlientChannel { }); } - async call(scope: ScopeRef, service: string, method: string, args: unknown[]): Promise { + async call( + scope: ScopeRef, + service: string, + method: string, + args: unknown[], + options?: CallOptions, + ): Promise { await this.ready; if (this.closed) throw new Error('ipc closed'); + const timeoutMs = options?.timeoutMs ?? this.callTimeoutMs; const id = this.nextId(); const promise = new Promise((resolve, reject) => { const timer = - this.callTimeoutMs > 0 + timeoutMs > 0 ? setTimeout(() => { this.pending.delete(id); - reject(new RPCError(50001, `call timed out after ${this.callTimeoutMs}ms`)); - }, this.callTimeoutMs) + reject(new RPCError(50001, `call timed out after ${timeoutMs}ms`)); + }, timeoutMs) : undefined; this.pending.set(id, { resolve, reject, timer }); }); diff --git a/packages/klient/src/transports/memory/dispatcher.ts b/packages/klient/src/transports/memory/dispatcher.ts index a967ffa913..3380ff0232 100644 --- a/packages/klient/src/transports/memory/dispatcher.ts +++ b/packages/klient/src/transports/memory/dispatcher.ts @@ -62,8 +62,14 @@ export interface MemoryDispatcher { const REQUEST_INVALID = 40001; const NOT_FOUND = 40404; +/** kap-server wire codes mirrored so memory/ipc surface the same numeric codes as `/api/v2/mcp`. */ +const MCP_SERVER_NOT_FOUND = 40408; +const MCP_OAUTH_FAILED = 40929; const PROMPT_ID_CONFLICT = 40927; +/** Wire name of the engine's `IMcpManagementService` decorator id. */ +const MCP_MANAGEMENT_SERVICE = 'mcpManagementService'; + /** * Engine file errors cross the facade as public `RPCError`s, never as the * engine's raw `Error2`. The dispatcher is shared by both transports, so @@ -77,6 +83,28 @@ function rethrowFileErrorAsRpc(error: unknown): never { throw error; } +/** + * Same treatment for the MCP management plane: its coded rejections cross as + * `RPCError`s carrying the kap-server wire codes, so memory and ipc behave + * identically (a raw `Error2` would cross ipc as a generic 50001) and both + * match `/api/v2/mcp` — `mcp.server_not_found` → 40408, `request.invalid` / + * `config.invalid` → 40001, `mcp.oauth_failed` → 40929. + */ +function rethrowMcpManagementErrorAsRpc(error: unknown): never { + if (error instanceof Error2) { + switch (error.code) { + case ErrorCodes.MCP_SERVER_NOT_FOUND: + throw new RPCError(MCP_SERVER_NOT_FOUND, error.message, error.details); + case ErrorCodes.REQUEST_INVALID: + case ErrorCodes.CONFIG_INVALID: + throw new RPCError(REQUEST_INVALID, error.message, error.details); + case ErrorCodes.MCP_OAUTH_FAILED: + throw new RPCError(MCP_OAUTH_FAILED, error.message, error.details); + } + } + throw error; +} + type ScopeKind = 'core' | 'workspace' | 'session' | 'agent'; interface ResolvedScope { @@ -229,6 +257,9 @@ export function createMemoryDispatcher(root: ScopeLike): MemoryDispatcher { const result = await (member as (...a: unknown[]) => unknown).apply(instance, clonedArgs); return wireClone(result); } catch (error) { + if (service === MCP_MANAGEMENT_SERVICE) { + rethrowMcpManagementErrorAsRpc(error); + } if (error instanceof Error2 && error.code === ErrorCodes.PROMPT_ID_CONFLICT) { throw new RPCError(PROMPT_ID_CONFLICT, error.message, error.details); } diff --git a/packages/klient/src/transports/memory/serviceRegistry.ts b/packages/klient/src/transports/memory/serviceRegistry.ts index 920aa8ca7b..87c8abdd1f 100644 --- a/packages/klient/src/transports/memory/serviceRegistry.ts +++ b/packages/klient/src/transports/memory/serviceRegistry.ts @@ -48,6 +48,7 @@ import { IAgentTaskService } from '@moonshot-ai/agent-core-v2/agent/task/task'; import { IAgentUsageService } from '@moonshot-ai/agent-core-v2/agent/usage/usage'; import { IAgentMcpService } from '@moonshot-ai/agent-core-v2/agent/mcp/mcp'; import { IAgentFullCompactionService } from '@moonshot-ai/agent-core-v2/agent/fullCompaction/fullCompaction'; +import { IMcpManagementService } from '@moonshot-ai/agent-core-v2/app/mcpManagement/mcpManagement'; /** Wire service name (decorator id string) → token. */ export const serviceTokens: Readonly>> = { @@ -90,6 +91,7 @@ export const serviceTokens: Readonly>> agentTaskService: IAgentTaskService, agentMcpService: IAgentMcpService, agentFullCompactionService: IAgentFullCompactionService, + mcpManagementService: IMcpManagementService, }; export { IEventService }; diff --git a/packages/klient/test/contract-parity.ts b/packages/klient/test/contract-parity.ts index 601fd28660..842dc34376 100644 --- a/packages/klient/test/contract-parity.ts +++ b/packages/klient/test/contract-parity.ts @@ -35,6 +35,26 @@ import type { PlanData } from '@moonshot-ai/agent-core-v2/features/plan/plan'; import type { UsageStatus } from '@moonshot-ai/agent-core-v2/agent/usage/usage'; import type { SkillSummary } from '@moonshot-ai/agent-core-v2/app/skillCatalog/types'; import type { McpServerEntry } from '@moonshot-ai/agent-core-v2/mcpCore/connection-manager'; +import type { + GlobalMcpServerConfig, + McpAuthStatusQuery, + McpManagedServer, + McpServerAuthBeginResult, + McpServerAuthFlowHandle, + McpServerAuthState, + McpServerAuthStatus, + McpServerInspection, + McpServerLocator, + McpServerTestResult, + McpServerTestTarget, +} from '@moonshot-ai/agent-core-v2/app/mcpManagement/mcpManagement'; +import type { + McpRegistryPluginOrigin, + McpRegistryQuery, + McpServerSource, +} from '@moonshot-ai/agent-core-v2/app/mcpRegistry/mcpRegistry'; +import type { McpServerConfig } from '@moonshot-ai/agent-core-v2/mcpCore/config-schema'; +import type { McpServerConfigView } from '@moonshot-ai/agent-core-v2/mcpCore/configView'; import type { FullCompactionInput } from '@moonshot-ai/agent-core-v2/agent/fullCompaction/fullCompaction'; import type { ISessionScopeHandle } from '@moonshot-ai/agent-core-v2/_base/di/scope'; import type { @@ -273,6 +293,23 @@ import { fsHomeResponseSchema, } from '../src/contract/global/hostFs.js'; import { modelConfigSchema } from '../src/contract/global/models.js'; +import { + globalMcpServerConfigSchema, + mcpAuthStatusQuerySchema, + mcpManagedServerSchema, + mcpRegistryPluginOriginSchema, + mcpRegistryQuerySchema, + mcpServerAuthBeginResultSchema, + mcpServerAuthFlowHandleSchema, + mcpServerAuthStateSchema, + mcpServerAuthStatusSchema, + mcpServerConfigDataSchema, + mcpServerInspectionSchema, + mcpServerLocatorSchema, + mcpServerSourceSchema, + mcpServerTestResultSchema, + mcpServerTestTargetSchema, +} from '../src/contract/global/mcpManagement.js'; import { getPluginInfoInputSchema, installPluginInputSchema, @@ -441,6 +478,50 @@ const _setPluginMcpServerEnabledInput: AssertWire< const _removePluginInput: AssertWire = true; const _getPluginInfoInput: AssertWire = true; +// global/mcpManagement.ts — the `McpServerConfig | McpServerConfigView` union +// a managed server's `config` carries (full for mutable entries, redacted for +// read-only ones) is mirrored by one schema covering both shapes; the +// inspection's `config` is always the redacted view, and both assignability +// directions hold against either engine type. +const _mcpServerSource: AssertWire = true; +const _mcpRegistryPluginOrigin: AssertWire< + typeof mcpRegistryPluginOriginSchema, + McpRegistryPluginOrigin +> = true; +const _mcpRegistryQuery: AssertWire = true; +const _mcpAuthStatusQuery: AssertWire = true; +const _globalMcpServerConfig: AssertWire< + typeof globalMcpServerConfigSchema, + GlobalMcpServerConfig +> = true; +const _mcpServerConfigData: AssertWire< + typeof mcpServerConfigDataSchema, + McpServerConfig | McpServerConfigView +> = true; +const _mcpServerConfigViewData: AssertWire< + typeof mcpServerConfigDataSchema, + McpServerConfigView +> = true; +const _mcpManagedServer: AssertWire = true; +const _mcpServerTestTarget: AssertWire = + true; +const _mcpServerTestResult: AssertWire = + true; +const _mcpServerLocator: AssertWire = true; +const _mcpServerAuthState: AssertWire = true; +const _mcpServerInspection: AssertWire = + true; +const _mcpServerAuthStatus: AssertWire = + true; +const _mcpServerAuthBeginResult: AssertWire< + typeof mcpServerAuthBeginResultSchema, + McpServerAuthBeginResult +> = true; +const _mcpServerAuthFlowHandle: AssertWire< + typeof mcpServerAuthFlowHandleSchema, + McpServerAuthFlowHandle +> = true; + // env.ts has no named schemas; `platform` narrows to `NodeJS.Platform` in the // engine — assert the bootstrap properties are all strings instead. The // object-typed `clientIdentity` is intentionally not in this list. diff --git a/packages/klient/test/contract.test.ts b/packages/klient/test/contract.test.ts index 0d3b1ae6ba..1571b6e68f 100644 --- a/packages/klient/test/contract.test.ts +++ b/packages/klient/test/contract.test.ts @@ -9,6 +9,7 @@ import { describe, expect, it } from 'vitest'; import { pluginManifestSchema } from '../src/contract/global/plugins.js'; +import { mcpServerAuthFlowHandleSchema } from '../src/contract/global/mcpManagement.js'; import { createSessionOptionsSchema } from '../src/contract/session/lifecycle.js'; import { promptPayloadSchema } from '../src/contract/agent/schemas.js'; @@ -65,6 +66,17 @@ describe('MCP timeout contract validation', () => { }); expect(parsed.success).toBe(false); }); + + it('completeAuth timeoutMs accepts the setTimeout maximum and rejects above it', () => { + expect( + mcpServerAuthFlowHandleSchema.safeParse({ flowId: 'flow-1', timeoutMs: 2_147_483_647 }) + .success, + ).toBe(true); + expect( + mcpServerAuthFlowHandleSchema.safeParse({ flowId: 'flow-1', timeoutMs: 2_147_483_648 }) + .success, + ).toBe(false); + }); }); describe('prompt contract validation', () => { diff --git a/packages/klient/test/helpers/conformance.ts b/packages/klient/test/helpers/conformance.ts index 531e1136df..ac62b6169c 100644 --- a/packages/klient/test/helpers/conformance.ts +++ b/packages/klient/test/helpers/conformance.ts @@ -322,6 +322,171 @@ export function defineKlientConformance( expect(typeof status.loggedIn).toBe('boolean'); }); + it('global mcp round-trips user-level server CRUD', async () => { + const mcp = target.klient.global.mcp; + const cwd = await mkdtemp(join(tmpdir(), 'klient-conf-mcp-crud-')); + try { + expect(await mcp.list({ cwd })).toEqual([]); + + const added = await mcp.add({ + cwd, + server: { + name: 'conf-mcp', + transport: 'stdio', + command: 'conf-command', + env: { TOKEN: 'secret' }, + }, + }); + const entry = added.find((server) => server.name === 'conf-mcp'); + // Mutable (user-level) entries carry the full config for edit prefill. + expect(entry).toMatchObject({ + name: 'conf-mcp', + source: 'global', + mutable: true, + config: { transport: 'stdio', command: 'conf-command', env: { TOKEN: 'secret' } }, + }); + + await mcp.update({ + cwd, + server: { name: 'conf-mcp', transport: 'stdio', command: 'conf-command-2' }, + }); + expect((await mcp.get({ name: 'conf-mcp', cwd })).config).toMatchObject({ + command: 'conf-command-2', + }); + + await mcp.remove({ name: 'conf-mcp', cwd }); + expect(await mcp.list({ cwd })).toEqual([]); + await expect(mcp.get({ name: 'conf-mcp', cwd })).rejects.toMatchObject({ + name: 'RPCError', + code: 40408, + }); + } finally { + await rm(cwd, { recursive: true, force: true }); + } + }); + + it('global mcp probes an inline server config without persisting it', async () => { + const mcp = target.klient.global.mcp; + // Inline probe against a scratch cwd: the binary runs but never + // speaks MCP, so the connection test reports a clean failure. + const probeCwd = await mkdtemp(join(tmpdir(), 'klient-conf-mcp-probe-')); + try { + const probe = await mcp.test({ + server: { + name: 'conf-probe', + transport: 'stdio', + command: process.execPath, + args: ['--version'], + startupTimeoutMs: 10_000, + }, + cwd: probeCwd, + }); + expect(probe.success).toBe(false); + expect(typeof probe.output).toBe('string'); + } finally { + await rm(probeCwd, { recursive: true, force: true }); + } + expect(await mcp.list()).toEqual([]); + }); + + it('global mcp resolves locators and classifies auth offline', async () => { + const mcp = target.klient.global.mcp; + await mcp.add({ + server: { name: 'conf-mcp', transport: 'stdio', command: 'conf-command' }, + }); + try { + // The locator surface: resolve a legacy name, inspect nothing/all. + expect(await mcp.resolveByName({ name: 'conf-mcp' })).toEqual({ + source: 'global', + name: 'conf-mcp', + }); + expect(await mcp.inspect({ targets: [] })).toEqual([]); + // An omitted `targets` ahead of a present options arg must mean "the + // whole catalog" on every transport (ipc encodes it as `null`). + expect((await mcp.inspect({})).map((i) => i.runtimeName)).toContain('conf-mcp'); + await expect( + mcp.inspect({ targets: [{ source: 'global', name: 'conf-missing' }] }), + ).rejects.toMatchObject({ name: 'RPCError', code: 40408 }); + + // Offline auth classification of a stdio server needs no probe, and + // an OAuth flow against a stdio target is request.invalid → 40001. + expect(await mcp.authStatuses()).toEqual([ + { name: 'conf-mcp', authStatus: 'not-applicable' }, + ]); + await expect( + mcp.beginAuth({ locator: { source: 'global', name: 'conf-mcp' } }), + ).rejects.toMatchObject({ name: 'RPCError', code: 40001 }); + } finally { + await mcp.remove({ name: 'conf-mcp' }); + } + }); + + it('global mcp completeAuth rejects an unknown flowId with 40001', async () => { + const mcp = target.klient.global.mcp; + await expect(mcp.completeAuth({ flowId: 'conf-unknown-flow' })).rejects.toMatchObject({ + name: 'RPCError', + code: 40001, + }); + }); + + it('global mcp OAuth failures map to the 40929 wire code on every transport', async () => { + const mcp = target.klient.global.mcp; + await mcp.add({ + server: { + name: 'conf-oauth-failure', + transport: 'http', + url: 'http://127.0.0.1:1/mcp', + auth: 'oauth', + }, + }); + try { + await expect( + mcp.beginAuth({ locator: { source: 'global', name: 'conf-oauth-failure' } }), + ).rejects.toMatchObject({ name: 'RPCError', code: 40929 }); + } finally { + await mcp.remove({ name: 'conf-oauth-failure' }); + } + }); + + it('global mcp cancelAuth ignores an unknown flowId', async () => { + const mcp = target.klient.global.mcp; + await expect(mcp.cancelAuth({ flowId: 'conf-unknown-flow' })).resolves.toBeUndefined(); + }); + + it('global mcp resetAuth clears a remote oauth server through the transport', async () => { + const mcp = target.klient.global.mcp; + await mcp.add({ + server: { + name: 'conf-oauth', + transport: 'http', + url: 'https://example.com/mcp', + auth: 'oauth', + }, + }); + try { + // Invalidate is offline: no stored grant and no network needed. + await expect( + mcp.resetAuth({ locator: { source: 'global', name: 'conf-oauth' } }), + ).resolves.toBeUndefined(); + } finally { + await mcp.remove({ name: 'conf-oauth' }); + } + }); + + it('global mcp resetAuth rejects a stdio locator with 40001', async () => { + const mcp = target.klient.global.mcp; + await mcp.add({ + server: { name: 'conf-stdio', transport: 'stdio', command: 'conf-command' }, + }); + try { + await expect( + mcp.resetAuth({ locator: { source: 'global', name: 'conf-stdio' } }), + ).rejects.toMatchObject({ name: 'RPCError', code: 40001 }); + } finally { + await mcp.remove({ name: 'conf-stdio' }); + } + }); + it('agent runtime binding is available through every transport', async () => { const created = await target.klient.global.sessions.create({ workDir: process.cwd(), diff --git a/packages/klient/test/ipc.test.ts b/packages/klient/test/ipc.test.ts index d0d91c5804..718bad4a12 100644 --- a/packages/klient/test/ipc.test.ts +++ b/packages/klient/test/ipc.test.ts @@ -2,7 +2,8 @@ import { rm } from 'node:fs/promises'; import { join } from 'node:path'; import { tmpdir } from 'node:os'; -import { describe, expect, it } from 'vitest'; +import { IMcpManagementService } from '@moonshot-ai/agent-core-v2'; +import { describe, expect, it, vi } from 'vitest'; import { defineKlientConformance } from './helpers/conformance.js'; import { createKlient, serveKlientIpc, type KlientIpcHost } from '../src/transports/ipc/index.js'; @@ -72,4 +73,62 @@ describe('ipc transport specifics', () => { await ok.close(); await teardown(); }); + + it('completeAuth outlives the channel default call timeout', async () => { + const socketPath = await setup(); + // A slow engine-side wait: without the facade's per-call deadline the + // channel's default would kill the long poll mid-flight. + const management = app.accessor.get(IMcpManagementService); + const completeSpy = vi + .spyOn(management, 'completeServerAuth') + .mockImplementation( + () => new Promise((resolve) => setTimeout(resolve, 200)), + ); + const cancelSpy = vi + .spyOn(management, 'cancelServerAuth') + .mockImplementation( + () => new Promise((resolve) => setTimeout(resolve, 200)), + ); + const klient = createKlient({ socketPath, callTimeoutMs: 25 }); + try { + // completeAuth passes the engine wait + margin as its per-call deadline, + // so the 200ms wait resolves instead of dying at the 25ms default. + await expect( + klient.global.mcp.completeAuth({ flowId: 'flow-1', timeoutMs: 100 }), + ).resolves.toBeUndefined(); + // Calls without the override still die at the channel default. + await expect(klient.global.mcp.cancelAuth({ flowId: 'flow-1' })).rejects.toThrow( + 'call timed out after 25ms', + ); + } finally { + completeSpy.mockRestore(); + cancelSpy.mockRestore(); + await klient.close(); + } + await teardown(); + }); + + it('completeAuth clamps a near-max timeoutMs instead of overflowing the call timer', async () => { + const socketPath = await setup(); + const management = app.accessor.get(IMcpManagementService); + const completeSpy = vi + .spyOn(management, 'completeServerAuth') + .mockImplementation( + () => new Promise((resolve) => setTimeout(resolve, 50)), + ); + const klient = createKlient({ socketPath, callTimeoutMs: 25 }); + try { + // timeoutMs at the contract max plus the facade margin would overflow + // Node's 32-bit setTimeout into ~1ms; the clamp keeps the call alive + // until the engine-side wait resolves. + await expect( + klient.global.mcp.completeAuth({ flowId: 'flow-1', timeoutMs: 2 ** 31 - 1 }), + ).resolves.toBeUndefined(); + expect(completeSpy).toHaveBeenCalled(); + } finally { + completeSpy.mockRestore(); + await klient.close(); + } + await teardown(); + }); }); diff --git a/packages/node-sdk/src/kimi-harness.ts b/packages/node-sdk/src/kimi-harness.ts index fb68348a28..48b2e3c1e9 100644 --- a/packages/node-sdk/src/kimi-harness.ts +++ b/packages/node-sdk/src/kimi-harness.ts @@ -501,27 +501,37 @@ export class KimiHarness { */ async inspectAppMcpServers( targets?: readonly McpServerLocator[], + options: { readonly cwd?: string } = {}, ): Promise { - return this.rpc.inspectAppMcpServers(targets); + return this.rpc.inspectAppMcpServers(targets, options); } - async addMcpServer(server: McpServerConfig): Promise { - return this.rpc.addGlobalMcpServer(server); + async addMcpServer( + server: McpServerConfig, + options: { readonly cwd?: string } = {}, + ): Promise { + return this.rpc.addGlobalMcpServer(server, options); } - async updateMcpServer(server: McpServerConfig): Promise { - return this.rpc.updateGlobalMcpServer(server); + async updateMcpServer( + server: McpServerConfig, + options: { readonly cwd?: string } = {}, + ): Promise { + return this.rpc.updateGlobalMcpServer(server, options); } - async removeMcpServer(name: string): Promise { - return this.rpc.removeGlobalMcpServer(name); + async removeMcpServer( + name: string, + options: { readonly cwd?: string } = {}, + ): Promise { + return this.rpc.removeGlobalMcpServer(name, options); } async authenticateMcpServer( name: string, options: AuthenticateMcpServerOptions, ): Promise { - const started = await this.rpc.beginGlobalMcpServerAuth(name); + const started = await this.rpc.beginGlobalMcpServerAuth(name, { cwd: options.cwd }); if (started.status === 'already-authorized') return; try { const opened = await options.onAuthorizationUrl(started.authorizationUrl); @@ -538,8 +548,11 @@ export class KimiHarness { } } - async resetMcpServerAuth(name: string): Promise { - return this.rpc.resetGlobalMcpServerAuth(name); + async resetMcpServerAuth( + name: string, + options: { readonly cwd?: string } = {}, + ): Promise { + return this.rpc.resetGlobalMcpServerAuth(name, options); } /** @@ -551,7 +564,7 @@ export class KimiHarness { locator: McpServerLocator, options: AuthenticateMcpServerOptions, ): Promise { - const started = await this.rpc.beginMcpServerAuth(locator); + const started = await this.rpc.beginMcpServerAuth(locator, { cwd: options.cwd }); if (started.status === 'already-authorized') return; try { const opened = await options.onAuthorizationUrl(started.authorizationUrl); @@ -569,8 +582,11 @@ export class KimiHarness { } /** The locator-addressed variant of {@link resetMcpServerAuth}. */ - async resetAppMcpServerAuth(locator: McpServerLocator): Promise { - return this.rpc.resetMcpServerAuth(locator); + async resetAppMcpServerAuth( + locator: McpServerLocator, + options: { readonly cwd?: string } = {}, + ): Promise { + return this.rpc.resetMcpServerAuth(locator, options); } async testMcpServer( diff --git a/packages/node-sdk/src/rpc.ts b/packages/node-sdk/src/rpc.ts index 44b35136cf..f5b96ef0fe 100644 --- a/packages/node-sdk/src/rpc.ts +++ b/packages/node-sdk/src/rpc.ts @@ -409,36 +409,50 @@ export abstract class SDKRpcClientBase { async inspectAppMcpServers( targets?: readonly McpServerLocator[], + options: { readonly cwd?: string } = {}, ): Promise { const rpc = await this.getRpc(); - return rpc.inspectAppMcpServers({ targets }); + return rpc.inspectAppMcpServers({ targets, cwd: options.cwd }); } - async addGlobalMcpServer(server: McpServerConfig): Promise { + async addGlobalMcpServer( + server: McpServerConfig, + options: { readonly cwd?: string } = {}, + ): Promise { const rpc = await this.getRpc(); - return rpc.addGlobalMcpServer({ server }); + return rpc.addGlobalMcpServer({ server, cwd: options.cwd }); } async updateGlobalMcpServer( server: McpServerConfig, + options: { readonly cwd?: string } = {}, ): Promise { const rpc = await this.getRpc(); - return rpc.updateGlobalMcpServer({ server }); + return rpc.updateGlobalMcpServer({ server, cwd: options.cwd }); } - async removeGlobalMcpServer(name: string): Promise { + async removeGlobalMcpServer( + name: string, + options: { readonly cwd?: string } = {}, + ): Promise { const rpc = await this.getRpc(); - return rpc.removeGlobalMcpServer({ name }); + return rpc.removeGlobalMcpServer({ name, cwd: options.cwd }); } - async beginGlobalMcpServerAuth(name: string): Promise { + async beginGlobalMcpServerAuth( + name: string, + options: { readonly cwd?: string } = {}, + ): Promise { const rpc = await this.getRpc(); - return rpc.beginGlobalMcpServerAuth({ name }); + return rpc.beginGlobalMcpServerAuth({ name, cwd: options.cwd }); } - async beginMcpServerAuth(locator: McpServerLocator): Promise { + async beginMcpServerAuth( + locator: McpServerLocator, + options: { readonly cwd?: string } = {}, + ): Promise { const rpc = await this.getRpc(); - return rpc.beginMcpServerAuth({ locator }); + return rpc.beginMcpServerAuth({ locator, cwd: options.cwd }); } async completeGlobalMcpServerAuth( @@ -467,14 +481,20 @@ export abstract class SDKRpcClientBase { return rpc.cancelMcpServerAuth({ flowId }); } - async resetGlobalMcpServerAuth(name: string): Promise { + async resetGlobalMcpServerAuth( + name: string, + options: { readonly cwd?: string } = {}, + ): Promise { const rpc = await this.getRpc(); - return rpc.resetGlobalMcpServerAuth({ name }); + return rpc.resetGlobalMcpServerAuth({ name, cwd: options.cwd }); } - async resetMcpServerAuth(locator: McpServerLocator): Promise { + async resetMcpServerAuth( + locator: McpServerLocator, + options: { readonly cwd?: string } = {}, + ): Promise { const rpc = await this.getRpc(); - return rpc.resetMcpServerAuth({ locator }); + return rpc.resetMcpServerAuth({ locator, cwd: options.cwd }); } async testGlobalMcpServer( diff --git a/packages/node-sdk/src/sdk-rpc-client-v2.ts b/packages/node-sdk/src/sdk-rpc-client-v2.ts index 5597a2dc00..89a6383ce6 100644 --- a/packages/node-sdk/src/sdk-rpc-client-v2.ts +++ b/packages/node-sdk/src/sdk-rpc-client-v2.ts @@ -83,23 +83,15 @@ * `resetGlobalMcpServerAuth` / `testGlobalMcpServer` / * `testGlobalMcpServerConfig` / `inspectAppMcpServers` / * `beginMcpServerAuth` / `completeMcpServerAuth` / `cancelMcpServerAuth` / - * `resetMcpServerAuth` → the v1 user-global - * MCP surface, rebuilt in `src/v2/global-mcp.ts`: agent-core-v2 only reads - * the user-global `mcp.json` (nothing in the engine writes it) and binds - * its OAuth orchestrator inside the session scope, so the file store and - * the `require*` guards are byte-identical ports, driven by the v2 - * engine's own `McpOAuthService` / `McpConnectionManager` (deep imports — - * the package index does not re-export them) over the app-scope - * `IAtomicDocumentStore`, whose on-disk layout - * (`/credentials/mcp/-*.json`) matches v1's. Every result is - * tagged `source: 'global'` / `mutable: true` with the file path as - * `origin` — plugin and project-layer entries in the unified view are a - * v1-only addition for now. The same gap shapes the locator-addressed - * app-level surface: the descriptor catalog holds global entries only - * (this engine has no `IPluginService.mcpServers` to flatten plugin - * manifests with), a plugin locator resolves to `mcp.server_not_found`, - * and credential resets do not notify live sessions (no - * `IMcpAuthCoordinator` here either). + * `resetMcpServerAuth` → the engine's App-scope `IMcpManagementService` + * (through {@link engineAccessor}; no klient facade exists): the unified + * MCP management plane over the `mcpRegistry` read view — user-level + * `mcp.json` CRUD guarded against read-only plugin / project-layer + * collisions, the standalone connection probe, the locator-addressed + * inspection catalog (plugin entries included), and locator-addressed + * OAuth flows keyed by flowId. The managed-server results are mapped back + * to the v1 wire shape (config flattened to the top level); the inspection + * and auth-status shapes are field-identical between the engines. * - `listMcpServers` / `getMcpStartupMetrics` / `reconnectMcpServer` / * `addSessionMcpServer` → * the seeded `ISessionMcpHandle.connectionManager` through the session @@ -137,7 +129,6 @@ * `toolCall` keeps the base class's "not supported" answer, which the * interaction bridge already relies on. */ -import { randomUUID } from 'node:crypto'; import { readdir } from 'node:fs/promises'; import { join } from 'node:path'; @@ -145,26 +136,18 @@ import { ensureConfigFile, ErrorCodes, HookDefSchema, + isKimiErrorCode, KimiError, limitAgentReplayByTurns, noopTelemetryClient, type AgentContextData, type BeginGlobalMcpServerAuthResult, type ExperimentalFeatureState, + type KimiErrorCode, } from '@moonshot-ai/agent-core'; import { encodeWorkDirKey } from '@moonshot-ai/agent-core-v2/_base/utils/workdir-slug'; -import { MCP_SECTION, type McpSection } from '@moonshot-ai/agent-core-v2/app/mcpConfig/configSection'; -import { IAgentIdentity } from '@moonshot-ai/agent-core-v2/app/agentIdentity/agentIdentity'; import { McpConnectionManager } from '@moonshot-ai/agent-core-v2/mcpCore/connection-manager'; -import { - AlreadyAuthorizedError, - McpOAuthService, - type BeginAuthorizationResult, -} from '@moonshot-ai/agent-core-v2/mcpCore/oauth/service'; -import { createMcpOAuthStore } from '@moonshot-ai/agent-core-v2/app/mcpConfig/oauthStore'; -import { canonicalMcpOAuthResource } from '@moonshot-ai/agent-core-v2/mcpCore/oauth/store'; -import { IAtomicDocumentStore } from '@moonshot-ai/agent-core-v2/persistence/interface/atomicDocumentStore'; -import { loadMcpServers } from '@moonshot-ai/agent-core-v2/workspace/workspaceMcpConfig/internal/config-loader'; +import { loadMcpServers } from '@moonshot-ai/agent-core-v2/app/mcpConfig/configLoader'; import type { McpServerConfig as WorkspaceMcpServerConfig } from '@moonshot-ai/agent-core-v2/mcpCore/config-schema'; import { bootstrap, @@ -197,6 +180,8 @@ import { IEventService, IHostEnvironment, IHostFileSystem, + IMcpManagementService, + IMcpOAuthService, IModelService, IProviderService, ISessionBtwService, @@ -215,11 +200,11 @@ import { ITelemetryService, IWorkspaceAliases, ISessionActivityView, - IRuntimeResolver, IWorkspaceInstanceManager, closeSessionById, followSessionLifecycles, getLiveSessionById, + isError2, programForSession, resumeSessionById, sessionDirOf, @@ -240,6 +225,7 @@ import { type IAgentScopeHandle, type IDisposable, type ISessionScopeHandle, + type McpManagedServer, type Scope, type ServicesAccessor, type SessionSummary as V2SessionSummary, @@ -289,7 +275,6 @@ import type { GenerateSessionTitleInput, GetConfigOptions, GetCronTasksResult, - GlobalMcpServerAuthState, GlobalMcpServerAuthStatus, GoalSnapshot, GoalToolResult, @@ -334,21 +319,10 @@ import { translateGlobalEvent } from '#/v2/event-mapper'; import { assertImportFits, buildImportContextMessage } from '#/v2/import-context'; import { foldAgentWireReplay } from '#/v2/resume-replay'; import { - GlobalMcpConfigStore, - configuredMcpAuthState, - isOAuthProbeCandidate, mcpConfigWithoutName, - mcpServerId, normalizeServerName, parseInlineMcpServer, parseReconnectMcpServerConfig, - requireOAuthMcpConfig, - requireRemoteMcpConfig, - sanitizeAppMcpServerInspection, - selectAppMcpServerDescriptors, - standaloneMcpTestResult, - type AppMcpServerRuntimeDescriptor, - type AppMcpServerRuntimeInspection, } from '#/v2/global-mcp'; import { normalizeWorkDir, @@ -380,9 +354,6 @@ export interface SDKRpcClientV2Options { */ const MAX_TIMER_DELAY_MS = 0x7fffffff; -/** v1's default for `completeGlobalMcpServerAuth` when the host gives no timeout. */ -const DEFAULT_GLOBAL_MCP_AUTH_TIMEOUT_MS = 15 * 60 * 1000; - export class SDKRpcClientV2 extends SDKRpcClientBase { readonly homeDir: string; readonly configPath: string; @@ -418,20 +389,6 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { * touching a profile. */ private readonly modelReady: Promise; - /** - * The user-global MCP file store (`/mcp.json`) — the SDK-side port in - * `src/v2/global-mcp.ts`, because agent-core-v2 only reads that file and - * has no write-side service for it. - */ - private readonly globalMcpConfig: GlobalMcpConfigStore; - /** - * v1's per-core global OAuth orchestrator, mirrored per client. Built - * lazily over the app-scope `IAtomicDocumentStore` (same on-disk layout as - * v1's `/credentials/mcp/` file store). - */ - private globalMcpOAuth: McpOAuthService | undefined; - /** In-flight OAuth flows keyed by the flowId handed to the host (v1 shape). */ - private readonly globalMcpOAuthFlows = new Map(); /** * Per-live-session event/interaction wirings (`src/v2/session-wiring.ts`): * created when a session materializes through this client (create / resume / @@ -493,7 +450,6 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { ); this.app = app; this.klient = createKlient({ scope: app }); - this.globalMcpConfig = new GlobalMcpConfigStore(this.homeDir); this.configReady = app.accessor.get(IConfigService).ready; this.installEngineTelemetry(options.telemetry); this.modelReady = Promise.all([ @@ -548,6 +504,11 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { // disposal fires — a host that removes homeDir right after close() must // not race an in-flight shard close (ENOTEMPTY on teardown). await this.app.accessor.get(ISessionIndexMirror).drain(); + // Await the OAuth service shutdown directly rather than after dispose(): + // its ledger-teardown dispose can queue behind slow async disposables, and + // the accessor throws once the scope is disposed. shutdown() is + // idempotent, so the ledger's own teardown turns into a no-op. + await this.app.accessor.get(IMcpOAuthService).shutdown(); this.app.dispose(); await drainSessionIndexMirror(); await drainQueryStoreDisposals(); @@ -956,29 +917,16 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { return handle; } - /** The live session's workspace cwd, resolved like `listSessions` maps it. */ - private async sessionWorkDir(sessionId: string): Promise { - const page = await this.klient.global.sessions.list({ sessionId, limit: 1 }); - const item = page.items[0]; - if (item === undefined) return undefined; - if (item.cwd !== undefined) return item.cwd; - const workspaces = await this.klient.global.workspaces.list(); - return workspaces.find((workspace) => workspace.id === item.workspaceId)?.root; - } - /** - * v1's persist-add guard ported to the workspace loader: a same-named - * project-layer entry wins over the user file, so persisting would write a - * shadow that never takes effect — and the direct workspace-manager upsert - * would displace the project config every live session runs. Reject like - * v1's read-only rule instead. + * v1's persist-add project guard ported to the workspace loader. This read + * deliberately includes the project layer even while the workspace is + * untrusted: a user-level write must not create a shadow that springs into + * conflict when the workspace is trusted later. */ private async rejectProjectLayerPersistedMcpAdd( - sessionId: string, + cwd: string, name: string, ): Promise { - const cwd = await this.sessionWorkDir(sessionId); - if (cwd === undefined) return; const fs = this.engineAccessor.get(IHostFileSystem); const [withProject, userOnly] = await Promise.all([ loadMcpServers({ fs, cwd, homeDir: this.homeDir, includeProject: true }), @@ -2343,179 +2291,145 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { } // ----------------------------------------------------------------------- - // MCP: the user-global surface is rebuilt over the SDK-side store port in - // `src/v2/global-mcp.ts` plus the v2 engine's own OAuth service and - // connection manager (agent-core-v2 has no app-scope MCP config service — - // it only reads `mcp.json`); the session-level reads go through the - // session scope's seeded `ISessionMcpHandle` (no klient facade exists for - // either group). + // MCP: the management plane (user-global CRUD, the standalone probe, the + // locator-addressed inspection catalog, OAuth flows) delegates to the + // engine's App-scope `IMcpManagementService`; the session-level reads go + // through the session scope's seeded `ISessionMcpHandle` (no klient facade + // exists for either group). // ----------------------------------------------------------------------- /** - * Configured custom identity announced to MCP servers, so these global flows - * match what the workspace-owned manager sends. Reads the frozen snapshot; - * every path that reaches it awaited `globalMcpOAuthService()` first. - */ - private resolveMcpClientName(): string | undefined { - return this.engineAccessor.get(IAgentIdentity).current().slug; - } - - /** - * v1's per-core `globalMcpOAuth`, built over the app-scope document store. - * - * Async on purpose: the service caches providers by store key and stamps the - * client name when it first builds one, so any path that can materialize a - * provider must not run before the identity snapshot froze. Guarding here - * rather than at each call site means a new entry point cannot forget to. + * The engine's management plane throws `Error2`; the SDK's public error + * contract is `KimiError` (what `isKimiError` branches on, and what the v1 + * client throws for the same failures). Restate so both engines surface + * the identical class — see `restateMcpManagementError`. */ - private async globalMcpOAuthService(): Promise { - await this.engineAccessor.get(IAgentIdentity).resolved(); - if (this.globalMcpOAuth === undefined) { - this.globalMcpOAuth = new McpOAuthService({ - store: createMcpOAuthStore(this.engineAccessor.get(IAtomicDocumentStore)), - resolveClientName: () => this.resolveMcpClientName(), - }); + private async mcpManagement( + call: (management: IMcpManagementService) => Promise, + ): Promise { + try { + return await call(this.engineAccessor.get(IMcpManagementService)); + } catch (error) { + throw restateMcpManagementError(error); } - return this.globalMcpOAuth; } - /** - * A fresh per-call service whose providers re-read the token store. The v2 - * providers snapshot tokens at construction, so the cached `globalMcpOAuth` - * can keep serving a grant another process has since removed — or stay blind - * to one just saved. Same options as the cached service. - */ - private async freshGlobalMcpOAuthService(): Promise { - await this.engineAccessor.get(IAgentIdentity).resolved(); - return new McpOAuthService({ - store: createMcpOAuthStore(this.engineAccessor.get(IAtomicDocumentStore)), - resolveClientName: () => this.resolveMcpClientName(), - }); - } - - /** - * The unified management view's `McpManagedServerInfo` tag. Plugin and - * project-layer entries are a v1-only addition for now — every entry on - * this surface comes from the user-level file, so all are `global`, - * mutable, and carry the file path as their origin. - */ - private managedGlobalMcpServer(server: McpServerConfig): McpManagedServerInfo { - return { ...server, source: 'global', origin: this.globalMcpConfig.path, mutable: true }; - } - - override async listGlobalMcpServers(): Promise { - return (await this.globalMcpConfig.list()).map((server) => this.managedGlobalMcpServer(server)); + override async listGlobalMcpServers( + options: { readonly cwd?: string } = {}, + ): Promise { + const servers = await this.mcpManagement((management) => + management.listServers({ cwd: options.cwd }), + ); + return servers.map(toManagedServerInfo); } - override async getGlobalMcpServer(name: string): Promise { - return this.managedGlobalMcpServer(await this.globalMcpConfig.get(name)); + override async getGlobalMcpServer( + name: string, + options: { readonly cwd?: string } = {}, + ): Promise { + const server = await this.mcpManagement((management) => + management.getServer(name, { cwd: options.cwd }), + ); + return toManagedServerInfo(server); } override async listGlobalMcpServerAuthStatuses( options: { readonly cwd?: string; readonly verify?: boolean } = {}, ): Promise { - const servers = await this.globalMcpConfig.list(); - const oauth = await this.freshGlobalMcpOAuthService(); - const verify = options.verify === true; - return Promise.all( - servers.map(async (server) => ({ - name: server.name, - authStatus: await this.globalMcpServerAuthState(server, oauth, options.cwd, verify), - })), + const statuses = await this.mcpManagement((management) => + management.listAuthStatuses({ cwd: options.cwd, verify: options.verify }), ); + // The legacy surface never reports `unavailable` (no ambiguity check + // here), so the engine's wider state union narrows to the v1 wire one. + return statuses as readonly GlobalMcpServerAuthStatus[]; } override async inspectAppMcpServers( targets?: readonly McpServerLocator[], + options: { readonly cwd?: string } = {}, ): Promise { - const catalog = await this.appMcpServerDescriptors(); - const descriptors = selectAppMcpServerDescriptors(catalog, targets); - const inspections = await this.inspectAppMcpServerDescriptors(descriptors, catalog); - return inspections.map(sanitizeAppMcpServerInspection); + const inspections = await this.mcpManagement((management) => + management.inspectServers(targets, { cwd: options.cwd }), + ); + // Field-identical with the v1 wire shape (the engines' locator / + // config-view / auth-state declarations match structurally). + return inspections as readonly AppMcpServerInspection[]; } override async addGlobalMcpServer( server: McpServerConfig, + options: { readonly cwd?: string } = {}, ): Promise { - return (await this.globalMcpConfig.add(server)).map((entry) => - this.managedGlobalMcpServer(entry), + const servers = await this.mcpManagement((management) => + management.addServer(server, { cwd: options.cwd }), ); + return servers.map(toManagedServerInfo); } override async updateGlobalMcpServer( server: McpServerConfig, + options: { readonly cwd?: string } = {}, ): Promise { - return (await this.globalMcpConfig.update(server)).map((entry) => - this.managedGlobalMcpServer(entry), + const servers = await this.mcpManagement((management) => + management.updateServer(server, { cwd: options.cwd }), ); + return servers.map(toManagedServerInfo); } - override async removeGlobalMcpServer(name: string): Promise { - return (await this.globalMcpConfig.remove(name)).map((entry) => - this.managedGlobalMcpServer(entry), + override async removeGlobalMcpServer( + name: string, + options: { readonly cwd?: string } = {}, + ): Promise { + const servers = await this.mcpManagement((management) => + management.removeServer(name, { cwd: options.cwd }), ); + return servers.map(toManagedServerInfo); } /** - * v1's flow state machine verbatim: the OAuth config guards - * (`requireOAuthMcpConfig`) run before any network I/O, a begun flow is held - * by flowId until complete/cancel, and an already-authorized server - * short-circuits without a flowId. The OAuth work itself is the v2 engine's - * `McpOAuthService` (same begin/complete/cancel contract as v1's). + * The legacy name-only entry point resolves its locator first: exactly one + * enabled entry may own the runtime name, so a global/plugin collision + * rejects instead of guessing which credential the flow acts on. */ - override async beginGlobalMcpServerAuth(name: string): Promise { - return this.beginMcpServerAuth({ source: 'global', name }); + override async beginGlobalMcpServerAuth( + name: string, + options: { readonly cwd?: string } = {}, + ): Promise { + return this.mcpManagement(async (management) => { + const query = { cwd: options.cwd }; + return management.beginServerAuth(await management.resolveServerByName(name, query), query); + }); } override async beginMcpServerAuth( locator: McpServerLocator, + options: { readonly cwd?: string } = {}, ): Promise { - const server = await this.resolveAppMcpServer(locator); - const config = requireOAuthMcpConfig(server.runtimeName, server.config); - try { - // Fresh service: a grant saved or reset by another process after the - // cached one was built must decide whether a browser flow is even - // needed. - const oauth = await this.freshGlobalMcpOAuthService(); - const flow = await oauth.beginAuthorization(server.runtimeName, config.url); - const flowId = randomUUID(); - this.globalMcpOAuthFlows.set(flowId, { flow }); - return { - status: 'authorization-required', - flowId, - authorizationUrl: flow.authorizationUrl.toString(), - }; - } catch (error) { - if (error instanceof AlreadyAuthorizedError) { - return { status: 'already-authorized' }; - } - throw error; - } + return this.mcpManagement((management) => + management.beginServerAuth(locator, { cwd: options.cwd }), + ); } override async completeGlobalMcpServerAuth( - input: { readonly flowId: string; readonly timeoutMs?: number }, + input: { + readonly flowId: string; + readonly timeoutMs?: number; + }, signal?: AbortSignal, ): Promise { return this.completeMcpServerAuth(input, signal); } override async completeMcpServerAuth( - input: { readonly flowId: string; readonly timeoutMs?: number }, + input: { + readonly flowId: string; + readonly timeoutMs?: number; + }, signal?: AbortSignal, ): Promise { - const active = this.globalMcpOAuthFlows.get(input.flowId); - if (active === undefined) { - throw new KimiError(ErrorCodes.REQUEST_INVALID, `Unknown MCP OAuth flow: ${input.flowId}`); - } - try { - await active.flow.complete({ - signal, - timeoutMs: input.timeoutMs ?? DEFAULT_GLOBAL_MCP_AUTH_TIMEOUT_MS, - }); - } finally { - this.globalMcpOAuthFlows.delete(input.flowId); - } + return this.mcpManagement((management) => + management.completeServerAuth(input, { signal }), + ); } override async cancelGlobalMcpServerAuth(flowId: string): Promise { @@ -2523,40 +2437,34 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { } override async cancelMcpServerAuth(flowId: string): Promise { - const active = this.globalMcpOAuthFlows.get(flowId); - if (active === undefined) return; - this.globalMcpOAuthFlows.delete(flowId); - await active.flow.cancel(); + return this.mcpManagement((management) => management.cancelServerAuth({ flowId })); } - override async resetGlobalMcpServerAuth(name: string): Promise { - return this.resetMcpServerAuth({ source: 'global', name }); + override async resetGlobalMcpServerAuth( + name: string, + options: { readonly cwd?: string } = {}, + ): Promise { + return this.mcpManagement(async (management) => { + const query = { cwd: options.cwd }; + return management.resetServerAuth(await management.resolveServerByName(name, query), query); + }); } - override async resetMcpServerAuth(locator: McpServerLocator): Promise { - const server = await this.resolveAppMcpServer(locator); - const config = requireRemoteMcpConfig(server.runtimeName, server.config); - const oauth = await this.globalMcpOAuthService(); - // No `IMcpAuthCoordinator` on this engine: live sessions are not - // notified about the invalidation (v1 propagates via its OAuth events). - await oauth.invalidate(server.runtimeName, config.url); + override async resetMcpServerAuth( + locator: McpServerLocator, + options: { readonly cwd?: string } = {}, + ): Promise { + return this.mcpManagement((management) => + management.resetServerAuth(locator, { cwd: options.cwd }), + ); } - /** - * v1's standalone probe: a throwaway connection manager over the global - * file entry, never the session's. The global default timeouts come from - * the live `[mcp]` config section (awaited first — the config service's - * reads are synchronous over pre-ready state, the same trap as the config - * batch); v1 resolves the same section with the same env bindings, so the - * precedence matches. - */ override async testGlobalMcpServer( name: string, options: { readonly cwd?: string } = {}, ): Promise { - const server = await this.globalMcpConfig.get(name); - return this.withGlobalMcpServerProbe(server, options.cwd, (manager) => - standaloneMcpTestResult(server.name, manager), + return this.mcpManagement((management) => + management.testServer({ name, cwd: options.cwd }), ); } @@ -2568,266 +2476,11 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { server: McpServerConfig, options: { readonly cwd?: string } = {}, ): Promise { - const target = parseInlineMcpServer(server); - return this.withGlobalMcpServerProbe(target, options.cwd, (manager) => - standaloneMcpTestResult(target.name, manager), + return this.mcpManagement((management) => + management.testServer({ server, cwd: options.cwd }), ); } - private async withGlobalMcpServerProbe( - server: McpServerConfig, - cwd: string | undefined, - inspect: (manager: McpConnectionManager) => T, - oauth?: McpOAuthService, - ): Promise { - await this.configReady; - const section = this.engineAccessor.get(IConfigService).get(MCP_SECTION); - const runtimeResolver = this.engineAccessor.get(IRuntimeResolver); - let workspaceId: string | undefined; - let stdioCwd = cwd; - if (server.transport === 'stdio') { - stdioCwd = normalizeWorkDir(cwd ?? process.cwd()); - const workspace = await this.engineAccessor - .get(IWorkspaceInstanceManager) - .getOrCreate({ root: stdioCwd }); - workspaceId = workspace.id; - } - const manager = new McpConnectionManager({ - stdioCwd, - runtimeResolver, - workspaceId, - runtimeId: workspaceId === undefined ? undefined : 'local', - // Callers that just read a fresh token snapshot pass their service in; - // the cached one may have been built before the grant landed on disk. - oauthService: oauth ?? (await this.globalMcpOAuthService()), - resolveClientName: () => this.resolveMcpClientName(), - resolveDefaultTimeouts: () => ({ - startupTimeoutMs: section?.startupTimeoutMs, - toolTimeoutMs: section?.toolTimeoutMs, - }), - }); - try { - await manager.connectAll({ [server.name]: mcpConfigWithoutName(server) }); - return inspect(manager); - } finally { - await manager.shutdown(); - } - } - - /** - * The locator-addressed app-level catalog. Globals only: this engine has - * no `IPluginService.mcpServers` to flatten plugin manifests with (and the - * project layer is a workspace concern), so plugin locators resolve to - * `mcp.server_not_found` here — a v1-only capability for now. - */ - private async appMcpServerDescriptors(): Promise { - return (await this.globalMcpConfig.list()).map((server) => { - const locator = { source: 'global', name: server.name } as const; - const config = mcpConfigWithoutName(server); - return { - serverId: mcpServerId(locator), - locator, - runtimeName: server.name, - canonicalUrl: - config.transport === 'stdio' ? undefined : canonicalMcpOAuthResource(config.url), - origin: 'global' as const, - config, - enabled: config.enabled !== false, - editable: true, - }; - }); - } - - private async resolveAppMcpServer( - locator: McpServerLocator, - ): Promise { - const catalog = await this.appMcpServerDescriptors(); - return selectAppMcpServerDescriptors(catalog, [locator])[0]!; - } - - /** - * A throwaway OAuth service for one-shot credential reads: the v2 providers - * cache tokens at construction, so the cached `globalMcpOAuth` service could - * keep serving a grant that was saved or invalidated since the first read. - */ - private async freshMcpOAuthService(): Promise { - await this.engineAccessor.get(IAgentIdentity).resolved(); - return new McpOAuthService({ - store: createMcpOAuthStore(this.engineAccessor.get(IAtomicDocumentStore)), - resolveClientName: () => this.resolveMcpClientName(), - }); - } - - /** - * v1's `inspectAppMcpServerDescriptors` verbatim: a batched real-connection - * probe of every OAuth candidate (one throwaway manager for all). A runtime - * name shared by two catalog entries cannot be probed unambiguously and is - * reported `unavailable`; a stored-but-rejected grant is `oauth-expired`. - */ - private async inspectAppMcpServerDescriptors( - descriptors: readonly AppMcpServerRuntimeDescriptor[], - catalog: readonly AppMcpServerRuntimeDescriptor[], - ): Promise { - await this.configReady; - const section = this.engineAccessor.get(IConfigService).get(MCP_SECTION); - const oauth = await this.freshMcpOAuthService(); - const runtimeNameCounts = new Map(); - for (const server of new Map(catalog.map((item) => [item.serverId, item])).values()) { - if (!server.enabled) continue; - runtimeNameCounts.set(server.runtimeName, (runtimeNameCounts.get(server.runtimeName) ?? 0) + 1); - } - const credentialStates = new Map(); - const probeConfigs: Record = {}; - for (const server of descriptors) { - if (!isOAuthProbeCandidate(server)) continue; - if (runtimeNameCounts.get(server.runtimeName) !== 1) continue; - const config = requireRemoteMcpConfig(server.runtimeName, server.config); - credentialStates.set( - server.serverId, - await this.globalMcpTokenState({ name: server.runtimeName, url: config.url }, oauth), - ); - probeConfigs[server.runtimeName] = server.config; - } - let manager: McpConnectionManager | undefined; - try { - if (Object.keys(probeConfigs).length > 0) { - manager = new McpConnectionManager({ - oauthService: oauth, - resolveClientName: () => this.resolveMcpClientName(), - resolveDefaultTimeouts: () => ({ - startupTimeoutMs: section?.startupTimeoutMs, - toolTimeoutMs: section?.toolTimeoutMs, - }), - }); - await manager.connectAll(probeConfigs); - } - const checkedAt = Date.now(); - return descriptors.map((server) => { - const configured = configuredMcpAuthState(server); - if (configured !== undefined) return { ...server, authStatus: configured }; - if (runtimeNameCounts.get(server.runtimeName) !== 1) { - return { - ...server, - authStatus: 'unavailable' as const, - checkedAt, - error: `MCP runtime name "${server.runtimeName}" is not unique`, - }; - } - const tokens = credentialStates.get(server.serverId); - const entry = manager?.get(server.runtimeName); - // A clean connect only proves OAuth-authorized when a grant exists; - // a server that never challenges is simply not applicable. - if (entry?.status === 'connected') { - return { - ...server, - authStatus: tokens?.hasTokens === true ? 'oauth-authorized' : 'not-applicable', - checkedAt, - }; - } - if (entry?.status === 'needs-auth') { - return { - ...server, - authStatus: tokens?.hasTokens === true ? 'oauth-expired' : 'oauth-required', - checkedAt, - }; - } - return { - ...server, - authStatus: 'unavailable' as const, - checkedAt, - error: entry?.error ?? `MCP server finished with status ${entry?.status ?? 'unknown'}`, - }; - }); - } finally { - await manager?.shutdown(); - } - } - - /** - * v1's `mcpServerAuthState` verbatim: the offline token view plus, when the - * offline view cannot settle the state, a real connection probe. The token - * view reads the v2 provider's async `tokens()`; the `obtained_at` stamp - * (written on every save) is read via a cast — a grant without it is - * treated as non-expiring, exactly like v1's `tokenState`. - */ - private async globalMcpServerAuthState( - server: McpServerConfig, - oauth: McpOAuthService, - cwd: string | undefined, - verify: boolean, - ): Promise { - // Parity with v1: a disabled server never participates in OAuth. - if (server.enabled === false) return 'not-applicable'; - if (server.transport === 'stdio') return 'not-applicable'; - if (server.bearerTokenEnvVar !== undefined) return 'bearer-token'; - // Keep status classification aligned with the existing connection manager: - // unmarked static headers are not treated as OAuth credentials. - if (server.headers !== undefined && server.auth !== 'oauth') return 'not-applicable'; - if (server.transport !== 'http' && server.auth !== 'oauth') return 'not-applicable'; - const tokens = await this.globalMcpTokenState(server, oauth); - const offline = (): GlobalMcpServerAuthState => { - if (tokens.hasTokens) { - // An expired grant with a refresh token recovers on the next connect; - // without one the credential is dead and must be re-created. - return !tokens.expired || tokens.hasRefreshToken ? 'oauth-authorized' : 'oauth-expired'; - } - return server.auth === 'oauth' ? 'oauth-required' : 'not-applicable'; - }; - const probe = (): Promise => - this.withGlobalMcpServerProbe( - server, - cwd, - (manager) => { - const status = manager.get(server.name)?.status; - // A clean connect only proves OAuth-authorized when a grant exists; - // a server that never challenges is simply not applicable. - if (status === 'connected') return tokens.hasTokens ? 'oauth-authorized' : 'not-applicable'; - if (status === 'needs-auth') return tokens.hasTokens ? 'oauth-expired' : 'oauth-required'; - return offline(); - }, - oauth, - ); - - if (verify) { - // Online verification: a real connection probe settles states the - // offline view cannot distinguish (revoked grant, dead refresh token). - return probe(); - } - if (tokens.hasTokens) return offline(); - if (server.auth === 'oauth') return 'oauth-required'; - // Unpinned auth with no stored grant: probe once to detect whether the - // server challenges at all. - return this.withGlobalMcpServerProbe( - server, - cwd, - (manager) => - manager.get(server.name)?.status === 'needs-auth' ? 'oauth-required' : 'not-applicable', - oauth, - ); - } - - /** v1's `McpOAuthService.tokenState` over the v2 provider's async tokens. */ - private async globalMcpTokenState( - server: { readonly name: string; readonly url: string }, - oauth: McpOAuthService, - ): Promise<{ hasTokens: boolean; hasRefreshToken: boolean; expired: boolean }> { - const tokens = (await oauth.getProvider(server.name, server.url).tokens()) as - | { obtained_at?: unknown; expires_in?: unknown; refresh_token?: unknown } - | undefined; - if (tokens === undefined) { - return { hasTokens: false, hasRefreshToken: false, expired: false }; - } - const expiresAt = - typeof tokens.obtained_at === 'number' && typeof tokens.expires_in === 'number' - ? tokens.obtained_at + tokens.expires_in * 1000 - : undefined; - return { - hasTokens: true, - hasRefreshToken: typeof tokens.refresh_token === 'string' && tokens.refresh_token.length > 0, - expired: expiresAt !== undefined && Date.now() >= expiresAt, - }; - } - /** * Through the session scope (the seeded `ISessionMcpHandle.connectionManager` * — the workspace handler's one shared manager). This is a live snapshot: @@ -2910,7 +2563,8 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { readonly server: McpServerConfig; readonly persist?: boolean; }): Promise { - const mcp = this.requireLiveSession(input.sessionId).accessor.get(ISessionMcpHandle); + const session = this.requireLiveSession(input.sessionId); + const mcp = session.accessor.get(ISessionMcpHandle); const manager = mcp.connectionManager; if (!(manager instanceof McpConnectionManager)) { throw new KimiError( @@ -2923,8 +2577,9 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { // the same normalized identity, like v1's addSessionMcpServer does. const target = { ...parsed, name: normalizeServerName(parsed.name) }; if (input.persist === true) { - await this.rejectProjectLayerPersistedMcpAdd(input.sessionId, target.name); - await this.globalMcpConfig.add(target); + const cwd = session.accessor.get(ISessionWorkspaceContext).workDir; + await this.rejectProjectLayerPersistedMcpAdd(cwd, target.name); + await this.mcpManagement((management) => management.addServer(target, { cwd })); } await manager.connect(target.name, mcpConfigWithoutName(target)); const entry = manager.get(target.name); @@ -2964,6 +2619,43 @@ function normalizeRequiredWorkDir(operation: string, workDir: string): string { return normalizeWorkDir(workDir); } +/** + * Restate an engine `Error2` in the SDK's public error shape (`KimiError`, + * what `isKimiError` branches on) so the delegated management plane throws + * the same class the v1 client throws for the same failure. Non-Error2 + * failures (DI resolution bugs, aborts) pass through untouched. + * + * An engine code this build's registry does not declare (a newer engine than + * the pinned SDK) restates as `internal` — stamping the unknown code would + * mint a `KimiError` that `toKimiErrorPayload` cannot serialize (its + * `KIMI_ERROR_INFO` lookup throws on undeclared codes). + */ +function restateMcpManagementError(error: unknown): unknown { + if (!isError2(error)) return error; + const code: KimiErrorCode = isKimiErrorCode(error.code) ? error.code : ErrorCodes.INTERNAL; + return new KimiError(code, error.message, { + details: error.details as Record | undefined, + cause: error.cause, + }); +} + +/** + * v1's `toManagedServerInfo` over the engine's managed view: flatten the + * config to the top level (mutable entries carry the full values, read-only + * entries the redacted `envKeys` / `headerKeys` lists) and tag it with the + * source metadata. + */ +function toManagedServerInfo(server: McpManagedServer): McpManagedServerInfo { + return { + name: server.name, + ...server.config, + source: server.source, + origin: server.origin, + mutable: server.mutable, + plugin: server.plugin, + } as McpManagedServerInfo; +} + function describeWorkspaceMcpServer( name: string, config: WorkspaceMcpServerConfig, diff --git a/packages/node-sdk/src/types.ts b/packages/node-sdk/src/types.ts index 7b474cae0a..082f1d5d56 100644 --- a/packages/node-sdk/src/types.ts +++ b/packages/node-sdk/src/types.ts @@ -289,6 +289,7 @@ export interface AuthenticateMcpServerOptions { ) => void | boolean | PromiseLike; readonly signal?: AbortSignal; readonly timeoutMs?: number; + readonly cwd?: string; } export interface TestMcpServerOptions { diff --git a/packages/node-sdk/src/v2/global-mcp.ts b/packages/node-sdk/src/v2/global-mcp.ts index 0ed00049a6..774612c7aa 100644 --- a/packages/node-sdk/src/v2/global-mcp.ts +++ b/packages/node-sdk/src/v2/global-mcp.ts @@ -1,262 +1,24 @@ /** - * The v1 user-global MCP surface (`/mcp.json` CRUD plus the - * standalone connection probe), rebuilt for the v2 client. + * The inline/reconnect MCP config validation for the v2 client's + * session-level MCP methods (`sdk-rpc-client-v2.ts`), which ride the + * session's own connection manager and have no engine service behind them. + * The validators keep v1's exact error text for the session RPC paths, and + * validation keeps using v1's own `McpServerConfigSchema`. * - * Why a replica exists: agent-core-v2 only READS the user-global file (its - * session config loader merges it with the project files); nothing in the - * engine writes it, there is no app-scope MCP config service, and the v2 - * OAuth orchestrator / connection manager live behind the session scope — so - * the store, the `require*` guards, and the probe result shaping are ported - * here byte-for-byte from v1 (`agent-core/src/mcp/global-config.ts` and the - * helpers at the bottom of `agent-core/src/rpc/core-impl.ts`). Validation - * keeps using v1's own `McpServerConfigSchema` (the v2 schema dropped the - * `auth: 'oauth'` marker field and would strip it on write). The moving - * parts that DO have v2 counterparts — the OAuth service and the connection - * manager — are the v2 engine's own classes, instantiated by the caller - * (`sdk-rpc-client-v2.ts`); the v2 credential store persists to the same - * on-disk layout as v1 (`/credentials/mcp/-*.json`). + * Persisting session-level adds to the user-level `mcp.json` no longer + * happens here: `addSessionMcpServer`'s `persist: true` path writes through + * the engine's App-scope `IMcpConfigStore` — the single writer of that file — + * and the unified management plane (CRUD facade, connection probe, + * inspection, OAuth orchestration) delegates to the engine's + * `IMcpManagementService`. */ -import { mkdir, readFile } from 'node:fs/promises'; -import { dirname, join } from 'node:path'; - import { ErrorCodes, KimiError, McpServerConfigSchema, type GlobalMcpServerConfig, - type McpRemoteServerConfig, type McpServerConfig, } from '@moonshot-ai/agent-core'; -import type { McpConnectionManager } from '@moonshot-ai/agent-core-v2/mcpCore/connection-manager'; -import { atomicWrite } from '@moonshot-ai/agent-core-v2/_base/utils/fs'; - -import type { - AppMcpServerConfig, - AppMcpServerDescriptor, - AppMcpServerInspection, - GlobalMcpServerAuthState, - McpServerLocator, - McpTestResult, -} from '#/types'; - -interface GlobalMcpConfigFile { - readonly raw: Record; - readonly rawServers: Record; - readonly servers: readonly GlobalMcpServerConfig[]; -} - -/** Byte-identical port of v1's `GlobalMcpConfigStore`. */ -export class GlobalMcpConfigStore { - readonly path: string; - - constructor(homeDir: string) { - this.path = join(homeDir, 'mcp.json'); - } - - async list(): Promise { - return (await this.read()).servers; - } - - async get(name: string): Promise { - const normalizedName = normalizeServerName(name); - const server = (await this.read()).servers.find((entry) => entry.name === normalizedName); - if (server !== undefined) return server; - throw serverNotFound(normalizedName); - } - - async add(server: GlobalMcpServerConfig): Promise { - const normalized = parseServerInput(server); - const file = await this.read(); - if (Object.hasOwn(file.rawServers, normalized.name)) { - throw new KimiError( - ErrorCodes.REQUEST_INVALID, - `MCP server "${normalized.name}" already exists`, - ); - } - await this.write(file, { - ...file.rawServers, - [normalized.name]: persistedEntry(normalized), - }); - return this.list(); - } - - async update(server: GlobalMcpServerConfig): Promise { - const normalized = parseServerInput(server); - const file = await this.read(); - if (!Object.hasOwn(file.rawServers, normalized.name)) { - throw serverNotFound(normalized.name); - } - await this.write(file, { - ...file.rawServers, - [normalized.name]: persistedEntry(normalized), - }); - return this.list(); - } - - async remove(name: string): Promise { - const normalizedName = normalizeServerName(name); - const file = await this.read(); - if (!Object.hasOwn(file.rawServers, normalizedName)) return file.servers; - const nextServers = Object.fromEntries( - Object.entries(file.rawServers).filter(([entryName]) => entryName !== normalizedName), - ); - await this.write(file, nextServers); - return this.list(); - } - - private async read(): Promise { - let text: string; - try { - text = await readFile(this.path, 'utf-8'); - } catch (error: unknown) { - if (errorCode(error) === 'ENOENT') { - return { raw: {}, rawServers: {}, servers: [] }; - } - throw configError(`Failed to read ${this.path}: ${describeError(error)}`, error); - } - - if (text.trim().length === 0) { - return { raw: {}, rawServers: {}, servers: [] }; - } - - let parsed: unknown; - try { - parsed = JSON.parse(text) as unknown; - } catch (error: unknown) { - throw configError(`Invalid JSON in ${this.path}: ${describeError(error)}`, error); - } - if (!isRecord(parsed)) { - throw configError(`Invalid MCP config in ${this.path}: expected a JSON object`); - } - const rawServersValue = parsed['mcpServers']; - if (rawServersValue !== undefined && !isRecord(rawServersValue)) { - throw configError(`Invalid MCP config in ${this.path}: "mcpServers" must be an object`); - } - const rawServers = rawServersValue ?? {}; - const servers = Object.entries(rawServers).map(([name, value]) => parseServer(name, value)); - return { raw: parsed, rawServers, servers }; - } - - private async write( - file: GlobalMcpConfigFile, - rawServers: Record, - ): Promise { - await mkdir(dirname(this.path), { recursive: true, mode: 0o700 }); - await atomicWrite( - this.path, - `${JSON.stringify({ ...file.raw, mcpServers: rawServers }, null, 2)}\n`, - ); - } -} - -/** Byte-identical port of v1's `requireRemoteMcpConfig` guard. */ -export function requireRemoteMcpConfig( - name: string, - config: McpServerConfig, -): McpRemoteServerConfig { - if (config.transport !== 'stdio') return config; - throw new KimiError( - ErrorCodes.REQUEST_INVALID, - `MCP server "${name}" does not use a remote transport`, - ); -} - -/** Byte-identical port of v1's `requireOAuthMcpConfig` guard. */ -export function requireOAuthMcpConfig( - name: string, - input: McpServerConfig, -): McpRemoteServerConfig { - const config = requireRemoteMcpConfig(name, input); - if (config.bearerTokenEnvVar !== undefined) { - throw new KimiError( - ErrorCodes.REQUEST_INVALID, - `MCP server "${name}" uses a static bearer token`, - ); - } - if (config.headers !== undefined && config.auth !== 'oauth') { - throw new KimiError( - ErrorCodes.REQUEST_INVALID, - `MCP server "${name}" uses static headers and is not marked for OAuth`, - ); - } - return config; -} - -/** Byte-identical port of v1's `mcpServerId`. */ -export function mcpServerId(locator: McpServerLocator): string { - if (locator.source === 'global') return `global:${encodeURIComponent(locator.name)}`; - return `plugin:${encodeURIComponent(locator.pluginId)}:${encodeURIComponent(locator.serverName)}`; -} - -/** Byte-identical port of v1's `describeMcpServerLocator`. */ -export function describeMcpServerLocator(locator: McpServerLocator): string { - if (locator.source === 'global') return locator.name; - return `${locator.pluginId}/${locator.serverName}`; -} - -/** Inspection-time descriptor: the wire shape but with the full config. */ -export type AppMcpServerRuntimeDescriptor = Omit & { - readonly config: McpServerConfig; -}; - -export type AppMcpServerRuntimeInspection = AppMcpServerRuntimeDescriptor & - Pick; - -/** Byte-identical port of v1's `sanitizeAppMcpServerInspection`. */ -export function sanitizeAppMcpServerInspection( - server: AppMcpServerRuntimeInspection, -): AppMcpServerInspection { - return { ...server, config: sanitizeAppMcpServerConfig(server.config) }; -} - -/** Byte-identical port of v1's `sanitizeAppMcpServerConfig`. */ -export function sanitizeAppMcpServerConfig(config: McpServerConfig): AppMcpServerConfig { - if (config.transport === 'stdio') { - const { env, ...safe } = config; - return env === undefined ? safe : { ...safe, envKeys: Object.keys(env).toSorted() }; - } - const { headers, ...safe } = config; - return headers === undefined ? safe : { ...safe, headerKeys: Object.keys(headers).toSorted() }; -} - -/** Byte-identical port of v1's `selectAppMcpServerDescriptors`. */ -export function selectAppMcpServerDescriptors( - catalog: readonly AppMcpServerRuntimeDescriptor[], - targets?: readonly McpServerLocator[], -): readonly AppMcpServerRuntimeDescriptor[] { - if (targets === undefined) return catalog; - const byId = new Map(catalog.map((server) => [server.serverId, server])); - return targets.map((target) => { - const server = byId.get(mcpServerId(target)); - if (server !== undefined) return server; - throw new KimiError( - ErrorCodes.MCP_SERVER_NOT_FOUND, - `MCP server "${describeMcpServerLocator(target)}" was not found`, - ); - }); -} - -/** - * Byte-identical port of v1's `configuredMcpAuthState`: states decidable - * without connecting — anything pinned (stdio, bearer token, static - * non-OAuth headers) or disabled never enters the OAuth probe. - */ -export function configuredMcpAuthState( - server: AppMcpServerRuntimeDescriptor, -): GlobalMcpServerAuthState | undefined { - if (!server.enabled || server.config.enabled === false) return 'not-applicable'; - if (server.config.transport === 'stdio') return 'not-applicable'; - if (server.config.bearerTokenEnvVar !== undefined) return 'bearer-token'; - if (server.config.headers !== undefined && server.config.auth !== 'oauth') { - return 'not-applicable'; - } - return undefined; -} - -/** Byte-identical port of v1's `isOAuthProbeCandidate`. */ -export function isOAuthProbeCandidate(server: AppMcpServerRuntimeDescriptor): boolean { - return configuredMcpAuthState(server) === undefined; -} /** Byte-identical port of v1's `mcpConfigWithoutName`. */ export function mcpConfigWithoutName(server: GlobalMcpServerConfig): McpServerConfig { @@ -298,75 +60,8 @@ export function parseReconnectMcpServerConfig( return parsed.data; } -/** - * Byte-identical port of v1's `standaloneMcpTestResult`, typed against the - * v2 engine's connection manager (the probe's `get` / `resolved` reads are - * the same on both ports of the manager). - */ -export function standaloneMcpTestResult( - name: string, - manager: McpConnectionManager, -): McpTestResult { - const entry = manager.get(name); - if (entry?.status !== 'connected') { - return { - success: false, - output: - entry?.error ?? `MCP server "${name}" finished with status ${entry?.status ?? 'unknown'}`, - }; - } - const tools = manager.resolved(name)?.rawTools ?? []; - const lines = [ - `Connected to MCP server "${name}".`, - `Available tools: ${tools.length}`, - ...tools.map((tool) => `- ${tool.name}${tool.description ? `: ${tool.description}` : ''}`), - ]; - return { success: true, output: lines.join('\n') }; -} - -function parseServerInput(server: GlobalMcpServerConfig): GlobalMcpServerConfig { - return parseServer(normalizeServerName(server.name), server); -} - -function parseServer(name: string, value: unknown): GlobalMcpServerConfig { - const result = McpServerConfigSchema.safeParse(value); - if (!result.success) { - throw configError( - `Invalid MCP server "${name}" in global config: ${result.error.message}`, - result.error, - ); - } - return { name, ...result.data }; -} - -function persistedEntry(server: GlobalMcpServerConfig): McpServerConfig { - const { name: _name, ...entry } = server; - return entry; -} - export function normalizeServerName(name: string): string { const normalized = name.trim(); if (normalized.length > 0) return normalized; throw new KimiError(ErrorCodes.REQUEST_INVALID, 'MCP server name cannot be empty'); } - -function serverNotFound(name: string): KimiError { - return new KimiError(ErrorCodes.MCP_SERVER_NOT_FOUND, `MCP server "${name}" was not found`); -} - -function configError(message: string, cause?: unknown): KimiError { - return new KimiError(ErrorCodes.CONFIG_INVALID, message, { cause }); -} - -function errorCode(error: unknown): unknown { - if (!isRecord(error)) return undefined; - return error['code']; -} - -function describeError(error: unknown): string { - return error instanceof Error ? error.message : String(error); -} - -function isRecord(value: unknown): value is Record { - return typeof value === 'object' && value !== null && !Array.isArray(value); -} diff --git a/packages/node-sdk/test/sdk-rpc-client-v2.test.ts b/packages/node-sdk/test/sdk-rpc-client-v2.test.ts index 4ddd1c1734..135e772d96 100644 --- a/packages/node-sdk/test/sdk-rpc-client-v2.test.ts +++ b/packages/node-sdk/test/sdk-rpc-client-v2.test.ts @@ -24,9 +24,11 @@ import { createKimiHarnessV2, ErrorCodes, isDaemonFileUrl, + isKimiError, KimiHarness, removeProviderFromConfig, SDKRpcClientV2, + toKimiErrorPayload, type Event, type KimiConfig, } from '#/index'; @@ -34,19 +36,22 @@ import { foldAgentWireReplay } from '#/v2/resume-replay'; import { drainQueryStoreDisposals, drainSessionIndexMirror, + Error2, getLiveSessionById, HostProcessError, IAgentLifecycleService, IHostRequestHeaders, + IMcpManagementService, + IMcpOAuthService, ISessionManager, ISessionTodoService, OsProcessErrors, } from '@moonshot-ai/agent-core-v2'; import { McpOAuthService } from '../../agent-core/src/mcp/oauth/service'; +import { McpOAuthService as McpOAuthServiceV2 } from '@moonshot-ai/agent-core-v2/mcpCore/oauth/service'; import { TEST_IDENTITY } from './test-identity'; -import { startMcpAuthStatusServer } from './mcp-auth-status-server'; import { recordingTelemetry, type TelemetryRecord } from './telemetry'; const hostEnvProbe = vi.hoisted(() => ({ failWithMissingShell: false })); @@ -132,10 +137,10 @@ describe('SDKRpcClientV2 (agent-core-v2 wiring)', () => { } }); - it('reports global MCP authorization from the persisted v2 credential store', async () => { + it('reports global MCP authorization without probing when verify is false', async () => { const homeDir = await mkdtemp(join(tmpdir(), 'kimi-sdk-v2-')); tempDirs.push(homeDir); - const statusServer = await startMcpAuthStatusServer(); + const implicitOAuthUrl = 'https://implicit-oauth.example.test/mcp'; const authorizedUrl = 'https://authorized.example.test/mcp'; const requiredUrl = 'https://required.example.test/mcp'; const externalOAuth = new McpOAuthService({ kimiHomeDir: homeDir }); @@ -143,17 +148,17 @@ describe('SDKRpcClientV2 (agent-core-v2 wiring)', () => { .getProvider('oauth-authorized', authorizedUrl) .saveTokens({ access_token: 'test-access-token', token_type: 'Bearer' }); await externalOAuth - .getProvider('sse', statusServer.oauthUrl) + .getProvider('sse', implicitOAuthUrl) .saveTokens({ access_token: 'stale-sse-token', token_type: 'Bearer' }); await writeFile( join(homeDir, 'mcp.json'), JSON.stringify({ mcpServers: { stdio: { command: 'local-command' }, - plain: { transport: 'http', url: statusServer.plainUrl }, - detected: { transport: 'http', url: statusServer.oauthUrl }, - sse: { transport: 'sse', url: statusServer.oauthUrl }, - 'sse-oauth': { transport: 'sse', url: statusServer.oauthUrl, auth: 'oauth' }, + plain: { transport: 'http', url: 'https://plain.example.test/mcp' }, + detected: { transport: 'http', url: implicitOAuthUrl }, + sse: { transport: 'sse', url: implicitOAuthUrl }, + 'sse-oauth': { transport: 'sse', url: implicitOAuthUrl, auth: 'oauth' }, bearer: { transport: 'http', url: 'https://bearer.example.test/mcp', @@ -176,10 +181,10 @@ describe('SDKRpcClientV2 (agent-core-v2 wiring)', () => { const harness = createKimiHarnessV2({ homeDir, identity: TEST_IDENTITY }); try { - await expect(harness.listMcpServerAuthStatuses()).resolves.toEqual([ + await expect(harness.listMcpServerAuthStatuses({ verify: false })).resolves.toEqual([ { name: 'stdio', authStatus: 'not-applicable' }, { name: 'plain', authStatus: 'not-applicable' }, - { name: 'detected', authStatus: 'oauth-required' }, + { name: 'detected', authStatus: 'not-applicable' }, { name: 'sse', authStatus: 'not-applicable' }, { name: 'sse-oauth', authStatus: 'oauth-required' }, { name: 'bearer', authStatus: 'bearer-token' }, @@ -192,10 +197,10 @@ describe('SDKRpcClientV2 (agent-core-v2 wiring)', () => { .saveTokens({ access_token: 'new-test-access-token', token_type: 'Bearer' }); await externalOAuth.invalidate('oauth-authorized', authorizedUrl, 'tokens'); - await expect(harness.listMcpServerAuthStatuses()).resolves.toEqual([ + await expect(harness.listMcpServerAuthStatuses({ verify: false })).resolves.toEqual([ { name: 'stdio', authStatus: 'not-applicable' }, { name: 'plain', authStatus: 'not-applicable' }, - { name: 'detected', authStatus: 'oauth-required' }, + { name: 'detected', authStatus: 'not-applicable' }, { name: 'sse', authStatus: 'not-applicable' }, { name: 'sse-oauth', authStatus: 'oauth-required' }, { name: 'bearer', authStatus: 'bearer-token' }, @@ -204,10 +209,116 @@ describe('SDKRpcClientV2 (agent-core-v2 wiring)', () => { ]); } finally { await harness.close(); - await statusServer.close(); } }, 15_000); + it('restates engine MCP management Error2s as KimiError, undeclared codes as internal', async () => { + const homeDir = await mkdtemp(join(tmpdir(), 'kimi-sdk-v2-')); + tempDirs.push(homeDir); + const client = new SDKRpcClientV2({ homeDir, identity: TEST_IDENTITY }); + try { + const management = client.engineAccessor.get(IMcpManagementService); + const listSpy = vi.spyOn(management, 'listServers'); + const captureRejection = async (promise: Promise): Promise => { + try { + await promise; + } catch (error) { + return error; + } + return expect.unreachable('expected the call to reject'); + }; + try { + // A declared engine code keeps its identity across the restate: the + // SDK throws the same class v1 throws, and the payload serializer + // accepts the code. + listSpy.mockRejectedValueOnce( + new Error2('mcp.oauth_failed', 'OAuth flow timed out', { + details: { flowId: 'flow-1' }, + }), + ); + const oauthError = await captureRejection(client.listGlobalMcpServers()); + expect(isKimiError(oauthError)).toBe(true); + expect(oauthError).toMatchObject({ + code: 'mcp.oauth_failed', + message: 'OAuth flow timed out', + details: { flowId: 'flow-1' }, + }); + expect(toKimiErrorPayload(oauthError)).toMatchObject({ + code: 'mcp.oauth_failed', + message: 'OAuth flow timed out', + }); + + // A code this build's registry does not declare (a newer engine than + // the pinned SDK) restates as `internal` instead of minting an + // undeclared KimiError code the serializer would reject. + listSpy.mockRejectedValueOnce(new Error2('mcp.future_code' as never, 'from a newer engine')); + const unknownError = await captureRejection(client.listGlobalMcpServers()); + expect(isKimiError(unknownError)).toBe(true); + expect(unknownError).toMatchObject({ + code: ErrorCodes.INTERNAL, + message: 'from a newer engine', + }); + expect(toKimiErrorPayload(unknownError)).toMatchObject({ + code: ErrorCodes.INTERNAL, + message: 'from a newer engine', + }); + } finally { + listSpy.mockRestore(); + } + } finally { + await client.close(); + } + }); + + it('close() awaits the MCP OAuth service shutdown', async () => { + const homeDir = await mkdtemp(join(tmpdir(), 'kimi-sdk-v2-')); + tempDirs.push(homeDir); + const client = new SDKRpcClientV2({ homeDir, identity: TEST_IDENTITY }); + // Activate the OnDemand OAuth service, then gate its shutdown behind a + // manual release: close() alone (no manual service.shutdown()) must + // trigger and await that shutdown, so a host removing homeDir right after + // close() cannot race in-flight token writes. + client.engineAccessor.get(IMcpOAuthService); + let releaseShutdown: () => void = () => undefined; + const gate = new Promise((resolve) => { + releaseShutdown = resolve; + }); + const baseShutdown = McpOAuthServiceV2.prototype.shutdown; + const shutdownSpy = vi + .spyOn(McpOAuthServiceV2.prototype, 'shutdown') + .mockImplementation(function (this: McpOAuthServiceV2) { + return baseShutdown.call(this).then(() => gate); + }); + try { + let closed = false; + const closePromise = client.close().then(() => { + closed = true; + }); + await vi.waitFor(() => { + expect(shutdownSpy).toHaveBeenCalled(); + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(closed).toBe(false); + releaseShutdown(); + await closePromise; + expect(closed).toBe(true); + } finally { + releaseShutdown(); + shutdownSpy.mockRestore(); + } + }); + + it('close() resolves promptly when the MCP OAuth service was never used', async () => { + const homeDir = await mkdtemp(join(tmpdir(), 'kimi-sdk-v2-')); + tempDirs.push(homeDir); + const client = new SDKRpcClientV2({ homeDir, identity: TEST_IDENTITY }); + // Nothing touched IMcpOAuthService: close() force-activates the OnDemand + // service only to shut it down, and that activate-then-shutdown cycle + // must be a clean no-op (the proactive-refresh sweep bows out on the + // shutdown flag). + await expect(client.close()).resolves.toBeUndefined(); + }); + it('seeds the host request headers (User-Agent + X-Msh-*) into the engine', async () => { const homeDir = await mkdtemp(join(tmpdir(), 'kimi-sdk-v2-')); tempDirs.push(homeDir); diff --git a/packages/node-sdk/test/v1-v2-parity.test.ts b/packages/node-sdk/test/v1-v2-parity.test.ts index 550752945c..e151e360e2 100644 --- a/packages/node-sdk/test/v1-v2-parity.test.ts +++ b/packages/node-sdk/test/v1-v2-parity.test.ts @@ -30,6 +30,7 @@ import { createKimiHarness, createKimiHarnessV2, ErrorCodes, + isKimiError, SDKRpcClient, SDKRpcClientV2, type ApprovalRequest, @@ -3713,8 +3714,9 @@ async function captureRejection(promise: Promise): Promise { } /** - * Both engines must reject with the same code and the same message (home - * prefixes scrubbed — the file-store errors embed the mcp.json path). + * Both engines must reject with the same class, code, and message (home + * prefixes scrubbed — the file-store errors embed the mcp.json path). The + * class is pinned because SDK consumers branch on `isKimiError`. */ async function expectSameMcpRejection( pair: GlobalMcpParityPair, @@ -3726,8 +3728,12 @@ async function expectSameMcpRejection( captureRejection(v2Call(pair.v2)), ]); const payload = (error: unknown): unknown => { - const err = error as { code?: unknown; message?: unknown }; - return { code: err.code ?? null, message: String(err.message ?? error) }; + const err = error as { name?: unknown; code?: unknown; message?: unknown }; + return { + name: err.name ?? null, + code: err.code ?? null, + message: String(err.message ?? error), + }; }; expect(scrubHomePrefixes(payload(v2Error), pair.v2Home)).toEqual( scrubHomePrefixes(payload(v1Error), pair.v1Home), @@ -3750,7 +3756,7 @@ function expectSameManagedServers( } describe('v1↔v2 global MCP parity', () => { - it('classifies global MCP authorization identically from persisted credentials', async () => { + it('detects implicit OAuth requirements identically by default', async () => { const statusServer = await startMcpAuthStatusServer(); const authorizedUrl = 'https://authorized.example.test/mcp'; const pair = await makeGlobalMcpParityPair({ @@ -3798,7 +3804,6 @@ describe('v1↔v2 global MCP parity', () => { pair.v1.listGlobalMcpServerAuthStatuses(), pair.v2.listGlobalMcpServerAuthStatuses(), ]); - expect(v2Statuses).toEqual(v1Statuses); expect(v1Statuses).toEqual([ { name: 'stdio', authStatus: 'not-applicable' }, { name: 'plain', authStatus: 'not-applicable' }, @@ -3810,6 +3815,7 @@ describe('v1↔v2 global MCP parity', () => { { name: 'oauth-authorized', authStatus: 'oauth-authorized' }, { name: 'disabled-oauth', authStatus: 'not-applicable' }, ]); + expect(v2Statuses).toEqual(v1Statuses); } finally { await closeGlobalMcpPair(pair); await statusServer.close(); @@ -3900,14 +3906,12 @@ describe('v1↔v2 global MCP parity', () => { { name: 'oauth-required', authStatus: 'oauth-required' }, ]); - // The legacy name-based list stays offline (verify is opt-in): a - // stored grant is `oauth-authorized` even when the server would reject - // it — the deliberate offline false positive. + // The name-based list preserves implicit no-grant detection when verify + // is omitted, while stored grants are classified from disk. const [v1LegacyStatuses, v2LegacyStatuses] = await Promise.all([ pair.v1.listGlobalMcpServerAuthStatuses(), pair.v2.listGlobalMcpServerAuthStatuses(), ]); - expect(v2LegacyStatuses).toEqual(v1LegacyStatuses); expect(v1LegacyStatuses).toEqual([ { name: 'stdio', authStatus: 'not-applicable' }, { name: 'plain', authStatus: 'not-applicable' }, @@ -3919,6 +3923,7 @@ describe('v1↔v2 global MCP parity', () => { { name: 'unavailable-explicit', authStatus: 'oauth-required' }, { name: 'unavailable-dynamic', authStatus: 'not-applicable' }, ]); + expect(v2LegacyStatuses).toEqual(v1LegacyStatuses); } finally { await closeGlobalMcpPair(pair); await statusServer.close(); @@ -4156,6 +4161,100 @@ describe('v1↔v2 global MCP parity', () => { } }); + it.each([ + [ + 'add', + (client: SDKRpcClient | SDKRpcClientV2, cwd: string) => + client.addGlobalMcpServer( + { name: 'project', transport: 'stdio', command: 'replacement' }, + { cwd }, + ), + ], + [ + 'update', + (client: SDKRpcClient | SDKRpcClientV2, cwd: string) => + client.updateGlobalMcpServer( + { name: 'project', transport: 'stdio', command: 'replacement' }, + { cwd }, + ), + ], + [ + 'remove', + (client: SDKRpcClient | SDKRpcClientV2, cwd: string) => + client.removeGlobalMcpServer('project', { cwd }), + ], + ])('%s rejects a trusted project-layer entry as read-only on both engines', async (_operation, mutate) => { + const pair = await makeGlobalMcpParityPair(); + const project = await makeTempDir('kimi-sdk-parity-mcp-project-'); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ mcpServers: { project: { command: 'project-command' } } }), + 'utf-8', + ); + try { + await pair.v2.trustWorkspace(project); + + await expectSameMcpRejection( + pair, + (client) => mutate(client, project), + (client) => mutate(client, project), + ); + } finally { + await closeGlobalMcpPair(pair); + } + }); + + it('threads cwd through project-layer inspection and authorization on both engines', async () => { + const pair = await makeGlobalMcpParityPair(); + const project = await makeTempDir('kimi-sdk-parity-mcp-auth-project-'); + await mkdir(join(project, '.kimi-code'), { recursive: true }); + await writeFile( + join(project, '.kimi-code', 'mcp.json'), + JSON.stringify({ + mcpServers: { + 'project-stdio': { command: 'project-command' }, + 'project-oauth': { + transport: 'http', + url: 'https://project.example.test/mcp', + auth: 'oauth', + }, + }, + }), + 'utf-8', + ); + try { + await pair.v2.trustWorkspace(project); + + const [v1Inspections, v2Inspections] = await Promise.all([ + pair.v1.inspectAppMcpServers([{ source: 'global', name: 'project-stdio' }], { + cwd: project, + }), + pair.v2.inspectAppMcpServers([{ source: 'global', name: 'project-stdio' }], { + cwd: project, + }), + ]); + const summarize = (inspections: typeof v1Inspections) => + inspections.map(({ runtimeName, authStatus }) => ({ runtimeName, authStatus })); + expect(summarize(v2Inspections)).toEqual(summarize(v1Inspections)); + expect(summarize(v1Inspections)).toEqual([ + { runtimeName: 'project-stdio', authStatus: 'not-applicable' }, + ]); + + await expectSameMcpRejection( + pair, + (client) => client.beginGlobalMcpServerAuth('project-stdio', { cwd: project }), + (client) => client.beginGlobalMcpServerAuth('project-stdio', { cwd: project }), + ); + await Promise.all([ + pair.v1.resetGlobalMcpServerAuth('project-oauth', { cwd: project }), + pair.v2.resetGlobalMcpServerAuth('project-oauth', { cwd: project }), + ]); + } finally { + await closeGlobalMcpPair(pair); + } + }); + it('a malformed mcp.json rejects every read with the same config.invalid', async () => { const pair = await makeGlobalMcpParityPair(); await writeFile(join(pair.v1HomeDir, 'mcp.json'), '{ not valid json', 'utf-8'); @@ -4647,6 +4746,50 @@ describe('v1↔v2 session MCP parity', () => { restoreEnv(); } }, 20_000); + + it('a persisted session add over an enabled plugin entry follows each engine\'s precedence', async () => { + const restoreEnv = scrubConfigEnv(); + const pair = await makeSessionMcpPair(); + const pluginSource = await makeTempDir('kimi-sdk-parity-mcp-plugin-src-'); + await writeFixturePlugin(pluginSource); + try { + await Promise.all([ + pair.v1.installPlugin(pluginSource), + pair.v2.installPlugin(pluginSource), + ]); + await createOnBoth(pair, { id: 'session_parity_mcp_plugin_shadow' }); + const input = { sessionId: 'session_parity_mcp_plugin_shadow' } as const; + const server: McpServerConfig = { + name: 'plugin-parity-plugin:parity-stdio', + transport: 'stdio', + command: process.execPath, + args: [MCP_STDIO_FIXTURE], + }; + + // Deliberate divergence: v1 ranks an enabled plugin above the file + // layers, so a persisted add could never take effect and rejects + // (`isKimiError` is the SDK's public error class on both engines). v2 + // keeps the file layers above plugins, so the same add is a real + // override: the user-level write lands and shadows the plugin. + const [v1Error, v2Info] = await Promise.all([ + captureRejection(pair.v1.addSessionMcpServer({ ...input, server, persist: true })), + pair.v2.addSessionMcpServer({ ...input, server, persist: true }), + ]); + expect(isKimiError(v1Error)).toBe(true); + expect(v1Error).toMatchObject({ code: 'request.invalid' }); + expect(v2Info.name).toBe('plugin-parity-plugin:parity-stdio'); + + const [v1File, v2File] = await Promise.all([ + readFile(join(pair.v1Home.raw, 'mcp.json'), 'utf-8').catch(() => ''), + readFile(join(pair.v2Home.raw, 'mcp.json'), 'utf-8').catch(() => ''), + ]); + expect(v1File).not.toContain('plugin-parity-plugin:parity-stdio'); + expect(v2File).toContain('plugin-parity-plugin:parity-stdio'); + } finally { + await closeSessionPair(pair); + restoreEnv(); + } + }, 20_000); }); // --------------------------------------------------------------------------- diff --git a/packages/oauth/src/oauth-token-transaction.ts b/packages/oauth/src/oauth-token-transaction.ts index 5062f4ec42..ac77faa5c3 100644 --- a/packages/oauth/src/oauth-token-transaction.ts +++ b/packages/oauth/src/oauth-token-transaction.ts @@ -21,6 +21,24 @@ export interface OAuthTokenTransactionOptions { readonly remove: () => Promise; readonly parse: (value: unknown) => T | undefined; readonly adopt?: (tokens: T | undefined) => void; + /** + * Strips caller-local fields before a remembered effect is compared with an + * incoming SDK save, so a durable-only stamp (e.g. `obtained_at`) does not + * defeat effect coalescing. The remembered effect keeps its original tokens. + */ + readonly normalize?: (tokens: T) => T; + /** + * Runs inside the exclusive lock right after a durable write or remove + * commits, so side effects (sidecar files, notifications) serialize with the + * token write instead of racing a concurrent clear. + */ + readonly afterCommit?: (tokens: T | undefined) => Promise; + /** + * Receives the in-flight promise of every token-grant request and of every + * exclusive save/invalidate commit, so callers can await transport-driven + * grants and their trailing durable commits during shutdown. + */ + readonly track?: (operation: Promise) => void; } /** @@ -45,25 +63,43 @@ export class OAuthTokenTransaction { ) { return fetchFn(input, init); } - return transactionLock.runExclusive(this.options.key, () => + const request = transactionLock.runExclusive(this.options.key, () => this.runTokenRequest(fetchFn, input, init, params, grantType), ); + // Track past the response headers until the body has fully arrived: the + // SDK's trailing parse + saveTokens continuation can only run after it + // reads the body, so a shutdown drain that outlives this promise also + // covers the durable commit (which is itself tracked from save()). + const operation = request.then((response) => + response.clone().arrayBuffer().catch(() => undefined), + ); + this.options.track?.(operation); + return request; }) as typeof fetch; } - async save(tokens: T): Promise { - await transactionLock.runExclusive(this.options.key, async () => { - if (this.consumeSave(tokens)) { - this.adopt(await this.options.read()); + async save(tokens: T): Promise { + let persisted: T | undefined; + const operation = transactionLock.runExclusive(this.options.key, async () => { + const pending = this.consumeSave(tokens); + if (pending !== undefined) { + persisted = await this.options.read(); + this.adopt(persisted); + if (persisted !== undefined) await this.options.afterCommit?.(persisted); return; } await this.options.write(tokens); this.adopt(tokens); + persisted = tokens; + await this.options.afterCommit?.(tokens); }); + this.options.track?.(operation); + await operation; + return persisted; } async invalidateFromSdk(scope: 'tokens' | 'all'): Promise { - return transactionLock.runExclusive(this.options.key, async () => { + const operation = transactionLock.runExclusive(this.options.key, async () => { const effect = this.takeInvalidate(scope); if (effect === undefined) return false; const current = await this.options.read(); @@ -82,14 +118,18 @@ export class OAuthTokenTransaction { } await this.options.remove(); this.adopt(undefined); + await this.options.afterCommit?.(undefined); return true; }); + this.options.track?.(operation); + return operation; } async clear(): Promise { await transactionLock.runExclusive(this.options.key, async () => { await this.options.remove(); this.adopt(undefined); + await this.options.afterCommit?.(undefined); }); } @@ -187,13 +227,18 @@ export class OAuthTokenTransaction { this.effects.push(effect); } - private consumeSave(tokens: T): boolean { + private consumeSave(tokens: T): T | undefined { + const normalize = this.options.normalize ?? ((value: T) => value); + const normalized = normalize(tokens); const index = this.effects.findIndex( - (effect) => effect.kind === 'save' && isDeepStrictEqual(effect.tokens, tokens), + (effect) => + effect.kind === 'save' && + (isDeepStrictEqual(normalize(effect.tokens), normalized) || + sameRefreshSave(normalize(effect.tokens), normalized)), ); - if (index === -1) return false; - this.effects.splice(index, 1); - return true; + if (index === -1) return undefined; + const effect = this.effects.splice(index, 1)[0] as Extract, { kind: 'save' }>; + return effect.tokens; } private takeInvalidate(scope: 'tokens' | 'all'): Extract, { kind: 'invalidate' }> | undefined { @@ -209,6 +254,20 @@ export class OAuthTokenTransaction { } } +function sameRefreshSave(expected: T, actual: T): boolean { + const expectedRecord = expected as Record; + const actualRecord = actual as Record; + if ( + refreshToken(expected) === undefined || + ('refresh_token' in actualRecord && actualRecord['refresh_token'] !== undefined) + ) { + return false; + } + const { refresh_token: _expectedRefreshToken, ...expectedWithoutRefresh } = expectedRecord; + const { refresh_token: _actualRefreshToken, ...actualWithoutRefresh } = actualRecord; + return isDeepStrictEqual(expectedWithoutRefresh, actualWithoutRefresh); +} + function refreshToken(tokens: object | undefined): string | undefined { if (tokens === undefined || !('refresh_token' in tokens)) return undefined; return typeof tokens.refresh_token === 'string' ? tokens.refresh_token : undefined; diff --git a/packages/oauth/test/oauth-token-transaction.test.ts b/packages/oauth/test/oauth-token-transaction.test.ts index 4adbc7bbb7..4cc9f3c575 100644 --- a/packages/oauth/test/oauth-token-transaction.test.ts +++ b/packages/oauth/test/oauth-token-transaction.test.ts @@ -5,6 +5,7 @@ import { OAuthTokenTransaction } from '../src/oauth-token-transaction'; interface TestTokens { access_token: string; refresh_token?: string; + obtained_at?: number; } describe('OAuthTokenTransaction', () => { @@ -41,6 +42,19 @@ describe('OAuthTokenTransaction', () => { expect(stored).toEqual(tokens('access-1', 'refresh-1')); }); + it('preserves a refresh token when the SDK save omits it after refresh', async () => { + let stored: TestTokens | undefined = tokens('access-0', 'refresh-0'); + const subject = transaction('same-server', () => stored, (value) => (stored = value)); + const response = await subject.createFetch(async () => json({ access_token: 'access-1' }))( + 'https://issuer.example.test/token', + refreshRequest('refresh-0'), + ); + + await subject.save((await response.json()) as TestTokens); + + expect(stored).toEqual(tokens('access-1', 'refresh-0')); + }); + it('preserves an access-only winner when an older refresh is queued', async () => { let stored: TestTokens | undefined = { access_token: 'access-from-login' }; const stale = transaction('same-server', () => stored, (value) => (stored = value)); @@ -110,12 +124,68 @@ describe('OAuthTokenTransaction', () => { await expect(subject.invalidateFromSdk('tokens')).resolves.toBe(false); expect(stored).toEqual(tokens('access-0', 'refresh-0')); }); + + it('coalesces a stripped SDK save against a stamped durable winner without a second write', async () => { + let stored: TestTokens | undefined = { + access_token: 'access-1', + refresh_token: 'refresh-1', + obtained_at: 123, + }; + let writes = 0; + const subject = transaction( + 'same-server', + () => stored, + (value) => { + writes += 1; + stored = value; + }, + { normalize: stripDurableStamp }, + ); + const tokenEndpoint = vi.fn(); + + const response = await subject.createFetch(tokenEndpoint)( + 'https://issuer.example.test/token', + refreshRequest('refresh-0'), + ); + + expect(response.status).toBe(200); + expect(tokenEndpoint).not.toHaveBeenCalled(); + const stripped = parseTokens(await response.json()); + if (stripped === undefined) throw new Error('invalid test token response'); + await subject.save(stripped); + + expect(writes).toBe(0); + expect(stored).toEqual({ access_token: 'access-1', refresh_token: 'refresh-1', obtained_at: 123 }); + }); + + it('does not resurrect a cleared credential when a stripped SDK save arrives late', async () => { + let stored: TestTokens | undefined = { + access_token: 'access-1', + refresh_token: 'refresh-1', + obtained_at: 123, + }; + const subject = transaction('same-server', () => stored, (value) => (stored = value), { + normalize: stripDurableStamp, + }); + + const response = await subject.createFetch(vi.fn())( + 'https://issuer.example.test/token', + refreshRequest('refresh-0'), + ); + const stripped = parseTokens(await response.json()); + if (stripped === undefined) throw new Error('invalid test token response'); + + await subject.clear(); + await subject.save(stripped); + expect(stored).toBeUndefined(); + }); }); function transaction( key: string, read: () => TestTokens | undefined, write: (tokens: TestTokens | undefined) => void, + extras?: { readonly normalize?: (tokens: TestTokens) => TestTokens }, ): OAuthTokenTransaction { return new OAuthTokenTransaction({ key, @@ -127,9 +197,15 @@ function transaction( write(undefined); }, parse: parseTokens, + normalize: extras?.normalize, }); } +function stripDurableStamp(value: TestTokens): TestTokens { + const { obtained_at: _omitted, ...stripped } = value; + return stripped; +} + async function sdkRefresh( transaction: OAuthTokenTransaction, fetchFn: typeof fetch,