hermes-agent/apps/desktop/electron/freemodel2api-ipc.test.ts

import assert from 'node:assert/strict'

import { test } from 'vitest'

import { registerFreeModel2ApiIpc } from './freemodel2api-ipc'
import type { FetchLike, FreeModel2ApiStoreIo } from './freemodel2api'

function fakeIo() {
  let fileText: string | null = null
  const io: FreeModel2ApiStoreIo = {
    encrypt: plaintext => ({ encoding: 'safeStorage', value: Buffer.from(plaintext, 'utf8').toString('base64') }),
    decrypt: secret => (secret?.encoding === 'safeStorage' ? Buffer.from(String(secret.value), 'base64').toString('utf8') : ''),
    readStoreText: () => {
      if (fileText === null) throw new Error('ENOENT')
      return fileText
    },
    writeStoreText: text => {
      fileText = text
    }
  }
  return { io, fileText: () => fileText }
}

function fakeIpcMain() {
  const handlers = new Map<string, (event: any, payload?: any) => Promise<any>>()
  return {
    handlers,
    ipcMain: { handle: (channel: string, fn: any) => handlers.set(channel, fn) }
  }
}

function jsonResponse(status: number, body: any) {
  return { ok: status >= 200 && status < 300, status, json: async () => body }
}

function makeFetcher(behavior: (url: string, init?: any) => { status: number; body: any }) {
  const calls: { url: string; init?: any }[] = []
  const fetcher: FetchLike = async (url, init) => {
    calls.push({ url, init })
    const { status, body } = behavior(url, init)
    return jsonResponse(status, body)
  }
  return { fetcher, calls }
}

const LOGIN_OK = {
  code: 0,
  data: { access_token: 'AT-1', expires_in: 3600, token_type: 'Bearer', user: { email: 'u@example.com', username: 'yb', balance: 9.5, frozen_balance: 12.5 } }
}

test('login persists credentials and returns balance; status reuses in-memory session', async () => {
  const { handlers, ipcMain } = fakeIpcMain()
  const { io } = fakeIo()
  const { fetcher, calls } = makeFetcher(url => (url.endsWith('/auth/login') ? { status: 200, body: LOGIN_OK } : { status: 200, body: { code: 0, data: { email: 'u@example.com', username: 'yb', balance: 9.5, frozen_balance: 12.5 } } }))
  registerFreeModel2ApiIpc({ ipcMain, io, fetcher })

  const login = await handlers.get('hermes:freemodel2api:login')!(null, { baseUrl: 'https://sub.example.com/', email: 'u@example.com', password: 'pw' })
  assert.deepEqual(login, { ok: true, email: 'u@example.com', balance: 9.5, username: 'yb', frozenBalance: 12.5 })
  assert.deepEqual(JSON.parse(io.decrypt(JSON.parse(io.readStoreText()).freemodel2api)), {
    baseUrl: 'https://sub.example.com', email: 'u@example.com', password: 'pw'
  })

  const status = await handlers.get('hermes:freemodel2api:status')!(null)
  assert.deepEqual(status, { configured: true, baseUrl: 'https://sub.example.com', email: 'u@example.com', balance: 9.5, username: 'yb', frozenBalance: 12.5 })
  // status 用内存里的 session,没有第二次 login 请求
  assert.equal(calls.filter(c => c.url.endsWith('/auth/login')).length, 1)
})

test('status with no stored credentials returns configured:false', async () => {
  const { handlers, ipcMain } = fakeIpcMain()
  const { io } = fakeIo()
  const { fetcher } = makeFetcher(() => ({ status: 200, body: LOGIN_OK }))
  registerFreeModel2ApiIpc({ ipcMain, io, fetcher })
  assert.deepEqual(await handlers.get('hermes:freemodel2api:status')!(null), { configured: false })
})

