Add a Design Chat to Your Product

import { readFileSync } from 'node:fs';

import {
  googleFonts,
  prepareCoDesign,
  remoteExecutors,
  type ToolCallEvent,
  type UploadExportRequest
} from '@imgly/codesign-sdk';

import { runYourModel } from './agent-loop.js';
import { brandTool } from './brand-tool.js';
import { designStoreFor } from './design-store.js';
import { executorPool } from './executor-pool.js';
import { waitForExecutor, type TurnEvent } from './jobs.js';
import {
  chatDesigns,
  planHasExports,
  rememberChatDesigns,
  saveRender,
  signedDownloadLink,
  tokenFor,
  uploadToStorage
} from './yours.js';

// The plugin module, sent as text. With a bundler: Vite's `?raw` import,
// esbuild's `--loader:.plugin.js=text`, or on Cloudflare Workers a wrangler
// rule `{ type = "Text", globs = ["**/*.plugin.js"] }`.
const AGENCY_PLUGIN = readFileSync(
  new URL('./agency.plugin.js', import.meta.url),
  'utf8'
);

// The executors your platform started, e.g. 'http://executor-1:8080,http://executor-2:8080'.
const pool = executorPool((process.env.EXECUTOR_URLS ?? '').split(','));

// One chat turn. `clientId` and `chatId` say who asks and in which chat;
// `designId` is the design on screen (an id your store holds, or a new
// crypto.randomUUID() for a chat that starts from nothing); `prompt` is what
// the person typed; `emit` sends one event to the browser.
export async function runTurn({
  clientId,
  chatId,
  designId,
  prompt,
  emit
}: {
  clientId: string;
  chatId: string;
  designId: string;
  prompt: string;
  emit: (event: TurnEvent) => void;
}): Promise<void> {
  // Every design a call opened (each holds an engine session on the executor,
  // released when the chat ends), and every design a call wrote this turn, to
  // render for the grid at the end.
  const opened = new Set<string>();
  const touched = new Set<string>();
  const followed = (e: ToolCallEvent) => {
    if (e.design === undefined) return;
    opened.add(e.design);
    if (e.wrote === undefined) return; // `wrote` is the version the call wrote, as your store returned it
    touched.add(e.design);
    emit({ type: 'design', design: e.design });
  };

  // Lets the agent hand people files: stores the rendered file and answers a link.
  async function uploadExport(req: UploadExportRequest) {
    const key = await uploadToStorage(clientId, req.bytes, req.mimeType);
    const downloadURL = await signedDownloadLink(clientId, key); // valid for a while
    emit({
      type: 'file',
      name: `${req.title}.${req.format}`,
      format: req.format,
      downloadURL
    });
    return { downloadURL }; // what the model tells the person
  }

  const codesign = prepareCoDesign({
    // Who the work is for: this chat, across all its messages. Every design it
    // touches opens in one executor that serves no other chat. The client id
    // in front keeps two clients' chat ids from ever naming the same executor.
    conversation: `${clientId}.${chatId}`,
    executors: remoteExecutors({
      // Asked with the conversation before each design's engine session opens,
      // on release, and when its executor was lost; answers this chat's
      // executor every time, never another's.
      acquire: ({ conversation }) => pool.acquire(conversation),
      token: process.env.EXECUTOR_TOKEN // the executors' EXECUTOR_TOKEN
    }),
    // The CE.SDK plugins every engine applies, in order. The sources they
    // register are everything the agent can find, place, or store into; you do
    // not list them, the executor reports them. Their order ranks search
    // results, the font lookup (the first source holding the exact family wins)
    // and which source is offered a new file first.
    plugins: [
      {
        module: AGENCY_PLUGIN,
        // This client's values for this turn.
        options: {
          api: `https://api.agency.example/clients/${clientId}`,
          token: tokenFor(clientId)
        }
      },
      // Not added for you. The font lookup finds Google families after the
      // brand's, and every Photoshop, InDesign, PDF, or PowerPoint import needs it.
      googleFonts()
    ],
    // Your design store, for this client: read and written here, in this
    // runner, where the model's code cannot reach it.
    designs: designStoreFor(clientId),
    // The design on screen: the model works on it without naming it.
    design: designId,
    // A plan without exports leaves it out, and the agent cannot export at
    // all: uploadExport would see a file only after it was rendered and charged.
    uploadExport: (await planHasExports(clientId)) ? uploadExport : undefined,
    // Tools of your own; they come after the SDK's.
    tools: [brandTool(clientId)],
    // Image generation, over IMG.LY's AI gateway with the credential IMG.LY gives
    // you; it stays in this process. `false` leaves it out.
    gateway: process.env.IMGLY_GATEWAY_API_KEY
      ? { bearer: async () => process.env.IMGLY_GATEWAY_API_KEY }
      : false,
    // Each tool call's event, after the call: which design, and what it wrote.
    events: [{ on: 'tool.call', run: followed }]
  });

  try {
    // Open the design on its engine before the model starts, so the first
    // tool call is fast; retries while the chat's executor starts, and tells
    // the person once. Saving as a scene is never charged.
    await waitForExecutor(
      () => codesign.export('scene'),
      () => emit({ type: 'status', text: 'waiting for an executor…' })
    );
    await runYourModel({
      tools: codesign.tools, // the SDK's, then yours
      instructions: codesign.instructions, // appended to your system prompt
      prompt,
      onProgress: emit
    });
    for (const id of touched) {
      await saveRender(id, await codesign.export('png', { design: id })); // a thumbnail for your grid
    }
  } finally {
    // Not released here: the chat's next message finds the engine still open.
    await rememberChatDesigns(chatId, [designId, ...opened]);
  }
}

