All files / src/logging connection-tracker.ts

98.11% Statements 52/53
88.88% Branches 16/18
100% Functions 16/16
97.95% Lines 48/49

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209                                115x 115x             115x 33x       33x             2x             82x                   58x   57x                 57x   57x             4x             42x 42x   41x             4x 4x   3x             4x 4x   3x 3x                     28x 28x   3x 3x   3x       25x     25x 1x       24x 24x         24x     24x     24x             3x             4x             36x                 34x 34x 8x               23x             115x         115x 115x 115x           115x 5x    
/**
 * Connection Tracker
 *
 * Tracks statistics for SSE/persistent connections and logs a summary
 * when connections close.
 *
 * Each connection (identified by session ID) has its own stats:
 * - Total requests
 * - Total tool invocations
 * - Total errors
 * - Last error (if any)
 *
 * Logs a single line when connection closes with reason and stats.
 */
 
import type { ConnectionStats, ConnectionCloseReason } from "./types.js";
import { formatConnectionClose, createConnectionCloseEntry } from "./access-log.js";
import { logInfo, logDebug, LOG_JSON } from "../logger.js";
 
/**
 * Connection tracker manages statistics for persistent connections.
 *
 * Connections are identified by session ID (MCP session ID).
 */
export class ConnectionTracker {
  private connections: Map<string, ConnectionStats> = new Map();
  private enabled: boolean;
 
  constructor(enabled = true) {
    this.enabled = enabled;
  }
 
  /**
   * Check if connection tracking is enabled
   */
  isEnabled(): boolean {
    return this.enabled;
  }
 
  /**
   * Enable or disable connection tracking
   */
  setEnabled(enabled: boolean): void {
    this.enabled = enabled;
  }
 
  /**
   * Register a new connection
   *
   * @param sessionId - MCP session ID
   * @param clientIp - Client IP address
   */
  openConnection(sessionId: string, clientIp: string): void {
    if (!this.enabled) return;
 
    const stats: ConnectionStats = {
      connectedAt: Date.now(),
      clientIp,
      sessionId,
      requestCount: 0,
      toolCount: 0,
      errorCount: 0,
    };
 
    this.connections.set(sessionId, stats);
 
    logDebug("Connection opened for tracking", { sessionId, clientIp });
  }
 
  /**
   * Get connection stats for a session
   */
  getStats(sessionId: string): ConnectionStats | undefined {
    return this.connections.get(sessionId);
  }
 
  /**
   * Increment request count for a connection
   */
  incrementRequests(sessionId: string): void {
    const stats = this.connections.get(sessionId);
    if (!stats) return;
 
    stats.requestCount++;
  }
 
  /**
   * Increment tool count for a connection
   */
  incrementTools(sessionId: string): void {
    const stats = this.connections.get(sessionId);
    if (!stats) return;
 
    stats.toolCount++;
  }
 
  /**
   * Record an error on the connection
   */
  recordError(sessionId: string, error: string): void {
    const stats = this.connections.get(sessionId);
    if (!stats) return;
 
    stats.errorCount++;
    stats.lastError = error;
  }
 
  /**
   * Close a connection and log the summary
   *
   * @param sessionId - Session ID
   * @param reason - Why the connection closed
   * @returns The formatted log line (for testing) or undefined if disabled/not found
   */
  closeConnection(sessionId: string, reason: ConnectionCloseReason): string | undefined {
    const stats = this.connections.get(sessionId);
    if (!stats) {
      // Only log debug when enabled to avoid noise in verbose mode
      Eif (this.enabled) {
        logDebug("Connection not found on close", { sessionId });
      }
      return undefined;
    }
 
    // Remove from map first to prevent leaking tracked connections
    this.connections.delete(sessionId);
 
    // When disabled, skip log emission but cleanup still occurred
    if (!this.enabled) {
      return undefined;
    }
 
    // Format and log the connection close entry
    const entry = createConnectionCloseEntry(stats, reason);
    const logLine = formatConnectionClose(entry);
 
    // Output connection close log at info level
    // JSON mode: include full connectionClose object for log aggregators
    // Plain mode: message only - prevents pino-pretty from outputting multiline JSON
    Iif (LOG_JSON) {
      logInfo(logLine, { connectionClose: entry });
    } else {
      logInfo(logLine);
    }
 
    return logLine;
  }
 
  /**
   * Check if a connection is being tracked
   */
  hasConnection(sessionId: string): boolean {
    return this.connections.has(sessionId);
  }
 
  /**
   * Get number of active connections being tracked
   */
  getActiveConnectionCount(): number {
    return this.connections.size;
  }
 
  /**
   * Get all session IDs (for shutdown handling)
   */
  getAllSessionIds(): string[] {
    return Array.from(this.connections.keys());
  }
 
  /**
   * Close all connections (for server shutdown)
   *
   * @param reason - Reason for closing all connections
   */
  closeAllConnections(reason: ConnectionCloseReason = "server_shutdown"): void {
    const sessionIds = this.getAllSessionIds();
    for (const sessionId of sessionIds) {
      this.closeConnection(sessionId, reason);
    }
  }
 
  /**
   * Clear all tracked connections (for testing)
   */
  clear(): void {
    this.connections.clear();
  }
}
 
/**
 * Singleton instance of ConnectionTracker
 */
let globalConnectionTracker: ConnectionTracker | null = null;
 
/**
 * Get the global ConnectionTracker instance
 */
export function getConnectionTracker(): ConnectionTracker {
  globalConnectionTracker ??= new ConnectionTracker();
  return globalConnectionTracker;
}
 
/**
 * Reset the global connection tracker (for testing)
 */
export function resetConnectionTracker(): void {
  globalConnectionTracker = null;
}