Build an AI Agent in Typescript / lesson 4 of 8

Task Memory

Fix AI coding agent drift on multi-step tasks with a built-in to-do manager.

Play
Transcript

[00:00] One of the problems we can run into with our coding agent is on multi-step tasks. So you’ve got five steps, right? Maybe it takes two, three refactors per step. And by the time we get to step three, the AI is wandering around. It doesn’t know what it’s doing. It’s starting to improvise because it’s kind of lost the thread. So what we’re going to do for Jeeves, our coding agent, is we’re going to create a to-do list.

[00:22] There’s our to-do TS. I’m going to drop in some initial setup stuff here. So we’ve got a to-do status, just a string union, a to-do item, kind of the standard elements of that. These little marks here are just to allow us to output this to the screen in a way that we can visually make sense of. We’re going to set up our to-do manager. We’re going to initialize our items. And we’re really only going to have two methods here. One’s going to be update, and the other’s going to be render. And at the end of every update, we’re going to call this.render.

[01:06] And into our update, we’re going to get our items, which could be a list of to-do items. Then we’re just going to have a couple guardrails in here. So if our item’s length is greater than 20, we’re going to throw a new error. Max 20 to-do’s allowed. And that’s going to keep us from building up a list of too many to-do’s. Now we’re going to take a look at how many items are in progress. We want to make sure that our AI is only doing one task at a time. That’s going to keep us, if we have a list of five tasks, it’s going to keep us moving through that list sequentially one at a time. So we’ll just filter that on the status of in progress. And if it’s greater than one, we’re also going to throw an error. We’ll break this out a little bit. And if all that’s good, we’ll just say this.items equals items. Okay, so now we can move on to our render method.

[02:15] So first, we’re going to check to see if, in fact, we have any items. And if we don’t, we’ll just say there are no to-do’s. Okay, now we’re going to build up our list for displaying in the UI. So list, we’ll map that guy. We’ll get our mark and our ID and our text. We’ll just join those with some line breaks. And now we’ll just get a quick summary. So we’ll say how many are done, right? We’ll filter. Uncompleted. And get the length of those. And then we’ll return our list of how many there are completed. Okay, so that’s our basic to-do manager. I’m going to save that. We’re going to jump over to our tools. We’ll bring in our to-do manager and our type of to-do item. And right here at the top, we’re going to initialize our new to-do manager.

[03:24] And then we’re not going to need to do like a run to-do, because all of that is managed in the to-do manager. So we’re just coming down here to our list of tools. And we’re going to have a new tool called to-do. It’s going to be a tool just like the rest of them. It’s going to have a description. That’ll do. It’s going to have an input schema, which I’m going to paste in here because I keep messing up the typing. But it’s our Zod schema with an object that is the items, an array, and then just the object of our to-do items. The ones that we just defined. So then we’re going to have our execute. It’s going to be an async. Take in our items. And then we’re going to grab this just so that we can console.log it. So output equals to-do.update items, which will trigger render. And then we’re just going to console.log to-do line break and output. And then we will return the output. Okay, so that’s our tool.

[04:39] We’re also going to jump over to our index and we’re going to update our system prompt. I’m just going to drop this in. So we still have your coding agent at work directory, but we’ve added this guy. Use the to-do tool to plan multi-step tasks. Mark in progress before starting. Completed when done. And then we still have our use the task tool to delegate, you know, create sub-agents later. Now one thing we are going to add here is a brand new hook, and that’s going to be on-step finish. Now there’s a bunch of things we’re going to get out of here, but what we want is tool calls. So that’s letting us know if the AI has called a tool. So now we need to look and see, did it use the to-do tool? So we’re going to say tool calls dot sum tool call dot tool name is equal to to-do. Otherwise, use to-do is false.

[05:32] So what I want to do here is keep track of how many times it’s used the to-do and remind it to update its to-do’s. So up here I’m going to say cycles since, or let’s say steps since to-do. Whatever. It doesn’t matter. It’s just an integer that we’re going to track. So now we’re going to say steps since to-do is equal to, if we just used it, it’s going to be zero because we just used it. So we’re trying to see since we’ve used it. So used to-do zero, otherwise we’ll say steps since to-do plus one. And if our steps since to-do, let’s say it’s greater than or equal to three, into our messages, we’re going to push a reminder. So again, our messages are just objects with a role and content. So our role is going to be user. We’re the ones reminding it. Oops. Our content is going to be reminder. Update your to-do’s. Okay. I think we’re looking good. We are going to try this guy out.

