@pompeii-labs/cli/dist/container-server/server.js

Magma CLI

import express from "express";
import { createServer } from "http";
import { WebSocketServer } from "ws";
import { Handlers } from "./handlers.js";
import fs from "fs";
import path from "path";
import cors from "cors";
import bodyParser from "body-parser";
import cron from "node-cron";
import { loadHooks, loadJobs } from "@pompeii-labs/magma";
class ContainerServer {
  constructor(args) {
    this.jobs = [];
    this.app = express();
    this.port = args.port || (process.env.PORT ? parseInt(process.env.PORT) : 3e3);
    this.host = args.host || process.env.HOST || "0.0.0.0";
    this.active = args.agent;
    this.app.use(cors());
    this.app.use(express.json());
    this.app.use(bodyParser.urlencoded({ extended: false }));
    this.httpServer = createServer(this.app);
    this.wsServer = new WebSocketServer({ noServer: true });
    this.setupWebSocket();
    this.setupRoutes();
  }
  setupRoutes() {
    this.app.get("/health", (req, res) => {
      res.sendStatus(200);
    });
    this.app.post("/v1/agents/:agentId/chat", async (req, res) => {
      try {
        let agent;
        if (this.active) {
          agent = this.active;
        } else {
          const AgentClass = await this.loadAgent();
          agent = new AgentClass();
        }
        const { message } = req.body;
        if (!message) {
          res.status(400).json({ message: "Message is required" });
          return;
        }
        if (!message.content) {
          res.status(400).json({ message: "Message content is required" });
          return;
        }
        agent.addMessage(message);
        const reply = await agent.main();
        if (!reply) {
          res.status(500).json({ message: "No reply from agent" });
          return;
        }
        res.status(201).json(reply);
      } catch (error) {
        console.error(error);
        res.status(500).json({ message: error.message });
      }
    });
    this.app.get("/v1/agents/:agentId/hooks/:hook", async (req, res) => {
      try {
        const hook = req.params.hook;
        let agent;
        const agentClass = await this.loadAgent();
        if (this.active) {
          agent = this.active;
        } else {
          agent = new agentClass();
        }
        let hookWrapper;
        hookWrapper = loadHooks(agentClass).find((h) => h.name === hook);
        if (!hookWrapper) {
          hookWrapper = loadHooks(agent).find((h) => h.name === hook);
        }
        if (!hookWrapper) {
          res.status(404).json({
            error: `No hook handler found for '${hook}'`
          });
          return;
        }
        try {
          if (req.method === "POST") {
            if ("raw" in req.body) {
              req["raw"] = req.body.raw;
            }
            if ("body" in req.body) {
              req.body = req.body.body;
            }
          }
          await hookWrapper.handler(req, res);
        } catch (error) {
          console.error(error);
          res.status(400).json({ error: error.message });
          return;
        }
        if (!res.headersSent) {
          res.status(200).json({ success: true });
        }
      } catch (error) {
        console.error(error);
        res.status(500).json({ error: error.message });
      }
    });
  }
  /**
   * Loads the agent specified as the `main` entrypoint in the package.json file
   *
   * @returns The agent class type
   */
  async loadAgent() {
    if (this.agent) {
      return this.agent;
    }
    const pkg = JSON.parse(fs.readFileSync(path.join(process.cwd(), "package.json"), "utf-8"));
    const entrypoint = pkg.main ?? "index.js";
    const entrypointPath = path.join(process.cwd(), entrypoint);
    const fileUrl = `file://${entrypointPath}`;
    try {
      let AgentClass = await import(fileUrl.toString());
      while (AgentClass.default) {
        AgentClass = AgentClass.default;
      }
      this.agent = AgentClass;
      return AgentClass;
    } catch (error) {
      console.error(error);
      throw error;
    }
  }
  setupWebSocket() {
    this.httpServer.on("upgrade", async (req, socket, head) => {
      try {
        const AgentClass = await this.loadAgent();
        this.wsServer.handleUpgrade(req, socket, head, (ws) => {
          this.wsServer.emit("connection", ws, req, AgentClass);
        });
      } catch (err) {
        console.error(`Error upgrading WebSocket connection: ${err.message ?? "Unknown"}`);
        socket.write("HTTP/1.1 500 Internal Server Error\r\n\r\n");
        socket.destroy();
      }
    });
    this.wsServer.on("connection", this.handleConnection.bind(this));
  }
  async handleConnection(ws, req, agentWs) {
    try {
      const agent = new agentWs();
      const prototype = Object.getPrototypeOf(agent);
      const propertyNames = Object.getOwnPropertyNames(prototype);
      if (!propertyNames.includes("receive")) {
        agent.receive = Handlers.receive.bind(agent);
      }
      if (!propertyNames.includes("send")) {
        agent.send = Handlers.send.bind(agent);
      }
      agent.ws = ws;
      ws.on("message", agent.receive.bind(agent));
      ws.on("close", (code) => this.handleClose(code, agent));
      ws.on("error", (error) => this.handleError(error, agent));
      ws.on("unexpected-response", (event) => console.log(event));
      await agent.setup();
    } catch (error) {
      console.error(error);
    }
  }
  async handleClose(code, agent) {
    if (code === 1e3) {
      console.info("WebSocket connection closed gracefully");
    }
    try {
      await agent.cleanup();
    } catch (error) {
      console.error(`Error closing WebSocket connection: ${error.message ?? "Unknown"}`);
    }
  }
  async handleError(err, agent) {
    try {
      await agent.onError(err);
      await agent.cleanup();
    } catch (error) {
      console.error(error);
    }
  }
  async scheduleAgentCronJobs() {
    try {
      const AgentClass = await this.loadAgent();
      const staticJobs = loadJobs(AgentClass);
      staticJobs.forEach((job) => {
        const wrappedHandler = async () => {
          try {
            await job.handler();
          } catch (error) {
            console.error(`Error in job '${job.schedule}': ${error.message}`);
          }
        };
        const scheduledJob = cron.schedule(job.schedule, wrappedHandler, job.options);
        this.jobs.push({ name: job.schedule, job: scheduledJob });
      });
      const jobs = loadJobs(this.active);
      jobs.forEach((job) => {
        const scheduledJob = cron.schedule(job.schedule, job.handler, job.options);
        this.jobs.push({ name: job.schedule, job: scheduledJob });
      });
    } catch (error) {
      console.error(`Error scheduling agent cron jobs: ${error.message ?? "Unknown"}`);
    }
  }
  async start() {
    this.httpServer.listen(this.port, this.host);
    await this.scheduleAgentCronJobs();
  }
  stop() {
    this.httpServer.close();
    this.wsServer.close();
  }
}
export {
  ContainerServer
};