next-tjwater-agent/src/routes/agent-message-routes.ts

import type { Express, Request, Response, Router } from "express";
import { z } from "zod";
import { getRuntimeSessionContext } from "../runtime/sessionContext.js";
import { isTrustedInternalBridgeRequest } from "../product/internal-bridge-auth.js";
import { getLocalAgentContext, getTrustedProductCredential, isTrustedTaskDispatchRequest } from "../context/localContext.js";
import type { SessionMetadataStore } from "../sessions/metadataStore.js";
import { readPurpose } from "../sessions/metadataStore.js";
import type { AgentMessageCoordinator } from "../chat/agent-message-coordinator.js";
import { logger } from "../logger.js";

const uuid = z.string().uuid();
const paging = z.object({ cursor: z.string().max(200).optional(), limit: z.number().int().min(1).max(50).optional() }).strict();
const discovery = z.object({ limit: z.number().int().min(1).max(50).optional(), offset: z.number().int().min(0).optional() }).strict();
const text = z.string().trim().min(1).max(8000);
const errorResponse = (res: Response, error: unknown) => {
  const status = error instanceof z.ZodError ? 400
    : error && typeof error === "object" && "status" in error && typeof error.status === "number" ? error.status : 503;
  res.status(status).json({ message: status === 503 ? "Agent 消息服务暂不可用" : error instanceof Error ? error.message : "Agent message rejected" });
};
const deny = (message: string) => Object.assign(new Error(message), { status: 403 });

// Called by the local MCP service, with a host-resolved native session only.
export function createAgentMessageQuery({ sessions, messages }: { sessions: SessionMetadataStore; messages: AgentMessageCoordinator }) {
  return async (sessionId: string, operation: "list" | "get" | "read", value: unknown) => {
    const context = getRuntimeSessionContext(sessionId);
    if (!context?.projectId || context.learningMode === "review") throw deny("project conversation required");
    const record = await sessions.get(context, sessionId);
    if (!record || !["ordinary_chat", "task_collaboration", "task_execution"].includes(readPurpose(record))) throw deny("conversation unavailable");
    const execution = readPurpose(record) === "task_execution";
    if (execution && (operation !== "read" || !messages.isActive(record.sessionId))) throw deny("active investigation may only read its own messages");
    if (operation === "list") {
      const input = discovery.parse(value);
      const query = new URLSearchParams({ task_type: "event_investigation", limit: String(input.limit ?? 20), offset: String(input.offset ?? 0) });
      return messages.client.request(record.sessionId, `/api/v1/operational-tasks?${query}`);
    }
    if (operation === "get") {
      const input = z.object({ task_id: uuid, run_id: uuid.optional() }).strict().parse(value);
      const task = await messages.client.request(record.sessionId, `/api/v1/operational-tasks/${input.task_id}/agent-context`);
      const runs = await messages.client.request(record.sessionId, `/api/v1/operational-tasks/${input.task_id}/runs${input.run_id ? `/${input.run_id}` : ""}`);
      return { context: task, ...(input.run_id ? { run: runs } : { runs }) };
    }
    const input = paging.parse(value);
    return messages.client.list(record.sessionId, record.sessionId, input.cursor, input.limit ?? 20);
  };
}

