From fb7575f850c6e4ad365bee7200f1e343dce133e0 Mon Sep 17 00:00:00 2001 From: William Valentin Date: Thu, 5 Feb 2026 00:43:09 -0800 Subject: [PATCH] refactor: integrate SessionManager into daemon and agent Co-Authored-By: Claude Opus 4.5 --- src/backends/native/agent.test.ts | 100 ++++++++++++++---------------- src/backends/native/agent.ts | 43 +++++++------ src/daemon/index.ts | 21 ++++--- 3 files changed, 78 insertions(+), 86 deletions(-) diff --git a/src/backends/native/agent.test.ts b/src/backends/native/agent.test.ts index 59a7c68..d17da0b 100644 --- a/src/backends/native/agent.test.ts +++ b/src/backends/native/agent.test.ts @@ -1,79 +1,69 @@ -import { describe, it, expect, vi } from 'vitest'; +import { describe, it, expect, vi, beforeEach } from 'vitest'; import { NativeAgent } from './agent.js'; import type { ModelClient, ChatResponse } from '../../models/types.js'; describe('NativeAgent', () => { - it('processes message and returns response', async () => { - const mockClient: ModelClient = { - chat: vi.fn().mockResolvedValue({ - content: 'Hello! How can I help you?', - stopReason: 'end_turn', - usage: { inputTokens: 10, outputTokens: 8 }, - } satisfies ChatResponse), - }; - - const agent = new NativeAgent({ - modelClient: mockClient, - systemPrompt: 'You are Flynn, a helpful assistant.', - }); - - const response = await agent.process('Hello'); - - expect(response).toBe('Hello! How can I help you?'); - expect(mockClient.chat).toHaveBeenCalledWith({ - messages: [{ role: 'user', content: 'Hello' }], - system: 'You are Flynn, a helpful assistant.', - }); + const createMockClient = (): ModelClient => ({ + chat: vi.fn().mockResolvedValue({ + content: 'Hello!', + stopReason: 'end_turn', + usage: { inputTokens: 10, outputTokens: 5 }, + } satisfies ChatResponse), }); - it('maintains conversation history', async () => { - const mockClient: ModelClient = { - chat: vi.fn().mockResolvedValue({ - content: 'Response', - stopReason: 'end_turn', - usage: { inputTokens: 10, outputTokens: 5 }, - } satisfies ChatResponse), - }; - + it('processes messages and maintains history', async () => { + const mockClient = createMockClient(); const agent = new NativeAgent({ modelClient: mockClient, - systemPrompt: 'System', + systemPrompt: 'You are helpful.', }); - await agent.process('First message'); - await agent.process('Second message'); + const response = await agent.process('Hi'); - expect(mockClient.chat).toHaveBeenLastCalledWith({ - messages: [ - { role: 'user', content: 'First message' }, - { role: 'assistant', content: 'Response' }, - { role: 'user', content: 'Second message' }, - ], - system: 'System', + expect(response).toBe('Hello!'); + expect(mockClient.chat).toHaveBeenCalledWith({ + messages: [{ role: 'user', content: 'Hi' }], + system: 'You are helpful.', }); + + const history = agent.getHistory(); + expect(history).toHaveLength(2); + expect(history[0]).toEqual({ role: 'user', content: 'Hi' }); + expect(history[1]).toEqual({ role: 'assistant', content: 'Hello!' }); }); it('resets conversation history', async () => { - const mockClient: ModelClient = { - chat: vi.fn().mockResolvedValue({ - content: 'Response', - stopReason: 'end_turn', - usage: { inputTokens: 10, outputTokens: 5 }, - } satisfies ChatResponse), + const mockClient = createMockClient(); + const agent = new NativeAgent({ + modelClient: mockClient, + systemPrompt: 'You are helpful.', + }); + + await agent.process('Hi'); + agent.reset(); + + expect(agent.getHistory()).toHaveLength(0); + }); + + it('uses session when provided', async () => { + const mockClient = createMockClient(); + const mockSession = { + id: 'test-session', + getHistory: vi.fn().mockReturnValue([]), + addMessage: vi.fn(), + clear: vi.fn(), }; const agent = new NativeAgent({ modelClient: mockClient, - systemPrompt: 'System', + systemPrompt: 'You are helpful.', + session: mockSession, }); - await agent.process('Message 1'); - agent.reset(); - await agent.process('Message 2'); + await agent.process('Hi'); - expect(mockClient.chat).toHaveBeenLastCalledWith({ - messages: [{ role: 'user', content: 'Message 2' }], - system: 'System', - }); + expect(mockSession.addMessage).toHaveBeenCalledTimes(2); + expect(mockSession.addMessage).toHaveBeenNthCalledWith(1, { role: 'user', content: 'Hi' }); + expect(mockSession.addMessage).toHaveBeenNthCalledWith(2, { role: 'assistant', content: 'Hello!' }); }); }); diff --git a/src/backends/native/agent.ts b/src/backends/native/agent.ts index 8630b54..a724b95 100644 --- a/src/backends/native/agent.ts +++ b/src/backends/native/agent.ts @@ -1,59 +1,58 @@ import type { ModelClient, Message } from '../../models/types.js'; -import type { SessionStore } from '../../session/index.js'; +import type { Session } from '../../session/index.js'; export interface NativeAgentConfig { modelClient: ModelClient; systemPrompt: string; - sessionStore?: SessionStore; - sessionId?: string; + session?: Session; } export class NativeAgent { private modelClient: ModelClient; private systemPrompt: string; - private sessionStore?: SessionStore; - private sessionId: string; - private history: Message[] = []; + private session?: Session; + private inMemoryHistory: Message[] = []; constructor(config: NativeAgentConfig) { this.modelClient = config.modelClient; this.systemPrompt = config.systemPrompt; - this.sessionStore = config.sessionStore; - this.sessionId = config.sessionId ?? 'default'; + this.session = config.session; + } - // Load existing history from store - if (this.sessionStore) { - this.history = this.sessionStore.getMessages(this.sessionId); - } + private get history(): Message[] { + return this.session?.getHistory() ?? [...this.inMemoryHistory]; } async process(userMessage: string): Promise { const userMsg: Message = { role: 'user', content: userMessage }; - this.history.push(userMsg); - if (this.sessionStore) { - this.sessionStore.addMessage(this.sessionId, userMsg); + if (this.session) { + this.session.addMessage(userMsg); + } else { + this.inMemoryHistory.push(userMsg); } const response = await this.modelClient.chat({ - messages: [...this.history], + messages: this.history, system: this.systemPrompt, }); const assistantMsg: Message = { role: 'assistant', content: response.content }; - this.history.push(assistantMsg); - if (this.sessionStore) { - this.sessionStore.addMessage(this.sessionId, assistantMsg); + if (this.session) { + this.session.addMessage(assistantMsg); + } else { + this.inMemoryHistory.push(assistantMsg); } return response.content; } reset(): void { - this.history = []; - if (this.sessionStore) { - this.sessionStore.clearSession(this.sessionId); + if (this.session) { + this.session.clear(); + } else { + this.inMemoryHistory = []; } } diff --git a/src/daemon/index.ts b/src/daemon/index.ts index b9dc44d..cb1e484 100644 --- a/src/daemon/index.ts +++ b/src/daemon/index.ts @@ -4,7 +4,7 @@ import type { Config } from '../config/index.js'; import { AnthropicClient, OpenAIClient, OllamaClient, ModelRouter } from '../models/index.js'; import { NativeAgent } from '../backends/index.js'; import { createTelegramBot } from '../frontends/telegram/index.js'; -import { SessionStore } from '../session/index.js'; +import { SessionStore, SessionManager } from '../session/index.js'; import { HookEngine } from '../hooks/index.js'; import { resolve } from 'path'; import { homedir } from 'os'; @@ -16,6 +16,7 @@ export interface DaemonContext { bot: Bot; agent: NativeAgent; sessionStore: SessionStore; + sessionManager: SessionManager; hookEngine: HookEngine; modelRouter: ModelRouter; } @@ -27,12 +28,10 @@ Keep responses focused and avoid unnecessary verbosity. Use markdown formatting function createModelRouter(config: Config): ModelRouter { const models = config.models; - // Create default client (required) const defaultClient = new AnthropicClient({ model: models.default.model, }); - // Create optional tier clients let fastClient; let complexClient; let localClient; @@ -54,7 +53,6 @@ function createModelRouter(config: Config): ModelRouter { } } - // Build fallback chain const fallbackChain = []; for (const providerName of models.fallback_chain) { if (providerName === 'openai') { @@ -80,8 +78,10 @@ export async function startDaemon(config: Config): Promise { const dataDir = resolve(homedir(), '.local/share/flynn'); mkdirSync(dataDir, { recursive: true }); - // Initialize session store + // Initialize session store and manager const sessionStore = new SessionStore(resolve(dataDir, 'sessions.db')); + const sessionManager = new SessionManager(sessionStore); + lifecycle.onShutdown(async () => { sessionStore.close(); console.log('Session store closed'); @@ -93,12 +93,15 @@ export async function startDaemon(config: Config): Promise { // Initialize model router const modelRouter = createModelRouter(config); - // Initialize native agent with session persistence + // Get Telegram session + const telegramUserId = String(config.telegram.allowed_chat_ids[0]); + const session = sessionManager.getSession('telegram', telegramUserId); + + // Initialize native agent with session const agent = new NativeAgent({ modelClient: modelRouter, systemPrompt: SYSTEM_PROMPT, - sessionStore, - sessionId: `telegram-${config.telegram.allowed_chat_ids[0]}`, + session, }); // Initialize Telegram bot with hook engine @@ -136,7 +139,7 @@ export async function startDaemon(config: Config): Promise { console.log('Flynn daemon started'); - return { config, lifecycle, bot, agent, sessionStore, hookEngine, modelRouter }; + return { config, lifecycle, bot, agent, sessionStore, sessionManager, hookEngine, modelRouter }; } export { Lifecycle } from './lifecycle.js';