mcpctl/src/mcplocal/src/router.ts

import type { UpstreamConnection, JsonRpcRequest, JsonRpcResponse, JsonRpcNotification } from './types.js';
import type { LlmProcessor } from './llm/processor.js';
import { ResponsePaginator } from './llm/pagination.js';
import type { McpdClient } from './http/mcpd-client.js';
import type { PromptIndexEntry } from './gate/tag-matcher.js';
import { LinkResolver } from './services/link-resolver.js';
import type { LLMProvider, CacheProvider, Section } from './proxymodel/types.js';
import { executePipeline } from './proxymodel/executor.js';
import { currentCorrelationId, currentScope } from './request-context.js';
import { getProxyModel } from './proxymodel/loader.js';
import type { ProxyModelPlugin, PluginSessionContext } from './proxymodel/plugin.js';
import { PluginContextImpl, type PluginContextDeps } from './proxymodel/plugin-context.js';
import type { AuditCollector } from './audit/collector.js';
import { pauseQueue } from './proxymodel/pause-queue.js';

export interface RouteContext {
  sessionId?: string;
  /** Correlation ID for traffic inspection (links upstream calls to client request) */
  correlationId?: string;
}

type ListCacheEntry =
  | { kind: 'ok'; result: unknown; fetchedAt: number }
  | { kind: 'err'; message: string; fetchedAt: number };

/**
 * Routes MCP requests to the appropriate upstream server.
 *
 * The proxy presents a unified MCP interface to clients. Tools, resources,
 * and prompts from all upstreams are merged under namespaced prefixes
 * (e.g., "slack/send_message").
 *
 * Routing is done by name prefix: "servername/toolname" -> upstream "servername".
 *
 * With no plugin attached, the router is a transparent MCP proxy — all requests
 * flow through unchanged. Plugins (gate, content-pipeline, etc.) add optional
 * interception on top.
 */
export class McpRouter {
  private upstreams = new Map<string, UpstreamConnection>();
  private toolToServer = new Map<string, string>();
  private resourceToServer = new Map<string, string>();
  private promptToServer = new Map<string, string>();
  private notificationHandler: ((notification: JsonRpcNotification) => void) | null = null;
  private llmProcessor: LlmProcessor | null = null;
  private instructions: string | null = null;
  private mcpdClient: McpdClient | null = null;
  private projectName: string | null = null;
  private mcpctlResourceContents = new Map<string, string>();
  private paginator: ResponsePaginator | null = null;
  private linkResolver: LinkResolver | null = null;
  private pendingNotifications = new Map<string, JsonRpcNotification[]>();

  // ProxyModel pipeline services (used by plugin context)
  private proxyModelName: string | null = null;
  private proxyModelLlm: LLMProvider | null = null;
  private proxyModelCache: CacheProvider | null = null;
  private auditCollector: AuditCollector | null = null;
  private serverProxyModels = new Map<string, { name: string; llm: LLMProvider; cache: CacheProvider }>();
  /** Per-server tool-call deadline overrides, in ms, keyed by server name. */
  private serverDeadlines = new Map<string, number>();

  // Prompt and system prompt caches (used by plugin context)
  private cachedPromptIndex: PromptIndexEntry[] | null = null;
  private promptIndexFetchedAt = 0;
  private readonly PROMPT_INDEX_TTL_MS = 60_000;
  private systemPromptCache = new Map<string, { content: string; fetchedAt: number }>();
  private readonly SYSTEM_PROMPT_TTL_MS = 300_000; // 5 minutes

  // Plugin system — when set, plugin hooks are dispatched for interception
  private plugin: ProxyModelPlugin | null = null;
  private pluginContexts = new Map<string, PluginContextImpl>();

  // Per-server discovery cache. Keyed `${serverName}:${method}`. Prevents every client
  // `tools/list` from re-hitting slow/dead upstreams and absorbs negative results so one
  // dead server only stalls the first POST, not every subsequent one.
  private listCache = new Map<string, ListCacheEntry>();
  /** Per-server failures from the most recent discoverTools run. */
  private toolDiscoveryFailures: Array<{ server: string; message: string }> = [];
  private readonly LIST_CACHE_POSITIVE_TTL_MS = 30_000;
  private readonly LIST_CACHE_NEGATIVE_TTL_MS = 30_000;

  /** Optional callback for traffic inspection — called after each upstream call completes. */
  onUpstreamCall: ((info: { upstream: string; method?: string; request: unknown; response: unknown; durationMs: number; correlationId?: string }) => void) | null = null;

  setPaginator(paginator: ResponsePaginator): void {
    this.paginator = paginator;
  }

  setProxyModel(name: string, llm: LLMProvider, cache: CacheProvider): void {
    this.proxyModelName = name;
    this.proxyModelLlm = llm;
    this.proxyModelCache = cache;
  }

  /**
   * Per-server override for the tool-call deadline, in milliseconds.
   *
   * Sourced from the server resource's `toolCallTimeoutSeconds`. Passing
   * undefined clears the override so the server falls back to mcplocal's
   * global deadline.
   */
  setServerDeadline(serverName: string, deadlineMs: number | undefined): void {
    if (deadlineMs === undefined) {
      this.serverDeadlines.delete(serverName);
      return;
    }
    this.serverDeadlines.set(serverName, deadlineMs);
  }

  /** Deadline for a canonical `server/tool` name, or undefined for the global. */
  getToolCallDeadlineMs(canonicalToolName?: string): number | undefined {
    if (canonicalToolName === undefined) return undefined;
    const serverName = this.toolToServer.get(canonicalToolName)
      ?? canonicalToolName.split('/')[0];
    return serverName === undefined ? undefined : this.serverDeadlines.get(serverName);
  }

  setServerProxyModel(serverName: string, name: string, llm: LLMProvider, cache: CacheProvider): void {
    this.serverProxyModels.set(serverName, { name, llm, cache });
  }

