ai-gateway-go/internal/trace/store.go

package trace

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"
	"unicode/utf8"

	platformid "aigateway.local/core/internal/platform/id"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

var ErrNotFound = errors.New("trace not found")

type Store struct{ pool *pgxpool.Pool }

func NewStore(pool *pgxpool.Pool) *Store { return &Store{pool: pool} }

type StartInput struct {
	RequestID      string
	APIKeyID       string
	TenantID       *string
	TraceType      string
	TargetID       string
	TargetCode     string
	ConversationID string
	Metadata       map[string]any
}

type Trace struct {
	ID             string          `json:"id"`
	RequestID      string          `json:"request_id"`
	APIKeyID       *string         `json:"api_key_id,omitempty"`
	TenantID       *string         `json:"tenant_id,omitempty"`
	TraceType      string          `json:"trace_type"`
	TargetID       *string         `json:"target_id,omitempty"`
	TargetCode     string          `json:"target_code"`
	ConversationID string          `json:"conversation_id"`
	Status         string          `json:"status"`
	StartedAt      time.Time       `json:"started_at"`
	FinishedAt     *time.Time      `json:"finished_at,omitempty"`
	LatencyMS      *int            `json:"latency_ms,omitempty"`
	RetrievalCount int             `json:"retrieval_count"`
	ModelCallCount int             `json:"model_call_count"`
	ToolCallCount  int             `json:"tool_call_count"`
	Error          string          `json:"error"`
	Metadata       json.RawMessage `json:"metadata"`
	Spans          []Span          `json:"spans,omitempty"`
}

type SpanInput struct {
	TraceID      string
	ParentID     string
	SpanType     string
	Name         string
	Round        int
	ProviderCode string
	Model        string
	Metadata     map[string]any
}

type Span struct {
	ID           string          `json:"id"`
	TraceID      string          `json:"trace_id"`
	ParentID     *string         `json:"parent_id,omitempty"`
	SpanType     string          `json:"span_type"`
	Name         string          `json:"name"`
	Status       string          `json:"status"`
	StartedAt    time.Time       `json:"started_at"`
	FinishedAt   *time.Time      `json:"finished_at,omitempty"`
	LatencyMS    *int            `json:"latency_ms,omitempty"`
	ProviderCode string          `json:"provider_code,omitempty"`
	Model        string          `json:"model,omitempty"`
	InputTokens  int64           `json:"input_tokens"`
	OutputTokens int64           `json:"output_tokens"`
	Round        int             `json:"round"`
	Error        string          `json:"error"`
	Metadata     json.RawMessage `json:"metadata"`
}

type FinishInput struct {
	Status         string
	Error          string
	RetrievalCount int
	ModelCallCount int
	ToolCallCount  int
	Metadata       map[string]any
}

type SpanFinishInput struct {
	Status       string
	Error        string
	InputTokens  int64
	OutputTokens int64
	ProviderCode string
	Model        string
	Metadata     map[string]any
}

type Filter struct {
	From       time.Time
	To         time.Time
	TraceType  string
	TargetCode string
	RequestID  string
	Status     string
	Limit      int
}

// Session is a metadata-only aggregation of traces that share a conversation
// ID. Stateless requests use a request-derived key so they remain visible in
// the session center without pretending to be part of a persistent chat.
type Session struct {
	ID             string    `json:"id"`
	TraceType      string    `json:"trace_type"`
	TargetCode     string    `json:"target_code"`
	TraceCount     int       `json:"trace_count"`
	LatestTraceID  string    `json:"latest_trace_id"`
	LatestStatus   string    `json:"latest_status"`
	StartedAt      time.Time `json:"started_at"`
	UpdatedAt      time.Time `json:"updated_at"`
	RetrievalCount int       `json:"retrieval_count"`
	ModelCallCount int       `json:"model_call_count"`
	ToolCallCount  int       `json:"tool_call_count"`
}

type SessionFilter struct {
	From       time.Time
	To         time.Time
	TraceType  string
	TargetCode string
	SessionID  string
	Limit      int
}

func metadataJSON(value map[string]any) []byte {
	if value == nil {
		return []byte(`{}`)
	}
	raw, err := json.Marshal(value)
	if err != nil {
		return []byte(`{}`)
	}
	return raw
}

func normalizeMetadata(raw []byte) json.RawMessage {
	if len(raw) == 0 || !json.Valid(raw) {
		return json.RawMessage(`{}`)
	}
	return raw
}

