forked from badlogic/pi-mono
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathagent-session.ts
More file actions
3289 lines (2892 loc) · 106 KB
/
agent-session.ts
File metadata and controls
3289 lines (2892 loc) · 106 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
/**
* AgentSession - Core abstraction for agent lifecycle and session management.
*
* This class is shared between all run modes (interactive, print, rpc).
* It encapsulates:
* - Agent state access
* - Event subscription with automatic session persistence
* - Model and thinking level management
* - Compaction (manual and auto)
* - Bash execution
* - Session switching and branching
*
* Modes use this class and add their own I/O layer on top.
*/
import { copyFileSync, existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { basename, dirname, join, resolve } from "node:path";
import type {
Agent,
AgentEvent,
AgentMessage,
AgentState,
AgentTool,
ThinkingLevel,
} from "@mariozechner/pi-agent-core";
import type { AssistantMessage, ImageContent, Message, Model, TextContent } from "@mariozechner/pi-ai";
import { isContextOverflow, modelsAreEqual, resetApiProviders, supportsXhigh } from "@mariozechner/pi-ai";
import { getDocsPath } from "../config.js";
import { theme } from "../modes/interactive/theme/theme.js";
import { stripFrontmatter } from "../utils/frontmatter.js";
import { sleep } from "../utils/sleep.js";
import { type BashResult, executeBash as executeBashCommand, executeBashWithOperations } from "./bash-executor.js";
import {
type CompactionResult,
calculateContextTokens,
collectEntriesForBranchSummary,
compact,
estimateContextTokens,
generateBranchSummary,
prepareCompaction,
shouldCompact,
} from "./compaction/index.js";
import { DEFAULT_THINKING_LEVEL } from "./defaults.js";
import { exportSessionToHtml, type ToolHtmlRenderer } from "./export-html/index.js";
import { createToolHtmlRenderer } from "./export-html/tool-renderer.js";
import {
type ContextUsage,
type ExtensionCommandContextActions,
type ExtensionErrorListener,
ExtensionRunner,
type ExtensionUIContext,
type InputSource,
type MessageEndEvent,
type MessageStartEvent,
type MessageUpdateEvent,
type SessionBeforeCompactResult,
type SessionBeforeForkResult,
type SessionBeforeSwitchResult,
type SessionBeforeTreeResult,
type ShutdownHandler,
type ToolDefinition,
type ToolExecutionEndEvent,
type ToolExecutionStartEvent,
type ToolExecutionUpdateEvent,
type ToolInfo,
type TreePreparation,
type TurnEndEvent,
type TurnStartEvent,
wrapRegisteredTools,
} from "./extensions/index.js";
import type { BashExecutionMessage, CustomMessage } from "./messages.js";
import type { ModelRegistry } from "./model-registry.js";
import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.js";
import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.js";
import type { BranchSummaryEntry, CompactionEntry, SessionManager } from "./session-manager.js";
import { CURRENT_SESSION_VERSION, getLatestCompactionEntry, type SessionHeader } from "./session-manager.js";
import type { SettingsManager } from "./settings-manager.js";
import type { SlashCommandInfo } from "./slash-commands.js";
import { createSyntheticSourceInfo, type SourceInfo } from "./source-info.js";
import { buildSystemPrompt } from "./system-prompt.js";
import type { BashOperations } from "./tools/bash.js";
import { createAllToolDefinitions } from "./tools/index.js";
import { createToolDefinitionFromAgentTool, wrapToolDefinition } from "./tools/tool-definition-wrapper.js";
// ============================================================================
// Skill Block Parsing
// ============================================================================
/** Parsed skill block from a user message */
export interface ParsedSkillBlock {
name: string;
location: string;
content: string;
userMessage: string | undefined;
}
/**
* Parse a skill block from message text.
* Returns null if the text doesn't contain a skill block.
*/
export function parseSkillBlock(text: string): ParsedSkillBlock | null {
const match = text.match(/^<skill name="([^"]+)" location="([^"]+)">\n([\s\S]*?)\n<\/skill>(?:\n\n([\s\S]+))?$/);
if (!match) return null;
return {
name: match[1],
location: match[2],
content: match[3],
userMessage: match[4]?.trim() || undefined,
};
}
/** Session-specific events that extend the core AgentEvent */
export type AgentSessionEvent =
| AgentEvent
| {
type: "queue_update";
steering: readonly string[];
followUp: readonly string[];
}
| { type: "compaction_start"; reason: "manual" | "threshold" | "overflow" }
| {
type: "compaction_end";
reason: "manual" | "threshold" | "overflow";
result: CompactionResult | undefined;
aborted: boolean;
willRetry: boolean;
errorMessage?: string;
}
| { type: "auto_retry_start"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string }
| { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string };
/** Listener function for agent session events */
export type AgentSessionEventListener = (event: AgentSessionEvent) => void;
// ============================================================================
// Types
// ============================================================================
export interface AgentSessionConfig {
agent: Agent;
sessionManager: SessionManager;
settingsManager: SettingsManager;
cwd: string;
/** Models to cycle through with Ctrl+P (from --models flag) */
scopedModels?: Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }>;
/** Resource loader for skills, prompts, themes, context files, system prompt */
resourceLoader: ResourceLoader;
/** SDK custom tools registered outside extensions */
customTools?: ToolDefinition[];
/** Model registry for API key resolution and model discovery */
modelRegistry: ModelRegistry;
/** Initial active built-in tool names. Default: [read, bash, edit, write] */
initialActiveToolNames?: string[];
/**
* Override base tools (useful for custom runtimes).
*
* These are synthesized into minimal ToolDefinitions internally so AgentSession can keep
* a definition-first registry even when callers provide plain AgentTool instances.
*/
baseToolsOverride?: Record<string, AgentTool>;
/** Mutable ref used by Agent to access the current ExtensionRunner */
extensionRunnerRef?: { current?: ExtensionRunner };
}
export interface ExtensionBindings {
uiContext?: ExtensionUIContext;
commandContextActions?: ExtensionCommandContextActions;
shutdownHandler?: ShutdownHandler;
onError?: ExtensionErrorListener;
}
/** Options for AgentSession.prompt() */
export interface PromptOptions {
/** Whether to expand file-based prompt templates (default: true) */
expandPromptTemplates?: boolean;
/** Image attachments */
images?: ImageContent[];
/** When streaming, how to queue the message: "steer" (interrupt) or "followUp" (wait). Required if streaming. */
streamingBehavior?: "steer" | "followUp";
/** Source of input for extension input event handlers. Defaults to "interactive". */
source?: InputSource;
}
/** Result from cycleModel() */
export interface ModelCycleResult {
model: Model<any>;
thinkingLevel: ThinkingLevel;
/** Whether cycling through scoped models (--models flag) or all available */
isScoped: boolean;
}
/** Session statistics for /session command */
export interface SessionStats {
sessionFile: string | undefined;
sessionId: string;
userMessages: number;
assistantMessages: number;
toolCalls: number;
toolResults: number;
totalMessages: number;
tokens: {
input: number;
output: number;
cacheRead: number;
cacheWrite: number;
total: number;
};
cost: number;
contextUsage?: ContextUsage;
}
interface ToolDefinitionEntry {
definition: ToolDefinition;
sourceInfo: SourceInfo;
}
// ============================================================================
// Constants
// ============================================================================
/** Standard thinking levels */
const THINKING_LEVELS: ThinkingLevel[] = ["off", "minimal", "low", "medium", "high"];
/** Thinking levels including xhigh (for supported models) */
const THINKING_LEVELS_WITH_XHIGH: ThinkingLevel[] = ["off", "minimal", "low", "medium", "high", "xhigh"];
// ============================================================================
// AgentSession Class
// ============================================================================
export class AgentSession {
readonly agent: Agent;
readonly sessionManager: SessionManager;
readonly settingsManager: SettingsManager;
private _scopedModels: Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }>;
// Event subscription state
private _unsubscribeAgent?: () => void;
private _eventListeners: AgentSessionEventListener[] = [];
private _agentEventQueue: Promise<void> = Promise.resolve();
/** Tracks pending steering messages for UI display. Removed when delivered. */
private _steeringMessages: string[] = [];
/** Tracks pending follow-up messages for UI display. Removed when delivered. */
private _followUpMessages: string[] = [];
/** Messages queued to be included with the next user prompt as context ("asides"). */
private _pendingNextTurnMessages: CustomMessage[] = [];
// Compaction state
private _compactionAbortController: AbortController | undefined = undefined;
private _autoCompactionAbortController: AbortController | undefined = undefined;
private _overflowRecoveryAttempted = false;
// Branch summarization state
private _branchSummaryAbortController: AbortController | undefined = undefined;
// Retry state
private _retryAbortController: AbortController | undefined = undefined;
private _retryAttempt = 0;
private _retryPromise: Promise<void> | undefined = undefined;
private _retryResolve: (() => void) | undefined = undefined;
// Bash execution state
private _bashAbortController: AbortController | undefined = undefined;
private _pendingBashMessages: BashExecutionMessage[] = [];
// Extension system
private _extensionRunner: ExtensionRunner | undefined = undefined;
private _turnIndex = 0;
private _resourceLoader: ResourceLoader;
private _customTools: ToolDefinition[];
private _baseToolDefinitions: Map<string, ToolDefinition> = new Map();
private _cwd: string;
private _extensionRunnerRef?: { current?: ExtensionRunner };
private _initialActiveToolNames?: string[];
private _baseToolsOverride?: Record<string, AgentTool>;
private _extensionUIContext?: ExtensionUIContext;
private _extensionCommandContextActions?: ExtensionCommandContextActions;
private _extensionShutdownHandler?: ShutdownHandler;
private _extensionErrorListener?: ExtensionErrorListener;
private _extensionErrorUnsubscriber?: () => void;
// Model registry for API key resolution
private _modelRegistry: ModelRegistry;
// Tool registry for extension getTools/setTools
private _toolRegistry: Map<string, AgentTool> = new Map();
private _toolDefinitions: Map<string, ToolDefinitionEntry> = new Map();
private _toolPromptSnippets: Map<string, string> = new Map();
private _toolPromptGuidelines: Map<string, string[]> = new Map();
// Base system prompt (without extension appends) - used to apply fresh appends each turn
private _baseSystemPrompt = "";
constructor(config: AgentSessionConfig) {
this.agent = config.agent;
this.sessionManager = config.sessionManager;
this.settingsManager = config.settingsManager;
this._scopedModels = config.scopedModels ?? [];
this._resourceLoader = config.resourceLoader;
this._customTools = config.customTools ?? [];
this._cwd = config.cwd;
this._modelRegistry = config.modelRegistry;
this._extensionRunnerRef = config.extensionRunnerRef;
this._initialActiveToolNames = config.initialActiveToolNames;
this._baseToolsOverride = config.baseToolsOverride;
// Always subscribe to agent events for internal handling
// (session persistence, extensions, auto-compaction, retry logic)
this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
this._installAgentToolHooks();
this._buildRuntime({
activeToolNames: this._initialActiveToolNames,
includeAllExtensionTools: true,
});
}
/** Model registry for API key resolution and model discovery */
get modelRegistry(): ModelRegistry {
return this._modelRegistry;
}
private async _getRequiredRequestAuth(model: Model<any>): Promise<{
apiKey: string;
headers?: Record<string, string>;
}> {
const result = await this._modelRegistry.getApiKeyAndHeaders(model);
if (!result.ok) {
throw new Error(result.error);
}
if (result.apiKey) {
return { apiKey: result.apiKey, headers: result.headers };
}
const isOAuth = this._modelRegistry.isUsingOAuth(model);
if (isOAuth) {
throw new Error(
`Authentication failed for "${model.provider}". ` +
`Credentials may have expired or network is unavailable. ` +
`Run '/login ${model.provider}' to re-authenticate.`,
);
}
throw new Error(
`No API key found for ${model.provider}.\n\n` +
`Use /login or set an API key environment variable. See ${join(getDocsPath(), "providers.md")}`,
);
}
/**
* Install tool hooks once on the Agent instance.
*
* The callbacks read `this._extensionRunner` at execution time, so extension reload swaps in the
* new runner without reinstalling hooks. Extension-specific tool wrappers are still used to adapt
* registered tool execution to the extension context. Tool call and tool result interception now
* happens here instead of in wrappers.
*/
private _installAgentToolHooks(): void {
this.agent.beforeToolCall = async ({ toolCall, args }) => {
const runner = this._extensionRunner;
if (!runner?.hasHandlers("tool_call")) {
return undefined;
}
await this._agentEventQueue;
try {
return await runner.emitToolCall({
type: "tool_call",
toolName: toolCall.name,
toolCallId: toolCall.id,
input: args as Record<string, unknown>,
});
} catch (err) {
if (err instanceof Error) {
throw err;
}
throw new Error(`Extension failed, blocking execution: ${String(err)}`);
}
};
this.agent.afterToolCall = async ({ toolCall, args, result, isError }) => {
const runner = this._extensionRunner;
if (!runner?.hasHandlers("tool_result")) {
return undefined;
}
const hookResult = await runner.emitToolResult({
type: "tool_result",
toolName: toolCall.name,
toolCallId: toolCall.id,
input: args as Record<string, unknown>,
content: result.content,
details: isError ? undefined : result.details,
isError,
});
if (!hookResult || isError) {
return undefined;
}
return {
content: hookResult.content,
details: hookResult.details,
};
};
}
// =========================================================================
// Event Subscription
// =========================================================================
/** Emit an event to all listeners */
private _emit(event: AgentSessionEvent): void {
for (const l of this._eventListeners) {
l(event);
}
}
private _emitQueueUpdate(): void {
this._emit({
type: "queue_update",
steering: [...this._steeringMessages],
followUp: [...this._followUpMessages],
});
}
// Track last assistant message for auto-compaction check
private _lastAssistantMessage: AssistantMessage | undefined = undefined;
/** Internal handler for agent events - shared by subscribe and reconnect */
private _handleAgentEvent = (event: AgentEvent): void => {
// Create retry promise synchronously before queueing async processing.
// Agent.emit() calls this handler synchronously, and prompt() calls waitForRetry()
// as soon as agent.prompt() resolves. If _retryPromise is created only inside
// _processAgentEvent, slow earlier queued events can delay agent_end processing
// and waitForRetry() can miss the in-flight retry.
this._createRetryPromiseForAgentEnd(event);
this._agentEventQueue = this._agentEventQueue.then(
() => this._processAgentEvent(event),
() => this._processAgentEvent(event),
);
// Keep queue alive if an event handler fails
this._agentEventQueue.catch(() => {});
};
private _createRetryPromiseForAgentEnd(event: AgentEvent): void {
if (event.type !== "agent_end" || this._retryPromise) {
return;
}
const settings = this.settingsManager.getRetrySettings();
if (!settings.enabled) {
return;
}
const lastAssistant = this._findLastAssistantInMessages(event.messages);
if (!lastAssistant || !this._isRetryableError(lastAssistant)) {
return;
}
this._retryPromise = new Promise((resolve) => {
this._retryResolve = resolve;
});
}
private _findLastAssistantInMessages(messages: AgentMessage[]): AssistantMessage | undefined {
for (let i = messages.length - 1; i >= 0; i--) {
const message = messages[i];
if (message.role === "assistant") {
return message as AssistantMessage;
}
}
return undefined;
}
private async _processAgentEvent(event: AgentEvent): Promise<void> {
// When a user message starts, check if it's from either queue and remove it BEFORE emitting
// This ensures the UI sees the updated queue state
if (event.type === "message_start" && event.message.role === "user") {
this._overflowRecoveryAttempted = false;
const messageText = this._getUserMessageText(event.message);
if (messageText) {
// Check steering queue first
const steeringIndex = this._steeringMessages.indexOf(messageText);
if (steeringIndex !== -1) {
this._steeringMessages.splice(steeringIndex, 1);
this._emitQueueUpdate();
} else {
// Check follow-up queue
const followUpIndex = this._followUpMessages.indexOf(messageText);
if (followUpIndex !== -1) {
this._followUpMessages.splice(followUpIndex, 1);
this._emitQueueUpdate();
}
}
}
}
// Emit to extensions first
await this._emitExtensionEvent(event);
// Notify all listeners
this._emit(event);
// Handle session persistence
if (event.type === "message_end") {
// Check if this is a custom message from extensions
if (event.message.role === "custom") {
// Persist as CustomMessageEntry
this.sessionManager.appendCustomMessageEntry(
event.message.customType,
event.message.content,
event.message.display,
event.message.details,
);
} else if (
event.message.role === "user" ||
event.message.role === "assistant" ||
event.message.role === "toolResult"
) {
// Regular LLM message - persist as SessionMessageEntry
this.sessionManager.appendMessage(event.message);
}
// Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere
// Track assistant message for auto-compaction (checked on agent_end)
if (event.message.role === "assistant") {
this._lastAssistantMessage = event.message;
const assistantMsg = event.message as AssistantMessage;
if (assistantMsg.stopReason !== "error") {
this._overflowRecoveryAttempted = false;
}
// Reset retry counter immediately on successful assistant response
// This prevents accumulation across multiple LLM calls within a turn
if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) {
this._emit({
type: "auto_retry_end",
success: true,
attempt: this._retryAttempt,
});
this._retryAttempt = 0;
}
}
}
// Check auto-retry and auto-compaction after agent completes
if (event.type === "agent_end" && this._lastAssistantMessage) {
const msg = this._lastAssistantMessage;
this._lastAssistantMessage = undefined;
// Check for retryable errors first (overloaded, rate limit, server errors)
if (this._isRetryableError(msg)) {
const didRetry = await this._handleRetryableError(msg);
if (didRetry) return; // Retry was initiated, don't proceed to compaction
}
this._resolveRetry();
await this._checkCompaction(msg);
}
}
/** Resolve the pending retry promise */
private _resolveRetry(): void {
if (this._retryResolve) {
this._retryResolve();
this._retryResolve = undefined;
this._retryPromise = undefined;
}
}
/** Extract text content from a message */
private _getUserMessageText(message: Message): string {
if (message.role !== "user") return "";
const content = message.content;
if (typeof content === "string") return content;
const textBlocks = content.filter((c) => c.type === "text");
return textBlocks.map((c) => (c as TextContent).text).join("");
}
/** Find the last assistant message in agent state (including aborted ones) */
private _findLastAssistantMessage(): AssistantMessage | undefined {
const messages = this.agent.state.messages;
for (let i = messages.length - 1; i >= 0; i--) {
const msg = messages[i];
if (msg.role === "assistant") {
return msg as AssistantMessage;
}
}
return undefined;
}
/** Emit extension events based on agent events */
private async _emitExtensionEvent(event: AgentEvent): Promise<void> {
if (!this._extensionRunner) return;
if (event.type === "agent_start") {
this._turnIndex = 0;
await this._extensionRunner.emit({ type: "agent_start" });
} else if (event.type === "agent_end") {
await this._extensionRunner.emit({ type: "agent_end", messages: event.messages });
} else if (event.type === "turn_start") {
const extensionEvent: TurnStartEvent = {
type: "turn_start",
turnIndex: this._turnIndex,
timestamp: Date.now(),
};
await this._extensionRunner.emit(extensionEvent);
} else if (event.type === "turn_end") {
const extensionEvent: TurnEndEvent = {
type: "turn_end",
turnIndex: this._turnIndex,
message: event.message,
toolResults: event.toolResults,
};
await this._extensionRunner.emit(extensionEvent);
this._turnIndex++;
} else if (event.type === "message_start") {
const extensionEvent: MessageStartEvent = {
type: "message_start",
message: event.message,
};
await this._extensionRunner.emit(extensionEvent);
} else if (event.type === "message_update") {
const extensionEvent: MessageUpdateEvent = {
type: "message_update",
message: event.message,
assistantMessageEvent: event.assistantMessageEvent,
};
await this._extensionRunner.emit(extensionEvent);
} else if (event.type === "message_end") {
const extensionEvent: MessageEndEvent = {
type: "message_end",
message: event.message,
};
await this._extensionRunner.emit(extensionEvent);
} else if (event.type === "tool_execution_start") {
const extensionEvent: ToolExecutionStartEvent = {
type: "tool_execution_start",
toolCallId: event.toolCallId,
toolName: event.toolName,
args: event.args,
};
await this._extensionRunner.emit(extensionEvent);
} else if (event.type === "tool_execution_update") {
const extensionEvent: ToolExecutionUpdateEvent = {
type: "tool_execution_update",
toolCallId: event.toolCallId,
toolName: event.toolName,
args: event.args,
partialResult: event.partialResult,
};
await this._extensionRunner.emit(extensionEvent);
} else if (event.type === "tool_execution_end") {
const extensionEvent: ToolExecutionEndEvent = {
type: "tool_execution_end",
toolCallId: event.toolCallId,
toolName: event.toolName,
result: event.result,
isError: event.isError,
};
await this._extensionRunner.emit(extensionEvent);
}
}
/**
* Subscribe to agent events.
* Session persistence is handled internally (saves messages on message_end).
* Multiple listeners can be added. Returns unsubscribe function for this listener.
*/
subscribe(listener: AgentSessionEventListener): () => void {
this._eventListeners.push(listener);
// Return unsubscribe function for this specific listener
return () => {
const index = this._eventListeners.indexOf(listener);
if (index !== -1) {
this._eventListeners.splice(index, 1);
}
};
}
/**
* Temporarily disconnect from agent events.
* User listeners are preserved and will receive events again after resubscribe().
* Used internally during operations that need to pause event processing.
*/
private _disconnectFromAgent(): void {
if (this._unsubscribeAgent) {
this._unsubscribeAgent();
this._unsubscribeAgent = undefined;
}
}
/**
* Reconnect to agent events after _disconnectFromAgent().
* Preserves all existing listeners.
*/
private _reconnectToAgent(): void {
if (this._unsubscribeAgent) return; // Already connected
this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
}
/**
* Remove all listeners and disconnect from agent.
* Call this when completely done with the session.
*/
dispose(): void {
this._disconnectFromAgent();
this._eventListeners = [];
}
// =========================================================================
// Read-only State Access
// =========================================================================
/** Full agent state */
get state(): AgentState {
return this.agent.state;
}
/** Current model (may be undefined if not yet selected) */
get model(): Model<any> | undefined {
return this.agent.state.model;
}
/** Current thinking level */
get thinkingLevel(): ThinkingLevel {
return this.agent.state.thinkingLevel;
}
/** Whether agent is currently streaming a response */
get isStreaming(): boolean {
return this.agent.state.isStreaming;
}
/** Current effective system prompt (includes any per-turn extension modifications) */
get systemPrompt(): string {
return this.agent.state.systemPrompt;
}
/** Current retry attempt (0 if not retrying) */
get retryAttempt(): number {
return this._retryAttempt;
}
/**
* Get the names of currently active tools.
* Returns the names of tools currently set on the agent.
*/
getActiveToolNames(): string[] {
return this.agent.state.tools.map((t) => t.name);
}
/**
* Get all configured tools with name, description, parameter schema, and source metadata.
*/
getAllTools(): ToolInfo[] {
return Array.from(this._toolDefinitions.values()).map(({ definition, sourceInfo }) => ({
name: definition.name,
description: definition.description,
parameters: definition.parameters,
sourceInfo,
}));
}
getToolDefinition(name: string): ToolDefinition | undefined {
return this._toolDefinitions.get(name)?.definition;
}
/**
* Set active tools by name.
* Only tools in the registry can be enabled. Unknown tool names are ignored.
* Also rebuilds the system prompt to reflect the new tool set.
* Changes take effect on the next agent turn.
*/
setActiveToolsByName(toolNames: string[]): void {
const tools: AgentTool[] = [];
const validToolNames: string[] = [];
for (const name of toolNames) {
const tool = this._toolRegistry.get(name);
if (tool) {
tools.push(tool);
validToolNames.push(name);
}
}
this.agent.state.tools = tools;
// Rebuild base system prompt with new tool set
this._baseSystemPrompt = this._rebuildSystemPrompt(validToolNames);
this.agent.state.systemPrompt = this._baseSystemPrompt;
}
/** Whether compaction or branch summarization is currently running */
get isCompacting(): boolean {
return (
this._autoCompactionAbortController !== undefined ||
this._compactionAbortController !== undefined ||
this._branchSummaryAbortController !== undefined
);
}
/** All messages including custom types like BashExecutionMessage */
get messages(): AgentMessage[] {
return this.agent.state.messages;
}
/** Current steering mode */
get steeringMode(): "all" | "one-at-a-time" {
return this.agent.steeringMode;
}
/** Current follow-up mode */
get followUpMode(): "all" | "one-at-a-time" {
return this.agent.followUpMode;
}
/** Current session file path, or undefined if sessions are disabled */
get sessionFile(): string | undefined {
return this.sessionManager.getSessionFile();
}
/** Current session ID */
get sessionId(): string {
return this.sessionManager.getSessionId();
}
/** Current session display name, if set */
get sessionName(): string | undefined {
return this.sessionManager.getSessionName();
}
/** Scoped models for cycling (from --models flag) */
get scopedModels(): ReadonlyArray<{ model: Model<any>; thinkingLevel?: ThinkingLevel }> {
return this._scopedModels;
}
/** Update scoped models for cycling */
setScopedModels(scopedModels: Array<{ model: Model<any>; thinkingLevel?: ThinkingLevel }>): void {
this._scopedModels = scopedModels;
}
/** File-based prompt templates */
get promptTemplates(): ReadonlyArray<PromptTemplate> {
return this._resourceLoader.getPrompts().prompts;
}
private _normalizePromptSnippet(text: string | undefined): string | undefined {
if (!text) return undefined;
const oneLine = text
.replace(/[\r\n]+/g, " ")
.replace(/\s+/g, " ")
.trim();
return oneLine.length > 0 ? oneLine : undefined;
}
private _normalizePromptGuidelines(guidelines: string[] | undefined): string[] {
if (!guidelines || guidelines.length === 0) {
return [];
}
const unique = new Set<string>();
for (const guideline of guidelines) {
const normalized = guideline.trim();
if (normalized.length > 0) {
unique.add(normalized);
}
}
return Array.from(unique);
}
private _rebuildSystemPrompt(toolNames: string[]): string {
const validToolNames = toolNames.filter((name) => this._toolRegistry.has(name));
const toolSnippets: Record<string, string> = {};
const promptGuidelines: string[] = [];
for (const name of validToolNames) {
const snippet = this._toolPromptSnippets.get(name);
if (snippet) {
toolSnippets[name] = snippet;
}
const toolGuidelines = this._toolPromptGuidelines.get(name);
if (toolGuidelines) {
promptGuidelines.push(...toolGuidelines);
}
}
const loaderSystemPrompt = this._resourceLoader.getSystemPrompt();
const loaderAppendSystemPrompt = this._resourceLoader.getAppendSystemPrompt();
const appendSystemPrompt =
loaderAppendSystemPrompt.length > 0 ? loaderAppendSystemPrompt.join("\n\n") : undefined;
const loadedSkills = this._resourceLoader.getSkills().skills;
const loadedContextFiles = this._resourceLoader.getAgentsFiles().agentsFiles;
return buildSystemPrompt({
cwd: this._cwd,
skills: loadedSkills,
contextFiles: loadedContextFiles,
customPrompt: loaderSystemPrompt,
appendSystemPrompt,
selectedTools: validToolNames,
toolSnippets,
promptGuidelines,
});
}
// =========================================================================
// Prompting
// =========================================================================
/**
* Send a prompt to the agent.
* - Handles extension commands (registered via pi.registerCommand) immediately, even during streaming
* - Expands file-based prompt templates by default
* - During streaming, queues via steer() or followUp() based on streamingBehavior option
* - Validates model and API key before sending (when not streaming)
* @throws Error if streaming and no streamingBehavior specified
* @throws Error if no model selected or no API key available (when not streaming)
*/
async prompt(text: string, options?: PromptOptions): Promise<void> {
const expandPromptTemplates = options?.expandPromptTemplates ?? true;
// Handle extension commands first (execute immediately, even during streaming)
// Extension commands manage their own LLM interaction via pi.sendMessage()
if (expandPromptTemplates && text.startsWith("/")) {
const handled = await this._tryExecuteExtensionCommand(text);
if (handled) {
// Extension command executed, no prompt to send
return;
}
}
// Emit input event for extension interception (before skill/template expansion)
let currentText = text;
let currentImages = options?.images;
if (this._extensionRunner?.hasHandlers("input")) {
const inputResult = await this._extensionRunner.emitInput(
currentText,
currentImages,
options?.source ?? "interactive",
);
if (inputResult.action === "handled") {
return;
}
if (inputResult.action === "transform") {
currentText = inputResult.text;
currentImages = inputResult.images ?? currentImages;
}
}
// Expand skill commands (/skill:name args) and prompt templates (/template args)
let expandedText = currentText;
if (expandPromptTemplates) {
expandedText = this._expandSkillCommand(expandedText);
expandedText = expandPromptTemplate(expandedText, [...this.promptTemplates]);
}
// If streaming, queue via steer() or followUp() based on option
if (this.isStreaming) {
if (!options?.streamingBehavior) {
throw new Error(
"Agent is already processing. Specify streamingBehavior ('steer' or 'followUp') to queue the message.",
);
}
if (options.streamingBehavior === "followUp") {
await this._queueFollowUp(expandedText, currentImages);
} else {
await this._queueSteer(expandedText, currentImages);
}
return;
}
// Flush any pending bash messages before the new prompt
this._flushPendingBashMessages();
// Validate model
if (!this.model) {
throw new Error(
"No model selected.\n\n" +
`Use /login or set an API key environment variable. See ${join(getDocsPath(), "providers.md")}\n\n` +
"Then use /model to select a model.",
);
}
if (!this._modelRegistry.hasConfiguredAuth(this.model)) {
const isOAuth = this._modelRegistry.isUsingOAuth(this.model);
if (isOAuth) {
throw new Error(
`Authentication failed for "${this.model.provider}". ` +
`Credentials may have expired or network is unavailable. ` +
`Run '/login ${this.model.provider}' to re-authenticate.`,
);
}
throw new Error(