Repository: open-telemetry/opamp-go
Pull request #148 (2022-11-11 15:39:18)

Permit per-connection Callbacks to allow per-connection state
This adds a ConnectionHandler field to the ConnectionResponse struct. If nil, the Callbacks provided with StartSettings would be used.

I don't have time at the moment to provide a working demonstration of maintaining per-connection state, but hopefully this is obvious.

The implementation would look something like this:


```go
// connectionState contains any state the implementer wishes to preserve per-connection.
// It also implements types.ConnectionCallbacks
type connectionState struct {
	connectedAt time.Time
}

func (c *connectionState) OnConnected(conn types.Connection) {
	c.connectedAt = time.Now()
}

func (c *connectionState) OnMessage(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
	// TODO: real message handling
	return &protobufs.ServerToAgent{}
}

func (c *connectionState) OnConnectionClose(conn types.Connection) {
}
```

The callbacks passed to StartSettings.Settings would construct a new connectionState in OnConnecting.

```go
callbacks := CallbacksStruct{
	OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
		return types.ConnectionResponse{
			Accept:             true,
			ConnectionHandler:  &connectionState{},
		}
	},
}
```

codecov[bot]: # [Codecov](https://codecov.io/gh/open-telemetry/opamp-go/pull/148?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) Report
Base: **75.05**% // Head: **76.01**% // Increases project coverage by **`+0.95%`** :tada:
> Coverage data is based on head [(`c85cee8`)](https://codecov.io/gh/open-telemetry/opamp-go/pull/148?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) compared to base [(`7462bc8`)](https://codecov.io/gh/open-telemetry/opamp-go/commit/7462bc8dafbd7ccca026c1ab7ea25a253e543cf7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry).
> Patch coverage: 69.23% of modified lines in pull request are covered.

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


```diff
@@            Coverage Diff             @@
##             main     #148      +/-   ##
==========================================
+ Coverage   75.05%   76.01%   +0.95%     
==========================================
  Files          23       23              
  Lines        1792     1801       +9     
==========================================
+ Hits         1345     1369      +24     
+ Misses        334      322      -12     
+ Partials      113      110       -3     
```


| [Impacted Files](https://codecov.io/gh/open-telemetry/opamp-go/pull/148?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry) | Coverage Δ | |
|---|---|---|
| [server/callbacks.go](https://codecov.io/gh/open-telemetry/opamp-go/pull/148/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-c2VydmVyL2NhbGxiYWNrcy5nbw==) | `72.72% <ø> (ø)` | |
| [server/serverimpl.go](https://codecov.io/gh/open-telemetry/opamp-go/pull/148/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-c2VydmVyL3NlcnZlcmltcGwuZ28=) | `58.26% <69.23%> (-0.12%)` | :arrow_down: |
| [client/wsclient.go](https://codecov.io/gh/open-telemetry/opamp-go/pull/148/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L3dzY2xpZW50Lmdv) | `84.75% <0.00%> (-2.44%)` | :arrow_down: |
| [client/internal/tcpproxy.go](https://codecov.io/gh/open-telemetry/opamp-go/pull/148/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL3RjcHByb3h5Lmdv) | | |
| [internal/testhelpers/tcpproxy.go](https://codecov.io/gh/open-telemetry/opamp-go/pull/148/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-aW50ZXJuYWwvdGVzdGhlbHBlcnMvdGNwcHJveHkuZ28=) | `76.19% <0.00%> (ø)` | |
| [client/internal/receivedprocessor.go](https://codecov.io/gh/open-telemetry/opamp-go/pull/148/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry#diff-Y2xpZW50L2ludGVybmFsL3JlY2VpdmVkcHJvY2Vzc29yLmdv) | `84.28% <0.00%> (+15.00%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=open-telemetry)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/open-telemetry/opamp-go/pull/148?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).

tigrannajaryan: I like this due to stronger typing.

Earlier I said for WebSocket we may want to avoid keeping this but I think in real world scenarios you inevitably need to keep carry this state for any transport. You typically authenticate in OnConnecting and need to save the tenant id that you get as a result of authenticating somewhere since it is necessary later in OnMessage.

Should we get rid of the callback in Settings and replace it with this per connection Callbacks?
andykellr: > I like this due to stronger typing.

I agree.

> Earlier I said for WebSocket we may want to avoid keeping this but I think in real world scenarios you inevitably need to keep carry this state for any transport. You typically authenticate in OnConnecting and need to save the tenant id that you get as a result of authenticating somewhere since it is necessary later in OnMessage.

Good example. I think it is critical that we at least figure out a way to connect OnConnecting and OnMessage and this approach feels right.

> Should we get rid of the callback in Settings and replace it with this per connection Callbacks?

I considered that, but wanted to make sure this PR was small and easy to understand. Changing the callbacks requires changes to the `CallbacksStruct`, etc.

One benefit of this implementation is compatibility with existing configurations, but there probably aren't many out there and if the same implementation is used for both callbacks, it's easy enough to return it in the ConnectionResponse.
andykellr: @tigrannajaryan I finally had a chance to refactor this PR to separate the Callbacks and ConnectionCallbacks and update the tests.

Patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 53bace9..9b035f7 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -2,7 +2,18 @@
 
 ## Generate Protobuf Go Files
 
-Make sure you have Docker installed.
+1. Make sure `opamp-go/internal/opamp-spec` submodule files exist. You can do this
+either by cloning the `opamp-go` repo with submodules the first time, e.g. 
+`git clone --recurse-submodules git@github.com:open-telemetry/opamp-go.git`.
+Alternatively, if you already cloned `opamp-go` without the submodule then
+execute `git submodule update --init` to fetch all files for `opamp-spec` submodule.
+Note that `opamp-spec` submodule requires ssh git cloning access to github and won't
+work if you only have https access.
 
-Run `make gen-proto`. This should compile `internal/proto/*.proto` files to 
+2. Make sure you have Docker installed and running.
+
+3. Run `make gen-proto`. This should compile `internal/proto/*.proto` files to 
 `internal/protobufs/*.pb.go` files.
+
+Note that this is tested on linux/amd64 and darwin/amd64 and is known to not work
+on darwin/arm64 (M1/M2 Macs). Fixes are welcome.
\ No newline at end of file
diff --git a/README.md b/README.md
index 57d6dd0..8bcc163 100644
--- a/README.md
+++ b/README.md
@@ -12,7 +12,7 @@
   <a href="https://pkg.go.dev/github.com/open-telemetry/opamp-go">
     <img alt="Go Docs" height="28" src="https://godoc.org/github.com/open-telemetry/opamp-go?status.svg"></a>
   <a href="https://github.com/open-telemetry/opamp-go/actions/workflows/build-and-test.yml?query=branch%3Amain">
-    <img alt="Build Status" src="https://img.shields.io/github/workflow/status/open-telemetry/opamp-go/build-and-test/main?style=for-the-badge"></a>
+    <img alt="Build Status" src="https://img.shields.io/github/actions/workflow/status/open-telemetry/opamp-go/build-and-test.yml?branch%3Amain&style=for-the-badge"></a>
   <a href="https://goreportcard.com/report/github.com/open-telemetry/opamp-go">
     <img alt="Go Report Card" src="https://goreportcard.com/badge/github.com/open-telemetry/opamp-go?style=for-the-badge"></a>
   <a href="https://codecov.io/gh/open-telemetry/opamp-go/branch/main/">
diff --git a/internal/examples/server/opampsrv/opampsrv.go b/internal/examples/server/opampsrv/opampsrv.go
index e30990e..082a4f4 100644
--- a/internal/examples/server/opampsrv/opampsrv.go
+++ b/internal/examples/server/opampsrv/opampsrv.go
@@ -3,6 +3,7 @@ package opampsrv
 import (
 	"context"
 	"log"
+	"net/http"
 
 	"github.com/open-telemetry/opamp-go/internal/examples/server/data"
 	"github.com/open-telemetry/opamp-go/protobufs"
@@ -35,8 +36,12 @@ func (srv *Server) Start() {
 	settings := server.StartSettings{
 		Settings: server.Settings{
 			Callbacks: server.CallbacksStruct{
-				OnMessageFunc:         srv.onMessage,
-				OnConnectionCloseFunc: srv.onDisconnect,
+				OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
+					return types.ConnectionResponse{Accept: true, ConnectionCallbacks: server.ConnectionCallbacksStruct{
+						OnMessageFunc:         srv.onMessage,
+						OnConnectionCloseFunc: srv.onDisconnect,
+					}}
+				},
 			},
 		},
 		ListenEndpoint: "127.0.0.1:4320",
diff --git a/makefile b/makefile
index 409882f..ee8e21b 100644
--- a/makefile
+++ b/makefile
@@ -48,6 +48,13 @@ OTEL_DOCKER_PROTOBUF ?= otel/build-protobuf:0.14.0
 .PHONY: gen-proto
 gen-proto:
 	mkdir -p ${PWD}/internal/proto/
+
+	@if [ ! -d "${PWD}/internal/opamp-spec/proto" ]; then \
+		echo "${PWD}/internal/opamp-spec/proto does not exist."; \
+		echo "Run \`git submodule update --init\` to fetch the submodule"; \
+		exit 1; \
+	fi
+
 	$(foreach file,$(BASELINE_PROTO_FILES),$(call exec-command,docker run --rm -v${PWD}:${PWD} \
         -w${PWD} $(OTEL_DOCKER_PROTOBUF) --proto_path=${PWD}/internal/opamp-spec/proto/ \
         --go_out=${PWD}/internal/proto/ -I${PWD}/internal/proto/ ${PWD}/$(file)))
diff --git a/server/callbacks.go b/server/callbacks.go
index 46619e5..c10a576 100644
--- a/server/callbacks.go
+++ b/server/callbacks.go
@@ -8,10 +8,7 @@ import (
 )
 
 type CallbacksStruct struct {
-	OnConnectingFunc      func(request *http.Request) types.ConnectionResponse
-	OnConnectedFunc       func(conn types.Connection)
-	OnMessageFunc         func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent
-	OnConnectionCloseFunc func(conn types.Connection)
+	OnConnectingFunc func(request *http.Request) types.ConnectionResponse
 }
 
 var _ types.Callbacks = (*CallbacksStruct)(nil)
@@ -23,13 +20,21 @@ func (c CallbacksStruct) OnConnecting(request *http.Request) types.ConnectionRes
 	return types.ConnectionResponse{Accept: true}
 }
 
-func (c CallbacksStruct) OnConnected(conn types.Connection) {
+type ConnectionCallbacksStruct struct {
+	OnConnectedFunc       func(conn types.Connection)
+	OnMessageFunc         func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent
+	OnConnectionCloseFunc func(conn types.Connection)
+}
+
+var _ types.ConnectionCallbacks = (*ConnectionCallbacksStruct)(nil)
+
+func (c ConnectionCallbacksStruct) OnConnected(conn types.Connection) {
 	if c.OnConnectedFunc != nil {
 		c.OnConnectedFunc(conn)
 	}
 }
 
-func (c CallbacksStruct) OnMessage(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
+func (c ConnectionCallbacksStruct) OnMessage(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
 	if c.OnMessageFunc != nil {
 		return c.OnMessageFunc(conn, message)
 	} else {
@@ -40,7 +45,7 @@ func (c CallbacksStruct) OnMessage(conn types.Connection, message *protobufs.Age
 	}
 }
 
-func (c CallbacksStruct) OnConnectionClose(conn types.Connection) {
+func (c ConnectionCallbacksStruct) OnConnectionClose(conn types.Connection) {
 	if c.OnConnectionCloseFunc != nil {
 		c.OnConnectionCloseFunc(conn)
 	}
diff --git a/server/serverimpl.go b/server/serverimpl.go
index 9758acd..5fc7cac 100644
--- a/server/serverimpl.go
+++ b/server/serverimpl.go
@@ -15,6 +15,7 @@ import (
 	"github.com/open-telemetry/opamp-go/client/types"
 	"github.com/open-telemetry/opamp-go/internal"
 	"github.com/open-telemetry/opamp-go/protobufs"
+	serverTypes "github.com/open-telemetry/opamp-go/server/types"
 )
 
 var (
@@ -141,6 +142,7 @@ func (s *server) Stop(ctx context.Context) error {
 }
 
 func (s *server) httpHandler(w http.ResponseWriter, req *http.Request) {
+	var connectionCallbacks serverTypes.ConnectionCallbacks
 	if s.settings.Callbacks != nil {
 		resp := s.settings.Callbacks.OnConnecting(req)
 		if !resp.Accept {
@@ -152,13 +154,15 @@ func (s *server) httpHandler(w http.ResponseWriter, req *http.Request) {
 			w.WriteHeader(resp.HTTPStatusCode)
 			return
 		}
+		// use connection-specific handler provided by ConnectionResponse
+		connectionCallbacks = resp.ConnectionCallbacks
 	}
 
 	// HTTP connection is accepted. Check if it is a plain HTTP request.
 
 	if req.Header.Get(headerContentType) == contentTypeProtobuf {
 		// Yes, a plain HTTP request.
-		s.handlePlainHTTPRequest(req, w)
+		s.handlePlainHTTPRequest(req, w, connectionCallbacks)
 		return
 	}
 
@@ -171,10 +175,10 @@ func (s *server) httpHandler(w http.ResponseWriter, req *http.Request) {
 
 	// Return from this func to reduce memory usage.
 	// Handle the connection on a separate goroutine.
-	go s.handleWSConnection(conn)
+	go s.handleWSConnection(conn, connectionCallbacks)
 }
 
-func (s *server) handleWSConnection(wsConn *websocket.Conn) {
+func (s *server) handleWSConnection(wsConn *websocket.Conn, connectionCallbacks serverTypes.ConnectionCallbacks) {
 	agentConn := wsConnection{wsConn: wsConn}
 
 	defer func() {
@@ -186,13 +190,13 @@ func (s *server) handleWSConnection(wsConn *websocket.Conn) {
 			}
 		}()
 
-		if s.settings.Callbacks != nil {
-			s.settings.Callbacks.OnConnectionClose(agentConn)
+		if connectionCallbacks != nil {
+			connectionCallbacks.OnConnectionClose(agentConn)
 		}
 	}()
 
-	if s.settings.Callbacks != nil {
-		s.settings.Callbacks.OnConnected(agentConn)
+	if connectionCallbacks != nil {
+		connectionCallbacks.OnConnected(agentConn)
 	}
 
 	// Loop until fail to read from the WebSocket connection.
@@ -221,8 +225,8 @@ func (s *server) handleWSConnection(wsConn *websocket.Conn) {
 			continue
 		}
 
-		if s.settings.Callbacks != nil {
-			response := s.settings.Callbacks.OnMessage(agentConn, &request)
+		if connectionCallbacks != nil {
+			response := connectionCallbacks.OnMessage(agentConn, &request)
 			if response.InstanceUid == "" {
 				response.InstanceUid = request.InstanceUid
 			}
@@ -271,7 +275,7 @@ func compressGzip(data []byte) ([]byte, error) {
 	return buf.Bytes(), nil
 }
 
-func (s *server) handlePlainHTTPRequest(req *http.Request, w http.ResponseWriter) {
+func (s *server) handlePlainHTTPRequest(req *http.Request, w http.ResponseWriter, connectionCallbacks serverTypes.ConnectionCallbacks) {
 	bytes, err := s.readReqBody(req)
 	if err != nil {
 		s.logger.Debugf("Cannot read HTTP body: %v", err)
@@ -292,22 +296,22 @@ func (s *server) handlePlainHTTPRequest(req *http.Request, w http.ResponseWriter
 		conn: connFromRequest(req),
 	}
 
-	if s.settings.Callbacks == nil {
+	if connectionCallbacks == nil {
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}
 
-	s.settings.Callbacks.OnConnected(agentConn)
+	connectionCallbacks.OnConnected(agentConn)
 
 	defer func() {
 		// Indicate via the callback that the OpAMP Connection is closed. From OpAMP
 		// perspective the connection represented by this http request
 		// is closed. It is not possible to send or receive more OpAMP messages
 		// via this agentConn.
-		s.settings.Callbacks.OnConnectionClose(agentConn)
+		connectionCallbacks.OnConnectionClose(agentConn)
 	}()
 
-	response := s.settings.Callbacks.OnMessage(agentConn, &request)
+	response := connectionCallbacks.OnMessage(agentConn, &request)
 
 	// Set the InstanceUid if it is not set by the callback.
 	if response.InstanceUid == "" {
diff --git a/server/types/callbacks.go b/server/types/callbacks.go
index 1681554..addac7e 100644
--- a/server/types/callbacks.go
+++ b/server/types/callbacks.go
@@ -8,26 +8,35 @@ import (
 
 // ConnectionResponse is the return type of the OnConnecting callback.
 type ConnectionResponse struct {
-	Accept             bool
-	HTTPStatusCode     int
-	HTTPResponseHeader map[string]string
+	Accept              bool
+	HTTPStatusCode      int
+	HTTPResponseHeader  map[string]string
+	ConnectionCallbacks ConnectionCallbacks
 }
 
 type Callbacks interface {
-	// The following callbacks will never be called concurrently for the same
-	// connection. They may be called concurrently for different connections.
-
 	// OnConnecting is called when there is a new incoming connection.
 	// The handler can examine the request and either accept or reject the connection.
 	// To accept:
-	//   Return ConnectionResponse with Accept=true.
+	//   Return ConnectionResponse with Accept=true. ConnectionCallbacks MUST be set to an
+	//   implementation of the ConnectionCallbacks interface to handle the connection.
 	//   HTTPStatusCode and HTTPResponseHeader are ignored.
 	//
 	// To reject:
 	//   Return ConnectionResponse with Accept=false. HTTPStatusCode MUST be set to
 	//   non-zero value to indicate the rejection reason (typically 401, 429 or 503).
 	//   HTTPResponseHeader may be optionally set (e.g. "Retry-After: 30").
+	//   ConnectionCallbacks is ignored.
 	OnConnecting(request *http.Request) ConnectionResponse
+}
+
+// ConnectionCallbacks receives callbacks for a specific connection. An implementation of
+// this interface MUST be set on the ConnectionResponse returned by the OnConnecting
+// callback if Accept=true. The implementation can be shared by all connections or can be
+// unique for each connection.
+type ConnectionCallbacks interface {
+	// The following callbacks will never be called concurrently for the same
+	// connection. They may be called concurrently for different connections.
 
 	// OnConnected is called when and incoming OpAMP connection is successfully
 	// established after OnConnecting() returns.


Test patch:
diff --git a/server/serverimpl_test.go b/server/serverimpl_test.go
index 129518e..fcd013b 100644
--- a/server/serverimpl_test.go
+++ b/server/serverimpl_test.go
@@ -93,15 +93,16 @@ func TestServerStartAcceptConnection(t *testing.T) {
 	var srvConn types.Connection
 	callbacks := CallbacksStruct{
 		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
-			return types.ConnectionResponse{Accept: true}
-		},
-		OnConnectedFunc: func(conn types.Connection) {
-			srvConn = conn
-			atomic.StoreInt32(&connectedCalled, 1)
-		},
-		OnConnectionCloseFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&connectionCloseCalled, 1)
-			assert.EqualValues(t, srvConn, conn)
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					srvConn = conn
+					atomic.StoreInt32(&connectedCalled, 1)
+				},
+				OnConnectionCloseFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&connectionCloseCalled, 1)
+					assert.EqualValues(t, srvConn, conn)
+				},
+			}}
 		},
 	}
 
@@ -141,8 +142,12 @@ func TestDisconnectHttpConnection(t *testing.T) {
 func TestDisconnectWSConnection(t *testing.T) {
 	connectionCloseCalled := int32(0)
 	callback := CallbacksStruct{
-		OnConnectionCloseFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&connectionCloseCalled, 1)
+		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectionCloseFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&connectionCloseCalled, 1)
+				},
+			}}
 		},
 	}
 
@@ -176,18 +181,19 @@ func TestServerReceiveSendMessage(t *testing.T) {
 	var rcvMsg atomic.Value
 	callbacks := CallbacksStruct{
 		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
-			return types.ConnectionResponse{Accept: true}
-		},
-		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
-			// Remember received message.
-			rcvMsg.Store(message)
-
-			// Send a response.
-			response := protobufs.ServerToAgent{
-				InstanceUid:  message.InstanceUid,
-				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
-			}
-			return &response
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
+					// Remember received message.
+					rcvMsg.Store(message)
+
+					// Send a response.
+					response := protobufs.ServerToAgent{
+						InstanceUid:  message.InstanceUid,
+						Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
+					}
+					return &response
+				},
+			}}
 		},
 	}
 
@@ -242,25 +248,26 @@ func TestServerReceiveSendMessageWithCompression(t *testing.T) {
 			var rcvMsg atomic.Value
 			callbacks := CallbacksStruct{
 				OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
-					return types.ConnectionResponse{Accept: true}
-				},
-				OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
-					// Remember received message.
-					rcvMsg.Store(message)
-
-					// Send a response.
-					response := protobufs.ServerToAgent{
-						InstanceUid:  message.InstanceUid,
-						Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
-						RemoteConfig: &protobufs.AgentRemoteConfig{
-							Config: &protobufs.AgentConfigMap{
-								ConfigMap: map[string]*protobufs.AgentConfigFile{
-									"": {Body: uncompressedCfg},
+					return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+						OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
+							// Remember received message.
+							rcvMsg.Store(message)
+
+							// Send a response.
+							response := protobufs.ServerToAgent{
+								InstanceUid:  message.InstanceUid,
+								Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
+								RemoteConfig: &protobufs.AgentRemoteConfig{
+									Config: &protobufs.AgentConfigMap{
+										ConfigMap: map[string]*protobufs.AgentConfigFile{
+											"": {Body: uncompressedCfg},
+										},
+									},
 								},
-							},
+							}
+							return &response
 						},
-					}
-					return &response
+					}}
 				},
 			}
 
@@ -339,24 +346,25 @@ func TestServerReceiveSendMessagePlainHTTP(t *testing.T) {
 	var onConnectedCalled, onCloseCalled int32
 	callbacks := CallbacksStruct{
 		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
-			return types.ConnectionResponse{Accept: true}
-		},
-		OnConnectedFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&onConnectedCalled, 1)
-		},
-		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
-			// Remember received message.
-			rcvMsg.Store(message)
-
-			// Send a response.
-			response := protobufs.ServerToAgent{
-				InstanceUid:  message.InstanceUid,
-				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
-			}
-			return &response
-		},
-		OnConnectionCloseFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&onCloseCalled, 1)
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&onConnectedCalled, 1)
+				},
+				OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
+					// Remember received message.
+					rcvMsg.Store(message)
+
+					// Send a response.
+					response := protobufs.ServerToAgent{
+						InstanceUid:  message.InstanceUid,
+						Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
+					}
+					return &response
+				},
+				OnConnectionCloseFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&onCloseCalled, 1)
+				},
+			}}
 		},
 	}
 
