import { calculateRetryDelay, sleep } from "./utils.js";
import { createMiddleware } from "../middleware.js";
import { RetrySchema } from "./constants.js";
import { InvalidRetryConfigError } from "./error.js";
import { AIMessage } from "@langchain/core/messages";
import { z } from "zod/v3";
//#region src/agents/middleware/modelRetry.ts
/**
* Model retry middleware for agents.
*/
/**
* Configuration options for the Model Retry Middleware.
*/
const ModelRetryMiddlewareOptionsSchema = z.object({ 
/**
* Behavior when all retries are exhausted. Options:
* - `"continue"` (default): Return an AIMessage with error details, allowing
*   the agent to potentially handle the failure gracefully.
* - `"error"`: Re-raise the exception, stopping agent execution.
* - Custom function: Function that takes the exception and returns a string
*   for the AIMessage content, allowing custom error formatting.
*/
onFailure: z.union([
	z.literal("error"),
	z.literal("continue"),
	z.function().args(z.instanceof(Error)).returns(z.string())
]).default("continue") }).merge(RetrySchema);
/**
* Middleware that automatically retries failed model calls with configurable backoff.
*
* Supports retrying on specific exceptions and exponential backoff.
*
* @example Basic usage with default settings (2 retries, exponential backoff)
* ```ts
* import { createAgent, modelRetryMiddleware } from "langchain";
*
* const agent = createAgent({
*   model: "openai:gpt-4o",
*   tools: [searchTool],
*   middleware: [modelRetryMiddleware()],
* });
* ```
*
* @example Retry specific exceptions only
* ```ts
* import { modelRetryMiddleware } from "langchain";
*
* const retry = modelRetryMiddleware({
*   maxRetries: 4,
*   retryOn: [TimeoutError, NetworkError],
*   backoffFactor: 1.5,
* });
* ```
*
* @example Custom exception filtering
* ```ts
* function shouldRetry(error: Error): boolean {
*   // Only retry on rate limit errors
*   if (error.name === "RateLimitError") {
*     return true;
*   }
*   // Or check for specific HTTP status codes
*   if (error.name === "HTTPError" && "statusCode" in error) {
*     const statusCode = (error as any).statusCode;
*     return statusCode === 429 || statusCode === 503;
*   }
*   return false;
* }
*
* const retry = modelRetryMiddleware({
*   maxRetries: 3,
*   retryOn: shouldRetry,
* });
* ```
*
* @example Return error message instead of raising
* ```ts
* const retry = modelRetryMiddleware({
*   maxRetries: 4,
*   onFailure: "continue", // Return AIMessage with error instead of throwing
* });
* ```
*
* @example Custom error message formatting
* ```ts
* const formatError = (error: Error) =>
*   `Model call failed: ${error.message}. Please try again later.`;
*
* const retry = modelRetryMiddleware({
*   maxRetries: 4,
*   onFailure: formatError,
* });
* ```
*
* @example Constant backoff (no exponential growth)
* ```ts
* const retry = modelRetryMiddleware({
*   maxRetries: 5,
*   backoffFactor: 0.0, // No exponential growth
*   initialDelayMs: 2000, // Always wait 2 seconds
* });
* ```
*
* @example Raise exception on failure
* ```ts
* const retry = modelRetryMiddleware({
*   maxRetries: 2,
*   onFailure: "error", // Re-raise exception instead of returning message
* });
* ```
*
* @param config - Configuration options for the retry middleware
* @returns A middleware instance that handles model failures with retries
*/
function modelRetryMiddleware(config = {}) {
	const { success, error, data } = ModelRetryMiddlewareOptionsSchema.safeParse(config);
	if (!success) throw new InvalidRetryConfigError(error);
	const { maxRetries, retryOn, onFailure, backoffFactor, initialDelayMs, maxDelayMs, jitter } = data;
	/**
	* Check if the exception should trigger a retry.
	*/
	const shouldRetryException = (error) => {
		if (typeof retryOn === "function") return retryOn(error);
		return retryOn.some((ErrorConstructor) => error.constructor === ErrorConstructor);
	};
	const delayConfig = {
		backoffFactor,
		initialDelayMs,
		maxDelayMs,
		jitter
	};
	/**
	* Format the failure message when retries are exhausted.
	*/
	const formatFailureMessage = (error, attemptsMade) => {
		const errorType = error.constructor.name;
		return `Model call failed after ${attemptsMade} ${attemptsMade === 1 ? "attempt" : "attempts"} with ${errorType}: ${error.message}`;
	};
	/**
	* Handle failure when all retries are exhausted.
	*/
	const handleFailure = (error, attemptsMade) => {
		if (onFailure === "error") throw error;
		let content;
		if (typeof onFailure === "function") content = onFailure(error);
		else content = formatFailureMessage(error, attemptsMade);
		return new AIMessage({ content });
	};
	return createMiddleware({
		name: "modelRetryMiddleware",
		contextSchema: ModelRetryMiddlewareOptionsSchema,
		wrapModelCall: async (request, handler) => {
			for (let attempt = 0; attempt <= maxRetries; attempt++) try {
				return await handler(request);
			} catch (error) {
				const attemptsMade = attempt + 1;
				const err = error && typeof error === "object" && "message" in error ? error : new Error(String(error));
				if (!shouldRetryException(err)) return handleFailure(err, attemptsMade);
				if (attempt < maxRetries) {
					const delay = calculateRetryDelay(delayConfig, attempt);
					if (delay > 0) await sleep(delay);
				} else return handleFailure(err, attemptsMade);
			}
			throw new Error("Unexpected: retry loop completed without returning");
		}
	});
}
//#endregion
export { modelRetryMiddleware };

//# sourceMappingURL=modelRetry.js.map