diff --git a/integrations/aws-strands/ARCHITECTURE.md b/integrations/aws-strands/ARCHITECTURE.md index b8b2992070..0daf322172 100644 --- a/integrations/aws-strands/ARCHITECTURE.md +++ b/integrations/aws-strands/ARCHITECTURE.md @@ -170,6 +170,7 @@ This document explains how the AWS Strands integration inside `integrations/aws- | `tool_behaviors: Dict[str, ToolBehavior]` | Per-tool overrides keyed by the Strands tool name. | | `state_context_builder` | Callable that enriches the outgoing prompt with the current shared state (useful for reiterating plan steps, recipes, etc.). | | `session_manager_provider` | Factory invoked once per thread to produce a per-thread `SessionManager`. | +| `template_tools_provider` | Per-request choice of which of the template agent's tools this request may see, applied to the live per-thread registry. TypeScript's equivalent field is `templateToolsProvider`. | | `thread_agent_kwargs` | Callable returning extra constructor kwargs for one thread's `StrandsAgentCore`. TypeScript's `threadAgentConfig` returns a partial `AgentConfig` instead. | | `emit_messages_snapshot` | Global opt-out of the four-point `MESSAGES_SNAPSHOT` emission. Default `True`. | | `replay_history_into_strands` | Global opt-out of the per-run Strands history reconciliation. Default `True`. | @@ -247,6 +248,7 @@ typescript/src/ ├── logger.ts ← injectable Logger interface + internal default ├── server.ts ← createStrandsApp factory + CORS/auth wiring ├── session-reconcile.ts ← port of session_reconcile.py, snapshot-shaped +├── template-tools.ts ← per-request filter over the template agent's tools ├── types.ts ← internal SeenToolCall bookkeeping ├── utils.ts ← content conversion + UrlFetchPolicy └── index.ts ← public exports diff --git a/integrations/aws-strands/error-codes.json b/integrations/aws-strands/error-codes.json index c29ac9ae86..90decee43f 100644 --- a/integrations/aws-strands/error-codes.json +++ b/integrations/aws-strands/error-codes.json @@ -170,6 +170,12 @@ "sides": ["python", "typescript"], "messages": ["{}"] }, + { + "code": "TEMPLATE_TOOLS_PROVIDER_ERROR", + "sides": ["python", "typescript"], + "messages": ["Failed to resolve the template tools for this request: {}"], + "note": "Reported when the per-request template-tools hook throws. Unlike the per-thread agent hook beside it, the option carries one name on both sides (template_tools_provider / templateToolsProvider) and returns the same thing, a selection of the template's tools, so there is one code and one sentence. Terminal rather than degrading to an unfiltered run: the hook exists to withhold tools from a caller, and failing open would hand the model exactly what the caller meant to keep back." + }, { "code": "THREAD_AGENT_CONFIG_ERROR", "sides": ["typescript"], diff --git a/integrations/aws-strands/python/README.md b/integrations/aws-strands/python/README.md index 21a8fbe257..3d78cab65c 100644 --- a/integrations/aws-strands/python/README.md +++ b/integrations/aws-strands/python/README.md @@ -129,6 +129,7 @@ both routes above. | `src/ag_ui_strands/a2ui_tool.py` | A2UI tool injection and the validate-and-retry recovery loop | | `src/ag_ui_strands/session_reconcile.py` | Frontend-result reconciliation against a persisted session | | `src/ag_ui_strands/client_proxy_tool.py` | Frontend tools registered into the Strands tool registry | +| `src/ag_ui_strands/template_tools.py` | Per-request filter over the template agent's own tools | | `src/ag_ui_strands/frontend_tool_interrupt.py` | The native checkpoint a waiting frontend tool parks in | | `examples/server/api/*.py` | Ready-to-run demo apps | @@ -251,6 +252,98 @@ trusted values from client-controlled `forwarded_props`; derive them from authenticated request context instead. Custom routes can pass the same state directly with `agent.run(input_data, invocation_state={...})`. +## Per-request tool filtering + +`StrandsAgentConfig.template_tools_provider` decides which of the template +agent's tools one request may see. It is called once per request with that +request's `RunAgentInput`, so the answer can vary turn by turn on a single +thread: + +```python +from ag_ui.core import RunAgentInput +from ag_ui_strands import StrandsAgent, StrandsAgentConfig + +READ_ONLY = ["search_docs", "get_order"] + +def tools_for(input_data: RunAgentInput): + # Derive the role from authenticated request context in production; + # forwarded_props is client-controlled. + if (input_data.forwarded_props or {}).get("role") == "admin": + return None # no filtering: every template tool stays available + return READ_ONLY + +agui_agent = StrandsAgent( + strands_agent, + name="assistant", + config=StrandsAgentConfig(template_tools_provider=tools_for), +) +``` + +Return the tools themselves or their names. `None` declines to filter; an empty +list is a real answer and withholds all of them. A name the template does not +contribute is dropped with a warning, because the hook narrows the wrapped +agent's tools and cannot add one. The provider may be async. + +Two boundary rules follow from that: + +- **The return value is checked, not merely iterated.** A string and a mapping + are both iterable and both mean something other than what iterating them + produces: a bare name would come apart into characters, and a permission map + would have its keys read as an allow-list while its values went unread, so a + name mapped to `False` would still be allowed. Both are refused with + `TEMPLATE_TOOLS_PROVIDER_ERROR`. Lists, tuples, sets and generators are all + accepted, and a generator that raises partway through iteration reports the + same code, because the answer is read inside the same guarded step that calls + the provider. +- **The filter reaches the registry, not only the advertised tool specs.** A + model that calls a withheld name anyway, primed by a stale turn or by the + visible history, is refused by the dispatcher rather than served. + +The filter is applied to the tool registry the thread's live Strands `Agent` +already owns, the same way client-declared tools are synchronised, and never by +rebuilding that agent. The per-thread instance holds the thread's +`SessionManager`, its native interrupt checkpoint and its history, so replacing +it to change a tool list would discard a conversation and any approval waiting +inside it. + +Three consequences follow from that: + +- **A parked call is never orphaned.** A tool in the batch a live interrupt + checkpoint would resume stays registered whatever the provider returns: the + human's answer is about to be routed back into that batch, and an absent tool + turns it into a "tool not found" the model re-fires. Filtering resumes once + the pause closes. This is the rule `sync_proxy_tools` already applies to a + proxy parked in a frontend-tool interrupt. +- **History is never rewritten.** A filtered-out tool's earlier calls and + results stay in the thread's messages, so the model can still read what it + did with a tool it can no longer call. +- **A failure is terminal.** If the provider raises, the run yields `RUN_ERROR` + with code `TEMPLATE_TOOLS_PROVIDER_ERROR` and stops, matching + `thread_agent_kwargs`. A filter that failed open would hand the model exactly + the tools the caller meant to withhold. + +The narrowing is also re-applied inside the run, once a tool batch has been +dispatched. The exemption above keeps a denied tool registered so a human's +answer can reach it, and Strands then carries on in the same run: it +re-dispatches the batch and makes its next model call from the same registry, +which would otherwise still be advertising what the request denied. The two +bridges hook different SDK events for this, because the SDKs read the tool +specs at different points relative to the events they dispatch; the effect is +the same on both. + +Scope is the template's own tools. Client-declared tools on +`RunAgentInput.tools` are re-synchronised from the request every turn already, +so a caller that wants fewer of those sends fewer. The hook is not applied on +the multi-agent orchestrator path, which has no template registry to filter. + +One deployment note. With an external per-thread agent map, a request-scoped +wrapper is rebuilt per request while the cached thread agent keeps the registry +it already had. If the template's tools are built per request too, the adapter +is handed equivalent but not identical objects, so which registry entry belongs +to the template is decided by name plus "not one of the adapter's other +producers" rather than by object identity alone. Stable tool objects are still +the simpler thing to hand it. + ## Human-in-the-loop (native Strands interrupts) Python frontend tools configured with diff --git a/integrations/aws-strands/python/src/ag_ui_strands/__init__.py b/integrations/aws-strands/python/src/ag_ui_strands/__init__.py index 8bf32ec3bf..b96b697e7b 100644 --- a/integrations/aws-strands/python/src/ag_ui_strands/__init__.py +++ b/integrations/aws-strands/python/src/ag_ui_strands/__init__.py @@ -18,6 +18,11 @@ ) from .citations import CITATIONS_METADATA_KEY from .client_proxy_tool import create_proxy_tool, sync_proxy_tools +from .template_tools import ( + EXEMPT_EVERY_TEMPLATE_TOOL, + TemplateToolsSelectionError, + sync_template_tools, +) from .utils import ( DEFAULT_URL_FETCH_POLICY, InvocationStateProvider, @@ -34,6 +39,7 @@ ToolStreamEventContext, PredictStateMapping, SessionManagerProvider, + TemplateToolsProvider, ToolStreamEventHandler, ) from ag_ui.core import ( @@ -57,6 +63,9 @@ "CITATIONS_METADATA_KEY", "create_proxy_tool", "sync_proxy_tools", + "sync_template_tools", + "TemplateToolsSelectionError", + "EXEMPT_EVERY_TEMPLATE_TOOL", "create_strands_app", "UrlFetchPolicy", "UrlFetchPolicyError", @@ -71,6 +80,7 @@ "ToolStreamEventContext", "PredictStateMapping", "SessionManagerProvider", + "TemplateToolsProvider", "ToolStreamEventHandler", "Interrupt", "ResumeEntry", diff --git a/integrations/aws-strands/python/src/ag_ui_strands/agent.py b/integrations/aws-strands/python/src/ag_ui_strands/agent.py index bd77d5c8d1..b5ce5dc532 100644 --- a/integrations/aws-strands/python/src/ag_ui_strands/agent.py +++ b/integrations/aws-strands/python/src/ag_ui_strands/agent.py @@ -1438,6 +1438,14 @@ def _error_events( sync_proxy_tools, waits_for_frontend_call, ) +from .template_tools import ( + TemplateToolsNarrowingHook, + apply_template_tool_selection, + index_template_tools, + parked_batch_tool_names, + record_template_tool_selection, + resolve_template_tool_selection, +) from .frontend_tool_interrupt import ( frontend_tool_response_schema, index_frontend_tool_interrupts, @@ -3383,6 +3391,14 @@ def __init__( if interrupt_tools: self._hooks = [StrandsInterruptHook(interrupt_tools), *self._hooks] + # Re-narrow the per-request tool filter before each model call. The + # parked-batch exemption holds a denied tool registered so a resume can + # reach it, and Strands then continues the same run against the same + # registry; without this the run would keep advertising what the + # request denied until the next request narrowed again. + if self.config.template_tools_provider is not None: + self._hooks = [*self._hooks, TemplateToolsNarrowingHook(self._tools)] + # Detect the common footgun: session_manager set on the template Agent # (stored as `_session_manager` by Strands) with no per-thread provider. # Forwarding it would make every AG-UI thread share one session_id. @@ -4486,6 +4502,65 @@ async def _run_raw( except Exception as e: logger.warning(f"Failed to set agui_context on strands_agent.state: {e}") + # Filter the tools the template contributed, per request. Applied to + # the registry this thread's live agent already owns: that instance + # carries the thread's session manager, its interrupt checkpoint and + # its history, so rebuilding it to change a tool list would discard a + # conversation and any approval waiting inside it. + if self.config.template_tools_provider is not None: + # Calling the provider and reading its answer are guarded + # together. Reading is where a mapping, a bare name or a generator + # that raises partway through is caught, and those are provider + # mistakes: leaving them outside this arm would let them bypass the + # documented code and, on the TypeScript side, end the stream with + # nothing terminal behind it. + try: + template_tool_allowed = resolve_template_tool_selection( + await maybe_await( + self.config.template_tools_provider(input_data) + ), + index_template_tools(self._tools), + ) + except Exception as e: # noqa: BLE001 - surfaced as RUN_ERROR + logger.error( + "template_tools_provider failed: %s", e, exc_info=True + ) + # Deliberately terminal rather than unfiltered: a filter that + # fails open hands the model tools the caller meant to withhold. + ev_started, ev_error = _error_events( + input_data, + "Failed to resolve the template tools for this request: " + f"{e}", + "TEMPLATE_TOOLS_PROVIDER_ERROR", + ) + yield ev_started + yield ev_error + return + # Guarded separately, and not as a provider error: past this point + # a failure is this adapter's, and it still must not escape as a + # stream that stops with nothing terminal behind it. + try: + apply_template_tool_selection( + strands_agent.tool_registry, + self._tools, + template_tool_allowed, + exempt_names=parked_batch_tool_names(strands_agent), + ) + except Exception as e: # noqa: BLE001 - surfaced as RUN_ERROR + logger.error( + "Applying the template tool filter failed: %s", e, exc_info=True + ) + ev_started, ev_error = _error_events( + input_data, str(e), _terminal_error_code(e) + ) + yield ev_started + yield ev_error + return + # Published for the re-narrowing hook: the exemption above holds a + # denied tool registered so a resume can reach it, and Strands then + # continues the same run from this registry. + record_template_tool_selection(strands_agent, template_tool_allowed) + # Sync proxy tools from client-defined tools. A proxy parked in a live # frontend-tool interrupt is exempt from removal: Strands is about to # resume that tool, and an absent registry entry turns the client's diff --git a/integrations/aws-strands/python/src/ag_ui_strands/config.py b/integrations/aws-strands/python/src/ag_ui_strands/config.py index 3099a06281..891f534c74 100644 --- a/integrations/aws-strands/python/src/ag_ui_strands/config.py +++ b/integrations/aws-strands/python/src/ag_ui_strands/config.py @@ -139,6 +139,16 @@ class ToolBehavior: """ +TemplateToolsProvider = Callable[ + ["RunAgentInput"], + Awaitable[Optional[Iterable[Any]]] | Optional[Iterable[Any]], +] +"""Chooses which of the template's tools one request may see. + +See :attr:`StrandsAgentConfig.template_tools_provider`. +""" + + @dataclass class StrandsAgentConfig: """Top-level configuration for the Strands agent adapter.""" @@ -169,6 +179,72 @@ class StrandsAgentConfig: the run yields ``RUN_ERROR`` and the thread is not cached, so the next request retries it. """ + template_tools_provider: Optional["TemplateToolsProvider"] = None + """Which of the template agent's tools this request may see. + + Called once per request with that request's ``RunAgentInput``, so the answer + can vary turn by turn on one thread: the caller's identity is in + ``forwarded_props`` or ``context``, and a tool the request must not reach is + simply left out of the returned iterable. May be async. + + Return the tools themselves or their names, whichever is to hand. Return + ``None`` to decline filtering, which leaves every template tool available; + an empty iterable is a real answer and leaves none of them. A name the + template does not contribute is dropped with a warning, because this hook + narrows the wrapped agent's tools and cannot add one. + + The container is checked rather than merely iterated. A ``str`` and a + ``Mapping`` are both refused: a bare name would come apart into characters, + and a permission map would have its keys read as an allow-list while its + values went unread, so a name mapped to ``False`` would still be allowed. + Lists, tuples, sets and generators are all accepted. + + Applied to the live per-thread agent's tool registry, never by rebuilding + that agent: the instance holds the thread's ``SessionManager``, its native + interrupt checkpoint and its history, so replacing it to change a tool list + would discard a conversation and any approval waiting inside it. + + Three consequences worth knowing: + + - A tool in the batch a live interrupt checkpoint would resume stays + registered whatever this returns. The human's answer is about to be + routed back into that batch, and an absent tool turns it into a "tool not + found" the model re-fires. This is the rule ``sync_proxy_tools`` already + applies to a proxy parked in a frontend-tool interrupt. The exemption + does not outlast what it is for: the narrowing is re-applied inside the + run once the batch has been dispatched, before the model is asked again. + - History is never rewritten. A filtered-out tool's earlier calls and + results stay in the thread's messages, so the model can still read what + it did with a tool it can no longer call, and a provider that returns + different sets across turns does not invalidate the transcript. + - If it raises, the run yields ``RUN_ERROR`` with code + ``TEMPLATE_TOOLS_PROVIDER_ERROR`` and stops, matching + ``thread_agent_kwargs``. A filter that fails open would hand the model + tools the caller meant to withhold. + + Client-declared tools on ``RunAgentInput.tools`` are outside this hook: + they are re-synchronised from the request every turn already, so a caller + that wants fewer of those sends fewer. Not applied on the multi-agent + orchestrator path, which has no template registry to filter. + + One deployment note. With an external ``agents_by_thread`` map a + request-scoped wrapper is rebuilt per request while the cached thread agent + keeps the registry it already had, so a template whose tools are built per + request hands the adapter equivalent but not identical objects. Ownership + of a registry entry therefore falls back from object identity to the tool's + name plus "not one of the adapter's other producers". Stable tool objects + are still the simpler thing to hand it. + + Example:: + + StrandsAgentConfig( + template_tools_provider=lambda input_data: ( + ["read_docs"] + if (input_data.forwarded_props or {}).get("role") != "admin" + else None + ) + ) + """ session_manager_provider: Optional[SessionManagerProvider] = None """Optional factory for creating per-thread SessionManager instances. diff --git a/integrations/aws-strands/python/src/ag_ui_strands/template_tools.py b/integrations/aws-strands/python/src/ag_ui_strands/template_tools.py new file mode 100644 index 0000000000..575449fc35 --- /dev/null +++ b/integrations/aws-strands/python/src/ag_ui_strands/template_tools.py @@ -0,0 +1,420 @@ +"""Per-request filtering of the tools the template agent contributed. + +The adapter builds one Strands ``Agent`` per thread and keeps it. That instance +is load-bearing: it holds the thread's ``SessionManager``, its native interrupt +checkpoint and its conversation history. Changing which tools a request sees +therefore has to be done to the registry the live instance already owns, the +way client-declared tools are already synchronised, and never by constructing a +replacement. + +Scope is the template's own tools. Client-declared tools arrive on +``RunAgentInput.tools`` every request and are synchronised by +:func:`~ag_ui_strands.client_proxy_tool.sync_proxy_tools`; a caller that wants +fewer of those sends fewer. Auto-injected A2UI tools are the adapter's and are +refreshed per turn. What no per-request channel reached until now is the set the +wrapped template contributed once, at construction. +""" + +from __future__ import annotations + +import logging +from typing import Any, Iterable, Mapping, Optional, Sequence, Set + +from strands.hooks import BeforeModelCallEvent, HookProvider, HookRegistry +from strands.tools.registry import ToolRegistry + +from .a2ui_tool import is_auto_injected_a2ui_tool +from .client_proxy_tool import _is_proxy + +logger = logging.getLogger(__name__) + +# The narrowed name set for the run in flight, stamped on the per-thread agent +# so the re-narrowing hook can read it back. Not routed through ``agent.state`` +# on purpose: that dict is persisted by a ``SessionManager``, and this is +# per-request scratch that must not outlive the process. +_ALLOWED_ATTR = "_ag_ui_template_tools_allowed" + +_UNSET = object() + + +class _ExemptEveryTemplateTool: + """Sentinel: hold every template tool, whatever the selection says.""" + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return "EXEMPT_EVERY_TEMPLATE_TOOL" + + +EXEMPT_EVERY_TEMPLATE_TOOL = _ExemptEveryTemplateTool() +"""What :func:`parked_batch_tool_names` answers for a checkpoint it cannot read. + +A distinct value rather than ``None``: ``None`` already means "no exemptions" +everywhere ``exempt_names`` is passed, and the two are opposites. +""" + + +class TemplateToolsSelectionError(ValueError): + """A provider answer this hook cannot read as a selection of tools. + + Raised at the boundary rather than absorbed, so the run reports + ``TEMPLATE_TOOLS_PROVIDER_ERROR`` and stops. The alternative, guessing what + an unreadable answer meant, is the one thing a hook that withholds + capability must not do. + """ + + +def index_template_tools(template_tools: Sequence[Any]) -> dict[str, Any]: + """Index the template's tools by the name a registry holds them under.""" + indexed: dict[str, Any] = {} + for tool in template_tools: + name = getattr(tool, "tool_name", None) + if isinstance(name, str) and name: + indexed[name] = tool + return indexed + + +def resolve_template_tool_selection( + selection: Optional[Iterable[Any]], + template_index: Mapping[str, Any], +) -> Optional[Set[str]]: + """Read one provider answer as the template tool names a request may see. + + Entries are either the template's own tool objects or their names, so a + caller can write the filter with whichever it has to hand. + + ``None`` means the provider declined to filter this request and every + template tool stays available. An empty iterable is a real answer and means + none of them do. + + The container is checked rather than merely iterated. A ``str`` and a + ``Mapping`` are both iterable and both mean something other than what + iterating them produces: a name would come apart into characters, and a + permission map would have its keys read as an allow-list while its values + went unread, which authorizes every key including the ones mapped to + ``False``. TypeScript's ``for...of`` already refuses a plain object, so + refusing a mapping here is also what keeps one return contract across the + two bridges rather than two. + + A name the template never contributed is dropped with a warning. This hook + narrows what the wrapped agent already gave the adapter; it cannot hand the + model a capability the template did not carry, so honouring an unknown name + is the one thing it must not do. + + Raises: + TemplateToolsSelectionError: If the answer is not a container of names + or tools. The run reports ``TEMPLATE_TOOLS_PROVIDER_ERROR``. + """ + if selection is None: + return None + + if isinstance(selection, (str, bytes, bytearray)): + raise TemplateToolsSelectionError( + "template_tools_provider returned a single " + f"{type(selection).__name__}, which iterates one character at a " + "time and would deny every tool. Return a container of the tool " + "names or tools this request may see, such as [\"a_tool\"]" + ) + if isinstance(selection, Mapping): + raise TemplateToolsSelectionError( + f"template_tools_provider returned a {type(selection).__name__}, " + "whose keys would be read as the allow-list while its values went " + "unread, so a name mapped to False would still be allowed. Return " + "a container holding only the tool names or tools this request may " + "see" + ) + try: + entries = list(selection) + except TypeError as exc: + raise TemplateToolsSelectionError( + f"template_tools_provider returned {type(selection).__name__}, " + "which is not a container of tool names or tools" + ) from exc + + allowed: Set[str] = set() + for entry in entries: + name = entry if isinstance(entry, str) else getattr(entry, "tool_name", None) + if not isinstance(name, str) or not name: + logger.warning( + "template_tools_provider returned an entry that names no tool: %r", + entry, + ) + continue + if name not in template_index: + logger.warning( + "template_tools_provider named %r, which the template agent does " + "not contribute; it stays unavailable. This hook filters the " + "template's tools and cannot add one.", + name, + ) + continue + allowed.add(name) + return allowed + + +def parked_batch_tool_names(agent: Any) -> "Set[str] | _ExemptEveryTemplateTool": + """Tool names in the batch a live interrupt checkpoint would resume. + + A parked run resumes into the tool batch it stopped inside: Strands + re-dispatches every ``toolUse`` in the assistant message it checkpointed, + answering the ones that already completed from the checkpoint and running + the one that is waiting. A tool absent from the registry at that moment + turns the human's answer into a "tool not found" the model then re-fires, + so nothing in that batch is filtered out while the pause is open. + + This is the same rule ``sync_proxy_tools`` applies through ``exempt_names`` + to a proxy parked in a frontend-tool interrupt, read off the checkpoint + instead of off the frontend-wait index because a template tool can park + through the approval hook, through an interrupt of its own, or not at all, + and the batch answers all three at once. + + Returns: + The names to hold registered; an empty set when nothing is parked; or + :data:`EXEMPT_EVERY_TEMPLATE_TOOL` for a checkpoint that is carrying a + tool batch this function cannot read, where holding everything costs + one unfiltered turn and the alternative breaks a resume. An activated + checkpoint with no ``tool_use_message`` at all is not that case: an + interrupt raised before any tool ran parks exactly that way, and it has + no batch to protect. + """ + state = getattr(agent, "_interrupt_state", None) + if state is None or getattr(state, "activated", False) is not True: + return set() + context = getattr(state, "context", None) + if not isinstance(context, Mapping): + logger.warning( + "An activated interrupt checkpoint carries no readable context; " + "holding every template tool registered rather than risk removing " + "one this thread's resume is about to re-dispatch." + ) + return EXEMPT_EVERY_TEMPLATE_TOOL + if "tool_use_message" not in context: + # A pause raised before any tool ran. Nothing is mid-dispatch, so + # nothing needs holding. + return set() + + message = context["tool_use_message"] + names: Set[str] = set() + if isinstance(message, Mapping): + for block in message.get("content") or []: + tool_use = block.get("toolUse") if isinstance(block, Mapping) else None + if not isinstance(tool_use, Mapping): + continue + name = tool_use.get("name") + if isinstance(name, str) and name: + names.add(name) + if not names: + logger.warning( + "An activated interrupt checkpoint carries a tool batch this " + "adapter cannot read (%s); holding every template tool registered " + "rather than risk removing one this thread's resume is about to " + "re-dispatch.", + type(message).__name__, + ) + return EXEMPT_EVERY_TEMPLATE_TOOL + return names + + +def _is_foreign_entry(entry: Any) -> bool: + """Whether a registry entry belongs to a producer other than the template. + + The adapter has exactly two others: proxies for client-declared tools, and + the A2UI tool it injects itself. Both are re-decided every turn by the code + that owns them, so neither is this filter's to touch. + """ + return _is_proxy(entry) or is_auto_injected_a2ui_tool(entry) + + +def _is_template_entry(entry: Any, template_tool: Any) -> bool: + """Whether a registry entry under a template tool's name is the template's. + + Identity settles it when it holds. It does not always hold: with an + external ``agents_by_thread`` map the wrapper is rebuilt per request while + the cached thread agent keeps the registry it already had, so a template + whose tools are built per request (a factory, or a closure over a + request-scoped handle) hands the adapter equivalent but not identical + objects. Reading a non-match as "someone else owns this name" would make a + deny-everything answer remove nothing, which is a silent failure in the + permissive direction on a hook whose whole job is withholding capability. + + So the fallback is ownership by elimination: the name is one the template + contributes, and the entry sitting on it is not one of the adapter's other + producers, therefore it is the template's. + """ + if entry is template_tool: + return True + return not _is_foreign_entry(entry) + + +def apply_template_tool_selection( + tool_registry: ToolRegistry, + template_tools: Sequence[Any], + allowed: Optional[Set[str]], + *, + exempt_names: "Set[str] | _ExemptEveryTemplateTool | None" = None, +) -> Set[str]: + """Make *tool_registry* hold exactly the template tools *allowed* permits. + + ``allowed`` is a resolved name set, ``None`` meaning no filtering. + ``exempt_names`` holds names to keep registered whatever ``allowed`` says; + ``None`` exempts nothing and :data:`EXEMPT_EVERY_TEMPLATE_TOOL` exempts + every name. + + Removal is not destructive. The template tool objects outlive the registry + entry, so a later request that allows a name again restores the same + instance, and history stays untouched throughout: a filtered-out tool's + earlier calls and results remain in the thread's messages, which is what + lets the model read what it already did with a tool it can no longer call. + + Returns the template tool names the registry holds after the call. + """ + template_index = index_template_tools(template_tools) + exempt_all = exempt_names is EXEMPT_EVERY_TEMPLATE_TOOL + exempt_set: Set[str] = set() if exempt_all or exempt_names is None else exempt_names # type: ignore[assignment] + + registered: Set[str] = set() + for name, tool in template_index.items(): + exempt = exempt_all or name in exempt_set + keep = allowed is None or name in allowed or exempt + existing = tool_registry.registry.get(name) + + if not keep: + if existing is not None and _is_template_entry(existing, tool): + del tool_registry.registry[name] + tool_registry.dynamic_tools.pop(name, None) + logger.debug("Filtered out template tool: %s", name) + elif existing is not None: + logger.debug( + "Template tool %s is held by another producer; the filter " + "leaves it in place", + name, + ) + continue + + if existing is tool: + registered.add(name) + continue + if existing is None: + _restore(tool_registry, name, tool) + registered.add(name) + continue + if not exempt and _is_proxy(existing): + # A client proxy took this name while the template tool was + # filtered out, and the provider now allows the template tool. A + # native tool wins a name collision, so hand the name back: the + # proxy sync runs after this and re-decides the client's side, + # skipping a name a native tool holds. Leaving the proxy would both + # shadow the allowed tool and, if the client has stopped declaring + # it, let the proxy sync drop the name outright. + # + # Guarded on ``not exempt`` rather than on the selection, because + # the keep branch is also reached by exemption, and a proxy the + # parked batch is answering keeps its name. + del tool_registry.registry[name] + tool_registry.dynamic_tools.pop(name, None) + _restore(tool_registry, name, tool) + registered.add(name) + logger.debug( + "Reclaimed template tool %s from a client proxy holding its name", + name, + ) + continue + if _is_template_entry(existing, tool): + # The template's, under a different object. Leave the entry the + # thread has been using rather than churn it. + registered.add(name) + continue + logger.debug( + "Template tool %s is held by another producer; leaving it in place", + name, + ) + return registered + + +def _restore(tool_registry: ToolRegistry, name: str, tool: Any) -> None: + """Put a template tool back under its name. + + By assignment rather than through ``register_tool``, which raises over a + name that normalizes onto an existing one. Nothing new is being registered + here: the entry is the one the template already put in this registry, so + validating it again could only fail a run over a collision the construction + it came from had accepted. + """ + tool_registry.registry[name] = tool + if getattr(tool, "is_dynamic", False): + tool_registry.dynamic_tools[name] = tool + + +def sync_template_tools( + tool_registry: ToolRegistry, + template_tools: Sequence[Any], + selection: Optional[Iterable[Any]], + *, + exempt_names: "Set[str] | _ExemptEveryTemplateTool | None" = None, +) -> Set[str]: + """Read *selection* and apply it to *tool_registry* in one call. + + The run path keeps the two halves apart, so that reading a provider's + answer fails as a provider error and applying it does not. This composes + them for a caller with an answer already in hand. + """ + return apply_template_tool_selection( + tool_registry, + template_tools, + resolve_template_tool_selection( + selection, index_template_tools(template_tools) + ), + exempt_names=exempt_names, + ) + + +def record_template_tool_selection(agent: Any, allowed: Optional[Set[str]]) -> None: + """Publish the run's narrowed name set for the re-narrowing hook to read.""" + setattr(agent, _ALLOWED_ATTR, allowed) + + +class TemplateToolsNarrowingHook(HookProvider): + """Re-narrow the filtered set once a resumed batch has been dispatched. + + The parked-batch exemption keeps a denied tool registered so a resume can + reach it. Strands then carries on inside the same run: it re-dispatches the + batch, clears the checkpoint, and makes its next model call from this same + registry, which would still be advertising what the request denied. Without + this hook the model could call a withheld tool for the rest of that run, + and only the next request would narrow again. + + Strands reads the registry fresh before every model call and announces that + read through ``BeforeModelCallEvent``, so that is where the narrowing is + re-applied. + + No exemption is passed. By the time this fires the parked batch has been + dispatched, which is the whole reason the exemption existed, so re-reading + the checkpoint would only ask a question whose answer no longer matters, + and the answer moved between SDK releases: the TypeScript SDK clears its + pending execution at different points in 1.1 and 1.16, so a hook that + reads it holds the exemption on one release and drops it on the other. + Asking nothing is both simpler and the same on every release. + + A run the SDK is cancelling can replay a skipped batch after this fires, + and a tool this narrowing removed would be missing from that replay. The + run is being torn down at that point, and the next request restores + whatever a live checkpoint still needs before anything reads the registry + again. + """ + + def __init__(self, template_tools: Sequence[Any]) -> None: + self._template_tools = template_tools + + def register_hooks(self, registry: HookRegistry, **_kwargs: Any) -> None: + registry.add_callback(BeforeModelCallEvent, self._narrow) + + def _narrow(self, event: Any) -> None: + agent = getattr(event, "agent", None) + if agent is None: + return + allowed = getattr(agent, _ALLOWED_ATTR, _UNSET) + if allowed is _UNSET or allowed is None: + return + apply_template_tool_selection( + agent.tool_registry, + self._template_tools, + allowed, + ) diff --git a/integrations/aws-strands/python/tests/test_template_tool_filtering.py b/integrations/aws-strands/python/tests/test_template_tool_filtering.py new file mode 100644 index 0000000000..519655b95d --- /dev/null +++ b/integrations/aws-strands/python/tests/test_template_tool_filtering.py @@ -0,0 +1,1147 @@ +"""Per-request filtering of the tools the template agent contributed. + +The adapter keeps one Strands ``Agent`` per thread and that instance is +load-bearing: it holds the thread's ``SessionManager``, its native interrupt +checkpoint and its conversation history. So the property these tests exist to +pin is not only that a filter takes effect, but that it takes effect on the +registry the live instance already owns. A filter applied by rebuilding the +thread's agent would pass a "the model saw fewer tools" assertion while +silently discarding a conversation and any approval waiting inside it, which is +why the identity of the cached agent is asserted alongside the tool specs. +""" + +from __future__ import annotations + +import logging + +import pytest +from ag_ui.core import EventType, RunAgentInput, Tool, UserMessage +from strands import Agent as StrandsAgentCore +from strands import tool +from strands.hooks import BeforeModelCallEvent, HookProvider +from strands.interrupt import Interrupt as StrandsInterrupt +from strands.models.model import Model as StrandsModel +from strands.tools.registry import ToolRegistry + +from ag_ui_strands.agent import StrandsAgent +from ag_ui_strands.config import StrandsAgentConfig, ToolBehavior +from ag_ui_strands.client_proxy_tool import _is_proxy, create_proxy_tool +from ag_ui_strands.template_tools import ( + EXEMPT_EVERY_TEMPLATE_TOOL, + TemplateToolsSelectionError, + index_template_tools, + parked_batch_tool_names, + resolve_template_tool_selection, + sync_template_tools, +) +from tests.interrupt_state_stub import InterruptStateStub + + +THREAD_ID = "template-filter-thread" + + +@tool +def read_docs(topic: str) -> str: + """Read the documentation for a topic.""" + return f"docs about {topic}" + + +@tool +def delete_record(record_id: str) -> str: + """Delete a record.""" + return f"deleted {record_id}" + + +class RecordingModel(StrandsModel): + """Answers with text and records the tool specs it was offered.""" + + def __init__(self) -> None: + self.offered_tool_names: list[set[str]] = [] + + def get_config(self): + return {} + + def update_config(self, **kwargs): + pass + + async def structured_output( + self, output_model, prompt=None, system_prompt=None, **kwargs + ): + raise NotImplementedError + yield # pragma: no cover + + async def stream(self, messages, tool_specs=None, system_prompt=None, **kwargs): + self.offered_tool_names.append( + {spec["name"] for spec in (tool_specs or [])} + ) + yield {"messageStart": {"role": "assistant"}} + yield {"contentBlockDelta": {"delta": {"text": "ok"}}} + yield {"contentBlockStop": {}} + yield {"messageStop": {"stopReason": "end_turn"}} + + +class DeleteThenAnswerModel(StrandsModel): + """Calls ``delete_record`` once, then answers with text.""" + + def __init__(self) -> None: + self.offered_tool_names: list[set[str]] = [] + self.called_delete = False + + def get_config(self): + return {} + + def update_config(self, **kwargs): + pass + + async def structured_output( + self, output_model, prompt=None, system_prompt=None, **kwargs + ): + raise NotImplementedError + yield # pragma: no cover + + async def stream(self, messages, tool_specs=None, system_prompt=None, **kwargs): + self.offered_tool_names.append( + {spec["name"] for spec in (tool_specs or [])} + ) + yield {"messageStart": {"role": "assistant"}} + if not self.called_delete: + self.called_delete = True + yield { + "contentBlockStart": { + "start": { + "toolUse": { + "toolUseId": "delete-1", + "name": "delete_record", + } + } + } + } + yield { + "contentBlockDelta": { + "delta": {"toolUse": {"input": '{"record_id": "r-1"}'}} + } + } + yield {"contentBlockStop": {}} + yield {"messageStop": {"stopReason": "tool_use"}} + return + yield {"contentBlockDelta": {"delta": {"text": "done"}}} + yield {"contentBlockStop": {}} + yield {"messageStop": {"stopReason": "end_turn"}} + + +def make_agent( + model: StrandsModel, + *, + config: StrandsAgentConfig | None = None, + tools=None, +) -> StrandsAgent: + core = StrandsAgentCore( + model=model, + tools=[read_docs, delete_record] if tools is None else tools, + system_prompt="Help the user.", + ) + return StrandsAgent( + core, + name="template-filter-test", + config=config or StrandsAgentConfig(), + ) + + +def run_input( + run_id: str, + *, + thread_id: str = THREAD_ID, + tools: list[Tool] | None = None, + forwarded_props: dict | None = None, + resume=None, +) -> RunAgentInput: + return RunAgentInput( + thread_id=thread_id, + run_id=run_id, + state={}, + messages=[ + UserMessage(id=f"user-{run_id}", role="user", content="Do the thing.") + ], + tools=tools or [], + context=[], + forwarded_props=forwarded_props or {}, + **({"resume": resume} if resume is not None else {}), + ) + + +async def drain(agent: StrandsAgent, request: RunAgentInput) -> list: + return [event async for event in agent.run(request)] + + +def registry_names(agent: StrandsAgent, thread_id: str = THREAD_ID) -> set[str]: + return set(agent._agents_by_thread[thread_id].tool_registry.registry) + + +# --------------------------------------------------------------------------- +# The rescope: filtering happens on the live agent, not by replacing it +# --------------------------------------------------------------------------- + + +class TestFilteringWithoutRebuildingTheThreadAgent: + async def test_the_filtered_set_varies_between_two_requests_on_one_thread(self): + allowed_by_run = {"r1": ["read_docs"], "r2": ["read_docs", "delete_record"]} + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: allowed_by_run[data.run_id] + ), + ) + + await drain(agent, run_input("r1")) + assert model.offered_tool_names[-1] == {"read_docs"} + assert registry_names(agent) == {"read_docs"} + + await drain(agent, run_input("r2")) + assert model.offered_tool_names[-1] == {"read_docs", "delete_record"} + assert registry_names(agent) == {"read_docs", "delete_record"} + + async def test_the_cached_thread_agent_is_the_same_instance_across_the_change( + self, + ): + """The whole point of applying the filter to the registry. + + Recreating the per-thread agent whenever the resolved tool set changed + would satisfy the assertion above and still be wrong, so identity is + asserted directly. + """ + allowed_by_run = {"r1": ["read_docs"], "r2": []} + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: allowed_by_run[data.run_id] + ), + ) + + await drain(agent, run_input("r1")) + first = agent._agents_by_thread[THREAD_ID] + + await drain(agent, run_input("r2")) + assert agent._agents_by_thread[THREAD_ID] is first + assert model.offered_tool_names[-1] == set() + + async def test_an_empty_selection_withholds_every_template_tool(self): + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig(template_tools_provider=lambda data: []), + ) + await drain(agent, run_input("r1")) + assert model.offered_tool_names[-1] == set() + + async def test_returning_none_declines_to_filter_that_request(self): + selection_by_run = {"r1": [], "r2": None} + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: selection_by_run[data.run_id] + ), + ) + + await drain(agent, run_input("r1")) + assert model.offered_tool_names[-1] == set() + + await drain(agent, run_input("r2")) + assert model.offered_tool_names[-1] == {"read_docs", "delete_record"} + + async def test_the_provider_may_be_async_and_may_read_the_caller_identity(self): + async def provider(data: RunAgentInput): + if (data.forwarded_props or {}).get("role") == "admin": + return None + return [read_docs] + + model = RecordingModel() + agent = make_agent( + model, config=StrandsAgentConfig(template_tools_provider=provider) + ) + + await drain(agent, run_input("r1", forwarded_props={"role": "reader"})) + assert model.offered_tool_names[-1] == {"read_docs"} + + await drain(agent, run_input("r2", forwarded_props={"role": "admin"})) + assert model.offered_tool_names[-1] == {"read_docs", "delete_record"} + + async def test_two_threads_can_see_different_tools_at_the_same_time(self): + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: ( + None if data.thread_id == "wide" else ["read_docs"] + ) + ), + ) + await drain(agent, run_input("r1", thread_id="wide")) + assert model.offered_tool_names[-1] == {"read_docs", "delete_record"} + await drain(agent, run_input("r2", thread_id="narrow")) + assert model.offered_tool_names[-1] == {"read_docs"} + assert registry_names(agent, "wide") == {"read_docs", "delete_record"} + assert registry_names(agent, "narrow") == {"read_docs"} + + +class TestWhatTheThreadKeepsAcrossAFilterChange: + async def test_the_session_manager_and_the_persisted_history_both_survive( + self, tmp_path + ): + from strands.session.file_session_manager import FileSessionManager + + session_manager = FileSessionManager( + session_id="template-filter-session", storage_dir=str(tmp_path) + ) + allowed_by_run = {"r1": None, "r2": ["read_docs"]} + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + session_manager_provider=lambda data: session_manager, + template_tools_provider=lambda data: allowed_by_run[data.run_id], + ), + ) + + await drain(agent, run_input("r1")) + core = agent._agents_by_thread[THREAD_ID] + turns_after_first_run = len(core.messages) + assert turns_after_first_run > 0 + + await drain(agent, run_input("r2")) + assert agent._agents_by_thread[THREAD_ID] is core + assert core._session_manager is session_manager + assert len(core.messages) > turns_after_first_run + assert model.offered_tool_names[-1] == {"read_docs"} + + async def test_a_filtered_out_tool_keeps_the_calls_it_already_made( + self, tmp_path + ): + """History is not rewritten, so the model still reads what it did. + + The filter answers "what may this request call", not "what happened on + this thread". Removing the record would leave an assistant tool-use + block with no result behind it. + + Backed by a real session manager on purpose: with none configured the + thread's history is reconciled against ``RunAgentInput.messages`` every + turn, so what a filter did to it could not be told apart from what the + replay did. + """ + from strands.session.file_session_manager import FileSessionManager + + session_manager = FileSessionManager( + session_id="template-filter-history", storage_dir=str(tmp_path) + ) + allowed_by_run = {"r1": None, "r2": ["read_docs"]} + model = DeleteThenAnswerModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + session_manager_provider=lambda data: session_manager, + template_tools_provider=lambda data: allowed_by_run[data.run_id], + ), + ) + + def delete_calls(messages): + found = [] + for message in messages: + for block in message.get("content") or []: + tool_use = block.get("toolUse") if isinstance(block, dict) else None + if isinstance(tool_use, dict) and tool_use.get("name") == "delete_record": + found.append(tool_use["toolUseId"]) + return found + + await drain(agent, run_input("r1")) + core = agent._agents_by_thread[THREAD_ID] + recorded = delete_calls(core.messages) + assert recorded, "the first run never called delete_record" + + await drain(agent, run_input("r2")) + assert agent._agents_by_thread[THREAD_ID] is core + assert delete_calls(core.messages) == recorded, ( + "filtering delete_record out erased the call it already made" + ) + assert model.offered_tool_names[-1] == {"read_docs"} + + +class TestAParkedCallIsNotOrphaned: + async def test_a_tool_awaiting_approval_stays_registered_while_filtered_out(self): + """The rule ``sync_proxy_tools`` already applies, read off the checkpoint. + + The human's answer is routed back into the tool batch the run stopped + inside. A tool absent from the registry at that moment turns the answer + into a "tool not found" the model then re-fires, so the batch is exempt + until the pause closes. + """ + model = DeleteThenAnswerModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + tool_behaviors={"delete_record": ToolBehavior(interrupt_on_call=True)}, + template_tools_provider=lambda data: ( + None if data.run_id == "r1" else ["read_docs"] + ), + ), + ) + + await drain(agent, run_input("r1")) + core = agent._agents_by_thread[THREAD_ID] + assert core._interrupt_state.activated, "the first run did not park" + parked = list(agent._pending_interrupts_by_thread[THREAD_ID]) + assert parked, "no AG-UI interrupt was recorded for the pause" + + from ag_ui_strands import ResumeEntry + + events = await drain( + agent, + run_input( + "r2", + resume=[ + ResumeEntry( + interrupt_id=parked[0], + status="resolved", + payload={"approved": True}, + ) + ], + ), + ) + + assert agent._agents_by_thread[THREAD_ID] is core + # The resume reaching the tool is the assertion. Dropping the exemption + # makes Strands report the tool as absent from the registry, which + # surfaces here as the approved tool never running. + assert [e for e in events if e.type == EventType.RUN_ERROR] == [] + assert [e for e in events if e.type == EventType.RUN_FINISHED] + assert model.called_delete + + async def test_a_plain_turn_against_a_pause_is_refused_before_the_filter_runs( + self, + ): + """The pause is not reachable by a filter at all on a non-resume turn. + + A turn that submits no answer is refused ahead of the tool sync, so a + provider that would have withheld the parked tool never runs and the + checkpoint and its AG-UI bookkeeping are still there for the resume + that follows. + """ + from ag_ui_strands import ResumeEntry + from tests.error_code_table import assert_contract_error + + consulted: list[str] = [] + + def provider(data: RunAgentInput): + consulted.append(data.run_id) + return None if data.run_id == "r1" else ["read_docs"] + + model = DeleteThenAnswerModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + tool_behaviors={"delete_record": ToolBehavior(interrupt_on_call=True)}, + template_tools_provider=provider, + ), + ) + + await drain(agent, run_input("r1")) + core = agent._agents_by_thread[THREAD_ID] + parked = dict(agent._pending_interrupts_by_thread[THREAD_ID]) + assert parked + + refused = await drain(agent, run_input("r2")) + assert_contract_error( + next(e for e in refused if e.type == EventType.RUN_ERROR), + "PENDING_INTERRUPTS", + ) + assert consulted == ["r1"], "the provider ran on a turn that was refused" + assert core._interrupt_state.activated + assert agent._pending_interrupts_by_thread[THREAD_ID] == parked + assert "delete_record" in core.tool_registry.registry + + events = await drain( + agent, + run_input( + "r3", + resume=[ + ResumeEntry( + interrupt_id=next(iter(parked)), + status="resolved", + payload={"approved": True}, + ) + ], + ), + ) + assert [e for e in events if e.type == EventType.RUN_ERROR] == [] + assert agent._agents_by_thread[THREAD_ID] is core + assert consulted == ["r1", "r3"] + + async def test_a_parked_frontend_call_survives_a_filter_that_allows_nothing(self): + """A template filter has no reach over client-declared tools. + + Client tools are re-synchronised from ``RunAgentInput.tools`` every + request and their proxies are a different producer's entries, so the + filter must not be able to remove one, parked or not. + """ + client_tool = Tool( + name="confirm_in_client", + description="Confirm in the client", + parameters={"type": "object", "properties": {}}, + ) + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + tool_behaviors={ + "confirm_in_client": ToolBehavior( + continue_after_frontend_call=False + ) + }, + template_tools_provider=lambda data: [], + ), + ) + + await drain(agent, run_input("r1", tools=[client_tool])) + core = agent._agents_by_thread[THREAD_ID] + assert "confirm_in_client" in core.tool_registry.registry + assert set(core.tool_registry.registry) == {"confirm_in_client"} + + await drain(agent, run_input("r2", tools=[client_tool])) + assert agent._agents_by_thread[THREAD_ID] is core + assert "confirm_in_client" in core.tool_registry.registry + + +class TestFilteringRemovesTheCapability: + async def test_a_filtered_out_tool_the_model_calls_anyway_does_not_run(self): + """The point of touching the registry rather than only the tool specs. + + Withholding a tool from the specs and leaving it registered would make + the filter advice a model can ignore. A model calling the name anyway, + because it was primed by a stale turn or by the visible history, has to + be refused by the dispatcher rather than served. + """ + executions: list[str] = [] + + @tool(name="delete_record") + def audited_delete(record_id: str) -> str: + """Delete a record, recording that it ran.""" + executions.append(record_id) + return f"deleted {record_id}" + + model = DeleteThenAnswerModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: ["read_docs"] + ), + tools=[read_docs, audited_delete], + ) + + events = await drain(agent, run_input("r1")) + + assert executions == [], "a filtered-out tool executed" + assert model.offered_tool_names[0] == {"read_docs"} + assert [e for e in events if e.type == EventType.RUN_ERROR] == [] + + +class TestTheReturnContract: + async def test_a_mapping_is_refused_rather_than_read_as_an_allow_list(self): + """The one failure that was silent and permissive. + + A permission map is a natural thing to reach for on a hook shaped like + this, and iterating it yields its keys: every name would be allowed, + including the ones mapped to False, and nothing would say so. + """ + from tests.error_code_table import assert_contract_error + + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: { + "read_docs": True, + "delete_record": False, + } + ), + ) + + events = await drain(agent, run_input("r1")) + + error = next(e for e in events if e.type == EventType.RUN_ERROR) + assert_contract_error(error, "TEMPLATE_TOOLS_PROVIDER_ERROR") + assert "values went unread" in error.message + assert model.offered_tool_names == [], "the model ran unfiltered" + + async def test_a_bare_name_is_refused_rather_than_read_one_character_at_a_time( + self, + ): + from tests.error_code_table import assert_contract_error + + agent = make_agent( + RecordingModel(), + config=StrandsAgentConfig( + template_tools_provider=lambda data: "read_docs" + ), + ) + error = next( + e + for e in await drain(agent, run_input("r1")) + if e.type == EventType.RUN_ERROR + ) + assert_contract_error(error, "TEMPLATE_TOOLS_PROVIDER_ERROR") + assert "one character at a time" in error.message + + async def test_a_non_container_is_refused(self): + from tests.error_code_table import assert_contract_error + + agent = make_agent( + RecordingModel(), + config=StrandsAgentConfig(template_tools_provider=lambda data: 42), + ) + assert_contract_error( + next( + e + for e in await drain(agent, run_input("r1")) + if e.type == EventType.RUN_ERROR + ), + "TEMPLATE_TOOLS_PROVIDER_ERROR", + ) + + async def test_a_generator_that_raises_partway_reports_the_provider_error(self): + """The provider's answer is read inside the guarded arm, not after it. + + A generator constructs without running its body, so a provider can hand + back something that fails only on first iteration. Reading it outside + the boundary bypassed the documented code entirely. + """ + from tests.error_code_table import assert_contract_error + + def provider(data: RunAgentInput): + def entries(): + yield "read_docs" + raise RuntimeError("directory lookup failed") + + return entries() + + model = RecordingModel() + agent = make_agent( + model, config=StrandsAgentConfig(template_tools_provider=provider) + ) + + events = await drain(agent, run_input("r1")) + error = next(e for e in events if e.type == EventType.RUN_ERROR) + assert_contract_error(error, "TEMPLATE_TOOLS_PROVIDER_ERROR") + assert "directory lookup failed" in error.message + assert model.offered_tool_names == [] + + async def test_a_generator_of_names_is_accepted(self): + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: (n for n in ["read_docs"]) + ), + ) + await drain(agent, run_input("r1")) + assert model.offered_tool_names[-1] == {"read_docs"} + + +class TestTheExemptionDoesNotOutliveTheCheckpoint: + async def test_the_denied_batch_is_narrowed_again_before_the_next_model_call( + self, + ): + """Strands keeps running after a resume, from this same registry. + + The exemption holds a denied tool registered so the human's answer can + reach it. Strands then re-dispatches the batch, clears the checkpoint + and calls the model again within the same run. Without a re-narrowing + that call would still advertise the tool the request denied. + """ + model = DeleteThenAnswerModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + tool_behaviors={"delete_record": ToolBehavior(interrupt_on_call=True)}, + template_tools_provider=lambda data: ( + None if data.run_id == "r1" else ["read_docs"] + ), + ), + ) + + await drain(agent, run_input("r1")) + parked = list(agent._pending_interrupts_by_thread[THREAD_ID]) + offered_before_resume = len(model.offered_tool_names) + + from ag_ui_strands import ResumeEntry + + events = await drain( + agent, + run_input( + "r2", + resume=[ + ResumeEntry( + interrupt_id=parked[0], + status="resolved", + payload={"approved": True}, + ) + ], + ), + ) + + assert [e for e in events if e.type == EventType.RUN_ERROR] == [] + assert len(model.offered_tool_names) > offered_before_resume, ( + "the resume never reached another model call, so this asserts nothing" + ) + for offered in model.offered_tool_names[offered_before_resume:]: + assert offered == {"read_docs"}, ( + "a model call after the resume still advertised the denied tool" + ) + + async def test_an_unreadable_parked_batch_holds_every_template_tool(self): + """The conservative direction, for a shape this adapter cannot read. + + An activated checkpoint whose tool batch does not decode is an SDK + shape this code does not know. Filtering anyway risks breaking a + resume a human is waiting on; holding everything costs one unfiltered + turn. + """ + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: ["read_docs"] + ), + ) + await drain(agent, run_input("r1")) + core = agent._agents_by_thread[THREAD_ID] + assert set(core.tool_registry.registry) == {"read_docs"} + + state = InterruptStateStub( + interrupts={"i1": StrandsInterrupt("i1", "generic")} + ) + state.activate({"tool_use_message": "not a message"}) + core._interrupt_state = state + + assert parked_batch_tool_names(core) is EXEMPT_EVERY_TEMPLATE_TOOL + sync_template_tools( + core.tool_registry, + agent._tools, + ["read_docs"], + exempt_names=parked_batch_tool_names(core), + ) + assert set(core.tool_registry.registry) == {"read_docs", "delete_record"} + + async def test_a_pause_with_no_tool_batch_holds_nothing(self): + """An interrupt raised before any tool ran has no batch to protect.""" + state = InterruptStateStub() + state.activate({"responses": []}) + + class _Agent: + _interrupt_state = state + + assert parked_batch_tool_names(_Agent()) == set() + + +class TestTheOrderingAgainstTheProxySync: + async def test_an_allowed_tool_survives_a_client_dropping_a_colliding_name( + self, + ): + """Neither producer ended up holding the name, for exactly one request. + + A client tool sharing a template tool's name takes the registry slot + while the template tool is filtered out. When the provider allows the + template tool again and the client has stopped declaring its own, the + template sync used to decline to touch the proxy and the proxy sync + then removed it as stale, leaving the allowed tool registered nowhere. + """ + client_tool = Tool( + name="delete_record", + description="A client tool of the same name", + parameters={"type": "object", "properties": {}}, + ) + allowed_by_run = {"r1": ["read_docs"], "r2": None} + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: allowed_by_run[data.run_id] + ), + ) + + await drain(agent, run_input("r1", tools=[client_tool])) + core = agent._agents_by_thread[THREAD_ID] + + await drain(agent, run_input("r2")) + assert agent._agents_by_thread[THREAD_ID] is core + assert "delete_record" in core.tool_registry.registry, ( + "the provider allowed the tool and it is registered nowhere" + ) + assert model.offered_tool_names[-1] == {"read_docs", "delete_record"} + + async def test_a_client_tool_cannot_shadow_an_allowed_template_tool(self): + """Native tools win a name collision, filter or no filter.""" + client_tool = Tool( + name="delete_record", + description="A client tool of the same name", + parameters={"type": "object", "properties": {}}, + ) + allowed_by_run = {"r1": ["read_docs"], "r2": None} + model = RecordingModel() + agent = make_agent( + model, + config=StrandsAgentConfig( + template_tools_provider=lambda data: allowed_by_run[data.run_id] + ), + ) + + await drain(agent, run_input("r1", tools=[client_tool])) + core = agent._agents_by_thread[THREAD_ID] + assert _is_proxy(core.tool_registry.registry["delete_record"]) + + await drain(agent, run_input("r2", tools=[client_tool])) + assert not _is_proxy(core.tool_registry.registry["delete_record"]) + + +class TestOwnershipIsNotPureIdentity: + async def test_a_rebuilt_wrapper_can_still_deny_a_cached_threads_tools(self): + """The shape ``agents_by_thread`` exists for. + + A request-scoped wrapper is rebuilt per request while the cached thread + agent keeps the registry it already had. A template whose tools are + built per request then hands each new wrapper equivalent but not + identical objects, and ownership by object identity alone would read + every one of them as another producer's entry, so a deny-everything + answer would remove nothing at all. + """ + agents_by_thread: dict = {} + allowed_by_run = {"r1": None, "r2": []} + # Registered ahead of the adapter's own re-narrowing hook, which is + # appended last, so this snapshot is the registry the request path left + # behind rather than the one the hook went on to correct. + seen_before_model_call: list[set[str]] = [] + + class _Snapshot(HookProvider): + def register_hooks(self, registry, **_kwargs): + registry.add_callback(BeforeModelCallEvent, self._record) + + def _record(self, event): + seen_before_model_call.append( + set(event.agent.tool_registry.registry) + ) + + def build_agent() -> StrandsAgent: + # Rebuilt per request, tools included: the case that breaks pure + # identity. Module-level tools would be the same objects every time + # and the bug would be invisible. + @tool(name="read_docs") + def read(topic: str) -> str: + """Read the documentation for a topic.""" + return topic + + @tool(name="delete_record") + def delete(record_id: str) -> str: + """Delete a record.""" + return record_id + + core = StrandsAgentCore( + model=RecordingModel(), + tools=[read, delete], + system_prompt="Help the user.", + ) + return StrandsAgent( + core, + name="rebuilt-wrapper", + agents_by_thread=agents_by_thread, + hooks=[_Snapshot()], + config=StrandsAgentConfig( + template_tools_provider=lambda data: allowed_by_run[data.run_id] + ), + ) + + first = build_agent() + await drain(first, run_input("r1")) + core = agents_by_thread[THREAD_ID] + assert set(core.tool_registry.registry) == {"read_docs", "delete_record"} + + second = build_agent() + assert all( + core.tool_registry.registry[t.tool_name] is not t for t in second._tools + ), "the rebuilt template reused its tool objects, so this asserts nothing" + + await drain(second, run_input("r2")) + assert agents_by_thread[THREAD_ID] is core + assert seen_before_model_call[-1] == set(), ( + "a deny-everything answer removed nothing from the cached thread" + ) + assert set(core.tool_registry.registry) == set() + + +class TestTheProviderFailureMode: + async def test_a_raising_provider_ends_the_run_rather_than_running_unfiltered( + self, caplog + ): + """Terminal, matching ``thread_agent_kwargs``. + + Degrading to an unfiltered run would hand the model exactly the tools + the caller meant to withhold, which is the one outcome this hook exists + to prevent. + """ + from tests.error_code_table import assert_contract_error + + def provider(data: RunAgentInput): + raise RuntimeError("authz lookup failed") + + model = RecordingModel() + agent = make_agent( + model, config=StrandsAgentConfig(template_tools_provider=provider) + ) + + with caplog.at_level(logging.ERROR, logger="ag_ui_strands.agent"): + events = await drain(agent, run_input("r1")) + + assert [e.type for e in events] == [ + EventType.RUN_STARTED, + EventType.RUN_ERROR, + ] + assert_contract_error(events[1], "TEMPLATE_TOOLS_PROVIDER_ERROR") + assert "authz lookup failed" in events[1].message + assert model.offered_tool_names == [], "the model ran despite the failure" + assert "template_tools_provider failed" in caplog.text + + async def test_the_run_error_is_bracketed_by_a_run_started(self): + """A client that brackets on the lifecycle events needs the opener.""" + agent = make_agent( + RecordingModel(), + config=StrandsAgentConfig( + template_tools_provider=lambda data: (_ for _ in ()).throw( + ValueError("boom") + ) + ), + ) + events = await drain(agent, run_input("r1")) + assert events[0].type == EventType.RUN_STARTED + assert events[0].run_id == "r1" + + +class TestNoProviderConfigured: + async def test_the_registry_is_untouched_when_no_provider_is_configured(self): + model = RecordingModel() + agent = make_agent(model) + await drain(agent, run_input("r1")) + await drain(agent, run_input("r2")) + assert registry_names(agent) == {"read_docs", "delete_record"} + assert model.offered_tool_names == [ + {"read_docs", "delete_record"}, + {"read_docs", "delete_record"}, + ] + + async def test_no_provider_never_reaches_the_sync(self, monkeypatch): + """Byte-identical behaviour, asserted as "the code did not run". + + A sync that happens to be a no-op today could stop being one; the + contract is that an unconfigured hook adds no step at all. + """ + import ag_ui_strands.agent as agent_module + + def explode(*args, **kwargs): # pragma: no cover - must not be reached + raise AssertionError( + "the template-tool sync ran with no provider configured" + ) + + monkeypatch.setattr(agent_module, "apply_template_tool_selection", explode) + monkeypatch.setattr(agent_module, "resolve_template_tool_selection", explode) + agent = make_agent(RecordingModel()) + events = await drain(agent, run_input("r1")) + assert [e for e in events if e.type == EventType.RUN_ERROR] == [] + + +# --------------------------------------------------------------------------- +# The pieces, in isolation +# --------------------------------------------------------------------------- + + +class TestSelectionResolution: + def test_names_and_tool_objects_are_both_accepted(self): + index = index_template_tools([read_docs, delete_record]) + assert resolve_template_tool_selection(["read_docs"], index) == {"read_docs"} + assert resolve_template_tool_selection([delete_record], index) == { + "delete_record" + } + + def test_none_and_empty_are_different_answers(self): + index = index_template_tools([read_docs]) + assert resolve_template_tool_selection(None, index) is None + assert resolve_template_tool_selection([], index) == set() + + def test_a_tool_the_template_never_contributed_is_refused(self, caplog): + @tool + def smuggled() -> str: + """Not on the template.""" + return "no" + + index = index_template_tools([read_docs]) + with caplog.at_level(logging.WARNING, logger="ag_ui_strands.template_tools"): + assert resolve_template_tool_selection( + ["read_docs", "smuggled", smuggled], index + ) == {"read_docs"} + assert caplog.text.count("does not contribute") == 2 + + def test_an_entry_that_names_no_tool_is_refused(self, caplog): + index = index_template_tools([read_docs]) + with caplog.at_level(logging.WARNING, logger="ag_ui_strands.template_tools"): + assert resolve_template_tool_selection([None, 7, ""], index) == set() + assert caplog.text.count("names no tool") == 3 + + +class TestSyncTemplateTools: + def _registry(self) -> ToolRegistry: + registry = ToolRegistry() + registry.process_tools([read_docs, delete_record]) + return registry + + def test_a_removed_tool_is_restored_as_the_same_instance(self): + registry = self._registry() + original = registry.registry["delete_record"] + + sync_template_tools(registry, [read_docs, delete_record], ["read_docs"]) + assert set(registry.registry) == {"read_docs"} + + sync_template_tools(registry, [read_docs, delete_record], None) + assert registry.registry["delete_record"] is original + + def test_a_denied_name_held_by_a_client_proxy_is_left_alone(self): + registry = self._registry() + proxy = create_proxy_tool( + Tool( + name="delete_record", + description="A client tool of the same name", + parameters={"type": "object", "properties": {}}, + ) + ) + registry.registry["delete_record"] = proxy + + kept = sync_template_tools(registry, [read_docs, delete_record], ["read_docs"]) + assert registry.registry["delete_record"] is proxy + assert kept == {"read_docs"} + + def test_an_allowed_name_held_by_a_client_proxy_is_reclaimed(self): + """Native tools win a name collision, and the proxy sync runs after. + + Leaving the proxy would shadow the tool the provider just allowed, and + if the client has stopped declaring it the proxy sync would then remove + the name outright, leaving an allowed tool registered nowhere. + """ + registry = self._registry() + template_tool = registry.registry["delete_record"] + registry.registry["delete_record"] = create_proxy_tool( + Tool( + name="delete_record", + description="A client tool of the same name", + parameters={"type": "object", "properties": {}}, + ) + ) + + kept = sync_template_tools(registry, [read_docs, delete_record], None) + assert registry.registry["delete_record"] is template_tool + assert kept == {"read_docs", "delete_record"} + + def test_a_parked_proxy_holding_a_name_is_not_reclaimed(self): + registry = self._registry() + proxy = create_proxy_tool( + Tool( + name="delete_record", + description="A client tool of the same name", + parameters={"type": "object", "properties": {}}, + ) + ) + registry.registry["delete_record"] = proxy + + sync_template_tools( + registry, + [read_docs, delete_record], + None, + exempt_names={"delete_record"}, + ) + assert registry.registry["delete_record"] is proxy + + def test_an_equivalent_but_not_identical_template_entry_is_still_ours(self): + """Ownership cannot rest on object identity alone. + + With an external ``agents_by_thread`` map the wrapper is rebuilt per + request while the cached thread agent keeps its registry, so a template + whose tools are built per request hands the adapter equivalent but not + identical objects. Reading that as another producer's entry would make + a deny-everything answer remove nothing. + """ + + @tool(name="delete_record") + def rebuilt(record_id: str) -> str: + """A second instance of the same template tool.""" + return record_id + + registry = self._registry() + assert registry.registry["delete_record"] is not rebuilt + + kept = sync_template_tools(registry, [read_docs, rebuilt], ["read_docs"]) + assert "delete_record" not in registry.registry + assert kept == {"read_docs"} + + def test_exempt_names_are_kept_however_the_selection_reads(self): + registry = self._registry() + kept = sync_template_tools( + registry, + [read_docs, delete_record], + [], + exempt_names={"delete_record"}, + ) + assert kept == {"delete_record"} + assert set(registry.registry) == {"delete_record"} + + def test_the_returned_set_is_what_the_registry_holds(self): + registry = self._registry() + assert sync_template_tools( + registry, [read_docs, delete_record], ["read_docs"] + ) == {"read_docs"} + assert sync_template_tools(registry, [read_docs, delete_record], None) == { + "read_docs", + "delete_record", + } + + +class TestParkedBatchToolNames: + def _agent(self, state): + class _Agent: + _interrupt_state = state + + return _Agent() + + def test_an_idle_agent_parks_nothing(self): + assert parked_batch_tool_names(self._agent(InterruptStateStub())) == set() + assert parked_batch_tool_names(object()) == set() + + def test_every_tool_in_the_parked_batch_is_named(self): + state = InterruptStateStub( + interrupts={"i1": StrandsInterrupt("i1", "ag_ui:tool_call:delete_record")}, + ) + state.activate( + { + "tool_use_message": { + "role": "assistant", + "content": [ + {"toolUse": {"toolUseId": "a", "name": "delete_record"}}, + {"toolUse": {"toolUseId": "b", "name": "read_docs"}}, + {"text": "thinking"}, + ], + } + } + ) + assert parked_batch_tool_names(self._agent(state)) == { + "delete_record", + "read_docs", + } + + def test_a_checkpoint_without_a_tool_batch_names_nothing(self): + state = InterruptStateStub() + state.activate({"tool_results": []}) + assert parked_batch_tool_names(self._agent(state)) == set() diff --git a/integrations/aws-strands/python/tests/test_terminal_error_paths.py b/integrations/aws-strands/python/tests/test_terminal_error_paths.py index 6026cf8bf4..e6e6f95002 100644 --- a/integrations/aws-strands/python/tests/test_terminal_error_paths.py +++ b/integrations/aws-strands/python/tests/test_terminal_error_paths.py @@ -241,6 +241,24 @@ def _kwargs(_input): assert_contract_error(_terminal_error(events), "THREAD_AGENT_KWARGS_ERROR") +@pytest.mark.asyncio +async def test_template_tools_provider_error_reports_a_hook_that_raised(): + def _tools(_input): + raise ValueError("authz lookup failed") + + events = await _drive( + _MockCore(), + _run_input(), + config=StrandsAgentConfig( + replay_history_into_strands=False, template_tools_provider=_tools + ), + ) + + assert_contract_error( + _terminal_error(events), "TEMPLATE_TOOLS_PROVIDER_ERROR" + ) + + # --------------------------------------------------------------------------- # Interrupt preflight # --------------------------------------------------------------------------- diff --git a/integrations/aws-strands/typescript/README.md b/integrations/aws-strands/typescript/README.md index 5783ac0f64..946bbc74b2 100644 --- a/integrations/aws-strands/typescript/README.md +++ b/integrations/aws-strands/typescript/README.md @@ -90,6 +90,7 @@ See [../ARCHITECTURE.md](../ARCHITECTURE.md) for diagrams and a deeper dive. | -------------------------- | ------------------------------------------------------------------------------- | | `src/agent.ts` | Core wrapper translating Strands streams into AG-UI events | | `src/config.ts` | Config primitives (`StrandsAgentConfig`, `ToolBehavior`, `PredictStateMapping`) | +| `src/template-tools.ts` | Per-request filter over the template agent's tools | | `src/server.ts` | `createStrandsApp` + Express transport (subpath: `@ag-ui/aws-strands/server`) | | `src/endpoint.ts` | Express endpoint helpers (used by `server.ts`) | | `src/utils.ts` | Multimodal content conversion and the `UrlFetchPolicy` that guards it | @@ -435,6 +436,96 @@ the streamable-HTTP one above is just the common case. See Strands' own [MCP tools guide](https://strandsagents.com/docs/user-guide/concepts/tools/mcp-tools/) for the transports and the elicitation callback. +## Per-request tool filtering + +`StrandsAgentConfig.templateToolsProvider` decides which of the template +agent's tools one request may see. It is called once per request with that +request's `RunAgentInput`, so the answer can vary turn by turn on a single +thread: + +```ts +const READ_ONLY = ["search_docs", "get_order"]; + +const aguiAgent = new StrandsAgent({ + agent, + name: "assistant", + config: { + templateToolsProvider: (input) => + // Derive the role from authenticated request context in production; + // forwardedProps is client-controlled. + (input.forwardedProps as { role?: string })?.role === "admin" + ? null // no filtering: every template tool stays available + : READ_ONLY, + }, +}); +``` + +Return the tools themselves or their names. `null` or `undefined` declines to +filter; an empty array is a real answer and withholds all of them. A name the +template does not contribute is dropped with a warning, because the hook +narrows the wrapped agent's tools and cannot add one. The provider may be +async. + +Two boundary rules follow from that: + +- **The return value is checked, not merely iterated.** A string and a mapping + are both iterable and both mean something other than what iterating them + produces: a bare name would come apart into characters, and a permission map + would have its keys read as an allow-list while its values went unread, so a + name mapped to `false` would still be allowed. Both are refused with + `TEMPLATE_TOOLS_PROVIDER_ERROR`. Arrays, sets and generators are all + accepted, and a generator that raises partway through iteration reports the + same code, because the answer is read inside the same guarded step that calls + the provider. +- **The filter reaches the registry, not only the advertised tool specs.** A + model that calls a withheld name anyway, primed by a stale turn or by the + visible history, is refused by the dispatcher rather than served. + +The filter is applied to the tool registry the thread's live Strands `Agent` +already owns, the same way client-declared tools are synchronised, and never by +rebuilding that agent. The per-thread instance holds the thread's +`SessionManager`, its native interrupt checkpoint and its history, so replacing +it to change a tool list would discard a conversation and any approval waiting +inside it. + +Three consequences follow from that: + +- **A parked call is never orphaned.** A tool in the batch a live interrupt + checkpoint would resume stays registered whatever the provider returns: the + human's answer is about to be routed back into that batch, and an absent tool + turns it into a "tool not found" the model re-fires. Filtering resumes once + the pause closes. This is the rule `syncProxyTools` already applies to a proxy + parked in a frontend-tool interrupt. +- **History is never rewritten.** A filtered-out tool's earlier calls and + results stay in the thread's messages, so the model can still read what it + did with a tool it can no longer call. +- **A failure is terminal.** If the provider throws, the run yields `RUN_ERROR` + with code `TEMPLATE_TOOLS_PROVIDER_ERROR` and stops, matching + `threadAgentConfig`. A filter that failed open would hand the model exactly + the tools the caller meant to withhold. + +The narrowing is also re-applied inside the run, once a tool batch has been +dispatched. The exemption above keeps a denied tool registered so a human's +answer can reach it, and Strands then carries on in the same run: it +re-dispatches the batch and makes its next model call from the same registry, +which would otherwise still be advertising what the request denied. The two +bridges hook different SDK events for this, because the SDKs read the tool +specs at different points relative to the events they dispatch; the effect is +the same on both. + +Scope is the template's own tools. Client-declared tools on +`RunAgentInput.tools` are re-synchronised from the request every turn already, +so a caller that wants fewer of those sends fewer. The hook is not applied on +the multi-agent orchestrator path, which has no template registry to filter. + +One deployment note. With an `agentsByThread` map, a request-scoped wrapper is +rebuilt per request while the cached thread agent keeps the registry it already +had. If the template's tools are built per request too, the adapter is handed +equivalent but not identical objects, so which registry entry belongs to the +template is decided by name plus "not one of the adapter's other producers" +rather than by object identity alone. Stable tool objects are still the simpler +thing to hand it. + ## Human-in-the-loop interrupts Two complementary patterns are supported: diff --git a/integrations/aws-strands/typescript/src/__tests__/exports.test.ts b/integrations/aws-strands/typescript/src/__tests__/exports.test.ts index efb2e6bddc..4a61d9a4cb 100644 --- a/integrations/aws-strands/typescript/src/__tests__/exports.test.ts +++ b/integrations/aws-strands/typescript/src/__tests__/exports.test.ts @@ -19,6 +19,8 @@ describe("public export surface", () => { "createProxyTool", "syncProxyTools", "isProxyTool", + "syncTemplateTools", + "parkedBatchToolNames", "DEFAULT_URL_FETCH_POLICY", "UrlFetchPolicyError", ]; diff --git a/integrations/aws-strands/typescript/src/__tests__/helpers.ts b/integrations/aws-strands/typescript/src/__tests__/helpers.ts index 0480ce7813..8ad5a6297d 100644 --- a/integrations/aws-strands/typescript/src/__tests__/helpers.ts +++ b/integrations/aws-strands/typescript/src/__tests__/helpers.ts @@ -359,6 +359,13 @@ export class ScriptedModel extends Model { */ public readonly handedMessages: StrandsMessage[][] = []; + /** + * The tool names offered on each turn, oldest first. What the registry held + * when the turn started, as the model itself saw it, which is the only place + * a per-request tool filter is observable from outside the adapter. + */ + public readonly offeredToolNames: Set[] = []; + private readonly config: Record = { modelId: "scripted-model", }; @@ -381,9 +388,15 @@ export class ScriptedModel extends Model { Object.assign(this.config, modelConfig); } - async *stream(messages: StrandsMessage[]): AsyncIterable { + async *stream( + messages: StrandsMessage[], + options?: { toolSpecs?: { name: string }[] }, + ): AsyncIterable { this.handedMessages.push(messages); this.seenMessages.push(messages.map(recordedCopy)); + this.offeredToolNames.push( + new Set((options?.toolSpecs ?? []).map((spec) => spec.name)), + ); if (this.throwOnCall !== undefined && this.calls + 1 === this.throwOnCall) { this.calls += 1; throw new Error("scripted provider failure"); diff --git a/integrations/aws-strands/typescript/src/__tests__/template-tool-filtering.test.ts b/integrations/aws-strands/typescript/src/__tests__/template-tool-filtering.test.ts new file mode 100644 index 0000000000..0b58d10a83 --- /dev/null +++ b/integrations/aws-strands/typescript/src/__tests__/template-tool-filtering.test.ts @@ -0,0 +1,972 @@ +/** + * Per-request filtering of the tools the template agent contributed. + * + * The adapter keeps one Strands `Agent` per thread and that instance is + * load-bearing: it holds the thread's `SessionManager`, its native interrupt + * checkpoint and its conversation history. So the property these tests exist to + * pin is not only that a filter takes effect, but that it takes effect on the + * registry the live instance already owns. A filter applied by rebuilding the + * thread's agent would pass a "the model saw fewer tools" assertion while + * silently discarding a conversation and any approval waiting inside it, which + * is why the identity of the cached agent is asserted alongside the tool specs. + * + * Driven through the real SDK: a genuine `Agent`, its real `ToolRegistry`, and + * a scripted model that records the tool specs each turn was offered. The + * offered specs are the only place the filter is observable from outside, so a + * registry assertion alone would not show the model was actually affected. + */ + +import { describe, it, expect, vi } from "vitest"; +import { EventType, type BaseEvent, type RunAgentInput } from "@ag-ui/core"; +import { + Agent as StrandsAgentCore, + BeforeModelCallEvent, + type Tool, +} from "@strands-agents/sdk"; + +import { StrandsAgent } from "../agent"; +import type { StrandsAgentConfig } from "../config"; +import { + EXEMPT_EVERY_TEMPLATE_TOOL, + indexTemplateTools, + parkedBatchToolNames, + resolveTemplateToolSelection, + syncTemplateTools, +} from "../template-tools"; +import { + createProxyTool, + isProxyTool, + type StrandsToolRegistry, +} from "../client-proxy-tool"; +import { + collect, + errorCodes, + fakeTool, + finishedOf, + interruptsOf, + minimalRunInput, + modelTurn, + realStrandsAgent, + recordingTool, + threadAgent, + ScriptedModel, +} from "./helpers"; + +const READ = "read_docs"; +const DELETE = "delete_record"; + +/** Two template tools and a two-turn script that answers with text. */ +function twoToolAgent(config: StrandsAgentConfig) { + const read = recordingTool(READ); + const del = recordingTool(DELETE); + const { agent, model } = realStrandsAgent( + [modelTurn.text("first"), modelTurn.text("second")], + { tools: [read.tool, del.tool], config }, + ); + return { agent, model, readCalls: read.calls, deleteCalls: del.calls }; +} + +function userTurn(overrides: Partial = {}): RunAgentInput { + return minimalRunInput({ + messages: [{ id: "u1", role: "user", content: "do the thing" } as never], + ...overrides, + }); +} + +function registryNames( + agent: StrandsAgent, + threadId = "thread-1", +): Set { + const core = threadAgent(agent, threadId); + if (!core) throw new Error(`no per-thread agent for "${threadId}"`); + return new Set(core.toolRegistry.list().map((t) => t.name)); +} + +function offered(model: ScriptedModel): Set { + const last = model.offeredToolNames.at(-1); + expect(last, "the model was never invoked").toBeDefined(); + return last!; +} + +// --------------------------------------------------------------------------- +// The rescope: filtering happens on the live agent, not by replacing it +// --------------------------------------------------------------------------- + +describe("filtering without rebuilding the thread agent", () => { + it("varies the filtered set between two requests on one thread", async () => { + const byRun: Record = { + "run-1": [READ], + "run-2": [READ, DELETE], + }; + const { agent, model } = twoToolAgent({ + templateToolsProvider: (input) => byRun[input.runId], + }); + + await collect(agent, userTurn()); + expect(offered(model)).toEqual(new Set([READ])); + expect(registryNames(agent)).toEqual(new Set([READ])); + + await collect(agent, userTurn({ runId: "run-2" })); + expect(offered(model)).toEqual(new Set([READ, DELETE])); + expect(registryNames(agent)).toEqual(new Set([READ, DELETE])); + }); + + it("keeps the same cached thread agent across the change", async () => { + // The whole point of applying the filter to the registry. Recreating the + // per-thread agent whenever the resolved tool set changed would satisfy the + // assertion above and still be wrong, so identity is asserted directly. + const byRun: Record = { "run-1": [READ], "run-2": [] }; + const { agent, model } = twoToolAgent({ + templateToolsProvider: (input) => byRun[input.runId], + }); + + await collect(agent, userTurn()); + const first = threadAgent(agent); + expect(first).toBeDefined(); + + await collect(agent, userTurn({ runId: "run-2" })); + expect(threadAgent(agent)).toBe(first); + expect(offered(model)).toEqual(new Set()); + }); + + it("withholds every template tool for an empty selection", async () => { + const { agent, model } = twoToolAgent({ templateToolsProvider: () => [] }); + await collect(agent, userTurn()); + expect(offered(model)).toEqual(new Set()); + }); + + it("treats null and undefined as declining to filter that request", async () => { + const byRun: Record = { + "run-1": [], + "run-2": null, + "run-3": undefined, + }; + const { agent, model } = twoToolAgent({ + templateToolsProvider: (input) => byRun[input.runId], + }); + + await collect(agent, userTurn()); + expect(offered(model)).toEqual(new Set()); + + await collect(agent, userTurn({ runId: "run-2" })); + expect(offered(model)).toEqual(new Set([READ, DELETE])); + + await collect(agent, userTurn({ runId: "run-3" })); + expect(offered(model)).toEqual(new Set([READ, DELETE])); + }); + + it("may be async and may read the caller identity off the request", async () => { + const { agent, model } = twoToolAgent({ + templateToolsProvider: async (input) => + (input.forwardedProps as { role?: string } | undefined)?.role === + "admin" + ? null + : [READ], + }); + + await collect(agent, userTurn({ forwardedProps: { role: "reader" } })); + expect(offered(model)).toEqual(new Set([READ])); + + await collect( + agent, + userTurn({ runId: "run-2", forwardedProps: { role: "admin" } }), + ); + expect(offered(model)).toEqual(new Set([READ, DELETE])); + }); + + it("lets two threads see different tools at the same time", async () => { + const { agent, model } = twoToolAgent({ + templateToolsProvider: (input) => + input.threadId === "wide" ? null : [READ], + }); + + await collect(agent, userTurn({ threadId: "wide" })); + expect(offered(model)).toEqual(new Set([READ, DELETE])); + + await collect(agent, userTurn({ threadId: "narrow", runId: "run-2" })); + expect(offered(model)).toEqual(new Set([READ])); + + expect(registryNames(agent, "wide")).toEqual(new Set([READ, DELETE])); + expect(registryNames(agent, "narrow")).toEqual(new Set([READ])); + }); +}); + +describe("what the thread keeps across a filter change", () => { + it("keeps its history and the calls a filtered-out tool already made", async () => { + // History is not rewritten: the filter answers "what may this request + // call", not "what happened on this thread". Removing the record would + // leave an assistant tool-use block with no result behind it. + const del = recordingTool(DELETE); + const read = recordingTool(READ); + const { agent, model } = realStrandsAgent( + [ + modelTurn.toolUse({ toolUseId: "tu-1", name: DELETE, input: {} }), + modelTurn.text("done"), + modelTurn.text("second turn"), + ], + { + tools: [read.tool, del.tool], + config: { + replayHistoryIntoStrands: false, + templateToolsProvider: (input) => + input.runId === "run-1" ? null : [READ], + }, + }, + ); + + await collect(agent, userTurn()); + const core = threadAgent(agent)!; + expect(del.calls, "the first run never called the tool").toHaveLength(1); + + // Live history carries `ToolUseBlock` instances, which name the tool + // directly; a serialized message nests the same fields under `toolUse`. + const deleteCallIds = (messages: typeof core.messages) => + messages.flatMap((message) => + (message.content ?? []) + .map((block) => { + const b = block as { + name?: string; + toolUseId?: string; + toolUse?: { name?: string; toolUseId?: string }; + }; + return b.toolUse ?? b; + }) + .filter((use) => use.name === DELETE) + .map((use) => use.toolUseId as string), + ); + const recorded = deleteCallIds(core.messages); + expect(recorded.length).toBeGreaterThan(0); + + await collect(agent, userTurn({ runId: "run-2" })); + expect(threadAgent(agent)).toBe(core); + expect( + deleteCallIds(core.messages), + "filtering the tool out erased the call it already made", + ).toEqual(recorded); + expect(offered(model)).toEqual(new Set([READ])); + }); + + it("keeps the session manager the thread was built with", async () => { + const sessionManager = { + initAgent: vi.fn(async () => {}), + appendMessage: vi.fn(async () => {}), + redactLatestMessage: vi.fn(async () => {}), + syncAgent: vi.fn(async () => {}), + }; + const { agent } = twoToolAgent({ + sessionManagerProvider: () => sessionManager as never, + templateToolsProvider: (input) => + input.runId === "run-1" ? null : [READ], + }); + + await collect(agent, userTurn()); + const core = threadAgent(agent)!; + const held = (core as unknown as { sessionManager?: unknown }) + .sessionManager; + expect(held).toBe(sessionManager); + + await collect(agent, userTurn({ runId: "run-2" })); + expect(threadAgent(agent)).toBe(core); + expect( + (core as unknown as { sessionManager?: unknown }).sessionManager, + ).toBe(sessionManager); + }); +}); + +describe("a parked call is not orphaned", () => { + it("keeps a tool awaiting approval registered while it is filtered out", async () => { + // The rule `syncProxyTools` already applies, read off the checkpoint. The + // human's answer is routed back into the tool batch the run stopped inside, + // so a tool absent from the registry at that moment turns the answer into a + // "tool not found" the model then re-fires. + const del = recordingTool(DELETE); + const read = recordingTool(READ); + const { agent } = realStrandsAgent( + [ + modelTurn.toolUse({ toolUseId: "tu-1", name: DELETE, input: {} }), + modelTurn.text("done"), + ], + { + tools: [read.tool, del.tool], + config: { + toolBehaviors: { [DELETE]: { interruptOnCall: true } }, + templateToolsProvider: (input) => + input.runId === "run-1" ? null : [READ], + }, + }, + ); + + const first = await collect(agent, userTurn()); + expect(errorCodes(first)).toEqual([]); + expect(finishedOf(first).outcome?.type).toBe("interrupt"); + const interruptId = interruptsOf(first)[0].id; + const core = threadAgent(agent)!; + + const second = await collect( + agent, + userTurn({ + runId: "run-2", + resume: [ + { interruptId, status: "resolved", payload: { approved: true } }, + ] as never, + }), + ); + + expect(threadAgent(agent)).toBe(core); + // The resume reaching the tool is the assertion. Dropping the exemption + // makes Strands refuse the tool as absent from the registry, which surfaces + // here as the approved tool never running. + expect(errorCodes(second)).toEqual([]); + expect(del.calls, "the approved tool never ran").toHaveLength(1); + }); + + it("leaves a parked frontend proxy alone when the filter allows nothing", async () => { + // A template filter has no reach over client-declared tools. Their proxies + // are a different producer's entries and are re-synchronised from + // `RunAgentInput.tools` every request, so the filter must not remove one. + const clientTool = { + name: "confirm_in_client", + description: "Confirm in the client", + parameters: { type: "object", properties: {} }, + }; + const { agent } = twoToolAgent({ + toolBehaviors: { + confirm_in_client: { continueAfterFrontendCall: false }, + }, + templateToolsProvider: () => [], + }); + + await collect(agent, userTurn({ tools: [clientTool] as never })); + const core = threadAgent(agent)!; + expect(new Set(core.toolRegistry.list().map((t) => t.name))).toEqual( + new Set(["confirm_in_client"]), + ); + + await collect( + agent, + userTurn({ runId: "run-2", tools: [clientTool] as never }), + ); + expect(threadAgent(agent)).toBe(core); + expect(core.toolRegistry.get("confirm_in_client")).toBeDefined(); + }); +}); + +describe("filtering removes the capability", () => { + it("does not run a filtered-out tool the model calls anyway", async () => { + // The point of touching the registry rather than only the tool specs. + // Withholding a tool from the specs and leaving it registered would make + // the filter advice a model can ignore. A model calling the name anyway, + // because it was primed by a stale turn or by the visible history, has to + // be refused by the dispatcher rather than served. + const del = recordingTool(DELETE); + const read = recordingTool(READ); + const { agent, model } = realStrandsAgent( + [ + modelTurn.toolUse({ toolUseId: "tu-1", name: DELETE, input: {} }), + modelTurn.text("done"), + ], + { + tools: [read.tool, del.tool], + config: { templateToolsProvider: () => [READ] }, + }, + ); + + const events = await collect(agent, userTurn()); + + expect(del.calls, "a filtered-out tool executed").toEqual([]); + expect(model.offeredToolNames[0]).toEqual(new Set([READ])); + expect(errorCodes(events)).toEqual([]); + }); +}); + +describe("a pause is out of the filter's reach on a plain turn", () => { + it("refuses the turn before the filter runs and keeps the pause intact", async () => { + // A turn that submits no answer is refused ahead of the tool sync, so a + // provider that would have withheld the parked tool never runs and the + // checkpoint is still there for the resume that follows. + const consulted: string[] = []; + const del = recordingTool(DELETE); + const read = recordingTool(READ); + const { agent } = realStrandsAgent( + [ + modelTurn.toolUse({ toolUseId: "tu-1", name: DELETE, input: {} }), + modelTurn.text("done"), + ], + { + tools: [read.tool, del.tool], + config: { + toolBehaviors: { [DELETE]: { interruptOnCall: true } }, + templateToolsProvider: (input) => { + consulted.push(input.runId); + return input.runId === "run-1" ? null : [READ]; + }, + logger: { debug() {}, warn() {}, error() {} }, + }, + }, + ); + + const first = await collect(agent, userTurn()); + const interruptId = interruptsOf(first)[0].id; + const core = threadAgent(agent)!; + + const refused = await collect(agent, userTurn({ runId: "run-2" })); + expect(errorCodes(refused)).toEqual(["PENDING_INTERRUPTS"]); + expect(consulted, "the provider ran on a turn that was refused").toEqual([ + "run-1", + ]); + expect(core.toolRegistry.get(DELETE)).toBeDefined(); + + const resumed = await collect( + agent, + userTurn({ + runId: "run-3", + resume: [ + { interruptId, status: "resolved", payload: { approved: true } }, + ] as never, + }), + ); + expect(errorCodes(resumed)).toEqual([]); + expect(threadAgent(agent)).toBe(core); + expect(del.calls).toHaveLength(1); + expect(consulted).toEqual(["run-1", "run-3"]); + }); +}); + +describe("the return contract", () => { + const quiet = { debug() {}, warn() {}, error() {} }; + + async function errorFor( + provider: StrandsAgentConfig["templateToolsProvider"], + ) { + const { agent, model } = twoToolAgent({ + templateToolsProvider: provider, + logger: quiet, + }); + const events = await collect(agent, userTurn()); + const error = events.find((e) => e.type === EventType.RUN_ERROR) as + | (BaseEvent & { code?: string; message?: string }) + | undefined; + return { error, model }; + } + + it("refuses a Map rather than reading it as an allow-list", async () => { + // The mistake Python failed silently and permissively on: iterating a + // permission map yields its keys, so every name would be allowed including + // the ones mapped to false. Refusing it is also what keeps one return + // contract across the two bridges. + const { error, model } = await errorFor( + () => + new Map([ + [READ, true], + [DELETE, false], + ]) as never, + ); + expect(error?.code).toBe("TEMPLATE_TOOLS_PROVIDER_ERROR"); + expect(error?.message).toContain("values went unread"); + expect(model.calls, "the model ran unfiltered").toBe(0); + }); + + it("refuses a plain object with the same error, not a bare TypeError", async () => { + const { error } = await errorFor( + () => ({ [READ]: true, [DELETE]: false }) as never, + ); + expect(error?.code).toBe("TEMPLATE_TOOLS_PROVIDER_ERROR"); + expect(error?.message).toContain("not a container"); + }); + + it("refuses a bare name rather than reading it one character at a time", async () => { + const { error } = await errorFor(() => READ as never); + expect(error?.code).toBe("TEMPLATE_TOOLS_PROVIDER_ERROR"); + expect(error?.message).toContain("one character at a time"); + }); + + it("refuses a non-container", async () => { + const { error } = await errorFor(() => 42 as never); + expect(error?.code).toBe("TEMPLATE_TOOLS_PROVIDER_ERROR"); + }); + + it("reports a generator that throws partway through iteration", async () => { + // The provider's answer is read inside the guarded arm, not after it. A + // generator constructs without running its body, so reading it outside the + // boundary skipped the documented code and ended the stream after + // RUN_STARTED with nothing terminal behind it. + const { error, model } = await errorFor(function* () { + yield READ; + throw new Error("directory lookup failed"); + }); + expect(error?.code).toBe("TEMPLATE_TOOLS_PROVIDER_ERROR"); + expect(error?.message).toContain("directory lookup failed"); + expect(model.calls).toBe(0); + }); + + it("accepts a generator of names", async () => { + const { agent, model } = twoToolAgent({ + templateToolsProvider: function* () { + yield READ; + }, + }); + await collect(agent, userTurn()); + expect(offered(model)).toEqual(new Set([READ])); + }); + + it("accepts a Set of names", async () => { + const { agent, model } = twoToolAgent({ + templateToolsProvider: () => new Set([READ]), + }); + await collect(agent, userTurn()); + expect(offered(model)).toEqual(new Set([READ])); + }); +}); + +describe("the exemption does not outlive the checkpoint", () => { + it("narrows again before the next model call in the same run", async () => { + // Strands keeps running after a resume, from this same registry: it + // re-dispatches the batch, clears the checkpoint and calls the model again + // within the same run. Without a re-narrowing that call would still + // advertise the tool the request denied. + const del = recordingTool(DELETE); + const read = recordingTool(READ); + const { agent, model } = realStrandsAgent( + [ + modelTurn.toolUse({ toolUseId: "tu-1", name: DELETE, input: {} }), + modelTurn.text("done"), + ], + { + tools: [read.tool, del.tool], + config: { + toolBehaviors: { [DELETE]: { interruptOnCall: true } }, + templateToolsProvider: (input) => + input.runId === "run-1" ? null : [READ], + }, + }, + ); + + const first = await collect(agent, userTurn()); + const interruptId = interruptsOf(first)[0].id; + const offeredBeforeResume = model.offeredToolNames.length; + + const resumed = await collect( + agent, + userTurn({ + runId: "run-2", + resume: [ + { interruptId, status: "resolved", payload: { approved: true } }, + ] as never, + }), + ); + + expect(errorCodes(resumed)).toEqual([]); + expect( + model.offeredToolNames.length, + "the resume never reached another model call, so this asserts nothing", + ).toBeGreaterThan(offeredBeforeResume); + for (const seen of model.offeredToolNames.slice(offeredBeforeResume)) { + expect( + seen, + "a model call after the resume still advertised the denied tool", + ).toEqual(new Set([READ])); + } + }); + + it("holds nothing for a pause raised before any tool ran", () => { + expect( + parkedBatchToolNames({ _interruptState: { activated: true } }), + ).toEqual(new Set()); + }); + + it("holds everything for a parked batch it cannot read", () => { + expect( + parkedBatchToolNames({ + _interruptState: { + activated: true, + pendingToolExecution: { assistantMessageData: "not a message" }, + }, + }), + ).toBe(EXEMPT_EVERY_TEMPLATE_TOOL); + }); +}); + +describe("the ordering against the proxy sync", () => { + const clientTool = { + name: DELETE, + description: "A client tool of the same name", + parameters: { type: "object", properties: {} }, + }; + + it("keeps an allowed tool when the client drops a colliding name", async () => { + // Neither producer ended up holding the name, for exactly one request. A + // client tool sharing a template tool's name takes the registry slot while + // the template tool is filtered out; when the provider allows the template + // tool again and the client has stopped declaring its own, the template + // sync used to decline to touch the proxy and the proxy sync then removed + // it as stale, leaving the allowed tool registered nowhere. + const byRun: Record = { + "run-1": [READ], + "run-2": null, + }; + const { agent, model } = twoToolAgent({ + templateToolsProvider: (input) => byRun[input.runId], + }); + + await collect(agent, userTurn({ tools: [clientTool] as never })); + const core = threadAgent(agent)!; + + await collect(agent, userTurn({ runId: "run-2" })); + expect(threadAgent(agent)).toBe(core); + expect( + core.toolRegistry.get(DELETE), + "the provider allowed the tool and it is registered nowhere", + ).toBeDefined(); + expect(offered(model)).toEqual(new Set([READ, DELETE])); + }); + + it("does not let a client tool shadow an allowed template tool", async () => { + const byRun: Record = { + "run-1": [READ], + "run-2": null, + }; + const { agent } = twoToolAgent({ + templateToolsProvider: (input) => byRun[input.runId], + logger: { debug() {}, warn() {}, error() {} }, + }); + + await collect(agent, userTurn({ tools: [clientTool] as never })); + const core = threadAgent(agent)!; + expect(isProxyTool(core.toolRegistry.get(DELETE))).toBe(true); + + await collect( + agent, + userTurn({ runId: "run-2", tools: [clientTool] as never }), + ); + expect(isProxyTool(core.toolRegistry.get(DELETE))).toBe(false); + }); +}); + +describe("ownership is not pure identity", () => { + it("lets a rebuilt wrapper deny a cached thread's tools", async () => { + // The shape `agentsByThread` exists for: a request-scoped wrapper is + // rebuilt per request while the cached thread agent keeps the registry it + // already had. A template whose tools are built per request then hands each + // new wrapper equivalent but not identical objects, and ownership by object + // identity alone would read every one of them as another producer's entry, + // so a deny-everything answer would remove nothing at all. + const agentsByThread = new Map(); + const byRun: Record = { + "run-1": null, + "run-2": [], + }; + // Registered ahead of the adapter's own re-narrowing hook, which is added + // after, so this snapshot is the registry the request path left behind + // rather than the one the hook went on to correct. + const seenBeforeModelCall: Set[] = []; + + function build(): StrandsAgent { + const template = new StrandsAgentCore({ + model: new ScriptedModel([modelTurn.text("hi")]), + tools: [recordingTool(READ).tool, recordingTool(DELETE).tool] as never, + printer: false, + }); + const wrapper = new StrandsAgent({ + agent: template, + name: "rebuilt-wrapper", + agentsByThread, + config: { + templateToolsProvider: (input) => byRun[input.runId], + }, + }); + return wrapper; + } + + const first = build(); + // The snapshot hook has to reach the per-thread agent before the adapter's + // own, which is added when that agent is built, so it goes on the template + // the first wrapper clones from. + const firstThreadAgentHook = (built: StrandsAgentCore) => + built.addHook(BeforeModelCallEvent, () => { + seenBeforeModelCall.push( + new Set(built.toolRegistry.list().map((t) => t.name)), + ); + }); + + await collect(first, userTurn()); + const core = agentsByThread.get("thread-1")!; + expect(new Set(core.toolRegistry.list().map((t) => t.name))).toEqual( + new Set([READ, DELETE]), + ); + firstThreadAgentHook(core); + + const second = build(); + const secondTools = ( + second as unknown as { _templateFields: { tools: Tool[] } } + )._templateFields.tools; + expect( + secondTools.every((t) => core.toolRegistry.get(t.name) !== t), + "the rebuilt template reused its tool objects, so this asserts nothing", + ).toBe(true); + + await collect(second, userTurn({ runId: "run-2" })); + expect(agentsByThread.get("thread-1")).toBe(core); + expect(core.toolRegistry.list()).toEqual([]); + }); +}); + +describe("the provider failure mode", () => { + it("ends the run rather than running unfiltered", async () => { + // Terminal, matching `threadAgentConfig`. Degrading to an unfiltered run + // would hand the model exactly the tools the caller meant to withhold, + // which is the one outcome this hook exists to prevent. + const { agent, model } = twoToolAgent({ + templateToolsProvider: () => { + throw new Error("authz lookup failed"); + }, + logger: { debug() {}, warn() {}, error() {} }, + }); + + const events = await collect(agent, userTurn()); + expect(events.map((e) => e.type)).toEqual([ + EventType.RUN_STARTED, + EventType.RUN_ERROR, + ]); + const error = events[1] as BaseEvent & { code?: string; message?: string }; + expect(error.code).toBe("TEMPLATE_TOOLS_PROVIDER_ERROR"); + expect(error.message).toContain("authz lookup failed"); + expect(model.calls, "the model ran despite the failure").toBe(0); + }); + + it("reports a rejected promise the same way", async () => { + const { agent } = twoToolAgent({ + templateToolsProvider: async () => { + throw new Error("network down"); + }, + logger: { debug() {}, warn() {}, error() {} }, + }); + expect(errorCodes(await collect(agent, userTurn()))).toEqual([ + "TEMPLATE_TOOLS_PROVIDER_ERROR", + ]); + }); +}); + +describe("no provider configured", () => { + it("leaves the registry and the offered specs untouched", async () => { + const { agent, model } = twoToolAgent({}); + await collect(agent, userTurn()); + // Identity, not just names: an unconfigured hook must add no step at all, + // so nothing may be removed and re-registered even to the same effect. + const held = new Map( + threadAgent(agent)! + .toolRegistry.list() + .map((t) => [t.name, t]), + ); + expect([...held.keys()].sort()).toEqual([DELETE, READ]); + + await collect(agent, userTurn({ runId: "run-2" })); + for (const [name, tool] of held) { + expect(threadAgent(agent)!.toolRegistry.get(name)).toBe(tool); + } + expect(model.offeredToolNames).toEqual([ + new Set([READ, DELETE]), + new Set([READ, DELETE]), + ]); + }); +}); + +// --------------------------------------------------------------------------- +// The pieces, in isolation +// --------------------------------------------------------------------------- + +describe("selection resolution", () => { + const read = fakeTool(READ) as unknown as Tool; + const del = fakeTool(DELETE) as unknown as Tool; + const index = indexTemplateTools([read, del]); + + it("accepts names and tool objects alike", () => { + expect(resolveTemplateToolSelection([READ], index)).toEqual( + new Set([READ]), + ); + expect(resolveTemplateToolSelection([del], index)).toEqual( + new Set([DELETE]), + ); + }); + + it("distinguishes declining to filter from filtering everything out", () => { + expect(resolveTemplateToolSelection(null, index)).toBeNull(); + expect(resolveTemplateToolSelection(undefined, index)).toBeNull(); + expect(resolveTemplateToolSelection([], index)).toEqual(new Set()); + }); + + it("refuses a tool the template never contributed", () => { + const warn = vi.fn(); + const log = { debug() {}, warn, error() {} }; + const smuggled = fakeTool( + "smuggled", + "not on the template", + ) as unknown as Tool; + expect( + resolveTemplateToolSelection([READ, "smuggled", smuggled], index, log), + ).toEqual(new Set([READ])); + expect(warn).toHaveBeenCalledTimes(2); + expect(warn.mock.calls.every(([m]) => /does not contribute/.test(m))).toBe( + true, + ); + }); + + it("refuses an entry that names no tool", () => { + const warn = vi.fn(); + const log = { debug() {}, warn, error() {} }; + expect( + resolveTemplateToolSelection([null, 7, ""] as never, index, log), + ).toEqual(new Set()); + expect(warn).toHaveBeenCalledTimes(3); + }); +}); + +describe("syncTemplateTools", () => { + function registry(): { registry: StrandsToolRegistry; tools: Tool[] } { + const core = new StrandsAgentCore({ + model: new ScriptedModel([]), + tools: [recordingTool(READ).tool, recordingTool(DELETE).tool] as never, + printer: false, + }); + return { registry: core.toolRegistry, tools: core.tools.slice() }; + } + + it("restores a removed tool as the same instance", () => { + const { registry: reg, tools } = registry(); + const original = reg.get(DELETE); + expect(original).toBeDefined(); + + syncTemplateTools(reg, tools, [READ]); + expect(reg.get(DELETE)).toBeUndefined(); + + syncTemplateTools(reg, tools, null); + expect(reg.get(DELETE)).toBe(original); + }); + + it("leaves a denied name held by a client proxy alone", () => { + const { registry: reg, tools } = registry(); + reg.remove(DELETE); + const proxy = createProxyTool({ + name: DELETE, + description: "A client tool of the same name", + parameters: { type: "object", properties: {} }, + } as never); + reg.add(proxy); + + expect(syncTemplateTools(reg, tools, [READ])).toEqual(new Set([READ])); + expect(reg.get(DELETE)).toBe(proxy); + }); + + it("reclaims an allowed name held by a client proxy", () => { + // Native tools win a name collision, and the proxy sync runs after this. + // Leaving the proxy would shadow the tool the provider just allowed, and if + // the client has stopped declaring it the proxy sync would then remove the + // name outright, leaving an allowed tool registered nowhere. + const { registry: reg, tools } = registry(); + const templateTool = reg.get(DELETE); + reg.remove(DELETE); + reg.add( + createProxyTool({ + name: DELETE, + description: "A client tool of the same name", + parameters: { type: "object", properties: {} }, + } as never), + ); + + expect(syncTemplateTools(reg, tools, null)).toEqual( + new Set([READ, DELETE]), + ); + expect(reg.get(DELETE)).toBe(templateTool); + }); + + it("does not reclaim a name the parked batch is answering", () => { + const { registry: reg, tools } = registry(); + reg.remove(DELETE); + const proxy = createProxyTool({ + name: DELETE, + description: "A client tool of the same name", + parameters: { type: "object", properties: {} }, + } as never); + reg.add(proxy); + + syncTemplateTools(reg, tools, null, { exemptNames: new Set([DELETE]) }); + expect(reg.get(DELETE)).toBe(proxy); + }); + + it("treats an equivalent but not identical template entry as ours", () => { + // Ownership cannot rest on object identity alone. With an external + // `agentsByThread` map the wrapper is rebuilt per request while the cached + // thread agent keeps its registry, so a template whose tools are built per + // request hands the adapter equivalent but not identical objects. Reading + // that as another producer's entry would make a deny-everything answer + // remove nothing. + const { registry: reg } = registry(); + const rebuiltRead = recordingTool(READ).tool; + const rebuiltDelete = recordingTool(DELETE).tool; + expect(reg.get(DELETE)).not.toBe(rebuiltDelete); + + const kept = syncTemplateTools(reg, [rebuiltRead, rebuiltDelete], [READ]); + expect(reg.get(DELETE)).toBeUndefined(); + expect(kept).toEqual(new Set([READ])); + }); + + it("holds every template tool for an unreadable parked batch", () => { + const { registry: reg, tools } = registry(); + expect(syncTemplateTools(reg, tools, [READ])).toEqual(new Set([READ])); + expect( + syncTemplateTools(reg, tools, [READ], { + exemptNames: EXEMPT_EVERY_TEMPLATE_TOOL, + }), + ).toEqual(new Set([READ, DELETE])); + }); + + it("keeps an exempt name however the selection reads", () => { + const { registry: reg, tools } = registry(); + expect( + syncTemplateTools(reg, tools, [], { exemptNames: new Set([DELETE]) }), + ).toEqual(new Set([DELETE])); + expect(new Set(reg.list().map((t) => t.name))).toEqual(new Set([DELETE])); + }); + + it("returns what the registry holds", () => { + const { registry: reg, tools } = registry(); + expect(syncTemplateTools(reg, tools, [READ])).toEqual(new Set([READ])); + expect(syncTemplateTools(reg, tools, null)).toEqual( + new Set([READ, DELETE]), + ); + }); +}); + +describe("parkedBatchToolNames", () => { + it("parks nothing for an idle agent", () => { + expect(parkedBatchToolNames(undefined)).toEqual(new Set()); + expect(parkedBatchToolNames({})).toEqual(new Set()); + expect( + parkedBatchToolNames({ _interruptState: { activated: false } }), + ).toEqual(new Set()); + }); + + it("names every tool in the parked batch", () => { + const agent = { + _interruptState: { + activated: true, + pendingToolExecution: { + assistantMessageData: { + role: "assistant", + content: [ + { toolUse: { toolUseId: "a", name: DELETE, input: {} } }, + { toolUse: { toolUseId: "b", name: READ, input: {} } }, + { text: "thinking" }, + ], + }, + }, + }, + }; + expect(parkedBatchToolNames(agent)).toEqual(new Set([DELETE, READ])); + }); + + it("names nothing for a checkpoint carrying no tool batch", () => { + expect( + parkedBatchToolNames({ _interruptState: { activated: true } }), + ).toEqual(new Set()); + }); +}); diff --git a/integrations/aws-strands/typescript/src/__tests__/terminal-error-paths.test.ts b/integrations/aws-strands/typescript/src/__tests__/terminal-error-paths.test.ts index 822cfc25d4..f84381f875 100644 --- a/integrations/aws-strands/typescript/src/__tests__/terminal-error-paths.test.ts +++ b/integrations/aws-strands/typescript/src/__tests__/terminal-error-paths.test.ts @@ -145,6 +145,21 @@ describe("terminal error paths", () => { expectContractError(terminalError(events), "THREAD_AGENT_CONFIG_ERROR"); }); + it("reports a templateToolsProvider hook that threw as TEMPLATE_TOOLS_PROVIDER_ERROR", async () => { + const agent = uncachedAgent({ + templateToolsProvider: () => { + throw new Error("authz lookup failed"); + }, + }); + + const events = await collect(agent, minimalRunInput()); + + expectContractError( + terminalError(events), + "TEMPLATE_TOOLS_PROVIDER_ERROR", + ); + }); + it("reports a seed this bridge cannot build as SEED_BUILD_ERROR", async () => { // The seed is built from the client's own messages, which nothing // validates. A message whose content cannot even be read is the diff --git a/integrations/aws-strands/typescript/src/agent.ts b/integrations/aws-strands/typescript/src/agent.ts index c50964a8e4..a332e7b7a1 100644 --- a/integrations/aws-strands/typescript/src/agent.ts +++ b/integrations/aws-strands/typescript/src/agent.ts @@ -8,6 +8,7 @@ import { createHash, randomUUID } from "crypto"; import { Agent as StrandsAgentCore, + AfterToolsEvent, BeforeToolCallEvent, InterruptResponseContent, Message as StrandsMessage, @@ -56,6 +57,15 @@ import { jsonRoundTrip, } from "./citations"; import { isProxyTool, syncProxyTools } from "./client-proxy-tool"; +import { + applyTemplateToolSelection, + indexTemplateTools, + parkedBatchToolNames, + recordTemplateToolSelection, + renarrowTemplateTools, + resolveTemplateToolSelection, +} from "./template-tools"; +import type { TemplateToolSelection } from "./template-tools"; import { planA2UIInjection, isAutoInjectedA2UITool, @@ -2376,6 +2386,29 @@ export class StrandsAgent { callerConfig, ), ); + // Re-narrow the per-request tool filter once a tool batch has run. The + // parked-batch exemption holds a denied tool registered so a resume can + // reach it, and Strands then continues the same run against the same + // registry; without this the run would keep advertising what the request + // denied until the next request narrowed again. + // + // `AfterToolsEvent`, not `BeforeModelCallEvent`: this SDK reads the tool + // specs off the registry as the first statement of its model call and + // dispatches `BeforeModelCallEvent` after, so a hook there would narrow + // the registry a moment too late to affect the specs it was narrowing + // for. Python's loop dispatches before that read, and its half of this + // hook uses `BeforeModelCallEvent`; it has no `AfterToolsEvent` to use. + if (this.config.templateToolsProvider) { + const built = strandsAgent; + built.addHook(AfterToolsEvent, () => { + renarrowTemplateTools( + built, + this._templateFields.tools ?? [], + this._log, + ); + }); + } + // Register interruptOnCall hooks on the per-thread agent. const behaviors = this.config.toolBehaviors; if (behaviors) { @@ -2831,6 +2864,70 @@ export class StrandsAgent { } const strandsAgent = agentResult.agent; + // Filter the tools the template contributed, per request. Applied to the + // registry this thread's live agent already owns: that instance carries the + // thread's session manager, its interrupt checkpoint and its history, so + // rebuilding it to change a tool list would discard a conversation and any + // approval waiting inside it. + if (this.config.templateToolsProvider) { + // Calling the provider and reading its answer are guarded together. + // Reading is where a Map, a bare name or a generator that throws partway + // through is caught, and those are provider mistakes: leaving them + // outside this arm let them bypass the documented code and end the + // stream after RUN_STARTED with nothing terminal behind it. + let allowed: Set | null; + try { + const selection: TemplateToolSelection = await maybeAwait( + this.config.templateToolsProvider(inputData), + ); + allowed = resolveTemplateToolSelection( + selection, + indexTemplateTools(this._templateFields.tools ?? []), + this._log, + ); + } catch (e) { + const msg = _errorMessage(e); + this._log.error( + `${LOG_PREFIX} templateToolsProvider failed: ${msg}`, + e, + ); + // Deliberately terminal rather than unfiltered: a filter that fails + // open hands the model tools the caller meant to withhold. + yield _runError( + `Failed to resolve the template tools for this request: ${msg}`, + "TEMPLATE_TOOLS_PROVIDER_ERROR", + ); + return; + } + // Guarded separately, and not as a provider error: past this point a + // failure is this adapter's, and this block runs before the main + // try/catch below, so it still must not escape as a stream that stops + // with nothing terminal behind it. + try { + applyTemplateToolSelection( + strandsAgent.toolRegistry, + this._templateFields.tools ?? [], + allowed, + { + exemptNames: parkedBatchToolNames(strandsAgent), + log: this._log, + }, + ); + // Published for the re-narrowing hook: the exemption above holds a + // denied tool registered so a resume can reach it, and Strands then + // continues the same run from this registry. + recordTemplateToolSelection(strandsAgent, allowed); + } catch (e) { + const msg = _errorMessage(e); + this._log.error( + `${LOG_PREFIX} applying the template tool filter failed: ${msg}`, + e, + ); + yield _runError(msg, _terminalErrorCode(e)); + return; + } + } + // Sync proxy tools from client-defined tools. if (inputData.tools && inputData.tools.length > 0) { const proxyNames = syncProxyTools( diff --git a/integrations/aws-strands/typescript/src/config.ts b/integrations/aws-strands/typescript/src/config.ts index afe73e5fb8..902e02ac93 100644 --- a/integrations/aws-strands/typescript/src/config.ts +++ b/integrations/aws-strands/typescript/src/config.ts @@ -3,6 +3,7 @@ import type { RunAgentInput, BaseEvent } from "@ag-ui/core"; import type { AgentConfig, SessionManager } from "@strands-agents/sdk"; import type { A2UIInjectConfig } from "./a2ui-tool"; +import type { TemplateToolSelectionEntry } from "./template-tools"; import type { Logger } from "./logger"; import type { UrlFetchPolicy } from "./utils"; @@ -141,6 +142,15 @@ export type ThreadAgentConfigProvider = ( input: RunAgentInput, ) => Partial | Promise>; +/** + * Chooses which of the template's tools one request may see. + * + * See {@link StrandsAgentConfig.templateToolsProvider}. + */ +export type TemplateToolsProvider = ( + input: RunAgentInput, +) => MaybePromise | null | undefined>; + /** Top-level configuration for the Strands agent adapter. */ export interface StrandsAgentConfig { /** Per-tool overrides keyed by the Strands tool name. */ @@ -207,6 +217,75 @@ export interface StrandsAgentConfig { * ``` */ threadAgentConfig?: ThreadAgentConfigProvider; + /** + * Which of the template agent's tools this request may see. + * + * Called once per request with that request's `RunAgentInput`, so the answer + * can vary turn by turn on one thread: the caller's identity is in + * `forwardedProps` or `context`, and a tool the request must not reach is + * simply left out of the returned iterable. May be async. + * + * Return the tools themselves or their names, whichever is to hand. Return + * `null` or `undefined` to decline filtering, which leaves every template + * tool available; an empty array is a real answer and leaves none of them. A + * name the template does not contribute is dropped with a warning, because + * this hook narrows the wrapped agent's tools and cannot add one. + * + * The container is checked rather than merely iterated. A `string` and a + * `Map` are both refused, as is a plain object: a bare name would come apart + * into characters, and a permission map would have its keys read as an + * allow-list while its values went unread, so a name mapped to `false` would + * still be allowed. Arrays, sets and generators are all accepted. + * + * Applied to the live per-thread agent's tool registry, never by rebuilding + * that agent: the instance holds the thread's `SessionManager`, its native + * interrupt checkpoint and its history, so replacing it to change a tool list + * would discard a conversation and any approval waiting inside it. + * + * Three consequences worth knowing: + * + * - A tool in the batch a live interrupt checkpoint would resume stays + * registered whatever this returns. The human's answer is about to be + * routed back into that batch, and an absent tool turns it into a "tool not + * found" the model re-fires. This is the rule `syncProxyTools` already + * applies to a proxy parked in a frontend-tool interrupt. The exemption + * does not outlast what it is for: the narrowing is re-applied inside the + * run once the batch has been dispatched, before the model is asked again. + * - History is never rewritten. A filtered-out tool's earlier calls and + * results stay in the thread's messages, so the model can still read what + * it did with a tool it can no longer call, and a provider that returns + * different sets across turns does not invalidate the transcript. + * - If it throws, the run yields `RUN_ERROR` with code + * `TEMPLATE_TOOLS_PROVIDER_ERROR` and stops, matching `threadAgentConfig`. + * A filter that fails open would hand the model tools the caller meant to + * withhold. + * + * Client-declared tools on `RunAgentInput.tools` are outside this hook: they + * are re-synchronised from the request every turn already, so a caller that + * wants fewer of those sends fewer. Not applied on the multi-agent + * orchestrator path, which has no template registry to filter. + * + * One deployment note. With an `agentsByThread` map a request-scoped wrapper + * is rebuilt per request while the cached thread agent keeps the registry it + * already had, so a template whose tools are built per request hands the + * adapter equivalent but not identical objects. Ownership of a registry + * entry therefore falls back from object identity to the tool's name plus + * "not one of the adapter's other producers". Stable tool objects are still + * the simpler thing to hand it. + * + * @example + * ```ts + * new StrandsAgent({ + * agent: template, + * name: "assistant", + * config: { + * templateToolsProvider: (input) => + * input.forwardedProps?.role === "admin" ? null : ["read_docs"], + * }, + * }) + * ``` + */ + templateToolsProvider?: TemplateToolsProvider; /** * Emit `MessagesSnapshotEvent` at lifecycle boundaries (after the initial * `STATE_SNAPSHOT`, after each `TOOL_CALL_END` / `TOOL_CALL_RESULT`, and diff --git a/integrations/aws-strands/typescript/src/index.ts b/integrations/aws-strands/typescript/src/index.ts index 118db5a8f8..fdb038a243 100644 --- a/integrations/aws-strands/typescript/src/index.ts +++ b/integrations/aws-strands/typescript/src/index.ts @@ -16,6 +16,9 @@ export { } from "./client-proxy-tool"; export type { StrandsToolRegistry } from "./client-proxy-tool"; +export { syncTemplateTools, parkedBatchToolNames } from "./template-tools"; +export type { TemplateToolSelectionEntry } from "./template-tools"; + export { CITATIONS_METADATA_KEY } from "./citations"; export type { AguiCitation, AguiCitationLocation } from "./citations"; @@ -70,6 +73,8 @@ export type { ToolStreamEventHandler, PredictStateMapping, SessionManagerProvider, + TemplateToolsProvider, + ThreadAgentConfigProvider, StateContextBuilder, StateFromArgs, StateFromResult, diff --git a/integrations/aws-strands/typescript/src/template-tools.ts b/integrations/aws-strands/typescript/src/template-tools.ts new file mode 100644 index 0000000000..9d69da3c7a --- /dev/null +++ b/integrations/aws-strands/typescript/src/template-tools.ts @@ -0,0 +1,437 @@ +/** + * Per-request filtering of the tools the template agent contributed. + * + * The adapter builds one Strands `Agent` per thread and keeps it. That instance + * is load-bearing: it holds the thread's `SessionManager`, its native interrupt + * checkpoint and its conversation history. Changing which tools a request sees + * therefore has to be done to the registry the live instance already owns, the + * way client-declared tools are already synchronised, and never by constructing + * a replacement. + * + * Scope is the template's own tools. Client-declared tools arrive on + * `RunAgentInput.tools` every request and are synchronised by + * {@link syncProxyTools}; a caller that wants fewer of those sends fewer. + * Auto-injected A2UI tools are the adapter's and are refreshed per turn. What + * no per-request channel reached until now is the set the wrapped template + * contributed once, at construction. + */ + +import type { Tool } from "@strands-agents/sdk"; + +import { isAutoInjectedA2UITool } from "./a2ui-tool"; +import { isProxyTool, type StrandsToolRegistry } from "./client-proxy-tool"; +import { DEFAULT_LOGGER, type Logger } from "./logger"; + +const LOG_PREFIX = "[@ag-ui/aws-strands]"; + +/** One entry a `templateToolsProvider` may return: a template tool or its name. */ +export type TemplateToolSelectionEntry = Tool | string; + +/** What a `templateToolsProvider` may answer with. */ +export type TemplateToolSelection = + | Iterable + | null + | undefined; + +/** + * The narrowed name set for the run in flight, stamped on the per-thread agent + * so the re-narrowing hook can read it back. Deliberately not routed through + * the agent's own state, which a `SessionManager` persists; this is per-request + * scratch that must not outlive the process. + */ +const ALLOWED_KEY = Symbol.for("@ag-ui/aws-strands.templateToolsAllowed"); + +/** A provider answer this hook cannot read as a selection of tools. */ +export class TemplateToolsSelectionError extends Error { + constructor(message: string) { + super(message); + this.name = "TemplateToolsSelectionError"; + } +} + +/** + * What {@link parkedBatchToolNames} answers for a checkpoint it cannot read. + * + * A distinct value rather than `undefined`: absent already means "no + * exemptions" wherever `exemptNames` is passed, and the two are opposites. + */ +export const EXEMPT_EVERY_TEMPLATE_TOOL = Symbol.for( + "@ag-ui/aws-strands.exemptEveryTemplateTool", +); + +/** Names to hold registered, or the sentinel meaning "hold all of them". */ +export type TemplateToolExemption = + | ReadonlySet + | typeof EXEMPT_EVERY_TEMPLATE_TOOL; + +/** Index the template's tools by the name a registry holds them under. */ +export function indexTemplateTools( + templateTools: readonly unknown[], +): Map { + const indexed = new Map(); + for (const tool of templateTools) { + const name = (tool as { name?: unknown } | null | undefined)?.name; + if (typeof name === "string" && name.length > 0) { + indexed.set(name, tool as Tool); + } + } + return indexed; +} + +/** + * Read one provider answer as the template tool names a request may see. + * + * Entries are either the template's own tool objects or their names, so a + * caller can write the filter with whichever it has to hand. + * + * `null`/`undefined` means the provider declined to filter this request and + * every template tool stays available. An empty array is a real answer and + * means none of them do. + * + * The container is checked rather than merely iterated. A `string` and a `Map` + * are both iterable and both mean something other than what iterating them + * produces: a name would come apart into characters, and a permission map would + * have its keys read as an allow-list while its values went unread. A plain + * object is refused for the same reason rather than left to throw a bare + * `TypeError`, which keeps one return contract across the two bridges rather + * than two. + * + * A name the template never contributed is dropped with a warning. This hook + * narrows what the wrapped agent already gave the adapter; it cannot hand the + * model a capability the template did not carry, so honouring an unknown name + * is the one thing it must not do. + * + * @throws {TemplateToolsSelectionError} If the answer is not a container of + * names or tools. The run reports `TEMPLATE_TOOLS_PROVIDER_ERROR`. + */ +export function resolveTemplateToolSelection( + selection: TemplateToolSelection, + templateIndex: ReadonlyMap, + log: Logger = DEFAULT_LOGGER, +): Set | null { + if (selection == null) return null; + + if (typeof selection === "string") { + throw new TemplateToolsSelectionError( + "templateToolsProvider returned a single string, which iterates one " + + "character at a time and would deny every tool. Return a container of " + + 'the tool names or tools this request may see, such as ["a_tool"]', + ); + } + if (selection instanceof Map) { + throw new TemplateToolsSelectionError( + "templateToolsProvider returned a Map, whose keys would be read as the " + + "allow-list while its values went unread, so a name mapped to false " + + "would still be allowed. Return a container holding only the tool " + + "names or tools this request may see", + ); + } + if (typeof (selection as Iterable)[Symbol.iterator] !== "function") { + throw new TemplateToolsSelectionError( + `templateToolsProvider returned ${describe(selection)}, which is not a ` + + "container of tool names or tools", + ); + } + + // Materialized here on purpose: a generator constructs without running its + // body, so a provider can hand back something that throws only on first + // iteration, and the caller guards this call. + const entries = Array.from(selection as Iterable); + + const allowed = new Set(); + for (const entry of entries) { + const name = + typeof entry === "string" + ? entry + : (entry as { name?: unknown } | null | undefined)?.name; + if (typeof name !== "string" || name.length === 0) { + log.warn( + `${LOG_PREFIX} templateToolsProvider returned an entry that names no tool`, + ); + continue; + } + if (!templateIndex.has(name)) { + log.warn( + `${LOG_PREFIX} templateToolsProvider named "${name}", which the template ` + + "agent does not contribute; it stays unavailable. This hook filters " + + "the template's tools and cannot add one.", + ); + continue; + } + allowed.add(name); + } + return allowed; +} + +function describe(value: unknown): string { + if (value === null) return "null"; + if (typeof value !== "object") return typeof value; + return (value as object).constructor?.name ?? "an object"; +} + +/** + * Tool names in the batch a live interrupt checkpoint would resume. + * + * A parked run resumes into the tool batch it stopped inside: Strands + * re-dispatches every `toolUse` in the assistant message it checkpointed, + * answering the ones that already completed from the checkpoint and running the + * one that is waiting. A tool absent from the registry at that moment turns the + * human's answer into a "tool not found" the model then re-fires, so nothing in + * that batch is filtered out while the pause is open. + * + * This is the same rule `syncProxyTools` applies to a proxy parked in a + * frontend-tool interrupt, read off the checkpoint instead of off a + * frontend-wait index because a template tool can park through the approval + * hook, through an interrupt of its own, or not at all, and the batch answers + * all three at once. + * + * Returns the names to hold registered; an empty set when nothing is parked; or + * {@link EXEMPT_EVERY_TEMPLATE_TOOL} for a checkpoint that is carrying a tool + * batch this function cannot read, where holding everything costs one + * unfiltered turn and the alternative breaks a resume. An activated checkpoint + * with no pending tool execution at all is not that case: an interrupt raised + * before any tool ran parks exactly that way, and it has no batch to protect. + */ +export function parkedBatchToolNames(agent: unknown): TemplateToolExemption { + const state = (agent as { _interruptState?: unknown } | null | undefined) + ?._interruptState as + | { + activated?: unknown; + pendingToolExecution?: unknown; + } + | undefined; + if (!state || state.activated !== true) return new Set(); + + const pending = state.pendingToolExecution; + if (pending == null) { + // A pause raised before any tool ran. Nothing is mid-dispatch, so nothing + // needs holding. + return new Set(); + } + + const message = (pending as { assistantMessageData?: unknown }) + .assistantMessageData as { content?: unknown } | undefined; + const names = new Set(); + if ( + message && + typeof message === "object" && + Array.isArray(message.content) + ) { + for (const block of message.content) { + const toolUse = (block as { toolUse?: unknown } | null | undefined) + ?.toolUse; + const name = (toolUse as { name?: unknown } | null | undefined)?.name; + if (typeof name === "string" && name.length > 0) names.add(name); + } + } + if (names.size === 0) { + DEFAULT_LOGGER.warn( + `${LOG_PREFIX} an activated interrupt checkpoint carries a tool batch ` + + "this adapter cannot read; holding every template tool registered " + + "rather than risk removing one this thread's resume is about to " + + "re-dispatch.", + ); + return EXEMPT_EVERY_TEMPLATE_TOOL; + } + return names; +} + +/** + * Whether a registry entry belongs to a producer other than the template. + * + * The adapter has exactly two others: proxies for client-declared tools, and + * the A2UI tool it injects itself. Both are re-decided every turn by the code + * that owns them, so neither is this filter's to touch. + */ +function isForeignEntry(entry: unknown): boolean { + return isProxyTool(entry) || isAutoInjectedA2UITool(entry); +} + +/** + * Whether a registry entry under a template tool's name is the template's. + * + * Identity settles it when it holds. It does not always hold: with an external + * `agentsByThread` map the wrapper is rebuilt per request while the cached + * thread agent keeps the registry it already had, so a template whose tools are + * built per request (a factory, or a closure over a request-scoped handle) + * hands the adapter equivalent but not identical objects. Reading a non-match + * as "someone else owns this name" would make a deny-everything answer remove + * nothing, which is a silent failure in the permissive direction on a hook + * whose whole job is withholding capability. + * + * So the fallback is ownership by elimination: the name is one the template + * contributes, and the entry sitting on it is not one of the adapter's other + * producers, therefore it is the template's. + */ +function isTemplateEntry(entry: unknown, templateTool: Tool): boolean { + if (entry === templateTool) return true; + return !isForeignEntry(entry); +} + +/** + * Make `toolRegistry` hold exactly the template tools `allowed` permits. + * + * `allowed` is a resolved name set, `null` meaning no filtering. `exemptNames` + * holds names to keep registered whatever `allowed` says; omitted exempts + * nothing and {@link EXEMPT_EVERY_TEMPLATE_TOOL} exempts every name. + * + * Removal is not destructive. The template tool objects outlive the registry + * entry, so a later request that allows a name again restores the same + * instance, and history stays untouched throughout: a filtered-out tool's + * earlier calls and results remain in the thread's messages, which is what lets + * the model read what it already did with a tool it can no longer call. + * + * Returns the template tool names the registry holds after the call. + */ +export function applyTemplateToolSelection( + toolRegistry: StrandsToolRegistry, + templateTools: readonly unknown[], + allowed: Set | null, + options: { exemptNames?: TemplateToolExemption; log?: Logger } = {}, +): Set { + const log = options.log ?? DEFAULT_LOGGER; + const templateIndex = indexTemplateTools(templateTools); + const exemptAll = options.exemptNames === EXEMPT_EVERY_TEMPLATE_TOOL; + const exemptSet: ReadonlySet = + exemptAll || options.exemptNames === undefined + ? new Set() + : (options.exemptNames as ReadonlySet); + + const registered = new Set(); + for (const [name, tool] of templateIndex) { + const exempt = exemptAll || exemptSet.has(name); + const keep = allowed === null || allowed.has(name) || exempt; + const existing = toolRegistry.get(name); + + if (!keep) { + if (existing !== undefined && isTemplateEntry(existing, tool)) { + toolRegistry.remove(name); + log.debug(`${LOG_PREFIX} Filtered out template tool: ${name}`); + } else if (existing !== undefined) { + log.debug( + `${LOG_PREFIX} Template tool ${name} is held by another producer; ` + + "the filter leaves it in place", + ); + } + continue; + } + + if (existing === tool) { + registered.add(name); + continue; + } + if (existing === undefined) { + toolRegistry.add(tool); + registered.add(name); + log.debug(`${LOG_PREFIX} Restored template tool: ${name}`); + continue; + } + if (!exempt && isProxyTool(existing)) { + // A client proxy took this name while the template tool was filtered + // out, and the provider now allows the template tool. A native tool wins + // a name collision, so hand the name back: the proxy sync runs after this + // and re-decides the client's side, skipping a name a native tool holds. + // Leaving the proxy would both shadow the allowed tool and, if the client + // has stopped declaring it, let the proxy sync drop the name outright. + // + // Guarded on `exempt` rather than on the selection, because the keep + // branch is also reached by exemption, and a proxy the parked batch is + // answering keeps its name. + toolRegistry.remove(name); + toolRegistry.add(tool); + registered.add(name); + log.debug( + `${LOG_PREFIX} Reclaimed template tool ${name} from a client proxy ` + + "holding its name", + ); + continue; + } + if (isTemplateEntry(existing, tool)) { + // The template's, under a different object. Leave the entry the thread + // has been using rather than churn it. + registered.add(name); + continue; + } + log.debug( + `${LOG_PREFIX} Template tool ${name} is held by another producer; ` + + "leaving it in place", + ); + } + return registered; +} + +/** + * Read `selection` and apply it to `toolRegistry` in one call. + * + * The run path keeps the two halves apart, so that reading a provider's answer + * fails as a provider error and applying it does not. This composes them for a + * caller with an answer already in hand. + */ +export function syncTemplateTools( + toolRegistry: StrandsToolRegistry, + templateTools: readonly unknown[], + selection: TemplateToolSelection, + options: { exemptNames?: TemplateToolExemption; log?: Logger } = {}, +): Set { + const log = options.log ?? DEFAULT_LOGGER; + return applyTemplateToolSelection( + toolRegistry, + templateTools, + resolveTemplateToolSelection( + selection, + indexTemplateTools(templateTools), + log, + ), + options, + ); +} + +/** Publish the run's narrowed name set for the re-narrowing hook to read. */ +export function recordTemplateToolSelection( + agent: unknown, + allowed: Set | null, +): void { + (agent as Record)[ALLOWED_KEY] = allowed; +} + +/** + * Re-narrow the filtered set once a resumed batch has been dispatched. + * + * The parked-batch exemption keeps a denied tool registered so a resume can + * reach it. Strands then carries on inside the same run: it re-dispatches the + * batch, clears the checkpoint, and makes its next model call from this same + * registry, which would still be advertising what the request denied. Without + * this the model could call a withheld tool for the rest of that run, and only + * the next request would narrow again. + * + * Strands reads the registry fresh for every model call, so this is called from + * a hook that fires between one tool batch and the next such read. Which hook + * that is differs by SDK, and the adapter picks it; see the call site. + * + * No exemption is passed. By the time this fires the parked batch has been + * dispatched, which is the whole reason the exemption existed, so re-reading + * the checkpoint would only ask a question whose answer no longer matters, and + * the answer moved between releases: this SDK clears its pending execution + * before the tool step in 1.1 and after it in 1.16, so a hook that reads it + * holds the exemption on one release and drops it on the other. Asking nothing + * is both simpler and the same on every release. + * + * A run the SDK is cancelling can replay a skipped batch after this fires, and + * a tool this narrowing removed would be missing from that replay. The run is + * being torn down at that point, and the next request restores whatever a live + * checkpoint still needs before anything reads the registry again. + */ +export function renarrowTemplateTools( + agent: unknown, + templateTools: readonly unknown[], + log: Logger = DEFAULT_LOGGER, +): void { + const allowed = (agent as Record)[ALLOWED_KEY]; + if (allowed == null) return; + applyTemplateToolSelection( + (agent as { toolRegistry: StrandsToolRegistry }).toolRegistry, + templateTools, + allowed as Set, + { log }, + ); +}