diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 1b8742ba4..3194a360f 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -725,6 +725,7 @@ def build_strix_agent( skills: list[str] | None = None, is_root: bool, scan_mode: str = "deep", + model: str | None = None, is_whitebox: bool = False, is_diff_scoped: bool = False, interactive: bool = False, @@ -809,7 +810,7 @@ def build_strix_agent( instructions=instructions, tools=tools, tool_use_behavior=_finish_tool_use_behavior, - model=None, + model=model, capabilities=[ Filesystem( configure_tools=_make_filesystem_configurator( @@ -838,6 +839,7 @@ def make_child_factory( chat_completions_tools: bool = False, strict_tool_schemas: bool = True, system_prompt_context: dict[str, Any] | None = None, + subagent_model: str | None = None, supports_images: bool = True, report_tool: Tool = create_vulnerability_report, coverage_tools: bool = True, @@ -854,6 +856,7 @@ def make_child_factory( name=name, skills=skills, is_root=False, + model=subagent_model, scan_mode=scan_mode, is_whitebox=is_whitebox, is_diff_scoped=is_diff_scoped, diff --git a/strix/core/runner.py b/strix/core/runner.py index 42d2c3e5e..1485d216d 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -233,6 +233,7 @@ async def run_strix_scan( status_sink: StatusSink | None = None, mcp_connection_requests: list[McpConnectionRequest] | None = None, mcp_status_sink: McpStatusSink | None = None, + subagent_model: str | None = None, root_finish_tool: Tool = finish_scan, report_tool: Tool = create_vulnerability_report, coverage_tools: bool = True, @@ -567,6 +568,7 @@ async def run_strix_scan( chat_completions_tools=chat_completions_tools, strict_tool_schemas=strict_tool_schemas, system_prompt_context=scope_context, + subagent_model=subagent_model or scan_config.get("subagent_model"), supports_images=supports_images, report_tool=report_tool, coverage_tools=coverage_tools, diff --git a/strix/interface/cli.py b/strix/interface/cli.py index a8992ce49..c28d385c3 100644 --- a/strix/interface/cli.py +++ b/strix/interface/cli.py @@ -88,8 +88,9 @@ async def run_cli(args: Any) -> None: # noqa: PLR0915 scan_config: dict[str, Any] = { "scan_id": args.run_name, "targets": args.targets_info, - "user_instructions": args.instruction or "", + "user_instruction": args.instruction or "", "run_name": args.run_name, + "subagent_model": getattr(args, "subagent_model", None), "diff_scope": getattr(args, "diff_scope", {"active": False}), "scan_mode": scan_mode, "non_interactive": bool(getattr(args, "non_interactive", False)), diff --git a/strix/interface/cli_args.py b/strix/interface/cli_args.py index c9b6951be..88ca29844 100644 --- a/strix/interface/cli_args.py +++ b/strix/interface/cli_args.py @@ -221,6 +221,17 @@ Strix Cloud: ), ) + parser.add_argument( + "--subagent-model", + dest="subagent_model", + type=str, + default=None, + help=( + "Model identifier to use for spawned subagents (e.g. claude-3-5-haiku, " + "gpt-4o-mini). If not specified, subagents inherit the root agent model." + ), + ) + parser.add_argument( "--scope-mode", type=str, @@ -532,6 +543,10 @@ def load_resume_state(args: argparse.Namespace) -> None: persisted_scan_mode = state.get("scan_mode") if persisted_scan_mode and args.scan_mode == "deep": args.scan_mode = persisted_scan_mode + persisted_subagent_model = state.get("subagent_model") + if not getattr(args, "subagent_model", None) and persisted_subagent_model: + args.subagent_model = persisted_subagent_model + agents_path = runtime_state_dir(run_dir) / "agents.json" if not agents_path.exists(): raise ResumeError( diff --git a/strix/report/state.py b/strix/report/state.py index 104e4af4c..3a7f9104f 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -1206,18 +1206,21 @@ def _estimate_response_cost(kwargs: Any, completion_response: Any) -> float | No candidates.append(model.rsplit("/", 1)[-1]) for candidate in candidates: + models_to_try = [candidate] resolved = resolve_litellm_model(candidate) - if not resolved: - continue - try: - value = completion_cost( - completion_response={"model": resolved, "usage": usage_payload}, - model=resolved, - ) - except Exception: # nosec B112 # noqa: BLE001, S112 - continue - if isinstance(value, int | float) and value > 0: - return float(value) + if resolved and resolved not in models_to_try: + models_to_try.append(resolved) + + for target_model in models_to_try: + try: + value = completion_cost( + completion_response={"model": target_model, "usage": usage_payload}, + model=target_model, + ) + except Exception: # nosec B112 # noqa: BLE001, S112 + continue + if isinstance(value, int | float) and value > 0: + return float(value) return None diff --git a/tests/test_cli_subagent_model.py b/tests/test_cli_subagent_model.py new file mode 100644 index 000000000..49e938058 --- /dev/null +++ b/tests/test_cli_subagent_model.py @@ -0,0 +1,48 @@ +import argparse +from pathlib import Path + +from strix.interface.cli_args import load_resume_state, parse_arguments + + +def test_subagent_model_default_none(monkeypatch): + """Verify --subagent-model defaults to None when omitted.""" + monkeypatch.setattr("sys.argv", ["strix", "--target", "https://example.com"]) + args = parse_arguments() + assert args.subagent_model is None + + +def test_subagent_model_flag_parsed(monkeypatch): + """Verify --subagent-model accepts a model name.""" + monkeypatch.setattr( + "sys.argv", + ["strix", "--target", "https://example.com", "--subagent-model", "claude-3-5-haiku"], + ) + args = parse_arguments() + assert args.subagent_model == "claude-3-5-haiku" + + +def test_resume_restores_subagent_model(monkeypatch, tmp_path): + """Verify _load_resume_state restores persisted subagent_model.""" + fake_state = { + "targets_info": [{"type": "domain", "original": "example.com", "details": {}}], + "subagent_model": "gpt-4o-mini", + "scan_mode": "deep", + } + + args = argparse.Namespace( + resume="test-run", + subagent_model=None, + scan_mode="deep", + diff_scope="auto", + targets_info=[], + workspace_mount=None, + instruction=None, + local_sources=[], + ) + + monkeypatch.setattr("strix.report.writer.read_run_record", lambda _run_dir: fake_state) + monkeypatch.setattr("strix.interface.cli_args.run_dir_for", lambda _r: tmp_path) + monkeypatch.setattr(Path, "exists", lambda _self: True) + + load_resume_state(args) + assert args.subagent_model == "gpt-4o-mini"