From 98b36616cbd65828604e1c19818eb335f906d69d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Greg=20Berg=C3=A9?= <berge.greg@gmail.com>
Date: Fri, 6 Mar 2026 16:58:48 +0100
Subject: [PATCH] Add interactive controls to docs assistant (#4073)

---
 bun.lock                                      |  15 +-
 packages/gitbook/e2e/internal.spec.ts         |   2 +-
 packages/gitbook/package.json                 |   1 +
 .../components/AI/controls/ConfirmControl.tsx |  55 +++++
 .../AI/controls/SingleChoiceControl.tsx       | 203 ++++++++++++++++++
 .../src/components/AI/controls/common.tsx     |  15 ++
 .../src/components/AI/controls/helpers.ts     |  80 +++++++
 .../src/components/AI/controls/index.ts       |  11 +
 packages/gitbook/src/components/AI/tools.ts   |   8 +
 .../gitbook/src/components/AI/useAIChat.tsx   | 156 ++++++++------
 .../gitbook/src/components/AIChat/AIChat.tsx  |  25 ++-
 .../src/components/AIChat/AIChatControl.tsx   |   8 +
 .../src/components/AIChat/AIChatMessages.tsx  |   5 +-
 .../AIChat/AIChatToolConfirmations.tsx        |  70 ------
 .../src/components/primitives/Input.tsx       |   4 +-
 15 files changed, 501 insertions(+), 157 deletions(-)
 create mode 100644 packages/gitbook/src/components/AI/controls/ConfirmControl.tsx
 create mode 100644 packages/gitbook/src/components/AI/controls/SingleChoiceControl.tsx
 create mode 100644 packages/gitbook/src/components/AI/controls/common.tsx
 create mode 100644 packages/gitbook/src/components/AI/controls/helpers.ts
 create mode 100644 packages/gitbook/src/components/AI/controls/index.ts
 create mode 100644 packages/gitbook/src/components/AI/tools.ts
 create mode 100644 packages/gitbook/src/components/AIChat/AIChatControl.tsx
 delete mode 100644 packages/gitbook/src/components/AIChat/AIChatToolConfirmations.tsx

diff --git a/bun.lock b/bun.lock
index e7186c086..d59fabc6c 100644
--- a/bun.lock
+++ b/bun.lock
@@ -6,9 +6,9 @@
       "name": "gitbook",
       "devDependencies": {
         "@biomejs/biome": "^1.9.4",
-        "@changesets/cli": "^2.29.8",
-        "turbo": "^2.8.10",
-        "vercel": "^50.23.2",
+        "@changesets/cli": "^2.30.0",
+        "turbo": "^2.8.13",
+        "vercel": "^50.26.1",
       },
     },
     "packages/browser-types": {
@@ -188,6 +188,7 @@
         "usehooks-ts": "catalog:",
         "warn-once": "^0.1.1",
         "zod": "^3",
+        "zod-to-json-schema": "^3.25.1",
         "zustand": "^5.0.3",
       },
       "devDependencies": {
@@ -3502,7 +3503,7 @@
 
     "zod": ["zod@3.24.2", "", {}, "sha512-lY7CDW43ECgW9u1TcT3IoXHflywfVqDYze4waEz812jR/bZ8FHDsl7pFQoSZTz5N+2NqRXs8GBwnAwo3ZNxqhQ=="],
 
-    "zod-to-json-schema": ["zod-to-json-schema@3.24.5", "", { "peerDependencies": { "zod": "^3.24.1" } }, "sha512-/AuWwMP+YqiPbsJx5D6TfgRTc4kTLjsh5SOcd4bLsfUg2RcEXrFMJl1DGgdHy2aCfsIA/cr/1JM0xcB2GZji8g=="],
+    "zod-to-json-schema": ["zod-to-json-schema@3.25.1", "", { "peerDependencies": { "zod": "^3.25 || ^4" } }, "sha512-pM/SU9d3YAggzi6MtR4h7ruuQlqKtad8e9S0fmxcMi+ueAK5Korys/aWcV9LIIHTVbj01NdzxcnXSN+O74ZIVA=="],
 
     "zustand": ["zustand@5.0.3", "", { "peerDependencies": { "@types/react": ">=18.0.0", "immer": ">=9.0.6", "react": ">=18.0.0", "use-sync-external-store": ">=1.2.0" }, "optionalPeers": ["@types/react", "immer", "react", "use-sync-external-store"] }, "sha512-14fwWQtU3pH4dE0dOpdMiWjddcH+QzKIgk1cl8epwSE7yag43k/AD/m4L6+K7DytAOr9gGBe3/EXj9g7cdostg=="],
 
@@ -3510,6 +3511,8 @@
 
     "@ai-sdk/provider-utils/nanoid": ["nanoid@3.3.8", "", { "bin": { "nanoid": "bin/nanoid.cjs" } }, "sha512-WNLf5Sd8oZxOm+TzppcYk8gVOgP+l58xNy58D0nbUnOxOWRWvlcCV4kUF7ltmI6PsrLl/BgKEyS4mqsGChFN0w=="],
 
+    "@ai-sdk/ui-utils/zod-to-json-schema": ["zod-to-json-schema@3.24.5", "", { "peerDependencies": { "zod": "^3.24.1" } }, "sha512-/AuWwMP+YqiPbsJx5D6TfgRTc4kTLjsh5SOcd4bLsfUg2RcEXrFMJl1DGgdHy2aCfsIA/cr/1JM0xcB2GZji8g=="],
+
     "@ampproject/remapping/@jridgewell/gen-mapping": ["@jridgewell/gen-mapping@0.3.5", "", { "dependencies": { "@jridgewell/set-array": "^1.2.1", "@jridgewell/sourcemap-codec": "^1.4.10", "@jridgewell/trace-mapping": "^0.3.24" } }, "sha512-IzL8ZoEDIBRWEzlCcRhOaCupYyN5gdIK+Q6fbFdPDg6HqX6jpkItn7DFIpW9LQzXG6Df9sA7+OKnq0qlz/GaQg=="],
 
     "@ampproject/remapping/@jridgewell/trace-mapping": ["@jridgewell/trace-mapping@0.3.25", "", { "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", "@jridgewell/sourcemap-codec": "^1.4.14" } }, "sha512-vNk6aEwybGtawWmy/PzwnGDOjCkLWSD2wqvjGGAgOAwCGWySYXfYoxt00IJkTF+8Lb57DwOb3Aa0o9CApepiYQ=="],
