feat!: replace msgpack with protobuf for all NATS messages

BREAKING CHANGE: All NATS message serialization now uses Protocol Buffers. - Added proto/messages/v1/messages.proto with 22 message types - Generated Go code at gen/messagespb/ - messages/ package now exports type aliases to proto types - natsutil.Publish/Request/Decode use proto.Marshal/Unmarshal - Removed legacy MessageHandler, OnMessage, wrapMapHandler - TypedMessageHandler now returns (proto.Message, error) - EffectiveQuery is now a free function: messages.EffectiveQuery(req) - Removed msgpack dependency entirely
2026-02-21 14:58:05 -05:00
parent 3585d81ff5
commit 13ef1df109
12 changed files with 3074 additions and 1293 deletions
--- a/handler/handler.go
+++ b/handler/handler.go
@@ -2,30 +2,27 @@
 package handler

 import (
-	"context"
-	"fmt"
-	"log/slog"
-	"os"
-	"os/signal"
-	"syscall"
+"context"
+"fmt"
+"log/slog"
+"os"
+"os/signal"
+"syscall"

-	"github.com/nats-io/nats.go"
+"github.com/nats-io/nats.go"
+"google.golang.org/protobuf/proto"

-	"git.daviestechlabs.io/daviestechlabs/handler-base/config"
-	"git.daviestechlabs.io/daviestechlabs/handler-base/health"
-	"git.daviestechlabs.io/daviestechlabs/handler-base/natsutil"
-	"git.daviestechlabs.io/daviestechlabs/handler-base/telemetry"
+"git.daviestechlabs.io/daviestechlabs/handler-base/config"
+pb "git.daviestechlabs.io/daviestechlabs/handler-base/gen/messagespb"
+"git.daviestechlabs.io/daviestechlabs/handler-base/health"
+"git.daviestechlabs.io/daviestechlabs/handler-base/natsutil"
+"git.daviestechlabs.io/daviestechlabs/handler-base/telemetry"
 )

-// MessageHandler is the callback for processing decoded NATS messages.
-// data is the msgpack-decoded map. Return a response map (or nil for no reply).
-type MessageHandler func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error)
-
-// TypedMessageHandler processes the raw NATS message without pre-decoding to
-// map[string]any. Services unmarshal msg.Data into their own typed structs,
-// avoiding the double-decode overhead. Return any msgpack-serialisable value
-// (a typed struct, map, or nil for no reply).
-type TypedMessageHandler func(ctx context.Context, msg *nats.Msg) (any, error)
+// TypedMessageHandler processes the raw NATS message.
+// Services unmarshal msg.Data into their own typed structs via natsutil.Decode.
+// Return a proto.Message (or nil for no reply).
+type TypedMessageHandler func(ctx context.Context, msg *nats.Msg) (proto.Message, error)

 // SetupFunc is called once before the handler starts processing messages.
 type SetupFunc func(ctx context.Context) error
@@ -35,37 +32,36 @@ type TeardownFunc func(ctx context.Context) error

 // Handler is the base service runner that wires NATS, health, and telemetry.
 type Handler struct {
-	Settings   *config.Settings
-	NATS       *natsutil.Client
-	Telemetry  *telemetry.Provider
-	Subject    string
-	QueueGroup string
+Settings   *config.Settings
+NATS       *natsutil.Client
+Telemetry  *telemetry.Provider
+Subject    string
+QueueGroup string

-	onSetup        SetupFunc
-	onTeardown     TeardownFunc
-	onMessage      MessageHandler
-	onTypedMessage TypedMessageHandler
-	running        bool
+onSetup        SetupFunc
+onTeardown     TeardownFunc
+onTypedMessage TypedMessageHandler
+running        bool
 }

 // New creates a Handler for the given NATS subject.
 func New(subject string, settings *config.Settings) *Handler {
-	if settings == nil {
-		settings = config.Load()
-	}
-	queueGroup := settings.NATSQueueGroup
+if settings == nil {
+settings = config.Load()
+}
+queueGroup := settings.NATSQueueGroup

-	natsOpts := []nats.Option{}
-	if settings.NATSUser != "" && settings.NATSPassword != "" {
-		natsOpts = append(natsOpts, nats.UserInfo(settings.NATSUser, settings.NATSPassword))
-	}
+natsOpts := []nats.Option{}
+if settings.NATSUser != "" && settings.NATSPassword != "" {
+natsOpts = append(natsOpts, nats.UserInfo(settings.NATSUser, settings.NATSPassword))
+}

-	return &Handler{
-		Settings:   settings,
-		Subject:    subject,
-		QueueGroup: queueGroup,
-		NATS:       natsutil.New(settings.NATSURL, natsOpts...),
-	}
+return &Handler{
+Settings:   settings,
+Subject:    subject,
+QueueGroup: queueGroup,
+NATS:       natsutil.New(settings.NATSURL, natsOpts...),
+}
 }

 // OnSetup registers the setup callback.
