Archived
feat: add hierarchical session tree navigator
This commit is contained in:
Federico Jaramillo Martinez
@@ -0,0 +1,395 @@
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import type { SessionTreeNavigateRequest, SessionTreeSummaryChoice } from "../../shared/apiTypes.js";
|
||||
import { PiSessionService, type PiAgentSession, type PiSessionManager, type PiSessionServiceDependencies } from "./piSessionService.js";
|
||||
import { CapturingSessionEventHub, emptyArchiveStore, fakeRuntime, fakeSessionManager, runtimeCreator, sessionGateway, sessionRecord, sessionRef, testModel, testModelRuntime, type TestSession } from "./piSessionService.testSupport.js";
|
||||
import type { ProjectableSessionTreeNode } from "./sessionTreeProjection.js";
|
||||
|
||||
const TEST_AGENT_DIR = "/tmp/pi-web-test-agent";
|
||||
const SESSION_ID = "tree-session";
|
||||
|
||||
type NavigateTree = NonNullable<PiAgentSession["navigateTree"]>;
|
||||
|
||||
function deferred<T>() {
|
||||
let resolve!: (value: T | PromiseLike<T>) => void;
|
||||
let reject!: (reason?: unknown) => void;
|
||||
const promise = new Promise<T>((promiseResolve, promiseReject) => {
|
||||
resolve = promiseResolve;
|
||||
reject = promiseReject;
|
||||
});
|
||||
return { promise, resolve, reject };
|
||||
}
|
||||
|
||||
function treeNode(entry: Record<string, unknown>, children: ProjectableSessionTreeNode[] = []): ProjectableSessionTreeNode {
|
||||
return { entry, children };
|
||||
}
|
||||
|
||||
function navigationRequest(
|
||||
summary: SessionTreeSummaryChoice = { mode: "none" },
|
||||
expectedLeafId: string | null = "leaf-1",
|
||||
): SessionTreeNavigateRequest {
|
||||
return { targetId: "target-1", expectedLeafId, summary };
|
||||
}
|
||||
|
||||
function treeHarness(
|
||||
managerPatch: Partial<PiSessionManager> = {},
|
||||
sessionPatch: Partial<TestSession> = {},
|
||||
dependenciesPatch: Partial<PiSessionServiceDependencies> = {},
|
||||
) {
|
||||
const hub = new CapturingSessionEventHub();
|
||||
const manager = fakeSessionManager("/workspace", {
|
||||
getSessionId: () => SESSION_ID,
|
||||
getLeafId: () => "leaf-1",
|
||||
...managerPatch,
|
||||
});
|
||||
const fake = fakeRuntime(SESSION_ID, { sessionManager: manager, ...sessionPatch });
|
||||
const service = new PiSessionService(hub, {
|
||||
agentDir: TEST_AGENT_DIR,
|
||||
modelRuntime: testModelRuntime,
|
||||
archiveStore: emptyArchiveStore(),
|
||||
createAgentRuntime: runtimeCreator(fake.runtime),
|
||||
sessionManager: sessionGateway([sessionRecord(SESSION_ID)]),
|
||||
heartbeatIntervalMs: 60_000,
|
||||
...dependenciesPatch,
|
||||
});
|
||||
return { service, fake, hub };
|
||||
}
|
||||
|
||||
describe("PiSessionService session-tree behavior", () => {
|
||||
it("opens /tree from the live manager through the safe projection boundary", async () => {
|
||||
const navigateTree = vi.fn<NavigateTree>(() => Promise.resolve({ cancelled: false }));
|
||||
const roots = [treeNode({
|
||||
type: "message",
|
||||
id: "leaf-1",
|
||||
parentId: null,
|
||||
timestamp: "2026-01-01T00:00:00.000Z",
|
||||
message: {
|
||||
role: "assistant",
|
||||
content: [
|
||||
{ type: "text", text: "safe answer" },
|
||||
{ type: "thinking", thinking: "private reasoning", thinkingSignature: "private signature" },
|
||||
],
|
||||
usage: { private: true },
|
||||
},
|
||||
})];
|
||||
const { service } = treeHarness({ getTree: () => roots }, { navigateTree });
|
||||
|
||||
await expect(service.runCommand(sessionRef(SESSION_ID), "/tree")).resolves.toEqual({
|
||||
type: "tree",
|
||||
tree: {
|
||||
nodes: [{
|
||||
id: "leaf-1",
|
||||
parentId: null,
|
||||
kind: "assistant",
|
||||
summary: "safe answer",
|
||||
timestamp: "2026-01-01T00:00:00.000Z",
|
||||
}],
|
||||
activeLeafId: "leaf-1",
|
||||
activePathIds: ["leaf-1"],
|
||||
},
|
||||
});
|
||||
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("maps none, default, and trimmed custom summary choices exactly and returns only editor text", async () => {
|
||||
const navigateTree = vi.fn<NavigateTree>();
|
||||
navigateTree
|
||||
.mockResolvedValueOnce({ cancelled: false })
|
||||
.mockResolvedValueOnce({ cancelled: false })
|
||||
.mockResolvedValueOnce({ cancelled: false, editorText: "exact user text", summaryEntry: { details: "must not escape" } });
|
||||
const { service, fake } = treeHarness({}, { navigateTree });
|
||||
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "none" }))).resolves.toEqual({ cancelled: false });
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "default" }))).resolves.toEqual({ cancelled: false });
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "custom", instructions: " focus on tests\nwithout losing context " }))).resolves.toEqual({
|
||||
cancelled: false,
|
||||
editorText: "exact user text",
|
||||
});
|
||||
|
||||
expect(navigateTree).toHaveBeenNthCalledWith(1, "target-1", { summarize: false });
|
||||
expect(navigateTree).toHaveBeenNthCalledWith(2, "target-1", { summarize: true });
|
||||
expect(navigateTree).toHaveBeenNthCalledWith(3, "target-1", { summarize: true, customInstructions: "focus on tests\nwithout losing context" });
|
||||
expect(service.activeCount()).toBe(1);
|
||||
expect(fake.calls.dispose).toBe(0);
|
||||
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("validates stale leaves, active work, unavailable runtimes, and custom instruction bounds", async () => {
|
||||
const navigateTree = vi.fn<NavigateTree>(() => Promise.resolve({ cancelled: false }));
|
||||
const { service, fake } = treeHarness({ getLeafId: () => "new-leaf" }, { navigateTree });
|
||||
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "none" }, "old-leaf"))).rejects.toThrow(
|
||||
"The session changed since /tree was opened. Reopen /tree and try again.",
|
||||
);
|
||||
expect(navigateTree).not.toHaveBeenCalled();
|
||||
|
||||
fake.session.isStreaming = true;
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "none" }, "new-leaf"))).rejects.toThrow(
|
||||
"Stop current session activity before navigating the session tree",
|
||||
);
|
||||
fake.session.isStreaming = false;
|
||||
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "custom", instructions: " " }, "new-leaf"))).rejects.toThrow(
|
||||
"Custom branch-summary instructions are required",
|
||||
);
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "custom", instructions: "x".repeat(10_001) }, "new-leaf"))).rejects.toThrow(
|
||||
"Custom branch-summary instructions must be at most 10000 characters",
|
||||
);
|
||||
expect(navigateTree).not.toHaveBeenCalled();
|
||||
await service.dispose();
|
||||
|
||||
const unavailable = treeHarness();
|
||||
await expect(unavailable.service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).rejects.toThrow(
|
||||
"Session tree navigation is not supported by this Pi runtime",
|
||||
);
|
||||
await unavailable.service.dispose();
|
||||
});
|
||||
|
||||
it("holds a per-runtime gate that rejects concurrent navigation and leaf-producing work", async () => {
|
||||
const navigation = deferred<Awaited<ReturnType<NavigateTree>>>();
|
||||
const navigateTree = vi.fn<NavigateTree>(() => navigation.promise);
|
||||
const { service, fake } = treeHarness({}, { navigateTree });
|
||||
|
||||
const firstNavigation = service.navigateTree(sessionRef(SESSION_ID), navigationRequest());
|
||||
await vi.waitFor(() => { expect(navigateTree).toHaveBeenCalledOnce(); });
|
||||
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).rejects.toThrow(
|
||||
"Stop current session activity before navigating the session tree",
|
||||
);
|
||||
await expect(service.prompt(sessionRef(SESSION_ID), "do not append yet")).rejects.toThrow(
|
||||
"Cannot send a prompt while session tree navigation is active",
|
||||
);
|
||||
await expect(service.shell(sessionRef(SESSION_ID), "!pwd")).rejects.toThrow(
|
||||
"Cannot run a shell command while session tree navigation is active",
|
||||
);
|
||||
await expect(service.setThinkingLevel(sessionRef(SESSION_ID), "off")).rejects.toThrow(
|
||||
"Cannot change the thinking level while session tree navigation is active",
|
||||
);
|
||||
const model = testModel();
|
||||
await expect(service.setModel(sessionRef(SESSION_ID), model.provider, model.id)).rejects.toThrow(
|
||||
"Cannot change models while session tree navigation is active",
|
||||
);
|
||||
await expect(service.cycleModel(sessionRef(SESSION_ID), "forward")).rejects.toThrow(
|
||||
"Cannot change models while session tree navigation is active",
|
||||
);
|
||||
await expect(service.runCommand(sessionRef(SESSION_ID), "/name blocked")).resolves.toEqual({
|
||||
type: "unsupported",
|
||||
message: "Cannot run commands while session tree navigation is active. Stop or finish the navigation first.",
|
||||
});
|
||||
await expect(service.archive(sessionRef(SESSION_ID))).rejects.toThrow("Stop current session activity before archiving");
|
||||
expect(fake.calls.prompt).toEqual([]);
|
||||
|
||||
navigation.resolve({ cancelled: false });
|
||||
await expect(firstNavigation).resolves.toEqual({ cancelled: false });
|
||||
await service.prompt(sessionRef(SESSION_ID), "now append");
|
||||
expect(fake.calls.prompt).toEqual([{ text: "now append", options: undefined }]);
|
||||
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("blocks navigation while prompt preflight can await before Pi reports streaming", async () => {
|
||||
const promptOperation = deferred<undefined>();
|
||||
const prompt = vi.fn(() => promptOperation.promise);
|
||||
const navigateTree = vi.fn<NavigateTree>(() => Promise.resolve({ cancelled: false }));
|
||||
const { service } = treeHarness({}, { prompt, navigateTree });
|
||||
|
||||
await service.prompt(sessionRef(SESSION_ID), "preflight is still running");
|
||||
await vi.waitFor(() => { expect(prompt).toHaveBeenCalledOnce(); });
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).rejects.toThrow(
|
||||
"Stop current session activity before navigating the session tree",
|
||||
);
|
||||
|
||||
promptOperation.resolve(undefined);
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).resolves.toEqual({ cancelled: false });
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("blocks navigation while an asynchronous model change can still append an entry", async () => {
|
||||
const modelChange = deferred<undefined>();
|
||||
const setModel = vi.fn(() => modelChange.promise);
|
||||
const navigateTree = vi.fn<NavigateTree>(() => Promise.resolve({ cancelled: false }));
|
||||
const { service } = treeHarness({}, { setModel, navigateTree });
|
||||
const model = testModel();
|
||||
|
||||
const changingModel = service.setModel(sessionRef(SESSION_ID), model.provider, model.id);
|
||||
await vi.waitFor(() => { expect(setModel).toHaveBeenCalledOnce(); });
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).rejects.toThrow(
|
||||
"Stop current session activity before navigating the session tree",
|
||||
);
|
||||
|
||||
modelChange.resolve(undefined);
|
||||
await changingModel;
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).resolves.toEqual({ cancelled: false });
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("blocks tree navigation while clone replaces and rebinds the runtime", async () => {
|
||||
const replacement = deferred<{ cancelled: boolean; selectedText?: string }>();
|
||||
const rebound = deferred<undefined>();
|
||||
const navigateTree = vi.fn<NavigateTree>(() => Promise.resolve({ cancelled: false }));
|
||||
const { service, fake } = treeHarness({}, { navigateTree });
|
||||
const replacementSessionId = "tree-session-replacement";
|
||||
const replacementFake = fakeRuntime(replacementSessionId, {
|
||||
sessionManager: fakeSessionManager("/workspace", {
|
||||
getSessionId: () => replacementSessionId,
|
||||
getLeafId: () => "replacement-leaf",
|
||||
}),
|
||||
navigateTree,
|
||||
});
|
||||
let rebindSession: ((session: PiAgentSession) => Promise<void>) | undefined;
|
||||
fake.runtime.setRebindSession = (callback) => { rebindSession = callback; };
|
||||
const fork = vi.fn(async () => {
|
||||
if (!Reflect.set(fake.runtime, "session", replacementFake.session)) throw new Error("Could not replace fake runtime session");
|
||||
await rebindSession?.(replacementFake.session);
|
||||
rebound.resolve(undefined);
|
||||
return replacement.promise;
|
||||
});
|
||||
fake.runtime.fork = fork;
|
||||
|
||||
const cloning = service.runCommand(sessionRef(SESSION_ID), "/clone");
|
||||
await rebound.promise;
|
||||
await expect(service.navigateTree(sessionRef(replacementSessionId), navigationRequest({ mode: "none" }, "replacement-leaf"))).rejects.toThrow(
|
||||
"Stop current session activity before navigating the session tree",
|
||||
);
|
||||
|
||||
replacement.resolve({ cancelled: false });
|
||||
await expect(cloning).resolves.toMatchObject({ type: "done", message: "Session cloned" });
|
||||
await expect(service.navigateTree(sessionRef(replacementSessionId), navigationRequest({ mode: "none" }, "replacement-leaf"))).resolves.toEqual({ cancelled: false });
|
||||
expect(fork).toHaveBeenCalledOnce();
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("blocks tree navigation while session resources reload", async () => {
|
||||
const reloadOperation = deferred<undefined>();
|
||||
const reload = vi.fn(() => reloadOperation.promise);
|
||||
const navigateTree = vi.fn<NavigateTree>(() => Promise.resolve({ cancelled: false }));
|
||||
const { service } = treeHarness({}, { navigateTree, reload });
|
||||
|
||||
const reloading = service.runCommand(sessionRef(SESSION_ID), "/reload");
|
||||
await vi.waitFor(() => { expect(reload).toHaveBeenCalledOnce(); });
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).rejects.toThrow(
|
||||
"Stop current session activity before navigating the session tree",
|
||||
);
|
||||
|
||||
reloadOperation.resolve(undefined);
|
||||
await expect(reloading).resolves.toMatchObject({ type: "done" });
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).resolves.toEqual({ cancelled: false });
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("blocks tree navigation while the route-level runtime reload disposes and reopens the session", async () => {
|
||||
const disposal = deferred<undefined>();
|
||||
const navigateTree = vi.fn<NavigateTree>(() => Promise.resolve({ cancelled: false }));
|
||||
const { service, fake } = treeHarness({}, { navigateTree });
|
||||
const dispose = vi.fn(() => disposal.promise);
|
||||
fake.runtime.dispose = dispose;
|
||||
|
||||
const reloading = service.reload(sessionRef(SESSION_ID));
|
||||
await vi.waitFor(() => { expect(dispose).toHaveBeenCalledOnce(); });
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).rejects.toThrow(
|
||||
"Stop current session activity before navigating the session tree",
|
||||
);
|
||||
|
||||
disposal.resolve(undefined);
|
||||
await reloading;
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("blocks tree navigation throughout an asynchronous archive operation", async () => {
|
||||
const archiveOperation = deferred<{ sessionId: string; cwd: string; archivedAt: string }>();
|
||||
const archive = vi.fn(() => archiveOperation.promise);
|
||||
const archiveStore = { ...emptyArchiveStore(), archive };
|
||||
const navigateTree = vi.fn<NavigateTree>(() => Promise.resolve({ cancelled: false }));
|
||||
const { service } = treeHarness({}, { navigateTree }, { archiveStore });
|
||||
|
||||
const archiving = service.archive(sessionRef(SESSION_ID));
|
||||
await vi.waitFor(() => { expect(archive).toHaveBeenCalledOnce(); });
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest())).rejects.toThrow(
|
||||
"Stop current session activity before navigating the session tree",
|
||||
);
|
||||
|
||||
archiveOperation.resolve({ sessionId: SESSION_ID, cwd: "/workspace", archivedAt: "2026-01-01T00:00:00.000Z" });
|
||||
await archiving;
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("aborts branch summarization through the existing abort path and reports cancellation", async () => {
|
||||
const navigation = deferred<Awaited<ReturnType<NavigateTree>>>();
|
||||
const navigateTree = vi.fn<NavigateTree>(() => navigation.promise);
|
||||
const abortBranchSummary = vi.fn(() => { navigation.resolve({ cancelled: true, aborted: true }); });
|
||||
const abort = vi.fn(() => Promise.resolve());
|
||||
const { service, hub } = treeHarness({}, { navigateTree, abortBranchSummary, abort });
|
||||
|
||||
const navigationResult = service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "default" }));
|
||||
await vi.waitFor(() => { expect(navigateTree).toHaveBeenCalledOnce(); });
|
||||
await service.abort(sessionRef(SESSION_ID));
|
||||
|
||||
await expect(navigationResult).resolves.toEqual({ cancelled: true, aborted: true });
|
||||
expect(abortBranchSummary).toHaveBeenCalledOnce();
|
||||
expect(abort).toHaveBeenCalledOnce();
|
||||
expect(abortBranchSummary.mock.invocationCallOrder[0]).toBeLessThan(abort.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY);
|
||||
expect(hub.sessionEvents.some(({ event }) => event.type === "activity.update" && event.activity.label === "branch summary aborted")).toBe(true);
|
||||
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("does not republish stale navigation state after stopping and disposing its runtime", async () => {
|
||||
const navigation = deferred<Awaited<ReturnType<NavigateTree>>>();
|
||||
const navigateTree = vi.fn<NavigateTree>(() => navigation.promise);
|
||||
const abortBranchSummary = vi.fn(() => { navigation.resolve({ cancelled: true, aborted: true }); });
|
||||
const { service, hub, fake } = treeHarness({}, { navigateTree, abortBranchSummary });
|
||||
|
||||
const navigationResult = service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "default" }));
|
||||
await vi.waitFor(() => { expect(navigateTree).toHaveBeenCalledOnce(); });
|
||||
hub.sessionEvents.length = 0;
|
||||
await service.stop(sessionRef(SESSION_ID));
|
||||
await expect(navigationResult).resolves.toEqual({ cancelled: true, aborted: true });
|
||||
|
||||
expect(service.activeCount()).toBe(0);
|
||||
expect(fake.calls.dispose).toBe(1);
|
||||
expect(hub.sessionEvents).toEqual([]);
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("still runs the normal abort and releases the gate when the branch-summary abort hook fails", async () => {
|
||||
const navigation = deferred<Awaited<ReturnType<NavigateTree>>>();
|
||||
const navigateTree = vi.fn<NavigateTree>(() => navigation.promise);
|
||||
const branchAbortFailure = new Error("branch abort hook failed");
|
||||
const abortBranchSummary = vi.fn<NonNullable<PiAgentSession["abortBranchSummary"]>>(() => { throw branchAbortFailure; });
|
||||
const abort = vi.fn(() => {
|
||||
navigation.resolve({ cancelled: true, aborted: true });
|
||||
return Promise.resolve();
|
||||
});
|
||||
const { service } = treeHarness({}, { navigateTree, abortBranchSummary, abort });
|
||||
|
||||
const navigationResult = service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "default" }));
|
||||
await vi.waitFor(() => { expect(navigateTree).toHaveBeenCalledOnce(); });
|
||||
await expect(service.abort(sessionRef(SESSION_ID))).rejects.toBe(branchAbortFailure);
|
||||
await expect(navigationResult).resolves.toEqual({ cancelled: true, aborted: true });
|
||||
|
||||
expect(abortBranchSummary).toHaveBeenCalledOnce();
|
||||
expect(abort).toHaveBeenCalledOnce();
|
||||
await expect(service.prompt(sessionRef(SESSION_ID), "gate released after abort failure")).resolves.toBeUndefined();
|
||||
abortBranchSummary.mockImplementation(() => undefined);
|
||||
await service.dispose();
|
||||
});
|
||||
|
||||
it("releases the gate and publishes final status after navigation failure", async () => {
|
||||
const failure = new Error("summary provider failed");
|
||||
const navigateTree = vi.fn<NavigateTree>();
|
||||
navigateTree.mockRejectedValueOnce(failure).mockResolvedValueOnce({ cancelled: false });
|
||||
const { service, hub } = treeHarness({}, { navigateTree });
|
||||
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "default" }))).rejects.toBe(failure);
|
||||
await expect(service.navigateTree(sessionRef(SESSION_ID), navigationRequest({ mode: "none" }))).resolves.toEqual({ cancelled: false });
|
||||
|
||||
expect(navigateTree).toHaveBeenCalledTimes(2);
|
||||
expect(hub.sessionEvents.some(({ event }) => event.type === "activity.update"
|
||||
&& event.activity.phase === "error"
|
||||
&& event.activity.detail === "summary provider failed")).toBe(true);
|
||||
expect(hub.sessionEvents.filter(({ event }) => event.type === "status.update").length).toBeGreaterThanOrEqual(4);
|
||||
|
||||
await service.dispose();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user