Add session affinity using cookie store

Author: federicociner

acp-web-transport/design.md

@@ -0,0 +1,117 @@
+import { describe, expect, it } from "vitest";
+
+import { MemoryAcpCookieStore } from "./cookie-store.js";
+
+describe("MemoryAcpCookieStore", () => {
+  it("stores single and multiple Set-Cookie values", () => {
+    const store = new MemoryAcpCookieStore();
+    store.store(headersWithSetCookie(["transport=alpha; Path=/"]));
+    store.store(
+      headersWithSetCookie(["route=bravo; Path=/", "affinity=charlie"]),
+    );
+
+    const headers = new Headers();
+    store.apply(headers);
+
+    expect(headers.get("Cookie")).toBe(
+      "transport=alpha; route=bravo; affinity=charlie",
+    );
+  });
+
+  it("splits combined Set-Cookie headers with Expires commas", () => {
+    const store = new MemoryAcpCookieStore();
+    store.store(
+      new Headers({
+        "Set-Cookie":
+          "transport=alpha; Expires=Wed, 21 Oct 2030 07:28:00 GMT, route=bravo; Path=/",
+      }),
+    );
+
+    const headers = new Headers();
+    store.apply(headers);
+
+    expect(headers.get("Cookie")).toBe("transport=alpha; route=bravo");
+  });
+
+  it("ignores malformed cookie headers", () => {
+    const store = new MemoryAcpCookieStore();
+    store.store(
+      headersWithSetCookie([
+        "missing-separator",
+        "=empty-name",
+        " =blank",
+        "ok=value",
+      ]),
+    );
+
+    const headers = new Headers();
+    store.apply(headers);
+
+    expect(headers.get("Cookie")).toBe("ok=value");
+  });
+
+  it("lets later cookies overwrite earlier cookies with the same name", () => {
+    const store = new MemoryAcpCookieStore();
+    store.store(headersWithSetCookie(["route=alpha", "route=bravo"]));
+
+    const headers = new Headers();
+    store.apply(headers);
+
+    expect(headers.get("Cookie")).toBe("route=bravo");
+  });
+
+  it("writes a Cookie header when managed cookies exist", () => {
+    const store = new MemoryAcpCookieStore();
+    store.store(headersWithSetCookie(["transport=alpha"]));
+
+    const headers = new Headers();
+    store.apply(headers);
+
+    expect(headers.get("Cookie")).toBe("transport=alpha");
+  });
+
+  it("merges managed cookies with caller-provided Cookie headers", () => {
+    const store = new MemoryAcpCookieStore();
+    store.store(headersWithSetCookie(["transport=alpha", "route=bravo"]));
+
+    const headers = new Headers({ Cookie: "caller=custom" });
+    store.apply(headers);
+
+    expect(headers.get("Cookie")).toBe(
+      "transport=alpha; route=bravo; caller=custom",
+    );
+  });
+
+  it("lets caller-provided cookie values override managed duplicate names", () => {
+    const store = new MemoryAcpCookieStore();
+    store.store(headersWithSetCookie(["transport=alpha", "route=bravo"]));
+
+    const headers = new Headers({ Cookie: "route=caller; caller=custom" });
+    store.apply(headers);
+
+    expect(headers.get("Cookie")).toBe(
+      "transport=alpha; route=caller; caller=custom",
+    );
+  });
+
+  it("clears managed cookies", () => {
+    const store = new MemoryAcpCookieStore();
+    store.store(headersWithSetCookie(["transport=alpha"]));
+    store.clear();
+
+    const headers = new Headers();
+    store.apply(headers);
+
+    expect(headers.get("Cookie")).toBeNull();
+  });
+});
+
+function headersWithSetCookie(values: readonly string[]): Headers {
+  const headers = new Headers();
+
+  Object.defineProperty(headers, "getSetCookie", {
+    value: () => values,
+  });
+
+  return headers;
+}

acp-web-transport/rfd.md

