Add /link and /unlink commands for managing agent tools (#59)

Co-authored-by: Letta <noreply@letta.com>

src/agent/modify.ts

@@ -2,6 +2,7 @@
 // Utilities for modifying agent configuration
 
 import type { LlmConfig } from "@letta-ai/letta-client/resources/models/models";
+import { getToolNames } from "../tools/manager";
 import { getClient } from "./client";
 
 /**
@@ -44,3 +45,149 @@ export async function updateAgentLLMConfig(
 
   return finalConfig;
 }
+
+export interface LinkResult {
+  success: boolean;
+  message: string;
+  addedCount?: number;
+}
+
+export interface UnlinkResult {
+  success: boolean;
+  message: string;
+  removedCount?: number;
+}
+
+/**
+ * Attach all Letta Code tools to an agent.
+ *
+ * @param agentId - The agent ID
+ * @returns Result with success status and message
+ */
+export async function linkToolsToAgent(agentId: string): Promise<LinkResult> {
+  try {
+    const client = await getClient();
+
+    // Get ALL agent tools from agent state
+    const agent = await client.agents.retrieve(agentId);
+    const currentTools = agent.tools || [];
+    const currentToolIds = currentTools
+      .map((t) => t.id)
+      .filter((id): id is string => typeof id === "string");
+    const currentToolNames = new Set(
+      currentTools
+        .map((t) => t.name)
+        .filter((name): name is string => typeof name === "string"),
+    );
+
+    // Get Letta Code tool names
+    const lettaCodeToolNames = getToolNames();
+
+    // Find tools to add (tools that aren't already attached)
+    const toolsToAdd = lettaCodeToolNames.filter(
+      (name) => !currentToolNames.has(name),
+    );
+
+    if (toolsToAdd.length === 0) {
+      return {
+        success: true,
+        message: "All Letta Code tools already attached",
+        addedCount: 0,
+      };
+    }
+
+    // Look up tool IDs from global tool list
+    const toolsToAddIds: string[] = [];
+    for (const toolName of toolsToAdd) {
+      const tools = await client.tools.list({ name: toolName });
+      const tool = tools[0];
+      if (tool?.id) {
+        toolsToAddIds.push(tool.id);
+      }
+    }
+
+    // Combine current tools with new tools
+    const newToolIds = [...currentToolIds, ...toolsToAddIds];
+
+    // Get current tool_rules and add requires_approval rules for new tools
+    const currentToolRules = agent.tool_rules || [];
+    const newToolRules = [
+      ...currentToolRules,
+      ...toolsToAdd.map((toolName) => ({
+        tool_name: toolName,
+        type: "requires_approval" as const,
+        prompt_template: null,
+      })),
+    ];
+
+    await client.agents.modify(agentId, {
+      tool_ids: newToolIds,
+      tool_rules: newToolRules,
+    });
+
+    return {
+      success: true,
+      message: `Attached ${toolsToAddIds.length} Letta Code tool(s) to agent`,
+      addedCount: toolsToAddIds.length,
+    };
+  } catch (error) {
+    return {
+      success: false,
+      message: `Failed: ${error instanceof Error ? error.message : String(error)}`,
+    };
+  }
+}
+
+/**
+ * Remove all Letta Code tools from an agent.
+ *
+ * @param agentId - The agent ID
+ * @returns Result with success status and message
+ */
+export async function unlinkToolsFromAgent(
+  agentId: string,
+): Promise<UnlinkResult> {
+  try {
+    const client = await getClient();
+
+    // Get ALL agent tools from agent state (not tools.list which may be incomplete)
+    const agent = await client.agents.retrieve(agentId);
+    const allTools = agent.tools || [];
+    const lettaCodeToolNames = new Set(getToolNames());
+
+    // Filter out Letta Code tools, keep everything else
+    const remainingTools = allTools.filter(
+      (t) => t.name && !lettaCodeToolNames.has(t.name),
+    );
+    const removedCount = allTools.length - remainingTools.length;
+
+    // Extract IDs from remaining tools (filter out any undefined IDs)
+    const remainingToolIds = remainingTools
+      .map((t) => t.id)
+      .filter((id): id is string => typeof id === "string");
+
+    // Remove approval rules for Letta Code tools being unlinked
+    const currentToolRules = agent.tool_rules || [];
+    const remainingToolRules = currentToolRules.filter(
+      (rule) =>
+        rule.type !== "requires_approval" ||
+        !lettaCodeToolNames.has(rule.tool_name),
+    );
+
+    await client.agents.modify(agentId, {
+      tool_ids: remainingToolIds,
+      tool_rules: remainingToolRules,
+    });
+
+    return {
+      success: true,
+      message: `Removed ${removedCount} Letta Code tool(s) from agent`,
+      removedCount,
+    };
+  } catch (error) {
+    return {
+      success: false,
+      message: `Failed: ${error instanceof Error ? error.message : String(error)}`,
+    };
+  }
+}

src/cli/App.tsx

@@ -15,6 +15,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { getResumeData } from "../agent/check-approval";
 import { getClient } from "../agent/client";
 import { sendMessageStream } from "../agent/message";
+import { linkToolsToAgent, unlinkToolsFromAgent } from "../agent/modify";
 import { SessionStats } from "../agent/stats";
 import type { ApprovalContext } from "../permissions/analyzer";
 import { permissionMode } from "../permissions/mode";
@@ -111,6 +112,8 @@ export default function App({
   loadingState?:
     | "assembling"
     | "upserting"
+    | "linking"
+    | "unlinking"
     | "initializing"
     | "checking"
     | "ready";
@@ -837,6 +840,92 @@ export default function App({
           return { submitted: true };
         }
 
+        // Special handling for /link command - attach Letta Code tools
+        if (msg.trim() === "/link") {
+          const cmdId = uid("cmd");
+          buffersRef.current.byId.set(cmdId, {
+            kind: "command",
+            id: cmdId,
+            input: msg,
+            output: "Attaching Letta Code tools to agent...",
+            phase: "running",
+          });
+          buffersRef.current.order.push(cmdId);
+          refreshDerived();
+
+          setCommandRunning(true);
+
+          try {
+            const result = await linkToolsToAgent(agentId);
+
+            buffersRef.current.byId.set(cmdId, {
+              kind: "command",
+              id: cmdId,
+              input: msg,
+              output: result.message,
+              phase: "finished",
+              success: result.success,
+            });
+            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 };
+        }
+
+        // Special handling for /unlink command - remove Letta Code tools
+        if (msg.trim() === "/unlink") {
+          const cmdId = uid("cmd");
+          buffersRef.current.byId.set(cmdId, {
+            kind: "command",
+            id: cmdId,
+            input: msg,
+            output: "Removing Letta Code tools from agent...",
+            phase: "running",
+          });
+          buffersRef.current.order.push(cmdId);
+          refreshDerived();
+
+          setCommandRunning(true);
+
+          try {
+            const result = await unlinkToolsFromAgent(agentId);
+
+            buffersRef.current.byId.set(cmdId, {
+              kind: "command",
+              id: cmdId,
+              input: msg,
+              output: result.message,
+              phase: "finished",
+              success: result.success,
+            });
+            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, {

src/cli/commands/registry.ts

@@ -50,6 +50,20 @@ export const commands: Record<string, Command> = {
       return "Clearing credentials...";
     },
   },
+  "/link": {
+    desc: "Attach Letta Code tools to current agent",
+    handler: () => {
+      // Handled specially in App.tsx to access agent ID and client
+      return "Attaching tools...";
+    },
+  },
+  "/unlink": {
+    desc: "Remove Letta Code tools from current agent",
+    handler: () => {
+      // Handled specially in App.tsx to access agent ID and client
+      return "Removing tools...";
+    },
+  },
 };
 
 /**

src/cli/components/CommandPreview.tsx

@@ -34,9 +34,10 @@ export function CommandPreview({
       paddingX={1}
     >
       {commandList.map((item) => (
-        <Box key={item.cmd} justifyContent="space-between" width={40}>
-          <Text>{item.cmd}</Text>
-          <Text dimColor>{item.desc}</Text>
+        <Box key={item.cmd}>
+          <Text>
+            {item.cmd.padEnd(15)} <Text dimColor>{item.desc}</Text>
+          </Text>
         </Box>
       ))}
       {showAgentUrl && (

src/cli/components/WelcomeScreen.tsx

@@ -9,6 +9,8 @@ import { colors } from "./colors";
 type LoadingState =
   | "assembling"
   | "upserting"
+  | "linking"
+  | "unlinking"
   | "initializing"
   | "checking"
   | "ready";
@@ -82,6 +84,12 @@ export function WelcomeScreen({
     if (loadingState === "upserting") {
       return "Upserting tools...";
     }
+    if (loadingState === "linking") {
+      return "Attaching Letta Code tools...";
+    }
+    if (loadingState === "unlinking") {
+      return "Removing Letta Code tools...";
+    }
     if (loadingState === "checking") {
       return "Checking for pending approvals...";
     }

src/index.ts

@@ -91,6 +91,8 @@ async function main() {
         "permission-mode": { type: "string" },
         yolo: { type: "boolean" },
         "output-format": { type: "string" },
+        link: { type: "boolean" },
+        unlink: { type: "boolean" },
       },
       strict: true,
       allowPositionals: true,
@@ -233,6 +235,21 @@ async function main() {
     }
   }
 
+  // Handle --link and --unlink flags (modify tools before starting session)
+  const shouldLink = values.link as boolean | undefined;
+  const shouldUnlink = values.unlink as boolean | undefined;
+
+  // Validate --link/--unlink flags require --agent
+  if (shouldLink || shouldUnlink) {
+    if (!specifiedAgentId) {
+      console.error(
+        `Error: --${shouldLink ? "link" : "unlink"} requires --agent <id>`,
+      );
+      process.exit(1);
+    }
+    // Implementation is in InteractiveSession init()
+  }
+
   if (isHeadless) {
     // For headless mode, load tools synchronously
     await loadTools();
@@ -263,7 +280,13 @@ async function main() {
     model?: string;
   }) {
     const [loadingState, setLoadingState] = useState<
-      "assembling" | "upserting" | "initializing" | "checking" | "ready"
+      | "assembling"
+      | "upserting"
+      | "linking"
+      | "unlinking"
+      | "initializing"
+      | "checking"
+      | "ready"
     >("assembling");
     const [agentId, setAgentId] = useState<string | null>(null);
     const [agentState, setAgentState] = useState<AgentState | null>(null);
@@ -279,6 +302,28 @@ async function main() {
         const client = await getClient();
         await upsertToolsToServer(client);
 
+        // Handle --link/--unlink after upserting tools
+        if (shouldLink || shouldUnlink) {
+          if (!agentIdArg) {
+            console.error("Error: --link/--unlink requires --agent <id>");
+            process.exit(1);
+          }
+
+          setLoadingState(shouldLink ? "linking" : "unlinking");
+          const { linkToolsToAgent, unlinkToolsFromAgent } = await import(
+            "./agent/modify"
+          );
+
+          const result = shouldLink
+            ? await linkToolsToAgent(agentIdArg)
+            : await unlinkToolsFromAgent(agentIdArg);
+
+          if (!result.success) {
+            console.error(`✗ ${result.message}`);
+            process.exit(1);
+          }
+        }
+
         setLoadingState("initializing");
         const { createAgent } = await import("./agent/create");
         const { getModelUpdateArgs } = await import("./agent/model");

src/tests/agent/link-unlink.test.ts

@@ -0,0 +1,205 @@
+import { afterAll, beforeAll, describe, expect, test } from "bun:test";
+import { Letta } from "@letta-ai/letta-client";
+import { linkToolsToAgent, unlinkToolsFromAgent } from "../../agent/modify";
+import { settingsManager } from "../../settings-manager";
+import { getToolNames, loadTools } from "../../tools/manager";
+
+// Skip these integration tests if LETTA_API_KEY is not set
+const shouldSkip = !process.env.LETTA_API_KEY;
+const describeOrSkip = shouldSkip ? describe.skip : describe;
+
+describeOrSkip("Link/Unlink Tools", () => {
+  let client: Letta;
+  let testAgentId: string;
+
+  beforeAll(async () => {
+    // Initialize settings and load tools
+    await settingsManager.initialize();
+    await loadTools();
+
+    // Create a test agent
+    const apiKey = process.env.LETTA_API_KEY;
+    if (!apiKey) {
+      throw new Error("LETTA_API_KEY required for tests");
+    }
+
+    client = new Letta({ apiKey });
+
+    const agent = await client.agents.create({
+      model: "openai/gpt-4o-mini",
+      embedding: "openai/text-embedding-3-small",
+      memory_blocks: [
+        { label: "human", value: "Test user" },
+        { label: "persona", value: "Test agent" },
+      ],
+      tools: [],
+    });
+
+    testAgentId = agent.id;
+  });
+
+  afterAll(async () => {
+    // Cleanup: delete test agent
+    if (testAgentId) {
+      try {
+        await client.agents.delete(testAgentId);
+      } catch (_error) {
+        // Ignore cleanup errors
+      }
+    }
+  });
+
+  test("linkToolsToAgent attaches all Letta Code tools", async () => {
+    // Reset: ensure tools are not already attached
+    await unlinkToolsFromAgent(testAgentId);
+
+    const result = await linkToolsToAgent(testAgentId);
+
+    expect(result.success).toBe(true);
+    expect(result.addedCount).toBeGreaterThan(0);
+
+    // Verify tools were attached
+    const agent = await client.agents.retrieve(testAgentId);
+    const toolNames = agent.tools?.map((t) => t.name) || [];
+    const lettaCodeTools = getToolNames();
+
+    for (const toolName of lettaCodeTools) {
+      expect(toolNames).toContain(toolName);
+    }
+  }, 30000);
+
+  test("linkToolsToAgent adds approval rules for all tools", async () => {
+    // First unlink to reset
+    await unlinkToolsFromAgent(testAgentId);
+
+    // Link tools
+    await linkToolsToAgent(testAgentId);
+
+    // Verify approval rules were added
+    const agent = await client.agents.retrieve(testAgentId);
+    const approvalRules = agent.tool_rules?.filter(
+      (rule) => rule.type === "requires_approval",
+    );
+
+    const lettaCodeTools = getToolNames();
+    expect(approvalRules?.length).toBe(lettaCodeTools.length);
+
+    // Check all Letta Code tools have approval rules
+    const rulesToolNames = approvalRules?.map((r) => r.tool_name) || [];
+    for (const toolName of lettaCodeTools) {
+      expect(rulesToolNames).toContain(toolName);
+    }
+  }, 30000);
+
+  test("linkToolsToAgent returns success when tools already attached", async () => {
+    // Reset and link once
+    await unlinkToolsFromAgent(testAgentId);
+    await linkToolsToAgent(testAgentId);
+
+    // Link again
+    const result = await linkToolsToAgent(testAgentId);
+
+    expect(result.success).toBe(true);
+    expect(result.addedCount).toBe(0);
+    expect(result.message).toContain("already attached");
+  }, 30000);
+
+  test("unlinkToolsFromAgent removes all Letta Code tools", async () => {
+    // First link tools
+    await linkToolsToAgent(testAgentId);
+
+    // Then unlink
+    const result = await unlinkToolsFromAgent(testAgentId);
+
+    expect(result.success).toBe(true);
+    expect(result.removedCount).toBeGreaterThan(0);
+
+    // Verify tools were removed
+    const agent = await client.agents.retrieve(testAgentId);
+    const toolNames = agent.tools?.map((t) => t.name) || [];
+    const lettaCodeTools = getToolNames();
+
+    for (const toolName of lettaCodeTools) {
+      expect(toolNames).not.toContain(toolName);
+    }
+  }, 30000);
+
+  test("unlinkToolsFromAgent removes approval rules", async () => {
+    // First link tools
+    await linkToolsToAgent(testAgentId);
+
+    // Then unlink
+    await unlinkToolsFromAgent(testAgentId);
+
+    // Verify approval rules were removed
+    const agent = await client.agents.retrieve(testAgentId);
+    const approvalRules = agent.tool_rules?.filter(
+      (rule) => rule.type === "requires_approval",
+    );
+
+    const lettaCodeTools = new Set(getToolNames());
+    const remainingApprovalRules = approvalRules?.filter((r) =>
+      lettaCodeTools.has(r.tool_name),
+    );
+
+    expect(remainingApprovalRules?.length || 0).toBe(0);
+  }, 30000);
+
+  test("unlinkToolsFromAgent preserves non-Letta-Code tools", async () => {
+    // Link Letta Code tools
+    await linkToolsToAgent(testAgentId);
+
+    // Attach memory tool
+    const memoryTools = await client.tools.list({ name: "memory" });
+    const memoryTool = memoryTools[0];
+    if (memoryTool?.id) {
+      await client.agents.tools.attach(memoryTool.id, {
+        agent_id: testAgentId,
+      });
+    }
+
+    // Unlink Letta Code tools
+    await unlinkToolsFromAgent(testAgentId);
+
+    // Verify memory tool is still there
+    const agent = await client.agents.retrieve(testAgentId);
+    const toolNames = agent.tools?.map((t) => t.name) || [];
+
+    expect(toolNames).toContain("memory");
+
+    // Verify Letta Code tools are gone
+    const lettaCodeTools = getToolNames();
+    for (const toolName of lettaCodeTools) {
+      expect(toolNames).not.toContain(toolName);
+    }
+  }, 30000);
+
+  test("unlinkToolsFromAgent preserves non-approval tool_rules", async () => {
+    // Link tools
+    await linkToolsToAgent(testAgentId);
+
+    // Add a continue_loop rule manually
+    const agent = await client.agents.retrieve(testAgentId);
+    const newToolRules = [
+      ...(agent.tool_rules || []),
+      {
+        tool_name: "memory",
+        type: "continue_loop" as const,
+        prompt_template: "Test rule",
+      },
+    ];
+
+    await client.agents.modify(testAgentId, { tool_rules: newToolRules });
+
+    // Unlink Letta Code tools
+    await unlinkToolsFromAgent(testAgentId);
+
+    // Verify continue_loop rule is still there
+    const updatedAgent = await client.agents.retrieve(testAgentId);
+    const continueLoopRules = updatedAgent.tool_rules?.filter(
+      (r) => r.type === "continue_loop" && r.tool_name === "memory",
+    );
+
+    expect(continueLoopRules?.length).toBe(1);
+  }, 30000);
+});