chat/api/server/routes/agents/index.js

const express = require('express');
const { isEnabled, GenerationJobManager } = require('@hanzochat/api');
const { logger } = require('@hanzochat/data-schemas');
const {
  uaParser,
  checkBan,
  requireJwtAuth,
  messageIpLimiter,
  configMiddleware,
  messageUserLimiter,
  enforceGuestScope,
  guestMessageLimiter,
  requireGuestOrJwtAuth,
} = require('~/server/middleware');
const { saveMessage } = require('~/models');
const openai = require('./openai');
const responses = require('./responses');
const { v1 } = require('./v1');
const cloud = require('./cloud');
const chat = require('./chat');

const { LIMIT_MESSAGE_IP, LIMIT_MESSAGE_USER } = process.env ?? {};

const router = express.Router();

/**
 * Open Responses API routes (API key authentication handled in route file)
 * Mounted at /agents/v1/responses (full path: /v1/chat/agents/v1/responses)
 * NOTE: Must be mounted BEFORE /v1 to avoid being caught by the less specific route
 * @see https://openresponses.org/specification
 */
router.use('/v1/responses', responses);

/**
 * OpenAI-compatible API routes (API key authentication handled in route file)
 * Mounted at /agents/v1 (full path: /v1/chat/agents/v1/chat/completions)
 */
router.use('/v1', openai);

/**
 * Guest-capable chat completion router.
 *
 * Auth here accepts either a real JWT or a guest token (`requireGuestOrJwtAuth`);
 * when guest chat is off, or the token isn't a valid guest token, it falls through
 * to the strict `requireJwtAuth`. `enforceGuestScope` then pins guests to the free
 * Zen endpoint/model and strips every other capability (agents/tools/files/spec/
 * preset), and `guestMessageLimiter` enforces the per-IP guest quota (real users
 * are skipped). A signed-in user still bills to their own org via their forwarded
 * bearer; a guest uses the shared, capped guest key (see custom/initialize.ts).
 * Every OTHER agents route (management, CRUD, files) stays gated by the strict
 * `requireJwtAuth` below, which rejects guest tokens.
 */
const RESERVED_CHAT_SUBPATHS = new Set(['stream', 'active', 'status', 'abort']);

const chatRouter = express.Router();
/**
 * Defer reserved management subpaths (stream/active/status/abort) to the
 * JWT-only handlers defined on the parent router. Without this guard the
 * completion route's `POST /:endpoint` would shadow `POST /chat/abort`.
 */
chatRouter.use((req, res, next) => {
  const subpath = req.path.split('/').filter(Boolean)[0];
  if (RESERVED_CHAT_SUBPATHS.has(subpath)) {
    return next('router');
  }
  return next();
});
chatRouter.use(requireGuestOrJwtAuth);
chatRouter.use(checkBan);
chatRouter.use(uaParser);
chatRouter.use(configMiddleware);
chatRouter.use(enforceGuestScope);
chatRouter.use(guestMessageLimiter);

if (isEnabled(LIMIT_MESSAGE_IP)) {
  chatRouter.use(messageIpLimiter);
}

if (isEnabled(LIMIT_MESSAGE_USER)) {
  chatRouter.use(messageUserLimiter);
}

chatRouter.use('/', chat);

router.use('/chat', chatRouter);

/**
 * Guest-safe active-jobs poll. Guests never own a generation job, so this returns
 * an empty set without a DB/user lookup. Registered BEFORE the strict JWT guard
 * below (and reachable because 'active' is a RESERVED_CHAT_SUBPATH deferred out of
 * chatRouter) so the composer's bootstrap poll doesn't 401-loop for guests; every
 * other agents route stays JWT-only and rejects guest tokens. Non-guests fall
 * through to the authoritative `/chat/active` handler further down.
 */
router.get('/chat/active', requireGuestOrJwtAuth, async (req, res, next) => {
  if (req.user?.guest === true) {
    return res.json({ activeJobIds: [] });
  }
  return next();
});

/**
 * Guest-safe SSE read-back: a guest MUST be able to read its OWN generation
 * stream. All chat streaming is unified under this route; reached via the
 * reserved-subpath deferral ('stream'), it is registered with
 * `requireGuestOrJwtAuth` ABOVE the strict JWT guard so a guest token is accepted
 * (the strict `jwt` strategy would 401 it → the streamed reply never reaches the
 * guest browser → empty bubble). `streamHandler` (hoisted, defined below) pins
 * every caller to their OWN job (`job.metadata.userId === req.user.id` → foreign
 * 403, missing 404), so a guest can never read another principal's stream.
 */
router.get('/chat/stream/:streamId', requireGuestOrJwtAuth, streamHandler);