  private getProxyModelForServer(serverName: string): { name: string; llm: LLMProvider; cache: CacheProvider } | null {
    const serverOverride = this.serverProxyModels.get(serverName);
    if (serverOverride) return serverOverride;
    if (this.proxyModelName && this.proxyModelLlm && this.proxyModelCache) {
      return { name: this.proxyModelName, llm: this.proxyModelLlm, cache: this.proxyModelCache };
    }
    return null;
  }

  setAuditCollector(collector: AuditCollector): void {
    this.auditCollector = collector;
  }

  getAuditCollector(): AuditCollector | null {
    return this.auditCollector;
  }

  setLlmProcessor(processor: LlmProcessor): void {
    this.llmProcessor = processor;
  }

  setInstructions(instructions: string): void {
    this.instructions = instructions;
  }

  setPromptConfig(mcpdClient: McpdClient, projectName: string): void {
    this.mcpdClient = mcpdClient;
    this.projectName = projectName;
    this.linkResolver = new LinkResolver(mcpdClient);
  }

  /** Return a bound system prompt fetcher for external consumers (e.g. ResponsePaginator). */
  getSystemPromptFn(): (name: string, fallback: string) => Promise<string> {
    return (name, fallback) => this.getSystemPrompt(name, fallback);
  }

  /** Set the plugin for this router. When set, plugin hooks are dispatched. */
  setPlugin(plugin: ProxyModelPlugin): void {
    this.plugin = plugin;
  }

  /** Get or create a plugin session context for the given session ID. */
  private async getOrCreatePluginContext(sessionId: string): Promise<PluginContextImpl> {
    let ctx = this.pluginContexts.get(sessionId);
    if (ctx) return ctx;

    const deps: PluginContextDeps = {
      sessionId,
      projectName: this.projectName ?? 'unknown',
      llm: this.proxyModelLlm ?? { complete: async () => '', available: () => false },
      cache: this.proxyModelCache ?? { getOrCompute: async (_k, fn) => fn(), hash: () => '', get: async () => null, set: async () => {} },
      log: {
        debug: (msg: string) => console.debug(`[plugin:${sessionId.slice(0, 8)}] ${msg}`),
        info: (msg: string) => console.info(`[plugin:${sessionId.slice(0, 8)}] ${msg}`),
        warn: (msg: string) => console.warn(`[plugin:${sessionId.slice(0, 8)}] ${msg}`),
        error: (msg: string) => console.error(`[plugin:${sessionId.slice(0, 8)}] ${msg}`),
      },
      discoverTools: () => this.discoverTools(),
      getToolDiscoveryFailures: () => this.getToolDiscoveryFailures(),
      routeToUpstream: (request) => this.routeNamespacedCall(request, 'name', this.toolToServer),
      fetchPromptIndex: () => this.fetchPromptIndex(),
      getSystemPrompt: (name, fallback) => this.getSystemPrompt(name, fallback),
      processContent: async (toolName, content, contentType) => {
        const serverName = this.toolToServer.get(toolName);
        const pmConfig = serverName ? this.getProxyModelForServer(serverName) : null;
        const effectiveName = pmConfig?.name ?? this.proxyModelName;
        const effectiveLlm = pmConfig?.llm ?? this.proxyModelLlm;
        const effectiveCache = pmConfig?.cache ?? this.proxyModelCache;
        const proxyModelDef = effectiveName ? await getProxyModel(effectiveName) : null;
        if (!proxyModelDef || !effectiveLlm || !effectiveCache) {
          return { content };
        }
        const result = await executePipeline({
          content,
          contentType,
          sourceName: toolName,
          projectName: this.projectName ?? 'unknown',
          sessionId,
          proxyModel: proxyModelDef,
          llm: effectiveLlm,
          cache: effectiveCache,
          getSystemPrompt: (name, fallback) => this.getSystemPrompt(name, fallback),
          ...(this.auditCollector ? { auditCollector: this.auditCollector } : {}),
          ...(serverName !== undefined ? { serverName } : {}),
          // processContent is a per-session closure with no RouteContext, so
          // the id comes from the request scope. Without this,
          // ExecuteOptions.correlationId was dead code and every
          // stage_execution / pipeline_execution row landed unjoinable.
          ...(currentCorrelationId() !== undefined ? { correlationId: currentCorrelationId()! } : {}),
        });

        // Pause queue: if paused, hold the result until released/edited/dropped.
        // This wait is deliberately unbounded — it ends when a human operator
        // acts — so the request deadline stops its clock rather than killing
        // it, and re-arms with whatever was left once the operator releases.
        if (pauseQueue.paused) {
          const enqueue = (): Promise<string> => pauseQueue.enqueue({
            sessionId,
            projectName: this.projectName ?? 'unknown',
            contentType,
            sourceName: toolName,
            original: content,
            transformed: result.content,
          });
          const deadline = currentScope()?.deadline;
          const pausedContent = deadline ? await deadline.suspended(enqueue) : await enqueue();
          return { ...result, content: pausedContent };
        }

        return result;
      },
      queueNotification: (notification) => this.queueNotification(sessionId, notification),
      postToMcpd: async (path, body) => {
        if (!this.mcpdClient) throw new Error('mcpd client not configured');
        return this.mcpdClient.post(path, body);
      },
      getFromMcpd: async (path) => {
        if (!this.mcpdClient) throw new Error('mcpd client not configured');
        return this.mcpdClient.get(path);
      },
      ...(this.auditCollector ? { auditCollector: this.auditCollector } : {}),
    };

    ctx = new PluginContextImpl(deps);
    this.pluginContexts.set(sessionId, ctx);

    // Initialize the session via plugin hook (registers virtual tools, etc.)
    if (this.plugin?.onSessionCreate) {
      await this.plugin.onSessionCreate(ctx);
    }

    return ctx;
  }

