From 4805c7bebe767e141aea2f8490cea6ee49b8d756 Mon Sep 17 00:00:00 2001 From: 7Sageer <7sageer@djwcb.cn> Date: Wed, 19 Aug 2026 21:12:01 +0800 Subject: [PATCH] fix(agent-core-v2): harden MCP management readiness --- .../respect-mcp-management-readiness.md | 5 + .../src/app/mcpConfig/configStore.ts | 1 - .../src/app/mcpConfig/oauthService.ts | 2 +- .../app/mcpManagement/mcpManagementService.ts | 47 ++----- .../src/app/mcpRegistry/mcpRegistryService.ts | 43 ++++-- .../src/mcpCore/oauth/provider.ts | 7 +- .../src/mcpCore/oauth/service.ts | 69 ++++++---- .../workspaceMcp/workspaceMcpService.ts | 7 - .../workspace/workspaceTrust/trustRecord.ts | 35 +++++ .../workspaceTrust/workspaceTrustService.ts | 25 +--- .../test/app/mcpConfig/oauthService.test.ts | 59 ++++++++ .../app/mcpManagement/mcpManagement.test.ts | 129 +++++++++++++++++- .../test/app/mcpRegistry/mcpRegistry.test.ts | 28 ++++ .../test/mcpCore/oauth/service.test.ts | 52 ++----- packages/agent-core-v2/test/mcpCore/stubs.ts | 39 ++++++ .../workspaceMcp/workspaceMcp.test.ts | 39 +++--- .../workspaceMcpConfig.test.ts | 24 +++- 17 files changed, 439 insertions(+), 172 deletions(-) create mode 100644 .changeset/respect-mcp-management-readiness.md create mode 100644 packages/agent-core-v2/src/workspace/workspaceTrust/trustRecord.ts create mode 100644 packages/agent-core-v2/test/app/mcpConfig/oauthService.test.ts diff --git a/.changeset/respect-mcp-management-readiness.md b/.changeset/respect-mcp-management-readiness.md new file mode 100644 index 000000000..98032fb55 --- /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/packages/agent-core-v2/src/app/mcpConfig/configStore.ts b/packages/agent-core-v2/src/app/mcpConfig/configStore.ts index 97c19607e..0b56c1e6b 100644 --- a/packages/agent-core-v2/src/app/mcpConfig/configStore.ts +++ b/packages/agent-core-v2/src/app/mcpConfig/configStore.ts @@ -45,7 +45,6 @@ export class McpConfigStore extends Disposable implements IMcpConfigStore { private readonly writeEmitter = this._register(new Emitter()); readonly onDidWrite: Event = this.writeEmitter.event; - /** Serializes the read-modify-write mutations so concurrent writes cannot lose updates. */ private mutationTail: Promise = Promise.resolve(); constructor( diff --git a/packages/agent-core-v2/src/app/mcpConfig/oauthService.ts b/packages/agent-core-v2/src/app/mcpConfig/oauthService.ts index 581159393..27459df2d 100644 --- a/packages/agent-core-v2/src/app/mcpConfig/oauthService.ts +++ b/packages/agent-core-v2/src/app/mcpConfig/oauthService.ts @@ -21,7 +21,7 @@ export class AppMcpOAuthService extends McpOAuthService { resolveClientName: () => identity.current().slug, log, }); - void this.sweepProactiveRefresh().catch((error: unknown) => { + void identity.resolved().then(() => this.sweepProactiveRefresh()).catch((error: unknown) => { log.warn(`mcp oauth proactive-refresh sweep failed: ${String(error)}`); }); } diff --git a/packages/agent-core-v2/src/app/mcpManagement/mcpManagementService.ts b/packages/agent-core-v2/src/app/mcpManagement/mcpManagementService.ts index 79ee7b936..8bb96dd7f 100644 --- a/packages/agent-core-v2/src/app/mcpManagement/mcpManagementService.ts +++ b/packages/agent-core-v2/src/app/mcpManagement/mcpManagementService.ts @@ -54,7 +54,6 @@ const DEFAULT_AUTH_TIMEOUT_MS = 15 * 60_000; export class McpManagementService extends Disposable implements IMcpManagementService { declare readonly _serviceBrand: undefined; - /** In-flight management-plane OAuth flows by flowId. */ private readonly authFlows = new Map(); constructor( @@ -109,18 +108,13 @@ export class McpManagementService extends Disposable implements IMcpManagementSe } 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), ); } - /** - * Mutation guard lookup: only a genuine not-found reads as "no collision". - * A plugin-state or config failure must abort the write — a user-level - * mutation guarded on a degraded view could shadow a read-only plugin - * server while the plugin contributions are unknown. - */ private async guardLookup(name: string): Promise { try { return await this.registry.get(name); @@ -130,11 +124,6 @@ export class McpManagementService extends Disposable implements IMcpManagementSe } } - /** - * Test target resolution: an inline `server` config probes as-is (nothing - * has to be saved first); a bare `name` goes through the unified registry, - * so plugin and project-layer servers are testable too. - */ private async resolveTestTarget(target: McpServerTestTarget): Promise { const { name, server, cwd } = target; if (server !== undefined) { @@ -179,6 +168,7 @@ export class McpManagementService extends Disposable implements IMcpManagementSe 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; @@ -209,6 +199,7 @@ export class McpManagementService extends Disposable implements IMcpManagementSe } async listAuthStatuses(query: McpAuthStatusQuery = {}): Promise { + await this.waitForReadiness(); const entries = await this.registry.list({ cwd: query.cwd }); const verify = query.verify === true; return Promise.all( @@ -222,6 +213,7 @@ export class McpManagementService extends Disposable implements IMcpManagementSe async inspectServers( targets?: readonly McpServerLocator[], ): Promise { + await this.waitForReadiness(); const catalog = await this.serverDescriptors(); const descriptors = selectServerDescriptors(catalog, targets); const inspections = await this.inspectServerDescriptors(descriptors, catalog); @@ -241,6 +233,7 @@ export class McpManagementService extends Disposable implements IMcpManagementSe } async beginServerAuth(locator: McpServerLocator): Promise { + await this.waitForReadiness(); const server = await this.resolveServer(locator); const config = requireOAuthMcpConfig(server.runtimeName, server.config); try { @@ -286,12 +279,12 @@ export class McpManagementService extends Disposable implements IMcpManagementSe } async resetServerAuth(locator: McpServerLocator): Promise { + await this.waitForReadiness(); const server = await this.resolveServer(locator); const config = requireRemoteMcpConfig(server.runtimeName, server.config); await this.oauth.invalidate(server.runtimeName, config.url); } - /** The registry catalog in the locator-addressed shape, with full configs. */ private async serverDescriptors(): Promise { return (await this.registry.list()).map((entry) => serverDescriptor(entry)); } @@ -305,10 +298,6 @@ export class McpManagementService extends Disposable implements IMcpManagementSe return server; } - /** - * A runtime name shared by another enabled entry makes the OAuth credential - * identity ambiguous; refuse to guess. - */ private requireUnambiguousRuntimeName( catalog: readonly McpServerRuntimeDescriptor[], server: McpServerRuntimeDescriptor, @@ -327,10 +316,6 @@ export class McpManagementService extends Disposable implements IMcpManagementSe } } - /** - * States decidable without connecting: anything pinned (stdio, bearer token, - * static non-OAuth headers) or disabled never enters the OAuth probe. - */ private async serverAuthState( entry: McpRegistryEntry, cwd: string | undefined, @@ -358,22 +343,9 @@ export class McpManagementService extends Disposable implements IMcpManagementSe return offline(); }); - if (verify) { - return probe(); - } - if (tokens.hasTokens) return offline(); - if (server.auth === 'oauth') return 'oauth-required'; - return this.withProbe({ name: entry.name, ...server }, cwd, (manager) => - manager.get(entry.name)?.status === 'needs-auth' ? 'oauth-required' : 'not-applicable', - ); + return verify ? probe() : offline(); } - /** - * Inspection = registry catalog + a batched real-connection probe of every - * OAuth candidate (one throwaway manager for all). A runtime name shared by - * a global and a plugin entry cannot be probed unambiguously and is - * reported `unavailable`; a stored-but-rejected grant is `oauth-expired`. - */ private async inspectServerDescriptors( descriptors: readonly McpServerRuntimeDescriptor[], catalog: readonly McpServerRuntimeDescriptor[], @@ -449,6 +421,11 @@ export class McpManagementService extends Disposable implements IMcpManagementSe await manager?.shutdown(); } } + + private async waitForReadiness(): Promise { + await this.config.ready; + await this.identity.resolved(); + } } function throwReadOnlyMcpServer(entry: McpRegistryEntry): void { diff --git a/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistryService.ts b/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistryService.ts index b3053bbd4..8b167cb46 100644 --- a/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistryService.ts +++ b/packages/agent-core-v2/src/app/mcpRegistry/mcpRegistryService.ts @@ -7,6 +7,8 @@ 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, @@ -22,6 +24,7 @@ export class McpRegistryService implements IMcpRegistryService { @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 { @@ -40,20 +43,34 @@ export class McpRegistryService implements IMcpRegistryService { }); } } else { - const detailed = await loadMcpServersDetailed({ - fs: this.fs, - cwd: query.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, + if (!(await readWorkspaceTrust(this.docs, query.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: query.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, + }); + } } } diff --git a/packages/agent-core-v2/src/mcpCore/oauth/provider.ts b/packages/agent-core-v2/src/mcpCore/oauth/provider.ts index 70736bf06..1b9f8da68 100644 --- a/packages/agent-core-v2/src/mcpCore/oauth/provider.ts +++ b/packages/agent-core-v2/src/mcpCore/oauth/provider.ts @@ -41,6 +41,7 @@ export interface McpOAuthProviderOptions { readonly store: McpOAuthStore; readonly clientLabel?: string; readonly clientName?: string; + 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. */ @@ -58,6 +59,7 @@ export class McpOAuthClientProvider implements OAuthClientProvider { 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; @@ -77,6 +79,7 @@ export class McpOAuthClientProvider implements OAuthClientProvider { `${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}`; this.tokenTransaction = new OAuthTokenTransaction({ key: this.storeKey, @@ -85,7 +88,7 @@ export class McpOAuthClientProvider implements OAuthClientProvider { const incoming = tokens as StoredMcpOAuthTokens; await this.store.write(tokensFile, { ...incoming, - obtained_at: incoming.obtained_at ?? Date.now(), + obtained_at: incoming.obtained_at ?? this.now(), }); }, remove: async () => { @@ -165,7 +168,7 @@ export class McpOAuthClientProvider implements OAuthClientProvider { await this.store.write(`${this.storeKey}${META_SUFFIX}`, meta); const stamped: StoredMcpOAuthTokens = { ...tokens, - obtained_at: (tokens as StoredMcpOAuthTokens).obtained_at ?? Date.now(), + obtained_at: (tokens as StoredMcpOAuthTokens).obtained_at ?? this.now(), }; this.onTokensSaved?.(stamped); } diff --git a/packages/agent-core-v2/src/mcpCore/oauth/service.ts b/packages/agent-core-v2/src/mcpCore/oauth/service.ts index 39dc4fa92..bee2c2d26 100644 --- a/packages/agent-core-v2/src/mcpCore/oauth/service.ts +++ b/packages/agent-core-v2/src/mcpCore/oauth/service.ts @@ -26,6 +26,16 @@ export interface McpOAuthServiceOptions { readonly clientLabel?: string; readonly resolveClientName?: () => string | undefined; readonly log?: Logger; + readonly scheduler?: McpOAuthScheduler; +} + +export interface McpOAuthScheduledTask { + cancel(): void; +} + +export interface McpOAuthScheduler { + now(): number; + schedule(delayMs: number, task: () => void | Promise): McpOAuthScheduledTask; } export interface BeginAuthorizationOptions { @@ -93,16 +103,25 @@ export interface McpOAuthTokenState { const REFRESH_AHEAD_MS = 120_000; const MAX_TIMER_DELAY_MS = 0x7fffffff; +const defaultScheduler: McpOAuthScheduler = { + now: () => Date.now(), + schedule: (delayMs, task) => { + const timer = setTimeout(() => void task(), delayMs); + timer.unref(); + return { cancel: () => clearTimeout(timer) }; + }, +}; + export class McpOAuthService extends Disposable { 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 providers = new Map(); private readonly listeners = new Set(); private readonly refreshes = new Map>(); - private readonly refreshTimers = new Map(); - /** In-flight interactive flows by credential store key; values resolve to the shared flow. */ + private readonly refreshTimers = new Map(); private readonly activeAuthorizations = new Map>(); constructor(options: McpOAuthServiceOptions) { @@ -111,6 +130,7 @@ export class McpOAuthService extends Disposable { this.clientLabel = options.clientLabel; this.resolveClientName = options.resolveClientName; this.log = options.log ?? defaultLog; + this.scheduler = options.scheduler ?? defaultScheduler; this._register({ dispose: () => { void this.shutdown(); @@ -154,7 +174,7 @@ export class McpOAuthService extends Disposable { hasTokens: true, hasRefreshToken: typeof tokens.refresh_token === 'string' && tokens.refresh_token.length > 0, expiresAt, - expired: expiresAt !== undefined && Date.now() >= expiresAt, + expired: expiresAt !== undefined && this.scheduler.now() >= expiresAt, }; } @@ -211,7 +231,7 @@ export class McpOAuthService extends Disposable { /** Clear every pending proactive-refresh timer (engine shutdown, tests). */ stopProactiveRefresh(): void { - for (const timer of this.refreshTimers.values()) clearTimeout(timer); + for (const timer of this.refreshTimers.values()) timer.cancel(); this.refreshTimers.clear(); } @@ -288,12 +308,6 @@ export class McpOAuthService extends Disposable { }; } - /** - * The initiating side of an interactive flow: start the callback listener, - * point the provider at it, and run `auth()` until it surfaces an - * authorization URL. The returned flow owns the single wait-for-callback + - * code exchange shared by every handle for this credential. - */ private async startAuthorizationFlow( serverName: string, serverUrl: string | URL, @@ -440,6 +454,7 @@ export class McpOAuthService extends Disposable { store: this.store, clientLabel: clientLabel ?? this.clientLabel, clientName: this.resolveClientName?.(), + now: () => this.scheduler.now(), onTokensSaved: (tokens) => { this.emit({ type: 'tokens-saved', serverName, serverUrl: canonicalUrl }); if (typeof tokens.obtained_at === 'number' && typeof tokens.expires_in === 'number') { @@ -491,39 +506,35 @@ export class McpOAuthService extends Disposable { const canonicalUrl = canonicalMcpOAuthResource(serverUrl); const storeKey = mcpOAuthStoreKey(serverName, canonicalUrl); this.cancelScheduledRefresh(serverName, canonicalUrl); - const now = Date.now(); + const now = this.scheduler.now(); if (expiresAt <= now) return; const delay = expiresAt - now - REFRESH_AHEAD_MS; - let timer: NodeJS.Timeout; + let timer: McpOAuthScheduledTask; if (delay > MAX_TIMER_DELAY_MS) { - timer = setTimeout(() => { + timer = this.scheduler.schedule(MAX_TIMER_DELAY_MS, () => { this.refreshTimers.delete(storeKey); this.scheduleRefresh(serverName, canonicalUrl, expiresAt); - }, MAX_TIMER_DELAY_MS); + }); } else { - timer = setTimeout( - () => { - this.refreshTimers.delete(storeKey); - void this.refresh(serverName, canonicalUrl).catch((error: unknown) => { - this.emit({ - type: 'refresh-failed', - serverName, - serverUrl: canonicalUrl, - error: error instanceof Error ? error.message : String(error), - }); + timer = this.scheduler.schedule(Math.max(delay, 0), 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), }); - }, - Math.max(delay, 0), - ); + }); + }); } - timer.unref(); this.refreshTimers.set(storeKey, timer); } private cancelScheduledRefresh(serverName: string, serverUrl: string | URL): void { const storeKey = mcpOAuthStoreKey(serverName, serverUrl); const timer = this.refreshTimers.get(storeKey); - if (timer !== undefined) clearTimeout(timer); + timer?.cancel(); this.refreshTimers.delete(storeKey); } diff --git a/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts b/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts index 70fac70c5..a6e889432 100644 --- a/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts +++ b/packages/agent-core-v2/src/workspace/workspaceMcp/workspaceMcpService.ts @@ -153,13 +153,6 @@ export class WorkspaceMcpService extends Disposable implements IWorkspaceMcpServ }; } - /** - * Subscribe a manager to the shared OAuth service's credential events, - * returning the unsubscribe. A completed login reconnects a `needs-auth` / - * `failed` entry; a reset or a failed proactive refresh flips a live - * connection back to `needs-auth` (the reconnect hits a 401) instead of - * leaving it doomed-but-connected. - */ private oauthEventSubscription(manager: McpConnectionManager): () => void { return this.oauthService.onEvent((event) => { void this.handleMcpOAuthEvent(manager, event).catch((error: unknown) => { 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 000000000..6c82da784 --- /dev/null +++ b/packages/agent-core-v2/src/workspace/workspaceTrust/trustRecord.ts @@ -0,0 +1,35 @@ +import { encodeWorkDirKey } from '#/_base/utils/workdir-slug'; +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 { + return (await docs.get(TRUST_SCOPE, encodeWorkDirKey(root))) !== undefined; + } catch { + return false; + } +} + +export function writeWorkspaceTrust( + docs: IAtomicDocumentStore, + root: string, + trustedAt: number, +): Promise { + return docs.set(TRUST_SCOPE, encodeWorkDirKey(root), { root, trustedAt }); +} + +export function deleteWorkspaceTrust( + docs: IAtomicDocumentStore, + root: string, +): Promise { + return docs.delete(TRUST_SCOPE, encodeWorkDirKey(root)); +} diff --git a/packages/agent-core-v2/src/workspace/workspaceTrust/workspaceTrustService.ts b/packages/agent-core-v2/src/workspace/workspaceTrust/workspaceTrustService.ts index 6ad2bd936..ed489d0c7 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/app/mcpConfig/oauthService.test.ts b/packages/agent-core-v2/test/app/mcpConfig/oauthService.test.ts new file mode 100644 index 000000000..e63882aa2 --- /dev/null +++ b/packages/agent-core-v2/test/app/mcpConfig/oauthService.test.ts @@ -0,0 +1,59 @@ +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 { IMcpOAuthService, AppMcpOAuthService } 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); + }); +}); diff --git a/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts b/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts index 347ad6c88..235f4d630 100644 --- a/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts +++ b/packages/agent-core-v2/test/app/mcpManagement/mcpManagement.test.ts @@ -12,6 +12,10 @@ 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'; @@ -31,6 +35,7 @@ import { HostFileSystem } from '#/os/backends/node-local/hostFsService'; import { HostProcessService } from '#/os/backends/node-local/hostProcessService'; 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'; import { FakeRuntime } from '#/runtime/fakeRuntime'; import type { WorkspaceInstance } from '#/workspace/workspaceInstance/workspaceInstance'; @@ -45,7 +50,7 @@ import { startInProcessHttpMcpServer, stdioFixture, } from '../../mcpCore/stubs'; -import { registerAgentIdentityStub } from '../agentIdentity/stubs'; +import { stubAgentIdentity } from '../agentIdentity/stubs'; function stdioServer(name: string, command = 'npx'): GlobalMcpServerConfig { return { name, transport: 'stdio', command }; @@ -66,6 +71,10 @@ describe('McpManagementService', () => { 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 management: IMcpManagementService; @@ -79,6 +88,10 @@ describe('McpManagementService', () => { 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, ); @@ -101,12 +114,22 @@ describe('McpManagementService', () => { }, }); reg.defineInstance(IHostFileSystem, new HostFileSystem()); + 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'], }); - registerAgentIdentityStub(reg); + reg.defineInstance(IAgentIdentity, { + _serviceBrand: undefined, + resolved: () => identityReady, + current: () => identitySnapshot, + }); reg.defineInstance(IRuntimeResolver, { _serviceBrand: undefined, inspect: () => runtime, @@ -135,6 +158,29 @@ describe('McpManagementService', () => { 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') { @@ -493,6 +539,27 @@ describe('McpManagementService', () => { 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', () => { @@ -566,6 +633,59 @@ describe('McpManagementService', () => { }); }); + 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(getOrCreate).not.toHaveBeenCalled(); + + releaseConfig(); + await Promise.resolve(); + expect(getOrCreate).not.toHaveBeenCalled(); + + releaseIdentity(); + await expect(probe).resolves.toMatchObject({ success: true }); + expect(getOrCreate).toHaveBeenCalledWith({ root: cwd }); + }, 20000); + it('rejects a name-only probe under an enabled runtime-name collision', async () => { pluginEntries = [ { @@ -693,8 +813,8 @@ describe('McpManagementService', () => { ]); }); - it('probes unpinned servers without a stored grant and classifies oauth-marked ones offline', async () => { - const server = await startHttpServer(); + it('classifies unpinned servers without a stored grant offline', async () => { + const server = await startCountingServer(); await management.addServer({ name: 'plain', transport: 'http', url: server.url }); await management.addServer({ name: 'challenged', @@ -707,6 +827,7 @@ describe('McpManagementService', () => { { name: 'plain', authStatus: 'not-applicable' }, { name: 'challenged', authStatus: 'oauth-required' }, ]); + expect(server.requestCount()).toBe(0); }, 20000); it('verify settles a stored-but-rejected grant as oauth-expired through a real probe', async () => { diff --git a/packages/agent-core-v2/test/app/mcpRegistry/mcpRegistry.test.ts b/packages/agent-core-v2/test/app/mcpRegistry/mcpRegistry.test.ts index 312e3764a..3f9c0cec7 100644 --- a/packages/agent-core-v2/test/app/mcpRegistry/mcpRegistry.test.ts +++ b/packages/agent-core-v2/test/app/mcpRegistry/mcpRegistry.test.ts @@ -21,6 +21,7 @@ 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 { @@ -47,6 +48,7 @@ describe('McpRegistryService', () => { let store: IMcpConfigStore; let pluginEntries: PluginMcpServerEntry[]; let pluginError: Error | undefined; + let trusted: boolean; let registry: IMcpRegistryService; beforeEach(() => { @@ -56,6 +58,7 @@ describe('McpRegistryService', () => { tempDirs = [home]; pluginEntries = []; pluginError = undefined; + trusted = true; const ix = createServices(disposables, { additionalServices: (reg) => { reg.defineInstance(IFileSystemStorageService, new InMemoryStorageService()); @@ -68,6 +71,9 @@ describe('McpRegistryService', () => { }, }); reg.defineInstance(IHostFileSystem, new HostFileSystem()); + reg.definePartialInstance(IAtomicDocumentStore, { + get: async () => (trusted ? ({} as T) : undefined), + }); reg.define(IMcpRegistryService, McpRegistryService); }, }); @@ -170,6 +176,28 @@ describe('McpRegistryService', () => { }); }); + 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('exposes plugin servers as read-only entries with their effective config', async () => { pluginEntries = [ pluginEntry('demo', 'finance', { diff --git a/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts b/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts index e83bcf4e1..6eedbcf8c 100644 --- a/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts +++ b/packages/agent-core-v2/test/mcpCore/oauth/service.test.ts @@ -17,7 +17,7 @@ import { } from '#/mcpCore/oauth/service'; import { mcpOAuthStoreKey, type McpOAuthStore } from '#/mcpCore/oauth/store'; -import { createMemoryMcpOAuthStore } from '../stubs'; +import { createMemoryMcpOAuthStore, ManualMcpOAuthScheduler } from '../stubs'; const SERVER_NAME = 'notion'; const SERVER_URL = 'https://mcp.example.test/mcp'; @@ -26,13 +26,15 @@ interface Fixture { readonly service: McpOAuthService; readonly store: McpOAuthStore; readonly events: McpOAuthEvent[]; + readonly scheduler: ManualMcpOAuthScheduler; } function makeFixture(store: McpOAuthStore = createMemoryMcpOAuthStore()): Fixture { const events: McpOAuthEvent[] = []; - const service = new McpOAuthService({ store }); + const scheduler = new ManualMcpOAuthScheduler(1_000_000); + const service = new McpOAuthService({ store, scheduler }); service.onEvent((event) => events.push(event)); - return { service, store, events }; + return { service, store, events, scheduler }; } const cleanups: Array<() => Promise | void> = []; @@ -147,20 +149,11 @@ async function deliverCallback(flow: BeginAuthorizationResult): Promise { await response.text(); } -async function waitFor(condition: () => boolean, description: string): Promise { - const deadline = Date.now() + 5000; - while (!condition()) { - if (Date.now() > deadline) throw new Error(`timed out waiting for ${description}`); - await new Promise((resolve) => setTimeout(resolve, 10)); - } -} - 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()); - const before = Date.now(); await fixture.service .getProvider(SERVER_NAME, SERVER_URL) .saveTokens({ access_token: 'a', token_type: 'Bearer', expires_in: 3600 }); @@ -168,9 +161,7 @@ describe('McpOAuthService credential bookkeeping', () => { const state = await fixture.service.tokenState(SERVER_NAME, SERVER_URL); expect(state.hasTokens).toBe(true); expect(state.expired).toBe(false); - expect(state.expiresAt).toBeDefined(); - expect(state.expiresAt!).toBeGreaterThanOrEqual(before + 3600_000); - expect(state.expiresAt!).toBeLessThanOrEqual(Date.now() + 3600_000); + expect(state.expiresAt).toBe(4_600_000); const metaFiles = await listMetaKeys(fixture.store); expect(metaFiles).toHaveLength(1); @@ -610,7 +601,7 @@ describe('McpOAuthService sweepProactiveRefresh resilience', () => { refresh_token: 'stale-refresh-token', token_type: 'Bearer', expires_in: 60, - obtained_at: Date.now(), + obtained_at: fixture.scheduler.now(), }); await fixture.store.write(`${storeKey}-meta.json`, { serverName: SERVER_NAME, @@ -623,10 +614,8 @@ describe('McpOAuthService sweepProactiveRefresh resilience', () => { await fixture.store.write('corrupt-meta.json', '{not json'); await expect(fixture.service.sweepProactiveRefresh()).resolves.toBeUndefined(); - await waitFor( - () => authServer.counts.refresh === 1, - 'the swept credential to refresh immediately', - ); + await fixture.scheduler.advanceBy(0); + expect(authServer.counts.refresh).toBe(1); }, 15000); }); @@ -647,17 +636,14 @@ describe('McpOAuthService proactive refresh scheduling', () => { expires_in: 60, }); - await waitFor(() => authServer.counts.refresh === 1, 'an immediate proactive refresh'); + await fixture.scheduler.advanceBy(0); + expect(authServer.counts.refresh).toBe(1); expect(fixture.events.filter((event) => event.type === 'tokens-saved')).toHaveLength(2); }, 15000); it('re-arms scheduling for expiries beyond the setTimeout limit', async () => { const fixture = makeFixture(); - cleanups.push(() => { - vi.useRealTimers(); - }); cleanups.push(() => fixture.service.dispose()); - vi.useFakeTimers(); const maxTimerDelayMs = 0x7fffffff; const refreshSpy = vi .spyOn(fixture.service, 'refresh') @@ -671,10 +657,10 @@ describe('McpOAuthService proactive refresh scheduling', () => { }); const expiresAt = (await fixture.service.tokenState(SERVER_NAME, SERVER_URL)).expiresAt!; - await vi.advanceTimersByTimeAsync(maxTimerDelayMs); + await fixture.scheduler.advanceBy(maxTimerDelayMs); expect(refreshSpy).not.toHaveBeenCalled(); - await vi.advanceTimersByTimeAsync(expiresAt - Date.now() - 120_000); + await fixture.scheduler.advanceBy(expiresAt - fixture.scheduler.now() - 120_000); expect(refreshSpy).toHaveBeenCalledTimes(1); expect(fixture.events).toContainEqual({ type: 'refresh-failed', @@ -686,11 +672,7 @@ describe('McpOAuthService proactive refresh scheduling', () => { it('does not proactively refresh an already-expired grant', async () => { const fixture = makeFixture(); - cleanups.push(() => { - vi.useRealTimers(); - }); cleanups.push(() => fixture.service.dispose()); - vi.useFakeTimers(); const refreshSpy = vi.spyOn(fixture.service, 'refresh'); await fixture.service.getProvider(SERVER_NAME, SERVER_URL).saveTokens({ @@ -700,7 +682,7 @@ describe('McpOAuthService proactive refresh scheduling', () => { expires_in: -60, }); - await vi.advanceTimersByTimeAsync(10_000); + await fixture.scheduler.advanceBy(10_000); expect(refreshSpy).not.toHaveBeenCalled(); }); }); @@ -746,11 +728,7 @@ describe('McpOAuthService shutdown', () => { it('clears pending proactive-refresh timers', async () => { const fixture = makeFixture(); - cleanups.push(() => { - vi.useRealTimers(); - }); cleanups.push(() => fixture.service.dispose()); - vi.useFakeTimers(); await fixture.service.getProvider(SERVER_NAME, SERVER_URL).saveTokens({ access_token: 'a', @@ -761,7 +739,7 @@ describe('McpOAuthService shutdown', () => { const refreshSpy = vi.spyOn(fixture.service, 'refresh'); await fixture.service.shutdown(); - await vi.advanceTimersByTimeAsync(3600_000); + await fixture.scheduler.advanceBy(3600_000); expect(refreshSpy).not.toHaveBeenCalled(); }); }); diff --git a/packages/agent-core-v2/test/mcpCore/stubs.ts b/packages/agent-core-v2/test/mcpCore/stubs.ts index d6013af99..957e9d9d8 100644 --- a/packages/agent-core-v2/test/mcpCore/stubs.ts +++ b/packages/agent-core-v2/test/mcpCore/stubs.ts @@ -9,6 +9,10 @@ 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, @@ -58,6 +62,41 @@ export function createMemoryMcpOAuthStore(): McpOAuthStore { }; } +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[] = [ { 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 b510fbcf0..79d4b3345 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceMcp/workspaceMcp.test.ts @@ -46,7 +46,12 @@ import { import { stubLog } from '../../_base/log/stubs'; import { registerAgentIdentityStub } from '../../app/agentIdentity/stubs'; -import { createMemoryMcpOAuthStore, startInProcessHttpMcpServer, stdioFixture } from '../../mcpCore/stubs'; +import { + createMemoryMcpOAuthStore, + ManualMcpOAuthScheduler, + startInProcessHttpMcpServer, + stdioFixture, +} from '../../mcpCore/stubs'; function stdioServer(): McpServerConfig { return { @@ -66,6 +71,7 @@ describe('WorkspaceMcpService', () => { let configChanges: Emitter; let assemblyEvents: Emitter; let oauthService: McpOAuthService; + let oauthScheduler: ManualMcpOAuthScheduler; let manager: InstanceType | undefined; beforeEach(() => { @@ -76,7 +82,11 @@ describe('WorkspaceMcpService', () => { tunablesFn = vi.fn(() => tunablesValue); configChanges = new Emitter(); assemblyEvents = disposables.add(new Emitter()); - oauthService = new McpOAuthService({ store: createMemoryMcpOAuthStore() }); + oauthScheduler = new ManualMcpOAuthScheduler(); + oauthService = new McpOAuthService({ + store: createMemoryMcpOAuthStore(), + scheduler: oauthScheduler, + }); manager = undefined; }); @@ -176,9 +186,14 @@ describe('WorkspaceMcpService', () => { 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; }), ); @@ -193,7 +208,7 @@ describe('WorkspaceMcpService', () => { manager = service.connectionManager(); configChanges.fire({ upsert: { beta: stdioServer() }, remove: ['alpha'] }); - await new Promise((resolvePromise) => setTimeout(resolvePromise, 300)); + await connectAllStarted; expect(connect).not.toHaveBeenCalled(); expect(markRemoved).not.toHaveBeenCalled(); @@ -544,7 +559,6 @@ describe('WorkspaceMcpService', () => { await oauthService .getProvider('notion', SERVER_URL) .saveTokens({ access_token: 'a', token_type: 'Bearer' }); - await new Promise((resolve) => setTimeout(resolve, 20)); expect(reconnectAndJoin).not.toHaveBeenCalled(); }); @@ -574,7 +588,6 @@ describe('WorkspaceMcpService', () => { await oauthService .getProvider('notion', SERVER_URL) .saveTokens({ access_token: 'a', token_type: 'Bearer' }); - await new Promise((resolve) => setTimeout(resolve, 20)); expect(reconnectAndJoin).not.toHaveBeenCalled(); expect(forgetProvider).not.toHaveBeenCalled(); @@ -599,7 +612,6 @@ describe('WorkspaceMcpService', () => { await oauthService .getProvider('notion', SERVER_URL) .saveTokens({ access_token: 'a', token_type: 'Bearer' }); - await new Promise((resolve) => setTimeout(resolve, 20)); expect(reconnectAfterCurrent).not.toHaveBeenCalled(); notifyStatus?.({ name: 'notion', transport: 'http', status: 'connected', toolCount: 0 }); @@ -618,7 +630,6 @@ describe('WorkspaceMcpService', () => { const provider = oauthService.getProvider('notion', SERVER_URL); await provider.ready; await provider.clearCredentials('client'); - await new Promise((resolve) => setTimeout(resolve, 20)); expect(reconnectAndJoin).not.toHaveBeenCalled(); expect(forgetProvider).not.toHaveBeenCalled(); @@ -634,7 +645,6 @@ describe('WorkspaceMcpService', () => { const provider = oauthService.getProvider('notion', SERVER_URL); await provider.ready; await provider.clearCredentials('discovery'); - await new Promise((resolve) => setTimeout(resolve, 20)); expect(reconnectAndJoin).not.toHaveBeenCalled(); expect(forgetProvider).not.toHaveBeenCalled(); @@ -655,9 +665,8 @@ describe('WorkspaceMcpService', () => { }); const reconnectAndJoin = mockManagerEntry('connected'); - await vi.waitFor(() => { - expect(reconnectAndJoin).toHaveBeenCalledWith('notion'); - }); + await oauthScheduler.advanceBy(0); + expect(reconnectAndJoin).toHaveBeenCalledWith('notion'); expect(authServer.counts.refresh).toBe(1); }); @@ -678,10 +687,8 @@ describe('WorkspaceMcpService', () => { }); const reconnectAndJoin = mockManagerEntry('needs-auth'); - await vi.waitFor(() => { - expect(events.some((event) => event.type === 'refresh-failed')).toBe(true); - }); - await new Promise((resolve) => setTimeout(resolve, 20)); + await oauthScheduler.advanceBy(0); + expect(events.some((event) => event.type === 'refresh-failed')).toBe(true); expect(reconnectAndJoin).not.toHaveBeenCalled(); }); @@ -694,7 +701,6 @@ describe('WorkspaceMcpService', () => { await oauthService .getProvider('notion', SERVER_URL) .saveTokens({ access_token: 'a', token_type: 'Bearer' }); - await new Promise((resolve) => setTimeout(resolve, 20)); expect(reconnectAndJoin).not.toHaveBeenCalled(); }); @@ -708,7 +714,6 @@ describe('WorkspaceMcpService', () => { await oauthService .getProvider('notion', SERVER_URL) .saveTokens({ access_token: 'a', token_type: 'Bearer' }); - await new Promise((resolve) => setTimeout(resolve, 20)); expect(reconnectAndJoin).not.toHaveBeenCalled(); }); 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 a44265842..c70a634a5 100644 --- a/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceMcpConfig/workspaceMcpConfig.test.ts @@ -224,18 +224,32 @@ describe('WorkspaceMcpConfigService', () => { }, 20000); it('keeps the winning plugin server when the same-named file entry vanishes', async () => { - const file = await writeProjectConfig({ shared: stdioConfig('file-version') }); + await writeProjectConfig({ shared: stdioConfig('file-version') }); pluginServers = { shared: stdioConfig('plugin-version') }; const service = createService(); await service.ready; expect(service.servers()).toEqual({ shared: stdioConfig('plugin-version') }); await writeProjectConfig({}); - watchFires.get(cwd)?.fire({ path: file, action: 'modified', kind: 'file' }); + storeWrites.fire(); + pluginServers = { + shared: stdioConfig('plugin-version'), + pluginOnly: stdioConfig('plugin'), + }; + pluginReloads.fire({ added: [], removed: [], errors: [] }); - await new Promise((resolvePromise) => setTimeout(resolvePromise, 500)); - expect(changes).toEqual([]); - expect(service.servers()).toEqual({ shared: stdioConfig('plugin-version') }); + await vi.waitFor( + () => { + expect(changes).toEqual([ + { 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 () => {