import { StrictMode, Suspense, act } from "react";
import { describe, expect, it, vi } from "vitest";
import { render } from "vitest-browser-react";
import type { UIMessage } from "ai";
import {
  useAgentChat,
  type PrepareSendMessagesRequestOptions,
  type PrepareSendMessagesRequestResult,
  type AITool
} from "../react";
import type { useAgent } from "agents/react";

function sleep(ms: number) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

function createAgent({
  name,
  url,
  send
}: {
  name: string;
  url: string;
  send?: (data: string) => void;
}) {
  const target = new EventTarget();
  const baseAgent = {
    _pkurl: url,
    _pk: name, // Use name as pk to distinguish agents
    _url: null as string | null,
    addEventListener: target.addEventListener.bind(target),
    agent: "Chat",
    close: () => {},
    id: "fake-agent",
    name,
    removeEventListener: target.removeEventListener.bind(target),
    send: send ?? (() => {}),
    dispatchEvent: target.dispatchEvent.bind(target)
  };
  return baseAgent as unknown as ReturnType<typeof useAgent>;
}

describe("useAgentChat", () => {
  it("should cache initial message responses across re-renders", async () => {
    const agent = createAgent({
      name: "thread-alpha",
      url: "ws://localhost:3000/agents/chat/thread-alpha?_pk=abc"
    });

    const testMessages = [
      {
        id: "1",
        role: "user" as const,
        parts: [{ type: "text" as const, text: "Hi" }]
      },
      {
        id: "2",
        role: "assistant" as const,
        parts: [{ type: "text" as const, text: "Hello" }]
      }
    ];

    const getInitialMessages = vi.fn(() => Promise.resolve(testMessages));

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages
      });
      return <div data-testid="messages">{JSON.stringify(chat.messages)}</div>;
    };

    const suspenseRendered = vi.fn();
    const SuspenseObserver = () => {
      suspenseRendered();
      return "Suspended";
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback={<SuspenseObserver />}>{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    await expect
      .element(screen.getByTestId("messages"))
      .toHaveTextContent(JSON.stringify(testMessages));

    expect(getInitialMessages).toHaveBeenCalledTimes(1);
    expect(suspenseRendered).toHaveBeenCalled();

    suspenseRendered.mockClear();

    await screen.rerender(<TestComponent />);

    await expect
      .element(screen.getByTestId("messages"))
      .toHaveTextContent(JSON.stringify(testMessages));

    expect(getInitialMessages).toHaveBeenCalledTimes(1);
    expect(suspenseRendered).not.toHaveBeenCalled();
  });

  it("should refetch initial messages when the agent name changes", async () => {
    const url = "ws://localhost:3000/agents/chat/thread-a?_pk=abc";
    const agentA = createAgent({ name: "thread-a", url });
    const agentB = createAgent({ name: "thread-b", url });

    const getInitialMessages = vi.fn(async ({ name }: { name: string }) => [
      {
        id: "1",
        role: "assistant" as const,
        parts: [{ type: "text" as const, text: `Hello from ${name}` }]
      }
    ]);

    const TestComponent = ({
      agent
    }: {
      agent: ReturnType<typeof useAgent>;
    }) => {
      const chat = useAgentChat({
        agent,
        getInitialMessages
      });
      return <div data-testid="messages">{JSON.stringify(chat.messages)}</div>;
    };

    const suspenseRendered = vi.fn();
    const SuspenseObserver = () => {
      suspenseRendered();
      return "Suspended";
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent agent={agentA} />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback={<SuspenseObserver />}>{children}</Suspense>
          </StrictMode>
        )
      });

      await sleep(10);
      return screen;
    });

    await expect
      .element(screen.getByTestId("messages"))
      .toHaveTextContent("Hello from thread-a");

    expect(getInitialMessages).toHaveBeenCalledTimes(1);
    expect(getInitialMessages).toHaveBeenNthCalledWith(
      1,
      expect.objectContaining({ name: "thread-a" })
    );

    suspenseRendered.mockClear();

    await act(async () => {
      screen.rerender(<TestComponent agent={agentB} />);
      await sleep(10);
    });

    await expect
      .element(screen.getByTestId("messages"))
      .toHaveTextContent("Hello from thread-b");

    expect(getInitialMessages).toHaveBeenCalledTimes(2);
    expect(getInitialMessages).toHaveBeenNthCalledWith(
      2,
      expect.objectContaining({ name: "thread-b" })
    );
  });

  it("should accept prepareSendMessagesRequest option without errors", async () => {
    const agent = createAgent({
      name: "thread-with-tools",
      url: "ws://localhost:3000/agents/chat/thread-with-tools?_pk=abc"
    });

    const prepareSendMessagesRequest = vi.fn(
      (
        _options: PrepareSendMessagesRequestOptions<UIMessage>
      ): PrepareSendMessagesRequestResult => ({
        body: {
          clientTools: [
            {
              name: "showAlert",
              description: "Shows an alert to the user",
              parameters: { message: { type: "string" } }
            }
          ]
        },
        headers: {
          "X-Client-Tool-Count": "1"
        }
      })
    );

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null, // Skip fetching initial messages
        prepareSendMessagesRequest
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    const screen = await act(() =>
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      })
    );

    // Verify component renders without errors
    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("0");
  });

  it("should handle async prepareSendMessagesRequest", async () => {
    const agent = createAgent({
      name: "thread-async-prepare",
      url: "ws://localhost:3000/agents/chat/thread-async-prepare?_pk=abc"
    });

    const prepareSendMessagesRequest = vi.fn(
      async (
        _options: PrepareSendMessagesRequestOptions<UIMessage>
      ): Promise<PrepareSendMessagesRequestResult> => {
        // Simulate async operation like fetching tool definitions
        await sleep(10);
        return {
          body: {
            clientTools: [
              { name: "navigateToPage", description: "Navigates to a page" }
            ]
          }
        };
      }
    );

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        prepareSendMessagesRequest
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    const screen = await act(() =>
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      })
    );

    // Verify component renders without errors
    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("0");
  });

  it("should auto-extract schemas from tools with execute functions", async () => {
    const agent = createAgent({
      name: "thread-client-tools",
      url: "ws://localhost:3000/agents/chat/thread-client-tools?_pk=abc"
    });

    // Tools with execute functions have their schemas auto-extracted and sent to server
    const tools: Record<string, AITool<unknown, unknown>> = {
      showAlert: {
        description: "Shows an alert dialog to the user",
        parameters: {
          type: "object",
          properties: {
            message: { type: "string", description: "The message to display" }
          },
          required: ["message"]
        },
        execute: async (input) => {
          // Client-side execution
          const { message } = input as { message: string };
          return { shown: true, message };
        }
      },
      changeBackgroundColor: {
        description: "Changes the page background color",
        parameters: {
          type: "object",
          properties: {
            color: { type: "string" }
          }
        },
        execute: async (input) => {
          const { color } = input as { color: string };
          return { success: true, color };
        }
      }
    };

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        tools
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    const screen = await act(() =>
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      })
    );

    // Verify component renders without errors
    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("0");
  });

  it("should combine auto-extracted tools with prepareSendMessagesRequest", async () => {
    const agent = createAgent({
      name: "thread-combined",
      url: "ws://localhost:3000/agents/chat/thread-combined?_pk=abc"
    });

    const tools: Record<string, AITool> = {
      showAlert: {
        description: "Shows an alert",
        execute: async () => ({ shown: true })
      }
    };

    const prepareSendMessagesRequest = vi.fn(
      (
        _options: PrepareSendMessagesRequestOptions<UIMessage>
      ): PrepareSendMessagesRequestResult => ({
        body: {
          customData: "extra-context",
          userTimezone: "America/New_York"
        },
        headers: {
          "X-Custom-Header": "custom-value"
        }
      })
    );

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        tools,
        prepareSendMessagesRequest
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    const screen = await act(() =>
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      })
    );

    // Verify component renders without errors
    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("0");
  });

  it("should work with tools that have execute functions for client-side execution", async () => {
    const agent = createAgent({
      name: "thread-tools-execution",
      url: "ws://localhost:3000/agents/chat/thread-tools-execution?_pk=abc"
    });

    const mockExecute = vi.fn().mockResolvedValue({ success: true });

    // Single unified tools object - schema + execute in one place
    const tools: Record<string, AITool> = {
      showAlert: {
        description: "Shows an alert",
        parameters: {
          type: "object",
          properties: { message: { type: "string" } }
        },
        execute: mockExecute
      }
    };

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        tools
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    const screen = await act(() =>
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      })
    );

    // Verify component renders without errors
    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("0");
  });
});