  addUpstream(connection: UpstreamConnection): void {
    this.upstreams.set(connection.name, connection);
    this.invalidateListCache(connection.name);
    if (this.notificationHandler && connection.onNotification) {
      const serverName = connection.name;
      const handler = this.notificationHandler;
      connection.onNotification((notification) => {
        handler({
          ...notification,
          params: {
            ...notification.params,
            _source: serverName,
          },
        });
      });
    }
  }

  removeUpstream(name: string): void {
    this.upstreams.delete(name);
    this.invalidateListCache(name);
    for (const map of [this.toolToServer, this.resourceToServer, this.promptToServer]) {
      for (const [key, server] of map) {
        if (server === name) {
          map.delete(key);
        }
      }
    }
  }

  /** Drop all discovery-cache entries for a server (called on register / remove). */
  private invalidateListCache(serverName: string): void {
    const prefix = `${serverName}:`;
    for (const key of this.listCache.keys()) {
      if (key.startsWith(prefix)) this.listCache.delete(key);
    }
  }

  private getListCacheEntry(serverName: string, method: string): ListCacheEntry | null {
    const entry = this.listCache.get(`${serverName}:${method}`);
    if (!entry) return null;
    const ttl = entry.kind === 'ok' ? this.LIST_CACHE_POSITIVE_TTL_MS : this.LIST_CACHE_NEGATIVE_TTL_MS;
    if (Date.now() - entry.fetchedAt >= ttl) return null;
    return entry;
  }

  private setListCacheEntry(serverName: string, method: string, entry: ListCacheEntry): void {
    this.listCache.set(`${serverName}:${method}`, entry);
  }

  setNotificationHandler(handler: (notification: JsonRpcNotification) => void): void {
    this.notificationHandler = handler;
    // Wire to all existing upstreams
    for (const [serverName, upstream] of this.upstreams) {
      if (upstream.onNotification) {
        upstream.onNotification((notification) => {
          handler({
            ...notification,
            params: {
              ...notification.params,
              _source: serverName,
            },
          });
        });
      }
    }
  }

  /**
   * Discover tools from all upstreams by calling tools/list on each.
   * Per-server results are cached (positive + negative) to absorb slow upstreams
   * and prevent repeated 30s timeouts on every client `tools/list`.
   */
  async discoverTools(correlationId?: string): Promise<Array<{ name: string; description?: string; inputSchema?: unknown }>> {
    const allTools: Array<{ name: string; description?: string; inputSchema?: unknown }> = [];
    const started = Date.now();
    let cachedCount = 0;
    let freshCount = 0;
    const failed: Array<{ server: string; message: string }> = [];

    // Discover tools from all servers in parallel so one slow server doesn't block the rest
    const entries = [...this.upstreams.entries()];
    const results = await Promise.allSettled(
      entries.map(async ([serverName, upstream]) => {
        const cached = this.getListCacheEntry(serverName, 'tools/list');
        if (cached) {
          return { serverName, upstream, cached };
        }
        const req = {
          jsonrpc: '2.0' as const,
          id: `discover-tools-${serverName}`,
          method: 'tools/list',
        };
        let response: JsonRpcResponse;
        if (this.onUpstreamCall) {
          const start = performance.now();
          response = await upstream.send(req);
          const durationMs = Math.round(performance.now() - start);
          this.onUpstreamCall({ upstream: serverName, method: req.method, request: req, response, durationMs, ...(correlationId ? { correlationId } : {}) });
        } else {
          response = await upstream.send(req);
        }
        return { serverName, upstream, response };
      }),
    );

    for (const result of results) {
      if (result.status === 'rejected') {
        console.warn(`[discoverTools] ${(result.reason as Error).message ?? 'unknown error'}`);
        continue;
      }
      const { serverName, upstream } = result.value;

      let response: JsonRpcResponse | null = null;
      if ('cached' in result.value) {
        const cached = result.value.cached;
        if (cached.kind === 'err') {
          cachedCount++;
          failed.push({ server: serverName, message: cached.message });
          continue;
        }
        response = { jsonrpc: '2.0', id: `cached-${serverName}`, result: cached.result };
        cachedCount++;
      } else {
        response = result.value.response;
        freshCount++;
        if (response.error) {
          const message = (response.error as { message?: string }).message ?? 'unknown error';
          this.setListCacheEntry(serverName, 'tools/list', { kind: 'err', message, fetchedAt: Date.now() });
          console.warn(`[discoverTools] ${serverName}: ${message}`);
          failed.push({ server: serverName, message });
          continue;
        }
        if (response.result !== undefined) {
          this.setListCacheEntry(serverName, 'tools/list', { kind: 'ok', result: response.result, fetchedAt: Date.now() });
        }
      }

      if (response.result && typeof response.result === 'object' && 'tools' in response.result) {
        const tools = (response.result as { tools: Array<{ name: string; description?: string; inputSchema?: unknown }> }).tools;
        for (const tool of tools) {
          const namespacedName = `${serverName}/${tool.name}`;
          this.toolToServer.set(namespacedName, serverName);
          // Enrich description with server context if available
          const entry: { name: string; description?: string; inputSchema?: unknown } = {
            ...tool,
            name: namespacedName,
          };
          if (upstream.description && tool.description) {
            entry.description = `[${upstream.description}] ${tool.description}`;
          } else if (upstream.description) {
            entry.description = `[${upstream.description}]`;
          }
          // If neither upstream.description nor tool.description, keep tool.description (may be undefined — that's fine, just don't set it)
          allTools.push(entry);
        }
      }
    }

    this.toolDiscoveryFailures = failed;

    if (entries.length > 0) {
      const elapsed = Date.now() - started;
      const project = this.projectName ? ` project=${this.projectName}` : '';
      const failedStr = failed.length > 0 ? ` failed=[${failed.map((f) => f.server).join(',')}]` : '';
      console.info(`[discoverTools]${project} fresh=${freshCount} cached=${cachedCount}${failedStr} elapsed=${elapsed}ms`);
    }

    return allTools;
  }

