opencode-morphllm 0.0.4 → 0.0.6

package/dist/shared/config.test.js

@@ -0,0 +1,204 @@
+import {
+  describe,
+  it,
+  expect,
+  beforeEach,
+  vi,
+  beforeAll,
+  afterAll,
+} from 'bun:test';
+import { existsSync, writeFileSync, rmSync, mkdirSync } from 'node:fs';
+import { join } from 'node:path';
+import { tmpdir } from 'node:os';
+// Mock the opencode-config-dir module before importing config
+const mockConfigDir = join(tmpdir(), 'mock-opencode-config-test');
+const mockMorphJson = join(mockConfigDir, 'morph.json');
+const mockMorphJsonc = join(mockConfigDir, 'morph.jsonc');
+vi.mock('./opencode-config-dir', () => ({
+  getOpenCodeConfigDir: vi.fn(() => mockConfigDir),
+}));
+import {
+  getMorphPluginConfigPath,
+  loadMorphPluginConfig,
+  loadMorphPluginConfigWithProjectOverride,
+} from './config';
+describe('config.ts', () => {
+  beforeAll(() => {
+    // Create mock config directory
+    if (!existsSync(mockConfigDir)) {
+      mkdirSync(mockConfigDir, { recursive: true });
+    }
+  });
+  afterAll(() => {
+    // Cleanup
+    try {
+      if (existsSync(mockMorphJson)) rmSync(mockMorphJson);
+      if (existsSync(mockMorphJsonc)) rmSync(mockMorphJsonc);
+      if (existsSync(mockConfigDir)) rmSync(mockConfigDir, { recursive: true });
+    } catch {
+      // Ignore cleanup errors
+    }
+  });
+  beforeEach(() => {
+    // Clean up any existing config files before each test
+    try {
+      if (existsSync(mockMorphJson)) rmSync(mockMorphJson);
+      if (existsSync(mockMorphJsonc)) rmSync(mockMorphJsonc);
+    } catch {
+      // Ignore cleanup errors
+    }
+  });
+  describe('getMorphPluginConfigPath', () => {
+    it('should return path to morph.json in config directory', () => {
+      const path = getMorphPluginConfigPath();
+      expect(path).toContain('morph.json');
+      expect(path).toContain(mockConfigDir);
+    });
+  });
+  describe('loadMorphPluginConfig', () => {
+    it('should return null when no config files exist', () => {
+      const result = loadMorphPluginConfig();
+      expect(result).toBeNull();
+    });
+    it('should load config from .jsonc file', () => {
+      const content = `{
+        "MORPH_API_KEY": "test-key-123",
+        "MORPH_ROUTER_ENABLED": false
+      }`;
+      writeFileSync(mockMorphJsonc, content);
+      const result = loadMorphPluginConfig();
+      expect(result).not.toBeNull();
+      expect(result?.MORPH_API_KEY).toBe('test-key-123');
+      expect(result?.MORPH_ROUTER_ENABLED).toBe(false);
+    });
+    it('should load config from .json file', () => {
+      const content = JSON.stringify({
+        MORPH_API_KEY: 'json-key-456',
+        MORPH_MODEL_EASY: 'provider/model-easy',
+      });
+      writeFileSync(mockMorphJson, content);
+      const result = loadMorphPluginConfig();
+      expect(result).not.toBeNull();
+      expect(result?.MORPH_API_KEY).toBe('json-key-456');
+      expect(result?.MORPH_MODEL_EASY).toBe('provider/model-easy');
+    });
+    it('should prefer .jsonc over .json when both exist', () => {
+      const jsoncContent = '{"MORPH_API_KEY": "from-jsonc"}';
+      const jsonContent = '{"MORPH_API_KEY": "from-json"}';
+      writeFileSync(mockMorphJsonc, jsoncContent);
+      writeFileSync(mockMorphJson, jsonContent);
+      const result = loadMorphPluginConfig();
+      expect(result?.MORPH_API_KEY).toBe('from-jsonc');
+    });
+    it('should return null for invalid .jsonc content', () => {
+      writeFileSync(mockMorphJsonc, 'invalid json content');
+      const result = loadMorphPluginConfig();
+      expect(result).toBeNull();
+    });
+    it('should return null for invalid .json content', () => {
+      writeFileSync(mockMorphJson, 'invalid json content');
+      const result = loadMorphPluginConfig();
+      expect(result).toBeNull();
+    });
+    it('should handle all MORPH config fields', () => {
+      const content = JSON.stringify({
+        MORPH_API_KEY: 'api-key',
+        MORPH_ROUTER_ENABLED: true,
+        MORPH_MODEL_EASY: 'easy/provider/model',
+        MORPH_MODEL_MEDIUM: 'medium/provider/model',
+        MORPH_MODEL_HARD: 'hard/provider/model',
+        MORPH_MODEL_DEFAULT: 'default/provider/model',
+      });
+      writeFileSync(mockMorphJson, content);
+      const result = loadMorphPluginConfig();
+      expect(result).toEqual({
+        MORPH_API_KEY: 'api-key',
+        MORPH_ROUTER_ENABLED: true,
+        MORPH_MODEL_EASY: 'easy/provider/model',
+        MORPH_MODEL_MEDIUM: 'medium/provider/model',
+        MORPH_MODEL_HARD: 'hard/provider/model',
+        MORPH_MODEL_DEFAULT: 'default/provider/model',
+      });
+    });
+    it('should handle nested MORPH_ROUTER_CONFIGS', () => {
+      const content = JSON.stringify({
+        MORPH_API_KEY: 'api-key',
+        MORPH_ROUTER_CONFIGS: {
+          MORPH_ROUTER_ENABLED: true,
+          MORPH_MODEL_EASY: 'easy/provider/model',
+          MORPH_MODEL_MEDIUM: 'medium/provider/model',
+          MORPH_MODEL_HARD: 'hard/provider/model',
+          MORPH_MODEL_DEFAULT: 'default/provider/model',
+        },
+      });
+      writeFileSync(mockMorphJson, content);
+      const result = loadMorphPluginConfig();
+      expect(result?.MORPH_API_KEY).toBe('api-key');
+      expect(result?.MORPH_ROUTER_CONFIGS).toEqual({
+        MORPH_ROUTER_ENABLED: true,
+        MORPH_MODEL_EASY: 'easy/provider/model',
+        MORPH_MODEL_MEDIUM: 'medium/provider/model',
+        MORPH_MODEL_HARD: 'hard/provider/model',
+        MORPH_MODEL_DEFAULT: 'default/provider/model',
+      });
+    });
+  });
+  describe('loadMorphPluginConfigWithProjectOverride', () => {
+    it('should return empty object when no config exists', () => {
+      const result =
+        loadMorphPluginConfigWithProjectOverride('/non-existent-path');
+      expect(result).toEqual({});
+    });
+    it('should merge user config with project config', () => {
+      // Create user config
+      writeFileSync(
+        mockMorphJson,
+        JSON.stringify({
+          MORPH_API_KEY: 'user-api-key',
+          MORPH_MODEL_EASY: 'user-easy-model',
+        })
+      );
+      // Create project config directory
+      const projectDir = join(tmpdir(), 'test-project');
+      const projectConfigDir = join(projectDir, '.opencode');
+      const projectConfigPath = join(projectConfigDir, 'morph.json');
+      mkdirSync(projectConfigDir, { recursive: true });
+      writeFileSync(
+        projectConfigPath,
+        JSON.stringify({
+          MORPH_API_KEY: 'project-api-key',
+          MORPH_MODEL_HARD: 'project-hard-model',
+        })
+      );
+      const result = loadMorphPluginConfigWithProjectOverride(projectDir);
+      // Project config should override user config
+      expect(result.MORPH_API_KEY).toBe('project-api-key');
+      expect(result.MORPH_MODEL_EASY).toBe('user-easy-model');
+      expect(result.MORPH_MODEL_HARD).toBe('project-hard-model');
+    });
+    it('should handle project .jsonc config', () => {
+      // Create user config
+      writeFileSync(
+        mockMorphJson,
+        JSON.stringify({
+          MORPH_API_KEY: 'user-key',
+        })
+      );
+      // Create project config with .jsonc
+      const projectDir = join(tmpdir(), 'test-project2');
+      const projectConfigDir = join(projectDir, '.opencode');
+      const projectConfigPath = join(projectConfigDir, 'morph.jsonc');
+      mkdirSync(projectConfigDir, { recursive: true });
+      const jsoncContent = `
+        // Project config with comments
+        {
+          "MORPH_MODEL_MEDIUM": "project-medium-model"
+        }
+      `;
+      writeFileSync(projectConfigPath, jsoncContent);
+      const result = loadMorphPluginConfigWithProjectOverride(projectDir);
+      expect(result.MORPH_API_KEY).toBe('user-key');
+      expect(result.MORPH_MODEL_MEDIUM).toBe('project-medium-model');
+    });
+  });
+});

