fix ctxwconn

pkg/codec/batch.go

-package codec
-
-// BatchElem is an element in a batch request.
-type BatchElem struct {
-	Method string
-	Params any
-
-	IsNotification bool
-
-	// The result is unmarshaled into this field. Result must be set to a
-	// non-nil pointer value of the desired type, otherwise the response will be
-	// discarded.
-	Result any
-	// Error is set if the server returns an error for this request, or if
-	// unmarshaling into Result fails. It is not set for I/O errors.
-	Error error
-}

pkg/codec/context.go

@@ -7,8 +7,7 @@ type clientContextKey struct{}
 // ClientFromContext retrieves the client from the context, if any. This can be used to perform
 // 'reverse calls' in a handler method.
 func ContextWithConn(ctx context.Context, c Conn) context.Context {
-	client, _ := ctx.Value(clientContextKey{}).(Conn)
-	return context.WithValue(ctx, clientContextKey{}, client)
+	return context.WithValue(ctx, clientContextKey{}, c)
 }
 
 // ClientFromContext retrieves the client from the context, if any. This can be used to perform

pkg/codec/context_test.go

+package codec
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestConnContext(t *testing.T) {
+	ctx := context.Background()
+
+	{
+		_, ok := ConnFromContext(ctx)
+		require.False(t, ok)
+	}
+	d := &DummyClient{}
+	ctx = ContextWithConn(ctx, d)
+
+	{
+		conn, ok := ConnFromContext(ctx)
+		require.True(t, ok)
+		require.Equal(t, conn, d)
+	}
+
+	{
+		conn, ok := StreamingConnFromContext(ctx)
+		require.True(t, ok)
+		require.Equal(t, conn, d)
+	}
+
+}

pkg/codec/dummy.go

+package codec
+
+import "context"
+
+type DummyClient struct{}
+
+func (d *DummyClient) Notify(ctx context.Context, method string, params any) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (d *DummyClient) Do(ctx context.Context, result any, method string, params any) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (d *DummyClient) BatchCall(ctx context.Context, b ...*BatchElem) error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (d *DummyClient) Close() error {
+	panic("not implemented") // TODO: Implement
+}
+
+func (d *DummyClient) Mount(_ Middleware) {
+	panic("not implemented") // TODO: Implement
+}

pkg/codec/reqresp.go

@@ -16,6 +16,22 @@ type ResponseWriter interface {
 	Notify(method string, v any) error
 }
 
+// BatchElem is an element in a batch request.
+type BatchElem struct {
+	Method string
+	Params any
+
+	IsNotification bool
+
+	// The result is unmarshaled into this field. Result must be set to a
+	// non-nil pointer value of the desired type, otherwise the response will be
+	// discarded.
+	Result any
+	// Error is set if the server returns an error for this request, or if
+	// unmarshaling into Result fails. It is not set for I/O errors.
+	Error error
+}
+
 type Response struct {
 	Version Version         `json:"jsonrpc,omitempty"`
 	ID      *ID             `json:"id,omitempty"`