Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions src/api/providers/__tests__/openai.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -986,6 +986,21 @@ describe("OpenAiHandler", () => {
const result = await handler.completePrompt("Test prompt")
expect(result).toBe("")
})

it("forwards abort signal and timeout to the SDK request", async () => {
const controller = new AbortController()
await handler.completePrompt("Test prompt", {
abortSignal: controller.signal,
timeoutMs: 15_000,
})
expect(mockCreate).toHaveBeenCalledWith(
expect.anything(),
expect.objectContaining({
signal: controller.signal,
timeout: 15_000,
}),
)
})
})

describe.each([
Expand Down
23 changes: 23 additions & 0 deletions src/api/providers/__tests__/openrouter.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -999,6 +999,29 @@ describe("OpenRouterHandler", () => {
)
})

it("forwards abort signal and timeout to the SDK request", async () => {
const handler = new OpenRouterHandler(mockOptions)
const mockResponse = { choices: [{ message: { content: "test completion" } }] }
// AGENTS.md: access private members with bracket notation without as any
const createSpy = vi
.spyOn(handler["client"].chat.completions, "create")
.mockResolvedValue(mockResponse as unknown as OpenAI.Chat.Completions.ChatCompletion)

const controller = new AbortController()
await handler.completePrompt("test prompt", {
abortSignal: controller.signal,
timeoutMs: 15_000,
})

expect(createSpy).toHaveBeenCalledWith(
expect.anything(),
expect.objectContaining({
signal: controller.signal,
timeout: 15_000,
}),
)
})

it("handles API errors and captures telemetry", async () => {
const handler = new OpenRouterHandler(mockOptions)
const mockError = {
Expand Down
8 changes: 6 additions & 2 deletions src/api/providers/base-openai-compatible-provider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 } from "./utils/abort-signal"

type BaseOpenAiCompatibleProviderOptions<ModelName extends string> = ApiHandlerOptions & {
providerName: string
Expand Down Expand Up @@ -212,7 +213,7 @@ export abstract class BaseOpenAiCompatibleProvider<ModelName extends string>
}
}

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

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

try {
const response = await this.client.chat.completions.create(params)
const response = await this.client.chat.completions.create(params, {
signal: options?.abortSignal,
...(hasUsableTimeout(options?.timeoutMs) ? { timeout: toSdkTimeoutMs(options.timeoutMs) } : {}),
})

// Check for provider-specific error responses (e.g., MiniMax base_resp)
const responseAny = response as any
Expand Down
8 changes: 7 additions & 1 deletion src/api/providers/openai-native.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ import { NOT_PROVIDED } from "./constants"
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata, CompletePromptOptions } from "../index"
import { isMcpTool } from "../../utils/mcp-name"
import { sanitizeOpenAiCallId } from "../../utils/tool-id"
import { mergeAbortSignalAndTimeout } from "./utils/abort-signal"

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

Expand Down Expand Up @@ -1643,8 +1644,13 @@ export class OpenAiNativeHandler extends BaseProvider implements SingleCompletio
}

// Make the non-streaming request
const requestSignal = mergeAbortSignalAndTimeout(options?.abortSignal, options?.timeoutMs)
const combinedSignal = requestSignal
? AbortSignal.any([this.abortController.signal, requestSignal])
: this.abortController.signal

const response = await (this.client as any).responses.create(requestBody, {
signal: this.abortController.signal,
signal: combinedSignal,
})

// Extract text from the response
Expand Down
12 changes: 7 additions & 5 deletions src/api/providers/openai.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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, toSdkTimeoutMs } from "./utils/abort-signal"

type OpenAiUsage =
| {
Expand Down Expand Up @@ -354,7 +355,7 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
}
}

async completePrompt(prompt: string, _options?: CompletePromptOptions): Promise<string> {
async completePrompt(prompt: string, options?: CompletePromptOptions): Promise<string> {
try {
const isAzureAiInference = this._isAzureAiInference(this.options.openAiBaseUrl)
const model = this.getModel()
Expand All @@ -372,10 +373,11 @@ 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 } : {},
)
response = await this.client.chat.completions.create(requestOptions, {
signal: options?.abortSignal,
...(hasUsableTimeout(options?.timeoutMs) ? { timeout: toSdkTimeoutMs(options.timeoutMs) } : {}),
...(isAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {}),
})
} catch (error) {
throw handleOpenAIError(error, this.providerName)
}
Expand Down
18 changes: 13 additions & 5 deletions src/api/providers/openrouter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ import type {
import { handleOpenAIError } from "./utils/error-handler"
import { generateImageWithProvider, ImageGenerationResult } from "./utils/image-generation"
import { applyRouterToolPreferences } from "./utils/router-tool-preferences"
import { hasUsableTimeout, toSdkTimeoutMs } from "./utils/abort-signal"

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

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

const completionParams: OpenRouterChatCompletionParams = {
Expand All @@ -622,14 +623,21 @@ 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 = {
...(modelId.startsWith("anthropic/")
? { headers: { "x-anthropic-beta": "fine-grained-tool-streaming-2025-05-14" } }
: {}),
...(options?.abortSignal ? { signal: options.abortSignal } : {}),
...(hasUsableTimeout(options?.timeoutMs) ? { timeout: toSdkTimeoutMs(options.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)
Expand Down
68 changes: 68 additions & 0 deletions src/core/task-persistence/TaskHistoryStore.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,14 @@ function mergeWithDisk(delta: Partial<HistoryItem>): (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. */
// Stryker disable next-line StringLiteral: diagnostic error message only, caught internally
class TitleWriteSkipped extends Error {
constructor(readonly reason: "missing" | "titled") {
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
Expand Down Expand Up @@ -996,6 +1004,66 @@ 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) {
throw new TitleWriteSkipped("titled")
}
written = mergeHistoryDelta(existing, incoming as HistoryItem, { title })
Comment thread
daewoongoh marked this conversation as resolved.
Resolved
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)
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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading
Loading