mcpctl/src/mcpd/src/services/instance.service.ts

import type { McpInstance, McpServer } from '@prisma/client';
import type { IMcpInstanceRepository, IMcpServerRepository } from '../repositories/interfaces.js';
import type { McpOrchestrator, ContainerSpec, ContainerInfo } from './orchestrator.js';
import { NotFoundError } from './mcp-server.service.js';
import {
  resolveServerEnv,
  partitionServerEnv,
  validateServerEnvRefs,
  type SecretResolver,
} from './env-resolver.js';

/**
 * What InstanceService needs from the per-server identity layer. Kept narrow
 * and local so the orchestration stays testable without an OpenBao backend.
 */
export interface ServerIdentityProvisioner {
  /** Converge the server's identity; resolves to the ServiceAccount/role name. */
  ensureFor(server: McpServer): Promise<string>;
  /** Injector pod annotations for these refs under that identity. */
  annotationsFor(
    identity: string,
    refs: Array<{ name: string; secretName: string; key: string }>,
  ): Record<string, string>;
  /** Rewrite argv to source the rendered secrets before exec'ing the server. */
  wrapCommand(server: McpServer, command: string[] | undefined): string[] | undefined;
}

/** Runner images for package-based MCP servers, keyed by runtime name. */
const RUNNER_IMAGES: Record<string, string> = {
  node: process.env['MCPD_NODE_RUNNER_IMAGE'] ?? 'mysources.co.uk/michal/mcpctl-node-runner:latest',
  python: process.env['MCPD_PYTHON_RUNNER_IMAGE'] ?? 'mysources.co.uk/michal/mcpctl-python-runner:latest',
};

/** Network for MCP server containers (matches docker-compose mcp-servers network). */
const MCP_SERVERS_NETWORK = process.env['MCPD_MCP_NETWORK'] ?? 'mcp-servers';

/**
 * Backoff schedule for instance startup failures (env resolution, container
 * creation, etc). Mirrors Kubernetes-style escalation: fast retries for
 * transient hiccups, then a longer pause once it's clear something is
 * persistently wrong.
 *
 * The retry state lives on `McpInstance.metadata` (no schema migration
 * needed) and is preserved across reconcile cycles by the in-place
 * `retryInstance` path so attemptCount actually accumulates.
 */
const FAST_RETRY_MS = 30_000;          // first 5 attempts: 30s apart
const SLOW_RETRY_MS = 5 * 60_000;      // afterwards: 5 minutes
const MAX_FAST_RETRIES = 5;

interface RetryMetadata {
  error?: string;
  attemptCount?: number;
  lastAttemptAt?: string;
  nextRetryAt?: string;
  /** containerStatuses[0].restartCount at last sync — a bump means every
   *  cached STDIO pipe to this (unchanged) containerId is dead. */
  lastRestartCount?: number;
  [k: string]: unknown;
}

function readRetryMeta(instance: McpInstance): RetryMetadata {
  return (instance.metadata ?? {}) as RetryMetadata;
}

function nextDelayMs(attemptCount: number): number {
  return attemptCount <= MAX_FAST_RETRIES ? FAST_RETRY_MS : SLOW_RETRY_MS;
}

export class InvalidStateError extends Error {
  readonly statusCode = 409;
  constructor(message: string) {
    super(message);
    this.name = 'InvalidStateError';
  }
}

export class InstanceService {
  private stdioInvalidator?: (containerId: string) => void;

  constructor(
    private instanceRepo: IMcpInstanceRepository,
    private serverRepo: IMcpServerRepository,
    private orchestrator: McpOrchestrator,
    private secretResolver?: SecretResolver,
    /** Provisions per-server OpenBao identities. Absent = injector unavailable. */
    private serverIdentity?: ServerIdentityProvisioner,
  ) {}

  /**
   * Hook for evicting cached STDIO clients (McpProxyService.removeClient).
   * Setter injection, matching serverService.setInstanceService in main.ts —
   * McpProxyService is constructed after this service and already imports
   * from this file, so a constructor arg would be a circular import.
   */
  setStdioInvalidator(fn: (containerId: string) => void): void {
    this.stdioInvalidator = fn;
  }

