Branded Types: adding feedback loops to my agents

Table of contents

Preamble

I left an agent refactoring a feature a few days ago. When I came back, I found a codebase with dozens of silent bugs that threw no type errors and looked almost perfect. They were all the same mistake: the agent decided to pass the ID of one resource where the ID of another resource was expected, but since both IDs are string, nothing stopped it. We can do better.

The refactor

At Fintoc, we recently decided to start offering our clients a product that had been internal-only until now. A kind of harness for the non-technical team, that works through the company’s chat. To offer it to other companies, I needed to change several things about how this product worked.

So I was doing a big refactor on the system. The idea was to split the concept of “user” into two distinct entities: an identity (the “profile”: name, timezone) and an authenticated user (the entity that the system’s resources are associated with). After the refactor, the profile became the central entity, but the resources (memories, commands, conversations) stayed associated with the user, not the profile. Due to the size of the system and the number of features, this change involves a ton of touch points. But conceptually it’s pretty simple.

I had already spent hours in conversations with agents for other parts of the changes. I had the refactor mapped out to the last detail, the plan ready, and a big chunk of the refactor executed, including the structural changes for the identity <-> user separation.

Up to this point, everything is still reasonable.

The context

To understand what kind of things broke across the entire application, I’m going to explain a specific case where the agent broke the system in a very subtle way, but this breakage pattern presented itself across the codebase pretty consistently.

The agent’s command system had an interface that looked like this:

export type CommandContext = {
  message: InboundMessage
  user: User
  conversation: Conversation
  reply: (text: string) => Promise<void>
};

export type Command = {
  pattern: string | RegExp
  execute: (ctx: CommandContext) => Promise<void | CommandExecutionResult>
}

And implementing a command looks like this:

import { commands } from '~/commands';

commands.register({ // type Command
  pattern: '!greet',
  execute: async ({ user, reply }) => {
    await reply(`Hey ${user.name}, how's it going?`)
  }
});

As extra context, many functions in the codebase received IDs from different entities to operate. And all these IDs were string:

export const createMemoryStore = async (options: {
  organizationId: string
  userId: string
  name: string
  // ...
}) => { /* ... */ };

So, there was this command to create an agent with its own memory that looked roughly like this:

commands.register({
  pattern: /^!agents\\s+create\\s+(.+)$/i,
  execute: async (ctx) => {
    const name = content.match(/^!agents\\s+create\\s+(.+)$/i)[1].trim();

    // ...

    await createMemoryStore({
      organizationId: ctx.conversation.organizationId,
      userId: ctx.user.id,
      name
    });

    // ...
  }
});

Also, the User and Identity types look like this:

type User = {
  id: string
  email: string
}

type Identity = {
  id: string
  userId?: string
  name: string
  timezone: string
}

OK, with this context, let’s see what happened to the agent.

The problem

Finally came the part of the refactor that involved actually using the identity in the logic.

The first thing my agent did was change CommandContext to use identity instead of user, which was the right call:

export type CommandContext = {
  message: InboundMessage
  identity: Identity // <- previously `user: User`
  conversation: Conversation
  reply: (text: string) => Promise<void>
};

So far so good. But now, the command code was throwing a type error, because it was calling createMemoryStore with userId: ctx.user.id, and ctx.user no longer exists.

And the agent decided to use the identity’s id.

await createMemoryStore({
  organizationId: ctx.conversation.organizationId,
  userId: ctx.identity.id, // <- here's the bug
  name
});

Technically, there are no type errors (both IDs are strings), and since both concepts are relatively similar, the model decided it was valid to make that change.

But it silently broke the codebase. And it did it across the entire codebase. Like it was nothing.

You can imagine I got pretty paranoid, and I went through the whole codebase. Every call that received a userId, by hand. I found more than 10 errors, all without type issues or runtime errors, just failing silently (memories weren’t found, the agent would respond that it couldn’t find commands, etc).

The root of the problem

The underlying problem is that all IDs in the system (and therefore for the agent) are the same thing. Every table has a string for the ID. But semantically, a user ID and an identity ID are completely different things, and treating them as the same type is what led the agent to mix up which ID to pass where.

Branded types

It occurred to me that if the agent had had some way to see that the IDs it used “didn’t match”, it probably would have corrected the error automatically.

I borrowed from Convex the idea that IDs should be typed by table, even though at runtime they’re simply strings. The way to do this in TypeScript is with branded types. The idea is to put an invisible “brand” on a type so that TypeScript treats them as incompatible, even though at runtime they’re still exactly the same type.

The complete type fits in 3 lines:

// TableName looks like: 'users' | 'identities' | ...
export type Id<TTable extends TableName> = string & {
  readonly __table: TTable
};

That’s it. Id<'users'> is a string that has a phantom property __table with value 'users'. This property doesn’t exist at runtime. It’s never set, never read, takes no memory, doesn’t show up in JSON.stringify. It’s purely a type system artifact.

But for TypeScript, Id<'users'> and Id<'identities'> are incompatible types. You can’t assign one to the other. You can’t pass one where the other is expected. And if you try, TypeScript throws an error. Beautiful.

The only annoying part is that you have to define which tables can be passed to the type and keep it in sync with the system’s tables (Convex solves this automatically with codegen, but I thought that was a bit overkill for my use case):

export const tableNames = [
  'users',
  'identities',
] as const;

export type TableName = (typeof tableNames)[number];

Branding the schema

The next step was to brand the columns in the database schema. This project uses Drizzle, which lets you specify the TypeScript type a column should have with .$type<T>():

export const users = pgTable('users', {
  id: text('id').$type<Id<'users'>>().primaryKey().$defaultFn(() => createId() as Id<'users'>),
  // ...
});

export const identities = pgTable('identities', {
  id: text('id').$type<Id<'identities'>>().primaryKey().$defaultFn(() => createId() as Id<'identities'>),
  userId: text('user_id').$type<Id<'users'>>().references(() => users.id).unique(),
  // ...
});

And the inferred types from the schema now carry the id as a branded type automatically:

export type User = RectifyIds<typeof users.$inferSelect>;
export type Identity = RectifyIds<typeof identities.$inferSelect>;

Bonus: RectifyIds

There’s something subtle going on here. When Drizzle infers the type from identities.$inferSelect, it “unravels” the Id<T> type. Instead of looking like this:

{
  id: Id<'identities'>
  userId: Id<'users'>
  // ...
}

It looks like this:

{
  id: string & { __table: 'identities' }
  userId: string & { __table: 'users' }
  // ...
}

That’s what RectifyIds is there for:

type RectifyIdField<TField> = TField extends Id<infer TTableName>
  ? Id<TTableName>
  : TField;

export type RectifyIds<TTable extends { id?: string }> = Prettify<{
  [TKey in keyof TTable]: RectifyIdField<TTable[TKey]>;
}>;

Two very simple type helpers here. RectifyIdField checks if a field extends the Id<T> type (string & { __table: T } return positive for this check), and if it does, it reassigns it to Id<T>, which “re-brands” the type, undoing the unraveling. Then, RectifyIds iterates over the table’s fields and applies this logic to all of them. It’s a trick to make the type “look nicer”.

This type uses Prettify, which is an internal utility type that “flattens” a type so it displays cleanly in the IDE:

export type Prettify<T> = {
  [K in keyof T]: T[K];
} & {};

Three lines. All it does is iterate over the type’s keys and reassign them. The & {} at the end is a trick to prevent TypeScript from collapsing the type back to the original. It’s pretty useful when you do Type & { key: 'value' }, so the final type doesn’t show the & but instead “absorbs” it.

The function signatures

Once the schema was branded, the service functions changed their signatures:

export const createMemoryStore = async (options: {
  organizationId: Id<'organizations'>
  userId: Id<'users'>
  name: string
  // ...
}) => { /* ... */ };

Now it’s impossible to pass an Id<'identities'> where an Id<'users'> is expected. TypeScript will throw an error at typecheck.

The agent that found more bugs

While the agent was changing all the function signatures to stop receiving string and receive the corresponding Id<T> instead, it kept running tsc after each change and fixing the type errors that came up.

And here’s where it gets good: it found errors that I had missed. Places where the wrong ID was being passed that I hadn’t caught in my manual review. The typecheck found them without having to understand a single line of business logic. Just because the types were no longer congruent.

And it’s not that the agent is smarter. It just has a better feedback loop. The more descriptive the types are when writing code, the less smart the agent needs to be to write software that works.

Why this matters

In my experience these agents are pretty wild in terms of productivity, but they’re very prone to certain kinds of errors. Particularly, errors that require understanding semantic context that isn’t expressed in the code. If two things are named differently but have the same type, an agent can confuse them, with so much confidence that sometimes it looks like it’s correct.

At the end of the day what I wanted to show was that by improving a little bit one part of the codebase, the agent became much more effective. Slightly better types here, slightly stricter linter there, and the agent gets back on track on its own after a while.

Branded types are a perfect example of a minimal complexity investment (literally 3 extra lines in the Id type) that significantly improves the system’s ability to catch errors. At runtime, the IDs are still just strings, but at typecheck, the compiler has enough information to distinguish between a user ID and an identity ID, and therefore so does the agent.

And they don’t just apply to strings. In theory, you can build a branded type for anything: numbers, objects, etc. By adding the “branding”, types with different brands become incongruent, even if the base type is exactly the same.

Well, one refactor led to another, but hopefully this was a fun read. I had a good time.

Peer Review