1import {

2 createInitialSystemMessage,

3 getCurrentSystemMessage,

4 getCurrentSystemPrompt,

5 type ImageContent,

6 type Message,

7 type Model,

8 type SimpleStreamOptions,

9 type TextContent,

10 type ThinkingBudgets,

11 type Transport,

12 toToolDeclaration,

13} from "@earendil-works/pi-ai";

14import { runAgentLoop, runAgentLoopContinue } from "./agent-loop.ts";

15import { getDefaultStreamFn } from "./stream-fn.ts";

16import type {

17 AfterToolCallContext,

18 AfterToolCallResult,

19 AgentContext,

20 AgentEvent,

21 AgentLoopConfig,

22 AgentLoopTurnUpdate,

23 AgentMessage,

24 AgentState,

25 AgentTool,

26 BeforeToolCallContext,

27 BeforeToolCallResult,

28 FinishTurn,

29 PrepareNextTurnContext,

30 PrepareRequest,

31 QueueMode,

32 StreamFn,

33 ToolExecutionMode,

34} from "./types.ts";

35

36export type { QueueMode } from "./types.ts";

37

38function defaultConvertToLlm(messages: AgentMessage[]): Message[] {

39 return messages.filter(

40 (message) =>

41 message.role === "system" ||

42 message.role === "user" ||

43 message.role === "assistant" ||

44 message.role === "toolResult",

45 );

46}

47

48const EMPTY_USAGE = {

49 input: 0,

50 output: 0,

51 cacheRead: 0,

52 cacheWrite: 0,

53 totalTokens: 0,

54 cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },

55};

56

57const DEFAULT_MODEL = {

58 id: "unknown",

59 name: "unknown",

60 api: "unknown",

61 provider: "unknown",

62 baseUrl: "",

63 reasoning: false,

64 input: [],

65 cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },

66 contextWindow: 0,

67 maxTokens: 0,

68} satisfies Model<any>;

69

70type MutableAgentState = Omit<AgentState, "isStreaming" | "streamingMessage" | "pendingToolCalls" | "errorMessage"> & {

71 isStreaming: boolean;

72 streamingMessage?: AgentMessage;

73 pendingToolCalls: Set<string>;

74 errorMessage?: string;

75};

76

77/** Initial state for {@link Agent}. `systemPrompt` and `tools` become the leading system message unless `messages` already starts with one. */

78export type AgentInitialState = Partial<

79 Omit<AgentState, "pendingToolCalls" | "isStreaming" | "streamingMessage" | "errorMessage">

80>;

81

82function createMutableAgentState(initialState?: AgentInitialState): MutableAgentState {

83 let tools = initialState?.tools?.slice() ?? [];

84 let messages = initialState?.messages?.slice() ?? [];

85 const initialMessage = createInitialSystemMessage(initialState?.systemPrompt, tools.map(toToolDeclaration));

86 if (messages[0]?.role !== "system" && initialMessage) messages.unshift(initialMessage);

87

88 return {

89 get systemPrompt() {

90 return getCurrentSystemPrompt(messages);

91 },

92 model: initialState?.model ?? DEFAULT_MODEL,

93 thinkingLevel: initialState?.thinkingLevel ?? "off",

94 get tools() {

95 return tools;

96 },

97 set tools(nextTools: AgentTool<any>[]) {

98 tools = nextTools.slice();

99 },

100 get messages() {

101 return messages;

102 },

103 set messages(nextMessages: AgentMessage[]) {

104 messages = nextMessages.slice();

105 },

106 isStreaming: false,

107 streamingMessage: undefined,

108 pendingToolCalls: new Set<string>(),

109 errorMessage: undefined,

110 };

111}

112

113/** Options for constructing an {@link Agent}. */

114export interface AgentOptions {

115 initialState?: AgentInitialState;

116 convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;

117 transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;

118 streamFn: StreamFn;

119 getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;

120 onPayload?: SimpleStreamOptions["onPayload"];

121 onResponse?: SimpleStreamOptions["onResponse"];

122 onProviderStreamEvent?: SimpleStreamOptions["onProviderStreamEvent"];

123 beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;

124 afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;

125 finishTurn?: FinishTurn;

126 prepareRequest?: PrepareRequest;

127 prepareNextTurn?: (

128 signal?: AbortSignal,

129 ) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;

130 prepareNextTurnWithContext?: (

131 context: PrepareNextTurnContext,

132 signal?: AbortSignal,

133 ) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;

134 steeringMode?: QueueMode;

135 followUpMode?: QueueMode;

136 sessionId?: string;

137 thinkingBudgets?: ThinkingBudgets;

138 transport?: Transport;

139 maxRetryDelayMs?: number;

140 toolExecution?: ToolExecutionMode;

141}

