pulse/internal/agentexec/server_websocket_test.go

package agentexec

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

	"github.com/gorilla/websocket"
	"github.com/rcourtman/pulse-go-rewrite/internal/operationreceipt"
	"github.com/rcourtman/pulse-go-rewrite/internal/securityutil"
)

type wsRawMessage struct {
	Type      MessageType      `json:"type"`
	ID        string           `json:"id,omitempty"`
	Timestamp time.Time        `json:"timestamp"`
	Payload   *json.RawMessage `json:"payload,omitempty"`
}

func TestOperationQueryInconclusiveAPTDriftPreservesAdmittedDigest(t *testing.T) {
	now := time.Now().UTC()
	updateReq := HostUpdatePayload{RequestID: "u-drift.dispatch.1", ActionID: "u-drift", Operation: HostUpdateOperationInstall, ExpectedInventoryHash: "sha256:" + strings.Repeat("a", 64)}
	if err := BindHostUpdatePayload(&updateReq); err != nil {
		t.Fatal(err)
	}
	updateIdentity := HostUpdateOperationIdentity("agent", updateReq)
	updateResult := HostUpdateResultPayload{
		RequestID: updateReq.RequestID, ActionID: updateReq.ActionID, ExecutionPhase: HostUpdatePhaseRefresh, Verification: HostUpdateVerificationInconclusive,
		Before: HostPackageUpdateSnapshot{Supported: true, Manager: "apt", InventoryHash: "sha256:" + strings.Repeat("b", 64), PendingCount: 2, CheckedAt: now.Add(-time.Second)},
		After:  HostPackageUpdateSnapshot{Supported: true, Manager: "apt", InventoryHash: "sha256:" + strings.Repeat("b", 64), PendingCount: 2, CheckedAt: now},
	}
	updateRaw, err := json.Marshal(updateResult)
	if err != nil {
		t.Fatal(err)
	}

	cleanupReq := HostStorageCleanupPayload{RequestID: "c-drift.dispatch.1", ActionID: "c-drift", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: "sha256:" + strings.Repeat("c", 64)}
	if err := BindHostStorageCleanupPayload(&cleanupReq); err != nil {
		t.Fatal(err)
	}
	cleanupIdentity := HostStorageCleanupOperationIdentity("agent", cleanupReq)
	cleanupResult := HostStorageCleanupResultPayload{
		RequestID: cleanupReq.RequestID, ActionID: cleanupReq.ActionID, ExecutionPhase: HostStorageCleanupPhasePreflight, Verification: HostStorageCleanupVerificationInconclusive,
		Before: HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: "sha256:" + strings.Repeat("d", 64), ReclaimableBytes: 10, CheckedAt: now.Add(-time.Second)},
		After:  HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: "sha256:" + strings.Repeat("d", 64), ReclaimableBytes: 10, CheckedAt: now},
	}
	cleanupRaw, err := json.Marshal(cleanupResult)
	if err != nil {
		t.Fatal(err)
	}

	for _, tc := range []struct {
		name     string
		identity operationreceipt.Identity
		kind     string
		payload  json.RawMessage
	}{
		{name: "update inventory drift", identity: updateIdentity, kind: HostUpdateReceiptKind, payload: updateRaw},
		{name: "cleanup fingerprint drift", identity: cleanupIdentity, kind: HostStorageCleanupReceiptKind, payload: cleanupRaw},
	} {
		t.Run(tc.name, func(t *testing.T) {
			record := operationreceipt.Record{Identity: tc.identity, State: operationreceipt.StateTerminal, AcceptedAt: now.Add(-3 * time.Second), StartedAt: now.Add(-2 * time.Second), TerminalAt: now.Add(time.Second), ResultKind: tc.kind, ResultVersion: HostAPTReceiptVersion, Result: tc.payload}
			query := operationreceipt.QueryResult{Version: operationreceipt.ProtocolVersion, Status: operationreceipt.QueryFoundTerminal, Record: &record}
			if err := ValidateOperationQueryResultForIdentity(query, tc.identity, now.Add(2*time.Second)); err != nil {
				t.Fatalf("bound inconclusive drift receipt rejected: %v", err)
			}
			tampered := tc.identity
			tampered.RequestDigest = "sha256:" + strings.Repeat("e", 64)
			if err := ValidateOperationQueryResultForIdentity(query, tampered, now.Add(2*time.Second)); err == nil {
				t.Fatal("wrong request digest accepted")
			}
		})
	}
}

func newWSServer(t *testing.T, s *Server) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		s.HandleWebSocket(w, r)
	}))
}

func wsURLForHTTP(serverURL string) string {
	return "ws" + strings.TrimPrefix(serverURL, "http")
}

func wsHeadersForHTTP(t *testing.T, serverURL string) http.Header {
	t.Helper()

	origin, err := securityutil.HTTPOriginForWebSocketBaseURL(serverURL)
	if err != nil {
		t.Fatalf("failed to derive websocket origin: %v", err)
	}

	headers := http.Header{}
	headers.Set("Origin", origin)
	return headers
}

func dialAgentExecWebSocket(t *testing.T, serverURL string) (*websocket.Conn, *http.Response, error) {
	t.Helper()
	return websocket.DefaultDialer.Dial(wsURLForHTTP(serverURL), wsHeadersForHTTP(t, serverURL))
}

func wsWriteMessage(t *testing.T, conn *websocket.Conn, msg Message) {
	t.Helper()
	_ = conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
	if err := conn.WriteJSON(msg); err != nil {
		t.Fatalf("WriteJSON: %v", err)
	}
}

func mustNewMessage(t *testing.T, msgType MessageType, id string, payload any) Message {
	t.Helper()
	msg, err := NewMessage(msgType, id, payload)
	if err != nil {
		t.Fatalf("NewMessage: %v", err)
	}
	return msg
}

func wsReadRawMessage(t *testing.T, conn *websocket.Conn) wsRawMessage {
	t.Helper()
	msg, err := wsReadRawMessageWithTimeout(conn, 2*time.Second)
	if err != nil {
		t.Fatalf("ReadMessage: %v", err)
	}
	return msg
}

