Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions integrations/aws-strands/ARCHITECTURE.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`. |
Expand Down Expand Up @@ -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
Expand Down
6 changes: 6 additions & 0 deletions integrations/aws-strands/error-codes.json
Original file line number Diff line number Diff line change
Expand Up @@ -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"],
Expand Down
93 changes: 93 additions & 0 deletions integrations/aws-strands/python/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 |

Expand Down Expand Up @@ -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
Expand Down
10 changes: 10 additions & 0 deletions integrations/aws-strands/python/src/ag_ui_strands/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -34,6 +39,7 @@
ToolStreamEventContext,
PredictStateMapping,
SessionManagerProvider,
TemplateToolsProvider,
ToolStreamEventHandler,
)
from ag_ui.core import (
Expand All @@ -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",
Expand All @@ -71,6 +80,7 @@
"ToolStreamEventContext",
"PredictStateMapping",
"SessionManagerProvider",
"TemplateToolsProvider",
"ToolStreamEventHandler",
"Interrupt",
"ResumeEntry",
Expand Down
75 changes: 75 additions & 0 deletions integrations/aws-strands/python/src/ag_ui_strands/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
76 changes: 76 additions & 0 deletions integrations/aws-strands/python/src/ag_ui_strands/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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.

Expand Down
Loading
Loading