package/dist/shared/opencode-config-dir.d.ts

@@ -0,0 +1,9 @@
+interface OpenCodeConfigDirOptions {
+  binary: 'opencode' | 'opencode-desktop';
+  version?: string | null;
+  checkExisting?: boolean;
+}
+export declare function getOpenCodeConfigDir(
+  options: OpenCodeConfigDirOptions
+): string;
+export {};

package/dist/shared/opencode-config-dir.js

@@ -0,0 +1,56 @@
+import { existsSync } from 'node:fs';
+import { homedir } from 'node:os';
+import { join, resolve } from 'node:path';
+function getTauriConfigDir(identifier) {
+  const platform = process.platform;
+  switch (platform) {
+    case 'darwin':
+      return join(homedir(), 'Library', 'Application Support', identifier);
+    case 'win32': {
+      const appData =
+        process.env.APPDATA || join(homedir(), 'AppData', 'Roaming');
+      return join(appData, identifier);
+    }
+    case 'linux':
+    default: {
+      const xdgConfig =
+        process.env.XDG_CONFIG_HOME || join(homedir(), '.config');
+      return join(xdgConfig, identifier);
+    }
+  }
+}
+function getCliConfigDir() {
+  const envConfigDir = process.env.OPENCODE_CONFIG_DIR?.trim();
+  if (envConfigDir) {
+    return resolve(envConfigDir);
+  }
+  if (process.platform === 'win32') {
+    const crossPlatformDir = join(homedir(), '.config', 'opencode');
+    const crossPlatformConfig = join(crossPlatformDir, 'opencode.json');
+    if (existsSync(crossPlatformConfig)) {
+      return crossPlatformDir;
+    }
+    const appData =
+      process.env.APPDATA || join(homedir(), 'AppData', 'Roaming');
+    const appdataDir = join(appData, 'opencode');
+    const appdataConfig = join(appdataDir, 'opencode.json');
+    if (existsSync(appdataConfig)) {
+      return appdataDir;
+    }
+    return crossPlatformDir;
+  }
+  const xdgConfig = process.env.XDG_CONFIG_HOME || join(homedir(), '.config');
+  return join(xdgConfig, 'opencode');
+}
+export function getOpenCodeConfigDir(options) {
+  if (options.binary === 'opencode-desktop') {
+    const version = options.version;
+    const isDev =
+      !!version && (version.includes('-dev') || version.includes('.dev'));
+    const identifier = isDev
+      ? 'ai.opencode.desktop.dev'
+      : 'ai.opencode.desktop';
+    return getTauriConfigDir(identifier);
+  }
+  return getCliConfigDir();
+}

