Skip to content

Commit 42101ae

Browse files
edelaunaedelauna
committed
fix(mcp): preserve selections across concurrent history saves
1 parent 335e78a commit 42101ae

8 files changed

Lines changed: 163 additions & 35 deletions

File tree

‎docs/architecture/mcp-dynamic-tool-loading.md‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@ allowed by the current mode and tool policy exceed `mcp.dynamicLoad.threshold`
77
tools because there is no initial text to rank.
88

99
The selection is stored in task history. Missing or invalid selections are
10-
initialized on demand; a valid empty selection remains initialized. MCP names
10+
initialized on demand; a valid empty selection remains initialized. Restored
11+
nonempty selections are checked against scoped definitions before use; selections
12+
with no available names are ranked again. MCP names
1113
already used in conversation history are retained for provider compatibility,
1214
so resumed tasks can require more than `initialK` declarations. This is an
1315
initial selection limit, not a lifetime cap. Historical names accumulate during
@@ -16,7 +18,9 @@ to reject conversation history. Current server, mode, and tool-policy
1618
restrictions still apply.
1719

1820
If saving fails, the task retains its in-memory selection and continues. Later
19-
normal history saves retry persistence. Storage errors do not trigger API
21+
normal history saves retry persistence. History writes union attached names under
22+
the disk lock, preserving newer historical names when an older save arrives late.
23+
Storage errors do not trigger API
2024
stream retries.
2125

2226
The gate limits advertised definitions; it does not add an execution permission

