feat(ipuaro): add session management (v0.10.0)

- Add ISessionStorage interface and RedisSessionStorage implementation - Add ContextManager for token budget and compression - Add StartSession, HandleMessage, UndoChange use cases - Update CHANGELOG and TODO documentation - 88 new tests (1174 total), 97.73% coverage
2025-12-28 07:16:53 +05:00 · 2025-12-01 12:27:22 +05:00
parent 56643d903f
commit 0f2ed5b301
22 changed files with 2798 additions and 261 deletions
--- a/packages/ipuaro/src/application/use-cases/ContextManager.ts
+++ b/packages/ipuaro/src/application/use-cases/ContextManager.ts
@@ -0,0 +1,229 @@
+import type { ContextState, Session } from "../../domain/entities/Session.js"
+import type { ILLMClient } from "../../domain/services/ILLMClient.js"
+import { type ChatMessage, createSystemMessage } from "../../domain/value-objects/ChatMessage.js"
+import { CONTEXT_COMPRESSION_THRESHOLD, CONTEXT_WINDOW_SIZE } from "../../domain/constants/index.js"
+
+/**
+ * File in context with token count.
+ */
+export interface FileContext {
+    path: string
+    tokens: number
+    addedAt: number
+}
+
+/**
+ * Compression result.
+ */
+export interface CompressionResult {
+    compressed: boolean
+    removedMessages: number
+    tokensSaved: number
+    summary?: string
+}
+
+const COMPRESSION_PROMPT = `Summarize the following conversation history in a concise way,
+preserving key information about:
+- What files were discussed or modified
+- What changes were made
+- Important decisions or context
+Keep the summary under 500 tokens.`
+
+const MESSAGES_TO_KEEP = 5
+const MIN_MESSAGES_FOR_COMPRESSION = 10
+
+/**
+ * Manages context window token budget and compression.
+ */
+export class ContextManager {
+    private readonly filesInContext = new Map<string, FileContext>()
+    private currentTokens = 0
+    private readonly contextWindowSize: number
+
+    constructor(contextWindowSize: number = CONTEXT_WINDOW_SIZE) {
+        this.contextWindowSize = contextWindowSize
+    }
+
+    /**
+     * Add a file to the context.
+     */
+    addToContext(file: string, tokens: number): void {
+        const existing = this.filesInContext.get(file)
+        if (existing) {
+            this.currentTokens -= existing.tokens
+        }
+
+        this.filesInContext.set(file, {
+            path: file,
+            tokens,
+            addedAt: Date.now(),
+        })
+        this.currentTokens += tokens
+    }
+
+    /**
+     * Remove a file from the context.
+     */
+    removeFromContext(file: string): void {
+        const existing = this.filesInContext.get(file)
+        if (existing) {
+            this.currentTokens -= existing.tokens
+            this.filesInContext.delete(file)
+        }
+    }
+
+    /**
+     * Get current token usage ratio (0-1).
+     */
+    getUsage(): number {
+        return this.currentTokens / this.contextWindowSize
+    }
+
+    /**
+     * Get current token count.
+     */
+    getTokenCount(): number {
+        return this.currentTokens
+    }
+
+    /**
+     * Get available tokens.
+     */
+    getAvailableTokens(): number {
+        return this.contextWindowSize - this.currentTokens
+    }
+
+    /**
+     * Check if compression is needed.
+     */
+    needsCompression(): boolean {
+        return this.getUsage() > CONTEXT_COMPRESSION_THRESHOLD
+    }
+
+    /**
+     * Update token count (e.g., after receiving a message).
+     */
+    addTokens(tokens: number): void {
+        this.currentTokens += tokens
+    }
+
+    /**
+     * Get files in context.
+     */
+    getFilesInContext(): string[] {
+        return Array.from(this.filesInContext.keys())
+    }
+
+    /**
+     * Sync context state from session.
+     */
+    syncFromSession(session: Session): void {
+        this.filesInContext.clear()
+        this.currentTokens = 0
+
+        for (const file of session.context.filesInContext) {
+            this.filesInContext.set(file, {
+                path: file,
+                tokens: 0,
+                addedAt: Date.now(),
+            })
+        }
+
+        this.currentTokens = Math.floor(session.context.tokenUsage * this.contextWindowSize)
+    }
+
+    /**
+     * Update session context state.
+     */
+    updateSession(session: Session): void {
+        session.context.filesInContext = this.getFilesInContext()
+        session.context.tokenUsage = this.getUsage()
+        session.context.needsCompression = this.needsCompression()
+    }
+
+    /**
+     * Compress context using LLM to summarize old messages.
+     */
+    async compress(session: Session, llm: ILLMClient): Promise<CompressionResult> {
+        const history = session.history
+        if (history.length < MIN_MESSAGES_FOR_COMPRESSION) {
+            return {
+                compressed: false,
+                removedMessages: 0,
+                tokensSaved: 0,
+            }
+        }
+
+        const messagesToCompress = history.slice(0, -MESSAGES_TO_KEEP)
+        const messagesToKeep = history.slice(-MESSAGES_TO_KEEP)
+
+        const tokensBeforeCompression = await this.countHistoryTokens(messagesToCompress, llm)
+
+        const summary = await this.summarizeMessages(messagesToCompress, llm)
+        const summaryTokens = await llm.countTokens(summary)
+
+        const summaryMessage = createSystemMessage(`[Previous conversation summary]\n${summary}`)
+
+        session.history = [summaryMessage, ...messagesToKeep]
+
+        const tokensSaved = tokensBeforeCompression - summaryTokens
+        this.currentTokens -= tokensSaved
+
+        this.updateSession(session)
+
+        return {
+            compressed: true,
+            removedMessages: messagesToCompress.length,
+            tokensSaved,
+            summary,
+        }
+    }
+
+    /**
+     * Create a new context state.
+     */
+    static createInitialState(): ContextState {
+        return {
+            filesInContext: [],
+            tokenUsage: 0,
+            needsCompression: false,
+        }
+    }
+
+    private async summarizeMessages(messages: ChatMessage[], llm: ILLMClient): Promise<string> {
+        const conversation = this.formatMessagesForSummary(messages)
+
+        const response = await llm.chat([
+            createSystemMessage(COMPRESSION_PROMPT),
+            createSystemMessage(conversation),
+        ])
+
+        return response.content
+    }
+
+    private formatMessagesForSummary(messages: ChatMessage[]): string {
+        return messages
+            .filter((m) => m.role !== "tool")
+            .map((m) => {
+                const role = m.role === "user" ? "User" : "Assistant"
+                const content = this.truncateContent(m.content, 500)
+                return `${role}: ${content}`
+            })
+            .join("\n\n")
+    }
+
+    private truncateContent(content: string, maxLength: number): string {
+        if (content.length <= maxLength) {
+            return content
+        }
+        return `${content.slice(0, maxLength)}...`
+    }
+
+    private async countHistoryTokens(messages: ChatMessage[], llm: ILLMClient): Promise<number> {
+        let total = 0
+        for (const message of messages) {
+            total += await llm.countTokens(message.content)
+        }
+        return total
+    }
+}