From 7b79117fbf5b0010d0785b80419c1ff6d5950683 Mon Sep 17 00:00:00 2001 From: pewdiepie-archdaemon Date: Fri, 18 Sep 2026 14:29:59 +0000 Subject: [PATCH] fix(agent): reserve artifact budget after scratch writes --- src/clean_agent_preview.py | 2 +- tests/test_clean_agent_preview.py | 18 +++++++++++++----- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/src/clean_agent_preview.py b/src/clean_agent_preview.py index 4a3e981ab..77ae86cf5 100644 --- a/src/clean_agent_preview.py +++ b/src/clean_agent_preview.py @@ -4268,7 +4268,7 @@ async def stream_preview(*, endpoint_url, model, messages, headers, turn_contrac yield event({'type': 'completion_recovery', 'reason': 'research_before_synthesis'}) if ( required_artifacts - and not successful_write + and not successful_artifact_write and not artifact_write_phase and calls >= min(NATIVE_ARTIFACT_RESEARCH_LIMIT, tool_call_limit - 1) ): diff --git a/tests/test_clean_agent_preview.py b/tests/test_clean_agent_preview.py index 94ad0ada6..3d72dad66 100644 --- a/tests/test_clean_agent_preview.py +++ b/tests/test_clean_agent_preview.py @@ -4551,13 +4551,19 @@ async def test_blocked_search_engine_browser_forces_native_web_search(monkeypatc async def test_native_stream_reserves_remaining_budget_for_required_artifact(monkeypatch): import src.clean_agent_preview as module - payloads = [ + payloads = [{"choices": [{"delta": {"tool_calls": [{ + "index": 0, "id": "scratch-write", + "function": {"name": "python", "arguments": json.dumps({ + "code": "open('/tmp_workspace/scratch.txt', 'w').write('notes')", + })}, + }]}}]}] + payloads.extend([ {"choices": [{"delta": {"tool_calls": [{ "index": 0, "id": f"search-{index}", "function": {"name": "web_search", "arguments": json.dumps({"query": f"topic {index}"})}, }]}}]} - for index in range(module.NATIVE_ARTIFACT_RESEARCH_LIMIT) - ] + for index in range(module.NATIVE_ARTIFACT_RESEARCH_LIMIT - 1) + ]) payloads.extend([ {"choices": [{"delta": {"tool_calls": [{ "index": 0, "id": "write", @@ -4594,7 +4600,7 @@ async def test_native_stream_reserves_remaining_budget_for_required_artifact(mon monkeypatch.setattr(module, "execute_tool_block", execute) schemas = [ item for item in FUNCTION_TOOL_SCHEMAS - if item["function"]["name"] in {"web_search", "write_file"} + if item["function"]["name"] in {"python", "web_search", "write_file"} ] contract = resolve_full_inventory_contract(schemas=schemas, policy=ToolPolicy()) raw = [chunk async for chunk in stream_preview( @@ -4616,7 +4622,9 @@ async def test_native_stream_reserves_remaining_budget_for_required_artifact(mon and event.get("reason") == "artifact_write_budget_reserved" ) assert recovery["calls_used"] == module.NATIVE_ARTIFACT_RESEARCH_LIMIT - assert [tool["function"]["name"] for tool in requests[12]["tools"]] == ["write_file"] + reserved_names = [tool["function"]["name"] for tool in requests[12]["tools"]] + assert "write_file" in reserved_names + assert "web_search" not in reserved_names assert requests[12]["tool_choice"] == { "type": "function", "function": {"name": "write_file"}, }