feat(ui): add UIAdapter.sanitize_messages and allowed_file_url_schemes#5228
feat(ui): add UIAdapter.sanitize_messages and allowed_file_url_schemes#5228
UIAdapter.sanitize_messages and allowed_file_url_schemes#5228Conversation
`UIAdapter.sanitize_messages()` is applied to protocol-sourced messages in
`run_stream_native` before passing them to the agent. Caller-supplied
`message_history` is unaffected (trusted as coming from server persistence).
Two sanitizations:
- `allowed_file_url_schemes: frozenset[str] = frozenset({'http', 'https'})`
drops `FileUrl` parts with disallowed schemes. Non-HTTP schemes like
`s3://` and `gs://` cause the provider to fetch the object using the
server-side IAM role or service account, so they should only be accepted
from trusted frontends. Apps that want to accept a scheme can add it to
the allowlist on the adapter.
- Dangling `ToolCallPart`s at the end of history are dropped unless matched
by `deferred_tool_results`. Tool calls are produced by the model on the
server side, so an unresolved tool call at the end of client-supplied
history does not correspond to a paused agent run. HITL resumption
continues to work when the caller passes explicit `deferred_tool_results`
(approvals / call results) matching the pending tool call IDs.
Both behaviors emit a `UserWarning` when they kick in, matching the
pattern introduced in #4087 for frontend `SystemPromptPart`s.
Adapter trust model is documented in `docs/ui/overview.md`, referenced
from both adapter pages, with a matching note in `docs/input.md` for the
non-adapter `FileUrl` case.
…sage-sanitization
github-actions
Bot
added
size: M
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| dangling_tool_call_names.append(part.tool_name) | ||
| continue | ||
| new_response_parts.append(part) | ||
| sanitized.append(replace(message, parts=new_response_parts)) |
There was a problem hiding this comment.
When all parts of the last ModelResponse are dangling tool calls (e.g. a client-submitted AssistantMessage with only fabricated tool calls and no text), this produces a ModelResponse(parts=[]) in the sanitized output. An empty response in the history could confuse downstream logic or the model.
Consider dropping the message entirely when no parts survive:
if new_response_parts:
sanitized.append(replace(message, parts=new_response_parts))The AG-UI test test_client_submitted_dangling_tool_calls_not_executed exercises exactly this scenario (an AssistantMessage with only a tool call and no text content) but doesn't assert on the resulting message structure, so the empty ModelResponse slips through.
| elif isinstance(message, ModelResponse) and index == last_index: | ||
| new_response_parts: list[ModelResponsePart] = [] | ||
| for part in message.parts: | ||
| if isinstance(part, ToolCallPart) and part.tool_call_id not in resolved_tool_call_ids: |
There was a problem hiding this comment.
The dangling tool call check only covers ToolCallPart, but BuiltinToolCallPart is a separate type and isn't filtered here. In AG-UI, a client can craft an AssistantMessage with the BUILTIN_TOOL_CALL_ID_PREFIX in the tool call ID, which load_messages will load as a BuiltinToolCallPart. If that fabricated builtin call has no matching BuiltinToolReturnPart in the same response, it would survive sanitization.
In practice this may be low-risk since builtin tool calls are typically paired with returns in the same ModelResponse, but it's worth considering whether the isinstance check should be isinstance(part, (ToolCallPart, BuiltinToolCallPart)) for consistency with the stated goal of stripping tool calls "produced by the model on the server side".
Docs Preview
|
- Match the dangling tool-call check on `BaseToolCallPart` so `BuiltinToolCallPart` is treated the same as `ToolCallPart`. - Drop the trailing `ModelResponse` entirely when sanitization leaves it with no parts, instead of leaving an empty placeholder in history.
Restructure the AG-UI and base-adapter end-to-end tests for dangling tool-call sanitization to assert via the captured model history rather than via a registered tool whose body never runs by design (the whole point of the test). Removes the unreachable lines that broke coverage.
|
|
||
| For stricter conversation integrity (e.g. ensuring prior assistant turns and tool returns match what the server actually produced), persist the history server-side keyed by the thread/session ID and pass it to the adapter via `message_history` — caller-supplied history is trusted as coming from server-side persistence and isn't subject to this sanitization. | ||
|
|
||
| Tool approval responses in the Vercel AI protocol are trusted from the request by design, matching the protocol's round-trip through `useChat`'s `addToolApprovalResponse` and the reference Next.js backend. If your application needs the approval decision tied to server-side state rather than the request, intercept [`DeferredToolRequests`][pydantic_ai.DeferredToolRequests], persist the approval IDs server-side, and pass explicit `deferred_tool_results` when resuming. |
There was a problem hiding this comment.
Move to Vercel AI docs
| deferred_tool_results = self.deferred_tool_results | ||
|
|
||
| frontend_messages = self.sanitize_messages(self.messages, deferred_tool_results=deferred_tool_results) | ||
| if self.manage_system_prompt == 'server' and any( |
There was a problem hiding this comment.
Would it make sense to move this to sanitize_messages?
There was a problem hiding this comment.
Where is the stripping implemented again?
- Move `SystemPromptPart` stripping (and the accompanying warning) into `sanitize_messages`, so the message claiming "they will be stripped" is truthful at the call site instead of relying on `ReinjectSystemPrompt` to strip later. The capability still runs with `replace_existing=True` to defend against caller-supplied `message_history` containing system prompts. - Move the Vercel-AI-specific paragraph about trusting tool approval responses out of `docs/ui/overview.md` into `docs/ui/vercel-ai.md`, next to the rest of the approval flow documentation. - Factor the per-message-kind sanitization into helpers to keep `sanitize_messages` under the cyclomatic complexity threshold.
Moving `SystemPromptPart` stripping into `sanitize_messages` left the "all parts stripped" branch in `ReinjectSystemPrompt._strip_system_prompts` unreachable from the UI tests that previously covered it, since the adapter now strips before the capability runs. Add a focused test that exercises the capability directly with a caller-supplied `message_history` containing a `ModelRequest` whose only part is a `SystemPromptPart`.
![devin-ai-integration[bot]](https://avatars.githubusercontent.com/in/811515?s=60&v=4){kind=small}
| !!! warning "Trust model for file URLs" | ||
| When URLs are forwarded to the provider, the provider fetches them under its own credentials. For cloud-storage schemes like `s3://` (Bedrock) and `gs://` (Vertex AI), those credentials are your server's IAM role or service account, so whoever controls the URL effectively controls what the provider can read on your behalf. | ||
|
|
||
| Don't construct [`ImageUrl`][pydantic_ai.messages.ImageUrl], [`AudioUrl`][pydantic_ai.messages.AudioUrl], [`VideoUrl`][pydantic_ai.messages.VideoUrl], or [`DocumentUrl`][pydantic_ai.messages.DocumentUrl] from untrusted user input without validating the scheme and scope. For frontend-initiated uploads to cloud storage, prefer pre-signed `https://` URLs over references like `s3://bucket/key`. If you do need to accept the raw cloud-storage scheme, set `force_download=True` to route the fetch through the library's HTTP client (which applies SSRF protection). |
There was a problem hiding this comment.
🚩 Misleading documentation advice: force_download=True with s3:// URLs
The new warning box at docs/input.md:172 says: "If you do need to accept the raw cloud-storage scheme, set force_download=True to route the fetch through the library's HTTP client (which applies SSRF protection)." However, download_item at pydantic_ai_slim/pydantic_ai/models/__init__.py:1395-1396 explicitly states "Only http:// and https:// protocols are allowed" and calls safe_download which rejects non-HTTP schemes. So force_download=True on an s3:// URL would raise a ValueError, not download the file. The phrasing starting with "If you do need to accept the raw cloud-storage scheme" implies the user wants it to work, but the suggested action would fail. This could confuse users. A more accurate recommendation would be to convert the s3:// reference to a pre-signed https:// URL before constructing the FileUrl.
Was this helpful? React with 👍 or 👎 to provide feedback.
Devin caught that the trust-model warning suggested setting `force_download=True` on cloud-storage URLs, but `download_item` only accepts `http(s)://` schemes. The right path for frontend-initiated cloud uploads is to convert to a pre-signed `https://` URL server-side before constructing the `FileUrl` part.
1 participant
Summary
Adds
UIAdapter.sanitize_messages(), applied to messages produced from the protocol-specific run input before they're passed to the agent. Caller-suppliedmessage_historyis not affected (trusted as coming from server-side persistence).Two sanitizations land in this PR:
allowed_file_url_schemes: frozenset[str] = frozenset({'http', 'https'})dropsFileUrlparts (ImageUrl/DocumentUrl/VideoUrl/AudioUrl) whose URL scheme isn't in the allowlist. Non-HTTP schemes likes3://andgs://cause the model provider to fetch the object using the server-side IAM role or service account, so they should only be accepted from trusted frontends. Apps that want a scheme can add it to the allowlist on the adapter.ToolCallParts at the end of the client-submitted history are dropped unless matched bydeferred_tool_results. Tool calls are produced by the model on the server side, so an unresolved tool call at the end of client-supplied history doesn't correspond to a paused agent run. HITL resumption still works when the caller passes explicitdeferred_tool_resultsmatching the pending tool call IDs.Both behaviors emit a
UserWarningwhen they trigger, matching the pattern from #4087 for frontendSystemPromptParts.The adapter trust model is documented in a new section of
docs/ui/overview.md, referenced from both adapter pages, with a matching note indocs/input.mdfor the non-adapterFileUrlcase.Checklist