Swap command (#84)

Co-authored-by: Shubham Naik <shub@memgpt.ai>
2025-11-17 11:26:48 -08:00
parent 9212dfb52b
commit d3b9eb7245
3 changed files with 423 additions and 2 deletions
--- a/src/cli/App.tsx
+++ b/src/cli/App.tsx
@@ -20,6 +20,7 @@ import { linkToolsToAgent, unlinkToolsFromAgent } from "../agent/modify";
 import { SessionStats } from "../agent/stats";
 import type { ApprovalContext } from "../permissions/analyzer";
 import { permissionMode } from "../permissions/mode";
+import { updateProjectSettings } from "../settings";
 import type { ToolExecutionResult } from "../tools/manager";
 import {
  analyzeToolApproval,
@@ -27,6 +28,7 @@ import {
  executeTool,
  savePermissionRule,
 } from "../tools/manager";
+import { AgentSelector } from "./components/AgentSelector";
 // import { ApprovalDialog } from "./components/ApprovalDialog";
 import { ApprovalDialog } from "./components/ApprovalDialogRich";
 // import { AssistantMessage } from "./components/AssistantMessage";
@@ -101,8 +103,8 @@ type StaticItem =
  | Line;

 export default function App({
-  agentId,
-  agentState,
+  agentId: initialAgentId,
+  agentState: initialAgentState,
  loadingState = "ready",
  continueSession = false,
  startupApproval = null,
@@ -126,6 +128,23 @@ export default function App({
  messageHistory?: LettaMessageUnion[];
  tokenStreaming?: boolean;
 }) {
+  // Track current agent (can change when swapping)
+  const [agentId, setAgentId] = useState(initialAgentId);
+  const [agentState, setAgentState] = useState(initialAgentState);
+
+  // Sync with prop changes (e.g., when parent updates from "loading" to actual ID)
+  useEffect(() => {
+    if (initialAgentId !== agentId) {
+      setAgentId(initialAgentId);
+    }
+  }, [initialAgentId, agentId]);
+
+  useEffect(() => {
+    if (initialAgentState !== agentState) {
+      setAgentState(initialAgentState);
+    }
+  }, [initialAgentState, agentState]);
+
  // Whether a stream is in flight (disables input)
  const [streaming, setStreaming] = useState(false);

@@ -178,6 +197,9 @@ export default function App({
  const [llmConfig, setLlmConfig] = useState<LlmConfig | null>(null);
  const [agentName, setAgentName] = useState<string | null>(null);

+  // Agent selector state
+  const [agentSelectorOpen, setAgentSelectorOpen] = useState(false);
+
  // Token streaming preference (can be toggled at runtime)
  const [tokenStreamingEnabled, setTokenStreamingEnabled] =
    useState(tokenStreaming);
@@ -1081,6 +1103,108 @@ export default function App({
          return { submitted: true };
        }

+        // Special handling for /swap command - switch to a different agent
+        if (msg.trim().startsWith("/swap")) {
+          const parts = msg.trim().split(/\s+/);
+          const targetAgentId = parts.slice(1).join(" ");
+
+          // If no agent ID provided, open agent selector
+          if (!targetAgentId) {
+            setAgentSelectorOpen(true);
+            return { submitted: true };
+          }
+
+          // Validate and swap to specified agent ID
+          const cmdId = uid("cmd");
+          buffersRef.current.byId.set(cmdId, {
+            kind: "command",
+            id: cmdId,
+            input: msg,
+            output: `Switching to agent ${targetAgentId}...`,
+            phase: "running",
+          });
+          buffersRef.current.order.push(cmdId);
+          refreshDerived();
+
+          setCommandRunning(true);
+
+          try {
+            const client = await getClient();
+            // Fetch new agent
+            const agent = await client.agents.retrieve(targetAgentId);
+
+            // Fetch agent's message history
+            const messagesPage =
+              await client.agents.messages.list(targetAgentId);
+            const messages = messagesPage.items;
+
+            // Update project settings with new agent
+            await updateProjectSettings({ lastAgent: targetAgentId });
+
+            // Clear current transcript
+            buffersRef.current.byId.clear();
+            buffersRef.current.order = [];
+            buffersRef.current.tokenCount = 0;
+            emittedIdsRef.current.clear();
+            setStaticItems([]);
+
+            // Update agent state
+            setAgentId(targetAgentId);
+            setAgentState(agent);
+            setAgentName(agent.name);
+            setLlmConfig(agent.llm_config);
+
+            // Add welcome screen for new agent
+            welcomeCommittedRef.current = false;
+            setStaticItems([
+              {
+                kind: "welcome",
+                id: `welcome-${Date.now().toString(36)}`,
+                snapshot: {
+                  continueSession: true,
+                  agentState: agent,
+                  terminalWidth: columns,
+                },
+              },
+            ]);
+
+            // Backfill message history
+            if (messages.length > 0) {
+              hasBackfilledRef.current = false;
+              backfillBuffers(buffersRef.current, messages);
+              refreshDerived();
+              commitEligibleLines(buffersRef.current);
+              hasBackfilledRef.current = true;
+            }
+
+            // Add success command to transcript
+            const successCmdId = uid("cmd");
+            buffersRef.current.byId.set(successCmdId, {
+              kind: "command",
+              id: successCmdId,
+              input: msg,
+              output: `✓ Switched to agent "${agent.name || targetAgentId}"`,
+              phase: "finished",
+              success: true,
+            });
+            buffersRef.current.order.push(successCmdId);
+            refreshDerived();
+          } catch (error) {
+            buffersRef.current.byId.set(cmdId, {
+              kind: "command",
+              id: cmdId,
+              input: msg,
+              output: `Failed: ${error instanceof Error ? error.message : String(error)}`,
+              phase: "finished",
+              success: false,
+            });
+            refreshDerived();
+          } finally {
+            setCommandRunning(false);
+          }
+          return { submitted: true };
+        }
+
        // Immediately add command to transcript with "running" phase
        const cmdId = uid("cmd");
        buffersRef.current.byId.set(cmdId, {
@@ -1220,6 +1344,8 @@ export default function App({
      refreshDerived,
      agentId,
      handleExit,
+      columns,
+      commitEligibleLines,
    ],
  );

@@ -1493,6 +1619,100 @@ export default function App({
    [agentId, refreshDerived],
  );

+  const handleAgentSelect = useCallback(
+    async (targetAgentId: string) => {
+      setAgentSelectorOpen(false);
+
+      const cmdId = uid("cmd");
+      buffersRef.current.byId.set(cmdId, {
+        kind: "command",
+        id: cmdId,
+        input: `/swap ${targetAgentId}`,
+        output: `Switching to agent ${targetAgentId}...`,
+        phase: "running",
+      });
+      buffersRef.current.order.push(cmdId);
+      refreshDerived();
+
+      setCommandRunning(true);
+
+      try {
+        const client = await getClient();
+        // Fetch new agent
+        const agent = await client.agents.retrieve(targetAgentId);
+
+        // Fetch agent's message history
+        const messagesPage = await client.agents.messages.list(targetAgentId);
+        const messages = messagesPage.items;
+
+        // Update project settings with new agent
+        await updateProjectSettings({ lastAgent: targetAgentId });
+
+        // Clear current transcript
+        buffersRef.current.byId.clear();
+        buffersRef.current.order = [];
+        buffersRef.current.tokenCount = 0;
+        emittedIdsRef.current.clear();
+        setStaticItems([]);
+
+        // Update agent state
+        setAgentId(targetAgentId);
+        setAgentState(agent);
+        setAgentName(agent.name);
+        setLlmConfig(agent.llm_config);
+
+        // Add welcome screen for new agent
+        welcomeCommittedRef.current = false;
+        setStaticItems([
+          {
+            kind: "welcome",
+            id: `welcome-${Date.now().toString(36)}`,
+            snapshot: {
+              continueSession: true,
+              agentState: agent,
+              terminalWidth: columns,
+            },
+          },
+        ]);
+
+        // Backfill message history
+        if (messages.length > 0) {
+          hasBackfilledRef.current = false;
+          backfillBuffers(buffersRef.current, messages);
+          refreshDerived();
+          commitEligibleLines(buffersRef.current);
+          hasBackfilledRef.current = true;
+        }
+
+        // Add success command to transcript
+        const successCmdId = uid("cmd");
+        buffersRef.current.byId.set(successCmdId, {
+          kind: "command",
+          id: successCmdId,
+          input: `/swap ${targetAgentId}`,
+          output: `✓ Switched to agent "${agent.name || targetAgentId}"`,
+          phase: "finished",
+          success: true,
+        });
+        buffersRef.current.order.push(successCmdId);
+        refreshDerived();
+      } catch (error) {
+        buffersRef.current.byId.set(cmdId, {
+          kind: "command",
+          id: cmdId,
+          input: `/swap ${targetAgentId}`,
+          output: `Failed: ${error instanceof Error ? error.message : String(error)}`,
+          phase: "finished",
+          success: false,
+        });
+        refreshDerived();
+      } finally {
+        setCommandRunning(false);
+      }
+    },
+    [refreshDerived, commitEligibleLines, columns],
+  );
+
  // Track permission mode changes for UI updates
  const [uiPermissionMode, setUiPermissionMode] = useState(
    permissionMode.getMode(),
@@ -1716,6 +1936,7 @@ export default function App({
                !showExitStats &&
                pendingApprovals.length === 0 &&
                !modelSelectorOpen &&
+                !agentSelectorOpen &&
                !planApprovalPending
              }
              streaming={streaming}
@@ -1741,6 +1962,15 @@ export default function App({
              />
            )}

+            {/* Agent Selector - conditionally mounted as overlay */}
+            {agentSelectorOpen && (
+              <AgentSelector
+                currentAgentId={agentId}
+                onSelect={handleAgentSelect}
+                onCancel={() => setAgentSelectorOpen(false)}
+              />
+            )}
+
            {/* Plan Mode Dialog - below live items */}
            {planApprovalPending && (
              <>
--- a/src/cli/commands/registry.ts
+++ b/src/cli/commands/registry.ts
@@ -71,6 +71,13 @@ export const commands: Record<string, Command> = {
      return "Renaming agent...";
    },
  },
+  "/swap": {
+    desc: "Switch to a different agent",
+    handler: () => {
+      // Handled specially in App.tsx to access agent list and client
+      return "Swapping agent...";
+    },
+  },
 };

 /**
--- a/src/cli/components/AgentSelector.tsx
+++ b/src/cli/components/AgentSelector.tsx
@@ -0,0 +1,184 @@
+import type { AgentState } from "@letta-ai/letta-client/resources/agents/agents";
+import { Box, Text, useInput } from "ink";
+import { useEffect, useState } from "react";
+import { getClient } from "../../agent/client";
+import { colors } from "./colors";
+
+interface AgentSelectorProps {
+  currentAgentId: string;
+  onSelect: (agentId: string) => void;
+  onCancel: () => void;
+}
+
+export function AgentSelector({
+  currentAgentId,
+  onSelect,
+  onCancel,
+}: AgentSelectorProps) {
+  const [agents, setAgents] = useState<AgentState[]>([]);
+  const [loading, setLoading] = useState(true);
+  const [error, setError] = useState<string | null>(null);
+  const [selectedIndex, setSelectedIndex] = useState(0);
+  const [searchQuery, setSearchQuery] = useState("");
+  const [debouncedQuery, setDebouncedQuery] = useState("");
+
+  useEffect(() => {
+    const fetchAgents = async () => {
+      try {
+        const client = await getClient();
+        const agentList = await client.agents.list();
+        setAgents(agentList.items);
+        setLoading(false);
+      } catch (err) {
+        setError(err instanceof Error ? err.message : String(err));
+        setLoading(false);
+      }
+    };
+    fetchAgents();
+  }, []);
+
+  // Debounce search query (300ms delay)
+  useEffect(() => {
+    const timer = setTimeout(() => {
+      setDebouncedQuery(searchQuery);
+    }, 300);
+
+    return () => clearTimeout(timer);
+  }, [searchQuery]);
+
+  // Filter agents based on debounced search query
+  const matchingAgents = agents.filter((agent) => {
+    if (!debouncedQuery) return true;
+    const query = debouncedQuery.toLowerCase();
+    const name = (agent.name || "").toLowerCase();
+    const id = (agent.id || "").toLowerCase();
+    return name.includes(query) || id.includes(query);
+  });
+
+  const filteredAgents = matchingAgents.slice(0, 10);
+
+  // Reset selected index when filtered list changes
+  useEffect(() => {
+    setSelectedIndex(0);
+  }, []);
+
+  useInput((input, key) => {
+    if (loading || error) return;
+
+    if (key.upArrow) {
+      setSelectedIndex((prev) => Math.max(0, prev - 1));
+    } else if (key.downArrow) {
+      setSelectedIndex((prev) => Math.min(filteredAgents.length - 1, prev + 1));
+    } else if (key.return) {
+      const selectedAgent = filteredAgents[selectedIndex];
+      if (selectedAgent?.id) {
+        onSelect(selectedAgent.id);
+      }
+    } else if (key.escape) {
+      onCancel();
+    } else if (key.backspace || key.delete) {
+      setSearchQuery((prev) => prev.slice(0, -1));
+    } else if (input && !key.ctrl && !key.meta) {
+      // Add regular characters to search query
+      setSearchQuery((prev) => prev + input);
+    }
+  });
+
+  if (loading) {
+    return (
+      <Box flexDirection="column">
+        <Text color={colors.selector.title}>Loading agents...</Text>
+      </Box>
+    );
+  }
+
+  if (error) {
+    return (
+      <Box flexDirection="column">
+        <Text color="red">Error loading agents: {error}</Text>
+        <Text dimColor>Press ESC to cancel</Text>
+      </Box>
+    );
+  }
+
+  if (agents.length === 0) {
+    return (
+      <Box flexDirection="column">
+        <Text color={colors.selector.title}>No agents found</Text>
+        <Text dimColor>Press ESC to cancel</Text>
+      </Box>
+    );
+  }
+
+  return (
+    <Box flexDirection="column" gap={1}>
+      <Box>
+        <Text bold color={colors.selector.title}>
+          Select Agent (↑↓ to navigate, Enter to select, ESC to cancel)
+        </Text>
+      </Box>
+
+      <Box>
+        <Text dimColor>Search: </Text>
+        <Text>{searchQuery || "_"}</Text>
+      </Box>
+
+      {filteredAgents.length === 0 && (
+        <Box>
+          <Text dimColor>No agents match your search</Text>
+        </Box>
+      )}
+
+      {filteredAgents.length > 0 && (
+        <Box>
+          <Text dimColor>
+            Showing {filteredAgents.length}
+            {matchingAgents.length > 10 ? ` of ${matchingAgents.length}` : ""}
+            {debouncedQuery ? " matching" : ""} agents
+          </Text>
+        </Box>
+      )}
+
+      <Box flexDirection="column">
+        {filteredAgents.map((agent, index) => {
+          const isSelected = index === selectedIndex;
+          const isCurrent = agent.id === currentAgentId;
+
+          const lastInteractedAt = agent.last_run_completion
+            ? new Date(agent.last_run_completion).toLocaleString()
+            : "Never";
+
+          return (
+            <Box key={agent.id} flexDirection="row" gap={1}>
+              <Text
+                color={isSelected ? colors.selector.itemHighlighted : undefined}
+              >
+                {isSelected ? "›" : " "}
+              </Text>
+              <Box flexDirection="row" gap={2}>
+                <Text
+                  bold={isSelected}
+                  color={
+                    isSelected ? colors.selector.itemHighlighted : undefined
+                  }
+                  wrap="truncate-end"
+                >
+                  {agent.name || "Unnamed"}
+                  {isCurrent && (
+                    <Text color={colors.selector.itemCurrent}> (current)</Text>
+                  )}
+                </Text>
+                <Text dimColor wrap="truncate-end">
+                  {agent.id}
+                </Text>
+                <Text dimColor wrap="truncate-end">
+                  {lastInteractedAt}
+                </Text>
+              </Box>
+            </Box>
+          );
+        })}
+      </Box>
+    </Box>
+  );
+}