export function registerAgentMessageBridge(app: Express, options: {
  internalToken: string; sessions: SessionMetadataStore; messages: AgentMessageCoordinator;
}) {
  const { sessions, messages } = options;
  app.post("/internal/agent-messages/:operation", async (req, res) => {
    if (!isTrustedInternalBridgeRequest(req, options.internalToken)) { res.status(403).json({ message: "forbidden" }); return; }
    try {
      const envelope = z.object({ session_id: z.string().min(1).max(128), input: z.unknown() }).strict().parse(req.body);
      const context = getRuntimeSessionContext(envelope.session_id);
      if (!context?.projectId || context.learningMode === "review") throw deny("project conversation required");
      const record = await sessions.get(context, envelope.session_id);
      if (!record || !["ordinary_chat", "task_collaboration", "task_execution"].includes(readPurpose(record))) throw deny("conversation unavailable");
      const sessionId = record.sessionId;
      const execution = readPurpose(record) === "task_execution";
      const operation = req.params.operation;
      if (execution && !["read", "reply"].includes(operation)) throw deny("investigation may only read and reply to its own messages");
      if (execution && !messages.isActive(sessionId)) throw deny("active investigation required");
      if (operation === "send") {
        const input = z.object({ task_id: uuid, run_id: uuid, message: text, request_id: uuid }).strict().parse(envelope.input);
        const item = await messages.client.send(sessionId, input);
        messages.notify(item);
        // Active target has its own current service credential. Never lend the
        // caller's credential to an investigation session.
        if (item.target_session_id && messages.isActive(item.target_session_id)) {
          void messages.deliver(item.target_session_id, item).catch(() => logger.warn("Agent message retained for worker delivery"));
        }
        res.json(item);
      } else if (operation === "reply") {
        const input = z.object({ message_id: uuid, message: text, request_id: uuid }).strict().parse(envelope.input);
        const original = await messages.client.get(sessionId, input.message_id);
        if (execution && (original.task_id !== record.taskId || original.run_id !== record.runId)) throw deny("message does not belong to this investigation run");
        if (!execution && (original.target_session_id !== sessionId || original.direction !== "to_chat")) throw deny("reply must belong to the calling conversation");
        const item = await messages.client.reply(sessionId, input);
        messages.notify(item);
        // Persisted already; receipt is not a promise that the receiver read it.
        const deliveryCredential = item.direction === "to_task" && item.target_session_id && messages.isActive(item.target_session_id)
          ? item.target_session_id : sessionId;
        void messages.deliver(deliveryCredential, item).catch(() => logger.warn("Agent reply retained for worker delivery"));
        res.json(item);
      } else res.status(404).json({ message: "unknown message operation" });
    } catch (error) { errorResponse(res, error); }
  });
}

export function registerAgentMessageRoutes(router: Router, options: {
  sessions: SessionMetadataStore; messages: AgentMessageCoordinator;
}) {
  const { sessions, messages } = options;
  const authorized = async (req: Request) => {
    const context = getLocalAgentContext(req);
    const credential = getTrustedProductCredential(req);
    if (!credential) throw deny("project authorization required");
    const id = z.string().min(1).max(128).parse(req.params.session_id);
    const record = await sessions.get(context, id);
    if (!record || !["ordinary_chat", "task_collaboration"].includes(readPurpose(record))) throw deny("conversation unavailable");
    return { record, credential };
  };
  router.get("/session/:session_id/agent-messages", async (req, res) => {
    try {
      const { record, credential } = await authorized(req);
      const query = z.object({ cursor: z.string().max(200).optional(), limit: z.coerce.number().int().min(1).max(50).default(50) }).parse(req.query);
      res.setHeader("Cache-Control", "no-store");
      res.json(await messages.client.list(credential, record.sessionId, query.cursor, query.limit));
    } catch (error) { errorResponse(res, error); }
  });
  router.get("/session/:session_id/agent-messages/stream", async (req, res) => {
    try {
      const { record, credential } = await authorized(req);
      // Also verify current Server permissions before creating the subscription.
      await messages.client.list(credential, record.sessionId, undefined, 1);
      res.set({ "Content-Type": "text/event-stream", "Cache-Control": "no-cache, no-transform", Connection: "keep-alive" });
      res.flushHeaders();
      const notify = () => { if (!res.destroyed) res.write(`event: agent_message\ndata: ${JSON.stringify({ session_id: record.sessionId })}\n\n`); };
      const unsubscribe = messages.subscribe(record.sessionId, notify);
      notify();
      const heartbeat = setInterval(() => { if (!res.destroyed) res.write(": heartbeat\n\n"); }, 15_000);
      const renewal = setTimeout(() => { res.write("event: reconnect\ndata: {}\n\n"); res.end(); }, 60_000);
      res.on("close", () => { unsubscribe(); clearInterval(heartbeat); clearTimeout(renewal); });
    } catch (error) { if (!res.headersSent) errorResponse(res, error); else res.end(); }
  });
  router.post("/task-messages/deliver", async (req, res) => {
    try {
      const credential = getTrustedProductCredential(req);
      if (!credential || !isTrustedTaskDispatchRequest(req)) throw deny("service task execution authorization required");
      const input = z.object({ message_id: uuid }).strict().parse(req.body);
      const item = await messages.client.get(credential, input.message_id);
      if (item.project_id !== credential.projectId) throw deny("message project mismatch");
      res.json(await messages.deliver(credential, item));
    } catch (error) { errorResponse(res, error); }
  });
}