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 8d1f2d234c..fd03db145f 100644 --- a/integrations/aws-strands/python/src/ag_ui_strands/agent.py +++ b/integrations/aws-strands/python/src/ag_ui_strands/agent.py @@ -67,6 +67,11 @@ def _extract_agent_kwargs(agent: StrandsAgentCore) -> dict: # outstanding frontend calls at once. _WIRE_MAP_MAX = 512 +# Upper bound on the per-agent tool-call metadata map held in session state. +# It bounds abandoned entries (tool calls whose result never returns) +# so state cannot grow without bound. +_TOOL_CALL_MAP_MAX = 512 + # Sentinel handed back to a paused ``tool_context.interrupt()`` when the client # cancels (``ResumeEntry.status == "cancelled"``) rather than resolving. The # tool receives this in place of a real answer and can treat it as a denial. @@ -205,6 +210,7 @@ def _extract_interrupts(agent: Any, terminal_result: Any) -> list: ) from .client_proxy_tool import sync_proxy_tools from .session_reconcile import ( + AG_UI_TOOL_CALL_MAP_STATE_KEY, AG_UI_WIRE_MAP_STATE_KEY, has_placeholder_results, reconcile_frontend_tool_results, @@ -1053,6 +1059,10 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]: # client dispatching its follow-up run before the backend results # reach it, narrowing the ConcurrencyException race window. deferred_frontend_tool_ends = [] + # Native ``toolUseId``s whose ``toolResult`` was processed this + # run. Drained after each result batch to prune the persisted + # tool-call meta map. + processed_result_native_ids: set[str] = set() # Terminal ``AgentResult`` from Strands (carried on the final # ``{"result": ...}`` stream event). Used after the loop to detect a # native interrupt pause (``stop_reason == "interrupt"``). @@ -1081,6 +1091,23 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]: wire_to_native = ( strands_agent.state.get(AG_UI_WIRE_MAP_STATE_KEY) or {} ) + + # The durable per-``toolUseId`` call metadata map recorded at + # emission (see the ``current_tool_use`` handler). On a RESUME + # run this is the ONLY source of ``{name, args, input, + # strands_tool_id}`` for the interrupted tool, since Strands does + # not re-emit ``current_tool_use`` events for it. Guarded because + # test doubles / stub agents may lack ``state`` entirely; a + # missing store just means "no persisted meta yet". + persisted_tool_call_meta: Dict[str, Dict[str, Any]] = {} + _agent_state = getattr(strands_agent, "state", None) + if _agent_state is not None: + try: + persisted_tool_call_meta = ( + _agent_state.get(AG_UI_TOOL_CALL_MAP_STATE_KEY) or {} + ) + except Exception: + persisted_tool_call_meta = {} # Scope to the TRAILING tool results (this continuation's just- # returned results). ``pending_tool_result_ids`` holds those ids; # without this, a multi-turn continuation re-sends already-reconciled @@ -1536,6 +1563,35 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]: if _data.get("strands_tool_id") == result_tool_id: call_info = _data break + # RESUME-run fallback: the interrupted tool never + # re-emits ``current_tool_use`` on resume, so + # ``tool_calls_seen`` is empty for it. The + # persisted meta map was populated when the call + # was originally streamed (possibly in a prior + # process). Direct native-id first, then scan by + # ``strands_tool_id`` to match the frontend-tool + # case. + if not call_info: + call_info = persisted_tool_call_meta.get( + result_tool_id, {} + ) + if not call_info: + for _pdata in persisted_tool_call_meta.values(): + if ( + isinstance(_pdata, dict) + and _pdata.get("strands_tool_id") + == result_tool_id + ): + call_info = _pdata + break + # Record consumption once the lookup is complete + # (even if it missed): the result was processed + # this turn, so any persisted entry keyed on this + # native id is safe to prune. Recording BEFORE the + # frontend-skip / behavior branches ensures a + # ``stop_streaming_after_result`` early break still + # flags this id for prune. + processed_result_native_ids.add(result_tool_id) tool_name = call_info.get("name") tool_args = call_info.get("args") tool_input = call_info.get("input") @@ -1667,6 +1723,33 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]: # Break inner loop — no further results should be emitted break + # Prune the persisted tool-call meta map for entries + # whose native id (or ``strands_tool_id`` for frontend + # tools stored under a wire key) was just consumed. + # The emission-time size cap (``_TOOL_CALL_MAP_MAX``) is + # only a backstop for abandoned entries. + if ( + persisted_tool_call_meta + and processed_result_native_ids + ): + _remaining = { + _k: _v + for _k, _v in persisted_tool_call_meta.items() + if _k not in processed_result_native_ids + and ( + not isinstance(_v, dict) + or _v.get("strands_tool_id") + not in processed_result_native_ids + ) + } + if len(_remaining) != len(persisted_tool_call_meta): + if _get_strands_session_manager(strands_agent): + strands_agent.state.set( + AG_UI_TOOL_CALL_MAP_STATE_KEY, _remaining + ) + persisted_tool_call_meta = _remaining + processed_result_native_ids.clear() + # Defer hand-off: now that this turn's backend # TOOL_CALL_RESULT(s) have been emitted above, flush the # buffered frontend-tool ToolCallEnd(s). Flushing here — @@ -1811,6 +1894,42 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]: "strands_tool_id": strands_tool_id, } + # Mirror the minimum-sufficient subset into durable + # session state so a RESUME run — which does not + # re-emit ``current_tool_use`` for the interrupted + # tool — can still resolve ``tool_name``/behavior/ + # context at the ``toolResult`` site. Gate on + # session_manager: only then does Strands durably flush state. + if _get_strands_session_manager(strands_agent): + _tc_meta = dict( + strands_agent.state.get(AG_UI_TOOL_CALL_MAP_STATE_KEY) + or {} + ) + # Key by the NATIVE ``toolUseId`` — that is what + # arrives on ``toolResult``. For backend tools + # this equals ``tool_use_id``; for frontend + # tools ``tool_use_id`` is a fresh wire UUID + # while ``strands_tool_id`` is native. + _tc_key = strands_tool_id or tool_use_id + _tc_meta[_tc_key] = { + "name": tool_name, + "args": args_str, + "input": tool_input, + "strands_tool_id": strands_tool_id, + } + if len(_tc_meta) > _TOOL_CALL_MAP_MAX: + for _stale in list(_tc_meta)[ + : len(_tc_meta) - _TOOL_CALL_MAP_MAX + ]: + _tc_meta.pop(_stale, None) + strands_agent.state.set( + AG_UI_TOOL_CALL_MAP_STATE_KEY, _tc_meta + ) + # Keep the in-run view aligned so downstream + # result lookups see the same entry a fresh + # process would restore from the store. + persisted_tool_call_meta = _tc_meta + if use_streaming: # Close any open assistant text turn so the # snapshot order matches the wire-event order @@ -1876,6 +1995,25 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[Any]: tool_calls_seen[tool_use_id]["args"] = args_str tool_calls_seen[tool_use_id]["raw"] = raw_str + # Keep the persisted meta in sync with the final + # streamed args. Without this refresh, resume runs + # would see the first partial-JSON delta rather + # than the complete args the model emitted. + if _get_strands_session_manager(strands_agent): + _tc_meta = dict( + strands_agent.state.get(AG_UI_TOOL_CALL_MAP_STATE_KEY) + or {} + ) + _tc_key = strands_tool_id or tool_use_id + _existing = _tc_meta.get(_tc_key) + if _existing is not None: + _existing["input"] = tool_input + _existing["args"] = args_str + strands_agent.state.set( + AG_UI_TOOL_CALL_MAP_STATE_KEY, _tc_meta + ) + persisted_tool_call_meta = _tc_meta + # Stream incremental ToolCallArgs deltas as the LLM # produces more characters of the JSON args. The FE # uses these to drive predictive state updates per the diff --git a/integrations/aws-strands/python/src/ag_ui_strands/session_reconcile.py b/integrations/aws-strands/python/src/ag_ui_strands/session_reconcile.py index 79c79b4cb3..4f53e485ee 100644 --- a/integrations/aws-strands/python/src/ag_ui_strands/session_reconcile.py +++ b/integrations/aws-strands/python/src/ag_ui_strands/session_reconcile.py @@ -23,6 +23,16 @@ # user-managed state keys. AG_UI_WIRE_MAP_STATE_KEY = "__ag_ui_wire_to_native__" +# Key under which the adapter stores every ``toolUseId`` tool call metadata +# (name, args, input, strands_tool_id) on the Strands agent's session state. +# On a native-interrupt RESUME run Strands does not re-invoke the model for the +# interrupted tool, so no ``current_tool_use`` events fire and the in-run +# ``tool_calls_seen`` dict is empty when the ``toolResult`` arrives. Reading +# from this durable map at that point restores ``tool_name`` (and thus every +# ``tool_behaviors`` gate + the frontend-placeholder skip) for the resumed +# tool. Namespaced to avoid clashing with user-managed state keys. +AG_UI_TOOL_CALL_MAP_STATE_KEY = "__ag_ui_tool_call_map__" + def resolve_native_ids( wire_to_native: Mapping[str, str],