  /**
   * Servers that contributed no tools on the most recent `discoverTools` run,
   * with the error that silenced them. A partial catalog must never look like
   * a healthy one: the gate reads this to warn the client in-band.
   */
  getToolDiscoveryFailures(): Array<{ server: string; message: string }> {
    return this.toolDiscoveryFailures;
  }

  /**
   * Discover resources from all upstreams by calling resources/list on each.
   * Shares the per-server list cache with `discoverTools`.
   */
  async discoverResources(correlationId?: string): Promise<Array<{ uri: string; name?: string; description?: string; mimeType?: string }>> {
    const allResources: Array<{ uri: string; name?: string; description?: string; mimeType?: string }> = [];

    // Discover resources from all servers in parallel
    const entries = [...this.upstreams.entries()];
    const results = await Promise.allSettled(
      entries.map(async ([serverName, upstream]) => {
        const cached = this.getListCacheEntry(serverName, 'resources/list');
        if (cached) return { serverName, cached };
        const req = {
          jsonrpc: '2.0' as const,
          id: `discover-resources-${serverName}`,
          method: 'resources/list',
        };
        let response: JsonRpcResponse;
        if (this.onUpstreamCall) {
          const start = performance.now();
          response = await upstream.send(req);
          const durationMs = Math.round(performance.now() - start);
          this.onUpstreamCall({ upstream: serverName, method: req.method, request: req, response, durationMs, ...(correlationId ? { correlationId } : {}) });
        } else {
          response = await upstream.send(req);
        }
        return { serverName, response };
      }),
    );

    for (const result of results) {
      if (result.status === 'rejected') continue;
      const { serverName } = result.value;

      let response: JsonRpcResponse | null = null;
      if ('cached' in result.value) {
        const cached = result.value.cached;
        if (cached.kind === 'err') continue;
        response = { jsonrpc: '2.0', id: `cached-${serverName}`, result: cached.result };
      } else {
        response = result.value.response;
        if (response.error) {
          const message = (response.error as { message?: string }).message ?? 'unknown error';
          this.setListCacheEntry(serverName, 'resources/list', { kind: 'err', message, fetchedAt: Date.now() });
          continue;
        }
        if (response.result !== undefined) {
          this.setListCacheEntry(serverName, 'resources/list', { kind: 'ok', result: response.result, fetchedAt: Date.now() });
        }
      }

      if (response.result && typeof response.result === 'object' && 'resources' in response.result) {
        const resources = (response.result as { resources: Array<{ uri: string; name?: string; description?: string; mimeType?: string }> }).resources;
        for (const resource of resources) {
          const namespacedUri = `${serverName}://${resource.uri}`;
          this.resourceToServer.set(namespacedUri, serverName);
          allResources.push({
            ...resource,
            uri: namespacedUri,
          });
        }
      }
    }

    return allResources;
  }

  /**
   * Discover prompts from all upstreams by calling prompts/list on each.
   */
  async discoverPrompts(correlationId?: string): Promise<Array<{ name: string; description?: string; arguments?: unknown[] }>> {
    const allPrompts: Array<{ name: string; description?: string; arguments?: unknown[] }> = [];

    for (const [serverName, upstream] of this.upstreams) {
      try {
        const req = {
          jsonrpc: '2.0' as const,
          id: `discover-prompts-${serverName}`,
          method: 'prompts/list',
        };
        let response: JsonRpcResponse;
        if (this.onUpstreamCall) {
          const start = performance.now();
          response = await upstream.send(req);
          const durationMs = Math.round(performance.now() - start);
          this.onUpstreamCall({ upstream: serverName, method: req.method, request: req, response, durationMs, ...(correlationId ? { correlationId } : {}) });
        } else {
          response = await upstream.send(req);
        }

        if (response.result && typeof response.result === 'object' && 'prompts' in response.result) {
          const prompts = (response.result as { prompts: Array<{ name: string; description?: string; arguments?: unknown[] }> }).prompts;
          for (const prompt of prompts) {
            const namespacedName = `${serverName}/${prompt.name}`;
            this.promptToServer.set(namespacedName, serverName);
            allPrompts.push({
              ...prompt,
              name: namespacedName,
            });
          }
        }
      } catch {
        // Server may be unavailable; skip its prompts
      }
    }

    return allPrompts;
  }

  /**
   * Route a namespaced call to the correct upstream, stripping the namespace prefix.
   */
  private async routeNamespacedCall(
    request: JsonRpcRequest,
    nameField: string,
    routingMap: Map<string, string>,
  ): Promise<JsonRpcResponse> {
    const params = request.params as Record<string, unknown> | undefined;
    const name = params?.[nameField] as string | undefined;
    if (!name) {
      return {
        jsonrpc: '2.0',
        id: request.id,
        error: { code: -32602, message: `Missing ${nameField} in params` },
      };
    }

    const serverName = routingMap.get(name);
    if (!serverName) {
      return {
        jsonrpc: '2.0',
        id: request.id,
        error: { code: -32601, message: `Unknown ${nameField}: ${name}` },
      };
    }

    const upstream = this.upstreams.get(serverName);
    if (!upstream || !upstream.isAlive()) {
      return {
        jsonrpc: '2.0',
        id: request.id,
        error: { code: -32603, message: `Upstream '${serverName}' is not available` },
      };
    }

    // Strip the namespace prefix
    const originalName = nameField === 'uri'
      ? name.slice(`${serverName}://`.length)
      : name.slice(serverName.length + 1);

    const upstreamRequest: JsonRpcRequest = {
      ...request,
      params: {
        ...params,
        [nameField]: originalName,
      },
    };

    if (this.onUpstreamCall) {
      const start = performance.now();
      const response = await upstream.send(upstreamRequest);
      const durationMs = Math.round(performance.now() - start);
      this.onUpstreamCall({ upstream: serverName, method: request.method, request: upstreamRequest, response, durationMs });
      return response;
    }

    return upstream.send(upstreamRequest);
  }

