Skip to content

feat(schema): add max_tokens field to AgentDef for per-agent override #471

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nskun (nskun) wants to merge 6 commits into
base: main
Choose a base branch
from
Open

feat(schema): add max_tokens field to AgentDef for per-agent override #471

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
cd9a25b
feat(schema): add max_tokens field to AgentDef for per-agent override
nskun Aug 20, 2026
2d5bf6e
style: format max_tokens schema tests
nskun Aug 20, 2026
c43e0f1
fix(config): reject unsupported per-agent max_tokens
nskun Aug 25, 2026
7268552
fix(config): complete per-agent max_tokens validation
nskun Aug 26, 2026
fb0cbb0
fix(config): clarify per-agent max token constraints
nskun Aug 26, 2026
40b3c1a
fix(claude): preserve explicit max token caps with reasoning
nskun Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions src/conductor/config/schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1529,6 +1529,13 @@ class AgentDef(BaseModel):
max_agent_iterations: 200 instead of using the default limit.
"""

max_tokens: int | None = Field(None, ge=1, le=200000)

@jrob5756 Jason Robert (jrob5756) Aug 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This validates clean on copilot, hermes, aca and claude-agent-sdk, and then has no effect on any of them. agent_builder.py:307 is the only read in the repo and it only runs under claude.

The two fields directly above this one are honoured by every provider (copilot:1296,1301, hermes:247,252, claude_agent_sdk:883,891, aca:763,764), so anyone reading AgentDef would expect the same treatment here.

Two ways to close it. Either wire up the remaining four providers, or declare the support explicitly and let validation refuse the rest. The second is what session_key and max_session_seconds already do:

# providers/capabilities.py
max_tokens: bool = False
"""``True`` when the provider applies a per-agent ``max_tokens`` output cap.

``False`` means the value would be silently ignored, so workflows that set it
fail validation instead."""
# config/validator.py, inside _check_agent_capabilities
if agent.max_tokens is not None and not caps.max_tokens:
    errors.append(
        f"Agent '{agent.name}' sets max_tokens={agent.max_tokens!r} but provider "
        f"'{provider_name}' does not apply per-agent output token caps "
        f"(capabilities.max_tokens=False). Remove it, use runtime.max_tokens where "
        f"the provider honours it, or override the agent to a provider that does."
    )

Then set max_tokens=True on ClaudeProvider and leave the default everywhere else. That also picks up claude-agent-sdk, which today refuses runtime.max_tokens at factory.py:195 but lets this one through.

"""Maximum output tokens per response for this agent.

Overrides the workflow-level runtime.max_tokens for this agent.
Only applies to provider-backed agents (not script or human_gate).

@jrob5756 Jason Robert (jrob5756) Aug 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

human_gate accepts max_tokens as things stand. I checked against a valid gate fixture and it takes the field without complaint, even though it rejects reasoning and session_key. questions accepts it too, which is odder still given it rejects model with the reason "no provider is invoked".

max_session_seconds and max_agent_iterations have the same hole, so this is inherited rather than introduced. But this is the line where max_tokens's contract gets written down, and it currently names the one type the code does not cover.

Suggested change
Overrides the workflow-level runtime.max_tokens for this agent.
Only applies to provider-backed agents (not script or human_gate).
Overrides the workflow-level runtime.max_tokens for this agent. Controls
response length, not the context window (that budget is context.max_tokens).
Rejected on script, workflow, wait, set, and terminate steps.

