mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-16 17:01:02 +00:00
Improve compaction hooks: add signal, no timeout, SessionManager cleanup, docs
This commit is contained in:
Mario Zechner
parent
a2664ba38a
commit
705ba5d4f2
19 changed files with 1236 additions and 207 deletions
|
|
@ -10,6 +10,11 @@
|
|||
- `resolveApiKey`: Function to resolve API keys for any model (checks settings, OAuth, env vars)
|
||||
- Removed `apiKey` string in favor of `resolveApiKey` for more flexibility
|
||||
|
||||
- **SessionManager API cleanup**:
|
||||
- Renamed `loadSessionFromEntries()` to `buildSessionContext()` (builds LLM context from entries, handling compaction)
|
||||
- Renamed `loadEntries()` to `getEntries()` (returns defensive copy of all session entries)
|
||||
- Added `buildSessionContext()` method to SessionManager
|
||||
|
||||
## [0.27.5] - 2025-12-24
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -151,7 +151,7 @@ pi.on("session", async (event, ctx) => {
|
|||
// event.entries: SessionEntry[] - all session entries
|
||||
// event.sessionFile: string | null - current session file (null with --no-session)
|
||||
// event.previousSessionFile: string | null - previous session file
|
||||
// event.reason: "start" | "before_switch" | "switch" | "before_clear" | "clear" |
|
||||
// event.reason: "start" | "before_switch" | "switch" | "before_clear" | "clear" |
|
||||
// "before_branch" | "branch" | "before_compact" | "compact" | "shutdown"
|
||||
// event.targetTurnIndex: number - only for "before_branch" and "branch"
|
||||
|
||||
|
|
@ -195,24 +195,26 @@ Legend:
|
|||
```
|
||||
Session entries (before compaction):
|
||||
|
||||
index: 0 1 2 3 4 5 6 7 8 9 10 11
|
||||
┌─────┬─────┬─────┬─────┬──────┬─────┬─────┬─────┬──────┬──────┬─────┬──────┐
|
||||
│ hdr │ cmp │ usr │ ass │ tool │ ass │ usr │ ass │ tool │ tool │ ass │ tool │
|
||||
└─────┴─────┴─────┴─────┴──────┴─────┴─────┴─────┴──────┴──────┴─────┴──────┘
|
||||
↑ └────────┬────────┘ └────────────┬────────────┘
|
||||
previousSummary messagesToSummarize messagesToKeep
|
||||
↑
|
||||
cutPoint.firstKeptEntryIndex = 6
|
||||
index: 0 1 2 3 4 5 6 7 8 9 10
|
||||
┌─────┬─────┬─────┬─────┬──────┬─────┬─────┬──────┬──────┬─────┬──────┐
|
||||
│ hdr │ cmp │ usr │ ass │ tool │ usr │ ass │ tool │ tool │ ass │ tool │
|
||||
└─────┴─────┴─────┴─────┴──────┴─────┴─────┴──────┴──────┴─────┴──────┘
|
||||
↑ └───────┬───────┘ └────────────┬────────────┘
|
||||
previousSummary messagesToSummarize messagesToKeep
|
||||
↑
|
||||
cutPoint.firstKeptEntryIndex = 5
|
||||
|
||||
After compaction (new entry appended):
|
||||
|
||||
index: 0 1 2 3 4 5 6 7 8 9 10 11 12
|
||||
┌─────┬─────┬─────┬─────┬──────┬─────┬─────┬─────┬──────┬──────┬─────┬──────┬─────┐
|
||||
│ hdr │ cmp │ usr │ ass │ tool │ ass │ usr │ ass │ tool │ tool │ ass │ tool │ cmp │
|
||||
└─────┴─────┴─────┴─────┴──────┴─────┴─────┴─────┴──────┴──────┴─────┴──────┴─────┘
|
||||
└────────┬────────┘ └────────────┬────────────┘ ↑
|
||||
ignored loaded firstKeptEntryIndex = 6
|
||||
on reload on reload (stored in new cmp)
|
||||
index: 0 1 2 3 4 5 6 7 8 9 10 11
|
||||
┌─────┬─────┬─────┬─────┬──────┬─────┬─────┬──────┬──────┬─────┬──────┬─────┐
|
||||
│ hdr │ cmp │ usr │ ass │ tool │ usr │ ass │ tool │ tool │ ass │ tool │ cmp │
|
||||
└─────┴─────┴─────┴─────┴──────┴─────┴─────┴──────┴──────┴─────┴──────┴─────┘
|
||||
└──────────┬───────────┘ └────────────────────────┬─────────────────┘
|
||||
not sent to LLM sent to LLM
|
||||
↑
|
||||
firstKeptEntryIndex = 5
|
||||
(stored in new cmp)
|
||||
```
|
||||
|
||||
The session file is append-only. When loading, the session loader finds the latest compaction entry, uses its summary, then loads messages starting from `firstKeptEntryIndex`. The cut point is always a user, assistant, or bashExecution message (never a tool result, which must stay with its tool call).
|
||||
|
|
@ -220,11 +222,13 @@ The session file is append-only. When loading, the session loader finds the late
|
|||
```
|
||||
What gets sent to the LLM as context:
|
||||
|
||||
┌──────────────────────────────────────────────────────────────────────────────────────┐
|
||||
│ [system] [summary] [usr idx 6] [ass idx 7] [tool idx 8] [tool idx 9] [ass idx 10] [tool idx 11] │
|
||||
└──────────────────────────────────────────────────────────────────────────────────────┘
|
||||
↑ └──────────────────────┬──────────────────────┘
|
||||
from new cmp's summary messages from firstKeptEntryIndex onwards
|
||||
5 6 7 8 9 10
|
||||
┌────────┬─────────┬─────┬─────┬──────┬──────┬─────┬──────┐
|
||||
│ system │ summary │ usr │ ass │ tool │ tool │ ass │ tool │
|
||||
└────────┴─────────┴─────┴─────┴──────┴──────┴─────┴──────┘
|
||||
↑ └─────────────────┬────────────────┘
|
||||
from new cmp's messages from
|
||||
summary firstKeptEntryIndex onwards
|
||||
```
|
||||
|
||||
**Split turns:** When a single turn is too large, the cut point may land mid-turn at an assistant message. In this case `cutPoint.isSplitTurn = true`:
|
||||
|
|
@ -265,37 +269,11 @@ See [src/core/compaction.ts](../src/core/compaction.ts) for the full implementat
|
|||
| `model` | Model to use for summarization. |
|
||||
| `resolveApiKey` | Function to resolve API key for any model: `await resolveApiKey(model)` |
|
||||
| `customInstructions` | Optional focus for summary (from `/compact <instructions>`). |
|
||||
| `signal` | AbortSignal for cancellation. Pass to LLM calls and check periodically. |
|
||||
|
||||
**Common patterns:**
|
||||
Custom compaction hooks should honor the abort signal by passing it to `complete()` calls. This allows users to cancel compaction (e.g., via Ctrl+C during `/compact`).
|
||||
|
||||
```typescript
|
||||
// Get all messages since last compaction
|
||||
const allMessages = [...event.messagesToSummarize, ...event.messagesToKeep];
|
||||
|
||||
// Default behavior: summarize old, keep recent
|
||||
return {
|
||||
compactionEntry: {
|
||||
type: "compaction",
|
||||
timestamp: new Date().toISOString(),
|
||||
summary: "Your custom summary...",
|
||||
firstKeptEntryIndex: event.cutPoint.firstKeptEntryIndex, // keep recent turns
|
||||
tokensBefore: event.tokensBefore,
|
||||
}
|
||||
};
|
||||
|
||||
// Full compaction: summarize everything, keep nothing
|
||||
return {
|
||||
compactionEntry: {
|
||||
type: "compaction",
|
||||
timestamp: new Date().toISOString(),
|
||||
summary: "Complete summary of all work...",
|
||||
firstKeptEntryIndex: event.entries.length, // discard all messages
|
||||
tokensBefore: event.tokensBefore,
|
||||
}
|
||||
};
|
||||
```
|
||||
|
||||
See [examples/hooks/full-compaction.ts](../examples/hooks/full-compaction.ts) for a complete example.
|
||||
See [examples/hooks/custom-compaction.ts](../examples/hooks/custom-compaction.ts) for a complete example.
|
||||
|
||||
**After compaction (`compact` event):**
|
||||
- `event.compactionEntry`: The saved compaction entry
|
||||
|
|
@ -367,7 +345,7 @@ pi.on("tool_result", async (event, ctx) => {
|
|||
// event.content: (TextContent | ImageContent)[]
|
||||
// event.details: tool-specific (see below)
|
||||
// event.isError: boolean
|
||||
|
||||
|
||||
// Return modified content/details, or undefined to keep original
|
||||
return { content: [...], details: {...} };
|
||||
});
|
||||
|
|
@ -420,7 +398,7 @@ Common fields in details:
|
|||
Custom tools use `CustomToolResultEvent` with `details: unknown`. Create your own type guard to get full type safety:
|
||||
|
||||
```typescript
|
||||
import {
|
||||
import {
|
||||
isBashToolResult,
|
||||
type CustomToolResultEvent,
|
||||
type HookAPI,
|
||||
|
|
@ -432,8 +410,8 @@ interface MyCustomToolDetails {
|
|||
}
|
||||
|
||||
// Type guard that narrows both toolName and details
|
||||
function isMyCustomToolResult(e: ToolResultEvent): e is CustomToolResultEvent & {
|
||||
toolName: "my-custom-tool";
|
||||
function isMyCustomToolResult(e: ToolResultEvent): e is CustomToolResultEvent & {
|
||||
toolName: "my-custom-tool";
|
||||
details: MyCustomToolDetails;
|
||||
} {
|
||||
return e.toolName === "my-custom-tool";
|
||||
|
|
@ -574,10 +552,10 @@ import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";
|
|||
export default function (pi: HookAPI) {
|
||||
pi.on("session", async (event, ctx) => {
|
||||
if (event.reason !== "start") return;
|
||||
|
||||
|
||||
// Watch a trigger file
|
||||
const triggerFile = "/tmp/agent-trigger.txt";
|
||||
|
||||
|
||||
fs.watch(triggerFile, () => {
|
||||
try {
|
||||
const content = fs.readFileSync(triggerFile, "utf-8").trim();
|
||||
|
|
@ -589,7 +567,7 @@ export default function (pi: HookAPI) {
|
|||
// File might not exist yet
|
||||
}
|
||||
});
|
||||
|
||||
|
||||
ctx.ui.notify("Watching /tmp/agent-trigger.txt", "info");
|
||||
});
|
||||
}
|
||||
|
|
@ -606,7 +584,7 @@ import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";
|
|||
export default function (pi: HookAPI) {
|
||||
pi.on("session", async (event, ctx) => {
|
||||
if (event.reason !== "start") return;
|
||||
|
||||
|
||||
const server = http.createServer((req, res) => {
|
||||
let body = "";
|
||||
req.on("data", chunk => body += chunk);
|
||||
|
|
@ -616,7 +594,7 @@ export default function (pi: HookAPI) {
|
|||
res.end("OK");
|
||||
});
|
||||
});
|
||||
|
||||
|
||||
server.listen(3333, () => {
|
||||
ctx.ui.notify("Webhook listening on http://localhost:3333", "info");
|
||||
});
|
||||
|
|
@ -842,7 +820,7 @@ The `HookRunner` class manages hook execution:
|
|||
```typescript
|
||||
class HookRunner {
|
||||
constructor(hooks: LoadedHook[], cwd: string, timeout?: number)
|
||||
|
||||
|
||||
setUIContext(ctx: HookUIContext, hasUI: boolean): void
|
||||
setSessionFile(path: string | null): void
|
||||
onError(listener): () => void
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ Prevents session changes when there are uncommitted git changes. Blocks clear/sw
|
|||
### auto-commit-on-exit.ts
|
||||
Automatically commits changes when the agent exits (shutdown event). Uses the last assistant message to generate a commit message.
|
||||
|
||||
### full-compaction.ts
|
||||
### custom-compaction.ts
|
||||
Custom context compaction that summarizes the entire conversation instead of keeping recent turns. Uses the `before_compact` hook event to intercept compaction and generate a comprehensive summary using `complete()` from the AI package. Useful when you want maximum context window space at the cost of losing exact conversation history.
|
||||
|
||||
## Usage
|
||||
|
|
|
|||
|
|
@ -1,41 +1,54 @@
|
|||
/**
|
||||
* Full Context Compaction Hook
|
||||
* Custom Compaction Hook
|
||||
*
|
||||
* Replaces the default compaction behavior with a full summary of the entire context.
|
||||
* Instead of keeping the last 20k tokens of conversation turns, this hook:
|
||||
* 1. Summarizes ALL messages (both messagesToSummarize and messagesToKeep)
|
||||
* 1. Summarizes ALL messages (both messagesToSummarize and messagesToKeep and previousSummary)
|
||||
* 2. Discards all old turns completely, keeping only the summary
|
||||
*
|
||||
* This is useful when you want maximum context window space for new work
|
||||
* at the cost of losing exact conversation history.
|
||||
* This example also demonstrates using a different model (Gemini Flash) for summarization,
|
||||
* which can be cheaper/faster than the main conversation model.
|
||||
*
|
||||
* Usage:
|
||||
* pi --hook examples/hooks/full-compaction.ts
|
||||
* pi --hook examples/hooks/custom-compaction.ts
|
||||
*/
|
||||
|
||||
import { complete } from "@mariozechner/pi-ai";
|
||||
import { messageTransformer } from "@mariozechner/pi-coding-agent";
|
||||
import { findModel, messageTransformer } from "@mariozechner/pi-coding-agent";
|
||||
import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";
|
||||
|
||||
export default function (pi: HookAPI) {
|
||||
pi.on("session", async (event, ctx) => {
|
||||
if (event.reason !== "before_compact") return;
|
||||
|
||||
const { messagesToSummarize, messagesToKeep, previousSummary, tokensBefore, model, resolveApiKey, entries } = event;
|
||||
ctx.ui.notify("Custom compaction hook triggered", "info");
|
||||
|
||||
// Combine all messages for full summary
|
||||
const allMessages = [...messagesToSummarize, ...messagesToKeep];
|
||||
const { messagesToSummarize, messagesToKeep, previousSummary, tokensBefore, resolveApiKey, entries, signal } = event;
|
||||
|
||||
ctx.ui.notify(`Full compaction: summarizing ${allMessages.length} messages (${tokensBefore.toLocaleString()} tokens)...`, "info");
|
||||
// Use Gemini Flash for summarization (cheaper/faster than most conversation models)
|
||||
// findModel searches both built-in models and custom models from models.json
|
||||
const { model, error } = findModel("google", "gemini-2.5-flash");
|
||||
if (error || !model) {
|
||||
ctx.ui.notify(`Could not find Gemini Flash model: ${error}, using default compaction`, "warning");
|
||||
return;
|
||||
}
|
||||
|
||||
// Resolve API key for the model
|
||||
// Resolve API key for the summarization model
|
||||
const apiKey = await resolveApiKey(model);
|
||||
if (!apiKey) {
|
||||
ctx.ui.notify(`No API key for ${model.provider}, using default compaction`, "warning");
|
||||
return;
|
||||
}
|
||||
|
||||
// Transform app messages to LLM-compatible format
|
||||
// Combine all messages for full summary
|
||||
const allMessages = [...messagesToSummarize, ...messagesToKeep];
|
||||
|
||||
ctx.ui.notify(
|
||||
`Custom compaction: summarizing ${allMessages.length} messages (${tokensBefore.toLocaleString()} tokens) with ${model.id}...`,
|
||||
"info",
|
||||
);
|
||||
|
||||
// Transform app messages to pi-ai package format
|
||||
const transformedMessages = messageTransformer(allMessages);
|
||||
|
||||
// Include previous summary context if available
|
||||
|
|
@ -68,8 +81,8 @@ Format the summary as structured markdown with clear sections.`,
|
|||
];
|
||||
|
||||
try {
|
||||
// Use the same model with resolved API key
|
||||
const response = await complete(model, { messages: summaryMessages }, { apiKey, maxTokens: 8192 });
|
||||
// Pass signal to honor abort requests (e.g., user cancels compaction)
|
||||
const response = await complete(model, { messages: summaryMessages }, { apiKey, maxTokens: 8192, signal });
|
||||
|
||||
const summary = response.content
|
||||
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
||||
|
|
@ -77,8 +90,8 @@ Format the summary as structured markdown with clear sections.`,
|
|||
.join("\n");
|
||||
|
||||
if (!summary.trim()) {
|
||||
ctx.ui.notify("Compaction summary was empty, using default compaction", "warning");
|
||||
return; // Fall back to default compaction
|
||||
if (!signal.aborted) ctx.ui.notify("Compaction summary was empty, using default compaction", "warning");
|
||||
return;
|
||||
}
|
||||
|
||||
// Return a compaction entry that discards ALL messages
|
||||
|
|
@ -24,7 +24,7 @@ import { exportSessionToHtml } from "./export-html.js";
|
|||
import type { HookRunner, SessionEventResult, TurnEndEvent, TurnStartEvent } from "./hooks/index.js";
|
||||
import type { BashExecutionMessage } from "./messages.js";
|
||||
import { getApiKeyForModel, getAvailableModels } from "./model-config.js";
|
||||
import { type CompactionEntry, loadSessionFromEntries, type SessionManager } from "./session-manager.js";
|
||||
import type { CompactionEntry, SessionManager } from "./session-manager.js";
|
||||
import type { SettingsManager, SkillsSettings } from "./settings-manager.js";
|
||||
import { expandSlashCommand, type FileSlashCommand } from "./slash-commands.js";
|
||||
|
||||
|
|
@ -510,7 +510,7 @@ export class AgentSession {
|
|||
*/
|
||||
async reset(): Promise<boolean> {
|
||||
const previousSessionFile = this.sessionFile;
|
||||
const entries = this.sessionManager.loadEntries();
|
||||
const entries = this.sessionManager.getEntries();
|
||||
|
||||
// Emit before_clear event (can be cancelled)
|
||||
if (this._hookRunner?.hasHandlers("session")) {
|
||||
|
|
@ -748,7 +748,7 @@ export class AgentSession {
|
|||
throw new Error(`No API key for ${this.model.provider}`);
|
||||
}
|
||||
|
||||
const entries = this.sessionManager.loadEntries();
|
||||
const entries = this.sessionManager.getEntries();
|
||||
const settings = this.settingsManager.getCompactionSettings();
|
||||
|
||||
const preparation = prepareCompaction(entries, settings);
|
||||
|
|
@ -783,6 +783,7 @@ export class AgentSession {
|
|||
customInstructions,
|
||||
model: this.model,
|
||||
resolveApiKey: this._resolveApiKey,
|
||||
signal: this._compactionAbortController.signal,
|
||||
})) as SessionEventResult | undefined;
|
||||
|
||||
if (result?.cancel) {
|
||||
|
|
@ -811,9 +812,9 @@ export class AgentSession {
|
|||
}
|
||||
|
||||
this.sessionManager.saveCompaction(compactionEntry);
|
||||
const newEntries = this.sessionManager.loadEntries();
|
||||
const loaded = loadSessionFromEntries(newEntries);
|
||||
this.agent.replaceMessages(loaded.messages);
|
||||
const newEntries = this.sessionManager.getEntries();
|
||||
const sessionContext = this.sessionManager.buildSessionContext();
|
||||
this.agent.replaceMessages(sessionContext.messages);
|
||||
|
||||
if (this._hookRunner) {
|
||||
await this._hookRunner.emit({
|
||||
|
|
@ -909,7 +910,7 @@ export class AgentSession {
|
|||
return;
|
||||
}
|
||||
|
||||
const entries = this.sessionManager.loadEntries();
|
||||
const entries = this.sessionManager.getEntries();
|
||||
|
||||
const preparation = prepareCompaction(entries, settings);
|
||||
if (!preparation) {
|
||||
|
|
@ -944,6 +945,7 @@ export class AgentSession {
|
|||
customInstructions: undefined,
|
||||
model: this.model,
|
||||
resolveApiKey: this._resolveApiKey,
|
||||
signal: this._autoCompactionAbortController.signal,
|
||||
})) as SessionEventResult | undefined;
|
||||
|
||||
if (hookResult?.cancel) {
|
||||
|
|
@ -973,9 +975,9 @@ export class AgentSession {
|
|||
}
|
||||
|
||||
this.sessionManager.saveCompaction(compactionEntry);
|
||||
const newEntries = this.sessionManager.loadEntries();
|
||||
const loaded = loadSessionFromEntries(newEntries);
|
||||
this.agent.replaceMessages(loaded.messages);
|
||||
const newEntries = this.sessionManager.getEntries();
|
||||
const sessionContext = this.sessionManager.buildSessionContext();
|
||||
this.agent.replaceMessages(sessionContext.messages);
|
||||
|
||||
if (this._hookRunner) {
|
||||
await this._hookRunner.emit({
|
||||
|
|
@ -1281,7 +1283,7 @@ export class AgentSession {
|
|||
*/
|
||||
async switchSession(sessionPath: string): Promise<boolean> {
|
||||
const previousSessionFile = this.sessionFile;
|
||||
const oldEntries = this.sessionManager.loadEntries();
|
||||
const oldEntries = this.sessionManager.getEntries();
|
||||
|
||||
// Emit before_switch event (can be cancelled)
|
||||
if (this._hookRunner?.hasHandlers("session")) {
|
||||
|
|
@ -1306,8 +1308,8 @@ export class AgentSession {
|
|||
this.sessionManager.setSessionFile(sessionPath);
|
||||
|
||||
// Reload messages
|
||||
const entries = this.sessionManager.loadEntries();
|
||||
const loaded = loadSessionFromEntries(entries);
|
||||
const entries = this.sessionManager.getEntries();
|
||||
const sessionContext = this.sessionManager.buildSessionContext();
|
||||
|
||||
// Emit session event to hooks
|
||||
if (this._hookRunner) {
|
||||
|
|
@ -1324,22 +1326,22 @@ export class AgentSession {
|
|||
// Emit session event to custom tools
|
||||
await this._emitToolSessionEvent("switch", previousSessionFile);
|
||||
|
||||
this.agent.replaceMessages(loaded.messages);
|
||||
this.agent.replaceMessages(sessionContext.messages);
|
||||
|
||||
// Restore model if saved
|
||||
const savedModel = this.sessionManager.loadModel();
|
||||
if (savedModel) {
|
||||
if (sessionContext.model) {
|
||||
const availableModels = (await getAvailableModels()).models;
|
||||
const match = availableModels.find((m) => m.provider === savedModel.provider && m.id === savedModel.modelId);
|
||||
const match = availableModels.find(
|
||||
(m) => m.provider === sessionContext.model!.provider && m.id === sessionContext.model!.modelId,
|
||||
);
|
||||
if (match) {
|
||||
this.agent.setModel(match);
|
||||
}
|
||||
}
|
||||
|
||||
// Restore thinking level if saved (setThinkingLevel clamps to model capabilities)
|
||||
const savedThinking = this.sessionManager.loadThinkingLevel();
|
||||
if (savedThinking) {
|
||||
this.setThinkingLevel(savedThinking as ThinkingLevel);
|
||||
if (sessionContext.thinkingLevel) {
|
||||
this.setThinkingLevel(sessionContext.thinkingLevel as ThinkingLevel);
|
||||
}
|
||||
|
||||
this._reconnectToAgent();
|
||||
|
|
@ -1357,7 +1359,7 @@ export class AgentSession {
|
|||
*/
|
||||
async branch(entryIndex: number): Promise<{ selectedText: string; cancelled: boolean }> {
|
||||
const previousSessionFile = this.sessionFile;
|
||||
const entries = this.sessionManager.loadEntries();
|
||||
const entries = this.sessionManager.getEntries();
|
||||
const selectedEntry = entries[entryIndex];
|
||||
|
||||
if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") {
|
||||
|
|
@ -1394,8 +1396,8 @@ export class AgentSession {
|
|||
}
|
||||
|
||||
// Reload messages from entries (works for both file and in-memory mode)
|
||||
const newEntries = this.sessionManager.loadEntries();
|
||||
const loaded = loadSessionFromEntries(newEntries);
|
||||
const newEntries = this.sessionManager.getEntries();
|
||||
const sessionContext = this.sessionManager.buildSessionContext();
|
||||
|
||||
// Emit branch event to hooks (after branch completes)
|
||||
if (this._hookRunner) {
|
||||
|
|
@ -1414,7 +1416,7 @@ export class AgentSession {
|
|||
await this._emitToolSessionEvent("branch", previousSessionFile);
|
||||
|
||||
if (!skipConversationRestore) {
|
||||
this.agent.replaceMessages(loaded.messages);
|
||||
this.agent.replaceMessages(sessionContext.messages);
|
||||
}
|
||||
|
||||
return { selectedText, cancelled: false };
|
||||
|
|
@ -1424,7 +1426,7 @@ export class AgentSession {
|
|||
* Get all user messages from session for branch selector.
|
||||
*/
|
||||
getUserMessagesForBranching(): Array<{ entryIndex: number; text: string }> {
|
||||
const entries = this.sessionManager.loadEntries();
|
||||
const entries = this.sessionManager.getEntries();
|
||||
const result: Array<{ entryIndex: number; text: string }> = [];
|
||||
|
||||
for (let i = 0; i < entries.length; i++) {
|
||||
|
|
@ -1570,7 +1572,7 @@ export class AgentSession {
|
|||
previousSessionFile: string | null,
|
||||
): Promise<void> {
|
||||
const event: ToolSessionEvent = {
|
||||
entries: this.sessionManager.loadEntries(),
|
||||
entries: this.sessionManager.getEntries(),
|
||||
sessionFile: this.sessionFile,
|
||||
previousSessionFile,
|
||||
reason,
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import type {
|
|||
HookEvent,
|
||||
HookEventContext,
|
||||
HookUIContext,
|
||||
SessionEvent,
|
||||
SessionEventResult,
|
||||
ToolCallEvent,
|
||||
ToolCallEventResult,
|
||||
|
|
@ -229,9 +230,17 @@ export class HookRunner {
|
|||
|
||||
for (const handler of handlers) {
|
||||
try {
|
||||
const timeout = createTimeout(this.timeout);
|
||||
const handlerResult = await Promise.race([handler(event, ctx), timeout.promise]);
|
||||
timeout.clear();
|
||||
// No timeout for before_compact events (like tool_call, they may take a while)
|
||||
const isBeforeCompact = event.type === "session" && (event as SessionEvent).reason === "before_compact";
|
||||
let handlerResult: unknown;
|
||||
|
||||
if (isBeforeCompact) {
|
||||
handlerResult = await handler(event, ctx);
|
||||
} else {
|
||||
const timeout = createTimeout(this.timeout);
|
||||
handlerResult = await Promise.race([handler(event, ctx), timeout.promise]);
|
||||
timeout.clear();
|
||||
}
|
||||
|
||||
// For session events, capture the result (for before_* cancellation)
|
||||
if (event.type === "session" && handlerResult) {
|
||||
|
|
|
|||
|
|
@ -141,6 +141,8 @@ export type SessionEvent =
|
|||
model: Model<any>;
|
||||
/** Resolve API key for any model (checks settings, OAuth, env vars) */
|
||||
resolveApiKey: (model: Model<any>) => Promise<string | undefined>;
|
||||
/** Abort signal - hooks should pass this to LLM calls and check it periodically */
|
||||
signal: AbortSignal;
|
||||
})
|
||||
| (SessionEventBase & {
|
||||
reason: "compact";
|
||||
|
|
|
|||
|
|
@ -380,8 +380,13 @@ export async function getAvailableModels(
|
|||
}
|
||||
|
||||
/**
|
||||
* Find a specific model by provider and ID
|
||||
* Returns { model, error } - either model or error message
|
||||
* Find a specific model by provider and ID.
|
||||
*
|
||||
* Searches models from:
|
||||
* 1. Built-in models from @mariozechner/pi-ai
|
||||
* 2. Custom models defined in ~/.pi/agent/models.json
|
||||
*
|
||||
* Returns { model, error } - either the model or an error message.
|
||||
*/
|
||||
export function findModel(
|
||||
provider: string,
|
||||
|
|
|
|||
|
|
@ -496,7 +496,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
|
|||
};
|
||||
|
||||
// Check if session has existing data to restore
|
||||
const existingSession = sessionManager.loadSession();
|
||||
const existingSession = sessionManager.buildSessionContext();
|
||||
time("loadSession");
|
||||
const hasExistingSession = existingSession.messages.length > 0;
|
||||
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ export type SessionEntry =
|
|||
| ModelChangeEntry
|
||||
| CompactionEntry;
|
||||
|
||||
export interface LoadedSession {
|
||||
export interface SessionContext {
|
||||
messages: AppMessage[];
|
||||
thinkingLevel: string;
|
||||
model: { provider: string; modelId: string } | null;
|
||||
|
|
@ -78,6 +78,7 @@ export const SUMMARY_PREFIX = `The conversation history before this point was co
|
|||
export const SUMMARY_SUFFIX = `
|
||||
</summary>`;
|
||||
|
||||
/** Exported for compaction.test.ts */
|
||||
export function createSummaryMessage(summary: string): AppMessage {
|
||||
return {
|
||||
role: "user",
|
||||
|
|
@ -86,6 +87,7 @@ export function createSummaryMessage(summary: string): AppMessage {
|
|||
};
|
||||
}
|
||||
|
||||
/** Exported for compaction.test.ts */
|
||||
export function parseSessionEntries(content: string): SessionEntry[] {
|
||||
const entries: SessionEntry[] = [];
|
||||
const lines = content.trim().split("\n");
|
||||
|
|
@ -112,7 +114,15 @@ export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEnt
|
|||
return null;
|
||||
}
|
||||
|
||||
export function loadSessionFromEntries(entries: SessionEntry[]): LoadedSession {
|
||||
/**
|
||||
* Build the session context from entries. This is what gets sent to the LLM.
|
||||
*
|
||||
* If there's a compaction entry, returns the summary message plus messages
|
||||
* from `firstKeptEntryIndex` onwards. Otherwise returns all messages.
|
||||
*
|
||||
* Also extracts the current thinking level and model from the entries.
|
||||
*/
|
||||
export function buildSessionContext(entries: SessionEntry[]): SessionContext {
|
||||
let thinkingLevel = "off";
|
||||
let model: { provider: string; modelId: string } | null = null;
|
||||
|
||||
|
|
@ -299,7 +309,7 @@ export class SessionManager {
|
|||
}
|
||||
}
|
||||
|
||||
saveMessage(message: any): void {
|
||||
saveMessage(message: AppMessage): void {
|
||||
const entry: SessionMessageEntry = {
|
||||
type: "message",
|
||||
timestamp: new Date().toISOString(),
|
||||
|
|
@ -335,29 +345,21 @@ export class SessionManager {
|
|||
this._persist(entry);
|
||||
}
|
||||
|
||||
loadSession(): LoadedSession {
|
||||
const entries = this.loadEntries();
|
||||
return loadSessionFromEntries(entries);
|
||||
/**
|
||||
* Build the session context (what gets sent to the LLM).
|
||||
* If compacted, returns summary + kept messages. Otherwise all messages.
|
||||
* Includes thinking level and model.
|
||||
*/
|
||||
buildSessionContext(): SessionContext {
|
||||
return buildSessionContext(this.getEntries());
|
||||
}
|
||||
|
||||
loadMessages(): AppMessage[] {
|
||||
return this.loadSession().messages;
|
||||
}
|
||||
|
||||
loadThinkingLevel(): string {
|
||||
return this.loadSession().thinkingLevel;
|
||||
}
|
||||
|
||||
loadModel(): { provider: string; modelId: string } | null {
|
||||
return this.loadSession().model;
|
||||
}
|
||||
|
||||
loadEntries(): SessionEntry[] {
|
||||
if (this.inMemoryEntries.length > 0) {
|
||||
return [...this.inMemoryEntries];
|
||||
} else {
|
||||
return loadEntriesFromFile(this.sessionFile);
|
||||
}
|
||||
/**
|
||||
* Get all session entries. Returns a defensive copy.
|
||||
* Use buildSessionContext() if you need the messages for the LLM.
|
||||
*/
|
||||
getEntries(): SessionEntry[] {
|
||||
return [...this.inMemoryEntries];
|
||||
}
|
||||
|
||||
createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null {
|
||||
|
|
|
|||
|
|
@ -119,13 +119,13 @@ export {
|
|||
readOnlyTools,
|
||||
} from "./core/sdk.js";
|
||||
export {
|
||||
buildSessionContext,
|
||||
type CompactionEntry,
|
||||
createSummaryMessage,
|
||||
getLatestCompactionEntry,
|
||||
type LoadedSession,
|
||||
loadSessionFromEntries,
|
||||
type ModelChangeEntry,
|
||||
parseSessionEntries,
|
||||
type SessionContext as LoadedSession,
|
||||
type SessionEntry,
|
||||
type SessionHeader,
|
||||
type SessionInfo,
|
||||
|
|
|
|||
|
|
@ -351,7 +351,7 @@ export class InteractiveMode {
|
|||
}
|
||||
|
||||
// Load session entries if any
|
||||
const entries = this.session.sessionManager.loadEntries();
|
||||
const entries = this.session.sessionManager.getEntries();
|
||||
|
||||
// Set TUI-based UI context for custom tools
|
||||
const uiContext = this.createHookUIContext();
|
||||
|
|
@ -1067,7 +1067,7 @@ export class InteractiveMode {
|
|||
this.updateEditorBorderColor();
|
||||
}
|
||||
|
||||
const compactionEntry = getLatestCompactionEntry(this.sessionManager.loadEntries());
|
||||
const compactionEntry = getLatestCompactionEntry(this.sessionManager.getEntries());
|
||||
|
||||
for (const message of messages) {
|
||||
if (isBashExecutionMessage(message)) {
|
||||
|
|
@ -1137,7 +1137,7 @@ export class InteractiveMode {
|
|||
this.renderMessages(state.messages, { updateFooter: true, populateHistory: true });
|
||||
|
||||
// Show compaction info if session was compacted
|
||||
const entries = this.sessionManager.loadEntries();
|
||||
const entries = this.sessionManager.getEntries();
|
||||
const compactionCount = entries.filter((e) => e.type === "compaction").length;
|
||||
if (compactionCount > 0) {
|
||||
const times = compactionCount === 1 ? "1 time" : `${compactionCount} times`;
|
||||
|
|
@ -1185,7 +1185,7 @@ export class InteractiveMode {
|
|||
// Emit shutdown event to hooks
|
||||
const hookRunner = this.session.hookRunner;
|
||||
if (hookRunner?.hasHandlers("session")) {
|
||||
const entries = this.sessionManager.loadEntries();
|
||||
const entries = this.sessionManager.getEntries();
|
||||
await hookRunner.emit({
|
||||
type: "session",
|
||||
entries,
|
||||
|
|
@ -1924,7 +1924,7 @@ export class InteractiveMode {
|
|||
}
|
||||
|
||||
private async handleCompactCommand(customInstructions?: string): Promise<void> {
|
||||
const entries = this.sessionManager.loadEntries();
|
||||
const entries = this.sessionManager.getEntries();
|
||||
const messageCount = entries.filter((e) => e.type === "message").length;
|
||||
|
||||
if (messageCount < 2) {
|
||||
|
|
|
|||
|
|
@ -28,7 +28,7 @@ export async function runPrintMode(
|
|||
initialAttachments?: Attachment[],
|
||||
): Promise<void> {
|
||||
// Load entries once for session start events
|
||||
const entries = session.sessionManager.loadEntries();
|
||||
const entries = session.sessionManager.getEntries();
|
||||
|
||||
// Hook runner already has no-op UI context by default (set in main.ts)
|
||||
// Set up hooks for print mode (no UI)
|
||||
|
|
|
|||
|
|
@ -121,7 +121,7 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
|
|||
});
|
||||
|
||||
// Load entries once for session start events
|
||||
const entries = session.sessionManager.loadEntries();
|
||||
const entries = session.sessionManager.getEntries();
|
||||
|
||||
// Set up hooks with RPC-based UI context
|
||||
const hookRunner = session.hookRunner;
|
||||
|
|
|
|||
|
|
@ -141,7 +141,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
|
|||
await session.compact();
|
||||
|
||||
// Load entries from session manager
|
||||
const entries = sessionManager.loadEntries();
|
||||
const entries = sessionManager.getEntries();
|
||||
|
||||
// Should have a compaction entry
|
||||
const compactionEntries = entries.filter((e) => e.type === "compaction");
|
||||
|
|
@ -200,7 +200,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
|
|||
expect(result.summary.length).toBeGreaterThan(0);
|
||||
|
||||
// In-memory entries should have the compaction
|
||||
const entries = noSessionManager.loadEntries();
|
||||
const entries = noSessionManager.getEntries();
|
||||
const compactionEntries = entries.filter((e) => e.type === "compaction");
|
||||
expect(compactionEntries.length).toBe(1);
|
||||
} finally {
|
||||
|
|
|
|||
|
|
@ -14,9 +14,9 @@ import {
|
|||
shouldCompact,
|
||||
} from "../src/core/compaction.js";
|
||||
import {
|
||||
buildSessionContext,
|
||||
type CompactionEntry,
|
||||
createSummaryMessage,
|
||||
loadSessionFromEntries,
|
||||
parseSessionEntries,
|
||||
type SessionEntry,
|
||||
type SessionMessageEntry,
|
||||
|
|
@ -226,7 +226,7 @@ describe("createSummaryMessage", () => {
|
|||
});
|
||||
});
|
||||
|
||||
describe("loadSessionFromEntries", () => {
|
||||
describe("buildSessionContext", () => {
|
||||
it("should load all messages when no compaction", () => {
|
||||
const entries: SessionEntry[] = [
|
||||
{
|
||||
|
|
@ -241,7 +241,7 @@ describe("loadSessionFromEntries", () => {
|
|||
createMessageEntry(createAssistantMessage("b")),
|
||||
];
|
||||
|
||||
const loaded = loadSessionFromEntries(entries);
|
||||
const loaded = buildSessionContext(entries);
|
||||
expect(loaded.messages.length).toBe(4);
|
||||
expect(loaded.thinkingLevel).toBe("off");
|
||||
expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" });
|
||||
|
|
@ -265,7 +265,7 @@ describe("loadSessionFromEntries", () => {
|
|||
createMessageEntry(createAssistantMessage("c")),
|
||||
];
|
||||
|
||||
const loaded = loadSessionFromEntries(entries);
|
||||
const loaded = buildSessionContext(entries);
|
||||
// summary + kept (u2,a2 from idx 3-4) + after (u3,a3 from idx 6-7) = 5
|
||||
expect(loaded.messages.length).toBe(5);
|
||||
expect(loaded.messages[0].role).toBe("user");
|
||||
|
|
@ -293,7 +293,7 @@ describe("loadSessionFromEntries", () => {
|
|||
createMessageEntry(createAssistantMessage("d")),
|
||||
];
|
||||
|
||||
const loaded = loadSessionFromEntries(entries);
|
||||
const loaded = buildSessionContext(entries);
|
||||
// summary + kept from idx 6 (u3,c) + after (u4,d) = 5
|
||||
expect(loaded.messages.length).toBe(5);
|
||||
expect((loaded.messages[0] as any).content).toContain("Second summary");
|
||||
|
|
@ -316,7 +316,7 @@ describe("loadSessionFromEntries", () => {
|
|||
createCompactionEntry("Second summary", 0), // index 0 is before compaction1, should still work
|
||||
];
|
||||
|
||||
const loaded = loadSessionFromEntries(entries);
|
||||
const loaded = buildSessionContext(entries);
|
||||
// Keeps from index 0, but compaction entries are skipped, so u1,a1,u2,b = 4 + summary = 5
|
||||
// Actually index 0 is session header, so messages are u1,a1,u2,b
|
||||
expect(loaded.messages.length).toBe(5); // summary + 4 messages
|
||||
|
|
@ -336,7 +336,7 @@ describe("loadSessionFromEntries", () => {
|
|||
{ type: "thinking_level_change", timestamp: "", thinkingLevel: "high" },
|
||||
];
|
||||
|
||||
const loaded = loadSessionFromEntries(entries);
|
||||
const loaded = buildSessionContext(entries);
|
||||
// model_change is later overwritten by assistant message's model info
|
||||
expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" });
|
||||
expect(loaded.thinkingLevel).toBe("high");
|
||||
|
|
@ -368,7 +368,7 @@ describe("Large session fixture", () => {
|
|||
|
||||
it("should load session correctly", () => {
|
||||
const entries = loadLargeSessionEntries();
|
||||
const loaded = loadSessionFromEntries(entries);
|
||||
const loaded = buildSessionContext(entries);
|
||||
|
||||
expect(loaded.messages.length).toBeGreaterThan(100);
|
||||
expect(loaded.model).not.toBeNull();
|
||||
|
|
@ -405,7 +405,7 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
|
|||
|
||||
it("should produce valid session after compaction", async () => {
|
||||
const entries = loadLargeSessionEntries();
|
||||
const loaded = loadSessionFromEntries(entries);
|
||||
const loaded = buildSessionContext(entries);
|
||||
const model = getModel("anthropic", "claude-sonnet-4-5")!;
|
||||
|
||||
const compactionEvent = await compact(
|
||||
|
|
@ -417,7 +417,7 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
|
|||
|
||||
// Simulate appending compaction to entries
|
||||
const newEntries = [...entries, compactionEvent];
|
||||
const reloaded = loadSessionFromEntries(newEntries);
|
||||
const reloaded = buildSessionContext(newEntries);
|
||||
|
||||
// Should have summary + kept messages
|
||||
expect(reloaded.messages.length).toBeLessThan(loaded.messages.length);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue