🤖 fix: storybook mocks for feature flags/telemetry

Change-Id: I31ce25f735298b7cf2c6d3e997eded7173771899
Signed-off-by: Thomas Kosiewski <tk@coder.com>

Author: ThomasK33

.storybook/mocks/orpc.ts

@@ -6,7 +6,11 @@
 import type { APIClient } from "@/browser/contexts/API";
 import type { FrontendWorkspaceMetadata } from "@/common/types/workspace";
 import type { ProjectConfig } from "@/node/config";
-import type { WorkspaceChatMessage, ProvidersConfigMap } from "@/common/orpc/types";
+import type {
+  WorkspaceChatMessage,
+  ProvidersConfigMap,
+  WorkspaceStatsSnapshot,
+} from "@/common/orpc/types";
 import type { ChatStats } from "@/common/types/chatStats";
 import { DEFAULT_RUNTIME_CONFIG } from "@/common/constants/workspace";
 import { createAsyncMessageQueue } from "@/common/utils/asyncMessageQueue";
@@ -69,6 +73,8 @@ export interface MockORPCClientOptions {
     }>
   >;
   /** Session usage data per workspace (for Costs tab) */
+  workspaceStatsSnapshots?: Map<string, WorkspaceStatsSnapshot>;
+  statsTabVariant?: "control" | "stats";
   sessionUsage?: Map<string, MockSessionUsage>;
   /** MCP server configuration per project */
   mcpServers?: Map<
@@ -112,11 +118,27 @@ export function createMockORPCClient(options: MockORPCClientOptions = {}): APICl
     onProjectRemove,
     backgroundProcesses = new Map(),
     sessionUsage = new Map(),
+    workspaceStatsSnapshots = new Map<string, WorkspaceStatsSnapshot>(),
+    statsTabVariant = "control",
     mcpServers = new Map(),
     mcpOverrides = new Map(),
     mcpTestResults = new Map(),
   } = options;
 
