Files
m 531dfa85ad
CI / test (push) Successful in 54s
CI / test (pull_request) Successful in 49s
feat(work): expose additional evaluation scheduling
2026-08-23 23:45:45 +03:00

197 lines
8.0 KiB
Go

package main
import (
"context"
"encoding/json"
"fmt"
"log"
"net/http"
"os"
"time"
"git.cin.su/m/Assistant/adapters/postgres"
secretcontracts "git.cin.su/m/Assistant/contracts/secrets"
"git.cin.su/m/Assistant/contracts/work"
"git.cin.su/m/Assistant/modules/work/application"
"git.cin.su/m/Assistant/platform/secrets"
"git.cin.su/m/Assistant/runtime/health"
)
func main() {
ctx := context.Background()
password, err := resolveDatabasePassword()
if err != nil {
log.Fatal(err)
}
repository, err := postgres.OpenWork(ctx, postgres.ContinuityConfig{Host: os.Getenv("POSTGRES_HOST"), Port: os.Getenv("POSTGRES_PORT"), Database: os.Getenv("POSTGRES_DATABASE"), User: os.Getenv("WORK_POSTGRES_USER"), Password: password})
if err != nil {
log.Fatal(err)
}
defer repository.Close()
evidenceSink, err := newCoreTaskEvaluationSink(envOr("CORE_RUNTIME_URL", "http://core-runtime:8080"), &http.Client{Timeout: 12 * time.Second})
if err != nil {
log.Fatal(err)
}
mux := http.NewServeMux()
mux.Handle("/healthz", health.Handler{Runtime: "work-graph-runtime", SemanticGeneration: 1, ArchitectureGeneration: 1})
mux.Handle("/readyz", health.ReadinessHandler{Runtime: "work-graph-runtime", Probes: []health.Probe{{Name: "work-postgres", Check: func(ctx context.Context) error { return repository.Ping(ctx) }}}})
mux.Handle("/v1/work-graphs/status", workGraphStatusHandler(repository))
mux.HandleFunc("POST /v1/work-graphs/schedule", func(w http.ResponseWriter, r *http.Request) {
var request application.GraphScheduleRequest
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4<<20)).Decode(&request); err != nil {
http.Error(w, "invalid graph schedule request", http.StatusBadRequest)
return
}
graph, ready, err := application.AdmitAndSchedule(r.Context(), repository, request)
if err != nil {
log.Printf("work graph schedule rejected: %v", err)
http.Error(w, err.Error(), http.StatusUnprocessableEntity)
return
}
writeJSON(w, map[string]any{"graph": graph, "scheduled_nodes": ready})
})
mux.HandleFunc("POST /v1/work-graphs/task-evaluations/schedule", func(w http.ResponseWriter, r *http.Request) {
var request application.TaskEvaluationScheduleRequest
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4<<20)).Decode(&request); err != nil {
http.Error(w, "invalid task evaluation graph request", http.StatusBadRequest)
return
}
graph, ready, err := application.AdmitAndScheduleTaskEvaluation(r.Context(), repository, request)
if err != nil {
log.Printf("task evaluation graph schedule rejected: %v", err)
http.Error(w, err.Error(), http.StatusUnprocessableEntity)
return
}
writeJSON(w, map[string]any{"graph": graph, "scheduled_nodes": ready})
})
mux.HandleFunc("POST /v1/work-graphs/task-evaluations/additional/schedule", func(w http.ResponseWriter, r *http.Request) {
var request application.TaskEvaluationExpansionScheduleRequest
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4<<20)).Decode(&request); err != nil {
http.Error(w, "invalid additional task evaluation request", http.StatusBadRequest)
return
}
graph, ready, err := application.ScheduleAdditionalTaskEvaluation(r.Context(), repository, request)
if err != nil {
log.Printf("additional task evaluation schedule rejected: %v", err)
http.Error(w, err.Error(), http.StatusUnprocessableEntity)
return
}
writeJSON(w, map[string]any{"graph": graph, "scheduled_nodes": ready})
})
mux.HandleFunc("POST /v1/work-graphs/orchestration/next", func(w http.ResponseWriter, r *http.Request) {
var request work.SolutionSearchInput
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 32<<10)).Decode(&request); err != nil {
http.Error(w, "invalid orchestration input", http.StatusBadRequest)
return
}
decision, err := application.DecideSolutionSearch(request)
if err != nil {
http.Error(w, "invalid orchestration input", http.StatusUnprocessableEntity)
return
}
writeJSON(w, decision)
})
mux.HandleFunc("POST /v1/work-graphs/orchestration/wave", func(w http.ResponseWriter, r *http.Request) {
var request work.SolutionSearchWaveRequest
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&request); err != nil {
http.Error(w, "invalid orchestration wave request", http.StatusBadRequest)
return
}
result, err := application.OrchestrateSolutionSearch(r.Context(), repository, request)
if err != nil {
http.Error(w, "invalid orchestration wave request", http.StatusUnprocessableEntity)
return
}
writeJSON(w, result)
})
mux.HandleFunc("POST /v1/work-graphs/result-callback", func(w http.ResponseWriter, r *http.Request) {
var request work.GraphResultCallback
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&request); err != nil {
http.Error(w, "invalid graph result callback", http.StatusBadRequest)
return
}
graph, err := application.RecordGraphResultCandidate(r.Context(), repository, repository, request)
if err != nil {
http.Error(w, "invalid graph result callback", http.StatusUnprocessableEntity)
return
}
writeJSON(w, graph)
})
mux.HandleFunc("POST /v1/work-graphs/complete", func(w http.ResponseWriter, r *http.Request) {
var request application.GraphCompleteRequest
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 64<<10)).Decode(&request); err != nil {
http.Error(w, "invalid graph completion request", http.StatusBadRequest)
return
}
result, err := application.CompleteGraphNodeWithReplan(r.Context(), repository, request)
if err != nil {
http.Error(w, err.Error(), http.StatusUnprocessableEntity)
return
}
writeJSON(w, result)
})
mux.HandleFunc("POST /v1/work-graphs/task-evaluations/callback", func(w http.ResponseWriter, r *http.Request) {
var request work.TaskEvaluationEvidenceCallback
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&request); err != nil {
http.Error(w, "invalid task evaluation callback", http.StatusBadRequest)
return
}
graph, err := application.CompleteTaskEvaluationNode(r.Context(), repository, evidenceSink, request)
if err != nil {
http.Error(w, err.Error(), http.StatusUnprocessableEntity)
return
}
writeJSON(w, graph)
})
mux.HandleFunc("POST /v1/work-graphs/synthesis", func(w http.ResponseWriter, r *http.Request) {
var request application.GraphSynthesisRequest
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4<<20)).Decode(&request); err != nil {
http.Error(w, "invalid graph synthesis request", http.StatusBadRequest)
return
}
conflict, synthesis, graph, err := application.SynthesizeAndCompleteGraph(r.Context(), repository, repository, request)
if err != nil {
http.Error(w, err.Error(), http.StatusUnprocessableEntity)
return
}
writeJSON(w, map[string]any{"conflict": conflict, "synthesis": synthesis, "graph": graph})
})
addr := os.Getenv("LISTEN_ADDR")
if addr == "" {
addr = ":8088"
}
log.Printf("work-graph-runtime listening on %s", addr)
log.Fatal(http.ListenAndServe(addr, mux))
}
func envOr(name, fallback string) string {
if value := os.Getenv(name); value != "" {
return value
}
return fallback
}
func writeJSON(w http.ResponseWriter, value any) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(value)
}
func resolveDatabasePassword() (string, error) {
if _, ok := os.LookupEnv("WORK_POSTGRES_PASSWORD"); ok {
return "", fmt.Errorf("WORK_POSTGRES_PASSWORD direct environment secret is forbidden")
}
path, ref := os.Getenv("WORK_POSTGRES_PASSWORD_FILE"), os.Getenv("WORK_POSTGRES_CREDENTIAL_REF")
if path == "" || ref == "" {
return "", fmt.Errorf("WORK_POSTGRES_PASSWORD_FILE and WORK_POSTGRES_CREDENTIAL_REF are required")
}
root := os.Getenv("SECRET_BACKEND_ROOT")
if root == "" {
root = "/run/secrets"
}
resolver, err := secrets.NewFileResolver(root)
if err != nil {
return "", err
}
return resolver.ResolveFile(secretcontracts.SecretFileRef{CredentialRef: secretcontracts.CredentialRef(ref), Purpose: secretcontracts.SecretPurposeDatabasePassword, Consumer: secretcontracts.RuntimeConsumerWorkWorker, Path: path})
}