From dc16598a4fda2a74bf3de99e206f3d0a33249ca7 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 26 Jul 2026 13:29:40 +0000 Subject: [PATCH] Budget the whole research prompt against the loaded context for PR #7219 Only the synthesis evidence was budgeted, so the budget could not prevent the overflow it existed to prevent. Measured at head with a realistic prompt (40-source catalog, 12-step plan): the untrimmable scaffolding is about 7,900 chars and the conversation context adds up to 12,000 more. On a 4096-token context, which is the GGUF auto-fit floor and the transformers default, the synthesis request came to about 1.7x the window. Worse, _synthesis_evidence_budget computed usable_tokens = 0 at or below the 4,096-token reserve and then returned the 1,500-char floor anyway, so it added evidence to a prompt that already did not fit. The decision prompt had no context awareness at all: a fixed evidence[-60000:], roughly ten times a small window, on every step rather than once at the end. Overflow is not cosmetic here. It either silently truncates and degenerates the report, as the comment above these constants already warned, or fails the run, and a failed run is only recoverable via retry, which deletes every plan step, source and document source and nulls the report. Both paths now share _prompt_char_budget plus _trimmable_budget: each trimmable section is measured against what the rest of the prompt leaves, and can reach 0 instead of a floor, because a shorter report beats a destroyed run. Evidence is budgeted before the chat history, since the evidence is the report. Unknown context still keeps the full cap. At 4096 tokens the synthesis prompt now fits (0.6x). Below that it is still over, since a 40-source catalog alone exceeds the window; that needs a smaller maxSources, and the context box does accept values down to 128. test_synthesis_evidence_budget_tracks_loaded_context asserted the old floor at 2048 tokens, which is the bug, so it now asserts 0 and that the rest of the prompt counts against the same budget. Verified: 2325 passed across the research/web/sandbox/chat-history/rag/tool suites. The test_mcp_stdio_sessions failure is pre-existing and fails identically with these changes stashed. --- studio/backend/core/research_runs.py | 100 +++++++++++++----- .../tests/test_research_runs_hardening.py | 17 +++ .../tests/test_research_runs_storage.py | 12 ++- 3 files changed, 100 insertions(+), 29 deletions(-) diff --git a/studio/backend/core/research_runs.py b/studio/backend/core/research_runs.py index 2e7437e970..959fbc0c19 100644 --- a/studio/backend/core/research_runs.py +++ b/studio/backend/core/research_runs.py @@ -90,10 +90,10 @@ _MAX_CONTEXT_CHARS = 12_000 _MAX_CONTEXT_MESSAGE_CHARS = 4_000 _MAX_SYNTHESIS_EVIDENCE_CHARS = 32_000 # The synthesis prompt must fit the loaded context or it is silently truncated and the report -# degenerates (echoes the evidence tail). Studio defaults context to 2048 tokens, far below the -# cap above, so the evidence budget adapts to the loaded context: reserve tokens for the prompt -# scaffolding (system prompt, plan, source catalogs) AND the generated report, then convert the -# remainder to chars. Unknown context keeps the full cap. +# degenerates (echoes the evidence tail). GGUF auto-fit floors at 4096 and transformers models +# default to 4096, but the context box accepts anything from 128 up, so the budget adapts: the +# reserve covers the generated report, and every trimmable section is measured against what the +# untrimmable scaffolding leaves. Unknown context keeps the full cap. _MIN_SYNTHESIS_EVIDENCE_CHARS = 1_500 _SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN = 3.0 _SYNTHESIS_CONTEXT_RESERVE_TOKENS = 4_096 @@ -537,15 +537,33 @@ def _local_model_ready() -> bool: return not probed -def _synthesis_evidence_budget() -> int: - """Char budget for synthesis evidence, sized to fit the loaded context (falls back to the - full cap when the context is unknown).""" +def _prompt_char_budget(reserve_tokens: int) -> int | None: + """Chars the whole prompt may occupy on the loaded context, or None when it is unknown.""" ctx = _loaded_context_length() if not ctx: - return _MAX_SYNTHESIS_EVIDENCE_CHARS - usable_tokens = max(0, ctx - _SYNTHESIS_CONTEXT_RESERVE_TOKENS) - budget = int(usable_tokens * _SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN) - return max(_MIN_SYNTHESIS_EVIDENCE_CHARS, min(budget, _MAX_SYNTHESIS_EVIDENCE_CHARS)) + return None + return int(max(0, ctx - reserve_tokens) * _SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN) + + +def _trimmable_budget(total: int | None, fixed_chars: int, hard_cap: int) -> int: + """Chars left for a trimmable section once the rest of the prompt is counted. + + Budgeting one section against the context while the others are unbounded does not stop an + overflow: at a 2048-token context the untrimmable scaffolding alone is several times the + window. Returns 0 rather than a floor, since a short report beats a failed run. + """ + if total is None: + return hard_cap + return max(0, min(hard_cap, total - fixed_chars)) + + +def _synthesis_evidence_budget(fixed_chars: int = 0) -> int: + """Char budget for synthesis evidence (full cap when the context is unknown).""" + return _trimmable_budget( + _prompt_char_budget(_SYNTHESIS_CONTEXT_RESERVE_TOKENS), + fixed_chars, + _MAX_SYNTHESIS_EVIDENCE_CHARS, + ) def _bounded_synthesis_evidence( @@ -1787,30 +1805,44 @@ class ResearchSupervisor: for source in sources ) evidence = "\n\n".join(decision_notes) + decision_system = _system_prompt_with_instructions( + _AGENT_SYSTEM_PROMPT + (f"\n\n{policy_prompt}" if policy_prompt else ""), + run["config"], + ) + # Same whole-prompt budget as synthesis: a fixed 60k evidence tail is many times a + # small loaded context, and this runs on every step, so an overflow here kills the + # run long before it can synthesize what it already gathered. + decision_plan_json = json.dumps(run["plan"], ensure_ascii = False) + decision_scaffold = ( + len(decision_system) + len(question) + len(decision_plan_json) + len(source_catalog) + ) + decision_total = _prompt_char_budget(_SYNTHESIS_CONTEXT_RESERVE_TOKENS) + evidence_chars = _trimmable_budget( + decision_total, decision_scaffold, _MAX_SYNTHESIS_EVIDENCE_CHARS + ) + decision_context = conversation_context[ + : _trimmable_budget( + decision_total, decision_scaffold + evidence_chars, _MAX_CONTEXT_CHARS + ) + ] decision, decision_reasoning, _finish_reason = await self._stream_completion( run, [ { "role": "system", - "content": ( - _system_prompt_with_instructions( - _AGENT_SYSTEM_PROMPT - + (f"\n\n{policy_prompt}" if policy_prompt else ""), - run["config"], - ) - ), + "content": decision_system, }, { "role": "user", "content": ( - f"Conversation context JSON:\n{_shield_untrusted(conversation_context)}\n\n" + f"Conversation context JSON:\n{_shield_untrusted(decision_context)}\n\n" f"Question:\n{_shield_untrusted(question)}\n\n" f"Approved plan (guidance only):\n" - f"{_shield_untrusted(json.dumps(run['plan'], ensure_ascii = False))}\n\n" + f"{_shield_untrusted(decision_plan_json)}\n\n" f"Actions remaining after this one: {max_steps - position - 1}\n" f"\n" f"Gathered sources:\n{_shield_untrusted(source_catalog) or '(none)'}\n\n" - f"{_shield_untrusted(evidence[-60000:]) or '(none)'}\n" + f"{_shield_untrusted(evidence[-evidence_chars:] if evidence_chars else '') or '(none)'}\n" f"" ), }, @@ -2070,16 +2102,32 @@ class ResearchSupervisor: f" Chunk ID: {source.get('chunkId') or '(unknown)'}" for index, source in enumerate(document_sources, 1) ) - evidence_text = _bounded_synthesis_evidence(notes, _synthesis_evidence_budget()) + # Budget the whole prompt, not just the evidence: trim the conversation context first, + # then the evidence, so the untrimmable scaffolding cannot push the request past the + # loaded context and turn a finished run into a failure. + report_system = _system_prompt_with_instructions(_REPORT_SYSTEM_PROMPT, run["config"]) + plan_json = json.dumps(run["plan"], ensure_ascii = False) + scaffold_chars = ( + len(report_system) + + len(question) + + len(plan_json) + + len(source_catalog) + + len(document_source_catalog) + ) + # Evidence is the report, so it is budgeted first and the chat history takes what is left. + total_budget = _prompt_char_budget(_SYNTHESIS_CONTEXT_RESERVE_TOKENS) + evidence_text = _bounded_synthesis_evidence( + notes, _synthesis_evidence_budget(scaffold_chars) + ) + conversation_context = conversation_context[ + : _trimmable_budget(total_budget, scaffold_chars + len(evidence_text), _MAX_CONTEXT_CHARS) + ] report, synthesis_reasoning, synthesis_finish_reason = await self._stream_completion( run, [ { "role": "system", - "content": _system_prompt_with_instructions( - _REPORT_SYSTEM_PROMPT, - run["config"], - ), + "content": report_system, }, { "role": "user", diff --git a/studio/backend/tests/test_research_runs_hardening.py b/studio/backend/tests/test_research_runs_hardening.py index 3af7be1453..d598e34dea 100644 --- a/studio/backend/tests/test_research_runs_hardening.py +++ b/studio/backend/tests/test_research_runs_hardening.py @@ -161,6 +161,23 @@ def test_citation_title_strips_brackets_for_catalog_and_citation(): assert _citation_title({}, "https://x/a") == "https://x/a" +def test_prompt_budget_counts_the_whole_prompt(monkeypatch): + # Budgeting only the evidence cannot prevent an overflow: at a small context the + # untrimmable scaffolding (system prompt, plan, source catalogs) is already several times + # the window, and the old floor added 1500 chars on top of that. + monkeypatch.setattr(research_runs, "_loaded_context_length", lambda: None) + assert research_runs._prompt_char_budget(4096) is None + assert research_runs._trimmable_budget(None, 99_999, 500) == 500 + + monkeypatch.setattr(research_runs, "_loaded_context_length", lambda: 16384) + total = research_runs._prompt_char_budget(4096) + assert total == int((16384 - 4096) * research_runs._SYNTHESIS_EVIDENCE_CHARS_PER_TOKEN) + # A trimmable section never exceeds what is left, and never goes negative. + assert research_runs._trimmable_budget(total, 0, 1_000) == 1_000 + assert research_runs._trimmable_budget(total, total - 10, 1_000) == 10 + assert research_runs._trimmable_budget(total, total + 5_000, 1_000) == 0 + + def _make_payload(**overrides) -> CreateResearchRun: payload = {"threadId": "t1", "userMessageId": "u1", "inferenceRequest": {"model": "m"}} payload.update(overrides) diff --git a/studio/backend/tests/test_research_runs_storage.py b/studio/backend/tests/test_research_runs_storage.py index c87ef586b8..161064c6ae 100644 --- a/studio/backend/tests/test_research_runs_storage.py +++ b/studio/backend/tests/test_research_runs_storage.py @@ -187,10 +187,16 @@ def test_synthesis_evidence_budget_tracks_loaded_context(monkeypatch): monkeypatch.setattr(worker, "_loaded_context_length", lambda: None) assert worker._synthesis_evidence_budget() == worker._MAX_SYNTHESIS_EVIDENCE_CHARS - # A small context (Studio's 2048 default) shrinks the budget so evidence fits. + # A small context (Studio's 2048 default) shrinks the budget so evidence fits. It must reach + # 0 rather than a floor: the old floor handed back 1500 chars even when the context left room + # for none, so the request still overflowed and failed the run after every search had run. monkeypatch.setattr(worker, "_loaded_context_length", lambda: 2048) - small = worker._synthesis_evidence_budget() - assert worker._MIN_SYNTHESIS_EVIDENCE_CHARS <= small < worker._MAX_SYNTHESIS_EVIDENCE_CHARS + assert worker._synthesis_evidence_budget() == 0 + + # The rest of the prompt counts against the same budget, not just the evidence. + monkeypatch.setattr(worker, "_loaded_context_length", lambda: 16384) + roomy = worker._synthesis_evidence_budget() + assert 0 < worker._synthesis_evidence_budget(8_000) < roomy # A large context uses (and clamps to) the full cap. monkeypatch.setattr(worker, "_loaded_context_length", lambda: 32768)