TJWaterAgent/src/server.ts

import { randomUUID } from "node:crypto";
import cors from "cors";
import express from "express";

import { SessionHistoryStore } from "./history/store.js";
import { ChatSessionBridge } from "./chat/sessionBridge.js";
import { config } from "./config.js";
import { logger } from "./logger.js";
import { LearningOrchestrator } from "./learning/orchestrator.js";
import { MemoryStore } from "./memory/store.js";
import { ResultReferenceResolver } from "./results/resolver.js";
import { ResultReferenceStore } from "./results/store.js";
import { buildChatRouter } from "./routes/chat.js";
import { opencodeRuntime } from "./runtime/opencode.js";
import { SessionRegistry } from "./session/registry.js";
import { ToolSessionContextStore } from "./session/toolContextStore.js";
import { DynamicHttpExecutor } from "./tools/dynamicHttpExecutor.js";

const app = express();
const registry = new SessionRegistry(config.SESSION_TTL_SECONDS);
const sessionBridge = new ChatSessionBridge(registry, opencodeRuntime);
const memoryStore = new MemoryStore();
const sessionHistoryStore = new SessionHistoryStore();
const toolContextStore = new ToolSessionContextStore();
const learningOrchestrator = new LearningOrchestrator(
  opencodeRuntime,
  memoryStore,
  sessionHistoryStore,
);
const resultReferenceStore = new ResultReferenceStore();
const resultReferenceResolver = new ResultReferenceResolver(resultReferenceStore);
const dynamicHttpExecutor = new DynamicHttpExecutor(resultReferenceStore);
const internalToken = config.AGENT_INTERNAL_TOKEN ?? randomUUID();

// 这个 token 只用于仍需服务端上下文的工具桥（dynamic_http_call / fetch_result_ref / store_render_ref）。
process.env.TJWATER_AGENT_INTERNAL_TOKEN = internalToken;

app.use(cors());
app.use(express.json({ limit: "1mb" }));

app.get("/health", async (_req, res) => {
  try {
    const runtime = await opencodeRuntime.health();
    res.json({
      ok: true,
      runtime,
      sessions: sessionBridge.count(),
    });
  } catch (error) {
    const detail = error instanceof Error ? error.message : String(error);
    res.status(503).json({
      ok: false,
      message: "opencode runtime unavailable",
      detail,
      sessions: sessionBridge.count(),
    });
  }
});

app.post("/internal/tools/dynamic-http-call", async (req, res) => {
  if (req.header("x-agent-internal-token") !== internalToken) {
    res.status(403).json({ message: "forbidden" });
    return;
  }

  const sessionId = typeof req.body?.sessionId === "string" ? req.body.sessionId : "";
  const context = sessionBridge.getSessionContext(sessionId);
  if (!context) {
    res.status(404).json({
      message: "session context not found",
      detail: sessionId,
    });
    return;
  }

  try {
    // opencode 工具运行在 .opencode 侧，这里负责把工具调用重新绑定到当前用户/项目上下文。
    const result = await dynamicHttpExecutor.execute(
      {
        reason: req.body?.reason,
        path: req.body?.path,
        method: req.body?.method,
        arguments: req.body?.arguments,
      },
      {
        accessToken: context.accessToken,
        actorKey: context.actorKey,
        clientSessionId: context.clientSessionId,
        projectId: context.projectId,
        projectKey: context.projectKey,
        sessionId,
        traceId: context.traceId,
      },
    );
    res.json(result);
  } catch (error) {
    const detail = error instanceof Error ? error.message : String(error);
    res.status(400).json({
      message: "dynamic http execution failed",
      detail,
    });
  }
});

