diff --git a/README.md b/README.md index 4ae2125..cb6a0c2 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,16 @@ Provides streaming orchestration that delivers LLM responses token-by-token for - Progressive rendering - Interruptible generation +### Provider budget preflight + +When a provider exposes the optional synchronous `request_budget` capability, +the loop checks the fully assembled request before dispatch. An oversized +request gets exactly one smaller, retention-aware context view; required +reminders and request-only injections are replayed without running hooks or +draining pending state again. If that view still cannot fit, the loop raises +locally and makes no SDK request. Providers without the capability retain the +existing request and dispatch behavior. + ## Configuration ```toml diff --git a/amplifier_module_loop_streaming/__init__.py b/amplifier_module_loop_streaming/__init__.py index ad070a2..d90299b 100644 --- a/amplifier_module_loop_streaming/__init__.py +++ b/amplifier_module_loop_streaming/__init__.py @@ -15,7 +15,7 @@ from collections.abc import AsyncIterator from typing import Any, ClassVar -from amplifier_core import HookRegistry, HookResult, ModuleCoordinator, ToolResult +from amplifier_core import ContextLengthError, HookRegistry, HookResult, ModuleCoordinator, ToolResult from amplifier_core.events import ( CANCEL_COMPLETED, CANCEL_REQUESTED, @@ -385,6 +385,70 @@ def _last_real_user_index(msgs: list[dict[str, Any]]) -> int | None: return None +def _replay_request_overlays( + base_messages: list[dict[str, Any]], + *, + turn_start_view_block: str | None, + request_injection: tuple[str, bool] | None, + pending_injections: list[dict[str, Any]], +) -> list[dict[str, Any]]: + """Apply already-decided request-only injections to a fresh context view. + + This is deliberately pure: a budget retry must not re-emit hooks or drain + pending state a second time. Callers snapshot the plan before the normal + assembly consumes it, then replay the same placement against one smaller + context view. + """ + messages = list(base_messages) + if turn_start_view_block is not None: + entry = { + "role": "user", + "content": turn_start_view_block, + "metadata": {"ephemeral": True, "reminder_placement": "pre_user"}, + } + splice_idx = _last_real_user_index(messages) + if splice_idx is None: + messages.append(entry) + else: + messages.insert(splice_idx, entry) + + def append_injection(body: str, append_to_tool_result: bool) -> None: + block = _wrap_reminders(body, tail=True) + if append_to_tool_result and messages and messages[-1].get("role") == "tool": + last = messages[-1] + messages[-1] = { + **last, + "content": f"{last.get('content', '')}\n\n{block}", + } + return + messages.append( + { + "role": "user", + "content": block, + "metadata": {"ephemeral": True, "reminder_placement": "tail"}, + } + ) + + if request_injection is not None: + append_injection(*request_injection) + + pending_to_tool_result = [ + injection["content"] + for injection in pending_injections + if injection.get("content") and injection.get("append_to_last_tool_result") + ] + pending_to_message = [ + injection["content"] + for injection in pending_injections + if injection.get("content") and not injection.get("append_to_last_tool_result") + ] + if pending_to_tool_result: + append_injection("\n\n".join(pending_to_tool_result), True) + if pending_to_message: + append_injection("\n\n".join(pending_to_message), False) + return messages + + class ConversationProviderPin: """Coordinator capability ``conversation.provider_pin`` -- pin, unpin, and read WHICH MOUNTED PROVIDER ANSWERS THE TOP-LEVEL CONVERSATION. @@ -3238,7 +3302,7 @@ async def _execute_stream( """ retaining_getter = None get_capability = getattr(coordinator, "get_capability", None) - if self._ephemeral_injection_mode == "persist" and callable(get_capability): + if callable(get_capability): candidate = get_capability("context.request_retention") if callable(candidate): retaining_getter = candidate @@ -3254,12 +3318,37 @@ async def _execute_stream( ) self._retention_capability_warned = True - async def request_messages(retain_contents: list[str]): - if retaining_getter is not None: - return await retaining_getter( - provider=provider, retain_contents=retain_contents - ) - return await context.get_messages_for_request(provider=provider) + async def request_messages( + retain_contents: list[str], *, token_budget: int | None = None + ): + if retaining_getter is not None and ( + self._ephemeral_injection_mode == "persist" or token_budget is not None + ): + kwargs: dict[str, Any] = { + "provider": provider, + "retain_contents": retain_contents, + } + if token_budget is not None: + kwargs["token_budget"] = token_budget + try: + return await retaining_getter(**kwargs) + except TypeError as exc: + if token_budget is not None: + raise ContextLengthError( + "context.request_retention does not accept token_budget" + ) from exc + raise + kwargs = {"provider": provider} + if token_budget is not None: + kwargs["token_budget"] = token_budget + try: + return await context.get_messages_for_request(**kwargs) + except TypeError as exc: + if token_budget is not None: + raise ContextLengthError( + "context request getter does not accept token_budget" + ) from exc + raise turn_start_retained_contents: list[str] = [] # Emit and process prompt submit (allows hooks to inject context before processing) @@ -3319,6 +3408,7 @@ async def request_messages(retain_contents: list[str]): if prov is provider: provider_name = name break + budget_capable = callable(getattr(provider, "request_budget", None)) # Pure observability. `basis` names WHY this provider won: # "pinned" when the conversation-scope pin decided it (capability @@ -3348,6 +3438,71 @@ async def request_messages(retain_contents: list[str]): }, ) + def build_chat_request( + message_dicts: list[dict[str, Any]], *, tool_choice: str | None = None + ) -> ChatRequest: + tools_list = [_build_tool_spec(tool) for tool in tools.values()] if tools else None + kwargs: dict[str, Any] = { + "messages": [Message(**message) for message in message_dicts], + "tools": tools_list, + "reasoning_effort": self.config.get("reasoning_effort"), + } + if tool_choice is not None: + kwargs["tool_choice"] = tool_choice + return ChatRequest( + **kwargs + ) + + async def check_request_budget( + request: ChatRequest, + base_messages: list[dict[str, Any]], + *, + attempt: int, + ) -> int | None: + """Return one requested smaller context budget, or ``None`` when it fits.""" + request_budget = getattr(provider, "request_budget", None) + if not budget_capable or not callable(request_budget): + return None + context_estimate = sum(len(str(message)) // 4 for message in base_messages) + decision = request_budget(request, context_estimate=context_estimate) + required = ( + "estimated_input_tokens", + "input_limit_tokens", + "context_token_budget", + ) + if not isinstance(decision, dict) or any( + isinstance(decision.get(key), bool) + or not isinstance(decision.get(key), int) + or decision[key] < 0 + for key in required + ): + raise ContextLengthError( + "Provider request_budget returned an invalid budget decision" + ) + + estimated = decision["estimated_input_tokens"] + allowance = decision["input_limit_tokens"] + target = decision["context_token_budget"] + fits = estimated <= allowance + await hooks.emit( + "orchestrator:provider_budget", + { + "attempt": attempt, + "context_estimate": context_estimate, + "estimated_input_tokens": estimated, + "input_limit_tokens": allowance, + "context_token_budget": target, + "result": "fits" if fits else "oversized", + }, + ) + if fits: + return None + if target <= 0: + raise ContextLengthError( + "Provider request exceeds its input budget and cannot retain a smaller context" + ) + return target + # --- Turn-start reminder assembly (reminder-redesign-spec.md, # W1.2, Option D). Hoists iteration 1's provider:request emit to # BEFORE the user prompt is appended below, so the merged @@ -3581,11 +3736,16 @@ async def exit_for_cancellation() -> None: retained_contents = ( list(turn_start_retained_contents) if iteration == 1 else [] ) + # Admit a persistent provider injection before requesting the view. + # That keeps the injected body in this request (and, on a budget + # rebuild, in the one smaller retained view) without a second + # destructive retention read. if ( result.action == "inject_context" and result.ephemeral and result.context_injection and retaining_getter is not None + and self._ephemeral_injection_mode == "persist" ): content, _ = await self._persist_reminder( context, @@ -3598,6 +3758,25 @@ async def exit_for_cancellation() -> None: # Pass provider for dynamic budget calculation based on model's context window message_dicts = await request_messages(retained_contents) message_dicts = list(message_dicts) # Convert to list for modification + base_message_dicts = list(message_dicts) + replay_turn_start_block = ( + self._turn_start_view_block + if self._ephemeral_injection_mode == "tail" and iteration == 1 + else None + ) + replay_request_injection = ( + (result.context_injection, result.append_to_last_tool_result) + if self._ephemeral_injection_mode == "tail" + and result.action == "inject_context" + and result.ephemeral + and result.context_injection + else None + ) + replay_pending_injections = ( + list(self._pending_ephemeral_injections) + if self._ephemeral_injection_mode == "tail" + else [] + ) # Splice the turn-start reminder block into the request VIEW # (reminder-redesign-spec.md, W1.2). Only reachable when @@ -3667,6 +3846,7 @@ async def exit_for_cancellation() -> None: ) if changed: message_dicts = list(await request_messages([])) + base_message_dicts = list(message_dicts) # Check if we should append to last tool result elif result.append_to_last_tool_result and len(message_dicts) > 0: last_msg = message_dicts[-1] @@ -3771,6 +3951,7 @@ async def exit_for_cancellation() -> None: message_dicts = list( await request_messages(retained_contents) ) + base_message_dicts = list(message_dicts) else: if pending_to_tool_result: tool_result_block = _wrap_reminders( @@ -3825,26 +4006,44 @@ async def exit_for_cancellation() -> None: # Clear pending injections after applying (both modes) self._pending_ephemeral_injections = [] - # Convert dicts to ChatRequest for provider - messages_objects = [Message(**msg) for msg in message_dicts] - - # Convert tools to ToolSpec format for ChatRequest - tools_list = None - if tools: - tools_list = [_build_tool_spec(t) for t in tools.values()] - - chat_request = ChatRequest( - messages=messages_objects, - tools=tools_list, - reasoning_effort=self.config.get("reasoning_effort"), - ) + chat_request = build_chat_request(message_dicts) logger.info( - f"[ORCHESTRATOR] ChatRequest created with {len(tools_list) if tools_list else 0} tools" + f"[ORCHESTRATOR] ChatRequest created with {len(tools) if tools else 0} tools" ) - if tools_list: + if tools: logger.debug( - f"[ORCHESTRATOR] Tool names: {[t.name for t in tools_list]}" + f"[ORCHESTRATOR] Tool names: {[t.name for t in tools.values()]}" + ) + + smaller_context_budget = await check_request_budget( + chat_request, base_message_dicts, attempt=0 + ) + if smaller_context_budget is not None: + rebuilt_base_messages = list( + await request_messages( + retained_contents, token_budget=smaller_context_budget + ) + ) + rebuilt_messages = ( + _replay_request_overlays( + rebuilt_base_messages, + turn_start_view_block=replay_turn_start_block, + request_injection=replay_request_injection, + pending_injections=replay_pending_injections, + ) + if self._ephemeral_injection_mode == "tail" + else rebuilt_base_messages ) + rebuilt_request = build_chat_request(rebuilt_messages) + if ( + await check_request_budget( + rebuilt_request, rebuilt_base_messages, attempt=1 + ) + ) is not None: + raise ContextLengthError( + "Provider request remains over budget after one context rebuild" + ) + chat_request = rebuilt_request # Apply rate limit delay before provider call await self._apply_rate_limit_delay(hooks, iteration) @@ -4381,8 +4580,25 @@ async def exit_for_cancellation() -> None: # Current provider-hook requirements and a queued tool-post # injection must survive the bounded finalization assembly too. final_retained_contents: list[str] = [] + final_replay_request_injection = ( + ( + finalization_result.context_injection, + finalization_result.append_to_last_tool_result, + ) + if self._ephemeral_injection_mode == "tail" + and finalization_result.action == "inject_context" + and finalization_result.ephemeral + and finalization_result.context_injection + else None + ) + final_replay_pending_injections = ( + list(self._pending_ephemeral_injections) + if self._ephemeral_injection_mode == "tail" + else [] + ) if ( retaining_getter is not None + and self._ephemeral_injection_mode == "persist" and finalization_result.action == "inject_context" and finalization_result.ephemeral and finalization_result.context_injection @@ -4394,7 +4610,11 @@ async def exit_for_cancellation() -> None: verify_admitted=True, ) final_retained_contents.append(content) - if retaining_getter is not None and self._pending_ephemeral_injections: + if ( + retaining_getter is not None + and self._ephemeral_injection_mode == "persist" + and self._pending_ephemeral_injections + ): pending_body = "\n\n".join( injection["content"] for injection in self._pending_ephemeral_injections @@ -4406,7 +4626,37 @@ async def exit_for_cancellation() -> None: ) final_retained_contents.append(content) self._pending_ephemeral_injections.clear() + if ( + retaining_getter is None + and self._ephemeral_injection_mode == "persist" + ): + if ( + finalization_result.action == "inject_context" + and finalization_result.ephemeral + and finalization_result.context_injection + ): + await self._persist_reminder( + context, finalization_result.context_injection, tail=True + ) + if self._pending_ephemeral_injections: + pending_body = "\n\n".join( + injection["content"] + for injection in self._pending_ephemeral_injections + if injection.get("content") + ) + if pending_body: + await self._persist_reminder(context, pending_body, tail=True) + self._pending_ephemeral_injections.clear() message_dicts = list(await request_messages(final_retained_contents)) + base_message_dicts = list(message_dicts) + if self._ephemeral_injection_mode == "tail": + message_dicts = _replay_request_overlays( + message_dicts, + turn_start_view_block=None, + request_injection=final_replay_request_injection, + pending_injections=final_replay_pending_injections, + ) + self._pending_ephemeral_injections.clear() # The finalization hook and context assembly both await. Check # again before contacting the provider so a concurrent # cancellation cannot buy an unrequested final provider call. @@ -4422,8 +4672,7 @@ async def exit_for_cancellation() -> None: # W5 (amplifier-module-provider-anthropic), an unstamped # trailing message here would be treated as stable and could # take a cache breakpoint on regenerated content. - message_dicts.append( - { + finalization_overlay = { "role": "user", "content": _wrap_reminders( """ @@ -4437,27 +4686,52 @@ async def exit_for_cancellation() -> None: "ephemeral": True, "reminder_placement": "tail", }, - } - ) + } + message_dicts.append(finalization_overlay) try: - # Convert dicts to ChatRequest - messages_objects = [Message(**msg) for msg in message_dicts] - # Preserve the normal declarations, including provider-native # specifications, so any assistant tool call and paired tool # result in the existing transcript stay valid. The portable # choice prevents new calls; this finalization path never # parses or dispatches a tool response. - tools_list = ( - [_build_tool_spec(tool) for tool in tools.values()] if tools else None + max_iter_chat_request = build_chat_request( + message_dicts, tool_choice="none" ) - max_iter_chat_request = ChatRequest( - messages=messages_objects, - tools=tools_list, - tool_choice="none", - reasoning_effort=self.config.get("reasoning_effort"), + smaller_context_budget = await check_request_budget( + max_iter_chat_request, base_message_dicts, attempt=0 ) + if smaller_context_budget is not None: + rebuilt_base_messages = list( + await request_messages( + final_retained_contents, + token_budget=smaller_context_budget, + ) + ) + rebuilt_messages = ( + _replay_request_overlays( + rebuilt_base_messages, + turn_start_view_block=None, + request_injection=final_replay_request_injection, + pending_injections=final_replay_pending_injections, + ) + if self._ephemeral_injection_mode == "tail" + else list(rebuilt_base_messages) + ) + rebuilt_messages.append(finalization_overlay) + rebuilt_request = build_chat_request( + rebuilt_messages, tool_choice="none" + ) + if ( + await check_request_budget( + rebuilt_request, rebuilt_base_messages, attempt=1 + ) + ) is not None: + raise ContextLengthError( + "Provider finalization request remains over budget " + "after one context rebuild" + ) + max_iter_chat_request = rebuilt_request kwargs = {} if self.extended_thinking: @@ -4539,6 +4813,8 @@ async def exit_for_cancellation() -> None: "The previous operation was cancelled. Results from completed tools have been preserved." ) raise + except ContextLengthError: + raise except LLMError as e: await hooks.emit( PROVIDER_ERROR, diff --git a/tests/test_provider_budget_guard.py b/tests/test_provider_budget_guard.py new file mode 100644 index 0000000..7db9703 --- /dev/null +++ b/tests/test_provider_budget_guard.py @@ -0,0 +1,493 @@ +"""Focused coverage for the optional provider request-budget preflight.""" + +from __future__ import annotations + +import pytest +from amplifier_core import ContextLengthError + +from amplifier_module_loop_streaming import ( + StreamingOrchestrator, + _replay_request_overlays, +) +from tests.test_ephemeral_cache_persist_mode import ( + MockContext, + MockCoordinator, + NRoundToolProvider, + OneShotTool, + RequestCapturingProvider, + ScriptedHookResult, + ScriptedHooks, +) + + +def _decision(estimated: int, limit: int, target: int) -> dict[str, int]: + return { + "estimated_input_tokens": estimated, + "input_limit_tokens": limit, + "context_token_budget": target, + } + + +class BudgetProvider(RequestCapturingProvider): + def __init__(self, decisions: list[dict[str, int]]) -> None: + super().__init__() + self.decisions = list(decisions) + self.budget_calls: list[tuple[object, int]] = [] + + def request_budget(self, request, *, context_estimate: int) -> dict[str, int]: + self.budget_calls.append((request, context_estimate)) + return self.decisions.pop(0) + + +class BudgetContext(MockContext): + """Context double that records ordinary and retention-budget requests.""" + + def __init__(self) -> None: + super().__init__() + self.request_calls: list[tuple[list[str], int | None]] = [] + self.legacy_calls: list[int | None] = [] + + async def get_messages(self) -> list[dict]: + return list(self._messages) + + async def get_messages_for_request( + self, provider=None, token_budget: int | None = None + ) -> list[dict]: + self.legacy_calls.append(token_budget) + if token_budget is None: + return list(self._messages) + return [ + message for message in self._messages if message.get("role") != "assistant" + ] + + async def retaining_view( + self, *, provider=None, retain_contents: list[str], token_budget: int | None = None + ) -> list[dict]: + self.request_calls.append((list(retain_contents), token_budget)) + if token_budget is None: + return list(self._messages) + return [ + message + for message in self._messages + if message.get("role") != "assistant" + or message.get("content") in retain_contents + ] + + +def _retaining_coordinator(context: BudgetContext) -> MockCoordinator: + coordinator = MockCoordinator() + coordinator.register_capability("context.request_retention", context.retaining_view) + return coordinator + + +def _injection(body: str) -> ScriptedHookResult: + return ScriptedHookResult( + action="inject_context", ephemeral=True, context_injection=body + ) + + +@pytest.mark.asyncio +async def test_provider_without_budget_capability_keeps_single_normal_dispatch() -> None: + context = MockContext() + provider = RequestCapturingProvider() + + await StreamingOrchestrator({}).execute( + "work", context, {"main": provider}, {}, ScriptedHooks({}), MockCoordinator() + ) + + assert len(provider.requests) == 1 + + +@pytest.mark.asyncio +async def test_provider_without_budget_capability_keeps_retained_injection() -> None: + context = BudgetContext() + provider = RequestCapturingProvider() + body = "LEGACY" + hooks = ScriptedHooks({"provider:request": _injection(body)}) + + await StreamingOrchestrator({}).execute( + "work", + context, + {"main": provider}, + {}, + hooks, + _retaining_coordinator(context), + ) + + assert body in "\n".join(message.content for message in provider.requests[0].messages) + assert len(provider.requests) == 1 + assert len(context.request_calls) == 1 + assert context.request_calls[0][1] is None + assert [name for name, _ in hooks.emitted].count("provider:request") == 1 + + +@pytest.mark.asyncio +async def test_tail_mode_without_budget_capability_stays_view_only() -> None: + context = BudgetContext() + provider = RequestCapturingProvider() + body = "TAIL-LEGACY" + hooks = ScriptedHooks({"provider:request": _injection(body)}) + + await StreamingOrchestrator( + {"ephemeral_injection_mode": "tail", "reminder_placement": "tail"} + ).execute( + "work", + context, + {"main": provider}, + {}, + hooks, + _retaining_coordinator(context), + ) + + assert context.legacy_calls == [None] + assert context.request_calls == [] + assert len(provider.requests) == 1 + assert "\n".join(message.content for message in provider.requests[0].messages).count(body) == 1 + assert [name for name, _ in hooks.emitted].count("provider:request") == 1 + + +@pytest.mark.asyncio +async def test_fitting_budget_dispatches_the_original_request_once() -> None: + context = BudgetContext() + provider = BudgetProvider([_decision(5, 5, 0)]) + + await StreamingOrchestrator({}).execute( + "work", + context, + {"main": provider}, + {}, + ScriptedHooks({}), + _retaining_coordinator(context), + ) + + assert len(provider.requests) == 1 + assert len(provider.budget_calls) == 1 + assert context.request_calls == [([], None)] + + +@pytest.mark.asyncio +async def test_one_smaller_retained_view_is_rechecked_and_dispatches_once() -> None: + context = BudgetContext() + context._messages.append({"role": "assistant", "content": "history" * 200}) + provider = BudgetProvider([_decision(100, 10, 7), _decision(9, 10, 0)]) + retained_body = "REQUIRED" + + await StreamingOrchestrator({}).execute( + "work", + context, + {"main": provider}, + {}, + ScriptedHooks({"provider:request": _injection(retained_body)}), + _retaining_coordinator(context), + ) + + assert len(provider.requests) == 1 + assert [budget for _, budget in context.request_calls] == [None, 7] + assert context.request_calls[0][0] == context.request_calls[1][0] + request_bodies = [message.content for message in provider.requests[0].messages] + assert retained_body in "\n".join(request_bodies) + assert "history" not in "\n".join(request_bodies) + + +class StreamingBudgetProvider(BudgetProvider): + async def stream(self, request, *, tools): + self.requests.append(request) + yield {"content": "streamed"} + + +@pytest.mark.asyncio +async def test_streaming_dispatch_uses_the_same_budget_rebuild() -> None: + context = BudgetContext() + context._messages.append({"role": "assistant", "content": "history" * 200}) + provider = StreamingBudgetProvider( + [_decision(100, 10, 7), _decision(9, 10, 0)] + ) + + await StreamingOrchestrator({}).execute( + "work", + context, + {"main": provider}, + {}, + ScriptedHooks({}), + _retaining_coordinator(context), + ) + + assert len(provider.requests) == 1 + assert len(provider.budget_calls) == 2 + # Ordinary initial view, followed by exactly one explicitly budgeted rebuild. + assert [budget for _, budget in context.request_calls] == [None, 7] + + +@pytest.mark.asyncio +async def test_irreducible_budget_makes_no_sdk_call() -> None: + context = BudgetContext() + provider = BudgetProvider([_decision(100, 10, 0)]) + + with pytest.raises(ContextLengthError, match="cannot retain"): + await StreamingOrchestrator({}).execute( + "work", + context, + {"main": provider}, + {}, + ScriptedHooks({}), + _retaining_coordinator(context), + ) + + assert provider.requests == [] + + +@pytest.mark.asyncio +async def test_second_oversize_after_one_rebuild_makes_no_sdk_call() -> None: + context = BudgetContext() + provider = BudgetProvider([_decision(100, 10, 7), _decision(50, 10, 1)]) + + with pytest.raises(ContextLengthError, match="remains over budget"): + await StreamingOrchestrator({}).execute( + "work", + context, + {"main": provider}, + {}, + ScriptedHooks({}), + _retaining_coordinator(context), + ) + + assert len(provider.budget_calls) == 2 + assert [budget for _, budget in context.request_calls] == [None, 7] + assert context.request_calls[0][0] == context.request_calls[1][0] + assert provider.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "decision", + [ + {"estimated_input_tokens": 1}, + {"estimated_input_tokens": True, "input_limit_tokens": 10, "context_token_budget": 7}, + {"estimated_input_tokens": -1, "input_limit_tokens": 10, "context_token_budget": 7}, + {"estimated_input_tokens": float("nan"), "input_limit_tokens": 10, "context_token_budget": 7}, + {"estimated_input_tokens": "1", "input_limit_tokens": 10, "context_token_budget": 7}, + None, + ], +) +async def test_malformed_budget_result_fails_before_dispatch(decision) -> None: + context = BudgetContext() + provider = BudgetProvider([decision]) + + with pytest.raises(ContextLengthError, match="invalid budget decision"): + await StreamingOrchestrator({}).execute( + "work", + context, + {"main": provider}, + {}, + ScriptedHooks({}), + _retaining_coordinator(context), + ) + + assert provider.requests == [] + + +@pytest.mark.asyncio +async def test_budget_replay_keeps_tail_overlay_once_without_rerunning_hooks() -> None: + context = BudgetContext() + context._messages.append({"role": "assistant", "content": "history" * 200}) + provider = BudgetProvider([_decision(100, 10, 7), _decision(9, 10, 0)]) + body = "ONCE" + hooks = ScriptedHooks({"provider:request": _injection(body)}) + + await StreamingOrchestrator( + {"ephemeral_injection_mode": "tail", "reminder_placement": "tail"} + ).execute( + "work", + context, + {"main": provider}, + {}, + hooks, + _retaining_coordinator(context), + ) + + request_bodies = [message.content for message in provider.requests[0].messages] + assert "\n".join(request_bodies).count(body) == 1 + provider_requests = [name for name, _ in hooks.emitted if name == "provider:request"] + assert provider_requests == ["provider:request"] + assert context.legacy_calls == [None] + assert [budget for _, budget in context.request_calls] == [7] + + +@pytest.mark.asyncio +async def test_budget_replay_keeps_pre_user_tail_overlay_in_its_original_position() -> None: + context = BudgetContext() + context._messages.append({"role": "assistant", "content": "history" * 200}) + provider = BudgetProvider([_decision(100, 10, 7), _decision(9, 10, 0)]) + body = "PRE-USER" + hooks = ScriptedHooks({"provider:request": _injection(body)}) + + await StreamingOrchestrator({"ephemeral_injection_mode": "tail"}).execute( + "work", + context, + {"main": provider}, + {}, + hooks, + _retaining_coordinator(context), + ) + + request_messages = provider.requests[0].messages + body_index = next( + index for index, message in enumerate(request_messages) if body in message.content + ) + assert request_messages[body_index + 1].content == "work" + assert sum(body in message.content for message in request_messages) == 1 + assert [name for name, _ in hooks.emitted].count("provider:request") == 1 + assert context.legacy_calls == [None] + assert [budget for _, budget in context.request_calls] == [7] + + +def test_replayed_pending_overlays_keep_tool_adjacency_and_bodies_once() -> None: + messages = _replay_request_overlays( + [{"role": "tool", "content": "tool output"}], + turn_start_view_block=None, + request_injection=None, + pending_injections=[ + {"content": "TOOL-REMINDER", "append_to_last_tool_result": True}, + {"content": "MESSAGE-REMINDER", "append_to_last_tool_result": False}, + ], + ) + + assert messages[0]["role"] == "tool" + assert messages[0]["content"].count("TOOL-REMINDER") == 1 + assert messages[1]["role"] == "user" + assert messages[1]["content"].count("MESSAGE-REMINDER") == 1 + + +class BudgetToolProvider(NRoundToolProvider): + def __init__(self, decisions: list[dict[str, int]]) -> None: + super().__init__(n_tool_rounds=1) + self.decisions = list(decisions) + self.budget_calls: list[object] = [] + + def request_budget(self, request, *, context_estimate: int) -> dict[str, int]: + self.budget_calls.append(request) + return self.decisions.pop(0) + + +@pytest.mark.asyncio +async def test_pending_tool_overlay_is_replayed_once_after_budget_rebuild() -> None: + context = BudgetContext() + provider = BudgetToolProvider( + [_decision(1, 10, 0), _decision(100, 10, 7), _decision(1, 10, 0)] + ) + pending = ScriptedHookResult( + action="inject_context", + ephemeral=True, + context_injection="PENDING-TOOL", + append_to_last_tool_result=True, + ) + + loop = StreamingOrchestrator( + {"ephemeral_injection_mode": "tail", "reminder_placement": "tail"} + ) + await loop.execute( + "work", + context, + {"main": provider}, + {"mock_tool": OneShotTool()}, + ScriptedHooks({"tool:post": pending}), + _retaining_coordinator(context), + ) + + second_request = provider.requests[1] + tool_messages = [message for message in second_request.messages if message.role == "tool"] + assert len(tool_messages) == 1 + assert tool_messages[0].content.count("PENDING-TOOL") == 1 + assert loop._pending_ephemeral_injections == [] + + +class FinalizingBudgetProvider(NRoundToolProvider): + def __init__(self, decisions: list[dict[str, int]] | None = None) -> None: + super().__init__(n_tool_rounds=1) + self.budget_calls: list[object] = [] + self.decisions = decisions or [_decision(1, 10, 0), _decision(1, 10, 0)] + + def request_budget(self, request, *, context_estimate: int) -> dict[str, int]: + self.budget_calls.append(request) + return self.decisions.pop(0) + + +@pytest.mark.asyncio +async def test_finalization_request_is_budget_checked_before_dispatch() -> None: + context = BudgetContext() + provider = FinalizingBudgetProvider() + + await StreamingOrchestrator({"max_iterations": 1}).execute( + "work", + context, + {"main": provider}, + {"mock_tool": OneShotTool()}, + ScriptedHooks({}), + _retaining_coordinator(context), + ) + + assert len(provider.requests) == 2 + assert len(provider.budget_calls) == 2 + assert provider.requests[-1].tool_choice == "none" + + +@pytest.mark.asyncio +async def test_finalization_irreducible_budget_skips_its_sdk_dispatch() -> None: + context = BudgetContext() + provider = FinalizingBudgetProvider( + [_decision(1, 10, 0), _decision(100, 10, 0)] + ) + + with pytest.raises(ContextLengthError, match="cannot retain"): + await StreamingOrchestrator({"max_iterations": 1}).execute( + "work", + context, + {"main": provider}, + {"mock_tool": OneShotTool()}, + ScriptedHooks({}), + _retaining_coordinator(context), + ) + + assert len(provider.requests) == 1 + assert len(provider.budget_calls) == 2 + + +@pytest.mark.asyncio +async def test_finalization_replays_current_and_pending_tail_overlays_once() -> None: + context = BudgetContext() + provider = FinalizingBudgetProvider( + [_decision(1, 10, 0), _decision(100, 10, 7), _decision(1, 10, 0)] + ) + direct = ScriptedHookResult( + action="inject_context", + ephemeral=True, + context_injection="FINAL-DIRECT", + append_to_last_tool_result=True, + ) + pending = ScriptedHookResult( + action="inject_context", + ephemeral=True, + context_injection="FINAL-PENDING", + append_to_last_tool_result=True, + ) + loop = StreamingOrchestrator( + { + "max_iterations": 1, + "ephemeral_injection_mode": "tail", + "reminder_placement": "tail", + } + ) + + await loop.execute( + "work", + context, + {"main": provider}, + {"mock_tool": OneShotTool()}, + ScriptedHooks({"provider:request": direct, "tool:post": pending}), + _retaining_coordinator(context), + ) + + final_bodies = "\n".join(message.content for message in provider.requests[-1].messages) + assert final_bodies.count("FINAL-DIRECT") == 1 + assert final_bodies.count("FINAL-PENDING") == 1 + assert loop._pending_ephemeral_injections == [] \ No newline at end of file