@marimo-team/frontend 0.23.9-dev4 → 0.23.9-dev41

package/src/components/app-config/ai-config.tsx

@@ -509,7 +509,7 @@ const ModelInfoCard = ({ model }: { model: AiModel }) => {
         <Tooltip content="Custom model">
           {model.custom && <BotIcon className="h-4 w-4" />}
         </Tooltip>
-        {model.thinking && (
+        {model.capabilities.includes("thinking") && (
           <div
             className={cn(
               "flex items-center gap-1 rounded px-1 py-0.5 w-fit",
@@ -1253,6 +1253,13 @@ export const AiAssistConfig: React.FC<AiConfigProps> = ({
   config,
   onSubmit,
 }) => {
+  // Tracked locally rather than derived from the field value so that clearing
+  // the input (a transient empty value, which commits `null`) does not disable
+  // the input mid-edit and force the user to re-tick the Override checkbox.
+  const [maxTokensEnabled, setMaxTokensEnabled] = useState(
+    config.ai?.max_tokens != null,
+  );
+
   return (
     <SettingGroup>
       <SettingSubtitle>AI Assistant</SettingSubtitle>
@@ -1279,6 +1286,71 @@ export const AiAssistConfig: React.FC<AiConfigProps> = ({
         )}
       />
 
+      <FormField
+        control={form.control}
+        name="ai.max_tokens"
+        render={({ field }) => {
+          return (
+            <div className="flex flex-col gap-y-1">
+              <div className="flex items-center gap-x-2">
+                <FormItem className={formItemClasses}>
+                  <FormLabel className="font-normal">
+                    Max output tokens
+                  </FormLabel>
+                  <FormControl>
+                    <Input
+                      data-testid="ai-max-tokens-input"
+                      type="number"
+                      min={1}
+                      disabled={!maxTokensEnabled}
+                      className="w-28 h-6"
+                      value={field.value ?? (maxTokensEnabled ? "" : 32768)}
+                      onChange={(e) => {
+                        const n = Number.parseInt(e.target.value, 10);
+                        field.onChange(Number.isFinite(n) && n > 0 ? n : null);
+                      }}
+                    />
+                  </FormControl>
+                </FormItem>
+                <FormItem className={formItemClasses}>
+                  <Checkbox
+                    data-testid="ai-max-tokens-checkbox"
+                    checked={maxTokensEnabled}
+                    onCheckedChange={(checked) => {
+                      const isChecked = checked === true;
+                      setMaxTokensEnabled(isChecked);
+                      // null signals delete to the server; cast because
+                      // UserConfig (OpenAPI-derived) types max_tokens as
+                      // `number | undefined`, but zod accepts `null`.
+                      const next = (
+                        isChecked ? (field.value ?? 32768) : null
+                      ) as number | undefined;
+                      // shouldDirty: true forces RHF to keep this in
+                      // dirtyFields even when `next` happens to equal the
+                      // form's defaultValue (e.g. untick → tick when disk
+                      // started with 32768). Otherwise getDirtyValues
+                      // would skip it and the save body would be empty.
+                      form.setValue("ai.max_tokens", next, {
+                        shouldDirty: true,
+                        shouldTouch: true,
+                      });
+                      onSubmit(form.getValues());
+                    }}
+                  />
+                  <FormLabel className="font-normal">Override</FormLabel>
+                </FormItem>
+              </div>
+
+              <FormDescription>
+                Each provider sets its own max output tokens (Anthropic uses a
+                recommended default). Adjust to control costs or enable more
+                output.
+              </FormDescription>
+            </div>
+          );
+        }}
+      />
+
       <FormErrorsBanner />
       <ModelSelector
         label="Chat Model"
@@ -1777,6 +1849,38 @@ export type AiSettingsSubTab =
   | "ai-models"
   | "mcp";
 
+const AiEnabledConfig: React.FC<AiConfigProps> = ({ form, config }) => {
+  return (
+    <SettingGroup>
+      <FormField
+        control={form.control}
+        name="ai.enabled"
+        render={({ field }) => (
+          <div className="flex flex-col gap-y-1">
+            <FormItem className={formItemClasses}>
+              <FormLabel className="font-normal">Enable AI features</FormLabel>
+              <FormControl>
+                <Checkbox
+                  data-testid="ai-enabled-checkbox"
+                  checked={field.value !== false}
+                  onCheckedChange={(checked) =>
+                    field.onChange(checked === true)
+                  }
+                />
+              </FormControl>
+              <IsOverridden userConfig={config} name="ai.enabled" />
+            </FormItem>
+            <FormDescription>
+              When disabled, AI actions and panels are hidden from the marimo
+              UI.
+            </FormDescription>
+          </div>
+        )}
+      />
+    </SettingGroup>
+  );
+};
+
 export const AiConfig: React.FC<AiConfigProps> = ({
   form,
   config,
@@ -1785,18 +1889,30 @@ export const AiConfig: React.FC<AiConfigProps> = ({
   // MCP is not supported in WASM
   const wasm = isWasm();
   const [activeTab, setActiveTab] = useAtom(aiSettingsSubTabAtom);
+  const aiEnabled = useWatch({
+    control: form.control,
+    name: "ai.enabled",
+  });
+  const activeVisibleTab =
+    aiEnabled === false && activeTab !== "ai-features"
+      ? "ai-features"
+      : activeTab;
 
   return (
     <Tabs
-      value={activeTab}
+      value={activeVisibleTab}
       onValueChange={(value) => setActiveTab(value as AiSettingsSubTab)}
       className="flex-1"
     >
       <TabsList className="mb-2">
         <TabsTrigger value="ai-features">AI Features</TabsTrigger>
-        <TabsTrigger value="ai-providers">AI Providers</TabsTrigger>
-        <TabsTrigger value="ai-models">AI Models</TabsTrigger>
-        {!wasm && <TabsTrigger value="mcp">MCP</TabsTrigger>}
+        {aiEnabled !== false && (
+          <>
+            <TabsTrigger value="ai-providers">AI Providers</TabsTrigger>
+            <TabsTrigger value="ai-models">AI Models</TabsTrigger>
+            {!wasm && <TabsTrigger value="mcp">MCP</TabsTrigger>}
+          </>
+        )}
       </TabsList>
 
       <TabsContent value="ai-features">
@@ -1805,18 +1921,33 @@ export const AiConfig: React.FC<AiConfigProps> = ({
           config={config}
           onSubmit={onSubmit}
         />
-        <AiAssistConfig form={form} config={config} onSubmit={onSubmit} />
-      </TabsContent>
-      <TabsContent value="ai-providers">
-        <AiProvidersConfig form={form} config={config} onSubmit={onSubmit} />
-      </TabsContent>
-      <TabsContent value="ai-models">
-        <AiModelDisplayConfig form={form} config={config} onSubmit={onSubmit} />
+        <AiEnabledConfig form={form} config={config} onSubmit={onSubmit} />
+        {aiEnabled !== false && (
+          <AiAssistConfig form={form} config={config} onSubmit={onSubmit} />
+        )}
       </TabsContent>
-      {!wasm && (
-        <TabsContent value="mcp">
-          <MCPConfig form={form} onSubmit={onSubmit} />
-        </TabsContent>
+      {aiEnabled !== false && (
+        <>
+          <TabsContent value="ai-providers">
+            <AiProvidersConfig
+              form={form}
+              config={config}
+              onSubmit={onSubmit}
+            />
+          </TabsContent>
+          <TabsContent value="ai-models">
+            <AiModelDisplayConfig
+              form={form}
+              config={config}
+              onSubmit={onSubmit}
+            />
+          </TabsContent>
+          {!wasm && (
+            <TabsContent value="mcp">
+              <MCPConfig form={form} onSubmit={onSubmit} />
+            </TabsContent>
+          )}
+        </>
       )}
     </Tabs>
   );

package/src/components/app-config/user-config-form.tsx

@@ -503,6 +503,42 @@ export const UserConfigForm: React.FC = () => {
                   </div>
                 )}
               />
+              <FormField
+                control={form.control}
+                name="completion.auto_close_pairs"
+                render={({ field }) => (
+                  <div className="flex flex-col space-y-1">
+                    <FormItem className={formItemClasses}>
+                      <FormLabel className="font-normal">
+                        Auto-close pairs
+                      </FormLabel>
+                      <FormControl>
+                        <Checkbox
+                          data-testid="auto-close-pairs-checkbox"
+                          checked={field.value ?? true}
+                          disabled={field.disabled}
+                          onCheckedChange={(checked) => {
+                            field.onChange(Boolean(checked));
+                          }}
+                        />
+                      </FormControl>
+                      <FormMessage />
+                      <IsOverridden
+                        userConfig={config}
+                        name="completion.auto_close_pairs"
+                      />
+                    </FormItem>
+                    <FormDescription>
+                      Automatically insert closing brackets{" "}
+                      <code className="text-xs">{"()"}</code>,{" "}
+                      <code className="text-xs">{"[]"}</code>,{" "}
+                      <code className="text-xs">{"{}"}</code>, and quotes{" "}
+                      <code className="text-xs">{`""`}</code>,{" "}
+                      <code className="text-xs">{`''`}</code> when opening one.
+                    </FormDescription>
+                  </div>
+                )}
+              />
             </SettingGroup>
             <SettingGroup title="Language Servers">
               <FormDescription>
@@ -1319,7 +1355,7 @@ export const UserConfigForm: React.FC = () => {
     <Form {...form}>
       <form
         ref={formElement}
-        onChange={form.handleSubmit(onSubmit)}
+        onChange={form.handleSubmit((values) => onSubmit(values))}
         className="flex text-pretty overflow-hidden"
       >
         <Tabs

package/src/components/chat/__tests__/chat-utils.test.ts

@@ -0,0 +1,269 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { UIMessage } from "ai";
+import { describe, expect, it } from "vitest";
+import { hasPendingToolCalls } from "../chat-utils";
+
+/**
+ * `hasPendingToolCalls` powers `sendAutomaticallyWhen` in `mo.ui.chat`:
+ * returns true only when the last assistant message *ends* with a tool
+ * call in a ready-to-round-trip state. Any trailing non-tool part (text,
+ * file, source-*, reasoning, data-*, new step-start) means the assistant
+ * has already answered and we leave the next turn to the user. The
+ * approval flow relies on this firing for `approval-responded`.
+ */
+
+const userMessage = (text: string): UIMessage => ({
+  id: `user-${text}`,
+  role: "user",
+  parts: [{ type: "text", text }],
+});
+
+const assistantToolMessage = (
+  parts: UIMessage["parts"],
+  id = "assistant-1",
+): UIMessage => ({
+  id,
+  role: "assistant",
+  parts,
+});
+
+describe("hasPendingToolCalls", () => {
+  it("returns false when there are no messages", () => {
+    expect(hasPendingToolCalls([])).toBe(false);
+  });
+
+  it("returns false when the last message is a user message", () => {
+    expect(hasPendingToolCalls([userMessage("hi")])).toBe(false);
+  });
+
+  it("returns false when the last assistant message has no tool parts", () => {
+    expect(
+      hasPendingToolCalls([
+        userMessage("hi"),
+        assistantToolMessage([{ type: "text", text: "hello!" }]),
+      ]),
+    ).toBe(false);
+  });
+
+  it("returns false while a tool is still streaming or awaiting approval", () => {
+    expect(
+      hasPendingToolCalls([
+        userMessage("delete it"),
+        assistantToolMessage([
+          {
+            type: "tool-delete_file",
+            toolCallId: "call-1",
+            state: "approval-requested",
+            input: { path: "secrets.env" },
+            approval: { id: "approval-1" },
+          } as unknown as UIMessage["parts"][number],
+        ]),
+      ]),
+    ).toBe(false);
+  });
+
+  it("returns true when the user has responded to an approval request", () => {
+    // The chat must auto-resume as soon as Approve/Deny is clicked.
+    expect(
+      hasPendingToolCalls([
+        userMessage("delete it"),
+        assistantToolMessage([
+          {
+            type: "tool-delete_file",
+            toolCallId: "call-1",
+            state: "approval-responded",
+            input: { path: "secrets.env" },
+            approval: { id: "approval-1", approved: true },
+          } as unknown as UIMessage["parts"][number],
+        ]),
+      ]),
+    ).toBe(true);
+  });
+
+  it("returns true when a tool reached a terminal output state", () => {
+    expect(
+      hasPendingToolCalls([
+        userMessage("run it"),
+        assistantToolMessage([
+          {
+            type: "tool-run_query",
+            toolCallId: "call-1",
+            state: "output-available",
+            input: { sql: "select 1" },
+            output: 1,
+          } as unknown as UIMessage["parts"][number],
+        ]),
+      ]),
+    ).toBe(true);
+  });
+
+  it("returns false when only some tool calls are ready", () => {
+    expect(
+      hasPendingToolCalls([
+        userMessage("two things"),
+        assistantToolMessage([
+          {
+            type: "tool-first",
+            toolCallId: "call-1",
+            state: "output-available",
+            input: {},
+            output: 1,
+          } as unknown as UIMessage["parts"][number],
+          {
+            type: "tool-second",
+            toolCallId: "call-2",
+            state: "input-available",
+            input: {},
+          } as unknown as UIMessage["parts"][number],
+        ]),
+      ]),
+    ).toBe(false);
+  });
+
+  it("returns false once the assistant has appended text after the tool result", () => {
+    expect(
+      hasPendingToolCalls([
+        userMessage("run it"),
+        assistantToolMessage([
+          {
+            type: "tool-run_query",
+            toolCallId: "call-1",
+            state: "output-available",
+            input: {},
+            output: 1,
+          } as unknown as UIMessage["parts"][number],
+          { type: "text", text: "The query returned 1." },
+        ]),
+      ]),
+    ).toBe(false);
+  });
+
+  it("returns false when a file part trails the completed tool call", () => {
+    // Regression: tool → text → file used to loop because only trailing
+    // text counted as "the assistant has answered".
+    expect(
+      hasPendingToolCalls([
+        userMessage("show me Starry Night"),
+        assistantToolMessage([
+          { type: "step-start" },
+          {
+            type: "tool-search_artwork",
+            toolCallId: "call-1",
+            state: "output-available",
+            input: { artist: "Van Gogh" },
+            output: { title: "The Starry Night" },
+          } as unknown as UIMessage["parts"][number],
+          { type: "text", text: "Here is the painting:" },
+          {
+            type: "file",
+            mediaType: "image/jpeg",
+            url: "https://example.com/starry-night.jpg",
+          } as unknown as UIMessage["parts"][number],
+        ]),
+      ]),
+    ).toBe(false);
+  });
+
+  it("returns false when a source-url part trails the completed tool call", () => {
+    expect(
+      hasPendingToolCalls([
+        userMessage("cite your sources"),
+        assistantToolMessage([
+          {
+            type: "tool-web_search",
+            toolCallId: "call-1",
+            state: "output-available",
+            input: { q: "marimo notebook" },
+            output: "found",
+          } as unknown as UIMessage["parts"][number],
+          { type: "text", text: "marimo is a reactive notebook." },
+          {
+            type: "source-url",
+            sourceId: "src-1",
+            url: "https://marimo.io",
+          } as unknown as UIMessage["parts"][number],
+        ]),
+      ]),
+    ).toBe(false);
+  });
+
+  it("returns false when a reasoning part trails the completed tool call", () => {
+    expect(
+      hasPendingToolCalls([
+        userMessage("explain"),
+        assistantToolMessage([
+          {
+            type: "tool-lookup",
+            toolCallId: "call-1",
+            state: "output-available",
+            input: {},
+            output: 1,
+          } as unknown as UIMessage["parts"][number],
+          {
+            type: "reasoning",
+            text: "Now I'll summarize.",
+          } as unknown as UIMessage["parts"][number],
+        ]),
+      ]),
+    ).toBe(false);
+  });
+
+  it("returns false when a new step-start follows the completed tool call", () => {
+    expect(
+      hasPendingToolCalls([
+        userMessage("multi-step"),
+        assistantToolMessage([
+          { type: "step-start" },
+          {
+            type: "tool-run_query",
+            toolCallId: "call-1",
+            state: "output-available",
+            input: {},
+            output: 1,
+          } as unknown as UIMessage["parts"][number],
+          { type: "step-start" },
+        ]),
+      ]),
+    ).toBe(false);
+  });
+
+  it("ignores providerExecuted tools", () => {
+    // Provider-side tools are resolved by the model, not the runtime, so
+    // they must not drive an auto-resume.
+    expect(
+      hasPendingToolCalls([
+        userMessage("hi"),
+        assistantToolMessage([
+          {
+            type: "tool-web_search",
+            toolCallId: "call-1",
+            state: "output-available",
+            input: {},
+            output: 1,
+            providerExecuted: true,
+          } as unknown as UIMessage["parts"][number],
+        ]),
+      ]),
+    ).toBe(false);
+  });
+
+  it("returns true for dynamic-tool parts in a terminal state", () => {
+    // `dynamic-tool` parts must drive auto-resume alongside `tool-*`.
+    expect(
+      hasPendingToolCalls([
+        userMessage("run it"),
+        assistantToolMessage([
+          {
+            type: "dynamic-tool",
+            toolName: "run_query",
+            toolCallId: "call-1",
+            state: "output-available",
+            input: {},
+            output: 1,
+          } as unknown as UIMessage["parts"][number],
+        ]),
+      ]),
+    ).toBe(true);
+  });
+});

package/src/components/chat/chat-panel.tsx

@@ -15,11 +15,13 @@ import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
 import {
   BotMessageSquareIcon,
   HatGlasses,
+  ArrowRightIcon,
   Loader2,
   type LucideIcon,
   MessageCircleIcon,
   NotebookText,
   PlusIcon,
+  SparklesIcon,
   SettingsIcon,
 } from "lucide-react";
 import { memo, useEffect, useRef, useState } from "react";
@@ -49,15 +51,17 @@ import {
   FRONTEND_TOOL_REGISTRY,
 } from "@/core/ai/tools/registry";
 import { useCellActions } from "@/core/cells/cells";
-import { aiAtom, aiEnabledAtom } from "@/core/config/config";
+import { aiAtom, aiModelConfiguredAtom } from "@/core/config/config";
 import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { useRequestClient } from "@/core/network/requests";
 import { useRuntimeManager } from "@/core/runtime/config";
+import { isWasm } from "@/core/wasm/utils";
 import { ErrorBanner } from "@/plugins/impl/common/error-banner";
 import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
 import { AIModelDropdown } from "../ai/ai-model-dropdown";
 import { useOpenSettingsToTab } from "../app-config/state";
+import { PairWithAgentModal } from "../editor/actions/pair-with-agent-modal";
 import { PromptInput } from "../editor/ai/add-cell-with-ai";
 import {
   addContextCompletion,
@@ -65,6 +69,7 @@ import {
 } from "../editor/ai/completion-utils";
 import { PanelEmptyState } from "../editor/chrome/panels/empty-state";
 import { CopyClipboardIcon } from "../icons/copy-icon";
+import { useImperativeModal } from "../modal/ImperativeModal";
 import { MCPStatusIndicator } from "../mcp/mcp-status-indicator";
 import { Tooltip, TooltipProvider } from "../ui/tooltip";
 import {
@@ -416,8 +421,28 @@ const ChatInput: React.FC<ChatInputProps> = memo(
 
 ChatInput.displayName = "ChatInput";
 
+const PairWithAgentCallout: React.FC<{
+  onPairWithAgent: () => void;
+}> = ({ onPairWithAgent }) => {
+  if (isWasm()) {
+    return null;
+  }
+
+  return (
+    <Button
+      variant="text"
+      className="gap-1.5 text-sm text-link hover:underline"
+      onClick={onPairWithAgent}
+    >
+      <SparklesIcon className="h-3.5 w-3.5 shrink-0" />
+      <span>Work on this notebook with your own agent</span>
+      <ArrowRightIcon className="h-3 w-3 shrink-0" />
+    </Button>
+  );
+};
+
 const ChatPanel = () => {
-  const aiConfigured = useAtomValue(aiEnabledAtom);
+  const aiConfigured = useAtomValue(aiModelConfiguredAtom);
   const { handleClick } = useOpenSettingsToTab();
 
   if (!aiConfigured) {
@@ -455,6 +480,7 @@ const ChatPanelBody = () => {
   const messagesEndRef = useRef<HTMLDivElement>(null);
   const runtimeManager = useRuntimeManager();
   const { invokeAiTool, sendRun } = useRequestClient();
+  const { openModal, closeModal } = useImperativeModal();
 
   const activeChatId = activeChat?.id;
   const store = useStore();
@@ -614,6 +640,10 @@ const ChatPanelBody = () => {
     clearFiles();
   });
 
+  const handlePairWithAgent = useEvent(() => {
+    openModal(<PairWithAgentModal onClose={closeModal} />);
+  });
+
   const handleMessageEdit = useEvent((index: number, newValue: string) => {
     const editedMessage = messages[index];
     const fileParts = editedMessage.parts?.filter((p) => p.type === "file");
@@ -724,9 +754,12 @@ const ChatPanelBody = () => {
         ref={scrollContainerRef}
       >
         {isNewThread && (
-          <div className="rounded-md border bg-background">
-            {filesPills}
-            {chatInput}
+          <div className="flex flex-col gap-2">
+            <div className="rounded-md border bg-background">
+              {filesPills}
+              {chatInput}
+            </div>
+            <PairWithAgentCallout onPairWithAgent={handlePairWithAgent} />
           </div>
         )}