diff --git a/features/update/async_accepted/feature.java b/features/update/async_accepted/feature.java
index 46ee48b1..8ed43af3 100644
--- a/features/update/async_accepted/feature.java
+++ b/features/update/async_accepted/feature.java
@@ -3,6 +3,7 @@
 import io.temporal.activity.ActivityInterface;
 import io.temporal.client.UpdateHandle;
 import io.temporal.client.UpdateOptions;
+import io.temporal.client.WorkflowUpdateTimeoutOrCancelledException;
 import io.temporal.client.WorkflowUpdateException;
 import io.temporal.client.WorkflowUpdateStage;
 import io.temporal.failure.ApplicationFailure;
@@ -18,7 +19,6 @@
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
 import org.junit.jupiter.api.Assertions;
 
 @ActivityInterface
@@ -125,7 +125,7 @@ public Run execute(Runner runner) throws Exception {
         timeoutHandle.getResultAsync(1, TimeUnit.SECONDS).get();
         Assertions.fail("unreachable");
       } catch (Exception e) {
-        Assertions.assertTrue(e.getCause() instanceof TimeoutException);
+        Assertions.assertTrue(e.getCause() instanceof WorkflowUpdateTimeoutOrCancelledException);
       }
 
       stub.finish();
diff --git a/features/update/client_interceptor/feature.java b/features/update/client_interceptor/feature.java
index 633dd098..e8cd5c98 100644
--- a/features/update/client_interceptor/feature.java
+++ b/features/update/client_interceptor/feature.java
@@ -1,6 +1,7 @@
 package update.client_interceptor;
 
 import io.temporal.client.WorkflowClientOptions;
+import io.temporal.client.UpdateHandle;
 import io.temporal.common.interceptors.WorkflowClientCallsInterceptor;
 import io.temporal.common.interceptors.WorkflowClientCallsInterceptorBase;
 import io.temporal.common.interceptors.WorkflowClientInterceptorBase;
@@ -53,7 +54,7 @@ public WorkflowClientCallsInterceptor workflowClientCallsInterceptor(
                 WorkflowClientCallsInterceptor next) {
               return new WorkflowClientCallsInterceptorBase(next) {
                 @Override
-                public <R> StartUpdateOutput<R> startUpdate(StartUpdateInput<R> input) {
+                public <R> UpdateHandle<R> startUpdate(StartUpdateInput<R> input) {
                   if (input.getUpdateName() == "update") {
                     input.getArguments()[0] = ((int) input.getArguments()[0]) + 1;
                   }