test(agent): replaced agent stream mocks with createMockModel responses
- Replaced custom MockAssistantStream helpers with createMockModel streams across agent tests. - Removed manual queueMicrotask stream-event scripting in favor of scripted mock responses. - Consolidated helper fixtures by deleting local aliases and reusing shared user-message/model helpers. - Updated test assertions to use mock.calls and mock.model metadata for call and context validation.
This commit is contained in:
can1357
@@ -9,36 +9,11 @@ import type {
|
||||
AgentToolContext,
|
||||
ToolCallContext,
|
||||
} from "@oh-my-pi/pi-agent-core/types";
|
||||
import type { AssistantMessage, Context, Message, Model, ToolResultMessage, UserMessage } from "@oh-my-pi/pi-ai";
|
||||
import type { AssistantMessage, Message, ToolResultMessage } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { createAssistantMessage, pushAlphaThenDoneEvent } from "./helpers";
|
||||
|
||||
// Mock stream for testing - uses actual AssistantMessageEventStream with throttling
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
function createModel(): Model<"openai-responses"> {
|
||||
return {
|
||||
id: "mock",
|
||||
name: "mock",
|
||||
api: "openai-responses",
|
||||
provider: "openai",
|
||||
baseUrl: "https://example.invalid",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 8192,
|
||||
maxTokens: 2048,
|
||||
};
|
||||
}
|
||||
|
||||
function createUserMessage(text: string): UserMessage {
|
||||
return {
|
||||
role: "user",
|
||||
content: text,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
import { createAssistantMessage, createUserMessage } from "./helpers";
|
||||
|
||||
// Simple identity converter for tests - just passes through standard messages
|
||||
function identityConverter(messages: AgentMessage[]): Message[] {
|
||||
@@ -53,24 +28,11 @@ describe("agentLoop with AgentMessage", () => {
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("Hello");
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
};
|
||||
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage([{ type: "text", text: "Hi there!" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
const mock = createMockModel({ responses: [{ content: ["Hi there!"] }] });
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter };
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
const stream = agentLoop([createUserMessage("Hello")], context, config, undefined, mock.stream);
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
@@ -99,16 +61,16 @@ describe("agentLoop with AgentMessage", () => {
|
||||
messages: [],
|
||||
tools: [],
|
||||
};
|
||||
const userPrompt: AgentMessage = createUserMessage("Hello");
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
};
|
||||
const mock = createMockModel();
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter };
|
||||
const controller = new AbortController();
|
||||
const streamFn = () => new MockAssistantStream();
|
||||
// The mock provider would reject without a configured response; we want the
|
||||
// agent's abort path to kick in before any event is emitted. Use a raw stream
|
||||
// that never emits anything.
|
||||
const streamFn = () => new AssistantMessageEventStream();
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, controller.signal, streamFn);
|
||||
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn);
|
||||
queueMicrotask(() => controller.abort());
|
||||
|
||||
for await (const event of stream) {
|
||||
@@ -144,11 +106,10 @@ describe("agentLoop with AgentMessage", () => {
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("Hello");
|
||||
|
||||
let convertedMessages: Message[] = [];
|
||||
const mock = createMockModel({ responses: [{ content: ["Response"] }] });
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: messages => {
|
||||
// Filter out notifications, convert rest
|
||||
convertedMessages = messages
|
||||
@@ -158,20 +119,9 @@ describe("agentLoop with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage([{ type: "text", text: "Response" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
const stream = agentLoop([createUserMessage("Hello")], context, config, undefined, mock.stream);
|
||||
for await (const _ of stream) {
|
||||
// drain
|
||||
}
|
||||
|
||||
// The notification should have been filtered out in convertToLlm
|
||||
@@ -191,13 +141,12 @@ describe("agentLoop with AgentMessage", () => {
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("new message");
|
||||
|
||||
let transformedMessages: AgentMessage[] = [];
|
||||
let convertedMessages: Message[] = [];
|
||||
|
||||
const mock = createMockModel({ responses: [{ content: ["Response"] }] });
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
transformContext: async messages => {
|
||||
// Keep only last 2 messages (prune old ones)
|
||||
transformedMessages = messages.slice(-2);
|
||||
@@ -211,19 +160,9 @@ describe("agentLoop with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage([{ type: "text", text: "Response" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
|
||||
const stream = agentLoop([createUserMessage("new message")], context, config, undefined, mock.stream);
|
||||
for await (const _ of stream) {
|
||||
// consume
|
||||
// drain
|
||||
}
|
||||
|
||||
// transformContext should have been called first, keeping only last 2
|
||||
@@ -252,46 +191,28 @@ describe("agentLoop with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: [""],
|
||||
messages: [],
|
||||
tools: [tool],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("echo something");
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [
|
||||
{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } },
|
||||
{ type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "world" } },
|
||||
],
|
||||
},
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
getToolContext: toolCall => ({ toolCall }) as AgentToolContext,
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const message = createAssistantMessage(
|
||||
[
|
||||
{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } },
|
||||
{ type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "world" } },
|
||||
],
|
||||
"toolUse",
|
||||
);
|
||||
stream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
|
||||
const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, mock.stream);
|
||||
for await (const _ of stream) {
|
||||
// consume
|
||||
// drain
|
||||
}
|
||||
|
||||
expect(contexts).toHaveLength(2);
|
||||
@@ -322,42 +243,18 @@ describe("agentLoop with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: [""],
|
||||
messages: [],
|
||||
tools: [tool],
|
||||
};
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("echo something");
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
// First call: return tool call
|
||||
const message = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }],
|
||||
"toolUse",
|
||||
);
|
||||
stream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
// Second call: return final response
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter };
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, mock.stream);
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
@@ -379,7 +276,6 @@ describe("agentLoop with AgentMessage", () => {
|
||||
it("injects and strips intent when intent tracing is enabled", async () => {
|
||||
const toolSchema = Type.Object({ value: Type.String() });
|
||||
const executedParams: Record<string, unknown>[] = [];
|
||||
let firstRequestToolSchema: Record<string, unknown> | undefined;
|
||||
const tool: AgentTool<typeof toolSchema, { value: string }> = {
|
||||
name: "echo",
|
||||
label: "Echo",
|
||||
@@ -394,50 +290,32 @@ describe("agentLoop with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: [""],
|
||||
messages: [],
|
||||
tools: [tool],
|
||||
};
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [
|
||||
{
|
||||
type: "toolCall",
|
||||
id: "tool-1",
|
||||
name: "echo",
|
||||
arguments: { value: "hello", [INTENT_FIELD]: "Read one file" },
|
||||
},
|
||||
],
|
||||
},
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
intentTracing: true,
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const streamFn = (_model: Model, llmContext: Context) => {
|
||||
if (callIndex === 0) {
|
||||
firstRequestToolSchema = llmContext.tools?.[0]?.parameters;
|
||||
}
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const message = createAssistantMessage(
|
||||
[
|
||||
{
|
||||
type: "toolCall",
|
||||
id: "tool-1",
|
||||
name: "echo",
|
||||
arguments: { value: "hello", [INTENT_FIELD]: "Read one file" },
|
||||
},
|
||||
],
|
||||
"toolUse",
|
||||
);
|
||||
stream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const stream = agentLoop([createUserMessage("run")], context, config, undefined, streamFn);
|
||||
const stream = agentLoop([createUserMessage("run")], context, config, undefined, mock.stream);
|
||||
for await (const _ of stream) {
|
||||
// consume
|
||||
// drain
|
||||
}
|
||||
const messages = await stream.result();
|
||||
const assistantWithToolCall = messages.find(
|
||||
@@ -445,6 +323,7 @@ describe("agentLoop with AgentMessage", () => {
|
||||
) as AssistantMessage | undefined;
|
||||
const tracedToolCall = assistantWithToolCall?.content.find(content => content.type === "toolCall");
|
||||
|
||||
const firstRequestToolSchema = mock.calls[0]?.context.tools?.[0]?.parameters;
|
||||
expect(firstRequestToolSchema?.properties).toMatchObject({
|
||||
value: { type: "string" },
|
||||
[INTENT_FIELD]: { type: "string" },
|
||||
@@ -489,43 +368,23 @@ describe("agentLoop with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: [""],
|
||||
messages: [],
|
||||
tools: [tool],
|
||||
};
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("start");
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const message = createAssistantMessage(
|
||||
[
|
||||
{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "slow" } },
|
||||
{ type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "fast" } },
|
||||
],
|
||||
"toolUse",
|
||||
);
|
||||
stream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [
|
||||
{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "slow" } },
|
||||
{ type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "fast" } },
|
||||
],
|
||||
},
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter };
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
const stream = agentLoop([createUserMessage("start")], context, config, undefined, mock.stream);
|
||||
const streamTask = (async () => {
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
@@ -564,15 +423,14 @@ describe("agentLoop with AgentMessage", () => {
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("start");
|
||||
const abortController = new AbortController();
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
};
|
||||
const mock = createMockModel();
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter };
|
||||
|
||||
// Custom stream: emit a partial start with a tool call, abort, then push done.
|
||||
// The mock provider doesn't model "abort between start and done"; do it inline.
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
queueMicrotask(() => {
|
||||
const partial = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tool-1", name: "yield", arguments: { data: { ok: true } } }],
|
||||
@@ -588,7 +446,7 @@ describe("agentLoop with AgentMessage", () => {
|
||||
};
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, abortController.signal, streamFn);
|
||||
const stream = agentLoop([createUserMessage("start")], context, config, abortController.signal, streamFn);
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
@@ -608,6 +466,7 @@ describe("agentLoop with AgentMessage", () => {
|
||||
expect(text).not.toContain("Tool execution was aborted.:");
|
||||
}
|
||||
});
|
||||
|
||||
it("should skip remaining tool calls when steering is queued", async () => {
|
||||
const toolSchema = Type.Object({ value: Type.String() });
|
||||
const executed: string[] = [];
|
||||
@@ -626,21 +485,25 @@ describe("agentLoop with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: [""],
|
||||
messages: [],
|
||||
tools: [tool],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("start");
|
||||
const queuedUserMessage: AgentMessage = createUserMessage("interrupt");
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
|
||||
const queuedUserMessage = createUserMessage("interrupt");
|
||||
let queuedDelivered = false;
|
||||
let callIndex = 0;
|
||||
let sawInterruptInContext = false;
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [
|
||||
{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "first" } },
|
||||
{ type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "second" } },
|
||||
],
|
||||
},
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
interruptMode: "immediate",
|
||||
getSteeringMessages: async () => {
|
||||
@@ -654,36 +517,7 @@ describe("agentLoop with AgentMessage", () => {
|
||||
};
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, (_model, ctx, _options) => {
|
||||
// Check if interrupt message is in context on second call
|
||||
if (callIndex === 1) {
|
||||
sawInterruptInContext = ctx.messages.some(
|
||||
m => m.role === "user" && typeof m.content === "string" && m.content === "interrupt",
|
||||
);
|
||||
}
|
||||
|
||||
const mockStream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
// First call: return two tool calls
|
||||
const message = createAssistantMessage(
|
||||
[
|
||||
{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "first" } },
|
||||
{ type: "toolCall", id: "tool-2", name: "echo", arguments: { value: "second" } },
|
||||
],
|
||||
"toolUse",
|
||||
);
|
||||
mockStream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
// Second call: return final response
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
mockStream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return mockStream;
|
||||
});
|
||||
|
||||
const stream = agentLoop([createUserMessage("start")], context, config, undefined, mock.stream);
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
@@ -715,6 +549,9 @@ describe("agentLoop with AgentMessage", () => {
|
||||
expect(eventSequence.indexOf("tool:tool-2")).toBeLessThan(eventSequence.indexOf("interrupt"));
|
||||
|
||||
// Interrupt message should be in context when second LLM call is made
|
||||
const sawInterruptInContext = mock.calls[1]?.context.messages.some(
|
||||
m => m.role === "user" && typeof m.content === "string" && m.content === "interrupt",
|
||||
);
|
||||
expect(sawInterruptInContext).toBe(true);
|
||||
});
|
||||
});
|
||||
@@ -756,9 +593,14 @@ it("refreshes tools and system prompt between same-turn model calls", async () =
|
||||
messages: [],
|
||||
tools: activeTools,
|
||||
};
|
||||
const userPrompt: AgentMessage = createUserMessage("refresh tools");
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tool-1", name: "alpha", arguments: { value: "hello" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
syncContextBeforeModelCall: async currentContext => {
|
||||
currentContext.systemPrompt = [activeSystemPrompt];
|
||||
@@ -766,28 +608,16 @@ it("refreshes tools and system prompt between same-turn model calls", async () =
|
||||
},
|
||||
};
|
||||
|
||||
const callContexts: Context[] = [];
|
||||
let callIndex = 0;
|
||||
const streamFn = (_model: Model, llmContext: Context) => {
|
||||
callContexts.push(llmContext);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
pushAlphaThenDoneEvent(stream, callIndex, createAssistantMessage);
|
||||
callIndex += 1;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
const stream = agentLoop([createUserMessage("refresh tools")], context, config, undefined, mock.stream);
|
||||
for await (const _ of stream) {
|
||||
// consume
|
||||
// drain
|
||||
}
|
||||
|
||||
expect(callContexts).toHaveLength(2);
|
||||
expect(callContexts[0]?.systemPrompt).toEqual(["prompt-one"]);
|
||||
expect(callContexts[0]?.tools?.map(tool => tool.name)).toEqual(["alpha"]);
|
||||
expect(callContexts[1]?.systemPrompt).toEqual(["prompt-two"]);
|
||||
expect(callContexts[1]?.tools?.map(tool => tool.name)).toEqual(["alpha", "beta"]);
|
||||
expect(mock.calls).toHaveLength(2);
|
||||
expect(mock.calls[0]?.context.systemPrompt).toEqual(["prompt-one"]);
|
||||
expect(mock.calls[0]?.context.tools?.map(tool => tool.name)).toEqual(["alpha"]);
|
||||
expect(mock.calls[1]?.context.systemPrompt).toEqual(["prompt-two"]);
|
||||
expect(mock.calls[1]?.context.tools?.map(tool => tool.name)).toEqual(["alpha", "beta"]);
|
||||
});
|
||||
|
||||
describe("agentLoopContinue with AgentMessage", () => {
|
||||
@@ -798,16 +628,14 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
};
|
||||
const mock = createMockModel();
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter };
|
||||
|
||||
expect(() => agentLoopContinue(context, config)).toThrow("Cannot continue: no messages in context");
|
||||
});
|
||||
|
||||
it("should continue from existing context without emitting user message events", async () => {
|
||||
const userMessage: AgentMessage = createUserMessage("Hello");
|
||||
const userMessage = createUserMessage("Hello");
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: ["You are helpful."],
|
||||
@@ -815,22 +643,11 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
};
|
||||
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage([{ type: "text", text: "Response" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
const mock = createMockModel({ responses: [{ content: ["Response"] }] });
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter };
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoopContinue(context, config, undefined, streamFn);
|
||||
const stream = agentLoopContinue(context, config, undefined, mock.stream);
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
@@ -845,7 +662,9 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
// Should NOT have user message events (that's the key difference from agentLoop)
|
||||
const messageEndEvents = events.filter(e => e.type === "message_end");
|
||||
expect(messageEndEvents.length).toBe(1);
|
||||
expect((messageEndEvents[0] as any).message.role).toBe("assistant");
|
||||
const firstEnd = messageEndEvents[0];
|
||||
if (firstEnd?.type !== "message_end") throw new Error("Expected message_end");
|
||||
expect(firstEnd.message.role).toBe("assistant");
|
||||
});
|
||||
|
||||
it("should allow custom message types as last message (caller responsibility)", async () => {
|
||||
@@ -868,17 +687,19 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const mock = createMockModel({ responses: [{ content: ["Response to hook"] }] });
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: messages => {
|
||||
// Convert hookMessage to user message
|
||||
return messages
|
||||
.map(m => {
|
||||
if ((m as any).role === "hookMessage") {
|
||||
const candidate = m as unknown as Partial<HookMessage>;
|
||||
if (candidate.role === "hookMessage") {
|
||||
return {
|
||||
role: "user" as const,
|
||||
content: (m as any).text,
|
||||
timestamp: m.timestamp,
|
||||
content: candidate.text ?? "",
|
||||
timestamp: candidate.timestamp ?? Date.now(),
|
||||
};
|
||||
}
|
||||
return m;
|
||||
@@ -887,17 +708,8 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage([{ type: "text", text: "Response to hook" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
// Should not throw - the hookMessage will be converted to user message
|
||||
const stream = agentLoopContinue(context, config, undefined, streamFn);
|
||||
const stream = agentLoopContinue(context, config, undefined, mock.stream);
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
for await (const event of stream) {
|
||||
@@ -908,6 +720,7 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
expect(messages.length).toBe(1);
|
||||
expect(messages[0].role).toBe("assistant");
|
||||
});
|
||||
|
||||
it("blocks tool execution when beforeToolCall returns block", async () => {
|
||||
const toolSchema = Type.Object({ value: Type.String() });
|
||||
const executed: string[] = [];
|
||||
@@ -926,35 +739,21 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
};
|
||||
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
const userPrompt: AgentMessage = createUserMessage("echo something");
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
beforeToolCall: async () => ({ block: true, reason: "policy: blocked" }),
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const message = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }],
|
||||
"toolUse",
|
||||
);
|
||||
stream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, mock.stream);
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
@@ -986,10 +785,15 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
};
|
||||
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
const userPrompt: AgentMessage = createUserMessage("echo something");
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
beforeToolCall: async ({ args }) => {
|
||||
(args as { value: string | number }).value = 123;
|
||||
@@ -997,27 +801,8 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const message = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }],
|
||||
"toolUse",
|
||||
);
|
||||
stream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
for await (const _event of stream) {
|
||||
const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, mock.stream);
|
||||
for await (const _ of stream) {
|
||||
// drain
|
||||
}
|
||||
|
||||
@@ -1040,11 +825,16 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
};
|
||||
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
const userPrompt: AgentMessage = createUserMessage("echo something");
|
||||
|
||||
const seen: Array<{ args: unknown; isError: boolean }> = [];
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
afterToolCall: async ({ args, isError }) => {
|
||||
seen.push({ args, isError });
|
||||
@@ -1055,27 +845,8 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
},
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const message = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }],
|
||||
"toolUse",
|
||||
);
|
||||
stream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, mock.stream);
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
@@ -1093,8 +864,8 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
|
||||
const toolResultMessage = events
|
||||
.filter(e => e.type === "message_start")
|
||||
.map(e => (e as { message: AgentMessage }).message)
|
||||
.find(m => m.role === "toolResult");
|
||||
.map(e => (e.type === "message_start" ? e.message : undefined))
|
||||
.find((m): m is AgentMessage => m !== undefined && m.role === "toolResult");
|
||||
expect(toolResultMessage).toBeDefined();
|
||||
if (toolResultMessage && toolResultMessage.role === "toolResult") {
|
||||
expect(toolResultMessage.isError).toBe(true);
|
||||
@@ -1118,37 +889,23 @@ describe("agentLoopContinue with AgentMessage", () => {
|
||||
};
|
||||
|
||||
const context: AgentContext = { systemPrompt: [""], messages: [], tools: [tool] };
|
||||
const userPrompt: AgentMessage = createUserMessage("echo");
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
afterToolCall: async () => {
|
||||
throw new Error("hook exploded");
|
||||
},
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const message = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }],
|
||||
"toolUse",
|
||||
);
|
||||
stream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
const stream = agentLoop([createUserMessage("echo")], context, config, undefined, mock.stream);
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
@@ -1,11 +1,9 @@
|
||||
import { describe, expect, it } from "bun:test";
|
||||
import { Agent, type AgentTool, ThinkingLevel } from "@oh-my-pi/pi-agent-core";
|
||||
import { getBundledModel, type SimpleStreamOptions } from "@oh-my-pi/pi-ai";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import type { SimpleStreamOptions } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { createAssistantMessage, pushAlphaThenDoneEvent } from "./helpers";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
import { createAssistantMessage } from "./helpers";
|
||||
|
||||
describe("Agent", () => {
|
||||
it("should support steering message queueing", async () => {
|
||||
@@ -19,19 +17,8 @@ describe("Agent", () => {
|
||||
});
|
||||
|
||||
it("continue() should process queued follow-up messages after an assistant turn", async () => {
|
||||
const agent = new Agent({
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({
|
||||
type: "done",
|
||||
reason: "stop",
|
||||
message: createAssistantMessage([{ type: "text", text: "Processed" }]),
|
||||
});
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
const mock = createMockModel({ responses: [{ content: ["Processed"] }] });
|
||||
const agent = new Agent({ streamFn: mock.stream });
|
||||
|
||||
agent.replaceMessages([
|
||||
{
|
||||
@@ -61,21 +48,10 @@ describe("Agent", () => {
|
||||
});
|
||||
|
||||
it("continue() should keep one-at-a-time steering semantics from assistant tail", async () => {
|
||||
let responseCount = 0;
|
||||
const agent = new Agent({
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
responseCount++;
|
||||
queueMicrotask(() => {
|
||||
stream.push({
|
||||
type: "done",
|
||||
reason: "stop",
|
||||
message: createAssistantMessage([{ type: "text", text: `Processed ${responseCount}` }]),
|
||||
});
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
const mock = createMockModel({
|
||||
responses: [{ content: ["Processed 1"] }, { content: ["Processed 2"] }],
|
||||
});
|
||||
const agent = new Agent({ streamFn: mock.stream });
|
||||
|
||||
agent.replaceMessages([
|
||||
{
|
||||
@@ -101,14 +77,12 @@ describe("Agent", () => {
|
||||
|
||||
const recentMessages = agent.state.messages.slice(-4);
|
||||
expect(recentMessages.map(m => m.role)).toEqual(["user", "assistant", "user", "assistant"]);
|
||||
expect(responseCount).toBe(2);
|
||||
expect(mock.calls.length).toBe(2);
|
||||
});
|
||||
|
||||
it("prompt() refreshes tools and system prompt between same-turn model calls", async () => {
|
||||
const toolSchema = Type.Object({ value: Type.String() });
|
||||
type Details = { value: string };
|
||||
let callIndex = 0;
|
||||
const callContexts: Array<{ systemPrompt: string; toolNames: string[] }> = [];
|
||||
|
||||
const betaTool: AgentTool<typeof toolSchema, Details> = {
|
||||
name: "beta",
|
||||
@@ -129,25 +103,21 @@ describe("Agent", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tool-1", name: "alpha", arguments: { value: "hello" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
model: getBundledModel("openai", "gpt-4o-mini"),
|
||||
model: mock.model,
|
||||
systemPrompt: ["prompt-one"],
|
||||
tools: [alphaTool],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: (_model, context) => {
|
||||
callContexts.push({
|
||||
systemPrompt: context.systemPrompt?.join("\n\n") ?? "",
|
||||
toolNames: (context.tools ?? []).map(tool => tool.name),
|
||||
});
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
pushAlphaThenDoneEvent(stream, callIndex, createAssistantMessage);
|
||||
callIndex += 1;
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
|
||||
const unsubscribe = agent.subscribe(event => {
|
||||
@@ -160,7 +130,11 @@ describe("Agent", () => {
|
||||
await agent.prompt("refresh tools");
|
||||
unsubscribe();
|
||||
|
||||
expect(callContexts).toEqual([
|
||||
const observed = mock.calls.map(call => ({
|
||||
systemPrompt: call.context.systemPrompt?.join("\n\n") ?? "",
|
||||
toolNames: (call.context.tools ?? []).map(tool => tool.name),
|
||||
}));
|
||||
expect(observed).toEqual([
|
||||
{ systemPrompt: "prompt-one", toolNames: ["alpha"] },
|
||||
{ systemPrompt: "prompt-two", toolNames: ["alpha", "beta"] },
|
||||
]);
|
||||
@@ -169,8 +143,6 @@ describe("Agent", () => {
|
||||
it("prompt() drops stale forced toolChoice after same-turn tool refresh", async () => {
|
||||
const toolSchema = Type.Object({ value: Type.String() });
|
||||
type Details = { value: string };
|
||||
let callIndex = 0;
|
||||
const providerCalls: Array<{ toolNames: string[]; toolChoice: SimpleStreamOptions["toolChoice"] }> = [];
|
||||
|
||||
const betaTool: AgentTool<typeof toolSchema, Details> = {
|
||||
name: "beta",
|
||||
@@ -191,24 +163,20 @@ describe("Agent", () => {
|
||||
},
|
||||
};
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tool-1", name: "alpha", arguments: { value: "hello" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
model: getBundledModel("openai", "gpt-4o-mini"),
|
||||
model: mock.model,
|
||||
tools: [alphaTool],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: (_model, context, options) => {
|
||||
providerCalls.push({
|
||||
toolNames: (context.tools ?? []).map(tool => tool.name),
|
||||
toolChoice: options?.toolChoice,
|
||||
});
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
pushAlphaThenDoneEvent(stream, callIndex, createAssistantMessage);
|
||||
callIndex += 1;
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
|
||||
const unsubscribe = agent.subscribe(event => {
|
||||
@@ -220,7 +188,11 @@ describe("Agent", () => {
|
||||
await agent.prompt("refresh tools", { toolChoice: { type: "function", name: "alpha" } });
|
||||
unsubscribe();
|
||||
|
||||
expect(providerCalls).toEqual([
|
||||
const observed = mock.calls.map(call => ({
|
||||
toolNames: (call.context.tools ?? []).map(tool => tool.name),
|
||||
toolChoice: call.options?.toolChoice,
|
||||
}));
|
||||
expect(observed).toEqual([
|
||||
{ toolNames: ["alpha"], toolChoice: { type: "function", name: "alpha" } },
|
||||
{ toolNames: ["beta"], toolChoice: undefined },
|
||||
]);
|
||||
@@ -239,25 +211,21 @@ describe("Agent", () => {
|
||||
},
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const reasoningPerCall: Array<SimpleStreamOptions["reasoning"]> = [];
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tool-1", name: "alpha", arguments: { value: "hello" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
model: getBundledModel("openai", "gpt-4o-mini"),
|
||||
model: mock.model,
|
||||
thinkingLevel: ThinkingLevel.Low,
|
||||
tools: [alphaTool],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: (_model, _context, options) => {
|
||||
reasoningPerCall.push(options?.reasoning);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
pushAlphaThenDoneEvent(stream, callIndex, createAssistantMessage);
|
||||
callIndex += 1;
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
|
||||
// Bump thinking level mid-run, after the first assistant turn finishes
|
||||
@@ -271,6 +239,7 @@ describe("Agent", () => {
|
||||
await agent.prompt("run");
|
||||
unsubscribe();
|
||||
|
||||
const reasoningPerCall: Array<SimpleStreamOptions["reasoning"]> = mock.calls.map(call => call.options?.reasoning);
|
||||
expect(reasoningPerCall).toEqual([ThinkingLevel.Low, ThinkingLevel.High]);
|
||||
});
|
||||
|
||||
|
||||
@@ -1,30 +1,23 @@
|
||||
import type { AssistantMessage, Usage } from "@oh-my-pi/pi-ai";
|
||||
import type { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import type { AssistantMessage, Usage, UserMessage } from "@oh-my-pi/pi-ai";
|
||||
|
||||
/**
|
||||
* Pushes the canonical two-call stream script used across agent tests:
|
||||
* - call 0: assistant emits a single `alpha` tool call with argument `{ value: "hello" }`
|
||||
* - call 1+: assistant emits a plain text "done" response
|
||||
* Caller is responsible for incrementing its own `callIndex` counter after invoking.
|
||||
*/
|
||||
export function pushAlphaThenDoneEvent(
|
||||
stream: AssistantMessageEventStream,
|
||||
callIndex: number,
|
||||
createAssistantMessage: (
|
||||
content: AssistantMessage["content"],
|
||||
stopReason?: AssistantMessage["stopReason"],
|
||||
) => AssistantMessage,
|
||||
): void {
|
||||
if (callIndex === 0) {
|
||||
const message = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tool-1", name: "alpha", arguments: { value: "hello" } }],
|
||||
"toolUse",
|
||||
);
|
||||
stream.push({ type: "done", reason: "toolUse", message });
|
||||
} else {
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
export function createUserMessage(text: string): UserMessage {
|
||||
return { role: "user", content: text, timestamp: Date.now() };
|
||||
}
|
||||
|
||||
export function createAssistantMessage(
|
||||
content: AssistantMessage["content"],
|
||||
stopReason: AssistantMessage["stopReason"] = "stop",
|
||||
): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content,
|
||||
api: "mock",
|
||||
provider: "mock",
|
||||
model: "mock-model",
|
||||
usage: createUsage(),
|
||||
stopReason,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function createUsage(): Usage {
|
||||
@@ -37,19 +30,3 @@ function createUsage(): Usage {
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
}
|
||||
|
||||
export function createAssistantMessage(
|
||||
content: AssistantMessage["content"],
|
||||
stopReason: AssistantMessage["stopReason"] = "stop",
|
||||
): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content,
|
||||
api: "openai-responses",
|
||||
provider: "openai",
|
||||
model: "mock",
|
||||
usage: createUsage(),
|
||||
stopReason,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
+291
-252
@@ -9,6 +9,7 @@ import { afterAll, afterEach, beforeAll, describe, expect, it } from "bun:test";
|
||||
import { agentLoop } from "@oh-my-pi/pi-agent-core/agent-loop";
|
||||
import {
|
||||
type AgentTelemetryConfig,
|
||||
type ChatUsageEvent,
|
||||
GenAIAttr,
|
||||
GenAIOperation,
|
||||
recordHandoff,
|
||||
@@ -17,8 +18,9 @@ import {
|
||||
type TelemetryHookContext,
|
||||
} from "@oh-my-pi/pi-agent-core/telemetry";
|
||||
import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "@oh-my-pi/pi-agent-core/types";
|
||||
import type { Message, Model, UserMessage } from "@oh-my-pi/pi-ai";
|
||||
import { AssistantMessageEventStream, type EventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import type { Message } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import type { EventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { context, SpanStatusCode, trace } from "@opentelemetry/api";
|
||||
import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks";
|
||||
import {
|
||||
@@ -28,9 +30,9 @@ import {
|
||||
SimpleSpanProcessor,
|
||||
} from "@opentelemetry/sdk-trace-base";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { createAssistantMessage } from "./helpers";
|
||||
import { createUserMessage } from "./helpers";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
const MOCK_IDENT = { id: "mock-model", provider: "mock-provider" } as const;
|
||||
|
||||
const exporter = new InMemorySpanExporter();
|
||||
let provider: BasicTracerProvider;
|
||||
@@ -52,25 +54,6 @@ afterAll(async () => {
|
||||
context.disable();
|
||||
});
|
||||
|
||||
function createModel(): Model<"openai-responses"> {
|
||||
return {
|
||||
id: "mock-model",
|
||||
name: "mock",
|
||||
api: "openai-responses",
|
||||
provider: "mock-provider",
|
||||
baseUrl: "https://example.invalid",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 8192,
|
||||
maxTokens: 2048,
|
||||
};
|
||||
}
|
||||
|
||||
function createUserMessage(text: string): UserMessage {
|
||||
return { role: "user", content: text, timestamp: Date.now() };
|
||||
}
|
||||
|
||||
function identityConverter(messages: AgentMessage[]): Message[] {
|
||||
return messages.filter(m => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[];
|
||||
}
|
||||
@@ -91,38 +74,34 @@ function spansByName(spans: ReadableSpan[], name: string): ReadableSpan[] {
|
||||
|
||||
describe("agent-loop OTEL instrumentation", () => {
|
||||
it("emits no spans when telemetry is unset (zero-cost path)", async () => {
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
s.push({ type: "done", reason: "stop", message: createAssistantMessage([{ type: "text", text: "ok" }]) });
|
||||
});
|
||||
return s;
|
||||
};
|
||||
const mock = createMockModel({ ...MOCK_IDENT, responses: [{ content: ["ok"] }] });
|
||||
const config: AgentLoopConfig = { model: mock.model, convertToLlm: identityConverter };
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
expect(exporter.getFinishedSpans()).toHaveLength(0);
|
||||
});
|
||||
|
||||
it("emits invoke_agent → chat hierarchy with full gen_ai.* attribute envelope", async () => {
|
||||
const finalMsg = createAssistantMessage([{ type: "text", text: "hello" }]);
|
||||
finalMsg.usage = {
|
||||
input: 12,
|
||||
output: 34,
|
||||
cacheRead: 5,
|
||||
cacheWrite: 7,
|
||||
totalTokens: 58,
|
||||
reasoningTokens: 11,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
finalMsg.stopReason = "stop";
|
||||
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [
|
||||
{
|
||||
content: ["hello"],
|
||||
stopReason: "stop",
|
||||
usage: {
|
||||
input: 12,
|
||||
output: 34,
|
||||
cacheRead: 5,
|
||||
cacheWrite: 7,
|
||||
totalTokens: 58,
|
||||
reasoningTokens: 11,
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
sessionId: "conv-42",
|
||||
temperature: 0.7,
|
||||
@@ -131,13 +110,8 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
presencePenalty: 0.1,
|
||||
telemetry: { agent: { id: "agent-1", name: "researcher", description: "test-agent" } },
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => s.push({ type: "done", reason: "stop", message: finalMsg }));
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: ["you are helpful"], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const finished = exporter.getFinishedSpans();
|
||||
const invoke = findSpan(finished, "invoke_agent researcher");
|
||||
@@ -179,9 +153,15 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("emits execute_tool spans parented to invoke_agent (not chat) per semconv", async () => {
|
||||
let callIndex = 0;
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tc-1", name: "alpha", arguments: { value: "x" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {},
|
||||
};
|
||||
@@ -193,28 +173,8 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
parameters: alphaSchema,
|
||||
execute: async () => ({ content: [{ type: "text", text: "alpha-result" }], details: {} }),
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const m = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tc-1", name: "alpha", arguments: { value: "x" } }],
|
||||
"toolUse",
|
||||
);
|
||||
s.push({ type: "done", reason: "toolUse", message: m });
|
||||
} else {
|
||||
s.push({
|
||||
type: "done",
|
||||
reason: "stop",
|
||||
message: createAssistantMessage([{ type: "text", text: "done" }]),
|
||||
});
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [alphaTool] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const finished = exporter.getFinishedSpans();
|
||||
const invoke = findSpan(finished, "invoke_agent");
|
||||
@@ -237,10 +197,16 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("parents downstream spans created during tool execution (active-context propagation)", async () => {
|
||||
let callIndex = 0;
|
||||
const userTracer = trace.getTracer("user-tool");
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tc-1", name: "probe", arguments: { value: "x" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {},
|
||||
};
|
||||
@@ -256,28 +222,8 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
return { content: [{ type: "text", text: "ok" }], details: {} };
|
||||
},
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const m = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tc-1", name: "probe", arguments: { value: "x" } }],
|
||||
"toolUse",
|
||||
);
|
||||
s.push({ type: "done", reason: "toolUse", message: m });
|
||||
} else {
|
||||
s.push({
|
||||
type: "done",
|
||||
reason: "stop",
|
||||
message: createAssistantMessage([{ type: "text", text: "done" }]),
|
||||
});
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [probeTool] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const finished = exporter.getFinishedSpans();
|
||||
const tool = findSpan(finished, "execute_tool probe");
|
||||
@@ -288,9 +234,15 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("records ERROR status + exception when a tool throws", async () => {
|
||||
let callIndex = 0;
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tc-1", name: "fail", arguments: { value: "x" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {},
|
||||
};
|
||||
@@ -304,28 +256,8 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
throw new Error("boom");
|
||||
},
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const m = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tc-1", name: "fail", arguments: { value: "x" } }],
|
||||
"toolUse",
|
||||
);
|
||||
s.push({ type: "done", reason: "toolUse", message: m });
|
||||
} else {
|
||||
s.push({
|
||||
type: "done",
|
||||
reason: "stop",
|
||||
message: createAssistantMessage([{ type: "text", text: "done" }]),
|
||||
});
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [failTool] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const tool = findSpan(exporter.getFinishedSpans(), "execute_tool fail");
|
||||
expect(tool).toBeDefined();
|
||||
@@ -335,21 +267,17 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("emits ERROR status on chat spans when stopReason is error", async () => {
|
||||
const errMsg = createAssistantMessage([{ type: "text", text: "" }], "error");
|
||||
errMsg.errorMessage = "provider returned 500";
|
||||
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [{ throw: "provider returned 500" }],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {},
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => s.push({ type: "error", reason: "error", error: errMsg }));
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const chat = findSpan(exporter.getFinishedSpans(), "chat mock-model");
|
||||
expect(chat).toBeDefined();
|
||||
@@ -359,21 +287,17 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("captures request/response content when captureMessageContent is true", async () => {
|
||||
const finalMsg = createAssistantMessage([{ type: "text", text: "hi back" }]);
|
||||
finalMsg.stopReason = "stop";
|
||||
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [{ content: ["hi back"], stopReason: "stop" }],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: { captureMessageContent: true },
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => s.push({ type: "done", reason: "stop", message: finalMsg }));
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: ["sys-instruction"], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const chat = findSpan(exporter.getFinishedSpans(), "chat mock-model");
|
||||
const inputs = chat?.attributes[GenAIAttr.InputMessages] as string | undefined;
|
||||
@@ -387,21 +311,17 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("captures bounded dashboard summary content when requested", async () => {
|
||||
const finalMsg = createAssistantMessage([{ type: "text", text: "hi back" }]);
|
||||
finalMsg.stopReason = "stop";
|
||||
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [{ content: ["hi back"], stopReason: "stop" }],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: { captureMessageContent: "summary" },
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => s.push({ type: "done", reason: "stop", message: finalMsg }));
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: ["sys-instruction"], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const chat = findSpan(exporter.getFinishedSpans(), "chat mock-model");
|
||||
const request = JSON.parse(chat?.attributes[GenAIAttr.RequestMessages] as string) as Array<{
|
||||
@@ -416,19 +336,24 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("invokes costEstimator and stamps gen_ai.cost.estimated_usd", async () => {
|
||||
const finalMsg = createAssistantMessage([{ type: "text", text: "ok" }]);
|
||||
finalMsg.usage = {
|
||||
input: 1000,
|
||||
output: 500,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 1500,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
finalMsg.stopReason = "stop";
|
||||
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [
|
||||
{
|
||||
content: ["ok"],
|
||||
stopReason: "stop",
|
||||
usage: {
|
||||
input: 1000,
|
||||
output: 500,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 1500,
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {
|
||||
costEstimator: input => ({
|
||||
@@ -438,13 +363,8 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
}),
|
||||
},
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => s.push({ type: "done", reason: "stop", message: finalMsg }));
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const chat = findSpan(exporter.getFinishedSpans(), "chat mock-model");
|
||||
expect(chat?.attributes[GenAIAttr.CostEstimatedUsd]).toBeCloseTo(0.0105, 6);
|
||||
@@ -453,16 +373,22 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("applies dynamic attributes, normalization hooks, and cost deltas", async () => {
|
||||
const finalMsg = createAssistantMessage([{ type: "text", text: "ok" }]);
|
||||
finalMsg.usage = {
|
||||
input: 200,
|
||||
output: 100,
|
||||
cacheRead: 5,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 305,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
finalMsg.stopReason = "stop";
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [
|
||||
{
|
||||
content: ["ok"],
|
||||
stopReason: "stop",
|
||||
usage: {
|
||||
input: 200,
|
||||
output: 100,
|
||||
cacheRead: 5,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 305,
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
const deltas: Array<{
|
||||
costUsd: number | undefined;
|
||||
model: string;
|
||||
@@ -471,25 +397,19 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
}> = [];
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {
|
||||
agent: { name: "prefix.worker" },
|
||||
normalizeAgentName: name => name?.replace(/^prefix\./, ""),
|
||||
normalizeProvider: provider =>
|
||||
provider === "mock-provider" || provider === "mock" ? "normalized-provider" : provider,
|
||||
normalizeProvider: provider => (provider === "mock-provider" ? "normalized-provider" : provider),
|
||||
resolveAttributes: ctx => ({ "tenant.id": "tenant-1", "telemetry.kind": ctx.kind }),
|
||||
costEstimator: () => ({ usd: 0.25 }),
|
||||
onCostDelta: delta => deltas.push(delta),
|
||||
},
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => s.push({ type: "done", reason: "stop", message: finalMsg }));
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const invoke = findSpan(exporter.getFinishedSpans(), "invoke_agent worker");
|
||||
const chat = findSpan(exporter.getFinishedSpans(), "chat mock-model");
|
||||
@@ -504,33 +424,188 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("emits gen_ai.cost.unavailable_reason when the estimator declines", async () => {
|
||||
const finalMsg = createAssistantMessage([{ type: "text", text: "ok" }]);
|
||||
finalMsg.stopReason = "stop";
|
||||
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [{ content: ["ok"], stopReason: "stop" }],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: { costEstimator: () => ({ unavailable: "unsupported_tier" }) },
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => s.push({ type: "done", reason: "stop", message: finalMsg }));
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const chat = findSpan(exporter.getFinishedSpans(), "chat mock-model");
|
||||
expect(chat?.attributes[GenAIAttr.CostUnavailableReason]).toBe("unsupported_tier");
|
||||
expect(chat?.attributes[GenAIAttr.CostEstimatedUsd]).toBeUndefined();
|
||||
});
|
||||
|
||||
it("fires onChatUsage for every chat step regardless of cost estimator", async () => {
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [
|
||||
{
|
||||
content: ["ok"],
|
||||
stopReason: "stop",
|
||||
usage: { input: 50, output: 25, cacheRead: 10, totalTokens: 85 },
|
||||
},
|
||||
],
|
||||
});
|
||||
const events: ChatUsageEvent[] = [];
|
||||
const config: AgentLoopConfig = {
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {
|
||||
agent: { id: "agent-1", name: "worker" },
|
||||
resolveAttributes: ctx => ({ "tenant.id": "tenant-7", "telemetry.kind": ctx.kind }),
|
||||
normalizeProvider: provider => (provider === "mock-provider" ? "normalized-provider" : provider),
|
||||
onChatUsage: event => {
|
||||
events.push(event);
|
||||
},
|
||||
},
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
expect(events).toHaveLength(1);
|
||||
const ev = events[0];
|
||||
expect(ev?.model).toBe("mock-model");
|
||||
expect(ev?.provider).toBe("normalized-provider");
|
||||
expect(ev?.stepNumber).toBe(0);
|
||||
expect(ev?.agent).toEqual({ id: "agent-1", name: "worker" });
|
||||
expect(ev?.usage.inputTokens).toBe(50);
|
||||
expect(ev?.usage.outputTokens).toBe(25);
|
||||
expect(ev?.usage.cachedInputTokens).toBe(10);
|
||||
expect(ev?.usage.totalTokens).toBe(85);
|
||||
expect(ev?.cost).toBeUndefined();
|
||||
expect(ev?.attributes?.["tenant.id"]).toBe("tenant-7");
|
||||
expect(ev?.attributes?.["telemetry.kind"]).toBe("chat");
|
||||
expect(ev?.span).toBeDefined();
|
||||
});
|
||||
|
||||
it("forwards cost estimate to onChatUsage when estimator is configured", async () => {
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [{ content: ["ok"], stopReason: "stop", usage: { input: 100, output: 50, totalTokens: 150 } }],
|
||||
});
|
||||
const events: ChatUsageEvent[] = [];
|
||||
const config: AgentLoopConfig = {
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {
|
||||
costEstimator: () => ({ usd: 0.05, inputUsd: 0.01, outputUsd: 0.04 }),
|
||||
onChatUsage: event => {
|
||||
events.push(event);
|
||||
},
|
||||
},
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
expect(events).toHaveLength(1);
|
||||
const cost = events[0]?.cost;
|
||||
expect(cost && "usd" in cost ? cost.usd : undefined).toBe(0.05);
|
||||
expect(cost && "usd" in cost ? cost.inputUsd : undefined).toBe(0.01);
|
||||
expect(cost && "usd" in cost ? cost.outputUsd : undefined).toBe(0.04);
|
||||
});
|
||||
|
||||
it("propagates unavailable cost reason to onChatUsage", async () => {
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [{ content: ["ok"], stopReason: "stop", usage: { input: 10, output: 5, totalTokens: 15 } }],
|
||||
});
|
||||
const events: ChatUsageEvent[] = [];
|
||||
const config: AgentLoopConfig = {
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {
|
||||
costEstimator: () => ({ unavailable: "unsupported_tier" }),
|
||||
onChatUsage: event => {
|
||||
events.push(event);
|
||||
},
|
||||
},
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
expect(events).toHaveLength(1);
|
||||
const cost = events[0]?.cost;
|
||||
expect(cost && "unavailable" in cost ? cost.unavailable : undefined).toBe("unsupported_tier");
|
||||
});
|
||||
|
||||
it("skips onChatUsage in recordManualChatTelemetry when usage is undefined", async () => {
|
||||
const events: ChatUsageEvent[] = [];
|
||||
const telemetry = resolveTelemetry(
|
||||
{
|
||||
onChatUsage: event => {
|
||||
events.push(event);
|
||||
},
|
||||
},
|
||||
undefined,
|
||||
);
|
||||
const mock = createMockModel({ ...MOCK_IDENT, responses: [] });
|
||||
await recordManualChatTelemetry(telemetry, {
|
||||
model: mock.model,
|
||||
responseModel: "manual-model",
|
||||
stepNumber: 0,
|
||||
});
|
||||
expect(events).toHaveLength(0);
|
||||
|
||||
await recordManualChatTelemetry(telemetry, {
|
||||
model: mock.model,
|
||||
responseModel: "manual-model",
|
||||
usage: {
|
||||
input: 7,
|
||||
output: 3,
|
||||
totalTokens: 10,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stepNumber: 1,
|
||||
});
|
||||
expect(events).toHaveLength(1);
|
||||
expect(events[0]?.model).toBe("manual-model");
|
||||
expect(events[0]?.stepNumber).toBe(1);
|
||||
expect(events[0]?.usage.inputTokens).toBe(7);
|
||||
});
|
||||
|
||||
it("captures async onChatUsage rejections via onTelemetryWarning", async () => {
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [{ content: ["ok"], stopReason: "stop", usage: { input: 1, output: 1, totalTokens: 2 } }],
|
||||
});
|
||||
const warnings: string[] = [];
|
||||
const config: AgentLoopConfig = {
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {
|
||||
onChatUsage: async () => {
|
||||
throw new Error("async boom");
|
||||
},
|
||||
onTelemetryWarning: warning => warnings.push(warning.code),
|
||||
},
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
await Bun.sleep(1);
|
||||
|
||||
expect(warnings).toContain("on_chat_usage_failed");
|
||||
});
|
||||
|
||||
it("fires onSpanStart and onSpanEnd for every kind", async () => {
|
||||
let callIndex = 0;
|
||||
const starts: TelemetryHookContext[] = [];
|
||||
const ends: TelemetryHookContext[] = [];
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tc-1", name: "echo", arguments: { value: "x" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {
|
||||
agent: { id: "a", name: "main" },
|
||||
@@ -546,28 +621,8 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
parameters: echoSchema,
|
||||
execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {} }),
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const m = createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tc-1", name: "echo", arguments: { value: "x" } }],
|
||||
"toolUse",
|
||||
);
|
||||
s.push({ type: "done", reason: "toolUse", message: m });
|
||||
} else {
|
||||
s.push({
|
||||
type: "done",
|
||||
reason: "stop",
|
||||
message: createAssistantMessage([{ type: "text", text: "done" }]),
|
||||
});
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [echoTool] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
const startKinds = starts.map(s => s.kind);
|
||||
const endKinds = ends.map(s => s.kind);
|
||||
@@ -593,7 +648,7 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
expect(span?.attributes[GenAIAttr.ConversationId]).toBe("conv-1");
|
||||
});
|
||||
|
||||
it("records manual chat telemetry for non-loop model calls", () => {
|
||||
it("records manual chat telemetry for non-loop model calls", async () => {
|
||||
const telemetry = resolveTelemetry(
|
||||
{
|
||||
costEstimator: () => ({ usd: 0.02 }),
|
||||
@@ -602,8 +657,9 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
);
|
||||
expect(telemetry).toBeDefined();
|
||||
|
||||
recordManualChatTelemetry(telemetry, {
|
||||
model: createModel(),
|
||||
const mock = createMockModel({ ...MOCK_IDENT });
|
||||
await recordManualChatTelemetry(telemetry, {
|
||||
model: mock.model,
|
||||
usage: {
|
||||
input: 10,
|
||||
output: 5,
|
||||
@@ -645,12 +701,18 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
});
|
||||
|
||||
it("attaches user-supplied attributes to every span", async () => {
|
||||
let callIndex = 0;
|
||||
const cfg: AgentTelemetryConfig = {
|
||||
attributes: { "deployment.environment": "prod", "service.name": "test-svc" },
|
||||
};
|
||||
const mock = createMockModel({
|
||||
...MOCK_IDENT,
|
||||
responses: [
|
||||
{ content: [{ type: "toolCall", id: "tc-1", name: "echo", arguments: { value: "x" } }] },
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: cfg,
|
||||
};
|
||||
@@ -662,31 +724,8 @@ describe("agent-loop OTEL instrumentation", () => {
|
||||
parameters: echoSchema,
|
||||
execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {} }),
|
||||
};
|
||||
const streamFn = () => {
|
||||
const s = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
s.push({
|
||||
type: "done",
|
||||
reason: "toolUse",
|
||||
message: createAssistantMessage(
|
||||
[{ type: "toolCall", id: "tc-1", name: "echo", arguments: { value: "x" } }],
|
||||
"toolUse",
|
||||
),
|
||||
});
|
||||
} else {
|
||||
s.push({
|
||||
type: "done",
|
||||
reason: "stop",
|
||||
message: createAssistantMessage([{ type: "text", text: "done" }]),
|
||||
});
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return s;
|
||||
};
|
||||
const ctx: AgentContext = { systemPrompt: [], messages: [], tools: [tool] };
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, streamFn));
|
||||
await runAndDrain(agentLoop([createUserMessage("hi")], ctx, config, undefined, mock.stream));
|
||||
|
||||
for (const span of exporter.getFinishedSpans()) {
|
||||
expect(span.attributes["deployment.environment"]).toBe("prod");
|
||||
|
||||
@@ -17,8 +17,8 @@ import {
|
||||
} from "@oh-my-pi/pi-agent-core/run-collector";
|
||||
import { AGGREGATE_ATTR, EXECUTE_TOOL_STATUS_ATTR, GenAIAttr } from "@oh-my-pi/pi-agent-core/telemetry";
|
||||
import type { AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "@oh-my-pi/pi-agent-core/types";
|
||||
import type { AssistantMessage, Context, Message, Model, UserMessage } from "@oh-my-pi/pi-ai";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import type { AssistantMessage, Message } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import type {
|
||||
AttributeValue,
|
||||
Context as OtelContext,
|
||||
@@ -29,9 +29,7 @@ import type {
|
||||
Tracer,
|
||||
} from "@opentelemetry/api";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { createAssistantMessage } from "./helpers";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
import { createUserMessage } from "./helpers";
|
||||
|
||||
interface RecordedSpan {
|
||||
readonly name: string;
|
||||
@@ -101,25 +99,6 @@ function makeFakeSpan(record: RecordedSpan): Span {
|
||||
return span;
|
||||
}
|
||||
|
||||
function createModel(): Model<"openai-responses"> {
|
||||
return {
|
||||
id: "mock-model",
|
||||
name: "mock",
|
||||
api: "openai-responses",
|
||||
provider: "mock-provider",
|
||||
baseUrl: "https://example.invalid",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 8192,
|
||||
maxTokens: 2048,
|
||||
};
|
||||
}
|
||||
|
||||
function createUserMessage(text: string): UserMessage {
|
||||
return { role: "user", content: text, timestamp: Date.now() };
|
||||
}
|
||||
|
||||
function identityConverter(messages: AgentMessage[]): Message[] {
|
||||
return messages.filter(m => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[];
|
||||
}
|
||||
@@ -141,10 +120,6 @@ function makeUsage(
|
||||
};
|
||||
}
|
||||
|
||||
function withUsage(message: AssistantMessage, usage: AssistantMessage["usage"]): AssistantMessage {
|
||||
return { ...message, usage };
|
||||
}
|
||||
|
||||
interface TestTool {
|
||||
readonly name: string;
|
||||
readonly behavior: "ok" | "throw" | "block";
|
||||
@@ -176,55 +151,14 @@ function buildTool(spec: TestTool): AgentTool {
|
||||
} satisfies AgentTool;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build a stream factory that walks the agent through `script` — one
|
||||
* assistant message per call. Each entry is either a final text response
|
||||
* (`{ text }`) or a tool-call message (`{ toolCalls }`).
|
||||
*/
|
||||
function scriptedStreamFn(
|
||||
script: readonly (
|
||||
| { readonly text: string; readonly usage?: AssistantMessage["usage"] }
|
||||
| {
|
||||
readonly toolCalls: readonly {
|
||||
readonly id: string;
|
||||
readonly name: string;
|
||||
readonly args?: Record<string, unknown>;
|
||||
}[];
|
||||
readonly usage?: AssistantMessage["usage"];
|
||||
}
|
||||
)[],
|
||||
) {
|
||||
let callIndex = 0;
|
||||
return (_model: Model, _ctx: Context) => {
|
||||
const stream = new MockAssistantStream();
|
||||
const entry = script[callIndex++] ?? script[script.length - 1];
|
||||
queueMicrotask(() => {
|
||||
const base =
|
||||
"text" in entry
|
||||
? createAssistantMessage([{ type: "text", text: entry.text }], "stop")
|
||||
: createAssistantMessage(
|
||||
entry.toolCalls.map(tc => ({
|
||||
type: "toolCall",
|
||||
id: tc.id,
|
||||
name: tc.name,
|
||||
arguments: tc.args ?? { value: "x" },
|
||||
})),
|
||||
"toolUse",
|
||||
);
|
||||
const message = entry.usage ? withUsage(base, entry.usage) : base;
|
||||
const reason =
|
||||
message.stopReason === "toolUse" || message.stopReason === "length" ? message.stopReason : "stop";
|
||||
stream.push({ type: "done", reason, message });
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
}
|
||||
|
||||
describe("AgentRunSummary delivery", () => {
|
||||
it("populates telemetry/coverage on agent_end when telemetry: {} is supplied", async () => {
|
||||
const tracer = new RecordingTracer();
|
||||
const mock = createMockModel({
|
||||
responses: [{ content: ["ok"], usage: makeUsage(7, 3) }],
|
||||
});
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: { tracer },
|
||||
};
|
||||
@@ -234,7 +168,7 @@ describe("AgentRunSummary delivery", () => {
|
||||
{ systemPrompt: ["sys"], messages: [], tools: [] },
|
||||
config,
|
||||
undefined,
|
||||
scriptedStreamFn([{ text: "ok", usage: makeUsage(7, 3) }]),
|
||||
mock.stream,
|
||||
);
|
||||
for await (const event of stream) events.push(event);
|
||||
const endEvent = events.find((e): e is Extract<AgentEvent, { type: "agent_end" }> => e.type === "agent_end");
|
||||
@@ -248,8 +182,9 @@ describe("AgentRunSummary delivery", () => {
|
||||
|
||||
it("emits no spans and no summary when telemetry is unset", async () => {
|
||||
const tracer = new RecordingTracer();
|
||||
const mock = createMockModel({ responses: [{ content: ["ok"] }] });
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
// telemetry intentionally unset.
|
||||
};
|
||||
@@ -259,7 +194,7 @@ describe("AgentRunSummary delivery", () => {
|
||||
{ systemPrompt: ["sys"], messages: [], tools: [] },
|
||||
config,
|
||||
undefined,
|
||||
scriptedStreamFn([{ text: "ok" }]),
|
||||
mock.stream,
|
||||
);
|
||||
for await (const event of stream) events.push(event);
|
||||
expect(tracer.spans.length).toBe(0);
|
||||
@@ -270,8 +205,9 @@ describe("AgentRunSummary delivery", () => {
|
||||
|
||||
it("preserves agentLoop().result() backwards-compat (still resolves to AgentMessage[])", async () => {
|
||||
const tracer = new RecordingTracer();
|
||||
const mock = createMockModel({ responses: [{ content: ["ok"] }] });
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: { tracer },
|
||||
};
|
||||
@@ -280,7 +216,7 @@ describe("AgentRunSummary delivery", () => {
|
||||
{ systemPrompt: ["sys"], messages: [], tools: [] },
|
||||
config,
|
||||
undefined,
|
||||
scriptedStreamFn([{ text: "ok" }]),
|
||||
mock.stream,
|
||||
);
|
||||
const messages = await stream.result();
|
||||
// 1 user prompt + 1 assistant message.
|
||||
@@ -294,11 +230,20 @@ describe("AgentRunSummary aggregation", () => {
|
||||
it("sums token + cost totals across multiple chats and counts stop_reasons", async () => {
|
||||
const tracer = new RecordingTracer();
|
||||
const tool = buildTool({ name: "alpha", behavior: "ok" });
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [{ type: "toolCall", id: "a-1", name: "alpha", arguments: { value: "x" } }],
|
||||
usage: makeUsage(5, 2),
|
||||
},
|
||||
{ content: ["wrap"], usage: makeUsage(8, 1) },
|
||||
],
|
||||
});
|
||||
const detailed = agentLoopDetailed(
|
||||
[createUserMessage("hi")],
|
||||
{ systemPrompt: ["sys"], messages: [], tools: [tool] },
|
||||
{
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {
|
||||
tracer,
|
||||
@@ -306,10 +251,7 @@ describe("AgentRunSummary aggregation", () => {
|
||||
},
|
||||
},
|
||||
undefined,
|
||||
scriptedStreamFn([
|
||||
{ toolCalls: [{ id: "a-1", name: "alpha" }], usage: makeUsage(5, 2) },
|
||||
{ text: "wrap", usage: makeUsage(8, 1) },
|
||||
]),
|
||||
mock.stream,
|
||||
);
|
||||
for await (const _ of detailed.stream) {
|
||||
// drain
|
||||
@@ -335,11 +277,23 @@ describe("AgentRunSummary aggregation", () => {
|
||||
buildTool({ name: "err-tool", behavior: "throw" }),
|
||||
buildTool({ name: "blocked-tool", behavior: "block" }),
|
||||
];
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [
|
||||
{ type: "toolCall", id: "t-1", name: "ok-tool", arguments: { value: "x" } },
|
||||
{ type: "toolCall", id: "t-2", name: "err-tool", arguments: { value: "x" } },
|
||||
{ type: "toolCall", id: "t-3", name: "blocked-tool", arguments: { value: "x" } },
|
||||
],
|
||||
},
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const detailed = agentLoopDetailed(
|
||||
[createUserMessage("hi")],
|
||||
{ systemPrompt: ["sys"], messages: [], tools },
|
||||
{
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: { tracer },
|
||||
beforeToolCall: async ctx => {
|
||||
@@ -348,16 +302,7 @@ describe("AgentRunSummary aggregation", () => {
|
||||
},
|
||||
},
|
||||
undefined,
|
||||
scriptedStreamFn([
|
||||
{
|
||||
toolCalls: [
|
||||
{ id: "t-1", name: "ok-tool" },
|
||||
{ id: "t-2", name: "err-tool" },
|
||||
{ id: "t-3", name: "blocked-tool" },
|
||||
],
|
||||
},
|
||||
{ text: "done" },
|
||||
]),
|
||||
mock.stream,
|
||||
);
|
||||
for await (const _ of detailed.stream) {
|
||||
// drain
|
||||
@@ -379,15 +324,21 @@ describe("AgentRunSummary aggregation", () => {
|
||||
it("populates aggregate gen_ai.agent.* attributes on the invoke_agent span", async () => {
|
||||
const tracer = new RecordingTracer();
|
||||
const tool = buildTool({ name: "alpha", behavior: "ok" });
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [{ type: "toolCall", id: "a-1", name: "alpha", arguments: { value: "x" } }],
|
||||
usage: makeUsage(4, 6),
|
||||
},
|
||||
{ content: ["done"], usage: makeUsage(2, 1) },
|
||||
],
|
||||
});
|
||||
const detailed = agentLoopDetailed(
|
||||
[createUserMessage("hi")],
|
||||
{ systemPrompt: ["sys"], messages: [], tools: [tool] },
|
||||
{ model: createModel(), convertToLlm: identityConverter, telemetry: { tracer } },
|
||||
{ model: mock.model, convertToLlm: identityConverter, telemetry: { tracer } },
|
||||
undefined,
|
||||
scriptedStreamFn([
|
||||
{ toolCalls: [{ id: "a-1", name: "alpha" }], usage: makeUsage(4, 6) },
|
||||
{ text: "done", usage: makeUsage(2, 1) },
|
||||
]),
|
||||
mock.stream,
|
||||
);
|
||||
for await (const _ of detailed.stream) {
|
||||
// drain
|
||||
@@ -412,22 +363,25 @@ describe("AgentRunCoverage", () => {
|
||||
buildTool({ name: "alpha", behavior: "ok" }),
|
||||
buildTool({ name: "mu", behavior: "ok" }),
|
||||
];
|
||||
const detailed = agentLoopDetailed(
|
||||
[createUserMessage("hi")],
|
||||
{ systemPrompt: ["sys"], messages: [], tools },
|
||||
{ model: createModel(), convertToLlm: identityConverter, telemetry: { tracer } },
|
||||
undefined,
|
||||
scriptedStreamFn([
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
// Step 1 invokes alpha and mu.
|
||||
{
|
||||
toolCalls: [
|
||||
{ id: "t-1", name: "alpha" },
|
||||
{ id: "t-2", name: "mu" },
|
||||
content: [
|
||||
{ type: "toolCall", id: "t-1", name: "alpha", arguments: { value: "x" } },
|
||||
{ type: "toolCall", id: "t-2", name: "mu", arguments: { value: "x" } },
|
||||
],
|
||||
},
|
||||
// Step 2 wraps up with a text response — zeta is never invoked.
|
||||
{ text: "done" },
|
||||
]),
|
||||
{ content: ["done"] },
|
||||
],
|
||||
});
|
||||
const detailed = agentLoopDetailed(
|
||||
[createUserMessage("hi")],
|
||||
{ systemPrompt: ["sys"], messages: [], tools },
|
||||
{ model: mock.model, convertToLlm: identityConverter, telemetry: { tracer } },
|
||||
undefined,
|
||||
mock.stream,
|
||||
);
|
||||
for await (const _ of detailed.stream) {
|
||||
// drain
|
||||
@@ -555,11 +509,12 @@ describe("onRunEnd is non-fatal", () => {
|
||||
warnings.push(args);
|
||||
};
|
||||
try {
|
||||
const mock = createMockModel({ responses: [{ content: ["ok"] }] });
|
||||
const stream = agentLoop(
|
||||
[createUserMessage("hi")],
|
||||
{ systemPrompt: ["sys"], messages: [], tools: [] },
|
||||
{
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: {
|
||||
tracer,
|
||||
@@ -569,7 +524,7 @@ describe("onRunEnd is non-fatal", () => {
|
||||
},
|
||||
},
|
||||
undefined,
|
||||
scriptedStreamFn([{ text: "ok" }]),
|
||||
mock.stream,
|
||||
);
|
||||
const messages = await stream.result();
|
||||
expect(messages.length).toBe(2);
|
||||
@@ -617,11 +572,22 @@ describe("skipped tools without spans", () => {
|
||||
};
|
||||
let triggered = false;
|
||||
let getSteeringCallCount = 0;
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [
|
||||
{ type: "toolCall", id: "tool-fast", name: "fast", arguments: { value: "x" } },
|
||||
{ type: "toolCall", id: "tool-slow", name: "slow", arguments: { value: "x" } },
|
||||
],
|
||||
},
|
||||
{ content: ["wrap"] },
|
||||
],
|
||||
});
|
||||
const detailed = agentLoopDetailed(
|
||||
[createUserMessage("hi")],
|
||||
{ systemPrompt: ["sys"], messages: [], tools: [fastTool, slowTool] },
|
||||
{
|
||||
model: createModel(),
|
||||
model: mock.model,
|
||||
convertToLlm: identityConverter,
|
||||
telemetry: { tracer },
|
||||
interruptMode: "immediate",
|
||||
@@ -637,15 +603,7 @@ describe("skipped tools without spans", () => {
|
||||
},
|
||||
},
|
||||
undefined,
|
||||
scriptedStreamFn([
|
||||
{
|
||||
toolCalls: [
|
||||
{ id: "tool-fast", name: "fast" },
|
||||
{ id: "tool-slow", name: "slow" },
|
||||
],
|
||||
},
|
||||
{ text: "wrap" },
|
||||
]),
|
||||
mock.stream,
|
||||
);
|
||||
for await (const _ of detailed.stream) {
|
||||
// drain
|
||||
|
||||
@@ -24,8 +24,6 @@ import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manage
|
||||
import { TempDir } from "@oh-my-pi/pi-utils";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shared setup
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -77,7 +75,7 @@ async function createSession(tools: AgentTool[], bridge?: ClientBridge): Promise
|
||||
messages: [],
|
||||
},
|
||||
convertToLlm,
|
||||
streamFn: () => new MockAssistantStream(),
|
||||
streamFn: () => new AssistantMessageEventStream(),
|
||||
});
|
||||
|
||||
const sess = new AgentSession({
|
||||
|
||||
@@ -42,8 +42,8 @@ import * as fs from "node:fs";
|
||||
import * as os from "node:os";
|
||||
import * as path from "node:path";
|
||||
import { Agent, type AgentMessage, type AgentTool } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AssistantMessage, getBundledModel, type ToolCall } from "@oh-my-pi/pi-ai";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { getBundledModel } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel, type MockResponse } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
||||
@@ -53,53 +53,19 @@ import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manage
|
||||
import { BashTool, type ToolSession } from "@oh-my-pi/pi-coding-agent/tools";
|
||||
import { Snowflake } from "@oh-my-pi/pi-utils";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
/** Build an assistant turn that issues a single `bash` tool call. */
|
||||
function bashCall(command: string, callId: string): AssistantMessage {
|
||||
const toolCall: ToolCall = {
|
||||
type: "toolCall",
|
||||
id: callId,
|
||||
name: "bash",
|
||||
arguments: { command, timeout: 10 },
|
||||
};
|
||||
/** Scripted assistant turn that issues a single `bash` tool call. */
|
||||
function bashCall(command: string, callId: string): MockResponse {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [toolCall],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "mock",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
content: [{ type: "toolCall", id: callId, name: "bash", arguments: { command, timeout: 10 } }],
|
||||
stopReason: "toolUse",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
/** Build a plain text assistant turn with `stopReason: "stop"`. */
|
||||
function stopReply(text: string): AssistantMessage {
|
||||
/** Scripted plain-text assistant turn with `stopReason: "stop"`. */
|
||||
function stopReply(text: string): MockResponse {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text }],
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
model: "mock",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
@@ -151,7 +117,7 @@ describe("BashTool through AgentSession runs children in their own session (e2e)
|
||||
let session: AgentSession;
|
||||
let tempDir: string;
|
||||
let authStorage: AuthStorage | undefined;
|
||||
let scriptedResponses: AssistantMessage[];
|
||||
let scriptedResponses: MockResponse[];
|
||||
let hostSid: number;
|
||||
|
||||
beforeAll(() => {
|
||||
@@ -201,6 +167,10 @@ describe("BashTool through AgentSession runs children in their own session (e2e)
|
||||
|
||||
scriptedResponses = [];
|
||||
|
||||
const mock = createMockModel({
|
||||
handler: () => scriptedResponses.shift() ?? stopReply("done"),
|
||||
});
|
||||
|
||||
const agent = new Agent({
|
||||
getApiKey: () => "test-key",
|
||||
initialState: {
|
||||
@@ -210,17 +180,7 @@ describe("BashTool through AgentSession runs children in their own session (e2e)
|
||||
messages: [],
|
||||
},
|
||||
convertToLlm,
|
||||
streamFn: () => {
|
||||
const response = scriptedResponses.shift() ?? stopReply("done");
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: response });
|
||||
const reason =
|
||||
response.stopReason === "toolUse" || response.stopReason === "length" ? response.stopReason : "stop";
|
||||
stream.push({ type: "done", reason, message: response });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
|
||||
session = new AgentSession({
|
||||
|
||||
@@ -3,7 +3,7 @@ import * as path from "node:path";
|
||||
import { Agent, type AgentMessage } from "@oh-my-pi/pi-agent-core";
|
||||
import { getBundledModel, type Message } from "@oh-my-pi/pi-ai";
|
||||
import { inferCopilotInitiator } from "@oh-my-pi/pi-ai/providers/github-copilot-headers";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import type { ExtensionRunner } from "@oh-my-pi/pi-coding-agent/extensibility/extensions";
|
||||
@@ -12,9 +12,6 @@ import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
||||
import { convertToLlm } from "@oh-my-pi/pi-coding-agent/session/messages";
|
||||
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
||||
import { TempDir } from "@oh-my-pi/pi-utils";
|
||||
import { createAssistantMessage } from "./helpers/agent-session-setup";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
describe("AgentSession before_agent_start attribution fallback", () => {
|
||||
let tempDir: TempDir;
|
||||
@@ -67,14 +64,7 @@ describe("AgentSession before_agent_start attribution fallback", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Done") });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: createMockModel({ responses: [{ content: ["Done"] }] }).stream,
|
||||
});
|
||||
|
||||
session = new AgentSession({
|
||||
|
||||
@@ -8,6 +8,7 @@ import * as os from "node:os";
|
||||
import * as path from "node:path";
|
||||
import { Agent, AgentBusyError, type AgentTool } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AssistantMessage, getBundledModel, type Message, type ToolCall } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import type { Rule } from "@oh-my-pi/pi-coding-agent/capability/rule";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
@@ -22,7 +23,6 @@ import { Type } from "@sinclair/typebox";
|
||||
import { createAssistantMessage } from "./helpers/agent-session-setup";
|
||||
|
||||
// Mock stream that mimics AssistantMessageEventStream
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
describe("AgentSession concurrent prompt guard", () => {
|
||||
let session: AgentSession;
|
||||
@@ -61,7 +61,7 @@ describe("AgentSession concurrent prompt guard", () => {
|
||||
},
|
||||
streamFn: (_model, _context, options) => {
|
||||
abortSignal = options?.signal;
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
const checkAbort = () => {
|
||||
@@ -158,7 +158,7 @@ describe("AgentSession concurrent prompt guard", () => {
|
||||
|
||||
it("delivers hidden nextTurn stop reactions through the next LLM call without exposing them in the visible queue", async () => {
|
||||
const model = getBundledModel("anthropic", "claude-sonnet-4-5")!;
|
||||
let firstStream: MockAssistantStream | undefined;
|
||||
let firstStream: AssistantMessageEventStream | undefined;
|
||||
const callMessages: Message[][] = [];
|
||||
|
||||
const agent = new Agent({
|
||||
@@ -171,7 +171,7 @@ describe("AgentSession concurrent prompt guard", () => {
|
||||
convertToLlm,
|
||||
streamFn: (_model, context) => {
|
||||
callMessages.push([...context.messages]);
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
if (callMessages.length > 1) {
|
||||
@@ -235,6 +235,7 @@ describe("AgentSession concurrent prompt guard", () => {
|
||||
it("should allow prompt() after previous completes", async () => {
|
||||
// Create session with a stream that completes immediately
|
||||
const model = getBundledModel("anthropic", "claude-sonnet-4-5")!;
|
||||
const mock = createMockModel({ handler: () => ({ content: ["Done"] }) });
|
||||
const agent = new Agent({
|
||||
getApiKey: () => "test-key",
|
||||
initialState: {
|
||||
@@ -242,14 +243,7 @@ describe("AgentSession concurrent prompt guard", () => {
|
||||
systemPrompt: ["Test"],
|
||||
tools: [],
|
||||
},
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Done") });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
|
||||
const sessionManager = SessionManager.inMemory();
|
||||
@@ -327,7 +321,7 @@ describe("AgentSession TTSR resume gate", () => {
|
||||
};
|
||||
}
|
||||
|
||||
function pushContinuationStream(stream: MockAssistantStream, onComplete: () => void): void {
|
||||
function pushContinuationStream(stream: AssistantMessageEventStream, onComplete: () => void): void {
|
||||
setTimeout(() => {
|
||||
const partial = makeMsg("");
|
||||
stream.push({ type: "start", partial });
|
||||
@@ -342,7 +336,7 @@ describe("AgentSession TTSR resume gate", () => {
|
||||
}, 10);
|
||||
}
|
||||
|
||||
function pushAbortableTtsrStream(stream: MockAssistantStream, signal: AbortSignal | undefined): void {
|
||||
function pushAbortableTtsrStream(stream: AssistantMessageEventStream, signal: AbortSignal | undefined): void {
|
||||
queueMicrotask(() => {
|
||||
const partial = makeMsg("");
|
||||
stream.push({ type: "start", partial });
|
||||
@@ -387,7 +381,7 @@ describe("AgentSession TTSR resume gate", () => {
|
||||
initialState: { model, systemPrompt: ["Test"], tools: [] },
|
||||
streamFn: (_model, _context, options) => {
|
||||
streamCallCount++;
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
const signal = options?.signal;
|
||||
|
||||
if (streamCallCount === 1) {
|
||||
@@ -448,7 +442,7 @@ describe("AgentSession TTSR resume gate", () => {
|
||||
initialState: { model, systemPrompt: ["Test"], tools: [] },
|
||||
streamFn: (_model, _context, _options) => {
|
||||
streamCallCount++;
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
|
||||
if (streamCallCount === 1) {
|
||||
// First stream: emit matching text and complete normally
|
||||
@@ -519,7 +513,7 @@ describe("AgentSession TTSR resume gate", () => {
|
||||
getApiKey: () => "test-key",
|
||||
initialState: { model, systemPrompt: ["Test"], tools: [] },
|
||||
streamFn: (_model, _context, options) => {
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
const signal = options?.signal;
|
||||
|
||||
queueMicrotask(() => {
|
||||
@@ -636,7 +630,7 @@ describe("AgentSession TTSR resume gate", () => {
|
||||
initialState: { model, systemPrompt: ["Test"], tools: [mockTool] },
|
||||
streamFn: (_model, _context, options) => {
|
||||
streamCallCount++;
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
const signal = options?.signal;
|
||||
|
||||
if (streamCallCount === 1) {
|
||||
@@ -746,7 +740,7 @@ describe("AgentSession TTSR resume gate", () => {
|
||||
initialState: { model: sparkModel, systemPrompt: ["Test"], tools: [] },
|
||||
streamFn: () => {
|
||||
streamCallCount++;
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
if (streamCallCount === 1) {
|
||||
queueMicrotask(() => {
|
||||
const message = makeOverflowMessage();
|
||||
|
||||
@@ -15,8 +15,6 @@ import { TempDir } from "@oh-my-pi/pi-utils";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { createAssistantMessage } from "./helpers/agent-session-setup";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
type ObservedPromptCall = {
|
||||
toolChoice: string | undefined;
|
||||
toolNames: string[];
|
||||
@@ -152,7 +150,7 @@ describe("AgentSession eager todo enforcement", () => {
|
||||
lastMessageText: getMessageText(lastMessage),
|
||||
});
|
||||
const response = scriptedResponses.shift() ?? createAssistantMessage("done");
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: response });
|
||||
const reason =
|
||||
|
||||
@@ -12,8 +12,6 @@ import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manage
|
||||
import { TempDir } from "@oh-my-pi/pi-utils";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
let tempDir: TempDir;
|
||||
let authStorage: AuthStorage | undefined;
|
||||
let session: AgentSession;
|
||||
@@ -53,7 +51,7 @@ beforeEach(async () => {
|
||||
messages: [],
|
||||
},
|
||||
convertToLlm,
|
||||
streamFn: () => new MockAssistantStream(),
|
||||
streamFn: () => new AssistantMessageEventStream(),
|
||||
});
|
||||
|
||||
session = new AgentSession({
|
||||
|
||||
@@ -3,7 +3,7 @@ import * as path from "node:path";
|
||||
import { Agent } from "@oh-my-pi/pi-agent-core";
|
||||
import type { AssistantMessage, ToolCall } from "@oh-my-pi/pi-ai";
|
||||
import { getBundledModel } from "@oh-my-pi/pi-ai/models";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import { ExtensionRunner, loadExtensions } from "@oh-my-pi/pi-coding-agent/extensibility/extensions";
|
||||
@@ -12,8 +12,6 @@ import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
||||
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
||||
import { TempDir } from "@oh-my-pi/pi-utils";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
describe("AgentSession handoff", () => {
|
||||
let tempDir: TempDir;
|
||||
let session: AgentSession;
|
||||
@@ -406,8 +404,6 @@ describe("AgentSession handoff", () => {
|
||||
await session.dispose();
|
||||
sessionManager = SessionManager.create(tempDir.path(), tempDir.path());
|
||||
events = [];
|
||||
let streamCallCount = 0;
|
||||
|
||||
sessionManager.appendMessage({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "seed" }],
|
||||
@@ -431,40 +427,34 @@ describe("AgentSession handoff", () => {
|
||||
timestamp: Date.now() - 1,
|
||||
});
|
||||
|
||||
const thresholdAssistant: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "maintenance trigger" }],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
stopReason: "stop",
|
||||
usage: {
|
||||
input: 190_000,
|
||||
output: 1_000,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 191_000,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
const handoffAssistant: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "## Goal\nContinue from here" }],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
stopReason: "stop",
|
||||
usage: {
|
||||
input: 8_000,
|
||||
output: 500,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 8_500,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
timestamp: Date.now() + 1,
|
||||
};
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [{ type: "text", text: "maintenance trigger" }],
|
||||
stopReason: "stop",
|
||||
usage: {
|
||||
input: 190_000,
|
||||
output: 1_000,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 191_000,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
},
|
||||
{
|
||||
content: [{ type: "text", text: "## Goal\nContinue from here" }],
|
||||
stopReason: "stop",
|
||||
usage: {
|
||||
input: 8_000,
|
||||
output: 500,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 8_500,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const agent = new Agent({
|
||||
getApiKey: () => "test-key",
|
||||
@@ -474,16 +464,7 @@ describe("AgentSession handoff", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: () => {
|
||||
streamCallCount++;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = streamCallCount === 1 ? thresholdAssistant : handoffAssistant;
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
|
||||
session = new AgentSession({
|
||||
@@ -504,7 +485,7 @@ describe("AgentSession handoff", () => {
|
||||
|
||||
await session.prompt("Trigger threshold handoff");
|
||||
|
||||
expect(streamCallCount).toBe(2);
|
||||
expect(mock.calls).toHaveLength(2);
|
||||
const endEvents = events.filter(event => event.type === "auto_compaction_end");
|
||||
expect(endEvents).toHaveLength(1);
|
||||
expect(endEvents[0]).toMatchObject({ type: "auto_compaction_end", action: "handoff", aborted: false });
|
||||
@@ -581,8 +562,9 @@ describe("AgentSession handoff", () => {
|
||||
});
|
||||
vi.spyOn(extensionRunner, "emit").mockResolvedValue(undefined);
|
||||
|
||||
const observedSystemPrompts: string[] = [];
|
||||
let streamCallCount = 0;
|
||||
const mock = createMockModel({
|
||||
responses: [{ content: ["normal response"] }, { content: ["## Goal\nContinue from here"] }],
|
||||
});
|
||||
const agent = new Agent({
|
||||
getApiKey: () => "test-key",
|
||||
initialState: {
|
||||
@@ -591,38 +573,7 @@ describe("AgentSession handoff", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: (_model, context) => {
|
||||
observedSystemPrompts.push(context.systemPrompt?.join("\n\n") ?? "");
|
||||
streamCallCount++;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: streamCallCount === 1 ? "normal response" : "## Goal\nContinue from here",
|
||||
},
|
||||
],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
stopReason: "stop",
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
|
||||
session = new AgentSession({
|
||||
@@ -659,7 +610,7 @@ describe("AgentSession handoff", () => {
|
||||
await session.handoff();
|
||||
|
||||
expect(emitBeforeAgentStart).toHaveBeenCalledTimes(1);
|
||||
expect(observedSystemPrompts).toEqual(["Hook override", "Test"]);
|
||||
expect(mock.calls.map(c => c.context.systemPrompt?.join("\n\n") ?? "")).toEqual(["Hook override", "Test"]);
|
||||
});
|
||||
|
||||
it("saves auto-handoff document to disk when enabled", async () => {
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
import { afterEach, beforeEach, describe, expect, it } from "bun:test";
|
||||
import * as path from "node:path";
|
||||
import { Agent } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AssistantMessage, getBundledModel, type Model } from "@oh-my-pi/pi-ai";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { type AssistantMessage, getBundledModel } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
||||
@@ -10,32 +10,6 @@ import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
||||
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
||||
import { TempDir } from "@oh-my-pi/pi-utils";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
function createAssistantMessage(
|
||||
model: Model,
|
||||
options: { text?: string; stopReason: "stop" | "error" | "aborted"; errorMessage?: string },
|
||||
): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: options.text ? [{ type: "text", text: options.text }] : [],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: options.stopReason,
|
||||
errorMessage: options.errorMessage,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
function lastAgentMessage(session: AgentSession): AssistantMessage {
|
||||
const message = session.agent.state.messages.at(-1);
|
||||
if (!message || message.role !== "assistant") {
|
||||
@@ -70,7 +44,12 @@ describe("AgentSession manual retry", () => {
|
||||
throw new Error("Expected bundled Anthropic test model to exist");
|
||||
}
|
||||
|
||||
let streamCalls = 0;
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{ throw: "manual retry test failure" },
|
||||
{ content: ["recovered after manual retry"], stopReason: "stop" },
|
||||
],
|
||||
});
|
||||
const agent = new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
initialState: {
|
||||
@@ -79,32 +58,7 @@ describe("AgentSession manual retry", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: requestedModel => {
|
||||
streamCalls += 1;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (streamCalls === 1) {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
stopReason: "error",
|
||||
errorMessage: "manual retry test failure",
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
return;
|
||||
}
|
||||
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
text: "recovered after manual retry",
|
||||
stopReason: "stop",
|
||||
});
|
||||
stream.push({
|
||||
type: "start",
|
||||
partial: createAssistantMessage(requestedModel, { text: "", stopReason: "stop" }),
|
||||
});
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
session = new AgentSession({
|
||||
agent,
|
||||
@@ -121,7 +75,7 @@ describe("AgentSession manual retry", () => {
|
||||
await expect(session.retry()).resolves.toBe(true);
|
||||
await session.waitForIdle();
|
||||
|
||||
expect(streamCalls).toBe(2);
|
||||
expect(mock.calls.length).toBe(2);
|
||||
expect(lastAgentMessage(session).stopReason).toBe("stop");
|
||||
expect(lastAgentMessage(session).content).toContainEqual({ type: "text", text: "recovered after manual retry" });
|
||||
});
|
||||
@@ -132,7 +86,9 @@ describe("AgentSession manual retry", () => {
|
||||
throw new Error("Expected bundled Anthropic test model to exist");
|
||||
}
|
||||
|
||||
let streamCalls = 0;
|
||||
const mock = createMockModel({
|
||||
responses: [{ content: ["already done"], stopReason: "stop" }],
|
||||
});
|
||||
const agent = new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
initialState: {
|
||||
@@ -141,19 +97,7 @@ describe("AgentSession manual retry", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: requestedModel => {
|
||||
streamCalls += 1;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage(requestedModel, { text: "already done", stopReason: "stop" });
|
||||
stream.push({
|
||||
type: "start",
|
||||
partial: createAssistantMessage(requestedModel, { text: "", stopReason: "stop" }),
|
||||
});
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
session = new AgentSession({
|
||||
agent,
|
||||
@@ -167,7 +111,7 @@ describe("AgentSession manual retry", () => {
|
||||
await session.waitForIdle();
|
||||
|
||||
await expect(session.retry()).resolves.toBe(false);
|
||||
expect(streamCalls).toBe(1);
|
||||
expect(mock.calls.length).toBe(1);
|
||||
expect(lastAgentMessage(session).content).toContainEqual({ type: "text", text: "already done" });
|
||||
});
|
||||
});
|
||||
|
||||
@@ -4,29 +4,26 @@ import * as os from "node:os";
|
||||
import * as path from "node:path";
|
||||
import { Agent } from "@oh-my-pi/pi-agent-core";
|
||||
import { getBundledModel } from "@oh-my-pi/pi-ai";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { createMockModel, type MockModelHandle } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
||||
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
||||
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
||||
import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools";
|
||||
import { queueResolveHandler } from "@oh-my-pi/pi-coding-agent/tools/resolve";
|
||||
import { buildNamedToolChoice } from "@oh-my-pi/pi-coding-agent/utils/tool-choice";
|
||||
import { Snowflake } from "@oh-my-pi/pi-utils";
|
||||
import { createAssistantMessage } from "./helpers/agent-session-setup";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
describe("AgentSession resolve reminder", () => {
|
||||
let session: AgentSession;
|
||||
let tempDir: string;
|
||||
let streamCallCount = 0;
|
||||
let mock: MockModelHandle;
|
||||
let authStorage: AuthStorage | undefined;
|
||||
|
||||
beforeEach(async () => {
|
||||
tempDir = path.join(os.tmpdir(), `pi-resolve-reminder-test-${Snowflake.next()}`);
|
||||
fs.mkdirSync(tempDir, { recursive: true });
|
||||
streamCallCount = 0;
|
||||
|
||||
const model = getBundledModel("anthropic", "claude-sonnet-4-5");
|
||||
if (!model) {
|
||||
@@ -37,16 +34,9 @@ describe("AgentSession resolve reminder", () => {
|
||||
authStorage.setRuntimeApiKey("anthropic", "test-key");
|
||||
const modelRegistry = new ModelRegistry(authStorage, path.join(tempDir, "models.yml"));
|
||||
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
model,
|
||||
systemPrompt: ["Test"],
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: () => {
|
||||
streamCallCount += 1;
|
||||
if (streamCallCount === 1) {
|
||||
mock = createMockModel({
|
||||
handler: () => {
|
||||
if (mock.calls.length === 1) {
|
||||
queueResolveHandler(
|
||||
{
|
||||
getToolChoiceQueue: () => session.toolChoiceQueue,
|
||||
@@ -61,7 +51,7 @@ describe("AgentSession resolve reminder", () => {
|
||||
attribution: "agent",
|
||||
timestamp: Date.now(),
|
||||
}),
|
||||
} as any,
|
||||
} as unknown as ToolSession,
|
||||
{
|
||||
label: "AST Edit: 1 replacement in 1 file",
|
||||
sourceToolName: "ast_edit",
|
||||
@@ -69,15 +59,20 @@ describe("AgentSession resolve reminder", () => {
|
||||
},
|
||||
);
|
||||
}
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("Done") });
|
||||
});
|
||||
return stream;
|
||||
return { content: ["Done"] };
|
||||
},
|
||||
});
|
||||
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
model,
|
||||
systemPrompt: ["Test"],
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: mock.stream,
|
||||
});
|
||||
|
||||
session = new AgentSession({
|
||||
agent,
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
@@ -99,7 +94,7 @@ describe("AgentSession resolve reminder", () => {
|
||||
it("forces an immediate steering turn and injects resolve reminder before second assistant response", async () => {
|
||||
await session.prompt("run preview");
|
||||
|
||||
expect(streamCallCount).toBe(2);
|
||||
expect(mock.calls).toHaveLength(2);
|
||||
|
||||
const messages = session.agent.state.messages;
|
||||
const assistantIndices = messages
|
||||
|
||||
@@ -2,7 +2,7 @@ import { afterEach, beforeEach, describe, expect, it } from "bun:test";
|
||||
import * as path from "node:path";
|
||||
import { Agent } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AssistantMessage, Effort, getBundledModel, type Model, writeModelCache } from "@oh-my-pi/pi-ai";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import { AgentSession, type AgentSessionEvent } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
||||
@@ -10,32 +10,6 @@ import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
||||
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
||||
import { TempDir } from "@oh-my-pi/pi-utils";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
function createAssistantMessage(
|
||||
model: Model,
|
||||
options: { text?: string; stopReason: "stop" | "error"; errorMessage?: string },
|
||||
): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: options.text ? [{ type: "text", text: options.text }] : [],
|
||||
api: model.api,
|
||||
provider: model.provider,
|
||||
model: model.id,
|
||||
usage: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
},
|
||||
stopReason: options.stopReason,
|
||||
errorMessage: options.errorMessage,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
type AutoRetryStartEvent = Extract<AgentSessionEvent, { type: "auto_retry_start" }>;
|
||||
type AutoRetryEndEvent = Extract<AgentSessionEvent, { type: "auto_retry_end" }>;
|
||||
|
||||
@@ -65,6 +39,7 @@ function getLastAssistantMessage(session: AgentSession): AssistantMessage {
|
||||
}
|
||||
|
||||
function createFallbackAgent(primaryModel: Model, requestedModels: string[]): Agent {
|
||||
const mock = createMockModel();
|
||||
let primaryAttempts = 0;
|
||||
return new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
@@ -74,28 +49,15 @@ function createFallbackAgent(primaryModel: Model, requestedModels: string[]): Ag
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: model => {
|
||||
streamFn: (model, context, options) => {
|
||||
requestedModels.push(`${model.provider}/${model.id}`);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (model.provider === primaryModel.provider && model.id === primaryModel.id && primaryAttempts === 0) {
|
||||
primaryAttempts += 1;
|
||||
const message = createAssistantMessage(model, {
|
||||
stopReason: "error",
|
||||
errorMessage: "rate limit exceeded retry-after-ms=200",
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
return;
|
||||
}
|
||||
const message = createAssistantMessage(model, {
|
||||
text: `ok:${model.provider}/${model.id}`,
|
||||
stopReason: "stop",
|
||||
});
|
||||
stream.push({ type: "start", partial: createAssistantMessage(model, { text: "", stopReason: "stop" }) });
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
if (model.provider === primaryModel.provider && model.id === primaryModel.id && primaryAttempts === 0) {
|
||||
primaryAttempts += 1;
|
||||
mock.push({ throw: "rate limit exceeded retry-after-ms=200" });
|
||||
} else {
|
||||
mock.push({ content: [`ok:${model.provider}/${model.id}`] });
|
||||
}
|
||||
return mock.stream(model, context, options);
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -137,6 +99,7 @@ describe("AgentSession retry fallback", () => {
|
||||
const fallbackAppliedEvents: Array<Extract<AgentSessionEvent, { type: "retry_fallback_applied" }>> = [];
|
||||
const fallbackSucceededEvents: Array<Extract<AgentSessionEvent, { type: "retry_fallback_succeeded" }>> = [];
|
||||
|
||||
const mock = createMockModel();
|
||||
const agent = new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
initialState: {
|
||||
@@ -145,43 +108,18 @@ describe("AgentSession retry fallback", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: model => {
|
||||
streamFn: (model, context, options) => {
|
||||
requestedModels.push(`${model.provider}/${model.id}`);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (model.provider === primaryModel.provider && model.id === primaryModel.id) {
|
||||
const message = createAssistantMessage(model, {
|
||||
stopReason: "error",
|
||||
errorMessage: "overloaded_error: provider returned error 503",
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
return;
|
||||
}
|
||||
if (model.provider === firstFallback.provider && model.id === firstFallback.id) {
|
||||
const message = createAssistantMessage(model, {
|
||||
stopReason: "error",
|
||||
errorMessage: "service unavailable: 503 overloaded",
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
return;
|
||||
}
|
||||
if (model.provider === secondFallback.provider && model.id === secondFallback.id) {
|
||||
const message = createAssistantMessage(model, {
|
||||
text: "Recovered on second fallback",
|
||||
stopReason: "stop",
|
||||
});
|
||||
stream.push({
|
||||
type: "start",
|
||||
partial: createAssistantMessage(model, { text: "", stopReason: "stop" }),
|
||||
});
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
return;
|
||||
}
|
||||
if (model.provider === primaryModel.provider && model.id === primaryModel.id) {
|
||||
mock.push({ throw: "overloaded_error: provider returned error 503" });
|
||||
} else if (model.provider === firstFallback.provider && model.id === firstFallback.id) {
|
||||
mock.push({ throw: "service unavailable: 503 overloaded" });
|
||||
} else if (model.provider === secondFallback.provider && model.id === secondFallback.id) {
|
||||
mock.push({ content: ["Recovered on second fallback"] });
|
||||
} else {
|
||||
throw new Error(`Unexpected model requested during retry fallback test: ${model.provider}/${model.id}`);
|
||||
});
|
||||
return stream;
|
||||
}
|
||||
return mock.stream(model, context, options);
|
||||
},
|
||||
});
|
||||
|
||||
@@ -263,8 +201,10 @@ describe("AgentSession retry fallback", () => {
|
||||
|
||||
const timeoutMessage = "OpenAI responses stream timed out while waiting for the first event";
|
||||
const requestedModels: string[] = [];
|
||||
let attemptCount = 0;
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [{ throw: timeoutMessage }, { content: ["Recovered after OpenAI timeout"] }],
|
||||
});
|
||||
const agent = new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
initialState: {
|
||||
@@ -273,35 +213,9 @@ describe("AgentSession retry fallback", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: requestedModel => {
|
||||
streamFn: (requestedModel, context, options) => {
|
||||
requestedModels.push(`${requestedModel.provider}/${requestedModel.id}`);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
attemptCount += 1;
|
||||
if (attemptCount === 1) {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
stopReason: "error",
|
||||
errorMessage: timeoutMessage,
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
return;
|
||||
}
|
||||
if (attemptCount === 2) {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
text: "Recovered after OpenAI timeout",
|
||||
stopReason: "stop",
|
||||
});
|
||||
stream.push({
|
||||
type: "start",
|
||||
partial: createAssistantMessage(requestedModel, { text: "", stopReason: "stop" }),
|
||||
});
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
return;
|
||||
}
|
||||
throw new Error(`Unexpected retry attempt in timeout test: ${attemptCount}`);
|
||||
});
|
||||
return stream;
|
||||
return mock.stream(requestedModel, context, options);
|
||||
},
|
||||
});
|
||||
|
||||
@@ -345,8 +259,10 @@ describe("AgentSession retry fallback", () => {
|
||||
|
||||
const stallMessage = "Provider stream stalled while waiting for the next event";
|
||||
const requestedModels: string[] = [];
|
||||
let attemptCount = 0;
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [{ throw: stallMessage }, { content: ["Recovered after stream stall"] }],
|
||||
});
|
||||
const agent = new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
initialState: {
|
||||
@@ -355,35 +271,9 @@ describe("AgentSession retry fallback", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: requestedModel => {
|
||||
streamFn: (requestedModel, context, options) => {
|
||||
requestedModels.push(`${requestedModel.provider}/${requestedModel.id}`);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
attemptCount += 1;
|
||||
if (attemptCount === 1) {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
stopReason: "error",
|
||||
errorMessage: stallMessage,
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
return;
|
||||
}
|
||||
if (attemptCount === 2) {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
text: "Recovered after stream stall",
|
||||
stopReason: "stop",
|
||||
});
|
||||
stream.push({
|
||||
type: "start",
|
||||
partial: createAssistantMessage(requestedModel, { text: "", stopReason: "stop" }),
|
||||
});
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
return;
|
||||
}
|
||||
throw new Error(`Unexpected retry attempt in stream stall test: ${attemptCount}`);
|
||||
});
|
||||
return stream;
|
||||
return mock.stream(requestedModel, context, options);
|
||||
},
|
||||
});
|
||||
|
||||
@@ -428,8 +318,10 @@ describe("AgentSession retry fallback", () => {
|
||||
const processingError =
|
||||
"An error occurred while processing your request. You can retry your request, or contact us through our help center at help.openai.com if the error persists. Please include the request ID 4a4c6b73-a07c-4de0-aaaf-82560f9f626a in your message.";
|
||||
const requestedModels: string[] = [];
|
||||
let attemptCount = 0;
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [{ throw: processingError }, { content: ["Recovered after OpenAI processing error"] }],
|
||||
});
|
||||
const agent = new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
initialState: {
|
||||
@@ -438,35 +330,9 @@ describe("AgentSession retry fallback", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: requestedModel => {
|
||||
streamFn: (requestedModel, context, options) => {
|
||||
requestedModels.push(`${requestedModel.provider}/${requestedModel.id}`);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
attemptCount += 1;
|
||||
if (attemptCount === 1) {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
stopReason: "error",
|
||||
errorMessage: processingError,
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
return;
|
||||
}
|
||||
if (attemptCount === 2) {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
text: "Recovered after OpenAI processing error",
|
||||
stopReason: "stop",
|
||||
});
|
||||
stream.push({
|
||||
type: "start",
|
||||
partial: createAssistantMessage(requestedModel, { text: "", stopReason: "stop" }),
|
||||
});
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
return;
|
||||
}
|
||||
throw new Error(`Unexpected retry attempt in processing-error test: ${attemptCount}`);
|
||||
});
|
||||
return stream;
|
||||
return mock.stream(requestedModel, context, options);
|
||||
},
|
||||
});
|
||||
|
||||
@@ -513,8 +379,10 @@ describe("AgentSession retry fallback", () => {
|
||||
|
||||
const envelopeError = "Anthropic stream envelope error: received content_block_start before message_start";
|
||||
const requestedModels: string[] = [];
|
||||
let attemptCount = 0;
|
||||
|
||||
const mock = createMockModel({
|
||||
responses: [{ throw: envelopeError }, { content: ["Recovered after Anthropic envelope retry"] }],
|
||||
});
|
||||
const agent = new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
initialState: {
|
||||
@@ -523,35 +391,9 @@ describe("AgentSession retry fallback", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: requestedModel => {
|
||||
streamFn: (requestedModel, context, options) => {
|
||||
requestedModels.push(`${requestedModel.provider}/${requestedModel.id}`);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
attemptCount += 1;
|
||||
if (attemptCount === 1) {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
stopReason: "error",
|
||||
errorMessage: envelopeError,
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
return;
|
||||
}
|
||||
if (attemptCount === 2) {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
text: "Recovered after Anthropic envelope retry",
|
||||
stopReason: "stop",
|
||||
});
|
||||
stream.push({
|
||||
type: "start",
|
||||
partial: createAssistantMessage(requestedModel, { text: "", stopReason: "stop" }),
|
||||
});
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
return;
|
||||
}
|
||||
throw new Error(`Unexpected retry attempt in Anthropic envelope test: ${attemptCount}`);
|
||||
});
|
||||
return stream;
|
||||
return mock.stream(requestedModel, context, options);
|
||||
},
|
||||
});
|
||||
|
||||
@@ -599,6 +441,7 @@ describe("AgentSession retry fallback", () => {
|
||||
const fallbackAppliedEvents: Array<Extract<AgentSessionEvent, { type: "retry_fallback_applied" }>> = [];
|
||||
const fallbackSucceededEvents: Array<Extract<AgentSessionEvent, { type: "retry_fallback_succeeded" }>> = [];
|
||||
|
||||
const mock = createMockModel({ handler: () => ({ throw: envelopeError }) });
|
||||
const agent = new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
initialState: {
|
||||
@@ -607,18 +450,9 @@ describe("AgentSession retry fallback", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: requestedModel => {
|
||||
streamFn: (requestedModel, context, options) => {
|
||||
requestedModels.push(`${requestedModel.provider}/${requestedModel.id}`);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
stopReason: "error",
|
||||
errorMessage: envelopeError,
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
});
|
||||
return stream;
|
||||
return mock.stream(requestedModel, context, options);
|
||||
},
|
||||
});
|
||||
|
||||
@@ -668,6 +502,7 @@ describe("AgentSession retry fallback", () => {
|
||||
}
|
||||
|
||||
const requestedModels: string[] = [];
|
||||
const mock = createMockModel({ handler: () => ({ throw: "Request was aborted." }) });
|
||||
const agent = new Agent({
|
||||
getApiKey: provider => `${provider}-test-key`,
|
||||
initialState: {
|
||||
@@ -676,18 +511,9 @@ describe("AgentSession retry fallback", () => {
|
||||
tools: [],
|
||||
messages: [],
|
||||
},
|
||||
streamFn: requestedModel => {
|
||||
streamFn: (requestedModel, context, options) => {
|
||||
requestedModels.push(`${requestedModel.provider}/${requestedModel.id}`);
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage(requestedModel, {
|
||||
stopReason: "error",
|
||||
errorMessage: "Request was aborted.",
|
||||
});
|
||||
stream.push({ type: "start", partial: message });
|
||||
stream.push({ type: "error", reason: "error", error: message });
|
||||
});
|
||||
return stream;
|
||||
return mock.stream(requestedModel, context, options);
|
||||
},
|
||||
});
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ import * as os from "node:os";
|
||||
import * as path from "node:path";
|
||||
import { Agent, type AgentTool } from "@oh-my-pi/pi-agent-core";
|
||||
import { type AssistantMessage, getBundledModel, type StopReason, type ToolCall } from "@oh-my-pi/pi-ai";
|
||||
import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock";
|
||||
import { AssistantMessageEventStream } from "@oh-my-pi/pi-ai/utils/event-stream";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
@@ -30,8 +31,6 @@ import { ToolError } from "@oh-my-pi/pi-coding-agent/tools/tool-errors";
|
||||
import { Snowflake } from "@oh-my-pi/pi-utils";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
function createAssistantMessage(content: AssistantMessage["content"], stopReason: StopReason): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
@@ -116,7 +115,7 @@ function streamForSingleToolCall(
|
||||
let callIndex = 0;
|
||||
return (_model, _context, options) => {
|
||||
abortSignalRef.current = options?.signal;
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
const toolCallId = "call_edit_regression";
|
||||
let aborted = false;
|
||||
|
||||
@@ -327,31 +326,17 @@ it("agent-loop propagates explicit isError from a tool result to the wire", asyn
|
||||
},
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const streamFn: Agent["streamFn"] = (_model, _context, _options) => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
const toolCall = createToolCall("call_self_error", {});
|
||||
const finalMessage = createAssistantMessage([toolCall], "toolUse");
|
||||
stream.push({ type: "start", partial: createAssistantMessage([], "stop") });
|
||||
stream.push({
|
||||
type: "toolcall_start",
|
||||
contentIndex: 0,
|
||||
partial: createAssistantMessage([toolCall], "stop"),
|
||||
});
|
||||
stream.push({ type: "toolcall_end", contentIndex: 0, toolCall, partial: finalMessage });
|
||||
stream.push({ type: "done", reason: "toolUse", message: finalMessage });
|
||||
} else {
|
||||
const finalMessage = createAssistantMessage([{ type: "text", text: "done" }], "stop");
|
||||
stream.push({ type: "done", reason: "stop", message: finalMessage });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
const mock = createMockModel({
|
||||
responses: [
|
||||
{
|
||||
content: [{ type: "toolCall", id: "call_self_error", name: "edit", arguments: {} }],
|
||||
stopReason: "toolUse",
|
||||
},
|
||||
{ content: ["done"], stopReason: "stop" },
|
||||
],
|
||||
});
|
||||
|
||||
const { session, authStorage } = await createSessionWith(tempDir, streamFn, errorTool);
|
||||
const { session, authStorage } = await createSessionWith(tempDir, mock.stream, errorTool);
|
||||
try {
|
||||
await session.prompt("trigger self-error");
|
||||
|
||||
|
||||
@@ -19,8 +19,6 @@ import { ToolError } from "@oh-my-pi/pi-coding-agent/tools/tool-errors";
|
||||
import { Snowflake } from "@oh-my-pi/pi-utils";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
|
||||
class MockAssistantStream extends AssistantMessageEventStream {}
|
||||
|
||||
function createAssistantMessage(content: AssistantMessage["content"], stopReason: StopReason): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
@@ -143,7 +141,7 @@ function createStreamForDiff(
|
||||
let callIndex = 0;
|
||||
return (_model, _context, options) => {
|
||||
abortSignalRef.current = options?.signal;
|
||||
const stream = new MockAssistantStream();
|
||||
const stream = new AssistantMessageEventStream();
|
||||
const toolCallId = "call_edit_1";
|
||||
let diffSoFar = "";
|
||||
let aborted = false;
|
||||
|
||||
Reference in New Issue
Block a user