fix: handle OpAMP AgentDisconnect message (#6792)

* fix: handle OpAMP AgentDisconnect message

When an enrolled agent sends an AgentDisconnect message, set its
last_checkin_status to "disconnected" via the bulk checkin system.
Return a BadRequest error if an unenrolled agent sends a disconnect
message.

Closes #6784

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Use const for status

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: michel-laterman

changelog/fragments/1775675165-handle-opamp-agent-disconnect.yaml

@@ -0,0 +1,4 @@
+kind: bug-fix
+summary: Handle OpAMP AgentDisconnect message by setting agent status to disconnected
+component: fleet-server
+issue: https://github.com/elastic/fleet-server/issues/6784

docs/opamp.md

@@ -27,9 +27,8 @@ Fleet-server operates in monitoring-only mode. The `ServerToAgent` response only
 
 ### Agent-to-server fields ignored
 
-Fleet-server reads `instance_uid`, `agent_description`, `capabilities`, `health`, `effective_config`, and `sequence_num` from `AgentToServer` messages. The following fields are ignored:
+The following fields are ignored:
 
-- `agent_disconnect` — The spec says this MUST be set in the agent's last message.
 - `remote_config_status`
 - `package_statuses`
 - `flags` (e.g., `RequestInstanceUid`)

internal/pkg/api/handleOpAMP.go

@@ -35,7 +35,8 @@ import (
 )
 
 const (
-	kOpAMPMod = "opAMP"
+	kOpAMPMod          = "opAMP"
+	statusDisconnected = "disconnected"
 )
 
 type OpAMPT struct {
@@ -177,6 +178,26 @@ func (oa *OpAMPT) handleMessage(zlog zerolog.Logger, apiKey *apikey.APIKey) func
 			Bool("is_enrolled", agent != nil).
 			Msg("agent enrollment status")
 
+		// Handle agent disconnect: set status to offline for enrolled agents,
+		// return an error for unenrolled agents.
+		if message.AgentDisconnect != nil {
+			if agent == nil {
+				zlog.Debug().Msg("agent disconnect received from unenrolled agent")
+				return &protobufs.ServerToAgent{
+					InstanceUid: instanceUID.Bytes(),
+					ErrorResponse: &protobufs.ServerErrorResponse{
+						Type:         protobufs.ServerErrorResponseType_ServerErrorResponseType_BadRequest,
+						ErrorMessage: "agent is not enrolled",
+					},
+				}
+			}
+			zlog.Debug().Msg("agent disconnect received")
+			_ = oa.bc.CheckIn(instanceUID.String(), checkin.WithStatus(statusDisconnected))
+			return &protobufs.ServerToAgent{
+				InstanceUid: instanceUID.Bytes(),
+			}
+		}
+
 		if agent == nil {
 			if agent, err = oa.enrollAgent(zlog, instanceUID.String(), message, apiKey); err != nil {
 				return &protobufs.ServerToAgent{

internal/pkg/api/handleOpAMP_test.go

@@ -12,18 +12,21 @@ import (
 	"testing"
 	"unsafe"
 
+	"github.com/gofrs/uuid/v5"
+
+	"github.com/open-telemetry/opamp-go/protobufs"
+	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	semconv "go.opentelemetry.io/otel/semconv/v1.37.0"
+
 	"github.com/elastic/fleet-server/v7/internal/pkg/apikey"
 	"github.com/elastic/fleet-server/v7/internal/pkg/checkin"
 	"github.com/elastic/fleet-server/v7/internal/pkg/config"
 	"github.com/elastic/fleet-server/v7/internal/pkg/dl"
 	"github.com/elastic/fleet-server/v7/internal/pkg/es"
 	"github.com/elastic/fleet-server/v7/internal/pkg/model"
 	ftesting "github.com/elastic/fleet-server/v7/internal/pkg/testing"
-	"github.com/open-telemetry/opamp-go/protobufs"
-	"github.com/rs/zerolog"
-	"github.com/stretchr/testify/mock"
-	"github.com/stretchr/testify/require"
-	semconv "go.opentelemetry.io/otel/semconv/v1.37.0"
 )
 
 func TestFeatureFlag(t *testing.T) {
@@ -275,6 +278,72 @@ func TestUpdateAgentWithAgentToServerMessage(t *testing.T) {
 	require.Equal(t, float64(2), config["num"])
 }
 
+func TestHandleMessageAgentDisconnect(t *testing.T) {
+	cases := []struct {
+		name      string
+		getBulker func(t *testing.T) *ftesting.MockBulk
+		wantError bool
+	}{
+		{
+			name: "enrolled agent sets status to offline",
+			getBulker: func(t *testing.T) *ftesting.MockBulk {
+				t.Helper()
+				bulker := ftesting.NewMockBulk()
+				agent := model.Agent{LastCheckinStatus: "online"}
+				agentBytes, err := json.Marshal(agent)
+				require.NoError(t, err)
+				bulker.On("Search", mock.Anything, dl.FleetAgents, mock.Anything, mock.Anything).
+					Return(&es.ResultT{HitsT: es.HitsT{Hits: []es.HitT{{ID: "agent-123", Source: agentBytes}}}}, nil)
+				return bulker
+			},
+			wantError: false,
+		},
+		{
+			name: "unenrolled agent returns error",
+			getBulker: func(_ *testing.T) *ftesting.MockBulk {
+				bulker := ftesting.NewMockBulk()
+				bulker.On("Search", mock.Anything, dl.FleetAgents, mock.Anything, mock.Anything).
+					Return(&es.ResultT{HitsT: es.HitsT{Hits: []es.HitT{}}}, nil)
+				return bulker
+			},
+			wantError: true,
+		},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			bulker := tc.getBulker(t)
+			checker := &mockCheckin{}
+			oa := &OpAMPT{bulk: bulker, bc: checker}
+
+			agentUID := uuid.Must(uuid.NewV7())
+			zlog := zerolog.New(io.Discard)
+			apiKey := &apikey.APIKey{ID: "test-key"}
+
+			handler := oa.handleMessage(zlog, apiKey)
+			msg := &protobufs.AgentToServer{
+				InstanceUid:     agentUID.Bytes(),
+				AgentDisconnect: &protobufs.AgentDisconnect{},
+			}
+
+			resp := handler(t.Context(), nil, msg)
+			require.Equal(t, agentUID.Bytes(), resp.InstanceUid)
+
+			if tc.wantError {
+				require.NotNil(t, resp.ErrorResponse)
+				require.Equal(t, protobufs.ServerErrorResponseType_ServerErrorResponseType_BadRequest, resp.ErrorResponse.Type)
+				require.Empty(t, checker.id, "CheckIn should not be called for unenrolled agent")
+			} else {
+				require.Nil(t, resp.ErrorResponse)
+				require.Equal(t, uint64(0), resp.Capabilities)
+				require.Equal(t, agentUID.String(), checker.id)
+
+				pending := pendingFromOptions(t, checker.opts)
+				require.Equal(t, statusDisconnected, getUnexportedField(pending, "status").String())
+			}
+		})
+	}
+}
+
 type mockCheckin struct {
 	id   string
 	opts []checkin.Option