Build an Agent Task State Machine Before You Build an Agent

“The agent is working” sounds like a status, but it hides several different situations.

The task may be waiting for a worker, running a command, blocked on a question, finished successfully, failed, or canceled. If all of those cases become one Boolean called is_running, the first retry or late network response will make the program confusing.

A small state machine is a useful learning project because it teaches an important AI application lesson without requiring an API key or a model. We will define the legal task states, reject impossible transitions, preserve a tiny event history, and test the tricky paths.

What is a state machine?

A state machine has:

• a finite list of states;


• a list of transitions that are allowed from each state;


• an event that explains why a transition happened.

Here is the model for this project:

queued -----> running -----> succeeded

   |             |  \

   |             |   +----> failed

   |             |

   |             +-------> needs_input --+

   |                                      |

   +------------> canceled <--------------+

succeeded, failed, and canceled are terminal states. Once a task reaches one of them, a late worker message must not silently restart it.

Implement the states

Create task_state.py:

from dataclasses import dataclass, field

from enum import Enum





class State(str, Enum):

    QUEUED = "queued"

    RUNNING = "running"

    NEEDS_INPUT = "needs_input"

    SUCCEEDED = "succeeded"

    FAILED = "failed"

    CANCELED = "canceled"





ALLOWED = {

    State.QUEUED: {State.RUNNING, State.CANCELED},

    State.RUNNING: {

        State.NEEDS_INPUT,

        State.SUCCEEDED,

        State.FAILED,

        State.CANCELED,

    },

    State.NEEDS_INPUT: {State.RUNNING, State.CANCELED},

    State.SUCCEEDED: set(),

    State.FAILED: set(),

    State.CANCELED: set(),

}





@dataclass

class Task:

    task_id: str

    state: State = State.QUEUED

    events: list[dict[str, str]] = field(default_factory=list)



    def transition(self, next_state: State, reason: str) -> None:

        if next_state not in ALLOWED[self.state]:

            raise ValueError(f"invalid transition: {self.state} -> {next_state}")

        self.events.append(

            {"from": self.state.value, "to": next_state.value, "reason": reason}

        )

        self.state = next_state

The transition table is data rather than a long chain of if statements. That makes it easy to review every legal path in one place.

Add a small example at the bottom:

if __name__ == "__main__":

    task = Task("demo-1")

    task.transition(State.RUNNING, "worker claimed task")

    task.transition(State.NEEDS_INPUT, "missing target branch")

    task.transition(State.RUNNING, "user selected main")

    task.transition(State.SUCCEEDED, "checks passed")

    print(task.state.value)

    for event in task.events:

        print(event)

Run it with Python 3.9 or newer:

python3 task_state.py

Expected final state:

succeeded

The four printed events should show the complete path, including the pause for user input.

Test the paths that usually break

Save this as test_task_state.py:

import unittest



from task_state import State, Task





class TaskStateTests(unittest.TestCase):

    def test_happy_path(self):

        task = Task("t-1")

        task.transition(State.RUNNING, "claimed")

        task.transition(State.SUCCEEDED, "verified")

        self.assertEqual(task.state, State.SUCCEEDED)

        self.assertEqual(len(task.events), 2)



    def test_terminal_state_cannot_restart(self):

        task = Task("t-2")

        task.transition(State.CANCELED, "user request")

        with self.assertRaisesRegex(ValueError, "invalid transition"):

            task.transition(State.RUNNING, "late worker message")



    def test_needs_input_can_resume(self):

        task = Task("t-3")

        task.transition(State.RUNNING, "claimed")

        task.transition(State.NEEDS_INPUT, "missing branch")

        task.transition(State.RUNNING, "answer received")

        self.assertEqual(task.state, State.RUNNING)





if __name__ == "__main__":

    unittest.main()

Run:

python3 -m unittest -v

I verified the three tests with Python 3.9.13. The important test is not the happy path. It is test_terminal_state_cannot_restart, because distributed programs often receive delayed or duplicate events.

What this model still lacks

This is an in-memory learning model. A real task system also needs timestamps, durable storage, optimistic concurrency or version numbers, authorization, retry policy, and a definition of who may perform each transition.

Try these extensions in order:

1. Add an integer version that increments on every transition.


2. Require an actor field in each event.


3. Add cancel_requested as a non-terminal state and decide what the worker may do next.


4. Serialize the task to JSON and load it again.


5. Reject an update whose expected version is stale.

Each extension teaches a different idea: observability, authorization, asynchronous cancellation, persistence, and concurrency control.

A real-world connection

The MonkeyCode repository describes AI task management as part of an open-source development platform. That makes it a useful source to review after finishing this exercise: look for how a production project names task states, represents user input, and separates task management from the development environment. I reviewed the public documentation; I did not run MonkeyCode for this Python example.

Disclosure: I contribute to the MonkeyCode project. The state-machine code and tests above are an independent teaching example, not a description of MonkeyCode internals.

Learners who want to discuss the project can join the MonkeyCode Discord. The team can also explain whether free model credits are currently available and what eligibility or usage limits apply.

The main lesson is small but powerful: progress is not a Boolean. Once a task can pause, retry, cancel, or receive a late message, its legal states belong in code and tests.