Harden worker federation and operator UI

2026-07-29 13:30:55 +04:00
parent 95a96d87a5
commit 1ca9d64e89
35 changed files with 1195 additions and 581 deletions
@@ -1,12 +1,19 @@
 package main

 import (
+	"bytes"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
 	"os"
 	"path/filepath"
 	"testing"
 	"time"

+	"orchestra/internal/authz"
+	"orchestra/internal/domain"
 	"orchestra/internal/registry"
+	"orchestra/internal/store"
 )

 type unreachable struct{}
@@ -33,3 +40,93 @@ func TestFederatedReachabilityDefersRemoteHerdrToWorkerHeartbeat(t *testing.T) {
 		t.Fatal("local herdr should still require its TCP probe")
 	}
 }
+
+func TestHarnessCompletionBuildsReceiptAndQuotaFromTranscript(t *testing.T) {
+	s, err := store.Open(t.TempDir())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := s.Append(domain.Event{ID: domain.NewID(), Type: "TaskCreated", TaskID: "done", Surface: string(authz.System), Payload: []byte(`{"source":"qa","external_id":"done","project":"p"}`)}); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := s.Lease("done", "local-claude", time.Minute); err != nil {
+		t.Fatal(err)
+	}
+	transcript := filepath.Join(t.TempDir(), "transcript.jsonl")
+	if err := os.WriteFile(transcript, []byte(`{"message":{"usage":{"input_tokens":100,"cache_read_input_tokens":20,"cache_creation_input_tokens":5,"output_tokens":7}}}`+"\n"), 0600); err != nil {
+		t.Fatal(err)
+	}
+	body, _ := json.Marshal(map[string]string{"task_id": "done", "transcript_path": transcript, "report": "# done"})
+	req := httptest.NewRequest(http.MethodPost, "/v1/harness/complete", bytes.NewReader(body))
+	res := httptest.NewRecorder()
+	harnessCompletion{store: s, token: "secret"}.ServeHTTP(res, req)
+	if res.Code != http.StatusUnauthorized {
+		t.Fatalf("missing token status = %d, want 401", res.Code)
+	}
+	req.Header.Set("Authorization", "Bearer secret")
+	req = httptest.NewRequest(http.MethodPost, "/v1/harness/complete", bytes.NewReader(body))
+	req.Header.Set("Authorization", "Bearer secret")
+	res = httptest.NewRecorder()
+	harnessCompletion{store: s, token: "secret"}.ServeHTTP(res, req)
+	if res.Code != http.StatusOK {
+		t.Fatalf("completion status = %d: %s", res.Code, res.Body.String())
+	}
+	task, ok := s.Task("done")
+	if !ok || task.State != domain.StateCompleted {
+		t.Fatalf("task after completion = %#v, present=%v", task, ok)
+	}
+	var completed struct {
+		Receipt struct {
+			Input      int `json:"input_tokens"`
+			CacheRead  int `json:"cache_read_tokens"`
+			CacheWrite int `json:"cache_write_tokens"`
+			Output     int `json:"output_tokens"`
+		} `json:"receipt"`
+	}
+	for _, e := range s.Events(0) {
+		if e.Type == "TaskCompleted" {
+			if err := json.Unmarshal(e.Payload, &completed); err != nil {
+				t.Fatal(err)
+			}
+		}
+	}
+	if completed.Receipt.Input != 100 || completed.Receipt.CacheRead != 20 || completed.Receipt.CacheWrite != 5 || completed.Receipt.Output != 7 {
+		t.Fatalf("receipt = %#v", completed.Receipt)
+	}
+	var quota struct {
+		Harness  string  `json:"harness_id"`
+		Consumed float64 `json:"consumed"`
+	}
+	found := false
+	for _, e := range s.Events(0) {
+		if e.Type == "QuotaReported" {
+			_ = json.Unmarshal(e.Payload, &quota)
+			found = true
+		}
+	}
+	if !found || quota.Harness != "local-claude" || quota.Consumed != 125 {
+		t.Fatalf("quota report = %#v, found=%v", quota, found)
+	}
+}
+
+func TestMultiMachineRegistryRequiresKnownLocalMachine(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "config.json")
+	if err := os.WriteFile(path, []byte(`{
+  "machines":[{"id":"homesrv","address":"192.168.1.104:9145"},{"id":"workpc","address":"192.168.1.105:9145"}]
+}`), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	r, err := registry.Load(path)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := validateLocalMachine(r, ""); err == nil {
+		t.Fatal("missing local machine accepted")
+	}
+	if err := validateLocalMachine(r, "missing"); err == nil {
+		t.Fatal("unknown local machine accepted")
+	}
+	if err := validateLocalMachine(r, "homesrv"); err != nil {
+		t.Fatalf("known local machine rejected: %v", err)
+	}
+}