Skip to content

Commit 5d11746

Browse files
Light2Darkclaude
Light2Dark
and
claude
authored
support tool approvals for pydantic-ai chatbot (#9621)
## 📝 Summary <!-- If this PR closes any issues, list them here by number (e.g., Closes #123). Detail the specific changes made in this pull request. Explain the problem addressed and how it was resolved. If applicable, provide before and after comparisons, screenshots, or any relevant details to help reviewers understand the changes easily. --> Supports more AI sdk parts, refactors logic to be more maintainable. - ChatMessage no longer loses unmodeled SDK fields. Approvals (and e.g. `callProviderMetadata`, `providerExecuted`, `preliminary`) live on AI SDK parts that marimo's typed dataclasses don't model. The message now snapshots the raw wire payload per-part in _raw_parts, so those fields survive every round-trip - `pydantic_ai._build_ui_messages` uses the raw payload. It now calls `message.raw_or_dumped_parts()` so the approval/tool state the frontend just sent us makes it into the agent run unmodified. The old `asdict` + `_remove_none_values` path was lossy. - `sanitize_part` strips keys the AI SDK's { ...part, state, ... } spread can leak from prior tool states (e.g. a stale output clinging to an approval-requested part). - hasPendingToolCalls (frontend) rewritten. The old predicate treated "every tool ready & no trailing text" silently looped whenever an assistant message ended in a non-text part (file, source-url, data-*, reasoning) after a completed tool call. The fix uses some native AI sdk logic and some testing. <img width="921" height="787" alt="image" src="https://github.com/user-attachments/assets/53f4a146-9554-4135-b9e9-459317a823dc" /> ## 📋 Pre-Review Checklist <!-- These checks need to be completed before a PR is reviewed --> - [x] For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on [Discord](https://marimo.io/discord?ref=pr), or the community [discussions](https://github.com/marimo-team/marimo/discussions) (Please provide a link if applicable). - [x] Any AI generated code has been reviewed line-by-line by the human PR author, who stands by it. - [x] Video or media evidence is provided for any visual changes (optional). <!-- PR is more likely to be merged if evidence is provided for changes made --> ## ✅ Merge Checklist - [x] I have read the [contributor guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md). - [ ] Documentation has been updated where applicable, including docstrings for API changes. - [x] Tests have been added for the changes made. --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
1 parent 65928cd commit 5d11746

11 files changed

Lines changed: 948 additions & 162 deletions

File tree

‎examples/ai/chat/pydantic-ai-chat.py‎

Lines changed: 273 additions & 42 deletions
Large diffs are not rendered by default.

‎frontend/src/components/chat/__tests__/chat-utils.test.ts‎

Lines changed: 269 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,269 @@
1+
/* Copyright 2026 Marimo. All rights reserved. */
2+
3+
import type { UIMessage } from "ai";
4+
import { describe, expect, it } from "vitest";
5+
import { hasPendingToolCalls } from "../chat-utils";
6+
7+
/**
8+
* `hasPendingToolCalls` powers `sendAutomaticallyWhen` in `mo.ui.chat`:
9+
* returns true only when the last assistant message *ends* with a tool
10+
* call in a ready-to-round-trip state. Any trailing non-tool part (text,
11+
* file, source-*, reasoning, data-*, new step-start) means the assistant
12+
* has already answered and we leave the next turn to the user. The
13+
* approval flow relies on this firing for `approval-responded`.
14+
*/
15+
16+
const userMessage = (text: string): UIMessage => ({
17+
id: `user-${text}`,
18+
role: "user",
19+
parts: [{ type: "text", text }],
20+
});
21+
22+
const assistantToolMessage = (
23+
parts: UIMessage["parts"],
24+
id = "assistant-1",
25+
): UIMessage => ({
26+
id,
27+
role: "assistant",
28+
parts,
29+
});
30+
31+
describe("hasPendingToolCalls", () => {
32+
it("returns false when there are no messages", () => {
33+
expect(hasPendingToolCalls([])).toBe(false);
34+
});
35+
36+
it("returns false when the last message is a user message", () => {
37+
expect(hasPendingToolCalls([userMessage("hi")])).toBe(false);
38+
});
39+
40+
it("returns false when the last assistant message has no tool parts", () => {
41+
expect(
42+
hasPendingToolCalls([
43+
userMessage("hi"),
44+
assistantToolMessage([{ type: "text", text: "hello!" }]),
45+
]),
46+
).toBe(false);
47+
});
48+
49+
it("returns false while a tool is still streaming or awaiting approval", () => {
50+
expect(
51+
hasPendingToolCalls([
52+
userMessage("delete it"),
53+
assistantToolMessage([
54+
{
55+
type: "tool-delete_file",
56+
toolCallId: "call-1",
57+
state: "approval-requested",
58+
input: { path: "secrets.env" },
59+
approval: { id: "approval-1" },
60+
} as unknown as UIMessage["parts"][number],
61+
]),
62+
]),
63+
).toBe(false);
64+
});
65+
66+
it("returns true when the user has responded to an approval request", () => {
67+
// The chat must auto-resume as soon as Approve/Deny is clicked.
68+
expect(
69+
hasPendingToolCalls([
70+
userMessage("delete it"),
71+
assistantToolMessage([
72+
{
73+
type: "tool-delete_file",
74+
toolCallId: "call-1",
75+
state: "approval-responded",
76+
input: { path: "secrets.env" },
77+
approval: { id: "approval-1", approved: true },
78+
} as unknown as UIMessage["parts"][number],
79+
]),
80+
]),
81+
).toBe(true);
82+
});
83+
84+
it("returns true when a tool reached a terminal output state", () => {
85+
expect(
86+
hasPendingToolCalls([
87+
userMessage("run it"),
88+
assistantToolMessage([
89+
{
90+
type: "tool-run_query",
91+
toolCallId: "call-1",
92+
state: "output-available",
93+
input: { sql: "select 1" },
94+
output: 1,
95+
} as unknown as UIMessage["parts"][number],
96+
]),
97+
]),
98+
).toBe(true);
99+
});
100+
101+
it("returns false when only some tool calls are ready", () => {
102+
expect(
103+
hasPendingToolCalls([
104+
userMessage("two things"),
105+
assistantToolMessage([
106+
{
107+
type: "tool-first",
108+
toolCallId: "call-1",
109+
state: "output-available",
110+
input: {},
111+
output: 1,
112+
} as unknown as UIMessage["parts"][number],
113+
{
114+
type: "tool-second",
115+
toolCallId: "call-2",
116+
state: "input-available",
117+
input: {},
118+
} as unknown as UIMessage["parts"][number],
119+
]),
120+
]),
121+
).toBe(false);
122+
});
123+
124+
it("returns false once the assistant has appended text after the tool result", () => {
125+
expect(
126+
hasPendingToolCalls([
127+
userMessage("run it"),
128+
assistantToolMessage([
129+
{
130+
type: "tool-run_query",
131+
toolCallId: "call-1",
132+
state: "output-available",
133+
input: {},
134+
output: 1,
135+
} as unknown as UIMessage["parts"][number],
136+
{ type: "text", text: "The query returned 1." },
137+
]),
138+
]),
139+
).toBe(false);
140+
});
141+
142+
it("returns false when a file part trails the completed tool call", () => {
143+
// Regression: tool → text → file used to loop because only trailing
144+
// text counted as "the assistant has answered".
145+
expect(
146+
hasPendingToolCalls([
147+
userMessage("show me Starry Night"),
148+
assistantToolMessage([
149+
{ type: "step-start" },
150+
{
151+
type: "tool-search_artwork",
152+
toolCallId: "call-1",
153+
state: "output-available",
154+
input: { artist: "Van Gogh" },
155+
output: { title: "The Starry Night" },
156+
} as unknown as UIMessage["parts"][number],
157+
{ type: "text", text: "Here is the painting:" },
158+
{
159+
type: "file",
160+
mediaType: "image/jpeg",
161+
url: "https://example.com/starry-night.jpg",
162+
} as unknown as UIMessage["parts"][number],
163+
]),
164+
]),
165+
).toBe(false);
166+
});
167+
168+
it("returns false when a source-url part trails the completed tool call", () => {
169+
expect(
170+
hasPendingToolCalls([
171+
userMessage("cite your sources"),
172+
assistantToolMessage([
173+
{
174+
type: "tool-web_search",
175+
toolCallId: "call-1",
176+
state: "output-available",
177+
input: { q: "marimo notebook" },
178+
output: "found",
179+
} as unknown as UIMessage["parts"][number],
180+
{ type: "text", text: "marimo is a reactive notebook." },
181+
{
182+
type: "source-url",
183+
sourceId: "src-1",
184+
url: "https://marimo.io",
185+
} as unknown as UIMessage["parts"][number],
186+
]),
187+
]),
188+
).toBe(false);
189+
});
190+
191+
it("returns false when a reasoning part trails the completed tool call", () => {
192+
expect(
193+
hasPendingToolCalls([
194+
userMessage("explain"),
195+
assistantToolMessage([
196+
{
197+
type: "tool-lookup",
198+
toolCallId: "call-1",
199+
state: "output-available",
200+
input: {},
201+
output: 1,
202+
} as unknown as UIMessage["parts"][number],
203+
{
204+
type: "reasoning",
205+
text: "Now I'll summarize.",
206+
} as unknown as UIMessage["parts"][number],
207+
]),
208+
]),
209+
).toBe(false);
210+
});
211+
212+
it("returns false when a new step-start follows the completed tool call", () => {
213+
expect(
214+
hasPendingToolCalls([
215+
userMessage("multi-step"),
216+
assistantToolMessage([
217+
{ type: "step-start" },
218+
{
219+
type: "tool-run_query",
220+
toolCallId: "call-1",
221+
state: "output-available",
222+
input: {},
223+
output: 1,
224+
} as unknown as UIMessage["parts"][number],
225+
{ type: "step-start" },
226+
]),
227+
]),
228+
).toBe(false);
229+
});
230+
231+
it("ignores providerExecuted tools", () => {
232+
// Provider-side tools are resolved by the model, not the runtime, so
233+
// they must not drive an auto-resume.
234+
expect(
235+
hasPendingToolCalls([
236+
userMessage("hi"),
237+
assistantToolMessage([
238+
{
239+
type: "tool-web_search",
240+
toolCallId: "call-1",
241+
state: "output-available",
242+
input: {},
243+
output: 1,
244+
providerExecuted: true,
245+
} as unknown as UIMessage["parts"][number],
246+
]),
247+
]),
248+
).toBe(false);
249+
});
250+
251+
it("returns true for dynamic-tool parts in a terminal state", () => {
252+
// `dynamic-tool` parts must drive auto-resume alongside `tool-*`.
253+
expect(
254+
hasPendingToolCalls([
255+
userMessage("run it"),
256+
assistantToolMessage([
257+
{
258+
type: "dynamic-tool",
259+
toolName: "run_query",
260+
toolCallId: "call-1",
261+
state: "output-available",
262+
input: {},
263+
output: 1,
264+
} as unknown as UIMessage["parts"][number],
265+
]),
266+
]),
267+
).toBe(true);
268+
});
269+
});

‎frontend/src/components/chat/chat-utils.ts‎

Lines changed: 14 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,8 @@ import {
55
type ChatAddToolOutputFunction,
66
type FileUIPart,
77
isToolUIPart,
8-
type ToolUIPart,
8+
lastAssistantMessageIsCompleteWithApprovalResponses,
9+
lastAssistantMessageIsCompleteWithToolCalls,
910
type UIMessage,
1011
} from "ai";
1112
import { useState } from "react";
@@ -17,7 +18,6 @@ import type {
1718
InvokeAiToolRequest,
1819
InvokeAiToolResponse,
1920
} from "@/core/network/types";
20-
import { logNever } from "@/utils/assertNever";
2121
import { blobToString } from "@/utils/fileToBase64";
2222
import { Logger } from "@/utils/Logger";
2323
import { getAICompletionBodyWithAttachments } from "../editor/ai/completion-utils";
@@ -169,69 +169,25 @@ export async function handleToolCall({
169169
}
170170

171171
/**
172-
* Returns true if a tool call is "ready to be sent back to the server" — i.e.
173-
* either it has reached a terminal output state, or the user has just supplied
174-
* an approval response that the server hasn't seen yet.
175-
*/
176-
function isToolCallReadyToSend(state: ToolUIPart["state"]): boolean {
177-
switch (state) {
178-
case "output-available":
179-
case "output-error":
180-
case "output-denied":
181-
case "approval-responded":
182-
return true;
183-
case "input-streaming":
184-
case "input-available":
185-
case "approval-requested":
186-
return false;
187-
default:
188-
logNever(state);
189-
return false;
190-
}
191-
}
192-
193-
/**
194-
* Checks if we should send a message automatically based on the messages.
195-
* We auto-send when every tool call on the last assistant message has either
196-
* finished (output-available/error/denied) or has just received a user
197-
* approval response, and the assistant hasn't replied yet.
172+
* Auto-send the next turn when the last assistant message ends with a
173+
* tool call ready to round-trip. Any non-tool trailing part (text, file,
174+
* source-*, reasoning, data-*, new step-start) means the assistant has
175+
* already answered, so we leave the next turn to the user. State checks
176+
* are delegated to the SDK to stay in sync with upstream.
198177
*/
199178
export function hasPendingToolCalls(messages: UIMessage[]): boolean {
200-
if (messages.length === 0) {
201-
return false;
202-
}
203-
204-
const lastMessage = messages[messages.length - 1];
205-
const parts = lastMessage.parts;
206-
207-
if (parts.length === 0) {
208-
return false;
209-
}
210-
211-
// Only auto-send if the last message is an assistant message
212-
// Because assistant messages are the ones that can have tool calls
213-
if (lastMessage.role !== "assistant") {
179+
const lastMessage = messages.at(-1);
180+
if (!lastMessage || lastMessage.role !== "assistant") {
214181
return false;
215182
}
216-
217-
const toolParts = parts.filter(isToolUIPart);
218-
219-
if (toolParts.length === 0) {
183+
const lastPart = lastMessage.parts.at(-1);
184+
if (!lastPart || !isToolUIPart(lastPart)) {
220185
return false;
221186
}
222-
223-
const allToolCallsReady = toolParts.every((part) =>
224-
isToolCallReadyToSend(part.state),
187+
return (
188+
lastAssistantMessageIsCompleteWithToolCalls({ messages }) ||
189+
lastAssistantMessageIsCompleteWithApprovalResponses({ messages })
225190
);
226-
227-
// Check if the last part has any text content
228-
const lastPart = parts[parts.length - 1];
229-
const hasTextContent =
230-
lastPart.type === "text" && lastPart.text?.trim().length > 0;
231-
232-
Logger.debug("All tool calls ready to send: %s", allToolCallsReady);
233-
234-
return allToolCallsReady && !hasTextContent;
235191
}
236192

237193
export function useFileState() {

0 commit comments

Comments
 (0)