@@ -406,15 +414,16 @@ func TestServerAttachAcceptConnection(t *testing.T) {
 	var srvConn types.Connection
 	callbacks := CallbacksStruct{
 		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
-			return types.ConnectionResponse{Accept: true}
-		},
-		OnConnectedFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&connectedCalled, 1)
-			srvConn = conn
-		},
-		OnConnectionCloseFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&connectionCloseCalled, 1)
-			assert.EqualValues(t, srvConn, conn)
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&connectedCalled, 1)
+					srvConn = conn
+				},
+				OnConnectionCloseFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&connectionCloseCalled, 1)
+					assert.EqualValues(t, srvConn, conn)
+				},
+			}}
 		},
 	}
 
@@ -455,26 +464,27 @@ func TestServerAttachSendMessagePlainHTTP(t *testing.T) {
 	var srvConn types.Connection
 	callbacks := CallbacksStruct{
 		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
-			return types.ConnectionResponse{Accept: true}
-		},
-		OnConnectedFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&connectedCalled, 1)
-			srvConn = conn
-		},
-		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
-			// Remember received message.
-			rcvMsg.Store(message)
-
-			// Send a response.
-			response := protobufs.ServerToAgent{
-				InstanceUid:  message.InstanceUid,
-				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
-			}
-			return &response
-		},
-		OnConnectionCloseFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&connectionCloseCalled, 1)
-			assert.EqualValues(t, srvConn, conn)
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&connectedCalled, 1)
+					srvConn = conn
+				},
+				OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
+					// Remember received message.
+					rcvMsg.Store(message)
+
+					// Send a response.
+					response := protobufs.ServerToAgent{
+						InstanceUid:  message.InstanceUid,
+						Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
+					}
+					return &response
+				},
+				OnConnectionCloseFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&connectionCloseCalled, 1)
+					assert.EqualValues(t, srvConn, conn)
+				},
+			}}
 		},
 	}
 