func validateStart(input StartInput) error {
	if strings.TrimSpace(input.RequestID) == "" || strings.TrimSpace(input.TargetCode) == "" {
		return errors.New("trace request_id 和 target_code 不能为空")
	}
	if input.TraceType != "application" && input.TraceType != "digital_employee" {
		return errors.New("trace 类型无效")
	}
	if len(input.TargetCode) > 128 || len(input.ConversationID) > 128 {
		return errors.New("trace 目标或会话 ID 过长")
	}
	return nil
}

func validateSpan(input SpanInput) error {
	if strings.TrimSpace(input.TraceID) == "" || strings.TrimSpace(input.Name) == "" {
		return errors.New("trace span 标识不能为空")
	}
	if input.SpanType != "model" && input.SpanType != "tool" && input.SpanType != "retrieval" {
		return errors.New("trace span 类型无效")
	}
	if input.Round < 0 || len(input.Name) > 256 {
		return errors.New("trace span 参数无效")
	}
	return nil
}

func (s *Store) Start(ctx context.Context, input StartInput) (Trace, error) {
	if s == nil || s.pool == nil {
		return Trace{}, errors.New("trace store unavailable")
	}
	input.RequestID = strings.TrimSpace(input.RequestID)
	input.TargetCode = strings.TrimSpace(input.TargetCode)
	input.ConversationID = strings.TrimSpace(input.ConversationID)
	if err := validateStart(input); err != nil {
		return Trace{}, err
	}
	id, err := platformid.NewUUID()
	if err != nil {
		return Trace{}, err
	}
	_, err = s.pool.Exec(ctx, `INSERT INTO gateway.agent_traces(id,request_id,api_key_id,tenant_id,trace_type,target_id,target_code,conversation_id,metadata) VALUES($1,$2,nullif($3,'')::uuid,nullif($4,'')::uuid,$5,nullif($6,'')::uuid,$7,$8,$9)`, id, input.RequestID, input.APIKeyID, valueOrEmpty(input.TenantID), input.TraceType, input.TargetID, input.TargetCode, input.ConversationID, metadataJSON(input.Metadata))
	if err != nil {
		return Trace{}, fmt.Errorf("start trace: %w", err)
	}
	return s.Get(ctx, id)
}

func valueOrEmpty(value *string) string {
	if value == nil {
		return ""
	}
	return *value
}

func (s *Store) StartSpan(ctx context.Context, input SpanInput) (Span, error) {
	if s == nil || s.pool == nil {
		return Span{}, errors.New("trace store unavailable")
	}
	input.Name = strings.TrimSpace(input.Name)
	if err := validateSpan(input); err != nil {
		return Span{}, err
	}
	id, err := platformid.NewUUID()
	if err != nil {
		return Span{}, err
	}
	_, err = s.pool.Exec(ctx, `INSERT INTO gateway.agent_trace_spans(id,trace_id,parent_id,span_type,name,round,provider_code,model,metadata) VALUES($1,$2,nullif($3,'')::uuid,$4,$5,$6,nullif($7,''),nullif($8,''),$9)`, id, input.TraceID, input.ParentID, input.SpanType, input.Name, input.Round, input.ProviderCode, input.Model, metadataJSON(input.Metadata))
	if err != nil {
		return Span{}, fmt.Errorf("start trace span: %w", err)
	}
	return s.GetSpan(ctx, id)
}

func (s *Store) Finish(ctx context.Context, id string, input FinishInput) error {
	if s == nil || s.pool == nil {
		return errors.New("trace store unavailable")
	}
	status := normalizeStatus(input.Status)
	errorText := truncate(input.Error, 4000)
	metadata := metadataJSON(input.Metadata)
	tag, err := s.pool.Exec(ctx, `UPDATE gateway.agent_traces SET status=$2,error=$3,retrieval_count=$4,model_call_count=$5,tool_call_count=$6,metadata=$7,finished_at=clock_timestamp(),latency_ms=(extract(epoch FROM (clock_timestamp()-started_at))*1000)::integer WHERE id=$1 AND status='running'`, id, status, errorText, max(input.RetrievalCount, 0), max(input.ModelCallCount, 0), max(input.ToolCallCount, 0), metadata)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return ErrNotFound
	}
	return nil
}

