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

feat: add river-backed async task queue

Browse Source
This commit is contained in:
wangbo
2026-05-12 10:11:54 +08:00
parent d69aaed444
commit 7e220b7477
30 changed files with 1342 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/postgres.go
+24 -5
View File
@@ -53,6 +53,10 @@ func (s *Store) Ping(ctx context.Context) error {
return s.pool.Ping(ctx)
}
func (s *Store) Pool() *pgxpool.Pool {
return s.pool
}
type Platform struct {
ID string `json:"id"`
Provider string `json:"provider"`
@@ -374,6 +378,7 @@ type CreateTaskInput struct {
Kind string `json:"kind"`
Model string `json:"model"`
RunMode string `json:"runMode"`
Async bool `json:"async"`
Request map[string]any `json:"request"`
}
@@ -398,7 +403,12 @@ type GatewayTask struct {
ResolvedModel string `json:"resolvedModel,omitempty"`
RequestID string `json:"requestId,omitempty"`
Request map[string]any `json:"request,omitempty"`
AsyncMode bool `json:"asyncMode"`
RiverJobID int64 `json:"riverJobId,omitempty"`
Status string `json:"status"`
AttemptCount int `json:"attemptCount"`
RemoteTaskID string `json:"remoteTaskId,omitempty"`
RemoteTaskPayload map[string]any `json:"remoteTaskPayload,omitempty"`
Result map[string]any `json:"result,omitempty"`
Billings []any `json:"billings,omitempty"`
Usage map[string]any `json:"usage"`
@@ -423,7 +433,9 @@ COALESCE(gateway_tenant_id::text, ''), COALESCE(tenant_id, ''), COALESCE(tenant_
COALESCE(api_key_id, ''), COALESCE(api_key_name, ''), COALESCE(api_key_prefix, ''),
COALESCE(user_group_id::text, ''), COALESCE(user_group_key, ''), model,
COALESCE(model_type, ''), COALESCE(requested_model, ''), COALESCE(resolved_model, ''), COALESCE(request_id, ''),
request, status, COALESCE(result, '{}'::jsonb), COALESCE(billings, '[]'::jsonb),
request, COALESCE(async_mode, false), COALESCE(river_job_id, 0), status, COALESCE(attempt_count, 0),
COALESCE(remote_task_id, ''), COALESCE(remote_task_payload, '{}'::jsonb),
COALESCE(result, '{}'::jsonb), COALESCE(billings, '[]'::jsonb),
COALESCE(usage, '{}'::jsonb), COALESCE(metrics, '{}'::jsonb), COALESCE(billing_summary, '{}'::jsonb),
COALESCE(final_charge_amount, 0)::float8, COALESCE(response_started_at::text, ''),
COALESCE(response_finished_at::text, ''), COALESCE(response_duration_ms, 0), COALESCE(error, ''),
@@ -1675,11 +1687,11 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
INSERT INTO gateway_tasks (
kind, run_mode, user_id, gateway_user_id, user_source, gateway_tenant_id, tenant_id, tenant_key,
api_key_id, api_key_name, api_key_prefix, user_group_id, user_group_key,
model, requested_model, request, status, result, billings, finished_at
model, requested_model, request, async_mode, status, result, billings, finished_at
)
VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, '')::uuid, NULLIF($13, ''), $14, $14, $15, $16, $17::jsonb, $18::jsonb, CASE WHEN $19 THEN now() ELSE NULL END)
VALUES ($1, $2, $3, NULLIF($4, '')::uuid, COALESCE(NULLIF($5, ''), 'gateway'), NULLIF($6, '')::uuid, NULLIF($7, ''), NULLIF($8, ''), NULLIF($9, ''), NULLIF($10, ''), NULLIF($11, ''), NULLIF($12, '')::uuid, NULLIF($13, ''), $14, $14, $15, $16, $17, $18::jsonb, $19::jsonb, CASE WHEN $20 THEN now() ELSE NULL END)
RETURNING `+gatewayTaskColumns,
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, status, resultBody, billingsBody, false,
input.Kind, runMode, user.ID, user.GatewayUserID, user.Source, user.GatewayTenantID, user.TenantID, user.TenantKey, user.APIKeyID, user.APIKeyName, user.APIKeyPrefix, user.UserGroupID, user.UserGroupKey, input.Model, requestBody, input.Async, status, resultBody, billingsBody, false,
))
if err != nil {
return GatewayTask{}, err
@@ -1689,7 +1701,7 @@ func (s *Store) CreateTask(ctx context.Context, input CreateTaskInput, user *aut
payload, _ := json.Marshal(event.Payload)
if _, err := tx.Exec(ctx, `
INSERT INTO gateway_task_events (task_id, seq, event_type, status, phase, progress, message, payload, simulated)
VALUES ($1::uuid, $2, $3, NULLIF($4, ''), NULLIF($5, ''), $6, NULLIF($7, ''), $8::jsonb, $9)`,
VALUES ($1::uuid, $2, $3::text, NULLIF($4::text, ''), NULLIF($5::text, ''), $6, NULLIF($7::text, ''), $8::jsonb, $9)`,
task.ID, event.Seq, event.EventType, event.Status, event.Phase, event.Progress, event.Message, string(payload), event.Simulated,
); err != nil {
return GatewayTask{}, err
@@ -1730,6 +1742,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
var usageBytes []byte
var metricsBytes []byte
var billingSummaryBytes []byte
var remoteTaskPayloadBytes []byte
if err := scanner.Scan(
&task.ID,
&task.Kind,
@@ -1751,7 +1764,12 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
&task.ResolvedModel,
&task.RequestID,
&requestBytes,
&task.AsyncMode,
&task.RiverJobID,
&task.Status,
&task.AttemptCount,
&task.RemoteTaskID,
&remoteTaskPayloadBytes,
&resultBytes,
&billingsBytes,
&usageBytes,
@@ -1771,6 +1789,7 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
return GatewayTask{}, err
}
task.Request = decodeObject(requestBytes)
task.RemoteTaskPayload = decodeObject(remoteTaskPayloadBytes)
task.Result = decodeObject(resultBytes)
task.Billings = decodeArray(billingsBytes)
task.Usage = decodeObject(usageBytes)