func wsReadRegisteredPayload(t *testing.T, conn *websocket.Conn) RegisteredPayload {
	t.Helper()
	msg := wsReadRawMessage(t, conn)
	if msg.Type != MsgTypeRegistered {
		t.Fatalf("message type = %q, want %q", msg.Type, MsgTypeRegistered)
	}
	if msg.Payload == nil {
		t.Fatalf("registered payload missing")
	}
	var payload RegisteredPayload
	if err := json.Unmarshal(*msg.Payload, &payload); err != nil {
		t.Fatalf("unmarshal registered payload: %v", err)
	}
	return payload
}

func wsReadRawMessageWithTimeout(conn *websocket.Conn, timeout time.Duration) (wsRawMessage, error) {
	_ = conn.SetReadDeadline(time.Now().Add(timeout))
	_, data, err := conn.ReadMessage()
	if err != nil {
		return wsRawMessage{}, err
	}
	var msg wsRawMessage
	if err := json.Unmarshal(data, &msg); err != nil {
		return wsRawMessage{}, err
	}
	return msg, nil
}

func waitFor(t *testing.T, timeout time.Duration, cond func() bool) {
	t.Helper()
	deadline := time.Now().Add(timeout)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(10 * time.Millisecond)
	}
	t.Fatalf("condition not met within %v", timeout)
}

func TestHandleWebSocket_RegistrationSuccessAndDisconnectRemovesAgent(t *testing.T) {
	s := NewServer(func(token string, agentID string, hostname string) bool { return token == "ok" })
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Tags:     []string{"tag1"},
		Token:    "ok",
	}))

	reg := wsReadRegisteredPayload(t, conn)
	if !reg.Success {
		t.Fatalf("registration failed: %q", reg.Message)
	}

	if !s.IsAgentConnected("a1") {
		t.Fatalf("expected agent to be connected")
	}

	conn.Close()

	waitFor(t, 2*time.Second, func() bool { return !s.IsAgentConnected("a1") })
}

func TestActionRunnerRegistrationRequiresCredentialBoundRuntimeRoleAndCapability(t *testing.T) {
	admission := AgentAdmission{
		OrganizationID:   "org-a",
		TokenID:          "runner-token",
		AgentID:          "machine-a",
		Hostname:         "node.example",
		RuntimeRole:      RuntimeRoleActionRunner,
		ActionCapability: ActionCapabilityTypedV1,
	}
	s := NewServerWithAdmissionValidator(func(token, _, _ string) (AgentAdmission, bool) {
		return admission, token == admission.TokenID
	}, func(candidate AgentAdmission) bool { return candidate == admission })
	ts := newWSServer(t, s)
	defer ts.Close()

	register := func(role, capability string) (*websocket.Conn, RegisteredPayload) {
		t.Helper()
		conn, _, err := dialAgentExecWebSocket(t, ts.URL)
		if err != nil {
			t.Fatalf("Dial: %v", err)
		}
		wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
			AgentID: admission.AgentID, Hostname: admission.Hostname, Token: admission.TokenID,
			RuntimeRole: role, ActionCapability: capability,
		}))
		return conn, wsReadRegisteredPayload(t, conn)
	}

	for _, invalid := range []struct{ role, capability string }{
		{"", ""},
		{RuntimeRoleActionRunner, ""},
		{RuntimeRoleActionRunner, "shell.v1"},
		{RuntimeRoleLegacyFullTrust, ActionCapabilityTypedV1},
	} {
		conn, ack := register(invalid.role, invalid.capability)
		conn.Close()
		if ack.Success {
			t.Fatalf("unbound runner assertion admitted: role=%q capability=%q", invalid.role, invalid.capability)
		}
	}

	conn, ack := register(RuntimeRoleActionRunner, ActionCapabilityTypedV1)
	defer conn.Close()
	if !ack.Success {
		t.Fatalf("bound action runner registration rejected: %s", ack.Message)
	}
	connected := s.GetConnectedAgentsForOrganization("org-a")
	if len(connected) != 1 || connected[0].RuntimeRole != RuntimeRoleActionRunner || connected[0].ActionCapability != ActionCapabilityTypedV1 {
		t.Fatalf("connected action runner = %#v", connected)
	}

	ctx := WithOrganizationID(context.Background(), "org-a")
	if _, err := s.ExecuteCommand(ctx, admission.AgentID, ExecuteCommandPayload{RequestID: "shell", Command: "true", TargetType: "agent", Trusted: true}); err == nil || !strings.Contains(err.Error(), "typed action-runner") {
		t.Fatalf("action runner accepted arbitrary command: %v", err)
	}
	if _, err := s.ReadFile(ctx, admission.AgentID, ReadFilePayload{RequestID: "read", Path: "/etc/hosts", TargetType: "agent"}); err == nil || !strings.Contains(err.Error(), "typed action-runner") {
		t.Fatalf("action runner accepted unrestricted read: %v", err)
	}
	if err := s.SendDeployCancel(ctx, admission.AgentID, DeployCancelPayload{RequestID: "deploy", JobID: "job"}); err == nil || !strings.Contains(err.Error(), "typed action-runner") {
		t.Fatalf("action runner accepted deploy protocol: %v", err)
	}
}

func TestLegacyCredentialCannotAssertActionRunnerRole(t *testing.T) {
	s := NewServerWithAdmissionValidator(func(token, _, _ string) (AgentAdmission, bool) {
		return AgentAdmission{TokenID: token, AgentID: "a1", Hostname: "host1", RuntimeRole: RuntimeRoleLegacyFullTrust}, token == "legacy"
	}, nil)
	ts := newWSServer(t, s)
	defer ts.Close()
	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "a1", Hostname: "host1", Token: "legacy", RuntimeRole: RuntimeRoleActionRunner, ActionCapability: ActionCapabilityTypedV1,
	}))
	if ack := wsReadRegisteredPayload(t, conn); ack.Success {
		t.Fatal("legacy credential self-promoted into action-runner role")
	}
}

