@tonyclaw/agent-inspector 3.0.44 → 3.0.46

package/src/proxy/handler.ts

@@ -1,7 +1,7 @@
-import { createLog, finalizeLogUpdate, getNextLogId, type CapturedLog } from "./store";
+import { Buffer } from "node:buffer";
+import { createLog, getNextLogId, persistFinalizedLogUpdate, type CapturedLog } from "./store";
 import { appendLogEntry, logger } from "./logger";
 import { extractRequestMetadata } from "./schemas";
-import { registry } from "./formats";
 import { findProviderByModel } from "./providers";
 import { getClientInfo } from "./socketTracker";
 import { formatForPath } from "./formats";
@@ -25,18 +25,26 @@ import { getConfig } from "./config";
 import { safeFetch } from "../lib/safeFetch";
 
 const ALLOW_LOOPBACK = process.env["AGENT_INSPECTOR_ALLOW_LOOPBACK"] === "1";
-import { stripClaudeCodeBillingHeader } from "./claudeCodeStrip";
+const EMPTY_UPSTREAM_BODY_ERROR = "Upstream response did not include a body";
+const UPSTREAM_STREAM_ERROR = "Upstream stream terminated unexpectedly";
+import { stripClaudeCodeBillingHeader } from "../lib/claudeCodeStrip";
 import { stripOpenAIOrphanToolMessages } from "./openaiOrphanToolStrip";
 import { maskApiKey } from "../lib/mask";
 import { buildFileLogEntry, type FinalizeLogJob } from "./logFinalizer";
 import { enqueueFinalizeLogJob } from "./sessionRuntime";
 import { RawStreamCapture } from "./rawStreamCapture";
-import { acquireRuntimeRequest, type RuntimeRequestLease } from "./runtimeAdmission";
+import {
+  acquireRuntimeRequest,
+  getRuntimeAdmissionState,
+  type RuntimeRequestLease,
+} from "./runtimeAdmission";
+import { resourceLimitResponse } from "../lib/resourceLimits";
 import {
   buildUpstreamUrl,
   describeApiRoute,
   getProxyApiPath,
   injectProviderAuth,
+  providerSupportsApiRoute,
   selectUpstreamBase,
   setUpstreamHost,
 } from "./upstream";
@@ -181,31 +189,81 @@ function handleNonStreamingResponse(
   });
 }
 