package/package.json

@@ -1,11 +1,11 @@
 {
   "name": "opencode-morphllm",
-  "version": "0.0.4",
+  "version": "0.0.6",
   "author": "Vito Lin",
   "main": "dist/index.js",
   "description": "OpenCode plugin for MorphLLM",
   "scripts": {
-    "build": "tsc -p tsconfig.json",
+    "build": "rm -rf dist && tsc -p tsconfig.json",
     "test": "bun test",
     "format": "prettier --write .",
     "format:check": "prettier --check .",

package/src/index.ts

@@ -1,9 +1,9 @@
 import { Plugin } from '@opencode-ai/plugin';
 import type { McpLocalConfig } from '@opencode-ai/sdk';
 
-import { createBuiltinMcps } from './mcps';
-import { createModelRouterHook } from './router';
-import { MORPH_ROUTER_ENABLED } from './const';
+import { createBuiltinMcps } from './morph/mcps';
+import { createModelRouterHook } from './morph/router';
+import { MORPH_ROUTER_ENABLED } from './shared/config';
 
 const MorphOpenCodePlugin: Plugin = async () => {
   const builtinMcps: Record<string, McpLocalConfig> = createBuiltinMcps();

package/src/morph/mcps.test.ts

@@ -0,0 +1,51 @@
+import { describe, it, expect, vi } from 'bun:test';
+
+vi.mock('../shared/config', () => ({
+  API_KEY: 'test-api-key-123',
+}));
+
+import { createBuiltinMcps } from './mcps';
+
+describe('mcps.ts', () => {
+  describe('createBuiltinMcps', () => {
+    it('should create morph_mcp configuration', () => {
+      const mcps = createBuiltinMcps();
+
+      expect(mcps).toHaveProperty('morph_mcp');
+      expect(mcps.morph_mcp.type).toBe('local');
+      expect(mcps.morph_mcp.enabled).toBe(true);
+    });
+
+    it('should set correct command for morph_mcp', () => {
+      const mcps = createBuiltinMcps();
+
+      expect(mcps.morph_mcp.command).toEqual([
+        'npx',
+        '-y',
+        '@morphllm/morphmcp',
+      ]);
+    });
+
+    it('should set MORPH_API_KEY in environment', () => {
+      const mcps = createBuiltinMcps();
+      const env = mcps.morph_mcp.environment;
+
+      expect(env).toBeDefined();
+      expect(env?.MORPH_API_KEY).toBe('test-api-key-123');
+    });
+
+    it('should set ENABLED_TOOLS environment variable', () => {
+      const mcps = createBuiltinMcps();
+      const env = mcps.morph_mcp.environment;
+
+      expect(env).toBeDefined();
+      expect(env?.ENABLED_TOOLS).toBe('edit_file,warpgrep_codebase_search');
+    });
+
+    it('should only create morph_mcp key', () => {
+      const mcps = createBuiltinMcps();
+
+      expect(Object.keys(mcps)).toEqual(['morph_mcp']);
+    });
+  });
+});

package/src/{mcps.ts → morph/mcps.ts}

@@ -1,5 +1,5 @@
 import type { McpLocalConfig } from '@opencode-ai/sdk';
-import { API_KEY } from './const';
+import { API_KEY } from '../shared/config';
 
 export function createBuiltinMcps(): Record<string, McpLocalConfig> {
   return {

package/src/morph/router.test.ts

@@ -0,0 +1,266 @@
+import { describe, it, expect, vi } from 'bun:test';
+import type { Part } from '@opencode-ai/sdk';
+
+vi.mock('@morphllm/morphsdk', () => ({
+  MorphClient: vi.fn().mockImplementation(() => ({
+    routers: {
+      raw: {
+        classify: vi.fn(),
+      },
+    },
+  })),
+}));
+
+vi.mock('../shared/config', () => ({
+  API_KEY: 'sk-test-api-key-123',
+  MORPH_MODEL_EASY: 'easy/easy',
+  MORPH_MODEL_MEDIUM: 'medium/medium',
+  MORPH_MODEL_HARD: 'hard/hard',
+  MORPH_MODEL_DEFAULT: 'default/default',
+  MORPH_ROUTER_ONLY_FIRST_MESSAGE: false,
+}));
+
+import { createModelRouterHook, extractPromptText } from './router';
+
+describe('router.ts', () => {
+  describe('extractPromptText', () => {
+    it('should extract text from parts', () => {
+      const parts: Part[] = [
+        { type: 'text', text: 'Hello' },
+        { type: 'text', text: 'World' },
+      ] as any;
+      expect(extractPromptText(parts)).toBe('Hello World');
+    });
+
+    it('should handle empty parts', () => {
+      const parts: Part[] = [];
+      expect(extractPromptText(parts)).toBe('');
+    });
+
+    it('should filter out non-text parts', () => {
+      const parts: Part[] = [
+        { type: 'text', text: 'Hello' },
+        { type: 'other', text: 'Ignore' },
+        { type: 'text', text: 'World' },
+      ] as any;
+      expect(extractPromptText(parts)).toBe('Hello World');
+    });
+
+    it('should handle parts with no text', () => {
+      const parts: Part[] = [
+        { type: 'text', text: 'Hello' },
+        { type: 'text' },
+        { type: 'text', text: 'World' },
+      ] as any;
+      expect(extractPromptText(parts)).toBe('Hello  World');
+    });
+
+    it('should handle all non-text parts', () => {
+      const parts: Part[] = [
+        { type: 'image', data: 'base64...' },
+        { type: 'tool_use', name: 'bash' },
+      ] as any;
+      expect(extractPromptText(parts)).toBe('');
+    });
+  });
+
+  describe('createModelRouterHook', () => {
+    it('should return a chat.message hook', () => {
+      const hook = createModelRouterHook();
+      expect('chat.message' in hook).toBe(true);
+      expect(typeof hook['chat.message']).toBe('function');
+    });
+
+    it('should call the classifier with the correct input', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockResolvedValue({ difficulty: 'easy' });
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+      await hook['chat.message'](input as any, output);
+      expect(classify).toHaveBeenCalledWith({ input: 'test prompt' });
+    });
+
+    it('should assign the correct model based on difficulty', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockResolvedValue({ difficulty: 'hard' });
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+      await hook['chat.message'](input as any, output);
+      expect(input).toHaveProperty('model');
+      expect((input as any).model.providerID).toBe('hard');
+      expect((input as any).model.modelID).toBe('hard');
+    });
+
+    it('should use the default model if no difficulty is returned', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockResolvedValue({});
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+      await hook['chat.message'](input as any, output);
+      expect(input).toHaveProperty('model');
+      expect((input as any).model.providerID).toBe('default');
+      expect((input as any).model.modelID).toBe('default');
+    });
+
+    it('should use default model when router returns nothing', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockResolvedValue({});
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+      await hook['chat.message'](input as any, output);
+      expect(input).toHaveProperty('model');
+      expect((input as any).model.providerID).toBe('default');
+      expect((input as any).model.modelID).toBe('default');
+    });
+
+    it('should handle classifier throwing an error gracefully', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockRejectedValue(new Error('API Error'));
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+
+      await expect(hook['chat.message'](input as any, output)).rejects.toThrow(
+        'API Error'
+      );
+    });
+
+    it('should handle classifier returning null', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockResolvedValue(null);
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+      await hook['chat.message'](input as any, output);
+      expect((input as any).model.providerID).toBe('default');
+      expect((input as any).model.modelID).toBe('default');
+    });
+
+    it('should handle classifier returning undefined difficulty', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockResolvedValue({ difficulty: undefined });
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+      await hook['chat.message'](input as any, output);
+      expect((input as any).model.providerID).toBe('default');
+      expect((input as any).model.modelID).toBe('default');
+    });
+
+    it('should default to default model for invalid difficulty values', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockResolvedValue({ difficulty: 'ultra-hard' });
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+      await hook['chat.message'](input as any, output);
+      expect((input as any).model.providerID).toBe('default');
+      expect((input as any).model.modelID).toBe('default');
+    });
+
+    it('should default to default model for empty string difficulty', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockResolvedValue({ difficulty: '' });
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+      await hook['chat.message'](input as any, output);
+      expect((input as any).model.providerID).toBe('default');
+      expect((input as any).model.modelID).toBe('default');
+    });
+
+    it('should handle case-insensitive difficulty matching', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi.fn().mockResolvedValue({ difficulty: 'HARD' });
+      const input = {
+        sessionID: '123',
+        classify,
+      };
+      const output = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'test prompt' }],
+      } as any;
+      await hook['chat.message'](input as any, output);
+      expect((input as any).model.providerID).toBe('hard');
+      expect((input as any).model.modelID).toBe('hard');
+    });
+
+    it('should route all messages when MORPH_ROUTER_ONLY_FIRST_MESSAGE is disabled', async () => {
+      const hook = createModelRouterHook();
+      const classify = vi
+        .fn()
+        .mockResolvedValueOnce({ difficulty: 'hard' })
+        .mockResolvedValueOnce({ difficulty: 'easy' });
+
+      const sessionID = 'session-456';
+      const input1 = { sessionID, classify };
+      const output1 = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'first message' }],
+      } as any;
+
+      await hook['chat.message'](input1 as any, output1);
+      expect(classify).toHaveBeenCalledTimes(1);
+      expect((input1 as any).model.providerID).toBe('hard');
+
+      const input2 = { sessionID, classify };
+      const output2 = {
+        message: {} as any,
+        parts: [{ type: 'text', text: 'second message' }],
+      } as any;
+
+      await hook['chat.message'](input2 as any, output2);
+      expect(classify).toHaveBeenCalledTimes(2);
+      expect((input2 as any).model.providerID).toBe('easy');
+    });
+  });
+});

