Write a custom wire codec

/** One message in your provider's world. */
export interface Turn {
  id: string;
  role: 'user' | 'assistant';
  text: string;
}

/** Client to agent. Every input carries a `kind` the codec dispatches on. */
export type MyInput = { kind: 'user-message'; message: Turn };

/** Agent to client. Text arrives as a start / delta / end trio. */
export type MyOutput =
  | { type: 'text-start'; id: string }
  | { type: 'text-delta'; id: string; delta: string }
  | { type: 'text-end'; id: string }
  | { type: 'tool-call'; callId: string; toolName: string; args: unknown }
  | { type: 'done'; reason: string };

import { defineCodec, jsonField, strField } from '@ably/ai-transport';
import type { MyInput, MyOutput } from './events';

// Each field binds one header key to one property, so encode and decode
// can never disagree about where a value lives.
const idField = strField('id');
const callIdField = strField('callId', '');
const toolNameField = strField('toolName', '');
const reasonField = strField('reason', '');
const argsField = jsonField('args');
const messageIdField = strField('messageId', '');

export const myCodec = defineCodec<MyInput, MyOutput>()({
  adapterTag: 'my-provider',

  output: ({ event, stream }) => [
    // One growing message per assistant reply. `streamId` is what ties the
    // three phases together on the wire.
    stream('text', {
      streamId: (chunk) => chunk.id,
      fields: [idField],
      start: { type: 'text-start' },
      delta: { type: 'text-delta', field: 'delta', decode: ({ rebuild }) => rebuild([idField]) },
      end: { type: 'text-end' },
    }),

    // Discrete: one publish each.
    event('tool-call', { fields: [callIdField, toolNameField, argsField] }),
    event('done', { fields: [reasonField] }),
  ],

  input: ({ batch }) => [
    // A user message fans out into one wire event per part. This provider has
    // one part kind, so the table has one entry.
    batch('user-message', {
      explode: (input) => [{ type: 'text', text: input.message.text }],
      partTypeOf: (part) => part.type,
      parts: (p) => [
        p('text', {
          data: { encode: (part) => part.text, decode: (d) => ({ text: String(d ?? '') }) },
        }),
      ],
      // Stamped on every part, so the decode side rebuilds the envelope from any one.
      messageHeaders: (input) => ({
        codecHeaders: { messageId: input.message.id },
        transportHeaders: { role: input.message.role },
      }),
      assemble: (part, ctx) => ({
        message: { id: messageIdField.read(ctx.codecHeaders), role: 'user', text: part.text },
      }),
    }),
  ],
});

// Agent-side. This runs on your own infrastructure, so it holds the key.
import { after } from 'next/server';
import * as Ably from 'ably';
import { createAgentTransport } from '@ably/ai-transport';
import { myCodec } from '../../codec';
import { appendTurn, loadConversation } from '../../store';

const ably = new Ably.Realtime({ key: process.env.ABLY_API_KEY });

// Your provider, reduced to what this guide needs: an async iterable of MyOutput.
async function* runModel(conversation, signal) {
  const id = crypto.randomUUID();
  yield { type: 'text-start', id };
  for await (const fragment of yourProvider(conversation, { signal })) {
    yield { type: 'text-delta', id, delta: fragment };
  }
  yield { type: 'text-end', id };
  yield { type: 'done', reason: 'stop' };
}

