mirror of
https://github.com/MoonshotAI/kimi-code.git
synced 2026-08-25 00:27:32 +00:00
fix(agent-core-v2): harden MCP management readiness
This commit is contained in:
parent
4715a8bda5
commit
4805c7bebe
17 changed files with 439 additions and 172 deletions
5
.changeset/respect-mcp-management-readiness.md
Normal file
5
.changeset/respect-mcp-management-readiness.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
---
|
||||
"@moonshot-ai/kimi-code": patch
|
||||
---
|
||||
|
||||
Respect workspace trust and configuration readiness when managing MCP servers.
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)}`);
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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) => {
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
}
|
||||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
@ -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 () => {
|
||||
|
|
|
|||
|
|
@ -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', {
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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[] = [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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 () => {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue