feat: enrich task record details

2026-05-10 22:33:58 +08:00
parent 205a4b625e
commit 53f8edfb67
19 changed files with 1781 additions and 165 deletions
@@ -41,18 +41,26 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger) *Service {
 }

 func (s *Service) Execute(ctx context.Context, task store.GatewayTask, user *auth.User) (Result, error) {
+	return s.execute(ctx, task, user, nil)
+}
+
+func (s *Service) ExecuteStream(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) {
+	return s.execute(ctx, task, user, onDelta)
+}
+
+func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) {
 	modelType := modelTypeFromKind(task.Kind)
 	body := normalizeRequest(task.Kind, task.Request)
 	if err := validateRequest(task.Kind, body); err != nil {
-		failed, finishErr := s.failTask(ctx, task.ID, "bad_request", err.Error(), task.RunMode == "simulation")
+		failed, finishErr := s.failTask(ctx, task.ID, "bad_request", err.Error(), task.RunMode == "simulation", err)
 		if finishErr != nil {
 			return Result{}, finishErr
 		}
 		return Result{Task: failed, Output: failed.Result}, err
 	}
-	candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType)
+	candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user)
 	if err != nil {
-		failed, finishErr := s.failTask(ctx, task.ID, "no_model_candidate", err.Error(), task.RunMode == "simulation")
+		failed, finishErr := s.failTask(ctx, task.ID, "no_model_candidate", err.Error(), task.RunMode == "simulation", err)
 		if finishErr != nil {
 			return Result{}, finishErr
 		}
@@ -72,14 +80,35 @@ func (s *Service) Execute(ctx context.Context, task store.GatewayTask, user *aut
 			break
 		}
 		attemptNo := index + 1
-		response, err := s.runCandidate(ctx, task, user, body, candidate, attemptNo)
+		response, err := s.runCandidate(ctx, task, user, body, candidate, attemptNo, onDelta)
 		if err == nil {
 			billings := s.billings(ctx, user, task.Kind, body, candidate, response, isSimulation(task, candidate))
-			finished, finishErr := s.store.FinishTaskSuccess(ctx, task.ID, response.Result, billings)
+			record := buildSuccessRecord(task, user, body, candidate, response, billings, isSimulation(task, candidate))
+			finished, finishErr := s.store.FinishTaskSuccess(ctx, store.FinishTaskSuccessInput{
+				TaskID:             task.ID,
+				Result:             response.Result,
+				Billings:           billings,
+				RequestID:          record.RequestID,
+				ResolvedModel:      record.ResolvedModel,
+				Usage:              record.Usage,
+				Metrics:            record.Metrics,
+				BillingSummary:     record.BillingSummary,
+				FinalChargeAmount:  record.FinalChargeAmount,
+				ResponseStartedAt:  record.ResponseStartedAt,
+				ResponseFinishedAt: record.ResponseFinishedAt,
+				ResponseDurationMS: record.ResponseDurationMS,
+			})
 			if finishErr != nil {
 				return Result{}, finishErr
 			}
-			if err := s.emit(ctx, task.ID, "task.completed", "succeeded", "completed", 1, "task completed", map[string]any{"result": response.Result, "billings": billings}, isSimulation(task, candidate)); err != nil {
+			if err := s.emit(ctx, task.ID, "task.completed", "succeeded", "completed", 1, "task completed", map[string]any{
+				"result":         response.Result,
+				"billings":       billings,
+				"usage":          record.Usage,
+				"metrics":        record.Metrics,
+				"billingSummary": record.BillingSummary,
+				"requestId":      record.RequestID,
+			}, isSimulation(task, candidate)); err != nil {
 				return Result{}, err
 			}
 			return Result{Task: finished, Output: response.Result}, nil
@@ -98,14 +127,14 @@ func (s *Service) Execute(ctx context.Context, task store.GatewayTask, user *aut
 	if lastErr != nil {
 		message = lastErr.Error()
 	}
-	failed, err := s.failTask(ctx, task.ID, code, message, task.RunMode == "simulation")
+	failed, err := s.failTask(ctx, task.ID, code, message, task.RunMode == "simulation", lastErr)
 	if err != nil {
 		return Result{}, err
 	}
 	return Result{Task: failed, Output: failed.Result}, lastErr
 }

-func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, candidate store.RuntimeModelCandidate, attemptNo int) (clients.Response, error) {
+func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta) (clients.Response, error) {
 	simulated := isSimulation(task, candidate)
 	if err := s.emit(ctx, task.ID, "task.attempt.started", "running", "submitting", 0.25, "client attempt started", map[string]any{"attempt": attemptNo, "clientId": candidate.ClientID}, simulated); err != nil {
 		return clients.Response{}, err
@@ -127,7 +156,14 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 	reservations := s.rateLimitReservations(ctx, user, candidate, body)
 	limitResult, err := s.store.ReserveRateLimits(ctx, task.ID, reservations)
 	if err != nil {
-		_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{AttemptID: attemptID, Status: "failed", Retryable: false, ErrorCode: "rate_limit", ErrorMessage: err.Error()})
+		_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
+			AttemptID:    attemptID,
+			Status:       "failed",
+			Retryable:    false,
+			Metrics:      map[string]any{"error": err.Error(), "candidateModel": candidate.ModelName, "clientId": candidate.ClientID},
+			ErrorCode:    "rate_limit",
+			ErrorMessage: err.Error(),
+		})
 		return clients.Response{}, &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: false}
 	}
 	defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.LeaseIDs)
