fix(beta9): 5 bugs from Cubic review of upstream PR — status-checking, idempotency, NodePort Service, SDK batch-embeddings

manifests/k3d/beta9.yaml

@@ -194,6 +194,36 @@ spec:
     targetPort: 9090
     protocol: TCP
 ---
+# External NodePort Service so Tailscale peers (SDK, workers, tests) can reach
+# the gateway via the host's stable :31994 (HTTP) / :31993 (gRPC) ports. The
+# internal ClusterIP Service `gateway` above stays for in-cluster traffic.
+# Cubic review finding: the previous manifest referenced 31994/31993 in docs
+# but never materialised a Service, so `ss -tlnp | grep 31994` returned empty
+# on the OCI host.
+apiVersion: v1
+kind: Service
+metadata:
+  name: gateway-nodeport
+  namespace: beta9
+  labels:
+    app: gateway
+    tier: external
+spec:
+  type: NodePort
+  selector:
+    app: gateway
+  ports:
+  - name: http
+    port: 1994
+    targetPort: 1994
+    nodePort: 31994
+    protocol: TCP
+  - name: grpc
+    port: 1993
+    targetPort: 1993
+    nodePort: 31993
+    protocol: TCP
+---
 # NOTE: IngressRouteTCP removed - using NodePort instead (no Traefik dependency)
 ---
 apiVersion: helm.cattle.io/v1

pkg/agent/control.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/rs/zerolog/log"
@@ -18,6 +19,18 @@ type ControlServer struct {
 	agent  *Agent
 	port   int
 	server *http.Server
+
+	// ollamaBaseURL lets tests redirect Ollama calls to an httptest server.
+	// Empty string means default: http://localhost:DefaultOllamaPort
+	ollamaBaseURL string
+}
+
+// ollamaURL returns the Ollama base URL honoring test overrides.
+func (c *ControlServer) ollamaURL() string {
+	if c.ollamaBaseURL != "" {
+		return c.ollamaBaseURL
+	}
+	return fmt.Sprintf("http://localhost:%d", DefaultOllamaPort)
 }
 
 // NewControlServer creates a new control server
@@ -117,13 +130,42 @@ func (c *ControlServer) handleInferenceStop(w http.ResponseWriter, r *http.Reque
 	})
 }
 
