Harnessing agents in Erlang - Part 2

This is part two of my harness in erlang. I’ll be turning what was a fairly rough harness with async tools and a client into something a bit more robust.

Queueing messages

In part one we built a very basic agent harness in Erlang. Our agent has a way of queuing messages you send while it’s still working, but our client had no ability to take them back.

My first thought was to use a simple slash command like /unqueue. But that alone won’t work. To understand why needs a little understanding of how the terminal works.

The terminal used to be a physical device connected with a cable to a computer. As such it has its own wire format, sending bytes that correspond to keypresses. Getting every program to reimplement the logic of turning these bytes into text would be wasteful, so most of the time the terminal is in what’s called cooked mode, where the OS handles the line editing (backspace, cursor movement, ctrl+c, etc) and feeds completed lines to the program. But cooked mode doesn’t give us any way to put a queued message back into the line you’re editing.

There’s an alternative, raw mode, and unfortunately we have to do it raw. Which means we have to handle the keystrokes ourselves. Erlang provides an easy way to switch to raw mode.

% harness.erl
main(_Args) ->
  ok = shell:start_interactive({noshell, raw}),
  io:setopts([binary]),
  try
    %% chatting loop
  after
    shell:start_interactive({noshell, cooked})
  end.

We’ve moved the previous /bin/ask escript into its own harness.erl file and swapping io:read_line for io:get_chars("> ", 1) to take one char at a time, let’s run it.

[michael@archlinux] ~/projects/erlang/harness ❯ ./bin/ask
> N what, guv’nor? Could you give me a bit more to go on?
                                                         > A, then! What’s the rest, guv’nor?
                                                                                             > E?

This doesn’t look right at all. Every char goes to the agent on its own. What’s more the output is staggered, because in raw mode a newline \n doesn’t move the cursor back to the left. That needs a carriage return \r too.

But there’s an even bigger problem. After sending a message the chat loop calls receive_messages, which blocks until the agent is done. We can’t process anything the user types while a message is in progress.

Luckily, Erlang gives us a very easy way to unblock this, processes. Taking some inspiration from the Elm architecture we will model our program in 3 processes:

read_loop -> model -> view

• read_loop: reads raw terminal input and translates it to structured messages
• model: owns the UI state, e.g. the text input
• view: draws whatever state the model sends it

% harness.erl
%% main(_Args) ->
  case application:ensure_all_started(harness) of
    {ok, _Started} ->
      process_flag(trap_exit, true),
      {ok, ViewPid} = create_view(),
      {ok, ModelPid} = create_model(ViewPid),
      Reader = spawn_link(fun() -> read_loop(ModelPid) end),
      Reason = wait(Reader),
      shell:start_interactive({noshell, cooked}),
      report(Reason);
%%% end

wait(Reader) ->
  receive
    {'EXIT', _Pid, Reason} ->
      exit(Reader, kill),
      Reason
  end.

All three are started with spawn_link, and main traps exits, so if any of them dies main gets an 'EXIT', puts the terminal back into cooked mode and quits.

The reader

% harness.erl
read_loop(ModelPid) ->
  case io:get_chars("", 1) of
    eof ->
      ok;
    {error, Reason} ->
      exit({io_error, Reason});
    <<27>> ->
      % TODO: fix. BAD - should parse until we get a terminating byte
      Seq = io:get_chars("", 2),
      ModelPid ! map_input(<<27, Seq/binary>>),
      read_loop(ModelPid);
    Chars ->
      ModelPid ! map_input(Chars),
      read_loop(ModelPid)
  end.

It reads a character, maps it to something like {key, enter} or {text, Chars} and sends it to the model. Keys like the arrows send escape sequences that start with ESC (<<27>>), so in that case we read 2 more bytes. That’s not really right, escape sequences can be much longer than this, but it’ll do for now.

The model

The model is where the state lives. It’s split into three functions:

• update(Msg, State) is a pure function returning the new state and a list of effects, so we can test it
• run_effect(Effect) does the side effects, e.g. sending a message to the agent
• loop ties them together

