flynn/src/gateway/handlers/system.ts

import type { GatewayRequest, OutboundMessage } from '../protocol.js';
import { makeResponse, makeError, ErrorCode } from '../protocol.js';
import type { MetricsSnapshot, EventEntry, ActiveRequestInfo } from '../metrics.js';
import type { ServiceInfo } from './services.js';
import type { NodeLocation } from './node.js';

/** Per-session token usage report returned by system.tokenUsage. */
export interface TokenUsageEntry {
  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 };
}

export interface PresenceEntry {
  channel: string;
  senderId: string;
  senderName?: string;
  firstSeenAt: number;
  lastSeenAt: number;
  messageCount: number;
  status: 'online' | 'offline';
}

export interface NodeLocationEntry {
  nodeId: string;
  role: string;
  connectionId: string;
  location: NodeLocation;
}

export interface SystemHandlerDeps {
  startTime: number;
  version: string;
  getSessionCount: () => number;
  getToolCount: () => number;
  getConnectionCount: () => number;
  /** Optional callback to trigger a graceful restart. If not provided, system.restart returns an error. */
  restart?: () => Promise<void>;
  getChannels?: () => Array<{ name: string; status: string }>;
  getUsage?: () => { totalSessions: number; activeConnections: number };
  /** Optional callback to retrieve per-session token usage data. */
  getTokenUsage?: () => TokenUsageEntry[];
  /** Optional callback to retrieve aggregated metrics snapshot. */
  getMetrics?: () => MetricsSnapshot;
  /** Optional callback to retrieve recent events. */
  getEvents?: (opts?: { level?: string; limit?: number }) => EventEntry[];
  /** Optional callback to retrieve active requests. */
  getActiveRequests?: () => ActiveRequestInfo[];
  /** Optional callback to retrieve all services (channels + automation). */
  getServices?: () => ServiceInfo[];
  /** Optional callback to retrieve tracked sender presence. */
  getPresence?: (opts?: { channel?: string; status?: 'online' | 'offline'; limit?: number }) => PresenceEntry[];
  /** Optional callback to retrieve latest node location data. */
  getNodeLocations?: (opts?: { role?: string; nodeId?: string; limit?: number }) => NodeLocationEntry[];
}

export function createSystemHandlers(deps: SystemHandlerDeps) {
  return {
    'system.health': async (request: GatewayRequest): Promise<OutboundMessage> => {
      return makeResponse(request.id, {
        status: 'ok',
        uptime: Math.floor((Date.now() - deps.startTime) / 1000),
        version: deps.version,
        sessions: deps.getSessionCount(),
        tools: deps.getToolCount(),
        connections: deps.getConnectionCount(),
      });
    },

    'system.restart': async (request: GatewayRequest): Promise<OutboundMessage> => {
      const restart = deps.restart;
      if (!restart) {
        return makeError(request.id, ErrorCode.InternalError, 'Restart not available in this environment');
      }

      // Send response before initiating restart (client receives confirmation)
      const response = makeResponse(request.id, { restarting: true });

      // Schedule restart after response is sent (next tick)
      queueMicrotask(() => {
        restart().catch((err) => {
          console.error('Restart failed:', err);
        });
      });

      return response;
    },

    'system.channels': async (request: GatewayRequest): Promise<OutboundMessage> => {
      if (!deps.getChannels) {
        return makeResponse(request.id, { channels: [] });
      }
      return makeResponse(request.id, { channels: deps.getChannels() });
    },

    'system.services': async (request: GatewayRequest): Promise<OutboundMessage> => {
      if (!deps.getServices) {
        return makeResponse(request.id, { services: [] });
      }
      return makeResponse(request.id, { services: deps.getServices() });
    },

    'system.presence': async (request: GatewayRequest): Promise<OutboundMessage> => {
      if (!deps.getPresence) {
        return makeResponse(request.id, { presence: [], summary: { total: 0, online: 0, offline: 0 } });
      }

      const params = request.params as { channel?: string; status?: 'online' | 'offline'; limit?: number } | undefined;
      const presence = deps.getPresence({
        channel: params?.channel,
        status: params?.status,
        limit: params?.limit,
      });
      const online = presence.filter((entry) => entry.status === 'online').length;
      return makeResponse(request.id, {
        presence,
        summary: {
          total: presence.length,
          online,
          offline: presence.length - online,
        },
      });
    },

    'system.location': async (request: GatewayRequest): Promise<OutboundMessage> => {
      if (!deps.getNodeLocations) {
        return makeResponse(request.id, { locations: [], summary: { total: 0 } });
      }

      const params = request.params as { role?: string; nodeId?: string; limit?: number } | undefined;
      const locations = deps.getNodeLocations({
        role: params?.role,
        nodeId: params?.nodeId,
        limit: params?.limit,
      });
      return makeResponse(request.id, {
        locations,
        summary: {
          total: locations.length,
        },
      });
    },

    'system.usage': async (request: GatewayRequest): Promise<OutboundMessage> => {
      const uptime = Math.floor((Date.now() - deps.startTime) / 1000);
      const usage = deps.getUsage?.() ?? { totalSessions: 0, activeConnections: 0 };
      return makeResponse(request.id, {
        uptime,
        ...usage,
        tools: deps.getToolCount(),
      });
    },

    'system.tokenUsage': async (request: GatewayRequest): Promise<OutboundMessage> => {
      const sessions = deps.getTokenUsage?.() ?? [];
      return makeResponse(request.id, { sessions });
    },

    'system.metrics': async (request: GatewayRequest): Promise<OutboundMessage> => {
      if (!deps.getMetrics) {
        return makeResponse(request.id, {});
      }
      return makeResponse(request.id, deps.getMetrics());
    },

    'system.events': async (request: GatewayRequest): Promise<OutboundMessage> => {
      if (!deps.getEvents) {
        return makeResponse(request.id, { events: [] });
      }
      const params = request.params as { level?: string; limit?: number } | undefined;
      const events = deps.getEvents({ level: params?.level, limit: params?.limit });
      return makeResponse(request.id, { events });
    },

    'system.activeRequests': async (request: GatewayRequest): Promise<OutboundMessage> => {
      if (!deps.getActiveRequests) {
        return makeResponse(request.id, { requests: [] });
      }
      return makeResponse(request.id, { requests: deps.getActiveRequests() });
    },
  };
}