fix(ai): honor embedded transport policy (#122946)

This commit is contained in:
Peter Steinberger
2026-08-12 20:20:18 -07:00
committed by GitHub
parent 7dbe21b9cf
commit 061c9c2f7f
2 changed files with 71 additions and 1 deletions
@@ -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");
});
});
@@ -629,6 +629,7 @@ export async function prepareEmbeddedAttemptTransport(input: {
`(${attempt.provider}/${attempt.modelId})`,
);
}
session.agent.transport = effectiveAgentTransport;
return {
effectiveAgentTransport,
effectiveExtraParams,