From 8f6fa744e3218c96ceef9537a39a14cc566db27e Mon Sep 17 00:00:00 2001 From: CodeReclaimers Date: Tue, 17 Feb 2026 12:56:00 -0500 Subject: [PATCH] Fix: preserve modelOverride in agent handler (#5369) Move session entry construction inside `updateSessionStore` mutator to fix race condition where sub-agent model overrides were silently ignored (~3% failure rate). The agent handler could read stale cached data and overwrite the fresh `modelOverride` when writing back to the session store. Building the entry inside the mutator ensures it always uses fresh store data (loaded with skipCache: true). Also moves the send-policy check inside the mutator for the same freshness guarantee, and integrates legacy key pruning via resolveGatewaySessionStoreTarget/pruneLegacyStoreKeys. Fixes #5369 Co-Authored-By: Claude Opus 4.6 --- src/gateway/server-methods/agent.test.ts | 314 +++++++++++++++++++++-- src/gateway/server-methods/agent.ts | 154 +++++++---- 2 files changed, 398 insertions(+), 70 deletions(-) diff --git a/src/gateway/server-methods/agent.test.ts b/src/gateway/server-methods/agent.test.ts index 06613d9e180..333a3990263 100644 --- a/src/gateway/server-methods/agent.test.ts +++ b/src/gateway/server-methods/agent.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, vi } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { BARE_SESSION_RESET_PROMPT } from "../../auto-reply/reply/session-reset-prompt.js"; import { agentHandlers } from "./agent.js"; import type { GatewayRequestContext } from "./types.js"; @@ -107,12 +107,13 @@ function mockMainSessionEntry(entry: Record, cfg: Record) { let capturedEntry: Record | undefined; mocks.updateSessionStore.mockImplementation(async (_path, updater) => { - const store: Record = {}; - await updater(store); + const store: Record = freshEntry ? { "agent:main:main": freshEntry } : {}; + const result = await updater(store); capturedEntry = store["agent:main:main"] as Record; + return result; }); return () => capturedEntry; } @@ -163,11 +164,9 @@ async function expectResetCall(expectedMessage: string) { } function primeMainAgentRun(params?: { sessionId?: string; cfg?: Record }) { - mockMainSessionEntry( - { sessionId: params?.sessionId ?? "existing-session-id" }, - params?.cfg ?? {}, - ); - mocks.updateSessionStore.mockResolvedValue(undefined); + const sessionId = params?.sessionId ?? "existing-session-id"; + mockMainSessionEntry({ sessionId }, params?.cfg ?? {}); + captureUpdatedMainEntry(buildExistingMainStoreEntry({ sessionId })); mocks.agentCommand.mockResolvedValue({ payloads: [{ text: "ok" }], meta: { durationMs: 100 }, @@ -267,6 +266,11 @@ async function invokeAgentIdentityGet( } describe("gateway agent handler", () => { + // Ensure fake timers are always restored, even if a test fails mid-execution + afterEach(() => { + vi.useRealTimers(); + }); + it("preserves ACP metadata from the current stored session entry", async () => { const existingAcpMeta = { backend: "acpx", @@ -281,15 +285,9 @@ describe("gateway agent handler", () => { acp: existingAcpMeta, }); - let capturedEntry: Record | undefined; - mocks.updateSessionStore.mockImplementation(async (_path, updater) => { - const store: Record = { - "agent:main:main": buildExistingMainStoreEntry({ acp: existingAcpMeta }), - }; - const result = await updater(store); - capturedEntry = store["agent:main:main"] as Record; - return result; - }); + const getCapturedEntry = captureUpdatedMainEntry( + buildExistingMainStoreEntry({ acp: existingAcpMeta }), + ); mocks.agentCommand.mockResolvedValue({ payloads: [{ text: "ok" }], @@ -299,6 +297,7 @@ describe("gateway agent handler", () => { await runMainAgent("test", "test-idem-acp-meta"); expect(mocks.updateSessionStore).toHaveBeenCalled(); + const capturedEntry = getCapturedEntry(); expect(capturedEntry).toBeDefined(); expect(capturedEntry?.acp).toEqual(existingAcpMeta); }); @@ -404,6 +403,258 @@ describe("gateway agent handler", () => { ); }); + /** + * Test for issue #5369: Verify that modelOverride from sessions.patch is preserved. + * + * When sessions.patch sets a modelOverride on a session, the agent handler + * must use the fresh store data (not potentially stale cached data) when + * building the session entry. + * + * This test simulates: + * 1. loadSessionEntry returning stale data (no modelOverride) - cache hit + * 2. updateSessionStore's store having fresh data (with modelOverride) + * 3. The mutator should use the fresh modelOverride, not the stale one + */ + it("issue #5369: agent handler preserves fresh modelOverride from store", async () => { + // Simulate stale cache: loadSessionEntry returns entry WITHOUT modelOverride + mocks.loadSessionEntry.mockReturnValue({ + cfg: {}, + storePath: "/tmp/sessions.json", + entry: { + sessionId: "subagent-session-id", + updatedAt: Date.now() - 1000, // Slightly older + // NO modelOverride - simulating stale cache read + }, + canonicalKey: "agent:main:subagent:test-uuid", + }); + + let capturedEntry: Record | undefined; + mocks.updateSessionStore.mockImplementation(async (_path, updater) => { + // Simulate fresh store read inside updateSessionStore (with skipCache: true) + // This store HAS the modelOverride that sessions.patch just wrote + const freshStore: Record> = { + "agent:main:subagent:test-uuid": { + sessionId: "subagent-session-id", + updatedAt: Date.now(), + modelOverride: "qwen3-coder:30b", // Fresh data from sessions.patch + providerOverride: "ollama", + }, + }; + const result = await updater(freshStore); + capturedEntry = freshStore["agent:main:subagent:test-uuid"]; + return result; + }); + + mocks.agentCommand.mockResolvedValue({ + payloads: [{ text: "ok" }], + meta: { durationMs: 100 }, + }); + + const respond = vi.fn(); + await agentHandlers.agent({ + params: { + message: "test subagent task", + agentId: "main", + sessionKey: "agent:main:subagent:test-uuid", + idempotencyKey: "test-model-override-race", + }, + respond, + context: makeContext(), + req: { type: "req", id: "race-1", method: "agent" }, + client: null, + isWebchatConnect: () => false, + }); + + expect(mocks.updateSessionStore).toHaveBeenCalled(); + expect(capturedEntry).toBeDefined(); + + // CORRECT BEHAVIOR: modelOverride should be preserved from fresh store + expect(capturedEntry?.modelOverride).toBe("qwen3-coder:30b"); + expect(capturedEntry?.providerOverride).toBe("ollama"); + }); + + /** + * Test that request values (label, spawnedBy) override store values. + * This ensures the fix correctly prioritizes request params over fresh store data. + */ + it("issue #5369: request params override fresh store values for label/spawnedBy", async () => { + mocks.loadSessionEntry.mockReturnValue({ + cfg: {}, + storePath: "/tmp/sessions.json", + entry: { + sessionId: "subagent-session-id", + updatedAt: Date.now() - 1000, + label: "old-label-from-cache", + spawnedBy: "old-spawner", + }, + canonicalKey: "agent:main:subagent:test-priority", + }); + + let capturedEntry: Record | undefined; + mocks.updateSessionStore.mockImplementation(async (_path, updater) => { + const freshStore: Record> = { + "agent:main:subagent:test-priority": { + sessionId: "subagent-session-id", + updatedAt: Date.now(), + label: "store-label", + spawnedBy: "store-spawner", + modelOverride: "gpt-4", // Should be preserved + }, + }; + const result = await updater(freshStore); + capturedEntry = freshStore["agent:main:subagent:test-priority"]; + return result; + }); + + mocks.agentCommand.mockResolvedValue({ + payloads: [{ text: "ok" }], + meta: { durationMs: 100 }, + }); + + const respond = vi.fn(); + await agentHandlers.agent({ + params: { + message: "test", + agentId: "main", + sessionKey: "agent:main:subagent:test-priority", + idempotencyKey: "test-priority", + label: "request-label", // Should take precedence + spawnedBy: "request-spawner", // Should take precedence + }, + respond, + context: makeContext(), + req: { type: "req", id: "priority-1", method: "agent" }, + client: null, + isWebchatConnect: () => false, + }); + + expect(capturedEntry).toBeDefined(); + // Request values should override store values + expect(capturedEntry?.label).toBe("request-label"); + expect(capturedEntry?.spawnedBy).toBe("agent:main:request-spawner"); + // But modelOverride should still come from fresh store + expect(capturedEntry?.modelOverride).toBe("gpt-4"); + }); + + /** + * Test that a new session entry is created correctly when store has no entry. + */ + it("issue #5369: creates new entry when store has no existing entry", async () => { + mocks.loadSessionEntry.mockReturnValue({ + cfg: {}, + storePath: "/tmp/sessions.json", + entry: undefined, // No existing entry + canonicalKey: "agent:main:subagent:new-session", + }); + + let capturedEntry: Record | undefined; + mocks.updateSessionStore.mockImplementation(async (_path, updater) => { + // Fresh store also has no entry - brand new session + const freshStore: Record> = {}; + const result = await updater(freshStore); + capturedEntry = freshStore["agent:main:subagent:new-session"]; + return result; + }); + + mocks.agentCommand.mockResolvedValue({ + payloads: [{ text: "ok" }], + meta: { durationMs: 100 }, + }); + + const respond = vi.fn(); + await agentHandlers.agent({ + params: { + message: "test new session", + agentId: "main", + sessionKey: "agent:main:subagent:new-session", + idempotencyKey: "test-new-session", + label: "new-label", + }, + respond, + context: makeContext(), + req: { type: "req", id: "new-1", method: "agent" }, + client: null, + isWebchatConnect: () => false, + }); + + expect(capturedEntry).toBeDefined(); + expect(capturedEntry?.sessionId).toBeDefined(); // Should generate a sessionId + expect(capturedEntry?.label).toBe("new-label"); + expect(capturedEntry?.modelOverride).toBeUndefined(); // No override for new session + }); + + /** + * Test that all important fields are preserved from fresh store. + */ + it("issue #5369: preserves all important fields from fresh store", async () => { + mocks.loadSessionEntry.mockReturnValue({ + cfg: {}, + storePath: "/tmp/sessions.json", + entry: { + sessionId: "session-id", + updatedAt: Date.now() - 1000, + // Stale data - missing all the fields + }, + canonicalKey: "agent:main:subagent:all-fields", + }); + + let capturedEntry: Record | undefined; + mocks.updateSessionStore.mockImplementation(async (_path, updater) => { + const freshStore: Record> = { + "agent:main:subagent:all-fields": { + sessionId: "session-id", + updatedAt: Date.now(), + thinkingLevel: "high", + verboseLevel: "detailed", + reasoningLevel: "on", + systemSent: true, + sendPolicy: "allow", + skillsSnapshot: { tools: ["bash"] }, + modelOverride: "claude-opus", + providerOverride: "anthropic", + cliSessionIds: { "claude-cli": "xyz" }, + claudeCliSessionId: "xyz", + }, + }; + const result = await updater(freshStore); + capturedEntry = freshStore["agent:main:subagent:all-fields"]; + return result; + }); + + mocks.agentCommand.mockResolvedValue({ + payloads: [{ text: "ok" }], + meta: { durationMs: 100 }, + }); + + const respond = vi.fn(); + await agentHandlers.agent({ + params: { + message: "test all fields", + agentId: "main", + sessionKey: "agent:main:subagent:all-fields", + idempotencyKey: "test-all-fields", + }, + respond, + context: makeContext(), + req: { type: "req", id: "all-1", method: "agent" }, + client: null, + isWebchatConnect: () => false, + }); + + expect(capturedEntry).toBeDefined(); + // All fields should be preserved from fresh store + expect(capturedEntry?.thinkingLevel).toBe("high"); + expect(capturedEntry?.verboseLevel).toBe("detailed"); + expect(capturedEntry?.reasoningLevel).toBe("on"); + expect(capturedEntry?.systemSent).toBe(true); + expect(capturedEntry?.sendPolicy).toBe("allow"); + expect(capturedEntry?.skillsSnapshot).toEqual({ tools: ["bash"] }); + expect(capturedEntry?.modelOverride).toBe("claude-opus"); + expect(capturedEntry?.providerOverride).toBe("anthropic"); + expect(capturedEntry?.cliSessionIds).toEqual({ "claude-cli": "xyz" }); + expect(capturedEntry?.claudeCliSessionId).toBe("xyz"); + }); + it("preserves cliSessionIds from existing session entry", async () => { const existingCliSessionIds = { "claude-cli": "abc-123-def" }; const existingClaudeCliSessionId = "abc-123-def"; @@ -413,7 +664,22 @@ describe("gateway agent handler", () => { claudeCliSessionId: existingClaudeCliSessionId, }); - const capturedEntry = await runMainAgentAndCaptureEntry("test-idem"); + const getCapturedEntry = captureUpdatedMainEntry( + buildExistingMainStoreEntry({ + cliSessionIds: existingCliSessionIds, + claudeCliSessionId: existingClaudeCliSessionId, + }), + ); + + mocks.agentCommand.mockResolvedValue({ + payloads: [{ text: "ok" }], + meta: { durationMs: 100 }, + }); + + await runMainAgent("test", "test-idem"); + + expect(mocks.updateSessionStore).toHaveBeenCalled(); + const capturedEntry = getCapturedEntry(); expect(capturedEntry).toBeDefined(); expect(capturedEntry?.cliSessionIds).toEqual(existingCliSessionIds); expect(capturedEntry?.claudeCliSessionId).toBe(existingClaudeCliSessionId); @@ -424,6 +690,13 @@ describe("gateway agent handler", () => { primeMainAgentRun({ cfg: mocks.loadConfigReturn }); + captureUpdatedMainEntry(buildExistingMainStoreEntry()); + + mocks.agentCommand.mockResolvedValue({ + payloads: [{ text: "ok" }], + meta: { durationMs: 100 }, + }); + await invokeAgent( { message: "Is it the weekend?", @@ -640,8 +913,9 @@ describe("gateway agent handler", () => { "agent:main:work": { sessionId: "existing-session-id", updatedAt: 10 }, "agent:main:MAIN": { sessionId: "legacy-session-id", updatedAt: 5 }, }; - await updater(store); + const result = await updater(store); capturedStore = store; + return result; }); mocks.agentCommand.mockResolvedValue({ diff --git a/src/gateway/server-methods/agent.ts b/src/gateway/server-methods/agent.ts index 9ab032a2edd..33782e639be 100644 --- a/src/gateway/server-methods/agent.ts +++ b/src/gateway/server-methods/agent.ts @@ -374,7 +374,7 @@ export const agentHandlers: GatewayRequestHandlers = { const { cfg, storePath, entry, canonicalKey } = loadSessionEntry(requestedSessionKey); cfgForAgent = cfg; const now = Date.now(); - const sessionId = entry?.sessionId ?? randomUUID(); + const generatedSessionId = entry?.sessionId ?? randomUUID(); const labelValue = request.label?.trim() || entry?.label; const sessionAgent = resolveAgentIdFromSessionKey(canonicalKey); spawnedByValue = canonicalizeSpawnedByForAgent(cfg, sessionAgent, entry?.spawnedBy); @@ -396,68 +396,122 @@ export const agentHandlers: GatewayRequestHandlers = { resolvedGroupId = resolvedGroupId || inheritedGroup?.groupId; resolvedGroupChannel = resolvedGroupChannel || inheritedGroup?.groupChannel; resolvedGroupSpace = resolvedGroupSpace || inheritedGroup?.groupSpace; - const deliveryFields = normalizeSessionDeliveryFields(entry); - const nextEntryPatch: SessionEntry = { - sessionId, - updatedAt: now, - thinkingLevel: entry?.thinkingLevel, - fastMode: entry?.fastMode, - verboseLevel: entry?.verboseLevel, - reasoningLevel: entry?.reasoningLevel, - systemSent: entry?.systemSent, - sendPolicy: entry?.sendPolicy, - skillsSnapshot: entry?.skillsSnapshot, - deliveryContext: deliveryFields.deliveryContext, - lastChannel: deliveryFields.lastChannel ?? entry?.lastChannel, - lastTo: deliveryFields.lastTo ?? entry?.lastTo, - lastAccountId: deliveryFields.lastAccountId ?? entry?.lastAccountId, - modelOverride: entry?.modelOverride, - providerOverride: entry?.providerOverride, - label: labelValue, - spawnedBy: spawnedByValue, - spawnedWorkspaceDir: entry?.spawnedWorkspaceDir, - spawnDepth: entry?.spawnDepth, - channel: entry?.channel ?? request.channel?.trim(), - groupId: resolvedGroupId ?? entry?.groupId, - groupChannel: resolvedGroupChannel ?? entry?.groupChannel, - space: resolvedGroupSpace ?? entry?.space, - cliSessionIds: entry?.cliSessionIds, - claudeCliSessionId: entry?.claudeCliSessionId, - }; - sessionEntry = mergeSessionEntry(entry, nextEntryPatch); - const sendPolicy = resolveSendPolicy({ - cfg, - entry, - sessionKey: canonicalKey, - channel: entry?.channel, - chatType: entry?.chatType, - }); - if (sendPolicy === "deny") { - respond( - false, - undefined, - errorShape(ErrorCodes.INVALID_REQUEST, "send blocked by session policy"), - ); - return; - } - resolvedSessionId = sessionId; const canonicalSessionKey = canonicalKey; resolvedSessionKey = canonicalSessionKey; const agentId = resolveAgentIdFromSessionKey(canonicalSessionKey); const mainSessionKey = resolveAgentMainSessionKey({ cfg, agentId }); if (storePath) { - const persisted = await updateSessionStore(storePath, (store) => { + // Build entry inside updateSessionStore to use fresh store data. + // This avoids race conditions where sessions.patch sets modelOverride + // or sendPolicy between our initial read and this write (issue #5369). + const storeResult = await updateSessionStore(storePath, (store) => { const { primaryKey } = migrateAndPruneGatewaySessionStoreKey({ cfg, key: requestedSessionKey, store, }); - const merged = mergeSessionEntry(store[primaryKey], nextEntryPatch); + const freshEntry = store[primaryKey]; + // Check send policy using fresh data to avoid stale policy decisions + const sendPolicy = resolveSendPolicy({ + cfg, + entry: freshEntry, + sessionKey: canonicalKey, + channel: freshEntry?.channel, + chatType: freshEntry?.chatType, + }); + if (sendPolicy === "deny") { + // Don't write to store, return denial indicator + return { denied: true as const }; + } + const deliveryFields = normalizeSessionDeliveryFields(freshEntry); + const nextEntry: SessionEntry = { + sessionId: freshEntry?.sessionId ?? generatedSessionId, + updatedAt: now, + thinkingLevel: freshEntry?.thinkingLevel, + fastMode: freshEntry?.fastMode, + verboseLevel: freshEntry?.verboseLevel, + reasoningLevel: freshEntry?.reasoningLevel, + systemSent: freshEntry?.systemSent, + sendPolicy: freshEntry?.sendPolicy, + skillsSnapshot: freshEntry?.skillsSnapshot, + deliveryContext: deliveryFields.deliveryContext, + lastChannel: deliveryFields.lastChannel ?? freshEntry?.lastChannel, + lastTo: deliveryFields.lastTo ?? freshEntry?.lastTo, + lastAccountId: deliveryFields.lastAccountId ?? freshEntry?.lastAccountId, + modelOverride: freshEntry?.modelOverride, + providerOverride: freshEntry?.providerOverride, + label: labelValue ?? freshEntry?.label, + spawnedBy: spawnedByValue ?? freshEntry?.spawnedBy, + spawnedWorkspaceDir: freshEntry?.spawnedWorkspaceDir, + spawnDepth: freshEntry?.spawnDepth, + channel: freshEntry?.channel ?? request.channel?.trim(), + groupId: resolvedGroupId ?? freshEntry?.groupId, + groupChannel: resolvedGroupChannel ?? freshEntry?.groupChannel, + space: resolvedGroupSpace ?? freshEntry?.space, + cliSessionIds: freshEntry?.cliSessionIds, + claudeCliSessionId: freshEntry?.claudeCliSessionId, + }; + // Use mergeSessionEntry to preserve extra fields (e.g. acp metadata) + const merged = mergeSessionEntry(freshEntry, nextEntry); store[primaryKey] = merged; - return merged; + return { denied: false as const, entry: merged }; + }); + if (storeResult.denied) { + respond( + false, + undefined, + errorShape(ErrorCodes.INVALID_REQUEST, "send blocked by session policy"), + ); + return; + } + sessionEntry = storeResult.entry; + } else { + // No store path - use initial entry for policy check and build (fallback) + const sendPolicy = resolveSendPolicy({ + cfg, + entry, + sessionKey: requestedSessionKey, + channel: entry?.channel, + chatType: entry?.chatType, + }); + if (sendPolicy === "deny") { + respond( + false, + undefined, + errorShape(ErrorCodes.INVALID_REQUEST, "send blocked by session policy"), + ); + return; + } + const deliveryFields = normalizeSessionDeliveryFields(entry); + sessionEntry = mergeSessionEntry(entry, { + sessionId: generatedSessionId, + updatedAt: now, + thinkingLevel: entry?.thinkingLevel, + fastMode: entry?.fastMode, + verboseLevel: entry?.verboseLevel, + reasoningLevel: entry?.reasoningLevel, + systemSent: entry?.systemSent, + sendPolicy: entry?.sendPolicy, + skillsSnapshot: entry?.skillsSnapshot, + deliveryContext: deliveryFields.deliveryContext, + lastChannel: deliveryFields.lastChannel ?? entry?.lastChannel, + lastTo: deliveryFields.lastTo ?? entry?.lastTo, + lastAccountId: deliveryFields.lastAccountId ?? entry?.lastAccountId, + modelOverride: entry?.modelOverride, + providerOverride: entry?.providerOverride, + label: labelValue, + spawnedBy: spawnedByValue, + spawnedWorkspaceDir: entry?.spawnedWorkspaceDir, + spawnDepth: entry?.spawnDepth, + channel: entry?.channel ?? request.channel?.trim(), + groupId: resolvedGroupId ?? entry?.groupId, + groupChannel: resolvedGroupChannel ?? entry?.groupChannel, + space: resolvedGroupSpace ?? entry?.space, + cliSessionIds: entry?.cliSessionIds, + claudeCliSessionId: entry?.claudeCliSessionId, }); - sessionEntry = persisted; } + resolvedSessionId = sessionEntry.sessionId; if (canonicalSessionKey === mainSessionKey || canonicalSessionKey === "global") { context.addChatRun(idem, { sessionKey: canonicalSessionKey,