boocode/apps/coder/src/services/__tests__/agent-pool.test.ts

import { describe, it, expect, vi } from 'vitest';
import { AgentPool, OPENCODE_POOL_KEY } from '../agent-pool.js';
import type {
  AgentBackend,
  AgentSessionHandle,
  EnsureSessionOpts,
  PromptCtx,
  TurnResult,
} from '../agent-backend.js';

/**
 * v2.6 Phase 3 — AgentPool lifecycle unit test (T.1). No DB / no child process:
 * a fake AgentBackend records dispose + reports busy/health, so we exercise
 * get-or-create, idle eviction, the LRU cap, the busy-never-evict rule, closeChat,
 * and dispose-drains directly. The pure decisions are covered separately in
 * backends/__tests__/lifecycle-decisions.test.ts; this verifies the wiring.
 */

class FakeBackend implements AgentBackend {
  disposed = 0;
  closedSessions = 0;
  private busyFlag = false;
  tickHealthCalls = 0;

  constructor(public readonly name = 'fake') {}

  setBusy(b: boolean): void {
    this.busyFlag = b;
  }

  // — AgentBackend —
  async ensureSession(sessionId: string, opts: EnsureSessionOpts): Promise<AgentSessionHandle> {
    return {
      sessionId,
      agent: opts.agent,
      backend: 'acp_warm',
      chatId: opts.chatId,
      worktreeId: opts.worktreeId,
      agentSessionId: 'fake-session',
      serverPort: null,
    };
  }
  async prompt(_h: AgentSessionHandle, _input: string, _ctx: PromptCtx): Promise<TurnResult> {
    return { ok: true };
  }
  async closeSession(): Promise<void> {
    this.closedSessions++;
  }
  async dispose(): Promise<void> {
    this.disposed++;
  }
  health(): 'up' | 'down' {
    return 'up';
  }
  isBusy(): boolean {
    return this.busyFlag;
  }
  async tickHealth(): Promise<void> {
    this.tickHealthCalls++;
  }
}

describe('AgentPool — get/register/touch (3.1)', () => {
  it('register then get returns the same backend', () => {
    const pool = new AgentPool();
    const b = new FakeBackend();
    pool.register('chat-1', 'goose', b);
    expect(pool.get('chat-1', 'goose')).toBe(b);
    expect(pool.get('chat-1', 'qwen')).toBeUndefined();
  });

  it('peek does NOT exist for a missing key', () => {
    const pool = new AgentPool();
    expect(pool.peek('nope', 'goose')).toBeUndefined();
  });

  it('health reports size + busy count', () => {
    const pool = new AgentPool();
    const a = new FakeBackend();
    const b = new FakeBackend();
    b.setBusy(true);
    pool.register('c1', 'goose', a);
    pool.register('c2', 'qwen', b);
    expect(pool.health()).toEqual({ size: 2, busy: 1 });
  });
});

describe('AgentPool.sweep — idle TTL eviction (3.1)', () => {
  it('evicts an idle backend past the TTL and disposes it', async () => {
    const pool = new AgentPool({ idleTtlMs: 1_000, maxLive: 100 });
    const b = new FakeBackend();
    pool.register('c1', 'goose', b);
    // Sweep with now far past the registration → idle → evicted.
    const { evicted } = await pool.sweep(Date.now() + 10_000);
    expect(evicted).toEqual(['c1:goose']);
    expect(b.disposed).toBe(1);
    expect(pool.get('c1', 'goose')).toBeUndefined();
  });

  it('never evicts a busy backend even past the TTL', async () => {
    const pool = new AgentPool({ idleTtlMs: 1_000, maxLive: 100 });
    const b = new FakeBackend();
    b.setBusy(true);
    pool.register('c1', 'goose', b);
    const { evicted } = await pool.sweep(Date.now() + 10_000);
    expect(evicted).toEqual([]);
    expect(b.disposed).toBe(0);
    expect(pool.get('c1', 'goose')).toBe(b);
  });

  it('touch keeps a backend warm so the TTL measures from the last turn', async () => {
    const pool = new AgentPool({ idleTtlMs: 5_000, maxLive: 100 });
    const b = new FakeBackend();
    pool.register('c1', 'goose', b);
    const base = Date.now();
    // 4s later, touch — resets activity. A sweep at +6s from base is only +2s from
    // the touch → still within TTL → not evicted.
    vi.spyOn(Date, 'now').mockReturnValue(base + 4_000);
    pool.touch('c1', 'goose');
    vi.restoreAllMocks();
    const { evicted } = await pool.sweep(base + 6_000);
    expect(evicted).toEqual([]);
  });
});

