Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/mcp-first-turn-ready.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@moonshot-ai/kimi-code": patch
---

Ensure the first request waits for MCP startup to finish while the interface still opens immediately.
2 changes: 1 addition & 1 deletion packages/acp-server/test/e2e-turn.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -720,7 +720,7 @@ describe('acp-server builtin slash commands (local execution, no LLM turn)', ()
const { chunk, stopReason } = await runSlash(c, created.sessionId, '/mcp');
expect(stopReason).toBe('end_turn');
expect(chunk).toContain('MCP servers (1):');
expect(chunk).toContain('- mock (stdio): connected,');
expect(chunk).toContain('- mock (stdio):');
expect(scripted!.callCount()).toBe(0);
}, 30_000);

Expand Down
14 changes: 12 additions & 2 deletions packages/agent-core-v2/src/agent/mcp/mcpService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,10 @@
* keeps them registered across reconnects, swaps in the OAuth tool for
* `needs-auth` servers, journals tool discoveries on the wire (queued until
* restore finishes), and publishes `mcp.server.status` / `tool.list.updated`
* events. The plain-data state (`mcpToolsByServer`, `discoveryWritesReady`)
* events. Sessions and agents construct without awaiting the manager's
* initial connect; each LLM step instead waits for it through a `loop`
* onWillBeginStep hook (a no-op once settled), with the per-execution
* `toolExecutor` onWillExecuteTool wait as the backstop. The plain-data state (`mcpToolsByServer`, `discoveryWritesReady`)
* is registered into `agentState` (`IAgentStateService`) and read/written
* through it; `mcpTools` stays a plain instance field (its values hold
* disposable resource handles, not plain data), as does `pendingDiscoveries`
Expand All @@ -31,6 +34,7 @@ import { ITelemetryService } from '#/app/telemetry/telemetry';
import { sessionMediaOriginalsDir } from '#/agent/media/image-originals';
import { IAgentToolExecutorService } from '#/agent/toolExecutor/toolExecutor';
import { IAgentToolRegistryService } from '#/agent/toolRegistry/toolRegistry';
import { IAgentLoopService } from '#/agent/loop/loop';
import { createMcpAuthTool } from '#/agent/mcp/tools/auth';
import { createMcpTool } from '#/agent/mcp/tools/mcp';
import { ISessionContext } from '#/session/sessionContext/sessionContext';
Expand Down Expand Up @@ -104,6 +108,7 @@ export class AgentMcpService extends Disposable implements IAgentMcpService {
@IAgentToolRegistryService private readonly registry: IAgentToolRegistryService,
@IEventBus private readonly eventBus: IEventBus,
@IAgentToolExecutorService toolExecutor: IAgentToolExecutorService,
@IAgentLoopService loop: IAgentLoopService,
@IWireService private readonly wire: IWireService,
@ITelemetryService private readonly telemetry: ITelemetryService,
@IAgentStateService private readonly states: IAgentStateService,
Expand All @@ -112,6 +117,10 @@ export class AgentMcpService extends Disposable implements IAgentMcpService {
this.states.register(mcpMcpToolsByServerKey);
this.states.register(mcpDiscoveryWritesReadyKey);
this.attachMcpTools();
loop.hooks.onWillBeginStep.register('mcp', async (ctx, next) => {
await this.waitForInitialLoad(ctx.signal);
await next();
});
this._register(
toolExecutor.onWillExecuteTool((event) => {
event.waitUntil(this.waitForInitialLoad(event.signal));
Expand Down Expand Up @@ -142,7 +151,8 @@ export class AgentMcpService extends Disposable implements IAgentMcpService {
}

waitForInitialLoad(signal?: AbortSignal): Promise<void> {
return this.mcpHandle.connectionManager.waitForInitialLoad(signal);
const ready = this.mcpHandle.ready;
return signal === undefined ? ready : abortable(ready, signal);
}

initialLoadDurationMs(): number {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,9 @@
*
* No agent id is special here: the main agent is simply the agent created
* with the conventional `MAIN_AGENT_ID`, and `fork` requires its source to
* exist. The workspace's shared MCP
* manager arrives through the seeded `ISessionMcpHandle`, whose initial
* connect this service awaits during creation.
* exist. MCP readiness is not awaited here: the workspace's shared manager
* connects in the background and the agent's LLM steps wait on it instead
* (see `AgentMcpService`).
*/

import { IInstantiationService } from '#/_base/di/instantiation';
Expand All @@ -40,7 +40,6 @@ import type { PermissionMode } from '#/agent/permissionPolicy/types';
import { IAgentTaskService } from '#/agent/task/task';
import { ISessionContext } from '#/session/sessionContext/sessionContext';
import { ISessionMetadata } from '#/session/sessionMetadata/sessionMetadata';
import { ISessionMcpHandle } from '#/session/mcp/sessionMcpHandle';
import { IAgentScopeContext, makeAgentScopeContext } from '#/agent/scopeContext/scopeContext';
import { IAgentLoopService } from '#/agent/loop/loop';
import { IAgentProfileService } from '#/agent/profile/profile';
Expand Down Expand Up @@ -82,7 +81,6 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle
@ISessionMetadata private readonly sessionMetadata: ISessionMetadata,
@IBootstrapService private readonly bootstrap: IBootstrapService,
@IConfigService private readonly config: IConfigService,
@ISessionMcpHandle private readonly mcpHandle: ISessionMcpHandle,
@ISessionInteractionService private readonly interaction: ISessionInteractionService,
@ITelemetryService private readonly telemetry: ITelemetryService,
) {
Expand Down Expand Up @@ -145,7 +143,6 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle
}

private async doCreate(agentId: string, opts: CreateAgentOptions): Promise<IAgentScopeHandle> {
const mcpReady = this.mcpHandle.ready;
const agentScope = this.ctx.scope(`agents/${agentId}`);
const agentHomedir = join(this.bootstrap.homeDir, agentScope);
const handle = createScopedChildHandle(
Expand All @@ -171,7 +168,6 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle
labels: opts.labels,
});
this.onDidCreateEmitter.fire(handle);
await mcpReady;
await wire.restore();
await this.bindBootstrap(handle, opts);
await handle.accessor.get(IAgentToolActivationService).activate();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,12 +45,15 @@
* with a fire-and-forget `reload()` so a fixed agent file unblocks later
* creates
* (the workspace skill catalog, by contrast, is kicked fire-and-forget).
* The handler's shared MCP manager is awaited before create/resume returns;
* a session created with ephemeral `mcpServers` additionally gets a session
* overlay from `workspaceMcp` (session-owned connections, seeded as a merged
* view, shut down when the session handle disposes — with a backstop in the
* service's own dispose for teardown paths that bypass the handle wrapper),
* whose initial connect is awaited here too.
* The handler's shared MCP manager is NOT awaited before create/resume
* returns — it connects fire-and-forget at Workspace scope, and the seeded
* handle's `ready` promise lets the agent's LLM steps wait on it instead
* (see `AgentMcpService`). A session created with ephemeral `mcpServers`
* additionally gets a session overlay from `workspaceMcp` (session-owned
* connections, seeded as a merged view, shut down when the session handle
* disposes — with a backstop in the service's own dispose for teardown
* paths that bypass the handle wrapper), likewise connected in the
* background.
* The session-level services whose subscriptions
* must exist before the first agent / turn (external hooks, cron, the
* secondary-model startup warning) opt into `OnScopeCreated` activation.
Expand Down Expand Up @@ -332,8 +335,6 @@ export class SessionLifecycleService extends Disposable implements ISessionLifec
this.userAgentProfileLoader.ready,
this.pluginAgentProfileLoader.ready,
]);
await this.mcp.ready;
await mcpOverlay?.handle.ready;
} catch (error) {
handle.dispose();
void this.explicitAgentProfileLoader.reload().catch(() => undefined);
Expand Down
37 changes: 31 additions & 6 deletions packages/agent-core-v2/test/agent/mcp/mcp.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -219,16 +219,18 @@ describe('AgentMcpService', () => {
disposables.dispose();
});

function createService(manager: FakeMcpManager): AgentMcpService {
function createService(
manager: FakeMcpManager,
ready: Promise<void> = Promise.resolve(),
): IAgentMcpService {
ix.stub(ISessionMcpHandle, {
_serviceBrand: undefined,
ready: Promise.resolve(),
ready,
connectionManager: manager as unknown as McpConnectionManager,
} satisfies ISessionMcpHandle);
ix.stub(ISessionContext, { sessionDir: '/tmp/kimi-code-mcp-test' });
const svc = ix.createInstance(AgentMcpService);
disposables.add(svc);
return svc;
ix.set(IAgentMcpService, new SyncDescriptor(AgentMcpService));
return ix.get(IAgentMcpService);
}

it('delegates list / status events to the connection manager', async () => {
Expand All @@ -249,9 +251,32 @@ describe('AgentMcpService', () => {
expect(statuses).toEqual(['s1:connected', 's2:connected', 's1:disabled']);
});

it('holds the LLM step until the session MCP handle is ready', async () => {
const manager = new FakeMcpManager();
let releaseReady!: () => void;
const ready = new Promise<void>((resolve) => {
releaseReady = resolve;
});
createService(manager, ready);

const loop = ix.get(IAgentLoopService);
let settled = false;
const step = loop.hooks.onWillBeginStep
.run({ turnId: 1, step: 1, signal: new AbortController().signal })
.then(() => {
settled = true;
});

await Promise.resolve();
expect(settled).toBe(false);

releaseReady();
await step;
expect(settled).toBe(true);
});

it('resolves through the IAgentMcpService binding with no manager', () => {
const created = createService(new FakeMcpManager());
ix.set(IAgentMcpService, created);
const svc = ix.get(IAgentMcpService);
expect(svc).toBe(created);
expect(svc.list()).toEqual([]);
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/**
* Scenario: session-owned agent creation, persistence, and MCP readiness.
* Scenario: session-owned agent creation, persistence, and MCP wiring.
*
* Exercises `AgentLifecycleService` through its DI contract with controlled
* persistence and MCP boundaries, including completion ordering.
Expand Down Expand Up @@ -648,7 +648,7 @@ describe('AgentLifecycleService', () => {
]);
});

it('waits for the MCP handle readiness before returning an agent', async () => {
it('returns an agent without waiting for the MCP handle readiness', async () => {
let releaseReady!: () => void;
const ready = new Promise<void>((resolve) => {
releaseReady = resolve;
Expand All @@ -660,21 +660,12 @@ describe('AgentLifecycleService', () => {
} satisfies ISessionMcpHandle);

const svc = ix.get(IAgentLifecycleService);
let settled = false;
const create = svc.create({ agentId: 'main' }).then(() => {
settled = true;
});

// The wire seal + registerAgent complete first; the create call then
// parks on the seeded MCP readiness promise.
await vi.waitFor(() => {
expect(registerAgent).toHaveBeenCalled();
});
expect(settled).toBe(false);
// MCP connects in the background; the agent's LLM steps wait on the
// seeded readiness promise instead of agent creation.
const handle = await svc.create({ agentId: 'main' });
expect(handle.id).toBe('main');

releaseReady();
await create;
expect(settled).toBe(true);
});

it('exposes the in-flight handle and joins it after bootstrap', async () => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1150,7 +1150,7 @@ describe('SessionLifecycleService', () => {
expect(recordedSessionHookEvents).toEqual(['create:startup:s1', 'close:exit:s1']);
});

it('waits for MCP initialization before create returns', async () => {
it('returns from create without waiting for MCP initialization', async () => {
let resolveMcpReady: (() => void) | undefined;
const mcpReady = new Promise<void>((resolve) => {
resolveMcpReady = resolve;
Expand All @@ -1159,17 +1159,13 @@ describe('SessionLifecycleService', () => {
stubPair(IWorkspaceMcpService, workspaceMcpServiceStub(mcpReady)),
]);

let settled = false;
const create = svc.create({ sessionId: 's1', workDir: '/tmp/proj' }).then(() => {
settled = true;
});

await tick();
expect(settled).toBe(false);
// Create resolves while the workspace MCP initial connect is still
// pending; the seeded handle carries the readiness promise so the agent's
// LLM steps can wait on it instead.
const handle = await svc.create({ sessionId: 's1', workDir: '/tmp/proj' });
expect(handle.accessor.get(ISessionMcpHandle).ready).toBe(mcpReady);

resolveMcpReady?.();
await create;
expect(settled).toBe(true);
});

function overlayStub(ready: Promise<void> = Promise.resolve()) {
Expand Down Expand Up @@ -1219,7 +1215,7 @@ describe('SessionLifecycleService', () => {
expect(handle?.accessor.get(ISessionMcpHandle)).toBe(overlayHandle);
});

it('waits for the session MCP overlay readiness before create returns', async () => {
it('returns from create without waiting for the session MCP overlay readiness', async () => {
let resolveOverlayReady: (() => void) | undefined;
const overlayReady = new Promise<void>((resolve) => {
resolveOverlayReady = resolve;
Expand All @@ -1229,23 +1225,17 @@ describe('SessionLifecycleService', () => {
stubPair(IWorkspaceMcpService, { ...workspaceMcpServiceStub(), sessionOverlay }),
]);

let settled = false;
const create = svc
.create({
sessionId: 's1',
workDir: '/tmp/proj',
mcpServers: { eph: { transport: 'stdio', command: 'node' } },
})
.then(() => {
settled = true;
});

await tick();
expect(settled).toBe(false);
// Create resolves while the overlay's initial connect is still pending;
// the seeded handle carries the readiness promise so the agent's LLM
// steps can wait on it instead.
const handle = await svc.create({
sessionId: 's1',
workDir: '/tmp/proj',
mcpServers: { eph: { transport: 'stdio', command: 'node' } },
});
expect(handle.accessor.get(ISessionMcpHandle).ready).toBe(overlayReady);

resolveOverlayReady?.();
await create;
expect(settled).toBe(true);
});

it('shuts the session MCP overlay down when create fails after materialization', async () => {
Expand Down Expand Up @@ -1323,14 +1313,26 @@ describe('SessionLifecycleService', () => {
});

it('hides a session from get/list until its resume finishes', async () => {
let resolveMcpReady: (() => void) | undefined;
const mcpReady = new Promise<void>((resolve) => {
resolveMcpReady = resolve;
let releaseMainAgent: ((handle: IAgentScopeHandle) => void) | undefined;
const mainAgent = new Promise<IAgentScopeHandle>((resolve) => {
releaseMainAgent = resolve;
});
const main = {
id: MAIN_AGENT_ID,
kind: LifecycleScope.Agent,
accessor: {
get: () => {
throw new Error('unexpected main agent service access');
},
},
dispose: () => {},
} as IAgentScopeHandle;
const svc = await build([
stubPair(ISessionIndex, sessionIndexWithSummary('s1', '/tmp/proj', 'wd_stub')),
stubPair(IAgentLifecycleService, agentLifecycleWithMainStub()),
stubPair(IWorkspaceMcpService, workspaceMcpServiceStub(mcpReady)),
stubPair(IAgentLifecycleService, {
...agentLifecycleStub(),
create: () => mainAgent,
}),
]);

const resumed = svc.resume('s1');
Expand All @@ -1339,7 +1341,7 @@ describe('SessionLifecycleService', () => {
expect(svc.get('s1')).toBeUndefined();
expect(svc.list()).toEqual([]);

resolveMcpReady?.();
releaseMainAgent?.(main);
const handle = await resumed;

expect(handle?.id).toBe('s1');
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -410,6 +410,9 @@ describe('workspace resource sharing (handler chain)', () => {
const m2 = s2.accessor.get(ISessionMcpHandle);
expect(m1.connectionManager).toBe(m2.connectionManager);
expect(connectAll).toHaveBeenCalledTimes(1);
// Session creation no longer waits for the initial connect; the seeded
// handle's readiness promise is the wait point.
await m1.ready;
expect(m1.connectionManager.get('alpha')?.status).toBe('connected');
}, 20000);

Expand Down
6 changes: 5 additions & 1 deletion packages/klient/src/core/facade/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,11 @@ export interface AgentFacade {
getTasks(input?: { activeOnly?: boolean; limit?: number }): Promise<readonly AgentTaskInfo[]>;
stopTask(input: { taskId: string; reason?: string }): Promise<void>;
getTaskOutput(input: { taskId: string; tail?: number }): Promise<string>;
/** Session-merged MCP server entries (workspace set + ephemeral session overlay). */
/**
* Session-merged MCP server entries (workspace set + ephemeral session
* overlay). This is a live snapshot, so entries may still be pending while
* the initial connection attempt runs.
*/
getMcpServers(): Promise<readonly McpServerEntry[]>;
/**
* Trigger a manual full compaction. Async: `true` means the compaction was
Expand Down
Loading
Loading