func TestHandleWebSocket_RejectsMissingOrigin(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, resp, err := websocket.DefaultDialer.Dial(wsURLForHTTP(ts.URL), nil)
	if err == nil {
		conn.Close()
		t.Fatalf("expected websocket upgrade to reject missing Origin")
	}
	if resp == nil {
		t.Fatalf("expected HTTP response for rejected websocket upgrade")
	}
	if resp.StatusCode != http.StatusForbidden {
		t.Fatalf("expected %d, got %d", http.StatusForbidden, resp.StatusCode)
	}
}

func TestHandleWebSocket_RejectsPerIPConnectionFlood(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	s.maxConnsPerIP = 1
	ts := newWSServer(t, s)
	defer ts.Close()

	firstConn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial first connection: %v", err)
	}
	defer firstConn.Close()

	wsWriteMessage(t, firstConn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	reg := wsReadRegisteredPayload(t, firstConn)
	if !reg.Success {
		t.Fatalf("first registration failed: %q", reg.Message)
	}

	secondConn, resp, err := dialAgentExecWebSocket(t, ts.URL)
	if err == nil {
		secondConn.Close()
		t.Fatalf("expected second websocket upgrade to be rejected")
	}
	if resp == nil {
		t.Fatalf("expected HTTP response for rejected websocket upgrade")
	}
	if resp.StatusCode != http.StatusTooManyRequests {
		t.Fatalf("expected %d, got %d", http.StatusTooManyRequests, resp.StatusCode)
	}

	firstConn.Close()
	waitFor(t, 2*time.Second, func() bool { return !s.IsAgentConnected("a1") })

	thirdConn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial third connection after release: %v", err)
	}
	defer thirdConn.Close()

	wsWriteMessage(t, thirdConn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a2",
		Hostname: "host2",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	reg = wsReadRegisteredPayload(t, thirdConn)
	if !reg.Success {
		t.Fatalf("third registration failed after slot release: %q", reg.Message)
	}
}

func TestHandleWebSocket_RegistrationFiresAgentRegisteredNotifier(t *testing.T) {
	s := NewServer(func(token string, agentID string, hostname string) bool { return token == "ok" })
	notified := make(chan string, 2)
	s.SetAgentRegisteredNotifier(func(agentID string) { notified <- agentID })
	ts := newWSServer(t, s)
	defer ts.Close()

	rejectedConn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer rejectedConn.Close()
	wsWriteMessage(t, rejectedConn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "a-rejected", Hostname: "host-rejected", Token: "bad",
	}))
	if reg := wsReadRegisteredPayload(t, rejectedConn); reg.Success {
		t.Fatalf("expected registration to be rejected")
	}

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "a1", Hostname: "host1", Token: "ok",
	}))
	if reg := wsReadRegisteredPayload(t, conn); !reg.Success {
		t.Fatalf("registration failed: %q", reg.Message)
	}

	select {
	case agentID := <-notified:
		if agentID != "a1" {
			t.Fatalf("notified agent = %q, want a1 (rejected registration must not notify)", agentID)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("agent-registered notifier did not fire")
	}
	select {
	case agentID := <-notified:
		t.Fatalf("unexpected second notification for agent %q", agentID)
	case <-time.After(100 * time.Millisecond):
	}
}

func TestHandleWebSocket_InvalidTokenRejected(t *testing.T) {
	s := NewServer(func(string, string, string) bool { return false })
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "bad",
	}))

	reg := wsReadRegisteredPayload(t, conn)
	if reg.Success {
		t.Fatalf("expected registration to be rejected")
	}

	waitFor(t, 2*time.Second, func() bool { return !s.IsAgentConnected("a1") })

	_ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
	_, _, err = conn.ReadMessage()
	if err == nil {
		t.Fatalf("expected connection to be closed by server")
	}
}

func TestHandleWebSocket_MissingAgentIDRejected(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "   ",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))

	reg := wsReadRegisteredPayload(t, conn)
	if reg.Success {
		t.Fatalf("expected registration to be rejected")
	}

	_ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
	_, _, err = conn.ReadMessage()
	if err == nil {
		t.Fatalf("expected connection to be closed by server")
	}
}

func TestHandleWebSocket_FirstMessageMustBeRegister(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentPing, "", nil))

	_ = conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, _, err = conn.ReadMessage()
	if err == nil {
		t.Fatalf("expected server to close connection")
	}
}

func TestHandleWebSocket_RejectsOversizedRegistrationMessage(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	oversized := bytes.Repeat([]byte("x"), int(maxWebSocketMessageBytes)+1)
	if err := conn.WriteMessage(websocket.TextMessage, oversized); err != nil {
		t.Fatalf("WriteMessage: %v", err)
	}

	_ = conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, _, err = conn.ReadMessage()
	if err == nil {
		t.Fatalf("expected server to close connection for oversized registration message")
	}
}

func TestHandleWebSocket_AgentPingRespondsWithPong(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	_ = wsReadRegisteredPayload(t, conn)

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentPing, "", nil))

	msg := wsReadRawMessage(t, conn)
	if msg.Type != MsgTypePong {
		t.Fatalf("message type = %q, want %q", msg.Type, MsgTypePong)
	}
}

