feat: Add onRunErrorCallback to ADK Plugin and Runner

PiperOrigin-RevId: 963992173

Author: google-genai-bot

core/src/main/java/com/google/adk/plugins/Plugin.java

@@ -87,6 +87,16 @@ default Completable afterRunCallback(InvocationContext invocationContext) {
     return Completable.complete();
   }
 
+  /**
+   * Callback executed when a run encounters an error.
+   *
+   * @param invocationContext The context for the entire invocation.
+   * @param error The exception that was raised.
+   */
+  default Completable onRunErrorCallback(InvocationContext invocationContext, Throwable error) {
+    return Completable.complete();
+  }
+
   /**
    * Method executed when the runner is closed.
    *

core/src/main/java/com/google/adk/plugins/PluginManager.java

@@ -143,6 +143,27 @@ public Completable afterRunCallback(InvocationContext invocationContext) {
         .compose(Tracing.withContext(capturedContext));
   }
 
+  public Completable runOnRunErrorCallback(InvocationContext invocationContext, Throwable error) {
+    return onRunErrorCallback(invocationContext, error);
+  }
+
+  @Override
+  public Completable onRunErrorCallback(InvocationContext invocationContext, Throwable error) {
+    Context capturedContext = Context.current();
+    return Flowable.fromIterable(plugins)
+        .concatMapCompletable(
+            plugin ->
+                plugin
+                    .onRunErrorCallback(invocationContext, error)
+                    .doOnError(
+                        e ->
+                            logger.error(
+                                "[{}] Error during callback 'onRunErrorCallback'",
+                                plugin.getName(),
+                                e)))
+        .compose(Tracing.withContext(capturedContext));
+  }
+
   @Override
   public Completable close() {
     Context capturedContext = Context.current();

core/src/main/java/com/google/adk/runner/Runner.java

@@ -557,7 +557,13 @@ protected Flowable<Event> runAsyncImpl(
                   .flatMapPublisher(
                       event ->
                           runAgentWithUpdatedSession(initialContext, session, event, rootAgent)
-                              .compose(Tracing.<Event>withContext(capturedContext)));
+                              .compose(Tracing.<Event>withContext(capturedContext)))
+                  .doOnError(
+                      throwable ->
+                          this.pluginManager
+                              .runOnRunErrorCallback(initialContext, throwable)
+                              .onErrorComplete()
+                              .subscribe());
             })
         .doOnError(
             throwable -> {
@@ -802,6 +808,10 @@ protected Flowable<Event> runLiveImpl(
                     Span span = Span.current();
                     span.setStatus(StatusCode.ERROR, "Error in runLive Flowable execution");
                     span.recordException(throwable);
+                    this.pluginManager
+                        .runOnRunErrorCallback(invocationContext, throwable)
+                        .onErrorComplete()
+                        .subscribe();
                   })
               .compose(Tracing.<Event>withContext(capturedContext));
         });

core/src/test/java/com/google/adk/runner/RunnerTest.java

@@ -31,6 +31,7 @@
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.CALLS_REAL_METHODS;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -370,6 +371,33 @@ public void afterRunCallback_error() {
     verify(plugin).afterRunCallback(any());
   }
 
+  @Test
+  public void onRunErrorCallback_isCalled() {
+    Exception exception = new Exception("test run error");
+    TestLlm failingTestLlm = createTestLlm(Flowable.error(exception));
+    LlmAgent failingAgent = createTestAgentBuilder(failingTestLlm).build();
+
+    Runner failingRunner =
+        Runner.builder()
+            .app(
+                App.builder()
+                    .name("test")
+                    .rootAgent(failingAgent)
+                    .plugins(ImmutableList.of(plugin))
+                    .build())
+            .sessionService(this.runner.sessionService())
+            .build();
+
+    when(plugin.onRunErrorCallback(any(), any())).thenReturn(Completable.complete());
+
+    failingRunner
+        .runAsync("user", session.id(), createContent("from user"))
+        .test()
+        .assertError(exception);
+
+    verify(plugin).onRunErrorCallback(any(), eq(exception));
+  }
+
   @Test
   public void onUserMessageCallback_success() {
     when(plugin.onUserMessageCallback(any(), any())).thenReturn(Maybe.just(pluginContent));