Skip to content

Commit 5ae40b8

Browse files
committed
fix(task): restore caller-layer cancellation for model-metadata fetches
Model-metadata fetches (ensureModelFetched) could outlive the request that started them: a canceled task or a timed-out prompt preview left the fetch awaited, with no signal to abort it and no check before its result was persisted. This restores cancellation handling at the caller layer: - The bounded preview timeout now aborts the metadata fetch it races, instead of leaving the fetcher's promise dangling after the timeout. - Condense paths now check abort/abandoned state before starting and before persisting summarized history, with an added guard before summarization so a canceled task cannot write summarize output. - cancelCurrentRequest aborts the in-flight metadata fetch and detaches waiters, so stale promises no longer retain task state. - Adds a standalone edit-tool coverage test for prompt-section rendering (coverage gap: the tool was only exercised via combined fixtures). Related to #505, #1240.
1 parent cc8ceaf commit 5ae40b8

8 files changed

Lines changed: 982 additions & 51 deletions

File tree

‎src/api/index.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,8 +130,15 @@ export interface ApiHandler {
130130
* Ensures model metadata has been fetched from the remote API so that getModel()
131131
* returns accurate info (context window, pricing, etc.) instead of hardcoded defaults.
132132
* Only router providers that discover models over the network implement this.
133+
*
134+
* `signal` bounds the caller's wait: when it aborts (e.g. the caller's bounded
135+
* metadata wait expired or the owning task was cancelled), the returned promise
136+
* settles with a rejection so no handler-side waiter outlives its caller.
137+
* Fetchers that observe the signal may also stop their network request; the
138+
* shared, de-duplicated catalog fetch may still complete and populate the model
139+
* cache, which is by design for concurrent waiters.
133140
*/
134-
ensureModelFetched?(): Promise<void>
141+
ensureModelFetched?(signal?: AbortSignal): Promise<void>
135142