  /**
   * Route a generic request. Handles protocol-level methods locally,
   * delegates tool/resource/prompt calls to upstreams.
   *
   * When a plugin is set, plugin hooks are dispatched at each stage.
   * Without a plugin, the router is a transparent MCP proxy.
   */
  async route(request: JsonRpcRequest, context?: RouteContext): Promise<JsonRpcResponse> {
    switch (request.method) {
      case 'initialize': {
        if (this.plugin && context?.sessionId) {
          const ctx = await this.getOrCreatePluginContext(context.sessionId);

          let instructions = this.instructions ?? '';
          if (this.plugin.onInitialize) {
            const hookResult = await this.plugin.onInitialize(request, ctx);
            if (hookResult?.instructions) {
              instructions = instructions ? `${instructions}\n${hookResult.instructions}` : hookResult.instructions;
            }
          }

          const result: Record<string, unknown> = {
            protocolVersion: '2024-11-05',
            serverInfo: { name: 'mcpctl-proxy', version: '0.0.1' },
            capabilities: { tools: { listChanged: true }, resources: {}, prompts: {} },
          };
          if (instructions) result['instructions'] = instructions;
          return { jsonrpc: '2.0', id: request.id, result };
        }

        // No plugin: transparent proxy — return basic init response
        const result: Record<string, unknown> = {
          protocolVersion: '2024-11-05',
          serverInfo: { name: 'mcpctl-proxy', version: '0.0.1' },
          capabilities: { tools: { listChanged: true }, resources: {}, prompts: {} },
        };
        if (this.instructions) result['instructions'] = this.instructions;
        return { jsonrpc: '2.0', id: request.id, result };
      }

      case 'tools/list': {
        if (this.plugin && context?.sessionId) {
          const ctx = await this.getOrCreatePluginContext(context.sessionId);
          let tools = await this.discoverTools(context?.correlationId);

          if (this.plugin.onToolsList) {
            tools = await this.plugin.onToolsList(tools, ctx);
          }

          return { jsonrpc: '2.0', id: request.id, result: { tools } };
        }

        // No plugin: return upstream tools only
        const tools = await this.discoverTools(context?.correlationId);
        return { jsonrpc: '2.0', id: request.id, result: { tools } };
      }

      case 'tools/call':
        return this.routeToolCall(request, context);

      case 'resources/list': {
        if (this.plugin?.onResourcesList && context?.sessionId) {
          const ctx = await this.getOrCreatePluginContext(context.sessionId);
          const resources = await this.discoverResources(context?.correlationId);
          const filtered = await this.plugin.onResourcesList(resources, ctx);
          return { jsonrpc: '2.0', id: request.id, result: { resources: filtered } };
        }

        const resources = await this.discoverResources(context?.correlationId);
        // Append mcpctl prompt resources
        const mcpdResources: Array<{ uri: string; name: string; description: string; mimeType: string }> = [];
        if (this.mcpdClient && this.projectName) {
          try {
            const sessionParam = context?.sessionId ? `?session=${encodeURIComponent(context.sessionId)}` : '';
            const visible = await this.mcpdClient.get<Array<{ name: string; content: string; type: string }>>(
              `/api/v1/projects/${encodeURIComponent(this.projectName)}/prompts/visible${sessionParam}`,
            );
            this.mcpctlResourceContents.clear();
            for (const p of visible) {
              const uri = `mcpctl://prompts/${p.name}`;
              const res = {
                uri,
                name: p.name,
                description: p.type === 'promptrequest' ? `[Pending proposal] ${p.name}` : `[Approved prompt] ${p.name}`,
                mimeType: 'text/plain',
              };
              resources.push(res);
              mcpdResources.push(res);
              this.mcpctlResourceContents.set(uri, p.content);
            }
          } catch {
            // Prompt resources are optional — don't fail discovery
          }
        }
        // Emit upstream event for mcpd-sourced resources so provenance view shows them
        if (mcpdResources.length > 0 && this.onUpstreamCall) {
          const mcpdResponse = { result: { resources: mcpdResources } };
          this.onUpstreamCall({
            upstream: 'mcpd',
            method: 'resources/list',
            request: { jsonrpc: '2.0', id: request.id, method: 'resources/list' },
            response: mcpdResponse,
            durationMs: 0,
            ...(context?.correlationId ? { correlationId: context.correlationId } : {}),
          });
        }
        return {
          jsonrpc: '2.0',
          id: request.id,
          result: { resources },
        };
      }

      case 'resources/read': {
        const params = request.params as Record<string, unknown> | undefined;
        const uri = params?.['uri'] as string | undefined;
        if (uri?.startsWith('mcpctl://prompts/') && this.mcpdClient && this.projectName) {
          const promptName = uri.slice('mcpctl://prompts/'.length);
          try {
            const sessionParam = context?.sessionId ? `?session=${encodeURIComponent(context.sessionId)}` : '';
            const visible = await this.mcpdClient.get<Array<{ name: string; content: string; type: string }>>(
              `/api/v1/projects/${encodeURIComponent(this.projectName)}/prompts/visible${sessionParam}`,
            );
            const found = visible.find((p) => p.name === promptName);
            if (found) {
              this.mcpctlResourceContents.set(uri, found.content);
              return {
                jsonrpc: '2.0',
                id: request.id,
                result: {
                  contents: [{ uri, mimeType: 'text/plain', text: found.content }],
                },
              };
            }
          } catch {
            // Fall through to cache
          }
          // Fallback to cache if mcpd is unreachable
          const cached = this.mcpctlResourceContents.get(uri);
          if (cached !== undefined) {
            return {
              jsonrpc: '2.0',
              id: request.id,
              result: {
                contents: [{ uri, mimeType: 'text/plain', text: cached }],
              },
            };
          }
          return {
            jsonrpc: '2.0',
            id: request.id,
            error: { code: -32602, message: `Resource not found: ${uri}` },
          };
        }
        if (uri?.startsWith('mcpctl://')) {
          const content = this.mcpctlResourceContents.get(uri);
          if (content !== undefined) {
            return {
              jsonrpc: '2.0',
              id: request.id,
              result: {
                contents: [{ uri, mimeType: 'text/plain', text: content }],
              },
            };
          }
          return {
            jsonrpc: '2.0',
            id: request.id,
            error: { code: -32602, message: `Resource not found: ${uri}` },
          };
        }
        return this.routeNamespacedCall(request, 'uri', this.resourceToServer);
      }

      case 'resources/subscribe':
      case 'resources/unsubscribe':
        return this.routeNamespacedCall(request, 'uri', this.resourceToServer);

      case 'prompts/list': {
        if (this.plugin?.onPromptsList && context?.sessionId) {
          const ctx = await this.getOrCreatePluginContext(context.sessionId);
          const upstreamPrompts = await this.discoverPrompts(context?.correlationId);
          const filtered = await this.plugin.onPromptsList(upstreamPrompts, ctx);
          return { jsonrpc: '2.0', id: request.id, result: { prompts: filtered } };
        }

        const upstreamPrompts = await this.discoverPrompts(context?.correlationId);
        // Include mcpctl-managed prompts from mcpd alongside upstream prompts
        const managedIndex = await this.fetchPromptIndex();
        const managedPrompts = managedIndex.map((p) => ({
          name: `mcpctl/${p.name}`,
          description: p.summary ?? `Priority ${p.priority} prompt`,
        }));
        // Emit upstream event for mcpd-sourced prompts so provenance view shows them
        if (managedPrompts.length > 0 && this.onUpstreamCall) {
          const mcpdResponse = { result: { prompts: managedPrompts } };
          this.onUpstreamCall({
            upstream: 'mcpd',
            method: 'prompts/list',
            request: { jsonrpc: '2.0', id: request.id, method: 'prompts/list' },
            response: mcpdResponse,
            durationMs: 0,
            ...(context?.correlationId ? { correlationId: context.correlationId } : {}),
          });
        }
        return {
          jsonrpc: '2.0',
          id: request.id,
          result: { prompts: [...upstreamPrompts, ...managedPrompts] },
        };
      }

      case 'prompts/get': {
        const promptName = (request.params as Record<string, unknown> | undefined)?.name as string | undefined;

        // Plugin hook: onPromptGet (section drill-down, etc.)
        if (this.plugin?.onPromptGet && context?.sessionId && promptName) {
          const ctx = await this.getOrCreatePluginContext(context.sessionId);
          const intercepted = await this.plugin.onPromptGet(promptName, request, ctx);
          if (intercepted) return intercepted;
        }

        let response: JsonRpcResponse;
        if (promptName?.startsWith('mcpctl/')) {
          const shortName = promptName.slice('mcpctl/'.length);
          const managedIndex = await this.fetchPromptIndex();
          const entry = managedIndex.find((p) => p.name === shortName);
          if (!entry) {
            return { jsonrpc: '2.0', id: request.id, error: { code: -32601, message: `Unknown name: ${promptName}` } };
          }
          response = {
            jsonrpc: '2.0',
            id: request.id,
            result: {
              prompt: {
                name: promptName,
                description: entry.summary ?? `Priority ${entry.priority} prompt`,
              },
              messages: [
                {
                  role: 'user',
                  content: { type: 'text', text: entry.content || '(empty)' },
                },
              ],
            },
          };
        } else {
          response = await this.routeNamespacedCall(request, 'name', this.promptToServer);
        }

        // Post-process: run large prompt content through the pipeline for section-split
        if (!response.error && this.plugin && context?.sessionId) {
          const ctx = await this.getOrCreatePluginContext(context.sessionId);
          const text = extractPromptText(response);
          if (text && text.length > 2000) {
            try {
              const result = await ctx.processContent(promptName ?? 'prompt', text, 'prompt');
              if (result.sections && result.sections.length > 0) {
                response = injectSectionsIntoPromptResponse(response, result.content, result.sections, ctx);
              } else if (result.content !== text) {
                response = replacePromptText(response, result.content);
              }
            } catch {
              // Pipeline failed — return original
            }
          }
        }

        return response;
      }

      // Handle MCP notifications (no response expected, but return empty result if called as request)
      case 'notifications/initialized':
      case 'notifications/cancelled':
      case 'notifications/progress':
      case 'notifications/roots/list_changed':
        return {
          jsonrpc: '2.0',
          id: request.id,
          result: {},
        };

      default:
        return {
          jsonrpc: '2.0',
          id: request.id,
          error: { code: -32601, message: `Method not found: ${request.method}` },
        };
    }
  }