  private invalidateStdio(containerId: string | null | undefined): void {
    if (!containerId) return;
    try {
      this.stdioInvalidator?.(containerId);
    } catch {
      /* best-effort */
    }
  }

  async list(serverId?: string): Promise<McpInstance[]> {
    return this.instanceRepo.findAll(serverId);
  }

  async getById(id: string): Promise<McpInstance> {
    const instance = await this.instanceRepo.findById(id);
    if (!instance) throw new NotFoundError(`Instance '${id}' not found`);
    return instance;
  }

  /**
   * Sync instance statuses with actual container state.
   *
   * Beyond marking crashed containers ERROR, this is the recovery path for
   * mcpctl#114: ERROR rows whose pod came back are re-adopted (the pod has
   * restartPolicy Always, so kubelet restarts the container in place and the
   * row must follow it back), and an in-place restart under an unchanged
   * containerId — visible only as a restartCount bump — invalidates any
   * cached STDIO pipe, which is dead by definition.
   *
   * Every metadata write here MERGES via readRetryMeta: updateStatus replaces
   * the JSON column wholesale, and clobbering nextRetryAt is what used to
   * make ERROR rows instantly dueForRetry and hot-loop against a 409.
   */
  async syncStatus(): Promise<void> {
    const instances = await this.instanceRepo.findAll();
    for (const inst of instances) {
      if (!inst.containerId) continue;
      if (inst.status !== 'RUNNING' && inst.status !== 'STARTING' && inst.status !== 'ERROR') {
        continue;
      }

      let info: ContainerInfo;
      try {
        info = await this.orchestrator.inspectContainer(inst.containerId);
      } catch {
        // Container gone entirely. ERROR rows with a missing pod stay as they
        // are — the retry/backoff path owns recreating them.
        if (inst.status !== 'ERROR') {
          await this.instanceRepo.updateStatus(inst.id, 'ERROR', {
            metadata: { ...readRetryMeta(inst), error: 'Container not found' },
          });
          this.invalidateStdio(inst.containerId);
        }
        continue;
      }

      const meta = readRetryMeta(inst);

      if (inst.status === 'ERROR') {
        // The pod outlived the ERROR verdict (kubelet restarted the container
        // in place). Re-adopt instead of leaving the row stuck forever.
        if (info.state === 'running') {
          const { error: _e, attemptCount: _a, lastAttemptAt: _l, nextRetryAt: _n, ...rest } = meta;
          await this.instanceRepo.updateStatus(inst.id, 'RUNNING', {
            metadata: { ...rest, lastRestartCount: info.restartCount ?? 0 },
          });
          this.invalidateStdio(inst.containerId);
        } else if (info.state === 'starting') {
          // Keep retry metadata until it is actually running.
          await this.instanceRepo.updateStatus(inst.id, 'STARTING', { metadata: meta });
          this.invalidateStdio(inst.containerId);
        }
        // stopped/error: leave for the backoff/retry path.
        continue;
      }

      if (info.state === 'stopped' || info.state === 'error') {
        // Container died — get last logs for error context
        let errorMsg = `Container ${info.state}`;
        try {
          const logs = await this.orchestrator.getContainerLogs(inst.containerId, { tail: 5 });
          const lastLog = (logs.stdout || logs.stderr).trim().split('\n').pop();
          if (lastLog) errorMsg = lastLog;
        } catch { /* best-effort */ }
        await this.instanceRepo.updateStatus(inst.id, 'ERROR', {
          metadata: { ...meta, error: errorMsg },
        });
        this.invalidateStdio(inst.containerId);
      } else if (info.state === 'starting' && inst.status === 'RUNNING') {
        // Pod went back to starting (e.g. CrashLoopBackOff restart)
        await this.instanceRepo.updateStatus(inst.id, 'STARTING', { metadata: meta });
        this.invalidateStdio(inst.containerId);
      } else if (info.state === 'running' && inst.status === 'STARTING') {
        // Pod became ready — promote to RUNNING and clear retry state.
        const { error: _e, attemptCount: _a, lastAttemptAt: _l, nextRetryAt: _n, ...rest } = meta;
        await this.instanceRepo.updateStatus(inst.id, 'RUNNING', {
          metadata: { ...rest, lastRestartCount: info.restartCount ?? 0 },
        });
        // A fresh/restarted pod under this name invalidates any cached pipe.
        this.invalidateStdio(inst.containerId);
      } else if (info.state === 'running' && info.restartCount !== undefined) {
        if (typeof meta.lastRestartCount === 'number' && info.restartCount > meta.lastRestartCount) {
          // In-place restart between polls: same pod name, dead pipes.
          await this.instanceRepo.updateStatus(inst.id, 'RUNNING', {
            metadata: { ...meta, lastRestartCount: info.restartCount },
          });
          this.invalidateStdio(inst.containerId);
        } else if (meta.lastRestartCount === undefined) {
          // First sighting: record the baseline without invalidating.
          await this.instanceRepo.updateStatus(inst.id, 'RUNNING', {
            metadata: { ...meta, lastRestartCount: info.restartCount },
          });
        }
      }
    }
  }