func TestExecuteCommand_RoundTripViaWebSocket(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	callerGrant := &CommandApprovalGrant{Signature: "caller-supplied"}
	s.SetCommandAuthorizationVerifier(func(req CommandAuthorizationRequest) error {
		if req.ApprovalID != "approval-1" || req.OrgID != "org-1" || req.ActionID != "action-1" {
			return fmt.Errorf("authorization mismatch: %+v", req)
		}
		return nil
	})
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	_ = wsReadRegisteredPayload(t, conn)

	agentDone := make(chan struct{})
	agentErr := make(chan error, 1)
	go func() {
		defer close(agentDone)
		for {
			msg, err := wsReadRawMessageWithTimeout(conn, 2*time.Second)
			if err != nil {
				agentErr <- err
				return
			}
			if msg.Type != MsgTypeExecuteCmd {
				continue
			}
			if msg.Payload == nil {
				agentErr <- nil
				return
			}
			var payload ExecuteCommandPayload
			if err := json.Unmarshal(*msg.Payload, &payload); err != nil {
				agentErr <- err
				return
			}
			if payload.ApprovalGrant == nil {
				agentErr <- fmt.Errorf("missing approval grant")
				return
			}
			if payload.ApprovalGrant.Signature == callerGrant.Signature {
				agentErr <- fmt.Errorf("caller-supplied approval grant was forwarded")
				return
			}
			if err := VerifyCommandApprovalGrant("any", "a1", payload, time.Now()); err != nil {
				agentErr <- err
				return
			}
			_ = conn.SetWriteDeadline(time.Now().Add(2 * time.Second))
			if err := conn.WriteJSON(mustNewMessage(t, MsgTypeCommandResult, "", CommandResultPayload{
				RequestID: payload.RequestID,
				Success:   true,
				Stdout:    "ok",
				ExitCode:  0,
				Duration:  1,
			})); err != nil {
				agentErr <- err
				return
			}
			agentErr <- nil
			return
		}
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	payload := ExecuteCommandPayload{
		RequestID:     "req1",
		Command:       "echo ok",
		ApprovalID:    "approval-1",
		ApprovalGrant: callerGrant,
		Timeout:       1,
	}
	payload.BindCommandAuthorization("org-1", "action-1")
	result, err := s.ExecuteCommand(ctx, "a1", payload)
	if err != nil {
		t.Fatalf("ExecuteCommand: %v", err)
	}
	if result == nil || !result.Success || result.Stdout != "ok" || result.ExitCode != 0 {
		t.Fatalf("unexpected result: %#v", result)
	}

	select {
	case <-agentDone:
	case <-time.After(2 * time.Second):
		t.Fatalf("agent goroutine did not finish")
	}

	if err := <-agentErr; err != nil {
		t.Fatalf("agent error: %v", err)
	}
}

func TestExecuteCommand_InvalidApprovalAuthorizationNeverMintsOrDispatches(t *testing.T) {
	cases := []struct {
		name string
		err  string
	}{
		{name: "nonexistent", err: "approval not found"},
		{name: "wrong-org", err: "approval belongs to another org"},
		{name: "expired", err: "approval expired"},
		{name: "consumed", err: "approval already consumed"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			s := NewServer(allowAllTestTokens)
			s.SetCommandAuthorizationVerifier(func(CommandAuthorizationRequest) error { return errors.New(tc.err) })
			grantCalls := 0
			s.newCommandApprovalGrant = func([]byte, string, ExecuteCommandPayload, time.Time, time.Duration) (*CommandApprovalGrant, error) {
				grantCalls++
				return nil, errors.New("grant must not be minted")
			}
			ts := newWSServer(t, s)
			defer ts.Close()

			conn, _, err := dialAgentExecWebSocket(t, ts.URL)
			if err != nil {
				t.Fatalf("Dial: %v", err)
			}
			defer conn.Close()
			wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
				AgentID: "a1", Hostname: "host1", Version: "1.2.3", Platform: "linux", Token: "any",
			}))
			_ = wsReadRegisteredPayload(t, conn)

			payload := ExecuteCommandPayload{
				RequestID: "req-invalid", Command: "echo rejected", ApprovalID: "approval-invalid", Timeout: 1,
			}
			payload.BindCommandAuthorization("org-1", "action-1")
			if _, err := s.ExecuteCommand(context.Background(), "a1", payload); err == nil || !strings.Contains(err.Error(), tc.err) {
				t.Fatalf("ExecuteCommand error = %v, want %q", err, tc.err)
			}
			if grantCalls != 0 {
				t.Fatalf("signed grant calls = %d, want 0", grantCalls)
			}
			if _, err := wsReadRawMessageWithTimeout(conn, 100*time.Millisecond); err == nil {
				t.Fatal("unexpected WebSocket dispatch for rejected approval")
			}
		})
	}
}

func TestExecuteHostUpdateRoundTripUsesTypedCommandFreeEnvelope(t *testing.T) {
	inventoryHash := "sha256:" + strings.Repeat("a", 64)
	emptyInventoryHash := "sha256:" + strings.Repeat("b", 64)
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()

	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "host-agent-1", Hostname: "host1", Version: "6.0.6", Platform: "linux", Token: "any", OperationReceiptVersion: 1,
	}))
	_ = wsReadRegisteredPayload(t, conn)

	agentErr := make(chan error, 1)
	go func() {
		msg, err := wsReadRawMessageWithTimeout(conn, 2*time.Second)
		if err != nil {
			agentErr <- err
			return
		}
		if msg.Type != MsgTypeHostUpdate || msg.Payload == nil {
			agentErr <- fmt.Errorf("message = %#v, want typed host update", msg)
			return
		}
		if bytes.Contains(*msg.Payload, []byte(`"command"`)) || bytes.Contains(*msg.Payload, []byte(`"packages"`)) {
			agentErr <- fmt.Errorf("host update request exposed command or package authority: %s", string(*msg.Payload))
			return
		}
		var payload HostUpdatePayload
		if err := json.Unmarshal(*msg.Payload, &payload); err != nil {
			agentErr <- err
			return
		}
		if payload.ActionID != "action-1" || payload.Operation != HostUpdateOperationInstall {
			agentErr <- fmt.Errorf("payload = %#v", payload)
			return
		}
		response := HostUpdateResultPayload{
			RequestID:      payload.RequestID,
			ActionID:       payload.ActionID,
			Success:        true,
			ExecutionPhase: HostUpdatePhaseComplete,
			Before:         HostPackageUpdateSnapshot{Supported: true, Manager: "apt", InventoryHash: inventoryHash, PendingCount: 2, CheckedAt: time.Now().UTC()},
			After:          HostPackageUpdateSnapshot{Supported: true, Manager: "apt", InventoryHash: emptyInventoryHash, PendingCount: 0, RebootRequired: true, CheckedAt: time.Now().UTC()},
			HealthChecked:  true, PackageManagerHealthy: true, Verification: HostUpdateVerificationVerified,
		}
		if err := conn.WriteJSON(mustNewMessage(t, MsgTypeHostUpdateResult, payload.RequestID, response)); err != nil {
			agentErr <- err
			return
		}
		agentErr <- nil
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	result, err := s.ExecuteHostUpdate(ctx, "host-agent-1", HostUpdatePayload{
		RequestID: "request-1", ActionID: "action-1", Operation: HostUpdateOperationInstall, ExpectedInventoryHash: inventoryHash, Timeout: 1,
	})
	if err != nil {
		t.Fatalf("ExecuteHostUpdate: %v", err)
	}
	if result == nil || !result.Success || result.Verification != HostUpdateVerificationVerified || result.After.PendingCount != 0 || !result.After.RebootRequired {
		t.Fatalf("result = %#v", result)
	}
	if err := <-agentErr; err != nil {
		t.Fatalf("agent: %v", err)
	}
}

