agent-orchestrator/ui-electrobun/tests/unit/agent-store.test.ts

// Tests for Electrobun agent-store — pure logic and data flow.
// Uses bun:test. Mocks appRpc since store depends on RPC calls.

import { describe, it, expect, beforeEach, mock, spyOn } from 'bun:test';

// ── Replicated types ──────────────────────────────────────────────────────────

type AgentStatus = 'idle' | 'running' | 'done' | 'error';
type MsgRole = 'user' | 'assistant' | 'tool-call' | 'tool-result' | 'thinking' | 'system';

interface AgentMessage {
  id: string;
  seqId: number;
  role: MsgRole;
  content: string;
  toolName?: string;
  toolInput?: string;
  toolPath?: string;
  timestamp: number;
}

interface AgentSession {
  sessionId: string;
  projectId: string;
  provider: string;
  status: AgentStatus;
  messages: AgentMessage[];
  costUsd: number;
  inputTokens: number;
  outputTokens: number;
  model: string;
  error?: string;
}

// ── Replicated pure functions ──────────────────────────────────────────────────

function normalizeStatus(status: string): AgentStatus {
  if (status === 'running' || status === 'idle' || status === 'done' || status === 'error') {
    return status;
  }
  return 'idle';
}

const BLOCKED_ENV_PREFIXES = ['CLAUDE', 'CODEX', 'OLLAMA', 'ANTHROPIC_'];

function validateExtraEnv(env: Record<string, string> | undefined): Record<string, string> | undefined {
  if (!env) return undefined;
  const clean: Record<string, string> = {};
  for (const [key, value] of Object.entries(env)) {
    const blocked = BLOCKED_ENV_PREFIXES.some(p => key.startsWith(p));
    if (blocked) continue;
    clean[key] = value;
  }
  return Object.keys(clean).length > 0 ? clean : undefined;
}

function extractToolPath(name: string, input: Record<string, unknown> | undefined): string | undefined {
  if (!input) return undefined;
  if (typeof input.file_path === 'string') return input.file_path;
  if (typeof input.path === 'string') return input.path;
  if (name === 'Bash' && typeof input.command === 'string') {
    return (input.command as string).length > 80 ? (input.command as string).slice(0, 80) + '...' : input.command as string;
  }
  return undefined;
}

function truncateOutput(text: string, maxLines: number): string {
  const lines = text.split('\n');
  if (lines.length <= maxLines) return text;
  return lines.slice(0, maxLines).join('\n') + `\n... (${lines.length - maxLines} more lines)`;
}

// ── seqId monotonic counter ─────────────────────────────────────────────────

function createSeqCounter() {
  const counters = new Map<string, number>();
  return {
    next(sessionId: string): number {
      const current = counters.get(sessionId) ?? 0;
      const next = current + 1;
      counters.set(sessionId, next);
      return next;
    },
    get(sessionId: string): number {
      return counters.get(sessionId) ?? 0;
    },
    set(sessionId: string, value: number): void {
      counters.set(sessionId, value);
    },
  };
}

// ── Tests ───────────────────────────────────────────────────────────────────

describe('normalizeStatus', () => {
  it('passes through valid statuses', () => {
    expect(normalizeStatus('running')).toBe('running');
    expect(normalizeStatus('idle')).toBe('idle');
    expect(normalizeStatus('done')).toBe('done');
    expect(normalizeStatus('error')).toBe('error');
  });

  it('normalizes unknown statuses to idle', () => {
    expect(normalizeStatus('starting')).toBe('idle');
    expect(normalizeStatus('unknown')).toBe('idle');
    expect(normalizeStatus('')).toBe('idle');
  });
});

describe('validateExtraEnv', () => {
  it('returns undefined for undefined input', () => {
    expect(validateExtraEnv(undefined)).toBeUndefined();
  });

  it('strips CLAUDE-prefixed keys', () => {
    const result = validateExtraEnv({ CLAUDE_API_KEY: 'secret', MY_VAR: 'ok' });
    expect(result).toEqual({ MY_VAR: 'ok' });
  });

  it('strips CODEX-prefixed keys', () => {
    const result = validateExtraEnv({ CODEX_TOKEN: 'secret', SAFE: '1' });
    expect(result).toEqual({ SAFE: '1' });
  });

  it('strips OLLAMA-prefixed keys', () => {
    const result = validateExtraEnv({ OLLAMA_HOST: 'localhost', PATH: '/usr/bin' });
    expect(result).toEqual({ PATH: '/usr/bin' });
  });

  it('strips ANTHROPIC_-prefixed keys', () => {
    const result = validateExtraEnv({ ANTHROPIC_API_KEY: 'sk-xxx' });
    expect(result).toBeUndefined(); // all stripped, empty → undefined
  });

  it('returns undefined when all keys blocked', () => {
    const result = validateExtraEnv({ CLAUDE_KEY: 'a', CODEX_KEY: 'b' });
    expect(result).toBeUndefined();
  });

  it('passes through safe keys', () => {
    const result = validateExtraEnv({ BTMSG_AGENT_ID: 'manager', NODE_ENV: 'test' });
    expect(result).toEqual({ BTMSG_AGENT_ID: 'manager', NODE_ENV: 'test' });
  });
});

