From 7782503a6cebb5dae7cf48fee452f990de1c2ecf Mon Sep 17 00:00:00 2001 From: ostwal99 Date: Wed, 13 May 2026 19:51:21 +0530 Subject: [PATCH] Add plugin hook model routing state + prompt-based Claude switching --- README.md | 13 +++ pre-tool-use/model-router.js | 167 +++++++++++++++++++++++++++------ test/model-router-hook.test.js | 154 ++++++++++++++++++++++++++++++ 3 files changed, 303 insertions(+), 31 deletions(-) create mode 100644 test/model-router-hook.test.js diff --git a/README.md b/README.md index 2b3c17a..a340f31 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,19 @@ nexus model-router disable nexus model-router enable ``` +When hooks run (`session-start`, `user-prompt-submit`, `pre-tool-use`), Nexus writes runtime model state to: + +```text +/.nexus/model-router-runtime.json +``` + +For Claude runtime, prompt-size routing runs automatically on `user-prompt-submit` and can switch between a default model and the small-edit model. Tune thresholds with: + +```bash +NEXUS_SMALL_PROMPT_MAX_TOKENS=180 +NEXUS_SMALL_PROMPT_MAX_WORDS=50 +``` + ## What Nexus Does In Your Terminal - Coordinate a multi-agent implementation by identifying which tasks can run in parallel. diff --git a/pre-tool-use/model-router.js b/pre-tool-use/model-router.js index 655e5f4..1790f80 100644 --- a/pre-tool-use/model-router.js +++ b/pre-tool-use/model-router.js @@ -3,16 +3,19 @@ const fs = require("fs"); const path = require("path"); const { runBootstrap } = require("../scripts/bootstrap-agent-docs"); -const { isModelRouterEnabled } = require("../scripts/model-router-state"); +const { isModelRouterEnabled, resolveRepoRoot } = require("../scripts/model-router-state"); const MODEL_BY_RUNTIME = { - claude: "claude-haiku-4-5-20251001", + claude: process.env.NEXUS_CLAUDE_DEFAULT_MODEL || "claude-sonnet-4-6-20250514", codex: "gpt-5.4-mini", }; const SETTINGS_PATH = path.join(process.env.HOME || "", ".claude", "settings.json"); const SMALL_EDIT_MODEL = "claude-haiku-4-5-20251001"; const MARKER_PATH = path.join(process.env.HOME || "", ".nexus", ".model-router-original"); +const RUNTIME_STATE_FILE = "model-router-runtime.json"; +const SMALL_PROMPT_MAX_TOKENS = Number(process.env.NEXUS_SMALL_PROMPT_MAX_TOKENS || 180); +const SMALL_PROMPT_MAX_WORDS = Number(process.env.NEXUS_SMALL_PROMPT_MAX_WORDS || 50); function detectRuntime(context) { const text = `${context?.session?.id || ""} ${process.env.CLAUDE_SESSION_ID || ""} ${process.env.CODEX_SESSION_ID || ""}`.toLowerCase(); @@ -33,11 +36,88 @@ function writeSettings(settings) { fs.writeFileSync(SETTINGS_PATH, JSON.stringify(settings, null, 4)); } -function getPromptWordCount(context) { +function getPromptMetrics(context) { const content = context?.messageHistory?.slice(-1)[0]?.content || ""; - return (typeof content === "string" ? content : JSON.stringify(content)) + const text = typeof content === "string" ? content : JSON.stringify(content); + const words = text .split(/\s+/) .filter(Boolean).length; + const chars = text.length; + const tokens = Math.ceil(chars / 4); + return { words, chars, tokens }; +} + +function getRuntimeStatePath(context) { + return path.join(resolveRepoRoot({ context }), ".nexus", RUNTIME_STATE_FILE); +} + +function writeRuntimeState(context, payload) { + const statePath = getRuntimeStatePath(context); + fs.mkdirSync(path.dirname(statePath), { recursive: true }); + fs.writeFileSync( + statePath, + JSON.stringify( + { + updatedAt: new Date().toISOString(), + ...payload, + }, + null, + 2 + ), + "utf8" + ); +} + +function readCurrentModel(runtime) { + if (runtime === "claude") { + return readSettings().model || process.env.CLAUDE_MODEL || MODEL_BY_RUNTIME.claude; + } + if (runtime === "codex") { + return process.env.CODEX_MODEL || MODEL_BY_RUNTIME.codex; + } + return "default"; +} + +function shouldHandlePromptRouting(toolName, trigger) { + const normalized = String(trigger || "").toLowerCase(); + if (normalized === "user-prompt-submit" || normalized === "userpromptsubmit") { + return true; + } + return ["Edit", "Write"].includes(toolName); +} + +function restoreLargeModel(settings) { + let desired = MODEL_BY_RUNTIME.claude; + try { + if (fs.existsSync(MARKER_PATH)) { + const fromMarker = fs.readFileSync(MARKER_PATH, "utf8").trim(); + if (fromMarker) { + desired = fromMarker; + } + fs.unlinkSync(MARKER_PATH); + } + } catch {} + + if (settings.model === SMALL_EDIT_MODEL && desired !== SMALL_EDIT_MODEL) { + settings.model = desired; + writeSettings(settings); + return { action: "restored-large", switchedTo: desired }; + } + return { action: "kept-large", switchedTo: settings.model || desired }; +} + +function switchToSmallModel(settings) { + const current = settings.model || process.env.CLAUDE_MODEL || MODEL_BY_RUNTIME.claude; + if (current && current !== SMALL_EDIT_MODEL) { + try { + fs.mkdirSync(path.dirname(MARKER_PATH), { recursive: true }); + fs.writeFileSync(MARKER_PATH, current); + } catch {} + settings.model = SMALL_EDIT_MODEL; + writeSettings(settings); + return { action: "switched-to-small", switchedTo: SMALL_EDIT_MODEL, previousModel: current }; + } + return { action: "kept-small", switchedTo: SMALL_EDIT_MODEL, previousModel: current }; } module.exports = async ({ toolName, context, hookEventName } = {}) => { @@ -51,41 +131,52 @@ module.exports = async ({ toolName, context, hookEventName } = {}) => { } const enabled = isModelRouterEnabled({ context }); - const configuredModel = runtime === "claude" ? readSettings().model || MODEL_BY_RUNTIME.claude : MODEL_BY_RUNTIME[runtime] || "default"; + const configuredModel = readCurrentModel(runtime); console.error(`[model-router] runtime=${runtime} trigger=${trigger} enabled=${enabled} model=${configuredModel}`); + const prompt = getPromptMetrics(context); + let routingAction = "observed"; + let switchedTo = configuredModel; + if (!enabled) { + writeRuntimeState(context, { + runtime, + trigger, + enabled, + currentModel: configuredModel, + prompt, + thresholds: { + smallPromptMaxWords: SMALL_PROMPT_MAX_WORDS, + smallPromptMaxTokens: SMALL_PROMPT_MAX_TOKENS, + }, + action: "disabled", + }); return { continue: true }; } - if (runtime === "claude" && ["Edit", "Write"].includes(toolName)) { - const wordCount = getPromptWordCount(context); - if (wordCount <= 50) { - const settings = readSettings(); - const current = settings.model; - if (current && current !== SMALL_EDIT_MODEL) { - try { - fs.mkdirSync(path.dirname(MARKER_PATH), { recursive: true }); - fs.writeFileSync(MARKER_PATH, current); - } catch {} - settings.model = SMALL_EDIT_MODEL; - writeSettings(settings); - console.error(`[model-router] Switched to ${SMALL_EDIT_MODEL} for small edit (was: ${current})`); - } + if (runtime === "claude" && shouldHandlePromptRouting(toolName, trigger)) { + const shouldUseSmall = prompt.tokens <= SMALL_PROMPT_MAX_TOKENS && prompt.words <= SMALL_PROMPT_MAX_WORDS; + const settings = readSettings(); + const result = shouldUseSmall ? switchToSmallModel(settings) : restoreLargeModel(settings); + routingAction = result.action; + switchedTo = result.switchedTo; + if (routingAction === "switched-to-small") { + console.error( + `[model-router] switched model ${result.previousModel} -> ${SMALL_EDIT_MODEL} (prompt words=${prompt.words} tokens~=${prompt.tokens})` + ); + } else if (routingAction === "restored-large") { + console.error( + `[model-router] restored model to ${result.switchedTo} (prompt words=${prompt.words} tokens~=${prompt.tokens})` + ); } else { - try { - if (fs.existsSync(MARKER_PATH)) { - const original = fs.readFileSync(MARKER_PATH, "utf8").trim(); - const settings = readSettings(); - if (settings.model === SMALL_EDIT_MODEL) { - settings.model = original; - writeSettings(settings); - fs.unlinkSync(MARKER_PATH); - console.error(`[model-router] Restored model to ${original}`); - } - } - } catch {} + console.error( + `[model-router] kept model ${result.switchedTo} (prompt words=${prompt.words} tokens~=${prompt.tokens})` + ); } + } else if (runtime === "codex" && shouldHandlePromptRouting(toolName, trigger)) { + routingAction = "codex-observe-only"; + switchedTo = readCurrentModel("codex"); + console.error("[model-router] codex runtime detected; model is observed and recorded, no dynamic switch is applied by this hook."); } if (toolName === "WebFetch" || toolName === "WebSearch") { @@ -102,5 +193,19 @@ module.exports = async ({ toolName, context, hookEventName } = {}) => { } } + writeRuntimeState(context, { + runtime, + trigger, + enabled, + currentModel: readCurrentModel(runtime), + prompt, + thresholds: { + smallPromptMaxWords: SMALL_PROMPT_MAX_WORDS, + smallPromptMaxTokens: SMALL_PROMPT_MAX_TOKENS, + }, + action: routingAction, + switchedTo, + }); + return { continue: true }; }; diff --git a/test/model-router-hook.test.js b/test/model-router-hook.test.js new file mode 100644 index 0000000..0a653d1 --- /dev/null +++ b/test/model-router-hook.test.js @@ -0,0 +1,154 @@ +"use strict"; + +const assert = require("node:assert/strict"); +const fs = require("node:fs"); +const os = require("node:os"); +const path = require("node:path"); +const test = require("node:test"); + +function makeTempDir(prefix) { + return fs.mkdtempSync(path.join(os.tmpdir(), prefix)); +} + +function loadRouter() { + const modulePath = require.resolve("../pre-tool-use/model-router"); + delete require.cache[modulePath]; + return require("../pre-tool-use/model-router"); +} + +function restoreEnv(name, value) { + if (typeof value === "undefined") { + delete process.env[name]; + return; + } + process.env[name] = value; +} + +async function withIsolatedEnv(fn) { + const originalHome = process.env.HOME; + const originalCwd = process.cwd(); + const originalClaudeModel = process.env.CLAUDE_MODEL; + const originalCodexModel = process.env.CODEX_MODEL; + const originalRuntime = process.env.NEXUS_AI_RUNTIME; + const originalSmallTokens = process.env.NEXUS_SMALL_PROMPT_MAX_TOKENS; + const originalSmallWords = process.env.NEXUS_SMALL_PROMPT_MAX_WORDS; + + try { + await fn(); + } finally { + process.chdir(originalCwd); + restoreEnv("HOME", originalHome); + restoreEnv("CLAUDE_MODEL", originalClaudeModel); + restoreEnv("CODEX_MODEL", originalCodexModel); + restoreEnv("NEXUS_AI_RUNTIME", originalRuntime); + restoreEnv("NEXUS_SMALL_PROMPT_MAX_TOKENS", originalSmallTokens); + restoreEnv("NEXUS_SMALL_PROMPT_MAX_WORDS", originalSmallWords); + } +} + +test("session-start writes runtime state file without CLI toggles", async () => { + await withIsolatedEnv(async () => { + const repoRoot = makeTempDir("nexus-router-repo-"); + const homeRoot = makeTempDir("nexus-router-home-"); + process.chdir(repoRoot); + process.env.HOME = homeRoot; + process.env.CLAUDE_MODEL = "claude-sonnet-4-6-20250514"; + + const claudeDir = path.join(homeRoot, ".claude"); + fs.mkdirSync(claudeDir, { recursive: true }); + fs.writeFileSync( + path.join(claudeDir, "settings.json"), + JSON.stringify({ model: "claude-sonnet-4-6-20250514" }, null, 2), + "utf8" + ); + + const router = loadRouter(); + await router({ + hookEventName: "session-start", + context: { cwd: repoRoot, session: { id: "claude-session-1" }, messageHistory: [] }, + }); + + const statePath = path.join(repoRoot, ".nexus", "model-router-runtime.json"); + assert.equal(fs.existsSync(statePath), true); + const state = JSON.parse(fs.readFileSync(statePath, "utf8")); + assert.equal(state.runtime, "claude"); + assert.equal(state.enabled, true); + assert.equal(state.action, "observed"); + }); +}); + +test("user prompt submit switches to small model and restores on larger prompt", async () => { + await withIsolatedEnv(async () => { + const repoRoot = makeTempDir("nexus-router-repo-"); + const homeRoot = makeTempDir("nexus-router-home-"); + process.chdir(repoRoot); + process.env.HOME = homeRoot; + process.env.NEXUS_SMALL_PROMPT_MAX_TOKENS = "120"; + process.env.NEXUS_SMALL_PROMPT_MAX_WORDS = "20"; + + const claudeDir = path.join(homeRoot, ".claude"); + fs.mkdirSync(claudeDir, { recursive: true }); + const settingsPath = path.join(claudeDir, "settings.json"); + fs.writeFileSync( + settingsPath, + JSON.stringify({ model: "claude-sonnet-4-6-20250514" }, null, 2), + "utf8" + ); + + const router = loadRouter(); + await router({ + hookEventName: "user-prompt-submit", + context: { + cwd: repoRoot, + session: { id: "claude-session-2" }, + messageHistory: [{ content: "Small ask: fix typo" }], + }, + }); + + let settings = JSON.parse(fs.readFileSync(settingsPath, "utf8")); + assert.equal(settings.model, "claude-haiku-4-5-20251001"); + + await router({ + hookEventName: "user-prompt-submit", + context: { + cwd: repoRoot, + session: { id: "claude-session-2" }, + messageHistory: [ + { + content: + "Please design and explain a production rollout plan with validation, rollback criteria, dependency analysis, and exhaustive acceptance checks for this feature set.", + }, + ], + }, + }); + + settings = JSON.parse(fs.readFileSync(settingsPath, "utf8")); + assert.equal(settings.model, "claude-sonnet-4-6-20250514"); + + const statePath = path.join(repoRoot, ".nexus", "model-router-runtime.json"); + const state = JSON.parse(fs.readFileSync(statePath, "utf8")); + assert.equal(state.action, "restored-large"); + assert.equal(state.currentModel, "claude-sonnet-4-6-20250514"); + }); +}); + +test("codex session-start writes observed model state", async () => { + await withIsolatedEnv(async () => { + const repoRoot = makeTempDir("nexus-router-repo-"); + const homeRoot = makeTempDir("nexus-router-home-"); + process.chdir(repoRoot); + process.env.HOME = homeRoot; + process.env.CODEX_MODEL = "gpt-5.5"; + + const router = loadRouter(); + await router({ + hookEventName: "session-start", + context: { cwd: repoRoot, session: { id: "codex-session-1" }, messageHistory: [] }, + }); + + const statePath = path.join(repoRoot, ".nexus", "model-router-runtime.json"); + const state = JSON.parse(fs.readFileSync(statePath, "utf8")); + assert.equal(state.runtime, "codex"); + assert.equal(state.currentModel, "gpt-5.5"); + }); +});