import {
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;
maxRetryDelayMs?: number;
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 maxRetryDelayMs?: number;
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.maxRetryDelayMs = options.maxRetryDelayMs;
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,
maxRetryDelayMs: this.maxRetryDelayMs,
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);
}
}
}