flynn/src/gateway/session-bridge.ts

import { randomUUID } from 'crypto';
import type { SessionManager } from '../session/manager.js';
import type { Session } from '../session/manager.js';
import type { ModelClient } from '../models/types.js';
import type { ModelRouter, ModelTier } from '../models/router.js';
import type { ToolRegistry } from '../tools/registry.js';
import type { ToolExecutor } from '../tools/executor.js';
import { NativeAgent } from '../backends/native/agent.js';
import type { ToolUseEvent } from '../backends/native/agent.js';

export interface SessionBridgeConfig {
  sessionManager: SessionManager;
  modelClient: ModelClient | ModelRouter;
  systemPrompt: string;
  toolRegistry: ToolRegistry;
  toolExecutor: ToolExecutor;
}

interface ClientEntry {
  connectionId: string;
  sessionId: string;
  agent: NativeAgent;
  busy: boolean;
}

export class SessionBridge {
  private clients: Map<string, ClientEntry> = new Map();
  private agents: Map<string, NativeAgent> = new Map();
  private config: SessionBridgeConfig;
  /** Tracks the current model tier so new agents inherit it and existing agents stay in sync. */
  private currentTier: ModelTier = 'default';

  constructor(config: SessionBridgeConfig) {
    this.config = config;

    // If the model client is a ModelRouter, subscribe to tier changes
    // so all WebChat agents stay in sync with TUI model switches.
    if ('getClient' in config.modelClient) {
      const router = config.modelClient as ModelRouter;
      this.currentTier = router.getTier();
      router.addOnTierChange((tier: ModelTier) => this.onTierChanged(tier));
    }
  }

  /** Called when the ModelRouter's active tier changes. Updates all existing agents. */
  private onTierChanged(tier: ModelTier): void {
    this.currentTier = tier;
    for (const agent of this.agents.values()) {
      agent.setModelTier(tier);
    }
  }

  /** Register a new WS connection. Returns the assigned connection ID. */
  connect(connectionId?: string): string {
    const id = connectionId ?? randomUUID();
    const sessionId = `ws:${id}`;
    const agent = this.getOrCreateAgent(sessionId);

    this.clients.set(id, {
      connectionId: id,
      sessionId,
      agent,
      busy: false,
    });

    return id;
  }

  /** Remove a WS connection. Does NOT destroy the session (persists in SQLite). */
  disconnect(connectionId: string): void {
    const client = this.clients.get(connectionId);
    if (client) {
      // Only remove the agent if no other clients share the session
      const otherClients = Array.from(this.clients.values())
        .filter(c => c.sessionId === client.sessionId && c.connectionId !== connectionId);
      if (otherClients.length === 0) {
        this.agents.delete(client.sessionId);
      }
      this.clients.delete(connectionId);
    }
  }

  /** Switch a connection to a different session (e.g. resuming an old session). */
  switchSession(connectionId: string, sessionId: string): void {
    const client = this.clients.get(connectionId);
    if (!client) throw new Error(`Unknown connection: ${connectionId}`);
    if (client.busy) throw new Error('Cannot switch session while agent is busy');

    const agent = this.getOrCreateAgent(sessionId);
    client.sessionId = sessionId;
    client.agent = agent;
  }

  /** Get the NativeAgent for a connection. */
  getAgent(connectionId: string): NativeAgent | undefined {
    return this.clients.get(connectionId)?.agent;
  }

  /** Get the session ID for a connection. */
  getSessionId(connectionId: string): string | undefined {
    return this.clients.get(connectionId)?.sessionId;
  }

  /** Check if a connection's agent is busy. */
  isBusy(connectionId: string): boolean {
    return this.clients.get(connectionId)?.busy ?? false;
  }

  /** Mark a connection's agent as busy/idle. */
  setBusy(connectionId: string, busy: boolean): void {
    const client = this.clients.get(connectionId);
    if (client) client.busy = busy;
  }

  /** Set onToolUse callback for a connection's agent. */
  setOnToolUse(connectionId: string, callback: ((event: ToolUseEvent) => void) | undefined): void {
    const client = this.clients.get(connectionId);
    if (client) client.agent.setOnToolUse(callback);
  }

  /** List all active sessions with connection counts. */
  listSessions(): Array<{ sessionId: string; connections: number }> {
    const sessionMap = new Map<string, number>();
    for (const client of this.clients.values()) {
      sessionMap.set(client.sessionId, (sessionMap.get(client.sessionId) ?? 0) + 1);
    }
    return Array.from(sessionMap.entries()).map(([sessionId, connections]) => ({
      sessionId,
      connections,
    }));
  }

  /** Get count of active connections. */
  get connectionCount(): number {
    return this.clients.size;
  }

  /** Get usage stats for a specific connection's agent. */
  getUsage(connectionId: string): { inputTokens: number; outputTokens: number; calls: number } | undefined {
    const agent = this.clients.get(connectionId)?.agent;
    return agent?.getUsage();
  }

  /** Get usage stats for all active sessions. Returns an array of per-session usage entries. */
  getAllUsage(): Array<{
    sessionId: string;
    primary: { inputTokens: number; outputTokens: number; calls: number };
    delegation: Record<string, { inputTokens: number; outputTokens: number; calls: number }>;
    total: { inputTokens: number; outputTokens: number; calls: number; estimatedCost: number };
  }> {
    const results: Array<{
      sessionId: string;
      primary: { inputTokens: number; outputTokens: number; calls: number };
      delegation: Record<string, { inputTokens: number; outputTokens: number; calls: number }>;
      total: { inputTokens: number; outputTokens: number; calls: number; estimatedCost: number };
    }> = [];

    // De-duplicate by sessionId (multiple connections may share a session)
    const seen = new Set<string>();

    for (const client of this.clients.values()) {
      if (seen.has(client.sessionId)) continue;
      seen.add(client.sessionId);

      const usage = client.agent.getUsage();
      results.push({
        sessionId: client.sessionId,
        primary: { inputTokens: usage.inputTokens, outputTokens: usage.outputTokens, calls: usage.calls },
        delegation: {} as Record<string, { inputTokens: number; outputTokens: number; calls: number }>,
        total: {
          inputTokens: usage.inputTokens,
          outputTokens: usage.outputTokens,
          calls: usage.calls,
          estimatedCost: 0, // NativeAgent doesn't track cost; only AgentOrchestrator does
        },
      });
    }

    return results;
  }

  private getOrCreateAgent(sessionId: string): NativeAgent {
    let agent = this.agents.get(sessionId);
    if (!agent) {
      const session = this.config.sessionManager.getSession('ws', sessionId);
      agent = new NativeAgent({
        modelClient: this.config.modelClient,
        systemPrompt: this.config.systemPrompt,
        session,
        toolRegistry: this.config.toolRegistry,
        toolExecutor: this.config.toolExecutor,
      });
      // Inherit the current model tier so the agent uses the same model as the TUI
      agent.setModelTier(this.currentTier);
      this.agents.set(sessionId, agent);
    }
    return agent;
  }
}