getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-21 05:02:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve compaction hooks: add signal, no timeout, SessionManager cleanup, docs

Browse source
This commit is contained in:
Mario Zechner 2025-12-24 13:54:05 +01:00
parent a2664ba38a
commit 705ba5d4f2
19 changed files with 1236 additions and 207 deletions
Download patch file Download diff file Expand all files Collapse all files

1013
compact.jsonl Normal file
View file

File diff suppressed because one or more lines are too long

102
packages/ai/src/models.generated.ts
View file

@ -6359,23 +6359,6 @@ export const MODELS = {
contextWindow: 128000, contextWindow: 128000,
maxTokens: 16384, maxTokens: 16384,
} satisfies Model<"openai-completions">, } 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": { "meta-llama/llama-3.1-8b-instruct": {
id: "meta-llama/llama-3.1-8b-instruct", id: "meta-llama/llama-3.1-8b-instruct",
name: "Meta: Llama 3.1 8B Instruct", name: "Meta: Llama 3.1 8B Instruct",
@ -6410,6 +6393,23 @@ export const MODELS = {
contextWindow: 10000, contextWindow: 10000,
maxTokens: 4096, maxTokens: 4096,
} satisfies Model<"openai-completions">, } 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": { "mistralai/mistral-nemo": {
id: "mistralai/mistral-nemo", id: "mistralai/mistral-nemo",
name: "Mistral: Mistral Nemo", name: "Mistral: Mistral Nemo",
@ -6546,6 +6546,23 @@ export const MODELS = {
contextWindow: 128000, contextWindow: 128000,
maxTokens: 4096, maxTokens: 4096,
} satisfies Model<"openai-completions">, } 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": { "openai/gpt-4o": {
id: "openai/gpt-4o", id: "openai/gpt-4o",
name: "OpenAI: GPT-4o", name: "OpenAI: GPT-4o",
@ -6580,23 +6597,6 @@ export const MODELS = {
contextWindow: 128000, contextWindow: 128000,
maxTokens: 64000, maxTokens: 64000,
} satisfies Model<"openai-completions">, } 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": { "meta-llama/llama-3-70b-instruct": {
id: "meta-llama/llama-3-70b-instruct", id: "meta-llama/llama-3-70b-instruct",
name: "Meta: Llama 3 70B Instruct", name: "Meta: Llama 3 70B Instruct",
@ -6716,23 +6716,6 @@ export const MODELS = {
contextWindow: 128000, contextWindow: 128000,
maxTokens: 4096, maxTokens: 4096,
} satisfies Model<"openai-completions">, } 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": { "openai/gpt-3.5-turbo-0613": {
id: "openai/gpt-3.5-turbo-0613", id: "openai/gpt-3.5-turbo-0613",
name: "OpenAI: GPT-3.5 Turbo (older v0613)", name: "OpenAI: GPT-3.5 Turbo (older v0613)",
@ -6750,6 +6733,23 @@ export const MODELS = {
contextWindow: 4095, contextWindow: 4095,
maxTokens: 4096, maxTokens: 4096,
} satisfies Model<"openai-completions">, } 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": { "mistralai/mistral-tiny": {
id: "mistralai/mistral-tiny", id: "mistralai/mistral-tiny",
name: "Mistral Tiny", name: "Mistral Tiny",

5
packages/coding-agent/CHANGELOG.md
View file

@ -10,6 +10,11 @@
- `resolveApiKey`: Function to resolve API keys for any model (checks settings, OAuth, env vars) - `resolveApiKey`: Function to resolve API keys for any model (checks settings, OAuth, env vars)
- Removed `apiKey` string in favor of `resolveApiKey` for more flexibility - 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 ## [0.27.5] - 2025-12-24
### Added ### Added

76
packages/coding-agent/docs/hooks.md
View file