func TestActionPreflightRoundTripIsReadOnlyAndDigestBound(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()
	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "preflight-agent", Hostname: "host1", Version: "6", Platform: "linux", Token: "any",
		OperationReceiptVersion: operationreceipt.ProtocolVersion, ActionPreflightVersion: ActionPreflightProtocolVersion,
	}))
	_ = wsReadRegisteredPayload(t, conn)
	req := boundHostUpdatePreflight(t)
	agentErr := make(chan error, 1)
	go func() {
		msg, readErr := wsReadRawMessageWithTimeout(conn, 2*time.Second)
		if readErr != nil {
			agentErr <- readErr
			return
		}
		if msg.Type != MsgTypeActionPreflight || msg.Payload == nil || bytes.Contains(*msg.Payload, []byte(`"command"`)) {
			agentErr <- fmt.Errorf("unexpected preflight envelope: %#v", msg)
			return
		}
		payload, decodeErr := DecodeActionPreflightPayload(*msg.Payload)
		if decodeErr != nil {
			agentErr <- decodeErr
			return
		}
		operation, version, digest := ActionPreflightBinding(payload)
		result := ActionPreflightResultPayload{
			RequestID: payload.RequestID, ProtocolVersion: payload.ProtocolVersion,
			Operation: operation, OperationVersion: version, RequestDigest: digest,
			ReasonCode: ActionRefusalPackageManagerUnhealthy, CheckedAt: time.Now().UTC(),
		}
		agentErr <- conn.WriteJSON(mustNewMessage(t, MsgTypeActionPreflightResult, payload.RequestID, result))
	}()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	result, err := s.PreflightAction(ctx, "preflight-agent", req)
	if err != nil {
		t.Fatal(err)
	}
	if result.Feasible || result.ReasonCode != ActionRefusalPackageManagerUnhealthy {
		t.Fatalf("result=%#v", result)
	}
	if err := <-agentErr; err != nil {
		t.Fatal(err)
	}
}

func TestDockerContainerObservationRoundTripIsReadOnlyAndActionBound(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()
	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "docker-observer", Hostname: "host1", Version: "6", Platform: "linux", Token: "any",
		DockerObservationVersion: DockerContainerObservationProtocolVersion,
	}))
	_ = wsReadRegisteredPayload(t, conn)
	containerID := strings.Repeat("a", 64)
	agentErr := make(chan error, 1)
	go func() {
		msg, readErr := wsReadRawMessageWithTimeout(conn, 2*time.Second)
		if readErr != nil {
			agentErr <- readErr
			return
		}
		if msg.Type != MsgTypeDockerContainerObserve || bytes.Contains(*msg.Payload, []byte(`"command"`)) || bytes.Contains(*msg.Payload, []byte(`"operation"`)) {
			agentErr <- fmt.Errorf("unexpected docker observation envelope: %#v", msg)
			return
		}
		payload, decodeErr := DecodeDockerContainerObservationPayload(*msg.Payload)
		if decodeErr != nil {
			agentErr <- decodeErr
			return
		}
		result := DockerContainerObservationResultPayload{
			RequestID: payload.RequestID, ActionID: payload.ActionID, ProtocolVersion: payload.ProtocolVersion, RequestDigest: payload.RequestDigest,
			Observed: true,
			Snapshot: DockerContainerObservationSnapshot{ContainerID: payload.ContainerID, State: "running", Running: true, Health: DockerContainerHealthHealthy, ObservedAt: time.Now().UTC()},
		}
		agentErr <- conn.WriteJSON(mustNewMessage(t, MsgTypeDockerContainerObserveResult, payload.RequestID, result))
	}()
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	result, err := s.ObserveDockerContainer(ctx, "docker-observer", DockerContainerObservationPayload{ActionID: "action-1", Runtime: "docker", ContainerID: containerID})
	if err != nil {
		t.Fatal(err)
	}
	if result.ActionID != "action-1" || result.Snapshot.ContainerID != containerID || !result.Snapshot.Running {
		t.Fatalf("result=%#v", result)
	}
	if err := <-agentErr; err != nil {
		t.Fatal(err)
	}
}

func TestDockerContainerObservationRejectsUnsupportedAgentBeforeDispatch(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()
	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "older-agent", Hostname: "host1", Version: "6.3", Platform: "linux", Token: "any",
		DockerObservationVersion: 1,
	}))
	_ = wsReadRegisteredPayload(t, conn)
	started := time.Now()
	_, err = s.ObserveDockerContainer(context.Background(), "older-agent", DockerContainerObservationPayload{
		ActionID: "action-1", Runtime: "docker", ContainerID: strings.Repeat("a", 64),
	})
	if err == nil || !strings.Contains(err.Error(), "does not support docker observation protocol") {
		t.Fatalf("error=%v", err)
	}
	if time.Since(started) > time.Second {
		t.Fatalf("unsupported agent rejection was not immediate: %s", time.Since(started))
	}
}

