DEV Community

Cover image for My LLM kept calling tools it shouldn't, so I built a state machine to stop it
roddcode
roddcode

Posted on

My LLM kept calling tools it shouldn't, so I built a state machine to stop it

A patient confirmed a booking for a slot that was already taken. The LLM said "confirmed" while the database said "sold." Not because the model was malicious — because it had access to confirm_booking during the identity verification step, and the user's message happened to contain the right trigger words.

The model cannot be the source of truth.

Why system prompts aren't enough

Every production AI agent I've built hits the same wall: the LLM calls tools it shouldn't have access to. The standard fix is a better system prompt. "Never confirm a booking before identity is verified. Never share internal data. Never skip a step."

Prompts leak. The model finds a path around them, especially under pressure or ambiguity. Add a second language, an unusual input format, or a user who sounds authoritative, and the guardrail dissolves.

The problem isn't the prompt — it's that the tool exists in the LLM's context at all. If confirm_booking is in the tools array during identity verification, the model can call it. Period. No amount of prompt engineering changes that.

This is why I built reactive-fsm: a zero-dependency TypeScript library that controls tool access per conversation state. The model doesn't get told "don't call this tool." The tool simply doesn't exist in its context until the state machine says it does.

The architecture

The core idea is declarative. You define states, which tools belong to each state, and how transitions happen. The LLM operates inside the current state and cannot escape it.

import { createFSM } from "reactive-fsm";

const fsm = createFSM({
  initialState: "TRIAGE",
  states: ["TRIAGE", "BOOKING", "DONE"],
  tools: {
    TRIAGE: ["check_service", "check_availability"],
    BOOKING: ["confirm", "invoice"],
    DONE: [],
  },
});

fsm.allowedTools; // ['check_service', 'check_availability']
fsm.transitionTo("BOOKING");
fsm.allowedTools; // ['confirm', 'invoice']
Enter fullscreen mode Exit fullscreen mode

In TRIAGE, confirm and invoice don't exist. The adapter passes only the allowed tools to the LLM. When a tool execution validates the user's identity and the system calls transitionTo("BOOKING"), those tools appear. check_service disappears. The guardrail is structural, not textual.

The adapters handle provider-specific formatting. Five are available: Vercel AI SDK, OpenAI, Anthropic, LangChain, and Google Gemini. Each one wraps the tools for the current state in the provider's format. The core has zero npm dependencies — only TypeScript types.

Loop shield: when the LLM won't stop

The most expensive bug in production AI systems isn't a wrong tool call — it's an infinite loop. The LLM calls a tool, gets a result it doesn't like, calls it again with slightly different parameters, and repeats. Each iteration costs tokens and latency. Without a circuit breaker, it runs until the request times out.

The loop shield detects consecutive tool calls and forces a stop. It has two modes: consecutive (any N tools in a row) and repeated (the same tool called N times).

const fsm = createFSM({
  initialState: "DIAGNOSIS",
  states: ["DIAGNOSIS", "ESCALATION"],
  tools: {
    DIAGNOSIS: ["search_kb", "check_status"],
    ESCALATION: ["transfer_to_human"],
  },
  loopShield: {
    enabled: true,
    maxConsecutiveTools: 3,
    mode: "repeated",
    fallbackState: "ESCALATION",
    onLoop: ({ consecutiveTools, maxAllowed }) => {
      console.warn(`Loop detected: ${consecutiveTools} tools in a row (max ${maxAllowed})`);
    },
  },
});
Enter fullscreen mode Exit fullscreen mode

When the shield triggers, the FSM auto-transitions to ESCALATION — bypassing any guards, since this is an emergency exit. The onLoop callback fires with metadata so you can hook in logging or alerts. The user never sees a hung agent.

Guards that can check the database

Transitions often depend on external state. You want to move to CONFIRMED only if the slot is still available. A synchronous guard function isn't enough — you need an async check against the database.

const fsm = createFSM({
  initialState: "BOOKING",
  states: ["BOOKING", "CONFIRMED", "UNAVAILABLE"],
  tools: {
    BOOKING: ["reserve_slot"],
    CONFIRMED: ["send_confirmation"],
    UNAVAILABLE: [],
  },
  context: { slotId: "slot-42" },
  guard: async (from, to, ctx) => {
    if (to === "CONFIRMED") {
      const available = await db.slots.isAvailable(ctx.slotId);
      return available;
    }
    return true;
  },
});

await fsm.transitionToAsync("CONFIRMED");
Enter fullscreen mode Exit fullscreen mode

The guard receives the shared context object — a mutable scratchpad for session state. Tools write to it during execution. The guard reads from it during transitions. If the slot was taken between the LLM's decision and the confirmation attempt, the guard blocks the transition. Use transitionTo() for sync guards, transitionToAsync() for async ones.

Serverless-friendly by default

Every request to a production agent in 2026 hits a different Lambda instance. The FSM needs to survive cold starts without an external state store.

// Request 1: user starts a booking
const fsm = createFSM(bookingConfig);
fsm.transitionTo("SCHEDULING");
await db.save({ sessionId, snapshot: fsm.snapshot() });
// → { state: 'SCHEDULING' }

// Request 2 (different instance, hours later)
const saved = await db.load(sessionId);
const fsm = createFSM({ ...bookingConfig, snapshot: saved });
fsm.currentState; // 'SCHEDULING'
Enter fullscreen mode Exit fullscreen mode

No Redis. No sticky sessions. snapshot() returns { state: string }. Store it anywhere. The next invocation picks up where the last one left off.

The source code is the documentation

The entire core is three files: machine.ts, tool-gating.ts, loop-shield.ts. The adapters share a common base. There are 96 tests covering edge cases I debugged at 2 AM in production — state validation on snapshot restore, Valibot output compatibility, consistent loop shield reset across all five adapters.

The library is MIT. Install it, read the README, open an issue if something breaks.

pnpm add reactive-fsm
Enter fullscreen mode Exit fullscreen mode

I built it because I was tired of fixing the same bug in four different projects. If your LLM is calling tools it shouldn't, you don't need a better prompt. You need the tool to not exist in its context.

Top comments (0)