@hraness/dawg 0.2.0 → 0.4.0

package/src/agent/gateway.ts

@@ -1,6 +1,9 @@
 import { readSseData } from "./sse.ts";
 
-/** The only model labels dawg accepts. Provider IDs are local configuration. */
+/**
+ * The two original model aliases. A model is either one of these or an exact
+ * `vendor/model` ID chosen from the model picker (see `models.ts`).
+ */
 export const GATEWAY_MODELS = Object.freeze(["opus-5.5", "sol-6.1"] as const);
 export type GatewayModel = (typeof GATEWAY_MODELS)[number];
 
@@ -26,15 +29,20 @@ export function isGatewayModel(value: unknown): value is GatewayModel {
   );
 }
 
-/** Resolve an allowlisted alias to its provider ID, rejecting anything else. */
+/**
+ * Resolve an alias (`opus-5.5`) or an exact `vendor/model` ID to the ID sent
+ * to the provider, rejecting anything else.
+ */
 export function resolveModelId(
   alias: string,
   overrides: Partial<Record<GatewayModel, string | undefined>> = {},
 ): string {
-  if (!isGatewayModel(alias))
+  if (!isGatewayModel(alias)) {
+    if (MODEL_ID_PATTERN.test(alias)) return alias;
     throw new Error(
-      `unknown model "${alias.slice(0, 32)}"; use ${GATEWAY_MODELS.join(" or ")}`,
+      `unknown model "${alias.slice(0, 32)}"; use ${GATEWAY_MODELS.join(" or ")} or a vendor/model ID`,
     );
+  }
   const id = overrides[alias] ?? DEFAULT_MODEL_IDS[alias];
   if (!MODEL_ID_PATTERN.test(id))
     throw new Error(`model ID for ${alias} must look like provider/model`);
@@ -68,7 +76,8 @@ export type ChatTool = {
 };
 
 export type ChatStreamRequest = {
-  model: GatewayModel;
+  /** An alias (`opus-5.5`) or an exact `vendor/model` ID. */
+  model: string;
   messages: readonly ChatMessage[];
   tools?: readonly ChatTool[];
   temperature?: number;
@@ -81,6 +90,8 @@ export type ChatStreamRequest = {
 /** Normalized stream events; provider chunk shapes never leave this module. */
 export type ChatStreamEvent =
   | { type: "text"; delta: string }
+  /** Progress worth a status line, such as a retry; never model output. */
+  | { type: "activity"; message: string }
   | {
       type: "tool-delta";
       index: number;
@@ -88,11 +99,25 @@ export type ChatStreamEvent =
       name?: string;
       arguments?: string;
     }
-  | { type: "finish"; reason: string };
+  | { type: "finish"; reason: string }
+  /**
+   * Token usage for one request, from the final chunk when the request asked
+   * for `stream_options.include_usage`. `costUsd` is the provider's own
+   * charge when it reports one (OpenRouter's `usage.cost`).
+   */
+  | {
+      type: "usage";
+      inputTokens: number;
+      outputTokens: number;
+      cachedInputTokens?: number;
+      costUsd?: number;
+    };
 
 export type GatewayClient = {
+  /** Which service this client talks to. */
+  readonly provider?: ApiProvider;
   /** The provider model ID that a turn will use (for diagnostics only). */
-  modelId(model: GatewayModel): string;
+  modelId(model: string): string;
   stream(
     request: ChatStreamRequest,
     signal?: AbortSignal,
@@ -114,21 +139,79 @@ type GatewayFetcher = (
   init?: RequestInit,
 ) => Promise<Response>;
 
+/** Retries after a failed request that has not yet produced a byte. */
+export const GATEWAY_RETRIES = 2;
+/** Base backoff; each retry doubles it, with +-50% jitter. */
+const RETRY_BASE_MS = 500;
+/** Response headers must arrive within this long, per attempt. */
+export const GATEWAY_HEADER_TIMEOUT_MS = 15_000;
+
+export type GatewayRetryOptions = Readonly<{
+  retries?: number;
+  baseMs?: number;
+  sleep?: (ms: number) => Promise<void>;
+  random?: () => number;
+}>;
+
+function retryableStatus(status: number): boolean {
+  return status === 408 || status === 429 || status >= 500;
+}
+
+/** The two OpenAI-compatible, key-based services. */
+export type ApiProvider = "gateway" | "openrouter";
+export const OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1";
+export const GATEWAY_BASE_URL = "https://ai-gateway.vercel.sh/v1";
+/** OpenRouter's app attribution headers (https://openrouter.ai/docs/app-attribution). */
+const OPENROUTER_HEADERS = Object.freeze({
+  "http-referer": "https://dawg.sh",
+  "x-title": "dawg",
+});
+
+export type ApiClientOptions = {
+  apiKey?: string;
+  baseUrl?: string;
+  modelIds?: Partial<Record<GatewayModel, string>>;
+  fetcher?: GatewayFetcher;
+  retry?: GatewayRetryOptions;
+  headerTimeoutMs?: number;
+};
+
+/** The streaming tool-calling client for OpenRouter's OpenAI-compatible API. */
+export function createOpenRouterClient(
+  options: ApiClientOptions = {},
+): GatewayClient {
+  return createApiClient("openrouter", {
+    ...options,
+    baseUrl:
+      options.baseUrl ?? process.env.OPENROUTER_BASE_URL ?? OPENROUTER_BASE_URL,
+    apiKey: options.apiKey ?? process.env.OPENROUTER_API_KEY,
+  });
+}
+
 export function createGatewayClient(
-  options: {
-    apiKey?: string;
-    baseUrl?: string;
-    modelIds?: Partial<Record<GatewayModel, string>>;
-    fetcher?: GatewayFetcher;
-  } = {},
+  options: ApiClientOptions = {},
+): GatewayClient {
+  return createApiClient("gateway", {
+    ...options,
+    baseUrl:
+      options.baseUrl ?? process.env.AI_GATEWAY_BASE_URL ?? GATEWAY_BASE_URL,
+    apiKey: options.apiKey ?? process.env.AI_GATEWAY_API_KEY,
+  });
+}
+
+function createApiClient(
+  provider: ApiProvider,
+  options: ApiClientOptions,
 ): GatewayClient {
   const fetcher = options.fetcher ?? fetch;
-  const baseUrl = (
-    options.baseUrl ??
-    process.env.AI_GATEWAY_BASE_URL ??
-    "https://ai-gateway.vercel.sh/v1"
-  ).replace(/\/$/, "");
-  const apiKey = options.apiKey ?? process.env.AI_GATEWAY_API_KEY;
+  const name = provider === "openrouter" ? "OpenRouter" : "AI Gateway";
+  const baseUrl = (options.baseUrl ?? GATEWAY_BASE_URL).replace(/\/$/, "");
+  const apiKey = options.apiKey;
+  const retries = options.retry?.retries ?? GATEWAY_RETRIES;
+  const retryBaseMs = options.retry?.baseMs ?? RETRY_BASE_MS;
+  const sleep = options.retry?.sleep ?? ((ms: number) => Bun.sleep(ms));
+  const random = options.retry?.random ?? Math.random;
+  const headerTimeoutMs = options.headerTimeoutMs ?? GATEWAY_HEADER_TIMEOUT_MS;
   const overrides: Partial<Record<GatewayModel, string | undefined>> = {
     "opus-5.5": process.env.DAWG_OPUS_MODEL || undefined,
     "sol-6.1": process.env.DAWG_SOL_MODEL || undefined,
@@ -138,6 +221,7 @@ export function createGatewayClient(
     apiKey && apiKey.length > 0 ? text.split(apiKey).join("[redacted]") : text;
 
   return {
+    provider,
     modelId: (model) => resolveModelId(model, overrides),
     async *stream(request, signal) {
       const model =
@@ -146,13 +230,18 @@ export function createGatewayClient(
           : resolveModelId(request.model, overrides);
       if (!apiKey)
         throw new GatewayError(
-          "no AI Gateway key; run `dawg login` or set AI_GATEWAY_API_KEY",
+          provider === "openrouter"
+            ? "no OpenRouter key; run `dawg login openrouter` or set OPENROUTER_API_KEY"
+            : "no AI Gateway key; run `dawg login` or set AI_GATEWAY_API_KEY",
         );
       const body: Record<string, unknown> = {
         model,
         messages: request.messages,
         stream: true,
+        // The final chunk then carries token usage (and OpenRouter's cost).
+        stream_options: { include_usage: true },
       };
+      if (provider === "openrouter") body.usage = { include: true };
       if (request.tools && request.tools.length > 0) {
         body.tools = request.tools;
         body.tool_choice = "auto";
@@ -165,26 +254,61 @@ export function createGatewayClient(
         request.maxTokens > 0
       )
         body.max_tokens = Math.min(request.maxTokens, 4096);
-      const init: RequestInit = {
-        method: "POST",
-        headers: {
-          authorization: `Bearer ${apiKey}`,
-          "content-type": "application/json",
-          accept: "text/event-stream",
-        },
-        body: JSON.stringify(body),
-      };
-      if (signal !== undefined) init.signal = signal;
-      const response = await fetcher(`${baseUrl}/chat/completions`, init);
-      if (!response.ok) {
-        const detail = redact(await boundedErrorDetail(response));
-        throw new GatewayError(
-          `AI Gateway request failed (${response.status})${detail ? `: ${detail}` : ""}`,
-          response.status,
-        );
+      const encoded = JSON.stringify(body);
+      // Retry only before any byte of a response has been consumed: a
+      // network failure, a header-phase timeout, or 408/429/5xx. Once the
+      // stream is flowing, a failure surfaces as-is (replaying could repeat
+      // tool calls the model already made).
+      let response: Response | undefined;
+      for (let attempt = 0; ; attempt += 1) {
+        if (signal?.aborted) throw abortError(signal);
+        const headerSignal = AbortSignal.any([
+          ...(signal === undefined ? [] : [signal]),
+          AbortSignal.timeout(headerTimeoutMs),
+        ]);
+        const init: RequestInit = {
+          method: "POST",
+          headers: {
+            authorization: `Bearer ${apiKey}`,
+            "content-type": "application/json",
+            accept: "text/event-stream",
+            ...(provider === "openrouter" ? OPENROUTER_HEADERS : {}),
+          },
+          body: encoded,
+          signal: headerSignal,
+        };
+        let reason: string;
+        try {
+          const candidate = await fetcher(`${baseUrl}/chat/completions`, init);
+          if (candidate.ok) {
+            response = candidate;
+            break;
+          }
+          const detail = redact(await boundedErrorDetail(candidate));
+          const error = new GatewayError(
+            `${name} request failed (${candidate.status})${detail ? `: ${detail}` : ""}`,
+            candidate.status,
+          );
+          if (!retryableStatus(candidate.status) || attempt >= retries)
+            throw error;
+          reason = String(candidate.status);
+        } catch (error) {
+          if (error instanceof GatewayError) throw error;
+          if (signal?.aborted) throw abortError(signal);
+          if (attempt >= retries)
+            throw new GatewayError(
+              headerSignal.aborted
+                ? `${name} did not respond within ${headerTimeoutMs} ms`
+                : `${name} request failed: ${redact(error instanceof Error ? error.message : String(error))}`,
+            );
+          reason = headerSignal.aborted ? "timeout" : "network";
+        }
+        yield { type: "activity", message: `retrying (${reason})…` };
+        const backoff = retryBaseMs * 2 ** attempt;
+        await sleep(Math.round(backoff * (0.5 + random())));
       }
       if (!response.body)
-        throw new GatewayError("AI Gateway returned an empty stream");
+        throw new GatewayError(`${name} returned an empty stream`);
       const streamOptions: { maxBytes: number; signal?: AbortSignal } = {
         maxBytes: request.maxResponseBytes,
       };
@@ -194,9 +318,9 @@ export function createGatewayClient(
         try {
           chunk = JSON.parse(data);
         } catch {
-          throw new GatewayError("AI Gateway sent a malformed stream chunk");
+          throw new GatewayError(`${name} sent a malformed stream chunk`);
         }
-        yield* normalizeChunk(chunk, redact);
+        yield* normalizeChunk(chunk, redact, name);
       }
     },
   };
@@ -205,6 +329,7 @@ export function createGatewayClient(
 function* normalizeChunk(
   chunk: unknown,
   redact: (text: string) => string,
+  name = "AI Gateway",
 ): Generator<ChatStreamEvent> {
   if (!isRecord(chunk)) return;
   if (isRecord(chunk.error)) {
@@ -212,8 +337,10 @@ function* normalizeChunk(
       typeof chunk.error.message === "string"
         ? chunk.error.message.slice(0, 300)
         : "unknown provider error";
-    throw new GatewayError(`AI Gateway stream error: ${redact(message)}`);
+    throw new GatewayError(`${name} stream error: ${redact(message)}`);
   }
+  const usage = parseUsage(chunk.usage);
+  if (usage) yield usage;
   const choice = Array.isArray(chunk.choices) ? chunk.choices[0] : undefined;
   if (!isRecord(choice)) return;
   const delta = isRecord(choice.delta) ? choice.delta : undefined;
@@ -242,6 +369,53 @@ function* normalizeChunk(
     yield { type: "finish", reason: choice.finish_reason };
 }
 
+/** Bounded, non-negative token counts; anything else is ignored. */
+function count(value: unknown): number | undefined {
+  return typeof value === "number" &&
+    Number.isFinite(value) &&
+    value >= 0 &&
+    value < 1e9
+    ? Math.floor(value)
+    : undefined;
+}
+
+export function parseUsage(
+  value: unknown,
+): Extract<ChatStreamEvent, { type: "usage" }> | undefined {
+  if (!isRecord(value)) return undefined;
+  const inputTokens = count(value.prompt_tokens);
+  const outputTokens = count(value.completion_tokens);
+  if (inputTokens === undefined || outputTokens === undefined) return undefined;
+  const event: Extract<ChatStreamEvent, { type: "usage" }> = {
+    type: "usage",
+    inputTokens,
+    outputTokens,
+  };
+  const details = isRecord(value.prompt_tokens_details)
+    ? value.prompt_tokens_details
+    : undefined;
+  const cached = count(details?.cached_tokens);
+  if (cached !== undefined) event.cachedInputTokens = cached;
+  const rawCost =
+    typeof value.cost === "string" ? Number(value.cost) : value.cost;
+  if (
+    typeof rawCost === "number" &&
+    Number.isFinite(rawCost) &&
+    rawCost >= 0 &&
+    rawCost < 1000
+  )
+    event.costUsd = rawCost;
+  return event;
+}
+
+function abortError(signal: AbortSignal): Error {
+  const reason: unknown = signal.reason;
+  if (reason instanceof Error) return reason;
+  const error = new Error("agent request was aborted");
+  error.name = "AbortError";
+  return error;
+}
+
 async function boundedErrorDetail(response: Response): Promise<string> {
   try {
     const text = (await response.text()).slice(0, MAX_ERROR_BODY_BYTES);