Collect a batch without losing individual failures

import { ApiClientError, ApiClientErrorCode, ApiClientErrorType, type CapabilityClient, type RuntimeBatchResult } from "@cavi-ai/api-client";

export async function collectBatch(client: CapabilityClient, batchId: string) {
  const status = await client.getBatch(batchId);
  if (!status.ok) return status;
  const batch = status.data;
  if (!status.data.resultsAvailable) {
    const kind = ["completed", "failed", "cancelled", "expired"].includes(batch.status)
      ? "terminal" as const : "pending" as const;
    return { ok: true as const, source: status.source, data: { kind, batch } };
  }
  const results = await client.getBatchResults(batchId);
  if (!results.ok) return results;
  const items = new Map<string, RuntimeBatchResult>();
  for (const item of results.data) {
    if (!item.customId.trim() || items.has(item.customId)) {
      throw new ApiClientError("Batch results have missing or duplicate correlation IDs.", {
        type: ApiClientErrorType.Validation, code: ApiClientErrorCode.ProtocolMismatch,
        cause: { batch, customId: item.customId },
      });
    }
    items.set(item.customId, item);
  }
  return {
    ok: true as const,
    source: results.source,
    data: {
      kind: "results" as const,
      batch,
      items,
    },
  };
}

import assert from "node:assert/strict";
import { test, type TestContext } from "node:test";
import { createApiClient, createRuntimeProviderRegistry, isApiClientError, isAuthError } from "@cavi-ai/api-client";
import { createClaudeProviderModule } from "@cavi-ai/api-client/providers/claude/messages";
import { createCodexProviderModule } from "@cavi-ai/api-client/providers/codex/runtime";
import { collectBatch } from "./batch-collector.js";

function fixture(t: TestContext, provider: "claude" | "codex", reply: (url: string) => Response) {
  const requests: string[] = [];
  const client = createApiClient(provider, {
    registry: createRuntimeProviderRegistry({ modules: [provider === "claude"
      ? createClaudeProviderModule({ apiKey: "fixture-key" })
      : createCodexProviderModule({ apiKey: "fixture-key" })] }),
    fetchImpl: async (url, init) => {
      assert.equal(init?.method ?? "GET", "GET", "resuming must never submit or cancel work");
      requests.push(String(url));
      return reply(String(url));
    },
  });
  t.after(() => client.dispose());
  return { client, requests };
}

test("keeps an active batch pending without fetching result files", async (t) => {
  const app = fixture(t, "codex", () => Response.json({ id: "saved-1", status: "in_progress" }));
  const result = await collectBatch(app.client, "saved-1");
  assert.equal(result.ok, true);
  if (!result.ok) return;
  assert.equal(result.data.kind, "pending");
  assert.equal(result.data.batch.batch_id, "saved-1");
  assert.equal(app.requests.length, 1);
});

test("stops rescheduling known terminal states with no available results", async (t) => {
  for (const status of ["failed", "cancelled", "expired", "completed"]) {
    const app = fixture(t, "codex", () => Response.json({ id: "saved-1", status }));
    const result = await collectBatch(app.client, "saved-1");
    assert.equal(result.ok, true);
    if (!result.ok) continue;
    assert.equal(result.data.kind, "terminal", status);
    assert.equal(result.data.batch.status, status);
    assert.equal(app.requests.length, 1);
  }
});

test("collects available results even when the batch was cancelled", async (t) => {
  const app = fixture(t, "codex", (url) => url.endsWith("/content")
    ? new Response(JSON.stringify({ custom_id: "one", response: { status_code: 200,
      body: { id: "answer-1", status: "completed", output_text: "Retained answer" },
    } }))
    : Response.json({ id: "saved-1", status: "cancelled", output_file_id: "output-1" }));
  const result = await collectBatch(app.client, "saved-1");
  assert.equal(result.ok, true);
  if (!result.ok || result.data.kind !== "results") return assert.fail("Expected available results");
  assert.equal(result.source, "live");
  assert.equal(result.data.batch.status, "cancelled");
  assert.equal(result.data.items.get("one")?.run?.output, "Retained answer");
  assert.ok(app.requests.every((url) => url.endsWith("/batches/saved-1") || url.endsWith("/files/output-1/content")));
});

function claudeResults(items: unknown[]) {
  return (url: string) => url.endsWith("/results")
    ? new Response(items.map((item) => JSON.stringify(item)).join("\n"))
    : Response.json({ id: "saved-1", processing_status: "ended" });
}

