Fix Codex continuation after workspace archive (#4512)

* Fix native writer ownership when restoring archived agents

* Keep archive QA mechanics in a scenario helper
This commit is contained in:
Mohamed Boudra
2026-09-09 08:51:55 +02:00
committed by GitHub
parent da8c1b5c94
commit 613cbbe9ef
3 changed files with 160 additions and 2 deletions
@@ -7765,6 +7765,53 @@ test("fires onAgentArchived for stored-only snapshot archives", async () => {
expect(archivedIds).toEqual([storedOnly.id]);
});
test("native archive and restore release the loaded session writer", async () => {
const workdir = mkdtempSync(join(tmpdir(), "agent-native-archive-writer-"));
class WriterClient extends NativeArchiveRecordingClient {
session: CloseRecordingTestAgentSession | undefined;
override async createSession(config: AgentSessionConfig): Promise<AgentSession> {
this.session = new CloseRecordingTestAgentSession(config);
return this.session;
}
override async resumeSession(
_handle: AgentPersistenceHandle,
config?: Partial<AgentSessionConfig>,
): Promise<AgentSession> {
return this.createSession({ provider: "codex", cwd: workdir, ...config });
}
override async archiveNativeSession(handle: AgentPersistenceHandle): Promise<void> {
if (!this.session?.closed) throw new Error("native thread has an active writer");
await super.archiveNativeSession(handle);
}
override async unarchiveNativeSession(handle: AgentPersistenceHandle): Promise<void> {
if (!this.session?.closed) throw new Error("native thread has an active writer");
await super.unarchiveNativeSession(handle);
}
}
const client = new WriterClient();
const storage = new AgentStorage(join(workdir, "agents"), logger);
const manager = new AgentManager({ clients: { codex: client }, registry: storage, logger });
try {
const agent = await manager.createAgent({ provider: "codex", cwd: workdir }, undefined, {
workspaceId: undefined,
});
await manager.archiveAgent(agent.id);
expect(client.archivedHandles).toHaveLength(1);
// Opening archived history can retain a writer, including records archived by older daemons.
await ensureAgentLoaded(agent.id, { agentManager: manager, agentStorage: storage, logger });
expect(client.session?.closed).toBe(false);
await manager.unarchiveSnapshot(agent.id);
expect(client.unarchivedHandles).toHaveLength(1);
expect((await storage.get(agent.id))?.archivedAt).toBeNull();
expect(client.session?.closed).toBe(true);
await ensureAgentLoaded(agent.id, { agentManager: manager, agentStorage: storage, logger });
expect(client.session?.closed).toBe(false);
} finally {
for (const agent of manager.listAgents()) await manager.closeAgent(agent.id);
rmSync(workdir, { recursive: true, force: true });
}
});
test("unarchiveSnapshot skips native provider unarchive for active records", async () => {
const workdir = mkdtempSync(join(tmpdir(), "agent-manager-unarchive-active-"));
const storagePath = join(workdir, "agents");
@@ -1703,6 +1703,7 @@ export class AgentManager {
const { archivedAt } = await this.markRecordArchived(stored);
agent.updatedAt = new Date(archivedAt);
await this.closeAgentRuntime(agentId);
await this.syncNativeArchiveState(stored.provider, stored.persistence, "archive");
this.discardRetainedAgentState(agentId);
await this.cascadeArchiveChildren(agentId);
@@ -1762,8 +1763,6 @@ export class AgentManager {
updatedAt: archivedAt,
});
await this.syncNativeArchiveState(record.provider, record.persistence, "archive");
if (this.agents.has(record.id)) {
this.notifyAgentState(record.id);
} else if (!archivedRecord.internal) {
@@ -2090,6 +2089,8 @@ export class AgentManager {
return false;
}
// Archived history may have loaded a runtime that still owns the native writer.
await this.closeAgent(agentId);
await this.syncNativeArchiveState(record.provider, record.persistence, "restore");
await registry.upsert({
@@ -0,0 +1,110 @@
import { execFileSync } from "node:child_process";
import { existsSync, mkdirSync, mkdtempSync, rmSync } from "node:fs";
import { tmpdir } from "node:os";
import path from "node:path";
import pino from "pino";
import { expect, test } from "vitest";
import { CodexAppServerAgentClient } from "../agent/providers/codex-app-server-agent.js";
import { DaemonClient } from "../test-utils/daemon-client.js";
import { createTestPaseoDaemon, type TestPaseoDaemon } from "../test-utils/paseo-daemon.js";
async function expectContextAfterWorkspaceRestore(legacyNativeArchive: boolean): Promise<void> {
const root = mkdtempSync(path.join(tmpdir(), "paseo-archive-codex-"));
const cwd = path.join(root, "repo");
mkdirSync(cwd);
execFileSync("git", ["init", "-b", "main", cwd]);
execFileSync(
"git",
[
"-c",
"user.name=QA",
"-c",
"user.email=qa@example.invalid",
"-c",
"commit.gpgsign=false",
"commit",
"--allow-empty",
"-m",
"QA",
],
{ cwd },
);
const logger = pino({ level: "warn" });
const provider = new CodexAppServerAgentClient(logger);
let daemon: TestPaseoDaemon | undefined;
let client: DaemonClient | undefined;
try {
daemon = await createTestPaseoDaemon({
agentClients: { codex: provider },
logger,
pluginsEnabled: false,
});
client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws`, appVersion: "0.8.0" });
await client.connect();
await client.fetchAgents({ subscribe: { subscriptionId: "archive-qa" } });
const agent = await client.createAgent({
config: {
provider: "codex",
cwd,
model: "gpt-5.6-sol",
modeId: "full-access",
thinkingOptionId: "low",
},
worktree: { mode: "branch-off", base: "main", newBranch: "archive-qa" },
});
const manager = daemon.daemon.agentManager;
const token = "ARCHIVE_CEDAR_7429";
const first = await manager.runAgent(
agent.id,
`Remember ${token}. Reply with only that token. Do not use tools.`,
);
expect(first.finalText).toContain(token);
const handle = manager.getAgent(agent.id)!.persistence!;
await client.archiveWorkspace(agent.workspaceId!);
expect(existsSync(agent.cwd)).toBe(false);
// Older daemons marked Paseo archived while native archive failed against its writer.
if (legacyNativeArchive) await provider.unarchiveNativeSession(handle);
await client.restoreWorkspace(agent.workspaceId!);
expect(existsSync(agent.cwd)).toBe(true);
await client.fetchAgentTimeline(agent.id, { direction: "tail", limit: 0 });
await client.sendAgentMessage(
agent.id,
"What token did I ask you to remember? Reply with RECALLED=<token>. Do not use tools.",
{ activeTurnBehavior: "steer" },
);
const connectedClient = client;
await expect
.poll(
async () => {
const timeline = await connectedClient.fetchAgentTimeline(agent.id, {
direction: "tail",
limit: 0,
projection: "canonical",
});
return timeline.entries
.map(({ item }) => (item.type === "assistant_message" ? item.text : ""))
.join("");
},
{ timeout: 30_000 },
)
.toContain(`RECALLED=${token}`);
expect(manager.getAgent(agent.id)?.persistence?.sessionId).toBe(handle.sessionId);
process.stdout.write(
`archived -> restored; same agent=${agent.id}; same session=${handle.sessionId}; two real replies recall ${token}; legacy=${legacyNativeArchive}\n`,
);
await client.archiveWorkspace(agent.workspaceId!);
} finally {
await client?.close();
await daemon?.close();
rmSync(root, { recursive: true, force: true });
}
}
// Opt in: real Codex requests using local authentication.
test.runIf(process.env.PASEO_NATIVE_ARCHIVE_QA === "1").each([false, true])(
"Codex recalls context after workspace restore (legacy native archive: %s)",
async (legacyNativeArchive) => {
await expectContextAfterWorkspaceRestore(legacyNativeArchive);
},
90_000,
);