// When the chat ends (the person closes it, or your product decides it is
// over): end the engine session of every design it opened. Releasing the last
// one retires the chat's executor about 30 seconds later, so never do this
// after every turn: the next message would wait for a fresh executor.
export async function endChat(clientId: string, chatId: string) {
  const codesign = prepareCoDesign({
    conversation: `${clientId}.${chatId}`,
    executors: remoteExecutors({
      acquire: ({ conversation }) => pool.acquire(conversation),
      token: process.env.EXECUTOR_TOKEN
    }),
    // The same module text as the turn's: an engine session is found by it
    // (the options may differ).
    plugins: [
      {
        module: AGENCY_PLUGIN,
        options: {
          api: `https://api.agency.example/clients/${clientId}`,
          token: tokenFor(clientId)
        }
      },
      googleFonts()
    ],
    designs: designStoreFor(clientId),
    gateway: false
  });
  await Promise.all(
    (await chatDesigns(chatId)).map((id) =>
      codesign.release({ design: id }).catch(() => undefined)
    )
  );
}

// The functions and clients this guide's code calls that are yours to write, over what your product already has.

/** Your database client, such as node-postgres's `Pool`. */
export const db = {
  async query(
    _sql: string,
    _params: readonly unknown[]
  ): Promise<{ rows: Record<string, any>[] }> {
    throw new Error('yours to write');
  }
};

/** Two calls over your object storage. */
export const bucket = {
  async get(_key: string): Promise<Uint8Array<ArrayBuffer>> {
    throw new Error('yours to write');
  },
  async put(_key: string, _bytes: Uint8Array): Promise<void> {
    throw new Error('yours to write');
  }
};

/** A short-lived token that opens only this client's files and routes, such as an HMAC of the client id and an expiry. */
export function tokenFor(_clientId: string): string {
  throw new Error('yours to write');
}

/** The client of the signed-in person, from your session cookie. */
export function clientOf(_request: Request): string {
  throw new Error('yours to write');
}

/** Stores a file for this client and returns its key. */
export async function uploadToStorage(
  _clientId: string,
  _bytes: Uint8Array,
  _mimeType: string
): Promise<string> {
  throw new Error('yours to write');
}

/** A download link to one stored file that expires after a while. */
export async function signedDownloadLink(
  _clientId: string,
  _key: string
): Promise<string> {
  throw new Error('yours to write');
}

