refactor: update registry persistence to use ManagedAgent

Remove AgentSnapshot from persistence layer. AgentRegistry now accepts
ManagedAgent and uses toStoredAgentRecord for atomic config + lifecycle
persistence. Deleted obsolete recordConfig method.

- Update applySnapshot to accept ManagedAgent and use toStoredAgentRecord
- Remove recordConfig method entirely (no longer needed)
- Remove sanitizeConfig helper (handled by projection)
- Update session to stop calling deleted recordConfig
- Add ManagedAgent test fixtures to registry and persistence-hook tests
- Test config persistence, title retention, and subscription forwarding

Registry and persistence-hooks now typecheck cleanly. Expected failures
in session/mcp-server/messages (still reference AgentSnapshot).

Task 4 of 8 in agent architecture refactor.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Mohamed Boudra
2025-11-30 00:03:00 +00:00
co-authored by Claude
parent 8a708aab8b
commit 7aec973323
4 changed files with 191 additions and 188 deletions
@@ -1,37 +1,84 @@
import { describe, expect, test, beforeEach, afterEach, vi } from "vitest"; import { describe, expect, test, beforeEach, afterEach } from "vitest";
import os from "node:os"; import os from "node:os";
import path from "node:path"; import path from "node:path";
import { mkdtempSync, rmSync, writeFileSync, readFileSync } from "node:fs"; import { mkdtempSync, rmSync, writeFileSync, readFileSync } from "node:fs";
import { AgentRegistry } from "./agent-registry.js"; import { AgentRegistry } from "./agent-registry.js";
import type { AgentSnapshot } from "./agent-manager.js"; import type { ManagedAgent } from "./agent-manager.js";
import type {
AgentPermissionRequest,
AgentSession,
AgentSessionConfig,
} from "./agent-sdk-types.js";
function createSnapshot(overrides?: Partial<AgentSnapshot>): AgentSnapshot { type ManagedAgentOverrides = Omit<
const now = new Date(); Partial<ManagedAgent>,
return { "config" | "pendingPermissions" | "session" | "pendingRun"
id: "agent-test", > & {
provider: "claude", config?: Partial<AgentSessionConfig>;
cwd: "/tmp/project", pendingPermissions?: Map<string, AgentPermissionRequest>;
model: null, session?: AgentSession | null;
createdAt: now, pendingRun?: ManagedAgent["pendingRun"];
updatedAt: now, };
lastUserMessageAt: null,
status: "idle", function createManagedAgent(
sessionId: null, overrides: ManagedAgentOverrides = {}
capabilities: { ): ManagedAgent {
supportsStreaming: true, const now = overrides.updatedAt ?? new Date("2025-01-01T00:00:00.000Z");
supportsSessionPersistence: true, const provider = overrides.provider ?? "claude";
supportsDynamicModes: true, const cwd = overrides.cwd ?? "/tmp/project";
supportsMcpServers: true, const lifecycle = overrides.lifecycle ?? "idle";
supportsReasoningStream: true, const configOverrides = overrides.config ?? {};
supportsToolInvocations: true, const config: AgentSessionConfig = {
}, provider,
currentModeId: null, cwd,
availableModes: [], modeId: configOverrides.modeId ?? "plan",
pendingPermissions: [], model: configOverrides.model ?? "gpt-5.1",
persistence: null, extra: configOverrides.extra ?? { claude: { maxThinkingTokens: 1024 } },
...overrides,
}; };
const session =
lifecycle === "closed"
? null
: overrides.session ?? ({} as AgentSession);
const pendingRun =
overrides.pendingRun ??
(lifecycle === "running" ? (async function* noop() {})() : null);
const agent: ManagedAgent = {
id: overrides.id ?? "agent-test",
provider,
cwd,
session,
sessionId: overrides.sessionId ?? "session-123",
capabilities:
overrides.capabilities ??
{
supportsStreaming: true,
supportsSessionPersistence: true,
supportsDynamicModes: true,
supportsMcpServers: true,
supportsReasoningStream: true,
supportsToolInvocations: true,
},
config,
lifecycle,
createdAt: overrides.createdAt ?? now,
updatedAt: overrides.updatedAt ?? now,
availableModes: overrides.availableModes ?? [],
currentModeId: overrides.currentModeId ?? config.modeId ?? null,
pendingPermissions:
overrides.pendingPermissions ??
new Map<string, AgentPermissionRequest>(),
pendingRun,
timeline: overrides.timeline ?? [],
persistence: overrides.persistence ?? null,
historyPrimed: overrides.historyPrimed ?? true,
lastUserMessageAt: overrides.lastUserMessageAt ?? now,
lastUsage: overrides.lastUsage,
lastError: overrides.lastError,
};
return agent;
} }
describe("AgentRegistry", () => { describe("AgentRegistry", () => {
@@ -49,27 +96,21 @@ describe("AgentRegistry", () => {
rmSync(tmpDir, { recursive: true, force: true }); rmSync(tmpDir, { recursive: true, force: true });
}); });
test("persists configs and snapshot metadata", async () => { test("applySnapshot persists configs and snapshot metadata", async () => {
await registry.applySnapshot( await registry.applySnapshot(
createSnapshot({ createManagedAgent({
id: "agent-1", id: "agent-1",
cwd: "/tmp/project", cwd: "/tmp/project",
currentModeId: "coding", currentModeId: "coding",
status: "idle", lifecycle: "idle",
config: {
modeId: "coding",
model: "gpt-5.1",
extra: { claude: { maxThinkingTokens: 1024 } },
},
}) })
); );
await registry.recordConfig(
"agent-1",
"claude",
"/tmp/project",
{
modeId: "coding",
model: "gpt-5.1",
extra: { claude: { maxThinkingTokens: 1024 } },
}
);
const records = await registry.list(); const records = await registry.list();
expect(records).toHaveLength(1); expect(records).toHaveLength(1);
const [record] = records; const [record] = records;
@@ -87,7 +128,7 @@ describe("AgentRegistry", () => {
test("stores titles independently of snapshots", async () => { test("stores titles independently of snapshots", async () => {
await registry.applySnapshot( await registry.applySnapshot(
createSnapshot({ createManagedAgent({
id: "agent-2", id: "agent-2",
provider: "codex", provider: "codex",
cwd: "/tmp/second", cwd: "/tmp/second",
@@ -103,21 +144,30 @@ describe("AgentRegistry", () => {
expect(persisted?.title).toBe("Fix Login Bug"); expect(persisted?.title).toBe("Fix Login Bug");
}); });
test("recordConfig warns if no snapshot exists", async () => { test("applySnapshot preserves custom titles while updating metadata", async () => {
const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); const agentId = "agent-3";
await registry.recordConfig( await registry.applySnapshot(
"agent-3", createManagedAgent({
"claude", id: agentId,
"/tmp/project", lifecycle: "idle",
{ currentModeId: "plan",
modeId: "plan", })
} );
await registry.setTitle(agentId, "Important Bug Fix");
await registry.applySnapshot(
createManagedAgent({
id: agentId,
lifecycle: "running",
currentModeId: "build",
updatedAt: new Date("2025-01-02T00:00:00.000Z"),
})
); );
const record = await registry.get("agent-3"); const record = await registry.get(agentId);
expect(record).toBeNull(); expect(record?.title).toBe("Important Bug Fix");
expect(warnSpy).toHaveBeenCalled(); expect(record?.lastModeId).toBe("build");
warnSpy.mockRestore(); expect(record?.lastStatus).toBe("running");
}); });
test("recovers from trailing garbage in agents.json", async () => { test("recovers from trailing garbage in agents.json", async () => {
@@ -6,8 +6,9 @@ import { fileURLToPath } from "node:url";
import { z } from "zod"; import { z } from "zod";
import { AgentStatusSchema } from "../messages.js"; import { AgentStatusSchema } from "../messages.js";
import type { AgentSnapshot } from "./agent-manager.js"; import { toStoredAgentRecord } from "./agent-projections.js";
import type { AgentProvider, AgentSessionConfig } from "./agent-sdk-types.js"; import type { ManagedAgent } from "./agent-manager.js";
import type { AgentSessionConfig } from "./agent-sdk-types.js";
const SERIALIZABLE_CONFIG_SCHEMA = z const SERIALIZABLE_CONFIG_SCHEMA = z
.object({ .object({
@@ -103,77 +104,13 @@ export class AgentRegistry {
await this.flush(); await this.flush();
} }
async recordConfig( async applySnapshot(agent: ManagedAgent): Promise<void> {
agentId: string,
provider: AgentProvider,
cwd: string,
config?: SerializableAgentConfig
): Promise<void> {
await this.load(); await this.load();
const existing = this.cache.get(agentId); const existing = this.cache.get(agent.id);
if (!existing) { const record = toStoredAgentRecord(agent, {
console.warn( title: existing?.title ?? null,
`[AgentRegistry] Cannot record config for ${agentId} because no snapshot has been persisted yet` });
); this.cache.set(agent.id, record);
return;
}
const now = new Date().toISOString();
const sanitizedConfig = config ? sanitizeConfig(config) : existing.config;
const nextModeId =
config?.modeId ?? existing.lastModeId ?? sanitizedConfig?.modeId ?? null;
const updated: StoredAgentRecord = {
...existing,
provider,
cwd,
updatedAt: now,
lastModeId: nextModeId,
config: sanitizedConfig,
};
this.cache.set(agentId, updated);
await this.flush();
}
async applySnapshot(snapshot: AgentSnapshot): Promise<void> {
await this.load();
const now = new Date().toISOString();
const existing = this.cache.get(snapshot.id);
if (!existing) {
const record: StoredAgentRecord = {
id: snapshot.id,
provider: snapshot.provider,
cwd: snapshot.cwd,
createdAt: now,
updatedAt: now,
lastActivityAt: snapshot.updatedAt.toISOString(),
lastUserMessageAt: snapshot.lastUserMessageAt
? snapshot.lastUserMessageAt.toISOString()
: null,
title: null,
lastStatus: snapshot.status,
lastModeId: snapshot.currentModeId ?? null,
config: null,
persistence: snapshot.persistence ?? null,
};
this.cache.set(snapshot.id, record);
await this.flush();
return;
}
const updated: StoredAgentRecord = {
...existing,
provider: snapshot.provider,
cwd: snapshot.cwd,
updatedAt: now,
lastActivityAt: snapshot.updatedAt.toISOString(),
lastUserMessageAt: snapshot.lastUserMessageAt
? snapshot.lastUserMessageAt.toISOString()
: existing.lastUserMessageAt ?? null,
lastStatus: snapshot.status,
lastModeId: snapshot.currentModeId ?? null,
persistence: snapshot.persistence ?? existing.persistence ?? null,
};
this.cache.set(snapshot.id, updated);
await this.flush(); await this.flush();
} }
@@ -250,19 +187,6 @@ export class AgentRegistry {
} }
} }
function sanitizeConfig(
config: SerializableAgentConfig | undefined
): SerializableAgentConfig | undefined {
if (!config) {
return undefined;
}
const cleaned: SerializableAgentConfig = {};
if (config.modeId) cleaned.modeId = config.modeId;
if (config.model) cleaned.model = config.model;
if (config.extra) cleaned.extra = JSON.parse(JSON.stringify(config.extra));
return cleaned;
}
function resolveServerPackageRoot(): string { function resolveServerPackageRoot(): string {
let currentDir = path.dirname(fileURLToPath(import.meta.url)); let currentDir = path.dirname(fileURLToPath(import.meta.url));
while (true) { while (true) {
@@ -1,38 +1,85 @@
import { describe, expect, test, vi } from "vitest"; import { describe, expect, test, vi } from "vitest";
import type { AgentSnapshot } from "./agent/agent-manager.js"; import type { ManagedAgent } from "./agent/agent-manager.js";
import type { StoredAgentRecord } from "./agent/agent-registry.js"; import type { StoredAgentRecord } from "./agent/agent-registry.js";
import { import {
attachAgentRegistryPersistence, attachAgentRegistryPersistence,
restorePersistedAgents, restorePersistedAgents,
} from "./persistence-hooks.js"; } from "./persistence-hooks.js";
import type {
AgentPermissionRequest,
AgentSession,
AgentSessionConfig,
} from "./agent/agent-sdk-types.js";
function createSnapshot(overrides?: Partial<AgentSnapshot>): AgentSnapshot { type ManagedAgentOverrides = Omit<
const now = new Date(); Partial<ManagedAgent>,
return { "config" | "pendingPermissions" | "session" | "pendingRun"
id: "agent-1", > & {
provider: "claude", config?: Partial<AgentSessionConfig>;
cwd: "/tmp/project", pendingPermissions?: Map<string, AgentPermissionRequest>;
model: null, session?: AgentSession | null;
createdAt: now, pendingRun?: ManagedAgent["pendingRun"];
updatedAt: now, };
lastUserMessageAt: null,
status: "idle", function createManagedAgent(
sessionId: null, overrides: ManagedAgentOverrides = {}
capabilities: { ): ManagedAgent {
supportsStreaming: true, const now = overrides.updatedAt ?? new Date("2025-01-01T00:00:00.000Z");
supportsSessionPersistence: true, const provider = overrides.provider ?? "claude";
supportsDynamicModes: true, const cwd = overrides.cwd ?? "/tmp/project";
supportsMcpServers: true, const lifecycle = overrides.lifecycle ?? "idle";
supportsReasoningStream: true, const configOverrides = overrides.config ?? {};
supportsToolInvocations: true, const config: AgentSessionConfig = {
}, provider,
currentModeId: "plan", cwd,
availableModes: [], modeId: configOverrides.modeId ?? "plan",
pendingPermissions: [], model: configOverrides.model ?? "claude-3.5-sonnet",
persistence: null, extra: configOverrides.extra ?? { claude: { tone: "focused" } },
...overrides,
}; };
const session =
lifecycle === "closed"
? null
: overrides.session ?? ({} as AgentSession);
const pendingRun =
overrides.pendingRun ??
(lifecycle === "running" ? (async function* noop() {})() : null);
const agent: ManagedAgent = {
id: overrides.id ?? "agent-1",
provider,
cwd,
session,
sessionId: overrides.sessionId ?? "session-123",
capabilities:
overrides.capabilities ??
{
supportsStreaming: true,
supportsSessionPersistence: true,
supportsDynamicModes: true,
supportsMcpServers: true,
supportsReasoningStream: true,
supportsToolInvocations: true,
},
config,
lifecycle,
createdAt: overrides.createdAt ?? now,
updatedAt: overrides.updatedAt ?? now,
availableModes: overrides.availableModes ?? [],
currentModeId: overrides.currentModeId ?? config.modeId ?? null,
pendingPermissions:
overrides.pendingPermissions ??
new Map<string, AgentPermissionRequest>(),
pendingRun,
timeline: overrides.timeline ?? [],
persistence: overrides.persistence ?? null,
historyPrimed: overrides.historyPrimed ?? true,
lastUserMessageAt: overrides.lastUserMessageAt ?? now,
lastUsage: overrides.lastUsage,
lastError: overrides.lastError,
};
return agent;
} }
function createRecord( function createRecord(
@@ -138,13 +185,13 @@ describe("persistence hooks", () => {
} as any); } as any);
expect(agentManager.subscribe).toHaveBeenCalledTimes(1); expect(agentManager.subscribe).toHaveBeenCalledTimes(1);
const snapshot = createSnapshot(); const agent = createManagedAgent();
subscriber({ type: "agent_state", agent: snapshot }); subscriber({ type: "agent_state", agent });
expect(applySnapshot).toHaveBeenCalledWith(snapshot); expect(applySnapshot).toHaveBeenCalledWith(agent);
subscriber({ subscriber({
type: "agent_stream", type: "agent_stream",
agentId: snapshot.id, agentId: agent.id,
event: { type: "timeline", item: { type: "assistant_message", text: "hi" } }, event: { type: "timeline", item: { type: "assistant_message", text: "hi" } },
}); });
expect(applySnapshot).toHaveBeenCalledTimes(1); expect(applySnapshot).toHaveBeenCalledTimes(1);
-18
View File
@@ -1237,24 +1237,6 @@ export class Session {
); );
const snapshot = await this.agentManager.createAgent(sessionConfig); const snapshot = await this.agentManager.createAgent(sessionConfig);
this.setCachedTitle(snapshot.id, null); this.setCachedTitle(snapshot.id, null);
try {
await this.agentRegistry.recordConfig(
snapshot.id,
snapshot.provider,
snapshot.cwd,
{
modeId: sessionConfig.modeId,
model: sessionConfig.model,
extra: sessionConfig.extra,
}
);
} catch (registryError) {
console.error(
`[Session ${this.clientId}] Failed to record agent config for ${snapshot.id}:`,
registryError
);
}
await this.forwardAgentState(snapshot); await this.forwardAgentState(snapshot);
const trimmedPrompt = initialPrompt?.trim(); const trimmedPrompt = initialPrompt?.trim();