Implement CallInfo usage in context and add integration tests

client.go

@@ -79,6 +79,10 @@ func NewClient[Req, Res any](httpClient HTTPClient, url string, options ...Clien
 		conn := client.protocolClient.NewConn(ctx, unarySpec, request.Header())
 		conn.onRequestSend(func(r *http.Request) {
 			request.setRequestMethod(r.Method)
+			callInfo, ok := getClientCallInfoFromContext(ctx)
+			if ok {
+				callInfo.method = r.Method
+			}
 		})
 		// Send always returns an io.EOF unless the error is from the client-side.
 		// We want the user to continue to call Receive in those cases to get the
@@ -100,6 +104,7 @@ func NewClient[Req, Res any](httpClient HTTPClient, url string, options ...Clien
 		return response, conn.CloseResponse()
 	})
 	if interceptor := config.Interceptor; interceptor != nil {
+		// interceptor here is the chain
 		unaryFunc = interceptor.WrapUnary(unaryFunc)
 	}
 	client.callUnary = func(ctx context.Context, request *Request[Req]) (*Response[Res], error) {
@@ -109,6 +114,23 @@ func NewClient[Req, Res any](httpClient HTTPClient, url string, options ...Clien
 		request.spec = unarySpec
 		request.peer = client.protocolClient.Peer()
 		protocolClient.WriteRequestHeader(StreamTypeUnary, request.Header())
+
+		// Also set them in the context if there's a call info present
+		callInfo, callInfoOk := getClientCallInfoFromContext(ctx)
+		if callInfoOk {
+			callInfo.peer = request.Peer()
+			callInfo.spec = request.Spec()
+			// A client could have set request headers in the call info OR the request wrapper
+			// So if a callInfo exists in context, merge any headers from there into the request wrapper
+			// so that all headers are sent in the request
+			mergeHeaders(request.Header(), callInfo.requestHeader)
+
+			// Copy the call info into a sentinel value. This is so we can compare
+			// the sentinel value against the call info in context. If they're different,
+			// we can stop the request. This protects against changing the context in interceptors.
+			ctx = context.WithValue(ctx, sentinelContextKey{}, callInfo)
+		}
+
 		response, err := unaryFunc(ctx, request)
 		if err != nil {
 			return nil, err
@@ -117,6 +139,12 @@ func NewClient[Req, Res any](httpClient HTTPClient, url string, options ...Clien
 		if !ok {
 			return nil, errorf(CodeInternal, "unexpected client response type %T", response)
 		}
+		if callInfoOk {
+			// Wrap the response and set it into the context callinfo
+			callInfo.responseSource = &responseWrapper[Res]{
+				response: typed,
+			}
+		}
 		return typed, nil
 	}
 	return client
@@ -135,8 +163,8 @@ func (c *Client[Req, Res]) CallUnary(ctx context.Context, request *Request[Req])
 //
 // This option eliminates the [Request] and [Response] wrappers, and instead uses the
 // context.Context to propagate information such as headers.
-func (c *Client[Req, Res]) CallUnarySimple(ctx context.Context, requestMsg *Req) (*Res, error) {
-	response, err := c.CallUnary(ctx, requestFromContext(ctx, requestMsg))
+func (c *Client[Req, Res]) CallUnarySimple(ctx context.Context, request *Req) (*Res, error) {
+	response, err := c.CallUnary(ctx, NewRequest(request))
 	if response != nil {
 		return response.Msg, err
 	}
@@ -159,12 +187,33 @@ func (c *Client[Req, Res]) CallServerStream(ctx context.Context, request *Reques
 	if c.err != nil {
 		return nil, c.err
 	}
+	callInfo, callInfoOk := getClientCallInfoFromContext(ctx)
+	// Set values in the context if there's a call info present
+	if callInfoOk {
+		// Copy the call info into a sentinel value. This is so we can compare
+		// the sentinel value against the call info in context. If they're different,
+		// we can stop the request. This protects against changing the context in interceptors.
+		ctx = context.WithValue(ctx, sentinelContextKey{}, callInfo)
+	}
 	conn := c.newConn(ctx, StreamTypeServer, func(r *http.Request) {
 		request.method = r.Method
 	})
-	request.spec = conn.Spec()
 	request.peer = conn.Peer()
+	request.spec = conn.Spec()
+
+	// Set values in the context if there's a call info present
+	if callInfoOk {
+		callInfo.peer = conn.Peer()
+		callInfo.spec = conn.Spec()
+		callInfo.responseSource = conn
+
+		// Merge any callInfo request headers first, then do the request.
+		// so that context headers show first in the list of headers
+		mergeHeaders(conn.RequestHeader(), callInfo.RequestHeader())
+	}
+
 	mergeHeaders(conn.RequestHeader(), request.header)
+
 	// Send always returns an io.EOF unless the error is from the client-side.
 	// We want the user to continue to call Receive in those cases to get the
 	// full error from the server-side.
@@ -188,7 +237,7 @@ func (c *Client[Req, Res]) CallServerStream(ctx context.Context, request *Reques
 // This option eliminates the [Request] wrapper, and instead uses the context.Context to
 // propagate information such as headers.
 func (c *Client[Req, Res]) CallServerStreamSimple(ctx context.Context, requestMsg *Req) (*ServerStreamForClient[Res], error) {
-	return c.CallServerStream(ctx, requestFromContext(ctx, requestMsg))
+	return c.CallServerStream(ctx, NewRequest(requestMsg))
 }
 
 // CallBidiStream calls a bidirectional streaming procedure.

connect.go

@@ -211,11 +211,6 @@ func (r *Request[_]) setRequestMethod(method string) {
 	r.method = method
 }
 
-// setHeader sets the request header to the given value.
-func (r *Request[_]) setHeader(header http.Header) {
-	r.header = header
-}
-
 // AnyRequest is the common method set of every [Request], regardless of type
 // parameter. It's used in unary interceptors.
 //
@@ -287,16 +282,6 @@ func (r *Response[_]) Trailer() http.Header {
 	return r.trailer
 }
 
-// setHeader sets the response header.
-func (r *Response[_]) setHeader(header http.Header) {
-	r.header = header
-}
-
-// setTrailer sets the response trailer.
-func (r *Response[_]) setTrailer(trailer http.Header) {
-	r.trailer = trailer
-}
-
 // internalOnly implements AnyResponse.
 func (r *Response[_]) internalOnly() {}
 
@@ -383,6 +368,47 @@ type hasHTTPMethod interface {
 	getHTTPMethod() string
 }
 
+type errStreamingClientConn struct {
+	StreamingClientConn
+	err error
+}
+
+func (c *errStreamingClientConn) Receive(msg any) error {
+	return c.err
+}
+
+func (c *errStreamingClientConn) Spec() Spec {
+	return Spec{}
+}
+
+func (c *errStreamingClientConn) Peer() Peer {
+	return Peer{}
+}
+
+func (c *errStreamingClientConn) Send(msg any) error {
+	return c.err
+}
+
+func (c *errStreamingClientConn) CloseRequest() error {
+	return c.err
+}
+
+func (c *errStreamingClientConn) CloseResponse() error {
+	return c.err
+}
+
+func (c *errStreamingClientConn) RequestHeader() http.Header {
+	return make(http.Header)
+}
+
+func (c *errStreamingClientConn) ResponseHeader() http.Header {
+	return make(http.Header)
+}
+
+func (c *errStreamingClientConn) ResponseTrailer() http.Header {
+	return make(http.Header)
+}
+
 // receiveUnaryResponse unmarshals a message from a StreamingClientConn, then
 // envelopes the message and attaches headers and trailers. It attempts to
 // consume the response stream and isn't appropriate when receiving multiple