diff options
Diffstat (limited to 'libpanto/src/provider_openai_chat.zig')
| -rw-r--r-- | libpanto/src/provider_openai_chat.zig | 572 |
1 files changed, 294 insertions, 278 deletions
diff --git a/libpanto/src/provider_openai_chat.zig b/libpanto/src/provider_openai_chat.zig index ca01d7c..30bccb6 100644 --- a/libpanto/src/provider_openai_chat.zig +++ b/libpanto/src/provider_openai_chat.zig @@ -18,11 +18,15 @@ const Uri = std.Uri; const conversation = @import("conversation.zig"); const provider_mod = @import("provider.zig"); +const stream_mod = @import("stream.zig"); const sse_mod = @import("sse.zig"); const json_mod = @import("openai_chat_json.zig"); const config_mod = @import("config.zig"); const tool_registry_mod = @import("tool_registry.zig"); +const Event = stream_mod.Event; +const EventQueue = stream_mod.EventQueue; + /// Decode a wire tool name (`__` -> `.`) in place within an assembled /// name buffer. Decoding only ever shrinks the buffer (reads stay ahead /// of writes), so aliasing src/dst is safe; we then truncate to the @@ -52,27 +56,31 @@ pub const OpenAIChatRequest = struct { /// they live only as long as this request object. diag: ?*provider_mod.ProviderDiagnostic = null, - pub fn streamStep( + /// Open the streaming HTTP request and return a heap-allocated + /// resumable response. Performs the POST and reads response headers + /// (classifying any >=400 status into a provider error), but does NOT + /// pump the body — that happens lazily in `ResumableResponse.produce`. + /// + /// On success the caller owns the returned `*ResumableResponse` and must + /// `deinit` it. On failure nothing is allocated. + pub fn open( self: *OpenAIChatRequest, conv: *conversation.Conversation, tools: *const provider_mod.ToolRegistry, - receiver: *provider_mod.Receiver, - ) !void { - // Outer wrapper guarantees `onError` is called exactly once if - // anything fails — whether the receiver, the HTTP transport, or the - // SSE/JSON parsers. The inner `streamStepInner` does the real work. - self.streamStepInner(conv, tools, receiver) catch |err| { - receiver.onError(err); - return err; + ) !*ResumableResponse { + const rr = try self.allocator.create(ResumableResponse); + errdefer self.allocator.destroy(rr); + rr.* = .{ + .allocator = self.allocator, + .conv = conv, + .parser = sse_mod.SSEParser.init(self.allocator), + .state = .init(self.allocator), }; - } + errdefer { + rr.parser.deinit(); + rr.state.deinit(); + } - fn streamStepInner( - self: *OpenAIChatRequest, - conv: *conversation.Conversation, - tools: *const provider_mod.ToolRegistry, - receiver: *provider_mod.Receiver, - ) !void { // Build URL: "{base_url}/chat/completions" const url = try std.fmt.allocPrint( self.allocator, @@ -102,8 +110,10 @@ pub const OpenAIChatRequest = struct { }; // Open the request. We can't use `fetch()` because it buffers the - // response; we want to stream the body as it arrives. - var req = try self.http_client.request(.POST, uri, .{ + // response; we want to stream the body as it arrives. The request + // is moved into the heap struct so the body reader (which borrows + // `&rr.response`) stays valid across `produce` calls. + rr.req = try self.http_client.request(.POST, uri, .{ .extra_headers = &extra_headers, // Disable compression: gzip buffers small SSE frames, defeating // the streaming property we paid for `stream: true` to get. @@ -111,24 +121,27 @@ pub const OpenAIChatRequest = struct { .keep_alive = false, .redirect_behavior = .not_allowed, }); - defer req.deinit(); + rr.req_open = true; + errdefer { + rr.req.deinit(); + rr.req_open = false; + } - req.transfer_encoding = .{ .content_length = body.len }; + rr.req.transfer_encoding = .{ .content_length = body.len }; var send_buf: [4096]u8 = undefined; - var bw = try req.sendBodyUnflushed(&send_buf); + var bw = try rr.req.sendBodyUnflushed(&send_buf); try bw.writer.writeAll(body); try bw.end(); - try req.connection.?.flush(); + try rr.req.connection.?.flush(); // Receive response headers. var redirect_buf: [1024]u8 = undefined; - var response = try req.receiveHead(&redirect_buf); + rr.response = try rr.req.receiveHead(&redirect_buf); - if (@intFromEnum(response.head.status) >= 400) { + if (@intFromEnum(rr.response.head.status) >= 400) { // Drain body for diagnostics. - var transfer_buf: [4096]u8 = undefined; - const body_reader = response.reader(&transfer_buf); + const body_reader = rr.response.reader(&rr.transfer_buf); var err_buf: std.ArrayList(u8) = .empty; defer err_buf.deinit(self.allocator); var tmp: [1024]u8 = undefined; @@ -138,7 +151,7 @@ pub const OpenAIChatRequest = struct { try err_buf.appendSlice(self.allocator, tmp[0..n]); if (err_buf.items.len > 16 * 1024) break; } - const status: u16 = @intFromEnum(response.head.status); + const status: u16 = @intFromEnum(rr.response.head.status); std.log.err("openai_chat HTTP {d}: {s}", .{ status, err_buf.items }); // Classify the status into a retryable/terminal provider error. // HTTP 400 with a context marker becomes `ContextOverflow` so the @@ -146,64 +159,127 @@ pub const OpenAIChatRequest = struct { const classified = provider_mod.classifyHttpStatus(status, err_buf.items); if (self.diag) |d| { d.status_code = status; - d.retry_after_ms = provider_mod.retryAfterFromHead(response.head); + d.retry_after_ms = provider_mod.retryAfterFromHead(rr.response.head); } return classified; } - // Stream the body through the SSE parser and event handler. - var transfer_buf: [4096]u8 = undefined; - const body_reader = response.reader(&transfer_buf); + // Bind the streaming body reader. Valid for the lifetime of `rr` + // (it borrows `&rr.response` and `&rr.transfer_buf`, both pinned). + rr.body_reader = rr.response.reader(&rr.transfer_buf); + return rr; + } +}; - var parser = sse_mod.SSEParser.init(self.allocator); - defer parser.deinit(); +/// A resumable OpenAI Chat streaming response. Owns the pinned HTTP +/// request/response, the body reader's transfer buffer, the `SSEParser`, +/// and the block-assembly `StreamState`. `produce` pumps just enough bytes +/// to emit one or more events into the queue, or reports the response is +/// complete (its assistant message committed to the conversation). +/// +/// Must be heap-allocated and never moved: `body_reader` borrows +/// `&self.response`. +pub const ResumableResponse = struct { + allocator: Allocator, + conv: *conversation.Conversation, + parser: sse_mod.SSEParser, + state: StreamState, - var state: StreamState = .init(self.allocator); - defer state.deinit(); + req: http.Client.Request = undefined, + response: http.Client.Response = undefined, + /// Transfer buffer backing `body_reader`. Pinned in the heap struct. + transfer_buf: [4096]u8 = undefined, + /// The streaming body reader, bound in `open` after a 2xx response. + body_reader: *std.Io.Reader = undefined, + /// Chunk scratch for `readVec`. + chunk: [4096]u8 = undefined, - // Use `readVec` so we return to the event loop as soon as *any* - // bytes arrive, rather than waiting for the buffer to fill. - // `readSliceShort` blocks until EOF or full, which defeats streaming. - // - // Per `std.Io.Reader.readVec` docs: a return value of 0 does NOT - // signal end-of-stream — it just means no new bytes were available - // this call, and the caller should try again. EOF is reported only - // via `error.EndOfStream`. Breaking on `n == 0` truncates the - // response and silently cuts off mid-stream. - var chunk: [4096]u8 = undefined; - var vecs: [1][]u8 = .{&chunk}; - while (true) { - const n = body_reader.readVec(&vecs) catch |err| switch (err) { - error.EndOfStream => break, - // A transport read failure mid-stream (reset, TLS, timeout) - // before `[DONE]` means no assistant message was committed. - // Surface it as a retryable malformed-stream error. - else => return error.ProviderStreamMalformed, - }; - if (n == 0) continue; + /// True once `req` has been initialized (so `deinit` knows to free it). + req_open: bool = false, + /// Set once the response is fully decoded and finalized. + done: bool = false, - const events = try parser.feed(chunk[0..n]); - defer parser.freeEvents(events); + pub const ProduceStatus = provider_mod.ProviderStream.ProduceStatus; - for (events) |ev_payload| { - std.log.debug("openai_chat <= {s}", .{ev_payload}); - if (std.mem.eql(u8, ev_payload, "[DONE]")) { - try state.finalize(receiver, conv); - return; - } - try handleEvent(self.allocator, ev_payload, &state, receiver); - // Note: we do NOT bail when state.end_of_stream is set. - // OpenAI emits the terminating `usage` chunk *after* the - // chunk carrying finish_reason, then sends `[DONE]`. If - // we returned on finish_reason we'd never capture usage. - // `[DONE]` is the authoritative end-of-stream marker. + /// Wrap this response in the provider-agnostic `ProviderStream` the agent + /// loop drives. + pub fn providerStream(self: *ResumableResponse) provider_mod.ProviderStream { + return .{ .ptr = self, .vtable = &vtable }; + } + + const vtable: provider_mod.ProviderStream.VTable = .{ + .produce = produceVT, + .deinit = deinitVT, + }; + + fn produceVT(ptr: *anyopaque, out: *EventQueue) anyerror!ProduceStatus { + const self: *ResumableResponse = @ptrCast(@alignCast(ptr)); + return self.produce(out); + } + + fn deinitVT(ptr: *anyopaque) void { + const self: *ResumableResponse = @ptrCast(@alignCast(ptr)); + self.deinit(); + } + + pub fn deinit(self: *ResumableResponse) void { + if (self.req_open) self.req.deinit(); + self.parser.deinit(); + self.state.deinit(); + self.allocator.destroy(self); + } + + /// Pump the response: read one chunk, feed it through the SSE parser, + /// and decode each SSE event into zero or more `Event`s appended to + /// `out`. Returns `.more` if the caller should pump again, or + /// `.response_complete` once the terminal (`[DONE]` or EOF) has been + /// reached and the assistant message has been committed + a final + /// `message_complete` pushed. + /// + /// Reading and finalizing here means a single `produce` call may push + /// several events; the `Stream` drains the queue before pumping again. + pub fn produce(self: *ResumableResponse, out: *EventQueue) !ProduceStatus { + if (self.done) return .response_complete; + + var vecs: [1][]u8 = .{&self.chunk}; + const n = self.body_reader.readVec(&vecs) catch |err| switch (err) { + // Stream ended without [DONE]. Some servers and proxies omit it + // (or drop the trailing usage chunk). Finalize with whatever + // we've got — usage will be null in that case, which is fine. + error.EndOfStream => { + try self.finishStream(out); + return .response_complete; + }, + // A transport read failure mid-stream (reset, TLS, timeout) + // before `[DONE]` means no assistant message was committed. + // Surface it as a retryable malformed-stream error. + else => return error.ProviderStreamMalformed, + }; + if (n == 0) return .more; + + const events = try self.parser.feed(self.chunk[0..n]); + defer self.parser.freeEvents(events); + + for (events) |ev_payload| { + std.log.debug("openai_chat <= {s}", .{ev_payload}); + if (std.mem.eql(u8, ev_payload, "[DONE]")) { + try self.finishStream(out); + return .response_complete; } + try handleEvent(self.allocator, ev_payload, &self.state, out); + // Note: we do NOT bail when state.end_of_stream is set. + // OpenAI emits the terminating `usage` chunk *after* the + // chunk carrying finish_reason, then sends `[DONE]`. If + // we returned on finish_reason we'd never capture usage. + // `[DONE]` is the authoritative end-of-stream marker. } + return .more; + } - // Stream ended without [DONE]. Some servers and proxies omit it - // (or drop the trailing usage chunk). Finalize with whatever we've - // got — usage will be null in that case, which is fine. - try state.finalize(receiver, conv); + fn finishStream(self: *ResumableResponse, out: *EventQueue) !void { + if (self.done) return; + self.done = true; + try self.state.finalize(out, self.conv); } }; @@ -310,7 +386,7 @@ const StreamState = struct { /// Close the active text/thinking block (if any) and emit /// onBlockComplete. Ownership of `current_buf` transfers into the /// appended block. - fn closeActive(self: *StreamState, receiver: *provider_mod.Receiver) !void { + fn closeActive(self: *StreamState, out: *EventQueue) !void { if (self.active == .none) return; const block: conversation.ContentBlock = switch (self.active) { @@ -321,7 +397,10 @@ const StreamState = struct { self.current_buf = .empty; try self.blocks.append(self.allocator, block); - try receiver.onBlockComplete(self.block_index, self.blocks.items[self.blocks.items.len - 1]); + try out.push(.{ .block_complete = .{ + .index = self.block_index, + .block = self.blocks.items[self.blocks.items.len - 1], + } }); self.active = .none; } @@ -330,12 +409,12 @@ const StreamState = struct { fn openBlock( self: *StreamState, new_active: ActiveBlock, - receiver: *provider_mod.Receiver, + out: *EventQueue, ) !void { std.debug.assert(new_active == .text or new_active == .thinking); if (self.active == new_active) return; if (self.active != .none) { - try self.closeActive(receiver); + try self.closeActive(out); self.block_index += 1; } self.active = new_active; @@ -344,16 +423,21 @@ const StreamState = struct { .thinking => .Thinking, .tool_use, .none => unreachable, }; - try receiver.onBlockStart(block_type, self.block_index); + try out.push(.{ .block_start = .{ .block_type = block_type, .index = self.block_index } }); } fn appendDelta( self: *StreamState, - receiver: *provider_mod.Receiver, + out: *EventQueue, delta: []const u8, ) !void { try self.current_buf.appendSlice(self.allocator, delta); - try receiver.onContentDelta(self.block_index, delta); + // Dupe into the queue arena: the raw `delta` borrows the transient + // SSE payload that `produce` frees before `next()` reads the queue. + try out.push(.{ .content_delta = .{ + .index = self.block_index, + .delta = try out.dupeBytes(delta), + } }); } /// Apply one streaming tool_call delta. Opens a new tool_use on the @@ -362,7 +446,7 @@ const StreamState = struct { /// a malformed stream — we log and drop it. fn applyToolCallDelta( self: *StreamState, - receiver: *provider_mod.Receiver, + out: *EventQueue, d: json_mod.ToolCallDelta, ) !void { // Degenerate backend: a delta arrived for an index whose block we @@ -382,14 +466,14 @@ const StreamState = struct { // the only signal openai_chat gives us for mid-stream tool_use // boundaries; see the StreamState doc-comment for the rationale. if (self.current_tool_index) |cur| { - if (cur != d.index) try self.closeActiveTool(receiver); + if (cur != d.index) try self.closeActiveTool(out); } if (self.active_tool == null) { // Opening a new tool_use. First close any open text/thinking // block so the tool_use gets its own block_index. if (self.active != .none) { - try self.closeActive(receiver); + try self.closeActive(out); self.block_index += 1; } self.active_tool = .{ .block_index = self.block_index }; @@ -410,32 +494,39 @@ const StreamState = struct { // to render. If the block closes before any args arrive (zero-arg // tool), `closeActiveTool` emits the start there. if (d.arguments) |a| { - try self.emitStartIfNeeded(receiver, tu); - // Fire `onToolDetails` as soon as both id and name are + try self.emitStartIfNeeded(out, tu); + // Fire `tool_details` as soon as both id and name are // known. We can't know identity is *final* until the block // closes (a later delta could append more bytes), but in // practice OpenAI sends each whole on the first delta. A // pathological backend that streams id/name across many // chunks would have us emit a truncated value here. We - // accept that trade-off: receivers that need the canonical + // accept that trade-off: consumers that need the canonical // value can read it from the assembled ContentBlock at - // onBlockComplete. - try self.emitDetailsIfReady(receiver, tu); + // block_complete. + try self.emitDetailsIfReady(out, tu); try tu.arguments.appendSlice(self.allocator, a); - try receiver.onContentDelta(tu.block_index, a); + // Dupe into the queue arena (the SSE payload is freed before + // `next()` reads the queue). + try out.push(.{ .content_delta = .{ + .index = tu.block_index, + .delta = try out.dupeBytes(a), + } }); } else { // Identity-only chunk (no args yet). Still try to emit // details, in case both fields are now populated. - if (tu.started) try self.emitDetailsIfReady(receiver, tu); + if (tu.started) try self.emitDetailsIfReady(out, tu); } } - /// Fire `onToolDetails` once both id and name are non-empty. No-op if + /// Fire `tool_details` once both id and name are non-empty. No-op if /// already fired or if either field is still empty. Requires that - /// `onBlockStart` has already been emitted. + /// `block_start` has already been emitted. Slices are duped into the + /// queue arena because `id_buf`/`name_buf` may still grow (and realloc) + /// on later fragments. fn emitDetailsIfReady( self: *StreamState, - receiver: *provider_mod.Receiver, + out: *EventQueue, tu: *ToolUseInProgress, ) !void { _ = self; @@ -443,13 +534,17 @@ const StreamState = struct { if (!tu.started) return; if (tu.id_buf.items.len == 0 or tu.name_buf.items.len == 0) return; tu.details_emitted = true; - try receiver.onToolDetails(tu.block_index, tu.id_buf.items, tu.name_buf.items); + try out.push(.{ .tool_details = .{ + .index = tu.block_index, + .id = try out.dupeBytes(tu.id_buf.items), + .name = try out.dupeBytes(tu.name_buf.items), + } }); } - /// Close the currently-active tool_use (if any), emitting onBlockStart - /// (if it wasn't already), onBlockComplete, and recording the wire + /// Close the currently-active tool_use (if any), emitting block_start + /// (if it wasn't already), block_complete, and recording the wire /// index as closed. No-op if there's no active tool_use. - fn closeActiveTool(self: *StreamState, receiver: *provider_mod.Receiver) !void { + fn closeActiveTool(self: *StreamState, out: *EventQueue) !void { var tu = self.active_tool orelse return; self.active_tool = null; const wire_index = self.current_tool_index.?; @@ -481,12 +576,12 @@ const StreamState = struct { // internal (dotted) name. Decoding never grows the buffer. decodeNameInPlace(&tu.name_buf); - // If no arguments ever arrived, we haven't emitted onBlockStart - // yet — do it now so the receiver sees a balanced start/complete. - try self.emitStartIfNeeded(receiver, &tu); + // If no arguments ever arrived, we haven't emitted block_start + // yet — do it now so the consumer sees a balanced start/complete. + try self.emitStartIfNeeded(out, &tu); // Last chance to fire details if a fragmented-identity provider // only finished id/name accumulation at the very end. - try self.emitDetailsIfReady(receiver, &tu); + try self.emitDetailsIfReady(out, &tu); const id_owned = try tu.id_buf.toOwnedSlice(self.allocator); const name_owned = try tu.name_buf.toOwnedSlice(self.allocator); @@ -500,39 +595,41 @@ const StreamState = struct { tu.arguments = .empty; try self.blocks.append(self.allocator, block); - try receiver.onBlockComplete(tu.block_index, self.blocks.items[self.blocks.items.len - 1]); + try out.push(.{ .block_complete = .{ + .index = tu.block_index, + .block = self.blocks.items[self.blocks.items.len - 1], + } }); } - /// Emit `onBlockStart(.ToolUse, ...)` once per in-progress tool use. - /// Callers must invoke this before the first `onContentDelta` or - /// `onBlockComplete` for the block. Identity (id/name) is *not* - /// passed at start — see provider.zig's ReceiverVTable docs for the - /// rationale. Receivers get identity from the assembled ContentBlock - /// at onBlockComplete time. + /// Emit `block_start(.ToolUse, ...)` once per in-progress tool use. + /// Callers must invoke this before the first `content_delta` or + /// `block_complete` for the block. Identity (id/name) is *not* passed at + /// start — consumers get identity from `tool_details` or the assembled + /// ContentBlock at block_complete time. fn emitStartIfNeeded( self: *StreamState, - receiver: *provider_mod.Receiver, + out: *EventQueue, tu: *ToolUseInProgress, ) !void { _ = self; if (tu.started) return; tu.started = true; - try receiver.onBlockStart(.ToolUse, tu.block_index); + try out.push(.{ .block_start = .{ .block_type = .ToolUse, .index = tu.block_index } }); } /// End the stream: close any open text/thinking block, close the still- /// active tool_use (if any), then commit the assembled assistant - /// Message to the conversation. + /// Message to the conversation and push the terminal `message_complete`. fn finalize( self: *StreamState, - receiver: *provider_mod.Receiver, + out: *EventQueue, conv: *conversation.Conversation, ) !void { if (self.finalized) return; self.finalized = true; - try self.closeActive(receiver); - try self.closeActiveTool(receiver); + try self.closeActive(out); + try self.closeActiveTool(out); // Move blocks into a fresh conversation message. const moved_blocks = try self.blocks.toOwnedSlice(self.allocator); @@ -541,7 +638,7 @@ const StreamState = struct { try conv.addAssistantMessageWithUsage(moved_blocks, self.usage); const msg = conv.messages.items[conv.messages.items.len - 1]; - try receiver.onMessageComplete(msg, self.usage); + try out.push(.{ .message_complete = .{ .message = msg, .usage = self.usage } }); } }; @@ -549,7 +646,7 @@ fn handleEvent( allocator: Allocator, payload: []const u8, state: *StreamState, - receiver: *provider_mod.Receiver, + out: *EventQueue, ) !void { var parsed = try json_mod.parseStreamEvent(allocator, payload); defer parsed.deinit(); @@ -587,33 +684,33 @@ fn handleEvent( if (!state.started and d.role != null) { state.started = true; - try receiver.onMessageStart(.assistant); + try out.push(.{ .message_start = .assistant }); } if (d.reasoning_content) |rc| { if (!state.started) { state.started = true; - try receiver.onMessageStart(.assistant); + try out.push(.{ .message_start = .assistant }); } - try state.openBlock(.thinking, receiver); - try state.appendDelta(receiver, rc); + try state.openBlock(.thinking, out); + try state.appendDelta(out, rc); } if (d.content) |c| { if (!state.started) { state.started = true; - try receiver.onMessageStart(.assistant); + try out.push(.{ .message_start = .assistant }); } - try state.openBlock(.text, receiver); - try state.appendDelta(receiver, c); + try state.openBlock(.text, out); + try state.appendDelta(out, c); } if (d.tool_calls.len > 0) { if (!state.started) { state.started = true; - try receiver.onMessageStart(.assistant); + try out.push(.{ .message_start = .assistant }); } - for (d.tool_calls) |tc| try state.applyToolCallDelta(receiver, tc); + for (d.tool_calls) |tc| try state.applyToolCallDelta(out, tc); } if (d.finish_reason) |_| { @@ -627,54 +724,76 @@ fn handleEvent( const testing = std.testing; -/// A no-op Receiver that drops every callback. Useful when the test cares -/// about post-stream conversation state rather than callback observability. -const NoopReceiver = struct { - fn make() provider_mod.Receiver { - return .{ .ptr = @ptrCast(@constCast(&dummy)), .vtable = &vt }; - } - var dummy: u8 = 0; - const vt: provider_mod.ReceiverVTable = .{ - .onMessageStart = noopMsgStart, - .onBlockStart = noopBlockStart, - .onToolDetails = noopToolDetails, - .onContentDelta = noopDelta, - .onBlockComplete = noopBlockComplete, - .onMessageComplete = noopMsgComplete, - .onError = noopErr, - .onProviderRetry = noopRetry, - }; - fn noopMsgStart(_: *anyopaque, _: conversation.MessageRole) anyerror!void {} - fn noopBlockStart(_: *anyopaque, _: provider_mod.ContentBlockType, _: usize) anyerror!void {} - fn noopToolDetails(_: *anyopaque, _: usize, _: []const u8, _: []const u8) anyerror!void {} - fn noopDelta(_: *anyopaque, _: usize, _: []const u8) anyerror!void {} - fn noopBlockComplete(_: *anyopaque, _: usize, _: conversation.ContentBlock) anyerror!void {} - fn noopMsgComplete(_: *anyopaque, _: conversation.Message, _: ?provider_mod.Usage) anyerror!void {} - fn noopErr(_: *anyopaque, _: anyerror) void {} - fn noopRetry(_: *anyopaque, _: provider_mod.ProviderRetryInfo) void {} -}; - /// Feed a sequence of SSE event payloads through the state machine as if -/// they had been delivered by the wire, finalizing into `conv`. +/// they had been delivered by the wire, finalizing into `conv`. The decoded +/// `Event`s are recorded as compact strings (the same schema the old +/// RecordingReceiver used) so callback-ordering assertions are preserved. fn runStreamedTurn( allocator: Allocator, conv: *conversation.Conversation, - receiver: *provider_mod.Receiver, + rec: ?*EventRecorder, events: []const []const u8, ) !void { var state: StreamState = .init(allocator); defer state.deinit(); + var queue = EventQueue.init(allocator); + defer queue.deinit(); + for (events) |payload| { if (std.mem.eql(u8, payload, "[DONE]")) break; // Process every chunk through to [DONE], including the - // post-finish_reason usage chunk. Mirrors the production loop - // in OpenAIChatRequest.streamStep. - try handleEvent(allocator, payload, &state, receiver); + // post-finish_reason usage chunk. Mirrors the production pump in + // ResumableResponse.produce. + try handleEvent(allocator, payload, &state, &queue); + } + try state.finalize(&queue, conv); + + // Drain into the recorder before the arena resets. The queue holds all + // events from this turn; popping records each, and the final null-pop + // resets the arena. + while (queue.pop()) |ev| { + if (rec) |r| try r.record(ev); } - try state.finalize(receiver, conv); } +/// Records decoded `Event`s as compact strings for ordering assertions. +const EventRecorder = struct { + allocator: Allocator, + events: std.ArrayList([]const u8) = .empty, + + fn deinit(self: *EventRecorder) void { + for (self.events.items) |e| self.allocator.free(e); + self.events.deinit(self.allocator); + } + + fn push(self: *EventRecorder, comptime fmt: []const u8, args: anytype) !void { + const owned = try std.fmt.allocPrint(self.allocator, fmt, args); + try self.events.append(self.allocator, owned); + } + + fn record(self: *EventRecorder, ev: Event) !void { + switch (ev) { + .message_start => try self.push("msg_start", .{}), + .block_start => |b| try self.push("block_start[{d}]:{s}", .{ b.index, @tagName(b.block_type) }), + .tool_details => |t| try self.push("tool_details[{d}]:{s}:{s}", .{ t.index, t.id, t.name }), + .content_delta => |d| try self.push("delta[{d}]:{s}", .{ d.index, d.delta }), + .block_complete => |b| try self.push("block_complete[{d}]", .{b.index}), + .message_complete => |m| { + if (m.usage) |u| { + try self.push( + "msg_complete[usage:in={d},out={d},cr={d},cw={d},rsn={d}]", + .{ u.input, u.output, u.cache_read, u.cache_write, u.reasoning }, + ); + } else { + try self.push("msg_complete[usage:null]", .{}); + } + }, + else => {}, + } + } +}; + test "two streamed turns persist assistant replies in the conversation" { // Regression test for the bug where `finish_reason` arrived before // `[DONE]` and `finalize` early-returned without appending the assistant @@ -688,8 +807,6 @@ test "two streamed turns persist assistant replies in the conversation" { try conv.addSystemMessage("You are a helpful assistant."); try conv.addUserMessage("hello!"); - var recv = NoopReceiver.make(); - const turn1 = [_][]const u8{ \\{"choices":[{"delta":{"role":"assistant"}}]} , @@ -701,7 +818,7 @@ test "two streamed turns persist assistant replies in the conversation" { , "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &turn1); + try runStreamedTurn(allocator, &conv, null, &turn1); try testing.expectEqual(@as(usize, 3), conv.messages.items.len); try testing.expectEqual(conversation.MessageRole.assistant, conv.messages.items[2].role); @@ -722,7 +839,7 @@ test "two streamed turns persist assistant replies in the conversation" { , "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &turn2); + try runStreamedTurn(allocator, &conv, null, &turn2); // System + user + assistant + user + assistant = 5 messages. try testing.expectEqual(@as(usize, 5), conv.messages.items.len); @@ -740,12 +857,8 @@ test "openai_chat: terminating usage chunk lands on message_complete with split defer conv.deinit(); try conv.addUserMessage("hi"); - var rec = RecordingReceiver{ .allocator = allocator }; - defer { - for (rec.events.items) |s| allocator.free(s); - rec.events.deinit(allocator); - } - var recv = rec.receiver(); + var rec = EventRecorder{ .allocator = allocator }; + defer rec.deinit(); const events = [_][]const u8{ \\{"choices":[{"delta":{"role":"assistant"}}]} @@ -759,7 +872,7 @@ test "openai_chat: terminating usage chunk lands on message_complete with split , "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &events); + try runStreamedTurn(allocator, &conv, &rec, &events); var found: ?[]const u8 = null; for (rec.events.items) |s| { @@ -777,12 +890,8 @@ test "openai_chat: omitted stream usage yields null on message_complete" { defer conv.deinit(); try conv.addUserMessage("hi"); - var rec = RecordingReceiver{ .allocator = allocator }; - defer { - for (rec.events.items) |s| allocator.free(s); - rec.events.deinit(allocator); - } - var recv = rec.receiver(); + var rec = EventRecorder{ .allocator = allocator }; + defer rec.deinit(); const events = [_][]const u8{ \\{"choices":[{"delta":{"role":"assistant"}}]} @@ -793,7 +902,7 @@ test "openai_chat: omitted stream usage yields null on message_complete" { , "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &events); + try runStreamedTurn(allocator, &conv, &rec, &events); var found: ?[]const u8 = null; for (rec.events.items) |s| { @@ -814,8 +923,6 @@ test "fragmented tool_call id and name are reassembled" { defer conv.deinit(); try conv.addUserMessage("call something"); - var recv = NoopReceiver.make(); - const events = [_][]const u8{ \\{"choices":[{"delta":{"role":"assistant"}}]} , @@ -830,7 +937,7 @@ test "fragmented tool_call id and name are reassembled" { "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &events); + try runStreamedTurn(allocator, &conv, null, &events); const asst = conv.messages.items[1]; try testing.expectEqual(@as(usize, 1), asst.content.items.len); @@ -851,8 +958,6 @@ test "inbound wire tool name is decoded to dotted form (even split across __)" { defer conv.deinit(); try conv.addUserMessage("read a file"); - var recv = NoopReceiver.make(); - const events = [_][]const u8{ \\{"choices":[{"delta":{"role":"assistant"}}]} , @@ -867,96 +972,12 @@ test "inbound wire tool name is decoded to dotted form (even split across __)" { "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &events); + try runStreamedTurn(allocator, &conv, null, &events); const tu = conv.messages.items[1].content.items[0].ToolUse; try testing.expectEqualStrings("std.read", tu.name); } -/// A Receiver that records the sequence of callback events as compact -/// strings. Useful for asserting per-block start/complete ordering. -const RecordingReceiver = struct { - allocator: Allocator, - events: std.ArrayList([]const u8) = .empty, - - fn receiver(self: *RecordingReceiver) provider_mod.Receiver { - return .{ .ptr = self, .vtable = &vt }; - } - - const vt: provider_mod.ReceiverVTable = .{ - .onMessageStart = onMessageStart, - .onBlockStart = onBlockStart, - .onToolDetails = onToolDetails, - .onContentDelta = onContentDelta, - .onBlockComplete = onBlockComplete, - .onMessageComplete = onMessageComplete, - .onError = onError, - .onProviderRetry = onProviderRetry, - }; - - fn onProviderRetry(_: *anyopaque, _: provider_mod.ProviderRetryInfo) void {} - - fn record(self: *RecordingReceiver, s: []const u8) !void { - const owned = try self.allocator.dupe(u8, s); - try self.events.append(self.allocator, owned); - } - - fn recordFmt(self: *RecordingReceiver, comptime fmt: []const u8, args: anytype) !void { - const owned = try std.fmt.allocPrint(self.allocator, fmt, args); - try self.events.append(self.allocator, owned); - } - - fn deinit(self: *RecordingReceiver) void { - for (self.events.items) |e| self.allocator.free(e); - self.events.deinit(self.allocator); - } - - fn onMessageStart(ptr: *anyopaque, _: conversation.MessageRole) anyerror!void { - const self: *RecordingReceiver = @ptrCast(@alignCast(ptr)); - try self.record("msg_start"); - } - fn onBlockStart( - ptr: *anyopaque, - bt: provider_mod.ContentBlockType, - idx: usize, - ) anyerror!void { - const self: *RecordingReceiver = @ptrCast(@alignCast(ptr)); - try self.recordFmt("block_start[{d}]:{s}", .{ idx, @tagName(bt) }); - } - fn onToolDetails( - ptr: *anyopaque, - idx: usize, - id: []const u8, - name: []const u8, - ) anyerror!void { - const self: *RecordingReceiver = @ptrCast(@alignCast(ptr)); - try self.recordFmt("tool_details[{d}]:{s}:{s}", .{ idx, id, name }); - } - fn onContentDelta(ptr: *anyopaque, idx: usize, delta: []const u8) anyerror!void { - const self: *RecordingReceiver = @ptrCast(@alignCast(ptr)); - try self.recordFmt("delta[{d}]:{s}", .{ idx, delta }); - } - fn onBlockComplete( - ptr: *anyopaque, - idx: usize, - _: conversation.ContentBlock, - ) anyerror!void { - const self: *RecordingReceiver = @ptrCast(@alignCast(ptr)); - try self.recordFmt("block_complete[{d}]", .{idx}); - } - fn onMessageComplete(ptr: *anyopaque, _: conversation.Message, usage: ?provider_mod.Usage) anyerror!void { - const self: *RecordingReceiver = @ptrCast(@alignCast(ptr)); - if (usage) |u| { - try self.recordFmt( - "msg_complete[usage:in={d},out={d},cr={d},cw={d},rsn={d}]", - .{ u.input, u.output, u.cache_read, u.cache_write, u.reasoning }, - ); - } else { - try self.record("msg_complete[usage:null]"); - } - } - fn onError(_: *anyopaque, _: anyerror) void {} -}; test "parallel tool_calls emit one complete start/delta/complete cycle per block" { // Regression test: previously, the OpenAI provider deferred ALL @@ -972,9 +993,8 @@ test "parallel tool_calls emit one complete start/delta/complete cycle per block defer conv.deinit(); try conv.addUserMessage("ping four hosts"); - var rec: RecordingReceiver = .{ .allocator = allocator }; + var rec: EventRecorder = .{ .allocator = allocator }; defer rec.deinit(); - var recv = rec.receiver(); const events = [_][]const u8{ \\{"choices":[{"delta":{"role":"assistant"}}]} @@ -992,7 +1012,7 @@ test "parallel tool_calls emit one complete start/delta/complete cycle per block "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &events); + try runStreamedTurn(allocator, &conv, &rec, &events); const expected = [_][]const u8{ "msg_start", @@ -1039,9 +1059,8 @@ test "non-contiguous tool_call deltas: re-emission of a closed index is dropped" defer conv.deinit(); try conv.addUserMessage("go"); - var rec: RecordingReceiver = .{ .allocator = allocator }; + var rec: EventRecorder = .{ .allocator = allocator }; defer rec.deinit(); - var recv = rec.receiver(); const events = [_][]const u8{ \\{"choices":[{"delta":{"role":"assistant"}}]} @@ -1058,7 +1077,7 @@ test "non-contiguous tool_call deltas: re-emission of a closed index is dropped" "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &events); + try runStreamedTurn(allocator, &conv, &rec, &events); // Two well-formed tool_use blocks in the final message, args unaffected // by the dropped fragment. @@ -1088,9 +1107,8 @@ test "onToolDetails fires after id+name complete, even mid-arg-stream" { defer conv.deinit(); try conv.addUserMessage("go"); - var rec: RecordingReceiver = .{ .allocator = allocator }; + var rec: EventRecorder = .{ .allocator = allocator }; defer rec.deinit(); - var recv = rec.receiver(); const events = [_][]const u8{ \\{"choices":[{"delta":{"role":"assistant"}}]} @@ -1114,7 +1132,7 @@ test "onToolDetails fires after id+name complete, even mid-arg-stream" { "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &events); + try runStreamedTurn(allocator, &conv, &rec, &events); // Exactly one tool_details event, fired with the id-prefix that was // current at first-args-arrival, and ordered between block_start and @@ -1155,8 +1173,6 @@ test "tool_call with no arguments still finalizes a well-formed ToolUse" { defer conv.deinit(); try conv.addUserMessage("ring it"); - var recv = NoopReceiver.make(); - const events = [_][]const u8{ \\{"choices":[{"delta":{"role":"assistant"}}]} , @@ -1167,7 +1183,7 @@ test "tool_call with no arguments still finalizes a well-formed ToolUse" { "[DONE]", }; - try runStreamedTurn(allocator, &conv, &recv, &events); + try runStreamedTurn(allocator, &conv, null, &events); const asst = conv.messages.items[1]; try testing.expectEqual(@as(usize, 1), asst.content.items.len); |