package/src/{router.ts → morph/router.ts}

@@ -6,7 +6,8 @@ import {
   MORPH_MODEL_MEDIUM,
   MORPH_MODEL_HARD,
   MORPH_MODEL_DEFAULT,
-} from './const';
+  MORPH_ROUTER_ONLY_FIRST_MESSAGE,
+} from '../shared/config';
 import type { Part, UserMessage } from '@opencode-ai/sdk';
 import type {
   RouterInput,
@@ -14,7 +15,17 @@ import type {
   ComplexityLevel,
 } from '@morphllm/morphsdk';
 
-const morph = new MorphClient({ apiKey: API_KEY });
+// Lazy initialization to allow mocking in tests
+let morph: MorphClient | null = null;
+
+const sessionsWithModelSelected = new Set<string>();
+
+function getMorphClient(): MorphClient {
+  if (!morph) {
+    morph = new MorphClient({ apiKey: API_KEY });
+  }
+  return morph;
+}
 
 function parseModel(s?: string): { providerID: string; modelID: string } {
   if (!s) return { providerID: '', modelID: '' };
@@ -54,11 +65,17 @@ export function createModelRouterHook() {
     ): Promise<void> => {
       input.model = input.model ?? { providerID: '', modelID: '' };
 
+      if (MORPH_ROUTER_ONLY_FIRST_MESSAGE) {
+        if (sessionsWithModelSelected.has(input.sessionID)) {
+          return;
+        }
+      }
+
       const promptText = extractPromptText(output.parts);
 
       const classifier =
         input.classify ??
-        ((args: RouterInput) => morph.routers.raw.classify(args));
+        ((args: RouterInput) => getMorphClient().routers.raw.classify(args));
 
       const classification: RawRouterResult = await classifier({
         input: promptText,
@@ -69,12 +86,12 @@ export function createModelRouterHook() {
       const finalProviderID = chosen.providerID || input.model.providerID;
       const finalModelID = chosen.modelID || input.model.modelID;
 
-      console.debug(
-        `[Morph Router] Prompt classified as difficulty: ${classification?.difficulty}. Routing to model: ${finalProviderID}/${finalModelID}`
-      );
-
       input.model.providerID = finalProviderID;
       input.model.modelID = finalModelID;
+
+      if (MORPH_ROUTER_ONLY_FIRST_MESSAGE) {
+        sessionsWithModelSelected.add(input.sessionID);
+      }
     },
   };
 }