@bman654/clodex 2.8.5 → 2.9.0

package/dist/cli.js

@@ -382,7 +382,7 @@ import { join } from "path";
 // package.json
 var package_default = {
   name: "@bman654/clodex",
-  version: "2.8.5",
+  version: "2.9.0",
   publishConfig: {
     access: "public"
   },
@@ -3423,7 +3423,7 @@ function savedStopsAfter(current, assignments) {
 }
 
 // src/patch-transforms.ts
-var PATCH_TRANSFORMS_VERSION = 11;
+var PATCH_TRANSFORMS_VERSION = 12;
 var NATIVE_EFFORT_LEVELS = ["low", "medium", "high", "xhigh", "max"];
 var BASE_EFFORT_LEVELS = ["low", "medium", "high"];
 function projectNativeEffort(effort) {
@@ -3812,7 +3812,7 @@ function applyClodexPatches(source, config) {
     }
     applyOnce(
       patchName,
-      /(function [\w$]+\(\)\{)(let (?:[^;{}]|\{[^;{}]*\})*?[\w$]+\(process\.env\.CLAUDE_CODE_REMOTE\)\?(?:(?!\}\s*function )[\s\S])*?\)return process\.env;let ([\w$]+)=\{(?:(?!\}\s*function )[\s\S])*?return \3)(\})/,
+      /(function [\w$]+\(\)\{)(let[ {[](?:[^;{}]|\{[^;{}]*\})*?(?:\{[^;{}]*)?(?:getAgentProxyEnv|[\w$]+\(process\.env\.CLAUDE_CODE_REMOTE\)\?)(?:(?!\}\s*function )[\s\S])*?\)return process\.env;let ([\w$]+)=\{(?:(?!\}\s*function )[\s\S])*?return \3)(\})/,
       (match, head, body, _copyVar, tail) => {
         const at = js.indexOf(match);
         const closingBrace = at < 0 ? -1 : blockEndIndex(js, at + head.length - 1);
@@ -10928,43 +10928,152 @@ function resolveUpstreamTools(tools, messages) {
 
 // src/upstream-retry.ts
 var UPSTREAM_MAX_RETRIES_ENV = "CLODEX_UPSTREAM_MAX_RETRIES";
-var MAX_UPSTREAM_MAX_RETRIES = 5;
+var UPSTREAM_IDLE_TIMEOUT_ENV = "CLODEX_UPSTREAM_IDLE_TIMEOUT_MS";
+var UPSTREAM_TOTAL_TIMEOUT_ENV = "CLODEX_UPSTREAM_TOTAL_TIMEOUT_MS";
+var DEFAULT_UPSTREAM_IDLE_TIMEOUT_MS = 12e4;
+var DEFAULT_UPSTREAM_TOTAL_TIMEOUT_MS = 10 * 6e4;
+var MIN_UPSTREAM_IDLE_TIMEOUT_MS = 1e4;
+var MAX_UPSTREAM_IDLE_TIMEOUT_MS = 60 * 6e4;
+var MIN_UPSTREAM_TOTAL_TIMEOUT_MS = 6e4;
+var MAX_UPSTREAM_TOTAL_TIMEOUT_MS = 6 * 60 * 6e4;
+var SDK_INITIAL_RETRY_DELAY_MS = 2e3;
 var reportedValues = /* @__PURE__ */ new Set();
-function reportOnce(raw, message, warn) {
-  if (reportedValues.has(raw)) return;
-  reportedValues.add(raw);
+var defaultWarn = (message) => emitParentNotice(`clodex: ${message}`);
+function reportOnce(key, message, warn) {
+  if (reportedValues.has(key)) return;
+  reportedValues.add(key);
   try {
     warn(message);
   } catch {
   }
 }
-function upstreamMaxRetries(env = process.env, warn = (message) => emitParentNotice(`clodex: ${message}`)) {
+function timeoutSetting(env, envName, fallback, min, max, warn) {
+  const raw = env[envName]?.trim();
+  if (raw === void 0 || raw === "") return { value: fallback, explicit: false };
+  const value = Number(raw);
+  if (!Number.isInteger(value) || value <= 0) {
+    reportOnce(
+      `${envName}=${raw}`,
+      `ignoring ${envName}=${raw} (expected a positive integer number of milliseconds)`,
+      warn
+    );
+    return { value: fallback, explicit: false };
+  }
+  if (value < min || value > max) {
+    const clamped = Math.min(max, Math.max(min, value));
+    reportOnce(
+      `${envName}=${raw}`,
+      `clamping ${envName}=${raw} to ${clamped}ms (supported range is ${min}-${max}ms)`,
+      warn
+    );
+    return { value: clamped, explicit: true };
+  }
+  return { value, explicit: true };
+}
+function maxRetriesForIdleTimeout(idleTimeoutMs) {
+  let retries = 0;
+  let elapsedMs = 0;
+  let delayMs = SDK_INITIAL_RETRY_DELAY_MS;
+  while (elapsedMs + delayMs < idleTimeoutMs) {
+    elapsedMs += delayMs;
+    delayMs *= 2;
+    retries += 1;
+  }
+  return retries;
+}
+var MAX_UPSTREAM_MAX_RETRIES = maxRetriesForIdleTimeout(
+  DEFAULT_UPSTREAM_IDLE_TIMEOUT_MS
+);
+var DEFAULT_UPSTREAM_MAX_RETRIES = MAX_UPSTREAM_MAX_RETRIES;
+function configuredUpstreamMaxRetries(env, warn) {
   const raw = env[UPSTREAM_MAX_RETRIES_ENV]?.trim();
   if (raw === void 0 || raw === "") return void 0;
   const value = Number(raw);
   if (!Number.isInteger(value) || value < 0) {
     reportOnce(
-      raw,
+      `${UPSTREAM_MAX_RETRIES_ENV}=${raw}`,
       `ignoring ${UPSTREAM_MAX_RETRIES_ENV}=${raw} (expected a non-negative integer)`,
       warn
     );
     return void 0;
   }
-  if (value > MAX_UPSTREAM_MAX_RETRIES) {
+  return value;
+}
+function resolveUpstreamMaxRetries(env, warn, idleTimeoutMs) {
+  const ceiling = maxRetriesForIdleTimeout(idleTimeoutMs);
+  const defaultRetries = Math.min(DEFAULT_UPSTREAM_MAX_RETRIES, ceiling);
+  const value = configuredUpstreamMaxRetries(env, warn);
+  if (value === void 0) return defaultRetries;
+  if (value > ceiling) {
     reportOnce(
-      raw,
-      `clamping ${UPSTREAM_MAX_RETRIES_ENV}=${raw} to ${MAX_UPSTREAM_MAX_RETRIES} (higher values exceed the 120s streaming idle budget)`,
+      `${UPSTREAM_MAX_RETRIES_ENV}=${value}:idle=${idleTimeoutMs}`,
+      `clamping ${UPSTREAM_MAX_RETRIES_ENV}=${value} to ${ceiling} (estimated from the SDK fallback backoff and resolved ${idleTimeoutMs}ms idle timeout; provider delays may allow fewer retries)`,
       warn
     );
-    return MAX_UPSTREAM_MAX_RETRIES;
+    return ceiling;
   }
   return value;
 }
+function upstreamRequestBudget(options = {}) {
+  const env = options.env ?? process.env;
+  const warn = options.warn ?? defaultWarn;
+  const hasIdleOverride = options.idleTimeoutMs !== void 0;
+  const configuredIdle = options.idleTimeoutMs !== void 0 ? { value: options.idleTimeoutMs, explicit: false } : timeoutSetting(
+    env,
+    UPSTREAM_IDLE_TIMEOUT_ENV,
+    DEFAULT_UPSTREAM_IDLE_TIMEOUT_MS,
+    MIN_UPSTREAM_IDLE_TIMEOUT_MS,
+    MAX_UPSTREAM_IDLE_TIMEOUT_MS,
+    warn
+  );
+  const configuredTotal = timeoutSetting(
+    env,
+    UPSTREAM_TOTAL_TIMEOUT_ENV,
+    DEFAULT_UPSTREAM_TOTAL_TIMEOUT_MS,
+    MIN_UPSTREAM_TOTAL_TIMEOUT_MS,
+    MAX_UPSTREAM_TOTAL_TIMEOUT_MS,
+    warn
+  );
+  let idleTimeoutMs = configuredIdle.value;
+  let totalTimeoutMs = configuredTotal.value;
+  if (totalTimeoutMs < idleTimeoutMs) {
+    if (hasIdleOverride) {
+      idleTimeoutMs = totalTimeoutMs;
+    } else if (configuredIdle.explicit && !configuredTotal.explicit) {
+      reportOnce(
+        `timeout-pair:raise-total:${idleTimeoutMs}:${totalTimeoutMs}`,
+        `raising the resolved total timeout from ${totalTimeoutMs}ms to ${idleTimeoutMs}ms so it is not shorter than ${UPSTREAM_IDLE_TIMEOUT_ENV}`,
+        warn
+      );
+      totalTimeoutMs = idleTimeoutMs;
+    } else {
+      reportOnce(
+        `timeout-pair:lower-idle:${idleTimeoutMs}:${totalTimeoutMs}`,
+        `lowering the resolved idle timeout from ${idleTimeoutMs}ms to ${totalTimeoutMs}ms because it cannot exceed ${UPSTREAM_TOTAL_TIMEOUT_ENV}`,
+        warn
+      );
+      idleTimeoutMs = totalTimeoutMs;
+    }
+  }
+  return {
+    idleTimeoutMs,
+    totalTimeoutMs,
+    maxRetries: resolveUpstreamMaxRetries(env, warn, idleTimeoutMs)
+  };
+}
 var CLIENT_MAX_RETRIES_ENV = "CLAUDE_CODE_MAX_RETRIES";
 var DEFAULT_PASSTHROUGH_RETRIES = 1;
-function passthroughUpstreamRetries(env = process.env) {
-  const explicit = upstreamMaxRetries(env);
-  if (explicit !== void 0) return explicit;
+function passthroughUpstreamRetries(env = process.env, warn = defaultWarn) {
+  const explicit = configuredUpstreamMaxRetries(env, warn);
+  if (explicit !== void 0) {
+    if (explicit <= MAX_UPSTREAM_MAX_RETRIES) return explicit;
+    reportOnce(
+      `${UPSTREAM_MAX_RETRIES_ENV}=${explicit}:passthrough`,
+      `clamping ${UPSTREAM_MAX_RETRIES_ENV}=${explicit} to ${MAX_UPSTREAM_MAX_RETRIES} (the raw HTTP MITM path supports at most this many replays)`,
+      warn
+    );
+    return MAX_UPSTREAM_MAX_RETRIES;
+  }
   const raw = env[CLIENT_MAX_RETRIES_ENV]?.trim();
   if (raw !== void 0 && raw !== "") {
     const clientRetries = Number(raw);
@@ -10973,6 +11082,49 @@ function passthroughUpstreamRetries(env = process.env) {
   return DEFAULT_PASSTHROUGH_RETRIES;
 }
 
+// src/upstream-attempts.ts
+import { RetryError as RetryError2, wrapLanguageModel as wrapLanguageModel2 } from "ai";
+function trackUpstreamAttempts(model) {
+  if (typeof model === "string") {
+    return { model, deadlineError: (timeoutError) => timeoutError };
+  }
+  const failedAttempts = [];
+  let waitingToRetry = false;
+  const track = async (call) => {
+    waitingToRetry = false;
+    try {
+      const result = await call();
+      failedAttempts.length = 0;
+      return result;
+    } catch (error) {
+      failedAttempts.push(error);
+      waitingToRetry = true;
+      throw error;
+    }
+  };
+  const middleware = {
+    specificationVersion: "v4",
+    wrapGenerate: ({ doGenerate }) => track(doGenerate),
+    wrapStream: ({ doStream }) => track(doStream)
+  };
+  return {
+    model: wrapLanguageModel2({ model, middleware }),
+    deadlineError: (timeoutError) => {
+      if (!waitingToRetry || failedAttempts.length === 0) return timeoutError;
+      const count = failedAttempts.length;
+      return new RetryError2({
+        message: [
+          "Provider retry interrupted by a request deadline after",
+          count,
+          `failed ${count === 1 ? "attempt" : "attempts"}`
+        ].join(" "),
+        reason: "abort",
+        errors: [...failedAttempts]
+      });
+    }
+  };
+}
+
 // src/sdk-adapter.ts
 function sdkTranslationErrorSignature(error) {
   const message = error instanceof Error ? error.message : typeof error === "string" ? error : void 0;
@@ -11347,8 +11499,6 @@ function toAnthropicUsage(u) {
     cache_read_input_tokens: cacheRead
   };
 }
-var SDK_STREAM_IDLE_TIMEOUT_MS = 12e4;
-var SDK_TOTAL_TIMEOUT_MS = 10 * 6e4;
 function streamAbortError(signal) {
   if (signal?.reason instanceof Error) return signal.reason;
   const error = new Error(
@@ -11569,42 +11719,44 @@ async function writeAnthropicStream(stream, modelId, write, log12, observer, too
   emit("message_stop", { type: "message_stop" });
 }
 async function streamAnthropicResponse(model, params, modelId, write, log12, observer) {
-  const idleTimeoutMs = observer?.idleTimeoutMs ?? SDK_STREAM_IDLE_TIMEOUT_MS;
+  const { idleTimeoutMs, totalTimeoutMs, maxRetries } = upstreamRequestBudget({
+    idleTimeoutMs: observer?.idleTimeoutMs
+  });
+  const attempts = trackUpstreamAttempts(model);
   const idleAbort = new AbortController();
   const stopForwardingAbort = forwardAbortSignal(observer?.abortSignal, idleAbort);
   const abortSignal = idleAbort.signal;
-  let idleTimer = setTimeout(
-    () => idleAbort.abort(new Error(`no data received from provider for ${Math.round(idleTimeoutMs / 1e3)}s`)),
-    idleTimeoutMs
+  const idleError = () => attempts.deadlineError(
+    new Error(`no data received from provider for ${Math.round(idleTimeoutMs / 1e3)}s`)
   );
+  let idleTimer = setTimeout(() => idleAbort.abort(idleError()), idleTimeoutMs);
   const totalTimer = setTimeout(
-    () => idleAbort.abort(new Error(`provider stream exceeded ${Math.round(SDK_TOTAL_TIMEOUT_MS / 1e3)}s`)),
-    SDK_TOTAL_TIMEOUT_MS
+    () => idleAbort.abort(attempts.deadlineError(
+      new Error(`provider stream exceeded ${Math.round(totalTimeoutMs / 1e3)}s`)
+    )),
+    totalTimeoutMs
   );
-  const result = streamText({
-    model,
-    ...params,
-    maxRetries: upstreamMaxRetries(),
-    abortSignal,
-    onError: () => {
-    },
-    onStepFinish: (step) => reportUnsupportedServiceTier(params, step.warnings)
-  });
-  const watchedStream = (async function* () {
-    try {
-      for await (const part of result.stream) {
+  try {
+    const result = streamText({
+      model: attempts.model,
+      ...params,
+      maxRetries,
+      abortSignal,
+      onError: () => {
+      },
+      onStepFinish: (step) => reportUnsupportedServiceTier(params, step.warnings)
+    });
+    const watchedStream = (async function* () {
+      try {
+        for await (const part of result.stream) {
+          clearTimeout(idleTimer);
+          idleTimer = setTimeout(() => idleAbort.abort(idleError()), idleTimeoutMs);
+          yield part;
+        }
+      } finally {
         clearTimeout(idleTimer);
-        idleTimer = setTimeout(
-          () => idleAbort.abort(new Error(`no data received from provider for ${Math.round(idleTimeoutMs / 1e3)}s`)),
-          idleTimeoutMs
-        );
-        yield part;
       }
-    } finally {
-      clearTimeout(idleTimer);
-    }
-  })();
-  try {
+    })();
     await writeAnthropicStream(watchedStream, modelId, write, log12, { ...observer, abortSignal }, params.tools);
   } finally {
     stopForwardingAbort();
@@ -11619,39 +11771,41 @@ async function generateAnthropicResponse(model, params, modelId, options) {
   let finishReason;
   let usage;
   let warnings;
+  const { idleTimeoutMs, totalTimeoutMs, maxRetries } = upstreamRequestBudget({
+    idleTimeoutMs: options?.forceStream ? options.idleTimeoutMs : void 0
+  });
+  const attempts = trackUpstreamAttempts(model);
   if (options?.forceStream) {
     const forceAbort = new AbortController();
     const stopForwardingAbort = forwardAbortSignal(options.abortSignal, forceAbort);
     const abortSignal = forceAbort.signal;
-    const idleTimeoutMs = options.idleTimeoutMs ?? SDK_STREAM_IDLE_TIMEOUT_MS;
-    let idleTimer = setTimeout(
-      () => forceAbort.abort(new Error(`no data received from provider for ${Math.round(idleTimeoutMs / 1e3)}s`)),
-      idleTimeoutMs
+    const idleError = () => attempts.deadlineError(
+      new Error(`no data received from provider for ${Math.round(idleTimeoutMs / 1e3)}s`)
     );
+    let idleTimer = setTimeout(() => forceAbort.abort(idleError()), idleTimeoutMs);
     const totalTimer = setTimeout(
-      () => forceAbort.abort(new Error(`provider stream exceeded ${Math.round(SDK_TOTAL_TIMEOUT_MS / 1e3)}s`)),
-      SDK_TOTAL_TIMEOUT_MS
+      () => forceAbort.abort(attempts.deadlineError(
+        new Error(`provider stream exceeded ${Math.round(totalTimeoutMs / 1e3)}s`)
+      )),
+      totalTimeoutMs
     );
-    const r = streamText({
-      model,
-      ...params,
-      maxRetries: upstreamMaxRetries(),
-      abortSignal,
-      onError: () => {
-      },
-      onStepFinish: (step) => reportUnsupportedServiceTier(params, step.warnings)
-    });
     const streamedText = [];
     const streamedToolCalls = [];
     let streamedFinishReason = "stop";
     let streamedUsage;
     try {
+      const r = streamText({
+        model: attempts.model,
+        ...params,
+        maxRetries,
+        abortSignal,
+        onError: () => {
+        },
+        onStepFinish: (step) => reportUnsupportedServiceTier(params, step.warnings)
+      });
       for await (const part of r.stream) {
         clearTimeout(idleTimer);
-        idleTimer = setTimeout(
-          () => forceAbort.abort(new Error(`no data received from provider for ${Math.round(idleTimeoutMs / 1e3)}s`)),
-          idleTimeoutMs
-        );
+        idleTimer = setTimeout(() => forceAbort.abort(idleError()), idleTimeoutMs);
         options.onPart?.(part.type);
         if (abortSignal.aborted || part.type === "abort") {
           throw streamAbortError(abortSignal);
@@ -11686,17 +11840,22 @@ async function generateAnthropicResponse(model, params, modelId, options) {
     const generateAbort = new AbortController();
     const stopForwardingAbort = forwardAbortSignal(options?.abortSignal, generateAbort);
     const totalTimer = setTimeout(
-      () => generateAbort.abort(new Error(`provider request exceeded ${Math.round(SDK_TOTAL_TIMEOUT_MS / 1e3)}s`)),
-      SDK_TOTAL_TIMEOUT_MS
+      () => generateAbort.abort(attempts.deadlineError(
+        new Error(`provider request exceeded ${Math.round(totalTimeoutMs / 1e3)}s`)
+      )),
+      totalTimeoutMs
     );
     try {
       const r = await generateText({
-        model,
+        model: attempts.model,
         ...params,
-        maxRetries: upstreamMaxRetries(),
+        maxRetries,
         abortSignal: generateAbort.signal
       });
       ({ text: text5, toolCalls, finishReason, usage, warnings } = r);
+    } catch (error) {
+      if (generateAbort.signal.aborted) throw streamAbortError(generateAbort.signal);
+      throw error;
     } finally {
       stopForwardingAbort();
       clearTimeout(totalTimer);
@@ -16670,24 +16829,81 @@ async function collectOpenAiStream(stream) {
   }
   return collected;
 }
+function startUpstreamBudget(model, streaming) {
+  const { idleTimeoutMs, totalTimeoutMs, maxRetries } = upstreamRequestBudget();
+  const attempts = trackUpstreamAttempts(model);
+  const abort = new AbortController();
+  const idleError = () => attempts.deadlineError(new Error(
+    `no data received from provider for ${Math.round(idleTimeoutMs / 1e3)}s`
+  ));
+  let idleTimer = streaming ? setTimeout(() => abort.abort(idleError()), idleTimeoutMs) : void 0;
+  const totalTimer = setTimeout(
+    () => abort.abort(attempts.deadlineError(new Error(
+      `provider ${streaming ? "stream" : "request"} exceeded ${Math.round(totalTimeoutMs / 1e3)}s`
+    ))),
+    totalTimeoutMs
+  );
+  return {
+    model: attempts.model,
+    maxRetries,
+    abortSignal: abort.signal,
+    onStreamPart: () => {
+      if (idleTimer === void 0) return;
+      clearTimeout(idleTimer);
+      idleTimer = setTimeout(() => abort.abort(idleError()), idleTimeoutMs);
+    },
+    close: () => {
+      if (idleTimer !== void 0) clearTimeout(idleTimer);
+      clearTimeout(totalTimer);
+      if (!abort.signal.aborted) abort.abort();
+    }
+  };
+}
+async function* watchOpenAiStream(stream, budget) {
+  for await (const part of stream) {
+    if (budget.abortSignal.aborted || part.type === "abort") {
+      throw budget.abortSignal.reason instanceof Error ? budget.abortSignal.reason : new Error("SDK stream aborted");
+    }
+    budget.onStreamPart();
+    yield part;
+  }
+  if (budget.abortSignal.aborted) {
+    throw budget.abortSignal.reason instanceof Error ? budget.abortSignal.reason : new Error("SDK stream aborted");
+  }
+}
 async function generateOpenAiResponse(model, params, responseModelId, options) {
   let result;
-  if (options?.forceStream) {
-    const { stream } = streamText2({
-      model,
-      ...params,
-      maxRetries: upstreamMaxRetries(),
-      onError: () => {
-      },
-      onStepFinish: (step) => reportUnsupportedServiceTier(params, step.warnings)
-    });
-    result = await collectOpenAiStream(stream);
-  } else {
-    result = await generateText2({
-      model,
-      ...params,
-      maxRetries: upstreamMaxRetries()
-    });
+  const streaming = options?.forceStream === true;
+  const budget = startUpstreamBudget(model, streaming);
+  try {
+    if (streaming) {
+      const { stream } = streamText2({
+        model: budget.model,
+        ...params,
+        maxRetries: budget.maxRetries,
+        abortSignal: budget.abortSignal,
+        onError: () => {
+        },
+        onStepFinish: (step) => reportUnsupportedServiceTier(params, step.warnings)
+      });
+      result = await collectOpenAiStream(watchOpenAiStream(stream, budget));
+    } else {
+      try {
+        result = await generateText2({
+          model: budget.model,
+          ...params,
+          maxRetries: budget.maxRetries,
+          abortSignal: budget.abortSignal
+        });
+      } catch (error) {
+        if (budget.abortSignal.aborted && budget.abortSignal.reason instanceof Error) {
+          throw budget.abortSignal.reason;
+        }
+        throw error;
+      }
+    }
+  } finally {
+    budget.close();
   }
   reportUnsupportedServiceTier(params, result.warnings);
   const message = { role: "assistant", content: result.text || null };
@@ -16712,41 +16928,47 @@ async function generateOpenAiResponse(model, params, responseModelId, options) {
   };
 }
 async function streamOpenAiResponse(model, params, responseModelId, onChunk) {
-  const { stream } = streamText2({
-    model,
-    ...params,
-    maxRetries: upstreamMaxRetries(),
-    onStepFinish: (step) => reportUnsupportedServiceTier(params, step.warnings)
-  });
-  const baseData = {
-    id: `chatcmpl-${Date.now()}`,
-    object: "chat.completion.chunk",
-    created: Math.floor(Date.now() / 1e3),
-    model: responseModelId
-  };
-  const send = (delta, finish_reason = null) => onChunk(`data: ${JSON.stringify({ ...baseData, choices: [{ index: 0, delta, finish_reason }] })}
+  const budget = startUpstreamBudget(model, true);
+  try {
+    const { stream } = streamText2({
+      model: budget.model,
+      ...params,
+      maxRetries: budget.maxRetries,
+      abortSignal: budget.abortSignal,
+      onStepFinish: (step) => reportUnsupportedServiceTier(params, step.warnings)
+    });
+    const baseData = {
+      id: `chatcmpl-${Date.now()}`,
+      object: "chat.completion.chunk",
+      created: Math.floor(Date.now() / 1e3),
+      model: responseModelId
+    };
+    const send = (delta, finish_reason = null) => onChunk(`data: ${JSON.stringify({ ...baseData, choices: [{ index: 0, delta, finish_reason }] })}
 
 `);
-  for await (const part of stream) {
-    const p13 = part;
-    switch (p13.type) {
-      case "text-delta":
-        send({ role: "assistant", content: p13.textDelta ?? p13.text ?? "" });
-        break;
-      case "tool-input-start":
-        send({ role: "assistant", tool_calls: [{ index: 0, id: p13.id ?? p13.toolCallId, type: "function", function: { name: p13.toolName, arguments: "" } }] });
-        break;
-      case "tool-input-delta":
-        send({ tool_calls: [{ index: 0, function: { arguments: p13.delta ?? p13.text ?? p13.argsTextDelta ?? "" } }] });
-        break;
-      case "finish":
-        send({}, p13.finishReason || "stop");
-        break;
-      case "error":
-        throw p13.error instanceof Error || p13.error && typeof p13.error === "object" ? p13.error : new Error(typeof p13.error === "string" ? p13.error : "Upstream stream failed");
+    for await (const part of watchOpenAiStream(stream, budget)) {
+      const p13 = part;
+      switch (p13.type) {
+        case "text-delta":
+          send({ role: "assistant", content: p13.textDelta ?? p13.text ?? "" });
+          break;
+        case "tool-input-start":
+          send({ role: "assistant", tool_calls: [{ index: 0, id: p13.id ?? p13.toolCallId, type: "function", function: { name: p13.toolName, arguments: "" } }] });
+          break;
+        case "tool-input-delta":
+          send({ tool_calls: [{ index: 0, function: { arguments: p13.delta ?? p13.text ?? p13.argsTextDelta ?? "" } }] });
+          break;
+        case "finish":
+          send({}, p13.finishReason || "stop");
+          break;
+        case "error":
+          throw p13.error instanceof Error || p13.error && typeof p13.error === "object" ? p13.error : new Error(typeof p13.error === "string" ? p13.error : "Upstream stream failed");
+      }
     }
+    onChunk("data: [DONE]\n\n");
+  } finally {
+    budget.close();
   }
-  onChunk("data: [DONE]\n\n");
 }
 
 // src/server/router.ts