fix(bigquery): address PR google#1344 review — flush/close coordination, HITL pair keys, fail-closed redaction, atomic finalize gate, per-operation tool spans

Resolves all five review findings on the previous commit:

1. BatchProcessor.close() now waits (bounded by the same shutdownTimeout
   deadline) for an in-flight flush to release the flush lock before
   tearing down the writer and Arrow resources, instead of treating an
   empty queue as proof that no flush is active.
2. HITL_*_COMPLETED rows (both the event path and the user-message path)
   now carry pause_kind and function_call_id, so completions are joinable
   to their HITL_*_REQUEST / TOOL_PAUSED rows when multiple HITL requests
   share an invocation.
3. New JsonFormatter.redactTree replaces the smartTruncate-based
   attributes boundary: raw containers are walked natively with keys
   redacted before any Jackson conversion, and an unserializable leaf
   becomes "[UNSERIALIZABLE]" instead of routing the whole map through
   the textual fallback that exposed sibling secrets. JsonNode leaves are
   handled before the Iterable branch (ObjectNode iterates its values).
4. The invocation-finalization gate is now atomic: the isProcessed check
   runs inside computeIfAbsent's mapping function under the per-key lock,
   paired with finalize's mark-before-remove ordering; the residual
   tombstone-cache-expiry window is documented.
5. Tool spans carry an operation identity (ToolContext.functionCallId)
   and a parent captured at push time: ADK runs an event's function calls
   concurrently by default within one branch, so TOOL_STARTING rows are
   stamped from the pushed record, completions/errors pop by identity
   rather than stack position, and parents skip concurrent siblings.

Tests: 160 focused BQAA tests pass (previous commit: 153), full core
suite 1,616 passed / 0 failed. New coverage: close-waits-for-in-flight-
flush, per-leaf fail-closed redaction (including nested containers and
converted POJO leaves), HITL completion pair keys on both paths, and
same-branch concurrent tools with out-of-order identity pops.

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

Author: caohy1988

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

@@ -325,6 +325,26 @@ public void close() {
         }
       }
     }
+    // An in-flight flush may have drained the queue (making the loop above exit immediately) while
+    // still blocked inside its bounded Storage Write append. Wait for it to release the flush lock
+    // within the same deadline, so the writer and Arrow resources are not torn down underneath an
+    // active flush.
+    while (flushLock.get() && Instant.now().isBefore(drainDeadline)) {
+      try {
+        Thread.sleep(10);
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        break;
+      }
+    }
+    if (flushLock.get()) {
+      // The in-flight append has its own shutdownTimeout bound and accounts its batch as
+      // append_error when it fails; resource teardown below may race it, but its catch blocks
+      // handle and count the failure.
+      logger.severe(
+          "Closing while a flush is still in flight; its batch is accounted by the bounded"
+              + " append.");
+    }
     int remaining = this.queue.size();
     if (remaining > 0) {
       droppedShutdownTimeout.addAndGet(remaining);

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

@@ -342,13 +342,12 @@ private Completable logEvent(
       row.put("latency_ms", convertToJsonNode(latencyMap));
     }
     // Redact the complete assembled attributes tree at the output boundary, regardless of which
-    // producer populated it (state deltas, custom tags, labels, extra attributes). Individual
-    // producers may also pre-truncate, but only this final pass guarantees sensitive keys never
-    // reach BigQuery. Integer.MAX_VALUE disables length truncation: this is a redaction-only walk,
-    // and redaction intentionally does not set is_truncated (parity with the Python plugin).
-    row.put(
-        "attributes",
-        smartTruncate(getAttributes(data, invocationContext), Integer.MAX_VALUE).node());
+    // producer populated it (state deltas, custom tags, labels, extra attributes). redactTree
+    // walks raw containers and fails CLOSED on unserializable values (one bad custom tag or
+    // session-state object must not route the whole map through a textual fallback that would
+    // expose sibling secrets). Redaction intentionally does not set is_truncated (parity with the
+    // Python plugin).
+    row.put("attributes", JsonFormatter.redactTree(getAttributes(data, invocationContext)));
 
     CompletableFuture<Void> parseFuture;
     if (content != null) {
@@ -414,6 +413,19 @@ private static String resolveAgentName(
     return eventData.flatMap(EventData::fallbackAgentName).orElse("unknown");
   }
 
+  /**
+   * 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
+   * completion to its HITL_*_REQUEST / TOOL_PAUSED rows.
+   */
+  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));
+    return keys.buildOrThrow();
+  }
+
   @CanIgnoreReturnValue
   private static EventData.Builder withFallbackAgent(
       EventData.Builder builder, @Nullable String author) {
@@ -569,15 +581,19 @@ public Maybe<Content> onUserMessageCallback(
             ImmutableMap.of("tool", responseName, "result", truncatedResult.node());
         if (HITL_EVENT_TYPES.containsKey(responseName)) {
           // HITL completions stay on the HITL_*_COMPLETED stream — they must not also emit
-          // TOOL_COMPLETED.
+          // TOOL_COMPLETED. The pair keys make the completion joinable to its HITL_*_REQUEST /
+          // TOOL_PAUSED rows even when multiple HITL requests share an invocation.
           logCompletable =
               logCompletable.andThen(
                   logEvent(
                       HITL_EVENT_TYPES.get(responseName) + "_COMPLETED",
                       invocationContext,
                       contentMap,
                       truncatedResult.isTruncated(),
-                      Optional.empty()));
+                      Optional.of(
+                          EventData.builder()
+                              .setExtraAttributes(hitlPairKeys(responseName, functionResponse.id()))
+                              .build())));
         } else {
           if (functionResponse.id().isEmpty()) {
             logger.fine(
@@ -683,21 +699,23 @@ public Maybe<Event> onEventCallback(InvocationContext invocationContext, Event e
         }
         if (part.functionResponse().isPresent()
             && HITL_EVENT_TYPES.containsKey(part.functionResponse().get().name().orElse(""))) {
-          String hitlEvent = HITL_EVENT_TYPES.get(part.functionResponse().get().name().get());
+          FunctionResponse hitlResponse = part.functionResponse().get();
+          String hitlEvent = HITL_EVENT_TYPES.get(hitlResponse.name().get());
           TruncationResult truncatedResult =
-              smartTruncate(part.functionResponse().get().response(), config.maxContentLength());
+              smartTruncate(hitlResponse.response(), config.maxContentLength());
           logCompletable =
               logCompletable.andThen(
                   logEvent(
                       hitlEvent + "_COMPLETED",
                       invocationContext,
                       ImmutableMap.of(
-                          "tool",
-                          part.functionResponse().get().name().get(),
-                          "response",
-                          truncatedResult.node()),
+                          "tool", hitlResponse.name().get(), "response", truncatedResult.node()),
                       truncatedResult.isTruncated(),
-                      Optional.empty()));
+                      Optional.of(
+                          EventData.builder()
+                              .setExtraAttributes(
+                                  hitlPairKeys(hitlResponse.name().get(), hitlResponse.id()))
+                              .build())));
         }
       }
     }
