vecna/pkg/server/handler.go

package server

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	"go.uber.org/zap"

	"github.com/Warky-Devs/vecna.git/pkg/adapter"
	"github.com/Warky-Devs/vecna.git/pkg/config"
	"github.com/Warky-Devs/vecna.git/pkg/embedclient"
	"github.com/Warky-Devs/vecna.git/pkg/forwardproxy"
)

// ExtraMap pairs a dimension adapter with an optional forward-target override.
type ExtraMap struct {
	Adapter       adapter.Adapter
	ForwardTarget string // named target in forward.targets; empty = model-based resolution
}

// handler holds shared dependencies for all HTTP handlers.
type handler struct {
	cfg        *config.Config
	clients    map[string]embedclient.Client
	adapter    adapter.Adapter
	extraMaps  map[string]ExtraMap
	forwarders map[string]*forwardproxy.Router
	logger     *zap.Logger
}

// resolveExtraMap returns the ExtraMap for the named extra_map entry.
func (h *handler) resolveExtraMap(name string) (ExtraMap, error) {
	em, ok := h.extraMaps[name]
	if !ok {
		return ExtraMap{}, fmt.Errorf("extra_map %q not configured", name)
	}
	return em, nil
}

// resolveClient selects the embed client for the given model name.
// Returns the client, target name, and first endpoint URL for tracing.
func (h *handler) resolveClient(model string) (embedclient.Client, string, string) {
	if c, ok := h.clients[model]; ok {
		url := firstEndpointURL(h.cfg, model)
		return c, model, url
	}
	name := h.cfg.Forward.Default
	c, ok := h.clients[name]
	if !ok {
		return &errClient{err: fmt.Errorf("no client configured for model %q and no default", model)}, name, ""
	}
	return c, name, firstEndpointURL(h.cfg, name)
}

// resolveClientOverride selects the client for targetOverride when set,
// otherwise falls back to model-based resolution.
func (h *handler) resolveClientOverride(targetOverride, model string) (embedclient.Client, string, string) {
	if targetOverride == "" {
		return h.resolveClient(model)
	}
	c, ok := h.clients[targetOverride]
	if !ok {
		return &errClient{err: fmt.Errorf("extra_map forward_target %q not configured", targetOverride)}, targetOverride, ""
	}
	return c, targetOverride, firstEndpointURL(h.cfg, targetOverride)
}

// resolveForwarder selects the raw forwarder for the given model, falling back
// to forward.default when the model doesn't name a configured target directly.
func (h *handler) resolveForwarder(model string) (*forwardproxy.Router, string, error) {
	if f, ok := h.forwarders[model]; ok {
		return f, model, nil
	}
	name := h.cfg.Forward.Default
	f, ok := h.forwarders[name]
	if !ok {
		return nil, name, fmt.Errorf("no forward target configured for model %q and no default", model)
	}
	return f, name, nil
}

// resolveForwarderOverride selects the forwarder for targetOverride when set,
// otherwise falls back to model-based resolution.
func (h *handler) resolveForwarderOverride(targetOverride, model string) (*forwardproxy.Router, string, error) {
	if targetOverride == "" {
		return h.resolveForwarder(model)
	}
	f, ok := h.forwarders[targetOverride]
	if !ok {
		return nil, targetOverride, fmt.Errorf("extra_map forward_target %q not configured", targetOverride)
	}
	return f, targetOverride, nil
}

func firstEndpointURL(cfg *config.Config, targetName string) string {
	t, ok := cfg.Forward.Targets[targetName]
	if !ok || len(t.Endpoints) == 0 {
		return ""
	}
	return t.Endpoints[0].URL
}

// writeJSON encodes v as JSON and writes it with the given status code.
func writeJSON(w http.ResponseWriter, status int, v interface{}) error {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	if err := json.NewEncoder(w).Encode(v); err != nil {
		return fmt.Errorf("writeJSON: %w", err)
	}
	return nil
}

// apiError is the OpenAI-compatible error envelope, so existing OpenAI SDKs
// (which read err.error.message) surface vecna's own errors correctly instead
// of failing to parse a flat {"error": "..."} string.
type apiError struct {
	Message string `json:"message"`
	Type    string `json:"type"`
	Code    int    `json:"code"`
}

type apiErrorEnvelope struct {
	Error apiError `json:"error"`
}

// writeError writes an OpenAI-compatible error response.
func writeError(w http.ResponseWriter, status int, message string) error {
	return writeJSON(w, status, apiErrorEnvelope{Error: apiError{
		Message: message,
		Type:    errorType(status),
		Code:    status,
	}})
}

// errorType maps an HTTP status to an OpenAI-style error type string.
func errorType(status int) string {
	switch status {
	case http.StatusUnauthorized:
		return "authentication_error"
	case http.StatusBadRequest, http.StatusNotFound:
		return "invalid_request_error"
	case http.StatusBadGateway:
		return "upstream_error"
	default:
		if status >= 500 {
			return "internal_error"
		}
		return "api_error"
	}
}

// writeTraceHeaders writes X-Vecna-* timing headers from the RequestTrace.
func writeTraceHeaders(w http.ResponseWriter, t *RequestTrace) {
	total := time.Since(t.Start)
	w.Header().Set("X-Vecna-Forward-Ms", fmt.Sprintf("%d", t.ForwardDuration.Milliseconds()))
	w.Header().Set("X-Vecna-Translate-Ms", fmt.Sprintf("%d", t.TranslateDuration.Milliseconds()))
	w.Header().Set("X-Vecna-Total-Ms", fmt.Sprintf("%d", total.Milliseconds()))
}

// errClient is a Client that always returns a fixed error (used as safe fallback).
type errClient struct {
	err error
}

func (e *errClient) Embed(_ context.Context, _ embedclient.Request) (embedclient.Response, error) {
	return embedclient.Response{}, e.err
}