better way to mock singleton or static keyed models

Author: sawka

.kilocode/skills/waveenv/SKILL.md

@@ -80,6 +80,14 @@ export type MyEnv = WaveEnvSubset<{
 }>;
 ```
 
+### Automatically Included Fields
+
+Every `WaveEnvSubset<T>` automatically includes the mock fields — you never need to declare them:
+
+- `isMock: boolean`
+- `mockSetWaveObj: <T extends WaveObj>(oref: string, obj: T) => void`
+- `mockModels?: Map<any, any>`
+
 ### Rules for Each Section
 
 | Section                    | Pattern                                                | Notes                                                                                              |

frontend/app/store/keymodel.ts

@@ -1,4 +1,4 @@
-// Copyright 2025, Command Line Inc.
+// Copyright 2026, Command Line Inc.
 // SPDX-License-Identifier: Apache-2.0
 
 import { WaveAIModel } from "@/app/aipanel/waveai-model";
@@ -490,7 +490,7 @@ function tryReinjectKey(event: WaveKeyboardEvent): boolean {
 function countTermBlocks(): number {
     const allBCMs = getAllBlockComponentModels();
     let count = 0;
-    let gsGetBound = globalStore.get.bind(globalStore);
+    const gsGetBound = globalStore.get.bind(globalStore);
     for (const bcm of allBCMs) {
         const viewModel = bcm.viewModel;
         if (viewModel.viewType == "term" && viewModel.isBasicTerm?.(gsGetBound)) {

frontend/app/store/tab-model.ts

@@ -1,24 +1,28 @@
 // Copyright 2026, Command Line Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-import { useWaveEnv, WaveEnv } from "@/app/waveenv/waveenv";
+import { WaveEnv, WaveEnvSubset } from "@/app/waveenv/waveenv";
 import { atom, Atom, PrimitiveAtom } from "jotai";
 import { createContext, useContext } from "react";
 import { globalStore } from "./jotaiStore";
 import * as WOS from "./wos";
 
+export type TabModelEnv = WaveEnvSubset<{
+    wos: WaveEnv["wos"];
+}>;
+
 const tabModelCache = new Map<string, TabModel>();
 export const activeTabIdAtom = atom<string>(null) as PrimitiveAtom<string>;
 
 export class TabModel {
     tabId: string;
-    waveEnv: WaveEnv;
+    waveEnv: TabModelEnv;
     tabAtom: Atom<Tab>;
     tabNumBlocksAtom: Atom<number>;
     isTermMultiInput = atom(false) as PrimitiveAtom<boolean>;
     metaCache: Map<string, Atom<any>> = new Map();
 
-    constructor(tabId: string, waveEnv?: WaveEnv) {
+    constructor(tabId: string, waveEnv?: TabModelEnv) {
         this.tabId = tabId;
         this.waveEnv = waveEnv;
         this.tabAtom = atom((get) => {
@@ -46,16 +50,25 @@ export class TabModel {
     }
 }
 
-export function getTabModelByTabId(tabId: string, waveEnv?: WaveEnv): TabModel {
-    let model = tabModelCache.get(tabId);
+export function getTabModelByTabId(tabId: string, waveEnv?: TabModelEnv): TabModel {
+    if (!waveEnv?.isMock) {
+        let model = tabModelCache.get(tabId);
+        if (model == null) {
+            model = new TabModel(tabId, waveEnv);
+            tabModelCache.set(tabId, model);
+        }
+        return model;
+    }
+    const key = `TabModel:${tabId}`;
+    let model = waveEnv.mockModels.get(key);
     if (model == null) {
         model = new TabModel(tabId, waveEnv);
-        tabModelCache.set(tabId, model);
+        waveEnv.mockModels.set(key, model);
     }
     return model;
 }
 
-export function getActiveTabModel(waveEnv?: WaveEnv): TabModel | null {
+export function getActiveTabModel(waveEnv?: TabModelEnv): TabModel | null {
     const activeTabId = globalStore.get(activeTabIdAtom);
     if (activeTabId == null) {
         return null;
@@ -66,22 +79,13 @@ export function getActiveTabModel(waveEnv?: WaveEnv): TabModel | null {
 export const TabModelContext = createContext<TabModel | undefined>(undefined);
 
 export function useTabModel(): TabModel {
-    const waveEnv = useWaveEnv();
     const ctxModel = useContext(TabModelContext);
-    if (waveEnv?.mockTabModel != null) {
-        return waveEnv.mockTabModel;
-    }
     if (ctxModel == null) {
         throw new Error("useTabModel must be used within a TabModelProvider");
     }
     return ctxModel;
 }
 
 export function useTabModelMaybe(): TabModel {
-    const waveEnv = useWaveEnv();
-    const ctxModel = useContext(TabModelContext);
-    if (waveEnv?.mockTabModel != null) {
-        return waveEnv.mockTabModel;
-    }
-    return ctxModel;
+    return useContext(TabModelContext);
 }

frontend/app/waveenv/waveenv.ts

@@ -1,7 +1,6 @@
 // Copyright 2026, Command Line Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-import type { TabModel } from "@/app/store/tab-model";
 import { RpcApiType } from "@/app/store/wshclientapi";
 import { Atom, PrimitiveAtom } from "jotai";
 import React from "react";
@@ -35,16 +34,24 @@ type ComplexWaveEnvKeys = {
     wos: WaveEnv["wos"];
 };
 
-export type WaveEnvSubset<T> = OmitNever<{
-    [K in keyof T]: K extends keyof ComplexWaveEnvKeys
-        ? Subset<T[K], ComplexWaveEnvKeys[K]>
-        : K extends keyof WaveEnv
-          ? T[K]
-          : never;
-}>;
+type WaveEnvMockFields = {
+    isMock: WaveEnv["isMock"];
+    mockSetWaveObj: WaveEnv["mockSetWaveObj"];
+    mockModels: WaveEnv["mockModels"];
+};
+
+export type WaveEnvSubset<T> = WaveEnvMockFields &
+    OmitNever<{
+        [K in keyof T]: K extends keyof ComplexWaveEnvKeys
+            ? Subset<T[K], ComplexWaveEnvKeys[K]>
+            : K extends keyof WaveEnv
+              ? T[K]
+              : never;
+    }>;
 
 // default implementation for production is in ./waveenvimpl.ts
 export type WaveEnv = {
+    isMock: boolean;
     electron: ElectronApi;
     rpc: RpcApiType;
     platform: NodeJS.Platform;
@@ -68,7 +75,7 @@ export type WaveEnv = {
 
     // the mock fields are only usable in the preview server (may be be null or throw errors in production)
     mockSetWaveObj: <T extends WaveObj>(oref: string, obj: T) => void;
-    mockTabModel?: TabModel;
+    mockModels?: Map<any, any>;
 };
 
 export const WaveEnvContext = React.createContext<WaveEnv>(null);

frontend/app/waveenv/waveenvimpl.ts

@@ -19,6 +19,7 @@ import { isMacOS, isWindows, PLATFORM } from "@/util/platformutil";
 
 export function makeWaveEnvImpl(): WaveEnv {
     return {
+        isMock: false,
         electron: (window as any).api,
         rpc: RpcApi,
         getSettingsKeyAtom,
@@ -41,8 +42,10 @@ export function makeWaveEnvImpl(): WaveEnv {
         },
         getBlockMetaKeyAtom,
         getConnConfigKeyAtom,
+
         mockSetWaveObj: <T extends WaveObj>(_oref: string, _obj: T) => {
             throw new Error("mockSetWaveObj is only available in the preview server");
         },
+        mockModels: new Map<any, any>(),
     };
 }

frontend/preview/mock/mockwaveenv.ts

@@ -3,7 +3,6 @@
 
 import { makeDefaultConnStatus } from "@/app/store/global";
 import { globalStore } from "@/app/store/jotaiStore";
-import { TabModel } from "@/app/store/tab-model";
 import { handleWaveEvent } from "@/app/store/wps";
 import { RpcApiType } from "@/app/store/wshclientapi";
 import { WaveEnv } from "@/app/waveenv/waveenv";
@@ -260,6 +259,7 @@ export function makeMockWaveEnv(mockEnv?: MockEnv): MockWaveEnv {
         },
     };
     const env = {
+        isMock: true,
         mockEnv: overrides,
         electron: {
             ...previewElectronApi,
@@ -318,7 +318,6 @@ export function makeMockWaveEnv(mockEnv?: MockEnv): MockWaveEnv {
                 return [useAtomValue(objAtom), false];
             },
         },
-        mockSetWaveObj: mockWosFns.mockSetWaveObj,
         getBlockMetaKeyAtom: <T extends keyof MetaType>(blockId: string, key: T) => {
             const cacheKey = blockId + "#meta-" + key;
             if (!blockMetaKeyAtomCache.has(cacheKey)) {
@@ -343,10 +342,8 @@ export function makeMockWaveEnv(mockEnv?: MockEnv): MockWaveEnv {
             }
             return connConfigKeyAtomCache.get(cacheKey) as Atom<ConnKeywords[T]>;
         },
-        mockTabModel: null as TabModel,
+        mockSetWaveObj: mockWosFns.mockSetWaveObj,
+        mockModels: new Map<any, any>(),
     } as MockWaveEnv;
-    if (overrides.tabId != null) {
-        env.mockTabModel = new TabModel(overrides.tabId, env);
-    }
     return env;
 }