fix(bigquery): address round-3 review — atomic append admission, single close-owned drain deadline

Resolves both round-3 P1 findings on PR google#1344:

1. Append admission is now atomic with lifecycle finalization:
   InvocationLifecycle.runIfActive executes the processor append inside
   the token's monitor, shared with markFinalized. A continuation that
   passed the processor-map gate can no longer get descheduled across
   finalization and then offer into a drained queue or record an
   after_close drop that the already-delivered final snapshot misses;
   finalization now waits for any in-flight admission (a short,
   non-blocking-offer critical section), and post-finalization
   admissions are refused and counted at the plugin level.
2. The explicit pre-close flush in invocation finalization is removed:
   it ran before close() published its deadline and consumed a separate
   full append budget, doubling the effective drain bound with multiple
   queued batches. closeAndFold now owns the complete drain under one
   shutdownTimeout deadline.

Cleanup per review: the stale tombstone-residual javadoc on
getOrCreateProcessorIfActive now describes the durable-token guarantee,
and the TraceManager class javadoc no longer claims within-branch
execution is sequential (tools run concurrently and use operation-
identity spans).

Tests: finalization blocks on an in-flight admission and refuses the
next one (latch-based interleaving), and a two-batch never-completing
drain finishes within a single shutdownTimeout bound. Focused BQAA
suite 165 passed (previous: 163); full core 1,621 passed / 0 failed.

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

Author: caohy1988

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

