fix(agent-core-v2): harden MCP management readiness

This commit is contained in:
7Sageer 2026-08-19 21:12:01 +08:00
parent 4715a8bda5
commit 4805c7bebe
17 changed files with 439 additions and 172 deletions

View file

@ -0,0 +1,5 @@
---
"@moonshot-ai/kimi-code": patch
---
Respect workspace trust and configuration readiness when managing MCP servers.

View file

@ -45,7 +45,6 @@ export class McpConfigStore extends Disposable implements IMcpConfigStore {
private readonly writeEmitter = this._register(new Emitter<void>());
readonly onDidWrite: Event<void> = this.writeEmitter.event;
/** Serializes the read-modify-write mutations so concurrent writes cannot lose updates. */
private mutationTail: Promise<void> = Promise.resolve();
constructor(

View file

@ -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)}`);
});
}

View file

@ -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<string, { flow: BeginAuthorizationResult }>();
constructor(
@ -109,18 +108,13 @@ export class McpManagementService extends Disposable implements IMcpManagementSe
}
async testServer(target: McpServerTestTarget): Promise<McpServerTestResult> {
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<McpRegistryEntry | undefined> {
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<GlobalMcpServerConfig> {
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<T> {
await this.waitForReadiness();
const section = this.config.get<McpSection | undefined>(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<readonly McpServerAuthStatus[]> {
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<readonly McpServerInspection[]> {
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<McpServerAuthBeginResult> {
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<void> {
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<readonly McpServerRuntimeDescriptor[]> {
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<void> {
await this.config.ready;
await this.identity.resolved();
}
}
function throwReadOnlyMcpServer(entry: McpRegistryEntry): void {

View file

@ -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<readonly McpRegistryEntry[]> {
@ -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,
});
}
}
}

View file

@ -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);
}

View file

@ -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<void>): 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<string, McpOAuthClientProvider>();
private readonly listeners = new Set<McpOAuthEventListener>();
private readonly refreshes = new Map<string, Promise<void>>();
private readonly refreshTimers = new Map<string, NodeJS.Timeout>();
/** In-flight interactive flows by credential store key; values resolve to the shared flow. */
private readonly refreshTimers = new Map<string, McpOAuthScheduledTask>();
private readonly activeAuthorizations = new Map<string, Promise<SharedAuthorizationFlow>>();
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);
}

View file

@ -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) => {

View file

@ -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<boolean> {
try {
return (await docs.get<TrustRecord>(TRUST_SCOPE, encodeWorkDirKey(root))) !== undefined;
} catch {
return false;
}
}
export function writeWorkspaceTrust(
docs: IAtomicDocumentStore,
root: string,
trustedAt: number,
): Promise<void> {
return docs.set(TRUST_SCOPE, encodeWorkDirKey(root), { root, trustedAt });
}
export function deleteWorkspaceTrust(
docs: IAtomicDocumentStore,
root: string,
): Promise<void> {
return docs.delete(TRUST_SCOPE, encodeWorkDirKey(root));
}

View file

@ -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<boolean>(
'workspaceTrust.trusted',
@ -25,7 +18,6 @@ export class WorkspaceTrustService extends Disposable implements IWorkspaceTrust
readonly ready: Promise<void>;
private readonly root: string;
private readonly storeKey: string;
private readonly changeEmitter = this._register(new Emitter<WorkspaceTrustChange>());
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<void> {
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<void> {
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<void> {
try {
this.trusted = (await this.docs.get<TrustRecord>(TRUST_SCOPE, this.storeKey)) !== undefined;
} catch {
this.trusted = false;
}
this.trusted = await readWorkspaceTrust(this.docs, this.root);
}
}

View file

@ -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<void>((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);
});
});

View file

@ -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<void>;
let identityReady: Promise<AgentIdentitySnapshot>;
let identitySnapshot: AgentIdentitySnapshot;
let trusted: boolean;
let getOrCreate: Mock<IWorkspaceInstanceManager['getOrCreate']>;
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<IWorkspaceInstanceManager['getOrCreate']>(async () =>
({ id: 'test-workspace' }) as unknown as WorkspaceInstance,
);
@ -101,12 +114,22 @@ describe('McpManagementService', () => {
},
});
reg.defineInstance(IHostFileSystem, new HostFileSystem());
reg.definePartialInstance(IAtomicDocumentStore, {
get: async <T>() => (trusted ? ({} as T) : undefined),
});
reg.define(IMcpRegistryService, McpRegistryService);
reg.defineInstance(IMcpOAuthService, oauth);
reg.definePartialInstance(IConfigService, {
get ready() {
return configReady;
},
get: (<T = unknown,>(_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<void>((resolve) => httpServer.listen(0, '127.0.0.1', resolve));
httpServers.push({
close: () =>
new Promise<void>((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<void>((resolve) => {
releaseConfig = resolve;
});
let releaseIdentity: () => void = () => undefined;
identityReady = new Promise<AgentIdentitySnapshot>((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 () => {

View file

@ -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 <T>() => (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', {

View file

@ -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> | void> = [];
@ -147,20 +149,11 @@ async function deliverCallback(flow: BeginAuthorizationResult): Promise<void> {
await response.text();
}
async function waitFor(condition: () => boolean, description: string): Promise<void> {
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();
});
});

View file

@ -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<void> }
>();
constructor(now = Date.now()) {
this.current = now;
}
now(): number {
return this.current;
}
schedule(delayMs: number, task: () => void | Promise<void>): 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<void> {
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[] = [
{

View file

@ -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<McpServersChange>;
let assemblyEvents: Emitter<SessionWillCreateEvent>;
let oauthService: McpOAuthService;
let oauthScheduler: ManualMcpOAuthScheduler;
let manager: InstanceType<typeof McpConnectionManager> | undefined;
beforeEach(() => {
@ -76,7 +82,11 @@ describe('WorkspaceMcpService', () => {
tunablesFn = vi.fn(() => tunablesValue);
configChanges = new Emitter<McpServersChange>();
assemblyEvents = disposables.add(new Emitter<SessionWillCreateEvent>());
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<void>((resolve) => {
signalConnectAllStarted = resolve;
});
vi.spyOn(McpConnectionManager.prototype, 'connectAll').mockImplementation(
() =>
new Promise<void>((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();
});

View file

@ -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 () => {