@@ -74,158 +70,106 @@ func (h *Handler) OnSetup(fn SetupFunc) { h.onSetup = fn }
 // OnTeardown registers the teardown callback.
 func (h *Handler) OnTeardown(fn TeardownFunc) { h.onTeardown = fn }

-// OnMessage registers the message handler callback.
-func (h *Handler) OnMessage(fn MessageHandler) { h.onMessage = fn }
-
-// OnTypedMessage registers a typed message handler. It replaces OnMessage —
-// wrapHandler will skip the map[string]any decode and let the callback
-// unmarshal msg.Data directly.
+// OnTypedMessage registers the message handler callback.
 func (h *Handler) OnTypedMessage(fn TypedMessageHandler) { h.onTypedMessage = fn }

 // Run starts the handler: telemetry, health server, NATS subscription, and blocks until SIGTERM/SIGINT.
 func (h *Handler) Run() error {
-	// Structured logging
-	slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo})))
-	slog.Info("starting service", "name", h.Settings.ServiceName, "version", h.Settings.ServiceVersion)
+// Structured logging
+slog.SetDefault(slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo})))
+slog.Info("starting service", "name", h.Settings.ServiceName, "version", h.Settings.ServiceVersion)

-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+ctx, cancel := context.WithCancel(context.Background())
+defer cancel()

-	// Telemetry
-	tp, shutdown, err := telemetry.Setup(ctx, telemetry.Config{
-		ServiceName:      h.Settings.ServiceName,
-		ServiceVersion:   h.Settings.ServiceVersion,
-		ServiceNamespace: h.Settings.ServiceNamespace,
-		DeploymentEnv:    h.Settings.DeploymentEnv,
-		Enabled:          h.Settings.OTELEnabled,
-		Endpoint:         h.Settings.OTELEndpoint,
-	})
-	if err != nil {
-		return fmt.Errorf("telemetry setup: %w", err)
-	}
-	defer shutdown(ctx)
-	h.Telemetry = tp
+// Telemetry
+tp, shutdown, err := telemetry.Setup(ctx, telemetry.Config{
+ServiceName:      h.Settings.ServiceName,
+ServiceVersion:   h.Settings.ServiceVersion,
+ServiceNamespace: h.Settings.ServiceNamespace,
+DeploymentEnv:    h.Settings.DeploymentEnv,
+Enabled:          h.Settings.OTELEnabled,
+Endpoint:         h.Settings.OTELEndpoint,
+})
+if err != nil {
+return fmt.Errorf("telemetry setup: %w", err)
+}
+defer shutdown(ctx)
+h.Telemetry = tp

-	// Health server
-	healthSrv := health.New(
-		h.Settings.HealthPort,
-		h.Settings.HealthPath,
-		h.Settings.ReadyPath,
-		func() bool { return h.running && h.NATS.IsConnected() },
-	)
-	healthSrv.Start()
-	defer healthSrv.Stop(ctx)
+// Health server
+healthSrv := health.New(
+h.Settings.HealthPort,
+h.Settings.HealthPath,
+h.Settings.ReadyPath,
+func() bool { return h.running && h.NATS.IsConnected() },
+)
+healthSrv.Start()
+defer healthSrv.Stop(ctx)

-	// Connect to NATS
-	if err := h.NATS.Connect(); err != nil {
-		return fmt.Errorf("nats: %w", err)
-	}
-	defer h.NATS.Close()
+// Connect to NATS
+if err := h.NATS.Connect(); err != nil {
+return fmt.Errorf("nats: %w", err)
+}
+defer h.NATS.Close()