func TestValidateHostUpdatePayloadRejectsOpenEndedAuthority(t *testing.T) {
	for _, req := range []HostUpdatePayload{
		{RequestID: "r1", ActionID: "a1", Operation: "run_command", ExpectedInventoryHash: "sha256:" + strings.Repeat("a", 64)},
		{RequestID: "r1", Operation: HostUpdateOperationInstall, ExpectedInventoryHash: "sha256:" + strings.Repeat("a", 64)},
		{RequestID: "r1", ActionID: "a1", Operation: HostUpdateOperationInstall, Timeout: 1801, ExpectedInventoryHash: "sha256:" + strings.Repeat("a", 64)},
		{RequestID: "r1", ActionID: "a1", Operation: HostUpdateOperationInstall},
	} {
		copy := req
		if err := validateHostUpdatePayload(&copy); err == nil {
			t.Fatalf("validateHostUpdatePayload(%#v) succeeded", req)
		}
	}
}

func TestValidateHostUpdateResultRejectsUnprovenVerifiedClaim(t *testing.T) {
	result := HostUpdateResultPayload{
		RequestID: "r1", Success: true, Verification: HostUpdateVerificationVerified,
		After: HostPackageUpdateSnapshot{
			Supported: true, Manager: "apt", InventoryHash: "sha256:" + strings.Repeat("a", 64), PendingCount: 1,
		},
	}
	if err := validateHostUpdateResultPayload(&result); err == nil {
		t.Fatal("verified result with pending packages must fail closed")
	}
}

func TestExecuteHostStorageCleanupRoundTripUsesPathAndCommandFreeEnvelope(t *testing.T) {
	fingerprint := "sha256:" + strings.Repeat("a", 64)
	afterFingerprint := "sha256:" + strings.Repeat("b", 64)
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: "host-agent-cleanup", Hostname: "host1", Version: "6.0.6", Platform: "linux", Token: "any", OperationReceiptVersion: 1,
	}))
	_ = wsReadRegisteredPayload(t, conn)

	agentErr := make(chan error, 1)
	go func() {
		msg, err := wsReadRawMessageWithTimeout(conn, 2*time.Second)
		if err != nil {
			agentErr <- err
			return
		}
		if msg.Type != MsgTypeHostStorageCleanup || msg.Payload == nil {
			agentErr <- fmt.Errorf("message = %#v, want typed host storage cleanup", msg)
			return
		}
		for _, forbidden := range []string{`"command"`, `"path"`, `"packages"`} {
			if bytes.Contains(*msg.Payload, []byte(forbidden)) {
				agentErr <- fmt.Errorf("storage cleanup request exposed forbidden authority %s: %s", forbidden, string(*msg.Payload))
				return
			}
		}
		var payload HostStorageCleanupPayload
		if err := json.Unmarshal(*msg.Payload, &payload); err != nil {
			agentErr <- err
			return
		}
		if payload.ActionID != "action-cleanup" || payload.Operation != HostStorageCleanupOperationPackageCache {
			agentErr <- fmt.Errorf("payload = %#v", payload)
			return
		}
		response := HostStorageCleanupResultPayload{
			RequestID:      payload.RequestID,
			ActionID:       payload.ActionID,
			ExecutionPhase: HostStorageCleanupPhaseComplete,
			Success:        true,
			Before:         HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: fingerprint, ReclaimableBytes: 500, CheckedAt: time.Now().UTC()},
			After:          HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: afterFingerprint, ReclaimableBytes: 20, CheckedAt: time.Now().UTC()},
			ReclaimedBytes: 480,
			Verification:   HostStorageCleanupVerificationVerified,
		}
		if err := conn.WriteJSON(mustNewMessage(t, MsgTypeHostStorageCleanupResult, payload.RequestID, response)); err != nil {
			agentErr <- err
			return
		}
		agentErr <- nil
	}()

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	result, err := s.ExecuteHostStorageCleanup(ctx, "host-agent-cleanup", HostStorageCleanupPayload{
		RequestID: "cleanup-1", ActionID: "action-cleanup", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: fingerprint, Timeout: 1,
	})
	if err != nil {
		t.Fatalf("ExecuteHostStorageCleanup: %v", err)
	}
	if result == nil || !result.Success || result.Verification != HostStorageCleanupVerificationVerified || result.ReclaimedBytes != 480 {
		t.Fatalf("result = %#v", result)
	}
	if err := <-agentErr; err != nil {
		t.Fatalf("agent: %v", err)
	}
}

func TestValidateHostStorageCleanupRejectsOpenEndedOrUnprovenClaims(t *testing.T) {
	fingerprint := "sha256:" + strings.Repeat("a", 64)
	for _, req := range []HostStorageCleanupPayload{
		{RequestID: "r1", ActionID: "a1", Operation: "delete_path", ExpectedFingerprint: fingerprint},
		{RequestID: "r1", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: fingerprint},
		{RequestID: "r1", ActionID: "a1", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: "bad"},
		{RequestID: "r1", ActionID: "a1", Operation: HostStorageCleanupOperationPackageCache, ExpectedFingerprint: fingerprint, Timeout: 901},
	} {
		copy := req
		if err := validateHostStorageCleanupPayload(&copy); err == nil {
			t.Fatalf("validateHostStorageCleanupPayload(%#v) succeeded", req)
		}
	}
	result := HostStorageCleanupResultPayload{
		RequestID: "r1", Success: true, Verification: HostStorageCleanupVerificationVerified,
		Before: HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: fingerprint, ReclaimableBytes: 500},
		After:  HostStorageCleanupSnapshot{Supported: true, Provider: "apt-package-cache", Fingerprint: fingerprint, ReclaimableBytes: 500},
	}
	if err := validateHostStorageCleanupResultPayload(&result); err == nil {
		t.Fatal("verified result without reclaimed bytes must fail closed")
	}
}