/** The client's brand kit, from your API. */
export async function brandKitOf(_clientId: string): Promise<{
  typefaces: { name: string }[];
  colors: { name: string; hex: string }[];
  rules: string;
}> {
  throw new Error('yours to write');
}

/** Keeps a rendered PNG of a design as its thumbnail in your design grid. */
export async function saveRender(
  _designId: string,
  _png: Uint8Array
): Promise<void> {
  throw new Error('yours to write');
}

/** Whether the client's plan lets the agent hand people files. */
export async function planHasExports(_clientId: string): Promise<boolean> {
  throw new Error('yours to write');
}

/** Adds designs to your record of the ones a chat opened, for `endChat`. */
export async function rememberChatDesigns(
  _chatId: string,
  _designIds: readonly string[]
): Promise<void> {
  throw new Error('yours to write');
}

/** Every design your record says a chat opened. */
export async function chatDesigns(_chatId: string): Promise<string[]> {
  throw new Error('yours to write');
}

import type { DesignStore } from '@imgly/codesign-sdk';

import { bucket, db } from './yours.js';

// The table, one row per design:
//   designs(client_id, id, title, group_name, revision integer, content_key,
//           is_archive boolean, cesdk_version, edit_key,
//           primary key (client_id, id))
export function designStoreFor(clientId: string): DesignStore {
  return {
    // The source in your plugin module that lists this client's designs.
    listSourceId: 'agency.designs',

    async read(id) {
      const { rows } = await db.query(
        'SELECT * FROM designs WHERE client_id = $1 AND id = $2',
        [clientId, id] // the client in every query: another client's id finds nothing
      );
      const row = rows[0];
      if (row === undefined) return undefined;
      const bytes = await bucket.get(row.content_key);
      return {
        content: row.is_archive
          ? { archive: bytes }
          : { scene: new TextDecoder().decode(bytes) },
        version: String(row.revision), // changes on every write, never reused
        cesdkVersion: row.cesdk_version ?? null, // the CE.SDK that saved it, as `write` got it; null: unknown
        title: row.title,
        editKey: row.edit_key ?? undefined // the edit that wrote this version
      };
    },

    async write(id, content, { ifMatch, cesdkVersion, editKey, title, group }) {
      // Every version's bytes under a key of their own, so a row never points
      // at bytes another write is still uploading. A write that loses leaves
      // its bytes behind: let a bucket lifecycle rule remove unreferenced keys.
      const key = `designs/${clientId}/${id}/${crypto.randomUUID()}`;
      await bucket.put(
        key,
        content.archive ?? new TextEncoder().encode(content.scene)
      );
      const fields = [
        clientId,
        id,
        key,
        content.archive !== undefined,
        cesdkVersion,
        editKey ?? null // every write replaces the key; none clears it
      ];
      const { rows } =
        ifMatch === undefined
          ? // A new design, whose id the SDK just minted: create it only if
            // the id is free. `title` and `group` come with this first write.
            await db.query(
              `INSERT INTO designs (client_id, id, content_key, is_archive,
                 cesdk_version, edit_key, title, group_name, revision)
               VALUES ($1, $2, $3, $4, $5, $6, $7, $8, 1)
               ON CONFLICT DO NOTHING RETURNING revision`,
              [...fields, title ?? 'Untitled', group ?? null]
            )
          : // An edit: store it only while the row is still at the revision
            // the SDK read (compare-and-set).
            await db.query(
              `UPDATE designs SET content_key = $3, is_archive = $4,
                 cesdk_version = $5, edit_key = $6, revision = revision + 1
               WHERE client_id = $1 AND id = $2 AND revision = $7
               RETURNING revision`,
              [...fields, Number(ifMatch)]
            );
      // No row: someone wrote first, or the id is taken. Nothing was stored.
      return rows[0] === undefined ? undefined : String(rows[0].revision);
    }
  };
}