-	// User setup
-	if h.onSetup != nil {
-		slog.Info("running service setup")
-		if err := h.onSetup(ctx); err != nil {
-			return fmt.Errorf("setup: %w", err)
-		}
-	}
+// User setup
+if h.onSetup != nil {
+slog.Info("running service setup")
+if err := h.onSetup(ctx); err != nil {
+return fmt.Errorf("setup: %w", err)
+}
+}

-	// Subscribe
-	if h.onMessage == nil && h.onTypedMessage == nil {
-		return fmt.Errorf("no message handler registered")
-	}
-	if err := h.NATS.Subscribe(h.Subject, h.wrapHandler(ctx), h.QueueGroup); err != nil {
-		return fmt.Errorf("subscribe: %w", err)
-	}
+// Subscribe
+if h.onTypedMessage == nil {
+return fmt.Errorf("no message handler registered")
+}
+if err := h.NATS.Subscribe(h.Subject, h.wrapHandler(ctx), h.QueueGroup); err != nil {
+return fmt.Errorf("subscribe: %w", err)
+}

-	h.running = true
-	slog.Info("handler ready", "subject", h.Subject)
+h.running = true
+slog.Info("handler ready", "subject", h.Subject)

-	// Wait for shutdown signal
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
-	<-sigCh
+// Wait for shutdown signal
+sigCh := make(chan os.Signal, 1)
+signal.Notify(sigCh, syscall.SIGTERM, syscall.SIGINT)
+<-sigCh

-	slog.Info("shutting down")
-	h.running = false
+slog.Info("shutting down")
+h.running = false

-	// Teardown
-	if h.onTeardown != nil {
-		if err := h.onTeardown(ctx); err != nil {
-			slog.Warn("teardown error", "error", err)
-		}
-	}
+// Teardown
+if h.onTeardown != nil {
+if err := h.onTeardown(ctx); err != nil {
+slog.Warn("teardown error", "error", err)
+}
+}

-	slog.Info("shutdown complete")
-	return nil
+slog.Info("shutdown complete")
+return nil
 }

 // wrapHandler creates a nats.MsgHandler that dispatches to the registered callback.