  /**
   * Route a tools/call request.
   *
   * Plugin path: virtual tools → onToolCallBefore → upstream → onToolCallAfter.
   * No-plugin path: pagination → LLM processing → upstream.
   */
  private async routeToolCall(request: JsonRpcRequest, context?: RouteContext): Promise<JsonRpcResponse> {
    const params = request.params as Record<string, unknown> | undefined;
    const toolName = params?.['name'] as string | undefined;
    const toolArgs = (params?.['arguments'] ?? {}) as Record<string, unknown>;
    const startMs = Date.now();

    const emitTrace = (response: JsonRpcResponse, serverName?: string): void => {
      if (this.auditCollector && context?.sessionId && toolName) {
        const durationMs = Date.now() - startMs;
        const resultSize = JSON.stringify(response.result ?? response.error ?? {}).length;
        const resolved = serverName ?? this.toolToServer.get(toolName);
        const base = {
          timestamp: new Date().toISOString(),
          sessionId: context.sessionId,
          eventKind: 'tool_call_trace' as const,
          source: 'mcplocal' as const,
          verified: true,
          // Already in scope and previously dropped, which left every durable
          // tool_call_trace row unjoinable to the rest of its request.
          ...(context.correlationId !== undefined ? { correlationId: context.correlationId } : {}),
          payload: {
            toolName,
            argKeys: Object.keys(toolArgs).join(', '),
            resultSizeBytes: resultSize,
            durationMs,
            error: response.error?.message ?? null,
          },
        };
        this.auditCollector.emit(resolved ? { ...base, serverName: resolved } : base);
      }
    };

    // Plugin path
    if (this.plugin && context?.sessionId) {
      const ctx = await this.getOrCreatePluginContext(context.sessionId);

      // Check virtual tools first (registered by plugin via ctx.registerTool)
      const virtualTool = ctx.virtualTools.get(toolName ?? '');
      if (virtualTool) {
        try {
          const result = await virtualTool.handler(toolArgs, ctx);
          const resp = { jsonrpc: '2.0' as const, id: request.id, result: result as JsonRpcResponse['result'] };
          emitTrace(resp, 'virtual');
          return resp;
        } catch (err) {
          const code = (err as { code?: number }).code ?? -32603;
          const resp = { jsonrpc: '2.0' as const, id: request.id, error: { code, message: err instanceof Error ? err.message : String(err) } };
          emitTrace(resp, 'virtual');
          return resp;
        }
      }

      // onToolCallBefore — can intercept and return a response directly
      if (this.plugin.onToolCallBefore) {
        const intercepted = await this.plugin.onToolCallBefore(toolName ?? '', toolArgs, request, ctx);
        if (intercepted) {
          // Intercepts were previously invisible in tool_call_trace, so every
          // call made while the session was gated produced no trace at all.
          emitTrace(intercepted, 'plugin');
          return intercepted;
        }
      }

      // Route to upstream
      let response = await this.routeNamespacedCall(request, 'name', this.toolToServer);

      // onToolCallAfter — can transform the response
      if (this.plugin.onToolCallAfter && toolName) {
        response = await this.plugin.onToolCallAfter(toolName, toolArgs, response, ctx);
      }

      emitTrace(response);
      return response;
    }

    // No plugin: transparent proxy with pagination and LLM processing

    // Intercept pagination page requests
    if (this.paginator) {
      const paginationReq = ResponsePaginator.extractPaginationParams(toolArgs);
      if (paginationReq) {
        const pageResult = this.paginator.getPage(paginationReq.resultId, paginationReq.page);
        if (pageResult) {
          return { jsonrpc: '2.0', id: request.id, result: pageResult };
        }
        return {
          jsonrpc: '2.0',
          id: request.id,
          result: {
            content: [{
              type: 'text',
              text: 'Cached result not found (expired or invalid _resultId). Please re-call the tool without _resultId/_page to get a fresh result.',
            }],
          },
        };
      }
    }

    // If no LLM processor or tool shouldn't be processed, route directly
    if (!this.llmProcessor || !toolName || !this.llmProcessor.shouldProcess('tools/call', toolName)) {
      const response = await this.routeNamespacedCall(request, 'name', this.toolToServer);
      const paginated = await this.maybePaginate(toolName, response);
      emitTrace(paginated);
      return paginated;
    }

    // Preprocess request params
    const processed = await this.llmProcessor.preprocessRequest(toolName, toolArgs);
    const processedRequest: JsonRpcRequest = processed.optimized
      ? { ...request, params: { ...params, arguments: processed.params } }
      : request;

    // Route to upstream
    const response = await this.routeNamespacedCall(processedRequest, 'name', this.toolToServer);

    // Try pagination
    const paginated = await this.maybePaginate(toolName, response);
    if (paginated !== response) {
      emitTrace(paginated);
      return paginated;
    }

    // Filter response
    if (response.error) {
      emitTrace(response);
      return response;
    }
    const filtered = await this.llmProcessor.filterResponse(toolName, response);
    if (filtered.filtered) {
      const filteredResp = { ...response, result: filtered.result };
      emitTrace(filteredResp);
      return filteredResp;
    }
    emitTrace(response);
    return response;
  }

