@@ -0,0 +1,167 @@
1+import type { StreamFn } from "@mariozechner/pi-agent-core";
2+import { beforeEach, describe, expect, it } from "vitest";
3+import {
4+onDiagnosticEvent,
5+resetDiagnosticEventsForTest,
6+type DiagnosticEventPayload,
7+} from "../../../infra/diagnostic-events.js";
8+import { createDiagnosticTraceContext } from "../../../infra/diagnostic-trace-context.js";
9+import { wrapStreamFnWithDiagnosticModelCallEvents } from "./attempt.model-diagnostic-events.js";
10+11+async function collectModelCallEvents(run: () => Promise<void>): Promise<DiagnosticEventPayload[]> {
12+const events: DiagnosticEventPayload[] = [];
13+const stop = onDiagnosticEvent((event) => {
14+if (event.type.startsWith("model.call.")) {
15+events.push(event);
16+}
17+});
18+try {
19+await run();
20+return events;
21+} finally {
22+stop();
23+}
24+}
25+26+async function drain(stream: AsyncIterable<unknown>): Promise<void> {
27+for await (const _ of stream) {
28+// drain
29+}
30+}
31+32+describe("wrapStreamFnWithDiagnosticModelCallEvents", () => {
33+beforeEach(() => {
34+resetDiagnosticEventsForTest();
35+});
36+37+it("emits started and completed events for async streams", async () => {
38+async function* stream() {
39+yield { type: "text", text: "ok" };
40+}
41+const originalStream = stream() as unknown as AsyncIterable<unknown> & {
42+result: () => Promise<string>;
43+};
44+originalStream.result = async () => "kept";
45+const wrapped = wrapStreamFnWithDiagnosticModelCallEvents(
46+(() => originalStream) as unknown as StreamFn,
47+{
48+runId: "run-1",
49+sessionKey: "session-key",
50+sessionId: "session-id",
51+provider: "openai",
52+model: "gpt-5.4",
53+api: "openai-responses",
54+transport: "http",
55+trace: createDiagnosticTraceContext({
56+traceId: "4bf92f3577b34da6a3ce929d0e0e4736",
57+spanId: "00f067aa0ba902b7",
58+}),
59+nextCallId: () => "call-1",
60+},
61+);
62+63+const events = await collectModelCallEvents(async () => {
64+const returned = wrapped(
65+{} as never,
66+{} as never,
67+{} as never,
68+) as unknown as typeof originalStream;
69+expect(returned).toBe(originalStream);
70+expect(await returned.result()).toBe("kept");
71+await drain(returned);
72+});
73+74+expect(events.map((event) => event.type)).toEqual([
75+"model.call.started",
76+"model.call.completed",
77+]);
78+expect(events[0]).toMatchObject({
79+type: "model.call.started",
80+runId: "run-1",
81+callId: "call-1",
82+sessionKey: "session-key",
83+sessionId: "session-id",
84+provider: "openai",
85+model: "gpt-5.4",
86+api: "openai-responses",
87+transport: "http",
88+});
89+expect(events[0]?.trace?.parentSpanId).toBe("00f067aa0ba902b7");
90+expect(events[1]).toMatchObject({
91+type: "model.call.completed",
92+callId: "call-1",
93+durationMs: expect.any(Number),
94+});
95+});
96+97+it("emits error events when stream iteration fails", async () => {
98+const stream = {
99+[Symbol.asyncIterator]() {
100+return {
101+async next(): Promise<IteratorResult<unknown>> {
102+throw new TypeError("provider failed");
103+},
104+};
105+},
106+};
107+const wrapped = wrapStreamFnWithDiagnosticModelCallEvents(
108+(() => stream) as unknown as StreamFn,
109+{
110+runId: "run-1",
111+provider: "anthropic",
112+model: "sonnet-4.6",
113+trace: createDiagnosticTraceContext(),
114+nextCallId: () => "call-err",
115+},
116+);
117+118+const events = await collectModelCallEvents(async () => {
119+await expect(
120+drain(wrapped({} as never, {} as never, {} as never) as AsyncIterable<unknown>),
121+).rejects.toThrow("provider failed");
122+});
123+124+expect(events.map((event) => event.type)).toEqual(["model.call.started", "model.call.error"]);
125+expect(events[1]).toMatchObject({
126+type: "model.call.error",
127+callId: "call-err",
128+errorCategory: "TypeError",
129+durationMs: expect.any(Number),
130+});
131+});
132+133+it("emits error events when stream consumption stops early", async () => {
134+async function* stream() {
135+yield { type: "text", text: "first" };
136+yield { type: "text", text: "second" };
137+}
138+const wrapped = wrapStreamFnWithDiagnosticModelCallEvents(
139+(() => stream()) as unknown as StreamFn,
140+{
141+runId: "run-1",
142+provider: "openai",
143+model: "gpt-5.4",
144+trace: createDiagnosticTraceContext(),
145+nextCallId: () => "call-abandoned",
146+},
147+);
148+149+const events = await collectModelCallEvents(async () => {
150+for await (const _ of wrapped(
151+{} as never,
152+{} as never,
153+{} as never,
154+) as AsyncIterable<unknown>) {
155+break;
156+}
157+});
158+159+expect(events.map((event) => event.type)).toEqual(["model.call.started", "model.call.error"]);
160+expect(events[1]).toMatchObject({
161+type: "model.call.error",
162+callId: "call-abandoned",
163+errorCategory: "StreamAbandoned",
164+durationMs: expect.any(Number),
165+});
166+});
167+});
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。