2 files changed, 36 insertions, 29 deletions

diff --git a/java/src/com/android/intentresolver/ChooserActivity.java b/java/src/com/android/intentresolver/ChooserActivity.java
index 2a73c42a..30ec2b91 100644
--- a/java/src/com/android/intentresolver/ChooserActivity.java
+++ b/java/src/com/android/intentresolver/ChooserActivity.java
@@ -279,10 +279,7 @@ public class ChooserActivity extends ResolverActivity implements
                         finish();
                     }
                 },
-                () -> {
-                    mRefinementManager.destroy();
-                    finish();
-                });
+                this::finish);
 
         mChooserContentPreviewUi = new ChooserContentPreviewUi(
                 mChooserRequest.getTargetIntent(),
diff --git a/java/src/com/android/intentresolver/ChooserRefinementManager.java b/java/src/com/android/intentresolver/ChooserRefinementManager.java
index 23cf6ba0..72b5305d 100644
--- a/java/src/com/android/intentresolver/ChooserRefinementManager.java
+++ b/java/src/com/android/intentresolver/ChooserRefinementManager.java
@@ -17,6 +17,7 @@
 package com.android.intentresolver;
 
 import android.annotation.Nullable;
+import android.annotation.UiThread;
 import android.app.Activity;
 import android.content.Context;
 import android.content.Intent;
@@ -25,7 +26,6 @@ import android.content.IntentSender.SendIntentException;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.Parcel;
-import android.os.Parcelable;
 import android.os.ResultReceiver;
 import android.util.Log;
 
@@ -41,6 +41,7 @@ import java.util.function.Consumer;
  * convert the format of the payload, or lazy-download some data that was deferred in the original
  * call).
  */
+@UiThread
 public final class ChooserRefinementManager {
     private static final String TAG = "ChooserRefinement";
 
@@ -51,7 +52,7 @@ public final class ChooserRefinementManager {
     private final Consumer<TargetInfo> mOnSelectionRefined;
     private final Runnable mOnRefinementCancelled;
 
-    @Nullable
+    @Nullable    // Non-null only during an active refinement session.
     private RefinementResultReceiver mRefinementResultReceiver;
 
     public ChooserRefinementManager(
@@ -98,7 +99,10 @@ public final class ChooserRefinementManager {
                         mOnRefinementCancelled.run();
                     }
                 },
-                mOnRefinementCancelled,
+                () -> {
+                    destroy();
+                    mOnRefinementCancelled.run();
+                },
                 mContext.getMainThreadHandler());
 
         Intent refinementRequest = makeRefinementRequest(mRefinementResultReceiver, selectedTarget);
@@ -114,7 +118,7 @@ public final class ChooserRefinementManager {
     /** Clean up any ongoing refinement session. */
     public void destroy() {
         if (mRefinementResultReceiver != null) {
-            mRefinementResultReceiver.destroy();
+            mRefinementResultReceiver.destroyReceiver();
             mRefinementResultReceiver = null;
         }
     }
@@ -151,7 +155,7 @@ public final class ChooserRefinementManager {
             mOnRefinementCancelled = onRefinementCancelled;
         }
 
-        public void destroy() {
+        public void destroyReceiver() {
             mDestroyed = true;
         }
 
@@ -161,27 +165,14 @@ public final class ChooserRefinementManager {
                 Log.e(TAG, "Destroyed RefinementResultReceiver received a result");
                 return;
             }
-            if (resultData == null) {
-                Log.e(TAG, "RefinementResultReceiver received null resultData");
-                // TODO: treat as cancellation?
-                return;
-            }
 
-            switch (resultCode) {
-                case Activity.RESULT_CANCELED:
-                    mOnRefinementCancelled.run();
-                    break;
-                case Activity.RESULT_OK:
-                    Parcelable intentParcelable = resultData.getParcelable(Intent.EXTRA_INTENT);
-                    if (intentParcelable instanceof Intent) {
-                        mOnSelectionRefined.accept((Intent) intentParcelable);
-                    } else {
-                        Log.e(TAG, "No valid Intent.EXTRA_INTENT in 'OK' refinement result data");
-                    }
-                    break;
-                default:
-                    Log.w(TAG, "Received unknown refinement result " + resultCode);
-                    break;
+            destroyReceiver();  // This is the single callback we'll accept from this session.
+
+            Intent refinedResult = tryToExtractRefinedResult(resultCode, resultData);
+            if (refinedResult == null) {
+                mOnRefinementCancelled.run();
+            } else {
+                mOnSelectionRefined.accept(refinedResult);
             }
         }
 
@@ -197,5 +188,24 @@ public final class ChooserRefinementManager {
             parcel.recycle();
             return receiverForSending;
         }
+
+        /**
+         * Get the refinement from the result data, if possible, or log diagnostics and return null.
+         */
+        @Nullable
+        private static Intent tryToExtractRefinedResult(int resultCode, Bundle resultData) {
+            if (Activity.RESULT_CANCELED == resultCode) {
+                Log.i(TAG, "Refinement canceled by caller");
+            } else if (Activity.RESULT_OK != resultCode) {
+                Log.w(TAG, "Canceling refinement on unrecognized result code " + resultCode);
+            } else if (resultData == null) {
+                Log.e(TAG, "RefinementResultReceiver received null resultData; canceling");
+            } else if (!(resultData.getParcelable(Intent.EXTRA_INTENT) instanceof Intent)) {
+                Log.e(TAG, "No valid Intent.EXTRA_INTENT in 'OK' refinement result data");
+            } else {
+                return resultData.getParcelable(Intent.EXTRA_INTENT, Intent.class);
+            }
+            return null;
+        }
     }
 }