@@ -138,34 +174,77 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 	defer s.store.RecordClientRelease(context.WithoutCancel(ctx), candidate.ClientID, "")

 	client := s.clientFor(candidate, simulated)
+	callStartedAt := time.Now()
 	response, err := client.Run(ctx, clients.Request{
-		Kind:      task.Kind,
-		ModelType: candidate.ModelType,
-		Model:     task.Model,
-		Body:      body,
-		Candidate: candidate,
-		Stream:    boolFromMap(body, "stream"),
+		Kind:        task.Kind,
+		ModelType:   candidate.ModelType,
+		Model:       task.Model,
+		Body:        body,
+		Candidate:   candidate,
+		Stream:      boolFromMap(body, "stream"),
+		StreamDelta: onDelta,
 	})
+	callFinishedAt := time.Now()
+	if response.ResponseStartedAt.IsZero() {
+		response.ResponseStartedAt = callStartedAt
+	}
+	if response.ResponseFinishedAt.IsZero() {
+		response.ResponseFinishedAt = callFinishedAt
+	}
+	if response.ResponseDurationMS == 0 {
+		response.ResponseDurationMS = response.ResponseFinishedAt.Sub(response.ResponseStartedAt).Milliseconds()
+		if response.ResponseDurationMS < 0 {
+			response.ResponseDurationMS = 0
+		}
+	}
 	if err != nil {
 		retryable := clients.IsRetryable(err)
+		requestID, metrics, responseStartedAt, responseFinishedAt, responseDurationMS := failureMetrics(err, simulated)
+		if responseStartedAt.IsZero() {
+			responseStartedAt = callStartedAt
+		}
+		if responseFinishedAt.IsZero() {
+			responseFinishedAt = callFinishedAt
+		}
+		if responseDurationMS == 0 {
+			responseDurationMS = responseFinishedAt.Sub(responseStartedAt).Milliseconds()
+			if responseDurationMS < 0 {
+				responseDurationMS = 0
+			}
+		}
 		_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
-			AttemptID:    attemptID,
-			Status:       "failed",
-			Retryable:    retryable,
-			ErrorCode:    clients.ErrorCode(err),
-			ErrorMessage: err.Error(),
+			AttemptID:          attemptID,
+			Status:             "failed",
+			Retryable:          retryable,
+			RequestID:          requestID,
+			Metrics:            metrics,
+			ResponseStartedAt:  responseStartedAt,
+			ResponseFinishedAt: responseFinishedAt,
+			ResponseDurationMS: responseDurationMS,
+			ErrorCode:          clients.ErrorCode(err),
+			ErrorMessage:       err.Error(),
 		})
