From 98650211ac1c811594b9ff6802c9978ee5c7ee55 Mon Sep 17 00:00:00 2001
From: Ryan McGuire <ryand_mcguire@sweetwater.com>
Date: Wed, 24 Dec 2025 10:47:28 -0500
Subject: [PATCH] enhance context management for custom types

---
 pkg/app/setup_custom.go | 11 ++++++-----
 pkg/config/ctx.go       | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/pkg/app/setup_custom.go b/pkg/app/setup_custom.go
index db20354..daacb9d 100644
--- a/pkg/app/setup_custom.go
+++ b/pkg/app/setup_custom.go
@@ -17,7 +17,7 @@ import (
 // that overloads *config.AppConfig. Will perform normal env
 // substitutions for AppConfig, but env overrides for custom type
 // are up to the caller.
-func MustLoadConfigInto[T any](ctx context.Context, into T) (context.Context, T) {
+func MustLoadConfigInto[T interface{ *config.AppConfig }](ctx context.Context, into T) (context.Context, T) {
 	// Step 1: Check our custom type for required *config.AppConfig
 	if err := HasAppConfig(into); err != nil {
 		panic(err)
@@ -50,13 +50,14 @@ func MustLoadConfigInto[T any](ctx context.Context, into T) (context.Context, T)
 	}
 
 	// Step 6: Update context, return custom type
+	ctx = config.AddToContextFor(ctx, into)
 	return ctx, into
 }
 
 func setAppConfig[T any](target T, appConfig *config.AppConfig) error {
 	// Ensure target is a pointer to a struct
 	v := reflect.ValueOf(target)
-	if v.Kind() != reflect.Ptr || v.IsNil() {
+	if v.Kind() != reflect.Pointer || v.IsNil() {
 		return errors.New("target must be a non-nil pointer to a struct")
 	}
 
@@ -68,7 +69,7 @@ func setAppConfig[T any](target T, appConfig *config.AppConfig) error {
 	// Replace *config.AppConfig
 	for i := range v.NumField() {
 		field := v.Field(i)
-		if field.Type() == reflect.TypeOf((*config.AppConfig)(nil)) {
+		if field.Type() == reflect.TypeFor[*config.AppConfig]() {
 			if !field.CanSet() {
 				return fmt.Errorf("field %q cannot be set", v.Type().Field(i).Name)
 			}
@@ -82,7 +83,7 @@ func setAppConfig[T any](target T, appConfig *config.AppConfig) error {
 
 func HasAppConfig[T any](target T) error {
 	v := reflect.ValueOf(target)
-	if v.Kind() != reflect.Ptr || v.IsNil() {
+	if v.Kind() != reflect.Pointer || v.IsNil() {
 		return errors.New("target must be a non-nil pointer to a struct")
 	}
 	v = v.Elem()
@@ -94,7 +95,7 @@ func HasAppConfig[T any](target T) error {
 	hasAppConfig := false
 	for i := range v.NumField() {
 		field := v.Type().Field(i)
-		if field.Type == reflect.TypeOf((*config.AppConfig)(nil)) {
+		if field.Type == reflect.TypeFor[*config.AppConfig]() {
 			hasAppConfig = true
 			break
 		}
diff --git a/pkg/config/ctx.go b/pkg/config/ctx.go
index 9345196..1b19ca7 100644
--- a/pkg/config/ctx.go
+++ b/pkg/config/ctx.go
@@ -14,6 +14,41 @@ func (ac *AppConfig) AddToCtx(ctx context.Context) context.Context {
 	return context.WithValue(ctx, appConfigCtxKey, ac)
 }
 
+// Add to Ctx for custom type that embeds AppConfig
+func AddToContextFor[T interface {
+	*AppConfig
+}](ctx context.Context, cfg T) context.Context {
+	return context.WithValue(ctx, appConfigCtxKey, cfg)
+}
+
+// FromCtxFor retrieves custom config that embeds AppConfig from context
+func FromCtxFor[T interface {
+	*AppConfig
+}](ctx context.Context) (T, error) {
+	ctxData := ctx.Value(appConfigCtxKey)
+	if ctxData == nil {
+		return nil, errors.New("no config found in context")
+	}
+
+	cfg, ok := ctxData.(T)
+	if !ok {
+		return nil, errors.New("invalid config stored in context")
+	}
+
+	return cfg, nil
+}
+
+// MustFromCtxFor retrieves custom config that embeds AppConfig from context, or panics if not found.
+func MustFromCtxFor[T interface {
+	*AppConfig
+}](ctx context.Context) T {
+	cfg, err := FromCtxFor[T](ctx)
+	if err != nil {
+		panic(err)
+	}
+	return cfg
+}
+
 // MustFromCtx retrieves the AppConfig from the context, or panics if not found.
 func MustFromCtx(ctx context.Context) *AppConfig {
 	cfg, err := FromCtx(ctx)