-// If OnTypedMessage was used, msg.Data is passed directly without map decode.
-// If OnMessage was used, msg.Data is decoded to map[string]any first.
 func (h *Handler) wrapHandler(ctx context.Context) nats.MsgHandler {
-	if h.onTypedMessage != nil {
-		return h.wrapTypedHandler(ctx)
-	}
-	return h.wrapMapHandler(ctx)
+return func(msg *nats.Msg) {
+response, err := h.onTypedMessage(ctx, msg)
+if err != nil {
+slog.Error("handler error", "subject", msg.Subject, "error", err)
+if msg.Reply != "" {
+_ = h.NATS.Publish(msg.Reply, &pb.ErrorResponse{
+Error:   true,
+Message: err.Error(),
+Type:    fmt.Sprintf("%T", err),
+})
+}
+return
+}
+if response != nil && msg.Reply != "" {
+if err := h.NATS.Publish(msg.Reply, response); err != nil {
+slog.Error("failed to publish reply", "error", err)
+}
 }
-
-// wrapTypedHandler dispatches to the TypedMessageHandler (no map decode).
-func (h *Handler) wrapTypedHandler(ctx context.Context) nats.MsgHandler {
-	return func(msg *nats.Msg) {
-		response, err := h.onTypedMessage(ctx, msg)
-		if err != nil {
-			slog.Error("handler error", "subject", msg.Subject, "error", err)
-			if msg.Reply != "" {
-				_ = h.NATS.Publish(msg.Reply, map[string]any{
-					"error":   true,
-					"message": err.Error(),
-					"type":    fmt.Sprintf("%T", err),
-				})
-			}
-			return
-		}
-		if response != nil && msg.Reply != "" {
-			if err := h.NATS.Publish(msg.Reply, response); err != nil {
-				slog.Error("failed to publish reply", "error", err)
-			}
-		}
-	}
 }
-
-// wrapMapHandler dispatches to the legacy MessageHandler (decodes to map first).
-func (h *Handler) wrapMapHandler(ctx context.Context) nats.MsgHandler {
-	return func(msg *nats.Msg) {
-		data, err := natsutil.DecodeMsgpackMap(msg.Data)
-		if err != nil {
-			slog.Error("failed to decode message", "subject", msg.Subject, "error", err)
-			if msg.Reply != "" {
-				_ = h.NATS.Publish(msg.Reply, map[string]any{
-					"error":   true,
-					"message": err.Error(),
-					"type":    "DecodeError",
-				})
-			}
-			return
-		}
-
-		response, err := h.onMessage(ctx, msg, data)
-		if err != nil {
-			slog.Error("handler error", "subject", msg.Subject, "error", err)
-			if msg.Reply != "" {
-				_ = h.NATS.Publish(msg.Reply, map[string]any{
-					"error":   true,
-					"message": err.Error(),
-					"type":    fmt.Sprintf("%T", err),
-				})
-			}
-			return
-		}
-
-		if response != nil && msg.Reply != "" {
-			if err := h.NATS.Publish(msg.Reply, response); err != nil {
-				slog.Error("failed to publish reply", "error", err)
-			}
-		}
-	}
 }
--- a/handler/handler_test.go
+++ b/handler/handler_test.go
@@ -1,13 +1,15 @@
 package handler

 import (
-	"context"
-	"testing"
+"context"
+"testing"

-	"github.com/nats-io/nats.go"
-	"github.com/vmihailenco/msgpack/v5"
+"github.com/nats-io/nats.go"
+"google.golang.org/protobuf/proto"

-	"git.daviestechlabs.io/daviestechlabs/handler-base/config"
+"git.daviestechlabs.io/daviestechlabs/handler-base/config"
+pb "git.daviestechlabs.io/daviestechlabs/handler-base/gen/messagespb"
+"git.daviestechlabs.io/daviestechlabs/handler-base/natsutil"
 )

 // ────────────────────────────────────────────────────────────────────────────
@@ -15,75 +17,75 @@ import (
 // ────────────────────────────────────────────────────────────────────────────

 func TestNewHandler(t *testing.T) {
-	cfg := config.Load()
-	cfg.ServiceName = "test-handler"
-	cfg.NATSQueueGroup = "test-group"
+cfg := config.Load()
+cfg.ServiceName = "test-handler"
+cfg.NATSQueueGroup = "test-group"

-	h := New("ai.test.subject", cfg)
-	if h.Subject != "ai.test.subject" {
-		t.Errorf("Subject = %q", h.Subject)
-	}
-	if h.QueueGroup != "test-group" {
-		t.Errorf("QueueGroup = %q", h.QueueGroup)
-	}
-	if h.Settings.ServiceName != "test-handler" {
-		t.Errorf("ServiceName = %q", h.Settings.ServiceName)
-	}
+h := New("ai.test.subject", cfg)
+if h.Subject != "ai.test.subject" {
+t.Errorf("Subject = %q", h.Subject)
+}
+if h.QueueGroup != "test-group" {
+t.Errorf("QueueGroup = %q", h.QueueGroup)
+}
+if h.Settings.ServiceName != "test-handler" {
+t.Errorf("ServiceName = %q", h.Settings.ServiceName)
+}
 }

 func TestNewHandlerNilSettings(t *testing.T) {
-	h := New("ai.test", nil)
-	if h.Settings == nil {
-		t.Fatal("Settings should be loaded automatically")
-	}
-	if h.Settings.ServiceName != "handler" {
-		t.Errorf("ServiceName = %q, want default", h.Settings.ServiceName)
-	}
+h := New("ai.test", nil)
+if h.Settings == nil {
+t.Fatal("Settings should be loaded automatically")
+}
+if h.Settings.ServiceName != "handler" {
+t.Errorf("ServiceName = %q, want default", h.Settings.ServiceName)
+}
 }

 func TestCallbackRegistration(t *testing.T) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
+cfg := config.Load()
+h := New("ai.test", cfg)

-	setupCalled := false
-	h.OnSetup(func(ctx context.Context) error {
-		setupCalled = true
-		return nil
-	})
+setupCalled := false
+h.OnSetup(func(ctx context.Context) error {
+setupCalled = true
+return nil
+})

-	teardownCalled := false
-	h.OnTeardown(func(ctx context.Context) error {
-		teardownCalled = true
-		return nil
-	})
+teardownCalled := false
+h.OnTeardown(func(ctx context.Context) error {
+teardownCalled = true
+return nil
+})

-	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
-		return nil, nil
-	})
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+return nil, nil
+})

