ironside 0.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,276 @@
1
+ import type { IronsideClient, TraceHandle } from "../client.js";
2
+ import { errorEndOptions, instrumentAsyncIterable } from "./streaming.js";
3
+
4
+ // Wraps the Anthropic Node SDK's messages.create() to automatically record
5
+ // a generation. Verified against @anthropic-ai/sdk@0.111.0: usage lives at
6
+ // message.usage.{input_tokens,output_tokens} (snake_case — same field
7
+ // names as OpenAI's Responses API, but this is Anthropic's only call
8
+ // path, unlike OpenAI's Chat-Completions-vs-Responses split). model lives
9
+ // at message.model. client.messages is a plain wrappable instance
10
+ // property; patched in place, same rationale as wrapOpenAI.
11
+ //
12
+ // STREAMING (M9-07) — `create({..., stream: true})` resolves to a Stream
13
+ // of RawMessageStreamEvents. The stream's asyncIterator is patched in
14
+ // place (wrappers/streaming.ts) and the Message is reassembled from the
15
+ // event protocol as the caller iterates: message_start carries the model
16
+ // and input_tokens, content_block_start/content_block_delta carry text
17
+ // and tool_use blocks (tool input arrives as partial_json string
18
+ // fragments), and the final message_delta carries the cumulative
19
+ // output_tokens and stop_reason — so unlike OpenAI (which hides usage
20
+ // behind stream_options.include_usage), a streamed Anthropic call always
21
+ // records full usage. The `messages.stream()` helper is NOT wrapped —
22
+ // it builds its own request path; only create() calls are traced.
23
+
24
+ export interface WrapAnthropicOptions {
25
+ /** Attach generations to an existing trace instead of creating a new standalone trace per call. */
26
+ trace?: TraceHandle;
27
+ }
28
+
29
+ // Request sampling parameters worth recording as modelParameters — verified
30
+ // against the Messages API request schema (@anthropic-ai/sdk@0.111.0).
31
+ // Only fields actually present on the request are ever recorded.
32
+ interface RequestModelParameters {
33
+ temperature?: number;
34
+ top_p?: number;
35
+ top_k?: number;
36
+ max_tokens?: number;
37
+ }
38
+
39
+ function extractModelParameters(
40
+ body: RequestModelParameters | undefined
41
+ ): Record<string, string | number | boolean | null> | undefined {
42
+ if (!body) return undefined;
43
+ const entries = (["temperature", "top_p", "top_k", "max_tokens"] as const)
44
+ .filter((key) => body[key] !== undefined)
45
+ .map((key) => [key, body[key]] as [string, number]);
46
+ return entries.length > 0 ? Object.fromEntries(entries) : undefined;
47
+ }
48
+
49
+ interface MessageUsage {
50
+ input_tokens?: number;
51
+ output_tokens?: number;
52
+ }
53
+
54
+ interface MessageLike {
55
+ model?: string;
56
+ usage?: MessageUsage;
57
+ }
58
+
59
+ // `never[]` params so the real Anthropic client's overloaded create()
60
+ // satisfies the constraint — same contravariance reasoning as wrapOpenAI.
61
+ interface MessagesLike {
62
+ create: (...args: never[]) => unknown;
63
+ }
64
+
65
+ interface AnthropicLike {
66
+ messages: MessagesLike;
67
+ }
68
+
69
+ function usageDetailsFrom(usage: MessageUsage | undefined) {
70
+ if (!usage) return undefined;
71
+ const details = {
72
+ ...(usage.input_tokens !== undefined && { input_tokens: usage.input_tokens }),
73
+ ...(usage.output_tokens !== undefined && { output_tokens: usage.output_tokens })
74
+ };
75
+ return Object.keys(details).length > 0 ? details : undefined;
76
+ }
77
+
78
+ // Streamed event protocol (RawMessageStreamEvent), verified against
79
+ // @anthropic-ai/sdk@0.111.0 + docs.anthropic.com/en/docs/build-with-claude/streaming.
80
+ interface StreamEventLike {
81
+ type?: string;
82
+ index?: number;
83
+ message?: { model?: string; usage?: MessageUsage };
84
+ content_block?: {
85
+ type?: string;
86
+ id?: string;
87
+ name?: string;
88
+ text?: string;
89
+ thinking?: string;
90
+ signature?: string;
91
+ };
92
+ delta?: {
93
+ type?: string;
94
+ text?: string;
95
+ partial_json?: string;
96
+ thinking?: string;
97
+ signature?: string;
98
+ stop_reason?: string;
99
+ };
100
+ usage?: MessageUsage;
101
+ }
102
+
103
+ type AccumulatedBlock =
104
+ | { type: "text"; text: string }
105
+ | { type: "tool_use"; id?: string; name?: string; partialJson: string }
106
+ | { type: "thinking"; thinking: string; signature: string }
107
+ | { type: string; [key: string]: unknown };
108
+
109
+ function createEventAccumulator() {
110
+ const blocks: AccumulatedBlock[] = [];
111
+ let responseModel: string | undefined;
112
+ let inputTokens: number | undefined;
113
+ let outputTokens: number | undefined;
114
+ let stopReason: string | undefined;
115
+
116
+ return {
117
+ onEvent(raw: unknown) {
118
+ const event = raw as StreamEventLike;
119
+ switch (event.type) {
120
+ case "message_start":
121
+ responseModel = event.message?.model;
122
+ inputTokens = event.message?.usage?.input_tokens;
123
+ outputTokens = event.message?.usage?.output_tokens;
124
+ break;
125
+ case "content_block_start": {
126
+ const block = event.content_block;
127
+ if (event.index === undefined || !block) break;
128
+ if (block.type === "text") {
129
+ blocks[event.index] = { type: "text", text: block.text ?? "" };
130
+ } else if (block.type === "tool_use") {
131
+ blocks[event.index] = {
132
+ type: "tool_use",
133
+ ...(block.id && { id: block.id }),
134
+ ...(block.name && { name: block.name }),
135
+ partialJson: ""
136
+ };
137
+ } else if (block.type === "thinking") {
138
+ // Extended-thinking blocks stream as thinking_delta text plus a
139
+ // final signature_delta (the signature is required to round-trip
140
+ // the block in later turns — dropping it would make the recorded
141
+ // output unusable as a replay input). Review finding, PR #39.
142
+ blocks[event.index] = {
143
+ type: "thinking",
144
+ thinking: block.thinking ?? "",
145
+ signature: block.signature ?? ""
146
+ };
147
+ } else if (block.type) {
148
+ // Other block kinds (redacted_thinking arrives complete at
149
+ // start; future API additions) are kept as-is rather than
150
+ // dropped — deltas for them aren't understood, but the block's
151
+ // existence is real data.
152
+ blocks[event.index] = { ...block, type: block.type };
153
+ }
154
+ break;
155
+ }
156
+ case "content_block_delta": {
157
+ if (event.index === undefined) break;
158
+ const block = blocks[event.index];
159
+ if (!block) break;
160
+ if (event.delta?.type === "text_delta" && block.type === "text") {
161
+ (block as { text: string }).text += event.delta.text ?? "";
162
+ } else if (event.delta?.type === "input_json_delta" && block.type === "tool_use") {
163
+ (block as { partialJson: string }).partialJson += event.delta.partial_json ?? "";
164
+ } else if (event.delta?.type === "thinking_delta" && block.type === "thinking") {
165
+ (block as { thinking: string }).thinking += event.delta.thinking ?? "";
166
+ } else if (event.delta?.type === "signature_delta" && block.type === "thinking") {
167
+ (block as { signature: string }).signature = event.delta.signature ?? "";
168
+ }
169
+ break;
170
+ }
171
+ case "message_delta":
172
+ if (event.usage?.output_tokens !== undefined) outputTokens = event.usage.output_tokens;
173
+ if (event.usage?.input_tokens !== undefined) inputTokens = event.usage.input_tokens;
174
+ if (event.delta?.stop_reason) stopReason = event.delta.stop_reason;
175
+ break;
176
+ }
177
+ },
178
+ endOptions() {
179
+ const content = blocks.filter(Boolean).map((block) => {
180
+ if (block.type === "tool_use" && "partialJson" in block) {
181
+ const { partialJson, ...rest } = block;
182
+ const rawJson = typeof partialJson === "string" ? partialJson : "";
183
+ let input: unknown;
184
+ try {
185
+ input = rawJson ? JSON.parse(rawJson) : {};
186
+ } catch {
187
+ // Truncated stream (early break mid-tool-call): keep the raw
188
+ // fragment rather than losing it or throwing in a finalizer.
189
+ input = { __partial_json: rawJson };
190
+ }
191
+ return { ...rest, input };
192
+ }
193
+ return block;
194
+ });
195
+ const usage = {
196
+ ...(inputTokens !== undefined && { input_tokens: inputTokens }),
197
+ ...(outputTokens !== undefined && { output_tokens: outputTokens })
198
+ };
199
+ return {
200
+ output: {
201
+ role: "assistant",
202
+ ...(responseModel && { model: responseModel }),
203
+ content,
204
+ ...(stopReason && { stop_reason: stopReason })
205
+ },
206
+ ...(Object.keys(usage).length > 0 && { usageDetails: usage }),
207
+ metadata: { streamed: "true" }
208
+ };
209
+ }
210
+ };
211
+ }
212
+
213
+ /**
214
+ * Wraps an Anthropic client instance so every messages.create() call —
215
+ * streaming or not — is automatically recorded as a generation. Mutates
216
+ * client.messages in place and returns the same client reference.
217
+ *
218
+ * Safe to call more than once on the same client — re-wrapping is detected
219
+ * and is a no-op, rather than nesting a second wrapper around the first
220
+ * (which would silently double-record every call).
221
+ */
222
+ export function wrapAnthropic<T extends AnthropicLike>(
223
+ client: T,
224
+ ironside: IronsideClient,
225
+ options: WrapAnthropicOptions = {}
226
+ ): T {
227
+ const messages = client.messages as MessagesLike & { __ironsideWrapped?: boolean };
228
+ if (messages.__ironsideWrapped) return client;
229
+
230
+ const originalCreate = messages.create.bind(messages) as (
231
+ ...args: unknown[]
232
+ ) => Promise<unknown>;
233
+
234
+ const wrappedCreate = async (...args: unknown[]): Promise<unknown> => {
235
+ const requestBody = args[0] as
236
+ | (RequestModelParameters & { model?: string; messages?: unknown; stream?: boolean })
237
+ | undefined;
238
+
239
+ const modelParameters = extractModelParameters(requestBody);
240
+ const trace = options.trace ?? ironside.trace({ name: "anthropic.messages.create" });
241
+ const generation = trace.generation({
242
+ name: "anthropic.messages.create",
243
+ ...(requestBody?.model && { model: requestBody.model }),
244
+ ...(modelParameters && { modelParameters }),
245
+ input: requestBody?.messages
246
+ });
247
+
248
+ try {
249
+ const result = await originalCreate(...args);
250
+
251
+ if (requestBody?.stream) {
252
+ const accumulator = createEventAccumulator();
253
+ return instrumentAsyncIterable(result, accumulator.onEvent, ({ error, consumed }) => {
254
+ if (error) generation.end(errorEndOptions(error));
255
+ else if (consumed) generation.end(accumulator.endOptions());
256
+ else generation.end({ metadata: { streamed: "true" } });
257
+ });
258
+ }
259
+
260
+ const message = result as MessageLike;
261
+ const usageDetails = usageDetailsFrom(message.usage);
262
+ generation.end({
263
+ output: message,
264
+ ...(usageDetails && { usageDetails })
265
+ });
266
+ return result;
267
+ } catch (error) {
268
+ generation.end(errorEndOptions(error));
269
+ throw error;
270
+ }
271
+ };
272
+ messages.create = wrappedCreate as typeof messages.create;
273
+ Object.defineProperty(messages, "__ironsideWrapped", { value: true, enumerable: false });
274
+
275
+ return client;
276
+ }
@@ -0,0 +1,283 @@
1
+ import type { IronsideClient, TraceHandle } from "../client.js";
2
+ import { errorEndOptions, instrumentAsyncIterable } from "./streaming.js";
3
+
4
+ // Wraps the OpenAI Node SDK's chat.completions.create() to automatically
5
+ // record a generation. Verified against openai@6.46.0: usage lives at
6
+ // completion.usage.{prompt_tokens,completion_tokens} for Chat Completions
7
+ // (the Responses API — client.responses.create — uses a different field
8
+ // vocabulary, input_tokens/output_tokens, and is not wrapped here; only
9
+ // Chat Completions, still the most common call site, is covered in this
10
+ // pass). client.chat.completions is a plain wrappable instance property,
11
+ // not a sealed/frozen internal — this is the same interception pattern
12
+ // OpenInference and Langfuse's own OpenAI integrations use.
13
+ //
14
+ // Patches the passed-in client's create() method in place and returns the
15
+ // SAME reference, rather than shallow-copying the client object. A shallow
16
+ // copy (`{...client}`) would drop every other resource (embeddings,
17
+ // images, ...) and any internal state accessed via `this` inside the SDK's
18
+ // own methods — patch-in-place is what real client instrumentation
19
+ // libraries (OpenInference, Langfuse) do for exactly this reason.
20
+ //
21
+ // STREAMING (M9-07) — `create({..., stream: true})` resolves to a Stream
22
+ // (async iterable of ChatCompletionChunk). The stream's asyncIterator is
23
+ // patched in place (wrappers/streaming.ts) so text deltas, streamed tool
24
+ // calls, and the final usage chunk are accumulated as the CALLER iterates,
25
+ // and the generation ends when iteration finishes (done, early break, or
26
+ // mid-stream error). Usage is only present on the final (empty-choices)
27
+ // chunk when the caller set `stream_options: {include_usage: true}` — the
28
+ // wrapper deliberately does NOT inject that option itself: mutating the
29
+ // wire request can break OpenAI-compatible backends that reject unknown
30
+ // fields, and a tracing wrapper silently changing what's sent to the
31
+ // provider is worse than missing token counts. Without it, the streamed
32
+ // generation records output text but no usageDetails (documented in the
33
+ // README quickstart).
34
+
35
+ export interface WrapOpenAIOptions {
36
+ /** Attach generations to an existing trace instead of creating a new standalone trace per call. */
37
+ trace?: TraceHandle;
38
+ }
39
+
40
+ // Request sampling parameters worth recording as modelParameters — verified
41
+ // against the Chat Completions request schema (openai@6.46.0). Kept to the
42
+ // common cross-provider set (temperature/top_p/max_tokens/penalties/seed);
43
+ // only fields actually present on the request are ever recorded.
44
+ interface RequestModelParameters {
45
+ temperature?: number;
46
+ top_p?: number;
47
+ max_tokens?: number;
48
+ max_completion_tokens?: number;
49
+ presence_penalty?: number;
50
+ frequency_penalty?: number;
51
+ seed?: number;
52
+ }
53
+
54
+ function extractModelParameters(
55
+ body: RequestModelParameters | undefined
56
+ ): Record<string, string | number | boolean | null> | undefined {
57
+ if (!body) return undefined;
58
+ const entries = (
59
+ ["temperature", "top_p", "max_tokens", "max_completion_tokens", "presence_penalty", "frequency_penalty", "seed"] as const
60
+ )
61
+ .filter((key) => body[key] !== undefined)
62
+ .map((key) => [key, body[key]] as [string, number]);
63
+ return entries.length > 0 ? Object.fromEntries(entries) : undefined;
64
+ }
65
+
66
+ interface ChatCompletionUsage {
67
+ prompt_tokens?: number;
68
+ completion_tokens?: number;
69
+ }
70
+
71
+ interface ChatCompletionLike {
72
+ model?: string;
73
+ usage?: ChatCompletionUsage;
74
+ }
75
+
76
+ function usageDetailsFrom(usage: ChatCompletionUsage | undefined) {
77
+ if (!usage) return undefined;
78
+ const details = {
79
+ ...(usage.prompt_tokens !== undefined && { input_tokens: usage.prompt_tokens }),
80
+ ...(usage.completion_tokens !== undefined && { output_tokens: usage.completion_tokens })
81
+ };
82
+ return Object.keys(details).length > 0 ? details : undefined;
83
+ }
84
+
85
+ // Streamed chunk shape (ChatCompletionChunk): delta carries content and
86
+ // incremental tool_calls (arguments arrive as string fragments to
87
+ // concatenate, keyed by index); usage arrives on a final chunk with empty
88
+ // choices, only under stream_options.include_usage.
89
+ interface ChunkToolCallDelta {
90
+ index: number;
91
+ id?: string;
92
+ type?: string;
93
+ function?: { name?: string; arguments?: string };
94
+ }
95
+
96
+ interface ChatCompletionChunkLike {
97
+ model?: string;
98
+ usage?: ChatCompletionUsage | null;
99
+ choices?: Array<{
100
+ index?: number;
101
+ delta?: { content?: string | null; tool_calls?: ChunkToolCallDelta[] };
102
+ finish_reason?: string | null;
103
+ }>;
104
+ }
105
+
106
+ interface AccumulatedToolCall {
107
+ id?: string;
108
+ type?: string;
109
+ function: { name?: string; arguments: string };
110
+ }
111
+
112
+ interface AccumulatedChoice {
113
+ content: string;
114
+ sawContent: boolean;
115
+ toolCalls: AccumulatedToolCall[];
116
+ finishReason?: string;
117
+ }
118
+
119
+ function assembleMessage(choice: AccumulatedChoice) {
120
+ const assembledToolCalls = choice.toolCalls.filter(Boolean);
121
+ return {
122
+ role: "assistant",
123
+ content: choice.sawContent ? choice.content : null,
124
+ ...(assembledToolCalls.length > 0 && { tool_calls: assembledToolCalls })
125
+ };
126
+ }
127
+
128
+ function createChunkAccumulator() {
129
+ // Keyed by choice.index — an n>1 request streams every choice's deltas
130
+ // interleaved in the same chunk sequence, so reading only choices[0]
131
+ // would silently truncate the response to a fraction of one choice
132
+ // (review finding, PR #39).
133
+ const choices = new Map<number, AccumulatedChoice>();
134
+ let usage: ChatCompletionUsage | undefined;
135
+ let responseModel: string | undefined;
136
+
137
+ return {
138
+ onChunk(raw: unknown) {
139
+ const chunk = raw as ChatCompletionChunkLike;
140
+ if (chunk.model) responseModel = chunk.model;
141
+ if (chunk.usage) usage = chunk.usage;
142
+ for (const choice of chunk.choices ?? []) {
143
+ const index = choice.index ?? 0;
144
+ let slot = choices.get(index);
145
+ if (!slot) {
146
+ slot = { content: "", sawContent: false, toolCalls: [] };
147
+ choices.set(index, slot);
148
+ }
149
+ if (choice.finish_reason) slot.finishReason = choice.finish_reason;
150
+ if (typeof choice.delta?.content === "string") {
151
+ slot.content += choice.delta.content;
152
+ slot.sawContent = true;
153
+ }
154
+ for (const tc of choice.delta?.tool_calls ?? []) {
155
+ const toolSlot = (slot.toolCalls[tc.index] ??= { function: { arguments: "" } });
156
+ if (tc.id) toolSlot.id = tc.id;
157
+ if (tc.type) toolSlot.type = tc.type;
158
+ if (tc.function?.name) toolSlot.function.name = tc.function.name;
159
+ if (tc.function?.arguments) toolSlot.function.arguments += tc.function.arguments;
160
+ }
161
+ }
162
+ },
163
+ endOptions() {
164
+ const sorted = [...choices.entries()].sort(([a], [b]) => a - b);
165
+ const usageDetails = usageDetailsFrom(usage);
166
+ // Single choice (the overwhelmingly common case): output is the
167
+ // assistant message itself. n>1: output mirrors the non-streaming
168
+ // completion's choices array so no choice is dropped.
169
+ const single =
170
+ sorted.length === 0
171
+ ? { content: "", sawContent: false, toolCalls: [] } // empty stream: a bare content-null message
172
+ : sorted.length === 1
173
+ ? sorted[0]![1]
174
+ : undefined;
175
+ const output = single
176
+ ? assembleMessage(single)
177
+ : {
178
+ choices: sorted.map(([index, choice]) => ({
179
+ index,
180
+ message: assembleMessage(choice),
181
+ ...(choice.finishReason && { finish_reason: choice.finishReason })
182
+ }))
183
+ };
184
+ const finishReason = single?.finishReason;
185
+ return {
186
+ output,
187
+ ...(usageDetails && { usageDetails }),
188
+ metadata: {
189
+ streamed: "true",
190
+ ...(responseModel && { response_model: responseModel }),
191
+ ...(finishReason && { finish_reason: finishReason })
192
+ }
193
+ };
194
+ }
195
+ };
196
+ }
197
+
198
+ // `never[]` parameters make this satisfiable by the real OpenAI client's
199
+ // overloaded create() — a `(...args: unknown[])` signature is
200
+ // contravariantly INCOMPATIBLE with any function taking typed params, so
201
+ // the real SDK class would fail the T constraint (caught by the real-SDK
202
+ // conformance tests, which the fake-client unit tests never exercised).
203
+ interface ChatCompletionsLike {
204
+ create: (...args: never[]) => unknown;
205
+ }
206
+
207
+ interface OpenAILike {
208
+ chat: { completions: ChatCompletionsLike };
209
+ }
210
+
211
+ /**
212
+ * Wraps an OpenAI client instance so every chat.completions.create() call
213
+ * — streaming or not — is automatically recorded as a generation. Mutates
214
+ * client.chat.completions in place and returns the same client reference
215
+ * for convenient chaining (`const client = wrapOpenAI(new OpenAI(...), ironside)`).
216
+ *
217
+ * Safe to call more than once on the same client — re-wrapping is detected
218
+ * and is a no-op, rather than nesting a second wrapper around the first
219
+ * (which would silently double-record every call: two traces, two
220
+ * generations, per real API call).
221
+ */
222
+ export function wrapOpenAI<T extends OpenAILike>(
223
+ client: T,
224
+ ironside: IronsideClient,
225
+ options: WrapOpenAIOptions = {}
226
+ ): T {
227
+ const completions = client.chat.completions as ChatCompletionsLike & {
228
+ __ironsideWrapped?: boolean;
229
+ };
230
+ if (completions.__ironsideWrapped) return client;
231
+
232
+ const originalCreate = completions.create.bind(completions) as (
233
+ ...args: unknown[]
234
+ ) => Promise<unknown>;
235
+
236
+ const wrappedCreate = async (...args: unknown[]): Promise<unknown> => {
237
+ const requestBody = args[0] as
238
+ | (RequestModelParameters & { model?: string; messages?: unknown; stream?: boolean })
239
+ | undefined;
240
+
241
+ const modelParameters = extractModelParameters(requestBody);
242
+ const trace = options.trace ?? ironside.trace({ name: "openai.chat.completions.create" });
243
+ const generation = trace.generation({
244
+ name: "openai.chat.completions.create",
245
+ ...(requestBody?.model && { model: requestBody.model }),
246
+ ...(modelParameters && { modelParameters }),
247
+ input: requestBody?.messages
248
+ });
249
+
250
+ try {
251
+ const result = await originalCreate(...args);
252
+
253
+ if (requestBody?.stream) {
254
+ const accumulator = createChunkAccumulator();
255
+ return instrumentAsyncIterable(result, accumulator.onChunk, ({ error, consumed }) => {
256
+ if (error) generation.end(errorEndOptions(error));
257
+ else if (consumed) generation.end(accumulator.endOptions());
258
+ // !consumed: the result wasn't iterable at all (unexpected SDK
259
+ // shape) — end with what we know rather than dangle forever.
260
+ else generation.end({ metadata: { streamed: "true" } });
261
+ });
262
+ }
263
+
264
+ const completion = result as ChatCompletionLike;
265
+ const usageDetails = usageDetailsFrom(completion.usage);
266
+ generation.end({
267
+ output: completion,
268
+ ...(usageDetails && { usageDetails })
269
+ });
270
+ return result;
271
+ } catch (error) {
272
+ generation.end(errorEndOptions(error));
273
+ throw error;
274
+ }
275
+ };
276
+ completions.create = wrappedCreate as typeof completions.create;
277
+ Object.defineProperty(completions, "__ironsideWrapped", {
278
+ value: true,
279
+ enumerable: false
280
+ });
281
+
282
+ return client;
283
+ }
@@ -0,0 +1,92 @@
1
+ // Shared machinery for tracing streamed responses (M9-07). Both provider
2
+ // SDKs return a Stream object that is an async iterable with extra API
3
+ // surface (.tee(), .controller, .toReadableStream(), ...). Returning our
4
+ // own wrapper generator would silently break every caller using that
5
+ // surface, so instead the stream's [Symbol.asyncIterator] is patched IN
6
+ // PLACE and the SAME object is returned — identical to how the wrappers
7
+ // patch client.chat.completions.create itself.
8
+ //
9
+ // The generation can only be finalized when the caller actually consumes
10
+ // the stream (that's when the text/usage exists at all). Three exits all
11
+ // funnel into one finalize call: normal completion (done), early break
12
+ // (the iterator's return()), and a mid-stream error (next() rejecting or
13
+ // throw()). A stream the caller never iterates records its start event
14
+ // but never ends — visible in the UI as a dangling in-progress
15
+ // generation, which is the honest representation of what happened.
16
+ //
17
+ // KNOWN LIMIT — .tee(): each branch iterates through the patched
18
+ // asyncIterator, so chunks from both branches feed one accumulator
19
+ // (double-counted text). finalize still fires exactly once. tee'd +
20
+ // wrapped is rare enough that correct-single-stream beats a per-iterator
21
+ // accumulator design that couldn't merge usage sanely anyway.
22
+
23
+ /**
24
+ * Patches `stream`'s async iterator in place so every yielded chunk feeds
25
+ * `onChunk` and exactly one of done/break/error triggers `onFinish`.
26
+ * Returns the same object. If the value isn't async-iterable at all
27
+ * (unexpected SDK shape), it is returned untouched and `onFinish` fires
28
+ * immediately with `{ consumed: false }` so the generation isn't left
29
+ * dangling by our own bug.
30
+ */
31
+ export function instrumentAsyncIterable<T>(
32
+ stream: T,
33
+ onChunk: (chunk: unknown) => void,
34
+ onFinish: (outcome: { error?: unknown; consumed: boolean }) => void
35
+ ): T {
36
+ const iterable = stream as T & {
37
+ [Symbol.asyncIterator]?: () => AsyncIterator<unknown>;
38
+ };
39
+ if (typeof iterable?.[Symbol.asyncIterator] !== "function") {
40
+ onFinish({ consumed: false });
41
+ return stream;
42
+ }
43
+
44
+ let finished = false;
45
+ const finishOnce = (outcome: { error?: unknown }) => {
46
+ if (finished) return;
47
+ finished = true;
48
+ onFinish({ ...outcome, consumed: true });
49
+ };
50
+
51
+ const originalFactory = iterable[Symbol.asyncIterator]!.bind(iterable);
52
+ iterable[Symbol.asyncIterator] = () => {
53
+ const inner = originalFactory();
54
+ return {
55
+ async next(): Promise<IteratorResult<unknown>> {
56
+ try {
57
+ const result = await inner.next();
58
+ if (result.done) finishOnce({});
59
+ else onChunk(result.value);
60
+ return result;
61
+ } catch (error) {
62
+ finishOnce({ error });
63
+ throw error;
64
+ }
65
+ },
66
+ // Called on `break`/`return` inside a for-await — the caller chose
67
+ // to stop early; what accumulated so far is the real output.
68
+ async return(value?: unknown): Promise<IteratorResult<unknown>> {
69
+ finishOnce({});
70
+ if (inner.return) return inner.return(value);
71
+ return { done: true, value: undefined };
72
+ },
73
+ async throw(error?: unknown): Promise<IteratorResult<unknown>> {
74
+ finishOnce({ error });
75
+ if (inner.throw) return inner.throw(error);
76
+ throw error;
77
+ }
78
+ };
79
+ };
80
+ return stream;
81
+ }
82
+
83
+ /** Shared error→end mapping so streamed and non-streamed failures record identically. */
84
+ export function errorEndOptions(error: unknown): {
85
+ level: "error";
86
+ statusMessage: string;
87
+ } {
88
+ return {
89
+ level: "error",
90
+ statusMessage: error instanceof Error ? error.message : String(error)
91
+ };
92
+ }