  /**
   * Reconcile instances for a server to match desired replica count.
   * - Syncs container statuses first (detect crashed containers)
   * - If fewer running instances than replicas: start new ones
   * - If more running instances than replicas: remove excess (oldest first)
   */
  async reconcile(serverId: string): Promise<McpInstance[]> {
    const server = await this.serverRepo.findById(serverId);
    if (!server) throw new NotFoundError(`McpServer '${serverId}' not found`);

    // Sync container statuses before counting active instances
    await this.syncStatus();

    const instances = await this.instanceRepo.findAll(serverId);
    const active = instances.filter((i) => i.status === 'RUNNING' || i.status === 'STARTING');
    const desired = server.replicas;

    if (active.length < desired) {
      // Scale up
      const toStart = desired - active.length;
      for (let i = 0; i < toStart; i++) {
        await this.startOne(serverId);
      }
    } else if (active.length > desired) {
      // Scale down — remove oldest first
      const excess = active
        .sort((a, b) => a.createdAt.getTime() - b.createdAt.getTime())
        .slice(0, active.length - desired);
      for (const inst of excess) {
        await this.removeOne(inst);
      }
    }

    return this.instanceRepo.findAll(serverId);
  }

  /**
   * Reconcile ALL servers — the operator loop.
   *
   * For every server with replicas > 0, ensures the correct number of
   * healthy instances exist. ERROR instances are not blindly recreated:
   * within their `nextRetryAt` window they're left alone (and counted
   * against the replica budget so we don't churn replacements while one
   * is in backoff); past their window they're retried in-place via
   * `retryInstance` so attemptCount accumulates and backoff escalates
   * correctly.
   */
  async reconcileAll(): Promise<{ reconciled: number; errors: string[] }> {
    await this.syncStatus();

    const servers = await this.serverRepo.findAll();
    let reconciled = 0;
    const errors: string[] = [];

    const now = Date.now();

    for (const server of servers) {
      if (server.replicas <= 0) continue;

      try {
        const instances = await this.instanceRepo.findAll(server.id);
        const active = instances.filter((i) => i.status === 'RUNNING' || i.status === 'STARTING');
        const errored = instances.filter((i) => i.status === 'ERROR');

        // Partition ERROR instances by whether their backoff window has elapsed.
        const dueForRetry: McpInstance[] = [];
        const stillWaiting: McpInstance[] = [];
        for (const inst of errored) {
          const meta = readRetryMeta(inst);
          const ts = meta.nextRetryAt ? Date.parse(meta.nextRetryAt) : 0;
          if (Number.isNaN(ts) || ts <= now) {
            dueForRetry.push(inst);
          } else {
            stillWaiting.push(inst);
          }
        }

        // Retry elapsed ones in-place. This preserves attemptCount across
        // attempts so the 30s × 5 → 5min schedule actually escalates.
        for (const inst of dueForRetry) {
          await this.retryInstance(inst);
        }

        // Scale up only if we don't already have enough live attempts.
        // Live attempts = currently-running OR -starting + still-waiting
        // (in backoff) + just-retried (now STARTING via retryInstance).
        // Counting waiting + retried against the budget prevents tight
        // create-fail-create churn while previous attempts work through
        // their backoff schedule.
        const toStart = server.replicas - active.length - stillWaiting.length - dueForRetry.length;
        if (toStart > 0) {
          for (let i = 0; i < toStart; i++) {
            await this.startOne(server.id);
          }
        }
        if (toStart > 0 || dueForRetry.length > 0) {
          reconciled++;
        }
      } catch (err) {
        errors.push(`${server.name}: ${err instanceof Error ? err.message : String(err)}`);
      }
    }

    return { reconciled, errors };
  }

