fix(agents): survive long overnight runs

- compact at memory-safe token budget (~20k) not the 131k window; gpt-oss
  hits oMLX's prefill memory guard ~47k tokens on a 24GB machine
- repair invalid_tool_calls (malformed JSON args) before echoing history —
  LangChain serializes the raw malformed string which oMLX 500s on
- retry invokes (4x) with an oMLX watchdog restart between attempts
- ALWAYS finalize: salvage picks from findings even after invoke failures

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_016DvdRZkyCy9NJpHYDdPqmS
This commit is contained in:
Aryan SinghandClaude Opus 4.8 committed 2026-07-02 12:22:29 -05:00
1 parent 3bf4a6f63c
commit 6b1c461a80
4 files changed
+104 -12

No files matched your search

+95 -10
View File
@@ -89,6 +89,55 @@ def _short(args: dict) -> str:
return str(args)[:60]
def _sanitize_tool_calls(messages: list) -> None:
"""Re-encode historical tool-call arguments to valid JSON, in place.
gpt-oss sometimes emits a tool call (especially run_analysis with a big multi-line code
string) whose raw `arguments` aren't valid JSON. It executes fine once, but when that
assistant turn is echoed back in the next request, oMLX's FunctionCall validator rejects
it with a 500 ('arguments must be valid JSON') and the whole agent dies. We rewrite the
stored raw arguments from the cleanly-parsed args dict so every echo is valid JSON.
"""
for m in messages:
ak = getattr(m, "additional_kwargs", None)
if not isinstance(ak, dict):
continue
raws = ak.get("tool_calls")
if not raws:
continue
parsed = getattr(m, "tool_calls", None) or []
for i, raw in enumerate(raws):
fn = raw.get("function")
if not isinstance(fn, dict):
continue
args = parsed[i]["args"] if i < len(parsed) and isinstance(parsed[i].get("args"), dict) else {}
fn["arguments"] = json.dumps(args)
def _repair_invalid_tool_calls(ai) -> list[tuple[str, str]]:
"""Convert `invalid_tool_calls` (malformed-JSON args) into valid empty-args calls, in place.
LangChain serializes invalid_tool_calls with their RAW malformed argument string — which
oMLX's validator 500s on at the next request (this bypasses additional_kwargs, so the
sanitizer above can't catch it). We convert each into a valid call with args={} and return
(id, name) pairs so the loop can attach an error ToolMessage telling the model to retry
with valid JSON.
"""
import uuid
repaired: list[tuple[str, str]] = []
invalid = getattr(ai, "invalid_tool_calls", None) or []
if not invalid:
return repaired
for itc in invalid:
name = ((itc.get("name") or "unknown").split("<|")[0].strip()) or "unknown"
tc_id = itc.get("id") or f"call_{uuid.uuid4().hex[:8]}"
ai.tool_calls.append({"name": name, "args": {}, "id": tc_id, "type": "tool_call"})
repaired.append((tc_id, name))
ai.invalid_tool_calls = []
return repaired
def _prompt_tokens(ai, messages: list) -> int:
"""Best available count of tokens the last prompt consumed (for compaction timing)."""
um = getattr(ai, "usage_metadata", None) or {}
@@ -197,7 +246,10 @@ def run_agent(
tools_by_name = {t.name: t for t in tools}
tool_llm = build_llm(tools=tools)
plain_llm = build_llm() # no tools — used for summarization
ctx_len = get_settings().llm_context_length
# Compaction ceiling is the MEMORY-safe KV budget, not the model's 131k window: gpt-oss
# hits oMLX's prefill memory guard around ~47k tokens on a 24GB Mac, so we compact well
# below that to keep long (2hr) runs alive.
compact_budget = get_settings().compaction_token_budget
date_str = datetime.now().astimezone().strftime("%A %Y-%m-%d")
messages: list = [
@@ -231,15 +283,43 @@ def run_agent(
ended_reason = "timeout"
break
steps_used += 1
try:
ai = tool_llm.invoke(messages)
except Exception as e:
reasoning_trace.append(f"[error] llm.invoke: {e}")
_sanitize_tool_calls(messages) # keep echoed tool-call args valid JSON (oMLX 500 guard)
ai = None
for attempt in range(4):
try:
ai = tool_llm.invoke(messages)
break
except Exception as e:
err = str(e)[:100]
reasoning_trace.append(f"[retry {attempt + 1}/4] invoke failed: {err}")
logger.warning("{}: invoke failed (attempt {}/4): {}", persona.id, attempt + 1, err)
time.sleep(4)
if attempt >= 1: # watchdog: restart oMLX if it's wedged, then keep going
from llm import server as llm_server
if not llm_server.server_up():
logger.warning("{}: oMLX down — watchdog restarting it", persona.id)
llm_server.ensure_ready(warm=False)
if ai is None:
reasoning_trace.append("[error] llm.invoke failed after 4 retries")
ended_reason = "error"
logger.warning("{}: llm.invoke failed: {}", persona.id, e)
logger.warning("{}: giving up research after 4 failed invokes (will still finalize)", persona.id)
break
messages.append(ai)
# malformed-JSON tool calls: make them serializable + tell the model to retry
repaired = _repair_invalid_tool_calls(ai)
repaired_ids = {tc_id for tc_id, _ in repaired}
for tc_id, name in repaired:
messages.append(
ToolMessage(
content=f"ERROR: your {name} call had malformed JSON arguments. "
"Call it again — arguments must be a valid JSON object.",
tool_call_id=tc_id,
)
)
reasoning_trace.append(f"[repair] malformed {name} args -> asked model to retry")
tool_calls = getattr(ai, "tool_calls", None)
if not tool_calls:
# Force depth: if the model tries to finalize before the TIME floor (or step floor),
@@ -259,6 +339,8 @@ def run_agent(
break
for tc in tool_calls:
if tc["id"] in repaired_ids: # already answered with a retry-error ToolMessage
continue
# gpt-oss can leak harmony channel tokens into the tool name (e.g. "get_bars<|channel|>")
name = tc["name"].split("<|")[0].strip()
args = tc.get("args", {}) or {}
@@ -270,7 +352,7 @@ def run_agent(
messages.append(ToolMessage(content=str(result), tool_call_id=tc["id"]))
reasoning_trace.append(f"{name}({_short(args)}) -> {str(result)[:160]}")
if compaction.should_compact(_prompt_tokens(ai, messages), ctx_len):
if compaction.should_compact(_prompt_tokens(ai, messages), compact_budget):
messages, running_summary = compaction.compact(plain_llm, messages, running_summary)
if checkpointer: # incremental save after each round
@@ -279,14 +361,17 @@ def run_agent(
if steps_used >= step_budget and ended_reason == "natural":
ended_reason = "budget"
# --- finalize (always attempted, even on timeout, to salvage picks) ---
# --- finalize (ALWAYS attempted — even after invoke failures, salvage the findings) ---
picks: list[Pick] = []
errors: list[str] = []
if ended_reason != "error":
try:
picks, errors = _finalize(persona, running_summary, reasoning_trace)
except Exception as e:
errors = [f"finalize crashed: {e}"]
if ended_reason == "error":
status = AgentStatus.ERROR
# salvaged picks still count — mark partial so execution can use them
status = AgentStatus.PARTIAL if picks else AgentStatus.ERROR
elif ended_reason == "timeout":
status = AgentStatus.TIMEOUT
elif picks: