From 53145ee76cb6e744548bac2defc36e1d7592c403 Mon Sep 17 00:00:00 2001 From: Slava Trofimov Date: Thu, 23 Jul 2026 18:41:08 -0400 Subject: [PATCH] auto-expand worker context budget --- src/octopal/runtime/workers/agent_worker.py | 56 +++++++++++++++++++-- tests/test_agent_worker_contracts.py | 35 ++++++++++++- 2 files changed, 86 insertions(+), 5 deletions(-) diff --git a/src/octopal/runtime/workers/agent_worker.py b/src/octopal/runtime/workers/agent_worker.py index e33a057c..a788292e 100644 --- a/src/octopal/runtime/workers/agent_worker.py +++ b/src/octopal/runtime/workers/agent_worker.py @@ -27,7 +27,13 @@ from octopal.infrastructure.config.settings import load_settings from octopal.infrastructure.providers.base import InferenceProvider from octopal.infrastructure.providers.factory import build_inference_provider -from octopal.runtime.context_compiler import ContextSection, compile_context +from octopal.runtime.context_compiler import ( + CompiledContext, + ContextBudgetExceededError, + ContextSection, + compile_context, + estimate_tokens, +) from octopal.runtime.memory.episodes import worker_task_fingerprint from octopal.runtime.temporal_context import format_temporal_context_prompt from octopal.runtime.tool_errors import ToolBridgeError @@ -57,6 +63,8 @@ _MAX_EMPTY_TURNS = 3 _MAX_MALFORMED_RESULT_TURNS = 2 _WORKER_SYSTEM_CONTEXT_TOKEN_BUDGET = 6000 +_WORKER_SYSTEM_CONTEXT_OPTIONAL_RESERVE = 2000 +_WORKER_SYSTEM_CONTEXT_HARD_LIMIT = 64_000 _TOOL_RESULT_READ_MAX_CHARS = 24_000 _TOOL_RESULT_SEARCH_MAX_MATCHES = 20 _TOOL_RESULT_ACCESS_TOOL_NAMES = frozenset({"tool_result_read", "tool_result_search"}) @@ -108,6 +116,47 @@ "bad gateway", "gateway timeout", ) + + +def _compile_worker_system_context(sections: list[ContextSection]) -> CompiledContext: + """Compile a worker prompt without making prompt growth an agent concern. + + The normal budget is a soft target for required worker instructions plus + optional runtime context. Required sections are never truncated. When they + outgrow that target, the runtime expands the effective budget automatically + and retains a small reserve for the highest-priority optional sections. + """ + + required_tokens = sum( + estimate_tokens(section.content) + for section in sections + if section.required and section.content + ) + if required_tokens > _WORKER_SYSTEM_CONTEXT_HARD_LIMIT: + raise ContextBudgetExceededError( + "required worker prompt context exceeds runtime safety limit " + f"({required_tokens} > {_WORKER_SYSTEM_CONTEXT_HARD_LIMIT})" + ) + + effective_budget = min( + _WORKER_SYSTEM_CONTEXT_HARD_LIMIT, + max( + _WORKER_SYSTEM_CONTEXT_TOKEN_BUDGET, + required_tokens + _WORKER_SYSTEM_CONTEXT_OPTIONAL_RESERVE, + ), + ) + compiled = compile_context(sections, token_budget=effective_budget) + return replace( + compiled, + manifest={ + **compiled.manifest, + "soft_token_budget": _WORKER_SYSTEM_CONTEXT_TOKEN_BUDGET, + "hard_token_limit": _WORKER_SYSTEM_CONTEXT_HARD_LIMIT, + "auto_expanded": effective_budget > _WORKER_SYSTEM_CONTEXT_TOKEN_BUDGET, + }, + ) + + _SYSTEMIC_TOOL_ERROR_CLASSIFICATIONS = {"schema_mismatch"} _RESULT_SCHEMA = { "type": "object", @@ -1325,7 +1374,7 @@ async def mcp_proxy_handler(args: dict, ctx: dict, s_id=s_id, t_name=t_name): spec.inputs, idempotency_key=spec.idempotency_key, ) - compiled_system_context = compile_context( + compiled_system_context = _compile_worker_system_context( [ ContextSection("worker_base", worker_base_prompt, required=True), ContextSection("template_role", f"Template role:\n{spec.system_prompt}", required=True), @@ -1335,8 +1384,7 @@ async def mcp_proxy_handler(args: dict, ctx: dict, s_id=s_id, t_name=t_name): ContextSection("tool_inventory", f"Available tools:\n{tool_inventory}", priority=100), ContextSection("guidance", guidance_prompt, required=True), ContextSection("completion_protocol", completion_protocol_prompt, required=True), - ], - token_budget=_WORKER_SYSTEM_CONTEXT_TOKEN_BUDGET, + ] ) system_prompt = compiled_system_context.content context_manifest = _build_worker_context_manifest( diff --git a/tests/test_agent_worker_contracts.py b/tests/test_agent_worker_contracts.py index eb7728d4..6c3d6b98 100644 --- a/tests/test_agent_worker_contracts.py +++ b/tests/test_agent_worker_contracts.py @@ -4,12 +4,18 @@ import json import re +import pytest + from octopal.infrastructure.config.models import LLMConfig from octopal.infrastructure.store.models import ( ProceduralRecipeContext, procedural_recipe_definition_fingerprint, ) -from octopal.runtime.context_compiler import ContextSection, compile_context +from octopal.runtime.context_compiler import ( + ContextBudgetExceededError, + ContextSection, + compile_context, +) from octopal.runtime.tool_result_store import ToolResultStore from octopal.runtime.workers.agent_worker import ( _build_procedural_recipe_prompt, @@ -19,6 +25,7 @@ _build_worker_skill_usage_prompt, _build_worker_task_prompt, _build_worker_tool_inventory_prompt, + _compile_worker_system_context, _force_tool_choice, _make_request_instruction_tool, _record_worker_llm_context_snapshot, @@ -400,6 +407,32 @@ def test_worker_context_manifest_records_compiler_decisions_without_content() -> assert "private memory value" not in str(manifest) +def test_worker_context_compiler_auto_expands_for_required_prompt_growth() -> None: + required_context = "r" * 6461 + + compiled = _compile_worker_system_context( + [ + ContextSection("required", required_context, required=True), + ContextSection("optional", "useful optional context", priority=100), + ] + ) + + assert compiled.sections["required"] == required_context + assert compiled.sections["optional"] == "useful optional context" + assert compiled.manifest["required_tokens"] == 6461 + assert compiled.manifest["token_budget"] == 8461 + assert compiled.manifest["soft_token_budget"] == 6000 + assert compiled.manifest["auto_expanded"] is True + + +def test_worker_context_compiler_keeps_a_runtime_safety_limit() -> None: + with pytest.raises( + ContextBudgetExceededError, + match=r"runtime safety limit \(64001 > 64000\)", + ): + _compile_worker_system_context([ContextSection("required", "r" * 64_001, required=True)]) + + def test_worker_recipe_context_is_bounded_advisory_and_manifest_is_content_free() -> None: definition = { "applicability_conditions": ["The fixture is local."],