  /**
   * Remove an instance (stop container + delete DB record).
   * Does NOT reconcile — caller should reconcile after if needed.
   */
  async remove(id: string): Promise<{ serverId: string }> {
    const instance = await this.getById(id);

    if (instance.containerId) {
      try {
        await this.orchestrator.stopContainer(instance.containerId);
      } catch {
        // Container may already be stopped
      }
      try {
        await this.orchestrator.removeContainer(instance.containerId, true);
      } catch {
        // Container may already be gone
      }
      this.invalidateStdio(instance.containerId);
    }

    await this.instanceRepo.delete(id);
    return { serverId: instance.serverId };
  }

  /**
   * Remove all instances for a server (used before server deletion).
   * Stops all containers so Prisma cascade only cleans up DB records.
   */
  async removeAllForServer(serverId: string): Promise<void> {
    const instances = await this.instanceRepo.findAll(serverId);
    for (const inst of instances) {
      if (inst.containerId) {
        try {
          await this.orchestrator.stopContainer(inst.containerId);
        } catch {
          // best-effort
        }
        try {
          await this.orchestrator.removeContainer(inst.containerId, true);
        } catch {
          // best-effort
        }
        this.invalidateStdio(inst.containerId);
      }
    }
  }

  async inspect(id: string): Promise<ContainerInfo> {
    const instance = await this.getById(id);
    if (!instance.containerId) {
      throw new InvalidStateError(`Instance '${id}' has no container`);
    }
    return this.orchestrator.inspectContainer(instance.containerId);
  }

  async getLogs(id: string, opts?: { tail?: number }): Promise<{ stdout: string; stderr: string }> {
    const instance = await this.getById(id);
    if (!instance.containerId) {
      return { stdout: '', stderr: '' };
    }

    return this.orchestrator.getContainerLogs(instance.containerId, opts);
  }

  /**
   * Start a single instance for a server. Creates a fresh `STARTING` row
   * and hands off to `attemptStart` for the env+container work. On
   * failure, `attemptStart` marks the row `ERROR` with a backoff-aware
   * `nextRetryAt`; the reconciler picks it up later via `retryInstance`.
   */
  private async startOne(serverId: string): Promise<McpInstance> {
    const server = await this.serverRepo.findById(serverId);
    if (!server) throw new NotFoundError(`McpServer '${serverId}' not found`);

    // External servers don't need container management
    if (server.externalUrl) {
      return this.instanceRepo.create({
        serverId,
        status: 'RUNNING',
        metadata: { external: true, url: server.externalUrl },
      });
    }

    const instance = await this.instanceRepo.create({
      serverId,
      status: 'STARTING',
    });
    return this.attemptStart(instance, server);
  }

  /**
   * Re-attempt a previously-errored instance in place, preserving its
   * `attemptCount` so the backoff schedule escalates correctly. Called
   * by `reconcileAll` for ERROR instances whose `nextRetryAt` has elapsed.
   */
  private async retryInstance(instance: McpInstance): Promise<McpInstance> {
    const server = await this.serverRepo.findById(instance.serverId);
    if (!server) {
      // Server was deleted underneath us — nothing to retry against.
      return this.markInstanceError(instance, 'Server no longer exists');
    }

    if (server.externalUrl) {
      // External servers don't need a container; the URL is the contract.
      return this.instanceRepo.updateStatus(instance.id, 'RUNNING', {
        metadata: { external: true, url: server.externalUrl },
      });
    }

    // Reset transient fields but keep retry counters via the metadata
    // passed through `attemptStart` → `markInstanceError`.
    await this.instanceRepo.updateStatus(instance.id, 'STARTING', {});
    const refreshed = (await this.instanceRepo.findById(instance.id)) ?? instance;
    return this.attemptStart(refreshed, server);
  }