% harness.erl
update({text, Chars}, State) ->
  Text = maps:get(text, State, <<"">>),
  {State#{text => <<Text/binary, Chars/binary>>}, []};

loop(ViewPid, State) ->
  receive
    Msg ->
      {NewState, Effects} = update(Msg, State),
      lists:foreach(fun run_effect/1, Effects),
      ViewPid ! {draw, NewState},
      loop(ViewPid, NewState)
  end.

For now the model builds its own copy of the conversation from events the agent sends it, like this:

% harness.erl
update({msg, delta, Thing}, State) ->
  Inpt = maps:get(in_progress_text, State, <<"">>),
  {State#{in_progress_text => <<Inpt/binary, Thing/binary>>}, []};
update({msg, done}, State) ->
  Inpt = maps:get(in_progress_text, State, <<"">>),
  History = maps:get(history, State, []),
  {State#{in_progress => false,
          in_progress_text => <<"">>,
          history => [{assistant, Inpt}] ++ History},
   []};

plus {msg, tool, Name} and {msg, error, Reason}. Remember that, it comes back to bite.

The view

Each draw call carries the whole state, so the view can skip any it’s fallen behind on:

% harness.erl
view_loop() ->
  receive
    {draw, State} ->
      NewState = latest_draw(State),
      io:put_chars(render(NewState)),
      view_loop()
  end.

latest_draw(State) ->
  receive
    {draw, NewerState} ->
      latest_draw(NewerState)
  after 0 ->
    State
  end.

The after 0 is what coalesces the draws. In Erlang receive only takes messages that match, and after 0 means don’t wait if there aren’t any. So:

mailbox before latest_draw:
[{draw, State}, {draw, State}, A, B, {draw, State}, C, {draw, State}]

mailbox after latest_draw:
[A, B, C]

render is a pure function from state to the bytes to print. It clears the screen and redraws everything, turning every \n into \r\n to fix the staggering. Clearing the whole screen on every draw will flicker, but this is a harness post not a TUI one, so I’m leaving it.

Unqueue and slash commands

You might have forgotten by now, but originally I was trying to implement unqueue. Now we finally can. Enter checks for a leading / and hands it off to command, and the up arrow is a shortcut for /unqueue.

% harness.erl
update({key, up}, State) ->
  command(<<"unqueue">>, State);

command(<<"quit">>, State) ->
  {State#{text => <<"">>}, [quit]};
command(<<"unqueue">>, State) ->
  {State#{text => <<"">>}, [unqueue]};
command(Cmd, State) when is_binary(Cmd) ->
  {State#{last_error => {no_command, Cmd}}, []}.

command doesn’t do anything itself, it just returns effects. The unqueue effect takes the agent’s pending messages:

% harness.erl
run_effect(unqueue) ->
  Self = self(),
  spawn(fun() ->
           {lease, Pending} = gen_server:call(agent, {user, lease_pending}),
           case Pending of
             [] -> ok;
             Other -> Self ! {pending_leased, Other}
           end
        end);

The gen_server:call happens in a throwaway process so the model never blocks waiting on the agent. The reply comes back as just another message, and update puts the messages back into the input line.

Cleaning up our state

Our agent state has turned into a pile of flags. api_call, response_complete, request_pid, reply_to and tool_calls all sit at the top level and every handler has to work out what’s going on by looking at a combination of them. Really the agent is only ever in one of three states: idle, streaming a response, or waiting on tools to finish. So let’s write down what every message should do in each of them.

The state now has a single turn key. It’s either idle or a map holding everything that only lives for the length of the turn, with a phase of streaming or awaiting_tools. Every handler matches on that, and the code falls straight out of the table. For example user messages always go into pending, and when we’re idle we also kick off a request:

% agent.erl
handle_cast({user, message, From, Msg}, #{turn := idle, pending := Pending} = State) ->
  NewState = State#{pending => [{user, Msg} | Pending]},
  {noreply, call_llm(From, NewState)};
handle_cast({user, message, _From, Msg}, #{pending := Pending} = State) ->
  NewState = State#{pending => [{user, Msg} | Pending]},
  {noreply, NewState};

That leaves the bit of the table I’d been putting off. Tools that crash or never finish.

Adding some robustness

Right now we spawn a process for each tool and forget about it. If it crashes, the turn waits forever.

We could put the tools under a supervisor, but that doesn’t really buy us a lot here. The tool’s lifetime belongs to the turn. The agent needs its result, has to kill it if the turn fails, and has to clean up after it. So instead the agent monitors each tool, and the tool hands back its result as its exit reason.

% agent.erl
dispatch_tool(Name, CallId, Args, #{turn := #{tool_calls := ToolCalls} = Turn} = State)->
  Agent = self(),
  OnCleanup = fun(Fun) -> Agent ! {tool_cleanup, self(), Fun} end,
  {Pid, Ref} =
    spawn_monitor(fun() ->
                     exit(try
                            call_tool(Name, Args, OnCleanup)
                          catch
                            Class:Reason -> {tool_crash, Class, Reason}
                          end)
                  end),
  erlang:send_after(?TOOL_TIMEOUT, self(), {tool_timeout, Pid}),
  State#{turn =>
           Turn#{tool_calls =>
                   ToolCalls#{Pid =>
                                #{ref => Ref,
                                  name => Name,
                                  call_id => CallId,
                                  cleanup => nocleanup}}}}.

There are a few things going on here:

• The result is the exit reason. If the function just returned, the process would exit with normal and the 'DOWN' message would carry nothing useful. With exit(...), success and failure both arrive as one 'DOWN' message, and one handler deals with both.
• The try turns a crash into {tool_crash, Class, Reason}. Without it a crash still comes through 'DOWN', but with a full stack trace, and that’s what the LLM would get to read.
• Calls are keyed on Pid, not CallId. 'DOWN', the timeout and the cleanup message all know the Pid, so every handler can look the call up directly. Each handler also checks the Pid is still in tool_calls, so a timeout that turns up after the turn has ended is just ignored.
• The timeout is a message to ourselves after 30 seconds. Its handler kills the tool with {timeout, ...}, which then comes back through 'DOWN' like any other crash.
• OnCleanup lets a tool hand the agent a function to run if it dies before it finishes. Only bash uses it, and we’ll see why in a bit.

If the request fails while tools are running, we cancel them:

% agent.erl
cancel_tools(#{tool_calls := Calls}) ->
  maps:foreach(fun(Pid, #{ref := Ref} = Call) ->
                  erlang:demonitor(Ref, [flush]),
                  exit(Pid, kill),
                  run_cleanup(Call)
               end,
               Calls);
cancel_tools(_) ->
  ok.

We demonitor (with [flush], in case the 'DOWN' is already sitting in our mailbox), kill them, then run whatever cleanup they registered.

There was one more hole. call_llm can fail before a request even starts, while the turn is still idle. But fail_turn only matched a turn map:

% agent.erl
fail_turn(Reason, #{turn := #{reply_to := To} = Turn} = S) ->
  cancel_tools(Turn),
  To ! {msg, error, Reason},
  end_turn(S).

So the agent crashed with a function_clause, and with intensity => 0 that takes the whole app down. Now fail_turn takes any turn, and cancel_tools falls through to its catch-all for idle:

% agent.erl
fail_turn(Reason, #{turn:=Turn}=S) ->
  cancel_tools(Turn),
  end_turn(S#{error=>Reason}).

The error goes on the state now rather than to whoever started the turn. More on that later.

You: please run sleep 3600

Tool Call: bash

Agent: The command was started, but the tool timed out after 30 seconds.

>

During the tool call

[michael@archlinux] ~ ❯ pgrep -a sleep
3105434 sleep 3600

After

[michael@archlinux] ~ ❯ pgrep -a sleep
[michael@archlinux] ~ ❯

Looks good right? I thought so too at first glance, but I was wrong.

Making sure bash actually dies

Killing the tool’s process on its own isn’t enough for bash. The Erlang process owns the port, so the port gets closed when it dies, but closing a port only closes bash’s stdin and stdout. Bash doesn’t care and carries on running. That’s what OnCleanup is for: once computer:bash has opened the port, it looks up the OS pid and hands the agent a function that kills it.

That got the sleep 3600 above, but only by luck. bash -c with a single simple command just replaces itself with that command, so the pid we killed happened to be the sleep. Give it anything more complicated and bash starts child processes, which outlive it when we kill bash. It gets worse when the command backgrounds something:

sleep 3600 >/dev/null 2>&1 &

Bash exits straight away with status 0, the tool succeeds, no cleanup runs, and the sleep lives forever. Leave the redirect off and it fails a different way. The sleep keeps hold of the pipe to our port, so the port never sees EOF and the tool sits there until the 30 second timeout. Then we kill -9 a bash that already exited, whose pid could belong to something else by now. Either way the sleep survives.

Running it under timeout doesn’t help, because timeout only does anything while its child is still running, and by then bash has already exited. Process groups don’t help either, since a command can always setsid its way out of one.

What does work is a PID namespace. The first process in a new PID namespace becomes its PID 1, and when PID 1 exits the kernel kills everything else in the namespace. However hard a command tries to detach, it can’t leave. unshare will set one up for us without root:

% computer.erl
bash(Command, OnCleanup) when is_binary(Command) ->
  Port =
    open_port({spawn_executable, os:find_executable("unshare")},
              [binary,
               exit_status,
               use_stdio,
               stderr_to_stdout,
               {args,
                ["--user",
                 "--map-current-user",
                 "--pid",
                 "--fork",
                 "--kill-child",
                 "bash",
                 "-c",
                 unicode:characters_to_list(Command)]}]),
  {os_pid, OsPid} = erlang:port_info(Port, os_pid),
  OnCleanup(fun() -> os:cmd("kill -9 " ++ integer_to_list(OsPid)) end),
  collect_output(Port, []).

• --user creates a user namespace, which is what lets us create the PID namespace without being root.
• --map-current-user keeps us as ourselves inside the namespace rather than root, so whoami and anything that acts differently as root see the real user. Files end up owned by us on disk either way. It needs util-linux 2.38 or newer.
• --pid --fork starts bash in a new PID namespace as its PID 1.
• --kill-child makes unshare kill bash, and so the whole namespace, if unshare itself gets killed. That’s the process our cleanup’s kill -9 hits now.

Now sleep 3600 & comes back straight away and the sleep goes with it, and the timeout and cancel paths take everything down too. Exit codes still come through as normal. The trade-off is that the agent can’t start anything that’s meant to outlive the command, like a dev server, but for a harness I think that’s the right default. It’s also Linux only, and needs unprivileged user namespaces, which some distros restrict (Ubuntu 24.04+ does, via AppArmor). Living in a namespace has other side effects too, sudo won’t work for one, but that’s a rabbit hole for another day. One quirk: bash is PID 1 now, and PID 1 ignores any signal it hasn’t installed a handler for, so a script that does kill $$ won’t die.

Unqueue, again

With the agent in better shape I went back to /unqueue, and it was worse than I thought.

• pending is newest first, so the messages came back into the input line in reverse.
• The UI puts a message into its history as soon as you hit enter. After you unqueue it, it’s still sitting there as You: ... like it was sent.
• After an error the queued messages just stay in pending, but on screen they look sent.

All three have the same cause. The UI has no idea which of its messages the agent has actually sent. It just assumes.

If we give every message an id, the agent could tell the UI the highest id that it commits to history. It’s a cumulative ack, the same idea as TCP sequence numbers: “everything up to N is in”. This works because of two things Erlang already gives us:

• Messages between two processes always arrive in the order they were sent, and the agent flushes all of pending at once. So whatever hasn’t been committed is always the newest messages, and “everything up to N” is always true.
• The agent is a single process, so an unqueue and a flush can never overlap. Either /unqueue gets there first and takes the messages, or the flush does, the commit goes out, and /unqueue comes back empty. A message can’t be both sent and unqueued.

So the UI hands out ids, keeps unsent messages in a queued list, and only moves them into its history when the agent commits them. The agent keeps the id in pending and drops it on the way into history:

% agent.erl
%% pending is newest first, so the head holds the highest id
commit([]) ->
  {[], none};
commit([{user, _, Last} | _] = Pending) ->
  {[{user, Msg} || {user, Msg, _Id} <- Pending], Last}.

My first go kept the id in history as well and pulled it out with {_, _, Id} = hd(History). That’s fine when a user message is at the head. When it’s only tool results being flushed, the head is {tool_result, CallId, Output}, which is also a 3-tuple, so the match worked and the UI got sent the bash output as an id.

A nice side effect of only moving messages into history on commit is that the UI’s history should end up in the same order as the agent’s.

Should. It has a bug.

The bug: committing nothing commits everything

On the UI side, a commit moves everything up to that id out of queued:

% harness.erl
update({msg, commit, Last}, State) ->
  {Done, Still} =
    lists:partition(fun({Id, _}) -> Id =< Last end, maps:get(queued, State, [])),

call_llm always sent the commit, even when pending was empty, which is the normal case when we’re sending tool results back. Then the UI gets {msg, commit, none}. In Erlang every number sorts before every atom, so Id =< none is always true and everything in queued gets committed.

Most of the time queued is empty there anyway, but not always. Type a message right as the tools finish and the cast can cross with the flush. The UI moves the message into history before the agent has even seen it, and /unqueue can’t take it back out. The fix is to not send anything when there’s nothing to commit.

Stepping back

The entire time I was writing the commit code, I had a nagging feeling about the duplication of state between the UI and the agent. I’d been trying to keep them decoupled, with the UI model building its own copy of the conversation out of the events the agent sends it: delta, tool, commit, done, error. Every bug in this section came from the UI reading an event slightly differently to what the agent meant by it. And an unqueue command had turned into a sync protocol.

The problem is two copies of the conversation with two different things writing to them. If the agent is the source of truth anyway, the UI shouldn’t be writing anything. So rather than sending events, the agent pushes its state and the UI overwrites its copy with it. It doesn’t have to be the whole map every time either, as long as two rules hold:

• Every key that gets sent is the complete value of that field. The UI overwrites it, it never appends to it.
• Fields that change together go in the same message. A flush moves messages from pending to history, and if those were two messages the view could draw in between them and show the message twice, or not at all.

The ids, queued, the commit and the unqueue reconciliation can all go. The UI model still owns everything that’s actually UI: the input line, which screen you’re on, errors from commands. It just defers to the agent for the conversation.

Pushing the state

First, what subscribers get to see:

% agent.erl
%% What subscribers see. Provider formats stay inside the agent.
-type conv() ::  #{history := [conv_entry()],
                  pending := [{user, binary()}],
                  streaming := [stream_item()],
                  tools := [binary()],
                  error := none | term(),
                  phase := idle | streaming | awaiting_tools}.

-type stream_item() :: {assistant, binary()} |
                       {tool_call, CallId :: binary(), Name :: binary(), Args :: binary()}.

-type conv_entry() :: {user, binary()} | stream_item() | {tool_result, CallId :: binary(), binary()}.

A few things changed in the agent to make this work:

• The openai module now turns output items into tagged tuples, like {function_call, CallId, Name, Args}, rather than raw JSON maps.
• The agent keeps the conversation twice. input is in OpenAI’s format, ready to send. history is the same conversation as conv_entry()s, and it’s what subscribers see.
• streaming is the list of items in the current response. A new item starts a new entry, and deltas get appended to the one at the head.
• Errors live on the state, so they’re still there after the turn goes back to idle.
• reply_to is gone. The agent no longer messages whoever started the turn, it publishes to everyone subscribed.

To publish, my handle_cast and handle_info clauses move into plain cast and info functions that just return the new state. The gen_server callbacks run them, then diff the subscriber view before and after:

% agent.erl
handle_cast(Msg, State) ->
  NewState = cast(Msg, State),
  publish(State,NewState),
  {noreply, NewState}.

handle_info(Msg, State) ->
  NewState = info(Msg, State),
  publish(State,NewState),
  {noreply, NewState}.

-spec publish(state(), state())->ok.
publish(OldState,#{subscribers:=Subs}=NewState) ->
  OldConv = conv(OldState),
  NewConv = conv(NewState),
  Changed = maps:filter(fun(K, V)-> V =/= maps:get(K,OldConv)  end, NewConv),
  case map_size(Changed) of
    0 -> ok;
    _ -> [P ! {conv, Changed} || P <- maps:keys(Subs)]
  end,
  ok.

Only the keys that changed get sent, and because each one is sent whole, both rules hold for free. Subscribing gets you the full state straight away, and a monitor drops you if you die:

% agent.erl
cast({subscribe, Pid}, #{subscribers:=Subscribers}=State) when is_map_key(Pid, Subscribers)->
  State;
cast({subscribe, Pid}, #{subscribers:=Subscribers}=State)->
  Pid ! {conv, conv(State)},
  State#{subscribers => Subscribers#{Pid=>erlang:monitor(process, Pid)}};

% subscriber crashed
info({'DOWN', _Ref, process, Pid, _Exit},
            #{subscribers := Subscribers} = State)
  when is_map_key(Pid, Subscribers) ->
  State#{subscribers => maps:remove(Pid, Subscribers)};

On the UI side, the model subscribes when it starts, and one update clause replaces every {msg, ...} event we had before. No deltas, no increments, just overwrite:

% harness.erl
update({conv, NewConv}, #{conv:=ConvState}=State)->
  Conv = case ConvState of
    none -> NewConv;
    OldConv -> maps:merge(OldConv, NewConv)
  end,
  {State#{conv=>Conv},[]};

create_model(ViewPid) ->
  Pid = spawn_link(fun() ->
                       gen_server:cast(agent, {subscribe, self()}),
                       loop(ViewPid, #{conv => none}) end),
  register(model, Pid),
  {ok, Pid}.

And unqueue is just a cast. The agent sends back whatever was pending and clears it, and the push takes care of the rest:

% harness.erl
run_effect(unqueue) ->
  gen_server:cast(agent, {user, unqueue, self()});

update({unqueued, []}, State)->
  {State, []};
update({unqueued, Unqueued}, State)->
  Taken = iolist_to_binary(lists:join(<<"\n">>, [M || {user, M} <- Unqueued])),
  {State#{text => Taken}, []};

% agent.erl
cast({user, unqueue, From}, #{pending := Pending} = State) ->
  From ! {unqueued, lists:reverse(Pending)},
  State#{pending => []};

Much simpler than the commit protocol. The downside is size: every push carries whole values, so a long conversation means big messages, and every delta resends the whole of streaming. For a chat in one terminal that’s fine.

That’s the turn handling sorted. In part 3 I’ll make conversations survive a restart with mnesia.