From 6c7a2efdbb2700763cac28ecfbf1f0135294b47c Mon Sep 17 00:00:00 2001 From: Charles Packer Date: Mon, 16 Mar 2026 12:53:17 -0700 Subject: [PATCH] fix(startup): clear Letta Code tool rules on boot (#1405) Co-authored-by: Jin Peng Co-authored-by: Letta Code --- src/headless.ts | 28 +++++- src/index.ts | 28 +++++- .../toolset-client-tool-rule-cleanup.test.ts | 85 +++++++++++++++++++ src/tests/tools/toolset-memfs-detach.test.ts | 1 + src/tools/toolset.ts | 48 +++++++++++ 5 files changed, 188 insertions(+), 2 deletions(-) create mode 100644 src/tests/tools/toolset-client-tool-rule-cleanup.test.ts diff --git a/src/headless.ts b/src/headless.ts index 5d1eb4c..d1a7bf2 100644 --- a/src/headless.ts +++ b/src/headless.ts @@ -101,6 +101,7 @@ import { registerExternalTools, setExternalToolExecutor, } from "./tools/manager"; +import { clearPersistedClientToolRules } from "./tools/toolset"; import type { AutoApprovalMessage, BootstrapSessionStateRequest, @@ -119,7 +120,7 @@ import type { StreamEvent, SystemInitMessage, } from "./types/protocol"; -import { debugLog } from "./utils/debug"; +import { debugLog, debugWarn } from "./utils/debug"; import { markMilestone, measureSinceMilestone, @@ -1042,6 +1043,31 @@ export async function handleHeadlessCommand( } } + const startupAgentId = agent.id; + void clearPersistedClientToolRules(startupAgentId) + .then((cleanup) => { + if (cleanup) { + const count = cleanup.removedToolNames.length; + const names = cleanup.removedToolNames.join(", "); + debugLog( + "headless startup", + `Cleared ${count} persisted client tool rule${count === 1 ? "" : "s"} for ${startupAgentId}${count > 0 ? `: ${names}` : ""}`, + ); + return; + } + + debugLog( + "headless startup", + `No persisted client tool rules to clear for ${startupAgentId}`, + ); + }) + .catch((error) => { + debugWarn( + "headless startup", + `Failed to clear persisted client tool rules for ${startupAgentId}: ${error instanceof Error ? error.message : String(error)}`, + ); + }); + try { effectiveReflectionSettings = await applyReflectionOverrides( agent.id, diff --git a/src/index.ts b/src/index.ts index 95d9db6..10283a9 100755 --- a/src/index.ts +++ b/src/index.ts @@ -48,7 +48,8 @@ import { settingsManager } from "./settings-manager"; import { startStartupAutoUpdateCheck } from "./startup-auto-update"; import { telemetry } from "./telemetry"; import { loadTools } from "./tools/manager"; -import { debugLog } from "./utils/debug"; +import { clearPersistedClientToolRules } from "./tools/toolset"; +import { debugLog, debugWarn } from "./utils/debug"; import { markMilestone } from "./utils/timing"; // Stable empty array constants to prevent new references on every render @@ -1793,6 +1794,31 @@ async function main(): Promise { } } + const startupAgentId = agent.id; + void clearPersistedClientToolRules(startupAgentId) + .then((cleanup) => { + if (cleanup) { + const count = cleanup.removedToolNames.length; + const names = cleanup.removedToolNames.join(", "); + debugLog( + "startup", + `Cleared ${count} persisted client tool rule${count === 1 ? "" : "s"} for ${startupAgentId}${count > 0 ? `: ${names}` : ""}`, + ); + return; + } + + debugLog( + "startup", + `No persisted client tool rules to clear for ${startupAgentId}`, + ); + }) + .catch((error) => { + debugWarn( + "startup", + `Failed to clear persisted client tool rules for ${startupAgentId}: ${error instanceof Error ? error.message : String(error)}`, + ); + }); + // Handle conversation: either resume existing or create new // Using definite assignment assertion - all branches below either set this or exit/throw let conversationIdToUse!: string; diff --git a/src/tests/tools/toolset-client-tool-rule-cleanup.test.ts b/src/tests/tools/toolset-client-tool-rule-cleanup.test.ts new file mode 100644 index 0000000..54ef1fd --- /dev/null +++ b/src/tests/tools/toolset-client-tool-rule-cleanup.test.ts @@ -0,0 +1,85 @@ +import { beforeEach, describe, expect, mock, test } from "bun:test"; + +const retrieveMock = mock((_agentId: string, _opts?: Record) => + Promise.resolve({ + tags: ["origin:letta-code"], + tool_rules: [ + { type: "requires_approval", tool_name: "Bash" }, + { type: "requires_approval", tool_name: "Task" }, + { type: "requires_approval", tool_name: "web_search" }, + { type: "continue", tool_name: "fetch_webpage" }, + ], + }), +); +const updateMock = mock((_agentId: string, _payload: Record) => + Promise.resolve({}), +); +const mockGetClient = mock(() => + Promise.resolve({ + agents: { + retrieve: retrieveMock, + update: updateMock, + }, + }), +); + +mock.module("../../agent/client", () => ({ + getClient: mockGetClient, + getServerUrl: () => "http://localhost:8283", +})); + +const { clearPersistedClientToolRules, shouldClearPersistedToolRules } = + await import("../../tools/toolset"); + +describe("client tool rule cleanup", () => { + beforeEach(() => { + retrieveMock.mockClear(); + updateMock.mockClear(); + mockGetClient.mockClear(); + }); + + test("marks Letta Code agents with any persisted tool rules for cleanup", () => { + expect( + shouldClearPersistedToolRules({ + tags: ["origin:letta-code"], + tool_rules: [{ type: "requires_approval", tool_name: "Bash" }], + }), + ).toBe(true); + }); + + test("clears all tool rules for Letta Code agents on startup", async () => { + const result = await clearPersistedClientToolRules("agent-123"); + + expect(result).toEqual({ + removedToolNames: ["Bash", "Task", "web_search", "fetch_webpage"], + }); + expect(updateMock).toHaveBeenCalledTimes(1); + expect(updateMock.mock.calls[0]?.[1]).toEqual({ + tool_rules: [], + }); + }); + + test("skips update when there are no persisted tool rules", async () => { + retrieveMock.mockResolvedValueOnce({ + tags: ["origin:letta-code"], + tool_rules: [], + }); + + const result = await clearPersistedClientToolRules("agent-123"); + + expect(result).toBeNull(); + expect(updateMock).not.toHaveBeenCalled(); + }); + + test("skips update for non-Letta Code agents", async () => { + retrieveMock.mockResolvedValueOnce({ + tags: ["some-other-tag"], + tool_rules: [{ type: "requires_approval", tool_name: "web_search" }], + }); + + const result = await clearPersistedClientToolRules("agent-123"); + + expect(result).toBeNull(); + expect(updateMock).not.toHaveBeenCalled(); + }); +}); diff --git a/src/tests/tools/toolset-memfs-detach.test.ts b/src/tests/tools/toolset-memfs-detach.test.ts index fc13427..8db5ab8 100644 --- a/src/tests/tools/toolset-memfs-detach.test.ts +++ b/src/tests/tools/toolset-memfs-detach.test.ts @@ -35,6 +35,7 @@ const mockGetClient = mock(() => mock.module("../../agent/client", () => ({ getClient: mockGetClient, + getServerUrl: () => "http://localhost:8283", })); const { detachMemoryTools } = await import("../../tools/toolset"); diff --git a/src/tools/toolset.ts b/src/tools/toolset.ts index d797604..c763c55 100644 --- a/src/tools/toolset.ts +++ b/src/tools/toolset.ts @@ -1,3 +1,4 @@ +import type { AgentState } from "@letta-ai/letta-client/resources/agents/agents"; import { getClient } from "../agent/client"; import { resolveModel } from "../agent/model"; import { toolFilter } from "./filter"; @@ -219,6 +220,53 @@ export async function reattachMemoryTool( } } +type PersistedToolRule = NonNullable[number]; + +interface AgentWithToolsAndRules { + tags?: string[] | null; + tool_rules?: PersistedToolRule[]; +} + +export function shouldClearPersistedToolRules( + agent: AgentWithToolsAndRules, +): boolean { + return ( + agent.tags?.includes("origin:letta-code") === true && + (agent.tool_rules?.length ?? 0) > 0 + ); +} + +export async function clearPersistedClientToolRules( + agentId: string, +): Promise<{ removedToolNames: string[] } | null> { + const client = await getClient(); + + try { + const agentWithTools = (await client.agents.retrieve(agentId, { + include: ["agent.tools"], + })) as AgentWithToolsAndRules; + if (!shouldClearPersistedToolRules(agentWithTools)) { + return null; + } + const existingRules = agentWithTools.tool_rules || []; + + await client.agents.update(agentId, { + tool_rules: [], + }); + + return { + removedToolNames: existingRules + .map((rule) => rule.tool_name) + .filter((name): name is string => typeof name === "string"), + }; + } catch (err) { + console.warn( + `Warning: Failed to clear persisted client tool rules: ${err instanceof Error ? err.message : String(err)}`, + ); + return null; + } +} + /** * Force switch to a specific toolset regardless of model. *