goodbuddy/src/agent-daemon/agent-model-gateway.test.ts

import { mkdirSync, mkdtempSync, readFileSync, rmSync } from 'node:fs'
import { tmpdir } from 'node:os'
import { join } from 'node:path'
import { afterEach, describe, expect, it, vi } from 'vitest'
import type { AgentPromptModelProfile } from '../shared/model-bridge-contracts'
import { REMOTE_MODEL_GATEWAY_LIMITS } from '../shared/remote-model-gateway-contracts'
import {
  AgentModelCallLedger,
  AgentModelGateway,
  type ModelCallLedger
} from './agent-model-gateway'

const temporary: string[] = []
const secret = 'provider-secret-never-persist'

function gate() {
  let resolve!: () => void
  const promise = new Promise<void>(done => { resolve = done })
  return { promise, resolve }
}

function setup(fetcher: typeof fetch) {
  const root = mkdtempSync(join(tmpdir(), 'goodbuddy-agent-model-'))
  temporary.push(root)
  mkdirSync(join(root, 'state'), { mode: 0o700 })
  const path = join(root, 'state', 'calls.sqlite')
  const ledger = new AgentModelCallLedger(path)
  return {
    path,
    ledger,
    gateway: new AgentModelGateway({ ledger, fetcher })
  }
}

const profile: AgentPromptModelProfile = {
  profileId: 'profile-1',
  modelProfileDigest: `sha256:${'a'.repeat(64)}`,
  provider: 'openai',
  baseUrl: 'https://provider.example/v1',
  model: 'model-1',
  protocol: 'openai-responses',
  authentication: 'api-key',
  apiKey: secret,
  capabilities: { imageInput: false },
  limits: {
    maximumOutputTokens: 4_096,
    requestTimeoutMilliseconds: 5_000
  }
}

const request = {
  method: 'POST' as const,
  path: '/v1/responses' as const,
  headers: { 'content-type': 'application/json' },
  bodyBase64: Buffer.from(JSON.stringify({ model: 'model-1' })).toString(
    'base64'
  )
}

const context = {
  bindingId: 'binding-1',
  operationId: 'operation-1',
  promptSequence: 0,
  roundIndex: 0,
  profileDigest: profile.modelProfileDigest,
  profile
}

afterEach(() => {
  for (const path of temporary.splice(0)) {
    rmSync(path, { recursive: true, force: true })
  }
})

