Skip to content

Commit d597356

Browse files
committed
feat(task): auto-generate session title from the first user message
- Start background title generation for root tasks without blocking task loop - Support positive rateLimitSeconds settings with model title generation - Enforce 15,000 ms timeout boundary with independent test coverage - Connect title generation to task and provider disposal to prevent orphaned requests - Propagate abort signal and timeout across completion providers - Dispose stalled single-completion handlers immediately upon abortion
1 parent bb6cbee commit d597356

15 files changed

Lines changed: 1572 additions & 18 deletions

‎src/api/providers/__tests__/openai.spec.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -986,6 +986,21 @@ describe("OpenAiHandler", () => {
986986
const result = await handler.completePrompt("Test prompt")
987987
expect(result).toBe("")
988988
})
989+
990+
it("forwards abort signal and timeout to the SDK request", async () => {
991+
const controller = new AbortController()
992+
await handler.completePrompt("Test prompt", {
993+
abortSignal: controller.signal,
994+
timeoutMs: 15_000,
995+
})
996+
expect(mockCreate).toHaveBeenCalledWith(
997+
expect.anything(),
998+
expect.objectContaining({
999+
signal: controller.signal,
1000+
timeout: 15_000,
1001+
}),
1002+
)
1003+
})
9891004
})
9901005