@ -195,24 +195,26 @@ Legend:
``` ```
Session entries (before compaction): Session entries (before compaction):
index: 0 1 2 3 4 5 6 7 8 9 10 11 index: 0 1 2 3 4 5 6 7 8 9 10
┌─────┬─────┬─────┬─────┬──────┬─────┬─────┬─────┬──────┬──────┬─────┬──────┐ ┌─────┬─────┬─────┬─────┬──────┬─────┬─────┬──────┬──────┬─────┬──────┐
│ hdr │ cmp │ usr │ ass │ tool │ ass │ usr │ ass │ tool │ tool │ ass │ tool │ │ hdr │ cmp │ usr │ ass │ tool │ usr │ ass │ tool │ tool │ ass │ tool │
└─────┴─────┴─────┴─────┴──────┴─────┴─────┴─────┴──────┴──────┴─────┴──────┘ └─────┴─────┴─────┴─────┴──────┴─────┴─────┴──────┴──────┴─────┴──────┘
↑ └────────┬────────┘ └────────────┬────────────┘ ↑ └───────┬───────┘ └────────────┬────────────┘
previousSummary messagesToSummarize messagesToKeep previousSummary messagesToSummarize messagesToKeep
↑ ↑
cutPoint.firstKeptEntryIndex = 6 cutPoint.firstKeptEntryIndex = 5
After compaction (new entry appended): After compaction (new entry appended):
index: 0 1 2 3 4 5 6 7 8 9 10 11 12 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 │ cmp │ │ hdr │ cmp │ usr │ ass │ tool │ usr │ ass │ tool │ tool │ ass │ tool │ cmp │
└─────┴─────┴─────┴─────┴──────┴─────┴─────┴─────┴──────┴──────┴─────┴──────┴─────┘ └─────┴─────┴─────┴─────┴──────┴─────┴─────┴──────┴──────┴─────┴──────┴─────┘
└────────┬────────┘ └────────────┬────────────┘ ↑ └──────────┬───────────┘ └────────────────────────┬─────────────────┘
ignored loaded firstKeptEntryIndex = 6 not sent to LLM sent to LLM
on reload on reload (stored in new cmp) ↑
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). 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: What gets sent to the LLM as context:
┌──────────────────────────────────────────────────────────────────────────────────────┐ 5 6 7 8 9 10
│ [system] [summary] [usr idx 6] [ass idx 7] [tool idx 8] [tool idx 9] [ass idx 10] [tool idx 11] │ ┌────────┬─────────┬─────┬─────┬──────┬──────┬─────┬──────┐
└──────────────────────────────────────────────────────────────────────────────────────┘ │ system │ summary │ usr │ ass │ tool │ tool │ ass │ tool │
↑ └──────────────────────┬──────────────────────┘ └────────┴─────────┴─────┴─────┴──────┴──────┴─────┴──────┘
from new cmp's summary messages from firstKeptEntryIndex onwards ↑ └─────────────────┬────────────────┘
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`: **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. | | `model` | Model to use for summarization. |
| `resolveApiKey` | Function to resolve API key for any model: `await resolveApiKey(model)` | | `resolveApiKey` | Function to resolve API key for any model: `await resolveApiKey(model)` |
| `customInstructions` | Optional focus for summary (from `/compact <instructions>`). | | `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 See [examples/hooks/custom-compaction.ts](../examples/hooks/custom-compaction.ts) for a complete example.
// 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.
**After compaction (`compact` event):** **After compaction (`compact` event):**
- `event.compactionEntry`: The saved compaction entry - `event.compactionEntry`: The saved compaction entry

2
packages/coding-agent/examples/hooks/README.md
View file

@ -25,7 +25,7 @@ Prevents session changes when there are uncommitted git changes. Blocks clear/sw
### auto-commit-on-exit.ts ### auto-commit-on-exit.ts
Automatically commits changes when the agent exits (shutdown event). Uses the last assistant message to generate a commit message. 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. 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 ## Usage

45
packages/coding-agent/examples/hooks/full-compaction.ts → packages/coding-agent/examples/hooks/custom-compaction.ts
View file

