From eb8ffeaca71eefb3f269a16bdebe42c845893fb9 Mon Sep 17 00:00:00 2001 From: Kiet Ho Date: Mon, 2 Mar 2026 17:58:38 -0800 Subject: [PATCH 1/7] feat: add agent launch orchestrator + superset chat --- .../NewWorkspaceModal/NewWorkspaceModal.tsx | 154 ++++++++++-- .../NewWorkspaceCreateFlow.tsx | 15 +- .../adapters/chat-adapter.ts | 138 +++++++++++ .../adapters/terminal-adapter.ts | 78 ++++++ .../agent-session-orchestrator.test.ts | 133 +++++++++++ .../agent-session-orchestrator.ts | 226 ++++++++++++++++++ .../lib/agent-session-orchestrator/index.ts | 11 + .../lib/agent-session-orchestrator/types.ts | 69 ++++++ .../workspaces/useCreateWorkspace.ts | 10 +- .../OpenInWorkspace/OpenInWorkspace.tsx | 154 ++++++++++-- .../tools/start-agent-session.ts | 202 ++++++++++++---- .../main/components/WorkspaceInitEffects.tsx | 107 ++++++++- .../src/renderer/stores/workspace-init.ts | 3 + .../start-agent-session.ts | 191 +++++++++------ packages/shared/package.json | 4 + packages/shared/src/agent-command.ts | 4 +- packages/shared/src/agent-launch.test.ts | 79 ++++++ packages/shared/src/agent-launch.ts | 171 +++++++++++++ packages/shared/src/constants.ts | 2 + 19 files changed, 1567 insertions(+), 184 deletions(-) create mode 100644 apps/desktop/src/renderer/lib/agent-session-orchestrator/adapters/chat-adapter.ts create mode 100644 apps/desktop/src/renderer/lib/agent-session-orchestrator/adapters/terminal-adapter.ts create mode 100644 apps/desktop/src/renderer/lib/agent-session-orchestrator/agent-session-orchestrator.test.ts create mode 100644 apps/desktop/src/renderer/lib/agent-session-orchestrator/agent-session-orchestrator.ts create mode 100644 apps/desktop/src/renderer/lib/agent-session-orchestrator/index.ts create mode 100644 apps/desktop/src/renderer/lib/agent-session-orchestrator/types.ts create mode 100644 packages/shared/src/agent-launch.test.ts create mode 100644 packages/shared/src/agent-launch.ts diff --git a/apps/desktop/src/renderer/components/NewWorkspaceModal/NewWorkspaceModal.tsx b/apps/desktop/src/renderer/components/NewWorkspaceModal/NewWorkspaceModal.tsx index dacdb8a5e09..0ac6a20901c 100644 --- a/apps/desktop/src/renderer/components/NewWorkspaceModal/NewWorkspaceModal.tsx +++ b/apps/desktop/src/renderer/components/NewWorkspaceModal/NewWorkspaceModal.tsx @@ -1,12 +1,19 @@ import { AGENT_PRESET_COMMANDS, - AGENT_TYPES, buildAgentPromptCommand, } from "@superset/shared/agent-command"; +import { + type AgentLaunchRequest, + STARTABLE_AGENT_TYPES, + type StartableAgentType, +} from "@superset/shared/agent-launch"; +import { FEATURE_FLAGS } from "@superset/shared/constants"; import { Dialog, DialogContent } from "@superset/ui/dialog"; import { toast } from "@superset/ui/sonner"; import { useNavigate } from "@tanstack/react-router"; +import { useFeatureFlagEnabled } from "posthog-js/react"; import { useEffect, useMemo, useRef, useState } from "react"; +import { launchAgentSession } from "renderer/lib/agent-session-orchestrator"; import { electronTrpc } from "renderer/lib/electron-trpc"; import { launchCommandInPane } from "renderer/lib/terminal/launch-command"; import { resolveEffectiveWorkspaceBaseBranch } from "renderer/lib/workspaceBaseBranch"; @@ -41,6 +48,9 @@ const WORKSPACE_AGENT_STORAGE_KEY = "lastSelectedWorkspaceCreateAgent"; export function NewWorkspaceModal() { const navigate = useNavigate(); + const isOrchestratorEnabled = useFeatureFlagEnabled( + FEATURE_FLAGS.DESKTOP_AGENT_LAUNCH_ORCHESTRATOR_V1, + ); const isOpen = useNewWorkspaceModalOpen(); const closeModal = useCloseNewWorkspaceModal(); const preSelectedProjectId = usePreSelectedProjectId(); @@ -62,7 +72,8 @@ export function NewWorkspaceModal() { if (typeof window === "undefined") return "none"; const stored = window.localStorage.getItem(WORKSPACE_AGENT_STORAGE_KEY); if (stored === "none") return "none"; - return stored && (AGENT_TYPES as readonly string[]).includes(stored) + return stored && + (STARTABLE_AGENT_TYPES as readonly string[]).includes(stored) ? (stored as WorkspaceCreateAgent) : "none"; }, @@ -103,6 +114,15 @@ export function NewWorkspaceModal() { const removePane = useTabsStore((s) => s.removePane); const setTabAutoTitle = useTabsStore((s) => s.setTabAutoTitle); const { openNew } = useOpenProject(); + const selectableAgents = useMemo( + () => + isOrchestratorEnabled + ? (STARTABLE_AGENT_TYPES as readonly StartableAgentType[]) + : (STARTABLE_AGENT_TYPES.filter( + (agent) => agent !== "superset-chat", + ) as readonly StartableAgentType[]), + [isOrchestratorEnabled], + ); const resolvedPrefix = useMemo(() => { const projectOverrides = project?.branchPrefixMode != null; @@ -136,6 +156,13 @@ export function NewWorkspaceModal() { } }, [isOpen]); + useEffect(() => { + if (selectedAgent === "none") return; + if (selectableAgents.includes(selectedAgent)) return; + setSelectedAgent("none"); + window.localStorage.setItem(WORKSPACE_AGENT_STORAGE_KEY, "none"); + }, [selectedAgent, selectableAgents]); + const effectiveBaseBranch = resolveEffectiveWorkspaceBaseBranch({ explicitBaseBranch: baseBranch, workspaceBaseBranch: project?.workspaceBaseBranch, @@ -249,6 +276,50 @@ export function NewWorkspaceModal() { /> ); const isCreateDisabled = createWorkspace.isPending || isBranchesError; + const buildLaunchRequestForWorkspace = ( + workspaceId: string, + prompt: string, + ): AgentLaunchRequest | null => { + if (selectedAgent === "none") { + return null; + } + + if (selectedAgent === "superset-chat") { + return { + kind: "chat", + workspaceId, + agentType: "superset-chat", + source: "new-workspace", + chat: { + initialPrompt: prompt || undefined, + retryCount: 1, + }, + }; + } + + const command = prompt + ? buildAgentPromptCommand({ + prompt, + randomId: window.crypto.randomUUID(), + agent: selectedAgent, + }) + : (AGENT_PRESET_COMMANDS[selectedAgent][0] ?? null); + + if (!command) { + return null; + } + + return { + kind: "terminal", + workspaceId, + agentType: selectedAgent, + source: "new-workspace", + terminal: { + command, + name: "Agent", + }, + }; + }; const handleCreateWorkspace = async () => { if (!selectedProjectId) return; @@ -256,16 +327,21 @@ export function NewWorkspaceModal() { const prompt = title.trim(); const workspaceName = deriveWorkspaceTitleFromPrompt(title) || undefined; - const agentCommand = - selectedAgent === "none" - ? null - : prompt + const launchRequestTemplate = isOrchestratorEnabled + ? buildLaunchRequestForWorkspace("pending-workspace", prompt) + : null; + const legacyAgentCommand = + !isOrchestratorEnabled && + selectedAgent !== "none" && + selectedAgent !== "superset-chat" + ? prompt ? buildAgentPromptCommand({ prompt, randomId: window.crypto.randomUUID(), agent: selectedAgent, }) - : (AGENT_PRESET_COMMANDS[selectedAgent][0] ?? null); + : (AGENT_PRESET_COMMANDS[selectedAgent][0] ?? null) + : null; closeModal(); @@ -278,34 +354,59 @@ export function NewWorkspaceModal() { baseBranch: baseBranch || undefined, applyPrefix, }, - agentCommand ? { agentCommand } : undefined, + launchRequestTemplate + ? { agentLaunchRequest: launchRequestTemplate } + : legacyAgentCommand + ? { agentCommand: legacyAgentCommand } + : undefined, ); - if (agentCommand) { + const launchRequest = launchRequestTemplate + ? { + ...launchRequestTemplate, + workspaceId: result.workspace.id, + } + : null; + + if (launchRequest && isOrchestratorEnabled) { if (result.wasExisting) { - const { tabId, paneId } = addTab(result.workspace.id); - setTabAutoTitle(tabId, "Agent"); - try { - await launchCommandInPane({ - paneId, - tabId, - workspaceId: result.workspace.id, - command: agentCommand, - createOrAttach: (input) => - terminalCreateOrAttach.mutateAsync(input), - write: (input) => terminalWrite.mutateAsync(input), - }); - } catch (error) { - removePane(paneId); + const launchResult = await launchAgentSession(launchRequest, { + source: "new-workspace", + createOrAttach: (input) => + terminalCreateOrAttach.mutateAsync(input), + write: (input) => terminalWrite.mutateAsync(input), + }); + if (launchResult.status === "failed") { toast.error("Failed to start agent", { description: - error instanceof Error - ? error.message - : "Failed to start agent terminal session.", + launchResult.error ?? "Failed to start agent session.", }); return; } } + } else if (legacyAgentCommand && result.wasExisting) { + const { tabId, paneId } = addTab(result.workspace.id); + setTabAutoTitle(tabId, "Agent"); + try { + await launchCommandInPane({ + paneId, + tabId, + workspaceId: result.workspace.id, + command: legacyAgentCommand, + createOrAttach: (input) => + terminalCreateOrAttach.mutateAsync(input), + write: (input) => terminalWrite.mutateAsync(input), + }); + } catch (error) { + removePane(paneId); + toast.error("Failed to start agent", { + description: + error instanceof Error + ? error.message + : "Failed to start agent terminal session.", + }); + return; + } } if (result.isInitializing) { @@ -383,6 +484,7 @@ export function NewWorkspaceModal() { void; title: string; onTitleChange: (value: string) => void; @@ -42,6 +42,7 @@ interface NewWorkspaceCreateFlowProps { export function NewWorkspaceCreateFlow({ projectSelector, selectedAgent, + agentOptions, onSelectedAgentChange, title, onTitleChange, @@ -80,7 +81,7 @@ export function NewWorkspaceCreateFlow({ No agent - {AGENT_TYPES.map((agent) => { + {agentOptions.map((agent) => { const icon = getPresetIcon(agent, isDark); return ( @@ -92,7 +93,7 @@ export function NewWorkspaceCreateFlow({ className="size-3.5 object-contain" /> )} - {AGENT_LABELS[agent]} + {STARTABLE_AGENT_LABELS[agent]} ); diff --git a/apps/desktop/src/renderer/lib/agent-session-orchestrator/adapters/chat-adapter.ts b/apps/desktop/src/renderer/lib/agent-session-orchestrator/adapters/chat-adapter.ts new file mode 100644 index 00000000000..3392e231be8 --- /dev/null +++ b/apps/desktop/src/renderer/lib/agent-session-orchestrator/adapters/chat-adapter.ts @@ -0,0 +1,138 @@ +import type { AgentLaunchRequest } from "@superset/shared/agent-launch"; +import type { AgentSessionLaunchContext, LaunchResultPayload } from "../types"; + +type ChatLaunchRequest = Extract; + +type ChatClient = { + session: { + sendMessage: { + mutate: (input: { + sessionId: string; + payload: { content: string }; + metadata?: { model?: string }; + }) => Promise; + }; + }; +}; + +let chatClientPromise: Promise | null = null; + +async function getChatClient(): Promise { + if (!chatClientPromise) { + chatClientPromise = import( + "renderer/screens/main/components/WorkspaceView/ContentView/TabsContent/TabView/ChatMastraPane/utils/chat-mastra-service-client" + ).then((module) => module.createChatMastraServiceIpcClient()); + } + return chatClientPromise; +} + +async function defaultSendChatMessage(input: { + sessionId: string; + prompt: string; + model?: string; +}) { + const chatClient = await getChatClient(); + await chatClient.session.sendMessage.mutate({ + sessionId: input.sessionId, + payload: { content: input.prompt }, + metadata: input.model ? { model: input.model } : undefined, + }); +} + +async function sendInitialPromptWithRetry({ + context, + sessionId, + prompt, + model, + retryCount, +}: { + context: AgentSessionLaunchContext; + sessionId: string; + prompt: string; + model?: string; + retryCount: number; +}) { + const send = context.sendChatMessage ?? defaultSendChatMessage; + let attempt = 0; + let lastError: unknown; + + while (attempt <= retryCount) { + try { + await send({ sessionId, prompt, model }); + return; + } catch (error) { + lastError = error; + attempt += 1; + if (attempt > retryCount) { + throw lastError; + } + } + } +} + +export async function launchChatAdapter( + request: ChatLaunchRequest, + context: AgentSessionLaunchContext, +): Promise { + const tabs = context.tabs; + if (!tabs) { + throw new Error("Missing tabs adapter"); + } + + let tabId: string; + let paneId: string; + + const targetPaneId = request.chat.paneId; + if (targetPaneId) { + const targetPane = tabs.getPane(targetPaneId); + if (!targetPane) { + throw new Error(`Pane not found: ${targetPaneId}`); + } + const tab = tabs.getTab(targetPane.tabId); + if (!tab || tab.workspaceId !== request.workspaceId) { + throw new Error(`Tab not found for pane: ${targetPaneId}`); + } + + if (targetPane.type === "chat-mastra") { + tabId = tab.id; + paneId = targetPane.id; + } else { + const created = tabs.addChatTab(request.workspaceId); + tabId = created.tabId; + paneId = created.paneId; + } + } else { + const created = tabs.addChatTab(request.workspaceId); + tabId = created.tabId; + paneId = created.paneId; + } + + tabs.setTabAutoTitle(tabId, "Superset Chat"); + + const pane = tabs.getPane(paneId); + let sessionId = request.chat.sessionId ?? pane?.chatMastra?.sessionId ?? null; + if (!sessionId) { + sessionId = crypto.randomUUID(); + } + + if (pane?.chatMastra?.sessionId !== sessionId) { + tabs.switchChatSession(paneId, sessionId); + } + + const initialPrompt = request.chat.initialPrompt?.trim(); + if (initialPrompt) { + await sendInitialPromptWithRetry({ + context, + sessionId, + prompt: initialPrompt, + model: request.chat.model, + retryCount: request.chat.retryCount ?? 0, + }); + } + + return { + tabId, + paneId, + sessionId, + }; +} diff --git a/apps/desktop/src/renderer/lib/agent-session-orchestrator/adapters/terminal-adapter.ts b/apps/desktop/src/renderer/lib/agent-session-orchestrator/adapters/terminal-adapter.ts new file mode 100644 index 00000000000..d2e8eb13998 --- /dev/null +++ b/apps/desktop/src/renderer/lib/agent-session-orchestrator/adapters/terminal-adapter.ts @@ -0,0 +1,78 @@ +import type { AgentLaunchRequest } from "@superset/shared/agent-launch"; +import { launchCommandInPane } from "renderer/lib/terminal/launch-command"; +import type { AgentSessionLaunchContext, LaunchResultPayload } from "../types"; + +type TerminalLaunchRequest = Extract; + +export async function launchTerminalAdapter( + request: TerminalLaunchRequest, + context: AgentSessionLaunchContext, +): Promise { + const tabs = context.tabs; + if (!tabs) { + throw new Error("Missing tabs adapter"); + } + + const { workspaceId } = request; + const targetPaneId = request.terminal.paneId; + + if (targetPaneId) { + const targetPane = tabs.getPane(targetPaneId); + if (!targetPane) { + throw new Error(`Pane not found: ${targetPaneId}`); + } + + const tab = tabs.getTab(targetPane.tabId); + if (!tab || tab.workspaceId !== workspaceId) { + throw new Error(`Tab not found for pane: ${targetPaneId}`); + } + + const newPaneId = tabs.addTerminalPane(tab.id); + if (!newPaneId) { + throw new Error("Failed to add pane"); + } + + try { + await launchCommandInPane({ + paneId: newPaneId, + tabId: tab.id, + workspaceId, + command: request.terminal.command, + createOrAttach: context.createOrAttach, + write: context.write, + }); + } catch (error) { + tabs.removePane(newPaneId); + throw error; + } + + return { + tabId: tab.id, + paneId: newPaneId, + sessionId: null, + }; + } + + const { tabId, paneId } = tabs.addTerminalTab(workspaceId); + tabs.setTabAutoTitle(tabId, request.terminal.name ?? "Agent"); + + try { + await launchCommandInPane({ + paneId, + tabId, + workspaceId, + command: request.terminal.command, + createOrAttach: context.createOrAttach, + write: context.write, + }); + } catch (error) { + tabs.removePane(paneId); + throw error; + } + + return { + tabId, + paneId, + sessionId: null, + }; +} diff --git a/apps/desktop/src/renderer/lib/agent-session-orchestrator/agent-session-orchestrator.test.ts b/apps/desktop/src/renderer/lib/agent-session-orchestrator/agent-session-orchestrator.test.ts new file mode 100644 index 00000000000..c1bfe471041 --- /dev/null +++ b/apps/desktop/src/renderer/lib/agent-session-orchestrator/agent-session-orchestrator.test.ts @@ -0,0 +1,133 @@ +import { describe, expect, it, mock } from "bun:test"; +import type { AgentLaunchRequest } from "@superset/shared/agent-launch"; +import { + launchAgentSession, + selectAgentLaunchAdapter, +} from "./agent-session-orchestrator"; +import type { AgentLaunchTabsAdapter } from "./types"; + +function createDeferred() { + let resolve!: () => void; + const promise = new Promise((res) => { + resolve = res; + }); + return { promise, resolve }; +} + +function createContext({ + tabs, + write, +}: { + tabs: AgentLaunchTabsAdapter; + write?: (input: { + paneId: string; + data: string; + throwOnError?: boolean; + }) => Promise; +}) { + return { + source: "command-watcher" as const, + tabs, + createOrAttach: mock(async () => ({})), + write: write ?? mock(async () => ({})), + captureEvent: mock(() => {}), + }; +} + +describe("selectAgentLaunchAdapter", () => { + it("picks terminal adapter for terminal requests", () => { + const request: AgentLaunchRequest = { + kind: "terminal", + workspaceId: "ws-1", + terminal: { command: "echo hello" }, + }; + + expect(selectAgentLaunchAdapter(request)).toBe("terminal"); + }); + + it("picks chat adapter for chat requests", () => { + const request: AgentLaunchRequest = { + kind: "chat", + workspaceId: "ws-1", + chat: {}, + }; + + expect(selectAgentLaunchAdapter(request)).toBe("chat"); + }); +}); + +describe("launchAgentSession", () => { + it("deduplicates concurrent launches with the same idempotency key", async () => { + const gate = createDeferred(); + const addTerminalTab = mock(() => ({ tabId: "tab-1", paneId: "pane-1" })); + const tabs: AgentLaunchTabsAdapter = { + getPane: mock(() => undefined), + getTab: mock(() => undefined), + addTerminalTab, + addTerminalPane: mock(() => "pane-2"), + removePane: mock(() => {}), + setTabAutoTitle: mock(() => {}), + addChatTab: mock(() => ({ tabId: "chat-tab", paneId: "chat-pane" })), + switchChatSession: mock(() => {}), + }; + + const context = createContext({ + tabs, + write: async () => { + await gate.promise; + }, + }); + const request: AgentLaunchRequest = { + kind: "terminal", + workspaceId: "ws-1", + idempotencyKey: "idem-concurrent", + terminal: { command: "echo hello" }, + }; + + const first = launchAgentSession(request, context); + const second = launchAgentSession(request, context); + + gate.resolve(); + const [firstResult, secondResult] = await Promise.all([first, second]); + + expect(addTerminalTab).toHaveBeenCalledTimes(1); + expect(firstResult.status).toBe("running"); + expect(secondResult.status).toBe("running"); + expect(firstResult.tabId).toBe("tab-1"); + expect(secondResult.tabId).toBe("tab-1"); + }); + + it("rolls back pane when terminal launch fails", async () => { + const removePane = mock(() => {}); + const tabs: AgentLaunchTabsAdapter = { + getPane: mock(() => undefined), + getTab: mock(() => undefined), + addTerminalTab: mock(() => ({ tabId: "tab-2", paneId: "pane-2" })), + addTerminalPane: mock(() => "pane-3"), + removePane, + setTabAutoTitle: mock(() => {}), + addChatTab: mock(() => ({ tabId: "chat-tab", paneId: "chat-pane" })), + switchChatSession: mock(() => {}), + }; + + const context = createContext({ + tabs, + write: async () => { + throw new Error("terminal write failed"); + }, + }); + + const result = await launchAgentSession( + { + kind: "terminal", + workspaceId: "ws-1", + terminal: { command: "echo fail" }, + }, + context, + ); + + expect(removePane).toHaveBeenCalledWith("pane-2"); + expect(result.status).toBe("failed"); + expect(result.error).toContain("terminal write failed"); + }); +}); diff --git a/apps/desktop/src/renderer/lib/agent-session-orchestrator/agent-session-orchestrator.ts b/apps/desktop/src/renderer/lib/agent-session-orchestrator/agent-session-orchestrator.ts new file mode 100644 index 00000000000..451f3729f3f --- /dev/null +++ b/apps/desktop/src/renderer/lib/agent-session-orchestrator/agent-session-orchestrator.ts @@ -0,0 +1,226 @@ +import type { + AgentLaunchRequest, + AgentLaunchResult, +} from "@superset/shared/agent-launch"; +import { normalizeAgentLaunchRequest } from "@superset/shared/agent-launch"; +import { posthog } from "renderer/lib/posthog"; +import { useWorkspaceInitStore } from "renderer/stores/workspace-init"; +import { launchChatAdapter } from "./adapters/chat-adapter"; +import { launchTerminalAdapter } from "./adapters/terminal-adapter"; +import type { + AgentLaunchTabsAdapter, + AgentSessionLaunchAdapterKind, + AgentSessionLaunchContext, + QueueAgentSessionLaunchInput, +} from "./types"; + +const inFlightByIdempotency = new Map>(); +const settledByIdempotency = new Map(); + +async function getDefaultTabsAdapter(): Promise { + const { useTabsStore } = await import("renderer/stores/tabs/store"); + return { + getPane: (paneId) => useTabsStore.getState().panes[paneId], + getTab: (tabId) => + useTabsStore.getState().tabs.find((tab) => tab.id === tabId), + addTerminalTab: (workspaceId) => + useTabsStore.getState().addTab(workspaceId), + addTerminalPane: (tabId) => useTabsStore.getState().addPane(tabId), + removePane: (paneId) => useTabsStore.getState().removePane(paneId), + setTabAutoTitle: (tabId, title) => + useTabsStore.getState().setTabAutoTitle(tabId, title), + addChatTab: (workspaceId) => + useTabsStore.getState().addChatMastraTab(workspaceId), + switchChatSession: (paneId, sessionId) => + useTabsStore.getState().switchChatMastraSession(paneId, sessionId), + }; +} + +function buildIdempotencyKey(request: AgentLaunchRequest): string | null { + if (!request.idempotencyKey) { + return null; + } + return `${request.workspaceId}:${request.idempotencyKey}`; +} + +function toErrorMessage(error: unknown): string { + return error instanceof Error + ? error.message + : "Failed to launch agent session"; +} + +function captureLaunchEvent({ + context, + request, + status, + latencyMs, + error, +}: { + context: AgentSessionLaunchContext; + request: AgentLaunchRequest; + status: AgentLaunchResult["status"]; + latencyMs: number; + error?: string; +}) { + const capture = + context.captureEvent ?? + (({ + event, + properties, + }: { + event: "agent_session_launch"; + properties: Record; + }) => { + posthog.capture(event, properties); + }); + + capture({ + event: "agent_session_launch", + properties: { + launch_source: request.source ?? context.source ?? "unknown", + request_kind: request.kind, + agent_type: request.agentType ?? null, + result: status, + latency_ms: latencyMs, + failure_reason: error ?? null, + }, + }); +} + +export function selectAgentLaunchAdapter( + request: AgentLaunchRequest, +): AgentSessionLaunchAdapterKind { + return request.kind === "chat" ? "chat" : "terminal"; +} + +export async function launchAgentSession( + requestInput: AgentLaunchRequest | unknown, + context: AgentSessionLaunchContext, +): Promise { + const normalized = normalizeAgentLaunchRequest(requestInput); + const request: AgentLaunchRequest = normalized.source + ? normalized + : { + ...normalized, + source: context.source, + }; + + const idempotencyKey = buildIdempotencyKey(request); + if (idempotencyKey) { + const settled = settledByIdempotency.get(idempotencyKey); + if (settled) { + return settled; + } + const inFlight = inFlightByIdempotency.get(idempotencyKey); + if (inFlight) { + return inFlight; + } + } + + const startedAt = Date.now(); + let phase: AgentLaunchResult["status"] = "queued"; + + const run = (async () => { + try { + const tabs = context.tabs ?? (await getDefaultTabsAdapter()); + const executionContext: AgentSessionLaunchContext = { + ...context, + tabs, + }; + phase = "launching"; + const payload = + request.kind === "terminal" + ? await launchTerminalAdapter(request, executionContext) + : await launchChatAdapter(request, executionContext); + phase = "running"; + const result: AgentLaunchResult = { + workspaceId: request.workspaceId, + tabId: payload.tabId ?? null, + paneId: payload.paneId ?? null, + sessionId: payload.sessionId ?? null, + status: phase, + error: null, + }; + captureLaunchEvent({ + context: executionContext, + request, + status: result.status, + latencyMs: Date.now() - startedAt, + }); + return result; + } catch (error) { + const executionContext: AgentSessionLaunchContext = { + ...context, + }; + phase = "failed"; + const errorMessage = toErrorMessage(error); + const result: AgentLaunchResult = { + workspaceId: request.workspaceId, + tabId: null, + paneId: null, + sessionId: null, + status: phase, + error: errorMessage, + }; + captureLaunchEvent({ + context: executionContext, + request, + status: result.status, + latencyMs: Date.now() - startedAt, + error: errorMessage, + }); + return result; + } + })(); + + if (idempotencyKey) { + inFlightByIdempotency.set(idempotencyKey, run); + } + + const result = await run; + + if (idempotencyKey) { + inFlightByIdempotency.delete(idempotencyKey); + settledByIdempotency.set(idempotencyKey, result); + } + + return result; +} + +export function queueAgentSessionLaunch( + input: QueueAgentSessionLaunchInput, +): AgentLaunchResult { + const request = normalizeAgentLaunchRequest(input.request); + const store = useWorkspaceInitStore.getState(); + const existing = store.pendingTerminalSetups[request.workspaceId]; + const projectId = input.projectId ?? existing?.projectId; + + if (!projectId) { + return { + workspaceId: request.workspaceId, + tabId: null, + paneId: null, + sessionId: null, + status: "failed", + error: `Project ID is required to queue launch for workspace ${request.workspaceId}`, + }; + } + + store.addPendingTerminalSetup({ + workspaceId: request.workspaceId, + projectId, + initialCommands: existing?.initialCommands ?? input.initialCommands ?? null, + defaultPresets: existing?.defaultPresets ?? input.defaultPresets, + agentCommand: existing?.agentCommand, + agentLaunchRequest: request, + }); + + return { + workspaceId: request.workspaceId, + tabId: null, + paneId: null, + sessionId: null, + status: "queued", + error: null, + }; +} diff --git a/apps/desktop/src/renderer/lib/agent-session-orchestrator/index.ts b/apps/desktop/src/renderer/lib/agent-session-orchestrator/index.ts new file mode 100644 index 00000000000..ffef684648e --- /dev/null +++ b/apps/desktop/src/renderer/lib/agent-session-orchestrator/index.ts @@ -0,0 +1,11 @@ +export { + launchAgentSession, + queueAgentSessionLaunch, + selectAgentLaunchAdapter, +} from "./agent-session-orchestrator"; +export type { + AgentLaunchTabsAdapter, + AgentSessionLaunchAdapterKind, + AgentSessionLaunchContext, + QueueAgentSessionLaunchInput, +} from "./types"; diff --git a/apps/desktop/src/renderer/lib/agent-session-orchestrator/types.ts b/apps/desktop/src/renderer/lib/agent-session-orchestrator/types.ts new file mode 100644 index 00000000000..a11bbe62d29 --- /dev/null +++ b/apps/desktop/src/renderer/lib/agent-session-orchestrator/types.ts @@ -0,0 +1,69 @@ +import type { TerminalPreset } from "@superset/local-db"; +import type { + AgentLaunchRequest, + AgentLaunchResult, + AgentLaunchSource, +} from "@superset/shared/agent-launch"; + +export interface AgentLaunchPane { + id: string; + tabId: string; + type: string; + chatMastra?: { + sessionId: string | null; + }; +} + +export interface AgentLaunchTab { + id: string; + workspaceId: string; +} + +export interface AgentLaunchTabsAdapter { + getPane: (paneId: string) => AgentLaunchPane | undefined; + getTab: (tabId: string) => AgentLaunchTab | undefined; + addTerminalTab: (workspaceId: string) => { tabId: string; paneId: string }; + addTerminalPane: (tabId: string) => string; + removePane: (paneId: string) => void; + setTabAutoTitle: (tabId: string, title: string) => void; + addChatTab: (workspaceId: string) => { tabId: string; paneId: string }; + switchChatSession: (paneId: string, sessionId: string | null) => void; +} + +export interface AgentSessionLaunchContext { + source?: AgentLaunchSource; + tabs?: AgentLaunchTabsAdapter; + createOrAttach: (input: { + paneId: string; + tabId: string; + workspaceId: string; + }) => Promise; + write: (input: { + paneId: string; + data: string; + throwOnError?: boolean; + }) => Promise; + sendChatMessage?: (input: { + sessionId: string; + prompt: string; + model?: string; + }) => Promise; + captureEvent?: (input: { + event: "agent_session_launch"; + properties: Record; + }) => void; +} + +export interface QueueAgentSessionLaunchInput { + request: AgentLaunchRequest | unknown; + projectId?: string; + initialCommands?: string[] | null; + defaultPresets?: TerminalPreset[]; +} + +export type AgentSessionLaunchAdapterKind = "terminal" | "chat"; + +export type LaunchResultPayload = Pick< + AgentLaunchResult, + "tabId" | "paneId" | "sessionId" +>; diff --git a/apps/desktop/src/renderer/react-query/workspaces/useCreateWorkspace.ts b/apps/desktop/src/renderer/react-query/workspaces/useCreateWorkspace.ts index e9a9fabd940..df4189e5e97 100644 --- a/apps/desktop/src/renderer/react-query/workspaces/useCreateWorkspace.ts +++ b/apps/desktop/src/renderer/react-query/workspaces/useCreateWorkspace.ts @@ -17,7 +17,7 @@ interface UseCreateWorkspaceOptions extends NonNullable { type PendingSetupOverrides = Pick< PendingTerminalSetup, - "defaultPresets" | "agentCommand" + "defaultPresets" | "agentCommand" | "agentLaunchRequest" >; export function useCreateWorkspace(options?: UseCreateWorkspaceOptions) { @@ -65,6 +65,13 @@ export function useCreateWorkspace(options?: UseCreateWorkspaceOptions) { } if (!data.wasExisting) { + const normalizedLaunchRequest = + pendingSetupOverrides?.agentLaunchRequest + ? { + ...pendingSetupOverrides.agentLaunchRequest, + workspaceId: data.workspace.id, + } + : undefined; addPendingTerminalSetup({ workspaceId: data.workspace.id, projectId: data.projectId, @@ -73,6 +80,7 @@ export function useCreateWorkspace(options?: UseCreateWorkspaceOptions) { : data.initialCommands, defaultPresets: pendingSetupOverrides?.defaultPresets, agentCommand: pendingSetupOverrides?.agentCommand, + agentLaunchRequest: normalizedLaunchRequest, }); } diff --git a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/tasks/$taskId/components/PropertiesSidebar/components/OpenInWorkspace/OpenInWorkspace.tsx b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/tasks/$taskId/components/PropertiesSidebar/components/OpenInWorkspace/OpenInWorkspace.tsx index 49050b16524..698f9994647 100644 --- a/apps/desktop/src/renderer/routes/_authenticated/_dashboard/tasks/$taskId/components/PropertiesSidebar/components/OpenInWorkspace/OpenInWorkspace.tsx +++ b/apps/desktop/src/renderer/routes/_authenticated/_dashboard/tasks/$taskId/components/PropertiesSidebar/components/OpenInWorkspace/OpenInWorkspace.tsx @@ -1,8 +1,11 @@ +import { buildAgentTaskPrompt } from "@superset/shared/agent-command"; import { - AGENT_LABELS, - AGENT_TYPES, - type AgentType, -} from "@superset/shared/agent-command"; + type AgentLaunchRequest, + STARTABLE_AGENT_LABELS, + STARTABLE_AGENT_TYPES, + type StartableAgentType, +} from "@superset/shared/agent-launch"; +import { FEATURE_FLAGS } from "@superset/shared/constants"; import { Button } from "@superset/ui/button"; import { DropdownMenu, @@ -18,12 +21,14 @@ import { SelectValue, } from "@superset/ui/select"; import { toast } from "@superset/ui/sonner"; +import { useFeatureFlagEnabled } from "posthog-js/react"; import { useEffect, useState } from "react"; import { HiArrowRight, HiChevronDown } from "react-icons/hi2"; import { getPresetIcon, useIsDarkTheme, } from "renderer/assets/app-icons/preset-icons"; +import { launchAgentSession } from "renderer/lib/agent-session-orchestrator"; import { electronTrpc } from "renderer/lib/electron-trpc"; import { launchCommandInPane } from "renderer/lib/terminal/launch-command"; import { useCreateWorkspace } from "renderer/react-query/workspaces"; @@ -38,6 +43,9 @@ interface OpenInWorkspaceProps { } export function OpenInWorkspace({ task }: OpenInWorkspaceProps) { + const isOrchestratorEnabled = useFeatureFlagEnabled( + FEATURE_FLAGS.DESKTOP_AGENT_LAUNCH_ORCHESTRATOR_V1, + ); const { data: recentProjects = [] } = electronTrpc.projects.getRecents.useQuery(); const createWorkspace = useCreateWorkspace(); @@ -48,13 +56,19 @@ export function OpenInWorkspace({ task }: OpenInWorkspaceProps) { const removePane = useTabsStore((s) => s.removePane); const setTabAutoTitle = useTabsStore((s) => s.setTabAutoTitle); const isDark = useIsDarkTheme(); + const selectableAgents = isOrchestratorEnabled + ? (STARTABLE_AGENT_TYPES as readonly StartableAgentType[]) + : (STARTABLE_AGENT_TYPES.filter( + (agent) => agent !== "superset-chat", + ) as readonly StartableAgentType[]); const [selectedProjectId, setSelectedProjectId] = useState( () => localStorage.getItem("lastOpenedInProjectId"), ); - const [selectedAgent, setSelectedAgent] = useState(() => { + const [selectedAgent, setSelectedAgent] = useState(() => { const stored = localStorage.getItem("lastSelectedAgent"); - return stored && (AGENT_TYPES as readonly string[]).includes(stored) - ? (stored as AgentType) + return stored && + (STARTABLE_AGENT_TYPES as readonly string[]).includes(stored) + ? (stored as StartableAgentType) : "claude"; }); @@ -70,29 +84,87 @@ export function OpenInWorkspace({ task }: OpenInWorkspaceProps) { } }, [selectedProjectId, recentProjects]); + useEffect(() => { + if (selectableAgents.includes(selectedAgent)) return; + setSelectedAgent("claude"); + localStorage.setItem("lastSelectedAgent", "claude"); + }, [selectedAgent, selectableAgents]); + const handleOpen = async () => { if (!effectiveProjectId) return; await handleSelectProject(effectiveProjectId); }; + const buildLaunchRequest = (workspaceId: string): AgentLaunchRequest => { + if (selectedAgent === "superset-chat") { + return { + kind: "chat", + workspaceId, + agentType: "superset-chat", + source: "open-in-workspace", + chat: { + initialPrompt: buildAgentTaskPrompt({ + id: task.id, + slug: task.slug, + title: task.title, + description: task.description, + priority: task.priority, + statusName: task.status.name, + labels: task.labels, + }), + retryCount: 1, + }, + }; + } + + return { + kind: "terminal", + workspaceId, + agentType: selectedAgent, + source: "open-in-workspace", + terminal: { + command: buildAgentCommand({ + task: { + id: task.id, + slug: task.slug, + title: task.title, + description: task.description, + priority: task.priority, + statusName: task.status.name, + labels: task.labels, + }, + randomId: window.crypto.randomUUID(), + agent: selectedAgent, + }), + name: task.slug, + }, + }; + }; + const handleSelectProject = async (projectId: string) => { const branchName = deriveBranchName({ slug: task.slug, title: task.title, }); - const command = buildAgentCommand({ - task: { - id: task.id, - slug: task.slug, - title: task.title, - description: task.description, - priority: task.priority, - statusName: task.status.name, - labels: task.labels, - }, - randomId: window.crypto.randomUUID(), - agent: selectedAgent, - }); + const launchRequestTemplate = isOrchestratorEnabled + ? buildLaunchRequest("pending-workspace") + : null; + const legacyCommand = + !isOrchestratorEnabled && selectedAgent !== "superset-chat" + ? buildAgentCommand({ + task: { + id: task.id, + slug: task.slug, + title: task.title, + description: task.description, + priority: task.priority, + statusName: task.status.name, + labels: task.labels, + }, + randomId: window.crypto.randomUUID(), + agent: selectedAgent, + }) + : null; try { const result = await createWorkspace.mutateAsyncWithPendingSetup( @@ -101,10 +173,40 @@ export function OpenInWorkspace({ task }: OpenInWorkspaceProps) { name: task.slug, branchName, }, - { agentCommand: command }, + launchRequestTemplate + ? { agentLaunchRequest: launchRequestTemplate } + : legacyCommand + ? { agentCommand: legacyCommand } + : undefined, ); - if (result.wasExisting) { + const launchRequest = launchRequestTemplate + ? { + ...launchRequestTemplate, + workspaceId: result.workspace.id, + } + : null; + + if (isOrchestratorEnabled) { + if (!launchRequest) { + return; + } + if (result.wasExisting) { + const launchResult = await launchAgentSession(launchRequest, { + source: "open-in-workspace", + createOrAttach: (input) => + terminalCreateOrAttach.mutateAsync(input), + write: (input) => terminalWrite.mutateAsync(input), + }); + if (launchResult.status === "failed") { + toast.error("Failed to start agent", { + description: + launchResult.error ?? "Failed to start agent session.", + }); + return; + } + } + } else if (legacyCommand && result.wasExisting) { const { tabId, paneId } = addTab(result.workspace.id); setTabAutoTitle(tabId, "Agent"); try { @@ -112,7 +214,7 @@ export function OpenInWorkspace({ task }: OpenInWorkspaceProps) { paneId, tabId, workspaceId: result.workspace.id, - command, + command: legacyCommand, createOrAttach: (input) => terminalCreateOrAttach.mutateAsync(input), write: (input) => terminalWrite.mutateAsync(input), @@ -215,7 +317,7 @@ export function OpenInWorkspace({ task }: OpenInWorkspaceProps) {