136143
/**
137144
* Optional context window for context-management / auto-condense when it must differ from

‎src/api/providers/__tests__/zoo-gateway.spec.ts‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -724,6 +724,72 @@ describe("ZooGatewayHandler", () => {
724724
expect(refreshModels).not.toHaveBeenCalled()
725725
})
726726

727+
it("settles the waiter with a rejection when the signal aborts mid-fetch", async () => {
728+
// A caller that gives up must not leave a handler-side waiter pending
729+
// on the (shared) catalog fetch: with an observing signal, the
730+
// ensureModelFetched promise rejects at abort time, while the
731+
// underlying fetch continues untouched for any other waiter.
732+
const { getModels } = await import("../fetchers/modelCache")
733+
vitest.mocked(getModels).mockImplementationOnce(() => new Promise(() => {}))
734+
735+
const handler = new ZooGatewayHandler(mockOptions)
736+
const controller = new AbortController()
737+
738+
const wait = handler.ensureModelFetched(controller.signal)
739+
// Let the waiter attach its abort listener before cancelling.
740+
await Promise.resolve()
741+
controller.abort()
742+
743+
await expect(wait).rejects.toThrow()
744+
})
745+
746+
it("settles the waiter when the fetch wins against a live signal and detaches the listener", async () => {
747+
// Fetch-wins branch: resolve() must settle the await (a dropped
748+
// resolve or a detached .then handler hangs this test), the abort
749+
// listener must be registered with the real { once: true } options
750+
// object, and the detach must target the *same* event name/handler
751+
// pair that was registered — a mutated event name detaches nothing.
752+
const handler = new ZooGatewayHandler(mockOptions)
753+
const controller = new AbortController()
754+
const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener")
755+
const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener")
756+
757+
await handler.ensureModelFetched(controller.signal)
758+
759+
expect(addEventListenerSpy).toHaveBeenCalledWith("abort", expect.any(Function), { once: true })
760+
const registered = addEventListenerSpy.mock.calls.find(([event]) => event === "abort")
761+
expect(registered).toBeDefined()
762+
expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", registered?.[1])
763+
})
764+
765+
it("rejects a signal-observing waiter with the fetch error and detaches the listener", async () => {
766+
// Rejection-branch twin of the fetch-wins test: reject(error) must
767+
// propagate the catalog failure to the waiter (a dropped reject hangs
768+
// this test) and the listener must be detached under the right event
769+
// name. The no-signal reject path cannot attach a listener, so this is the only
770+
// coverage of the reject-side detach.
771+
const { getModels } = await import("../fetchers/modelCache")
772+
vitest.mocked(getModels).mockRejectedValueOnce(new Error("network down"))
773+
774+
const handler = new ZooGatewayHandler(mockOptions)
775+
const controller = new AbortController()
776+
const addEventListenerSpy = vitest.spyOn(controller.signal, "addEventListener")
777+
const removeEventListenerSpy = vitest.spyOn(controller.signal, "removeEventListener")
778+
779+
await expect(handler.ensureModelFetched(controller.signal)).rejects.toThrow("network down")
780+
const registered = addEventListenerSpy.mock.calls.find(([event]) => event === "abort")
781+
expect(registered).toBeDefined()
782+
expect(removeEventListenerSpy).toHaveBeenCalledWith("abort", registered?.[1])
783+
})
784+
785+
it("never starts a wait when the signal is already aborted", async () => {
786+
const handler = new ZooGatewayHandler(mockOptions)
787+
const controller = new AbortController()
788+
controller.abort()
789+
790+
await expect(handler.ensureModelFetched(controller.signal)).rejects.toThrow()
791+
})
792+
727793
it("skips the fetch when models are already populated", async () => {
728794
const handler = new ZooGatewayHandler(mockOptions)
729795
const { getModels, refreshModels } = await import("../fetchers/modelCache")

‎src/api/providers/router-provider.ts‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,8 +108,35 @@ export abstract class RouterProvider extends BaseProvider {
108108
return this.modelFetchPromise
109109
}
110110

111-
async ensureModelFetched(): Promise<void> {
112-
await this.fetchModel()
111+
async ensureModelFetched(signal?: AbortSignal): Promise<void> {
112+
// A caller that already gave up must not start (or keep) a wait on the
113+
// shared catalog fetch.
114+
if (signal?.aborted) {
115+
throw signal.reason
116+
}
117+
118+
const fetch = this.fetchModel()
119+
if (!signal) {
120+
await fetch
121+
return
122+
}
123+
124+
// Detach this waiter as soon as the signal aborts; the shared in-flight
125+
// fetch continues for any other waiter and still populates the cache.
126+
await new Promise<void>((resolve, reject) => {
127+
const onAbort = () => reject(signal.reason)
128+
signal.addEventListener("abort", onAbort, { once: true })
129+
fetch.then(
130+
() => {
131+
signal.removeEventListener("abort", onAbort)
132+
resolve()
133+
},
134+
(error: unknown) => {
135+
signal.removeEventListener("abort", onAbort)
136+
reject(error)
137+
},
138+
)
139+
})
113140
}
114141

115142
override getModel(): { id: string; info: ModelInfo } {

‎src/core/prompts/__tests__/sections.spec.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,21 @@ describe("getCapabilitiesSection", () => {
165165
expect(result).not.toContain("only files matching")
166166
})
167167

168+
it("keeps the capability clause and restriction for a model-included standalone edit tool", () => {
169+
// The restricted edit group's default edit tools are disabled, but the
170+
// model catalog re-adds the standalone `edit` tool: it is still an edit
171+
// capability, so the clause and the file restriction both render.
172+
const result = getCapabilitiesSection(
173+
policyFor([["edit", { fileRegex: "\\.md$", description: "Markdown files only" }]], {
174+
disabledTools: ["write_to_file", "apply_diff"],
175+
modelInfo: { contextWindow: 128_000, supportsPromptCache: true, includedTools: ["edit"] },
176+
}),
177+
)
178+
179+
expect(result).toContain("write and edit files")
180+
expect(result).toContain("(in this mode only files matching '\\.md$' can be edited — Markdown files only)")
181+
})
182+
168183
it("lists files guidance only when list_files is available", () => {
169184
const withListFiles = getCapabilitiesSection(policyFor(["read"]))
170185
expect(withListFiles).toContain("you can use the list_files tool")

‎src/core/task/Task.ts‎

Lines changed: 86 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,17 @@ import { type TaskExecutionContext } from "./providerHandoff"
142142

143143
const MAX_EXPONENTIAL_BACKOFF_SECONDS = 600 // 10 minutes
144144
const DEFAULT_USAGE_COLLECTION_TIMEOUT_MS = 5000 // 5 seconds
145+
// Upper bound on awaiting lazily loaded model metadata. Some model-catalog
146+
// fetchers (e.g. OpenRouter's bare axios GET) have no request timeout, so a
147+
// hung endpoint must not stall streaming, condense, or context-window
148+
// handling. On expiry the caller aborts its per-call AbortSignal (detaching
149+
// the provider-side waiter) and falls back to the handler's existing
150+
// getModel().info metadata — the same degradation a rejected fetch produces.
151+
// Kept in sync with PREVIEW_MODEL_FETCH_TIMEOUT_MS in the preview path
152+
// (src/core/webview/generateSystemPrompt.ts). Deliberately duplicated, not
153+
// shared: importing from the webview layer would close a
154+
// Task -> generateSystemPrompt -> ClineProvider -> Task circular import.
155+
export const MODEL_FETCH_TIMEOUT_MS = 5_000
145156
const QUEUED_FEEDBACK_SAVE_RETRY_DELAYS_MS = [250, 1_000, 4_000] as const
146157

147158
type QueuedAskResolution = { response: ClineAskResponse; requiresDurableAck: boolean }
@@ -313,6 +324,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
313324
private readonly globalStoragePath: string
314325
abort: boolean = false
315326
currentRequestAbortController?: AbortController
327+
/**
328+
* Controller for the waiter on an in-flight `ensureModelFetched()` call (see
329+
* safeEnsureModelFetched). Aborting it detaches this task from the provider-side
330+
* metadata fetch; it is aborted when the bounded wait expires and when the task's
331+
* current request is cancelled (cancel/dispose path in cancelCurrentRequest).
332+
*/
333+
metadataFetchAbortController?: AbortController
316334
skipPrevResponseIdOnce: boolean = false
317335

