From e09d1119a024ed1fb5bcf4bd5fff329fe54ba844 Mon Sep 17 00:00:00 2001
From: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com>
Date: Mon, 14 Sep 2026 22:43:38 -0700
Subject: [PATCH 1/3] fix: forward provider hard-fit budget requests
Generated with Amplifier
Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com>
---
README.md | 10 +
amplifier_module_loop_streaming/__init__.py | 49 ++-
..._provider_budget_context_simple_runtime.py | 254 ++++++++++++++++
tests/test_provider_budget_guard.py | 285 +++++++++++++++++-
tests/test_request_retention.py | 4 +-
5 files changed, 582 insertions(+), 20 deletions(-)
create mode 100644 tests/test_provider_budget_context_simple_runtime.py
diff --git a/README.md b/README.md
index cb6a0c2..d36eac7 100644
--- a/README.md
+++ b/README.md
@@ -46,6 +46,16 @@ 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.
+When `context.request_retention` advertises its optional `hard_fit` keyword,
+that one provider-forced rebuild forwards `hard_fit=True`, allowing the context
+to target the provider's requested budget directly. Older retention
+capabilities, uninspectable dynamic callables, and the generic context fallback
+keep their existing `provider`/`retain_contents`/`token_budget` assembly; the
+second preflight and provider's final payload guard remain the safety boundary.
+The existing `orchestrator:provider_budget` event exposes each preflight's
+attempt, result, estimate, allowance, and requested context budget to mounted
+observability consumers.
+
## Configuration
```toml
diff --git a/amplifier_module_loop_streaming/__init__.py b/amplifier_module_loop_streaming/__init__.py
index d90299b..80a04a4 100644
--- a/amplifier_module_loop_streaming/__init__.py
+++ b/amplifier_module_loop_streaming/__init__.py
@@ -8,6 +8,7 @@
import asyncio
import fnmatch
+import inspect
import json
import logging
import re
@@ -746,6 +747,7 @@ async def mount(coordinator: ModuleCoordinator, config: dict[str, Any] | None =
"orchestrator:steering_injected", # When a steer message is injected mid-turn
"orchestrator:goal_progress", # /goal auto-continue loop progress (see docs/designs/goal-command.md)
"orchestrator:budget_warning", # Layer 1 call budget at budget_warn_ratio (see _execute_stream)
+ "orchestrator:provider_budget", # Provider request-budget preflight result (see _execute_stream)
],
)
@@ -3318,26 +3320,44 @@ async def _execute_stream(
)
self._retention_capability_warned = True
+ def retention_accepts_hard_fit() -> bool:
+ """Whether the optional retention capability supports ``hard_fit``.
+
+ The capability is module-owned and evolves independently of Core.
+ Signature inspection is deliberately separate from invocation:
+ an implementation ``TypeError`` must reach the caller rather than
+ being mistaken for an old capability and invoked a second time.
+ """
+ if retaining_getter is None:
+ return False
+ try:
+ parameters = inspect.signature(retaining_getter).parameters.values()
+ except (TypeError, ValueError):
+ # Some dynamic callables cannot expose a signature. Keep their
+ # legacy behavior rather than claiming hard-fit support.
+ return False
+ return any(
+ parameter.name == "hard_fit"
+ or parameter.kind is inspect.Parameter.VAR_KEYWORD
+ for parameter in parameters
+ )
+
async def request_messages(
- retain_contents: list[str], *, token_budget: int | None = None
+ retain_contents: list[str],
+ *,
+ token_budget: int | None = None,
+ hard_fit: bool = False,
):
- if retaining_getter is not None and (
- self._ephemeral_injection_mode == "persist" or token_budget is not None
- ):
+ if retaining_getter 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
+ if hard_fit and retention_accepts_hard_fit():
+ kwargs["hard_fit"] = True
+ return await retaining_getter(**kwargs)
kwargs = {"provider": provider}
if token_budget is not None:
kwargs["token_budget"] = token_budget
@@ -4021,7 +4041,9 @@ async def exit_for_cancellation() -> None:
if smaller_context_budget is not None:
rebuilt_base_messages = list(
await request_messages(
- retained_contents, token_budget=smaller_context_budget
+ retained_contents,
+ token_budget=smaller_context_budget,
+ hard_fit=True,
)
)
rebuilt_messages = (
@@ -4706,6 +4728,7 @@ async def exit_for_cancellation() -> None:
await request_messages(
final_retained_contents,
token_budget=smaller_context_budget,
+ hard_fit=True,
)
)
rebuilt_messages = (
diff --git a/tests/test_provider_budget_context_simple_runtime.py b/tests/test_provider_budget_context_simple_runtime.py
new file mode 100644
index 0000000..230d2c2
--- /dev/null
+++ b/tests/test_provider_budget_context_simple_runtime.py
@@ -0,0 +1,254 @@
+"""Joint request-budget regression with real loop, context-simple, and OpenAI code.
+
+This intentionally has no local path manipulation: repository-only runs skip when
+the sibling modules are not installed. The root DTU installs all three local
+checkouts, where this is required to execute without a skip.
+"""
+
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace
+
+import pytest
+from amplifier_core import ContextLengthError
+from amplifier_core.message_models import ChatRequest, Message
+
+pytest.importorskip("amplifier_module_context_simple")
+pytest.importorskip("amplifier_module_provider_openai")
+
+from amplifier_module_context_simple import SimpleContextManager
+from amplifier_module_loop_streaming import StreamingOrchestrator
+from amplifier_module_provider_openai import OpenAIProvider
+
+
+class _Cancellation:
+ is_cancelled = False
+ is_immediate = False
+ state = "running"
+
+
+class _StableReminderHooks:
+ """A real provider sees a stable persisted reminder on every assembled view."""
+
+ def __init__(self) -> None:
+ self.events: list[tuple[str, dict]] = []
+
+ async def emit(self, event: str, payload: dict | None = None):
+ self.events.append((event, payload or {}))
+ if event == "provider:request":
+ return SimpleNamespace(
+ action="inject_context",
+ ephemeral=True,
+ context_injection="REQUIRED-REMINDER",
+ context_injection_role="user",
+ append_to_last_tool_result=False,
+ data=None,
+ reason=None,
+ )
+ return SimpleNamespace(
+ action="continue",
+ ephemeral=False,
+ context_injection=None,
+ context_injection_role="system",
+ append_to_last_tool_result=False,
+ data=None,
+ reason=None,
+ )
+
+
+class _Coordinator:
+ def __init__(self, hooks: _StableReminderHooks) -> None:
+ self.hooks = hooks
+ self.cancellation = _Cancellation()
+ self.session_state: dict = {}
+ self._capabilities: dict[str, object] = {}
+
+ def register_capability(self, name: str, capability: object) -> None:
+ self._capabilities[name] = capability
+
+ def get_capability(self, name: str):
+ return self._capabilities.get(name)
+
+ async def process_hook_result(self, result, *_args):
+ return result
+
+
+class _RecordingContext(SimpleContextManager):
+ """The actual context-simple algorithm, with only seam-call observation added."""
+
+ def __init__(self, **kwargs) -> None:
+ super().__init__(**kwargs)
+ self.hard_fit_calls: list[bool] = []
+
+ async def get_messages_for_request_retaining(
+ self,
+ *,
+ retain_contents: list[str],
+ provider=None,
+ token_budget: int | None = None,
+ hard_fit: bool = False,
+ ) -> list[dict]:
+ self.hard_fit_calls.append(hard_fit)
+ return await super().get_messages_for_request_retaining(
+ retain_contents=retain_contents,
+ provider=provider,
+ token_budget=token_budget,
+ hard_fit=hard_fit,
+ )
+
+
+class _InMemoryResponses:
+ """Completed Responses SDK fake whose usage is derived from received params."""
+
+ def __init__(self) -> None:
+ self.calls: list[dict] = []
+
+ async def create(self, **params):
+ serialized = json.dumps(
+ params, ensure_ascii=False, sort_keys=True, separators=(",", ":"), default=str
+ )
+ # Record exactly the SDK payload, and derive usage from that same payload
+ # rather than a fixture outcome or a desired compaction size.
+ self.calls.append(json.loads(serialized))
+ input_tokens = len(serialized.encode("utf-8"))
+ return SimpleNamespace(
+ id=f"fake-{len(self.calls)}",
+ status="completed",
+ model=params["model"],
+ output=[
+ {
+ "type": "message",
+ "content": [{"type": "output_text", "text": "accepted"}],
+ }
+ ],
+ usage=SimpleNamespace(input_tokens=input_tokens, output_tokens=1),
+ )
+
+
+class _InMemoryClient:
+ def __init__(self) -> None:
+ self.responses = _InMemoryResponses()
+
+
+def _payload_text(params: dict) -> str:
+ return json.dumps(params, ensure_ascii=False, sort_keys=True)
+
+
+@pytest.mark.asyncio
+async def test_hard_fit_stays_compacted_across_real_openai_dispatches() -> None:
+ hooks = _StableReminderHooks()
+ coordinator = _Coordinator(hooks)
+ context = _RecordingContext(
+ # Initial ordinary assembly is deliberately below this context's own
+ # threshold. The provider's serialized-payload preflight forces the
+ # first fit, exercising the optional hard-fit seam.
+ max_tokens=500_000,
+ compact_threshold=0.99,
+ target_usage=0.50,
+ protected_recent=0.10,
+ protected_tool_results=1,
+ truncate_chars=64,
+ compaction_notice_enabled=True,
+ )
+ coordinator.register_capability(
+ "context.request_retention", context.get_messages_for_request_retaining
+ )
+ client = _InMemoryClient()
+ provider = OpenAIProvider(
+ api_key="test-key",
+ client=client,
+ coordinator=coordinator,
+ config={
+ "default_model": "gpt-5-mini",
+ "max_output_tokens": 1024,
+ "max_retries": 0,
+ "use_streaming": False,
+ },
+ )
+ loop = StreamingOrchestrator({})
+
+ bulk = "REMOVED-BULK-MARKER:" + ("x" * 800_000)
+ await context.add_message({"role": "assistant", "content": bulk})
+
+ await loop.execute(
+ "ORIGINAL-HUMAN", context, {"openai": provider}, {}, hooks, coordinator
+ )
+ # New growth after the forced rebuild: a complete assistant/tool pair and
+ # a current tool result that context-simple's protection floor must keep
+ # complete.
+ await context.add_message(
+ {
+ "role": "assistant",
+ "content": "Calling the current tool.",
+ "tool_calls": [
+ {"id": "runtime-tool-1", "name": "current_tool", "arguments": {}}
+ ],
+ }
+ )
+ await context.add_message(
+ {
+ "role": "tool",
+ "name": "current_tool",
+ "tool_call_id": "runtime-tool-1",
+ "content": "CURRENT-PROTECTED-TOOL-RESULT",
+ }
+ )
+ for prompt in (
+ "SECOND-HUMAN",
+ "POST-FORCE-ORDINARY-ONE",
+ "POST-FORCE-ORDINARY-TWO",
+ "POST-FORCE-ORDINARY-THREE",
+ ):
+ await loop.execute(
+ prompt, context, {"openai": provider}, {}, hooks, coordinator
+ )
+
+ # Every call in this list reached the fake SDK. The first accepted call
+ # followed the provider-directed rebuild; the four subsequent calls prove
+ # ordinary fetches do not resurrect the canonical bulk history.
+ assert len(client.responses.calls) >= 5
+ assert context.hard_fit_calls.count(True) == 1
+ post_force_payloads = [
+ _payload_text(params) for params in client.responses.calls[1:]
+ ]
+ assert post_force_payloads
+ assert all("REMOVED-BULK-MARKER" not in payload for payload in post_force_payloads)
+ assert all("ORIGINAL-HUMAN" in payload for payload in post_force_payloads)
+ assert "SECOND-HUMAN" in post_force_payloads[-1]
+ assert "REQUIRED-REMINDER" in post_force_payloads[-1]
+ assert "CURRENT-PROTECTED-TOOL-RESULT" in post_force_payloads[-1]
+ assert 'source=\\"context-compaction\\"' in post_force_payloads[-1]
+
+ final_input = client.responses.calls[-1]["input"]
+ tool_call_index = next(
+ index
+ for index, item in enumerate(final_input)
+ if item.get("type") == "function_call"
+ and item.get("call_id") == "runtime-tool-1"
+ )
+ assert final_input[tool_call_index + 1] == {
+ "type": "function_call_output",
+ "call_id": "runtime-tool-1",
+ "output": "CURRENT-PROTECTED-TOOL-RESULT",
+ }
+
+ canonical = await context.get_messages()
+ assert any(message.get("content") == bulk for message in canonical)
+ assert any(message.get("content") == "ORIGINAL-HUMAN" for message in canonical)
+ assert any(message.get("content") == "SECOND-HUMAN" for message in canonical)
+ assert any(
+ message.get("content") == "CURRENT-PROTECTED-TOOL-RESULT"
+ for message in canonical
+ )
+
+ # OpenAI's direct final assembled-payload guard remains the final boundary:
+ # an impossible protected payload never reaches the SDK fake.
+ accepted_before_guard = len(client.responses.calls)
+ with pytest.raises(ContextLengthError):
+ await provider.complete(
+ ChatRequest(
+ messages=[Message(role="user", content="IMPOSSIBLE" * 200_000)]
+ )
+ )
+ assert len(client.responses.calls) == accepted_before_guard
diff --git a/tests/test_provider_budget_guard.py b/tests/test_provider_budget_guard.py
index 7db9703..4e6fa84 100644
--- a/tests/test_provider_budget_guard.py
+++ b/tests/test_provider_budget_guard.py
@@ -8,6 +8,8 @@
from amplifier_module_loop_streaming import (
StreamingOrchestrator,
_replay_request_overlays,
+ _wrap_reminders,
+ mount,
)
from tests.test_ephemeral_cache_persist_mode import (
MockContext,
@@ -74,6 +76,63 @@ async def retaining_view(
]
+class HardFitBudgetContext(BudgetContext):
+ """Modern retention capability that records the optional hard-fit signal."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.hard_fit_calls: list[bool] = []
+
+ async def retaining_view(
+ self,
+ *,
+ provider=None,
+ retain_contents: list[str],
+ token_budget: int | None = None,
+ hard_fit: bool = False,
+ ) -> list[dict]:
+ self.hard_fit_calls.append(hard_fit)
+ return await super().retaining_view(
+ provider=provider,
+ retain_contents=retain_contents,
+ token_budget=token_budget,
+ )
+
+
+class KwargsBudgetContext(BudgetContext):
+ """Modern retention capability accepting future keywords through ``**kwargs``."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.hard_fit_calls: list[bool] = []
+
+ async def retaining_view(self, **kwargs) -> list[dict]:
+ self.hard_fit_calls.append(kwargs.get("hard_fit", False))
+ return await super().retaining_view(
+ provider=kwargs["provider"],
+ retain_contents=kwargs["retain_contents"],
+ token_budget=kwargs.get("token_budget"),
+ )
+
+
+class OldSignatureBudgetContext(BudgetContext):
+ """Pre-hard-fit retention capability; its call shape is the compatibility check."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.legacy_retention_calls: list[tuple[object, list[str], int | None]] = []
+
+ async def retaining_view(
+ self, *, provider=None, retain_contents: list[str], token_budget: int | None = None
+ ) -> list[dict]:
+ self.legacy_retention_calls.append((provider, list(retain_contents), token_budget))
+ return await super().retaining_view(
+ provider=provider,
+ retain_contents=retain_contents,
+ token_budget=token_budget,
+ )
+
+
def _retaining_coordinator(context: BudgetContext) -> MockCoordinator:
coordinator = MockCoordinator()
coordinator.register_capability("context.request_retention", context.retaining_view)
@@ -139,8 +198,8 @@ async def test_tail_mode_without_budget_capability_stays_view_only() -> None:
_retaining_coordinator(context),
)
- assert context.legacy_calls == [None]
- assert context.request_calls == []
+ assert context.legacy_calls == []
+ assert context.request_calls == [([], None)]
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
@@ -189,6 +248,136 @@ async def test_one_smaller_retained_view_is_rechecked_and_dispatches_once() -> N
assert "history" not in "\n".join(request_bodies)
+@pytest.mark.asyncio
+async def test_forced_normal_rebuild_forwards_hard_fit_only_to_modern_retention() -> None:
+ context = HardFitBudgetContext()
+ context._messages.append({"role": "assistant", "content": "history" * 200})
+ provider = BudgetProvider([_decision(100, 10, 7), _decision(9, 10, 0)])
+
+ await StreamingOrchestrator({}).execute(
+ "work",
+ context,
+ {"main": provider},
+ {},
+ ScriptedHooks({}),
+ _retaining_coordinator(context),
+ )
+
+ # The first ordinary request remains legacy/default behavior; exactly the
+ # forced provider-directed rebuild opts into hard fitting.
+ assert context.hard_fit_calls == [False, True]
+
+
+@pytest.mark.asyncio
+async def test_forced_rebuild_forwards_hard_fit_to_kwargs_retention() -> None:
+ context = KwargsBudgetContext()
+ context._messages.append({"role": "assistant", "content": "history" * 200})
+ provider = BudgetProvider([_decision(100, 10, 7), _decision(9, 10, 0)])
+
+ await StreamingOrchestrator({}).execute(
+ "work",
+ context,
+ {"main": provider},
+ {},
+ ScriptedHooks({}),
+ _retaining_coordinator(context),
+ )
+
+ assert context.hard_fit_calls == [False, True]
+
+
+@pytest.mark.asyncio
+async def test_forced_rebuild_preserves_old_retention_call_signature() -> None:
+ context = OldSignatureBudgetContext()
+ context._messages.append({"role": "assistant", "content": "history" * 200})
+ provider = BudgetProvider([_decision(100, 10, 7), _decision(9, 10, 0)])
+
+ await StreamingOrchestrator({}).execute(
+ "work",
+ context,
+ {"main": provider},
+ {},
+ ScriptedHooks({}),
+ _retaining_coordinator(context),
+ )
+
+ assert context.legacy_retention_calls == [
+ (provider, [], None),
+ (provider, [], 7),
+ ]
+
+
+@pytest.mark.asyncio
+async def test_forced_rebuild_uses_generic_legacy_context_when_retention_is_absent() -> None:
+ context = BudgetContext()
+ context._messages.append({"role": "assistant", "content": "history" * 200})
+ provider = BudgetProvider([_decision(100, 10, 7), _decision(9, 10, 0)])
+
+ await StreamingOrchestrator({}).execute(
+ "work", context, {"main": provider}, {}, ScriptedHooks({}), MockCoordinator()
+ )
+
+ assert context.request_calls == []
+ assert context.legacy_calls == [None, 7]
+
+
+@pytest.mark.asyncio
+async def test_signature_inspection_failure_keeps_safe_legacy_retention_call(monkeypatch) -> None:
+ context = HardFitBudgetContext()
+ context._messages.append({"role": "assistant", "content": "history" * 200})
+ provider = BudgetProvider([_decision(100, 10, 7), _decision(9, 10, 0)])
+
+ def unavailable_signature(_callable):
+ raise ValueError("signature unavailable")
+
+ monkeypatch.setattr(
+ "amplifier_module_loop_streaming.inspect.signature", unavailable_signature
+ )
+
+ await StreamingOrchestrator({}).execute(
+ "work",
+ context,
+ {"main": provider},
+ {},
+ ScriptedHooks({}),
+ _retaining_coordinator(context),
+ )
+
+ assert context.hard_fit_calls == [False, False]
+
+
+@pytest.mark.asyncio
+async def test_retention_type_error_is_not_mistaken_for_a_signature_mismatch() -> None:
+ class ExplodingHardFitContext(HardFitBudgetContext):
+ async def retaining_view(self, **kwargs) -> list[dict]:
+ self.hard_fit_calls.append(kwargs.get("hard_fit", False))
+ if kwargs.get("hard_fit"):
+ raise TypeError("retention implementation exploded")
+ return await BudgetContext.retaining_view(
+ self,
+ provider=kwargs["provider"],
+ retain_contents=kwargs["retain_contents"],
+ token_budget=kwargs.get("token_budget"),
+ )
+
+ context = ExplodingHardFitContext()
+ context._messages.append({"role": "assistant", "content": "history" * 200})
+ provider = BudgetProvider([_decision(100, 10, 7)])
+
+ with pytest.raises(TypeError, match="retention implementation exploded"):
+ await StreamingOrchestrator({}).execute(
+ "work",
+ context,
+ {"main": provider},
+ {},
+ ScriptedHooks({}),
+ _retaining_coordinator(context),
+ )
+
+ assert context.hard_fit_calls == [False, True]
+ assert provider.requests == []
+
+
class StreamingBudgetProvider(BudgetProvider):
async def stream(self, request, *, tools):
self.requests.append(request)
@@ -309,8 +498,8 @@ async def test_budget_replay_keeps_tail_overlay_once_without_rerunning_hooks() -
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]
+ assert context.legacy_calls == []
+ assert [budget for _, budget in context.request_calls] == [None, 7]
@pytest.mark.asyncio
@@ -337,8 +526,8 @@ async def test_budget_replay_keeps_pre_user_tail_overlay_in_its_original_positio
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]
+ assert context.legacy_calls == []
+ assert [budget for _, budget in context.request_calls] == [None, 7]
def test_replayed_pending_overlays_keep_tool_adjacency_and_bodies_once() -> None:
@@ -431,6 +620,90 @@ async def test_finalization_request_is_budget_checked_before_dispatch() -> None:
assert provider.requests[-1].tool_choice == "none"
+@pytest.mark.asyncio
+async def test_forced_finalization_rebuild_forwards_hard_fit_only_at_rebuild() -> None:
+ context = HardFitBudgetContext()
+ provider = FinalizingBudgetProvider(
+ [_decision(1, 10, 0), _decision(100, 10, 7), _decision(1, 10, 0)]
+ )
+
+ await StreamingOrchestrator({"max_iterations": 1}).execute(
+ "work",
+ context,
+ {"main": provider},
+ {"mock_tool": OneShotTool()},
+ ScriptedHooks({}),
+ _retaining_coordinator(context),
+ )
+
+ assert context.hard_fit_calls == [False, False, True]
+
+
+class AnthropicStyleAssemblyProvider(RequestCapturingProvider):
+ """Non-budget control: preserve ordinary assembled requests for other providers."""
+
+ name = "anthropic"
+
+
+@pytest.mark.asyncio
+async def test_non_budget_anthropic_style_provider_keeps_ordinary_request_assembly() -> None:
+ context = BudgetContext()
+ provider = AnthropicStyleAssemblyProvider()
+ body = "ASSEMBLY-CONTROL"
+
+ await StreamingOrchestrator({}).execute(
+ "work",
+ context,
+ {"main": provider},
+ {},
+ ScriptedHooks({"provider:request": _injection(body)}),
+ _retaining_coordinator(context),
+ )
+
+ assert len(provider.requests) == 1
+ assert context.request_calls == [([_wrap_reminders(body, tail=False)], None)]
+ assert body in "\n".join(message.content for message in provider.requests[0].messages)
+
+
+class ContributorSpy:
+ """Mount-level coordinator double retaining real contributor callables."""
+
+ def __init__(self) -> None:
+ self.contributors: list[tuple[str, str, object]] = []
+ self.capabilities: dict[str, object] = {}
+
+ def register_contributor(self, channel: str, name: str, callback) -> None:
+ self.contributors.append((channel, name, callback))
+
+ async def mount(self, _name: str, _module: object) -> None:
+ pass
+
+ def register_capability(self, name: str, capability: object) -> None:
+ self.capabilities[name] = capability
+
+
+@pytest.mark.asyncio
+async def test_mount_discovers_provider_budget_observability_event() -> None:
+ coordinator = ContributorSpy()
+
+ await mount(coordinator, {})
+
+ events_contributor = next(
+ callback
+ for channel, name, callback in coordinator.contributors
+ if (channel, name) == ("observability.events", "loop-streaming")
+ )
+ events = events_contributor()
+ assert events.count("orchestrator:provider_budget") == 1
+ assert {
+ "execution:start",
+ "execution:end",
+ "orchestrator:steering_injected",
+ "orchestrator:goal_progress",
+ "orchestrator:budget_warning",
+ }.issubset(events)
+
+
@pytest.mark.asyncio
async def test_finalization_irreducible_budget_skips_its_sdk_dispatch() -> None:
context = BudgetContext()
diff --git a/tests/test_request_retention.py b/tests/test_request_retention.py
index 6bff389..4f3f8b1 100644
--- a/tests/test_request_retention.py
+++ b/tests/test_request_retention.py
@@ -268,7 +268,9 @@ async def test_missing_retention_capability_warns_once_and_tail_remains_unmodifi
ScriptedHooks({"provider:request": injection(body)}),
coordinator_for(tail_context),
)
- assert tail_context.requirements == []
+ # A mounted retention capability owns every request view, even in explicit
+ # tail mode; tail still remains view-only and persists nothing.
+ assert tail_context.requirements == [[]]
assert persisted(tail_context) == []
assert reminder_contents(tail_provider.requests[0]) == [_wrap_reminders(body, tail=False)]
From d0b8ca82d2e1f9c1c44115b066aaab791e317a27 Mon Sep 17 00:00:00 2001
From: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com>
Date: Mon, 14 Sep 2026 22:53:19 -0700
Subject: [PATCH 2/3] fix: guard positional-only hard-fit retention
Generated with Amplifier
Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com>
---
amplifier_module_loop_streaming/__init__.py | 5 +-
..._provider_budget_context_simple_runtime.py | 23 ++++++----
tests/test_provider_budget_guard.py | 46 +++++++++++++++++++
3 files changed, 65 insertions(+), 9 deletions(-)
diff --git a/amplifier_module_loop_streaming/__init__.py b/amplifier_module_loop_streaming/__init__.py
index 80a04a4..c4d0cc7 100644
--- a/amplifier_module_loop_streaming/__init__.py
+++ b/amplifier_module_loop_streaming/__init__.py
@@ -3337,7 +3337,10 @@ def retention_accepts_hard_fit() -> bool:
# legacy behavior rather than claiming hard-fit support.
return False
return any(
- parameter.name == "hard_fit"
+ (
+ parameter.name == "hard_fit"
+ and parameter.kind is not inspect.Parameter.POSITIONAL_ONLY
+ )
or parameter.kind is inspect.Parameter.VAR_KEYWORD
for parameter in parameters
)
diff --git a/tests/test_provider_budget_context_simple_runtime.py b/tests/test_provider_budget_context_simple_runtime.py
index 230d2c2..fc4ff70 100644
--- a/tests/test_provider_budget_context_simple_runtime.py
+++ b/tests/test_provider_budget_context_simple_runtime.py
@@ -101,8 +101,10 @@ async def get_messages_for_request_retaining(
class _InMemoryResponses:
"""Completed Responses SDK fake whose usage is derived from received params."""
- def __init__(self) -> None:
+ def __init__(self, hard_fit_calls: list[bool]) -> None:
self.calls: list[dict] = []
+ self.hard_fit_counts_at_dispatch: list[int] = []
+ self._hard_fit_calls = hard_fit_calls
async def create(self, **params):
serialized = json.dumps(
@@ -111,6 +113,7 @@ async def create(self, **params):
# Record exactly the SDK payload, and derive usage from that same payload
# rather than a fixture outcome or a desired compaction size.
self.calls.append(json.loads(serialized))
+ self.hard_fit_counts_at_dispatch.append(self._hard_fit_calls.count(True))
input_tokens = len(serialized.encode("utf-8"))
return SimpleNamespace(
id=f"fake-{len(self.calls)}",
@@ -127,8 +130,8 @@ async def create(self, **params):
class _InMemoryClient:
- def __init__(self) -> None:
- self.responses = _InMemoryResponses()
+ def __init__(self, hard_fit_calls: list[bool]) -> None:
+ self.responses = _InMemoryResponses(hard_fit_calls)
def _payload_text(params: dict) -> str:
@@ -154,7 +157,7 @@ async def test_hard_fit_stays_compacted_across_real_openai_dispatches() -> None:
coordinator.register_capability(
"context.request_retention", context.get_messages_for_request_retaining
)
- client = _InMemoryClient()
+ client = _InMemoryClient(context.hard_fit_calls)
provider = OpenAIProvider(
api_key="test-key",
client=client,
@@ -208,12 +211,16 @@ async def test_hard_fit_stays_compacted_across_real_openai_dispatches() -> None:
# followed the provider-directed rebuild; the four subsequent calls prove
# ordinary fetches do not resurrect the canonical bulk history.
assert len(client.responses.calls) >= 5
+ payloads = [_payload_text(params) for params in client.responses.calls]
+ assert all("REMOVED-BULK-MARKER" not in payload for payload in payloads)
+ first_payload = payloads[0]
+ assert "ORIGINAL-HUMAN" in first_payload
+ assert "REQUIRED-REMINDER" in first_payload
+ assert context.hard_fit_calls[:2] == [False, True]
assert context.hard_fit_calls.count(True) == 1
- post_force_payloads = [
- _payload_text(params) for params in client.responses.calls[1:]
- ]
+ assert client.responses.hard_fit_counts_at_dispatch == [1] * len(payloads)
+ post_force_payloads = payloads[1:]
assert post_force_payloads
- assert all("REMOVED-BULK-MARKER" not in payload for payload in post_force_payloads)
assert all("ORIGINAL-HUMAN" in payload for payload in post_force_payloads)
assert "SECOND-HUMAN" in post_force_payloads[-1]
assert "REQUIRED-REMINDER" in post_force_payloads[-1]
diff --git a/tests/test_provider_budget_guard.py b/tests/test_provider_budget_guard.py
index 4e6fa84..4e756e8 100644
--- a/tests/test_provider_budget_guard.py
+++ b/tests/test_provider_budget_guard.py
@@ -115,6 +115,30 @@ async def retaining_view(self, **kwargs) -> list[dict]:
)
+class PositionalOnlyHardFitBudgetContext(BudgetContext):
+ """Legacy retention callable whose similarly named parameter is positional-only."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.hard_fit_values: list[bool] = []
+
+ async def retaining_view(
+ self,
+ hard_fit: bool = False,
+ /,
+ *,
+ provider=None,
+ retain_contents: list[str],
+ token_budget: int | None = None,
+ ) -> list[dict]:
+ self.hard_fit_values.append(hard_fit)
+ return await super().retaining_view(
+ provider=provider,
+ retain_contents=retain_contents,
+ token_budget=token_budget,
+ )
+
+
class OldSignatureBudgetContext(BudgetContext):
"""Pre-hard-fit retention capability; its call shape is the compatibility check."""
@@ -286,6 +310,28 @@ async def test_forced_rebuild_forwards_hard_fit_to_kwargs_retention() -> None:
assert context.hard_fit_calls == [False, True]
+@pytest.mark.asyncio
+async def test_forced_rebuild_treats_positional_only_hard_fit_as_legacy() -> None:
+ context = PositionalOnlyHardFitBudgetContext()
+ context._messages.append({"role": "assistant", "content": "history" * 200})
+ provider = BudgetProvider([_decision(100, 10, 7), _decision(9, 10, 0)])
+
+ await StreamingOrchestrator({}).execute(
+ "work",
+ context,
+ {"main": provider},
+ {},
+ ScriptedHooks({}),
+ _retaining_coordinator(context),
+ )
+
+ # Passing hard_fit by keyword would raise for this positional-only callable.
+ # Its default on both legacy-shaped calls proves the guard withheld that keyword.
+ assert context.hard_fit_values == [False, False]
+ assert [budget for _, budget in context.request_calls] == [None, 7]
+ assert len(provider.requests) == 1
+
+
@pytest.mark.asyncio
async def test_forced_rebuild_preserves_old_retention_call_signature() -> None:
context = OldSignatureBudgetContext()
From 7452bc9f77b4a8dfb5595da305cf4887baa8045f Mon Sep 17 00:00:00 2001
From: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com>
Date: Mon, 14 Sep 2026 23:11:47 -0700
Subject: [PATCH 3/3] fix(test): use UTF-8 provider budget fixture
Generated with Amplifier
Co-Authored-By: Amplifier <240397093+microsoft-amplifier@users.noreply.github.com>
---
tests/test_provider_budget_context_simple_runtime.py | 9 +++++----
1 file changed, 5 insertions(+), 4 deletions(-)
diff --git a/tests/test_provider_budget_context_simple_runtime.py b/tests/test_provider_budget_context_simple_runtime.py
index fc4ff70..d4b6e94 100644
--- a/tests/test_provider_budget_context_simple_runtime.py
+++ b/tests/test_provider_budget_context_simple_runtime.py
@@ -143,9 +143,10 @@ async def test_hard_fit_stays_compacted_across_real_openai_dispatches() -> None:
hooks = _StableReminderHooks()
coordinator = _Coordinator(hooks)
context = _RecordingContext(
- # Initial ordinary assembly is deliberately below this context's own
- # threshold. The provider's serialized-payload preflight forces the
- # first fit, exercising the optional hard-fit seam.
+ # The character-based ordinary estimate stays below context-simple's
+ # real provider-derived budget; max_tokens is only a fallback here.
+ # Each supplementary Han character serializes as four UTF-8 bytes, so
+ # the provider's payload preflight forces the first hard-fit rebuild.
max_tokens=500_000,
compact_threshold=0.99,
target_usage=0.50,
@@ -171,7 +172,7 @@ async def test_hard_fit_stays_compacted_across_real_openai_dispatches() -> None:
)
loop = StreamingOrchestrator({})
- bulk = "REMOVED-BULK-MARKER:" + ("x" * 800_000)
+ bulk = "REMOVED-BULK-MARKER:" + ("\U00020000" * 40_000)
await context.add_message({"role": "assistant", "content": bulk})
await loop.execute(