Moltbot/src/commands/agent/session.test.ts

import { beforeEach, describe, expect, it, vi } from "vitest";
import type { OpenClawConfig } from "../../config/config.js";

const mocks = vi.hoisted(() => ({
  loadSessionStore: vi.fn(),
  resolveStorePath: vi.fn(),
  listAgentIds: vi.fn(),
}));

vi.mock("../../config/sessions.js", async () => {
  const actual = await vi.importActual<typeof import("../../config/sessions.js")>(
    "../../config/sessions.js",
  );
  return {
    ...actual,
    loadSessionStore: mocks.loadSessionStore,
    resolveStorePath: mocks.resolveStorePath,
  };
});

vi.mock("../../agents/agent-scope.js", () => ({
  listAgentIds: mocks.listAgentIds,
}));

describe("resolveSessionKeyForRequest", () => {
  beforeEach(() => {
    vi.clearAllMocks();
    mocks.listAgentIds.mockReturnValue(["main"]);
  });

  async function importFresh() {
    return await import("./session.js");
  }

  const baseCfg: OpenClawConfig = {};

  it("returns sessionKey when --to resolves a session key via context", async () => {
    const { resolveSessionKeyForRequest } = await importFresh();

    mocks.resolveStorePath.mockReturnValue("/tmp/main-store.json");
    mocks.loadSessionStore.mockReturnValue({
      "agent:main:main": { sessionId: "sess-1", updatedAt: 0 },
    });

    const result = resolveSessionKeyForRequest({
      cfg: baseCfg,
      to: "+15551234567",
    });
    expect(result.sessionKey).toBe("agent:main:main");
  });

  it("finds session by sessionId via reverse lookup in primary store", async () => {
    const { resolveSessionKeyForRequest } = await importFresh();

    mocks.resolveStorePath.mockReturnValue("/tmp/main-store.json");
    mocks.loadSessionStore.mockReturnValue({
      "agent:main:main": { sessionId: "target-session-id", updatedAt: 0 },
    });

    const result = resolveSessionKeyForRequest({
      cfg: baseCfg,
      sessionId: "target-session-id",
    });
    expect(result.sessionKey).toBe("agent:main:main");
  });

  it("finds session by sessionId in non-primary agent store", async () => {
    const { resolveSessionKeyForRequest } = await importFresh();

    mocks.listAgentIds.mockReturnValue(["main", "mybot"]);
    mocks.resolveStorePath.mockImplementation(
      (_store: string | undefined, opts?: { agentId?: string }) => {
        if (opts?.agentId === "mybot") {
          return "/tmp/mybot-store.json";
        }
        return "/tmp/main-store.json";
      },
    );
    mocks.loadSessionStore.mockImplementation((storePath: string) => {
      if (storePath === "/tmp/mybot-store.json") {
        return {
          "agent:mybot:main": { sessionId: "target-session-id", updatedAt: 0 },
        };
      }
      return {};
    });

    const result = resolveSessionKeyForRequest({
      cfg: baseCfg,
      sessionId: "target-session-id",
    });
    expect(result.sessionKey).toBe("agent:mybot:main");
    expect(result.storePath).toBe("/tmp/mybot-store.json");
  });

  it("returns correct sessionStore when session found in non-primary agent store", async () => {
    const { resolveSessionKeyForRequest } = await importFresh();

    const mybotStore = {
      "agent:mybot:main": { sessionId: "target-session-id", updatedAt: 0 },
    };
    mocks.listAgentIds.mockReturnValue(["main", "mybot"]);
    mocks.resolveStorePath.mockImplementation(
      (_store: string | undefined, opts?: { agentId?: string }) => {
        if (opts?.agentId === "mybot") {
          return "/tmp/mybot-store.json";
        }
        return "/tmp/main-store.json";
      },
    );
    mocks.loadSessionStore.mockImplementation((storePath: string) => {
      if (storePath === "/tmp/mybot-store.json") {
        return { ...mybotStore };
      }
      return {};
    });

    const result = resolveSessionKeyForRequest({
      cfg: baseCfg,
      sessionId: "target-session-id",
    });
    expect(result.sessionStore["agent:mybot:main"]?.sessionId).toBe("target-session-id");
  });

  it("returns undefined sessionKey when sessionId not found in any store", async () => {
    const { resolveSessionKeyForRequest } = await importFresh();

    mocks.listAgentIds.mockReturnValue(["main", "mybot"]);
    mocks.resolveStorePath.mockImplementation(
      (_store: string | undefined, opts?: { agentId?: string }) => {
        if (opts?.agentId === "mybot") {
          return "/tmp/mybot-store.json";
        }
        return "/tmp/main-store.json";
      },
    );
    mocks.loadSessionStore.mockReturnValue({});

    const result = resolveSessionKeyForRequest({
      cfg: baseCfg,
      sessionId: "nonexistent-id",
    });
    expect(result.sessionKey).toBeUndefined();
  });

  it("does not search other stores when explicitSessionKey is set", async () => {
    const { resolveSessionKeyForRequest } = await importFresh();

    mocks.listAgentIds.mockReturnValue(["main", "mybot"]);
    mocks.resolveStorePath.mockReturnValue("/tmp/main-store.json");
    mocks.loadSessionStore.mockReturnValue({
      "agent:main:main": { sessionId: "other-id", updatedAt: 0 },
    });

    const result = resolveSessionKeyForRequest({
      cfg: baseCfg,
      sessionKey: "agent:main:main",
      sessionId: "target-session-id",
    });
    // explicitSessionKey is set, so sessionKey comes from it, not from sessionId lookup
    expect(result.sessionKey).toBe("agent:main:main");
  });

  it("searches other stores when --to derives a key that does not match --session-id", async () => {
    const { resolveSessionKeyForRequest } = await importFresh();

    mocks.listAgentIds.mockReturnValue(["main", "mybot"]);
    mocks.resolveStorePath.mockImplementation(
      (_store: string | undefined, opts?: { agentId?: string }) => {
        if (opts?.agentId === "mybot") {
          return "/tmp/mybot-store.json";
        }
        return "/tmp/main-store.json";
      },
    );
    mocks.loadSessionStore.mockImplementation((storePath: string) => {
      if (storePath === "/tmp/main-store.json") {
        return {
          "agent:main:main": { sessionId: "other-session-id", updatedAt: 0 },
        };
      }
      if (storePath === "/tmp/mybot-store.json") {
        return {
          "agent:mybot:main": { sessionId: "target-session-id", updatedAt: 0 },
        };
      }
      return {};
    });

    const result = resolveSessionKeyForRequest({
      cfg: baseCfg,
      to: "+15551234567",
      sessionId: "target-session-id",
    });
    // --to derives agent:main:main, but its sessionId doesn't match target-session-id,
    // so the cross-store search finds it in the mybot store
    expect(result.sessionKey).toBe("agent:mybot:main");
    expect(result.storePath).toBe("/tmp/mybot-store.json");
  });

  it("skips already-searched primary store when iterating agents", async () => {
    const { resolveSessionKeyForRequest } = await importFresh();

    mocks.listAgentIds.mockReturnValue(["main", "mybot"]);
    mocks.resolveStorePath.mockImplementation(
      (_store: string | undefined, opts?: { agentId?: string }) => {
        if (opts?.agentId === "mybot") {
          return "/tmp/mybot-store.json";
        }
        return "/tmp/main-store.json";
      },
    );
    mocks.loadSessionStore.mockReturnValue({});

    resolveSessionKeyForRequest({
      cfg: baseCfg,
      sessionId: "nonexistent-id",
    });

    // loadSessionStore should be called twice: once for main, once for mybot
    // (not twice for main)
    const storePaths = mocks.loadSessionStore.mock.calls.map((call: [string]) => call[0]);
    expect(storePaths).toHaveLength(2);
    expect(storePaths).toContain("/tmp/main-store.json");
    expect(storePaths).toContain("/tmp/mybot-store.json");
  });
});