bkinnightskytw
/
agent-orchestrator
mirror of https://github.com/ComposioHQ/agent-orchestrator.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix native restore fallback for Claude and Codex sessions (#1602) 

* Fix native session restore fallback for Claude and Codex

* Address restore metadata review comments

* Fix metadata normalization lint

* Address PR metadata review feedback

* Prevent fresh restore fallback for native agents
This commit is contained in:
Ashish Huddar 2026-05-01 21:27:17 +05:30 committed by GitHub
parent 2306078761
commit ab65d12356
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 261 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

139
packages/core/src/__tests__/session-manager/restore.test.ts
View File

@ -547,6 +547,145 @@ describe("restore", () => {
expect(mockAgent.getLaunchCommand).toHaveBeenCalled();
const createCall = (mockRuntime.create as ReturnType<typeof vi.fn>).mock.calls[0][0];
expect(createCall.launchCommand).toBe("mock-agent --start");
const meta = readMetadataRaw(sessionsDir, "app-1");
expect(meta!["restoreFallbackReason"]).toBe("mock-agent.getRestoreCommand returned null");
});
it("does not launch a fresh chat when a native-restore agent cannot build restore command", async () => {
const wsPath = join(tmpDir, "ws-app-native-restore-missing");
mkdirSync(wsPath, { recursive: true });
const mockNativeRestoreAgent: Agent = {
...mockAgent,
name: "codex",
getRestoreCommand: vi.fn().mockResolvedValue(null),
};
const registryWithNativeRestoreAgent: PluginRegistry = {
...mockRegistry,
get: vi.fn().mockImplementation((slot: string) => {
if (slot === "runtime") return mockRuntime;
if (slot === "agent") return mockNativeRestoreAgent;
if (slot === "workspace") return mockWorkspace;
return null;
}),
};
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-1",
status: "killed",
project: "my-app",
runtimeHandle: makeHandle("rt-old"),
});
const sm = createSessionManager({ config, registry: registryWithNativeRestoreAgent });
await expect(sm.restore("app-1")).rejects.toThrow(SessionNotRestorableError);
expect(mockRuntime.create).not.toHaveBeenCalled();
const meta = readMetadataRaw(sessionsDir, "app-1");
expect(meta!["restoreFallbackReason"]).toBe("codex.getRestoreCommand returned null");
});
it("clears restore fallback reason when getRestoreCommand succeeds", async () => {
const wsPath = join(tmpDir, "ws-app-restore-clears-fallback");
mkdirSync(wsPath, { recursive: true });
const mockAgentWithRestore: Agent = {
...mockAgent,
getRestoreCommand: vi.fn().mockResolvedValue("claude --resume abc123"),
};
const registryWithAgentRestore: PluginRegistry = {
...mockRegistry,
get: vi.fn().mockImplementation((slot: string) => {
if (slot === "runtime") return mockRuntime;
if (slot === "agent") return mockAgentWithRestore;
if (slot === "workspace") return mockWorkspace;
return null;
}),
};
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-1",
status: "killed",
project: "my-app",
runtimeHandle: makeHandle("rt-old"),
restoreFallbackReason: "previous fallback",
});
const sm = createSessionManager({ config, registry: registryWithAgentRestore });
await sm.restore("app-1");
const meta = readMetadataRaw(sessionsDir, "app-1");
expect(meta!["restoreFallbackReason"]).toBeUndefined();
});
it("normalizes agent metadata empty strings in memory like metadata persistence", async () => {
const wsPath = join(tmpDir, "ws-app-agent-metadata-normalize");
mkdirSync(wsPath, { recursive: true });
const mockAgentWithMetadata: Agent = {
...mockAgent,
getSessionInfo: vi.fn().mockResolvedValue({
summary: null,
agentSessionId: "native-1",
metadata: {
codexThreadId: "thread-1",
restoreFallbackReason: "",
},
}),
};
const registryWithAgentMetadata: PluginRegistry = {
...mockRegistry,
get: vi.fn().mockImplementation((slot: string) => {
if (slot === "runtime") return mockRuntime;
if (slot === "agent") return mockAgentWithMetadata;
if (slot === "workspace") return mockWorkspace;
return null;
}),
};
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-1",
status: "killed",
project: "my-app",
runtimeHandle: makeHandle("rt-old"),
restoreFallbackReason: "previous fallback",
});
const sm = createSessionManager({ config, registry: registryWithAgentMetadata });
const restored = await sm.restore("app-1");
expect(restored.metadata["codexThreadId"]).toBe("thread-1");
expect(restored.metadata["restoreFallbackReason"]).toBeUndefined();
const meta = readMetadataRaw(sessionsDir, "app-1");
expect(meta!["codexThreadId"]).toBe("thread-1");
expect(meta!["restoreFallbackReason"]).toBeUndefined();
});
it("clears restore fallback reason when agent has no restore command", async () => {
const wsPath = join(tmpDir, "ws-app-no-restore-clears-fallback");
mkdirSync(wsPath, { recursive: true });
writeMetadata(sessionsDir, "app-1", {
worktree: wsPath,
branch: "feat/TEST-1",
status: "killed",
project: "my-app",
runtimeHandle: makeHandle("rt-old"),
restoreFallbackReason: "previous fallback",
});
const sm = createSessionManager({ config, registry: mockRegistry });
await sm.restore("app-1");
const meta = readMetadataRaw(sessionsDir, "app-1");
expect(meta!["restoreFallbackReason"]).toBeUndefined();
});
it("does not inject OPENCODE_CONFIG when restoring OpenCode orchestrators", async () => {

2
packages/core/src/metadata.ts
View File

@ -299,7 +299,7 @@ export function updateMetadata(
}, { createIfMissing: true });
}
function applyMetadataUpdates(
export function applyMetadataUpdates(
existing: Record<string, string>,
updates: Partial<Record<string, string>>,
): Record<string, string> {

51
packages/core/src/session-manager.ts
View File

@ -53,6 +53,7 @@ import {
readMetadataRaw,
writeMetadata,
updateMetadata,
applyMetadataUpdates,
mutateMetadata,
deleteMetadata,
listMetadata,
@ -428,6 +429,15 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM
return isOrchestratorSessionRecord(sessionId, metadata ?? {}, project.sessionPrefix);
}
function requiresNativeRestore(agentName: string): boolean {
return (
agentName === "claude-code" ||
agentName === "codex" ||
agentName === "kimicode" ||
agentName === "opencode"
);
}
function applyMetadataUpdatesToRaw(
raw: Record<string, string>,
updates: Partial<Record<string, string>>,
@ -972,7 +982,7 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM
data: {},
};
}
await enrichSessionWithRuntimeState(session, plugins, handleFromMetadata);
await enrichSessionWithRuntimeState(session, plugins, handleFromMetadata, sessionsDir);
}
/**
@ -985,6 +995,7 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM
session: Session,
plugins: ReturnType<typeof resolvePlugins>,
handleFromMetadata: boolean,
sessionsDir: string,
): Promise<void> {
// Check runtime liveness first — for all statuses except "spawning".
// Skip spawning sessions because tmux may not be fully initialized yet,
@ -1058,13 +1069,26 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM
}
// Enrich with live agent session info (summary, cost).
let info: Awaited<ReturnType<Agent["getSessionInfo"]>>;
try {
const info = await plugins.agent.getSessionInfo(session);
if (info) {
session.agentInfo = info;
}
info = await plugins.agent.getSessionInfo(session);
} catch {
// Can't get session info — keep existing values
info = null;
}
if (info) {
session.agentInfo = info;
const metadataUpdates = info.metadata ?? {};
if (Object.keys(metadataUpdates).length > 0) {
try {
updateMetadata(sessionsDir, session.id, metadataUpdates);
session.metadata = applyMetadataUpdates(session.metadata, metadataUpdates);
invalidateCache();
} catch {
// Persisting agent metadata is best-effort; keep live agent info.
}
}
}
}
}
@ -2811,7 +2835,7 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM
// and isRestorable would reject it.
const session = metadataToSession(sessionId, raw, projectId, project.sessionPrefix);
const plugins = resolvePlugins(project, selection.agentName);
await enrichSessionWithRuntimeState(session, plugins, true);
await enrichSessionWithRuntimeState(session, plugins, true, sessionsDir);
// 3. Validate restorability
if (!isRestorable(session)) {
@ -2924,9 +2948,22 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM
if (plugins.agent.getRestoreCommand) {
const restoreCmd = await plugins.agent.getRestoreCommand(session, projectConfigForLaunch);
launchCommand = restoreCmd ?? plugins.agent.getLaunchCommand(agentLaunchConfig);
if (restoreCmd) {
launchCommand = restoreCmd;
updateMetadata(sessionsDir, sessionId, { restoreFallbackReason: "" });
} else {
const reason = `${plugins.agent.name}.getRestoreCommand returned null`;
updateMetadata(sessionsDir, sessionId, {
restoreFallbackReason: reason,
});
if (requiresNativeRestore(plugins.agent.name)) {
throw new SessionNotRestorableError(sessionId, reason);
}
launchCommand = plugins.agent.getLaunchCommand(agentLaunchConfig);
}
} else {
launchCommand = plugins.agent.getLaunchCommand(agentLaunchConfig);
updateMetadata(sessionsDir, sessionId, { restoreFallbackReason: "" });
}
const environment = plugins.agent.getEnvironment(agentLaunchConfig);

6
packages/core/src/types.ts
View File

@ -601,6 +601,8 @@ export interface AgentSessionInfo {
summaryIsFallback?: boolean;
/** Agent's internal session ID (for resume) */
agentSessionId: string | null;
/** Agent-owned metadata worth persisting for later restore. */
metadata?: Record<string, string>;
/** Estimated cost so far */
cost?: CostEstimate;
}
@ -1695,6 +1697,10 @@ export interface SessionMetadata {
directTerminalWsPort?: number;
};
opencodeSessionId?: string;
claudeSessionUuid?: string;
codexThreadId?: string;
codexModel?: string;
restoreFallbackReason?: string;
pinnedSummary?: string; // First quality summary, pinned for display stability
userPrompt?: string; // Prompt used when spawning without a tracker issue
/**

22
packages/plugins/agent-claude-code/src/index.test.ts
View File

@ -566,6 +566,28 @@ describe("getSessionInfo", () => {
mockJsonlFiles('{"type":"user","message":{"content":"hi"}}', ["abc-def-123.jsonl"]);
const result = await agent.getSessionInfo(makeSession());
expect(result?.agentSessionId).toBe("abc-def-123");
expect(result?.metadata?.claudeSessionUuid).toBe("abc-def-123");
});
});
describe("getRestoreCommand metadata", () => {
it("uses persisted Claude session UUID without scanning project files", async () => {
const agent = create();
const session = makeSession({
workspacePath: "/workspace/test-project",
metadata: { claudeSessionUuid: "persisted-uuid" },
});
const command = await agent.getRestoreCommand!(session, {
name: "test-project",
repo: "owner/repo",
path: "/workspace/test-project",
defaultBranch: "main",
sessionPrefix: "test",
});
expect(command).toBe("claude --resume 'persisted-uuid'");
expect(mockReaddir).not.toHaveBeenCalled();
});
});

22
packages/plugins/agent-claude-code/src/index.ts
View File

@ -832,23 +832,27 @@ function createClaudeCodeAgent(): Agent {
summary: summaryResult?.summary ?? null,
summaryIsFallback: summaryResult?.isFallback,
agentSessionId,
metadata: { claudeSessionUuid: agentSessionId },
cost: extractCost(lines),
};
},
async getRestoreCommand(session: Session, project: ProjectConfig): Promise<string | null> {
if (!session.workspacePath) return null;
let sessionUuid = session.metadata?.["claudeSessionUuid"]?.trim();
if (!sessionUuid) {
if (!session.workspacePath) return null;
// Find Claude's project directory for this workspace
const projectPath = toClaudeProjectPath(session.workspacePath);
const projectDir = join(homedir(), ".claude", "projects", projectPath);
// Find Claude's project directory for this workspace
const projectPath = toClaudeProjectPath(session.workspacePath);
const projectDir = join(homedir(), ".claude", "projects", projectPath);
// Find the latest session JSONL file
const sessionFile = await findLatestSessionFile(projectDir);
if (!sessionFile) return null;
// Find the latest session JSONL file
const sessionFile = await findLatestSessionFile(projectDir);
if (!sessionFile) return null;
// Extract session UUID from filename (e.g. "abc123-def456.jsonl" → "abc123-def456")
const sessionUuid = basename(sessionFile, ".jsonl");
// Extract session UUID from filename (e.g. "abc123-def456.jsonl" → "abc123-def456")
sessionUuid = basename(sessionFile, ".jsonl");
}
if (!sessionUuid) return null;
// Build resume command

14
packages/plugins/agent-codex/src/index.test.ts
View File

@ -1328,6 +1328,20 @@ describe("getRestoreCommand", () => {
expect(cmd).toContain("thread-abc-123");
});
it("uses persisted Codex thread ID without scanning session files", async () => {
const session = makeSession({
workspacePath: "/workspace/test",
metadata: { codexThreadId: "persisted-thread", codexModel: "gpt-5.3-codex" },
});
const cmd = await agent.getRestoreCommand!(session, makeProjectConfig());
expect(cmd).toContain("'codex' resume");
expect(cmd).toContain("--model 'gpt-5.3-codex'");
expect(cmd).toContain("persisted-thread");
expect(mockReaddir).not.toHaveBeenCalled();
});
it("builds native resume command from payload-wrapped Codex session id", async () => {
const content = jsonl(
{

32
packages/plugins/agent-codex/src/index.ts
View File

@ -703,21 +703,33 @@ function createCodexAgent(): Agent {
summary: data.model ? `Codex session (${data.model})` : null,
summaryIsFallback: true,
agentSessionId,
metadata: data.threadId
? {
codexThreadId: data.threadId,
...(data.model ? { codexModel: data.model } : {}),
}
: undefined,
cost,
};
},
async getRestoreCommand(session: Session, project: ProjectConfig): Promise<string | null> {
if (!session.workspacePath) return null;
let threadId = session.metadata?.["codexThreadId"]?.trim();
let model: string | null = session.metadata?.["codexModel"]?.trim() || null;
if (!threadId) {
if (!session.workspacePath) return null;
// Find the Codex session file for this workspace
const sessionFile = await findCodexSessionFileCached(session.workspacePath);
if (!sessionFile) return null;
// Find the Codex session file for this workspace
const sessionFile = await findCodexSessionFileCached(session.workspacePath);
if (!sessionFile) return null;
// Stream the file line-by-line to avoid loading potentially huge
// rollout files (100 MB+) entirely into memory.
const data = await streamCodexSessionData(sessionFile);
if (!data?.threadId) return null;
// Stream the file line-by-line to avoid loading potentially huge
// rollout files (100 MB+) entirely into memory.
const data = await streamCodexSessionData(sessionFile);
if (!data?.threadId) return null;
threadId = data.threadId;
model = data.model;
}
// Use Codex's native `resume` subcommand for proper conversation resume.
// This restores the full thread state, not just a text prompt re-injection.
@ -727,11 +739,11 @@ function createCodexAgent(): Agent {
appendNoUpdateCheckFlag(parts);
appendApprovalFlags(parts, project.agentConfig?.permissions);
const effectiveModel = (project.agentConfig?.model ?? data.model) as string | undefined;
const effectiveModel = (project.agentConfig?.model ?? model) as string | undefined;
appendModelFlags(parts, effectiveModel ?? undefined);
// Positional threadId goes last, after all flags
parts.push(shellEscape(data.threadId));
parts.push(shellEscape(threadId));
return parts.join(" ");
},