paif/server.js

'use strict';

// Load persisted config from ~/.memory-bridge/.env (populated by auth.init on first run)
const path = require('path');
const os = require('os');
require('dotenv').config({ path: path.join(os.homedir(), '.memory-bridge', '.env') });

const express = require('express');
const auth = require('./auth');
const embedder = require('./embed');
const store = require('./store');
const identity = require('./identity');
const checkpoint = require('./checkpoint');
const gravity = require('./gravity');

const PORT = Number(process.env.MEMORY_BRIDGE_PORT) || 3722;
const VERSION = '2.0.0'; // Major version bump for PAIF multi-tenant support

// Validate that agent request matches authenticated context
function validateAgentAccess(req, res, requestedAgentId) {
  const { agentId, isMaster } = req.agentContext;

  // Master token can access any agent
  if (isMaster) {
    return true;
  }

  // Agent token can only access their own namespace
  if (agentId !== requestedAgentId) {
    res.status(403).json({
      error: 'Access denied',
      message: `Token for agent '${agentId}' cannot access namespace '${requestedAgentId}'`
    });
    return false;
  }

  return true;
}

async function main() {
  auth.init();
  await store.init();
  const embeddingMode = await embedder.init();

  const app = express();
  app.use(express.json({ limit: '10mb' }));

  // ── GET /health (public) ──────────────────────────────────────────────────
  app.get('/health', async (_req, res) => {
    try {
      const memoriesCount = await store.count();
      const stats = await store.stats();
      res.json({
        status: 'ok',
        version: VERSION,
        embeddingMode,
        memoriesCount,
        registeredAgents: stats.totalAgents,
        paifMode: true
      });
    } catch (err) {
      res.status(500).json({ status: 'error', error: err.message });
    }
  });

  // ── POST /register-agent (master only) ──────────────────────────────────────
  app.post('/register-agent', auth.middleware, (req, res) => {
    if (!req.agentContext.isMaster) {
      return res.status(403).json({ error: 'Only master token can register agents' });
    }

    const { agent_id } = req.body || {};
    if (!agent_id || typeof agent_id !== 'string') {
      return res.status(400).json({ error: 'agent_id is required' });
    }

    try {
      const token = auth.registerAgent(agent_id);
      res.status(201).json({
        agent_id,
        token,
        message: `Agent '${agent_id}' registered successfully`
      });
    } catch (err) {
      res.status(400).json({ error: err.message });
    }
  });

  // ── GET /agents (master only) ───────────────────────────────────────────────
  app.get('/agents', auth.middleware, (req, res) => {
    if (!req.agentContext.isMaster) {
      return res.status(403).json({ error: 'Only master token can list agents' });
    }
    res.json({ agents: auth.listAgents() });
  });

  // ── POST /identity (master only - create identity) ──────────────────────────
  app.post('/identity', auth.middleware, (req, res) => {
    if (!req.agentContext.isMaster) {
      return res.status(403).json({ error: 'Only master token can create identities' });
    }

    const { agent_id, origin, purpose, values, voice, lineage } = req.body || {};

    if (!agent_id) {
      return res.status(400).json({ error: 'agent_id is required' });
    }

    try {
      const newIdentity = identity.createIdentity(agent_id, {
        createdBy: origin?.created_by || 'unknown',
        purpose,
        values,
        voice,
        lineage
      });
      res.status(201).json({
        agent_id,
        identity: newIdentity,
        message: `Identity created for agent '${agent_id}'`
      });
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── GET /identity/:agent_id ─────────────────────────────────────────────────
  app.get('/identity/:agent_id', auth.middleware, (req, res) => {
    const { agent_id } = req.params;

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    try {
      const id = identity.loadIdentity(agent_id);
      if (!id) {
        return res.status(404).json({ error: `Identity not found for agent '${agent_id}'` });
      }
      res.json({ agent_id, identity: id });
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── POST /identity/:agent_id/validate ─────────────────────────────────────────
  app.post('/identity/:agent_id/validate', auth.middleware, (req, res) => {
    const { agent_id } = req.params;
    const { text } = req.body || {};

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    if (!text) {
      return res.status(400).json({ error: 'text is required for validation' });
    }

    try {
      const result = identity.validateAlignment(agent_id, text);
      res.json({ agent_id, validation: result });
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── POST /identity/:agent_id/update ─────────────────────────────────────────
  // Update mutable state (master only or own agent)
  app.post('/identity/:agent_id/update', auth.middleware, (req, res) => {
    const { agent_id } = req.params;
    const { updates } = req.body || {};

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    if (!updates) {
      return res.status(400).json({ error: 'updates object is required' });
    }

    try {
      const updated = identity.updateMutableState(agent_id, updates);
      res.json({ agent_id, identity: updated, message: 'Mutable state updated' });
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── POST /checkpoint/:agent_id ───────────────────────────────────────────────
  app.post('/checkpoint/:agent_id', auth.middleware, (req, res) => {
    const { agent_id } = req.params;
    const sessionData = req.body || {};

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    try {
      const cp = checkpoint.createCheckpoint(agent_id, sessionData);
      res.status(201).json({
        agent_id,
        session_id: cp.metadata.session_id,
        checkpointed_at: cp.metadata.ended_at,
        message: 'Session checkpointed successfully'
      });
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── GET /checkpoint/:agent_id/latest ──────────────────────────────────────────
  app.get('/checkpoint/:agent_id/latest', auth.middleware, (req, res) => {
    const { agent_id } = req.params;

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    try {
      const cp = checkpoint.getLatestCheckpoint(agent_id);
      if (!cp) {
        return res.status(404).json({ error: `No checkpoint found for agent '${agent_id}'` });
      }
      res.json({ agent_id, checkpoint: cp });
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── POST /restore/:agent_id ──────────────────────────────────────────────────
  app.post('/restore/:agent_id', auth.middleware, (req, res) => {
    const { agent_id } = req.params;

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    try {
      const result = checkpoint.restoreSession(agent_id);
      res.json({
        agent_id,
        restored: result.restored,
        checkpoint_found: !!result.checkpoint,
        identity_updated: result.identity_updated,
        message: result.restored ? 'Session restored from checkpoint' : 'No checkpoint found - fresh session started'
      });
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── GET /sessions/:agent_id ───────────────────────────────────────────────────
  app.get('/sessions/:agent_id', auth.middleware, (req, res) => {
    const { agent_id } = req.params;

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    try {
      const checkpoints = checkpoint.listCheckpoints(agent_id);
      const stats = checkpoint.getSessionStats(agent_id);
      res.json({ agent_id, checkpoints, stats });
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── POST /gravity-check/:agent_id ───────────────────────────────────────────
  app.post('/gravity-check/:agent_id', auth.middleware, (req, res) => {
    const { agent_id } = req.params;
    const { text, auto_realign = false } = req.body || {};

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    if (!text) {
      return res.status(400).json({ error: 'text is required for gravity check' });
    }

    try {
      const result = gravity.gravityCheck(agent_id, text, {
        autoRealign: auto_realign,
        updateStats: true
      });
      res.json(result);
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── GET /gravity-check/:agent_id/quick ────────────────────────────────────────
  app.get('/gravity-check/:agent_id/quick', auth.middleware, (req, res) => {
    const { agent_id } = req.params;
    const { text } = req.query;

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    if (!text) {
      return res.status(400).json({ error: 'text query param is required' });
    }

    try {
      const result = gravity.quickCheck(text, agent_id);
      res.json({ agent_id, ...result });
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  // ── GET /gravity-patterns (public) ─────────────────────────────────────────────
  app.get('/gravity-patterns', (_req, res) => {
    res.json({
      patterns: Object.keys(gravity.DRIFT_PATTERNS),
      description: 'Categories of drift patterns detected by persona gravity system'
    });
  });

  // ── Auth guard for all remaining endpoints ─────────────────────────────────
  app.use(auth.middleware);

  // ── POST /store ────────────────────────────────────────────────────────────
  app.post('/store', async (req, res) => {
    const { text, tags = [], source = '', agent_id } = req.body || {};

    if (!text || typeof text !== 'string' || !text.trim()) {
      return res.status(400).json({ error: '`text` is required' });
    }

    if (!agent_id || typeof agent_id !== 'string') {
      return res.status(400).json({ error: '`agent_id` is required' });
    }

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    try {
      const vector = await embedder.embed(text.trim());
      const result = await store.add(agent_id, { text: text.trim(), tags, source, vector });
      res.status(201).json(result);
    } catch (err) {
      console.error('POST /store:', err.message);
      res.status(500).json({ error: err.message });
    }
  });

  // ── POST /recall ────────────────────────────────────────────────────────────
  app.post('/recall', async (req, res) => {
    const { query, limit = 10, tags = [], agent_id } = req.body || {};

    if (!query || typeof query !== 'string' || !query.trim()) {
      return res.status(400).json({ error: '`query` is required' });
    }

    if (!agent_id || typeof agent_id !== 'string') {
      return res.status(400).json({ error: '`agent_id` is required' });
    }

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    try {
      const vector = await embedder.embed(query.trim());
      const memories = await store.query(agent_id, {
        vector,
        limit: Math.min(Number(limit) || 10, 100),
        tags,
      });
      res.json({ memories });
    } catch (err) {
      console.error('POST /recall:', err.message);
      res.status(500).json({ error: err.message });
    }
  });

  // ── POST /forget ──────────────────────────────────────────────────────────
  app.post('/forget', async (req, res) => {
    const { id, agent_id } = req.body || {};

    if (!id) return res.status(400).json({ error: '`id` is required' });
    if (!agent_id) return res.status(400).json({ error: '`agent_id` is required' });

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    try {
      await store.remove(agent_id, id);
      res.json({ message: 'forgotten', id });
    } catch (err) {
      console.error('POST /forget:', err.message);
      res.status(500).json({ error: err.message });
    }
  });

  // ── GET /list ──────────────────────────────────────────────────────────────
  app.get('/list', async (req, res) => {
    const { agent_id, page = 1, pageSize = 50 } = req.query;

    if (!agent_id) return res.status(400).json({ error: '`agent_id` query param is required' });

    if (!validateAgentAccess(req, res, agent_id)) {
      return;
    }

    try {
      const result = await store.list(agent_id, {
        page: Math.max(1, parseInt(page) || 1),
        pageSize: Math.min(100, Math.max(1, parseInt(pageSize) || 50))
      });
      res.json(result);
    } catch (err) {
      console.error('GET /list:', err.message);
      res.status(500).json({ error: err.message });
    }
  });

  // ── GET /stats (master only) ────────────────────────────────────────────────
  app.get('/stats', auth.middleware, async (req, res) => {
    if (!req.agentContext.isMaster) {
      return res.status(403).json({ error: 'Only master token can view stats' });
    }
    try {
      const stats = await store.stats();
      res.json(stats);
    } catch (err) {
      res.status(500).json({ error: err.message });
    }
  });

  app.listen(PORT, () => {
    console.log(`\nmemory-bridge v${VERSION} listening on http://localhost:${PORT}`);
    console.log(`Embedding       : ${embeddingMode}`);
    console.log(`Storage         : ${path.join(os.homedir(), '.memory-bridge')}`);
    console.log(`PAIF Mode       : Multi-tenant agent isolation enabled\n`);
  });
}

main().catch((err) => {
  console.error('Failed to start memory-bridge:', err);
  process.exit(1);
});