-		_ = s.emit(ctx, task.ID, "task.attempt.failed", "running", "attempt_failed", 0.45, err.Error(), map[string]any{"attempt": attemptNo, "retryable": retryable}, simulated)
+		_ = s.emit(ctx, task.ID, "task.attempt.failed", "running", "attempt_failed", 0.45, err.Error(), map[string]any{"attempt": attemptNo, "retryable": retryable, "requestId": requestID, "metrics": metrics}, simulated)
 		return clients.Response{}, err
 	}
 	uploadedResult, err := s.uploadGeneratedAssets(ctx, response.Result)
 	if err != nil {
+		metrics := taskMetrics(task, user, body, candidate, response, simulated)
+		metrics["error"] = err.Error()
+		metrics["retryable"] = clients.IsRetryable(err)
 		_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
-			AttemptID:    attemptID,
-			Status:       "failed",
-			Retryable:    clients.IsRetryable(err),
-			ErrorCode:    clients.ErrorCode(err),
-			ErrorMessage: err.Error(),
+			AttemptID:          attemptID,
+			Status:             "failed",
+			Retryable:          clients.IsRetryable(err),
+			RequestID:          response.RequestID,
+			Usage:              usageToMap(response.Usage),
+			Metrics:            metrics,
+			ResponseSnapshot:   response.Result,
+			ResponseStartedAt:  response.ResponseStartedAt,
+			ResponseFinishedAt: response.ResponseFinishedAt,
+			ResponseDurationMS: response.ResponseDurationMS,
+			ErrorCode:          clients.ErrorCode(err),
+			ErrorMessage:       err.Error(),
 		})
 		return clients.Response{}, err
 	}
@@ -176,9 +255,15 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 		}
 	}
 	if err := s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
-		AttemptID:        attemptID,
-		Status:           "succeeded",
-		ResponseSnapshot: response.Result,
+		AttemptID:          attemptID,
+		Status:             "succeeded",
+		RequestID:          response.RequestID,
+		Usage:              usageToMap(response.Usage),
+		Metrics:            taskMetrics(task, user, body, candidate, response, simulated),
+		ResponseSnapshot:   response.Result,
+		ResponseStartedAt:  response.ResponseStartedAt,
+		ResponseFinishedAt: response.ResponseFinishedAt,
+		ResponseDurationMS: response.ResponseDurationMS,
 	}); err != nil {
 		return clients.Response{}, err
 	}
@@ -189,19 +274,32 @@ func (s *Service) clientFor(candidate store.RuntimeModelCandidate, simulated boo
 	if simulated {
 		return s.clients["simulation"]
 	}
-	key := strings.ToLower(candidate.Provider)
+	key := strings.ToLower(strings.TrimSpace(candidate.SpecType))
+	if key == "" {
+		key = strings.ToLower(strings.TrimSpace(candidate.Provider))
+	}
 	if client, ok := s.clients[key]; ok {
 		return client
 	}
 	return s.clients["openai"]
 }

-func (s *Service) failTask(ctx context.Context, taskID string, code string, message string, simulated bool) (store.GatewayTask, error) {
-	failed, err := s.store.FinishTaskFailure(ctx, taskID, code, message)
+func (s *Service) failTask(ctx context.Context, taskID string, code string, message string, simulated bool, cause error) (store.GatewayTask, error) {
+	requestID, metrics, responseStartedAt, responseFinishedAt, responseDurationMS := failureMetrics(cause, simulated)
+	failed, err := s.store.FinishTaskFailure(ctx, store.FinishTaskFailureInput{
+		TaskID:             taskID,
+		Code:               code,
+		Message:            message,
+		RequestID:          requestID,
+		Metrics:            metrics,
+		ResponseStartedAt:  responseStartedAt,
+		ResponseFinishedAt: responseFinishedAt,
+		ResponseDurationMS: responseDurationMS,
+	})
 	if err != nil {
 		return store.GatewayTask{}, err
 	}
-	if eventErr := s.emit(ctx, taskID, "task.failed", "failed", "failed", 1, message, map[string]any{"code": code}, simulated); eventErr != nil {
+	if eventErr := s.emit(ctx, taskID, "task.failed", "failed", "failed", 1, message, map[string]any{"code": code, "requestId": requestID, "metrics": metrics}, simulated); eventErr != nil {
 		return store.GatewayTask{}, eventErr
 	}
 	return failed, nil
@@ -221,14 +319,23 @@ func (s *Service) emit(ctx context.Context, taskID string, eventType string, sta
 func modelTypeFromKind(kind string) string {
 	switch kind {
 	case "chat.completions", "responses":
-		return "chat"
+		return "text_generate"
 	case "images.generations", "images.edits":
-		return "image"
+		if kind == "images.edits" {
+			return "image_edit"
+		}
+		return "image_generate"
+	case "videos.generations":
+		return "video_generate"
 	default:
 		return "task"
 	}
 }

+func isTextGenerationKind(kind string) bool {
+	return kind == "chat.completions" || kind == "responses"
+}
+
 func isSimulation(task store.GatewayTask, candidate store.RuntimeModelCandidate) bool {
 	if task.RunMode == "simulation" {
 		return true