1import type {

2 Api,

3 AssistantMessage,

4 AssistantMessageEvent,

5 AssistantMessageEventStream,

6 ImageContent,

7 JsonValue,

8 Message,

9 Model,

10 SimpleStreamOptions,

11 TextContent,

12 Tool,

13 ToolResultMessage,

14 TranscriptContext,

15 Usage,

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

17import type { Static, TSchema } from "typebox";

18

19/**

20 * Stream function used by the agent loop. `Models.streamSimple` satisfies

21 * this shape.

22 *

23 * The loop passes a normalized transcript: the system prompt and tool

24 * declarations are carried by the transcript's system messages, never by

25 * `context.systemPrompt` or `context.tools`.

26 *

27 * Contract:

28 * - Must not throw or return a rejected promise for request/model/runtime failures.

29 * - Must return an AssistantMessageEventStream.

30 * - Failures must be encoded in the returned stream via protocol events and a

31 * final AssistantMessage with stopReason "error" or "aborted" and errorMessage.

32 */

33export type StreamFn = (

34 model: Model<Api>,

35 context: TranscriptContext,

36 options?: SimpleStreamOptions,

37) => AssistantMessageEventStream | Promise<AssistantMessageEventStream>;

38

39/**

40 * Configuration for how tool calls from a single assistant message are executed.

41 *

42 * - "sequential": each tool call is prepared, executed, and finalized before the next one starts.

43 * - "parallel": tool calls are prepared sequentially, then allowed tools execute concurrently.

44 * `tool_execution_end` is emitted in tool completion order after each tool is finalized,

45 * while tool-result message artifacts are emitted later in assistant source order.

46 */

47export type ToolExecutionMode = "sequential" | "parallel";

48

49/**

50 * Controls how many queued user messages are injected when the agent loop reaches a queue drain point.

51 *

52 * - "all": drain and inject every queued message at that point.

53 * - "one-at-a-time": drain and inject only the oldest queued message, leaving the rest queued for later drain points.

54 */

55export type QueueMode = "all" | "one-at-a-time";

56

57/** A single tool call content block emitted by an assistant message. */

58export type AgentToolCall = Extract<AssistantMessage["content"][number], { type: "toolCall" }>;

59

60/**

61 * Result returned from `beforeToolCall`.

62 *

63 * Returning `{ block: true }` prevents the tool from executing. The loop emits an error tool result instead.

64 * `reason` becomes the text shown in that error result. If omitted, a default blocked message is used.

65 */

66export interface BeforeToolCallResult {

67 block?: boolean;

68 reason?: string;

69 /**

70 * Hint that the agent should stop after the current tool batch when this call is blocked.

71 * Early termination only happens when every finalized tool result in the batch sets this to true.

72 */

73 terminate?: boolean;

74}

75

76/**

77 * Partial override returned from `afterToolCall`.

78 *

79 * Merge semantics are field-by-field:

80 * - `content`: if provided, replaces the tool result content array in full

81 * - `details`: if provided, replaces the tool result details value in full

82 * - `isError`: if provided, replaces the tool result error flag

83 * - `usage`: if provided, replaces the tool result usage

84 * - `terminate`: if provided, replaces the early-termination hint

85 * - `structuredContent`: if provided, replaces the structured content. If `content` is provided

86 * without it, the structured content is dropped, because it may no longer match the content.

87 * Return it along with `content` to keep it.

88 *

89 * Other omitted fields keep the original executed tool result values.

90 * There is no deep merge for `content`, `details`, or `usage`.

91 */

92export interface AfterToolCallResult {

93 content?: (TextContent | ImageContent)[];

94 details?: unknown;

95 structuredContent?: JsonValue;

96 isError?: boolean;

97 /** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */

98 usage?: Usage;

99 /**

100 * Hint that the agent should stop after the current tool batch.

101 * Early termination only happens when every finalized tool result in the batch sets this to true.

102 */

103 terminate?: boolean;

104}

105

106/** Context passed to `beforeToolCall`. */

107export interface BeforeToolCallContext {

108 /** The assistant message that requested the tool call. */

109 assistantMessage: AssistantMessage;

110 /** The raw tool call block from `assistantMessage.content`. */

111 toolCall: AgentToolCall;

112 /** Validated tool arguments for the target tool schema. */

113 args: unknown;

114 /** Current agent context at the time the tool call is prepared. */

115 context: AgentContext;

116}

117

118/** Context passed to `afterToolCall`. */

119export interface AfterToolCallContext {

120 /** The assistant message that requested the tool call. */

121 assistantMessage: AssistantMessage;

122 /** The raw tool call block from `assistantMessage.content`. */

123 toolCall: AgentToolCall;

124 /** Validated tool arguments for the target tool schema. */

125 args: unknown;

126 /** The executed tool result before any `afterToolCall` overrides are applied. */

127 result: AgentToolResult<any>;

128 /** Whether the executed tool result is currently treated as an error. */

129 isError: boolean;

130 /** Current agent context at the time the tool call is finalized. */

131 context: AgentContext;

132}

133

134/** Context passed to completed-turn callbacks. */

135export interface AgentTurnContext {

136 /** The assistant message that completed the turn. */

137 message: AssistantMessage;

138 /** Tool result messages emitted for the completed turn. */

139 toolResults: ToolResultMessage[];

140 /** Current agent context after the turn's assistant message and tool results have been appended. */

141 context: AgentContext;

142 /** Messages that this loop invocation will return if it exits at this point. Prompt runs include the initial prompt messages; continuation runs do not include pre-existing context messages. */

143 newMessages: AgentMessage[];

144}

145

146/** Decision returned by {@link FinishTurn}. Returning undefined preserves normal scheduling. */

147export type AgentTurnDecision = { action: "continue" } | { action: "end" };

148

149/**

150 * Called after a completed assistant turn and all of its tool-result messages, but before `turn_end`.

151 * On a normal turn, `{ action: "continue" }` ensures one next provider request. Tool-result, steering, or

152 * follow-up scheduling can satisfy that request and adds no extra request; otherwise the loop continues once

153 * with the current context. Error and aborted responses remain hard exits.

154 */

155export type FinishTurn = (

156 turn: AgentTurnContext,

157 signal?: AbortSignal,

158) => AgentTurnDecision | void | Promise<AgentTurnDecision | undefined> | Promise<void>;

159

160/** Replacement runtime state used by the agent loop before starting another provider request. */

161export interface AgentLoopTurnUpdate {

162 /** Context for the next provider request. */

163 context?: AgentContext;

164 /** Messages to append before the next provider request, with normal lifecycle events. */

165 messages?: AgentMessage[];

166 /** Model for the next provider request. */

167 model?: Model<any>;

168 /** Thinking level for the next provider request. */

169 thinkingLevel?: ThinkingLevel;

170}

171

172/** Runtime state available immediately before a conversational provider request. */

173export interface PrepareRequestContext {

174 context: AgentContext;

175 model: Model<any>;

176 thinkingLevel: ThinkingLevel;

177}

178

179/** Replacement runtime state for the provider request being prepared. */

180export type AgentRequestUpdate = Omit<AgentLoopTurnUpdate, "messages">;

181

182/**

183 * Called immediately before every conversational provider request, including the first.

184 * Pending messages have already been appended and emitted when this callback runs.

185 */

186export type PrepareRequest = (

187 request: PrepareRequestContext,

188 signal?: AbortSignal,

189) => AgentRequestUpdate | void | Promise<AgentRequestUpdate | undefined> | Promise<void>;

190

191export interface PrepareNextTurnContext extends AgentTurnContext {}

192

193export interface AgentLoopConfig extends SimpleStreamOptions {

194 model: Model<any>;

195

196 /**

197 * Converts AgentMessage[] to LLM-compatible Message[] before each LLM call.

198 *

199 * Each AgentMessage must be converted to a SystemMessage, UserMessage, AssistantMessage, or ToolResultMessage

200 * that the LLM can understand. AgentMessages that cannot be converted (e.g., UI-only notifications,

201 * status messages) should be filtered out.

202 *

203 * Contract: must not throw or reject. Return a safe fallback value instead.

204 * Throwing interrupts the low-level agent loop without producing a normal event sequence.

205 *

206 * @example

207 * ```typescript

208 * convertToLlm: (messages) => messages.flatMap(m => {

209 * if (m.role === "custom") {

210 * // Convert custom message to user message

211 * return [{ role: "user", content: m.content, timestamp: m.timestamp }];

212 * }

213 * if (m.role === "notification") {

214 * // Filter out UI-only messages

215 * return [];

216 * }

217 * // Pass through standard LLM messages

218 * return [m];

219 * })

220 * ```

221 */

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

223

224 /**

225 * Optional transform applied to the context before `convertToLlm`.

226 *

227 * Use this for operations that work at the AgentMessage level:

228 * - Context window management (pruning old messages)

229 * - Injecting context from external sources

230 *

231 * Contract: must not throw or reject. Return the original messages or another

232 * safe fallback value instead.

233 *

234 * @example

235 * ```typescript

236 * transformContext: async (messages) => {

237 * if (estimateTokens(messages) > MAX_TOKENS) {

238 * return pruneOldMessages(messages);

239 * }

240 * return messages;

241 * }

242 * ```

243 */

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

245

246 /**

247 * Resolves an API key dynamically for each LLM call.

248 *

249 * Useful for short-lived OAuth tokens (e.g., GitHub Copilot) that may expire

250 * during long-running tool execution phases.

251 *

252 * Contract: must not throw or reject. Return undefined when no key is available.

253 */

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

255

256 /**

257 * Called after the assistant message and all tool-result messages have been emitted, immediately before `turn_end`.

258 * `{ action: "end" }` ends the run without polling queues or preparing another request.

259 * On a normal turn, `{ action: "continue" }` ensures one next provider request. Tool-result, steering, or

260 * follow-up scheduling can satisfy that request and adds no extra request; otherwise the loop continues once

261 * with the current context. Returning undefined preserves normal scheduling. Error and aborted responses remain

262 * hard exits.

263 */

264 finishTurn?: FinishTurn;

265

266 /**

267 * Called immediately before every conversational provider request, including the first.

268 * Pending messages have already been appended. The returned context, model, and thinking level

269 * replace the runtime values for this and later requests in the run. This hook does not poll queues.

270 */

271 prepareRequest?: PrepareRequest;

272

273 /**

274 * Called after `turn_end` when the loop will continue, immediately before the next turn starts.

275 * Return replacement context/model/thinking state or messages to append to affect that turn.

276 * Return undefined to keep using the current context/config.

277 */

278 prepareNextTurn?: (

279 context: PrepareNextTurnContext,

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

281

282 /**

283 * Returns steering messages to inject into the conversation mid-run.

284 *

285 * Called after the current assistant turn finishes executing its tool calls, unless `finishTurn` ends the run.

286 * If messages are returned, they are added to the context before the next LLM call.

287 * Tool calls from the current assistant message are not skipped.

288 *

289 * Use this for "steering" the agent while it's working.

290 *

291 * Contract: must not throw or reject. Return [] when no steering messages are available.

292 */

293 getSteeringMessages?: () => Promise<AgentMessage[]>;

294

295 /**

296 * Returns follow-up messages to process after the agent would otherwise stop.

297 *

298 * Called when the agent has no more tool calls and no steering messages.

299 * If messages are returned, they're added to the context and the agent

300 * continues with another turn.

301 *

302 * Use this for follow-up messages that should wait until the agent finishes.

303 *

304 * Contract: must not throw or reject. Return [] when no follow-up messages are available.

305 */

306 getFollowUpMessages?: () => Promise<AgentMessage[]>;

307

308 /**

309 * Tool execution mode.

310 * - "sequential": execute tool calls one by one

311 * - "parallel": preflight tool calls sequentially, then execute allowed tools concurrently;

312 * emit `tool_execution_end` in tool completion order after each tool is finalized,

313 * then emit tool-result message artifacts later in assistant source order

314 *

315 * Default: "parallel"

316 */

317 toolExecution?: ToolExecutionMode;

318

319 /**

320 * Called before a tool is executed, after arguments have been validated.

321 *

322 * Return `{ block: true }` to prevent execution. The loop emits an error tool result instead.

323 * A blocked result can also set `terminate: true` to participate in the batch early-termination rule.

324 * The hook receives the agent abort signal and is responsible for honoring it.

325 */

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

327

328 /**

329 * Called after a tool finishes executing, before `tool_execution_end` and tool-result message events are emitted.

330 *

331 * Return an `AfterToolCallResult` to override parts of the executed tool result:

332 * - `content` replaces the full content array

333 * - `details` replaces the full details payload

334 * - `isError` replaces the error flag

335 * - `usage` replaces the tool result usage

336 * - `terminate` replaces the early-termination hint

337 *

338 * Any omitted fields keep their original values. No deep merge is performed.

339 * The hook receives the agent abort signal and is responsible for honoring it.

340 */

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

342}

343

344/**

345 * Thinking/reasoning level for models that support it.

346 * Note: "xhigh" and "max" are only supported by selected model families. Use model

347 * thinking-level metadata from @earendil-works/pi-ai to detect support for a concrete model.

348 */

349export type ThinkingLevel = "off" | "minimal" | "low" | "medium" | "high" | "xhigh" | "max";

350

351/**

352 * Extensible interface for custom app messages.

353 * Apps can extend via declaration merging:

354 *

355 * @example

356 * ```typescript

357 * declare module "@mariozechner/agent" {

358 * interface CustomAgentMessages {

359 * artifact: ArtifactMessage;

360 * notification: NotificationMessage;

361 * }

362 * }

363 * ```

364 */

365export interface CustomAgentMessages {

366 // Empty by default - apps extend via declaration merging

367}

368

369/**

370 * AgentMessage: Union of LLM messages + custom messages.

371 * This abstraction allows apps to add custom message types while maintaining

372 * type safety and compatibility with the base LLM messages.

373 */

374export type AgentMessage = Message | CustomAgentMessages[keyof CustomAgentMessages];

375

376/**

377 * Public agent state.

378 *

379 * `tools` and `messages` use accessor properties so implementations can copy

380 * assigned arrays before storing them.

381 */

382export interface AgentState {

383 /**

384 * Current system prompt, replayed from the transcript's system messages.

385 *

386 * Read-only: to change the prompt, append a system message with `content` or `sections`.

387 * In `initialState`, this seeds the leading system message.

388 */

389 readonly systemPrompt: string;

390 /** Active model used for future turns. */

391 model: Model<any>;

392 /** Requested reasoning level for future turns. */

393 thinkingLevel: ThinkingLevel;

394 /**

395 * Executable tools. Assigning a new array copies the top-level array.

396 *

397 * Differences from the tools declared in the transcript are announced to the model

398 * with a system message before the next request.

399 */

400 set tools(tools: AgentTool<any>[]);

401 get tools(): AgentTool<any>[];

402 /**

403 * Conversation transcript. Assigning a new array copies the top-level array.

404 *

405 * System messages in the transcript carry the prompt and tool declarations.

406 */

407 set messages(messages: AgentMessage[]);

408 get messages(): AgentMessage[];

409 /**

410 * True while the agent is processing a prompt or continuation.

411 *

412 * This remains true until awaited `agent_end` listeners settle.

413 */

414 readonly isStreaming: boolean;

415 /** Partial assistant message for the current streamed response, if any. */

416 readonly streamingMessage?: AgentMessage;

417 /** Tool call ids currently executing. */

418 readonly pendingToolCalls: ReadonlySet<string>;

419 /** Error message from the most recent failed or aborted assistant turn, if any. */

420 readonly errorMessage?: string;

421}

422

423/** Final or partial result produced by a tool. */

424export interface AgentToolResult<T = JsonValue | undefined> {

425 /** Text or image content returned to the model. */

426 content: (TextContent | ImageContent)[];

427 /** Arbitrary structured details for logs or UI rendering. */

428 details: T;

429 /**

430 * Machine-readable result matching the tool's `outputSchema`, for programmatic callers. Not sent

431 * to the model; `content` remains the model-facing result.

432 */

433 structuredContent?: JsonValue;

434 /** Usage from the final tool execution itself, if available. Not used for main LLM context accounting. */

435 usage?: Usage;

436 /**

437 * Report a failure without throwing. The model sees `content` as an error result, like a thrown

438 * error, but `details` and `structuredContent` are kept for the UI and programmatic callers.

439 */

440 isError?: boolean;

441 /**

442 * Hint that the agent should stop after the current tool batch.

443 * Early termination only happens when every finalized tool result in the batch sets this to true.

444 */

445 terminate?: boolean;

446}

447

448/** Final outcome of a tool call after hooks ran. */

449export interface AgentToolCallOutcome {

450 toolCall: AgentToolCall;

451 result: AgentToolResult<any>;

452 isError: boolean;

453}

454

455/**

456 * Callback used by tools to stream partial execution updates.

457 *

458 * The callback is scoped to the current `execute()` invocation. Calls made after

459 * the tool promise settles are ignored.

460 */

461export type AgentToolUpdateCallback<T = any> = (partialResult: AgentToolResult<T>) => void;

462

463/** Tool definition used by the agent runtime. */

464export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any> extends Tool<TParameters> {

465 /** Human-readable label for UI display. */

466 label: string;

467 /**

468 * Optional compatibility shim for raw tool-call arguments before schema validation.

469 * Must return an object that matches `TParameters`.

470 */

471 prepareArguments?: (args: unknown) => Static<TParameters>;

472 /**

473 * JSON Schema of `structuredContent` in successful results. Tools that declare it should always

474 * set `structuredContent`.

475 */

476 outputSchema?: TSchema;

477 /**

478 * Execute the tool call. Throw on failure, or return a result with `isError: true`; do not only

479 * describe the failure in `content`.

480 */

481 execute: (

482 toolCallId: string,

483 params: Static<TParameters>,

484 signal?: AbortSignal,

485 onUpdate?: AgentToolUpdateCallback<TDetails>,

486 ) => Promise<AgentToolResult<TDetails>>;

487 /** Recovery policy for an effect whose durable intent exists but whose outcome is unknown. */

488 replay?: "never" | "safe";

489 /**

490 * Per-tool execution mode override.

491 * - "sequential": this tool must execute one at a time with other tool calls.

492 * - "parallel": this tool can execute concurrently with other tool calls.

493 *

494 * If omitted, the default execution mode applies.

495 */

496 executionMode?: ToolExecutionMode;

497}

498

499/** Context snapshot passed into the low-level agent loop. */

500export interface AgentContext {

501 /** Transcript visible to the model. */

502 messages: AgentMessage[];

503 /** Tools available for execution in this run. */

504 tools?: AgentTool<any>[];

505}

506

507/**

508 * Events emitted by the Agent for UI updates.

509 *

510 * `agent_end` is the last event emitted for a run, but awaited `Agent.subscribe()`

511 * listeners for that event are still part of run settlement. The agent becomes

512 * idle only after those listeners finish.

513 */

514export type AgentEvent =

515 // Agent lifecycle

516 | { type: "agent_start" }

517 | { type: "agent_end"; messages: AgentMessage[] }

518 // Turn lifecycle - a turn is one assistant response + any tool calls/results

519 | { type: "turn_start" }

520 | { type: "turn_end"; message: AgentMessage; toolResults: ToolResultMessage[] }

521 // Message lifecycle - emitted for system, user, assistant, and toolResult messages

522 | { type: "message_start"; message: AgentMessage }

523 // Only emitted for assistant messages during streaming

524 | { type: "message_update"; message: AgentMessage; assistantMessageEvent: AssistantMessageEvent }

525 | { type: "message_end"; message: AgentMessage }

526 // Tool execution lifecycle

527 | { type: "tool_execution_start"; toolCallId: string; toolName: string; args: any }

528 | { type: "tool_execution_update"; toolCallId: string; toolName: string; args: any; partialResult: any }

529 | { type: "tool_execution_end"; toolCallId: string; toolName: string; result: any; isError: boolean };