Add Go worker agent prototype

2026-08-02 16:27:58 +03:00
parent f20cc7fe00
commit 644c287002
15 changed files with 2012 additions and 1 deletions
@@ -0,0 +1,205 @@
+// Package agent implements a Go worker agent: a coordinator client and
+// task-lifecycle supervisor that executes SDK workloads in a Python
+// subprocess. It mirrors the Python worker's v1 wire contract exactly; the
+// Python worker remains the reference implementation.
+package agent
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/url"
+	"regexp"
+	"strings"
+	"time"
+)
+
+var (
+	uuidPattern     = regexp.MustCompile(`^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$`)
+	sha256Pattern   = regexp.MustCompile(`^[0-9a-fA-F]{64}$`)
+	workloadPattern = regexp.MustCompile(`^[a-z][a-z0-9]*(?:[_-][a-z0-9]+)*$`)
+)
+
+// RegisteredWorker is the coordinator's answer to /workers/register.
+type RegisteredWorker struct {
+	WorkerID                 string
+	HeartbeatIntervalSeconds float64
+}
+
+// Input is the claimed task's input artifact.
+type Input struct {
+	URI    string
+	SHA256 string
+}
+
+// Task is one claimed, leased task.
+type Task struct {
+	TaskID          string
+	Attempt         int
+	LeaseExpiresAt  time.Time
+	Workload        string
+	Input           Input
+	Parameters      map[string]any
+	leaseExpiresRaw string
+}
+
+// Uploaded is the coordinator-owned metadata returned after artifact upload.
+type Uploaded struct {
+	ArtifactID string
+	URI        string
+	SHA256     string
+	SizeBytes  int64
+}
+
+func requireString(value any, field string) (string, error) {
+	text, ok := value.(string)
+	if !ok || strings.TrimSpace(text) == "" {
+		return "", fmt.Errorf("%s must be a non-empty string", field)
+	}
+	return text, nil
+}
+
+func safeCoordinatorURI(value any, field string) (string, error) {
+	uri, err := requireString(value, field)
+	if err != nil {
+		return "", err
+	}
+	if strings.HasPrefix(uri, "/") {
+		// A network-path reference (//host/path) or dot segments would
+		// resolve to another origin; reject both.
+		if strings.HasPrefix(uri, "//") {
+			return "", fmt.Errorf("%s must be a safe coordinator path", field)
+		}
+		for _, segment := range strings.Split(uri, "/") {
+			if segment == ".." {
+				return "", fmt.Errorf("%s must be a safe coordinator path", field)
+			}
+		}
+		return uri, nil
+	}
+	parsed, err := url.Parse(uri)
+	if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" {
+		return "", fmt.Errorf("%s must be an absolute HTTP(S) URL or coordinator path", field)
+	}
+	return uri, nil
+}
+
+func sha256Hex(value any, field string) (string, error) {
+	digest, err := requireString(value, field)
+	if err != nil {
+		return "", err
+	}
+	digest = strings.ToLower(digest)
+	if !sha256Pattern.MatchString(digest) {
+		return "", fmt.Errorf("%s must be a SHA-256 hex digest", field)
+	}
+	return digest, nil
+}
+
+func uuid(value any, field string) (string, error) {
+	text, err := requireString(value, field)
+	if err != nil {
+		return "", err
+	}
+	if !uuidPattern.MatchString(text) {
+		return "", fmt.Errorf("%s must be a UUID", field)
+	}
+	return strings.ToLower(text), nil
+}
+
+// ParseTask validates a claimed-task response with the same strictness as the
+// Python worker's ClaimedTask.from_json.
+func ParseTask(payload map[string]any) (*Task, error) {
+	rawInput, ok := payload["input"].(map[string]any)
+	if !ok {
+		return nil, fmt.Errorf("input must be an object")
+	}
+	rawAttempt, ok := payload["attempt"].(float64)
+	if !ok || rawAttempt < 1 || rawAttempt != float64(int(rawAttempt)) {
+		return nil, fmt.Errorf("attempt must be a positive integer")
+	}
+	taskID, err := uuid(payload["task_id"], "task_id")
+	if err != nil {
+		return nil, fmt.Errorf("invalid claimed-task response: %w", err)
+	}
+	rawLease, err := requireString(payload["lease_expires_at"], "lease_expires_at")
+	if err != nil {
+		return nil, fmt.Errorf("invalid claimed-task response: %w", err)
+	}
+	lease, err := time.Parse(time.RFC3339, rawLease)
+	if err != nil || lease.Location() == nil {
+		return nil, fmt.Errorf("lease_expires_at must include a timezone")
+	}
+	workload, err := requireString(payload["workload"], "workload")
+	if err != nil || !workloadPattern.MatchString(workload) {
+		return nil, fmt.Errorf("workload must be a canonical name")
+	}
+	uri, err := safeCoordinatorURI(rawInput["uri"], "input.uri")
+	if err != nil {
+		return nil, fmt.Errorf("invalid claimed-task response: %w", err)
+	}
+	digest, err := sha256Hex(rawInput["sha256"], "input.sha256")
+	if err != nil {
+		return nil, fmt.Errorf("invalid claimed-task response: %w", err)
+	}
+	parameters, ok := payload["parameters"].(map[string]any)
+	if !ok {
+		parameters = map[string]any{}
+	}
+	return &Task{
+		TaskID:          taskID,
+		Attempt:         int(rawAttempt),
+		LeaseExpiresAt:  lease,
+		leaseExpiresRaw: rawLease,
+		Workload:        workload,
+		Input:           Input{URI: uri, SHA256: digest},
+		Parameters:      parameters,
+	}, nil
+}
+
+// LeaseExpiresRaw returns the original lease timestamp string for
+// round-tripping in heartbeat deadlines.
+func (t *Task) LeaseExpiresRaw() string { return t.leaseExpiresRaw }
+
+// ParseRegistered validates a registration response.
+func ParseRegistered(payload map[string]any) (*RegisteredWorker, error) {
+	workerID, err := uuid(payload["worker_id"], "worker_id")
+	if err != nil {
+		return nil, fmt.Errorf("invalid worker registration response: %w", err)
+	}
+	interval, ok := payload["heartbeat_interval_seconds"].(float64)
+	if !ok || interval <= 0 {
+		return nil, fmt.Errorf("heartbeat_interval_seconds must be positive")
+	}
+	return &RegisteredWorker{WorkerID: workerID, HeartbeatIntervalSeconds: interval}, nil
+}
+
+// ParseUploaded validates an artifact upload response.
+func ParseUploaded(payload map[string]any) (*Uploaded, error) {
+	artifactID, err := uuid(payload["artifact_id"], "artifact_id")
+	if err != nil {
+		return nil, fmt.Errorf("invalid artifact upload response: %w", err)
+	}
+	uri, err := safeCoordinatorURI(payload["uri"], "uri")
+	if err != nil {
+		return nil, fmt.Errorf("invalid artifact upload response: %w", err)
+	}
+	digest, err := sha256Hex(payload["sha256"], "sha256")
+	if err != nil {
+		return nil, fmt.Errorf("invalid artifact upload response: %w", err)
+	}
+	rawSize, ok := payload["size_bytes"].(float64)
+	if !ok || rawSize < 0 || rawSize != float64(int64(rawSize)) {
+		return nil, fmt.Errorf("artifact size_bytes must be a non-negative integer")
+	}
+	return &Uploaded{ArtifactID: artifactID, URI: uri, SHA256: digest, SizeBytes: int64(rawSize)}, nil
+}
+
+// TaskRunnerManifest is what the Python task entry writes on success.
+type TaskRunnerManifest struct {
+	ArtifactPath string         `json:"artifact_path"`
+	ContentType  string         `json:"content_type"`
+	Metrics      map[string]any `json:"metrics"`
+}
+
+// Encode serializes a claim payload for /tasks/claim.
+func Encode(v any) ([]byte, error) { return json.Marshal(v) }