diff --git a/src/cmdforge/cli/__init__.py b/src/cmdforge/cli/__init__.py index f773fcf..545fbd8 100644 --- a/src/cmdforge/cli/__init__.py +++ b/src/cmdforge/cli/__init__.py @@ -97,6 +97,7 @@ def main(): p_run.add_argument("--data-classification", choices=["public", "internal", "private"], help="Input data classification") p_run.add_argument("--require-model-identity", action="store_true") p_run.add_argument("--require-model-digest", action="store_true") + p_run.add_argument("--provider-timeout", type=int, help="Seconds each AI provider call may take") p_run.add_argument("--result-envelope", choices=["json"], help="Wrap output with verified execution provenance") p_run.add_argument("tool_args", nargs=argparse.REMAINDER, help="Additional tool-specific arguments (use -- to separate)") p_run.set_defaults(func=cmd_run) @@ -123,6 +124,7 @@ def main(): p_run_once.add_argument("--data-classification", choices=["public", "internal", "private"]) p_run_once.add_argument("--require-model-identity", action="store_true") p_run_once.add_argument("--require-model-digest", action="store_true") + p_run_once.add_argument("--provider-timeout", type=int, help="Seconds each AI provider call may take") p_run_once.add_argument("--result-envelope", choices=["json"]) p_run_once.set_defaults(func=cmd_run_once) diff --git a/src/cmdforge/cli/tool_commands.py b/src/cmdforge/cli/tool_commands.py index 8c2ad35..e3816a4 100644 --- a/src/cmdforge/cli/tool_commands.py +++ b/src/cmdforge/cli/tool_commands.py @@ -81,6 +81,11 @@ def _execution_policy_from_args(*namespaces): getattr(namespace, "require_model_digest", False) for namespace in namespaces ), + provider_timeout=next(( + getattr(namespace, "provider_timeout", None) + for namespace in namespaces + if getattr(namespace, "provider_timeout", None) + ), None), ) diff --git a/src/cmdforge/providers.py b/src/cmdforge/providers.py index b8211b8..06b30bf 100644 --- a/src/cmdforge/providers.py +++ b/src/cmdforge/providers.py @@ -426,8 +426,11 @@ class ProviderExecutionPolicy: data_classification: str = "public" require_model_identity: bool = False require_model_digest: bool = False + provider_timeout: Optional[int] = None # the caller's limit in seconds; overrides a step's own timeout def __post_init__(self) -> None: + if self.provider_timeout is not None and self.provider_timeout <= 0: + raise ValueError("provider_timeout must be a positive number of seconds") if self.required_locality not in {None, "local", "remote"}: raise ValueError("required_locality must be local, remote, or null") if self.data_classification not in {"public", "internal", "private"}: diff --git a/src/cmdforge/runner.py b/src/cmdforge/runner.py index 1184bee..049f129 100644 --- a/src/cmdforge/runner.py +++ b/src/cmdforge/runner.py @@ -559,6 +559,8 @@ def execute_prompt_step( # Determine provider provider = provider_override or substitute_variables(step.provider, variables, warn_non_scalar=verbose) + # How long the provider may take: the caller's --provider-timeout, else the step's timeout, else 300 s + step_timeout = effective_policy.provider_timeout or step.timeout or 300 # Build context from the inside out so the final order is: # profile system prompt -> skills -> user prompt. @@ -580,7 +582,7 @@ def execute_prompt_step( result = mock_provider(prompt) else: result = call_provider( - provider, prompt, max_tokens=step.max_tokens, + provider, prompt, timeout=step_timeout, max_tokens=step.max_tokens, execution_policy=effective_policy, ) @@ -629,7 +631,7 @@ Please try again with valid JSON matching the schema exactly.""" result = mock_provider(current_prompt) else: result = call_provider( - provider, current_prompt, max_tokens=step.max_tokens, + provider, current_prompt, timeout=step_timeout, max_tokens=step.max_tokens, execution_policy=effective_policy, ) @@ -1186,6 +1188,8 @@ def create_argument_parser(tool: Tool) -> argparse.ArgumentParser: parser.add_argument("--require-model-identity", action="store_true") parser.add_argument("--require-model-digest", action="store_true") parser.add_argument("--result-envelope", choices=["json"]) + parser.add_argument("--provider-timeout", type=int, + help="Seconds each AI provider call may take (overrides the tool's own step timeouts)") # Tool-specific flags from arguments for arg in tool.arguments: @@ -1315,6 +1319,7 @@ def main(): data_classification=args.data_classification or "public", require_model_identity=args.require_model_identity, require_model_digest=args.require_model_digest, + provider_timeout=args.provider_timeout, ) provider_calls = [] diff --git a/src/cmdforge/tool.py b/src/cmdforge/tool.py index b46588a..d8af4a5 100644 --- a/src/cmdforge/tool.py +++ b/src/cmdforge/tool.py @@ -188,11 +188,14 @@ class PromptStep: max_tokens: Optional[int] = None # Max output tokens (provider-dependent) skills: Optional[List[str]] = None # Skill names to enable for this step fallback_policy: str = "allow" # "allow" | "deny" + timeout: Optional[int] = None # Seconds the provider may take (None: the default 300; --provider-timeout wins) def __post_init__(self) -> None: _validate_skill_selection(self.skills) if self.fallback_policy not in {"allow", "deny"}: raise ValueError("PromptStep fallback_policy must be allow or deny") + if self.timeout is not None and self.timeout <= 0: + raise ValueError("PromptStep timeout must be a positive number of seconds") def to_dict(self) -> dict: d = { @@ -221,6 +224,8 @@ class PromptStep: d["skills"] = self.skills if self.fallback_policy != "allow": d["fallback_policy"] = self.fallback_policy + if self.timeout is not None: + d["timeout"] = self.timeout return d @classmethod @@ -233,6 +238,12 @@ class PromptStep: raise ValueError("max_tokens must be an integer") if max_tokens <= 0 or max_tokens > 1_000_000: raise ValueError("max_tokens must be between 1 and 1000000") + timeout = data.get("timeout") + if timeout is not None: + try: + timeout = int(timeout) + except (TypeError, ValueError): + raise ValueError("timeout must be a whole number of seconds") return cls( prompt=data["prompt"], @@ -248,6 +259,7 @@ class PromptStep: max_tokens=max_tokens, skills=data.get("skills"), fallback_policy=data.get("fallback_policy", "allow"), + timeout=timeout, ) diff --git a/tests/test_runner.py b/tests/test_runner.py index 55b1887..4b918b9 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1256,3 +1256,35 @@ class TestProviderExecutionPolicies: assert success is False resolver.assert_not_called() + + +class TestPromptStepTimeout: + """X-8a: the provider's time limit comes from the caller or the step, not a fixed 300 s.""" + + @patch('cmdforge.runner.call_provider') + def test_default_stays_300_seconds(self, mock_call): + mock_call.return_value = ProviderResult(text="ok", success=True) + execute_prompt_step(PromptStep(prompt="x", provider="claude", output_var="o", plain_text=True), {}) + assert mock_call.call_args.kwargs["timeout"] == 300 + + @patch('cmdforge.runner.call_provider') + def test_step_timeout_is_used(self, mock_call): + mock_call.return_value = ProviderResult(text="ok", success=True) + step = PromptStep(prompt="x", provider="claude", output_var="o", plain_text=True, timeout=1800) + execute_prompt_step(step, {}) + assert mock_call.call_args.kwargs["timeout"] == 1800 + + @patch('cmdforge.runner.call_provider') + def test_callers_provider_timeout_overrides_the_step(self, mock_call): + from cmdforge.providers import ProviderExecutionPolicy + mock_call.return_value = ProviderResult(text="ok", success=True) + step = PromptStep(prompt="x", provider="claude", output_var="o", plain_text=True, timeout=1800) + execute_prompt_step(step, {}, execution_policy=ProviderExecutionPolicy(provider_timeout=120)) + assert mock_call.call_args.kwargs["timeout"] == 120 + + def test_step_timeout_round_trips_and_is_validated(self): + step = PromptStep.from_dict({"prompt": "x", "provider": "claude", "output_var": "o", "timeout": "900"}) + assert step.timeout == 900 and step.to_dict()["timeout"] == 900 + assert "timeout" not in PromptStep(prompt="x", provider="claude", output_var="o").to_dict() + with pytest.raises(ValueError): + PromptStep(prompt="x", provider="claude", output_var="o", timeout=0)