Merge pull request #2616 from ag-ui-protocol/claude/objective-cohen-fb1f7d

feat(aws-strands): filter the template agent's tools per request

Author: ranst91

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

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"],

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

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",

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

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.