From 510b1489e4427e9a885060c505e461d800d29cce Mon Sep 17 00:00:00 2001
From: Alex Smolya <alexsmolya@gmail.com>
Date: Thu, 3 Sep 2026 08:59:11 +0200
Subject: [PATCH] fix: preserve tool budget on unknown tools

---
 src/actions/action-executor.ts |  3 +--
 tests/action-executor.test.ts  | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/src/actions/action-executor.ts b/src/actions/action-executor.ts
index 13cffc2..eaa78a8 100644
--- a/src/actions/action-executor.ts
+++ b/src/actions/action-executor.ts
@@ -44,9 +44,8 @@ export class ActionExecutor {
       assertMaxToolCalls(currentCount, this.maxToolCallsPerTask);
     }
 
-    this.toolCallCounts.set(context.taskId, currentCount + 1);
-
     const tool = this.toolRegistry.get(toolName);
+    this.toolCallCounts.set(context.taskId, currentCount + 1);
     const startedAt = Date.now();
 
     this.eventBus?.emit({
diff --git a/tests/action-executor.test.ts b/tests/action-executor.test.ts
index e5ff93c..82d66b2 100644
--- a/tests/action-executor.test.ts
+++ b/tests/action-executor.test.ts
@@ -44,6 +44,40 @@ describe("ActionExecutor", () => {
     expect(executor.getToolCallCount("task-1")).toBe(2);
   });
 
+  it("does not consume a call when the tool is not registered", async () => {
+    const registry = new ToolRegistry();
+    const executor = new ActionExecutor(registry);
+    const ctx = createMockContext("task-missing-tool");
+
+    await expect(executor.execute("missing", {}, ctx)).rejects.toMatchObject({
+      name: "RuntimeError",
+      code: "TOOL_NOT_FOUND",
+      details: { toolName: "missing" }
+    });
+
+    expect(executor.getToolCallCount(ctx.taskId)).toBe(0);
+  });
+
+  it("allows a valid call after an unknown tool with a one-call limit", async () => {
+    const registry = new ToolRegistry();
+    registry.register({
+      name: "ping",
+      description: "Ping tool",
+      execute() {
+        return "pong";
+      }
+    });
+
+    const executor = new ActionExecutor(registry, 1);
+    const ctx = createMockContext("task-retry-after-missing-tool");
+
+    await expect(executor.execute("missing", {}, ctx)).rejects.toMatchObject({
+      code: "TOOL_NOT_FOUND"
+    });
+    await expect(executor.execute("ping", {}, ctx)).resolves.toBe("pong");
+    expect(executor.getToolCallCount(ctx.taskId)).toBe(1);
+  });
+
   it("enforces maxToolCallsPerTask policy limit", async () => {
     const registry = new ToolRegistry();
     registry.register({