feat(agent): improve investigation contracts and runtime readiness

src/product/mcp-auth-discovery.ts

@@ -1,10 +1,11 @@
 import { discoverAuthorizationServerMetadata, discoverOAuthProtectedResourceMetadata, extractWWWAuthenticateParams } from "@modelcontextprotocol/client";
 import { secureMcpUrl, validateDevHttpIssuer, type McpUrlPolicy } from "./mcp-url-policy.js";
+import { isTransientMcpFailure, transientMcpHttpStatus } from "./mcp-query-retry.js";
 
 export type McpHttpFetch = (input: Parameters<typeof fetch>[0], init?: Parameters<typeof fetch>[1]) => ReturnType<typeof fetch>;
 type Settings = McpUrlPolicy & { resource: string; issuer: string };
 type Endpoint = { tokenEndpoint: string; authMethod: "client_secret_basic" | "client_secret_post" };
-const unavailable = (message: string) => Object.assign(new Error(message), { status: 503 });
+const unavailable = (message: string, retryable = false) => Object.assign(new Error(message), { status: 503, retryable });
 
 function oauthChallenge(response: Response) {
   const header = response.headers.get("WWW-Authenticate") ?? "";
@@ -61,6 +62,10 @@ export class McpAuthDiscovery {
       const headers = new Headers(init?.headers);
       headers.delete("Authorization"); headers.delete("Cookie");
       const response = await this.http(url, { ...init, headers, credentials: "omit", redirect: "error", signal: deadline });
+      if (transientMcpHttpStatus(response.status)) {
+        await response.body?.cancel();
+        throw unavailable("MCP OAuth metadata temporarily unavailable", true);
+      }
       // SDK parsers receive bounded public documents; no token or client secret is attached.
       const reader = response.body?.getReader();
       const chunks: Uint8Array[] = [];
@@ -107,8 +112,9 @@ export class McpAuthDiscovery {
       const authMethod = methods.includes("client_secret_post") ? "client_secret_post" : methods.includes("client_secret_basic") ? "client_secret_basic" : undefined;
       if (!authMethod) throw unavailable("MCP OAuth server does not support the configured client secret");
       return { tokenEndpoint: tokenEndpoint.href, authMethod };
-    } catch {
-      throw unavailable("MCP OAuth discovery failed or is incompatible with the configured resource and issuer");
+    } catch (cause) {
+      const transient = isTransientMcpFailure(cause);
+      throw unavailable(transient ? "MCP OAuth discovery temporarily unavailable" : "MCP OAuth discovery failed or is incompatible with the configured resource and issuer", transient);
     }
   }
 }

src/product/mcp-query-client.ts

@@ -7,12 +7,13 @@ import type { RuntimeSessionContext } from "../runtime/sessionContext.js";
 import { PRODUCT_MCP_TOOLS } from "../runtime/product-mcp-tools.js";
 import { McpAuthDiscovery, type McpHttpFetch } from "./mcp-auth-discovery.js";
 import { secureMcpUrl, validateDevHttpIssuer, type McpUrlPolicy } from "./mcp-url-policy.js";
+import { isMcpAuthenticationFailure, isTransientMcpFailure, mayRetryMcpRead, mcpToolFailure, transientMcpHttpStatus, waitForMcpRetry } from "./mcp-query-retry.js";
 
 export type { McpHttpFetch } from "./mcp-auth-discovery.js";
 export type QueryOwner = { context: RuntimeSessionContext; runtimeSessionId: string; taskId?: string };
 export type McpQuerySettings = McpUrlPolicy & { issuer: string; resource: string; audience: string; clientId: string; clientSecret: string };
 type Token = { token: string; expires: number; source: string };
