From 3b1a7eb990014f722da163f78d63480d0f19cb14 Mon Sep 17 00:00:00 2001 From: Ayush Patel Date: Fri, 14 Aug 2026 17:41:04 -0400 Subject: [PATCH] fix(search): resolve server-tool generation ids --- src/benchmarks/search/core/solver.test.ts | 21 +++ src/benchmarks/search/core/solver.ts | 3 + src/judge/judge.test.ts | 1 + src/providers/openrouter-model.test.ts | 2 + src/providers/responses-client.test.ts | 3 +- src/providers/responses-client.ts | 4 +- src/runtime/generation-ids.test.ts | 22 +++ src/runtime/generation-ids.ts | 18 ++- src/runtime/generation-resolver.test.ts | 84 +++++++++-- src/runtime/generation-resolver.ts | 166 ++++++++++++++-------- 10 files changed, 253 insertions(+), 71 deletions(-) diff --git a/src/benchmarks/search/core/solver.test.ts b/src/benchmarks/search/core/solver.test.ts index 86762ff..0e18bc0 100644 --- a/src/benchmarks/search/core/solver.test.ts +++ b/src/benchmarks/search/core/solver.test.ts @@ -293,6 +293,27 @@ describe("searchSolver", () => { solver(initialTaskState({ id: "s", input: "q", target: { text: "t" } })) ); expect(sentOptions?.extraHeaders).toBeUndefined(); + expect(sentOptions?.resolveGenerationChildren).toBe(true); + }); + it("does not resolve child generations for plugin search", async () => { + let sentOptions: ResponsesSendOptions | undefined; + const solver = searchSolver( + { + send: (_body, options) => { + sentOptions = options; + return effectSucceed(fixtureResult({ text: "x" })); + }, + }, + { + model: "m", + instructions: "i", + lane: makeLane({ webSearch: "plugin" }), + } + ); + await runSolver( + solver(initialTaskState({ id: "s", input: "q", target: { text: "t" } })) + ); + expect(sentOptions?.resolveGenerationChildren).toBeUndefined(); }); it("forwards provider flags without an endpoint id", async () => { let sentOptions: ResponsesSendOptions | undefined; diff --git a/src/benchmarks/search/core/solver.ts b/src/benchmarks/search/core/solver.ts index cfcfd6e..d2b0299 100644 --- a/src/benchmarks/search/core/solver.ts +++ b/src/benchmarks/search/core/solver.ts @@ -125,6 +125,9 @@ export function searchSolver( }; const sendOptions = (): ResponsesSendOptions => ({ timeoutMs: opts.timeoutMs ?? DEFAULT_SEARCH_TIMEOUT_MS, + ...(opts.lane.webSearch !== "plugin" && { + resolveGenerationChildren: true, + }), ...(extraHeaders !== undefined && { extraHeaders }), ...(opts.versionOverride !== undefined && { versionOverride: opts.versionOverride, diff --git a/src/judge/judge.test.ts b/src/judge/judge.test.ts index b8e78bf..40ccee4 100644 --- a/src/judge/judge.test.ts +++ b/src/judge/judge.test.ts @@ -119,6 +119,7 @@ describe("judgeCall", () => { isCacheHit: true, countsTowardUsage: true, isResolvedSource: false, + shouldResolveChildren: false, }, ]); }); diff --git a/src/providers/openrouter-model.test.ts b/src/providers/openrouter-model.test.ts index 7764a6e..e90e38e 100644 --- a/src/providers/openrouter-model.test.ts +++ b/src/providers/openrouter-model.test.ts @@ -182,6 +182,7 @@ describe("openrouter-model request parity", () => { isCacheHit: true, countsTowardUsage: true, isResolvedSource: true, + shouldResolveChildren: false, }, ]); } finally { @@ -221,6 +222,7 @@ describe("openrouter-model request parity", () => { isCacheHit: true, countsTowardUsage: true, isResolvedSource: false, + shouldResolveChildren: false, }, ]); } finally { diff --git a/src/providers/responses-client.test.ts b/src/providers/responses-client.test.ts index f1173d2..5956068 100644 --- a/src/providers/responses-client.test.ts +++ b/src/providers/responses-client.test.ts @@ -453,7 +453,7 @@ describe("makeResponsesLayer", () => { const responses = yield* Responses; yield* responses.send( { model: "m", input: [] }, - { timeoutMs: 1000 } + { timeoutMs: 1000, resolveGenerationChildren: true } ); }) ), @@ -472,6 +472,7 @@ describe("makeResponsesLayer", () => { isCacheHit: true, countsTowardUsage: true, isResolvedSource: true, + shouldResolveChildren: true, }, ]); } finally { diff --git a/src/providers/responses-client.ts b/src/providers/responses-client.ts index 191b4c3..5f33339 100644 --- a/src/providers/responses-client.ts +++ b/src/providers/responses-client.ts @@ -94,6 +94,7 @@ export interface ResponsesSendOptions { readonly extraBody?: Readonly>; readonly onResponseIdentifiers?: (identifiers: ModelErrorIdentifiers) => void; readonly onStreamEvent?: (event: StreamEvents) => void; + readonly resolveGenerationChildren?: boolean; } export interface ResponsesConfig { @@ -232,7 +233,8 @@ export function makeResponsesLayer(config: ResponsesConfig): Layer { ? cacheSourceId : result.generationId, isCacheHit, - isCacheHit && cacheSourceId !== undefined + isCacheHit && cacheSourceId !== undefined, + options.resolveGenerationChildren ?? false ).pipe(map(() => result)) : fail( new ResponsesError({ diff --git a/src/runtime/generation-ids.test.ts b/src/runtime/generation-ids.test.ts index f2423e7..4886ee1 100644 --- a/src/runtime/generation-ids.test.ts +++ b/src/runtime/generation-ids.test.ts @@ -26,18 +26,21 @@ describe("generation id collector", () => { isCacheHit: true, countsTowardUsage: true, isResolvedSource: false, + shouldResolveChildren: false, }, { id: "gen-real", isCacheHit: false, countsTowardUsage: true, isResolvedSource: false, + shouldResolveChildren: false, }, { id: "gen-source", isCacheHit: true, countsTowardUsage: true, isResolvedSource: true, + shouldResolveChildren: false, }, ]); }); @@ -58,12 +61,31 @@ describe("generation id collector", () => { isCacheHit: true, countsTowardUsage: false, isResolvedSource: false, + shouldResolveChildren: false, }, { id: "gen-solver", isCacheHit: true, countsTowardUsage: true, isResolvedSource: false, + shouldResolveChildren: false, + }, + ]); + }); + it("flags generation roots whose child ids should be resolved", async () => { + const entries = await runPromise( + resetGenerationIds.pipe( + flatMap(() => recordGenerationId("gen-root", false, false, true)), + flatMap(() => getCollectedGenerationIdEntries) + ) + ); + expect(entries).toEqual([ + { + id: "gen-root", + isCacheHit: false, + countsTowardUsage: true, + isResolvedSource: false, + shouldResolveChildren: true, }, ]); }); diff --git a/src/runtime/generation-ids.ts b/src/runtime/generation-ids.ts index 3199225..4f9d0ff 100644 --- a/src/runtime/generation-ids.ts +++ b/src/runtime/generation-ids.ts @@ -20,6 +20,8 @@ export const auxiliaryUsageGenerationIdCollector = export const resolvedSourceGenerationIdCollector = unsafeMakeHashSet(empty()); +export const childGenerationRootCollector = unsafeMakeHashSet(empty()); + export const auxiliaryUsageRef: FiberRef = unsafeMake(false); export function withAuxiliaryUsage( @@ -31,7 +33,8 @@ export function withAuxiliaryUsage( export function recordGenerationId( id: string | null | undefined, isCacheHit = false, - isResolvedSource = false + isResolvedSource = false, + shouldResolveChildren = false ): Effect { if (id === null || id === undefined || id.length === 0) { return succeed(undefined); @@ -48,6 +51,11 @@ export function recordGenerationId( ); } } + if (shouldResolveChildren) { + record = record.pipe( + zipRight(update(childGenerationRootCollector, add(generationId))) + ); + } return get(auxiliaryUsageRef).pipe( flatMap((isAuxiliary) => isAuxiliary @@ -67,7 +75,8 @@ export const resetGenerationIds: Effect = set( ).pipe( zipRight(set(cacheHitGenerationIdCollector, empty())), zipRight(set(auxiliaryUsageGenerationIdCollector, empty())), - zipRight(set(resolvedSourceGenerationIdCollector, empty())) + zipRight(set(resolvedSourceGenerationIdCollector, empty())), + zipRight(set(childGenerationRootCollector, empty())) ); export const getCollectedGenerationIds: Effect = get( @@ -79,6 +88,7 @@ export interface GenerationIdEntry { readonly isCacheHit: boolean; readonly countsTowardUsage: boolean; readonly isResolvedSource: boolean; + readonly shouldResolveChildren: boolean; } export const getCollectedGenerationIdEntries: Effect< @@ -88,13 +98,15 @@ export const getCollectedGenerationIdEntries: Effect< get(cacheHitGenerationIdCollector), get(auxiliaryUsageGenerationIdCollector), get(resolvedSourceGenerationIdCollector), + get(childGenerationRootCollector), ]).pipe( - map(([ids, cacheHitIds, auxiliaryIds, resolvedSourceIds]) => + map(([ids, cacheHitIds, auxiliaryIds, resolvedSourceIds, childRootIds]) => [...ids].map((id) => ({ id, isCacheHit: has(cacheHitIds, id), countsTowardUsage: !has(auxiliaryIds, id), isResolvedSource: has(resolvedSourceIds, id), + shouldResolveChildren: has(childRootIds, id), })) ) ); diff --git a/src/runtime/generation-resolver.test.ts b/src/runtime/generation-resolver.test.ts index 6c2a15c..b2467eb 100644 --- a/src/runtime/generation-resolver.test.ts +++ b/src/runtime/generation-resolver.test.ts @@ -69,7 +69,7 @@ describe("resolveCollectedGenerations", () => { resolveSourceGeneration: (generationId) => { requested.push(generationId); return succeed({ - sourceId: `source-of-${generationId}`, + sourceIds: [`source-of-${generationId}`], usage: SOURCE_USAGE, }); }, @@ -107,7 +107,7 @@ describe("resolveCollectedGenerations", () => { includeUsage: options?.includeUsage, }); return succeed({ - sourceId: `source-of-${generationId}`, + sourceIds: [`source-of-${generationId}`], ...(options?.includeUsage === false ? {} : { usage: SOURCE_USAGE }), }); }, @@ -154,7 +154,7 @@ describe("resolveCollectedGenerations", () => { resolveSourceGeneration: (generationId, options) => { requested.push(generationId); return succeed({ - sourceId: generationId, + sourceIds: [generationId], ...(options?.includeUsage === false ? {} : { usage: SOURCE_USAGE }), }); }, @@ -179,7 +179,7 @@ describe("resolveCollectedGenerations", () => { it("omits usage for entries that resolve without usage", async () => { const resolver: GenerationResolverService = { resolveSourceGeneration: (generationId) => - succeed({ sourceId: `source-of-${generationId}` }), + succeed({ sourceIds: [`source-of-${generationId}`] }), }; const resolved = await runPromise( resetGenerationIds.pipe( @@ -191,6 +191,31 @@ describe("resolveCollectedGenerations", () => { expect(resolved.ids).toEqual(["source-of-gen-dummy"]); expect(resolved.replayedUsage).toBeUndefined(); }); + it("replaces marked server-tools roots with their child generation ids", async () => { + const requested: { + id: string; + includeRelated: boolean | undefined; + }[] = []; + const resolver: GenerationResolverService = { + resolveSourceGeneration: (generationId, options) => { + requested.push({ + id: generationId, + includeRelated: options?.includeRelated, + }); + return succeed({ sourceIds: ["gen-child-1", "gen-child-2"] }); + }, + }; + const resolved = await runPromise( + resetGenerationIds.pipe( + flatMap(() => recordGenerationId("gen-root", false, false, true)), + flatMap(() => resolveCollectedGenerations), + provideService(GenerationResolver, resolver) + ) + ); + expect(requested).toEqual([{ id: "gen-root", includeRelated: true }]); + expect(resolved.ids).toEqual(["gen-child-1", "gen-child-2"]); + expect(resolved.replayedUsage).toBeUndefined(); + }); }); describe("makeOpenRouterGenerationResolver", () => { @@ -219,7 +244,7 @@ describe("makeOpenRouterGenerationResolver", () => { resolver.resolveSourceGeneration("gen-dummy") ); expect(resolved).toEqual({ - sourceId: "gen-original", + sourceIds: ["gen-original"], usage: SOURCE_USAGE, }); expect(getCalls()).toEqual([ @@ -240,7 +265,7 @@ describe("makeOpenRouterGenerationResolver", () => { const resolved = await runPromise( resolver.resolveSourceGeneration("gen-dummy", { includeUsage: false }) ); - expect(resolved).toEqual({ sourceId: "gen-original" }); + expect(resolved).toEqual({ sourceIds: ["gen-original"] }); expect(getCalls()).toEqual([ "https://example.com/api/v1/generation?id=gen-dummy", ]); @@ -265,7 +290,7 @@ describe("makeOpenRouterGenerationResolver", () => { const resolved = await runPromise( resolver.resolveSourceGeneration("gen-dummy") ); - expect(resolved?.sourceId).toBe("gen-original"); + expect(resolved?.sourceIds).toEqual(["gen-original"]); expect(getCalls().filter((url) => url.includes("gen-dummy")).length).toBe( 3 ); @@ -306,7 +331,10 @@ describe("makeOpenRouterGenerationResolver", () => { const resolved = await runPromise( resolver.resolveSourceGeneration("gen-source") ); - expect(resolved).toEqual({ sourceId: "gen-source", usage: SOURCE_USAGE }); + expect(resolved).toEqual({ + sourceIds: ["gen-source"], + usage: SOURCE_USAGE, + }); expect(getCalls().length).toBe(1); }); it("warns when the source generation has no usage fields", async () => { @@ -371,7 +399,45 @@ describe("makeOpenRouterGenerationResolver", () => { const resolved = await runPromise( resolver.resolveSourceGeneration("gen-dummy") ); - expect(resolved).toEqual({ sourceId: "gen-original" }); + expect(resolved).toEqual({ sourceIds: ["gen-original"] }); expect(getCalls().length).toBe(2); }); + it("resolves a server-tools root to its leaf generations", async () => { + const getCalls = mockFetch((url) => { + if (url.includes("gen-root")) { + return jsonResponse({ + data: { + response_cache_source_id: null, + related_generation_ids: ["gen-child-1", "gen-child-2"], + }, + }); + } + return jsonResponse({ + data: { + response_cache_source_id: null, + related_generation_ids: [], + }, + }); + }); + const resolver = makeOpenRouterGenerationResolver({ + apiKey: "test-key", + baseUrl: "https://example.com", + pollIntervalMs: 1, + maxAttempts: 1, + }); + const resolved = await runPromise( + resolver.resolveSourceGeneration("gen-root", { + includeRelated: true, + includeUsage: false, + }) + ); + expect(resolved).toEqual({ + sourceIds: ["gen-child-1", "gen-child-2"], + }); + expect(getCalls()).toEqual([ + "https://example.com/api/v1/generation?id=gen-root&include_related=true", + "https://example.com/api/v1/generation?id=gen-child-1&include_related=true", + "https://example.com/api/v1/generation?id=gen-child-2&include_related=true", + ]); + }); }); diff --git a/src/runtime/generation-resolver.ts b/src/runtime/generation-resolver.ts index aeaff1f..cd8cb2d 100644 --- a/src/runtime/generation-resolver.ts +++ b/src/runtime/generation-resolver.ts @@ -35,14 +35,17 @@ export interface ReplayedUsage { } export interface ResolvedSourceGeneration { - readonly sourceId: string; + readonly sourceIds: readonly string[]; readonly usage?: ReplayedUsage; } export interface GenerationResolverService { readonly resolveSourceGeneration: ( generationId: string, - options?: { readonly includeUsage?: boolean } + options?: { + readonly includeUsage?: boolean; + readonly includeRelated?: boolean; + } ) => Effect; } @@ -58,6 +61,7 @@ const GenerationLookupSchema = z.object({ native_tokens_reasoning: z.number().nullish(), total_cost: z.number().nullish(), generation_time: z.number().nullish(), + related_generation_ids: z.array(z.string()).nullish(), }), }); @@ -82,6 +86,7 @@ export interface GenerationResolverConfig { const DEFAULT_POLL_INTERVAL_MS = 5000; const DEFAULT_MAX_ATTEMPTS = 12; const RESOLVE_CONCURRENCY = 8; +const MAX_RELATED_DEPTH = 16; function normalizeBaseUrl(baseUrl: string): string { const trimmed = baseUrl.replace(/\/+$/u, ""); @@ -125,14 +130,15 @@ export function makeOpenRouterGenerationResolver( whileInput((error: GenerationLookupError) => error.retryable) ); const lookupOnce = ( - generationId: string + generationId: string, + includeRelated: boolean ): Effect => tryPromise({ try: async ( signal ): Promise<{ readonly status: number } | { readonly json: unknown }> => { const response = await fetch( - `${baseUrl}/generation?id=${encodeURIComponent(generationId)}`, + `${baseUrl}/generation?id=${encodeURIComponent(generationId)}${includeRelated ? "&include_related=true" : ""}`, { headers: { Authorization: `Bearer ${config.apiKey}` }, signal, @@ -171,61 +177,101 @@ export function makeOpenRouterGenerationResolver( }) ); const lookupGeneration = ( - generationId: string + generationId: string, + includeRelated = false ): Effect => - lookupOnce(generationId).pipe(retry(pollSchedule(maxAttempts))); - const lookupSourceUsage = ( - sourceId: string - ): Effect => - lookupGeneration(sourceId).pipe( - map((data): ReplayedUsage | undefined => usageFromLookup(data, sourceId)), + lookupOnce(generationId, includeRelated).pipe( + retry(pollSchedule(maxAttempts)) + ); + const resolveSourceGeneration = ( + generationId: string, + options: + | { + readonly includeUsage?: boolean; + readonly includeRelated?: boolean; + } + | undefined, + visited: ReadonlySet, + depth: number + ): Effect => { + if (depth >= MAX_RELATED_DEPTH || visited.has(generationId)) { + return sync(() => { + wLog("Generation relation resolution reached a cycle or depth limit", { + generation_id: generationId, + }); + return undefined; + }); + } + const nextVisited = new Set(visited).add(generationId); + const includeRelated = options?.includeRelated === true; + return lookupGeneration(generationId, includeRelated).pipe( + flatMap((data) => { + const cacheSourceId = data.response_cache_source_id; + if (cacheSourceId !== null && cacheSourceId !== undefined) { + if (options?.includeUsage === false && !includeRelated) { + return succeed({ + sourceIds: [cacheSourceId], + }); + } + return resolveSourceGeneration( + cacheSourceId, + options, + nextVisited, + depth + 1 + ).pipe( + map( + (resolved): ResolvedSourceGeneration => + resolved ?? { sourceIds: [cacheSourceId] } + ) + ); + } + const relatedIds = includeRelated + ? (data.related_generation_ids ?? []) + : []; + if (relatedIds.length === 0) { + return succeed({ + sourceIds: [generationId], + ...(options?.includeUsage === false + ? {} + : { usage: usageFromLookup(data, generationId) }), + }); + } + return forEach( + relatedIds, + (relatedId) => + resolveSourceGeneration(relatedId, options, nextVisited, depth + 1), + { concurrency: RESOLVE_CONCURRENCY } + ).pipe( + map((resolved) => { + if (resolved.some((entry) => entry === undefined)) { + return undefined; + } + const entries = resolved.filter(isDefinedAndNotNull); + let usage: ReplayedUsage | undefined; + for (const entry of entries) { + usage = sumReplayedUsage(usage, entry.usage); + } + return { + sourceIds: entries.flatMap((entry) => entry.sourceIds), + ...(usage !== undefined && { usage }), + } satisfies ResolvedSourceGeneration; + }) + ); + }), catchAll((error) => sync(() => { - wLog("Failed to fetch source generation usage", { - generation_id: sourceId, + wLog("Failed to resolve generation relations", { + generation_id: generationId, error: error.message, }); return undefined; }) ) ); + }; return { resolveSourceGeneration: (generationId, options) => - lookupGeneration(generationId).pipe( - flatMap((data) => { - const dummySourceId = data.response_cache_source_id; - const sourceId = - dummySourceId === null || - dummySourceId === undefined || - dummySourceId.length === 0 - ? generationId - : dummySourceId; - if (options?.includeUsage === false) { - return succeed({ sourceId }); - } - if (sourceId === generationId) { - return succeed({ - sourceId, - usage: usageFromLookup(data, generationId), - }); - } - return lookupSourceUsage(sourceId).pipe( - map((usage): ResolvedSourceGeneration => ({ - sourceId, - ...(usage !== undefined && { usage }), - })) - ); - }), - catchAll((error) => - sync(() => { - wLog("Failed to resolve cache-hit source generation", { - generation_id: generationId, - error: error.message, - }); - return undefined; - }) - ) - ), + resolveSourceGeneration(generationId, options, new Set(), 0), }; } @@ -235,7 +281,7 @@ export interface ResolvedGenerations { } interface ResolvedEntry { - readonly id: string; + readonly ids: readonly string[]; readonly usage?: ReplayedUsage; } @@ -244,24 +290,25 @@ function resolveEntry( resolver: GenerationResolverService ): Effect { if (entry.isResolvedSource && !entry.countsTowardUsage) { - return succeed({ id: entry.id }); + return succeed({ ids: [entry.id] }); } - return entry.isCacheHit + return entry.isCacheHit || entry.shouldResolveChildren ? resolver .resolveSourceGeneration(entry.id, { - includeUsage: entry.countsTowardUsage, + includeUsage: entry.isCacheHit && entry.countsTowardUsage, + includeRelated: entry.shouldResolveChildren, }) .pipe( map((resolved): ResolvedEntry => resolved === undefined - ? { id: entry.id } + ? { ids: [entry.id] } : { - id: resolved.sourceId, + ids: resolved.sourceIds, ...(resolved.usage && { usage: resolved.usage }), } ) ) - : succeed({ id: entry.id }); + : succeed({ ids: [entry.id] }); } function sumReplayedUsage( @@ -288,7 +335,12 @@ export const resolveCollectedGenerations: Effect = gen( function* () { const entries = yield* getCollectedGenerationIdEntries; const resolver = yield* serviceOption(GenerationResolver); - if (isNone(resolver) || entries.every((entry) => !entry.isCacheHit)) { + if ( + isNone(resolver) || + entries.every( + (entry) => !entry.isCacheHit && !entry.shouldResolveChildren + ) + ) { return { ids: entries.map((entry) => entry.id) }; } const resolved = yield* forEach( @@ -301,7 +353,7 @@ export const resolveCollectedGenerations: Effect = gen( replayedUsage = sumReplayedUsage(replayedUsage, entry.usage); } return { - ids: resolved.map((entry) => entry.id), + ids: resolved.flatMap((entry) => entry.ids), ...(replayedUsage !== undefined && { replayedUsage }), }; }