describe('AgentModelGateway', () => {
  it('awaits host claim, completion and delivery before their dependent actions', async () => {
    const claimGate = gate()
    const completeGate = gate()
    const deliveryGate = gate()
    const ledger: ModelCallLedger = {
      claim: vi.fn(() => claimGate.promise), complete: vi.fn(() => completeGate.promise),
      delivered: vi.fn(() => deliveryGate.promise), outcomeUnknown: vi.fn(), get: vi.fn(), close: vi.fn()
    }
    const fetcher = vi.fn<typeof fetch>(async () => new Response('{}'))
    const gateway = new AgentModelGateway({ ledger, fetcher })
    const settled = vi.fn()
    const pending = gateway.dispatch(context, request, new AbortController().signal).then(result => { settled(); return result })
    await Promise.resolve()
    expect(fetcher).not.toHaveBeenCalled()
    claimGate.resolve()
    await vi.waitFor(() => expect(ledger.complete).toHaveBeenCalledOnce())
    expect(settled).not.toHaveBeenCalled()
    completeGate.resolve()
    const result = await pending
    const acknowledged = vi.fn()
    const delivery = result.acknowledgeDelivery().then(acknowledged)
    await Promise.resolve()
    expect(acknowledged).not.toHaveBeenCalled()
    deliveryGate.resolve()
    await delivery
    expect(acknowledged).toHaveBeenCalledOnce()
  })

  it.each(['claim', 'complete', 'delivered'] as const)('propagates a rejected host %s without claiming success', async method => {
    const failure = new Error('storage unavailable')
    const ledger: ModelCallLedger = {
      claim: vi.fn(), complete: vi.fn(), delivered: vi.fn(), outcomeUnknown: vi.fn(), get: vi.fn(), close: vi.fn(),
      [method]: vi.fn(async () => { throw failure })
    }
    const fetcher = vi.fn<typeof fetch>(async () => new Response('{}'))
    const gateway = new AgentModelGateway({ ledger, fetcher })
    const pending = gateway.dispatch(context, request, new AbortController().signal)
    if (method === 'delivered') await expect((await pending).acknowledgeDelivery()).rejects.toBe(failure)
    else await expect(pending).rejects.toBe(failure)
    expect(fetcher).toHaveBeenCalledTimes(method === 'claim' ? 0 : 1)
  })

  it('settles cancellation during claim only after recording uncertainty, without dispatch', async () => {
    const claimGate = gate()
    const unknownGate = gate()
    const ledger: ModelCallLedger = {
      claim: () => claimGate.promise, complete: vi.fn(), delivered: vi.fn(),
      outcomeUnknown: vi.fn(() => unknownGate.promise), get: vi.fn(), close: vi.fn()
    }
    const fetcher = vi.fn<typeof fetch>()
    const gateway = new AgentModelGateway({ ledger, fetcher })
    const controller = new AbortController()
    const settled = vi.fn()
    const pending = gateway.dispatch(context, request, controller.signal).catch(error => { settled(); throw error })
    const rejected = expect(pending).rejects.toMatchObject({ code: 'cancelled' })
    controller.abort()
    claimGate.resolve()
    await vi.waitFor(() => expect(ledger.outcomeUnknown).toHaveBeenCalledOnce())
    expect(fetcher).not.toHaveBeenCalled()
    expect(settled).not.toHaveBeenCalled()
    unknownGate.resolve()
    await rejected
  })

  it('keeps a slow provider request alive past 60 seconds and still cancels it', async () => {
    vi.useFakeTimers()
    let providerSignal: AbortSignal | undefined
    const { gateway, ledger } = setup(vi.fn<typeof fetch>(async (_url, init) => {
      providerSignal = init!.signal!
      return await new Promise<Response>((_resolve, reject) => {
        providerSignal!.addEventListener('abort', () => reject(providerSignal!.reason), { once: true })
      })
    }))
    try {
      const controller = new AbortController()
      const pending = gateway.dispatch({ ...context, profile: { ...profile, limits: {} } }, request, controller.signal)
      const rejected = expect(pending).rejects.toMatchObject({ code: 'cancelled' })
      await vi.advanceTimersByTimeAsync(600_000)
      expect(providerSignal?.aborted).toBe(false)
      controller.abort()
      await rejected
    } finally {
      ledger.close()
      vi.useRealTimers()
    }
  })

  it.each([undefined, 128_000])('preserves runtime output hints without an invented token default: %s', async (tokens) => {
    const payload = { model: 'model-1', ...(tokens === undefined ? {} : { max_output_tokens: tokens }) }
    const fetcher = vi.fn<typeof fetch>(async (_url, init) => {
      expect(JSON.parse(String(init?.body))).toEqual(payload)
      return new Response('x'.repeat(1024 * 1024))
    })
    const { gateway, ledger } = setup(fetcher)
    try {
      const result = await gateway.dispatch(
        { ...context, profile: { ...profile, limits: {} } },
        { ...request, bodyBase64: Buffer.from(JSON.stringify(payload)).toString('base64') },
        new AbortController().signal
      )
      expect(Buffer.from(result.response.bodyBase64, 'base64').length).toBe(1024 * 1024)
      await result.acknowledgeDelivery()
    } finally {
      ledger.close()
    }
  })

  it('dispatches a stable call exactly once and records delivery ACK without persisting credentials', async () => {
    const fetcher = vi.fn<typeof fetch>(async (_url, init) => {
      expect(new Headers(init?.headers).get('authorization')).toBe(
        `Bearer ${secret}`
      )
      expect(JSON.parse(String(init?.body))).toMatchObject({
        max_output_tokens: 4_096
      })
      return new Response(JSON.stringify({ id: 'response-1' }), {
        status: 200,
        headers: { 'content-type': 'application/json' }
      })
    })
    const { gateway, ledger, path } = setup(fetcher)
    const result = await gateway.dispatch(
      context,
      request,
      new AbortController().signal
    )
    await result.acknowledgeDelivery()
    await expect(
      gateway.dispatch(context, request, new AbortController().signal)
    ).rejects.toMatchObject({ code: 'already-dispatched' })
    expect(fetcher).toHaveBeenCalledOnce()
    ledger.close()
    expect(readFileSync(path).includes(Buffer.from(secret))).toBe(false)
  })

  it('caps provider output hints at the prompt-scoped limit', async () => {
    const fetcher = vi.fn<typeof fetch>(async (_url, init) => {
      expect(JSON.parse(String(init?.body))).toMatchObject({
        max_output_tokens: 4_096
      })
      return new Response(JSON.stringify({ id: 'response-1' }), {
        status: 200,
        headers: { 'content-type': 'application/json' }
      })
    })
    const { gateway, ledger } = setup(fetcher)
    const result = await gateway.dispatch(
      context,
      {
        ...request,
        bodyBase64: Buffer.from(
          JSON.stringify({
            model: 'model-1',
            max_output_tokens: 32_000
          })
        ).toString('base64')
      },
      new AbortController().signal
    )
    await result.acknowledgeDelivery()
    expect(fetcher).toHaveBeenCalledOnce()
    ledger.close()
  })

  it('merges trusted profile request customization without overriding Runtime fields', async () => {
    const fetcher = vi.fn<typeof fetch>(async (_url, init) => {
      const headers = new Headers(init?.headers)
      expect(headers.get('x-tenant-id')).toBe('agent-tenant')
      expect(headers.get('authorization')).toBe(`Bearer ${secret}`)
      expect(JSON.parse(String(init?.body))).toEqual({
        max_output_tokens: 4_096,
        metadata: { source: 'goodbuddy' },
        model: 'model-1',
        temperature: 0.8
      })
      return new Response('{}')
    })
    const { gateway, ledger } = setup(fetcher)
    const customizedProfile: AgentPromptModelProfile = {
      ...profile,
      requestHeaders: {
        'x-tenant-id': 'agent-tenant'
      },
      requestBody: {
        temperature: 0.2,
        metadata: { source: 'goodbuddy' }
      }
    }
    const result = await gateway.dispatch(
      {
        ...context,
        profile: customizedProfile
      },
      {
        ...request,
        bodyBase64: Buffer.from(
          JSON.stringify({
            model: 'model-1',
            temperature: 0.8
          })
        ).toString('base64')
      },
      new AbortController().signal
    )

    await result.acknowledgeDelivery()
    expect(fetcher).toHaveBeenCalledOnce()
    ledger.close()
  })

  it('does not impose prompt-wide model-call or output-token quotas', async () => {
    const fetcher = vi.fn<typeof fetch>(async () =>
      new Response(
        JSON.stringify({
          id: 'response',
          usage: { output_tokens: 4_096 }
        }),
        {
          status: 200,
          headers: { 'content-type': 'application/json' }
        }
      )
    )
    const { gateway, ledger } = setup(fetcher)

    for (let roundIndex = 0; roundIndex < 101; roundIndex += 1) {
      const result = await gateway.dispatch(
        { ...context, roundIndex },
        request,
        new AbortController().signal
      )
      await result.acknowledgeDelivery()
    }

    expect(fetcher).toHaveBeenCalledTimes(101)
    ledger.close()
  })

  it('caps OpenCode chat-completions max_tokens before dispatch', async () => {
    const fetcher = vi.fn<typeof fetch>(async (_url, init) => {
      expect(JSON.parse(String(init?.body))).toMatchObject({
        max_tokens: 4_096
      })
      return new Response(JSON.stringify({ id: 'response-1' }), {
        status: 200,
        headers: { 'content-type': 'application/json' }
      })
    })
    const { gateway, ledger } = setup(fetcher)
    const result = await gateway.dispatch(
      {
        ...context,
        profile: {
          ...profile,
          protocol: 'openai-chat-completions'
        }
      },
      {
        method: 'POST',
        path: '/v1/chat/completions',
        headers: { 'content-type': 'application/json' },
        bodyBase64: Buffer.from(
          JSON.stringify({
            model: 'model-1',
            max_tokens: 32_000
          })
        ).toString('base64')
      },
      new AbortController().signal
    )
    await result.acknowledgeDelivery()
    expect(fetcher).toHaveBeenCalledOnce()
    ledger.close()
  })

  it.each([
    { background: true },
    { store: true },
    { web_search_options: {} },
    { mcp_servers: [] },
    { previous_response_id: 'response-previous' },
    { tools: [{ type: 'web_search' }] },
    { include: ['web_search_call.results'] },
    { modalities: ['text', 'audio'] }
  ])('rejects independently billed provider capabilities %#', async (extra) => {
    const fetcher = vi.fn<typeof fetch>()
    const { gateway, ledger } = setup(fetcher)
    await expect(
      gateway.dispatch(
        context,
        {
          ...request,
          bodyBase64: Buffer.from(
            JSON.stringify({ model: 'model-1', ...extra })
          ).toString('base64')
        },
        new AbortController().signal
      )
    ).rejects.toMatchObject({ code: 'policy' })
    expect(fetcher).not.toHaveBeenCalled()
    ledger.close()
  })

  it('marks an in-flight provider failure outcome-unknown and never retries it', async () => {
    const fetcher = vi.fn<typeof fetch>(async () => {
      throw new Error('connection lost after write')
    })
    const { gateway, ledger } = setup(fetcher)
    await expect(
      gateway.dispatch(context, request, new AbortController().signal)
    ).rejects.toMatchObject({ code: 'outcome-unknown' })
    await expect(
      gateway.dispatch(context, request, new AbortController().signal)
    ).rejects.toMatchObject({ code: 'outcome-unknown' })
    expect(fetcher).toHaveBeenCalledOnce()
    ledger.close()
  })

  it('rejects an oversized declared provider response before buffering it', async () => {
    const fetcher = vi.fn<typeof fetch>(async () =>
      new Response('{}', {
        status: 200,
        headers: {
          'content-length': String(REMOTE_MODEL_GATEWAY_LIMITS.maximumResponseBodyBytes + 1)
        }
      })
    )
    const { gateway, ledger } = setup(fetcher)
    await expect(
      gateway.dispatch(context, request, new AbortController().signal)
    ).rejects.toMatchObject({ code: 'response-too-large' })
    ledger.close()
  })

  it('marks a truncated provider response outcome-unknown', async () => {
    const fetcher = vi.fn<typeof fetch>(async () =>
      new Response('{}', {
        status: 200,
        headers: { 'content-length': '100' }
      })
    )
    const { gateway, ledger } = setup(fetcher)
    await expect(
      gateway.dispatch(context, request, new AbortController().signal)
    ).rejects.toMatchObject({ code: 'outcome-unknown' })
    ledger.close()
  })
})

describe('AgentModelCallLedger', () => {
  it('prunes delivered terminal calls while the daemon remains running', () => {
    const root = mkdtempSync(join(tmpdir(), 'goodbuddy-agent-ledger-'))
    temporary.push(root)
    const ledger = new AgentModelCallLedger(
      join(root, 'calls.sqlite'),
      { maximumRetainedTerminalCalls: 2 }
    )
    for (let index = 1; index <= 3; index += 1) {
      const callId = `call-${index}`
      ledger.claim({
        callId,
        bindingId: 'binding-1',
        operationId: `operation-${index}`,
        promptSequence: index,
        roundIndex: 0,
        profileDigest: `sha256:${'a'.repeat(64)}`,
        requestDigest: `sha256:${String(index).repeat(64)}`
      })
      ledger.complete(callId)
      ledger.delivered(callId)
    }

    expect(ledger.get('call-1')).toBeUndefined()
    expect(ledger.get('call-2')).toMatchObject({ status: 'completed' })
    expect(ledger.get('call-3')).toMatchObject({ status: 'completed' })
    ledger.close()
  })
})