fix(bigquery): address round-4 review — shared shutdown deadline, collision-proof tool identity, HITL content parity

Resolves both round-4 P1 findings and the P2 on PR google#1344:

1. One absolute deadline per shutdown operation: ensureInvocationCompleted
   computes its deadline at subscription (Completable.defer) and shares it
   between the pending-task wait and the processor drain; plugin-wide
   close() shares one deadline across the task wait, every processor's
   drain, and both executor terminations. BatchProcessor gains
   deadline-accepting closeAndFold/close variants so each drain consumes
   only the remaining budget — a stuck parse plus a queued processor, or
   N stuck processors, are all bounded by a single shutdownTimeout.
2. Collision-proof tool operation identity: the framework materializes an
   absent function-call ID as "", so two concurrent id-less calls
   collided and cross-popped. toolOperationId now uses the real ID when
   non-empty, else a synthetic ID keyed by ToolContext instance identity
   (weak-keys cache; the same instance flows through before/after/error
   callbacks of one call). hitlPairKeys also filters empty-string IDs,
   which are useless as join keys.
3. HITL completion content normalized to "result" on the event path,
   matching the user-message path and the Python plugin, so one event
   type has one queryable content shape.

Tests: concurrent id-less tools keep span ownership out of completion
order (plugin-level, span-id pairing per tool); multi-batch drain bound
tightened to <900ms so the old two-deadline behavior fails; new
pending-task-plus-drain and multi-stuck-processor cases assert the
shared bound; HITL event-path completion asserts content.result.
Focused BQAA suite 168 passed (previous: 165); full core 1,624 passed /
0 failed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: caohy1988

core/src/main/java/com/google/adk/plugins/agentanalytics/BatchProcessor.java

