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

import type { TaskCalculationWaiter } from "../mcp/task-calculation-waiter.js";
import { registerInteractionRequestRoutes } from "./interaction-request-routes.js";
import type { InteractionRequestStore } from "../sessions/interaction-request-store.js";
import { Router, type NextFunction, type Request, type Response } from "express";

import { agentModelOptions, resolveDefaultModel } from "../chat/models.js";
import { type ChatSessionBridge } from "../chat/sessionBridge.js";
import { config } from "../config.js";
import { type FrontendActionCoordinator } from "../frontendAction/coordinator.js";
import { frontendActionRegistryResponse } from "../frontendAction/registry.js";
import { type LearningOrchestrator } from "../learning/orchestrator.js";
import { type MemoryStore } from "../memory/store.js";
import { type OpencodeRuntimeAdapter } from "../runtime/opencode.js";
import { type SessionMetadataStore } from "../sessions/metadataStore.js";
import { SessionContextMismatchError } from "../sessions/metadataStore.js";
import { type SessionTranscriptStore } from "../sessions/transcriptStore.js";
import { type SessionUiStateStore } from "../sessions/uiStateStore.js";
import { uiRegistryResponse } from "../uiEnvelope/registry.js";
import {
  attachTrustedAgentContext,
  getLocalAgentContext,
  ProductContextError,
} from "../context/localContext.js";
import { registerChatAuxiliaryRoutes } from "./chatAuxiliaryRoutes.js";
import { registerChatInteractionRoutes } from "./chatInteractionRoutes.js";
import { createChatRunState } from "./chatRunState.js";
import { registerChatSessionRoutes } from "./chatSessionRoutes.js";
import { registerChatStreamRoute } from "./chatStreamRoute.js";
import { registerResponseInputRoutes } from "./response-input-routes.js";
import { registerMcpStatusRoutes } from "./mcpRoutes.js";
import { registerChatSuggestionRoutes } from "./chatSuggestionRoutes.js";
import { registerAgentMessageRoutes } from "./agent-message-routes.js";
import { registerArtifactRoutes } from "./artifact-routes.js";
import type { AgentMessageCoordinator } from "../chat/agent-message-coordinator.js";

export type BuildChatRouterOptions = {
  waitTaskCalculations?: TaskCalculationWaiter;
  interactionRequestStore?: InteractionRequestStore;
  agentMessages?: AgentMessageCoordinator;
  frontendActionCoordinator: FrontendActionCoordinator;
  learningOrchestrator: LearningOrchestrator;
  memoryStore: MemoryStore;
  runtime: OpencodeRuntimeAdapter;
  sessionBridge: ChatSessionBridge;
  sessionMetadataStore: SessionMetadataStore;
  sessionTranscriptStore: SessionTranscriptStore;
  sessionUiStateStore: SessionUiStateStore;
};

const runState = createChatRunState();

export const buildChatRouter = (options: BuildChatRouterOptions) => {
  const {
    frontendActionCoordinator,
    learningOrchestrator,
    memoryStore,
    runtime,
    sessionBridge,
    sessionMetadataStore,
    sessionTranscriptStore,
    sessionUiStateStore,
    agentMessages,
  } = options;
  const chatRouter = Router();
  const { activeRuns, lastRunStatuses } = runState;

  chatRouter.use((req, res, next) => {
    void attachTrustedAgentContext(req)
      .then(() => next())
      .catch((error) => {
        const status = error instanceof ProductContextError ? error.status : 503;
        res.status(status).json({
          code: "PRODUCT_CONTEXT_UNAVAILABLE",
          message: error instanceof Error ? error.message : "product identity verification failed",
        });
      });
  });

  chatRouter.get("/models", (_req, res) => {
    res.json({
      default_model: resolveDefaultModel(config.OPENCODE_MODEL),
      models: agentModelOptions,
    });
  });
  if (options.interactionRequestStore) registerInteractionRequestRoutes(chatRouter, { sessions: sessionMetadataStore, interactions: options.interactionRequestStore,
    runtime, isRunning: (sessionId) => activeRuns.get(sessionId)?.status === "running" });
  registerArtifactRoutes(chatRouter, { sessionMetadataStore, runtime });

  chatRouter.get("/ui-registry", (_req, res) => {
    res.json(uiRegistryResponse);
  });

  chatRouter.get("/frontend-action-registry", (_req, res) => {
    res.json(frontendActionRegistryResponse);
  });

  chatRouter.post("/frontend-actions/:action_id/result", async (req, res) => {
    const sessionId = req.header("x-agent-session-id")?.trim() ?? "";
    if (!sessionId || req.params.action_id !== req.body?.actionId) {
      res.status(400).json({ message: "invalid frontend action result" });
      return;
    }
    try {
      const context = getLocalAgentContext(req);
      const sessionRecord = await sessionMetadataStore.get(
        {
          actorKey: context.actorKey,
          projectId: context.projectId,
          projectKey: context.projectKey,
          userId: context.userId,
        },
        sessionId,
      );
      if (!sessionRecord) {
        res.status(404).json({ message: "session not found" });
        return;
      }
      const outcome = frontendActionCoordinator.submitResult(sessionId, req.body);
      res.json({
        accepted: true,
        duplicate: outcome.duplicate,
        result: outcome.result,
      });
    } catch (error) {
      const message =
        error instanceof Error
          ? error.message
          : "frontend action result rejected";
      const status =
        typeof error === "object" &&
        error &&
        "status" in error &&
        typeof error.status === "number"
          ? error.status
          : 400;
      res.status(status).json({ message });
    }
  });

  registerChatSessionRoutes(chatRouter, {
    interactionRequestStore: options.interactionRequestStore,
    learningOrchestrator,
    frontendActionCoordinator,
    runState,
    runtime,
    sessionBridge,
    sessionMetadataStore,
    sessionTranscriptStore,
    sessionUiStateStore,
  });
  registerMcpStatusRoutes(chatRouter, { sessionMetadataStore, runtime });
  registerChatSuggestionRoutes(chatRouter, runtime);
  if (agentMessages) registerAgentMessageRoutes(chatRouter, { sessions: sessionMetadataStore, messages: agentMessages });

  registerChatAuxiliaryRoutes(chatRouter, {
    interactionRequestStore: options.interactionRequestStore,
    runtime,
    activeRuns,
    lastRunStatuses,
    sessionBridge,
    sessionMetadataStore,
    sessionUiStateStore,
    frontendActionCoordinator,
  });

  registerChatInteractionRoutes(chatRouter, {
    activeRuns,
    runtime,
    sessionMetadataStore,
    sessionUiStateStore,
  });

  registerResponseInputRoutes(chatRouter, { runtime, sessionMetadataStore, runState });
  registerChatStreamRoute(chatRouter, {
    waitTaskCalculations: options.waitTaskCalculations,
    interactionRequestStore: options.interactionRequestStore,
    agentMessages,
    frontendActionCoordinator,
    learningOrchestrator,
    memoryStore,
    runState,
    runtime,
    sessionBridge,
    sessionMetadataStore,
    sessionTranscriptStore,
    sessionUiStateStore,
  });

  chatRouter.use((error: unknown, _req: Request, res: Response, next: NextFunction) => {
    if (error instanceof SessionContextMismatchError && !res.headersSent) {
      res.status(404).json({ message: "session not found" });
      return;
    }
    next(error);
  });
  return chatRouter;
};