"""

session_key: Annotated[str, StringConstraints(strip_whitespace=True, min_length=1)] | None = (
None
)
Expand Down Expand Up @@ -2083,6 +2090,8 @@ def validate_agent_type(self) -> AgentDef:
raise ValueError("script agents cannot have 'max_session_seconds'")
if self.max_agent_iterations is not None:
raise ValueError("script agents cannot have 'max_agent_iterations'")
if self.max_tokens is not None:
raise ValueError("script agents cannot have 'max_tokens'")

@jrob5756 Jason Robert (jrob5756) Aug 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This block is repeated verbatim five times, and the copying is what let human_gate and questions slip through.

validate_agent_type already has a standalone-guard idiom for this, at lines 1929, 1944 and 1956. The comment on the stdin one states the reasoning outright: being a standalone guard rather than a per-branch check, it also covers the types that have no branch of their own.

The same shape replaces all five and closes the gap:

if self.type not in (None, "agent") and self.max_tokens is not None:
    raise ValueError(
        f"'{self.type}' agents cannot have 'max_tokens' "
        "(only provider-backed agents support this field)"
    )

It is a net reduction in lines, and nothing existing can break: extra="forbid" meant no workflow could carry the field at all before this PR, so the new rejection is strictly tighter than nothing. validator.py:1733 already defines _LLM_AGENT_TYPES = frozenset({None, "agent"}) if you would rather have one source of truth, and schema.py does not import validator.py, so there is no cycle.

Your five existing rejection tests still pass against this, since they only match on the field name.

if self.session_key is not None:
raise ValueError("script agents cannot have 'session_key'")
if self.retry is not None:
Expand Down Expand Up @@ -2139,6 +2148,8 @@ def validate_agent_type(self) -> AgentDef:
raise ValueError("workflow agents cannot have 'max_session_seconds'")
if self.max_agent_iterations is not None:
raise ValueError("workflow agents cannot have 'max_agent_iterations'")
if self.max_tokens is not None:
raise ValueError("workflow agents cannot have 'max_tokens'")
if self.session_key is not None:
raise ValueError("workflow agents cannot have 'session_key'")
if self.retry is not None:
Expand Down Expand Up @@ -2196,6 +2207,8 @@ def validate_agent_type(self) -> AgentDef:
raise ValueError("wait agents cannot have 'max_session_seconds'")
if self.max_agent_iterations is not None:
raise ValueError("wait agents cannot have 'max_agent_iterations'")
if self.max_tokens is not None:
raise ValueError("wait agents cannot have 'max_tokens'")
if self.session_key is not None:
raise ValueError("wait agents cannot have 'session_key'")
if self.retry is not None:
Expand Down Expand Up @@ -2269,6 +2282,8 @@ def validate_agent_type(self) -> AgentDef:
raise ValueError("set agents cannot have 'max_session_seconds'")
if self.max_agent_iterations is not None:
raise ValueError("set agents cannot have 'max_agent_iterations'")
if self.max_tokens is not None:
raise ValueError("set agents cannot have 'max_tokens'")
if self.session_key is not None:
raise ValueError("set agents cannot have 'session_key'")
if self.retry is not None:
Expand Down Expand Up @@ -2339,6 +2354,8 @@ def validate_agent_type(self) -> AgentDef:
raise ValueError("terminate agents cannot have 'max_session_seconds'")
if self.max_agent_iterations is not None:
raise ValueError("terminate agents cannot have 'max_agent_iterations'")
if self.max_tokens is not None:
raise ValueError("terminate agents cannot have 'max_tokens'")
if self.session_key is not None:
raise ValueError("terminate agents cannot have 'session_key'")
if self.max_depth is not None:
Expand Down
67 changes: 67 additions & 0 deletions tests/test_config/test_schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -667,6 +667,73 @@ def test_allowed_on_regular_agent(self) -> None:
assert agent.max_session_seconds == 90.0


class TestAgentDefMaxTokens:

@jrob5756 Jason Robert (jrob5756) Aug 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All eleven tests here exercise AgentDef.__init__. None of them reach the line that consumes the field, which was dead code until this PR. Delete agent_builder.py:307-308 and this suite still goes green.

tests/test_integration/test_parameter_flow_verification.py was written for this exact worry (its module docstring names it) and already covers the workflow-level value. The cheapest addition is two synchronous tests in test_pydantic_ai_agent_builder.py::TestSamplingSettings:

def test_agent_max_tokens_overrides_workflow_default(self) -> None:
    """A per-agent max_tokens must win over the workflow-level default."""
    agent_def = AgentDef(name="sampler", max_tokens=1000)
    pydantic_agent = build_agent(
        agent_def, system_prompt="", rendered_prompt="", default_max_tokens=4096
    )
    assert pydantic_agent.model_settings["max_tokens"] == 1000


def test_workflow_default_used_when_agent_max_tokens_unset(self) -> None:
    """With no per-agent override the workflow default still applies."""
    agent_def = AgentDef(name="sampler")
    pydantic_agent = build_agent(
        agent_def, system_prompt="", rendered_prompt="", default_max_tokens=4096
    )
    assert pydantic_agent.model_settings["max_tokens"] == 4096

I ran both against this branch and they pass.

One gap in the range coverage too: 0, -100 and 200001 are all tested, but the accepted endpoints 1 and 200000 are not, so swapping ge/le for gt/lt would go unnoticed. TestAgentDefMaxSessionSeconds has test_minimum_boundary for the same reason.

Worth adding a case for reasoning as well. max_tokens=1000 with reasoning.effort=low currently produces model_settings["max_tokens"] == 6144, which may be correct for the Anthropic API but is worth pinning so it cannot drift unnoticed.

"""Tests for max_tokens on AgentDef."""

def test_default_is_none(self) -> None:
"""Test that max_tokens defaults to None."""
agent = AgentDef(name="a", model="gpt-4", prompt="test")
assert agent.max_tokens is None

def test_accepts_valid_value(self) -> None:
"""Test that max_tokens accepts a valid value."""
agent = AgentDef(name="a", model="gpt-4", prompt="test", max_tokens=8192)
assert agent.max_tokens == 8192

def test_rejects_zero(self) -> None:
"""Test that max_tokens rejects zero."""
with pytest.raises(ValidationError) as exc_info:
AgentDef(name="a", model="gpt-4", prompt="test", max_tokens=0)
assert "greater than or equal to 1" in str(exc_info.value)

def test_rejects_negative(self) -> None:
"""Test that max_tokens rejects negative values."""
with pytest.raises(ValidationError) as exc_info:
AgentDef(name="a", model="gpt-4", prompt="test", max_tokens=-100)
assert "greater than or equal to 1" in str(exc_info.value)

def test_allowed_on_regular_agent(self) -> None:
"""Test that regular agents can have max_tokens."""
agent = AgentDef(name="a", type="agent", model="gpt-4", prompt="test", max_tokens=32768)
assert agent.max_tokens == 32768

def test_rejects_over_200000(self) -> None:
"""Test that max_tokens rejects values above 200000."""
with pytest.raises(ValidationError) as exc_info:
AgentDef(name="a", model="gpt-4", prompt="test", max_tokens=200001)
assert "less than or equal to 200000" in str(exc_info.value)

def test_rejected_on_script_agent(self) -> None:
"""Test that script agents cannot have max_tokens."""
with pytest.raises(ValidationError) as exc_info:
AgentDef(name="s", type="script", command="echo hi", max_tokens=8192)
assert "max_tokens" in str(exc_info.value)

@jrob5756 Jason Robert (jrob5756) Aug 22, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This assertion cannot fail for the reason it looks like it is checking. Pydantic v2 echoes the input dict into the error message, so "max_tokens" in str(exc_info.value) is true for any ValidationError raised on this input. Drop the command kwarg and it still passes, on a completely unrelated "script agents require 'command'" error.

pytest.raises is carrying the test on its own here. The range tests further up (line 687) already match on real message text, so this is inconsistent within the same class rather than a house style:

Suggested change
assert "max_tokens" in str(exc_info.value)
assert "script agents cannot have 'max_tokens'" in str(exc_info.value)

Same applies to the workflow, wait, set and terminate cases below.


def test_rejected_on_workflow_agent(self) -> None:
"""Test that workflow agents cannot have max_tokens."""
with pytest.raises(ValidationError) as exc_info:
AgentDef(name="w", type="workflow", workflow="./sub.yaml", max_tokens=8192)
assert "max_tokens" in str(exc_info.value)

def test_rejected_on_wait_agent(self) -> None:
"""Test that wait agents cannot have max_tokens."""
with pytest.raises(ValidationError) as exc_info:
AgentDef(name="w", type="wait", duration="5s", max_tokens=8192)
assert "max_tokens" in str(exc_info.value)

def test_rejected_on_set_agent(self) -> None:
"""Test that set agents cannot have max_tokens."""
with pytest.raises(ValidationError) as exc_info:
AgentDef(name="s", type="set", value="x", max_tokens=8192)
assert "max_tokens" in str(exc_info.value)

def test_rejected_on_terminate_agent(self) -> None:
"""Test that terminate agents cannot have max_tokens."""
with pytest.raises(ValidationError) as exc_info:
AgentDef(name="t", type="terminate", status="success", reason="done", max_tokens=8192)
assert "max_tokens" in str(exc_info.value)


class TestRuntimeConfig:
"""Tests for RuntimeConfig model."""

Expand Down
Loading