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                                105x 105x             105x 33x       33x             2x             55x                   39x   38x                 38x   38x             4x             23x 23x   22x             4x 4x   3x             4x 4x   3x 3x                     21x 21x   3x 3x   3x       18x     18x 1x       17x 17x         17x     17x     17x             3x             4x             36x                 34x 34x 8x               23x             105x         105x 88x 88x           105x 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;
}