diff --git a/src/config/sessions/targets.test.ts b/src/config/sessions/targets.test.ts index 012302c323e..f67bc10f8b9 100644 --- a/src/config/sessions/targets.test.ts +++ b/src/config/sessions/targets.test.ts @@ -3,7 +3,71 @@ import path from "node:path"; import { describe, expect, it } from "vitest"; import { withTempHome } from "../../../test/helpers/temp-home.js"; import type { OpenClawConfig } from "../config.js"; -import { resolveAllAgentSessionStoreTargets } from "./targets.js"; +import { resolveAllAgentSessionStoreTargets, resolveSessionStoreTargets } from "./targets.js"; + +describe("resolveSessionStoreTargets", () => { + it("resolves all configured agent stores", () => { + const cfg: OpenClawConfig = { + session: { + store: "~/.openclaw/agents/{agentId}/sessions/sessions.json", + }, + agents: { + list: [{ id: "main", default: true }, { id: "work" }], + }, + }; + + const targets = resolveSessionStoreTargets(cfg, { allAgents: true }); + + expect(targets).toEqual([ + { + agentId: "main", + storePath: path.resolve( + path.join(process.env.HOME ?? "", ".openclaw/agents/main/sessions/sessions.json"), + ), + }, + { + agentId: "work", + storePath: path.resolve( + path.join(process.env.HOME ?? "", ".openclaw/agents/work/sessions/sessions.json"), + ), + }, + ]); + }); + + it("dedupes shared store paths for --all-agents", () => { + const cfg: OpenClawConfig = { + session: { + store: "/tmp/shared-sessions.json", + }, + agents: { + list: [{ id: "main", default: true }, { id: "work" }], + }, + }; + + expect(resolveSessionStoreTargets(cfg, { allAgents: true })).toEqual([ + { agentId: "main", storePath: path.resolve("/tmp/shared-sessions.json") }, + ]); + }); + + it("rejects unknown agent ids", () => { + const cfg: OpenClawConfig = { + agents: { + list: [{ id: "main", default: true }, { id: "work" }], + }, + }; + + expect(() => resolveSessionStoreTargets(cfg, { agent: "ghost" })).toThrow(/Unknown agent id/); + }); + + it("rejects conflicting selectors", () => { + expect(() => resolveSessionStoreTargets({}, { agent: "main", allAgents: true })).toThrow( + /cannot be used together/i, + ); + expect(() => + resolveSessionStoreTargets({}, { store: "/tmp/sessions.json", allAgents: true }), + ).toThrow(/cannot be combined/i); + }); +}); describe("resolveAllAgentSessionStoreTargets", () => { it("includes discovered on-disk agent stores alongside configured targets", async () => { diff --git a/src/gateway/server-session-key.test.ts b/src/gateway/server-session-key.test.ts index 0baac9d3dc1..dfdec62f488 100644 --- a/src/gateway/server-session-key.test.ts +++ b/src/gateway/server-session-key.test.ts @@ -1,54 +1,70 @@ -import fs from "node:fs"; -import path from "node:path"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import type { OpenClawConfig } from "../config/config.js"; import { resetAgentRunContextForTest } from "../infra/agent-events.js"; -import { withStateDirEnv } from "../test-helpers/state-dir-env.js"; -const loadConfigMock = vi.hoisted(() => vi.fn<() => OpenClawConfig>()); - -vi.mock("../config/config.js", () => ({ - loadConfig: () => loadConfigMock(), +const hoisted = vi.hoisted(() => ({ + loadConfigMock: vi.fn<() => OpenClawConfig>(), + loadCombinedSessionStoreForGatewayMock: vi.fn(), })); -const { resolveSessionKeyForRun } = await import("./server-session-key.js"); +vi.mock("../config/config.js", () => ({ + loadConfig: () => hoisted.loadConfigMock(), +})); + +vi.mock("./session-utils.js", async () => { + const actual = await vi.importActual("./session-utils.js"); + return { + ...actual, + loadCombinedSessionStoreForGateway: (cfg: OpenClawConfig) => + hoisted.loadCombinedSessionStoreForGatewayMock(cfg), + }; +}); + +const { resolveSessionKeyForRun, resetResolvedSessionKeyForRunCacheForTest } = + await import("./server-session-key.js"); describe("resolveSessionKeyForRun", () => { beforeEach(() => { - loadConfigMock.mockReset(); + hoisted.loadConfigMock.mockReset(); + hoisted.loadCombinedSessionStoreForGatewayMock.mockReset(); resetAgentRunContextForTest(); + resetResolvedSessionKeyForRunCacheForTest(); }); afterEach(() => { resetAgentRunContextForTest(); + resetResolvedSessionKeyForRunCacheForTest(); }); - it("finds run ids in disk-only agent stores under a custom session root", async () => { - await withStateDirEnv("openclaw-run-key-", async ({ stateDir }) => { - const customRoot = path.join(stateDir, "custom-state"); - const retiredSessionsDir = path.join(customRoot, "agents", "retired", "sessions"); - fs.mkdirSync(retiredSessionsDir, { recursive: true }); - fs.writeFileSync( - path.join(retiredSessionsDir, "sessions.json"), - JSON.stringify({ - "agent:retired:acp:run-1": { sessionId: "run-1", updatedAt: 123 }, - }), - "utf8", - ); - - loadConfigMock.mockReturnValue({ - session: { - store: path.join(customRoot, "agents", "{agentId}", "sessions", "sessions.json"), - }, - agents: { - list: [{ id: "main", default: true }], - }, - }); - - expect(resolveSessionKeyForRun("run-1")).toBe("acp:run-1"); - - fs.rmSync(customRoot, { recursive: true, force: true }); - expect(resolveSessionKeyForRun("run-1")).toBe("acp:run-1"); + it("resolves run ids from the combined gateway store and caches the result", () => { + const cfg: OpenClawConfig = { + session: { + store: "/custom/root/agents/{agentId}/sessions/sessions.json", + }, + }; + hoisted.loadConfigMock.mockReturnValue(cfg); + hoisted.loadCombinedSessionStoreForGatewayMock.mockReturnValue({ + storePath: "(multiple)", + store: { + "agent:retired:acp:run-1": { sessionId: "run-1", updatedAt: 123 }, + }, }); + + expect(resolveSessionKeyForRun("run-1")).toBe("acp:run-1"); + expect(resolveSessionKeyForRun("run-1")).toBe("acp:run-1"); + expect(hoisted.loadCombinedSessionStoreForGatewayMock).toHaveBeenCalledTimes(1); + expect(hoisted.loadCombinedSessionStoreForGatewayMock).toHaveBeenCalledWith(cfg); + }); + + it("caches misses so repeated lookups do not rebuild the combined store", () => { + hoisted.loadConfigMock.mockReturnValue({}); + hoisted.loadCombinedSessionStoreForGatewayMock.mockReturnValue({ + storePath: "(multiple)", + store: {}, + }); + + expect(resolveSessionKeyForRun("missing-run")).toBeUndefined(); + expect(resolveSessionKeyForRun("missing-run")).toBeUndefined(); + expect(hoisted.loadCombinedSessionStoreForGatewayMock).toHaveBeenCalledTimes(1); }); }); diff --git a/src/gateway/server-session-key.ts b/src/gateway/server-session-key.ts index c0eb51d8ca3..aaa742acfa8 100644 --- a/src/gateway/server-session-key.ts +++ b/src/gateway/server-session-key.ts @@ -3,11 +3,34 @@ import { getAgentRunContext, registerAgentRunContext } from "../infra/agent-even import { toAgentRequestSessionKey } from "../routing/session-key.js"; import { loadCombinedSessionStoreForGateway } from "./session-utils.js"; +const RUN_LOOKUP_CACHE_LIMIT = 256; +const resolvedSessionKeyByRunId = new Map(); + +function setResolvedSessionKeyCache(runId: string, sessionKey: string | null): void { + if (!runId) { + return; + } + if ( + !resolvedSessionKeyByRunId.has(runId) && + resolvedSessionKeyByRunId.size >= RUN_LOOKUP_CACHE_LIMIT + ) { + const oldest = resolvedSessionKeyByRunId.keys().next().value; + if (oldest) { + resolvedSessionKeyByRunId.delete(oldest); + } + } + resolvedSessionKeyByRunId.set(runId, sessionKey); +} + export function resolveSessionKeyForRun(runId: string) { const cached = getAgentRunContext(runId)?.sessionKey; if (cached) { return cached; } + const cachedLookup = resolvedSessionKeyByRunId.get(runId); + if (cachedLookup !== undefined) { + return cachedLookup ?? undefined; + } const cfg = loadConfig(); const { store } = loadCombinedSessionStoreForGateway(cfg); const found = Object.entries(store).find(([, entry]) => entry?.sessionId === runId); @@ -15,7 +38,13 @@ export function resolveSessionKeyForRun(runId: string) { if (storeKey) { const sessionKey = toAgentRequestSessionKey(storeKey) ?? storeKey; registerAgentRunContext(runId, { sessionKey }); + setResolvedSessionKeyCache(runId, sessionKey); return sessionKey; } + setResolvedSessionKeyCache(runId, null); return undefined; } + +export function resetResolvedSessionKeyForRunCacheForTest(): void { + resolvedSessionKeyByRunId.clear(); +}