Files
TJWaterAgent/src/server.ts
T

728 lines
20 KiB
TypeScript

import { randomUUID } from "node:crypto";
import { spawn } from "node:child_process";
import cors from "cors";
import express from "express";
import { requireAgentAuth } from "./auth/agentAuth.js";
import { buildBackendContextHeaders } from "./auth/backendContextHeaders.js";
import {
CredentialRefreshError,
CredentialRefreshCoordinator,
runWithCredentialRefresh,
} from "./auth/credentialRefresh.js";
import { SessionTranscriptStore } from "./sessions/transcriptStore.js";
import { ChatSessionBridge } from "./chat/sessionBridge.js";
import { config } from "./config.js";
import { SessionUiStateStore } from "./sessions/uiStateStore.js";
import { SessionMetadataStore } from "./sessions/metadataStore.js";
import { logger } from "./logger.js";
import { LearningOrchestrator } from "./learning/orchestrator.js";
import {
executeMemoryManager,
executeSkillManager,
type MemoryManagerInput,
type SkillManagerInput,
} from "./learning/toolManagers.js";
import { MemoryStore } from "./memory/store.js";
import { ResultReferenceResolver } from "./results/resolver.js";
import {
RESULT_REFERENCE_SOURCE,
ResultReferenceStore,
} from "./results/store.js";
import { buildChatRouter } from "./routes/chat.js";
import { buildAgentPublicRouter } from "./routes/publicApi.js";
import { opencodeRuntime } from "./runtime/opencode.js";
import {
getRuntimeSessionContext,
markRuntimeSessionAuthExpired,
type RuntimeSessionContext,
} from "./runtime/sessionContext.js";
import { ensureDirectory } from "./utils/fileStore.js";
import { SkillStore } from "./skills/store.js";
const app = express();
// 这里集中组装 Agent 服务的运行期依赖,路由层只通过接口调用,便于测试时替换实现。
const sessionBridge = new ChatSessionBridge(opencodeRuntime);
const sessionMetadataStore = new SessionMetadataStore();
const sessionUiStateStore = new SessionUiStateStore();
const memoryStore = new MemoryStore();
const skillStore = new SkillStore();
const sessionTranscriptStore = new SessionTranscriptStore();
const learningOrchestrator = new LearningOrchestrator(
opencodeRuntime,
memoryStore,
sessionTranscriptStore,
skillStore,
);
const resultReferenceStore = new ResultReferenceStore();
const resultReferenceResolver = new ResultReferenceResolver(
resultReferenceStore,
config.RESULT_REF_IMPORT_DIR,
config.RESULT_REF_IMPORT_MAX_BYTES,
);
const internalToken = config.AGENT_INTERNAL_TOKEN ?? randomUUID();
const credentialRefreshCoordinator = new CredentialRefreshCoordinator();
// 这个 token 只用于 OpenCode 子进程回调本服务的内部工具桥。
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,
ready: true,
warmed_up: true,
runtime,
sessions: sessionBridge.count(),
});
} catch (error) {
const detail = error instanceof Error ? error.message : String(error);
res.status(503).json({
ok: false,
ready: false,
warmed_up: true,
message: "opencode runtime unavailable",
detail,
sessions: sessionBridge.count(),
});
}
});
app.post("/internal/tools/memory-manager", async (req, res) => {
if (req.header("x-agent-internal-token") !== internalToken) {
res.status(403).json({ message: "forbidden" });
return;
}
const sessionId =
typeof req.body?.session_id === "string" ? req.body.session_id.trim() : "";
const context = sessionId ? getRuntimeSessionContext(sessionId) : null;
if (!context) {
res.status(404).json({
message: "session context not found",
detail: sessionId,
});
return;
}
const action = req.body?.action;
if (
typeof action !== "string" ||
!["add", "list", "replace", "remove"].includes(action) ||
typeof req.body?.scope !== "string"
) {
res.status(400).json({ message: "invalid memory manager request" });
return;
}
try {
res.json(
await executeMemoryManager(memoryStore, context, {
action: action as MemoryManagerInput["action"],
content:
typeof req.body?.content === "string" ? req.body.content : undefined,
scope: req.body.scope,
target_id:
typeof req.body?.target_id === "string" ? req.body.target_id : undefined,
}),
);
} catch (error) {
res.status(500).json({
message: "memory manager failed",
detail: error instanceof Error ? error.message : String(error),
});
}
});
app.post("/internal/tools/skill-manager", async (req, res) => {
if (req.header("x-agent-internal-token") !== internalToken) {
res.status(403).json({ message: "forbidden" });
return;
}
const sessionId =
typeof req.body?.session_id === "string" ? req.body.session_id.trim() : "";
const context = sessionId ? getRuntimeSessionContext(sessionId) : null;
if (!context) {
res.status(404).json({ message: "session context not found", detail: sessionId });
return;
}
const action = req.body?.action;
if (
typeof action !== "string" ||
![
"list",
"write_skill",
"remove_skill",
"append_pattern",
"remove_pattern",
"write_reference",
"remove_reference",
"write_script",
"remove_script",
].includes(action) ||
typeof req.body?.skill_path !== "string"
) {
res.status(400).json({ message: "invalid skill manager request" });
return;
}
try {
res.json(
await executeSkillManager(skillStore, context, {
action: action as SkillManagerInput["action"],
content:
typeof req.body?.content === "string" ? req.body.content : undefined,
file_path:
typeof req.body?.file_path === "string" ? req.body.file_path : undefined,
pattern:
typeof req.body?.pattern === "string" ? req.body.pattern : undefined,
skill_path: req.body.skill_path,
target_id:
typeof req.body?.target_id === "string" ? req.body.target_id : undefined,
}),
);
} catch (error) {
res.status(500).json({
message: "skill manager failed",
detail: error instanceof Error ? error.message : String(error),
});
}
});
app.post("/internal/tools/tjwater-cli-call", async (req, res) => {
if (req.header("x-agent-internal-token") !== internalToken) {
res.status(403).json({ message: "forbidden" });
return;
}
const sessionId =
typeof req.body?.session_id === "string" ? req.body.session_id.trim() : "";
const context = sessionId ? getRuntimeSessionContext(sessionId) : null;
if (!context) {
res.status(404).json({
message: "session context not found",
detail: sessionId,
});
return;
}
const command = typeof req.body?.command === "string" ? req.body.command.trim() : "";
if (!command) {
res.status(400).json({ message: "command is required" });
return;
}
const timeoutSec =
typeof req.body?.timeout === "number" && req.body.timeout > 0 ? req.body.timeout : 120;
if (!context.network) {
res.status(400).json({
message: "runtime network missing; refresh chat context",
detail: sessionId,
});
return;
}
let result;
try {
result = await runWithCredentialRefresh(
credentialRefreshCoordinator,
context,
(activeContext) => executeCliCommand(activeContext, command, timeoutSec),
);
} catch (error) {
if (!(error instanceof CredentialRefreshError)) {
const detail = error instanceof Error ? error.message : String(error);
res.status(502).json({
message: "CLI execution failed",
detail,
});
return;
}
if (error.code === "cancelled") {
res.status(409).json({ message: "agent run was aborted" });
return;
}
markAuthExpired(context, "access_token_expired");
res.status(401).json({
message: "credential refresh failed",
detail: error instanceof Error ? error.message : String(error),
});
return;
}
if (result.status === 504) {
res.status(504).json({
ok: false,
schema_version: "tjwater-cli/v1",
summary: "命令超时",
error: {
code: "TIMEOUT",
message: `command timed out after ${timeoutSec}s`,
retryable: true,
},
});
return;
}
if (result.status === 401) {
markAuthExpired(
getRuntimeSessionContext(sessionId) ?? context,
"access_token_rejected",
);
}
if (result.exitCode !== 0) {
res
.status(result.status)
.type("application/json")
.send(
result.stdout ||
JSON.stringify({
ok: false,
exit_code: result.exitCode,
stderr: result.stderr.slice(0, 2000),
message: `CLI exited with code ${result.exitCode}`,
}),
);
return;
}
try {
res.json(JSON.parse(result.stdout));
} catch {
res.json({
ok: true,
schema_version: "tjwater-cli/v1",
raw: result.stdout,
stderr: result.stderr || undefined,
});
}
});
const executeCliCommand = async (
context: RuntimeSessionContext,
command: string,
timeoutSec: number,
) => {
const child = spawn(
config.TJWATER_CLI_PATH,
["--auth-stdin", ...command.split(/\s+/).filter(Boolean)],
{ stdio: ["pipe", "pipe", "pipe"] },
);
let stdout = "";
let stderr = "";
child.stdout.on("data", (data: Buffer) => {
stdout += data.toString("utf-8");
});
child.stderr.on("data", (data: Buffer) => {
stderr += data.toString("utf-8");
});
child.stdin.write(
JSON.stringify({
server: config.TJWATER_API_BASE_URL,
access_token: context.accessToken,
project_id: context.projectId,
}),
);
child.stdin.end();
const exitCode = await new Promise<number | null>((resolve, reject) => {
const timer = setTimeout(() => {
child.kill("SIGTERM");
resolve(-1);
}, timeoutSec * 1000);
child.on("close", (code) => {
clearTimeout(timer);
resolve(code);
});
child.on("error", (error) => {
clearTimeout(timer);
reject(error);
});
});
let errorCode = "";
try {
const payload = JSON.parse(stdout) as { error?: { code?: unknown } };
errorCode =
typeof payload.error?.code === "string" ? payload.error.code : "";
} catch {
errorCode = "";
}
const status =
exitCode === -1
? 504
: errorCode === "HTTP_401" || errorCode === "UNAUTHENTICATED"
? 401
: errorCode === "HTTP_403"
? 403
: exitCode === 0
? 200
: 502;
return { exitCode, status, stderr, stdout };
};
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?.session_id === "string" ? req.body.session_id.trim() : "";
const filePath = typeof req.body?.file_path === "string" ? req.body.file_path.trim() : "";
const context = sessionId ? getRuntimeSessionContext(sessionId) : null;
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: context.clientSessionId,
source: RESULT_REFERENCE_SOURCE.agentGenerated,
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?.session_id === "string" ? req.body.session_id.trim() : "";
const query = typeof req.body?.query === "string" ? req.body.query : "";
const context = sessionId ? getRuntimeSessionContext(sessionId) : null;
if (!context) {
res.status(404).json({
message: "session context not found",
detail: sessionId,
});
return;
}
if (!query.trim()) {
res.status(400).json({ message: "query is required" });
return;
}
const hits = await sessionTranscriptStore.search(
{
actorKey: context.actorKey,
projectKey: context.projectKey,
},
query,
typeof req.body?.max_results === "number" ? req.body.max_results : undefined,
);
res.json({
hits,
query,
});
});
const callBackendJson = async (
path: string,
context: RuntimeSessionContext,
payload: unknown,
) => {
try {
const result = await runWithCredentialRefresh(
credentialRefreshCoordinator,
context,
async (activeContext) => {
const controller = new AbortController();
const timer = setTimeout(
() => controller.abort(),
config.TJWATER_API_TIMEOUT_MS,
);
try {
const response = await fetch(
new URL(path, config.TJWATER_API_BASE_URL),
{
method: "POST",
headers: buildBackendContextHeaders(activeContext),
body: JSON.stringify(payload),
signal: controller.signal,
},
);
return {
ok: response.ok,
status: response.status,
text: await response.text(),
};
} finally {
clearTimeout(timer);
}
},
);
if (result.status === 401) {
markAuthExpired(
getRuntimeSessionContext(context.sessionId) ?? context,
"access_token_rejected",
);
}
return result;
} catch (error) {
if (!(error instanceof CredentialRefreshError)) {
throw error;
}
if (error.code === "cancelled") {
throw error;
}
markAuthExpired(context, "access_token_expired");
return {
ok: false,
status: 401,
text: JSON.stringify({
message: "credential refresh failed",
detail: error instanceof Error ? error.message : String(error),
}),
};
}
};
const parseStringArray = (value: unknown) =>
Array.isArray(value)
? value.filter((item): item is string => typeof item === "string")
: undefined;
const webSearchFreshnessMap: Record<string, string> = {
no_limit: "noLimit",
one_day: "oneDay",
one_week: "oneWeek",
one_month: "oneMonth",
one_year: "oneYear",
};
const normalizeWebSearchFreshness = (value: unknown) => {
if (typeof value !== "string") {
return undefined;
}
return webSearchFreshnessMap[value] ?? value;
};
function markAuthExpired(
context: RuntimeSessionContext,
reason: NonNullable<RuntimeSessionContext["authExpired"]>["reason"],
) {
markRuntimeSessionAuthExpired(context.sessionId, reason);
void opencodeRuntime.abortSession(context.sessionId).catch((error) => {
logger.warn(
{ err: error, sessionId: context.sessionId },
"failed to abort runtime after auth expired",
);
});
}
app.post("/internal/tools/web-search", async (req, res) => {
if (req.header("x-agent-internal-token") !== internalToken) {
res.status(403).json({ message: "forbidden" });
return;
}
const sessionId =
typeof req.body?.session_id === "string" ? req.body.session_id.trim() : "";
const context = sessionId ? getRuntimeSessionContext(sessionId) : null;
if (!context) {
res.status(404).json({
message: "session context not found",
detail: sessionId,
});
return;
}
const query = typeof req.body?.query === "string" ? req.body.query.trim() : "";
if (!query) {
res.status(400).json({ message: "query is required" });
return;
}
const count =
typeof req.body?.count === "number" && Number.isFinite(req.body.count)
? Math.trunc(req.body.count)
: undefined;
const payload = {
query,
freshness: normalizeWebSearchFreshness(req.body?.freshness),
summary:
typeof req.body?.summary === "boolean" ? req.body.summary : undefined,
count,
include: parseStringArray(req.body?.include),
exclude: parseStringArray(req.body?.exclude),
};
try {
const response = await callBackendJson(
"/api/v1/web-searches",
context,
payload,
);
res
.status(response.ok ? 200 : response.status)
.type("application/json")
.send(response.text);
} catch (error) {
const detail = error instanceof Error ? error.message : String(error);
res.status(503).json({
message: "web search service unavailable",
detail,
});
}
});
app.post("/internal/tools/geocode", async (req, res) => {
if (req.header("x-agent-internal-token") !== internalToken) {
res.status(403).json({ message: "forbidden" });
return;
}
const sessionId =
typeof req.body?.session_id === "string" ? req.body.session_id.trim() : "";
const context = sessionId ? getRuntimeSessionContext(sessionId) : null;
if (!context) {
res.status(404).json({
message: "session context not found",
detail: sessionId,
});
return;
}
const keyword =
typeof req.body?.keyword === "string" ? req.body.keyword.trim() : "";
if (!keyword) {
res.status(400).json({ message: "keyword is required" });
return;
}
try {
const response = await callBackendJson(
"/api/v1/geocoding-requests",
context,
{ keyword },
);
res
.status(response.ok ? 200 : response.status)
.type("application/json")
.send(response.text);
} catch (error) {
const detail = error instanceof Error ? error.message : String(error);
res.status(503).json({
message: "geocoding service unavailable",
detail,
});
}
});
const chatRouter = buildChatRouter(
sessionBridge,
opencodeRuntime,
sessionMetadataStore,
sessionUiStateStore,
memoryStore,
sessionTranscriptStore,
learningOrchestrator,
resultReferenceResolver,
credentialRefreshCoordinator,
);
const authenticatedChatRouter = express.Router();
authenticatedChatRouter.use(requireAgentAuth, chatRouter);
app.use(
"/api/v1/agent",
buildAgentPublicRouter(authenticatedChatRouter),
);
const bootstrap = async () => {
await Promise.all([
sessionMetadataStore.initialize(),
sessionUiStateStore.initialize(),
learningOrchestrator.initialize(),
memoryStore.initialize(),
resultReferenceStore.initialize(),
ensureDirectory(config.RESULT_REF_IMPORT_DIR),
sessionTranscriptStore.initialize(),
]);
};
const warmupOpencodeRuntime = async () => {
const startedAt = Date.now();
try {
await opencodeRuntime.warmup();
logger.info(
{
elapsedMs: Math.max(0, Date.now() - startedAt),
mode: config.OPENCODE_MODE,
},
"opencode runtime warmed up",
);
} catch (error) {
logger.error(
{
err: error,
elapsedMs: Math.max(0, Date.now() - startedAt),
mode: config.OPENCODE_MODE,
},
"failed to warm up opencode runtime",
);
throw error;
}
};
await bootstrap();
await warmupOpencodeRuntime();
resultReferenceStore.startCleanupLoop();
const server = app.listen(config.PORT, config.HOST, () => {
logger.info(
{
host: config.HOST,
port: config.PORT,
ready: true,
warmedUp: true,
},
"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();
});