@@ -4132,6 +4135,8 @@
 
     "@modelcontextprotocol/sdk/cross-spawn": ["cross-spawn@7.0.6", "", { "dependencies": { "path-key": "^3.1.0", "shebang-command": "^2.0.0", "which": "^2.0.1" } }, "sha512-uV2QOWP2nWzsy2aMp8aRibhi9dlzF5Hgh5SHaB9OiTGEyDTiJJyx0uy51QXdyWbtAHNua4XJzUKca3OzKUd3vA=="],
 
+    "@modelcontextprotocol/sdk/zod-to-json-schema": ["zod-to-json-schema@3.24.5", "", { "peerDependencies": { "zod": "^3.24.1" } }, "sha512-/AuWwMP+YqiPbsJx5D6TfgRTc4kTLjsh5SOcd4bLsfUg2RcEXrFMJl1DGgdHy2aCfsIA/cr/1JM0xcB2GZji8g=="],
+
     "@napi-rs/wasm-runtime/@emnapi/runtime": ["@emnapi/runtime@1.5.0", "", { "dependencies": { "tslib": "^2.4.0" } }, "sha512-97/BJ3iXHww3djw6hYIfErCZFee7qCtrneuLa20UXFCOTCfBM2cvQHjWJ2EG0s0MtdNwInarqCTz35i4wWXHsQ=="],
 
     "@node-minify/core/glob": ["glob@9.3.5", "", { "dependencies": { "fs.realpath": "^1.0.0", "minimatch": "^8.0.2", "minipass": "^4.2.4", "path-scurry": "^1.6.1" } }, "sha512-e1LleDykUz2Iu+MTYdkSsuWX8lvAjAcs0Xef0lNIu0S2wOAzuTxCJtcd9S3cijlwYF18EsU3rzb8jPVobxDh9Q=="],
@@ -4818,6 +4823,8 @@
 
     "yargs/string-width": ["string-width@7.2.0", "", { "dependencies": { "emoji-regex": "^10.3.0", "get-east-asian-width": "^1.0.0", "strip-ansi": "^7.1.0" } }, "sha512-tsaTIkKW9b4N+AEj+SVA+WhJzV7/zMhcSu78mLKWSk7cXMOSHsBKFWUs0fWwq8QyK3MgJBQRX6Gbi4kYbdvGkQ=="],
 
+    "zod-to-json-schema/zod": ["zod@4.3.6", "", {}, "sha512-rftlrkhHZOcjDwkGlnUtZZkvaPHCsDATp4pGpuOOMDaTdDDXF91wuVDJoWoPsKX/3YPQ5fHuF3STjcYyKr+Qhg=="],
+
     "@ampproject/remapping/@jridgewell/gen-mapping/@jridgewell/sourcemap-codec": ["@jridgewell/sourcemap-codec@1.5.0", "", {}, "sha512-gv3ZRaISU3fjPAgNsriBRqGWQL6quFx04YMPW/zD8XMLsU32mhCCbfbO6KZFLjvYpCZ8zyDEgqsgf+PwPaM7GQ=="],
 
     "@ampproject/remapping/@jridgewell/trace-mapping/@jridgewell/sourcemap-codec": ["@jridgewell/sourcemap-codec@1.5.0", "", {}, "sha512-gv3ZRaISU3fjPAgNsriBRqGWQL6quFx04YMPW/zD8XMLsU32mhCCbfbO6KZFLjvYpCZ8zyDEgqsgf+PwPaM7GQ=="],
diff --git a/packages/gitbook/e2e/internal.spec.ts b/packages/gitbook/e2e/internal.spec.ts
index 4b52103c8..0f36b3882 100644
--- a/packages/gitbook/e2e/internal.spec.ts
+++ b/packages/gitbook/e2e/internal.spec.ts
@@ -2315,7 +2315,7 @@ const testCases: TestsCase[] = [
                         'I want to contact support. Call the tool directly without a preamble. Do not respond with anything else.'
                     );
                     const toolConfirmation = iframe