describe("useAgentChat client-side tool execution (issue #728)", () => {
  it("should update tool part state from input-available to output-available when addToolResult is called", async () => {
    const agent = createAgent({
      name: "tool-state-test",
      url: "ws://localhost:3000/agents/chat/tool-state-test?_pk=abc"
    });

    const mockExecute = vi.fn().mockResolvedValue({ location: "New York" });

    // Initial messages with a tool call in input-available state
    const initialMessages: UIMessage[] = [
      {
        id: "msg-1",
        role: "user",
        parts: [{ type: "text", text: "Where am I?" }]
      },
      {
        id: "msg-2",
        role: "assistant",
        parts: [
          {
            type: "tool-getLocation",
            toolCallId: "tool-call-1",
            state: "input-available",
            input: {}
          }
        ]
      }
    ];

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: () => Promise.resolve(initialMessages),
        experimental_automaticToolResolution: true,
        tools: {
          getLocation: {
            execute: mockExecute
          }
        }
      });

      // Find the tool part to check its state
      const assistantMsg = chat.messages.find((m) => m.role === "assistant");
      const toolPart = assistantMsg?.parts.find(
        (p) => "toolCallId" in p && p.toolCallId === "tool-call-1"
      );
      const toolState =
        toolPart && "state" in toolPart ? toolPart.state : "not-found";

      return (
        <div>
          <div data-testid="messages-count">{chat.messages.length}</div>
          <div data-testid="tool-state">{toolState}</div>
        </div>
      );
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      // The tool should have been automatically executed
      await sleep(10);
      return screen;
    });

    // Wait for initial messages to load
    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("2");

    // Verify the tool execute was called
    expect(mockExecute).toHaveBeenCalled();

    // the tool part should be updated to output-available
    // in the SAME message (msg-2), not in a new message
    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("2"); // Should still be 2 messages, not 3

    // The tool state should be output-available after addToolResult
    await expect
      .element(screen.getByTestId("tool-state"))
      .toHaveTextContent("output-available");
  });

  it("should not create duplicate tool parts when client executes tool", async () => {
    const agent = createAgent({
      name: "duplicate-test",
      url: "ws://localhost:3000/agents/chat/duplicate-test?_pk=abc"
    });

    const mockExecute = vi.fn().mockResolvedValue({ confirmed: true });

    const initialMessages: UIMessage[] = [
      {
        id: "msg-1",
        role: "assistant",
        parts: [
          { type: "text", text: "Should I proceed?" },
          {
            type: "tool-askForConfirmation",
            toolCallId: "confirm-1",
            state: "input-available",
            input: { message: "Proceed with action?" }
          }
        ]
      }
    ];

    let chatInstance: ReturnType<typeof useAgentChat> | null = null;

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: () => Promise.resolve(initialMessages),
        tools: {
          askForConfirmation: {
            execute: mockExecute
          }
        }
      });
      chatInstance = chat;

      // Count tool parts with this toolCallId
      const toolPartsCount = chat.messages.reduce((count, msg) => {
        return (
          count +
          msg.parts.filter(
            (p) => "toolCallId" in p && p.toolCallId === "confirm-1"
          ).length
        );
      }, 0);

      // Get the tool state
      const toolPart = chat.messages
        .flatMap((m) => m.parts)
        .find((p) => "toolCallId" in p && p.toolCallId === "confirm-1");
      const toolState =
        toolPart && "state" in toolPart ? toolPart.state : "not-found";

      return (
        <div>
          <div data-testid="messages-count">{chat.messages.length}</div>
          <div data-testid="tool-parts-count">{toolPartsCount}</div>
          <div data-testid="tool-state">{toolState}</div>
        </div>
      );
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("1");

    // Manually trigger addToolResult to simulate user confirming
    await act(async () => {
      if (chatInstance) {
        await chatInstance.addToolResult({
          tool: "askForConfirmation",
          toolCallId: "confirm-1",
          output: { confirmed: true }
        });
      }
    });

    // There should still be exactly ONE tool part with this toolCallId
    await expect
      .element(screen.getByTestId("tool-parts-count"))
      .toHaveTextContent("1");

    // The tool state should be updated to output-available
    await expect
      .element(screen.getByTestId("tool-state"))
      .toHaveTextContent("output-available");
  });
});

