feat(eval): add batch-evaluation simulate (dataset simulations)

src/core/eval/simulate.ts

@@ -0,0 +1,120 @@
+import type { SessionMetadataShape } from "@aws-sdk/client-bedrock-agentcore";
+import { InputValidationError } from "../../errors";
+
+// Scenario is one dataset row for replay. Local to this module on purpose — it is
+// not a handler contract, so it does not live in handlers/eval/types.tsx. Field
+// names mirror the dataset JSONL (snake_case in, camelCase here).
+export type Scenario = {
+  scenarioId: string;
+  turns: { input: string; expectedResponse?: string }[];
+  assertions?: string[];
+  expectedTrajectory?: string[];
+};
+
+// parseScenarios reads dataset JSONL (one scenario per line) into Scenario records.
+// Each scenario needs a non-empty, unique `scenario_id` — the id is the join key
+// between the session created for it and its ground truth, so a missing/duplicate
+// id silently misassigns ground truth to the wrong session.
+export function parseScenarios(text: string): Scenario[] {
+  const scenarios: Scenario[] = [];
+  const seen = new Set<string>();
+  for (const line of text.split("\n")) {
+    const trimmed = line.trim();
+    if (!trimmed) continue;
+    let row: Record<string, unknown>;
+    try {
+      row = JSON.parse(trimmed) as Record<string, unknown>;
+    } catch {
+      throw new InputValidationError("dataset contains a line that is not valid JSON");
+    }
+    const scenario = toScenario(row);
+    if (!scenario.scenarioId) {
+      throw new InputValidationError("dataset scenario is missing 'scenario_id'");
+    }
+    if (seen.has(scenario.scenarioId)) {
+      throw new InputValidationError(
+        `dataset has a duplicate scenario_id: "${scenario.scenarioId}"`,
+      );
+    }
+    if (scenario.turns.length === 0) {
+      throw new InputValidationError(`scenario "${scenario.scenarioId}" has no turns`);
+    }
+    seen.add(scenario.scenarioId);
+    scenarios.push(scenario);
+  }
+  if (scenarios.length === 0) throw new InputValidationError("dataset has no scenarios");
+  return scenarios;
+}
+
+function toScenario(row: Record<string, unknown>): Scenario {
+  const turns = Array.isArray(row.turns) ? row.turns : [];
+  return {
+    scenarioId: String(row.scenario_id ?? ""),
+    turns: turns.map((t: Record<string, unknown>) => ({
+      input: String(t.input ?? ""),
+      expectedResponse: t.expected_response as string | undefined,
+    })),
+    assertions: row.assertions as string[] | undefined,
+    expectedTrajectory: row.expected_trajectory as string[] | undefined,
+  };
+}
+
+// loadDatasetFile reads scenarios from a local JSONL path. The dataset-id path is
+// handled by the caller (downloadDataset to a temp file, then this).
+export async function loadDatasetFile(path: string): Promise<Scenario[]> {
+  return parseScenarios(await Bun.file(path).text());
+}
+
+// runScenarios runs `worker` over every scenario with bounded concurrency. A failed
+// worker doesn't sink the run — the failure is captured (so the caller can report
+// on all-failed) but drops that scenario. Returns ok results + the first error we
+// saw, which the caller can surface to explain a total failure.
+export type ScenarioRun<T> = { ok: T[]; failed: number; firstError?: Error };
+export async function runScenarios<T>(
+  scenarios: Scenario[],
+  worker: (scenario: Scenario) => Promise<T>,
+  concurrency = 5,
+): Promise<ScenarioRun<T>> {
+  const ok: T[] = [];
+  let failed = 0;
+  let firstError: Error | undefined;
+  let next = 0;
+  const run = async (): Promise<void> => {
+    while (next < scenarios.length) {
+      const scenario = scenarios[next++]!;
+      try {
+        ok.push(await worker(scenario));
+      } catch (error) {
+        failed++;
+        if (!firstError) firstError = error instanceof Error ? error : new Error(String(error));
+      }
+    }
+  };
+  await Promise.all(Array.from({ length: Math.min(concurrency, scenarios.length) }, run));
+  return { ok, failed, firstError };
+}
+
+// toSessionMetadata maps a scenario's ground truth onto the session the replay
+// created — the batch service's per-session ground-truth shape (inline arm).
+// Omit array fields when empty: the service rejects zero-length `turns` /
+// `assertions` (`length >= 1`) rather than treating an empty array as "no data".
+export function toSessionMetadata(scenario: Scenario, sessionId: string): SessionMetadataShape {
+  const turns = scenario.turns
+    .filter((t) => t.expectedResponse !== undefined)
+    .map((t) => ({ expectedResponse: { text: t.expectedResponse! } }));
+  const assertions = scenario.assertions?.map((text) => ({ text }));
+  return {
+    sessionId,
+    testScenarioId: scenario.scenarioId,
+    groundTruth: {
+      inline: {
+        ...(assertions && assertions.length > 0 && { assertions }),
+        ...(scenario.expectedTrajectory &&
+          scenario.expectedTrajectory.length > 0 && {
+            expectedTrajectory: { toolNames: scenario.expectedTrajectory },
+          }),
+        ...(turns.length > 0 && { turns }),
+      },
+    },
+  };
+}

