fix(search): resolve server-tool generation ids

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;

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,

src/judge/judge.test.ts

@@ -119,6 +119,7 @@ describe("judgeCall", () => {
         isCacheHit: true,
         countsTowardUsage: true,
         isResolvedSource: false,
+        shouldResolveChildren: false,
       },
     ]);
   });

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 {

src/providers/responses-client.test.ts

@@ -345,7 +345,7 @@
     expect(initialIdentifiers.generationId).toBe("gen-789");
   });
   it("includes response headers once when stream iteration raises an SDK error", async () => {
     async function* stream(): AsyncGenerator<StreamEvents> {
       throw new OpenRouterError("upstream", {
         response: new Response("failure", {
           status: 503,
@@ -453,7 +453,7 @@
               const responses = yield* Responses;
               yield* responses.send(
                 { model: "m", input: [] },
-                { timeoutMs: 1000 }
+                { timeoutMs: 1000, resolveGenerationChildren: true }
               );
             })
           ),
@@ -472,6 +472,7 @@
           isCacheHit: true,
           countsTowardUsage: true,
           isResolvedSource: true,
+          shouldResolveChildren: true,
         },
       ]);
     } finally {

src/providers/responses-client.ts

@@ -94,6 +94,7 @@ export interface ResponsesSendOptions {
   readonly extraBody?: Readonly<Record<string, unknown>>;
   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<Responses> {
                 ? cacheSourceId
                 : result.generationId,
               isCacheHit,
-              isCacheHit && cacheSourceId !== undefined
+              isCacheHit && cacheSourceId !== undefined,
+              options.resolveGenerationChildren ?? false
             ).pipe(map(() => result))
           : fail(
               new ResponsesError({

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,
       },
     ]);
   });

src/runtime/generation-ids.ts

@@ -20,6 +20,8 @@ export const auxiliaryUsageGenerationIdCollector =
 export const resolvedSourceGenerationIdCollector =
   unsafeMakeHashSet<string>(empty());
 
+export const childGenerationRootCollector = unsafeMakeHashSet<string>(empty());
+
 export const auxiliaryUsageRef: FiberRef<boolean> = unsafeMake(false);
 
 export function withAuxiliaryUsage<A, E, R>(
@@ -31,7 +33,8 @@ export function withAuxiliaryUsage<A, E, R>(
 export function recordGenerationId(
   id: string | null | undefined,
   isCacheHit = false,
-  isResolvedSource = false
+  isResolvedSource = false,
+  shouldResolveChildren = false
 ): Effect<void> {
   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<void> = set(
 ).pipe(
   zipRight(set(cacheHitGenerationIdCollector, empty<string>())),
   zipRight(set(auxiliaryUsageGenerationIdCollector, empty<string>())),
-  zipRight(set(resolvedSourceGenerationIdCollector, empty<string>()))
+  zipRight(set(resolvedSourceGenerationIdCollector, empty<string>())),
+  zipRight(set(childGenerationRootCollector, empty<string>()))
 );
 
 export const getCollectedGenerationIds: Effect<readonly string[]> = 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),
     }))
   )
 );

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",
+    ]);
+  });
 });