query-orchestration/internal/bot/fastapi.go

package bot

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"

	"queryorchestration/internal/serviceconfig/chatbot"
)

// fastAPIMaxAttempts is the hard upper bound on Chat retry attempts per
// plan §7 ("Up to 3 attempts"). Exposed as a constant so the retry
// matrix tests reference a single source of truth.
const fastAPIMaxAttempts = 3

// fastAPIChatPath is the FastAPI agent service's chat endpoint per the
// swagger. Concatenated to cfg.ServiceURL at request time.
const fastAPIChatPath = "/agent/chat"

// fastAPIHealthPath is the standard health probe path. Plan §7 uses it
// for the boot-time smoke check.
const fastAPIHealthPath = "/health"

// apiKeyHeader is the literal header name plan §11 mandates. The value
// is whatever cfg.APIKey carries verbatim — including the JSON-encoded
// string the live deployment expects. The constant is the header name
// only; the credential lives in cfg.APIKey.
const apiKeyHeader = "X-API-Key" //nolint:gosec // header name, not a credential

// FastAPIClient is the queryAPI-side FastAPI agent service client. Plan
// §7 specifies the request/response mapping and retry policy this struct
// implements.
type FastAPIClient struct {
	cfg        chatbot.ChatbotConfig
	baseURL    string
	httpClient *http.Client
}

// FastAPIOption mutates a FastAPIClient at construction time. Used by
// tests to override the baseURL (httptest.Server.URL) or substitute a
// custom http.Client (e.g., to inject a transport).
type FastAPIOption func(*FastAPIClient)

// WithBaseURL overrides cfg.ServiceURL. Tests set this to the
// httptest.Server.URL so requests dial the local stub.
func WithBaseURL(url string) FastAPIOption {
	return func(c *FastAPIClient) {
		c.baseURL = url
	}
}

// WithHTTPClient substitutes the default http.Client. Tests use this
// when they need full control over the round-tripper.
func WithHTTPClient(client *http.Client) FastAPIOption {
	return func(c *FastAPIClient) {
		if client != nil {
			c.httpClient = client
		}
	}
}

// NewFastAPIClient constructs a FastAPIClient. The default http.Client
// has no per-call timeout because Chat applies a per-attempt deadline
// derived from cfg.RequestTimeoutSeconds via context.WithTimeout, and
// imposing a Client.Timeout on top would interfere with the parent
// context cancellation tests.
func NewFastAPIClient(cfg chatbot.ChatbotConfig, opts ...FastAPIOption) *FastAPIClient {
	c := &FastAPIClient{
		cfg:        cfg,
		baseURL:    strings.TrimRight(cfg.ServiceURL, "/"),
		httpClient: &http.Client{},
	}
	for _, opt := range opts {
		opt(c)
	}
	c.baseURL = strings.TrimRight(c.baseURL, "/")
	return c
}