/**
 * The other two halves of reading and stopping one's OWN generation, on the same
 * terms as the stream above.
 *
 * All four reserved subpaths are deferred out of the guest-capable chatRouter,
 * but only `active` and `stream` were given a handler above the strict guard, so
 * `status` and `abort` reached `requireJwtAuth` and answered a guest 401. That is
 * a guest who cannot learn whether their own reply is still running — so
 * resuming after a reload never worked for them — and cannot stop it once it is.
 *
 * Both handlers already pin every caller to their own job by the same test the
 * stream uses (`job.metadata.userId !== req.user.id` → 403), and a guest's id is
 * stable for the length of an exchange, so admitting them here widens no scope.
 */
router.get('/chat/status/:conversationId', requireGuestOrJwtAuth, statusHandler);
router.post('/chat/abort', requireGuestOrJwtAuth, abortHandler);

router.use(requireJwtAuth);
router.use(checkBan);
router.use(uaParser);

/**
 * Canonical Hanzo Cloud agents (`/v1/agents`). Mounted BEFORE the legacy `/`
 * (v1) router so `/cloud/*` is not shadowed by v1's `GET /:id`. The cloud router
 * carries its own `requireJwtAuth` and forwards the user's hanzo.id bearer to
 * cloud server-side (see cloud.js). Legacy `/v1/chat/agents` CRUD stays untouched.
 */
router.use('/cloud', cloud);

router.use('/', v1);

/**
 * Stream endpoints - mounted before chatRouter to bypass rate limiters
 * These are GET requests and don't need message body validation or rate limiting
 */

/**
 * @route GET /chat/stream/:streamId
 * @desc Subscribe to an ongoing generation job's SSE stream with replay support
 * @access Private
 * @description Sends sync event with resume state, replays missed chunks, then streams live
 * @query resume=true - Indicates this is a reconnection (sends sync event)
 */
async function streamHandler(req, res) {
  const { streamId } = req.params;
  const isResume = req.query.resume === 'true';

  const job = await GenerationJobManager.getJob(streamId);
  if (!job) {
    return res.status(404).json({
      error: 'Stream not found',
      message: 'The generation job does not exist or has expired.',
    });
  }

  if (job.metadata?.userId && job.metadata.userId !== req.user.id) {
    return res.status(403).json({ error: 'Unauthorized' });
  }

  res.setHeader('Content-Encoding', 'identity');
  res.setHeader('Content-Type', 'text/event-stream');
  res.setHeader('Cache-Control', 'no-cache, no-transform');
  res.setHeader('Connection', 'keep-alive');
  res.setHeader('X-Accel-Buffering', 'no');
  res.flushHeaders();

  logger.debug(`[AgentStream] Client subscribed to ${streamId}, resume: ${isResume}`);

  // Send sync event with resume state for ALL reconnecting clients
  // This supports multi-tab scenarios where each tab needs run step data
  if (isResume) {
    const resumeState = await GenerationJobManager.getResumeState(streamId);
    if (resumeState && !res.writableEnded) {
      // Send sync event with run steps AND aggregatedContent
      // Client will use aggregatedContent to initialize message state
      res.write(`event: message\ndata: ${JSON.stringify({ sync: true, resumeState })}\n\n`);
      if (typeof res.flush === 'function') {
        res.flush();
      }
      logger.debug(
        `[AgentStream] Sent sync event for ${streamId} with ${resumeState.runSteps.length} run steps`,
      );
    }
  }

  const result = await GenerationJobManager.subscribe(
    streamId,
    (event) => {
      if (!res.writableEnded) {
        res.write(`event: message\ndata: ${JSON.stringify(event)}\n\n`);
        if (typeof res.flush === 'function') {
          res.flush();
        }
      }
    },
    (event) => {
      if (!res.writableEnded) {
        res.write(`event: message\ndata: ${JSON.stringify(event)}\n\n`);
        if (typeof res.flush === 'function') {
          res.flush();
        }
        res.end();
      }
    },
    (error) => {
      if (!res.writableEnded) {
        res.write(`event: error\ndata: ${JSON.stringify({ error })}\n\n`);
        if (typeof res.flush === 'function') {
          res.flush();
        }
        res.end();
      }
    },
  );

  if (!result) {
    return res.status(404).json({ error: 'Failed to subscribe to stream' });
  }

  req.on('close', () => {
    logger.debug(`[AgentStream] Client disconnected from ${streamId}`);
    result.unsubscribe();
  });
}

/**
 * @route GET /chat/active
 * @desc Get all active generation job IDs for the current user
 * @access Private
 * @returns { activeJobIds: string[] }
 */
router.get('/chat/active', async (req, res) => {
  const activeJobIds = await GenerationJobManager.getActiveJobIdsForUser(req.user.id);
  res.json({ activeJobIds });
});

