Files
openclaw/src/agents/subagent-registry.persistence.test.ts
2026-03-03 00:15:00 +00:00

438 lines
15 KiB
TypeScript

import fs from "node:fs/promises";
import os from "node:os";
import path from "node:path";
import { afterEach, describe, expect, it, vi } from "vitest";
import "./subagent-registry.mocks.shared.js";
import { captureEnv } from "../test-utils/env.js";
import {
addSubagentRunForTests,
clearSubagentRunSteerRestart,
initSubagentRegistry,
listSubagentRunsForRequester,
registerSubagentRun,
resetSubagentRegistryForTests,
} from "./subagent-registry.js";
import { loadSubagentRegistryFromDisk } from "./subagent-registry.store.js";
const { announceSpy } = vi.hoisted(() => ({
announceSpy: vi.fn(async () => true),
}));
vi.mock("./subagent-announce.js", () => ({
runSubagentAnnounceFlow: announceSpy,
}));
describe("subagent registry persistence", () => {
const envSnapshot = captureEnv(["OPENCLAW_STATE_DIR"]);
let tempStateDir: string | null = null;
const resolveAgentIdFromSessionKey = (sessionKey: string) => {
const match = sessionKey.match(/^agent:([^:]+):/i);
return (match?.[1] ?? "main").trim().toLowerCase() || "main";
};
const resolveSessionStorePath = (stateDir: string, agentId: string) =>
path.join(stateDir, "agents", agentId, "sessions", "sessions.json");
const readSessionStore = async (storePath: string) => {
try {
const raw = await fs.readFile(storePath, "utf8");
const parsed = JSON.parse(raw) as unknown;
if (parsed && typeof parsed === "object" && !Array.isArray(parsed)) {
return parsed as Record<string, Record<string, unknown>>;
}
} catch {
// ignore
}
return {} as Record<string, Record<string, unknown>>;
};
const writeChildSessionEntry = async (params: {
sessionKey: string;
sessionId?: string;
updatedAt?: number;
}) => {
if (!tempStateDir) {
throw new Error("tempStateDir not initialized");
}
const agentId = resolveAgentIdFromSessionKey(params.sessionKey);
const storePath = resolveSessionStorePath(tempStateDir, agentId);
const store = await readSessionStore(storePath);
store[params.sessionKey] = {
...store[params.sessionKey],
sessionId: params.sessionId ?? `sess-${agentId}-${Date.now()}`,
updatedAt: params.updatedAt ?? Date.now(),
};
await fs.mkdir(path.dirname(storePath), { recursive: true });
await fs.writeFile(storePath, `${JSON.stringify(store)}\n`, "utf8");
return storePath;
};
const removeChildSessionEntry = async (sessionKey: string) => {
if (!tempStateDir) {
throw new Error("tempStateDir not initialized");
}
const agentId = resolveAgentIdFromSessionKey(sessionKey);
const storePath = resolveSessionStorePath(tempStateDir, agentId);
const store = await readSessionStore(storePath);
delete store[sessionKey];
await fs.mkdir(path.dirname(storePath), { recursive: true });
await fs.writeFile(storePath, `${JSON.stringify(store)}\n`, "utf8");
return storePath;
};
const seedChildSessionsForPersistedRuns = async (persisted: Record<string, unknown>) => {
const runs = (persisted.runs ?? {}) as Record<
string,
{
runId?: string;
childSessionKey?: string;
}
>;
for (const [runId, run] of Object.entries(runs)) {
const childSessionKey = run?.childSessionKey?.trim();
if (!childSessionKey) {
continue;
}
await writeChildSessionEntry({
sessionKey: childSessionKey,
sessionId: `sess-${run.runId ?? runId}`,
});
}
};
const writePersistedRegistry = async (
persisted: Record<string, unknown>,
opts?: { seedChildSessions?: boolean },
) => {
tempStateDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-subagent-"));
process.env.OPENCLAW_STATE_DIR = tempStateDir;
const registryPath = path.join(tempStateDir, "subagents", "runs.json");
await fs.mkdir(path.dirname(registryPath), { recursive: true });
await fs.writeFile(registryPath, `${JSON.stringify(persisted)}\n`, "utf8");
if (opts?.seedChildSessions !== false) {
await seedChildSessionsForPersistedRuns(persisted);
}
return registryPath;
};
const readPersistedRun = async <T>(
registryPath: string,
runId: string,
): Promise<T | undefined> => {
const parsed = JSON.parse(await fs.readFile(registryPath, "utf8")) as {
runs?: Record<string, unknown>;
};
return parsed.runs?.[runId] as T | undefined;
};
const createPersistedEndedRun = (params: {
runId: string;
childSessionKey: string;
task: string;
cleanup: "keep" | "delete";
}) => {
const now = Date.now();
return {
version: 2,
runs: {
[params.runId]: {
runId: params.runId,
childSessionKey: params.childSessionKey,
requesterSessionKey: "agent:main:main",
requesterDisplayKey: "main",
task: params.task,
cleanup: params.cleanup,
createdAt: now - 2,
startedAt: now - 1,
endedAt: now,
},
},
};
};
const flushQueuedRegistryWork = async () => {
await Promise.resolve();
await Promise.resolve();
};
const restartRegistryAndFlush = async () => {
resetSubagentRegistryForTests({ persist: false });
initSubagentRegistry();
await flushQueuedRegistryWork();
};
afterEach(async () => {
announceSpy.mockClear();
resetSubagentRegistryForTests({ persist: false });
if (tempStateDir) {
await fs.rm(tempStateDir, { recursive: true, force: true });
tempStateDir = null;
}
envSnapshot.restore();
});
it("persists runs to disk and resumes after restart", async () => {
tempStateDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-subagent-"));
process.env.OPENCLAW_STATE_DIR = tempStateDir;
registerSubagentRun({
runId: "run-1",
childSessionKey: "agent:main:subagent:test",
requesterSessionKey: "agent:main:main",
requesterOrigin: { channel: " whatsapp ", accountId: " acct-main " },
requesterDisplayKey: "main",
task: "do the thing",
cleanup: "keep",
});
await writeChildSessionEntry({
sessionKey: "agent:main:subagent:test",
sessionId: "sess-test",
});
const registryPath = path.join(tempStateDir, "subagents", "runs.json");
const raw = await fs.readFile(registryPath, "utf8");
const parsed = JSON.parse(raw) as { runs?: Record<string, unknown> };
expect(parsed.runs && Object.keys(parsed.runs)).toContain("run-1");
const run = parsed.runs?.["run-1"] as
| {
requesterOrigin?: { channel?: string; accountId?: string };
}
| undefined;
expect(run).toBeDefined();
if (run) {
expect("requesterAccountId" in run).toBe(false);
expect("requesterChannel" in run).toBe(false);
}
expect(run?.requesterOrigin?.channel).toBe("whatsapp");
expect(run?.requesterOrigin?.accountId).toBe("acct-main");
// Simulate a process restart: module re-import should load persisted runs
// and trigger the announce flow once the run resolves.
resetSubagentRegistryForTests({ persist: false });
initSubagentRegistry();
// allow queued async wait/cleanup to execute
await flushQueuedRegistryWork();
expect(announceSpy).toHaveBeenCalled();
type AnnounceParams = {
childSessionKey: string;
childRunId: string;
requesterSessionKey: string;
requesterOrigin?: { channel?: string; accountId?: string };
task: string;
cleanup: string;
label?: string;
};
const first = (announceSpy.mock.calls as unknown as Array<[unknown]>)[0]?.[0] as
| AnnounceParams
| undefined;
if (!first) {
throw new Error("expected announce call");
}
expect(first.childSessionKey).toBe("agent:main:subagent:test");
expect(first.requesterOrigin?.channel).toBe("whatsapp");
expect(first.requesterOrigin?.accountId).toBe("acct-main");
});
it("skips cleanup when cleanupHandled was persisted", async () => {
tempStateDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-subagent-"));
process.env.OPENCLAW_STATE_DIR = tempStateDir;
const registryPath = path.join(tempStateDir, "subagents", "runs.json");
const persisted = {
version: 2,
runs: {
"run-2": {
runId: "run-2",
childSessionKey: "agent:main:subagent:two",
requesterSessionKey: "agent:main:main",
requesterDisplayKey: "main",
task: "do the other thing",
cleanup: "keep",
createdAt: 1,
startedAt: 1,
endedAt: 2,
cleanupHandled: true, // Already handled - should be skipped
},
},
};
await fs.mkdir(path.dirname(registryPath), { recursive: true });
await fs.writeFile(registryPath, `${JSON.stringify(persisted)}\n`, "utf8");
await writeChildSessionEntry({
sessionKey: "agent:main:subagent:two",
sessionId: "sess-two",
});
resetSubagentRegistryForTests({ persist: false });
initSubagentRegistry();
await flushQueuedRegistryWork();
// announce should NOT be called since cleanupHandled was true
const calls = (announceSpy.mock.calls as unknown as Array<[unknown]>).map((call) => call[0]);
const match = calls.find(
(params) =>
(params as { childSessionKey?: string }).childSessionKey === "agent:main:subagent:two",
);
expect(match).toBeFalsy();
});
it("maps legacy announce fields into cleanup state", async () => {
const persisted = {
version: 1,
runs: {
"run-legacy": {
runId: "run-legacy",
childSessionKey: "agent:main:subagent:legacy",
requesterSessionKey: "agent:main:main",
requesterDisplayKey: "main",
task: "legacy announce",
cleanup: "keep",
createdAt: 1,
startedAt: 1,
endedAt: 2,
announceCompletedAt: 9,
announceHandled: true,
requesterChannel: "whatsapp",
requesterAccountId: "legacy-account",
},
},
};
const registryPath = await writePersistedRegistry(persisted);
const runs = loadSubagentRegistryFromDisk();
const entry = runs.get("run-legacy");
expect(entry?.cleanupHandled).toBe(true);
expect(entry?.cleanupCompletedAt).toBe(9);
expect(entry?.requesterOrigin?.channel).toBe("whatsapp");
expect(entry?.requesterOrigin?.accountId).toBe("legacy-account");
const after = JSON.parse(await fs.readFile(registryPath, "utf8")) as { version?: number };
expect(after.version).toBe(2);
});
it("retries cleanup announce after a failed announce", async () => {
const persisted = createPersistedEndedRun({
runId: "run-3",
childSessionKey: "agent:main:subagent:three",
task: "retry announce",
cleanup: "keep",
});
const registryPath = await writePersistedRegistry(persisted);
announceSpy.mockResolvedValueOnce(false);
await restartRegistryAndFlush();
expect(announceSpy).toHaveBeenCalledTimes(1);
const afterFirst = await readPersistedRun<{
cleanupHandled?: boolean;
cleanupCompletedAt?: number;
}>(registryPath, "run-3");
expect(afterFirst?.cleanupHandled).toBe(false);
expect(afterFirst?.cleanupCompletedAt).toBeUndefined();
announceSpy.mockResolvedValueOnce(true);
await restartRegistryAndFlush();
expect(announceSpy).toHaveBeenCalledTimes(2);
const afterSecond = JSON.parse(await fs.readFile(registryPath, "utf8")) as {
runs: Record<string, { cleanupCompletedAt?: number }>;
};
expect(afterSecond.runs["run-3"].cleanupCompletedAt).toBeDefined();
});
it("keeps delete-mode runs retryable when announce is deferred", async () => {
const persisted = createPersistedEndedRun({
runId: "run-4",
childSessionKey: "agent:main:subagent:four",
task: "deferred announce",
cleanup: "delete",
});
const registryPath = await writePersistedRegistry(persisted);
announceSpy.mockResolvedValueOnce(false);
await restartRegistryAndFlush();
expect(announceSpy).toHaveBeenCalledTimes(1);
const afterFirst = await readPersistedRun<{ cleanupHandled?: boolean }>(registryPath, "run-4");
expect(afterFirst?.cleanupHandled).toBe(false);
announceSpy.mockResolvedValueOnce(true);
await restartRegistryAndFlush();
expect(announceSpy).toHaveBeenCalledTimes(2);
const afterSecond = JSON.parse(await fs.readFile(registryPath, "utf8")) as {
runs?: Record<string, unknown>;
};
expect(afterSecond.runs?.["run-4"]).toBeUndefined();
});
it("reconciles orphaned restored runs by pruning them from registry", async () => {
const persisted = createPersistedEndedRun({
runId: "run-orphan-restore",
childSessionKey: "agent:main:subagent:ghost-restore",
task: "orphan restore",
cleanup: "keep",
});
const registryPath = await writePersistedRegistry(persisted, {
seedChildSessions: false,
});
await restartRegistryAndFlush();
expect(announceSpy).not.toHaveBeenCalled();
const after = JSON.parse(await fs.readFile(registryPath, "utf8")) as {
runs?: Record<string, unknown>;
};
expect(after.runs?.["run-orphan-restore"]).toBeUndefined();
expect(listSubagentRunsForRequester("agent:main:main")).toHaveLength(0);
});
it("resume guard prunes orphan runs before announce retry", async () => {
tempStateDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-subagent-"));
process.env.OPENCLAW_STATE_DIR = tempStateDir;
const runId = "run-orphan-resume-guard";
const childSessionKey = "agent:main:subagent:ghost-resume";
const now = Date.now();
await writeChildSessionEntry({
sessionKey: childSessionKey,
sessionId: "sess-resume-guard",
updatedAt: now,
});
addSubagentRunForTests({
runId,
childSessionKey,
requesterSessionKey: "agent:main:main",
requesterDisplayKey: "main",
task: "resume orphan guard",
cleanup: "keep",
createdAt: now - 50,
startedAt: now - 25,
endedAt: now,
suppressAnnounceReason: "steer-restart",
cleanupHandled: false,
});
await removeChildSessionEntry(childSessionKey);
const changed = clearSubagentRunSteerRestart(runId);
expect(changed).toBe(true);
await flushQueuedRegistryWork();
expect(announceSpy).not.toHaveBeenCalled();
expect(listSubagentRunsForRequester("agent:main:main")).toHaveLength(0);
const persisted = loadSubagentRegistryFromDisk();
expect(persisted.has(runId)).toBe(false);
});
it("uses isolated temp state when OPENCLAW_STATE_DIR is unset in tests", async () => {
delete process.env.OPENCLAW_STATE_DIR;
vi.resetModules();
const { resolveSubagentRegistryPath } = await import("./subagent-registry.store.js");
const registryPath = resolveSubagentRegistryPath();
expect(registryPath).toContain(path.join(os.tmpdir(), "openclaw-test-state"));
});
});