// Chat sends one chat request to FastAPI with retry. Plan §7 retry
// policy:
//   - Up to 3 attempts.
//   - Retry HTTP 429, 500, 502, 503, 504, and transport timeouts.
//   - Do not retry other 4xx responses.
//   - Do not retry HTTP 200 with status:"error".
//   - Reuse the same request_id for every attempt (caller responsibility).
//   - One parent context enforces total wall time.
//
// Returns ChatCallResult on success. On failure returns an attemptError
// implementing AttemptCounter so the caller can extract the count via
// errors.As(err, &counter).
func (c *FastAPIClient) Chat(ctx context.Context, req ChatRequest) (ChatCallResult, error) {
	body, err := json.Marshal(req)
	if err != nil {
		return ChatCallResult{}, &attemptError{
			sentinel: ErrAgentInvalidResponse,
			attempts: 0,
			cause:    fmt.Errorf("marshal request: %w", err),
		}
	}

	url := c.baseURL + fastAPIChatPath
	timeout := time.Duration(c.cfg.RequestTimeoutSeconds) * time.Second
	sentAt := time.Now()
	var lastErr error
	var lastResponseBody []byte
	var lastStatus int

	for attempt := 1; attempt <= fastAPIMaxAttempts; attempt++ {
		// Bail before the next attempt when the parent ctx has expired.
		// errors.Is preserves context.Canceled / context.DeadlineExceeded
		// so the test suite can assert on the wrapped sentinel.
		if ctxErr := ctx.Err(); ctxErr != nil {
			slog.Warn("bot.fastapi.context_canceled",
				"request_id", req.RequestID,
				"attempts_completed", attempt-1,
				"cause", ctxErr.Error(),
			)
			return ChatCallResult{}, &attemptError{
				sentinel: nil,
				attempts: attempt - 1,
				cause:    ctxErr,
			}
		}

		attemptStart := time.Now()
		result, err := c.attemptChat(ctx, url, body, timeout, sentAt, attempt)
		latencyMs := time.Since(attemptStart).Milliseconds()
		if err == nil {
			slog.Debug("bot.fastapi.attempt",
				"request_id", req.RequestID,
				"attempt", attempt,
				"status_code", http.StatusOK,
				"latency_ms", latencyMs,
				"outcome", "success",
			)
			return result, nil
		}

		// The attempt classified its own outcome; pull the classification
		// off the error so the loop can decide whether to retry.
		var classified *attemptError
		if !errors.As(err, &classified) {
			// Defensive: every attempt path returns *attemptError. If
			// that invariant ever breaks, surface the raw error with
			// the current attempt count rather than panicking.
			slog.Error("bot.fastapi.attempt",
				"request_id", req.RequestID,
				"attempt", attempt,
				"latency_ms", latencyMs,
				"outcome", "nonretryable_failure",
				"error", err.Error(),
			)
			return ChatCallResult{}, &attemptError{
				sentinel: nil,
				attempts: attempt,
				cause:    err,
			}
		}

		logAttemptOutcome(req.RequestID, attempt, latencyMs, classified)

		if !classified.retryable {
			return ChatCallResult{}, classified
		}

		lastErr = classified
		lastResponseBody = classified.responseBody
		lastStatus = classified.responseStatus
	}

	// All attempts exhausted as retryable failures.
	lastErrString := ""
	if lastErr != nil {
		lastErrString = lastErr.Error()
	}
	slog.Error("bot.fastapi.exhausted",
		"request_id", req.RequestID,
		"attempts", fastAPIMaxAttempts,
		"last_status", lastStatus,
		"last_error", lastErrString,
	)
	if lastErr != nil {
		return ChatCallResult{}, &attemptError{
			sentinel:       nil,
			attempts:       fastAPIMaxAttempts,
			cause:          lastErr,
			responseBody:   lastResponseBody,
			responseStatus: lastStatus,
		}
	}
	return ChatCallResult{}, &attemptError{
		sentinel: nil,
		attempts: fastAPIMaxAttempts,
		cause:    errors.New("bot: chat retries exhausted"),
	}
}

// logAttemptOutcome emits a single bot.fastapi.attempt slog record
// summarizing one *attemptError. WARN for retryable failures, ERROR
// otherwise. status_code is omitted when the failure occurred before
// receiving an HTTP response (transport error, parent ctx canceled).
func logAttemptOutcome(requestID string, attempt int, latencyMs int64, classified *attemptError) {
	outcome := "retryable_failure"
	if !classified.retryable {
		outcome = "nonretryable_failure"
	}
	cause := ""
	if classified.cause != nil {
		cause = classified.cause.Error()
	}
	attrs := []any{
		"request_id", requestID,
		"attempt", attempt,
		"latency_ms", latencyMs,
		"outcome", outcome,
		"error", cause,
	}
	if classified.responseStatus != 0 {
		attrs = append(attrs, "status_code", classified.responseStatus)
	}
	if classified.retryable {
		slog.Warn("bot.fastapi.attempt", attrs...)
		return
	}
	slog.Error("bot.fastapi.attempt", attrs...)
}

