feat(aws-strands): forward RunAgentInput.context to the model in the TypeScript bridge - #2619
Conversation
Python Preview PackagesVersion
Install with uvAdd the TestPyPI index to your [[tool.uv.index]]
name = "testpypi"
url = "https://test.pypi.org/simple/"
explicit = trueThen install the packages you need: # Core SDK
uv add 'ag-ui-protocol==0.0.0.dev1788528809' --index testpypi
# Integrations (each already depends on the matching ag-ui-protocol preview)
uv add 'ag-ui-langgraph==0.0.0.dev1788528809' --index testpypi
uv add 'ag-ui-crewai==0.0.0.dev1788528809' --index testpypi
# NOTE: ag-ui-agent-spec depends on pyagentspec (git-only, not on PyPI).
# You will need to install pyagentspec separately from its git repo.
uv add 'ag-ui-agent-spec==0.0.0.dev1788528809' --index testpypi
uv add 'ag_ui_adk==0.0.0.dev1788528809' --index testpypi
uv add 'ag_ui_strands==0.0.0.dev1788528809' --index testpypiInstall with pippip install \
--index-url https://test.pypi.org/simple/ \
--extra-index-url https://pypi.org/simple/ \
ag-ui-protocol==0.0.0.dev1788528809
Commit: dff84e2 |
@ag-ui/a2a-middleware
@ag-ui/a2ui-middleware
@ag-ui/event-throttle-middleware
@ag-ui/mcp-apps-middleware
@ag-ui/mcp-middleware
@ag-ui/a2a
@ag-ui/adk
@ag-ui/ag2
@ag-ui/agno
@ag-ui/aws-strands
@ag-ui/claude-agent-sdk
@ag-ui/claude-managed-agents
@ag-ui/crewai
@ag-ui/langchain
@ag-ui/langgraph
@ag-ui/llamaindex
@ag-ui/mastra
@ag-ui/pydantic-ai
@ag-ui/vercel-ai-sdk
@ag-ui/watsonx
@ag-ui/a2ui-toolkit
create-ag-ui-app
@ag-ui/client
@ag-ui/core
@ag-ui/encoder
@ag-ui/proto
commit: |
757b070 to
bf0ed39
Compare
mme
left a comment
mme
left a comment
There was a problem hiding this comment.
Nice work on this, the design is really carefully thought through and the parity with the Python side is easy to follow. I ran a codex review over the diff and then checked its one finding against the Strands SDK source, and I think it's real. Details inline, it's a small reorder in two places.
A few things I checked independently that all hold up: the SDK guarantees the before/after model-call pair even when the model throws, so the "not restored" guard can't misfire on a retry; the model reads a clone of agent.messages after the before-hook runs, so the splice is visible; and the AsyncLocalStorage scoping around each pull does cover the hooks.
| // the block it was shown would otherwise stay in `agent.messages` and | ||
| // be the next thing persisted. Idempotent when the hook already ran. | ||
| // Python restores from the same teardown `finally`. | ||
| restoreTransientModelContext(strandsAgent); |
There was a problem hiding this comment.
I think this restore needs to happen before agentStream.return() rather than after.
Returning the Strands stream isn't a silent close. In SDK 1.8.0 the stream's own finally runs the AfterInvocationEvent callbacks even on consumer break (the comment in stream() says exactly that), and the session manager's default trigger is saveLatestOn: 'invocation', so it writes a full snapshot of agent.messages right there, while the context block is still in the list. The in-memory removal on this line runs afterwards and can't undo a snapshot that's already in storage, so the next restore of that session replays the app context as a real user turn.
1.1.0 (what the lockfile pins) is a bit different: it yields the AfterInvocationEvent from inside a finally, but it also appends a "Cancelled by user" message from that same finally, which fires MessageAddedEvent and saves under the per-message strategy. The package accepts >=1.1.0 so users will hit 1.8.0 behaviour anyway.
Moving restoreTransientModelContext(strandsAgent) above the .return() is safe: the after-hook doesn't fire on consumer break, and the restore is idempotent regardless. The cancellation case in context-forwarding.test.ts would be a natural place to assert on what a session manager sees at AfterInvocationEvent.
| // A leaf abandoned mid-model-call never reached its after-hook; take | ||
| // the block back off every hooked leaf, as Python's | ||
| // `_restore_orchestrator_context` does from the same teardown. | ||
| restoreOrchestratorContext(this._orchestrator); |
There was a problem hiding this comment.
Same ordering point as the single-agent teardown above: each leaf's stream runs its after-invocation callbacks during orchestratorStream.return(), so this should go before the .return() so a session manager on a leaf doesn't snapshot the block.
…TypeScript bridge The Python bridge already renders the application's context entries as a text block and shows that block to the model for exactly one call, without ever persisting it. The TypeScript bridge handed the same context to tools, hooks and the A2UI subagent, but never to the model, so anything the app published through useCopilotReadable was invisible to the LLM. This brings the TypeScript bridge to the same observable behaviour: - The block is rendered byte for byte as Python renders it, with the A2UI component-schema entry excluded through the shared toolkit split. - A BeforeModelCallEvent hook splices the block into agent.messages immediately before the latest user turn, into that turn's content when it carries a tool result, or appended when there is no user turn at all. The paired AfterModelCallEvent hook splices it back out, and the run-loop teardown restores again for the cancellation path. - The block is carried in an AsyncLocalStorage set around each pull of the Strands stream, the direct analogue of Python's ContextVar, so concurrent runs cannot see each other's context. - The orchestrator path installs the hook on every reachable leaf agent and falls back to prefixing the prompt when none is reachable, as Python does. - The middleware's usage guide for a render tool that A2UI auto-injection replaced is dropped from both the model block and the recovery subagent's context, closing a prompt-parity gap the TypeScript bridge had. A run that carries context on an agent with no hook registry ends with a RUN_ERROR, matching Python. The shared scripted stub in the test helpers gains an inert addHook so existing stub-driven tests keep their behaviour.
bf0ed39 to
a32b3d3
Compare
2 participants
Summary
Brings the TypeScript AWS Strands bridge to parity with the Python bridge on forwarding
RunAgentInput.contextto the model.Until now the TypeScript bridge handed context to tools and hooks (
buildContextExtras) and to the A2UI subagent, but nothing put it in front of the model, so whatever the application published throughuseCopilotReadablewas invisible to the LLM. Python has done this for a while; this PR implements the same contract in TypeScript, matching Python's observable behaviour rather than its plumbing.What the model sees
The context entries are rendered as one text block:
The A2UI component-schema entry is excluded through the shared toolkit's
splitA2UISchemaContext, because thegenerate_a2uipath already carries it. The block is shown to the model for exactly one call and is never persisted: it does not reachagent.messagesafter the call, aMESSAGES_SNAPSHOT, or the session store.Placement mirrors Python case for case:
toolUseand itstoolResultis a shape providers reject);How it is carried
model-context.ts(new): block rendering, the request-scoped carrier, theBeforeModelCallEvent/AfterModelCallEventhook pair, restore, and the orchestrator walkers.AsyncLocalStorage, the Node analogue of Python'sContextVar. It is entered around each pull of the Strands stream (pullWithModelContext), exactly as Python's_stream_with_model_contextsets and resets its token around each__anext__, so two runs in flight in one process cannot see each other's block. Covered by an interleaved two-thread test.WeakMapand undone from the after-hook on the normal path and from both run loops' teardownfinallyon the cancellation path, where the after-hook never fires. The teardown restore runs before the stream is returned, because returning the Strands stream runs the after-invocation hooks and a session manager snapshotsagent.messagesfrom there; a restore after the drain would come too late for the store. A session-manager cancellation test covers that ordering. The restore runs before the session manager can save: the SDK runsAfterModelCallEventcallbacks in reverse registration order, the bridge registers its pair after the session manager (at first run rather than at construction), and the assistant message is only appended, and so only persisted, after the after-hooks have returned.RUN_ERROR, as in Python.AgentNode.agent, recursing throughMultiAgentNode.orchestrator, depth-capped like_MAX_MULTIAGENT_NESTING). An orchestrator with no reachable leaf gets the block prefixed onto its prompt instead, which is Python's fallback too.Also ported: A2UI render-guide exclusion
While comparing prompts I found a pre-existing drift in the same area. Once A2UI auto-injection replaces the middleware's
render_a2uiproxy withgenerate_a2ui, Python drops the middleware's usage guide for that proxy from both the model context and the recovery subagent's context (_without_a2ui_render_guides). TypeScript had no equivalent. This PR addswithoutA2UIRenderGuidesand applies it at both sites, so the two bridges produce the same prompt with A2UI on as well as off. TheplanA2UIInjectionstate build is extracted intobuildA2UISubagentStatefor that.Where the two bridges still cannot produce the same prompt
Context.valueas a string, so these only arrive from lax callers.JSON.stringifyandjson.dumpsagree on strings, integers, booleans andnull. They disagree on objects and arrays (json.dumpsinserts a space after,and:), integral floats (1.0against1) and non-ASCII text (json.dumpsescapes it). Left to differ rather than emulating Python's serializer in TypeScript; noted onformatAguiContext.contentlist and swaps the original back; TypeScript unshifts aTextBlockinto the existing array and removes it by identity. The model sees the same bytes.Tests
context-forwarding.test.ts(24 tests) mirrorstests/test_context_forwarding.pycase for case at the real model boundary (a scripted model records the exact history the SDK handed it), then covers the arms Python leaves to other suites: the tool-result turn, cancellation mid-model-call (in memory and at the session manager's save point), two interleaved runs, the orchestrator leaf path and its prompt fallback, the refusal on an agent with no hook registry, and the render-guide exclusion.integrations/aws-strands/typescript:pnpm test77 files, 1635 tests passed;pnpm typecheckclean;pnpm buildclean.integrations/aws-strands/python:pytest tests/ -q1270 passed, no Python file touched.ARCHITECTURE.mdgains a "Model context" bullet that applies to both bridges and records divergence 2.