describe("useAgentChat setMessages", () => {
  it("should handle functional updater and sync resolved messages to server", async () => {
    const sentMessages: string[] = [];
    const agent = createAgent({
      name: "set-messages-test",
      url: "ws://localhost:3000/agents/chat/set-messages-test?_pk=abc",
      send: (data: string) => sentMessages.push(data)
    });

    const initialMessages: UIMessage[] = [
      {
        id: "msg-1",
        role: "user",
        parts: [{ type: "text", text: "Hello" }]
      },
      {
        id: "msg-2",
        role: "assistant",
        parts: [{ type: "text", text: "Hi there!" }]
      }
    ];

    let chatInstance: ReturnType<typeof useAgentChat> | null = null;

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: () => Promise.resolve(initialMessages)
      });
      chatInstance = chat;
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("2");

    // Use functional updater to append a message
    const newMessage: UIMessage = {
      id: "msg-3",
      role: "user",
      parts: [{ type: "text", text: "Follow up" }]
    };

    await act(async () => {
      chatInstance!.setMessages((prev) => [...prev, newMessage]);
      await sleep(10);
    });

    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("3");

    // Verify the server received the RESOLVED messages (not empty array)
    const chatMessagesSent = sentMessages
      .map((m) => JSON.parse(m))
      .filter((m) => m.type === "cf_agent_chat_messages");

    expect(chatMessagesSent.length).toBeGreaterThan(0);
    const lastSent = chatMessagesSent[chatMessagesSent.length - 1];
    // Should have the full 3 messages, NOT an empty array
    expect(lastSent.messages.length).toBe(3);
    expect(lastSent.messages[2].id).toBe("msg-3");
  });

  it("should handle array setMessages and sync to server", async () => {
    const sentMessages: string[] = [];
    const agent = createAgent({
      name: "set-messages-array-test",
      url: "ws://localhost:3000/agents/chat/set-messages-array-test?_pk=abc",
      send: (data: string) => sentMessages.push(data)
    });

    let chatInstance: ReturnType<typeof useAgentChat> | null = null;

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: [] as UIMessage[]
      });
      chatInstance = chat;
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
    });

    // Set messages with an array directly
    const newMessages: UIMessage[] = [
      {
        id: "arr-1",
        role: "user",
        parts: [{ type: "text", text: "Replaced" }]
      }
    ];

    await act(async () => {
      chatInstance!.setMessages(newMessages);
      await sleep(10);
    });

    // Verify the server received the array
    const chatMessagesSent = sentMessages
      .map((m) => JSON.parse(m))
      .filter((m) => m.type === "cf_agent_chat_messages");

    expect(chatMessagesSent.length).toBeGreaterThan(0);
    const lastSent = chatMessagesSent[chatMessagesSent.length - 1];
    expect(lastSent.messages.length).toBe(1);
    expect(lastSent.messages[0].id).toBe("arr-1");
  });
});