// attemptChat performs one HTTP attempt and classifies the outcome.
// Returns nil err + ChatCallResult on success; otherwise returns a
// non-nil *attemptError carrying the retryable flag.
func (c *FastAPIClient) attemptChat(parent context.Context, url string, body []byte, timeout time.Duration, sentAt time.Time, attempt int) (ChatCallResult, error) {
	attemptCtx, cancel := context.WithTimeout(parent, timeout)
	defer cancel()

	httpReq, err := http.NewRequestWithContext(attemptCtx, http.MethodPost, url, bytes.NewReader(body))
	if err != nil {
		return ChatCallResult{}, &attemptError{
			sentinel:  ErrAgentInvalidResponse,
			attempts:  attempt,
			cause:     err,
			retryable: false,
		}
	}
	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set(apiKeyHeader, c.cfg.APIKey)

	resp, err := c.httpClient.Do(httpReq)
	receivedAt := time.Now()
	if err != nil {
		// Distinguish parent ctx cancellation from a transient network
		// failure. Parent-ctx-canceled aborts the loop; a per-attempt
		// timeout (attemptCtx done while parent still alive) is a
		// retryable transport failure per plan §7.
		if parent.Err() != nil {
			return ChatCallResult{}, &attemptError{
				sentinel:  nil,
				attempts:  attempt,
				cause:     parent.Err(),
				retryable: false,
			}
		}
		return ChatCallResult{}, &attemptError{
			sentinel:  nil,
			attempts:  attempt,
			cause:     err,
			retryable: true,
		}
	}
	defer resp.Body.Close()

	respBody, readErr := io.ReadAll(resp.Body)
	if readErr != nil {
		return ChatCallResult{}, &attemptError{
			sentinel:  nil,
			attempts:  attempt,
			cause:     fmt.Errorf("read response body: %w", readErr),
			retryable: true,
		}
	}

	if resp.StatusCode != http.StatusOK {
		return ChatCallResult{}, classifyNon200(resp.StatusCode, respBody, attempt)
	}

	return classify200(respBody, sentAt, receivedAt, attempt)
}

// classifyNon200 builds the *attemptError for a non-2xx response. The
// returned error is always *attemptError; the caller wraps it in a
// dummy non-nil error sentinel so the calling code path treats it as
// "attempt failed".
func classifyNon200(status int, body []byte, attempt int) *attemptError {
	retryable := status == http.StatusTooManyRequests ||
		status == http.StatusInternalServerError ||
		status == http.StatusBadGateway ||
		status == http.StatusServiceUnavailable ||
		status == http.StatusGatewayTimeout
	return &attemptError{
		sentinel:       nil,
		attempts:       attempt,
		cause:          fmt.Errorf("FastAPI returned status %d: %s", status, truncateBody(body)),
		responseStatus: status,
		responseBody:   body,
		retryable:      retryable,
	}
}

// classify200 parses a 200 response body and applies plan §7's
// success-path validation rules. Returns ChatCallResult + nil on
// success; otherwise *attemptError with retryable=false (plan §7
// mandates that all 200-classified failures are not retried).
func classify200(body []byte, sentAt, receivedAt time.Time, attempt int) (ChatCallResult, error) {
	var resp ChatResponse
	if err := json.Unmarshal(body, &resp); err != nil {
		return ChatCallResult{}, &attemptError{
			sentinel:       ErrAgentInvalidResponse,
			attempts:       attempt,
			cause:          fmt.Errorf("parse response: %w", err),
			responseStatus: http.StatusOK,
			responseBody:   body,
			retryable:      false,
		}
	}

	if resp.Status == "error" {
		code, message := errorPayloadFields(resp.Error)
		return ChatCallResult{}, &attemptError{
			sentinel:       ErrAgentApplicationError,
			attempts:       attempt,
			cause:          fmt.Errorf("FastAPI status=error: code=%s message=%s", code, message),
			responseStatus: http.StatusOK,
			responseBody:   body,
			retryable:      false,
		}
	}

	if resp.Answer == nil || strings.TrimSpace(resp.Answer.Text) == "" {
		return ChatCallResult{}, &attemptError{
			sentinel:       ErrAgentMissingAnswer,
			attempts:       attempt,
			cause:          errors.New("FastAPI 200 response has empty/null answer.text"),
			responseStatus: http.StatusOK,
			responseBody:   body,
			retryable:      false,
		}
	}

	return ChatCallResult{
		Response:     &resp,
		SentAt:       sentAt,
		ReceivedAt:   receivedAt,
		AttemptCount: attempt,
	}, nil
}