src/core/invokeRuntime.ts

@@ -0,0 +1,177 @@
+import { randomUUID } from "node:crypto";
+import { InvokeAgentRuntimeCommand } from "@aws-sdk/client-bedrock-agentcore";
+import type { RuntimeInvokeRequest, RuntimeInvokeResponse } from "../handlers/runtime/types";
+import type { Logger } from "../logging";
+import type { AwsClients, CoreFetch, CoreOptions } from "./types";
+import { abortable } from "./abortable";
+import { toClientConfig } from "./utils";
+
+// InvokeRuntimeDeps is the slice of a Core client an invoke needs. Passed in as a
+// bag (not a sibling client) so both RuntimeClient and EvalClient.simulate can call
+// these free functions off their own `this.clients`/`this.fetch`/`this.logger`.
+export type InvokeRuntimeDeps = {
+  clients: AwsClients;
+  fetch: CoreFetch;
+  logger: Logger;
+};
+
+async function* emptyBody(): AsyncGenerator<Uint8Array> {}
+
+// invokeRuntime dispatches by auth mode: a bearer token routes to the CUSTOM_JWT
+// (raw fetch) path, otherwise the SigV4 SDK path.
+export async function invokeRuntime(
+  deps: InvokeRuntimeDeps,
+  request: RuntimeInvokeRequest,
+  options: CoreOptions,
+  signal?: AbortSignal,
+): Promise<RuntimeInvokeResponse> {
+  const { runtimeId, bearerToken } = request;
+  if (bearerToken !== undefined) {
+    const logger = deps.logger.child({
+      operation: "invokeRuntime",
+      authMode: "CUSTOM_JWT",
+      runtimeId,
+      qualifier: request.qualifier,
+      region: options.region,
+    });
+    return invokeRuntimeWithCustomJwt(deps, request, bearerToken, options, logger, signal);
+  }
+  return invokeRuntimeWithIam(deps, request, options, signal);
+}
+
+async function invokeRuntimeWithCustomJwt(
+  deps: InvokeRuntimeDeps,
+  request: RuntimeInvokeRequest,
+  bearerToken: string,
+  options: CoreOptions,
+  logger: Logger,
+  signal?: AbortSignal,
+): Promise<RuntimeInvokeResponse> {
+  const client = deps.clients.data(toClientConfig(options));
+  const endpoint = client.config.endpointProvider({
+    Region: options.region,
+    Endpoint: options.endpointUrl,
+  });
+  const url = new URL(endpoint.url);
+  if (url.protocol !== "https:") {
+    throw new TypeError("CUSTOM_JWT requires an HTTPS endpoint");
+  }
+  url.pathname = `${url.pathname.replace(/\/?$/, "/")}runtimes/${encodeURIComponent(request.runtimeId)}/invocations`;
+  url.search = new URLSearchParams({
+    accountId: request.accountId,
+    qualifier: request.qualifier,
+  }).toString();
+  const headers = new Headers(request.applicationHeaders);
+  try {
+    headers.set("Authorization", `Bearer ${bearerToken}`);
+  } catch {
+    throw new TypeError("Invalid bearer token");
+  }
+  try {
+    for (const [name, value] of [
+      ["Content-Type", request.contentType],
+      ["Accept", request.accept],
+      ["Mcp-Session-Id", request.mcpSessionId],
+      ["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id", request.runtimeSessionId ?? randomUUID()],
+      ["Mcp-Protocol-Version", request.mcpProtocolVersion],
+      ["Mcp-Method", request.mcpMethod],
+      ["Mcp-Name", request.mcpName],
+      ["X-Amzn-Bedrock-AgentCore-Runtime-User-Id", request.runtimeUserId],
+      ["X-Amzn-Trace-Id", request.traceId],
+      ["traceparent", request.traceParent],
+      ["tracestate", request.traceState],
+      ["baggage", request.baggage],
+    ] as const) {
+      if (value !== undefined) headers.set(name, value);
+    }
+  } catch {
+    throw new TypeError("Invalid Runtime request header");
+  }
+  let response: Response;
+  try {
+    response = await deps.fetch(url, {
+      method: "POST",
+      redirect: "error",
+      headers,
+      body: request.payload as RequestInit["body"],
+      signal,
+    });
+  } catch (error) {
+    if (signal?.aborted) throw signal.reason ?? error;
+    logger
+      .child({
+        errorName:
+          error instanceof TypeError
+            ? "TypeError"
+            : error instanceof Error
+              ? "Error"
+              : typeof error,
+      })
+      .debug("Runtime invocation transport failed");
+    throw new Error("Runtime invocation failed");
+  }
+  if (!response.ok) {
+    logger
+      .child({ httpStatusCode: response.status })
+      .debug("Runtime invocation returned a non-success response");
+    await response.body?.cancel().catch(() => undefined);
+    throw new Error(`HTTP ${response.status}`);
+  }
+  const body = (response.body as AsyncIterable<Uint8Array> | null) ?? emptyBody();
+  return {
+    statusCode: response.status,
+    contentType: response.headers.get("content-type") ?? "",
+    runtimeSessionId:
+      response.headers.get("x-amzn-bedrock-agentcore-runtime-session-id") ?? undefined,
+    mcpSessionId: response.headers.get("mcp-session-id") ?? undefined,
+    mcpProtocolVersion: response.headers.get("mcp-protocol-version") ?? undefined,
+    traceId: response.headers.get("x-amzn-trace-id") ?? undefined,
+    traceParent: response.headers.get("traceparent") ?? undefined,
+    traceState: response.headers.get("tracestate") ?? undefined,
+    baggage: response.headers.get("baggage") ?? undefined,
+    body: signal ? abortable(body, signal) : body,
+  };
+}
+
+async function invokeRuntimeWithIam(
+  deps: InvokeRuntimeDeps,
+  request: RuntimeInvokeRequest,
+  options: CoreOptions,
+  signal?: AbortSignal,
+): Promise<RuntimeInvokeResponse> {
+  const { runtimeId, applicationHeaders, bearerToken: _bearerToken, ...input } = request;
+  const command = new InvokeAgentRuntimeCommand({ ...input, agentRuntimeArn: runtimeId });
+  if (applicationHeaders?.length) {
+    command.middlewareStack.add(
+      (next) => async (args) => {
+        const sdkRequest = args.request as { headers: Record<string, string> };
+        for (const [name, value] of applicationHeaders) sdkRequest.headers[name] = value;
+        return next(args);
+      },
+      { step: "build", name: "runtimeApplicationHeaders" },
+    );
+  }
+  let response;
+  try {
+    response = await deps.clients.data(toClientConfig(options)).send(command, {
+      abortSignal: signal,
+    });
+  } catch (error) {
+    if (signal?.aborted) throw signal.reason ?? error;
+    throw error;
+  }
+
+  const body = (response.response as AsyncIterable<Uint8Array> | undefined) ?? emptyBody();
+  return {
+    statusCode: response.statusCode ?? 0,
+    contentType: response.contentType ?? "",
+    runtimeSessionId: response.runtimeSessionId,
+    mcpSessionId: response.mcpSessionId,
+    mcpProtocolVersion: response.mcpProtocolVersion,
+    traceId: response.traceId,
+    traceParent: response.traceParent,
+    traceState: response.traceState,
+    baggage: response.baggage,
+    body: signal ? abortable(body, signal) : body,
+  };
+}