fix(zed): bound buffered translation and cancel delegated streams

scripts/test-layout/layout.json

@@ -2,6 +2,8 @@
   "version": 1,
   "root": "tests",
   "explicit": {
+    "compaction-event-ownership.test.ts": "responses",
+    "zed-retained-budget.test.ts": "providers",
     "claude-intercept-on-demand.test.ts": "claude-integration",
     "main-account-hard-lock-thresholds.test.ts": "codex-integration",
     "main-account-external-usage.test.ts": "codex-integration",

src/adapters/openai-responses/passthrough.ts

@@ -8,7 +8,7 @@ import { Buffer } from "node:buffer";
 import type { IncomingMeta, ProviderAdapter } from "../base";
 import { namespacedToolName, type AdapterEvent, type OcxParsedRequest, type OcxProviderConfig, type OcxUsage, type TierDecision } from "../../types";
 import { applyCodexRoutingHint, CODEX_RESPONSES_LITE_HEADER, CODEX_ROUTING_HINT_HEADER } from "../../codex/forward-transport-headers";
-import { COMPACT_PROMPT, compactionItemToText, decodeCompactionSummary, isCompactionItemType } from "../../responses/compaction";
+import { COMPACT_PROMPT, compactionItemToText, decodeCompactionSummary, isCompactionItemType, transferCompactionCiphertextLease } from "../../responses/compaction";
 import { decodeServerSentEvents } from "../../lib/sse-decoder";
 import {
   CODEX_FORWARD_BASE_URL,
@@ -590,112 +590,121 @@ export function createResponsesPassthroughAdapter(provider: OcxProviderConfig):
       let compactionEncryptedContent: string | undefined;
       let compactionEncryptedContentBytes = 0;
       let completedSeen = false;
-      for await (const event of decodeServerSentEvents(response.body, { translatorBudget: budget })) {
-        let payload: unknown;
-        try { payload = JSON.parse(event.data); } catch { continue; }
-        if (!isPlainObject(payload)) continue;
-        switch (payload.type) {
-          case "response.output_text.delta":
-            if (typeof payload.delta === "string") {
-              const next = deltas + payload.delta;
-              const nextBytes = appendedUtf8Bytes(deltasBytes, deltasLastCodeUnit, payload.delta);
-              const reservation = budget.reserveTransient(nextBytes, { kind: "retained_collectors" });
-              deltas = next;
-              reservation.commitRetained();
-              budget.releaseRetained(deltasBytes, { kind: "retained_collectors" });
-              deltasBytes = nextBytes;
-              if (payload.delta.length > 0) deltasLastCodeUnit = payload.delta.charCodeAt(payload.delta.length - 1);
-            }
-            break;
-          case "response.output_text.done":
-            if (typeof payload.text === "string") {
-              const next = doneText + payload.text;
-              const nextBytes = appendedUtf8Bytes(doneTextBytes, doneTextLastCodeUnit, payload.text);
-              const reservation = budget.reserveTransient(nextBytes, { kind: "retained_collectors" });
-              doneText = next;
-              reservation.commitRetained();
-              budget.releaseRetained(doneTextBytes, { kind: "retained_collectors" });
-              doneTextBytes = nextBytes;
-              if (payload.text.length > 0) doneTextLastCodeUnit = payload.text.charCodeAt(payload.text.length - 1);
-            }
-            break;
-          case "response.failed":
-          case "error":
-            yield { type: "error", message: responsesErrorMessage(payload.response ?? payload) };
-            return;
-          case "response.incomplete":
-            yield { type: "incomplete", reason: responsesErrorMessage(payload.response ?? payload) };
-            return;
-          case "response.completed":
-            {
-              completedSeen = true;
-              const responsePayload = isPlainObject(payload.response) ? payload.response : undefined;
-              const output = Array.isArray(responsePayload?.output) ? responsePayload.output : [];
-              const compaction = output.find(item => isPlainObject(item) && item.type === "compaction");
-              if (isPlainObject(compaction) && typeof compaction.encrypted_content === "string") {
-                const nextEncryptedContent = compaction.encrypted_content;
-                const nextEncryptedContentBytes = Buffer.byteLength(nextEncryptedContent, "utf8");
-                const reservation = budget.reserveTransient(nextEncryptedContentBytes, { kind: "retained_collectors" });
-                compactionEncryptedContent = nextEncryptedContent;
+      let ciphertextTransferred = false;
+      try {
+        for await (const event of decodeServerSentEvents(response.body, { translatorBudget: budget })) {
+          let payload: unknown;
+          try { payload = JSON.parse(event.data); } catch { continue; }
+          if (!isPlainObject(payload)) continue;
+          switch (payload.type) {
+            case "response.output_text.delta":
+              if (typeof payload.delta === "string") {
+                const next = deltas + payload.delta;
+                const nextBytes = appendedUtf8Bytes(deltasBytes, deltasLastCodeUnit, payload.delta);
+                const reservation = budget.reserveTransient(nextBytes, { kind: "retained_collectors" });
+                deltas = next;
                 reservation.commitRetained();
-                budget.releaseRetained(compactionEncryptedContentBytes, { kind: "retained_collectors" });
-                compactionEncryptedContentBytes = nextEncryptedContentBytes;
+                budget.releaseRetained(deltasBytes, { kind: "retained_collectors" });
+                deltasBytes = nextBytes;
+                if (payload.delta.length > 0) deltasLastCodeUnit = payload.delta.charCodeAt(payload.delta.length - 1);
               }
-              const next = responsesPayloadText(payload.response);
-              const nextBytes = Buffer.byteLength(next, "utf8");
-              const reservation = budget.reserveTransient(nextBytes, { kind: "retained_collectors" });
-              snapshot = next;
-              reservation.commitRetained();
-              budget.releaseRetained(snapshotBytes, { kind: "retained_collectors" });
-              snapshotBytes = nextBytes;
-            }
-            {
-              const nextUsage = usageFromResponsesPayload(payload.response);
-              // The attached raw usage object can be event-sized (unknown keys carry arbitrary
-              // values); it stays reachable until the terminal yields, so charge it like the
-              // adjacent retained collectors or it would defeat the per-request memory cap.
-              const nextRawBytes = nextUsage?.rawUsage === undefined ? 0
-                : Buffer.byteLength(JSON.stringify(nextUsage.rawUsage), "utf8");
-              if (nextRawBytes > 0) {
-                const reservation = budget.reserveTransient(nextRawBytes, { kind: "retained_collectors" });
-                usage = nextUsage;
+              break;
+            case "response.output_text.done":
+              if (typeof payload.text === "string") {
+                const next = doneText + payload.text;
+                const nextBytes = appendedUtf8Bytes(doneTextBytes, doneTextLastCodeUnit, payload.text);
+                const reservation = budget.reserveTransient(nextBytes, { kind: "retained_collectors" });
+                doneText = next;
                 reservation.commitRetained();
-              } else {
-                usage = nextUsage;
+                budget.releaseRetained(doneTextBytes, { kind: "retained_collectors" });
+                doneTextBytes = nextBytes;
+                if (payload.text.length > 0) doneTextLastCodeUnit = payload.text.charCodeAt(payload.text.length - 1);
               }
-              if (usageRawBytes > 0) {
-                budget.releaseRetained(usageRawBytes, { kind: "retained_collectors" });
+              break;
+            case "response.failed":
+            case "error":
+              yield { type: "error", message: responsesErrorMessage(payload.response ?? payload) };
+              return;
+            case "response.incomplete":
+              yield { type: "incomplete", reason: responsesErrorMessage(payload.response ?? payload) };
+              return;
+            case "response.completed":
+              {
+                completedSeen = true;
+                const responsePayload = isPlainObject(payload.response) ? payload.response : undefined;
+                const output = Array.isArray(responsePayload?.output) ? responsePayload.output : [];
+                const compaction = output.find(item => isPlainObject(item) && item.type === "compaction");
+                if (isPlainObject(compaction) && typeof compaction.encrypted_content === "string") {
+                  const nextEncryptedContent = compaction.encrypted_content;
+                  const nextEncryptedContentBytes = Buffer.byteLength(nextEncryptedContent, "utf8");
+                  const reservation = budget.reserveTransient(nextEncryptedContentBytes, { kind: "retained_collectors" });
+                  compactionEncryptedContent = nextEncryptedContent;
+                  reservation.commitRetained();
+                  budget.releaseRetained(compactionEncryptedContentBytes, { kind: "retained_collectors" });
+                  compactionEncryptedContentBytes = nextEncryptedContentBytes;
+                }
+                const next = responsesPayloadText(payload.response);
+                const nextBytes = Buffer.byteLength(next, "utf8");
+                const reservation = budget.reserveTransient(nextBytes, { kind: "retained_collectors" });
+                snapshot = next;
+                reservation.commitRetained();
+                budget.releaseRetained(snapshotBytes, { kind: "retained_collectors" });
+                snapshotBytes = nextBytes;
               }
-              usageRawBytes = nextRawBytes;
-            }
-            break;
-        }
-        // Buffered text is still upstream progress, but gateway keepalives are not.
-        // Yield after accounting, directly to the consumer: no progress queue or content leak.
-        if (
-          !completedSeen
-          && (payload.type === "response.output_text.delta"
-            || payload.type === "response.reasoning_summary_text.delta"
-            || payload.type === "response.reasoning_text.delta")
-          && typeof payload.delta === "string"
-          && payload.delta.length > 0
-        ) {
-          yield { type: "heartbeat" };
+              {
+                const nextUsage = usageFromResponsesPayload(payload.response);
+                // The attached raw usage object can be event-sized (unknown keys carry arbitrary
+                // values); it stays reachable until the terminal yields, so charge it like the
+                // adjacent retained collectors or it would defeat the per-request memory cap.
+                const nextRawBytes = nextUsage?.rawUsage === undefined ? 0
+                  : Buffer.byteLength(JSON.stringify(nextUsage.rawUsage), "utf8");
+                if (nextRawBytes > 0) {
+                  const reservation = budget.reserveTransient(nextRawBytes, { kind: "retained_collectors" });
+                  usage = nextUsage;
+                  reservation.commitRetained();
+                } else {
+                  usage = nextUsage;
+                }
+                if (usageRawBytes > 0) {
+                  budget.releaseRetained(usageRawBytes, { kind: "retained_collectors" });
+                }
+                usageRawBytes = nextRawBytes;
+              }
+              break;
+          }
+          // Buffered text is still upstream progress, but gateway keepalives are not.
+          // Yield after accounting, directly to the consumer: no progress queue or content leak.
+          if (
+            !completedSeen
+            && (payload.type === "response.output_text.delta"
+              || payload.type === "response.reasoning_summary_text.delta"
+              || payload.type === "response.reasoning_text.delta")
+            && typeof payload.delta === "string"
+            && payload.delta.length > 0
+          ) {
+            yield { type: "heartbeat" };
+          }
         }
+        // Gateways differ in which of these they emit; prefer the authoritative
+        // completed snapshot so text is never double-counted.
+        const text = snapshot || doneText || deltas;
+        if (text) yield { type: "text_delta", text };
+        const done: AdapterEvent = {
+          type: "done",
+          ...(usage ? { usage } : {}),
+          ...(compactionEncryptedContent ? { compactionEncryptedContent } : {}),
+        };
+        // Exact event ownership survives buffered collection without charging a second copy.
+        transferCompactionCiphertextLease(done, budget, compactionEncryptedContentBytes);
+        ciphertextTransferred = true;
+        yield done;
+      } finally {
+        budget.releaseRetained(
+          deltasBytes + doneTextBytes + snapshotBytes + usageRawBytes
+            + (ciphertextTransferred ? 0 : compactionEncryptedContentBytes),
+          { kind: "retained_collectors" },
+        );
       }
-      // Gateways differ in which of these they emit; prefer the authoritative
-      // completed snapshot so text is never double-counted.
-      const text = snapshot || doneText || deltas;
-      if (text) yield { type: "text_delta", text };
-      budget.releaseRetained(
-        deltasBytes + doneTextBytes + snapshotBytes + usageRawBytes,
-        { kind: "retained_collectors" },
-      );
-      yield {
-        type: "done",
-        ...(usage ? { usage } : {}),
-        ...(compactionEncryptedContent ? { compactionEncryptedContent } : {}),
-      };
     },
 
     async parseResponse(response: Response, budget: TranslatorBudget): Promise<AdapterEvent[]> {

src/adapters/zed.ts

@@ -5,7 +5,8 @@ import { createGoogleAdapter } from "./google";
 import { createOpenAIChatAdapter } from "./openai-chat";
 import { createResponsesPassthroughAdapter } from "./openai-responses";
 import type { AdapterEvent, OcxParsedRequest, OcxProviderConfig } from "../types";
-import { createTranslatorBudget, type TranslatorBudget } from "../lib/translator-budget";
+import { createTranslatorBudget, releaseTranslatedEvent, retainTranslatedEvent, type TranslatorBudget } from "../lib/translator-budget";
+import { releaseCompactionCiphertextLease } from "../responses/compaction";
 import { redactSecretString } from "../lib/redact";
 import {
   normalizeZedProvider,
@@ -365,12 +366,33 @@ export function createZedAdapter(provider: OcxProviderConfig): ProviderAdapter {
         status: response.status,
         headers: { "Content-Type": "text/event-stream" },
       });
-      yield* delegate.adapter.parseStream(translated, budget);
+      try {
+        yield* delegate.adapter.parseStream(translated, budget);
+      } finally {
+        // Some delegates only unlock their reader when a collector stops early.
+        try { await translated.body?.cancel(); } catch { /* already closed/errored */ }
+      }
     },
     async parseResponse(response: Response, budget: TranslatorBudget): Promise<AdapterEvent[]> {
       const events: AdapterEvent[] = [];
-      for await (const event of adapter.parseStream(response, budget)) events.push(event);
-      return events;
+      try {
+        for await (const event of adapter.parseStream(response, budget)) {
+          try {
+            retainTranslatedEvent(event, budget, events.at(-1));
+          } catch (error) {
+            releaseCompactionCiphertextLease(event, budget);
+            throw error;
+          }
+          events.push(event);
+        }
+        return events;
+      } catch (error) {
+        for (const event of events) {
+          releaseTranslatedEvent(event, budget);
+          releaseCompactionCiphertextLease(event, budget);
+        }
+        throw error;
+      }
     },
   };
   return adapter;

src/bridge/response-json.ts

@@ -17,7 +17,7 @@ import {
   type OcxErrorPayload,
 } from "../lib/errors";
 import { mayBecomePatchEnvelope, repairFreeformToolInput } from "../responses/apply-patch-envelope";
-import { encodeCompactionSummary } from "../responses/compaction";
+import { encodeCompactionSummary, releaseCompactionCiphertextLease } from "../responses/compaction";
 import { compileCodeModeHelperInput, resolveCodeModeHelperName } from "../responses/code-mode-helper-compat";
 import { isTruncatedStopReason, truncationReasonFor } from "../responses/truncated-stop-reason";
 import { encodeReasoningEnvelope, type ReasoningEnvelope } from "../responses/reasoning-envelope";
@@ -59,21 +59,21 @@ export function buildResponseJSON(
 ): Record<string, unknown> {
   // Default-budget safety net: a caller that omits the budget gets a bounded
   // default (disposed with the call), never the unbounded append path.
-  if (options?.translatorBudget) {
-    const body = buildResponseJSONWithBudget(events, modelId, options);
-    // A buffered turn delivers its whole answer as one body, so nothing calls the per-frame
-    // recorder on the SSE bridge. Without this the attempt would persist adapter events with
-    // zero relayed ones, which is the loss signal -- raised on every non-streaming request.
-    if (options.recordBufferedDelivery !== false) {
-      attemptDeliveryRecorder(options.translatorBudget)?.noteBufferedDelivery(body);
+  const ownsBudget = !options?.translatorBudget;
+  const budget = options?.translatorBudget ?? createTranslatorBudget();
+  try {
+    const body = buildResponseJSONWithBudget(events, modelId, { ...options, translatorBudget: budget });
+    // Buffered delivery has no per-frame recorder; retain its existing one-shot attribution.
+    if (options?.translatorBudget && options.recordBufferedDelivery !== false) {
+      attemptDeliveryRecorder(budget)?.noteBufferedDelivery(body);
     }
     return body;
-  }
-  const budget = createTranslatorBudget();
-  try {
-    return buildResponseJSONWithBudget(events, modelId, { ...options, translatorBudget: budget });
   } finally {
-    budget.dispose();
+    for (const event of events) {
+      releaseTranslatedEvent(event, budget);
+      releaseCompactionCiphertextLease(event, budget);
+    }
+    if (ownsBudget) budget.dispose();
   }
 }
 
@@ -618,7 +618,7 @@ function buildResponseJSONWithBudget(
       type: "compaction", id: `cmp_${uuid()}`,
       encrypted_content: compactionEncryptedContent ?? encodeCompactionSummary(joinChunks(batchCompaction)),
     };
-    pushOutput(item, compactionEncryptedContent ? bytesOf(compactionEncryptedContent) : batchCompaction.bytes);
+    pushOutput(item, compactionEncryptedContent ? 0 : batchCompaction.bytes);
   }
 
   const failure = errorEvent ? adapterFailureFromEvent(errorEvent) : undefined;