sync: update from private repo (c8b6d29)

src/worker/sticky-backend.test.ts

@@ -13,11 +13,11 @@ function makeLogger(): StickyBackendLogger & {
   };
 }
 
-describe('createStickyBackendResolver', () => {
-  it('persists the first backend, sets sticky, and logs at info', async () => {
+describe('createStickyBackendResolver (follow-current semantics)', () => {
+  it('persists the first backend, advances current, and logs at info', async () => {
     const logger = makeLogger();
     const persist = vi.fn().mockResolvedValue(undefined);
-    const resolve = createStickyBackendResolver({
+    const tracker = createStickyBackendResolver({
       initial: null,
       persist,
       logger,
@@ -25,18 +25,19 @@ describe('createStickyBackendResolver', () => {
       jobId: 'j1',
     });
 
-    await resolve({ backendId: 'gpu-a', cacheKey: null });
+    await tracker.onEvent({ backendId: 'gpu-a', cacheKey: null });
 
     expect(persist).toHaveBeenCalledTimes(1);
     expect(persist).toHaveBeenCalledWith('gpu-a');
+    expect(tracker.current()).toBe('gpu-a');
     expect(logger.calls.info).toHaveLength(1);
     expect(logger.calls.info[0]).toContain('gpu-a');
   });
 
-  it('short-circuits subsequent events once sticky is set', async () => {
+  it('follows backend switches: each CHANGE persists; repeats do not', async () => {
     const logger = makeLogger();
     const persist = vi.fn().mockResolvedValue(undefined);
-    const resolve = createStickyBackendResolver({
+    const tracker = createStickyBackendResolver({
       initial: null,
       persist,
       logger,
@@ -44,20 +45,24 @@ describe('createStickyBackendResolver', () => {
       jobId: 'j1',
     });
 
-    await resolve({ backendId: 'gpu-a', cacheKey: null });
-    await resolve({ backendId: 'gpu-b', cacheKey: 'sha:xyz' });
-    await resolve({ backendId: 'gpu-a', cacheKey: null });
+    await tracker.onEvent({ backendId: 'gpu-a', cacheKey: null });
+    await tracker.onEvent({ backendId: 'gpu-a', cacheKey: 'sha:1' }); // same — no persist
+    await tracker.onEvent({ backendId: 'gpu-b', cacheKey: null });    // switch — persists
+    await tracker.onEvent({ backendId: 'gpu-b', cacheKey: null });    // same — no persist
 
-    expect(persist).toHaveBeenCalledTimes(1);
-    // gpu-b ≠ sticky → debug; gpu-a == sticky → no log
-    expect(logger.calls.debug).toHaveLength(1);
-    expect(logger.calls.debug[0]).toContain('gpu-b');
+    expect(persist).toHaveBeenCalledTimes(2);
+    expect(persist).toHaveBeenNthCalledWith(1, 'gpu-a');
+    expect(persist).toHaveBeenNthCalledWith(2, 'gpu-b');
+    expect(tracker.current()).toBe('gpu-b');
+    // The switch log mentions both ends
+    expect(logger.calls.info[1]).toContain('gpu-a');
+    expect(logger.calls.info[1]).toContain('gpu-b');
   });
 
-  it('honors initial sticky from DB without re-persisting', async () => {
+  it('honors initial value from DB: same backend does not re-persist, a switch does', async () => {
     const logger = makeLogger();
     const persist = vi.fn().mockResolvedValue(undefined);
-    const resolve = createStickyBackendResolver({
+    const tracker = createStickyBackendResolver({
       initial: 'gpu-seed',
       persist,
       logger,
@@ -65,21 +70,22 @@ describe('createStickyBackendResolver', () => {
       jobId: 'j1',
     });
 
-    await resolve({ backendId: 'gpu-other', cacheKey: null });
-    await resolve({ backendId: 'gpu-seed', cacheKey: null });
-
+    expect(tracker.current()).toBe('gpu-seed');
+    await tracker.onEvent({ backendId: 'gpu-seed', cacheKey: null });
     expect(persist).not.toHaveBeenCalled();
-    expect(logger.calls.debug).toHaveLength(1);
-    expect(logger.calls.debug[0]).toContain('gpu-other');
+
+    await tracker.onEvent({ backendId: 'gpu-other', cacheKey: null });
+    expect(persist).toHaveBeenCalledTimes(1);
+    expect(tracker.current()).toBe('gpu-other');
   });
 
-  it('does NOT set sticky when persist fails — next event retries', async () => {
+  it('does NOT advance when persist fails — next event retries', async () => {
     const logger = makeLogger();
     const persist = vi
       .fn()
       .mockRejectedValueOnce(new Error('SQLITE_BUSY'))
       .mockResolvedValueOnce(undefined);
-    const resolve = createStickyBackendResolver({
+    const tracker = createStickyBackendResolver({
       initial: null,
       persist,
       logger,
@@ -87,23 +93,19 @@ describe('createStickyBackendResolver', () => {
       jobId: 'j1',
     });
 
-    // First call: DB write fails → sticky unset → warn logged
-    await resolve({ backendId: 'gpu-a', cacheKey: null });
+    // First call: DB write fails → current stays null → warn logged
+    await tracker.onEvent({ backendId: 'gpu-a', cacheKey: null });
     expect(persist).toHaveBeenCalledTimes(1);
+    expect(tracker.current()).toBeNull();
     expect(logger.calls.warn).toHaveLength(1);
     expect(logger.calls.warn[0]).toContain('SQLITE_BUSY');
     expect(logger.calls.info).toHaveLength(0);
 
-    // Second call: DB write succeeds → sticky set
-    await resolve({ backendId: 'gpu-b', cacheKey: null });
+    // Second call (same backend again): retries because current ≠ backendId
+    await tracker.onEvent({ backendId: 'gpu-a', cacheKey: null });
     expect(persist).toHaveBeenCalledTimes(2);
-    expect(persist).toHaveBeenLastCalledWith('gpu-b');
+    expect(tracker.current()).toBe('gpu-a');
     expect(logger.calls.info).toHaveLength(1);
-    expect(logger.calls.info[0]).toContain('gpu-b');
-
-    // Third call: sticky is now set → no further persist
-    await resolve({ backendId: 'gpu-c', cacheKey: null });
-    expect(persist).toHaveBeenCalledTimes(2);
   });
 
   it('retries on every event until persist succeeds (multiple failures)', async () => {
@@ -113,7 +115,7 @@ describe('createStickyBackendResolver', () => {
       .mockRejectedValueOnce(new Error('fail 1'))
       .mockRejectedValueOnce(new Error('fail 2'))
       .mockResolvedValueOnce(undefined);
-    const resolve = createStickyBackendResolver({
+    const tracker = createStickyBackendResolver({
       initial: null,
       persist,
       logger,
@@ -121,13 +123,13 @@ describe('createStickyBackendResolver', () => {
       jobId: 'j1',
     });
 
-    await resolve({ backendId: 'gpu-a', cacheKey: null });
-    await resolve({ backendId: 'gpu-b', cacheKey: null });
-    await resolve({ backendId: 'gpu-c', cacheKey: null });
+    await tracker.onEvent({ backendId: 'gpu-a', cacheKey: null });
+    await tracker.onEvent({ backendId: 'gpu-b', cacheKey: null });
+    await tracker.onEvent({ backendId: 'gpu-c', cacheKey: null });
 
     expect(persist).toHaveBeenCalledTimes(3);
     expect(logger.calls.warn).toHaveLength(2);
     expect(logger.calls.info).toHaveLength(1);
-    expect(logger.calls.info[0]).toContain('gpu-c');
+    expect(tracker.current()).toBe('gpu-c');
   });
 });

src/worker/sticky-backend.ts

@@ -1,24 +1,23 @@
 /**
- * Sticky-backend resolver for proxy worker jobs.
+ * Backend tracker for proxy worker jobs.
  *
- * Design (per
- *   docs/superpowers/specs/2026-05-18-multi-team-gpu-pool-and-node-status-design.md
- * Open Question #3 case 1, "最初に確定したら以後 update しない"):
+ * History: originally "first backend wins" (2026-05-18 design, Open Question
+ * #3 case 1) to keep the UI pet from flickering while the gateway rebalanced
+ * every request. As of 2026-06 the gateway honors `x-aao-preferred-backend`
+ * (client-side sticky routing for KV-cache reuse), so backend switches are
+ * RARE — they only happen when the preferred backend goes offline or
+ * saturates. The tracker therefore now follows the CURRENT backend:
  *
- * - For a proxy worker, every LLM call may resolve to a different
- *   physical backend (LiteLLM rebalances per request). The UI pet should
- *   not flicker, so we record only the FIRST backend a job sees.
- * - Persistence happens via `updateJob({ lastBackendId })`. If that DB
- *   write FAILS, we must remain in the unset state so the next
- *   `onBackendResolved` event has a chance to retry. If we set the local
- *   sticky variable BEFORE persisting, a transient DB error would lose
- *   the binding permanently for the lifetime of the job (sticky check
- *   short-circuits all subsequent events) and the UI would never see
- *   the worker → backend mapping.
+ * - `jobs.last_backend_id` is updated whenever the resolved backend CHANGES,
+ *   so the UI (pet, badges) tracks where the job actually runs.
+ * - Persistence happens via `updateJob({ lastBackendId })`. If that DB write
+ *   FAILS, the in-memory value is left unchanged so the next
+ *   `onBackendResolved` event retries the persist (a transient DB error must
+ *   not permanently lose the worker → backend mapping).
  *
- * This module isolates the "set sticky only after persist succeeds"
- * invariant from `Worker.buildPieceRunCallbacks`, which already has a
- * dozen other concerns and is hard to unit-test in isolation.
+ * This module isolates the "advance only after persist succeeds" invariant
+ * from `Worker.buildPieceCallbacks`, which already has a dozen other
+ * concerns and is hard to unit-test in isolation.
  */
 
 export interface StickyBackendLogger {
@@ -32,16 +31,20 @@ export interface StickyBackendEvent {
   cacheKey: string | null;
 }
 
+export interface BackendTracker {
+  /** The onBackendResolved callback for the agent loop (fire-and-forget safe). */
+  onEvent: (event: StickyBackendEvent) => Promise<void>;
+  /**
+   * The most recently persisted backend id (or the initial DB value).
+   * Used as the `x-aao-preferred-backend` hint on the next LLM request.
+   */
+  current: () => string | null;
+}
+
 /**
- * Build the `onBackendResolved` callback. `persist(backendId)` is the
- * DB write (typically `repo.updateJob(jobId, { lastBackendId })`); it
- * must reject on failure so we can keep sticky unset for retry.
- *
- * Returns an async function the agent-loop can call without awaiting —
- * errors are caught internally and logged. The function resolves once
- * either: (a) sticky was already set and we short-circuited, (b) the
- * persist succeeded and sticky is now set, or (c) the persist failed
- * and sticky remains unset for the next event to retry.
+ * Build the backend tracker. `persist(backendId)` is the DB write
+ * (typically `repo.updateJob(jobId, { lastBackendId })`); it must reject on
+ * failure so the in-memory value stays put for retry.
  */
 export function createStickyBackendResolver(opts: {
   initial: string | null;
@@ -49,33 +52,31 @@ export function createStickyBackendResolver(opts: {
   logger: StickyBackendLogger;
   workerId: string;
   jobId: string;
-}): (event: StickyBackendEvent) => Promise<void> {
+}): BackendTracker {
   const { initial, persist, logger, workerId, jobId } = opts;
-  let sticky: string | null = initial;
+  let current: string | null = initial;
 
-  return async function onBackendResolved({ backendId, cacheKey }: StickyBackendEvent): Promise<void> {
-    if (sticky) {
-      if (sticky !== backendId) {
-        logger.debug(
-          `[worker:${workerId}] job ${jobId} backend re-resolved to ${backendId} (sticky=${sticky}, cache=${cacheKey ?? 'miss'}); keeping sticky`,
-        );
-      }
-      return;
-    }
+  async function onEvent({ backendId, cacheKey }: StickyBackendEvent): Promise<void> {
+    if (current === backendId) return; // unchanged — nothing to persist
     try {
       await persist(backendId);
-      // Only set sticky AFTER persist succeeds. If we set first and persist
-      // failed, the next event would short-circuit on the sticky check and
-      // we'd never recover — the UI would render "no backend" forever.
-      sticky = backendId;
+      // Only advance AFTER persist succeeds. If we advanced first and the
+      // persist failed, the next identical event would short-circuit on the
+      // equality check and the DB would stay stale forever.
+      const previous = current;
+      current = backendId;
       logger.info(
-        `[worker:${workerId}] job ${jobId} backend resolved: ${backendId} cache=${cacheKey ?? 'miss'}`,
+        previous
+          ? `[worker:${workerId}] job ${jobId} backend switched: ${previous} → ${backendId} cache=${cacheKey ?? 'miss'}`
+          : `[worker:${workerId}] job ${jobId} backend resolved: ${backendId} cache=${cacheKey ?? 'miss'}`,
       );
     } catch (err) {
       logger.warn(
-        `[worker:${workerId}] failed to persist lastBackendId for job ${jobId}: ${err} — sticky left unset for retry`,
+        `[worker:${workerId}] failed to persist lastBackendId for job ${jobId}: ${err} — keeping ${current ?? 'unset'} for retry`,
       );
-      // Intentionally do NOT set sticky. Next event retries.
+      // Intentionally do NOT advance. Next event retries.
     }
-  };
+  }
+
+  return { onEvent, current: () => current };
 }