src/mcp/createServer.ts

import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { ReadResourceRequestSchema } from "@modelcontextprotocol/sdk/types.js";
import { z } from "zod";
import type { SessionManager } from "../runtime/SessionManager.js";
import { ToolError } from "../utils/errors.js";
import { SERVER_INSTRUCTIONS } from "./instructions.js";
import { newId } from "../utils/ids.js";
import type { ExecResult } from "../worker/protocol.js";
import { encodePathForUri, decodePathFromUri } from "../utils/paths.js";

export const SERVER_NAME = "just-bash-mcp";
export const SERVER_VERSION = "0.1.0";

const PathSchema = z.string().min(1);

function toolErrorResult(err: unknown) {
  const message =
    err instanceof ToolError
      ? `[${err.code}] ${err.message}`
      : err instanceof Error
        ? err.message
        : String(err);
  return {
    isError: true,
    content: [{ type: "text" as const, text: message }],
    structuredContent: {
      error: {
        code: err instanceof ToolError ? err.code : "INTERNAL",
        message,
      },
    },
  };
}

function jsonOk(structured: unknown) {
  return {
    content: [{ type: "text" as const, text: JSON.stringify(structured) }],
    structuredContent: structured as Record<string, unknown>,
  };
}

export function createServer(sm: SessionManager) {
  const server = new McpServer(
    { name: SERVER_NAME, version: SERVER_VERSION },
    {
      instructions: SERVER_INSTRUCTIONS,
      capabilities: {
        tools: { listChanged: false },
        resources: { listChanged: false },
      },
    },
  );

  const cfg = sm.config;

  server.registerTool(
    "session_create",
    {
      title: "Create a just-bash workspace session",
      description:
        "Creates an isolated workspace with an in-memory filesystem. Returns a session_id that must be passed to other tools.",
      inputSchema: {
        initial_files: z
          .record(z.string())
          .optional()
          .describe("Map of absolute path -> utf8 content to seed the workspace"),
        cwd: z.string().optional(),
        ttl_seconds: z.number().int().positive().optional(),
        capabilities: z
          .object({
            network: z.boolean().optional(),
            python: z.boolean().optional(),
            javascript: z.boolean().optional(),
          })
          .optional(),
      },
    },
    async (args, extra) => {
      try {
        const tokenHash = (extra as any)?.authInfo?.tokenHash as string | undefined;
        const rec = await sm.createSession({
          initialFiles: args.initial_files,
          cwd: args.cwd,
          ttlSeconds: args.ttl_seconds,
          capabilities: args.capabilities,
          ownerTokenHash: tokenHash,
        });
        return jsonOk({
          session_id: rec.id,
          cwd: rec.cwd,
          expires_at: new Date(rec.expiresAt).toISOString(),
          capabilities: rec.capabilities,
          limits: {
            max_session_fs_bytes: cfg.maxSessionFsBytes,
            max_single_file_bytes: cfg.maxSingleFileBytes,
            max_write_batch_bytes: cfg.maxWriteBatchBytes,
            max_read_bytes_per_call: cfg.maxReadBytesPerCall,
            default_timeout_ms: cfg.defaultTimeoutMs,
            max_timeout_ms: cfg.maxTimeoutMs,
          },
        });
      } catch (e) {
        return toolErrorResult(e);
      }
    },
  );

  server.registerTool(
    "session_reset",
    {
      title: "Reset a session",
      description:
        "Terminates the existing worker, starts a fresh one with the same session_id. mode=empty starts empty; mode=initial_files restores the files passed to session_create.",
      inputSchema: {
        session_id: z.string(),
        mode: z.enum(["empty", "initial_files"]).default("empty"),
      },
    },
    async (args) => {
      try {
        const rec = await sm.reset(args.session_id, args.mode);
        return jsonOk({
          session_id: rec.id,
          reset: true,
          cwd: rec.cwd,
          expires_at: new Date(rec.expiresAt).toISOString(),
        });
      } catch (e) {
        return toolErrorResult(e);
      }
    },
  );

  server.registerTool(
    "session_delete",
    {
      title: "Delete a session",
      description: "Terminates the worker and removes session metadata. Idempotent.",
      inputSchema: {
        session_id: z.string(),
      },
    },
    async (args) => {
      try {
        const ok = await sm.delete(args.session_id, "user_delete");
        return jsonOk({ deleted: ok });
      } catch (e) {
        return toolErrorResult(e);
      }
    },
  );

  server.registerTool(
    "bash_exec",
    {
      title: "Run a command in a session",
      description:
        "Runs a script inside an isolated worker. NOTE: files persist across calls, but cwd/env/aliases/functions do NOT persist between bash_exec calls unless passed in again.",
      inputSchema: {
        session_id: z.string(),
        script: z.string(),
        stdin: z.string().optional(),
        cwd: z.string().optional(),
        env: z.record(z.string()).optional(),
        replace_env: z.boolean().optional(),
        timeout_ms: z.number().int().positive().optional(),
        max_stdout_bytes: z.number().int().nonnegative().optional(),
        max_stderr_bytes: z.number().int().nonnegative().optional(),
      },
    },
    async (args) => {
      try {
        const rec = sm.require(args.session_id);
        const timeoutMs = Math.min(
          Math.max(args.timeout_ms ?? cfg.defaultTimeoutMs, 1),
          cfg.maxTimeoutMs,
        );
        const hardTimeoutMs = timeoutMs + cfg.watchdogGraceMs;
        const maxStdoutBytes = Math.min(
          args.max_stdout_bytes ?? cfg.maxEmbeddedStdoutBytes,
          cfg.maxEmbeddedStdoutBytes,
        );
        const maxStderrBytes = Math.min(
          args.max_stderr_bytes ?? cfg.maxEmbeddedStderrBytes,
          cfg.maxEmbeddedStderrBytes,
        );
        const execId = newId("ex");
        const outcome = await sm.runOnSessionQueue(args.session_id, () =>
          rec.worker.exec(
            {
              execId,
              script: args.script,
              stdin: args.stdin,
              cwd: args.cwd,
              env: args.env,
              replaceEnv: args.replace_env,
              timeoutMs,
              maxStdoutBytes,
              maxStderrBytes,
            },
            hardTimeoutMs,
          ),
        );
        sm.touch(args.session_id);
        if (outcome.hardKilled) {
          sm.invalidate(args.session_id, "hard_timeout");
          const body = {
            exec_id: execId,
            exit_code: 124,
            stdout: "",
            stderr: "execution exceeded hard timeout; session worker was terminated",
            duration_ms: hardTimeoutMs,
            timed_out: true,
            hard_killed: true,
            session_valid: false,
            truncated: { stdout: false, stderr: false },
            resource_links: [],
            metadata: { commands: [] },
          };
          return jsonOk(body);
        }
        const res = outcome.result as ExecResult;

        const resourceLinks = collectResourceLinks(sm, rec.id, execId, res, cfg.maxResourceBytesSession, rec.idleTtlMs);

        const body = {
          exec_id: execId,
          exit_code: res.exitCode,
          stdout: res.stdout,
          stderr: res.stderr,
          duration_ms: res.durationMs,
          timed_out: res.timedOut,
          hard_killed: false,
          session_valid: true,
          truncated: { stdout: res.truncatedStdout, stderr: res.truncatedStderr },
          resource_links: resourceLinks,
          metadata: { commands: res.commands },
          session_fs_bytes: res.sessionFsBytes,
        };
        return {
          content: [
            { type: "text" as const, text: JSON.stringify(body) },
            ...resourceLinks.map((rl) => ({
              type: "resource_link" as const,
              uri: rl.uri,
              name: rl.name,
              mimeType: rl.mime_type,
            })),
          ],
          structuredContent: body,
        };
      } catch (e) {
        return toolErrorResult(e);
      }
    },
  );

  server.registerTool(
    "files_write",
    {
      title: "Write files into the session",
      description: "Batch write files into the session virtual filesystem.",
      inputSchema: {
        session_id: z.string(),
        files: z
          .array(
            z.object({
              path: PathSchema,
              content: z.string(),
              encoding: z.enum(["utf8", "base64"]).default("utf8"),
            }),
          )
          .min(1),
        overwrite: z.boolean().default(true),
        create_dirs: z.boolean().default(true),
      },
    },
    async (args) => {
      try {
        const rec = sm.require(args.session_id);
        const result = await rec.worker.filesWrite({
          files: args.files.map((f) => ({
            path: f.path,
            content: f.content,
            encoding: f.encoding,
          })),
          overwrite: args.overwrite,
          createDirs: args.create_dirs,
        });
        sm.touch(args.session_id);
        return jsonOk({
          written: result.written,
          session_fs_bytes: result.sessionFsBytes,
        });
      } catch (e) {
        if (e instanceof Error && (e as any).code) {
          return toolErrorResult(new ToolError((e as any).code as any, e.message));
        }
        return toolErrorResult(e);
      }
    },
  );

  server.registerTool(
    "files_read",
    {
      title: "Read a file from the session",
      description: "Reads file content with offset-based pagination. Returns base64 for binary.",
      inputSchema: {
        session_id: z.string(),
        path: PathSchema,
        offset: z.number().int().nonnegative().default(0),
        max_bytes: z.number().int().positive().optional(),
        encoding: z.enum(["utf8", "base64", "auto"]).default("auto"),
      },
    },
    async (args) => {
      try {
        const rec = sm.require(args.session_id);
        const maxBytes = Math.min(
          args.max_bytes ?? cfg.maxReadBytesPerCall,
          cfg.maxReadBytesPerCall,
        );
        const result = await rec.worker.filesRead({
          path: args.path,
          offset: args.offset,
          maxBytes,
          encoding: args.encoding,
        });
        sm.touch(args.session_id);
        return jsonOk({
          path: result.path,
          content: result.content,
          encoding: result.encoding,
          bytes_read: result.bytesRead,
          next_offset: result.nextOffset,
        });
      } catch (e) {
        if (e instanceof Error && (e as any).code) {
          return toolErrorResult(new ToolError((e as any).code as any, e.message));
        }
        return toolErrorResult(e);
      }
    },
  );

  server.registerTool(
    "files_list",
    {
      title: "List files/directories in the session",
      description: "Lists files and directories with optional glob and recursive traversal.",
      inputSchema: {
        session_id: z.string(),
        path: PathSchema,
        recursive: z.boolean().default(false),
        glob: z.string().optional(),
        max_entries: z.number().int().positive().default(200),
        cursor: z.string().nullable().optional(),
      },
    },
    async (args) => {
      try {
        const rec = sm.require(args.session_id);
        const result = await rec.worker.filesList({
          path: args.path,
          recursive: args.recursive,
          glob: args.glob,
          maxEntries: args.max_entries,
          cursor: args.cursor ?? undefined,
        });
        sm.touch(args.session_id);
        return jsonOk({
          entries: result.entries,
          next_cursor: result.nextCursor,
        });
      } catch (e) {
        if (e instanceof Error && (e as any).code) {
          return toolErrorResult(new ToolError((e as any).code as any, e.message));
        }
        return toolErrorResult(e);
      }
    },
  );

  // Dynamic resource reads for justbash:// URIs.
  server.server.setRequestHandler(
    ReadResourceRequestSchema,
    async (req) => {
      const uri: string = req.params.uri;
      const parsed = parseResourceUri(uri);
      if (!parsed) {
        throw new Error(`Unknown resource URI: ${uri}`);
      }
      if (parsed.kind === "exec_stream") {
        const buf = sm.resourceStore.get(parsed.sessionId, parsed.execId, parsed.stream);
        if (!buf) {
          throw new Error(`Resource expired or evicted: ${uri}`);
        }
        return {
          contents: [
            {
              uri,
              mimeType: "text/plain",
              text: buf.bytes.toString("utf8"),
            },
          ],
        };
      }
      if (parsed.kind === "file") {
        const rec = sm.require(parsed.sessionId);
        const result = await rec.worker.filesRead({
          path: parsed.path,
          offset: 0,
          maxBytes: cfg.maxReadBytesPerCall,
          encoding: "auto",
        });
        return {
          contents: [
            {
              uri,
              mimeType: result.encoding === "base64" ? "application/octet-stream" : "text/plain",
              text: result.content,
            },
          ],
        };
      }
      throw new Error(`Unknown resource URI: ${uri}`);
    },
  );

  return server;
}

