fix(agents): bound cold avatar admission and retain company drafts

Co-Authored-By: Paperclip <noreply@paperclip.ing>

server/src/__tests__/agent-avatars.test.ts

@@ -8,7 +8,7 @@ import path from "node:path";
 import express from "express";
 import type { Server } from "node:http";
 import sharp from "sharp";
-import { appearanceForPalette } from "@paperclipai/shared";
+import { AGENT_PALETTE_IDS, appearanceForPalette } from "@paperclipai/shared";
 import { createLocalDiskStorageProvider } from "../storage/local-disk-provider.js";
 import { createAgentAvatarService, avatarCacheKey, type AgentAvatarRequest } from "../services/agent-avatars.js";
 import { createAgentAvatarPool } from "../services/agent-avatar-pool.js";
@@ -43,6 +43,39 @@ describe("on-demand agent avatars", () => {
     (await service.get(request)).stream.destroy();
     expect(render).toHaveBeenCalledTimes(2);
   });
+  it("limits cold keys per client while admitting warm hits, joiners and other clients", async () => {
+    let release!: () => void;
+    const blocked = new Promise<void>(resolve => { release = resolve; });
+    const render = vi.fn(async () => { await blocked; return Buffer.from("png"); });
+    const provider = await storage();
+    const service = createAgentAvatarService(provider, render);
+    const keys = AGENT_PALETTE_IDS.flatMap(palette => ([16, 20, 24] as const).map(size => ({ ...request, appearance: appearanceForPalette(palette), size })));
+    const pending = keys.slice(0, 32).map(key => service.get(key, "one"));
+    try {
+      await vi.waitFor(() => expect(render).toHaveBeenCalledTimes(32));
+      await expect(service.get(keys[32], "one")).rejects.toThrow("Too many cold avatar requests");
+      pending.push(service.get(keys[0], "one")); // Same cold key is free.
+      pending.push(service.get(keys[32], "two")); // Another client still has room.
+      await vi.waitFor(() => expect(render).toHaveBeenCalledTimes(33));
+    } finally { release(); }
+    for (const result of await Promise.all(pending)) result.stream.destroy();
+    (await service.get(keys[0], "one")).stream.destroy();
+    expect(render).toHaveBeenCalledTimes(33);
+    (await service.get(keys[33], "one")).stream.destroy(); // Completed renders release slots.
+    expect(render).toHaveBeenCalledTimes(34);
+  });
+  it("returns retryable admission errors without caching them", async () => {
+    const service = createAgentAvatarService(await storage(), async () => Buffer.from("png"));
+    const { AvatarAdmissionError } = await import("../services/agent-avatars.js");
+    vi.spyOn(service, "get").mockRejectedValueOnce(new AvatarAdmissionError(12));
+    const url = await serve(service);
+    const denied = await fetch(url);
+    expect(denied.status).toBe(429);
+    expect(denied.headers.get("retry-after")).toBe("12");
+    expect(denied.headers.get("cache-control")).toBe("no-store");
+    await denied.text();
+    const retry = await fetch(url); expect(retry.status).toBe(200); await retry.arrayBuffer();
+  });
   it("uses the configured S3 prefix and reuses bytes across service instances", async () => {
     const objects = new Map<string, Buffer>();
     const send = vi.spyOn(S3Client.prototype, "send").mockImplementation(async (command: any) => {

server/src/routes/agent-avatars.ts

@@ -2,7 +2,7 @@ import { pipeline } from "node:stream/promises";
 import { Router } from "express";
 import { z } from "zod";
 import { AGENT_PALETTE_IDS, AGENT_AVATAR_SIZES, CHARACTER_STATES, appearanceForPalette, type AgentAvatarSize } from "@paperclipai/shared";
-import { createAgentAvatarService } from "../services/agent-avatars.js";
+import { AvatarAdmissionError, createAgentAvatarService } from "../services/agent-avatars.js";
 import { createStorageProviderFromConfig } from "../storage/provider-registry.js";
 import { loadConfig } from "../config.js";
 import { logger } from "../middleware/logger.js";
@@ -29,7 +29,7 @@ export function agentAvatarRoutes(injected?: ReturnType<typeof createAgentAvatar
     const { palette, pose, size, scale } = parsed.data;
     try {
       service ??= createAgentAvatarService(createStorageProviderFromConfig(loadConfig()));
-      const { stream, byteSize, etag } = await service.get({ appearance: appearanceForPalette(palette === "muted-dream" ? AGENT_PALETTE_IDS[0] : palette), muted: palette === "muted-dream", pose, size: size as AgentAvatarSize, scale: Number(scale) as 1 | 2 });
+      const { stream, byteSize, etag } = await service.get({ appearance: appearanceForPalette(palette === "muted-dream" ? AGENT_PALETTE_IDS[0] : palette), muted: palette === "muted-dream", pose, size: size as AgentAvatarSize, scale: Number(scale) as 1 | 2 }, req.ip || req.socket.remoteAddress || "unknown");
       res.setHeader("Cache-Control", "public, max-age=31536000, immutable");
       res.setHeader("ETag", etag);
       res.setHeader("X-Content-Type-Options", "nosniff");
@@ -39,13 +39,13 @@ export function agentAvatarRoutes(injected?: ReturnType<typeof createAgentAvatar
       res.setHeader("Content-Length", byteSize);
       await pipeline(stream, res);
     } catch (error) {
-      logger.warn({ err: error }, "Could not render agent avatar");
+      if (!(error instanceof AvatarAdmissionError)) logger.warn({ err: error }, "Could not render agent avatar");
       if (res.headersSent || res.destroyed) { res.destroy(); return; }
       res.removeHeader("Content-Length");
       res.removeHeader("ETag");
       res.setHeader("Cache-Control", "no-store");
-      res.setHeader("Retry-After", "5");
-      res.status(503).json({ error: "Avatar temporarily unavailable" });
+      res.setHeader("Retry-After", String(error instanceof AvatarAdmissionError ? error.retryAfterSeconds : 5));
+      res.status(error instanceof AvatarAdmissionError ? 429 : 503).json({ error: "Avatar temporarily unavailable" });
     }
   });
   return router;

server/src/routes/openapi.ts

@@ -2544,6 +2544,7 @@ registry.registerPath({
     200: { description: "PNG portrait; Cache-Control: public, max-age=31536000, immutable; ETag: SHA-256 of PNG bytes", content: { "image/png": { schema: { type: "string", format: "binary" } } } },
     304: { description: "If-None-Match matches the cached content ETag" },
     400: r.badRequest,
+    429: { description: "Cold-render admission limit for this client; Cache-Control: no-store; Retry-After in seconds. Cached portraits remain available." },
     503: { description: "Retryable rendering/storage failure; Cache-Control: no-store; Retry-After: 5" },
   },
 });

server/src/services/agent-avatar-pool.ts

@@ -11,7 +11,7 @@ export function createAgentAvatarPool(concurrency = 2, maxQueue = 64) {
   const workers = new Set<Worker>();
   let closed = false;
   function spawn() {
-    const source = import.meta.url.endsWith(".ts");
+    const source = new URL(import.meta.url).pathname.endsWith(".ts");
     const url = new URL(source ? "./agent-avatar-worker.ts" : "./agent-avatar-worker.js", import.meta.url);
     const worker = source
       ? new Worker(`import(${JSON.stringify(import.meta.resolve('tsx/esm/api'))}).then(({tsImport}) => tsImport(${JSON.stringify(url.href)}, ${JSON.stringify(import.meta.url)}));`, { eval: true })

server/src/services/agent-avatars.ts

@@ -1,6 +1,7 @@
 import { createHash } from "node:crypto";
 import type { AgentAppearance, AgentAvatarSize, CharacterState } from "@paperclipai/shared";
 import type { StorageProvider } from "../storage/types.js";
+import { createInviteRateLimiter } from "./invite-rate-limit.js";
 import { createAgentAvatarPool } from "./agent-avatar-pool.js";
 
 export interface AgentAvatarRequest {
@@ -14,11 +15,18 @@ export function avatarCacheKey(request: AgentAvatarRequest) {
   const { appearance, size, scale, pose, muted } = request;
   return `generated-agent-avatars/${appearance.characterVersion}/${muted ? "muted-dream" : appearance.paletteId}/${pose}-${size}-${scale}.png`;
 }
+export class AvatarAdmissionError extends Error {
+  constructor(public readonly retryAfterSeconds: number) { super("Too many cold avatar requests"); }
+}
 type CacheMetadata = { sha256: string; byteSize: number };
 export function createAgentAvatarService(storage: StorageProvider, render?: (request: AgentAvatarRequest) => Promise<Buffer>) {
   const pool = render ? undefined : createAgentAvatarPool();
   const pending = new Map<string, Promise<CacheMetadata>>();
-  async function ensure(request: AgentAvatarRequest, key: string): Promise<CacheMetadata> {
+  // Only cold keys consume admission; warm images and single-flight joiners
+  // remain available. Leave at least half the pool queue for other clients.
+  const activeByClient = new Map<string, number>();
+  const limiter = createInviteRateLimiter({ maxRequests: 256 });
+  async function ensure(request: AgentAvatarRequest, key: string, client: string): Promise<CacheMetadata> {
     const metadataKey = `${key}.json`;
     const [image, metadata] = await Promise.all([
       storage.headObject({ objectKey: key }), storage.headObject({ objectKey: metadataKey }),
@@ -32,21 +40,32 @@ export function createAgentAvatarService(storage: StorageProvider, render?: (req
         if (/^[a-f0-9]{64}$/.test(cached.sha256) && cached.byteSize > 0 && cached.byteSize === image.contentLength) return cached;
       } catch { /* Disposable metadata: regenerate a corrupt or old cache entry. */ }
     }
-    const bytes = await (render ?? pool!.render)(request);
-    const result = { sha256: createHash("sha256").update(bytes).digest("hex"), byteSize: bytes.length };
-    // Both providers publish whole objects atomically. Publish metadata last so
-    // readers never consider an unfinished image a completed cache entry.
-    await storage.putObject({ objectKey: key, body: bytes, contentLength: bytes.length, contentType: "image/png" });
-    const encoded = Buffer.from(JSON.stringify(result));
-    await storage.putObject({ objectKey: metadataKey, body: encoded, contentLength: encoded.length, contentType: "application/json" });
-    return result;
+    const active = activeByClient.get(client) ?? 0;
+    if (active >= 32) throw new AvatarAdmissionError(5);
+    const admission = limiter.consume(client);
+    if (!admission.allowed) throw new AvatarAdmissionError(admission.retryAfterSeconds);
+    activeByClient.set(client, active + 1);
+    try {
+      const bytes = await (render ?? pool!.render)(request);
+      const result = { sha256: createHash("sha256").update(bytes).digest("hex"), byteSize: bytes.length };
+      // Both providers publish whole objects atomically. Publish metadata last so
+      // readers never consider an unfinished image a completed cache entry.
+      await storage.putObject({ objectKey: key, body: bytes, contentLength: bytes.length, contentType: "image/png" });
+      const encoded = Buffer.from(JSON.stringify(result));
+      await storage.putObject({ objectKey: metadataKey, body: encoded, contentLength: encoded.length, contentType: "application/json" });
+      return result;
+    } finally {
+      const remaining = (activeByClient.get(client) ?? 1) - 1;
+      if (remaining) activeByClient.set(client, remaining);
+      else activeByClient.delete(client);
+    }
   }
   return {
-    async get(request: AgentAvatarRequest) {
+    async get(request: AgentAvatarRequest, client = "unknown") {
       const key = avatarCacheKey(request);
       let result = pending.get(key);
       if (!result) {
-        result = ensure(request, key).finally(() => pending.delete(key));
+        result = ensure(request, key, client).finally(() => pending.delete(key));
         pending.set(key, result);
       }
       const metadata = await result;

ui/src/components/AgentPersona.test.tsx

@@ -63,6 +63,17 @@ describe("agent persona presentation", () => {
     expect(createCharacter).not.toHaveBeenCalled();
     expect(host.querySelectorAll("img")).toHaveLength(2);
   });
+  it("selects the correct saved draft when the company changes without remounting", async () => {
+    let draft!: ReturnType<typeof useAgentAppearanceDraft>;
+    function Draft({ company }: { company: string }) { draft = useAgentAppearanceDraft(`${company}:new-agent`); return null; }
+    sessionStorage.setItem("paperclip.agent-appearance.two:new-agent", JSON.stringify(appearance));
+    await act(async () => root.render(<Draft company="one" />));
+    const first = draft.appearance;
+    await act(async () => root.render(<Draft company="two" />));
+    expect(draft.appearance).toEqual(appearance);
+    await act(async () => root.render(<Draft company="one" />));
+    expect(draft.appearance).toEqual(first);
+  });
   it("retains the draft assignment across remounts and clears it only after creation", async () => {
     let draft!: ReturnType<typeof useAgentAppearanceDraft>;
     function Draft() { draft = useAgentAppearanceDraft("company:new-agent"); return null; }

ui/src/hooks/useAgentAppearanceDraft.ts

@@ -1,10 +1,10 @@
 import { useState } from "react";
 import { agentAppearanceSchema, randomAgentAppearance } from "@paperclipai/shared";
 
-/** Non-secret visual identity only. The caller remounts when its draft key changes. */
+/** Non-secret visual identity, retained across navigation and company changes. */
 export function useAgentAppearanceDraft(draftKey: string) {
   const key = `paperclip.agent-appearance.${draftKey}`;
-  const [appearance] = useState(() => {
+  function readDraft() {
     try {
       const stored = agentAppearanceSchema.safeParse(JSON.parse(sessionStorage.getItem(key) ?? "null"));
       if (stored.success) return stored.data;
@@ -12,6 +12,8 @@ export function useAgentAppearanceDraft(draftKey: string) {
     const value = randomAgentAppearance();
     try { sessionStorage.setItem(key, JSON.stringify(value)); } catch { /* In-memory draft still works. */ }
     return value;
-  });
-  return { appearance, clear() { try { sessionStorage.removeItem(key); } catch { /* Best effort. */ } } };
+  }
+  const [draft, setDraft] = useState(() => ({ key, appearance: readDraft() }));
+  if (draft.key !== key) setDraft({ key, appearance: readDraft() });
+  return { appearance: draft.appearance, clear() { try { sessionStorage.removeItem(key); } catch { /* Best effort. */ } } };
 }