oto/services/core/internal/httpserver/server_test.go

package httpserver

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/toki/oto/services/core/internal/cicdstate"
	"github.com/toki/oto/services/core/internal/runnerregistry"
	otopb "github.com/toki/oto/services/core/oto"
)

func TestHandleHealthz(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
	rr := httptest.NewRecorder()

	handleHealthz(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Errorf("handleHealthz returned wrong status code: got %v want %v", status, http.StatusOK)
	}

	expected := "OK"
	if body := rr.Body.String(); body != expected {
		t.Errorf("handleHealthz returned unexpected body: got %v want %v", body, expected)
	}
}

func TestHandleHealthz_MethodNotAllowed(t *testing.T) {
	req := httptest.NewRequest(http.MethodPost, "/healthz", nil)
	rr := httptest.NewRecorder()

	handleHealthz(rr, req)

	if status := rr.Code; status != http.StatusMethodNotAllowed {
		t.Errorf("handleHealthz for POST returned wrong status code: got %v want %v", status, http.StatusMethodNotAllowed)
	}
}

func TestHandleReadyz(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/readyz", nil)
	rr := httptest.NewRecorder()

	handleReadyz(rr, req)

	if status := rr.Code; status != http.StatusOK {
		t.Errorf("handleReadyz returned wrong status code: got %v want %v", status, http.StatusOK)
	}

	expected := "OK"
	if body := rr.Body.String(); body != expected {
		t.Errorf("handleReadyz returned unexpected body: got %v want %v", body, expected)
	}
}

func TestHandleReadyz_MethodNotAllowed(t *testing.T) {
	req := httptest.NewRequest(http.MethodPost, "/readyz", nil)
	rr := httptest.NewRecorder()

	handleReadyz(rr, req)

	if status := rr.Code; status != http.StatusMethodNotAllowed {
		t.Errorf("handleReadyz for POST returned wrong status code: got %v want %v", status, http.StatusMethodNotAllowed)
	}
}

func TestServerMux(t *testing.T) {
	registry := runnerregistry.New()
	server := NewServerWithRegistry(":0", registry)
	mux := server.httpServer.Handler.(*http.ServeMux)

	// Test healthz routing
	reqHealth := httptest.NewRequest(http.MethodGet, "/healthz", nil)
	rrHealth := httptest.NewRecorder()
	mux.ServeHTTP(rrHealth, reqHealth)
	if rrHealth.Code != http.StatusOK {
		t.Errorf("mux /healthz failed: got status %v", rrHealth.Code)
	}

	// Test readyz routing
	reqReady := httptest.NewRequest(http.MethodGet, "/readyz", nil)
	rrReady := httptest.NewRecorder()
	mux.ServeHTTP(rrReady, reqReady)
	if rrReady.Code != http.StatusOK {
		t.Errorf("mux /readyz failed: got status %v", rrReady.Code)
	}

	// Test fallback/not found routing
	reqNotFound := httptest.NewRequest(http.MethodGet, "/unknown", nil)
	rrNotFound := httptest.NewRecorder()
	mux.ServeHTTP(rrNotFound, reqNotFound)
	if rrNotFound.Code != http.StatusNotFound {
		t.Errorf("mux /unknown routing status code: got %v want %v", rrNotFound.Code, http.StatusNotFound)
	}
}

func TestHandleRunnerRegisterAcceptsAndStoresRunner(t *testing.T) {
	registry := runnerregistry.New()
	body := bytes.NewBufferString(`{
		"enrollment_token":"token-123",
		"runner_id":"runner-123",
		"alias":"linux-build",
		"protocol_version":"oto.runner.v1",
		"capability":{"name":"oto-runner","version":"1.0.0"},
		"command_catalog":{"command_types":["Shell","Git"]}
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/register", body)
	rr := httptest.NewRecorder()

	handleRunnerRegister(registry)(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String())
	}
	var response otopb.RegisterRunnerResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if !response.GetAccepted() {
		t.Fatalf("accepted=false, reason=%q", response.GetRejectReason())
	}
	record, ok := registry.Snapshot("runner-123")
	if !ok {
		t.Fatal("registered runner was not stored")
	}
	if record.Alias != "linux-build" {
		t.Fatalf("alias = %q, want linux-build", record.Alias)
	}
	if got := record.CommandTypes; len(got) != 2 || got[0] != "Shell" || got[1] != "Git" {
		t.Fatalf("command types = %#v, want [Shell Git]", got)
	}
}

func TestHandleRunnerRegisterRejectsMissingToken(t *testing.T) {
	registry := runnerregistry.New()
	body := bytes.NewBufferString(`{
		"runner_id":"runner-123",
		"protocol_version":"oto.runner.v1"
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/register", body)
	rr := httptest.NewRecorder()

	handleRunnerRegister(registry)(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}
	var response otopb.RegisterRunnerResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if response.GetAccepted() {
		t.Fatal("missing token registration was accepted")
	}
	if response.GetRejectReason() != "missing enrollment token" {
		t.Fatalf("reject reason = %q, want missing enrollment token", response.GetRejectReason())
	}
	if response.GetError() == nil {
		t.Fatal("missing structured error")
	}
	if response.GetError().GetCode() != "registration_rejected" {
		t.Fatalf("error code = %q, want registration_rejected", response.GetError().GetCode())
	}
}

func TestHandleRunnerRegisterMethodNotAllowed(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/api/v1/runners/register", nil)
	rr := httptest.NewRecorder()

	handleRunnerRegister(runnerregistry.New())(rr, req)

	if rr.Code != http.StatusMethodNotAllowed {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusMethodNotAllowed)
	}
}

func TestHandleRunnerHeartbeat(t *testing.T) {
	registry := runnerregistry.New()

	// 1. Heartbeat on unregistered runner should fail (404)
	body := bytes.NewBufferString(`{"runner_id":"runner-123", "status":1}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/heartbeat", body)
	rr := httptest.NewRecorder()

	handleRunnerHeartbeat(registry)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("status = %v, want %v for unregistered runner", rr.Code, http.StatusNotFound)
	}

	// 2. Register runner first
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	// 3. Valid Heartbeat on registered runner should succeed
	body = bytes.NewBufferString(`{"runner_id":"runner-123", "status":1}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/heartbeat", body)
	rr = httptest.NewRecorder()

	handleRunnerHeartbeat(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	var response otopb.HeartbeatResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode heartbeat response: %v", err)
	}
	if !response.GetSuccess() {
		t.Fatal("heartbeat failed")
	}

	// Verify status is online
	record, _ := registry.Snapshot("runner-123")
	if record.Status != runnerregistry.StatusOnline {
		t.Fatalf("status = %q, want %q", record.Status, runnerregistry.StatusOnline)
	}
}

func TestHandleRunnerDisconnect(t *testing.T) {
	registry := runnerregistry.New()

	// 1. Disconnect unregistered runner should fail (404)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/disconnect", nil)
	rr := httptest.NewRecorder()

	handleRunnerDisconnect(registry)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusNotFound)
	}

	// 2. Register runner
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	// 3. Disconnect registered runner should succeed
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/disconnect", nil)
	rr = httptest.NewRecorder()

	handleRunnerDisconnect(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	var res struct {
		Success bool `json:"success"`
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &res); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if !res.Success {
		t.Fatal("expected success=true")
	}

	record, _ := registry.Snapshot("runner-123")
	if record.Status != runnerregistry.StatusDisconnected {
		t.Fatalf("status = %q, want %q", record.Status, runnerregistry.StatusDisconnected)
	}
}

func TestHandleGetRunner(t *testing.T) {
	registry := runnerregistry.New()

	// 1. Get unregistered runner should return 404
	req := httptest.NewRequest(http.MethodGet, "/api/v1/runners/runner-123", nil)
	rr := httptest.NewRecorder()

	handleGetRunner(registry)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusNotFound)
	}

	// 2. Register runner
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		Alias:           "build-linux",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	// 3. Get registered runner should succeed (200) and return JSON record status
	req = httptest.NewRequest(http.MethodGet, "/api/v1/runners/runner-123", nil)
	rr = httptest.NewRecorder()

	handleGetRunner(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	var res map[string]interface{}
	if err := json.Unmarshal(rr.Body.Bytes(), &res); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if res["runner_id"] != "runner-123" {
		t.Fatalf("runner_id = %v", res["runner_id"])
	}
	if res["status"] != runnerregistry.StatusAccepted {
		t.Fatalf("status = %v", res["status"])
	}
}