interface ResourceLink {
  uri: string;
  name: string;
  mime_type: string;
  bytes_available: number;
  complete: boolean;
  expires_at: string;
}

function collectResourceLinks(
  sm: SessionManager,
  sessionId: string,
  execId: string,
  res: ExecResult,
  ttlMs: number,
  sessionTtlMs: number,
): ResourceLink[] {
  void ttlMs;
  const links: ResourceLink[] = [];
  if (res.retainedStdout) {
    const ok = sm.resourceStore.store(
      sessionId,
      execId,
      "stdout",
      res.retainedStdout,
      res.capturedStdoutBytes <= Buffer.byteLength(res.retainedStdout, "utf8"),
      sessionTtlMs,
    );
    if (ok) {
      const buf = sm.resourceStore.get(sessionId, execId, "stdout");
      if (buf) {
        links.push({
          uri: `justbash://session/${sessionId}/exec/${execId}/stdout`,
          name: "stdout",
          mime_type: "text/plain",
          bytes_available: buf.bytes.length,
          complete: buf.complete,
          expires_at: new Date(buf.expiresAt).toISOString(),
        });
      }
    }
  }
  if (res.retainedStderr) {
    const ok = sm.resourceStore.store(
      sessionId,
      execId,
      "stderr",
      res.retainedStderr,
      res.capturedStderrBytes <= Buffer.byteLength(res.retainedStderr, "utf8"),
      sessionTtlMs,
    );
    if (ok) {
      const buf = sm.resourceStore.get(sessionId, execId, "stderr");
      if (buf) {
        links.push({
          uri: `justbash://session/${sessionId}/exec/${execId}/stderr`,
          name: "stderr",
          mime_type: "text/plain",
          bytes_available: buf.bytes.length,
          complete: buf.complete,
          expires_at: new Date(buf.expiresAt).toISOString(),
        });
      }
    }
  }
  return links;
}

type ParsedResourceUri =
  | { kind: "exec_stream"; sessionId: string; execId: string; stream: "stdout" | "stderr" }
  | { kind: "file"; sessionId: string; path: string };

function parseResourceUri(uri: string): ParsedResourceUri | null {
  const m1 = /^justbash:\/\/session\/([^/]+)\/exec\/([^/]+)\/(stdout|stderr)$/.exec(uri);
  if (m1) {
    return {
      kind: "exec_stream",
      sessionId: m1[1],
      execId: m1[2],
      stream: m1[3] as "stdout" | "stderr",
    };
  }
  const m2 = /^justbash:\/\/session\/([^/]+)\/file\/(.+)$/.exec(uri);
  if (m2) {
    return { kind: "file", sessionId: m2[1], path: decodePathFromUri(m2[2]) };
  }
  return null;
}

void encodePathForUri;