@@ -89,15 +89,36 @@ class PluginState {
   private final ConcurrentHashMap<String, InvocationLifecycle> lifecycles =
       new ConcurrentHashMap<>();
 
-  /** Terminal-state token for one invocation; see {@link #lifecycles}. */
+  /**
+   * Terminal-state token for one invocation; see {@link #lifecycles}.
+   *
+   * <p>Admission and finalization share this token's monitor: {@link #runIfActive} executes an
+   * append admission atomically against {@link #markFinalized}, so a continuation cannot pass the
+   * gate, get descheduled across finalization (processor close + final stats delivery), and then
+   * append into a torn-down processor or record a loss the final snapshot has already missed.
+   * Critical sections are short: the append side only holds the monitor through a non-blocking
+   * queue offer.
+   */
   static final class InvocationLifecycle {
-    private volatile boolean finalized;
+    private boolean finalized;
+
+    /**
+     * Runs {@code action} iff the invocation is not finalized, atomically with {@link
+     * #markFinalized}. Returns whether the action ran.
+     */
+    synchronized boolean runIfActive(Runnable action) {
+      if (finalized) {
+        return false;
+      }
+      action.run();
+      return true;
+    }
 
-    void markFinalized() {
+    synchronized void markFinalized() {
       finalized = true;
     }
 
-    boolean isFinalized() {
+    synchronized boolean isFinalized() {
       return finalized;
     }
   }
@@ -269,7 +290,18 @@ void appendRow(InvocationLifecycle lifecycle, String invocationId, Map<String, O
           "Dropping late analytics row: invocation " + invocationId + " is already finalized.");
       return;
     }
-    processor.append(row);
+    // Admit the row atomically with finalization: without this, a continuation could pass the
+    // gate above, get descheduled while finalization closes the processor and delivers its final
+    // stats, and then either offer into a drained queue or record an after_close drop the folded
+    // snapshot has already missed.
+    boolean admitted = lifecycle.runIfActive(() -> processor.append(row));
+    if (!admitted) {
+      droppedAfterFinalize.incrementAndGet();
+      logger.warning(
+          "Dropping late analytics row: invocation "
+              + invocationId
+              + " finalized during admission.");
+    }
   }
 
   /**
@@ -285,9 +317,9 @@ void appendRow(InvocationLifecycle lifecycle, String invocationId, Map<String, O
    * 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.
+   * <p>Cache eviction is not a correctness hole: every continuation checks its captured durable
+   * {@link InvocationLifecycle} token first, which — unlike the bounded tombstone cache — cannot be
+   * evicted, so a finalized invocation's processor cannot be resurrected regardless of cache state.
    */
   private @Nullable BatchProcessor getOrCreateProcessorIfActive(
       InvocationLifecycle lifecycle, String invocationId) {
@@ -405,8 +437,10 @@ Completable ensureInvocationCompleted(String invocationId) {
               markProcessed(invocationId);
               BatchProcessor processor = removeProcessor(invocationId);
               if (processor != null) {
-                processor.flush();
-                // Fold via the teardown callback: it fires when teardown ACTUALLY completes
+                // 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);

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

@@ -44,12 +44,15 @@
  * Ambient OpenTelemetry context is still consulted for the {@code trace_id} (and the invocation
  * root's {@code span_id}) so BigQuery rows stay joinable to Cloud Trace.
  *
- * <p>Span records are kept in per-branch stacks keyed by {@link InvocationContext#branch()}. Within
- * one branch execution is sequential, so per-branch LIFO is sound; concurrently scheduled {@code
- * ParallelAgent} branches (which share an invocation ID but carry distinct branch strings) never
- * touch each other's stacks, so a branch completing first can no longer pop another branch's span.
- * Pops additionally verify the record's {@code kind}, so an error callback firing without its
- * matching push cannot pop an unrelated record.
+ * <p>Span records are kept in per-branch stacks keyed by {@link InvocationContext#branch()}.
+ * Concurrently scheduled {@code ParallelAgent} branches (which share an invocation ID but carry
+ * distinct branch strings) never touch each other's stacks, so a branch completing first can no
+ * longer pop another branch's span. Within one branch, agent and model spans execute sequentially
+ * and use top-of-stack semantics, but ADK executes an event's function calls CONCURRENTLY by
+ * default: tool spans therefore carry an operation identity (the function-call ID) plus a parent
+ * captured at push time, and are popped by identity rather than stack position. Pops additionally
+ * verify the record's {@code kind}, so an error callback firing without its matching push cannot
+ * pop an unrelated record.
  */
 public final class TraceManager {
   private static final Logger logger = Logger.getLogger(TraceManager.class.getName());

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

@@ -408,4 +408,73 @@ public void manyCompletedInvocations_leaveNoRetainedProcessorsOrTraceManagers()
     assertTrue(pluginState.getBatchProcessors().isEmpty());
     assertTrue(pluginState.getTraceManagers().isEmpty());
   }
+
+  @Test
+  public void appendRow_admissionIsAtomicWithFinalization() throws Exception {
+    String invocationId = "inv-atomic";
+    PluginState.InvocationLifecycle lifecycle = pluginState.getLifecycle(invocationId);
+
+    // Model a continuation inside the admission critical section (post token-gate, mid-append):
+    // finalization must block on the token monitor until the admission completes, so the admitted
+    // row is drained by close rather than stranded or double-counted after the final snapshot.
+    java.util.concurrent.CountDownLatch inAdmission = new java.util.concurrent.CountDownLatch(1);
+    Thread admitting =
+        new Thread(
+            () ->
+                lifecycle.runIfActive(
+                    () -> {
+                      inAdmission.countDown();
+                      Uninterruptibles.sleepUninterruptibly(java.time.Duration.ofMillis(300));
+                    }));
+    admitting.start();
+    assertTrue(inAdmission.await(2, TimeUnit.SECONDS));
+
+    long start = System.nanoTime();
+    pluginState.ensureInvocationCompleted(invocationId).blockingAwait();
+    long elapsedMs = (System.nanoTime() - start) / 1_000_000;
+    admitting.join(2000);
+
+    assertTrue(
+        "finalization must wait for the in-flight admission, took " + elapsedMs + "ms",
+        elapsedMs >= 250);
+
+    // A continuation arriving after finalization is refused atomically and accounted.
+    pluginState.appendRow(lifecycle, invocationId, ImmutableMap.of("event_type", "LLM_REQUEST"));
+    assertEquals(1L, (long) pluginState.getDropStats().get("late_after_finalize"));
+  }
+
+  @Test
+  public void ensureInvocationCompleted_multiBatchDrain_boundedByOneShutdownTimeout()
+      throws IOException {
+    // Every queued batch must drain under ONE close-owned deadline; a pre-close flush would grant
+    // the first batch a separate full append budget, doubling the effective bound.
+    config = config.toBuilder().shutdownTimeout(Duration.ofMillis(500)).build();
+    TestPluginState slowState =
+        new TestPluginState(config) {
+          @Override
+          protected StreamWriter createWriter() {
+            StreamWriter writer = mock(StreamWriter.class);
+            // Appends never complete in time; each get() must be capped by the REMAINING budget.
+            when(writer.append(any(ArrowRecordBatch.class)))
+                .thenReturn(com.google.api.core.SettableApiFuture.create());
+            return writer;
+          }
+        };
+    String invocationId = "inv-multibatch";
+    BatchProcessor processor = slowState.getBatchProcessor(invocationId);
+    java.util.Map<String, Object> row1 = new java.util.HashMap<>();
+    row1.put("event_type", "A");
+    java.util.Map<String, Object> row2 = new java.util.HashMap<>();
+    row2.put("event_type", "B");
+    processor.queue.offer(row1);
+    processor.queue.offer(row2);
+
+    long start = System.nanoTime();
+    slowState.ensureInvocationCompleted(invocationId).blockingAwait();
+    long elapsedMs = (System.nanoTime() - start) / 1_000_000;
+
+    assertTrue(
+        "multi-batch drain must fit one shutdownTimeout bound, took " + elapsedMs + "ms",
+        elapsedMs < 1_200);
+  }
 }