  /**
   * Apply pagination to a tool response if it exceeds the size threshold.
   */
  private async maybePaginate(toolName: string | undefined, response: JsonRpcResponse): Promise<JsonRpcResponse> {
    if (!toolName || !this.paginator || response.error) return response;
    const raw = JSON.stringify(response.result);
    if (!this.paginator.shouldPaginate(raw)) return response;
    const paginated = await this.paginator.paginate(toolName, raw);
    if (!paginated) return response;
    return { jsonrpc: '2.0', id: response.id, result: paginated };
  }

  // ── Prompt and system prompt fetchers (used by plugin context) ──

  /**
   * Fetch prompt index from mcpd with caching.
   */
  private async fetchPromptIndex(): Promise<PromptIndexEntry[]> {
    const now = Date.now();
    if (this.cachedPromptIndex && (now - this.promptIndexFetchedAt) < this.PROMPT_INDEX_TTL_MS) {
      return this.cachedPromptIndex;
    }

    if (!this.mcpdClient || !this.projectName) {
      return [];
    }

    const index = await this.mcpdClient.get<Array<{
      name: string;
      priority: number;
      summary: string | null;
      chapters: string[] | null;
      content?: string;
      linkTarget?: string | null;
    }>>(
      `/api/v1/projects/${encodeURIComponent(this.projectName)}/prompts/visible`,
    );

    // Resolve linked prompts: fetch fresh content from linked MCP resources
    const entries: PromptIndexEntry[] = [];
    for (const p of index) {
      let content = p.content ?? '';
      if (p.linkTarget && this.linkResolver) {
        try {
          const resolution = await this.linkResolver.resolve(p.linkTarget);
          if (resolution.status === 'alive' && resolution.content) {
            content = resolution.content;
          }
        } catch {
          // Keep static content as fallback
        }
      }
      entries.push({
        name: p.name,
        priority: p.priority,
        summary: p.summary,
        chapters: p.chapters,
        content,
      });
    }

    this.cachedPromptIndex = entries;
    this.promptIndexFetchedAt = now;
    return this.cachedPromptIndex;
  }

