Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
60 kB · 2034 lines
TypeScript
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035import { type AssistantMessage, type AssistantMessageEvent, EventStream, type Message, type Model, type UserMessage,} from "@earendil-works/pi-ai";import { Type } from "typebox";import { describe, expect, it, vi } from "vitest";import { agentLoop, agentLoopContinue, runAgentLoop } from "../src/agent-loop.js";import type { AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentTool } from "../src/types.js";
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"); }, ); }}
class DelayedResultStream extends MockAssistantStream { constructor(private readonly getDelayedResult: () => Promise<AssistantMessage>) { super(); }
override result(): Promise<AssistantMessage> { return this.getDelayedResult(); }}
class ThrowingResultStream extends MockAssistantStream { constructor( private readonly onResult: () => void, private readonly error: Error, ) { super(); }
override [Symbol.asyncIterator](): AsyncIterator<AssistantMessageEvent> { return { next: async () => ({ done: true, value: undefined as never }), }; }
override result(): Promise<AssistantMessage> { this.onResult(); throw this.error; }}
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(), };}
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 preserve a terminal response when abort fires after done", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const finalMessage = createAssistantMessage([{ type: "text", text: "complete" }]); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const streamFn = () => { const stream = new DelayedResultStream(() => { controller.abort(); return Promise.resolve(finalMessage); }); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: finalMessage }); }); return stream; };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const _event of stream) { }
const messages = await stream.result(); const assistant = messages.find((message) => message.role === "assistant"); expect(assistant?.role).toBe("assistant"); if (assistant?.role === "assistant") { expect(assistant.stopReason).toBe("stop"); expect(assistant.content).toEqual([{ type: "text", text: "complete" }]); } });
it("should not wait for a pending terminal result after abort", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const finalMessage = createAssistantMessage([{ type: "text", text: "complete" }]); let resolveResult: ((message: AssistantMessage) => void) | undefined; const pendingResult = new Promise<AssistantMessage>((resolve) => { resolveResult = resolve; }); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const streamFn = () => { const stream = new DelayedResultStream(() => { controller.abort(); return pendingResult; }); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: finalMessage }); }); return stream; };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); const events: AgentEvent[] = []; const consume = (async () => { for await (const event of stream) { events.push(event); } return stream.result(); })(); const result = await Promise.race([ consume, new Promise<"timeout">((resolve) => { setTimeout(() => resolve("timeout"), 50); }), ]); if (result === "timeout") { resolveResult?.(finalMessage); await consume; throw new Error("agent loop waited for a pending terminal result after abort"); }
const assistant = result.find((message) => message.role === "assistant"); expect(assistant?.role).toBe("assistant"); if (assistant?.role === "assistant") { expect(assistant.stopReason).toBe("stop"); expect(assistant.content).toEqual([{ type: "text", text: "complete" }]); } expect(events.some((event) => event.type === "agent_end")).toBe(true); });
it("should not mask stream result errors when abort is already signaled", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const providerError = new Error("provider parse failed"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const events: AgentEvent[] = [];
await expect( runAgentLoop( [createUserMessage("Hello")], context, config, (event) => { events.push(event); }, controller.signal, () => new ThrowingResultStream(() => { controller.abort(); }, providerError), ), ).rejects.toThrow("provider parse failed"); expect( events.some( (event) => event.type === "message_end" && event.message.role === "assistant" && event.message.stopReason === "aborted", ), ).toBe(false); });
it("should return an aborted assistant when abort fires before the stream starts", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const streamFn = vi.fn(() => new MockAssistantStream()); const config: AgentLoopConfig = { model: createModel(), convertToLlm: (messages) => { controller.abort(); return identityConverter(messages); }, };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const _event of stream) { }
const messages = await stream.result(); const assistant = messages.find((message) => message.role === "assistant"); expect(streamFn).not.toHaveBeenCalled(); expect(assistant?.role).toBe("assistant"); if (assistant?.role === "assistant") { expect(assistant.stopReason).toBe("aborted"); expect(assistant.errorMessage).toBe("Request was aborted"); } });
it("should end without adding an aborted assistant when abort fires after turn completion", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const finalMessage = createAssistantMessage([{ type: "text", text: "complete" }]); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, shouldStopAfterTurn: async () => { controller.abort(); await new Promise((resolve) => setTimeout(resolve, 0)); return false; }, }; const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "stop", message: finalMessage }); }); return stream; };
const events: AgentEvent[] = []; const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const event of stream) { events.push(event); }
const messages = await stream.result(); const assistantMessages = messages.filter((message) => message.role === "assistant");
expect(assistantMessages).toHaveLength(1); expect(assistantMessages[0]?.role).toBe("assistant"); if (assistantMessages[0]?.role === "assistant") { expect(assistantMessages[0].stopReason).toBe("stop"); expect(assistantMessages[0].content).toEqual([{ type: "text", text: "complete" }]); } expect(events.filter((event) => event.type === "agent_end")).toHaveLength(1); });
it("should not invoke tools when the signal aborts before tool execution", async () => { const controller = new AbortController(); const toolExecute: AgentTool["execute"] = vi.fn(async () => ({ content: [{ type: "text" as const, text: "should not run" }], details: {}, })); const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [ { name: "wait", label: "Wait", description: "Wait", parameters: Type.Object({}), execute: toolExecute, }, ], }; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, beforeToolCall: async () => { controller.abort(); return undefined; }, }; const assistantMessage = createAssistantMessage( [{ type: "toolCall", id: "tool_1", name: "wait", arguments: {} }], "toolUse", ); const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "toolUse", message: assistantMessage }); }); return stream; };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const _event of stream) { } await stream.result();
expect(toolExecute).not.toHaveBeenCalled(); });
it("should stop a sequential tool batch after aborting a tool call", async () => { const controller = new AbortController(); const toolSchema = Type.Object({ value: Type.String() }); const executed: string[] = []; const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "work", label: "Work", description: "Work", parameters: toolSchema, execute: async (_toolCallId, params) => { executed.push(params.value); if (params.value === "first") { controller.abort(); } return { content: [{ type: "text", text: `done:${params.value}` }], details: { value: params.value }, }; }, }; const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [tool], }; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, toolExecution: "sequential", }; const assistantMessage = createAssistantMessage( [ { type: "toolCall", id: "tool_1", name: "work", arguments: { value: "first" } }, { type: "toolCall", id: "tool_2", name: "work", arguments: { value: "second" } }, ], "toolUse", ); const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "toolUse", message: assistantMessage }); }); return stream; }; const events: AgentEvent[] = [];
await runAgentLoop( [createUserMessage("Hello")], context, config, (event) => { events.push(event); }, controller.signal, streamFn, );
const toolStartIds = events.flatMap((event) => (event.type === "tool_execution_start" ? [event.toolCallId] : [])); const toolResultIds = events.flatMap((event) => event.type === "message_end" && event.message.role === "toolResult" ? [event.message.toolCallId] : [], );
expect(executed).toEqual(["first"]); expect(toolStartIds).toEqual(["tool_1"]); expect(toolResultIds).toEqual(["tool_1"]); expect(events.some((event) => event.type === "agent_end")).toBe(true); });
it("should preserve a successful tool result when abort fires during update flush", async () => { const controller = new AbortController(); const toolSchema = Type.Object({}); const tool: AgentTool<typeof toolSchema, Record<string, never>> = { name: "work", label: "Work", description: "Work", parameters: toolSchema, execute: async (_toolCallId, _params, _signal, onUpdate) => { onUpdate?.({ content: [{ type: "text", text: "progress" }], details: {} }); return { content: [{ type: "text", text: "done" }], details: {}, }; }, }; const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [tool], }; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, toolExecution: "sequential", }; const assistantMessage = createAssistantMessage( [{ type: "toolCall", id: "tool_1", name: "work", arguments: {} }], "toolUse", ); const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "done", reason: "toolUse", message: assistantMessage }); }); return stream; }; const events: AgentEvent[] = [];
const messages = await runAgentLoop( [createUserMessage("Hello")], context, config, (event) => { events.push(event); if (event.type === "tool_execution_update") { return new Promise<void>(() => { setTimeout(() => controller.abort(), 0); }); } }, controller.signal, streamFn, );
const toolResult = messages.find((message) => message.role === "toolResult"); expect(toolResult?.role).toBe("toolResult"); if (toolResult?.role === "toolResult") { expect(toolResult.isError).toBe(false); expect(toolResult.content).toEqual([{ type: "text", text: "done" }]); } expect(events.some((event) => event.type === "agent_end")).toBe(true); });
it("should freeze synthetic aborted messages against later partial mutation", async () => { const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [], }; const controller = new AbortController(); const partialMessage = createAssistantMessage([{ type: "text", text: "partial" }]); partialMessage.usage.cost.total = 1; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { stream.push({ type: "start", partial: partialMessage }); }); return stream; };
const stream = agentLoop([createUserMessage("Hello")], context, config, controller.signal, streamFn); for await (const event of stream) { if (event.type === "message_start" && event.message.role === "assistant") { controller.abort(); } }
const messages = await stream.result(); const assistant = messages.find((message) => message.role === "assistant"); const text = partialMessage.content[0]; if (text?.type === "text") { text.text = "mutated"; } partialMessage.usage.cost.total = 99;
expect(assistant?.role).toBe("assistant"); if (assistant?.role === "assistant") { expect(assistant.stopReason).toBe("aborted"); expect(assistant.content).toEqual([{ type: "text", text: "partial" }]); expect(assistant.usage.cost.total).toBe(1); } });
it("should handle custom message types via convertToLlm", async () => { 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) => { 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); }
expect(convertedMessages.length).toBe(1); // Only user message expect(convertedMessages[0].role).toBe("user"); });
it("should resolve the system prompt after asynchronous API key lookup", async () => { let systemPrompt = "before lookup"; let resolveApiKey: ((apiKey: string) => void) | undefined; const apiKey = new Promise<string>((resolve) => { resolveApiKey = resolve; }); let markLookupStarted: (() => void) | undefined; const lookupStarted = new Promise<void>((resolve) => { markLookupStarted = resolve; }); let providerSystemPrompt: string | undefined; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, getApiKey: async () => { markLookupStarted?.(); return apiKey; }, getSystemPrompt: () => systemPrompt, }; const stream = agentLoop( [createUserMessage("Hello")], { systemPrompt: "fallback", messages: [], tools: [] }, config, undefined, (_model, context, options) => { providerSystemPrompt = context.systemPrompt; expect(options?.apiKey).toBe("resolved key"); const mockStream = new MockAssistantStream(); queueMicrotask(() => { mockStream.push({ type: "done", reason: "stop", message: createAssistantMessage([{ type: "text", text: "Response" }]), }); }); return mockStream; }, ); const consume = (async () => { for await (const _event of stream) { } })();
await lookupStarted; systemPrompt = "after lookup"; resolveApiKey?.("resolved key"); await consume;
expect(providerSystemPrompt).toBe("after lookup"); });
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) => { 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) { }
expect(transformedMessages.length).toBe(2); expect(convertedMessages.length).toBe(2); });
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) { }
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) { }
expect(executed).toEqual([[{ oldText: "before", newText: "after" }]]); });
it("should emit tool_execution_end in completion order but persist 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 toolExecutionEndIds = events.flatMap((event) => { if (event.type !== "tool_execution_end") { return []; } return [event.toolCallId]; }); const toolResultIds = events.flatMap((event) => { if (event.type !== "message_end" || event.message.role !== "toolResult") { return []; } return [event.message.toolCallId]; }); const turnToolResultIds = events.flatMap((event) => { if (event.type !== "turn_end") { return []; } return event.toolResults.map((toolResult) => toolResult.toolCallId); });
expect(parallelObserved).toBe(true); expect(toolExecutionEndIds).toEqual(["tool-2", "tool-1"]); expect(toolResultIds).toEqual(["tool-1", "tool-2"]); expect(turnToolResultIds).toEqual(["tool-1", "tool-2"]); });
it("should measure and persist tool call duration", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool<typeof toolSchema> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute() { await new Promise((resolve) => setTimeout(resolve, 20)); return { content: [{ type: "text", text: "ok" }], details: {}, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const userPrompt: AgentMessage = createUserMessage("measure"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
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: "x" } }], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); } 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 endEvents = events.filter( (e): e is Extract<AgentEvent, { type: "tool_execution_end" }> => e.type === "tool_execution_end", ); expect(endEvents.length).toBe(1); expect(endEvents[0].durationMs).toBeGreaterThanOrEqual(15);
const toolResults = events.flatMap((event) => { if (event.type !== "message_end" || event.message.role !== "toolResult") { return []; } return [event.message]; }); expect(toolResults.length).toBe(1); expect(toolResults[0].durationMs).toBeGreaterThanOrEqual(15); });
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 () => { if (executed.length >= 1 && !queuedDelivered) { queuedDelivered = true; return [queuedUserMessage]; } return []; }, };
const events: AgentEvent[] = []; const stream = agentLoop([userPrompt], context, config, undefined, (_model, ctx, _options) => { 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) { 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 { 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); }
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);
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"));
expect(sawInterruptInContext).toBe(true); });
it("should inject continuation messages when the agent would otherwise stop", async () => { const context: AgentContext = { systemPrompt: "", messages: [], tools: [], };
let continuationPolls = 0; let sawContinuationInContext = false; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, getContinuationMessages: async ({ message }) => { continuationPolls++; expect(message.role).toBe("assistant"); if (continuationPolls === 1) { return [createUserMessage("continue")]; } return []; }, };
let callIndex = 0; const stream = agentLoop([createUserMessage("start")], context, config, undefined, (_model, ctx) => { if (callIndex === 1) { sawContinuationInContext = ctx.messages.some( (message) => message.role === "user" && message.content === "continue", ); } const mockStream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: callIndex === 0 ? "paused" : "done" }]); mockStream.push({ type: "done", reason: "stop", message }); callIndex++; }); return mockStream; });
for await (const _event of stream) { }
const messages = await stream.result(); expect(callIndex).toBe(2); expect(continuationPolls).toBe(2); expect(sawContinuationInContext).toBe(true); expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "user", "assistant"]); });
it("should prefer explicit follow-up messages before continuation messages", async () => { const context: AgentContext = { systemPrompt: "", messages: [], tools: [], };
let followUpDelivered = false; let continuationPolls = 0; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, getFollowUpMessages: async () => { if (followUpDelivered) { return []; } followUpDelivered = true; return [createUserMessage("follow up")]; }, getContinuationMessages: async () => { continuationPolls++; return []; }, };
let callIndex = 0; let sawFollowUpInContext = false; const stream = agentLoop([createUserMessage("start")], context, config, undefined, (_model, ctx) => { if (callIndex === 1) { sawFollowUpInContext = ctx.messages.some( (message) => message.role === "user" && message.content === "follow up", ); } const mockStream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage([{ type: "text", text: callIndex === 0 ? "paused" : "done" }]); mockStream.push({ type: "done", reason: "stop", message }); callIndex++; }); return mockStream; });
for await (const _event of stream) { }
expect(callIndex).toBe(2); expect(continuationPolls).toBe(1); expect(sawFollowUpInContext).toBe(true); });
it("should force sequential execution when a tool has executionMode=sequential even with default parallel config", 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 slowTool: AgentTool<typeof toolSchema, { value: string }> = { name: "slow", label: "Slow", description: "Slow tool", parameters: toolSchema, executionMode: "sequential", async execute(_toolCallId, params) { if (params.value === "first") { await firstDone; firstResolved = true; } if (params.value === "second" && !firstResolved) { parallelObserved = true; } return { content: [{ type: "text", text: `slow: ${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [slowTool], };
const userPrompt: AgentMessage = createUserMessage("run both"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
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: "slow", arguments: { value: "first" } }, { type: "toolCall", id: "tool-2", name: "slow", 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); }
expect(parallelObserved).toBe(false);
const toolResultIds = events.flatMap((event) => { if (event.type !== "message_end" || event.message.role !== "toolResult") { return []; } return [event.message.toolCallId]; }); expect(toolResultIds).toEqual(["tool-1", "tool-2"]); });
it("should force sequential execution when one of multiple tools has executionMode=sequential", async () => { const toolSchema = Type.Object({ value: Type.String() }); const executionOrder: string[] = []; let releaseSlow: (() => void) | undefined; const slowDone = new Promise<void>((resolve) => { releaseSlow = resolve; });
const slowTool: AgentTool<typeof toolSchema, { value: string }> = { name: "slow", label: "Slow", description: "Slow tool", parameters: toolSchema, executionMode: "sequential", async execute(_toolCallId, params) { executionOrder.push(`slow:${params.value}`); if (params.value === "a") { await slowDone; } return { content: [{ type: "text", text: `slow: ${params.value}` }], details: { value: params.value }, }; }, };
const fastTool: AgentTool<typeof toolSchema, { value: string }> = { name: "fast", label: "Fast", description: "Fast tool", parameters: toolSchema, async execute(_toolCallId, params) { executionOrder.push(`fast:${params.value}`); return { content: [{ type: "text", text: `fast: ${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [slowTool, fastTool], };
const userPrompt: AgentMessage = createUserMessage("run both"); const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
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: "slow", arguments: { value: "a" } }, { type: "toolCall", id: "tool-2", name: "fast", arguments: { value: "b" } }, ], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); setTimeout(() => releaseSlow?.(), 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); }
expect(executionOrder[0]).toBe("slow:a"); expect(executionOrder).toContain("fast:b"); });
it("should allow parallel execution when all tools have executionMode=parallel", 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, executionMode: "parallel", 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, };
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); }
expect(parallelObserved).toBe(true); });
const naturalCompletionCases: Array<{ name: string; stopBefore: "always" | "afterTurn" | "never"; steerOnce?: boolean; naturalWork?: boolean; expected: { llmCalls: number; stopChecks?: number; steeringPolls?: number; naturalPolls?: number }; }> = [ { name: "stops only after completing the initial turn", stopBefore: "always", expected: { llmCalls: 1, stopChecks: 1 }, }, { name: "observes a stop requested after the turn before polling steering", stopBefore: "afterTurn", expected: { llmCalls: 1, steeringPolls: 1 }, }, { name: "stops before polling natural-completion work", stopBefore: "always", naturalWork: true, expected: { llmCalls: 1, naturalPolls: 0 }, }, { name: "continues with steering when no stop is requested", stopBefore: "never", steerOnce: true, expected: { llmCalls: 2 }, }, ];
it.each(naturalCompletionCases)("$name", async ({ stopBefore, steerOnce, naturalWork, expected }) => { let llmCalls = 0; let stopRequested = false; let steeringDelivered = false; const shouldStopBeforeTurn = vi.fn(() => stopBefore === "always" || stopRequested); const getSteeringMessages = vi.fn(async () => { if (steerOnce && llmCalls === 1 && !steeringDelivered) { steeringDelivered = true; return [createUserMessage("steer")]; } return []; }); const getFollowUpMessages = vi.fn(async () => (naturalWork ? [createUserMessage("follow up")] : [])); const getContinuationMessages = vi.fn(async () => (naturalWork ? [createUserMessage("continue")] : [])); const stream = agentLoop( [createUserMessage("start")], { systemPrompt: "", messages: [], tools: [] }, { model: createModel(), convertToLlm: identityConverter, shouldStopBeforeTurn, getSteeringMessages, getFollowUpMessages, getContinuationMessages, ...(stopBefore === "afterTurn" && { shouldStopAfterTurn: async () => { await Promise.resolve(); stopRequested = true; return false; }, }), }, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => mockStream.push({ type: "done", reason: "stop", message: createAssistantMessage([{ type: "text", text: "done" }]), }), ); return mockStream; }, ); for await (const _event of stream) { }
expect(llmCalls).toBe(expected.llmCalls); if (expected.stopChecks !== undefined) expect(shouldStopBeforeTurn).toHaveBeenCalledTimes(expected.stopChecks); if (expected.steeringPolls !== undefined) expect(getSteeringMessages).toHaveBeenCalledTimes(expected.steeringPolls); if (expected.naturalPolls !== undefined) { expect(getFollowUpMessages).toHaveBeenCalledTimes(expected.naturalPolls); expect(getContinuationMessages).toHaveBeenCalledTimes(expected.naturalPolls); } });
it("rechecks a tool-cycle stop after polling steering", async () => { const execute = vi.fn(async () => ({ content: [{ type: "text" as const, text: "done" }], details: {} })); let stopRequested = false; let llmCalls = 0; const stream = agentLoop( [createUserMessage("start")], { systemPrompt: "", messages: [], tools: [{ name: "work", label: "Work", description: "Work", parameters: Type.Object({}), execute }], }, { model: createModel(), convertToLlm: identityConverter, shouldStopBeforeTurn: () => stopRequested, getSteeringMessages: async () => { if (llmCalls > 0) stopRequested = true; return []; }, }, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => mockStream.push({ type: "done", reason: "toolUse", message: createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "work", arguments: {} }], "toolUse", ), }), ); return mockStream; }, ); for await (const _event of stream) { } expect(llmCalls).toBe(1); });
it("injects steering drained by the poll even when a stop flips during it", async () => { const execute = vi.fn(async () => ({ content: [{ type: "text" as const, text: "done" }], details: {} })); let stopRequested = false; const injected: string[] = []; let llmCalls = 0; const stream = agentLoop( [createUserMessage("start")], { systemPrompt: "", messages: [], tools: [{ name: "work", label: "Work", description: "Work", parameters: Type.Object({}), execute }], }, { model: createModel(), convertToLlm: identityConverter, shouldStopBeforeTurn: () => stopRequested, getSteeringMessages: async () => { if (llmCalls !== 1 || stopRequested) return []; stopRequested = true; return [createUserMessage("late steer")]; }, }, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); const calls = llmCalls; queueMicrotask(() => mockStream.push({ type: "done", reason: calls === 1 ? "toolUse" : "stop", message: calls === 1 ? createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "work", arguments: {} }], "toolUse", ) : createAssistantMessage([{ type: "text", text: "steered reply" }], "stop"), }), ); return mockStream; }, ); for await (const event of stream) { if (event.type === "message_end" && event.message.role === "user") { const text = event.message.content; if (typeof text !== "string") continue; injected.push(text); } } expect(injected).toContain("late steer"); expect(llmCalls).toBe(2); });
it("checks shouldStopBeforeTurn after a tool batch and before another model call", async () => { const execute = vi.fn(async () => ({ content: [{ type: "text" as const, text: "done" }], details: {} })); const context: AgentContext = { systemPrompt: "", messages: [], tools: [ { name: "work", label: "Work", description: "Work", parameters: Type.Object({}), execute, }, ], }; let llmCalls = 0; const stream = agentLoop( [createUserMessage("start")], context, { model: createModel(), convertToLlm: identityConverter, shouldStopBeforeTurn: () => true }, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => mockStream.push({ type: "done", reason: "toolUse", message: createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "work", arguments: {} }], "toolUse", ), }), ); return mockStream; }, ); for await (const _event of stream) { }
expect(execute).toHaveBeenCalledOnce(); expect(llmCalls).toBe(1); });
it("should stop after the current turn when shouldStopAfterTurn returns true", 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], };
let steeringPolls = 0; let followUpPolls = 0; let callbackToolResultIds: string[] = []; let callbackContextRoles: string[] = []; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, getSteeringMessages: async () => { steeringPolls++; return []; }, getFollowUpMessages: async () => { followUpPolls++; return [createUserMessage("follow up should stay queued")]; }, shouldStopAfterTurn: async ({ message, toolResults, context }) => { expect(message.role).toBe("assistant"); callbackToolResultIds = toolResults.map((toolResult) => toolResult.toolCallId); callbackContextRoles = context.messages.map((contextMessage) => contextMessage.role); return true; }, };
let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => { if (llmCalls === 1) { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); } else { mockStream.push({ type: "done", reason: "stop", message: createAssistantMessage([{ type: "text", text: "should not run" }]), }); } }); return mockStream; });
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
const messages = await stream.result(); expect(llmCalls).toBe(1); expect(executed).toEqual(["hello"]); expect(steeringPolls).toBe(1); expect(followUpPolls).toBe(0); expect(callbackToolResultIds).toEqual(["tool-1"]); expect(callbackContextRoles).toEqual(["user", "assistant", "toolResult"]); expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "toolResult"]); expect(events.map((event) => event.type)).toEqual([ "agent_start", "turn_start", "message_start", "message_end", "message_start", "message_end", "tool_execution_start", "tool_execution_end", "message_start", "message_end", "turn_end", "agent_end", ]); });
it("should stop after a tool batch when every tool result sets terminate=true", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, terminate: true, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, };
let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); }); return mockStream; });
const events: AgentEvent[] = []; for await (const event of stream) { events.push(event); }
const messages = await stream.result(); expect(llmCalls).toBe(1); expect(messages.map((message) => message.role)).toEqual(["user", "assistant", "toolResult"]); expect(events.filter((event) => event.type === "turn_end")).toHaveLength(1); });
it("should continue after parallel tool calls when not all tool results terminate", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, terminate: params.value === "first", }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, toolExecution: "parallel", };
let callIndex = 0; const stream = agentLoop([createUserMessage("echo both")], 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 }); } else { const message = createAssistantMessage([{ type: "text", text: "done" }]); mockStream.push({ type: "done", reason: "stop", message }); } callIndex++; }); return mockStream; });
for await (const _event of stream) { }
const messages = await stream.result(); expect(callIndex).toBe(2); expect(messages.map((message) => message.role)).toEqual([ "user", "assistant", "toolResult", "toolResult", "assistant", ]); });
it("should allow afterToolCall to mark a tool batch as terminating", async () => { const toolSchema = Type.Object({ value: Type.String() }); const tool: AgentTool<typeof toolSchema, { value: string }> = { name: "echo", label: "Echo", description: "Echo tool", parameters: toolSchema, async execute(_toolCallId, params) { return { content: [{ type: "text", text: `echoed: ${params.value}` }], details: { value: params.value }, }; }, };
const context: AgentContext = { systemPrompt: "", messages: [], tools: [tool], };
const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, afterToolCall: async () => ({ terminate: true }), };
let llmCalls = 0; const stream = agentLoop([createUserMessage("echo something")], context, config, undefined, () => { llmCalls++; const mockStream = new MockAssistantStream(); queueMicrotask(() => { const message = createAssistantMessage( [{ type: "toolCall", id: "tool-1", name: "echo", arguments: { value: "hello" } }], "toolUse", ); mockStream.push({ type: "done", reason: "toolUse", message }); }); return mockStream; });
for await (const _event of stream) { }
expect(llmCalls).toBe(1); });});
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();
expect(messages.length).toBe(1); expect(messages[0].role).toBe("assistant");
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 () => { 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) => { 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; };
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"); });});