chore: refactor the result registry logic

Author: kikin81

nibel-runtime/src/main/kotlin/nibel/runtime/ComposableFragment.kt

@@ -34,11 +34,13 @@ abstract class ComposableFragment : Fragment() {
         setViewCompositionStrategy(ViewCompositionStrategy.DisposeOnViewTreeLifecycleDestroyed)
         @Suppress("DEPRECATION")
         val args = arguments?.get(Nibel.argsKey) as? Parcelable?
+        val requestKey = arguments?.getRequestKey()
 
         setContent {
             Nibel.RootDelegate.Content {
                 CompositionLocalProvider(
                     LocalImplementationType provides ImplementationType.Fragment,
+                    LocalRequestKey provides requestKey,
                 ) {
                     Nibel.NavigationDelegate.Content(rootArgs = args) {
                         ComposableContent()

nibel-runtime/src/main/kotlin/nibel/runtime/ComposeNavigationDelegate.kt

@@ -56,6 +56,12 @@ open class ComposeNavigationDelegate : NavigationDelegate<NavigationControllerAr
             ),
         )
 
+        // Set request key from composition local if available (for Fragment-based navigation)
+        val requestKey = LocalRequestKey.current
+        if (navController is NibelNavigationController) {
+            navController.setRequestKeyFromFragment(requestKey)
+        }
+
         CompositionLocalProvider(LocalNavigationController provides navController) {
             NavHost(internalNavController, startDestination = ROOT_ROUTE) {
                 registerComposable(

nibel-runtime/src/main/kotlin/nibel/runtime/CompositionLocals.kt

@@ -25,3 +25,10 @@ val LocalImplementationType = compositionLocalOf<ImplementationType?> {
 val LocalArgs = compositionLocalOf<Parcelable?> {
     null
 }
+
+/**
+ * Provides request key for result-based navigation or `null` if not navigated via navigateForResult.
+ */
+val LocalRequestKey = compositionLocalOf<String?> {
+    null
+}

nibel-runtime/src/main/kotlin/nibel/runtime/NibelArgs.kt

@@ -11,16 +11,33 @@ import androidx.lifecycle.SavedStateHandle
  */
 const val NIBEL_ARGS = "nibel_args"
 
+/**
+ * Key for request ID used in result-based navigation.
+ */
+const val NIBEL_REQUEST_KEY = "nibel_request_key"
+
 /**
  * Convert [Parcelable] args to a [Bundle].
  */
 fun Parcelable.asNibelArgs() = Bundle().apply { putParcelable(Nibel.argsKey, this@asNibelArgs) }
 
+/**
+ * Add request key to a [Bundle] for result-based navigation.
+ */
+fun Bundle.putRequestKey(key: String?) {
+    key?.let { putString(NIBEL_REQUEST_KEY, it) }
+}
+
 /**
  * Retrieve args from a [SavedStateHandle].
  */
 fun <A : Parcelable> SavedStateHandle.getNibelArgs(): A? = get<A>(Nibel.argsKey)
 
+/**
+ * Retrieve request key from a [SavedStateHandle].
+ */
+fun SavedStateHandle.getRequestKey(): String? = get<String>(NIBEL_REQUEST_KEY)
+
 /**
  * Retrieve args from a [Bundle].
  */
@@ -31,3 +48,8 @@ inline fun <reified A : Parcelable> Bundle.getNibelArgs(): A? =
         @Suppress("DEPRECATION")
         getParcelable(Nibel.argsKey)
     }
+
+/**
+ * Retrieve request key from a [Bundle].
+ */
+fun Bundle.getRequestKey(): String? = getString(NIBEL_REQUEST_KEY)

nibel-runtime/src/main/kotlin/nibel/runtime/NibelNavigationController.kt

@@ -28,18 +28,23 @@ open class NibelNavigationController(
     protected val composeNavigationContext =
         ComposeNavigationContext(internalNavController, exploredEntries)
 
-    /**
-     * Storage for result callbacks, keyed by unique request ID.
-     * Callbacks are invoked when a screen returns a result or cancels.
-     */
-    private val resultCallbacks = mutableMapOf<String, (Any?) -> Unit>()
+    // Note: Result callbacks are now stored in ResultCallbackRegistry singleton
+    // to persist across Fragment/NavigationController instance recreation
 
     /**
      * Storage for the current screen's request key (if navigated via navigateForResult).
      * Used to deliver results back to the caller.
      */
     private var currentRequestKey: String? = null
 
+    /**
+     * Sets the request key for this navigation controller instance.
+     * Used when a Fragment is navigated to via navigateForResult.
+     */
+    fun setRequestKeyFromFragment(key: String?) {
+        currentRequestKey = key
+    }
+
     override fun navigateBack() {
         onBackPressedDispatcher.onBackPressed()
     }
@@ -97,25 +102,32 @@ open class NibelNavigationController(
 
         // Store the callback with type erasure (we trust compile-time type safety)
         @Suppress("UNCHECKED_CAST")
-        resultCallbacks[requestKey] = callback as (Any?) -> Unit
+        ResultCallbackRegistry.storeCallback(requestKey, callback as (Any?) -> Unit)
 
         // Store request key so the callee can use it to return results
-        // Note: This will be passed through SavedStateHandle for Fragment entries
-        // and through the entry itself for Composable entries
+        // For Fragment entries, pass it through Bundle arguments
+        // For Composable entries, store it in currentRequestKey
         val previousRequestKey = currentRequestKey
         currentRequestKey = requestKey
 
+        // Inject request key into Fragment arguments for result-based navigation
+        if (entry is FragmentEntry) {
+            val args = entry.fragment.arguments ?: android.os.Bundle()
+            args.putRequestKey(requestKey)
+            entry.fragment.arguments = args
+        }
+
         try {
             // Perform the navigation
             navigateTo(entry, fragmentSpec, composeSpec)
         } catch (e: IllegalStateException) {
             // If navigation fails, clean up and restore state
-            resultCallbacks.remove(requestKey)
+            ResultCallbackRegistry.removeCallback(requestKey)
             currentRequestKey = previousRequestKey
             throw e
         } catch (e: IllegalArgumentException) {
             // If navigation fails, clean up and restore state
-            resultCallbacks.remove(requestKey)
+            ResultCallbackRegistry.removeCallback(requestKey)
             currentRequestKey = previousRequestKey
             throw e
         }
@@ -167,10 +179,9 @@ open class NibelNavigationController(
      * Cleans up the callback after delivery.
      */
     private fun deliverResult(requestKey: String, result: Any?) {
-        val callback = resultCallbacks.remove(requestKey)
-        if (callback != null) {
-            callback(result)
-        }
+        val callback = ResultCallbackRegistry.removeCallback(requestKey)
+        callback?.invoke(result)
+
         // Clear current request key if it matches
         if (currentRequestKey == requestKey) {
             currentRequestKey = null

nibel-runtime/src/main/kotlin/nibel/runtime/ResultCallbackRegistry.kt

@@ -0,0 +1,80 @@
+package nibel.runtime
+
+import java.util.concurrent.ConcurrentHashMap
+
+/**
+ * Global registry for result callbacks that persists across Fragment/Activity recreation.
+ * This ensures callbacks survive navigation and configuration changes.
+ *
+ * Production-hardened with:
+ * - TTL-based cleanup (5-minute expiry) to prevent memory leaks
+ * - Max callback limit (50) to prevent DoS attacks
+ * - Thread-safe implementation using ConcurrentHashMap
+ */
+internal object ResultCallbackRegistry {
+    private const val MAX_CALLBACKS = 50
+    private const val TTL_MILLIS = 5 * 60 * 1000L // 5 minutes
+
+    private val callbacks = ConcurrentHashMap<String, CallbackEntry>()
+
+    /**
+     * Stores a callback with timestamp for TTL-based cleanup.
+     */
+    fun storeCallback(requestKey: String, callback: (Any?) -> Unit) {
+        // Cleanup stale callbacks if we're at the limit
+        if (callbacks.size >= MAX_CALLBACKS) {
+            cleanupStaleCallbacks()
+
+            // If still at limit after cleanup, drop the oldest callback
+            if (callbacks.size >= MAX_CALLBACKS) {
+                val oldestKey = callbacks.entries
+                    .minByOrNull { it.value.timestamp }
+                    ?.key
+                oldestKey?.let { callbacks.remove(it) }
+            }
+        }
+
+        callbacks[requestKey] = CallbackEntry(callback, System.currentTimeMillis())
+    }
+
+    /**
+     * Removes and returns a callback, or null if not found.
+     */
+    fun removeCallback(requestKey: String): ((Any?) -> Unit)? {
+        return callbacks.remove(requestKey)?.callback
+    }
+
+    /**
+     * Checks if a callback exists for the given request key.
+     */
+    fun hasCallback(requestKey: String): Boolean {
+        return callbacks.containsKey(requestKey)
+    }
+
+    /**
+     * Clears all callbacks. Should only be used in tests or emergency cleanup.
+     */
+    fun clear() {
+        callbacks.clear()
+    }
+
+    /**
+     * Removes callbacks older than TTL_MILLIS to prevent memory leaks
+     * from abandoned navigation flows.
+     */
+    private fun cleanupStaleCallbacks() {
+        val now = System.currentTimeMillis()
+        val staleKeys = callbacks.filterValues { entry ->
+            now - entry.timestamp > TTL_MILLIS
+        }.keys
+        staleKeys.forEach { callbacks.remove(it) }
+    }
+
+    /**
+     * Internal data class to track callback with creation timestamp.
+     */
+    private data class CallbackEntry(
+        val callback: (Any?) -> Unit,
+        val timestamp: Long,
+    )
+}

sample/feature-A/src/main/kotlin/com/turo/nibel/sample/featureA/secondscreen/SecondScreen.kt

@@ -43,7 +43,10 @@ private fun SideEffectHandler(sideEffects: Flow<SecondSideEffect>) {
             is SecondSideEffect.NavigateBack -> navigateBack()
             is SecondSideEffect.NavigateToThirdScreen -> {
                 val args = ThirdArgs(it.inputText)
-                navigateTo(ThirdScreenDestination(args))
+                navigateForResult(
+                    externalDestination = ThirdScreenDestination(args),
+                    callback = it.callback,
+                )
             }
         }
     }

sample/feature-A/src/main/kotlin/com/turo/nibel/sample/featureA/secondscreen/SecondSideEffect.kt

@@ -1,8 +1,13 @@
 package com.turo.nibel.sample.featureA.secondscreen
 
+import com.turo.nibel.sample.navigation.ThirdScreenResult
+
 sealed interface SecondSideEffect {
 
     object NavigateBack : SecondSideEffect
 
-    data class NavigateToThirdScreen(val inputText: String) : SecondSideEffect
+    data class NavigateToThirdScreen(
+        val inputText: String,
+        val callback: (ThirdScreenResult?) -> Unit,
+    ) : SecondSideEffect
 }

sample/feature-A/src/main/kotlin/com/turo/nibel/sample/featureA/secondscreen/SecondViewModel.kt

@@ -28,12 +28,28 @@ class SecondViewModel @Inject constructor(
 
     fun onContinue(nextButton: SecondNextButton) {
         when (nextButton) {
-            SecondNextButton.SecondScreen ->
-                _sideEffects.tryEmit(SecondSideEffect.NavigateToThirdScreen(state.value.inputText))
+            SecondNextButton.SecondScreen -> {
+                _sideEffects.tryEmit(
+                    SecondSideEffect.NavigateToThirdScreen(
+                        inputText = state.value.inputText,
+                        callback = { result ->
+                            // Handle result from ThirdScreen
+                            result?.let {
+                                onResultFromThirdScreen(it.inputText)
+                            }
+                        },
+                    ),
+                )
+            }
         }
     }
 
     fun onInputTextChanged(inputText: String) {
         _state.value = _state.value.copy(inputText = inputText)
     }
+
+    private fun onResultFromThirdScreen(inputText: String) {
+        // Update state with the result from ThirdScreen
+        _state.value = _state.value.copy(inputText = inputText)
+    }
 }

sample/feature-B/src/main/kotlin/com/turo/nibel/sample/featureB/thirdscreen/ThirdScreen.kt

@@ -10,6 +10,7 @@ import com.turo.nibel.sample.navigation.FirstScreenDestination
 import com.turo.nibel.sample.navigation.FourthArgs
 import com.turo.nibel.sample.navigation.FourthScreenDestination
 import com.turo.nibel.sample.navigation.ThirdScreenDestination
+import com.turo.nibel.sample.navigation.ThirdScreenResult
 import kotlinx.coroutines.flow.Flow
 import nibel.annotations.ImplementationType
 import nibel.annotations.UiExternalEntry
@@ -18,6 +19,7 @@ import nibel.runtime.LocalImplementationType
 @UiExternalEntry(
     type = ImplementationType.Fragment,
     destination = ThirdScreenDestination::class,
+    result = ThirdScreenResult::class,
 )
 @Composable
 fun ThirdScreen(viewModel: ThirdViewModel = hiltViewModel()) {
@@ -40,14 +42,23 @@ fun ThirdScreen(viewModel: ThirdViewModel = hiltViewModel()) {
 private fun SideEffectHandler(sideEffects: Flow<ThirdSideEffect>) {
     SideEffectHandler(sideEffects) {
         when (it) {
-            is ThirdSideEffect.NavigateBack -> navigateBack()
+            is ThirdSideEffect.NavigateBack -> {
+                cancelResultAndNavigateBack()
+            }
+
             is ThirdSideEffect.NavigateToFourthScreen -> {
                 val args = FourthArgs(it.inputText)
                 navigateTo(FourthScreenDestination(args))
             }
 
             is ThirdSideEffect.NavigateToFirstScreen ->
                 navigateTo(FirstScreenDestination)
+
+            is ThirdSideEffect.SetResultAndNavigateBack ->
+                setResultAndNavigateBack(it.result)
+
+            is ThirdSideEffect.CancelResultAndNavigateBack ->
+                cancelResultAndNavigateBack()
         }
     }
 }