func TestHandleRunnerHeartbeatMismatch(t *testing.T) {
	registry := runnerregistry.New()

	// Register runner-a
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-a",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})
	// Register runner-b
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-b",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})

	// Heartbeat with mismatched runner ID in body vs path
	body := bytes.NewBufferString(`{"runner_id":"runner-b", "status":1}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-a/heartbeat", body)
	rr := httptest.NewRecorder()

	handleRunnerHeartbeat(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusBadRequest)
	}

	var response otopb.HeartbeatResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if response.GetSuccess() {
		t.Fatal("expected heartbeat to fail due to ID mismatch")
	}
	if response.GetErrorMessage() != "runner id mismatch between path and body" {
		t.Fatalf("error message = %q", response.GetErrorMessage())
	}

	// Verify runner-b remains accepted (not revived/online)
	rec, _ := registry.Snapshot("runner-b")
	if rec.Status != runnerregistry.StatusAccepted {
		t.Fatalf("runner-b status = %q, want %q", rec.Status, runnerregistry.StatusAccepted)
	}
}

func TestHandleRunnerBootstrapCommand(t *testing.T) {
	registry := runnerregistry.New()

	// 1. Missing runner_id or enrollment_token should be rejected
	bodyMissing := bytes.NewBufferString(`{"runner_id":"","enrollment_token":""}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyMissing)
	rr := httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("status = %v, want %v for missing runner_id/token", rr.Code, http.StatusBadRequest)
	}

	// 2. Missing token only should be rejected
	bodyNoToken := bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":""}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyNoToken)
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("status = %v, want %v for missing token", rr.Code, http.StatusBadRequest)
	}

	// 3. Request should be rejected if server is HTTP and OTO_RUNNER_RELEASE_BASE_URL is not set
	t.Setenv("OTO_RUNNER_RELEASE_BASE_URL", "")
	bodyValid := bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyValid)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 Bad Request when HTTP and no release override, got %v", rr.Code)
	}

	// 4. Request should be rejected if OTO_RUNNER_RELEASE_BASE_URL is not HTTPS
	t.Setenv("OTO_RUNNER_RELEASE_BASE_URL", "http://example.com/releases")
	bodyValid = bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyValid)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 Bad Request when override is not HTTPS, got %v", rr.Code)
	}

	// 5. Valid request should return correct escaped command when OTO_RUNNER_RELEASE_BASE_URL is set to HTTPS
	t.Setenv("OTO_RUNNER_RELEASE_BASE_URL", "https://example.com/releases")
	bodyValid = bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyValid)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	var response otopb.BootstrapCommandResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode bootstrap command response: %v", err)
	}

	expectedCmd := "curl -fsSL 'http://localhost:8080/bootstrap/oto-agent.sh' | bash -s -- --server-url 'http://localhost:8080' --agent-id 'runner-123' --enrollment-token 'token-123' --release-base-url 'https://example.com/releases'"
	if response.GetBootstrapCommand() != expectedCmd {
		t.Fatalf("bootstrap_command = %q, want %q", response.GetBootstrapCommand(), expectedCmd)
	}

	// 6. Request with shell metacharacters in runner_id/token should be safely escaped
	bodyMalicious := bytes.NewBufferString(`{"runner_id":"runner; rm -rf /","enrollment_token":"token'$(say hello)'"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyMalicious)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}

	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode bootstrap command response: %v", err)
	}

	expectedEscapedCmd := "curl -fsSL 'http://localhost:8080/bootstrap/oto-agent.sh' | bash -s -- --server-url 'http://localhost:8080' --agent-id 'runner; rm -rf /' --enrollment-token 'token'\\''$(say hello)'\\''' --release-base-url 'https://example.com/releases'"
	if response.GetBootstrapCommand() != expectedEscapedCmd {
		t.Fatalf("bootstrap_command = %q, want %q", response.GetBootstrapCommand(), expectedEscapedCmd)
	}

	// 7. Malicious Host header with invalid characters should be rejected
	bodyValid = bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyValid)
	req.Host = "localhost; rm -rf /"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 Bad Request for malicious Host header, got %v", rr.Code)
	}

	// 8. Host with shell metacharacter but no space (e.g. localhost;rm) should also be rejected
	bodyValid = bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyValid)
	req.Host = "localhost;rm"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("expected 400 Bad Request for Host with shell metacharacter (no space), got %v", rr.Code)
	}

	// 9. target="linux" should work and return Unix shell command
	bodyLinux := bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123","target":"linux"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyLinux)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v for target=linux", rr.Code, http.StatusOK)
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode bootstrap command response: %v", err)
	}
	expectedLinuxCmd := "curl -fsSL 'http://localhost:8080/bootstrap/oto-agent.sh' | bash -s -- --server-url 'http://localhost:8080' --agent-id 'runner-123' --enrollment-token 'token-123' --release-base-url 'https://example.com/releases'"
	if response.GetBootstrapCommand() != expectedLinuxCmd {
		t.Fatalf("bootstrap_command for linux = %q, want %q", response.GetBootstrapCommand(), expectedLinuxCmd)
	}

	// 10. target="macos" should work and return Unix shell command
	bodyMacos := bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123","target":"macos"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyMacos)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v for target=macos", rr.Code, http.StatusOK)
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode bootstrap command response: %v", err)
	}
	expectedMacosCmd := "curl -fsSL 'http://localhost:8080/bootstrap/oto-agent.sh' | bash -s -- --server-url 'http://localhost:8080' --agent-id 'runner-123' --enrollment-token 'token-123' --release-base-url 'https://example.com/releases'"
	if response.GetBootstrapCommand() != expectedMacosCmd {
		t.Fatalf("bootstrap_command for macos = %q, want %q", response.GetBootstrapCommand(), expectedMacosCmd)
	}

	// 11. target="windows" should work and return PowerShell command
	bodyWindows := bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123","target":"windows"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyWindows)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v for target=windows", rr.Code, http.StatusOK)
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode bootstrap command response: %v", err)
	}
	expectedWindowsCmd := "powershell -ExecutionPolicy Bypass -Command \"& ([scriptblock]::Create((irm 'http://localhost:8080/bootstrap/oto-agent.ps1'))) -- --server-url 'http://localhost:8080' --agent-id 'runner-123' --enrollment-token 'token-123' --release-base-url 'https://example.com/releases'\""
	if response.GetBootstrapCommand() != expectedWindowsCmd {
		t.Fatalf("bootstrap_command for windows = %q, want %q", response.GetBootstrapCommand(), expectedWindowsCmd)
	}

	// 12. target="windows" with malicious runner_id/token should escape properly for PowerShell
	bodyWindowsMalicious := bytes.NewBufferString(`{"runner_id":"runner; rm -rf /","enrollment_token":"token'$(say hello)'","target":"windows"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyWindowsMalicious)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v for target=windows with malicious fields", rr.Code, http.StatusOK)
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode bootstrap command response: %v", err)
	}
	expectedWindowsEscapedCmd := "powershell -ExecutionPolicy Bypass -Command \"& ([scriptblock]::Create((irm 'http://localhost:8080/bootstrap/oto-agent.ps1'))) -- --server-url 'http://localhost:8080' --agent-id 'runner; rm -rf /' --enrollment-token 'token''$(say hello)''' --release-base-url 'https://example.com/releases'\""
	if response.GetBootstrapCommand() != expectedWindowsEscapedCmd {
		t.Fatalf("bootstrap_command for windows malicious = %q, want %q", response.GetBootstrapCommand(), expectedWindowsEscapedCmd)
	}

	// 13. unsupported target should be rejected with 400 Bad Request
	bodyUnsupported := bytes.NewBufferString(`{"runner_id":"runner-123","enrollment_token":"token-123","target":"freebsd"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/bootstrap-command", bodyUnsupported)
	req.Host = "localhost:8080"
	rr = httptest.NewRecorder()

	handleRunnerBootstrapCommand(registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("status = %v, want %v for unsupported target", rr.Code, http.StatusBadRequest)
	}
	var errResponse map[string]interface{}
	if err := json.Unmarshal(rr.Body.Bytes(), &errResponse); err != nil {
		t.Fatalf("decode error response: %v", err)
	}
	expectedErr := "unsupported bootstrap target: freebsd"
	if errResponse["error"] != expectedErr {
		t.Fatalf("error = %v, want %q", errResponse["error"], expectedErr)
	}
}

// staticProvider is a test-only bootstrapScriptProvider with fixed content.
type staticProvider struct {
	content []byte
	err     error
}

func (p staticProvider) bootstrapScript() ([]byte, error) { return p.content, p.err }

func TestHandleServeBootstrapScript(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/bootstrap/oto-agent.sh", nil)
	rr := httptest.NewRecorder()

	handleServeBootstrapScript(staticProvider{content: []byte("#!/usr/bin/env bash\n")})(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v; body = %s", rr.Code, http.StatusOK, rr.Body.String())
	}

	if contentType := rr.Header().Get("Content-Type"); contentType != "application/x-sh" {
		t.Fatalf("Content-Type = %q, want application/x-sh", contentType)
	}

	if len(rr.Body.Bytes()) == 0 {
		t.Fatal("empty bootstrap script returned")
	}
}

func TestHandleServeBootstrapScript_MethodNotAllowed(t *testing.T) {
	req := httptest.NewRequest(http.MethodPost, "/bootstrap/oto-agent.sh", nil)
	rr := httptest.NewRecorder()

	handleServeBootstrapScript(staticProvider{content: []byte("#!/usr/bin/env bash\n")})(rr, req)
	if rr.Code != http.StatusMethodNotAllowed {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusMethodNotAllowed)
	}
}

func TestHandleServeBootstrapScript_ProviderError(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/bootstrap/oto-agent.sh", nil)
	rr := httptest.NewRecorder()

	handleServeBootstrapScript(staticProvider{err: fmt.Errorf("script unavailable")})(rr, req)
	if rr.Code != http.StatusInternalServerError {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusInternalServerError)
	}
}

func TestEmbeddedBootstrapScriptMatchesRunnerAsset(t *testing.T) {
	// Drift check: embedded copy must be byte-identical to the runner asset.
	// On failure: cp apps/runner/assets/script/shell/oto_agent_bootstrap.sh \
	//               services/core/internal/httpserver/oto_agent_bootstrap.sh
	embedded, err := embeddedBootstrapProvider{}.bootstrapScript()
	if err != nil {
		t.Fatalf("embeddedBootstrapProvider.bootstrapScript() failed: %v", err)
	}

	const runnerAsset = "../../../../apps/runner/assets/script/shell/oto_agent_bootstrap.sh"
	asset, err := os.ReadFile(runnerAsset)
	if err != nil {
		t.Fatalf("failed to read runner asset %s: %v", runnerAsset, err)
	}

	if !bytes.Equal(embedded, asset) {
		t.Fatalf("embedded bootstrap script differs from runner asset %s - "+
			"run: cp %s services/core/internal/httpserver/oto_agent_bootstrap.sh",
			runnerAsset, runnerAsset)
	}
}

func TestEmbeddedBootstrapScriptNonEmpty(t *testing.T) {
	provider := embeddedBootstrapProvider{}
	content, err := provider.bootstrapScript()
	if err != nil {
		t.Fatalf("bootstrapScript() error: %v", err)
	}
	if len(content) == 0 {
		t.Fatal("embedded bootstrap script content is empty")
	}
}

func TestHandleCreateJob(t *testing.T) {
	store := cicdstate.NewStore()

	// Valid job creation should succeed
	body := bytes.NewBufferString(`{"id":"job-1","name":"build"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/jobs", body)
	rr := httptest.NewRecorder()

	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusCreated {
		t.Fatalf("status = %v, want %v; body=%s", rr.Code, http.StatusCreated, rr.Body.String())
	}

	var res map[string]interface{}
	if err := json.Unmarshal(rr.Body.Bytes(), &res); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if res["id"] != "job-1" {
		t.Fatalf("id = %v, want job-1", res["id"])
	}
	if res["name"] != "build" {
		t.Fatalf("name = %v, want build", res["name"])
	}
	if res["state"] != cicdstate.StateQueued {
		t.Fatalf("state = %v, want %s", res["state"], cicdstate.StateQueued)
	}

	// Get job by ID
	req = httptest.NewRequest(http.MethodGet, "/api/v1/jobs/job-1", nil)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("GET status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String())
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &res); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if res["id"] != "job-1" {
		t.Fatalf("GET job id = %v, want job-1", res["id"])
	}

	// Unknown job ID should return 404
	req = httptest.NewRequest(http.MethodGet, "/api/v1/jobs/nonexistent", nil)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("GET unknown job status = %v, want %v", rr.Code, http.StatusNotFound)
	}

	// Duplicate job should return 409
	body = bytes.NewBufferString(`{"id":"job-1","name":"duplicate"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/jobs", body)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusConflict {
		t.Fatalf("duplicate job status = %v, want %v", rr.Code, http.StatusConflict)
	}

	// Missing fields should return 400
	body = bytes.NewBufferString(`{"id":"job-2"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/jobs", body)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("missing fields status = %v, want %v", rr.Code, http.StatusBadRequest)
	}
}

func TestHandleCreateJobStoresRunRequest(t *testing.T) {
	store := cicdstate.NewStore()

	body := bytes.NewBufferString(`{
		"id":"job-run",
		"name":"build",
		"run_request":{
			"pipeline_yaml":"commands:\n  - type: Shell",
			"variables":{"FLAVOR":"release"},
			"command_types":["Shell","Git"]
		}
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/jobs", body)
	rr := httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusCreated {
		t.Fatalf("create status = %v, want %v; body=%s", rr.Code, http.StatusCreated, rr.Body.String())
	}

	// Stored input must survive on the store.
	job, err := store.GetJob("job-run")
	if err != nil {
		t.Fatalf("GetJob failed: %v", err)
	}
	if job.RunInput == nil || job.RunInput.PipelineYAML != "commands:\n  - type: Shell" {
		t.Fatalf("stored RunInput = %+v", job.RunInput)
	}
	if job.RunInput.Variables["FLAVOR"] != "release" {
		t.Fatalf("stored variables = %+v", job.RunInput.Variables)
	}

	// GET /jobs/{id} must expose the stored input.
	req = httptest.NewRequest(http.MethodGet, "/api/v1/jobs/job-run", nil)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("get status = %v, want %v", rr.Code, http.StatusOK)
	}
	var res map[string]interface{}
	if err := json.Unmarshal(rr.Body.Bytes(), &res); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	runReq, ok := res["run_request"].(map[string]interface{})
	if !ok {
		t.Fatalf("run_request missing in GET job: %v", res["run_request"])
	}
	if runReq["pipeline_yaml"] != "commands:\n  - type: Shell" {
		t.Fatalf("GET run_request pipeline_yaml = %v", runReq["pipeline_yaml"])
	}
}

func TestHandleRunnerCicdClaimReturnsRunRequest(t *testing.T) {
	store := cicdstate.NewStore()
	registry := runnerregistry.New()
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})
	if _, err := store.CreateJob("job-1", "build", &cicdstate.RunInput{
		PipelineYAMLPath: "/pipelines/build.yaml",
		CommandTypes:     []string{"Shell"},
	}); err != nil {
		t.Fatalf("CreateJob failed: %v", err)
	}

	claimBody := bytes.NewBufferString(`{
		"runner_id":"runner-123",
		"job_id":"job-1",
		"execution_id":"exec-1"
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/jobs/claim", claimBody)
	rr := httptest.NewRecorder()
	handleRouter(store, registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("claim status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String())
	}
	var claim otopb.JobClaimResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &claim); err != nil {
		t.Fatalf("decode claim response: %v", err)
	}
	if claim.GetRunRequest() == nil {
		t.Fatal("claim run_request is nil")
	}
	if claim.GetRunRequest().GetPipelineYamlPath() != "/pipelines/build.yaml" {
		t.Fatalf("claim run_request pipeline_yaml_path = %q", claim.GetRunRequest().GetPipelineYamlPath())
	}
	if claim.GetRunRequest().GetRunnerId() != "runner-123" ||
		claim.GetRunRequest().GetJobId() != "job-1" ||
		claim.GetRunRequest().GetExecutionId() != "exec-1" {
		t.Fatalf("claim run_request ids = %+v", claim.GetRunRequest())
	}
}

func TestHandleRunnerCicdNextJobClaim(t *testing.T) {
	store := cicdstate.NewStore()
	registry := runnerregistry.New()
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})

	// Scenario 1: Claim when no job is queued.
	claimBodyNoJob := bytes.NewBufferString(`{
		"runner_id":"runner-123",
		"job_id":"",
		"execution_id":""
	}`)
	req1 := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/jobs/claim", claimBodyNoJob)
	rr1 := httptest.NewRecorder()
	handleRouter(store, registry)(rr1, req1)
	if rr1.Code != http.StatusOK {
		t.Fatalf("claim next job status = %v, want %v", rr1.Code, http.StatusOK)
	}
	var claim1 otopb.JobClaimResponse
	if err := json.Unmarshal(rr1.Body.Bytes(), &claim1); err != nil {
		t.Fatalf("decode claim response: %v", err)
	}
	if claim1.GetAccepted() {
		t.Fatal("expected claim to not be accepted when no jobs are queued")
	}

	// Scenario 2: Claim when a job is queued.
	if _, err := store.CreateJob("job-next", "build", &cicdstate.RunInput{
		PipelineYAMLPath: "/pipelines/next.yaml",
		CommandTypes:     []string{"Shell"},
	}); err != nil {
		t.Fatalf("CreateJob failed: %v", err)
	}

	claimBodyQueued := bytes.NewBufferString(`{
		"runner_id":"runner-123",
		"job_id":"",
		"execution_id":""
	}`)
	req2 := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/jobs/claim", claimBodyQueued)
	rr2 := httptest.NewRecorder()
	handleRouter(store, registry)(rr2, req2)
	if rr2.Code != http.StatusOK {
		t.Fatalf("claim next job status = %v, want %v; body=%s", rr2.Code, http.StatusOK, rr2.Body.String())
	}
	var claim2 otopb.JobClaimResponse
	if err := json.Unmarshal(rr2.Body.Bytes(), &claim2); err != nil {
		t.Fatalf("decode claim response: %v", err)
	}
	if !claim2.GetAccepted() {
		t.Fatal("expected claim to be accepted")
	}
	if claim2.GetJobId() != "job-next" {
		t.Fatalf("expected job id 'job-next', got %q", claim2.GetJobId())
	}
	if claim2.GetExecutionId() == "" {
		t.Fatal("expected non-empty generated execution id")
	}
	if claim2.GetRunRequest().GetPipelineYamlPath() != "/pipelines/next.yaml" {
		t.Fatalf("expected run request pipeline_yaml_path '/pipelines/next.yaml', got %q", claim2.GetRunRequest().GetPipelineYamlPath())
	}

	// Check if the job state has changed in the store.
	job, err := store.GetJob("job-next")
	if err != nil {
		t.Fatalf("GetJob failed: %v", err)
	}
	if job.State != cicdstate.StateRunning {
		t.Fatalf("expected job state running, got %q", job.State)
	}
	if job.ExecutionID != claim2.GetExecutionId() {
		t.Fatalf("expected job execution id %q, got %q", claim2.GetExecutionId(), job.ExecutionID)
	}
}

func TestHandleRunnerCicdRejectsInvalidRunInput(t *testing.T) {
	store := cicdstate.NewStore()
	registry := runnerregistry.New()
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})

	// Job with no run request → claim is rejected with 409.
	if _, err := store.CreateJob("job-missing", "build", nil); err != nil {
		t.Fatalf("CreateJob failed: %v", err)
	}
	missingBody := bytes.NewBufferString(`{"runner_id":"runner-123","job_id":"job-missing","execution_id":"exec-1"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/jobs/claim", missingBody)
	rr := httptest.NewRecorder()
	handleRouter(store, registry)(rr, req)
	if rr.Code != http.StatusConflict {
		t.Fatalf("missing run input claim status = %v, want %v; body=%s", rr.Code, http.StatusConflict, rr.Body.String())
	}

	// Job with both pipeline_yaml and pipeline_yaml_path → claim is rejected with 400.
	if _, err := store.CreateJob("job-both", "build", &cicdstate.RunInput{
		PipelineYAML:     "commands: []",
		PipelineYAMLPath: "/pipelines/build.yaml",
	}); err != nil {
		t.Fatalf("CreateJob failed: %v", err)
	}
	bothBody := bytes.NewBufferString(`{"runner_id":"runner-123","job_id":"job-both","execution_id":"exec-2"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/jobs/claim", bothBody)
	rr = httptest.NewRecorder()
	handleRouter(store, registry)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("invalid run input claim status = %v, want %v; body=%s", rr.Code, http.StatusBadRequest, rr.Body.String())
	}

	// The rejected job must remain queued (no state mutation / execution created).
	job, err := store.GetJob("job-both")
	if err != nil {
		t.Fatalf("GetJob failed: %v", err)
	}
	if job.State != cicdstate.StateQueued || job.ExecutionID != "" {
		t.Fatalf("rejected job mutated state = %+v", job)
	}
}