func TestHandleWebSocket_ReconnectSameAgentIDClosesOldConnection(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	dial := func() *websocket.Conn {
		t.Helper()
		conn, _, err := dialAgentExecWebSocket(t, ts.URL)
		if err != nil {
			t.Fatalf("Dial: %v", err)
		}
		return conn
	}

	c1 := dial()
	defer c1.Close()
	wsWriteMessage(t, c1, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	_ = wsReadRegisteredPayload(t, c1)

	progressCh := s.SubscribeDeployProgress("a1", "job-reconnect", 1)
	defer s.UnsubscribeDeployProgress("a1", "job-reconnect")
	commandDone := make(chan error, 1)
	go func() {
		_, err := s.ExecuteCommand(context.Background(), "a1", ExecuteCommandPayload{
			RequestID: "command-before-reconnect",
			Command:   "true",
			Timeout:   10,
			Trusted:   true,
		})
		commandDone <- err
	}()
	if command := wsReadRawMessage(t, c1); command.Type != MsgTypeExecuteCmd {
		t.Fatalf("old session received %q, want %q", command.Type, MsgTypeExecuteCmd)
	}

	c2 := dial()
	defer c2.Close()
	wsWriteMessage(t, c2, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	_ = wsReadRegisteredPayload(t, c2)

	_ = c1.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
	_, _, err := c1.ReadMessage()
	if err == nil {
		t.Fatalf("expected old connection to be closed")
	}

	select {
	case commandErr := <-commandDone:
		if commandErr == nil || !strings.Contains(commandErr.Error(), "disconnected") {
			t.Fatalf("in-flight command reconnect result = %v, want disconnected", commandErr)
		}
	case <-time.After(time.Second):
		t.Fatal("in-flight command did not stop when its session was replaced")
	}

	progress := DeployProgressPayload{
		RequestID: "deploy-after-reconnect",
		JobID:     "job-reconnect",
		Phase:     DeployPhasePreflightSSH,
		Status:    DeployStepOK,
	}
	wsWriteMessage(t, c2, mustNewMessage(t, MsgTypeDeployProgress, progress.RequestID, progress))
	select {
	case received, ok := <-progressCh:
		if !ok {
			t.Fatal("replacement cleanup closed the active deploy subscription")
		}
		if received.RequestID != progress.RequestID {
			t.Fatalf("deploy progress request id = %q, want %q", received.RequestID, progress.RequestID)
		}
	case <-time.After(time.Second):
		t.Fatal("replacement session did not retain deploy progress subscription")
	}
}

func TestCommandSessionsAreTenantScopedAndDuplicateIdentityFailsClosed(t *testing.T) {
	admissions := map[string]AgentAdmission{
		"token-a": {OrganizationID: "org-a", TokenID: "token-a", AgentID: "shared", Hostname: "host-a"},
		"token-b": {OrganizationID: "org-b", TokenID: "token-b", AgentID: "shared", Hostname: "host-b"},
		"token-c": {OrganizationID: "org-a", TokenID: "token-c", AgentID: "shared", Hostname: "other-host"},
		"token-d": {OrganizationID: "org-a", TokenID: "token-d", AgentID: "other-id", Hostname: "host-a"},
	}
	s := NewServerWithAdmissionValidator(func(token, _, _ string) (AgentAdmission, bool) {
		admission, ok := admissions[token]
		return admission, ok
	}, func(AgentAdmission) bool { return true })
	ts := newWSServer(t, s)
	defer ts.Close()

	register := func(token, agentID, hostname string) (*websocket.Conn, RegisteredPayload) {
		t.Helper()
		conn, _, err := dialAgentExecWebSocket(t, ts.URL)
		if err != nil {
			t.Fatalf("Dial: %v", err)
		}
		wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
			AgentID: agentID, Hostname: hostname, Token: token,
		}))
		return conn, wsReadRegisteredPayload(t, conn)
	}

	orgA, ack := register("token-a", "shared", "host-a")
	defer orgA.Close()
	if !ack.Success {
		t.Fatalf("org-a registration failed: %s", ack.Message)
	}
	orgB, ack := register("token-b", "shared", "host-b")
	defer orgB.Close()
	if !ack.Success {
		t.Fatalf("org-b registration failed: %s", ack.Message)
	}
	if !s.IsAgentConnectedForOrganization("org-a", "shared") ||
		!s.IsAgentConnectedForOrganization("org-b", "shared") {
		t.Fatal("same agent id must remain independently connected in both organizations")
	}
	if s.IsAgentConnected("shared") {
		t.Fatal("tenant-scoped sessions must not leak into the default organization")
	}
	orgAView := s.ForOrganization("org-a")
	orgAAgents := orgAView.GetConnectedAgents()
	if len(orgAAgents) != 1 || orgAAgents[0].Hostname != "host-a" {
		t.Fatalf("org-a server view leaked another tenant: %#v", orgAAgents)
	}
	orgBAgents := s.ForOrganization("org-b").GetConnectedAgents()
	if len(orgBAgents) != 1 || orgBAgents[0].Hostname != "host-b" {
		t.Fatalf("org-b server view leaked another tenant: %#v", orgBAgents)
	}

	duplicate, ack := register("token-c", "shared", "other-host")
	defer duplicate.Close()
	if ack.Success {
		t.Fatal("same-tenant duplicate identity from another hostname was admitted")
	}
	if !s.IsAgentConnectedForOrganization("org-a", "shared") {
		t.Fatal("rejected duplicate identity evicted the original session")
	}

	duplicateHost, ack := register("token-d", "other-id", "host-a")
	defer duplicateHost.Close()
	if ack.Success {
		t.Fatal("same-tenant hostname was admitted under a second identity")
	}
	if !s.IsAgentConnectedForOrganization("org-a", "shared") {
		t.Fatal("rejected duplicate hostname evicted the original session")
	}
}

