EJClaw/src/token-refresh.test.ts

import { describe, expect, it } from 'vitest';

import {
  applyUpdatedTokensToEnvContent,
  collectSessionCredentialPaths,
  pickFreshestOAuth,
  shouldAdoptSessionOAuth,
  shouldStartTokenRefreshLoop,
} from './token-refresh.js';

function creds(expiresAt: number, refreshToken = 'rt') {
  return {
    claudeAiOauth: {
      accessToken: `at-${expiresAt}`,
      refreshToken,
      expiresAt,
      scopes: [] as string[],
    },
  };
}

describe('shouldStartTokenRefreshLoop', () => {
  it('starts refresh for the Claude service', () => {
    expect(shouldStartTokenRefreshLoop('claude-code')).toBe(true);
  });

  it('skips refresh for the Codex service', () => {
    expect(shouldStartTokenRefreshLoop('codex')).toBe(false);
  });

  it('updates both multi-token and single-token env vars when present', () => {
    const next = applyUpdatedTokensToEnvContent(
      [
        'CLAUDE_CODE_OAUTH_TOKEN=old-primary',
        'CLAUDE_CODE_OAUTH_TOKENS=old-primary,old-secondary',
        'OTHER=value',
      ].join('\n'),
      ['new-primary', 'new-secondary'],
    );

    expect(next).toContain('CLAUDE_CODE_OAUTH_TOKEN=new-primary');
    expect(next).toContain(
      'CLAUDE_CODE_OAUTH_TOKENS=new-primary,new-secondary',
    );
  });

  it('adds the multi-token env var when it is missing', () => {
    const next = applyUpdatedTokensToEnvContent(
      ['CLAUDE_CODE_OAUTH_TOKEN=old-primary', 'OTHER=value'].join('\n'),
      ['new-primary', 'new-secondary'],
    );

    expect(next).toContain('CLAUDE_CODE_OAUTH_TOKEN=new-primary');
    expect(next).toContain(
      'CLAUDE_CODE_OAUTH_TOKENS=new-primary,new-secondary',
    );
  });
});

describe('pickFreshestOAuth', () => {
  it('adopts a session copy that expires later than the source', () => {
    const best = pickFreshestOAuth([
      { label: 'source', creds: creds(1000) },
      { label: 'sessionA', creds: creds(5000) },
      { label: 'sessionB', creds: creds(3000) },
    ]);
    expect(best?.label).toBe('sessionA');
  });

  it('keeps the source on ties so we do not churn copies needlessly', () => {
    const best = pickFreshestOAuth([
      { label: 'source', creds: creds(5000) },
      { label: 'sessionA', creds: creds(5000) },
    ]);
    expect(best?.label).toBe('source');
  });

  it('ignores candidates without a refresh token or numeric expiry', () => {
    const best = pickFreshestOAuth([
      { label: 'source', creds: creds(1000) },
      { label: 'noRefresh', creds: creds(9000, '') },
      { label: 'missing', creds: null },
    ]);
    expect(best?.label).toBe('source');
  });

  it('returns null when no candidate is usable', () => {
    expect(
      pickFreshestOAuth([
        { label: 'source', creds: null },
        { label: 'noRefresh', creds: creds(9000, '') },
      ]),
    ).toBeNull();
  });
});

describe('shouldAdoptSessionOAuth (session→canonical write-back)', () => {
  const oauth = (expiresAt: number, extra: Record<string, unknown> = {}) => ({
    accessToken: `at-${expiresAt}`,
    refreshToken: `rt-${expiresAt}`,
    expiresAt,
    scopes: [] as string[],
    ...extra,
  });

  it('adopts a session token strictly newer than canonical', () => {
    expect(shouldAdoptSessionOAuth(oauth(2000), oauth(1000))).toBe(true);
  });

  it('does not regress to an older or equal session token', () => {
    expect(shouldAdoptSessionOAuth(oauth(1000), oauth(2000))).toBe(false);
    expect(shouldAdoptSessionOAuth(oauth(2000), oauth(2000))).toBe(false);
  });

  it('adopts when canonical is missing/invalid', () => {
    expect(shouldAdoptSessionOAuth(oauth(1000), null)).toBe(true);
  });

  it('rejects an invalid session token', () => {
    expect(shouldAdoptSessionOAuth(null, oauth(1000))).toBe(false);
    expect(
      shouldAdoptSessionOAuth(
        { accessToken: '', refreshToken: '', expiresAt: 5000, scopes: [] },
        oauth(1000),
      ),
    ).toBe(false);
  });

  it('refuses to cross subscription types', () => {
    expect(
      shouldAdoptSessionOAuth(
        oauth(2000, { subscriptionType: 'pro' }),
        oauth(1000, { subscriptionType: 'max' }),
      ),
    ).toBe(false);
  });
});

describe('collectSessionCredentialPaths (real session layout)', () => {
  it('builds service group + task credential paths, not <group>/.claude', () => {
    const tree: Record<string, string[]> = {
      '/s': ['iger', 'javis_bot'],
      '/s/iger/services': ['claude'],
      '/s/iger/services/claude/tasks': ['task-1'],
      '/s/javis_bot/services': ['claude'],
      '/s/javis_bot/services/claude/tasks': [],
    };
    const paths = collectSessionCredentialPaths('/s', (dir) => tree[dir] ?? []);
    expect(paths).toContain(
      '/s/iger/services/claude/.claude/.credentials.json',
    );
    expect(paths).toContain(
      '/s/iger/services/claude/tasks/task-1/.claude/.credentials.json',
    );
    expect(paths).toContain(
      '/s/javis_bot/services/claude/.claude/.credentials.json',
    );
    // Must NOT use the old broken <group>/.claude path that matched nothing.
    expect(paths).not.toContain('/s/iger/.claude/.credentials.json');
  });

  it('returns nothing when there are no session folders', () => {
    expect(collectSessionCredentialPaths('/s', () => [])).toEqual([]);
  });
});