Skip to content

Commit 6f89a35

Browse files
Tom-RyderTom-Ryder
committed
fix(streaming): reject incomplete sse streams
1 parent 04ac149 commit 6f89a35

17 files changed

Lines changed: 444 additions & 38 deletions

‎Sources/AgentRunKit/Core/Streaming/StreamProcessor.swift‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@ private struct StreamAccumulation {
9393
var usage: TokenUsage?
9494
var continuity: AssistantContinuity?
9595
var yieldedEvent = false
96+
var sawFinished = false
9697

9798
mutating func apply(
9899
_ input: RunStreamElement,
@@ -141,6 +142,7 @@ private struct StreamAccumulation {
141142
yieldedEvent = true
142143
continuation.yield(eventFactory.make(.audioTranscript(text)))
143144
case let .finished(iterationUsage):
145+
sawFinished = true
144146
guard let iterationUsage else { return }
145147
totalUsage += iterationUsage
146148
usage = iterationUsage
@@ -259,6 +261,10 @@ struct StreamProcessor {
259261
emittedOutput = state.yieldedEvent
260262
throw AgentError.malformedStream(.orphanedToolCallArguments(indices: state.pendingArguments.keys.sorted()))
261263
}
264+
guard state.sawFinished else {
265+
emittedOutput = state.yieldedEvent
266+
throw AgentError.llmError(.streamStalled)
267+
}
262268

263269
emittedOutput = true
264270
state.finishAudio(continuation: continuation)

‎Sources/AgentRunKit/LLM/Providers/Anthropic/AnthropicClientStreaming.swift‎

Lines changed: 12 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,17 @@ extension AnthropicClient {
2222

2323
let state = AnthropicStreamState()
2424

25-
try await processSSEStream(
25+
let completed = try await processSSEStream(
2626
bytes: bytes,
2727
stallTimeout: retryPolicy.streamStallTimeout
2828
) { event in
2929
try await self.handleSSEEvent(
3030
event, state: state, continuation: continuation
3131
)
3232
}
33+
guard completed else {
34+
throw AgentError.llmError(.streamStalled)
35+
}
3336
continuation.finish()
3437
}
3538

@@ -54,14 +57,17 @@ extension AnthropicClient {
5457

5558
let state = AnthropicStreamState()
5659

57-
try await processSSEStream(
60+
let completed = try await processSSEStream(
5861
bytes: bytes,
5962
stallTimeout: retryPolicy.streamStallTimeout
6063
) { event in
6164
try await self.handleSSEEvent(event, state: state) { delta in
6265
continuation.yield(.delta(delta))
6366
}
6467
}
68+
guard completed else {
69+
throw AgentError.llmError(.streamStalled)
70+
}
6571

6672
if await state.isCompleted {
6773
let blocks = try await state.finalizedBlocks()
@@ -129,9 +135,10 @@ extension AnthropicClient {
129135
data: data, state: state, yield: yield
130136
)
131137
case .messageDelta:
132-
try await handleMessageDelta(data: data, state: state, yield: yield)
138+
try await handleMessageDelta(data: data, state: state)
133139
case .messageStop:
134140
await state.markCompleted()
141+
await yield(.finished(usage: state.finalUsage()))
135142
return true
136143
case .error:
137144
try handleError(data: data)
@@ -244,20 +251,10 @@ extension AnthropicClient {
244251

245252
private func handleMessageDelta(
246253
data: Data,
247-
state: AnthropicStreamState,
248-
yield: @Sendable (StreamDelta) -> Void
254+
state: AnthropicStreamState
249255
) async throws {
250256
let event = try decodeEvent(AnthropicMessageDeltaEvent.self, from: data)
251-
let outputTokens = event.usage?.outputTokens ?? 0
252-
let inputUsage = await state.inputUsage
253-
254-
let usage = TokenUsage(
255-
input: inputUsage?.inputTokens ?? 0,
256-
output: outputTokens,
257-
cacheRead: inputUsage?.cacheReadInputTokens,
258-
cacheWrite: inputUsage?.cacheCreationInputTokens
259-
)
260-
yield(.finished(usage: usage))
257+
await state.setOutputTokens(event.usage?.outputTokens ?? 0)
261258
}
262259

263260
private func handleError(data: Data) throws {

‎Sources/AgentRunKit/LLM/Providers/Anthropic/AnthropicStreamState.swift‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ actor AnthropicStreamState {
1616
private var hasOpaqueDeltaBlocks = false
1717
private(set) var toolCallCount: Int = 0
1818
private(set) var inputUsage: AnthropicUsage?
19+
private var outputTokens: Int?
1920
private(set) var isCompleted: Bool = false
2021
private var maxBlockIndex: Int = -1
2122

@@ -27,6 +28,19 @@ actor AnthropicStreamState {
2728
inputUsage = usage
2829
}
2930

31+
func setOutputTokens(_ tokens: Int) {
32+
outputTokens = tokens
33+
}
34+
35+
func finalUsage() -> TokenUsage {
36+
TokenUsage(
37+
input: inputUsage?.inputTokens ?? 0,
38+
output: outputTokens ?? 0,
39+
cacheRead: inputUsage?.cacheReadInputTokens,
40+
cacheWrite: inputUsage?.cacheCreationInputTokens
41+
)
42+
}
43+
3044
func setBlockType(_ index: Int, _ type: BlockType) {
3145
blockTypes[index] = type
3246
maxBlockIndex = max(maxBlockIndex, index)

‎Sources/AgentRunKit/LLM/Providers/Gemini/GeminiClientStreaming.swift‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,14 +26,17 @@ extension GeminiClient {
2626

2727
let state = GeminiStreamState()
2828

29-
try await processSSEStream(
29+
let completed = try await processSSEStream(
3030
bytes: bytes,
3131
stallTimeout: retryPolicy.streamStallTimeout
3232
) { event in
3333
try await self.handleSSEEvent(
3434
event, state: state, continuation: continuation
3535
)
3636
}
37+
guard completed else {
38+
throw AgentError.llmError(.streamStalled)
39+
}
3740
continuation.finish()
3841
}
3942

‎Sources/AgentRunKit/LLM/Providers/OpenAIChat/OpenAIClientStreaming.swift‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,15 @@ extension OpenAIClient {
2222
urlRequest: urlRequest, session: session, retryPolicy: retryPolicy
2323
)
2424
onResponse?(httpResponse)
25-
try await processSSEStream(
25+
let completed = try await processSSEStream(
2626
bytes: bytes,
2727
stallTimeout: retryPolicy.streamStallTimeout
2828
) { [self] event in
2929
try handleSSEEvent(event, continuation: continuation)
3030
}
31+
guard completed else {
32+
throw AgentError.llmError(.streamStalled)
33+
}
3134
continuation.finish()
3235
}
3336

‎Sources/AgentRunKit/LLM/Providers/Vertex/VertexAnthropicClient.swift‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -161,14 +161,17 @@ public struct VertexAnthropicClient: LLMClient, Sendable {
161161

162162
let state = AnthropicStreamState()
163163

164-
try await processSSEStream(
164+
let completed = try await processSSEStream(
165165
bytes: bytes,
166166
stallTimeout: retryPolicy.streamStallTimeout
167167
) { event in
168168
try await anthropic.handleSSEEvent(
169169
event, state: state, continuation: continuation
170170
)
171171
}
172+
guard completed else {
173+
throw AgentError.llmError(.streamStalled)
174+
}
172175
continuation.finish()
173176
}
174177

@@ -196,14 +199,17 @@ public struct VertexAnthropicClient: LLMClient, Sendable {
196199

197200
let state = AnthropicStreamState()
198201

199-
try await processSSEStream(
202+
let completed = try await processSSEStream(
200203
bytes: bytes,
201204
stallTimeout: retryPolicy.streamStallTimeout
202205
) { event in
203206
try await anthropic.handleSSEEvent(event, state: state) { delta in
204207
continuation.yield(.delta(delta))
205208
}
206209
}
210+
guard completed else {
211+
throw AgentError.llmError(.streamStalled)
212+
}
207213

208214
if await state.isCompleted {
209215
let blocks = try await state.finalizedBlocks()

‎Sources/AgentRunKit/LLM/Providers/Vertex/VertexGoogleClient.swift‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,14 +150,17 @@ public struct VertexGoogleClient: LLMClient, Sendable {
150150

151151
let state = GeminiStreamState()
152152

153-
try await processSSEStream(
153+
let completed = try await processSSEStream(
154154
bytes: bytes,
155155
stallTimeout: retryPolicy.streamStallTimeout
156156
) { event in
157157
try await gemini.handleSSEEvent(
158158
event, state: state, continuation: continuation
159159
)
160160
}
161+
guard completed else {
162+
throw AgentError.llmError(.streamStalled)
163+
}
161164
continuation.finish()
162165
}
163166

‎Sources/AgentRunKit/LLM/Transport/SSEParser.swift‎

Lines changed: 19 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -88,13 +88,14 @@ func buildJSONPostRequest(
8888
return request
8989
}
9090

91+
@discardableResult
9192
func processSSEStream<S: AsyncSequence & Sendable>(
9293
bytes: S,
9394
stallTimeout: Duration?,
9495
handler: @escaping @Sendable (SSEEvent) async throws -> Bool
95-
) async throws where S.Element == UInt8 {
96+
) async throws -> Bool where S.Element == UInt8 {
9697
if let stallTimeout {
97-
try await withThrowingTaskGroup(of: Void.self) { group in
98+
return try await withThrowingTaskGroup(of: Bool.self) { group in
9899
let watchdog = StallWatchdog()
99100

100101
group.addTask {
@@ -106,6 +107,7 @@ func processSSEStream<S: AsyncSequence & Sendable>(
106107
throw AgentError.llmError(.streamStalled)
107108
}
108109
}
110+
return false
109111
}
110112

111113
group.addTask {
@@ -114,27 +116,31 @@ func processSSEStream<S: AsyncSequence & Sendable>(
114116
await watchdog.recordActivity()
115117
if let event = parser.appendLine(line),
116118
try await handler(event) {
117-
return
119+
return true
118120
}
119121
}
120122
if let event = parser.finish() {
121-
_ = try await handler(event)
123+
return try await handler(event)
122124
}
125+
return false
123126
}
124127

125-
try await group.next()
128+
let completed = try await group.next() ?? false
126129
group.cancelAll()
130+
return completed
127131
}
128-
} else {
129-
var parser = SSEEventParser()
130-
for try await line in UnboundedLines(source: bytes) {
131-
guard let event = parser.appendLine(line) else { continue }
132-
guard try await !handler(event) else { return }
133-
}
134-
if let event = parser.finish() {
135-
_ = try await handler(event)
132+
}
133+
var parser = SSEEventParser()
134+
for try await line in UnboundedLines(source: bytes) {
135+
guard let event = parser.appendLine(line) else { continue }
136+
guard try await !handler(event) else {
137+
return true
136138
}
137139
}
140+
if let event = parser.finish() {
141+
return try await handler(event)
142+
}
143+
return false
138144
}
139145

140146
actor StallWatchdog {

‎Sources/AgentRunKit/LLM/Transport/TransportError.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ public enum TransportError: Error, Sendable, Equatable, CustomStringConvertible
4040
case let .encodingFailed(description): "Encoding failed: \(description)"
4141
case let .decodingFailed(description): "Decoding failed: \(description)"
4242
case .noChoices: "No choices in response"
43-
case .streamStalled: "Stream stalled (no data received within timeout)"
43+
case .streamStalled: "Stream stalled or ended before completion"
4444
case let .capabilityMismatch(model, requirement):
4545
"Capability mismatch for model '\(model)': \(requirement)"
4646
case let .featureUnsupported(provider, feature):

‎Tests/AgentRunKitTests/Core/Streaming/StreamProcessorTests.swift‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,39 @@ struct StreamProcessorEmittedOutputTests {
131131
}
132132
}
133133

134+
struct StreamProcessorCompletionTests {
135+
@Test
136+
func streamEndingWithoutFinishedDeltaThrowsStreamStalled() async {
137+
let client = ScriptedStreamClient(
138+
deltas: [.content("partial")],
139+
error: nil
140+
)
141+
let processor = StreamProcessor(client: client, toolDefinitions: [], policy: .chat, eventFactory: testFactory)
142+
let (_, eventContinuation) = AsyncThrowingStream<StreamEvent, Error>.makeStream()
143+
var totalUsage = TokenUsage()
144+
var emittedOutput = false
145+
146+
do {
147+
_ = try await processor.process(
148+
messages: [.user("Hi")],
149+
totalUsage: &totalUsage,
150+
emittedOutput: &emittedOutput,
151+
continuation: eventContinuation
152+
)
153+
Issue.record("Expected streamStalled error")
154+
} catch let error as AgentError {
155+
guard case let .llmError(transport) = error else {
156+
Issue.record("Expected llmError, got \(error)")
157+
return
158+
}
159+
#expect(transport == .streamStalled)
160+
#expect(emittedOutput)
161+
} catch {
162+
Issue.record("Expected AgentError, got \(error)")
163+
}
164+
}
165+
}
166+
134167
struct StreamProcessorToolCallAccumulationTests {
135168
@Test
136169
func duplicateToolCallStartDoesNotResetAccumulatedArguments() async throws {

0 commit comments

Comments
 (0)