feat: add max_retries config for PydanticAI agent retries

.gitignore

@@ -47,6 +47,7 @@ tests/*
 !tests/test_branding.py
 !tests/test_artifact_analyzer.py
 !tests/test_scala_analyzer.py
+!tests/test_max_retries.py
 
 # Jupyter
 *.ipynb

codewiki/cli/adapters/doc_generator.py

@@ -140,6 +140,7 @@ def generate(self) -> DocumentationJob:
                 fallback_model=self.config.get("fallback_model"),
                 provider=self.config.get("provider", "openai-compatible"),
                 aws_region=self.config.get("aws_region", "us-east-1"),
+                max_retries=self.config.get("max_retries", 3),
                 max_tokens=self.config.get("max_tokens", 32768),
                 max_token_per_module=self.config.get("max_token_per_module", 36369),
                 max_token_per_leaf_module=self.config.get("max_token_per_leaf_module", 16000),

codewiki/cli/commands/config.py

@@ -39,6 +39,7 @@ def config_group():
 @click.option("--main-model", type=str, help="Primary model for documentation generation")
 @click.option("--cluster-model", type=str, help="Model for module clustering (recommend top-tier)")
 @click.option("--fallback-model", type=str, help="Fallback model for documentation generation")
+@click.option("--max-retries", type=int, help="Maximum number of retries for failed agent tool calls (default: 3)")
 @click.option("--max-tokens", type=int, help="Maximum tokens for LLM response (default: 32768)")
 @click.option(
     "--max-token-per-module",
@@ -93,6 +94,7 @@ def config_set(
     main_model: Optional[str],
     cluster_model: Optional[str],
     fallback_model: Optional[str],
+    max_retries: Optional[int],
     max_tokens: Optional[int],
     max_token_per_module: Optional[int],
     max_token_per_leaf_module: Optional[int],
@@ -106,12 +108,12 @@ def config_set(
 ):
     """
     Set configuration values for CodeWiki.
-    
+
     API keys are stored securely in your system keychain:
       • macOS: Keychain Access
-      • Windows: Credential Manager  
+      • Windows: Credential Manager
       • Linux: Secret Service (GNOME Keyring, KWallet)
-    
+
     Examples:
 
     \b
@@ -162,6 +164,7 @@ def config_set(
                 main_model,
                 cluster_model,
                 fallback_model,
+                max_retries,
                 max_tokens,
                 max_token_per_module,
                 max_token_per_leaf_module,
@@ -206,6 +209,11 @@ def config_set(
         if fallback_model:
             validated_data["fallback_model"] = validate_model_name(fallback_model)
 
+        if max_retries is not None:
+            if max_retries < 0:
+                raise ConfigurationError("max_retries must be a non-negative integer")
+            validated_data["max_retries"] = max_retries
+
         if max_tokens is not None:
             if max_tokens < 1:
                 raise ConfigurationError("max_tokens must be a positive integer")
@@ -254,6 +262,7 @@ def config_set(
             main_model=validated_data.get("main_model"),
             cluster_model=validated_data.get("cluster_model"),
             fallback_model=validated_data.get("fallback_model"),
+            max_retries=validated_data.get("max_retries"),
             max_tokens=validated_data.get("max_tokens"),
             max_token_per_module=validated_data.get("max_token_per_module"),
             max_token_per_leaf_module=validated_data.get("max_token_per_leaf_module"),
@@ -380,6 +389,7 @@ def config_show(output_json: bool):
                 "cluster_model": config.cluster_model if config else "",
                 "fallback_model": config.fallback_model if config else "glm-4p5",
                 "default_output": config.default_output if config else "docs",
+                "max_retries": config.max_retries if config else 3,
                 "max_tokens": config.max_tokens if config else 32768,
                 "max_token_per_module": config.max_token_per_module if config else 36369,
                 "max_token_per_leaf_module": config.max_token_per_leaf_module if config else 16000,
@@ -441,6 +451,7 @@ def config_show(output_json: bool):
             click.echo()
             click.secho("Token Settings", fg="cyan", bold=True)
             if config:
+                click.echo(f"  Max Retries:             {config.max_retries}")
                 click.echo(f"  Max Tokens:              {config.max_tokens}")
                 click.echo(f"  Max Token/Module:        {config.max_token_per_module}")
                 click.echo(f"  Max Token/Leaf Module:   {config.max_token_per_leaf_module}")

codewiki/cli/commands/generate.py

@@ -272,6 +272,12 @@ def _find_affected(tree, parent_names=None):
     is_flag=True,
     help="Show detailed progress and debug information",
 )
+@click.option(
+    "--max-retries",
+    type=int,
+    default=None,
+    help="Maximum retries for failed LLM requests (overrides config)",
+)
 @click.option(
     "--max-tokens",
     type=int,
@@ -394,6 +400,7 @@ def generate_command(
     instructions: str | None,
     use_gitignore: bool | None,
     verbose: bool,
+    max_retries: int | None,
     max_tokens: int | None,
     max_token_per_module: int | None,
     max_token_per_leaf_module: int | None,
@@ -707,6 +714,8 @@ def generate_command(
                 "provider": getattr(config, "provider", "openai-compatible"),
                 "aws_region": getattr(config, "aws_region", "us-east-1"),
                 "agent_instructions": agent_instructions_dict,
+                # Max retries setting (runtime override takes precedence)
+                "max_retries": max_retries if max_retries is not None else config.max_retries,
                 # Max token settings (runtime overrides take precedence)
                 "max_tokens": max_tokens if max_tokens is not None else config.max_tokens,
                 "max_token_per_module": max_token_per_module

codewiki/cli/config_manager.py

@@ -132,6 +132,7 @@ def save(
         cluster_model: Optional[str] = None,
         fallback_model: Optional[str] = None,
         default_output: Optional[str] = None,
+        max_retries: Optional[int] = None,
         max_tokens: Optional[int] = None,
         max_token_per_module: Optional[int] = None,
         max_token_per_leaf_module: Optional[int] = None,
@@ -153,6 +154,7 @@ def save(
             cluster_model: Clustering model
             fallback_model: Fallback model
             default_output: Default output directory
+            max_retries: Maximum number of retries for failed agent tool calls
             max_tokens: Maximum tokens for LLM response
             max_token_per_module: Maximum tokens per module for clustering
             max_token_per_leaf_module: Maximum tokens per leaf module
@@ -197,6 +199,8 @@ def save(
             self._config.fallback_model = fallback_model
         if default_output is not None:
             self._config.default_output = default_output
+        if max_retries is not None:
+            self._config.max_retries = max_retries
         if max_tokens is not None:
             self._config.max_tokens = max_tokens
         if max_token_per_module is not None:

codewiki/cli/models/config.py

@@ -126,6 +126,7 @@ class Configuration:
         aws_region: AWS region for Bedrock provider
         api_version: Azure OpenAI API version
         azure_deployment: Azure OpenAI deployment name
+        max_retries: Maximum number of retries for failed agent tool calls (default: 3)
         max_tokens: Maximum tokens for LLM response (default: 32768)
         max_token_per_module: Maximum tokens per module for clustering (default: 36369)
         max_token_per_leaf_module: Maximum tokens per leaf module (default: 16000)
@@ -144,6 +145,7 @@ class Configuration:
     aws_region: str = "us-east-1"
     api_version: str = "2024-12-01-preview"
     azure_deployment: str = ""
+    max_retries: int = 3
     max_tokens: int = 32768
     max_token_per_module: int = 36369
     max_token_per_leaf_module: int = 16000
@@ -183,6 +185,7 @@ def to_dict(self) -> dict:
             "aws_region": self.aws_region,
             "api_version": self.api_version,
             "azure_deployment": self.azure_deployment,
+            "max_retries": self.max_retries,
             "max_tokens": self.max_tokens,
             "max_token_per_module": self.max_token_per_module,
             "max_token_per_leaf_module": self.max_token_per_leaf_module,
@@ -220,6 +223,7 @@ def from_dict(cls, data: dict) -> "Configuration":
             aws_region=data.get("aws_region", "us-east-1"),
             api_version=data.get("api_version", "2024-12-01-preview"),
             azure_deployment=data.get("azure_deployment", ""),
+            max_retries=data.get("max_retries", 3),
             max_tokens=data.get("max_tokens", 32768),
             max_token_per_module=data.get("max_token_per_module", 36369),
             max_token_per_leaf_module=data.get("max_token_per_leaf_module", 16000),
@@ -296,6 +300,7 @@ def to_backend_config(
             aws_region=self.aws_region,
             api_version=self.api_version,
             azure_deployment=self.azure_deployment,
+            max_retries=self.max_retries,
             max_tokens=self.max_tokens,
             max_token_per_module=self.max_token_per_module,
             max_token_per_leaf_module=self.max_token_per_leaf_module,

codewiki/mcp/server.py

@@ -534,6 +534,7 @@ async def _legacy_generate_docs(arguments: dict[str, Any]) -> list[TextContent]:
         fallback_model=config.fallback_model,
         provider=getattr(config, "provider", "openai-compatible"),
         aws_region=getattr(config, "aws_region", "us-east-1"),
+        max_retries=config.max_retries,
         max_tokens=config.max_tokens,
         agent_instructions=agent_instructions or None,
         use_gitignore=arguments.get("use_gitignore", True),

codewiki/src/be/agent_tools/generate_sub_module_documentations.py

@@ -91,6 +91,7 @@ async def generate_sub_module_documentation(
                     str_replace_editor_tool,
                     generate_sub_module_documentation_tool,
                 ],
+                retries=ctx.deps.config.max_retries,
             )
         else:
             sub_agent = Agent(
@@ -101,6 +102,7 @@ async def generate_sub_module_documentation(
                     module_name=sub_module_name, custom_instructions=ctx.deps.custom_instructions
                 ),
                 tools=[read_code_components_tool, str_replace_editor_tool],
+                retries=ctx.deps.config.max_retries,
             )
 
         deps.current_module_name = sub_module_name

codewiki/src/be/pydantic_ai_backend.py

@@ -81,6 +81,7 @@ async def run_update_agent(
             deps_type=CodeWikiDeps,
             tools=[read_code_components_tool, str_replace_editor_tool],
             system_prompt=system_prompt,
+            retries=self._config.max_retries,
         )
         started = time.time()
         result = await agent.run(user_prompt, deps=deps)
@@ -127,6 +128,7 @@ async def run_module_agent(
                     generate_sub_module_documentation_tool,
                 ],
                 system_prompt=format_system_prompt(module_name, self._custom_instructions),
+                retries=config.max_retries,
             )
         else:
             agent = Agent(
@@ -135,6 +137,7 @@ async def run_module_agent(
                 deps_type=CodeWikiDeps,
                 tools=[read_code_components_tool, str_replace_editor_tool],
                 system_prompt=format_leaf_system_prompt(module_name, self._custom_instructions),
+                retries=config.max_retries,
             )
 
         deps = CodeWikiDeps(

codewiki/src/config.py

@@ -85,6 +85,9 @@ class Config:
     aws_region: str = "us-east-1"
     api_version: str = "2024-12-01-preview"  # Azure OpenAI API version
     azure_deployment: str = ""  # Azure OpenAI deployment name
+    # Number of times an agent may retry a tool call whose arguments fail
+    # validation before giving up (default: 3)
+    max_retries: int = 3
     # Max token settings
     max_tokens: int = DEFAULT_MAX_TOKENS
     max_token_per_module: int = DEFAULT_MAX_TOKEN_PER_MODULE
@@ -211,6 +214,7 @@ def from_cli(
         aws_region: str = "us-east-1",
         api_version: str = "2024-12-01-preview",
         azure_deployment: str = "",
+        max_retries: int = 3,
         max_tokens: int = DEFAULT_MAX_TOKENS,
         max_token_per_module: int = DEFAULT_MAX_TOKEN_PER_MODULE,
         max_token_per_leaf_module: int = DEFAULT_MAX_TOKEN_PER_LEAF_MODULE,
@@ -239,6 +243,7 @@ def from_cli(
             aws_region: AWS region for Bedrock provider
             api_version: Azure OpenAI API version
             azure_deployment: Azure OpenAI deployment name
+            max_retries: Maximum number of retries for failed agent tool calls (default: 3)
             max_tokens: Maximum tokens for LLM response
             max_token_per_module: Maximum tokens per module for clustering
             max_token_per_leaf_module: Maximum tokens per leaf module
@@ -276,6 +281,7 @@ def from_cli(
             aws_region=aws_region,
             api_version=api_version,
             azure_deployment=azure_deployment,
+            max_retries=max_retries,
             max_tokens=max_tokens,
             max_token_per_module=max_token_per_module,
             max_token_per_leaf_module=max_token_per_leaf_module,

tests/test_max_retries.py

@@ -0,0 +1,95 @@
+"""Tests for max_retries PydanticAI agent configuration."""
+
+from codewiki.cli.models.config import Configuration
+from codewiki.src.config import Config
+
+from pydantic_ai import Agent, ModelRetry
+from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart
+from pydantic_ai.models.function import FunctionModel
+
+
+def _make_config(max_retries: int | None = None) -> Config:
+    """Build a minimal Config for testing."""
+    kwargs = {
+        "repo_path": "/tmp/repo",
+        "output_dir": "/tmp/out",
+        "llm_base_url": "http://localhost:1/v1",
+        "llm_api_key": "test-key",
+        "main_model": "test-model",
+        "cluster_model": "test-model",
+    }
+
+    if max_retries is not None:
+        kwargs["max_retries"] = max_retries
+
+    return Config.from_cli(**kwargs)
+
+
+def test_config_max_retries_defaults_to_three():
+    """Config defaults max_retries to 3."""
+    config = _make_config()
+    assert config.max_retries == 3
+
+
+def test_config_accepts_max_retries():
+    """Config accepts a custom max_retries value."""
+    config = _make_config(max_retries=5)
+    assert config.max_retries == 5
+
+
+def test_configuration_round_trip_preserves_max_retries():
+    """Configuration serialization preserves max_retries."""
+    config = Configuration(
+        base_url="http://localhost:1/v1",
+        main_model="test-model",
+        cluster_model="test-model",
+        max_retries=7,
+    )
+
+    restored = Configuration.from_dict(config.to_dict())
+    assert restored.max_retries == 7
+
+
+def test_agent_retries_failed_tool_call():
+    """Agent retries a tool call after ModelRetry."""
+    tool_attempts = 0
+    model_calls = 0
+
+    def model_function(messages, info):
+        nonlocal model_calls
+        model_calls += 1
+
+        if model_calls <= 2:
+            return ModelResponse(
+                parts=[
+                    ToolCallPart(
+                        tool_name="test_tool",
+                        args={},
+                    )
+                ]
+            )
+
+        return ModelResponse(
+            parts=[TextPart("success")],
+        )
+
+    agent = Agent(
+        FunctionModel(model_function),
+        retries=1,
+    )
+
+    @agent.tool_plain
+    def test_tool() -> str:
+        nonlocal tool_attempts
+        tool_attempts += 1
+
+        if tool_attempts == 1:
+            raise ModelRetry("Intentional failure")
+
+        return "success"
+
+    result = agent.run_sync("Call the test tool")
+
+    assert result.output == "success"
+    assert tool_attempts == 2
+    assert model_calls == 3

tests/test_sub_module_dedupe.py

@@ -53,7 +53,7 @@ def _deps(tmp_path) -> CodeWikiDeps:
         module_tree={},
         max_depth=2,
         current_depth=1,
-        config=SimpleNamespace(max_token_per_leaf_module=4000),
+        config=SimpleNamespace(max_token_per_leaf_module=4000, max_retries=3),
         custom_instructions="",
     )