@@ -536,24 +546,25 @@ func TestServerHonoursClientRequestContentEncoding(t *testing.T) {
 	var onConnectedCalled, onCloseCalled int32
 	callbacks := CallbacksStruct{
 		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
-			return types.ConnectionResponse{Accept: true}
-		},
-		OnConnectedFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&onConnectedCalled, 1)
-		},
-		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
-			// Remember received message.
-			rcvMsg.Store(message)
-
-			// Send a response.
-			response := protobufs.ServerToAgent{
-				InstanceUid:  message.InstanceUid,
-				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
-			}
-			return &response
-		},
-		OnConnectionCloseFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&onCloseCalled, 1)
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&onConnectedCalled, 1)
+				},
+				OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
+					// Remember received message.
+					rcvMsg.Store(message)
+
+					// Send a response.
+					response := protobufs.ServerToAgent{
+						InstanceUid:  message.InstanceUid,
+						Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
+					}
+					return &response
+				},
+				OnConnectionCloseFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&onCloseCalled, 1)
+				},
+			}}
 		},
 	}
 
@@ -612,24 +623,25 @@ func TestServerHonoursAcceptEncoding(t *testing.T) {
 	var onConnectedCalled, onCloseCalled int32
 	callbacks := CallbacksStruct{
 		OnConnectingFunc: func(request *http.Request) types.ConnectionResponse {
-			return types.ConnectionResponse{Accept: true}
-		},
-		OnConnectedFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&onConnectedCalled, 1)
-		},
-		OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
-			// Remember received message.
-			rcvMsg.Store(message)
-
-			// Send a response.
-			response := protobufs.ServerToAgent{
-				InstanceUid:  message.InstanceUid,
-				Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
-			}
-			return &response
-		},
-		OnConnectionCloseFunc: func(conn types.Connection) {
-			atomic.StoreInt32(&onCloseCalled, 1)
+			return types.ConnectionResponse{Accept: true, ConnectionCallbacks: ConnectionCallbacksStruct{
+				OnConnectedFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&onConnectedCalled, 1)
+				},
+				OnMessageFunc: func(conn types.Connection, message *protobufs.AgentToServer) *protobufs.ServerToAgent {
+					// Remember received message.
+					rcvMsg.Store(message)
+
+					// Send a response.
+					response := protobufs.ServerToAgent{
+						InstanceUid:  message.InstanceUid,
+						Capabilities: uint64(protobufs.ServerCapabilities_ServerCapabilities_AcceptsStatus),
+					}
+					return &response
+				},
+				OnConnectionCloseFunc: func(conn types.Connection) {
+					atomic.StoreInt32(&onCloseCalled, 1)
+				},
+			}}
 		},
 	}