+  // Feature flags
+  let statsTabOverride: "default" | "on" | "off" = "default";
+
+  const getStatsTabState = () => {
+    const enabled =
+      statsTabOverride === "on"
+        ? true
+        : statsTabOverride === "off"
+          ? false
+          : statsTabVariant === "stats";
+
+    return { enabled, variant: statsTabVariant, override: statsTabOverride } as const;
+  };
+
   const workspaceMap = new Map(workspaces.map((w) => [w.id, w]));
 
   const mockStats: ChatStats = {
@@ -136,8 +158,14 @@ export function createMockORPCClient(options: MockORPCClientOptions = {}): APICl
       calculateStats: async () => mockStats,
     },
     features: {
-      getStatsTabState: async () => ({ enabled: false, variant: "control", override: "default" }),
-      setStatsTabOverride: async () => ({ enabled: false, variant: "control", override: "default" }),
+      getStatsTabState: async () => getStatsTabState(),
+      setStatsTabOverride: async (input: { override: "default" | "on" | "off" }) => {
+        statsTabOverride = input.override;
+        return getStatsTabState();
+      },
+    },
+    telemetry: {
+      track: async () => undefined,
     },
     server: {
       getLaunchProject: async () => null,
@@ -276,11 +304,16 @@ export function createMockORPCClient(options: MockORPCClientOptions = {}): APICl
         sendToBackground: async () => ({ success: true, data: undefined }),
       },
       stats: {
-        subscribe: async function* () {
-          // Empty generator - no stats updates in mock
-          return;
+        subscribe: async function* (input: { workspaceId: string }) {
+          const snapshot = workspaceStatsSnapshots.get(input.workspaceId);
+          if (snapshot) {
+            yield snapshot;
+          }
+        },
+        clear: async (input: { workspaceId: string }) => {
+          workspaceStatsSnapshots.delete(input.workspaceId);
+          return { success: true, data: undefined };
         },
-        clear: async () => ({ success: true, data: undefined }),
       },
       getSessionUsage: async (input: { workspaceId: string }) => sessionUsage.get(input.workspaceId),
       mcp: {

src/browser/contexts/FeatureFlagsContext.tsx

@@ -1,4 +1,9 @@
 import React, { createContext, useContext, useEffect, useState, type ReactNode } from "react";
+import { useAPI } from "@/browser/contexts/API";
+
+function isStorybookIframe(): boolean {
+  return typeof window !== "undefined" && window.location.pathname.endsWith("iframe.html");
+}
 
 export type StatsTabVariant = "control" | "stats";
 export type StatsTabOverride = "default" | "on" | "off";
@@ -23,41 +28,55 @@ export function useFeatureFlags(): FeatureFlagsContextValue {
   return ctx;
 }
 
-async function fetchStatsTabState(): Promise<StatsTabState> {
-  const client = window.__ORPC_CLIENT__;
-  if (!client) {
-    throw new Error("ORPC client not initialized");
-  }
-
-  return await client.features.getStatsTabState();
-}
-
 export function FeatureFlagsProvider(props: { children: ReactNode }) {
-  const [statsTabState, setStatsTabState] = useState<StatsTabState | null>(null);
+  const { api } = useAPI();
+  const [statsTabState, setStatsTabState] = useState<StatsTabState | null>(() => {
+    if (isStorybookIframe()) {
+      return { enabled: true, variant: "stats", override: "default" };
+    }
+
+    return null;
+  });
 
   const refreshStatsTabState = async (): Promise<void> => {
-    const state = await fetchStatsTabState();
+    if (!api) {
+      setStatsTabState({ enabled: false, variant: "control", override: "default" });
+      return;
+    }
+
+    const state = await api.features.getStatsTabState();
     setStatsTabState(state);
   };
 
   const setStatsTabOverride = async (override: StatsTabOverride): Promise<void> => {
-    const client = window.__ORPC_CLIENT__;
-    if (!client) throw new Error("ORPC client not initialized");
-    const state = await client.features.setStatsTabOverride({ override });
+    if (!api) {
+      throw new Error("ORPC client not initialized");
+    }
+
+    const state = await api.features.setStatsTabOverride({ override });
     setStatsTabState(state);
   };
 
   useEffect(() => {
+    if (isStorybookIframe()) {
+      return;
+    }
+
     (async () => {
       try {
-        const state = await fetchStatsTabState();
+        if (!api) {
+          setStatsTabState({ enabled: false, variant: "control", override: "default" });
+          return;
+        }
+
+        const state = await api.features.getStatsTabState();
         setStatsTabState(state);
       } catch {
         // Treat as disabled if we can't fetch.
         setStatsTabState({ enabled: false, variant: "control", override: "default" });
       }
     })();
-  }, []);
+  }, [api]);
 
   return (
     <FeatureFlagsContext.Provider

src/browser/stories/App.rightsidebar.stories.tsx

@@ -226,15 +226,10 @@ export const StatsTabIdle: AppStory = {
   play: async ({ canvasElement }) => {
     const canvas = within(canvasElement);
 
-    // Wait for Stats tab to be visible and selected
-    await waitFor(
-      () => {
-        canvas.getByRole("tab", { name: /^stats$/i, selected: true });
-      },
-      { timeout: 5000 }
-    );
+    // Feature flags are async, so allow more time.
+    const statsTab = await canvas.findByRole("tab", { name: /^stats$/i }, { timeout: 3000 });
+    await userEvent.click(statsTab);
 
-    // Verify idle state message is shown
     await waitFor(() => {
       canvas.getByText(/no timing data yet/i);
     });
@@ -254,6 +249,7 @@ export const StatsTabStreaming: AppStory = {
           workspaceId: "ws-stats-streaming",
           workspaceName: "feature/streaming",
           projectName: "my-app",
+          statsTabEnabled: true,
           messages: [
             createUserMessage("msg-1", "Write a comprehensive test suite", { historySequence: 1 }),
           ],
@@ -267,7 +263,10 @@ export const StatsTabStreaming: AppStory = {
   play: async ({ canvasElement }) => {
     const canvas = within(canvasElement);
 
-    // Wait for Stats tab to be visible and selected
+    // Feature flags are async; wait for Stats tab to appear, then select it.
+    const statsTab = await canvas.findByRole("tab", { name: /^stats$/i }, { timeout: 5000 });
+    await userEvent.click(statsTab);
+
     await waitFor(
       () => {
         canvas.getByRole("tab", { name: /^stats$/i, selected: true });

src/browser/stories/meta.tsx

@@ -38,8 +38,26 @@ interface AppWithMocksProps {
 }
 
 /** Wrapper that runs setup once and passes the client to AppLoader */
+function getStorybookStoryId(): string | null {
+  if (typeof window === "undefined") {
+    return null;
+  }
+
+  const params = new URLSearchParams(window.location.search);
+  return params.get("id") ?? params.get("path");
+}
+
 export const AppWithMocks: FC<AppWithMocksProps> = ({ setup }) => {
+  const lastStoryIdRef = useRef<string | null>(null);
   const clientRef = useRef<APIClient | null>(null);
+
+  const storyId = getStorybookStoryId();
+  if (lastStoryIdRef.current !== storyId) {
+    lastStoryIdRef.current = storyId;
+    clientRef.current = setup();
+  }
+
   clientRef.current ??= setup();
+
   return <AppLoader client={clientRef.current} />;
 };

src/browser/stories/storyHelpers.ts

@@ -6,7 +6,12 @@
  */
 
 import type { FrontendWorkspaceMetadata } from "@/common/types/workspace";
-import type { WorkspaceChatMessage, ChatMuxMessage, ProvidersConfigMap } from "@/common/orpc/types";
+import type {
+  WorkspaceChatMessage,
+  ChatMuxMessage,
+  ProvidersConfigMap,
+  WorkspaceStatsSnapshot,
+} from "@/common/orpc/types";
 import type { APIClient } from "@/browser/contexts/API";
 import {
   SELECTED_WORKSPACE_KEY,
@@ -161,6 +166,7 @@ export interface SimpleChatSetupOptions {
   providersConfig?: ProvidersConfigMap;
   backgroundProcesses?: BackgroundProcessFixture[];
   /** Session usage data for Costs tab */
+  statsTabEnabled?: boolean;
   sessionUsage?: MockSessionUsage;
   /** Optional custom chat handler for emitting additional events (e.g., queued-message-changed) */
   onChat?: (workspaceId: string, emit: (msg: WorkspaceChatMessage) => void) => void;
@@ -216,6 +222,7 @@ export function setupSimpleChatStory(opts: SimpleChatSetupOptions): APIClient {
     executeBash: createGitStatusExecutor(gitStatus),
     providersConfig: opts.providersConfig,
     backgroundProcesses: bgProcesses,
+    statsTabVariant: opts.statsTabEnabled ? "stats" : "control",
     sessionUsage: sessionUsageMap,
   });
 }
@@ -235,6 +242,7 @@ export interface StreamingChatSetupOptions {
   streamText?: string;
   pendingTool?: { toolCallId: string; toolName: string; args: object };
   gitStatus?: GitStatusFixture;
+  statsTabEnabled?: boolean;
 }
 
 /**
@@ -272,12 +280,48 @@ export function setupStreamingChatStory(opts: StreamingChatSetupOptions): APICli
   // Set localStorage for workspace selection
   selectWorkspace(workspaces[0]);
 
+  const workspaceStatsSnapshots = new Map<string, WorkspaceStatsSnapshot>();
+  if (opts.statsTabEnabled) {
+    workspaceStatsSnapshots.set(workspaceId, {
+      workspaceId,
+      generatedAt: Date.now(),
+      active: {
+        messageId: opts.streamingMessageId,
+        model: "openai:gpt-4o",
+        elapsedMs: 2000,
+        ttftMs: 200,
+        toolExecutionMs: 0,
+        modelTimeMs: 2000,
+        streamingMs: 1800,
+        outputTokens: 100,
+        reasoningTokens: 0,
+        liveTokenCount: 100,
+        liveTPS: 50,
+        invalid: false,
+        anomalies: [],
+      },
+      session: {
+        totalDurationMs: 0,
+        totalToolExecutionMs: 0,
+        totalStreamingMs: 0,
+        totalTtftMs: 0,
+        ttftCount: 0,
+        responseCount: 0,
+        totalOutputTokens: 0,
+        totalReasoningTokens: 0,
+        byModel: {},
+      },
+    });
+  }
+
   // Return ORPC client
   return createMockORPCClient({
     projects: groupWorkspacesByProject(workspaces),
     workspaces,
     onChat: createOnChatAdapter(chatHandlers),
     executeBash: createGitStatusExecutor(gitStatus),
+    workspaceStatsSnapshots,
+    statsTabVariant: opts.statsTabEnabled ? "stats" : "control",
   });
 }
 

src/common/telemetry/client.ts

@@ -89,7 +89,7 @@ export function trackEvent(payload: TelemetryEventPayload): void {
   }
 
   const client = window.__ORPC_CLIENT__;
-  if (!client) {
+  if (!client?.telemetry?.track) {
     return;
   }