cc-libs/apis/libmcpcomputer.lua

local function defaultOs()
  return os;
end

local function formatLabel(label)
  if type(label) ~= 'string' or label == '' then
    return 'null';
  end
  return label;
end

local function appendOutput(output, value)
  output[#output + 1] = tostring(value);
end

local function serializeReturn(value)
  local valueType = type(value);
  if valueType == 'nil' then
    return { type = 'nil' };
  end

  if valueType == 'string' or valueType == 'number' or valueType == 'boolean' then
    return { type = valueType, value = value };
  end

  local ok, serialized = pcall(textutils.serialize, value);
  if ok then
    return { type = valueType, repr = serialized };
  end
  return { type = valueType, repr = tostring(value) };
end

local function createExecEnv(output)
  local env = setmetatable({}, { __index = _G });

  env.write = function(value)
    appendOutput(output, value);
  end;

  env.print = function(...)
    local values = table.pack(...);
    for i = 1, values.n do
      if i > 1 then
        appendOutput(output, '\t');
      end
      appendOutput(output, values[i]);
    end
    appendOutput(output, '\n');
  end;

  return env;
end

local function createMcpComputer()
  local api = {};

  api.formatLabel = formatLabel;

  function api.formatPong(osLike)
    osLike = osLike or defaultOs();
    return 'pong from ' .. tostring(osLike.getComputerID())
      .. ' (Label: ' .. formatLabel(osLike.getComputerLabel()) .. ')';
  end

  function api.parseArgs(args)
    args = args or {};
    local count = args.n or #args;
    local url = nil;

    if count == 0 then
      return nil, 'missing websocket URL';
    end

    local i = 1;
    while i <= count do
      local arg = args[i];
      if arg == '-url' then
        if not args[i + 1] or args[i + 1] == '' then
          return nil, 'missing value for -url';
        end
        url = args[i + 1];
        i = i + 1;
      elseif string.sub(tostring(arg), 1, 1) == '-' then
        return nil, 'unknown option: ' .. tostring(arg);
      elseif not url then
        url = arg;
      else
        return nil, 'unexpected argument: ' .. tostring(arg);
      end
      i = i + 1;
    end

    if not url or url == '' then
      return nil, 'missing websocket URL';
    end
    return { url = url };
  end

  function api.hello(osLike)
    osLike = osLike or defaultOs();
    return {
      type = 'hello',
      computerId = osLike.getComputerID(),
      computerLabel = osLike.getComputerLabel(),
    };
  end

  function api.executeLua(code)
    local output = {};
    if type(code) ~= 'string' or code == '' then
      return { ok = false, error = 'code must be a non-empty string', output = '' };
    end

    local fn, loadErr = load(code, 'mcp-exec', 't', createExecEnv(output));
    if not fn then
      return { ok = false, error = tostring(loadErr), output = table.concat(output) };
    end

    local values = table.pack(pcall(fn));
    if not values[1] then
      return { ok = false, error = tostring(values[2]), output = table.concat(output) };
    end

    local returns = {};
    for i = 2, values.n do
      returns[#returns + 1] = serializeReturn(values[i]);
    end
    return { ok = true, returns = returns, output = table.concat(output) };
  end

  function api.handleRequest(request, osLike)
    if type(request) ~= 'table' or request.type ~= 'request' or type(request.id) ~= 'string' then
      return nil;
    end

    if request.method == 'ping' then
      return {
        type = 'response',
        id = request.id,
        ok = true,
        result = api.formatPong(osLike),
      };
    end

    if request.method == 'exec-lua' then
      local params = request.params;
      local result = api.executeLua(type(params) == 'table' and params.code or nil);
      return {
        type = 'response',
        id = request.id,
        ok = result.ok,
        result = {
          returns = result.returns or {},
          output = result.output or '',
        },
        error = result.error,
      };
    end

    return {
      type = 'response',
      id = request.id,
      ok = false,
      error = 'unknown method',
    };
  end

  -- Resolve the websocket URL from CLI args first, then fall back to the
  -- 'mcp-computer.ws-url' setting when no argument is provided.
  function api.resolveUrl(args, settingsLib)
    args = args or {};
    local count = args.n or #args;

    if count > 0 then
      return api.parseArgs(args);
    end

    local url = settingsLib and settingsLib.get('mcp-computer.ws-url');
    if type(url) ~= 'string' or url == '' then
      return nil, 'missing websocket URL (pass a URL or set mcp-computer.ws-url)';
    end
    return { url = url };
  end

  -- Decode a raw websocket frame and dispatch it. Non-request frames (e.g.
  -- 'hello-ok') and malformed payloads produce no response.
  function api.onMessage(content, osLike, sendFn, decode)
    decode = decode or textutils.unserializeJSON;

    local ok, frame = pcall(decode, content);
    if not ok then
      return;
    end

    local response = api.handleRequest(frame, osLike);
    if response and sendFn then
      sendFn(response);
    end
  end

  -- Wire an event-driven MCP session onto an eventloop and connect. Registers
  -- websocket handlers and returns immediately, so it never blocks the boot
  -- sequence. Auto-reconnects on failure or close.
  function api.startSession(opts)
    opts = opts or {};
    local el = assert(opts.eventloop, 'startSession requires an eventloop');
    local url = assert(opts.url, 'startSession requires a url');
    local osLike = opts.os or defaultOs();
    local httpLike = opts.http or http;
    local reconnectDelay = opts.reconnectDelay or 5;
    local encode = opts.encode or textutils.serializeJSON;
    local decode = opts.decode or textutils.unserializeJSON;

    local activeWs = nil;

    local function connect()
      httpLike.websocketAsync(url);
    end

    local function scheduleReconnect()
      activeWs = nil;
      el.setTimeout(connect, reconnectDelay);
    end

    el.register('websocket_success', function(eventUrl, ws)
      if eventUrl ~= url then return; end
      activeWs = ws;
      ws.send(encode(api.hello(osLike)));
    end);

    el.register('websocket_message', function(eventUrl, content)
      if eventUrl ~= url then return; end
      api.onMessage(content, osLike, function(response)
        if activeWs then
          activeWs.send(encode(response));
        end
      end, decode);
    end);

    el.register('websocket_failure', function(eventUrl)
      if eventUrl ~= url then return; end
      scheduleReconnect();
    end);

    el.register('websocket_closed', function(eventUrl)
      if eventUrl ~= url then return; end
      scheduleReconnect();
    end);

    connect();

    return {
      isConnected = function() return activeWs ~= nil; end,
    };
  end

  return api;
end

return createMcpComputer;