+// isAllowedModelName enforces a narrow allowlist pattern for model names
+// accepted by /inference/pull. Ollama model names are of the form
+// "<namespace>/<name>:<tag>" or "<name>:<tag>" — only A-Z a-z 0-9 . _ - / :
+// are permitted, length <= 128. Explicitly rejects ".." to block path
+// traversal payloads from reaching the Ollama daemon.
+func isAllowedModelName(name string) bool {
+	if name == "" || len(name) > 128 {
+		return false
+	}
+	// Block path-traversal sequences outright.
+	if strings.Contains(name, "..") {
+		return false
+	}
+	for _, ch := range name {
+		switch {
+		case ch >= 'a' && ch <= 'z':
+		case ch >= 'A' && ch <= 'Z':
+		case ch >= '0' && ch <= '9':
+		case ch == '.' || ch == '_' || ch == '-' || ch == '/' || ch == ':':
+		default:
+			return false
+		}
+	}
+	return true
+}
+
 // handleInferencePull pulls a model and streams progress to TUI logs
 func (c *ControlServer) handleInferencePull(w http.ResponseWriter, r *http.Request) {
 	if r.Method != http.MethodPost {
 		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
 		return
 	}
 
+	// Cap request body to prevent memory-exhaustion payloads (P0-A hardening).
+	r.Body = http.MaxBytesReader(w, r.Body, 4096)
+
 	// Parse request body
 	var req struct {
 		Model string `json:"model"`
@@ -144,6 +186,16 @@ func (c *ControlServer) handleInferencePull(w http.ResponseWriter, r *http.Reque
 		return
 	}
 
+	// P0-A model-name allowlist: reject anything with shell metacharacters
+	// or path-traversal sequences before forwarding to the Ollama daemon.
+	if !isAllowedModelName(req.Model) {
+		writeJSON(w, http.StatusBadRequest, map[string]any{
+			"status": "error",
+			"error":  "Invalid model name",
+		})
+		return
+	}
+
 	log.Info().Str("model", req.Model).Msg("Control: pull model command received")
 	c.agent.state.AddLog(fmt.Sprintf("Pulling model: %s", req.Model))
 
@@ -153,7 +205,7 @@ func (c *ControlServer) handleInferencePull(w http.ResponseWriter, r *http.Reque
 
 	client := &http.Client{Timeout: 30 * time.Minute} // Model pulls can take a while
 	resp, err := client.Post(
-		fmt.Sprintf("http://localhost:%d/api/pull", DefaultOllamaPort),
+		fmt.Sprintf("%s/api/pull", c.ollamaURL()),
 		"application/json",
 		bytes.NewReader(body),
 	)
@@ -167,9 +219,30 @@ func (c *ControlServer) handleInferencePull(w http.ResponseWriter, r *http.Reque
 	}
 	defer resp.Body.Close()
 
-	// Stream progress to logs
+	// Check Ollama's status BEFORE streaming the response body. Previously a
+	// non-200 from Ollama still produced StatusOK downstream because the
+	// decode loop swallowed errors and the handler unconditionally returned
+	// ok (Cubic finding on pkg/agent/control.go:227).
+	if resp.StatusCode != http.StatusOK {
+		bodyPreview := make([]byte, 512)
+		n, _ := resp.Body.Read(bodyPreview)
+		errMsg := string(bodyPreview[:n])
+		c.agent.state.AddLog(fmt.Sprintf("Pull failed: Ollama returned HTTP %d", resp.StatusCode))
+		writeJSON(w, http.StatusBadGateway, map[string]any{
+			"status":        "error",
+			"error":         fmt.Sprintf("Ollama returned HTTP %d", resp.StatusCode),
+			"ollama_body":   errMsg,
+			"ollama_status": resp.StatusCode,
+		})
+		return
+	}
+
+	// Stream progress to logs. Track decode errors so a malformed stream
+	// surfaces as a 502 rather than silent success.
 	decoder := json.NewDecoder(resp.Body)
 	lastStatus := ""
+	var decodeErr error
+	sawAny := false
 	for {
 		var progress struct {
 			Status    string `json:"status"`
@@ -178,8 +251,14 @@ func (c *ControlServer) handleInferencePull(w http.ResponseWriter, r *http.Reque
 			Completed int64  `json:"completed"`
 		}
 		if err := decoder.Decode(&progress); err != nil {
+			// io.EOF is a clean end of stream; any other error is a decode
+			// failure that should be reported upstream.
+			if err.Error() != "EOF" {
+				decodeErr = err
+			}
 			break
 		}
+		sawAny = true
 
 		// Update logs with progress (avoid duplicate messages)
 		status := progress.Status
@@ -193,6 +272,19 @@ func (c *ControlServer) handleInferencePull(w http.ResponseWriter, r *http.Reque
 		}
 	}
 
+	if decodeErr != nil || !sawAny {
+		reason := "no progress frames received"
+		if decodeErr != nil {
+			reason = decodeErr.Error()
+		}
+		c.agent.state.AddLog(fmt.Sprintf("Pull failed: %s", reason))
+		writeJSON(w, http.StatusBadGateway, map[string]any{
+			"status": "error",
+			"error":  fmt.Sprintf("failed to decode Ollama response: %s", reason),
+		})
+		return
+	}
+
 	c.agent.state.AddLog(fmt.Sprintf("Model %s ready", req.Model))
 
 	// Update models list
@@ -201,7 +293,7 @@ func (c *ControlServer) handleInferencePull(w http.ResponseWriter, r *http.Reque
 		c.agent.state.UpdateInference("running", c.agent.ollama.TailscaleIP(), DefaultOllamaPort, models)
 	}
 
-	writeJSON(w, http.StatusOK, map[string]any{
+	writeJSON(w, http.StatusCreated, map[string]any{
 		"status":  "ok",
 		"message": fmt.Sprintf("Model %s pulled successfully", req.Model),
 	})

pkg/agent/control_test.go

@@ -0,0 +1,120 @@
+package agent
+
+import (
+	"bytes"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+// newTestControlServer wires a minimal Agent with state + OllamaManager so the
+// /inference/pull handler's state touches are safe without starting a TUI or
+// keepalive loop.
+func newTestControlServer(t *testing.T, ollamaBaseURL string) *ControlServer {
+	t.Helper()
+	state := NewAgentState("test-machine", "test-pool", "http://gateway")
+	agent := &Agent{
+		state:  state,
+		ollama: NewOllamaManager("127.0.0.1", 0),
+	}
+	cs := NewControlServer(agent, 0)
+	cs.ollamaBaseURL = ollamaBaseURL
+	return cs
+}
+
+// TestInferencePullStatus verifies that the handler only reports success when
+// Ollama returns HTTP 200 AND the NDJSON stream decodes cleanly with at least
+// one progress frame. Any non-200 upstream must become HTTP 502 at the
+// control-plane boundary (Cubic finding on pkg/agent/control.go:227).
+func TestInferencePullStatus(t *testing.T) {
+	// Valid NDJSON progress-stream frames.
+	okFrames := "{\"status\":\"pulling manifest\"}\n{\"status\":\"success\"}\n"
+
+	cases := []struct {
+		name             string
+		upstreamStatus   int
+		upstreamBody     string
+		wantHandlerCode  int
+	}{
+		{"ollama 200 returns 201", http.StatusOK, okFrames, http.StatusCreated},
+		{"ollama 404 surfaces 502", http.StatusNotFound, "{\"error\":\"model not found\"}", http.StatusBadGateway},
+		{"ollama 500 surfaces 502", http.StatusInternalServerError, "{\"error\":\"boom\"}", http.StatusBadGateway},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			ollama := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(tc.upstreamStatus)
+				_, _ = w.Write([]byte(tc.upstreamBody))
+			}))
+			defer ollama.Close()
+
+			cs := newTestControlServer(t, ollama.URL)
+
+			body, _ := json.Marshal(map[string]string{"model": "llama3:8b"})
+			req := httptest.NewRequest(http.MethodPost, "/inference/pull", bytes.NewReader(body))
+			rr := httptest.NewRecorder()
+			cs.handleInferencePull(rr, req)
+
+			if rr.Code != tc.wantHandlerCode {
+				t.Fatalf("handler code = %d, want %d; body=%s", rr.Code, tc.wantHandlerCode, rr.Body.String())
+			}
+		})
+	}
+}
+
+// TestInferencePullDecodeFailureReturnsBadGateway covers the case where
+// Ollama returns 200 but the body is not NDJSON. The handler must not
+// pretend success — it must surface a 502.
+func TestInferencePullDecodeFailureReturnsBadGateway(t *testing.T) {
+	ollama := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte("this is not valid json at all"))
+	}))
+	defer ollama.Close()
+
+	cs := newTestControlServer(t, ollama.URL)
+
+	body, _ := json.Marshal(map[string]string{"model": "llama3:8b"})
+	req := httptest.NewRequest(http.MethodPost, "/inference/pull", bytes.NewReader(body))
+	rr := httptest.NewRecorder()
+	cs.handleInferencePull(rr, req)
+
+	if rr.Code != http.StatusBadGateway {
+		t.Fatalf("decode failure should produce 502, got %d; body=%s", rr.Code, rr.Body.String())
+	}
+}
+
+// TestInferencePullRejectsInvalidModelName asserts the P0-A allowlist blocks
+// shell-metacharacter payloads before they reach Ollama.
+func TestInferencePullRejectsInvalidModelName(t *testing.T) {
+	cs := newTestControlServer(t, "http://127.0.0.1:1") // unreachable on purpose
+
+	payloads := []string{
+		"llama3; rm -rf /",
+		"../../etc/passwd",
+		"model`whoami`",
+		"",
+	}
+	for _, p := range payloads {
+		body, _ := json.Marshal(map[string]string{"model": p})
+		req := httptest.NewRequest(http.MethodPost, "/inference/pull", bytes.NewReader(body))
+		rr := httptest.NewRecorder()
+		cs.handleInferencePull(rr, req)
+		if rr.Code != http.StatusBadRequest {
+			t.Fatalf("payload %q should be rejected with 400, got %d", p, rr.Code)
+		}
+	}
+}
+
+// TestInferencePullMethodNotAllowed checks GET is refused.
+func TestInferencePullMethodNotAllowed(t *testing.T) {
+	cs := newTestControlServer(t, "")
+	req := httptest.NewRequest(http.MethodGet, "/inference/pull", nil)
+	rr := httptest.NewRecorder()
+	cs.handleInferencePull(rr, req)
+	if rr.Code != http.StatusMethodNotAllowed {
+		t.Fatalf("GET should be 405, got %d", rr.Code)
+	}
+}

