mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-17 11:04:51 +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
1013
compact.jsonl
Normal file
1013
compact.jsonl
Normal file
File diff suppressed because one or more lines are too long
|
|
@ -6359,23 +6359,6 @@ export const MODELS = {
|
|||
contextWindow: 128000,
|
||||
maxTokens: 16384,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"meta-llama/llama-3.1-70b-instruct": {
|
||||
id: "meta-llama/llama-3.1-70b-instruct",
|
||||
name: "Meta: Llama 3.1 70B Instruct",
|
||||
api: "openai-completions",
|
||||
provider: "openrouter",
|
||||
baseUrl: "https://openrouter.ai/api/v1",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0.39999999999999997,
|
||||
output: 0.39999999999999997,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 131072,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"meta-llama/llama-3.1-8b-instruct": {
|
||||
id: "meta-llama/llama-3.1-8b-instruct",
|
||||
name: "Meta: Llama 3.1 8B Instruct",
|
||||
|
|
@ -6410,6 +6393,23 @@ export const MODELS = {
|
|||
contextWindow: 10000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"meta-llama/llama-3.1-70b-instruct": {
|
||||
id: "meta-llama/llama-3.1-70b-instruct",
|
||||
name: "Meta: Llama 3.1 70B Instruct",
|
||||
api: "openai-completions",
|
||||
provider: "openrouter",
|
||||
baseUrl: "https://openrouter.ai/api/v1",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0.39999999999999997,
|
||||
output: 0.39999999999999997,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 131072,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"mistralai/mistral-nemo": {
|
||||
id: "mistralai/mistral-nemo",
|
||||
name: "Mistral: Mistral Nemo",
|
||||
|
|
@ -6546,6 +6546,23 @@ export const MODELS = {
|
|||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"openai/gpt-4o-2024-05-13": {
|
||||
id: "openai/gpt-4o-2024-05-13",
|
||||
name: "OpenAI: GPT-4o (2024-05-13)",
|
||||
api: "openai-completions",
|
||||
provider: "openrouter",
|
||||
baseUrl: "https://openrouter.ai/api/v1",
|
||||
reasoning: false,
|
||||
input: ["text", "image"],
|
||||
cost: {
|
||||
input: 5,
|
||||
output: 15,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"openai/gpt-4o": {
|
||||
id: "openai/gpt-4o",
|
||||
name: "OpenAI: GPT-4o",
|
||||
|
|
@ -6580,23 +6597,6 @@ export const MODELS = {
|
|||
contextWindow: 128000,
|
||||
maxTokens: 64000,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"openai/gpt-4o-2024-05-13": {
|
||||
id: "openai/gpt-4o-2024-05-13",
|
||||
name: "OpenAI: GPT-4o (2024-05-13)",
|
||||
api: "openai-completions",
|
||||
provider: "openrouter",
|
||||
baseUrl: "https://openrouter.ai/api/v1",
|
||||
reasoning: false,
|
||||
input: ["text", "image"],
|
||||
cost: {
|
||||
input: 5,
|
||||
output: 15,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"meta-llama/llama-3-70b-instruct": {
|
||||
id: "meta-llama/llama-3-70b-instruct",
|
||||
name: "Meta: Llama 3 70B Instruct",
|
||||
|
|
@ -6716,23 +6716,6 @@ export const MODELS = {
|
|||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"openai/gpt-4-turbo-preview": {
|
||||
id: "openai/gpt-4-turbo-preview",
|
||||
name: "OpenAI: GPT-4 Turbo Preview",
|
||||
api: "openai-completions",
|
||||
provider: "openrouter",
|
||||
baseUrl: "https://openrouter.ai/api/v1",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 10,
|
||||
output: 30,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"openai/gpt-3.5-turbo-0613": {
|
||||
id: "openai/gpt-3.5-turbo-0613",
|
||||
name: "OpenAI: GPT-3.5 Turbo (older v0613)",
|
||||
|
|
@ -6750,6 +6733,23 @@ export const MODELS = {
|
|||
contextWindow: 4095,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"openai/gpt-4-turbo-preview": {
|
||||
id: "openai/gpt-4-turbo-preview",
|
||||
name: "OpenAI: GPT-4 Turbo Preview",
|
||||
api: "openai-completions",
|
||||
provider: "openrouter",
|
||||
baseUrl: "https://openrouter.ai/api/v1",
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 10,
|
||||
output: 30,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"mistralai/mistral-tiny": {
|
||||
id: "mistralai/mistral-tiny",
|
||||
name: "Mistral Tiny",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 │
|
||||
└─────┴─────┴─────┴─────┴──────┴─────┴─────┴─────┴──────┴──────┴─────┴──────┘
|
||||
↑ └────────┬────────┘ └────────────┬────────────┘
|
||||
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 = 6
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
// 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);
|
||||
const handlerResult = await Promise.race([handler(event, ctx), timeout.promise]);
|
||||
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) {
|
||||
/**
|
||||
* Get all session entries. Returns a defensive copy.
|
||||
* Use buildSessionContext() if you need the messages for the LLM.
|
||||
*/
|
||||
getEntries(): SessionEntry[] {
|
||||
return [...this.inMemoryEntries];
|
||||
} else {
|
||||
return loadEntriesFromFile(this.sessionFile);
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
|
|
|
|||
|
|
@ -12,9 +12,9 @@
|
|||
|
||||
import type { AppMessage } from "@mariozechner/pi-agent-core";
|
||||
import {
|
||||
buildSessionContext,
|
||||
type CompactionEntry,
|
||||
type LoadedSession,
|
||||
loadSessionFromEntries,
|
||||
type ModelChangeEntry,
|
||||
type SessionEntry,
|
||||
type SessionMessageEntry,
|
||||
|
|
@ -285,7 +285,7 @@ export class MomSessionManager {
|
|||
/** Load session with compaction support */
|
||||
loadSession(): LoadedSession {
|
||||
const entries = this.loadEntries();
|
||||
return loadSessionFromEntries(entries);
|
||||
return buildSessionContext(entries);
|
||||
}
|
||||
|
||||
loadEntries(): SessionEntry[] {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue