Branded Types: agregando feedback loops a mis agentes

Dejé un agente refactoreando una feature hace unos días. Al rato llegué a un codebase con decenas de bugs silenciosos que no tiraban type errors y se veían casi perfecto. Todos eran el mismo error: el agente decidió pasar el ID de un recurso donde iba el ID de otro recurso, pero como ambos IDs son string, nada lo detuvo. Pero podemos hacerlo mejor.

El refactor

Hace poco decidimos empezar a ofrecer a nuestros clientes un producto que hasta ahora había sido solamente interno. Una especie de harness para el equipo no técnico, que funciona a través del chat de la empresa. Para poder ofrecerlo a otras empresas, necesitaba cambiar varias cosas de cómo funcionaba este producto.

Estaba haciendo entonces un refactor grande en el sistema. La idea era separar el concepto de "usuario" en dos entidades distintas: una identidad (perfil: nombre, timezone) y un usuario autenticado (la entidad con la que se relacionan los recursos del sistema). Después del refactor, el perfil pasó a ser la entidad central, pero los recursos (memorias, comandos, conversaciones) siguieron asociados al usuario, no al perfil. Debido al tamaño del sistema y a la cantidad de features, este cambio supone muchísimos puntos de contacto. Pero conceptualmente es bien simple.

Ya llevaba horas de conversaciones con agentes para otras partes de los cambios. Tenía el refactor mapeadísimo, el plan armado, y un gran chunk del refactor ejecutado, incluidos los cambios estructurales de la separación identidad <-> usuario.

Hasta acá, todo razonable.

El contexto

Para entender qué tipo de cosas se rompieron en toda la aplicación, voy a explicar un caso muy puntual donde el agente rompió el sistema de forma muy sutil, pero esto se repitió a través del codebase de manera bastante consistente.

El sistema de comandos del agente tenía una interfaz que se veía así:

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>
}

E implementar un comando se ve así:

import { commands } from '~/commands';

commands.register({ // type Command
  pattern: '!greet',
  execute: async ({ user, reply }) => {
    await reply(`Wenas ${user.name}, qué tal?`)
  }
});

Como contexto extra, muchas funciones del codebase recibían IDs de distintas entidades para operar. Y todos estos IDs eran string:

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

Entonces, había un comando para crear un agente con su propia memoria que se veía más o menos así:

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
    });

    // ...
  }
});

Además, los tipos de User y Identity se ven así:

type User = {
  id: string
  email: string
}

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

OK, con este contexto, veamos qué le pasó al agente.

El problema

Finalmente llegó la parte del refactor que implicaba efectivamente usar la identidad en la lógica.

Lo primero que hizo mi agente fue cambiar CommandContext para que use identity en vez de user, que fue el cambio adecuado:

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

Hasta ahí todo OK. Pero ahora, el código del comando tiraba un type error, porque estaba llamando createMemoryStore con userId: ctx.user.id, y ctx.user ya no existe.

Y el agente decidió usar el id de la identidad.

await createMemoryStore({
  organizationId: ctx.conversation.organizationId,
  userId: ctx.identity.id, // <- acá está el bug
  name
});

Técnicamente, no hay type errors (ambos IDs son strings), y como ambos conceptos son relativamente parecidos, el modelo decidió que era válido hacer ese cambio.

Pero rompe el codebase silenciosamente. Y lo hizo en todo el codebase. Como si nada.

Te imaginarás que quedé bastante paranoico, y me puse a revisar todo el codebase. Cada llamada que recibía un userId, a mano. Encontré más de 10 errores, todos sin problemas de tipos ni runtime, solamente fallando en silencio (no se encontraban memorias, el agente respondía que no encontraba comandos, etc).

La raíz del problema

El problema de fondo es que todos los IDs en el sistema (y por consiguiente para el agente) son lo mismo. Todas las tablas tienen un string para el ID. Pero semánticamente, el ID de un usuario y el ID de una identidad son cosas completamente distintas, y tratarlos como el mismo tipo es lo que resultó en el agente confundiendo qué ID pasar en qué lugar.

Branded types

Se me ocurrió que si el agente hubiera tenido alguna forma de ver que los IDs que usó "no calzaban", probablemente habría corregido el error automáticamente.

Tomé prestado de Convex la idea de que los IDs están tipados por tabla, aunque en runtime sean simplemente strings. La forma de hacerlo en TypeScript es con branded types. La idea es ponerle una "marca" invisible a un tipo para que TypeScript los trate como incompatibles, aunque en runtime sigan siendo exactamente el mismo tipo.

El tipo completo cabe en 3 líneas:

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

Eso es todo. Id<'users'> es un string que tiene una propiedad fantasma __table con valor 'users'. Esta propiedad no existe en runtime. Nunca se setea, nunca se lee, no ocupa memoria, no aparece en JSON.stringify. Es puramente un artefacto del sistema de tipos.

Pero para TypeScript, Id<'users'> y Id<'identities'> son tipos incompatibles. No se pueden asignar uno al otro. No se pueden pasar donde se espera el otro. Y si uno trata de hacer, TypeScript tira un error. Hermoso.

Lo único "fome" es que hay que definir qué tablas pueden pasarse al tipo y mantenerlo en sync con las tablas del sistema (Convex resuelve esto automáticamente con codegen, pero creo que era un poco overkill para mi caso de uso):

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

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

Brandeando el schema

El siguiente paso fue marcar las columnas en el schema de la base de datos. Este proyecto usa Drizzle, que permite especificar el tipo de TypeScript que debería tener una columna con .$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(),
  // ...
});

Y los tipos inferidos del schema ahora cargan el id como branded type automáticamente:

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

Bonus: RectifyId

Acá pasa algo sutil. Cuando Drizzle infiere el tipo de users.$inferSelect, "desenreda" el tipo Id<T>. En vez de verse así:

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

Se ve así:

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

Para eso está ese RectifyId ahí metido:

type ExtendTableId<TTable extends { id?: string }, TId extends string> = Prettify<
  Omit<TTable, 'id'>
  & (TTable extends { id: string }
    ? { id: TId }
    : { id?: TId })
>;

export type RectifyId<TTable extends { id?: string }> = TTable extends { id?: infer TId }
  ? TId extends Id<infer TTableName>
    ? ExtendTableId<TTable, Id<TTableName>>
    : TTable
  : TTable;

Se puede ver medio complejo, pero básicamente lo que hace es ver si el tipo recibido (el tipo de la tabla) extiende id: Id<'table'> (que va a estar "desenredado" en string & { __table: 'table' }), y si lo extiende, lo re-asigna a Id<'table'>, que es exactamente lo mismo, pero sin desenredar. Es un truco para que el tipo se "vea mejor".

Este tipo usa Prettify, que es un utility type interno que "aplana" un tipo para que se muestre limpio en el IDE:

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

Tres líneas. Lo único que hace es iterar sobre las keys del tipo y reasignarlas. El & {} al final es un truco para que TypeScript no colapse el tipo de vuelta al original. Es bien útil para cuando uno hace Type & { key: 'value' }, para que el tipo final no muestre el &, sino que lo "absorba".

Las firmas de las funciones

Una vez que el schema estaba brandeado, las funciones de servicio cambiaron sus firmas:

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

Ahora es imposible pasar un Id<'identities'> donde se espera un Id<'users'>. TypeScript va a tirar un error en typecheck.

El agente que encontró más bugs

Mientras el agente cambiaba todas las firmas de las funciones para que dejaran de recibir string y recibieran los Id<T> correspondientes, iba corriendo tsc después de cada cambio y arreglando los errores de tipo que iban saliendo.

Y acá viene lo bueno: encontró errores que a mí se me habían pasado. Lugares donde se estaba pasando el ID equivocado y que yo no había detectado en mi revisión manual. El typecheck los encontró sin tener que entender una sola línea de lógica de negocio. Solo por el hecho de que los tipos ya no eran congruentes.

Y no es que el agente sea más inteligente. Solo tiene un mejor feedback loop. Mientras más descriptivos son los tipos al momento de escribir el código, menos inteligente necesita ser el agente para escribir software que funciona.

Por qué importa esto

En mi experiencia estos agentes son bien locos en términos de productividad, pero son muy propensos a ciertos tipos de errores. Particularmente, errores que requieren entender contexto semántico que no está expresado en el código. Si dos cosas se llaman distinto pero tienen el mismo tipo, un agente puede confundirlas, con tanta confianza que a veces parece que estuviera bien.

Al final lo que quería mostrar era que mejorando un poco una parte del codebase, el agente fue mucho más efectivo. Tipos un poco mejores por acá, linter un poco más estricto por allá, y el agente solito se encausa después de un rato.

Los branded types son un ejemplo perfecto de una inversión de complejidad mínima (literalmente 3 líneas extra en el tipo Id) que mejora harto la capacidad del sistema para detectar errores. En runtime, los IDs siguen siendo simplemente strings, pero en typecheck, el compilador tiene la información suficiente para distinguir entre un ID de usuario y un ID de identidad, y por ende el agente también.

Y no aplican solo a strings. En teoría, uno puede armar un branded type de lo que sea: números, objetos, etc. Con agregarle el "branding", quedan incongruentes los tipos con distinto brand, aunque el tipo base sea exactamente el mismo.

Bueno, un refactor llevó a otro, pero ojalá que haya sido entrete. Yo lo pasé bien.

Peer Review