1 files changed, 9 insertions, 1 deletions
diff --git a/notification/src/com/android/textclassifier/notification/SmartSuggestionsHelper.java b/notification/src/com/android/textclassifier/notification/SmartSuggestionsHelper.java
index 0fee3b3..44af148 100644
--- a/notification/src/com/android/textclassifier/notification/SmartSuggestionsHelper.java
+++ b/notification/src/com/android/textclassifier/notification/SmartSuggestionsHelper.java
@@ -312,7 +312,15 @@ public class SmartSuggestionsHelper {
    */
   public void onNotificationExpansionChanged(
       StatusBarNotification statusBarNotification, boolean isExpanded) {
-    SmartSuggestionsLogSession session = sessionCache.get(statusBarNotification.getKey());
+    onNotificationExpansionChanged(statusBarNotification.getKey(), isExpanded);
+  }
+
+  /**
+   * Similar to {@link onNotificationExpansionChanged}, except that this takes the notificataion key
+   * as input.
+   */
+  public void onNotificationExpansionChanged(String key, boolean isExpanded) {
+    SmartSuggestionsLogSession session = sessionCache.get(key);
     if (session == null) {
       return;
     }