/**
 * @route GET /chat/status/:conversationId
 * @desc Check if there's an active generation job for a conversation
 * @access Private
 * @returns { active, streamId, status, aggregatedContent, createdAt, resumeState }
 */
async function statusHandler(req, res) {
  const { conversationId } = req.params;

  // streamId === conversationId, so we can use getJob directly
  const job = await GenerationJobManager.getJob(conversationId);

  if (!job) {
    return res.json({ active: false });
  }

  if (job.metadata.userId !== req.user.id) {
    return res.status(403).json({ error: 'Unauthorized' });
  }

  // Get resume state which contains aggregatedContent
  // Avoid calling both getStreamInfo and getResumeState (both fetch content)
  const resumeState = await GenerationJobManager.getResumeState(conversationId);
  const isActive = job.status === 'running';

  res.json({
    active: isActive,
    streamId: conversationId,
    status: job.status,
    aggregatedContent: resumeState?.aggregatedContent ?? [],
    createdAt: job.createdAt,
    resumeState,
  });
}

/**
 * @route POST /chat/abort
 * @desc Abort an ongoing generation job
 * @access Private
 * @description Mounted before chatRouter to bypass buildEndpointOption middleware
 */
async function abortHandler(req, res) {
  logger.debug(`[AgentStream] ========== ABORT ENDPOINT HIT ==========`);
  logger.debug(`[AgentStream] Method: ${req.method}, Path: ${req.path}`);
  logger.debug(`[AgentStream] Body:`, req.body);

  const { streamId, conversationId, abortKey } = req.body;
  const userId = req.user?.id;

  // streamId === conversationId, so try any of the provided IDs
  // Skip "new" as it's a placeholder for new conversations, not an actual ID
  let jobStreamId =
    streamId || (conversationId !== 'new' ? conversationId : null) || abortKey?.split(':')[0];
  let job = jobStreamId ? await GenerationJobManager.getJob(jobStreamId) : null;

  // Fallback: if job not found and we have a userId, look up active jobs for user
  // This handles the case where frontend sends "new" but job was created with a UUID
  if (!job && userId) {
    logger.debug(`[AgentStream] Job not found by ID, checking active jobs for user: ${userId}`);
    const activeJobIds = await GenerationJobManager.getActiveJobIdsForUser(userId);
    if (activeJobIds.length > 0) {
      // Abort the most recent active job for this user
      jobStreamId = activeJobIds[0];
      job = await GenerationJobManager.getJob(jobStreamId);
      logger.debug(`[AgentStream] Found active job for user: ${jobStreamId}`);
    }
  }

  logger.debug(`[AgentStream] Computed jobStreamId: ${jobStreamId}`);

  if (job && jobStreamId) {
    if (job.metadata?.userId && job.metadata.userId !== userId) {
      logger.warn(`[AgentStream] Unauthorized abort attempt for ${jobStreamId} by user ${userId}`);
      return res.status(403).json({ error: 'Unauthorized' });
    }

    logger.debug(`[AgentStream] Job found, aborting: ${jobStreamId}`);
    const abortResult = await GenerationJobManager.abortJob(jobStreamId);
    logger.debug(`[AgentStream] Job aborted successfully: ${jobStreamId}`, {
      abortResultSuccess: abortResult.success,
      abortResultUserMessageId: abortResult.jobData?.userMessage?.messageId,
      abortResultResponseMessageId: abortResult.jobData?.responseMessageId,
    });

    // CRITICAL: Save partial response BEFORE returning to prevent race condition.
    // If user sends a follow-up immediately after abort, the parentMessageId must exist in DB.
    // Only save if we have a valid responseMessageId (skip early aborts before generation started)
    if (
      abortResult.success &&
      abortResult.jobData?.userMessage?.messageId &&
      abortResult.jobData?.responseMessageId
    ) {
      const { jobData, content, text } = abortResult;
      const responseMessage = {
        messageId: jobData.responseMessageId,
        parentMessageId: jobData.userMessage.messageId,
        conversationId: jobData.conversationId,
        content: content || [],
        text: text || '',
        sender: jobData.sender || 'AI',
        endpoint: jobData.endpoint,
        model: jobData.model,
        unfinished: true,
        error: false,
        isCreatedByUser: false,
        user: userId,
      };

      try {
        await saveMessage(req, responseMessage, {
          context: 'api/server/routes/agents/index.js - abort endpoint',
        });
        logger.debug(`[AgentStream] Saved partial response for: ${jobStreamId}`);
      } catch (saveError) {
        logger.error(`[AgentStream] Failed to save partial response: ${saveError.message}`);
      }
    }

    return res.json({ success: true, aborted: jobStreamId });
  }

  logger.warn(`[AgentStream] Job not found for streamId: ${jobStreamId}`);
  return res.status(404).json({ error: 'Job not found', streamId: jobStreamId });
}

module.exports = router;