danielgtaylor · danielgtaylor · Mar 11, 2024 · Mar 9, 2024 · Mar 11, 2024 · Mar 11, 2024
diff --git a/api.go b/api.go
@@ -36,8 +36,10 @@ type ResolverWithPath interface {
 	Resolve(ctx Context, prefix *PathBuffer) []error
 }
 
-var resolverType = reflect.TypeOf((*Resolver)(nil)).Elem()
-var resolverWithPathType = reflect.TypeOf((*ResolverWithPath)(nil)).Elem()
+var (
+	resolverType         = reflect.TypeOf((*Resolver)(nil)).Elem()
+	resolverWithPathType = reflect.TypeOf((*ResolverWithPath)(nil)).Elem()
+)
 
 // Adapter is an interface that allows the API to be used with different HTTP
 // routers and frameworks. It is designed to work with the standard library
@@ -108,6 +110,32 @@ type Context interface {
 	BodyWriter() io.Writer
 }
 
+type (
+	humaContext Context
+	subContext  struct {
+		humaContext
+		override context.Context
+	}
+)
+
+func (c subContext) Context() context.Context {
+	return c.override
+}
+
+// WithContext returns a new `huma.Context` with the underlying `context.Context`
+// replaced with the given one. This is useful for middleware that needs to
+// modify the request context.
+func WithContext(ctx Context, override context.Context) Context {
+	return subContext{humaContext: ctx, override: override}
+}
+
+// WithValue returns a new `huma.Context` with the given key and value set in
+// the underlying `context.Context`. This is useful for middleware that needs to
+// set request-scoped values.
+func WithValue(ctx Context, key, value any) Context {
+	return WithContext(ctx, context.WithValue(ctx.Context(), key, value))
+}
+
 // Transformer is a function that can modify a response body before it is
 // serialized. The `status` is the HTTP status code for the response and `v` is
 // the value to be serialized. The return value is the new value to be

diff --git a/api_test.go b/api_test.go
@@ -1,6 +1,7 @@
 package huma_test
 
 import (
+	"context"
 	"net/http"
 	"testing"
 
@@ -67,3 +68,22 @@ func ExampleAdapter_handle() {
 		ctx.BodyWriter().Write([]byte("Hello, " + name))
 	})
 }
+
+func TestContextValue(t *testing.T) {
+	_, api := humatest.New(t)
+
+	api.UseMiddleware(func(ctx huma.Context, next func(huma.Context)) {
+		// Make an updated context available to the handler.
+		ctx = huma.WithValue(ctx, "foo", "bar")
+		next(ctx)
+	})
+
+	// Register a simple hello world operation in the API.
+	huma.Get(api, "/test", func(ctx context.Context, input *struct{}) (*struct{}, error) {
+		assert.Equal(t, "bar", ctx.Value("foo"))
+		return nil, nil
+	})
+
+	resp := api.Get("/test")
+	assert.Equal(t, http.StatusNoContent, resp.Code)
+}
diff --git a/docs/docs/features/middleware.md b/docs/docs/features/middleware.md
@@ -73,6 +73,29 @@ func NewHumaAPI() huma.API {
 }
 ```
 
+### Context Values
+
+The `huma.Context` interface provides a `Context()` method to retrieve the underlying request `context.Context` value. This can be used to retrieve context values in middleware and operation handlers, such as request-scoped loggers, metrics, or user information.
+
+```go title="code.go"
+if v, ok := ctx.Context().Value("some-key").(string); ok {
+	// Do something with `v`!
+}
+```
+
+You can also wrap the `huma.Context` to provide additional or override functionality. Some utilities are provided for this, including [`huma.WithValue`](https://pkg.go.dev/github.com/danielgtaylor/huma/v2#WithValue):
+
+```go title="code.go"
+func MyMiddleware(ctx huma.Context, next func(huma.Context)) {
+	// Wrap the context to add a value.
+	ctx = huma.WithValue(ctx, "some-key", "some-value")
+
+	// Call the next middleware in the chain. This eventually calls the
+	// operation handler as well.
+	next(ctx)
+}
+```
+
 ### Cookies
 
 You can use the `huma.Context` interface along with [`huma.ReadCookie`](https://pkg.go.dev/github.com/danielgtaylor/huma/v2#ReadCookie) or [`huma.ReadCookies`](https://pkg.go.dev/github.com/danielgtaylor/huma/v2#ReadCookies) to access cookies from middleware, and can also write cookies by adding `Set-Cookie` headers in the response: