diff --git a/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java b/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java
index eeff202842..629e477e6c 100755
--- a/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java
+++ b/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java
@@ -33,6 +33,7 @@
 import static com.uber.nullaway.NullAway.getTreesInstance;
 
 import com.google.common.base.Joiner;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Iterables;
@@ -133,7 +134,7 @@ public Description createErrorDescription(
     }
 
     if (config.suggestSuppressions() && suggestTree != null) {
-      builder = addSuggestedSuppression(errorMessage, suggestTree, builder);
+      builder = addSuggestedSuppression(errorMessage, suggestTree, builder, state);
     }
 
     if (config.serializationIsActive()) {
@@ -187,7 +188,10 @@ private boolean hasPathSuppression(TreePath treePath, String subcheckerName) {
   }
 
   private Description.Builder addSuggestedSuppression(
-      ErrorMessage errorMessage, Tree suggestTree, Description.Builder builder) {
+      ErrorMessage errorMessage,
+      Tree suggestTree,
+      Description.Builder builder,
+      VisitorState state) {
     switch (errorMessage.messageType) {
       case DEREFERENCE_NULLABLE:
       case RETURN_NULLABLE:
@@ -201,7 +205,7 @@ private Description.Builder addSuggestedSuppression(
         }
         break;
       case CAST_TO_NONNULL_ARG_NONNULL:
-        builder = removeCastToNonNullFix(suggestTree, builder);
+        builder = removeCastToNonNullFix(suggestTree, builder, state);
         break;
       case WRONG_OVERRIDE_RETURN:
         builder = addSuppressWarningsFix(suggestTree, builder, suppressionName);
@@ -334,20 +338,23 @@ private Description.Builder addCastToNonNullFix(Tree suggestTree, Description.Bu
   }
 
   private Description.Builder removeCastToNonNullFix(
-      Tree suggestTree, Description.Builder builder) {
-    assert suggestTree.getKind() == Tree.Kind.METHOD_INVOCATION;
-    final MethodInvocationTree invTree = (MethodInvocationTree) suggestTree;
-    final Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(invTree);
-    final String qualifiedName =
-        ASTHelpers.enclosingClass(methodSymbol) + "." + methodSymbol.getSimpleName().toString();
-    if (!qualifiedName.equals(config.getCastToNonNullMethod())) {
-      throw new RuntimeException("suggestTree should point to the castToNonNull invocation.");
-    }
+      Tree suggestTree, Description.Builder builder, VisitorState state) {
+    // Note: Here suggestTree refers to the argument being cast. We need to find the
+    // castToNonNull(...) invocation to be replaced by it. Fortunately, state.getPath()
+    // should be currently pointing at said call.
+    Tree currTree = state.getPath().getLeaf();
+    Preconditions.checkArgument(
+        currTree.getKind() == Tree.Kind.METHOD_INVOCATION,
+        String.format("Expected castToNonNull invocation expression, found:\n%s", currTree));
+    final MethodInvocationTree invTree = (MethodInvocationTree) currTree;
+    Preconditions.checkArgument(
+        invTree.getArguments().contains(suggestTree),
+        String.format(
+            "Method invocation tree %s does not contain the expression %s as an argument being cast",
+            invTree, suggestTree));
     // Remove the call to castToNonNull:
     final SuggestedFix fix =
-        SuggestedFix.builder()
-            .replace(suggestTree, invTree.getArguments().get(0).toString())
-            .build();
+        SuggestedFix.builder().replace(invTree, suggestTree.toString()).build();
     return builder.addFix(fix);
   }
 
diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index b78c228657..d7c0afb42c 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -1712,7 +1712,9 @@ private Description checkCastToNonNullTakesNullable(
                 + "at the invocation site, but which are known not to be null at runtime.";
         return errorBuilder.createErrorDescription(
             new ErrorMessage(MessageTypes.CAST_TO_NONNULL_ARG_NONNULL, message),
-            tree,
+            // The Tree passed as suggestTree is the expression being cast
+            // to avoid recomputing the arg index:
+            actual,
             buildDescription(tree),
             state,
             null);
diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayAutoSuggestTest.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayAutoSuggestTest.java
index 5aec8370a2..92138b2f79 100644
--- a/nullaway/src/test/java/com/uber/nullaway/NullAwayAutoSuggestTest.java
+++ b/nullaway/src/test/java/com/uber/nullaway/NullAwayAutoSuggestTest.java
@@ -24,6 +24,7 @@
 
 import com.google.errorprone.BugCheckerRefactoringTestHelper;
 import com.google.errorprone.ErrorProneFlags;
+import com.uber.nullaway.testlibrarymodels.TestLibraryModels;
 import java.io.IOException;
 import org.junit.Before;
 import org.junit.Rule;
@@ -57,6 +58,8 @@ private BugCheckerRefactoringTestHelper makeTestHelper() {
         .setArgs(
             "-d",
             temporaryFolder.getRoot().getAbsolutePath(),
+            "-processorpath",
+            TestLibraryModels.class.getProtectionDomain().getCodeSource().getLocation().getPath(),
             // the remaining args are not needed right now, but they will be necessary when we
             // switch to the more modern newInstance() API
             "-XepOpt:NullAway:AnnotatedPackages=com.uber,com.ubercab,io.reactivex",
@@ -132,6 +135,32 @@ public void removeUnnecessaryCastToNonNull() throws IOException {
         .doTest();
   }
 
+  @Test
+  public void removeUnnecessaryCastToNonNullFromLibraryModel() throws IOException {
+    makeTestHelper()
+        .addInputLines(
+            "Test.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "import static com.uber.nullaway.testdata.Util.castToNonNull;",
+            "class Test {",
+            "  Object test1(Object o) {",
+            "    return castToNonNull(\"CAST_REASON\",o,42);",
+            "  }",
+            "}")
+        .addOutputLines(
+            "out/Test.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "import static com.uber.nullaway.testdata.Util.castToNonNull;",
+            "class Test {",
+            "  Object test1(Object o) {",
+            "    return o;",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void suggestSuppressionOnMethodRef() throws IOException {
     makeTestHelper()