flynn/src/session/manager.ts

import type { Message } from '../models/types.js';
import type { SessionStore } from './store.js';
import { auditLogger } from '../audit/index.js';
import { SessionIndexer } from './indexer.js';
import { SessionSearch, type HistorySearchResult } from './search.js';

export interface Session {
  id: string;
  addMessage(message: Message): void;
  getHistory(): Message[];
  clear(): void;
  replaceHistory(messages: Message[]): void;
  getConfig(key: string): string | undefined;
  setConfig(key: string, value: string): void;
  deleteConfig(key: string): void;
}

export class ManagedSession implements Session {
  constructor(
    public readonly id: string,
    private store: SessionStore,
    private history: Message[] = [],
    private indexer?: SessionIndexer,
  ) {}

  addMessage(message: Message): Message {
    const messageWithTimestamp: Message = {
      ...message,
      timestamp: Date.now(),
    };
    this.history.push(messageWithTimestamp);
    const content = typeof message.content === 'string'
      ? message.content
      : JSON.stringify(message.content);
    const metadata = this.indexer?.indexText(content);
    this.store.addMessage(this.id, messageWithTimestamp, metadata);

    auditLogger?.sessionMessage({
      session_id: this.id,
      role: message.role,
      content_length: typeof message.content === 'string'
        ? message.content.length
        : JSON.stringify(message.content).length,
    });

    return messageWithTimestamp;
  }

  getHistory(): Message[] {
    return [...this.history];
  }

  clear(): void {
    const messageCount = this.history.length;
    this.history = [];
    this.store.clearSession(this.id);

    auditLogger?.sessionDelete({
      session_id: this.id,
      message_count: messageCount,
    });
  }

  /**
   * Replace the entire session history with new messages.
   * Used after compaction to persist the compacted state.
   * Updates both in-memory history and SQLite storage atomically.
   */
  replaceHistory(messages: Message[]): void {
    this.history = [...messages];
    this.store.replaceMessages(this.id, messages);
  }

  setHistory(messages: Message[]): void {
    this.history = [...messages];
  }

  getConfig(key: string): string | undefined {
    return this.store.getSessionConfig(this.id, key);
  }

  setConfig(key: string, value: string): void {
    this.store.setSessionConfig(this.id, key, value);
  }

  deleteConfig(key: string): void {
    this.store.deleteSessionConfig(this.id, key);
  }
}

export class SessionManager {
  private sessions: Map<string, ManagedSession> = new Map();
  private indexer?: SessionIndexer;
  private search?: SessionSearch;

  constructor(private store: SessionStore, historyIndexConfig?: {
    enabled: boolean;
    maxKeywords: number;
    searchLimit: number;
    minScore: number;
  }) {
    if (historyIndexConfig?.enabled) {
      this.indexer = new SessionIndexer({
        maxKeywords: historyIndexConfig.maxKeywords,
      });
      this.search = new SessionSearch(store, {
        limit: historyIndexConfig.searchLimit,
        minScore: historyIndexConfig.minScore,
      });
    }
  }

  private makeSessionId(frontend: string, userId: string): string {
    return `${frontend}:${userId}`;
  }

  getSession(frontend: string, userId: string): ManagedSession {
    const id = this.makeSessionId(frontend, userId);

    let session = this.sessions.get(id);
    if (!session) {
      const history = this.store.getMessages(id);
      session = new ManagedSession(id, this.store, history, this.indexer);
      this.sessions.set(id, session);

      auditLogger?.sessionCreate({
        session_id: id,
        frontend,
        user_id: userId,
      });
    }

    return session;
  }

  transferSession(
    fromFrontend: string,
    fromUserId: string,
    toFrontend: string,
    toUserId: string,
  ): void {
    const fromSession = this.getSession(fromFrontend, fromUserId);
    const toSession = this.getSession(toFrontend, toUserId);

    const history = fromSession.getHistory();

    // Clear target and copy history
    toSession.clear();
    for (const message of history) {
      toSession.addMessage(message);
    }

    auditLogger?.sessionTransfer(fromSession.id, toSession.id, history.length);
  }

  listSessions(): string[] {
    return Array.from(this.sessions.keys());
  }

  closeSession(frontend: string, userId: string): void {
    const id = this.makeSessionId(frontend, userId);
    this.sessions.delete(id);
  }

  /** Remove sessions from the in-memory cache by their IDs. */
  evictSessions(sessionIds: string[]): void {
    for (const id of sessionIds) {
      this.sessions.delete(id);
    }
  }

  /** Get a session config value. */
  getSessionConfig(frontend: string, userId: string, key: string): string | undefined {
    const session = this.getSession(frontend, userId);
    return session.getConfig(key);
  }

  /** Set a session config value. */
  setSessionConfig(frontend: string, userId: string, key: string, value: string): void {
    const session = this.getSession(frontend, userId);
    session.setConfig(key, value);
  }

  /** Delete a session config value. */
  deleteSessionConfig(frontend: string, userId: string, key: string): void {
    const session = this.getSession(frontend, userId);
    session.deleteConfig(key);
  }

  searchHistory(query: string, opts?: { limit?: number; sessionId?: string }): HistorySearchResult[] {
    if (!this.search) {
      return [];
    }
    return this.search.search(query, opts);
  }

  reindexHistory(): number {
    if (!this.indexer) {
      return 0;
    }

    const rows = this.store.getAllMessagesWithMetadata();
    for (const row of rows) {
      const metadata = this.indexer.indexText(row.content);
      this.store.updateMessageMetadata(row.id, metadata);
    }
    return rows.length;
  }
}