@illuma-ai/agents 1.4.0-alpha.6 → 1.5.1

package/src/agents/__tests__/AgentContext.crossUserCache.live.test.ts

@@ -0,0 +1,342 @@
+/* eslint-disable no-console */
+/**
+ * LIVE Bedrock prompt-cache verification: cross-user, account-level.
+ *
+ * What this test proves:
+ *
+ *   1. The `cachePoint` block emitted inline by `buildSystemRunnable` for
+ *      Bedrock actually causes Bedrock to write a cache entry.
+ *   2. The cache entry is keyed off the STABLE prefix only — two simulated
+ *      users (Alice, Bob) with the SAME `agent.instructions` but DIFFERENT
+ *      `additional_instructions` both hit the same cache entry.
+ *   3. Bedrock's prompt cache is account-scoped: a second AgentContext (a
+ *      separate "user") in the SAME AWS account reads the cache that the
+ *      first user populated.
+ *
+ * What it does NOT prove (out of scope for a unit-runner):
+ *   - True concurrent multi-tenant load (this runs sequentially)
+ *   - Cache TTL / eviction behavior
+ *   - Behavior across AWS accounts (cache is account-scoped by design)
+ *
+ * Run with:
+ *   BEDROCK_AWS_REGION=us-east-1 \
+ *     BEDROCK_AWS_ACCESS_KEY_ID=... \
+ *     BEDROCK_AWS_SECRET_ACCESS_KEY=... \
+ *     npx jest src/agents/__tests__/AgentContext.crossUserCache.live.test.ts
+ *
+ * Skipped automatically if Bedrock credentials are not present.
+ */
+
+import { config } from 'dotenv';
+config();
+
+import { HumanMessage, SystemMessage } from '@langchain/core/messages';
+import type { BaseMessage, AIMessage } from '@langchain/core/messages';
+import { IllumaBedrockConverse } from '@/llm/bedrock';
+import { addBedrockCacheControl } from '@/messages/cache';
+import { AgentContext } from '../AgentContext';
+import { Providers } from '@/common';
+import type * as t from '@/types';
+
+const region =
+  process.env.BEDROCK_AWS_REGION ??
+  process.env.BEDROCK_AWS_DEFAULT_REGION ??
+  process.env.AWS_REGION ??
+  'us-east-1';
+const accessKeyId =
+  process.env.BEDROCK_AWS_ACCESS_KEY_ID ?? process.env.AWS_ACCESS_KEY_ID;
+const secretAccessKey =
+  process.env.BEDROCK_AWS_SECRET_ACCESS_KEY ??
+  process.env.AWS_SECRET_ACCESS_KEY;
+
+const haveBedrock =
+  accessKeyId !== undefined &&
+  accessKeyId !== '' &&
+  secretAccessKey !== undefined &&
+  secretAccessKey !== '';
+
+/**
+ * Build a stable system prompt LARGER than the Bedrock cache write minimum.
+ * Anthropic Claude on Bedrock generally needs ≥ 1024 tokens before cache
+ * writes are honoured; we pad to ~2k tokens of stable text to be safe.
+ */
+const buildLargeStableInstructions = (): string => {
+  const header =
+    'You are a helpful assistant working as part of a multi-user platform.';
+  const padding = Array.from({ length: 80 }, (_, i) =>
+    `Rule ${i + 1}: When responding, prefer concise and accurate answers. Cite assumptions explicitly. Treat all users with respect. Do not hallucinate facts. If unsure, say so.`
+  ).join('\n');
+  return `${header}\n\n${padding}`;
+};
+
+const STABLE_INSTRUCTIONS = buildLargeStableInstructions();
+
+const describeFn = haveBedrock ? describe : describe.skip;
+
+/**
+ * Run one Bedrock call through the real AgentContext + IllumaBedrockConverse
+ * pipeline and return the raw usage_metadata so we can read
+ * `input_token_details.cache_read` / `cache_creation`.
+ */
+async function runOneTurn({
+  agentInstructions,
+  additionalInstructions,
+  userQuery,
+  agentId,
+}: {
+  agentInstructions: string;
+  additionalInstructions: string;
+  userQuery: string;
+  agentId: string;
+}): Promise<{
+  context: AgentContext;
+  usage: {
+    input_tokens?: number;
+    output_tokens?: number;
+    cache_read_input_tokens?: number;
+    cache_creation_input_tokens?: number;
+  };
+  systemPrefixSnapshot: string;
+}> {
+  const ctx = AgentContext.fromConfig({
+    agentId,
+    provider: Providers.BEDROCK,
+    instructions: agentInstructions,
+    additional_instructions: additionalInstructions,
+    clientOptions: {
+      promptCache: true,
+      model: 'us.anthropic.claude-sonnet-4-5-20250929-v1:0',
+    } as t.BedrockAnthropicClientOptions,
+  });
+
+  const runnable = ctx.systemRunnable;
+  expect(runnable).toBeDefined();
+
+  /* Capture the cacheable prefix as a deterministic string snapshot so we
+   * can compare two contexts for byte-equality (cache stability rule). */
+  const systemArr = (await runnable!.invoke([])) as BaseMessage[];
+  const sysContent = systemArr[0].content as Array<Record<string, unknown>>;
+  /* Anything up to and including the cachePoint block forms the cache
+   * prefix — we serialize that slice for comparison. */
+  const cpIdx = sysContent.findIndex((b) => 'cachePoint' in b);
+  const prefixForCacheKey = JSON.stringify(
+    cpIdx >= 0 ? sysContent.slice(0, cpIdx + 1) : sysContent
+  );
+
+  const userMessages: BaseMessage[] = [new HumanMessage(userQuery)];
+  const allMessages = [...systemArr, ...userMessages];
+  const finalMessages = addBedrockCacheControl<BaseMessage>(allMessages);
+
+  /* System message must come through with cachePoint intact. Verify
+   * defensively before paying for the API call. */
+  const sys = finalMessages[0];
+  const sysOut = sys.content as Array<Record<string, unknown>>;
+  expect(sysOut.some((b) => 'cachePoint' in b)).toBe(true);
+
+  const llm = new IllumaBedrockConverse({
+    model: 'us.anthropic.claude-sonnet-4-5-20250929-v1:0',
+    region,
+    credentials: {
+      accessKeyId: accessKeyId as string,
+      secretAccessKey: secretAccessKey as string,
+    },
+    promptCache: true,
+    streaming: false,
+  });
+
+  /* Use the non-streaming path so we can read usage metadata directly. */
+  const aiMessage = (await llm.invoke(finalMessages)) as AIMessage;
+
+  type UsageMeta = {
+    input_tokens?: number;
+    output_tokens?: number;
+    cache_read_input_tokens?: number;
+    cache_creation_input_tokens?: number;
+    input_token_details?: {
+      cache_read?: number;
+      cache_creation?: number;
+    };
+  };
+  type RespMeta = {
+    usage?: {
+      inputTokens?: number;
+      outputTokens?: number;
+      cacheReadInputTokens?: number;
+      cacheWriteInputTokens?: number;
+    };
+  };
+  const u = (aiMessage.usage_metadata ?? {}) as UsageMeta;
+  const rmUsage = ((aiMessage.response_metadata ?? {}) as RespMeta).usage ?? {};
+
+  /* LangChain Bedrock currently only surfaces basic input/output tokens in
+   * usage_metadata. Cache fields are in response_metadata.usage as
+   * cacheReadInputTokens / cacheWriteInputTokens. Read from both. */
+  const cache_read_input_tokens =
+    u.cache_read_input_tokens ??
+    u.input_token_details?.cache_read ??
+    rmUsage.cacheReadInputTokens ??
+    0;
+  const cache_creation_input_tokens =
+    u.cache_creation_input_tokens ??
+    u.input_token_details?.cache_creation ??
+    rmUsage.cacheWriteInputTokens ??
+    0;
+
+  return {
+    context: ctx,
+    usage: {
+      input_tokens: u.input_tokens,
+      output_tokens: u.output_tokens,
+      cache_read_input_tokens,
+      cache_creation_input_tokens,
+    },
+    systemPrefixSnapshot: prefixForCacheKey,
+  };
+}
+
+describeFn('LIVE: Bedrock cross-user prompt cache (account-scoped)', () => {
+  jest.setTimeout(120000);
+
+  test('two distinct users with same agent + different memory hit the SAME cache entry', async () => {
+    /* Step 1 — user "Alice" makes the first call. Cache is cold so this
+     * call writes the cache entry and pays cache_creation_input_tokens. */
+    const alice = await runOneTurn({
+      agentId: 'shared-agent',
+      agentInstructions: STABLE_INSTRUCTIONS,
+      additionalInstructions: 'User: Alice. Memories: meetings on Mondays.',
+      userQuery: 'Say hi in five words.',
+    });
+
+    console.log('Alice usage:', alice.usage);
+
+    /* Step 2 — user "Bob" makes the same agent's first call with DIFFERENT
+     * dynamic context (different additional_instructions). Bob's stable
+     * prefix is byte-identical to Alice's, so Bedrock should serve a cache
+     * read for the system prefix even though Bob has never called before. */
+    const bob = await runOneTurn({
+      agentId: 'shared-agent',
+      agentInstructions: STABLE_INSTRUCTIONS,
+      additionalInstructions: 'User: Bob. Memories: lives in Berlin.',
+      userQuery: 'Say hi in five words.',
+    });
+
+    console.log('Bob usage:', bob.usage);
+
+    /* CONTRACT 1: stable prefix is byte-identical across the two users.
+     * If this fails, the cross-user cache cannot work — the cache key
+     * would differ. */
+    expect(bob.systemPrefixSnapshot).toBe(alice.systemPrefixSnapshot);
+
+    /* CONTRACT 2: Bob's stable prefix is served from cache. We accept
+     * either cache_read > 0 (the strict assertion) OR a sane fallback if
+     * Bedrock cache is in cooldown — in which case we at least require
+     * Alice's call wrote one. */
+    const aliceWroteCache =
+      (alice.usage.cache_creation_input_tokens ?? 0) > 0 ||
+      (alice.usage.cache_read_input_tokens ?? 0) > 0;
+    const bobReadCache = (bob.usage.cache_read_input_tokens ?? 0) > 0;
+
+    if (!bobReadCache) {
+      console.warn(
+        'Bob did not see a cache_read. This usually means: ' +
+          '(a) cache write was below Bedrock minimum (try larger STABLE_INSTRUCTIONS), or ' +
+          '(b) the cachePoint reached Bedrock without an inference profile that supports caching. ' +
+          'Alice cache_creation_input_tokens=' +
+          alice.usage.cache_creation_input_tokens
+      );
+    }
+
+    expect(aliceWroteCache).toBe(true);
+    expect(bobReadCache).toBe(true);
+
+    /* CONTRACT 3: cost optimization is real. Bob's effective uncached
+     * input should be much smaller than Alice's first uncached input,
+     * because most of Bob's input was served from cache. */
+    const aliceUncached = alice.usage.input_tokens ?? 0;
+    const bobCacheRead = bob.usage.cache_read_input_tokens ?? 0;
+    expect(bobCacheRead).toBeGreaterThan(aliceUncached / 2);
+  });
+
+  test('5 concurrent users hitting same agent → all read from shared cache entry', async () => {
+    /* Simulates 5 concurrent users on the platform, each with their own
+     * AgentContext (different agentId, different additional_instructions
+     * a.k.a. memory) but identical agent.instructions. They fire requests
+     * in parallel via Promise.all — every one of them must read the
+     * shared system-prefix cache entry. */
+    const distinctUsers = [
+      { id: 'user-1', memory: 'Alice — likes morning meetings.' },
+      { id: 'user-2', memory: 'Bob — based in Berlin.' },
+      { id: 'user-3', memory: 'Carol — works in finance.' },
+      { id: 'user-4', memory: 'Dave — prefers detailed responses.' },
+      { id: 'user-5', memory: 'Eve — uses voice transcription.' },
+    ];
+
+    const results = await Promise.all(
+      distinctUsers.map((u) =>
+        runOneTurn({
+          agentId: u.id,
+          agentInstructions: STABLE_INSTRUCTIONS,
+          additionalInstructions: u.memory,
+          userQuery: 'reply with the word "ok"',
+        })
+      )
+    );
+
+    /* Each user got a usage breakdown — print so we have evidence. */
+    results.forEach((r, i) => {
+      console.log(
+        `Concurrent user ${distinctUsers[i].id}: cache_read=${r.usage.cache_read_input_tokens}, ` +
+          `cache_create=${r.usage.cache_creation_input_tokens}, input=${r.usage.input_tokens}`
+      );
+    });
+
+    /* Stable prefix must be byte-identical across ALL users — same cache key. */
+    const firstSnapshot = results[0].systemPrefixSnapshot;
+    for (let i = 1; i < results.length; i++) {
+      expect(results[i].systemPrefixSnapshot).toBe(firstSnapshot);
+    }
+
+    /* At least one user wrote cache OR all users read from a previously
+     * populated cache. (Re-running the test rapidly may show all 5 reading
+     * the cache from a prior run; first run will show one writing.) */
+    const totalReads = results.reduce(
+      (acc, r) => acc + (r.usage.cache_read_input_tokens ?? 0),
+      0
+    );
+    const totalWrites = results.reduce(
+      (acc, r) => acc + (r.usage.cache_creation_input_tokens ?? 0),
+      0
+    );
+    expect(totalReads + totalWrites).toBeGreaterThan(0);
+
+    /* Most users should read from cache. With 5 concurrent requests on a
+     * cold cache, exactly 1 may write and 4 may read; on a warm cache, 5
+     * read. Either way, at least 4 of 5 must show a cache_read. */
+    const usersWithCacheRead = results.filter(
+      (r) => (r.usage.cache_read_input_tokens ?? 0) > 0
+    ).length;
+    expect(usersWithCacheRead).toBeGreaterThanOrEqual(4);
+  });
+
+  test('byte-identical agent.instructions across runs → cross-run cache hit (single user, two requests)', async () => {
+    /* Same user, two requests with same stable prefix and same dynamic
+     * context. Second request must read from cache. */
+    const r1 = await runOneTurn({
+      agentId: 'sticky-agent',
+      agentInstructions: STABLE_INSTRUCTIONS,
+      additionalInstructions: 'Session 1',
+      userQuery: 'one word: hi',
+    });
+    const r2 = await runOneTurn({
+      agentId: 'sticky-agent',
+      agentInstructions: STABLE_INSTRUCTIONS,
+      additionalInstructions: 'Session 1',
+      userQuery: 'one word: bye',
+    });
+
+    console.log('Run1 usage:', r1.usage);
+    console.log('Run2 usage:', r2.usage);
+
+    expect((r2.usage.cache_read_input_tokens ?? 0) > 0).toBe(true);
+  });
+});