diff --git a/.secrets.baseline b/.secrets.baseline index 07483c87f47..76b0c59820f 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -3,7 +3,7 @@ "files": "package-lock.json|Cargo.lock|^.secrets.baseline$|scripts/sign_image.sh|scripts/zap|sonar-project.properties|^/Users/brian/dev/github.ibm.com/contextforge-org/sps-pipeline-config/.secrets.baseline$|^./.secrets.baseline$", "lines": null }, - "generated_at": "2026-03-27T22:09:20Z", + "generated_at": "2026-03-28T14:11:28Z", "plugins_used": [ { "name": "AWSKeyDetector" @@ -10080,7 +10080,7 @@ "hashed_secret": "d63b39580934e062f89aae63426d2f2c77c3e258", "is_secret": false, "is_verified": false, - "line_number": 501, + "line_number": 503, "type": "Base64 High Entropy String", "verified_result": null }, @@ -10088,7 +10088,7 @@ "hashed_secret": "586a55a9b8b97f0cd88e24ce8279ebc955949688", "is_secret": false, "is_verified": false, - "line_number": 502, + "line_number": 504, "type": "Secret Keyword", "verified_result": null }, @@ -10096,7 +10096,7 @@ "hashed_secret": "00cafd126182e8a9e7c01bb2f0dfd00496be724f", "is_secret": false, "is_verified": false, - "line_number": 518, + "line_number": 520, "type": "Secret Keyword", "verified_result": null }, @@ -10104,7 +10104,7 @@ "hashed_secret": "7b1552c7c7ffb8bd70b5666e5997c8e017630aab", "is_secret": false, "is_verified": false, - "line_number": 1936, + "line_number": 1934, "type": "Base64 High Entropy String", "verified_result": null }, @@ -10112,7 +10112,7 @@ "hashed_secret": "9fb7fe1217aed442b04c0f5e43b5d5a7d3287097", "is_secret": false, "is_verified": false, - "line_number": 2704, + "line_number": 2702, "type": "Secret Keyword", "verified_result": null }, @@ -10120,7 +10120,7 @@ "hashed_secret": "72cb70dbbafe97e5ea13ad88acd65d08389439b0", "is_secret": false, "is_verified": false, - "line_number": 3332, + "line_number": 3330, "type": "Secret Keyword", "verified_result": null }, @@ -10128,7 +10128,7 @@ "hashed_secret": "ee977806d7286510da8b9a7492ba58e2484c0ecc", "is_secret": false, "is_verified": false, - "line_number": 5366, + "line_number": 5609, "type": "Secret Keyword", "verified_result": null }, @@ -10136,7 +10136,7 @@ "hashed_secret": "f2e7745f43b0ef0e2c2faf61d6c6a28be2965750", "is_secret": false, "is_verified": false, - "line_number": 5858, + "line_number": 6101, "type": "Secret Keyword", "verified_result": null }, @@ -10144,7 +10144,7 @@ "hashed_secret": "4a249743d4d2241bd2ae085b4fe654d089488295", "is_secret": false, "is_verified": false, - "line_number": 7025, + "line_number": 7268, "type": "Secret Keyword", "verified_result": null }, @@ -10152,7 +10152,7 @@ "hashed_secret": "0c8d051d3c7eada5d31b53d9936fce6bcc232ae2", "is_secret": false, "is_verified": false, - "line_number": 7163, + "line_number": 7406, "type": "Secret Keyword", "verified_result": null }, @@ -10160,7 +10160,7 @@ "hashed_secret": "f2b14f68eb995facb3a1c35287b778d5bd785511", "is_secret": false, "is_verified": false, - "line_number": 7492, + "line_number": 7735, "type": "Secret Keyword", "verified_result": null } @@ -10202,7 +10202,7 @@ "hashed_secret": "99834bc4eff3f1e1c1e4692d2476b593b501d045", "is_secret": false, "is_verified": false, - "line_number": 5481, + "line_number": 5483, "type": "Secret Keyword", "verified_result": null }, @@ -10210,7 +10210,7 @@ "hashed_secret": "f2b14f68eb995facb3a1c35287b778d5bd785511", "is_secret": false, "is_verified": false, - "line_number": 5499, + "line_number": 5501, "type": "Secret Keyword", "verified_result": null }, @@ -10218,7 +10218,7 @@ "hashed_secret": "b8cec023ad982a1355abb9e7c7700e42d7e6fac3", "is_secret": false, "is_verified": false, - "line_number": 5523, + "line_number": 5525, "type": "Secret Keyword", "verified_result": null }, @@ -10226,7 +10226,7 @@ "hashed_secret": "0614eb27c6e0d2f4ed9a1cce2a5bcbbbc17aa556", "is_secret": false, "is_verified": false, - "line_number": 5587, + "line_number": 5589, "type": "Secret Keyword", "verified_result": null }, @@ -10234,7 +10234,7 @@ "hashed_secret": "a38fea79a043f0c9a62f851659d459dc3b716ea9", "is_secret": false, "is_verified": false, - "line_number": 5598, + "line_number": 5600, "type": "Secret Keyword", "verified_result": null }, @@ -10242,7 +10242,7 @@ "hashed_secret": "a50da1fb101595ac5158f2c9394a65a84061b56c", "is_secret": false, "is_verified": false, - "line_number": 5639, + "line_number": 5641, "type": "Secret Keyword", "verified_result": null }, @@ -10250,7 +10250,7 @@ "hashed_secret": "ed3e0017cb8e4b06a59af1a441f62cbe58d2ef59", "is_secret": false, "is_verified": false, - "line_number": 5645, + "line_number": 5647, "type": "Secret Keyword", "verified_result": null } diff --git a/mcpgateway/plugins/framework/manager.py b/mcpgateway/plugins/framework/manager.py index 0191af1e1aa..636c2b8865c 100644 --- a/mcpgateway/plugins/framework/manager.py +++ b/mcpgateway/plugins/framework/manager.py @@ -170,6 +170,7 @@ async def execute( combined_metadata: dict[str, Any] = {} current_payload: PluginPayload | None = None decision_plugin_name: Optional[str] = None + max_retry_delay_ms: int = 0 for hook_ref in hook_refs: # Skip disabled plugins @@ -224,6 +225,9 @@ async def execute( combined_metadata, ) + # Propagate retry signal — take the largest delay requested by any plugin + max_retry_delay_ms = max(max_retry_delay_ms, result.retry_delay_ms) + # Apply policy-based controlled merge (per-plugin) if result.modified_payload is not None: if policy: @@ -292,7 +296,7 @@ async def execute( if hook_type == HTTP_AUTH_CHECK_PERMISSION_HOOK and decision_plugin_name: combined_metadata[DECISION_PLUGIN_METADATA_KEY] = decision_plugin_name - return (PluginResult(continue_processing=True, modified_payload=current_payload, violation=None, metadata=combined_metadata), res_local_contexts) + return (PluginResult(continue_processing=True, modified_payload=current_payload, violation=None, metadata=combined_metadata, retry_delay_ms=max_retry_delay_ms), res_local_contexts) async def execute_plugin( self, diff --git a/mcpgateway/plugins/framework/models.py b/mcpgateway/plugins/framework/models.py index 5892279989c..7e550c91126 100644 --- a/mcpgateway/plugins/framework/models.py +++ b/mcpgateway/plugins/framework/models.py @@ -1336,6 +1336,9 @@ class PluginResult(BaseModel, Generic[T]): violation (Optional[PluginViolation]): violation object. metadata (Optional[dict[str, Any]]): additional metadata. http_headers (Optional[dict[str, str]]): HTTP headers to include in successful responses. + retry_delay_ms (int): Milliseconds the gateway should wait before retrying the tool call. + 0 (default) means no retry. Set by retry_with_backoff plugin to request + a delayed re-execution of the tool. Examples: >>> result = PluginResult() @@ -1358,6 +1361,9 @@ class PluginResult(BaseModel, Generic[T]): >>> r2 = PluginResult(continue_processing=False) >>> r2.continue_processing False + >>> r3 = PluginResult(retry_delay_ms=500) + >>> r3.retry_delay_ms + 500 """ continue_processing: bool = True @@ -1365,6 +1371,7 @@ class PluginResult(BaseModel, Generic[T]): violation: Optional[PluginViolation] = None metadata: Optional[dict[str, Any]] = Field(default_factory=dict) http_headers: Optional[dict[str, str]] = None + retry_delay_ms: int = 0 class GlobalContext(BaseModel): diff --git a/mcpgateway/services/tool_service.py b/mcpgateway/services/tool_service.py index de3fd08eef0..87685ef09ce 100644 --- a/mcpgateway/services/tool_service.py +++ b/mcpgateway/services/tool_service.py @@ -543,8 +543,24 @@ class ToolTimeoutError(ToolInvocationError): This subclass is used to distinguish timeout errors from other invocation errors. Timeout handlers call tool_post_invoke before raising this, so the generic exception handler should skip calling post_invoke again to avoid double-counting failures. + + Attributes: + retry_delay_ms: Delay in milliseconds requested by the retry plugin. + 0 (default) means no retry. Set by the timeout handler after + invoking the post-invoke hook so the outer catch block can honour + the signal without calling post_invoke a second time. """ + def __init__(self, message: str, retry_delay_ms: int = 0) -> None: + """Initialise with an optional retry delay from the post-invoke hook. + + Args: + message: Human-readable error description. + retry_delay_ms: Milliseconds the gateway should wait before retrying. + """ + super().__init__(message) + self.retry_delay_ms = retry_delay_ms + class ToolService(BaseService): """Service for managing and invoking tools. @@ -3216,6 +3232,121 @@ def _load_invocable_tools(self, db: Session, name: str, server_id: Optional[str] query = query.join(server_tool_association, DbTool.id == server_tool_association.c.tool_id).where(server_tool_association.c.server_id == server_id) return db.execute(query).scalars().all() + # ------------------------------------------------------------------ + # Retry helpers (used by invoke_tool) + # ------------------------------------------------------------------ + + async def _run_timeout_post_invoke( + self, + name: str, + effective_timeout: float, + global_context: Any, + context_table: Any, + ) -> None: + """Invoke post-invoke plugins after a timeout and raise with retry signal if requested. + + Called from each transport-specific timeout handler so the retry plugin + can record the failure and (optionally) request a retry. If the plugin + sets ``retry_delay_ms > 0``, a ``ToolTimeoutError`` carrying the delay + is raised immediately; otherwise control returns to the caller which + raises a plain ``ToolTimeoutError``. + + Args: + name: Tool name. + effective_timeout: Timeout duration in seconds. + global_context: Plugin global context for cross-hook state. + context_table: Plugin local context table for per-plugin state. + + Raises: + ToolTimeoutError: When the retry plugin requests a delayed retry. + """ + if context_table: + for ctx in context_table.values(): + ctx.set_state("cb_timeout_failure", True) + + if not self._plugin_manager: + return + + if self._plugin_manager.has_hooks_for(ToolHookType.TOOL_POST_INVOKE): + timeout_error_result = ToolResult(content=[TextContent(type="text", text=f"Tool invocation timed out after {effective_timeout}s")], is_error=True) + timeout_post_result, _ = await self._plugin_manager.invoke_hook( + ToolHookType.TOOL_POST_INVOKE, + payload=ToolPostInvokePayload(name=name, result=timeout_error_result.model_dump(by_alias=True)), + global_context=global_context, + local_contexts=context_table, + violations_as_exceptions=False, + ) + if timeout_post_result and timeout_post_result.retry_delay_ms > 0: + raise ToolTimeoutError(f"Tool invocation timed out after {effective_timeout}s", retry_delay_ms=timeout_post_result.retry_delay_ms) + + async def _retry_tool_invocation( + self, + delay_ms: int, + retry_attempt: int, + name: str, + arguments: Dict[str, Any], + request_headers: Any, + app_user_email: Optional[str], + user_email: Optional[str], + token_teams: Optional[List[str]], + server_id: Optional[str], + context_table: Any, + global_context: Any, + meta_data: Optional[Dict[str, Any]], + skip_pre_invoke: bool, + path_label: str, + ) -> "ToolResult": + """Sleep for the plugin-requested delay, then recursively re-invoke the tool. + + The sleep is cancellation-aware: if the calling task is cancelled (e.g. + client disconnect) the ``CancelledError`` propagates immediately instead + of wasting time on a retry that nobody will consume. + + Args: + delay_ms: Backoff delay in milliseconds before retrying. + retry_attempt: Current zero-based retry counter. + name: Tool name to re-invoke. + arguments: Tool arguments to forward. + request_headers: Original request headers. + app_user_email: ContextForge user email for OAuth. + user_email: User email for authorization. + token_teams: Team IDs from JWT token. + server_id: Virtual server ID for scoping. + context_table: Plugin local context table. + global_context: Plugin global context. + meta_data: Optional metadata dictionary. + skip_pre_invoke: Whether to skip pre-invoke hooks. + path_label: Label for log messages (success/timeout/exception). + + Returns: + ToolResult from the retried invocation. + """ + logger.debug( + "tool_service: retry requested (%s) for tool=%s attempt=%d/%d delay_ms=%d", + path_label, + name, + retry_attempt + 1, + settings.max_tool_retries, + delay_ms, + ) + await asyncio.sleep(delay_ms / 1000) + with fresh_db_session() as retry_db: + return await self.invoke_tool( + db=retry_db, + name=name, + arguments=arguments, + request_headers=request_headers, + app_user_email=app_user_email, + user_email=user_email, + token_teams=token_teams, + server_id=server_id, + plugin_context_table=context_table, + plugin_global_context=global_context, + meta_data=meta_data, + skip_pre_invoke=skip_pre_invoke, + retry_attempt=retry_attempt + 1, + ) + async def invoke_tool( self, db: Session, @@ -3230,6 +3361,7 @@ async def invoke_tool( plugin_global_context: Optional[GlobalContext] = None, meta_data: Optional[Dict[str, Any]] = None, skip_pre_invoke: bool = False, + retry_attempt: int = 0, ) -> ToolResult: """ Invoke a registered tool and record execution metrics. @@ -3252,6 +3384,8 @@ async def invoke_tool( plugin_global_context: Optional global context from middleware for consistency across hooks. meta_data: Optional metadata dictionary for additional context (e.g., request ID). skip_pre_invoke: When True, skip TOOL_PRE_INVOKE hooks (used by trusted Rust fallback path). + retry_attempt: Zero-based retry counter; 0 = original call. Incremented by the retry + loop and compared against ``settings.max_tool_retries``. Returns: Tool invocation result. @@ -3801,19 +3935,7 @@ async def invoke_tool( ) if self._plugin_manager: - if context_table: - for ctx in context_table.values(): - ctx.set_state("cb_timeout_failure", True) - - if self._plugin_manager.has_hooks_for(ToolHookType.TOOL_POST_INVOKE): - timeout_error_result = ToolResult(content=[TextContent(type="text", text=f"Tool invocation timed out after {effective_timeout}s")], is_error=True) - await self._plugin_manager.invoke_hook( - ToolHookType.TOOL_POST_INVOKE, - payload=ToolPostInvokePayload(name=name, result=timeout_error_result.model_dump(by_alias=True)), - global_context=global_context, - local_contexts=context_table, - violations_as_exceptions=False, - ) + await self._run_timeout_post_invoke(name, effective_timeout, global_context, context_table) raise ToolTimeoutError(f"Tool invocation timed out after {effective_timeout}s") response.raise_for_status() @@ -4137,19 +4259,7 @@ async def connect_to_sse_server(server_url: str, headers: dict = headers): ) if self._plugin_manager: - if context_table: - for ctx in context_table.values(): - ctx.set_state("cb_timeout_failure", True) - - if self._plugin_manager.has_hooks_for(ToolHookType.TOOL_POST_INVOKE): - timeout_error_result = ToolResult(content=[TextContent(type="text", text=f"Tool invocation timed out after {effective_timeout}s")], is_error=True) - await self._plugin_manager.invoke_hook( - ToolHookType.TOOL_POST_INVOKE, - payload=ToolPostInvokePayload(name=name, result=timeout_error_result.model_dump(by_alias=True)), - global_context=global_context, - local_contexts=context_table, - violations_as_exceptions=False, - ) + await self._run_timeout_post_invoke(name, effective_timeout, global_context, context_table) raise ToolTimeoutError(f"Tool invocation timed out after {effective_timeout}s") except BaseException as e: @@ -4288,19 +4398,7 @@ async def connect_to_streamablehttp_server(server_url: str, headers: dict = head ) if self._plugin_manager: - if context_table: - for ctx in context_table.values(): - ctx.set_state("cb_timeout_failure", True) - - if self._plugin_manager.has_hooks_for(ToolHookType.TOOL_POST_INVOKE): - timeout_error_result = ToolResult(content=[TextContent(type="text", text=f"Tool invocation timed out after {effective_timeout}s")], is_error=True) - await self._plugin_manager.invoke_hook( - ToolHookType.TOOL_POST_INVOKE, - payload=ToolPostInvokePayload(name=name, result=timeout_error_result.model_dump(by_alias=True)), - global_context=global_context, - local_contexts=context_table, - violations_as_exceptions=False, - ) + await self._run_timeout_post_invoke(name, effective_timeout, global_context, context_table) raise ToolTimeoutError(f"Tool invocation timed out after {effective_timeout}s") except BaseException as e: @@ -4466,19 +4564,7 @@ async def connect_to_streamablehttp_server(server_url: str, headers: dict = head # Trigger circuit breaker on timeout if self._plugin_manager: - if context_table: - for ctx in context_table.values(): - ctx.set_state("cb_timeout_failure", True) - - if self._plugin_manager.has_hooks_for(ToolHookType.TOOL_POST_INVOKE): - timeout_error_result = ToolResult(content=[TextContent(type="text", text=f"Tool invocation timed out after {effective_timeout}s")], is_error=True) - await self._plugin_manager.invoke_hook( - ToolHookType.TOOL_POST_INVOKE, - payload=ToolPostInvokePayload(name=name, result=timeout_error_result.model_dump(by_alias=True)), - global_context=global_context, - local_contexts=context_table, - violations_as_exceptions=False, - ) + await self._run_timeout_post_invoke(name, effective_timeout, global_context, context_table) raise ToolTimeoutError(f"Tool invocation timed out after {effective_timeout}s") @@ -4524,17 +4610,57 @@ async def connect_to_streamablehttp_server(server_url: str, headers: dict = head except Exception: tool_result = ToolResult(content=[TextContent(type="text", text=str(modified_result))]) + # Retry: if the plugin requested a delayed retry and we haven't hit the gateway ceiling. + # retry_attempt is 0-based (0 = original call). The condition allows retry_attempt + # values 0..max_tool_retries-1, meaning up to max_tool_retries *retry* attempts on + # top of the original call (total attempts = max_tool_retries + 1). + if post_result.retry_delay_ms > 0 and retry_attempt < settings.max_tool_retries: + return await self._retry_tool_invocation( + post_result.retry_delay_ms, + retry_attempt, + name, + arguments, + request_headers, + app_user_email, + user_email, + token_teams, + server_id, + context_table, + global_context, + meta_data, + skip_pre_invoke, + "success", + ) + return tool_result except (PluginError, PluginViolationError): raise except ToolTimeoutError as e: - # ToolTimeoutError is raised by timeout handlers which already called tool_post_invoke - # Re-raise without calling post_invoke again to avoid double-counting failures - # But DO set error_message and span attributes for observability + # ToolTimeoutError is raised by timeout handlers which already called tool_post_invoke. + # Do NOT call post_invoke again — the retry_delay_ms signal is carried on the exception. error_message = str(e) if span: span.set_attribute("error", True) span.set_attribute("error.message", error_message) + + # Retry if the post-invoke hook (called by the timeout handler) requested it. + if e.retry_delay_ms > 0 and retry_attempt < settings.max_tool_retries: + return await self._retry_tool_invocation( + e.retry_delay_ms, + retry_attempt, + name, + arguments, + request_headers, + app_user_email, + user_email, + token_teams, + server_id, + context_table, + global_context, + meta_data, + skip_pre_invoke, + "timeout", + ) raise except BaseException as e: # Extract root cause from ExceptionGroup (Python 3.11+) @@ -4549,12 +4675,19 @@ async def connect_to_streamablehttp_server(server_url: str, headers: dict = head span.set_attribute("error", True) span.set_attribute("error.message", error_message) - # Notify plugins of the failure so circuit breaker can track it - # This ensures HTTP 4xx/5xx errors and MCP failures are counted + # Notify plugins of the failure so circuit breaker / retry plugin can track it. + # Capture the result so we can honour a retry_delay_ms signal from the retry plugin. + # When the exception carries an HTTP status code (e.g. httpx.HTTPStatusError), + # include it in structuredContent so the retry plugin can honour retry_on_status + # instead of blindly retrying every exception. + exc_post_result = None if self._plugin_manager and self._plugin_manager.has_hooks_for(ToolHookType.TOOL_POST_INVOKE): try: - exception_error_result = ToolResult(content=[TextContent(type="text", text=f"Tool invocation failed: {error_message}")], is_error=True) - await self._plugin_manager.invoke_hook( + exc_structured: Optional[Dict[str, Any]] = None + if isinstance(root_cause, httpx.HTTPStatusError): + exc_structured = {"status_code": root_cause.response.status_code} + exception_error_result = ToolResult(content=[TextContent(type="text", text=f"Tool invocation failed: {error_message}")], is_error=True, structured_content=exc_structured) + exc_post_result, _ = await self._plugin_manager.invoke_hook( ToolHookType.TOOL_POST_INVOKE, payload=ToolPostInvokePayload(name=name, result=exception_error_result.model_dump(by_alias=True)), global_context=global_context, @@ -4564,6 +4697,27 @@ async def connect_to_streamablehttp_server(server_url: str, headers: dict = head except Exception as plugin_exc: logger.debug("Failed to invoke post-invoke plugins on exception: %s", plugin_exc) + # Retry if the plugin requested a delayed retry and we haven't hit the ceiling. + # Same counting convention as the success path: retry_attempt is 0-based, + # so this allows up to max_tool_retries retry attempts beyond the original call. + if exc_post_result is not None and exc_post_result.retry_delay_ms > 0 and retry_attempt < settings.max_tool_retries: + return await self._retry_tool_invocation( + exc_post_result.retry_delay_ms, + retry_attempt, + name, + arguments, + request_headers, + app_user_email, + user_email, + token_teams, + server_id, + context_table, + global_context, + meta_data, + skip_pre_invoke, + "exception", + ) + raise ToolInvocationError(f"Tool invocation failed: {error_message}") finally: # Calculate duration diff --git a/plugins/config.yaml b/plugins/config.yaml index 1a69d2f28c3..d38d6cf092b 100644 --- a/plugins/config.yaml +++ b/plugins/config.yaml @@ -320,15 +320,15 @@ plugins: allowed_mime_types: ["text/plain", "text/markdown", "text/html", "application/json"] allowed_extensions: [".md", ".txt", ".html", ".json"] - # Retry policy annotations + # Retry with exponential backoff — triggers real gateway re-invocation - name: "RetryWithBackoffPlugin" kind: "plugins.retry_with_backoff.retry_with_backoff.RetryWithBackoffPlugin" - description: "Annotates retry/backoff policy in metadata" + description: "Detects transient failures and asks the gateway to re-invoke the tool after a jittered exponential backoff delay" version: "0.1.0" author: "Mihai Criveti" - hooks: ["tool_post_invoke", "resource_post_fetch"] - tags: ["reliability", "retry"] - mode: "disabled" + hooks: ["tool_post_invoke"] + tags: ["retry", "backoff", "resilience"] + mode: "permissive" priority: 170 conditions: [] config: @@ -336,6 +336,9 @@ plugins: backoff_base_ms: 200 max_backoff_ms: 5000 retry_on_status: [429, 500, 502, 503, 504] + jitter: true + check_text_content: false + tool_overrides: {} # Markdown cleaner - name: "MarkdownCleanerPlugin" diff --git a/plugins/retry_with_backoff/README.md b/plugins/retry_with_backoff/README.md index d141231bd80..1cfd9625b35 100644 --- a/plugins/retry_with_backoff/README.md +++ b/plugins/retry_with_backoff/README.md @@ -3,11 +3,11 @@ > Author: Mihai Criveti > Version: 0.1.0 -Annotates retry/backoff policy in metadata for downstream orchestration. Does not re-execute tools. +Exponential backoff retry plugin: detects transient failures and asks the gateway to re-invoke the tool after a jittered delay. The gateway owns the sleep and the retry loop (see `tool_service.py`); this plugin owns the failure detection and delay calculation. ## Hooks -- tool_post_invoke -- resource_post_fetch +- `tool_post_invoke` — active retry (detects failure, computes delay, requests re-invocation) +- `resource_post_fetch` — advisory only (attaches retry policy metadata; see Limitations) ## Config ```yaml @@ -16,16 +16,54 @@ config: backoff_base_ms: 200 max_backoff_ms: 5000 retry_on_status: [429, 500, 502, 503, 504] + jitter: true + check_text_content: false + tool_overrides: {} ``` ## Design -- Adds a retry policy descriptor to metadata for tools and resource fetches; no side effects. -- Fields include max retries and exponential backoff parameters; downstream decides how to apply. + +The plugin checks three failure signals in order: + +1. **`isError`** — set to `true` when the tool raises an exception. When the + gateway can determine the HTTP status code of the failure (e.g. from an + `httpx.HTTPStatusError`), it includes the code in `structuredContent`. + If a status code is present, `retry_on_status` is checked — non-transient + errors like 400 or 404 are **not** retried. Generic exceptions without a + status code (connection resets, timeouts) are always retried. +2. **`structuredContent.status_code`** — for tools on MCP spec 2025-03-26+; the gateway places a plain dict in `structuredContent`. +3. **Text content JSON parsing** — opt-in (`check_text_content: true`) for older MCP servers that return HTTP-style error dicts as serialised JSON in text content instead of raising exceptions. Disabled by default to avoid false-positives. + +Backoff uses full-jitter exponential delay: + +``` +delay = random(0, min(max_backoff_ms, backoff_base_ms × 2^attempt)) +``` + +A Rust extension (`retry_with_backoff_rust`) is used when available for signals 1 and 2, falling back to the pure-Python implementation otherwise. + +### State Management + +Per-invocation retry state is keyed by `(tool_name, request_id)` and cleaned up +on success or budget exhaustion. Entries orphaned by cancelled retries (e.g. +client disconnect during the backoff sleep) are automatically evicted after a +5-minute TTL. + +## Tool-Level Overrides + +Individual tools can override any config field: + +```yaml +config: + tool_overrides: + my_flaky_tool: + max_retries: 4 + backoff_base_ms: 500 +``` ## Limitations -- Purely advisory; does not perform any retry logic. -- No per-tool/resource overrides in current version. -## TODOs -- Add per-tool/resource overrides; dynamic tuning based on payload size. -- Include jitter strategy hints and orchestration examples. +- `max_retries` is clamped to the gateway-level `max_tool_retries` setting. +- `check_text_content` is off by default to avoid false-positives on tools that legitimately return status codes as informational data. +- Per-tool overrides are also clamped to the gateway ceiling. +- **Resource retry is not yet implemented.** The `resource_post_fetch` hook registers successfully and returns retry policy metadata, but does not trigger actual retries. Resource fetch failures raise exceptions before the post-fetch hook fires, so transient resource errors are not retried. Only `tool_post_invoke` performs active retry logic. diff --git a/plugins/retry_with_backoff/__init__.py b/plugins/retry_with_backoff/__init__.py index c5b63aad44a..c56a0434c13 100644 --- a/plugins/retry_with_backoff/__init__.py +++ b/plugins/retry_with_backoff/__init__.py @@ -1,9 +1,7 @@ # -*- coding: utf-8 -*- -"""Module Description. -Location: ./plugins/retry_with_backoff/__init__.py +"""Location: ./plugins/retry_with_backoff/__init__.py Copyright 2025 SPDX-License-Identifier: Apache-2.0 -Authors: Mihai Criveti -Module documentation... +Retry with Backoff Plugin package. """ diff --git a/plugins/retry_with_backoff/plugin-manifest.yaml b/plugins/retry_with_backoff/plugin-manifest.yaml index 1517d53de57..1aac272dce0 100644 --- a/plugins/retry_with_backoff/plugin-manifest.yaml +++ b/plugins/retry_with_backoff/plugin-manifest.yaml @@ -1,11 +1,20 @@ -description: "Annotate retry/backoff policy in metadata" +description: "Exponential backoff retry plugin: detects transient failures and asks the gateway to re-invoke the tool after a jittered delay." author: "Mihai Criveti" version: "0.1.0" +tags: ["retry", "backoff", "resilience"] available_hooks: - "tool_post_invoke" - "resource_post_fetch" -default_configs: +default_config: max_retries: 2 backoff_base_ms: 200 max_backoff_ms: 5000 - retry_on_status: [429, 500, 502, 503, 504] + retry_on_status: + - 429 + - 500 + - 502 + - 503 + - 504 + jitter: true + check_text_content: false + tool_overrides: {} diff --git a/plugins/retry_with_backoff/retry_with_backoff.py b/plugins/retry_with_backoff/retry_with_backoff.py index ef63ee87fa5..bc6c93d62f4 100644 --- a/plugins/retry_with_backoff/retry_with_backoff.py +++ b/plugins/retry_with_backoff/retry_with_backoff.py @@ -2,20 +2,34 @@ """Location: ./plugins/retry_with_backoff/retry_with_backoff.py Copyright 2025 SPDX-License-Identifier: Apache-2.0 -Authors: Mihai Criveti -Retry With Backoff Plugin. -Advisory plugin that annotates retry policy metadata for downstream systems. -Note: The framework cannot re-execute tools/resources; this provides guidance only. +Retry with Backoff Plugin. + +Uses the retry_delay_ms field on PluginResult to ask the gateway to +re-execute the tool after a computed delay. The gateway owns the sleep +and the retry loop (see tool_service.py); this plugin owns the failure +detection and the delay calculation. + +Hooks: tool_post_invoke """ # Future from __future__ import annotations +# Standard +import json +import logging +import math +import random +import time +from dataclasses import dataclass +from typing import Any, Dict + # Third-Party from pydantic import BaseModel, Field # First-Party +from mcpgateway.config import get_settings from mcpgateway.plugins.framework import ( Plugin, PluginConfig, @@ -26,64 +40,448 @@ ToolPostInvokeResult, ) +log = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Optional Rust accelerator +# +# If the compiled extension is installed, a RetryStateManager instance is +# created per plugin instance (in __init__), with config baked in so the +# hot check_and_update call only passes the four dynamic args: +# tool, request_id, is_error, status_code +# +# If it's absent (e.g. dev machine without the Rust toolchain, or a +# pure-Python wheel), the plugin silently falls back to the Python +# implementation below. +# +# The try/except ImportError pattern is the standard Python idiom for +# optional compiled extensions. It imposes zero cost when Rust IS present +# (the import succeeds and _RUST_AVAILABLE is set once at module load), and +# makes the plugin fully portable when it is NOT. +# --------------------------------------------------------------------------- +try: + from retry_with_backoff_rust import RetryStateManager as _RustRetryStateManager + + _RUST_AVAILABLE = True + log.debug("retry_with_backoff: Rust extension loaded") +except ImportError: + _RustRetryStateManager = None # type: ignore[assignment,misc] + _RUST_AVAILABLE = False + log.debug("retry_with_backoff: Rust extension not available, using Python fallback") + + +# --------------------------------------------------------------------------- +# Per-tool runtime state +# --------------------------------------------------------------------------- + + +@dataclass +class _ToolRetryState: + """Mutable retry state for a single tool.""" + + consecutive_failures: int = 0 + last_failure_at: float = 0.0 + + +# Module-level dict — one entry per (tool_name, request_id). +# Each independent tool invocation gets its own fresh state; retries of the +# same invocation share state because the gateway passes the same global_context +# (and therefore the same request_id) on every retry attempt. +_STATE: Dict[str, _ToolRetryState] = {} + +# Entries older than this are considered orphaned (e.g. the retry sleep was +# cancelled by a client disconnect) and are evicted on the next _get_state call. +_STATE_TTL_SECONDS: float = 300.0 + + +def _evict_stale_entries() -> None: + """Remove state entries whose last failure is older than the TTL. + + Called from _get_state on every access. The dict is typically very small + (one entry per in-flight retry chain) so the scan is negligible. + """ + cutoff = time.monotonic() - _STATE_TTL_SECONDS + stale = [k for k, v in _STATE.items() if v.last_failure_at > 0 and v.last_failure_at < cutoff] + for k in stale: + del _STATE[k] + + +def _get_state(tool: str, request_id: str) -> _ToolRetryState: + """Return the retry state entry for a given (tool, request_id) pair, creating it if absent. + + Evicts stale entries on every call to prevent unbounded growth from + cancelled retries (e.g. client disconnects during the backoff sleep). + + Args: + tool: Tool name. + request_id: Unique request identifier. + + Returns: + The mutable retry state for this (tool, request_id) pair. + """ + _evict_stale_entries() + key = f"{tool}:{request_id}" + if key not in _STATE: + _STATE[key] = _ToolRetryState() + return _STATE[key] + + +def _del_state(tool: str, request_id: str) -> None: + """Remove the retry state entry for a given (tool, request_id) pair, if it exists. + + Args: + tool: Tool name. + request_id: Unique request identifier. + """ + _STATE.pop(f"{tool}:{request_id}", None) + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +class RetryConfig(BaseModel): + """Per-plugin configuration, read from config.yaml under the plugin's config: key.""" + + max_retries: int = Field(default=2, ge=0, description="Max consecutive retries before giving up") + backoff_base_ms: int = Field(default=200, ge=1, description="Initial backoff in milliseconds") + max_backoff_ms: int = Field(default=5000, ge=1, description="Ceiling for computed backoff in milliseconds") + retry_on_status: list[int] = Field( + default_factory=lambda: [429, 500, 502, 503, 504], + description="HTTP-style status codes in tool result that count as transient failures", + ) + jitter: bool = Field(default=True, description="Apply full-jitter to avoid thundering-herd") + check_text_content: bool = Field( + default=False, + description=( + "Parse text content as JSON and check for status_code when structuredContent is absent. " + "Enable only for tools on older MCP servers (pre-2025 spec) that return HTTP-style error " + "dicts in text content instead of raising exceptions. OFF by default because it can " + "false-positive on tools that legitimately return status codes as informational data." + ), + ) + tool_overrides: Dict[str, Dict[str, Any]] = Field( + default_factory=dict, + description="Per-tool config overrides; key = tool name, value = subset of above fields", + ) + + +def _cfg_for(cfg: RetryConfig, tool: str) -> RetryConfig: + """Return config merged with any per-tool overrides. + + Args: + cfg: Base plugin configuration. + tool: Tool name to look up overrides for. + + Returns: + Merged config if overrides exist, otherwise the original config. + """ + overrides = cfg.tool_overrides.get(tool) + if not overrides: + return cfg + merged = cfg.model_dump() + merged.update(overrides) + merged.pop("tool_overrides", None) + return RetryConfig(**merged) + + +# --------------------------------------------------------------------------- +# Backoff calculation +# --------------------------------------------------------------------------- + + +def _compute_delay_ms(attempt: int, cfg: RetryConfig) -> int: + """Return jittered exponential backoff delay in milliseconds. -class RetryPolicyConfig(BaseModel): - """Configuration for retry policy. + Uses full-jitter: random value between 0 and min(cap, base * 2^attempt). + This prevents thundering-herd when many tools fail at the same time. - Attributes: - max_retries: Maximum number of retry attempts. - backoff_base_ms: Base backoff duration in milliseconds. - max_backoff_ms: Maximum backoff duration in milliseconds. - retry_on_status: HTTP status codes that trigger retries. + Args: + attempt: Zero-based retry attempt index. + cfg: Retry configuration with backoff parameters. + + Returns: + Delay in milliseconds. """ + cap = cfg.max_backoff_ms + base = cfg.backoff_base_ms + ceiling = min(cap, base * (2**attempt)) + if cfg.jitter: + return math.ceil(random.uniform(0, ceiling)) # nosec B311 # noqa: DUO102 - timing jitter, not security + return ceiling + + +# --------------------------------------------------------------------------- +# Failure detection +# --------------------------------------------------------------------------- + + +def _is_failure(result: Any, cfg: RetryConfig) -> bool: + """Return True if the tool result should trigger a retry. + + The plugin receives result = ToolResult.model_dump(by_alias=True), which has + the shape: {"content": [...], "isError": bool, "structuredContent": {...}}. - max_retries: int = Field(default=2, ge=0) - backoff_base_ms: int = Field(default=200, ge=0) - max_backoff_ms: int = Field(default=5000, ge=0) - retry_on_status: list[int] = Field(default_factory=lambda: [429, 500, 502, 503, 504]) + Three failure signals are checked, in order: + + 1. Outer ``isError`` — set to True by the gateway when the tool raises an + exception. When the gateway can determine the HTTP status code + (e.g. from ``httpx.HTTPStatusError``), it includes the code in + ``structuredContent``. If a status code is present, ``retry_on_status`` + is checked — non-transient errors (400, 401, 404 …) are skipped. + Generic exceptions without a status code are always retried. + + 2. ``structuredContent.status_code`` — when a tool returns a plain dict, the + gateway places it in structuredContent (MCP spec 2025-03-26+ only). Older + servers leave structuredContent=None and this check silently does nothing. + + 3. Text content JSON parsing (opt-in, ``check_text_content: true``) — for + tools on older MCP servers that return HTTP-style error dicts as serialized + JSON in text content instead of raising exceptions. Disabled by default + because it can false-positive on tools that legitimately return status codes + as informational data (e.g. a monitoring tool reporting downstream statuses). + + Args: + result: Serialised ToolResult dict (via model_dump with by_alias=True). + cfg: Retry configuration with retry_on_status list. + + Returns: + True if the result indicates a transient failure that should be retried. + """ + if not isinstance(result, dict): + return False + + # Signal 1: outer MCP-level isError (tool raised an exception). + # Works on all MCP spec versions. When the gateway can determine the + # HTTP status code of the failure (e.g. httpx.HTTPStatusError), it places + # the code in structuredContent so we can honour retry_on_status. + # Generic exceptions with no status code (connection errors, timeouts) + # are always considered retryable. + if result.get("isError") is True: + structured = result.get("structuredContent") + if isinstance(structured, dict): + sc = structured.get("status_code") + if isinstance(sc, int): + return sc in cfg.retry_on_status + return True + + # Signal 2: structuredContent — only populated on MCP spec 2025-03-26+. + structured = result.get("structuredContent") + if isinstance(structured, dict): + if structured.get("isError") is True: + return True + sc_status = structured.get("status_code") + if isinstance(sc_status, int) and sc_status in cfg.retry_on_status: + return True + + # Signal 3: opt-in text content parsing for older MCP servers. + # Only runs when structuredContent was absent (None) — not a double-check. + if cfg.check_text_content and structured is None: + for item in result.get("content", []): + if not isinstance(item, dict) or item.get("type") != "text": + continue + try: + parsed = json.loads(item["text"]) + except (json.JSONDecodeError, KeyError, TypeError): + continue + if not isinstance(parsed, dict): + continue + if parsed.get("isError") is True: + return True + txt_status = parsed.get("status_code") + if isinstance(txt_status, int) and txt_status in cfg.retry_on_status: + return True + + return False + + +# --------------------------------------------------------------------------- +# Plugin +# --------------------------------------------------------------------------- class RetryWithBackoffPlugin(Plugin): - """Attach retry/backoff policy in metadata for observability/orchestration.""" + """Active retry-with-backoff plugin. + + On failure, returns retry_delay_ms > 0 in PluginResult to ask the + gateway to re-invoke the tool after the computed delay. + On success, resets the per-tool failure counter. + """ def __init__(self, config: PluginConfig) -> None: - """Initialize the retry with backoff plugin. + """Initialise the plugin, clamp max_retries to the gateway ceiling, and prepare Rust state managers. Args: - config: Plugin configuration. + config: Plugin configuration from the gateway plugin framework. """ super().__init__(config) - self._cfg = RetryPolicyConfig(**(config.config or {})) + raw_cfg = RetryConfig(**(config.config or {})) + + # Clamp max_retries to the gateway hard ceiling + ceiling = get_settings().max_tool_retries + if raw_cfg.max_retries > ceiling: + log.warning( + "retry_with_backoff: max_retries=%d exceeds gateway ceiling=%d, clamping", + raw_cfg.max_retries, + ceiling, + ) + raw_cfg = raw_cfg.model_copy(update={"max_retries": ceiling}) + + # Clamp per-tool overrides too + for tool_name, overrides in raw_cfg.tool_overrides.items(): + if overrides.get("max_retries", 0) > ceiling: + log.warning( + "retry_with_backoff: tool_overrides[%s].max_retries=%d exceeds ceiling=%d, clamping", + tool_name, + overrides["max_retries"], + ceiling, + ) + overrides["max_retries"] = ceiling + + self._cfg = raw_cfg + + # Build Rust instances with config baked in so the hot check_and_update + # call only crosses the FFI boundary with 4 dynamic args instead of 9. + # One instance per unique config: base + one per tool override. + if _RUST_AVAILABLE: + self._rust: Any = _RustRetryStateManager( + self._cfg.max_retries, + self._cfg.backoff_base_ms, + self._cfg.max_backoff_ms, + self._cfg.jitter, + self._cfg.retry_on_status, + ) + self._rust_overrides: Dict[str, Any] = { + tool_name: _RustRetryStateManager( + overrides.get("max_retries", self._cfg.max_retries), + overrides.get("backoff_base_ms", self._cfg.backoff_base_ms), + overrides.get("max_backoff_ms", self._cfg.max_backoff_ms), + overrides.get("jitter", self._cfg.jitter), + overrides.get("retry_on_status", self._cfg.retry_on_status), + ) + for tool_name, overrides in self._cfg.tool_overrides.items() + } + else: + self._rust = None + self._rust_overrides = {} async def tool_post_invoke(self, payload: ToolPostInvokePayload, context: PluginContext) -> ToolPostInvokeResult: - """Attach retry policy metadata after tool invocation. + """Detect failure and return retry_delay_ms > 0 to request a retry. + + Also attaches retry_policy metadata on every response so downstream + clients and orchestrators can observe the active policy. Args: - payload: Tool invocation result payload. - context: Plugin execution context. + payload: Post-invoke payload containing the tool name and result. + context: Plugin execution context with request_id for state isolation. Returns: - Result with retry policy metadata. + Result with retry_delay_ms set and retry_policy metadata attached. """ - return ToolPostInvokeResult( - metadata={ - "retry_policy": { - "max_retries": self._cfg.max_retries, - "backoff_base_ms": self._cfg.backoff_base_ms, - "max_backoff_ms": self._cfg.max_backoff_ms, - } + tool = payload.name + cfg = _cfg_for(self._cfg, tool) + request_id = context.global_context.request_id + result = payload.result + + retry_policy_meta = { + "retry_policy": { + "max_retries": cfg.max_retries, + "backoff_base_ms": cfg.backoff_base_ms, + "max_backoff_ms": cfg.max_backoff_ms, + "retry_on_status": cfg.retry_on_status, } - ) + } + + # ------------------------------------------------------------------ + # Fast path: delegate to Rust when the extension is available and + # check_text_content is off (Rust handles signals 1 and 2 only). + # + # We pre-extract the two typed signals Python-side before crossing + # the FFI boundary. Passing a raw PyDict into Rust would be slower + # and require more PyO3 boilerplate. Two attribute lookups in Python + # are cheap and keep the Rust code purely typed. + # + # Config is already baked into the Rust instance — the hot call only + # passes the four dynamic args: tool, request_id, is_error, status_code. + # ------------------------------------------------------------------ + if self._rust is not None and not cfg.check_text_content: + is_error: bool = isinstance(result, dict) and result.get("isError") is True + status_code: int | None = None + if isinstance(result, dict): + structured = result.get("structuredContent") + if isinstance(structured, dict): + if structured.get("isError") is True: + is_error = True + sc = structured.get("status_code") + if isinstance(sc, int): + status_code = sc + + rust_inst = self._rust_overrides.get(tool, self._rust) + should_retry, delay_ms = rust_inst.check_and_update( + tool, + request_id, + is_error, + status_code, + ) + if should_retry: + log.debug( + "retry_with_backoff (rust): tool=%s delay_ms=%d", + tool, + delay_ms, + ) + else: + log.debug("retry_with_backoff (rust): tool=%s success/exhausted", tool) + return ToolPostInvokeResult(retry_delay_ms=delay_ms, metadata=retry_policy_meta) + + # ------------------------------------------------------------------ + # Python fallback path — used when: + # * Rust extension is not installed, OR + # * check_text_content=True (signal 3 requires Python dict parsing) + # This path is identical to the pre-Rust implementation. + # ------------------------------------------------------------------ + st = _get_state(tool, request_id) + + if _is_failure(result, cfg): + st.consecutive_failures += 1 + st.last_failure_at = time.monotonic() + + if st.consecutive_failures <= cfg.max_retries: + delay_ms = _compute_delay_ms(st.consecutive_failures - 1, cfg) + log.debug( + "retry_with_backoff: tool=%s failure=%d/%d delay_ms=%d", + tool, + st.consecutive_failures, + cfg.max_retries, + delay_ms, + ) + return ToolPostInvokeResult(retry_delay_ms=delay_ms, metadata=retry_policy_meta) + + # Max retries exhausted — give up, clean up state for this invocation. + log.warning( + "retry_with_backoff: tool=%s exhausted %d retries, returning failure", + tool, + cfg.max_retries, + ) + _del_state(tool, request_id) + return ToolPostInvokeResult(retry_delay_ms=0, metadata=retry_policy_meta) + + # Success — log recovery, clean up state for this invocation. + if st.consecutive_failures > 0: + log.debug("retry_with_backoff: tool=%s recovered after %d failure(s)", tool, st.consecutive_failures) + _del_state(tool, request_id) + return ToolPostInvokeResult(retry_delay_ms=0, metadata=retry_policy_meta) - async def resource_post_fetch(self, payload: ResourcePostFetchPayload, context: PluginContext) -> ResourcePostFetchResult: + async def resource_post_fetch(self, payload: ResourcePostFetchPayload, context: PluginContext) -> ResourcePostFetchResult: # pylint: disable=unused-argument """Attach retry policy metadata after resource fetch. Args: - payload: Resource fetch payload. + payload: Resource fetch payload with URI and content. context: Plugin execution context. Returns: - Result with retry policy metadata. + Result with retry_policy metadata (advisory only, no active retry). """ return ResourcePostFetchResult( metadata={ diff --git a/plugins_rust/encoded_exfil_detection/python/encoded_exfil_detection_rust/__init__.pyi b/plugins_rust/encoded_exfil_detection/python/encoded_exfil_detection_rust/__init__.pyi index c676a2edb41..93e144b9d46 100644 --- a/plugins_rust/encoded_exfil_detection/python/encoded_exfil_detection_rust/__init__.pyi +++ b/plugins_rust/encoded_exfil_detection/python/encoded_exfil_detection_rust/__init__.pyi @@ -3,6 +3,7 @@ import builtins import typing + __all__ = [ "py_scan_container", ] diff --git a/plugins_rust/pii_filter/python/pii_filter_rust/__init__.pyi b/plugins_rust/pii_filter/python/pii_filter_rust/__init__.pyi index 6ca733e1afc..7ee3abf6c23 100644 --- a/plugins_rust/pii_filter/python/pii_filter_rust/__init__.pyi +++ b/plugins_rust/pii_filter/python/pii_filter_rust/__init__.pyi @@ -3,6 +3,7 @@ import builtins import typing + __all__ = [ "PIIDetectorRust", ] @@ -27,6 +28,7 @@ class PIIDetectorRust: print(masked) # "My SSN is [REDACTED] and email is [REDACTED]" ``` """ + def __new__(cls, config: typing.Any) -> PIIDetectorRust: r""" Create a new PII detector @@ -52,6 +54,7 @@ class PIIDetectorRust: * `block_on_detection` (bool): Whether to block on detection * `whitelist_patterns` (list[str]): Regex patterns to exclude from detection """ + def detect(self, text: builtins.str) -> typing.Any: r""" Detect PII in text @@ -72,6 +75,7 @@ class PIIDetectorRust: } ``` """ + def mask(self, text: builtins.str, detections: typing.Any) -> builtins.str: r""" Mask detected PII in text @@ -83,6 +87,7 @@ class PIIDetectorRust: # Returns Masked text with PII replaced """ + def process_nested(self, data: typing.Any, path: builtins.str) -> tuple[builtins.bool, typing.Any, typing.Any]: r""" Process nested data structures (dicts, lists, strings) diff --git a/plugins_rust/retry_with_backoff/Cargo.lock b/plugins_rust/retry_with_backoff/Cargo.lock new file mode 100644 index 00000000000..999acde1f4b --- /dev/null +++ b/plugins_rust/retry_with_backoff/Cargo.lock @@ -0,0 +1,1056 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[package]] +name = "arc-swap" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a07d1f37ff60921c83bdfc7407723bdefe89b44b98a9b772f225c8f9d67141a6" +dependencies = [ + "rustversion", +] + +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + +[[package]] +name = "bumpalo" +version = "3.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" + +[[package]] +name = "cc" +version = "1.2.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a0dd1ca384932ff3641c8718a02769f1698e7563dc6974ffd03346116310423" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chrono" +version = "0.4.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "wasm-bindgen", + "windows-link", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "powerfmt", +] + +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "getopts" +version = "0.2.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df" +dependencies = [ + "unicode-width", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "indexmap" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "inventory" +version = "0.3.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "009ae045c87e7082cb72dab0ccd01ae075dd00141ddc108f43a0ea150a9e7227" +dependencies = [ + "rustversion", +] + +[[package]] +name = "is-macro" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d57a3e447e24c22647738e4607f1df1e0ec6f72e16182c4cd199f647cdfb0e4" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "itertools" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.91" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b49715b7073f385ba4bc528e5747d02e66cb39c6146efb66b781f131f0fb399c" +dependencies = [ + "once_cell", + "wasm-bindgen", +] + +[[package]] +name = "lalrpop-util" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507460a910eb7b32ee961886ff48539633b788a36b65692b95f225b844c82553" + +[[package]] +name = "libc" +version = "0.2.183" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" + +[[package]] +name = "maplit" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e2e65a1a2e43cfcb47a895c4c8b10d1f4a61097f9f254f183aee60cad9c651d" + +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + +[[package]] +name = "memchr" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" + +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-bigint" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-conv" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6673768db2d862beb9b39a78fdcb1a69439615d5794a1be50caa9bc92c81967" + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "numpy" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "778da78c64ddc928ebf5ad9df5edf0789410ff3bdbf3619aed51cd789a6af1e2" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash 2.1.1", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "ordered-float" +version = "5.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f4779c6901a562440c3786d08192c6fbda7c1c2060edd10006b05ee35d10f2d" +dependencies = [ + "num-traits", +] + +[[package]] +name = "phf" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_codegen" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a" +dependencies = [ + "phf_generator", + "phf_shared", +] + +[[package]] +name = "phf_generator" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" +dependencies = [ + "phf_shared", + "rand", +] + +[[package]] +name = "phf_shared" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5" +dependencies = [ + "siphasher", +] + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "091397be61a01d4be58e7841595bd4bfedb15f1cd54977d79b8271e94ed799a3" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.28.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf85e27e86080aafd5a22eae58a162e133a589551542b3e5cee4beb27e54f8e1" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.28.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8bf94ee265674bf76c09fa430b0e99c26e319c945d96ca0d5a8215f31bf81cf7" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.28.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "491aa5fc66d8059dd44a75f4580a2962c1862a1c2945359db36f6c2818b748dc" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-log" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26c2ec80932c5c3b2d4fbc578c9b56b2d4502098587edb8bef5b6bfcad43682e" +dependencies = [ + "arc-swap", + "log", + "pyo3", +] + +[[package]] +name = "pyo3-macros" +version = "0.28.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5d671734e9d7a43449f8480f8b38115df67bef8d21f76837fa75ee7aaa5e52e" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.28.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22faaa1ce6c430a1f71658760497291065e6450d7b5dc2bcf254d49f66ee700a" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + +[[package]] +name = "pyo3-stub-gen" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b159f7704044f57d058f528a6f1f22a0a0a327dcb595c5fb38beae658e0338d6" +dependencies = [ + "anyhow", + "chrono", + "either", + "indexmap", + "inventory", + "itertools 0.14.0", + "log", + "maplit", + "num-complex", + "numpy", + "ordered-float", + "pyo3", + "pyo3-stub-gen-derive", + "rustpython-parser", + "serde", + "serde_json", + "time", + "toml", +] + +[[package]] +name = "pyo3-stub-gen-derive" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8c79e7c5b1fcec7c39ab186594658a971c59911eb6fbab5a5932cf2318534be" +dependencies = [ + "heck", + "indexmap", + "proc-macro2", + "quote", + "rustpython-parser", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "retry_with_backoff" +version = "0.1.0" +dependencies = [ + "log", + "pyo3", + "pyo3-log", + "pyo3-stub-gen", + "rand", +] + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + +[[package]] +name = "rustpython-ast" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cdaf8ee5c1473b993b398c174641d3aa9da847af36e8d5eb8291930b72f31a5" +dependencies = [ + "is-macro", + "num-bigint", + "rustpython-parser-core", + "static_assertions", +] + +[[package]] +name = "rustpython-parser" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "868f724daac0caf9bd36d38caf45819905193a901e8f1c983345a68e18fb2abb" +dependencies = [ + "anyhow", + "is-macro", + "itertools 0.11.0", + "lalrpop-util", + "log", + "num-bigint", + "num-traits", + "phf", + "phf_codegen", + "rustc-hash 1.1.0", + "rustpython-ast", + "rustpython-parser-core", + "tiny-keccak", + "unic-emoji-char", + "unic-ucd-ident", + "unicode_names2", +] + +[[package]] +name = "rustpython-parser-core" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4b6c12fa273825edc7bccd9a734f0ad5ba4b8a2f4da5ff7efe946f066d0f4ad" +dependencies = [ + "is-macro", + "memchr", + "rustpython-parser-vendored", +] + +[[package]] +name = "rustpython-parser-vendored" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04fcea49a4630a3a5d940f4d514dc4f575ed63c14c3e3ed07146634aed7f67a6" +dependencies = [ + "memchr", + "once_cell", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_spanned" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "876ac351060d4f882bb1032b6369eb0aef79ad9df1ea8bc404874d8cc3d0cd98" +dependencies = [ + "serde_core", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "siphasher" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "time" +version = "0.3.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c" +dependencies = [ + "deranged", + "num-conv", + "powerfmt", + "serde_core", + "time-core", +] + +[[package]] +name = "time-core" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca" + +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + +[[package]] +name = "toml" +version = "1.1.0+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8195ca05e4eb728f4ba94f3e3291661320af739c4e43779cbdfae82ab239fcc" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime", + "toml_parser", + "toml_writer", + "winnow", +] + +[[package]] +name = "toml_datetime" +version = "1.1.0+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97251a7c317e03ad83774a8752a7e81fb6067740609f75ea2b585b569a59198f" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_parser" +version = "1.1.0+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2334f11ee363607eb04df9b8fc8a13ca1715a72ba8662a26ac285c98aabb4011" +dependencies = [ + "winnow", +] + +[[package]] +name = "toml_writer" +version = "1.1.0+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d282ade6016312faf3e41e57ebbba0c073e4056dab1232ab1cb624199648f8ed" + +[[package]] +name = "unic-char-property" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8c57a407d9b6fa02b4795eb81c5b6652060a15a7903ea981f3d723e6c0be221" +dependencies = [ + "unic-char-range", +] + +[[package]] +name = "unic-char-range" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0398022d5f700414f6b899e10b8348231abf9173fa93144cbc1a43b9793c1fbc" + +[[package]] +name = "unic-common" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80d7ff825a6a654ee85a63e80f92f054f904f21e7d12da4e22f9834a4aaa35bc" + +[[package]] +name = "unic-emoji-char" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b07221e68897210270a38bde4babb655869637af0f69407f96053a34f76494d" +dependencies = [ + "unic-char-property", + "unic-char-range", + "unic-ucd-version", +] + +[[package]] +name = "unic-ucd-ident" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e230a37c0381caa9219d67cf063aa3a375ffed5bf541a452db16e744bdab6987" +dependencies = [ + "unic-char-property", + "unic-char-range", + "unic-ucd-version", +] + +[[package]] +name = "unic-ucd-version" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96bd2f2237fe450fcd0a1d2f5f4e91711124f7857ba2e964247776ebeeb7b0c4" +dependencies = [ + "unic-common", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode_names2" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1673eca9782c84de5f81b82e4109dcfb3611c8ba0d52930ec4a9478f547b2dd" +dependencies = [ + "phf", + "unicode_names2_generator", +] + +[[package]] +name = "unicode_names2_generator" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b91e5b84611016120197efd7dc93ef76774f4e084cd73c9fb3ea4a86c570c56e" +dependencies = [ + "getopts", + "log", + "phf_codegen", + "rand", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasm-bindgen" +version = "0.2.114" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6532f9a5c1ece3798cb1c2cfdba640b9b3ba884f5db45973a6f442510a87d38e" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.114" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18a2d50fcf105fb33bb15f00e7a77b772945a2ee45dcf454961fd843e74c18e6" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.114" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03ce4caeaac547cdf713d280eda22a730824dd11e6b8c3ca9e42247b25c631e3" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.114" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75a326b8c223ee17883a4251907455a2431acc2791c98c26279376490c378c16" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "winnow" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a90e88e4667264a994d34e6d1ab2d26d398dcdca8b7f52bec8668957517fc7d8" + +[[package]] +name = "zerocopy" +version = "0.8.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efbb2a062be311f2ba113ce66f697a4dc589f85e78a4aea276200804cea0ed87" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e8bc7269b54418e7aeeef514aa68f8690b8c0489a06b0136e5f57c4c5ccab89" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/plugins_rust/retry_with_backoff/Cargo.toml b/plugins_rust/retry_with_backoff/Cargo.toml new file mode 100644 index 00000000000..62120527035 --- /dev/null +++ b/plugins_rust/retry_with_backoff/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "retry_with_backoff" +version = "0.1.0" +edition = "2024" +authors = ["ContextForge Contributors"] +license = "Apache-2.0" +description = "Rust-accelerated retry state manager for the retry-with-backoff plugin" + +[lib] +# cdylib = C-compatible dynamic library (what Python loads as a .so) +# rlib = Rust static library (needed for tests / benchmarks) +name = "retry_with_backoff_rust" +crate-type = ["cdylib", "rlib"] + +[[bin]] +name = "stub_gen" +path = "src/bin/stub_gen.rs" + +[dependencies] +log = "0.4.29" +# pyo3: the bridge between Python and Rust +# abi3-py311 = compile once, works on Python 3.11+ +pyo3 = { version = "0.28.2", features = ["abi3-py311"] } +pyo3-log = "0.13.3" +pyo3-stub-gen = "0.19" + +# rand: random number generation — needed for full-jitter backoff +# ≈ Python's random.uniform() +rand = "0.8" + +[profile.release] +opt-level = 3 +lto = "fat" +codegen-units = 1 +strip = true diff --git a/plugins_rust/retry_with_backoff/Makefile b/plugins_rust/retry_with_backoff/Makefile new file mode 100644 index 00000000000..69b3f981cba --- /dev/null +++ b/plugins_rust/retry_with_backoff/Makefile @@ -0,0 +1,236 @@ +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# 🦀 RETRY-WITH-BACKOFF - Makefile +# Rust-accelerated retry state manager for the retry-with-backoff plugin +# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +# +# Usage: make or just `make help` +# +# help: 🦀 RETRY-WITH-BACKOFF (Rust + Python extension build & automation) +# ───────────────────────────────────────────────────────────────────────── + +# ============================================================================= +# 📖 DYNAMIC HELP +# ============================================================================= +.PHONY: help +help: + @grep '^# help\:' $(firstword $(MAKEFILE_LIST)) | sed 's/^# help\: //' + +# ============================================================================= +# 📦 PROJECT METADATA +# ============================================================================= +PACKAGE_NAME := retry_with_backoff +VERSION ?= $(shell cargo metadata --format-version 1 --no-deps 2>/dev/null | jq -r '.packages[0].version' || echo "0.1.0") +DIST_DIR := target +PYTHON ?= python3 + +# Colors for output +BLUE := \033[0;34m +GREEN := \033[0;32m +YELLOW:= \033[0;33m +RED := \033[0;31m +NC := \033[0m # No Color + +# ============================================================================= +# 🔍 LINTING & FORMAT +# ============================================================================= +# help: 🔍 LINTING & FORMAT +# help: fmt - Format Rust code with rustfmt +# help: fmt-check - Check Rust code formatting (CI) +# help: clippy - Run clippy lints +# help: cargo-check - Run cargo check +.PHONY: fmt fmt-check clippy cargo-check + +fmt: + @echo "$(GREEN)Formatting code...$(NC)" + cargo fmt + +fmt-check: + @echo "$(GREEN)Checking code format...$(NC)" + cargo fmt -- --check + +clippy: + @echo "$(GREEN)Running clippy...$(NC)" + cargo clippy -- -D warnings + +cargo-check: + @cargo check + +# ============================================================================= +# 🧪 TESTS +# ============================================================================= +# help: 🧪 TESTS +# help: test - Run Rust unit tests +# help: test-verbose - Run Rust tests with verbose output +# help: test-python - Run Python integration tests +# help: test-all - Run both Rust and Python tests +.PHONY: test test-verbose test-python test-all + +test: + @echo "$(GREEN)Running retry_with_backoff tests...$(NC)" + cargo test + +test-verbose: + @echo "$(GREEN)Running retry_with_backoff tests (verbose)...$(NC)" + cargo test -- --nocapture + +test-python: + @echo "$(GREEN)Running Python unit tests...$(NC)" + cd ../.. && uv run pytest -k retry_with_backoff -v + +test-all: test test-python + +# ============================================================================= +# 📊 COVERAGE +# ============================================================================= +# help: 📊 COVERAGE +# help: coverage-tools - Install cargo-llvm-cov and llvm-tools-preview +# help: coverage - Generate code coverage report with llvm-cov +# help: coverage-summary - Generate code coverage summary (aliases: cs, cov-sum) +.PHONY: coverage-tools coverage cov cv coverage-summary cs cov-sum + +coverage-tools: + @echo "$(GREEN)Installing cargo-llvm-cov and llvm-tools-preview...$(NC)" + @cargo install cargo-llvm-cov + @rustup component add llvm-tools-preview + @echo "$(GREEN)Coverage tools installed$(NC)" + +coverage cov cv: coverage-tools + @echo "$(GREEN)Generating code coverage report...$(NC)" + @cargo llvm-cov --html --ignore-filename-regex "src/bin/.*" + @echo "$(GREEN)Coverage report generated at target/llvm-cov/html/index.html$(NC)" + @xdg-open target/llvm-cov/html/index.html 2>/dev/null || true + +coverage-summary cs cov-sum: coverage-tools + @echo "$(GREEN)Generating code coverage summary...$(NC)" + @cargo llvm-cov --summary-only --ignore-filename-regex "src/bin/.*" + +# ============================================================================= +# 🛠 BUILD (maturin for Python extension) +# ============================================================================= +# help: 🛠 BUILD +# help: build - Build release extension (no install) +# help: build-target - Build for specific target (use TARGET=...) +# help: stub-gen - Generate Python type stubs (.pyi files) +# help: install - Build and install wheel +.PHONY: stub-gen build build-target install + +stub-gen: + @echo "$(GREEN)Generating Python type stubs...$(NC)" + @cargo run --bin stub_gen + @echo "$(GREEN)Type stubs generated$(NC)" + +build: stub-gen + @echo "$(GREEN)Building $(PACKAGE_NAME)...$(NC)" + @cd ../.. && uv run maturin build --release --manifest-path plugins_rust/retry_with_backoff/Cargo.toml + @echo "$(GREEN)Build complete$(NC)" + +build-target: stub-gen + @echo "$(GREEN)Building for target: $(TARGET)...$(NC)" + @uv run maturin build --release --target $(TARGET) + @echo "$(GREEN)Build complete for $(TARGET)$(NC)" + +install: stub-gen + @echo "$(GREEN)Installing $(PACKAGE_NAME) plugin...$(NC)" + @cd ../.. && uv run maturin develop --release --manifest-path plugins_rust/retry_with_backoff/Cargo.toml + @echo "$(GREEN)Installation complete$(NC)" + +# ============================================================================= +# 📊 BENCHMARKS +# ============================================================================= +# help: 📊 BENCHMARKS +# help: bench - Run Criterion benchmarks (skipped — no benches) +# help: bench-compare - Compare against baseline +# help: benchmark - Run Rust vs Python performance comparison +# help: benchmark-quick - Quick benchmark (100 iterations, 10 warmup) +.PHONY: bench bench-baseline bench-compare compare benchmark benchmark-quick + +bench: + @echo "$(YELLOW)No Criterion benchmarks defined for $(PACKAGE_NAME) — skipping$(NC)" + +bench-baseline: + @echo "$(YELLOW)No benchmarks defined for $(PACKAGE_NAME) — skipping$(NC)" + +bench-compare: + @echo "$(YELLOW)No benchmarks defined for $(PACKAGE_NAME) — skipping$(NC)" + +compare: install + @echo "$(GREEN)Running performance comparison (Python vs Rust)...$(NC)" + @echo "$(YELLOW)Installing plugin first...$(NC)" + @$(MAKE) --no-print-directory install + @echo "" + @echo "$(YELLOW)Running comparison script...$(NC)" + cd ../../ && uv run python3 plugins_rust/retry_with_backoff/compare_performance.py + +benchmark: install + @echo "$(GREEN)Running Rust vs Python performance comparison...$(NC)" + @echo "" + cd ../../ && uv run python3 plugins_rust/retry_with_backoff/compare_performance.py + +benchmark-quick: install + @echo "$(GREEN)Running quick performance comparison (100 iterations)...$(NC)" + @echo "" + cd ../../ && uv run python3 plugins_rust/retry_with_backoff/compare_performance.py --iterations 100 --warmup 10 + +# ============================================================================= +# 🧹 CLEANUP +# ============================================================================= +# help: 🧹 CLEANUP +# help: clean - Remove build artifacts +# help: clean-all - Remove all build artifacts including wheels +# help: uninstall - Uninstall plugin from Python environment +.PHONY: clean clean-all uninstall + +uninstall: + @echo "$(YELLOW)Uninstalling $(PACKAGE_NAME)...$(NC)" + @uv pip uninstall -y $(PACKAGE_NAME) 2>/dev/null || pip uninstall -y $(PACKAGE_NAME) 2>/dev/null || true + @echo "$(GREEN)$(PACKAGE_NAME) uninstalled$(NC)" + +clean: + @echo "$(YELLOW)Cleaning build artifacts...$(NC)" + cargo clean + rm -rf target/ + rm -rf coverage/ + find . -type f -name "*.whl" -delete + find . -type f -name "*.pyc" -delete + find . -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null || true + +clean-all: clean + @echo "$(RED)Cleaning all generated files...$(NC)" + rm -rf ~/.cargo/registry/cache/ + rm -rf ~/.cargo/git/db/ + +# ============================================================================= +# 📚 DOCUMENTATION +# ============================================================================= +# help: 📚 DOCUMENTATION +# help: doc - Generate Rust documentation +# help: doc-open - Generate and open documentation +.PHONY: doc doc-open + +doc: + @echo "$(GREEN)Building documentation...$(NC)" + cargo doc --no-deps --document-private-items + +doc-open: doc + @echo "$(GREEN)Opening documentation...$(NC)" + cargo doc --no-deps --document-private-items --open + +# ============================================================================= +# 🔧 DEVELOPMENT HELPERS +# ============================================================================= +# help: 🔧 DEVELOPMENT HELPERS +# help: verify - Verify plugin installation +# help: check-all - Run all checks (fmt, clippy, test) +.PHONY: verify check-all + +verify: + @echo "$(GREEN)Verifying $(PACKAGE_NAME) installation...$(NC)" + @uv run python -c "import retry_with_backoff_rust; print('✅ retry_with_backoff_rust available')" || echo "⚠️ retry_with_backoff_rust not installed" + +check-all: fmt-check clippy test + @echo "$(GREEN)✔ All checks passed$(NC)" + +# --------------------------------------------------------------------------- +# Default goal +# --------------------------------------------------------------------------- +.DEFAULT_GOAL := help diff --git a/plugins_rust/retry_with_backoff/compare_performance.py b/plugins_rust/retry_with_backoff/compare_performance.py new file mode 100755 index 00000000000..b10495e7a6c --- /dev/null +++ b/plugins_rust/retry_with_backoff/compare_performance.py @@ -0,0 +1,624 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Performance comparison: Rust vs Python for retry_with_backoff plugin. + +This benchmark provides a fair apples-to-apples comparison by using native +Python objects for both implementations, eliminating JSON serialization overhead. + +Measurements: +- Python (native): Baseline Python implementation +- Rust (native): High-performance Rust implementation via PyO3 + +Usage: + python compare_performance.py + python compare_performance.py --iterations 10000 --warmup 100 +""" + +import argparse +import random +import statistics +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +# --------------------------------------------------------------------------- +# Import Rust implementation +# --------------------------------------------------------------------------- +try: + from retry_with_backoff_rust.retry_with_backoff_rust import ( + RetryStateManager as RustRetryStateManager, + ) + + RUST_AVAILABLE = True +except ImportError: + RUST_AVAILABLE = False + RustRetryStateManager = None + print("⚠️ Rust implementation not available. Build it with:") + print(" cd plugins_rust/retry_with_backoff && make install") + print() + +# --------------------------------------------------------------------------- +# Import Python implementation +# --------------------------------------------------------------------------- +plugins_path = Path(__file__).parent.parent.parent / "plugins" / "retry_with_backoff" +if plugins_path.exists(): + sys.path.insert(0, str(plugins_path)) + mcpgateway_path = Path(__file__).parent.parent.parent / "mcpgateway" + if mcpgateway_path.exists(): + sys.path.insert(0, str(mcpgateway_path)) +else: + print(f"⚠️ Warning: Python implementation path not found: {plugins_path}") + print() + +try: + from retry_with_backoff import ( + RetryConfig, + _ToolRetryState, + _compute_delay_ms, + _del_state, + _get_state, + _is_failure, + ) + + PYTHON_AVAILABLE = True +except ImportError as e: + PYTHON_AVAILABLE = False + print(f"⚠️ Python implementation not available: {e}") + print(" Make sure pydantic is installed: uv pip install pydantic") + print() + + +# --------------------------------------------------------------------------- +# Python fallback implementation (mirrors Rust logic) +# --------------------------------------------------------------------------- + + +class PythonRetryStateManager: + """Python implementation mirroring Rust RetryStateManager API.""" + + def __init__( + self, + max_retries: int, + base_ms: int, + max_ms: int, + jitter: bool, + retry_on_status: List[int], + ): + self.max_retries = max_retries + self.base_ms = base_ms + self.max_ms = max_ms + self.jitter = jitter + self.retry_on_status = set(retry_on_status) + self._state: Dict[str, _ToolRetryState] = {} + + def _make_key(self, tool: str, request_id: str) -> str: + return f"{tool}:{request_id}" + + def check_and_update( + self, + tool: str, + request_id: str, + is_error: bool, + status_code: Optional[int], + ) -> Tuple[bool, int]: + """Check failure and update state, returning (should_retry, delay_ms).""" + failed = is_error or (status_code is not None and status_code in self.retry_on_status) + + key = self._make_key(tool, request_id) + + if failed: + state = self._state.get(key) + if state is None: + state = _ToolRetryState() + self._state[key] = state + + state.consecutive_failures += 1 + state.last_failure_at = time.monotonic() + + if state.consecutive_failures <= self.max_retries: + attempt = state.consecutive_failures - 1 + delay = _compute_delay_ms( + attempt, + RetryConfig( + max_retries=self.max_retries, + backoff_base_ms=self.base_ms, + max_backoff_ms=self.max_ms, + jitter=self.jitter, + retry_on_status=list(self.retry_on_status), + ), + ) + return (True, delay) + else: + self._del_state(tool, request_id) + return (False, 0) + else: + self._del_state(tool, request_id) + return (False, 0) + + def _del_state(self, tool: str, request_id: str) -> None: + key = self._make_key(tool, request_id) + self._state.pop(key, None) + + +# --------------------------------------------------------------------------- +# Benchmark functions +# --------------------------------------------------------------------------- + + +def benchmark_implementation( + impl: Any, + tool_names: List[str], + request_ids: List[str], + failure_rate: float, + iterations: int, + warmup: int = 5, +) -> Tuple[List[float], int]: + """Benchmark an implementation (Python or Rust). + + Args: + impl: RetryStateManager instance (Python or Rust) + tool_names: List of tool names to simulate + request_ids: List of request IDs to simulate + failure_rate: Fraction of calls that should fail (0.0-1.0) + iterations: Number of benchmark iterations + warmup: Number of warmup iterations + + Returns: + Tuple of (list of times in seconds, retry count) + """ + # Warmup phase + for _ in range(warmup): + for tool in tool_names[:3]: + for req_id in request_ids[:3]: + is_error = random.random() < failure_rate + status_code = 503 if is_error and random.random() < 0.7 else None + impl.check_and_update(tool, req_id, is_error, status_code) + + # Benchmark phase + times = [] + retry_count = 0 + + random.seed(42) # Reproducible results + + for _ in range(iterations): + start = time.perf_counter() + + for tool in tool_names: + for req_id in request_ids: + is_error = random.random() < failure_rate + if is_error: + status_code = 503 if random.random() < 0.7 else None + else: + status_code = None + + should_retry, _ = impl.check_and_update(tool, req_id, is_error, status_code) + if should_retry: + retry_count += 1 + + times.append(time.perf_counter() - start) + + return times, retry_count + + +def run_load_scenario( + name: str, + tool_count: int, + request_count: int, + failure_rate: float, + iterations: int, + warmup: int, +) -> Optional[Dict[str, Any]]: + """Run load-based benchmark scenario. + + Returns: + Dictionary with scenario results, or None if benchmark failed. + """ + print(f"\n{'=' * 70}") + print(f"Scenario: {name}") + print(f" Tools: {tool_count}, Requests: {request_count}, Failure rate: {failure_rate:.0%}") + print(f"{'=' * 70}") + + tool_names = [f"tool_{i}" for i in range(tool_count)] + request_ids = [f"req_{i}" for i in range(request_count)] + + config = { + "max_retries": 3, + "base_ms": 100, + "max_ms": 5000, + "jitter": False, # Disabled for consistent measurements + "retry_on_status": [429, 500, 502, 503, 504], + } + + results = {} + + # Benchmark Python + if PYTHON_AVAILABLE: + print("Running Python...", end=" ", flush=True) + py_impl = PythonRetryStateManager(**config) + py_times, py_retries = benchmark_implementation(py_impl, tool_names, request_ids, failure_rate, iterations, warmup) + py_mean = statistics.mean(py_times) * 1_000_000 # Convert to microseconds + py_median = statistics.median(py_times) * 1_000_000 + py_stdev = statistics.stdev(py_times) * 1_000_000 if len(py_times) > 1 else 0 + results["python"] = { + "mean": py_mean, + "median": py_median, + "stdev": py_stdev, + "retries": py_retries, + } + print(f"✓ ({py_mean:.3f} µs/iter, {py_retries} retries)") + else: + print("Running Python... ✗ (not available)") + + # Benchmark Rust + if RUST_AVAILABLE: + print("Running Rust...", end=" ", flush=True) + rust_impl = RustRetryStateManager(**config) + rust_times, rust_retries = benchmark_implementation(rust_impl, tool_names, request_ids, failure_rate, iterations, warmup) + rust_mean = statistics.mean(rust_times) * 1_000_000 # Convert to microseconds + rust_median = statistics.median(rust_times) * 1_000_000 + rust_stdev = statistics.stdev(rust_times) * 1_000_000 if len(rust_times) > 1 else 0 + results["rust"] = { + "mean": rust_mean, + "median": rust_median, + "stdev": rust_stdev, + "retries": rust_retries, + } + print(f"✓ ({rust_mean:.3f} µs/iter, {rust_retries} retries)") + else: + print("Running Rust... ✗ (not available)") + + # Calculate and display results + print("\n📊 Results:") + + if "python" in results: + py = results["python"] + print(f" Python: {py['mean']:>10.3f} µs ±{py['stdev']:>8.3f} (median: {py['median']:>10.3f})") + + if "rust" in results: + rust = results["rust"] + print(f" Rust: {rust['mean']:>10.3f} µs ±{rust['stdev']:>8.3f} (median: {rust['median']:>10.3f})") + + speedup = None + if "python" in results and "rust" in results: + speedup = py_mean / rust_mean if rust_mean > 0 else 0 + print(f"\n 🚀 Speedup: {speedup:.2f}x faster with Rust") + + if abs(py_retries - rust_retries) > py_retries * 0.01: + print(f"\n ⚠️ WARNING: Different retry counts! Python={py_retries}, Rust={rust_retries}") + else: + print(f" ✓ Retry counts match (Python={py_retries}, Rust={rust_retries})") + + return { + "name": name, + "type": "load", + "config": {"tools": tool_count, "requests": request_count, "failure_rate": failure_rate}, + "python": results.get("python"), + "rust": results.get("rust"), + "speedup": speedup, + } + + +def reset_state(): + """Reset Python state between benchmarks.""" + if PYTHON_AVAILABLE: + from retry_with_backoff import _STATE + + _STATE.clear() + + +def benchmark_sequential( + impl_name: str, + calls: List[Dict[str, Any]], + config: RetryConfig, + rust_mgr: Optional[Any] = None, + warmup: int = 5, +) -> Tuple[List[float], int]: + """Benchmark sequential call pattern (original compare_performance.py style).""" + reset_state() + + times = [] + retry_count = 0 + tool = "test_tool" + + # Warmup + for i in range(warmup): + if impl_name == "python": + _get_state(tool, f"warmup_{i}") + _del_state(tool, f"warmup_{i}") + else: + pass # Rust handles state internally + + for i, call in enumerate(calls): + req_id = f"seq_{i}" + start = time.perf_counter() + + if impl_name == "python": + result = {"isError": call["is_error"], "content": [], "structuredContent": None} + if call["status_code"] is not None: + result["structuredContent"] = {"status_code": call["status_code"]} + + st = _get_state(tool, req_id) + if _is_failure(result, config): + st.consecutive_failures += 1 + if st.consecutive_failures <= config.max_retries: + _compute_delay_ms(st.consecutive_failures - 1, config) + retry_count += 1 + _del_state(tool, req_id) + else: + _del_state(tool, req_id) + else: + _del_state(tool, req_id) + else: + # Rust + should_retry, _ = rust_mgr.check_and_update(tool, req_id, call["is_error"], call["status_code"]) + if should_retry: + retry_count += 1 + + times.append(time.perf_counter() - start) + + return times, retry_count + + +def run_sequential_scenario( + name: str, + calls: List[Dict[str, Any]], + config: RetryConfig, + warmup: int, +) -> Optional[Dict[str, Any]]: + """Run sequential call pattern scenario. + + Returns: + Dictionary with scenario results, or None if benchmark failed. + """ + print(f"\n{'=' * 70}") + print(f"Scenario: {name}") + print(f"{'=' * 70}") + + results = {} + + # Python + if PYTHON_AVAILABLE: + print("Running Python...", end=" ", flush=True) + py_times, py_count = benchmark_sequential("python", calls, config, warmup) + py_mean = statistics.mean(py_times) * 1_000_000 + py_median = statistics.median(py_times) * 1_000_000 + py_stdev = statistics.stdev(py_times) * 1_000_000 if len(py_times) > 1 else 0 + results["python"] = { + "mean": py_mean, + "median": py_median, + "stdev": py_stdev, + "retries": py_count, + } + print(f"✓ ({py_mean:.3f} µs/call, {py_count} retries)") + else: + print("Running Python... ✗ (not available)") + + # Rust + if RUST_AVAILABLE: + rust_mgr = RustRetryStateManager( + config.max_retries, + config.backoff_base_ms, + config.max_backoff_ms, + config.jitter, + config.retry_on_status, + ) + print("Running Rust...", end=" ", flush=True) + rust_times, rust_count = benchmark_sequential("rust", calls, config, rust_mgr=rust_mgr, warmup=warmup) + rust_mean = statistics.mean(rust_times) * 1_000_000 + rust_median = statistics.median(rust_times) * 1_000_000 + rust_stdev = statistics.stdev(rust_times) * 1_000_000 if len(rust_times) > 1 else 0 + results["rust"] = { + "mean": rust_mean, + "median": rust_median, + "stdev": rust_stdev, + "retries": rust_count, + } + print(f"✓ ({rust_mean:.3f} µs/call, {rust_count} retries)") + else: + print("Running Rust... ✗ (not available)") + + print("\n📊 Results:") + if "python" in results: + py = results["python"] + print(f" Python: {py['mean']:>10.3f} µs ±{py['stdev']:>8.3f} (median: {py['median']:>10.3f})") + if "rust" in results: + rust = results["rust"] + print(f" Rust: {rust['mean']:>10.3f} µs ±{rust['stdev']:>8.3f} (median: {rust['median']:>10.3f})") + + speedup = None + if "python" in results and "rust" in results: + speedup = py_mean / rust_mean if rust_mean > 0 else 0 + print(f"\n 🚀 Speedup: {speedup:.2f}x faster with Rust") + if py_count != rust_count: + print(f"\n ⚠️ WARNING: Different retry counts! Python={py_count}, Rust={rust_count}") + + return { + "name": name, + "type": "sequential", + "python": results.get("python"), + "rust": results.get("rust"), + "speedup": speedup, + } + + +def generate_sequential_scenarios(iterations: int) -> List[Dict[str, Any]]: + """Generate sequential test scenarios with different failure patterns.""" + scenarios = [] + + scenarios.append( + { + "name": "All successes", + "calls": [{"is_error": False, "status_code": None} for _ in range(iterations)], + } + ) + + scenarios.append( + { + "name": "All failures (exhaust)", + "calls": [{"is_error": True, "status_code": None} for _ in range(iterations)], + } + ) + + scenarios.append( + { + "name": "Mixed (50/50)", + "calls": [{"is_error": (i % 2 == 0), "status_code": None} for i in range(iterations)], + } + ) + + scenarios.append( + { + "name": "Rate limiting (429)", + "calls": [{"is_error": False, "status_code": 429 if i % 3 == 0 else 200} for i in range(iterations)], + } + ) + + scenarios.append( + { + "name": "Server errors", + "calls": [ + { + "is_error": False, + "status_code": [500, 502, 503][i % 3] if i % 4 != 0 else 200, + } + for i in range(iterations) + ], + } + ) + + return scenarios + + +def main(): + """Run performance comparison benchmarks.""" + parser = argparse.ArgumentParser(description="Rust vs Python performance comparison for retry_with_backoff") + parser.add_argument( + "--iterations", + type=int, + default=10000, + help="Benchmark iterations per scenario", + ) + parser.add_argument( + "--warmup", + type=int, + default=100, + help="Warmup iterations", + ) + args = parser.parse_args() + + print("🔄 Retry With Backoff Performance Comparison") + print(f"{'=' * 70}") + print(f"Iterations: {args.iterations} (+ {args.warmup} warmup)") + print(f"Rust available: {'✓' if RUST_AVAILABLE else '✗'}") + print(f"Python available: {'✓' if PYTHON_AVAILABLE else '✗'}") + + if not RUST_AVAILABLE and not PYTHON_AVAILABLE: + print("\n❌ Error: Neither implementation is available!") + print(" Install at least one implementation to run benchmarks.") + sys.exit(1) + + config = RetryConfig( + max_retries=3, + backoff_base_ms=200, + max_backoff_ms=5000, + jitter=False, + retry_on_status=[429, 500, 502, 503, 504], + ) + + all_results: List[Dict[str, Any]] = [] + + print("\n" + "=" * 70) + print("PART 1: Sequential Call Patterns (per-call overhead)") + print("=" * 70) + + sequential_scenarios = generate_sequential_scenarios(min(args.iterations, 1000)) + for scenario in sequential_scenarios: + result = run_sequential_scenario( + scenario["name"], + scenario["calls"], + config, + args.warmup, + ) + if result: + all_results.append(result) + + print("\n" + "=" * 70) + print("PART 2: Concurrent Load Patterns (batch throughput)") + print("=" * 70) + + load_scenarios = [ + ("Low load (1 tool, 10 reqs, 10% fail)", 1, 10, 0.1), + ("Medium load (5 tools, 50 reqs, 30% fail)", 5, 50, 0.3), + ("High load (10 tools, 100 reqs, 50% fail)", 10, 100, 0.5), + ("Stress test (20 tools, 200 reqs, 70% fail)", 20, 200, 0.7), + ] + + for name, tool_count, request_count, failure_rate in load_scenarios: + result = run_load_scenario( + name, + tool_count, + request_count, + failure_rate, + args.iterations, + args.warmup, + ) + if result: + all_results.append(result) + + # Display summary + print(f"\n{'=' * 70}") + print("📊 PERFORMANCE SUMMARY") + print(f"{'=' * 70}") + + if all_results: + # Calculate overall statistics + speedups = [r["speedup"] for r in all_results if r["speedup"] is not None] + + if speedups: + avg_speedup = statistics.mean(speedups) + min_speedup = min(speedups) + max_speedup = max(speedups) + + print("\n🚀 Rust Speedup Overview:") + print(f" Average: {avg_speedup:.2f}x faster") + print(f" Min: {min_speedup:.2f}x faster") + print(f" Max: {max_speedup:.2f}x faster") + + # Per-scenario breakdown + print("\n📈 Scenario Breakdown:") + print(f"{'Scenario':<45} {'Python (µs)':>14} {'Rust (µs)':>12} {'Speedup':>10}") + print(f"{'-' * 45} {'-' * 14} {'-' * 12} {'-' * 10}") + + for result in all_results: + name = result["name"][:44] + py_time = result["python"]["mean"] if result["python"] else float("inf") + rust_time = result["rust"]["mean"] if result["rust"] else float("inf") + speedup = result["speedup"] if result["speedup"] else float("inf") + + py_str = f"{py_time:>14.3f}" if py_time != float("inf") else "N/A" + rust_str = f"{rust_time:>12.3f}" if rust_time != float("inf") else "N/A" + speedup_str = f"{speedup:>10.2f}x" if speedup != float("inf") else "N/A" + + print(f"{name:<45} {py_str:>14} {rust_str:>12} {speedup_str:>10}") + + # Overall recommendation + print("\n💡 Recommendation:") + if speedups: + if avg_speedup >= 2.0: + print(" Rust implementation provides significant performance benefits") + print(f" ({avg_speedup:.1f}x average speedup). Recommended for production.") + elif avg_speedup >= 1.5: + print(f" Rust implementation offers moderate speedup ({avg_speedup:.1f}x).") + print(" Consider using for high-throughput scenarios.") + else: + print(f" Rust implementation shows marginal improvement ({avg_speedup:.1f}x).") + print(" Python may be sufficient for low-load use cases.") + + print(f"\n{'=' * 70}") + print("✅ Benchmark complete!") + print(f"{'=' * 70}\n") + + +if __name__ == "__main__": + main() diff --git a/plugins_rust/retry_with_backoff/deny.toml b/plugins_rust/retry_with_backoff/deny.toml new file mode 100644 index 00000000000..142f5157ff8 --- /dev/null +++ b/plugins_rust/retry_with_backoff/deny.toml @@ -0,0 +1,27 @@ +# Cargo-deny config: license and policy checks for this crate. +# See https://embarkstudios.github.io/cargo-deny/ + +[licenses] +unused-allowed-license = "allow" +confidence-threshold = 0.95 +allow = [ + # Currently used across our Rust projects + "Apache-2.0", + "BSD-2-Clause", + "BSD-3-Clause", + "BSL-1.0", + "CC0-1.0", + "ISC", + "LGPL-2.1-or-later", + "MIT", + "MIT-0", + "OpenSSL", + "Unicode-3.0", + "Unicode-DFS-2016", + "Unlicense", + "Zlib", + # Common safe licenses in the Rust ecosystem + "0BSD", + "Apache-2.0 WITH LLVM-exception", + "Unicode-DFS-2015", +] diff --git a/plugins_rust/retry_with_backoff/pyproject.toml b/plugins_rust/retry_with_backoff/pyproject.toml new file mode 100644 index 00000000000..c5c835571cc --- /dev/null +++ b/plugins_rust/retry_with_backoff/pyproject.toml @@ -0,0 +1,16 @@ +[build-system] +requires = ["maturin>=1.4,<2.0"] +build-backend = "maturin" + +[project] +name = "mcpgateway-retry-with-backoff" +version = "0.1.0" +description = "Rust-accelerated retry state manager for ContextForge" +license = { text = "Apache-2.0" } +requires-python = ">=3.11" + +[tool.maturin] +# The Python module name: `import retry_with_backoff_rust` +module-name = "retry_with_backoff_rust" +python-source = "python" +features = ["pyo3/extension-module"] diff --git a/plugins_rust/retry_with_backoff/python/retry_with_backoff_rust/__init__.pyi b/plugins_rust/retry_with_backoff/python/retry_with_backoff_rust/__init__.pyi new file mode 100644 index 00000000000..7eae370797e --- /dev/null +++ b/plugins_rust/retry_with_backoff/python/retry_with_backoff_rust/__init__.pyi @@ -0,0 +1,22 @@ +# This file is automatically generated by pyo3_stub_gen +# ruff: noqa: E501, F401, F403, F405 + +import builtins +import typing + +__all__ = [ + "RetryStateManager", +] + +@typing.final +class RetryStateManager: + def __new__(cls, max_retries: builtins.int, base_ms: builtins.int, max_ms: builtins.int, jitter: builtins.bool, retry_on_status: typing.Sequence[builtins.int]) -> RetryStateManager: ... + def ping(self) -> builtins.str: ... + def get_failures(self, tool: builtins.str, request_id: builtins.str) -> builtins.int: ... + def record_failure(self, tool: builtins.str, request_id: builtins.str) -> builtins.int: ... + def record_success(self, tool: builtins.str, request_id: builtins.str) -> None: ... + def delete_state(self, tool: builtins.str, request_id: builtins.str) -> None: ... + def state_count(self) -> builtins.int: ... + def compute_delay(self, attempt: builtins.int) -> builtins.int: ... + def check_failure(self, is_error: builtins.bool, status_code: typing.Optional[builtins.int]) -> builtins.bool: ... + def check_and_update(self, tool: builtins.str, request_id: builtins.str, is_error: builtins.bool, status_code: typing.Optional[builtins.int]) -> tuple[builtins.bool, builtins.int]: ... diff --git a/plugins_rust/retry_with_backoff/src/bin/stub_gen.rs b/plugins_rust/retry_with_backoff/src/bin/stub_gen.rs new file mode 100644 index 00000000000..c5cf230bc45 --- /dev/null +++ b/plugins_rust/retry_with_backoff/src/bin/stub_gen.rs @@ -0,0 +1,15 @@ +// Copyright 2025 +// SPDX-License-Identifier: Apache-2.0 +// +// Stub file generator for retry_with_backoff module +// +// This binary generates Python type stub files (.pyi) for the retry_with_backoff module. +// Run with: cargo run --bin stub_gen + +use retry_with_backoff_rust::stub_info; + +fn main() { + let stub_info = stub_info().expect("Failed to get stub info"); + stub_info.generate().expect("Failed to generate stub file"); + println!("✓ Generated stub files successfully"); +} diff --git a/plugins_rust/retry_with_backoff/src/lib.rs b/plugins_rust/retry_with_backoff/src/lib.rs new file mode 100644 index 00000000000..53caf132b74 --- /dev/null +++ b/plugins_rust/retry_with_backoff/src/lib.rs @@ -0,0 +1,532 @@ +// Copyright 2025 +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{HashMap, HashSet}; +use std::sync::{Mutex, OnceLock}; +use std::time::Instant; + +use log::{debug, warn}; +use pyo3::prelude::*; +use pyo3_stub_gen::define_stub_info_gatherer; +use pyo3_stub_gen::derive::*; +use rand::Rng; + +// --------------------------------------------------------------------------- +// State struct — mirrors Python's _ToolRetryState dataclass. +// --------------------------------------------------------------------------- +pub struct ToolRetryState { + pub consecutive_failures: u32, + pub last_failure_at: f64, +} + +impl ToolRetryState { + fn new() -> Self { + ToolRetryState { + consecutive_failures: 0, + last_failure_at: 0.0, + } + } +} + +// --------------------------------------------------------------------------- +// Global state — mirrors Python's module-level _STATE dict. +// Mutex protects concurrent access; OnceLock ensures single initialisation. +// +// NOTE: This map is process-global and shared across all RetryStateManager +// instances. The gateway creates a single plugin instance so this is fine +// in practice, but if multiple instances with different configs are ever +// constructed they will share state entries. +// --------------------------------------------------------------------------- +static STATE: OnceLock>> = OnceLock::new(); + +/// Monotonic reference point — equivalent of Python's `time.monotonic()` epoch. +/// Using `Instant` instead of `SystemTime` ensures TTL eviction is immune to +/// wall-clock jumps from NTP sync or manual adjustment. +static MONO_EPOCH: OnceLock = OnceLock::new(); + +/// Entries older than this (in seconds) are considered orphaned — e.g. the +/// retry sleep was cancelled by a client disconnect — and are evicted. +const STATE_TTL_SECS: f64 = 300.0; + +/// Return a monotonic timestamp in seconds, analogous to Python's `time.monotonic()`. +fn monotonic_secs() -> f64 { + let epoch = MONO_EPOCH.get_or_init(Instant::now); + epoch.elapsed().as_secs_f64() +} + +fn state_map() -> &'static Mutex> { + STATE.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// Remove entries whose `last_failure_at` is older than `STATE_TTL_SECS`. +/// Called under an already-held lock by `check_and_update`. +fn evict_stale(map: &mut HashMap) { + let cutoff = monotonic_secs() - STATE_TTL_SECS; + map.retain(|_, v| v.last_failure_at <= 0.0 || v.last_failure_at >= cutoff); +} + +fn make_key(tool: &str, request_id: &str) -> String { + format!("{tool}:{request_id}") +} + +// --------------------------------------------------------------------------- +// Pure functions — no PyO3 types. Called from RetryStateManager methods. +// --------------------------------------------------------------------------- + +// Exponential backoff with optional jitter, capped at max_ms. +fn compute_delay_ms(attempt: u32, base_ms: u64, max_ms: u64, jitter: bool) -> u64 { + let ceiling = base_ms + .saturating_mul(2u64.saturating_pow(attempt)) + .min(max_ms); + if jitter { + rand::thread_rng().gen_range(0..=ceiling) + } else { + ceiling + } +} + +// Checks the two pre-extracted failure signals: outer isError flag and status code. +// Text-content parsing (signal 3) is handled entirely in Python. +// +// When isError is true AND a status code is present (e.g. the gateway extracted +// it from an httpx.HTTPStatusError), the code is checked against retry_on_status +// so that non-transient HTTP errors (400, 401, 404 …) are not wastefully retried. +// Generic exceptions with no status code (connection errors, timeouts) are always +// considered retryable. +fn is_failure_from_signals( + is_error: bool, + status_code: Option, + retry_on_status: &HashSet, +) -> bool { + if is_error { + return match status_code { + Some(sc) => retry_on_status.contains(&sc), + None => true, + }; + } + if let Some(sc) = status_code { + return retry_on_status.contains(&sc); + } + false +} + +// --------------------------------------------------------------------------- +// Python-visible class. +// Config is stored in the struct — set once at construction, never +// re-allocated on the hot path. retry_on_status is kept as a HashSet for +// O(1) membership tests instead of the O(n) Vec scan it replaced. +// --------------------------------------------------------------------------- +#[gen_stub_pyclass] +#[pyclass] +pub struct RetryStateManager { + max_retries: u32, + base_ms: u64, + max_ms: u64, + jitter: bool, + retry_on_status: HashSet, +} + +#[gen_stub_pymethods] +#[pymethods] +impl RetryStateManager { + #[new] + fn new( + max_retries: u32, + base_ms: u64, + max_ms: u64, + jitter: bool, + retry_on_status: Vec, + ) -> Self { + debug!( + "RetryStateManager created: max_retries={max_retries} base_ms={base_ms} \ + max_ms={max_ms} jitter={jitter} retry_on_status={retry_on_status:?}" + ); + RetryStateManager { + max_retries, + base_ms, + max_ms, + jitter, + // Vec → HashSet: one-time allocation at construction; O(1) lookups thereafter. + retry_on_status: retry_on_status.into_iter().collect(), + } + } + + fn ping(&self) -> &str { + "retry_with_backoff_rust is alive" + } + + // Returns consecutive_failures for (tool, request_id), or 0 if absent. + fn get_failures(&self, tool: &str, request_id: &str) -> u32 { + let map = state_map().lock().unwrap(); + let key = make_key(tool, request_id); + map.get(&key).map(|s| s.consecutive_failures).unwrap_or(0) + } + + // Increments consecutive_failures and records the current timestamp. + fn record_failure(&self, tool: &str, request_id: &str) -> u32 { + let mut map = state_map().lock().unwrap(); + let key = make_key(tool, request_id); + let state = map.entry(key).or_insert_with(ToolRetryState::new); + state.consecutive_failures += 1; + state.last_failure_at = monotonic_secs(); + debug!( + "record_failure: tool={tool} request_id={request_id} consecutive_failures={}", + state.consecutive_failures + ); + state.consecutive_failures + } + + // Resets consecutive_failures to 0 without removing the entry. + fn record_success(&self, tool: &str, request_id: &str) { + let mut map = state_map().lock().unwrap(); + let key = make_key(tool, request_id); + if let Some(state) = map.get_mut(&key) { + state.consecutive_failures = 0; + debug!("record_success: tool={tool} request_id={request_id} — failure counter reset"); + } + } + + // Removes the state entry for a completed invocation (success or exhausted). + fn delete_state(&self, tool: &str, request_id: &str) { + let mut map = state_map().lock().unwrap(); + let key = make_key(tool, request_id); + let _ = map.remove(&key); + } + + // Number of active (tool, request_id) entries — useful for tests/debugging. + fn state_count(&self) -> usize { + state_map().lock().unwrap().len() + } + + fn compute_delay(&self, attempt: u32) -> u64 { + compute_delay_ms(attempt, self.base_ms, self.max_ms, self.jitter) + } + + fn check_failure(&self, is_error: bool, status_code: Option) -> bool { + is_failure_from_signals(is_error, status_code, &self.retry_on_status) + } + + // ----------------------------------------------------------------------- + // Main API called by the Python plugin on every post-invoke hook. + // + // Config (max_retries, base_ms, max_ms, jitter, retry_on_status) lives in + // self — no per-call allocations or list conversions cross the FFI boundary. + // Only the four truly dynamic arguments are passed. + // + // Returns (should_retry, delay_ms): + // (true, delay) — failure within budget; caller should schedule retry + // (false, 0) — success OR retries exhausted; caller propagates result + // + // The Mutex is held for the entire method to make the check-then-act + // sequence atomic. + // ----------------------------------------------------------------------- + fn check_and_update( + &self, + tool: &str, + request_id: &str, + is_error: bool, + status_code: Option, + ) -> (bool, u64) { + let failed = is_failure_from_signals(is_error, status_code, &self.retry_on_status); + + // Acquire the lock once for the entire check-then-act sequence. + let mut map = state_map().lock().unwrap(); + evict_stale(&mut map); + let key = make_key(tool, request_id); + + if failed { + let state = map.entry(key.clone()).or_insert_with(ToolRetryState::new); + state.consecutive_failures += 1; + state.last_failure_at = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs_f64(); + + if state.consecutive_failures <= self.max_retries { + // attempt index is 0-based; saturating_sub guards against underflow. + let attempt = state.consecutive_failures.saturating_sub(1); + let delay = compute_delay_ms(attempt, self.base_ms, self.max_ms, self.jitter); + debug!( + "check_and_update: tool={tool} request_id={request_id} \ + failure={}/{} — retry in {delay}ms", + state.consecutive_failures, self.max_retries + ); + (true, delay) + } else { + warn!( + "check_and_update: tool={tool} request_id={request_id} \ + retry budget exhausted ({} failures) — propagating error", + state.consecutive_failures + ); + map.remove(&key); + (false, 0) + } + } else { + debug!( + "check_and_update: tool={tool} request_id={request_id} — success, state cleared" + ); + let _ = map.remove(&key); + (false, 0) + } + } +} + +// --------------------------------------------------------------------------- +// Module entry point. +// --------------------------------------------------------------------------- +#[pymodule] +fn retry_with_backoff_rust(m: &Bound<'_, PyModule>) -> PyResult<()> { + // Bridge Rust log records to Python's logging module. + // After this call, log::debug!/warn!/etc. emit to the logger named + // "retry_with_backoff_rust" in Python, respecting the Python log level. + pyo3_log::init(); + m.add_class::()?; + Ok(()) +} + +define_stub_info_gatherer!(stub_info); + +// --------------------------------------------------------------------------- +// Unit tests — pure Rust, no Python interpreter required. +// Run with: cargo test (or: make test) +// --------------------------------------------------------------------------- +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; + + fn status_set(codes: &[i32]) -> HashSet { + codes.iter().cloned().collect() + } + + fn manager_with(base_ms: u64, max_ms: u64) -> RetryStateManager { + RetryStateManager::new(2, base_ms, max_ms, false, vec![500, 503]) + } + + // ── compute_delay_ms ──────────────────────────────────────────────────── + + #[test] + fn delay_attempt_zero_returns_base() { + assert_eq!(compute_delay_ms(0, 100, 10_000, false), 100); + } + + #[test] + fn delay_doubles_each_attempt() { + assert_eq!(compute_delay_ms(1, 100, 10_000, false), 200); + assert_eq!(compute_delay_ms(2, 100, 10_000, false), 400); + assert_eq!(compute_delay_ms(3, 100, 10_000, false), 800); + } + + #[test] + fn delay_is_capped_at_max_ms() { + assert_eq!(compute_delay_ms(10, 100, 500, false), 500); + } + + #[test] + fn delay_no_overflow_on_extreme_attempt() { + // saturating_pow / saturating_mul must not panic — result must equal cap. + let d = compute_delay_ms(63, 100, 5_000, false); + assert_eq!(d, 5_000, "expected cap, got {d}"); + } + + // ── is_failure_from_signals ───────────────────────────────────────────── + + #[test] + fn failure_when_is_error_true() { + assert!(is_failure_from_signals(true, None, &status_set(&[]))); + } + + #[test] + fn no_failure_when_is_error_false_and_no_status() { + assert!(!is_failure_from_signals(false, None, &status_set(&[]))); + } + + #[test] + fn failure_when_status_code_in_retry_set() { + assert!(is_failure_from_signals( + false, + Some(500), + &status_set(&[500, 503]) + )); + assert!(is_failure_from_signals( + false, + Some(503), + &status_set(&[500, 503]) + )); + } + + #[test] + fn no_failure_when_status_code_not_in_retry_set() { + assert!(!is_failure_from_signals( + false, + Some(200), + &status_set(&[500, 503]) + )); + assert!(!is_failure_from_signals( + false, + Some(404), + &status_set(&[500, 503]) + )); + } + + #[test] + fn is_error_with_non_retryable_status_does_not_retry() { + // isError=true + status_code NOT in retry set → not retryable. + assert!(!is_failure_from_signals(true, Some(200), &status_set(&[]))); + assert!(!is_failure_from_signals( + true, + Some(400), + &status_set(&[500, 503]) + )); + assert!(!is_failure_from_signals( + true, + Some(404), + &status_set(&[500, 503]) + )); + } + + #[test] + fn is_error_with_retryable_status_retries() { + assert!(is_failure_from_signals( + true, + Some(500), + &status_set(&[500, 503]) + )); + assert!(is_failure_from_signals( + true, + Some(503), + &status_set(&[500, 503]) + )); + } + + #[test] + fn is_error_without_status_always_retries() { + // isError=true with no status code → generic exception → always retry. + assert!(is_failure_from_signals(true, None, &status_set(&[]))); + assert!(is_failure_from_signals(true, None, &status_set(&[500]))); + } + + // ── make_key ──────────────────────────────────────────────────────────── + + #[test] + fn key_format_is_tool_colon_request() { + assert_eq!(make_key("my_tool", "req-123"), "my_tool:req-123"); + } + + #[test] + fn key_with_empty_parts() { + assert_eq!(make_key("", ""), ":"); + } + + // ── RetryStateManager: state helpers ──────────────────────────────────── + + #[test] + fn get_failures_returns_zero_for_unknown_key() { + let m = manager_with(100, 10_000); + assert_eq!(m.get_failures("unknown_t", "unknown_r"), 0); + } + + #[test] + fn record_failure_increments_and_get_failures_reads_back() { + let m = manager_with(100, 10_000); + let (tool, req) = ("state_rf_t", "state_rf_r"); + m.delete_state(tool, req); // clean slate + assert_eq!(m.record_failure(tool, req), 1); + assert_eq!(m.record_failure(tool, req), 2); + assert_eq!(m.get_failures(tool, req), 2); + m.delete_state(tool, req); + } + + #[test] + fn record_success_resets_failure_counter_to_zero() { + let m = manager_with(100, 10_000); + let (tool, req) = ("state_rs_t", "state_rs_r"); + m.delete_state(tool, req); + m.record_failure(tool, req); + m.record_failure(tool, req); + m.record_success(tool, req); + assert_eq!(m.get_failures(tool, req), 0); + m.delete_state(tool, req); + } + + #[test] + fn delete_state_removes_entry() { + let m = manager_with(100, 10_000); + let (tool, req) = ("state_del_t", "state_del_r"); + m.record_failure(tool, req); + m.delete_state(tool, req); + assert_eq!(m.get_failures(tool, req), 0); + } + + // ── RetryStateManager: check_and_update ───────────────────────────────── + + #[test] + fn check_and_update_success_returns_no_retry() { + let m = manager_with(100, 10_000); + let (retry, delay) = m.check_and_update("cau_ok_t", "cau_ok_r", false, None); + assert!(!retry); + assert_eq!(delay, 0); + } + + #[test] + fn check_and_update_first_failure_triggers_retry_with_base_delay() { + let m = manager_with(100, 10_000); + let (tool, req) = ("cau_f1_t", "cau_f1_r"); + m.delete_state(tool, req); + let (retry, delay) = m.check_and_update(tool, req, true, None); + assert!(retry, "expected retry on first failure"); + assert_eq!(delay, 100, "first failure should use base_ms (attempt 0)"); + m.delete_state(tool, req); + } + + #[test] + fn check_and_update_status_code_match_triggers_retry() { + let m = manager_with(100, 10_000); + let (tool, req) = ("cau_sc_t", "cau_sc_r"); + m.delete_state(tool, req); + let (retry, _) = m.check_and_update(tool, req, false, Some(500)); + assert!(retry, "status 500 should trigger retry"); + m.delete_state(tool, req); + } + + #[test] + fn check_and_update_delay_doubles_with_successive_failures() { + let m = RetryStateManager::new(2, 100, 10_000, false, vec![]); + let (tool, req) = ("cau_exp_t", "cau_exp_r"); + m.delete_state(tool, req); + let (_, d1) = m.check_and_update(tool, req, true, None); // attempt 0 → 100 + let (_, d2) = m.check_and_update(tool, req, true, None); // attempt 1 → 200 + assert_eq!(d1, 100); + assert_eq!(d2, 200); + m.delete_state(tool, req); + } + + #[test] + fn check_and_update_exhausts_budget_then_stops() { + let m = manager_with(100, 10_000); // max_retries = 2 + let (tool, req) = ("cau_ex_t", "cau_ex_r"); + m.delete_state(tool, req); + // First 2 failures are within budget. + for i in 1..=2 { + let (retry, _) = m.check_and_update(tool, req, true, None); + assert!(retry, "failure {i} should still be within retry budget"); + } + // 3rd failure exceeds max_retries — entry is removed, no retry. + let (retry, delay) = m.check_and_update(tool, req, true, None); + assert!(!retry, "retry budget exhausted — should not retry"); + assert_eq!(delay, 0); + } + + #[test] + fn check_and_update_success_after_failures_clears_state() { + let m = manager_with(100, 10_000); + let (tool, req) = ("cau_clr_t", "cau_clr_r"); + m.delete_state(tool, req); + m.check_and_update(tool, req, true, None); // record a failure + let (retry, delay) = m.check_and_update(tool, req, false, None); // success + assert!(!retry); + assert_eq!(delay, 0); + // Entry should be gone — subsequent get_failures returns 0. + assert_eq!(m.get_failures(tool, req), 0); + } +} diff --git a/plugins_rust/secrets_detection/python/secrets_detection_rust/__init__.pyi b/plugins_rust/secrets_detection/python/secrets_detection_rust/__init__.pyi index a8edcaa3a22..693236e1fc9 100644 --- a/plugins_rust/secrets_detection/python/secrets_detection_rust/__init__.pyi +++ b/plugins_rust/secrets_detection/python/secrets_detection_rust/__init__.pyi @@ -3,6 +3,7 @@ import builtins import typing + __all__ = [ "py_scan_container", ] diff --git a/tests/performance/plugins/config.yaml b/tests/performance/plugins/config.yaml index 81fecef8151..0c12ba91671 100644 --- a/tests/performance/plugins/config.yaml +++ b/tests/performance/plugins/config.yaml @@ -351,14 +351,14 @@ plugins: - application/json allowed_extensions: [.md, .txt, .html, .json] - # Retry policy annotations + # Retry with exponential backoff — triggers real gateway re-invocation - name: RetryWithBackoffPlugin kind: plugins.retry_with_backoff.retry_with_backoff.RetryWithBackoffPlugin - description: Annotates retry/backoff policy in metadata + description: Detects transient failures and asks the gateway to re-invoke the tool after a jittered exponential backoff delay version: 0.1.0 author: Mihai Criveti - hooks: [tool_post_invoke, resource_post_fetch] - tags: [reliability, retry] + hooks: [tool_post_invoke] + tags: [retry, backoff, resilience] mode: permissive priority: 170 conditions: [] @@ -367,6 +367,9 @@ plugins: backoff_base_ms: 200 max_backoff_ms: 5000 retry_on_status: [429, 500, 502, 503, 504] + jitter: true + check_text_content: false + tool_overrides: {} # Markdown cleaner - name: MarkdownCleanerPlugin diff --git a/tests/unit/mcpgateway/services/test_tool_service.py b/tests/unit/mcpgateway/services/test_tool_service.py index 9b45e8a5378..9a5889fd8bf 100644 --- a/tests/unit/mcpgateway/services/test_tool_service.py +++ b/tests/unit/mcpgateway/services/test_tool_service.py @@ -31,6 +31,7 @@ from mcpgateway.db import Tool as DbTool from mcpgateway.plugins.framework import PluginManager from mcpgateway.plugins.framework.hooks.tools import ToolHookType +from mcpgateway.plugins.framework.models import PluginResult from mcpgateway.schemas import AuthenticationValues, ToolCreate, ToolRead, ToolUpdate from mcpgateway.services.tool_service import ( _decrypt_tool_header_value, @@ -47,6 +48,7 @@ ToolNotFoundError, ToolResult, ToolService, + ToolTimeoutError, ToolValidationError, ) from mcpgateway.utils.pagination import decode_cursor @@ -1381,11 +1383,7 @@ async def test_list_server_tools_with_include_metrics_true(self): service.convert_tool_to_read = Mock(return_value="converted_tool_with_metrics") # Call with include_metrics=True to trigger eager loading code path - tools = await service.list_server_tools( - mock_db, - server_id="server123", - include_metrics=True - ) + tools = await service.list_server_tools(mock_db, server_id="server123", include_metrics=True) assert tools == ["converted_tool_with_metrics"] # Verify convert_tool_to_read was called with include_metrics=True @@ -1395,7 +1393,7 @@ async def test_list_server_tools_with_include_metrics_true(self): include_auth=False, requesting_user_email=None, requesting_user_is_admin=False, - requesting_user_team_roles=None + requesting_user_team_roles=None, ) @pytest.mark.asyncio @@ -3514,6 +3512,7 @@ async def test_invoke_tool_with_plugin_post_invoke_success(self, tool_service, m mock_post_result.continue_processing = True mock_post_result.violation = None mock_post_result.modified_payload = None + mock_post_result.retry_delay_ms = 0 tool_service._plugin_manager = Mock() @@ -3562,6 +3561,7 @@ async def test_invoke_tool_with_plugin_post_invoke_modified_payload(self, tool_s mock_post_result.continue_processing = True mock_post_result.violation = None mock_post_result.modified_payload = mock_modified_payload + mock_post_result.retry_delay_ms = 0 # First-Party from mcpgateway.plugins.framework import PluginResult, ToolHookType @@ -3613,6 +3613,7 @@ async def test_invoke_tool_with_plugin_post_invoke_invalid_modified_payload(self mock_post_result.continue_processing = True mock_post_result.violation = None mock_post_result.modified_payload = mock_modified_payload + mock_post_result.retry_delay_ms = 0 # First-Party from mcpgateway.plugins.framework import ToolHookType @@ -3776,6 +3777,248 @@ async def test_invoke_tool_with_plugin_metadata_sse(self, tool_service, mock_too await tool_service._plugin_manager.shutdown() + @pytest.mark.asyncio + async def test_invoke_tool_plugin_retry_delay_triggers_retry(self, tool_service, mock_tool, mock_global_config_obj, test_db): + """Test that when plugin returns retry_delay_ms > 0 the tool is retried after the delay.""" + mock_tool.integration_type = "REST" + mock_tool.request_type = "POST" + mock_tool.auth_value = None + + setup_db_execute_mock(test_db, mock_tool, mock_global_config_obj) + + mock_response = AsyncMock() + mock_response.raise_for_status = Mock() + mock_response.status_code = 200 + mock_response.json = Mock(return_value={"result": "first attempt"}) + tool_service._http_client.request.return_value = mock_response + + tool_service._plugin_manager = Mock() + + post_invoke_count = [0] + + def invoke_hook_side_effect(hook_type, payload, global_context, local_contexts=None, **kwargs): + if hook_type == ToolHookType.TOOL_PRE_INVOKE: + return (PluginResult(continue_processing=True, violation=None, modified_payload=None), None) + # POST_INVOKE: first call requests a retry, subsequent calls succeed + post_invoke_count[0] += 1 + delay = 100 if post_invoke_count[0] == 1 else 0 + return (PluginResult(continue_processing=True, violation=None, modified_payload=None, retry_delay_ms=delay), None) + + tool_service._plugin_manager.invoke_hook = AsyncMock(side_effect=invoke_hook_side_effect) + + # Intercept the recursive retry call so it returns immediately on retry_attempt > 0 + retry_result = ToolResult(content=[TextContent(type="text", text="retry succeeded")]) + original_invoke = tool_service.invoke_tool + + async def spy_invoke(db, name, arguments=None, **kwargs): + if kwargs.get("retry_attempt", 0) > 0: + return retry_result + return await original_invoke(db, name, arguments, **kwargs) + + with ( + patch("mcpgateway.services.tool_service.decode_auth", return_value={}), + patch("mcpgateway.services.tool_service.extract_using_jq", return_value={"result": "first attempt"}), + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + patch.object(tool_service, "invoke_tool", side_effect=spy_invoke), + ): + result = await tool_service.invoke_tool(test_db, "test_tool", {"param": "value"}, request_headers=None) + + # Verify the retry delay sleep was triggered + mock_sleep.assert_awaited_once() + assert mock_sleep.call_args[0][0] == pytest.approx(0.1) # 100ms → 0.1s + # Verify the result is from the retry call + assert result.content[0].text == "retry succeeded" + + @pytest.mark.asyncio + async def test_invoke_tool_exception_path_retry_fires(self, tool_service, mock_tool, mock_global_config_obj, test_db): + """When a tool raises an exception the retry plugin's delay_ms must also trigger a retry.""" + mock_tool.integration_type = "REST" + mock_tool.request_type = "POST" + mock_tool.auth_value = None + + setup_db_execute_mock(test_db, mock_tool, mock_global_config_obj) + + # HTTP client raises an exception (simulates network error) + tool_service._http_client.request.side_effect = RuntimeError("connection reset") + + tool_service._plugin_manager = Mock() + tool_service._plugin_manager.has_hooks_for.return_value = True + + # Plugin returns delay > 0 on first post-invoke, 0 on subsequent + post_invoke_count = [0] + + def invoke_hook_side_effect(hook_type, payload, global_context, local_contexts=None, **kwargs): + if hook_type == ToolHookType.TOOL_PRE_INVOKE: + return (PluginResult(continue_processing=True, violation=None, modified_payload=None), None) + post_invoke_count[0] += 1 + delay = 50 if post_invoke_count[0] == 1 else 0 + return (PluginResult(continue_processing=True, violation=None, modified_payload=None, retry_delay_ms=delay), None) + + tool_service._plugin_manager.invoke_hook = AsyncMock(side_effect=invoke_hook_side_effect) + + retry_result = ToolResult(content=[TextContent(type="text", text="recovered after exception")]) + original_invoke = tool_service.invoke_tool + + async def spy_invoke(db, name, arguments=None, **kwargs): + if kwargs.get("retry_attempt", 0) > 0: + return retry_result + return await original_invoke(db, name, arguments, **kwargs) + + with ( + patch("mcpgateway.services.tool_service.decode_auth", return_value={}), + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + patch.object(tool_service, "invoke_tool", side_effect=spy_invoke), + ): + result = await tool_service.invoke_tool(test_db, "test_tool", {"param": "value"}, request_headers=None) + + mock_sleep.assert_awaited_once() + assert mock_sleep.call_args[0][0] == pytest.approx(0.05) # 50ms → 0.05s + assert result.content[0].text == "recovered after exception" + + @pytest.mark.asyncio + async def test_invoke_tool_http_status_error_passes_status_to_plugin(self, tool_service, mock_tool, mock_global_config_obj, test_db): + """When raise_for_status() raises HTTPStatusError the status code must be forwarded to the plugin in structuredContent.""" + import httpx + + mock_tool.integration_type = "REST" + mock_tool.request_type = "POST" + mock_tool.auth_value = None + + setup_db_execute_mock(test_db, mock_tool, mock_global_config_obj) + + # Simulate httpx.HTTPStatusError (e.g. a 404 from raise_for_status) + mock_response = MagicMock() + mock_response.status_code = 404 + mock_request = MagicMock() + tool_service._http_client.request.side_effect = httpx.HTTPStatusError("Not Found", request=mock_request, response=mock_response) + + tool_service._plugin_manager = Mock() + tool_service._plugin_manager.has_hooks_for.return_value = True + + captured_payloads = [] + + def invoke_hook_side_effect(hook_type, payload, global_context, local_contexts=None, **kwargs): + if hook_type == ToolHookType.TOOL_POST_INVOKE: + captured_payloads.append(payload) + return (PluginResult(continue_processing=True, violation=None, modified_payload=None, retry_delay_ms=0), None) + + tool_service._plugin_manager.invoke_hook = AsyncMock(side_effect=invoke_hook_side_effect) + + with patch("mcpgateway.services.tool_service.decode_auth", return_value={}): + with pytest.raises(ToolInvocationError): + await tool_service.invoke_tool(test_db, "test_tool", {"param": "value"}, request_headers=None) + + # The exception-path post-invoke payload must include the HTTP status code in structuredContent + assert len(captured_payloads) >= 1 + post_invoke_result = captured_payloads[-1].result + assert post_invoke_result["isError"] is True + assert post_invoke_result["structuredContent"] == {"status_code": 404} + + @pytest.mark.asyncio + async def test_invoke_tool_timeout_path_retry_fires(self, tool_service, mock_tool, mock_global_config_obj, test_db): + """When a tool times out and the retry plugin requests a retry via ToolTimeoutError.retry_delay_ms, the gateway must retry.""" + mock_tool.integration_type = "REST" + mock_tool.request_type = "POST" + mock_tool.auth_value = None + + setup_db_execute_mock(test_db, mock_tool, mock_global_config_obj) + + # HTTP client raises ToolTimeoutError with retry_delay_ms to simulate + # the timeout handler having already called _run_timeout_post_invoke. + tool_service._http_client.request.side_effect = ToolTimeoutError("timed out after 30s", retry_delay_ms=75) + + tool_service._plugin_manager = Mock() + tool_service._plugin_manager.has_hooks_for.return_value = True + + def invoke_hook_side_effect(hook_type, payload, global_context, local_contexts=None, **kwargs): + if hook_type == ToolHookType.TOOL_PRE_INVOKE: + return (PluginResult(continue_processing=True, violation=None, modified_payload=None), None) + return (PluginResult(continue_processing=True, violation=None, modified_payload=None, retry_delay_ms=0), None) + + tool_service._plugin_manager.invoke_hook = AsyncMock(side_effect=invoke_hook_side_effect) + + retry_result = ToolResult(content=[TextContent(type="text", text="recovered after timeout")]) + original_invoke = tool_service.invoke_tool + + async def spy_invoke(db, name, arguments=None, **kwargs): + if kwargs.get("retry_attempt", 0) > 0: + return retry_result + return await original_invoke(db, name, arguments, **kwargs) + + with ( + patch("mcpgateway.services.tool_service.decode_auth", return_value={}), + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + patch.object(tool_service, "invoke_tool", side_effect=spy_invoke), + ): + result = await tool_service.invoke_tool(test_db, "test_tool", {"param": "value"}, request_headers=None) + + mock_sleep.assert_awaited_once() + assert mock_sleep.call_args[0][0] == pytest.approx(0.075) # 75ms → 0.075s + assert result.content[0].text == "recovered after timeout" + + @pytest.mark.asyncio + async def test_invoke_tool_timeout_no_retry_reraises(self, tool_service, mock_tool, mock_global_config_obj, test_db): + """ToolTimeoutError with retry_delay_ms=0 must re-raise without retrying.""" + mock_tool.integration_type = "REST" + mock_tool.request_type = "POST" + mock_tool.auth_value = None + + setup_db_execute_mock(test_db, mock_tool, mock_global_config_obj) + + tool_service._http_client.request.side_effect = ToolTimeoutError("timed out after 30s") + + tool_service._plugin_manager = Mock() + tool_service._plugin_manager.has_hooks_for.return_value = True + + def invoke_hook_side_effect(hook_type, payload, global_context, local_contexts=None, **kwargs): + if hook_type == ToolHookType.TOOL_PRE_INVOKE: + return (PluginResult(continue_processing=True, violation=None, modified_payload=None), None) + return (PluginResult(continue_processing=True, violation=None, modified_payload=None, retry_delay_ms=0), None) + + tool_service._plugin_manager.invoke_hook = AsyncMock(side_effect=invoke_hook_side_effect) + + with ( + patch("mcpgateway.services.tool_service.decode_auth", return_value={}), + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + ): + with pytest.raises(ToolTimeoutError): + await tool_service.invoke_tool(test_db, "test_tool", {"param": "value"}, request_headers=None) + + mock_sleep.assert_not_awaited() + + @pytest.mark.asyncio + async def test_run_timeout_post_invoke_calls_hook(self, tool_service): + """_run_timeout_post_invoke must invoke the post-invoke hook when the plugin manager has hooks.""" + tool_service._plugin_manager = Mock() + tool_service._plugin_manager.has_hooks_for.return_value = True + tool_service._plugin_manager.invoke_hook = AsyncMock(return_value=(PluginResult(retry_delay_ms=0), None)) + + await tool_service._run_timeout_post_invoke("test_tool", 30.0, None, None) + + tool_service._plugin_manager.invoke_hook.assert_awaited_once() + call_kwargs = tool_service._plugin_manager.invoke_hook.call_args + assert call_kwargs[1]["payload"].name == "test_tool" + assert call_kwargs[1]["payload"].result["isError"] is True + + @pytest.mark.asyncio + async def test_run_timeout_post_invoke_raises_on_retry_signal(self, tool_service): + """_run_timeout_post_invoke must raise ToolTimeoutError with retry_delay_ms when the plugin requests a retry.""" + tool_service._plugin_manager = Mock() + tool_service._plugin_manager.has_hooks_for.return_value = True + tool_service._plugin_manager.invoke_hook = AsyncMock(return_value=(PluginResult(retry_delay_ms=250), None)) + + with pytest.raises(ToolTimeoutError) as exc_info: + await tool_service._run_timeout_post_invoke("test_tool", 30.0, None, None) + + assert exc_info.value.retry_delay_ms == 250 + + @pytest.mark.asyncio + async def test_run_timeout_post_invoke_noop_without_plugin_manager(self, tool_service): + """_run_timeout_post_invoke must return immediately when plugin_manager is None.""" + tool_service._plugin_manager = None + # Should not raise + await tool_service._run_timeout_post_invoke("test_tool", 30.0, None, None) + # --------------------------------------------------------------------------- # # extract_using_jq # diff --git a/tests/unit/mcpgateway/services/test_tool_service_coverage.py b/tests/unit/mcpgateway/services/test_tool_service_coverage.py index 535b4f1bf8e..bde3295acbf 100644 --- a/tests/unit/mcpgateway/services/test_tool_service_coverage.py +++ b/tests/unit/mcpgateway/services/test_tool_service_coverage.py @@ -4804,7 +4804,9 @@ def _make_tool_payload( } -def _make_gateway_payload(*, auth_type=None, auth_value=None, auth_query_params=None, oauth_config=None, ca_certificate=None, ca_certificate_sig=None, passthrough_headers=None, url=None, client_cert=None, client_key=None): +def _make_gateway_payload( + *, auth_type=None, auth_value=None, auth_query_params=None, oauth_config=None, ca_certificate=None, ca_certificate_sig=None, passthrough_headers=None, url=None, client_cert=None, client_key=None +): return { "id": "gw-uuid-1", "name": "test_gw", @@ -4915,7 +4917,7 @@ async def test_rest_timeout_triggers_cb_and_post_hook_and_metrics_counter_failur plugin_manager.invoke_hook = AsyncMock( side_effect=[ (SimpleNamespace(modified_payload=None), context_table), # pre-invoke - (SimpleNamespace(modified_payload=None), context_table), # post-invoke (timeout handler) + (SimpleNamespace(modified_payload=None, retry_delay_ms=0), context_table), # post-invoke (timeout handler) ] ) tool_service._plugin_manager = plugin_manager @@ -5827,9 +5829,6 @@ class TestInvokeToolPluginPostInvokeSerialization: @pytest.mark.asyncio async def test_plugin_post_invoke_dict_result_serialized_as_json(self, tool_service): """When plugin post-invoke returns a dict without 'content' key, it should be serialized as valid JSON.""" - # First-Party - from mcpgateway.plugins.framework import ToolHookType - tp = _make_tool_payload(integration_type="REST", request_type="GET") db = MagicMock() @@ -5846,8 +5845,8 @@ async def fake_get(*a, **kw): plugin_manager.has_hooks_for = MagicMock(return_value=True) plugin_manager.invoke_hook = AsyncMock( side_effect=[ - (SimpleNamespace(modified_payload=None), {}), # pre-invoke - (SimpleNamespace(modified_payload=SimpleNamespace(result={"status": "transformed", "valid": False})), {}), # post-invoke + (SimpleNamespace(modified_payload=None, retry_delay_ms=0), {}), # pre-invoke + (SimpleNamespace(modified_payload=SimpleNamespace(result={"status": "transformed", "valid": False}), retry_delay_ms=0), {}), # post-invoke ] ) tool_service._plugin_manager = plugin_manager @@ -5886,9 +5885,6 @@ async def fake_get(*a, **kw): @pytest.mark.asyncio async def test_plugin_post_invoke_unserializable_result_falls_back_to_str(self, tool_service): """When plugin post-invoke returns an unserializable value (e.g. set), it should fall back to str() instead of crashing.""" - # First-Party - from mcpgateway.plugins.framework import ToolHookType - tp = _make_tool_payload(integration_type="REST", request_type="GET") db = MagicMock() @@ -5905,8 +5901,8 @@ async def fake_get(*a, **kw): plugin_manager.has_hooks_for = MagicMock(return_value=True) plugin_manager.invoke_hook = AsyncMock( side_effect=[ - (SimpleNamespace(modified_payload=None), {}), # pre-invoke - (SimpleNamespace(modified_payload=SimpleNamespace(result={"unserializable", "set", "values"})), {}), # post-invoke + (SimpleNamespace(modified_payload=None, retry_delay_ms=0), {}), # pre-invoke + (SimpleNamespace(modified_payload=SimpleNamespace(result={"unserializable", "set", "values"}), retry_delay_ms=0), {}), # post-invoke ] ) tool_service._plugin_manager = plugin_manager @@ -6684,7 +6680,7 @@ def _has_hooks_for(hook_type): return hook_type == ToolHookType.TOOL_POST_INVOKE plugin_manager.has_hooks_for = MagicMock(side_effect=_has_hooks_for) - plugin_manager.invoke_hook = AsyncMock(return_value=(SimpleNamespace(modified_payload=None), context_table)) + plugin_manager.invoke_hook = AsyncMock(return_value=(SimpleNamespace(modified_payload=None, retry_delay_ms=0), context_table)) tool_service._plugin_manager = plugin_manager with ( @@ -6858,6 +6854,7 @@ async def test_mcp_gateway_oauth_client_credentials_error_raises(self, tool_serv with pytest.raises(ToolInvocationError, match="OAuth authentication failed for gateway"): await tool_service.invoke_tool(db, "test_tool", {}) + @pytest.mark.asyncio async def test_mcp_http_url_bypasses_ssl_context_creation(self, tool_service): """HTTP URLs should skip SSL context creation entirely (line 3986-3988).""" @@ -6930,7 +6927,7 @@ async def test_mcp_https_url_with_mtls_creates_ssl_context(self, tool_service): auth_type="basic", ca_certificate="dummy-ca", client_cert="client-cert-data", # mTLS client cert - client_key="client-key-data", # mTLS client key + client_key="client-key-data", # mTLS client key ) db = MagicMock() @@ -7357,7 +7354,7 @@ def _has_hooks_for(hook_type): return hook_type == ToolHookType.TOOL_POST_INVOKE plugin_manager.has_hooks_for = MagicMock(side_effect=_has_hooks_for) - plugin_manager.invoke_hook = AsyncMock(return_value=(SimpleNamespace(modified_payload=None), context_table)) + plugin_manager.invoke_hook = AsyncMock(return_value=(SimpleNamespace(modified_payload=None, retry_delay_ms=0), context_table)) tool_service._plugin_manager = plugin_manager def fake_sse_client(*, url=None, headers=None, httpx_client_factory=None, **_kw): @@ -7612,7 +7609,7 @@ def _has_hooks_for(hook_type): return hook_type == ToolHookType.TOOL_POST_INVOKE plugin_manager.has_hooks_for = MagicMock(side_effect=_has_hooks_for) - plugin_manager.invoke_hook = AsyncMock(return_value=(SimpleNamespace(modified_payload=None), None)) + plugin_manager.invoke_hook = AsyncMock(return_value=(SimpleNamespace(modified_payload=None, retry_delay_ms=0), None)) tool_service._plugin_manager = plugin_manager def fake_streamablehttp_client(*, url=None, headers=None, httpx_client_factory=None, **_kw): diff --git a/tests/unit/mcpgateway/test_coverage_push.py b/tests/unit/mcpgateway/test_coverage_push.py index f3bc7135b46..dae963d74e3 100644 --- a/tests/unit/mcpgateway/test_coverage_push.py +++ b/tests/unit/mcpgateway/test_coverage_push.py @@ -93,6 +93,28 @@ def test_mcp_path_rewrite_middleware(): assert middleware.application == app_mock +@pytest.mark.asyncio +async def test_mcp_path_rewrite_non_servers_passthrough(): + """Non-/servers/ paths ending in /mcp must pass through without rewriting.""" + from unittest.mock import AsyncMock + + from mcpgateway.main import MCPPathRewriteMiddleware + + app_mock = AsyncMock() + middleware = MCPPathRewriteMiddleware(app_mock) + + scope = {"type": "http", "path": "/custom/mcp", "headers": []} + receive = AsyncMock() + send = AsyncMock() + + with patch("mcpgateway.main.streamable_http_auth", new_callable=AsyncMock, return_value=True): + await middleware._call_streamable_http(scope, receive, send) + + app_mock.assert_awaited_once_with(scope, receive, send) + # Path must NOT be rewritten to /mcp/ + assert scope["path"] == "/custom/mcp" + + def test_service_instances(): """Test that service instances exist.""" # First-Party diff --git a/tests/unit/plugins/test_retry_with_backoff.py b/tests/unit/plugins/test_retry_with_backoff.py new file mode 100644 index 00000000000..cf774ad1bc4 --- /dev/null +++ b/tests/unit/plugins/test_retry_with_backoff.py @@ -0,0 +1,613 @@ +# -*- coding: utf-8 -*- +"""Tests for Retry With Backoff Plugin. + +Verifies: +1. _compute_delay_ms — no-jitter exact values, jitter range, exponential growth, cap +2. _is_failure — isError flag, status_code variants, non-retriable codes, non-dict +3. _cfg_for — base config passthrough, per-tool override merging +4. RetryWithBackoffPlugin.__init__ — max_retries clamping, tool_overrides clamping +5. tool_post_invoke — first failure signals retry, exhaustion gives up, success resets state +6. State isolation — unique request_id per make_context() call ensures natural key isolation +7. Rust / Python path selection — Rust fast path taken when available, Python fallback when absent +8. retry_policy metadata — all return paths include advisory policy dict; resource_post_fetch hook +""" + +import logging +import uuid +import pytest +from unittest.mock import MagicMock, patch + +from plugins.retry_with_backoff.retry_with_backoff import ( + RetryWithBackoffPlugin, + RetryConfig, + _STATE, + _STATE_TTL_SECONDS, + _get_state, + _del_state, + _cfg_for, + _compute_delay_ms, + _is_failure, +) +from mcpgateway.plugins.framework import ( + PluginConfig, + PluginContext, + ResourcePostFetchPayload, + ToolPostInvokePayload, + GlobalContext, +) +from mcpgateway.common.models import ResourceContent + +# --------------------------------------------------------------------------- +# Helpers / fixtures +# --------------------------------------------------------------------------- + + +def make_plugin(config_overrides: dict | None = None) -> RetryWithBackoffPlugin: + """Build a plugin with default config, optionally overriding fields. + + Args: + config_overrides: Optional dict of config fields to override. + + Returns: + Configured RetryWithBackoffPlugin instance. + """ + cfg = { + "max_retries": 3, + "backoff_base_ms": 200, + "max_backoff_ms": 5000, + "jitter": False, # deterministic by default in tests + "retry_on_status": [429, 500, 502, 503, 504], + "tool_overrides": {}, + } + if config_overrides: + cfg.update(config_overrides) + plugin_config = PluginConfig( + id="test-retry", + kind="retry_with_backoff", + name="Test Retry Plugin", + enabled=True, + order=0, + config=cfg, + ) + return RetryWithBackoffPlugin(plugin_config) + + +def make_context() -> PluginContext: + # Unique request_id per call ensures each test's state entries never + # collide with another test's, giving natural isolation without clearing. + return PluginContext(plugin_id="test-retry", global_context=GlobalContext(request_id=str(uuid.uuid4()))) + + +def make_payload(tool: str, result: dict) -> ToolPostInvokePayload: + return ToolPostInvokePayload(name=tool, result=result) + + +# --------------------------------------------------------------------------- +# 1. _compute_delay_ms +# --------------------------------------------------------------------------- + + +class TestComputeDelayMs: + def test_no_jitter_returns_exact_ceiling(self): + cfg = RetryConfig(backoff_base_ms=200, max_backoff_ms=5000, jitter=False) + assert _compute_delay_ms(0, cfg) == 200 # base * 2^0 + assert _compute_delay_ms(1, cfg) == 400 # base * 2^1 + assert _compute_delay_ms(2, cfg) == 800 # base * 2^2 + + def test_no_jitter_caps_at_max_backoff(self): + cfg = RetryConfig(backoff_base_ms=200, max_backoff_ms=500, jitter=False) + assert _compute_delay_ms(0, cfg) == 200 + assert _compute_delay_ms(1, cfg) == 400 + assert _compute_delay_ms(2, cfg) == 500 # capped — 800 > 500 + assert _compute_delay_ms(10, cfg) == 500 # still capped + + def test_jitter_returns_value_within_cap(self): + cfg = RetryConfig(backoff_base_ms=200, max_backoff_ms=300, jitter=True) + delay = _compute_delay_ms(5, cfg) + assert 0 <= delay <= 300 + + def test_exponential_growth_without_jitter(self): + cfg = RetryConfig(backoff_base_ms=100, max_backoff_ms=100_000, jitter=False) + delays = [_compute_delay_ms(i, cfg) for i in range(5)] + assert delays == [100, 200, 400, 800, 1600] + + +# --------------------------------------------------------------------------- +# 2. _is_failure +# --------------------------------------------------------------------------- + + +class TestIsFailure: + def setup_method(self): + self.cfg = RetryConfig() + + def test_is_error_true_triggers_failure(self): + assert _is_failure({"isError": True}, self.cfg) is True + + def test_is_error_false_is_not_failure(self): + assert _is_failure({"isError": False}, self.cfg) is False + + def test_status_code_500_in_structured_content_is_failure(self): + assert _is_failure({"isError": False, "structuredContent": {"status_code": 500}}, self.cfg) is True + + def test_status_400_in_structured_content_is_not_retriable(self): + assert _is_failure({"isError": False, "structuredContent": {"status_code": 400}}, self.cfg) is False + + def test_status_200_in_structured_content_is_not_failure(self): + assert _is_failure({"isError": False, "structuredContent": {"status_code": 200}}, self.cfg) is False + + def test_check_text_content_disabled_by_default(self): + # check_text_content=false (default): text content with status_code NOT checked + result = { + "isError": False, + "structuredContent": None, + "content": [{"type": "text", "text": '{"status_code": 503, "message": "downstream down"}'}], + } + assert _is_failure(result, self.cfg) is False + + def test_check_text_content_enabled_retryable_status(self): + cfg = RetryConfig(check_text_content=True) + result = { + "isError": False, + "structuredContent": None, + "content": [{"type": "text", "text": '{"status_code": 503, "message": "downstream down"}'}], + } + assert _is_failure(result, cfg) is True + + def test_check_text_content_enabled_non_retryable_status(self): + cfg = RetryConfig(check_text_content=True) + result = { + "isError": False, + "structuredContent": None, + "content": [{"type": "text", "text": '{"status_code": 400, "message": "bad request"}'}], + } + assert _is_failure(result, cfg) is False + + def test_check_text_content_skipped_when_structured_content_present(self): + # Text content parsing only runs when structuredContent is absent (None) + # If structuredContent is present but has no failure, text content is NOT parsed + cfg = RetryConfig(check_text_content=True) + result = { + "isError": False, + "structuredContent": {"status_code": 200}, # present, not a failure + "content": [{"type": "text", "text": '{"status_code": 503}'}], # would be a failure if parsed + } + assert _is_failure(result, cfg) is False + + def test_check_text_content_invalid_json_ignored(self): + cfg = RetryConfig(check_text_content=True) + result = { + "isError": False, + "structuredContent": None, + "content": [{"type": "text", "text": "not json at all"}], + } + assert _is_failure(result, cfg) is False + + def test_check_text_content_is_error_in_text(self): + cfg = RetryConfig(check_text_content=True) + result = { + "isError": False, + "structuredContent": None, + "content": [{"type": "text", "text": '{"isError": true, "message": "failed"}'}], + } + assert _is_failure(result, cfg) is True + + def test_non_dict_result_is_not_failure(self): + assert _is_failure("error string", self.cfg) is False + assert _is_failure(None, self.cfg) is False + assert _is_failure(42, self.cfg) is False + + def test_empty_dict_is_not_failure(self): + assert _is_failure({}, self.cfg) is False + + def test_custom_retry_on_status(self): + cfg = RetryConfig(retry_on_status=[408]) + assert _is_failure({"structuredContent": {"status_code": 408}}, cfg) is True + assert _is_failure({"structuredContent": {"status_code": 500}}, cfg) is False # not in custom list + + def test_structured_content_is_error_true_triggers_failure(self): + # Signal 2: structuredContent.isError=True counts as a failure even + # when the outer isError is False (nested error from a structured response). + assert _is_failure({"isError": False, "structuredContent": {"isError": True}}, self.cfg) is True + + def test_structured_content_without_status_code_is_not_failure(self): + # structuredContent present but contains neither isError nor status_code → not a failure. + assert _is_failure({"isError": False, "structuredContent": {"result": "ok"}}, self.cfg) is False + + # -- Signal 1 status-code-aware tests -- + + def test_is_error_with_retryable_status_triggers_failure(self): + # isError=True + structuredContent.status_code in retry_on_status → retry. + result = {"isError": True, "structuredContent": {"status_code": 503}} + assert _is_failure(result, self.cfg) is True + + def test_is_error_with_non_retryable_status_skips_retry(self): + # isError=True + status_code NOT in retry_on_status → not retryable. + result = {"isError": True, "structuredContent": {"status_code": 400}} + assert _is_failure(result, self.cfg) is False + + def test_is_error_with_404_skips_retry(self): + result = {"isError": True, "structuredContent": {"status_code": 404}} + assert _is_failure(result, self.cfg) is False + + def test_is_error_with_401_skips_retry(self): + result = {"isError": True, "structuredContent": {"status_code": 401}} + assert _is_failure(result, self.cfg) is False + + def test_is_error_without_status_code_always_retries(self): + # isError=True with no structuredContent → generic exception → always retry. + assert _is_failure({"isError": True}, self.cfg) is True + assert _is_failure({"isError": True, "structuredContent": None}, self.cfg) is True + + def test_is_error_with_empty_structured_content_always_retries(self): + # isError=True + structuredContent without status_code → still a generic error. + result = {"isError": True, "structuredContent": {}} + assert _is_failure(result, self.cfg) is True + + +# --------------------------------------------------------------------------- +# 3. _cfg_for +# --------------------------------------------------------------------------- + + +class TestCfgFor: + def test_no_override_returns_same_object(self): + cfg = RetryConfig() + result = _cfg_for(cfg, "unknown_tool") + assert result is cfg + + def test_override_merges_max_retries(self): + cfg = RetryConfig(max_retries=3, tool_overrides={"my_tool": {"max_retries": 1}}) + merged = _cfg_for(cfg, "my_tool") + assert merged.max_retries == 1 + assert merged.backoff_base_ms == cfg.backoff_base_ms # base fields preserved + + def test_override_does_not_include_tool_overrides(self): + cfg = RetryConfig(tool_overrides={"my_tool": {"max_retries": 1}}) + merged = _cfg_for(cfg, "my_tool") + assert merged.tool_overrides == {} + + def test_other_tool_not_affected_by_override(self): + cfg = RetryConfig(max_retries=3, tool_overrides={"tool_a": {"max_retries": 1}}) + result = _cfg_for(cfg, "tool_b") + assert result is cfg + assert result.max_retries == 3 + + +# --------------------------------------------------------------------------- +# 4. Plugin __init__ — clamping +# --------------------------------------------------------------------------- + + +class TestPluginInit: + def test_max_retries_not_clamped_when_within_ceiling(self): + with patch("plugins.retry_with_backoff.retry_with_backoff.get_settings") as mock_settings: + mock_settings.return_value.max_tool_retries = 5 + plugin = make_plugin({"max_retries": 3}) + assert plugin._cfg.max_retries == 3 + + def test_max_retries_clamped_to_gateway_ceiling(self): + with patch("plugins.retry_with_backoff.retry_with_backoff.get_settings") as mock_settings: + mock_settings.return_value.max_tool_retries = 2 + plugin = make_plugin({"max_retries": 5}) + assert plugin._cfg.max_retries == 2 + + def test_tool_override_max_retries_clamped(self): + with patch("plugins.retry_with_backoff.retry_with_backoff.get_settings") as mock_settings: + mock_settings.return_value.max_tool_retries = 2 + plugin = make_plugin( + { + "max_retries": 2, + "tool_overrides": {"slow_api": {"max_retries": 10}}, + } + ) + assert plugin._cfg.tool_overrides["slow_api"]["max_retries"] == 2 + + def test_clamping_emits_warning(self, caplog): + with patch("plugins.retry_with_backoff.retry_with_backoff.get_settings") as mock_settings: + mock_settings.return_value.max_tool_retries = 1 + with caplog.at_level(logging.WARNING): + make_plugin({"max_retries": 5}) + assert any(r.getMessage() == "retry_with_backoff: max_retries=5 exceeds gateway ceiling=1, clamping" for r in caplog.records) + + def test_max_retries_equal_ceiling_not_clamped(self): + """max_retries exactly equal to the gateway ceiling must not be clamped.""" + with patch("plugins.retry_with_backoff.retry_with_backoff.get_settings") as mock_settings: + mock_settings.return_value.max_tool_retries = 3 + plugin = make_plugin({"max_retries": 3}) + assert plugin._cfg.max_retries == 3 + + +# --------------------------------------------------------------------------- +# 5. tool_post_invoke — core behaviour +# --------------------------------------------------------------------------- + + +class TestToolPostInvoke: + @pytest.mark.asyncio + async def test_success_returns_no_retry(self): + plugin = make_plugin() + ctx = make_context() + payload = make_payload("tool_a", {"result": "ok"}) + result = await plugin.tool_post_invoke(payload, ctx) + assert result.retry_delay_ms == 0 + + @pytest.mark.asyncio + async def test_first_failure_requests_retry(self): + plugin = make_plugin() + ctx = make_context() + payload = make_payload("tool_a", {"isError": True}) + result = await plugin.tool_post_invoke(payload, ctx) + assert result.retry_delay_ms > 0 + + @pytest.mark.asyncio + async def test_delay_grows_on_consecutive_failures(self): + plugin = make_plugin({"jitter": False, "backoff_base_ms": 100}) + ctx = make_context() + # failure 1: attempt=0 → 100ms + r1 = await plugin.tool_post_invoke(make_payload("t", {"isError": True}), ctx) + # failure 2: attempt=1 → 200ms + r2 = await plugin.tool_post_invoke(make_payload("t", {"isError": True}), ctx) + assert r2.retry_delay_ms > r1.retry_delay_ms + + @pytest.mark.asyncio + async def test_exhausted_retries_returns_zero_delay(self): + """After max_retries failures, plugin gives up (retry_delay_ms=0).""" + plugin = make_plugin({"max_retries": 2}) + ctx = make_context() + payload = make_payload("tool_a", {"isError": True}) + # 2 failures → still within budget + await plugin.tool_post_invoke(payload, ctx) + await plugin.tool_post_invoke(payload, ctx) + # 3rd failure → exhausted + result = await plugin.tool_post_invoke(payload, ctx) + assert result.retry_delay_ms == 0 + + @pytest.mark.asyncio + async def test_exhaustion_resets_counter_for_next_call(self): + """After exhaustion the counter resets so the next independent call gets a fresh retry budget.""" + plugin = make_plugin({"max_retries": 2, "jitter": False}) + ctx = make_context() + payload = make_payload("tool_x", {"isError": True}) + # Exhaust: 3 failures (original + 2 retries) + await plugin.tool_post_invoke(payload, ctx) + await plugin.tool_post_invoke(payload, ctx) + await plugin.tool_post_invoke(payload, ctx) # exhausted, returns 0 + # Counter must be reset — next independent call should retry again + r = await plugin.tool_post_invoke(payload, ctx) + assert r.retry_delay_ms > 0, "next independent call must get a fresh retry, not be blocked by previous exhaustion" + + @pytest.mark.asyncio + async def test_success_resets_failure_counter(self): + """After a partial failure run followed by a success the state is cleared, + so a subsequent failure on the same invocation must retry again.""" + plugin = make_plugin({"max_retries": 1, "jitter": False}) + ctx = make_context() + # First failure — within budget + r1 = await plugin.tool_post_invoke(make_payload("t", {"isError": True}), ctx) + assert r1.retry_delay_ms > 0 + # Success — state is deleted (reset) + await plugin.tool_post_invoke(make_payload("t", {"result": "ok"}), ctx) + # Next failure on the same invocation must get a fresh retry budget + r3 = await plugin.tool_post_invoke(make_payload("t", {"isError": True}), ctx) + assert r3.retry_delay_ms > 0, "success must reset state so the next failure retries" + + @pytest.mark.asyncio + async def test_per_tool_override_is_applied(self): + """A tool with max_retries=1 override should exhaust after 1 retry.""" + plugin = make_plugin( + { + "max_retries": 3, + "tool_overrides": {"fragile_tool": {"max_retries": 1}}, + } + ) + ctx = make_context() + # 1st failure: within budget + r1 = await plugin.tool_post_invoke(make_payload("fragile_tool", {"isError": True}), ctx) + assert r1.retry_delay_ms > 0 + # 2nd failure: exhausted (override max_retries=1) + r2 = await plugin.tool_post_invoke(make_payload("fragile_tool", {"isError": True}), ctx) + assert r2.retry_delay_ms == 0 + + @pytest.mark.asyncio + async def test_different_tools_have_independent_state(self): + plugin = make_plugin({"max_retries": 1}) + ctx = make_context() + # tool_a exhausts retries + await plugin.tool_post_invoke(make_payload("tool_a", {"isError": True}), ctx) + await plugin.tool_post_invoke(make_payload("tool_a", {"isError": True}), ctx) + # tool_b is unaffected + r = await plugin.tool_post_invoke(make_payload("tool_b", {"isError": True}), ctx) + assert r.retry_delay_ms > 0 + + @pytest.mark.asyncio + async def test_status_code_failure_triggers_retry(self): + plugin = make_plugin() + ctx = make_context() + result = await plugin.tool_post_invoke(make_payload("t", {"structuredContent": {"status_code": 503}}), ctx) + assert result.retry_delay_ms > 0 + + @pytest.mark.asyncio + async def test_non_retriable_status_does_not_retry(self): + plugin = make_plugin() + ctx = make_context() + result = await plugin.tool_post_invoke(make_payload("t", {"structuredContent": {"status_code": 400}}), ctx) + assert result.retry_delay_ms == 0 + + @pytest.mark.asyncio + async def test_max_retries_zero_gives_up_immediately(self): + """max_retries=0 means no retries at all — every failure must return delay=0.""" + plugin = make_plugin({"max_retries": 0}) + ctx = make_context() + result = await plugin.tool_post_invoke(make_payload("t", {"isError": True}), ctx) + assert result.retry_delay_ms == 0 + + +# --------------------------------------------------------------------------- +# 6. _get_state +# --------------------------------------------------------------------------- + + +class TestGetState: + def test_creates_fresh_state_for_new_tool(self): + st = _get_state("brand_new_tool", "req-fresh") + assert st.consecutive_failures == 0 + assert st.last_failure_at == 0.0 + + def test_returns_same_object_on_second_call(self): + s1 = _get_state("tool_x", "req-same") + s1.consecutive_failures = 7 + s2 = _get_state("tool_x", "req-same") + assert s2.consecutive_failures == 7 + assert s1 is s2 + + def test_ttl_eviction_removes_stale_entries(self): + """Entries whose last_failure_at is older than _STATE_TTL_SECONDS are evicted.""" + import time + + key = "evict_tool:evict_req" + # Inject a stale entry directly into _STATE + from plugins.retry_with_backoff.retry_with_backoff import _ToolRetryState + + _STATE[key] = _ToolRetryState(consecutive_failures=3, last_failure_at=time.monotonic() - _STATE_TTL_SECONDS - 1) + assert key in _STATE + # _get_state triggers eviction + _get_state("other_tool", "other_req") + assert key not in _STATE, "stale entry should have been evicted" + # Clean up + _del_state("other_tool", "other_req") + + def test_ttl_eviction_preserves_fresh_entries(self): + """Entries within the TTL window are not evicted.""" + import time + + key = "fresh_tool:fresh_req" + from plugins.retry_with_backoff.retry_with_backoff import _ToolRetryState + + _STATE[key] = _ToolRetryState(consecutive_failures=1, last_failure_at=time.monotonic()) + _get_state("other_tool2", "other_req2") + assert key in _STATE, "fresh entry should not be evicted" + # Clean up + _STATE.pop(key, None) + _del_state("other_tool2", "other_req2") + + +# --------------------------------------------------------------------------- +# 7. Rust / Python path selection +# --------------------------------------------------------------------------- + + +class TestRustFallback: + """Verify that the plugin behaves identically whether the Rust extension is + present or absent, and that the correct code path is selected in each case. + """ + + @pytest.mark.asyncio + async def test_python_fallback_when_rust_unavailable(self): + """With _rust patched to None the Python path must still retry correctly.""" + plugin = make_plugin() + ctx = make_context() + + with patch.object(plugin, "_rust", None): + r1 = await plugin.tool_post_invoke(make_payload("t", {"isError": True}), ctx) + assert r1.retry_delay_ms > 0, "Python fallback should request a retry on first failure" + + r2 = await plugin.tool_post_invoke(make_payload("t", {"result": "ok"}), ctx) + assert r2.retry_delay_ms == 0, "Python fallback should return 0 on success" + + @pytest.mark.asyncio + async def test_rust_path_taken_when_available(self): + """When _RUST is not None and check_text_content=False, check_and_update + must be called instead of the Python state functions.""" + plugin = make_plugin() + ctx = make_context() + + # Mock the instance-level _rust so we can assert it was called, + # regardless of whether the .so is present. + mock_rust = MagicMock() + mock_rust.check_and_update.return_value = (True, 300) + + with patch.object(plugin, "_rust", mock_rust): + r = await plugin.tool_post_invoke(make_payload("t", {"isError": True}), ctx) + + mock_rust.check_and_update.assert_called_once() + assert r.retry_delay_ms == 300 + + @pytest.mark.asyncio + async def test_rust_path_bypassed_for_check_text_content(self): + """When check_text_content=True the plugin must use the Python path + even if _RUST is present, because signal 3 isn't implemented in Rust.""" + plugin = make_plugin({"check_text_content": True}) + ctx = make_context() + + mock_rust = MagicMock() + mock_rust.check_and_update.return_value = (True, 300) + + with patch.object(plugin, "_rust", mock_rust): + await plugin.tool_post_invoke(make_payload("t", {"isError": True}), ctx) + + mock_rust.check_and_update.assert_not_called() + + +# --------------------------------------------------------------------------- +# 8. retry_policy metadata +# --------------------------------------------------------------------------- + + +class TestRetryPolicyMetadata: + """Verify that retry_policy metadata is attached on every tool_post_invoke + return path and on resource_post_fetch.""" + + @pytest.mark.asyncio + async def test_success_path_includes_policy_metadata(self): + plugin = make_plugin({"max_retries": 2, "backoff_base_ms": 100, "max_backoff_ms": 1000, "retry_on_status": [429, 503]}) + ctx = make_context() + result = await plugin.tool_post_invoke(make_payload("t", {"result": "ok"}), ctx) + assert result.metadata["retry_policy"] == { + "max_retries": 2, + "backoff_base_ms": 100, + "max_backoff_ms": 1000, + "retry_on_status": [429, 503], + } + + @pytest.mark.asyncio + async def test_failure_retry_path_includes_policy_metadata(self): + plugin = make_plugin({"max_retries": 3, "backoff_base_ms": 200, "max_backoff_ms": 5000, "retry_on_status": [500]}) + ctx = make_context() + result = await plugin.tool_post_invoke(make_payload("t", {"isError": True}), ctx) + assert result.retry_delay_ms > 0 + assert result.metadata["retry_policy"] == { + "max_retries": 3, + "backoff_base_ms": 200, + "max_backoff_ms": 5000, + "retry_on_status": [500], + } + + @pytest.mark.asyncio + async def test_exhaustion_path_includes_policy_metadata(self): + plugin = make_plugin({"max_retries": 1, "backoff_base_ms": 200, "max_backoff_ms": 5000, "retry_on_status": [503]}) + ctx = make_context() + payload = make_payload("t", {"isError": True}) + await plugin.tool_post_invoke(payload, ctx) # failure 1 — within budget + result = await plugin.tool_post_invoke(payload, ctx) # failure 2 — exhausted + assert result.retry_delay_ms == 0 + assert result.metadata["retry_policy"] == { + "max_retries": 1, + "backoff_base_ms": 200, + "max_backoff_ms": 5000, + "retry_on_status": [503], + } + + @pytest.mark.asyncio + async def test_resource_post_fetch_returns_policy_metadata(self): + plugin = make_plugin({"max_retries": 2, "backoff_base_ms": 150, "max_backoff_ms": 3000, "retry_on_status": [503]}) + ctx = make_context() + content = ResourceContent(type="resource", id="r1", uri="file:///data.txt", text="hello") + payload = ResourcePostFetchPayload(uri="file:///data.txt", content=content) + result = await plugin.resource_post_fetch(payload, ctx) + assert result.metadata["retry_policy"] == { + "max_retries": 2, + "backoff_base_ms": 150, + "max_backoff_ms": 3000, + "retry_on_status": [503], + }