func (s *Store) FinishSpan(ctx context.Context, id string, input SpanFinishInput) error {
	if s == nil || s.pool == nil {
		return errors.New("trace store unavailable")
	}
	status := normalizeStatus(input.Status)
	tag, err := s.pool.Exec(ctx, `UPDATE gateway.agent_trace_spans SET status=$2,error=$3,input_tokens=$4,output_tokens=$5,provider_code=coalesce(nullif($6,''),provider_code),model=coalesce(nullif($7,''),model),metadata=$8,finished_at=clock_timestamp(),latency_ms=(extract(epoch FROM (clock_timestamp()-started_at))*1000)::integer WHERE id=$1 AND status='running'`, id, status, truncate(input.Error, 4000), max(input.InputTokens, 0), max(input.OutputTokens, 0), input.ProviderCode, input.Model, metadataJSON(input.Metadata))
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return ErrNotFound
	}
	return nil
}

func normalizeStatus(value string) string {
	if value == "success" {
		return "success"
	}
	return "error"
}

func truncate(value string, limit int) string {
	value = strings.TrimSpace(value)
	if len(value) <= limit {
		return value
	}
	cut := value[:limit]
	// 按字节截断可能切半多字节 rune,产生无效 UTF-8 使 trace 落库失败;
	// 回退到最近一个完整 rune 的边界。
	for len(cut) > 0 && !utf8.RuneStart(cut[len(cut)-1]) {
		cut = cut[:len(cut)-1]
	}
	return cut
}

const traceSelect = `SELECT id::text,request_id,api_key_id::text,tenant_id::text,trace_type,target_id::text,target_code,conversation_id,status,started_at,finished_at,latency_ms,retrieval_count,model_call_count,tool_call_count,error,metadata FROM gateway.agent_traces`

func scanTrace(row pgx.Row) (Trace, error) {
	var item Trace
	err := row.Scan(&item.ID, &item.RequestID, &item.APIKeyID, &item.TenantID, &item.TraceType, &item.TargetID, &item.TargetCode, &item.ConversationID, &item.Status, &item.StartedAt, &item.FinishedAt, &item.LatencyMS, &item.RetrievalCount, &item.ModelCallCount, &item.ToolCallCount, &item.Error, &item.Metadata)
	if errors.Is(err, pgx.ErrNoRows) {
		return Trace{}, ErrNotFound
	}
	item.Metadata = normalizeMetadata(item.Metadata)
	return item, err
}

func scanSpan(row pgx.Row) (Span, error) {
	var item Span
	err := row.Scan(&item.ID, &item.TraceID, &item.ParentID, &item.SpanType, &item.Name, &item.Status, &item.StartedAt, &item.FinishedAt, &item.LatencyMS, &item.ProviderCode, &item.Model, &item.InputTokens, &item.OutputTokens, &item.Round, &item.Error, &item.Metadata)
	if errors.Is(err, pgx.ErrNoRows) {
		return Span{}, ErrNotFound
	}
	item.Metadata = normalizeMetadata(item.Metadata)
	return item, err
}

func (s *Store) Get(ctx context.Context, id string) (Trace, error) {
	if s == nil || s.pool == nil {
		return Trace{}, errors.New("trace store unavailable")
	}
	item, err := scanTrace(s.pool.QueryRow(ctx, traceSelect+` WHERE id=$1`, id))
	if err != nil {
		return Trace{}, err
	}
	spans, err := s.listSpans(ctx, id)
	if err != nil {
		return Trace{}, err
	}
	item.Spans = spans
	return item, nil
}

func (s *Store) GetSpan(ctx context.Context, id string) (Span, error) {
	if s == nil || s.pool == nil {
		return Span{}, errors.New("trace store unavailable")
	}
	return scanSpan(s.pool.QueryRow(ctx, `SELECT id::text,trace_id::text,parent_id::text,span_type,name,status,started_at,finished_at,latency_ms,coalesce(provider_code,''),coalesce(model,''),input_tokens,output_tokens,round,error,metadata FROM gateway.agent_trace_spans WHERE id=$1`, id))
}