describe("useAgentChat clearHistory", () => {
  it("should clear local state and send CF_AGENT_CHAT_CLEAR to server", async () => {
    const sentMessages: string[] = [];
    const agent = createAgent({
      name: "clear-test",
      url: "ws://localhost:3000/agents/chat/clear-test?_pk=abc",
      send: (data: string) => sentMessages.push(data)
    });

    const initialMessages: UIMessage[] = [
      {
        id: "clear-1",
        role: "user",
        parts: [{ type: "text", text: "Hello" }]
      }
    ];

    let chatInstance: ReturnType<typeof useAgentChat> | null = null;

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: () => Promise.resolve(initialMessages)
      });
      chatInstance = chat;
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("1");

    await act(async () => {
      chatInstance!.clearHistory();
      await sleep(10);
    });

    await expect
      .element(screen.getByTestId("messages-count"))
      .toHaveTextContent("0");

    // Verify CF_AGENT_CHAT_CLEAR was sent
    const clearMessages = sentMessages
      .map((m) => JSON.parse(m))
      .filter((m) => m.type === "cf_agent_chat_clear");
    expect(clearMessages.length).toBe(1);
  });
});

describe("useAgentChat autoContinueAfterToolResult default", () => {
  it("should send autoContinue: true by default with tool results", async () => {
    const sentMessages: string[] = [];
    const agent = createAgent({
      name: "auto-continue-default",
      url: "ws://localhost:3000/agents/chat/auto-continue-default?_pk=abc",
      send: (data: string) => sentMessages.push(data)
    });

    const initialMessages: UIMessage[] = [
      {
        id: "msg-1",
        role: "assistant",
        parts: [
          {
            type: "tool-getLocation",
            toolCallId: "tc-default-1",
            state: "input-available",
            input: {}
          }
        ]
      }
    ];

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: () => Promise.resolve(initialMessages),
        // No explicit autoContinueAfterToolResult — should default to true
        onToolCall: ({ toolCall, addToolOutput }) => {
          addToolOutput({
            toolCallId: toolCall.toolCallId,
            output: { lat: 51.5, lng: -0.1 }
          });
        }
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(50);
    });

    // Find the CF_AGENT_TOOL_RESULT message
    const toolResultMessages = sentMessages
      .map((m) => JSON.parse(m))
      .filter((m) => m.type === "cf_agent_tool_result");

    expect(toolResultMessages.length).toBeGreaterThanOrEqual(1);
    // Default should be autoContinue: true
    expect(toolResultMessages[0].autoContinue).toBe(true);
  });

  it("should send autoContinue: false when explicitly disabled", async () => {
    const sentMessages: string[] = [];
    const agent = createAgent({
      name: "auto-continue-disabled",
      url: "ws://localhost:3000/agents/chat/auto-continue-disabled?_pk=abc",
      send: (data: string) => sentMessages.push(data)
    });

    const initialMessages: UIMessage[] = [
      {
        id: "msg-1",
        role: "assistant",
        parts: [
          {
            type: "tool-getLocation",
            toolCallId: "tc-disabled-1",
            state: "input-available",
            input: {}
          }
        ]
      }
    ];

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: () => Promise.resolve(initialMessages),
        autoContinueAfterToolResult: false, // Explicitly disabled
        onToolCall: ({ toolCall, addToolOutput }) => {
          addToolOutput({
            toolCallId: toolCall.toolCallId,
            output: { lat: 51.5, lng: -0.1 }
          });
        }
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(50);
    });

    const toolResultMessages = sentMessages
      .map((m) => JSON.parse(m))
      .filter((m) => m.type === "cf_agent_tool_result");

    expect(toolResultMessages.length).toBeGreaterThanOrEqual(1);
    expect(toolResultMessages[0].autoContinue).toBe(false);
  });

  it("should send autoContinue: true by default with tool approvals", async () => {
    const sentMessages: string[] = [];
    const agent = createAgent({
      name: "auto-continue-approval",
      url: "ws://localhost:3000/agents/chat/auto-continue-approval?_pk=abc",
      send: (data: string) => sentMessages.push(data)
    });

    // Tool part must have approval.id so the wrapper can find the toolCallId
    const initialMessages: UIMessage[] = [
      {
        id: "msg-1",
        role: "assistant",
        parts: [
          {
            type: "tool-dangerousAction",
            toolCallId: "tc-approval-1",
            state: "approval-requested",
            input: { action: "delete" },
            approval: { id: "approval-req-1" }
          }
        ]
      }
    ];

    let chatInstance: ReturnType<typeof useAgentChat> | null = null;

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: () => Promise.resolve(initialMessages)
        // No explicit autoContinueAfterToolResult — should default to true
      });
      chatInstance = chat;
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(50);
    });

    // Send approval via the hook using the approval request ID
    await act(async () => {
      if (chatInstance) {
        chatInstance.addToolApprovalResponse({
          id: "approval-req-1",
          approved: true
        });
      }
      await sleep(10);
    });

    // Find the CF_AGENT_TOOL_APPROVAL message
    const approvalMessages = sentMessages
      .map((m) => JSON.parse(m))
      .filter((m) => m.type === "cf_agent_tool_approval");

    expect(approvalMessages.length).toBeGreaterThanOrEqual(1);
    expect(approvalMessages[0].autoContinue).toBe(true);
    expect(approvalMessages[0].approved).toBe(true);
  });
});