func TestHandleCreateExecution(t *testing.T) {
	store := cicdstate.NewStore()
	store.CreateJob("job-1", "build", nil)

	body := bytes.NewBufferString(`{"id":"exec-1"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/jobs/job-1/executions", body)
	rr := httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusCreated {
		t.Fatalf("status = %v, want %v; body=%s", rr.Code, http.StatusCreated, rr.Body.String())
	}

	var execRes map[string]interface{}
	if err := json.Unmarshal(rr.Body.Bytes(), &execRes); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if execRes["id"] != "exec-1" {
		t.Fatalf("execution id = %v, want exec-1", execRes["id"])
	}

	// Get execution by ID
	req = httptest.NewRequest(http.MethodGet, "/api/v1/executions/exec-1", nil)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("GET execution status = %v, want %v", rr.Code, http.StatusOK)
	}
	if err := json.Unmarshal(rr.Body.Bytes(), &execRes); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if execRes["id"] != "exec-1" {
		t.Fatalf("GET execution id = %v, want exec-1", execRes["id"])
	}

	// Unknown execution should return 404
	req = httptest.NewRequest(http.MethodGet, "/api/v1/executions/nonexistent", nil)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("GET unknown execution status = %v, want %v", rr.Code, http.StatusNotFound)
	}

	// Execution for unknown job should return 400
	body = bytes.NewBufferString(`{"id":"exec-99"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/jobs/unknown/executions", body)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("unknown job execution status = %v, want %v", rr.Code, http.StatusBadRequest)
	}
}

func TestHandleExecutionLogsAndArtifacts(t *testing.T) {
	store := cicdstate.NewStore()
	store.CreateJob("job-1", "build", nil)
	store.CreateExecution("job-1", "exec-1")

	// Append log
	logBody := bytes.NewBufferString(`{"line":"build starting"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/executions/exec-1/logs", logBody)
	rr := httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusCreated {
		t.Fatalf("append log status = %v, want %v; body=%s", rr.Code, http.StatusCreated, rr.Body.String())
	}

	// Append another log
	logBody = bytes.NewBufferString(`{"line":"build done"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/executions/exec-1/logs", logBody)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusCreated {
		t.Fatalf("append log 2 status = %v, want %v", rr.Code, http.StatusCreated)
	}

	// Get logs
	req = httptest.NewRequest(http.MethodGet, "/api/v1/executions/exec-1/logs", nil)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("get logs status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String())
	}
	var logRes map[string]interface{}
	if err := json.Unmarshal(rr.Body.Bytes(), &logRes); err != nil {
		t.Fatalf("decode log response: %v", err)
	}
	logsArr, ok := logRes["logs"].([]interface{})
	if !ok {
		t.Fatalf("logs key not an array")
	}
	if len(logsArr) != 2 {
		t.Fatalf("log count = %d, want 2", len(logsArr))
	}

	// Append artifact
	artifactBody := bytes.NewBufferString(`{"name":"binary","path":"/dist/app"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/executions/exec-1/artifacts", artifactBody)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusCreated {
		t.Fatalf("append artifact status = %v, want %v; body=%s", rr.Code, http.StatusCreated, rr.Body.String())
	}

	// Get artifacts
	req = httptest.NewRequest(http.MethodGet, "/api/v1/executions/exec-1/artifacts", nil)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("get artifacts status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String())
	}
	var artRes map[string]interface{}
	if err := json.Unmarshal(rr.Body.Bytes(), &artRes); err != nil {
		t.Fatalf("decode artifact response: %v", err)
	}
	artifactsArr, ok := artRes["artifacts"].([]interface{})
	if !ok {
		t.Fatalf("artifacts key not an array")
	}
	if len(artifactsArr) != 1 {
		t.Fatalf("artifact count = %d, want 1", len(artifactsArr))
	}

	// Unknown execution logs/artifacts should return 404
	logBody2 := bytes.NewBufferString(`{"line":"second log line"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/executions/unknown/logs", logBody2)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("unknown log append status = %v, want %v", rr.Code, http.StatusNotFound)
	}

	artifactBody2 := bytes.NewBufferString(`{"name":"binary","path":"/dist/app"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/executions/unknown/artifacts", artifactBody2)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("unknown artifact append status = %v, want %v", rr.Code, http.StatusNotFound)
	}
}

func TestHandleCicdUnknownIds(t *testing.T) {
	store := cicdstate.NewStore()

	// Get unknown job -> 404
	req := httptest.NewRequest(http.MethodGet, "/api/v1/jobs/nonexistent", nil)
	rr := httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("GET unknown job = %v, want %v", rr.Code, http.StatusNotFound)
	}

	// Create execution for unknown job -> 400
	body := bytes.NewBufferString(`{"id":"exec-1"}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/jobs/nonexistent/executions", body)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusBadRequest {
		t.Fatalf("exec unknown job = %v, want %v", rr.Code, http.StatusBadRequest)
	}

	// Get unknown execution -> 404
	req = httptest.NewRequest(http.MethodGet, "/api/v1/executions/nonexistent", nil)
	rr = httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("GET unknown exec = %v, want %v", rr.Code, http.StatusNotFound)
	}
}

func TestHandleCicdMethodNotAllowed(t *testing.T) {
	store := cicdstate.NewStore()

	// GET on /api/v1/jobs (create endpoint) should return method not allowed
	req := httptest.NewRequest(http.MethodGet, "/api/v1/jobs", nil)
	rr := httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusMethodNotAllowed {
		t.Fatalf("GET /api/v1/jobs = %v, want %v", rr.Code, http.StatusMethodNotAllowed)
	}
}

func TestHandleCicdNotFound(t *testing.T) {
	store := cicdstate.NewStore()

	// Unknown path should return 404
	req := httptest.NewRequest(http.MethodGet, "/api/v1/unknown", nil)
	rr := httptest.NewRecorder()
	handleRouter(store, nil)(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("unknown path = %v, want %v", rr.Code, http.StatusNotFound)
	}
}

func TestHandleRunnerCicdClaimReportLogsAndArtifacts(t *testing.T) {
	store := cicdstate.NewStore()
	registry := runnerregistry.New()
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-123",
		ProtocolVersion: "oto.runner.v1",
		Capability: &otopb.RunnerCapability{
			Name:    "oto-runner",
			Version: "1.0.0",
		},
	})
	if _, err := store.CreateJob("job-1", "build", &cicdstate.RunInput{
		PipelineYAML: "commands:\n  - type: Shell",
		Variables:    map[string]string{"FLAVOR": "release"},
		CommandTypes: []string{"Shell", "Git"},
	}); err != nil {
		t.Fatalf("CreateJob failed: %v", err)
	}

	claimBody := bytes.NewBufferString(`{
		"runner_id":"runner-123",
		"job_id":"job-1",
		"execution_id":"exec-1"
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/jobs/claim", claimBody)
	rr := httptest.NewRecorder()
	handleRouter(store, registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("claim status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String())
	}
	var claim otopb.JobClaimResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &claim); err != nil {
		t.Fatalf("decode claim response: %v", err)
	}
	if !claim.GetAccepted() || claim.GetState() != cicdstate.StateRunning {
		t.Fatalf("claim response = %+v, want accepted running", claim)
	}
	if claim.GetRunRequest().GetPipelineYaml() != "commands:\n  - type: Shell" {
		t.Fatalf("claim run_request pipeline yaml = %q", claim.GetRunRequest().GetPipelineYaml())
	}
	if claim.GetRunRequest().GetJobId() != "job-1" || claim.GetRunRequest().GetExecutionId() != "exec-1" || claim.GetRunRequest().GetRunnerId() != "runner-123" {
		t.Fatalf("claim run_request ids = %+v", claim.GetRunRequest())
	}

	job, err := store.GetJob("job-1")
	if err != nil {
		t.Fatalf("GetJob after claim failed: %v", err)
	}
	if job.State != cicdstate.StateRunning || job.ExecutionID != "exec-1" {
		t.Fatalf("job after claim = %+v", job)
	}

	logBody := bytes.NewBufferString(`{
		"runner_id":"runner-123",
		"execution_id":"exec-1",
		"line":"build started"
	}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/executions/exec-1/logs", logBody)
	rr = httptest.NewRecorder()
	handleRouter(store, registry)(rr, req)
	if rr.Code != http.StatusCreated {
		t.Fatalf("append log status = %v, want %v; body=%s", rr.Code, http.StatusCreated, rr.Body.String())
	}

	artifactBody := bytes.NewBufferString(`{
		"runner_id":"runner-123",
		"execution_id":"exec-1",
		"name":"binary",
		"path":"/dist/app"
	}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/executions/exec-1/artifacts", artifactBody)
	rr = httptest.NewRecorder()
	handleRouter(store, registry)(rr, req)
	if rr.Code != http.StatusCreated {
		t.Fatalf("append artifact status = %v, want %v; body=%s", rr.Code, http.StatusCreated, rr.Body.String())
	}

	reportBody := bytes.NewBufferString(`{
		"runner_id":"runner-123",
		"job_id":"job-1",
		"execution_id":"exec-1",
		"success":true,
		"exit_code":0,
		"message":"Build completed successfully.",
		"step_events":[{"event":"completed"}]
	}`)
	req = httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-123/executions/exec-1/report", reportBody)
	rr = httptest.NewRecorder()
	handleRouter(store, registry)(rr, req)
	if rr.Code != http.StatusOK {
		t.Fatalf("report status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String())
	}
	var report otopb.ExecutionReportResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &report); err != nil {
		t.Fatalf("decode report response: %v", err)
	}
	if !report.GetAccepted() || report.GetState() != cicdstate.StateSucceeded {
		t.Fatalf("report response = %+v, want accepted succeeded", report)
	}

	exec, err := store.GetExecution("exec-1")
	if err != nil {
		t.Fatalf("GetExecution after report failed: %v", err)
	}
	if exec.State != cicdstate.StateSucceeded {
		t.Fatalf("execution state = %q, want %q", exec.State, cicdstate.StateSucceeded)
	}
	if len(exec.Logs) != 2 {
		t.Fatalf("log count = %d, want explicit log plus report message", len(exec.Logs))
	}
	if len(exec.Artifacts) != 1 || exec.Artifacts[0].Name != "binary" {
		t.Fatalf("artifacts = %+v, want binary artifact", exec.Artifacts)
	}

	job, err = store.GetJob("job-1")
	if err != nil {
		t.Fatalf("GetJob after report failed: %v", err)
	}
	if job.State != cicdstate.StateSucceeded {
		t.Fatalf("job state = %q, want %q", job.State, cicdstate.StateSucceeded)
	}
}

func TestHandleRunnerCicdRejectsUnknownRunner(t *testing.T) {
	store := cicdstate.NewStore()
	if _, err := store.CreateJob("job-1", "build", nil); err != nil {
		t.Fatalf("CreateJob failed: %v", err)
	}

	body := bytes.NewBufferString(`{
		"runner_id":"missing-runner",
		"job_id":"job-1",
		"execution_id":"exec-1"
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/missing-runner/jobs/claim", body)
	rr := httptest.NewRecorder()
	handleRouter(store, runnerregistry.New())(rr, req)
	if rr.Code != http.StatusNotFound {
		t.Fatalf("status = %v, want %v; body=%s", rr.Code, http.StatusNotFound, rr.Body.String())
	}
}

func TestHandleRunnerRegisterRejectsIncompatibleProtocol(t *testing.T) {
	registry := runnerregistry.New()
	body := bytes.NewBufferString(`{
		"enrollment_token":"token-123",
		"runner_id":"runner-123",
		"protocol_version":"oto.runner.v2",
		"capability":{"name":"oto-runner","version":"1.0.0"}
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/register", body)
	rr := httptest.NewRecorder()

	handleRunnerRegister(registry)(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}
	var response otopb.RegisterRunnerResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if response.GetAccepted() {
		t.Fatal("incompatible protocol registration was accepted")
	}
	if response.GetRejectReason() != `unsupported protocol version: "oto.runner.v2"` {
		t.Fatalf("reject reason = %q", response.GetRejectReason())
	}
	if response.GetError() == nil {
		t.Fatal("missing structured error")
	}
	if response.GetError().GetCode() != "incompatible_runner" {
		t.Fatalf("error code = %q, want incompatible_runner", response.GetError().GetCode())
	}
}

func TestHandleRunnerRegisterRejectsIncompatibleCapability(t *testing.T) {
	registry := runnerregistry.New()
	body := bytes.NewBufferString(`{
		"enrollment_token":"token-123",
		"runner_id":"runner-123",
		"protocol_version":"oto.runner.v1",
		"capability":{"name":"oto-runner","version":"2.0.0"}
	}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/register", body)
	rr := httptest.NewRecorder()

	handleRunnerRegister(registry)(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status = %v, want %v", rr.Code, http.StatusOK)
	}
	var response otopb.RegisterRunnerResponse
	if err := json.Unmarshal(rr.Body.Bytes(), &response); err != nil {
		t.Fatalf("decode response: %v", err)
	}
	if response.GetAccepted() {
		t.Fatal("incompatible capability version registration was accepted")
	}
	if response.GetRejectReason() != `incompatible capability major version: "2.0.0", want major version 1` {
		t.Fatalf("reject reason = %q", response.GetRejectReason())
	}
	if response.GetError() == nil {
		t.Fatal("missing structured error")
	}
	if response.GetError().GetCode() != "incompatible_runner" {
		t.Fatalf("error code = %q, want incompatible_runner", response.GetError().GetCode())
	}
}

func TestHandleRunnerCancelExecution(t *testing.T) {
	store := cicdstate.NewStore()
	registry := runnerregistry.New()

	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-1",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})

	store.CreateJob("job-1", "build", nil)
	store.CreateExecution("job-1", "exec-1")
	store.SetExecutionRunnerID("exec-1", "runner-1")
	store.TransitionJob("job-1", cicdstate.StateRunning)
	store.TransitionExecution("exec-1", cicdstate.StateRunning)

	// 1. Success cancellation
	body := bytes.NewBufferString(`{"runner_id":"runner-1", "execution_id":"exec-1", "reason":"user request"}`)
	req := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-1/executions/exec-1/cancel", body)
	rr := httptest.NewRecorder()
	handleRouter(store, registry)(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("cancel status = %v, want %v; body=%s", rr.Code, http.StatusOK, rr.Body.String())
	}
	var cancelResp otopb.CancelRunResponse
	json.Unmarshal(rr.Body.Bytes(), &cancelResp)
	if !cancelResp.Success {
		t.Fatal("expected success cancel response")
	}

	job, _ := store.GetJob("job-1")
	if job.State != cicdstate.StateCanceled {
		t.Fatalf("job state = %q, want canceled", job.State)
	}

	// 2. Mismatch cancellation
	bodyMismatch := bytes.NewBufferString(`{"runner_id":"runner-wrong", "execution_id":"exec-1"}`)
	reqMismatch := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-1/executions/exec-1/cancel", bodyMismatch)
	rrMismatch := httptest.NewRecorder()
	handleRouter(store, registry)(rrMismatch, reqMismatch)
	if rrMismatch.Code != http.StatusBadRequest {
		t.Fatalf("mismatch cancel status = %v, want %v", rrMismatch.Code, http.StatusBadRequest)
	}

	// 3. Terminal state cancellation rejection
	store.CreateJob("job-2", "build", nil)
	store.CreateExecution("job-2", "exec-2")
	store.SetExecutionRunnerID("exec-2", "runner-1")
	store.TransitionJob("job-2", cicdstate.StateRunning)
	store.TransitionExecution("exec-2", cicdstate.StateRunning)
	store.TransitionJob("job-2", cicdstate.StateSucceeded)
	store.TransitionExecution("exec-2", cicdstate.StateSucceeded)

	bodyTerminal := bytes.NewBufferString(`{"runner_id":"runner-1", "execution_id":"exec-2"}`)
	reqTerminal := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-1/executions/exec-2/cancel", bodyTerminal)
	rrTerminal := httptest.NewRecorder()
	handleRouter(store, registry)(rrTerminal, reqTerminal)
	if rrTerminal.Code != http.StatusBadRequest {
		t.Fatalf("terminal cancel status = %v, want %v; body=%s", rrTerminal.Code, http.StatusBadRequest, rrTerminal.Body.String())
	}

	// 4. Cross-runner execution cancel rejection
	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-2",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})

	store.CreateJob("job-3", "build", nil)
	store.CreateExecution("job-3", "exec-3")
	store.SetExecutionRunnerID("exec-3", "runner-1")
	store.TransitionJob("job-3", cicdstate.StateRunning)
	store.TransitionExecution("exec-3", cicdstate.StateRunning)

	bodyCross := bytes.NewBufferString(`{"runner_id":"runner-2", "execution_id":"exec-3"}`)
	reqCross := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-2/executions/exec-3/cancel", bodyCross)
	rrCross := httptest.NewRecorder()
	handleRouter(store, registry)(rrCross, reqCross)
	if rrCross.Code != http.StatusBadRequest {
		t.Fatalf("cross-runner cancel status = %v, want %v; body=%s", rrCross.Code, http.StatusBadRequest, rrCross.Body.String())
	}

	// Verify job and execution remain running
	job3, _ := store.GetJob("job-3")
	if job3.State != cicdstate.StateRunning {
		t.Fatalf("expected job state to remain running, got %q", job3.State)
	}
	exec3, _ := store.GetExecution("exec-3")
	if exec3.State != cicdstate.StateRunning {
		t.Fatalf("expected execution state to remain running, got %q", exec3.State)
	}
}

func TestHandleRunnerStatus(t *testing.T) {
	store := cicdstate.NewStore()
	registry := runnerregistry.New()

	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-1",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})

	store.CreateJob("job-1", "build", nil)
	store.CreateExecution("job-1", "exec-1")
	store.SetExecutionRunnerID("exec-1", "runner-1")

	req := httptest.NewRequest(http.MethodGet, "/api/v1/runners/runner-1/status", nil)
	rr := httptest.NewRecorder()
	handleRouter(store, registry)(rr, req)

	if rr.Code != http.StatusOK {
		t.Fatalf("status code = %v, want %v", rr.Code, http.StatusOK)
	}

	var res map[string]interface{}
	json.Unmarshal(rr.Body.Bytes(), &res)
	if res["runner_id"] != "runner-1" {
		t.Fatalf("runner_id = %v", res["runner_id"])
	}
	if res["current_execution_id"] != "exec-1" {
		t.Fatalf("current_execution_id = %v", res["current_execution_id"])
	}
	if res["current_job_id"] != "job-1" {
		t.Fatalf("current_job_id = %v", res["current_job_id"])
	}
}

func TestHandleRunnerSelfUpdate(t *testing.T) {
	store := cicdstate.NewStore()
	registry := runnerregistry.New()

	registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "token-123",
		RunnerId:        "runner-1",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})

	// 1. Invalid URL (not HTTPS)
	bodyHttp := bytes.NewBufferString(`{"runner_id":"runner-1", "version":"2.0.0", "download_url":"http://example.com/binary"}`)
	reqHttp := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-1/self-update", bodyHttp)
	rrHttp := httptest.NewRecorder()
	handleRouter(store, registry)(rrHttp, reqHttp)
	if rrHttp.Code != http.StatusBadRequest {
		t.Fatalf("http url update status = %v, want %v", rrHttp.Code, http.StatusBadRequest)
	}

	// 2. Idle accepted
	bodyOk := bytes.NewBufferString(`{"runner_id":"runner-1", "version":"2.0.0", "download_url":"https://example.com/binary"}`)
	reqOk := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-1/self-update", bodyOk)
	rrOk := httptest.NewRecorder()
	handleRouter(store, registry)(rrOk, reqOk)
	if rrOk.Code != http.StatusOK {
		t.Fatalf("idle update status = %v, want %v", rrOk.Code, http.StatusOK)
	}
	var resOk map[string]interface{}
	json.Unmarshal(rrOk.Body.Bytes(), &resOk)
	if resOk["accepted"] != true || resOk["deferred"] != false {
		t.Fatalf("idle update response = %+v", resOk)
	}

	// 3. Active execution defer
	store.CreateJob("job-1", "build", nil)
	store.CreateExecution("job-1", "exec-1")
	store.SetExecutionRunnerID("exec-1", "runner-1")
	store.TransitionJob("job-1", cicdstate.StateRunning)
	store.TransitionExecution("exec-1", cicdstate.StateRunning)

	bodyActive := bytes.NewBufferString(`{"runner_id":"runner-1", "version":"2.0.0", "download_url":"https://example.com/binary"}`)
	reqActive := httptest.NewRequest(http.MethodPost, "/api/v1/runners/runner-1/self-update", bodyActive)
	rrActive := httptest.NewRecorder()
	handleRouter(store, registry)(rrActive, reqActive)
	if rrActive.Code != http.StatusOK {
		t.Fatalf("active update status = %v, want %v", rrActive.Code, http.StatusOK)
	}
	var resActive map[string]interface{}
	json.Unmarshal(rrActive.Body.Bytes(), &resActive)
	if resActive["accepted"] != false || resActive["deferred"] != true {
		t.Fatalf("active update response = %+v", resActive)
	}
}

func waitServerReady(t *testing.T, baseURL string) {
	t.Helper()
	for i := 0; i < 40; i++ {
		resp, err := http.Get(baseURL + "/healthz")
		if err == nil && resp.StatusCode == http.StatusOK {
			resp.Body.Close()
			return
		}
		if resp != nil {
			resp.Body.Close()
		}
		time.Sleep(5 * time.Millisecond)
	}
	t.Fatal("server did not become ready")
}

func TestServerStartFailureStopsTimeoutLoop(t *testing.T) {
	// Occupy a port so ListenAndServe fails immediately with "address already in use".
	blocker, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("blocker listen: %v", err)
	}
	defer blocker.Close()

	cfg := ServerConfig{
		HeartbeatTimeout: time.Hour,
		ScanInterval:     10 * time.Millisecond,
	}
	server := NewServerWithConfig(blocker.Addr().String(), runnerregistry.New(), cicdstate.NewStore(), cfg)

	startErr := make(chan error, 1)
	go func() { startErr <- server.Start() }()

	// Start must return quickly with a non-nil error (port in use).
	select {
	case err := <-startErr:
		if err == nil || err == http.ErrServerClosed {
			t.Fatalf("expected non-nil, non-ErrServerClosed error, got %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("Start did not return within 2 seconds with occupied port")
	}

	// Directly observe loopDone under server.mu — proves Start() cleaned up
	// the loop on its own return path, without relying on Shutdown().
	server.mu.Lock()
	done := server.loopDone
	server.mu.Unlock()
	if done == nil {
		t.Fatal("loopDone is nil after Start returned — startTimeoutLoop was never called")
	}
	select {
	case <-done:
		// loop goroutine exited — Start() cleaned up correctly
	default:
		t.Fatal("loopDone is still open after Start returned — Start() did not stop the timeout loop")
	}
}

func TestServerStartListenerFailureStopsTimeoutLoop(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	ln.Close() // cause immediate serve failure

	cfg := ServerConfig{
		HeartbeatTimeout: time.Hour,
		ScanInterval:     10 * time.Millisecond,
	}
	server := NewServerWithConfig(ln.Addr().String(), runnerregistry.New(), cicdstate.NewStore(), cfg)

	serveErr := make(chan error, 1)
	go func() { serveErr <- server.StartListener(ln) }()

	// StartListener must return quickly with an error (closed listener).
	select {
	case err := <-serveErr:
		if err == nil || err == http.ErrServerClosed {
			t.Fatalf("expected non-nil, non-ErrServerClosed error, got %v", err)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("StartListener did not return within 2 seconds after closed listener")
	}

	// Directly observe loopDone under server.mu — proves StartListener() cleaned up
	// the loop on its own return path, without relying on Shutdown().
	server.mu.Lock()
	done := server.loopDone
	server.mu.Unlock()
	if done == nil {
		t.Fatal("loopDone is nil after StartListener returned — startTimeoutLoop was never called")
	}
	select {
	case <-done:
		// loop goroutine exited — StartListener() cleaned up correctly
	default:
		t.Fatal("loopDone is still open after StartListener returned — StartListener() did not stop the timeout loop")
	}
}

func TestServerShutdownWithoutStartDoesNotHang(t *testing.T) {
	server := NewServerWithRegistry(":0", runnerregistry.New())
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	_ = server.Shutdown(ctx)
}

func TestServerTimeoutLoopTransitionsRunnerOffline(t *testing.T) {
	registry := runnerregistry.New()
	cfg := ServerConfig{
		HeartbeatTimeout: 50 * time.Millisecond,
		ScanInterval:     10 * time.Millisecond,
	}

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}

	server := NewServerWithConfig(ln.Addr().String(), registry, cicdstate.NewStore(), cfg)
	go func() { server.StartListener(ln) }()

	baseURL := "http://" + ln.Addr().String()
	waitServerReady(t, baseURL)

	regBody := `{"enrollment_token":"token-loop","runner_id":"runner-loop","protocol_version":"oto.runner.v1","capability":{"name":"oto-runner","version":"1.0.0"}}`
	regResp, err := http.Post(baseURL+"/api/v1/runners/register", "application/json", strings.NewReader(regBody))
	if err != nil {
		t.Fatalf("register: %v", err)
	}
	regResp.Body.Close()
	if regResp.StatusCode != http.StatusOK {
		t.Fatalf("register status = %v", regResp.StatusCode)
	}

	hbBody := `{"runner_id":"runner-loop","status":1}`
	hbResp, err := http.Post(baseURL+"/api/v1/runners/runner-loop/heartbeat", "application/json", strings.NewReader(hbBody))
	if err != nil {
		t.Fatalf("heartbeat: %v", err)
	}
	hbResp.Body.Close()
	if hbResp.StatusCode != http.StatusOK {
		t.Fatalf("heartbeat status = %v", hbResp.StatusCode)
	}

	record, _ := registry.Snapshot("runner-loop")
	if record.Status != runnerregistry.StatusOnline {
		t.Fatalf("expected online before timeout, got %q", record.Status)
	}

	// Poll until the timeout loop marks the runner offline (up to 300ms).
	deadline := time.Now().Add(300 * time.Millisecond)
	var ok bool
	for time.Now().Before(deadline) {
		record, ok = registry.Snapshot("runner-loop")
		if ok && record.Status == runnerregistry.StatusHeartbeatTimeout {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}
	if !ok {
		t.Fatal("runner not found after timeout")
	}
	if record.Status != runnerregistry.StatusHeartbeatTimeout {
		t.Fatalf("status = %q, want %q after polling deadline", record.Status, runnerregistry.StatusHeartbeatTimeout)
	}

	shutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := server.Shutdown(shutCtx); err != nil {
		t.Fatalf("shutdown: %v", err)
	}
}

func TestServerShutdownStopsTimeoutLoop(t *testing.T) {
	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}

	cfg := ServerConfig{
		HeartbeatTimeout: time.Hour,
		ScanInterval:     10 * time.Millisecond,
	}
	server := NewServerWithConfig(ln.Addr().String(), runnerregistry.New(), cicdstate.NewStore(), cfg)
	go func() { server.StartListener(ln) }()

	waitServerReady(t, "http://"+ln.Addr().String())

	done := make(chan error, 1)
	go func() { done <- server.Shutdown(context.Background()) }()

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("shutdown error: %v", err)
		}
	case <-time.After(5 * time.Second):
		t.Fatal("shutdown did not complete within 5 seconds")
	}
}