BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: record task attempt chains

Browse Source
This commit is contained in:
wangbo
2026-05-11 22:39:45 +08:00
parent 4c2de4b4c9
commit 0049b246c1
9 changed files with 492 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/tasks_runtime.go
+138 -2
View File
@@ -132,6 +132,10 @@ LIMIT $8 OFFSET $9`, queryArgs...)
if err := rows.Err(); err != nil {
return TaskListResult{}, err
}
items, err = s.attachTaskAttempts(ctx, items)
if err != nil {
return TaskListResult{}, err
}
return TaskListResult{
Items: items,
Total: total,
@@ -163,6 +167,7 @@ WHERE id = $1::uuid`, taskID, modelType, string(normalizedJSON))
func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptInput) (string, error) {
requestJSON, _ := json.Marshal(emptyObjectIfNil(input.RequestSnapshot))
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
tx, err := s.pool.Begin(ctx)
if err != nil {
return "", err
@@ -173,11 +178,11 @@ func (s *Store) CreateTaskAttempt(ctx context.Context, input CreateTaskAttemptIn
err = tx.QueryRow(ctx, `
INSERT INTO gateway_task_attempts (
task_id, attempt_no, platform_id, platform_model_id, client_id, queue_key,
status, simulated, request_snapshot
status, simulated, request_snapshot, metrics
)
VALUES (
$1::uuid, $2, NULLIF($3, '')::uuid, NULLIF($4, '')::uuid, NULLIF($5, ''), $6,
$7, $8, $9::jsonb
$7, $8, $9::jsonb, $10::jsonb
)
RETURNING id::text`,
input.TaskID,
@@ -189,6 +194,7 @@ RETURNING id::text`,
firstNonEmpty(input.Status, "running"),
input.Simulated,
string(requestJSON),
string(metricsJSON),
).Scan(&attemptID)
if err != nil {
return "", err
@@ -202,6 +208,136 @@ WHERE id = $1::uuid`, input.TaskID, input.AttemptNo); err != nil {
return attemptID, tx.Commit(ctx)
}
func (s *Store) attachTaskAttempts(ctx context.Context, items []GatewayTask) ([]GatewayTask, error) {
if len(items) == 0 {
return items, nil
}
taskIDs := make([]string, 0, len(items))
for _, item := range items {
taskIDs = append(taskIDs, item.ID)
}
attemptsByTaskID, err := s.listTaskAttemptsByTaskIDs(ctx, taskIDs)
if err != nil {
return nil, err
}
for index := range items {
items[index].Attempts = attemptsByTaskID[items[index].ID]
}
return items, nil
}
func (s *Store) ListTaskAttempts(ctx context.Context, taskID string) ([]TaskAttempt, error) {
attemptsByTaskID, err := s.listTaskAttemptsByTaskIDs(ctx, []string{taskID})
if err != nil {
return nil, err
}
return attemptsByTaskID[taskID], nil
}
func (s *Store) listTaskAttemptsByTaskIDs(ctx context.Context, taskIDs []string) (map[string][]TaskAttempt, error) {
itemsByTaskID := map[string][]TaskAttempt{}
if len(taskIDs) == 0 {
return itemsByTaskID, nil
}
rows, err := s.pool.Query(ctx, `
SELECT a.id::text, a.task_id::text, a.attempt_no,
COALESCE(a.platform_id::text, ''), COALESCE(p.name, ''), COALESCE(p.provider, ''),
COALESCE(a.platform_model_id::text, ''), COALESCE(pm.model_name, ''),
COALESCE(NULLIF(pm.provider_model_name, ''), pm.model_name, ''),
COALESCE(pm.model_alias, ''),
COALESCE(a.client_id, ''), a.queue_key, a.status, a.retryable, a.simulated,
COALESCE(a.request_id, ''), COALESCE(a.usage, '{}'::jsonb), COALESCE(a.metrics, '{}'::jsonb),
a.request_snapshot, COALESCE(a.response_snapshot, '{}'::jsonb),
COALESCE(a.response_started_at::text, ''), COALESCE(a.response_finished_at::text, ''),
COALESCE(a.response_duration_ms, 0), COALESCE(a.error_code, ''), COALESCE(a.error_message, ''),
a.started_at, COALESCE(a.finished_at::text, '')
FROM gateway_task_attempts a
LEFT JOIN integration_platforms p ON p.id = a.platform_id
LEFT JOIN platform_models pm ON pm.id = a.platform_model_id
WHERE a.task_id::text = ANY($1)
ORDER BY a.task_id, a.attempt_no`, taskIDs)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
item, err := scanTaskAttempt(rows)
if err != nil {
return nil, err
}
itemsByTaskID[item.TaskID] = append(itemsByTaskID[item.TaskID], item)
}
if err := rows.Err(); err != nil {
return nil, err
}
return itemsByTaskID, nil
}
func scanTaskAttempt(scanner taskScanner) (TaskAttempt, error) {
var item TaskAttempt
var usageBytes []byte
var metricsBytes []byte
var requestBytes []byte
var responseBytes []byte
if err := scanner.Scan(
&item.ID,
&item.TaskID,
&item.AttemptNo,
&item.PlatformID,
&item.PlatformName,
&item.Provider,
&item.PlatformModelID,
&item.ModelName,
&item.ProviderModelName,
&item.ModelAlias,
&item.ClientID,
&item.QueueKey,
&item.Status,
&item.Retryable,
&item.Simulated,
&item.RequestID,
&usageBytes,
&metricsBytes,
&requestBytes,
&responseBytes,
&item.ResponseStartedAt,
&item.ResponseFinishedAt,
&item.ResponseDurationMS,
&item.ErrorCode,
&item.ErrorMessage,
&item.StartedAt,
&item.FinishedAt,
); err != nil {
return TaskAttempt{}, err
}
item.Usage = decodeObject(usageBytes)
item.Metrics = decodeObject(metricsBytes)
item.RequestSnapshot = decodeObject(requestBytes)
item.ResponseSnapshot = decodeObject(responseBytes)
enrichTaskAttemptFromMetrics(&item)
return item, nil
}
func enrichTaskAttemptFromMetrics(item *TaskAttempt) {
if item == nil || len(item.Metrics) == 0 {
return
}
item.PlatformID = firstNonEmpty(item.PlatformID, taskAttemptMetricString(item.Metrics, "platformId"))
item.PlatformName = firstNonEmpty(item.PlatformName, taskAttemptMetricString(item.Metrics, "platformName"))
item.Provider = firstNonEmpty(item.Provider, taskAttemptMetricString(item.Metrics, "provider"))
item.PlatformModelID = firstNonEmpty(item.PlatformModelID, taskAttemptMetricString(item.Metrics, "platformModelId"))
item.ModelName = firstNonEmpty(item.ModelName, taskAttemptMetricString(item.Metrics, "resolvedModel"), taskAttemptMetricString(item.Metrics, "modelName"))
item.ProviderModelName = firstNonEmpty(item.ProviderModelName, taskAttemptMetricString(item.Metrics, "providerModel"))
item.ModelAlias = firstNonEmpty(item.ModelAlias, taskAttemptMetricString(item.Metrics, "modelAlias"))
item.ModelType = firstNonEmpty(item.ModelType, taskAttemptMetricString(item.Metrics, "modelType"))
item.ClientID = firstNonEmpty(item.ClientID, taskAttemptMetricString(item.Metrics, "clientId"))
}
func taskAttemptMetricString(metrics map[string]any, key string) string {
value, _ := metrics[key].(string)
return strings.TrimSpace(value)
}
func (s *Store) FinishTaskAttempt(ctx context.Context, input FinishTaskAttemptInput) error {
responseJSON, _ := json.Marshal(emptyObjectIfNil(input.ResponseSnapshot))
usageJSON, _ := json.Marshal(emptyObjectIfNil(input.Usage))