app.post("/internal/tools/fetch-result-ref", async (req, res) => {
  if (req.header("x-agent-internal-token") !== internalToken) {
    res.status(403).json({ message: "forbidden" });
    return;
  }

  const sessionId = typeof req.body?.sessionId === "string" ? req.body.sessionId : "";
  const resultRef = typeof req.body?.result_ref === "string" ? req.body.result_ref : "";
  const context = sessionBridge.getSessionContext(sessionId);
  if (!context) {
    res.status(404).json({
      message: "session context not found",
      detail: sessionId,
    });
    return;
  }
  if (!resultRef) {
    res.status(400).json({ message: "result_ref is required" });
    return;
  }

  const result = await resultReferenceResolver.getAuthorized(
    resultRef,
    {
      actorKey: context.actorKey,
      projectId: context.projectId,
    },
    {
      maxItems:
        typeof req.body?.max_items === "number" ? req.body.max_items : undefined,
    },
  );

  if (!result) {
    res.status(404).json({ message: "result_ref not found" });
    return;
  }

  res.json(result);
});

app.post("/internal/tools/store-render-ref", async (req, res) => {
  if (req.header("x-agent-internal-token") !== internalToken) {
    res.status(403).json({ message: "forbidden" });
    return;
  }

  const sessionId = typeof req.body?.sessionId === "string" ? req.body.sessionId : "";
  const filePath = typeof req.body?.file_path === "string" ? req.body.file_path.trim() : "";
  const context = sessionBridge.getSessionContext(sessionId);
  if (!context) {
    res.status(404).json({
      message: "session context not found",
      detail: sessionId,
    });
    return;
  }
  if (!filePath) {
    res.status(400).json({ message: "file_path is required" });
    return;
  }

  try {
    const record = await resultReferenceResolver.registerRenderPayloadFile(filePath, {
      actorKey: context.actorKey,
      clientSessionId: context.clientSessionId,
      projectId: context.projectId,
      projectKey: context.projectKey,
      sessionId,
      source: "migration",
      traceId: context.traceId,
    });
    res.json({
      ok: true,
      render_ref: record.resultRef,
      stored_at: record.createdAt,
      preview: record.preview,
      kind: record.kind,
      schema_version: record.schemaVersion,
      source: record.source,
    });
  } catch (error) {
    const detail = error instanceof Error ? error.message : String(error);
    res.status(400).json({
      message: "store render ref failed",
      detail,
    });
  }
});

app.post("/internal/tools/session-search", async (req, res) => {
  if (req.header("x-agent-internal-token") !== internalToken) {
    res.status(403).json({ message: "forbidden" });
    return;
  }

  const sessionId = typeof req.body?.sessionId === "string" ? req.body.sessionId : "";
  const query = typeof req.body?.query === "string" ? req.body.query : "";
  const context = await toolContextStore.read(sessionId);
  if (!context) {
    res.status(404).json({
      message: "tool session context not found",
      detail: sessionId,
    });
    return;
  }
  if (!query.trim()) {
    res.status(400).json({ message: "query is required" });
    return;
  }
  const hits = await sessionHistoryStore.search(
    {
      actorKey: context.actorKey,
      projectKey: context.projectKey,
    },
    query,
    typeof req.body?.max_results === "number" ? req.body.max_results : undefined,
  );
  res.json({
    hits,
    query,
  });
});

app.use(
  "/api/v1/agent/chat",
  buildChatRouter(
    sessionBridge,
    opencodeRuntime,
    memoryStore,
    learningOrchestrator,
    resultReferenceResolver,
  ),
);

const bootstrap = async () => {
  await Promise.all([
    learningOrchestrator.initialize(),
    memoryStore.initialize(),
    resultReferenceStore.initialize(),
    sessionHistoryStore.initialize(),
    toolContextStore.initialize(),
  ]);
  resultReferenceStore.startCleanupLoop();
};

await bootstrap();

const server = app.listen(config.PORT, config.HOST, () => {
  logger.info(
    { host: config.HOST, port: config.PORT },
    "TJWaterAgent listening",
  );
});

const shutdown = async () => {
  logger.info("shutting down TJWaterAgent");
  server.close();
  resultReferenceStore.stopCleanupLoop();
  // 同步关闭 embedded opencode server，避免本服务退出后留下孤儿进程。
  await opencodeRuntime.dispose();
};

process.on("SIGINT", () => {
  void shutdown();
});

process.on("SIGTERM", () => {
  void shutdown();
});