ai-agent-session-center 2.0.2 → 2.0.3

package/server/teamManager.ts

@@ -0,0 +1,289 @@
+/**
+ * @module teamManager
+ * Manages team/subagent hierarchies. Teams are auto-created when a SubagentStart event
+ * matches a new session by working directory, or directly linked via CLAUDE_CODE_PARENT_SESSION_ID
+ * env var (Priority 0). Tracks parent-child relationships and handles cleanup when team
+ * members end (with 15s delay for the parent).
+ */
+import { readFileSync } from 'fs';
+import { join } from 'path';
+import { homedir } from 'os';
+import log from './logger.js';
+import { SESSION_STATUS } from './constants.js';
+import type { Session } from '../src/types/session.js';
+import type { Team, TeamSerialized, PendingSubagent, TeamConfig, TeamLinkResult } from '../src/types/team.js';
+
+const teams = new Map<string, Team>();            // teamId -> Team
+const sessionToTeam = new Map<string, string>();   // sessionId -> teamId
+const pendingSubagents: PendingSubagent[] = [];    // { parentSessionId, parentCwd, agentType, timestamp }
+
+/**
+ * Find a pending subagent entry matching a new child session.
+ * Matches by cwd (exact or parent/child path relationship).
+ */
+export function findPendingSubagentMatch(
+  childSessionId: string,
+  childCwd: string,
+  sessions: Map<string, Session>,
+): TeamLinkResult | null {
+  const now = Date.now();
+  // Clean stale entries (>10s old)
+  while (pendingSubagents.length > 0 && now - pendingSubagents[0].timestamp > 10000) {
+    pendingSubagents.shift();
+  }
+  if (!childCwd || pendingSubagents.length === 0) return null;
+
+  // Match by cwd — exact match or parent/child path relationship
+  for (let i = pendingSubagents.length - 1; i >= 0; i--) {
+    const pending = pendingSubagents[i];
+    if (pending.parentSessionId === childSessionId) continue; // skip self
+    const parentCwd = pending.parentCwd;
+    if (parentCwd && (childCwd === parentCwd || childCwd.startsWith(parentCwd + '/') || parentCwd.startsWith(childCwd + '/'))) {
+      // Found match — consume it
+      pendingSubagents.splice(i, 1);
+      return linkSessionToTeam(pending.parentSessionId, childSessionId, pending.agentType, sessions);
+    }
+  }
+  return null;
+}
+
+/**
+ * Directly link a child session to its parent using CLAUDE_CODE_PARENT_SESSION_ID.
+ * This is Priority 0 matching — no path guessing needed.
+ */
+export function linkByParentSessionId(
+  childSessionId: string,
+  parentSessionId: string,
+  agentType: string,
+  agentName: string | null,
+  teamName: string | null,
+  sessions: Map<string, Session>,
+): TeamLinkResult | null {
+  if (!parentSessionId || !childSessionId) return null;
+  if (parentSessionId === childSessionId) return null;
+
+  const parentSession = sessions.get(parentSessionId);
+  if (!parentSession) {
+    log.debug('session', `linkByParentSessionId: parent ${parentSessionId?.slice(0, 8)} not found in sessions`);
+    return null;
+  }
+
+  const result = linkSessionToTeam(parentSessionId, childSessionId, agentType, sessions);
+
+  // Apply team name from env var if available
+  if (teamName && result) {
+    const team = teams.get(result.teamId);
+    if (team) {
+      team.teamName = teamName;
+      result.team = serializeTeam(team)!;
+    }
+  }
+
+  // Store agent name on the child session
+  const childSession = sessions.get(childSessionId);
+  if (childSession && agentName) {
+    childSession.agentName = agentName;
+  }
+
+  // Try to read team config for additional member metadata
+  const effectiveTeamName = teamName || (result ? teams.get(result.teamId)?.teamName : null);
+  if (effectiveTeamName && childSession) {
+    const config = readTeamConfig(effectiveTeamName);
+    if (config && agentName && config.members) {
+      const memberConfig = config.members[agentName];
+      if (memberConfig) {
+        if (memberConfig.tmuxPaneId) childSession.tmuxPaneId = memberConfig.tmuxPaneId;
+        if (memberConfig.backendType) childSession.backendType = memberConfig.backendType;
+        if (memberConfig.color) childSession.agentColor = memberConfig.color;
+      }
+    }
+  }
+
+  log.info('session', `linkByParentSessionId: ${childSessionId?.slice(0, 8)} -> parent ${parentSessionId?.slice(0, 8)} (agent=${agentName || 'unknown'}, team=${teamName || 'auto'})`);
+  return result;
+}
+
+/**
+ * Read team configuration from ~/.claude/teams/{teamName}/config.json.
+ * Returns null if file doesn't exist or is invalid.
+ */
+export function readTeamConfig(teamName: string): TeamConfig | null {
+  if (!teamName || typeof teamName !== 'string') return null;
+  // Sanitize team name to prevent path traversal
+  const safeName = teamName.replace(/[^a-zA-Z0-9_\-. ]/g, '');
+  if (!safeName) return null;
+
+  const configPath = join(homedir(), '.claude', 'teams', safeName, 'config.json');
+  try {
+    const raw = readFileSync(configPath, 'utf8');
+    return JSON.parse(raw) as TeamConfig;
+  } catch {
+    // Config file doesn't exist or is invalid — that's fine
+    return null;
+  }
+}
+
+/**
+ * Link a child session to a team (creating the team if needed).
+ */
+function linkSessionToTeam(
+  parentId: string,
+  childId: string,
+  agentType: string,
+  sessions: Map<string, Session>,
+): TeamLinkResult {
+  const teamId = `team-${parentId}`;
+  let team = teams.get(teamId);
+
+  if (!team) {
+    team = {
+      teamId,
+      parentSessionId: parentId,
+      childSessionIds: new Set<string>(),
+      teamName: null,
+      createdAt: Date.now(),
+    };
+    teams.set(teamId, team);
+
+    // Set team name from parent's project name
+    const parentSession = sessions.get(parentId);
+    if (parentSession) {
+      team.teamName = `${parentSession.projectName} Team`;
+      parentSession.teamId = teamId;
+      parentSession.teamRole = 'leader';
+      sessionToTeam.set(parentId, teamId);
+    }
+  }
+
+  // Link child
+  team.childSessionIds.add(childId);
+  const childSession = sessions.get(childId);
+  if (childSession) {
+    childSession.teamId = teamId;
+    childSession.teamRole = 'member';
+    childSession.agentType = agentType;
+  }
+  sessionToTeam.set(childId, teamId);
+
+  log.info('session', `Linked session ${childId} to team ${teamId} as ${agentType}`);
+  return { teamId, team: serializeTeam(team)! };
+}
+
+/**
+ * Handle team cleanup when a member session ends.
+ */
+export function handleTeamMemberEnd(
+  sessionId: string,
+  sessions: Map<string, Session>,
+): TeamLinkResult | null {
+  const teamId = sessionToTeam.get(sessionId);
+  if (!teamId) return null;
+
+  const team = teams.get(teamId);
+  if (!team) return null;
+
+  team.childSessionIds.delete(sessionId);
+  sessionToTeam.delete(sessionId);
+
+  // If parent ended and all children ended, clean up the team
+  if (sessionId === team.parentSessionId) {
+    const allChildrenEnded = [...team.childSessionIds].every(cid => {
+      const s = sessions.get(cid);
+      return !s || s.status === SESSION_STATUS.ENDED;
+    });
+    if (allChildrenEnded) {
+      // Clean up team after a delay
+      setTimeout(() => {
+        teams.delete(teamId);
+        sessionToTeam.delete(team.parentSessionId);
+        for (const cid of team.childSessionIds) {
+          sessionToTeam.delete(cid);
+        }
+      }, 15000);
+    }
+  }
+
+  return { teamId, team: serializeTeam(team)! };
+}
+
+/**
+ * Add a pending subagent entry for team auto-detection.
+ */
+export function addPendingSubagent(
+  parentSessionId: string,
+  parentCwd: string,
+  agentType: string | undefined,
+  agentId: string | undefined,
+): void {
+  pendingSubagents.push({
+    parentSessionId,
+    parentCwd,
+    agentType: agentType || 'unknown',
+    agentId: agentId || null,
+    timestamp: Date.now(),
+  });
+  // Prune stale entries (>30s old)
+  const now = Date.now();
+  while (pendingSubagents.length > 0 && now - pendingSubagents[0].timestamp > 30000) {
+    pendingSubagents.shift();
+  }
+}
+
+function serializeTeam(team: Team): TeamSerialized | null {
+  if (!team) return null;
+  return {
+    teamId: team.teamId,
+    parentSessionId: team.parentSessionId,
+    childSessionIds: [...team.childSessionIds],
+    teamName: team.teamName,
+    createdAt: team.createdAt,
+  };
+}
+
+export function getTeam(teamId: string): TeamSerialized | null {
+  const team = teams.get(teamId);
+  return team ? serializeTeam(team) : null;
+}
+
+export function getAllTeams(): Record<string, TeamSerialized> {
+  const result: Record<string, TeamSerialized> = {};
+  for (const [id, team] of teams) {
+    const serialized = serializeTeam(team);
+    if (serialized) result[id] = serialized;
+  }
+  return result;
+}
+
+export function getTeamForSession(sessionId: string): TeamSerialized | null {
+  const teamId = sessionToTeam.get(sessionId);
+  if (!teamId) return null;
+  return getTeam(teamId);
+}
+
+export function getTeamIdForSession(sessionId: string): string | null {
+  return sessionToTeam.get(sessionId) || null;
+}
+
+/**
+ * Get the tmux pane ID for a team member session.
+ * Looks up the session's stored tmuxPaneId field (set during linkByParentSessionId).
+ */
+export function getMemberTmuxPaneId(
+  teamId: string,
+  sessionId: string,
+  sessions: Map<string, Session>,
+): string | null {
+  const team = teams.get(teamId);
+  if (!team) return null;
+  if (sessionId !== team.parentSessionId && !team.childSessionIds.has(sessionId)) return null;
+  const session = sessions.get(sessionId);
+  return session?.tmuxPaneId || null;
+}
+
+/**
+ * Get the config file path for a team.
+ */
+export function getTeamConfigPath(teamName: string): string {
+  const safeName = (teamName || '').replace(/[^a-zA-Z0-9_\-. ]/g, '');
+  return join(homedir(), '.claude', 'teams', safeName, 'config.json');
+}

