diff --git a/src/jumar/cli.py b/src/jumar/cli.py index feac634..6b47bf4 100644 --- a/src/jumar/cli.py +++ b/src/jumar/cli.py @@ -670,6 +670,7 @@ def run_item( base_url=judge_resolved.base_url, api_key_env=judge_resolved.api_key_env, reasoning_effort=judge_resolved.reasoning_effort, + max_tokens=judge_resolved.max_tokens, commands_allow=config.commands.allow, commands_deny=config.commands.deny, ) diff --git a/src/jumar/config.py b/src/jumar/config.py index a9a8820..d1f5bfb 100644 --- a/src/jumar/config.py +++ b/src/jumar/config.py @@ -19,6 +19,7 @@ import tomllib from collections.abc import Mapping from dataclasses import dataclass, field +from dataclasses import replace as dc_replace from enum import StrEnum from pathlib import Path, PurePath from typing import Any, cast @@ -177,7 +178,7 @@ def _system_timezone() -> str: # Scalar keys valid at [harness], inside a stage table, and inside a profile. _HARNESS_SCALAR_KEYS: frozenset[str] = frozenset( - {"agent", "model", "base_url", "api_key_env", "reasoning_effort"} + {"agent", "model", "base_url", "api_key_env", "reasoning_effort", "max_tokens"} ) # Sub-table under [harness] holding named alternative harnesses: @@ -222,6 +223,10 @@ class HarnessConfig: model is: decompose and judge are single bounded calls, while execute runs a tool loop where the cost of deliberation is paid on every turn. + ``max_tokens`` caps a single generation. Unset, the only bound is the + context window, so one runaway response can consume an entire subtask + deadline and take the whole attempt down with it. + ``None`` for a per-stage field means "inherit the top-level value". A ``HarnessConfig`` is always fully resolved by the time it reaches this @@ -235,22 +240,26 @@ class HarnessConfig: base_url: str | None = None api_key_env: str | None = None reasoning_effort: str | None = None + max_tokens: int | None = None # Per-stage overrides — None inherits the top-level value. decompose_agent: str | None = None decompose_model: str | None = None decompose_base_url: str | None = None decompose_api_key_env: str | None = None decompose_reasoning_effort: str | None = None + decompose_max_tokens: int | None = None execute_agent: str | None = None execute_model: str | None = None execute_base_url: str | None = None execute_api_key_env: str | None = None execute_reasoning_effort: str | None = None + execute_max_tokens: int | None = None judge_agent: str | None = None judge_model: str | None = None judge_base_url: str | None = None judge_api_key_env: str | None = None judge_reasoning_effort: str | None = None + judge_max_tokens: int | None = None def for_stage(self, stage: str) -> HarnessConfig: """Return the resolved ``HarnessConfig`` for *stage*. @@ -268,12 +277,14 @@ def for_stage(self, stage: str) -> HarnessConfig: b: str | None = getattr(self, f"{stage}_base_url") k: str | None = getattr(self, f"{stage}_api_key_env") r: str | None = getattr(self, f"{stage}_reasoning_effort") + x: int | None = getattr(self, f"{stage}_max_tokens") return HarnessConfig( agent=a if a is not None else self.agent, model=m if m is not None else self.model, base_url=b if b is not None else self.base_url, api_key_env=k if k is not None else self.api_key_env, reasoning_effort=r if r is not None else self.reasoning_effort, + max_tokens=x if x is not None else self.max_tokens, ) @@ -528,14 +539,18 @@ def _first(*values: str | None) -> str | None: ) profile_raw = cast(dict[str, Any], profiles_raw[harness_profile]) - def _validate_reasoning_efforts(resolved: HarnessConfig, label: str) -> None: - """Reject an unrecognised reasoning_effort at load, not at request time. + def _finalise(resolved: HarnessConfig, label: str) -> HarnessConfig: + """Validate reasoning_effort and coerce max_tokens, or raise ConfigError. + + Both are checked at load rather than at request time. A bad value would + otherwise be forwarded verbatim and ignored by the endpoint, leaving a + run that looks configured and is not. Every profile is finalised, not + just the selected one, so a typo in an unused profile fails the next + load rather than the next scheduled run that happens to select it. - A bad value would otherwise be forwarded verbatim and silently ignored - by the endpoint, leaving a run that looks configured and is not. Every - profile is checked, not just the selected one, so a typo in an unused - profile still fails the next load rather than the next scheduled run - that selects it. + max_tokens arrives as a string because the layering above resolves + every scalar uniformly; it is coerced back to int here, which is also + where a non-numeric or non-positive value is rejected. """ stage_attrs = (f"{s}_reasoning_effort" for s in sorted(_VALID_HARNESS_STAGES)) for attr in ("reasoning_effort", *stage_attrs): @@ -547,6 +562,28 @@ def _validate_reasoning_efforts(resolved: HarnessConfig, label: str) -> None: f"Valid values: {sorted(_VALID_REASONING_EFFORTS)}." ) + # Any, not int: the values land in fields the dataclass types + # int | None, and **kwargs into dc_replace cannot be narrowed per key. + coerced: dict[str, Any] = {} + token_attrs = (f"{s}_max_tokens" for s in sorted(_VALID_HARNESS_STAGES)) + for attr in ("max_tokens", *token_attrs): + value = getattr(resolved, attr) + if value is None: + continue + where = label if attr == "max_tokens" else f"{label}.{attr.split('_')[0]}" + try: + as_int = int(str(value)) + except ValueError: + raise ConfigError( + f"Invalid max_tokens {value!r} under [{where}]: expected a positive integer." + ) from None + if as_int <= 0: + raise ConfigError( + f"Invalid max_tokens {as_int} under [{where}]: expected a positive integer." + ) + coerced[attr] = as_int + return dc_replace(resolved, **coerced) if coerced else resolved + # Resolution order, highest first: # [harness.profiles..] → [harness.profiles.] # → [harness.] → [harness] @@ -557,7 +594,9 @@ def _validate_reasoning_efforts(resolved: HarnessConfig, label: str) -> None: def _build(profile_table: dict[str, Any], profile_name: str | None) -> HarnessConfig: """Layer *profile_table* over [harness] into one flat HarnessConfig.""" label = f"harness.{_HARNESS_PROFILES_KEY}.{profile_name}" - stage_kwargs: dict[str, str | None] = {} + # Any because max_tokens lands in an int | None field while every + # other scalar is str | None; _finalise coerces it after layering. + stage_kwargs: dict[str, Any] = {} for stage in sorted(_VALID_HARNESS_STAGES): base_stage = _stage_table(harness_raw, stage, f"harness.{stage}") profile_stage = ( @@ -585,19 +624,24 @@ def _build(profile_table: dict[str, Any], profile_name: str | None) -> HarnessCo _scalar(profile_table, "reasoning_effort"), _scalar(harness_raw, "reasoning_effort"), ), + # Still a string here; _finalise coerces it to int after layering. + max_tokens=_first( + _scalar(profile_table, "max_tokens"), + _scalar(harness_raw, "max_tokens"), + ), # type: ignore[arg-type] **stage_kwargs, ) - harness = _build(profile_raw, harness_profile) - _validate_reasoning_efforts(harness, "harness") + harness = _finalise(_build(profile_raw, harness_profile), "harness") # Every profile is resolved, not just the selected one, so an item can # name one with @harness= after load without re-reading the file. harness_profiles: dict[str, HarnessConfig] = { - name: _build(cast(dict[str, Any], profiles_raw[name]), name) + name: _finalise( + _build(cast(dict[str, Any], profiles_raw[name]), name), + f"harness.{_HARNESS_PROFILES_KEY}.{name}", + ) for name in sorted(profiles_raw) } - for _name, _resolved in harness_profiles.items(): - _validate_reasoning_efforts(_resolved, f"harness.{_HARNESS_PROFILES_KEY}.{_name}") commands_raw: dict[str, Any] = raw.get("commands", {}) commands = CommandPolicy( diff --git a/src/jumar/decompose.py b/src/jumar/decompose.py index 40f9529..ce55cab 100644 --- a/src/jumar/decompose.py +++ b/src/jumar/decompose.py @@ -438,6 +438,7 @@ def decompose( base_url=resolved.base_url, api_key_env=resolved.api_key_env, reasoning_effort=resolved.reasoning_effort, + max_tokens=resolved.max_tokens, commands_allow=config.commands.allow, commands_deny=config.commands.deny, ) diff --git a/src/jumar/execute.py b/src/jumar/execute.py index 8bcc1c8..99ec1d2 100644 --- a/src/jumar/execute.py +++ b/src/jumar/execute.py @@ -141,6 +141,7 @@ def execute( base_url=resolved.base_url, api_key_env=resolved.api_key_env, reasoning_effort=resolved.reasoning_effort, + max_tokens=resolved.max_tokens, commands_allow=config.commands.allow, commands_deny=config.commands.deny, ) diff --git a/src/jumar/models.py b/src/jumar/models.py index da1dd33..d142c5e 100644 --- a/src/jumar/models.py +++ b/src/jumar/models.py @@ -204,6 +204,7 @@ class HarnessInfo: base_url: str | None = None api_key_env: str | None = None reasoning_effort: str | None = None + max_tokens: int | None = None commands_allow: tuple[str, ...] = () commands_deny: tuple[str, ...] = () diff --git a/src/jumar/openai_agent.py b/src/jumar/openai_agent.py index 85fd650..5615681 100644 --- a/src/jumar/openai_agent.py +++ b/src/jumar/openai_agent.py @@ -389,6 +389,8 @@ def _rate() -> dict[str, Any]: # is unaffected by the default. if harness.reasoning_effort: payload["reasoning_effort"] = harness.reasoning_effort + if harness.max_tokens: + payload["max_tokens"] = harness.max_tokens if allow_tools: payload["tools"] = list(_TOOLS) @@ -455,6 +457,28 @@ def _rate() -> dict[str, Any]: if content: transcript.append(content) + # A generation stopped by `max_tokens` is not an answer. Without this + # branch it reads as one: `finish_reason` is "length", there are no + # tool calls to continue the loop with, and the fall-through below + # would return exit_status=0 with whatever half-sentence the model got + # to — a silent bad result, which is worse than the timeout the cap + # exists to prevent. Tool calls truncated mid-emission are refused for + # the same reason: the arguments JSON may be incomplete, so dispatching + # them would act on a half-read instruction. + if choices[0].get("finish_reason") == "length": + return AgentResult( + exit_status=-1, + stdout="\n".join(transcript), + stderr=( + f"generation hit max_tokens ({harness.max_tokens}) and was truncated; " + "raise [harness] max_tokens, or lower reasoning_effort so less of the " + "budget is spent before the answer" + ), + timed_out=False, + agent_claim=None, + **_rate(), + ) + if not tool_calls or not allow_tools: stdout = "\n".join(transcript) lines = [ln.strip() for ln in stdout.splitlines() if ln.strip()] diff --git a/tests/test_config.py b/tests/test_config.py index 1525142..78f2470 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1036,3 +1036,44 @@ def test_invalid_reasoning_effort_in_an_unselected_profile_still_fails(tmp_path: ) with pytest.raises(ConfigError, match="Invalid reasoning_effort"): load_config(tmp_path) + + +def test_max_tokens_resolves_per_stage_as_an_int(tmp_path: Path) -> None: + (tmp_path / "jumar.toml").write_text( + "[jumar.harness]\n" + 'agent = "openai"\n' + "max_tokens = 4096\n" + "[jumar.harness.execute]\n" + "max_tokens = 8192\n" + ) + config = load_config(tmp_path) + assert config.harness.for_stage("judge").max_tokens == 4096 + assert config.harness.for_stage("execute").max_tokens == 8192 + # Coerced back to int, not left as the string the layering produces. + assert isinstance(config.harness.for_stage("execute").max_tokens, int) + + +def test_max_tokens_defaults_to_none(tmp_path: Path) -> None: + (tmp_path / "jumar.toml").write_text('[jumar.harness]\nagent = "openai"\n') + assert load_config(tmp_path).harness.for_stage("execute").max_tokens is None + + +def test_non_positive_max_tokens_is_a_startup_error(tmp_path: Path) -> None: + (tmp_path / "jumar.toml").write_text('[jumar.harness]\nagent = "openai"\nmax_tokens = 0\n') + with pytest.raises(ConfigError, match="positive integer"): + load_config(tmp_path) + + +def test_non_numeric_max_tokens_is_a_startup_error(tmp_path: Path) -> None: + (tmp_path / "jumar.toml").write_text('[jumar.harness]\nagent = "openai"\nmax_tokens = "lots"\n') + with pytest.raises(ConfigError, match="positive integer"): + load_config(tmp_path) + + +def test_non_numeric_max_tokens_in_an_unselected_profile_still_fails(tmp_path: Path) -> None: + """Every profile is finalised, not only the one selected for this run.""" + (tmp_path / "jumar.toml").write_text( + '[jumar.harness]\nagent = "openai"\n[jumar.harness.profiles.heavy]\nmax_tokens = "heaps"\n' + ) + with pytest.raises(ConfigError, match="positive integer"): + load_config(tmp_path) diff --git a/tests/test_openai_agent.py b/tests/test_openai_agent.py index ab7a52e..3063bb3 100644 --- a/tests/test_openai_agent.py +++ b/tests/test_openai_agent.py @@ -963,3 +963,92 @@ def test_reasoning_effort_is_sent_on_every_turn_of_the_loop( assert result.exit_status == 0 assert len(handler.received) == 2 assert all(r["reasoning_effort"] == "medium" for r in handler.received) + + +# --------------------------------------------------------------------------- +# max_tokens / finish_reason +# --------------------------------------------------------------------------- + + +def _finished(response: dict[str, Any], reason: str) -> dict[str, Any]: + out = json.loads(json.dumps(response)) + out["choices"][0]["finish_reason"] = reason + return out + + +def test_max_tokens_is_sent_when_configured(tmp_path: Path, chat_server: Any) -> None: + base_url, handler = chat_server([_message("All done.")]) + result = run_openai_agent( + "do the thing", + cwd=tmp_path, + capabilities=_ALL_CAPS, + timeout_s=30, + harness=replace(_harness(base_url), max_tokens=4096), + ) + assert result.exit_status == 0 + assert handler.received[0]["max_tokens"] == 4096 + + +def test_max_tokens_is_omitted_when_unset(tmp_path: Path, chat_server: Any) -> None: + base_url, handler = chat_server([_message("All done.")]) + run_openai_agent( + "do the thing", + cwd=tmp_path, + capabilities=_ALL_CAPS, + timeout_s=30, + harness=_harness(base_url), + ) + assert "max_tokens" not in handler.received[0] + + +def test_truncated_generation_fails_rather_than_passing_as_an_answer( + tmp_path: Path, chat_server: Any +) -> None: + """finish_reason=length is a cut-off mid-sentence, not a final answer.""" + base_url, _handler = chat_server([_finished(_message("I have started to writ"), "length")]) + result = run_openai_agent( + "do the thing", + cwd=tmp_path, + capabilities=_ALL_CAPS, + timeout_s=30, + harness=replace(_harness(base_url), max_tokens=8), + ) + assert result.exit_status == -1 + assert result.timed_out is False + assert "max_tokens" in (result.stderr or "") + # The partial text is kept — it is evidence, not a result. + assert "I have started to writ" in result.stdout + assert result.agent_claim is None + + +def test_truncated_tool_call_is_not_dispatched(tmp_path: Path, chat_server: Any) -> None: + """A tool call cut off mid-emission may carry incomplete arguments.""" + (tmp_path / "foo.txt").write_text("hello world") + base_url, handler = chat_server( + [ + _finished(_message("", [_tool_call("1", "read_file", {"path": "foo.txt"})]), "length"), + _message("should never be reached."), + ] + ) + result = run_openai_agent( + "read foo.txt", + cwd=tmp_path, + capabilities=_ALL_CAPS, + timeout_s=30, + harness=replace(_harness(base_url), max_tokens=8), + ) + assert result.exit_status == -1 + assert len(handler.received) == 1 + + +def test_normal_stop_is_unaffected(tmp_path: Path, chat_server: Any) -> None: + base_url, _handler = chat_server([_finished(_message("All done."), "stop")]) + result = run_openai_agent( + "do the thing", + cwd=tmp_path, + capabilities=_ALL_CAPS, + timeout_s=30, + harness=replace(_harness(base_url), max_tokens=4096), + ) + assert result.exit_status == 0 + assert result.agent_claim == "All done."