-                        .getByTestId('ai-chat-tool-confirmation')
+                        .getByTestId('ai-chat-tool-confirm-accept')
                         .first();
                     await expect(toolConfirmation).toBeVisible({
                         timeout: 30000,
diff --git a/packages/gitbook/package.json b/packages/gitbook/package.json
index 8f77eb000..a91c0ef2a 100644
--- a/packages/gitbook/package.json
+++ b/packages/gitbook/package.json
@@ -80,6 +80,7 @@
         "usehooks-ts": "catalog:",
         "warn-once": "^0.1.1",
         "zod": "^3",
+        "zod-to-json-schema": "^3.25.1",
         "zustand": "^5.0.3"
     },
     "devDependencies": {
diff --git a/packages/gitbook/src/components/AI/controls/ConfirmControl.tsx b/packages/gitbook/src/components/AI/controls/ConfirmControl.tsx
new file mode 100644
index 000000000..46421f874
--- /dev/null
+++ b/packages/gitbook/src/components/AI/controls/ConfirmControl.tsx
@@ -0,0 +1,55 @@
+'use client';
+
+import { Button } from '@/components/primitives';
+import { z } from 'zod';
+import { AIToolContainer } from './common';
+import { type GetAIControlProps, createAIControl } from './helpers';
+
+export const ConfirmControlOutputSchema = z.object({
+    result: z.enum(['confirmed', 'cancelled']),
+});
+
+export const ConfirmControlDef = createAIControl({
+    name: 'confirm',
+    exposeAsTool: false,
+    description:
+        'Display a confirmation prompt to the user (Confirm / Cancel) to approve or abort a pending action. Use this when an operation is irreversible, sensitive, or should only proceed with explicit user consent. Returns either a `confirmed` or `cancelled` result based on the user’s click.',
+    inputSchema: z.object({
+        icon: z
+            .string()
+            .optional()
+            .describe('Icon to be shown in the confirm button (any Fontawesome icon name).'),
+        label: z.string().describe('Label to be shown in the confirm button.'),
+    }),
+    outputSchema: ConfirmControlOutputSchema,
+    render: (props) => {
+        return <ConfirmControl {...props} />;
+    },
+});
+
+function ConfirmControl(props: GetAIControlProps<typeof ConfirmControlDef>) {
+    const { label, icon, onSubmit } = props;
+    return (
+        <AIToolContainer className="flex w-full flex-col gap-2">
+            <Button
+                data-testid="ai-chat-tool-confirm-cancel"
+                onClick={() => {
+                    onSubmit({ result: 'cancelled' });
+                }}
+                variant="secondary"
+                icon="xmark"
+                label="Cancel"
+            />
+            <Button
+                data-testid="ai-chat-tool-confirm-accept"
+                onClick={() => {
+                    onSubmit({ result: 'confirmed' });
+                }}
+                className="flex-1"
+                variant="primary"
+                icon={icon}
+                label={label}
+            />
+        </AIToolContainer>
+    );
+}
diff --git a/packages/gitbook/src/components/AI/controls/SingleChoiceControl.tsx b/packages/gitbook/src/components/AI/controls/SingleChoiceControl.tsx
new file mode 100644
index 000000000..c51b3fc32
--- /dev/null
+++ b/packages/gitbook/src/components/AI/controls/SingleChoiceControl.tsx
@@ -0,0 +1,203 @@
+'use client';
+
+import { Button, Input } from '@/components/primitives';
+import { tcls } from '@/lib/tailwind';
+import * as React from 'react';
+import { z } from 'zod';
+import { AIToolContainer } from './common';
+import { type GetAIControlProps, createAIControl } from './helpers';
+
+const OTHER_OPTION_ID = '$other';
+
+export const SingleChoiceControlOutputSchema = z.object({
+    id: z.string().describe('The identifier of the option selected by the user.'),
+    input: z
+        .string()
+        .optional()
+        .describe('The custom text entered by the user when "Other" is enabled and selected.'),
+});
+
+export const SingleChoiceControlDef = createAIControl({
+    name: 'single-choice',
+    exposeAsTool: true,
+    description:
+        'Use this control when you need the user to choose exactly one option from a predefined list.',
+    inputSchema: z.object({
+        prompt: z
+            .string()
+            .describe(
+                'Provide the question or instruction that tells the user what single choice they need to make.'
+            ),
+        options: z
+            .array(
+                z
+                    .object({
+                        id: z
+                            .string()
+                            .describe(
+                                'Provide a unique, stable identifier for this option. This is the value returned to the agent when the user selects it.'
+                            ),
+                        label: z
+                            .string()
+                            .describe('Provide the short label the user sees for this option.'),
+                        description: z
+                            .string()
+                            .optional()
+                            .describe(
+                                'Optionally provide supporting details to help the user understand this option.'
+                            ),
+                    })
+                    .describe('Define one selectable option the user can pick.')
+            )
+            .describe(
+                'Provide the list of options the user can choose from. The user must select exactly one.'
+            ),
+        allowOther: z
+            .boolean()
+            .optional()
+            .describe(
+                'Set to true to let the user select an "Other" option and enter a custom text response.'
+            ),
+    }),
+    outputSchema: SingleChoiceControlOutputSchema,
+    render: (props) => {
+        return <SingleChoiceControl {...props} />;
+    },
+});
+
+function SingleChoiceControl(props: GetAIControlProps<typeof SingleChoiceControlDef>) {
+    const { prompt, options, allowOther, onSubmit } = props;
+    const [selectedId, setSelectedId] = React.useState<string | null>(null);
+    const [otherInput, setOtherInput] = React.useState('');
+
+    const canSubmit =
+        selectedId !== null &&
+        (selectedId !== OTHER_OPTION_ID || (allowOther && otherInput.trim().length > 0));
+
+    return (
+        <AIToolContainer className="flex w-full flex-col gap-2">
+            <div className="no-scrollbar flex flex-1 flex-col gap-1 overflow-auto">
+                <p className="mb-1 font-semibold text-sm">{prompt}</p>
+                {options.map((option) => {
+                    const isSelected = selectedId === option.id;
+                    return (
+                        <button
+                            key={option.id}
+                            type="button"
+                            data-testid={`ai-chat-tool-single-choice-option-${option.id}`}
+                            onClick={() => {
+                                setSelectedId(option.id);
+                            }}
+                            className={tcls(
+                                'circular-corners:rounded-3xl rounded-corners:rounded-xl border px-3 py-2 text-left transition-colors',
+                                isSelected
+                                    ? 'border-primary-original bg-primary-subtle text-tint-strong'
+                                    : 'border-tint bg-tint-base hover:bg-tint-subtle'
+                            )}
+                        >
+                            <div className="flex items-center gap-3">
+                                <span
+                                    aria-hidden
+                                    className={tcls(
+                                        'mt-0.5 flex size-5 shrink-0 items-center justify-center rounded-full border transition-colors',
+                                        isSelected
+                                            ? 'border-primary-original'
+                                            : 'border-tint-subtle'
+                                    )}
+                                >
+                                    <span
+                                        className={tcls(
+                                            'size-2.5 rounded-full transition-colors',
+                                            isSelected ? 'bg-primary-original' : 'bg-transparent'
+                                        )}
+                                    />
+                                </span>
+                                <span className="min-w-0">
+                                    <p className="font-medium text-sm">{option.label}</p>
+                                    {option.description ? (
+                                        <p className="mt-0.5 text-sm text-tint-subtle">
+                                            {option.description}
+                                        </p>
+                                    ) : null}
+                                </span>
+                            </div>
+                        </button>
+                    );
+                })}
+
+                {allowOther ? (
+                    <button
+                        type="button"
+                        data-testid="ai-chat-tool-single-choice-option-other"
+                        onClick={() => {
+                            setSelectedId(OTHER_OPTION_ID);
+                        }}
+                        className={tcls(
+                            'circular-corners:rounded-3xl rounded-corners:rounded-xl border px-3 py-2 text-left transition-colors',
+                            selectedId === OTHER_OPTION_ID
+                                ? 'border-primary-original bg-primary-subtle text-tint-strong'
+                                : 'border-tint bg-tint-base hover:bg-tint-subtle'
+                        )}
+                    >
+                        <div className="flex items-center gap-3">
+                            <span
+                                aria-hidden
+                                className={tcls(
+                                    'mt-0.5 flex size-5 shrink-0 items-center justify-center rounded-full border transition-colors',
+                                    selectedId === OTHER_OPTION_ID
+                                        ? 'border-primary-original'
+                                        : 'border-tint-subtle'
+                                )}
+                            >
+                                <span
+                                    className={tcls(
+                                        'size-2.5 rounded-full transition-colors',
+                                        selectedId === OTHER_OPTION_ID
+                                            ? 'bg-primary-original'
+                                            : 'bg-transparent'
+                                    )}
+                                />
+                            </span>
+                            <div className="flex flex-1 flex-col gap-1">
+                                <p className="font-medium text-sm">Other</p>
+                                {allowOther && selectedId === OTHER_OPTION_ID ? (
+                                    <Input
+                                        label="Enter your answer"
+                                        value={otherInput}
+                                        onValueChange={setOtherInput}
+                                        data-testid="ai-chat-tool-single-choice-other-input"
+                                        autoFocus
+                                        className="w-full"
+                                        sizing="small"
+                                    />
+                                ) : null}
+                            </div>
+                        </div>
+                    </button>
+                ) : null}
+            </div>
+
+            <Button
+                data-testid="ai-chat-tool-single-choice-submit"
+                variant="primary"
+                label="Submit answer"
+                disabled={!canSubmit}
+                onClick={() => {
+                    if (!canSubmit || !selectedId) {
+                        return;
+                    }
+
+                    if (selectedId === OTHER_OPTION_ID) {
+                        onSubmit({
+                            id: OTHER_OPTION_ID,
+                            input: otherInput.trim(),
+                        });
+                        return;
+                    }
+
+                    onSubmit({ id: selectedId });
+                }}
+            />
+        </AIToolContainer>
+    );
+}
diff --git a/packages/gitbook/src/components/AI/controls/common.tsx b/packages/gitbook/src/components/AI/controls/common.tsx
new file mode 100644
index 000000000..e1a6e7f9f
--- /dev/null
+++ b/packages/gitbook/src/components/AI/controls/common.tsx
@@ -0,0 +1,15 @@
+import { tcls } from '@/lib/tailwind';
+import type { ComponentPropsWithRef } from 'react';
+
+export function AIToolContainer(props: ComponentPropsWithRef<'div'>) {
+    return (
+        <div
+            {...props}
+            className={tcls(
+                'min-h-0 animate-present-slow circular-corners:rounded-3xl rounded-corners:rounded-xl border border-tint bg-tint-subtle p-2',
+                props.className
+            )}
+            style={{ animationDelay: '0.5s', ...props.style }}
+        />
+    );
+}
diff --git a/packages/gitbook/src/components/AI/controls/helpers.ts b/packages/gitbook/src/components/AI/controls/helpers.ts
new file mode 100644
index 000000000..65a348528
--- /dev/null
+++ b/packages/gitbook/src/components/AI/controls/helpers.ts
@@ -0,0 +1,80 @@
+import type {
+    AIStreamResponseToolCallPending,
+    AIToolCallResult,
+    AIToolDefinition,
+} from '@gitbook/api';
+import type { ZodType, z } from 'zod';
+import { zodToJsonSchema } from 'zod-to-json-schema';
+
+type AIUIToolContext = Pick<AIStreamResponseToolCallPending, 'toolCall' | 'toolCallId'>;
+
+type AIControlDefinition<
+    Name extends string = any,
+    Input extends Record<string, unknown> = any,
+    Output extends Record<string, unknown> = any,
+> = AIToolDefinition & {
+    createControl: (args: {
+        context: AIUIToolContext;
+        input: Input;
+        send: (result: Pick<AIToolCallResult, 'output' | 'summary'>) => Promise<void>;
+    }) => AIControl<Name, Input, Output>;
+    exposeAsTool: boolean;
+};
+
+type AIControlProps<Input = Record<string, unknown>, Output = Record<string, unknown>> = Input & {
+    onSubmit: (output: Output) => Promise<void>;
+};
+
+type AIControl<Name = string, Input = Record<string, unknown>, Output = Record<string, unknown>> = {
+    props: AIControlProps<Input, Output>;
+    name: Name;
+    context: AIUIToolContext;
+    render: () => React.ReactNode;
+};
+
+type AIControlFromDef<T extends AIControlDefinition> = ReturnType<T['createControl']>;
+
+export type GetAIControlProps<T extends AIControlDefinition> = AIControlFromDef<T>['props'];
+
+export function createAIControl<
+    Name extends string,
+    InputSchema extends ZodType<Record<string, unknown>>,
+    OutputSchema extends ZodType<Record<string, unknown>>,
+>(def: {
+    name: Name;
+    description: string;
+    inputSchema: InputSchema;
+    outputSchema: OutputSchema;
+    render: (props: AIControlProps<z.infer<InputSchema>, z.infer<OutputSchema>>) => React.ReactNode;
+    /**
+     * Indicates if the control should be exposed as a tool or not.
+     */
+    exposeAsTool: boolean;
+}): AIControlDefinition<Name, z.infer<InputSchema>, z.infer<OutputSchema>> {
+    return {
+        name: `ui--${def.name}`,
+        description: def.description,
+        inputSchema: zodToJsonSchema(def.inputSchema as any) as AIToolDefinition['inputSchema'],
+        createControl: ({ context, input, send }) => {
+            const props: AIControlProps<z.infer<InputSchema>, z.infer<OutputSchema>> = {
+                ...input,
+                onSubmit: async (output) => {
+                    await send({
+                        output,
+                        summary: {
+                            icon: 'check',
+                            text: 'Submitted',
+                        },
+                    });
+                },
+            };
+            return {
+                props,
+                name: def.name,
+                context,
+                render: () => def.render(props),
+            };
+        },
+        exposeAsTool: def.exposeAsTool,
+    };
+}
diff --git a/packages/gitbook/src/components/AI/controls/index.ts b/packages/gitbook/src/components/AI/controls/index.ts
new file mode 100644
index 000000000..afa29757c
--- /dev/null
+++ b/packages/gitbook/src/components/AI/controls/index.ts
@@ -0,0 +1,11 @@
+import { ConfirmControlDef } from './ConfirmControl';
+import { SingleChoiceControlDef } from './SingleChoiceControl';
+
+const CONTROLS = [SingleChoiceControlDef, ConfirmControlDef];
+
+export type AnyAIControlTool = (typeof CONTROLS)[number];
+export type AnyAIControl = ReturnType<AnyAIControlTool['createControl']>;
+
+export function getControlTools(): AnyAIControlTool[] {
+    return CONTROLS.filter((control) => control.exposeAsTool);
+}
diff --git a/packages/gitbook/src/components/AI/tools.ts b/packages/gitbook/src/components/AI/tools.ts
new file mode 100644
index 000000000..b21ba03bb
--- /dev/null
+++ b/packages/gitbook/src/components/AI/tools.ts
@@ -0,0 +1,8 @@
+import type { GitBookIntegrationTool } from '@gitbook/browser-types';
+import { integrationsAssistantTools } from '../Integrations';
+import { type AnyAIControlTool, getControlTools } from './controls';
+
+export function getTools(): (GitBookIntegrationTool | AnyAIControlTool)[] {
+    const integrationTools = integrationsAssistantTools.getState().tools;
+    return [...getControlTools(), ...integrationTools];
+}
diff --git a/packages/gitbook/src/components/AI/useAIChat.tsx b/packages/gitbook/src/components/AI/useAIChat.tsx
index 5d92dfc40..dd68b2c9e 100644
--- a/packages/gitbook/src/components/AI/useAIChat.tsx
+++ b/packages/gitbook/src/components/AI/useAIChat.tsx
@@ -9,13 +9,15 @@ import {
     type AIStreamResponseToolCallPending,
     type AIToolCallResult,
 } from '@gitbook/api';
-import type { IconName } from '@gitbook/icons';
+import assertNever from 'assert-never';
 import * as React from 'react';
 import { getVisitor, useTrackEvent } from '../Insights';
 import { getSession } from '../Insights/sessions';
-import { integrationsAssistantTools } from '../Integrations';
 import { useSetSearchState } from '../Search';
+import type { AnyAIControl } from './controls';
+import { ConfirmControlDef, ConfirmControlOutputSchema } from './controls/ConfirmControl';
 import { type RenderAIMessageOptions, streamAIChatResponse } from './server-actions';
+import { getTools } from './tools';
 import { useAIMessageContextRef } from './useAIMessageContext';
 
 export type AIChatMessage = {
@@ -24,21 +26,6 @@ export type AIChatMessage = {
     query?: string;
 };
 
-export type AIChatPendingTool = {
-    icon?: IconName;
-    label: string;
-
-    /**
-     * Confirm the tool call by calling this function.
-     */
-    confirm: () => Promise<void>;
-
-    /**
-     * Tool call result to cancel it.
-     */
-    cancelToolCall: AIToolCallResult;
-};
-
 export type AIChatState = {
     /**
      * If true, the chat is open.
@@ -71,9 +58,9 @@ export type AIChatState = {
     followUpSuggestions: string[];
 
     /**
-     * Tools that are pending confirmation to be executed.
+     * Control to be displayed to the user.
      */
-    pendingTools: AIChatPendingTool[];
+    control: AnyAIControl | null;
 
     /**
      * If true, the session is in progress.
@@ -127,7 +114,7 @@ const globalState = zustand.create<AIChatState>(() => {
         messages: [],
         query: null,
         followUpSuggestions: [],
-        pendingTools: [],
+        control: null,
         loading: false,
         error: false,
         initialQuery: null,
@@ -213,7 +200,7 @@ export function AIChatProvider(props: {
                 return {
                     ...state,
                     followUpSuggestions: [],
-                    pendingTools: [],
+                    control: null,
                     loading: true,
                     error: false,
                     messages: [
@@ -228,10 +215,10 @@ export function AIChatProvider(props: {
 
             // Execute a tool call
             const executeToolCall = async (event: AIStreamResponseToolCallPending) => {
-                const integrationTools = integrationsAssistantTools.getState().tools;
-                const toolDef = integrationTools.find((tool) => tool.name === event.toolCall.tool);
+                const tools = getTools();
+                const toolDef = tools.find((tool) => tool.name === event.toolCall.tool);
 
-                if (!toolDef) {
+                if (!toolDef || !('execute' in toolDef)) {
                     throw new Error(`Tool ${event.toolCall.tool} not found`);
                 }
 
@@ -264,7 +251,7 @@ export function AIChatProvider(props: {
 
             let toolToExecute: AIStreamResponseToolCallPending | null = null;
             try {
-                const integrationTools = integrationsAssistantTools.getState().tools;
+                const tools = getTools();
                 const stream = await streamAIChatResponse({
                     message: input.message,
                     toolCall: input.toolCall,
@@ -274,9 +261,10 @@ export function AIChatProvider(props: {
                         sessionId: getSession().id,
                         visitorId: (await getVisitor()).deviceId,
                     },
-                    tools: integrationTools.map((tool) => ({
+                    tools: tools.map((tool) => ({
                         name: tool.name,
                         description: tool.description,
+                        // Issue with the schema generated by Zod and Next.js serialization.
                         inputSchema: tool.inputSchema,
                     })),
                     options: {
@@ -320,46 +308,84 @@ export function AIChatProvider(props: {
                             break;
                         }
                         case 'response_tool_call_pending': {
-                            const toolDef = integrationTools.find(
-                                (tool) => tool.name === event.toolCall.tool
-                            );
+                            const toolDef = tools.find((tool) => tool.name === event.toolCall.tool);
                             if (!toolDef) {
                                 throw new Error(`Tool ${event.toolCall.tool} not found`);
                             }
 
-                            const confirmation = toolDef.confirmation;
+                            if ('createControl' in toolDef) {
+                                globalState.setState((state) => ({
+                                    ...state,
+                                    control: toolDef.createControl({
+                                        context: {
+                                            toolCall: event.toolCall,
+                                            toolCallId: event.toolCallId,
+                                        },
+                                        input: event.toolCall.input as any,
+                                        send: async (result) => {
+                                            await streamResponse({
+                                                toolCall: {
+                                                    tool: event.toolCall.tool,
+                                                    toolCallId: event.toolCallId,
+                                                    output: result.output,
+                                                    summary: result.summary,
+                                                },
+                                            });
+                                        },
+                                    }),
+                                }));
+                                break;
+                            }
+
+                            const confirmation = 'confirmation' in toolDef && toolDef.confirmation;
                             if (confirmation) {
                                 globalState.setState((state) => ({
                                     ...state,
-                                    pendingTools: [
-                                        ...state.pendingTools,
-                                        {
-                                            icon: confirmation.icon,
-                                            label: confirmation.label,
-                                            cancelToolCall: {
-                                                tool: event.toolCall.tool,
-                                                toolCallId: event.toolCallId,
-                                                output: {
-                                                    cancelled: 'User did not confirm the tool call',
-                                                },
-                                                summary: {
-                                                    icon: 'forward',
-                                                    text: tString(
-                                                        language,
-                                                        'tool_call_skipped',
-                                                        confirmation.label
-                                                    ),
-                                                },
-                                            },
-                                            confirm: async () => {
-                                                await executeToolCall(event);
-                                            },
+                                    control: ConfirmControlDef.createControl({
+                                        context: {
+                                            toolCall: event.toolCall,
+                                            toolCallId: event.toolCallId,
                                         },
-                                    ],
+                                        input: {
+                                            label: confirmation.label,
+                                            icon: confirmation.icon,
+                                        },
+                                        send: async (result) => {
+                                            const output = ConfirmControlOutputSchema.parse(
+                                                result.output
+                                            );
+                                            switch (output.result) {
+                                                case 'cancelled': {
+                                                    await streamResponse({
+                                                        toolCall: {
+                                                            tool: event.toolCall.tool,
+                                                            toolCallId: event.toolCallId,
+                                                            output: { cancelled: true },
+                                                            summary: {
+                                                                icon: 'forward',
+                                                                text: tString(
+                                                                    language,
+                                                                    'tool_call_skipped',
+                                                                    confirmation.label
+                                                                ),
+                                                            },
+                                                        },
+                                                    });
+                                                    break;
+                                                }
+                                                case 'confirmed':
+                                                    await executeToolCall(event);
+                                                    break;
+                                                default:
+                                                    assertNever(output.result);
+                                            }
+                                        },
+                                    }),
                                 }));
-                            } else {
-                                toolToExecute = event;
+                                break;
                             }
+
+                            toolToExecute = event;
                             break;
                         }
                     }
@@ -408,7 +434,11 @@ export function AIChatProvider(props: {
     // Post a message to the AI chat
     const onPostMessage = React.useCallback(
         async (input: { message: string }) => {
-            const { query, messages, pendingTools } = globalState.getState();
+            const { query, messages, control } = globalState.getState();
+
+            if (control) {
+                throw new Error("We can't post a message when a control is active");
+            }
 
             // For first message, update the ask parameter in URL
             if (messages.length === 0) {
@@ -453,15 +483,9 @@ export function AIChatProvider(props: {
                 };
             });
 
-            const pendingTool = pendingTools[0];
-            streamResponse({
-                message: input.message,
-                // If we had a pending tool call, we need to send it as being cancelled
-                // otherwise the AI will fail to process the message
-                ...(pendingTool ? { toolCall: pendingTool.cancelToolCall } : {}),
-            });
+            streamResponse({ message: input.message });
         },
-        [setSearchState, trackEvent, streamResponse]
+        [setSearchState, trackEvent, streamResponse, language]
     );
 
     // Clear the conversation and reset ask parameter
@@ -472,7 +496,7 @@ export function AIChatProvider(props: {
             messages: [],
             query: null,
             followUpSuggestions: [],
-            pendingTools: [],
+            control: null,
             responseId: null,
             error: false,
             initialQuery: null,
diff --git a/packages/gitbook/src/components/AIChat/AIChat.tsx b/packages/gitbook/src/components/AIChat/AIChat.tsx
index 762b055ff..28ef35af0 100644
--- a/packages/gitbook/src/components/AIChat/AIChat.tsx
+++ b/packages/gitbook/src/components/AIChat/AIChat.tsx
@@ -28,6 +28,7 @@ import { useNow } from '../hooks';
 import { Button } from '../primitives';
 import { ScrollContainer } from '../primitives/ScrollContainer';
 import { SideSheet } from '../primitives/SideSheet';
+import { AIChatControl } from './AIChatControl';
 import { AIChatControlButton } from './AIChatControlButton';
 import { AIChatIcon } from './AIChatIcon';
 import { AIChatInput } from './AIChatInput';
@@ -142,7 +143,7 @@ export function AIChatDynamicIcon(props: {
                           ? 'working'
                           : 'thinking'
                       : chat.messages.length > 0
-                        ? chat.pendingTools.length > 0
+                        ? chat.control
                             ? 'confirm'
                             : 'done'
                         : 'default'
@@ -224,7 +225,7 @@ export function AIChatBody(props: {
     return (
         <>
             <ScrollContainer
-                className="shrink grow basis-80 animate-fade-in-slow [container-type:size]"
+                className="min-h-[20%] shrink grow animate-fade-in-slow [container-type:size]"
                 contentClassName="py-4 gutter-stable flex flex-col gap-4"
                 orientation="vertical"
                 trailing={{ fade: false, button: true }}
@@ -273,17 +274,21 @@ export function AIChatBody(props: {
                 )}
             </ScrollContainer>
 
-            <div className="flex flex-col gap-2 pb-4">
+            <div className="flex min-h-0 flex-col gap-2 pb-4">
                 {/* Display an error banner when something went wrong. */}
                 {chat.error ? <AIChatError chatController={chatController} /> : null}
 
-                <AIChatInput
-                    loading={chat.loading}
-                    disabled={chat.loading || chat.error}
-                    onSubmit={(value) => {
-                        chatController.postMessage({ message: value });
-                    }}
-                />
+                {chat.control ? (
+                    <AIChatControl control={chat.control} />
+                ) : (
+                    <AIChatInput
+                        loading={chat.loading}
+                        disabled={chat.loading || chat.error}
+                        onSubmit={(value) => {
+                            chatController.postMessage({ message: value });
+                        }}
+                    />
+                )}
             </div>
         </>
     );
diff --git a/packages/gitbook/src/components/AIChat/AIChatControl.tsx b/packages/gitbook/src/components/AIChat/AIChatControl.tsx
new file mode 100644
index 000000000..0b350f7ee
--- /dev/null
+++ b/packages/gitbook/src/components/AIChat/AIChatControl.tsx
@@ -0,0 +1,8 @@
+'use client';
+
+import type { AnyAIControl } from '../AI/controls';
+
+export function AIChatControl(props: { control: AnyAIControl }) {
+    const { control } = props;
+    return control.render();
+}
diff --git a/packages/gitbook/src/components/AIChat/AIChatMessages.tsx b/packages/gitbook/src/components/AIChat/AIChatMessages.tsx
index e120ed8e6..abdaed9ba 100644
--- a/packages/gitbook/src/components/AIChat/AIChatMessages.tsx
+++ b/packages/gitbook/src/components/AIChat/AIChatMessages.tsx
@@ -3,7 +3,6 @@ import { tString } from '@/intl/translate';
 import { tcls } from '@/lib/tailwind';
 import { AIMessageRole } from '@gitbook/api';
 import type { AIChatController, AIChatState } from '../AI';
-import { AIChatToolConfirmations } from './AIChatToolConfirmations';
 import { AIResponseFeedback } from './AIResponseFeedback';
 import { AIChatFollowupSuggestions } from './AiChatFollowupSuggestions';
 
@@ -67,7 +66,6 @@ export function AIChatMessages(props: {
                                 'break-words',
                                 'group/message',
                                 'animate-blur-in-slow',
-                                isLastMessage ? 'basis-full' : '',
                                 message.role === AIMessageRole.User
                                     ? 'max-w-[80%] origin-top-right self-end circular-corners:rounded-2xl rounded-corners:rounded-md bg-tint px-4 py-2'
                                     : 'origin-top-left text-tint-strong'
@@ -96,14 +94,13 @@ export function AIChatMessages(props: {
                                     !chat.error &&
                                     chat.query &&
                                     chat.responseId &&
-                                    chat.pendingTools.length === 0 ? (
+                                    chat.control ? (
                                         <AIResponseFeedback
                                             responseId={chat.responseId}
                                             query={chat.query}
                                             className="-ml-1 -mt-4"
                                         />
                                     ) : null}
-                                    <AIChatToolConfirmations chat={chat} />
                                     <AIChatFollowupSuggestions
                                         chat={chat}
                                         chatController={chatController}
diff --git a/packages/gitbook/src/components/AIChat/AIChatToolConfirmations.tsx b/packages/gitbook/src/components/AIChat/AIChatToolConfirmations.tsx
deleted file mode 100644
index d61585999..000000000
--- a/packages/gitbook/src/components/AIChat/AIChatToolConfirmations.tsx
+++ /dev/null
@@ -1,70 +0,0 @@
-'use client';
-import { useLanguage } from '@/intl/client';
-import { t } from '@/intl/translate';
-import { useHotkeys } from 'react-hotkeys-hook';
-import type { AIChatState } from '../AI';
-import { Button } from '../primitives';
-import { KeyboardShortcut } from '../primitives/KeyboardShortcut';
-
-/**
- * Display buttons to confirm tool calls.
- */
-export function AIChatToolConfirmations(props: {
-    chat: AIChatState;
-}) {
-    const { chat } = props;
-
-    const language = useLanguage();
-
-    useHotkeys(
-        'mod+enter',
-        (e) => {
-            e.preventDefault();
-            chat.pendingTools[0]?.confirm();
-        },
-        {
-            enableOnFormTags: true,
-        },
-        [chat.pendingTools]
-    );
-
-    return (
-        <div className="flex w-full flex-wrap justify-end gap-2 empty:hidden">
-            {chat.pendingTools.map((tool, index) => (
-                <div
-                    className="flex w-full animate-present-slow flex-col gap-1"
-                    key={index}
-                    style={{ animationDelay: `${0.5 + index * 0.1}s` }}
-                >
-                    <Button
-                        data-testid="ai-chat-tool-confirmation"
-                        onClick={() => {
-                            tool.confirm();
-                        }}
-                        tabIndex={index}
-                        label={tool.label}
-                        className="w-full justify-center"
-                        size={index === 0 ? 'medium' : 'small'}
-                        variant={index === 0 ? 'primary' : 'secondary'}
-                        icon={tool.icon}
-                    />
-                    {index === 0 && (
-                        <div
-                            className="flex pointer-none:hidden w-full animate-fade-in-slow items-center justify-end gap-2 text-tint text-xs"
-                            style={{ animationDelay: '1000ms' }}
-                        >
-                            {t(
-                                language,
-                                'press_to_confirm',
-                                <KeyboardShortcut
-                                    keys={['mod', 'enter']}
-                                    className="mx-0 text-tint"
-                                />
-                            )}
-                        </div>
-                    )}
-                </div>
-            ))}
-        </div>
-    );
-}
diff --git a/packages/gitbook/src/components/primitives/Input.tsx b/packages/gitbook/src/components/primitives/Input.tsx
index 8a3f03f3f..c8ffdb3c1 100644
--- a/packages/gitbook/src/components/primitives/Input.tsx
+++ b/packages/gitbook/src/components/primitives/Input.tsx
@@ -57,11 +57,11 @@ export type InputProps = CustomInputProps &
     (
         | ({ multiline?: false } & Omit<
               React.InputHTMLAttributes<HTMLInputElement>,
-              'value' | 'onChange'
+              'value' | 'onChange' | 'onSubmit'
           >)
         | ({ multiline: true } & Omit<
               React.TextareaHTMLAttributes<HTMLTextAreaElement>,
-              'value' | 'onChange'
+              'value' | 'onChange' | 'onSubmit'
           >)
     );