feat(agent): add credential refresh and unify learning tools

2026-08-06 10:16:50 +08:00
parent 2dc37e3fd8
commit 258f4996eb
21 changed files with 1270 additions and 770 deletions
@@ -5,6 +5,11 @@ 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";
@@ -12,6 +17,12 @@ 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 {
@@ -21,12 +32,12 @@ import {
 import { buildChatRouter } from "./routes/chat.js";
 import { buildAgentPublicRouter } from "./routes/publicApi.js";
 import { opencodeRuntime } from "./runtime/opencode.js";
-import { serializeRuntimeSessionContext } from "./runtime/internalSessionContextBridge.js";
 import {
  getRuntimeSessionContext,
  markRuntimeSessionAuthExpired,
  type RuntimeSessionContext,
 } from "./runtime/sessionContext.js";
+import { SkillStore } from "./skills/store.js";

 const app = express();

@@ -35,15 +46,18 @@ 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);
 const internalToken = config.AGENT_INTERNAL_TOKEN ?? randomUUID();
+const credentialRefreshCoordinator = new CredentialRefreshCoordinator();

 // 这个 token 只用于 OpenCode 子进程回调本服务的内部工具桥。
 process.env.TJWATER_AGENT_INTERNAL_TOKEN = internalToken;
@@ -74,7 +88,7 @@ app.get("/health", async (_req, res) => {
  }
 });

-app.post("/internal/tools/session-context", (req, res) => {
+app.post("/internal/tools/memory-manager", async (req, res) => {
  if (req.header("x-agent-internal-token") !== internalToken) {
    res.status(403).json({ message: "forbidden" });
    return;
@@ -90,8 +104,86 @@ app.post("/internal/tools/session-context", (req, res) => {
    });
    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),
+    });
+  }
+});

-  res.json(serializeRuntimeSessionContext(context));
+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) => {
@@ -110,15 +202,6 @@ app.post("/internal/tools/tjwater-cli-call", async (req, res) => {
    });
    return;
  }
-  if (isRuntimeAuthExpired(context)) {
-    markAuthExpired(context, "access_token_expired");
-    res.status(401).json({
-      message: "access token expired; refresh chat context",
-      detail: sessionId,
-    });
-    return;
-  }
-
  const command = typeof req.body?.command === "string" ? req.body.command.trim() : "";
  if (!command) {
    res.status(400).json({ message: "command is required" });
@@ -136,46 +219,35 @@ app.post("/internal/tools/tjwater-cli-call", async (req, res) => {
    return;
  }

-  const authJson = JSON.stringify({
-    server: config.TJWATER_API_BASE_URL,
-    access_token: context.accessToken,
-    project_id: context.projectId,
-  });
-
-  const cliArgs = ["--auth-stdin", ...command.split(/\s+/).filter(Boolean)];
-
-  const child = spawn(config.TJWATER_CLI_PATH, cliArgs, {
-    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(authJson);
-  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);
+  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),
    });
-    child.on("error", (err) => {
-      clearTimeout(timer);
-      reject(err);
-    });
-  });
+    return;
+  }

-  if (exitCode === -1) {
+  if (result.status === 504) {
    res.status(504).json({
      ok: false,
      schema_version: "tjwater-cli/v1",
@@ -189,29 +261,102 @@ app.post("/internal/tools/tjwater-cli-call", async (req, res) => {
    return;
  }

-  if (exitCode !== 0) {
-    res.status(502).json({
-      ok: false,
-      exit_code: exitCode,
-      stderr: stderr.slice(0, 2000),
-      stdout: stdout.slice(0, 2000),
-      message: `CLI exited with code ${exitCode}`,
-    });
+  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(stdout));
+    res.json(JSON.parse(result.stdout));
  } catch {
    res.json({
      ok: true,
      schema_version: "tjwater-cli/v1",
-      raw: stdout,
-      stderr: stderr || undefined,
+      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" });
@@ -302,39 +447,60 @@ const callBackendJson = async (
  context: RuntimeSessionContext,
  payload: unknown,
 ) => {
-  if (isRuntimeAuthExpired(context)) {
+  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: "access token expired; refresh chat context",
+        message: "credential refresh failed",
+        detail: error instanceof Error ? error.message : String(error),
      }),
    };
  }
-
-  const controller = new AbortController();
-  const timer = setTimeout(() => controller.abort(), config.TJWATER_API_TIMEOUT_MS);
-  try {
-    const headers = buildBackendContextHeaders(context);
-    const response = await fetch(new URL(path, config.TJWATER_API_BASE_URL), {
-      method: "POST",
-      headers,
-      body: JSON.stringify(payload),
-      signal: controller.signal,
-    });
-    const text = await response.text();
-    if (response.status === 401) {
-      markAuthExpired(context, "access_token_rejected");
-    }
-    return {
-      ok: response.ok,
-      status: response.status,
-      text,
-    };
-  } finally {
-    clearTimeout(timer);
-  }
 };

 const parseStringArray = (value: unknown) =>
@@ -357,19 +523,6 @@ const normalizeWebSearchFreshness = (value: unknown) => {
  return webSearchFreshnessMap[value] ?? value;
 };

-const AUTH_EXPIRY_SKEW_MS = 30_000;
-
-function isRuntimeAuthExpired(context: RuntimeSessionContext) {
-  if (!context.tokenExpiresAt) {
-    return false;
-  }
-  const expiresAt = Date.parse(context.tokenExpiresAt);
-  if (!Number.isFinite(expiresAt)) {
-    return false;
-  }
-  return Date.now() >= expiresAt - AUTH_EXPIRY_SKEW_MS;
-}
-
 function markAuthExpired(
  context: RuntimeSessionContext,
  reason: NonNullable<RuntimeSessionContext["authExpired"]>["reason"],
@@ -491,6 +644,7 @@ const chatRouter = buildChatRouter(
  sessionTranscriptStore,
  learningOrchestrator,
  resultReferenceResolver,
+  credentialRefreshCoordinator,
 );
 const authenticatedChatRouter = express.Router();
 authenticatedChatRouter.use(requireAgentAuth, chatRouter);