ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Budget the whole research prompt against the loaded context for PR #7219

Browse source 
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.
This commit is contained in:
danielhanchen 2026-07-26 13:29:40 +00:00
commit dc16598a4f
3 changed files with 100 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

100
studio/backend/core/research_runs.py
View file

@ -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"<untrusted_web_evidence>\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"</untrusted_web_evidence>"
),
},
@ -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",

17
studio/backend/tests/test_research_runs_hardening.py
View file

@ -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)

12
studio/backend/tests/test_research_runs_storage.py
View file

@ -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)