9911006
describe.each([

‎src/api/providers/__tests__/openrouter.spec.ts‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -999,6 +999,29 @@ describe("OpenRouterHandler", () => {
999999
)
10001000
})
10011001

1002+
it("forwards abort signal and timeout to the SDK request", async () => {
1003+
const handler = new OpenRouterHandler(mockOptions)
1004+
const mockResponse = { choices: [{ message: { content: "test completion" } }] }
1005+
// AGENTS.md: access private members with bracket notation without as any
1006+
const createSpy = vi
1007+
.spyOn(handler["client"].chat.completions, "create")
1008+
.mockResolvedValue(mockResponse as unknown as OpenAI.Chat.Completions.ChatCompletion)
1009+
1010+
const controller = new AbortController()
1011+
await handler.completePrompt("test prompt", {
1012+
abortSignal: controller.signal,
1013+
timeoutMs: 15_000,
1014+
})
1015+
1016+
expect(createSpy).toHaveBeenCalledWith(
1017+
expect.anything(),
1018+
expect.objectContaining({
1019+
signal: controller.signal,
1020+
timeout: 15_000,
1021+
}),
1022+
)
1023+
})
1024+
10021025
it("handles API errors and captures telemetry", async () => {
10031026
const handler = new OpenRouterHandler(mockOptions)
10041027
const mockError = {

‎src/api/providers/base-openai-compatible-provider.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import { BaseProvider } from "./base-provider"
1414
import { handleOpenAIError } from "./utils/error-handler"
1515
import { calculateApiCostOpenAI } from "../../shared/cost"
1616
import { extractReasoningFromDelta } from "./utils/extract-reasoning"
17+
import { hasUsableTimeout, toSdkTimeoutMs } from "./utils/abort-signal"
1718

1819
type BaseOpenAiCompatibleProviderOptions<ModelName extends string> = ApiHandlerOptions & {
1920
providerName: string
@@ -212,7 +213,7 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
212213
}
213214
}
214215

215-
async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise<string> {
216+
async completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
216217
const { id: modelId, info: modelInfo } = this.getModel()
217218

218219
const params: OpenAI.Chat.Completions.ChatCompletionCreateParams = {
@@ -226,7 +227,10 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
226227
}
227228

228229
try {
229-
const response = await this.client.chat.completions.create(params)
230+
const response = await this.client.chat.completions.create(params, {
231+
signal: options?.abortSignal,
232+
...(hasUsableTimeout(options?.timeoutMs) ? { timeout: toSdkTimeoutMs(options.timeoutMs) } : {}),
233+
})
230234

231235
// Check for provider-specific error responses (e.g., MiniMax base_resp)
232236
const responseAny = response as any

‎src/api/providers/openai-native.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import { NOT_PROVIDED } from "./constants"
3434
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
3535
import { isMcpTool } from "../../utils/mcp-name"
3636
import { sanitizeOpenAiCallId } from "../../utils/tool-id"
37+
import { mergeAbortSignalAndTimeout } from "./utils/abort-signal"
3738

3839
export type OpenAiNativeModel = ReturnType<OpenAiNativeHandler["getModel"]>
3940

@@ -1643,8 +1644,13 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio
16431644
}
16441645

16451646
// Make the non-streaming request
1647+
const requestSignal = mergeAbortSignalAndTimeout(options?.abortSignal, options?.timeoutMs)
1648+
const combinedSignal = requestSignal
1649+
? AbortSignal.any([this.abortController.signal, requestSignal])
1650+
: this.abortController.signal
1651+
16461652
const response = await (this.client as any).responses.create(requestBody, {
1647-
signal: this.abortController.signal,
1653+
signal: combinedSignal,
16481654
})
16491655

16501656
// Extract text from the response

‎src/api/providers/openai.ts‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ import { BaseProvider } from "./base-provider"
2929
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
3030
import { handleOpenAIError } from "./utils/error-handler"
3131
import { extractReasoningFromDelta } from "./utils/extract-reasoning"
32+
import { hasUsableTimeout, toSdkTimeoutMs } from "./utils/abort-signal"
3233

3334
type OpenAiUsage =
3435
| {
@@ -354,7 +355,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
354355
}
355356
}
356357

357-
async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise<string> {
358+
async completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
358359
try {
359360
const isAzureAiInference = this._isAzureAiInference(this.options.openAiBaseUrl)
360361
const model = this.getModel()
@@ -372,10 +373,11 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
372373

373374
let response
374375
try {
375-
response = await this.client.chat.completions.create(
376-
requestOptions,
377-
isAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {},
378-
)
376+
response = await this.client.chat.completions.create(requestOptions, {
377+
signal: options?.abortSignal,
378+
...(hasUsableTimeout(options?.timeoutMs) ? { timeout: toSdkTimeoutMs(options.timeoutMs) } : {}),
379+
...(isAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {}),
380+
})
379381
} catch (error) {
380382
throw handleOpenAIError(error, this.providerName)
381383
}

‎src/api/providers/openrouter.ts‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ import type {
4545
import { handleOpenAIError } from "./utils/error-handler"
4646
import { generateImageWithProvider, ImageGenerationResult } from "./utils/image-generation"
4747
import { applyRouterToolPreferences } from "./utils/router-tool-preferences"
48+
import { hasUsableTimeout, toSdkTimeoutMs } from "./utils/abort-signal"
4849

4950
// Add custom interface for OpenRouter params.
5051
type OpenRouterChatCompletionParams = OpenAI.Chat.ChatCompletionCreateParams & {
@@ -600,7 +601,7 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
600601
return { id, info, topP: isDeepSeekR1 ? 0.95 : undefined, ...params }
601602
}
602603

603-
async completePrompt(prompt: string, _options?: CompletePromptOptions) {
604+
async completePrompt(prompt: string, options?: CompletePromptOptions) {
604605
const { id: modelId, maxTokens, temperature, reasoning } = await this.fetchModel()
605606

606607
const completionParams: OpenRouterChatCompletionParams = {
@@ -622,14 +623,21 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH
622623
}
623624

624625
// Add Anthropic beta header for fine-grained tool streaming when using Anthropic models
625-
const requestOptions = modelId.startsWith("anthropic/")
626-
? { headers: { "x-anthropic-beta": "fine-grained-tool-streaming-2025-05-14" } }
627-
: undefined
626+
const requestOptions: OpenAI.RequestOptions = {
627+
...(modelId.startsWith("anthropic/")
628+
? { headers: { "x-anthropic-beta": "fine-grained-tool-streaming-2025-05-14" } }
629+
: {}),
630+
...(options?.abortSignal ? { signal: options.abortSignal } : {}),
631+
...(hasUsableTimeout(options?.timeoutMs) ? { timeout: toSdkTimeoutMs(options.timeoutMs) } : {}),
632+
}
628633

629634
let response
630635

631636
try {
632-
response = await this.client.chat.completions.create(completionParams, requestOptions)
637+
response = await this.client.chat.completions.create(
638+
completionParams,
639+
Object.keys(requestOptions).length > 0 ? requestOptions : undefined,
640+
)
633641
} catch (error) {
634642
// Try to parse as OpenRouter error structure using Zod
635643
const parseResult = OpenRouterErrorResponseSchema.safeParse(error)

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

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,14 @@ function mergeWithDisk(delta: Partial<HistoryItem>): (existing: unknown, incomin
2626
return (existing, incoming) => mergeHistoryDelta(existing, incoming as HistoryItem, delta)
2727
}
2828

29+
/** Thrown inside the `setTitleIfUnset` merge callback to abort the write without touching the file. */
30+
// Stryker disable next-line StringLiteral: diagnostic error message only, caught internally
31+
class TitleWriteSkipped extends Error {
32+
constructor(readonly reason: "missing" | "titled") {
33+
super(`title write skipped: ${reason}`)
34+
}
35+
}
36+
2937
/**
3038
* Durable intent for the one repair that spans an active delegated child and
3139
* its parent. Task files remain authoritative; this file only records the
@@ -996,6 +1004,66 @@ export class TaskHistoryStore {
9961004

9971005
// ────────────────────────────── Atomic read-modify-write ──────────────────────────────
9981006

1007+
/**
1008+
* Set `title` only while the item has none, checking the title on disk inside the
1009+
* per-file lock. The cache can be stale when another store (for example a second
1010+
* provider on the same storage path) saved a manual title, and a delta merge would
1011+
* otherwise overwrite it.
1012+
*
1013+
* @returns "conflict" when a title already exists or the record was deleted on disk (the cache is refreshed).
1014+
* @throws If the task ID is not present in the cache.
1015+
*/
1016+
public setTitleIfUnset(taskId: string, title: string): Promise<"applied" | "conflict"> {
1017+
return this.withLock(async () => {
1018+
const current = this.cache.get(taskId)
1019+
if (!current) {
1020+
throw new Error(`[TaskHistoryStore] setTitleIfUnset: task ${taskId} not found in cache`)
1021+
}
1022+
if (current.title) {
1023+
return "conflict"
1024+
}
1025+
1026+
const filePath = await this.getTaskFilePath(taskId)
1027+
let written: HistoryItem = { ...current, title }
1028+
try {
1029+
await safeWriteJson(filePath, written, {
1030+
merge: (existing, incoming) => {
1031+
// Throwing aborts the write, so a deleted record is never recreated from the cache
1032+
// and an existing title is never rewritten.
1033+
// Stryker disable next-line ConditionalExpression: defense-in-depth, corrupted disk content cannot be merged
1034+
if (!existing || typeof existing !== "object" || !("id" in existing)) {
1035+
throw new TitleWriteSkipped("missing")
1036+
}
1037+
if ((existing as HistoryItem).title) {
1038+
throw new TitleWriteSkipped("titled")
1039+
}
1040+
written = mergeHistoryDelta(existing, incoming as HistoryItem, { title })
1041+
return written
1042+
},
1043+
})
1044+
} catch (error) {
1045+
if (!(error instanceof TitleWriteSkipped)) {
1046+
throw error
1047+
}
1048+
if (error.reason === "missing") {
1049+
this.cache.delete(taskId)
1050+
} else {
1051+
const diskItem = await this.readTaskFile(taskId)
1052+
if (diskItem) {
1053+
this.cache.set(taskId, diskItem)
1054+
}
1055+
}
1056+
return "conflict"
1057+
}
1058+
1059+
this.cache.set(taskId, written)
1060+
if (this.onWrite) {
1061+
await this.onWrite(this.getAll())
1062+
}
1063+
return "applied"
1064+
})
1065+
}
1066+
9991067
/**
10001068
* Read a HistoryItem from the in-memory cache and write back an updated version,
10011069
* all within a single lock acquisition so no concurrent writer can interleave

‎src/core/task-persistence/__tests__/TaskHistoryStore.crossInstance.spec.ts‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,104 @@ describe("TaskHistoryStore cross-instance safety", () => {
7373
await fs.rm(tmpDir, { recursive: true, force: true }).catch(() => {})
7474
})
7575

76+
describe("setTitleIfUnset", () => {
77+
it("keeps a manual title that another instance wrote after this cache was loaded", async () => {
78+
await storeA.initialize()
79+
await storeB.initialize()
80+
await storeA.upsert(makeHistoryItem({ id: "title-task", task: "Original prompt" }))
81+
await storeB.reconcile()
82+
expect(storeB.get("title-task")?.title).toBeUndefined()
83+
84+
// Instance A still caches no title when B renames the task on disk.
85+
await storeB.atomicReadAndUpdate("title-task", (item) => ({ ...item, title: "Manual title" }))
86+
expect(storeA.get("title-task")?.title).toBeUndefined()
87+
88+
await expect(storeA.setTitleIfUnset("title-task", "Auto title")).resolves.toBe("conflict")
89+
90+
expect(storeA.get("title-task")?.title).toBe("Manual title")
91+
await storeB.reconcile()
92+
expect(storeB.get("title-task")?.title).toBe("Manual title")
93+
})
94+
95+
it("does not recreate a task that another instance deleted", async () => {
96+
await storeA.initialize()
97+
await storeB.initialize()
98+
await storeA.upsert(makeHistoryItem({ id: "title-deleted", task: "Original prompt" }))
99+
await storeB.reconcile()
100+
101+
await storeB.delete("title-deleted")
102+
expect(storeA.get("title-deleted")).toBeDefined()
103+
104+
await expect(storeA.setTitleIfUnset("title-deleted", "Auto title")).resolves.toBe("conflict")
105+
106+
expect(storeA.get("title-deleted")).toBeUndefined()
107+
await storeB.reconcile()
108+
expect(storeB.get("title-deleted")).toBeUndefined()
109+
})
110+
111+
it("applies the title when no instance has set one", async () => {
112+
await storeA.initialize()
113+
await storeA.upsert(makeHistoryItem({ id: "title-free", task: "Original prompt" }))
114+
115+
await expect(storeA.setTitleIfUnset("title-free", "Auto title")).resolves.toBe("applied")
116+
117+
expect(storeA.get("title-free")).toMatchObject({ title: "Auto title", task: "Original prompt" })
118+
})
119+
120+
it("persists the title to disk so reloading from storage preserves it", async () => {
121+
await storeA.initialize()
122+
await storeA.upsert(makeHistoryItem({ id: "persist-check", task: "Original task prompt" }))
123+
124+
await expect(storeA.setTitleIfUnset("persist-check", "Persisted title")).resolves.toBe("applied")
125+
126+
// Reload a fresh store instance pointing to the same disk directory
127+
const storeReloaded = new TaskHistoryStore(tmpDir)
128+
await storeReloaded.initialize()
129+
130+
const diskItem = storeReloaded.get("persist-check")
131+
expect(diskItem).toBeDefined()
132+
expect(diskItem?.title).toBe("Persisted title")
133+
expect(diskItem?.task).toBe("Original task prompt")
134+
storeReloaded.dispose()
135+
})
136+
137+
it("throws when the task is not in the cache", async () => {
138+
await storeA.initialize()
139+
await expect(storeA.setTitleIfUnset("missing-id", "Auto title")).rejects.toThrow("not found in cache")
140+
})
141+
142+
it("returns conflict immediately when cached item already has a title", async () => {
143+
await storeA.initialize()
144+
await storeA.upsert(makeHistoryItem({ id: "already-titled", title: "Existing", task: "Prompt" }))
145+
await expect(storeA.setTitleIfUnset("already-titled", "New title")).resolves.toBe("conflict")
146+
})
147+
148+
it("notifies onWrite callback when title is successfully applied", async () => {
149+
const onWriteSpy = vi.fn().mockResolvedValue(undefined)
150+
const storeWithOnWrite = new TaskHistoryStore(tmpDir, { onWrite: onWriteSpy })
151+
await storeWithOnWrite.initialize()
152+
await storeWithOnWrite.upsert(makeHistoryItem({ id: "title-onwrite", task: "Prompt" }))
153+
onWriteSpy.mockClear()
154+
155+
await expect(storeWithOnWrite.setTitleIfUnset("title-onwrite", "Notified title")).resolves.toBe("applied")
156+
expect(onWriteSpy).toHaveBeenCalledWith(
157+
expect.arrayContaining([expect.objectContaining({ id: "title-onwrite", title: "Notified title" })]),
158+
)
159+
storeWithOnWrite.dispose()
160+
})
161+
162+
it("re-throws unexpected errors from safeWriteJson", async () => {
163+
await storeA.initialize()
164+
await storeA.upsert(makeHistoryItem({ id: "title-io-err", task: "Prompt" }))
165+
166+
const { safeWriteJson } = await import("../../../utils/safeWriteJson")
167+
const mockedSafeWrite = vi.mocked(safeWriteJson)
168+
mockedSafeWrite.mockRejectedValueOnce(new Error("disk IO error"))
169+
170+
await expect(storeA.setTitleIfUnset("title-io-err", "Auto title")).rejects.toThrow("disk IO error")
171+
})
172+
})
173+
76174
it("two instances can write different tasks without conflict", async () => {
77175
await storeA.initialize()
78176
await storeB.initialize()

0 commit comments

Comments
 (0)