/** * Tests for AgentSession forking behavior. * * These tests verify: * - Forking from a single message works * - Forking in --no-session mode (in-memory only) * - getUserMessagesForForking returns correct entries */ import { existsSync, mkdirSync, rmSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { Agent } from "@mariozechner/companion-agent-core"; import { getModel } from "@mariozechner/companion-ai"; import { afterEach, beforeEach, describe, expect, it } from "vitest"; import { AgentSession } from "../src/core/agent-session.js"; import { AuthStorage } from "../src/core/auth-storage.js"; import { ModelRegistry } from "../src/core/model-registry.js"; import { SessionManager } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; import { codingTools } from "../src/core/tools/index.js"; import { API_KEY, createTestResourceLoader } from "./utilities.js"; describe.skipIf(!API_KEY)("AgentSession forking", () => { let session: AgentSession; let tempDir: string; let sessionManager: SessionManager; beforeEach(() => { // Create temp directory for session files tempDir = join(tmpdir(), `companion-branching-test-${Date.now()}`); mkdirSync(tempDir, { recursive: true }); }); afterEach(async () => { if (session) { session.dispose(); } if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } }); function createSession(noSession: boolean = false) { const model = getModel("anthropic", "claude-sonnet-4-5")!; const agent = new Agent({ getApiKey: () => API_KEY, initialState: { model, systemPrompt: "You are a helpful assistant. Be extremely concise, reply with just a few words.", tools: codingTools, }, }); sessionManager = noSession ? SessionManager.inMemory() : SessionManager.create(tempDir); const settingsManager = SettingsManager.create(tempDir, tempDir); const authStorage = AuthStorage.create(join(tempDir, "auth.json")); const modelRegistry = new ModelRegistry(authStorage, tempDir); session = new AgentSession({ agent, sessionManager, settingsManager, cwd: tempDir, modelRegistry, resourceLoader: createTestResourceLoader(), }); // Must subscribe to enable session persistence session.subscribe(() => {}); return session; } it("should allow forking from single message", async () => { createSession(); // Send one message await session.prompt("Say hello"); await session.agent.waitForIdle(); // Should have exactly 1 user message available for forking const userMessages = session.getUserMessagesForForking(); expect(userMessages.length).toBe(1); expect(userMessages[0].text).toBe("Say hello"); // Fork from the first message const result = await session.fork(userMessages[0].entryId); expect(result.selectedText).toBe("Say hello"); expect(result.cancelled).toBe(false); // After forking, conversation should be empty (forked before the first message) expect(session.messages.length).toBe(0); // Session file path should be set, but file is created lazily after first assistant message expect(session.sessionFile).not.toBeNull(); expect(existsSync(session.sessionFile!)).toBe(false); }); it("should support in-memory forking in --no-session mode", async () => { createSession(true); // Verify sessions are disabled expect(session.sessionFile).toBeUndefined(); // Send one message await session.prompt("Say hi"); await session.agent.waitForIdle(); // Should have 1 user message const userMessages = session.getUserMessagesForForking(); expect(userMessages.length).toBe(1); // Verify we have messages before forking expect(session.messages.length).toBeGreaterThan(0); // Fork from the first message const result = await session.fork(userMessages[0].entryId); expect(result.selectedText).toBe("Say hi"); expect(result.cancelled).toBe(false); // After forking, conversation should be empty expect(session.messages.length).toBe(0); // Session file should still be undefined (no file created) expect(session.sessionFile).toBeUndefined(); }); it("should fork from middle of conversation", async () => { createSession(); // Send multiple messages await session.prompt("Say one"); await session.agent.waitForIdle(); await session.prompt("Say two"); await session.agent.waitForIdle(); await session.prompt("Say three"); await session.agent.waitForIdle(); // Should have 3 user messages const userMessages = session.getUserMessagesForForking(); expect(userMessages.length).toBe(3); // Fork from second message (keeps first message + response) const secondMessage = userMessages[1]; const result = await session.fork(secondMessage.entryId); expect(result.selectedText).toBe("Say two"); // After forking, should have first user message + assistant response expect(session.messages.length).toBe(2); expect(session.messages[0].role).toBe("user"); expect(session.messages[1].role).toBe("assistant"); }); });