story/internal/agenteval/registrar.go

package agenteval

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

type Registrar struct {
	baseURL string
	token   string
	client  *http.Client
}

type RegistrationResult struct {
	StatusCode   int
	Replayed     bool
	RequestID    string
	EvaluationID string
}

func ValidateArtifact(suite Suite, run ObservationRun, artifact Artifact, allowFailed bool) error {
	if err := validateSuite(suite); err != nil {
		return err
	}
	rebuiltReport, err := Evaluate(suite, run, artifact.Report.EvaluatedAt)
	if err != nil {
		return fmt.Errorf("rebuild report from observations: %w", err)
	}
	rebuiltJSON, _ := json.Marshal(rebuiltReport)
	actualReportJSON, _ := json.Marshal(artifact.Report)
	if !bytes.Equal(rebuiltJSON, actualReportJSON) {
		return errors.New("report does not match the supplied observations")
	}
	digest, err := SuiteSHA256(suite)
	if err != nil {
		return err
	}
	report := artifact.Report
	if report.Schema != ReportSchema || report.SuiteVersion != suite.SuiteVersion || report.ModelConfigRef != suite.ModelConfigRef || report.SuiteSHA256 != digest {
		return errors.New("report does not match the supplied evaluation suite")
	}
	if report.Result != "PASSED" && report.Result != "FAILED" {
		return errors.New("report result must be PASSED or FAILED")
	}
	if report.Result == "FAILED" && !allowFailed {
		return errors.New("FAILED report registration requires explicit allow-failed")
	}
	if report.Metrics.CasesTotal != len(suite.Cases) || report.Metrics.CasesPassed < 0 || report.Metrics.CasesPassed > report.Metrics.CasesTotal {
		return errors.New("report case counts do not match the evaluation suite")
	}
	if report.Result == "PASSED" && (len(report.FailedChecks) != 0 || report.Metrics.CasesPassed != report.Metrics.CasesTotal) {
		return errors.New("PASSED report contains failed checks or incomplete cases")
	}
	if report.Result == "FAILED" && len(report.FailedChecks) == 0 {
		return errors.New("FAILED report does not contain failure evidence")
	}
	expected, err := NewArtifact(report, artifact.RegistrationPayload.ReportRef)
	if err != nil {
		return err
	}
	expectedJSON, _ := json.Marshal(expected.RegistrationPayload)
	actualJSON, _ := json.Marshal(artifact.RegistrationPayload)
	if !bytes.Equal(expectedJSON, actualJSON) {
		return errors.New("registration_payload does not match report contents")
	}
	return nil
}

func RegistrationIdempotencyKey(configID string, payload RegistrationPayload) (string, error) {
	configID = strings.TrimSpace(configID)
	if configID == "" {
		return "", errors.New("configuration id is required")
	}
	payload = BindRegistrationPayload(configID, payload)
	encoded, err := json.Marshal(payload)
	if err != nil {
		return "", fmt.Errorf("encode registration payload: %w", err)
	}
	digest := sha256.Sum256(append(append([]byte(configID), '\n'), encoded...))
	return "agent-eval-" + hex.EncodeToString(digest[:16]), nil
}

func BindRegistrationPayload(configID string, payload RegistrationPayload) RegistrationPayload {
	metrics := make(map[string]any, len(payload.Metrics)+1)
	for key, value := range payload.Metrics {
		metrics[key] = value
	}
	metrics["config_version_id"] = strings.TrimSpace(configID)
	payload.Metrics = metrics
	return payload
}

func NewRegistrar(baseURL, accessToken string, client *http.Client) (*Registrar, error) {
	baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
	accessToken = strings.TrimSpace(accessToken)
	parsed, err := url.Parse(baseURL)
	if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" || accessToken == "" {
		return nil, errors.New("valid API base URL and access token are required")
	}
	if client == nil {
		client = &http.Client{Timeout: 15 * time.Second}
	}
	return &Registrar{baseURL: baseURL, token: accessToken, client: client}, nil
}

func (r *Registrar) Register(ctx context.Context, tenantID, configID string, payload RegistrationPayload) (RegistrationResult, error) {
	tenantID, configID = strings.TrimSpace(tenantID), strings.TrimSpace(configID)
	if tenantID == "" || configID == "" {
		return RegistrationResult{}, errors.New("tenant and configuration ids are required")
	}
	payload = BindRegistrationPayload(configID, payload)
	encoded, err := json.Marshal(payload)
	if err != nil {
		return RegistrationResult{}, fmt.Errorf("encode registration payload: %w", err)
	}
	key, err := RegistrationIdempotencyKey(configID, payload)
	if err != nil {
		return RegistrationResult{}, err
	}
	endpoint := r.baseURL + "/api/v1/ops/config-versions/" + url.PathEscape(configID) + "/evaluations"
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, bytes.NewReader(encoded))
	if err != nil {
		return RegistrationResult{}, err
	}
	request.Header.Set("Authorization", "Bearer "+r.token)
	request.Header.Set("Content-Type", "application/json")
	request.Header.Set("Accept", "application/json")
	request.Header.Set("X-Tenant-Id", tenantID)
	request.Header.Set("Idempotency-Key", key)
	response, err := r.client.Do(request)
	if err != nil {
		return RegistrationResult{}, err
	}
	defer response.Body.Close()
	result := RegistrationResult{StatusCode: response.StatusCode, Replayed: response.Header.Get("Idempotency-Replayed") == "true", RequestID: response.Header.Get("X-Request-Id")}
	if response.StatusCode != http.StatusCreated && response.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, io.LimitReader(response.Body, 1<<20))
		return result, fmt.Errorf("evaluation registration returned HTTP %d", response.StatusCode)
	}
	var envelope struct {
		RequestID string `json:"request_id"`
		Data      struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	decoder := json.NewDecoder(io.LimitReader(response.Body, 1<<20))
	if err := decoder.Decode(&envelope); err != nil {
		return result, fmt.Errorf("decode registration response: %w", err)
	}
	if envelope.RequestID != "" {
		result.RequestID = envelope.RequestID
	}
	result.EvaluationID = envelope.Data.ID
	if result.EvaluationID == "" {
		return result, errors.New("registration response did not contain an evaluation id")
	}
	return result, nil
}