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

src/routes/chat.ts

@@ -2,6 +2,7 @@ import { Router } from "express";
 import { z } from "zod";
 
 import { getAgentAuthContext } from "../auth/agentAuth.js";
+import { type CredentialRefreshCoordinator } from "../auth/credentialRefresh.js";
 import {
   agentModelOptions,
   isSupportedModel,
@@ -121,6 +122,7 @@ export const buildChatRouter = (
   sessionTranscriptStore: SessionTranscriptStore,
   learningOrchestrator: LearningOrchestrator,
   resultReferenceResolver: ResultReferenceResolver,
+  credentialRefreshCoordinator: CredentialRefreshCoordinator,
 ) => {
   const chatRouter = Router();
 
@@ -295,6 +297,16 @@ export const buildChatRouter = (
       },
     };
     run.subscribers.add(subscriber);
+    const pendingCredentialRefresh =
+      credentialRefreshCoordinator.getPendingEvent(sessionRecord.sessionId);
+    if (pendingCredentialRefresh) {
+      subscriber.write(pendingCredentialRefresh.type, {
+        session_id: sessionRecord.sessionId,
+        request_id: pendingCredentialRefresh.requestId,
+        reason: pendingCredentialRefresh.reason,
+        timeout_ms: pendingCredentialRefresh.timeoutMs,
+      });
+    }
 
     const cleanup = () => {
       run.subscribers.delete(subscriber);
@@ -390,6 +402,7 @@ export const buildChatRouter = (
 
   registerChatInteractionRoutes(chatRouter, {
     activeRuns,
+    credentialRefreshCoordinator,
     runtime,
     sessionMetadataStore,
     sessionUiStateStore,
@@ -803,6 +816,35 @@ export const buildChatRouter = (
           logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
         });
       };
+      const unsubscribeCredentialRefresh = credentialRefreshCoordinator.subscribe(
+        binding.sessionId,
+        (event) => {
+          publish(event.type, {
+            session_id: clientSessionId,
+            request_id: event.requestId,
+            ...(event.type === "credential_refresh_required"
+              ? {
+                  reason: event.reason,
+                  timeout_ms: event.timeoutMs,
+                }
+              : {}),
+            ...(event.type === "credential_refresh_failed"
+              ? { message: event.message }
+              : {}),
+          });
+        },
+      );
+      const cancelCredentialRefreshOnAbort = () => {
+        credentialRefreshCoordinator.cancelSession(
+          binding.sessionId,
+          "credential refresh cancelled because the agent run was aborted",
+        );
+      };
+      abortController.signal.addEventListener(
+        "abort",
+        cancelCredentialRefreshOnAbort,
+        { once: true },
+      );
 
       try {
         const preparedMessage = await buildPromptWithLearningContext(
@@ -925,6 +967,12 @@ export const buildChatRouter = (
           logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
         });
         sessionBridge.finalizeRequest(clientSessionId);
+        abortController.signal.removeEventListener(
+          "abort",
+          cancelCredentialRefreshOnAbort,
+        );
+        credentialRefreshCoordinator.cancelSession(binding.sessionId);
+        unsubscribeCredentialRefresh();
         activeRun.status = abortController.signal.aborted
           ? activeRun.status === "aborted"
             ? "aborted"

src/routes/chatInteractionRoutes.ts

@@ -2,8 +2,13 @@ import { type Router } from "express";
 import { z } from "zod";
 
 import { getAgentAuthContext } from "../auth/agentAuth.js";
+import { type CredentialRefreshCoordinator } from "../auth/credentialRefresh.js";
 import { logger } from "../logger.js";
 import { type OpencodeRuntimeAdapter } from "../runtime/opencode.js";
+import {
+  getRuntimeSessionContext,
+  setRuntimeSessionContext,
+} from "../runtime/sessionContext.js";
 import { type SessionMetadataStore } from "../sessions/metadataStore.js";
 import { type SessionUiStateStore } from "../sessions/uiStateStore.js";
 import { toActorKey, toProjectKey } from "../utils/fileStore.js";
@@ -26,8 +31,13 @@ const questionReplyPayloadSchema = z.object({
   answers: z.array(z.array(z.string().max(2000))).default([]),
 });
 
+const credentialRefreshPayloadSchema = z.object({
+  request_id: z.string().min(1).max(128),
+});
+
 type RegisterInteractionRoutesOptions = {
   activeRuns: Map<string, ActiveRun>;
+  credentialRefreshCoordinator: CredentialRefreshCoordinator;
   runtime: OpencodeRuntimeAdapter;
   sessionMetadataStore: SessionMetadataStore;
   sessionUiStateStore: SessionUiStateStore;
@@ -41,11 +51,73 @@ export const registerChatInteractionRoutes = (
   chatRouter: Router,
   {
     activeRuns,
+    credentialRefreshCoordinator,
     runtime,
     sessionMetadataStore,
     sessionUiStateStore,
   }: RegisterInteractionRoutesOptions,
 ) => {
+  chatRouter.post("/sessions/:session_id/credential-refreshes", async (req, res) => {
+    const parsed = credentialRefreshPayloadSchema.safeParse(req.body);
+    if (!parsed.success) {
+      res.status(400).json({
+        message: "invalid request payload",
+        detail: parsed.error.flatten(),
+      });
+      return;
+    }
+
+    const authContext = getAgentAuthContext(req);
+    const actorKey = toActorKey(authContext.userId);
+    const projectKey = toProjectKey(authContext.projectId);
+    const sessionRecord = await sessionMetadataStore.get(
+      {
+        actorKey,
+        projectId: authContext.projectId,
+        projectKey,
+        userId: authContext.userId,
+      },
+      req.params.session_id,
+    );
+    if (!sessionRecord) {
+      res.status(404).json({ message: "session not found" });
+      return;
+    }
+
+    const current = getRuntimeSessionContext(sessionRecord.sessionId);
+    if (!current || current.actorKey !== actorKey || current.projectKey !== projectKey) {
+      res.status(409).json({ message: "runtime session context unavailable" });
+      return;
+    }
+    if (
+      credentialRefreshCoordinator.getPendingRequestId(sessionRecord.sessionId) !==
+      parsed.data.request_id
+    ) {
+      res.status(409).json({ message: "credential refresh request is no longer pending" });
+      return;
+    }
+    const refreshedContext = {
+      ...current,
+      accessToken: authContext.accessToken,
+      authExpired: undefined,
+      network: authContext.network,
+      projectId: authContext.projectId,
+      tokenExpiresAt: authContext.tokenExpiresAt,
+      traceId: req.header("x-trace-id")?.trim() || current.traceId,
+    };
+    setRuntimeSessionContext(refreshedContext);
+    credentialRefreshCoordinator.resolve(
+      sessionRecord.sessionId,
+      parsed.data.request_id,
+      refreshedContext,
+    );
+    res.status(202).json({
+      session_id: sessionRecord.sessionId,
+      request_id: parsed.data.request_id,
+      status: "accepted",
+    });
+  });
+
   chatRouter.post("/sessions/:session_id/permission-responses", async (req, res) => {
     const parsed = permissionReplyPayloadSchema.safeParse(req.body);
     if (!parsed.success) {