-	if h.onSetup == nil || h.onTeardown == nil || h.onMessage == nil {
-		t.Error("callbacks should not be nil after registration")
-	}
+if h.onSetup == nil || h.onTeardown == nil || h.onTypedMessage == nil {
+t.Error("callbacks should not be nil after registration")
+}

-	// Verify setup/teardown work when called directly.
-	_ = h.onSetup(context.Background())
-	_ = h.onTeardown(context.Background())
-	if !setupCalled || !teardownCalled {
-		t.Error("callbacks should have been invoked")
-	}
+// Verify setup/teardown work when called directly.
+_ = h.onSetup(context.Background())
+_ = h.onTeardown(context.Background())
+if !setupCalled || !teardownCalled {
+t.Error("callbacks should have been invoked")
+}
 }

 func TestTypedMessageRegistration(t *testing.T) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
+cfg := config.Load()
+h := New("ai.test", cfg)

-	h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (any, error) {
-		return map[string]any{"ok": true}, nil
-	})
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+		return &pb.ChatResponse{Response: "ok"}, nil
+})

-	if h.onTypedMessage == nil {
-		t.Error("onTypedMessage should not be nil after registration")
-	}
+if h.onTypedMessage == nil {
+t.Error("onTypedMessage should not be nil after registration")
+}
 }

 // ────────────────────────────────────────────────────────────────────────────
@@ -91,164 +93,161 @@ func TestTypedMessageRegistration(t *testing.T) {
 // ────────────────────────────────────────────────────────────────────────────

 func TestWrapHandler_ValidMessage(t *testing.T) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
+cfg := config.Load()
+h := New("ai.test", cfg)

-	var receivedData map[string]any
-	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
-		receivedData = data
-		return map[string]any{"status": "ok"}, nil
-	})
+var receivedReq pb.ChatRequest
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+if err := natsutil.Decode(msg.Data, &receivedReq); err != nil {
+return nil, err
+}
+		return &pb.ChatResponse{Response: "ok", UserId: receivedReq.GetUserId()}, nil
+})

-	// Encode a message the same way services would.
-	payload := map[string]any{
-		"request_id": "test-001",
-		"message":    "hello",
-		"premium":    true,
-	}
-	encoded, err := msgpack.Marshal(payload)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	// Call wrapHandler directly without NATS.
-	handler := h.wrapHandler(context.Background())
-	handler(&nats.Msg{
-		Subject: "ai.test.user.42.message",
-		Data:    encoded,
-	})
-
-	if receivedData == nil {
-		t.Fatal("handler was not called")
-	}
-	if receivedData["request_id"] != "test-001" {
-		t.Errorf("request_id = %v", receivedData["request_id"])
-	}
-	if receivedData["premium"] != true {
-		t.Errorf("premium = %v", receivedData["premium"])
-	}
+// Encode a message the same way services would.
+encoded, err := proto.Marshal(&pb.ChatRequest{
+RequestId: "test-001",
+Message:   "hello",
+Premium:   true,
+})
+if err != nil {
+t.Fatal(err)
 }

-func TestWrapHandler_InvalidMsgpack(t *testing.T) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
+// Call wrapHandler directly without NATS.
+handler := h.wrapHandler(context.Background())
+handler(&nats.Msg{
+Subject: "ai.test.user.42.message",
+Data:    encoded,
+})

-	handlerCalled := false
-	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
-		handlerCalled = true
-		return nil, nil
-	})
+if receivedReq.GetRequestId() != "test-001" {
+t.Errorf("request_id = %v", receivedReq.GetRequestId())
+}
+if receivedReq.GetPremium() != true {
+t.Errorf("premium = %v", receivedReq.GetPremium())
+}
+}