‎packages/types/src/__tests__/history.test.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@ describe("historyItemSchema", () => {
7272
{ value: {}, expected: undefined },
7373
{ value: [42], expected: undefined },
7474
{ value: [], expected: [] },
75+
{ value: ["invalid"], expected: undefined },
76+
{ value: ["mcp__a__b"], expected: undefined },
7577
{ value: ["mcp--a--b", null, 42], expected: ["mcp--a--b"] },
7678
])("normalizes persisted attachments $value without losing the task", ({ value, expected }) => {
7779
const parsed = historyItemSchema.parse({ ...baseItem, attachedMcpToolNames: value })

‎packages/types/src/history.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@ export type PendingTaskAction = z.infer<typeof pendingTaskActionSchema>
3030
// discard the task record or mark a non-array value as an initialized selection.
3131
export const attachedMcpToolNamesSchema = z.preprocess((value) => {
3232
if (!Array.isArray(value)) return undefined
33-
const names = value.filter((name): name is string => typeof name === "string")
33+
const names = value.filter(
34+
(name): name is string => typeof name === "string" && /^mcp--[A-Za-z_][\w-]*--[A-Za-z_][\w-]*$/.test(name),
35+
)
3436
return value.length > 0 && names.length === 0 ? undefined : names
3537
}, z.array(z.string()).optional())
3638

‎src/core/task-persistence/__tests__/mcpAttachments.integration.spec.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,55 @@ import { TaskHistoryStore } from "../TaskHistoryStore"
55
import { taskMetadata } from "../taskMetadata"
66

77
describe("MCP attachment history round trip", () => {
8+
it.each([false, true])(
9+
"retains historical names after a delayed stale save (separate store: %s)",
10+
async (separateStore) => {
11+
const directory = await fs.mkdtemp(path.join(os.tmpdir(), "mcp-attachments-race-"))
12+
const store = new TaskHistoryStore(directory)
13+
const peer = new TaskHistoryStore(directory)
14+
const restored = new TaskHistoryStore(directory)
15+
try {
16+
const { historyItem } = await taskMetadata({
17+
taskId: "attachment-task",
18+
taskNumber: 1,
19+
messages: [],
20+
workspace: directory,
21+
globalStoragePath: directory,
22+
attachedMcpToolNames: ["mcp--catalog--initial"],
23+
})
24+
await store.initialize()
25+
await store.upsert(historyItem)
26+
await peer.initialize()
27+
let release = () => {}
28+
const gate = new Promise<void>((resolve) => {
29+
release = resolve
30+
})
31+
const olderSave = gate.then(() =>
32+
(separateStore ? peer : store).upsert({
33+
...historyItem,
34+
attachedMcpToolNames: separateStore ? [] : historyItem.attachedMcpToolNames,
35+
}),
36+
)
37+
await store.upsert({
38+
...historyItem,
39+
attachedMcpToolNames: ["mcp--catalog--initial", "mcp--catalog--historical"],
40+
})
41+
release()
42+
await olderSave
43+
await restored.initialize()
44+
expect(restored.get(historyItem.id)?.attachedMcpToolNames).toEqual([
45+
"mcp--catalog--initial",
46+
"mcp--catalog--historical",
47+
])
48+
} finally {
49+
store.dispose()
50+
peer.dispose()
51+
restored.dispose()
52+
await fs.rm(directory, { recursive: true, force: true })
53+
}
54+
},
55+
)
56+
857
it.each([{ names: undefined }, { names: [] }, { names: ["mcp--github--issues"] }])(
958
"preserves $names across a store restart",
1059
async ({ names }) => {

‎src/core/task-persistence/taskStoreConcurrency.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,13 @@ export function mergeHistoryDelta(existing: unknown, incoming: HistoryItem, delt
3838
normalizedDelta.status = attemptedStatus
3939
}
4040
const merged = { ...disk, ...normalizedDelta }
41+
// Attachments accumulate for history compatibility. Merge under the disk lock
42+
// so a delayed metadata snapshot cannot remove a newer historical name.
43+
if (normalizedDelta.attachedMcpToolNames && disk.attachedMcpToolNames) {
44+
merged.attachedMcpToolNames = [
45+
...new Set([...disk.attachedMcpToolNames, ...normalizedDelta.attachedMcpToolNames]),
46+
]
47+
}
4148
if (normalizedDelta.childIds && disk.childIds) {
4249
merged.childIds = [...new Set([...disk.childIds, ...normalizedDelta.childIds])]
4350
}

‎src/core/task/Task.ts‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -297,6 +297,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
297297
readonly metadata: TaskMetadata
298298
readonly attachedMcpToolNames = new Set<string>()
299299
private mcpToolsInitialized = false
300+
private restoredMcpToolsNeedValidation = false
300301
private mcpToolsInitialization?: Promise<void>
301302

302303
todoList?: TodoItem[]
@@ -757,6 +758,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
757758
if (savedMcpToolNames !== undefined) {
758759
for (const name of savedMcpToolNames) this.attachedMcpToolNames.add(name)
759760
this.mcpToolsInitialized = true
761+
this.restoredMcpToolsNeedValidation = savedMcpToolNames.length > 0
760762
}
761763

762764
// Normal use-case is usually retry similar history task with new workspace.
@@ -3033,6 +3035,15 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
30333035
if (this.mcpToolsInitialization) return this.mcpToolsInitialization
30343036

30353037
this.mcpToolsInitialization = (async () => {
3038+
if (this.restoredMcpToolsNeedValidation) {
3039+
const available = new Set(scopedTools.map((tool) => buildMcpToolName(tool.serverName, tool.toolName)))
3040+
for (const name of this.attachedMcpToolNames) {
3041+
if (!available.has(name)) this.attachedMcpToolNames.delete(name)
3042+
}
3043+
this.mcpToolsInitialized = this.attachedMcpToolNames.size > 0
3044+
this.restoredMcpToolsNeedValidation = false
3045+
}
3046+
30363047
const historicalNames = new Set<string>()
30373048
for (const message of this.apiConversationHistory) {
30383049
if (message.role !== "assistant" || !Array.isArray(message.content)) continue
@@ -5968,6 +5979,12 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
59685979
allowedFunctionNames = toolsResult.allowedFunctionNames
59695980
}
59705981

5982+
if (this.abort || this.abandoned) {
5983+
throw new Error(
5984+
`[Task#attemptApiRequest] task ${this.taskId}.${this.instanceId} aborted during request construction`,
5985+
)
5986+
}
5987+
59715988
const shouldIncludeTools = allTools.length > 0
59725989

59735990
// Create an AbortController to allow cancelling the request mid-stream

‎src/core/task/__tests__/Task.persistence.spec.ts‎

Lines changed: 28 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -326,31 +326,32 @@ describe("Task persistence", () => {
326326
},
327327
)
328328

329-
it.each([{ names: [] }, { names: ["mcp--github--issues"] }])(
330-
"restores initialized selections without reranking: $names",
331-
async ({ names }) => {
332-
const rank = vi.spyOn(ToolRouter.getInstance(), "attachInitial")
333-
const task = new Task({
334-
provider: mockProvider,
335-
apiConfiguration: mockApiConfig,
336-
startTask: false,
337-
historyItem: {
338-
id: "restored",
339-
number: 1,
340-
ts: 1,
341-
task: "find tickets",
342-
tokensIn: 0,
343-
tokensOut: 0,
344-
totalCost: 0,
345-
attachedMcpToolNames: names,
346-
},
347-
})
348-
await task.initializeAttachedMcpTools([], 30)
349-
expect([...task.attachedMcpToolNames]).toEqual(names)
350-
expect(rank).not.toHaveBeenCalled()
351-
rank.mockRestore()
352-
},
353-
)
329+
it.each([
330+
{ names: [], savedNames: [] },
331+
{ names: ["mcp--github--issues"], savedNames: ["mcp--github--issues"] },
332+
{ names: ["mcp--github--issues"], savedNames: ["mcp--github--issues", "mcp--missing--tool"] },
333+
])("restores initialized selections without reranking: $names", async ({ names, savedNames }) => {
334+
const rank = vi.spyOn(ToolRouter.getInstance(), "attachInitial")
335+
const task = new Task({
336+
provider: mockProvider,
337+
apiConfiguration: mockApiConfig,
338+
startTask: false,
339+
historyItem: {
340+
id: "restored",
341+
number: 1,
342+
ts: 1,
343+
task: "find tickets",
344+
tokensIn: 0,
345+
tokensOut: 0,
346+
totalCost: 0,
347+
attachedMcpToolNames: savedNames,
348+
},
349+
})
350+
await task.initializeAttachedMcpTools([{ serverName: "github", toolName: "issues", description: "" }], 30)
351+
expect([...task.attachedMcpToolNames]).toEqual(names)
352+
expect(rank).not.toHaveBeenCalled()
353+
rank.mockRestore()
354+
})
354355

355356
it.each(["history", "messages", "provider"] as const)(
356357
"keeps selection without retrying on a %s save failure",
@@ -409,7 +410,7 @@ describe("Task persistence", () => {
409410
}
410411
}
411412

412-
it.each([undefined, null, "abc", 12, {}, [12]])(
413+
it.each([undefined, null, "abc", 12, {}, [12], ["invalid"], ["mcp__github__issues"], ["mcp--missing--tool"]])(
413414
"treats invalid or absent persisted selection %j as uninitialized",
414415
async (value) => {
415416
const historyItem = legacyHistory()
@@ -496,7 +497,7 @@ describe("Task persistence", () => {
496497
],
497498
},
498499
]
499-
await task.initializeAttachedMcpTools([], 1)
500+
await task.initializeAttachedMcpTools([{ serverName: "catalog", toolName: "top", description: "" }], 1)
500501
expect(task.attachedMcpToolNames).toEqual(new Set(["mcp--catalog--previous", "mcp--catalog--top"]))
501502
expect(rank).toHaveBeenCalledTimes(initialized ? 0 : 1)
502503
expect(mockTaskMetadata).toHaveBeenLastCalledWith(

‎src/core/task/__tests__/Task.spec.ts‎

Lines changed: 51 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4725,7 +4725,7 @@ describe("Cline", () => {
47254725
expect(abortSpy).toHaveBeenCalled()
47264726

47274727
// Verify the controller was cleared
4728-
expect(task.currentRequestAbortController).toBeUndefined()
4728+
expect(task["currentRequestAbortController"]).toBeUndefined()
47294729

47304730
// Verify logging
47314731
expect(consoleLogSpy).toHaveBeenCalledWith(expect.stringContaining("Aborting current HTTP request"))
@@ -5036,6 +5036,7 @@ describe("Cline", () => {
50365036
"function" in tool ? [tool.function.name] : [],
50375037
)
50385038
expect(names.filter((name) => name.startsWith("mcp--"))).toHaveLength(expected)
5039+
if (historyCall) expect(names).toContain("mcp--catalog--tool_99")
50395040
expect(names).toContain("read_file")
50405041
expect(names).not.toContain("mcp_load")
50415042
if (provider === providerIdentifiers.gemini)
@@ -5045,6 +5046,51 @@ describe("Cline", () => {
50455046
},
50465047
)
50475048

5049+
it.each(["abort", "abandoned"] as const)(
5050+
"stops a request canceled during tool building: %s",
5051+
async (flag) => {
5052+
const task = new Task({
5053+
provider: mockProvider,
5054+
apiConfiguration: mockApiConfig,
5055+
task: "search tickets",
5056+
startTask: false,
5057+
})
5058+
vi.spyOn(getTaskTestAccess(task), "getSystemPrompt").mockResolvedValue("mock prompt")
5059+
vi.spyOn(task.api, "getModel").mockReturnValue({
5060+
id: "test-model",
5061+
info: { contextWindow: 200000, maxTokens: 4096, supportsPromptCache: false },
5062+
})
5063+
const createMessage = vi.spyOn(task.api, "createMessage")
5064+
let release = () => {}
5065+
const gate = new Promise<void>((resolve) => {
5066+
release = resolve
5067+
})
5068+
let entered = () => {}
5069+
const building = new Promise<void>((resolve) => {
5070+
entered = resolve
5071+
})
5072+
const build = vi
5073+
.spyOn(buildTools, "buildNativeToolsArrayWithRestrictions")
5074+
.mockImplementation(async () => {
5075+
entered()
5076+
await gate
5077+
return { tools: [] }
5078+
})
5079+
try {
5080+
const request = task.attemptApiRequest(0).next()
5081+
await building
5082+
task[flag] = true
5083+
release()
5084+
await expect(request).rejects.toThrow("aborted during request construction")
5085+
expect(createMessage).not.toHaveBeenCalled()
5086+
expect(task["currentRequestAbortController"]).toBeUndefined()
5087+
} finally {
5088+
release()
5089+
build.mockRestore()
5090+
}
5091+
},
5092+
)
5093+
50485094
it.each(["manual", "recovery", "automatic"] as const)(
50495095
"passes task routing settings to the %s condensation build",
50505096
async (path) => {
@@ -5121,7 +5167,7 @@ describe("Cline", () => {
51215167
task.cancelCurrentRequest()
51225168

51235169
expect(abortSpy).toHaveBeenCalledTimes(1)
5124-
expect(task.currentRequestAbortController).toBeUndefined()
5170+
expect(task["currentRequestAbortController"]).toBeUndefined()
51255171
})
51265172

51275173
it("should reject streaming consumption when aborted between chunks", async () => {
@@ -5330,7 +5376,7 @@ describe("Cline", () => {
53305376
},
53315377
]
53325378

5333-
expect(task.currentRequestAbortController).toBeUndefined()
5379+
expect(task["currentRequestAbortController"]).toBeUndefined()
53345380

53355381
const iterator = task.attemptApiRequest(0)
53365382
await iterator.next()
@@ -6066,7 +6112,7 @@ describe("Cline", () => {
60666112
expect(createMessageSpy).not.toHaveBeenCalled()
60676113
// The per-request controller is only created once the request is
60686114
// committed, so a cancelled construction never reaches it.
6069-
expect(task.currentRequestAbortController).toBeUndefined()
6115+
expect(task["currentRequestAbortController"]).toBeUndefined()
60706116
await cancelling
60716117
} finally {
60726118
vi.useRealTimers()
@@ -6126,7 +6172,7 @@ describe("Cline", () => {
61266172
expect(createMessageSpy).not.toHaveBeenCalled()
61276173
// The per-request controller is only created once the request is
61286174
// committed, so a disposed construction never reaches it.
6129-
expect(task.currentRequestAbortController).toBeUndefined()
6175+
expect(task["currentRequestAbortController"]).toBeUndefined()
61306176
await disposal
61316177
} finally {
61326178
vi.useRealTimers()

0 commit comments

Comments
 (0)