mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-21 02:04:32 +00:00
WP7: Add AgentSession compaction (manual + auto), fix listener preservation
This commit is contained in:
parent
0119d7610b
commit
8d6d2dd72b
2 changed files with 177 additions and 24 deletions
|
|
@ -341,8 +341,9 @@ private unsubscribeAll(): void {
|
||||||
1. `npm run check` passes
|
1. `npm run check` passes
|
||||||
|
|
||||||
- [x] Add `subscribe()` method to AgentSession
|
- [x] Add `subscribe()` method to AgentSession
|
||||||
- [x] Add `unsubscribeAll()` method
|
- [x] Add `_disconnectFromAgent()` private method (renamed from unsubscribeAll)
|
||||||
- [x] Add `resubscribe()` method
|
- [x] Add `_reconnectToAgent()` private method (renamed from resubscribe)
|
||||||
|
- [x] Add `dispose()` public method for full cleanup
|
||||||
- [x] Verify with `npm run check`
|
- [x] Verify with `npm run check`
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
@ -808,12 +809,12 @@ get autoCompactionEnabled(): boolean {
|
||||||
**Verification:**
|
**Verification:**
|
||||||
1. `npm run check` passes
|
1. `npm run check` passes
|
||||||
|
|
||||||
- [ ] Add `CompactionResult` interface
|
- [x] Add `CompactionResult` interface
|
||||||
- [ ] Add `compact()` method
|
- [x] Add `compact()` method
|
||||||
- [ ] Add `abortCompaction()` method
|
- [x] Add `abortCompaction()` method
|
||||||
- [ ] Add `checkAutoCompaction()` method
|
- [x] Add `checkAutoCompaction()` method
|
||||||
- [ ] Add `setAutoCompactionEnabled()` and getter
|
- [x] Add `setAutoCompactionEnabled()` and getter
|
||||||
- [ ] Verify with `npm run check`
|
- [x] Verify with `npm run check`
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -14,10 +14,11 @@
|
||||||
*/
|
*/
|
||||||
|
|
||||||
import type { Agent, AgentEvent, AgentState, AppMessage, Attachment, ThinkingLevel } from "@mariozechner/pi-agent-core";
|
import type { Agent, AgentEvent, AgentState, AppMessage, Attachment, ThinkingLevel } from "@mariozechner/pi-agent-core";
|
||||||
import type { Model } from "@mariozechner/pi-ai";
|
import type { AssistantMessage, Model } from "@mariozechner/pi-ai";
|
||||||
|
import { calculateContextTokens, compact, shouldCompact } from "../compaction.js";
|
||||||
import { getModelsPath } from "../config.js";
|
import { getModelsPath } from "../config.js";
|
||||||
import { getApiKeyForModel, getAvailableModels } from "../model-config.js";
|
import { getApiKeyForModel, getAvailableModels } from "../model-config.js";
|
||||||
import type { SessionManager } from "../session-manager.js";
|
import { loadSessionFromEntries, type SessionManager } from "../session-manager.js";
|
||||||
import type { SettingsManager } from "../settings-manager.js";
|
import type { SettingsManager } from "../settings-manager.js";
|
||||||
import { expandSlashCommand, type FileSlashCommand } from "../slash-commands.js";
|
import { expandSlashCommand, type FileSlashCommand } from "../slash-commands.js";
|
||||||
|
|
||||||
|
|
@ -54,6 +55,12 @@ export interface ModelCycleResult {
|
||||||
isScoped: boolean;
|
isScoped: boolean;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/** Result from compact() or checkAutoCompaction() */
|
||||||
|
export interface CompactionResult {
|
||||||
|
tokensBefore: number;
|
||||||
|
summary: string;
|
||||||
|
}
|
||||||
|
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
// AgentSession Class
|
// AgentSession Class
|
||||||
// ============================================================================
|
// ============================================================================
|
||||||
|
|
@ -73,6 +80,9 @@ export class AgentSession {
|
||||||
// Message queue state
|
// Message queue state
|
||||||
private _queuedMessages: string[] = [];
|
private _queuedMessages: string[] = [];
|
||||||
|
|
||||||
|
// Compaction state
|
||||||
|
private _compactionAbortController: AbortController | null = null;
|
||||||
|
|
||||||
constructor(config: AgentSessionConfig) {
|
constructor(config: AgentSessionConfig) {
|
||||||
this.agent = config.agent;
|
this.agent = config.agent;
|
||||||
this.sessionManager = config.sessionManager;
|
this.sessionManager = config.sessionManager;
|
||||||
|
|
@ -111,10 +121,9 @@ export class AgentSession {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check auto-compaction after assistant messages
|
// Check auto-compaction after assistant messages
|
||||||
// (will be implemented in WP7)
|
if (event.message.role === "assistant") {
|
||||||
// if (event.message.role === "assistant") {
|
await this.checkAutoCompaction();
|
||||||
// await this.checkAutoCompaction();
|
}
|
||||||
// }
|
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
@ -129,23 +138,23 @@ export class AgentSession {
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Unsubscribe from agent entirely and clear all listeners.
|
* Temporarily disconnect from agent events.
|
||||||
* Used during reset/cleanup operations.
|
* User listeners are preserved and will receive events again after resubscribe().
|
||||||
|
* Used internally during operations that need to pause event processing.
|
||||||
*/
|
*/
|
||||||
unsubscribeAll(): void {
|
private _disconnectFromAgent(): void {
|
||||||
if (this._unsubscribeAgent) {
|
if (this._unsubscribeAgent) {
|
||||||
this._unsubscribeAgent();
|
this._unsubscribeAgent();
|
||||||
this._unsubscribeAgent = undefined;
|
this._unsubscribeAgent = undefined;
|
||||||
}
|
}
|
||||||
this._eventListeners = [];
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Re-subscribe to agent after unsubscribeAll.
|
* Reconnect to agent events after _disconnectFromAgent().
|
||||||
* Call this after operations that require temporary unsubscription.
|
* Preserves all existing listeners.
|
||||||
*/
|
*/
|
||||||
resubscribe(): void {
|
private _reconnectToAgent(): void {
|
||||||
if (this._unsubscribeAgent) return; // Already subscribed
|
if (this._unsubscribeAgent) return; // Already connected
|
||||||
|
|
||||||
this._unsubscribeAgent = this.agent.subscribe(async (event) => {
|
this._unsubscribeAgent = this.agent.subscribe(async (event) => {
|
||||||
for (const l of this._eventListeners) {
|
for (const l of this._eventListeners) {
|
||||||
|
|
@ -158,10 +167,24 @@ export class AgentSession {
|
||||||
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
|
if (this.sessionManager.shouldInitializeSession(this.agent.state.messages)) {
|
||||||
this.sessionManager.startSession(this.agent.state);
|
this.sessionManager.startSession(this.agent.state);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Check auto-compaction after assistant messages
|
||||||
|
if (event.message.role === "assistant") {
|
||||||
|
await this.checkAutoCompaction();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Remove all listeners and disconnect from agent.
|
||||||
|
* Call this when completely done with the session.
|
||||||
|
*/
|
||||||
|
dispose(): void {
|
||||||
|
this._disconnectFromAgent();
|
||||||
|
this._eventListeners = [];
|
||||||
|
}
|
||||||
|
|
||||||
// =========================================================================
|
// =========================================================================
|
||||||
// Read-only State Access
|
// Read-only State Access
|
||||||
// =========================================================================
|
// =========================================================================
|
||||||
|
|
@ -299,14 +322,15 @@ export class AgentSession {
|
||||||
/**
|
/**
|
||||||
* Reset agent and session to start fresh.
|
* Reset agent and session to start fresh.
|
||||||
* Clears all messages and starts a new session.
|
* Clears all messages and starts a new session.
|
||||||
|
* Listeners are preserved and will continue receiving events.
|
||||||
*/
|
*/
|
||||||
async reset(): Promise<void> {
|
async reset(): Promise<void> {
|
||||||
this.unsubscribeAll();
|
this._disconnectFromAgent();
|
||||||
await this.abort();
|
await this.abort();
|
||||||
this.agent.reset();
|
this.agent.reset();
|
||||||
this.sessionManager.reset();
|
this.sessionManager.reset();
|
||||||
this._queuedMessages = [];
|
this._queuedMessages = [];
|
||||||
// Note: caller should re-subscribe after reset if needed
|
this._reconnectToAgent();
|
||||||
}
|
}
|
||||||
|
|
||||||
// =========================================================================
|
// =========================================================================
|
||||||
|
|
@ -464,4 +488,132 @@ export class AgentSession {
|
||||||
this.agent.setQueueMode(mode);
|
this.agent.setQueueMode(mode);
|
||||||
this.settingsManager.setQueueMode(mode);
|
this.settingsManager.setQueueMode(mode);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// =========================================================================
|
||||||
|
// Compaction
|
||||||
|
// =========================================================================
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Manually compact the session context.
|
||||||
|
* Aborts current agent operation first.
|
||||||
|
* @param customInstructions Optional instructions for the compaction summary
|
||||||
|
*/
|
||||||
|
async compact(customInstructions?: string): Promise<CompactionResult> {
|
||||||
|
// Abort any running operation
|
||||||
|
this._disconnectFromAgent();
|
||||||
|
await this.abort();
|
||||||
|
|
||||||
|
// Create abort controller
|
||||||
|
this._compactionAbortController = new AbortController();
|
||||||
|
|
||||||
|
try {
|
||||||
|
if (!this.model) {
|
||||||
|
throw new Error("No model selected");
|
||||||
|
}
|
||||||
|
|
||||||
|
const apiKey = await getApiKeyForModel(this.model);
|
||||||
|
if (!apiKey) {
|
||||||
|
throw new Error(`No API key for ${this.model.provider}`);
|
||||||
|
}
|
||||||
|
|
||||||
|
const entries = this.sessionManager.loadEntries();
|
||||||
|
const settings = this.settingsManager.getCompactionSettings();
|
||||||
|
const compactionEntry = await compact(
|
||||||
|
entries,
|
||||||
|
this.model,
|
||||||
|
settings,
|
||||||
|
apiKey,
|
||||||
|
this._compactionAbortController.signal,
|
||||||
|
customInstructions,
|
||||||
|
);
|
||||||
|
|
||||||
|
if (this._compactionAbortController.signal.aborted) {
|
||||||
|
throw new Error("Compaction cancelled");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Save and reload
|
||||||
|
this.sessionManager.saveCompaction(compactionEntry);
|
||||||
|
const loaded = loadSessionFromEntries(this.sessionManager.loadEntries());
|
||||||
|
this.agent.replaceMessages(loaded.messages);
|
||||||
|
|
||||||
|
return {
|
||||||
|
tokensBefore: compactionEntry.tokensBefore,
|
||||||
|
summary: compactionEntry.summary,
|
||||||
|
};
|
||||||
|
} finally {
|
||||||
|
this._compactionAbortController = null;
|
||||||
|
this._reconnectToAgent();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Cancel in-progress compaction.
|
||||||
|
*/
|
||||||
|
abortCompaction(): void {
|
||||||
|
this._compactionAbortController?.abort();
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Check if auto-compaction should run, and run it if so.
|
||||||
|
* Called internally after assistant messages.
|
||||||
|
* @returns Result if compaction occurred, null otherwise
|
||||||
|
*/
|
||||||
|
async checkAutoCompaction(): Promise<CompactionResult | null> {
|
||||||
|
const settings = this.settingsManager.getCompactionSettings();
|
||||||
|
if (!settings.enabled) return null;
|
||||||
|
|
||||||
|
// Get last non-aborted assistant message
|
||||||
|
const messages = this.messages;
|
||||||
|
let lastAssistant: AssistantMessage | null = null;
|
||||||
|
for (let i = messages.length - 1; i >= 0; i--) {
|
||||||
|
const msg = messages[i];
|
||||||
|
if (msg.role === "assistant") {
|
||||||
|
const assistantMsg = msg as AssistantMessage;
|
||||||
|
if (assistantMsg.stopReason !== "aborted") {
|
||||||
|
lastAssistant = assistantMsg;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if (!lastAssistant) return null;
|
||||||
|
|
||||||
|
const contextTokens = calculateContextTokens(lastAssistant.usage);
|
||||||
|
const contextWindow = this.model?.contextWindow ?? 0;
|
||||||
|
|
||||||
|
if (!shouldCompact(contextTokens, contextWindow, settings)) return null;
|
||||||
|
|
||||||
|
// Perform auto-compaction (don't abort current operation for auto)
|
||||||
|
try {
|
||||||
|
if (!this.model) return null;
|
||||||
|
|
||||||
|
const apiKey = await getApiKeyForModel(this.model);
|
||||||
|
if (!apiKey) return null;
|
||||||
|
|
||||||
|
const entries = this.sessionManager.loadEntries();
|
||||||
|
const compactionEntry = await compact(entries, this.model, settings, apiKey);
|
||||||
|
|
||||||
|
this.sessionManager.saveCompaction(compactionEntry);
|
||||||
|
const loaded = loadSessionFromEntries(this.sessionManager.loadEntries());
|
||||||
|
this.agent.replaceMessages(loaded.messages);
|
||||||
|
|
||||||
|
return {
|
||||||
|
tokensBefore: compactionEntry.tokensBefore,
|
||||||
|
summary: compactionEntry.summary,
|
||||||
|
};
|
||||||
|
} catch {
|
||||||
|
return null; // Silently fail auto-compaction
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Toggle auto-compaction setting.
|
||||||
|
*/
|
||||||
|
setAutoCompactionEnabled(enabled: boolean): void {
|
||||||
|
this.settingsManager.setCompactionEnabled(enabled);
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Whether auto-compaction is enabled */
|
||||||
|
get autoCompactionEnabled(): boolean {
|
||||||
|
return this.settingsManager.getCompactionEnabled();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue