385 lines
13 KiB
TypeScript
385 lines
13 KiB
TypeScript
import type { AgentMessage } from "@earendil-works/pi-agent-core";
|
|
import { createAssistantMessageEventStream, normalizeContext } from "@earendil-works/pi-ai";
|
|
import { describe, expect, it } from "vitest";
|
|
import {
|
|
describeToolActivity,
|
|
estimateAgentContextTokens,
|
|
maxInputTokensForWindow,
|
|
parseContextOverflowDetails,
|
|
pruneContextToWindow,
|
|
raceSignal,
|
|
streamWithContextOverflowRecovery,
|
|
} from "./pi-runtime.js";
|
|
|
|
describe("describeToolActivity", () => {
|
|
it("summarizes builtin tools with their most informative argument", () => {
|
|
expect(describeToolActivity("shell", { command: "pnpm test --filter web" })).toBe(
|
|
"Running: pnpm test --filter web",
|
|
);
|
|
expect(describeToolActivity("read_file", { path: "notes/plan.md" })).toBe(
|
|
"Reading notes/plan.md",
|
|
);
|
|
expect(describeToolActivity("write_file", { path: "out.csv", content: "…" })).toBe(
|
|
"Writing out.csv",
|
|
);
|
|
expect(describeToolActivity("render_plot", { spec: {} })).toBe("Rendering a chart");
|
|
expect(describeToolActivity("add_mcp_server", { name: "Linear" })).toBe(
|
|
"Connecting MCP server: Linear",
|
|
);
|
|
expect(describeToolActivity("run_subagent", { name: "scout", task: "…" })).toBe(
|
|
"Delegating to helper: scout",
|
|
);
|
|
expect(describeToolActivity("create_space", { name: "Customer support" })).toBe(
|
|
"Creating space: Customer support",
|
|
);
|
|
expect(describeToolActivity("web_search", { query: "rakazo agents" })).toBe(
|
|
"Searching the web: rakazo agents",
|
|
);
|
|
expect(describeToolActivity("web_fetch", { url: "https://example.com" })).toBe(
|
|
"Reading page: https://example.com/",
|
|
);
|
|
});
|
|
|
|
it("names MCP server and remote tool", () => {
|
|
expect(describeToolActivity("mcp__brex__list_expenses", {})).toBe("Using brex: list_expenses");
|
|
expect(describeToolActivity("mcp__demo-oauth__greet", {})).toBe("Using demo-oauth: greet");
|
|
});
|
|
|
|
it("truncates long details and collapses whitespace", () => {
|
|
const long = `x${"y".repeat(200)}`;
|
|
const line = describeToolActivity("shell", { command: `a\n\t${long}` });
|
|
expect(line.length).toBeLessThanOrEqual("Running: ".length + 91);
|
|
expect(line).toContain("…");
|
|
expect(line).not.toContain("\n");
|
|
expect(line).toMatch(/^Running: a x/);
|
|
});
|
|
|
|
it("redacts credentials from activity details", () => {
|
|
const token = "fake-token";
|
|
const line = describeToolActivity("shell", {
|
|
command: `curl -H 'Authorization: Bearer ${token}' https://example.test?api_key=fake-key password=fake-password`,
|
|
});
|
|
|
|
expect(line).toContain("Bearer [redacted]");
|
|
expect(line).toContain("api_key=[redacted]");
|
|
expect(line).not.toContain(token);
|
|
expect(line).not.toContain("fake-key");
|
|
expect(line).not.toContain("fake-password");
|
|
});
|
|
|
|
it("strips signed-URL query and fragment from web_fetch activity", () => {
|
|
const line = describeToolActivity("web_fetch", {
|
|
url: "https://user:secret@cdn.example.test/doc.pdf?X-Amz-Signature=abc123&token=leak#frag",
|
|
});
|
|
expect(line).toBe("Reading page: https://cdn.example.test/doc.pdf");
|
|
expect(line).not.toContain("secret");
|
|
expect(line).not.toContain("X-Amz-Signature");
|
|
expect(line).not.toContain("token=");
|
|
expect(line).not.toContain("abc123");
|
|
expect(line).not.toContain("frag");
|
|
});
|
|
|
|
it("does not echo secrets from malformed web_fetch URLs", () => {
|
|
const line = describeToolActivity("web_fetch", {
|
|
url: "https://user:secret@[",
|
|
});
|
|
expect(line).toBe("Reading page: [invalid URL]");
|
|
expect(line).not.toContain("secret");
|
|
expect(line).not.toContain("user:");
|
|
});
|
|
|
|
it("falls back to the tool name", () => {
|
|
expect(describeToolActivity("destination_write", undefined)).toBe("Using destination_write");
|
|
});
|
|
});
|
|
|
|
describe("raceSignal", () => {
|
|
it("resolves when the promise resolves before abort", async () => {
|
|
const controller = new AbortController();
|
|
const result = await raceSignal(Promise.resolve("ok"), controller.signal);
|
|
expect(result).toBe("ok");
|
|
});
|
|
|
|
it("rejects when the promise rejects before abort", async () => {
|
|
const controller = new AbortController();
|
|
await expect(
|
|
raceSignal(Promise.reject(new Error("tool failed")), controller.signal),
|
|
).rejects.toThrow("tool failed");
|
|
});
|
|
|
|
it("rejects with abort reason when aborted while promise is pending", async () => {
|
|
const controller = new AbortController();
|
|
let resolvePending: (val: string) => void;
|
|
const pending = new Promise<string>((resolve) => {
|
|
resolvePending = resolve;
|
|
});
|
|
|
|
const promise = raceSignal(pending, controller.signal);
|
|
controller.abort(new Error("Interrupted by user speech"));
|
|
|
|
await expect(promise).rejects.toThrow("Interrupted by user speech");
|
|
resolvePending!("late");
|
|
});
|
|
|
|
it("rejects immediately if signal is already aborted", async () => {
|
|
const controller = new AbortController();
|
|
controller.abort(new Error("Already cancelled"));
|
|
await expect(raceSignal(Promise.resolve("ok"), controller.signal)).rejects.toThrow(
|
|
"Already cancelled",
|
|
);
|
|
});
|
|
});
|
|
|
|
describe("context window pruning and overflow recovery", () => {
|
|
it("parses llama.cpp exceed_context_size_error and standard overflow errors", () => {
|
|
const llamaCppError =
|
|
'400: {"code":400,"message":"request (133699 tokens) exceeds the available context size (131072 tokens), try increasing it","type":"exceed_context_size_error","n_prompt_tokens":133699,"n_ctx":131072}';
|
|
expect(parseContextOverflowDetails(llamaCppError)).toEqual({
|
|
limitTokens: 131_072,
|
|
requestedTokens: 133_699,
|
|
});
|
|
|
|
const plainLlamaError =
|
|
"request (133699 tokens) exceeds the available context size (131072 tokens), try increasing it";
|
|
expect(parseContextOverflowDetails(plainLlamaError)).toEqual({
|
|
limitTokens: 131_072,
|
|
requestedTokens: 133_699,
|
|
});
|
|
|
|
const anthropicError = "prompt is too long: 213462 tokens > 200000 maximum";
|
|
expect(parseContextOverflowDetails(anthropicError)).toEqual({
|
|
limitTokens: 200_000,
|
|
requestedTokens: 213_462,
|
|
});
|
|
|
|
expect(parseContextOverflowDetails("429 Too many requests: rate limit exceeded")).toBeNull();
|
|
});
|
|
|
|
it("compacts older tool results and drops oldest history/tool rounds while keeping pairing intact", () => {
|
|
const messages: AgentMessage[] = [
|
|
{
|
|
role: "system",
|
|
content: "You are Trading-Bot.",
|
|
toolsAdded: [{ name: "shell", description: "run shell", parameters: {} as never }],
|
|
timestamp: 1,
|
|
},
|
|
{
|
|
role: "user",
|
|
content: `Old history 1: ${"h".repeat(9_000)}`,
|
|
timestamp: 2,
|
|
},
|
|
{
|
|
role: "user",
|
|
content: `Old history 2: ${"h".repeat(9_000)}`,
|
|
timestamp: 3,
|
|
},
|
|
{
|
|
role: "user",
|
|
content: "Fix the index error in backtest_p2.py.",
|
|
timestamp: 4,
|
|
},
|
|
{
|
|
role: "assistant",
|
|
content: [
|
|
{ type: "thinking", thinking: "t".repeat(6_000) },
|
|
{
|
|
type: "toolCall",
|
|
id: "call_1",
|
|
name: "write_file",
|
|
arguments: { path: "backtest_p2.py", content: "c".repeat(12_000) },
|
|
},
|
|
],
|
|
api: "openai-completions",
|
|
provider: "local",
|
|
model: "qwen",
|
|
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: 5,
|
|
},
|
|
{
|
|
role: "toolResult",
|
|
toolCallId: "call_1",
|
|
toolName: "write_file",
|
|
content: [{ type: "text", text: "r".repeat(12_000) }],
|
|
isError: false,
|
|
timestamp: 6,
|
|
},
|
|
{
|
|
role: "assistant",
|
|
content: [
|
|
{
|
|
type: "toolCall",
|
|
id: "call_2",
|
|
name: "shell",
|
|
arguments: { command: "python3 backtest_p2.py" },
|
|
},
|
|
],
|
|
api: "openai-completions",
|
|
provider: "local",
|
|
model: "qwen",
|
|
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: 7,
|
|
},
|
|
{
|
|
role: "toolResult",
|
|
toolCallId: "call_2",
|
|
toolName: "shell",
|
|
content: [{ type: "text", text: "latest output: IndexError fixed" }],
|
|
isError: false,
|
|
timestamp: 8,
|
|
},
|
|
];
|
|
|
|
const pruned = pruneContextToWindow(messages, {
|
|
contextWindow: 4_096,
|
|
maxTokens: 1_024,
|
|
});
|
|
|
|
expect(estimateAgentContextTokens(pruned)).toBeLessThanOrEqual(
|
|
maxInputTokensForWindow(4_096, 1_024),
|
|
);
|
|
expect(pruned[0]?.role).toBe("system");
|
|
expect(
|
|
pruned.some(
|
|
(m) => m.role === "user" && m.content === "Fix the index error in backtest_p2.py.",
|
|
),
|
|
).toBe(true);
|
|
const remainingToolCallIds = new Set(
|
|
pruned
|
|
.filter((m): m is Extract<AgentMessage, { role: "assistant" }> => m.role === "assistant")
|
|
.flatMap((m) =>
|
|
m.content
|
|
.filter(
|
|
(part): part is Extract<(typeof m.content)[number], { type: "toolCall" }> =>
|
|
part.type === "toolCall",
|
|
)
|
|
.map((part) => part.id),
|
|
),
|
|
);
|
|
for (const msg of pruned) {
|
|
if (msg.role === "toolResult") {
|
|
expect(remainingToolCallIds.has(msg.toolCallId)).toBe(true);
|
|
}
|
|
}
|
|
});
|
|
|
|
it("retries transparently after a provider exceed_context_size_error", async () => {
|
|
let effectiveWindow = 131_072;
|
|
let calls = 0;
|
|
|
|
const initialMessages: AgentMessage[] = [
|
|
{ role: "system", content: "You are a bot.", timestamp: 1 },
|
|
{ role: "user", content: `Prior history: ${"x".repeat(60_000)}`, timestamp: 2 },
|
|
{ role: "user", content: "Run the backtest.", timestamp: 3 },
|
|
];
|
|
|
|
const stream = streamWithContextOverflowRecovery({
|
|
models: {
|
|
streamSimple: (_model, ctx) => {
|
|
calls += 1;
|
|
expect(ctx.messages.length).toBeGreaterThanOrEqual(2);
|
|
const inner = createAssistantMessageEventStream();
|
|
queueMicrotask(() => {
|
|
if (calls === 1) {
|
|
inner.push({
|
|
type: "error",
|
|
reason: "error",
|
|
error: {
|
|
role: "assistant",
|
|
content: [],
|
|
api: "openai-completions",
|
|
provider: "local",
|
|
model: "qwen",
|
|
usage: {
|
|
input: 0,
|
|
output: 0,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 0,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
},
|
|
stopReason: "error",
|
|
errorMessage:
|
|
'400: {"code":400,"message":"request (133699 tokens) exceeds the available context size (131072 tokens), try increasing it","type":"exceed_context_size_error","n_prompt_tokens":133699,"n_ctx":131072}',
|
|
timestamp: Date.now(),
|
|
},
|
|
});
|
|
inner.end();
|
|
return;
|
|
}
|
|
const okMessage = {
|
|
role: "assistant" as const,
|
|
content: [{ type: "text" as const, text: "Backtest completed." }],
|
|
api: "openai-completions" as const,
|
|
provider: "local",
|
|
model: "qwen",
|
|
usage: {
|
|
input: 1_000,
|
|
output: 10,
|
|
cacheRead: 0,
|
|
cacheWrite: 0,
|
|
totalTokens: 1_010,
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
|
},
|
|
stopReason: "stop" as const,
|
|
timestamp: Date.now(),
|
|
};
|
|
inner.push({ type: "start", partial: okMessage });
|
|
inner.push({
|
|
type: "text_delta",
|
|
contentIndex: 0,
|
|
delta: "Backtest completed.",
|
|
partial: okMessage,
|
|
});
|
|
inner.push({ type: "done", reason: "stop", message: okMessage });
|
|
inner.end();
|
|
});
|
|
return inner;
|
|
},
|
|
},
|
|
model: {
|
|
id: "qwen",
|
|
name: "qwen",
|
|
api: "openai-completions",
|
|
provider: "local",
|
|
baseUrl: "http://127.0.0.1:11434/v1",
|
|
reasoning: false,
|
|
input: ["text"],
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: 131_072,
|
|
maxTokens: 4_096,
|
|
},
|
|
context: normalizeContext({ messages: initialMessages as never }),
|
|
getEffectiveContextWindow: () => effectiveWindow,
|
|
setEffectiveContextWindow: (next) => {
|
|
effectiveWindow = next;
|
|
},
|
|
});
|
|
|
|
const events = [];
|
|
for await (const event of stream) {
|
|
events.push(event.type);
|
|
}
|
|
const finalMsg = await stream.result();
|
|
expect(calls).toBe(2);
|
|
expect(effectiveWindow).toBeLessThan(131_072);
|
|
expect(events).toEqual(["start", "text_delta", "done"]);
|
|
expect(finalMsg.stopReason).toBe("stop");
|
|
});
|
|
});
|