From 061c9c2f7f5899255b057f828a1f82a122152f24 Mon Sep 17 00:00:00 2001
From: Peter Steinberger <steipete@gmail.com>
Date: Wed, 12 Aug 2026 20:20:18 -0700
Subject: [PATCH] fix(ai): honor embedded transport policy (#122946)

---
 .../run/attempt-stream-settle.test.ts         | 71 ++++++++++++++++++-
 .../run/attempt-stream-settle.ts              |  1 +
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/agents/embedded-agent-runner/run/attempt-stream-settle.test.ts b/src/agents/embedded-agent-runner/run/attempt-stream-settle.test.ts
index f8d8c9e9741c..0ebe08151c20 100644
--- a/src/agents/embedded-agent-runner/run/attempt-stream-settle.test.ts
+++ b/src/agents/embedded-agent-runner/run/attempt-stream-settle.test.ts
@@ -1,10 +1,15 @@
 // Settlement liveness: a wedged block-reply flush must not park the turn.
 import { afterEach, describe, expect, it, vi } from "vitest";
+import { bindStreamLlmRuntime } from "../../../llm/model-runtime-binding.js";
 import { SessionManager } from "../../sessions/index.js";
 import { RUN_LIVENESS_JOIN_TIMEOUT_MS } from "./abortable.js";
-import { settleEmbeddedAttemptStream } from "./attempt-stream-settle.js";
+import {
+  prepareEmbeddedAttemptTransport,
+  settleEmbeddedAttemptStream,
+} from "./attempt-stream-settle.js";
 
 type SettleInput = Parameters<typeof settleEmbeddedAttemptStream>[0];
+type PrepareTransportInput = Parameters<typeof prepareEmbeddedAttemptTransport>[0];
 
 function createSettleFixture(overrides?: Partial<SettleInput>): SettleInput {
   const sessionManager = SessionManager.inMemory();
@@ -99,3 +104,67 @@ describe("settleEmbeddedAttemptStream liveness", () => {
     expect(result.sessionIdUsed).toBe("sess-settle-1");
   });
 });
+
+describe("prepareEmbeddedAttemptTransport", () => {
+  it("applies the prepared transport to the live agent owner", async () => {
+    const streamFn = vi.fn();
+    bindStreamLlmRuntime(streamFn, {
+      streamSimple: streamFn,
+      registry: { getApiProvider: () => undefined },
+    } as never);
+    const session = {
+      agent: {
+        streamFn,
+        transport: "auto",
+      },
+    };
+    const input = {
+      attempt: {
+        config: {},
+        model: {
+          api: "test-api",
+          provider: "test-provider",
+          id: "test-model",
+        },
+        modelId: "test-model",
+        provider: "test-provider",
+        promptCacheKey: undefined,
+        resolvedApiKey: undefined,
+        runId: "run-transport-1",
+        runtimePlan: {
+          auth: { forwardedAuthProfileId: undefined },
+          transport: {
+            resolveExtraParams: () => ({ transport: "sse" }),
+          },
+        },
+        sessionId: "sess-transport-1",
+      },
+      session,
+      settingsManager: {
+        getGlobalSettings: () => ({}),
+        getProjectSettings: () => ({}),
+      },
+      providerThinkingLevel: undefined,
+      sessionAgentId: "main",
+      workspaceDir: "/workspace",
+      workspaceOnly: false,
+      agentDir: "/agent",
+      abortSignal: new AbortController().signal,
+      getProviderRuntimeHandle: () => ({
+        provider: "test-provider",
+        modelId: "test-model",
+      }),
+      sandboxSessionKey: "agent:main:test",
+      codeModeControlsEnabled: false,
+      providerPromptState: {
+        state: {},
+        effectiveContextTokenBudget: 128_000,
+      },
+    } as unknown as PrepareTransportInput;
+
+    const result = await prepareEmbeddedAttemptTransport(input);
+
+    expect(result.effectiveAgentTransport).toBe("sse");
+    expect(session.agent.transport).toBe("sse");
+  });
+});
diff --git a/src/agents/embedded-agent-runner/run/attempt-stream-settle.ts b/src/agents/embedded-agent-runner/run/attempt-stream-settle.ts
index cc07d4af66a7..bcbca7aa248c 100644
--- a/src/agents/embedded-agent-runner/run/attempt-stream-settle.ts
+++ b/src/agents/embedded-agent-runner/run/attempt-stream-settle.ts
@@ -629,6 +629,7 @@ export async function prepareEmbeddedAttemptTransport(input: {
         `(${attempt.provider}/${attempt.modelId})`,
     );
   }
+  session.agent.transport = effectiveAgentTransport;
   return {
     effectiveAgentTransport,
     effectiveExtraParams,