-const error = (status: number, message: string) => Object.assign(new Error(message), { status });
+const error = (status: number, message: string, retryable = false) => Object.assign(new Error(message), { status, retryable });
 
 export function assertExchangedIdentity(source: string, exchanged: string, settings: McpQuerySettings) {
   const payload = (token: string): Record<string, unknown> => {
@@ -78,7 +79,7 @@ export class McpQueryClient {
     return response;
   }
   private async parseToken(response: Response, source: string): Promise<Token> {
-    if (!response.ok) { await response.body?.cancel(); throw error(503, "MCP token exchange unavailable"); }
+    if (!response.ok) { await response.body?.cancel(); throw error(503, "MCP token exchange unavailable", transientMcpHttpStatus(response.status)); }
     const data = await response.json() as Record<string, unknown>;
     if (typeof data.access_token !== "string" || !data.access_token || data.access_token.length > 16384 ||
       typeof data.expires_in !== "number" || data.expires_in < 30 || typeof data.token_type !== "string" || data.token_type.toLowerCase() !== "bearer") throw error(502, "MCP token response invalid");
@@ -165,18 +166,33 @@ export class McpQueryClient {
     const group = this.active.get(owner.runtimeSessionId) ?? new Set<AbortController>();
     group.add(controller); this.active.set(owner.runtimeSessionId, group);
     const combined = AbortSignal.any([controller.signal, AbortSignal.timeout(budget.requestMs), ...(signal ? [signal] : [])]);
+    const retryRead = mayRetryMcpRead(definition);
+    let refreshed = false;
     try {
-      for (let attempt = 0; attempt < 2; attempt++) {
-        try { return await boundedQuery(() => this.queryOnce(sessionId, owner, name, args, callId, combined), combined); }
-        catch (cause) {
-          const status = cause && typeof cause === "object" && "status" in cause ? cause.status : undefined;
-          const code = cause && typeof cause === "object" && "code" in cause ? cause.code : undefined;
-          if (attempt || !(status === 401 || code === "UNAUTHORIZED" || cause instanceof UnauthorizedError)) throw cause;
+      for (let attempt = 0; ; attempt++) {
+        let result: Awaited<ReturnType<typeof this.queryOnce>> | undefined;
+        let failure: unknown;
+        let returned = false;
+        try {
+          result = await boundedQuery(() => this.queryOnce(sessionId, owner, name, args, callId, combined), combined);
+          returned = true;
+          failure = mcpToolFailure(result);
+        } catch (cause) { failure = cause; }
+        combined.throwIfAborted();
+        if (returned && !failure) return result!;
+        const renew = !refreshed && (isMcpAuthenticationFailure(failure) || failure instanceof UnauthorizedError);
+        if (attempt >= 2 || !renew && !(retryRead && isTransientMcpFailure(failure))) {
+          if (returned) return result!;
+          throw failure;
+        }
+        if (renew) {
+          refreshed = true;
           this.tokens.delete(owner.runtimeSessionId);
           this.authDiscovery.invalidate();
+        } else {
+          await waitForMcpRetry(attempt, combined);
         }
       }
-      throw error(401, "MCP authentication required");
     } finally {
       group.delete(controller);
       if (!group.size) this.active.delete(owner.runtimeSessionId);
@@ -201,7 +217,7 @@ export class McpQueryClient {
     const binding = await this.http(bindUrl, { method: "POST", headers, redirect: "error",
       signal: signal ? AbortSignal.any([signal, AbortSignal.timeout(5000)]) : AbortSignal.timeout(5000),
       body: JSON.stringify(command) });
-    if (!binding.ok) { await binding.body?.cancel(); if (binding.status === 401) this.tokens.delete(owner.runtimeSessionId); throw error(binding.status, "MCP query context rejected"); }
+    if (!binding.ok) { await binding.body?.cancel(); if (binding.status === 401) this.tokens.delete(owner.runtimeSessionId); throw error(binding.status, "MCP query context rejected", transientMcpHttpStatus(binding.status)); }
     const body = await binding.json() as { context_handle?: unknown };
     if (typeof body.context_handle !== "string" || !/^[A-Za-z0-9_-]{32,128}$/.test(body.context_handle)) throw error(502, "MCP query context response invalid");
     const handle = body.context_handle;

src/product/mcp-query-retry.ts

@@ -0,0 +1,43 @@
+import { setTimeout as delay } from "node:timers/promises";
+import type { Tool } from "@modelcontextprotocol/client";
+import { PRODUCT_MCP_READ_TOOLS } from "../runtime/product-mcp-tools.js";
+
+const record = (value: unknown): value is Record<string, unknown> => !!value && typeof value === "object" && !Array.isArray(value);
+const transientCodes = new Set(["SOURCE_UNAVAILABLE", "QUERY_TIMEOUT", "RATE_LIMIT", "RATE_LIMITED", "REQUEST_TIMEOUT"]);
+const networkCodes = new Set(["ECONNRESET", "ECONNREFUSED", "ETIMEDOUT", "EAI_AGAIN", "ENOTFOUND", "EPIPE", "UND_ERR_SOCKET", "UND_ERR_CONNECT_TIMEOUT"]);
+
+export const transientMcpHttpStatus = (status: unknown) => status === 429 || status === 502 || status === 503 || status === 504;
+
+export function mayRetryMcpRead(tool: Tool): boolean {
+  // The calculation waiter already owns four bounded status attempts. Do not
+  // multiply its polling budget, or trust readOnlyHint as an independent grant.
+  return tool.name !== "list_calculation_jobs" && (PRODUCT_MCP_READ_TOOLS as readonly string[]).includes(tool.name) &&
+    tool.annotations?.readOnlyHint === true && tool.annotations.destructiveHint !== true;
+}
+
+export function mcpToolFailure(result: unknown): Record<string, unknown> | undefined {
+  if (!record(result) || result.isError !== true || !record(result.structuredContent) || !record(result.structuredContent.error)) return;
+  return result.structuredContent.error;
+}
+
+export function isMcpAuthenticationFailure(value: unknown): boolean {
+  if (!record(value)) return false;
+  const status = value.status ?? value.statusCode;
+  return typeof status === "number" ? status === 401 : value.code === "UNAUTHORIZED" || value.code === "AUTH_REQUIRED";
+}
+
+export function isTransientMcpFailure(value: unknown): boolean {
+  if (!record(value) || value.retryable === false) return false;
+  const status = value.status ?? value.statusCode;
+  if (typeof status === "number") return transientMcpHttpStatus(status);
+  if (typeof value.code === "string" && (transientCodes.has(value.code) || networkCodes.has(value.code))) return true;
+  if (value.name === "TimeoutError" || value.code === -32001) return true;
+  if (value.name === "TypeError" && typeof value.message === "string" && /fetch|network|connection|connect|load failed|socket/i.test(value.message)) return true;
+  return record(value.cause) && typeof value.cause.code === "string" && networkCodes.has(value.cause.code);
+}
+
+export async function waitForMcpRetry(attempt: number, signal: AbortSignal): Promise<void> {
+  signal.throwIfAborted();
+  try { await delay(attempt === 0 ? 100 : 300, undefined, { signal }); }
+  catch (cause) { signal.throwIfAborted(); throw cause; }
+}