Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,7 @@
"test:model-pool": "pnpm run test:model-pool:vitest && npx tsx test/continuous-test-suite-model-pool.ts",
"test:model-capabilities:vitest": "pnpm exec vitest run test/modelCapabilities.test.ts",
"test:agent-runtime:vitest": "pnpm exec vitest run test/agentRuntime.test.ts test/agentDelegation.test.ts test/agentPlumbing.test.ts test/toolExecutionRecorder.test.ts test/samplingParams.test.ts test/structuredRecovery.test.ts",
"test:retry-after:vitest": "pnpm exec vitest run test/retryAfter.test.ts",
"test:ci": "pnpm run test && pnpm run test:client && pnpm run test:hitl",
"// CI tier — fast, no live AI calls, safe for every commit": "",
"test:unit:vitest": "pnpm exec vitest run test/toolRouting.test.ts",
Expand All @@ -152,7 +153,7 @@
"test:system-messages:vitest": "pnpm exec vitest run test/systemMessages.test.ts",
"test:tool-routing-semantic:vitest": "pnpm exec vitest run test/toolRoutingSemantic.test.ts",
"test:tool-routing-semantic": "pnpm run test:tool-routing-semantic:vitest && npx tsx test/continuous-test-suite-tool-routing-semantic.ts",
"test:unit": "pnpm run test:envguard && pnpm run test:bugfixes && pnpm run test:file-detector-extension && pnpm run test:file-detector-magic-bytes && pnpm run test:mcp:infra && pnpm run test:mcp:bash && pnpm run test:mcp:limits && pnpm run test:mcp:spans && pnpm run test:autoresearch:redis && pnpm run test:unit:vitest && pnpm run test:tool-routing-cli:vitest && pnpm run test:tool-dedup:vitest && pnpm run test:model-pool:vitest && pnpm run test:litellm-context:vitest && pnpm run test:step-budget-guard:vitest && pnpm run test:system-messages:vitest && pnpm run test:tool-routing-semantic:vitest && pnpm run test:anthropic-tools-policy && pnpm run test:sagemaker-tools && pnpm run test:anthropic-multimodal && pnpm run test:excel-interop && pnpm run test:model-capabilities:vitest && pnpm run test:agent-runtime:vitest",
"test:unit": "pnpm run test:envguard && pnpm run test:bugfixes && pnpm run test:file-detector-extension && pnpm run test:file-detector-magic-bytes && pnpm run test:mcp:infra && pnpm run test:mcp:bash && pnpm run test:mcp:limits && pnpm run test:mcp:spans && pnpm run test:autoresearch:redis && pnpm run test:unit:vitest && pnpm run test:tool-routing-cli:vitest && pnpm run test:tool-dedup:vitest && pnpm run test:model-pool:vitest && pnpm run test:litellm-context:vitest && pnpm run test:step-budget-guard:vitest && pnpm run test:system-messages:vitest && pnpm run test:tool-routing-semantic:vitest && pnpm run test:anthropic-tools-policy && pnpm run test:sagemaker-tools && pnpm run test:anthropic-multimodal && pnpm run test:excel-interop && pnpm run test:model-capabilities:vitest && pnpm run test:agent-runtime:vitest && pnpm run test:retry-after:vitest",
"// CI tier — live providers, runs only when API keys are present (test:credentials and test:dynamic make real provider calls when keys are set, so they live here, not in test:unit)": "",
"test:live": "pnpm run test:providers && pnpm run test:mcp:http && pnpm run test:mcp:sdk && pnpm run test:mcp:cli && pnpm run test:observability && pnpm run test:context && pnpm run test:memory && pnpm run test:tool-reliability && pnpm run test:evaluation && pnpm run test:autoresearch && pnpm run test:credentials && pnpm run test:dynamic",
"// CI tier — product output (image/video/TTS/PPT) — costs $$ per run": "",
Expand Down
102 changes: 58 additions & 44 deletions src/lib/adapters/replicate/predictionLifecycle.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@ import { ErrorCategory, ErrorSeverity } from "../../constants/enums.js";
import { logger } from "../../utils/logger.js";
import { NeuroLinkError, ERROR_CODES } from "../../utils/errorHandling.js";
import { sanitizeForLog } from "../../utils/logSanitize.js";
import { withProviderRetry } from "../../utils/providerRetry.js";
import { parseRetryAfterMs } from "../../utils/retryAfter.js";
import { safeDownload } from "../../utils/safeFetch.js";
import { MAX_VIDEO_BYTES } from "../../utils/sizeGuard.js";
import type {
Expand All @@ -43,6 +45,7 @@ const DEFAULT_TOTAL_TIMEOUT_MS = 5 * 60_000;
export async function createPrediction(
auth: ReplicateAuth,
input: ReplicateCreatePredictionInput,
sleep?: (delayMs: number) => Promise<void>,
): Promise<ReplicatePrediction> {
const baseUrl = auth.baseUrl ?? "https://api.replicate.com";
const [modelPath, version] = input.model.split(":", 2);
Expand All @@ -62,52 +65,62 @@ export async function createPrediction(
body.webhook_events_filter = input.webhookEventsFilter;
}

const controller = new AbortController();
const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
return withProviderRetry(
async () => {
const controller = new AbortController();
const timeoutId = setTimeout(
() => controller.abort(),
REQUEST_TIMEOUT_MS,
);

let response: Response;
try {
response = await fetch(endpoint, {
method: "POST",
headers: {
Authorization: `Token ${auth.apiToken}`,
"Content-Type": "application/json",
Prefer: "wait=60",
},
body: JSON.stringify(body),
signal: controller.signal,
});
} catch (err: unknown) {
if (err instanceof NeuroLinkError) {
throw err;
}
if (err instanceof Error && err.name === "AbortError") {
throw new NeuroLinkError({
code: ERROR_CODES.OPERATION_ABORTED,
message: `Replicate predictions submit timed out after ${REQUEST_TIMEOUT_MS / 1000}s`,
category: ErrorCategory.TIMEOUT,
severity: ErrorSeverity.HIGH,
retriable: true,
originalError: err,
});
}
throw err;
} finally {
clearTimeout(timeoutId);
}
try {
const response = await fetch(endpoint, {
method: "POST",
headers: {
Authorization: `Token ${auth.apiToken}`,
"Content-Type": "application/json",
Prefer: "wait=60",
},
body: JSON.stringify(body),
signal: controller.signal,
});

if (!response.ok) {
const raw = await response.text();
throw new NeuroLinkError({
code: ERROR_CODES.PROVIDER_NOT_AVAILABLE,
message: `Replicate predictions submit failed: ${response.status} — ${sanitizeForLog(raw, 500)}`,
category: ErrorCategory.NETWORK,
severity: ErrorSeverity.HIGH,
retriable: response.status >= 500,
});
}
if (!response.ok) {
const raw = await response.text();
throw new NeuroLinkError({
code: ERROR_CODES.PROVIDER_NOT_AVAILABLE,
message: `Replicate predictions submit failed: ${response.status} — ${sanitizeForLog(raw, 500)}`,
category: ErrorCategory.NETWORK,
severity: ErrorSeverity.HIGH,
retriable: response.status >= 500 || response.status === 429,
retryAfterMs: parseRetryAfterMs(response.headers),
});
}

return (await response.json()) as ReplicatePrediction;
return (await response.json()) as ReplicatePrediction;
} catch (error: unknown) {
if (error instanceof NeuroLinkError) {
throw error;
}
if (error instanceof Error && error.name === "AbortError") {
throw new NeuroLinkError({
code: ERROR_CODES.OPERATION_ABORTED,
message: `Replicate predictions submit timed out after ${REQUEST_TIMEOUT_MS / 1000}s`,
category: ErrorCategory.TIMEOUT,
severity: ErrorSeverity.HIGH,
retriable: true,
originalError: error,
});
}
throw error;
} finally {
clearTimeout(timeoutId);
}
},
undefined,
"Replicate prediction submission",
sleep,
);
}

/**
Expand Down Expand Up @@ -188,7 +201,8 @@ export async function pollPrediction(
message: `Replicate poll failed: ${response.status} — ${sanitizeForLog(raw, 500)}`,
category: ErrorCategory.NETWORK,
severity: ErrorSeverity.HIGH,
retriable: response.status >= 500,
retriable: response.status >= 500 || response.status === 429,
retryAfterMs: parseRetryAfterMs(response.headers),
});
}

Expand Down
36 changes: 15 additions & 21 deletions src/lib/mcp/httpRateLimiter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
*/

import { mcpLogger } from "../utils/logger.js";
import { parseRetryAfterMs } from "../utils/retryAfter.js";
import type {
TokenBucketRateLimitConfig,
RateLimiterStats,
Expand Down Expand Up @@ -229,61 +230,54 @@ export class HTTPRateLimiter {
* @returns Wait time in milliseconds, or 0 if no rate limit headers found
*/
handleRateLimitResponse(headers: Headers): number {
// Check for Retry-After header (standard HTTP 429 response)
const parsedWaitTimeMs = parseRetryAfterMs(headers);

// Keep the existing source-specific logs while delegating delay parsing.
const retryAfter = headers.get("Retry-After");

if (retryAfter) {
// Retry-After can be either a number of seconds or an HTTP-date
if (retryAfter && parsedWaitTimeMs !== undefined) {
const seconds = parseInt(retryAfter, 10);

if (!isNaN(seconds)) {
// It's a number of seconds
const waitTimeMs = seconds * 1000;
mcpLogger.info(
`[HTTPRateLimiter] Server requested retry after ${seconds} seconds`,
);
return waitTimeMs;
return parsedWaitTimeMs;
} else {
// Try to parse as HTTP-date
const retryDate = new Date(retryAfter);
if (!isNaN(retryDate.getTime())) {
const waitTimeMs = Math.max(0, retryDate.getTime() - Date.now());
mcpLogger.info(
`[HTTPRateLimiter] Server requested retry at ${retryDate.toISOString()} (${waitTimeMs}ms)`,
`[HTTPRateLimiter] Server requested retry at ${retryDate.toISOString()} (${parsedWaitTimeMs}ms)`,
);
return waitTimeMs;
return parsedWaitTimeMs;
}
}
}

// Check for X-RateLimit-Reset header (common non-standard header)
const rateLimitReset = headers.get("X-RateLimit-Reset");
if (rateLimitReset) {
if (rateLimitReset && parsedWaitTimeMs !== undefined) {
const resetTimestamp = parseInt(rateLimitReset, 10);
if (!isNaN(resetTimestamp)) {
// Could be Unix timestamp (seconds) or milliseconds
const resetTime =
resetTimestamp > 1e12 ? resetTimestamp : resetTimestamp * 1000;
const waitTimeMs = Math.max(0, resetTime - Date.now());
mcpLogger.info(
`[HTTPRateLimiter] Rate limit resets at ${new Date(resetTime).toISOString()} (${waitTimeMs}ms)`,
`[HTTPRateLimiter] Rate limit resets at ${new Date(resetTime).toISOString()} (${parsedWaitTimeMs}ms)`,
);
return waitTimeMs;
return parsedWaitTimeMs;
}
}

// Check for X-RateLimit-Remaining header
const remaining = headers.get("X-RateLimit-Remaining");
if (remaining === "0") {
// No remaining requests, use default backoff
const defaultBackoffMs = 1000;
if (remaining === "0" && parsedWaitTimeMs !== undefined) {
mcpLogger.info(
`[HTTPRateLimiter] Rate limit exhausted, using default backoff: ${defaultBackoffMs}ms`,
`[HTTPRateLimiter] Rate limit exhausted, using default backoff: ${parsedWaitTimeMs}ms`,
);
return defaultBackoffMs;
return parsedWaitTimeMs;
}

return 0;
return parsedWaitTimeMs ?? 0;
}

/**
Expand Down
3 changes: 3 additions & 0 deletions src/lib/utils/errorHandling.ts
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ export class NeuroLinkError extends Error {
public readonly category: ErrorCategory;
public readonly severity: ErrorSeverity;
public readonly retriable: boolean;
public readonly retryAfterMs?: number;
public readonly context: Record<string, unknown>;
public readonly timestamp: Date;
public readonly toolName?: string;
Expand All @@ -114,6 +115,7 @@ export class NeuroLinkError extends Error {
category: ErrorCategory;
severity: ErrorSeverity;
retriable: boolean;
retryAfterMs?: number;
context?: Record<string, unknown>;
originalError?: Error;
toolName?: string;
Expand All @@ -125,6 +127,7 @@ export class NeuroLinkError extends Error {
this.category = options.category;
this.severity = options.severity;
this.retriable = options.retriable;
this.retryAfterMs = options.retryAfterMs;
this.context = options.context || {};
this.timestamp = new Date();
this.toolName = options.toolName;
Expand Down
53 changes: 45 additions & 8 deletions src/lib/utils/providerRetry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,26 @@
*/

import { type Span } from "@opentelemetry/api";
import { NeuroLinkError } from "./errorHandling.js";
import { logger } from "./logger.js";
import { APICallError } from "./generationErrors.js";
import { parseRetryAfterMs } from "./retryAfter.js";

/** Maximum number of retry attempts after the initial call (total = 1 + MAX_PROVIDER_RETRIES). */
export const MAX_PROVIDER_RETRIES = 2;

/** Base delay in ms for exponential backoff between retries. */
export const BASE_RETRY_DELAY_MS = 1000;

/** Minimum delay in ms when a retryable response provides no retry timing. */
export const NO_HINT_FLOOR_MS = 10_000;

/** Maximum server-requested retry delay honored by provider retries. */
export const MAX_RETRY_AFTER_MS = 120_000;

const sleepWithTimeout = (delayMs: number): Promise<void> =>
new Promise((resolve) => setTimeout(resolve, delayMs));

/**
* Check whether an error thrown by the AI SDK is retryable.
*
Expand All @@ -37,6 +48,10 @@ export function isRetryableProviderError(error: unknown): boolean {
return error.isRetryable;
}

if (error instanceof NeuroLinkError) {
return error.retriable;
}

// Fallback: duck-type for status codes on errors that aren't APICallError
if (error && typeof error === "object" && "statusCode" in error) {
const statusCode = (error as { statusCode: number }).statusCode;
Expand All @@ -59,6 +74,21 @@ export function getErrorStatusCode(error: unknown): number | undefined {
return undefined;
}

function getRetryAfterMs(error: unknown): number | undefined {
if (APICallError.isInstance(error) && error.responseHeaders) {
const parsedDelay = parseRetryAfterMs(error.responseHeaders);
if (parsedDelay !== undefined) {
return parsedDelay;
}
}

if (error instanceof NeuroLinkError && error.retryAfterMs !== undefined) {
return error.retryAfterMs;
}

return undefined;
}

/**
* Execute a provider call with instrumented retry logic.
*
Expand All @@ -69,15 +99,16 @@ export function getErrorStatusCode(error: unknown): number | undefined {
*/
export async function withProviderRetry<T>(
operation: () => Promise<T>,
span: Span,
span: Span | undefined,
label: string,
sleep: (delayMs: number) => Promise<void> = sleepWithTimeout,
): Promise<T> {
for (let attempt = 0; attempt <= MAX_PROVIDER_RETRIES; attempt++) {
try {
const result = await operation();

// Record how many attempts it took on the span
span.setAttribute("gen_ai.provider.total_attempts", attempt + 1);
span?.setAttribute("gen_ai.provider.total_attempts", attempt + 1);

if (attempt > 0) {
logger.info(
Expand All @@ -94,9 +125,9 @@ export async function withProviderRetry<T>(

if (!retryable || attempt === MAX_PROVIDER_RETRIES) {
// Record failure details before re-throwing
span.setAttribute("gen_ai.provider.total_attempts", attempt + 1);
span?.setAttribute("gen_ai.provider.total_attempts", attempt + 1);
if (attempt > 0) {
span.setAttribute("gen_ai.provider.retries_exhausted", true);
span?.setAttribute("gen_ai.provider.retries_exhausted", true);
}

logger.warn(
Expand All @@ -112,11 +143,17 @@ export async function withProviderRetry<T>(
throw error;
}

// Calculate exponential backoff delay
const delay = BASE_RETRY_DELAY_MS * Math.pow(2, attempt);
const retryAfterMs = getRetryAfterMs(error);
const boundedRetryAfterMs =
retryAfterMs === undefined
? undefined
: Math.min(MAX_RETRY_AFTER_MS, Math.max(0, retryAfterMs));
const delay =
boundedRetryAfterMs ??
Math.max(BASE_RETRY_DELAY_MS * Math.pow(2, attempt), NO_HINT_FLOOR_MS);

// Record retry event on the OTel span
span.addEvent("gen_ai.provider.retry", {
span?.addEvent("gen_ai.provider.retry", {
Comment on lines +146 to +156

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Reject non-finite retry metadata before clamping.

If retryAfterMs is NaN, the clamp also becomes NaN; the default timer then retries immediately instead of applying NO_HINT_FLOOR_MS. Treat non-finite values as absent hints.

Proposed fix
       const retryAfterMs = getRetryAfterMs(error);
       const boundedRetryAfterMs =
-        retryAfterMs === undefined
-          ? undefined
+        retryAfterMs !== undefined && Number.isFinite(retryAfterMs)
+          ? Math.min(MAX_RETRY_AFTER_MS, Math.max(0, retryAfterMs))
           : undefined;
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
const retryAfterMs = getRetryAfterMs(error);
const boundedRetryAfterMs =
retryAfterMs === undefined
? undefined
: Math.min(MAX_RETRY_AFTER_MS, Math.max(0, retryAfterMs));
const delay =
boundedRetryAfterMs ??
Math.max(BASE_RETRY_DELAY_MS * Math.pow(2, attempt), NO_HINT_FLOOR_MS);
// Record retry event on the OTel span
span.addEvent("gen_ai.provider.retry", {
span?.addEvent("gen_ai.provider.retry", {
const retryAfterMs = getRetryAfterMs(error);
const boundedRetryAfterMs =
retryAfterMs !== undefined && Number.isFinite(retryAfterMs)
? Math.min(MAX_RETRY_AFTER_MS, Math.max(0, retryAfterMs))
: undefined;
const delay =
boundedRetryAfterMs ??
Math.max(BASE_RETRY_DELAY_MS * Math.pow(2, attempt), NO_HINT_FLOOR_MS);
// Record retry event on the OTel span
span?.addEvent("gen_ai.provider.retry", {
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/lib/utils/providerRetry.ts` around lines 146 - 156, Update the retry
delay calculation around getRetryAfterMs and boundedRetryAfterMs to treat
non-finite retryAfterMs values, including NaN and infinities, as absent hints.
Validate finiteness before clamping so the nullish fallback applies
NO_HINT_FLOOR_MS (or the exponential delay) rather than propagating NaN or an
invalid bounded value.

"retry.attempt": attempt + 1,
"retry.delay_ms": delay,
...(statusCode !== undefined && { "retry.status_code": statusCode }),
Expand All @@ -134,7 +171,7 @@ export async function withProviderRetry<T>(
},
);

await new Promise((r) => setTimeout(r, delay));
await sleep(delay);
}
}

Expand Down
Loading
Loading