@ -1,41 +1,54 @@
/** /**
* Full Context Compaction Hook * Custom Compaction Hook
* *
* Replaces the default compaction behavior with a full summary of the entire context. * 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: * 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 * 2. Discards all old turns completely, keeping only the summary
* *
* This is useful when you want maximum context window space for new work * This example also demonstrates using a different model (Gemini Flash) for summarization,
* at the cost of losing exact conversation history. * which can be cheaper/faster than the main conversation model.
* *
* Usage: * Usage:
* pi --hook examples/hooks/full-compaction.ts * pi --hook examples/hooks/custom-compaction.ts
*/ */
import { complete } from "@mariozechner/pi-ai"; 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"; import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks";
export default function (pi: HookAPI) { export default function (pi: HookAPI) {
pi.on("session", async (event, ctx) => { pi.on("session", async (event, ctx) => {
if (event.reason !== "before_compact") return; 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 { messagesToSummarize, messagesToKeep, previousSummary, tokensBefore, resolveApiKey, entries, signal } = event;
const allMessages = [...messagesToSummarize, ...messagesToKeep];
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); const apiKey = await resolveApiKey(model);
if (!apiKey) { if (!apiKey) {
ctx.ui.notify(`No API key for ${model.provider}, using default compaction`, "warning"); ctx.ui.notify(`No API key for ${model.provider}, using default compaction`, "warning");
return; 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); const transformedMessages = messageTransformer(allMessages);
// Include previous summary context if available // Include previous summary context if available
@ -68,8 +81,8 @@ Format the summary as structured markdown with clear sections.`,
]; ];
try { try {
// Use the same model with resolved API key // Pass signal to honor abort requests (e.g., user cancels compaction)
const response = await complete(model, { messages: summaryMessages }, { apiKey, maxTokens: 8192 }); const response = await complete(model, { messages: summaryMessages }, { apiKey, maxTokens: 8192, signal });
const summary = response.content const summary = response.content
.filter((c): c is { type: "text"; text: string } => c.type === "text") .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"); .join("\n");
if (!summary.trim()) { if (!summary.trim()) {
ctx.ui.notify("Compaction summary was empty, using default compaction", "warning"); if (!signal.aborted) ctx.ui.notify("Compaction summary was empty, using default compaction", "warning");
return; // Fall back to default compaction return;
} }
// Return a compaction entry that discards ALL messages // Return a compaction entry that discards ALL messages

54
packages/coding-agent/src/core/agent-session.ts
View file

@ -24,7 +24,7 @@ import { exportSessionToHtml } from "./export-html.js";
import type { HookRunner, SessionEventResult, TurnEndEvent, TurnStartEvent } from "./hooks/index.js"; import type { HookRunner, SessionEventResult, TurnEndEvent, TurnStartEvent } from "./hooks/index.js";
import type { BashExecutionMessage } from "./messages.js"; import type { BashExecutionMessage } from "./messages.js";
import { getApiKeyForModel, getAvailableModels } from "./model-config.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 type { SettingsManager, SkillsSettings } from "./settings-manager.js";
import { expandSlashCommand, type FileSlashCommand } from "./slash-commands.js"; import { expandSlashCommand, type FileSlashCommand } from "./slash-commands.js";
@ -510,7 +510,7 @@ export class AgentSession {
*/ */
async reset(): Promise<boolean> { async reset(): Promise<boolean> {
const previousSessionFile = this.sessionFile; const previousSessionFile = this.sessionFile;
const entries = this.sessionManager.loadEntries(); const entries = this.sessionManager.getEntries();
// Emit before_clear event (can be cancelled) // Emit before_clear event (can be cancelled)
if (this._hookRunner?.hasHandlers("session")) { if (this._hookRunner?.hasHandlers("session")) {
@ -748,7 +748,7 @@ export class AgentSession {
throw new Error(`No API key for ${this.model.provider}`); 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 settings = this.settingsManager.getCompactionSettings();
const preparation = prepareCompaction(entries, settings); const preparation = prepareCompaction(entries, settings);
@ -783,6 +783,7 @@ export class AgentSession {
customInstructions, customInstructions,
model: this.model, model: this.model,
resolveApiKey: this._resolveApiKey, resolveApiKey: this._resolveApiKey,
signal: this._compactionAbortController.signal,
})) as SessionEventResult | undefined; })) as SessionEventResult | undefined;
if (result?.cancel) { if (result?.cancel) {
@ -811,9 +812,9 @@ export class AgentSession {
} }
this.sessionManager.saveCompaction(compactionEntry); this.sessionManager.saveCompaction(compactionEntry);
const newEntries = this.sessionManager.loadEntries(); const newEntries = this.sessionManager.getEntries();
const loaded = loadSessionFromEntries(newEntries); const sessionContext = this.sessionManager.buildSessionContext();
this.agent.replaceMessages(loaded.messages); this.agent.replaceMessages(sessionContext.messages);
if (this._hookRunner) { if (this._hookRunner) {
await this._hookRunner.emit({ await this._hookRunner.emit({
@ -909,7 +910,7 @@ export class AgentSession {
return; return;
} }
const entries = this.sessionManager.loadEntries(); const entries = this.sessionManager.getEntries();
const preparation = prepareCompaction(entries, settings); const preparation = prepareCompaction(entries, settings);
if (!preparation) { if (!preparation) {
@ -944,6 +945,7 @@ export class AgentSession {
customInstructions: undefined, customInstructions: undefined,
model: this.model, model: this.model,
resolveApiKey: this._resolveApiKey, resolveApiKey: this._resolveApiKey,
signal: this._autoCompactionAbortController.signal,
})) as SessionEventResult | undefined; })) as SessionEventResult | undefined;
if (hookResult?.cancel) { if (hookResult?.cancel) {
@ -973,9 +975,9 @@ export class AgentSession {
} }
this.sessionManager.saveCompaction(compactionEntry); this.sessionManager.saveCompaction(compactionEntry);
const newEntries = this.sessionManager.loadEntries(); const newEntries = this.sessionManager.getEntries();
const loaded = loadSessionFromEntries(newEntries); const sessionContext = this.sessionManager.buildSessionContext();
this.agent.replaceMessages(loaded.messages); this.agent.replaceMessages(sessionContext.messages);
if (this._hookRunner) { if (this._hookRunner) {
await this._hookRunner.emit({ await this._hookRunner.emit({
@ -1281,7 +1283,7 @@ export class AgentSession {
*/ */
async switchSession(sessionPath: string): Promise<boolean> { async switchSession(sessionPath: string): Promise<boolean> {
const previousSessionFile = this.sessionFile; const previousSessionFile = this.sessionFile;
const oldEntries = this.sessionManager.loadEntries(); const oldEntries = this.sessionManager.getEntries();
// Emit before_switch event (can be cancelled) // Emit before_switch event (can be cancelled)
if (this._hookRunner?.hasHandlers("session")) { if (this._hookRunner?.hasHandlers("session")) {
@ -1306,8 +1308,8 @@ export class AgentSession {
this.sessionManager.setSessionFile(sessionPath); this.sessionManager.setSessionFile(sessionPath);
// Reload messages // Reload messages
const entries = this.sessionManager.loadEntries(); const entries = this.sessionManager.getEntries();
const loaded = loadSessionFromEntries(entries); const sessionContext = this.sessionManager.buildSessionContext();
// Emit session event to hooks // Emit session event to hooks
if (this._hookRunner) { if (this._hookRunner) {
@ -1324,22 +1326,22 @@ export class AgentSession {
// Emit session event to custom tools // Emit session event to custom tools
await this._emitToolSessionEvent("switch", previousSessionFile); await this._emitToolSessionEvent("switch", previousSessionFile);
this.agent.replaceMessages(loaded.messages); this.agent.replaceMessages(sessionContext.messages);
// Restore model if saved // Restore model if saved
const savedModel = this.sessionManager.loadModel(); if (sessionContext.model) {
if (savedModel) {
const availableModels = (await getAvailableModels()).models; 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) { if (match) {
this.agent.setModel(match); this.agent.setModel(match);
} }
} }
// Restore thinking level if saved (setThinkingLevel clamps to model capabilities) // Restore thinking level if saved (setThinkingLevel clamps to model capabilities)
const savedThinking = this.sessionManager.loadThinkingLevel(); if (sessionContext.thinkingLevel) {
if (savedThinking) { this.setThinkingLevel(sessionContext.thinkingLevel as ThinkingLevel);
this.setThinkingLevel(savedThinking as ThinkingLevel);
} }
this._reconnectToAgent(); this._reconnectToAgent();
@ -1357,7 +1359,7 @@ export class AgentSession {
*/ */
async branch(entryIndex: number): Promise<{ selectedText: string; cancelled: boolean }> { async branch(entryIndex: number): Promise<{ selectedText: string; cancelled: boolean }> {
const previousSessionFile = this.sessionFile; const previousSessionFile = this.sessionFile;
const entries = this.sessionManager.loadEntries(); const entries = this.sessionManager.getEntries();
const selectedEntry = entries[entryIndex]; const selectedEntry = entries[entryIndex];
if (!selectedEntry || selectedEntry.type !== "message" || selectedEntry.message.role !== "user") { 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) // Reload messages from entries (works for both file and in-memory mode)
const newEntries = this.sessionManager.loadEntries(); const newEntries = this.sessionManager.getEntries();
const loaded = loadSessionFromEntries(newEntries); const sessionContext = this.sessionManager.buildSessionContext();
// Emit branch event to hooks (after branch completes) // Emit branch event to hooks (after branch completes)
if (this._hookRunner) { if (this._hookRunner) {
@ -1414,7 +1416,7 @@ export class AgentSession {
await this._emitToolSessionEvent("branch", previousSessionFile); await this._emitToolSessionEvent("branch", previousSessionFile);
if (!skipConversationRestore) { if (!skipConversationRestore) {
this.agent.replaceMessages(loaded.messages); this.agent.replaceMessages(sessionContext.messages);
} }
return { selectedText, cancelled: false }; return { selectedText, cancelled: false };
@ -1424,7 +1426,7 @@ export class AgentSession {
* Get all user messages from session for branch selector. * Get all user messages from session for branch selector.
*/ */
getUserMessagesForBranching(): Array<{ entryIndex: number; text: string }> { getUserMessagesForBranching(): Array<{ entryIndex: number; text: string }> {
const entries = this.sessionManager.loadEntries(); const entries = this.sessionManager.getEntries();
const result: Array<{ entryIndex: number; text: string }> = []; const result: Array<{ entryIndex: number; text: string }> = [];
for (let i = 0; i < entries.length; i++) { for (let i = 0; i < entries.length; i++) {
@ -1570,7 +1572,7 @@ export class AgentSession {
previousSessionFile: string | null, previousSessionFile: string | null,
): Promise<void> { ): Promise<void> {
const event: ToolSessionEvent = { const event: ToolSessionEvent = {
entries: this.sessionManager.loadEntries(), entries: this.sessionManager.getEntries(),
sessionFile: this.sessionFile, sessionFile: this.sessionFile,
previousSessionFile, previousSessionFile,
reason, reason,

15
packages/coding-agent/src/core/hooks/runner.ts
View file

@ -11,6 +11,7 @@ import type {
HookEvent, HookEvent,
HookEventContext, HookEventContext,
HookUIContext, HookUIContext,
SessionEvent,
SessionEventResult, SessionEventResult,
ToolCallEvent, ToolCallEvent,
ToolCallEventResult, ToolCallEventResult,
@ -229,9 +230,17 @@ export class HookRunner {
for (const handler of handlers) { for (const handler of handlers) {
try { try {
const timeout = createTimeout(this.timeout); // No timeout for before_compact events (like tool_call, they may take a while)
const handlerResult = await Promise.race([handler(event, ctx), timeout.promise]); const isBeforeCompact = event.type === "session" && (event as SessionEvent).reason === "before_compact";
timeout.clear(); 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) // For session events, capture the result (for before_* cancellation)
if (event.type === "session" && handlerResult) { if (event.type === "session" && handlerResult) {

2
packages/coding-agent/src/core/hooks/types.ts
View file

@ -141,6 +141,8 @@ export type SessionEvent =
model: Model<any>; model: Model<any>;
/** Resolve API key for any model (checks settings, OAuth, env vars) */ /** Resolve API key for any model (checks settings, OAuth, env vars) */
resolveApiKey: (model: Model<any>) => Promise<string | undefined>; resolveApiKey: (model: Model<any>) => Promise<string | undefined>;
/** Abort signal - hooks should pass this to LLM calls and check it periodically */
signal: AbortSignal;
}) })
| (SessionEventBase & { | (SessionEventBase & {
reason: "compact"; reason: "compact";

9
packages/coding-agent/src/core/model-config.ts
View file

@ -380,8 +380,13 @@ export async function getAvailableModels(
} }
/** /**
* Find a specific model by provider and ID * Find a specific model by provider and ID.
* Returns { model, error } - either model or error message *
* 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( export function findModel(
provider: string, provider: string,

2
packages/coding-agent/src/core/sdk.ts
View file

@ -496,7 +496,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
}; };
// Check if session has existing data to restore // Check if session has existing data to restore
const existingSession = sessionManager.loadSession(); const existingSession = sessionManager.buildSessionContext();
time("loadSession"); time("loadSession");
const hasExistingSession = existingSession.messages.length > 0; const hasExistingSession = existingSession.messages.length > 0;

50
packages/coding-agent/src/core/session-manager.ts
View file

@ -54,7 +54,7 @@ export type SessionEntry =
| ModelChangeEntry | ModelChangeEntry
| CompactionEntry; | CompactionEntry;
export interface LoadedSession { export interface SessionContext {
messages: AppMessage[]; messages: AppMessage[];
thinkingLevel: string; thinkingLevel: string;
model: { provider: string; modelId: string } | null; 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 = ` export const SUMMARY_SUFFIX = `
</summary>`; </summary>`;
/** Exported for compaction.test.ts */
export function createSummaryMessage(summary: string): AppMessage { export function createSummaryMessage(summary: string): AppMessage {
return { return {
role: "user", role: "user",
@ -86,6 +87,7 @@ export function createSummaryMessage(summary: string): AppMessage {
}; };
} }
/** Exported for compaction.test.ts */
export function parseSessionEntries(content: string): SessionEntry[] { export function parseSessionEntries(content: string): SessionEntry[] {
const entries: SessionEntry[] = []; const entries: SessionEntry[] = [];
const lines = content.trim().split("\n"); const lines = content.trim().split("\n");
@ -112,7 +114,15 @@ export function getLatestCompactionEntry(entries: SessionEntry[]): CompactionEnt
return null; 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 thinkingLevel = "off";
let model: { provider: string; modelId: string } | null = null; 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 = { const entry: SessionMessageEntry = {
type: "message", type: "message",
timestamp: new Date().toISOString(), timestamp: new Date().toISOString(),
@ -335,29 +345,21 @@ export class SessionManager {
this._persist(entry); this._persist(entry);
} }
loadSession(): LoadedSession { /**
const entries = this.loadEntries(); * Build the session context (what gets sent to the LLM).
return loadSessionFromEntries(entries); * 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; * Get all session entries. Returns a defensive copy.
} * Use buildSessionContext() if you need the messages for the LLM.
*/
loadThinkingLevel(): string { getEntries(): SessionEntry[] {
return this.loadSession().thinkingLevel; return [...this.inMemoryEntries];
}
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);
}
} }
createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null { createBranchedSessionFromEntries(entries: SessionEntry[], branchBeforeIndex: number): string | null {

4
packages/coding-agent/src/index.ts
View file

@ -119,13 +119,13 @@ export {
readOnlyTools, readOnlyTools,
} from "./core/sdk.js"; } from "./core/sdk.js";
export { export {
buildSessionContext,
type CompactionEntry, type CompactionEntry,
createSummaryMessage, createSummaryMessage,
getLatestCompactionEntry, getLatestCompactionEntry,
type LoadedSession,
loadSessionFromEntries,
type ModelChangeEntry, type ModelChangeEntry,
parseSessionEntries, parseSessionEntries,
type SessionContext as LoadedSession,
type SessionEntry, type SessionEntry,
type SessionHeader, type SessionHeader,
type SessionInfo, type SessionInfo,

10
packages/coding-agent/src/modes/interactive/interactive-mode.ts
View file

@ -351,7 +351,7 @@ export class InteractiveMode {
} }
// Load session entries if any // 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 // Set TUI-based UI context for custom tools
const uiContext = this.createHookUIContext(); const uiContext = this.createHookUIContext();
@ -1067,7 +1067,7 @@ export class InteractiveMode {
this.updateEditorBorderColor(); this.updateEditorBorderColor();
} }
const compactionEntry = getLatestCompactionEntry(this.sessionManager.loadEntries()); const compactionEntry = getLatestCompactionEntry(this.sessionManager.getEntries());
for (const message of messages) { for (const message of messages) {
if (isBashExecutionMessage(message)) { if (isBashExecutionMessage(message)) {
@ -1137,7 +1137,7 @@ export class InteractiveMode {
this.renderMessages(state.messages, { updateFooter: true, populateHistory: true }); this.renderMessages(state.messages, { updateFooter: true, populateHistory: true });
// Show compaction info if session was compacted // 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; const compactionCount = entries.filter((e) => e.type === "compaction").length;
if (compactionCount > 0) { if (compactionCount > 0) {
const times = compactionCount === 1 ? "1 time" : `${compactionCount} times`; const times = compactionCount === 1 ? "1 time" : `${compactionCount} times`;
@ -1185,7 +1185,7 @@ export class InteractiveMode {
// Emit shutdown event to hooks // Emit shutdown event to hooks
const hookRunner = this.session.hookRunner; const hookRunner = this.session.hookRunner;
if (hookRunner?.hasHandlers("session")) { if (hookRunner?.hasHandlers("session")) {
const entries = this.sessionManager.loadEntries(); const entries = this.sessionManager.getEntries();
await hookRunner.emit({ await hookRunner.emit({
type: "session", type: "session",
entries, entries,
@ -1924,7 +1924,7 @@ export class InteractiveMode {
} }
private async handleCompactCommand(customInstructions?: string): Promise<void> { 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; const messageCount = entries.filter((e) => e.type === "message").length;
if (messageCount < 2) { if (messageCount < 2) {

2
packages/coding-agent/src/modes/print-mode.ts
View file

@ -28,7 +28,7 @@ export async function runPrintMode(
initialAttachments?: Attachment[], initialAttachments?: Attachment[],
): Promise<void> { ): Promise<void> {
// Load entries once for session start events // 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) // Hook runner already has no-op UI context by default (set in main.ts)
// Set up hooks for print mode (no UI) // Set up hooks for print mode (no UI)

2
packages/coding-agent/src/modes/rpc/rpc-mode.ts
View file

@ -121,7 +121,7 @@ export async function runRpcMode(session: AgentSession): Promise<never> {
}); });
// Load entries once for session start events // 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 // Set up hooks with RPC-based UI context
const hookRunner = session.hookRunner; const hookRunner = session.hookRunner;

4
packages/coding-agent/test/agent-session-compaction.test.ts
View file

@ -141,7 +141,7 @@ describe.skipIf(!API_KEY)("AgentSession compaction e2e", () => {
await session.compact(); await session.compact();
// Load entries from session manager // Load entries from session manager
const entries = sessionManager.loadEntries(); const entries = sessionManager.getEntries();
// Should have a compaction entry // Should have a compaction entry
const compactionEntries = entries.filter((e) => e.type === "compaction"); 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); expect(result.summary.length).toBeGreaterThan(0);
// In-memory entries should have the compaction // In-memory entries should have the compaction
const entries = noSessionManager.loadEntries(); const entries = noSessionManager.getEntries();
const compactionEntries = entries.filter((e) => e.type === "compaction"); const compactionEntries = entries.filter((e) => e.type === "compaction");
expect(compactionEntries.length).toBe(1); expect(compactionEntries.length).toBe(1);
} finally { } finally {

20
packages/coding-agent/test/compaction.test.ts
View file

@ -14,9 +14,9 @@ import {
shouldCompact, shouldCompact,
} from "../src/core/compaction.js"; } from "../src/core/compaction.js";
import { import {
buildSessionContext,
type CompactionEntry, type CompactionEntry,
createSummaryMessage, createSummaryMessage,
loadSessionFromEntries,
parseSessionEntries, parseSessionEntries,
type SessionEntry, type SessionEntry,
type SessionMessageEntry, type SessionMessageEntry,
@ -226,7 +226,7 @@ describe("createSummaryMessage", () => {
}); });
}); });
describe("loadSessionFromEntries", () => { describe("buildSessionContext", () => {
it("should load all messages when no compaction", () => { it("should load all messages when no compaction", () => {
const entries: SessionEntry[] = [ const entries: SessionEntry[] = [
{ {
@ -241,7 +241,7 @@ describe("loadSessionFromEntries", () => {
createMessageEntry(createAssistantMessage("b")), createMessageEntry(createAssistantMessage("b")),
]; ];
const loaded = loadSessionFromEntries(entries); const loaded = buildSessionContext(entries);
expect(loaded.messages.length).toBe(4); expect(loaded.messages.length).toBe(4);
expect(loaded.thinkingLevel).toBe("off"); expect(loaded.thinkingLevel).toBe("off");
expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" }); expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" });
@ -265,7 +265,7 @@ describe("loadSessionFromEntries", () => {
createMessageEntry(createAssistantMessage("c")), 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 // 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.length).toBe(5);
expect(loaded.messages[0].role).toBe("user"); expect(loaded.messages[0].role).toBe("user");
@ -293,7 +293,7 @@ describe("loadSessionFromEntries", () => {
createMessageEntry(createAssistantMessage("d")), createMessageEntry(createAssistantMessage("d")),
]; ];
const loaded = loadSessionFromEntries(entries); const loaded = buildSessionContext(entries);
// summary + kept from idx 6 (u3,c) + after (u4,d) = 5 // summary + kept from idx 6 (u3,c) + after (u4,d) = 5
expect(loaded.messages.length).toBe(5); expect(loaded.messages.length).toBe(5);
expect((loaded.messages[0] as any).content).toContain("Second summary"); 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 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 // 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 // Actually index 0 is session header, so messages are u1,a1,u2,b
expect(loaded.messages.length).toBe(5); // summary + 4 messages expect(loaded.messages.length).toBe(5); // summary + 4 messages
@ -336,7 +336,7 @@ describe("loadSessionFromEntries", () => {
{ type: "thinking_level_change", timestamp: "", thinkingLevel: "high" }, { 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 // model_change is later overwritten by assistant message's model info
expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" }); expect(loaded.model).toEqual({ provider: "anthropic", modelId: "claude-sonnet-4-5" });
expect(loaded.thinkingLevel).toBe("high"); expect(loaded.thinkingLevel).toBe("high");
@ -368,7 +368,7 @@ describe("Large session fixture", () => {
it("should load session correctly", () => { it("should load session correctly", () => {
const entries = loadLargeSessionEntries(); const entries = loadLargeSessionEntries();
const loaded = loadSessionFromEntries(entries); const loaded = buildSessionContext(entries);
expect(loaded.messages.length).toBeGreaterThan(100); expect(loaded.messages.length).toBeGreaterThan(100);
expect(loaded.model).not.toBeNull(); 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 () => { it("should produce valid session after compaction", async () => {
const entries = loadLargeSessionEntries(); const entries = loadLargeSessionEntries();
const loaded = loadSessionFromEntries(entries); const loaded = buildSessionContext(entries);
const model = getModel("anthropic", "claude-sonnet-4-5")!; const model = getModel("anthropic", "claude-sonnet-4-5")!;
const compactionEvent = await compact( const compactionEvent = await compact(
@ -417,7 +417,7 @@ describe.skipIf(!process.env.ANTHROPIC_OAUTH_TOKEN)("LLM summarization", () => {
// Simulate appending compaction to entries // Simulate appending compaction to entries
const newEntries = [...entries, compactionEvent]; const newEntries = [...entries, compactionEvent];
const reloaded = loadSessionFromEntries(newEntries); const reloaded = buildSessionContext(newEntries);
// Should have summary + kept messages // Should have summary + kept messages
expect(reloaded.messages.length).toBeLessThan(loaded.messages.length); expect(reloaded.messages.length).toBeLessThan(loaded.messages.length);

4
packages/mom/src/context.ts
View file

@ -12,9 +12,9 @@
import type { AppMessage } from "@mariozechner/pi-agent-core"; import type { AppMessage } from "@mariozechner/pi-agent-core";
import { import {
buildSessionContext,
type CompactionEntry, type CompactionEntry,
type LoadedSession, type LoadedSession,
loadSessionFromEntries,
type ModelChangeEntry, type ModelChangeEntry,
type SessionEntry, type SessionEntry,
type SessionMessageEntry, type SessionMessageEntry,
@ -285,7 +285,7 @@ export class MomSessionManager {
/** Load session with compaction support */ /** Load session with compaction support */
loadSession(): LoadedSession { loadSession(): LoadedSession {
const entries = this.loadEntries(); const entries = this.loadEntries();
return loadSessionFromEntries(entries); return buildSessionContext(entries);
} }
loadEntries(): SessionEntry[] { loadEntries(): SessionEntry[] {