core_service.go

@@ -18,6 +18,7 @@ type CoreService struct {
 
 	missingKeys   []MissingKey
 	missingKeysMu sync.Mutex
+	hookInstalled bool
 }
 
 // ServiceOptions configures the i18n Core service.
@@ -81,11 +82,19 @@ func NewCoreService(opts ServiceOptions) func(*core.Core) (any, error) {
 // OnStartup initialises the i18n service.
 func (s *CoreService) OnStartup(_ context.Context) error {
 	if s.svc.Mode() == ModeCollect {
-		OnMissingKey(s.handleMissingKey)
+		s.ensureMissingKeyCollector()
 	}
 	return nil
 }
 
+func (s *CoreService) ensureMissingKeyCollector() {
+	if s.hookInstalled {
+		return
+	}
+	appendMissingKeyHandler(s.handleMissingKey)
+	s.hookInstalled = true
+}
+
 func (s *CoreService) handleMissingKey(mk MissingKey) {
 	s.missingKeysMu.Lock()
 	defer s.missingKeysMu.Unlock()
@@ -112,9 +121,7 @@ func (s *CoreService) ClearMissingKeys() {
 func (s *CoreService) SetMode(mode Mode) {
 	s.svc.SetMode(mode)
 	if mode == ModeCollect {
-		OnMissingKey(s.handleMissingKey)
-	} else {
-		OnMissingKey(nil)
+		s.ensureMissingKeyCollector()
 	}
 }
 

hooks.go

@@ -11,6 +11,10 @@ import (
 
 var missingKeyHandler atomic.Value
 
+type missingKeyHandlersState struct {
+	handlers []MissingKeyHandler
+}
+
 type localeRegistration struct {
 	fsys fs.FS
 	dir  string
@@ -55,20 +59,46 @@ func loadRegisteredLocales(svc *Service) {
 
 // OnMissingKey registers a handler for missing translation keys.
 func OnMissingKey(h MissingKeyHandler) {
-	missingKeyHandler.Store(h)
+	if h == nil {
+		missingKeyHandler.Store(missingKeyHandlersState{})
+		return
+	}
+	missingKeyHandler.Store(missingKeyHandlersState{handlers: []MissingKeyHandler{h}})
+}
+
+func appendMissingKeyHandler(h MissingKeyHandler) {
+	if h == nil {
+		return
+	}
+	current := missingKeyHandlers()
+	current.handlers = append(current.handlers, h)
+	missingKeyHandler.Store(current)
+}
+
+func missingKeyHandlers() missingKeyHandlersState {
+	v := missingKeyHandler.Load()
+	if v == nil {
+		return missingKeyHandlersState{}
+	}
+	state, ok := v.(missingKeyHandlersState)
+	if !ok {
+		return missingKeyHandlersState{}
+	}
+	return state
 }
 
 func dispatchMissingKey(key string, args map[string]any) {
-	v := missingKeyHandler.Load()
-	if v == nil {
-		return
-	}
-	h, ok := v.(MissingKeyHandler)
-	if !ok || h == nil {
+	state := missingKeyHandlers()
+	if len(state.handlers) == 0 {
 		return
 	}
 	file, line := missingKeyCaller()
-	h(MissingKey{Key: key, Args: args, CallerFile: file, CallerLine: line})
+	mk := MissingKey{Key: key, Args: args, CallerFile: file, CallerLine: line}
+	for _, h := range state.handlers {
+		if h != nil {
+			h(mk)
+		}
+	}
 }
 
 func missingKeyCaller() (string, int) {

hooks_test.go

@@ -324,9 +324,44 @@ func TestOnMissingKey_Good_TranslationContextArgs(t *testing.T) {
 }
 
 func TestDispatchMissingKey_Good_NoHandler(t *testing.T) {
-	// Store nil handler (using correct type)
-	missingKeyHandler.Store(MissingKeyHandler(nil))
+	// Reset to the empty handler set.
+	OnMissingKey(nil)
 
 	// Should not panic when dispatching with nil handler
 	dispatchMissingKey("test.key", nil)
 }
+
+func TestCoreServiceSetMode_Good_PreservesMissingKeyHandlers(t *testing.T) {
+	svc, err := New()
+	require.NoError(t, err)
+
+	prev := missingKeyHandlers()
+	t.Cleanup(func() {
+		missingKeyHandler.Store(prev)
+	})
+
+	var observed int
+	OnMissingKey(func(MissingKey) {
+		observed++
+	})
+	t.Cleanup(func() {
+		OnMissingKey(nil)
+	})
+
+	coreSvc := &CoreService{svc: svc}
+	coreSvc.SetMode(ModeCollect)
+
+	_ = svc.T("missing.core.service.key")
+
+	if observed != 1 {
+		t.Fatalf("custom missing key handler called %d times, want 1", observed)
+	}
+
+	missing := coreSvc.MissingKeys()
+	if len(missing) != 1 {
+		t.Fatalf("CoreService captured %d missing keys, want 1", len(missing))
+	}
+	if missing[0].Key != "missing.core.service.key" {
+		t.Fatalf("captured missing key = %q, want %q", missing[0].Key, "missing.core.service.key")
+	}
+}