chore: drop non-code junk (pi-mono-reference, runtime state, root duplicates)
Remove files that should never have been tracked: - pi-mono-reference/: 688-file copy of an upstream monorepo, zero imports, used only as offline human reference - build/electron-builder.runtime.json: environment-specific build artifact with hardcoded local paths - state/cancel-signal-state.json: runtime state written by the app - prompt.txt, logo.svg (root): exact duplicates of the canonical copies under .planning/archive/ and assets/ Add state/ to .gitignore so future runtime state stays out.
This commit is contained in:
ZacharyZhang-NY committed
1 parent
34ca5b77f1
commit
1f7230fe87
691 files changed
+3
-190434
No files matched your search
@@ -1,758 +0,0 @@
|
||||
import {
|
||||
type AssistantMessage,
|
||||
type AssistantMessageEvent,
|
||||
EventStream,
|
||||
type Message,
|
||||
type Model,
|
||||
type UserMessage,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import { Type } from "@sinclair/typebox";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { agentLoop, agentLoopContinue } from "../src/agent-loop.js";
|
||||
import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "../src/types.js";
|
||||
|
||||
// Mock stream for testing - mimics MockAssistantStream
|
||||
class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
|
||||
constructor() {
|
||||
super(
|
||||
(event) => event.type === "done" || event.type === "error",
|
||||
(event) => {
|
||||
if (event.type === "done") return event.message;
|
||||
if (event.type === "error") return event.error;
|
||||
throw new Error("Unexpected event type");
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
function createUsage() {
|
||||
return {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
totalTokens: 0,
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||
};
|
||||
}
|
||||
|
||||
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 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(),
|
||||
};
|
||||
}
|
||||
|
||||
function createUserMessage(text: string): UserMessage {
|
||||
return {
|
||||
role: "user",
|
||||
content: text,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
}
|
||||
|
||||
// Simple identity converter for tests - just passes through standard messages
|
||||
function identityConverter(messages: AgentMessage[]): Message[] {
|
||||
return messages.filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[];
|
||||
}
|
||||
|
||||
describe("agentLoop with AgentMessage", () => {
|
||||
it("should emit events with AgentMessage types", async () => {
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "You are helpful.",
|
||||
messages: [],
|
||||
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 events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
const messages = await stream.result();
|
||||
|
||||
// Should have user message and assistant message
|
||||
expect(messages.length).toBe(2);
|
||||
expect(messages[0].role).toBe("user");
|
||||
expect(messages[1].role).toBe("assistant");
|
||||
|
||||
// Verify event sequence
|
||||
const eventTypes = events.map((e) => e.type);
|
||||
expect(eventTypes).toContain("agent_start");
|
||||
expect(eventTypes).toContain("turn_start");
|
||||
expect(eventTypes).toContain("message_start");
|
||||
expect(eventTypes).toContain("message_end");
|
||||
expect(eventTypes).toContain("turn_end");
|
||||
expect(eventTypes).toContain("agent_end");
|
||||
});
|
||||
|
||||
it("should handle custom message types via convertToLlm", async () => {
|
||||
// Create a custom message type
|
||||
interface CustomNotification {
|
||||
role: "notification";
|
||||
text: string;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
const notification: CustomNotification = {
|
||||
role: "notification",
|
||||
text: "This is a notification",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "You are helpful.",
|
||||
messages: [notification as unknown as AgentMessage], // Custom message in context
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("Hello");
|
||||
|
||||
let convertedMessages: Message[] = [];
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: (messages) => {
|
||||
// Filter out notifications, convert rest
|
||||
convertedMessages = messages
|
||||
.filter((m) => (m as { role: string }).role !== "notification")
|
||||
.filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[];
|
||||
return convertedMessages;
|
||||
},
|
||||
};
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
// The notification should have been filtered out in convertToLlm
|
||||
expect(convertedMessages.length).toBe(1); // Only user message
|
||||
expect(convertedMessages[0].role).toBe("user");
|
||||
});
|
||||
|
||||
it("should apply transformContext before convertToLlm", async () => {
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "You are helpful.",
|
||||
messages: [
|
||||
createUserMessage("old message 1"),
|
||||
createAssistantMessage([{ type: "text", text: "old response 1" }]),
|
||||
createUserMessage("old message 2"),
|
||||
createAssistantMessage([{ type: "text", text: "old response 2" }]),
|
||||
],
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("new message");
|
||||
|
||||
let transformedMessages: AgentMessage[] = [];
|
||||
let convertedMessages: Message[] = [];
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
transformContext: async (messages) => {
|
||||
// Keep only last 2 messages (prune old ones)
|
||||
transformedMessages = messages.slice(-2);
|
||||
return transformedMessages;
|
||||
},
|
||||
convertToLlm: (messages) => {
|
||||
convertedMessages = messages.filter(
|
||||
(m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult",
|
||||
) as Message[];
|
||||
return convertedMessages;
|
||||
},
|
||||
};
|
||||
|
||||
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);
|
||||
|
||||
for await (const _ of stream) {
|
||||
// consume
|
||||
}
|
||||
|
||||
// transformContext should have been called first, keeping only last 2
|
||||
expect(transformedMessages.length).toBe(2);
|
||||
// Then convertToLlm receives the pruned messages
|
||||
expect(convertedMessages.length).toBe(2);
|
||||
});
|
||||
|
||||
it("should handle tool calls and results", async () => {
|
||||
const toolSchema = Type.Object({ value: Type.String() });
|
||||
const executed: string[] = [];
|
||||
const tool: AgentTool<typeof toolSchema, { value: string }> = {
|
||||
name: "echo",
|
||||
label: "Echo",
|
||||
description: "Echo tool",
|
||||
parameters: toolSchema,
|
||||
async execute(_toolCallId, params) {
|
||||
executed.push(params.value);
|
||||
return {
|
||||
content: [{ type: "text", text: `echoed: ${params.value}` }],
|
||||
details: { value: params.value },
|
||||
};
|
||||
},
|
||||
};
|
||||
|
||||
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 events: AgentEvent[] = [];
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, streamFn);
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
// Tool should have been executed
|
||||
expect(executed).toEqual(["hello"]);
|
||||
|
||||
// Should have tool execution events
|
||||
const toolStart = events.find((e) => e.type === "tool_execution_start");
|
||||
const toolEnd = events.find((e) => e.type === "tool_execution_end");
|
||||
expect(toolStart).toBeDefined();
|
||||
expect(toolEnd).toBeDefined();
|
||||
if (toolEnd?.type === "tool_execution_end") {
|
||||
expect(toolEnd.isError).toBe(false);
|
||||
}
|
||||
});
|
||||
|
||||
it("should execute mutated beforeToolCall args without revalidation", async () => {
|
||||
const toolSchema = Type.Object({ value: Type.String() });
|
||||
const executed: Array<string | number> = [];
|
||||
const tool: AgentTool<typeof toolSchema, { value: string | number }> = {
|
||||
name: "echo",
|
||||
label: "Echo",
|
||||
description: "Echo tool",
|
||||
parameters: toolSchema,
|
||||
async execute(_toolCallId, params) {
|
||||
executed.push(params.value as string | number);
|
||||
return {
|
||||
content: [{ type: "text", text: `echoed: ${String(params.value)}` }],
|
||||
details: { value: params.value as string | number },
|
||||
};
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "",
|
||||
messages: [],
|
||||
tools: [tool],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("echo something");
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
beforeToolCall: async ({ args }) => {
|
||||
const mutableArgs = args as { value: string | number };
|
||||
mutableArgs.value = 123;
|
||||
return undefined;
|
||||
},
|
||||
};
|
||||
|
||||
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) {
|
||||
// consume
|
||||
}
|
||||
|
||||
expect(executed).toEqual([123]);
|
||||
});
|
||||
|
||||
it("should prepare tool arguments for validation", async () => {
|
||||
const replaceSchema = Type.Object({ oldText: Type.String(), newText: Type.String() });
|
||||
const toolSchema = Type.Object({ edits: Type.Array(replaceSchema) });
|
||||
const executed: Array<Array<{ oldText: string; newText: string }>> = [];
|
||||
const tool: AgentTool<typeof toolSchema, { count: number }> = {
|
||||
name: "edit",
|
||||
label: "Edit",
|
||||
description: "Edit tool",
|
||||
parameters: toolSchema,
|
||||
prepareArguments(args) {
|
||||
if (!args || typeof args !== "object") {
|
||||
return args as { edits: { oldText: string; newText: string }[] };
|
||||
}
|
||||
const input = args as {
|
||||
edits?: Array<{ oldText: string; newText: string }>;
|
||||
oldText?: string;
|
||||
newText?: string;
|
||||
};
|
||||
if (typeof input.oldText !== "string" || typeof input.newText !== "string") {
|
||||
return args as { edits: { oldText: string; newText: string }[] };
|
||||
}
|
||||
return {
|
||||
edits: [...(input.edits ?? []), { oldText: input.oldText, newText: input.newText }],
|
||||
};
|
||||
},
|
||||
async execute(_toolCallId, params) {
|
||||
executed.push(params.edits);
|
||||
return {
|
||||
content: [{ type: "text", text: `edited ${params.edits.length}` }],
|
||||
details: { count: params.edits.length },
|
||||
};
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "",
|
||||
messages: [],
|
||||
tools: [tool],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("edit something");
|
||||
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: "edit",
|
||||
arguments: { oldText: "before", newText: "after" },
|
||||
},
|
||||
],
|
||||
"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) {
|
||||
// consume
|
||||
}
|
||||
|
||||
expect(executed).toEqual([[{ oldText: "before", newText: "after" }]]);
|
||||
});
|
||||
|
||||
it("should execute tool calls in parallel and emit tool results in source order", async () => {
|
||||
const toolSchema = Type.Object({ value: Type.String() });
|
||||
let firstResolved = false;
|
||||
let parallelObserved = false;
|
||||
let releaseFirst: (() => void) | undefined;
|
||||
const firstDone = new Promise<void>((resolve) => {
|
||||
releaseFirst = resolve;
|
||||
});
|
||||
|
||||
const tool: AgentTool<typeof toolSchema, { value: string }> = {
|
||||
name: "echo",
|
||||
label: "Echo",
|
||||
description: "Echo tool",
|
||||
parameters: toolSchema,
|
||||
async execute(_toolCallId, params) {
|
||||
if (params.value === "first") {
|
||||
await firstDone;
|
||||
firstResolved = true;
|
||||
}
|
||||
if (params.value === "second" && !firstResolved) {
|
||||
parallelObserved = true;
|
||||
}
|
||||
return {
|
||||
content: [{ type: "text", text: `echoed: ${params.value}` }],
|
||||
details: { value: params.value },
|
||||
};
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "",
|
||||
messages: [],
|
||||
tools: [tool],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("echo both");
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
toolExecution: "parallel",
|
||||
};
|
||||
|
||||
let callIndex = 0;
|
||||
const stream = agentLoop([userPrompt], context, config, undefined, () => {
|
||||
const mockStream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
if (callIndex === 0) {
|
||||
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 });
|
||||
setTimeout(() => releaseFirst?.(), 20);
|
||||
} else {
|
||||
const message = createAssistantMessage([{ type: "text", text: "done" }]);
|
||||
mockStream.push({ type: "done", reason: "stop", message });
|
||||
}
|
||||
callIndex++;
|
||||
});
|
||||
return mockStream;
|
||||
});
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
const toolResultIds = events.flatMap((event) => {
|
||||
if (event.type !== "message_end" || event.message.role !== "toolResult") {
|
||||
return [];
|
||||
}
|
||||
return [event.message.toolCallId];
|
||||
});
|
||||
|
||||
expect(parallelObserved).toBe(true);
|
||||
expect(toolResultIds).toEqual(["tool-1", "tool-2"]);
|
||||
});
|
||||
|
||||
it("should inject queued messages after all tool calls complete", async () => {
|
||||
const toolSchema = Type.Object({ value: Type.String() });
|
||||
const executed: string[] = [];
|
||||
const tool: AgentTool<typeof toolSchema, { value: string }> = {
|
||||
name: "echo",
|
||||
label: "Echo",
|
||||
description: "Echo tool",
|
||||
parameters: toolSchema,
|
||||
async execute(_toolCallId, params) {
|
||||
executed.push(params.value);
|
||||
return {
|
||||
content: [{ type: "text", text: `ok:${params.value}` }],
|
||||
details: { value: params.value },
|
||||
};
|
||||
},
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "",
|
||||
messages: [],
|
||||
tools: [tool],
|
||||
};
|
||||
|
||||
const userPrompt: AgentMessage = createUserMessage("start");
|
||||
const queuedUserMessage: AgentMessage = createUserMessage("interrupt");
|
||||
|
||||
let queuedDelivered = false;
|
||||
let callIndex = 0;
|
||||
let sawInterruptInContext = false;
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: identityConverter,
|
||||
toolExecution: "sequential",
|
||||
getSteeringMessages: async () => {
|
||||
// Return steering message after tool execution has started.
|
||||
if (executed.length >= 1 && !queuedDelivered) {
|
||||
queuedDelivered = true;
|
||||
return [queuedUserMessage];
|
||||
}
|
||||
return [];
|
||||
},
|
||||
};
|
||||
|
||||
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;
|
||||
});
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
// Both tools should execute before steering is injected
|
||||
expect(executed).toEqual(["first", "second"]);
|
||||
|
||||
const toolEnds = events.filter(
|
||||
(e): e is Extract<AgentEvent, { type: "tool_execution_end" }> => e.type === "tool_execution_end",
|
||||
);
|
||||
expect(toolEnds.length).toBe(2);
|
||||
expect(toolEnds[0].isError).toBe(false);
|
||||
expect(toolEnds[1].isError).toBe(false);
|
||||
|
||||
// Queued message should appear in events after both tool result messages
|
||||
const eventSequence = events.flatMap((event) => {
|
||||
if (event.type !== "message_start") return [];
|
||||
if (event.message.role === "toolResult") return [`tool:${event.message.toolCallId}`];
|
||||
if (event.message.role === "user" && typeof event.message.content === "string") {
|
||||
return [event.message.content];
|
||||
}
|
||||
return [];
|
||||
});
|
||||
expect(eventSequence).toContain("interrupt");
|
||||
expect(eventSequence.indexOf("tool:tool-1")).toBeLessThan(eventSequence.indexOf("interrupt"));
|
||||
expect(eventSequence.indexOf("tool:tool-2")).toBeLessThan(eventSequence.indexOf("interrupt"));
|
||||
|
||||
// Interrupt message should be in context when second LLM call is made
|
||||
expect(sawInterruptInContext).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe("agentLoopContinue with AgentMessage", () => {
|
||||
it("should throw when context has no messages", () => {
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "You are helpful.",
|
||||
messages: [],
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
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 context: AgentContext = {
|
||||
systemPrompt: "You are helpful.",
|
||||
messages: [userMessage],
|
||||
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 events: AgentEvent[] = [];
|
||||
const stream = agentLoopContinue(context, config, undefined, streamFn);
|
||||
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
const messages = await stream.result();
|
||||
|
||||
// Should only return the new assistant message (not the existing user message)
|
||||
expect(messages.length).toBe(1);
|
||||
expect(messages[0].role).toBe("assistant");
|
||||
|
||||
// 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");
|
||||
});
|
||||
|
||||
it("should allow custom message types as last message (caller responsibility)", async () => {
|
||||
// Custom message that will be converted to user message by convertToLlm
|
||||
interface CustomMessage {
|
||||
role: "custom";
|
||||
text: string;
|
||||
timestamp: number;
|
||||
}
|
||||
|
||||
const customMessage: CustomMessage = {
|
||||
role: "custom",
|
||||
text: "Hook content",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const context: AgentContext = {
|
||||
systemPrompt: "You are helpful.",
|
||||
messages: [customMessage as unknown as AgentMessage],
|
||||
tools: [],
|
||||
};
|
||||
|
||||
const config: AgentLoopConfig = {
|
||||
model: createModel(),
|
||||
convertToLlm: (messages) => {
|
||||
// Convert custom to user message
|
||||
return messages
|
||||
.map((m) => {
|
||||
if ((m as any).role === "custom") {
|
||||
return {
|
||||
role: "user" as const,
|
||||
content: (m as any).text,
|
||||
timestamp: m.timestamp,
|
||||
};
|
||||
}
|
||||
return m;
|
||||
})
|
||||
.filter((m) => m.role === "user" || m.role === "assistant" || m.role === "toolResult") as Message[];
|
||||
},
|
||||
};
|
||||
|
||||
const streamFn = () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage([{ type: "text", text: "Response to custom message" }]);
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
};
|
||||
|
||||
// Should not throw - the custom message will be converted to user message
|
||||
const stream = agentLoopContinue(context, config, undefined, streamFn);
|
||||
|
||||
const events: AgentEvent[] = [];
|
||||
for await (const event of stream) {
|
||||
events.push(event);
|
||||
}
|
||||
|
||||
const messages = await stream.result();
|
||||
expect(messages.length).toBe(1);
|
||||
expect(messages[0].role).toBe("assistant");
|
||||
});
|
||||
});
|
||||
@@ -1,463 +0,0 @@
|
||||
import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@mariozechner/pi-ai";
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { Agent } from "../src/index.js";
|
||||
|
||||
// Mock stream that mimics AssistantMessageEventStream
|
||||
class MockAssistantStream extends EventStream<AssistantMessageEvent, AssistantMessage> {
|
||||
constructor() {
|
||||
super(
|
||||
(event) => event.type === "done" || event.type === "error",
|
||||
(event) => {
|
||||
if (event.type === "done") return event.message;
|
||||
if (event.type === "error") return event.error;
|
||||
throw new Error("Unexpected event type");
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
function createAssistantMessage(text: string): AssistantMessage {
|
||||
return {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text }],
|
||||
api: "openai-responses",
|
||||
provider: "openai",
|
||||
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(),
|
||||
};
|
||||
}
|
||||
|
||||
function createDeferred(): {
|
||||
promise: Promise<void>;
|
||||
resolve: () => void;
|
||||
} {
|
||||
let resolve = () => {};
|
||||
const promise = new Promise<void>((resolvePromise) => {
|
||||
resolve = resolvePromise;
|
||||
});
|
||||
return { promise, resolve };
|
||||
}
|
||||
|
||||
describe("Agent", () => {
|
||||
it("should create an agent instance with default state", () => {
|
||||
const agent = new Agent();
|
||||
|
||||
expect(agent.state).toBeDefined();
|
||||
expect(agent.state.systemPrompt).toBe("");
|
||||
expect(agent.state.model).toBeDefined();
|
||||
expect(agent.state.thinkingLevel).toBe("off");
|
||||
expect(agent.state.tools).toEqual([]);
|
||||
expect(agent.state.messages).toEqual([]);
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.streamingMessage).toBe(undefined);
|
||||
expect(agent.state.pendingToolCalls).toEqual(new Set());
|
||||
expect(agent.state.errorMessage).toBeUndefined();
|
||||
});
|
||||
|
||||
it("should create an agent instance with custom initial state", () => {
|
||||
const customModel = getModel("openai", "gpt-4o-mini");
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model: customModel,
|
||||
thinkingLevel: "low",
|
||||
},
|
||||
});
|
||||
|
||||
expect(agent.state.systemPrompt).toBe("You are a helpful assistant.");
|
||||
expect(agent.state.model).toBe(customModel);
|
||||
expect(agent.state.thinkingLevel).toBe("low");
|
||||
});
|
||||
|
||||
it("should subscribe to events", () => {
|
||||
const agent = new Agent();
|
||||
|
||||
let eventCount = 0;
|
||||
const unsubscribe = agent.subscribe((_event) => {
|
||||
eventCount++;
|
||||
});
|
||||
|
||||
// No initial event on subscribe
|
||||
expect(eventCount).toBe(0);
|
||||
|
||||
// State mutators don't emit events
|
||||
agent.state.systemPrompt = "Test prompt";
|
||||
expect(eventCount).toBe(0);
|
||||
expect(agent.state.systemPrompt).toBe("Test prompt");
|
||||
|
||||
// Unsubscribe should work
|
||||
unsubscribe();
|
||||
agent.state.systemPrompt = "Another prompt";
|
||||
expect(eventCount).toBe(0); // Should not increase
|
||||
});
|
||||
|
||||
it("should await async subscribers before prompt resolves", async () => {
|
||||
const barrier = createDeferred();
|
||||
const agent = new Agent({
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
let listenerFinished = false;
|
||||
agent.subscribe(async (event) => {
|
||||
if (event.type === "agent_end") {
|
||||
await barrier.promise;
|
||||
listenerFinished = true;
|
||||
}
|
||||
});
|
||||
|
||||
let promptResolved = false;
|
||||
const promptPromise = agent.prompt("hello").then(() => {
|
||||
promptResolved = true;
|
||||
});
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
expect(promptResolved).toBe(false);
|
||||
expect(listenerFinished).toBe(false);
|
||||
expect(agent.state.isStreaming).toBe(true);
|
||||
|
||||
barrier.resolve();
|
||||
await promptPromise;
|
||||
|
||||
expect(listenerFinished).toBe(true);
|
||||
expect(promptResolved).toBe(true);
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
});
|
||||
|
||||
it("waitForIdle should wait for async subscribers", async () => {
|
||||
const barrier = createDeferred();
|
||||
const agent = new Agent({
|
||||
streamFn: () => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "done", reason: "stop", message: createAssistantMessage("ok") });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
agent.subscribe(async (event) => {
|
||||
if (event.type === "message_end" && event.message.role === "assistant") {
|
||||
await barrier.promise;
|
||||
}
|
||||
});
|
||||
|
||||
const promptPromise = agent.prompt("hello");
|
||||
let idleResolved = false;
|
||||
const idlePromise = agent.waitForIdle().then(() => {
|
||||
idleResolved = true;
|
||||
});
|
||||
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
expect(idleResolved).toBe(false);
|
||||
expect(agent.state.isStreaming).toBe(true);
|
||||
|
||||
barrier.resolve();
|
||||
await Promise.all([promptPromise, idlePromise]);
|
||||
|
||||
expect(idleResolved).toBe(true);
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
});
|
||||
|
||||
it("should pass the active abort signal to subscribers", async () => {
|
||||
let receivedSignal: AbortSignal | undefined;
|
||||
const agent = new Agent({
|
||||
streamFn: (_model, _context, options) => {
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
const checkAbort = () => {
|
||||
if (options?.signal?.aborted) {
|
||||
stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
|
||||
} else {
|
||||
setTimeout(checkAbort, 5);
|
||||
}
|
||||
};
|
||||
checkAbort();
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
agent.subscribe((event, signal) => {
|
||||
if (event.type === "agent_start") {
|
||||
receivedSignal = signal;
|
||||
}
|
||||
});
|
||||
|
||||
const promptPromise = agent.prompt("hello");
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
|
||||
expect(receivedSignal).toBeDefined();
|
||||
expect(receivedSignal?.aborted).toBe(false);
|
||||
|
||||
agent.abort();
|
||||
await promptPromise;
|
||||
|
||||
expect(receivedSignal?.aborted).toBe(true);
|
||||
});
|
||||
|
||||
it("should update state with mutators", () => {
|
||||
const agent = new Agent();
|
||||
|
||||
// Test setSystemPrompt
|
||||
agent.state.systemPrompt = "Custom prompt";
|
||||
expect(agent.state.systemPrompt).toBe("Custom prompt");
|
||||
|
||||
// Test setModel
|
||||
const newModel = getModel("google", "gemini-2.5-flash");
|
||||
agent.state.model = newModel;
|
||||
expect(agent.state.model).toBe(newModel);
|
||||
|
||||
// Test setThinkingLevel
|
||||
agent.state.thinkingLevel = "high";
|
||||
expect(agent.state.thinkingLevel).toBe("high");
|
||||
|
||||
// Test setTools
|
||||
const tools = [{ name: "test", description: "test tool" } as any];
|
||||
agent.state.tools = tools;
|
||||
expect(agent.state.tools).toEqual(tools);
|
||||
expect(agent.state.tools).not.toBe(tools); // Should be a copy
|
||||
|
||||
// Test replaceMessages
|
||||
const messages = [{ role: "user" as const, content: "Hello", timestamp: Date.now() }];
|
||||
agent.state.messages = messages;
|
||||
expect(agent.state.messages).toEqual(messages);
|
||||
expect(agent.state.messages).not.toBe(messages); // Should be a copy
|
||||
|
||||
// Test appendMessage
|
||||
const newMessage = { role: "assistant" as const, content: [{ type: "text" as const, text: "Hi" }] };
|
||||
agent.state.messages.push(newMessage as any);
|
||||
expect(agent.state.messages).toHaveLength(2);
|
||||
expect(agent.state.messages[1]).toBe(newMessage);
|
||||
|
||||
// Test clearMessages
|
||||
agent.state.messages = [];
|
||||
expect(agent.state.messages).toEqual([]);
|
||||
});
|
||||
|
||||
it("should support steering message queue", async () => {
|
||||
const agent = new Agent();
|
||||
|
||||
const message = { role: "user" as const, content: "Steering message", timestamp: Date.now() };
|
||||
agent.steer(message);
|
||||
|
||||
// The message is queued but not yet in state.messages
|
||||
expect(agent.state.messages).not.toContainEqual(message);
|
||||
});
|
||||
|
||||
it("should support follow-up message queue", async () => {
|
||||
const agent = new Agent();
|
||||
|
||||
const message = { role: "user" as const, content: "Follow-up message", timestamp: Date.now() };
|
||||
agent.followUp(message);
|
||||
|
||||
// The message is queued but not yet in state.messages
|
||||
expect(agent.state.messages).not.toContainEqual(message);
|
||||
});
|
||||
|
||||
it("should handle abort controller", () => {
|
||||
const agent = new Agent();
|
||||
|
||||
// Should not throw even if nothing is running
|
||||
expect(() => agent.abort()).not.toThrow();
|
||||
});
|
||||
|
||||
it("should throw when prompt() called while streaming", async () => {
|
||||
let abortSignal: AbortSignal | undefined;
|
||||
const agent = new Agent({
|
||||
// Use a stream function that responds to abort
|
||||
streamFn: (_model, _context, options) => {
|
||||
abortSignal = options?.signal;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
// Check abort signal periodically
|
||||
const checkAbort = () => {
|
||||
if (abortSignal?.aborted) {
|
||||
stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
|
||||
} else {
|
||||
setTimeout(checkAbort, 5);
|
||||
}
|
||||
};
|
||||
checkAbort();
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
// Start first prompt (don't await, it will block until abort)
|
||||
const firstPrompt = agent.prompt("First message");
|
||||
|
||||
// Wait a tick for isStreaming to be set
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
expect(agent.state.isStreaming).toBe(true);
|
||||
|
||||
// Second prompt should reject
|
||||
await expect(agent.prompt("Second message")).rejects.toThrow(
|
||||
"Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.",
|
||||
);
|
||||
|
||||
// Cleanup - abort to stop the stream
|
||||
agent.abort();
|
||||
await firstPrompt.catch(() => {}); // Ignore abort error
|
||||
});
|
||||
|
||||
it("should throw when continue() called while streaming", async () => {
|
||||
let abortSignal: AbortSignal | undefined;
|
||||
const agent = new Agent({
|
||||
streamFn: (_model, _context, options) => {
|
||||
abortSignal = options?.signal;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
stream.push({ type: "start", partial: createAssistantMessage("") });
|
||||
const checkAbort = () => {
|
||||
if (abortSignal?.aborted) {
|
||||
stream.push({ type: "error", reason: "aborted", error: createAssistantMessage("Aborted") });
|
||||
} else {
|
||||
setTimeout(checkAbort, 5);
|
||||
}
|
||||
};
|
||||
checkAbort();
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
// Start first prompt
|
||||
const firstPrompt = agent.prompt("First message");
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
expect(agent.state.isStreaming).toBe(true);
|
||||
|
||||
// continue() should reject
|
||||
await expect(agent.continue()).rejects.toThrow(
|
||||
"Agent is already processing. Wait for completion before continuing.",
|
||||
);
|
||||
|
||||
// Cleanup
|
||||
agent.abort();
|
||||
await firstPrompt.catch(() => {});
|
||||
});
|
||||
|
||||
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("Processed") });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
agent.state.messages = [
|
||||
{
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Initial" }],
|
||||
timestamp: Date.now() - 10,
|
||||
},
|
||||
createAssistantMessage("Initial response"),
|
||||
];
|
||||
|
||||
agent.followUp({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Queued follow-up" }],
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
|
||||
await expect(agent.continue()).resolves.toBeUndefined();
|
||||
|
||||
const hasQueuedFollowUp = agent.state.messages.some((message) => {
|
||||
if (message.role !== "user") return false;
|
||||
if (typeof message.content === "string") return message.content === "Queued follow-up";
|
||||
return message.content.some((part) => part.type === "text" && part.text === "Queued follow-up");
|
||||
});
|
||||
|
||||
expect(hasQueuedFollowUp).toBe(true);
|
||||
expect(agent.state.messages[agent.state.messages.length - 1].role).toBe("assistant");
|
||||
});
|
||||
|
||||
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(`Processed ${responseCount}`),
|
||||
});
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
agent.state.messages = [
|
||||
{
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Initial" }],
|
||||
timestamp: Date.now() - 10,
|
||||
},
|
||||
createAssistantMessage("Initial response"),
|
||||
];
|
||||
|
||||
agent.steer({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Steering 1" }],
|
||||
timestamp: Date.now(),
|
||||
});
|
||||
agent.steer({
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Steering 2" }],
|
||||
timestamp: Date.now() + 1,
|
||||
});
|
||||
|
||||
await expect(agent.continue()).resolves.toBeUndefined();
|
||||
|
||||
const recentMessages = agent.state.messages.slice(-4);
|
||||
expect(recentMessages.map((m) => m.role)).toEqual(["user", "assistant", "user", "assistant"]);
|
||||
expect(responseCount).toBe(2);
|
||||
});
|
||||
|
||||
it("forwards sessionId to streamFn options", async () => {
|
||||
let receivedSessionId: string | undefined;
|
||||
const agent = new Agent({
|
||||
sessionId: "session-abc",
|
||||
streamFn: (_model, _context, options) => {
|
||||
receivedSessionId = options?.sessionId;
|
||||
const stream = new MockAssistantStream();
|
||||
queueMicrotask(() => {
|
||||
const message = createAssistantMessage("ok");
|
||||
stream.push({ type: "done", reason: "stop", message });
|
||||
});
|
||||
return stream;
|
||||
},
|
||||
});
|
||||
|
||||
await agent.prompt("hello");
|
||||
expect(receivedSessionId).toBe("session-abc");
|
||||
|
||||
// Test setter
|
||||
agent.sessionId = "session-def";
|
||||
expect(agent.sessionId).toBe("session-def");
|
||||
|
||||
await agent.prompt("hello again");
|
||||
expect(receivedSessionId).toBe("session-def");
|
||||
});
|
||||
});
|
||||
@@ -1,404 +0,0 @@
|
||||
import {
|
||||
type AssistantMessage,
|
||||
type FauxProviderRegistration,
|
||||
fauxAssistantMessage,
|
||||
fauxText,
|
||||
fauxThinking,
|
||||
fauxToolCall,
|
||||
type Model,
|
||||
registerFauxProvider,
|
||||
type ToolResultMessage,
|
||||
type UserMessage,
|
||||
} from "@mariozechner/pi-ai";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { Agent, type AgentEvent } from "../src/index.js";
|
||||
import { calculateTool } from "./utils/calculate.js";
|
||||
|
||||
const registrations: FauxProviderRegistration[] = [];
|
||||
|
||||
function createFauxRegistration(options: Parameters<typeof registerFauxProvider>[0] = {}): FauxProviderRegistration {
|
||||
const registration = registerFauxProvider(options);
|
||||
registrations.push(registration);
|
||||
return registration;
|
||||
}
|
||||
|
||||
function getTextContent(message: AssistantMessage | ToolResultMessage): string {
|
||||
return message.content
|
||||
.filter((block) => block.type === "text")
|
||||
.map((block) => block.text)
|
||||
.join("\n");
|
||||
}
|
||||
|
||||
afterEach(() => {
|
||||
while (registrations.length > 0) {
|
||||
registrations.pop()?.unregister();
|
||||
}
|
||||
});
|
||||
|
||||
async function basicPrompt(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant. Keep your responses concise.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
await agent.prompt("What is 2+2? Answer with just the number.");
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBe(2);
|
||||
expect(agent.state.messages[0].role).toBe("user");
|
||||
expect(agent.state.messages[1].role).toBe("assistant");
|
||||
|
||||
const assistantMessage = agent.state.messages[1];
|
||||
if (assistantMessage.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(getTextContent(assistantMessage)).toContain("4");
|
||||
}
|
||||
|
||||
async function toolExecution(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant. Always use the calculator tool for math.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [calculateTool],
|
||||
},
|
||||
});
|
||||
|
||||
const pendingToolCallsDuringEvents: Array<{ type: AgentEvent["type"]; ids: string[] }> = [];
|
||||
agent.subscribe((event) => {
|
||||
if (event.type === "tool_execution_start" || event.type === "tool_execution_end") {
|
||||
pendingToolCallsDuringEvents.push({
|
||||
type: event.type,
|
||||
ids: [...agent.state.pendingToolCalls],
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
await agent.prompt("Calculate 123 * 456 using the calculator tool.");
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBeGreaterThanOrEqual(4);
|
||||
const toolResultMsg = agent.state.messages.find((message) => message.role === "toolResult");
|
||||
expect(toolResultMsg).toBeDefined();
|
||||
if (toolResultMsg?.role !== "toolResult") throw new Error("Expected tool result message");
|
||||
expect(getTextContent(toolResultMsg)).toContain("123 * 456 = 56088");
|
||||
|
||||
const finalMessage = agent.state.messages[agent.state.messages.length - 1];
|
||||
if (finalMessage.role !== "assistant") throw new Error("Expected final assistant message");
|
||||
expect(getTextContent(finalMessage)).toContain("56088");
|
||||
expect(agent.state.pendingToolCalls.size).toBe(0);
|
||||
expect(pendingToolCallsDuringEvents).toEqual([
|
||||
{ type: "tool_execution_start", ids: ["calc-1"] },
|
||||
{ type: "tool_execution_end", ids: [] },
|
||||
]);
|
||||
}
|
||||
|
||||
async function abortExecution(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
const promptPromise = agent.prompt("Count slowly from 1 to 20.");
|
||||
setTimeout(() => {
|
||||
agent.abort();
|
||||
}, 30);
|
||||
|
||||
await promptPromise;
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBeGreaterThanOrEqual(2);
|
||||
|
||||
const lastMessage = agent.state.messages[agent.state.messages.length - 1];
|
||||
if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(lastMessage.stopReason).toBe("aborted");
|
||||
expect(lastMessage.errorMessage).toBeDefined();
|
||||
expect(agent.state.errorMessage).toBe(lastMessage.errorMessage);
|
||||
}
|
||||
|
||||
async function stateUpdates(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
const events: AgentEvent["type"][] = [];
|
||||
agent.subscribe((event) => {
|
||||
events.push(event.type);
|
||||
});
|
||||
|
||||
await agent.prompt("Count from 1 to 5.");
|
||||
|
||||
expect(events).toContain("agent_start");
|
||||
expect(events).toContain("turn_start");
|
||||
expect(events).toContain("message_start");
|
||||
expect(events).toContain("message_update");
|
||||
expect(events).toContain("message_end");
|
||||
expect(events).toContain("turn_end");
|
||||
expect(events).toContain("agent_end");
|
||||
expect(events.indexOf("agent_start")).toBeLessThan(events.indexOf("message_start"));
|
||||
expect(events.indexOf("message_start")).toBeLessThan(events.indexOf("message_end"));
|
||||
expect(events.indexOf("message_end")).toBeLessThan(events.lastIndexOf("agent_end"));
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBe(2);
|
||||
}
|
||||
|
||||
async function multiTurnConversation(model: Model<string>) {
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
await agent.prompt("My name is Alice.");
|
||||
expect(agent.state.messages.length).toBe(2);
|
||||
|
||||
await agent.prompt("What is my name?");
|
||||
expect(agent.state.messages.length).toBe(4);
|
||||
|
||||
const lastMessage = agent.state.messages[3];
|
||||
if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(getTextContent(lastMessage).toLowerCase()).toContain("alice");
|
||||
}
|
||||
|
||||
describe("Agent integration with faux provider", () => {
|
||||
it("handles a basic text prompt", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
faux.setResponses([fauxAssistantMessage("4")]);
|
||||
await basicPrompt(faux.getModel());
|
||||
});
|
||||
|
||||
it("executes tools and tracks pending tool calls", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
faux.setResponses([
|
||||
fauxAssistantMessage(
|
||||
[
|
||||
fauxText("Let me calculate that."),
|
||||
fauxToolCall("calculate", { expression: "123 * 456" }, { id: "calc-1" }),
|
||||
],
|
||||
{ stopReason: "toolUse" },
|
||||
),
|
||||
fauxAssistantMessage("The result is 56088."),
|
||||
]);
|
||||
await toolExecution(faux.getModel());
|
||||
});
|
||||
|
||||
it("handles abort during streaming", async () => {
|
||||
const faux = createFauxRegistration({
|
||||
tokensPerSecond: 20,
|
||||
tokenSize: { min: 2, max: 2 },
|
||||
});
|
||||
faux.setResponses([
|
||||
fauxAssistantMessage(
|
||||
"one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen",
|
||||
),
|
||||
]);
|
||||
await abortExecution(faux.getModel());
|
||||
});
|
||||
|
||||
it("emits lifecycle updates while streaming", async () => {
|
||||
const faux = createFauxRegistration({ tokenSize: { min: 1, max: 1 } });
|
||||
faux.setResponses([fauxAssistantMessage("1 2 3 4 5")]);
|
||||
await stateUpdates(faux.getModel());
|
||||
});
|
||||
|
||||
it("maintains context across multiple turns", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
faux.setResponses([
|
||||
fauxAssistantMessage("Nice to meet you, Alice."),
|
||||
(context) => {
|
||||
const hasAlice = context.messages.some((message) => {
|
||||
if (message.role !== "user") return false;
|
||||
if (typeof message.content === "string") return message.content.includes("Alice");
|
||||
return message.content.some((block) => block.type === "text" && block.text.includes("Alice"));
|
||||
});
|
||||
return fauxAssistantMessage(hasAlice ? "Your name is Alice." : "I do not know your name.");
|
||||
},
|
||||
]);
|
||||
await multiTurnConversation(faux.getModel());
|
||||
});
|
||||
|
||||
it("preserves thinking content blocks", async () => {
|
||||
const faux = createFauxRegistration({ models: [{ id: "faux-reasoning", reasoning: true }] });
|
||||
faux.setResponses([fauxAssistantMessage([fauxThinking("step by step"), fauxText("4")])]);
|
||||
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant.",
|
||||
model: faux.getModel(),
|
||||
thinkingLevel: "low",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
await agent.prompt("What is 2+2?");
|
||||
|
||||
const assistantMessage = agent.state.messages[1];
|
||||
if (assistantMessage?.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(assistantMessage.content).toEqual([
|
||||
{ type: "thinking", thinking: "step by step" },
|
||||
{ type: "text", text: "4" },
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Agent.continue() with faux provider", () => {
|
||||
describe("validation", () => {
|
||||
it("throws when no messages in context", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "Test",
|
||||
model: faux.getModel(),
|
||||
},
|
||||
});
|
||||
|
||||
await expect(agent.continue()).rejects.toThrow("No messages to continue from");
|
||||
});
|
||||
|
||||
it("throws when last message is assistant", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
const model = faux.getModel();
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "Test",
|
||||
model,
|
||||
},
|
||||
});
|
||||
|
||||
const assistantMessage: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [{ type: "text", text: "Hello" }],
|
||||
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: "stop",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
agent.state.messages = [assistantMessage];
|
||||
|
||||
await expect(agent.continue()).rejects.toThrow("Cannot continue from message role: assistant");
|
||||
});
|
||||
});
|
||||
|
||||
describe("continue from user message", () => {
|
||||
it("continues and gets a response when last message is user", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
faux.setResponses([fauxAssistantMessage("HELLO WORLD")]);
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt: "You are a helpful assistant. Follow instructions exactly.",
|
||||
model: faux.getModel(),
|
||||
thinkingLevel: "off",
|
||||
tools: [],
|
||||
},
|
||||
});
|
||||
|
||||
const userMessage: UserMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "Say exactly: HELLO WORLD" }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
agent.state.messages = [userMessage];
|
||||
|
||||
await agent.continue();
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBe(2);
|
||||
expect(agent.state.messages[0].role).toBe("user");
|
||||
expect(agent.state.messages[1].role).toBe("assistant");
|
||||
|
||||
const assistantMsg = agent.state.messages[1];
|
||||
if (assistantMsg.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(getTextContent(assistantMsg).toUpperCase()).toContain("HELLO WORLD");
|
||||
});
|
||||
});
|
||||
|
||||
describe("continue from tool result", () => {
|
||||
it("continues and processes tool results", async () => {
|
||||
const faux = createFauxRegistration();
|
||||
const model = faux.getModel();
|
||||
faux.setResponses([fauxAssistantMessage("The answer is 8.")]);
|
||||
const agent = new Agent({
|
||||
initialState: {
|
||||
systemPrompt:
|
||||
"You are a helpful assistant. After getting a calculation result, state the answer clearly.",
|
||||
model,
|
||||
thinkingLevel: "off",
|
||||
tools: [calculateTool],
|
||||
},
|
||||
});
|
||||
|
||||
const userMessage: UserMessage = {
|
||||
role: "user",
|
||||
content: [{ type: "text", text: "What is 5 + 3?" }],
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const assistantMessage: AssistantMessage = {
|
||||
role: "assistant",
|
||||
content: [
|
||||
{ type: "text", text: "Let me calculate that." },
|
||||
{ type: "toolCall", id: "calc-1", name: "calculate", arguments: { expression: "5 + 3" } },
|
||||
],
|
||||
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: "toolUse",
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
const toolResult: ToolResultMessage = {
|
||||
role: "toolResult",
|
||||
toolCallId: "calc-1",
|
||||
toolName: "calculate",
|
||||
content: [{ type: "text", text: "5 + 3 = 8" }],
|
||||
isError: false,
|
||||
timestamp: Date.now(),
|
||||
};
|
||||
|
||||
agent.state.messages = [userMessage, assistantMessage, toolResult];
|
||||
|
||||
await agent.continue();
|
||||
|
||||
expect(agent.state.isStreaming).toBe(false);
|
||||
expect(agent.state.messages.length).toBeGreaterThanOrEqual(4);
|
||||
|
||||
const lastMessage = agent.state.messages[agent.state.messages.length - 1];
|
||||
expect(lastMessage.role).toBe("assistant");
|
||||
if (lastMessage.role !== "assistant") throw new Error("Expected assistant message");
|
||||
expect(getTextContent(lastMessage)).toContain("8");
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,32 +0,0 @@
|
||||
import { type Static, Type } from "@sinclair/typebox";
|
||||
import type { AgentTool, AgentToolResult } from "../../src/types.js";
|
||||
|
||||
export interface CalculateResult extends AgentToolResult<undefined> {
|
||||
content: Array<{ type: "text"; text: string }>;
|
||||
details: undefined;
|
||||
}
|
||||
|
||||
export function calculate(expression: string): CalculateResult {
|
||||
try {
|
||||
const result = new Function(`return ${expression}`)();
|
||||
return { content: [{ type: "text", text: `${expression} = ${result}` }], details: undefined };
|
||||
} catch (e: any) {
|
||||
throw new Error(e.message || String(e));
|
||||
}
|
||||
}
|
||||
|
||||
const calculateSchema = Type.Object({
|
||||
expression: Type.String({ description: "The mathematical expression to evaluate" }),
|
||||
});
|
||||
|
||||
type CalculateParams = Static<typeof calculateSchema>;
|
||||
|
||||
export const calculateTool: AgentTool<typeof calculateSchema, undefined> = {
|
||||
label: "Calculator",
|
||||
name: "calculate",
|
||||
description: "Evaluate mathematical expressions",
|
||||
parameters: calculateSchema,
|
||||
execute: async (_toolCallId: string, args: CalculateParams) => {
|
||||
return calculate(args.expression);
|
||||
},
|
||||
};
|
||||
@@ -1,46 +0,0 @@
|
||||
import { type Static, Type } from "@sinclair/typebox";
|
||||
import type { AgentTool, AgentToolResult } from "../../src/types.js";
|
||||
|
||||
export interface GetCurrentTimeResult extends AgentToolResult<{ utcTimestamp: number }> {}
|
||||
|
||||
export async function getCurrentTime(timezone?: string): Promise<GetCurrentTimeResult> {
|
||||
const date = new Date();
|
||||
if (timezone) {
|
||||
try {
|
||||
const timeStr = date.toLocaleString("en-US", {
|
||||
timeZone: timezone,
|
||||
dateStyle: "full",
|
||||
timeStyle: "long",
|
||||
});
|
||||
return {
|
||||
content: [{ type: "text", text: timeStr }],
|
||||
details: { utcTimestamp: date.getTime() },
|
||||
};
|
||||
} catch (_e) {
|
||||
throw new Error(`Invalid timezone: ${timezone}. Current UTC time: ${date.toISOString()}`);
|
||||
}
|
||||
}
|
||||
const timeStr = date.toLocaleString("en-US", { dateStyle: "full", timeStyle: "long" });
|
||||
return {
|
||||
content: [{ type: "text", text: timeStr }],
|
||||
details: { utcTimestamp: date.getTime() },
|
||||
};
|
||||
}
|
||||
|
||||
const getCurrentTimeSchema = Type.Object({
|
||||
timezone: Type.Optional(
|
||||
Type.String({ description: "Optional timezone (e.g., 'America/New_York', 'Europe/London')" }),
|
||||
),
|
||||
});
|
||||
|
||||
type GetCurrentTimeParams = Static<typeof getCurrentTimeSchema>;
|
||||
|
||||
export const getCurrentTimeTool: AgentTool<typeof getCurrentTimeSchema, { utcTimestamp: number }> = {
|
||||
label: "Current Time",
|
||||
name: "get_current_time",
|
||||
description: "Get the current date and time",
|
||||
parameters: getCurrentTimeSchema,
|
||||
execute: async (_toolCallId: string, args: GetCurrentTimeParams) => {
|
||||
return getCurrentTime(args.timezone);
|
||||
},
|
||||
};
|
||||
Reference in new issue
Block a user