fix custom type management

pkg/app/setup_custom.go

@@ -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
 		}

pkg/config/ctx.go

@@ -3,6 +3,7 @@ package config
 import (
 	"context"
 	"errors"
+	"reflect"
 )
 
 type appConfigKey uint8
@@ -14,6 +15,35 @@ 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 any](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 any](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 any](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)
@@ -24,16 +54,38 @@ func MustFromCtx(ctx context.Context) *AppConfig {
 }
 
 // FromCtx retrieves the AppConfig from the context.
+// It handles both direct *AppConfig and custom types that embed *AppConfig.
 func FromCtx(ctx context.Context) (*AppConfig, error) {
 	ctxData := ctx.Value(appConfigCtxKey)
 	if ctxData == nil {
 		return nil, errors.New("no config found in context")
 	}
 
-	cfg, ok := ctxData.(*AppConfig)
-	if !ok {
-		return nil, errors.New("invalid config stored in context")
-	}
-
+	// Try direct type assertion first
+	if cfg, ok := ctxData.(*AppConfig); ok {
 		return cfg, nil
 	}
+
+	// Try to extract *AppConfig from a custom type that embeds it
+	v := reflect.ValueOf(ctxData)
+	if v.Kind() != reflect.Pointer || v.IsNil() {
+		return nil, errors.New("invalid config stored in context: must be a non-nil pointer")
+	}
+
+	v = v.Elem()
+	if v.Kind() != reflect.Struct {
+		return nil, errors.New("invalid config stored in context: must be a pointer to a struct")
+	}
+
+	// Look for *AppConfig field
+	for i := range v.NumField() {
+		field := v.Field(i)
+		if field.Type() == reflect.TypeFor[*AppConfig]() {
+			if appConfig, ok := field.Interface().(*AppConfig); ok {
+				return appConfig, nil
+			}
+		}
+	}
+
+	return nil, errors.New("no *AppConfig found in context value")
+}