// Runs on the executor and in your editor. No imports: the executor loads
// this text as it is (bundle it into one file if it uses a library).
export default function agencyPlugin({ api, token }) {
  const headers = { authorization: `Bearer ${token}` };
  // undefined: 404, not this token's.
  const call = async (path, init) => {
    const response = await fetch(`${api}${path}`, {
      ...init,
      headers: { ...headers, ...(init?.headers ?? {}) }
    });
    if (response.status === 404) return undefined;
    if (!response.ok) throw new Error(`${path} answered ${response.status}`);
    return response;
  };
  const json = async (path) => (await call(path))?.json();
  // A file of this client as an asset. `uri` is what the design keeps for good.
  const asFile = (file) => ({
    id: file.id,
    label: file.name,
    meta: {
      uri: `agency-file://${file.id}`, // durable, never a signed link
      thumbUri: `agency-file://${file.thumbId}`,
      mimeType: file.mimeType,
      width: file.width,
      height: file.height
    }
  });
  // query.page counts from 0.
  const paged = (items, query) => {
    const start = query.page * query.perPage;
    const page = items.slice(start, start + query.perPage);
    return {
      assets: page,
      currentPage: query.page,
      ...(start + page.length < items.length
        ? { nextPage: query.page + 1 }
        : {}),
      total: items.length
    };
  };
  const matches = (query, text) =>
    !query.query || text.toLowerCase().includes(query.query.toLowerCase());

  return {
    name: 'agency.client',
    version: '1.0.0',
    initialize({ engine }) {
      // The client's designs, by title: your editor's "your designs" panel,
      // and the list the model searches for saved designs. It only lists: the
      // design store in your runner reads and writes the designs, so this
      // source holds no file types and the SDK never offers it a file.
      engine.asset.addSource({
        id: 'agency.designs',
        description: "This client's saved designs, by title.",
        getSupportedMimeTypes: () => [],
        async findAssets(query) {
          const params = new URLSearchParams({
            q: query.query ?? '',
            page: String(query.page),
            perPage: String(query.perPage)
          });
          const page = await json(`/designs?${params}`); // the query behind your design grid
          return {
            // One asset per design: `id` the design id, `label` its title.
            assets: page.items.map((row) => ({ id: row.id, label: row.title })),
            currentPage: query.page,
            ...(page.hasMore ? { nextPage: query.page + 1 } : {}),
            total: page.total
          };
        }
      });

      // Sources are asked in the order they are registered: the brand's typefaces come first.
      engine.asset.addSource({
        id: 'agency.fonts',
        description: "The brand's typefaces. Use these before any other font.",
        getSupportedMimeTypes: () => ['application/ubq-typeface'],
        async findAssets(query) {
          const families = (await json('/brand/typefaces')).filter((f) =>
            matches(query, f.name)
          );
          // One asset per family, every font in `fonts`, `name` exactly the family's.
          return paged(
            families.map((family) => ({
              id: family.id,
              label: family.name,
              payload: {
                typeface: {
                  name: family.name,
                  fonts: family.files.map((file) => ({
                    uri: file.url, // a stable https url, or a uri `resolve` answers
                    subFamily: file.subFamily,
                    weight: file.weight, // 'normal', 'bold', …
                    style: file.style // 'normal' or 'italic'
                  }))
                }
              }
            })),
            query
          );
        }
      });

      engine.asset.addSource({
        id: 'agency.colors',
        description: "The brand's colors.",
        getSupportedMimeTypes: () => ['application/ubq-color'],
        async findAssets(query) {
          const colors = (await json('/brand/colors')).filter((c) =>
            matches(query, c.name)
          );
          return paged(
            colors.map((c) => ({
              id: c.id,
              label: c.name,
              payload: { color: { colorSpace: 'sRGB', r: c.r, g: c.g, b: c.b } } // channels 0 to 1
            })),
            query
          );
        }
      });

      engine.asset.addSource({
        id: 'agency.photos',
        // The model reads this line in search results.
        description:
          "This client's logos and photos, including what people uploaded.",
        getSupportedMimeTypes: () => [
          'image/png',
          'image/jpeg',
          'image/webp',
          'image/gif',
          'image/svg+xml'
        ],
        // Live: a photo uploaded a second ago is found.
        async findAssets(query) {
          const files = (await json('/photos')).filter((f) =>
            matches(query, f.name)
          );
          return paged(files.map(asFile), query);
        },
        // null: no such photo for this client. This keeps another client's ids private.
        async fetchAsset(id) {
          const file = await json(`/photos/${encodeURIComponent(id)}`);
          return file === undefined ? null : asFile(file);
        },
        // agency-file://3f9a → a link the engine can fetch now; undefined: not this source's scheme.
        async resolve(uri) {
          if (!uri.startsWith('agency-file://')) return undefined;
          const id = encodeURIComponent(uri.slice('agency-file://'.length));
          return (await json(`/files/${id}/signed-url`)).url;
        },
        // Having this makes the source writable: it is offered every new
        // image (a generated one, one a person uploads in the editor, a
        // picture inside an imported file). Return its asset to keep it,
        // undefined to decline.
        async uploadFile(file, { action, label, width, height }) {
          // Pictures inside an imported file belong to that design only:
          // decline, and they stay inside it.
          if (action === 'import') return undefined;
          const query = new URLSearchParams({
            name: label ?? file.name,
            action
          });
          // The image's size in pixels, read by the SDK from the bytes. Your
          // route keeps it, and /photos answers it as `width` and `height`.
          if (width !== undefined) query.set('width', String(width));
          if (height !== undefined) query.set('height', String(height));
          const response = await call(`/photos?${query}`, {
            method: 'POST',
            headers: { 'content-type': file.type },
            body: file
          });
          if (response === undefined) throw new Error('/photos answered 404');
          return asFile(await response.json());
        }
      });

      engine.asset.addSource({
        id: 'agency.files',
        description:
          'Files the person attached to the chat, to open as designs.',
        getSupportedMimeTypes: () => [
          'application/pdf',
          'image/vnd.adobe.photoshop',
          'application/vnd.adobe.indesign-idml-package'
        ],
        async findAssets(query) {
          const files = (await json('/attachments')).filter((f) =>
            matches(query, f.name)
          );
          return paged(files.map(asFile), query);
        },
        async fetchAsset(id) {
          const file = await json(`/attachments/${encodeURIComponent(id)}`);
          return file === undefined ? null : asFile(file);
        }
      });
    }
  };
}

