Merge pull request #430 from opticon454/custom-model-run-menu

src/web/server.ts

@@ -71,7 +71,7 @@ import {
 import { imageWatcher } from '../image-watcher.js';
 import { workflowRunWatcher, summarizeRun } from '../workflow-run-watcher.js';
 import { attachmentRegistry, buildFileThumbnailRoute, registerExternalAttachment } from '../attachment-registry.js';
-import { getCli } from '../config/cli-registry/registry.js';
+import { getCli, enabledClis } from '../config/cli-registry/registry.js';
 import { readCustomModelHosts } from '../custom-model-hosts.js';
 import { applyCustomModelInjection, customModelConfigDir, removeConfigDir } from '../custom-model-injection-apply.js';
 import type { CustomModelBookkeeping } from '../types/session.js';
@@ -195,6 +195,10 @@ import {
   registerWebviewRoutes,
   registerTabLayoutRoutes,
   registerCustomModelRoutes,
+  refreshAllCustomModelHosts,
+  readCustomModelEndpointsEnabled,
+  detectCustomModelSwapDisplacements,
+  pruneIdleLlamaSwapLogTails,
   tryWebviewRefererFallback,
 } from './routes/index.js';
 import { isLostWebviewFrameNavigation } from './webview-proxy.js';
@@ -207,11 +211,32 @@ const __dirname = dirname(fileURLToPath(import.meta.url));
 // while capping growth of `sseClientsById` and blocking pathological inputs.
 const SSE_CLIENT_ID_RE = /^[A-Za-z0-9_-]{8,64}$/;
 const CODEX_USAGE_POLL_INTERVAL_MS = 5 * 60_000;
+const CUSTOM_MODEL_REDISCOVER_INTERVAL_MS = 5 * 60_000;
+// Much shorter than the model-LIST refresh above on purpose: this catches an actual
+// eviction (a session's model no longer loaded, silently swapped out by another
+// session's use), which the user wants to know about promptly, not once every 5
+// minutes. Cheap either way — one /running GET per distinct endpoint with at least
+// one live custom-model session, not per session.
+const CUSTOM_MODEL_SWAP_CHECK_INTERVAL_MS = 20_000;
 
 function escapeHtmlText(value: string): string {
   return value.replaceAll('&', '&amp;').replaceAll('<', '&lt;').replaceAll('>', '&gt;');
 }
 
+/**
+ * Escapes a JSON string for safe embedding as the body of an inline `<script>`
+ * tag: `<` becomes the six-character sequence `<`, which both a JSON
+ * parser and a plain JS string literal decode back to `<` (both treat
+ * `\uXXXX` identically), but which can never itself form the two literal
+ * characters `<` `/` a browser's HTML tokenizer looks for to end the tag. A
+ * value containing a literal `</script>` would otherwise close the tag early
+ * and turn the rest of the document into inert script-body text. Exported so
+ * it unit-tests without constructing a WebServer (which needs a real tmux).
+ */
+export function escapeScriptJson(json: string): string {
+  return json.replace(/</g, '\\u003c');
+}
+
 import {
   SESSIONS_LIST_CACHE_TTL,
   SCHEDULED_CLEANUP_INTERVAL,
@@ -270,6 +295,8 @@ export class WebServer extends EventEmitter {
   // Store session listener references for explicit cleanup (prevents memory leaks)
   private sessionListenerRefs: Map<string, SessionListenerRefs> = new Map();
   private scheduledRuns: Map<string, ScheduledRun> = new Map();
+  /** De-dupe state for the swap-displacement sweep — see detectCustomModelSwapDisplacements. */
+  private _customModelDisplacedNotified: Set<string> = new Set();
   /** Cron service (assigned in setupRoutes). */
   private cronService!: CronService;
   /**
@@ -1224,7 +1251,13 @@ export class WebServer extends EventEmitter {
       return undefined;
     }
     try {
-      return applyCustomModelInjection(entry, endpoint, saved.modelId, session.id)?.envOverrides;
+      return applyCustomModelInjection(
+        entry,
+        endpoint,
+        saved.modelId,
+        session.id,
+        endpoint.modelContextLengths?.[saved.modelId]
+      )?.envOverrides;
     } catch (err) {
       console.warn('[WebServer] Failed to rebuild custom-model env on recovery:', err);
       return undefined;
@@ -1319,6 +1352,10 @@ export class WebServer extends EventEmitter {
       session.ralphTracker.stopWatchingFixPlan();
     }
 
+    // Custom Model Endpoint Profiles: drop this session's swap-displacement notify flag
+    // (see _checkCustomModelSwapDisplacements below) so it can't linger in that Set forever.
+    this._customModelDisplacedNotified.delete(sessionId);
+
     // Kill all subagents spawned by this session (scoped to sessionId to avoid cross-session kills)
     if (session && killMux) {
       try {
@@ -1619,6 +1656,23 @@ export class WebServer extends EventEmitter {
         '</head>',
         `<script>window.__codemanCliAvailable=${JSON.stringify(available)};</script>\n</head>`
       );
+      // Which run modes the Run-menu picker (docs/custom-model-endpoints-plan.md) may
+      // generate an entry for: read generically off the registry's `capabilities`
+      // (never an id list here) so a CLI whose customModelInjection lands later shows
+      // up in the picker with no frontend change, and one that ships `unsupported`
+      // (antigravity, and `shell`'s `kind !== 'agent'`) never does.
+      const customModelClis = enabledClis()
+        .filter((entry) => entry.kind === 'agent' && entry.capabilities.customModelInjection.kind !== 'unsupported')
+        .map((entry) => ({ id: entry.id, label: entry.label }));
+      // Unlike the boolean-only __codemanCliAvailable above, this payload carries
+      // `label`, a string a user's own clis.json can set (CliEntry.label, up to 60
+      // chars) — see escapeScriptJson's own doc comment for why that needs escaping
+      // and __codemanCliAvailable's booleans never did.
+      const customModelClisJson = escapeScriptJson(JSON.stringify(customModelClis));
+      html = html.replace(
+        '</head>',
+        `<script>window.__codemanCustomModelClis=${customModelClisJson};</script>\n</head>`
+      );
     }
     if (!soloSessionId && process.env.CODEMAN_GESTURE === '1') {
       html = html.replace('</head>', `<script>window.__codemanGestureAvailable=true;</script>\n</head>`);
@@ -2350,6 +2404,7 @@ export class WebServer extends EventEmitter {
       'team:',
       'case:',
       'remote:',
+      'custom-model:',
     ];
     if (SESSION_PREFIXES.some((p) => event.startsWith(p))) {
       const d = (data ?? {}) as { sessionId?: string; id?: string; session?: { id?: string }; username?: string };
@@ -2734,6 +2789,64 @@ export class WebServer extends EventEmitter {
       });
     }
 
+    // Custom Model Endpoint Profiles (docs/custom-model-endpoints-plan.md): keeps
+    // each saved endpoint's discovered model list current with no manual
+    // "Discover" click, so a model added on the server side (or one that drops
+    // off) shows up in the Run-menu picker within one cycle. Best-effort per
+    // endpoint (refreshAllCustomModelHosts skips one that's unreachable rather
+    // than failing the sweep) and off in tests for the same reason the Codex
+    // poll above is — no real network to hit, no server instance to keep alive.
+    if (!this.testMode) {
+      this.cleanup.setInterval(
+        () => {
+          // Reads the setting fresh on every tick, same reasoning as
+          // readPlanUsageTelemetryEnabled() beside it: a live toggle takes effect
+          // on the very next cycle, not just at server boot, and turning the
+          // feature off actually stops the polling instead of only hiding the UI.
+          void readCustomModelEndpointsEnabled()
+            .then((enabled) => {
+              if (!enabled) return;
+              return refreshAllCustomModelHosts();
+            })
+            .catch((err) => {
+              console.error('[custom-model] periodic re-discovery failed:', getErrorMessage(err));
+            });
+        },
+        CUSTOM_MODEL_REDISCOVER_INTERVAL_MS,
+        { description: 'custom model endpoint re-discovery' }
+      );
+    }
+
+    // Custom Model Endpoint Profiles: the swap-conflict check on the apply/create routes
+    // only ever runs at THAT session's own launch/apply moment — it cannot catch a LATER
+    // eviction triggered by a different session's normal use, since llama-swap has no push
+    // notification of its own and only swaps in response to a real inference request
+    // (confirmed live: a session created while nothing else conflicted at that instant can
+    // still get silently displaced afterward). This periodic sweep is what catches that
+    // case after the fact and tells the displaced session's user, rather than leaving them
+    // to discover it only when their next prompt behaves unexpectedly.
+    if (!this.testMode) {
+      this.cleanup.setInterval(
+        () => {
+          detectCustomModelSwapDisplacements(this.sessions.values(), this._customModelDisplacedNotified)
+            .then((displacements) => {
+              for (const displacement of displacements) {
+                this.broadcast(SseEvent.CustomModelSwappedOut, displacement);
+              }
+            })
+            .catch((err) => {
+              console.error('[custom-model] swap-displacement check failed:', getErrorMessage(err));
+            });
+          // Same cadence, unrelated concern: close any /logs tail (see
+          // getLatestLlamaSwapLogLine) nothing has polled in a while, so a loading banner
+          // that finished (or was abandoned) doesn't leave a connection open forever.
+          pruneIdleLlamaSwapLogTails();
+        },
+        CUSTOM_MODEL_SWAP_CHECK_INTERVAL_MS,
+        { description: 'custom model swap-displacement check' }
+      );
+    }
+
     // Start scheduled runs cleanup timer
     this.cleanup.setInterval(
       () => {