diff --git a/.changeset/mcp-first-turn-ready.md b/.changeset/mcp-first-turn-ready.md new file mode 100644 index 0000000000..419372e4e3 --- /dev/null +++ b/.changeset/mcp-first-turn-ready.md @@ -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. diff --git a/packages/acp-server/test/e2e-turn.test.ts b/packages/acp-server/test/e2e-turn.test.ts index 1bef51982a..ec1f5686e9 100644 --- a/packages/acp-server/test/e2e-turn.test.ts +++ b/packages/acp-server/test/e2e-turn.test.ts @@ -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); diff --git a/packages/agent-core-v2/src/agent/mcp/mcpService.ts b/packages/agent-core-v2/src/agent/mcp/mcpService.ts index 3dbe3193f7..2815e286f8 100644 --- a/packages/agent-core-v2/src/agent/mcp/mcpService.ts +++ b/packages/agent-core-v2/src/agent/mcp/mcpService.ts @@ -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` @@ -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'; @@ -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, @@ -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)); @@ -142,7 +151,8 @@ export class AgentMcpService extends Disposable implements IAgentMcpService { } waitForInitialLoad(signal?: AbortSignal): Promise { - return this.mcpHandle.connectionManager.waitForInitialLoad(signal); + const ready = this.mcpHandle.ready; + return signal === undefined ? ready : abortable(ready, signal); } initialLoadDurationMs(): number { diff --git a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts index b83c91b9c4..794f6a364f 100644 --- a/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts +++ b/packages/agent-core-v2/src/session/agentLifecycle/agentLifecycleService.ts @@ -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'; @@ -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'; @@ -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, ) { @@ -145,7 +143,6 @@ export class AgentLifecycleService extends Disposable implements IAgentLifecycle } private async doCreate(agentId: string, opts: CreateAgentOptions): Promise { - const mcpReady = this.mcpHandle.ready; const agentScope = this.ctx.scope(`agents/${agentId}`); const agentHomedir = join(this.bootstrap.homeDir, agentScope); const handle = createScopedChildHandle( @@ -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(); diff --git a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts index f4931503e5..0ac6170e9d 100644 --- a/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts +++ b/packages/agent-core-v2/src/workspace/sessionLifecycle/sessionLifecycleService.ts @@ -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. @@ -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); diff --git a/packages/agent-core-v2/test/agent/mcp/mcp.test.ts b/packages/agent-core-v2/test/agent/mcp/mcp.test.ts index dd9ca41a7e..7754432d5e 100644 --- a/packages/agent-core-v2/test/agent/mcp/mcp.test.ts +++ b/packages/agent-core-v2/test/agent/mcp/mcp.test.ts @@ -219,16 +219,18 @@ describe('AgentMcpService', () => { disposables.dispose(); }); - function createService(manager: FakeMcpManager): AgentMcpService { + function createService( + manager: FakeMcpManager, + ready: Promise = 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 () => { @@ -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((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([]); diff --git a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts index 804585ff9a..f3e30df449 100644 --- a/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts +++ b/packages/agent-core-v2/test/session/agentLifecycle/agentLifecycle.test.ts @@ -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. @@ -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((resolve) => { releaseReady = resolve; @@ -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 () => { diff --git a/packages/agent-core-v2/test/workspace/sessionLifecycle/sessionLifecycle.test.ts b/packages/agent-core-v2/test/workspace/sessionLifecycle/sessionLifecycle.test.ts index d6b6ab5281..2466cf4490 100644 --- a/packages/agent-core-v2/test/workspace/sessionLifecycle/sessionLifecycle.test.ts +++ b/packages/agent-core-v2/test/workspace/sessionLifecycle/sessionLifecycle.test.ts @@ -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((resolve) => { resolveMcpReady = resolve; @@ -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 = Promise.resolve()) { @@ -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((resolve) => { resolveOverlayReady = resolve; @@ -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 () => { @@ -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((resolve) => { - resolveMcpReady = resolve; + let releaseMainAgent: ((handle: IAgentScopeHandle) => void) | undefined; + const mainAgent = new Promise((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'); @@ -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'); diff --git a/packages/agent-core-v2/test/workspace/workspaceResources.test.ts b/packages/agent-core-v2/test/workspace/workspaceResources.test.ts index b4b4f5f837..c6bad21ef4 100644 --- a/packages/agent-core-v2/test/workspace/workspaceResources.test.ts +++ b/packages/agent-core-v2/test/workspace/workspaceResources.test.ts @@ -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); diff --git a/packages/klient/src/core/facade/agent.ts b/packages/klient/src/core/facade/agent.ts index 5c9ac4bde6..f22c4cdc81 100644 --- a/packages/klient/src/core/facade/agent.ts +++ b/packages/klient/src/core/facade/agent.ts @@ -62,7 +62,11 @@ export interface AgentFacade { getTasks(input?: { activeOnly?: boolean; limit?: number }): Promise; stopTask(input: { taskId: string; reason?: string }): Promise; getTaskOutput(input: { taskId: string; tail?: number }): Promise; - /** 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; /** * Trigger a manual full compaction. Async: `true` means the compaction was diff --git a/packages/klient/test/facade.test.ts b/packages/klient/test/facade.test.ts index 04db76c885..77f9f6d9ef 100644 --- a/packages/klient/test/facade.test.ts +++ b/packages/klient/test/facade.test.ts @@ -214,15 +214,16 @@ describe('session skills routing', () => { }); }); -describe('agent mcp / compaction routing', () => { it('getMcpServers routes to agentMcpService.list with the agent scope', async () => { +describe('agent mcp / compaction routing', () => { + it('getMcpServers returns the live snapshot with the agent scope', async () => { const channel = new FakeChannel(); const klient = createKlientFromChannel(channel); const agent = klient.session('s1').agent('main'); const entries = [ - { name: 'mock', transport: 'stdio', status: 'connected', toolCount: 2 }, + { name: 'mock', transport: 'stdio', status: 'pending', toolCount: 0 }, ]; - channel.result = entries; + channel.results.set('agentMcpService.list', entries); await expect(agent.getMcpServers()).resolves.toEqual(entries); expect(channel.calls[0]).toEqual({ scope: { sessionId: 's1', agentId: 'main' }, @@ -230,6 +231,7 @@ describe('agent mcp / compaction routing', () => { it('getMcpServers routes to method: 'list', args: [], }); + expect(channel.calls).toHaveLength(1); }); it('compact issues a manual begin with the optional instruction', async () => { diff --git a/packages/node-sdk/src/sdk-rpc-client-v2.ts b/packages/node-sdk/src/sdk-rpc-client-v2.ts index 3e3dfda2b3..81b2c980a3 100644 --- a/packages/node-sdk/src/sdk-rpc-client-v2.ts +++ b/packages/node-sdk/src/sdk-rpc-client-v2.ts @@ -2151,9 +2151,9 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { /** * Through the session scope (the seeded `ISessionMcpHandle.connectionManager` - * — the workspace handler's one shared manager). Both engines settle the - * initial connect before create/resume returns, so the entry list is final - * here; the v2 `McpServerEntry` is field-identical with v1's + * — the workspace handler's one shared manager). This is a live snapshot: + * create/resume no longer waits for MCP startup, so entries may still be + * pending. The v2 `McpServerEntry` is field-identical with v1's * `McpServerInfo` (the cast bridges the two packages' type declarations). */ override async listMcpServers(input: SessionIdRpcInput): Promise { @@ -2163,7 +2163,7 @@ export class SDKRpcClientV2 extends SDKRpcClientBase { override async getMcpStartupMetrics(input: SessionIdRpcInput): Promise { const mcp = this.requireLiveSession(input.sessionId).accessor.get(ISessionMcpHandle); - await mcp.connectionManager.waitForInitialLoad(); + await mcp.ready; return { durationMs: mcp.connectionManager.initialLoadDurationMs() }; }