[07:07] I’m going to come down here. Now I’ve got what I hope is going to be a good prompt. So we’re going to run our index. Now here’s my prompt. Create a file called hello, add types, comments, guards, then create greet, add types, comments, guards, that uses hello TS, along with a node reline to accept the name for the hello function. This might not make a bunch of sense to the AI. I’m not sure what it’s going to output, but I just wanted to have something that seemed to have a bunch of steps. All right. I’m going to run this and hope that I didn’t screw something up. Oh, I got an undefined in there on the to-do. I’m actually going to go ahead and stop this and look at where we might be getting that undefined from. So that was when we were logging it out, which is going to be in our tools. Our output was on the to-do update. Come back here. Ah, this should be return this dot render. Okay. I think we’re good now. Try this one more time.

[08:33] And to be clear, that probably wasn’t going to be hurting anything. I’m not sure, but better safe than sorry. Okay. So we’ve got our list of two to-dos. It’s working on number one. And one is completed. And now it’s working on task two. So we’ll look at this. Yeah, it created a pretty reasonable, kind of overly engineered version of that function. And now task two of two is completed. And then it’s going to come back with a summary. So there we go. Our AI, or our coding agent, now has a basic state slash task management system that will help keep it on track.

Fix AI coding agent drift on multi-step tasks with a built-in to-do manager. This agentic AI solution enforces sequential task execution and state tracking to prevent AI improvisation and keep your coding agent focused on long workflows.

The To-Do Manager

Let’s start by creating a new file called todos.ts.

todos.ts
export type TodoStatus = "pending" | "in_progress" | "completed";
export type TodoItem = { id: string; text: string; status: TodoStatus; };
export const MARKS: Record<TodoStatus, string> = {
pending: "[ ]",
in_progress: "[>]",
completed: "[x]",
};
export class TodoManager {
private items: TodoItem[] = [];
update(items: TodoItem[]) {
if (items.length > 20) throw new Error("Max 20 todos allowed");
if (items.filter(i => i.status === 'in_progress').length > 1) {
throw new Error("Only one task can be in progress at a time");
}
this.items = items;
return this.render();
}
render() {
if (!this.items.length) return "no todos";
const list = this.items
.map(i => `${MARKS[i.status]} ${i.id} ${i.text}`)
.join('\n');
const done = this.items.filter(i => i.status === 'completed').length;
return `${list}\n\n(${done}/${this.items.length} completed)`;
}
}

Add Todo to Tools

Now let’s add our TodoManager tool to tools.ts.

tools.ts
import { tool, zodSchema, generateText, isLoopFinished } from "ai";
import {z} from "zod":
import { spawnSync } from "node: child_process";
import { api, MODEL, WORKDIR } from './config'
import { TodoManager } from "./todo";
const TODO = new TodoManager();
...
/** TOOLS */
const runBash = (command: string): string => {/* ... */};
const runRead = (filePath: string, limit?: number) => {/* ... */};
const runWrite = (filePath: string, content: string) => {/* ... */};
const runEdit = (filePath: string, oldText: string, newText: string) => {/* ... */};
export const TOOLS = {
bash: tool({/* ... */}),
read_file: tool({/* ... */}),
edit_file: tool({/* ... */}),
write_file: tool({/* ... */}),
todo: tool({
description: "Update task list. Track progress on multi-step tasks.",
inputSchema: zodSchema(
z.object({
items: z.array(
z.object({
id: z.string(),
text: z.string(),
status: z.enum(["pending", "in_progress", "completed"]),
}),
),
}),
),
execute: async ({ items }) => {
const output = TODO.update(items);
console.log(`> todo:\n${output}`);
return output;
},
}),
};

Update Agent Loop

Now let’s update index.ts to use the todo tool. Here will utilize the hook onStepFinish to track when the todo tool has been called.

index.ts
const agentLoop = async (messages: ModelMessage[]): Promise<string> => {
let stepsSinceTodo = 0;
const { text } = await generateText({
model: api.chat(MODEL),
system: `You are a coding agent at ${WORKDIR}. Use bash to solve tasks. Act, dont explain.`,
system: `You are a coding agent at $(WORKDIR}.
Use the todo tool to plan multi-step tasks. Mark in_progress before starting, completed when done.
Use the task tool to delegate exploration or subtasks.`
messages,
tools: PARENT_TOOLS,
stopWhen: isLoopFinished(),
onStepFinish: ({ toolCalls }) => {
const usedTodo = toolCalls.some((tc) => tc.toolName === "todo") ?? false;
stepsSinceTodo = usedTodo ? 0 : stepsSinceTodo + 1;
if (stepsSinceTodo >= 3) {
messages.push({
role: "user",
content: "<reminder>Updated your todos.</reminder>",
});
}
},
});
return text;
};

Try it out

Terminal window
bun run index.ts
input » Create a file called hello.ts (add types, JSDOC comments, and guards), then create greet.ts (add types, JSDOC comments, and guards) that uses hello.ts along with a node readLine to accept a name for the hello function

Share this post on:

Previous
Subagents
Next
Task Persistence