318336
// TaskStatus
@@ -1858,8 +1876,8 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
18581876
const state = await this.providerRef.deref()?.getState()
18591877
const requestModelInfo = await this.safeEnsureModelFetched()
18601878

1861-
// A cancellation landing during the metadata wait must stop manual
1862-
// condensation before any prompt build or summarization request.
1879+
// A cancellation landing during the bounded metadata wait must stop
1880+
// manual condensation before any prompt build or summarization request.
18631881
if (this.abort || this.abandoned) {
18641882
return
18651883
}
@@ -1920,6 +1938,12 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
19201938

19211939
const filesReadByRoo = await this.getFilesReadByRooSafely("condenseContext")
19221940

1941+
// A cancellation landing while the summarization inputs are gathered must
1942+
// stop manual condensation before any summarization request is issued.
1943+
if (this.abort || this.abandoned) {
1944+
return
1945+
}
1946+
19231947
const {
19241948
messages,
19251949
summary,
@@ -1941,6 +1965,11 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
19411965
cwd: this.cwd,
19421966
rooIgnoreController: this.rooIgnoreController,
19431967
})
1968+
// A cancellation landing during the summarization request must stop
1969+
// manual condensation before it replaces and persists the history.
1970+
if (this.abort || this.abandoned) {
1971+
return
1972+
}
19441973
if (error) {
19451974
await this.say(
19461975
"condense_context_error",
@@ -2629,6 +2658,12 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
26292658
this.currentRequestAbortController.abort()
26302659
this.currentRequestAbortController = undefined
26312660
}
2661+
// A metadata-fetch waiter still in flight is as stale as an abandoned
2662+
// stream: detach it from the provider-side fetch on cancel/dispose.
2663+
if (this.metadataFetchAbortController) {
2664+
this.metadataFetchAbortController.abort()
2665+
this.metadataFetchAbortController = undefined
2666+
}
26322667
}
26332668