pkg/agent/keepalive.go

@@ -41,6 +41,14 @@ type KeepaliveLoop struct {
 	maxFailures         int32
 	stopCh              chan struct{}
 	doneCh              chan struct{}
+
+	// startOnce guards Start so a second Start() after Stop is a no-op
+	// rather than re-launching the goroutine against an already-closed
+	// stopCh. stopOnce guards Stop so a double-Stop does not panic with
+	// "close of closed channel" (Cubic finding, pkg/agent/keepalive.go).
+	startOnce sync.Once
+	stopOnce  sync.Once
+	started   int32 // atomic flag: 1 after Start has been invoked
 }
 
 // NewKeepaliveLoop creates a new keepalive loop
@@ -76,9 +84,14 @@ func (k *KeepaliveLoop) GetLastMetrics() *MachineMetrics {
 	return &MachineMetrics{}
 }
 
-// Start begins the keepalive loop in a goroutine
+// Start begins the keepalive loop in a goroutine. Calling Start more than
+// once is a no-op; Start after Stop refuses to re-launch (safe on concurrent
+// invocation).
 func (k *KeepaliveLoop) Start(ctx context.Context) {
-	go k.run(ctx)
+	k.startOnce.Do(func() {
+		atomic.StoreInt32(&k.started, 1)
+		go k.run(ctx)
+	})
 }
 
 func (k *KeepaliveLoop) run(ctx context.Context) {
@@ -108,15 +121,22 @@ func (k *KeepaliveLoop) run(ctx context.Context) {
 	}
 }
 
-// Stop signals the keepalive loop to stop and waits for it to finish
+// Stop signals the keepalive loop to stop and waits for it to finish.
+// Safe to call multiple times concurrently; subsequent calls are no-ops.
 func (k *KeepaliveLoop) Stop() {
-	close(k.stopCh)
-	// Wait for loop to finish with timeout
-	select {
-	case <-k.doneCh:
-	case <-time.After(5 * time.Second):
-		log.Warn().Msg("Keepalive loop did not stop within timeout")
-	}
+	k.stopOnce.Do(func() {
+		close(k.stopCh)
+		// Wait for loop to finish with timeout. If Start was never called
+		// there is no goroutine to wait for; skip the wait in that case.
+		if atomic.LoadInt32(&k.started) == 0 {
+			return
+		}
+		select {
+		case <-k.doneCh:
+		case <-time.After(5 * time.Second):
+			log.Warn().Msg("Keepalive loop did not stop within timeout")
+		}
+	})
 }
 
 // IsHealthy returns true if recent keepalives succeeded