@@ -1040,10 +1058,24 @@ public Maybe<Map<String, Object>> beforeToolCallback(
     }
     ImmutableMap<String, Object> contentMap =
         ImmutableMap.of("tool_origin", getToolOrigin(tool), "tool", tool.name(), "args", toolArgs);
-    state
-        .getTraceManager(toolContext.invocationContext().invocationId())
-        .pushSpan(toolContext.invocationContext(), "tool");
-    return logEvent("TOOL_STARTING", toolContext.invocationContext(), contentMap, Optional.empty())
+    // Push with the function-call identity: ADK executes an event's function calls concurrently by
+    // default within one branch, so tool spans are created, stamped, and popped by operation
+    // identity rather than stack position. Stamp the row directly from the pushed record so a
+    // sibling tool pushing in between cannot divert the row's span IDs.
+    TraceManager.SpanRecord toolSpan =
+        state
+            .getTraceManager(toolContext.invocationContext().invocationId())
+            .pushSpanRecord(
+                toolContext.invocationContext(), "tool", toolContext.functionCallId().orElse(null));
+    EventData.Builder startingData = EventData.builder().setSpanIdOverride(toolSpan.spanId());
+    if (toolSpan.parentSpanId() != null) {
+      startingData.setParentSpanIdOverride(toolSpan.parentSpanId());
+    }
+    return logEvent(
+            "TOOL_STARTING",
+            toolContext.invocationContext(),
+            contentMap,
+            Optional.of(startingData.build()))
         .andThen(Maybe.empty());
   }
 