describe("useAgentChat onToolCall", () => {
  it("should fire onToolCall for input-available tool parts", async () => {
    const agent = createAgent({
      name: "ontoolcall-test",
      url: "ws://localhost:3000/agents/chat/ontoolcall-test?_pk=abc",
      send: () => {}
    });

    const toolCallReceived = vi.fn();

    const initialMessages: UIMessage[] = [
      {
        id: "msg-tool-1",
        role: "assistant",
        parts: [
          {
            type: "tool-getLocation",
            toolCallId: "tc-1",
            state: "input-available",
            input: { query: "current" }
          }
        ]
      }
    ];

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: () => Promise.resolve(initialMessages),
        onToolCall: ({ toolCall, addToolOutput }) => {
          toolCallReceived(toolCall);
          addToolOutput({
            toolCallId: toolCall.toolCallId,
            output: { lat: 40.7, lng: -74.0 }
          });
        }
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(50);
    });

    // onToolCall should have been called with the tool call details
    expect(toolCallReceived).toHaveBeenCalledWith(
      expect.objectContaining({
        toolCallId: "tc-1",
        toolName: "getLocation",
        input: { query: "current" }
      })
    );
  });
});

describe("useAgentChat re-render stability", () => {
  it("should not cause infinite re-renders when idle", async () => {
    const agent = createAgent({
      name: "rerender-idle",
      url: "ws://localhost:3000/agents/chat/rerender-idle?_pk=abc"
    });

    let renderCount = 0;

    const TestComponent = () => {
      renderCount++;
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: []
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
    });

    // Capture render count after initial mount
    const afterMountCount = renderCount;

    // Wait to see if more renders happen (would indicate an infinite loop)
    await sleep(200);

    // In Strict Mode, React double-renders. After mount stabilizes,
    // there should be NO additional renders (no infinite loop).
    expect(renderCount).toBe(afterMountCount);
  });

  it("should not re-render excessively when messages are set", async () => {
    const agent = createAgent({
      name: "rerender-messages",
      url: "ws://localhost:3000/agents/chat/rerender-messages?_pk=abc"
    });

    let renderCount = 0;
    let chatInstance: ReturnType<typeof useAgentChat> | null = null;

    const TestComponent = () => {
      renderCount++;
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: [] as UIMessage[]
      });
      chatInstance = chat;
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
    });

    const beforeSetMessages = renderCount;

    // Set messages
    await act(async () => {
      chatInstance!.setMessages([
        {
          id: "msg-1",
          role: "user",
          parts: [{ type: "text", text: "Hello" }]
        }
      ]);
      await sleep(10);
    });

    const afterSetMessages = renderCount;

    // Wait to see if renders stabilize
    await sleep(200);

    // Should have re-rendered for the setMessages call but then stopped.
    // Allow some re-renders (React batching, state updates) but not infinite.
    const rendersFromSetMessages = afterSetMessages - beforeSetMessages;
    expect(rendersFromSetMessages).toBeLessThan(10);

    // No additional renders after stabilizing
    expect(renderCount).toBe(afterSetMessages);
  });

  it("should stabilize after receiving a broadcast message", async () => {
    const target = new EventTarget();
    const agent = createAgent({
      name: "rerender-broadcast",
      url: "ws://localhost:3000/agents/chat/rerender-broadcast?_pk=abc"
    });
    // Override addEventListener/removeEventListener to use our target
    (agent as unknown as Record<string, unknown>).addEventListener =
      target.addEventListener.bind(target);
    (agent as unknown as Record<string, unknown>).removeEventListener =
      target.removeEventListener.bind(target);

    let renderCount = 0;

    const TestComponent = () => {
      renderCount++;
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: []
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
    });

    const beforeBroadcast = renderCount;

    // Simulate a server broadcast (CF_AGENT_CHAT_MESSAGES)
    await act(async () => {
      target.dispatchEvent(
        new MessageEvent("message", {
          data: JSON.stringify({
            type: "cf_agent_chat_messages",
            messages: [
              {
                id: "broadcast-1",
                role: "user",
                parts: [{ type: "text", text: "From other tab" }]
              }
            ]
          })
        })
      );
      await sleep(10);
    });

    const afterBroadcast = renderCount;

    // Wait for stabilization
    await sleep(200);

    // Should have re-rendered for the broadcast but then stopped
    const rendersFromBroadcast = afterBroadcast - beforeBroadcast;
    expect(rendersFromBroadcast).toBeGreaterThan(0); // Must have re-rendered
    expect(rendersFromBroadcast).toBeLessThan(10); // But not infinitely

    // No additional renders after stabilizing
    expect(renderCount).toBe(afterBroadcast);
  });
});

