fix(llm): reject unterminated provider streams (#36881)

This commit is contained in:
Kit Langton 2026-07-14 12:44:45 -04:00 committed by GitHub
parent cd9be63484
commit 5c5579e90c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 45 additions and 4 deletions

View file

@ -10,7 +10,7 @@ import { WebSocketExecutor } from "./transport"
import type { Protocol } from "./protocol"
import { applyCachePolicy } from "../cache-policy"
import * as ProviderShared from "../protocols/shared"
import type { LLMError, LLMEvent, PreparedRequestOf, ProtocolID, ProviderOptions } from "../schema"
import type { LLMError, PreparedRequestOf, ProtocolID, ProviderOptions } from "../schema"
import {
GenerationOptions,
HttpOptions,
@ -19,6 +19,7 @@ import {
Model,
ModelLimits,
LLMError as LLMErrorClass,
LLMEvent,
PreparedRequest,
ProviderID,
mergeGenerationOptions,
@ -229,6 +230,28 @@ const streamError = (route: string, message: string, cause: Cause.Cause<unknown>
return ProviderShared.eventError(route, message, Cause.pretty(cause))
}
const requireTerminalEvent = (route: string) => (events: Stream.Stream<LLMEvent, LLMError>) =>
Stream.suspend(() => {
let terminal = false
return events.pipe(
Stream.mapEffect((event) => {
if (terminal)
return Effect.fail(
ProviderShared.eventError(route, `Provider emitted ${event.type} after the terminal event`),
)
if (LLMEvent.is.finish(event) || LLMEvent.is.providerError(event)) terminal = true
return Effect.succeed(event)
}),
Stream.onEnd(
Effect.suspend(() =>
terminal
? Effect.void
: Effect.fail(ProviderShared.eventError(route, "Provider stream ended without a terminal finish event")),
),
),
)
})
function makeFromTransport<Body, Prepared, Frame, Event, State>(
input: MakeTransportInput<Body, Prepared, Frame, Event, State>,
): Route<Body, Prepared> {
@ -298,6 +321,7 @@ function makeFromTransport<Body, Prepared, Frame, Event, State>(
protocol.stream.onHalt ? { onHalt: protocol.stream.onHalt } : undefined,
),
Stream.catchCause((cause) => Stream.fail(streamError(route, `Failed to read ${route} stream`, cause))),
requireTerminalEvent(route),
)
},
} satisfies Route<Body, Prepared>

View file

@ -105,6 +105,9 @@ const echoLayer = dynamicResponse(({ text, respond }) =>
)
const it = testEffect(echoLayer)
const unterminated = testEffect(
dynamicResponse(({ respond }) => Effect.succeed(respond(encodeJson([{ type: "text", text: "partial" }])))),
)
describe("llm route", () => {
it.effect("stream and generate use the route pipeline", () =>
@ -125,6 +128,15 @@ describe("llm route", () => {
}),
)
unterminated.effect("fails when the normalized stream ends without a terminal event", () =>
Effect.gen(function* () {
const error = yield* (yield* LLMClient.Service).stream(request).pipe(Stream.runDrain, Effect.flip)
expect(error.reason).toMatchObject({ _tag: "InvalidProviderOutput" })
expect(error.message).toContain("Provider stream ended without a terminal finish event")
}),
)
it.effect("selects routes by model route value", () =>
Effect.gen(function* () {
const llm = yield* LLMClient.Service

View file

@ -602,7 +602,7 @@ describe("OpenAI Chat route", () => {
}),
)
it.effect("does not finalize streamed tool calls without a finish reason", () =>
it.effect("fails a streamed tool call when the provider ends without a finish reason", () =>
Effect.gen(function* () {
const body = sseEvents(
deltaChunk({
@ -614,8 +614,11 @@ describe("OpenAI Chat route", () => {
const input = LLM.updateRequest(request, {
tools: [{ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } }],
})
const events = Array.from(
yield* LLMClient.stream(input).pipe(Stream.runCollect, Effect.provide(fixedResponse(body))),
const events: LLMEvent[] = []
const streamError = yield* LLMClient.stream(input).pipe(
Stream.runForEach((event) => Effect.sync(() => events.push(event))),
Effect.flip,
Effect.provide(fixedResponse(body)),
)
const error = yield* LLMClient.generate(input).pipe(Effect.provide(fixedResponse(body)), Effect.flip)
@ -626,6 +629,8 @@ describe("OpenAI Chat route", () => {
{ type: "tool-input-delta", id: "call_1", name: "lookup", text: ':"weather"}' },
])
expect(events.filter(LLMEvent.is.toolCall)).toEqual([])
expect(streamError.reason).toMatchObject({ _tag: "InvalidProviderOutput" })
expect(streamError.message).toContain("Provider stream ended without a terminal finish event")
expect(error.message).toContain("Provider stream ended without a terminal finish event")
}),
)