@copilotkit/vue 2.0.0-rc.3 → 2.0.0-rc.30

package/src/v2/components/chat/CopilotSidebarViewInternal.vue

@@ -1,20 +1,29 @@
 <script setup lang="ts">
 import {
   computed,
-  nextTick,
   onBeforeUnmount,
   onMounted,
   ref,
   useAttrs,
   watch,
 } from "vue";
+import {
+  ɵSIDEBAR_DEFAULT_WIDTH,
+  ɵSIDEBAR_TRANSITION_MS,
+  ɵcreateDockedSidebar,
+  ɵdimensionToCss,
+  ɵwatchElementWidth,
+} from "@copilotkit/core";
+import type { ɵDockedSidebar } from "@copilotkit/core";
 import { CopilotChatDefaultLabels } from "../../providers/types";
 import { useCopilotChatConfiguration } from "../../providers/useCopilotChatConfiguration";
 import CopilotChatToggleButton from "./CopilotChatToggleButton.vue";
 import CopilotChatView from "./CopilotChatView.vue";
 import CopilotModalHeader from "./CopilotModalHeader.vue";
+import { forwardedSlotNames } from "./forwarded-slots";
 import type {
   CopilotChatFeatherSlotProps,
+  CopilotChatInterruptSlotProps,
   CopilotChatMessageViewSlotProps,
   CopilotChatScrollToBottomButtonSlotProps,
   CopilotChatScrollViewSlotProps,
@@ -28,9 +37,6 @@ import type {
 
 defineOptions({ inheritAttrs: false });
 
-const DEFAULT_SIDEBAR_WIDTH = 480;
-const SIDEBAR_TRANSITION_MS = 260;
-
 const props = withDefaults(
   defineProps<Omit<CopilotSidebarViewProps, "defaultOpen">>(),
   {
@@ -59,14 +65,22 @@ defineSlots<{
   "scroll-to-bottom-button"?: (
     props: CopilotChatScrollToBottomButtonSlotProps,
   ) => unknown;
+  interrupt?: (props: CopilotChatInterruptSlotProps) => unknown;
   input?: (props: CopilotSidebarWelcomeScreenInputSlotProps) => unknown;
   "suggestion-view"?: (
     props: CopilotSidebarWelcomeScreenSuggestionViewSlotProps,
   ) => unknown;
   "welcome-screen"?: (props: CopilotChatWelcomeScreenSlotProps) => unknown;
   "welcome-message"?: () => unknown;
+  // Any other slot goes to the chat view, as on CopilotChat.
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  [key: string]: ((props: any) => unknown) | undefined;
 }>();
 
+// Slots this component renders itself. It passes every other slot to the
+// chat view.
+const OWN_SLOTS = ["header", "toggle-button"] as const;
+
 const emit = defineEmits<{
   "submit-message": [value: string];
   stop: [];
@@ -86,9 +100,8 @@ const config = useCopilotChatConfiguration();
 
 const sidebarRef = ref<HTMLElement | null>(null);
 const measuredSidebarWidth = ref<number | string>(
-  props.width ?? DEFAULT_SIDEBAR_WIDTH,
+  props.width ?? ɵSIDEBAR_DEFAULT_WIDTH,
 );
-let resizeObserver: ResizeObserver | null = null;
 
 const isSidebarOpen = computed(() => config.value?.isModalOpen ?? false);
 const resolvedSidebarWidth = computed(
@@ -112,43 +125,23 @@ const asideAttrs = computed(() => {
 const asideStyle = computed(
   () =>
     ({
-      "--sidebar-width": widthToCss(resolvedSidebarWidth.value),
+      "--sidebar-width": widthCss.value,
       paddingTop: "env(safe-area-inset-top)",
       paddingBottom: "env(safe-area-inset-bottom)",
     }) as Record<string, string>,
 );
-// The page makes room on the docked edge (desktop only).
-const bodyMarginStyle = computed(() => {
-  const margin =
-    props.position === "left" ? "margin-inline-start" : "margin-inline-end";
-  return `
-@media (min-width: 768px) {
-  body {
-    ${margin}: ${widthToMargin(resolvedSidebarWidth.value)};
-    transition: ${margin} ${SIDEBAR_TRANSITION_MS}ms ease;
-  }
-}
-`;
-});
-
-function widthToCss(width: number | string): string {
-  return typeof width === "number" ? `${width}px` : width;
-}
-
-function widthToMargin(width: number | string): string {
-  return typeof width === "number" ? `${width}px` : width;
-}
-
-function updateMeasuredWidth() {
-  if (props.width !== undefined || !sidebarRef.value) {
-    return;
-  }
+const widthCss = computed(() =>
+  ɵdimensionToCss(resolvedSidebarWidth.value, ɵSIDEBAR_DEFAULT_WIDTH),
+);
 
-  const nextWidth = sidebarRef.value.getBoundingClientRect().width;
-  if (nextWidth > 0) {
-    measuredSidebarWidth.value = nextWidth;
-  }
-}
+// The page makes room on the docked edge (desktop only).
+let dock: ɵDockedSidebar | undefined;
+watch(
+  [isSidebarOpen, () => props.position, widthCss],
+  ([open, position, width]) => {
+    dock?.set({ open, position, width });
+  },
+);
 
 function closeSidebar() {
   config.value?.setModalOpen?.(false);
@@ -223,46 +216,39 @@ const chatViewEventProps = computed(() => {
   return listeners;
 });
 
+// Without a `width` prop, the page margin follows the width the CSS gives
+// the sidebar.
 watch(
-  () => props.width,
-  (nextWidth) => {
-    if (nextWidth !== undefined) {
-      measuredSidebarWidth.value = nextWidth;
-      return;
-    }
-    updateMeasuredWidth();
+  [sidebarRef, () => props.width],
+  ([element, width], _old, onCleanup) => {
+    if (width !== undefined || !element) return;
+    onCleanup(
+      ɵwatchElementWidth(element, (next) => {
+        measuredSidebarWidth.value = next;
+      }),
+    );
   },
+  { flush: "post" },
 );
 
-onMounted(async () => {
-  await nextTick();
-  updateMeasuredWidth();
-
-  if (
-    props.width !== undefined ||
-    !sidebarRef.value ||
-    typeof ResizeObserver === "undefined"
-  ) {
-    return;
-  }
-
-  resizeObserver = new ResizeObserver(() => {
-    updateMeasuredWidth();
+onMounted(() => {
+  dock = ɵcreateDockedSidebar(document, {
+    transitionMs: ɵSIDEBAR_TRANSITION_MS,
+  });
+  dock.set({
+    open: isSidebarOpen.value,
+    position: props.position,
+    width: widthCss.value,
   });
-  resizeObserver.observe(sidebarRef.value);
 });
 
 onBeforeUnmount(() => {
-  resizeObserver?.disconnect();
-  resizeObserver = null;
+  dock?.dispose();
+  dock = undefined;
 });
 </script>
 
 <template>
-  <component :is="'style'" v-if="isSidebarOpen">
-    {{ bodyMarginStyle }}
-  </component>
-
   <slot
     name="toggle-button"
     :is-open="isSidebarOpen"
@@ -317,42 +303,12 @@ onBeforeUnmount(() => {
           :on-finish-transcribe-with-audio="onFinishTranscribeWithAudio"
           v-bind="chatViewEventProps"
         >
-          <template v-if="$slots['message-view']" #message-view="slotProps">
-            <slot name="message-view" v-bind="slotProps" />
-          </template>
-
-          <template v-if="$slots['scroll-view']" #scroll-view="slotProps">
-            <slot name="scroll-view" v-bind="slotProps" />
-          </template>
-
-          <template v-if="$slots.feather" #feather>
-            <slot name="feather" />
-          </template>
-
           <template
-            v-if="$slots['scroll-to-bottom-button']"
-            #scroll-to-bottom-button="slotProps"
+            v-for="slotName in forwardedSlotNames($slots, OWN_SLOTS)"
+            :key="slotName"
+            #[slotName]="slotProps"
           >
-            <slot name="scroll-to-bottom-button" v-bind="slotProps" />
-          </template>
-
-          <template v-if="$slots.input" #input="slotProps">
-            <slot name="input" v-bind="slotProps" />
-          </template>
-
-          <template
-            v-if="$slots['suggestion-view']"
-            #suggestion-view="slotProps"
-          >
-            <slot name="suggestion-view" v-bind="slotProps" />
-          </template>
-
-          <template v-if="$slots['welcome-screen']" #welcome-screen="slotProps">
-            <slot name="welcome-screen" v-bind="slotProps" />
-          </template>
-
-          <template v-if="$slots['welcome-message']" #welcome-message>
-            <slot name="welcome-message" />
+            <slot :name="slotName" v-bind="slotProps" />
           </template>
         </CopilotChatView>
       </div>

package/src/v2/components/chat/CopilotThreadsDrawer.vue

@@ -8,7 +8,6 @@ import { computed, onMounted, onScopeDispose, ref, watchEffect } from "vue";
 // below is kept at module scope.
 import type {
   CopilotKitThreadsDrawer as CopilotKitThreadsDrawerElement,
-  DrawerThread,
   ThreadSelectedDetail,
   ArchiveDetail,
   UnarchiveDetail,
@@ -22,7 +21,13 @@ import type {
 // lands and is published; declared locally here because the built element types
 // in this worktree predate it.
 type CollapseChangeDetail = { collapsed: boolean };
-import { DEFAULT_AGENT_ID } from "@copilotkit/shared";
+import {
+  ɵcreateThreadsDrawerActions,
+  ɵgetThreadsDrawerElementProps,
+  ɵresolveAgentId,
+  ɵresolveThreadsLicense,
+  ɵtoDrawerThread,
+} from "@copilotkit/core";
 import { useThreads } from "../../hooks/use-threads";
 import type { Thread } from "../../hooks/use-threads";
 import { useCopilotChatConfiguration } from "../../providers/useCopilotChatConfiguration";
@@ -71,18 +76,17 @@ const emit = defineEmits<{
 const config = useCopilotChatConfiguration();
 const license = useLicenseContext();
 
-const licensed = computed(() => {
-  const status = license.value.status;
-  // `unknown` is an unrecovered retryable entitlement lookup, not a settled
-  // negative: fetch the list and let the threads endpoint decide.
-  if (status === "unknown") return true;
-  const present = status === "valid" || status === "expiring";
-  return present && license.value.checkFeature("threads");
-});
-const licensePending = computed(() => license.value.status === null);
+// `unknown` (an unrecovered entitlement lookup) counts as licensed so the
+// threads endpoint decides; `null` (not reported yet) is pending.
+const threadsLicense = computed(() => ɵresolveThreadsLicense(license.value));
+const licensed = computed(() => threadsLicense.value.licensed);
+const licensePending = computed(() => threadsLicense.value.licensePending);
 
-const resolvedAgentId = computed(
-  () => props.agentId ?? config.value?.agentId ?? DEFAULT_AGENT_ID,
+const resolvedAgentId = computed(() =>
+  ɵresolveAgentId({
+    agentId: props.agentId,
+    chatAgentId: config.value?.agentId,
+  }),
 );
 const activeThreadId = computed(() => config.value?.threadId ?? null);
 
@@ -107,18 +111,8 @@ const threadsApi = useThreads({
   limit: () => props.limit,
 });
 
-function toDrawerThread(t: Thread): DrawerThread {
-  return {
-    id: t.id,
-    name: t.name,
-    archived: t.archived,
-    createdAt: t.createdAt,
-    updatedAt: t.updatedAt,
-    ...(t.lastRunAt !== undefined ? { lastRunAt: t.lastRunAt } : {}),
-  };
-}
 const drawerThreads = computed(() =>
-  threadsApi.threads.value.map(toDrawerThread),
+  threadsApi.threads.value.map(ɵtoDrawerThread),
 );
 
 // Client-only element registration (SSR-safe). The element module is imported
@@ -144,10 +138,15 @@ watchEffect(
     const el = elRef.value;
     if (!el) return;
     el.threads = drawerThreads.value;
-    el.loading = threadsApi.isLoading.value || licensePending.value;
+    const elementProps = ɵgetThreadsDrawerElementProps({
+      loading: threadsApi.isLoading.value,
+      licensed: licensed.value,
+      licensePending: licensePending.value,
+    });
+    el.loading = elementProps.loading;
     el.error = threadsApi.listError.value?.message ?? null;
     el.activeThreadId = activeThreadId.value;
-    el.licensed = licensed.value || licensePending.value;
+    el.licensed = elementProps.licensed;
     el.hasMore = threadsApi.hasMoreThreads.value;
     el.fetchingMore = threadsApi.isFetchingMoreThreads.value;
     // Dedicated fetch-more error channel: drives the element's inline
@@ -177,101 +176,40 @@ const unregisterDrawer = config.value?.registerDrawer?.();
 if (unregisterDrawer) onScopeDispose(unregisterDrawer);
 
 // --- Outbound event handlers ------------------------------------------------
-/** The chat input textarea's documented `data-testid`. */
-const CHAT_INPUT_TESTID = "copilot-chat-input-textarea";
-/** The chat view container's documented `data-testid`. */
-const CHAT_CONTAINER_TESTID = "copilot-chat-view";
-
-/**
- * Returns the chat input element for focus-return after a thread is selected.
- *
- * Best-effort and SCOPED: walks up from the drawer's custom element looking
- * for an ancestor that contains a chat-view container
- * (`data-testid="copilot-chat-view"`), then returns the chat input within
- * that subtree. This avoids focusing the wrong composer on a page hosting
- * more than one chat (multi-chat dashboards), where a document-global lookup
- * would grab whichever input appears first in DOM order rather than the one
- * this drawer drives.
- *
- * Falls back to a document-global lookup when no scoping ancestor is found
- * (e.g. the drawer and chat share no common container, or headless usage).
- */
-function findChatInput(origin: Element | null): HTMLElement | null {
-  if (typeof document === "undefined") return null;
-  const container = origin?.closest?.(
-    `[data-testid="${CHAT_CONTAINER_TESTID}"]`,
-  );
-  if (container) {
-    const scoped = container.querySelector<HTMLElement>(
-      `[data-testid="${CHAT_INPUT_TESTID}"]`,
-    );
-    if (scoped) return scoped;
-  }
-  return document.querySelector<HTMLElement>(
-    `[data-testid="${CHAT_INPUT_TESTID}"]`,
-  );
-}
-
-function focusChatInput() {
-  findChatInput(elRef.value)?.focus();
-}
+// Core owns the drawer actions; they read the latest values at call time.
+const actions = ɵcreateThreadsDrawerActions({
+  getThreads: () => threadsApi,
+  getConfig: () => config.value,
+  getCallbacks: () => ({
+    onThreadSelect: props.onThreadSelect,
+    onNewThread: props.onNewThread,
+  }),
+  getElement: () => elRef.value,
+});
 
 function onThreadSelected(event: Event) {
   const { threadId } = (event as CustomEvent<ThreadSelectedDetail>).detail;
-  if (props.onThreadSelect) props.onThreadSelect(threadId);
-  else config.value?.setActiveThreadId?.(threadId, { explicit: true });
-  focusChatInput();
-}
-function handleNewThread() {
-  threadsApi.startNewThread();
-  if (props.onNewThread) props.onNewThread();
-  else config.value?.startNewThread?.();
+  actions.selectThread(threadId);
 }
+const handleNewThread = actions.newThread;
 function onArchive(event: Event) {
-  const { threadId } = (event as CustomEvent<ArchiveDetail>).detail;
-  void threadsApi
-    .archiveThread(threadId)
-    .catch((e) =>
-      console.error("CopilotThreadsDrawer: archiveThread failed", e),
-    );
+  actions.archive((event as CustomEvent<ArchiveDetail>).detail.threadId);
 }
 function onUnarchive(event: Event) {
-  const { threadId } = (event as CustomEvent<UnarchiveDetail>).detail;
-  void threadsApi
-    .unarchiveThread(threadId)
-    .catch((e) =>
-      console.error("CopilotThreadsDrawer: unarchiveThread failed", e),
-    );
+  actions.unarchive((event as CustomEvent<UnarchiveDetail>).detail.threadId);
 }
 function onDelete(event: Event) {
-  const { threadId } = (event as CustomEvent<DeleteDetail>).detail;
-  const wasActive = threadId === activeThreadId.value;
-  void threadsApi
-    .deleteThread(threadId)
-    .then(() => {
-      if (wasActive) {
-        handleNewThread();
-      }
-    })
-    .catch((e) =>
-      console.error("CopilotThreadsDrawer: deleteThread failed", e),
-    );
-}
-function onFilterChange() {
-  threadsApi.refetchThreads();
+  actions.remove((event as CustomEvent<DeleteDetail>).detail.threadId);
 }
+const onFilterChange = actions.refetch;
 function onCollapseChange(event: Event) {
   const { collapsed } = (event as CustomEvent<CollapseChangeDetail>).detail;
   emit("collapse-change", collapsed);
 }
 function onRetry(event: Event) {
-  const { scope } = (event as CustomEvent<RetryDetail>).detail;
-  if (scope === "fetch-more") threadsApi.fetchMoreThreads();
-  else threadsApi.refetchThreads();
-}
-function onLoadMore() {
-  threadsApi.fetchMoreThreads();
+  actions.retry((event as CustomEvent<RetryDetail>).detail.scope);
 }
+const onLoadMore = actions.loadMore;
 function onOpenChange(event: Event) {
   const { open } = (event as CustomEvent<OpenChangeDetail>).detail;
   setDrawerOpen(open);

package/src/v2/components/chat/StreamingMarkdownDefault.vue

@@ -1,43 +1,26 @@
 <script setup lang="ts">
-import { computed, h, watchEffect } from "vue";
-import { StreamingMarkdownRenderer } from "@copilotkit/markdown-renderer/vue";
-import type { VueStreamingMarkdownNodeRenderer } from "@copilotkit/markdown-renderer/vue";
-import { warnUnsupportedRichSyntaxOnce } from "@copilotkit/markdown-renderer";
-import type { StreamingMarkdownParserOptions } from "@copilotkit/markdown-renderer";
+import { computed, watchEffect } from "vue";
+import { StreamingMarkdownRenderer } from "@copilotkit/markdown/vue";
+import {
+  CHAT_MARKDOWN_NODE_CLASS_NAMES,
+  warnUnsupportedRichSyntaxOnce,
+} from "@copilotkit/markdown";
+import type { StreamingMarkdownParserOptions } from "@copilotkit/markdown";
 import type { DefaultMarkdownRendererProps } from "../../providers/markdown-renderer";
 
 const props = withDefaults(
   defineProps<
     { content: string; isStreaming?: boolean } & DefaultMarkdownRendererProps
   >(),
-  { isStreaming: false },
+  { isStreaming: false, caret: undefined },
 );
 
+// vue-tsc reads the renderer's `options` prop as the full options type; the
+// renderer itself accepts a partial object.
 const resolvedOptions = computed(
   () => props.options as StreamingMarkdownParserOptions | undefined,
 );
 
-const defaultCodeBlockRenderer: VueStreamingMarkdownNodeRenderer = (
-  node,
-  defaultVNode,
-) => {
-  if (node.type !== "code-block") return defaultVNode;
-  return h(
-    "pre",
-    {
-      "data-streaming-markdown-node": "code-block",
-      "data-node-open": String(!node.closed),
-      class: "cpk-markdown-code-block",
-    },
-    [h("code", { "data-code-info": node.info ?? undefined }, node.text)],
-  );
-};
-
-const mergedNodeRenderers = computed(() => ({
-  codeBlock: defaultCodeBlockRenderer,
-  ...props.nodeRenderers,
-}));
-
 // Dev-only: nudge upgraders from the bundled Streamdown default when their
 // content needs math/syntax highlighting the built-in renderer doesn't do.
 watchEffect(() => warnUnsupportedRichSyntaxOnce(props.content));
@@ -45,11 +28,15 @@ watchEffect(() => warnUnsupportedRichSyntaxOnce(props.content));
 
 <template>
   <StreamingMarkdownRenderer
+    v-if="props.content"
     :content="props.content"
     :is-complete="!props.isStreaming"
-    :node-renderers="mergedNodeRenderers"
+    :node-renderers="props.nodeRenderers"
+    :node-class-names="CHAT_MARKDOWN_NODE_CLASS_NAMES"
     :caret="props.caret ?? props.isStreaming"
     :options="resolvedOptions"
     :class="props.class"
+    :on-link-click="props.onLinkClick"
+    :on-citation-click="props.onCitationClick"
   />
 </template>

package/src/v2/components/chat/__tests__/CopilotChat.labels.test.ts

@@ -1,7 +1,7 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { defineComponent, nextTick, onMounted } from "vue";
 import { fireEvent, render, screen } from "@testing-library/vue";
-import type { Message, ReasoningMessage } from "@ag-ui/core";
+import type { ReasoningMessage } from "@ag-ui/core";
 import type { Attachment } from "@copilotkit/shared";
 import CopilotKitProvider from "../../../providers/CopilotKitProvider.vue";
 import CopilotChatConfigurationProvider from "../../../providers/CopilotChatConfigurationProvider.vue";
@@ -16,6 +16,8 @@ import { CopilotModalHeader, CopilotSidebarView } from "../index";
 // hard-coded English.
 const germanLabels: Partial<Record<keyof CopilotChatLabels, string>> = {
   reasoningMessageThinkingLabel: "Denkt nach…",
+  reasoningMessageThoughtForLabel: "Gedacht für {duration}",
+  reasoningMessageDurationFewSeconds: "ein paar Sekunden",
   chatInputSlashMenuLabel: "Befehle",
   chatInputSlashMenuEmptyText: "Keine Befehle gefunden",
   attachmentRemoveButtonLabel: "Anhang entfernen",
@@ -26,11 +28,11 @@ const germanLabels: Partial<Record<keyof CopilotChatLabels, string>> = {
   sidebarLabel: "Chat-Seitenleiste",
 };
 
-const reasoning = {
+const reasoning: ReasoningMessage = {
   id: "reasoning-1",
   role: "reasoning",
   content: "",
-} as ReasoningMessage;
+};
 
 const video: Attachment = {
   id: "clip",
@@ -66,7 +68,7 @@ function renderWithLabels(template: string) {
       return {
         labels: germanLabels,
         reasoning,
-        reasoningMessages: [reasoning] as Message[],
+        reasoningMessages: [reasoning],
         attachments: [video],
         toolsMenu: [{ label: "Say hi", action: vi.fn() }],
       };
@@ -112,6 +114,19 @@ describe("chat strings from labels", () => {
     expect(screen.queryByText("Thinking…")).toBeNull();
   });
 
+  it("uses the configured duration labels once the model is done", () => {
+    renderWithLabels(`
+      <CopilotChatReasoningMessage
+        :message="reasoning"
+        :messages="reasoningMessages"
+        :is-running="false"
+      />
+    `);
+
+    expect(screen.getByText("Gedacht für ein paar Sekunden")).toBeDefined();
+    expect(screen.queryByText(/Thought for/)).toBeNull();
+  });
+
   it("passes the elapsed reasoning time to a custom header", async () => {
     vi.useFakeTimers();
     renderWithLabels(`

package/src/v2/components/chat/__tests__/CopilotChat.test.ts

@@ -60,6 +60,11 @@ describe("CopilotChat", () => {
       }) => void | Promise<void>,
       errorMessage: string,
     ) => Promise<void>;
+    emitError: (params: {
+      error: Error;
+      code: CopilotKitCoreErrorCode;
+      context?: Record<string, unknown>;
+    }) => Promise<void>;
   }
 
   beforeEach(() => {
@@ -405,16 +410,13 @@ describe("CopilotChat", () => {
       },
     });
 
-    await (core as unknown as CopilotKitCoreTestAccess).notifySubscribers(
-      (subscriber) =>
-        subscriber.onError?.({
-          copilotkit: core!,
-          error: new Error("matching agent"),
-          code: CopilotKitCoreErrorCode.AGENT_RUN_FAILED,
-          context: { agentId: "default" },
-        }),
-      "test chat onError",
-    );
+    // Through core's error path: the provider's onError is core's handler,
+    // the chat's is a subscriber.
+    await (core as unknown as CopilotKitCoreTestAccess).emitError({
+      error: new Error("matching agent"),
+      code: CopilotKitCoreErrorCode.AGENT_RUN_FAILED,
+      context: { agentId: "default" },
+    });
 
     await nextTick();
 

package/src/v2/components/chat/__tests__/CopilotChatActivityRendering.e2e.test.ts

@@ -7,6 +7,7 @@ import {
 } from "@testing-library/vue";
 import { defineComponent, ref, toRaw } from "vue";
 import { afterEach, describe, expect, it, vi } from "vitest";
+import type * as CoreModule from "@copilotkit/core";
 import type { AbstractAgent } from "@ag-ui/client";
 import { EventType } from "@ag-ui/client";
 import type { A2UITheme as Theme } from "../../../types";
@@ -160,11 +161,25 @@ vi.mock("phoenix", () => ({
   Socket: MockPhoenixSocket,
 }));
 
-vi.mock("@jetbrains/websandbox", () => ({
-  default: {
-    create: (...args: unknown[]) => mockWebsandboxCreate(...args),
-  },
-}));
+// Run the real core controller against a fake websandbox
+vi.mock("@copilotkit/core", async (importOriginal) => {
+  const actual = await importOriginal<typeof CoreModule>();
+  class TestOpenGenerativeUIController
+    extends actual.ɵOpenGenerativeUIController
+  {
+    constructor() {
+      super({
+        loadWebsandbox: async () => ({
+          create: (...args) => mockWebsandboxCreate(...args),
+        }),
+      });
+    }
+  }
+  return {
+    ...actual,
+    ɵOpenGenerativeUIController: TestOpenGenerativeUIController,
+  };
+});
 
 function jsonResponse(body: unknown, status = 200): Response {
   return {
@@ -667,13 +682,15 @@ describe("CopilotChat activity message rendering", () => {
 
     renderOpenGenerativeUIChat();
 
+    // The restored chat can remount the activity while it reconnects, so wait
+    // until exactly one sandbox, the restored one, is alive.
     await waitFor(() => {
-      expect(mockWebsandboxCreate).toHaveBeenCalledTimes(2);
+      const created = mockWebsandboxCreate.mock.calls.length;
+      expect(created).toBeGreaterThanOrEqual(2);
+      expect(mockWebsandboxDestroy).toHaveBeenCalledTimes(created - 1);
     });
-    expect(mockWebsandboxCreate.mock.calls[1]?.[1]).toMatchObject({
+    expect(mockWebsandboxCreate.mock.calls.at(-1)?.[1]).toMatchObject({
       frameContent: restoredHtml,
     });
-
-    expect(mockWebsandboxDestroy).toHaveBeenCalledTimes(1);
   });
 });