From fbf4c73823dd8e44f9b0f7d99c7a9142a209be04 Mon Sep 17 00:00:00 2001 From: daewoongoh Date: Sat, 7 Nov 2026 07:54:05 +0900 Subject: [PATCH 1/7] feat(task): generate a session title from the first message - Generate a concise 5-10 word title from the root user message in the background - Fall back to the first 42 characters of the sanitized user message when generation fails - Enforce 15-second timeout on title generation without blocking the task loop - Persist title via setTitleIfUnset inside the atomic per-task store lock - Preserve manual renames and manual title clears over auto-generated titles - Skip session title generation in Test extension mode and for resumed tasks or subtasks --- src/core/task-persistence/TaskHistoryStore.ts | 70 +++ .../TaskHistoryStore.crossInstance.spec.ts | 98 +++ src/core/task/Task.ts | 33 +- .../task/__tests__/Task.session-title.spec.ts | 295 +++++++++ src/core/task/__tests__/session-title.spec.ts | 584 ++++++++++++++++++ src/core/task/session-title.ts | 233 +++++++ src/core/webview/ClineProvider.ts | 42 +- .../ClineProvider.taskHistory.spec.ts | 128 ++++ 8 files changed, 1478 insertions(+), 5 deletions(-) create mode 100644 src/core/task/__tests__/Task.session-title.spec.ts create mode 100644 src/core/task/__tests__/session-title.spec.ts create mode 100644 src/core/task/session-title.ts diff --git a/src/core/task-persistence/TaskHistoryStore.ts b/src/core/task-persistence/TaskHistoryStore.ts index c1c50cca1e..a9e0243823 100644 --- a/src/core/task-persistence/TaskHistoryStore.ts +++ b/src/core/task-persistence/TaskHistoryStore.ts @@ -26,6 +26,14 @@ function mergeWithDisk(delta: Partial): (existing: unknown, incomin return (existing, incoming) => mergeHistoryDelta(existing, incoming as HistoryItem, delta) } +/** Thrown inside the `setTitleIfUnset` merge callback to abort the write without touching the file. */ +class TitleWriteSkipped extends Error { + constructor(readonly reason: "missing" | "titled") { + // Stryker disable next-line StringLiteral: diagnostic error message only, caught internally + super(`title write skipped: ${reason}`) + } +} + /** * Durable intent for the one repair that spans an active delegated child and * its parent. Task files remain authoritative; this file only records the @@ -996,6 +1004,68 @@ export class TaskHistoryStore { // ────────────────────────────── Atomic read-modify-write ────────────────────────────── + /** + * Set `title` only while the item has none, checking the title on disk inside the + * per-file lock. The cache can be stale when another store (for example a second + * provider on the same storage path) saved a manual title, and a delta merge would + * otherwise overwrite it. + * + * @returns "conflict" when a title already exists or the record was deleted on disk (the cache is refreshed). + * @throws If the task ID is not present in the cache. + */ + public setTitleIfUnset(taskId: string, title: string): Promise<"applied" | "conflict"> { + return this.withLock(async () => { + const current = this.cache.get(taskId) + if (!current) { + throw new Error(`[TaskHistoryStore] setTitleIfUnset: task ${taskId} not found in cache`) + } + if (current.title) { + return "conflict" + } + + const filePath = await this.getTaskFilePath(taskId) + let written: HistoryItem = { ...current, title } + try { + await safeWriteJson(filePath, written, { + merge: (existing, incoming) => { + // Throwing aborts the write, so a deleted record is never recreated from the cache + // and an existing title is never rewritten. + // Stryker disable next-line ConditionalExpression: defense-in-depth, corrupted disk content cannot be merged + if (!existing || typeof existing !== "object" || !("id" in existing)) { + throw new TitleWriteSkipped("missing") + } + if ((existing as HistoryItem).title) { + // Stryker disable next-line StringLiteral: internal discrimination literal between missing and titled reasons + throw new TitleWriteSkipped("titled") + } + written = mergeHistoryDelta(existing, incoming as HistoryItem, { title }) + return written + }, + }) + } catch (error) { + if (!(error instanceof TitleWriteSkipped)) { + throw error + } + if (error.reason === "missing") { + this.cache.delete(taskId) + } else { + const diskItem = await this.readTaskFile(taskId) + // Stryker disable next-line ConditionalExpression: defense-in-depth, disk item is always present when merge found titled record + if (diskItem) { + this.cache.set(taskId, diskItem) + } + } + return "conflict" + } + + this.cache.set(taskId, written) + if (this.onWrite) { + await this.onWrite(this.getAll()) + } + return "applied" + }) + } + /** * Read a HistoryItem from the in-memory cache and write back an updated version, * all within a single lock acquisition so no concurrent writer can interleave diff --git a/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts b/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts index 5b85c9b4e2..98efed9a61 100644 --- a/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts +++ b/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts @@ -73,6 +73,104 @@ describe("TaskHistoryStore cross-instance safety", () => { await fs.rm(tmpDir, { recursive: true, force: true }).catch(() => {}) }) + describe("setTitleIfUnset", () => { + it("keeps a manual title that another instance wrote after this cache was loaded", async () => { + await storeA.initialize() + await storeB.initialize() + await storeA.upsert(makeHistoryItem({ id: "title-task", task: "Original prompt" })) + await storeB.reconcile() + expect(storeB.get("title-task")?.title).toBeUndefined() + + // Instance A still caches no title when B renames the task on disk. + await storeB.atomicReadAndUpdate("title-task", (item) => ({ ...item, title: "Manual title" })) + expect(storeA.get("title-task")?.title).toBeUndefined() + + await expect(storeA.setTitleIfUnset("title-task", "Auto title")).resolves.toBe("conflict") + + expect(storeA.get("title-task")?.title).toBe("Manual title") + await storeB.reconcile() + expect(storeB.get("title-task")?.title).toBe("Manual title") + }) + + it("does not recreate a task that another instance deleted", async () => { + await storeA.initialize() + await storeB.initialize() + await storeA.upsert(makeHistoryItem({ id: "title-deleted", task: "Original prompt" })) + await storeB.reconcile() + + await storeB.delete("title-deleted") + expect(storeA.get("title-deleted")).toBeDefined() + + await expect(storeA.setTitleIfUnset("title-deleted", "Auto title")).resolves.toBe("conflict") + + expect(storeA.get("title-deleted")).toBeUndefined() + await storeB.reconcile() + expect(storeB.get("title-deleted")).toBeUndefined() + }) + + it("applies the title when no instance has set one", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "title-free", task: "Original prompt" })) + + await expect(storeA.setTitleIfUnset("title-free", "Auto title")).resolves.toBe("applied") + + expect(storeA.get("title-free")).toMatchObject({ title: "Auto title", task: "Original prompt" }) + }) + + it("persists the title to disk so reloading from storage preserves it", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "persist-check", task: "Original task prompt" })) + + await expect(storeA.setTitleIfUnset("persist-check", "Persisted title")).resolves.toBe("applied") + + // Reload a fresh store instance pointing to the same disk directory + const storeReloaded = new TaskHistoryStore(tmpDir) + await storeReloaded.initialize() + + const diskItem = storeReloaded.get("persist-check") + expect(diskItem).toBeDefined() + expect(diskItem?.title).toBe("Persisted title") + expect(diskItem?.task).toBe("Original task prompt") + storeReloaded.dispose() + }) + + it("throws when the task is not in the cache", async () => { + await storeA.initialize() + await expect(storeA.setTitleIfUnset("missing-id", "Auto title")).rejects.toThrow("not found in cache") + }) + + it("returns conflict immediately when cached item already has a title", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "already-titled", title: "Existing", task: "Prompt" })) + await expect(storeA.setTitleIfUnset("already-titled", "New title")).resolves.toBe("conflict") + }) + + it("notifies onWrite callback when title is successfully applied", async () => { + const onWriteSpy = vi.fn().mockResolvedValue(undefined) + const storeWithOnWrite = new TaskHistoryStore(tmpDir, { onWrite: onWriteSpy }) + await storeWithOnWrite.initialize() + await storeWithOnWrite.upsert(makeHistoryItem({ id: "title-onwrite", task: "Prompt" })) + onWriteSpy.mockClear() + + await expect(storeWithOnWrite.setTitleIfUnset("title-onwrite", "Notified title")).resolves.toBe("applied") + expect(onWriteSpy).toHaveBeenCalledWith( + expect.arrayContaining([expect.objectContaining({ id: "title-onwrite", title: "Notified title" })]), + ) + storeWithOnWrite.dispose() + }) + + it("re-throws unexpected errors from safeWriteJson", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "title-io-err", task: "Prompt" })) + + const { safeWriteJson } = await import("../../../utils/safeWriteJson") + const mockedSafeWrite = vi.mocked(safeWriteJson) + mockedSafeWrite.mockRejectedValueOnce(new Error("disk IO error")) + + await expect(storeA.setTitleIfUnset("title-io-err", "Auto title")).rejects.toThrow("disk IO error") + }) + }) + it("two instances can write different tasks without conflict", async () => { await storeA.initialize() await storeB.initialize() diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index b911479bc6..4b108221b1 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -145,6 +145,7 @@ import { prepareApiConversationMessage } from "./apiConversationHistory" import { shouldAddUserMessageToHistory } from "./messageCounting" import { type TaskExecutionContext } from "./providerHandoff" import { initialRunState, type Ask, type RunEvent, type RunState } from "./run-state/runState" +import { generateSessionTitle } from "./session-title" import { applyRunEvent, type RunRejection } from "./run-state/runStateShadow" import { divergenceKey, @@ -286,6 +287,9 @@ export class PendingActionSettlementError extends Error { } } +// Literal, not the vscode enum: specs that mock "vscode" without ExtensionMode would throw on the read. +const EXTENSION_MODE_TEST = 3 + export class Task extends EventEmitter implements TaskLike { readonly taskId: string readonly rootTaskId?: string @@ -3110,6 +3114,31 @@ export class Task extends EventEmitter implements TaskLike { await this.say("text", task, images) + // Fire-and-forget session title generation: the run must never wait + // on it and no title failure may surface in the chat. Only a root + // task's first message triggers this — resumes go through + // resumeTaskFromHistory, subtask continuations re-enter the task + // loop directly, and both must keep their existing names. `say` + // above persisted the history item, so the conditional write will + // find it. + // The extra provider request would break E2E and visual tests that assert + // exact request sequences and rendered task names, so Test mode skips it. + if (!this.parentTask && task) { + const provider = this.providerRef.deref() + if (provider && provider.context.extensionMode !== EXTENSION_MODE_TEST) { + void generateSessionTitle({ + provider, + taskId: this.taskId, + userMessage: task, + apiConfiguration: this.apiConfiguration, + }).catch((error) => { + // generateSessionTitle never rejects by contract; this + // arm is defense-in-depth for future regressions. + console.error(`[Task#startTask] session title generation failed: ${error}`) + }) + } + } + // Check for too many MCP tools and warn the user const { enabledToolCount, enabledServerCount } = await this.getEnabledMcpToolsCount() if (enabledToolCount > MAX_MCP_TOOLS_THRESHOLD) { @@ -4419,7 +4448,7 @@ export class Task extends EventEmitter implements TaskLike { nativeToolCallParserScope, ) - // Stryker disable next-line ConditionalExpression, EqualityOperator: equivalent mutants, streamingKey is always defined here, so both conditions reduce to the partialToolUse check + // Stryker disable next-line ConditionalExpression,EqualityOperator: equivalent mutants, streamingKey is always defined here, so both conditions reduce to the partialToolUse check if (partialToolUse && streamingKey !== undefined) { // Reuse the compound key built above: it is exactly the key the start event // registered under, so the delta lookup needs no second encoding. @@ -4855,7 +4884,7 @@ export class Task extends EventEmitter implements TaskLike { // name (provider quirk), the compound-key lookup above misses and both the // parser entry and the dedup tracking would stay stale. Resolve the tracked // entry by id and target the real compound key for finalization and cleanup. - // Stryker disable next-line ConditionalExpression, LogicalOperator, EqualityOperator: equivalent mutants, the both-miss case is unreachable via the raw-chunk pipeline + // Stryker disable next-line ConditionalExpression,LogicalOperator,EqualityOperator: equivalent mutants, the both-miss case is unreachable via the raw-chunk pipeline if (finalToolUse === null && toolUseIndex === undefined) { const resolved = NativeToolCallParser.getStreamingToolCallById( event.id, diff --git a/src/core/task/__tests__/Task.session-title.spec.ts b/src/core/task/__tests__/Task.session-title.spec.ts new file mode 100644 index 0000000000..719e62fa27 --- /dev/null +++ b/src/core/task/__tests__/Task.session-title.spec.ts @@ -0,0 +1,295 @@ +// npx vitest run core/task/__tests__/Task.session-title.spec.ts + +import type { HistoryItem, ProviderSettings } from "@roo-code/types" +import { providerIdentifiers } from "@roo-code/types/provider-identifiers" + +import { Task } from "../Task" +import { ClineProvider } from "../../webview/ClineProvider" +import type { GenerateSessionTitleOptions } from "../session-title" + +vi.mock("../session-title", () => ({ + generateSessionTitle: vi.fn().mockResolvedValue(undefined), +})) + +// Import the mocked module for spy access. +import { generateSessionTitle } from "../session-title" + +vi.mock("@roo-code/telemetry", () => ({ + TelemetryService: { + hasInstance: vi.fn().mockReturnValue(true), + createInstance: vi.fn(), + get instance() { + return { + captureTaskCreated: vi.fn(), + captureTaskRestarted: vi.fn(), + captureModeSwitch: vi.fn(), + captureConversationMessage: vi.fn(), + captureLlmCompletion: vi.fn(), + captureConsecutiveMistakeError: vi.fn(), + captureCodeActionUsed: vi.fn(), + setProvider: vi.fn(), + } + }, + }, +})) + +vi.mock("vscode", () => { + const mockDisposable = { dispose: vi.fn() } + + return { + TabInputTextDiff: vi.fn(), + CodeActionKind: { + QuickFix: { value: "quickfix" }, + RefactorRewrite: { value: "refactor.rewrite" }, + }, + window: { + createTextEditorDecorationType: vi.fn().mockReturnValue({ dispose: vi.fn() }), + showErrorMessage: vi.fn(), + }, + workspace: { + getConfiguration: vi.fn(() => ({ get: (_k: string, d: unknown) => d })), + workspaceFolders: [{ uri: { fsPath: "/mock/workspace/path" }, name: "mock-workspace", index: 0 }], + createFileSystemWatcher: vi.fn(() => ({ + onDidCreate: vi.fn(() => mockDisposable), + onDidDelete: vi.fn(() => mockDisposable), + onDidChange: vi.fn(() => mockDisposable), + dispose: vi.fn(), + })), + onDidSaveTextDocument: vi.fn(() => mockDisposable), + }, + env: { uriScheme: "vscode", language: "en" }, + EventEmitter: vi.fn(), + Disposable: { from: vi.fn() }, + TabInputText: vi.fn(), + version: "1.85.0", + } +}) + +vi.mock("../../environment/getEnvironmentDetails", () => ({ + getEnvironmentDetails: vi.fn().mockResolvedValue(""), +})) + +vi.mock("../../ignore/RooIgnoreController") + +vi.mock("p-wait-for", () => ({ + default: vi.fn().mockImplementation(async () => Promise.resolve()), +})) + +vi.mock("delay", () => ({ + __esModule: true, + default: vi.fn().mockResolvedValue(undefined), +})) + +type StartTaskAccess = { + startTask: (task?: string, images?: string[]) => Promise + resumeTaskFromHistory: () => Promise + getEnabledMcpToolsCount: () => Promise<{ enabledToolCount: number; enabledServerCount: number }> + initiateTaskLoop: (userContent: unknown[]) => Promise +} + +// The private-member double assertion is the established Task-spec pattern +// (see Task.spec.ts getTaskTestAccess): the spies below must reach private +// startup collaborators without widening their production visibility. +function getStartTaskAccess(task: Task): StartTaskAccess { + return task as unknown as StartTaskAccess +} + +type ProviderStateShaped = Partial>> + +/** + * The trigger path only reads provider startup collaborators; the double is + * cast once here (AGENTS.md) naming the members tests rely on: postStateToWebviewWithoutTaskHistory, + * getState, updateTaskHistory, log, on, off. + */ +function createMockProvider(): ClineProvider { + return { + context: { globalStorageUri: { fsPath: "/test/storage" } }, + getState: vi.fn().mockResolvedValue({} as ProviderStateShaped), + log: vi.fn(), + on: vi.fn(), + off: vi.fn(), + postStateToWebview: vi.fn().mockResolvedValue(undefined), + postStateToWebviewWithoutTaskHistory: vi.fn().mockResolvedValue(undefined), + updateTaskHistory: vi.fn().mockResolvedValue(undefined), + } as unknown as ClineProvider +} + +function createTestTask(provider: ClineProvider, parentTask?: Task, historyItem?: HistoryItem): Task { + const apiConfiguration: ProviderSettings = { + apiProvider: providerIdentifiers.anthropic, + apiModelId: "claude-3-5-sonnet-20241022", + apiKey: "test-api-key", + } + + return new Task({ + provider, + apiConfiguration, + // A history item must be the only task source or it is not a history + // task; omit the task text for that case. + task: historyItem ? undefined : "example root task", + parentTask, + historyItem, + startTask: false, + }) +} + +function stubStartup(task: Task): void { + vi.spyOn(task, "say").mockResolvedValue(undefined) + const taskAccess = getStartTaskAccess(task) + vi.spyOn(taskAccess, "getEnabledMcpToolsCount").mockResolvedValue({ enabledToolCount: 0, enabledServerCount: 0 }) + vi.spyOn(taskAccess, "initiateTaskLoop").mockResolvedValue(undefined) +} + +function lastTitleCall(): GenerateSessionTitleOptions | undefined { + const calls = vi.mocked(generateSessionTitle).mock.calls + return calls.length > 0 ? calls[calls.length - 1][0] : undefined +} + +describe("Task#startTask session title trigger", () => { + beforeEach(() => { + vi.mocked(generateSessionTitle).mockClear() + }) + + it("fires title generation for a root task without awaiting it", async () => { + const provider = createMockProvider() + const task = createTestTask(provider) + stubStartup(task) + + let resolved = false + vi.mocked(generateSessionTitle).mockImplementation(async () => { + resolved = true + }) + + await getStartTaskAccess(task).startTask("title trigger message") + + const call = lastTitleCall() + expect(call).toBeDefined() + expect(call?.taskId).toBe(task.taskId) + expect(call?.userMessage).toBe("title trigger message") + expect(call?.provider).toBe(provider) + // The trigger is fire-and-forget: startTask resolved and the generation + // ran to completion without startTask ever awaiting it beyond its own + // synchronous promise chain. + expect(resolved).toBe(true) + }) + + it("starts title generation only after the first message is stored", async () => { + const provider = createMockProvider() + const task = createTestTask(provider) + stubStartup(task) + + let releaseSay: () => void = () => {} + vi.spyOn(task, "say").mockImplementation( + () => + new Promise((resolve) => { + releaseSay = () => resolve(undefined) + }), + ) + + const started = getStartTaskAccess(task).startTask("stored first") + await Promise.resolve() + expect(vi.mocked(generateSessionTitle)).not.toHaveBeenCalled() + + releaseSay() + await started + + expect(lastTitleCall()?.userMessage).toBe("stored first") + }) + + it("does not wait for a pending title before starting the task loop", async () => { + const provider = createMockProvider() + const task = createTestTask(provider) + stubStartup(task) + + let releaseTitle: () => void = () => {} + vi.mocked(generateSessionTitle).mockImplementation( + () => + new Promise((resolve) => { + releaseTitle = resolve + }), + ) + + // startTask resolving while the title promise is still pending proves it is not awaited. + await getStartTaskAccess(task).startTask("slow title message") + + expect(getStartTaskAccess(task).initiateTaskLoop).toHaveBeenCalled() + expect(lastTitleCall()?.userMessage).toBe("slow title message") + releaseTitle() + }) + + it.each([ + ["Production", 1, true], + ["Development", 2, true], + ["Test", 3, false], + ])("in %s extension mode fires title generation: %s", async (_name, extensionMode, shouldFire) => { + const provider = createMockProvider() + // The double only carries the context members the trigger reads. + Object.assign(provider.context, { extensionMode }) + const task = createTestTask(provider) + stubStartup(task) + + await getStartTaskAccess(task).startTask("mode gated message") + + expect(vi.mocked(generateSessionTitle).mock.calls.length > 0).toBe(shouldFire) + }) + + it("does not fire for a subtask", async () => { + const provider = createMockProvider() + const parentTask = createTestTask(provider) + const childTask = createTestTask(provider, parentTask) + stubStartup(childTask) + + await getStartTaskAccess(childTask).startTask("child task message") + + expect(vi.mocked(generateSessionTitle)).not.toHaveBeenCalled() + }) + + it("does not fire when the root message is empty", async () => { + const provider = createMockProvider() + const task = createTestTask(provider) + stubStartup(task) + + await getStartTaskAccess(task).startTask(undefined) + + expect(vi.mocked(generateSessionTitle)).not.toHaveBeenCalled() + }) + + it("routes a history task away from startTask, so resumes never trigger", async () => { + const provider = createMockProvider() + const historyItem: HistoryItem = { + id: "resume-title-task", + number: 3, + task: "historical task", + ts: Date.now() - 60_000, + tokensIn: 10, + tokensOut: 5, + totalCost: 0.01, + } + // A history item with no task text marks the task as a history task, so + // run() routes through resumeTaskFromHistory and never reaches the + // startTask trigger (Task.ts `_isHistoryTask` routing). + const task = createTestTask(provider, undefined, historyItem) + const taskAccess = getStartTaskAccess(task) + const startTaskSpy = vi.spyOn(taskAccess, "startTask") + vi.spyOn(taskAccess, "resumeTaskFromHistory").mockResolvedValue(undefined) + + await task.run() + + expect(startTaskSpy).not.toHaveBeenCalled() + expect(vi.mocked(generateSessionTitle)).not.toHaveBeenCalled() + }) + + it("keeps the run alive when generation rejects", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + const provider = createMockProvider() + const task = createTestTask(provider) + stubStartup(task) + vi.mocked(generateSessionTitle).mockRejectedValue(new Error("boom")) + + await expect(getStartTaskAccess(task).startTask("survives rejection")).resolves.toBeUndefined() + expect(consoleErrorSpy).toHaveBeenCalledWith( + expect.stringContaining("[Task#startTask] session title generation failed: Error: boom"), + ) + consoleErrorSpy.mockRestore() + }) +}) diff --git a/src/core/task/__tests__/session-title.spec.ts b/src/core/task/__tests__/session-title.spec.ts new file mode 100644 index 0000000000..4d5e43ff57 --- /dev/null +++ b/src/core/task/__tests__/session-title.spec.ts @@ -0,0 +1,584 @@ +// npx vitest run core/task/__tests__/session-title.spec.ts + +import { providerIdentifiers, type ProviderSettings } from "@roo-code/types" + +import type { CompletePromptOptions } from "../../../api" +import type { ClineProvider } from "../../webview/ClineProvider" +import { + FALLBACK_TITLE_MAX_CHARS, + SESSION_TITLE_TIMEOUT_MS, + fallbackTitle, + generateSessionTitle, + cancelSessionTitle, + sanitizeUserMessageForTitle, + extractLlmTitle, + type SessionTitleApplyResult, +} from "../session-title" + +type CompletePromptCall = { + prompt: string + apiConfiguration: ProviderSettings + options: CompletePromptOptions | undefined +} + +type ApplyTitleCall = { taskId: string; title: string } + +/** + * The orchestration only needs the one provider method it calls. The double is + * typed structurally and cast once at the construction site (AGENTS.md): the + * test never reads other provider members. + */ +function createTitleProvider( + overrides: { + applySessionTitle?: (taskId: string, title: string) => Promise + } = {}, +): ClineProvider { + const provider = { + applySessionTitle: overrides.applySessionTitle ?? (async () => "applied" as const), + } + return provider as unknown as ClineProvider +} + +function createCompletePromptRecorder(impl: (call: CompletePromptCall) => Promise): { + fn: (config: ProviderSettings, prompt: string, options?: CompletePromptOptions) => Promise + calls: CompletePromptCall[] +} { + const calls: CompletePromptCall[] = [] + const fn = (apiConfiguration: ProviderSettings, prompt: string, options: CompletePromptCall["options"]) => { + const call = { prompt, apiConfiguration, options } + calls.push(call) + return impl(call) + } + return { fn, calls } +} + +const testApiConfiguration: ProviderSettings = { + apiProvider: providerIdentifiers.anthropic, + apiModelId: "claude-3-5-sonnet-20241022", + apiKey: "test-api-key", +} + +describe("sanitizeUserMessageForTitle", () => { + it.each([ + ["plain message", "plain message"], + [" padded ", "padded"], + ["before hidden after", "before after"], + ["a\nbreal", "real"], + ["start unterminated", "start"], + ["only system", ""], + ["axbyc", "abc"], + ["keep stray close", "keep stray close"], + ])("sanitizes %j", (input, expected) => { + expect(sanitizeUserMessageForTitle(input)).toBe(expected) + }) + + it("handles many unterminated openings in linear time", () => { + const input = "".repeat(50_000) + const startedAt = performance.now() + + expect(sanitizeUserMessageForTitle(input)).toBe("") + expect(performance.now() - startedAt).toBeLessThan(500) + }) +}) + +describe("extractLlmTitle", () => { + it.each([ + ["Fix login bug in auth flow", "Fix login bug in auth flow"], + ["Fix Title: parsing in task history", "Fix Title: parsing in task history"], + ['Fix "quoted" search terms', 'Fix "quoted" search terms'], + ['"Leading quote only', "Leading quote only"], + ['Trailing quote only"', "Trailing quote only"], + ['"""Multiple quoted"""', "Multiple quoted"], + ["' Spaced quotes '", "Spaced quotes"], + ["Title: Fix login bug", "Fix login bug"], + ["title: lowercase prefix", "lowercase prefix"], + ['"Quoted title"', "Quoted title"], + ["'Single quoted'", "Single quoted"], + ["“Curly quoted”", "Curly quoted"], + [" Title: Fix login bug", "Fix login bug"], + [" Fix login bug\nSecond line", "Fix login bug"], + ["First line\nSecond line", "First line"], + ['Title: "Nested cleanup"', "Nested cleanup"], + ])("extracts %j", (input, expected) => { + expect(extractLlmTitle(input)).toBe(expected) + }) + + it.each([[""], [" "], ["\n\n"], ['""'], ["Title:"]])("returns undefined for %j", (input) => { + expect(extractLlmTitle(input)).toBeUndefined() + }) +}) + +describe("fallbackTitle", () => { + it.each([ + ["\n\nsecond\nthird", "second third"], + ["Fix\nlogin bug", "Fix login bug"], + [" leading space ", "leading space"], + ])("collapses whitespace in %j", (input, expected) => { + expect(fallbackTitle(input)).toBe(expected) + }) + + it("cuts at the character limit before collapsing whitespace", () => { + const input = "a".repeat(FALLBACK_TITLE_MAX_CHARS - 2) + " ".repeat(5) + "bcd" + expect(fallbackTitle(input)).toBe("a".repeat(FALLBACK_TITLE_MAX_CHARS - 2)) + expect(fallbackTitle("ab" + " ".repeat(60) + "cd")).toBe("ab") + }) + + it("slices the first 42 characters before trimming or collapsing leading whitespace", () => { + const input = " ".repeat(40) + "ABCDEF" + expect(fallbackTitle(input)).toBe("AB") + }) + + it("caps the fallback at the character limit", () => { + const long = "a".repeat(FALLBACK_TITLE_MAX_CHARS + 10) + expect(fallbackTitle(long)).toBe("a".repeat(FALLBACK_TITLE_MAX_CHARS)) + expect(fallbackTitle(long)).toHaveLength(FALLBACK_TITLE_MAX_CHARS) + }) + + it.each([[""], [" \n \t"]])("returns undefined for %j", (input) => { + expect(fallbackTitle(input)).toBeUndefined() + }) +}) + +describe("generateSessionTitle", () => { + it("applies the cleaned LLM title", async () => { + const recorder = createCompletePromptRecorder(async () => 'Title: "Fix Login Bug"\nextra') + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-1", + userMessage: "please fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(recorder.calls).toHaveLength(1) + expect(recorder.calls[0].prompt).toContain("please fix the login bug") + expect(recorder.calls[0].prompt).toContain("Create a descriptive 5-10 word title") + expect(recorder.calls[0].options?.taskId).toBe("task-1") + expect(recorder.calls[0].options?.timeoutMs).toBe(SESSION_TITLE_TIMEOUT_MS) + expect(recorder.calls[0].apiConfiguration).toBe(testApiConfiguration) + expect(applied).toEqual([{ taskId: "task-1", title: "Fix Login Bug" }]) + }) + + it("falls back to the start of the message when the LLM fails", async () => { + const recorder = createCompletePromptRecorder(async () => { + throw new Error("api down") + }) + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-2", + userMessage: "first line of the request\nsecond line", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(applied).toEqual([{ taskId: "task-2", title: "first line of the request second line" }]) + }) + + it("strips system-reminder blocks before deriving the fallback title", async () => { + const recorder = createCompletePromptRecorder(async () => { + throw new Error("api down") + }) + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-sanitized-fallback", + userMessage: "internal instructionsPlease fix the login button", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(applied).toEqual([{ taskId: "task-sanitized-fallback", title: "Please fix the login button" }]) + }) + + it.each([[""], [" "], ['""'], ["Title:"]])( + "falls back to the message start when the LLM returns the unusable response %j", + async (response) => { + const recorder = createCompletePromptRecorder(async () => response) + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-unusable", + userMessage: "Short first line\nwith more text after it that goes on and on", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(applied).toEqual([{ taskId: "task-unusable", title: "Short first line with more text after it t" }]) + }, + ) + + it("falls back when the LLM times out", async () => { + vi.useFakeTimers() + try { + const recorder = createCompletePromptRecorder(() => new Promise(() => {})) + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + applied.push({ taskId, title }) + return "applied" + }, + }) + + const generation = generateSessionTitle({ + provider, + taskId: "task-timeout", + userMessage: "timeout message body", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + // Drain microtasks so the raced completion settles first. + await vi.advanceTimersByTimeAsync(SESSION_TITLE_TIMEOUT_MS + 1) + await generation + + expect(recorder.calls[0].options?.abortSignal?.aborted).toBe(true) + expect(applied).toEqual([{ taskId: "task-timeout", title: "timeout message body" }]) + } finally { + vi.useRealTimers() + } + }) + + it.each([[""], [" "], ["only system content"]])( + "skips everything for sanitized-empty message %j", + async (userMessage) => { + const recorder = createCompletePromptRecorder(async () => "should not be called") + const provider = createTitleProvider() + + await generateSessionTitle({ + provider, + taskId: "task-skip", + userMessage, + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(recorder.calls).toHaveLength(0) + }, + ) + + it("stops immediately on conflict so a manual title always wins", async () => { + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let applyCalls = 0 + const provider = createTitleProvider({ + applySessionTitle: async () => { + applyCalls++ + return "conflict" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-conflict", + userMessage: "some root message", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(applyCalls).toBe(1) + }) + + it("retries a failing write up to the retry cap, then gives up", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let applyCalls = 0 + const provider = createTitleProvider({ + applySessionTitle: async () => { + applyCalls++ + throw new Error("disk on fire") + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-retry", + userMessage: "some root message", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + // 4 attempts for LLM title + 4 attempts for fallback title = 8 total attempts. + expect(applyCalls).toBe(8) + expect(consoleErrorSpy).toHaveBeenCalledWith( + expect.stringContaining("[session-title] giving up on task task-retry after 4 write attempts:"), + expect.any(Error), + ) + consoleErrorSpy.mockRestore() + }) + + it("attempts the fallback title when LLM title write retries fail", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + const recorder = createCompletePromptRecorder(async () => "A".repeat(500)) + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + if (title.length > 50) { + throw new Error("payload too large") + } + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-llm-write-fail", + userMessage: "Fix login issue with auth flow", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(applied).toEqual([{ taskId: "task-llm-write-fail", title: "Fix login issue with auth flow" }]) + consoleErrorSpy.mockRestore() + }) + + it("recovers when a retried write succeeds", async () => { + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let applyCalls = 0 + const provider = createTitleProvider({ + applySessionTitle: async () => { + applyCalls++ + if (applyCalls < 3) { + throw new Error("transient") + } + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-retry-recover", + userMessage: "some root message", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(applyCalls).toBe(3) + }) + + it("runs only one generation per task at a time", async () => { + let release: (() => void) | undefined + const gate = new Promise((resolve) => { + release = resolve + }) + const recorder = createCompletePromptRecorder(async () => { + await gate + return "First Title" + }) + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + applied.push({ taskId, title }) + return "applied" + }, + }) + + const first = generateSessionTitle({ + provider, + taskId: "task-dedup", + userMessage: "root message", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + const second = generateSessionTitle({ + provider, + taskId: "task-dedup", + userMessage: "root message", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + release!() + await Promise.all([first, second]) + + expect(applied).toHaveLength(1) + expect(applied[0].title).toBe("First Title") + }) + + it("aborts in-flight generation and does not apply title when cancelled", async () => { + let abortSignal: AbortSignal | undefined + let resolveCompletion!: (val: string) => void + const completionPromise = new Promise((resolve) => { + resolveCompletion = resolve + }) + const recorder = createCompletePromptRecorder(async ({ options }) => { + abortSignal = options?.abortSignal + return completionPromise + }) + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + applied.push({ taskId, title }) + return "applied" + }, + }) + + const generation = generateSessionTitle({ + provider, + taskId: "task-cancel", + userMessage: "User prompt to name", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(recorder.calls).toHaveLength(1) + expect(abortSignal?.aborted).toBe(false) + + cancelSessionTitle("task-cancel") + + expect(abortSignal?.aborted).toBe(true) + resolveCompletion("Generated title") + await generation + + expect(applied).toHaveLength(0) + }) + + it("skips LLM completion and uses fallback title when rateLimitSeconds is set", async () => { + const recorder = createCompletePromptRecorder(async () => "LLM Title") + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-ratelimit", + userMessage: "Fix urgent production issue in login flow", + apiConfiguration: { ...testApiConfiguration, rateLimitSeconds: 5 }, + completePrompt: recorder.fn, + }) + + expect(recorder.calls).toHaveLength(0) + expect(applied).toEqual([{ taskId: "task-ratelimit", title: "Fix urgent production issue in login flow" }]) + }) + + it("settles promptly via abortPromise when cancelled without completion settling", async () => { + const recorder = createCompletePromptRecorder(() => new Promise(() => {})) + const applySpy = vi.fn().mockResolvedValue("applied" as const) + const provider = createTitleProvider({ applySessionTitle: applySpy }) + + const generation = generateSessionTitle({ + provider, + taskId: "task-abort-prompt", + userMessage: "Fix login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + await Promise.resolve() + cancelSessionTitle("task-abort-prompt") + await generation + + expect(applySpy).not.toHaveBeenCalled() + }) + + it("stops silently when cancelled during write retries", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let applyCalls = 0 + const provider = createTitleProvider({ + applySessionTitle: async () => { + applyCalls++ + cancelSessionTitle("task-cancel-retry") + throw new Error("transient failure") + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-cancel-retry", + userMessage: "some root message", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(applyCalls).toBe(1) + expect(consoleErrorSpy).not.toHaveBeenCalled() + consoleErrorSpy.mockRestore() + }) + + it("stops fallback retry immediately if the generation is cancelled", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + const recorder = createCompletePromptRecorder(async () => { + throw new Error("LLM failed") + }) + let applyCalls = 0 + const provider = createTitleProvider({ + applySessionTitle: async () => { + applyCalls++ + cancelSessionTitle("task-cancel-fallback-retry") + throw new Error("transient failure") + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-cancel-fallback-retry", + userMessage: "Fix login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(applyCalls).toBe(1) + expect(consoleErrorSpy).toHaveBeenCalledTimes(1) + expect(consoleErrorSpy).toHaveBeenCalledWith( + expect.stringContaining("LLM call failed for task task-cancel-fallback-retry"), + expect.anything(), + ) + consoleErrorSpy.mockRestore() + }) + + it("safely ignores cancellation for an untracked task", () => { + expect(() => cancelSessionTitle("non-existent-task")).not.toThrow() + }) + + it("never rejects even when everything fails", async () => { + const provider = createTitleProvider({ + applySessionTitle: async () => { + throw new Error("write failed") + }, + }) + + await expect( + generateSessionTitle({ + provider, + taskId: "task-quiet", + userMessage: "root message", + apiConfiguration: testApiConfiguration, + completePrompt: async () => { + throw new Error("llm failed") + }, + }), + ).resolves.toBeUndefined() + }) +}) diff --git a/src/core/task/session-title.ts b/src/core/task/session-title.ts new file mode 100644 index 0000000000..612746c0ca --- /dev/null +++ b/src/core/task/session-title.ts @@ -0,0 +1,233 @@ +import type { ProviderSettings } from "@roo-code/types" + +import type { ClineProvider } from "../webview/ClineProvider" +import { singleCompletionHandler } from "../../utils/single-completion-handler" + +/** Hard bound on the single-shot title request (15 seconds). */ +export const SESSION_TITLE_TIMEOUT_MS = 15_000 +/** Character cap for the fallback title taken from the user message (42 chars). */ +export const FALLBACK_TITLE_MAX_CHARS = 42 +/** Write attempts given up on after this many retries (3 retries). */ +const TITLE_WRITE_MAX_RETRIES = 3 + +// Stryker disable next-line StringLiteral: diagnostic label, does not alter control flow +const TITLE_PROMPT_PREFIX = `Create a descriptive 5-10 word title for the user message below. Use the user's language. +Output only the title. + +` + +/** Outcome of `ClineProvider.applySessionTitle`. */ +export type SessionTitleApplyResult = "applied" | "conflict" + +/** Removes system-injected `` blocks and trims. */ +export function sanitizeUserMessageForTitle(raw: string): string { + const open = "" + const close = "" + let result = "" + let cursor = 0 + // Stryker disable next-line EqualityOperator: cursor reaches raw.length on trailing text; <= is equivalent as indexOf returns -1 at length + while (cursor < raw.length) { + const start = raw.indexOf(open, cursor) + if (start === -1) { + result += raw.slice(cursor) + break + } + result += raw.slice(cursor, start) + // Stryker disable next-line ArithmeticOperator: equivalent mutant, (18 chars) cannot start in the [start-17, start+17] window around (17 chars) + const end = raw.indexOf(close, start + open.length) + // An unterminated block would otherwise leak the rest of the message as user text; drop it. + if (end === -1) { + break + } + cursor = end + close.length + } + return result.trim() +} + +/** First line, `Title:` prefix removed, surrounding quotes removed. */ +export function extractLlmTitle(raw: string): string | undefined { + const firstLine = raw.split("\n")[0].trim() + const withoutPrefix = firstLine.replace(/^title:\s*/i, "") + const withoutQuotes = withoutPrefix.replace(/^["'“”‘’]+|["'“”‘’]+$/g, "") + const title = withoutQuotes.trim() + return title.length > 0 ? title : undefined +} + +/** First 42 characters of the message, then whitespace collapsed. */ +export function fallbackTitle(rawMessage: string): string | undefined { + const head = rawMessage.slice(0, FALLBACK_TITLE_MAX_CHARS) + return head.replace(/\s+/g, " ").trim() || undefined +} + +type CompletePromptFn = typeof singleCompletionHandler + +export interface GenerateSessionTitleOptions { + provider: ClineProvider + taskId: string + /** The raw root user message that started the run. */ + userMessage: string + /** The API configuration the session itself uses (same model reuse). */ + apiConfiguration: ProviderSettings + /** Injectable for tests; defaults to the shared single-shot completion helper. */ + completePrompt?: CompletePromptFn +} + +interface InFlightTitle { + controller: AbortController + userCancelled: boolean +} + +/** In-flight title generation abort controllers by task ID. */ +const inFlightTitles = new Map() + +/** + * Cancels any pending or in-flight title generation for a task. + * Called when a task is manually renamed, cleared, or deleted so an auto-generated title + * never overwrites the user's manual intent. + */ +export function cancelSessionTitle(taskId: string): void { + const entry = inFlightTitles.get(taskId) + if (entry) { + entry.userCancelled = true + entry.controller.abort() + } +} + +/** + * Generates and applies a session title from the root user message. + * + * Fire-and-forget by contract: never rejects, never blocks a run, and every + * failure degrades quietly (fallback title, or no write at all). A manual + * rename always wins over the generated title. + */ +export async function generateSessionTitle(options: GenerateSessionTitleOptions): Promise { + const { provider, taskId, userMessage, apiConfiguration, completePrompt = singleCompletionHandler } = options + + const sanitized = sanitizeUserMessageForTitle(userMessage) + if (!sanitized) { + return + } + + if (inFlightTitles.has(taskId)) { + return + } + const inFlight: InFlightTitle = { + controller: new AbortController(), + userCancelled: false, + } + inFlightTitles.set(taskId, inFlight) + + try { + let title: string | undefined + const rateLimitSeconds = apiConfiguration.rateLimitSeconds ?? 0 + if (rateLimitSeconds <= 0 && !inFlight.userCancelled) { + try { + title = await completeTitleViaLlm( + completePrompt, + apiConfiguration, + taskId, + sanitized, + inFlight.controller, + ) + } catch (error) { + // LLM failure, cancellation, or timeout falls back to the message's first line. + // Stryker disable next-line BlockStatement,BooleanLiteral,ConditionalExpression,StringLiteral: diagnostic logging only, does not alter control flow + if (!inFlight.userCancelled) { + console.error(`[session-title] LLM call failed for task ${taskId}:`, error) + } + } + } + let applied = false + if (title) { + applied = await applyTitleWithRetry(provider, taskId, title, () => inFlight.userCancelled) + } + if (!applied && !inFlight.userCancelled) { + const fallback = fallbackTitle(sanitized) + // Stryker disable next-line ConditionalExpression,LogicalOperator: defense-in-depth, fallback is always truthy for non-empty sanitized messages and distinct from failed title + if (fallback && fallback !== title) { + await applyTitleWithRetry(provider, taskId, fallback, () => inFlight.userCancelled) + } + } + } catch (error) { + // Best-effort by design: a failed title must never surface in the chat. + // Stryker disable next-line BlockStatement,BooleanLiteral,ConditionalExpression: defense-in-depth, applyTitleWithRetry handles all expected errors internally + if (!inFlight.userCancelled) { + // Stryker disable next-line CallExpression,StringLiteral: diagnostic error logging only + console.error(`[session-title] generation failed for task ${taskId}:`, error) + } + } finally { + inFlightTitles.delete(taskId) + } +} + +async function completeTitleViaLlm( + completePrompt: CompletePromptFn, + apiConfiguration: ProviderSettings, + taskId: string, + sanitized: string, + controller: AbortController, +): Promise { + let timeoutId: ReturnType | undefined + try { + const completion = completePrompt(apiConfiguration, TITLE_PROMPT_PREFIX + sanitized, { + taskId, + timeoutMs: SESSION_TITLE_TIMEOUT_MS, + abortSignal: controller.signal, + }) + const timeout = new Promise((_, reject) => { + timeoutId = setTimeout(() => { + controller.abort() + // Stryker disable next-line CallExpression,StringLiteral: internal timeout error rejection, caught internally + reject(new Error(`[session-title] timed out after ${SESSION_TITLE_TIMEOUT_MS}ms`)) + }, SESSION_TITLE_TIMEOUT_MS) + }) + // Stryker disable next-line BlockStatement: executor registers abort listener + const abortPromise = new Promise((_, reject) => { + // Stryker disable next-line BlockStatement,BooleanLiteral,ConditionalExpression: defense-in-depth, signal is fresh and rarely pre-aborted + if (controller.signal.aborted) { + // Stryker disable next-line CallExpression,StringLiteral: internal abort error rejection, caught internally + reject(new Error("[session-title] aborted")) + return + } + // Stryker disable next-line ArrowFunction,CallExpression,ObjectLiteral,StringLiteral: equivalent mutant, abort event fires at most once per controller + controller.signal.addEventListener("abort", () => reject(new Error("[session-title] aborted")), { + // Stryker disable next-line BooleanLiteral: equivalent mutant, abort controller aborts at most once + once: true, + }) + }) + const raw = await Promise.race([completion, timeout, abortPromise]) + return extractLlmTitle(raw) + } finally { + // Stryker disable next-line CallExpression: timer cleanup does not affect the resolved return value + clearTimeout(timeoutId) + } +} + +async function applyTitleWithRetry( + provider: ClineProvider, + taskId: string, + title: string, + isCancelled: () => boolean, +): Promise { + for (let attempt = 0; ; attempt++) { + // Stryker disable next-line BooleanLiteral,ConditionalExpression: defense-in-depth, caller checks cancellation before initiating write + if (isCancelled()) { + // Stryker disable next-line BooleanLiteral: return value indicating write was not applied + return false + } + try { + await provider.applySessionTitle(taskId, title) + return true + } catch (error) { + if (isCancelled() || attempt >= TITLE_WRITE_MAX_RETRIES) { + if (!isCancelled()) { + console.error( + `[session-title] giving up on task ${taskId} after ${attempt + 1} write attempts:`, + error, + ) + } + return false + } + } + } +} diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 8fff156425..46e0e35ce1 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -76,6 +76,7 @@ import { providerIdentifiers, } from "@roo-code/types" import { RateLimitClock, createRateLimitClock } from "../task/RateLimitClock" +import { cancelSessionTitle, type SessionTitleApplyResult } from "../task/session-title" import { TaskRegistry } from "../task/TaskRegistry" import { TaskScheduler } from "../task/TaskScheduler" import { @@ -2809,6 +2810,7 @@ export class ClineProvider * clears it so the history falls back to the original prompt. */ async renameTaskWithId(id: string, title: string | undefined) { + cancelSessionTitle(id) if (!this.taskHistoryStore.get(id)) { this.log(`[renameTaskWithId] task ${id} not found in history`) return @@ -2831,6 +2833,38 @@ export class ClineProvider } } + /* + * Applies an auto-generated session title, but only while the item still + * carries its default name. The title check runs inside the store's + * lock and re-reads the title on disk, so a manual rename that lands while a title is + * generating always wins and is never overwritten. Returns "conflict" + * when nothing was applied (an existing title won, or the title was + * blank); throws on transient write failures so the caller can retry. + */ + async applySessionTitle(id: string, title: string): Promise { + const trimmed = title.trim() + if (!trimmed) { + return "conflict" + } + + const result = await this.taskHistoryStore.setTitleIfUnset(id, trimmed) + if (result === "conflict") { + return "conflict" + } + + this.recentTasksCache = undefined + + if (this.isViewLaunched) { + const updatedItem = this.taskHistoryStore.get(id) + // Stryker disable next-line ConditionalExpression: equivalent mutant, setTitleIfUnset guarantees the item is in the cache + if (updatedItem) { + await this.postMessageToWebview({ type: "taskHistoryItemUpdated", taskHistoryItem: updatedItem }) + } + } + + return "applied" + } + /* Condenses a task's message history to use fewer tokens. */ async condenseTaskContext(taskId: string) { const task = this.taskRegistry.getById(taskId) @@ -2844,6 +2878,7 @@ export class ClineProvider // this function deletes a task from task history, and deletes its checkpoints and delete the task folder // If the task has subtasks (childIds), they will also be deleted recursively async deleteTaskWithId(id: string, cascadeSubtasks: boolean = true) { + cancelSessionTitle(id) try { await this.getTaskWithId(id) @@ -3637,9 +3672,10 @@ export class ClineProvider * Updates a task in the task history and optionally broadcasts the updated history to the webview. * Now delegates to TaskHistoryStore for per-task file persistence. * - * This writer never owns the `title` field: `title` is written only by the rename path - * (`renameTaskWithId` via `taskHistoryStore.atomicReadAndUpdate`). Callers here build items - * from snapshots read before taking the store lock, so any title they carry can be stale. + * This writer never owns the `title` field: `title` is written only by the title-owning paths + * (`renameTaskWithId` via `atomicReadAndUpdate` and `applySessionTitle` via `setTitleIfUnset`). + * Callers here build items from snapshots read before taking the store lock, so any title they + * carry can be stale. * * @param item The history item to update or add * @param options.broadcast Whether to broadcast the updated history to the webview (default: true) diff --git a/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts b/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts index 7b6223e2e2..a0910800e0 100644 --- a/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts +++ b/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts @@ -11,6 +11,8 @@ import { ContextProxy } from "../../config/ContextProxy" import { Task } from "../../task/Task" import { ProfileValidator } from "../../../shared/ProfileValidator" import { ClineProvider } from "../ClineProvider" +import { generateSessionTitle } from "../../task/session-title" +import * as sessionTitleModule from "../../task/session-title" // Mock setup vi.mock("p-wait-for", () => ({ @@ -653,6 +655,35 @@ describe("ClineProvider Task History Synchronization", () => { expect(getStored("task-race")).toMatchObject({ title: "Renamed", tokensIn: 500 }) }) + it("cancels in-flight session title generation and preserves a manual title clear", async () => { + await provider.updateTaskHistory(createHistoryItem({ id: "task-manual-clear", task: "Original prompt" }), { + broadcast: false, + }) + let resolveCompletion!: (title: string) => void + const completionPromise = new Promise((resolve) => { + resolveCompletion = resolve + }) + const generation = generateSessionTitle({ + provider, + taskId: "task-manual-clear", + userMessage: "Original prompt", + apiConfiguration: { + apiProvider: providerIdentifiers.anthropic, + apiModelId: "claude-3-5-sonnet-20241022", + apiKey: "test-api-key", + }, + completePrompt: async () => completionPromise, + }) + + await provider.renameTaskWithId("task-manual-clear", "") + expect(getStored("task-manual-clear")?.title).toBeUndefined() + + resolveCompletion("Generated Title") + await generation + + expect(getStored("task-manual-clear")?.title).toBeUndefined() + }) + it("does not write a title through a full-item update", async () => { await provider.updateTaskHistory( createHistoryItem({ id: "task-title-owner", task: "Original prompt", title: "Stale title" }), @@ -690,6 +721,103 @@ describe("ClineProvider Task History Synchronization", () => { }) }) + describe("applySessionTitle", () => { + const getStored = (id: string) => provider.taskHistoryStore.get(id) + + it("stores a trimmed title and broadcasts the updated item", async () => { + await provider.resolveWebviewView(mockWebviewView) + provider.isViewLaunched = true + await provider.updateTaskHistory(createHistoryItem({ id: "task-auto-title", task: "Original prompt" })) + mockPostMessage.mockClear() + + const result = await provider.applySessionTitle("task-auto-title", " Auto Title ") + + expect(result).toBe("applied") + expect(getStored("task-auto-title")).toMatchObject({ title: "Auto Title", task: "Original prompt" }) + const updates = findCallsByType(mockPostMessage.mock.calls, "taskHistoryItemUpdated") + expect(updates.at(-1)?.[0].taskHistoryItem).toMatchObject({ id: "task-auto-title", title: "Auto Title" }) + }) + + it("returns conflict and keeps the existing title when one is already set", async () => { + await provider.updateTaskHistory(createHistoryItem({ id: "task-auto-conflict", task: "Original prompt" }), { + broadcast: false, + }) + await provider.renameTaskWithId("task-auto-conflict", "Manual title") + + const result = await provider.applySessionTitle("task-auto-conflict", "Auto Title") + + expect(result).toBe("conflict") + expect(getStored("task-auto-conflict")).toMatchObject({ title: "Manual title" }) + }) + + it.each([ + ["an empty string", ""], + ["whitespace", " "], + ])("returns conflict without writing for %s", async (_label, title) => { + await provider.updateTaskHistory(createHistoryItem({ id: "task-auto-blank", task: "Original prompt" }), { + broadcast: false, + }) + + const result = await provider.applySessionTitle("task-auto-blank", title) + + expect(result).toBe("conflict") + expect(getStored("task-auto-blank")?.title).toBeUndefined() + }) + + it("throws for an unknown task so the caller can give up", async () => { + await expect(provider.applySessionTitle("missing-task", "Auto Title")).rejects.toThrow("not found in cache") + }) + + it("keeps the auto title when the task later saves metadata without one", async () => { + const item = createHistoryItem({ id: "task-auto-keep", task: "Original prompt", status: "active" }) + await provider.updateTaskHistory(item, { broadcast: false }) + await provider.applySessionTitle("task-auto-keep", "Auto Title") + + // Mirrors Task.saveClineMessages(), whose taskMetadata() result never carries a title. + await provider.updateTaskHistory({ ...item, tokensIn: 500 }, { broadcast: false }) + + expect(getStored("task-auto-keep")).toMatchObject({ title: "Auto Title", tokensIn: 500 }) + }) + + it("keeps a concurrent task save that is queued before the auto title", async () => { + const item = createHistoryItem({ id: "task-auto-race", task: "Original prompt", status: "active" }) + await provider.updateTaskHistory(item, { broadcast: false }) + + // Start the save without awaiting it, so the auto title runs while the save is queued. + const save = provider.updateTaskHistory({ ...item, tokensIn: 500 }, { broadcast: false }) + const autoTitle = provider.applySessionTitle("task-auto-race", "Auto Title") + await Promise.all([save, autoTitle]) + + expect(getStored("task-auto-race")).toMatchObject({ title: "Auto Title", tokensIn: 500 }) + }) + + it("does not broadcast when the view is not launched", async () => { + await provider.resolveWebviewView(mockWebviewView) + provider.isViewLaunched = false + await provider.updateTaskHistory(createHistoryItem({ id: "task-auto-silent", task: "Original prompt" }), { + broadcast: false, + }) + mockPostMessage.mockClear() + + const result = await provider.applySessionTitle("task-auto-silent", "Auto Title") + + expect(result).toBe("applied") + expect(getStored("task-auto-silent")).toMatchObject({ title: "Auto Title" }) + expect(findCallsByType(mockPostMessage.mock.calls, "taskHistoryItemUpdated")).toHaveLength(0) + }) + + it("cancels pending session title generation when deleting a task", async () => { + const cancelSpy = vi.spyOn(sessionTitleModule, "cancelSessionTitle") + try { + await provider.deleteTaskWithId("task-cancel-on-delete") + } catch { + // Task doesn't exist on disk, delete will catch/proceed + } + expect(cancelSpy).toHaveBeenCalledWith("task-cancel-on-delete") + cancelSpy.mockRestore() + }) + }) + describe("broadcastTaskHistoryUpdate", () => { it("sends taskHistoryUpdated message with sorted history", async () => { await provider.resolveWebviewView(mockWebviewView) From 344ef96af9f909c17287f70860c65bd6cc9e8b1a Mon Sep 17 00:00:00 2001 From: daewoongoh Date: Sat, 7 Nov 2026 23:01:43 +0900 Subject: [PATCH 2/7] fix(task): address CodeRabbit review feedback for session title generation --- packages/types/src/history.ts | 1 + .../base-openai-compatible-provider.spec.ts | 20 ++ src/api/providers/__tests__/kenari.spec.ts | 19 ++ src/api/providers/__tests__/openai.spec.ts | 19 ++ .../providers/__tests__/openrouter.spec.ts | 35 +++ src/api/providers/__tests__/poe.spec.ts | 23 ++ src/api/providers/__tests__/requesty.spec.ts | 19 ++ src/api/providers/__tests__/unbound.spec.ts | 23 ++ .../__tests__/vercel-ai-gateway.spec.ts | 17 ++ .../providers/__tests__/zoo-gateway.spec.ts | 17 ++ .../base-openai-compatible-provider.ts | 17 +- src/api/providers/kenari.ts | 20 +- src/api/providers/openai.ts | 22 +- src/api/providers/openrouter.ts | 34 ++- src/api/providers/poe.ts | 8 +- src/api/providers/requesty.ts | 19 +- src/api/providers/unbound.ts | 19 +- src/api/providers/utils/abort-signal.ts | 12 + src/api/providers/vercel-ai-gateway.ts | 20 +- src/api/providers/zoo-gateway.ts | 20 +- src/core/task-persistence/TaskHistoryStore.ts | 51 +++- .../TaskHistoryStore.crossInstance.spec.ts | 53 +++- .../TaskHistoryStore.realConcurrency.spec.ts | 65 ++++ src/core/task/Task.ts | 6 +- .../task/__tests__/Task.session-title.spec.ts | 57 +++- src/core/task/__tests__/session-title.spec.ts | 278 ++++++++++++++++++ src/core/task/session-title.ts | 130 +++++++- src/core/webview/ClineProvider.ts | 49 ++- .../ClineProvider.taskHistory.spec.ts | 99 ++++++- 29 files changed, 1105 insertions(+), 67 deletions(-) diff --git a/packages/types/src/history.ts b/packages/types/src/history.ts index 131194212e..e381dc5991 100644 --- a/packages/types/src/history.ts +++ b/packages/types/src/history.ts @@ -34,6 +34,7 @@ export const historyItemSchema = z.object({ ts: z.number(), task: z.string(), title: z.string().optional(), // User-assigned display name; `task` keeps the original prompt + titleCleared: z.boolean().optional(), // Persisted manual clear: blocks an automatically generated title tokensIn: z.number(), tokensOut: z.number(), cacheWrites: z.number().optional(), diff --git a/src/api/providers/__tests__/base-openai-compatible-provider.spec.ts b/src/api/providers/__tests__/base-openai-compatible-provider.spec.ts index 1ffd63aa12..bbda1503d8 100644 --- a/src/api/providers/__tests__/base-openai-compatible-provider.spec.ts +++ b/src/api/providers/__tests__/base-openai-compatible-provider.spec.ts @@ -355,6 +355,26 @@ describe("BaseOpenAiCompatibleProvider", () => { expect(mockCreate.mock.calls[0][0].max_tokens).toBeUndefined() }) + it("forwards the abort signal and deadline to the completion request", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + const { signal } = new AbortController() + + await new TestOpenAiCompatibleProvider("test-api-key", {}).completePrompt("hello", { + abortSignal: signal, + timeoutMs: 5_000, + }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { signal, timeout: 5_000 }) + }) + + it("sends a single-argument request when no cancellation or deadline is given", async () => { + mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) + + await new TestOpenAiCompatibleProvider("test-api-key", {}).completePrompt("hello", { timeoutMs: 0 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything()) + }) + it("completePrompt sends the custom model id verbatim to the API", async () => { mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: "ok" } }] }) diff --git a/src/api/providers/__tests__/kenari.spec.ts b/src/api/providers/__tests__/kenari.spec.ts index 0ba360a069..cab9495a6d 100644 --- a/src/api/providers/__tests__/kenari.spec.ts +++ b/src/api/providers/__tests__/kenari.spec.ts @@ -345,6 +345,25 @@ describe("KenariHandler", () => { await expect(handler.completePrompt("ping")).rejects.toThrow("Kenari completion error: boom") }) + it("forwards the abort signal and deadline to the completion request", async () => { + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + const handler = new KenariHandler(mockOptions) + const { signal } = new AbortController() + + await handler.completePrompt("ping", { abortSignal: signal, timeoutMs: 5_000 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { signal, timeout: 5_000 }) + }) + + it("sends a single-argument request when no cancellation or deadline is given", async () => { + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + const handler = new KenariHandler(mockOptions) + + await handler.completePrompt("ping", { timeoutMs: Number.NaN }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything()) + }) + it("re-throws non-Error rejections unchanged", async () => { mockCreate.mockRejectedValue("string failure") const handler = new KenariHandler(mockOptions) diff --git a/src/api/providers/__tests__/openai.spec.ts b/src/api/providers/__tests__/openai.spec.ts index 7c45f22180..d5e18f0eec 100644 --- a/src/api/providers/__tests__/openai.spec.ts +++ b/src/api/providers/__tests__/openai.spec.ts @@ -942,6 +942,25 @@ describe("OpenAiHandler", () => { }) describe("completePrompt", () => { + it("forwards the abort signal and deadline to the completion request", async () => { + const { signal } = new AbortController() + + await handler.completePrompt("Test prompt", { abortSignal: signal, timeoutMs: 5_000 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { signal, timeout: 5_000 }) + }) + + it("does not start a request when the caller already aborted", async () => { + const controller = new AbortController() + controller.abort() + mockCreate.mockClear() + + await expect(handler.completePrompt("Test prompt", { abortSignal: controller.signal })).rejects.toThrow( + "This operation was aborted", + ) + expect(mockCreate).not.toHaveBeenCalled() + }) + it("should complete prompt successfully", async () => { const result = await handler.completePrompt("Test prompt") expect(result).toBe("Test response") diff --git a/src/api/providers/__tests__/openrouter.spec.ts b/src/api/providers/__tests__/openrouter.spec.ts index fec3e4a2e2..5057fc09cd 100644 --- a/src/api/providers/__tests__/openrouter.spec.ts +++ b/src/api/providers/__tests__/openrouter.spec.ts @@ -973,7 +973,42 @@ describe("OpenRouterHandler", () => { }) }) + // Installs a typed completion double without widening the client to `any`. + const installMockClient = (create: ReturnType) => { + ;(OpenAI as unknown as { prototype: { chat: unknown } }).prototype.chat = { + completions: { create }, + } + } + describe("completePrompt", () => { + it("forwards the abort signal and deadline to the completion request", async () => { + const handler = new OpenRouterHandler(mockOptions) + const { signal } = new AbortController() + const mockCreate = vitest.fn().mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + installMockClient(mockCreate) + + await handler.completePrompt("test prompt", { abortSignal: signal, timeoutMs: 5_000 }) + + expect(mockCreate).toHaveBeenCalledWith(expect.objectContaining({ model: mockOptions.openRouterModelId }), { + headers: { "x-anthropic-beta": "fine-grained-tool-streaming-2025-05-14" }, + signal, + timeout: 5_000, + }) + }) + + it("does not start a request when the caller already aborted", async () => { + const handler = new OpenRouterHandler(mockOptions) + const mockCreate = vitest.fn() + installMockClient(mockCreate) + const controller = new AbortController() + controller.abort() + + await expect( + handler.completePrompt("test prompt", { abortSignal: controller.signal }), + ).rejects.toMatchObject({ name: "AbortError" }) + expect(mockCreate).not.toHaveBeenCalled() + }) + it("returns correct response", async () => { const handler = new OpenRouterHandler(mockOptions) const mockResponse = { choices: [{ message: { content: "test completion" } }] } diff --git a/src/api/providers/__tests__/poe.spec.ts b/src/api/providers/__tests__/poe.spec.ts index 9d9698aa0f..31d62281a4 100644 --- a/src/api/providers/__tests__/poe.spec.ts +++ b/src/api/providers/__tests__/poe.spec.ts @@ -402,6 +402,29 @@ describe("PoeHandler", () => { }) describe("completePrompt", () => { + it("forwards the abort signal and deadline to the generation", async () => { + const handler = new PoeHandler({ poeApiKey: "key", apiModelId: "openai/gpt-4o" }) + mockGenerateText.mockResolvedValue({ text: "generated response" }) + const { signal } = new AbortController() + + await handler.completePrompt("complete this", { abortSignal: signal, timeoutMs: 5_000 }) + + const options = mockGenerateText.mock.calls[0][0] + // The deadline is merged into a per-request signal that still follows the caller's stop. + expect(options.abortSignal).not.toBe(signal) + expect(options.abortSignal?.aborted).toBe(false) + }) + + it("does not start a generation when the caller already aborted", async () => { + const handler = new PoeHandler({ poeApiKey: "key", apiModelId: "openai/gpt-4o" }) + mockGenerateText.mockClear() + const controller = new AbortController() + controller.abort() + + await expect(handler.completePrompt("complete this", { abortSignal: controller.signal })).rejects.toThrow() + expect(mockGenerateText).not.toHaveBeenCalled() + }) + it("returns generated text", async () => { const handler = new PoeHandler({ poeApiKey: "key", apiModelId: "openai/gpt-4o" }) diff --git a/src/api/providers/__tests__/requesty.spec.ts b/src/api/providers/__tests__/requesty.spec.ts index 346ea4f14c..97b27dd19a 100644 --- a/src/api/providers/__tests__/requesty.spec.ts +++ b/src/api/providers/__tests__/requesty.spec.ts @@ -634,6 +634,25 @@ describe("RequestyHandler", () => { }) }) + it("forwards the abort signal and deadline to the completion request", async () => { + const handler = new RequestyHandler(mockOptions) + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + const { signal } = new AbortController() + + await handler.completePrompt("test prompt", { abortSignal: signal, timeoutMs: 5_000 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { signal, timeout: 5_000 }) + }) + + it("sends a single-argument request when no cancellation or deadline is given", async () => { + const handler = new RequestyHandler(mockOptions) + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + + await handler.completePrompt("test prompt", { timeoutMs: 0 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything()) + }) + it("omits temperature for Claude Fable 5 in completePrompt", async () => { const handler = new RequestyHandler( makeApiHandlerOptions({ diff --git a/src/api/providers/__tests__/unbound.spec.ts b/src/api/providers/__tests__/unbound.spec.ts index c0f142d8ad..d3d9c9b010 100644 --- a/src/api/providers/__tests__/unbound.spec.ts +++ b/src/api/providers/__tests__/unbound.spec.ts @@ -204,6 +204,29 @@ describe("UnboundHandler", () => { ) }) + type OpenAiClientStub = { chat: { completions: { create: ReturnType } } } + + it("forwards the abort signal and deadline to the completion request", async () => { + const mockCreate = (OpenAI as unknown as () => OpenAiClientStub)().chat.completions.create + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + const handler = new UnboundHandler({ unboundApiKey: "test-key", unboundModelId: "openai/gpt-4o" }) + const { signal } = new AbortController() + + await handler.completePrompt("Write a haiku", { abortSignal: signal, timeoutMs: 5_000 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { signal, timeout: 5_000 }) + }) + + it("sends a single-argument request when no cancellation or deadline is given", async () => { + const mockCreate = (OpenAI as unknown as () => OpenAiClientStub)().chat.completions.create + mockCreate.mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + const handler = new UnboundHandler({ unboundApiKey: "test-key", unboundModelId: "openai/gpt-4o" }) + + await handler.completePrompt("Write a haiku", { timeoutMs: -1 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything()) + }) + describe("model catalog", () => { it("reports its cache scope", () => { const handler = new UnboundHandler({ unboundApiKey: "test-key", unboundModelId: "openai/gpt-4o" }) diff --git a/src/api/providers/__tests__/vercel-ai-gateway.spec.ts b/src/api/providers/__tests__/vercel-ai-gateway.spec.ts index 5bafa25c78..064778ce4f 100644 --- a/src/api/providers/__tests__/vercel-ai-gateway.spec.ts +++ b/src/api/providers/__tests__/vercel-ai-gateway.spec.ts @@ -763,6 +763,23 @@ describe("VercelAiGatewayHandler", () => { })) }) + it("forwards the abort signal and deadline to the completion request", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + const { signal } = new AbortController() + + await handler.completePrompt("Complete this: Hello", { abortSignal: signal, timeoutMs: 5_000 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { signal, timeout: 5_000 }) + }) + + it("sends a single-argument request when no cancellation or deadline is given", async () => { + const handler = new VercelAiGatewayHandler(mockOptions) + + await handler.completePrompt("Complete this: Hello", { timeoutMs: Number.POSITIVE_INFINITY }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything()) + }) + it("completes prompt correctly", async () => { const handler = new VercelAiGatewayHandler(mockOptions) const prompt = "Complete this: Hello" diff --git a/src/api/providers/__tests__/zoo-gateway.spec.ts b/src/api/providers/__tests__/zoo-gateway.spec.ts index fe088bb5c9..31f2e3e828 100644 --- a/src/api/providers/__tests__/zoo-gateway.spec.ts +++ b/src/api/providers/__tests__/zoo-gateway.spec.ts @@ -506,6 +506,23 @@ describe("ZooGatewayHandler", () => { })) }) + it("forwards the abort signal and deadline to the completion request", async () => { + const handler = new ZooGatewayHandler(mockOptions) + const { signal } = new AbortController() + + await handler.completePrompt("Complete this: Hello", { abortSignal: signal, timeoutMs: 5_000 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything(), { signal, timeout: 5_000 }) + }) + + it("sends a single-argument request when no cancellation or deadline is given", async () => { + const handler = new ZooGatewayHandler(mockOptions) + + await handler.completePrompt("Complete this: Hello", { timeoutMs: -1 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.anything()) + }) + it("returns completion text from the gateway", async () => { const handler = new ZooGatewayHandler(mockOptions) diff --git a/src/api/providers/base-openai-compatible-provider.ts b/src/api/providers/base-openai-compatible-provider.ts index b177726f85..9167594690 100644 --- a/src/api/providers/base-openai-compatible-provider.ts +++ b/src/api/providers/base-openai-compatible-provider.ts @@ -14,6 +14,7 @@ import { BaseProvider } from "./base-provider" import { handleOpenAIError } from "./utils/error-handler" import { calculateApiCostOpenAI } from "../../shared/cost" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +import { hasUsableTimeout, toSdkTimeoutMs, withSdkRequestOptions } from "./utils/abort-signal" type BaseOpenAiCompatibleProviderOptions = ApiHandlerOptions & { providerName: string @@ -212,7 +213,7 @@ export abstract class BaseOpenAiCompatibleProvider } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id: modelId, info: modelInfo } = this.getModel() const params: OpenAI.Chat.Completions.ChatCompletionCreateParams = { @@ -226,7 +227,19 @@ export abstract class BaseOpenAiCompatibleProvider } try { - const response = await this.client.chat.completions.create(params) + // Forward the caller's cancellation and deadline to the request itself, so a + // superseded title request stops instead of running to the client default timeout. + // The OpenAI SDK rejects a present-but-undefined or unusable `timeout`, so only a + // real value is forwarded, rounded and capped like mergeAbortSignalAndTimeout. + const requestOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal !== undefined) { + requestOptions.signal = options.abortSignal + } + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + requestOptions.timeout = toSdkTimeoutMs(timeoutMs) + } + const response = await this.client.chat.completions.create(params, ...withSdkRequestOptions(requestOptions)) // Check for provider-specific error responses (e.g., MiniMax base_resp) const responseAny = response as any diff --git a/src/api/providers/kenari.ts b/src/api/providers/kenari.ts index 22ca50c916..ad91bef29e 100644 --- a/src/api/providers/kenari.ts +++ b/src/api/providers/kenari.ts @@ -24,6 +24,7 @@ import type { import { RouterProvider } from "./router-provider" import { getKenariModels } from "./fetchers/kenari" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +import { hasUsableTimeout, throwIfAborted, toSdkTimeoutMs, withSdkRequestOptions } from "./utils/abort-signal" /** * API handler for Kenari, an Indonesian OpenAI-compatible AI gateway billed @@ -140,7 +141,9 @@ export class KenariHandler extends RouterProvider implements SingleCompletionHan * @returns The model's reply text, or an empty string if no content is returned. * @throws Error with a Kenari-specific prefix if the request fails. */ - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // An already-cancelled operation must not start a request. + throwIfAborted(options?.abortSignal) const { id: modelId, info } = await this.fetchModel() try { @@ -156,7 +159,20 @@ export class KenariHandler extends RouterProvider implements SingleCompletionHan requestOptions.max_completion_tokens = info.maxTokens - const response = await this.client.chat.completions.create(requestOptions) + // Forward the caller's cancellation and deadline to the request itself. + const sdkRequestOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal !== undefined) { + sdkRequestOptions.signal = options.abortSignal + } + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + sdkRequestOptions.timeout = toSdkTimeoutMs(timeoutMs) + } + + const response = await this.client.chat.completions.create( + requestOptions, + ...withSdkRequestOptions(sdkRequestOptions), + ) return response.choices[0]?.message.content || "" } catch (error) { if (error instanceof Error) { diff --git a/src/api/providers/openai.ts b/src/api/providers/openai.ts index 98ea18df8a..acba39eed8 100644 --- a/src/api/providers/openai.ts +++ b/src/api/providers/openai.ts @@ -29,6 +29,7 @@ import { BaseProvider } from "./base-provider" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index" import { handleOpenAIError } from "./utils/error-handler" import { extractReasoningFromDelta } from "./utils/extract-reasoning" +import { hasUsableTimeout, throwIfAborted, toSdkTimeoutMs } from "./utils/abort-signal" type OpenAiUsage = | { @@ -354,8 +355,10 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { try { + // An already-cancelled operation must not start a request. + throwIfAborted(options?.abortSignal) const isAzureAiInference = this._isAzureAiInference(this.options.openAiBaseUrl) const model = this.getModel() const modelInfo = model.info @@ -372,10 +375,19 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl let response try { - response = await this.client.chat.completions.create( - requestOptions, - isAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {}, - ) + // Forward the caller's cancellation and deadline to the request itself, so a + // superseded title request stops instead of running to the client default timeout. + const sdkRequestOptions: OpenAI.RequestOptions = isAzureAiInference + ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } + : {} + if (options?.abortSignal !== undefined) { + sdkRequestOptions.signal = options.abortSignal + } + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + sdkRequestOptions.timeout = toSdkTimeoutMs(timeoutMs) + } + response = await this.client.chat.completions.create(requestOptions, sdkRequestOptions) } catch (error) { throw handleOpenAIError(error, this.providerName) } diff --git a/src/api/providers/openrouter.ts b/src/api/providers/openrouter.ts index 6d2057a7a3..ca97b98b65 100644 --- a/src/api/providers/openrouter.ts +++ b/src/api/providers/openrouter.ts @@ -43,6 +43,7 @@ import type { SingleCompletionHandler, } from "../index" import { handleOpenAIError } from "./utils/error-handler" +import { hasUsableTimeout, resolveModelWithAbort, throwIfAborted, toSdkTimeoutMs } from "./utils/abort-signal" import { generateImageWithProvider, ImageGenerationResult } from "./utils/image-generation" import { applyRouterToolPreferences } from "./utils/router-tool-preferences" @@ -600,8 +601,16 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH return { id, info, topP: isDeepSeekR1 ? 0.95 : undefined, ...params } } - async completePrompt(prompt: string, _options?: CompletePromptOptions) { - const { id: modelId, maxTokens, temperature, reasoning } = await this.fetchModel() + async completePrompt(prompt: string, options?: CompletePromptOptions) { + // A cancelled operation must not start a request, and model discovery must not outlive + // the caller's cancellation: both run under the caller's signal. + throwIfAborted(options?.abortSignal) + const { + id: modelId, + maxTokens, + temperature, + reasoning, + } = await resolveModelWithAbort(() => this.fetchModel(), options?.abortSignal, this.providerName) const completionParams: OpenRouterChatCompletionParams = { model: modelId, @@ -622,14 +631,27 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH } // Add Anthropic beta header for fine-grained tool streaming when using Anthropic models - const requestOptions = modelId.startsWith("anthropic/") - ? { headers: { "x-anthropic-beta": "fine-grained-tool-streaming-2025-05-14" } } - : undefined + const requestOptions: OpenAI.RequestOptions = {} + if (modelId.startsWith("anthropic/")) { + requestOptions.headers = { "x-anthropic-beta": "fine-grained-tool-streaming-2025-05-14" } + } + // Propagate the caller's cancellation and deadline to the request itself, so a + // superseded title request stops instead of running to the client default timeout. + if (options?.abortSignal !== undefined) { + requestOptions.signal = options.abortSignal + } + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + requestOptions.timeout = toSdkTimeoutMs(timeoutMs) + } let response try { - response = await this.client.chat.completions.create(completionParams, requestOptions) + response = await this.client.chat.completions.create( + completionParams, + Object.keys(requestOptions).length > 0 ? requestOptions : undefined, + ) } catch (error) { // Try to parse as OpenRouter error structure using Zod const parseResult = OpenRouterErrorResponseSchema.safeParse(error) diff --git a/src/api/providers/poe.ts b/src/api/providers/poe.ts index 29333aa056..75ff138fca 100644 --- a/src/api/providers/poe.ts +++ b/src/api/providers/poe.ts @@ -25,6 +25,7 @@ import { convertToAiSdkMessages, convertToolsForAiSdk, processAiSdkStreamPart } import { ApiStream } from "../transform/stream" import { BaseProvider } from "./base-provider" +import { mergeAbortSignalAndTimeout, throwIfAborted } from "./utils/abort-signal" import { NOT_PROVIDED } from "./constants" import type { SingleCompletionHandler, @@ -184,12 +185,17 @@ export class PoeHandler extends BaseProvider implements SingleCompletionHandler } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { const { id } = this.getModel() try { + // An already-cancelled operation must not start a request. + throwIfAborted(options?.abortSignal) const { text } = await generateText({ model: this.poe(id), prompt, + // Forward the caller's cancellation and deadline to the request itself so a + // superseded title request stops instead of running to the default timeout. + abortSignal: mergeAbortSignalAndTimeout(options?.abortSignal, options?.timeoutMs), }) return text } catch (error) { diff --git a/src/api/providers/requesty.ts b/src/api/providers/requesty.ts index f8d23c0ce6..1c96f50fd0 100644 --- a/src/api/providers/requesty.ts +++ b/src/api/providers/requesty.ts @@ -29,6 +29,7 @@ import type { } from "../index" import { toRequestyServiceUrl } from "../../shared/utils/requesty" import { handleOpenAIError } from "./utils/error-handler" +import { hasUsableTimeout, throwIfAborted, toSdkTimeoutMs, withSdkRequestOptions } from "./utils/abort-signal" import { applyRouterToolPreferences } from "./utils/router-tool-preferences" import { extractReasoningFromDelta } from "./utils/extract-reasoning" @@ -249,7 +250,9 @@ export class RequestyHandler extends BaseProvider implements SingleCompletionHan } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // An already-cancelled operation must not start a request. + throwIfAborted(options?.abortSignal) const { id: model, maxTokens: max_tokens, temperature } = await this.fetchModel() const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [{ role: "system", content: prompt }] @@ -263,7 +266,19 @@ export class RequestyHandler extends BaseProvider implements SingleCompletionHan let response: OpenAI.Chat.ChatCompletion try { - response = await this.client.chat.completions.create(completionParams) + // Forward the caller's cancellation and deadline to the request itself. + const sdkRequestOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal !== undefined) { + sdkRequestOptions.signal = options.abortSignal + } + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + sdkRequestOptions.timeout = toSdkTimeoutMs(timeoutMs) + } + response = await this.client.chat.completions.create( + completionParams, + ...withSdkRequestOptions(sdkRequestOptions), + ) } catch (error) { throw handleOpenAIError(error, this.providerName) } diff --git a/src/api/providers/unbound.ts b/src/api/providers/unbound.ts index 59c7bb3b66..641b9e0bc2 100644 --- a/src/api/providers/unbound.ts +++ b/src/api/providers/unbound.ts @@ -28,6 +28,7 @@ import type { ModelCacheScope, } from "../index" import { handleOpenAIError } from "./utils/error-handler" +import { hasUsableTimeout, throwIfAborted, toSdkTimeoutMs, withSdkRequestOptions } from "./utils/abort-signal" import { applyRouterToolPreferences } from "./utils/router-tool-preferences" import { extractReasoningFromDelta } from "./utils/extract-reasoning" @@ -228,7 +229,9 @@ export class UnboundHandler extends BaseProvider implements SingleCompletionHand } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // An already-cancelled operation must not start a request. + throwIfAborted(options?.abortSignal) const { id: model, maxTokens: max_tokens, temperature } = await this.fetchModel() const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = [{ role: "system", content: prompt }] @@ -242,7 +245,19 @@ export class UnboundHandler extends BaseProvider implements SingleCompletionHand let response: OpenAI.Chat.ChatCompletion try { - response = await this.client.chat.completions.create(completionParams) + // Forward the caller's cancellation and deadline to the request itself. + const sdkRequestOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal !== undefined) { + sdkRequestOptions.signal = options.abortSignal + } + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + sdkRequestOptions.timeout = toSdkTimeoutMs(timeoutMs) + } + response = await this.client.chat.completions.create( + completionParams, + ...withSdkRequestOptions(sdkRequestOptions), + ) } catch (error) { throw handleOpenAIError(error, this.providerName) } diff --git a/src/api/providers/utils/abort-signal.ts b/src/api/providers/utils/abort-signal.ts index 6abbc12437..cf6c858c08 100644 --- a/src/api/providers/utils/abort-signal.ts +++ b/src/api/providers/utils/abort-signal.ts @@ -134,6 +134,18 @@ export type OpenAiRequestOptions = { signal?: AbortSignal } +/** + * Build the trailing request-options argument for an SDK call: the options + * object when it carries anything, or no argument at all. + * + * Call sites that pass nothing keep their original single-argument shape, so a + * completion without cancellation or a deadline sends exactly the request it + * sent before these options existed. + */ +export function withSdkRequestOptions(options: T): [T] | [] { + return Object.keys(options).length > 0 ? [options] : [] +} + /** * Whether a failure indicates an aborted request: the caller's signal fired, * an `Error` carries a native abort error name (`AbortError`, diff --git a/src/api/providers/vercel-ai-gateway.ts b/src/api/providers/vercel-ai-gateway.ts index 135ac8bef6..efb9443f31 100644 --- a/src/api/providers/vercel-ai-gateway.ts +++ b/src/api/providers/vercel-ai-gateway.ts @@ -27,6 +27,7 @@ import type { } from "../index" import { RouterProvider } from "./router-provider" import { getVercelAiGatewayModels } from "./fetchers/vercel-ai-gateway" +import { hasUsableTimeout, throwIfAborted, toSdkTimeoutMs, withSdkRequestOptions } from "./utils/abort-signal" // Extend OpenAI's CompletionUsage to include Vercel AI Gateway specific fields interface VercelAiGatewayUsage extends OpenAI.CompletionUsage { @@ -155,7 +156,9 @@ export class VercelAiGatewayHandler extends RouterProvider implements SingleComp } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // An already-cancelled operation must not start a request. + throwIfAborted(options?.abortSignal) const { id: modelId, info } = await this.fetchModel() try { @@ -175,7 +178,20 @@ export class VercelAiGatewayHandler extends RouterProvider implements SingleComp requestOptions.max_completion_tokens = info.maxTokens - const response = await this.client.chat.completions.create(requestOptions) + // Forward the caller's cancellation and deadline to the request itself. + const sdkRequestOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal !== undefined) { + sdkRequestOptions.signal = options.abortSignal + } + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + sdkRequestOptions.timeout = toSdkTimeoutMs(timeoutMs) + } + + const response = await this.client.chat.completions.create( + requestOptions, + ...withSdkRequestOptions(sdkRequestOptions), + ) return response.choices[0]?.message.content || "" } catch (error) { if (error instanceof Error) { diff --git a/src/api/providers/zoo-gateway.ts b/src/api/providers/zoo-gateway.ts index 5f33ee48fe..bafef79ce0 100644 --- a/src/api/providers/zoo-gateway.ts +++ b/src/api/providers/zoo-gateway.ts @@ -30,6 +30,7 @@ import type { import { NOT_PROVIDED } from "./constants" import { RouterProvider } from "./router-provider" import { getZooGatewayModels } from "./fetchers/zoo-gateway" +import { hasUsableTimeout, throwIfAborted, toSdkTimeoutMs, withSdkRequestOptions } from "./utils/abort-signal" function getApiErrorStatus(error: unknown): number | undefined { if (typeof error === "object" && error !== null && "status" in error) { @@ -294,7 +295,9 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio } } - async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise { + async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // An already-cancelled operation must not start a request. + throwIfAborted(options?.abortSignal) this.ensureAuthenticated() const { id: modelId, info } = await this.fetchModel() @@ -312,7 +315,20 @@ export class ZooGatewayHandler extends RouterProvider implements SingleCompletio requestOptions.max_completion_tokens = info.maxTokens - const response = await this.client.chat.completions.create(requestOptions) + // Forward the caller's cancellation and deadline to the request itself. + const sdkRequestOptions: OpenAI.RequestOptions = {} + if (options?.abortSignal !== undefined) { + sdkRequestOptions.signal = options.abortSignal + } + const timeoutMs = options?.timeoutMs + if (hasUsableTimeout(timeoutMs)) { + sdkRequestOptions.timeout = toSdkTimeoutMs(timeoutMs) + } + + const response = await this.client.chat.completions.create( + requestOptions, + ...withSdkRequestOptions(sdkRequestOptions), + ) return response.choices[0]?.message.content || "" } catch (error) { try { diff --git a/src/core/task-persistence/TaskHistoryStore.ts b/src/core/task-persistence/TaskHistoryStore.ts index a9e0243823..2c5464fad4 100644 --- a/src/core/task-persistence/TaskHistoryStore.ts +++ b/src/core/task-persistence/TaskHistoryStore.ts @@ -28,7 +28,7 @@ function mergeWithDisk(delta: Partial): (existing: unknown, incomin /** Thrown inside the `setTitleIfUnset` merge callback to abort the write without touching the file. */ class TitleWriteSkipped extends Error { - constructor(readonly reason: "missing" | "titled") { + constructor(readonly reason: "missing" | "titled" | "invalid" | "mismatched") { // Stryker disable next-line StringLiteral: diagnostic error message only, caught internally super(`title write skipped: ${reason}`) } @@ -1010,21 +1010,26 @@ export class TaskHistoryStore { * provider on the same storage path) saved a manual title, and a delta merge would * otherwise overwrite it. * - * @returns "conflict" when a title already exists or the record was deleted on disk (the cache is refreshed). - * @throws If the task ID is not present in the cache. + * The locked record must be a valid task-history item carrying this task ID: a + * malformed or mismatched record is never merged, rewritten, or cached. + * + * @returns "missing" when the task has no history item yet (the caller defers the title), + * "conflict" when a title or a manual clear already owns the record, or when the + * record was deleted or is invalid on disk (the cache is refreshed or dropped). */ - public setTitleIfUnset(taskId: string, title: string): Promise<"applied" | "conflict"> { + public setTitleIfUnset(taskId: string, title: string): Promise<"applied" | "conflict" | "missing"> { return this.withLock(async () => { const current = this.cache.get(taskId) if (!current) { - throw new Error(`[TaskHistoryStore] setTitleIfUnset: task ${taskId} not found in cache`) + return "missing" } - if (current.title) { + if (current.title || current.titleCleared) { return "conflict" } const filePath = await this.getTaskFilePath(taskId) let written: HistoryItem = { ...current, title } + let diskRecord: HistoryItem | undefined try { await safeWriteJson(filePath, written, { merge: (existing, incoming) => { @@ -1034,11 +1039,25 @@ export class TaskHistoryStore { if (!existing || typeof existing !== "object" || !("id" in existing)) { throw new TitleWriteSkipped("missing") } - if ((existing as HistoryItem).title) { + // Validate the locked disk record with the canonical task-history schema and + // the expected task ID, so a malformed record or another task's record is + // never merged into this task or sent to the history view. + const parsed = historyItemSchema.safeParse(existing) + if (!parsed.success) { + // Stryker disable next-line StringLiteral: internal discrimination literal between skip reasons + throw new TitleWriteSkipped("invalid") + } + if (parsed.data.id !== taskId) { + // Stryker disable next-line StringLiteral: internal discrimination literal between skip reasons + throw new TitleWriteSkipped("mismatched") + } + // Keep the validated record so a losing race refreshes the cache with it. + diskRecord = parsed.data + if (parsed.data.title || parsed.data.titleCleared) { // Stryker disable next-line StringLiteral: internal discrimination literal between missing and titled reasons throw new TitleWriteSkipped("titled") } - written = mergeHistoryDelta(existing, incoming as HistoryItem, { title }) + written = mergeHistoryDelta(parsed.data, incoming as HistoryItem, { title }) return written }, }) @@ -1046,15 +1065,17 @@ export class TaskHistoryStore { if (!(error instanceof TitleWriteSkipped)) { throw error } - if (error.reason === "missing") { - this.cache.delete(taskId) - } else { - const diskItem = await this.readTaskFile(taskId) - // Stryker disable next-line ConditionalExpression: defense-in-depth, disk item is always present when merge found titled record - if (diskItem) { - this.cache.set(taskId, diskItem) + if (error.reason === "titled") { + // Stryker disable next-line ConditionalExpression: defense-in-depth, a titled record is validated before the write + if (diskRecord) { + this.cache.set(taskId, diskRecord) } + return "conflict" } + // The record is gone or unusable, so the stale cache entry is dropped instead of + // being written back; the disk record itself is left untouched. + this.cache.delete(taskId) + this.taskFileMtimes.delete(taskId) return "conflict" } diff --git a/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts b/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts index 98efed9a61..d4a8b8dada 100644 --- a/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts +++ b/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts @@ -134,9 +134,58 @@ describe("TaskHistoryStore cross-instance safety", () => { storeReloaded.dispose() }) - it("throws when the task is not in the cache", async () => { + it("reports missing when the task is not in the cache", async () => { await storeA.initialize() - await expect(storeA.setTitleIfUnset("missing-id", "Auto title")).rejects.toThrow("not found in cache") + await expect(storeA.setTitleIfUnset("missing-id", "Auto title")).resolves.toBe("missing") + }) + + it("reports conflict for a malformed disk record without rewriting it", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "title-malformed", task: "Prompt" })) + const filePath = path.join(tmpDir, "tasks", "title-malformed", GlobalFileNames.historyItem) + // Valid JSON that fails the canonical task-history schema: id must be a string. + const malformed = JSON.stringify({ id: 123 }) + await fs.writeFile(filePath, malformed, "utf8") + + await expect(storeA.setTitleIfUnset("title-malformed", "Auto title")).resolves.toBe("conflict") + + // The stale cache entry is dropped and the unusable record stays untouched. + expect(storeA.get("title-malformed")).toBeUndefined() + expect(await fs.readFile(filePath, "utf8")).toBe(malformed) + }) + + it("reports conflict when the disk record belongs to a different task", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "title-mismatched", task: "Prompt" })) + const filePath = path.join(tmpDir, "tasks", "title-mismatched", GlobalFileNames.historyItem) + const otherTask = JSON.stringify(makeHistoryItem({ id: "other-task" }), null, "\t") + await fs.writeFile(filePath, otherTask, "utf8") + + await expect(storeA.setTitleIfUnset("title-mismatched", "Auto title")).resolves.toBe("conflict") + + expect(storeA.get("title-mismatched")).toBeUndefined() + // The other task's record is never merged into this task or overwritten. + expect(await fs.readFile(filePath, "utf8")).toBe(otherTask) + }) + + it("keeps a manual clear that another instance persisted", async () => { + await storeA.initialize() + await storeB.initialize() + await storeA.upsert(makeHistoryItem({ id: "title-cleared", task: "Original prompt" })) + await storeB.reconcile() + + // Instance B clears the title (rename with an empty name) on disk. + await storeB.atomicReadAndUpdate("title-cleared", (item) => ({ + ...item, + title: undefined, + titleCleared: true, + })) + + await expect(storeA.setTitleIfUnset("title-cleared", "Auto title")).resolves.toBe("conflict") + + expect(storeA.get("title-cleared")?.title).toBeUndefined() + await storeB.reconcile() + expect(storeB.get("title-cleared")?.title).toBeUndefined() }) it("returns conflict immediately when cached item already has a title", async () => { diff --git a/src/core/task-persistence/__tests__/TaskHistoryStore.realConcurrency.spec.ts b/src/core/task-persistence/__tests__/TaskHistoryStore.realConcurrency.spec.ts index e3e4445f35..5574f352ad 100644 --- a/src/core/task-persistence/__tests__/TaskHistoryStore.realConcurrency.spec.ts +++ b/src/core/task-persistence/__tests__/TaskHistoryStore.realConcurrency.spec.ts @@ -137,6 +137,71 @@ describe("TaskHistoryStore real cross-host locking", () => { } }) + it("keeps a manual rename that commits while the automatic title waits on the file lock", async () => { + const storagePath = await fs.mkdtemp(path.join(os.tmpdir(), "task-history-title-race-")) + // storeA generates the session title; storeB is a second window that renames the task. + const storeA = new TaskHistoryStore(storagePath) + const storeB = new TaskHistoryStore(storagePath) + const filePath = path.join(storagePath, "tasks", "shared-task", "history_item.json") + const actualFs = await vi.importActual("fs/promises") + let signalRenameRead!: () => void + const renameRead = new Promise((resolve) => { + signalRenameRead = resolve + }) + let releaseRename!: () => void + const renameCanCommit = new Promise((resolve) => { + releaseRename = resolve + }) + let pauseRenameRead = false + + try { + await storeA.initialize() + await storeA.upsert(item("shared-task")) + await storeB.initialize() + + // Pause the rename inside the per-file lock, after its disk read and before + // its commit, so the automatic title write is queued behind the rename. + vi.mocked(fs.readFile).mockImplementation(async (...args: Parameters) => { + const result = await actualFs.readFile(...args) + if (pauseRenameRead && args[0] === filePath) { + pauseRenameRead = false + signalRenameRead() + await renameCanCommit + } + return result + }) + + pauseRenameRead = true + const rename = storeB.atomicReadAndUpdate("shared-task", (current) => ({ + ...current, + title: "Manual title", + })) + await renameRead + + const autoTitle = storeA.setTitleIfUnset("shared-task", "Auto title") + let autoSettled = false + void autoTitle.finally(() => { + autoSettled = true + }) + await new Promise((resolve) => setTimeout(resolve, 25)) + // The automatic write must stay queued while the rename holds the lock. + expect(autoSettled).toBe(false) + + releaseRename() + await rename + await expect(autoTitle).resolves.toBe("conflict") + + const onDisk = JSON.parse(await fs.readFile(filePath, "utf8")) as HistoryItem + expect(onDisk.title).toBe("Manual title") + expect(storeA.get("shared-task")?.title).toBe("Manual title") + } finally { + vi.mocked(fs.readFile).mockImplementation(actualFs.readFile) + storeA.dispose() + storeB.dispose() + await fs.rm(storagePath, { recursive: true, force: true }) + } + }) + it("reports a bounded error when one store never reaches the write barrier", async () => { const storagePath = await fs.mkdtemp(path.join(os.tmpdir(), "task-history-missed-barrier-")) const storeA = new TaskHistoryStore(storagePath) diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts index 4b108221b1..32fde1ddd7 100644 --- a/src/core/task/Task.ts +++ b/src/core/task/Task.ts @@ -145,7 +145,7 @@ import { prepareApiConversationMessage } from "./apiConversationHistory" import { shouldAddUserMessageToHistory } from "./messageCounting" import { type TaskExecutionContext } from "./providerHandoff" import { initialRunState, type Ask, type RunEvent, type RunState } from "./run-state/runState" -import { generateSessionTitle } from "./session-title" +import { flushPendingSessionTitle, generateSessionTitle } from "./session-title" import { applyRunEvent, type RunRejection } from "./run-state/runStateShadow" import { divergenceKey, @@ -1893,6 +1893,10 @@ export class Task extends EventEmitter implements TaskLike { const provider = this.providerRef.deref() const existingStatus = provider?.taskHistoryStore.get(this.taskId)?.status await provider?.updateTaskHistory(existingStatus ? { ...historyItem, status: existingStatus } : historyItem) + + // A title generated before this task's history item existed is retried now that the + // record is durable. The call is a no-op when no title is pending. + void flushPendingSessionTitle(this.taskId) } catch (error) { // The message array was persisted above; a metadata or task-history failure must // not mask that write (see the method docs). The next saveClineMessages() call diff --git a/src/core/task/__tests__/Task.session-title.spec.ts b/src/core/task/__tests__/Task.session-title.spec.ts index 719e62fa27..1c2d3b8737 100644 --- a/src/core/task/__tests__/Task.session-title.spec.ts +++ b/src/core/task/__tests__/Task.session-title.spec.ts @@ -9,10 +9,26 @@ import type { GenerateSessionTitleOptions } from "../session-title" vi.mock("../session-title", () => ({ generateSessionTitle: vi.fn().mockResolvedValue(undefined), + flushPendingSessionTitle: vi.fn().mockResolvedValue(undefined), })) // Import the mocked module for spy access. -import { generateSessionTitle } from "../session-title" +import { flushPendingSessionTitle, generateSessionTitle } from "../session-title" + +// The task save path is mocked so the retry trigger can be observed without touching disk. +const persistenceMocks = vi.hoisted(() => ({ + saveTaskMessages: vi.fn(), + taskMetadata: vi.fn(), +})) + +vi.mock("../../task-persistence", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + saveTaskMessages: persistenceMocks.saveTaskMessages, + taskMetadata: persistenceMocks.taskMetadata, + } +}) vi.mock("@roo-code/telemetry", () => ({ TelemetryService: { @@ -111,6 +127,8 @@ function createMockProvider(): ClineProvider { postStateToWebview: vi.fn().mockResolvedValue(undefined), postStateToWebviewWithoutTaskHistory: vi.fn().mockResolvedValue(undefined), updateTaskHistory: vi.fn().mockResolvedValue(undefined), + // `saveClineMessages` reads the stored status before writing the item back. + taskHistoryStore: { get: vi.fn().mockReturnValue(undefined) }, } as unknown as ClineProvider } @@ -148,6 +166,20 @@ function lastTitleCall(): GenerateSessionTitleOptions | undefined { describe("Task#startTask session title trigger", () => { beforeEach(() => { vi.mocked(generateSessionTitle).mockClear() + vi.mocked(flushPendingSessionTitle).mockClear() + persistenceMocks.saveTaskMessages.mockReset().mockResolvedValue(undefined) + persistenceMocks.taskMetadata.mockReset().mockImplementation(async ({ taskId }: { taskId: string }) => ({ + historyItem: { + id: taskId, + number: 1, + ts: Date.now(), + task: "example root task", + tokensIn: 0, + tokensOut: 0, + totalCost: 0, + }, + tokenUsage: { totalTokensIn: 0, totalTokensOut: 0, totalCacheWrites: 0, totalCacheReads: 0, totalCost: 0 }, + })) }) it("fires title generation for a root task without awaiting it", async () => { @@ -292,4 +324,27 @@ describe("Task#startTask session title trigger", () => { ) consoleErrorSpy.mockRestore() }) + + it("retries a pending title once the history item is saved", async () => { + const provider = createMockProvider() + const task = createTestTask(provider) + + await (task as unknown as { saveClineMessages: (merge?: boolean) => Promise }).saveClineMessages() + + expect(persistenceMocks.saveTaskMessages).toHaveBeenCalled() + expect(vi.mocked(flushPendingSessionTitle)).toHaveBeenCalledWith(task.taskId) + }) + + it("does not retry a pending title when the message save fails", async () => { + const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) + const provider = createMockProvider() + const task = createTestTask(provider) + persistenceMocks.saveTaskMessages.mockRejectedValue(new Error("disk full")) + + await expect( + (task as unknown as { saveClineMessages: (merge?: boolean) => Promise }).saveClineMessages(), + ).resolves.toBe(false) + expect(vi.mocked(flushPendingSessionTitle)).not.toHaveBeenCalled() + consoleErrorSpy.mockRestore() + }) }) diff --git a/src/core/task/__tests__/session-title.spec.ts b/src/core/task/__tests__/session-title.spec.ts index 4d5e43ff57..3cde79d675 100644 --- a/src/core/task/__tests__/session-title.spec.ts +++ b/src/core/task/__tests__/session-title.spec.ts @@ -10,6 +10,9 @@ import { fallbackTitle, generateSessionTitle, cancelSessionTitle, + cancelSessionTitlesForProvider, + drainSessionTitles, + flushPendingSessionTitle, sanitizeUserMessageForTitle, extractLlmTitle, type SessionTitleApplyResult, @@ -31,10 +34,12 @@ type ApplyTitleCall = { taskId: string; title: string } function createTitleProvider( overrides: { applySessionTitle?: (taskId: string, title: string) => Promise + isDisposed?: boolean } = {}, ): ClineProvider { const provider = { applySessionTitle: overrides.applySessionTitle ?? (async () => "applied" as const), + isDisposed: overrides.isDisposed ?? false, } return provider as unknown as ClineProvider } @@ -527,6 +532,30 @@ describe("generateSessionTitle", () => { consoleErrorSpy.mockRestore() }) + it("does not write the fallback title once the generation was cancelled", async () => { + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let applyCalls = 0 + const provider = createTitleProvider({ + applySessionTitle: async () => { + applyCalls++ + // A manual rename lands while this write is in flight and the write fails. + cancelSessionTitle("task-cancel-before-fallback") + throw new Error("transient failure") + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-cancel-before-fallback", + userMessage: "Fix login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + // The fallback write is skipped at the cancellation check instead of being attempted. + expect(applyCalls).toBe(1) + }) + it("stops fallback retry immediately if the generation is cancelled", async () => { const consoleErrorSpy = vi.spyOn(console, "error").mockImplementation(() => {}) const recorder = createCompletePromptRecorder(async () => { @@ -581,4 +610,253 @@ describe("generateSessionTitle", () => { }), ).resolves.toBeUndefined() }) + + it("sends only the user text to the model when the message carries a system reminder", async () => { + const recorder = createCompletePromptRecorder(async () => "Login fix title") + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-reminder", + userMessage: + "Environment details: workspace /home/user/projectPlease fix the login button", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(recorder.calls).toHaveLength(1) + expect(recorder.calls[0].prompt).toContain("Please fix the login button") + expect(recorder.calls[0].prompt).not.toContain("") + expect(recorder.calls[0].prompt).not.toContain("") + expect(recorder.calls[0].prompt).not.toContain("Environment details") + expect(applied).toEqual([{ taskId: "task-reminder", title: "Login fix title" }]) + }) + + it("does not start a request for a provider that is already disposed", async () => { + const recorder = createCompletePromptRecorder(async () => "Should not be requested") + const provider = createTitleProvider({ isDisposed: true }) + + await generateSessionTitle({ + provider, + taskId: "task-disposed-start", + userMessage: "Fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(recorder.calls).toHaveLength(0) + }) + + it("cancels and drains the generation a disposing provider owns", async () => { + let abortSignal: AbortSignal | undefined + const recorder = createCompletePromptRecorder(async ({ options }) => { + abortSignal = options?.abortSignal + return new Promise(() => {}) + }) + const provider = createTitleProvider() + + const generation = generateSessionTitle({ + provider, + taskId: "task-dispose", + userMessage: "Fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + expect(abortSignal?.aborted).toBe(false) + cancelSessionTitlesForProvider(provider) + expect(abortSignal?.aborted).toBe(true) + + // Draining waits for the aborted generation to settle instead of leaving the write behind. + await drainSessionTitles(provider) + await generation + }) + + it("cancels only the disposing provider's generation", async () => { + const signals: Record = {} + const recorder = createCompletePromptRecorder(async ({ prompt, options }) => { + signals[prompt] = options?.abortSignal + return new Promise(() => {}) + }) + const providerA = createTitleProvider() + const providerB = createTitleProvider() + + const generationA = generateSessionTitle({ + provider: providerA, + taskId: "task-dispose-a", + userMessage: "message for A", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + const generationB = generateSessionTitle({ + provider: providerB, + taskId: "task-dispose-b", + userMessage: "message for B", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + cancelSessionTitlesForProvider(providerA) + + const signalForA = Object.entries(signals).find(([prompt]) => prompt.includes("message for A"))?.[1] + const signalForB = Object.entries(signals).find(([prompt]) => prompt.includes("message for B"))?.[1] + expect(signalForA?.aborted).toBe(true) + expect(signalForB?.aborted).toBe(false) + + await generationA + cancelSessionTitle("task-dispose-b") + await generationB + }) +}) + +describe("flushPendingSessionTitle", () => { + it("applies a title that was generated before the history item existed", async () => { + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let historyItemPersisted = false + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + if (!historyItemPersisted) { + return "not-ready" + } + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-deferred", + userMessage: "Fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + expect(applied).toHaveLength(0) + + // The first write attempted before the record existed is retried until it lands. + await flushPendingSessionTitle("task-deferred") + expect(applied).toHaveLength(0) + + historyItemPersisted = true + await flushPendingSessionTitle("task-deferred") + + expect(applied).toEqual([{ taskId: "task-deferred", title: "Generated Title" }]) + + // The title is applied once; a later save does not repeat the write. + await flushPendingSessionTitle("task-deferred") + expect(applied).toHaveLength(1) + }) + + it("applies the fallback title once the history item exists", async () => { + const recorder = createCompletePromptRecorder(async () => { + throw new Error("api down") + }) + let historyItemPersisted = false + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + if (!historyItemPersisted) { + return "not-ready" + } + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-deferred-fallback", + userMessage: "Fix login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + historyItemPersisted = true + await flushPendingSessionTitle("task-deferred-fallback") + + expect(applied).toEqual([{ taskId: "task-deferred-fallback", title: "Fix login bug" }]) + }) + + it("keeps the pending title when the retry fails", async () => { + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let attempt = 0 + const applied: ApplyTitleCall[] = [] + const provider = createTitleProvider({ + applySessionTitle: async (taskId, title) => { + attempt++ + if (attempt === 1) { + return "not-ready" + } + if (attempt === 2) { + throw new Error("transient write failure") + } + applied.push({ taskId, title }) + return "applied" + }, + }) + + await generateSessionTitle({ + provider, + taskId: "task-deferred-retry", + userMessage: "Fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + await flushPendingSessionTitle("task-deferred-retry") + expect(applied).toHaveLength(0) + + await flushPendingSessionTitle("task-deferred-retry") + expect(applied).toEqual([{ taskId: "task-deferred-retry", title: "Generated Title" }]) + }) + + it("drops the pending title of a disposed provider", async () => { + const recorder = createCompletePromptRecorder(async () => "Generated Title") + const applySpy = vi.fn().mockResolvedValue("not-ready" as const) + const provider = createTitleProvider({ applySessionTitle: applySpy }) + + await generateSessionTitle({ + provider, + taskId: "task-deferred-disposed", + userMessage: "Fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + Object.assign(provider, { isDisposed: true }) + await flushPendingSessionTitle("task-deferred-disposed") + + expect(applySpy).toHaveBeenCalledTimes(1) + }) + + it("drops the pending title when the task is cancelled", async () => { + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let historyItemPersisted = false + const applySpy = vi.fn().mockImplementation(async () => (historyItemPersisted ? "applied" : "not-ready")) + const provider = createTitleProvider({ applySessionTitle: applySpy }) + + await generateSessionTitle({ + provider, + taskId: "task-deferred-cancelled", + userMessage: "Fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + cancelSessionTitle("task-deferred-cancelled") + historyItemPersisted = true + await flushPendingSessionTitle("task-deferred-cancelled") + + expect(applySpy).toHaveBeenCalledTimes(1) + }) + + it("does nothing when no title is pending", async () => { + await expect(flushPendingSessionTitle("unknown-task")).resolves.toBeUndefined() + }) }) diff --git a/src/core/task/session-title.ts b/src/core/task/session-title.ts index 612746c0ca..d13080b359 100644 --- a/src/core/task/session-title.ts +++ b/src/core/task/session-title.ts @@ -16,8 +16,17 @@ Output only the title. ` -/** Outcome of `ClineProvider.applySessionTitle`. */ -export type SessionTitleApplyResult = "applied" | "conflict" +/** + * Outcome of `ClineProvider.applySessionTitle`. + * + * - `applied`: the generated title was written. + * - `conflict`: a manual title (or a manual clear) already owns the task, so nothing was written. + * - `not-ready`: the task's history item does not exist yet; the caller keeps the title for later. + */ +export type SessionTitleApplyResult = "applied" | "conflict" | "not-ready" + +/** Outcome of one title write, including the local give-up case. */ +type TitleWriteOutcome = SessionTitleApplyResult | "failed" /** Removes system-injected `` blocks and trims. */ export function sanitizeUserMessageForTitle(raw: string): string { @@ -75,17 +84,29 @@ export interface GenerateSessionTitleOptions { interface InFlightTitle { controller: AbortController userCancelled: boolean + /** Provider that owns this generation, so disposal can cancel and drain it. */ + provider: ClineProvider + /** Settles once the generation has fully finished, including its write attempts. */ + settled: Promise } /** In-flight title generation abort controllers by task ID. */ const inFlightTitles = new Map() +/** + * Generated titles whose history item did not exist yet, keyed by task ID. + * `Task` flushes these once the item is persisted, so a transient initial save + * failure cannot lose a title that was already generated. + */ +const pendingTitleWrites = new Map() + /** * Cancels any pending or in-flight title generation for a task. * Called when a task is manually renamed, cleared, or deleted so an auto-generated title * never overwrites the user's manual intent. */ export function cancelSessionTitle(taskId: string): void { + pendingTitleWrites.delete(taskId) const entry = inFlightTitles.get(taskId) if (entry) { entry.userCancelled = true @@ -93,6 +114,59 @@ export function cancelSessionTitle(taskId: string): void { } } +/** + * Cancels every pending or in-flight title generation owned by `provider`. + * Called from `ClineProvider#dispose()` so a request and its timer cannot + * outlive the provider and write through the disposed history store. + */ +export function cancelSessionTitlesForProvider(provider: ClineProvider): void { + for (const [taskId, entry] of inFlightTitles) { + if (entry.provider !== provider) { + continue + } + pendingTitleWrites.delete(taskId) + entry.userCancelled = true + entry.controller.abort() + } +} + +/** + * Waits for `provider`'s in-flight title generation to settle, so disposal does + * not release the history store underneath a title that is still being written. + */ +export async function drainSessionTitles(provider: ClineProvider): Promise { + const settled = Array.from(inFlightTitles.values()) + .filter((entry) => entry.provider === provider) + .map((entry) => entry.settled) + await Promise.allSettled(settled) +} + +/** + * Retries a title whose history item did not exist when the title was ready. + * + * The title write is deferred rather than dropped: `Task` calls this after a + * save that persisted the history item, so the task still receives the + * generated title (or the fallback) once its record exists. + */ +export async function flushPendingSessionTitle(taskId: string): Promise { + const pending = pendingTitleWrites.get(taskId) + if (!pending) { + return + } + if (pending.provider.isDisposed) { + pendingTitleWrites.delete(taskId) + return + } + try { + const result = await pending.provider.applySessionTitle(taskId, pending.title) + if (result === "applied" || result === "conflict") { + pendingTitleWrites.delete(taskId) + } + } catch { + // Keep the title for the next save instead of losing it to a transient write failure. + } +} + /** * Generates and applies a session title from the root user message. * @@ -103,6 +177,12 @@ export function cancelSessionTitle(taskId: string): void { export async function generateSessionTitle(options: GenerateSessionTitleOptions): Promise { const { provider, taskId, userMessage, apiConfiguration, completePrompt = singleCompletionHandler } = options + // A disposed provider has no history store to write to; the task started after disposal + // must not launch a request that would outlive it. + if (provider.isDisposed) { + return + } + const sanitized = sanitizeUserMessageForTitle(userMessage) if (!sanitized) { return @@ -111,9 +191,15 @@ export async function generateSessionTitle(options: GenerateSessionTitleOptions) if (inFlightTitles.has(taskId)) { return } + let settle!: () => void + const settled = new Promise((resolve) => { + settle = resolve + }) const inFlight: InFlightTitle = { controller: new AbortController(), userCancelled: false, + provider, + settled, } inFlightTitles.set(taskId, inFlight) @@ -137,15 +223,27 @@ export async function generateSessionTitle(options: GenerateSessionTitleOptions) } } } - let applied = false if (title) { - applied = await applyTitleWithRetry(provider, taskId, title, () => inFlight.userCancelled) + const outcome = await applyTitleWithRetry(provider, taskId, title, () => inFlight.userCancelled) + if (outcome === "not-ready") { + // The history item is not persisted yet; `Task` retries this write once it is. + pendingTitleWrites.set(taskId, { provider, title }) + return + } + // "applied" means the title is durable, and "conflict" means manual intent owns the + // task; neither may be followed by the fallback write. + if (outcome !== "failed") { + return + } } - if (!applied && !inFlight.userCancelled) { - const fallback = fallbackTitle(sanitized) - // Stryker disable next-line ConditionalExpression,LogicalOperator: defense-in-depth, fallback is always truthy for non-empty sanitized messages and distinct from failed title - if (fallback && fallback !== title) { - await applyTitleWithRetry(provider, taskId, fallback, () => inFlight.userCancelled) + // `applyTitleWithRetry` owns the cancellation check, so a generation cancelled during + // the LLM title write never reaches the fallback write. + const fallback = fallbackTitle(sanitized) + // Stryker disable next-line ConditionalExpression,LogicalOperator: defense-in-depth, fallback is always truthy for non-empty sanitized messages and distinct from failed title + if (fallback && fallback !== title) { + const outcome = await applyTitleWithRetry(provider, taskId, fallback, () => inFlight.userCancelled) + if (outcome === "not-ready") { + pendingTitleWrites.set(taskId, { provider, title: fallback }) } } } catch (error) { @@ -157,6 +255,7 @@ export async function generateSessionTitle(options: GenerateSessionTitleOptions) } } finally { inFlightTitles.delete(taskId) + settle() } } @@ -208,16 +307,17 @@ async function applyTitleWithRetry( taskId: string, title: string, isCancelled: () => boolean, -): Promise { +): Promise { for (let attempt = 0; ; attempt++) { // Stryker disable next-line BooleanLiteral,ConditionalExpression: defense-in-depth, caller checks cancellation before initiating write if (isCancelled()) { - // Stryker disable next-line BooleanLiteral: return value indicating write was not applied - return false + // Stryker disable next-line StringLiteral: return value indicating write was not applied + return "failed" } try { - await provider.applySessionTitle(taskId, title) - return true + // "not-ready" is a result, not an error: the history item does not exist yet, so an + // immediate retry cannot help. The caller keeps the title for a later apply. + return await provider.applySessionTitle(taskId, title) } catch (error) { if (isCancelled() || attempt >= TITLE_WRITE_MAX_RETRIES) { if (!isCancelled()) { @@ -226,7 +326,7 @@ async function applyTitleWithRetry( error, ) } - return false + return "failed" } } } diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 46e0e35ce1..cb53ea189b 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -76,7 +76,12 @@ import { providerIdentifiers, } from "@roo-code/types" import { RateLimitClock, createRateLimitClock } from "../task/RateLimitClock" -import { cancelSessionTitle, type SessionTitleApplyResult } from "../task/session-title" +import { + cancelSessionTitle, + cancelSessionTitlesForProvider, + drainSessionTitles, + type SessionTitleApplyResult, +} from "../task/session-title" import { TaskRegistry } from "../task/TaskRegistry" import { TaskScheduler } from "../task/TaskScheduler" import { @@ -1001,6 +1006,11 @@ export class ClineProvider this.skillsManager = undefined await this.marketplaceManager?.cleanup() this.customModesManager?.dispose() + // Cancel and drain title generation this provider owns before releasing the + // history store, so an in-flight request or timer cannot write through the + // disposed store afterwards. + cancelSessionTitlesForProvider(this) + await drainSessionTitles(this) this.taskHistoryStore.dispose() this.log("Disposed all disposables") ClineProvider.activeInstances.delete(this) @@ -1016,6 +1026,15 @@ export class ClineProvider return this.view?.visible === true } + /** + * True once disposal has started. Work started after this point (such as a + * session title for a task still registering) must not touch the provider + * or its history store. + */ + public get isDisposed(): boolean { + return this._disposed + } + public static getVisibleInstance(): ClineProvider | undefined { return findLast(Array.from(this.activeInstances), (instance) => instance.isViewVisible) } @@ -2822,6 +2841,9 @@ export class ClineProvider await this.taskHistoryStore.atomicReadAndUpdate(id, (historyItem) => ({ ...historyItem, title: trimmed ? trimmed : undefined, + // Persist the clear intent: an empty title is otherwise indistinguishable from a + // task that never had one, which lets another host's pending auto title win. + titleCleared: trimmed ? undefined : true, })) this.recentTasksCache = undefined @@ -2848,7 +2870,14 @@ export class ClineProvider } const result = await this.taskHistoryStore.setTitleIfUnset(id, trimmed) + if (result === "missing") { + // The history item has not been persisted yet; the caller keeps the title and retries. + return "not-ready" + } if (result === "conflict") { + // The store refreshed its cache with the winning manual record. Push it so this + // window shows the manual title instead of waiting for another history refresh. + await this.broadcastCachedTaskHistoryItem(id) return "conflict" } @@ -2865,6 +2894,21 @@ export class ClineProvider return "applied" } + /** + * Broadcasts the cached record for a task when one exists. Used when a manual + * write won a race and the webview still shows the pre-race record. + */ + private async broadcastCachedTaskHistoryItem(id: string): Promise { + if (!this.isViewLaunched) { + return + } + const cachedItem = this.taskHistoryStore.get(id) + if (cachedItem) { + this.recentTasksCache = undefined + await this.postMessageToWebview({ type: "taskHistoryItemUpdated", taskHistoryItem: cachedItem }) + } + } + /* Condenses a task's message history to use fewer tokens. */ async condenseTaskContext(taskId: string) { const task = this.taskRegistry.getById(taskId) @@ -3686,7 +3730,8 @@ export class ClineProvider // Drop the snapshot's title so upsertCore preserves the value held under the store lock // instead of treating a stale title as a changed field and reverting a concurrent rename. - const { title: _title, ...itemWithoutTitle } = item + // The manual-clear intent is owned by `renameTaskWithId` for the same reason. + const { title: _title, titleCleared: _titleCleared, ...itemWithoutTitle } = item const history = await this.taskHistoryStore.upsert(itemWithoutTitle) this.recentTasksCache = undefined diff --git a/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts b/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts index a0910800e0..75742438af 100644 --- a/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts +++ b/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts @@ -764,8 +764,35 @@ describe("ClineProvider Task History Synchronization", () => { expect(getStored("task-auto-blank")?.title).toBeUndefined() }) - it("throws for an unknown task so the caller can give up", async () => { - await expect(provider.applySessionTitle("missing-task", "Auto Title")).rejects.toThrow("not found in cache") + it("returns not-ready for a task whose history item does not exist yet", async () => { + const result = await provider.applySessionTitle("missing-task", "Auto Title") + + expect(result).toBe("not-ready") + expect(getStored("missing-task")).toBeUndefined() + }) + + it("broadcasts the winning manual title when the automatic title conflicts", async () => { + await provider.resolveWebviewView(mockWebviewView) + provider.isViewLaunched = true + await provider.updateTaskHistory(createHistoryItem({ id: "task-auto-lost", task: "Original prompt" }), { + broadcast: false, + }) + // Another window renames the task while this window's title is generating. + await provider.taskHistoryStore.atomicReadAndUpdate("task-auto-lost", (item) => ({ + ...item, + title: "Manual title", + })) + mockPostMessage.mockClear() + + const result = await provider.applySessionTitle("task-auto-lost", "Auto Title") + + expect(result).toBe("conflict") + expect(getStored("task-auto-lost")).toMatchObject({ title: "Manual title" }) + const updates = findCallsByType(mockPostMessage.mock.calls, "taskHistoryItemUpdated") + expect(updates.at(-1)?.[0].taskHistoryItem).toMatchObject({ + id: "task-auto-lost", + title: "Manual title", + }) }) it("keeps the auto title when the task later saves metadata without one", async () => { @@ -806,15 +833,73 @@ describe("ClineProvider Task History Synchronization", () => { expect(findCallsByType(mockPostMessage.mock.calls, "taskHistoryItemUpdated")).toHaveLength(0) }) - it("cancels pending session title generation when deleting a task", async () => { + it("cancels pending generation on dispose so no title is written afterwards", async () => { + await provider.updateTaskHistory(createHistoryItem({ id: "task-dispose", task: "Original prompt" }), { + broadcast: false, + }) + let resolveCompletion!: (title: string) => void + const completionPromise = new Promise((resolve) => { + resolveCompletion = resolve + }) + const generation = generateSessionTitle({ + provider, + taskId: "task-dispose", + userMessage: "Original prompt", + apiConfiguration: { + apiProvider: providerIdentifiers.anthropic, + apiModelId: "claude-3-5-sonnet-20241022", + apiKey: "test-api-key", + }, + completePrompt: async () => completionPromise, + }) + + // Disposal cancels and drains the generation before the history store is released. + await provider.dispose() + expect(provider.isDisposed).toBe(true) + + resolveCompletion("Late Title") + await generation + + expect(getStored("task-dispose")?.title).toBeUndefined() + }) + + it("cancels pending generation and writes no late title when deleting a task", async () => { + await provider.updateTaskHistory( + createHistoryItem({ id: "task-cancel-on-delete", task: "Original prompt" }), + { + broadcast: false, + }, + ) + let resolveCompletion!: (title: string) => void + const completionPromise = new Promise((resolve) => { + resolveCompletion = resolve + }) + const generation = generateSessionTitle({ + provider, + taskId: "task-cancel-on-delete", + userMessage: "Original prompt", + apiConfiguration: { + apiProvider: providerIdentifiers.anthropic, + apiModelId: "claude-3-5-sonnet-20241022", + apiKey: "test-api-key", + }, + completePrompt: async () => completionPromise, + }) + const cancelSpy = vi.spyOn(sessionTitleModule, "cancelSessionTitle") try { await provider.deleteTaskWithId("task-cancel-on-delete") - } catch { - // Task doesn't exist on disk, delete will catch/proceed + expect(cancelSpy).toHaveBeenCalledWith("task-cancel-on-delete") + // The deletion succeeded, so the task is gone from the store. + expect(getStored("task-cancel-on-delete")).toBeUndefined() + + // A completion that settles after the deletion must not resurrect the task or its title. + resolveCompletion("Late Title") + await generation + expect(getStored("task-cancel-on-delete")).toBeUndefined() + } finally { + cancelSpy.mockRestore() } - expect(cancelSpy).toHaveBeenCalledWith("task-cancel-on-delete") - cancelSpy.mockRestore() }) }) From 7b7f795ee54b49d074bd90c429e8f88f73b27551 Mon Sep 17 00:00:00 2001 From: daewoongoh Date: Sat, 7 Nov 2026 23:25:20 +0900 Subject: [PATCH 3/7] test(task): strengthen title request-option coverage and keep unknown disk fields --- .../providers/__tests__/openrouter.spec.ts | 48 +++++++++++++++++++ src/api/providers/__tests__/poe.spec.ts | 28 ++++++++--- src/core/task-persistence/TaskHistoryStore.ts | 4 +- .../TaskHistoryStore.crossInstance.spec.ts | 16 +++++++ 4 files changed, 88 insertions(+), 8 deletions(-) diff --git a/src/api/providers/__tests__/openrouter.spec.ts b/src/api/providers/__tests__/openrouter.spec.ts index 5057fc09cd..ab189baebd 100644 --- a/src/api/providers/__tests__/openrouter.spec.ts +++ b/src/api/providers/__tests__/openrouter.spec.ts @@ -996,6 +996,54 @@ describe("OpenRouterHandler", () => { }) }) + it("sends no request options when the caller passes none", async () => { + const handler = new OpenRouterHandler( + makeApiHandlerOptions({ + openRouterApiKey: "test-key", + openRouterModelId: "openai/gpt-4o", + }), + ) + const mockCreate = vitest.fn().mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + installMockClient(mockCreate) + + await handler.completePrompt("test prompt") + + expect(mockCreate).toHaveBeenLastCalledWith(expect.objectContaining({ model: "openai/gpt-4o" }), undefined) + }) + + it("forwards only the abort signal when no deadline is given", async () => { + const handler = new OpenRouterHandler( + makeApiHandlerOptions({ + openRouterApiKey: "test-key", + openRouterModelId: "openai/gpt-4o", + }), + ) + const mockCreate = vitest.fn().mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + installMockClient(mockCreate) + const { signal } = new AbortController() + + await handler.completePrompt("test prompt", { abortSignal: signal }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.objectContaining({ model: "openai/gpt-4o" }), { signal }) + }) + + it("forwards only the deadline when no signal is given", async () => { + const handler = new OpenRouterHandler( + makeApiHandlerOptions({ + openRouterApiKey: "test-key", + openRouterModelId: "openai/gpt-4o", + }), + ) + const mockCreate = vitest.fn().mockResolvedValue({ choices: [{ message: { content: "ok" } }] }) + installMockClient(mockCreate) + + await handler.completePrompt("test prompt", { timeoutMs: 5_000 }) + + expect(mockCreate).toHaveBeenLastCalledWith(expect.objectContaining({ model: "openai/gpt-4o" }), { + timeout: 5_000, + }) + }) + it("does not start a request when the caller already aborted", async () => { const handler = new OpenRouterHandler(mockOptions) const mockCreate = vitest.fn() diff --git a/src/api/providers/__tests__/poe.spec.ts b/src/api/providers/__tests__/poe.spec.ts index 31d62281a4..b779d7fe97 100644 --- a/src/api/providers/__tests__/poe.spec.ts +++ b/src/api/providers/__tests__/poe.spec.ts @@ -402,17 +402,31 @@ describe("PoeHandler", () => { }) describe("completePrompt", () => { - it("forwards the abort signal and deadline to the generation", async () => { + it("links the caller's stop into the merged generation signal", async () => { const handler = new PoeHandler({ poeApiKey: "key", apiModelId: "openai/gpt-4o" }) mockGenerateText.mockResolvedValue({ text: "generated response" }) - const { signal } = new AbortController() + const abortController = new AbortController() - await handler.completePrompt("complete this", { abortSignal: signal, timeoutMs: 5_000 }) + await handler.completePrompt("complete this", { abortSignal: abortController.signal, timeoutMs: 5_000 }) - const options = mockGenerateText.mock.calls[0][0] - // The deadline is merged into a per-request signal that still follows the caller's stop. - expect(options.abortSignal).not.toBe(signal) - expect(options.abortSignal?.aborted).toBe(false) + const mergedSignal = mockGenerateText.mock.calls[0][0].abortSignal + // A per-request signal carries both the caller's stop and the deadline. + expect(mergedSignal).not.toBe(abortController.signal) + expect(mergedSignal?.aborted).toBe(false) + + abortController.abort() + expect(mergedSignal?.aborted).toBe(true) + }) + + it("aborts the merged generation signal when the deadline passes", async () => { + const handler = new PoeHandler({ poeApiKey: "key", apiModelId: "openai/gpt-4o" }) + mockGenerateText.mockResolvedValue({ text: "generated response" }) + + await handler.completePrompt("complete this", { timeoutMs: 200 }) + + const mergedSignal = mockGenerateText.mock.calls[0][0].abortSignal + expect(mergedSignal?.aborted).toBe(false) + await vi.waitFor(() => expect(mergedSignal?.aborted).toBe(true)) }) it("does not start a generation when the caller already aborted", async () => { diff --git a/src/core/task-persistence/TaskHistoryStore.ts b/src/core/task-persistence/TaskHistoryStore.ts index 2c5464fad4..06a6c07fc8 100644 --- a/src/core/task-persistence/TaskHistoryStore.ts +++ b/src/core/task-persistence/TaskHistoryStore.ts @@ -1057,7 +1057,9 @@ export class TaskHistoryStore { // Stryker disable next-line StringLiteral: internal discrimination literal between missing and titled reasons throw new TitleWriteSkipped("titled") } - written = mergeHistoryDelta(parsed.data, incoming as HistoryItem, { title }) + // Merge into the raw record, not the parsed copy: Zod strips unknown keys, and a + // field written by another extension version must not be deleted by a title write. + written = mergeHistoryDelta(existing, incoming as HistoryItem, { title }) return written }, }) diff --git a/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts b/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts index d4a8b8dada..ff4259cbb9 100644 --- a/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts +++ b/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts @@ -208,6 +208,22 @@ describe("TaskHistoryStore cross-instance safety", () => { storeWithOnWrite.dispose() }) + it("preserves fields written by another extension version when applying a title", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "title-extra-fields", task: "Prompt" })) + const filePath = path.join(tmpDir, "tasks", "title-extra-fields", GlobalFileNames.historyItem) + const onDisk = JSON.parse(await fs.readFile(filePath, "utf8")) + // A field this version does not know about, written by another build. + onDisk.futureField = "kept" + await fs.writeFile(filePath, JSON.stringify(onDisk), "utf8") + + await expect(storeA.setTitleIfUnset("title-extra-fields", "Auto title")).resolves.toBe("applied") + + const final = JSON.parse(await fs.readFile(filePath, "utf8")) + expect(final.title).toBe("Auto title") + expect(final.futureField).toBe("kept") + }) + it("re-throws unexpected errors from safeWriteJson", async () => { await storeA.initialize() await storeA.upsert(makeHistoryItem({ id: "title-io-err", task: "Prompt" })) From 7ad2959b704f312da7c65bfc6d7c1b50658b1f61 Mon Sep 17 00:00:00 2001 From: daewoongoh Date: Sat, 7 Nov 2026 23:44:41 +0900 Subject: [PATCH 4/7] test(poe): observe cancellation while the generation is still in flight --- src/api/providers/__tests__/poe.spec.ts | 53 ++++++++++++++++++------- 1 file changed, 38 insertions(+), 15 deletions(-) diff --git a/src/api/providers/__tests__/poe.spec.ts b/src/api/providers/__tests__/poe.spec.ts index b779d7fe97..9b3eef5c33 100644 --- a/src/api/providers/__tests__/poe.spec.ts +++ b/src/api/providers/__tests__/poe.spec.ts @@ -401,32 +401,55 @@ describe("PoeHandler", () => { }) }) + // Holds generateText pending so a test can observe the request signal while it is in + // flight; `release` settles the call. + function holdGeneration(generateTextMock: { + mockImplementation: (fn: (options: { abortSignal?: AbortSignal }) => Promise<{ text: string }>) => unknown + }) { + let captured: { abortSignal?: AbortSignal } | undefined + let release: () => void = () => {} + generateTextMock.mockImplementation( + (options) => + new Promise<{ text: string }>((resolve) => { + captured = options + release = () => resolve({ text: "generated response" }) + }), + ) + return { release: () => release(), abortSignal: () => captured?.abortSignal } + } + describe("completePrompt", () => { - it("links the caller's stop into the merged generation signal", async () => { + it("links the caller's stop into the in-flight generation signal", async () => { const handler = new PoeHandler({ poeApiKey: "key", apiModelId: "openai/gpt-4o" }) - mockGenerateText.mockResolvedValue({ text: "generated response" }) + const { release, abortSignal } = holdGeneration(mockGenerateText) const abortController = new AbortController() - await handler.completePrompt("complete this", { abortSignal: abortController.signal, timeoutMs: 5_000 }) - - const mergedSignal = mockGenerateText.mock.calls[0][0].abortSignal - // A per-request signal carries both the caller's stop and the deadline. - expect(mergedSignal).not.toBe(abortController.signal) - expect(mergedSignal?.aborted).toBe(false) + const generation = handler.completePrompt("complete this", { + abortSignal: abortController.signal, + timeoutMs: 5_000, + }) + await vi.waitFor(() => expect(abortSignal()).toBeDefined()) + // A per-request signal carries the caller's stop while the request is still running. + expect(abortSignal()?.aborted).toBe(false) abortController.abort() - expect(mergedSignal?.aborted).toBe(true) + expect(abortSignal()?.aborted).toBe(true) + // The generation only settles when the provider call is released. + release() + await expect(generation).resolves.toBe("generated response") }) - it("aborts the merged generation signal when the deadline passes", async () => { + it("aborts the in-flight generation signal when the deadline passes", async () => { const handler = new PoeHandler({ poeApiKey: "key", apiModelId: "openai/gpt-4o" }) - mockGenerateText.mockResolvedValue({ text: "generated response" }) + const { release, abortSignal } = holdGeneration(mockGenerateText) - await handler.completePrompt("complete this", { timeoutMs: 200 }) + const generation = handler.completePrompt("complete this", { timeoutMs: 200 }) + await vi.waitFor(() => expect(abortSignal()).toBeDefined()) - const mergedSignal = mockGenerateText.mock.calls[0][0].abortSignal - expect(mergedSignal?.aborted).toBe(false) - await vi.waitFor(() => expect(mergedSignal?.aborted).toBe(true)) + // The deadline aborts the request while the provider call is still pending. + await vi.waitFor(() => expect(abortSignal()?.aborted).toBe(true)) + release() + await expect(generation).resolves.toBe("generated response") }) it("does not start a generation when the caller already aborted", async () => { From 1b741fb1844a7633504e1b9cba4a5e01b7ef8130 Mon Sep 17 00:00:00 2001 From: daewoongoh Date: Sun, 8 Nov 2026 00:25:58 +0900 Subject: [PATCH 5/7] fix(task): make manual title writes disk-authoritative and drain deferred title writes --- src/api/providers/__tests__/kimi-code.spec.ts | 27 ++++ src/api/providers/kimi-code.ts | 20 ++- src/core/task-persistence/TaskHistoryStore.ts | 58 +++++++++ .../TaskHistoryStore.crossInstance.spec.ts | 120 ++++++++++++++++++ src/core/task/__tests__/session-title.spec.ts | 65 ++++++++++ src/core/task/session-title.ts | 26 +++- src/core/webview/ClineProvider.ts | 21 ++- .../ClineProvider.taskHistory.spec.ts | 17 +++ 8 files changed, 330 insertions(+), 24 deletions(-) diff --git a/src/api/providers/__tests__/kimi-code.spec.ts b/src/api/providers/__tests__/kimi-code.spec.ts index c8e1130b30..2fb77296e8 100644 --- a/src/api/providers/__tests__/kimi-code.spec.ts +++ b/src/api/providers/__tests__/kimi-code.spec.ts @@ -131,6 +131,33 @@ describe("KimiCodeHandler", () => { expect(createCompletion).toHaveBeenCalledTimes(2) }) + it("forwards the caller's cancellation and deadline to each completion attempt", async () => { + const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" }) + const unauthorized = Object.assign(new Error("Unauthorized"), { status: 401 }) + const client = ( + handler as unknown as { client: { chat: { completions: { create: ReturnType } } } } + ).client + const createCompletion = vi + .spyOn(client.chat.completions, "create") + .mockRejectedValueOnce(unauthorized) + .mockResolvedValueOnce({ choices: [{ message: { content: "retried" } }] }) + const { signal } = new AbortController() + + await expect(handler.completePrompt("test", { abortSignal: signal, timeoutMs: 5_000 })).resolves.toBe("retried") + + // The retry after the forced refresh keeps the same options. + expect(createCompletion).toHaveBeenNthCalledWith(2, expect.anything(), { signal, timeout: 5_000 }) + }) + + it("does not start preparation when the caller already aborted", async () => { + const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" }) + const controller = new AbortController() + controller.abort() + + await expect(handler.completePrompt("test", { abortSignal: controller.signal })).rejects.toThrow() + expect(mockGetAccessToken).not.toHaveBeenCalled() + }) + it("does not retry on 401 when using API key auth", async () => { const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" }) vi.spyOn(globalThis, "fetch").mockResolvedValueOnce(new Response(null, { status: 401 })) diff --git a/src/api/providers/kimi-code.ts b/src/api/providers/kimi-code.ts index 54e48b29c0..371fa79a83 100644 --- a/src/api/providers/kimi-code.ts +++ b/src/api/providers/kimi-code.ts @@ -13,7 +13,7 @@ import { import type { ApiHandlerOptions, GetModelsOptions } from "../../shared/api" import { kimiCodeOAuthManager } from "../../integrations/kimi-code/oauth" -import type { ApiHandlerCreateMessageMetadata, ModelCacheScope } from "../index" +import type { ApiHandlerCreateMessageMetadata, CompletePromptOptions, ModelCacheScope } from "../index" import type { ApiStream } from "../transform/stream" import { getModelParams } from "../transform/model-params" @@ -21,6 +21,7 @@ import { OpenAiHandler } from "./openai" import { NOT_PROVIDED } from "./constants" import { getModels } from "./fetchers/modelCache" import { getKimiCodeModels } from "./fetchers/kimi-code" +import { throwIfAborted } from "./utils/abort-signal" const OAUTH_AUTH_METHOD: KimiCodeAuthMethod = "oauth" const API_KEY_AUTH_METHOD: KimiCodeAuthMethod = "api-key" @@ -72,8 +73,11 @@ export class KimiCodeHandler extends OpenAiHandler { return token } - private async prepareRequest(forceRefresh = false): Promise { + private async prepareRequest(forceRefresh = false, abortSignal?: AbortSignal): Promise { + throwIfAborted(abortSignal) const accessToken = await this.resolveAccessToken(forceRefresh) + // Token resolution (including an OAuth refresh) may outlive the caller's cancellation. + throwIfAborted(abortSignal) this.client.apiKey = accessToken if (!this.modelDiscoveryAttempted) { this.modelDiscoveryAttempted = true @@ -107,14 +111,16 @@ export class KimiCodeHandler extends OpenAiHandler { } } - override async completePrompt(prompt: string): Promise { - await this.prepareRequest() + override async completePrompt(prompt: string, options?: CompletePromptOptions): Promise { + // Preparation and both attempts run under the caller's cancellation and deadline, so a + // superseded title request stops instead of finishing its token and model work. + await this.prepareRequest(false, options?.abortSignal) try { - return await super.completePrompt(prompt) + return await super.completePrompt(prompt, options) } catch (error) { if (getHttpStatus(error) !== 401 || !this.canRefreshOAuth()) throw error - await this.prepareRequest(true) - return super.completePrompt(prompt) + await this.prepareRequest(true, options?.abortSignal) + return super.completePrompt(prompt, options) } } diff --git a/src/core/task-persistence/TaskHistoryStore.ts b/src/core/task-persistence/TaskHistoryStore.ts index 06a6c07fc8..d252390b62 100644 --- a/src/core/task-persistence/TaskHistoryStore.ts +++ b/src/core/task-persistence/TaskHistoryStore.ts @@ -1089,6 +1089,64 @@ export class TaskHistoryStore { }) } + /** + * Apply a user's manual title (or clear it) to the record persisted on disk. + * + * The write runs inside the per-file lock and merges into the disk record, so a + * clear removes a title another instance wrote even when this cache never saw + * it, and a manual title is never reverted by a stale cached snapshot. + * + * @returns The written item, or undefined when the record was deleted or is unusable + * on disk (the stale cache entry is dropped and the disk record is left untouched). + */ + public setTitleManually(taskId: string, title: string | undefined): Promise { + return this.withLock(async () => { + const cached = this.cache.get(taskId) + if (!cached) { + return undefined + } + + // A clear is persisted as an explicit intent, so a later automatic title cannot + // mistake "no title" for "never titled". + const delta: Partial = title ? { title } : { title: undefined, titleCleared: true } + const filePath = await this.getTaskFilePath(taskId) + let written: HistoryItem = { ...cached, ...delta } + try { + await safeWriteJson(filePath, written, { + merge: (existing, incoming) => { + if (!existing || typeof existing !== "object" || !("id" in existing)) { + throw new TitleWriteSkipped("missing") + } + // The manual write targets the record that is actually persisted: a cached + // snapshot does not know which of the fields on disk are newer. + const parsed = historyItemSchema.safeParse(existing) + if (!parsed.success) { + throw new TitleWriteSkipped("invalid") + } + if (parsed.data.id !== taskId) { + throw new TitleWriteSkipped("mismatched") + } + written = mergeHistoryDelta(existing, incoming as HistoryItem, delta) + return written + }, + }) + } catch (error) { + if (!(error instanceof TitleWriteSkipped)) { + throw error + } + this.cache.delete(taskId) + this.taskFileMtimes.delete(taskId) + return undefined + } + + this.cache.set(taskId, written) + if (this.onWrite) { + await this.onWrite(this.getAll()) + } + return written + }) + } + /** * Read a HistoryItem from the in-memory cache and write back an updated version, * all within a single lock acquisition so no concurrent writer can interleave diff --git a/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts b/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts index ff4259cbb9..4c1ebb35f5 100644 --- a/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts +++ b/src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts @@ -236,6 +236,126 @@ describe("TaskHistoryStore cross-instance safety", () => { }) }) + describe("setTitleManually", () => { + it("removes a title another instance wrote even when this cache has none", async () => { + await storeA.initialize() + await storeB.initialize() + await storeA.upsert(makeHistoryItem({ id: "manual-clear", task: "Prompt" })) + await storeB.reconcile() + + // Instance B renames the task on disk; instance A's cache still has no title. + await storeB.atomicReadAndUpdate("manual-clear", (item) => ({ ...item, title: "Manual title" })) + expect(storeA.get("manual-clear")?.title).toBeUndefined() + + // A clear from instance A must remove the disk title, not just the cached one. + await expect(storeA.setTitleManually("manual-clear", undefined)).resolves.toMatchObject({ + id: "manual-clear", + titleCleared: true, + }) + + const filePath = path.join(tmpDir, "tasks", "manual-clear", GlobalFileNames.historyItem) + const onDisk = JSON.parse(await fs.readFile(filePath, "utf8")) + expect(onDisk.title).toBeUndefined() + expect(onDisk.titleCleared).toBe(true) + }) + + it("keeps a manual title that commits while the automatic title is generating", async () => { + await storeA.initialize() + await storeB.initialize() + await storeA.upsert(makeHistoryItem({ id: "manual-wins", task: "Prompt" })) + await storeB.reconcile() + + await storeB.setTitleManually("manual-wins", "Manual title") + // The generating host still has the untitled record cached. + expect(storeA.get("manual-wins")?.title).toBeUndefined() + + await expect(storeA.setTitleIfUnset("manual-wins", "Auto title")).resolves.toBe("conflict") + expect(storeA.get("manual-wins")?.title).toBe("Manual title") + }) + + it("returns undefined for a task that is not in the cache", async () => { + await storeA.initialize() + await expect(storeA.setTitleManually("missing-id", "Manual title")).resolves.toBeUndefined() + }) + + it("drops the cache entry for a malformed disk record without rewriting it", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "manual-malformed", task: "Prompt" })) + const filePath = path.join(tmpDir, "tasks", "manual-malformed", GlobalFileNames.historyItem) + const malformed = JSON.stringify({ id: 123 }) + await fs.writeFile(filePath, malformed, "utf8") + + await expect(storeA.setTitleManually("manual-malformed", "Manual title")).resolves.toBeUndefined() + + expect(storeA.get("manual-malformed")).toBeUndefined() + expect(await fs.readFile(filePath, "utf8")).toBe(malformed) + }) + + it("drops the cache entry when the disk record belongs to another task", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "manual-mismatched", task: "Prompt" })) + const filePath = path.join(tmpDir, "tasks", "manual-mismatched", GlobalFileNames.historyItem) + const otherTask = JSON.stringify(makeHistoryItem({ id: "other-task" }), null, "\t") + await fs.writeFile(filePath, otherTask, "utf8") + + await expect(storeA.setTitleManually("manual-mismatched", "Manual title")).resolves.toBeUndefined() + + expect(storeA.get("manual-mismatched")).toBeUndefined() + expect(await fs.readFile(filePath, "utf8")).toBe(otherTask) + }) + + it("re-throws unexpected errors from safeWriteJson", async () => { + await storeA.initialize() + await storeA.upsert(makeHistoryItem({ id: "manual-io-err", task: "Prompt" })) + const { safeWriteJson } = await import("../../../utils/safeWriteJson") + vi.mocked(safeWriteJson).mockRejectedValueOnce(new Error("disk IO error")) + + await expect(storeA.setTitleManually("manual-io-err", "Manual title")).rejects.toThrow("disk IO error") + }) + + it("notifies onWrite when a manual title is applied", async () => { + const onWriteSpy = vi.fn().mockResolvedValue(undefined) + const storeWithOnWrite = new TaskHistoryStore(tmpDir, { onWrite: onWriteSpy }) + await storeWithOnWrite.initialize() + await storeWithOnWrite.upsert(makeHistoryItem({ id: "manual-onwrite", task: "Prompt" })) + onWriteSpy.mockClear() + + await expect(storeWithOnWrite.setTitleManually("manual-onwrite", "Notified title")).resolves.toMatchObject({ + id: "manual-onwrite", + title: "Notified title", + }) + expect(onWriteSpy).toHaveBeenCalledWith( + expect.arrayContaining([expect.objectContaining({ id: "manual-onwrite", title: "Notified title" })]), + ) + storeWithOnWrite.dispose() + }) + + it("writes a title on a record this cache has never seen the latest fields of", async () => { + await storeA.initialize() + await storeB.initialize() + await storeA.upsert(makeHistoryItem({ id: "manual-set", task: "Prompt", tokensIn: 1 })) + await storeB.reconcile() + + await storeB.atomicReadAndUpdate("manual-set", (item) => ({ ...item, tokensIn: 99 })) + await storeA.setTitleManually("manual-set", "Manual title") + + const filePath = path.join(tmpDir, "tasks", "manual-set", GlobalFileNames.historyItem) + const onDisk = JSON.parse(await fs.readFile(filePath, "utf8")) + expect(onDisk).toMatchObject({ title: "Manual title", tokensIn: 99 }) + }) + + it("drops the task when another instance deleted it", async () => { + await storeA.initialize() + await storeB.initialize() + await storeA.upsert(makeHistoryItem({ id: "manual-deleted", task: "Prompt" })) + await storeB.reconcile() + await storeB.delete("manual-deleted") + + await expect(storeA.setTitleManually("manual-deleted", "Manual title")).resolves.toBeUndefined() + expect(storeA.get("manual-deleted")).toBeUndefined() + }) + }) + it("two instances can write different tasks without conflict", async () => { await storeA.initialize() await storeB.initialize() diff --git a/src/core/task/__tests__/session-title.spec.ts b/src/core/task/__tests__/session-title.spec.ts index 3cde79d675..9bd09d40f9 100644 --- a/src/core/task/__tests__/session-title.spec.ts +++ b/src/core/task/__tests__/session-title.spec.ts @@ -856,6 +856,71 @@ describe("flushPendingSessionTitle", () => { expect(applySpy).toHaveBeenCalledTimes(1) }) + it("does not restore a pending title when the generation is cancelled before the retry", async () => { + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let historyItemPersisted = false + const applySpy = vi.fn().mockImplementation(async () => { + // The task is deleted or renamed while the write is in flight. + cancelSessionTitle("task-not-ready-cancelled") + return historyItemPersisted ? "applied" : "not-ready" + }) + const provider = createTitleProvider({ applySessionTitle: applySpy }) + + await generateSessionTitle({ + provider, + taskId: "task-not-ready-cancelled", + userMessage: "Fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + // The cancelled generation must not queue a title that a later save would apply. + historyItemPersisted = true + await flushPendingSessionTitle("task-not-ready-cancelled") + expect(applySpy).toHaveBeenCalledTimes(1) + }) + + it("drains a deferred title write before the store is released", async () => { + let releaseWrite!: () => void + const writeGate = new Promise((resolve) => { + releaseWrite = resolve + }) + let historyItemPersisted = false + const provider = createTitleProvider({ + applySessionTitle: async () => { + if (!historyItemPersisted) { + return "not-ready" + } + await writeGate + return "applied" + }, + }) + const recorder = createCompletePromptRecorder(async () => "Generated Title") + + await generateSessionTitle({ + provider, + taskId: "task-drain-flush", + userMessage: "Fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + historyItemPersisted = true + void flushPendingSessionTitle("task-drain-flush") + + let drained = false + const drain = drainSessionTitles(provider).then(() => { + drained = true + }) + await new Promise((resolve) => setTimeout(resolve, 10)) + // The drain waits for the deferred write instead of releasing the store underneath it. + expect(drained).toBe(false) + + releaseWrite() + await drain + expect(drained).toBe(true) + }) + it("does nothing when no title is pending", async () => { await expect(flushPendingSessionTitle("unknown-task")).resolves.toBeUndefined() }) diff --git a/src/core/task/session-title.ts b/src/core/task/session-title.ts index d13080b359..b7c36b63b1 100644 --- a/src/core/task/session-title.ts +++ b/src/core/task/session-title.ts @@ -100,6 +100,9 @@ const inFlightTitles = new Map() */ const pendingTitleWrites = new Map() +/** In-flight deferred-title writes, so disposal can wait for them to finish. */ +const inFlightFlushes = new Set<{ provider: ClineProvider; settled: Promise }>() + /** * Cancels any pending or in-flight title generation for a task. * Called when a task is manually renamed, cleared, or deleted so an auto-generated title @@ -135,10 +138,13 @@ export function cancelSessionTitlesForProvider(provider: ClineProvider): void { * not release the history store underneath a title that is still being written. */ export async function drainSessionTitles(provider: ClineProvider): Promise { - const settled = Array.from(inFlightTitles.values()) + const generations = Array.from(inFlightTitles.values()) + .filter((entry) => entry.provider === provider) + .map((entry) => entry.settled) + const flushes = Array.from(inFlightFlushes) .filter((entry) => entry.provider === provider) .map((entry) => entry.settled) - await Promise.allSettled(settled) + await Promise.allSettled([...generations, ...flushes]) } /** @@ -157,6 +163,12 @@ export async function flushPendingSessionTitle(taskId: string): Promise { pendingTitleWrites.delete(taskId) return } + let settle!: () => void + const settled = new Promise((resolve) => { + settle = resolve + }) + const entry = { provider: pending.provider, settled } + inFlightFlushes.add(entry) try { const result = await pending.provider.applySessionTitle(taskId, pending.title) if (result === "applied" || result === "conflict") { @@ -164,6 +176,9 @@ export async function flushPendingSessionTitle(taskId: string): Promise { } } catch { // Keep the title for the next save instead of losing it to a transient write failure. + } finally { + inFlightFlushes.delete(entry) + settle() } } @@ -227,7 +242,10 @@ export async function generateSessionTitle(options: GenerateSessionTitleOptions) const outcome = await applyTitleWithRetry(provider, taskId, title, () => inFlight.userCancelled) if (outcome === "not-ready") { // The history item is not persisted yet; `Task` retries this write once it is. - pendingTitleWrites.set(taskId, { provider, title }) + // A cancelled generation must not restore the entry `cancelSessionTitle` dropped. + if (!inFlight.userCancelled) { + pendingTitleWrites.set(taskId, { provider, title }) + } return } // "applied" means the title is durable, and "conflict" means manual intent owns the @@ -242,7 +260,7 @@ export async function generateSessionTitle(options: GenerateSessionTitleOptions) // Stryker disable next-line ConditionalExpression,LogicalOperator: defense-in-depth, fallback is always truthy for non-empty sanitized messages and distinct from failed title if (fallback && fallback !== title) { const outcome = await applyTitleWithRetry(provider, taskId, fallback, () => inFlight.userCancelled) - if (outcome === "not-ready") { + if (outcome === "not-ready" && !inFlight.userCancelled) { pendingTitleWrites.set(taskId, { provider, title: fallback }) } } diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index cb53ea189b..2287fe5d57 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -2835,23 +2835,18 @@ export class ClineProvider return } - // Read and write under the store lock so a queued task save or delegation - // transition cannot be overwritten by a stale snapshot of the item. + // Apply the manual title to the record persisted on disk under the file lock, so a + // clear removes a title another window wrote even when this cache has none. const trimmed = title?.trim() - await this.taskHistoryStore.atomicReadAndUpdate(id, (historyItem) => ({ - ...historyItem, - title: trimmed ? trimmed : undefined, - // Persist the clear intent: an empty title is otherwise indistinguishable from a - // task that never had one, which lets another host's pending auto title win. - titleCleared: trimmed ? undefined : true, - })) + const updatedItem = await this.taskHistoryStore.setTitleManually(id, trimmed ? trimmed : undefined) + if (!updatedItem) { + this.log(`[renameTaskWithId] task ${id} is no longer on disk`) + return + } this.recentTasksCache = undefined if (this.isViewLaunched) { - const updatedItem = this.taskHistoryStore.get(id) - if (updatedItem) { - await this.postMessageToWebview({ type: "taskHistoryItemUpdated", taskHistoryItem: updatedItem }) - } + await this.postMessageToWebview({ type: "taskHistoryItemUpdated", taskHistoryItem: updatedItem }) } } diff --git a/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts b/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts index 75742438af..290742b05b 100644 --- a/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts +++ b/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts @@ -711,6 +711,23 @@ describe("ClineProvider Task History Synchronization", () => { expect(getStored("task-stale-write")).toMatchObject({ title: "Renamed", status: "completed" }) }) + it("logs and skips the broadcast when the record is gone from disk", async () => { + await provider.resolveWebviewView(mockWebviewView) + provider.isViewLaunched = true + await provider.updateTaskHistory(createHistoryItem({ id: "task-vanished", task: "Original prompt" }), { + broadcast: false, + }) + const logSpy = vi.spyOn(provider, "log") + // Another window deletes the record between the cache check and the write. + vi.spyOn(provider.taskHistoryStore, "setTitleManually").mockResolvedValueOnce(undefined) + mockPostMessage.mockClear() + + await provider.renameTaskWithId("task-vanished", "Renamed") + + expect(logSpy).toHaveBeenCalledWith("[renameTaskWithId] task task-vanished is no longer on disk") + expect(findCallsByType(mockPostMessage.mock.calls, "taskHistoryItemUpdated")).toHaveLength(0) + }) + it("does nothing for an unknown task", async () => { const logSpy = vi.spyOn(provider, "log") From 9910985473251ea88a25c71baffb0008f50def20 Mon Sep 17 00:00:00 2001 From: daewoongoh Date: Sun, 8 Nov 2026 00:47:17 +0900 Subject: [PATCH 6/7] test(task): cover disposal drain and first-attempt option forwarding --- src/api/providers/__tests__/kimi-code.spec.ts | 3 +- src/core/task/__tests__/session-title.spec.ts | 22 +++++++ src/core/task/session-title.ts | 7 +++ .../ClineProvider.taskHistory.spec.ts | 63 ++++++++++++++++++- 4 files changed, 93 insertions(+), 2 deletions(-) diff --git a/src/api/providers/__tests__/kimi-code.spec.ts b/src/api/providers/__tests__/kimi-code.spec.ts index 2fb77296e8..ceeb9a88e0 100644 --- a/src/api/providers/__tests__/kimi-code.spec.ts +++ b/src/api/providers/__tests__/kimi-code.spec.ts @@ -145,7 +145,8 @@ describe("KimiCodeHandler", () => { await expect(handler.completePrompt("test", { abortSignal: signal, timeoutMs: 5_000 })).resolves.toBe("retried") - // The retry after the forced refresh keeps the same options. + // Both the initial attempt and the retry after the forced refresh carry the options. + expect(createCompletion).toHaveBeenNthCalledWith(1, expect.anything(), { signal, timeout: 5_000 }) expect(createCompletion).toHaveBeenNthCalledWith(2, expect.anything(), { signal, timeout: 5_000 }) }) diff --git a/src/core/task/__tests__/session-title.spec.ts b/src/core/task/__tests__/session-title.spec.ts index 9bd09d40f9..4cdaee26a6 100644 --- a/src/core/task/__tests__/session-title.spec.ts +++ b/src/core/task/__tests__/session-title.spec.ts @@ -921,6 +921,28 @@ describe("flushPendingSessionTitle", () => { expect(drained).toBe(true) }) + it("drops a settled pending title when the provider is disposed", async () => { + const recorder = createCompletePromptRecorder(async () => "Generated Title") + let historyItemPersisted = false + const applySpy = vi.fn().mockImplementation(async () => (historyItemPersisted ? "applied" : "not-ready")) + const provider = createTitleProvider({ applySessionTitle: applySpy }) + + await generateSessionTitle({ + provider, + taskId: "task-dispose-pending", + userMessage: "Fix the login bug", + apiConfiguration: testApiConfiguration, + completePrompt: recorder.fn, + }) + + // The generation already finished, so only the pending entry remains for disposal to drop. + cancelSessionTitlesForProvider(provider) + historyItemPersisted = true + await flushPendingSessionTitle("task-dispose-pending") + + expect(applySpy).toHaveBeenCalledTimes(1) + }) + it("does nothing when no title is pending", async () => { await expect(flushPendingSessionTitle("unknown-task")).resolves.toBeUndefined() }) diff --git a/src/core/task/session-title.ts b/src/core/task/session-title.ts index b7c36b63b1..bd9cdfac37 100644 --- a/src/core/task/session-title.ts +++ b/src/core/task/session-title.ts @@ -131,6 +131,13 @@ export function cancelSessionTitlesForProvider(provider: ClineProvider): void { entry.userCancelled = true entry.controller.abort() } + // Titles whose generation already finished keep their entry until a save flushes or a cancel + // drops it; disposal drops them too, so the released store is never written afterwards. + for (const [taskId, pending] of pendingTitleWrites) { + if (pending.provider === provider) { + pendingTitleWrites.delete(taskId) + } + } } /** diff --git a/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts b/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts index 290742b05b..e6903ad9a0 100644 --- a/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts +++ b/src/core/webview/__tests__/ClineProvider.taskHistory.spec.ts @@ -11,7 +11,7 @@ import { ContextProxy } from "../../config/ContextProxy" import { Task } from "../../task/Task" import { ProfileValidator } from "../../../shared/ProfileValidator" import { ClineProvider } from "../ClineProvider" -import { generateSessionTitle } from "../../task/session-title" +import { flushPendingSessionTitle, generateSessionTitle } from "../../task/session-title" import * as sessionTitleModule from "../../task/session-title" // Mock setup @@ -880,6 +880,67 @@ describe("ClineProvider Task History Synchronization", () => { expect(getStored("task-dispose")?.title).toBeUndefined() }) + it("waits for a deferred title write before releasing the history store", async () => { + const taskId = "task-dispose-deferred" + let resolveCompletion!: (title: string) => void + const completion = new Promise((resolve) => { + resolveCompletion = resolve + }) + const generation = generateSessionTitle({ + provider, + taskId, + userMessage: "Original prompt", + apiConfiguration: { + apiProvider: providerIdentifiers.anthropic, + apiModelId: "claude-3-5-sonnet-20241022", + apiKey: "test-api-key", + }, + completePrompt: async () => completion, + }) + // The title is generated before the history item exists, so it is deferred. + resolveCompletion("Generated Title") + await generation + await provider.updateTaskHistory(createHistoryItem({ id: taskId, task: "Original prompt" }), { + broadcast: false, + }) + + // Pause the deferred write inside the store. + let releaseWrite!: () => void + const writeGate = new Promise((resolve) => { + releaseWrite = resolve + }) + const originalSetTitle = provider.taskHistoryStore.setTitleIfUnset.bind(provider.taskHistoryStore) + vi.spyOn(provider.taskHistoryStore, "setTitleIfUnset").mockImplementation(async (id, title) => { + await writeGate + return originalSetTitle(id, title) + }) + const storeEvents: string[] = [] + const originalDispose = provider.taskHistoryStore.dispose.bind(provider.taskHistoryStore) + vi.spyOn(provider.taskHistoryStore, "dispose").mockImplementation(() => { + storeEvents.push("store-dispose") + originalDispose() + }) + + const flush = flushPendingSessionTitle(taskId) + await new Promise((resolve) => setTimeout(resolve, 10)) + + // Disposal is held open while the deferred write is still running. + let disposed = false + const disposal = provider.dispose().then(() => { + disposed = true + }) + await new Promise((resolve) => setTimeout(resolve, 10)) + expect(disposed).toBe(false) + expect(storeEvents).toEqual([]) + + releaseWrite() + await flush + await disposal + + expect(storeEvents).toEqual(["store-dispose"]) + expect(getStored(taskId)?.title).toBe("Generated Title") + }) + it("cancels pending generation and writes no late title when deleting a task", async () => { await provider.updateTaskHistory( createHistoryItem({ id: "task-cancel-on-delete", task: "Original prompt" }), From 2df5449c6b8e305de5f5abf50f1a435e2ad6b94b Mon Sep 17 00:00:00 2001 From: daewoongoh Date: Sun, 8 Nov 2026 01:11:06 +0900 Subject: [PATCH 7/7] fix(api): release the completion handler when cancellation outlives model selection --- src/api/providers/__tests__/kimi-code.spec.ts | 28 +++++++++++++++++-- src/utils/__tests__/enhance-prompt.spec.ts | 17 +++++++++++ src/utils/single-completion-handler.ts | 7 ++++- 3 files changed, 48 insertions(+), 4 deletions(-) diff --git a/src/api/providers/__tests__/kimi-code.spec.ts b/src/api/providers/__tests__/kimi-code.spec.ts index ceeb9a88e0..a2cb8a1d99 100644 --- a/src/api/providers/__tests__/kimi-code.spec.ts +++ b/src/api/providers/__tests__/kimi-code.spec.ts @@ -145,9 +145,11 @@ describe("KimiCodeHandler", () => { await expect(handler.completePrompt("test", { abortSignal: signal, timeoutMs: 5_000 })).resolves.toBe("retried") - // Both the initial attempt and the retry after the forced refresh carry the options. - expect(createCompletion).toHaveBeenNthCalledWith(1, expect.anything(), { signal, timeout: 5_000 }) - expect(createCompletion).toHaveBeenNthCalledWith(2, expect.anything(), { signal, timeout: 5_000 }) + // Both the initial attempt and the retry after the forced refresh carry the caller's own signal. + expect(createCompletion.mock.calls[0][1].signal).toBe(signal) + expect(createCompletion.mock.calls[0][1].timeout).toBe(5_000) + expect(createCompletion.mock.calls[1][1].signal).toBe(signal) + expect(createCompletion.mock.calls[1][1].timeout).toBe(5_000) }) it("does not start preparation when the caller already aborted", async () => { @@ -159,6 +161,26 @@ describe("KimiCodeHandler", () => { expect(mockGetAccessToken).not.toHaveBeenCalled() }) + it("stops before model discovery and completion when aborted during token resolution", async () => { + const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "oauth" }) + let resolveToken!: (token: string) => void + mockGetAccessToken.mockReturnValueOnce( + new Promise((resolve) => { + resolveToken = resolve + }), + ) + const controller = new AbortController() + + const completion = handler.completePrompt("test", { abortSignal: controller.signal }) + await vi.waitFor(() => expect(resolveToken).toBeDefined()) + controller.abort() + resolveToken("oauth-token") + + await expect(completion).rejects.toMatchObject({ name: "AbortError" }) + // The cancelled request never reaches model discovery or the provider. + expect(mockGetModels).not.toHaveBeenCalled() + }) + it("does not retry on 401 when using API key auth", async () => { const handler = new KimiCodeHandler({ kimiCodeAuthMethod: "api-key", kimiCodeApiKey: "key" }) vi.spyOn(globalThis, "fetch").mockResolvedValueOnce(new Response(null, { status: 401 })) diff --git a/src/utils/__tests__/enhance-prompt.spec.ts b/src/utils/__tests__/enhance-prompt.spec.ts index 47149ae4a7..5463a5390e 100644 --- a/src/utils/__tests__/enhance-prompt.spec.ts +++ b/src/utils/__tests__/enhance-prompt.spec.ts @@ -118,6 +118,23 @@ describe("enhancePrompt", () => { expect(succeeding.dispose).toHaveBeenCalledOnce() }) + it("settles on abort and disposes the handler while the provider request is still pending", async () => { + // Model discovery or provider preparation may never settle; the caller's stop must still + // release the handler and its listeners. + const pendingHandler = createCompletionHandlerDouble({ + completePrompt: vi.fn(() => new Promise(() => {})), + dispose: vi.fn(), + }) + vi.mocked(buildApiHandler).mockReturnValueOnce(pendingHandler) + const abortController = new AbortController() + + const completion = singleCompletionHandler(mockApiConfig, "a", { abortSignal: abortController.signal }) + abortController.abort() + + await expect(completion).rejects.toMatchObject({ name: "AbortError" }) + expect(pendingHandler.dispose).toHaveBeenCalledOnce() + }) + it("disposes the handler when the provider does not support prompt enhancement", async () => { const unsupported = createApiHandlerDouble({ dispose: vi.fn() }) vi.mocked(buildApiHandler).mockReturnValueOnce(unsupported) diff --git a/src/utils/single-completion-handler.ts b/src/utils/single-completion-handler.ts index 148f440461..9cef87b9a6 100644 --- a/src/utils/single-completion-handler.ts +++ b/src/utils/single-completion-handler.ts @@ -1,6 +1,7 @@ import type { ProviderSettings } from "@roo-code/types" import { buildApiHandler, disposeApiHandler, SingleCompletionHandler, type CompletePromptOptions } from "../api" +import { settleOnAbort } from "../api/providers/utils/abort-signal" /** * Enhances a prompt using the configured API without creating a full Cline instance or task history. @@ -26,7 +27,11 @@ export async function singleCompletionHandler( throw new Error("The selected API provider does not support prompt enhancement") } - return await (handler as SingleCompletionHandler).completePrompt(promptText, options) + const completion = (handler as SingleCompletionHandler).completePrompt(promptText, options) + // Settle on the caller's stop even when the provider never returns (model discovery or a + // preparation step can outlive it), so the handler below is disposed and its listeners are + // released instead of being retained by a request nobody waits for any more. + return await settleOnAbort(completion, options?.abortSignal, apiConfiguration.apiProvider) } finally { // Handlers such as VS Code LM / GitHub Copilot hold listeners that must be released. disposeApiHandler(handler)