From 4e9caeb85ce0116c9f4171435fbf79ba5672fb50 Mon Sep 17 00:00:00 2001 From: Brian Krabach Date: Mon, 19 Oct 2026 08:39:03 -0700 Subject: [PATCH 1/2] Add shared scoped dispatch and attributed utility inference --- amplifier_module_loop_streaming/__init__.py | 56 +- amplifier_module_loop_streaming/dispatch.py | 534 ++++++++++++++++ docs/programmatic-capabilities.md | 38 ++ tests/test_approved_dispatch.py | 643 ++++++++++++++++++++ 4 files changed, 1262 insertions(+), 9 deletions(-) create mode 100644 amplifier_module_loop_streaming/dispatch.py create mode 100644 docs/programmatic-capabilities.md create mode 100644 tests/test_approved_dispatch.py diff --git a/amplifier_module_loop_streaming/__init__.py b/amplifier_module_loop_streaming/__init__.py index b1f488f..44d7d6a 100644 --- a/amplifier_module_loop_streaming/__init__.py +++ b/amplifier_module_loop_streaming/__init__.py @@ -762,12 +762,16 @@ async def mount(coordinator: ModuleCoordinator, config: dict[str, Any] | None = "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) "orchestrator:provider_overflow_recovery", # Bounded provider-authorized input-overflow retry + "model:complete", # Attributed stateless inference by an approved caller + "tool:dispatch_cancelled", # Nested execution outcome may be unknown ], ) orchestrator = StreamingOrchestrator(config) await coordinator.mount("orchestrator", orchestrator) coordinator.register_capability("session.steer", orchestrator.steer) + coordinator.register_capability("tools.dispatch", orchestrator.approved_dispatch) + coordinator.register_capability("model.complete", orchestrator.approved_dispatch.complete) # Conversation-scope provider pin (see ConversationProviderPin). Registered # as a capability, NOT read from config, so an app can detect its absence # (`get_capability` -> None) under a different orchestrator and refuse @@ -1088,6 +1092,8 @@ class StreamingOrchestrator: def __init__(self, config: dict[str, Any]): self.config = config + from .dispatch import ApprovedDispatch + self.approved_dispatch = ApprovedDispatch(self) # -1 means unlimited iterations (default) max_iter_config = config.get("max_iterations", -1) self.max_iterations = int(max_iter_config) if max_iter_config != -1 else -1 @@ -1379,7 +1385,15 @@ async def _drain_steering(self, context, hooks, iteration: int) -> int: ) return total - async def execute( + async def execute(self, prompt, context, providers, tools, hooks, coordinator=None): + # One execution authority is shared by foreground and nested tools. + # The lease expires even when the turn fails or is cancelled. + with self.approved_dispatch.run_scope(coordinator): + return await self._execute_with_goals( + prompt, context, providers, tools, hooks, coordinator + ) + + async def _execute_with_goals( self, prompt: str, context, @@ -1809,6 +1823,7 @@ async def _execute_one_turn( self._tool_calls_this_turn = 0 self._goal_turn_evidence = {"tools": []} if coordinator is not None else None self._llm_calls_this_turn = 0 + self.approved_dispatch.reset_model_budget() # Layer 1 call-budget bookkeeping, reset per turn alongside # _tool_calls_this_turn above (spec: 298-replacement). self._budget_exhausted = False @@ -6394,7 +6409,19 @@ async def _execute_tool( tool_call, tools, context, hooks, coordinator ) - async def _execute_tool_only( + async def _execute_tool_only(self, tool_call, tools, hooks, parallel_group_id, coordinator=None): + """Run the ordinary tool pipeline in an active dispatch scope. + + Tool errors become results; asyncio.CancelledError still propagates so + the caller can preserve completed sibling results during cancellation. + """ + if not self.approved_dispatch.enabled: + return await self._execute_tool_unscoped(tool_call, tools, hooks, parallel_group_id, coordinator) + tools = self.approved_dispatch.wrap_tools(tools) + with self.approved_dispatch.call_scope(tool_call, tools, hooks, parallel_group_id): + return await self._execute_tool_unscoped(tool_call, tools, hooks, parallel_group_id, coordinator) + + async def _execute_tool_unscoped( self, tool_call, tools: dict[str, Any], @@ -6486,6 +6513,7 @@ async def _execute_tool_only( if not hasattr(coordinator, "_tool_dispatch_contexts"): coordinator._tool_dispatch_contexts = {} _dispatch_task = asyncio.current_task() + _previous_dispatch = coordinator._tool_dispatch_contexts.get(_dispatch_task) coordinator._tool_dispatch_contexts[_dispatch_task] = { "tool_call_id": tool_call.id, "parallel_group_id": parallel_group_id, @@ -6506,9 +6534,10 @@ async def _execute_tool_only( coordinator.cancellation.register_tool_complete(tool_call.id) # Clear per-task dispatch context so completed tasks don't linger if hasattr(coordinator, "_tool_dispatch_contexts"): - coordinator._tool_dispatch_contexts.pop( - asyncio.current_task(), None - ) + if _previous_dispatch is None: + coordinator._tool_dispatch_contexts.pop(asyncio.current_task(), None) + else: + coordinator._tool_dispatch_contexts[asyncio.current_task()] = _previous_dispatch # Serialize result for logging result_data = ( @@ -6583,7 +6612,14 @@ async def _execute_tool_only( ) return (tool_call.id, tool_call.name, error_msg) - async def _execute_tool_with_result( + async def _execute_tool_with_result(self, tool_call, tools, context, hooks, coordinator=None): + if not self.approved_dispatch.enabled: + return await self._execute_tool_with_result_unscoped(tool_call, tools, context, hooks, coordinator) + tools = self.approved_dispatch.wrap_tools(tools) + with self.approved_dispatch.call_scope(tool_call, tools, hooks, None): + return await self._execute_tool_with_result_unscoped(tool_call, tools, context, hooks, coordinator) + + async def _execute_tool_with_result_unscoped( self, tool_call, tools: dict[str, Any], @@ -6647,6 +6683,7 @@ async def _execute_tool_with_result( if not hasattr(coordinator, "_tool_dispatch_contexts"): coordinator._tool_dispatch_contexts = {} _dispatch_task = asyncio.current_task() + _previous_dispatch = coordinator._tool_dispatch_contexts.get(_dispatch_task) coordinator._tool_dispatch_contexts[_dispatch_task] = { "tool_call_id": tool_call.id, "parallel_group_id": None, @@ -6664,9 +6701,10 @@ async def _execute_tool_with_result( finally: # Clear per-task dispatch context so completed tasks don't linger if coordinator and hasattr(coordinator, "_tool_dispatch_contexts"): - coordinator._tool_dispatch_contexts.pop( - asyncio.current_task(), None - ) + if _previous_dispatch is None: + coordinator._tool_dispatch_contexts.pop(asyncio.current_task(), None) + else: + coordinator._tool_dispatch_contexts[asyncio.current_task()] = _previous_dispatch # Serialize result for logging result_data = ( diff --git a/amplifier_module_loop_streaming/dispatch.py b/amplifier_module_loop_streaming/dispatch.py new file mode 100644 index 0000000..63ce3d6 --- /dev/null +++ b/amplifier_module_loop_streaming/dispatch.py @@ -0,0 +1,534 @@ +"""Active-call dispatch leases for portable programmatic orchestration. + +The ordinary streaming engine remains the execution and hook authority. This +adapter adds scope, nested attribution, and structured outcomes, never a second +unchecked tool execution path. +""" + +import asyncio +import json +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from uuid import uuid4 + +from amplifier_core.message_models import ToolCall + +_RUN = ContextVar("approved_dispatch_run", default=None) +_CALL = ContextVar("approved_dispatch_call", default=None) +FORBIDDEN = frozenset({"tool_exec", "delegate", "task", "live_job"}) + + +@dataclass +class Run: + owner: object + coordinator: object + identity: str + active: bool = True + model_calls: int = 0 + + +@dataclass +class CallScope: + run: Run + call: object + tools: dict + hooks: object + parallel_group_id: str | None + active: bool = True + + +class GuardedTool: + """Serialize actual execution, after the engine has processed approvals.""" + + def __init__(self, tool, lock, observation=None, authority=None): + self.tool, self.lock, self.observation = tool, lock, observation + self.authority = authority + + def __getattr__(self, name): + return getattr(self.tool, name) + + async def execute(self, arguments): + async with self.lock: + if self.authority: + self.authority() + if self.observation is not None: + self.observation["started"] = True + result = await self.tool.execute(arguments) + if self.observation is not None: + self.observation["raw_success"] = getattr(result, "success", None) + return result + + +class AttributedHooks: + def __init__(self, hooks, source, observation): + self.hooks, self.source, self.observation = hooks, source, observation + + async def emit(self, event, data): + enriched = { + **data, + "parent_tool_call_id": self.source["parent_tool_call_id"], + "dispatch_source": dict(self.source), + } + result = await self.hooks.emit(event, enriched) + if event == "tool:post": + self.observation["post_seen"] = True + returned = result.data.get("result") if result and result.data else None + self.observation["post_result"] = ( + returned if returned is not None else data.get("result") + ) + elif event == "tool:error": + self.observation["engine_error"] = True + return result + + def __getattr__(self, name): + return getattr(self.hooks, name) + + +class AttributedCoordinator: + def __init__(self, coordinator, observation): + object.__setattr__(self, "coordinator", coordinator) + object.__setattr__(self, "observation", observation) + + def __getattr__(self, name): + return getattr(self.coordinator, name) + + def __setattr__(self, name, value): + setattr(self.coordinator, name, value) + + async def process_hook_result(self, result, event, name): + result = await self.coordinator.process_hook_result(result, event, name) + if result.action == "deny": + self.observation["denied"] = True + return result + + +class DispatchLease: + def __init__(self, dispatcher, scope): + self.dispatcher, self.scope = dispatcher, scope + self.closed = False + allowlist = dispatcher.allowed_tools(scope.call.name) + self.tools = tuple( + name + for name in scope.tools + if name not in FORBIDDEN + and name not in dispatcher.callers + and getattr(scope.tools[name], "name", name) not in FORBIDDEN + and getattr(scope.tools[name], "name", name) not in dispatcher.callers + and (allowlist is None or name in allowlist) + ) + self.context = { + "session_id": str(scope.run.coordinator.session_id), + "run_id": scope.run.identity, + "parent_tool_call_id": scope.call.id, + } + + def close(self): + self.closed = True + + def _valid(self, expected_scope=None): + run = _RUN.get() + scope = _CALL.get() + if ( + self.closed + or not self.dispatcher.enabled + or self.scope.call.name not in self.dispatcher.callers + or not self.scope.active + or not self.scope.run.active + or run is not self.scope.run + or scope is not (expected_scope or self.scope) + or run.owner is not self.dispatcher + or run.coordinator is not self.scope.run.coordinator + ): + raise RuntimeError( + "Dispatch lease is not active in this session/run/parent call" + ) + if run.coordinator.cancellation.is_immediate: + raise RuntimeError("Immediate host cancellation is active") + validate = self.dispatcher.validate_context + if validate is not None and validate(run.coordinator) is not True: + raise RuntimeError("Host activation is no longer valid") + + async def wait_cancelled(self): + """Observe the real host token even across an async foreign-runtime bridge.""" + while True: + if self.scope.run.coordinator.cancellation.is_immediate: + return "host_cancel_requested" + try: + self._valid() + except RuntimeError: + return "lease_revoked" + await asyncio.sleep(0.05) + + async def call(self, name, arguments, *, request_id=None): + self._valid() + allowed = self.dispatcher.allowed_tools(self.scope.call.name) + if ( + name not in self.tools + or name in FORBIDDEN + or (allowed is not None and name not in allowed) + ): + raise ValueError("Tool is not available through this lease") + if not isinstance(arguments, dict) or arguments.get("async") is True: + raise ValueError("Background tool dispatch is unsupported") + + def still_authorized(child_scope): + self._valid(child_scope) + current = self.dispatcher.allowed_tools(self.scope.call.name) + if current is not None and name not in current: + raise RuntimeError("Nested tool authorization was revoked") + + call = ToolCall(id=str(uuid4()), name=name, arguments=arguments) + source = { + **self.context, + "tool_call_id": call.id, + "tool_name": name, + "request_id": request_id, + } + observation = {} + tools = self.dispatcher.wrap_tools( + self.scope.tools, observed_name=name, observation=observation + ) + hooks = AttributedHooks(self.scope.hooks, source, observation) + coordinator = AttributedCoordinator(self.scope.run.coordinator, observation) + # Always take the existing ordinary engine path. In particular, never + # call a mounted tool directly when hooks or authority are unavailable. + try: + with self.dispatcher.call_scope( + call, tools, hooks, self.scope.parallel_group_id + ) as child_scope: + tools = self.dispatcher.wrap_tools( + tools, + observed_name=name, + observation=observation, + authority=lambda: still_authorized(child_scope), + ) + _, _, content = await self.dispatcher.owner._execute_tool_unscoped( + call, + tools, + hooks, + self.scope.parallel_group_id or self.scope.call.id, + coordinator, + ) + except asyncio.CancelledError: + await self.scope.hooks.emit( + "tool:dispatch_cancelled", + { + "tool_name": name, + "tool_call_id": call.id, + "parent_tool_call_id": self.scope.call.id, + "dispatch_source": source, + "outcome": "unknown", + "effects": "not_rolled_back", + }, + ) + raise + self._valid() + if observation.get("denied"): + status, success, output = ( + "denied", + False, + "Tool execution denied by host policy", + ) + else: + post = observation.get("post_result") + confirmed = post.get("success") if isinstance(post, dict) else None + if observation.get("engine_error"): + status, success = "failed", False + elif observation.get("post_seen") and confirmed is True: + status, success = "completed", True + elif observation.get("post_seen") and confirmed is False: + status, success = "failed", False + else: + status, success = "unknown", None + try: + output = json.loads(content) + # The core hook transport may copy result dictionaries, making + # the ordinary engine serialize the full post-hook ToolResult. + # Unwrap only when it is exactly the captured confirmed envelope. + if ( + isinstance(post, dict) + and isinstance(post.get("success"), bool) + and output == post + ): + output = post.get("output") + except (ValueError, TypeError): + output = content + return { + "status": status, + "success": success, + "output": output, + "source": source, + } + + +class ApprovedDispatch: + version = 1 + + def __init__(self, owner): + self.owner = owner + self.locks = {} + self.validate_context = None + + @property + def enabled(self): + explicit = self.owner.config.get("programmatic_dispatch") + if explicit is not None: + return explicit is True + return bool(self.registrations) + + @property + def registrations(self): + run = _RUN.get() + if not run or run.owner is not self or not run.coordinator: + return {} + if not callable(getattr(run.coordinator, "get_capability", None)): + return {} + rows = run.coordinator.get_capability("tools.dispatch.registrations") + if not isinstance(rows, dict): + return {} + return { + name: row + for name, row in rows.items() + if isinstance(name, str) + and name not in FORBIDDEN + and isinstance(row, dict) + and isinstance(row.get("tools"), list) + and all(isinstance(tool, str) for tool in row["tools"]) + } + + @property + def callers(self): + """Host-selected controllers; ordinary configuration permits only tool_exec.""" + names = self.owner.config.get("programmatic_callers") + if names is None: + names = ["tool_exec", *self.registrations] + if not isinstance(names, list) or any( + not isinstance(name, str) for name in names + ): + return frozenset() + return frozenset( + name for name in names if name == "tool_exec" or name not in FORBIDDEN + ) + + def allowed_tools(self, caller): + host = self.owner.config.get("programmatic_allowed_tools") + declared = self.registrations.get(caller, {}).get("tools") + if declared is None: + return host + return set(declared) if host is None else set(declared).intersection(host) + + def reset_model_budget(self): + run = _RUN.get() + if run and run.owner is self: + run.model_calls = 0 + + @contextmanager + def run_scope(self, coordinator): + existing = _RUN.get() + if ( + existing + and existing.owner is self + and existing.coordinator is coordinator + and existing.active + ): + # loop-live surrounds its finite streaming turns with a longer run. + yield + return + run = Run(self, coordinator, str(uuid4())) + token = _RUN.set(run) + try: + yield + finally: + run.active = False + _RUN.reset(token) + + @contextmanager + def call_scope(self, call, tools, hooks, parallel_group_id): + run = _RUN.get() + scope = ( + CallScope(run, call, tools, hooks, parallel_group_id) + if run and run.owner is self + else None + ) + token = _CALL.set(scope) + try: + yield scope + finally: + if scope: + scope.active = False + _CALL.reset(token) + + def bind(self): + scope = _CALL.get() + if ( + not self.enabled + or scope is None + or scope.call.name not in self.callers + or not scope.active + or not scope.run.active + or scope.run.owner is not self + or scope.run.coordinator is None + ): + raise RuntimeError("No active authorized programmatic caller") + lease = DispatchLease(self, scope) + lease._valid() + return lease + + async def complete(self, request): + """A bounded, attributed ordinary inference call for an active caller. + + This is a stateless utility request, not another foreground turn: it + never mutates conversation history or its context occupancy meter. + Provider-owned usage/cost hooks remain active, with streaming disabled. + """ + from amplifier_core.message_models import ChatRequest + + lease = self.bind() + scope = lease.scope + coordinator = scope.run.coordinator + try: + if not isinstance(request, ChatRequest) or request.tools: + raise ValueError("Nested inference requires a tool-free ChatRequest") + model_limit = self.owner.config.get("programmatic_max_model_calls", 24) + if type(model_limit) is not int or model_limit < 0: + raise ValueError("Invalid nested inference budget") + # Serialize admission so concurrent callers cannot overspend this turn. + lock = self.locks.setdefault("model.complete", asyncio.Lock()) + async with lock: + lease._valid() + if scope.run.model_calls >= model_limit: + raise RuntimeError("Nested inference budget exhausted") + provider = self.owner._select_provider(coordinator.get("providers")) + if provider is None: + raise RuntimeError("No ordinary provider is configured") + model = self.owner._provider_default_model(provider) + # Reuse the current provider/account/model. The consumer cannot + # secretly override model routing or turn on provider streaming. + request = request.model_copy( + update={ + "model": model, + "metadata": {**(request.metadata or {}), "stream": False}, + "max_output_tokens": min( + request.max_output_tokens or 2048, 4096 + ), + }, + deep=True, + ) + source = { + **lease.context, + "caller": scope.call.name, + "purpose": "nested_inference", + } + result = await scope.hooks.emit( + "provider:request", + { + "provider": getattr(provider, "name", "unknown"), + "model": model, + "iteration": 0, + **source, + }, + ) + result = await coordinator.process_hook_result( + result, "provider:request", scope.call.name + ) + if result.action == "deny": + raise RuntimeError("Nested inference denied by host policy") + lease._valid() + options = {"extended_thinking": False} + if model: + options["model"] = model + budget = getattr(provider, "request_budget", None) + if callable(budget): + import inspect + + from ._text_estimate import estimate_messages + + kwargs = { + "context_estimate": estimate_messages( + [ + message.model_dump(exclude_none=True) + for message in request.messages + ] + ).tokens + } + if "request_options" in inspect.signature(budget).parameters: + kwargs["request_options"] = options + measured = budget(request, **kwargs) + if inspect.isawaitable(measured): + measured = await measured + if isinstance(measured, dict): + count = measured.get("estimated_input_tokens") + limit = measured.get("input_limit_tokens") + if type(count) is int and type(limit) is int and count > limit: + raise RuntimeError( + "Nested inference exceeds provider input budget" + ) + lease._valid() + scope.run.model_calls += 1 + self.owner._llm_calls_this_turn += 1 + pending = asyncio.create_task(provider.complete(request, **options)) + try: + while not pending.done(): + await asyncio.wait({pending}, timeout=0.05) + if coordinator.cancellation.is_cancelled: + raise asyncio.CancelledError + lease._valid() + response = await pending + finally: + if not pending.done(): + pending.cancel() + await asyncio.gather(pending, return_exceptions=True) + lease._valid() + usage = getattr(response, "usage", None) + await scope.hooks.emit( + "model:complete", + { + **source, + "provider": getattr(provider, "name", "unknown"), + "model": model, + "usage": usage.model_dump() + if hasattr(usage, "model_dump") + else usage, + "nested_model_calls": scope.run.model_calls, + }, + ) + return response + except asyncio.CancelledError: + raise + except Exception as exc: + await scope.hooks.emit( + "provider:error", + { + **lease.context, + "purpose": "nested_inference", + "error": { + "type": type(exc).__name__, + "msg": "Nested inference failed", + }, + }, + ) + raise + finally: + lease.close() + + def wrap_tools(self, tools, observed_name=None, observation=None, authority=None): + wrapped = {} + for name, tool in tools.items(): + raw = tool.tool if isinstance(tool, GuardedTool) else tool + allowlist = self.owner.config.get("programmatic_allowed_tools") + if name != "tool_exec" and ( + name in FORBIDDEN + or getattr(raw, "name", name) in FORBIDDEN + or (allowlist is not None and name not in allowlist) + ): + wrapped[name] = raw + continue + lock = self.locks.setdefault(id(raw), asyncio.Lock()) + wrapped[name] = GuardedTool( + raw, + lock, + observation if name == observed_name else None, + authority if name == observed_name else None, + ) + return wrapped diff --git a/docs/programmatic-capabilities.md b/docs/programmatic-capabilities.md new file mode 100644 index 0000000..32856c4 --- /dev/null +++ b/docs/programmatic-capabilities.md @@ -0,0 +1,38 @@ +# Programmatic execution capabilities (development candidate) + +`tools.dispatch` and `model.complete` are optional capabilities mounted by the +streaming loop and reused by loop-live. They contain no UI-use or vendor logic. +Consumers bind during an admitted parent tool call. The lease is scoped to that +session, run and parent, expires on completion, and cannot be used from an idle +session. Nested actions re-enter the ordinary tool hooks and approval processor. +They retain parent/call identities, serialize actual execution, and report denied, +failed and unknown outcomes separately. Cancellation never implies rollback. + +Root configuration can explicitly enable `programmatic_dispatch` and select +`programmatic_callers` / `programmatic_allowed_tools`, as before. A portable tool +behavior can instead register a named entry in `tools.dispatch.registrations`: + +```python +{"controller_name": {"tools": ["permitted_surface_tool"]}} +``` + +This is a trusted module declaration made at mount, not an argument supplied by +a model. Explicit host `programmatic_dispatch: false` disables it; an explicit +caller list constrains it; a host tool allowlist intersects it. Removing the +declaration revokes it. A module cleans up only the declaration it owns. Nested +delegation, background execution and recursive controllers remain unavailable. + +`model.complete(ChatRequest)` accepts a stateless tool-free request only during +such an active call. The selected ordinary provider and its account are reused; +the caller cannot select a different model through this service. Provider request +hooks may deny the call, input budgets are checked, streaming is disabled, and +cancellation reaches the in-flight provider task. Provider usage/cost reporting +remains active. A `model:complete` receipt adds parent attribution without copying +the prompt or result into diagnostics. These calls increment the turn's actual +`llm_calls`; they do not mutate history or its context-occupancy estimate. + +`programmatic_max_model_calls` bounds these utility calls per turn (default 24, +zero disables them). It is separate from the existing foreground iteration cap, +not a claim that routing, goal evaluation, decision endpoints and financial +budgets are already governed by a single meter. A broader inference service can +be extracted later without putting that policy in a UI controller. diff --git a/tests/test_approved_dispatch.py b/tests/test_approved_dispatch.py new file mode 100644 index 0000000..bb5fb5f --- /dev/null +++ b/tests/test_approved_dispatch.py @@ -0,0 +1,643 @@ +"""Contract tests against the real coordinator and ordinary engine hook path.""" + +import asyncio +import json +from types import SimpleNamespace + +import pytest +import pytest_asyncio +from amplifier_core import AmplifierSession, HookResult, ToolResult +from amplifier_core.message_models import ChatResponse, TextBlock, ToolCall + + +pytestmark = pytest.mark.asyncio + + +class Echo: + name = "read" + description = "Fixture read" + input_schema = {"type": "object", "properties": {}} + + def __init__(self): + self.calls = 0 + self.active = self.peak = 0 + self.started = asyncio.Event() + self.release = None + + async def execute(self, arguments): + self.calls += 1 + self.active += 1 + self.peak = max(self.peak, self.active) + self.started.set() + try: + if self.release: + await self.release.wait() + return ToolResult( + success=True, output={"value": arguments.get("value", 42)} + ) + finally: + self.active -= 1 + + +class Client: + name = "tool_exec" + description = "Fixture orchestration client" + input_schema = {"type": "object", "properties": {}} + + def __init__(self, coordinator, handler=None): + self.coordinator, self.handler = coordinator, handler + self.lease = None + + async def execute(self, arguments): + self.lease = self.coordinator.get_capability("tools.dispatch").bind() + result = await ( + self.handler(self.lease) + if self.handler + else self.lease.call("read", arguments, request_id="1") + ) + return ToolResult(success=True, output=result) + + +@pytest_asyncio.fixture +async def environment(): + class Approval: + def __init__(self): + self.calls = [] + self.waiting = asyncio.Event() + self.release = None + self.answer = "deny" + + async def request_approval(self, prompt, options, timeout, default): + self.calls.append(prompt) + self.waiting.set() + if self.release: + await self.release.wait() + return self.answer + + approval = Approval() + session = AmplifierSession( + { + "session": { + "orchestrator": { + "module": "loop-streaming", + "config": { + "min_delay_between_calls_ms": 0, + "programmatic_dispatch": True, + }, + }, + "context": {"module": "context-simple"}, + }, + "providers": [], + }, + approval_system=approval, + ) + await session.initialize() + coordinator = session.coordinator + coordinator._test_approval = approval + loop = coordinator.get("orchestrator") + echo = Echo() + await coordinator.mount("tools", echo, name="read") + try: + yield session, coordinator, loop, echo + finally: + await session.cleanup() + + +async def call_client(environment, client=None): + _, coordinator, loop, _ = environment + client = client or Client(coordinator) + await coordinator.mount("tools", client, name="tool_exec") + with loop.approved_dispatch.run_scope(coordinator): + _, _, result = await loop._execute_tool_only( + ToolCall(id="parent", name="tool_exec", arguments={}), + coordinator.get("tools"), + coordinator.hooks, + "group", + coordinator, + ) + return json.loads(result)["output"], client + + +async def test_nested_uses_standard_hooks_once_and_keeps_identity(environment): + session, coordinator, loop, echo = environment + events = [] + + async def capture(event, data): + events.append((event, dict(data))) + return HookResult() + + coordinator.hooks.register("tool:pre", capture) + coordinator.hooks.register("tool:post", capture) + result, client = await call_client(environment) + assert result["status"] == "completed" + assert result["success"] is True + assert result["output"] == {"value": 42} + assert echo.calls == 1 + nested = [(event, data) for event, data in events if data["tool_name"] == "read"] + assert [event for event, _ in nested] == ["tool:pre", "tool:post"] + source = result["source"] + assert source["session_id"] == session.session_id + assert source["parent_tool_call_id"] == "parent" + assert source["request_id"] == "1" + assert source["tool_call_id"] != "parent" + assert all(data["dispatch_source"] == source for _, data in nested) + assert loop._tool_calls_this_turn == 2 + with pytest.raises(RuntimeError): + await client.lease.call("read", {}) + + +async def test_hook_denial_never_executes_nested_tool(environment): + _, coordinator, _, echo = environment + + async def deny(event, data): + return ( + HookResult(action="deny", reason="Policy rejected") + if data["tool_name"] == "read" + else HookResult() + ) + + coordinator.hooks.register("tool:pre", deny) + result, _ = await call_client(environment) + assert result["status"] == "denied" + assert result["success"] is False + assert echo.calls == 0 + + +async def test_coordinator_approval_rejection_is_preserved(environment): + _, coordinator, _, echo = environment + + async def policy(event, data): + if data["tool_name"] == "read": + return HookResult( + action="ask_user", + approval_prompt="Allow nested read?", + approval_options=["allow", "deny"], + approval_default="deny", + ) + return HookResult() + + coordinator.hooks.register("tool:pre", policy) + result, _ = await call_client(environment) + assert coordinator._test_approval.calls == ["Allow nested read?"] + assert result["status"] == "denied" + assert echo.calls == 0 + + +async def test_post_hook_modified_output_and_unknown_outcome_are_preserved(environment): + _, coordinator, _, _ = environment + + async def change(event, data): + if data["tool_name"] == "read": + return HookResult( + action="modify", data={**data, "result": {"output": "redacted"}} + ) + return HookResult() + + coordinator.hooks.register("tool:post", change) + result, _ = await call_client(environment) + assert result["status"] == "unknown" + assert result["success"] is None + assert result["output"] == {"output": "redacted"} + + +async def test_bind_rejects_idle_and_wrong_parent(environment): + _, coordinator, loop, echo = environment + dispatch = coordinator.get_capability("tools.dispatch") + with pytest.raises(RuntimeError): + dispatch.bind() + with dispatch.run_scope(coordinator): + with dispatch.call_scope( + ToolCall(id="ordinary", name="read", arguments={}), + {"read": echo}, + coordinator.hooks, + None, + ): + with pytest.raises(RuntimeError): + dispatch.bind() + + +async def test_lease_cannot_cross_run_session_or_parent_scope(environment): + _, coordinator, loop, echo = environment + dispatch = loop.approved_dispatch + parent = ToolCall(id="parent", name="tool_exec", arguments={}) + with dispatch.run_scope(coordinator): + with dispatch.call_scope(parent, {"read": echo}, coordinator.hooks, None): + lease = dispatch.bind() + with dispatch.run_scope(SimpleNamespace(session_id="other")): + with pytest.raises(RuntimeError): + await lease.call("read", {}) + with dispatch.call_scope( + ToolCall(id="other-parent", name="tool_exec", arguments={}), + {"read": echo}, + coordinator.hooks, + None, + ): + with pytest.raises(RuntimeError): + await lease.call("read", {}) + lease.close() + with pytest.raises(RuntimeError): + await lease.call("read", {}) + assert echo.calls == 0 + + +async def test_nested_cancellation_emits_unknown_outcome_with_identity(environment): + _, coordinator, _, echo = environment + echo.release = asyncio.Event() + seen = [] + + async def capture(event, data): + seen.append(data) + return HookResult() + + coordinator.hooks.register("tool:dispatch_cancelled", capture) + task = asyncio.create_task(call_client(environment)) + await asyncio.wait_for(echo.started.wait(), 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert echo.active == 0 + assert seen[0]["outcome"] == "unknown" + assert seen[0]["effects"] == "not_rolled_back" + assert seen[0]["dispatch_source"]["parent_tool_call_id"] == "parent" + + +async def test_pending_approval_holds_no_tool_execution_lock(environment): + _, coordinator, _, echo = environment + approval = coordinator._test_approval + approval.release = asyncio.Event() + approval.answer = "allow" + + async def policy(event, data): + if data["tool_name"] == "read" and data["tool_input"].get("approval"): + return HookResult( + action="ask_user", + approval_prompt="Allow slow read?", + approval_options=["allow", "deny"], + approval_default="deny", + ) + return HookResult() + + coordinator.hooks.register("tool:pre", policy) + + async def program(lease): + pending = asyncio.create_task(lease.call("read", {"approval": True})) + await asyncio.wait_for(approval.waiting.wait(), 1) + immediate = await asyncio.wait_for( + lease.call("read", {"value": "unrelated"}), 1 + ) + assert immediate["success"] is True + approval.release.set() + await pending + return immediate + + result, _ = await call_client(environment, Client(coordinator, program)) + assert result["success"] is True + assert echo.calls == 2 + + +async def test_same_tool_executes_serially_in_parallel_nested_calls(environment): + _, coordinator, _, echo = environment + echo.release = asyncio.Event() + + async def program(lease): + one = asyncio.create_task(lease.call("read", {})) + await echo.started.wait() + two = asyncio.create_task(lease.call("read", {})) + await asyncio.sleep(0) + assert echo.calls == 1 + echo.release.set() + return await asyncio.gather(one, two) + + result, _ = await call_client(environment, Client(coordinator, program)) + assert all(row["success"] for row in result) + assert echo.peak == 1 + + +async def test_ordinary_tool_is_invoked_exactly_once(environment): + _, coordinator, loop, echo = environment + with loop.approved_dispatch.run_scope(coordinator): + result = await loop._execute_tool_only( + ToolCall(id="ordinary", name="read", arguments={}), + coordinator.get("tools"), + coordinator.hooks, + None, + coordinator, + ) + assert json.loads(result[2])["output"] == {"value": 42} + assert echo.calls == 1 + + +async def test_real_finite_session_enters_and_expires_lease_scope(environment): + session, coordinator, _, echo = environment + client = Client(coordinator) + await coordinator.mount("tools", client, name="tool_exec") + + class Provider: + name = "fixture" + config = {} + count = 0 + + def get_info(self): + return SimpleNamespace(default_model="fixture") + + async def complete(self, request, **kwargs): + self.count += 1 + if self.count == 1: + return ChatResponse( + content=[], + tool_calls=[ + ToolCall(id="actual-parent", name="tool_exec", arguments={}) + ], + ) + assert "actual-parent" in str(request.messages) + return ChatResponse(content=[TextBlock(text="done")]) + + def parse_tool_calls(self, response): + return response.tool_calls or [] + + await coordinator.mount("providers", Provider(), name="fixture") + assert await session.execute("Use the programmatic tool") == "done" + assert echo.calls == 1 + with pytest.raises(RuntimeError): + await client.lease.call("read", {}) + + +async def test_revoked_lease_cannot_execute_after_delayed_approval(environment): + _, coordinator, loop, echo = environment + approval = coordinator._test_approval + approval.release = asyncio.Event() + approval.answer = "allow" + + async def policy(event, data): + return HookResult( + action="ask_user", + approval_prompt="Delayed approval", + approval_options=["allow", "deny"], + ) + + coordinator.hooks.register("tool:pre", policy) + dispatch = loop.approved_dispatch + with dispatch.run_scope(coordinator): + with dispatch.call_scope( + ToolCall(id="parent", name="tool_exec", arguments={}), + {"read": echo}, + coordinator.hooks, + None, + ): + lease = dispatch.bind() + task = asyncio.create_task(lease.call("read", {})) + await asyncio.wait_for(approval.waiting.wait(), 1) + lease.close() + approval.release.set() + with pytest.raises(RuntimeError): + await task + assert echo.calls == 0 + + +async def test_disabled_dispatch_preserves_concurrent_ordinary_calls(environment): + _, coordinator, loop, echo = environment + loop.config.pop("programmatic_dispatch") + echo.release = asyncio.Event() + with loop.approved_dispatch.run_scope(coordinator): + with pytest.raises(RuntimeError): + loop.approved_dispatch.bind() + tasks = [ + asyncio.create_task( + loop._execute_tool_only( + ToolCall(id=f"ordinary-{i}", name="read", arguments={}), + coordinator.get("tools"), + coordinator.hooks, + "ordinary", + coordinator, + ) + ) + for i in range(2) + ] + for _ in range(20): + if echo.calls == 2: + break + await asyncio.sleep(0.01) + assert echo.calls == 2 + assert echo.peak == 2 + echo.release.set() + await asyncio.gather(*tasks) + + +async def test_lease_wait_observes_real_coordinator_immediate_cancellation(environment): + _, coordinator, loop, echo = environment + dispatch = loop.approved_dispatch + with dispatch.run_scope(coordinator): + with dispatch.call_scope( + ToolCall(id="parent", name="tool_exec", arguments={}), + {"read": echo}, + coordinator.hooks, + None, + ): + lease = dispatch.bind() + waiting = asyncio.create_task(lease.wait_cancelled()) + await coordinator.request_cancel(immediate=True) + assert await asyncio.wait_for(waiting, 1) == "host_cancel_requested" + with pytest.raises(RuntimeError): + await lease.call("read", {}) + assert echo.calls == 0 + + +async def test_delegation_aliases_are_excluded_and_keep_ordinary_concurrency( + environment, +): + _, coordinator, loop, echo = environment + echo.name = "delegate" + tools = {"alias": echo} + dispatch = loop.approved_dispatch + with dispatch.run_scope(coordinator): + with dispatch.call_scope( + ToolCall(id="parent", name="tool_exec", arguments={}), + tools, + coordinator.hooks, + None, + ): + lease = dispatch.bind() + assert lease.tools == () + with pytest.raises(ValueError): + await lease.call("alias", {}) + assert dispatch.wrap_tools(tools)["alias"] is echo + + +async def test_portable_registration_requires_no_orchestrator_configuration( + environment, +): + _, coordinator, loop, echo = environment + loop.config.pop("programmatic_dispatch") + coordinator.register_capability( + "tools.dispatch.registrations", {"ui_use": {"tools": ["read"]}} + ) + client = Client(coordinator) + client.name = "ui_use" + await coordinator.mount("tools", client, name="ui_use") + with loop.approved_dispatch.run_scope(coordinator): + _, _, raw = await loop._execute_tool_only( + ToolCall(id="ui", name="ui_use", arguments={}), + coordinator.get("tools"), + coordinator.hooks, + "ui", + coordinator, + ) + assert json.loads(raw)["output"]["status"] == "completed" and echo.calls == 1 + assert client.lease.tools == ("read",) + loop.config["programmatic_dispatch"] = False + client.lease = None + with loop.approved_dispatch.run_scope(coordinator): + await loop._execute_tool_only( + ToolCall(id="off", name="ui_use", arguments={}), + coordinator.get("tools"), + coordinator.hooks, + "off", + coordinator, + ) + assert client.lease is None and echo.calls == 1 + + +async def test_nested_inference_reuses_selection_counts_and_enforces_policy( + environment, +): + from amplifier_core.message_models import ChatRequest, Message + + _, coordinator, loop, _ = environment + seen, events = [], [] + + class Provider: + name = "ordinary" + + def get_info(self): + return SimpleNamespace(defaults={"model": "selected-model"}) + + def request_budget(self, request, *, context_estimate, request_options): + assert context_estimate > 0 and request_options["model"] == "selected-model" + return {"estimated_input_tokens": 10, "input_limit_tokens": 100} + + async def complete(self, request, **kwargs): + seen.append((request, kwargs)) + return ChatResponse(content=[TextBlock(text="chosen")]) + + await coordinator.mount("providers", Provider(), name="ordinary") + request = ChatRequest( + messages=[Message(role="user", content="Pick a candidate")], model="unselected" + ) + + async def capture(event, data): + events.append((event, data)) + return HookResult() + + coordinator.hooks.register("provider:request", capture) + coordinator.hooks.register("model:complete", capture) + loop.config["programmatic_max_model_calls"] = 1 + + async def program(lease): + response = await coordinator.get_capability("model.complete")(request) + with pytest.raises(RuntimeError, match="budget exhausted"): + await coordinator.get_capability("model.complete")(request) + return {"text": response.content[0].text} + + result, _ = await call_client(environment, Client(coordinator, program)) + assert isinstance(result, dict), result + assert result["text"] == "chosen" and loop._llm_calls_this_turn == 1 + assert ( + seen[0][0].model == "selected-model" and seen[0][0].metadata["stream"] is False + ) + assert request.model == "unselected" and request.metadata is None + assert [event for event, _ in events] == ["provider:request", "model:complete"] + assert all(data["parent_tool_call_id"] == "parent" for _, data in events) + with pytest.raises(RuntimeError, match="authorized"): + await coordinator.get_capability("model.complete")(request) + + +async def test_nested_inference_denial_does_not_call_provider(environment): + from amplifier_core.message_models import ChatRequest, Message + + _, coordinator, loop, _ = environment + + class Provider: + name = "ordinary" + + def get_info(self): + return SimpleNamespace(defaults={"model": "selected"}) + + async def complete(self, request, **kwargs): + raise AssertionError("Policy denied inference") + + await coordinator.mount("providers", Provider(), name="ordinary") + + async def deny(event, data): + return HookResult(action="deny", reason="test policy") + + coordinator.hooks.register("provider:request", deny) + + async def program(lease): + with pytest.raises(RuntimeError, match="denied"): + await coordinator.get_capability("model.complete")( + ChatRequest(messages=[Message(role="user", content="private")]) + ) + return {"denied": True} + + result, _ = await call_client(environment, Client(coordinator, program)) + assert result["denied"] and loop._llm_calls_this_turn == 0 + + +async def test_nested_dispatch_restores_parent_task_identity(environment): + _, coordinator, _, _ = environment + + async def program(lease): + before = dict(coordinator._tool_dispatch_contexts[asyncio.current_task()]) + result = await lease.call("read", {}) + assert coordinator._tool_dispatch_contexts[asyncio.current_task()] == before + return result + + result, _ = await call_client(environment, Client(coordinator, program)) + assert result["status"] == "completed" + + +async def test_nested_inference_observes_host_cancel_and_cancels_provider(environment): + from amplifier_core.message_models import ChatRequest, Message + + _, coordinator, _, _ = environment + started, cancelled = asyncio.Event(), asyncio.Event() + + class Provider: + name = "ordinary" + + def get_info(self): + return SimpleNamespace(defaults={"model": "selected"}) + + async def complete(self, request, **kwargs): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + await coordinator.mount("providers", Provider(), name="ordinary") + + async def program(lease): + return await coordinator.get_capability("model.complete")( + ChatRequest(messages=[Message(role="user", content="private")]) + ) + + pending = asyncio.create_task( + call_client(environment, Client(coordinator, program)) + ) + await asyncio.wait_for(started.wait(), 1) + await coordinator.request_cancel(immediate=True) + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, 1) + assert cancelled.is_set() + + +async def test_host_validity_callback_revokes_without_live_dependency(environment): + _, coordinator, loop, echo = environment + async def program(lease): + loop.approved_dispatch.validate_context = lambda actual: False + with pytest.raises(RuntimeError, match='no longer valid'): + await lease.call('read', {}) + return {'revoked': True} + result, _ = await call_client(environment, Client(coordinator, program)) + assert result['revoked'] and echo.calls == 0 From 20792f1cc7afe0945a789db61a3fa476306efd69 Mon Sep 17 00:00:00 2001 From: Brian Krabach Date: Wed, 21 Oct 2026 21:00:58 -0700 Subject: [PATCH 2/2] Preserve legacy live dispatcher ownership during independent upgrades --- .github/workflows/tests.yml | 1 + amplifier_module_loop_streaming/__init__.py | 30 +++++++--- amplifier_module_loop_streaming/dispatch.py | 6 +- docs/programmatic-capabilities.md | 4 ++ pyproject.toml | 2 +- tests/test_approved_dispatch.py | 62 +++++++++++++++++++++ uv.lock | 9 ++- 7 files changed, 102 insertions(+), 12 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index c418f39..98a47f3 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -30,6 +30,7 @@ jobs: uv pip install --no-config --python .venv/bin/python \ "amplifier-core @ git+https://github.com/microsoft/amplifier-core@main" \ "amplifier-foundation @ git+https://github.com/microsoft/amplifier-foundation@main" \ + "amplifier-module-context-simple @ git+https://github.com/microsoft/amplifier-module-context-simple@main" \ "pytest==9.0.3" \ "pytest-asyncio==1.3.0" uv pip install --no-config --no-deps --python .venv/bin/python . diff --git a/amplifier_module_loop_streaming/__init__.py b/amplifier_module_loop_streaming/__init__.py index 44d7d6a..f6294ef 100644 --- a/amplifier_module_loop_streaming/__init__.py +++ b/amplifier_module_loop_streaming/__init__.py @@ -1094,6 +1094,7 @@ def __init__(self, config: dict[str, Any]): self.config = config from .dispatch import ApprovedDispatch self.approved_dispatch = ApprovedDispatch(self) + self._streaming_dispatch = self.approved_dispatch # -1 means unlimited iterations (default) max_iter_config = config.get("max_iterations", -1) self.max_iterations = int(max_iter_config) if max_iter_config != -1 else -1 @@ -1388,10 +1389,22 @@ async def _drain_steering(self, context, hooks, iteration: int) -> int: async def execute(self, prompt, context, providers, tools, hooks, coordinator=None): # One execution authority is shared by foreground and nested tools. # The lease expires even when the turn fails or is cancelled. - with self.approved_dispatch.run_scope(coordinator): + if not self._owns_dispatch_pipeline: return await self._execute_with_goals( prompt, context, providers, tools, hooks, coordinator ) + with self._streaming_dispatch.run_scope(coordinator): + return await self._execute_with_goals( + prompt, context, providers, tools, hooks, coordinator + ) + + @property + def _owns_dispatch_pipeline(self): + # Older loop-live versions replace the public dispatcher and already + # establish their own scopes and guarded tools. Preserve that extension + # boundary during independent module upgrades: never double-wrap their + # pipeline or call new dispatcher methods on their legacy instance. + return self.approved_dispatch is self._streaming_dispatch async def _execute_with_goals( self, @@ -1823,7 +1836,8 @@ async def _execute_one_turn( self._tool_calls_this_turn = 0 self._goal_turn_evidence = {"tools": []} if coordinator is not None else None self._llm_calls_this_turn = 0 - self.approved_dispatch.reset_model_budget() + if self._owns_dispatch_pipeline: + self._streaming_dispatch.reset_model_budget() # Layer 1 call-budget bookkeeping, reset per turn alongside # _tool_calls_this_turn above (spec: 298-replacement). self._budget_exhausted = False @@ -6415,10 +6429,10 @@ async def _execute_tool_only(self, tool_call, tools, hooks, parallel_group_id, c Tool errors become results; asyncio.CancelledError still propagates so the caller can preserve completed sibling results during cancellation. """ - if not self.approved_dispatch.enabled: + if not self._owns_dispatch_pipeline or not self._streaming_dispatch.enabled: return await self._execute_tool_unscoped(tool_call, tools, hooks, parallel_group_id, coordinator) - tools = self.approved_dispatch.wrap_tools(tools) - with self.approved_dispatch.call_scope(tool_call, tools, hooks, parallel_group_id): + tools = self._streaming_dispatch.wrap_tools(tools) + with self._streaming_dispatch.call_scope(tool_call, tools, hooks, parallel_group_id): return await self._execute_tool_unscoped(tool_call, tools, hooks, parallel_group_id, coordinator) async def _execute_tool_unscoped( @@ -6613,10 +6627,10 @@ async def _execute_tool_unscoped( return (tool_call.id, tool_call.name, error_msg) async def _execute_tool_with_result(self, tool_call, tools, context, hooks, coordinator=None): - if not self.approved_dispatch.enabled: + if not self._owns_dispatch_pipeline or not self._streaming_dispatch.enabled: return await self._execute_tool_with_result_unscoped(tool_call, tools, context, hooks, coordinator) - tools = self.approved_dispatch.wrap_tools(tools) - with self.approved_dispatch.call_scope(tool_call, tools, hooks, None): + tools = self._streaming_dispatch.wrap_tools(tools) + with self._streaming_dispatch.call_scope(tool_call, tools, hooks, None): return await self._execute_tool_with_result_unscoped(tool_call, tools, context, hooks, coordinator) async def _execute_tool_with_result_unscoped( diff --git a/amplifier_module_loop_streaming/dispatch.py b/amplifier_module_loop_streaming/dispatch.py index 63ce3d6..17686e8 100644 --- a/amplifier_module_loop_streaming/dispatch.py +++ b/amplifier_module_loop_streaming/dispatch.py @@ -299,10 +299,12 @@ def registrations(self): @property def callers(self): - """Host-selected controllers; ordinary configuration permits only tool_exec.""" + """Admit declared controllers, or the explicitly enabled legacy caller.""" names = self.owner.config.get("programmatic_callers") if names is None: - names = ["tool_exec", *self.registrations] + names = list(self.registrations) + if self.owner.config.get("programmatic_dispatch") is True: + names.append("tool_exec") if not isinstance(names, list) or any( not isinstance(name, str) for name in names ): diff --git a/docs/programmatic-capabilities.md b/docs/programmatic-capabilities.md index 32856c4..b5bd789 100644 --- a/docs/programmatic-capabilities.md +++ b/docs/programmatic-capabilities.md @@ -21,6 +21,10 @@ a model. Explicit host `programmatic_dispatch: false` disables it; an explicit caller list constrains it; a host tool allowlist intersects it. Removing the declaration revokes it. A module cleans up only the declaration it owns. Nested delegation, background execution and recursive controllers remain unavailable. +Automatic enablement admits only registered callers; it does not enable an +unregistered `tool_exec`. That legacy caller still requires the host's explicit +`programmatic_dispatch: true`, or inclusion in a host caller list when dispatch +is enabled. `model.complete(ChatRequest)` accepts a stateless tool-free request only during such an active call. The selected ordinary provider and its account are reused; diff --git a/pyproject.toml b/pyproject.toml index 6826251..1d7a171 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,7 @@ allow-direct-references = true [dependency-groups] dev = [ "amplifier-core", + "amplifier-module-context-simple @ git+https://github.com/microsoft/amplifier-module-context-simple@main", "pytest>=9.0.3", "pytest-asyncio>=1.0.0", ] @@ -45,4 +46,3 @@ amplifier-core = { git = "https://github.com/microsoft/amplifier-core", branch = testpaths = ["tests"] addopts = "--import-mode=importlib" asyncio_mode = "strict" - diff --git a/tests/test_approved_dispatch.py b/tests/test_approved_dispatch.py index bb5fb5f..d5f87c4 100644 --- a/tests/test_approved_dispatch.py +++ b/tests/test_approved_dispatch.py @@ -497,6 +497,27 @@ async def test_portable_registration_requires_no_orchestrator_configuration( assert client.lease is None and echo.calls == 1 +async def test_portable_registration_does_not_enable_an_unregistered_legacy_caller(environment): + _, coordinator, loop, _ = environment + loop.config.pop("programmatic_dispatch") + coordinator.register_capability( + "tools.dispatch.registrations", {"ui_use": {"tools": ["read"]}} + ) + dispatch = loop.approved_dispatch + with dispatch.run_scope(coordinator): + with dispatch.call_scope( + ToolCall(id="legacy", name="tool_exec", arguments={}), + coordinator.get("tools"), coordinator.hooks, None, + ): + with pytest.raises(RuntimeError, match="authorized"): + dispatch.bind() + # Preserve the existing host's explicit legacy opt-in. + loop.config["programmatic_dispatch"] = True + lease = dispatch.bind() + assert lease.tools == ("read",) + lease.close() + + async def test_nested_inference_reuses_selection_counts_and_enforces_policy( environment, ): @@ -641,3 +662,44 @@ async def program(lease): return {'revoked': True} result, _ = await call_client(environment, Client(coordinator, program)) assert result['revoked'] and echo.calls == 0 + + +@pytest.mark.parametrize("path", ["finite-turn", "parallel-tool", "serial-tool"]) +async def test_legacy_subclass_dispatcher_keeps_ownership_of_its_pipeline(environment, path): + """Existing loop-live replaces this field and supplies its own wrappers. + + That legacy interface has no model-budget method. The base engine must not + call or wrap the replacement: doing so both breaks ordinary finite turns + and can discard the legacy lease's post-approval authority check. + """ + session, coordinator, loop, echo = environment + loop.approved_dispatch = SimpleNamespace(enabled=True) + context = coordinator.get("context") + if path == "finite-turn": + class Provider: + name = "fixture" + config = {} + + def get_info(self): + return SimpleNamespace(defaults={"model": "fixture"}) + + async def complete(self, request, **kwargs): + return ChatResponse(content=[TextBlock(text="ordinary result")]) + + def parse_tool_calls(self, response): + return [] + + await coordinator.mount("providers", Provider(), name="fixture") + assert await session.execute("ordinary turn") == "ordinary result" + elif path == "parallel-tool": + await loop._execute_tool_only( + ToolCall(id="legacy", name="read", arguments={}), + coordinator.get("tools"), coordinator.hooks, None, coordinator, + ) + assert echo.calls == 1 + else: + await loop._execute_tool_with_result( + ToolCall(id="legacy", name="read", arguments={}), + coordinator.get("tools"), context, coordinator.hooks, coordinator, + ) + assert echo.calls == 1 diff --git a/uv.lock b/uv.lock index e673424..35c86b7 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 5 requires-python = ">=3.11" [[package]] @@ -14,6 +14,11 @@ dependencies = [ { name = "typing-extensions" }, ] +[[package]] +name = "amplifier-module-context-simple" +version = "1.0.0" +source = { git = "https://github.com/microsoft/amplifier-module-context-simple?rev=main#db641363013ff906b656f49498a87ab01f5c506f" } + [[package]] name = "amplifier-module-loop-streaming" version = "1.0.0" @@ -22,6 +27,7 @@ source = { editable = "." } [package.dev-dependencies] dev = [ { name = "amplifier-core" }, + { name = "amplifier-module-context-simple" }, { name = "pytest" }, { name = "pytest-asyncio" }, ] @@ -31,6 +37,7 @@ dev = [ [package.metadata.requires-dev] dev = [ { name = "amplifier-core", git = "https://github.com/microsoft/amplifier-core?branch=main" }, + { name = "amplifier-module-context-simple", git = "https://github.com/microsoft/amplifier-module-context-simple?rev=main" }, { name = "pytest", specifier = ">=9.0.3" }, { name = "pytest-asyncio", specifier = ">=1.0.0" }, ]