orchestra/internal/federation/client.go

package federation

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"orchestra/internal/domain"
	"strings"
)

// Client is the worker-side protocol client. It carries no task state: the
// homesrv event log remains authoritative and workers only persist their
// local execution session/checkouts.
type Client struct {
	BaseURL    string
	WorkerID   string
	Token      string
	AdmitToken string
	HTTP       *http.Client
}

func (c Client) Register(ctx context.Context, w Worker) error {
	b, err := json.Marshal(map[string]any{"id": w.ID, "address": w.Address, "capacity": w.Capacity, "token": c.Token})
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/v1/federation/workers", bytes.NewReader(b))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	if c.AdmitToken != "" {
		req.Header.Set("Authorization", "Bearer "+c.AdmitToken)
	}
	h := c.HTTP
	if h == nil {
		h = http.DefaultClient
	}
	resp, err := h.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		msg, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("federation register: %s: %s", resp.Status, strings.TrimSpace(string(msg)))
	}
	return nil
}

func (c Client) request(ctx context.Context, method, path string, body any) (*http.Response, error) {
	var r io.Reader
	if body != nil {
		b, err := json.Marshal(body)
		if err != nil {
			return nil, err
		}
		r = bytes.NewReader(b)
	}
	req, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(c.BaseURL, "/")+path, r)
	if err != nil {
		return nil, err
	}
	req.Header.Set("X-Orchestra-Worker", c.WorkerID)
	req.Header.Set("Authorization", "Bearer "+c.Token)
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	h := c.HTTP
	if h == nil {
		h = http.DefaultClient
	}
	resp, err := h.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode/100 != 2 {
		defer resp.Body.Close()
		b, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("federation: %s: %s", resp.Status, strings.TrimSpace(string(b)))
	}
	return resp, nil
}

func (c Client) Events(ctx context.Context, since uint64) ([]domain.Event, uint64, error) {
	resp, err := c.request(ctx, http.MethodGet, "/v1/federation/events?since="+fmt.Sprint(since), nil)
	if err != nil {
		return nil, 0, err
	}
	defer resp.Body.Close()
	var out struct {
		Cursor uint64         `json:"cursor"`
		Events []domain.Event `json:"events"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return nil, 0, err
	}
	return out.Events, out.Cursor, nil
}

// Tasks hydrates the worker's cache when its local state predates the
// coordinator's event-retention window. The coordinator remains authoritative
// for the task projection.
func (c Client) Tasks(ctx context.Context) ([]domain.Task, error) {
	resp, err := c.request(ctx, http.MethodGet, "/v1/tasks", nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	var tasks []domain.Task
	if err := json.NewDecoder(resp.Body).Decode(&tasks); err != nil {
		return nil, err
	}
	return tasks, nil
}
func (c Client) Ack(ctx context.Context, cursor uint64) error {
	resp, err := c.request(ctx, http.MethodPost, "/v1/federation/events/ack", map[string]uint64{"cursor": cursor})
	if resp != nil {
		resp.Body.Close()
	}
	return err
}
func (c Client) Heartbeat(ctx context.Context, health WorkerHealth) error {
	resp, err := c.request(ctx, http.MethodPost, "/v1/federation/workers/"+url.PathEscape(c.WorkerID)+"/heartbeat", health)
	if resp != nil {
		resp.Body.Close()
	}
	return err
}
func (c Client) Artifact(ctx context.Context, ref string) ([]byte, error) {
	resp, err := c.request(ctx, http.MethodGet, "/v1/artifacts/"+url.PathEscape(ref), nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return io.ReadAll(resp.Body)
}

func (c Client) PutArtifact(ctx context.Context, b []byte) (string, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.BaseURL, "/")+"/v1/artifacts", bytes.NewReader(b))
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", "application/octet-stream")
	req.Header.Set("X-Orchestra-Worker", c.WorkerID)
	req.Header.Set("Authorization", "Bearer "+c.Token)
	h := c.HTTP
	if h == nil {
		h = http.DefaultClient
	}
	resp, err := h.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		msg, _ := io.ReadAll(resp.Body)
		return "", fmt.Errorf("artifact upload: %s: %s", resp.Status, strings.TrimSpace(string(msg)))
	}
	var out struct {
		Ref string `json:"ref"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return "", err
	}
	return out.Ref, nil
}
func (c Client) Release(ctx context.Context, taskID, ref, anchor string) error {
	resp, err := c.request(ctx, http.MethodPost, "/v1/federation/workers/"+url.PathEscape(c.WorkerID)+"/handoff", map[string]string{"task_id": taskID, "handoff_ref": ref, "anchor_sha": anchor})
	if resp != nil {
		resp.Body.Close()
	}
	return err
}
func (c Client) Complete(ctx context.Context, taskID, reportRef string) error {
	resp, err := c.request(ctx, http.MethodPost, "/v1/federation/workers/"+url.PathEscape(c.WorkerID)+"/complete", map[string]string{"task_id": taskID, "handoff_ref": reportRef})
	if resp != nil {
		resp.Body.Close()
	}
	return err
}

func (c Client) PublishCapture(ctx context.Context, capture Capture) (Capture, error) {
	resp, err := c.request(ctx, http.MethodPost, "/v1/federation/workers/"+url.PathEscape(c.WorkerID)+"/captures", capture)
	if err != nil {
		return Capture{}, err
	}
	defer resp.Body.Close()
	var out Capture
	err = json.NewDecoder(resp.Body).Decode(&out)
	return out, err
}
func (c Client) Commands(ctx context.Context) ([]Command, error) {
	resp, err := c.request(ctx, http.MethodGet, "/v1/federation/commands", nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	var out []Command
	err = json.NewDecoder(resp.Body).Decode(&out)
	return out, err
}
func (c Client) ResolveCommand(ctx context.Context, id, status, message string) error {
	resp, err := c.request(ctx, http.MethodPost, "/v1/federation/commands/"+url.PathEscape(id), map[string]string{"status": status, "message": message})
	if resp != nil {
		resp.Body.Close()
	}
	return err
}