  /**
   * Run the env-resolution + container-creation steps for a STARTING
   * instance. On any failure, mark the instance `ERROR` with structured
   * retry metadata. Used by both initial start (`startOne`) and retry
   * (`retryInstance`).
   */
  private async attemptStart(
    instance: McpInstance,
    server: McpServer,
  ): Promise<McpInstance> {
    // Determine image + command based on server config:
    // 1. Explicit dockerImage → use as-is
    // 2. packageName → use runtime-specific runner image (node/python/go/...)
    // 3. Fallback → server name (legacy)
    let image: string;
    let pkgCommand: string[] | undefined;

    if (server.dockerImage) {
      image = server.dockerImage;
    } else if (server.packageName) {
      const runtime = (server.runtime as string | null) ?? 'node';
      image = RUNNER_IMAGES[runtime] ?? RUNNER_IMAGES['node']!;
      // Runner entrypoint handles package execution (npx -y / uvx / go run)
      const serverCommand = server.command as string[] | null;
      pkgCommand = [server.packageName, ...(serverCommand ?? [])];
    } else {
      image = server.name;
    }

    try {
      const spec: ContainerSpec = {
        image,
        name: `mcpctl-${server.name}-${instance.id}`,
        hostPort: null,
        network: MCP_SERVERS_NETWORK,
        labels: {
          'mcpctl.server-id': server.id,
          'mcpctl.instance-id': instance.id,
        },
      };
      if (server.transport === 'SSE' || server.transport === 'STREAMABLE_HTTP') {
        spec.containerPort = server.containerPort ?? 3000;
      }

      // Volumes are keyed on the server, not this instance: `spec.name` carries
      // the instance id and a claim named after it would be torn down on the
      // next server edit — exactly when the data has to survive.
      const volumes = (server.volumes ?? []) as Array<{
        name: string;
        mountPath: string;
        sizeGb?: number;
        storageClass?: string;
        fsGroup?: number;
      }>;
      if (volumes.length > 0) {
        spec.volumes = volumes.map((v) => ({
          claimName: `mcpctl-${server.name}-${v.name}`,
          mountPath: v.mountPath,
          sizeGb: v.sizeGb ?? 10,
          ...(v.storageClass !== undefined && v.storageClass !== ''
            ? { storageClass: v.storageClass }
            : {}),
          ...(v.fsGroup !== undefined ? { fsGroup: v.fsGroup } : {}),
        }));
      }
      // Package-based servers: command = [packageName, ...args] (entrypoint handles execution)
      // Docker-image servers: use explicit command if provided
      if (pkgCommand) {
        spec.command = pkgCommand;
      } else {
        const command = server.command as string[] | null;
        if (command) {
          spec.command = command;
        }
      }

      // Resolve env vars from inline values and secret refs.
      //
      // Failure here is FATAL for the start attempt: a container that
      // boots without its declared secrets will silently mis-behave (we
      // saw this with gitea-mcp-server starting up with an empty
      // GITEA_ACCESS_TOKEN when OpenBao was unreachable, then reporting
      // "healthy" while every authed call failed). Marking the instance
      // ERROR with a backoff-aware nextRetryAt is honest; the reconciler
      // will retry it in-place on the next tick whose nextRetryAt has
      // elapsed. Optional/missing env vars should be modeled as `value: ""`
      // entries on the server, not as silent secret-resolution failures.
      // Injected delivery needs BOTH an orchestrator that can honour references
      // (Docker cannot) and a provisioner for the per-server identity. Falling
      // back to inline resolution when either is missing keeps local dev and
      // plaintext-backend setups working unchanged.
      const useInjector = server.secretDelivery === 'injector'
        && this.orchestrator.capabilities?.secretRefs === true
        && this.serverIdentity !== undefined;

      if (this.secretResolver) {
        try {
          if (useInjector) {
            // Resolve and DISCARD, purely so a missing secret still fails here
            // with a clear message instead of becoming a vault-agent-init
            // crashloop that mcpd reports as a generic pod failure. See
            // validateServerEnvRefs.
            await validateServerEnvRefs(server, this.secretResolver);

            const { inline, refs } = partitionServerEnv(server);
            if (Object.keys(inline).length > 0) spec.env = inline;
            spec.envFromSecret = refs;
          } else {
            const resolvedEnv = await resolveServerEnv(server, this.secretResolver);
            if (Object.keys(resolvedEnv).length > 0) {
              spec.env = resolvedEnv;
            }
          }
        } catch (envErr) {
          const msg = envErr instanceof Error ? envErr.message : String(envErr);
          return this.markInstanceError(instance, `secret resolution failed: ${msg}`);
        }
      }

      if (useInjector && spec.envFromSecret !== undefined && spec.envFromSecret.length > 0) {
        try {
          // Converge the identity BEFORE the pod exists: a pod whose role is
          // not yet written logs in, gets no capabilities, and crashloops.
          const identity = await this.serverIdentity!.ensureFor(server);
          spec.serviceAccountName = identity;
          spec.automountServiceAccountToken = true;
          spec.annotations = this.serverIdentity!.annotationsFor(identity, spec.envFromSecret);
          const wrapped = this.serverIdentity!.wrapCommand(server, spec.command);
          if (wrapped !== undefined) spec.command = wrapped;
        } catch (idErr) {
          const msg = idErr instanceof Error ? idErr.message : String(idErr);
          return this.markInstanceError(instance, `secret identity provisioning failed: ${msg}`);
        }
      }

      // Pull image if not available locally
      try {
        await this.orchestrator.pullImage(image);
      } catch {
        // Image may already be available locally
      }

      const containerInfo = await this.orchestrator.createContainer(spec);

      const updateFields: { containerId: string; port?: number } = {
        containerId: containerInfo.containerId,
      };
      if (containerInfo.port !== undefined) {
        updateFields.port = containerInfo.port;
      }

      // Set STARTING — syncStatus will promote to RUNNING once the container is actually ready
      return this.instanceRepo.updateStatus(instance.id, 'STARTING', updateFields);
    } catch (err) {
      return this.markInstanceError(
        instance,
        err instanceof Error ? err.message : String(err),
      );
    }
  }