26342669
/**
@@ -4283,28 +4318,63 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
42834318
/**
42844319
* Ensures router-provider model metadata is loaded before getModel() is used for
42854320
* context management or streaming. Failures fall back to hardcoded defaults rather
4286-
* than aborting the task.
4321+
* than aborting the task; the wait is bounded by MODEL_FETCH_TIMEOUT_MS (see the
4322+
* constant for the degradation semantics). On expiry or cancellation the per-call
4323+
* AbortSignal detaches this task's waiter from the provider-side fetch instead of
4324+
* leaving a handler-side promise waiting on it indefinitely.
42874325
*
42884326
* The return value is the settled post-wait read of getModel().info: request
42894327
* entry points (attemptApiRequest, condenseContext) await this once before
42904328
* prompt generation and share the returned snapshot between getSystemPrompt
4291-
* and every tool-array build of the same request, so a fetch that resolves
4292-
* mid-request cannot move model-specific tool policy between prompt time and
4293-
* request time. Callers that do not thread a snapshot keep awaiting this
4294-
* immediately before their own getModel() read; that per-site guard remains the
4295-
* standalone/fallback read path, and repeat awaits stay cheap once a fetch has
4296-
* succeeded because the provider caches successes. RouterProvider already
4297-
* negative-caches catalog misses with a TTL (`missingModelRefreshAt`).
4329+
* and every tool-array build of the same request, so a fetch that resolves after
4330+
* the bounded wait was abandoned cannot move model-specific tool policy between
4331+
* prompt time and request time. Callers that do not thread a snapshot keep
4332+
* awaiting this immediately before their own getModel() read; that per-site guard
4333+
* remains the standalone/fallback read path, and repeat awaits stay cheap once a
4334+
* fetch has succeeded because the provider caches successes. RouterProvider
4335+
* already negative-caches catalog misses with a TTL (`missingModelRefreshAt`).
42984336
*/
42994337
private async safeEnsureModelFetched(): Promise<ModelInfo> {
4338+
// Per-call controller: its signal makes ensureModelFetched() settle on
4339+
// abort, so neither this race nor the provider-side waiter outlives the
4340+
// bounded wait. cancel/dispose aborts it early via cancelCurrentRequest.
4341+
const controller = new AbortController()
4342+
this.metadataFetchAbortController = controller
4343+
// Promise.race attaches handlers to both inputs, so the fetch rejecting
4344+
// after the abort is already considered handled — no extra .catch needed.
4345+
let timeoutId: ReturnType<typeof setTimeout> | undefined
4346+
let timedOut = false
43004347
try {
4301-
await this.api.ensureModelFetched?.()
4348+
await Promise.race([
4349+
this.api.ensureModelFetched?.(controller.signal),
4350+
new Promise<void>((resolve) => {
4351+
timeoutId = setTimeout(() => {
4352+
timedOut = true
4353+
resolve()
4354+
}, MODEL_FETCH_TIMEOUT_MS)
4355+
}),
4356+
])
4357+
if (timedOut) {
4358+
console.warn(
4359+
`[Task#${this.taskId}] Timed out after ${MODEL_FETCH_TIMEOUT_MS}ms fetching model metadata; using fallback model info.`,
4360+
)
4361+
}
43024362
} catch (error) {
43034363
console.error(
43044364
`[Task#${this.taskId}] Failed to fetch model metadata:`,
43054365
error instanceof Error ? error.message : error,
43064366
)
4367+
} finally {
4368+
if (timeoutId) {
4369+
clearTimeout(timeoutId)
4370+
}
4371+
if (this.metadataFetchAbortController === controller) {
4372+
this.metadataFetchAbortController = undefined
4373+
}
4374+
// Unconditional: settles any waiter still attached to this call.
4375+
controller.abort()
43074376
}
4377+
// The post-wait read is the settled snapshot callers must share per request.
43084378
return this.api.getModel().info
43094379
}
43104380

@@ -4518,9 +4588,9 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
45184588
// Thread the request state snapshot into prompt generation so the prompt
45194589
// and the runtime tools (built below from the same `state`) stay aligned
45204590
// even if settings change while this method waits on MCP or rate limits.
4521-
// Capture one model-info snapshot per request, shared by the prompt and
4522-
// every tool array built below; prefer the caller's snapshot when one was
4523-
// threaded.
4591+
// Capture one bounded-wait model-info snapshot per request, shared by the
4592+
// prompt and every tool array built below; prefer the caller's snapshot
4593+
// when one was threaded.
45244594
const requestModelInfo = options.requestModelInfo ?? (await this.safeEnsureModelFetched())
45254595
// Retry recursions must reuse this snapshot instead of re-deriving it: a
45264596
// metadata fetch landing between attempts would otherwise move
@@ -4530,8 +4600,8 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
45304600
const retryOptions = options.requestModelInfo === undefined ? { ...options, requestModelInfo } : options
45314601
const systemPrompt = await this.getSystemPrompt(state, requestModelInfo)
45324602

4533-
// A cancellation landing during the rate-limit countdown, the metadata wait,
4534-
// or the MCP wait inside getSystemPrompt must stop this request before any
4603+
// A cancellation landing during the rate-limit countdown, the bounded metadata
4604+
// wait, or the MCP wait inside getSystemPrompt must stop this request before any
45354605
// tool array, AbortController, or createMessage call is issued for it.
45364606
if (this.abort || this.abandoned) {
45374607
throw new Error(

0 commit comments

Comments
 (0)