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.
18 kB · 600 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601import { createAssistantMessageDiagnostic, type ImageContent, type Message, type Model, type SimpleStreamOptions, streamSimple, type TextContent, type ThinkingBudgets, type Transport,} from "@earendil-works/pi-ai";import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.js";import type { AfterToolCallContext, AfterToolCallResult, AgentContext, AgentEvent, AgentLoopConfig, AgentMessage, AgentState, AgentTool, BeforeToolCallContext, BeforeToolCallResult, GetContinuationMessagesContext, ShouldStopAfterTurnContext, StreamFn, ToolExecutionMode,} from "./types.js";
function defaultConvertToLlm(messages: AgentMessage[]): Message[] { return messages.filter( (message) => message.role === "user" || message.role === "assistant" || message.role === "toolResult", );}
const EMPTY_USAGE = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },};
const DEFAULT_MODEL = { id: "unknown", name: "unknown", api: "unknown", provider: "unknown", baseUrl: "", reasoning: false, input: [], cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, contextWindow: 0, maxTokens: 0,} satisfies Model<any>;
type QueueMode = "all" | "one-at-a-time";
type MutableAgentState = Omit<AgentState, "isStreaming" | "streamingMessage" | "pendingToolCalls" | "errorMessage"> & { isStreaming: boolean; streamingMessage?: AgentMessage; pendingToolCalls: Set<string>; errorMessage?: string;};
function createMutableAgentState( initialState?: Partial<Omit<AgentState, "pendingToolCalls" | "isStreaming" | "streamingMessage" | "errorMessage">>,): MutableAgentState { let tools = initialState?.tools?.slice() ?? []; let messages = initialState?.messages?.slice() ?? [];
return { systemPrompt: initialState?.systemPrompt ?? "", model: initialState?.model ?? DEFAULT_MODEL, thinkingLevel: initialState?.thinkingLevel ?? "off", serviceTier: initialState?.serviceTier ?? "default", get tools() { return tools; }, set tools(nextTools: AgentTool<any>[]) { tools = nextTools.slice(); }, get messages() { return messages; }, set messages(nextMessages: AgentMessage[]) { messages = nextMessages.slice(); }, isStreaming: false, streamingMessage: undefined, pendingToolCalls: new Set<string>(), errorMessage: undefined, };}
export interface AgentOptions { initialState?: Partial<Omit<AgentState, "pendingToolCalls" | "isStreaming" | "streamingMessage" | "errorMessage">>; convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise<Message[]>; transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>; streamFn?: StreamFn; getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined; onPayload?: SimpleStreamOptions["onPayload"]; onResponse?: SimpleStreamOptions["onResponse"]; beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>; afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>; shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise<boolean>; shouldStopBeforeTurn?: () => boolean; getContinuationMessages?: (context: GetContinuationMessagesContext, signal?: AbortSignal) => Promise<AgentMessage[]>; steeringMode?: QueueMode; followUpMode?: QueueMode; sessionId?: string; thinkingBudgets?: ThinkingBudgets; transport?: Transport; toolExecution?: ToolExecutionMode;}
class PendingMessageQueue { private batches: AgentMessage[][] = [];
constructor(public mode: QueueMode) {}
enqueue(message: AgentMessage | AgentMessage[]): void { const batch = Array.isArray(message) ? message.slice() : [message]; if (batch.length > 0) { this.batches.push(batch); } }
hasItems(): boolean { return this.batches.length > 0; }
drain(): AgentMessage[] { if (this.mode === "all") { const drained = this.batches.flat(); this.batches = []; return drained; }
const first = this.batches[0]; if (!first) { return []; } this.batches = this.batches.slice(1); return first; }
clear(): void { this.batches = []; }
removeWhere(predicate: (message: AgentMessage) => boolean): AgentMessage[] { const removed: AgentMessage[] = []; const retained: AgentMessage[][] = []; for (const batch of this.batches) { if (batch.some(predicate)) { removed.push(...batch); } else { retained.push(batch); } } this.batches = retained; return removed; }}
type ActiveRun = { promise: Promise<void>; resolve: () => void; abortController: AbortController;};
/** Why {@link Agent.continue} refused to start a continuation. */export type AgentContinueErrorCode = "busy" | "nothing-to-continue";
/** Typed precondition failure from {@link Agent.continue}, so callers classify by code instead of message text. */export class AgentContinueError extends Error { constructor( readonly code: AgentContinueErrorCode, message: string, ) { super(message); this.name = "AgentContinueError"; }}
export class Agent { private _state: MutableAgentState; private readonly listeners = new Set<(event: AgentEvent, signal: AbortSignal) => Promise<void> | void>(); private readonly steeringQueue: PendingMessageQueue; private readonly followUpQueue: PendingMessageQueue;
public convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>; public transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>; public streamFn: StreamFn; public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined; public onPayload?: SimpleStreamOptions["onPayload"]; public onResponse?: SimpleStreamOptions["onResponse"]; public beforeToolCall?: ( context: BeforeToolCallContext, signal?: AbortSignal, ) => Promise<BeforeToolCallResult | undefined>; public afterToolCall?: ( context: AfterToolCallContext, signal?: AbortSignal, ) => Promise<AfterToolCallResult | undefined>; public shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise<boolean>; public shouldStopBeforeTurn?: () => boolean; public getContinuationMessages?: ( context: GetContinuationMessagesContext, signal?: AbortSignal, ) => Promise<AgentMessage[]>; private activeRun?: ActiveRun; public sessionId?: string; public thinkingBudgets?: ThinkingBudgets; public transport: Transport; public toolExecution: ToolExecutionMode;
constructor(options: AgentOptions = {}) { this._state = createMutableAgentState(options.initialState); this.convertToLlm = options.convertToLlm ?? defaultConvertToLlm; this.transformContext = options.transformContext; this.streamFn = options.streamFn ?? streamSimple; this.getApiKey = options.getApiKey; this.onPayload = options.onPayload; this.onResponse = options.onResponse; this.beforeToolCall = options.beforeToolCall; this.afterToolCall = options.afterToolCall; this.shouldStopAfterTurn = options.shouldStopAfterTurn; this.shouldStopBeforeTurn = options.shouldStopBeforeTurn; this.getContinuationMessages = options.getContinuationMessages; this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time"); this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time"); this.sessionId = options.sessionId; this.thinkingBudgets = options.thinkingBudgets; this.transport = options.transport ?? "auto"; this.toolExecution = options.toolExecution ?? "parallel"; }
/** * Subscribe to agent lifecycle events. * * Listener promises are awaited in subscription order and are included in * the current run's settlement. Listeners also receive the active abort * signal for the current run. * * `agent_end` is the final emitted event for a run, but the agent does not * become idle until all awaited listeners for that event have settled. */ subscribe(listener: (event: AgentEvent, signal: AbortSignal) => Promise<void> | void): () => void { this.listeners.add(listener); return () => this.listeners.delete(listener); }
/** * Current agent state. * * Assigning `state.tools` or `state.messages` copies the provided top-level array. */ get state(): AgentState { return this._state; }
set steeringMode(mode: QueueMode) { this.steeringQueue.mode = mode; }
get steeringMode(): QueueMode { return this.steeringQueue.mode; }
set followUpMode(mode: QueueMode) { this.followUpQueue.mode = mode; }
get followUpMode(): QueueMode { return this.followUpQueue.mode; }
/** Queue a message batch to be injected after the current assistant turn finishes. */ steer(message: AgentMessage | AgentMessage[]): void { this.steeringQueue.enqueue(message); }
/** Queue a message batch to run only after the agent would otherwise stop. */ followUp(message: AgentMessage | AgentMessage[]): void { this.followUpQueue.enqueue(message); }
clearSteeringQueue(): void { this.steeringQueue.clear(); }
clearFollowUpQueue(): void { this.followUpQueue.clear(); }
clearAllQueues(): void { this.clearSteeringQueue(); this.clearFollowUpQueue(); }
removeQueuedMessages(predicate: (message: AgentMessage) => boolean): AgentMessage[] { return [...this.steeringQueue.removeWhere(predicate), ...this.followUpQueue.removeWhere(predicate)]; }
hasQueuedMessages(): boolean { return this.steeringQueue.hasItems() || this.followUpQueue.hasItems(); }
get signal(): AbortSignal | undefined { return this.activeRun?.abortController.signal; }
abort(): void { this.activeRun?.abortController.abort(); }
/** * Resolve when the current run and all awaited event listeners have finished. * * This resolves after `agent_end` listeners settle. */ waitForIdle(): Promise<void> { return this.activeRun?.promise ?? Promise.resolve(); }
reset(): void { this._state.messages = []; this._state.isStreaming = false; this._state.streamingMessage = undefined; this._state.pendingToolCalls = new Set<string>(); this._state.errorMessage = undefined; this.clearFollowUpQueue(); this.clearSteeringQueue(); }
async prompt(message: AgentMessage | AgentMessage[]): Promise<void>; async prompt(input: string, images?: ImageContent[]): Promise<void>; async prompt(input: string | AgentMessage | AgentMessage[], images?: ImageContent[]): Promise<void> { if (this.activeRun) { throw new Error( "Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.", ); } const messages = this.normalizePromptInput(input, images); await this.runPromptMessages(messages); }
/** The last message must convert to a user or tool-result message. */ async continue(): Promise<void> { if (this.activeRun) { throw new AgentContinueError("busy", "Agent is already processing. Wait for completion before continuing."); }
const runQueuedMessages = (): Promise<void> | undefined => { const queuedSteering = this.steeringQueue.drain(); if (queuedSteering.length > 0) { return this.runPromptMessages(queuedSteering, { skipInitialSteeringPoll: true }); }
const queuedFollowUps = this.followUpQueue.drain(); if (queuedFollowUps.length > 0) { return this.runPromptMessages(queuedFollowUps); }
return undefined; };
const lastMessage = this._state.messages[this._state.messages.length - 1]; if (!lastMessage) { const queuedRun = runQueuedMessages(); if (queuedRun) { await queuedRun; return; }
throw new AgentContinueError("nothing-to-continue", "No messages to continue from"); }
if (lastMessage.role === "assistant") { const queuedRun = runQueuedMessages(); if (queuedRun) { await queuedRun; return; }
throw new AgentContinueError("nothing-to-continue", "Cannot continue from message role: assistant"); }
const lastMessageRole: string = lastMessage.role; if (lastMessageRole === "custom") { const queuedRun = runQueuedMessages(); if (queuedRun) { await queuedRun; return; } }
await this.runContinuation(); }
private normalizePromptInput( input: string | AgentMessage | AgentMessage[], images?: ImageContent[], ): AgentMessage[] { if (Array.isArray(input)) { return input; }
if (typeof input !== "string") { return [input]; }
const content: Array<TextContent | ImageContent> = [{ type: "text", text: input }]; if (images && images.length > 0) { content.push(...images); } return [{ role: "user", content, timestamp: Date.now() }]; }
private async runPromptMessages( messages: AgentMessage[], options: { skipInitialSteeringPoll?: boolean } = {}, ): Promise<void> { await this.runWithLifecycle(async (signal) => { await runAgentLoop( messages, this.createContextSnapshot(), this.createLoopConfig(options), (event) => this.processEvents(event), signal, this.streamFn, ); }); }
private async runContinuation(): Promise<void> { await this.runWithLifecycle(async (signal) => { await runAgentLoopContinue( this.createContextSnapshot(), this.createLoopConfig(), (event) => this.processEvents(event), signal, this.streamFn, ); }); }
private createContextSnapshot(): AgentContext { return { systemPrompt: this._state.systemPrompt, messages: this._state.messages.slice(), tools: this._state.tools.slice(), }; }
private createLoopConfig(options: { skipInitialSteeringPoll?: boolean } = {}): AgentLoopConfig { let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true; return { model: this._state.model, reasoning: this._state.thinkingLevel, serviceTier: this._state.serviceTier, sessionId: this.sessionId, onPayload: this.onPayload, onResponse: this.onResponse, transport: this.transport, thinkingBudgets: this.thinkingBudgets, toolExecution: this.toolExecution, beforeToolCall: this.beforeToolCall, afterToolCall: this.afterToolCall, shouldStopAfterTurn: async (context) => this.shouldStopAfterTurn?.(context) ?? false, shouldStopBeforeTurn: () => this.shouldStopBeforeTurn?.() ?? false, convertToLlm: this.convertToLlm, transformContext: this.transformContext, getSystemPrompt: () => this._state.systemPrompt, getApiKey: this.getApiKey, getSteeringMessages: async () => { if (skipInitialSteeringPoll) { skipInitialSteeringPoll = false; return []; } return this.steeringQueue.drain(); }, getFollowUpMessages: async () => this.followUpQueue.drain(), getContinuationMessages: async (context, signal) => this.getContinuationMessages?.(context, signal) ?? [], }; }
private async runWithLifecycle(executor: (signal: AbortSignal) => Promise<void>): Promise<void> { if (this.activeRun) { throw new Error("Agent is already processing."); }
const abortController = new AbortController(); let resolvePromise = () => {}; const promise = new Promise<void>((resolve) => { resolvePromise = resolve; }); this.activeRun = { promise, resolve: resolvePromise, abortController };
this._state.isStreaming = true; this._state.streamingMessage = undefined; this._state.errorMessage = undefined;
try { await executor(abortController.signal); } catch (error) { await this.handleRunFailure(error, abortController.signal.aborted); } finally { this.finishRun(); } }
private async handleRunFailure(error: unknown, aborted: boolean): Promise<void> { const failureMessage = { role: "assistant", content: [{ type: "text", text: "" }], api: this._state.model.api, provider: this._state.model.provider, model: this._state.model.id, usage: EMPTY_USAGE, stopReason: aborted ? "aborted" : "error", errorMessage: error instanceof Error ? error.message : String(error), diagnostics: aborted ? undefined : [createAssistantMessageDiagnostic("agent_lifecycle_failure", error, { source: "run_with_lifecycle" })], timestamp: Date.now(), } satisfies AgentMessage; this._state.errorMessage = failureMessage.errorMessage; await this.processEvents({ type: "message_start", message: failureMessage }).catch(() => undefined); await this.processEvents({ type: "message_end", message: failureMessage }).catch(() => undefined); await this.processEvents({ type: "agent_end", messages: [failureMessage] }).catch(() => undefined); }
private finishRun(): void { this._state.isStreaming = false; this._state.streamingMessage = undefined; this._state.pendingToolCalls = new Set<string>(); this.activeRun?.resolve(); this.activeRun = undefined; }
/** * Reduce internal state for a loop event, then await listeners. * * `agent_end` only means no further loop events will be emitted. The run is * considered idle later, after all awaited listeners for `agent_end` finish * and `finishRun()` clears runtime-owned state. */ private async processEvents(event: AgentEvent): Promise<void> { switch (event.type) { case "message_start": this._state.streamingMessage = event.message; break;
case "message_update": this._state.streamingMessage = event.message; break;
case "message_end": this._state.streamingMessage = undefined; this._state.messages.push(event.message); break;
case "tool_execution_start": { const pendingToolCalls = new Set(this._state.pendingToolCalls); pendingToolCalls.add(event.toolCallId); this._state.pendingToolCalls = pendingToolCalls; break; }
case "tool_execution_end": { const pendingToolCalls = new Set(this._state.pendingToolCalls); pendingToolCalls.delete(event.toolCallId); this._state.pendingToolCalls = pendingToolCalls; break; }
case "turn_end": if (event.message.role === "assistant" && event.message.errorMessage) { this._state.errorMessage = event.message.errorMessage; } break;
case "agent_end": this._state.streamingMessage = undefined; break; }
const signal = this.activeRun?.abortController.signal; if (!signal) { throw new Error("Agent listener invoked outside active run"); } for (const listener of this.listeners) { await listener(event, signal); } }}