export async function POST(req) {
  const { conversationId, eventId } = await req.json();

  const transport = createAgentTransport({
    channel: ably.channels.get(conversationId),
    codec: myCodec,
    clientId: 'agent',
  });
  await transport.connect();

  after(async () => {
    try {
      // locateInput scans channel history for the input carrying this event id.
      // It is the one history read the SDK does for you.
      const trigger = await transport.locateInput(eventId);
      const prompt = trigger?.inputs.find((input) => input.kind === 'user-message');

      // Passing inputCodecMessageId lets a cancel keyed on the input find this
      // run, including one that arrived before the run opened.
      const run = transport.openRun({ inputCodecMessageId: trigger?.meta.codecMessageId });

      const conversation = [...loadConversation(conversationId), prompt.message];
      appendTurn(conversationId, prompt.message);

      const reply = { id: crypto.randomUUID(), role: 'assistant', text: '' };
      const { reason } = await run.pipe(
        (async function* () {
          for await (const chunk of runModel(conversation, run.abortSignal)) {
            if (chunk.type === 'text-delta') reply.text += chunk.delta;
            yield chunk;
          }
        })(),
      );
      await run.end({ reason });

      if (reason === 'complete') appendTurn(conversationId, reply);
    } finally {
      transport.close();
    }
  });

  return Response.json({ ok: true });
}

'use client';

// Client-side. The browser fetches a token; never put an API key here.
import { useEffect, useRef, useState } from 'react';
import * as Ably from 'ably';
import { createClientTransport } from '@ably/ai-transport';
import { myCodec } from './codec';
import { createFold } from './fold';

export function Chat({ conversationId, seed }) {
  const [messages, setMessages] = useState(seed);
  const [input, setInput] = useState('');
  const [activeRunId, setActiveRunId] = useState(undefined);
  const transportRef = useRef(null);

  useEffect(() => {
    const ably = new Ably.Realtime({ authUrl: '/api/auth/token' });
    const fold = createFold();
    const transport = createClientTransport({
      channel: ably.channels.get(conversationId),
      codec: myCodec,
    });
    transportRef.current = transport;

    transport.subscribe((event) => {
      if (event.kind === 'run-lifecycle') {
        setActiveRunId(event.event.type === 'start' ? event.event.runId : undefined);
        return;
      }
      fold.apply(event);
      setMessages([...seed, ...fold.render()]);
    });

    void transport.connect();
    return () => {
      transport.close();
      ably.close();
    };
  }, [conversationId]);

  const send = async (text) => {
    const sent = await transportRef.current.publishInput({
      kind: 'user-message',
      message: { id: crypto.randomUUID(), role: 'user', text },
    });
    // The SDK publishes; waking the agent is yours.
    await fetch('/api/chat', {
      method: 'POST',
      body: JSON.stringify({ conversationId, eventId: sent.eventId }),
    });
  };

  return (
    <div>
      {messages.map((m, i) => (
        <div key={i}><strong>{m.role}:</strong> {m.text}</div>
      ))}
      <form onSubmit={(e) => { e.preventDefault(); void send(input); setInput(''); }}>
        <input value={input} onChange={(e) => setInput(e.target.value)} />
        {activeRunId ? (
          <button type="button" onClick={() => transportRef.current.cancel(activeRunId)}>Stop</button>
        ) : (
          <button type="submit">Send</button>
        )}
      </form>
    </div>
  );
}

// Client-side. Fold the transport's event stream into a renderable list.
export function createFold() {
  const messages = new Map();     // codecMessageId -> { role, text, stepId, stepStartSerial }
  const canonical = new Map();    // stepId -> the newest attempt's start serial

  function apply(event) {
    if (event.kind !== 'message') return;
    const { codecMessageId, stepId, stepStartSerial, role } = event.meta;
    if (!codecMessageId) return;

    // A retry publishes under the same step id with a higher start serial.
    // Record the newest, so render drops the attempt it replaced.
    if (stepId && stepStartSerial && (canonical.get(stepId) ?? '') < stepStartSerial) {
      canonical.set(stepId, stepStartSerial);
    }

    const entry = messages.get(codecMessageId) ?? { role, text: '', stepId, stepStartSerial };
    for (const input of event.inputs) {
      if (input.kind === 'user-message') entry.text = input.message.text;
    }
    for (const output of event.outputs) {
      if (output.type === 'text-delta') entry.text += output.delta;
    }
    messages.set(codecMessageId, entry);
  }

  function render() {
    return [...messages.values()].filter(
      (m) => !m.stepId || canonical.get(m.stepId) === m.stepStartSerial,
    );
  }

  return { apply, render };
}