  /**
   * Mark an instance ERROR with a backoff-aware retry schedule. The
   * `attemptCount` accumulates across retries (preserved by
   * `retryInstance` which reuses the same row), so the schedule
   * actually escalates: 30s × 5 → 5min thereafter.
   */
  private async markInstanceError(
    instance: McpInstance,
    error: string,
  ): Promise<McpInstance> {
    const meta = readRetryMeta(instance);
    const attemptCount = (typeof meta.attemptCount === 'number' ? meta.attemptCount : 0) + 1;
    const delayMs = nextDelayMs(attemptCount);
    const now = new Date();
    const nextRetryAt = new Date(now.getTime() + delayMs).toISOString();
    return this.instanceRepo.updateStatus(instance.id, 'ERROR', {
      metadata: {
        ...meta,
        error,
        attemptCount,
        lastAttemptAt: now.toISOString(),
        nextRetryAt,
      },
    });
  }

  /** Stop and remove a single instance. */
  private async removeOne(instance: McpInstance): Promise<void> {
    if (instance.containerId) {
      try {
        await this.orchestrator.stopContainer(instance.containerId);
      } catch { /* best-effort */ }
      try {
        await this.orchestrator.removeContainer(instance.containerId, true);
      } catch { /* best-effort */ }
      this.invalidateStdio(instance.containerId);
    }
    await this.instanceRepo.delete(instance.id);
  }
}