describe("useAgentChat body option", () => {
  it("should include static body fields in sent messages", async () => {
    const sentMessages: string[] = [];
    const agent = createAgent({
      name: "body-static-test",
      url: "ws://localhost:3000/agents/chat/body-static-test?_pk=abc",
      send: (data: string) => sentMessages.push(data)
    });

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: [],
        body: { timezone: "America/New_York", userId: "user-123" }
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
    });

    // The body fields should be included when the transport sends messages
    // We can verify by checking that the component rendered without errors
    // (the actual body merging is tested via the sent WS messages)
    expect(sentMessages).toBeDefined();
  });

  it("should include dynamic body fields from function", async () => {
    const sentMessages: string[] = [];
    let callCount = 0;
    const agent = createAgent({
      name: "body-dynamic-test",
      url: "ws://localhost:3000/agents/chat/body-dynamic-test?_pk=abc",
      send: (data: string) => sentMessages.push(data)
    });

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: [],
        body: () => {
          callCount++;
          return { timestamp: Date.now(), requestNumber: callCount };
        }
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
    });

    // Component should render without errors with function body
    expect(callCount).toBeDefined();
  });

  it("should work alongside prepareSendMessagesRequest", async () => {
    const sentMessages: string[] = [];
    const agent = createAgent({
      name: "body-combined-test",
      url: "ws://localhost:3000/agents/chat/body-combined-test?_pk=abc",
      send: (data: string) => sentMessages.push(data)
    });

    const prepareSendMessagesRequest = vi.fn(() => ({
      body: { fromPrepare: true }
    }));

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: [],
        body: { fromBody: true },
        prepareSendMessagesRequest
      });
      return <div data-testid="messages-count">{chat.messages.length}</div>;
    };

    await act(async () => {
      render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
    });

    // Both body and prepareSendMessagesRequest should coexist without errors
    expect(sentMessages).toBeDefined();
  });
});

describe("useAgentChat stale agent ref (issue #929)", () => {
  it("should use the new agent's send method after agent switch, not the old one", async () => {
    const oldSend = vi.fn();
    const newSend = vi.fn();

    const agentOld = createAgent({
      name: "thread-old",
      url: "ws://localhost:3000/agents/chat/thread-old?_pk=old",
      send: oldSend
    });

    const agentNew = createAgent({
      name: "thread-new",
      url: "ws://localhost:3000/agents/chat/thread-new?_pk=new",
      send: newSend
    });

    let chatInstance: ReturnType<typeof useAgentChat> | null = null;

    const TestComponent = ({
      agent
    }: {
      agent: ReturnType<typeof useAgent>;
    }) => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: [] as UIMessage[]
      });
      chatInstance = chat;
      return <div data-testid="status">{chat.status}</div>;
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent agent={agentOld} />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    // Switch to the new agent
    await act(async () => {
      screen.rerender(<TestComponent agent={agentNew} />);
      await sleep(10);
    });

    // Clear any sends that happened during setup (e.g., stream resume requests)
    oldSend.mockClear();
    newSend.mockClear();

    // Clear history triggers agent.send() — this should go to the NEW agent
    await act(async () => {
      chatInstance!.clearHistory();
      await sleep(10);
    });

    // The clear message should have been sent to the NEW agent, not the old one
    const newSendCalls = newSend.mock.calls
      .map((args) => JSON.parse(args[0] as string))
      .filter((m: Record<string, unknown>) => m.type === "cf_agent_chat_clear");
    expect(newSendCalls.length).toBe(1);

    // The old agent should NOT have received the clear message
    const oldSendCalls = oldSend.mock.calls
      .map((args) => JSON.parse(args[0] as string))
      .filter((m: Record<string, unknown>) => m.type === "cf_agent_chat_clear");
    expect(oldSendCalls.length).toBe(0);
  });
});

