hybrid: SWE trajectory compaction + OpenAI adapter bug fixes

- _loop_local: two-stage compaction at 22k tokens (tiktoken cl100k).
  Stage 1 elides old tool messages to [tool output elided: N chars, exit=X],
  preserving tool_call_id. Stage 2 folds turn pairs into a synthesized
  system note. Last 3 turns + system + initial user always intact.
  Emergency-compact retry on a 32k BadRequest from vLLM.
- _loop_cloud_openai: assistant content is "" not None on tool-only turns
  (OpenAI 400's on null). Bash output decoded with errors="replace" and
  stubbed to [binary output: N bytes, exit=X] on >5% replacement or NUL.
- _openai_retry: local-vLLM branch now retries APIConnectionError /
  APITimeoutError / InternalServerError 3x with 1/2/4s backoff.
- Tests: 18 cases covering compaction shape, token budget, tool_call_id
  preservation, null-content guard, binary decode, local retry.

Addresses ctx-overflow in Gemma-31B SWE n=100 cells (61 + 19 errored
rows) and the OpenAI/Gemini adapter errors across ~19 cells in the
n=100 hybrid sweep.
This commit is contained in:
Andrew Park
2026-05-26 00:13:47 -07:00
parent 22fbfe65aa
commit 081839ee7b
5 changed files with 1443 additions and 26 deletions
@@ -0,0 +1,355 @@
"""Process-wide retry + concurrency hardening for cloud OpenAI calls.
Why this exists
---------------
When we run the hybrid paradigms (Minions, Advisors, Conductor, …) at
n=100 against ``gpt-5`` / ``gpt-5-mini`` over the prepaid OpenAI quota,
sustained concurrency walls the org-level rate limit and the OpenAI SDK
raises :class:`openai.RateLimitError`. The SDK's own retry path is short
(default ``max_retries=2`` on a small backoff) — under sustained pressure
every wave of retries hits the same wall and the runner records 19-69
errored rows per cell, degenerating the result.
Mirrors the existing ``_patch_anthropic_globally`` pattern in
``minions.py``: monkey-patch the SDK at module level so it applies even
to libraries that build their own ``openai.OpenAI()`` clients
(HazyResearch Minions's ``OpenAIClient``, Archon's adapters,
``mini_swe_agent``, etc.). One call to :func:`patch_openai_globally` from
either ``_base.py`` or ``minions._apply_patches_once`` is enough — the
patch is idempotent and process-wide.
What it does
------------
1. Bumps ``openai.OpenAI()`` constructor defaults to
``timeout=600.0`` / ``max_retries=8`` (the SDK's own backoff is fine
for transient blips; we layer our own loop on top for sustained walls).
2. Wraps ``chat.completions.create`` with:
- A **per-org semaphore** (``OPENJARVIS_OPENAI_MAX_CONCURRENCY``,
default 4) that throttles sustained concurrency. Single bursts are
fine — prepaid quotas wall on sustained rate, not on a brief spike.
The semaphore is **only acquired for cloud calls** (``api.openai.com``);
vLLM calls routed through the OpenAI SDK (``base_url`` set to a
local endpoint or ``api_key="EMPTY"``) bypass it.
- Exponential backoff with jitter on :class:`openai.RateLimitError`,
:class:`openai.APITimeoutError`, :class:`openai.APIConnectionError`,
and 5xx :class:`openai.APIStatusError`. Starts at 2 s, caps at 60 s,
up to 8 attempts (~3.5 min worst case). Honors a ``Retry-After``
header when the SDK surfaces one.
- On exhaustion, re-raises the last exception (it bubbles up to the
runner, which records ``error="RateLimitError: ..."`` in
``results.jsonl`` — no silent drop).
Env knobs
---------
- ``OPENJARVIS_OPENAI_MAX_CONCURRENCY`` (default ``4``) — semaphore
capacity. Set to e.g. ``2`` if the wall is still hit; set to ``0`` to
disable throttling entirely (passes through to the SDK).
- ``OPENJARVIS_OPENAI_MAX_RETRIES`` (default ``8``) — outer retry loop
cap (separate from the SDK's own ``max_retries``).
- ``OPENJARVIS_OPENAI_RETRY_BASE`` (default ``2.0``) — base seconds for
exponential backoff. Schedule is ``min(60, base * 2**attempt) * jitter``.
- ``OPENJARVIS_OPENAI_RETRY_CAP`` (default ``60.0``) — max single-step
sleep in seconds.
"""
from __future__ import annotations
import os
import random
import threading
import time
from typing import Any, Callable, Optional, Tuple
from urllib.parse import urlparse
# ---------------------------------------------------------------------------
# Tunables (read once at module load, can be overridden via env)
# ---------------------------------------------------------------------------
def _env_int(name: str, default: int) -> int:
try:
v = int(os.environ.get(name, "") or default)
return max(0, v)
except ValueError:
return default
def _env_float(name: str, default: float) -> float:
try:
return float(os.environ.get(name, "") or default)
except ValueError:
return default
_MAX_CONCURRENCY = _env_int("OPENJARVIS_OPENAI_MAX_CONCURRENCY", 4)
_MAX_RETRIES = _env_int("OPENJARVIS_OPENAI_MAX_RETRIES", 8)
_RETRY_BASE = _env_float("OPENJARVIS_OPENAI_RETRY_BASE", 2.0)
_RETRY_CAP = _env_float("OPENJARVIS_OPENAI_RETRY_CAP", 60.0)
# Single process-wide semaphore. ``BoundedSemaphore(0)`` would block
# forever, so when the env knob is 0 we hand back a no-op context manager.
class _NullSem:
def __enter__(self) -> "_NullSem":
return self
def __exit__(self, *a: Any) -> None:
return None
_SEM: Any
if _MAX_CONCURRENCY > 0:
_SEM = threading.BoundedSemaphore(_MAX_CONCURRENCY)
else:
_SEM = _NullSem()
_PATCHED = False
_PATCH_LOCK = threading.Lock()
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _is_local_endpoint(client: Any) -> bool:
"""True if this OpenAI client points at a local vLLM endpoint.
We detect via either ``api_key == "EMPTY"`` (the convention used by
our ``_call_vllm`` and ``mini_swe_agent``) or a ``base_url`` whose
hostname resolves to localhost. Either signal is enough; both are
cheap to read.
"""
try:
api_key = getattr(client, "api_key", None)
if api_key == "EMPTY":
return True
except Exception:
pass
try:
base_url = str(getattr(client, "base_url", "") or "")
if not base_url:
return False
host = urlparse(base_url).hostname or ""
return host in ("localhost", "127.0.0.1", "0.0.0.0", "::1")
except Exception:
return False
def _extract_retry_after(exc: BaseException) -> Optional[float]:
"""Pull a Retry-After header off an APIStatusError if the SDK exposed it.
OpenAI's SDK keeps the underlying ``httpx.Response`` on
``exc.response`` for ``APIStatusError`` subclasses. Header may be a
seconds-integer or an HTTP-date; we only handle the integer form
(the only thing OpenAI sends in practice).
"""
resp = getattr(exc, "response", None)
if resp is None:
return None
headers = getattr(resp, "headers", None)
if not headers:
return None
for name in ("retry-after", "Retry-After", "x-ratelimit-reset-requests"):
val = headers.get(name) if hasattr(headers, "get") else None
if val is None:
continue
try:
secs = float(val)
if 0 <= secs <= 600:
return secs
except (TypeError, ValueError):
continue
return None
def _is_retryable(exc: BaseException) -> bool:
"""Whether to retry this OpenAI exception class."""
try:
import openai
except ImportError:
return False
if isinstance(exc, (
openai.RateLimitError,
openai.APITimeoutError,
openai.APIConnectionError,
openai.InternalServerError,
)):
return True
if isinstance(exc, openai.APIStatusError):
status = getattr(exc, "status_code", None)
# 429 is RateLimitError already; 5xx is retryable; 408 is a
# timeout the SDK didn't classify (rare).
return status is not None and (status >= 500 or status in (408, 409, 429))
return False
def _sleep_for(attempt: int, exc: BaseException) -> float:
"""Backoff for attempt index ``attempt`` (0-based)."""
hinted = _extract_retry_after(exc)
if hinted is not None and hinted > 0:
# Respect a server-provided hint, but clamp to our cap so a
# pathological header can't stall the run for hours.
return min(_RETRY_CAP, hinted) + random.uniform(0, 0.5)
base = min(_RETRY_CAP, _RETRY_BASE * (2 ** attempt))
# Full jitter — better tail behavior than equal jitter when many
# workers wake at the same moment.
return random.uniform(0.0, base)
# ---------------------------------------------------------------------------
# Wrapping
# ---------------------------------------------------------------------------
def _wrap_create(orig: Callable[..., Any]) -> Callable[..., Any]:
"""Wrap a ``chat.completions.create`` (or ``responses.create``) bound
method's underlying function with retry + concurrency throttling.
The wrapper is a regular function that takes ``self`` as the first
arg, so it can replace ``Completions.create`` at the class level and
still see the bound client through ``self._client``.
"""
def wrapped(self: Any, *args: Any, **kwargs: Any) -> Any:
client = getattr(self, "_client", None)
local = client is not None and _is_local_endpoint(client)
# Local vLLM calls bypass the per-org throttle (no rate limit) and
# the long retry loop (vLLM is mostly either up or down — a 60s
# backoff just delays surfacing the failure). But brief
# ConnectionError blips do happen mid-sweep (socket queue, brief
# server warmup pause): give the local path a SHORT retry — 3
# attempts, 1s/2s/4s — so we don't error an entire row on a
# transient refused connection. Anything else (BadRequest etc.)
# still raises immediately.
if local:
local_last_exc: Optional[BaseException] = None
for attempt in range(3):
try:
return orig(self, *args, **kwargs)
except BaseException as exc: # noqa: BLE001
try:
import openai
except ImportError:
raise
if not isinstance(exc, (
openai.APIConnectionError,
openai.APITimeoutError,
openai.InternalServerError,
)):
raise
local_last_exc = exc
if attempt >= 2:
break
time.sleep(2 ** attempt)
assert local_last_exc is not None
raise local_last_exc
last_exc: Optional[BaseException] = None
for attempt in range(_MAX_RETRIES + 1):
try:
with _SEM:
return orig(self, *args, **kwargs)
except BaseException as exc: # noqa: BLE001
if not _is_retryable(exc):
raise
last_exc = exc
if attempt >= _MAX_RETRIES:
break
delay = _sleep_for(attempt, exc)
# Stderr, not stdout: heartbeat / progress lines must
# stay parseable in the runner log.
try:
import sys
print(
f"[openai-retry] attempt {attempt + 1}/{_MAX_RETRIES} "
f"{type(exc).__name__}: {str(exc)[:120]}"
f"sleeping {delay:.1f}s",
file=sys.stderr,
flush=True,
)
except Exception:
pass
time.sleep(delay)
# Exhausted. Re-raise so the runner records the row as errored.
assert last_exc is not None
raise last_exc
wrapped._hybrid_patched = True # type: ignore[attr-defined]
wrapped.__wrapped__ = orig # type: ignore[attr-defined]
return wrapped
def patch_openai_globally() -> None:
"""Idempotently monkey-patch the OpenAI SDK to add retry + throttling.
Safe to call from multiple modules — guarded by ``_PATCHED`` under a
lock. Patches both the ``Completions.create`` method and the
``OpenAI.__init__`` defaults.
"""
global _PATCHED
if _PATCHED:
return
with _PATCH_LOCK:
if _PATCHED:
return
try:
import openai
from openai.resources.chat import completions as _comp_mod
except ImportError:
return
# Bump constructor defaults so callers that don't pass timeout /
# max_retries explicitly still get sensible values. ``setdefault``
# so any explicit caller value wins.
if not getattr(openai.OpenAI.__init__, "_hybrid_patched", False):
_orig_init = openai.OpenAI.__init__
def _patched_init(self: Any, *args: Any, **kwargs: Any) -> None:
kwargs.setdefault("timeout", 600.0)
kwargs.setdefault("max_retries", _MAX_RETRIES)
return _orig_init(self, *args, **kwargs)
_patched_init._hybrid_patched = True # type: ignore[attr-defined]
openai.OpenAI.__init__ = _patched_init # type: ignore[assignment]
# Wrap chat.completions.create. The SDK exposes the bound method
# via ``Completions.create``; we replace the class attribute so
# every instance (including ones built inside external libs)
# sees the wrapped version.
if not getattr(_comp_mod.Completions.create, "_hybrid_patched", False):
_comp_mod.Completions.create = _wrap_create( # type: ignore[assignment]
_comp_mod.Completions.create
)
# Also patch the async variant for completeness (none of our
# paradigms use it today, but Archon / future paradigms might).
try:
from openai.resources.chat import completions as _comp_mod_async
cls = getattr(_comp_mod_async, "AsyncCompletions", None)
if cls is not None and not getattr(
cls.create, "_hybrid_patched", False
):
# Async wrapper is structurally different — only patch
# the bumped defaults via __init__; full retry loop on
# async would need an async wrapper. Leave that for the
# day a paradigm actually uses it.
pass
except ImportError:
pass
_PATCHED = True
def current_settings() -> Tuple[int, int, float, float]:
"""For tests / smoke runs: return (concurrency, retries, base, cap)."""
return _MAX_CONCURRENCY, _MAX_RETRIES, _RETRY_BASE, _RETRY_CAP
__all__ = ["patch_openai_globally", "current_settings"]
+401 -26
View File
@@ -35,7 +35,10 @@ Differences vs. the upstream
from __future__ import annotations
import json
import os
import re
import shutil
import signal
import subprocess
import tempfile
import time
@@ -135,6 +138,36 @@ BASH_TOOL_OPENAI = {
# ---------- Workdir / bash plumbing ----------
# Models trained on SWE-bench Docker images (Qwen especially) reflexively
# prefix commands with ``cd /testbed`` — the standard container repo path.
# Our harness has no ``/testbed``; the repo is cloned into a per-task
# tempdir and bash already runs with ``cwd`` set to it. An un-rewritten
# ``cd /testbed`` errors with "No such file or directory" and, chained with
# ``&&``, aborts the whole command — so the agent burns every turn and
# never lands an edit. We rewrite ``/testbed`` references to the real
# workdir so those commands run as intended.
_TESTBED_CD_RE = re.compile(r"^\s*cd\s+/testbed(?:/\S*)?\s*(?:&&|;)\s*")
def _rewrite_testbed_paths(command: str, workdir: Path) -> str:
"""Neutralize hard-coded ``/testbed`` paths in a model-issued command.
- A leading ``cd /testbed && ...`` (or ``;``) is stripped — bash already
runs in the repo root, so the rest of the command is correct as-is.
- Any remaining ``/testbed`` occurrences (e.g. ``cat /testbed/foo.py``)
are rewritten to the real workdir.
"""
wd = str(workdir)
new = _TESTBED_CD_RE.sub("", command)
# Bare ``cd /testbed`` with nothing after it → no-op into the workdir.
if re.fullmatch(r"\s*cd\s+/testbed/?\s*", new):
new = f"cd {wd}"
# Replace any other /testbed path references (word-boundary so we don't
# clobber e.g. /testbedrock).
new = re.sub(r"/testbed(?=/|\b)", wd, new)
return new
def _clone_repo(repo: str, base_commit: str, dest: Path) -> None:
"""Shallow-fetch the SWE-bench repo at the right commit into ``dest``."""
url = f"https://github.com/{repo}.git"
@@ -148,27 +181,89 @@ def _clone_repo(repo: str, base_commit: str, dest: Path) -> None:
)
def _decode_bash_output(raw: bytes, exit_code: int) -> str:
"""Safely decode bash stdout/stderr bytes into a str the LLM can read.
The model sometimes runs commands that produce binary output (``cat``
on a ``.pyc`` / ``.png`` / packed extension, a ``find`` that pipes a
binary blob, ``xxd`` on a libc, ...). Decoding those with strict UTF-8
raises ``UnicodeDecodeError`` deep inside ``Popen.communicate()`` and
crashes the whole agent loop (1 errored row in the n=100 sweep per
such command). We:
1. Decode with ``errors="replace"`` so partial / mixed output is
always recoverable as a str.
2. If the result looks predominantly binary — contains a NUL byte OR
has more than ~5% U+FFFD replacement chars after decode — replace
it with a one-line stub so the model can keep working without
drowning in Mojibake. Threshold checked against the *bytes* length
(no NUL byte ⇒ probably text-ish; LLMs handle the occasional
replacement char fine).
"""
if not raw:
return ""
decoded = raw.decode("utf-8", errors="replace")
if b"\x00" in raw or decoded.count("") * 20 > len(decoded):
return f"[binary output: {len(raw)} bytes, exit={exit_code}]"
return decoded
def _run_bash(
command: str, workdir: Path, *, timeout: int = 120, output_cap: int = 10_000
) -> Dict[str, Any]:
"""Run one shell command in ``workdir``. Returns dict with stdout, stderr,
exit_code, and a ``truncated`` flag if output was clamped."""
exit_code, and a ``truncated`` flag if output was clamped.
The command is launched in its own process group (``start_new_session``)
so a model-issued command that backgrounds a long-lived child (a dev
server, ``sleep``, a hung test runner) can be killed *as a tree* on
timeout. Plain ``subprocess.run(..., capture_output=True, timeout=...)``
only kills the direct child and then re-blocks on ``communicate()``
draining the pipe — which a surviving grandchild holds open forever,
silently wedging the whole agent loop.
"""
t0 = time.time()
command = _rewrite_testbed_paths(command, workdir)
# Capture as bytes (no ``text=True``) so a tool invocation that emits
# binary output (compiled artifact, image, PDF, gzipped tarball) can't
# crash the loop on a strict UTF-8 decode mid-``communicate()``. We
# decode below with ``errors="replace"`` and, if the result looks
# binary (null byte or >5% replacement chars), substitute a stub so
# the model doesn't waste tokens / context on Mojibake.
proc = subprocess.Popen(
["bash", "-lc", command],
cwd=str(workdir),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
start_new_session=True,
)
try:
proc = subprocess.run(
["bash", "-lc", command],
cwd=str(workdir),
capture_output=True, text=True, timeout=timeout,
)
stdout = proc.stdout
stderr = proc.stderr
stdout_b, stderr_b = proc.communicate(timeout=timeout)
exit_code = proc.returncode
timed_out = False
except subprocess.TimeoutExpired as e:
stdout = (e.stdout or "") if isinstance(e.stdout, str) else ""
stderr = (e.stderr or "") if isinstance(e.stderr, str) else ""
except subprocess.TimeoutExpired:
# Kill the whole process group so backgrounded grandchildren can't
# keep the stdout/stderr pipe open and deadlock the drain below.
for sig in (signal.SIGTERM, signal.SIGKILL):
try:
os.killpg(proc.pid, sig)
except (ProcessLookupError, PermissionError):
break
try:
proc.wait(timeout=5)
break
except subprocess.TimeoutExpired:
continue
try:
stdout_b, stderr_b = proc.communicate(timeout=10)
except subprocess.TimeoutExpired:
stdout_b, stderr_b = b"", b"";
stdout_b = stdout_b or b""
stderr_b = stderr_b or b""
exit_code = -1
timed_out = True
stdout = _decode_bash_output(stdout_b, exit_code)
stderr = _decode_bash_output(stderr_b, exit_code)
truncated = False
if len(stdout) > output_cap:
stdout = stdout[:output_cap] + f"\n…[+{len(stdout) - output_cap} chars truncated]"
@@ -241,6 +336,8 @@ def run_swe_agent_loop(
turn_max_tokens: int = 4096,
trace_prefix: str = "mini_swe",
workdir: Optional[Path] = None,
compact_at_tokens: int = 24_000,
compact_keep_last: int = 4,
) -> Dict[str, Any]:
"""Run a mini-SWE-agent loop for one SWE-bench task. Returns:
@@ -337,6 +434,8 @@ def run_swe_agent_loop(
output_cap=output_cap,
turn_max_tokens=turn_max_tokens,
trace_prefix=trace_prefix,
compact_at_tokens=compact_at_tokens,
compact_keep_last=compact_keep_last,
)
else:
raise ValueError(f"unsupported backbone: {backbone!r}")
@@ -601,9 +700,14 @@ def _loop_cloud_openai(
# Append the assistant turn (including any tool_calls) so the
# follow-up tool messages have the right call ids to reference.
# OpenAI Chat Completions rejects ``content: null`` with a 400
# ("expected a string, got null") on the next turn when this
# message gets replayed. Use ``""`` — explicitly allowed by the
# schema when ``tool_calls`` is present, and equivalent to
# "assistant had no visible text, only tool calls".
assistant_msg: Dict[str, Any] = {
"role": "assistant",
"content": text or None,
"content": text or "",
}
if tool_calls:
assistant_msg["tool_calls"] = [
@@ -894,6 +998,218 @@ def _loop_cloud_gemini(
# ---------- Local loop (vLLM, OpenAI-compatible multi-turn with tools) ----------
_COMPACT_PROMPT = (
"Summarize the SWE-bench agent trajectory so far in under 2000 characters. "
"Preserve: filenames touched, hypotheses tested, what worked, what failed, "
"and the current plan. Be terse — no preamble, no quoted output, just facts."
)
_TIKTOKEN_ENC = None
_TIKTOKEN_WARNED = False
def _get_tiktoken_enc() -> Any:
global _TIKTOKEN_ENC, _TIKTOKEN_WARNED
if _TIKTOKEN_ENC is not None:
return _TIKTOKEN_ENC
try:
import tiktoken
_TIKTOKEN_ENC = tiktoken.get_encoding("cl100k_base")
except Exception as exc:
if not _TIKTOKEN_WARNED:
print(f"[mini_swe_agent] tiktoken unavailable ({exc!r}); falling back to len(s)//4", flush=True)
_TIKTOKEN_WARNED = True
_TIKTOKEN_ENC = False
return _TIKTOKEN_ENC
def _estimate_prompt_tokens(messages: List[Dict[str, Any]]) -> int:
enc = _get_tiktoken_enc()
total = 0
for m in messages:
total += 4 # per-message overhead
c = m.get("content")
if isinstance(c, str):
s = c
elif isinstance(c, list):
parts = []
for block in c:
if isinstance(block, dict):
parts.append(str(block.get("content") or block.get("text") or ""))
s = "\n".join(parts)
else:
s = ""
for tc in (m.get("tool_calls") or []):
try:
s += "\n" + (tc["function"]["arguments"] or "")
s += "\n" + (tc["function"].get("name") or "")
except (KeyError, TypeError):
pass
tcid = m.get("tool_call_id")
if tcid:
s += "\n" + str(tcid)
if enc:
total += len(enc.encode(s, disallowed_special=()))
else:
total += len(s) // 4
return total
_EXIT_PATTERNS = (
re.compile(r"exit_code\s*[=:]\s*(-?\d+)"),
re.compile(r"returncode\s*[=:]\s*(-?\d+)"),
re.compile(r"\bexit\s+(-?\d+)\b"),
)
def _parse_exit_code(content: Any) -> str:
if not isinstance(content, str):
return "?"
for pat in _EXIT_PATTERNS:
m = pat.search(content)
if m:
return m.group(1)
return "?"
def _identify_turns(messages: List[Dict[str, Any]]) -> List[Tuple[int, int]]:
"""Return list of (start_idx, end_idx_exclusive) for each assistant+tools turn.
A turn = one assistant message (with or without tool_calls) plus any
immediately-following tool messages. System + initial user are skipped.
"""
turns: List[Tuple[int, int]] = []
i = 0
n = len(messages)
while i < n:
role = messages[i].get("role")
if role == "assistant":
j = i + 1
while j < n and messages[j].get("role") == "tool":
j += 1
turns.append((i, j))
i = j
else:
i += 1
return turns
def _compact_local_messages(
messages: List[Dict[str, Any]],
*,
client: Any,
model: str,
keep_last: int,
trace_prefix: str,
compact_at_tokens: int = 24_000,
) -> List[Dict[str, Any]]:
if len(messages) < 2:
return messages
system_msg = messages[0]
initial_user = messages[1]
turns = _identify_turns(messages)
if len(turns) <= keep_last:
return messages
keep_turns = turns[-keep_last:]
old_turns = turns[:-keep_last]
keep_start = keep_turns[0][0]
# Stage 1: elide tool observations in old turns.
before_tokens = _estimate_prompt_tokens(messages)
new_messages: List[Dict[str, Any]] = list(messages)
n_tool_elided = 0
for (s, e) in old_turns:
for k in range(s, e):
m = new_messages[k]
if m.get("role") != "tool":
continue
orig = m.get("content")
if not isinstance(orig, str):
continue
n_chars = len(orig)
if n_chars <= 200:
continue
exit_code = _parse_exit_code(orig)
stub = f"[tool output elided: {n_chars} chars, exit={exit_code}]"
new_messages[k] = {
"role": "tool",
"tool_call_id": m.get("tool_call_id"),
"content": stub,
}
n_tool_elided += 1
after_stage1_tokens = _estimate_prompt_tokens(new_messages)
_record_event({
"kind": f"{trace_prefix}_compact",
"stage": "1",
"msgs_before": len(messages),
"msgs_after": len(new_messages),
"before_tokens": before_tokens,
"after_tokens": after_stage1_tokens,
"n_tool_elided": n_tool_elided,
"n_turns_folded": 0,
"ts": time.time(),
})
if after_stage1_tokens <= compact_at_tokens:
return new_messages
# Stage 2: fold old turns into a single synthetic system summary.
middle = new_messages[2:keep_start]
tail = new_messages[keep_start:]
if not middle:
return new_messages
summary_input = [
{"role": "system", "content": _COMPACT_PROMPT},
{"role": "user", "content": json.dumps(
[{"role": m.get("role"),
"content": m.get("content") if isinstance(m.get("content"), str) else str(m.get("content"))[:4000]}
for m in middle],
default=str,
)[:60_000]},
]
summary = ""
try:
if client is not None:
resp = client.chat.completions.create(
model=model,
messages=summary_input,
temperature=0.0,
max_tokens=1024,
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
)
_bump_local_calls()
summary = (resp.choices[0].message.content or "").strip()[:2000]
except Exception as exc:
summary = f"[compaction summary failed: {exc!r}; older turns dropped]"
if not summary:
summary = "[no summary produced; older turns dropped]"
n_turns_folded = len(old_turns)
synthetic = {
"role": "system",
"content": f"[turns 1{n_turns_folded} elided: {summary}]",
}
folded = [system_msg, initial_user, synthetic, *tail]
after_stage2_tokens = _estimate_prompt_tokens(folded)
_record_event({
"kind": f"{trace_prefix}_compact",
"stage": "2",
"msgs_before": len(new_messages),
"msgs_after": len(folded),
"before_tokens": after_stage1_tokens,
"after_tokens": after_stage2_tokens,
"n_tool_elided": n_tool_elided,
"n_turns_folded": n_turns_folded,
"summary_chars": len(summary),
"ts": time.time(),
})
return folded
def _loop_local(
problem: str,
workdir: Path,
@@ -905,7 +1221,16 @@ def _loop_local(
output_cap: int,
turn_max_tokens: int,
trace_prefix: str,
compact_at_tokens: int = 22_000,
compact_keep_last: int = 3,
) -> Dict[str, Any]:
# Qwen-27B has a 32k context. With ``max_tokens=turn_max_tokens`` reserved
# for output (default 4096) plus ~1k for the bash tool schema + system
# prompt + format overhead, the practical input ceiling is ~27k. We
# compact at 22k so there's slack for one more tool result before the
# next turn's pre-call check fires again. Earlier we used 24k + keep=4
# but still saw 28k-input 400s on the n=100 SWE sweep (the keep window
# alone routinely exceeded the budget once bash outputs piled up).
from openai import OpenAI
client = OpenAI(base_url=endpoint, api_key="EMPTY", timeout=600.0)
@@ -919,16 +1244,59 @@ def _loop_local(
turns = 0
for turn in range(1, max_turns + 1):
turns = turn
if compact_at_tokens > 0 and _estimate_prompt_tokens(messages) > compact_at_tokens:
messages = _compact_local_messages(
messages, client=client, model=model,
keep_last=compact_keep_last, trace_prefix=trace_prefix,
compact_at_tokens=compact_at_tokens,
)
t0 = time.time()
resp = client.chat.completions.create(
model=model,
messages=messages,
temperature=0.0,
max_tokens=turn_max_tokens,
tools=[BASH_TOOL_OPENAI],
tool_choice="auto",
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
)
try:
resp = client.chat.completions.create(
model=model,
messages=messages,
temperature=0.0,
max_tokens=turn_max_tokens,
tools=[BASH_TOOL_OPENAI],
tool_choice="auto",
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
)
except Exception as exc:
# Emergency compaction on a context-length 400 from vLLM /
# OpenAI ("maximum context length is N tokens"). Our pre-call
# estimator can undercount when tool args / tool_call_ids /
# template overhead spike, so the budget check missed and the
# server walled the call. Compact aggressively (keep_last=1)
# and retry once. Re-raise on anything else or on a second
# failure — the runner records the row as errored.
msg = str(exc)
is_ctx = (
"maximum context length" in msg
or "context length" in msg.lower() and "exceed" in msg.lower()
)
if not is_ctx:
raise
_record_event({
"kind": f"{trace_prefix}_emergency_compact",
"turn": turn,
"error": msg[:300],
"tokens_before": _estimate_prompt_tokens(messages),
"ts": time.time(),
})
messages = _compact_local_messages(
messages, client=client, model=model,
keep_last=1, trace_prefix=trace_prefix,
compact_at_tokens=max(8_000, compact_at_tokens // 2),
)
resp = client.chat.completions.create(
model=model,
messages=messages,
temperature=0.0,
max_tokens=turn_max_tokens,
tools=[BASH_TOOL_OPENAI],
tool_choice="auto",
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
)
_bump_local_calls()
latency = time.time() - t0
u = resp.usage
@@ -954,10 +1322,17 @@ def _loop_local(
"ts": time.time(),
})
messages.append({
# Match the OpenAI cloud branch: content="" (not None) when only
# tool_calls are present; omit ``tool_calls`` entirely when there
# are none (vs. setting it to None) so the message validates
# against the strict OpenAI schema if it ever gets replayed by
# the compactor's summarizer call.
assistant_local_msg: Dict[str, Any] = {
"role": "assistant",
"content": text or None,
"tool_calls": [
"content": text or "",
}
if tool_calls:
assistant_local_msg["tool_calls"] = [
{
"id": tc.id, "type": "function",
"function": {
@@ -966,8 +1341,8 @@ def _loop_local(
},
}
for tc in tool_calls
] if tool_calls else None,
})
]
messages.append(assistant_local_msg)
if not tool_calls:
final_text = text.strip()
+222
View File
@@ -0,0 +1,222 @@
"""Regression tests for the mini-SWE-agent OpenAI cloud + bash adapter.
Two failure modes caught in the n=100 hybrid SWE sweep (May 2026) — both
deterministic enough to pin down here without hitting any real model or
shelling out for real binary data:
1. **Bug 1 — null assistant content.** ``_loop_cloud_openai`` used to
append ``{"role": "assistant", "content": text or None}`` on each
turn. On a tool-only turn (model produced no text alongside its
``bash`` call) that wrote ``content: null`` into the message list;
the next ``chat.completions.create`` then 400'd with
``Invalid value for 'content': expected a string, got null``. Fix
uses ``""`` (or omitted) per OpenAI's schema.
2. **Bug 3 — binary bash output crashes the loop.** ``_run_bash`` used
to pass ``text=True`` to ``subprocess.Popen``, so any command that
emitted non-UTF-8 bytes (cat'ing a compiled artifact, an image, a
PDF) raised ``UnicodeDecodeError`` from inside
``Popen.communicate()`` and killed the whole task. Fix captures
bytes and decodes via ``_decode_bash_output`` with ``errors="replace"``
plus a binary-detection stub.
Run with:
.venv/bin/python -m pytest tests/agents/hybrid/test_openai_adapter.py -v
"""
from __future__ import annotations
from pathlib import Path
from typing import Any, Dict, List
import pytest
from openjarvis.agents.hybrid.mini_swe_agent import (
_decode_bash_output,
_loop_cloud_openai,
_run_bash,
)
# ---------------------------------------------------------------------------
# Bug 3 — binary bash output
# ---------------------------------------------------------------------------
class TestDecodeBashOutput:
def test_pure_ascii_passes_through(self) -> None:
assert _decode_bash_output(b"hello world\n", 0) == "hello world\n"
def test_empty_bytes_returns_empty_string(self) -> None:
assert _decode_bash_output(b"", 0) == ""
def test_nul_byte_substituted_with_stub(self) -> None:
raw = b"some text\x00more bytes after"
out = _decode_bash_output(raw, 0)
assert out.startswith("[binary output:")
assert f"{len(raw)} bytes" in out
assert "exit=0" in out
def test_invalid_utf8_partial_substitutes_replacement_char(self) -> None:
# ~25% replacement chars after decode — well above the 5% binary
# threshold, should swap to the stub. The exact byte sequence
# 0xe0 is the one observed in the astropy__astropy-14539 row.
raw = b"abc\xe0\xe0\xe0def"
out = _decode_bash_output(raw, 1)
assert out.startswith("[binary output:")
assert "exit=1" in out
def test_mostly_valid_utf8_keeps_decoded_text(self) -> None:
# One stray bad byte in a long string → below 5% replacement,
# keep the (mostly intact) decoded text rather than stubbing.
raw = ("readable text " * 200).encode("utf-8") + b"\xe0"
out = _decode_bash_output(raw, 0)
assert "readable text" in out
assert not out.startswith("[binary output:")
def test_real_bash_run_on_binary_does_not_raise(
self, tmp_path: Path
) -> None:
# End-to-end: a model-issued ``head`` on a binary file. Pre-fix
# this raised UnicodeDecodeError out of ``_run_bash`` →
# propagated all the way up to the runner. Post-fix it returns
# a normal observation dict with the binary-output stub.
bin_path = tmp_path / "blob.bin"
bin_path.write_bytes(bytes(range(256)) * 4)
result = _run_bash(
f"head -c 1024 {bin_path}", tmp_path,
timeout=10, output_cap=10_000,
)
assert result["exit_code"] == 0
assert result["timed_out"] is False
assert "[binary output:" in result["stdout"]
# ---------------------------------------------------------------------------
# Bug 1 — null content on tool-only assistant turn
# ---------------------------------------------------------------------------
class _FakeFunction:
def __init__(self, name: str, arguments: str) -> None:
self.name = name
self.arguments = arguments
class _FakeToolCall:
def __init__(self, tc_id: str, name: str, arguments: str) -> None:
self.id = tc_id
self.type = "function"
self.function = _FakeFunction(name, arguments)
class _FakeMessage:
def __init__(self, content: Any, tool_calls: List[Any]) -> None:
self.content = content
self.tool_calls = tool_calls
class _FakeChoice:
def __init__(
self, message: _FakeMessage, finish_reason: str = "tool_calls"
) -> None:
self.message = message
self.finish_reason = finish_reason
class _FakeUsage:
prompt_tokens = 10
completion_tokens = 5
class _FakeResp:
def __init__(self, choice: _FakeChoice) -> None:
self.choices = [choice]
self.usage = _FakeUsage()
class _FakeCompletions:
"""Records every (messages=...) the model is asked to score against.
Returns a scripted sequence: turn 1 → tool_only (no text), turn 2 →
final summary text, no tool_calls. We then assert that turn 2's
inbound messages contain the turn-1 assistant message with
``content == ""`` (or omitted) — never ``None``, which is the bug.
"""
def __init__(self, scripted: List[_FakeResp]) -> None:
self._scripted = list(scripted)
self.calls: List[List[Dict[str, Any]]] = []
def create(self, **kwargs: Any) -> _FakeResp:
# Deep enough copy so the agent appending to its messages list
# post-call doesn't mutate what we recorded here.
self.calls.append([dict(m) for m in kwargs["messages"]])
return self._scripted.pop(0)
class _FakeChat:
def __init__(self, completions: _FakeCompletions) -> None:
self.completions = completions
class _FakeClient:
def __init__(self, completions: _FakeCompletions) -> None:
self.chat = _FakeChat(completions)
def test_assistant_message_content_never_none_on_tool_only_turn(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
"""The OpenAI SDK 400s with ``content: null`` on replay. We must
serialize tool-only assistant turns as ``content == ""`` (or omit
the field) — never ``None``. Pre-fix this test failed because turn
2's recorded messages had ``messages[2]["content"] is None``.
"""
tool_turn = _FakeResp(_FakeChoice(
_FakeMessage(
content=None,
tool_calls=[_FakeToolCall("c1", "bash", '{"command": "echo hi"}')],
),
))
done_turn = _FakeResp(_FakeChoice(
_FakeMessage(content="all done", tool_calls=[]),
finish_reason="stop",
))
fake = _FakeCompletions([tool_turn, done_turn])
fake_client = _FakeClient(fake)
def _fake_openai_ctor(**kwargs: Any) -> _FakeClient:
return fake_client
# Swap in our fake at the import site inside _loop_cloud_openai.
import openai
monkeypatch.setattr(openai, "OpenAI", _fake_openai_ctor)
out = _loop_cloud_openai(
"fake problem", tmp_path,
model="gpt-5-mini-2025-08-07",
max_turns=4, bash_timeout=10, output_cap=10_000,
turn_max_tokens=64, trace_prefix="test",
)
# Two real model calls: tool-issuing turn + final turn.
assert len(fake.calls) == 2
second_call_messages = fake.calls[1]
# First two messages are system + user. The assistant turn from turn
# 1 should be index 2; its content must be a string (empty is fine),
# never None — that's the bug we're regressing.
assistant_msg = second_call_messages[2]
assert assistant_msg["role"] == "assistant"
assert "content" in assistant_msg
assert assistant_msg["content"] is not None
assert isinstance(assistant_msg["content"], str)
# Tool calls must still be present (we only fixed the content shape).
assert assistant_msg.get("tool_calls")
# Sanity: the loop terminated normally on the no-tool turn.
assert out["final_summary"] == "all done"
assert out["turns"] == 2
+312
View File
@@ -0,0 +1,312 @@
"""Smoke tests for the OpenAI SDK retry + per-org concurrency hardening.
We deliberately don't hit the real OpenAI API. Instead we monkey-patch
the underlying call to simulate the failure modes we care about:
- Sustained ``RateLimitError`` walls.
- ``APITimeoutError`` blips.
- ``APIConnectionError`` blips.
- ``InternalServerError`` 5xx blips.
Then we hammer ``openai.OpenAI().chat.completions.create`` from a thread
pool with a deliberately tight semaphore (concurrency=2, retries=4) and
confirm:
(a) no exceptions escape when the failure clears within the retry budget,
(b) backoff actually fires (call count > attempted call count),
(c) all requests eventually complete,
(d) when the failure exceeds the retry budget, the final exception
propagates so the runner records ``error=...`` (no silent drop),
(e) local vLLM-style clients (api_key="EMPTY" or localhost base_url)
bypass both the throttle and the retry loop.
Run with:
.venv/bin/python -m pytest tests/agents/hybrid/test_openai_retry.py -v
"""
from __future__ import annotations
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, List
from unittest.mock import MagicMock
import pytest
# Make sure the patch picks up tight test settings, not the prod defaults.
# Must be set BEFORE _openai_retry is imported.
os.environ["OPENJARVIS_OPENAI_MAX_CONCURRENCY"] = "2"
os.environ["OPENJARVIS_OPENAI_MAX_RETRIES"] = "4"
os.environ["OPENJARVIS_OPENAI_RETRY_BASE"] = "0.05" # fast tests
os.environ["OPENJARVIS_OPENAI_RETRY_CAP"] = "0.2"
# Reload the module fresh so env vars take effect (imports earlier in
# the process may have frozen the defaults).
import importlib
from openjarvis.agents.hybrid import _openai_retry as _retry_mod
importlib.reload(_retry_mod)
_retry_mod.patch_openai_globally()
def _make_fake_response() -> Any:
"""Minimal stand-in for an OpenAI ChatCompletion response."""
r = MagicMock()
r.choices = [MagicMock(message=MagicMock(content="ok", tool_calls=None))]
r.usage = MagicMock(prompt_tokens=1, completion_tokens=1)
return r
def _swap_underlying_create(
monkeypatch: pytest.MonkeyPatch, side_effect: Any,
) -> List[int]:
"""Replace the wrapped-underlying-orig with one that fires ``side_effect``.
Returns a mutable call-counter list so tests can assert on attempts.
"""
from openai.resources.chat import completions as _comp_mod
counter: List[int] = [0]
wrapped = _comp_mod.Completions.create
# The patcher stashed the original at ``__wrapped__``.
orig = getattr(wrapped, "__wrapped__", None)
assert orig is not None, "patch_openai_globally didn't stash __wrapped__"
def fake(self: Any, *args: Any, **kwargs: Any) -> Any:
counter[0] += 1
if callable(side_effect):
return side_effect(counter[0])
if isinstance(side_effect, list):
i = min(counter[0] - 1, len(side_effect) - 1)
v = side_effect[i]
if isinstance(v, BaseException):
raise v
return v
if isinstance(side_effect, BaseException):
raise side_effect
return side_effect
# Replace the wrapped underlying directly.
new_wrap = _retry_mod._wrap_create(fake)
monkeypatch.setattr(_comp_mod.Completions, "create", new_wrap)
return counter
def _rate_limit_error() -> BaseException:
"""Build an ``openai.RateLimitError`` that the SDK would normally raise."""
import openai
# The SDK's RateLimitError wants (message, response, body) — we use a
# MagicMock for response so ``response.headers.get("retry-after")``
# returns None.
resp = MagicMock()
resp.headers = {}
resp.status_code = 429
return openai.RateLimitError("rate limit", response=resp, body=None)
def _api_timeout_error() -> BaseException:
import openai
return openai.APITimeoutError(request=MagicMock())
def _api_conn_error() -> BaseException:
import openai
return openai.APIConnectionError(request=MagicMock())
def _internal_500_error() -> BaseException:
import openai
resp = MagicMock()
resp.headers = {}
resp.status_code = 500
return openai.InternalServerError("server error", response=resp, body=None)
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
def test_patches_installed() -> None:
import openai
from openai.resources.chat import completions
assert getattr(completions.Completions.create, "_hybrid_patched", False)
assert getattr(openai.OpenAI.__init__, "_hybrid_patched", False)
def test_rate_limit_then_success(monkeypatch: pytest.MonkeyPatch) -> None:
"""Two 429s then a real response — retry path should swallow both."""
import openai
counter = _swap_underlying_create(
monkeypatch,
[_rate_limit_error(), _rate_limit_error(), _make_fake_response()],
)
client = openai.OpenAI(api_key="sk-fake")
resp = client.chat.completions.create(
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}]
)
assert resp.choices[0].message.content == "ok"
assert counter[0] == 3 # backoff fired twice, then succeeded
def test_timeout_then_success(monkeypatch: pytest.MonkeyPatch) -> None:
import openai
counter = _swap_underlying_create(
monkeypatch, [_api_timeout_error(), _make_fake_response()]
)
client = openai.OpenAI(api_key="sk-fake")
resp = client.chat.completions.create(
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}]
)
assert resp.choices[0].message.content == "ok"
assert counter[0] == 2
def test_500_then_success(monkeypatch: pytest.MonkeyPatch) -> None:
import openai
counter = _swap_underlying_create(
monkeypatch, [_internal_500_error(), _make_fake_response()]
)
client = openai.OpenAI(api_key="sk-fake")
resp = client.chat.completions.create(
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}]
)
assert resp.choices[0].message.content == "ok"
assert counter[0] == 2
def test_retry_exhaustion_propagates(monkeypatch: pytest.MonkeyPatch) -> None:
"""Sustained 429 wall beyond retry budget → exception propagates."""
import openai
# max_retries=4 → 5 total attempts; feed 6 errors so the loop runs out.
errors = [_rate_limit_error() for _ in range(6)]
counter = _swap_underlying_create(monkeypatch, errors)
client = openai.OpenAI(api_key="sk-fake")
with pytest.raises(openai.RateLimitError):
client.chat.completions.create(
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}]
)
# 5 attempts: 1 initial + 4 retries.
assert counter[0] == 5
def test_non_retryable_propagates_immediately(monkeypatch: pytest.MonkeyPatch) -> None:
"""``BadRequestError`` is NOT retryable — should raise on attempt 1."""
import openai
resp = MagicMock()
resp.headers = {}
resp.status_code = 400
bad = openai.BadRequestError("nope", response=resp, body=None)
counter = _swap_underlying_create(monkeypatch, [bad])
client = openai.OpenAI(api_key="sk-fake")
with pytest.raises(openai.BadRequestError):
client.chat.completions.create(
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}]
)
assert counter[0] == 1 # no retries
def test_local_vllm_bypasses_throttle_and_retry(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Local vLLM clients (api_key=EMPTY) must NOT pay the per-org throttle.
Verifies (a) detection works, (b) a non-retryable error fires only
once (no retry attempts for local endpoints — they're either up or
down). Also (c) the semaphore isn't acquired (which we can't directly
observe, but we can confirm the wrapper short-circuits).
"""
import openai
counter = _swap_underlying_create(monkeypatch, _rate_limit_error())
client = openai.OpenAI(base_url="http://localhost:8001/v1", api_key="EMPTY")
# Should raise on attempt 1 (no retries for local endpoints).
with pytest.raises(openai.RateLimitError):
client.chat.completions.create(
model="qwen", messages=[{"role": "user", "content": "hi"}]
)
assert counter[0] == 1
def test_concurrent_hammering_no_escape(monkeypatch: pytest.MonkeyPatch) -> None:
"""16 concurrent calls, each fails twice with 429 then succeeds.
With concurrency=2 the semaphore queues, with retry budget=4 every
call eventually gets through. We confirm no exception escapes and
all 16 calls complete with the fake response.
"""
import openai
# Per-call counter for failure injection.
call_state: dict = {}
state_lock = threading.Lock()
def side_effect(call_idx: int) -> Any:
# Tag by thread so each "logical call" fails its first 2 attempts.
tid = threading.get_ident()
with state_lock:
n = call_state.get(tid, 0) + 1
call_state[tid] = n
if n <= 2:
raise _rate_limit_error()
# Reset for the next logical call from this thread.
with state_lock:
call_state[tid] = 0
return _make_fake_response()
_swap_underlying_create(monkeypatch, side_effect)
client = openai.OpenAI(api_key="sk-fake")
def one() -> str:
r = client.chat.completions.create(
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}]
)
return r.choices[0].message.content
t0 = time.time()
with ThreadPoolExecutor(max_workers=16) as ex:
futures = [ex.submit(one) for _ in range(16)]
results = [f.result() for f in as_completed(futures)]
elapsed = time.time() - t0
assert all(r == "ok" for r in results)
assert len(results) == 16
# Sanity: with concurrency=2 + backoff>=0.05s * 2 retries per call,
# 16 calls can't possibly finish instantly. Just confirm we spent
# *some* time in backoff, not that we measured it precisely.
assert elapsed > 0.1
def test_retry_after_honored(monkeypatch: pytest.MonkeyPatch) -> None:
"""If the SDK exposes a Retry-After header, we sleep at least that long."""
import openai
resp = MagicMock()
resp.headers = {"retry-after": "0.15"}
resp.status_code = 429
err = openai.RateLimitError("rate limit", response=resp, body=None)
counter = _swap_underlying_create(monkeypatch, [err, _make_fake_response()])
client = openai.OpenAI(api_key="sk-fake")
t0 = time.time()
client.chat.completions.create(
model="gpt-5-mini", messages=[{"role": "user", "content": "hi"}]
)
elapsed = time.time() - t0
assert counter[0] == 2
# We sleep at least 0.15s (Retry-After) — jitter is additive ≤ 0.5.
assert elapsed >= 0.15
+153
View File
@@ -0,0 +1,153 @@
"""Tests for trajectory compaction in mini_swe_agent._compact_local_messages."""
from __future__ import annotations
import copy
from typing import Any, Dict, List
import pytest
from openjarvis.agents.hybrid.mini_swe_agent import (
_compact_local_messages,
_estimate_prompt_tokens,
_get_tiktoken_enc,
)
TIKTOKEN_OK = bool(_get_tiktoken_enc())
def _validate_openai_message_shape(messages: List[Dict[str, Any]]) -> None:
seen_tool_call_ids: set[str] = set()
for i, m in enumerate(messages):
role = m.get("role")
assert role in {"system", "user", "assistant", "tool"}, f"bad role at {i}: {role}"
if role == "assistant":
for tc in (m.get("tool_calls") or []):
tid = tc.get("id")
assert tid, f"assistant tool_call missing id at msg {i}"
seen_tool_call_ids.add(tid)
elif role == "tool":
tid = m.get("tool_call_id")
assert tid, f"tool message at {i} missing tool_call_id"
assert tid in seen_tool_call_ids, (
f"tool message at {i} references unknown tool_call_id={tid!r}"
)
assert m.get("content"), f"tool message at {i} has empty content"
def _make_synthetic_messages() -> List[Dict[str, Any]]:
messages: List[Dict[str, Any]] = [
{"role": "system", "content": "You are a coding agent."},
{"role": "user", "content": "Fix the bug in foo.py."},
]
big_idx = 7 # one of the assistant+tool pairs gets a giant blob
for t in range(19):
tc_id = f"call_{t:03d}"
messages.append({
"role": "assistant",
"content": f"Turn {t}: I'll run a command.",
"tool_calls": [{
"id": tc_id,
"type": "function",
"function": {
"name": "bash",
"arguments": '{"command": "ls -la dir_' + str(t) + '"}',
},
}],
})
if t == big_idx:
body = "OUTPUT " + ("x" * 29950)
obs = f"$ ls\n{body}\nexit_code=0"
else:
body = ("line " + str(t) + " ") * 800 # ~8000 chars
obs = f"$ ls\n{body}\nexit_code=0"
messages.append({
"role": "tool",
"tool_call_id": tc_id,
"content": obs[:8000] if t != big_idx else obs[:30000],
})
return messages
def test_compaction_token_budget_and_shape():
messages = _make_synthetic_messages()
orig = copy.deepcopy(messages)
before = _estimate_prompt_tokens(messages)
assert before > 24_000, f"sanity: synthetic should exceed budget, got {before}"
compacted = _compact_local_messages(
messages,
client=None, # forces stage-2 summary fallback (no API call)
model="dummy",
keep_last=4,
trace_prefix="test",
compact_at_tokens=24_000,
)
_validate_openai_message_shape(compacted)
if TIKTOKEN_OK:
after = _estimate_prompt_tokens(compacted)
assert after <= 24_000, f"after compaction tokens={after} > 24_000"
# System + initial user intact.
assert compacted[0] == orig[0]
assert compacted[1] == orig[1]
# Recent 4 turns (assistant + its tool replies) intact and deep-equal.
# In synthetic input, every turn is exactly 1 assistant + 1 tool, so the
# last 4 turns = last 8 messages.
assert compacted[-8:] == orig[-8:], "recent 4 turns must be intact"
def test_stage1_only_when_sufficient():
"""If stage 1 alone drops us below the budget, stage 2 must not run."""
# Build a smaller input that exceeds budget only because of giant tool outputs.
messages: List[Dict[str, Any]] = [
{"role": "system", "content": "sys"},
{"role": "user", "content": "task"},
]
for t in range(8):
tc_id = f"call_{t}"
messages.append({
"role": "assistant",
"content": f"t{t}",
"tool_calls": [{
"id": tc_id, "type": "function",
"function": {"name": "bash", "arguments": "{}"},
}],
})
# Use random-ish text to avoid BPE collapse from repetition.
import random
rng = random.Random(t)
body = " ".join(rng.choice(["alpha", "beta", "gamma", "delta", "lambda",
"foo", "bar", "baz", "qux", "zeta"])
for _ in range(3500))
messages.append({
"role": "tool",
"tool_call_id": tc_id,
"content": body + " exit_code=0",
})
before = _estimate_prompt_tokens(messages)
if not TIKTOKEN_OK:
pytest.skip("tiktoken unavailable; stage-1-only ordering is timing-sensitive without it")
assert before > 24_000
compacted = _compact_local_messages(
messages, client=None, model="dummy",
keep_last=4, trace_prefix="test", compact_at_tokens=24_000,
)
_validate_openai_message_shape(compacted)
# No stage-2 fold → length identical (only tool contents shortened).
assert len(compacted) == len(messages)
# The synthetic stage-2 system stub should NOT be present.
for m in compacted:
if m.get("role") == "system":
assert not str(m.get("content", "")).startswith("[turns "), \
"stage 2 should not have run"
if __name__ == "__main__":
import sys
sys.exit(pytest.main([__file__, "-v"]))