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

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

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);
+});