FusionAGI/tests/load/k6_prompt.js

/**
 * k6 load test for FusionAGI prompt endpoint.
 *
 * Run:
 *   k6 run tests/load/k6_prompt.js
 *
 * Options:
 *   k6 run --vus 10 --duration 30s tests/load/k6_prompt.js
 *   k6 run --vus 50 --duration 2m tests/load/k6_prompt.js
 *
 * Requires:
 *   - FusionAGI API running at http://localhost:8000
 *   - k6 installed (https://k6.io/docs/getting-started/installation/)
 */

import http from 'k6/http'
import { check, sleep } from 'k6'
import { Rate, Trend } from 'k6/metrics'

// Custom metrics
const errorRate = new Rate('errors')
const promptDuration = new Trend('prompt_duration', true)
const sessionDuration = new Trend('session_duration', true)

// Test configuration
export const options = {
  stages: [
    { duration: '10s', target: 5 },    // ramp up
    { duration: '30s', target: 10 },   // steady
    { duration: '10s', target: 20 },   // spike
    { duration: '10s', target: 0 },    // ramp down
  ],
  thresholds: {
    http_req_duration: ['p(95)<5000'],  // 95% under 5s
    errors: ['rate<0.1'],              // <10% error rate
  },
}

const BASE_URL = __ENV.API_URL || 'http://localhost:8000'
const API_KEY = __ENV.API_KEY || ''

const PROMPTS = [
  'Explain the concept of recursion',
  'What are the benefits of microservices?',
  'Design a rate limiter',
  'Compare SQL and NoSQL databases',
  'Explain the CAP theorem',
  'What is eventual consistency?',
  'How does garbage collection work?',
  'Explain WebSocket vs HTTP polling',
]

function getHeaders() {
  const headers = { 'Content-Type': 'application/json' }
  if (API_KEY) {
    headers['Authorization'] = `Bearer ${API_KEY}`
  }
  return headers
}

export default function () {
  const headers = getHeaders()

  // 1. Create session
  const sessionStart = Date.now()
  const sessionRes = http.post(`${BASE_URL}/v1/sessions`, null, { headers })
  sessionDuration.add(Date.now() - sessionStart)

  const sessionOk = check(sessionRes, {
    'session created': (r) => r.status === 200 || r.status === 201,
    'session has id': (r) => {
      try { return !!JSON.parse(r.body).session_id } catch { return false }
    },
  })

  if (!sessionOk) {
    errorRate.add(1)
    sleep(1)
    return
  }

  const sessionId = JSON.parse(sessionRes.body).session_id
  const prompt = PROMPTS[Math.floor(Math.random() * PROMPTS.length)]

  // 2. Send prompt
  const promptStart = Date.now()
  const promptRes = http.post(
    `${BASE_URL}/v1/sessions/${sessionId}/prompt`,
    JSON.stringify({ prompt }),
    { headers, timeout: '30s' },
  )
  promptDuration.add(Date.now() - promptStart)

  const promptOk = check(promptRes, {
    'prompt success': (r) => r.status === 200,
    'has final_answer': (r) => {
      try { return !!JSON.parse(r.body).final_answer } catch { return false }
    },
  })

  if (!promptOk) {
    errorRate.add(1)
  }

  // 3. Health check
  const healthRes = http.get(`${BASE_URL}/health`, { headers })
  check(healthRes, {
    'health ok': (r) => r.status === 200,
  })

  sleep(0.5 + Math.random())
}

export function handleSummary(data) {
  return {
    stdout: JSON.stringify({
      total_requests: data.metrics.http_reqs.values.count,
      avg_duration_ms: Math.round(data.metrics.http_req_duration.values.avg),
      p95_duration_ms: Math.round(data.metrics.http_req_duration.values['p(95)']),
      error_rate: data.metrics.errors ? data.metrics.errors.values.rate : 0,
      avg_prompt_ms: data.metrics.prompt_duration ? Math.round(data.metrics.prompt_duration.values.avg) : 0,
    }, null, 2),
  }
}