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.
17 kB · 430 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431import type { AssistantMessage, AssistantMessageEvent, ImageContent, Message, Model, ServiceTier, SimpleStreamOptions, streamSimple, TextContent, Tool, ToolResultMessage,} from "@earendil-works/pi-ai";import type { Static, TSchema } from "typebox";
/** * Stream function used by the agent loop. * * Contract: * - Must not throw or return a rejected promise for request/model/runtime failures. * - Must return an AssistantMessageEventStream. * - Failures must be encoded in the returned stream via protocol events and a * final AssistantMessage with stopReason "error" or "aborted" and errorMessage. */export type StreamFn = ( ...args: Parameters<typeof streamSimple>) => ReturnType<typeof streamSimple> | Promise<ReturnType<typeof streamSimple>>;
/** * Configuration for how tool calls from a single assistant message are executed. * * - "sequential": each tool call is prepared, executed, and finalized before the next one starts. * - "parallel": tool calls are prepared sequentially, then allowed tools execute concurrently. * `tool_execution_end` is emitted in tool completion order after each tool is finalized, * while tool-result message artifacts are emitted later in assistant source order. */export type ToolExecutionMode = "sequential" | "parallel";
/** A tool-call content block emitted by an assistant message. */export type AgentToolCall = Extract<AssistantMessage["content"][number], { type: "toolCall" }>;
/** * Result returned from `beforeToolCall`. * * Returning `{ block: true }` prevents the tool from executing. The loop emits an error tool result instead. * `reason` becomes the text shown in that error result. If omitted, a default blocked message is used. */export interface BeforeToolCallResult { block?: boolean; reason?: string;}
/** * Partial override returned from `afterToolCall`. * * Merge semantics are field-by-field: * - `content`: if provided, replaces the tool result content array in full * - `details`: if provided, replaces the tool result details value in full * - `isError`: if provided, replaces the tool result error flag * - `terminate`: if provided, replaces the early-termination hint * * Omitted fields keep the original executed tool result values. * There is no deep merge for `content` or `details`. */export interface AfterToolCallResult { content?: (TextContent | ImageContent)[]; details?: unknown; isError?: boolean; /** * Hint that the agent should stop after the current tool batch. * Early termination only happens when every finalized tool result in the batch sets this to true. */ terminate?: boolean;}
/** Context passed to `beforeToolCall` after arguments are validated. */export interface BeforeToolCallContext { /** Assistant message that requested the call. */ assistantMessage: AssistantMessage; /** Raw tool-call block from `assistantMessage.content`. */ toolCall: AgentToolCall; /** Validated arguments for the target tool schema. */ args: unknown; /** Agent context when this call is prepared. */ context: AgentContext;}
/** Context passed to `afterToolCall`. */export interface AfterToolCallContext { /** Assistant message that requested the call. */ assistantMessage: AssistantMessage; /** Raw tool-call block from `assistantMessage.content`. */ toolCall: AgentToolCall; /** Validated arguments for the target tool schema. */ args: unknown; /** Executed result before any `afterToolCall` overrides. */ result: AgentToolResult<any>; /** Whether the executed result is currently treated as an error. */ isError: boolean; /** Agent context when this call is finalized. */ context: AgentContext;}
/** Context passed to `shouldStopAfterTurn` and `getContinuationMessages`. */export interface ShouldStopAfterTurnContext { /** Assistant message that completed the turn. */ message: AssistantMessage; /** Tool-result messages included in the preceding `turn_end` event. */ toolResults: ToolResultMessage[]; /** Context after appending the turn's assistant message and tool results. */ context: AgentContext; /** Messages returned by this invocation; prompts include initial prompts, continuations exclude prior context. */ newMessages: AgentMessage[];}
export type GetContinuationMessagesContext = ShouldStopAfterTurnContext;
export interface AgentLoopConfig extends SimpleStreamOptions { model: Model<any>;
/** * Converts AgentMessage[] to LLM-compatible Message[] before each LLM call. * * Each AgentMessage must be converted to a UserMessage, AssistantMessage, or ToolResultMessage * that the LLM can understand. AgentMessages that cannot be converted (e.g., UI-only notifications, * status messages) should be filtered out. * * Contract: must not throw or reject. Return a safe fallback value instead. * Throwing interrupts the low-level agent loop without producing a normal event sequence. * * @example * ```typescript * convertToLlm: (messages) => messages.flatMap(m => { * if (m.role === "custom") { * // Convert custom message to user message * return [{ role: "user", content: m.content, timestamp: m.timestamp }]; * } * if (m.role === "notification") { * // Filter out UI-only messages * return []; * } * // Pass through standard LLM messages * return [m]; * }) * ``` */ convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
/** * Optional transform applied to the context before `convertToLlm`. * * Use this for operations that work at the AgentMessage level: * - Context window management (pruning old messages) * - Injecting context from external sources * * Contract: must not throw or reject. Return the original messages or another * safe fallback value instead. * * @example * ```typescript * transformContext: async (messages) => { * if (estimateTokens(messages) > MAX_TOKENS) { * return pruneOldMessages(messages); * } * return messages; * } * ``` */ transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
/** Resolves the system prompt immediately before each LLM call. */ getSystemPrompt?: () => string;
/** * Resolves the model immediately before each LLM call. * * Lets a caller retarget a loop that is already running, for example when an * extension routes every turn by policy. Defaults to `model`. */ getModel?: () => Model<any>;
/** * Resolves an API key dynamically for each LLM call. * * Useful for short-lived OAuth tokens (e.g., GitHub Copilot) that may expire * during long-running tool execution phases. * * Contract: must not throw or reject. Return undefined when no key is available. */ getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
/** * Called after each turn fully completes and `turn_end` has been emitted. * * If it returns true, the loop emits `agent_end` and exits before polling steering or follow-up queues, * without starting another LLM call. The current assistant response and any tool executions finish normally. * * Use this to request a graceful stop after the current turn, e.g. before context gets too full. * * Contract: must not throw or reject. Throwing interrupts the low-level agent loop without producing a normal event sequence. */ shouldStopAfterTurn?: (context: ShouldStopAfterTurnContext) => boolean | Promise<boolean>;
/** * Called synchronously after a completed turn and before polling work for another turn. * Return true to emit `agent_end` without starting another provider call. Work returned by * an asynchronous poll owns that boundary; queue owners must suppress stale continuation * results if their higher-level stop condition changes while generating them. * The hook is never checked before the initial assistant turn. */ shouldStopBeforeTurn?: () => boolean;
/** * Returns steering messages to inject into the conversation mid-run. * * Called after the current assistant turn finishes executing its tool calls, unless `shouldStopAfterTurn` exits first. * If messages are returned, they are added to the context before the next LLM call. * Tool calls from the current assistant message are not skipped. * * Use this for "steering" the agent while it's working. * * Contract: must not throw or reject. Return [] when no steering messages are available. */ getSteeringMessages?: () => Promise<AgentMessage[]>;
/** * Returns follow-up messages to process after the agent would otherwise stop. * * Called when the agent has no more tool calls and no steering messages. * If messages are returned, they're added to the context and the agent * continues with another turn. * * Use this for follow-up messages that should wait until the agent finishes. * * Contract: must not throw or reject. Return [] when no follow-up messages are available. */ getFollowUpMessages?: () => Promise<AgentMessage[]>;
/** * Returns continuation messages when the agent would otherwise stop. * * Called after follow-up messages have been polled and none are available. * If messages are returned, they're added to the context and the agent * continues with another turn. * * Use this for host-owned continuation policies such as long-running goals. * Explicit follow-up messages always take precedence over continuation messages. * * Contract: must not throw or reject. Return [] when no continuation should run. */ getContinuationMessages?: (context: GetContinuationMessagesContext, signal?: AbortSignal) => Promise<AgentMessage[]>;
/** * Tool execution mode. Defaults to `"parallel"`. * Parallel mode preflights calls sequentially, executes allowed calls concurrently, emits * `tool_execution_end` in completion order, then emits tool-result messages in assistant source order. */ toolExecution?: ToolExecutionMode;
/** * Called before a tool is executed, after arguments have been validated. * * Return `{ block: true }` to prevent execution. The loop emits an error tool result instead. * The hook receives the agent abort signal and is responsible for honoring it. */ beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
/** * Called after a tool finishes executing, before `tool_execution_end` and tool-result message events are emitted. * * Return an `AfterToolCallResult` to override parts of the executed tool result: * - `content` replaces the full content array * - `details` replaces the full details payload * - `isError` replaces the error flag * - `terminate` replaces the early-termination hint * * Any omitted fields keep their original values. No deep merge is performed. * The hook receives the agent abort signal and is responsible for honoring it. */ afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;}
/** * Thinking/reasoning level for models that support it. * Note: "xhigh" and "max" are only supported by selected model families. Use model * thinking-level metadata from @earendil-works/pi-ai to detect support for a concrete model. */export type ThinkingLevel = "off" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max";
/** * Extensible interface for custom app messages. * Apps can extend via declaration merging: * * @example * ```typescript * declare module "@mariozechner/agent" { * interface CustomAgentMessages { * artifact: ArtifactMessage; * notification: NotificationMessage; * } * } * ``` */export interface CustomAgentMessages {}
export type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages];
/** * Public agent state. * * `tools` and `messages` use accessor properties so implementations can copy * assigned arrays before storing them. */export interface AgentState { /** System prompt sent with each model request. */ systemPrompt: string; /** Model used for future turns. */ model: Model<any>; /** Requested reasoning level for future turns. */ thinkingLevel: ThinkingLevel; /** Requested provider service tier for future turns. */ serviceTier: ServiceTier; /** Available tools. Assigning a new array copies its top-level array. */ set tools(tools: AgentTool<any>[]); get tools(): AgentTool<any>[]; /** Conversation transcript. Assigning a new array copies its top-level array. */ set messages(messages: AgentMessage[]); get messages(): AgentMessage[]; /** True while processing a prompt or continuation, including awaited `agent_end` listeners. */ readonly isStreaming: boolean; /** Partial assistant message for the active streamed response, if any. */ readonly streamingMessage?: AgentMessage; /** Tool-call IDs currently executing. */ readonly pendingToolCalls: ReadonlySet<string>; /** Error from the most recent failed or aborted assistant turn, if any. */ readonly errorMessage?: string;}
/** Final or partial result produced by a tool. */export interface AgentToolResult<T> { /** Text or image content returned to the model. */ content: (TextContent | ImageContent)[]; /** Structured details for logs or UI rendering. */ details: T; /** * Tool-declared failure flag. The loop reports the call as an error (tool result * message and `tool_execution_end`) when a tool sets it, which is the only way a * tool that returns a result instead of throwing can fail visibly. Throwing remains * the primary failure channel; `afterToolCall` can still override this flag. */ isError?: boolean; /** * Hint that the agent should stop after the current tool batch. * Early termination only happens when every finalized tool result in the batch sets this to true. */ terminate?: boolean;}
/** * Callback used by tools to publish coalesced, rate-limited execution updates. * Only the newest partial is guaranteed delivery; tools must not rely on every partial reaching clients. */export type AgentToolUpdateCallback<T = any> = (partialResult: AgentToolResult<T>) => void;
/** Tool definition used by the agent runtime. */export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any> extends Tool<TParameters> { /** Human-readable label for UI display. */ label: string; /** * Optional compatibility shim for raw tool-call arguments before schema validation. * Must return an object that matches `TParameters`. */ prepareArguments?: (args: unknown) => Static<TParameters>; /** Execute the tool call. Throw on failure instead of encoding errors in `content`. */ execute: ( toolCallId: string, params: Static<TParameters>, signal?: AbortSignal, onUpdate?: AgentToolUpdateCallback<TDetails>, ) => Promise<AgentToolResult<TDetails>>; /** * Per-tool execution mode override. * - "sequential": this tool must execute one at a time with other tool calls. * - "parallel": this tool can execute concurrently with other tool calls. * * If omitted, the default execution mode applies. */ executionMode?: ToolExecutionMode;}
/** Context snapshot passed to the low-level agent loop and tool hooks. */export interface AgentContext { /** System prompt included with the request. */ systemPrompt: string; /** Transcript visible to the model. */ messages: AgentMessage[]; /** Tools available for this run. */ tools?: AgentTool<any>[];}
/** * Events emitted by the Agent for UI updates. * * `agent_end` is the last event emitted for a run, but awaited `Agent.subscribe()` * listeners for that event are still part of run settlement. The agent becomes * idle only after those listeners finish. */export type AgentEvent = /** Starts and ends one agent run; `agent_end` carries all messages produced by that run. */ | { type: "agent_start" } | { type: "agent_end"; messages: AgentMessage[] } /** One assistant response and its resulting tool calls. */ | { type: "turn_start" } | { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[] } /** Lifecycle events for user, assistant, and tool-result messages. */ | { type: "message_start"; message: AgentMessage } /** Only emitted for assistant messages during streaming. */ | { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent } | { type: "message_end"; message: AgentMessage } /** Tool execution events; parallel calls may end in completion rather than source order. */ | { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any } | { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any } | { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean; durationMs?: number; };