describe('AgentPool.sweep — LRU cap (3.4)', () => {
  it('evicts the least-recently-used beyond the cap', async () => {
    const pool = new AgentPool({ idleTtlMs: 1_000_000, maxLive: 2 });
    const base = 1_000_000;
    const mk = (key: string, regAt: number) => {
      vi.spyOn(Date, 'now').mockReturnValue(regAt);
      const b = new FakeBackend(key);
      const [chat, agent] = key.split(':');
      pool.register(chat!, agent!, b);
      vi.restoreAllMocks();
      return b;
    };
    const a = mk('c1:goose', base + 100);
    const b = mk('c2:goose', base + 300);
    const c = mk('c3:goose', base + 200);
    // 3 entries, cap 2, all within idle TTL → LRU (oldest = a@+100) evicted.
    const { evicted } = await pool.sweep(base + 1_000);
    expect(evicted).toEqual(['c1:goose']);
    expect(a.disposed).toBe(1);
    expect(b.disposed).toBe(0);
    expect(c.disposed).toBe(0);
  });
});

describe('AgentPool.sweep — proactive health probe (3.2)', () => {
  it('drives each backend tickHealth before eviction', async () => {
    const pool = new AgentPool({ idleTtlMs: 1_000_000, maxLive: 100 });
    const b = new FakeBackend();
    pool.register('c1', 'opencode', b);
    await pool.sweep(Date.now());
    expect(b.tickHealthCalls).toBe(1);
  });
});

describe('AgentPool.closeChat — chat-close teardown (3.3)', () => {
  it('disposes only the matching chat keys, leaving others + the shared server', async () => {
    const pool = new AgentPool();
    const goose = new FakeBackend('goose');
    const qwen = new FakeBackend('qwen');
    const other = new FakeBackend('other-chat');
    const ocServer = new FakeBackend('opencode-server');
    pool.register('chat-1', 'goose', goose);
    pool.register('chat-1', 'qwen', qwen);
    pool.register('chat-2', 'goose', other);
    pool.register(OPENCODE_POOL_KEY, 'opencode', ocServer);

    const removed = await pool.closeChat('chat-1');
    expect(removed.sort()).toEqual(['chat-1:goose', 'chat-1:qwen']);
    expect(goose.disposed).toBe(1);
    expect(qwen.disposed).toBe(1);
    // other chat + shared opencode server untouched.
    expect(other.disposed).toBe(0);
    expect(ocServer.disposed).toBe(0);
    expect(pool.peek('chat-2', 'goose')).toBe(other);
    expect(pool.peek(OPENCODE_POOL_KEY, 'opencode')).toBe(ocServer);
  });

  it('does not dispose a busy backend on closeChat', async () => {
    const pool = new AgentPool();
    const b = new FakeBackend();
    b.setBusy(true);
    pool.register('chat-1', 'goose', b);
    const removed = await pool.closeChat('chat-1');
    expect(removed).toEqual([]);
    expect(b.disposed).toBe(0);
  });

  it('does not match a chat id that is a prefix of another', async () => {
    // 'chat-1' must not match 'chat-10' — keys are `${chatId}:${agent}` so the
    // colon delimiter prevents the prefix collision.
    const pool = new AgentPool();
    const a = new FakeBackend();
    const b = new FakeBackend();
    pool.register('chat-1', 'goose', a);
    pool.register('chat-10', 'goose', b);
    await pool.closeChat('chat-1');
    expect(a.disposed).toBe(1);
    expect(b.disposed).toBe(0);
    expect(pool.peek('chat-10', 'goose')).toBe(b);
  });
});

describe('AgentPool.dispose — drain all (T.1)', () => {
  it('disposes every backend and clears the map', async () => {
    const pool = new AgentPool();
    const a = new FakeBackend();
    const b = new FakeBackend();
    pool.register('c1', 'goose', a);
    pool.register('c2', 'qwen', b);
    await pool.dispose();
    expect(a.disposed).toBe(1);
    expect(b.disposed).toBe(1);
    expect(pool.health()).toEqual({ size: 0, busy: 0 });
  });

  it('tolerates a backend whose dispose throws', async () => {
    const pool = new AgentPool();
    const good = new FakeBackend();
    const bad = new FakeBackend();
    bad.dispose = async () => {
      throw new Error('boom');
    };
    pool.register('c1', 'goose', bad);
    pool.register('c2', 'qwen', good);
    await expect(pool.dispose()).resolves.toBeUndefined();
    expect(good.disposed).toBe(1);
  });
});