142

143class PendingMessageQueue {

144 private messages: AgentMessage[] = [];

145 public mode: QueueMode;

146

147 constructor(mode: QueueMode) {

148 this.mode = mode;

149 }

150

151 enqueue(message: AgentMessage): void {

152 this.messages.push(message);

153 }

154

155 hasItems(): boolean {

156 return this.messages.length > 0;

157 }

158

159 peek(): AgentMessage[] {

160 if (this.mode === "all") return this.messages.slice();

161 const first = this.messages[0];

162 return first ? [first] : [];

163 }

164

165 drain(): AgentMessage[] {

166 const drained = this.peek();

167 this.messages = this.messages.slice(drained.length);

168 return drained;

169 }

170

171 clear(): void {

172 this.messages = [];

173 }

174}

175

176type ActiveRun = {

177 promise: Promise<void>;

178 resolve: () => void;

179 abortController: AbortController;

180};

181

182/**

183 * Stateful wrapper around the low-level agent loop.

184 *

185 * `Agent` owns the current transcript, emits lifecycle events, executes tools,

186 * and exposes queueing APIs for steering and follow-up messages.

187 */

188export class Agent {

189 private _state: MutableAgentState;

190 private readonly listeners = new Set<(event: AgentEvent, signal: AbortSignal) => Promise<void> | void>();

191 private readonly steeringQueue: PendingMessageQueue;

192 private readonly followUpQueue: PendingMessageQueue;

193

194 public convertToLlm: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;

195 public transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;

196 public streamFunction: StreamFn;

197 public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;

198 public onPayload?: SimpleStreamOptions["onPayload"];

199 public onResponse?: SimpleStreamOptions["onResponse"];

200 public onProviderStreamEvent?: SimpleStreamOptions["onProviderStreamEvent"];

201 public beforeToolCall?: (

202 context: BeforeToolCallContext,

203 signal?: AbortSignal,

204 ) => Promise<BeforeToolCallResult | undefined>;

205 public afterToolCall?: (

206 context: AfterToolCallContext,

207 signal?: AbortSignal,

208 ) => Promise<AfterToolCallResult | undefined>;

209 public finishTurn?: FinishTurn;

210 public prepareRequest?: PrepareRequest;

211 public prepareNextTurn?: (

212 signal?: AbortSignal,

213 ) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;

214 public prepareNextTurnWithContext?: (

215 context: PrepareNextTurnContext,

216 signal?: AbortSignal,

217 ) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;

218 private activeRun?: ActiveRun;

219 /** Session identifier forwarded to providers for cache-aware backends. */

220 public sessionId?: string;

221 /** Optional per-level thinking token budgets forwarded to the stream function. */

222 public thinkingBudgets?: ThinkingBudgets;

223 /** Preferred transport forwarded to the stream function. */

224 public transport: Transport;

225 /** Optional cap for provider-requested retry delays. */

226 public maxRetryDelayMs?: number;

227 /** Tool execution strategy for assistant messages that contain multiple tool calls. */

228 public toolExecution: ToolExecutionMode;

229

230 constructor(options: AgentOptions) {

231 // Older compiled consumers may omit options or streamFn even though the current API requires them.

232 const runtimeOptions: Partial<AgentOptions> = options ?? {};

233 this._state = createMutableAgentState(runtimeOptions.initialState);

234 this.convertToLlm = runtimeOptions.convertToLlm ?? defaultConvertToLlm;

235 this.transformContext = runtimeOptions.transformContext;

236 this.streamFunction = runtimeOptions.streamFn ?? getDefaultStreamFn();

237 this.getApiKey = runtimeOptions.getApiKey;

238 this.onPayload = runtimeOptions.onPayload;

239 this.onResponse = runtimeOptions.onResponse;

240 this.onProviderStreamEvent = runtimeOptions.onProviderStreamEvent;

241 this.beforeToolCall = runtimeOptions.beforeToolCall;

242 this.afterToolCall = runtimeOptions.afterToolCall;

243 this.finishTurn = runtimeOptions.finishTurn;

244 this.prepareRequest = runtimeOptions.prepareRequest;

245 this.prepareNextTurn = runtimeOptions.prepareNextTurn;

246 this.prepareNextTurnWithContext = runtimeOptions.prepareNextTurnWithContext;

247 this.steeringQueue = new PendingMessageQueue(runtimeOptions.steeringMode ?? "one-at-a-time");

248 this.followUpQueue = new PendingMessageQueue(runtimeOptions.followUpMode ?? "one-at-a-time");

249 this.sessionId = runtimeOptions.sessionId;

250 this.thinkingBudgets = runtimeOptions.thinkingBudgets;

251 this.transport = runtimeOptions.transport ?? "auto";

252 this.maxRetryDelayMs = runtimeOptions.maxRetryDelayMs;

253 this.toolExecution = runtimeOptions.toolExecution ?? "parallel";

254 }

255

256 /**

257 * Subscribe to agent lifecycle events.

258 *

259 * Listener promises are awaited in subscription order and are included in

260 * the current run's settlement. Listeners also receive the active abort

261 * signal for the current run.

262 *

263 * `agent_end` is the final emitted event for a run, but the agent does not

264 * become idle until all awaited listeners for that event have settled.

265 */

266 subscribe(listener: (event: AgentEvent, signal: AbortSignal) => Promise<void> | void): () => void {

267 this.listeners.add(listener);

268 return () => this.listeners.delete(listener);

269 }

270

271 /**

272 * Current agent state.

273 *

274 * Assigning `state.tools` or `state.messages` copies the provided top-level array.

275 */

276 get state(): AgentState {

277 return this._state;

278 }

279

280 /** Controls how queued steering messages are drained. */

281 set steeringMode(mode: QueueMode) {

282 this.steeringQueue.mode = mode;

283 }

284

285 get steeringMode(): QueueMode {

286 return this.steeringQueue.mode;

287 }

288

289 /** Controls how queued follow-up messages are drained. */

290 set followUpMode(mode: QueueMode) {

291 this.followUpQueue.mode = mode;

292 }

293

294 get followUpMode(): QueueMode {

295 return this.followUpQueue.mode;

296 }

297

298 /** Queue a message to be injected after the current assistant turn finishes. */

299 steer(message: AgentMessage): void {

300 this.steeringQueue.enqueue(message);

301 }

302

303 /** Queue a message to run only after the agent would otherwise stop. */

304 followUp(message: AgentMessage): void {

305 this.followUpQueue.enqueue(message);

306 }

307

308 /** Remove all queued steering messages. */

309 clearSteeringQueue(): void {

310 this.steeringQueue.clear();

311 }

312

313 /** Remove all queued follow-up messages. */

314 clearFollowUpQueue(): void {

315 this.followUpQueue.clear();

316 }

317

318 /** Remove all queued steering and follow-up messages. */

319 clearAllQueues(): void {

320 this.clearSteeringQueue();

321 this.clearFollowUpQueue();

322 }

323

324 /** Returns true when either queue still contains pending messages. */

325 hasQueuedMessages(): boolean {

326 return this.steeringQueue.hasItems() || this.followUpQueue.hasItems();

327 }

328

329 /** Preview the messages selected for the next turn without consuming them. */

330 peekQueuedMessages(): AgentMessage[] {

331 const steering = this.steeringQueue.peek();

332 return steering.length > 0 ? steering : this.followUpQueue.peek();

333 }

334

335 /** Active abort signal for the current run, if any. */

336 get signal(): AbortSignal | undefined {

337 return this.activeRun?.abortController.signal;

338 }

339

340 /** Abort the current run, if one is active. */

341 abort(): void {

342 this.activeRun?.abortController.abort();

343 }

344

345 /**

346 * Resolve when the current run and all awaited event listeners have finished.

347 *

348 * This resolves after `agent_end` listeners settle.

349 */

350 waitForIdle(): Promise<void> {

351 return this.activeRun?.promise ?? Promise.resolve();

352 }

353

354 /** Clear conversation state and queues while retaining the replayed prompt/tool baseline. */

355 reset(): void {

356 if (this.activeRun) {

357 throw new Error("Agent is already processing. Wait for completion before resetting.");

358 }

359

360 const baseline = getCurrentSystemMessage(this._state.messages);

361 this._state.messages = baseline ? [baseline] : [];

362 this._state.isStreaming = false;

363 this._state.streamingMessage = undefined;

364 this._state.pendingToolCalls = new Set<string>();

365 this._state.errorMessage = undefined;

366 this.clearFollowUpQueue();

367 this.clearSteeringQueue();

368 }

369

370 /** Start a new prompt from text, a single message, or a batch of messages. */

371 async prompt(message: AgentMessage | AgentMessage[]): Promise<void>;

372 async prompt(input: string, images?: ImageContent[]): Promise<void>;

373 async prompt(input: string | AgentMessage | AgentMessage[], images?: ImageContent[]): Promise<void> {

374 if (this.activeRun) {

375 throw new Error(

376 "Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion.",

377 );

378 }

379 const messages = this.normalizePromptInput(input, images);

380 await this.runPromptMessages(messages);

381 }

382

383 /** Continue from the current transcript. The last message must be a user or tool-result message. */

384 async continue(): Promise<void> {

385 if (this.activeRun) {

386 throw new Error("Agent is already processing. Wait for completion before continuing.");

387 }

388

389 const lastMessage = this._state.messages[this._state.messages.length - 1];

390 if (!lastMessage || this._state.messages.every((message) => message.role === "system")) {

391 throw new Error("No messages to continue from");

392 }

393

394 if (lastMessage.role === "assistant") {

395 const queuedSteering = this.steeringQueue.drain();

396 if (queuedSteering.length > 0) {

397 await this.runPromptMessages(queuedSteering, { skipInitialSteeringPoll: true });

398 return;

399 }

400

401 const queuedFollowUps = this.followUpQueue.drain();

402 if (queuedFollowUps.length > 0) {

403 await this.runPromptMessages(queuedFollowUps);

404 return;

405 }

406

407 throw new Error("Cannot continue from message role: assistant");

408 }

409

410 await this.runContinuation();

411 }

412

413 private normalizePromptInput(

414 input: string | AgentMessage | AgentMessage[],

415 images?: ImageContent[],

416 ): AgentMessage[] {

417 if (Array.isArray(input)) {

418 return input;

419 }

420

421 if (typeof input !== "string") {

422 return [input];

423 }

424

425 const content: Array<TextContent | ImageContent> = [{ type: "text", text: input }];

426 if (images && images.length > 0) {

427 content.push(...images);

428 }

429 return [{ role: "user", content, timestamp: Date.now() }];

430 }

431

432 private async runPromptMessages(

433 messages: AgentMessage[],

434 options: { skipInitialSteeringPoll?: boolean } = {},

435 ): Promise<void> {

436 await this.runWithLifecycle(async (signal) => {

437 await runAgentLoop(

438 messages,

439 this.createContextSnapshot(),

440 this.createLoopConfig(options),

441 (event) => this.processEvents(event),

442 signal,

443 this.streamFunction,

444 );

445 });

446 }

447

448 private async runContinuation(): Promise<void> {

449 await this.runWithLifecycle(async (signal) => {

450 await runAgentLoopContinue(

451 this.createContextSnapshot(),

452 this.createLoopConfig(),

453 (event) => this.processEvents(event),

454 signal,

455 this.streamFunction,

456 );

457 });

458 }

459

460 private createContextSnapshot(): AgentContext {

461 return {

462 messages: this._state.messages.slice(),

463 tools: this._state.tools.slice(),

464 };

465 }

466

467 private createLoopConfig(options: { skipInitialSteeringPoll?: boolean } = {}): AgentLoopConfig {

468 let skipInitialSteeringPoll = options.skipInitialSteeringPoll === true;

469 return {

470 model: this._state.model,

471 reasoning: this._state.thinkingLevel === "off" ? undefined : this._state.thinkingLevel,

472 sessionId: this.sessionId,

473 onPayload: this.onPayload,

474 onResponse: this.onResponse,

475 onProviderStreamEvent: this.onProviderStreamEvent,

476 transport: this.transport,

477 thinkingBudgets: this.thinkingBudgets,

478 maxRetryDelayMs: this.maxRetryDelayMs,

479 toolExecution: this.toolExecution,

480 beforeToolCall: this.beforeToolCall,

481 afterToolCall: this.afterToolCall,

482 finishTurn: this.finishTurn,

483 prepareRequest: this.prepareRequest,

484 prepareNextTurn:

485 this.prepareNextTurnWithContext || this.prepareNextTurn

486 ? async (context) => {

487 if (this.prepareNextTurnWithContext) {

488 return await this.prepareNextTurnWithContext(context, this.signal);

489 }

490 return await this.prepareNextTurn?.(this.signal);

491 }

492 : undefined,

493 convertToLlm: this.convertToLlm,

494 transformContext: this.transformContext,

495 getApiKey: this.getApiKey,

496 getSteeringMessages: async () => {

497 if (skipInitialSteeringPoll) {

498 skipInitialSteeringPoll = false;

499 return [];

500 }

501 return this.steeringQueue.drain();

502 },

503 getFollowUpMessages: async () => this.followUpQueue.drain(),

504 };

505 }

506

507 private async runWithLifecycle(executor: (signal: AbortSignal) => Promise<void>): Promise<void> {

508 if (this.activeRun) {

509 throw new Error("Agent is already processing.");

510 }

511

512 const abortController = new AbortController();

513 let resolvePromise = () => {};

514 const promise = new Promise<void>((resolve) => {

515 resolvePromise = resolve;

516 });

517 this.activeRun = { promise, resolve: resolvePromise, abortController };

518

519 this._state.isStreaming = true;

520 this._state.streamingMessage = undefined;

521 this._state.errorMessage = undefined;

522

523 try {

524 await executor(abortController.signal);

525 } catch (error) {

526 await this.handleRunFailure(error, abortController.signal.aborted);

527 } finally {

528 this.finishRun();

529 }

530 }

531

532 private async handleRunFailure(error: unknown, aborted: boolean): Promise<void> {

533 const failureMessage = {

534 role: "assistant",

535 content: [{ type: "text", text: "" }],

536 api: this._state.model.api,

537 provider: this._state.model.provider,

538 model: this._state.model.id,

539 usage: EMPTY_USAGE,

540 stopReason: aborted ? "aborted" : "error",

541 errorMessage: error instanceof Error ? error.message : String(error),

542 timestamp: Date.now(),

543 } satisfies AgentMessage;

544 await this.processEvents({ type: "message_start", message: failureMessage });

545 await this.processEvents({ type: "message_end", message: failureMessage });

546 await this.processEvents({ type: "turn_end", message: failureMessage, toolResults: [] });

547 await this.processEvents({ type: "agent_end", messages: [failureMessage] });

548 }

549

550 private finishRun(): void {

551 this._state.isStreaming = false;

552 this._state.streamingMessage = undefined;

553 this._state.pendingToolCalls = new Set<string>();

554 this.activeRun?.resolve();

555 this.activeRun = undefined;

556 }

557

558 /**

559 * Reduce internal state for a loop event, then await listeners.

560 *

561 * `agent_end` only means no further loop events will be emitted. The run is

562 * considered idle later, after all awaited listeners for `agent_end` finish

563 * and `finishRun()` clears runtime-owned state.

564 */

565 private async processEvents(event: AgentEvent): Promise<void> {

566 switch (event.type) {

567 case "message_start":

568 this._state.streamingMessage = event.message;

569 break;

570

571 case "message_update":

572 this._state.streamingMessage = event.message;

573 break;

574

575 case "message_end":

576 this._state.streamingMessage = undefined;

577 this._state.messages.push(event.message);

578 break;

579

580 case "tool_execution_start": {

581 const pendingToolCalls = new Set(this._state.pendingToolCalls);

582 pendingToolCalls.add(event.toolCallId);

583 this._state.pendingToolCalls = pendingToolCalls;

584 break;

585 }

586

587 case "tool_execution_end": {

588 const pendingToolCalls = new Set(this._state.pendingToolCalls);

589 pendingToolCalls.delete(event.toolCallId);

590 this._state.pendingToolCalls = pendingToolCalls;

591 break;

592 }

593

594 case "turn_end":

595 if (event.message.role === "assistant" && event.message.errorMessage) {

596 this._state.errorMessage = event.message.errorMessage;

597 }

598 break;

599

600 case "agent_end":

601 this._state.streamingMessage = undefined;

602 break;

603 }

604

605 const signal = this.activeRun?.abortController.signal;

606 if (!signal) {

607 throw new Error("Agent listener invoked outside active run");

608 }

609 for (const listener of this.listeners) {

610 await listener(event, signal);

611 }

612 }

613}