@@ -0,0 +1,147 @@
+/**
+ * Minimal ACP affinity cookie store.
+ *
+ * This helper stores cookie name/value pairs from `Set-Cookie` response
+ * headers and applies them to outgoing `Cookie` request headers. It is meant
+ * for ACP routing affinity across reconnects, not authentication or
+ * authorization, and not as a general-purpose browser cookie jar: it
+ * intentionally does not implement domain/path matching,
+ * expiry, `Secure`, `HttpOnly`, or `SameSite` handling.
+ */
+export interface AcpCookieStore {
+  /** Stores cookies from response headers. */
+  store(headers: Headers): void;
+  /** Applies stored cookies to outgoing request headers. */
+  apply(headers: Headers): void;
+  /** Clears all stored cookies. */
+  clear(): void;
+}
+
+/** In-memory implementation of {@link AcpCookieStore}. */
+export class MemoryAcpCookieStore implements AcpCookieStore {
+  private readonly cookies = new Map<string, string>();
+
+  store(headers: Headers): void {
+    for (const value of setCookieHeaders(headers)) {
+      const cookie = parseSetCookie(value);
+      if (!cookie) {
+        continue;
+      }
+
+      this.cookies.set(cookie.name, cookie.value);
+    }
+  }
+
+  apply(headers: Headers): void {
+    const merged = mergeCookieHeaders(
+      this.cookieHeader(),
+      headers.get("Cookie"),
+    );
+    if (merged) {
+      headers.set("Cookie", merged);
+    }
+  }
+
+  clear(): void {
+    this.cookies.clear();
+  }
+
+  private cookieHeader(): string | undefined {
+    return this.cookies.size === 0
+      ? undefined
+      : Array.from(this.cookies)
+          .map(([name, value]) => `${name}=${value}`)
+          .join("; ");
+  }
+}
+
+interface CookiePair {
+  readonly name: string;
+  readonly value: string;
+}
+
+function setCookieHeaders(headers: Headers): string[] {
+  const getSetCookie = headers.getSetCookie;
+  if (typeof getSetCookie === "function") {
+    return getSetCookie.call(headers).flatMap(splitSetCookieHeader);
+  }
+
+  const setCookie = headers.get("Set-Cookie");
+  return setCookie ? splitSetCookieHeader(setCookie) : [];
+}
+
+function splitSetCookieHeader(header: string): string[] {
+  return header
+    .split(/,(?=\s*[^;,\s]+=)/)
+    .map((value) => value.trim())
+    .filter((value) => value.length > 0);
+}
+
+function parseSetCookie(header: string): CookiePair | undefined {
+  const pair = header.split(";", 1)[0];
+  const separator = pair.indexOf("=");
+
+  if (separator <= 0) {
+    return undefined;
+  }
+
+  const name = pair.slice(0, separator).trim();
+  if (!name) {
+    return undefined;
+  }
+
+  return {
+    name,
+    value: pair.slice(separator + 1).trim(),
+  };
+}
+
+function mergeCookieHeaders(
+  managedCookieHeader: string | undefined,
+  callerCookieHeader: string | null,
+): string | undefined {
+  const cookies = new Map<string, string>();
+
+  for (const cookie of parseCookieHeader(managedCookieHeader)) {
+    cookies.set(cookie.name, cookie.value);
+  }
+
+  for (const cookie of parseCookieHeader(callerCookieHeader ?? undefined)) {
+    cookies.set(cookie.name, cookie.value);
+  }
+
+  return cookies.size === 0
+    ? undefined
+    : Array.from(cookies)
+        .map(([name, value]) => `${name}=${value}`)
+        .join("; ");
+}
+
+function parseCookieHeader(header: string | undefined): CookiePair[] {
+  if (!header) {
+    return [];
+  }
+
+  return header
+    .split(";")
+    .map(parseCookiePair)
+    .filter((cookie): cookie is CookiePair => cookie !== undefined);
+}
+
+function parseCookiePair(value: string): CookiePair | undefined {
+  const separator = value.indexOf("=");
+
+  if (separator <= 0) {
+    return undefined;
+  }
+
+  const name = value.slice(0, separator).trim();
+  if (!name) {
+    return undefined;
+  }
+
+  return {
+    name,
+    value: value.slice(separator + 1).trim(),
+  };
+}

acp-web-transport/tasks.md

@@ -1,7 +1,10 @@
 #!/usr/bin/env node
 
 import * as acp from "@agentclientprotocol/sdk";
-import { createHttpStream } from "@agentclientprotocol/sdk/http-client";
+import {
+  MemoryAcpCookieStore,
+  createHttpStream,
+} from "@agentclientprotocol/sdk/http-client";
 
 class HttpExampleClient implements acp.Client {
   async requestPermission(
@@ -30,19 +33,36 @@ class HttpExampleClient implements acp.Client {
 }
 
 const serverUrl = process.env.ACP_HTTP_URL ?? "http://127.0.0.1:7331/acp";
-const stream = createHttpStream(serverUrl, {
-  headers: {
-    Authorization: "Bearer example-token",
-  },
-  // Cookies are included by default and scoped to this stream. Use `cookies: "omit"` for stateless requests.
-});
-const connection = new acp.ClientSideConnection(
-  (_agent) => new HttpExampleClient(),
-  stream,
-);
+const authHeaders = {
+  Authorization: "Bearer example-token",
+};
+
+// Keep reconnect state outside individual stream instances. Reuse this store
+// across fresh streams so an external affinity layer can route reconnects.
+const cookieStore = new MemoryAcpCookieStore();
+let savedSessionId: string | undefined;
+
+function connect(): {
+  readonly stream: acp.Stream;
+  readonly connection: acp.ClientSideConnection;
+} {
+  const stream = createHttpStream(serverUrl, {
+    headers: authHeaders,
+    cookieStore,
+    // Cookies are included by default. Use `cookies: "omit"` for stateless requests.
+  });
+  const connection = new acp.ClientSideConnection(
+    (_agent) => new HttpExampleClient(),
+    stream,
+  );
+
+  return { stream, connection };
+}
+
+const { stream, connection } = connect();
 
 try {
-  await connection.initialize({
+  const initialized = await connection.initialize({
     protocolVersion: acp.PROTOCOL_VERSION,
     clientCapabilities: {},
   });
@@ -51,6 +71,7 @@ try {
     cwd: process.cwd(),
     mcpServers: [],
   });
+  savedSessionId = session.sessionId;
 
   const result = await connection.prompt({
     sessionId: session.sessionId,
@@ -63,6 +84,22 @@ try {
   });
 
   console.log(`\nDone: ${result.stopReason}`);
+
+  console.log(
+    `Saved session ${savedSessionId}; loadSession=${initialized.agentCapabilities?.loadSession === true}`,
+  );
+
+  // Reconnect flow sketch:
+  // 1. Save `sessionId`, auth headers, cwd, MCP servers, and `cookieStore`.
+  // 2. Create a fresh stream with the same auth headers and cookie store.
+  // 3. Call initialize and require `agentCapabilities.loadSession`.
+  // 4. Call session/load for the saved session ID.
+  // Production agents must authorize session/load for the authenticated user.
+  // ACP v1 does not replay in-flight transport messages emitted while disconnected.
+  // Example:
+  // const next = connect();
+  // await next.connection.initialize({ protocolVersion: acp.PROTOCOL_VERSION, clientCapabilities: {} });
+  // await next.connection.loadSession({ sessionId: savedSessionId, cwd: process.cwd(), mcpServers: [] });
 } finally {
   await stream.writable.close();
 }