mcpctl/src/mcplocal/src/proxymodel/plugin-context.ts

/**
 * PluginSessionContext implementation.
 *
 * Wraps router internals and exposes them via the clean PluginSessionContext
 * interface. Each session gets its own context instance.
 */
import type { JsonRpcRequest, JsonRpcResponse, JsonRpcNotification } from '../types.js';
import type { LLMProvider, CacheProvider, StageLogger, Section, ToolDefinition, ContentType } from './types.js';
import type { PluginSessionContext, VirtualToolHandler, VirtualServer, PromptIndexEntry } from './plugin.js';
import type { AuditCollector } from '../audit/collector.js';
import type { AuditEvent } from '../audit/types.js';
import { currentCorrelationId } from '../request-context.js';

/** Dependencies injected from the router into each context. */
export interface PluginContextDeps {
  sessionId: string;
  projectName: string;
  llm: LLMProvider;
  cache: CacheProvider;
  log: StageLogger;
  discoverTools: () => Promise<ToolDefinition[]>;
  getToolDiscoveryFailures?: () => Array<{ server: string; message: string }>;
  routeToUpstream: (request: JsonRpcRequest) => Promise<JsonRpcResponse>;
  fetchPromptIndex: () => Promise<PromptIndexEntry[]>;
  getSystemPrompt: (name: string, fallback: string) => Promise<string>;
  processContent: (toolName: string, content: string, contentType: ContentType) => Promise<{ content: string; sections?: Section[] }>;
  queueNotification: (notification: JsonRpcNotification) => void;
  postToMcpd: (path: string, body: Record<string, unknown>) => Promise<unknown>;
  getFromMcpd: (path: string) => Promise<unknown>;
  auditCollector?: AuditCollector;
}

/**
 * Concrete PluginSessionContext. One per session.
 */
export class PluginContextImpl implements PluginSessionContext {
  readonly sessionId: string;
  readonly projectName: string;
  readonly state = new Map<string, unknown>();
  readonly llm: LLMProvider;
  readonly cache: CacheProvider;
  readonly log: StageLogger;

  /** Virtual tools registered by plugins (name → handler). Tool name is non-namespaced. */
  readonly virtualTools = new Map<string, { definition: ToolDefinition; handler: VirtualToolHandler }>();

  /** Virtual servers registered by plugins. */
  readonly virtualServers = new Map<string, VirtualServer>();

  private readonly deps: PluginContextDeps;

  constructor(deps: PluginContextDeps) {
    this.sessionId = deps.sessionId;
    this.projectName = deps.projectName;
    this.llm = deps.llm;
    this.cache = deps.cache;
    this.log = deps.log;
    this.deps = deps;
  }

  registerTool(tool: ToolDefinition, handler: VirtualToolHandler): void {
    this.virtualTools.set(tool.name, { definition: tool, handler });
  }

  unregisterTool(name: string): void {
    this.virtualTools.delete(name);
  }

  registerServer(server: VirtualServer): void {
    this.virtualServers.set(server.name, server);
    // Register each tool under the server namespace
    for (const t of server.tools) {
      const namespacedName = `${server.name}/${t.definition.name}`;
      this.virtualTools.set(namespacedName, { definition: { ...t.definition, name: namespacedName }, handler: t.handler });
    }
  }

  unregisterServer(name: string): void {
    const server = this.virtualServers.get(name);
    if (server) {
      for (const t of server.tools) {
        this.virtualTools.delete(`${name}/${t.definition.name}`);
      }
      this.virtualServers.delete(name);
    }
  }

  queueNotification(method: string, params?: unknown): void {
    const notification: JsonRpcNotification = { jsonrpc: '2.0', method };
    if (params !== undefined) {
      notification.params = params as Record<string, unknown>;
    }
    this.deps.queueNotification(notification);
  }

  discoverTools(): Promise<ToolDefinition[]> {
    return this.deps.discoverTools();
  }

  getToolDiscoveryFailures(): Array<{ server: string; message: string }> {
    return this.deps.getToolDiscoveryFailures?.() ?? [];
  }

  routeToUpstream(request: JsonRpcRequest): Promise<JsonRpcResponse> {
    return this.deps.routeToUpstream(request);
  }

  fetchPromptIndex(): Promise<PromptIndexEntry[]> {
    return this.deps.fetchPromptIndex();
  }

  getSystemPrompt(name: string, fallback: string): Promise<string> {
    return this.deps.getSystemPrompt(name, fallback);
  }

  processContent(toolName: string, content: string, contentType: ContentType): Promise<{ content: string; sections?: Section[] }> {
    return this.deps.processContent(toolName, content, contentType);
  }

  postToMcpd(path: string, body: Record<string, unknown>): Promise<unknown> {
    return this.deps.postToMcpd(path, body);
  }

  getFromMcpd(path: string): Promise<unknown> {
    return this.deps.getFromMcpd(path);
  }

  /** Emit an audit event, auto-filling sessionId, projectName and correlationId. */
  emitAuditEvent(event: Omit<AuditEvent, 'sessionId' | 'projectName'>): void {
    // Auto-filled here rather than at each call site: this covers every gate
    // event and every plugin written later, with no per-callsite edit.
    const correlationId = event.correlationId ?? currentCorrelationId();
    this.deps.auditCollector?.emit({
      ...event,
      sessionId: this.sessionId,
      ...(correlationId !== undefined ? { correlationId } : {}),
    });
  }
}