-	handler := h.wrapHandler(context.Background())
-	handler(&nats.Msg{
-		Subject: "ai.test",
-		Data:    []byte{0xFF, 0xFE, 0xFD}, // invalid msgpack
-	})
+func TestWrapHandler_InvalidMessage(t *testing.T) {
+cfg := config.Load()
+h := New("ai.test", cfg)

-	if handlerCalled {
-		t.Error("handler should not be called for invalid msgpack")
-	}
+handlerCalled := false
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+handlerCalled = true
+var req pb.ChatRequest
+if err := natsutil.Decode(msg.Data, &req); err != nil {
+return nil, err
+}
+return &pb.ChatResponse{}, nil
+})
+
+handler := h.wrapHandler(context.Background())
+handler(&nats.Msg{
+Subject: "ai.test",
+Data:    []byte{0xFF, 0xFE, 0xFD}, // invalid protobuf
+})
+
+// The handler IS called (wrapHandler doesn't pre-decode), but it should
+// return an error from Decode. Either way no panic.
+_ = handlerCalled
 }

 func TestWrapHandler_HandlerError(t *testing.T) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
+cfg := config.Load()
+h := New("ai.test", cfg)

-	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
-		return nil, context.DeadlineExceeded
-	})
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+return nil, context.DeadlineExceeded
+})

-	encoded, _ := msgpack.Marshal(map[string]any{"key": "val"})
-	handler := h.wrapHandler(context.Background())
+encoded, _ := proto.Marshal(&pb.ChatRequest{RequestId: "err-test"})
+handler := h.wrapHandler(context.Background())

-	// Should not panic even when handler returns error.
-	handler(&nats.Msg{
-		Subject: "ai.test",
-		Data:    encoded,
-	})
+// Should not panic even when handler returns error.
+handler(&nats.Msg{
+Subject: "ai.test",
+Data:    encoded,
+})
 }

 func TestWrapHandler_NilResponse(t *testing.T) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
+cfg := config.Load()
+h := New("ai.test", cfg)

-	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
-		return nil, nil // fire-and-forget style
-	})
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+return nil, nil // fire-and-forget style
+})

-	encoded, _ := msgpack.Marshal(map[string]any{"x": 1})
-	handler := h.wrapHandler(context.Background())
+encoded, _ := proto.Marshal(&pb.ChatRequest{RequestId: "nil-resp"})
+handler := h.wrapHandler(context.Background())

-	// Should not panic with nil response and no reply subject.
-	handler(&nats.Msg{
-		Subject: "ai.test",
-		Data:    encoded,
-	})
+// Should not panic with nil response and no reply subject.
+handler(&nats.Msg{
+Subject: "ai.test",
+Data:    encoded,
+})
 }

 // ────────────────────────────────────────────────────────────────────────────
 // wrapHandler dispatch tests — typed handler path
 // ────────────────────────────────────────────────────────────────────────────

-func TestWrapTypedHandler_ValidMessage(t *testing.T) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
+func TestWrapHandler_Typed(t *testing.T) {
+cfg := config.Load()
+h := New("ai.test", cfg)

-	type testReq struct {
-		RequestID string `msgpack:"request_id"`
-		Message   string `msgpack:"message"`
-	}
+var received pb.ChatRequest
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+if err := natsutil.Decode(msg.Data, &received); err != nil {
+return nil, err
+}
+		return &pb.ChatResponse{UserId: received.GetUserId(), Response: "ok"}, nil
+})

-	var received testReq
-	h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (any, error) {
-		if err := msgpack.Unmarshal(msg.Data, &received); err != nil {
-			return nil, err
-		}
-		return map[string]any{"status": "ok"}, nil
-	})
+encoded, _ := proto.Marshal(&pb.ChatRequest{
+RequestId: "typed-001",
+Message:   "hello typed",
+})

-	encoded, _ := msgpack.Marshal(map[string]any{
-		"request_id": "typed-001",
-		"message":    "hello typed",
-	})
+handler := h.wrapHandler(context.Background())
+handler(&nats.Msg{Subject: "ai.test", Data: encoded})

-	handler := h.wrapHandler(context.Background())
-	handler(&nats.Msg{Subject: "ai.test", Data: encoded})
-
-	if received.RequestID != "typed-001" {
-		t.Errorf("RequestID = %q", received.RequestID)
-	}
-	if received.Message != "hello typed" {
-		t.Errorf("Message = %q", received.Message)
-	}
+if received.GetRequestId() != "typed-001" {
+t.Errorf("RequestId = %q", received.GetRequestId())
+}
+if received.GetMessage() != "hello typed" {
+t.Errorf("Message = %q", received.GetMessage())
+}
 }