  /**
   * Fetch a system prompt from mcpctl-system project, with caching and fallback.
   */
  private async getSystemPrompt(name: string, fallback: string): Promise<string> {
    const now = Date.now();
    const cached = this.systemPromptCache.get(name);
    if (cached && (now - cached.fetchedAt) < this.SYSTEM_PROMPT_TTL_MS) {
      return cached.content;
    }

    if (!this.mcpdClient) return fallback;

    try {
      const visible = await this.mcpdClient.get<Array<{ name: string; content: string }>>(
        '/api/v1/projects/mcpctl-system/prompts/visible',
      );
      // Cache all system prompts from the response
      for (const p of visible) {
        this.systemPromptCache.set(p.name, { content: p.content, fetchedAt: now });
      }
      const found = visible.find((p) => p.name === name);
      return found?.content ?? fallback;
    } catch {
      return fallback;
    }
  }

  // ── Notification queue ──

  private queueNotification(sessionId: string | undefined, notification: JsonRpcNotification): void {
    if (!sessionId) return;
    const queue = this.pendingNotifications.get(sessionId) ?? [];
    queue.push(notification);
    this.pendingNotifications.set(sessionId, queue);
  }

  /** Consume and return any pending notifications for a session (e.g., tools/list_changed after ungating). */
  consumeNotifications(sessionId: string): JsonRpcNotification[] {
    const notifications = this.pendingNotifications.get(sessionId) ?? [];
    this.pendingNotifications.delete(sessionId);
    return notifications;
  }

  // ── Session cleanup ──

  cleanupSession(sessionId: string): void {
    const pluginCtx = this.pluginContexts.get(sessionId);
    if (pluginCtx && this.plugin?.onSessionDestroy) {
      void this.plugin.onSessionDestroy(pluginCtx);
    }
    this.pluginContexts.delete(sessionId);
    this.pendingNotifications.delete(sessionId);
  }

  getUpstreamNames(): string[] {
    return [...this.upstreams.keys()];
  }

  getProxyModelInfo(): { projectDefault: string | null; serverOverrides: Record<string, string> } {
    const serverOverrides: Record<string, string> = {};
    for (const [server, config] of this.serverProxyModels) {
      serverOverrides[server] = config.name;
    }
    return { projectDefault: this.proxyModelName, serverOverrides };
  }

  async closeAll(): Promise<void> {
    for (const upstream of this.upstreams.values()) {
      await upstream.close();
    }
    this.upstreams.clear();
    this.toolToServer.clear();
    this.resourceToServer.clear();
    this.promptToServer.clear();
    this.pluginContexts.clear();
  }
}

// ── Prompt section drill-down helpers ──

// Same key used by content-pipeline plugin so drill-down lookups work across tools and prompts
const SECTION_STORE_KEY = '_contentPipeline_sections';
const SECTION_TTL_MS = 300_000; // 5 minutes

interface SectionStoreEntry {
  sections: Section[];
  createdAt: number;
}

/** Extract text from all messages in a prompts/get response. */
function extractPromptText(response: JsonRpcResponse): string | null {
  if (!response.result || typeof response.result !== 'object') return null;
  const result = response.result as Record<string, unknown>;
  const messages = result['messages'] as Array<{ role: string; content: unknown }> | undefined;
  if (!Array.isArray(messages)) return null;
  const texts: string[] = [];
  for (const msg of messages) {
    if (typeof msg.content === 'string') {
      texts.push(msg.content);
    } else if (msg.content && typeof msg.content === 'object') {
      const c = msg.content as { type?: string; text?: string };
      if (c.type === 'text' && c.text) texts.push(c.text);
    }
  }
  return texts.length > 0 ? texts.join('\n') : null;
}

/** Replace prompt message text with processed content. */
function replacePromptText(response: JsonRpcResponse, newText: string): JsonRpcResponse {
  const result = response.result as Record<string, unknown>;
  const messages = result['messages'] as Array<{ role: string; content: unknown }>;
  const newMessages = messages.map((msg) => {
    if (typeof msg.content === 'string') {
      return { ...msg, content: newText };
    }
    if (msg.content && typeof msg.content === 'object') {
      return { ...msg, content: { type: 'text', text: newText } };
    }
    return msg;
  });
  return { ...response, result: { ...result, messages: newMessages } };
}

/** Inject sections TOC into prompt response and store sections for drill-down. */
function injectSectionsIntoPromptResponse(
  response: JsonRpcResponse,
  tocContent: string,
  sections: Section[],
  ctx: PluginSessionContext,
): JsonRpcResponse {
  const resultId = `pm-${Date.now().toString(36)}`;

  // Store sections (same store as tool results so drill-down works uniformly)
  let store = ctx.state.get(SECTION_STORE_KEY) as Map<string, SectionStoreEntry> | undefined;
  if (!store) {
    store = new Map();
    ctx.state.set(SECTION_STORE_KEY, store);
  }
  store.set(resultId, { sections, createdAt: Date.now() });
  // Evict stale
  const now = Date.now();
  for (const [key, entry] of store) {
    if (now - entry.createdAt > SECTION_TTL_MS) store.delete(key);
  }

  const text = `${tocContent}\n\nTo read one, request this prompt again with _resultId="${resultId}" and _section="<id>". Both parameters are required together.`;
  return replacePromptText(response, text);
}