diff --git a/apps/web/lib/layout/panel-portal-manager.test.ts b/apps/web/lib/layout/panel-portal-manager.test.ts new file mode 100644 index 0000000000..f63d507c29 --- /dev/null +++ b/apps/web/lib/layout/panel-portal-manager.test.ts @@ -0,0 +1,45 @@ +import { describe, it, expect, vi } from "vitest"; +import { PanelPortalManager } from "./panel-portal-manager"; + +function mockApi() { + return { title: "", setTitle: vi.fn() } as never; +} + +describe("PanelPortalManager.reconcile", () => { + it("removes portals whose panel is no longer live", () => { + const mgr = new PanelPortalManager(); + mgr.acquire("panel-a", "chat", {}, mockApi()); + mgr.acquire("panel-b", "terminal", {}, mockApi()); + + mgr.reconcile(new Set(["panel-a"])); + + expect(mgr.has("panel-a")).toBe(true); + expect(mgr.has("panel-b")).toBe(false); + }); + + it("no-ops when all portals are live", () => { + const mgr = new PanelPortalManager(); + const listener = vi.fn(); + mgr.acquire("panel-a", "chat", {}, mockApi()); + mgr.subscribe(listener); + listener.mockClear(); + + mgr.reconcile(new Set(["panel-a"])); + + expect(mgr.has("panel-a")).toBe(true); + expect(listener).not.toHaveBeenCalled(); + }); + + it("notifies listeners when portals are removed", () => { + const mgr = new PanelPortalManager(); + const listener = vi.fn(); + mgr.acquire("panel-a", "chat", {}, mockApi()); + mgr.acquire("panel-b", "terminal", {}, mockApi()); + mgr.subscribe(listener); + listener.mockClear(); + + mgr.reconcile(new Set(["panel-a"])); + + expect(listener).toHaveBeenCalledOnce(); + }); +}); diff --git a/apps/web/lib/layout/panel-portal-manager.ts b/apps/web/lib/layout/panel-portal-manager.ts index 3707beb85a..67a616e519 100644 --- a/apps/web/lib/layout/panel-portal-manager.ts +++ b/apps/web/lib/layout/panel-portal-manager.ts @@ -52,7 +52,7 @@ export type PortalEntry = { type Listener = () => void; -class PanelPortalManager { +export class PanelPortalManager { private entries = new Map(); private listeners = new Set(); @@ -112,6 +112,28 @@ class PanelPortalManager { this.notify(); } + /** + * Release portals whose panel no longer exists in dockview. + * Call after fast-path session switches where `isRestoringLayout` blocked + * the normal `onDidRemovePanel` cleanup. + */ + reconcile(livePanelIds: Set): void { + const toRemove: string[] = []; + for (const panelId of this.entries.keys()) { + if (!livePanelIds.has(panelId)) { + toRemove.push(panelId); + } + } + if (toRemove.length === 0) return; + for (const panelId of toRemove) { + const entry = this.entries.get(panelId)!; + entry.element.remove(); + entry.api = null; + this.entries.delete(panelId); + } + this.notify(); + } + /** Release all portals (e.g. when the layout component unmounts entirely). */ releaseAll(): void { if (this.entries.size === 0) return; diff --git a/apps/web/lib/state/dockview-session-switch.test.ts b/apps/web/lib/state/dockview-session-switch.test.ts index 5a25f5a6de..054aa09a4a 100644 --- a/apps/web/lib/state/dockview-session-switch.test.ts +++ b/apps/web/lib/state/dockview-session-switch.test.ts @@ -29,7 +29,8 @@ function makeMockApi() { panels: [], layout: vi.fn(), fromJSON: vi.fn(), - getPanel: vi.fn(), + getPanel: vi.fn(() => null), + addPanel: vi.fn(), } as unknown as SessionSwitchParams["api"]; } @@ -73,6 +74,36 @@ describe("performSessionSwitch", () => { expect(params.api.fromJSON).not.toHaveBeenCalled(); }); + it("creates session panel inline on the fast path when it does not exist", () => { + vi.mocked(layoutStructuresMatch).mockReturnValueOnce(true); + const params = makeParams(); + + performSessionSwitch(params); + + expect(params.api.addPanel).toHaveBeenCalledWith( + expect.objectContaining({ + id: "session:new-session", + component: "chat", + params: { sessionId: "new-session" }, + }), + ); + }); + + it("skips addPanel on the fast path when the session panel already exists", () => { + vi.mocked(layoutStructuresMatch).mockReturnValueOnce(true); + const panel = { id: "session:new-session", api: { component: "chat" }, group: { id: "g1" } }; + const params = makeParams({ + api: { + ...makeMockApi(), + getPanel: vi.fn((id: string) => (id === "session:new-session" ? panel : null)), + } as unknown as SessionSwitchParams["api"], + }); + + performSessionSwitch(params); + + expect(params.api.addPanel).not.toHaveBeenCalled(); + }); + it("calls api.layout on the slow path (buildDefault fallback)", () => { const params = makeParams(); diff --git a/apps/web/lib/state/dockview-session-switch.ts b/apps/web/lib/state/dockview-session-switch.ts index a904ccbea5..a0b36f943a 100644 --- a/apps/web/lib/state/dockview-session-switch.ts +++ b/apps/web/lib/state/dockview-session-switch.ts @@ -83,6 +83,31 @@ function tryFastSessionSwitch(params: SessionSwitchParams): LayoutGroupIds | nul // Session-scoped portals (browser, vscode, etc.) will be re-acquired // via usePortalSlot's sessionId dependency change. removeEphemeralPanels(api, newSessionId); + + // Create the new session tab inline so there's no gap between removing the + // old tab and creating the new one. useAutoSessionTab will detect the panel + // already exists and skip creation. + if (!api.getPanel(`session:${newSessionId}`)) { + const chatPanel = api.panels.find( + (p) => p.id.startsWith("session:") || p.api.component === "chat", + ); + const sidebarPanel = api.getPanel("sidebar"); + let position: import("dockview-react").AddPanelOptions["position"]; + if (chatPanel) { + position = { referenceGroup: chatPanel.group.id }; + } else if (sidebarPanel) { + position = { direction: "right" as const, referencePanel: "sidebar" }; + } + api.addPanel({ + id: `session:${newSessionId}`, + component: "chat", + tabComponent: "sessionTab", + title: "Agent", + params: { sessionId: newSessionId }, + position, + }); + } + api.layout(params.safeWidth, params.safeHeight); return applyLayoutFixups(api); } diff --git a/apps/web/lib/state/dockview-store.ts b/apps/web/lib/state/dockview-store.ts index fd109f46c6..4e71569d4a 100644 --- a/apps/web/lib/state/dockview-store.ts +++ b/apps/web/lib/state/dockview-store.ts @@ -472,6 +472,7 @@ function buildSessionSwitchAction(set: StoreSet, get: StoreGet) { }); set(ids); set({ isRestoringLayout: false }); + panelPortalManager.reconcile(new Set(api.panels.map((p) => p.id))); } catch { set({ isRestoringLayout: false }); } diff --git a/apps/web/lib/state/slices/session-runtime/migrate-env-keyed-data.test.ts b/apps/web/lib/state/slices/session-runtime/migrate-env-keyed-data.test.ts index fb1e5efb07..d51b7df3cf 100644 --- a/apps/web/lib/state/slices/session-runtime/migrate-env-keyed-data.test.ts +++ b/apps/web/lib/state/slices/session-runtime/migrate-env-keyed-data.test.ts @@ -2,12 +2,27 @@ import { describe, it, expect } from "vitest"; import { create } from "zustand"; import { immer } from "zustand/middleware/immer"; import { createSessionRuntimeSlice } from "./session-runtime-slice"; +import { createSessionSlice } from "../session/session-slice"; import type { SessionRuntimeSlice } from "./types"; +import type { SessionSlice } from "../session/types"; function makeStore() { return create()(immer(createSessionRuntimeSlice)); } +type CombinedSlice = SessionRuntimeSlice & SessionSlice; + +function makeCombinedStore() { + return create()( + immer((set, get, api) => ({ + // eslint-disable-next-line @typescript-eslint/no-explicit-any + ...createSessionSlice(set as any), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + ...createSessionRuntimeSlice(set as any, get as any, api as any), + })), + ); +} + describe("registerSessionEnvironment — migrateEnvKeyedData", () => { it("migrates data from sessionId key to environmentId key", () => { const store = makeStore(); @@ -89,3 +104,83 @@ describe("registerSessionEnvironment — migrateEnvKeyedData", () => { expect(state.sessionCommits.byEnvironmentId["sess-3"]).toBeUndefined(); }); }); + +describe("setTaskSession — cross-slice migration", () => { + it("migrates env-keyed data when task_environment_id is set", () => { + const store = makeCombinedStore(); + + // Data arrives under fallback sessionId key before the mapping is known + store.getState().setGitStatus("sess-1", { branch: "main" } as never); + store.getState().setSessionCommits("sess-1", [{ commit_sha: "abc" } as never]); + + // Session update arrives with task_environment_id + store.getState().setTaskSession({ + id: "sess-1", + task_id: "task-1", + state: "RUNNING", + task_environment_id: "env-1", + started_at: "", + updated_at: "", + }); + + const state = store.getState(); + expect(state.environmentIdBySessionId["sess-1"]).toBe("env-1"); + expect(state.gitStatus.byEnvironmentId["env-1"]).toEqual({ branch: "main" }); + expect(state.gitStatus.byEnvironmentId["sess-1"]).toBeUndefined(); + expect(state.sessionCommits.byEnvironmentId["env-1"]).toEqual([{ commit_sha: "abc" }]); + expect(state.sessionCommits.byEnvironmentId["sess-1"]).toBeUndefined(); + }); + + it("does not migrate when task_environment_id is absent", () => { + const store = makeCombinedStore(); + + store.getState().setGitStatus("sess-2", { branch: "dev" } as never); + + store.getState().setTaskSession({ + id: "sess-2", + task_id: "task-2", + state: "RUNNING", + started_at: "", + updated_at: "", + }); + + const state = store.getState(); + expect(state.environmentIdBySessionId["sess-2"]).toBeUndefined(); + expect(state.gitStatus.byEnvironmentId["sess-2"]).toEqual({ branch: "dev" }); + }); +}); + +describe("setTaskSessionsForTask — bulk cross-slice migration", () => { + it("migrates env-keyed data for all sessions with task_environment_id", () => { + const store = makeCombinedStore(); + + // Data under fallback keys + store.getState().setGitStatus("sess-a", { branch: "a" } as never); + store.getState().appendShellOutput("sess-b", "hello"); + + store.getState().setTaskSessionsForTask("task-1", [ + { + id: "sess-a", + task_id: "task-1", + state: "COMPLETED", + task_environment_id: "env-x", + started_at: "", + updated_at: "", + }, + { + id: "sess-b", + task_id: "task-1", + state: "RUNNING", + task_environment_id: "env-y", + started_at: "", + updated_at: "", + }, + ]); + + const state = store.getState(); + expect(state.gitStatus.byEnvironmentId["env-x"]).toEqual({ branch: "a" }); + expect(state.gitStatus.byEnvironmentId["sess-a"]).toBeUndefined(); + expect(state.shell.outputs["env-y"]).toBe("hello"); + expect(state.shell.outputs["sess-b"]).toBeUndefined(); + }); +}); diff --git a/apps/web/lib/state/slices/session-runtime/session-runtime-slice.ts b/apps/web/lib/state/slices/session-runtime/session-runtime-slice.ts index baadeee62c..0217735f8c 100644 --- a/apps/web/lib/state/slices/session-runtime/session-runtime-slice.ts +++ b/apps/web/lib/state/slices/session-runtime/session-runtime-slice.ts @@ -190,7 +190,7 @@ function buildUserShellActions(set: ImmerSet) { * proper `environmentId` key so selectors don't see stale data after the * session→environment mapping is registered. */ -function migrateEnvKeyedData( +export function migrateEnvKeyedData( draft: SessionRuntimeSliceState, sessionId: string, environmentId: string, diff --git a/apps/web/lib/state/slices/session/session-slice.ts b/apps/web/lib/state/slices/session/session-slice.ts index b6b5808810..54e55a8180 100644 --- a/apps/web/lib/state/slices/session/session-slice.ts +++ b/apps/web/lib/state/slices/session/session-slice.ts @@ -1,6 +1,7 @@ import type { StateCreator } from "zustand"; import type { TaskSession } from "@/lib/types/http"; import type { SessionSlice, SessionSliceState } from "./types"; +import { migrateEnvKeyedData } from "@/lib/state/slices/session-runtime/session-runtime-slice"; /** Ensure message metadata exists for a session, initializing with defaults if needed. */ function ensureMessageMeta( @@ -37,6 +38,15 @@ function mergeMessageFields(target: Record, source: Record s.id === session.id); if (sessionIndex >= 0) sessionsByTask[sessionIndex] = mergedSession; } - // Eagerly populate session→environment mapping (cross-slice access to session-runtime) - if (mergedSession.task_environment_id) { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (draft as any).environmentIdBySessionId[session.id] = mergedSession.task_environment_id; - } + syncEnvironmentMapping(draft, session.id, mergedSession.task_environment_id); }), removeTaskSession: (taskId, sessionId) => set((draft) => { @@ -230,11 +236,7 @@ export const createSessionSlice: StateCreator< draft.taskSessionsByTask.loadedByTaskId[taskId] = true; for (const session of sessions) { draft.taskSessions.items[session.id] = session; - // Eagerly populate session→environment mapping (cross-slice access to session-runtime) - if (session.task_environment_id) { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (draft as any).environmentIdBySessionId[session.id] = session.task_environment_id; - } + syncEnvironmentMapping(draft, session.id, session.task_environment_id); } }), setTaskSessionsLoading: (taskId, loading) =>