feat: add river-backed async task queue
This commit is contained in:
@@ -542,11 +542,13 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||
writeError(w, http.StatusForbidden, "api key scope does not allow this capability")
|
||||
return
|
||||
}
|
||||
asyncMode := asyncRequest(r)
|
||||
|
||||
task, err := s.store.CreateTask(r.Context(), store.CreateTaskInput{
|
||||
Kind: kind,
|
||||
Model: model,
|
||||
RunMode: runModeFromRequest(body),
|
||||
Async: asyncMode,
|
||||
Request: body,
|
||||
}, user)
|
||||
if err != nil {
|
||||
@@ -554,6 +556,14 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||
writeError(w, http.StatusInternalServerError, "create task failed")
|
||||
return
|
||||
}
|
||||
if asyncMode {
|
||||
if err := s.runner.EnqueueAsyncTask(r.Context(), task); err != nil {
|
||||
writeError(w, http.StatusInternalServerError, err.Error(), "enqueue_failed")
|
||||
return
|
||||
}
|
||||
writeTaskAccepted(w, task)
|
||||
return
|
||||
}
|
||||
if compatible {
|
||||
if boolValue(body, "stream") {
|
||||
flusher := prepareCompatibleStream(w)
|
||||
@@ -602,13 +612,23 @@ func (s *Server) createTask(kind string, compatible bool) http.Handler {
|
||||
s.logger.Warn("task completed with failure", "kind", kind, "taskId", task.ID, "error", runErr)
|
||||
}
|
||||
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{
|
||||
"task": result.Task,
|
||||
"next": map[string]string{
|
||||
"events": fmt.Sprintf("/api/v1/tasks/%s/events", task.ID),
|
||||
"detail": fmt.Sprintf("/api/v1/tasks/%s", task.ID),
|
||||
},
|
||||
})
|
||||
writeTaskAccepted(w, result.Task)
|
||||
})
|
||||
}
|
||||
|
||||
func asyncRequest(r *http.Request) bool {
|
||||
value := strings.TrimSpace(strings.ToLower(r.Header.Get("x-async")))
|
||||
return value == "1" || value == "true" || value == "yes" || value == "on"
|
||||
}
|
||||
|
||||
func writeTaskAccepted(w http.ResponseWriter, task store.GatewayTask) {
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{
|
||||
"taskId": task.ID,
|
||||
"task": task,
|
||||
"next": map[string]string{
|
||||
"events": fmt.Sprintf("/api/v1/tasks/%s/events", task.ID),
|
||||
"detail": fmt.Sprintf("/api/v1/tasks/%s", task.ID),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user