@@ -1061,7 +1093,9 @@ public Maybe<Map<String, Object>> afterToolCallback(
         .ensureInvocationSpan(toolContext.invocationContext());
     TraceManager traceManager =
         state.getTraceManager(toolContext.invocationContext().invocationId());
-    Optional<RecordData> popped = traceManager.popSpan(toolContext.invocationContext(), "tool");
+    Optional<RecordData> popped =
+        traceManager.popSpan(
+            toolContext.invocationContext(), "tool", toolContext.functionCallId().orElse(null));
     TruncationResult truncationResult = smartTruncate(result, config.maxContentLength());
     ImmutableMap<String, Object> contentMap =
         ImmutableMap.of(
@@ -1072,15 +1106,15 @@ public Maybe<Map<String, Object>> afterToolCallback(
             "tool_origin",
             getToolOrigin(tool));
 
-    SpanIds spanIds = traceManager.getCurrentSpanAndParent(toolContext.invocationContext());
-
     EventData.Builder eventDataBuilder = EventData.builder();
     if (popped.isPresent()) {
       eventDataBuilder.setLatency(popped.get().duration());
     }
     // Always record the internal execution-tree span so parent_span_id references a logged row.
+    // The parent comes from the popped record (captured at push time): under concurrent tool
+    // execution the branch's stack top may be a sibling tool, not this span's parent.
     popped.ifPresent(p -> eventDataBuilder.setSpanIdOverride(p.spanId()));
-    spanIds.spanId().ifPresent(eventDataBuilder::setParentSpanIdOverride);
+    popped.flatMap(RecordData::parentSpanId).ifPresent(eventDataBuilder::setParentSpanIdOverride);
 
     return logEvent(
             "TOOL_COMPLETED",
@@ -1102,7 +1136,9 @@ public Maybe<Map<String, Object>> onToolErrorCallback(
         .ensureInvocationSpan(toolContext.invocationContext());
     TraceManager traceManager =
         state.getTraceManager(toolContext.invocationContext().invocationId());
-    Optional<RecordData> popped = traceManager.popSpan(toolContext.invocationContext(), "tool");
+    Optional<RecordData> popped =
+        traceManager.popSpan(
+            toolContext.invocationContext(), "tool", toolContext.functionCallId().orElse(null));
 
     TruncationResult truncationResult = smartTruncate(toolArgs, config.maxContentLength());
     ImmutableMap<String, Object> contentMap =
@@ -1112,16 +1148,16 @@ public Maybe<Map<String, Object>> onToolErrorCallback(
             .put("tool_origin", getToolOrigin(tool))
             .buildOrThrow();
 
-    SpanIds spanIds = traceManager.getCurrentSpanAndParent(toolContext.invocationContext());
-
     EventData.Builder eventDataBuilder =
         EventData.builder().setStatus("ERROR").setErrorMessage(error.getMessage());
     if (popped.isPresent()) {
       eventDataBuilder.setLatency(popped.get().duration());
     }
     // Always record the internal execution-tree span so parent_span_id references a logged row.
+    // The parent comes from the popped record (captured at push time): under concurrent tool
+    // execution the branch's stack top may be a sibling tool, not this span's parent.
     popped.ifPresent(p -> eventDataBuilder.setSpanIdOverride(p.spanId()));
-    spanIds.spanId().ifPresent(eventDataBuilder::setParentSpanIdOverride);
+    popped.flatMap(RecordData::parentSpanId).ifPresent(eventDataBuilder::setParentSpanIdOverride);
 
     return logEvent(
             "TOOL_ERROR",

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

@@ -40,6 +40,7 @@ final class JsonFormatter {
   static final String CYCLE_DETECTED_MESSAGE = "[cycle detected]";
   static final String MAX_DEPTH_MESSAGE = "[max depth exceeded]";
   static final String REDACTED_MESSAGE = "[REDACTED]";
+  static final String UNSERIALIZABLE_MESSAGE = "[UNSERIALIZABLE]";
   // Guard against unbounded recursion on deeply nested (non-cyclic) payloads.
   static final int MAX_TRUNCATE_DEPTH = 200;
 
@@ -86,6 +87,83 @@ static TruncationResult smartTruncate(Object obj, int maxLength) {
     }
   }
 
+  /**
+   * Redacts sensitive keys across an attributes tree, failing closed on unserializable values.
+   *
+   * <p>Unlike {@link #smartTruncate}, which converts the whole object to JSON first (so one
+   * unsupported value routes the ENTIRE tree through the textual {@code safeToString} fallback,
+   * exposing sibling secrets as plain text), this walks raw Java containers natively: keys are
+   * redacted before any Jackson conversion, and only leaf values are converted individually. A leaf
+   * that cannot be converted becomes {@value #UNSERIALIZABLE_MESSAGE} without affecting its
+   * siblings. No length truncation is applied.
+   */
+  static JsonNode redactTree(Object obj) {
+    return redactTreeInternal(obj, newSetFromMap(new IdentityHashMap<>()), 0);
+  }
+
+  private static JsonNode redactTreeInternal(Object obj, Set<Object> visited, int depth) {
+    if (obj == null) {
+      return mapper.nullNode();
+    }
+    if (depth > MAX_TRUNCATE_DEPTH) {
+      return mapper.valueToTree(MAX_DEPTH_MESSAGE);
+    }
+    // JsonNode must be handled before the Iterable branch: ObjectNode implements
+    // Iterable<JsonNode> over its VALUES, so the generic Iterable walk would flatten a JSON
+    // object into an array and lose its keys.
+    if (obj instanceof JsonNode jsonNode) {
+      return recursiveSmartTruncate(
+              jsonNode, Integer.MAX_VALUE, newSetFromMap(new IdentityHashMap<>()), depth)
+          .node();
+    }
+    if (obj instanceof Map<?, ?> map) {
+      if (!visited.add(obj)) {
+        return mapper.valueToTree(CYCLE_DETECTED_MESSAGE);
+      }
+      try {
+        ObjectNode node = mapper.createObjectNode();
+        for (Map.Entry<?, ?> entry : map.entrySet()) {
+          String key = String.valueOf(entry.getKey());
+          if (isSensitiveKey(key)) {
+            node.set(key, mapper.valueToTree(REDACTED_MESSAGE));
+            continue;
+          }
+          node.set(key, redactTreeInternal(entry.getValue(), visited, depth + 1));
+        }
+        return node;
+      } finally {
+        visited.remove(obj);
+      }
+    }
+    if (obj instanceof Iterable<?> iterable) {
+      if (!visited.add(obj)) {
+        return mapper.valueToTree(CYCLE_DETECTED_MESSAGE);
+      }
+      try {
+        ArrayNode node = mapper.createArrayNode();
+        for (Object element : iterable) {
+          node.add(redactTreeInternal(element, visited, depth + 1));
+        }
+        return node;
+      } finally {
+        visited.remove(obj);
+      }
+    }
+    try {
+      // A converted leaf may itself be a container (e.g. a POJO serialized to an object): run the
+      // JSON-level redacting walk over it with truncation disabled.
+      return recursiveSmartTruncate(
+              mapper.valueToTree(obj),
+              Integer.MAX_VALUE,
+              newSetFromMap(new IdentityHashMap<>()),
+              depth)
+          .node();
+    } catch (IllegalArgumentException e) {
+      logger.fine("redactTree replacing unserializable value: " + e.getMessage());
+      return mapper.valueToTree(UNSERIALIZABLE_MESSAGE);
+    }
+  }
+
   static JsonNode convertToJsonNode(Object obj) {
     if (obj == null) {
       return mapper.nullNode();

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

@@ -223,15 +223,9 @@ BatchProcessor getBatchProcessor(String invocationId) {
    * </ul>
    */
   void appendRow(String invocationId, Map<String, Object> row) {
-    if (isProcessed(invocationId)) {
-      droppedAfterFinalize.incrementAndGet();
-      logger.warning(
-          "Dropping late analytics row: invocation " + invocationId + " is already finalized.");
-      return;
-    }
     BatchProcessor processor;
     try {
-      processor = getBatchProcessor(invocationId);
+      processor = getOrCreateProcessorIfActive(invocationId);
     } catch (RuntimeException e) {
       droppedWriterCreateError.incrementAndGet();
       logger.log(
@@ -240,9 +234,52 @@ void appendRow(String invocationId, Map<String, Object> row) {
           e);
       return;
     }
+    if (processor == null) {
+      droppedAfterFinalize.incrementAndGet();
+      logger.warning(
+          "Dropping late analytics row: invocation " + invocationId + " is already finalized.");
+      return;
+    }
     processor.append(row);
   }
 
+  /**
+   * Atomically returns the invocation's processor, creating it only while the invocation is still
+   * active.
+   *
+   * <p>The {@code isProcessed} check runs INSIDE the {@code computeIfAbsent} mapping function, i.e.
+   * under the map's per-key lock, closing the check-then-act race with {@code
+   * ensureInvocationCompleted}: finalization marks the invocation processed BEFORE removing the
+   * processor mapping, so a continuation that finds no mapping after removal is guaranteed to
+   * observe {@code isProcessed == true} and installs nothing (a null mapping-function result adds
+   * no entry). If a continuation instead wins the key lock first and creates the processor,
+   * finalization subsequently removes and closes that same processor, and the late row is accounted
+   * by {@link BatchProcessor#append}'s closed gate.
+   *
+   * <p>Residual limitation: the processed-invocations tombstone cache is bounded (size and TTL), so
+   * a continuation arriving after its tombstone was evicted could still recreate a processor; such
+   * processors are drained and closed at plugin-wide shutdown.
+   */
+  private @Nullable BatchProcessor getOrCreateProcessorIfActive(String invocationId) {
+    return batchProcessors.computeIfAbsent(
+        invocationId,
+        id -> {
+          if (isProcessed(id)) {
+            return null;
+          }
+          BatchProcessor p =
+              new BatchProcessor(
+                  createWriter(),
+                  config.batchSize(),
+                  config.batchFlushInterval(),
+                  config.queueMaxSize(),
+                  executor,
+                  config.shutdownTimeout());
+          p.start();
+          return p;
+        });
+  }
+
   protected @Nullable GcsOffloader getGcsOffloader(BigQueryLoggerConfig config) {
     if (config.gcsBucketName().isEmpty()) {
       return null;