Repository: open-telemetry/opamp-go
Pull request #64 (2022-04-13 23:15:50)

Implement plain HTTP Transport in OpAMPClient
- Added httpClient implementation of OpAMP client.
- Renamed client to wsClient to indicates it is a WebSocket client.
- Split Receiver into wsReceiver that contains WebSocket specific code and
  receiveProcessor that contains transport-agnostic code and used by HTTP transport
  implementation too.
- Adjusted most existing client tests to run twice: once for wsClient and once
  for httpClient.
- Added a couple transport specific tests to httpclient_test.go and wsclient_test.go.


Patch:
diff --git a/client/client.go b/client/client.go
index 7851767..2a51938 100644
--- a/client/client.go
+++ b/client/client.go
@@ -10,9 +10,15 @@ import (
 
 type StartSettings struct {
 	// Connection parameters.
-	OpAMPServerURL      string
+
+	// Server URL. MUST be set.
+	OpAMPServerURL string
+
+	// Optional value of "Authorization" HTTP header in the HTTP request.
 	AuthorizationHeader string
-	TLSConfig           *tls.Config
+
+	// Optional TLS config for HTTP connection.
+	TLSConfig *tls.Config
 
 	// Agent information.
 	InstanceUid string
@@ -24,7 +30,7 @@ type StartSettings struct {
 	// Callbacks that the client will call after Start() returns nil.
 	Callbacks types.Callbacks
 
-	// Previously saved state. This will be reported to the server immediately
+	// Previously saved state. These will be reported to the server immediately
 	// after the connection is established.
 
 	// The hash of the last received remote config. If nil is passed it will force
@@ -81,9 +87,11 @@ type OpAMPClient interface {
 
 	// SetAgentDescription sets attributes of the Agent. The attributes will be included
 	// in the next status report sent to the server.
-	// MUST NOT be called before Start().
-	// MAY be called after Start(), in which case the attributes will be included
-	// in the next outgoing status report.
+	// May be called after Start(), in which case the attributes will be included
+	// in the next outgoing status report. This is typically used by Agents which allow
+	// their AgentDescription to change dynamically while the OpAMPClient is started.
+	// To define the initial agent description that is included in the first status report
+	// set StartSettings.AgentDescription field.
 	SetAgentDescription(descr *protobufs.AgentDescription) error
 
 	// SetEffectiveConfig sets the effective config of the Agent. The config will be
diff --git a/client/httpclient.go b/client/httpclient.go
new file mode 100644
index 0000000..1d17e98
--- /dev/null
+++ b/client/httpclient.go
@@ -0,0 +1,188 @@
+package client
+
+import (
+	"context"
+	"net/url"
+	"sync"
+
+	"github.com/open-telemetry/opamp-go/client/internal"
+	"github.com/open-telemetry/opamp-go/client/types"
+	sharedinternal "github.com/open-telemetry/opamp-go/internal"
+	"github.com/open-telemetry/opamp-go/protobufs"
+)
+
+// httpClient is an OpAMP Client implementation for plain HTTP transport.
+// See specification: https://github.com/open-telemetry/opamp-spec/blob/main/specification.md#plain-http-transport
+type httpClient struct {
+	logger   types.Logger
+	settings StartSettings
+
+	// OpAMP Server URL.
+	url *url.URL
+
+	// True if Start() is successful.
+	isStarted bool
+
+	// Cancellation func for background go routines.
+	runCancel context.CancelFunc
+
+	// True when stopping is in progress.
+	isStoppingFlag  bool
+	isStoppingMutex sync.RWMutex
+
+	// Indicates that the Client is fully stopped.
+	stoppedSignal chan struct{}
+
+	// The looper performs HTTP request/response loop.
+	looper *internal.HTTPLooper
+}
+
+var _ OpAMPClient = (*httpClient)(nil)
+
+func NewHTTP(logger types.Logger) *httpClient {
+	if logger == nil {
+		logger = &sharedinternal.NopLogger{}
+	}
+
+	w := &httpClient{
+		logger:        logger,
+		stoppedSignal: make(chan struct{}, 1),
+		looper:        internal.NewHTTPLooper(logger),
+	}
+	return w
+}
+
+func (w *httpClient) Start(settings StartSettings) error {
+	if w.isStarted {
+		return errAlreadyStarted
+	}
+	if settings.AgentDescription == nil {
+		return errAgentDescriptionMissing
+	}
+
+	w.settings = settings
+	if w.settings.Callbacks == nil {
+		// Make sure it is always safe to call Callbacks.
+		w.settings.Callbacks = types.CallbacksStruct{}
+	}
+
+	if err := w.looper.SetInstanceUid(w.settings.InstanceUid); err != nil {
+		return err
+	}
+
+	// Prepare server connection settings.
+	var err error
+	w.url, err = url.Parse(w.settings.OpAMPServerURL)
+	if err != nil {
+		return err
+	}
+
+	if w.settings.TLSConfig != nil {
+		w.url.Scheme = "https"
+	}
+
+	if w.settings.AuthorizationHeader != "" {
+		w.looper.SetRequestHeader("Authorization", w.settings.AuthorizationHeader)
+	}
+
+	// Prepare the first message to send.
+	w.looper.NextMessage().Update(
+		func(msg *protobufs.AgentToServer) {
+			if msg.StatusReport == nil {
+				msg.StatusReport = &protobufs.StatusReport{}
+			}
+			msg.StatusReport.AgentDescription = w.settings.AgentDescription
+			if msg.AddonStatuses == nil {
+				msg.AddonStatuses = &protobufs.AgentAddonStatuses{}
+			}
+			msg.AddonStatuses.ServerProvidedAllAddonsHash = w.settings.LastServerProvidedAllAddonsHash
+		},
+	)
+	w.looper.ScheduleSend()
+
+	w.startConnectAndRun()
+
+	w.isStarted = true
+
+	return nil
+}
+
+func (w *httpClient) Stop(ctx context.Context) error {
+	if !w.isStarted {
+		return errCannotStopNotStarted
+	}
+
+	w.isStoppingMutex.Lock()
+	cancelFunc := w.runCancel
+	w.isStoppingFlag = true
+	w.isStoppingMutex.Unlock()
+
+	cancelFunc()
+
+	// Wait until stopping is finished.
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-w.stoppedSignal:
+	}
+	return nil
+}
+
+func (w *httpClient) SetAgentDescription(descr *protobufs.AgentDescription) error {
+	if descr == nil {
+		return errAgentDescriptionMissing
+	}
+
+	// store the agent description to send on reconnect
+	w.settings.AgentDescription = descr
+	w.looper.NextMessage().UpdateStatus(func(statusReport *protobufs.StatusReport) {
+		statusReport.AgentDescription = descr
+	})
+	w.looper.ScheduleSend()
+	return nil
+}
+
+func (w *httpClient) SetEffectiveConfig(config *protobufs.EffectiveConfig) error {
+	w.looper.NextMessage().UpdateStatus(func(statusReport *protobufs.StatusReport) {
+		statusReport.EffectiveConfig = config
+	})
+	w.looper.ScheduleSend()
+	return nil
+}
+
+func (w *httpClient) isStopping() bool {
+	w.isStoppingMutex.RLock()
+	defer w.isStoppingMutex.RUnlock()
+	return w.isStoppingFlag
+}
+
+func (w *httpClient) startConnectAndRun() {
+	// Create a cancellable context.
+	runCtx, runCancel := context.WithCancel(context.Background())
+
+	w.isStoppingMutex.Lock()
+	defer w.isStoppingMutex.Unlock()
+
+	if w.isStoppingFlag {
+		// Stop() was called. Don't connect.
+		runCancel()
+		return
+	}
+	w.runCancel = runCancel
+
+	go w.runUntilStopped(runCtx)
+}
+
+func (w *httpClient) runUntilStopped(ctx context.Context) {
+
+	defer func() {
+		// We only return from runUntilStopped when we are instructed to stop.
+		// When returning signal that we stopped.
+		w.stoppedSignal <- struct{}{}
+	}()
+
+	// Start the HTTP looper. This will make request/responses with retries for
+	// failures and will wait with configured polling interval if there is nothing
+	// to send.
+	w.looper.Run(ctx, w.settings.OpAMPServerURL, w.settings.Callbacks)
+}
diff --git a/client/internal/httplooper.go b/client/internal/httplooper.go
new file mode 100644
index 0000000..762b16b
--- /dev/null
+++ b/client/internal/httplooper.go
@@ -0,0 +1,257 @@
+package internal
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"sync/atomic"
+	"time"
+
+	"github.com/cenkalti/backoff/v4"
+	"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"
+)
+
+const OpAMPPlainHTTPMethod = "POST"
+const defaultPollingIntervalMs = 30 * 1000 // default interval is 30 seconds.
+
+// HTTPLooper allows scheduling messages to send. Once run, it will loop through
+// a request/response cycle for each message to send and will process all received
+// responses using a receivedProcessor. If there are no pending messages to send
+// the HTTPLooper will wait for the configured polling interval.
+type HTTPLooper struct {
+	url               string
+	instanceUid       atomic.Value
+	logger            types.Logger
+	client            *http.Client
+	callbacks         types.Callbacks
+	pollingIntervalMs int64
+
+	// Headers to send with all requests.
+	requestHeader http.Header
+
+	// Indicates that there is a pending message to send.
+	hasPendingMessage chan struct{}
+
+	// The next message to send.
+	nextMessage nextMessage
+
+	// Processor to handle received messages.
+	receiveProcessor receivedProcessor
+}
+
+func NewHTTPLooper(logger types.Logger) *HTTPLooper {
+	h := &HTTPLooper{
+		logger:            logger,
+		client:            http.DefaultClient,
+		hasPendingMessage: make(chan struct{}, 1),
+		requestHeader:     http.Header{},
+		pollingIntervalMs: defaultPollingIntervalMs,
+	}
+	h.requestHeader.Set(headerContentType, contentTypeProtobuf)
+	return h
+}
+
+// Run starts the processing loop that will perform the HTTP request/response.
+// When there are no more messages to send Run will suspend until either there is
+// a new message to send or the polling interval elapses.
+// Should not be called concurrently with itself. Can be called concurrently with
+// modifying nextMessage().
+// Run continues until ctx is cancelled.
+func (h *HTTPLooper) Run(ctx context.Context, url string, callbacks types.Callbacks) {
+	h.url = url
+	h.callbacks = callbacks
+	h.receiveProcessor = receivedProcessor{
+		callbacks: callbacks,
+		logger:    h.logger,
+		sender:    h,
+	}
+
+	for {
+		pollingTimer := time.NewTimer(time.Millisecond * time.Duration(atomic.LoadInt64(&h.pollingIntervalMs)))
+		select {
+		case <-h.hasPendingMessage:
+			// Have something to send. Stop the polling timer and send what we have.
+			pollingTimer.Stop()
+			h.makeOneRequestRoundtrip(ctx)
+
+		case <-pollingTimer.C:
+			// Polling interval has passed. Force a status update.
+			h.NextMessage().UpdateStatus(func(statusReport *protobufs.StatusReport) {})
+			// This will make hasPendingMessage channel readable, so we will enter
+			// the case above on the next iteration of the loop.
+			h.ScheduleSend()
+			break
+
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
+// ScheduleSend signals to HTTPLooper that the message in nextMessage struct
+// is now ready to be sent. If there is no pending message (e.g. the nextMessage was
+// already sent and "pending" flag is reset) then no message will be sent.
+func (h *HTTPLooper) ScheduleSend() {
+	// Set pending flag. Don't block on writing to channel.
+	select {
+	case h.hasPendingMessage <- struct{}{}:
+	default:
+		break
+	}
+}
+
+// SetInstanceUid sets a new instanceUid to be used for all subsequent messages to be sent.
+// Can be called concurrently, normally is called when a message is received from the
+// Server that instructs us to change our instance UID.
+func (h *HTTPLooper) SetInstanceUid(instanceUid string) error {
+	if instanceUid == "" {
+		return fmt.Errorf("cannot set instance uid to empty value")
+	}
+	h.instanceUid.Store(instanceUid)
+	return nil
+}
+
+// SetRequestHeader sets an additional HTTP header to send with all future requests.
+// Should not be called concurrently with any other method.
+func (h *HTTPLooper) SetRequestHeader(key, value string) {
+	h.requestHeader.Set(key, value)
+}
+
+// nextMessage gives access to the next message that will be sent by this looper.
+// Can be called concurrently with any other method.
+func (h *HTTPLooper) NextMessage() *nextMessage {
+	return &h.nextMessage
+}
+
+func (h *HTTPLooper) makeOneRequestRoundtrip(ctx context.Context) {
+	resp, err := h.sendRequestWithRetries(ctx)
+	if err != nil {
+		return
+	}
+	h.receiveResponse(ctx, resp)
+}
+
+func (h *HTTPLooper) sendRequestWithRetries(ctx context.Context) (*http.Response, error) {
+	req, err := h.prepareRequest(ctx)
+	if err != nil {
+		if errors.Is(err, context.Canceled) {
+			h.logger.Debugf("Client is stopped, will not try anymore.")
+		} else {
+			h.logger.Errorf("Failed prepare request (%v), will not try anymore.", err)
+		}
+		return nil, err
+	}
+
+	// Repeatedly try requests with a backoff strategy.
+	infiniteBackoff := backoff.NewExponentialBackOff()
+	// Make backoff run forever.
+	infiniteBackoff.MaxElapsedTime = 0
+
+	interval := time.Duration(0)
+
+	for {
+		timer := time.NewTimer(interval)
+		interval = infiniteBackoff.NextBackOff()
+
+		select {
+		case <-timer.C:
+			{
+				resp, err := h.client.Do(req)
+				if err == nil {
+					switch resp.StatusCode {
+					case http.StatusOK:
+						// We consider it connected if we receive 200 status from the Server.
+						h.callbacks.OnConnect()
+						return resp, nil
+
+					case http.StatusTooManyRequests:
+					case http.StatusServiceUnavailable:
+						interval = recalculateInterval(interval, resp)
+						err = fmt.Errorf("server response code=%d", resp.StatusCode)
+
+					default:
+						return nil, fmt.Errorf("invalid response from server: %d", resp.StatusCode)
+					}
+				} else if errors.Is(err, context.Canceled) {
+					h.logger.Debugf("Client is stopped, will not try anymore.")
+					return nil, err
+				}
+
+				h.logger.Errorf("Failed to do HTTP request (%v), will retry", err)
+				h.callbacks.OnConnectFailed(err)
+			}
+
+		case <-ctx.Done():
+			h.logger.Debugf("Client is stopped, will not try anymore.")
+			return nil, ctx.Err()
+		}
+	}
+}
+
+func recalculateInterval(interval time.Duration, resp *http.Response) time.Duration {
+	retryAfter := internal.ExtractRetryAfterHeader(resp)
+	if retryAfter.Defined && retryAfter.Duration > interval {
+		// If the server suggested connecting later than our interval
+		// then honour server's request, otherwise wait at least
+		// as much as we calculated.
+		interval = retryAfter.Duration
+	}
+	return interval
+}
+
+func (h *HTTPLooper) prepareRequest(ctx context.Context) (*http.Request, error) {
+	msgToSend := h.nextMessage.PopPending()
+
+	if msgToSend == nil || proto.Equal(msgToSend, &protobufs.AgentToServer{}) {
+		// There is no pending message or the message is empty.
+		// Nothing to send.
+		return nil, nil
+	}
+
+	msgToSend.InstanceUid = h.instanceUid.Load().(string)
+
+	data, err := proto.Marshal(msgToSend)
+	if err != nil {
+		return nil, err
+	}
+
+	body := bytes.NewReader(data)
+	req, err := http.NewRequestWithContext(ctx, OpAMPPlainHTTPMethod, h.url, body)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header = h.requestHeader
+	return req, nil
+}
+
+func (h *HTTPLooper) receiveResponse(ctx context.Context, resp *http.Response) {
+	msgBytes, err := io.ReadAll(resp.Body)
+	if err != nil {
+		_ = resp.Body.Close()
+		h.logger.Errorf("cannot read response body: %v", err)
+		return
+	}
+	_ = resp.Body.Close()
+
+	var response protobufs.ServerToAgent
+	if err := proto.Unmarshal(msgBytes, &response); err != nil {
+		h.logger.Errorf("cannot unmarshal response: %v", err)
+		return
+	}
+
+	h.receiveProcessor.ProcessReceivedMessage(ctx, &response)
+}
+
+// SetPollingInterval sets the interval between polling. Has effect starting from the
+// next polling cycle.
+func (h *HTTPLooper) SetPollingInterval(duration time.Duration) {
+	atomic.StoreInt64(&h.pollingIntervalMs, duration.Milliseconds())
+}
diff --git a/client/internal/mockserver.go b/client/internal/mockserver.go
index ce64886..46a57ab 100644
--- a/client/internal/mockserver.go
+++ b/client/internal/mockserver.go
@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"io"
 	"log"
 	"net/http"
 	"net/http/httptest"
@@ -16,13 +17,17 @@ import (
 )
 
 type MockServer struct {
-	Endpoint  string
-	OnRequest func(w http.ResponseWriter, r *http.Request)
-	OnConnect func(r *http.Request, conn *websocket.Conn)
-	OnMessage func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent
-	srv       *httptest.Server
+	Endpoint    string
+	OnRequest   func(w http.ResponseWriter, r *http.Request)
+	OnConnect   func(r *http.Request)
+	OnWSConnect func(conn *websocket.Conn)
+	OnMessage   func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent
+	srv         *httptest.Server
 }
 
+const headerContentType = "Content-Type"
+const contentTypeProtobuf = "application/x-protobuf"
+
 var upgrader = websocket.Upgrader{}
 
 func StartMockServer(t *testing.T) *MockServer {
@@ -36,41 +41,16 @@ func StartMockServer(t *testing.T) *MockServer {
 				return
 			}
 
-			conn, err := upgrader.Upgrade(w, r, nil)
-			if err != nil {
-				return
-			}
 			if srv.OnConnect != nil {
-				srv.OnConnect(r, conn)
+				srv.OnConnect(r)
 			}
-			for {
-				var messageType int
-				var msgBytes []byte
-				if messageType, msgBytes, err = conn.ReadMessage(); err != nil {
-					return
-				}
-				assert.EqualValues(t, websocket.BinaryMessage, messageType)
-
-				var dest protobufs.AgentToServer
-				err := proto.Unmarshal(msgBytes, &dest)
-				if err != nil {
-					log.Fatal("cannot decode:", err)
-				}
 
-				if srv.OnMessage != nil {
-					response := srv.OnMessage(&dest)
-					if response != nil {
-						msgBytes, err := proto.Marshal(response)
-						if err != nil {
-							log.Fatal("cannot encode:", err)
-						}
-						err = conn.WriteMessage(websocket.BinaryMessage, msgBytes)
-						if err != nil {
-							log.Fatal("cannot send:", err)
-						}
-					}
-				}
+			if r.Header.Get(headerContentType) == contentTypeProtobuf {
+				srv.handlePlainHttp(w, r)
+				return
 			}
+
+			srv.handleWebSocket(t, w, r)
 		},
 	)
 
@@ -87,6 +67,76 @@ func StartMockServer(t *testing.T) *MockServer {
 	return srv
 }
 
+func (srv *MockServer) handlePlainHttp(w http.ResponseWriter, r *http.Request) {
+	msgBytes, err := io.ReadAll(r.Body)
+
+	var request protobufs.AgentToServer
+	err = proto.Unmarshal(msgBytes, &request)
+	if err != nil {
+		log.Fatal("cannot decode:", err)
+	}
+
+	var response *protobufs.ServerToAgent
+	if srv.OnMessage != nil {
+		response = srv.OnMessage(&request)
+	}
+	if response == nil {
+		response = &protobufs.ServerToAgent{
+			InstanceUid: request.InstanceUid,
+		}
+	}
+	msgBytes, err = proto.Marshal(response)
+	if err != nil {
+		log.Fatal("cannot encode:", err)
+	}
+
+	// Send the response.
+	w.Header().Set(headerContentType, contentTypeProtobuf)
+	_, err = w.Write(msgBytes)
+	if err != nil {
+		log.Fatal("cannot send:", err)
+	}
+
+}
+
+func (srv *MockServer) handleWebSocket(t *testing.T, w http.ResponseWriter, r *http.Request) {
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		return
+	}
+	if srv.OnWSConnect != nil {
+		srv.OnWSConnect(conn)
+	}
+	for {
+		var messageType int
+		var msgBytes []byte
+		if messageType, msgBytes, err = conn.ReadMessage(); err != nil {
+			return
+		}
+		assert.EqualValues(t, websocket.BinaryMessage, messageType)
+
+		var dest protobufs.AgentToServer
+		err := proto.Unmarshal(msgBytes, &dest)
+		if err != nil {
+			log.Fatal("cannot decode:", err)
+		}
+
+		if srv.OnMessage != nil {
+			response := srv.OnMessage(&dest)
+			if response != nil {
+				msgBytes, err := proto.Marshal(response)
+				if err != nil {
+					log.Fatal("cannot encode:", err)
+				}
+				err = conn.WriteMessage(websocket.BinaryMessage, msgBytes)
+				if err != nil {
+					log.Fatal("cannot send:", err)
+				}
+			}
+		}
+	}
+}
+
 func (m *MockServer) Close() {
 	m.srv.Close()
 }
diff --git a/client/internal/nextmessage.go b/client/internal/nextmessage.go
new file mode 100644
index 0000000..44c5233
--- /dev/null
+++ b/client/internal/nextmessage.go
@@ -0,0 +1,60 @@
+package internal
+
+import (
+	"sync"
+
+	"github.com/open-telemetry/opamp-go/protobufs"
+	"google.golang.org/protobuf/proto"
+)
+
+// nextMessage encapsulates the next message to be sent and provides a
+// concurrency-safe interface to work with the message.
+type nextMessage struct {
+	// The next message to send.
+	nextMessage protobufs.AgentToServer
+	// Indicates that nextMessage is pending to be sent.
+	messagePending bool
+	// Mutex to protect the above 2 fields.
+	messageMutex sync.Mutex
+}
+
+// Update applies the specified modifier function to the next message that
+// will be sent and marks the message as pending to be sent.
+func (s *nextMessage) Update(modifier func(msg *protobufs.AgentToServer)) {
+	s.messageMutex.Lock()
+	modifier(&s.nextMessage)
+	s.messagePending = true
+	s.messageMutex.Unlock()
+}
+
+// UpdateStatus applies the specified modifier function to the status report that
+// will be sent next and marks the status report as pending to be sent.
+func (s *nextMessage) UpdateStatus(modifier func(statusReport *protobufs.StatusReport)) {
+	s.Update(
+		func(msg *protobufs.AgentToServer) {
+			if s.nextMessage.StatusReport == nil {
+				s.nextMessage.StatusReport = &protobufs.StatusReport{}
+			}
+			modifier(s.nextMessage.StatusReport)
+		},
+	)
+}
+
+// PopPending returns the next message to be sent, if it is pending or nil otherwise.
+// Clears the "pending" flag.
+func (s *nextMessage) PopPending() *protobufs.AgentToServer {
+	var msgToSend *protobufs.AgentToServer
+	s.messageMutex.Lock()
+	if s.messagePending {
+		// Clone the message to have a copy for sending and avoid blocking
+		// future updates to s.nextMessage field.
+		msgToSend = proto.Clone(&s.nextMessage).(*protobufs.AgentToServer)
+		s.messagePending = false
+
+		// Reset fields that we do not have to send unless they change before the
+		// next report after this one.
+		s.nextMessage = protobufs.AgentToServer{}
+	}
+	s.messageMutex.Unlock()
+	return msgToSend
+}
diff --git a/client/internal/receiver.go b/client/internal/receivedprocessor.go
similarity index 62%
rename from client/internal/receiver.go
rename to client/internal/receivedprocessor.go
index 94f3c0d..010d3f0 100644
--- a/client/internal/receiver.go
+++ b/client/internal/receivedprocessor.go
@@ -2,64 +2,38 @@ package internal
 
 import (
 	"context"
-	"fmt"
-
-	"github.com/gorilla/websocket"
-	"google.golang.org/protobuf/proto"
 
 	"github.com/open-telemetry/opamp-go/client/types"
 	"github.com/open-telemetry/opamp-go/protobufs"
 )
 
-// Receiver implements the client's receiving portion of OpAMP protocol.
-type Receiver struct {
-	conn      *websocket.Conn
-	logger    types.Logger
-	sender    *Sender
+// receivedProcessor handles the processing of messages received from the Server.
+type receivedProcessor struct {
+	logger types.Logger
+
+	// Callbacks to call for corresponding messages.
 	callbacks types.Callbacks
-}
 
-func NewReceiver(logger types.Logger, callbacks types.Callbacks, conn *websocket.Conn, sender *Sender) *Receiver {
-	return &Receiver{
-		conn:      conn,
-		logger:    logger,
-		sender:    sender,
-		callbacks: callbacks,
-	}
+	// A sender to cooperate with when the received message has an impact on
+	// what will be sent later.
+	sender sender
 }
 
-func (r *Receiver) ReceiverLoop(ctx context.Context) {
-	runContext, cancelFunc := context.WithCancel(ctx)
+type sender interface {
+	// NextMessage provides access to modify the next message that will be sent.
+	NextMessage() *nextMessage
 
-out:
-	for {
-		var message protobufs.ServerToAgent
-		if err := r.receiveMessage(&message); err != nil {
-			if ctx.Err() == nil && !websocket.IsCloseError(err, websocket.CloseNormalClosure) {
-				r.logger.Errorf("Unexpected error while receiving: %v", err)
-			}
-			break out
-		} else {
-			r.processReceivedMessage(runContext, &message)
-		}
-	}
+	// ScheduleSend signals the sender to send a pending next message.
+	ScheduleSend()
 
-	cancelFunc()
-}
-
-func (r *Receiver) receiveMessage(msg *protobufs.ServerToAgent) error {
-	_, bytes, err := r.conn.ReadMessage()
-	if err != nil {
-		return err
-	}
-	err = proto.Unmarshal(bytes, msg)
-	if err != nil {
-		return fmt.Errorf("cannot decode received message: %w", err)
-	}
-	return err
+	// SetInstanceUid sets a new instanceUid to be used when next message is being sent.
+	SetInstanceUid(instanceUid string) error
 }
 
-func (r *Receiver) processReceivedMessage(ctx context.Context, msg *protobufs.ServerToAgent) {
+// ProcessReceivedMessage is the entry point into the processing routine. It examines
+// the received message and performs any processing necessary based on what fields are set.
+// This function will call any relevant callbacks.
+func (r *receivedProcessor) ProcessReceivedMessage(ctx context.Context, msg *protobufs.ServerToAgent) {
 	if r.callbacks != nil {
 		// If a command message exists, other messages will be ignored
 		if msg.Command != nil {
@@ -84,7 +58,7 @@ func (r *Receiver) processReceivedMessage(ctx context.Context, msg *protobufs.Se
 	}
 }
 
-func (r *Receiver) rcvRemoteConfig(
+func (r *receivedProcessor) rcvRemoteConfig(
 	ctx context.Context,
 	config *protobufs.AgentRemoteConfig,
 	flags protobufs.ServerToAgent_Flags,
@@ -92,7 +66,7 @@ func (r *Receiver) rcvRemoteConfig(
 	effective, changed, err := r.callbacks.OnRemoteConfig(ctx, config)
 	if err != nil {
 		// Return a status message with status failed.
-		r.sender.UpdateNextStatus(func(statusReport *protobufs.StatusReport) {
+		r.sender.NextMessage().UpdateStatus(func(statusReport *protobufs.StatusReport) {
 			statusReport.RemoteConfigStatus = &protobufs.RemoteConfigStatus{
 				Status:       protobufs.RemoteConfigStatus_Failed,
 				ErrorMessage: err.Error(),
@@ -106,7 +80,7 @@ func (r *Receiver) rcvRemoteConfig(
 	reportEffective := changed || (flags&protobufs.ServerToAgent_ReportEffectiveConfig != 0)
 
 	if reportEffective {
-		r.sender.UpdateNextStatus(func(statusReport *protobufs.StatusReport) {
+		r.sender.NextMessage().UpdateStatus(func(statusReport *protobufs.StatusReport) {
 			statusReport.EffectiveConfig = effective
 		})
 		return true
@@ -114,7 +88,7 @@ func (r *Receiver) rcvRemoteConfig(
 	return false
 }
 
-func (r *Receiver) rcvConnectionSettings(ctx context.Context, settings *protobufs.ConnectionSettingsOffers) {
+func (r *receivedProcessor) rcvConnectionSettings(ctx context.Context, settings *protobufs.ConnectionSettingsOffers) {
 	if settings == nil {
 		return
 	}
@@ -153,7 +127,7 @@ func (r *Receiver) rcvConnectionSettings(ctx context.Context, settings *protobuf
 	}
 }
 
-func (r *Receiver) rcvOwnTelemetryConnectionSettings(
+func (r *receivedProcessor) rcvOwnTelemetryConnectionSettings(
 	ctx context.Context,
 	settings *protobufs.ConnectionSettings,
 	telemetryType types.OwnTelemetryType,
@@ -164,7 +138,7 @@ func (r *Receiver) rcvOwnTelemetryConnectionSettings(
 	}
 }
 
-func (r *Receiver) rcvOtherConnectionSettings(
+func (r *receivedProcessor) rcvOtherConnectionSettings(
 	ctx context.Context,
 	settings *protobufs.ConnectionSettings,
 	name string,
@@ -175,15 +149,16 @@ func (r *Receiver) rcvOtherConnectionSettings(
 	}
 }
 
-func (r *Receiver) processErrorResponse(body *protobufs.ServerErrorResponse) {
+func (r *receivedProcessor) processErrorResponse(body *protobufs.ServerErrorResponse) {
 	// TODO: implement this.
+	r.logger.Errorf("received an error from server: %s", body.ErrorMessage)
 }
 
-func (r *Receiver) rcvAddonsAvailable(addons *protobufs.AddonsAvailable) {
+func (r *receivedProcessor) rcvAddonsAvailable(addons *protobufs.AddonsAvailable) {
 	// TODO: implement this.
 }
 
-func (r *Receiver) rcvAgentIdentification(ctx context.Context, agentId *protobufs.AgentIdentification) {
+func (r *receivedProcessor) rcvAgentIdentification(ctx context.Context, agentId *protobufs.AgentIdentification) {
 	if agentId == nil {
 		return
 	}
@@ -205,7 +180,7 @@ func (r *Receiver) rcvAgentIdentification(ctx context.Context, agentId *protobuf
 	}
 }
 
-func (r *Receiver) rcvCommand(command *protobufs.ServerToAgentCommand) {
+func (r *receivedProcessor) rcvCommand(command *protobufs.ServerToAgentCommand) {
 	if command != nil {
 		r.callbacks.OnCommand(command)
 	}
diff --git a/client/internal/wsreceiver.go b/client/internal/wsreceiver.go
new file mode 100644
index 0000000..8689594
--- /dev/null
+++ b/client/internal/wsreceiver.go
@@ -0,0 +1,66 @@
+package internal
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/gorilla/websocket"
+	"google.golang.org/protobuf/proto"
+
+	"github.com/open-telemetry/opamp-go/client/types"
+	"github.com/open-telemetry/opamp-go/protobufs"
+)
+
+// wsReceiver implements the WebSocket client's receiving portion of OpAMP protocol.
+type wsReceiver struct {
+	conn      *websocket.Conn
+	logger    types.Logger
+	sender    *WSSender
+	callbacks types.Callbacks
+	processor receivedProcessor
+}
+
+func NewWSReceiver(logger types.Logger, callbacks types.Callbacks, conn *websocket.Conn, sender *WSSender) *wsReceiver {
+	return &wsReceiver{
+		conn:      conn,
+		logger:    logger,
+		sender:    sender,
+		callbacks: callbacks,
+		processor: receivedProcessor{
+			callbacks: callbacks,
+			logger:    logger,
+			sender:    sender,
+		},
+	}
+}
+
+func (r *wsReceiver) ReceiverLoop(ctx context.Context) {
+	runContext, cancelFunc := context.WithCancel(ctx)
+
+out:
+	for {
+		var message protobufs.ServerToAgent
+		if err := r.receiveMessage(&message); err != nil {
+			if ctx.Err() == nil && !websocket.IsCloseError(err, websocket.CloseNormalClosure) {
+				r.logger.Errorf("Unexpected error while receiving: %v", err)
+			}
+			break out
+		} else {
+			r.processor.ProcessReceivedMessage(runContext, &message)
+		}
+	}
+
+	cancelFunc()
+}
+
+func (r *wsReceiver) receiveMessage(msg *protobufs.ServerToAgent) error {
+	_, bytes, err := r.conn.ReadMessage()
+	if err != nil {
+		return err
+	}
+	err = proto.Unmarshal(bytes, msg)
+	if err != nil {
+		return fmt.Errorf("cannot decode received message: %w", err)
+	}
+	return err
+}
diff --git a/client/internal/sender.go b/client/internal/wssender.go
similarity index 52%
rename from client/internal/sender.go
rename to client/internal/wssender.go
index 8810c43..2edbc54 100644
--- a/client/internal/sender.go
+++ b/client/internal/wssender.go
@@ -3,7 +3,6 @@ package internal
 import (
 	"context"
 	"fmt"
-	"sync"
 	"sync/atomic"
 
 	"github.com/gorilla/websocket"
@@ -13,8 +12,8 @@ import (
 	"github.com/open-telemetry/opamp-go/protobufs"
 )
 
-// Sender implements the client's sending portion of OpAMP protocol.
-type Sender struct {
+// WSSender implements the WebSocket client's sending portion of OpAMP protocol.
+type WSSender struct {
 	instanceUid atomic.Value
 	conn        *websocket.Conn
 
@@ -24,26 +23,22 @@ type Sender struct {
 	hasMessages chan struct{}
 
 	// The next message to send.
-	nextMessage protobufs.AgentToServer
-	// Indicates that nextMessage is pending to be sent.
-	messagePending bool
-	// Mutex to protect the above 2 fields.
-	messageMutex sync.Mutex
+	nextMessage nextMessage
 
 	// Indicates that the sender has fully stopped.
 	stopped chan struct{}
 }
 
-func NewSender(logger types.Logger) *Sender {
-	return &Sender{
+func NewSender(logger types.Logger) *WSSender {
+	return &WSSender{
 		logger:      logger,
 		hasMessages: make(chan struct{}, 1),
 	}
 }
 
-// Start the sender and send the first message that was set via UpdateNextMessage()
-// earlier. To stop the Sender cancel the ctx.
-func (s *Sender) Start(ctx context.Context, instanceUid string, conn *websocket.Conn) error {
+// Start the sender and send the first message that was set via NextMessage().Update()
+// earlier. To stop the WSSender cancel the ctx.
+func (s *WSSender) Start(ctx context.Context, instanceUid string, conn *websocket.Conn) error {
 	s.conn = conn
 	s.instanceUid.Store(instanceUid)
 	err := s.sendNextMessage()
@@ -57,7 +52,7 @@ func (s *Sender) Start(ctx context.Context, instanceUid string, conn *websocket.
 
 // SetInstanceUid sets a new instanceUid without closing and reopening the connection. It will be used
 // when next message is being sent.
-func (s *Sender) SetInstanceUid(instanceUid string) error {
+func (s *WSSender) SetInstanceUid(instanceUid string) error {
 	if instanceUid == "" {
 		return fmt.Errorf("cannot set instance uid to empty value")
 	}
@@ -65,38 +60,20 @@ func (s *Sender) SetInstanceUid(instanceUid string) error {
 	return nil
 }
 
+func (s *WSSender) NextMessage() *nextMessage {
+	return &s.nextMessage
+}
+
 // WaitToStop blocks until the sender is stopped. To stop the sender cancel the context
 // that was passed to Start().
-func (s *Sender) WaitToStop() {
+func (s *WSSender) WaitToStop() {
 	<-s.stopped
 }
 
-// UpdateNextMessage applies the specified modifier function to the next message that
-// will be sent and marks the message as pending to be sent.
-func (s *Sender) UpdateNextMessage(modifier func(msg *protobufs.AgentToServer)) {
-	s.messageMutex.Lock()
-	modifier(&s.nextMessage)
-	s.messagePending = true
-	s.messageMutex.Unlock()
-}
-
-// UpdateNextStatus applies the specified modifier function to the status report that
-// will be sent next and marks the status report as pending to be sent.
-func (s *Sender) UpdateNextStatus(modifier func(statusReport *protobufs.StatusReport)) {
-	s.UpdateNextMessage(
-		func(msg *protobufs.AgentToServer) {
-			if s.nextMessage.StatusReport == nil {
-				s.nextMessage.StatusReport = &protobufs.StatusReport{}
-			}
-			modifier(s.nextMessage.StatusReport)
-		},
-	)
-}
-
 // ScheduleSend signals to the sending goroutine to send the next message
 // if it is pending. If there is no pending message (e.g. the message was
 // already sent and "pending" flag is reset) then no message will be be sent.
-func (s *Sender) ScheduleSend() {
+func (s *WSSender) ScheduleSend() {
 	select {
 	case s.hasMessages <- struct{}{}:
 	default:
@@ -104,7 +81,7 @@ func (s *Sender) ScheduleSend() {
 	}
 }
 
-func (s *Sender) run(ctx context.Context) {
+func (s *WSSender) run(ctx context.Context) {
 out:
 	for {
 		select {
@@ -119,21 +96,8 @@ out:
 	close(s.stopped)
 }
 
-func (s *Sender) sendNextMessage() error {
-	var msgToSend *protobufs.AgentToServer
-	s.messageMutex.Lock()
-	if s.messagePending {
-		// Clone the message to have a copy for sending and avoid blocking
-		// future updates to s.nextMessage field.
-		msgToSend = proto.Clone(&s.nextMessage).(*protobufs.AgentToServer)
-		s.messagePending = false
-
-		// Reset fields that we do not have to send unless they change before the
-		// next report after this one.
-		s.nextMessage = protobufs.AgentToServer{}
-	}
-	s.messageMutex.Unlock()
-
+func (s *WSSender) sendNextMessage() error {
+	msgToSend := s.nextMessage.PopPending()
 	if msgToSend != nil && !proto.Equal(msgToSend, &protobufs.AgentToServer{}) {
 		// There is a pending message and the message has some fields populated.
 		// Set the InstanceUid field and send it.
@@ -143,7 +107,7 @@ func (s *Sender) sendNextMessage() error {
 	return nil
 }
 
-func (s *Sender) sendMessage(msg *protobufs.AgentToServer) error {
+func (s *WSSender) sendMessage(msg *protobufs.AgentToServer) error {
 	data, err := proto.Marshal(msg)
 	if err != nil {
 		s.logger.Errorf("Cannot marshal data: %v", err)
diff --git a/client/clientimpl.go b/client/wsclient.go
similarity index 75%
rename from client/clientimpl.go
rename to client/wsclient.go
index 38d91ff..12da213 100644
--- a/client/clientimpl.go
+++ b/client/wsclient.go
@@ -5,8 +5,6 @@ import (
 	"errors"
 	"net/http"
 	"net/url"
-	"strconv"
-	"strings"
 	"sync"
 	"time"
 
@@ -25,8 +23,9 @@ var (
 	errAgentDescriptionMissing = errors.New("AgentDescription is not set")
 )
 
-// client is a Client implementation.
-type client struct {
+// wsClient is an OpAMP Client implementation for WebSocket transport.
+// See specification: https://github.com/open-telemetry/opamp-spec/blob/main/specification.md#websocket-transport
+type wsClient struct {
 	logger   types.Logger
 	settings StartSettings
 
@@ -55,19 +54,17 @@ type client struct {
 	stoppedSignal chan struct{}
 
 	// The sender is responsible for sending portion of the OpAMP protocol.
-	sender *internal.Sender
+	sender *internal.WSSender
 }
 
-var _ OpAMPClient = (*client)(nil)
+var _ OpAMPClient = (*wsClient)(nil)
 
-const retryAfterHTTPHeader = "Retry-After"
-
-func New(logger types.Logger) *client {
+func NewWebSocket(logger types.Logger) *wsClient {
 	if logger == nil {
 		logger = &sharedinternal.NopLogger{}
 	}
 
-	w := &client{
+	w := &wsClient{
 		logger:        logger,
 		stoppedSignal: make(chan struct{}, 1),
 		sender:        internal.NewSender(logger),
@@ -75,7 +72,7 @@ func New(logger types.Logger) *client {
 	return w
 }
 
-func (w *client) Start(settings StartSettings) error {
+func (w *wsClient) Start(settings StartSettings) error {
 	if w.isStarted {
 		return errAlreadyStarted
 	}
@@ -113,7 +110,7 @@ func (w *client) Start(settings StartSettings) error {
 	return nil
 }
 
-func (w *client) Stop(ctx context.Context) error {
+func (w *wsClient) Stop(ctx context.Context) error {
 	if !w.isStarted {
 		return errCannotStopNotStarted
 	}
@@ -147,39 +144,32 @@ func (w *client) Stop(ctx context.Context) error {
 	return nil
 }
 
-func (w *client) SetAgentDescription(descr *protobufs.AgentDescription) error {
+func (w *wsClient) SetAgentDescription(descr *protobufs.AgentDescription) error {
 	if descr == nil {
 		return errAgentDescriptionMissing
 	}
 
 	// store the agent description to send on reconnect
 	w.settings.AgentDescription = descr
-
-	w.sender.UpdateNextStatus(func(statusReport *protobufs.StatusReport) {
+	w.sender.NextMessage().UpdateStatus(func(statusReport *protobufs.StatusReport) {
 		statusReport.AgentDescription = descr
 	})
 	w.sender.ScheduleSend()
 	return nil
 }
 
-func (w *client) SetEffectiveConfig(config *protobufs.EffectiveConfig) error {
-	w.sender.UpdateNextStatus(func(statusReport *protobufs.StatusReport) {
+func (w *wsClient) SetEffectiveConfig(config *protobufs.EffectiveConfig) error {
+	w.sender.NextMessage().UpdateStatus(func(statusReport *protobufs.StatusReport) {
 		statusReport.EffectiveConfig = config
 	})
 	w.sender.ScheduleSend()
 	return nil
 }
 
-type optionalDuration struct {
-	duration time.Duration
-	// true if duration field is defined.
-	defined bool
-}
-
 // Try to connect once. Returns an error if connection fails and optional retryAfter
 // duration to indicate to the caller to retry after the specified time as instructed
 // by the server.
-func (w *client) tryConnectOnce(ctx context.Context) (err error, retryAfter optionalDuration) {
+func (w *wsClient) tryConnectOnce(ctx context.Context) (err error, retryAfter sharedinternal.OptionalDuration) {
 	var resp *http.Response
 	conn, resp, err := w.dialer.DialContext(ctx, w.url.String(), w.requestHeader)
 	if err != nil {
@@ -188,10 +178,10 @@ func (w *client) tryConnectOnce(ctx context.Context) (err error, retryAfter opti
 		}
 		if resp != nil {
 			w.logger.Errorf("Server responded with status=%v", resp.Status)
-			duration := extractRetryAfterHeader(resp)
+			duration := sharedinternal.ExtractRetryAfterHeader(resp)
 			return err, duration
 		}
-		return err, optionalDuration{defined: false}
+		return err, sharedinternal.OptionalDuration{Defined: false}
 	}
 
 	// Successfully connected.
@@ -202,29 +192,12 @@ func (w *client) tryConnectOnce(ctx context.Context) (err error, retryAfter opti
 		w.settings.Callbacks.OnConnect()
 	}
 
-	return nil, optionalDuration{defined: false}
-}
-
-// Extract Retry-After response header if the status is 503 or 429. Returns
-// 0 duration if the header is not found or the status is different.
-func extractRetryAfterHeader(resp *http.Response) optionalDuration {
-	if resp.StatusCode == http.StatusServiceUnavailable ||
-		resp.StatusCode == http.StatusTooManyRequests {
-		retryAfter := strings.TrimSpace(resp.Header.Get(retryAfterHTTPHeader))
-		if retryAfter != "" {
-			retryIntervalSec, err := strconv.Atoi(retryAfter)
-			if err == nil {
-				retryInterval := time.Duration(retryIntervalSec) * time.Second
-				return optionalDuration{defined: true, duration: retryInterval}
-			}
-		}
-	}
-	return optionalDuration{defined: false}
+	return nil, sharedinternal.OptionalDuration{Defined: false}
 }
 
 // Continuously try until connected. Will return nil when successfully
 // connected. Will return error if it is cancelled via context.
-func (w *client) ensureConnected(ctx context.Context) error {
+func (w *wsClient) ensureConnected(ctx context.Context) error {
 	infiniteBackoff := backoff.NewExponentialBackOff()
 
 	// Make ticker run forever.
@@ -248,11 +221,11 @@ func (w *client) ensureConnected(ctx context.Context) error {
 					}
 					// Retry again a bit later.
 
-					if retryAfter.defined && retryAfter.duration > interval {
+					if retryAfter.Defined && retryAfter.Duration > interval {
 						// If the server suggested connecting later than our interval
 						// then honour server's request, otherwise wait at least
 						// as much as we calculated.
-						interval = retryAfter.duration
+						interval = retryAfter.Duration
 					}
 
 					continue
@@ -269,13 +242,13 @@ func (w *client) ensureConnected(ctx context.Context) error {
 	}
 }
 
-func (w *client) isStopping() bool {
+func (w *wsClient) isStopping() bool {
 	w.isStoppingMutex.RLock()
 	defer w.isStoppingMutex.RUnlock()
 	return w.isStoppingFlag
 }
 
-func (w *client) startConnectAndRun() {
+func (w *wsClient) startConnectAndRun() {
 	// Create a cancellable context.
 	runCtx, runCancel := context.WithCancel(context.Background())
 
@@ -297,7 +270,7 @@ func (w *client) startConnectAndRun() {
 //   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 (w *client) runOneCycle(ctx context.Context) {
+func (w *wsClient) runOneCycle(ctx context.Context) {
 	if err := w.ensureConnected(ctx); err != nil {
 		// Can't connect, so can't move forward. This currently happens when we
 		// are being stopped.
@@ -313,13 +286,11 @@ func (w *client) runOneCycle(ctx context.Context) {
 	procCtx, procCancel := context.WithCancel(ctx)
 
 	// Prepare the first status report.
-	w.sender.UpdateNextStatus(
-		func(statusReport *protobufs.StatusReport) {
-			statusReport.AgentDescription = w.settings.AgentDescription
-		},
-	)
+	w.sender.NextMessage().UpdateStatus(func(statusReport *protobufs.StatusReport) {
+		statusReport.AgentDescription = w.settings.AgentDescription
+	})
 
-	w.sender.UpdateNextMessage(
+	w.sender.NextMessage().Update(
 		func(msg *protobufs.AgentToServer) {
 			if msg.AddonStatuses == nil {
 				msg.AddonStatuses = &protobufs.AgentAddonStatuses{}
@@ -338,7 +309,7 @@ func (w *client) runOneCycle(ctx context.Context) {
 	}
 
 	// First status report sent. Now loop to receive and process messages.
-	r := internal.NewReceiver(w.logger, w.settings.Callbacks, w.conn, w.sender)
+	r := internal.NewWSReceiver(w.logger, w.settings.Callbacks, w.conn, w.sender)
 	r.ReceiverLoop(ctx)
 
 	// Stop the background processors.
@@ -347,14 +318,14 @@ func (w *client) runOneCycle(ctx context.Context) {
 	// If we exited receiverLoop it means there is a connection error, we cannot
 	// read messages anymore. We need to start over.
 
-	// Close the connection to unblock the Sender as well.
+	// Close the connection to unblock the WSSender as well.
 	w.conn.Close()
 
-	// Wait for Sender to stop.
+	// Wait for WSSender to stop.
 	w.sender.WaitToStop()
 }
 
-func (w *client) runUntilStopped(ctx context.Context) {
+func (w *wsClient) runUntilStopped(ctx context.Context) {
 
 	defer func() {
 		// We only return from runUntilStopped when we are instructed to stop.
diff --git a/internal/examples/agent/agent/agent.go b/internal/examples/agent/agent/agent.go
index fe1f64c..c706c4e 100644
--- a/internal/examples/agent/agent/agent.go
+++ b/internal/examples/agent/agent/agent.go
@@ -81,7 +81,7 @@ func NewAgent(logger types.Logger, agentType string, agentVersion string) *Agent
 }
 
 func (agent *Agent) start() error {
-	agent.opampClient = client.New(agent.logger)
+	agent.opampClient = client.NewWebSocket(agent.logger)
 
 	settings := client.StartSettings{
 		OpAMPServerURL:   "ws://127.0.0.1:4320/v1/opamp",
diff --git a/internal/retryafter.go b/internal/retryafter.go
new file mode 100644
index 0000000..70eb831
--- /dev/null
+++ b/internal/retryafter.go
@@ -0,0 +1,34 @@
+package internal
+
+import (
+	"net/http"
+	"strconv"
+	"strings"
+	"time"
+)
+
+const retryAfterHTTPHeader = "Retry-After"
+
+type OptionalDuration struct {
+	Duration time.Duration
+	// true if duration field is defined.
+	Defined bool
+}
+
+// ExtractRetryAfterHeader extracts Retry-After response header if the status
+// is 503 or 429. Returns 0 duration if the header is not found or the status
+// is different.
+func ExtractRetryAfterHeader(resp *http.Response) OptionalDuration {
+	if resp.StatusCode == http.StatusServiceUnavailable ||
+		resp.StatusCode == http.StatusTooManyRequests {
+		retryAfter := strings.TrimSpace(resp.Header.Get(retryAfterHTTPHeader))
+		if retryAfter != "" {
+			retryIntervalSec, err := strconv.Atoi(retryAfter)
+			if err == nil {
+				retryInterval := time.Duration(retryIntervalSec) * time.Second
+				return OptionalDuration{Defined: true, Duration: retryInterval}
+			}
+		}
+	}
+	return OptionalDuration{Defined: false}
+}


Test patch:
diff --git a/client/clientimpl_test.go b/client/clientimpl_test.go
index 3f93b42..05e19db 100644
--- a/client/clientimpl_test.go
+++ b/client/clientimpl_test.go
@@ -4,14 +4,15 @@ import (
 	"context"
 	"math/rand"
 	"net/http"
+	"net/url"
 	"sync/atomic"
 	"testing"
 	"time"
 
-	"github.com/gorilla/websocket"
 	ulid "github.com/oklog/ulid/v2"
 	"github.com/open-telemetry/opamp-go/protobufshelpers"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/protobuf/proto"
 
 	"github.com/open-telemetry/opamp-go/client/internal"
@@ -20,286 +21,296 @@ import (
 	"github.com/open-telemetry/opamp-go/protobufs"
 )
 
-func TestConnectInvalidURL(t *testing.T) {
-	settings := StartSettings{
-		OpAMPServerURL: ":not a url",
+const retryAfterHTTPHeader = "Retry-After"
+
+func testClients(t *testing.T, f func(client OpAMPClient)) {
+	// Run the test defined by f() for WebSocket and HTTP clients.
+	tests := []struct {
+		name   string
+		client OpAMPClient
+	}{
+		{
+			name:   "http",
+			client: NewHTTP(nil),
+		},
+		{
+			name:   "ws",
+			client: NewWebSocket(nil),
+		},
 	}
 
-	client := New(nil)
-	err := client.Start(settings)
-	assert.Error(t, err)
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			f(test.client)
+		})
+	}
+}
+
+func TestConnectInvalidURL(t *testing.T) {
+	testClients(t, func(client OpAMPClient) {
+		settings := StartSettings{
+			OpAMPServerURL: ":not a url",
+		}
+
+		err := client.Start(settings)
+		assert.Error(t, err)
+	})
 }
 
 func eventually(t *testing.T, f func() bool) {
 	assert.Eventually(t, f, 5*time.Second, 10*time.Millisecond)
 }
 
-func prepareClient(settings StartSettings) *client {
+func prepareClient(t *testing.T, settings *StartSettings, c OpAMPClient) {
 	// Autogenerate instance id.
 	entropy := ulid.Monotonic(rand.New(rand.NewSource(99)), 0)
 	settings.InstanceUid = ulid.MustNew(ulid.Timestamp(time.Now()), entropy).String()
 
-	return New(nil)
+	// Make sure correct URL scheme is used, based on the type of the OpAMP client.
+	u, err := url.Parse(settings.OpAMPServerURL)
+	require.NoError(t, err)
+	switch c.(type) {
+	case *httpClient:
+		u.Scheme = "http"
+	case *wsClient:
+		u.Scheme = "ws"
+	}
+	settings.OpAMPServerURL = u.String()
+
+	settings.AgentDescription = &protobufs.AgentDescription{}
 }
 
-func startClient(t *testing.T, settings StartSettings) *client {
-	client := prepareClient(settings)
+func startClient(t *testing.T, settings StartSettings, client OpAMPClient) {
+	prepareClient(t, &settings, client)
 	err := client.Start(settings)
 	assert.NoError(t, err)
-	return client
 }
 
 // Create start settings that point to a non-existing server.
 func createNoServerSettings() StartSettings {
 	return StartSettings{
-		OpAMPServerURL:   "ws://" + testhelpers.GetAvailableLocalAddress(),
-		AgentDescription: &protobufs.AgentDescription{},
+		OpAMPServerURL: "ws://" + testhelpers.GetAvailableLocalAddress(),
 	}
 }
 
 func TestConnectNoServer(t *testing.T) {
-	client := startClient(t, createNoServerSettings())
-
-	err := client.Stop(context.Background())
-	assert.NoError(t, err)
+	testClients(t, func(client OpAMPClient) {
+		startClient(t, createNoServerSettings(), client)
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
+	})
 }
 
 func TestOnConnectFail(t *testing.T) {
-	var connectErr atomic.Value
-	settings := createNoServerSettings()
-	settings.Callbacks = types.CallbacksStruct{
-		OnConnectFailedFunc: func(err error) {
-			connectErr.Store(err)
-		},
-	}
+	testClients(t, func(client OpAMPClient) {
+		var connectErr atomic.Value
+		settings := createNoServerSettings()
+		settings.Callbacks = types.CallbacksStruct{
+			OnConnectFailedFunc: func(err error) {
+				connectErr.Store(err)
+			},
+		}
 
-	client := startClient(t, settings)
+		startClient(t, settings, client)
 
-	eventually(t, func() bool { return connectErr.Load() != nil })
+		eventually(t, func() bool { return connectErr.Load() != nil })
 
-	err := client.Stop(context.Background())
-	assert.NoError(t, err)
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
+	})
 }
 
 func TestStartStarted(t *testing.T) {
-	settings := createNoServerSettings()
-	client := startClient(t, settings)
+	testClients(t, func(client OpAMPClient) {
+		settings := createNoServerSettings()
+		startClient(t, settings, client)
 
-	// Try to start again.
-	err := client.Start(settings)
-	assert.Error(t, err)
+		// Try to start again.
+		err := client.Start(settings)
+		assert.Error(t, err)
 
-	err = client.Stop(context.Background())
-	assert.NoError(t, err)
+		err = client.Stop(context.Background())
+		assert.NoError(t, err)
+	})
 }
 
 func TestStopWithoutStart(t *testing.T) {
-	client := New(nil)
-	err := client.Stop(context.Background())
-	assert.Error(t, err)
+	testClients(t, func(client OpAMPClient) {
+		err := client.Stop(context.Background())
+		assert.Error(t, err)
+	})
 }
 
 func TestStopCancellation(t *testing.T) {
-	client := startClient(t, createNoServerSettings())
-
-	ctx, cancel := context.WithCancel(context.Background())
-	cancel()
-	err := client.Stop(ctx)
-	if err != nil {
-		assert.ErrorIs(t, err, context.Canceled)
-	}
+	testClients(t, func(client OpAMPClient) {
+		startClient(t, createNoServerSettings(), client)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		err := client.Stop(ctx)
+		if err != nil {
+			assert.ErrorIs(t, err, context.Canceled)
+		}
+	})
 }
 
 func TestConnectWithServer(t *testing.T) {
-	// Start a server.
-	srv := internal.StartMockServer(t)
-
-	// Start a client.
-	var connected int64
-	settings := StartSettings{
-		Callbacks: types.CallbacksStruct{
-			OnConnectFunc: func() {
-				atomic.StoreInt64(&connected, 1)
+	testClients(t, func(client OpAMPClient) {
+		// Start a server.
+		srv := internal.StartMockServer(t)
+
+		// Start a client.
+		var connected int64
+		settings := StartSettings{
+			Callbacks: types.CallbacksStruct{
+				OnConnectFunc: func() {
+					atomic.StoreInt64(&connected, 1)
+				},
 			},
-		},
-		AgentDescription: &protobufs.AgentDescription{},
-	}
-	settings.OpAMPServerURL = "ws://" + srv.Endpoint
-	client := startClient(t, settings)
+		}
+		settings.OpAMPServerURL = "ws://" + srv.Endpoint
+		startClient(t, settings, client)
 
-	// Wait for connection to be established.
-	eventually(t, func() bool { return atomic.LoadInt64(&connected) != 0 })
+		// Wait for connection to be established.
+		eventually(t, func() bool { return atomic.LoadInt64(&connected) != 0 })
 
-	// Shutdown the server.
-	srv.Close()
+		// Shutdown the server.
+		srv.Close()
 
-	// Shutdown the client.
-	err := client.Stop(context.Background())
-	assert.NoError(t, err)
+		// Shutdown the client.
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
+	})
 }
 
 func TestConnectWithServer503(t *testing.T) {
-	// Start a server.
-	var connectionAttempts int64
-	srv := internal.StartMockServer(t)
-	srv.OnRequest = func(w http.ResponseWriter, r *http.Request) {
-		atomic.StoreInt64(&connectionAttempts, 1)
-
-		// Always respond with an error to the client.
-		w.Header().Set(retryAfterHTTPHeader, "30")
-		w.WriteHeader(http.StatusServiceUnavailable)
-	}
+	testClients(t, func(client OpAMPClient) {
+		// Start a server.
+		var connectionAttempts int64
+		srv := internal.StartMockServer(t)
+		srv.OnRequest = func(w http.ResponseWriter, r *http.Request) {
+			atomic.StoreInt64(&connectionAttempts, 1)
+
+			// Always respond with an error to the client.
+			w.Header().Set(retryAfterHTTPHeader, "30")
+			w.WriteHeader(http.StatusServiceUnavailable)
+		}
 
-	// Start a client.
-	var clientConnected int64
-	var connectErr atomic.Value
-	settings := StartSettings{
-		Callbacks: types.CallbacksStruct{
-			OnConnectFunc: func() {
-				atomic.StoreInt64(&clientConnected, 1)
-				assert.Fail(t, "Client should not be able to connect")
-			},
-			OnConnectFailedFunc: func(err error) {
-				connectErr.Store(err)
+		// Start a client.
+		var clientConnected int64
+		var connectErr atomic.Value
+		settings := StartSettings{
+			Callbacks: types.CallbacksStruct{
+				OnConnectFunc: func() {
+					atomic.StoreInt64(&clientConnected, 1)
+					assert.Fail(t, "Client should not be able to connect")
+				},
+				OnConnectFailedFunc: func(err error) {
+					connectErr.Store(err)
+				},
 			},
-		},
-		AgentDescription: &protobufs.AgentDescription{},
-	}
-	settings.OpAMPServerURL = "ws://" + srv.Endpoint
-	client := startClient(t, settings)
+		}
+		settings.OpAMPServerURL = "ws://" + srv.Endpoint
+		startClient(t, settings, client)
 
-	// Wait for connection to fail.
-	eventually(t, func() bool { return connectErr.Load() != nil })
+		// Wait for connection to fail.
+		eventually(t, func() bool { return connectErr.Load() != nil })
 
-	assert.EqualValues(t, 1, atomic.LoadInt64(&connectionAttempts))
-	assert.EqualValues(t, 0, atomic.LoadInt64(&clientConnected))
+		assert.EqualValues(t, 1, atomic.LoadInt64(&connectionAttempts))
+		assert.EqualValues(t, 0, atomic.LoadInt64(&clientConnected))
 
-	// Shutdown the server.
-	srv.Close()
-	client.Stop(context.Background())
+		// Shutdown the server.
+		srv.Close()
+		client.Stop(context.Background())
+	})
 }
 
 func TestConnectWithHeader(t *testing.T) {
-	// Start a server.
-	srv := internal.StartMockServer(t)
-	var conn atomic.Value
-	srv.OnConnect = func(r *http.Request, c *websocket.Conn) {
-		authHdr := r.Header.Get("Authorization")
-		assert.EqualValues(t, "Bearer 12345678", authHdr)
-		conn.Store(c)
-	}
+	testClients(t, func(client OpAMPClient) {
+		// Start a server.
+		srv := internal.StartMockServer(t)
+		var conn atomic.Value
+		srv.OnConnect = func(r *http.Request) {
+			authHdr := r.Header.Get("Authorization")
+			assert.EqualValues(t, "Bearer 12345678", authHdr)
+			conn.Store(true)
+		}
 
-	// Start a client.
-	settings := StartSettings{
-		OpAMPServerURL:      "ws://" + srv.Endpoint,
-		AuthorizationHeader: "Bearer 12345678",
-		AgentDescription:    &protobufs.AgentDescription{},
-	}
-	client := startClient(t, settings)
+		// Start a client.
+		settings := StartSettings{
+			OpAMPServerURL:      "ws://" + srv.Endpoint,
+			AuthorizationHeader: "Bearer 12345678",
+		}
+		startClient(t, settings, client)
 
-	// Wait for connection to be established.
-	eventually(t, func() bool { return conn.Load() != nil })
+		// Wait for connection to be established.
+		eventually(t, func() bool { return conn.Load() != nil })
 
-	// Shutdown the server and the client.
-	srv.Close()
-	client.Stop(context.Background())
+		// Shutdown the server and the client.
+		srv.Close()
+		client.Stop(context.Background())
+	})
 }
 
-func TestDisconnectByServer(t *testing.T) {
-	// Start a server.
-	srv := internal.StartMockServer(t)
-	var conn atomic.Value
-	srv.OnConnect = func(r *http.Request, c *websocket.Conn) {
-		conn.Store(c)
-	}
+func TestFirstStatusReport(t *testing.T) {
+	testClients(t, func(client OpAMPClient) {
 
-	// Start a client.
-	var connected int64
-	var connectErr atomic.Value
-	settings := StartSettings{
-		Callbacks: types.CallbacksStruct{
-			OnConnectFunc: func() {
-				atomic.StoreInt64(&connected, 1)
+		remoteConfig := &protobufs.AgentRemoteConfig{
+			Config: &protobufs.AgentConfigMap{
+				ConfigMap: map[string]*protobufs.AgentConfigFile{},
 			},
-			OnConnectFailedFunc: func(err error) {
-				connectErr.Store(err)
-			},
-		},
-		AgentDescription: &protobufs.AgentDescription{},
-	}
-	settings.OpAMPServerURL = "ws://" + srv.Endpoint
-	client := startClient(t, settings)
-
-	// Wait for connection to be established.
-	eventually(t, func() bool { return conn.Load() != nil })
-	assert.True(t, connectErr.Load() == nil)
-
-	// Close the server and forcefully disconnect.
-	srv.Close()
-	conn.Load().(*websocket.Conn).Close()
-
-	// The client must retry and must fail now.
-	eventually(t, func() bool { return connectErr.Load() != nil })
-
-	// Stop the client.
-	err := client.Stop(context.Background())
-	assert.NoError(t, err)
-}
-
-func TestFirstStatusReport(t *testing.T) {
-	remoteConfig := &protobufs.AgentRemoteConfig{
-		Config: &protobufs.AgentConfigMap{
-			ConfigMap: map[string]*protobufs.AgentConfigFile{},
-		},
-		ConfigHash: []byte{1, 2, 3, 4},
-	}
+			ConfigHash: []byte{1, 2, 3, 4},
+		}
 
-	// Start a server.
-	srv := internal.StartMockServer(t)
-	srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
-		return &protobufs.ServerToAgent{
-			InstanceUid:  msg.InstanceUid,
-			RemoteConfig: remoteConfig,
+		// Start a server.
+		srv := internal.StartMockServer(t)
+		srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
+			return &protobufs.ServerToAgent{
+				InstanceUid:  msg.InstanceUid,
+				RemoteConfig: remoteConfig,
+			}
 		}
-	}
 
-	// Start a client.
-	var connected, remoteConfigReceived int64
-	settings := StartSettings{
-		Callbacks: types.CallbacksStruct{
-			OnConnectFunc: func() {
-				atomic.AddInt64(&connected, 1)
-			},
-			OnRemoteConfigFunc: func(
-				ctx context.Context,
-				config *protobufs.AgentRemoteConfig,
-			) (effectiveConfig *protobufs.EffectiveConfig, configChanged bool, err error) {
-				// Verify that the client received exactly the remote config that
-				// the server sent.
-				assert.True(t, proto.Equal(remoteConfig, config))
-				atomic.AddInt64(&remoteConfigReceived, 1)
-				return &protobufs.EffectiveConfig{
-					ConfigMap: remoteConfig.Config,
-				}, true, nil
+		// Start a client.
+		var connected, remoteConfigReceived int64
+		settings := StartSettings{
+			Callbacks: types.CallbacksStruct{
+				OnConnectFunc: func() {
+					atomic.AddInt64(&connected, 1)
+				},
+				OnRemoteConfigFunc: func(
+					ctx context.Context,
+					config *protobufs.AgentRemoteConfig,
+				) (
+					effectiveConfig *protobufs.EffectiveConfig, configChanged bool,
+					err error,
+				) {
+					// Verify that the client received exactly the remote config that
+					// the server sent.
+					assert.True(t, proto.Equal(remoteConfig, config))
+					atomic.AddInt64(&remoteConfigReceived, 1)
+					return &protobufs.EffectiveConfig{
+						ConfigMap: remoteConfig.Config,
+					}, true, nil
+				},
 			},
-		},
-		AgentDescription: &protobufs.AgentDescription{},
-	}
-	settings.OpAMPServerURL = "ws://" + srv.Endpoint
-	client := startClient(t, settings)
+		}
+		settings.OpAMPServerURL = "ws://" + srv.Endpoint
+		startClient(t, settings, client)
 
-	// Wait for connection to be established.
-	eventually(t, func() bool { return atomic.LoadInt64(&connected) != 0 })
+		// Wait for connection to be established.
+		eventually(t, func() bool { return atomic.LoadInt64(&connected) != 0 })
 
-	// Wait to receive remote config.
-	eventually(t, func() bool { return atomic.LoadInt64(&remoteConfigReceived) != 0 })
+		// Wait to receive remote config.
+		eventually(t, func() bool { return atomic.LoadInt64(&remoteConfigReceived) != 0 })
 
-	// Shutdown the server.
-	srv.Close()
+		// Shutdown the server.
+		srv.Close()
 
-	// Shutdown the client.
-	err := client.Stop(context.Background())
-	assert.NoError(t, err)
+		// Shutdown the client.
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
+	})
 }
 
 func TestIncludesDetailsOnReconnect(t *testing.T) {
@@ -334,7 +345,8 @@ func TestIncludesDetailsOnReconnect(t *testing.T) {
 	}
 
 	settings.OpAMPServerURL = "ws://" + srv.Endpoint
-	client := startClient(t, settings)
+	client := NewWebSocket(nil)
+	startClient(t, settings, client)
 
 	eventually(t, func() bool { return atomic.LoadInt64(&connected) == 1 })
 	eventually(t, func() bool { return atomic.LoadInt64(&receivedDetails) == 1 })
@@ -351,53 +363,67 @@ func TestIncludesDetailsOnReconnect(t *testing.T) {
 }
 
 func TestSetEffectiveConfig(t *testing.T) {
-	// Start a server.
-	srv := internal.StartMockServer(t)
-	var rcvConfig atomic.Value
-	srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
-		if statusReport := msg.GetStatusReport(); statusReport != nil {
-			if statusReport.EffectiveConfig != nil {
-				rcvConfig.Store(statusReport.EffectiveConfig)
+	testClients(t, func(client OpAMPClient) {
+		// Start a server.
+		srv := internal.StartMockServer(t)
+		var rcvConfig atomic.Value
+		srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
+			if statusReport := msg.GetStatusReport(); statusReport != nil {
+				if statusReport.EffectiveConfig != nil {
+					rcvConfig.Store(statusReport.EffectiveConfig)
+				}
 			}
+			return nil
 		}
-		return nil
-	}
 
-	// Start a client.
-	settings := StartSettings{}
-	settings.OpAMPServerURL = "ws://" + srv.Endpoint
-	settings.AgentDescription = &protobufs.AgentDescription{}
-	client := prepareClient(settings)
+		// Start a client.
+		settings := StartSettings{}
+		settings.OpAMPServerURL = "ws://" + srv.Endpoint
+		prepareClient(t, &settings, client)
 
-	sendConfig := &protobufs.EffectiveConfig{
-		ConfigMap: &protobufs.AgentConfigMap{
-			ConfigMap: map[string]*protobufs.AgentConfigFile{
-				"key": {},
+		sendConfig := &protobufs.EffectiveConfig{
+			ConfigMap: &protobufs.AgentConfigMap{
+				ConfigMap: map[string]*protobufs.AgentConfigFile{
+					"key": {},
+				},
 			},
-		},
-	}
-
-	// Set before start. It should be delivered immediately after Start().
-	client.SetEffectiveConfig(sendConfig)
+		}
 
-	assert.NoError(t, client.Start(settings))
+		// Set before start. This should be delivered immediately after Start().
+		client.SetEffectiveConfig(sendConfig)
 
-	// Verify it is delivered.
-	eventually(t, func() bool { return proto.Equal(sendConfig, rcvConfig.Load().(*protobufs.EffectiveConfig)) })
+		assert.NoError(t, client.Start(settings))
 
-	// Now change again.
-	sendConfig.ConfigMap.ConfigMap["key2"] = &protobufs.AgentConfigFile{}
-	client.SetEffectiveConfig(sendConfig)
+		// Verify it is delivered.
+		eventually(
+			t,
+			func() bool {
+				return rcvConfig.Load() != nil &&
+					proto.Equal(sendConfig, rcvConfig.Load().(*protobufs.EffectiveConfig))
+			},
+		)
+
+		// Now change again.
+		sendConfig.ConfigMap.ConfigMap["key2"] = &protobufs.AgentConfigFile{}
+		client.SetEffectiveConfig(sendConfig)
+
+		// Verify change is delivered.
+		eventually(
+			t,
+			func() bool {
+				return rcvConfig.Load() != nil &&
+					proto.Equal(sendConfig, rcvConfig.Load().(*protobufs.EffectiveConfig))
+			},
+		)
 
-	// Verify change is delivered.
-	eventually(t, func() bool { return proto.Equal(sendConfig, rcvConfig.Load().(*protobufs.EffectiveConfig)) })
+		// Shutdown the server.
+		srv.Close()
 
-	// Shutdown the server.
-	srv.Close()
+		// Shutdown the client.
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
 
-	// Shutdown the client.
-	err := client.Stop(context.Background())
-	assert.NoError(t, err)
+	})
 }
 
 func isEqualAgentAttrs(sent []*protobufs.KeyValue, received []*protobufs.KeyValue) bool {
@@ -416,222 +442,246 @@ func isEqualAgentAttrs(sent []*protobufs.KeyValue, received []*protobufs.KeyValu
 }
 
 func TestSetAgentAttributes(t *testing.T) {
-	// Start a server.
-	srv := internal.StartMockServer(t)
-	var rcvAgentDescr atomic.Value
-	srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
-		if statusReport := msg.GetStatusReport(); statusReport != nil {
-			if statusReport.AgentDescription != nil {
-				rcvAgentDescr.Store(statusReport.AgentDescription)
+	testClients(t, func(client OpAMPClient) {
+
+		// Start a server.
+		srv := internal.StartMockServer(t)
+		var rcvAgentDescr atomic.Value
+		srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
+			if statusReport := msg.GetStatusReport(); statusReport != nil {
+				if statusReport.AgentDescription != nil {
+					rcvAgentDescr.Store(statusReport.AgentDescription)
+				}
 			}
+			return nil
 		}
-		return nil
-	}
 
-	// Start a client.
-	settings := StartSettings{
-		OpAMPServerURL:   "ws://" + srv.Endpoint,
-		AgentDescription: &protobufs.AgentDescription{},
-	}
-	client := prepareClient(settings)
+		// Start a client.
+		settings := StartSettings{
+			OpAMPServerURL: "ws://" + srv.Endpoint,
+		}
+		prepareClient(t, &settings, client)
 
-	settings.AgentDescription = &protobufs.AgentDescription{
-		IdentifyingAttributes: []*protobufs.KeyValue{
-			{
-				Key:   "host.name",
-				Value: &protobufs.AnyValue{Value: &protobufs.AnyValue_StringValue{StringValue: "somehost"}},
+		agentDescr := &protobufs.AgentDescription{
+			IdentifyingAttributes: []*protobufs.KeyValue{
+				{
+					Key:   "host.name",
+					Value: &protobufs.AnyValue{Value: &protobufs.AnyValue_StringValue{StringValue: "somehost"}},
+				},
 			},
-		},
-	}
-
-	assert.NoError(t, client.Start(settings))
-
-	// Verify it is delivered.
-	eventually(t, func() bool {
-		agentDescr, ok := rcvAgentDescr.Load().(*protobufs.AgentDescription)
-		if !ok || agentDescr == nil {
-			return false
 		}
-		return isEqualAgentAttrs(settings.AgentDescription.IdentifyingAttributes, agentDescr.IdentifyingAttributes)
-	})
+		client.SetAgentDescription(agentDescr)
 
-	// Now change again.
-	settings.AgentDescription.NonIdentifyingAttributes = []*protobufs.KeyValue{
-		{
-			Key:   "os.name",
-			Value: &protobufs.AnyValue{Value: &protobufs.AnyValue_StringValue{StringValue: "linux"}},
-		},
-	}
+		assert.NoError(t, client.Start(settings))
 
-	client.SetAgentDescription(settings.AgentDescription)
+		// Verify it is delivered.
+		eventually(
+			t,
+			func() bool {
+				agentDescr, ok := rcvAgentDescr.Load().(*protobufs.AgentDescription)
+				if !ok || agentDescr == nil {
+					return false
+				}
+				return isEqualAgentAttrs(
+					agentDescr.IdentifyingAttributes, agentDescr.IdentifyingAttributes,
+				)
+			},
+		)
 
-	// Verify change is delivered.
-	eventually(t, func() bool {
-		agentDescr := rcvAgentDescr.Load().(*protobufs.AgentDescription)
-		if agentDescr == nil {
-			return false
+		// Now change again.
+		agentDescr.NonIdentifyingAttributes = []*protobufs.KeyValue{
+			{
+				Key:   "os.name",
+				Value: &protobufs.AnyValue{Value: &protobufs.AnyValue_StringValue{StringValue: "linux"}},
+			},
 		}
-		return isEqualAgentAttrs(settings.AgentDescription.IdentifyingAttributes, agentDescr.IdentifyingAttributes) &&
-			isEqualAgentAttrs(settings.AgentDescription.NonIdentifyingAttributes, agentDescr.NonIdentifyingAttributes)
-	})
-
-	// Shutdown the server.
-	srv.Close()
 
-	// Shutdown the client.
-	err := client.Stop(context.Background())
-	assert.NoError(t, err)
-}
+		client.SetAgentDescription(agentDescr)
 
-func TestAgentIdentification(t *testing.T) {
-	// Start a server.
-	srv := internal.StartMockServer(t)
-	newInstanceUid := ulid.MustNew(ulid.Timestamp(time.Now()), ulid.Monotonic(rand.New(rand.NewSource(0)), 0))
-	var rcvAgentInstanceUid atomic.Value
-	srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
-		rcvAgentInstanceUid.Store(msg.InstanceUid)
-		return &protobufs.ServerToAgent{
-			InstanceUid: msg.InstanceUid,
-			AgentIdentification: &protobufs.AgentIdentification{
-				NewInstanceUid: newInstanceUid.String(),
+		// Verify change is delivered.
+		eventually(
+			t,
+			func() bool {
+				agentDescr := rcvAgentDescr.Load().(*protobufs.AgentDescription)
+				if agentDescr == nil {
+					return false
+				}
+				return isEqualAgentAttrs(agentDescr.IdentifyingAttributes, agentDescr.IdentifyingAttributes) &&
+					isEqualAgentAttrs(agentDescr.NonIdentifyingAttributes, agentDescr.NonIdentifyingAttributes)
 			},
-		}
-	}
+		)
 
-	// Start a client.
-	settings := StartSettings{}
-	settings.OpAMPServerURL = "ws://" + srv.Endpoint
-	settings.AgentDescription = &protobufs.AgentDescription{}
-	settings.Callbacks = types.CallbacksStruct{
-		OnAgentIdentificationFunc: func(
-			ctx context.Context,
-			agentId *protobufs.AgentIdentification,
-		) error {
-			return nil
-		},
-	}
-	client := prepareClient(settings)
+		// Shutdown the server.
+		srv.Close()
 
-	oldInstanceUid := settings.InstanceUid
-	assert.NoError(t, client.Start(settings))
+		// Shutdown the client.
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
 
-	// First, server gets the original instanceId
-	eventually(t, func() bool {
-		instanceUid, ok := rcvAgentInstanceUid.Load().(string)
-		if !ok {
-			return false
-		}
-		return instanceUid == oldInstanceUid
 	})
+}
 
-	// Send a dummy message
-	client.SetAgentDescription(&protobufs.AgentDescription{})
+func TestAgentIdentification(t *testing.T) {
+	testClients(t, func(client OpAMPClient) {
+		// Start a server.
+		srv := internal.StartMockServer(t)
+		newInstanceUid := ulid.MustNew(
+			ulid.Timestamp(time.Now()), ulid.Monotonic(rand.New(rand.NewSource(0)), 0),
+		)
+		var rcvAgentInstanceUid atomic.Value
+		srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
+			rcvAgentInstanceUid.Store(msg.InstanceUid)
+			return &protobufs.ServerToAgent{
+				InstanceUid: msg.InstanceUid,
+				AgentIdentification: &protobufs.AgentIdentification{
+					NewInstanceUid: newInstanceUid.String(),
+				},
+			}
+		}
 
-	// When it was sent, the new instance uid should have been used, which should
-	// have been observed by the server
-	eventually(t, func() bool {
-		instanceUid, ok := rcvAgentInstanceUid.Load().(string)
-		if !ok {
-			return false
+		// Start a client.
+		settings := StartSettings{}
+		settings.OpAMPServerURL = "ws://" + srv.Endpoint
+		settings.AgentDescription = &protobufs.AgentDescription{}
+		settings.Callbacks = types.CallbacksStruct{
+			OnAgentIdentificationFunc: func(
+				ctx context.Context,
+				agentId *protobufs.AgentIdentification,
+			) error {
+				return nil
+			},
 		}
-		return instanceUid == newInstanceUid.String()
-	})
+		prepareClient(t, &settings, client)
+
+		oldInstanceUid := settings.InstanceUid
+		assert.NoError(t, client.Start(settings))
+
+		// First, server gets the original instanceId
+		eventually(
+			t,
+			func() bool {
+				instanceUid, ok := rcvAgentInstanceUid.Load().(string)
+				if !ok {
+					return false
+				}
+				return instanceUid == oldInstanceUid
+			},
+		)
+
+		// Send a dummy message
+		client.SetAgentDescription(&protobufs.AgentDescription{})
+
+		// When it was sent, the new instance uid should have been used, which should
+		// have been observed by the server
+		eventually(
+			t,
+			func() bool {
+				instanceUid, ok := rcvAgentInstanceUid.Load().(string)
+				if !ok {
+					return false
+				}
+				return instanceUid == newInstanceUid.String()
+			},
+		)
 
-	// Shutdown the server.
-	srv.Close()
+		// Shutdown the server.
+		srv.Close()
 
-	// Shutdown the client.
-	err := client.Stop(context.Background())
-	assert.NoError(t, err)
+		// Shutdown the client.
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
+	})
 }
 
 func TestConnectionSettings(t *testing.T) {
-	hash := []byte{1, 2, 3}
-	opampSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://opamp.com"}
-	metricsSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://metrics.com"}
-	tracesSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://traces.com"}
-	logsSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://logs.com"}
-	otherSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://other.com"}
-
-	var rcvStatus int64
-	// Start a server.
-	srv := internal.StartMockServer(t)
-	srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
-		if statusReport := msg.GetStatusReport(); statusReport != nil {
-			atomic.AddInt64(&rcvStatus, 1)
-
-			return &protobufs.ServerToAgent{
-				ConnectionSettings: &protobufs.ConnectionSettingsOffers{
-					Hash:       hash,
-					Opamp:      opampSettings,
-					OwnMetrics: metricsSettings,
-					OwnTraces:  tracesSettings,
-					OwnLogs:    logsSettings,
-					OtherConnections: map[string]*protobufs.ConnectionSettings{
-						"other": otherSettings,
+	testClients(t, func(client OpAMPClient) {
+		hash := []byte{1, 2, 3}
+		opampSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://opamp.com"}
+		metricsSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://metrics.com"}
+		tracesSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://traces.com"}
+		logsSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://logs.com"}
+		otherSettings := &protobufs.ConnectionSettings{DestinationEndpoint: "http://other.com"}
+
+		var rcvStatus int64
+		// Start a server.
+		srv := internal.StartMockServer(t)
+		srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
+			if statusReport := msg.GetStatusReport(); statusReport != nil {
+				atomic.AddInt64(&rcvStatus, 1)
+
+				return &protobufs.ServerToAgent{
+					ConnectionSettings: &protobufs.ConnectionSettingsOffers{
+						Hash:       hash,
+						Opamp:      opampSettings,
+						OwnMetrics: metricsSettings,
+						OwnTraces:  tracesSettings,
+						OwnLogs:    logsSettings,
+						OtherConnections: map[string]*protobufs.ConnectionSettings{
+							"other": otherSettings,
+						},
 					},
-				},
+				}
 			}
+			return nil
 		}
-		return nil
-	}
-
-	var gotOpampSettings int64
-	var gotOwnSettings int64
-	var gotOtherSettings int64
 
-	// Start a client.
-	settings := StartSettings{
-		Callbacks: types.CallbacksStruct{
-			OnOpampConnectionSettingsFunc: func(
-				ctx context.Context, settings *protobufs.ConnectionSettings,
-			) error {
-				assert.True(t, proto.Equal(opampSettings, settings))
-				atomic.AddInt64(&gotOpampSettings, 1)
-				return nil
-			},
+		var gotOpampSettings int64
+		var gotOwnSettings int64
+		var gotOtherSettings int64
+
+		// Start a client.
+		settings := StartSettings{
+			Callbacks: types.CallbacksStruct{
+				OnOpampConnectionSettingsFunc: func(
+					ctx context.Context, settings *protobufs.ConnectionSettings,
+				) error {
+					assert.True(t, proto.Equal(opampSettings, settings))
+					atomic.AddInt64(&gotOpampSettings, 1)
+					return nil
+				},
 
-			OnOwnTelemetryConnectionSettingsFunc: func(
-				ctx context.Context, telemetryType types.OwnTelemetryType,
-				settings *protobufs.ConnectionSettings,
-			) error {
-				switch telemetryType {
-				case types.OwnMetrics:
-					assert.True(t, proto.Equal(metricsSettings, settings))
-				case types.OwnTraces:
-					assert.True(t, proto.Equal(tracesSettings, settings))
-				case types.OwnLogs:
-					assert.True(t, proto.Equal(logsSettings, settings))
-				}
-				atomic.AddInt64(&gotOwnSettings, 1)
-				return nil
-			},
+				OnOwnTelemetryConnectionSettingsFunc: func(
+					ctx context.Context, telemetryType types.OwnTelemetryType,
+					settings *protobufs.ConnectionSettings,
+				) error {
+					switch telemetryType {
+					case types.OwnMetrics:
+						assert.True(t, proto.Equal(metricsSettings, settings))
+					case types.OwnTraces:
+						assert.True(t, proto.Equal(tracesSettings, settings))
+					case types.OwnLogs:
+						assert.True(t, proto.Equal(logsSettings, settings))
+					}
+					atomic.AddInt64(&gotOwnSettings, 1)
+					return nil
+				},
 
-			OnOtherConnectionSettingsFunc: func(
-				ctx context.Context, name string, settings *protobufs.ConnectionSettings,
-			) error {
-				assert.EqualValues(t, "other", name)
-				assert.True(t, proto.Equal(otherSettings, settings))
-				atomic.AddInt64(&gotOtherSettings, 1)
-				return nil
+				OnOtherConnectionSettingsFunc: func(
+					ctx context.Context, name string,
+					settings *protobufs.ConnectionSettings,
+				) error {
+					assert.EqualValues(t, "other", name)
+					assert.True(t, proto.Equal(otherSettings, settings))
+					atomic.AddInt64(&gotOtherSettings, 1)
+					return nil
+				},
 			},
-		},
-		AgentDescription: &protobufs.AgentDescription{},
-	}
-	settings.OpAMPServerURL = "ws://" + srv.Endpoint
-	client := prepareClient(settings)
+		}
+		settings.OpAMPServerURL = "ws://" + srv.Endpoint
+		prepareClient(t, &settings, client)
 
-	assert.NoError(t, client.Start(settings))
+		assert.NoError(t, client.Start(settings))
 
-	eventually(t, func() bool { return atomic.LoadInt64(&gotOpampSettings) == 1 })
-	eventually(t, func() bool { return atomic.LoadInt64(&gotOwnSettings) == 3 })
-	eventually(t, func() bool { return atomic.LoadInt64(&gotOtherSettings) == 1 })
-	eventually(t, func() bool { return atomic.LoadInt64(&rcvStatus) == 1 })
+		eventually(t, func() bool { return atomic.LoadInt64(&gotOpampSettings) == 1 })
+		eventually(t, func() bool { return atomic.LoadInt64(&gotOwnSettings) == 3 })
+		eventually(t, func() bool { return atomic.LoadInt64(&gotOtherSettings) == 1 })
+		eventually(t, func() bool { return atomic.LoadInt64(&rcvStatus) == 1 })
 
-	// Shutdown the server.
-	srv.Close()
+		// Shutdown the server.
+		srv.Close()
 
-	// Shutdown the client.
-	err := client.Stop(context.Background())
-	assert.NoError(t, err)
+		// Shutdown the client.
+		err := client.Stop(context.Background())
+		assert.NoError(t, err)
+	})
 }
diff --git a/client/httpclient_test.go b/client/httpclient_test.go
new file mode 100644
index 0000000..d49d36b
--- /dev/null
+++ b/client/httpclient_test.go
@@ -0,0 +1,48 @@
+package client
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/open-telemetry/opamp-go/client/internal"
+	"github.com/open-telemetry/opamp-go/protobufs"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHTTPPolling(t *testing.T) {
+	// Start a server.
+	srv := internal.StartMockServer(t)
+	var rcvCounter int64
+	srv.OnMessage = func(msg *protobufs.AgentToServer) *protobufs.ServerToAgent {
+		if statusReport := msg.GetStatusReport(); statusReport != nil {
+			atomic.AddInt64(&rcvCounter, 1)
+		}
+		return nil
+	}
+
+	// Start a client.
+	settings := StartSettings{}
+	settings.OpAMPServerURL = "http://" + srv.Endpoint
+	client := NewHTTP(nil)
+	prepareClient(t, &settings, client)
+
+	// Shorten the polling interval to speed up the test.
+	client.looper.SetPollingInterval(time.Millisecond * 10)
+
+	assert.NoError(t, client.Start(settings))
+
+	// Verify that status report is delivered.
+	eventually(t, func() bool { return atomic.LoadInt64(&rcvCounter) == 1 })
+
+	// Verify that status report is delivered again. Polling should ensure this.
+	eventually(t, func() bool { return atomic.LoadInt64(&rcvCounter) == 2 })
+
+	// Shutdown the server.
+	srv.Close()
+
+	// Shutdown the client.
+	err := client.Stop(context.Background())
+	assert.NoError(t, err)
+}
diff --git a/client/internal/receiver_test.go b/client/internal/wsreceiver_test.go
similarity index 81%
rename from client/internal/receiver_test.go
rename to client/internal/wsreceiver_test.go
index 1c12300..24ccd4e 100644
--- a/client/internal/receiver_test.go
+++ b/client/internal/wsreceiver_test.go
@@ -71,8 +71,8 @@ func TestServerToAgentCommand(t *testing.T) {
 					return nil
 				},
 			}
-			receiver := NewReceiver(TestLogger{t}, callbacks, nil, nil)
-			receiver.processReceivedMessage(context.Background(), &protobufs.ServerToAgent{
+			receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, nil)
+			receiver.processor.ProcessReceivedMessage(context.Background(), &protobufs.ServerToAgent{
 				Command: test.command,
 			})
 			assert.Equal(t, test.action, action, test.message)
@@ -94,8 +94,8 @@ func TestServerToAgentCommandExclusive(t *testing.T) {
 			return nil, false, nil
 		},
 	}
-	receiver := NewReceiver(TestLogger{t}, callbacks, nil, nil)
-	receiver.processReceivedMessage(context.Background(), &protobufs.ServerToAgent{
+	receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, nil)
+	receiver.processor.ProcessReceivedMessage(context.Background(), &protobufs.ServerToAgent{
 		Command: &protobufs.ServerToAgentCommand{
 			Type: protobufs.ServerToAgentCommand_Restart,
 		},
@@ -116,14 +116,14 @@ func TestOnRemoteConfigReportsError(t *testing.T) {
 	}
 
 	sender := NewSender(TestLogger{t})
-	receiver := NewReceiver(TestLogger{t}, callbacks, nil, sender)
+	receiver := NewWSReceiver(TestLogger{t}, callbacks, nil, sender)
 
-	receiver.processReceivedMessage(context.Background(), &protobufs.ServerToAgent{
+	receiver.processor.ProcessReceivedMessage(context.Background(), &protobufs.ServerToAgent{
 		RemoteConfig: &protobufs.AgentRemoteConfig{},
 	})
 
-	gotStatus := receiver.sender.nextMessage.StatusReport.RemoteConfigStatus.Status
-	gotMsg := receiver.sender.nextMessage.StatusReport.RemoteConfigStatus.ErrorMessage
+	gotStatus := receiver.sender.NextMessage().nextMessage.StatusReport.RemoteConfigStatus.Status
+	gotMsg := receiver.sender.NextMessage().nextMessage.StatusReport.RemoteConfigStatus.ErrorMessage
 
 	assert.Equal(t, expectStatus, gotStatus)
 	assert.Equal(t, expectMsg, gotMsg)
diff --git a/client/wsclient_test.go b/client/wsclient_test.go
new file mode 100644
index 0000000..0f69732
--- /dev/null
+++ b/client/wsclient_test.go
@@ -0,0 +1,53 @@
+package client
+
+import (
+	"context"
+	"sync/atomic"
+	"testing"
+
+	"github.com/gorilla/websocket"
+	"github.com/open-telemetry/opamp-go/client/internal"
+	"github.com/open-telemetry/opamp-go/client/types"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDisconnectWSByServer(t *testing.T) {
+	// Start a server.
+	srv := internal.StartMockServer(t)
+	var conn atomic.Value
+	srv.OnWSConnect = func(c *websocket.Conn) {
+		conn.Store(c)
+	}
+
+	// Start an OpAMP/WebSocket client.
+	var connected int64
+	var connectErr atomic.Value
+	settings := StartSettings{
+		Callbacks: types.CallbacksStruct{
+			OnConnectFunc: func() {
+				atomic.StoreInt64(&connected, 1)
+			},
+			OnConnectFailedFunc: func(err error) {
+				connectErr.Store(err)
+			},
+		},
+	}
+	settings.OpAMPServerURL = "ws://" + srv.Endpoint
+	client := NewWebSocket(nil)
+	startClient(t, settings, client)
+
+	// Wait for connection to be established.
+	eventually(t, func() bool { return conn.Load() != nil })
+	assert.True(t, connectErr.Load() == nil)
+
+	// Close the server and forcefully disconnect.
+	srv.Close()
+	conn.Load().(*websocket.Conn).Close()
+
+	// The client must retry and must fail now.
+	eventually(t, func() bool { return connectErr.Load() != nil })
+
+	// Stop the client.
+	err := client.Stop(context.Background())
+	assert.NoError(t, err)
+}