A self-hosted orchestration system crafted for power users. composia.xyz
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153package agent
import ( "context" "fmt" "log" "time"
"connectrpc.com/connect" agentv1 "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/agent/v1" "forgejo.alexma.top/alexma233/composia/gen/go/proto/composia/agent/v1/agentv1connect" "forgejo.alexma.top/alexma233/composia/internal/core/config" "forgejo.alexma.top/alexma233/composia/internal/core/task" "forgejo.alexma.top/alexma233/composia/internal/platform/store" "google.golang.org/protobuf/types/known/timestamppb")
func executeTaskStep(ctx context.Context, client agentv1connect.AgentReportServiceClient, logUploader *taskLogUploader, taskID string, stepName task.StepName, execute func() error) error { stepStartedAt := time.Now() startedAt := timestamppb.New(stepStartedAt) stepDuration := func() time.Duration { return time.Since(stepStartedAt).Round(time.Millisecond) } logStepFailed := func(err error) { log.Printf("agent task step failed: task_id=%s step=%s duration=%s error=%v", taskID, stepName, stepDuration(), err) }
log.Printf("agent task step started: task_id=%s step=%s", taskID, stepName) if err := reportTaskStepStateWithTimeout(ctx, client, &agentv1.ReportTaskStepStateRequest{TaskId: taskID, StepName: protoAgentTaskStepName(stepName), Status: protoAgentTaskStatus(task.StatusRunning), StartedAt: startedAt}, taskReportTimeout); err != nil { logStepFailed(err) return fmt.Errorf("report running step %s: %w", stepName, err) } if err := uploadTaskLog(ctx, logUploader, fmt.Sprintf("step %s started\n", stepName)); err != nil { logStepFailed(err) return err } if err := execute(); err != nil { finishedAt := timestamppb.Now() _ = reportTaskStepStateWithTimeout(ctx, client, &agentv1.ReportTaskStepStateRequest{TaskId: taskID, StepName: protoAgentTaskStepName(stepName), Status: protoAgentTaskStatus(task.StatusFailed), StartedAt: startedAt, FinishedAt: finishedAt}, taskReportTimeout) _ = uploadTaskLog(ctx, logUploader, fmt.Sprintf("step %s failed: %v\n", stepName, err)) logStepFailed(err) return err } finishedAt := timestamppb.Now() if err := reportTaskStepStateWithTimeout(ctx, client, &agentv1.ReportTaskStepStateRequest{TaskId: taskID, StepName: protoAgentTaskStepName(stepName), Status: protoAgentTaskStatus(task.StatusSucceeded), StartedAt: startedAt, FinishedAt: finishedAt}, taskReportTimeout); err != nil { logStepFailed(err) log.Printf("agent task step succeeded but its status report failed: task_id=%s step=%s error=%v", taskID, stepName, err) } if err := uploadTaskLog(ctx, logUploader, fmt.Sprintf("step %s succeeded\n", stepName)); err != nil { logStepFailed(err) return err } log.Printf("agent task step succeeded: task_id=%s step=%s duration=%s", taskID, stepName, stepDuration()) return nil}
func reportTaskCompletion(ctx context.Context, client agentv1connect.AgentReportServiceClient, taskID string, status task.Status, errorSummary string) error { return reportTaskCompletionWithTimeout(ctx, client, taskID, status, errorSummary, taskReportTimeout)}
func reportTaskCompletionWithTimeout(ctx context.Context, client agentv1connect.AgentReportServiceClient, taskID string, status task.Status, errorSummary string, timeout time.Duration) error { finishedAt := time.Now().UTC() if runtime := taskExecutionFromContext(ctx); runtime != nil { if err := runtime.storeCompletion(completion{Status: status, ErrorSummary: errorSummary, FinishedAt: finishedAt}); err != nil { return err } } callCtx := ctx var cancel context.CancelFunc if timeout > 0 { callCtx, cancel = context.WithTimeout(ctx, timeout) defer cancel() } _, err := client.ReportTaskState(callCtx, connect.NewRequest(&agentv1.ReportTaskStateRequest{TaskId: taskID, ExecutionId: taskExecutionID(ctx), Status: protoAgentTaskStatus(status), ErrorSummary: errorSummary, FinishedAt: timestamppb.New(finishedAt)})) if err != nil { return fmt.Errorf("report task completion: %w", err) } if runtime := taskExecutionFromContext(ctx); runtime != nil { if err := runtime.clear(); err != nil { return fmt.Errorf("clear task execution state: %w", err) } } return nil}
func reportTaskStepStateWithTimeout(ctx context.Context, client agentv1connect.AgentReportServiceClient, request *agentv1.ReportTaskStepStateRequest, timeout time.Duration) error { callCtx := ctx var cancel context.CancelFunc if timeout > 0 { callCtx, cancel = context.WithTimeout(ctx, timeout) defer cancel() } request.ExecutionId = taskExecutionID(ctx) _, err := client.ReportTaskStepState(callCtx, connect.NewRequest(request)) if err != nil { return err } return nil}
func failTask(ctx context.Context, client agentv1connect.AgentReportServiceClient, taskID string, err error) error { _ = reportTaskCompletion(ctx, client, taskID, task.StatusFailed, err.Error()) return err}
func failServiceTask(ctx context.Context, client agentv1connect.AgentReportServiceClient, cfg *config.AgentConfig, pulledTask *agentv1.AgentTask, err error) error { _ = reportServiceStatus(ctx, client, cfg, pulledTask.GetServiceName(), store.ServiceRuntimeError) return failTask(ctx, client, pulledTask.GetTaskId(), err)}
func uploadTaskLog(ctx context.Context, logUploader *taskLogUploader, content string) error { if logUploader == nil { return nil } if err := logUploader.Upload(ctx, content); err != nil { log.Printf("upload task logs failed: %v", err) } return nil}
func reportBackupResult(ctx context.Context, client agentv1connect.AgentReportServiceClient, taskID, serviceName, dataName, artifactRef string, status task.Status, startedAt, finishedAt time.Time, errorSummary string) error { result := persistedBackupResult{ServiceName: serviceName, DataName: dataName, ArtifactRef: artifactRef, Status: status, StartedAt: startedAt, FinishedAt: finishedAt, ErrorSummary: errorSummary} runtime := taskExecutionFromContext(ctx) if runtime != nil { if err := runtime.storeBackupResult(result); err != nil { return err } } for { err := sendPersistedBackupResult(ctx, client, taskID, taskExecutionID(ctx), result) if err == nil { if runtime != nil { return runtime.removeFirstBackupResult() } return nil } timer := time.NewTimer(500 * time.Millisecond) select { case <-ctx.Done(): timer.Stop() return fmt.Errorf("report backup result: %w", err) case <-timer.C: } }}
func sendPersistedBackupResult(ctx context.Context, client agentv1connect.AgentReportServiceClient, taskID, executionID string, result persistedBackupResult) error { callCtx, cancel := context.WithTimeout(ctx, taskReportTimeout) defer cancel() _, err := client.ReportBackupResult(callCtx, connect.NewRequest(&agentv1.ReportBackupResultRequest{TaskId: taskID, ServiceName: result.ServiceName, DataName: result.DataName, Status: string(result.Status), StartedAt: timestamppb.New(result.StartedAt), FinishedAt: timestamppb.New(result.FinishedAt), ArtifactRef: result.ArtifactRef, ErrorSummary: result.ErrorSummary, ExecutionId: executionID})) return err}