diff --git a/src/custom-model-injection-apply.ts b/src/custom-model-injection-apply.ts index 51ca6fbc..43cfd5ad 100644 --- a/src/custom-model-injection-apply.ts +++ b/src/custom-model-injection-apply.ts @@ -93,6 +93,15 @@ function linkSharedProjectsDir(isolatedDir: string): void { * keys), and a corrupt or partially-written file (a crash mid-write) is treated as absent * rather than failing the whole apply over a nice-to-have. */ +/** + * The form Claude Code actually stores an approved key in: the trimmed last 20 + * characters. Mirrors the CLI's own `e.trim().slice(-20)`, which is applied on BOTH + * the write and the lookup, so anything else never matches. + */ +export function truncateApiKeyForTrustFile(apiKey: string): string { + return apiKey.trim().slice(-20); +} + function seedApiKeyTrustFile( configDir: string, trustFile: { relPath: string; shape: 'claude-api-key-responses' }, @@ -107,7 +116,18 @@ function seedApiKeyTrustFile( } const responses = (existing.customApiKeyResponses ?? {}) as { approved?: unknown; rejected?: unknown }; const approved = new Set(Array.isArray(responses.approved) ? (responses.approved as string[]) : []); - approved.add(apiKey); + // ⚠ Claude Code stores and compares only the LAST 20 CHARACTERS of a key, never the + // whole thing: its lookup is `approved.includes(key.trim().slice(-20))` (decompiled + // from the 2.1.278 bundle, and corroborated by real `~/.claude.json` files, whose + // customApiKeyResponses entries are all exactly 20 characters). Seeding the full key + // therefore never matches for a REAL key, and claude stops at the interactive + // "Detected a custom API key in your environment" prompt, whose default is + // "No (recommended)" — so the launch hangs or silently refuses the key this feature + // just injected. It went unnoticed because a keyless llama.cpp/llama-swap endpoint + // uses DEFAULT_API_KEY ('local-dummy-key', 15 chars), where slice(-20) is the whole + // string and the seed matches by accident. Truncating here also keeps a full + // third-party credential from being written into a second file on disk. + approved.add(truncateApiKeyForTrustFile(apiKey)); const rejected = Array.isArray(responses.rejected) ? responses.rejected : []; existing.customApiKeyResponses = { approved: [...approved], rejected }; try { diff --git a/src/web/public/session-ui.js b/src/web/public/session-ui.js index 90fe768c..07d3e1eb 100644 --- a/src/web/public/session-ui.js +++ b/src/web/public/session-ui.js @@ -812,7 +812,7 @@ Object.assign(CodemanApp.prototype, { * POSTs a /api/quick-start body already carrying `customModel` (see the run() * call sites below), showing the same llama-swap "this will unload it for session X" * warning the restart path's `_applyCustomModelToSession` shows when the route asks - * for confirmation, and retrying with `confirmed: true` on accept. Stashes the final + * for confirmation, and retrying with that question's own flag on accept. Stashes the final * response's payload on `_lastCustomModelLaunchResult` for * `_runCustomModelEntryOneShot` to read `modelSwapInProgress` off afterward — run()'s * eleven per-mode dispatch targets have no shared return-value contract of their own, @@ -827,6 +827,12 @@ Object.assign(CodemanApp.prototype, { }); return res.json(); }; + // Each question is answered with its OWN flag, and the answer accumulates, so the + // second POST still carries the first answer. Never the blanket `confirmed`: the two + // questions are about different people (a context window too small is the caller's + // problem, unloading a model is another session's), and while they shared one flag + // clicking past the context warning silently answered the swap question too. + let answered = {}; let data = await post(bodyObj); if (data?.data?.requiresContextWarning) { const { modelId, contextLength, minSafeContextTokens } = data.data; @@ -835,21 +841,27 @@ Object.assign(CodemanApp.prototype, { this._lastCustomModelLaunchResult = undefined; return { success: false, error: 'Launch cancelled — context window too small' }; } - data = await post({ ...bodyObj, customModel: { ...bodyObj.customModel, confirmed: true } }); + answered = { ...answered, confirmedContext: true }; + data = await post({ ...bodyObj, customModel: { ...bodyObj.customModel, ...answered } }); } if (data?.data?.requiresConfirmation) { const { currentlyLoadedModel, affectedSessions } = data.data; + // The swap is blocked regardless of ownership, but multi-user mode scopes which + // sessions get NAMED, so this list can be empty while the conflict is real. const names = affectedSessions.map((s) => s.name || s.id).join(', '); + const who = names + ? `${names} ${affectedSessions.length === 1 ? 'is' : 'are'} currently using` + : 'Another session on this endpoint is currently using'; + const them = names && affectedSessions.length > 1 ? 'those sessions' : 'that session'; const proceed = await this._confirmModelSwap( - `${names} ${affectedSessions.length === 1 ? 'is' : 'are'} currently using ` + - `${currentlyLoadedModel} on this endpoint. Switching will unload it for ` + - `${affectedSessions.length === 1 ? 'that session' : 'those sessions'} too. Continue?` + `${who} ${currentlyLoadedModel} on this endpoint. Switching will unload it for ` + `${them} too. Continue?` ); if (!proceed) { this._lastCustomModelLaunchResult = undefined; return { success: false, error: 'Model switch cancelled' }; } - data = await post({ ...bodyObj, customModel: { ...bodyObj.customModel, confirmed: true } }); + answered = { ...answered, confirmedSwap: true }; + data = await post({ ...bodyObj, customModel: { ...bodyObj.customModel, ...answered } }); } this._lastCustomModelLaunchResult = data?.success !== false ? data?.data : undefined; return data; @@ -919,11 +931,15 @@ Object.assign(CodemanApp.prototype, { // once `data.success !== false`. let payload = data?.success !== false ? data?.data : undefined; + // Accumulates the questions the user has answered, so a second retry still carries + // the first answer. See _applyCustomModelToSession for why these are per-question. + let answered = {}; + // This CLI's own fixed overhead (system prompt + tool schemas) may exceed the // model's real discovered context outright — no context-length declaration can // fix that, since compaction only trims conversation history and there is none // on message 1. Warn and let the user decide whether to launch anyway, same - // confirmed:true re-send pattern as the swap check below. + // re-send pattern as the swap check below. if (ok && payload?.requiresContextWarning) { const proceed = await this._confirmContextWarning( payload.modelId, @@ -935,27 +951,34 @@ Object.assign(CodemanApp.prototype, { this.showToast('Kept the native backend — context window too small', 'info'); return; } - ({ ok, data, res } = await this._applyCustomModelToSession(sessionId, endpointId, modelId, true)); + answered = { ...answered, confirmedContext: true }; + ({ ok, data, res } = await this._applyCustomModelToSession(sessionId, endpointId, modelId, answered)); payload = data?.success !== false ? data?.data : undefined; } // llama-swap runs one model at a time: switching would unload it out from under // another session actively using it. The route only asks when that's actually true - // (never just because a swap is needed at all) — confirming re-sends the exact same - // call with `confirmed: true` so the route skips the check the second time. + // (never just because a swap is needed at all) — confirming re-sends the same call + // with `confirmedSwap` so the route skips THIS check, and only this one, next time. if (ok && payload?.requiresConfirmation) { + // See the one-shot path above: an empty list means the conflicting sessions are + // ones this caller may not be told about, not that there is no conflict. const names = payload.affectedSessions.map((s) => s.name || s.id).join(', '); + const who = names + ? `${names} ${payload.affectedSessions.length === 1 ? 'is' : 'are'} currently using` + : 'Another session on this endpoint is currently using'; + const them = names && payload.affectedSessions.length > 1 ? 'those sessions' : 'that session'; const proceed = await this._confirmModelSwap( - `${names} ${payload.affectedSessions.length === 1 ? 'is' : 'are'} currently using ` + - `${payload.currentlyLoadedModel} on this endpoint. Switching to ${modelId} will unload it ` + - `for ${payload.affectedSessions.length === 1 ? 'that session' : 'those sessions'} too. Continue?` + `${who} ${payload.currentlyLoadedModel} on this endpoint. Switching to ${modelId} will unload it ` + + `for ${them} too. Continue?` ); if (!proceed) { switchingToast?.dismiss(); this.showToast('Kept the native backend — model switch cancelled', 'info'); return; } - ({ ok, data, res } = await this._applyCustomModelToSession(sessionId, endpointId, modelId, true)); + answered = { ...answered, confirmedSwap: true }; + ({ ok, data, res } = await this._applyCustomModelToSession(sessionId, endpointId, modelId, answered)); payload = data?.success !== false ? data?.data : undefined; } @@ -987,10 +1010,16 @@ Object.assign(CodemanApp.prototype, { /** POST /api/sessions/:id/custom-model, returning {ok, data, res} rather than throwing — * see runCustomModelEntry's own comment for why this goes through `_api()` (raw fetch) * rather than `_apiJson()`: a failure's `error` detail must survive to the caller. */ - async _applyCustomModelToSession(sessionId, endpointId, modelId, confirmed) { + /** + * `answered` carries the questions the user has ALREADY said yes to, as the route's own + * per-question flags (`confirmedContext`, `confirmedSwap`). Never the blanket + * `confirmed`: the two questions are about different people, so answering one must not + * answer the other. It accumulates, so the second retry still carries the first answer. + */ + async _applyCustomModelToSession(sessionId, endpointId, modelId, answered) { const res = await this._api(`/api/sessions/${sessionId}/custom-model`, { method: 'POST', - body: confirmed ? { endpointId, modelId, confirmed } : { endpointId, modelId }, + body: { endpointId, modelId, ...(answered || {}) }, }); const data = res ? await res.json().catch(() => null) : null; return { ok: !!res, data, res }; @@ -1874,7 +1903,7 @@ Object.assign(CodemanApp.prototype, { const body = buildBody(sessionName); const data = await this._quickStartWithCustomModelConfirm( customModelAnswered && body.customModel - ? { ...body, customModel: { ...body.customModel, confirmed: true } } + ? { ...body, customModel: { ...body.customModel, confirmedContext: true, confirmedSwap: true } } : body ); if (!data.success) throw new Error(data.error || `Failed to start ${label}`); diff --git a/src/web/routes/custom-model-routes.ts b/src/web/routes/custom-model-routes.ts index fc86d883..7f96a7c6 100644 --- a/src/web/routes/custom-model-routes.ts +++ b/src/web/routes/custom-model-routes.ts @@ -507,6 +507,18 @@ export function getLatestLlamaSwapLogLine( return entry.latestLine; } +/** + * Closes EVERY open log tail. The idle sweep above only runs on server.ts's periodic + * interval, and that interval is disposed on shutdown, so without this an outbound + * stream outlives `WebServer.stop()` against CLAUDE.md's "clear Maps in stop()" rule. + * Harmless today only because `cli.ts`'s shutdown handler reaches `process.exit(0)`, + * which is not a property to rely on: tests and any in-process restart do not. + */ +export function closeAllLlamaSwapLogTails(): void { + for (const entry of llamaSwapLogTails.values()) entry.controller.abort(); + llamaSwapLogTails.clear(); +} + /** * Closes any log tail nothing has called `getLatestLlamaSwapLogLine` about in * `LOG_TAIL_IDLE_MS` — a stream nobody is polling is an open connection with nothing to diff --git a/src/web/routes/index.ts b/src/web/routes/index.ts index c022e6b3..8a3f166f 100644 --- a/src/web/routes/index.ts +++ b/src/web/routes/index.ts @@ -32,6 +32,7 @@ export { registerCustomModelRoutes, refreshAllCustomModelHosts, readCustomModelEndpointsEnabled, + closeAllLlamaSwapLogTails, detectCustomModelSwapDisplacements, pruneIdleLlamaSwapLogTails, type CustomModelSessionLike, diff --git a/src/web/routes/session-routes.ts b/src/web/routes/session-routes.ts index f929ee2b..729c8983 100644 --- a/src/web/routes/session-routes.ts +++ b/src/web/routes/session-routes.ts @@ -1245,9 +1245,11 @@ export function registerSessionRoutes( // no matter what CLAUDE_CODE_MAX_CONTEXT_TOKENS says — confirmed live at ~36.4K tokens // against a model configured with a real 16384-token context. Warn before committing // to a restart that's certain to fail, rather than letting the user discover it via a - // cryptic 400 from the CLI itself. `confirmed` (already used for the swap-conflict - // warning below) skips this too — the user has already said "launch anyway" once. - if (!body.confirmed && exceedsSafeContextFloor(entry, contextLength)) { + // cryptic 400 from the CLI itself. Answered by `confirmedContext` (or the legacy + // `confirmed`, which still means both) — NOT by `confirmedSwap`: this warning is + // about the caller's own session, and the swap warning below is about someone + // else's, so an answer to one is not consent to the other. + if (!(body.confirmed || body.confirmedContext) && exceedsSafeContextFloor(entry, contextLength)) { return { requiresContextWarning: true, modelId: body.modelId, @@ -1275,9 +1277,12 @@ export function registerSessionRoutes( const targetReady = swapStatus.running.some((r) => r.model === body.modelId && r.state === 'ready'); // Only ask when switching would actually take the model away from another session - // that is currently using it — never just because a swap is needed at all. `confirmed` - // (set by the caller after showing that warning once) skips asking again. - if (swapNeeded && !body.confirmed) { + // that is currently using it — never just because a swap is needed at all. Answered + // by `confirmedSwap` (or the legacy `confirmed`). ⚠ It must NOT read + // `confirmedContext`: this check runs second, and while the two shared one flag a + // user who clicked past a too-small-context warning had already, silently, agreed to + // evict another session's model. + if (swapNeeded && !(body.confirmed || body.confirmedSwap)) { const conflicting = [...ctx.sessions.values()].filter( (s) => s.id !== session.id && s.customModel?.endpointId === endpoint.id && s.customModel?.modelId === currentlyLoaded @@ -3793,7 +3798,12 @@ export function registerSessionRoutes( // overhead can exceed a small enough real context on the very first message, // regardless of contextLengthVar. Warn before creating a session that's certain to // fail immediately. - if (!customModel.confirmed && exceedsSafeContextFloor(cmEntry, cmContextLength)) { + // See the dedicated route above for why this reads `confirmedContext` and never + // `confirmedSwap`. + if ( + !(customModel.confirmed || customModel.confirmedContext) && + exceedsSafeContextFloor(cmEntry, cmContextLength) + ) { return { requiresContextWarning: true, modelId: customModel.modelId, @@ -3816,7 +3826,7 @@ export function registerSessionRoutes( // is loaded at all yet. Drives the actual load trigger below. const cmTargetReady = cmSwapStatus.running.some((r) => r.model === customModel.modelId && r.state === 'ready'); qsCustomModelSwapInProgress = cmSwapStatus.isLlamaSwap && !cmTargetReady; - if (cmSwapNeeded && !customModel.confirmed) { + if (cmSwapNeeded && !(customModel.confirmed || customModel.confirmedSwap)) { const cmConflicting = [...ctx.sessions.values()].filter( (s) => s.customModel?.endpointId === cmEndpoint.id && s.customModel?.modelId === cmCurrentlyLoaded ); diff --git a/src/web/schemas.ts b/src/web/schemas.ts index f387a091..f1b1fc22 100644 --- a/src/web/schemas.ts +++ b/src/web/schemas.ts @@ -1071,15 +1071,17 @@ export const QuickStartSchema = z.object({ * model` does — never trusting raw env values from the client. One-shot, launch-time * equivalent of that route: no restart, so no visible relaunch (that route's restart-in- * place is still what an ALREADY-RUNNING session uses to switch later). Rejected for - * remote/docker cases, same reasoning as `envOverrides` above. `confirmed` mirrors that - * route's field: skips the llama-swap "this will unload it for another session" check on - * a deliberate retry. + * remote/docker cases, same reasoning as `envOverrides` above. The three confirmation + * flags mirror that route's fields; see `SessionCustomModelSchema` for why there are + * two specific ones rather than the single legacy `confirmed`. */ customModel: z .object({ endpointId: z.string().regex(/^[a-zA-Z0-9_-]+$/, 'Invalid endpoint id'), modelId: z.string().min(1).max(200), confirmed: z.boolean().optional(), + confirmedContext: z.boolean().optional(), + confirmedSwap: z.boolean().optional(), }) .strict() .optional(), @@ -2000,10 +2002,22 @@ export const CustomModelSelectionSchema = z.union([ z.object({ endpointId: z.string().regex(/^[a-zA-Z0-9_-]+$/, 'Invalid endpoint id'), modelId: z.string().min(1).max(200), - // Set once the caller has already shown the "this will unload for session(s) - // X" warning (see session-routes.ts's llama-swap conflict check) and the user chose to - // proceed anyway — skips that check on this call instead of asking again. + /** + * Two DIFFERENT questions can block a launch, and answering one is not consent to + * the other: `confirmedContext` answers "this model's context window is below the + * floor for this CLI", which affects only the caller, while `confirmedSwap` answers + * "loading this will unload the model another session is using", which affects + * someone else. They were one flag until the context check (which runs first) + * silently spent the swap answer too, so a user clicking "launch anyway" past a + * too-small context evicted another session's model without ever being asked. + * + * `confirmed` is the original single flag and still means BOTH, because it shipped + * in the HTTP-API-only cut of this feature and an existing caller must keep working. + * New callers should send the specific one they actually asked about. + */ confirmed: z.boolean().optional(), + confirmedContext: z.boolean().optional(), + confirmedSwap: z.boolean().optional(), }), z.object({ clear: z.literal(true) }), ]); diff --git a/src/web/server.ts b/src/web/server.ts index 40161727..f2f9b419 100644 --- a/src/web/server.ts +++ b/src/web/server.ts @@ -197,6 +197,7 @@ import { registerCustomModelRoutes, refreshAllCustomModelHosts, readCustomModelEndpointsEnabled, + closeAllLlamaSwapLogTails, detectCustomModelSwapDisplacements, pruneIdleLlamaSwapLogTails, tryWebviewRefererFallback, @@ -3719,6 +3720,10 @@ export class WebServer extends EventEmitter { // got wrong once. void stopDeepSeekWeb(); + // Same teardown rule: the per-endpoint llama-swap log tails are otherwise closed + // only by the periodic idle sweep, whose interval is disposed just below. + closeAllLlamaSwapLogTails(); + // Dispose all managed timers (intervals + resettable timeouts) this.cleanup.dispose(); diff --git a/test/custom-model-injection-apply.test.ts b/test/custom-model-injection-apply.test.ts index dfc8e793..9f870100 100644 --- a/test/custom-model-injection-apply.test.ts +++ b/test/custom-model-injection-apply.test.ts @@ -123,6 +123,34 @@ describe('applyCustomModelInjection: apiKeyTrustFile (pre-approves the injected expect(written.customApiKeyResponses.rejected).toEqual([]); }); + // ⚠ Claude Code stores and looks up only the LAST 20 CHARACTERS of a key + // (`key.trim().slice(-20)`, applied on both write and read), so seeding the whole key + // never matches for a REAL one and the launch stops at the interactive "Detected a + // custom API key" prompt whose default is "No (recommended)". Every other test here + // uses a key shorter than 20 characters, where slice(-20) is the whole string and the + // bug is invisible, which is exactly how it survived review. + it('claude: seeds a REAL-length key in the truncated form the CLI actually matches on', () => { + const sessionId = 'sess-trust-long'; + sessionsToClean.push(sessionId); + const longKey = 'sk-or-v1-0123456789abcdef0123456789abcdef0123456789abcdef'; + expect(longKey.length).toBeGreaterThan(20); + + const applied = applyCustomModelInjection( + entryOrThrow('claude'), + { ...endpoint, apiKey: longKey }, + 'qwen3', + sessionId + ); + const written = JSON.parse(readFileSync(join(applied!.configDir!, '.claude.json'), 'utf8')) as { + customApiKeyResponses: { approved: string[] }; + }; + + expect(written.customApiKeyResponses.approved).toEqual(['cdef0123456789abcdef']); + expect(written.customApiKeyResponses.approved[0]).toHaveLength(20); + // and the full credential is not written into this second file at all + expect(readFileSync(join(applied!.configDir!, '.claude.json'), 'utf8')).not.toContain(longKey); + }); + it('claude: falls back to the dummy key when the endpoint has none, and still seeds it', () => { const sessionId = 'sess-trust-2'; sessionsToClean.push(sessionId); diff --git a/test/routes/session-custom-model.test.ts b/test/routes/session-custom-model.test.ts index 32cbc7e7..6dd25c6f 100644 --- a/test/routes/session-custom-model.test.ts +++ b/test/routes/session-custom-model.test.ts @@ -463,6 +463,99 @@ describe('POST /api/sessions/:id/custom-model', () => { expect(session.restartCli).toHaveBeenCalledTimes(1); }); + // The two questions are about DIFFERENT people: a context window below the floor is + // the caller's own problem, while unloading a model takes it away from someone else's + // session. They shared one `confirmed` flag until this release, and because the + // context check runs first, clicking "launch anyway" past the context warning silently + // answered the swap question too and evicted another session's model unasked. + describe('answering one question is not consent to the other', () => { + async function bothConditions() { + const { app, ctx } = await setup(); + await writeCustomModelHosts(getDataDir(), [CLAUDE_ENDPOINT, SMALL_CTX_ENDPOINT]); + const session = ctx.sessions.get('test-session-1')!; + session.mode = 'claude'; + // another session is actively on the model this endpoint currently has loaded + const other = createMockSession('other-session'); + other.name = 'w2-otherbox'; + other.customModel = { endpointId: 'ep-small', modelId: 'llama3' }; + ctx.sessions.set('other-session', other); + mockRunning([{ model: 'llama3', state: 'ready' }]); + return { app, session }; + } + + it('still asks about the swap after the context warning was confirmed', async () => { + const { app, session } = await bothConditions(); + + const res = await app.inject({ + method: 'POST', + url: '/api/sessions/test-session-1/custom-model', + payload: { + endpointId: 'ep-small', + modelId: 'qwen3.8-27b-ud-q4_k_xl', + confirmedContext: true, + }, + }); + + const body = res.json(); + expect(body.requiresContextWarning).toBeUndefined(); + expect(body.requiresConfirmation).toBe(true); + expect(body.currentlyLoadedModel).toBe('llama3'); + // and crucially nothing was applied: the other session keeps its model + expect(session.setCustomModel).not.toHaveBeenCalled(); + expect(session.restartCli).not.toHaveBeenCalled(); + }); + + it('applies once BOTH questions are answered', async () => { + const { app, session } = await bothConditions(); + + const res = await app.inject({ + method: 'POST', + url: '/api/sessions/test-session-1/custom-model', + payload: { + endpointId: 'ep-small', + modelId: 'qwen3.8-27b-ud-q4_k_xl', + confirmedContext: true, + confirmedSwap: true, + }, + }); + + const body = res.json(); + expect(body.requiresContextWarning).toBeUndefined(); + expect(body.requiresConfirmation).toBeUndefined(); + expect(session.setCustomModel).toHaveBeenCalledTimes(1); + }); + + it('confirmedSwap alone does not silence the context warning either', async () => { + const { app, session } = await bothConditions(); + + const res = await app.inject({ + method: 'POST', + url: '/api/sessions/test-session-1/custom-model', + payload: { endpointId: 'ep-small', modelId: 'qwen3.8-27b-ud-q4_k_xl', confirmedSwap: true }, + }); + + expect(res.json().requiresContextWarning).toBe(true); + expect(session.setCustomModel).not.toHaveBeenCalled(); + }); + + // `confirmed` shipped in the HTTP-API-only cut of this feature, so a caller written + // against that must keep working: it means both, exactly as it used to. + it('keeps the legacy blanket `confirmed` meaning both', async () => { + const { app, session } = await bothConditions(); + + const res = await app.inject({ + method: 'POST', + url: '/api/sessions/test-session-1/custom-model', + payload: { endpointId: 'ep-small', modelId: 'qwen3.8-27b-ud-q4_k_xl', confirmed: true }, + }); + + const body = res.json(); + expect(body.requiresContextWarning).toBeUndefined(); + expect(body.requiresConfirmation).toBeUndefined(); + expect(session.setCustomModel).toHaveBeenCalledTimes(1); + }); + }); + it('does not warn when the discovered context is comfortably above the floor', async () => { const { app, ctx } = await setup(); const roomyEndpoint: CustomModelHost = {