-function handleStreamingResponse(
+function applyRawStreamCaptureStatus(log: CapturedLog, capture: RawStreamCapture): void {
+  const status = capture.status;
+  log.captureIncomplete = status.incomplete;
+  log.droppedChunks = status.droppedChunks;
+  log.droppedBytes = status.droppedBytes;
+  log.captureIncompleteReason = status.reason;
+}
+
+function applyUpstreamStreamError(log: CapturedLog): void {
+  log.captureIncomplete = true;
+  log.captureIncompleteReason = "upstream-stream-error";
+  log.error = UPSTREAM_STREAM_ERROR;
+}
+
+async function handleStreamingResponse(
   upstreamRes: Response,
   req: Request,
   startTime: number,
   upstreamUrl: string,
   log: CapturedLog,
   collectStreamingChunks: boolean,
-): Response {
+): Promise<Response> {
   log.streaming = true;
   log.responseStatus = upstreamRes.status;
 
+  if (upstreamRes.body === null) {
+    log.elapsedMs = Date.now() - startTime;
+    log.responseStatus = STATUS_BAD_GATEWAY;
+    log.responseText = EMPTY_UPSTREAM_BODY_ERROR;
+    log.captureIncomplete = true;
+    log.droppedChunks = 0;
+    log.droppedBytes = 0;
+    log.captureIncompleteReason = "upstream-stream-error";
+    log.error = EMPTY_UPSTREAM_BODY_ERROR;
+    appendLogEntry({ ...buildFileLogEntry(log, upstreamUrl), error: EMPTY_UPSTREAM_BODY_ERROR });
+    await persistFinalizedLogUpdate(log);
+    return new Response("No response body", { status: STATUS_BAD_GATEWAY });
+  }
+
   const capture = new RawStreamCapture();
   const decoder = new TextDecoder();
   let finalizationScheduled = false;
   let firstChunkMs: number | null = null;
+  let pendingDecoderBytes = 0;
+  let downstreamCancelled = false;
+
+  const submitDecodedText = (text: string, flush: boolean = false): void => {
+    const observedBytes = flush
+      ? pendingDecoderBytes
+      : Math.min(pendingDecoderBytes, Buffer.byteLength(text, "utf-8"));
+    pendingDecoderBytes -= observedBytes;
+    if (text === "" && observedBytes === 0) return;
+    capture.tryAppend(text, observedBytes);
+    applyRawStreamCaptureStatus(log, capture);
+  };
 
-  const scheduleStreamingFinalization = (type: "streaming" | "stream-abort"): void => {
+  const scheduleStreamingFinalization = (
+    type: "streaming" | "stream-abort",
+    upstreamStreamFailed: boolean = false,
+  ): void => {
     if (finalizationScheduled) return;
     finalizationScheduled = true;
 
+    if (pendingDecoderBytes > 0) {
+      capture.recordDroppedObservation("pipeline-closed", pendingDecoderBytes);
+      pendingDecoderBytes = 0;
+    }
+
     const totalStreamMs = Date.now() - startTime;
     const elapsedMs = totalStreamMs;
     void capture
       .finalize()
       .then((rawStream) => {
+        applyRawStreamCaptureStatus(log, capture);
+        if (upstreamStreamFailed) applyUpstreamStreamError(log);
         switch (type) {
           case "streaming":
             scheduleLogFinalization({
@@ -218,6 +276,7 @@ function handleStreamingResponse(
               collectStreamingChunks,
               firstChunkMs,
               totalStreamMs,
+              terminalError: upstreamStreamFailed ? UPSTREAM_STREAM_ERROR : null,
             });
             return;
           case "stream-abort":
@@ -237,6 +296,10 @@ function handleStreamingResponse(
       })
       .catch((error: unknown) => {
         const message = error instanceof Error ? error.message : String(error);
+        applyRawStreamCaptureStatus(log, capture);
+        if (upstreamStreamFailed) applyUpstreamStreamError(log);
+        log.captureIncomplete = true;
+        log.captureIncompleteReason ??= "finalize-failed";
         logger.error(
           `[handler] Raw stream finalization failed for log #${String(log.id)}:`,
           message,
@@ -256,28 +319,48 @@ function handleStreamingResponse(
   };
 
   const transform = new TransformStream<Uint8Array, Uint8Array>({
-    async transform(chunk, controller) {
+    transform(chunk, controller) {
       if (firstChunkMs === null) firstChunkMs = Date.now() - startTime;
       controller.enqueue(chunk);
-      await capture.append(decoder.decode(chunk, { stream: true }));
+      pendingDecoderBytes += chunk.byteLength;
+      submitDecodedText(decoder.decode(chunk, { stream: true }));
     },
-    async flush() {
-      await capture.append(decoder.decode());
+    flush() {
+      submitDecodedText(decoder.decode(), true);
       scheduleStreamingFinalization("streaming");
     },
   });
 
-  if (upstreamRes.body === null) {
-    return new Response("No response body", { status: STATUS_BAD_GATEWAY });
-  }
-
   // pipeTo returns a promise that rejects when either side errors. Without
   // this catch, an upstream connection drop skips flush() and leaves the log
   // permanently un-finalized with a leaked temp file.
   void upstreamRes.body.pipeTo(transform.writable).catch(() => {
-    scheduleStreamingFinalization("streaming");
+    if (downstreamCancelled || req.signal?.aborted === true) {
+      scheduleStreamingFinalization("stream-abort");
+      return;
+    }
+    scheduleStreamingFinalization("streaming", true);
+  });
+  const transformedReader = transform.readable.getReader();
+  const loggedStream = new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      try {
+        const next = await transformedReader.read();
+        if (next.done) {
+          controller.close();
+          return;
+        }
+        controller.enqueue(next.value);
+      } catch (error) {
+        controller.error(error);
+      }
+    },
+    async cancel(reason) {
+      downstreamCancelled = true;
+      scheduleStreamingFinalization("stream-abort");
+      await transformedReader.cancel(reason);
+    },
   });
-  const loggedStream = transform.readable;
 
   if (req.signal === undefined) {
     logger.warn("[handler] req.signal is undefined; abort detection disabled for this request");
@@ -361,6 +444,19 @@ async function handleAdmittedProxy(req: Request): Promise<Response> {
   // Find provider config using already-extracted model (not calling extractModelFromBody again)
   const matchedProviderConfig = model !== null ? findProviderByModel(model) : null;
   const route = describeApiRoute(parsed.apiPath);
+
+  // Model matching is necessary but not sufficient: the Provider must explicitly expose the
+  // protocol selected by the request path. This guard stays before log creation and fetch so a
+  // protocol mismatch has the same zero-side-effect contract as an unconfigured model.
+  if (
+    model === null ||
+    matchedProviderConfig === null ||
+    !providerSupportsApiRoute(route, matchedProviderConfig)
+  ) {
+    logger.warn(`[handler] Unsupported provider: model=${model}`);
+    return new Response("Forbidden: unsupported provider", { status: STATUS_FORBIDDEN });
+  }
+
   const upstreamBase = selectUpstreamBase(route, matchedProviderConfig);
   const upstreamUrl = buildUpstreamUrl(upstreamBase, parsed.normalizedPath);
 
@@ -374,15 +470,7 @@ async function handleAdmittedProxy(req: Request): Promise<Response> {
   setUpstreamHost(upstreamHeaders, upstreamBase);
   injectProviderAuth(upstreamHeaders, matchedProviderConfig);
 
-  const provider = model !== null ? registry.findProvider(model) : null;
-
-  // Only proxy requests matching a registered provider
-  if (model === null || provider === null) {
-    logger.warn(`[handler] Unsupported provider: model=${model}`);
-    return new Response("Forbidden: unsupported provider", { status: STATUS_FORBIDDEN });
-  }
-
-  // Pick the format handler from the request path. provider.format is only a
+  // Pick the format handler from the request path. ProviderConfig.format is only a
   // hint used as a fallback when no anthropicBaseUrl/openaiBaseUrl is set, and
   // does NOT describe what the upstream actually returns: a provider configured
   // with format="openai" may still expose an Anthropic-compatible endpoint
@@ -441,7 +529,7 @@ async function handleAdmittedProxy(req: Request): Promise<Response> {
       log.responseStatus = status;
       log.responseText = fetchResult.reason;
       appendLogEntry({ ...buildFileLogEntry(log, upstreamUrl), error: fetchResult.reason });
-      finalizeLogUpdate(log);
+      await persistFinalizedLogUpdate(log);
       return new Response(
         fetchResult.kind === "network-policy"
           ? "Forbidden: upstream URL is not permitted"
@@ -458,14 +546,14 @@ async function handleAdmittedProxy(req: Request): Promise<Response> {
       log.responseStatus = 499; // Client Closed Request (non-standard but descriptive)
       log.responseText = "Client aborted";
       appendLogEntry({ ...buildFileLogEntry(log, upstreamUrl), error: "Client aborted" });
-      finalizeLogUpdate(log);
+      await persistFinalizedLogUpdate(log);
       return new Response("Client aborted", { status: 499 });
     }
     logger.error(`[handler] Proxy error: ${req.method} ${parsed.apiPath}`, String(err));
     log.responseStatus = STATUS_BAD_GATEWAY;
     log.responseText = String(err);
     appendLogEntry({ ...buildFileLogEntry(log, upstreamUrl), error: String(err) });
-    finalizeLogUpdate(log);
+    await persistFinalizedLogUpdate(log);
     return new Response(`Proxy error: ${err}`, { status: STATUS_BAD_GATEWAY });
   }
 
@@ -527,6 +615,18 @@ function responseWithRuntimeLease(response: Response, lease: RuntimeRequestLease
 export async function handleProxy(req: Request): Promise<Response> {
   const lease = acquireRuntimeRequest();
   if (lease === null) {
+    const admission = getRuntimeAdmissionState();
+    if (admission.acceptingRequests) {
+      return resourceLimitResponse({
+        error: `Proxy request capacity reached. At most ${String(admission.capacity)} active requests are allowed; retry after a response completes.`,
+        code: "proxy-capacity-exceeded",
+        status: 503,
+        resource: "proxy-requests",
+        limit: admission.capacity,
+        retryable: true,
+        retryAfterSeconds: 1,
+      });
+    }
     return Response.json(
       { error: "Agent Inspector is shutting down" },
       { status: 503, headers: { "cache-control": "no-store", "retry-after": "5" } },