-func TestWrapTypedHandler_Error(t *testing.T) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
+func TestWrapHandler_TypedError(t *testing.T) {
+cfg := config.Load()
+h := New("ai.test", cfg)

-	h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (any, error) {
-		return nil, context.DeadlineExceeded
-	})
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+return nil, context.DeadlineExceeded
+})

-	encoded, _ := msgpack.Marshal(map[string]any{"key": "val"})
-	handler := h.wrapHandler(context.Background())
+encoded, _ := proto.Marshal(&pb.ChatRequest{RequestId: "err"})
+handler := h.wrapHandler(context.Background())

-	// Should not panic.
-	handler(&nats.Msg{Subject: "ai.test", Data: encoded})
+// Should not panic.
+handler(&nats.Msg{Subject: "ai.test", Data: encoded})
 }

-func TestWrapTypedHandler_NilResponse(t *testing.T) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
+func TestWrapHandler_TypedNilResponse(t *testing.T) {
+cfg := config.Load()
+h := New("ai.test", cfg)

-	h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (any, error) {
-		return nil, nil
-	})
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+return nil, nil
+})

-	encoded, _ := msgpack.Marshal(map[string]any{"x": 1})
-	handler := h.wrapHandler(context.Background())
-	handler(&nats.Msg{Subject: "ai.test", Data: encoded})
+encoded, _ := proto.Marshal(&pb.ChatRequest{RequestId: "nil"})
+handler := h.wrapHandler(context.Background())
+handler(&nats.Msg{Subject: "ai.test", Data: encoded})
 }

 // ────────────────────────────────────────────────────────────────────────────
@@ -256,56 +255,25 @@ func TestWrapTypedHandler_NilResponse(t *testing.T) {
 // ────────────────────────────────────────────────────────────────────────────

 func BenchmarkWrapHandler(b *testing.B) {
-	cfg := config.Load()
-	h := New("ai.test", cfg)
-	h.OnMessage(func(ctx context.Context, msg *nats.Msg, data map[string]any) (map[string]any, error) {
-		return map[string]any{"ok": true}, nil
-	})
+cfg := config.Load()
+h := New("ai.test", cfg)
+h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (proto.Message, error) {
+var req pb.ChatRequest
+_ = natsutil.Decode(msg.Data, &req)
+return &pb.ChatResponse{Response: "ok"}, nil
+})

-	payload := map[string]any{
-		"request_id": "bench-001",
-		"message":    "What is the capital of France?",
-		"premium":    true,
-		"top_k":      10,
-	}
-	encoded, _ := msgpack.Marshal(payload)
-	handler := h.wrapHandler(context.Background())
-	msg := &nats.Msg{Subject: "ai.test", Data: encoded}
+encoded, _ := proto.Marshal(&pb.ChatRequest{
+RequestId: "bench-001",
+Message:   "What is the capital of France?",
+Premium:   true,
+TopK:      10,
+})
+handler := h.wrapHandler(context.Background())
+msg := &nats.Msg{Subject: "ai.test", Data: encoded}

-	b.ResetTimer()
-	for b.Loop() {
-		handler(msg)
-	}
+b.ResetTimer()
+for b.Loop() {
+handler(msg)
 }
-
-func BenchmarkWrapTypedHandler(b *testing.B) {
-	type benchReq struct {
-		RequestID string `msgpack:"request_id"`
-		Message   string `msgpack:"message"`
-		Premium   bool   `msgpack:"premium"`
-		TopK      int    `msgpack:"top_k"`
-	}
-
-	cfg := config.Load()
-	h := New("ai.test", cfg)
-	h.OnTypedMessage(func(ctx context.Context, msg *nats.Msg) (any, error) {
-		var req benchReq
-		_ = msgpack.Unmarshal(msg.Data, &req)
-		return map[string]any{"ok": true}, nil
-	})
-
-	payload := map[string]any{
-		"request_id": "bench-001",
-		"message":    "What is the capital of France?",
-		"premium":    true,
-		"top_k":      10,
-	}
-	encoded, _ := msgpack.Marshal(payload)
-	handler := h.wrapHandler(context.Background())
-	msg := &nats.Msg{Subject: "ai.test", Data: encoded}
-
-	b.ResetTimer()
-	for b.Loop() {
-		handler(msg)
-	}
 }