From 915ecfbefbaf1b0a7f5876dedbca6c79008ee771 Mon Sep 17 00:00:00 2001 From: marius-kilocode Date: Wed, 10 Jun 2026 14:39:07 +0200 Subject: [PATCH] perf(agent-manager): overlap MCP startup with worktree creation --- .changeset/warm-agent-manager-mcp.md | 5 ++ .../src/agent-manager/AgentManagerProvider.ts | 9 +- .../src/agent-manager/mcp-warmup.ts | 15 ++++ .../tests/unit/agent-manager-arch.test.ts | 10 +++ .../unit/agent-manager-mcp-warmup.test.ts | 84 +++++++++++++++++++ 5 files changed, 120 insertions(+), 3 deletions(-) create mode 100644 .changeset/warm-agent-manager-mcp.md create mode 100644 packages/kilo-vscode/src/agent-manager/mcp-warmup.ts create mode 100644 packages/kilo-vscode/tests/unit/agent-manager-mcp-warmup.test.ts diff --git a/.changeset/warm-agent-manager-mcp.md b/.changeset/warm-agent-manager-mcp.md new file mode 100644 index 00000000000..1a052339eef --- /dev/null +++ b/.changeset/warm-agent-manager-mcp.md @@ -0,0 +1,5 @@ +--- +"kilo-code": patch +--- + +Start MCP servers while Agent Manager worktree sessions initialize to reduce the delay before the first response. diff --git a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts index dd16ecc6098..8356a371e63 100644 --- a/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts +++ b/packages/kilo-vscode/src/agent-manager/AgentManagerProvider.ts @@ -35,6 +35,7 @@ import { diffSummary as localDiffSummary, diffFile as localDiffFile } from "./lo import { parseToolRequest, startFromTool, type ToolRequest } from "./tool-start" import { stopSessionProcesses } from "../kilo-provider/background-process" +import { startSession } from "./mcp-warmup" import { buildKeybindingMap } from "./format-keybinding" import { resolveVersionModels, buildInitialMessages, type CreatedVersion } from "./multi-version" import { Semaphore } from "./semaphore" @@ -810,9 +811,11 @@ export class AgentManagerProvider implements Disposable { }) try { - const { data: session } = await client.session.create( - { directory: worktreePath, platform: PLATFORM }, - { throwOnError: true }, + const { data: session } = await startSession( + client, + worktreePath, + () => client.session.create({ directory: worktreePath, platform: PLATFORM }, { throwOnError: true }), + (...args) => this.log(...args), ) return session } catch (error) { diff --git a/packages/kilo-vscode/src/agent-manager/mcp-warmup.ts b/packages/kilo-vscode/src/agent-manager/mcp-warmup.ts new file mode 100644 index 00000000000..60f95f836e6 --- /dev/null +++ b/packages/kilo-vscode/src/agent-manager/mcp-warmup.ts @@ -0,0 +1,15 @@ +import type { KiloClient } from "@kilocode/sdk/v2/client" + +type Client = Pick +type Log = (...args: unknown[]) => void + +async function warm(client: Client, dir: string, log: Log): Promise { + log(`[MCPWarmup] Starting for ${dir}`) + await client.mcp.status({ directory: dir }, { throwOnError: true }) + log(`[MCPWarmup] Completed for ${dir}`) +} + +export function startSession(client: Client, dir: string, create: () => Promise, log: Log): Promise { + void warm(client, dir, log).catch((err) => log(`[MCPWarmup] Failed for ${dir}:`, err)) + return create() +} diff --git a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts b/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts index 9abcb7bb52a..71733ccdd94 100644 --- a/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts +++ b/packages/kilo-vscode/tests/unit/agent-manager-arch.test.ts @@ -163,6 +163,16 @@ describe("Agent Manager Provider Messages", () => { expect(body).toContain("agentManager.sessionAdded") }) + it("warms MCP before creating every new worktree session", () => { + const body = getMethodBody("createSessionInWorktree") + const warmup = body.indexOf("startSession(") + const create = body.indexOf("client.session.create(") + + expect(warmup).toBeGreaterThanOrEqual(0) + expect(create).toBeGreaterThanOrEqual(0) + expect(warmup).toBeLessThan(create) + }) + it("state-mutating messages wait for state initialization", () => { const body = getMethodBody("shouldWaitForState") const messages = [ diff --git a/packages/kilo-vscode/tests/unit/agent-manager-mcp-warmup.test.ts b/packages/kilo-vscode/tests/unit/agent-manager-mcp-warmup.test.ts new file mode 100644 index 00000000000..4e718980b18 --- /dev/null +++ b/packages/kilo-vscode/tests/unit/agent-manager-mcp-warmup.test.ts @@ -0,0 +1,84 @@ +import { describe, expect, it } from "bun:test" +import { startSession } from "../../src/agent-manager/mcp-warmup" + +function tick(): Promise { + return new Promise((resolve) => queueMicrotask(resolve)) +} + +describe("Agent Manager MCP warmup", () => { + it("starts MCP status for the worktree directory before session creation", async () => { + const calls: unknown[][] = [] + const client = { + mcp: { + status: (input: unknown, opts: unknown) => { + calls.push(["warm", input, opts]) + return Promise.resolve({ data: {} }) + }, + }, + } + + const result = await startSession( + client as never, + "/repo/.kilo/worktrees/feature", + async () => { + calls.push(["session"]) + return "created" + }, + () => {}, + ) + + expect(result).toBe("created") + expect(calls[0]).toEqual(["warm", { directory: "/repo/.kilo/worktrees/feature" }, { throwOnError: true }]) + expect(calls[1]).toEqual(["session"]) + }) + + it("does not wait for MCP warmup before creating the session", async () => { + const calls: string[] = [] + const warmup = new Promise(() => {}) + const client = { + mcp: { + status: () => { + calls.push("warm") + return warmup + }, + }, + } + + const result = await startSession( + client as never, + "/repo/.kilo/worktrees/feature", + async () => { + calls.push("session") + return "created" + }, + () => {}, + ) + + expect(result).toBe("created") + expect(calls).toEqual(["warm", "session"]) + }) + + it("logs and contains MCP warmup failures", async () => { + const logs: unknown[][] = [] + const client = { + mcp: { + status: () => { + throw new Error("connection failed") + }, + }, + } + + const result = await startSession( + client as never, + "/repo/.kilo/worktrees/feature", + async () => "created", + (...args) => logs.push(args), + ) + await tick() + + expect(result).toBe("created") + expect(logs[0]).toEqual(["[MCPWarmup] Starting for /repo/.kilo/worktrees/feature"]) + expect(logs[1]?.[0]).toBe("[MCPWarmup] Failed for /repo/.kilo/worktrees/feature:") + expect(logs[1]?.[1]).toBeInstanceOf(Error) + }) +})