From c042050469f220cb13cc1fcd2073c4e8ea20c7b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20Flor=C3=AAncio?= Date: Mon, 4 May 2026 18:53:33 +0100 Subject: [PATCH] fix: stabilize task switch flow and session recovery (#818) --- apps/backend/cmd/kandev/helpers.go | 3 + .../internal/agent/lifecycle/event_types.go | 15 +- .../internal/agent/lifecycle/events.go | 15 +- .../task/dockview-desktop-layout.test.ts | 34 ++++ .../task/dockview-desktop-layout.tsx | 19 +- .../task/dockview-layout-restore.test.ts | 49 +++++ .../task/dockview-layout-restore.ts | 39 ++-- .../web/components/task/file-browser-hooks.ts | 4 +- .../domains/session/use-session-agentctl.ts | 4 +- .../session/use-session-resumption.test.ts | 36 ++++ .../domains/session/use-session-resumption.ts | 18 +- apps/web/hooks/use-task-sessions.ts | 3 +- apps/web/lib/state/dockview-env-switch.ts | 3 +- .../lib/state/slices/kanban/kanban-slice.ts | 15 +- .../lib/state/slices/session/session-slice.ts | 113 +++++++---- .../session/session-slice.upsert.test.ts | 124 ++++++++++++ apps/web/lib/state/slices/session/types.ts | 1 + apps/web/lib/state/store.ts | 1 + apps/web/lib/types/backend.ts | 1 + .../web/lib/ws/handlers/agent-session.test.ts | 189 ++++++++++++++++++ apps/web/lib/ws/handlers/agent-session.ts | 125 ++++++------ 21 files changed, 665 insertions(+), 146 deletions(-) create mode 100644 apps/web/components/task/dockview-desktop-layout.test.ts create mode 100644 apps/web/lib/state/slices/session/session-slice.upsert.test.ts diff --git a/apps/backend/cmd/kandev/helpers.go b/apps/backend/cmd/kandev/helpers.go index 3c62ae128d..d00c074c66 100644 --- a/apps/backend/cmd/kandev/helpers.go +++ b/apps/backend/cmd/kandev/helpers.go @@ -120,6 +120,9 @@ func appendAgentctlStatusMessage( sessionIDPayloadKey: sessionID, "agent_execution_id": execution.ID, } + if execution.TaskEnvironmentID != "" { + payload["task_environment_id"] = execution.TaskEnvironmentID + } action := ws.ActionSessionAgentctlStarting if execution.GetWorkspaceStream() != nil { action = ws.ActionSessionAgentctlReady diff --git a/apps/backend/internal/agent/lifecycle/event_types.go b/apps/backend/internal/agent/lifecycle/event_types.go index d8885c9426..bfdc6ad49c 100644 --- a/apps/backend/internal/agent/lifecycle/event_types.go +++ b/apps/backend/internal/agent/lifecycle/event_types.go @@ -23,13 +23,14 @@ type AgentEventPayload struct { // AgentctlEventPayload is the payload for agentctl lifecycle events (starting, ready, error). type AgentctlEventPayload struct { - TaskID string `json:"task_id"` - SessionID string `json:"session_id"` - AgentExecutionID string `json:"agent_execution_id"` - ErrorMessage string `json:"error_message,omitempty"` - WorktreeID string `json:"worktree_id,omitempty"` - WorktreePath string `json:"worktree_path,omitempty"` - WorktreeBranch string `json:"worktree_branch,omitempty"` + TaskID string `json:"task_id"` + SessionID string `json:"session_id"` + TaskEnvironmentID string `json:"task_environment_id,omitempty"` + AgentExecutionID string `json:"agent_execution_id"` + ErrorMessage string `json:"error_message,omitempty"` + WorktreeID string `json:"worktree_id,omitempty"` + WorktreePath string `json:"worktree_path,omitempty"` + WorktreeBranch string `json:"worktree_branch,omitempty"` } // ACPSessionCreatedPayload is the payload when an ACP session is created. diff --git a/apps/backend/internal/agent/lifecycle/events.go b/apps/backend/internal/agent/lifecycle/events.go index d147c94160..3951d9995c 100644 --- a/apps/backend/internal/agent/lifecycle/events.go +++ b/apps/backend/internal/agent/lifecycle/events.go @@ -80,13 +80,14 @@ func (p *EventPublisher) PublishAgentctlEvent(ctx context.Context, eventType str } payload := AgentctlEventPayload{ - TaskID: execution.TaskID, - SessionID: execution.SessionID, - AgentExecutionID: execution.ID, - ErrorMessage: errMsg, - WorktreeID: worktreeID, - WorktreePath: execution.WorkspacePath, - WorktreeBranch: worktreeBranch, + TaskID: execution.TaskID, + SessionID: execution.SessionID, + TaskEnvironmentID: execution.TaskEnvironmentID, + AgentExecutionID: execution.ID, + ErrorMessage: errMsg, + WorktreeID: worktreeID, + WorktreePath: execution.WorkspacePath, + WorktreeBranch: worktreeBranch, } event := bus.NewEvent(eventType, "agent-manager", payload) diff --git a/apps/web/components/task/dockview-desktop-layout.test.ts b/apps/web/components/task/dockview-desktop-layout.test.ts new file mode 100644 index 0000000000..4adf55ca77 --- /dev/null +++ b/apps/web/components/task/dockview-desktop-layout.test.ts @@ -0,0 +1,34 @@ +import { describe, it, expect } from "vitest"; +import { CHAT_PANEL_FALLBACK_LABEL, resolveChatPanelTitle } from "./dockview-desktop-layout"; + +/** + * Regression: the generic "chat" placeholder dockview panel used to fall back + * to the literal "Agent" label even when the active session's agent profile + * was loaded (e.g. "Opus"). The bug was a stale `isSessionTab && agentLabel` + * gate inside `useChatSessionTitle` that suppressed the agent label for the + * non-session-scoped placeholder. The pure resolver below is the single place + * the gate would have to be re-introduced, so this test pins the behavior. + */ +describe("resolveChatPanelTitle", () => { + it("returns the agent label when one is provided", () => { + expect(resolveChatPanelTitle("Opus")).toBe("Opus"); + }); + + it("falls back to the generic 'Agent' label when null", () => { + expect(resolveChatPanelTitle(null)).toBe(CHAT_PANEL_FALLBACK_LABEL); + }); + + it("falls back to the generic 'Agent' label when undefined", () => { + expect(resolveChatPanelTitle(undefined)).toBe(CHAT_PANEL_FALLBACK_LABEL); + }); + + it("falls back to the generic 'Agent' label when the agent label is empty", () => { + expect(resolveChatPanelTitle("")).toBe(CHAT_PANEL_FALLBACK_LABEL); + }); + + it("uses the agent label verbatim — does not coerce or relabel valid names", () => { + for (const name of ["Mock", "Claude Code", "GPT-5", "amp"]) { + expect(resolveChatPanelTitle(name)).toBe(name); + } + }); +}); diff --git a/apps/web/components/task/dockview-desktop-layout.tsx b/apps/web/components/task/dockview-desktop-layout.tsx index 6ae8ed5f33..fa5d302aee 100644 --- a/apps/web/components/task/dockview-desktop-layout.tsx +++ b/apps/web/components/task/dockview-desktop-layout.tsx @@ -201,7 +201,13 @@ function SidebarContent({ panelId }: { panelId: string }) { return ; } -function useChatSessionTitle(panelId: string, sessionId: string | null, isSessionTab: boolean) { +export const CHAT_PANEL_FALLBACK_LABEL = "Agent"; + +export function resolveChatPanelTitle(agentLabel: string | null | undefined): string { + return agentLabel || CHAT_PANEL_FALLBACK_LABEL; +} + +function useChatSessionTitle(panelId: string, sessionId: string | null) { const agentLabel = useAppStore((state) => { if (!sessionId) return null; const session = state.taskSessions.items[sessionId]; @@ -214,12 +220,8 @@ function useChatSessionTitle(panelId: string, sessionId: string | null, isSessio return parts[1] || parts[0] || profile.label; }); useEffect(() => { - let label = "Agent"; - if (isSessionTab && agentLabel) { - label = agentLabel; - } - setPanelTitle(panelId, label); - }, [panelId, isSessionTab, agentLabel]); + setPanelTitle(panelId, resolveChatPanelTitle(agentLabel)); + }, [panelId, agentLabel]); } function ChatContent({ panelId, params }: { panelId: string; params: Record }) { @@ -230,7 +232,7 @@ function ChatContent({ panelId, params }: { panelId: string; params: Record sessionId ? state.taskSessions.items[sessionId]?.is_passthrough === true : false, ); - useChatSessionTitle(panelId, sessionId, !!paramSessionId); + useChatSessionTitle(panelId, sessionId); if (isPassthrough) { return ( @@ -405,7 +407,6 @@ function useEnvSwitchCleanup(effectiveSessionId: string | null, effectiveEnvId: const prevEnvRef = useRef(undefined); useEffect(() => { const newEnvId = effectiveEnvId; - if (prevEnvRef.current === undefined) { prevEnvRef.current = newEnvId; return; diff --git a/apps/web/components/task/dockview-layout-restore.test.ts b/apps/web/components/task/dockview-layout-restore.test.ts index f0a3e5f980..26cba61145 100644 --- a/apps/web/components/task/dockview-layout-restore.test.ts +++ b/apps/web/components/task/dockview-layout-restore.test.ts @@ -138,3 +138,52 @@ describe("sanitizeLayout - size validation", () => { expect(Object.keys(result.panels)).not.toContain("unknown"); }); }); + +describe("sanitizeLayout - session panel handling", () => { + const SESSION_PANEL_ID = "session:abc-123"; + + function buildLayoutWithSession() { + const layout = buildLayout(); + layout.grid.root.data[1].data.views = ["chat", SESSION_PANEL_ID]; + layout.panels = { + ...layout.panels, + // @ts-expect-error - injecting a session panel without a matching valid component + [SESSION_PANEL_ID]: { id: SESSION_PANEL_ID, contentComponent: "session-panel" }, + }; + return layout; + } + + it("keeps session:* panels by default (per-env restore)", () => { + const result = sanitizeLayout(buildLayoutWithSession(), VALID_COMPONENTS); + expect(result).not.toBeNull(); + expect(Object.keys(result.panels)).toContain(SESSION_PANEL_ID); + }); + + it("strips session:* panels when stripSessionPanels=true (global fallback)", () => { + const result = sanitizeLayout(buildLayoutWithSession(), VALID_COMPONENTS, { + stripSessionPanels: true, + }); + expect(result).not.toBeNull(); + expect(Object.keys(result.panels)).not.toContain(SESSION_PANEL_ID); + // The non-session panels stay so the rest of the layout is still usable. + expect(Object.keys(result.panels)).toContain("chat"); + expect(Object.keys(result.panels)).toContain("files"); + }); + + it("strips session:* panels even when contentComponent is a valid component (e.g. 'chat')", () => { + // dockview serializes panels added via api.addPanel({ component: "chat" }) + // with contentComponent: "chat", which is in VALID_COMPONENTS. The strip + // logic must key off the panel id, not the component name. + const layout = buildLayout(); + layout.grid.root.data[1].data.views = ["chat", SESSION_PANEL_ID]; + layout.panels = { + ...layout.panels, + // @ts-expect-error - injecting a session panel with the chat component + [SESSION_PANEL_ID]: { id: SESSION_PANEL_ID, contentComponent: "chat" }, + }; + const result = sanitizeLayout(layout, VALID_COMPONENTS, { stripSessionPanels: true }); + expect(result).not.toBeNull(); + expect(Object.keys(result.panels)).not.toContain(SESSION_PANEL_ID); + expect(Object.keys(result.panels)).toContain("chat"); + }); +}); diff --git a/apps/web/components/task/dockview-layout-restore.ts b/apps/web/components/task/dockview-layout-restore.ts index 7c1ba1dc37..6b0656f292 100644 --- a/apps/web/components/task/dockview-layout-restore.ts +++ b/apps/web/components/task/dockview-layout-restore.ts @@ -6,17 +6,31 @@ import { getEnvLayout, getEnvMaximizeState } from "@/lib/local-storage"; const LAYOUT_STORAGE_KEY = "dockview-layout-v1"; -// eslint-disable-next-line @typescript-eslint/no-explicit-any -export function sanitizeLayout(layout: any, validComponents: Set): any { +/* eslint-disable @typescript-eslint/no-explicit-any */ +export function sanitizeLayout( + layout: any, + validComponents: Set, + options: { stripSessionPanels?: boolean } = {}, +): any { if (!isLayoutShapeHealthy(layout)) return null; const invalidIds = new Set(); - // eslint-disable-next-line @typescript-eslint/no-explicit-any const validPanels: Record = {}; for (const [id, panel] of Object.entries(layout.panels)) { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const comp = (panel as any).contentComponent; - if (comp && (validComponents.has(comp) || id.startsWith("session:"))) { + // Session panels are scoped to a specific environment; when restoring the + // global fallback (no envId yet), they belong to the previous task and + // would leak in as duplicate tabs. Strip them in that case. The session + // check must happen before component-validity, since session panels are + // serialized with contentComponent: "chat" (a valid component) and would + // otherwise short-circuit the strip guard. + if (id.startsWith("session:")) { + if (options.stripSessionPanels) { + invalidIds.add(id); + } else { + validPanels[id] = panel; + } + } else if (comp && validComponents.has(comp)) { validPanels[id] = panel; } else { invalidIds.add(id); @@ -25,7 +39,6 @@ export function sanitizeLayout(layout: any, validComponents: Set): any { if (invalidIds.size === 0) return layout; - // eslint-disable-next-line @typescript-eslint/no-explicit-any function cleanNode(node: any): any { if (node.type === "leaf") { const views = (node.data.views as string[]).filter((v) => !invalidIds.has(v)); @@ -34,7 +47,6 @@ export function sanitizeLayout(layout: any, validComponents: Set): any { return { ...node, data: { ...node.data, views, activeView } }; } if (node.type === "branch") { - // eslint-disable-next-line @typescript-eslint/no-explicit-any const children = (node.data as any[]).map(cleanNode).filter(Boolean); if (children.length === 0) return null; return { ...node, data: children }; @@ -51,6 +63,7 @@ export function sanitizeLayout(layout: any, validComponents: Set): any { panels: validPanels, }; } +/* eslint-enable @typescript-eslint/no-explicit-any */ function applyFixupsWithMaximize(api: DockviewReadyEvent["api"], envId: string | null): void { const savedMax = envId ? getEnvMaximizeState(envId) : null; @@ -97,15 +110,13 @@ export function tryRestoreLayout( const envLayout = getEnvLayout(currentEnvId); if (envLayout) { const sanitized = sanitizeLayout(envLayout, validComponents); - if (!sanitized) { - return false; - } + if (!sanitized) return false; api.fromJSON(sanitized as SerializedDockview); applyFixupsWithMaximize(api, currentEnvId); return true; } } catch { - // Per-env restore failed, try global + // fall through to maximize-only / global fallback } if (tryRestoreMaximizeOnly(api, currentEnvId)) return true; } @@ -114,14 +125,16 @@ export function tryRestoreLayout( try { const saved = localStorage.getItem(LAYOUT_STORAGE_KEY); if (saved) { - const layout = sanitizeLayout(JSON.parse(saved), validComponents); + const layout = sanitizeLayout(JSON.parse(saved), validComponents, { + stripSessionPanels: true, + }); if (!layout) return false; api.fromJSON(layout); useDockviewStore.setState(applyLayoutFixups(api)); return true; } } catch { - // Global restore failed, build default + // fall through to default-build } } diff --git a/apps/web/components/task/file-browser-hooks.ts b/apps/web/components/task/file-browser-hooks.ts index f8d1e81ed7..e93df8eb72 100644 --- a/apps/web/components/task/file-browser-hooks.ts +++ b/apps/web/components/task/file-browser-hooks.ts @@ -373,7 +373,7 @@ export function useFileBrowserTree(sessionId: string, resetKey?: string) { return () => { clearRetryTimer(); }; - }, [clearRetryTimer, loadTree, effectiveResetKey]); + }, [clearRetryTimer, loadTree, effectiveResetKey, sessionId]); // Fire the initial load on the waiting → ready transition. `loadState` is // read via a ref so a failed load (which sets state back to "waiting") does @@ -385,7 +385,7 @@ export function useFileBrowserTree(sessionId: string, resetKey?: string) { if (loadStateRef.current === "loading") return; if (loadStateRef.current === "loaded" && treeRef.current) return; void loadTree({ resetRetry: true }); - }, [agentctlStatus.isReady, loadTree]); + }, [agentctlStatus.isReady, loadTree, sessionId]); useEffect(() => { if (!tree || isLoadingTree || hasInitializedExpandedRef.current === effectiveResetKey) return; diff --git a/apps/web/hooks/domains/session/use-session-agentctl.ts b/apps/web/hooks/domains/session/use-session-agentctl.ts index 30847c935f..6b1d1adbc7 100644 --- a/apps/web/hooks/domains/session/use-session-agentctl.ts +++ b/apps/web/hooks/domains/session/use-session-agentctl.ts @@ -19,11 +19,13 @@ export function useSessionAgentctl(sessionId: string | null) { return client.subscribeSession(session.id); }, [session?.id, connectionStatus]); + const isReady = status?.status === "ready"; + return { status: status?.status ?? "starting", errorMessage: status?.errorMessage, agentExecutionId: status?.agentExecutionId, - isReady: status?.status === "ready", + isReady, isStarting: status?.status === "starting" || !status, isError: status?.status === "error", }; diff --git a/apps/web/hooks/domains/session/use-session-resumption.test.ts b/apps/web/hooks/domains/session/use-session-resumption.test.ts index b991e499b2..1977f97638 100644 --- a/apps/web/hooks/domains/session/use-session-resumption.test.ts +++ b/apps/web/hooks/domains/session/use-session-resumption.test.ts @@ -48,6 +48,7 @@ function createSetters(): { setters: ResumeStateSetter; calls: SetterCalls } { return { setters, calls }; } +// eslint-disable-next-line max-lines-per-function -- test describe block, splitting hurts readability describe("resumeWithSilentFallback", () => { beforeEach(() => { vi.clearAllMocks(); @@ -155,4 +156,39 @@ describe("resumeWithSilentFallback", () => { "Failed to resume session — workspace restore also unavailable", ); }); + + it("seeds agentctl ready when restore_workspace fallback succeeds", async () => { + mockRequest + .mockResolvedValueOnce({ success: false, task_id: "t1", state: "FAILED" }) + .mockResolvedValueOnce({ + success: true, + task_id: "t1", + session_id: "s1", + state: "FAILED", + }); + const { setters } = createSetters(); + const setAgentctlReady = vi.fn(); + setters.setAgentctlReady = setAgentctlReady; + + await resumeWithSilentFallback("t1", "s1", null, setters); + + expect(setAgentctlReady).toHaveBeenCalledTimes(1); + expect(setAgentctlReady).toHaveBeenCalledWith("s1"); + }); + + it("does not seed agentctl ready when resume succeeds (new execution will emit its own events)", async () => { + mockRequest.mockResolvedValueOnce({ + success: true, + task_id: "t1", + session_id: "s1", + state: "STARTING", + }); + const { setters } = createSetters(); + const setAgentctlReady = vi.fn(); + setters.setAgentctlReady = setAgentctlReady; + + await resumeWithSilentFallback("t1", "s1", null, setters); + + expect(setAgentctlReady).not.toHaveBeenCalled(); + }); }); diff --git a/apps/web/hooks/domains/session/use-session-resumption.ts b/apps/web/hooks/domains/session/use-session-resumption.ts index e373a23ed2..38d755a046 100644 --- a/apps/web/hooks/domains/session/use-session-resumption.ts +++ b/apps/web/hooks/domains/session/use-session-resumption.ts @@ -103,7 +103,7 @@ async function resumeViaLaunch( setters.setResumptionState("resuming"); const { request } = buildRequest(taskId, sessionId); const launchResp = await launchSession(request); - applyResumeResponse( + const ok = applyResumeResponse( { success: launchResp.success, state: launchResp.state, @@ -115,6 +115,16 @@ async function resumeViaLaunch( session, setters, ); + // restore_workspace's whole purpose is to bring up agentctl HTTP for an + // otherwise-idle session — when it returns success the workspace+agentctl + // is up by definition. The backend's cached agentctl status snapshot uses + // "workspace stream attached" as its readiness signal, which is wrong on + // WS reconnect (stream detaches but agentctl HTTP keeps running) and the + // existing execution does not re-emit agentctl_ready, so the FileBrowser + // would otherwise stay stuck on "Preparing workspace". + if (ok && request.intent === "restore_workspace" && setters.setAgentctlReady) { + setters.setAgentctlReady(sessionId); + } } /** Attempt resume, silently falling back to restore_workspace on any failure. @@ -183,6 +193,12 @@ async function tryLaunch( session, setters, ); + // See comment in resumeViaLaunch — restore_workspace success implies + // agentctl HTTP is ready, but the existing execution may not re-emit the + // agentctl_ready WS event after WS reconnect. + if (request.intent === "restore_workspace" && setters.setAgentctlReady) { + setters.setAgentctlReady(sessionId); + } return true; } catch (err) { console.error("[tryLaunch] session launch failed", { diff --git a/apps/web/hooks/use-task-sessions.ts b/apps/web/hooks/use-task-sessions.ts index 0969c98b2d..37204ec8d1 100644 --- a/apps/web/hooks/use-task-sessions.ts +++ b/apps/web/hooks/use-task-sessions.ts @@ -25,7 +25,8 @@ export function useTaskSessions(taskId: string | null) { setTaskSessionsLoading(taskId, true); try { const response = await listTaskSessions(taskId, { cache: "no-store" }); - setTaskSessionsForTask(taskId, response.sessions ?? []); + const sessions = response.sessions ?? []; + setTaskSessionsForTask(taskId, sessions); } catch { setTaskSessionsForTask(taskId, []); } finally { diff --git a/apps/web/lib/state/dockview-env-switch.ts b/apps/web/lib/state/dockview-env-switch.ts index fdac03f2b4..744d9820e9 100644 --- a/apps/web/lib/state/dockview-env-switch.ts +++ b/apps/web/lib/state/dockview-env-switch.ts @@ -104,7 +104,6 @@ function tryFastEnvSwitch(params: EnvSwitchParams): LayoutGroupIds | null { } if (!structuresMatch) return null; - if (saved && savedLayoutHasEphemeralPanels(saved as SerializedDockview)) return null; const outgoingSessionPanel = api.panels.find( @@ -158,7 +157,7 @@ export function performEnvSwitch(params: EnvSwitchParams): LayoutGroupIds { api.layout(safeWidth, safeHeight); return applyLayoutFixups(api); } catch { - /* fall through */ + /* fall through to default layout build */ } } buildDefault(api); diff --git a/apps/web/lib/state/slices/kanban/kanban-slice.ts b/apps/web/lib/state/slices/kanban/kanban-slice.ts index 152515ee01..02bc40c301 100644 --- a/apps/web/lib/state/slices/kanban/kanban-slice.ts +++ b/apps/web/lib/state/slices/kanban/kanban-slice.ts @@ -42,29 +42,26 @@ export const createKanbanSlice: StateCreator< } draft.workflows.items = reordered; }), - setActiveTask: (taskId) => { + setActiveTask: (taskId) => set((draft) => { draft.tasks.activeTaskId = taskId; draft.tasks.activeSessionId = null; // New task → drop any pin; the pin only applies within a single task. draft.tasks.pinnedSessionId = null; - }); - }, - setActiveSession: (taskId, sessionId) => { + }), + setActiveSession: (taskId, sessionId) => set((draft) => { draft.tasks.activeTaskId = taskId; draft.tasks.activeSessionId = sessionId; // User-initiated selection: pin so WS auto-replace handoff respects it. draft.tasks.pinnedSessionId = sessionId; - }); - }, - setActiveSessionAuto: (taskId, sessionId) => { + }), + setActiveSessionAuto: (taskId, sessionId) => set((draft) => { draft.tasks.activeTaskId = taskId; draft.tasks.activeSessionId = sessionId; // Auto-driven (WS) selection: don't touch pinnedSessionId. - }); - }, + }), clearActiveSession: () => set((draft) => { draft.tasks.activeSessionId = null; diff --git a/apps/web/lib/state/slices/session/session-slice.ts b/apps/web/lib/state/slices/session/session-slice.ts index 1ade07a095..79a88d5ac3 100644 --- a/apps/web/lib/state/slices/session/session-slice.ts +++ b/apps/web/lib/state/slices/session/session-slice.ts @@ -58,6 +58,7 @@ function mergeTaskSession(existing: TaskSession, incoming: TaskSession): TaskSes worktree_branch: incoming.worktree_branch ?? existing.worktree_branch, repository_id: incoming.repository_id ?? existing.repository_id, base_branch: incoming.base_branch ?? existing.base_branch, + task_environment_id: incoming.task_environment_id ?? existing.task_environment_id, }; } @@ -288,6 +289,79 @@ function nextPair( return [current[1], revisionId]; } +function buildTaskSessionActions(set: ImmerSet) { + return { + setTaskSession: (session: Parameters[0]) => + set((draft) => { + const existingSession = draft.taskSessions.items[session.id]; + const mergedSession = existingSession + ? mergeTaskSession(existingSession, session) + : session; + draft.taskSessions.items[session.id] = mergedSession; + const sessionsByTask = draft.taskSessionsByTask.itemsByTaskId[session.task_id]; + if (sessionsByTask) { + const sessionIndex = sessionsByTask.findIndex((s) => s.id === session.id); + if (sessionIndex >= 0) sessionsByTask[sessionIndex] = mergedSession; + } + syncEnvironmentMapping(draft, session.id, mergedSession.task_environment_id); + }), + removeTaskSession: (taskId: string, sessionId: string) => + set((draft) => { + delete draft.taskSessions.items[sessionId]; + const sessionsByTask = draft.taskSessionsByTask.itemsByTaskId[taskId]; + if (sessionsByTask) { + draft.taskSessionsByTask.itemsByTaskId[taskId] = sessionsByTask.filter( + (s) => s.id !== sessionId, + ); + } + // eslint-disable-next-line @typescript-eslint/no-explicit-any + delete (draft as any).environmentIdBySessionId[sessionId]; + }), + setTaskSessionsForTask: ( + taskId: string, + sessions: Parameters[1], + ) => + set((draft) => { + const merged = sessions.map((session) => { + const existing = draft.taskSessions.items[session.id]; + return existing ? mergeTaskSession(existing, session) : session; + }); + draft.taskSessionsByTask.itemsByTaskId[taskId] = merged; + draft.taskSessionsByTask.loadingByTaskId[taskId] = false; + draft.taskSessionsByTask.loadedByTaskId[taskId] = true; + for (const session of merged) { + draft.taskSessions.items[session.id] = session; + syncEnvironmentMapping(draft, session.id, session.task_environment_id); + } + }), + // Upsert a session from a WS event without flipping the per-task `loadedByTaskId` + // flag — partial event-driven records must not gate the API hydration that + // fills in fields like agent_profile_id / repository_id / worktree_path. + upsertTaskSessionFromEvent: ( + taskId: string, + session: Parameters[1], + ) => + set((draft) => { + const existing = draft.taskSessions.items[session.id]; + const merged = existing ? mergeTaskSession(existing, session) : session; + draft.taskSessions.items[session.id] = merged; + const list = draft.taskSessionsByTask.itemsByTaskId[taskId]; + if (list) { + const idx = list.findIndex((s) => s.id === session.id); + if (idx >= 0) list[idx] = merged; + else list.push(merged); + } else { + draft.taskSessionsByTask.itemsByTaskId[taskId] = [merged]; + } + syncEnvironmentMapping(draft, session.id, merged.task_environment_id); + }), + setTaskSessionsLoading: (taskId: string, loading: boolean) => + set((draft) => { + draft.taskSessionsByTask.loadingByTaskId[taskId] = loading; + }), + }; +} + export const createSessionSlice: StateCreator< SessionSlice, [["zustand/immer", never]], @@ -313,44 +387,7 @@ export const createSessionSlice: StateCreator< set((draft) => { draft.turns.activeBySession[sessionId] = turnId; }), - setTaskSession: (session) => - set((draft) => { - const existingSession = draft.taskSessions.items[session.id]; - const mergedSession = existingSession ? mergeTaskSession(existingSession, session) : session; - draft.taskSessions.items[session.id] = mergedSession; - const sessionsByTask = draft.taskSessionsByTask.itemsByTaskId[session.task_id]; - if (sessionsByTask) { - const sessionIndex = sessionsByTask.findIndex((s) => s.id === session.id); - if (sessionIndex >= 0) sessionsByTask[sessionIndex] = mergedSession; - } - syncEnvironmentMapping(draft, session.id, mergedSession.task_environment_id); - }), - removeTaskSession: (taskId, sessionId) => - set((draft) => { - delete draft.taskSessions.items[sessionId]; - const sessionsByTask = draft.taskSessionsByTask.itemsByTaskId[taskId]; - if (sessionsByTask) { - draft.taskSessionsByTask.itemsByTaskId[taskId] = sessionsByTask.filter( - (s) => s.id !== sessionId, - ); - } - // eslint-disable-next-line @typescript-eslint/no-explicit-any - delete (draft as any).environmentIdBySessionId[sessionId]; - }), - setTaskSessionsForTask: (taskId, sessions) => - set((draft) => { - draft.taskSessionsByTask.itemsByTaskId[taskId] = sessions; - draft.taskSessionsByTask.loadingByTaskId[taskId] = false; - draft.taskSessionsByTask.loadedByTaskId[taskId] = true; - for (const session of sessions) { - draft.taskSessions.items[session.id] = session; - syncEnvironmentMapping(draft, session.id, session.task_environment_id); - } - }), - setTaskSessionsLoading: (taskId, loading) => - set((draft) => { - draft.taskSessionsByTask.loadingByTaskId[taskId] = loading; - }), + ...buildTaskSessionActions(set), setSessionAgentctlStatus: (sessionId, status) => set((draft) => { draft.sessionAgentctl.itemsBySessionId[sessionId] = status; diff --git a/apps/web/lib/state/slices/session/session-slice.upsert.test.ts b/apps/web/lib/state/slices/session/session-slice.upsert.test.ts new file mode 100644 index 0000000000..1a1fb0443e --- /dev/null +++ b/apps/web/lib/state/slices/session/session-slice.upsert.test.ts @@ -0,0 +1,124 @@ +import { describe, it, expect } from "vitest"; +import { create } from "zustand"; +import { immer } from "zustand/middleware/immer"; +import { createSessionSlice } from "./session-slice"; +import { createSessionRuntimeSlice } from "../session-runtime/session-runtime-slice"; +import type { SessionSlice } from "./types"; +import type { SessionRuntimeSlice } from "../session-runtime/types"; +import type { TaskSession } from "@/lib/types/http"; + +type CombinedSlice = SessionSlice & SessionRuntimeSlice; + +function makeStore() { + return create()( + immer((set) => ({ + // eslint-disable-next-line @typescript-eslint/no-explicit-any + ...(createSessionSlice as any)(set), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + ...(createSessionRuntimeSlice as any)(set), + })), + ); +} + +const TASK_ID = "task-1"; +const SESSION_ID = "session-1"; +const TS = "2026-04-20T00:00:00Z"; + +function makeSession(overrides: Partial = {}): TaskSession { + return { + id: SESSION_ID, + task_id: TASK_ID, + state: "RUNNING", + started_at: TS, + updated_at: TS, + ...overrides, + }; +} + +describe("upsertTaskSessionFromEvent", () => { + it("does not flip loadedByTaskId so API hydration can still run", () => { + const store = makeStore(); + + store.getState().upsertTaskSessionFromEvent(TASK_ID, makeSession()); + + expect(store.getState().taskSessionsByTask.loadedByTaskId[TASK_ID]).toBeFalsy(); + }); + + it("merges fields on a second call rather than replacing the row", () => { + const store = makeStore(); + + store + .getState() + .upsertTaskSessionFromEvent( + TASK_ID, + makeSession({ agent_profile_id: "profile-1", repository_id: "repo-1" }), + ); + // Second event omits fields that were set by the first + store.getState().upsertTaskSessionFromEvent(TASK_ID, makeSession({ state: "COMPLETED" })); + + const session = store.getState().taskSessions.items[SESSION_ID]; + expect(session.state).toBe("COMPLETED"); + expect(session.agent_profile_id).toBe("profile-1"); + expect(session.repository_id).toBe("repo-1"); + }); + + it("seeds environmentIdBySessionId when task_environment_id is present", () => { + const store = makeStore(); + + store + .getState() + .upsertTaskSessionFromEvent(TASK_ID, makeSession({ task_environment_id: "env-1" })); + + expect(store.getState().environmentIdBySessionId[SESSION_ID]).toBe("env-1"); + }); + + it("does not seed environmentIdBySessionId when task_environment_id is absent", () => { + const store = makeStore(); + + store.getState().upsertTaskSessionFromEvent(TASK_ID, makeSession()); + + expect(store.getState().environmentIdBySessionId[SESSION_ID]).toBeUndefined(); + }); + + it("appends to itemsByTaskId when the list already exists", () => { + const store = makeStore(); + const other = makeSession({ id: "session-other" }); + + store.getState().upsertTaskSessionFromEvent(TASK_ID, other); + store.getState().upsertTaskSessionFromEvent(TASK_ID, makeSession()); + + const list = store.getState().taskSessionsByTask.itemsByTaskId[TASK_ID]; + expect(list.map((s) => s.id)).toEqual(["session-other", SESSION_ID]); + }); +}); + +describe("setTaskSessionsForTask preserves WS-seeded fields", () => { + it("merges incoming sessions with existing rows so task_environment_id is not clobbered", () => { + const store = makeStore(); + + // WS event arrives first and seeds task_environment_id + agent_profile_id + store + .getState() + .upsertTaskSessionFromEvent( + TASK_ID, + makeSession({ task_environment_id: "env-1", agent_profile_id: "profile-1" }), + ); + + // API hydration arrives next without task_environment_id (race window) + store.getState().setTaskSessionsForTask(TASK_ID, [makeSession({ repository_id: "repo-1" })]); + + const session = store.getState().taskSessions.items[SESSION_ID]; + expect(session.task_environment_id).toBe("env-1"); + expect(session.agent_profile_id).toBe("profile-1"); + expect(session.repository_id).toBe("repo-1"); + expect(store.getState().environmentIdBySessionId[SESSION_ID]).toBe("env-1"); + }); + + it("flips loadedByTaskId to true (unlike upsertTaskSessionFromEvent)", () => { + const store = makeStore(); + + store.getState().setTaskSessionsForTask(TASK_ID, [makeSession()]); + + expect(store.getState().taskSessionsByTask.loadedByTaskId[TASK_ID]).toBe(true); + }); +}); diff --git a/apps/web/lib/state/slices/session/types.ts b/apps/web/lib/state/slices/session/types.ts index cde68ad071..6cdc91f108 100644 --- a/apps/web/lib/state/slices/session/types.ts +++ b/apps/web/lib/state/slices/session/types.ts @@ -143,6 +143,7 @@ export type SessionSliceActions = { setTaskSession: (session: TaskSession) => void; removeTaskSession: (taskId: string, sessionId: string) => void; setTaskSessionsForTask: (taskId: string, sessions: TaskSession[]) => void; + upsertTaskSessionFromEvent: (taskId: string, session: TaskSession) => void; setTaskSessionsLoading: (taskId: string, loading: boolean) => void; setSessionAgentctlStatus: (sessionId: string, status: SessionAgentctlStatus) => void; setWorktree: (worktree: Worktree) => void; diff --git a/apps/web/lib/state/store.ts b/apps/web/lib/state/store.ts index 0f3b1988a2..c8944b6f2e 100644 --- a/apps/web/lib/state/store.ts +++ b/apps/web/lib/state/store.ts @@ -397,6 +397,7 @@ export type AppState = { setTaskSession: (session: TaskSession) => void; removeTaskSession: (taskId: string, sessionId: string) => void; setTaskSessionsForTask: (taskId: string, sessions: TaskSession[]) => void; + upsertTaskSessionFromEvent: (taskId: string, session: TaskSession) => void; setTaskSessionsLoading: (taskId: string, loading: boolean) => void; setSessionAgentctlStatus: (sessionId: string, status: SessionAgentctlStatus) => void; setWorktree: (worktree: Worktree) => void; diff --git a/apps/web/lib/types/backend.ts b/apps/web/lib/types/backend.ts index 70f93f7cfa..fdb011af44 100644 --- a/apps/web/lib/types/backend.ts +++ b/apps/web/lib/types/backend.ts @@ -252,6 +252,7 @@ export type TaskSessionWaitingForInputPayload = { export type TaskSessionAgentctlPayload = { task_id: string; session_id: string; + task_environment_id?: string; agent_execution_id?: string; error_message?: string; worktree_id?: string; diff --git a/apps/web/lib/ws/handlers/agent-session.test.ts b/apps/web/lib/ws/handlers/agent-session.test.ts index ac49b8221b..f869add3cd 100644 --- a/apps/web/lib/ws/handlers/agent-session.test.ts +++ b/apps/web/lib/ws/handlers/agent-session.test.ts @@ -16,6 +16,7 @@ function makeStore(overrides: Record = {}) { taskSessionsByTask: { itemsByTaskId: {} }, setTaskSession: vi.fn(), setTaskSessionsForTask: vi.fn(), + upsertTaskSessionFromEvent: vi.fn(), setActiveSession: vi.fn(), setActiveSessionAuto: vi.fn(), setSessionFailureNotification: vi.fn(), @@ -453,3 +454,191 @@ describe("session.state_changed → respects user-pinned session", () => { expect(store.getState().setActiveSession).not.toHaveBeenCalled(); }); }); + +// eslint-disable-next-line max-lines-per-function -- test describe block, splitting hurts readability +describe("session.state_changed → agentctl ready fallback", () => { + const TS = "2026-05-04T00:00:00Z"; + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("promotes agentctl status to 'ready' when session enters RUNNING and ready event was missed", () => { + const setSessionAgentctlStatus = vi.fn(); + const store = makeStore({ + taskSessions: { + items: { "s-1": { id: "s-1", task_id: "t-1", state: "STARTING" } }, + }, + sessionAgentctl: { itemsBySessionId: { "s-1": { status: "starting" } } }, + setSessionAgentctlStatus, + }); + const handler = registerTaskSessionHandlers(store)[STATE_CHANGED_EVENT]!; + + handler({ + id: "m", + type: "notification", + action: STATE_CHANGED_EVENT, + timestamp: TS, + payload: { task_id: "t-1", session_id: "s-1", new_state: "RUNNING" }, + }); + + expect(setSessionAgentctlStatus).toHaveBeenCalledWith( + "s-1", + expect.objectContaining({ status: "ready" }), + ); + }); + + it("promotes agentctl status to 'ready' on WAITING_FOR_INPUT even when no prior entry exists", () => { + const setSessionAgentctlStatus = vi.fn(); + const store = makeStore({ + taskSessions: { + items: { "s-1": { id: "s-1", task_id: "t-1", state: "STARTING" } }, + }, + sessionAgentctl: { itemsBySessionId: {} }, + setSessionAgentctlStatus, + }); + const handler = registerTaskSessionHandlers(store)[STATE_CHANGED_EVENT]!; + + handler({ + id: "m", + type: "notification", + action: STATE_CHANGED_EVENT, + timestamp: TS, + payload: { task_id: "t-1", session_id: "s-1", new_state: "WAITING_FOR_INPUT" }, + }); + + expect(setSessionAgentctlStatus).toHaveBeenCalledWith( + "s-1", + expect.objectContaining({ status: "ready" }), + ); + }); + + it("does not re-set 'ready' when the session is already ready", () => { + const setSessionAgentctlStatus = vi.fn(); + const store = makeStore({ + taskSessions: { + items: { "s-1": { id: "s-1", task_id: "t-1", state: "RUNNING" } }, + }, + sessionAgentctl: { itemsBySessionId: { "s-1": { status: "ready" } } }, + setSessionAgentctlStatus, + }); + const handler = registerTaskSessionHandlers(store)[STATE_CHANGED_EVENT]!; + + handler({ + id: "m", + type: "notification", + action: STATE_CHANGED_EVENT, + timestamp: TS, + payload: { task_id: "t-1", session_id: "s-1", new_state: "WAITING_FOR_INPUT" }, + }); + + expect(setSessionAgentctlStatus).not.toHaveBeenCalled(); + }); + + it("seeds env mapping from agentctl_starting payload via upsertTaskSessionFromEvent", () => { + const upsertTaskSessionFromEvent = vi.fn(); + const store = makeStore({ + taskSessions: { + items: { "s-1": { id: "s-1", task_id: "t-1", state: "CREATED" } }, + }, + sessionAgentctl: { itemsBySessionId: {} }, + setSessionAgentctlStatus: vi.fn(), + upsertTaskSessionFromEvent, + }); + const handler = registerTaskSessionHandlers(store)["session.agentctl_starting"]!; + + handler({ + id: "m", + type: "notification", + action: "session.agentctl_starting", + timestamp: TS, + payload: { + task_id: "t-1", + session_id: "s-1", + agent_execution_id: "ae-1", + task_environment_id: "env-1", + }, + }); + + expect(upsertTaskSessionFromEvent).toHaveBeenCalledWith( + "t-1", + expect.objectContaining({ id: "s-1", task_environment_id: "env-1" }), + ); + }); + + it("seeds env mapping from agentctl_ready payload", () => { + const upsertTaskSessionFromEvent = vi.fn(); + const store = makeStore({ + taskSessions: { items: {} }, + sessionAgentctl: { itemsBySessionId: {} }, + setSessionAgentctlStatus: vi.fn(), + upsertTaskSessionFromEvent, + setWorktree: vi.fn(), + sessionWorktreesBySessionId: { itemsBySessionId: {} }, + setSessionWorktrees: vi.fn(), + }); + const handler = registerTaskSessionHandlers(store)["session.agentctl_ready"]!; + + handler({ + id: "m", + type: "notification", + action: "session.agentctl_ready", + timestamp: TS, + payload: { + task_id: "t-1", + session_id: "s-1", + agent_execution_id: "ae-1", + task_environment_id: "env-1", + }, + }); + + expect(upsertTaskSessionFromEvent).toHaveBeenCalledWith( + "t-1", + expect.objectContaining({ id: "s-1", task_environment_id: "env-1" }), + ); + }); + + it("does not call upsertTaskSessionFromEvent when agentctl payload omits task_environment_id", () => { + const upsertTaskSessionFromEvent = vi.fn(); + const store = makeStore({ + taskSessions: { items: {} }, + sessionAgentctl: { itemsBySessionId: {} }, + setSessionAgentctlStatus: vi.fn(), + upsertTaskSessionFromEvent, + }); + const handler = registerTaskSessionHandlers(store)["session.agentctl_starting"]!; + + handler({ + id: "m", + type: "notification", + action: "session.agentctl_starting", + timestamp: TS, + payload: { task_id: "t-1", session_id: "s-1", agent_execution_id: "ae-1" }, + }); + + expect(upsertTaskSessionFromEvent).not.toHaveBeenCalled(); + }); + + it("does not promote on non-live states (STARTING, COMPLETED, FAILED)", () => { + const setSessionAgentctlStatus = vi.fn(); + const store = makeStore({ + taskSessions: { + items: { "s-1": { id: "s-1", task_id: "t-1", state: "CREATED" } }, + }, + sessionAgentctl: { itemsBySessionId: {} }, + setSessionAgentctlStatus, + }); + const handler = registerTaskSessionHandlers(store)[STATE_CHANGED_EVENT]!; + + for (const newState of ["STARTING", "COMPLETED", "FAILED", "CANCELLED"]) { + handler({ + id: "m", + type: "notification", + action: STATE_CHANGED_EVENT, + timestamp: TS, + payload: { task_id: "t-1", session_id: "s-1", new_state: newState }, + }); + } + + expect(setSessionAgentctlStatus).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/lib/ws/handlers/agent-session.ts b/apps/web/lib/ws/handlers/agent-session.ts index 8f119ddb96..5fa14ee84a 100644 --- a/apps/web/lib/ws/handlers/agent-session.ts +++ b/apps/web/lib/ws/handlers/agent-session.ts @@ -10,10 +10,35 @@ const TERMINAL_SESSION_STATES: ReadonlySet = new Set([ "FAILED", ]); +// States that imply agentctl is up and processing (or has fully processed) input. +// If we observe a session in one of these, agentctl must be ready even if we +// missed the agentctl_ready WS event (e.g. it fired before our subscription). +const AGENT_LIVE_STATES: ReadonlySet = new Set(["RUNNING", "WAITING_FOR_INPUT"]); + export function isTerminalSessionState(state: TaskSessionState | undefined): boolean { return !!state && TERMINAL_SESSION_STATES.has(state); } +/** Promote agentctl status to "ready" when the session enters a live state. + * Acts as a fallback for missed/late agentctl_ready WS events — the backend + * cannot reach RUNNING/WAITING_FOR_INPUT without a live agentctl. Never + * downgrades an existing "ready" entry. */ +function maybePromoteAgentctlReady( + store: StoreApi, + sessionId: string, + newState: TaskSessionState | undefined, + timestamp: string | undefined, +): void { + if (!newState || !AGENT_LIVE_STATES.has(newState)) return; + const current = store.getState().sessionAgentctl?.itemsBySessionId?.[sessionId]; + if (current?.status === "ready") return; + store.getState().setSessionAgentctlStatus(sessionId, { + status: "ready", + agentExecutionId: current?.agentExecutionId, + updatedAt: timestamp, + }); +} + /** * When the backend creates a new session for the active task (e.g., due to a * workflow step transition with a different agent profile), the chat UI should @@ -73,7 +98,10 @@ function buildSessionUpdate(payload: any): Record { return update; } -/** Upsert the session in the per-task sessions list. */ +/** Upsert the session in the per-task sessions list from a WS event. + * Uses `upsertTaskSessionFromEvent` so the per-task list is not marked as + * fully loaded — partial event payloads must not gate the API hydration that + * fills in fields like agent_profile_id / repository_id / worktree_path. */ function upsertTaskSessionList( store: StoreApi, taskId: string, @@ -82,49 +110,19 @@ function upsertTaskSessionList( payload: any, sessionUpdate: Record, ): void { - const sessionsByTask = store.getState().taskSessionsByTask.itemsByTaskId[taskId]; const newState = payload.new_state as TaskSessionState | undefined; + const existing = store.getState().taskSessions.items[sessionId]; + if (!existing && !newState) return; - if (!sessionsByTask) { - // Initialize per-task sessions from WS event so the sidebar can - // display up-to-date status even for tasks not yet clicked. - if (newState) { - store.getState().setTaskSessionsForTask(taskId, [ - { - id: sessionId, - task_id: taskId, - state: newState, - started_at: "", - updated_at: "", - ...(payload.agent_profile_id ? { agent_profile_id: payload.agent_profile_id } : {}), - ...sessionUpdate, - }, - ]); - } - return; - } - - const hasSession = sessionsByTask.some((s) => s.id === sessionId); - - if (!hasSession && newState) { - store.getState().setTaskSessionsForTask(taskId, [ - ...sessionsByTask, - { - id: sessionId, - task_id: taskId, - state: newState, - started_at: "", - updated_at: "", - ...(payload.agent_profile_id ? { agent_profile_id: payload.agent_profile_id } : {}), - ...sessionUpdate, - }, - ]); - } else if (hasSession) { - store.getState().setTaskSessionsForTask( - taskId, - sessionsByTask.map((s) => (s.id === sessionId ? { ...s, ...sessionUpdate } : s)), - ); - } + store.getState().upsertTaskSessionFromEvent(taskId, { + id: sessionId, + task_id: taskId, + state: (newState ?? existing?.state) as TaskSessionState, + started_at: existing?.started_at ?? "", + updated_at: existing?.updated_at ?? "", + ...(payload.agent_profile_id ? { agent_profile_id: payload.agent_profile_id } : {}), + ...sessionUpdate, + }); } /** Extract context window data from payload metadata and store it. */ @@ -197,6 +195,32 @@ function maybeAdoptSessionOnTransition( } } +/** Seed the session→environment mapping from an agentctl event payload. + * Brand-new CREATED sessions never receive a `state_changed` carrying + * `task_environment_id` (no transition fires), so the env mapping would + * otherwise stay empty and env-routed shell terminals stall on + * "Connecting terminal...". Routes through `upsertTaskSessionFromEvent` + * which calls `syncEnvironmentMapping` and merges with any existing row. */ +function syncEnvFromAgentctlPayload( + store: StoreApi, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + payload: any, +): void { + const taskId = payload?.task_id; + const sessionId = payload?.session_id; + const envId = payload?.task_environment_id; + if (!taskId || !sessionId || !envId) return; + const existing = store.getState().taskSessions.items[sessionId]; + store.getState().upsertTaskSessionFromEvent(taskId, { + id: sessionId, + task_id: taskId, + state: existing?.state ?? "CREATED", + started_at: existing?.started_at ?? "", + updated_at: existing?.updated_at ?? "", + task_environment_id: envId, + }); +} + /** Handle the agentctl_ready event: update session worktree info. */ // eslint-disable-next-line @typescript-eslint/no-explicit-any function handleAgentctlReady(store: StoreApi, payload: any): void { @@ -285,22 +309,9 @@ export function registerTaskSessionHandlers(store: StoreApi): WsHandle const sessionUpdate = buildSessionUpdate(payload); const existingSession = store.getState().taskSessions.items[sessionId]; - if (existingSession) { - store.getState().setTaskSession({ ...existingSession, ...sessionUpdate }); - } else if (newState) { - store.getState().setTaskSession({ - id: sessionId, - task_id: taskId, - state: newState as TaskSessionState, - started_at: "", - updated_at: "", - ...(payload.agent_profile_id ? { agent_profile_id: payload.agent_profile_id } : {}), - ...sessionUpdate, - }); - } - upsertTaskSessionList(store, taskId, sessionId, payload, sessionUpdate); extractContextWindow(store, sessionId, payload); + maybePromoteAgentctlReady(store, sessionId, newState, message.timestamp); maybeAdoptSessionOnTransition(store, taskId, sessionId, newState, !!existingSession); @@ -320,6 +331,7 @@ export function registerTaskSessionHandlers(store: StoreApi): WsHandle agentExecutionId: payload.agent_execution_id, updatedAt: message.timestamp, }); + syncEnvFromAgentctlPayload(store, payload); }, "session.agentctl_ready": (message) => { const payload = message.payload; @@ -329,6 +341,7 @@ export function registerTaskSessionHandlers(store: StoreApi): WsHandle agentExecutionId: payload.agent_execution_id, updatedAt: message.timestamp, }); + syncEnvFromAgentctlPayload(store, payload); handleAgentctlReady(store, payload); }, "session.agentctl_error": (message) => {