describe("useAgentChat stream resumption (issue #896)", () => {
  function createAgentWithTarget({ name, url }: { name: string; url: string }) {
    const target = new EventTarget();
    const sentMessages: string[] = [];
    const agent = createAgent({
      name,
      url,
      send: (data: string) => sentMessages.push(data)
    });
    // Wire up the target so we can dispatch messages to the hook
    (agent as unknown as Record<string, unknown>).addEventListener =
      target.addEventListener.bind(target);
    (agent as unknown as Record<string, unknown>).removeEventListener =
      target.removeEventListener.bind(target);
    return { agent, target, sentMessages };
  }

  function dispatch(target: EventTarget, data: Record<string, unknown>) {
    target.dispatchEvent(
      new MessageEvent("message", { data: JSON.stringify(data) })
    );
  }

  it("should process resumed stream chunks progressively and update status", async () => {
    const { agent, target } = createAgentWithTarget({
      name: "replay-complete-test",
      url: "ws://localhost:3000/agents/chat/replay-complete-test?_pk=abc"
    });

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: [] as UIMessage[]
      });
      const assistantMsg = chat.messages.find(
        (m: UIMessage) => m.role === "assistant"
      );
      const textPart = assistantMsg?.parts.find(
        (p: UIMessage["parts"][number]) => p.type === "text"
      ) as { text?: string } | undefined;
      return (
        <div>
          <div data-testid="count">{chat.messages.length}</div>
          <div data-testid="text">{textPart?.text ?? ""}</div>
          <div data-testid="status">{chat.status}</div>
        </div>
      );
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    // Initially no messages
    await expect.element(screen.getByTestId("count")).toHaveTextContent("0");

    // Simulate server sending CF_AGENT_STREAM_RESUMING
    // The transport's reconnectToStream picks this up and returns a ReadableStream
    await act(async () => {
      dispatch(target, {
        type: "cf_agent_stream_resuming",
        id: "req-1"
      });
      await sleep(10);
    });

    // Simulate replay chunks — now processed progressively by useChat's pipeline
    await act(async () => {
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-1",
        body: '{"type":"text-start","id":"t1"}',
        done: false,
        replay: true
      });
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-1",
        body: '{"type":"text-delta","id":"t1","delta":"Hello world"}',
        done: false,
        replay: true
      });
      await sleep(10);
    });

    // Chunks are processed progressively by useChat — message appears immediately
    await expect.element(screen.getByTestId("count")).toHaveTextContent("1");
    await expect
      .element(screen.getByTestId("text"))
      .toHaveTextContent("Hello world");
  });

  it("should flush and finalize after done:true for orphaned streams", async () => {
    const { agent, target } = createAgentWithTarget({
      name: "orphaned-done-test",
      url: "ws://localhost:3000/agents/chat/orphaned-done-test?_pk=abc"
    });

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: [] as UIMessage[]
      });
      const assistantMsg = chat.messages.find(
        (m: UIMessage) => m.role === "assistant"
      );
      const textPart = assistantMsg?.parts.find(
        (p: UIMessage["parts"][number]) => p.type === "text"
      ) as { text?: string } | undefined;
      return (
        <div>
          <div data-testid="count">{chat.messages.length}</div>
          <div data-testid="text">{textPart?.text ?? ""}</div>
          <div data-testid="status">{chat.status}</div>
        </div>
      );
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    // Simulate resume + replay + done (orphaned stream path)
    await act(async () => {
      dispatch(target, {
        type: "cf_agent_stream_resuming",
        id: "req-orphaned"
      });
      await sleep(5);

      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-orphaned",
        body: '{"type":"text-start","id":"t1"}',
        done: false,
        replay: true
      });
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-orphaned",
        body: '{"type":"text-delta","id":"t1","delta":"partial from hibernation"}',
        done: false,
        replay: true
      });

      // done:true signals orphaned stream is finalized
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-orphaned",
        body: "",
        done: true,
        replay: true
      });
      await sleep(10);
    });

    // Message should be flushed with the accumulated text
    await expect.element(screen.getByTestId("count")).toHaveTextContent("1");
    await expect
      .element(screen.getByTestId("text"))
      .toHaveTextContent("partial from hibernation");
  });

  it("should continue receiving live chunks after replayComplete", async () => {
    const { agent, target } = createAgentWithTarget({
      name: "replay-then-live-test",
      url: "ws://localhost:3000/agents/chat/replay-then-live-test?_pk=abc"
    });

    const TestComponent = () => {
      const chat = useAgentChat({
        agent,
        getInitialMessages: null,
        messages: [] as UIMessage[]
      });
      const assistantMsg = chat.messages.find(
        (m: UIMessage) => m.role === "assistant"
      );
      const textPart = assistantMsg?.parts.find(
        (p: UIMessage["parts"][number]) => p.type === "text"
      ) as { text?: string } | undefined;
      return (
        <div>
          <div data-testid="count">{chat.messages.length}</div>
          <div data-testid="text">{textPart?.text ?? ""}</div>
        </div>
      );
    };

    const screen = await act(async () => {
      const screen = render(<TestComponent />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    // Replay phase
    await act(async () => {
      dispatch(target, {
        type: "cf_agent_stream_resuming",
        id: "req-live"
      });
      await sleep(5);

      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-live",
        body: '{"type":"text-start","id":"t1"}',
        done: false,
        replay: true
      });
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-live",
        body: '{"type":"text-delta","id":"t1","delta":"replayed-"}',
        done: false,
        replay: true
      });
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-live",
        body: "",
        done: false,
        replay: true,
        replayComplete: true
      });
      await sleep(10);
    });

    // After replay, message should show replayed text
    await expect.element(screen.getByTestId("count")).toHaveTextContent("1");
    await expect
      .element(screen.getByTestId("text"))
      .toHaveTextContent("replayed-");

    // Now simulate a live chunk arriving (no replay flag)
    await act(async () => {
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-live",
        body: '{"type":"text-delta","id":"t1","delta":"and live!"}',
        done: false
      });
      await sleep(10);
    });

    // The live chunk should append to the same message
    await expect.element(screen.getByTestId("count")).toHaveTextContent("1");
    await expect
      .element(screen.getByTestId("text"))
      .toHaveTextContent("replayed-and live!");
  });
});