package/server/wsManager.ts

@@ -0,0 +1,200 @@
+// wsManager.ts — WebSocket broadcast manager with bidirectional terminal support
+import { getAllSessions, getAllTeams, getEventSeq, getEventsSince, updateQueueCount } from './sessionStore.js';
+import { writeToTerminal, resizeTerminal, closeTerminal, setWsClient } from './sshManager.js';
+import { WS_TYPES } from './constants.js';
+import log from './logger.js';
+import type WebSocket from 'ws';
+
+interface WsClient extends WebSocket {
+  _terminalIds: Set<string>;
+  _isAlive: boolean;
+}
+
+const clients = new Set<WsClient>();
+
+// Heartbeat: ping every 30s, terminate connections that don't pong within 10s
+const HEARTBEAT_INTERVAL_MS = 30000;
+let heartbeatTimer: ReturnType<typeof setInterval> | null = null;
+
+// Backpressure: skip non-critical updates if client buffer exceeds 1MB
+const MAX_BUFFERED_AMOUNT = 1 * 1024 * 1024;
+
+// Throttle hook_stats broadcasts to once per second max
+let lastHookStatsBroadcastAt = 0;
+let pendingHookStats: unknown = null;
+let hookStatsTimer: ReturnType<typeof setTimeout> | null = null;
+const HOOK_STATS_THROTTLE_MS = 1000;
+
+function startHeartbeat(): void {
+  if (heartbeatTimer) return;
+  heartbeatTimer = setInterval(() => {
+    for (const ws of clients) {
+      if (ws._isAlive === false) {
+        // Didn't respond to last ping — terminate
+        log.info('ws', 'Terminating unresponsive client');
+        ws.terminate();
+        clients.delete(ws);
+        continue;
+      }
+      ws._isAlive = false;
+      ws.ping();
+    }
+  }, HEARTBEAT_INTERVAL_MS);
+}
+
+export function stopHeartbeat(): void {
+  if (heartbeatTimer) {
+    clearInterval(heartbeatTimer);
+    heartbeatTimer = null;
+  }
+  if (hookStatsTimer) {
+    clearTimeout(hookStatsTimer);
+    hookStatsTimer = null;
+  }
+}
+
+/**
+ * Handle a new WebSocket connection: send snapshot and wire up message/close handlers.
+ */
+export function handleConnection(ws: WebSocket): void {
+  const client = ws as WsClient;
+  clients.add(client);
+  client._terminalIds = new Set();
+  client._isAlive = true;
+  log.info('ws', `Client connected (total: ${clients.size})`);
+
+  // Start heartbeat on first connection
+  startHeartbeat();
+
+  // Handle pong responses
+  client.on('pong', () => {
+    client._isAlive = true;
+  });
+
+  // Send full snapshot on connect (includes teams + event sequence for replay)
+  const sessions = getAllSessions();
+  const teams = getAllTeams();
+  const seq = getEventSeq();
+  log.debug('ws', `Sending snapshot: ${Object.keys(sessions).length} sessions, ${Object.keys(teams).length} teams, seq=${seq}`);
+  client.send(JSON.stringify({ type: WS_TYPES.SNAPSHOT, sessions, teams, seq }));
+
+  // Handle incoming messages (terminal input, resize, etc.)
+  client.on('message', (raw: WebSocket.RawData) => {
+    try {
+      const msg = JSON.parse(raw.toString());
+      switch (msg.type) {
+        case WS_TYPES.TERMINAL_INPUT:
+          if (msg.terminalId && msg.data) {
+            writeToTerminal(msg.terminalId, msg.data);
+          }
+          break;
+        case WS_TYPES.TERMINAL_RESIZE:
+          if (msg.terminalId && msg.cols && msg.rows) {
+            resizeTerminal(msg.terminalId, msg.cols, msg.rows);
+          }
+          break;
+        case WS_TYPES.TERMINAL_DISCONNECT:
+          if (msg.terminalId) {
+            closeTerminal(msg.terminalId);
+            client._terminalIds.delete(msg.terminalId);
+          }
+          break;
+        case WS_TYPES.TERMINAL_SUBSCRIBE:
+          if (msg.terminalId) {
+            client._terminalIds.add(msg.terminalId);
+            setWsClient(msg.terminalId, client);
+          }
+          break;
+        case WS_TYPES.UPDATE_QUEUE_COUNT:
+          if (msg.sessionId != null && msg.count != null) {
+            const updated = updateQueueCount(msg.sessionId, msg.count);
+            if (updated) {
+              broadcast({ type: WS_TYPES.SESSION_UPDATE, session: updated });
+            }
+          }
+          break;
+        case WS_TYPES.REPLAY:
+          // Client reconnected and wants events since a certain sequence number
+          if (typeof msg.sinceSeq === 'number') {
+            const missed = getEventsSince(msg.sinceSeq);
+            log.debug('ws', `Replaying ${missed.length} events since seq=${msg.sinceSeq}`);
+            for (const evt of missed) {
+              client.send(JSON.stringify(evt.data));
+            }
+          }
+          break;
+        default:
+          log.debug('ws', `Unknown message type: ${msg.type}`);
+      }
+    } catch (e: unknown) {
+      const msg = e instanceof Error ? e.message : String(e);
+      log.debug('ws', `Invalid WS message: ${msg}`);
+    }
+  });
+
+  client.on('close', () => {
+    clients.delete(client);
+    log.info('ws', `Client disconnected (total: ${clients.size})`);
+    // Stop heartbeat if no clients remain
+    if (clients.size === 0) {
+      stopHeartbeat();
+    }
+  });
+  client.on('error', (err: Error) => {
+    clients.delete(client);
+    log.error('ws', 'Client error:', err.message);
+  });
+}
+
+/**
+ * Check if a broadcast type is critical (must not be skipped under backpressure).
+ * Session updates and snapshots are critical; hook_stats are not.
+ */
+function isCriticalBroadcast(data: { type: string }): boolean {
+  return data.type !== WS_TYPES.HOOK_STATS;
+}
+
+/**
+ * Broadcast a message to all connected WebSocket clients.
+ * Throttles hook_stats to once per second; applies backpressure for non-critical messages.
+ */
+export function broadcast(data: { type: string; [key: string]: unknown }): void {
+  // Throttle hook_stats broadcasts to once per second max
+  if (data.type === WS_TYPES.HOOK_STATS) {
+    const now = Date.now();
+    if (now - lastHookStatsBroadcastAt < HOOK_STATS_THROTTLE_MS) {
+      // Store for deferred send
+      pendingHookStats = data;
+      if (!hookStatsTimer) {
+        hookStatsTimer = setTimeout(() => {
+          hookStatsTimer = null;
+          if (pendingHookStats) {
+            const deferred = pendingHookStats as { type: string; [key: string]: unknown };
+            pendingHookStats = null;
+            lastHookStatsBroadcastAt = Date.now();
+            broadcastToClients(deferred, false);
+          }
+        }, HOOK_STATS_THROTTLE_MS - (now - lastHookStatsBroadcastAt));
+      }
+      return;
+    }
+    lastHookStatsBroadcastAt = now;
+  }
+
+  const critical = isCriticalBroadcast(data);
+  broadcastToClients(data, critical);
+}
+
+function broadcastToClients(data: { type: string; [key: string]: unknown }, critical: boolean): void {
+  const msg = JSON.stringify(data);
+  log.debug('ws', `Broadcasting ${data.type} to ${clients.size} clients`);
+  for (const client of clients) {
+    if (client.readyState !== 1) continue;
+    // Backpressure: skip non-critical updates if buffer is too large
+    if (!critical && client.bufferedAmount > MAX_BUFFERED_AMOUNT) {
+      log.debug('ws', `Skipping ${data.type} for client (buffered=${client.bufferedAmount})`);
+      continue;
+    }
+    client.send(msg);
+  }
+}