test('key-balance queries by api key without a session; non-sub2api answers null', async () => {
  const { handlers, ipcMain } = fakeIpcMain()
  const { io } = fakeIo()
  const { fetcher, calls } = makeFetcher(url =>
    url.endsWith('/v1/sub2api/balance')
      ? { status: 200, body: { object: 'sub2api.key_balance', billing_mode: 'balance', balance: 9181.34, unit: 'USD' } }
      : { status: 200, body: LOGIN_OK })
  registerFreeModel2ApiIpc({ ipcMain, io, fetcher })

  // baseUrl 带旧 /v1 也要归到同一个端点;凭 key,不 login
  assert.deepEqual(
    await handlers.get('hermes:freemodel2api:key-balance')!(null, { baseUrl: 'http://h:3020/v1', apiKey: 'sk-k' }),
    { balance: 9181.34 }
  )
  assert.equal(calls[0].url, 'http://h:3020/v1/sub2api/balance')
  assert.equal(calls[0].init?.headers?.Authorization, 'Bearer sk-k')
  assert.equal(calls.filter(c => c.url.endsWith('/auth/login')).length, 0)

  // 非 sub2api 端点 404:不是错误,只是查不到
  const { fetcher: fetcher404 } = makeFetcher(() => ({ status: 404, body: { message: 'nope' } }))
  const second = fakeIpcMain()
  registerFreeModel2ApiIpc({ ipcMain: second.ipcMain, io, fetcher: fetcher404 })
  assert.deepEqual(
    await second.handlers.get('hermes:freemodel2api:key-balance')!(null, { baseUrl: 'http://h:3020', apiKey: 'sk-k' }),
    { balance: null }
  )
})

test('keys re-login once on 401 then succeed', async () => {
  const { handlers, ipcMain } = fakeIpcMain()
  const { io } = fakeIo()
  let keysCalls = 0
  const { fetcher, calls } = makeFetcher(url => {
    if (url.endsWith('/auth/login')) return { status: 200, body: LOGIN_OK }
    if (url.endsWith('/keys')) {
      keysCalls += 1
      if (keysCalls === 1) return { status: 401, body: { code: 401, message: 'expired', data: null } }
      return { status: 200, body: { code: 0, data: { items: [], total: 0, page: 1, page_size: 20, pages: 1 } } }
    }
    return { status: 200, body: { code: 0, data: {} } }
  })
  registerFreeModel2ApiIpc({ ipcMain, io, fetcher })
  await handlers.get('hermes:freemodel2api:login')!(null, { baseUrl: 'https://sub.example.com', email: 'u@example.com', password: 'pw' })
  const result = await handlers.get('hermes:freemodel2api:keys')!(null)
  assert.deepEqual(result, { keys: [] })
  assert.equal(calls.filter(c => c.url.endsWith('/auth/login')).length, 2)
})

test('logout clears stored credentials and session', async () => {
  const { handlers, ipcMain } = fakeIpcMain()
  const { io } = fakeIo()
  const { fetcher } = makeFetcher(() => ({ status: 200, body: LOGIN_OK }))
  registerFreeModel2ApiIpc({ ipcMain, io, fetcher })
  await handlers.get('hermes:freemodel2api:login')!(null, { baseUrl: 'https://sub.example.com', email: 'u@example.com', password: 'pw' })
  await handlers.get('hermes:freemodel2api:logout')!(null)
  assert.deepEqual(await handlers.get('hermes:freemodel2api:status')!(null), { configured: false })
})

test('remembered save/load round trips and survives logout', async () => {
  const { handlers, ipcMain } = fakeIpcMain()
  const { io, fileText } = fakeIo()
  const { fetcher } = makeFetcher(() => ({ status: 200, body: LOGIN_OK }))
  registerFreeModel2ApiIpc({ ipcMain, io, fetcher })

  assert.equal(await handlers.get('hermes:freemodel2api:remembered:load')!(null), null)

  const saved = await handlers.get('hermes:freemodel2api:remembered:save')!(null, {
    baseUrl: 'https://sub.example.com/', email: 'u@example.com', password: 'pw-123'
  })
  assert.deepEqual(saved, { ok: true })
  assert.ok(!fileText()!.includes('pw-123'))

  await handlers.get('hermes:freemodel2api:login')!(null, { baseUrl: 'https://sub.example.com', email: 'u@example.com', password: 'pw' })
  await handlers.get('hermes:freemodel2api:logout')!(null)
  assert.deepEqual(await handlers.get('hermes:freemodel2api:remembered:load')!(null), {
    baseUrl: 'https://sub.example.com', email: 'u@example.com', password: 'pw-123'
  })

  await handlers.get('hermes:freemodel2api:remembered:save')!(null, null)
  assert.equal(await handlers.get('hermes:freemodel2api:remembered:load')!(null), null)
})