forked from molecule-ai/molecule-core
feat(issue-652): wire effort and task_budget to claude sdk output_config
Adds _load_config_dict() helper to ClaudeSDKExecutor and wires the new effort and task_budget config fields into _build_options() before the Anthropic API call: - effort (str): low|medium|high|xhigh|max — populates output_config.effort - task_budget (int): advisory total-token budget; must be >= 20000 when set; automatically adds task-budgets-2026-03-13 beta header Also adds WorkspaceConfig.effort and WorkspaceConfig.task_budget fields in config.py and 5 acceptance tests covering all code paths. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
1ffa33cf61
commit
cf5428664b
@ -33,6 +33,8 @@ from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import yaml
|
||||
|
||||
import claude_agent_sdk as sdk
|
||||
|
||||
from a2a.server.agent_execution import AgentExecutor, RequestContext
|
||||
@ -233,6 +235,19 @@ class ClaudeSDKExecutor(AgentExecutor):
|
||||
return prompt
|
||||
return f"[Prior context from memory]\n{memories}\n\n{prompt}"
|
||||
|
||||
def _load_config_dict(self) -> dict:
|
||||
"""Read config.yaml as a raw dict for field-level inspection.
|
||||
|
||||
Returns an empty dict on any I/O or parse error so callers can
|
||||
always use ``.get()`` without guards.
|
||||
"""
|
||||
try:
|
||||
config_file = os.path.join(self.config_path, "config.yaml")
|
||||
with open(config_file) as f:
|
||||
return yaml.safe_load(f) or {}
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
def _build_options(self) -> Any:
|
||||
"""Build ClaudeAgentOptions.
|
||||
|
||||
@ -243,6 +258,18 @@ class ClaudeSDKExecutor(AgentExecutor):
|
||||
|
||||
The MCP server launcher uses `sys.executable` so tests and alternate
|
||||
virtual-env layouts don't depend on a `python3` shim being on PATH.
|
||||
|
||||
output_config wiring (issue #652)
|
||||
----------------------------------
|
||||
Reads ``effort`` and ``task_budget`` from config.yaml and populates
|
||||
``output_config`` on the SDK options before the API call:
|
||||
|
||||
- ``effort`` (str): one of low|medium|high|xhigh|max. xhigh is the
|
||||
Opus 4.7 recommended default for long agentic tasks.
|
||||
- ``task_budget`` (int): advisory total-token budget across the full
|
||||
agentic loop. Must be >= 20000 (API minimum) or 0/absent (unset).
|
||||
When set, the ``task-budgets-2026-03-13`` beta header is added so
|
||||
the API accepts the field.
|
||||
"""
|
||||
mcp_servers = {
|
||||
"a2a": {
|
||||
@ -250,7 +277,8 @@ class ClaudeSDKExecutor(AgentExecutor):
|
||||
"args": [get_mcp_server_path()],
|
||||
}
|
||||
}
|
||||
return sdk.ClaudeAgentOptions(
|
||||
|
||||
create_kwargs: dict = dict(
|
||||
model=self.model,
|
||||
permission_mode="bypassPermissions",
|
||||
cwd=self._resolve_cwd(),
|
||||
@ -259,6 +287,35 @@ class ClaudeSDKExecutor(AgentExecutor):
|
||||
resume=self._session_id,
|
||||
)
|
||||
|
||||
# --- output_config: effort + task_budget (issue #652) ---
|
||||
config = self._load_config_dict()
|
||||
output_config: dict = {}
|
||||
effort = config.get("effort", "")
|
||||
task_budget = config.get("task_budget", 0)
|
||||
|
||||
if effort:
|
||||
output_config["effort"] = effort # "low"|"medium"|"high"|"xhigh"|"max"
|
||||
|
||||
if task_budget and int(task_budget) >= 20000:
|
||||
output_config["task_budget"] = {
|
||||
"type": "tokens",
|
||||
"total": int(task_budget),
|
||||
}
|
||||
betas = list(create_kwargs.get("betas", []))
|
||||
if "task-budgets-2026-03-13" not in betas:
|
||||
betas.append("task-budgets-2026-03-13")
|
||||
create_kwargs["betas"] = betas
|
||||
elif task_budget and int(task_budget) > 0:
|
||||
# Below minimum — reject clearly before any API call is made.
|
||||
raise ValueError(
|
||||
f"task_budget must be >= 20000 tokens (got {task_budget})"
|
||||
)
|
||||
|
||||
if output_config:
|
||||
create_kwargs["output_config"] = output_config
|
||||
|
||||
return sdk.ClaudeAgentOptions(**create_kwargs)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Query streaming
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@ -228,6 +228,14 @@ class WorkspaceConfig:
|
||||
security_scan: SecurityScanConfig = field(default_factory=SecurityScanConfig)
|
||||
compliance: ComplianceConfig = field(default_factory=ComplianceConfig)
|
||||
sub_workspaces: list[dict] = field(default_factory=list)
|
||||
effort: str = ""
|
||||
"""Claude output effort level for the agentic loop: low | medium | high | xhigh | max.
|
||||
Empty string = not set (model default applies). xhigh is the Opus 4.7 recommended
|
||||
default for long agentic tasks. Passed as ``output_config.effort`` by ClaudeSDKExecutor."""
|
||||
task_budget: int = 0
|
||||
"""Advisory total-token budget across the full agentic loop. 0 = not set.
|
||||
Must be >= 20000 when non-zero (API minimum). When set, ClaudeSDKExecutor
|
||||
automatically adds the ``task-budgets-2026-03-13`` beta header."""
|
||||
|
||||
|
||||
def load_config(config_path: Optional[str] = None) -> WorkspaceConfig:
|
||||
@ -346,4 +354,6 @@ def load_config(config_path: Optional[str] = None) -> WorkspaceConfig:
|
||||
max_task_duration_seconds=int(compliance_raw.get("max_task_duration_seconds", 300)),
|
||||
),
|
||||
sub_workspaces=raw.get("sub_workspaces", []),
|
||||
effort=str(raw.get("effort", "")),
|
||||
task_budget=int(raw.get("task_budget", 0)),
|
||||
)
|
||||
|
||||
@ -1071,3 +1071,85 @@ def test_execute_clears_session_between_retries_on_process_error(caplog):
|
||||
# INFO log confirms the reset fired
|
||||
info_messages = " | ".join(r.message for r in caplog.records if r.levelname == "INFO")
|
||||
assert "SDK session reset after FakeProcessError" in info_messages
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _build_options — issue #652: effort + task_budget output_config wiring
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_options_with_config(config: dict):
|
||||
"""Helper: build ClaudeAgentOptions with the given config.yaml values.
|
||||
|
||||
Stubs out all I/O helpers so only the output_config wiring logic is tested.
|
||||
"""
|
||||
e = ClaudeSDKExecutor(system_prompt=None, config_path="/tmp", heartbeat=None)
|
||||
with patch.object(e, "_load_config_dict", return_value=config), \
|
||||
patch.object(e, "_resolve_cwd", return_value="/workspace"), \
|
||||
patch.object(e, "_build_system_prompt", return_value=None), \
|
||||
patch("claude_sdk_executor.get_mcp_server_path", return_value="/mcp.py"):
|
||||
return e._build_options()
|
||||
|
||||
|
||||
def test_build_options_effort_only_sets_output_config_no_beta():
|
||||
"""effort='xhigh', no task_budget → output_config={'effort':'xhigh'}, no betas.
|
||||
|
||||
Acceptance criterion: effort field wired into output_config without adding
|
||||
the task-budgets beta header (beta is only required for task_budget).
|
||||
"""
|
||||
opts = _build_options_with_config({"effort": "xhigh"})
|
||||
assert opts.kwargs.get("output_config") == {"effort": "xhigh"}
|
||||
assert "betas" not in opts.kwargs
|
||||
|
||||
|
||||
def test_build_options_task_budget_sets_output_config_and_beta():
|
||||
"""task_budget=128000 → output_config with token budget struct + beta header.
|
||||
|
||||
Acceptance criterion: task_budget >= 20000 writes the nested
|
||||
{'type':'tokens','total':N} struct and adds 'task-budgets-2026-03-13' to betas.
|
||||
"""
|
||||
opts = _build_options_with_config({"task_budget": 128000})
|
||||
assert opts.kwargs.get("output_config") == {
|
||||
"task_budget": {"type": "tokens", "total": 128000}
|
||||
}
|
||||
assert "task-budgets-2026-03-13" in opts.kwargs.get("betas", [])
|
||||
|
||||
|
||||
def test_build_options_both_effort_and_task_budget():
|
||||
"""Both effort and task_budget → combined output_config + beta header.
|
||||
|
||||
Acceptance criterion: both keys present in the single output_config dict;
|
||||
betas includes the task-budget feature flag.
|
||||
"""
|
||||
opts = _build_options_with_config({"effort": "high", "task_budget": 50000})
|
||||
assert opts.kwargs.get("output_config") == {
|
||||
"effort": "high",
|
||||
"task_budget": {"type": "tokens", "total": 50000},
|
||||
}
|
||||
assert "task-budgets-2026-03-13" in opts.kwargs.get("betas", [])
|
||||
|
||||
|
||||
def test_build_options_neither_effort_nor_task_budget_no_output_config():
|
||||
"""Empty config (effort='', task_budget=0) → output_config absent, no betas.
|
||||
|
||||
Acceptance criterion: when neither field is configured the SDK options
|
||||
are unchanged — no spurious output_config or betas keys.
|
||||
"""
|
||||
opts = _build_options_with_config({})
|
||||
assert "output_config" not in opts.kwargs
|
||||
assert "betas" not in opts.kwargs
|
||||
|
||||
|
||||
def test_build_options_task_budget_below_minimum_raises_value_error():
|
||||
"""task_budget=5000 (below 20000 API minimum) → ValueError before any API call.
|
||||
|
||||
Acceptance criterion: the executor must refuse to build options when
|
||||
task_budget is set but too small, so no invalid request reaches the API.
|
||||
"""
|
||||
e = ClaudeSDKExecutor(system_prompt=None, config_path="/tmp", heartbeat=None)
|
||||
with patch.object(e, "_load_config_dict", return_value={"task_budget": 5000}), \
|
||||
patch.object(e, "_resolve_cwd", return_value="/workspace"), \
|
||||
patch.object(e, "_build_system_prompt", return_value=None), \
|
||||
patch("claude_sdk_executor.get_mcp_server_path", return_value="/mcp.py"):
|
||||
with pytest.raises(ValueError, match="task_budget must be >= 20000"):
|
||||
e._build_options()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user