fix(runtime): complete logical LLM calls after acceptance

Signed-off-by: Alex Fournier <afournier@nvidia.com>
This commit is contained in:
Alex Fournier 2026-07-20 11:01:14 -07:00
parent 0ddbb1f2bf
commit a15b98f414
4 changed files with 103 additions and 29 deletions

View file

@ -2440,6 +2440,7 @@ def _relay_sync_completion(
name=provider_name,
model_name=str(kwargs.get("model") or fallback_model),
metadata=metadata,
defer_logical_completion=True,
)
@ -2462,6 +2463,7 @@ async def _relay_async_completion(
name=provider_name,
model_name=str(kwargs.get("model") or fallback_model),
metadata=metadata,
defer_logical_completion=True,
)
@ -7017,6 +7019,7 @@ def _validate_llm_response(
except (AttributeError, TypeError, IndexError) as exc:
recovered = _recover_aux_response_message(response)
if recovered is not None:
_complete_relay_auxiliary_call()
return recovered
response_type = type(response).__name__
response_preview = str(response)[:120]
@ -7026,9 +7029,23 @@ def _validate_llm_response(
f"Expected object with .choices[0].message — check provider "
f"adapter or custom endpoint compatibility."
) from exc
_complete_relay_auxiliary_call()
return response
def _complete_relay_auxiliary_call() -> None:
"""Close one auxiliary logical call only after Hermes accepts its response."""
context = _RELAY_AUX_CALL_CONTEXT.get()
if context is None:
return
from agent import relay_llm
relay_llm.complete_logical_call(
str(context.get("request_id") or ""),
outcome="success",
)
def _recover_aux_response_message(response: Any) -> Optional[Any]:
"""Synthesize chat-completions shape from Responses-style text fields.

View file

@ -1746,13 +1746,6 @@ def run_conversation(
)
continue # Retry the API call
from agent import relay_llm
relay_llm.complete_logical_call(
api_request_id,
outcome="success",
)
# Check finish_reason before proceeding
if agent.api_mode == "codex_responses":
status = getattr(response, "status", None)
@ -2445,6 +2438,12 @@ def run_conversation(
clear_nous_rate_limit()
except Exception:
pass
from agent import relay_llm
relay_llm.complete_logical_call(
api_request_id,
outcome="success",
)
agent._touch_activity(f"API call #{api_call_count} completed")
break # Success, exit retry loop

View file

@ -29,10 +29,17 @@ def relay_turn(tmp_path, monkeypatch):
def test_auxiliary_retries_share_logical_relay_identity(monkeypatch):
attempts = []
logical_completions = []
responses = iter([
SimpleNamespace(choices=[]),
SimpleNamespace(
choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))]
),
])
client = SimpleNamespace(
chat=SimpleNamespace(
completions=SimpleNamespace(
create=lambda **kwargs: {"request": kwargs},
create=lambda **_kwargs: next(responses),
)
)
)
@ -42,6 +49,13 @@ def test_auxiliary_retries_share_logical_relay_identity(monkeypatch):
return callback(request)
monkeypatch.setattr(relay_llm, "execute_current", execute_current)
monkeypatch.setattr(
relay_llm,
"complete_logical_call",
lambda request_id, *, outcome: logical_completions.append(
(request_id, outcome)
),
)
@auxiliary_client._relay_auxiliary_call
def run(task):
@ -50,33 +64,46 @@ def test_auxiliary_retries_share_logical_relay_identity(monkeypatch):
"test-model",
"chat_completions",
)
first = auxiliary_client._relay_sync_completion(
client,
{"model": "test-model", "messages": []},
with pytest.raises(RuntimeError, match="invalid response"):
auxiliary_client._validate_llm_response(
auxiliary_client._relay_sync_completion(
client,
{"model": "test-model", "messages": []},
),
task,
)
return auxiliary_client._validate_llm_response(
auxiliary_client._relay_sync_completion(
client,
{"model": "test-model", "messages": []},
),
task,
)
second = auxiliary_client._relay_sync_completion(
client,
{"model": "test-model", "messages": []},
)
return first, second
first, second = run("compression")
result = run("compression")
assert first["request"]["model"] == "test-model"
assert second["request"]["model"] == "test-model"
assert result.choices[0].message.content == "ok"
assert attempts[0]["metadata"]["api_request_id"] == (
attempts[1]["metadata"]["api_request_id"]
)
assert [attempt["metadata"]["retry_count"] for attempt in attempts] == [0, 1]
assert attempts[0]["metadata"]["call_role"] == "auxiliary:compression"
assert all(attempt["defer_logical_completion"] is True for attempt in attempts)
assert logical_completions == [
(attempts[0]["metadata"]["api_request_id"], "success")
]
@pytest.mark.asyncio
async def test_async_auxiliary_attempt_uses_inherited_relay_adapter(monkeypatch):
captured = {}
logical_completions = []
async def create(**kwargs):
return {"request": kwargs}
return SimpleNamespace(
request=kwargs,
choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))],
)
client = SimpleNamespace(
chat=SimpleNamespace(completions=SimpleNamespace(create=create))
@ -91,6 +118,13 @@ async def test_async_auxiliary_attempt_uses_inherited_relay_adapter(monkeypatch)
"execute_current_async",
execute_current_async,
)
monkeypatch.setattr(
relay_llm,
"complete_logical_call",
lambda request_id, *, outcome: logical_completions.append(
(request_id, outcome)
),
)
@auxiliary_client._relay_auxiliary_call_async
async def run(task):
@ -99,16 +133,23 @@ async def test_async_auxiliary_attempt_uses_inherited_relay_adapter(monkeypatch)
"claude-test",
"chat_completions",
)
return await auxiliary_client._relay_async_completion(
client,
{"model": "claude-test", "messages": []},
return auxiliary_client._validate_llm_response(
await auxiliary_client._relay_async_completion(
client,
{"model": "claude-test", "messages": []},
),
task,
)
result = await run("title_generation")
assert result["request"]["model"] == "claude-test"
assert result.request["model"] == "claude-test"
assert captured["name"] == "anthropic"
assert captured["metadata"]["call_role"] == "auxiliary:title_generation"
assert captured["defer_logical_completion"] is True
assert logical_completions == [
(captured["metadata"]["api_request_id"], "success")
]
def test_auxiliary_stream_uses_streaming_relay_primitive(monkeypatch):
@ -149,7 +190,11 @@ def test_auxiliary_attempt_uses_real_relay_request_intercepts(relay_turn):
chat=SimpleNamespace(
completions=SimpleNamespace(
create=lambda **kwargs: captured_requests.append(kwargs)
or {"content": "ok"},
or SimpleNamespace(
choices=[
SimpleNamespace(message=SimpleNamespace(content="ok"))
]
),
)
)
)
@ -172,15 +217,18 @@ def test_auxiliary_attempt_uses_real_relay_request_intercepts(relay_turn):
"test-model",
"chat_completions",
)
return auxiliary_client._relay_sync_completion(
client,
{"model": "test-model", "messages": []},
return auxiliary_client._validate_llm_response(
auxiliary_client._relay_sync_completion(
client,
{"model": "test-model", "messages": []},
),
task,
)
result = run("compression")
finally:
relay.intercepts.deregister_llm_request("hermes-auxiliary-request")
assert result == {"content": "ok"}
assert result.choices[0].message.content == "ok"
assert captured_requests[0]["temperature"] == 0.25
assert turn.logical_llm_calls == {}

View file

@ -4245,6 +4245,7 @@ class TestRunConversation:
usage=None,
)
hook_events = []
logical_completions = []
def _fake_activate(reason=None):
agent._fallback_index = len(agent._fallback_chain)
@ -4256,6 +4257,12 @@ class TestRunConversation:
patch.object(agent, "_run_codex_stream", side_effect=[content_filter_response, fallback_response]) as mock_run_codex_stream,
patch.object(agent, "_try_activate_fallback", side_effect=_fake_activate) as mock_try_activate_fallback,
patch.object(agent, "_invoke_api_request_error_hook", side_effect=lambda **kw: hook_events.append(kw)),
patch(
"agent.relay_llm.complete_logical_call",
side_effect=lambda request_id, *, outcome: logical_completions.append(
(request_id, outcome)
),
),
patch.object(agent, "_persist_session"),
patch.object(agent, "_save_trajectory"),
patch.object(agent, "_cleanup_task_resources"),
@ -4269,6 +4276,9 @@ class TestRunConversation:
assert hook_events[0]["error_type"] == "ContentPolicyBlocked"
assert hook_events[0]["retryable"] is False
assert hook_events[0]["reason"] == FailoverReason.content_policy_blocked.value
assert logical_completions == [
(hook_events[0]["api_request_id"], "success")
]
def test_ollama_small_runtime_context_fails_before_api_call(self, agent, caplog):
self._setup_agent(agent)