import type { AcquiredExecutor } from '@imgly/codesign-sdk';

// `urls` are the addresses of executors your platform started, as the runner reaches them.
export function executorPool(urls: readonly string[]) {
  const held = new Map<string, string>(); // conversation → executor url
  // Free: no conversation has bound it yet, and it accepts new sessions.
  const unbound = async (url: string): Promise<boolean> => {
    try {
      const health = (await (await fetch(`${url}/health`)).json()) as {
        result?: { bound?: boolean; accepting?: boolean };
      };
      return health.result?.bound === false && health.result.accepting === true;
    } catch {
      return false; // not answering: not a candidate now
    }
  };
  const pick = async (
    conversation: string
  ): Promise<AcquiredExecutor | undefined> => {
    const mine = held.get(conversation);
    if (mine !== undefined) return { url: mine }; // the chat's own, every time
    const taken = new Set(held.values());
    for (const url of urls) {
      if (taken.has(url) || !(await unbound(url))) continue;
      held.set(conversation, url);
      return { url };
    }
    return undefined; // every executor serves another chat: the call fails as unreachable
  };
  // One pick at a time, so two chats never claim the same executor between
  // its health check and the claim.
  let queue: Promise<unknown> = Promise.resolve();
  return {
    acquire(conversation: string): Promise<AcquiredExecutor | undefined> {
      const picked = queue.then(() => pick(conversation));
      queue = picked.catch(() => undefined);
      return picked;
    },
    // Once the chat is over and its container replaced, its address is free again.
    forget(conversation: string): void {
      held.delete(conversation);
    }
  };
}