Add runtime restore and temp asset cleanup

This commit is contained in:
2026-06-05 20:28:26 +08:00
parent 644a6f9d17
commit d41d9482c7
22 changed files with 1837 additions and 50 deletions
+26 -9
View File
@@ -87,9 +87,13 @@ func (s *Service) ExecuteStream(ctx context.Context, task store.GatewayTask, use
func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta) (Result, error) {
executeStartedAt := time.Now()
body := normalizeRequest(task.Kind, task.Request)
restoredRequest, err := s.restoreTaskRequestReferences(ctx, task)
if err != nil {
return Result{}, err
}
body := normalizeRequest(task.Kind, restoredRequest)
modelType := modelTypeFromKind(task.Kind, body)
if err := s.store.MarkTaskRunning(ctx, task.ID, modelType, body); err != nil {
if err := s.store.MarkTaskRunning(ctx, task.ID, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
return Result{}, err
}
if task.Status != "running" {
@@ -194,7 +198,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
}
return Result{Task: failed, Output: failed.Result}, clientErr
}
if err := s.store.MarkTaskRunning(ctx, task.ID, candidates[0].ModelType, firstCandidateBody); err != nil {
if err := s.store.MarkTaskRunning(ctx, task.ID, candidates[0].ModelType, s.slimTaskRequestSnapshot(task, firstCandidateBody)); err != nil {
return Result{}, err
}
estimatedBillings := s.estimatedBillings(ctx, user, task.Kind, firstCandidateBody, candidates[0])
@@ -508,13 +512,13 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
QueueKey: candidate.QueueKey,
Status: "running",
Simulated: simulated,
RequestSnapshot: body,
RequestSnapshot: s.slimTaskRequestSnapshot(task, body),
Metrics: baseAttemptMetrics,
})
if err != nil {
return clients.Response{}, fmt.Errorf("create task attempt: %w", err)
}
if err := s.recordTaskParameterPreprocessing(ctx, task.ID, attemptID, attemptNo, candidate, preprocessing); err != nil {
if err := s.recordTaskParameterPreprocessing(ctx, task.ID, attemptID, attemptNo, candidate, s.slimParameterPreprocessingLog(task, preprocessing)); err != nil {
clientErr := &clients.ClientError{Code: "runtime_error", Message: err.Error(), Retryable: false}
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
AttemptID: attemptID,
@@ -560,12 +564,24 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
return clients.Response{}, fmt.Errorf("prepare http client: %w", err)
}
client := s.clientFor(candidate, simulated)
providerBody, err := s.hydrateProviderRequestAssets(ctx, body)
if err != nil {
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
AttemptID: attemptID,
Status: "failed",
Retryable: false,
Metrics: mergeMetrics(baseAttemptMetrics, map[string]any{"error": err.Error(), "retryable": false, "trace": []any{failureTraceEntry(err, false)}}),
ErrorCode: clients.ErrorCode(err),
ErrorMessage: err.Error(),
})
return clients.Response{}, err
}
callStartedAt := time.Now()
response, err := client.Run(ctx, clients.Request{
Kind: task.Kind,
ModelType: candidate.ModelType,
Model: task.Model,
Body: body,
Body: providerBody,
Candidate: candidate,
HTTPClient: requestHTTPClient,
RemoteTaskID: task.RemoteTaskID,
@@ -576,7 +592,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
}
return s.store.SetTaskRemoteTask(context.WithoutCancel(ctx), task.ID, attemptID, remoteTaskID, payload)
},
Stream: boolFromMap(body, "stream"),
Stream: boolFromMap(providerBody, "stream"),
StreamDelta: onDelta,
})
callFinishedAt := time.Now()
@@ -826,7 +842,7 @@ func (s *Service) recordFailedAttempt(ctx context.Context, input failedAttemptRe
QueueKey: queueKey,
Status: "running",
Simulated: input.Simulated,
RequestSnapshot: input.Body,
RequestSnapshot: s.slimTaskRequestSnapshot(input.Task, input.Body),
Metrics: metrics,
})
if err != nil {
@@ -834,7 +850,8 @@ func (s *Service) recordFailedAttempt(ctx context.Context, input failedAttemptRe
return attemptNo
}
if input.Preprocessing != nil && input.Candidate != nil {
if err := s.recordTaskParameterPreprocessing(ctx, input.Task.ID, attemptID, attemptNo, *input.Candidate, *input.Preprocessing); err != nil {
preprocessing := s.slimParameterPreprocessingLog(input.Task, *input.Preprocessing)
if err := s.recordTaskParameterPreprocessing(ctx, input.Task.ID, attemptID, attemptNo, *input.Candidate, preprocessing); err != nil {
s.logger.Warn("record failed attempt parameter preprocessing failed", "taskID", input.Task.ID, "attempt", attemptNo, "error", err)
}
}