@@ -338,12 +338,26 @@ ImmutableMap<String, Long> getDropStats() {
    * failure) are included in the delivered snapshot exactly once.
    */
   void closeAndFold(Consumer<ImmutableMap<String, Long>> statsConsumer) {
+    closeAndFold(statsConsumer, Instant.now().plus(shutdownTimeout));
+  }
+
+  /**
+   * Deadline-accepting variant: the caller passes ONE absolute deadline shared across a larger
+   * shutdown operation (pending-task waits, sibling processors, executor termination), so this
+   * processor's drain consumes only the remaining budget instead of restarting a fresh
+   * shutdownTimeout.
+   */
+  void closeAndFold(Consumer<ImmutableMap<String, Long>> statsConsumer, Instant deadline) {
     this.onFinalStats = statsConsumer;
-    close();
+    close(deadline);
   }
 
   @Override
   public void close() {
+    close(Instant.now().plus(shutdownTimeout));
+  }
+
+  private void close(Instant drainDeadline) {
     // Idempotent: ensureInvocationCompleted and plugin-wide shutdown may both close a processor.
     if (!closed.compareAndSet(false, true)) {
       return;
@@ -354,9 +368,8 @@ public void close() {
     if (task != null) {
       task.cancel(false);
     }
-    // Final drain, bounded by shutdownTimeout rather than looping until empty. Publishing the
-    // deadline caps every drain/in-flight append to the remaining close budget.
-    Instant drainDeadline = Instant.now().plus(shutdownTimeout);
+    // Final drain, bounded by the caller's absolute deadline rather than looping until empty.
+    // Publishing the deadline caps every drain/in-flight append to the remaining close budget.
     this.closeDeadline = drainDeadline;
     while (!this.queue.isEmpty() && Instant.now().isBefore(drainDeadline)) {
       this.flush();

core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java

@@ -53,6 +53,8 @@
 import com.google.cloud.bigquery.TableInfo;
 import com.google.cloud.bigquery.TimePartitioning;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
@@ -74,6 +76,7 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
@@ -418,6 +421,29 @@ private static String resolveAgentName(
     return eventData.flatMap(EventData::fallbackAgentName).orElse("unknown");
   }
 
+  // Synthetic operation identities for tool calls whose function-call ID is absent: the framework
+  // materializes ToolContext.functionCallId as "" when the model omitted the ID, so two concurrent
+  // id-less calls would collide on "" and cross-pop each other's spans. The same ToolContext
+  // instance flows through the before/after/error callbacks of one call, so it keys a unique
+  // synthetic ID; weakKeys gives identity semantics plus GC-based cleanup for calls whose
+  // completion callback never fires.
+  private final Cache<ToolContext, String> syntheticToolCallIds =
+      CacheBuilder.newBuilder().weakKeys().build();
+
+  /**
+   * Operation identity for a tool call's span: the real function-call ID when present and
+   * non-empty, else a per-ToolContext synthetic ID (see {@link #syntheticToolCallIds}).
+   */
+  private String toolOperationId(ToolContext toolContext) {
+    String id = toolContext.functionCallId().orElse("");
+    if (!id.isEmpty()) {
+      return id;
+    }
+    return syntheticToolCallIds
+        .asMap()
+        .computeIfAbsent(toolContext, tc -> "tool-ctx-" + UUID.randomUUID());
+  }
+
   /**
    * Pause/resume pair keys for a HITL completion row: {@code pause_kind} derived from the synthetic
    * call name and, when the response carries one, the {@code function_call_id} that joins the
@@ -427,7 +453,8 @@ private static ImmutableMap<String, Object> hitlPairKeys(
       String hitlName, Optional<String> functionCallId) {
     ImmutableMap.Builder<String, Object> keys = ImmutableMap.builder();
     keys.put("pause_kind", HITL_PAUSE_KIND_MAP.getOrDefault(hitlName, "tool"));
-    functionCallId.ifPresent(id -> keys.put("function_call_id", id));
+    // The framework materializes an absent ID as "", which is useless as a join key.
+    functionCallId.filter(id -> !id.isEmpty()).ifPresent(id -> keys.put("function_call_id", id));
     return keys.buildOrThrow();
   }
 
@@ -719,8 +746,10 @@ public Maybe<Event> onEventCallback(InvocationContext invocationContext, Event e
                   logEvent(
                       hitlEvent + "_COMPLETED",
                       invocationContext,
+                      // "result" matches the Python plugin's HITL completion content on BOTH
+                      // producer paths, so one event type has one queryable content shape.
                       ImmutableMap.of(
-                          "tool", hitlResponse.name().get(), "response", truncatedResult.node()),
+                          "tool", hitlResponse.name().get(), "result", truncatedResult.node()),
                       truncatedResult.isTruncated(),
                       Optional.of(
                           EventData.builder()
@@ -1076,8 +1105,7 @@ public Maybe<Map<String, Object>> beforeToolCallback(
     TraceManager.SpanRecord toolSpan =
         state
             .getTraceManager(toolContext.invocationContext().invocationId())
-            .pushSpanRecord(
-                toolContext.invocationContext(), "tool", toolContext.functionCallId().orElse(null));
+            .pushSpanRecord(toolContext.invocationContext(), "tool", toolOperationId(toolContext));
     EventData.Builder startingData = EventData.builder().setSpanIdOverride(toolSpan.spanId());
     if (toolSpan.parentSpanId() != null) {
       startingData.setParentSpanIdOverride(toolSpan.parentSpanId());
@@ -1105,8 +1133,7 @@ public Maybe<Map<String, Object>> afterToolCallback(
     TraceManager traceManager =
         state.getTraceManager(toolContext.invocationContext().invocationId());
     Optional<RecordData> popped =
-        traceManager.popSpan(
-            toolContext.invocationContext(), "tool", toolContext.functionCallId().orElse(null));
+        traceManager.popSpan(toolContext.invocationContext(), "tool", toolOperationId(toolContext));
     TruncationResult truncationResult = smartTruncate(result, config.maxContentLength());
     ImmutableMap<String, Object> contentMap =
         ImmutableMap.of(
@@ -1148,8 +1175,7 @@ public Maybe<Map<String, Object>> onToolErrorCallback(
     TraceManager traceManager =
         state.getTraceManager(toolContext.invocationContext().invocationId());
     Optional<RecordData> popped =
-        traceManager.popSpan(
-            toolContext.invocationContext(), "tool", toolContext.functionCallId().orElse(null));
+        traceManager.popSpan(toolContext.invocationContext(), "tool", toolOperationId(toolContext));
 
     TruncationResult truncationResult = smartTruncate(toolArgs, config.maxContentLength());
     ImmutableMap<String, Object> contentMap =

core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java

@@ -52,7 +52,6 @@
 import java.util.logging.Logger;
 import org.jspecify.annotations.Nullable;
 import org.threeten.bp.Duration;
-import org.threeten.bp.Instant;
 
 /** Manages state for the BigQueryAgentAnalyticsPlugin. */
 class PluginState {
@@ -402,6 +401,19 @@ void addPendingTask(String invocationId, CompletableFuture<Void> task) {
   }
 
   Completable ensureInvocationCompleted(String invocationId) {
+    // ONE absolute deadline for the whole finalization: waiting for pending tasks and draining
+    // the processor share it, so shutdownTimeout is the total bound rather than restarting per
+    // phase (a stuck parse consuming one full timeout must not grant the drain another).
+    // Deferred so the budget starts at subscription, not assembly.
+    return Completable.defer(
+        () -> {
+          java.time.Instant finalizeDeadline =
+              java.time.Instant.now().plus(config.shutdownTimeout());
+          return finalizeInvocation(invocationId, finalizeDeadline);
+        });
+  }
+
+  private Completable finalizeInvocation(String invocationId, java.time.Instant finalizeDeadline) {
     Set<CompletableFuture<Void>> tasks = pendingTasks.get(invocationId);
     Completable tasksState = Completable.complete();
     if (tasks != null && !tasks.isEmpty()) {
@@ -437,13 +449,13 @@ Completable ensureInvocationCompleted(String invocationId) {
               markProcessed(invocationId);
               BatchProcessor processor = removeProcessor(invocationId);
               if (processor != null) {
-                // closeAndFold owns the COMPLETE drain under one shutdownTimeout deadline. An
-                // explicit flush() here would run before the close deadline is published and
-                // consume a separate full append budget, doubling the effective bound. Folding
-                // happens via the teardown callback, which fires when teardown ACTUALLY completes
-                // (possibly after close() returns, if an in-flight flush owns the resources past
-                // the deadline), so counters recorded by that last flush are never lost.
-                processor.closeAndFold(this::foldStats);
+                // closeAndFold drains under the SAME absolute deadline the pending-task wait
+                // consumed from, so the total finalization is bounded by one shutdownTimeout.
+                // Folding happens via the teardown callback, which fires when teardown ACTUALLY
+                // completes (possibly after close() returns, if an in-flight flush owns the
+                // resources past the deadline), so counters recorded by that last flush are never
+                // lost.
+                processor.closeAndFold(this::foldStats, finalizeDeadline);
               }
               TraceManager traceManager = removeTraceManager(invocationId);
               if (traceManager != null) {
@@ -494,6 +506,15 @@ ImmutableMap<String, Long> getDropStats() {
   }
 
   Completable close() {
+    // ONE absolute deadline for the whole plugin shutdown: the pending-task wait, every
+    // processor's drain, and executor termination all consume from the same shutdownTimeout
+    // budget, so total shutdown is bounded by one timeout rather than one per phase/processor.
+    // Deferred so the budget starts at subscription, not assembly.
+    return Completable.defer(
+        () -> closeInternal(java.time.Instant.now().plus(config.shutdownTimeout())));
+  }
+
+  private Completable closeInternal(java.time.Instant closeDeadline) {
     ImmutableList<CompletableFuture<Void>> tasks =
         pendingTasks.values().stream().flatMap(Set::stream).collect(toImmutableList());
     Completable tasksState = Completable.complete();
@@ -519,8 +540,9 @@ Completable close() {
               }
               lifecycles.clear();
               for (BatchProcessor processor : getBatchProcessors()) {
-                // Fold via the teardown callback; see ensureInvocationCompleted.
-                processor.closeAndFold(this::foldStats);
+                // Fold via the teardown callback; each drain consumes only the REMAINING shared
+                // budget, so N processors cannot take N timeouts.
+                processor.closeAndFold(this::foldStats, closeDeadline);
               }
               for (TraceManager traceManager : getTraceManagers()) {
                 traceManager.clearStack();
@@ -538,13 +560,14 @@ Completable close() {
               try {
                 executor.shutdown();
                 offloadExecutor.shutdown();
-                long totalTimeoutMillis = config.shutdownTimeout().toMillis();
-                Instant startTime = Instant.now();
-                if (!executor.awaitTermination(totalTimeoutMillis, MILLISECONDS)) {
+                long remainingMillis =
+                    java.time.Duration.between(java.time.Instant.now(), closeDeadline).toMillis();
+                if (remainingMillis <= 0
+                    || !executor.awaitTermination(remainingMillis, MILLISECONDS)) {
                   executor.shutdownNow();
                 }
-                long elapsedTimeMillis = Duration.between(startTime, Instant.now()).toMillis();
-                long remainingMillis = totalTimeoutMillis - elapsedTimeMillis;
+                remainingMillis =
+                    java.time.Duration.between(java.time.Instant.now(), closeDeadline).toMillis();
                 if (remainingMillis > 0) {
                   if (!offloadExecutor.awaitTermination(remainingMillis, MILLISECONDS)) {
                     offloadExecutor.shutdownNow();

core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java

@@ -1037,14 +1037,22 @@ public void afterToolCallback_populatesCorrectFields() throws Exception {
     ImmutableMap<String, Object> toolArgs = ImmutableMap.of("arg1", "value1");
     ImmutableMap<String, Object> result = ImmutableMap.of("res1", "value2");
 
-    // Establish the invocation before pushing the tool span, so afterToolCallback's
-    // ensureInvocationSpan does not clear the pushed span as stale state from a prior invocation.
+    // Mirror the production flow: beforeToolCallback pushes the tool span with the SAME
+    // operation identity (from the ToolContext) that afterToolCallback pops with.
     state.getTraceManager("invocation_id").ensureInvocationSpan(mockInvocationContext);
-    state.getTraceManager("invocation_id").pushSpan(mockInvocationContext, "tool_request");
+    plugin.beforeToolCallback(mockTool, toolArgs, mockToolContext).blockingSubscribe();
     plugin.afterToolCallback(mockTool, toolArgs, mockToolContext, result).blockingSubscribe();
 
-    Map<String, Object> row = state.getBatchProcessor("invocation_id").queue.poll();
-    assertNotNull("Row not found in queue", row);
+    CompletableFuture.allOf(
+            state
+                .getPendingTasksForInvocation("invocation_id")
+                .toArray(new CompletableFuture<?>[0]))
+        .join();
+    Map<String, Object> row;
+    do {
+      row = state.getBatchProcessor("invocation_id").queue.poll();
+      assertNotNull("TOOL_COMPLETED row not found in queue", row);
+    } while (!"TOOL_COMPLETED".equals(row.get("event_type")));
     assertEquals("TOOL_COMPLETED", row.get("event_type"));
     assertEquals("agent_name", row.get("agent"));
     ObjectNode contentMap = (ObjectNode) row.get("content");
@@ -2034,6 +2042,61 @@ public void onEventCallback_hitlFunctionResponse_completionCarriesPairKeys() thr
     JsonNode attributes = (JsonNode) completed.get("attributes");
     assertEquals("hitl_confirmation", attributes.get("pause_kind").asText());
     assertEquals("fc-7", attributes.get("function_call_id").asText());
+    // Content key parity: both HITL completion producer paths (event and user-message) use
+    // "result", matching the Python plugin, so one event type has one queryable content shape.
+    JsonNode content = (JsonNode) completed.get("content");
+    assertNotNull("content.result must be present on the event path", content.get("result"));
+  }
+
+  @Test
+  public void concurrentIdLessTools_keepSpanOwnership() throws Exception {
+    // The framework materializes an absent function-call ID as "" — two concurrent id-less calls
+    // must not collide on it and cross-pop each other's spans.
+    BaseTool toolA = mock(BaseTool.class);
+    when(toolA.name()).thenReturn("tool_a");
+    BaseTool toolB = mock(BaseTool.class);
+    when(toolB.name()).thenReturn("tool_b");
+    ToolContext contextA = mock(ToolContext.class);
+    when(contextA.invocationContext()).thenReturn(mockInvocationContext);
+    when(contextA.functionCallId()).thenReturn(Optional.of(""));
+    ToolContext contextB = mock(ToolContext.class);
+    when(contextB.invocationContext()).thenReturn(mockInvocationContext);
+    when(contextB.functionCallId()).thenReturn(Optional.of(""));
+
+    state.getTraceManager("invocation_id").ensureInvocationSpan(mockInvocationContext);
+    plugin.beforeToolCallback(toolA, ImmutableMap.of(), contextA).blockingSubscribe();
+    plugin.beforeToolCallback(toolB, ImmutableMap.of(), contextB).blockingSubscribe();
+    // A completes FIRST even though B's record sits above it.
+    plugin
+        .afterToolCallback(toolA, ImmutableMap.of(), contextA, ImmutableMap.of())
+        .blockingSubscribe();
+    plugin
+        .afterToolCallback(toolB, ImmutableMap.of(), contextB, ImmutableMap.of())
+        .blockingSubscribe();
+
+    CompletableFuture.allOf(
+            state
+                .getPendingTasksForInvocation("invocation_id")
+                .toArray(new CompletableFuture<?>[0]))
+        .join();
+    Map<String, String> startingSpanByTool = new HashMap<>();
+    Map<String, String> completedSpanByTool = new HashMap<>();
+    Map<String, Object> row;
+    while ((row = state.getBatchProcessor("invocation_id").queue.poll()) != null) {
+      String tool = ((JsonNode) row.get("content")).get("tool").asText();
+      if ("TOOL_STARTING".equals(row.get("event_type"))) {
+        startingSpanByTool.put(tool, (String) row.get("span_id"));
+      } else if ("TOOL_COMPLETED".equals(row.get("event_type"))) {
+        completedSpanByTool.put(tool, (String) row.get("span_id"));
+      }
+    }
+
+    // Each tool's completion row references ITS OWN starting span, not the sibling's.
+    assertEquals(startingSpanByTool.get("tool_a"), completedSpanByTool.get("tool_a"));
+    assertEquals(startingSpanByTool.get("tool_b"), completedSpanByTool.get("tool_b"));
+    assertFalse(
+        "sibling id-less tools must not share a span",
+        startingSpanByTool.get("tool_a").equals(startingSpanByTool.get("tool_b")));
   }
 
   @Test