describe('seqId counter', () => {
  it('starts at 1', () => {
    const counter = createSeqCounter();
    expect(counter.next('s1')).toBe(1);
  });

  it('increments monotonically', () => {
    const counter = createSeqCounter();
    expect(counter.next('s1')).toBe(1);
    expect(counter.next('s1')).toBe(2);
    expect(counter.next('s1')).toBe(3);
  });

  it('tracks separate sessions independently', () => {
    const counter = createSeqCounter();
    expect(counter.next('s1')).toBe(1);
    expect(counter.next('s2')).toBe(1);
    expect(counter.next('s1')).toBe(2);
    expect(counter.next('s2')).toBe(2);
  });

  it('resumes from set value', () => {
    const counter = createSeqCounter();
    counter.set('s1', 42);
    expect(counter.next('s1')).toBe(43);
  });
});

describe('deduplication on restore', () => {
  it('removes duplicate seqIds', () => {
    const messages = [
      { msgId: 'a', seqId: 1, role: 'user', content: 'hello', timestamp: 1000 },
      { msgId: 'b', seqId: 2, role: 'assistant', content: 'hi', timestamp: 1001 },
      { msgId: 'a-dup', seqId: 1, role: 'user', content: 'hello', timestamp: 1002 }, // duplicate seqId
      { msgId: 'c', seqId: 3, role: 'assistant', content: 'ok', timestamp: 1003 },
    ];

    const seqIdSet = new Set<number>();
    const deduplicated: Array<{ msgId: string; seqId: number }> = [];
    let maxSeqId = 0;

    for (const m of messages) {
      const sid = m.seqId ?? 0;
      if (sid > 0 && seqIdSet.has(sid)) continue;
      if (sid > 0) seqIdSet.add(sid);
      if (sid > maxSeqId) maxSeqId = sid;
      deduplicated.push({ msgId: m.msgId, seqId: sid });
    }

    expect(deduplicated).toHaveLength(3);
    expect(maxSeqId).toBe(3);
    expect(deduplicated.map(m => m.msgId)).toEqual(['a', 'b', 'c']);
  });

  it('counter resumes from max seqId after restore', () => {
    const counter = createSeqCounter();
    // Simulate restore
    const maxSeqId = 15;
    counter.set('s1', maxSeqId);
    // Next should be 16
    expect(counter.next('s1')).toBe(16);
  });
});

describe('extractToolPath', () => {
  it('extracts file_path from input', () => {
    expect(extractToolPath('Read', { file_path: '/src/main.ts' })).toBe('/src/main.ts');
  });

  it('extracts path from input', () => {
    expect(extractToolPath('Glob', { path: '/src', pattern: '*.ts' })).toBe('/src');
  });

  it('extracts command from Bash tool', () => {
    expect(extractToolPath('Bash', { command: 'ls -la' })).toBe('ls -la');
  });

  it('truncates long Bash commands at 80 chars', () => {
    const longCmd = 'a'.repeat(100);
    const result = extractToolPath('Bash', { command: longCmd });
    expect(result!.length).toBe(83); // 80 + '...'
  });

  it('returns undefined for unknown tool without paths', () => {
    expect(extractToolPath('Custom', { data: 'value' })).toBeUndefined();
  });
});

describe('truncateOutput', () => {
  it('returns short text unchanged', () => {
    expect(truncateOutput('hello\nworld', 10)).toBe('hello\nworld');
  });

  it('truncates text exceeding maxLines', () => {
    const lines = Array.from({ length: 20 }, (_, i) => `line ${i}`);
    const result = truncateOutput(lines.join('\n'), 5);
    expect(result.split('\n')).toHaveLength(6); // 5 lines + truncation message
    expect(result).toContain('15 more lines');
  });
});

describe('double-start guard', () => {
  it('startingProjects Set prevents concurrent starts', () => {
    const startingProjects = new Set<string>();

    // First start: succeeds
    expect(startingProjects.has('proj-1')).toBe(false);
    startingProjects.add('proj-1');
    expect(startingProjects.has('proj-1')).toBe(true);

    // Second start: blocked
    const blocked = startingProjects.has('proj-1');
    expect(blocked).toBe(true);

    // After completion: removed
    startingProjects.delete('proj-1');
    expect(startingProjects.has('proj-1')).toBe(false);
  });
});

describe('persistMessages — lastPersistedIndex', () => {
  it('only saves new messages from lastPersistedIndex', () => {
    const allMessages = [
      { id: 'm1', content: 'a' },
      { id: 'm2', content: 'b' },
      { id: 'm3', content: 'c' },
      { id: 'm4', content: 'd' },
    ];

    const lastPersistedIndex = 2;
    const newMsgs = allMessages.slice(lastPersistedIndex);
    expect(newMsgs).toHaveLength(2);
    expect(newMsgs[0].id).toBe('m3');
    expect(newMsgs[1].id).toBe('m4');
  });

  it('returns empty when nothing new', () => {
    const allMessages = [{ id: 'm1', content: 'a' }];
    const lastPersistedIndex = 1;
    const newMsgs = allMessages.slice(lastPersistedIndex);
    expect(newMsgs).toHaveLength(0);
  });
});

describe('loadLastSession — active session guard', () => {
  it('skips restore if project has running session', () => {
    const sessions: Record<string, AgentSession> = {
      's1': {
        sessionId: 's1', projectId: 'p1', provider: 'claude',
        status: 'running', messages: [], costUsd: 0,
        inputTokens: 0, outputTokens: 0, model: 'claude-opus-4-5',
      },
    };
    const projectSessionMap = new Map([['p1', 's1']]);
    const startingProjects = new Set<string>();

    const existingSessionId = projectSessionMap.get('p1');
    let shouldSkip = false;
    if (existingSessionId) {
      const existing = sessions[existingSessionId];
      if (existing && (existing.status === 'running' || startingProjects.has('p1'))) {
        shouldSkip = true;
      }
    }
    expect(shouldSkip).toBe(true);
  });
});