describe("useAgentChat tool approval continuations (issue #1108)", () => {
  function createAgentWithTarget({ name, url }: { name: string; url: string }) {
    const target = new EventTarget();
    const agent = createAgent({ name, url });
    (agent as unknown as Record<string, unknown>).addEventListener =
      target.addEventListener.bind(target);
    (agent as unknown as Record<string, unknown>).removeEventListener =
      target.removeEventListener.bind(target);
    return { agent, target };
  }

  function dispatch(target: EventTarget, data: Record<string, unknown>) {
    target.dispatchEvent(
      new MessageEvent("message", { data: JSON.stringify(data) })
    );
  }

  const initialMessages: UIMessage[] = [
    {
      id: "assistant-local",
      role: "assistant",
      parts: [
        {
          type: "tool-dangerousAction",
          toolCallId: "tc-approval-1",
          state: "approval-responded",
          input: { action: "delete" },
          approval: { id: "approval-req-1", approved: true }
        }
      ]
    }
  ];

  function TestComponent({ agent }: { agent: ReturnType<typeof useAgent> }) {
    const chat = useAgentChat({
      agent,
      getInitialMessages: () => Promise.resolve(initialMessages)
    });
    const assistantMessages = chat.messages.filter(
      (message) => message.role === "assistant"
    );
    const textPart = assistantMessages
      .flatMap((message) => message.parts)
      .find((part) => part.type === "text") as { text?: string } | undefined;
    const toolPartsCount = assistantMessages.reduce((count, message) => {
      return (
        count +
        message.parts.filter(
          (part) => "toolCallId" in part && part.toolCallId === "tc-approval-1"
        ).length
      );
    }, 0);

    return (
      <div>
        <div data-testid="assistant-count">{assistantMessages.length}</div>
        <div data-testid="assistant-ids">
          {assistantMessages.map((message) => message.id).join(",")}
        </div>
        <div data-testid="tool-parts-count">{toolPartsCount}</div>
        <div data-testid="text">{textPart?.text ?? ""}</div>
      </div>
    );
  }

  it("keeps the existing assistant id for continuation start chunks", async () => {
    const { agent, target } = createAgentWithTarget({
      name: "continuation-start-id",
      url: "ws://localhost:3000/agents/chat/continuation-start-id?_pk=abc"
    });

    const screen = await act(async () => {
      const screen = render(<TestComponent agent={agent} />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    await expect
      .element(screen.getByTestId("assistant-count"))
      .toHaveTextContent("1");

    await act(async () => {
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-continuation-start",
        continuation: true,
        body: '{"type":"start","messageId":"assistant-stream"}',
        done: false
      });
      await sleep(10);
    });

    await expect
      .element(screen.getByTestId("assistant-count"))
      .toHaveTextContent("1");
    await expect
      .element(screen.getByTestId("assistant-ids"))
      .toHaveTextContent("assistant-local");
    await expect
      .element(screen.getByTestId("tool-parts-count"))
      .toHaveTextContent("1");
  });

  it("keeps merging continuations when broadcasts replace assistant ids mid-stream", async () => {
    const { agent, target } = createAgentWithTarget({
      name: "continuation-remap-id",
      url: "ws://localhost:3000/agents/chat/continuation-remap-id?_pk=abc"
    });

    const screen = await act(async () => {
      const screen = render(<TestComponent agent={agent} />, {
        wrapper: ({ children }) => (
          <StrictMode>
            <Suspense fallback="Loading...">{children}</Suspense>
          </StrictMode>
        )
      });
      await sleep(10);
      return screen;
    });

    await act(async () => {
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-continuation-remap",
        continuation: true,
        body: '{"type":"start","messageId":"assistant-stream"}',
        done: false
      });
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-continuation-remap",
        continuation: true,
        body: '{"type":"text-start","id":"text-1"}',
        done: false
      });
      await sleep(10);
    });

    await act(async () => {
      dispatch(target, {
        type: "cf_agent_chat_messages",
        messages: [
          {
            id: "assistant-server",
            role: "assistant",
            parts: [
              {
                type: "tool-dangerousAction",
                toolCallId: "tc-approval-1",
                state: "approval-responded",
                input: { action: "delete" },
                approval: { id: "approval-req-1", approved: true }
              }
            ]
          }
        ]
      });
      dispatch(target, {
        type: "cf_agent_use_chat_response",
        id: "req-continuation-remap",
        continuation: true,
        body: '{"type":"text-delta","id":"text-1","delta":"done"}',
        done: false
      });
      await sleep(10);
    });

    await expect
      .element(screen.getByTestId("assistant-count"))
      .toHaveTextContent("1");
    await expect
      .element(screen.getByTestId("assistant-ids"))
      .toHaveTextContent("assistant-server");
    await expect
      .element(screen.getByTestId("tool-parts-count"))
      .toHaveTextContent("1");
    await expect.element(screen.getByTestId("text")).toHaveTextContent("done");
  });
});