func (s *Store) listSpans(ctx context.Context, traceID string) ([]Span, error) {
	if s == nil || s.pool == nil {
		return nil, errors.New("trace store unavailable")
	}
	rows, err := s.pool.Query(ctx, `SELECT id::text,trace_id::text,parent_id::text,span_type,name,status,started_at,finished_at,latency_ms,coalesce(provider_code,''),coalesce(model,''),input_tokens,output_tokens,round,error,metadata FROM gateway.agent_trace_spans WHERE trace_id=$1 ORDER BY started_at,id`, traceID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []Span{}
	for rows.Next() {
		item, scanErr := scanSpan(rows)
		if scanErr != nil {
			return nil, scanErr
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

func (s *Store) List(ctx context.Context, filter Filter) ([]Trace, error) {
	if s == nil || s.pool == nil {
		return nil, errors.New("trace store unavailable")
	}
	if filter.Limit < 1 || filter.Limit > 200 {
		filter.Limit = 50
	}
	where := []string{"started_at >= $1", "started_at < $2"}
	args := []any{filter.From, filter.To}
	add := func(condition string, value any) {
		args = append(args, value)
		where = append(where, fmt.Sprintf(condition, len(args)))
	}
	if filter.TraceType != "" {
		add("trace_type = $%d", filter.TraceType)
	}
	if filter.TargetCode != "" {
		add("target_code = $%d", filter.TargetCode)
	}
	if filter.RequestID != "" {
		add("request_id = $%d", filter.RequestID)
	}
	if filter.Status != "" {
		add("status = $%d", filter.Status)
	}
	args = append(args, filter.Limit)
	query := traceSelect + ` WHERE ` + strings.Join(where, " AND ") + ` ORDER BY started_at DESC,id DESC LIMIT $` + strconv.Itoa(len(args))
	rows, err := s.pool.Query(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []Trace{}
	for rows.Next() {
		item, scanErr := scanTrace(rows)
		if scanErr != nil {
			return nil, scanErr
		}
		items = append(items, item)
	}
	return items, rows.Err()
}

const sessionGroupSelect = `
WITH grouped AS (
	SELECT trace_type,
		target_code,
		coalesce(nullif(conversation_id, ''), 'request:' || request_id) AS session_key,
		count(*)::int AS trace_count,
		(array_agg(id::text ORDER BY started_at DESC, id DESC))[1] AS latest_trace_id,
		(array_agg(status ORDER BY started_at DESC, id DESC))[1] AS latest_status,
		min(started_at) AS started_at,
		max(coalesce(finished_at, started_at)) AS updated_at,
		sum(retrieval_count)::int AS retrieval_count,
		sum(model_call_count)::int AS model_call_count,
		sum(tool_call_count)::int AS tool_call_count
	FROM gateway.agent_traces`

func scanSession(row pgx.Row) (Session, error) {
	var item Session
	err := row.Scan(&item.ID, &item.TraceType, &item.TargetCode, &item.TraceCount, &item.LatestTraceID, &item.LatestStatus, &item.StartedAt, &item.UpdatedAt, &item.RetrievalCount, &item.ModelCallCount, &item.ToolCallCount)
	return item, err
}

func (s *Store) ListSessions(ctx context.Context, filter SessionFilter) ([]Session, error) {
	if s == nil || s.pool == nil {
		return nil, errors.New("trace store unavailable")
	}
	if filter.Limit < 1 || filter.Limit > 200 {
		filter.Limit = 50
	}
	innerWhere := []string{"started_at >= $1", "started_at < $2"}
	args := []any{filter.From, filter.To}
	addInner := func(condition string, value any) {
		args = append(args, value)
		innerWhere = append(innerWhere, fmt.Sprintf(condition, len(args)))
	}
	if filter.TraceType != "" {
		addInner("trace_type = $%d", filter.TraceType)
	}
	if filter.TargetCode != "" {
		addInner("target_code = $%d", filter.TargetCode)
	}
	outerWhere := []string{}
	if filter.SessionID != "" {
		args = append(args, filter.SessionID)
		outerWhere = append(outerWhere, fmt.Sprintf("trace_type || ':' || target_code || ':' || session_key = $%d", len(args)))
	}
	args = append(args, filter.Limit)
	limitArg := strconv.Itoa(len(args))
	query := sessionGroupSelect + ` WHERE ` + strings.Join(innerWhere, " AND ") + ` GROUP BY trace_type,target_code,session_key) SELECT trace_type || ':' || target_code || ':' || session_key AS id,trace_type,target_code,trace_count,latest_trace_id,latest_status,started_at,updated_at,retrieval_count,model_call_count,tool_call_count FROM grouped`
	if len(outerWhere) > 0 {
		query += ` WHERE ` + strings.Join(outerWhere, " AND ")
	}
	query += ` ORDER BY updated_at DESC,id DESC LIMIT $` + limitArg
	rows, err := s.pool.Query(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []Session{}
	for rows.Next() {
		item, scanErr := scanSession(rows)
		if scanErr != nil {
			return nil, scanErr
		}
		items = append(items, item)
	}
	return items, rows.Err()
}