diff --git a/benchmarks/gdpval/config.yaml b/benchmarks/gdpval/config.yaml
index 04030acb9a9..c22c63fc69c 100644
--- a/benchmarks/gdpval/config.yaml
+++ b/benchmarks/gdpval/config.yaml
@@ -103,6 +103,15 @@ gdpval_stirrup_agent:
agent_max_turns: 250
concurrency: 32
temperature: 1.0
+ # GDPVal tasks routinely need multi-thousand-token scripts. A 1,024-token
+ # completion cannot form a useful tool call and can create a sticky
+ # code_exec({}) loop after long reasoning turns.
+ context_window_tokens: 262144
+ min_completion_tokens: 8192
+ # Reject the observed near-empty compaction outputs that erase progress
+ # on long GDPVal sessions. Generic Stirrup agents accept any non-empty
+ # summary for backward compatibility.
+ min_compaction_summary_words: 50
system_prompt_template: ${oc.env:SYSTEM_PROMPT_TEMPLATE,null}
user_prompt_template: ${oc.env:USER_PROMPT_TEMPLATE,null}
gdpval_container_path: ${oc.env:GDPVAL_CONTAINER_PATH,null}
diff --git a/responses_api_agents/stirrup_agent/README.md b/responses_api_agents/stirrup_agent/README.md
index 68b71fe0603..87b34d88800 100644
--- a/responses_api_agents/stirrup_agent/README.md
+++ b/responses_api_agents/stirrup_agent/README.md
@@ -131,6 +131,10 @@ The agent reads its Hydra config at `configs/stirrup_agent.yaml`. Notable keys:
| `rerun_incomplete` | `false` | If true, skip the rollout for tasks that already **finished** (a finish marker is cached) and only re-run the ones that didn't (see [Task Re-run Mode](#task-re-run-mode)). Composes with `execute_only` and `judge_only`. |
| `model_id` | `null` | HF model id or local path used to load a tokenizer for dynamic output sizing. |
| `completion_token_buffer` | `1000` | Safety margin (in tokens) reserved when sizing `max_completion_tokens` per call. |
+| `context_window_tokens` | `262144` | Model context-window size used for dynamic completion budgeting. Independent of a request's `max_output_tokens`. |
+| `min_completion_tokens` | `1024` | Target useful per-call completion floor. The hard cap always applies; estimated remaining context is also a strict bound when the tokenizer successfully renders the complete prompt. Approximate fallbacks preserve this floor even when their estimate exceeds context. GDPVal raises it to `8192`. |
+| `prompt_estimator_truncate_history_thinking` | `null` | Optional estimator-only switch for templates that omit completed historical reasoning. It is never sent to the provider. |
+| `min_compaction_summary_words` | `1` | Minimum summary length accepted during context compaction. GDPVal uses `50`; invalid summaries retry up to three times and never replace the live history. |
Env vars honored: `TAVILY_API_KEY`, `HF_TOKEN`, `OPENAI_API_KEY`.
@@ -155,13 +159,28 @@ The wrapper ships a `DynamicMaxTokensChatCompletionsClient`
Set `model_id` to the same checkpoint (or HF id) you are serving via
vLLM and the tokeniser match is exact. Leave it unset and the client
-falls back to a conservative character-count estimate — slower to
-allocate completion budget but always safe. `completion_token_buffer`
+falls back to a character-count estimate. When
+`prompt_estimator_truncate_history_thinking=true`, that fallback mirrors
+Nemotron's prior-assistant transformation without changing the retained
+trajectory or provider request. `completion_token_buffer`
absorbs the residual gap between our estimate and the exact prompt the
server renders (chat-template wrappers, tool-schema injection). The
default `1000` works in practice; raise it (e.g. 2000–5000) if you see
sporadic HTTP 400 responses at the vLLM proxy.
+When the loaded tokenizer successfully renders the complete prompt, estimated
+remaining context is a strict output bound and a prompt with no remaining room
+fails before dispatch. If full rendering is unsupported, JSON and character
+estimates can substantially overcount retained reasoning, so the client still
+dispatches the configured completion floor; the hard completion cap applies in
+both modes.
+
+GDPVal uses an `8192`-token minimum because a 1024-token truncation cannot
+reliably produce a complete script or tool call. Generic Stirrup agents keep
+the historical 1024-token default. Launchers serving a local checkpoint should
+still pass that checkpoint as `model_id`; the fallback is a resilience path,
+not a replacement for template-exact tokenization.
+
## Advanced Features
NeMo Gym splits a GDPVal evaluation into two halves — *executing* the task
diff --git a/responses_api_agents/stirrup_agent/app.py b/responses_api_agents/stirrup_agent/app.py
index 9937b103206..5d376e3b053 100644
--- a/responses_api_agents/stirrup_agent/app.py
+++ b/responses_api_agents/stirrup_agent/app.py
@@ -475,7 +475,7 @@ async def _run_stirrup_agent(
api_key: str = "dummy",
max_turns: int = 250,
temperature: float = 0.6,
- max_tokens: int = 262144,
+ context_window_tokens: int = 262144,
reference_files: Optional[list] = None,
reference_file_urls: Optional[list] = None,
exec_provider_class: Optional[str] = None,
@@ -489,7 +489,10 @@ async def _run_stirrup_agent(
top_p: float = 0.95,
enable_thinking: bool = True,
max_completion_tokens_cap: int = 64000,
+ min_completion_tokens: int = 1024,
+ prompt_estimator_truncate_history_thinking: Optional[bool] = None,
truncation_recovery: bool = True,
+ min_compaction_summary_words: int = 1,
tavily_api_key: Optional[Union[str, List[str]]] = None,
tavily_max_sweeps: int = 1,
) -> Dict[str, Any]:
@@ -502,6 +505,9 @@ async def _run_stirrup_agent(
a tokenizer that sizes ``max_completion_tokens`` dynamically per call
(see ``DynamicMaxTokensChatCompletionsClient``). When unset, a
character-count fallback is used.
+
+ *context_window_tokens* describes model capacity. It is deliberately
+ separate from the outer Responses API request's ``max_output_tokens``.
"""
from stirrup.tools import DEFAULT_TOOLS
from stirrup.tools.code_backends.base import SHELL_TIMEOUT, CodeExecToolProvider, CommandResult
@@ -550,13 +556,15 @@ async def run_command(self, cmd: str, *, timeout: int = SHELL_TIMEOUT) -> Comman
model=model_name,
base_url=model_base_url,
api_key=api_key,
- max_tokens=max_tokens,
+ max_tokens=context_window_tokens,
model_id=model_id,
completion_token_buffer=completion_token_buffer,
temperature=temperature,
top_p=top_p,
enable_thinking=enable_thinking,
max_completion_tokens_cap=max_completion_tokens_cap,
+ min_completion_tokens=min_completion_tokens,
+ prompt_estimator_truncate_history_thinking=prompt_estimator_truncate_history_thinking,
truncation_recovery=truncation_recovery,
)
@@ -609,6 +617,7 @@ async def run_command(self, cmd: str, *, timeout: int = SHELL_TIMEOUT) -> Comman
"tools": tools,
"tool_response_as_user": True,
"skip_input_file_listing": is_gdpval,
+ "min_compaction_summary_words": min_compaction_summary_words,
}
if system_prompt:
agent_kwargs["system_prompt"] = system_prompt
@@ -969,8 +978,8 @@ class StirrupAgentWrapperConfig(BaseResponsesAPIAgentConfig):
default=None,
description="HuggingFace model ID (or local checkpoint path) used to load a tokenizer "
"for dynamic max_completion_tokens sizing. E.g. 'Qwen/Qwen3-Coder-30B-A3B-Instruct'. "
- "When None, a character-count fallback is used (conservative, but slightly over-allocates "
- "input tokens). See ``nemo_client.DynamicMaxTokensChatCompletionsClient``.",
+ "When None, an approximate character-count fallback is used; it can substantially "
+ "overcount retained reasoning. See ``nemo_client.DynamicMaxTokensChatCompletionsClient``.",
)
completion_token_buffer: int = Field(
default=1000,
@@ -979,6 +988,13 @@ class StirrupAgentWrapperConfig(BaseResponsesAPIAgentConfig):
"(messages + tool-schema JSON) and the exact prompt the server sees after chat-template "
"rendering. See ``nemo_client.DynamicMaxTokensChatCompletionsClient``.",
)
+ context_window_tokens: int = Field(
+ default=262144,
+ ge=1,
+ description="Model context-window size used for dynamic per-call completion budgeting. "
+ "This is independent of the Responses API request's max_output_tokens, which limits "
+ "output rather than describing model capacity.",
+ )
top_p: float = Field(
default=0.95,
description="Top-p sampling cutoff for the policy model. Forwarded to the LLM client.",
@@ -990,10 +1006,25 @@ class StirrupAgentWrapperConfig(BaseResponsesAPIAgentConfig):
)
max_completion_tokens_cap: int = Field(
default=64000,
+ ge=1,
description="Hard ceiling on per-call ``max_completion_tokens``. Dynamic sizing computes "
"context_window - input_tokens - completion_token_buffer, then caps to this value. "
"Set to match the training-side response-length budget for RL.",
)
+ min_completion_tokens: int = Field(
+ default=1024,
+ ge=1,
+ description="Target minimum per-call ``max_completion_tokens``. The hard cap always applies; "
+ "estimated remaining context is also a strict bound when the loaded tokenizer successfully "
+ "renders the complete prompt. The "
+ "approximate fallback preserves this floor even when its estimate exceeds context. "
+ "Long-horizon benchmarks can opt into a larger usable floor.",
+ )
+ prompt_estimator_truncate_history_thinking: Optional[bool] = Field(
+ default=None,
+ description="Optional prompt-estimator setting for checkpoints whose template omits historical "
+ "reasoning before the last user turn. This is never forwarded to the model request.",
+ )
truncation_recovery: bool = Field(
default=True,
description="When a model call spends its entire completion budget without emitting a "
@@ -1004,6 +1035,13 @@ class StirrupAgentWrapperConfig(BaseResponsesAPIAgentConfig):
"deliverable writes (median 34.6k tokens). The instruction is sent to the server but not "
"recorded in the trajectory; set False for RL rollouts that require an unsteered policy.",
)
+ min_compaction_summary_words: int = Field(
+ default=1,
+ ge=1,
+ description="Minimum whitespace-delimited word count accepted from context compaction. "
+ "The generic default rejects only empty output; long-horizon benchmarks can require a "
+ "more complete summary.",
+ )
tavily_api_key: Optional[Union[str, List[str]]] = Field(
default=None,
description="Tavily API key(s) for the ``web_search`` / ``fetch_web_page`` tools. "
@@ -1163,7 +1201,10 @@ async def responses(
model_name = getattr(body, "model", None) or "default"
temperature = getattr(body, "temperature", None) or self.config.temperature
- max_tokens = getattr(body, "max_output_tokens", 262144) or 262144
+ requested_output_tokens = getattr(body, "max_output_tokens", None)
+ max_completion_tokens_cap = self.config.max_completion_tokens_cap
+ if requested_output_tokens is not None:
+ max_completion_tokens_cap = min(max_completion_tokens_cap, requested_output_tokens)
exec_provider = self.task_strategy.get_exec_provider(task_info, self.config)
exec_provider_class = None
@@ -1183,7 +1224,7 @@ async def responses(
"api_key": "dummy", # pragma: allowlist secret
"max_turns": self.config.agent_max_turns,
"temperature": temperature,
- "max_tokens": max_tokens,
+ "context_window_tokens": self.config.context_window_tokens,
"reference_files": task_info.get("reference_files") if self.config.task == "gdpval" else None,
"reference_file_urls": task_info.get("reference_file_urls") if self.config.task == "gdpval" else None,
"exec_provider_class": exec_provider_class,
@@ -1196,8 +1237,11 @@ async def responses(
"completion_token_buffer": self.config.completion_token_buffer,
"top_p": getattr(body, "top_p", None) or self.config.top_p,
"enable_thinking": self.config.enable_thinking,
- "max_completion_tokens_cap": self.config.max_completion_tokens_cap,
+ "max_completion_tokens_cap": max_completion_tokens_cap,
+ "min_completion_tokens": self.config.min_completion_tokens,
+ "prompt_estimator_truncate_history_thinking": (self.config.prompt_estimator_truncate_history_thinking),
"truncation_recovery": self.config.truncation_recovery,
+ "min_compaction_summary_words": self.config.min_compaction_summary_words,
"tavily_api_key": self.config.tavily_api_key,
"tavily_max_sweeps": self.config.tavily_max_sweeps,
}
diff --git a/responses_api_agents/stirrup_agent/nemo_agent.py b/responses_api_agents/stirrup_agent/nemo_agent.py
index 20d83562a8e..e432ce54204 100644
--- a/responses_api_agents/stirrup_agent/nemo_agent.py
+++ b/responses_api_agents/stirrup_agent/nemo_agent.py
@@ -11,7 +11,7 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
-"""NeMoAgent — Stirrup ``Agent`` subclass with two behavioural overrides.
+"""NeMoAgent — Stirrup ``Agent`` subclass with Gym-specific behaviour.
1. **tool_response_as_user** — Return a ``UserMessage`` (not ``ToolMessage``)
from ``run_tool()`` so the conversation history presents tool results
@@ -23,6 +23,10 @@
injects into the system prompt; useful when a task prompt already lists
its own reference files (e.g. GDPVal).
+3. **safe context compaction** — Removes inline reasoning from generated
+ summaries and retries empty or degenerate summaries before replacing the
+ live history.
+
To preserve tool-call metadata that ``Agent.step()`` reads immediately after
``run_tool()`` returns (``.name``, ``.success``, ``.tool_call_id``), the
conversion uses :class:`NeMoUserMessage` — a ``UserMessage`` subclass with
@@ -32,12 +36,69 @@
from __future__ import annotations
+import logging
+import re
+from itertools import takewhile
from typing import Any, Optional
from pydantic import ConfigDict
from stirrup import Agent
from stirrup.core.agent import SessionAgent
-from stirrup.core.models import ToolCall, ToolMessage, UserMessage
+from stirrup.core.exceptions import ContextOverflowError
+from stirrup.core.models import AssistantMessage, ChatMessage, SummaryMessage, ToolCall, ToolMessage, UserMessage
+from stirrup.prompts import MESSAGE_SUMMARIZER, MESSAGE_SUMMARIZER_BRIDGE_TEMPLATE
+
+
+LOGGER = logging.getLogger(__name__)
+
+_MAX_COMPACTION_SUMMARY_ATTEMPTS = 3
+_LEADING_THINK_OPEN_RE = re.compile(r"\A\s*", re.IGNORECASE)
+_THINK_CLOSE_RE = re.compile(r"", re.IGNORECASE)
+
+
+def _sanitize_compaction_summary(content: Any) -> str:
+ """Remove leading inline-thinking blocks from a generated summary.
+
+ Only a leading block is special: a later literal ```` may be part
+ of task content and must survive. An unclosed leading block has no safe
+ answer portion, so it is treated as empty and retried.
+ """
+ if not isinstance(content, str):
+ return ""
+
+ sanitized = content.strip()
+ while match := _LEADING_THINK_OPEN_RE.match(sanitized):
+ closing = _THINK_CLOSE_RE.search(sanitized, match.end())
+ if closing is None:
+ return ""
+ sanitized = sanitized[closing.end() :].lstrip()
+ return sanitized.strip()
+
+
+def _is_meaningful_compaction_summary(summary: str, *, min_words: int = 1) -> bool:
+ """Return whether a summary is large enough to replace a full context."""
+ return len(summary.split()) >= min_words
+
+
+def _summary_text_only_instruction(min_words: int) -> str:
+ minimum = (
+ f"The summary must contain at least {min_words} words."
+ if min_words > 1
+ else "The summary must contain useful, non-empty text."
+ )
+ return (
+ "Return a concise but comprehensive summary as text only. "
+ "Do not call tools. Do not include analysis, reasoning, or tags. "
+ f"{minimum}"
+ )
+
+
+def _summary_rejection_reason(content: Any, sanitized: str, min_words: int) -> str:
+ if not isinstance(content, str):
+ return f"non-text content ({type(content).__name__})"
+ if not sanitized:
+ return "empty after removing leading reasoning"
+ return f"{len(sanitized.split())} words, below configured minimum {min_words}"
class NeMoUserMessage(UserMessage):
@@ -76,11 +137,15 @@ def __init__(
*,
tool_response_as_user: bool = False,
skip_input_file_listing: bool = False,
+ min_compaction_summary_words: int = 1,
**kwargs: Any,
) -> None:
+ if min_compaction_summary_words < 1:
+ raise ValueError("min_compaction_summary_words must be at least 1")
super().__init__(**kwargs)
self._tool_response_as_user = tool_response_as_user
self._skip_input_file_listing = skip_input_file_listing
+ self._min_compaction_summary_words = min_compaction_summary_words
def _build_system_prompt(self) -> str:
"""Override to optionally skip the input file listing."""
@@ -103,6 +168,101 @@ def _build_system_prompt(self) -> str:
return result
+ async def summarize_messages(
+ self,
+ messages: list[ChatMessage],
+ run_metadata_by_turn: dict[str, dict[str, list[Any]]],
+ ) -> tuple[list[ChatMessage], list[ChatMessage]]:
+ """Compact history without replaying reasoning or erasing progress.
+
+ Stirrup 0.1 inserts ``summary.content`` into the next context verbatim.
+ Some reasoning-model endpoints return a leading ```` block in
+ that field, including think-only responses. Sanitize at the adapter
+ boundary and retry bad summaries before replacing any history.
+ """
+ current_messages = messages
+ while True:
+ try:
+ task_context: list[ChatMessage] = list(
+ takewhile(
+ lambda message: not isinstance(message, (AssistantMessage, SummaryMessage)),
+ current_messages,
+ )
+ )
+
+ tool_docs = "\n".join(f"- {tool.name}: {tool.description}" for tool in self._active_tools.values())
+ strict_prompt = (
+ f"{MESSAGE_SUMMARIZER}\n\n{_summary_text_only_instruction(self._min_compaction_summary_words)}"
+ )
+ no_tools_prompt = f"{strict_prompt}\n\nTools are disabled for this response." + (
+ f" Tools used earlier were:\n{tool_docs}" if tool_docs else ""
+ )
+ attempts = (
+ (MESSAGE_SUMMARIZER, self._active_tools),
+ (strict_prompt, self._active_tools),
+ (no_tools_prompt, {}),
+ )
+
+ summary_content = ""
+ for attempt, (prompt, tools) in enumerate(attempts, start=1):
+ # Compaction uses the policy client internally, but its
+ # one-shot truncation recovery belongs to policy turns.
+ # Do not let a summary consume an armed recovery or arm a
+ # recovery that leaks into the next policy turn.
+ sentinel = object()
+ saved_recovery = getattr(self._client, "_recover_from_truncation", sentinel)
+ saved_overruns = getattr(self._client, "_truncation_overruns", sentinel)
+ if saved_recovery is not sentinel:
+ self._client._recover_from_truncation = False
+ try:
+ summary = await self._client.generate(
+ [*current_messages, UserMessage(content=prompt)],
+ tools,
+ )
+ finally:
+ if saved_recovery is not sentinel:
+ self._client._recover_from_truncation = saved_recovery
+ if saved_overruns is not sentinel:
+ self._client._truncation_overruns = saved_overruns
+ summary_content = _sanitize_compaction_summary(summary.content)
+ if _is_meaningful_compaction_summary(
+ summary_content,
+ min_words=self._min_compaction_summary_words,
+ ):
+ break
+ usage = getattr(summary, "token_usage", None)
+ LOGGER.warning(
+ "Compaction summary attempt %d/%d rejected: %s (input=%s answer=%s reasoning=%s); %s",
+ attempt,
+ _MAX_COMPACTION_SUMMARY_ATTEMPTS,
+ _summary_rejection_reason(
+ summary.content,
+ summary_content,
+ self._min_compaction_summary_words,
+ ),
+ getattr(usage, "input", None),
+ getattr(usage, "answer", None),
+ getattr(usage, "reasoning", None),
+ "retrying" if attempt < _MAX_COMPACTION_SUMMARY_ATTEMPTS else "no attempts remain",
+ )
+ else:
+ raise RuntimeError(
+ "Compaction summarizer produced no usable text after "
+ f"{_MAX_COMPACTION_SUMMARY_ATTEMPTS} attempts; refusing to erase context"
+ )
+
+ summary_bridge_prompt = MESSAGE_SUMMARIZER_BRIDGE_TEMPLATE.format(summary=summary_content)
+ summary_bridge = SummaryMessage(content=summary_bridge_prompt)
+ acknowledgement_msg = UserMessage(content="Got it, thanks!")
+
+ self._logger.context_summarization_complete(summary_content, summary_bridge_prompt)
+ return current_messages, [*task_context, summary_bridge, acknowledgement_msg]
+ except ContextOverflowError:
+ # Preserve Stirrup's recovery contract: discard the latest
+ # completed turn, including its per-turn metadata, and retry.
+ current_messages, dropped_turn_id = self._unwind_context_overflow(current_messages)
+ run_metadata_by_turn.pop(dropped_turn_id, None)
+
async def run_tool(self, tool_call: ToolCall, run_metadata: dict[str, list[Any]]) -> ToolMessage:
"""Run a tool and optionally return a ``NeMoUserMessage`` instead of ``ToolMessage``.
diff --git a/responses_api_agents/stirrup_agent/nemo_client.py b/responses_api_agents/stirrup_agent/nemo_client.py
index 790ec2df482..1749044d2ba 100644
--- a/responses_api_agents/stirrup_agent/nemo_client.py
+++ b/responses_api_agents/stirrup_agent/nemo_client.py
@@ -32,18 +32,22 @@
context_window − tokenized(messages) − completion_token_buffer
-clamped to a minimum of ``_MIN_COMPLETION_TOKENS``. On the response
+raised toward a configurable minimum (historically 1,024 tokens), subject to
+the hard cap. Estimated remaining context is also a strict bound when the
+tokenizer successfully renders the complete prompt; approximate fallbacks keep
+the configured floor because they can substantially overcount retained reasoning. On the response
side, we replicate Stirrup parsing but do *not* raise on
``finish_reason=length`` — the agent loop will either terminate when the
model invokes the ``finish`` tool or exhaust ``max_turns``, yielding a
clean timeout instead of a crash.
``model_id`` selects the HuggingFace tokenizer (or local checkpoint path).
-When unset, a conservative character-count fallback is used.
+When unset, an approximate character-count fallback is used.
"""
from __future__ import annotations
+import json
import logging
from time import perf_counter
from typing import Any, Optional
@@ -52,6 +56,7 @@
from pydantic import ValidationError as _PydanticValidationError
from stirrup.clients.chat_completions_client import ChatCompletionsClient
from stirrup.clients.utils import to_openai_tools
+from stirrup.core.exceptions import ContextOverflowError
from stirrup.core.models import (
AssistantMessage,
ChatMessage,
@@ -153,8 +158,9 @@ def _install_coercing_finish_tool() -> None:
_install_coercing_finish_tool()
-# Floor for per-call max_completion_tokens. Below this the model basically
-# cannot produce a useful answer — treat as a hard minimum.
+# Target floor for per-call max_completion_tokens. Below this the model usually
+# cannot produce a useful answer. The hard cap always takes precedence; exact
+# tokenizer estimates also enforce remaining context as a strict bound.
_MIN_COMPLETION_TOKENS = 1024
# Hard cap on per-call max_completion_tokens. Oversized completion budgets
@@ -231,15 +237,23 @@ def __init__(
top_p: float = 0.95,
enable_thinking: bool = True,
max_completion_tokens_cap: int = _DEFAULT_MAX_COMPLETION_TOKENS_CAP,
+ min_completion_tokens: int = _MIN_COMPLETION_TOKENS,
+ prompt_estimator_truncate_history_thinking: Optional[bool] = None,
truncation_recovery: bool = True,
**kwargs: Any,
) -> None:
super().__init__(*args, **kwargs)
+ if max_completion_tokens_cap < 1:
+ raise ValueError("max_completion_tokens_cap must be at least 1")
+ if min_completion_tokens < 1:
+ raise ValueError("min_completion_tokens must be at least 1")
self._completion_token_buffer = completion_token_buffer
self._temperature = temperature
self._top_p = top_p
self._enable_thinking = enable_thinking
self._max_completion_tokens_cap = max_completion_tokens_cap
+ self._min_completion_tokens = min_completion_tokens
+ self._prompt_estimator_truncate_history_thinking = prompt_estimator_truncate_history_thinking
self._truncation_recovery = truncation_recovery
# Set when the previous call exhausted its completion budget without
# emitting a tool call; consumed by the very next generate().
@@ -256,11 +270,141 @@ def __init__(
# Token counting
# ------------------------------------------------------------------
+ @staticmethod
+ def _truncate_prior_assistant_text_for_estimate(content: str, *, has_tool_calls: bool) -> str:
+ """Mirror Nemotron's historical-assistant rendering for one text value."""
+ # The template first inserts an empty reasoning block when neither tag
+ # is present. This matters for both its normal and tool-call branches.
+ if "" not in content and "" not in content:
+ content = f"{content}"
+
+ if has_tool_calls:
+ if not content.strip():
+ return ""
+ if "" in content:
+ content = content.rsplit("", 1)[-1]
+ elif "" in content:
+ content = content.split("", 1)[0]
+ return f"{content}"
+
+ if "" in content and "" in content:
+ content = f"{content.rsplit('', 1)[-1]}"
+ return content.strip()
+
+ def _messages_for_estimator_count(
+ self,
+ messages: list[NeMoGymChatCompletionMessageParam],
+ ) -> list[NeMoGymChatCompletionMessageParam]:
+ """Return the history shape rendered by Nemotron's truncation mode.
+
+ This is estimator-only: it never changes the retained trajectory or
+ request. Tool calls and every other serialized field stay intact.
+ """
+ if self._prompt_estimator_truncate_history_thinking is not True:
+ return messages
+
+ last_user_index = max(
+ (index for index, message in enumerate(messages) if message.get("role") == "user"),
+ default=-1,
+ )
+ if last_user_index < 0:
+ return messages
+
+ counted_messages: list[NeMoGymChatCompletionMessageParam] = []
+ for index, message in enumerate(messages):
+ if index >= last_user_index or message.get("role") != "assistant":
+ counted_messages.append(message)
+ continue
+
+ content = message.get("content")
+ estimated_content: Any = content
+ has_tool_calls = bool(message.get("tool_calls"))
+ if isinstance(content, str):
+ estimated_content = self._truncate_prior_assistant_text_for_estimate(
+ content,
+ has_tool_calls=has_tool_calls,
+ )
+ elif isinstance(content, list):
+ estimated_parts = []
+ changed = False
+ for part in content:
+ if isinstance(part, dict) and isinstance(part.get("text"), str):
+ estimated_text = self._truncate_prior_assistant_text_for_estimate(
+ part["text"],
+ has_tool_calls=has_tool_calls,
+ )
+ if estimated_text != part["text"]:
+ part = {**part, "text": estimated_text}
+ changed = True
+ estimated_parts.append(part)
+ if changed:
+ estimated_content = estimated_parts
+
+ if estimated_content is content or estimated_content == content:
+ counted_messages.append(message)
+ else:
+ counted_messages.append({**message, "content": estimated_content})
+ return counted_messages
+
+ def _messages_for_template_count(
+ self,
+ messages: list[NeMoGymChatCompletionMessageParam],
+ ) -> list[NeMoGymChatCompletionMessageParam]:
+ """Mirror vLLM's tool-argument decoding before template rendering.
+
+ OpenAI history carries ``function.arguments`` as a JSON string, while
+ Nemotron's Jinja template iterates it as a mapping. vLLM decodes that
+ field before rendering; HuggingFace ``apply_chat_template`` does not.
+ Keep this as a fallback template-input variant so templates that
+ natively accept OpenAI's string representation retain their existing
+ behavior. The provider payload is never mutated.
+ """
+ normalized_messages: list[NeMoGymChatCompletionMessageParam] = []
+ messages_changed = False
+ for message in messages:
+ tool_calls = message.get("tool_calls")
+ if not isinstance(tool_calls, list):
+ normalized_messages.append(message)
+ continue
+
+ normalized_calls = []
+ changed = False
+ for tool_call in tool_calls:
+ if not isinstance(tool_call, dict):
+ normalized_calls.append(tool_call)
+ continue
+ function = tool_call.get("function")
+ if not isinstance(function, dict) or not isinstance(function.get("arguments"), str):
+ normalized_calls.append(tool_call)
+ continue
+ try:
+ arguments = json.loads(function["arguments"])
+ except json.JSONDecodeError:
+ normalized_calls.append(tool_call)
+ continue
+ if not isinstance(arguments, dict):
+ normalized_calls.append(tool_call)
+ continue
+ normalized_calls.append({**tool_call, "function": {**function, "arguments": arguments}})
+ changed = True
+
+ normalized_messages.append({**message, "tool_calls": normalized_calls} if changed else message)
+ messages_changed = messages_changed or changed
+ return normalized_messages if messages_changed else messages
+
def _count_input_tokens(
self,
messages: list[NeMoGymChatCompletionMessageParam],
tools: Optional[dict[str, Tool]] = None,
) -> int:
+ """Return the best available prompt-token estimate."""
+ return self._count_input_tokens_with_confidence(messages, tools)[0]
+
+ def _count_input_tokens_with_confidence(
+ self,
+ messages: list[NeMoGymChatCompletionMessageParam],
+ tools: Optional[dict[str, Tool]] = None,
+ ) -> tuple[int, bool]:
"""Estimate the full prompt token count the server will see.
``messages`` must already be serialized for the provider. This keeps
@@ -279,19 +423,26 @@ def _count_input_tokens(
rough but serialises everything.
4. Character-count fallback when no tokenizer is present.
+ Returns ``(count, exact_template_render)``. Only a successful render of
+ the complete prompt is exact enough to impose a hard context bound.
+ JSON and character fallbacks remain useful for budget sizing, but must
+ not turn an approximate over-count into a false context-overflow error.
+
Any residual gap is absorbed by ``completion_token_buffer``.
"""
import json as _json
if self._tokenizer is None:
- # Pure character-count fallback.
- total = sum(len(str(m.get("content") or "")) for m in messages) // 3
+ # Pure character-count fallback. Count the complete serialized
+ # payload, including assistant tool-call names and arguments.
+ counted_messages = self._messages_for_estimator_count(messages)
+ total = len(_json.dumps(counted_messages, ensure_ascii=False)) // 3
if tools:
try:
total += len(_json.dumps(to_openai_tools(tools))) // 3
except Exception:
pass
- return total
+ return total, False
oai_tools = None
if tools:
@@ -301,38 +452,68 @@ def _count_input_tokens(
LOGGER.warning(f"to_openai_tools failed ({exc}).")
# Strategy 1: apply_chat_template with tools=
+ template_kwargs: dict[str, Any] = {}
+ if self._prompt_estimator_truncate_history_thinking is not None:
+ template_kwargs["truncate_history_thinking"] = self._prompt_estimator_truncate_history_thinking
+ normalized_template_messages = self._messages_for_template_count(messages)
+ template_message_variants = [messages]
+ if normalized_template_messages is not messages:
+ template_message_variants.append(normalized_template_messages)
if oai_tools is not None:
+ last_template_error = None
+ for template_messages in template_message_variants:
+ try:
+ text = self._tokenizer.apply_chat_template(
+ template_messages,
+ tools=oai_tools,
+ tokenize=False,
+ add_generation_prompt=True,
+ **template_kwargs,
+ )
+ return len(self._tokenizer(text, add_special_tokens=False)["input_ids"]), True
+ except Exception as exc:
+ last_template_error = exc
+ LOGGER.debug(
+ "apply_chat_template(tools=) unsupported (%s); trying separate tool count.",
+ last_template_error,
+ )
+
+ # Strategy 2: apply_chat_template on messages only + separate tool JSON count
+ last_template_error = None
+ for template_messages in template_message_variants:
try:
text = self._tokenizer.apply_chat_template(
- messages, tools=oai_tools, tokenize=False, add_generation_prompt=True
+ template_messages,
+ tokenize=False,
+ add_generation_prompt=True,
+ **template_kwargs,
)
- return len(self._tokenizer(text, add_special_tokens=False)["input_ids"])
+ total = len(self._tokenizer(text, add_special_tokens=False)["input_ids"])
+ if oai_tools is not None:
+ total += len(self._tokenizer(_json.dumps(oai_tools), add_special_tokens=False)["input_ids"])
+ # With no tools this is the complete rendered prompt. When tools
+ # exist, tokenizing their JSON separately does not reproduce the
+ # template's wrappers or token boundaries and remains approximate.
+ return total, oai_tools is None
except Exception as exc:
- LOGGER.debug(f"apply_chat_template(tools=) unsupported ({exc}); trying separate tool count.")
-
- # Strategy 2: apply_chat_template on messages only + separate tool JSON count
- try:
- text = self._tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
- total = len(self._tokenizer(text, add_special_tokens=False)["input_ids"])
- if oai_tools is not None:
- total += len(self._tokenizer(_json.dumps(oai_tools), add_special_tokens=False)["input_ids"])
- return total
- except Exception as exc:
- LOGGER.warning(f"apply_chat_template(messages) failed ({exc}); falling back to JSON count.")
+ last_template_error = exc
+ LOGGER.warning("apply_chat_template(messages) failed (%s); falling back to JSON count.", last_template_error)
# Strategy 3: tokenise the full JSON payload
try:
- blob = _json.dumps(messages)
+ counted_messages = self._messages_for_estimator_count(messages)
+ blob = _json.dumps(counted_messages)
total = len(self._tokenizer(blob, add_special_tokens=False)["input_ids"])
if oai_tools is not None:
total += len(self._tokenizer(_json.dumps(oai_tools), add_special_tokens=False)["input_ids"])
- return total
+ return total, False
except Exception as exc:
LOGGER.warning(f"JSON tokenisation failed ({exc}); falling back to character count.")
# Strategy 4: character count
- total = sum(len(str(m.get("content") or "")) for m in messages) // 3
- return total
+ counted_messages = self._messages_for_estimator_count(messages)
+ total = len(_json.dumps(counted_messages, ensure_ascii=False)) // 3
+ return total, False
async def generate(
self,
@@ -352,13 +533,20 @@ async def generate(
if recovering:
provider_messages = [*provider_messages, {"role": "user", "content": _TRUNCATION_RECOVERY_NUDGE}]
- input_tokens = self._count_input_tokens(provider_messages, tools)
+ input_tokens, exact_template_render = self._count_input_tokens_with_confidence(provider_messages, tools)
context_window = self._max_tokens
+ estimated_remaining_context = context_window - input_tokens
+ if exact_template_render and estimated_remaining_context <= 0:
+ raise ContextOverflowError(
+ f"Estimated prompt ({input_tokens} tokens) leaves no room in the {context_window}-token context window"
+ )
dynamic_max = max(
- context_window - input_tokens - self._completion_token_buffer,
- _MIN_COMPLETION_TOKENS,
+ estimated_remaining_context - self._completion_token_buffer,
+ self._min_completion_tokens,
)
capped_max = min(dynamic_max, self._max_completion_tokens_cap)
+ if exact_template_render:
+ capped_max = min(capped_max, estimated_remaining_context)
# ``self._kwargs`` is spread last so explicit per-request kwargs override
# the agent-level defaults.
@@ -447,15 +635,30 @@ async def generate(
for tc in (msg.tool_calls or [])
]
- # A call that exhausted its completion budget *and* produced no tool call
- # advanced the task by nothing. Arm the recovery turn (see
- # _TRUNCATION_RECOVERY_NUDGE). A truncated call that still carried a tool
- # call made progress and is left alone.
- if choice.finish_reason in ("length", "max_tokens") and not tool_calls:
+ # A call that exhausted its completion budget without a schema-valid
+ # tool call advanced the task by nothing. Parser-truncated calls often
+ # surface as ``code_exec({})``; treating their mere presence as progress
+ # creates a sticky invalid-call loop.
+ usable_tool_call = False
+ for tool_call in tool_calls:
+ tool = tools.get(tool_call.name)
+ if tool is None:
+ # We cannot safely validate an unknown provider-specific tool,
+ # so preserve the historical assumption that it made progress.
+ usable_tool_call = True
+ break
+ try:
+ tool.parameters.model_validate_json(tool_call.arguments or "{}")
+ except _PydanticValidationError:
+ continue
+ usable_tool_call = True
+ break
+
+ if choice.finish_reason in ("length", "max_tokens") and not usable_tool_call:
self._truncation_overruns += 1
self._recover_from_truncation = True
LOGGER.warning(
- "completion budget (%d tokens) exhausted with no tool call [overrun #%d]; %s",
+ "completion budget (%d tokens) exhausted with no usable tool call [overrun #%d]; %s",
capped_max,
self._truncation_overruns,
"next turn will run with thinking disabled"
diff --git a/responses_api_agents/stirrup_agent/tests/test_app.py b/responses_api_agents/stirrup_agent/tests/test_app.py
index 0a5e99ec166..8e75b4514e5 100644
--- a/responses_api_agents/stirrup_agent/tests/test_app.py
+++ b/responses_api_agents/stirrup_agent/tests/test_app.py
@@ -17,6 +17,7 @@
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
+from omegaconf import OmegaConf
from stirrup.core.models import AssistantMessage, TokenUsage, ToolCall
from nemo_gym.config_types import ModelServerRef, ResourcesServerRef
@@ -129,6 +130,63 @@ def test_sanity(self) -> None:
)
StirrupAgentWrapper(config=config, server_client=MagicMock(spec=ServerClient))
+ # Generic Stirrup users retain the historical behavior. GDPVal opts in
+ # to its larger floor and matching template semantics in its benchmark
+ # config below.
+ assert config.min_completion_tokens == 1024
+ assert config.context_window_tokens == 262144
+ assert config.prompt_estimator_truncate_history_thinking is None
+ assert config.min_compaction_summary_words == 1
+
+ def test_gdpval_config_opts_into_safe_context_budget(self) -> None:
+ repo_root = STIRRUP_AGENT_DIR.parents[1]
+ config = OmegaConf.load(repo_root / "benchmarks" / "gdpval" / "config.yaml")
+ agent = config.gdpval_stirrup_agent.responses_api_agents.stirrup_agent
+
+ assert agent.min_completion_tokens == 8192
+ assert agent.context_window_tokens == 262144
+ assert "prompt_estimator_truncate_history_thinking" not in agent
+ assert agent.min_compaction_summary_words == 50
+
+ @pytest.mark.asyncio
+ async def test_request_output_cap_does_not_replace_model_context_window(self) -> None:
+ config = _make_config()
+ config.context_window_tokens = 262144
+ config.max_completion_tokens_cap = 64000
+ wrapper = StirrupAgentWrapper(config=config, server_client=MagicMock(spec=ServerClient))
+ body = NeMoGymResponseCreateParamsNonStreaming(
+ input="ignored",
+ model="policy",
+ max_output_tokens=8192,
+ metadata={"task_id": "task-1", "prompt": "do the thing"},
+ )
+ captured = {}
+
+ async def completed_rollout():
+ return {
+ "input_items": [],
+ "output_items": [],
+ "deliverable_text": "",
+ "elapsed_seconds": 0,
+ }
+
+ def launch(params):
+ captured.update(params)
+ return completed_rollout()
+
+ with (
+ patch.object(StirrupAgentWrapper, "resolve_model_base_url", return_value="http://policy.invalid/v1"),
+ patch.object(wrapper.task_strategy, "get_exec_provider", return_value=None),
+ patch(
+ "responses_api_agents.stirrup_agent.app.run_stirrup_agent_remote.remote",
+ side_effect=launch,
+ ),
+ ):
+ await wrapper.responses(body)
+
+ assert captured["context_window_tokens"] == 262144
+ assert captured["max_completion_tokens_cap"] == 8192
+
def test_output_history_preserves_nemo_user_tool_results(self) -> None:
"""Run-history export should keep NeMo user-role tool results as tool outputs."""
history = [
diff --git a/responses_api_agents/stirrup_agent/tests/test_nemo_agent_compaction.py b/responses_api_agents/stirrup_agent/tests/test_nemo_agent_compaction.py
new file mode 100644
index 00000000000..af860a976d3
--- /dev/null
+++ b/responses_api_agents/stirrup_agent/tests/test_nemo_agent_compaction.py
@@ -0,0 +1,303 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Compaction hardening for the Gym Stirrup adapter."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+from stirrup.core.exceptions import ContextOverflowError
+from stirrup.core.models import (
+ AssistantMessage,
+ Reasoning,
+ SummaryMessage,
+ SystemMessage,
+ TokenUsage,
+ ToolMessage,
+ UserMessage,
+)
+from stirrup.prompts import MESSAGE_SUMMARIZER
+
+from responses_api_agents.stirrup_agent.nemo_agent import (
+ NeMoAgent,
+ NeMoSessionAgent,
+ _is_meaningful_compaction_summary,
+ _sanitize_compaction_summary,
+)
+
+
+def _words(count: int) -> str:
+ return " ".join(f"word{index}" for index in range(count))
+
+
+def _assistant(content: str, *, message_id: str = "assistant-1") -> AssistantMessage:
+ return AssistantMessage(
+ id=message_id,
+ content=content,
+ token_usage=TokenUsage(input=1, answer=1, reasoning=0),
+ )
+
+
+class _FakeClient:
+ model_slug = "fake-model"
+ max_tokens = 262_144
+
+ def __init__(self, *results: AssistantMessage | Exception) -> None:
+ self.results = list(results)
+ self.calls: list[tuple[list, dict]] = []
+
+ async def generate(self, messages, tools):
+ self.calls.append((messages, tools))
+ result = self.results.pop(0)
+ if isinstance(result, Exception):
+ raise result
+ return result
+
+
+class _RecoveryAwareFakeClient(_FakeClient):
+ def __init__(
+ self,
+ *results: AssistantMessage | Exception,
+ recovery: bool,
+ overruns: int,
+ mutate_during_generate: bool = False,
+ ) -> None:
+ super().__init__(*results)
+ self._recover_from_truncation = recovery
+ self._truncation_overruns = overruns
+ self.mutate_during_generate = mutate_during_generate
+ self.recovery_seen: list[bool] = []
+
+ async def generate(self, messages, tools):
+ self.recovery_seen.append(self._recover_from_truncation)
+ if self.mutate_during_generate:
+ self._recover_from_truncation = True
+ self._truncation_overruns += 1
+ return await super().generate(messages, tools)
+
+
+def _agent(
+ client: _FakeClient,
+ logger: MagicMock | None = None,
+ *,
+ min_words: int = 1,
+) -> NeMoAgent:
+ return NeMoAgent(
+ client=client,
+ name="test_agent",
+ tools=[],
+ logger=logger or MagicMock(),
+ min_compaction_summary_words=min_words,
+ )
+
+
+def _history() -> list:
+ return [
+ SystemMessage(content="system"),
+ UserMessage(content="original task"),
+ _assistant("worked", message_id="assistant-1"),
+ ToolMessage(content="tool result", tool_call_id="call-1", name="code_exec", success=True),
+ ]
+
+
+def test_sanitize_compaction_summary_removes_only_leading_think_blocks() -> None:
+ summary = " private reasoning\n\n" + _words(50) + " literal tag"
+
+ sanitized = _sanitize_compaction_summary(summary)
+
+ assert sanitized.startswith("word0 word1")
+ assert "private reasoning" not in sanitized
+ assert sanitized.endswith("literal tag")
+
+
+@pytest.mark.parametrize(
+ "summary",
+ [
+ "",
+ " ",
+ "reasoning without a closing tag",
+ "reasoning\n\n",
+ "first second but unclosed",
+ ],
+)
+def test_sanitize_compaction_summary_rejects_empty_or_unclosed_think(summary: str) -> None:
+ assert _sanitize_compaction_summary(summary) == ""
+
+
+def test_meaningful_compaction_summary_uses_configured_minimum() -> None:
+ assert not _is_meaningful_compaction_summary(_words(49), min_words=50)
+ assert _is_meaningful_compaction_summary(_words(50), min_words=50)
+
+
+@pytest.mark.parametrize("summary", ["Done.", "已完成并保留所有关键进展。"])
+def test_generic_compaction_accepts_concise_and_cjk_summaries(summary: str) -> None:
+ assert _is_meaningful_compaction_summary(summary)
+
+
+def test_compaction_minimum_must_be_positive() -> None:
+ with pytest.raises(ValueError, match="at least 1"):
+ _agent(_FakeClient(), min_words=0)
+
+
+def test_session_agent_uses_the_hardened_compaction_override() -> None:
+ assert NeMoSessionAgent.summarize_messages is NeMoAgent.summarize_messages
+
+
+@pytest.mark.asyncio
+async def test_summarization_retries_degenerate_outputs_and_persists_only_sanitized_text() -> None:
+ good_summary = _words(60)
+ client = _FakeClient(
+ _assistant("reasoning only"),
+ _assistant("too short to preserve a long conversation"),
+ AssistantMessage(
+ content=f"private chain of thought\n{good_summary}",
+ reasoning=Reasoning(content="separate private reasoning"),
+ ),
+ )
+ logger = MagicMock()
+ agent = _agent(client, logger, min_words=50)
+ history = _history()
+
+ archived, compacted = await agent.summarize_messages(history, {"assistant-1": {}})
+
+ assert archived == history
+ assert len(client.calls) == 3
+ assert client.calls[0][0][-1].content == MESSAGE_SUMMARIZER
+ assert client.calls[0][1]
+ assert "text only" in client.calls[1][0][-1].content.lower()
+ assert client.calls[1][1]
+ assert client.calls[2][1] == {}
+
+ summaries = [message for message in compacted if isinstance(message, SummaryMessage)]
+ assert len(summaries) == 1
+ persisted = summaries[0].content
+ assert good_summary in persisted
+ assert "" not in persisted.lower()
+ assert "private" not in persisted.lower()
+ assert compacted[-1] == UserMessage(content="Got it, thanks!")
+
+ logger.context_summarization_complete.assert_called_once()
+ logged_summary, logged_bridge = logger.context_summarization_complete.call_args.args
+ assert logged_summary == good_summary
+ assert "" not in logged_bridge.lower()
+
+
+@pytest.mark.asyncio
+async def test_summarization_fails_closed_after_three_degenerate_outputs() -> None:
+ client = _FakeClient(
+ _assistant("one"),
+ _assistant(_words(10)),
+ _assistant("unclosed"),
+ )
+ logger = MagicMock()
+ agent = _agent(client, logger, min_words=50)
+ metadata = {"assistant-1": {"code_exec": []}}
+
+ with pytest.raises(RuntimeError, match="usable text after 3 attempts"):
+ await agent.summarize_messages(_history(), metadata)
+
+ assert len(client.calls) == 3
+ assert metadata == {"assistant-1": {"code_exec": []}}
+ logger.context_summarization_complete.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_summarization_preserves_context_overflow_unwind_semantics() -> None:
+ history = [
+ SystemMessage(content="system"),
+ UserMessage(content="original task"),
+ _assistant("first turn", message_id="assistant-1"),
+ ToolMessage(content="first result", tool_call_id="call-1", name="code_exec", success=True),
+ _assistant("second turn", message_id="assistant-2"),
+ ToolMessage(content="second result", tool_call_id="call-2", name="code_exec", success=True),
+ ]
+ client = _FakeClient(ContextOverflowError("too large"), _assistant(_words(60)))
+ agent = _agent(client, min_words=50)
+ metadata = {"assistant-1": {"first": []}, "assistant-2": {"second": []}}
+
+ archived, compacted = await agent.summarize_messages(history, metadata)
+
+ assert archived == history[:4]
+ assert "assistant-1" in metadata
+ assert "assistant-2" not in metadata
+ assert any(isinstance(message, SummaryMessage) for message in compacted)
+ assert len(client.calls) == 2
+
+
+@pytest.mark.asyncio
+async def test_summary_calls_do_not_consume_an_armed_policy_recovery() -> None:
+ client = _RecoveryAwareFakeClient(_assistant(_words(60)), recovery=True, overruns=7)
+ agent = _agent(client, min_words=50)
+
+ await agent.summarize_messages(_history(), {})
+
+ assert client.recovery_seen == [False]
+ assert client._recover_from_truncation is True
+ assert client._truncation_overruns == 7
+
+
+@pytest.mark.asyncio
+async def test_summary_calls_do_not_arm_policy_recovery_or_increment_overruns() -> None:
+ client = _RecoveryAwareFakeClient(
+ _assistant("too short"),
+ _assistant("still too short"),
+ _assistant(_words(60)),
+ recovery=False,
+ overruns=2,
+ mutate_during_generate=True,
+ )
+ agent = _agent(client, min_words=50)
+
+ await agent.summarize_messages(_history(), {})
+
+ assert client.recovery_seen == [False, False, False]
+ assert client._recover_from_truncation is False
+ assert client._truncation_overruns == 2
+
+
+@pytest.mark.asyncio
+async def test_session_run_replays_only_the_sanitized_summary_after_compaction() -> None:
+ good_summary = _words(60)
+ client = _FakeClient(
+ AssistantMessage(
+ content="first policy turn",
+ token_usage=TokenUsage(input=190_000, answer=10, reasoning=0),
+ ),
+ _assistant(f"private summary reasoning\n{good_summary}"),
+ _assistant("second policy turn"),
+ )
+ agent = NeMoAgent(
+ client=client,
+ name="test_agent",
+ tools=[],
+ logger=MagicMock(),
+ max_turns=2,
+ context_summarization_cutoff=0.7,
+ min_compaction_summary_words=50,
+ )
+
+ async with agent.session(cache_on_interrupt=False) as session:
+ assert isinstance(session, NeMoSessionAgent)
+ _, history, _ = await session.run("original task")
+
+ assert len(history) == 2
+ assert any(isinstance(message, SummaryMessage) for message in history[1])
+ second_policy_messages = client.calls[2][0]
+ serialized_context = "\n".join(
+ message.content for message in second_policy_messages if isinstance(message.content, str)
+ )
+ assert good_summary in serialized_context
+ assert "private summary reasoning" not in serialized_context
+ assert "" not in serialized_context.lower()
diff --git a/responses_api_agents/stirrup_agent/tests/test_nemo_client.py b/responses_api_agents/stirrup_agent/tests/test_nemo_client.py
index 41b8833fe14..2b2690aea0d 100644
--- a/responses_api_agents/stirrup_agent/tests/test_nemo_client.py
+++ b/responses_api_agents/stirrup_agent/tests/test_nemo_client.py
@@ -20,10 +20,14 @@
from __future__ import annotations
+import json
+from copy import deepcopy
from unittest.mock import AsyncMock, MagicMock
import pytest
-from stirrup.core.models import AssistantMessage, SystemMessage, TokenUsage, ToolCall, ToolMessage, UserMessage
+from pydantic import BaseModel
+from stirrup.core.exceptions import ContextOverflowError
+from stirrup.core.models import AssistantMessage, SystemMessage, TokenUsage, Tool, ToolCall, ToolMessage, UserMessage
from responses_api_agents.stirrup_agent.nemo_agent import NeMoUserMessage
from responses_api_agents.stirrup_agent.nemo_client import (
@@ -80,6 +84,7 @@ async def test_generate_forwards_configured_sampling_kwargs() -> None:
assert sent["temperature"] == pytest.approx(0.42)
assert sent["top_p"] == pytest.approx(0.7)
assert sent["extra_body"]["chat_template_kwargs"]["enable_thinking"] is False
+ assert "truncate_history_thinking" not in sent["extra_body"]["chat_template_kwargs"]
assert sent["max_completion_tokens"] <= 2048
@@ -209,9 +214,396 @@ def test_defaults_match_pre_lift_behaviour() -> None:
assert client._top_p == 0.95
assert client._enable_thinking is True
assert client._max_completion_tokens_cap == 64000
+ assert client._min_completion_tokens == 1024
+ assert client._prompt_estimator_truncate_history_thinking is None
assert client._truncation_recovery is True
+def _long_thinking_history(*, closed: bool = True) -> list:
+ messages = [UserMessage(content="do the task")]
+ for _ in range(4):
+ reasoning = "r" * 195_000
+ content = f"{reasoning}visible" if closed else f"{reasoning}"
+ messages.extend(
+ [
+ AssistantMessage(content=content, token_usage=TokenUsage(input=1, answer=1, reasoning=0)),
+ UserMessage(content="Please continue the task"),
+ ]
+ )
+ return messages
+
+
+@pytest.mark.asyncio
+async def test_fallback_ignores_closed_historical_thinking_when_template_truncates_it() -> None:
+ """Regression: raw 64k thinking turns must not collapse later calls to the floor.
+
+ Nemotron's chat template removes each completed reasoning block once a later
+ user message exists. The tokenizer-free fallback must estimate that rendered
+ prompt, not the much larger forensic history retained by Stirrup.
+ """
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=262_144,
+ base_url="http://test",
+ api_key="k",
+ min_completion_tokens=8192,
+ max_completion_tokens_cap=64_000,
+ prompt_estimator_truncate_history_thinking=True,
+ )
+ fake_create = AsyncMock(return_value=_make_response())
+ client._client = MagicMock()
+ client._client.chat.completions.create = fake_create
+
+ history = _long_thinking_history()
+ original_content = history[1].content
+ await client.generate(history, tools={})
+
+ sent = fake_create.await_args.kwargs
+ assert sent["max_completion_tokens"] == 64_000
+ assert "truncate_history_thinking" not in sent["extra_body"]["chat_template_kwargs"]
+ assert sent["messages"][1]["content"][0]["text"] == original_content
+ assert history[1].content == original_content
+
+
+def test_tokenizer_receives_estimator_history_truncation_setting() -> None:
+ class RecordingTokenizer:
+ def __init__(self) -> None:
+ self.template_kwargs = None
+
+ def apply_chat_template(self, messages, **kwargs):
+ self.template_kwargs = kwargs
+ return "rendered"
+
+ def __call__(self, text, **kwargs):
+ return {"input_ids": [1, 2, 3]}
+
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=262_144,
+ base_url="http://test",
+ api_key="k",
+ prompt_estimator_truncate_history_thinking=True,
+ )
+ tokenizer = RecordingTokenizer()
+ client._tokenizer = tokenizer
+
+ assert client._count_input_tokens(to_provider_openai_messages([UserMessage(content="hi")]), tools={}) == 3
+ assert tokenizer.template_kwargs["truncate_history_thinking"] is True
+
+
+def test_token_count_is_exact_only_when_the_complete_prompt_renders() -> None:
+ class MessagesOnlyTokenizer:
+ def apply_chat_template(self, messages, **kwargs):
+ if "tools" in kwargs:
+ raise TypeError("template does not accept tools")
+ return "rendered messages"
+
+ def __call__(self, text, **kwargs):
+ return {"input_ids": [1, 2, 3]}
+
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=262_144,
+ base_url="http://test",
+ api_key="k",
+ )
+ client._tokenizer = MessagesOnlyTokenizer()
+ messages = to_provider_openai_messages([UserMessage(content="hi")])
+
+ count_without_tools, exact_without_tools = client._count_input_tokens_with_confidence(messages, tools={})
+ count_with_tools, exact_with_tools = client._count_input_tokens_with_confidence(
+ messages,
+ tools={"code_exec": _code_exec_tool()},
+ )
+
+ assert count_without_tools == 3
+ assert exact_without_tools is True
+ assert count_with_tools > 3
+ assert exact_with_tools is False
+
+
+def test_nemotron_estimator_decodes_tool_arguments_for_exact_template_render() -> None:
+ class MappingArgumentsTokenizer:
+ def __init__(self) -> None:
+ self.rendered_messages = None
+
+ def apply_chat_template(self, messages, **kwargs):
+ arguments = messages[1]["tool_calls"][0]["function"]["arguments"]
+ if not isinstance(arguments, dict):
+ raise TypeError("tool arguments must be a mapping")
+ self.rendered_messages = messages
+ return "complete rendered prompt"
+
+ def __call__(self, text, **kwargs):
+ return {"input_ids": [1, 2, 3, 4]}
+
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=262_144,
+ base_url="http://test",
+ api_key="k",
+ prompt_estimator_truncate_history_thinking=True,
+ )
+ tokenizer = MappingArgumentsTokenizer()
+ client._tokenizer = tokenizer
+ messages = _serialized_history_with_prior_assistant("oldvisible", with_tool_call=True)
+ original = deepcopy(messages)
+
+ count, exact = client._count_input_tokens_with_confidence(
+ messages,
+ tools={"code_exec": _code_exec_tool()},
+ )
+
+ assert (count, exact) == (4, True)
+ assert tokenizer.rendered_messages[1]["tool_calls"][0]["function"]["arguments"] == {"cmd": "echo hello"}
+ assert messages == original
+
+
+def test_nemotron_estimator_does_not_claim_exact_render_for_invalid_tool_arguments() -> None:
+ class MappingArgumentsTokenizer:
+ def apply_chat_template(self, messages, **kwargs):
+ arguments = messages[1]["tool_calls"][0]["function"]["arguments"]
+ if not isinstance(arguments, dict):
+ raise TypeError("tool arguments must be a mapping")
+ return "complete rendered prompt"
+
+ def __call__(self, text, **kwargs):
+ return {"input_ids": [1, 2, 3]}
+
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=262_144,
+ base_url="http://test",
+ api_key="k",
+ prompt_estimator_truncate_history_thinking=True,
+ )
+ client._tokenizer = MappingArgumentsTokenizer()
+ messages = _serialized_history_with_prior_assistant("visible", with_tool_call=True)
+ messages[1]["tool_calls"][0]["function"]["arguments"] = "not-json"
+
+ _, exact = client._count_input_tokens_with_confidence(
+ messages,
+ tools={"code_exec": _code_exec_tool()},
+ )
+
+ assert exact is False
+
+
+def _serialized_history_with_prior_assistant(content: str, *, with_tool_call: bool) -> list:
+ tool_calls = (
+ [ToolCall(tool_call_id="call_1", name="code_exec", arguments='{"cmd":"echo hello"}')] if with_tool_call else []
+ )
+ return to_provider_openai_messages(
+ [
+ UserMessage(content="start"),
+ AssistantMessage(
+ content=content,
+ tool_calls=tool_calls,
+ token_usage=TokenUsage(input=1, answer=1, reasoning=0),
+ ),
+ UserMessage(content="continue"),
+ ]
+ )
+
+
+@pytest.mark.parametrize(
+ ("content", "with_tool_call", "expected"),
+ [
+ ("oldvisible", False, "visible"),
+ ("unfinished", False, "unfinished"),
+ ("prefixoldvisible", True, "visible"),
+ ("prefixunfinished", True, "prefix"),
+ ],
+ ids=("closed-no-tool", "unclosed-no-tool", "closed-tool", "unclosed-tool"),
+)
+def test_estimator_mirrors_nemotron_prior_assistant_semantics(
+ content: str,
+ with_tool_call: bool,
+ expected: str,
+) -> None:
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=262_144,
+ base_url="http://test",
+ api_key="k",
+ prompt_estimator_truncate_history_thinking=True,
+ )
+ serialized = _serialized_history_with_prior_assistant(content, with_tool_call=with_tool_call)
+ original = deepcopy(serialized)
+
+ estimated = client._messages_for_estimator_count(serialized)
+
+ assert estimated[1]["content"][0]["text"] == expected
+ assert estimated[1].get("tool_calls") == serialized[1].get("tool_calls")
+ assert serialized == original
+
+
+def test_estimator_uses_content_after_last_closed_think_block() -> None:
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=262_144,
+ base_url="http://test",
+ api_key="k",
+ prompt_estimator_truncate_history_thinking=True,
+ )
+ serialized = _serialized_history_with_prior_assistant(
+ "firstmiddlesecondtail",
+ with_tool_call=False,
+ )
+
+ estimated = client._messages_for_estimator_count(serialized)
+
+ assert estimated[1]["content"][0]["text"] == "tail"
+
+
+def test_character_fallback_counts_complete_serialized_messages() -> None:
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=262_144,
+ base_url="http://test",
+ api_key="k",
+ )
+ serialized = _serialized_history_with_prior_assistant("same content", with_tool_call=True)
+ expected = len(json.dumps(serialized, ensure_ascii=False)) // 3
+
+ count = client._count_input_tokens(serialized, tools={})
+
+ assert count == expected
+ long_arguments = deepcopy(serialized)
+ long_arguments[1]["tool_calls"][0]["function"]["arguments"] = "x" * 3000
+ assert client._count_input_tokens(long_arguments, tools={}) >= count + 990
+
+
+@pytest.mark.asyncio
+async def test_completion_budget_is_bounded_by_positive_estimated_remaining_context() -> None:
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=10_000,
+ base_url="http://test",
+ api_key="k",
+ completion_token_buffer=1000,
+ min_completion_tokens=8192,
+ max_completion_tokens_cap=64_000,
+ )
+ client._count_input_tokens_with_confidence = MagicMock(return_value=(9500, True))
+ client._tokenizer = object()
+ fake_create = AsyncMock(return_value=_make_response())
+ client._client = MagicMock()
+ client._client.chat.completions.create = fake_create
+
+ await client.generate([UserMessage(content="hi")], tools={})
+
+ assert fake_create.await_args.kwargs["max_completion_tokens"] == 500
+
+
+@pytest.mark.asyncio
+async def test_no_estimated_context_remaining_raises_before_dispatch() -> None:
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=10_000,
+ base_url="http://test",
+ api_key="k",
+ min_completion_tokens=8192,
+ )
+ client._count_input_tokens_with_confidence = MagicMock(return_value=(10_000, True))
+ client._tokenizer = object()
+ fake_create = AsyncMock(return_value=_make_response())
+ client._client = MagicMock()
+ client._client.chat.completions.create = fake_create
+
+ with pytest.raises(ContextOverflowError, match="leaves no room"):
+ await client.generate([UserMessage(content="hi")], tools={})
+
+ fake_create.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_output_cap_smaller_than_prompt_does_not_shrink_context_window() -> None:
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=262_144,
+ base_url="http://test",
+ api_key="k",
+ min_completion_tokens=8192,
+ max_completion_tokens_cap=8192,
+ )
+ client._count_input_tokens_with_confidence = MagicMock(return_value=(10_000, True))
+ client._tokenizer = object()
+ fake_create = AsyncMock(return_value=_make_response())
+ client._client = MagicMock()
+ client._client.chat.completions.create = fake_create
+
+ await client.generate([UserMessage(content="hi")], tools={})
+
+ assert fake_create.await_args.kwargs["max_completion_tokens"] == 8192
+
+
+@pytest.mark.asyncio
+async def test_fallback_overestimate_still_dispatches_configured_floor() -> None:
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=10_000,
+ base_url="http://test",
+ api_key="k",
+ min_completion_tokens=8192,
+ max_completion_tokens_cap=64_000,
+ )
+ client._count_input_tokens_with_confidence = MagicMock(return_value=(20_000, False))
+ fake_create = AsyncMock(return_value=_make_response())
+ client._client = MagicMock()
+ client._client.chat.completions.create = fake_create
+
+ await client.generate([UserMessage(content="hi")], tools={})
+
+ assert client._tokenizer is None
+ assert fake_create.await_args.kwargs["max_completion_tokens"] == 8192
+
+
+@pytest.mark.asyncio
+async def test_loaded_tokenizer_approximation_does_not_impose_a_false_context_bound() -> None:
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=10_000,
+ base_url="http://test",
+ api_key="k",
+ min_completion_tokens=8192,
+ max_completion_tokens_cap=64_000,
+ )
+ client._tokenizer = object()
+ client._count_input_tokens_with_confidence = MagicMock(return_value=(20_000, False))
+ fake_create = AsyncMock(return_value=_make_response())
+ client._client = MagicMock()
+ client._client.chat.completions.create = fake_create
+
+ await client.generate([UserMessage(content="hi")], tools={})
+
+ assert fake_create.await_args.kwargs["max_completion_tokens"] == 8192
+
+
+@pytest.mark.asyncio
+async def test_minimum_completion_floor_never_consumes_more_than_remaining_context() -> None:
+ client = DynamicMaxTokensChatCompletionsClient(
+ model="m",
+ max_tokens=10_000,
+ base_url="http://test",
+ api_key="k",
+ completion_token_buffer=1000,
+ min_completion_tokens=8192,
+ max_completion_tokens_cap=64_000,
+ )
+ client._count_input_tokens_with_confidence = MagicMock(return_value=(1500, True))
+ client._tokenizer = object()
+ fake_create = AsyncMock(return_value=_make_response())
+ client._client = MagicMock()
+ client._client.chat.completions.create = fake_create
+
+ await client.generate([UserMessage(content="hi")], tools={})
+
+ # 8,500 raw tokens remain. The floor may reclaim part of the 1,000-token
+ # safety buffer, but it may not exceed the actual estimated remainder.
+ assert fake_create.await_args.kwargs["max_completion_tokens"] == 8192
+
+
# ---------------------------------------------------------------------------
# Thinking-overrun recovery
# ---------------------------------------------------------------------------
@@ -249,6 +641,19 @@ def _tool_call_response(finish_reason: str = "stop"):
return response
+class _CodeExecParams(BaseModel):
+ cmd: str
+
+
+def _code_exec_tool() -> Tool:
+ return Tool(
+ name="code_exec",
+ description="Run a command.",
+ parameters=_CodeExecParams,
+ executor=lambda _: None,
+ )
+
+
@pytest.mark.asyncio
async def test_budget_overrun_without_tool_call_arms_recovery() -> None:
"""The next call drops thinking and carries an explicit instruction."""
@@ -325,6 +730,21 @@ async def test_truncated_call_that_still_made_a_tool_call_is_left_alone() -> Non
assert len(second["messages"]) == 1
+@pytest.mark.asyncio
+@pytest.mark.parametrize("arguments", ["", "{}", '{"wrong":"field"}'])
+async def test_truncated_schema_invalid_tool_call_arms_recovery(arguments: str) -> None:
+ """A parser-emitted call with unusable arguments made no task progress."""
+ client = _make_client()
+ response = _tool_call_response(finish_reason="length")
+ response.choices[0].message.tool_calls[0].function.arguments = arguments
+ client._client.chat.completions.create = AsyncMock(return_value=response)
+
+ await client.generate([UserMessage(content="hi")], tools={"code_exec": _code_exec_tool()})
+
+ assert client._recover_from_truncation is True
+ assert client._truncation_overruns == 1
+
+
@pytest.mark.asyncio
async def test_text_only_response_that_finished_normally_is_left_alone() -> None:
"""A short text-only turn is a different problem; don't touch it."""