// errorPayloadFields safely extracts (code, message) from a *ErrorPayload
// pointer that may be nil. Used so the wrapped error message always
// surfaces something concrete instead of a Go nil-pointer panic.
func errorPayloadFields(p *ErrorPayload) (string, string) {
	if p == nil {
		return "", ""
	}
	return p.Code, p.Message
}

// truncateBody returns at most 256 bytes of the supplied response body
// for inclusion in error messages. Avoids dumping multi-megabyte HTML
// error pages into logs.
func truncateBody(b []byte) string {
	const max = 256
	if len(b) <= max {
		return string(b)
	}
	return string(b[:max]) + "..."
}

// Health performs the /health smoke-check defined in plan §7. Returns
// nil on success; ErrHealthStatusNotOK or ErrHealthMissingService on
// the documented failure modes; a wrapped non-sentinel error on
// transport / parse failures.
func (c *FastAPIClient) Health(ctx context.Context) error {
	url := c.baseURL + fastAPIHealthPath
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return fmt.Errorf("bot.Health: build request: %w", err)
	}
	req.Header.Set(apiKeyHeader, c.cfg.APIKey)

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("bot.Health: do: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("bot.Health: read body: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("bot.Health: status %d: %s", resp.StatusCode, truncateBody(body))
	}

	var health HealthResponse
	if err := json.Unmarshal(body, &health); err != nil {
		return fmt.Errorf("bot.Health: parse body: %w", err)
	}
	if health.Status != "ok" {
		return fmt.Errorf("%w: got %q", ErrHealthStatusNotOK, health.Status)
	}
	if health.Service == "" || health.Version == "" {
		return fmt.Errorf("%w: service=%q version=%q", ErrHealthMissingService, health.Service, health.Version)
	}
	return nil
}

// attemptError is the single error type Chat surfaces. It implements
// AttemptCounter and supports errors.Is on the wrapped sentinel + the
// underlying cause so callers can switch on either layer.
type attemptError struct {
	sentinel       error
	attempts       int
	cause          error
	responseStatus int
	responseBody   []byte
	retryable      bool
}

// Error renders a single-line summary that includes the attempt count
// and either the sentinel name or the cause message — whichever is
// available.
func (e *attemptError) Error() string {
	parts := []string{fmt.Sprintf("attempts=%d", e.attempts)}
	if e.sentinel != nil {
		parts = append(parts, e.sentinel.Error())
	}
	if e.cause != nil {
		parts = append(parts, e.cause.Error())
	}
	if e.responseStatus != 0 {
		parts = append(parts, fmt.Sprintf("status=%d", e.responseStatus))
	}
	return "bot.Chat: " + strings.Join(parts, "; ")
}

// Is supports errors.Is(err, sentinel) directly so callers do not have
// to navigate the cause chain manually. The custom Is method also
// dispatches to the wrapped cause so errors.Is(err, context.Canceled)
// works without a separate Unwrap method on this type.
func (e *attemptError) Is(target error) bool {
	if e.sentinel != nil && errors.Is(e.sentinel, target) {
		return true
	}
	if e.cause != nil && errors.Is(e.cause, target) {
		return true
	}
	return false
}

// AttemptCount satisfies the AttemptCounter interface so the service
// layer can persist this scalar in bot_turns.error per plan §6.
func (e *attemptError) AttemptCount() int {
	return e.attempts
}