func TestRevokedAdmissionInvalidatesStaleSocketBeforeDispatch(t *testing.T) {
	valid := true
	admission := AgentAdmission{
		OrganizationID: "org-a",
		TokenID:        "token-a",
		AgentID:        "agent-a",
		Hostname:       "host-a",
	}
	s := NewServerWithAdmissionValidator(func(token, _, _ string) (AgentAdmission, bool) {
		return admission, token == admission.TokenID
	}, func(candidate AgentAdmission) bool {
		return valid && candidate == admission
	})
	ts := newWSServer(t, s)
	defer ts.Close()

	conn, _, err := dialAgentExecWebSocket(t, ts.URL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	defer conn.Close()
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID: admission.AgentID, Hostname: admission.Hostname, Token: admission.TokenID,
	}))
	if ack := wsReadRegisteredPayload(t, conn); !ack.Success {
		t.Fatalf("registration failed: %s", ack.Message)
	}
	if !s.IsAgentConnectedForOrganization("org-a", "agent-a") {
		t.Fatal("expected admitted session")
	}

	valid = false
	ctx := WithOrganizationID(context.Background(), "org-a")
	if _, err := s.ExecuteCommand(ctx, "agent-a", ExecuteCommandPayload{
		RequestID:  "after-revocation",
		Command:    "true",
		TargetType: "agent",
		Trusted:    true,
	}); err == nil || !strings.Contains(err.Error(), "not connected") {
		t.Fatalf("revoked stale socket remained dispatchable: %v", err)
	}
	if s.IsAgentConnectedForOrganization("org-a", "agent-a") {
		t.Fatal("revoked session remained visible as connected")
	}
}

// registerCancelTestAgent registers agent "a1" over the websocket harness and
// returns after the registration ack has been read.
func registerCancelTestAgent(t *testing.T, s *Server, tsURL string) *cancelTestConn {
	t.Helper()
	conn, _, err := dialAgentExecWebSocket(t, tsURL)
	if err != nil {
		t.Fatalf("Dial: %v", err)
	}
	t.Cleanup(func() { conn.Close() })
	wsWriteMessage(t, conn, mustNewMessage(t, MsgTypeAgentRegister, "", AgentRegisterPayload{
		AgentID:  "a1",
		Hostname: "host1",
		Version:  "1.2.3",
		Platform: "linux",
		Token:    "any",
	}))
	_ = wsReadRegisteredPayload(t, conn)
	return &cancelTestConn{t: t, conn: conn}
}

type cancelTestConn struct {
	t    *testing.T
	conn *websocket.Conn
}

// nextMessage returns the next message from the server, or ok=false when
// nothing arrives before the timeout.
func (c *cancelTestConn) nextMessage(timeout time.Duration) (wsRawMessage, bool) {
	c.t.Helper()
	msg, err := wsReadRawMessageWithTimeout(c.conn, timeout)
	if err != nil {
		return wsRawMessage{}, false
	}
	return msg, true
}

// The probe-storm incident (minipc, 2026-08-20) started with the server
// dispatching commands under a parent context that had already expired: the
// send succeeded, ExecuteCommand returned "context deadline exceeded,
// duration 0.05", and the agent was left running the command. An expired
// context must fail the call before anything reaches the agent.
func TestExecuteCommand_ExpiredContextNeverDispatches(t *testing.T) {
	s := NewServer(allowAllTestTokens)
	ts := newWSServer(t, s)
	defer ts.Close()

	agent := registerCancelTestAgent(t, s, ts.URL)

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	_, err := s.ExecuteCommand(ctx, "a1", ExecuteCommandPayload{
		RequestID: "req-expired",
		Command:   "echo hi",
		Timeout:   5,
		Trusted:   true,
	})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("ExecuteCommand error = %v, want context.Canceled", err)
	}
	if !strings.Contains(err.Error(), "not dispatched") {
		t.Fatalf("error should say the command was not dispatched, got %q", err)
	}

	// The agent must never see the command.
	if msg, ok := agent.nextMessage(500 * time.Millisecond); ok && msg.Type == MsgTypeExecuteCmd {
		t.Fatalf("agent received execute_command despite expired context")
	}
}

// When the server stops waiting for a dispatched command (its own timeout or
// the caller's context), it must tell the agent to abort the execution so the
// process tree is reaped instead of running on for minutes.
func TestExecuteCommand_AbandonedCommandSendsCancel(t *testing.T) {
	cases := []struct {
		name    string
		timeout int // ExecuteCommandPayload.Timeout in seconds
		abandon func(cancel context.CancelFunc)
		wantErr string
	}{
		{
			name:    "server timeout",
			timeout: 1,
			abandon: func(context.CancelFunc) {}, // let the 1s timer fire
			wantErr: "timed out",
		},
		{
			name:    "caller context canceled",
			timeout: 30,
			abandon: func(cancel context.CancelFunc) {
				time.Sleep(200 * time.Millisecond)
				cancel()
			},
			wantErr: "context canceled",
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			s := NewServer(allowAllTestTokens)
			ts := newWSServer(t, s)
			defer ts.Close()

			agent := registerCancelTestAgent(t, s, ts.URL)

			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			go tc.abandon(cancel)

			requestID := "req-" + strings.ReplaceAll(tc.name, " ", "-")
			_, err := s.ExecuteCommand(ctx, "a1", ExecuteCommandPayload{
				RequestID: requestID,
				Command:   "sleep 300",
				Timeout:   tc.timeout,
				Trusted:   true,
			})
			if err == nil || !strings.Contains(err.Error(), tc.wantErr) {
				t.Fatalf("ExecuteCommand error = %v, want containing %q", err, tc.wantErr)
			}

			// The agent first receives the command, then the cancellation.
			sawExecute := false
			sawCancel := false
			deadline := time.Now().Add(3 * time.Second)
			for time.Now().Before(deadline) && !sawCancel {
				msg, ok := agent.nextMessage(time.Until(deadline))
				if !ok {
					break
				}
				switch msg.Type {
				case MsgTypeExecuteCmd:
					sawExecute = true
				case MsgTypeCancelCmd:
					if msg.Payload == nil {
						t.Fatalf("cancel_command payload missing")
					}
					var payload CancelCommandPayload
					if err := json.Unmarshal(*msg.Payload, &payload); err != nil {
						t.Fatalf("unmarshal cancel_command payload: %v", err)
					}
					if payload.RequestID != requestID {
						t.Fatalf("cancel_command request_id = %q, want %q", payload.RequestID, requestID)
					}
					sawCancel = true
				}
			}
			if !sawExecute {
				t.Fatalf("agent never received execute_command")
			}
			if !sawCancel {
				t.Fatalf("agent never received cancel_command after the server abandoned the request")
			}
		})
	}
}