Repository: open-telemetry/opamp-go
Pull request #184 (2023-07-03 17:32:15)

Stop fully when accepting OpAMP connection settings in the agent example
Fixes https://github.com/open-telemetry/opamp-go/issues/178

I reverted https://github.com/open-telemetry/opamp-go/pull/170 which was unnecessary.
We don't really need to call Stop() from callbacks. In fact doing
stopping on a separate goroutine is preferable, properly waiting until
stopping is complete and then trying new connection settings.

I kept the reverting commit separate to make reviewing changes easier.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) Report
Patch coverage: **`89.65`**% and project coverage change: **`-0.77`** :warning:
> Comparison is base [(`efddaa2`)](https://app.codecov.io/gh/open-telemetry/opamp-go/commit/efddaa260895e7ebb2431deeefef3a3380d8d902?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) 76.11% compared to head [(`74af2c9`)](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) 75.34%.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #184      +/-   ##
==========================================
- Coverage   76.11%   75.34%   -0.77%     
==========================================
  Files          24       24              
  Lines        1834     1866      +32     
==========================================
+ Hits         1396     1406      +10     
- Misses        326      342      +16     
- Partials      112      118       +6     
```


| [Impacted Files](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) | Coverage Δ | |
|---|---|---|
| [client/types/callbacks.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L3R5cGVzL2NhbGxiYWNrcy5nbw==) | `66.66% <ø> (ø)` | |
| [server/httpconnection.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-c2VydmVyL2h0dHBjb25uZWN0aW9uLmdv) | `33.33% <0.00%> (ø)` | |
| [client/internal/mockserver.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL21vY2tzZXJ2ZXIuZ28=) | `82.35% <85.71%> (+0.31%)` | :arrow_up: |
| [server/serverimpl.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-c2VydmVyL3NlcnZlcmltcGwuZ28=) | `58.92% <94.73%> (+0.37%)` | :arrow_up: |
| [client/httpclient.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2h0dHBjbGllbnQuZ28=) | `95.00% <100.00%> (+0.26%)` | :arrow_up: |
| [client/internal/clientcommon.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL2NsaWVudGNvbW1vbi5nbw==) | `79.80% <100.00%> (ø)` | |
| [client/internal/httpsender.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL2h0dHBzZW5kZXIuZ28=) | `72.22% <100.00%> (+1.25%)` | :arrow_up: |
| [server/callbacks.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-c2VydmVyL2NhbGxiYWNrcy5nbw==) | `68.18% <100.00%> (-4.55%)` | :arrow_down: |
| [server/wsconnection.go](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-c2VydmVyL3dzY29ubmVjdGlvbi5nbw==) | `100.00% <100.00%> (ø)` | |

... and [2 files with indirect coverage changes](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/open-telemetry/opamp-go/pull/184?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry).


Patch:
diff --git a/client/client.go b/client/client.go
index c7ddc03..4c0ac37 100644
--- a/client/client.go
+++ b/client/client.go
@@ -37,11 +37,9 @@ type OpAMPClient interface {
 	// May be called only once.
 	// After this call returns successfully it is guaranteed that no
 	// callbacks will be called. Stop() will cancel context of any in-fly
-	// callbacks.
-	//
-	// If a callback is in progress (e.g. OnMessage is called but not finished)
-	// Stop() initiates stopping and returns without waiting for stopping to finish.
-	//
+	// callbacks, but will wait until such in-fly callbacks are returned before
+	// Stop returns, so make sure the callbacks don't block infinitely and react
+	// promptly to context cancellations.
 	// Once stopped OpAMPClient cannot be started again.
 	Stop(ctx context.Context) error
 
diff --git a/client/internal/clientcommon.go b/client/internal/clientcommon.go
index ff8a2b9..7c2d55d 100644
--- a/client/internal/clientcommon.go
+++ b/client/internal/clientcommon.go
@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"sync"
-	"sync/atomic"
 
 	"google.golang.org/protobuf/proto"
 
@@ -27,105 +26,11 @@ var (
 	errReportsPackageStatusesNotSet = errors.New("ReportsPackageStatuses capability is not set")
 )
 
-// CallbacksWrapper wraps Callbacks such that it is possible to query if any callback
-// function is in progress (called, but not yet returned). This is necessary for
-// safe handling of certain ClientCommon methods when they are called from the callbacks.
-// See for example Stop() implementation.
-type CallbacksWrapper struct {
-	wrapped types.Callbacks
-	// Greater than zero if currently processing a callback.
-	inCallback *int64
-}
-
-func (cc *CallbacksWrapper) OnConnect() {
-	cc.EnterCallback()
-	defer cc.LeaveCallback()
-	cc.wrapped.OnConnect()
-}
-
-func (cc *CallbacksWrapper) OnConnectFailed(err error) {
-	cc.EnterCallback()
-	defer cc.LeaveCallback()
-	cc.wrapped.OnConnectFailed(err)
-}
-
-func (cc *CallbacksWrapper) OnError(err *protobufs.ServerErrorResponse) {
-	cc.EnterCallback()
-	defer cc.LeaveCallback()
-	cc.wrapped.OnError(err)
-}
-
-func (cc *CallbacksWrapper) OnMessage(ctx context.Context, msg *types.MessageData) {
-	cc.EnterCallback()
-	defer cc.LeaveCallback()
-	cc.wrapped.OnMessage(ctx, msg)
-}
-
-func (cc *CallbacksWrapper) OnOpampConnectionSettings(
-	ctx context.Context, settings *protobufs.OpAMPConnectionSettings,
-) error {
-	cc.EnterCallback()
-	defer cc.LeaveCallback()
-	return cc.wrapped.OnOpampConnectionSettings(ctx, settings)
-}
-
-func (cc *CallbacksWrapper) OnOpampConnectionSettingsAccepted(settings *protobufs.OpAMPConnectionSettings) {
-	cc.EnterCallback()
-	defer cc.LeaveCallback()
-	cc.wrapped.OnOpampConnectionSettingsAccepted(settings)
-}
-
-func (cc *CallbacksWrapper) SaveRemoteConfigStatus(ctx context.Context, status *protobufs.RemoteConfigStatus) {
-	cc.EnterCallback()
-	defer cc.LeaveCallback()
-	cc.wrapped.SaveRemoteConfigStatus(ctx, status)
-}
-
-func (cc *CallbacksWrapper) GetEffectiveConfig(ctx context.Context) (*protobufs.EffectiveConfig, error) {
-	cc.EnterCallback()
-	defer cc.LeaveCallback()
-	return cc.wrapped.GetEffectiveConfig(ctx)
-}
-
-func (cc *CallbacksWrapper) OnCommand(command *protobufs.ServerToAgentCommand) error {
-	cc.EnterCallback()
-	defer cc.LeaveCallback()
-	return cc.wrapped.OnCommand(command)
-}
-
-var _ types.Callbacks = (*CallbacksWrapper)(nil)
-
-func NewCallbacksWrapper(wrapped types.Callbacks) *CallbacksWrapper {
-	zero := int64(0)
-
-	if wrapped == nil {
-		// Make sure it is always safe to call Callbacks.
-		wrapped = types.CallbacksStruct{}
-	}
-
-	return &CallbacksWrapper{
-		wrapped:    wrapped,
-		inCallback: &zero,
-	}
-}
-
-func (cc *CallbacksWrapper) EnterCallback() {
-	atomic.AddInt64(cc.inCallback, 1)
-}
-
-func (cc *CallbacksWrapper) LeaveCallback() {
-	atomic.AddInt64(cc.inCallback, -1)
-}
-
-func (cc *CallbacksWrapper) InCallback() bool {
-	return atomic.LoadInt64(cc.inCallback) != 0
-}
-
 // ClientCommon contains the OpAMP logic that is common between WebSocket and
 // plain HTTP transports.
 type ClientCommon struct {
 	Logger    types.Logger
-	Callbacks *CallbacksWrapper
+	Callbacks types.Callbacks
 
 	// Agent's capabilities defined at Start() time.
 	Capabilities protobufs.AgentCapabilities
@@ -225,7 +130,11 @@ func (c *ClientCommon) PrepareStart(
 	}
 
 	// Prepare callbacks.
-	c.Callbacks = NewCallbacksWrapper(settings.Callbacks)
+	c.Callbacks = settings.Callbacks
+	if c.Callbacks == nil {
+		// Make sure it is always safe to call Callbacks.
+		c.Callbacks = types.CallbacksStruct{}
+	}
 
 	if err := c.sender.SetInstanceUid(settings.InstanceUid); err != nil {
 		return err
@@ -247,16 +156,6 @@ func (c *ClientCommon) Stop(ctx context.Context) error {
 
 	cancelFunc()
 
-	if c.Callbacks.InCallback() {
-		// Stop() is called from a callback. We cannot wait and block here
-		// because the c.stoppedSignal may not be set until the callback
-		// returns. This is the case for example when OnMessage callback is
-		// called. So, for this case we return immediately and the caller
-		// needs to be aware that Stop() does not wait for stopping to
-		// finish in this case.
-		return nil
-	}
-
 	// Wait until stopping is finished.
 	select {
 	case <-ctx.Done():
diff --git a/client/internal/httpsender.go b/client/internal/httpsender.go
index 417e6f1..ccb6caf 100644
--- a/client/internal/httpsender.go
+++ b/client/internal/httpsender.go
@@ -13,9 +13,8 @@ import (
 	"time"
 
 	"github.com/cenkalti/backoff/v4"
-	"google.golang.org/protobuf/proto"
-
 	"github.com/open-telemetry/opamp-go/internal"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/open-telemetry/opamp-go/client/types"
 	"github.com/open-telemetry/opamp-go/protobufs"
@@ -37,7 +36,7 @@ type HTTPSender struct {
 	url                string
 	logger             types.Logger
 	client             *http.Client
-	callbacks          *CallbacksWrapper
+	callbacks          types.Callbacks
 	pollingIntervalMs  int64
 	compressionEnabled bool
 
@@ -71,7 +70,7 @@ func NewHTTPSender(logger types.Logger) *HTTPSender {
 func (h *HTTPSender) Run(
 	ctx context.Context,
 	url string,
-	callbacks *CallbacksWrapper,
+	callbacks types.Callbacks,
 	clientSyncedState *ClientSyncedState,
 	packagesStateProvider types.PackagesStateProvider,
 	capabilities protobufs.AgentCapabilities,
diff --git a/client/internal/receivedprocessor.go b/client/internal/receivedprocessor.go
index 1fd4c81..942505a 100644
--- a/client/internal/receivedprocessor.go
+++ b/client/internal/receivedprocessor.go
@@ -13,7 +13,7 @@ type receivedProcessor struct {
 	logger types.Logger
 
 	// Callbacks to call for corresponding messages.
-	callbacks *CallbacksWrapper
+	callbacks types.Callbacks
 
 	// A sender to cooperate with when the received message has an impact on
 	// what will be sent later.
@@ -30,7 +30,7 @@ type receivedProcessor struct {
 
 func newReceivedProcessor(
 	logger types.Logger,
-	callbacks *CallbacksWrapper,
+	callbacks types.Callbacks,
 	sender Sender,
 	clientSyncedState *ClientSyncedState,
 	packagesStateProvider types.PackagesStateProvider,
diff --git a/client/internal/wsreceiver.go b/client/internal/wsreceiver.go
index 1c62cf8..f4a3be5 100644
--- a/client/internal/wsreceiver.go
+++ b/client/internal/wsreceiver.go
@@ -5,7 +5,6 @@ import (
 	"fmt"
 
 	"github.com/gorilla/websocket"
-
 	"github.com/open-telemetry/opamp-go/client/types"
 	"github.com/open-telemetry/opamp-go/internal"
 	"github.com/open-telemetry/opamp-go/protobufs"
@@ -16,7 +15,7 @@ type wsReceiver struct {
 	conn      *websocket.Conn
 	logger    types.Logger
 	sender    *WSSender
-	callbacks *CallbacksWrapper
+	callbacks types.Callbacks
 	processor receivedProcessor
 }
 
@@ -24,7 +23,7 @@ type wsReceiver struct {
 // messages from the server.
 func NewWSReceiver(
 	logger types.Logger,
-	callbacks *CallbacksWrapper,
+	callbacks types.Callbacks,
 	conn *websocket.Conn,
 	sender *WSSender,
 	clientSyncedState *ClientSyncedState,
diff --git a/client/wsclient.go b/client/wsclient.go
index 6a189ae..d07f7bd 100644
--- a/client/wsclient.go
+++ b/client/wsclient.go
@@ -123,7 +123,7 @@ func (c *wsClient) tryConnectOnce(ctx context.Context) (err error, retryAfter sh
 	var resp *http.Response
 	conn, resp, err := c.dialer.DialContext(ctx, c.url.String(), c.requestHeader)
 	if err != nil {
-		if !c.common.IsStopping() {
+		if c.common.Callbacks != nil && !c.common.IsStopping() {
 			c.common.Callbacks.OnConnectFailed(err)
 		}
 		if resp != nil {
@@ -138,7 +138,9 @@ func (c *wsClient) tryConnectOnce(ctx context.Context) (err error, retryAfter sh
 	c.connMutex.Lock()
 	c.conn = conn
 	c.connMutex.Unlock()
-	c.common.Callbacks.OnConnect()
+	if c.common.Callbacks != nil {
+		c.common.Callbacks.OnConnect()
+	}
 
 	return nil, sharedinternal.OptionalDuration{Defined: false}
 }
@@ -191,10 +193,9 @@ func (c *wsClient) ensureConnected(ctx context.Context) error {
 }
 
 // runOneCycle performs the following actions:
-//  1. connect (try until succeeds).
-//  2. send first status report.
-//  3. receive and process messages until error happens.
-//
+//   1. connect (try until succeeds).
+//   2. send first status report.
+//   3. receive and process messages until error happens.
 // If it encounters an error it closes the connection and returns.
 // Will stop and return if Stop() is called (ctx is cancelled, isStopping is set).
 func (c *wsClient) runOneCycle(ctx context.Context) {
diff --git a/internal/examples/agent/agent/agent.go b/internal/examples/agent/agent/agent.go
index 8865136..c82d96c 100644
--- a/internal/examples/agent/agent/agent.go
+++ b/internal/examples/agent/agent/agent.go
@@ -395,24 +395,11 @@ func (agent *Agent) onMessage(ctx context.Context, msg *types.MessageData) {
 	}
 }
 
-func (agent *Agent) onOpampConnectionSettings(ctx context.Context, settings *protobufs.OpAMPConnectionSettings) error {
-	if settings == nil || settings.Certificate == nil {
-		agent.logger.Debugf("Received nil certificate offer, ignoring.\n")
-		return nil
-	}
-
-	cert, err := agent.getCertFromSettings(settings.Certificate)
-	if err != nil {
-		return err
-	}
-
+func (agent *Agent) tryChangeOpAMPCert(cert *tls.Certificate) {
 	agent.logger.Debugf("Reconnecting to verify offered client certificate.\n")
 
 	agent.disconnect()
 
-	// TODO: wait for disconnect() to complete the stopping of old connection before
-	// we begin connecting with the new settings.
-
 	agent.opampClientCert = cert
 	if err := agent.connect(); err != nil {
 		agent.logger.Errorf("Cannot connect using offered certificate: %s. Ignoring the offer\n", err)
@@ -420,7 +407,6 @@ func (agent *Agent) onOpampConnectionSettings(ctx context.Context, settings *pro
 
 		if err := agent.connect(); err != nil {
 			agent.logger.Errorf("Unable to reconnect after restoring client certificate: %v\n", err)
-			return err
 		}
 	}
 
@@ -428,8 +414,21 @@ func (agent *Agent) onOpampConnectionSettings(ctx context.Context, settings *pro
 
 	// TODO: we can also persist the successfully accepted certificate and use it when the
 	// agent connects to the server after the restart.
+}
+
+func (agent *Agent) onOpampConnectionSettings(ctx context.Context, settings *protobufs.OpAMPConnectionSettings) error {
+	if settings == nil || settings.Certificate == nil {
+		agent.logger.Debugf("Received nil certificate offer, ignoring.\n")
+		return nil
+	}
+
+	cert, err := agent.getCertFromSettings(settings.Certificate)
+	if err != nil {
+		return err
+	}
 
 	// TODO: also use settings.DestinationEndpoint and settings.Headers for future connections.
+	go agent.tryChangeOpAMPCert(cert)
 
 	return nil
 }


Test patch:
diff --git a/client/clientimpl_test.go b/client/clientimpl_test.go
index 84ff779..16f66dd 100644
--- a/client/clientimpl_test.go
+++ b/client/clientimpl_test.go
@@ -483,85 +483,6 @@ func TestIncludesDetailsOnReconnect(t *testing.T) {
 	assert.NoError(t, err)
 }
 
-func TestStopFromCallback(t *testing.T) {
-	// This test verifies calling Stop() from a callback. We had a bug previously
-	// where Stop() would hang if called from a callback.
-
-	callbacksToTest := []string{"connect", "opamp", "message"}
-	for _, callbackToTest := range callbacksToTest {
-		t.Run(
-			callbackToTest, func(t *testing.T) {
-
-				testClients(
-					t, func(t *testing.T, client OpAMPClient) {
-						var called int64
-
-						hash := []byte{1, 2, 3}
-						opampSettings := &protobufs.OpAMPConnectionSettings{DestinationEndpoint: "http://opamp.com"}
-
-						// Start a Server.
-						srv := internal.StartMockServer(t)
-						srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
-							if msg != nil {
-								return &protobufs.ServerToAgent{
-									ConnectionSettings: &protobufs.ConnectionSettingsOffers{
-										Hash:  hash,
-										Opamp: opampSettings,
-									},
-								}
-							}
-							return nil
-						}
-
-						// Start a client.
-						settings := types.StartSettings{
-							Callbacks: types.CallbacksStruct{
-								OnConnectFunc: func() {
-									if callbackToTest == "connect" {
-										client.Stop(context.Background())
-										atomic.StoreInt64(&called, 1)
-									}
-								},
-								OnOpampConnectionSettingsFunc: func(
-									ctx context.Context, settings *protobufs.OpAMPConnectionSettings,
-								) error {
-									if callbackToTest == "opamp" {
-										client.Stop(context.Background())
-										atomic.StoreInt64(&called, 1)
-									}
-									return nil
-								},
-								OnMessageFunc: func(ctx context.Context, msg *types.MessageData) {
-									if callbackToTest == "message" {
-										client.Stop(context.Background())
-										atomic.StoreInt64(&called, 1)
-									}
-								},
-							},
-							Capabilities: protobufs.AgentCapabilities_AgentCapabilities_AcceptsOpAMPConnectionSettings,
-						}
-						settings.OpAMPServerURL = "ws://" + srv.Endpoint
-						prepareClient(t, &settings, client)
-
-						assert.NoError(t, client.Start(context.Background(), settings))
-
-						eventually(
-							t, func() bool {
-								return atomic.LoadInt64(&called) == 1
-							},
-						)
-
-						// Shutdown the Server.
-						srv.Close()
-
-						// Shutdown the client.
-						err := client.Stop(context.Background())
-						assert.NoError(t, err)
-					})
-			})
-	}
-}
-
 func createEffectiveConfig() *protobufs.EffectiveConfig {
 	cfg := &protobufs.EffectiveConfig{
 		ConfigMap: &protobufs.AgentConfigMap{
diff --git a/client/internal/httpsender_test.go b/client/internal/httpsender_test.go
index 3a984e1..58a5268 100644
--- a/client/internal/httpsender_test.go
+++ b/client/internal/httpsender_test.go
@@ -8,11 +8,10 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
-
 	"github.com/open-telemetry/opamp-go/client/types"
 	sharedinternal "github.com/open-telemetry/opamp-go/internal"
 	"github.com/open-telemetry/opamp-go/protobufs"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestHTTPSenderRetryForStatusTooManyRequests(t *testing.T) {
@@ -42,12 +41,12 @@ func TestHTTPSenderRetryForStatusTooManyRequests(t *testing.T) {
 			}},
 		}
 	})
-	sender.callbacks = NewCallbacksWrapper(types.CallbacksStruct{
+	sender.callbacks = types.CallbacksStruct{
 		OnConnectFunc: func() {
 		},
 		OnConnectFailedFunc: func(_ error) {
 		},
-	})
+	}
 	sender.url = url
 	start := time.Now()
 	resp, err := sender.sendRequestWithRetries(ctx)
diff --git a/client/internal/wsreceiver_test.go b/client/internal/wsreceiver_test.go
index 462ba17..44cdfdb 100644
--- a/client/internal/wsreceiver_test.go
+++ b/client/internal/wsreceiver_test.go
@@ -77,7 +77,7 @@ func TestServerToAgentCommand(t *testing.T) {
 				remoteConfigStatus: &protobufs.RemoteConfigStatus{},
 			}
 			sender := WSSender{}
-			receiver := NewWSReceiver(TestLogger{t}, NewCallbacksWrapper(callbacks), nil, &sender, &clientSyncedState, nil, 0)
+			receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, &sender, &clientSyncedState, nil, 0)
 			receiver.processor.ProcessReceivedMessage(context.Background(), &protobufs.ServerToAgent{
 				Command: test.command,
 			})
@@ -100,7 +100,7 @@ func TestServerToAgentCommandExclusive(t *testing.T) {
 		},
 	}
 	clientSyncedState := ClientSyncedState{}
-	receiver := NewWSReceiver(TestLogger{t}, NewCallbacksWrapper(callbacks), nil, nil, &clientSyncedState, nil, 0)
+	receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, nil, &clientSyncedState, nil, 0)
 	receiver.processor.ProcessReceivedMessage(context.Background(), &protobufs.ServerToAgent{
 		Command: &protobufs.ServerToAgentCommand{
 			Type: protobufs.CommandType_CommandType_Restart,