test("correlates out-of-order items and retains every outcome", async (t) => {
  const app = fixture(t, "claude", claudeResults([
    { custom_id: "two", result: { type: "errored", error: { message: "Private item diagnostic" } } },
    { custom_id: "one", result: { type: "succeeded", message: { id: "answer-1", content: [{ type: "text", text: "Answer" }], stop_reason: "end_turn" } } },
    { custom_id: "three", result: { type: "canceled" } },
    { custom_id: "four", result: { type: "expired" } },
    { custom_id: "five", result: { type: "future-outcome" } },
  ]));
  const result = await collectBatch(app.client, "saved-1");
  assert.equal(result.ok, true);
  if (!result.ok || result.data.kind !== "results") return assert.fail("Expected results");
  assert.equal(result.data.batch.batch_id, "saved-1");
  assert.equal(result.data.items.size, 5);
  assert.equal(result.data.items.get("one")?.run?.output, "Answer");
  assert.equal(result.data.items.get("two")?.outcome, "errored");
  assert.equal(result.data.items.get("three")?.outcome, "canceled");
  assert.equal(result.data.items.get("four")?.outcome, "expired");
  assert.equal(result.data.items.get("five")?.outcome, "future-outcome");
});

test("rejects missing or duplicated correlation IDs without overwriting items", async (t) => {
  for (const ids of [["same", "same"], [""], [undefined]]) {
    const app = fixture(t, "claude", claudeResults(ids.map((id) => ({ custom_id: id, result: { type: "canceled" } }))));
    await assert.rejects(collectBatch(app.client, "saved-1"), (error) => {
      assert.ok(isApiClientError(error));
      assert.equal(error.code, "protocol_mismatch");
      assert.equal((error.cause as { batch: { batch_id: string } }).batch.batch_id, "saved-1");
      return true;
    });
  }
});

test("preserves a status retrieval gap", async (t) => {
  const app = fixture(t, "codex", () => Response.json({}, { status: 503 }));
  const result = await collectBatch(app.client, "saved-1");
  assert.equal(result.ok, false);
  if (result.ok) return;
  assert.equal(result.gap.reason, "backend-unavailable");
  assert.equal(app.requests.length, 1);
});

test("preserves a result retrieval gap without reporting empty success", async (t) => {
  const app = fixture(t, "claude", (url) => url.endsWith("/results")
    ? Response.json({}, { status: 503 })
    : Response.json({ id: "saved-1", processing_status: "ended" }));
  const result = await collectBatch(app.client, "saved-1");
  assert.equal(result.ok, false);
  if (result.ok) return;
  assert.equal(result.gap.reason, "backend-unavailable");
  assert.equal(app.requests.length, 2);
});

test("keeps authentication failures as exceptions", async (t) => {
  const app = fixture(t, "codex", () => Response.json({}, { status: 401 }));
  await assert.rejects(collectBatch(app.client, "saved-1"), isAuthError);
  assert.equal(app.requests.length, 1);
});

import { ApiClientError, ApiClientErrorCode, ApiClientErrorType, type CapabilityClient, type RuntimeBatchRequest } from "@cavi-ai/api-client";

export async function submitBatchAndCollect(
  client: CapabilityClient,
  requests: RuntimeBatchRequest[],
  options: { maxPolls?: number; pollIntervalMs?: number } = {},
) {
  const { maxPolls = 60, pollIntervalMs = 60_000 } = options;
  if (!Number.isInteger(maxPolls) || maxPolls < 0 || !Number.isFinite(pollIntervalMs) || pollIntervalMs < 0) {
    throw new ApiClientError("Use a non-negative poll count and interval.", {
      type: ApiClientErrorType.Validation, code: ApiClientErrorCode.InvalidRequest,
    });
  }
  const submitted = await client.submitBatch(requests);
  if (!submitted.ok) throw new ApiClientError(submitted.gap.note, {
    code: ApiClientErrorCode.RequestFailed, cause: submitted.gap,
  });
  let batch = submitted.data;
  for (let poll = 0; !batch.resultsAvailable && poll < maxPolls; poll += 1) {
    if (["failed", "cancelled"].includes(batch.status)) break;
    await new Promise((resolve) => setTimeout(resolve, pollIntervalMs));
    const next = await client.getBatch(batch.batch_id);
    if (!next.ok) throw new ApiClientError(next.gap.note, {
      code: ApiClientErrorCode.RequestFailed, cause: next.gap,
    });
    batch = next.data;
  }
  if (!batch.resultsAvailable) {
    // Keep the ID: stopping this local wait does not cancel the backend batch.
    throw new ApiClientError(`Batch ${batch.batch_id} is not ready (status: ${batch.status}); retrieve it later or cancel explicitly.`, {
      code: ApiClientErrorCode.RequestFailed, cause: batch,
    });
  }
  const results = await client.getBatchResults(batch.batch_id);
  if (!results.ok) throw new ApiClientError(results.gap.note, {
    code: ApiClientErrorCode.RequestFailed, cause: results.gap,
  });
  // Callers must inspect each item's outcome, including errored/canceled/expired.
  return results.data;
}

import { runtimeSupports, type RuntimeClient, type RuntimeBatchRequest } from "@cavi-ai/api-client";

export async function submitWhenBatchIsAvailable(client: RuntimeClient, requests: RuntimeBatchRequest[]) {
  const capabilities = await client.getRuntimeCapabilities();
  if (!runtimeSupports(capabilities, "batch") || !client.submitBatch) return null;
  return client.submitBatch(requests);
}