From 70d7d54794f879f8d181d1376a16c20477a26fd4 Mon Sep 17 00:00:00 2001 From: Haiyuan Cao Date: Sat, 11 Jul 2026 01:38:02 -0700 Subject: [PATCH] =?UTF-8?q?fix(bigquery):=20BQAA=20P1=20preview=20blockers?= =?UTF-8?q?=20=E2=80=94=20branch-safe=20tracing,=20bounded=20lifecycle,=20?= =?UTF-8?q?boundary=20redaction,=20HITL=20pause/resume?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolves the plugin-local P1 items from the BQAA preview readiness tracker (#1316) across seven review rounds. Framework-level terminal agent/run error callbacks remain intentionally deferred to a separate framework API change. Tracing correctness: - TraceManager keeps ID-only span records (no plugin-owned OTel spans are created or exported; ambient trace IDs are still inherited) in per-branch stacks keyed by InvocationContext.branch(), with kind-checked pops, so concurrent ParallelAgent branches cannot pop each other's spans. Tool spans additionally carry an operation identity — the function-call ID when present, else a per-ToolContext synthetic ID (the framework materializes absent IDs as "") — plus a push-time parent, because ADK executes an event's function calls concurrently within one branch. Privacy boundary: - JsonFormatter.redactTree walks raw containers with sensitive-key redaction BEFORE any Jackson conversion and fails closed per leaf ("[UNSERIALIZABLE]"), and the assembled attributes tree passes through it at the output boundary. Session state is redacted before truncation so smartTruncate's whole-map textual fallback can never expose embedded secrets. Lifecycle and shutdown: - Per-invocation lifecycle tokens are durable (captured by every append continuation; immune to tombstone-cache eviction) and append admission is atomic with finalization via the token's monitor. - One absolute deadline bounds each shutdown operation end-to-end: pending-task waits, every processor drain, and executor termination share it; finalization is completion-ordered (cleanup runs before the returned Completable completes, since RxJava's doFinally notifies downstream first) with disposal covered by the same idempotent action. - BatchProcessor teardown is ownership-transferred (ReentrantLock; an in-flight flush past the deadline performs the teardown), appends are bounded (get(timeout) plus LimitExceededBehavior.ThrowException so Storage quota saturation cannot park append() for minutes), and the final drop-stat snapshot is delivered at true teardown completion. - Live StreamWriters are bounded by admission permits acquired BEFORE construction (each writer owns an internal client and a non-daemon append thread; MAX_LIVE_WRITERS <= closer queue capacity, so a constructed writer can never lose its cleanup owner). Detached closes run on a bounded plugin-owned closer service; at plugin shutdown the unstarted queue drains to a bounded reclaim owner without interrupting active closes. Rows refused at the cap are accounted ("writer_permit_exhausted"), alongside "after_close", "shutdown_timeout", "writer_create_error", and "late_after_finalize". Event contract: - Synthetic adk_request_* function calls emit HITL_*_REQUEST (not _COMPLETED); long-running calls emit pairable TOOL_PAUSED rows with pause_kind and function_call_id; user-message handling inspects FunctionResponse parts, routing HITL responses to HITL_*_COMPLETED (with pair keys, content.result on both producer paths, Python parity) and non-HITL responses to TOOL_COMPLETED with pair keys. Adds the TOOL_PAUSED view and pair-key columns on TOOL_COMPLETED. Tests: focused BQAA suite 173 passed (upstream baseline: 135); full core main profile 1,629 passed / 0 failed / 24 skipped. New regressions cover parallel-branch and concurrent id-less tool span ownership, zero exported spans, nested and fail-closed redaction (attributes and session state), atomic admission/finalization interleavings, durable tokens across cache eviction, single-deadline multi-batch / task-plus-drain / multi-processor shutdown bounds, completion-ordered cleanup, blocked-append and blocked-close teardown ownership, writer permit exhaustion, and shutdownNow-leftover reclaim. Co-Authored-By: Claude Fable 5 --- .../agentanalytics/BatchProcessor.java | 206 ++++++- .../BigQueryAgentAnalyticsPlugin.java | 292 +++++++-- .../plugins/agentanalytics/BigQueryUtils.java | 11 + .../plugins/agentanalytics/JsonFormatter.java | 78 +++ .../plugins/agentanalytics/PluginState.java | 580 ++++++++++++++++-- .../plugins/agentanalytics/TraceManager.java | 382 +++++++++--- .../agentanalytics/BatchProcessorTest.java | 276 ++++++++- .../BigQueryAgentAnalyticsPluginTest.java | 345 ++++++++++- .../agentanalytics/JsonFormatterTest.java | 52 ++ .../agentanalytics/PluginStateTest.java | 544 +++++++++++++++- .../agentanalytics/TraceManagerTest.java | 219 +++++-- 11 files changed, 2688 insertions(+), 297 deletions(-) diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/BatchProcessor.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/BatchProcessor.java index 8278e0ad6..7743819b7 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/BatchProcessor.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/BatchProcessor.java @@ -24,6 +24,7 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.node.ArrayNode; import com.fasterxml.jackson.databind.node.ObjectNode; +import com.google.api.core.ApiFuture; import com.google.cloud.bigquery.storage.v1.AppendRowsResponse; import com.google.cloud.bigquery.storage.v1.Exceptions.AppendSerializationError; import com.google.cloud.bigquery.storage.v1.StreamWriter; @@ -37,8 +38,12 @@ import java.util.concurrent.BlockingQueue; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.arrow.memory.BufferAllocator; @@ -55,6 +60,7 @@ import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; import org.apache.arrow.vector.types.pojo.Field; import org.apache.arrow.vector.types.pojo.Schema; +import org.jspecify.annotations.Nullable; /** Handles asynchronous batching and writing of events to BigQuery. */ class BatchProcessor implements AutoCloseable { @@ -63,10 +69,30 @@ class BatchProcessor implements AutoCloseable { private final StreamWriter writer; private final int batchSize; private final Duration flushInterval; + private final Duration shutdownTimeout; @VisibleForTesting final BlockingQueue> queue; private final ScheduledExecutorService executor; @VisibleForTesting final BufferAllocator allocator; - final AtomicBoolean flushLock = new AtomicBoolean(false); + // Mutual exclusion for flush; a ReentrantLock (not a CAS flag) so close() can WAIT, bounded, + // for an in-flight flush instead of guessing from queue emptiness. + private final ReentrantLock flushMutex = new ReentrantLock(); + private final AtomicBoolean closed = new AtomicBoolean(false); + // Set by close() before it waits for the flush mutex: whichever party releases the mutex last + // (close, or an in-flight flush that outlived close's deadline) performs the actual teardown. + private final AtomicBoolean teardownRequested = new AtomicBoolean(false); + private final AtomicBoolean tornDown = new AtomicBoolean(false); + // While closing, bounds every drain/in-flight append to the remaining close budget. + private volatile @Nullable Instant closeDeadline; + // Delivered the FINAL drop-stat snapshot when teardown actually completes; see closeAndFold. + private volatile @Nullable Consumer> onFinalStats; + // Owner-preserving detached close for the StreamWriter (see teardownOnce). PluginState provides + // an implementation that is bounded, never blocks, and guarantees every writer's close + // eventually runs (a StreamWriter owns an internal client and a NON-DAEMON append thread that + // only ConnectionWorker.close() stops — abandoning one leaks process-level resources). + private final Consumer writerCloser; + // The periodic flush task; stored so per-invocation close() can cancel it instead of leaving a + // scheduled task retaining this (closed) processor until plugin-wide shutdown. + private volatile @Nullable ScheduledFuture flushTask; private final Schema arrowSchema; private final VectorSchemaRoot root; @@ -74,16 +100,22 @@ class BatchProcessor implements AutoCloseable { private final AtomicLong droppedQueueFull = new AtomicLong(); private final AtomicLong droppedAppendError = new AtomicLong(); private final AtomicLong droppedSerializationError = new AtomicLong(); + private final AtomicLong droppedAfterClose = new AtomicLong(); + private final AtomicLong droppedShutdownTimeout = new AtomicLong(); public BatchProcessor( StreamWriter writer, int batchSize, Duration flushInterval, int queueMaxSize, - ScheduledExecutorService executor) { + ScheduledExecutorService executor, + Duration shutdownTimeout, + Consumer writerCloser) { this.writer = writer; + this.writerCloser = writerCloser; this.batchSize = batchSize; this.flushInterval = flushInterval; + this.shutdownTimeout = shutdownTimeout; this.queue = new LinkedBlockingQueue<>(queueMaxSize); this.executor = executor; // It's safe to use Long.MAX_VALUE here as this is a top-level RootAllocator, @@ -95,8 +127,7 @@ public BatchProcessor( } public void start() { - @SuppressWarnings("unused") - var unused = + this.flushTask = executor.scheduleWithFixedDelay( () -> { try { @@ -111,19 +142,26 @@ public void start() { } public void append(Map row) { + if (closed.get()) { + // The owning invocation has already been finalized; accept-and-drop with accounting rather + // than silently enqueueing into a processor whose final drain has already run. + droppedAfterClose.incrementAndGet(); + logger.warning("BatchProcessor is closed, dropping late event."); + return; + } if (!queue.offer(row)) { droppedQueueFull.incrementAndGet(); logger.warning("BigQuery event queue is full, dropping event."); return; } - if (queue.size() >= batchSize && !flushLock.get()) { + if (queue.size() >= batchSize && !flushMutex.isLocked()) { executor.execute(this::flush); } } public void flush() { - // Acquire the flushLock. If another flush is already in progress, return immediately. - if (!flushLock.compareAndSet(false, true)) { + // Acquire the flush mutex. If another flush is already in progress, return immediately. + if (!flushMutex.tryLock()) { return; } try { @@ -145,7 +183,16 @@ public void flush() { } root.setRowCount(batch.size()); try (ArrowRecordBatch recordBatch = new VectorUnloader(root).getRecordBatch()) { - AppendRowsResponse result = writer.append(recordBatch).get(); + // Bound the append so one stuck Storage Write RPC cannot block the flush path (and, + // during close(), the final drain) indefinitely. + ApiFuture appendFuture = writer.append(recordBatch); + AppendRowsResponse result; + try { + result = appendFuture.get(appendTimeoutMillis(), MILLISECONDS); + } catch (TimeoutException e) { + appendFuture.cancel(true); + throw e; + } if (result.hasError()) { droppedAppendError.addAndGet(batch.size()); logger.severe("BigQuery append error: " + result.getError().getMessage()); @@ -185,13 +232,32 @@ public void flush() { root.clear(); } } finally { - flushLock.set(false); - if (queue.size() >= batchSize && !flushLock.get()) { + flushMutex.unlock(); + // Deferred teardown: close() timed out waiting for this flush, transferring ownership of + // the final resource teardown (and drop-stat delivery) to us. + if (teardownRequested.get()) { + teardownOnce(); + } + if (queue.size() >= batchSize && !flushMutex.isLocked()) { executor.execute(this::flush); } } } + /** + * Per-append deadline: normally {@code shutdownTimeout}; once close() has started, capped to the + * remaining close budget so the final drain cannot exceed the caller's bound. + */ + private long appendTimeoutMillis() { + long timeoutMillis = shutdownTimeout.toMillis(); + Instant deadline = this.closeDeadline; + if (deadline != null) { + long remaining = Duration.between(Instant.now(), deadline).toMillis(); + timeoutMillis = Math.max(1, Math.min(timeoutMillis, remaining)); + } + return timeoutMillis; + } + private void populateVector(FieldVector vector, int index, Object value) { if (value == null || (value instanceof JsonNode jsonNode && jsonNode.isNull())) { vector.setNull(index); @@ -266,13 +332,113 @@ ImmutableMap getDropStats() { return ImmutableMap.of( "queue_full", droppedQueueFull.get(), "append_error", droppedAppendError.get(), - "serialization_error", droppedSerializationError.get()); + "serialization_error", droppedSerializationError.get(), + "after_close", droppedAfterClose.get(), + "shutdown_timeout", droppedShutdownTimeout.get()); + } + + /** + * Closes the processor and delivers the FINAL drop-stat snapshot to {@code statsConsumer} when + * teardown actually completes — which may be after this call returns, if an in-flight flush still + * owns the resources when the shutdownTimeout deadline expires (ownership of the teardown then + * transfers to that flush). This guarantees counters recorded by that last flush (e.g. its append + * failure) are included in the delivered snapshot exactly once. + */ + void closeAndFold(Consumer> 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> statsConsumer, Instant deadline) { + this.onFinalStats = statsConsumer; + close(deadline); } @Override public void close() { - while (this.queue != null && !this.queue.isEmpty()) { + 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; + } + // Cancel the periodic flush task so a completed invocation does not leave a scheduled task + // retaining this processor (and its writer) until plugin-wide shutdown. + ScheduledFuture task = this.flushTask; + if (task != null) { + task.cancel(false); + } + // 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(); + if (!this.queue.isEmpty()) { + // Another thread may hold the flush mutex; back off briefly instead of spinning. + try { + Thread.sleep(10); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + break; + } + } + } + int remaining = this.queue.size(); + if (remaining > 0) { + droppedShutdownTimeout.addAndGet(remaining); + this.queue.clear(); + logger.severe( + "Dropping " + remaining + " rows: final drain did not complete within shutdownTimeout."); + } + // Teardown ownership: request it, then try to acquire the flush mutex within the remaining + // budget. If acquired, no flush is active and we tear down here. If the wait expires, the + // in-flight flush performs the teardown (and final stats delivery) when it releases the mutex + // — resources are never destroyed underneath an active flush, and counters that flush records + // are still included in the final snapshot. + teardownRequested.set(true); + boolean acquired = false; + long waitMillis = Math.max(1, Duration.between(Instant.now(), drainDeadline).toMillis()); + try { + acquired = flushMutex.tryLock(waitMillis, MILLISECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + if (acquired) { + try { + teardownOnce(); + } finally { + flushMutex.unlock(); + } + } else { + logger.severe( + "Deferring resource teardown to the in-flight flush: it did not release the flush mutex" + + " within shutdownTimeout."); + // The flush may have released the mutex between the timed wait expiring and the request + // flag becoming visible to it; re-check so the teardown is never lost. + if (flushMutex.tryLock()) { + try { + teardownOnce(); + } finally { + flushMutex.unlock(); + } + } + } + } + + /** + * Tears down Arrow and writer resources and delivers the final drop-stat snapshot, exactly once, + * regardless of whether close() or a deferred in-flight flush gets here first. + */ + private void teardownOnce() { + if (!tornDown.compareAndSet(false, true)) { + return; } if (this.allocator != null) { try { @@ -289,10 +455,22 @@ public void close() { } } if (this.writer != null) { + // StreamWriter.close() can block far beyond any shutdownTimeout (it joins the writer's + // internal non-daemon append thread, then may wait minutes on its internal client and + // callback pools). Delegate to the plugin-owned closer, which detaches the close without + // ever abandoning the writer: cleanup ownership is guaranteed by writer admission permits + // acquired before construction. No further work touches the writer here: teardown only + // runs after appends have stopped (closed gate + flush-mutex ownership), and drop counters + // are final below. + writerCloser.accept(this.writer); + } + // Deliver the final snapshot only now, when no flush can mutate the counters anymore. + Consumer> statsConsumer = this.onFinalStats; + if (statsConsumer != null) { try { - this.writer.close(); + statsConsumer.accept(getDropStats()); } catch (RuntimeException e) { - logger.log(Level.SEVERE, "Failed to close BigQuery writer", e); + logger.log(Level.WARNING, "Failed to deliver final drop stats", e); } } } diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java index ba1ccda18..a71785c49 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java @@ -53,11 +53,16 @@ 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; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; import com.google.genai.types.CustomMetadata; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Maybe; @@ -71,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; @@ -96,6 +102,18 @@ public class BigQueryAgentAnalyticsPlugin extends BasePlugin { "adk_request_input", "HITL_INPUT_REQUEST"); + // pause_kind discriminator for TOOL_PAUSED rows, keyed by the synthetic HITL function-call NAME + // (mirrors the Python plugin's _HITL_PAUSE_KIND_MAP). Long-running calls that are not HITL carry + // pause_kind = "tool". + private static final ImmutableMap HITL_PAUSE_KIND_MAP = + ImmutableMap.of( + "adk_request_credential", + "hitl_credential", + "adk_request_confirmation", + "hitl_confirmation", + "adk_request_input", + "hitl_input"); + private final BigQueryLoggerConfig config; private final BigQuery bigQuery; private final Object tableEnsuredLock = new Object(); @@ -326,7 +344,13 @@ private Completable logEvent( if (latencyMap != null) { row.put("latency_ms", convertToJsonNode(latencyMap)); } - row.put("attributes", convertToJsonNode(getAttributes(data, invocationContext))); + // Redact the complete assembled attributes tree at the output boundary, regardless of which + // 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 parseFuture; if (content != null) { @@ -361,12 +385,17 @@ private Completable logEvent( parseFuture = CompletableFuture.completedFuture(null); } + // Capture the durable lifecycle token NOW, while the invocation is active: the continuation + // below may complete arbitrarily late (e.g. after a parse/offload timeout), and the captured + // token — unlike the bounded processed-invocations cache — cannot be evicted, so a late + // completion can never resurrect a processor for a finalized invocation. + PluginState.InvocationLifecycle lifecycle = + state.getLifecycle(invocationContext.invocationId()); CompletableFuture appendFuture = parseFuture.thenRun( - () -> { - BatchProcessor processor = state.getBatchProcessor(invocationContext.invocationId()); - processor.append(row); - }); + // appendRow enforces the invocation lifecycle gate and accounts for writer + // construction failures instead of losing the row silently. + () -> state.appendRow(lifecycle, invocationContext.invocationId(), row)); state.addPendingTask(invocationContext.invocationId(), appendFuture); return Completable.complete(); } @@ -392,6 +421,43 @@ 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 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 + * completion to its HITL_*_REQUEST / TOOL_PAUSED rows. + */ + private static ImmutableMap hitlPairKeys( + String hitlName, Optional functionCallId) { + ImmutableMap.Builder keys = ImmutableMap.builder(); + keys.put("pause_kind", HITL_PAUSE_KIND_MAP.getOrDefault(hitlName, "tool")); + // 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(); + } + @CanIgnoreReturnValue private static EventData.Builder withFallbackAgent( EventData.Builder builder, @Nullable String author) { @@ -412,7 +478,7 @@ private ResolvedTraceIds getResolvedTraceIds( // written as rows), not an ambient OpenTelemetry framework span that is never logged as a row. // Otherwise parent_span_id would dangle. Ambient OTel still governs trace_id (via getTraceId) // for cross-system correlation. - SpanIds spanIds = traceManager.getCurrentSpanAndParent(); + SpanIds spanIds = traceManager.getCurrentSpanAndParent(invocationContext); return new ResolvedTraceIds( traceId, @@ -461,7 +527,13 @@ private Map getAttributes( sessionMeta.put("user_id", session.userId()); if (!session.state().isEmpty()) { - TruncationResult result = smartTruncate(session.state(), config.maxContentLength()); + // Redact BEFORE truncating: smartTruncate's whole-object fallback stringifies the map + // when one value is unserializable, which would put embedded secrets beyond the reach + // of the final redaction boundary (a string leaf has no keys to redact). redactTree is + // fail-closed per leaf and returns a JsonNode, which smartTruncate then length-bounds + // without ever hitting the textual fallback. + TruncationResult result = + smartTruncate(JsonFormatter.redactTree(session.state()), config.maxContentLength()); sessionMeta.put("state", toJavaObject(result.node())); } attributes.put("session_metadata", sessionMeta); @@ -492,17 +564,18 @@ PluginState getState() { return state; } - private Optional getCompletedEventData(InvocationContext invocationContext) { + private Optional getCompletedEventData( + InvocationContext invocationContext, String expectedKindPrefix) { TraceManager traceManager = state.getTraceManager(invocationContext.invocationId()); String traceId = traceManager.getTraceId(invocationContext); - // Pop the invocation span from the trace manager. - Optional popped = traceManager.popSpan(); + // Pop the completed span (of the expected kind) from the trace manager. + Optional popped = traceManager.popSpan(invocationContext, expectedKindPrefix); if (popped.isEmpty()) { - // No invocation span to pop. - logger.info("No invocation span to pop."); + // No matching span to pop. + logger.info("No span with kind prefix '" + expectedKindPrefix + "' to pop."); return Optional.empty(); } - Optional parentSpanId = traceManager.getCurrentSpanId(); + Optional parentSpanId = traceManager.getCurrentSpanId(invocationContext); EventData.Builder eventDataBuilder = EventData.builder(); eventDataBuilder.setTraceIdOverride(traceId); @@ -528,24 +601,59 @@ public Maybe onUserMessageCallback( Completable logCompletable = logEvent("USER_MESSAGE_RECEIVED", invocationContext, userMessage, Optional.empty()); + // Resumed input arrives in the user message as FunctionResponse parts (a FunctionCall never + // appears here): HITL responses complete their HITL_*_REQUEST / TOOL_PAUSED pair, and a + // non-HITL FunctionResponse is by construction the resume side of a paused long-running tool + // (regular tools complete inside the agent run via afterToolCallback), so it emits + // TOOL_COMPLETED carrying the pause pair keys. if (userMessage.parts().isPresent()) { for (Part part : userMessage.parts().get()) { - if (part.functionCall().isPresent() - && HITL_EVENT_TYPES.containsKey(part.functionCall().get().name().orElse(""))) { - String hitlEvent = HITL_EVENT_TYPES.get(part.functionCall().get().name().get()); - TruncationResult truncatedResult = smartTruncate(part, config.maxContentLength()); + if (part.functionResponse().isEmpty()) { + continue; + } + FunctionResponse functionResponse = part.functionResponse().get(); + String responseName = functionResponse.name().orElse(""); + TruncationResult truncatedResult = + smartTruncate(functionResponse.response(), config.maxContentLength()); + ImmutableMap contentMap = + 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. 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( - hitlEvent + "_COMPLETED", + HITL_EVENT_TYPES.get(responseName) + "_COMPLETED", invocationContext, - ImmutableMap.of( - "tool", - part.functionCall().get().name().get(), - "result", - truncatedResult.node()), + contentMap, + truncatedResult.isTruncated(), + Optional.of( + EventData.builder() + .setExtraAttributes(hitlPairKeys(responseName, functionResponse.id())) + .build()))); + } else { + if (functionResponse.id().isEmpty()) { + logger.fine( + "User-message function response for tool " + + responseName + + " has no id; the resulting TOOL_COMPLETED row cannot pair with a TOOL_PAUSED" + + " row."); + } + ImmutableMap.Builder pairKeys = ImmutableMap.builder(); + pairKeys.put("pause_kind", "tool"); + functionResponse.id().ifPresent(id -> pairKeys.put("function_call_id", id)); + logCompletable = + logCompletable.andThen( + logEvent( + "TOOL_COMPLETED", + invocationContext, + contentMap, truncatedResult.isTruncated(), - Optional.empty())); + Optional.of( + EventData.builder() + .setExtraAttributes(pairKeys.buildOrThrow()) + .build()))); } } } @@ -579,42 +687,75 @@ public Maybe onEventCallback(InvocationContext invocationContext, Event e } if (event.content().isPresent() && event.content().get().parts().isPresent()) { + Set longRunningIds = event.longRunningToolIds().orElse(ImmutableSet.of()); for (Part part : event.content().get().parts().get()) { - if (part.functionCall().isPresent() - && HITL_EVENT_TYPES.containsKey(part.functionCall().get().name().orElse(""))) { - String hitlEvent = HITL_EVENT_TYPES.get(part.functionCall().get().name().get()); - TruncationResult truncatedResult = - smartTruncate(part.functionCall().get().args(), config.maxContentLength()); - logCompletable = - logCompletable.andThen( - logEvent( - hitlEvent + "_COMPLETED", - invocationContext, - ImmutableMap.of( - "tool", - part.functionCall().get().name().get(), - "args", - truncatedResult.node()), - truncatedResult.isTruncated(), - Optional.empty())); + if (part.functionCall().isPresent()) { + FunctionCall functionCall = part.functionCall().get(); + String callName = functionCall.name().orElse(""); + // A synthetic adk_request_* function call is the HITL *request* (the pause side), not a + // completion: emit the plain HITL_*_REQUEST event. The response side emits _COMPLETED. + if (HITL_EVENT_TYPES.containsKey(callName)) { + String hitlEvent = HITL_EVENT_TYPES.get(callName); + TruncationResult truncatedResult = + smartTruncate(functionCall.args(), config.maxContentLength()); + logCompletable = + logCompletable.andThen( + logEvent( + hitlEvent, + invocationContext, + ImmutableMap.of("tool", callName, "args", truncatedResult.node()), + truncatedResult.isTruncated(), + Optional.empty())); + } + // Any long-running function call (HITL or ordinary) suspends awaiting resumption: emit + // a pairable TOOL_PAUSED row. pause_kind derives from the call NAME so HITL pauses read + // hitl_* and ordinary long-running tools read "tool"; function_call_id joins the pair + // to the later resumed completion row. + if (functionCall.id().isPresent() && longRunningIds.contains(functionCall.id().get())) { + TruncationResult truncatedResult = + smartTruncate(functionCall.args(), config.maxContentLength()); + EventData.Builder pausedData = + withFallbackAgent( + EventData.builder() + .setExtraAttributes( + ImmutableMap.builder() + .put( + "pause_kind", + HITL_PAUSE_KIND_MAP.getOrDefault(callName, "tool")) + .put("function_call_id", functionCall.id().get()) + .buildOrThrow()), + event.author()); + logCompletable = + logCompletable.andThen( + logEvent( + "TOOL_PAUSED", + invocationContext, + ImmutableMap.of("tool", callName, "args", truncatedResult.node()), + truncatedResult.isTruncated(), + Optional.of(pausedData.build()))); + } } 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, + // "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", - part.functionResponse().get().name().get(), - "response", - truncatedResult.node()), + "tool", hitlResponse.name().get(), "result", truncatedResult.node()), truncatedResult.isTruncated(), - Optional.empty())); + Optional.of( + EventData.builder() + .setExtraAttributes( + hitlPairKeys(hitlResponse.name().get(), hitlResponse.id())) + .build()))); } } } @@ -723,7 +864,7 @@ public Completable afterRunCallback(InvocationContext invocationContext) { "INVOCATION_COMPLETED", invocationContext, null, - getCompletedEventData(invocationContext)) + getCompletedEventData(invocationContext, "invocation")) .andThen(state.ensureInvocationCompleted(invocationContext.invocationId())); } @@ -734,7 +875,7 @@ public Maybe beforeAgentCallback(BaseAgent agent, CallbackContext callb } state .getTraceManager(callbackContext.invocationContext().invocationId()) - .pushSpan("agent:" + agent.name()); + .pushSpan(callbackContext.invocationContext(), "agent:" + agent.name()); return logEvent("AGENT_STARTING", callbackContext.invocationContext(), null, Optional.empty()) .andThen(Maybe.empty()); } @@ -745,7 +886,7 @@ public Maybe afterAgentCallback(BaseAgent agent, CallbackContext callba "AGENT_COMPLETED", callbackContext.invocationContext(), null, - getCompletedEventData(callbackContext.invocationContext())) + getCompletedEventData(callbackContext.invocationContext(), "agent:")) .andThen(Maybe.empty()); } @@ -821,7 +962,7 @@ public Maybe beforeModelCallback( EventData.builder().setModel(req.model().orElse("")).setExtraAttributes(attributes).build(); state .getTraceManager(callbackContext.invocationContext().invocationId()) - .pushSpan("llm_request"); + .pushSpan(callbackContext.invocationContext(), "llm_request"); return logEvent("LLM_REQUEST", callbackContext.invocationContext(), req, Optional.of(eventData)) .andThen(Maybe.empty()); } @@ -849,8 +990,8 @@ public Maybe afterModelCallback( }); InvocationContext invocationContext = callbackContext.invocationContext(); - Optional spanId = traceManager.getCurrentSpanId(); - SpanIds spanIds = traceManager.getCurrentSpanAndParent(); + Optional spanId = traceManager.getCurrentSpanId(invocationContext); + SpanIds spanIds = traceManager.getCurrentSpanAndParent(invocationContext); String parentSpanId = spanIds.parentSpanId().orElse(null); boolean isPopped = false; @@ -875,7 +1016,7 @@ public Maybe afterModelCallback( } } else { // Final response - pop span - Optional popped = traceManager.popSpan(); + Optional popped = traceManager.popSpan(invocationContext, "llm_request"); if (popped.isPresent()) { spanId = Optional.of(popped.get().spanId()); duration = popped.get().duration(); @@ -927,10 +1068,10 @@ public Maybe onModelErrorCallback( TraceManager traceManager = state.getTraceManager(callbackContext.invocationContext().invocationId()); InvocationContext invocationContext = callbackContext.invocationContext(); - Optional popped = traceManager.popSpan(); + Optional popped = traceManager.popSpan(invocationContext, "llm_request"); String spanId = popped.map(RecordData::spanId).orElse(null); - SpanIds spanIds = traceManager.getCurrentSpanAndParent(); + SpanIds spanIds = traceManager.getCurrentSpanAndParent(invocationContext); String parentSpanId = spanIds.spanId().orElse(null); EventData.Builder eventDataBuilder = @@ -957,8 +1098,23 @@ public Maybe> beforeToolCallback( } ImmutableMap contentMap = ImmutableMap.of("tool_origin", getToolOrigin(tool), "tool", tool.name(), "args", toolArgs); - state.getTraceManager(toolContext.invocationContext().invocationId()).pushSpan("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", toolOperationId(toolContext)); + 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()); } @@ -976,7 +1132,8 @@ public Maybe> afterToolCallback( .ensureInvocationSpan(toolContext.invocationContext()); TraceManager traceManager = state.getTraceManager(toolContext.invocationContext().invocationId()); - Optional popped = traceManager.popSpan(); + Optional popped = + traceManager.popSpan(toolContext.invocationContext(), "tool", toolOperationId(toolContext)); TruncationResult truncationResult = smartTruncate(result, config.maxContentLength()); ImmutableMap contentMap = ImmutableMap.of( @@ -987,15 +1144,15 @@ public Maybe> afterToolCallback( "tool_origin", getToolOrigin(tool)); - SpanIds spanIds = traceManager.getCurrentSpanAndParent(); - 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", @@ -1017,7 +1174,8 @@ public Maybe> onToolErrorCallback( .ensureInvocationSpan(toolContext.invocationContext()); TraceManager traceManager = state.getTraceManager(toolContext.invocationContext().invocationId()); - Optional popped = traceManager.popSpan(); + Optional popped = + traceManager.popSpan(toolContext.invocationContext(), "tool", toolOperationId(toolContext)); TruncationResult truncationResult = smartTruncate(toolArgs, config.maxContentLength()); ImmutableMap contentMap = @@ -1027,16 +1185,16 @@ public Maybe> onToolErrorCallback( .put("tool_origin", getToolOrigin(tool)) .buildOrThrow(); - SpanIds spanIds = traceManager.getCurrentSpanAndParent(); - 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", diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryUtils.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryUtils.java index 3439606e1..58aa66bcc 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryUtils.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryUtils.java @@ -113,7 +113,18 @@ final class BigQueryUtils { "JSON_VALUE(content, '$.tool') AS tool_name", "JSON_QUERY(content, '$.result') AS tool_result", "JSON_VALUE(content, '$.tool_origin') AS tool_origin", + // Pause pair keys: present on resumed long-running tool completions so + // consumers can do the TOOL_PAUSED <-> TOOL_COMPLETED join end-to-end. + "JSON_VALUE(attributes, '$.pause_kind') AS pause_kind", + "JSON_VALUE(attributes, '$.function_call_id') AS function_call_id", "CAST(JSON_VALUE(latency_ms, '$.total_ms') AS INT64) AS total_ms")) + .put( + "TOOL_PAUSED", + ImmutableList.of( + "JSON_VALUE(content, '$.tool') AS tool_name", + "JSON_QUERY(content, '$.args') AS tool_args", + "JSON_VALUE(attributes, '$.pause_kind') AS pause_kind", + "JSON_VALUE(attributes, '$.function_call_id') AS function_call_id")) .put( "TOOL_ERROR", ImmutableList.of( diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/JsonFormatter.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/JsonFormatter.java index 922e03c02..2566781a0 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/JsonFormatter.java +++ b/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. + * + *

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 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 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(); diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java index afb3a1ea3..392135067 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java @@ -21,6 +21,7 @@ import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; +import com.google.api.gax.batching.FlowController; import com.google.api.gax.core.FixedCredentialsProvider; import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.rpc.FixedHeaderProvider; @@ -34,8 +35,10 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.functions.Action; import java.io.IOException; import java.util.Collection; +import java.util.Map; import java.util.Set; import java.util.UUID; import java.util.concurrent.ArrayBlockingQueue; @@ -51,7 +54,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 { @@ -64,9 +66,21 @@ class PluginState { // Idle time before threads are terminated. private static final int GCS_OFFLOAD_IDLE_TIME_SECONDS = 30; + // Bounded detached-close service shared by all BatchProcessors. + private static final int WRITER_CLOSE_MAX_THREADS = 2; + private static final int WRITER_CLOSE_QUEUE_SIZE = 256; + // Hard cap on LIVE StreamWriters (each owns an internal client and a NON-DAEMON append thread + // that only its own close() stops). A permit is acquired BEFORE construction and released only + // when that writer's close has run, so a constructed writer can never lose its cleanup owner + // and process-level resource growth is bounded even under a sustained Storage outage. + // INVARIANT: must be <= WRITER_CLOSE_QUEUE_SIZE so a pending close always has a queue slot and + // pre-shutdown closer rejection is impossible. + @VisibleForTesting static final int MAX_LIVE_WRITERS = 64; + private final BigQueryLoggerConfig config; private final ScheduledExecutorService executor; private final ExecutorService offloadExecutor; + private final ThreadPoolExecutor writerCloseExecutor; private final BigQueryWriteClient writeClient; private static final AtomicLong threadCounter = new AtomicLong(0); // Map of invocation ID to BatchProcessor. @@ -80,11 +94,113 @@ class PluginState { private final Parser parser; private final ConcurrentHashMap>> pendingTasks = new ConcurrentHashMap<>(); + // Durable per-invocation lifecycle tokens. Unlike the bounded processedInvocations cache (whose + // entries can be evicted by size or TTL), a token is captured by every append continuation at + // logEvent time and stays reachable through that reference even after removal from this map, so + // a continuation completing arbitrarily late still observes the invocation's terminal state and + // cannot resurrect a processor. + private final ConcurrentHashMap lifecycles = + new ConcurrentHashMap<>(); + + /** + * Terminal-state token for one invocation; see {@link #lifecycles}. + * + *

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 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; + } + + synchronized void markFinalized() { + finalized = true; + } + + synchronized boolean isFinalized() { + return finalized; + } + } + // Drop counters accumulated from BatchProcessors that have already been closed/removed, so the // aggregate survives per-invocation processor churn. private final AtomicLong droppedQueueFull = new AtomicLong(); private final AtomicLong droppedAppendError = new AtomicLong(); private final AtomicLong droppedSerializationError = new AtomicLong(); + private final AtomicLong droppedAfterClose = new AtomicLong(); + private final AtomicLong droppedShutdownTimeout = new AtomicLong(); + // Rows lost before a BatchProcessor existed (StreamWriter construction failed) or because their + // continuation completed after the invocation was already finalized. + private final AtomicLong droppedWriterCreateError = new AtomicLong(); + private final AtomicLong droppedAfterFinalize = new AtomicLong(); + // Rows dropped because the live-writer permit cap was exhausted (sustained Storage outage). + private final AtomicLong droppedWriterPermitExhausted = new AtomicLong(); + private final java.util.concurrent.Semaphore writerPermits = + new java.util.concurrent.Semaphore(MAX_LIVE_WRITERS); + // Cleanup owners for admitted writers, registered BEFORE StreamWriter construction so a writer + // can never exist without an owner (a construction/startup failure or a plugin close racing + // admission would otherwise abandon its internal client and non-daemon append thread). + private final Set liveLeases = ConcurrentHashMap.newKeySet(); + // Plugin-wide closing gate, set by closeInternal's cleanup before it drains leases and + // processors; creators racing it re-check after publication and self-close. + private volatile boolean closing = false; + + /** + * Cleanup owner for one admitted writer, alive from permit acquisition until the writer's + * detached close task has run (which is the single permit-release point). State transitions are + * monitor-guarded so "dispatch the close exactly once" holds no matter whether the creator, the + * processor's teardown, or plugin close gets there first. + */ + static final class WriterLease { + private @Nullable StreamWriter writer; + private boolean closeRequested; + private boolean closeDispatched; + + /** + * Attaches the constructed writer. Returns true if a close was already requested (plugin + * closing raced admission): the creator must dispatch the close and must not publish. + */ + synchronized boolean attachWriter(StreamWriter writer) { + this.writer = writer; + return closeRequested; + } + + /** + * Marks the lease close-requested and returns the writer to dispatch, or null if none is + * attached yet (the creator will dispatch on attach) or the close was already dispatched. + */ + synchronized @Nullable StreamWriter requestClose() { + closeRequested = true; + return takeForCloseLocked(); + } + + /** Returns the writer to dispatch exactly once, or null. */ + synchronized @Nullable StreamWriter takeForClose() { + return takeForCloseLocked(); + } + + private @Nullable StreamWriter takeForCloseLocked() { + if (writer == null || closeDispatched) { + return null; + } + closeDispatched = true; + return writer; + } + } PluginState(BigQueryLoggerConfig config) throws IOException { this.config = config; @@ -92,6 +208,21 @@ class PluginState { Executors.newScheduledThreadPool( 2, r -> new Thread(r, "bq-analytics-plugin-" + threadCounter.getAndIncrement())); this.offloadExecutor = createGcsOffloadThreadPool(); + this.writerCloseExecutor = + new ThreadPoolExecutor( + WRITER_CLOSE_MAX_THREADS, + WRITER_CLOSE_MAX_THREADS, + 30, + SECONDS, + new ArrayBlockingQueue<>(WRITER_CLOSE_QUEUE_SIZE), + r -> { + Thread t = + new Thread(r, "bq-analytics-writer-close-" + threadCounter.getAndIncrement()); + t.setDaemon(true); + return t; + }, + new ThreadPoolExecutor.AbortPolicy()); + this.writerCloseExecutor.allowCoreThreadTimeOut(true); // One write client per plugin instance, shared by all invocations. this.writeClient = createWriteClient(config); this.processedInvocations = @@ -161,6 +292,11 @@ protected StreamWriter createWriter() { try { return StreamWriter.newBuilder(streamName, writeClient) .setTraceId(BigQueryUtils.getVersionHeaderValue() + ":" + UUID.randomUUID()) + // Nonblocking admission: the default LimitExceededBehavior.Block parks append() for up + // to five minutes waiting on inflight quota, escaping every shutdownTimeout bound this + // plugin enforces. ThrowException surfaces quota saturation as an append failure, which + // the flush path catches and accounts as dropped rows. + .setLimitExceededBehavior(FlowController.LimitExceededBehavior.ThrowException) .setRetrySettings(retrySettings) .setWriterSchema(BigQuerySchema.getArrowSchema()) // Route Storage Write append RPCs to the dataset's region. Without this, appends to any @@ -179,6 +315,14 @@ String getStreamName(BigQueryLoggerConfig config) { config.projectId(), config.datasetId(), config.tableName()); } + /** + * Returns (creating if absent) the invocation's lifecycle token. Called at logEvent time, while + * the invocation is active, so the token each continuation captures predates finalization. + */ + InvocationLifecycle getLifecycle(String invocationId) { + return lifecycles.computeIfAbsent(invocationId, id -> new InvocationLifecycle()); + } + @VisibleForTesting TraceManager getTraceManager(String invocationId) { return traceManagers.computeIfAbsent(invocationId, id -> new TraceManager()); @@ -189,18 +333,210 @@ BatchProcessor getBatchProcessor(String invocationId) { return batchProcessors.computeIfAbsent( invocationId, id -> { - BatchProcessor p = - new BatchProcessor( - createWriter(), - config.batchSize(), - config.batchFlushInterval(), - config.queueMaxSize(), - executor); - p.start(); + BatchProcessor p = tryCreateProcessor(); + if (p == null) { + throw new IllegalStateException( + "Writer admission refused (permit cap exhausted or plugin closing)"); + } return p; }); } + /** + * Creates and starts a processor, or returns null when admission is refused (permit cap + * exhausted, or plugin close raced admission). Accounting for refusals happens here. + * + *

Ownership protocol: the permit is acquired and a {@link WriterLease} registered BEFORE + * {@link #createWriter()}, so from the instant a StreamWriter exists it has a cleanup owner. Any + * failure after construction — including {@code start()} rejection when the shared scheduler has + * concurrently shut down — routes the writer to the detached closer through the lease rather than + * releasing the permit directly; the close task's completion remains the single permit-release + * point. + */ + private @Nullable BatchProcessor tryCreateProcessor() { + if (!writerPermits.tryAcquire()) { + droppedWriterPermitExhausted.incrementAndGet(); + logger.severe( + "Dropping analytics row: live-writer permit cap exhausted (pending StreamWriter closes" + + " have not completed; likely a Storage outage)."); + return null; + } + WriterLease lease = new WriterLease(); + liveLeases.add(lease); + BatchProcessor p = null; + try { + StreamWriter writer = createWriter(); + boolean closeAlreadyRequested = lease.attachWriter(writer); + if (closeAlreadyRequested || closing) { + // Plugin close raced admission: do not publish; the writer goes straight to cleanup. + dispatchLeaseClose(lease); + droppedAfterFinalize.incrementAndGet(); + logger.warning("Dropping analytics row: plugin is closing."); + return null; + } + p = + new BatchProcessor( + writer, + config.batchSize(), + config.batchFlushInterval(), + config.queueMaxSize(), + executor, + config.shutdownTimeout(), + unusedWriter -> dispatchLeaseClose(lease)); + p.start(); + return p; + } catch (RuntimeException e) { + if (p != null) { + // Constructed but start() failed: processor teardown closes the Arrow resources and + // dispatches the writer through the lease. + p.close(); + } else if (!dispatchLeaseClose(lease)) { + // No writer was ever attached (createWriter itself failed): release directly. + liveLeases.remove(lease); + writerPermits.release(); + } + throw e; + } + } + + /** + * Dispatches the lease's writer to the detached closer exactly once; returns whether a writer was + * dispatched. The close task retires the lease and releases the permit. + */ + private boolean dispatchLeaseClose(WriterLease lease) { + StreamWriter writer = lease.takeForClose(); + if (writer == null) { + return false; + } + submitWriterClose(writer, lease); + return true; + } + + /** + * Detached, owner-preserving StreamWriter close: never blocks the caller; the close task retires + * the lease and releases the permit — the ONLY permit-release point for an attached writer. + * Pre-shutdown rejection is impossible (MAX_LIVE_WRITERS <= closer queue capacity); a + * rejection can therefore only mean the closer service was already shut down, in which case + * ownership transfers to a bounded daemon reclaim thread rather than abandoning the writer. + */ + private void submitWriterClose(StreamWriter writer, WriterLease lease) { + Runnable closeTask = + () -> { + try { + writer.close(); + } catch (RuntimeException e) { + logger.log(Level.SEVERE, "Failed to close BigQuery writer", e); + } finally { + liveLeases.remove(lease); + writerPermits.release(); + } + }; + try { + writerCloseExecutor.execute(closeTask); + } catch (java.util.concurrent.RejectedExecutionException e) { + // Closer already shut down (this close raced plugin close). Bounded by the permit cap. + Thread reclaim = + new Thread( + closeTask, "bq-analytics-writer-close-reclaim-" + threadCounter.getAndIncrement()); + reclaim.setDaemon(true); + reclaim.start(); + } + } + + /** + * Appends a row for the given invocation, honoring the invocation lifecycle and accounting for + * every loss mode that can occur before a {@link BatchProcessor} accepts the row: + * + *

    + *
  • A continuation (late parse/offload) completing after {@code ensureInvocationCompleted} + * finalized the invocation must not recreate a processor that nothing will ever close; the + * row is dropped and counted under {@code late_after_finalize}. + *
  • A {@link StreamWriter} construction failure must not silently discard the row; it is + * counted under {@code writer_create_error} and surfaced in the log. The processor mapping + * is not populated on failure, so a later event retries construction. + *
+ */ + void appendRow(InvocationLifecycle lifecycle, String invocationId, Map row) { + BatchProcessor processor; + try { + processor = getOrCreateProcessorIfActive(lifecycle, invocationId); + } catch (RuntimeException e) { + droppedWriterCreateError.incrementAndGet(); + logger.log( + Level.SEVERE, + "Dropping analytics row: failed to create BigQuery writer for invocation " + invocationId, + e); + return; + } + if (processor == null) { + // Accounted inside the creation gate (finalized invocation or permit exhaustion). + return; + } + // 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."); + } + } + + /** + * Atomically returns the invocation's processor, creating it only while the invocation is still + * active. + * + *

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. + * + *

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) { + BatchProcessor processor = + batchProcessors.computeIfAbsent( + invocationId, + id -> { + // The durable token is the primary gate (it survives processedInvocations cache + // eviction); the cache check additionally covers callers holding a token created + // after + // an evicted invocation's finalization. + if (lifecycle.isFinalized() || isProcessed(id)) { + droppedAfterFinalize.incrementAndGet(); + logger.warning( + "Dropping late analytics row: invocation " + id + " is already finalized."); + return null; + } + // Refusal accounting (permit exhaustion / closing race) happens inside. + return tryCreateProcessor(); + }); + if (processor != null && closing) { + // Plugin close may have iterated the processor map before this publication became visible. + // Exactly one party wins the identity-remove: either the close iteration owns the + // processor, or we self-close it here — never neither. + if (batchProcessors.remove(invocationId, processor)) { + processor.closeAndFold( + this::foldStats, java.time.Instant.now().plus(config.shutdownTimeout())); + } + droppedAfterFinalize.incrementAndGet(); + logger.warning("Dropping analytics row: plugin is closing."); + return null; + } + return processor; + } + protected @Nullable GcsOffloader getGcsOffloader(BigQueryLoggerConfig config) { if (config.gcsBucketName().isEmpty()) { return null; @@ -258,6 +594,19 @@ void addPendingTask(String invocationId, CompletableFuture 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> tasks = pendingTasks.get(invocationId); Completable tasksState = Completable.complete(); if (tasks != null && !tasks.isEmpty()) { @@ -266,29 +615,33 @@ Completable ensureInvocationCompleted(String invocationId) { CompletableFuture.allOf(tasks.toArray(new CompletableFuture[0]))); } logger.fine("Waiting for pending tasks to complete for invocation ID: " + invocationId); - return tasksState - .timeout(config.shutdownTimeout().toMillis(), MILLISECONDS) - .doOnError( - e -> { - if (e instanceof TimeoutException) { - logger.log( - Level.WARNING, - "Timeout while waiting for pending tasks to complete for invocation ID: " - + invocationId, - e); - } - }) - .onErrorComplete() - .doFinally( + // Idempotent cleanup shared by the completion-ordered andThen (normal path) and doFinally + // (disposal path): RxJava's doFinally notifies the downstream FIRST and runs its action + // afterwards, so relying on it alone would let blockingAwait()/subscribers observe success + // while finalization is still running. andThen(fromAction) runs the cleanup BEFORE the + // returned Completable completes. + Action cleanup = + runOnce( () -> { + // Mark the durable lifecycle token FIRST (before removing the processor), so any + // continuation that later finds no mapping is guaranteed to observe the terminal + // state. The map entry is removed for memory bounds; outstanding continuations keep + // the token reachable through their captured reference. + InvocationLifecycle lifecycle = lifecycles.remove(invocationId); + if (lifecycle != null) { + lifecycle.markFinalized(); + } // Mark invocation ID as processed to avoid memory leaks. markProcessed(invocationId); BatchProcessor processor = removeProcessor(invocationId); if (processor != null) { - processor.flush(); - processor.close(); - // Fold after close() so rows dropped during the final drain are also counted. - foldDropStats(processor); + // 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) { @@ -297,36 +650,108 @@ Completable ensureInvocationCompleted(String invocationId) { logger.fine("Removing pending tasks for invocation ID: " + invocationId); pendingTasks.remove(invocationId); }); + return tasksState + .timeout(config.shutdownTimeout().toMillis(), MILLISECONDS) + .doOnError( + e -> { + if (e instanceof TimeoutException) { + logger.log( + Level.WARNING, + "Timeout while waiting for pending tasks to complete for invocation ID: " + + invocationId, + e); + } + }) + .onErrorComplete() + .andThen(Completable.fromAction(cleanup)) + .doFinally(cleanup); + } + + /** Wraps an action so repeated invocations (completion path + disposal path) run it once. */ + private static Action runOnce(Action delegate) { + java.util.concurrent.atomic.AtomicBoolean ran = + new java.util.concurrent.atomic.AtomicBoolean(false); + return () -> { + if (ran.compareAndSet(false, true)) { + delegate.run(); + } + }; + } + + /** + * Drains the closer service's unstarted queue to a bounded reclaim owner, WITHOUT interrupting + * active closes (they finish naturally on their daemon workers). The drained tasks' writers (each + * holding an internal client and non-daemon append thread) must still be closed; one daemon + * thread runs them sequentially, and the backlog is bounded by the writer permit cap. + */ + private void drainQueuedWriterCloses() { + java.util.List pending = new java.util.ArrayList<>(); + writerCloseExecutor.getQueue().drainTo(pending); + reclaimPendingWriterCloses(pending); + } + + private void reclaimPendingWriterCloses(java.util.List pending) { + if (pending.isEmpty()) { + return; + } + Thread reclaim = + new Thread( + () -> pending.forEach(Runnable::run), + "bq-analytics-writer-close-reclaim-" + threadCounter.getAndIncrement()); + reclaim.setDaemon(true); + reclaim.start(); } - private void foldDropStats(BatchProcessor processor) { - ImmutableMap stats = processor.getDropStats(); + private void foldStats(ImmutableMap stats) { droppedQueueFull.addAndGet(stats.getOrDefault("queue_full", 0L)); droppedAppendError.addAndGet(stats.getOrDefault("append_error", 0L)); droppedSerializationError.addAndGet(stats.getOrDefault("serialization_error", 0L)); + droppedAfterClose.addAndGet(stats.getOrDefault("after_close", 0L)); + droppedShutdownTimeout.addAndGet(stats.getOrDefault("shutdown_timeout", 0L)); } /** - * Aggregated dropped-row counters across closed and still-live BatchProcessors. Non-zero values - * indicate analytics rows that never reached BigQuery. + * Aggregated dropped-row counters across closed and still-live BatchProcessors, plus rows lost + * before a processor existed ({@code writer_create_error}) or after their invocation was + * finalized ({@code late_after_finalize}). Non-zero values indicate analytics rows that never + * reached BigQuery. */ ImmutableMap getDropStats() { long queueFull = droppedQueueFull.get(); long appendError = droppedAppendError.get(); long serializationError = droppedSerializationError.get(); + long afterClose = droppedAfterClose.get(); + long shutdownTimeout = droppedShutdownTimeout.get(); for (BatchProcessor processor : getBatchProcessors()) { ImmutableMap stats = processor.getDropStats(); queueFull += stats.getOrDefault("queue_full", 0L); appendError += stats.getOrDefault("append_error", 0L); serializationError += stats.getOrDefault("serialization_error", 0L); + afterClose += stats.getOrDefault("after_close", 0L); + shutdownTimeout += stats.getOrDefault("shutdown_timeout", 0L); } - return ImmutableMap.of( - "queue_full", queueFull, - "append_error", appendError, - "serialization_error", serializationError); + return ImmutableMap.builder() + .put("queue_full", queueFull) + .put("append_error", appendError) + .put("serialization_error", serializationError) + .put("after_close", afterClose) + .put("shutdown_timeout", shutdownTimeout) + .put("writer_permit_exhausted", droppedWriterPermitExhausted.get()) + .put("writer_create_error", droppedWriterCreateError.get()) + .put("late_after_finalize", droppedAfterFinalize.get()) + .buildOrThrow(); } 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> tasks = pendingTasks.values().stream().flatMap(Set::stream).collect(toImmutableList()); Completable tasksState = Completable.complete(); @@ -335,27 +760,41 @@ Completable close() { Completable.fromCompletionStage( CompletableFuture.allOf(tasks.toArray(new CompletableFuture[0]))); } - return tasksState - .timeout(config.shutdownTimeout().toMillis(), MILLISECONDS) - .doOnError( - e -> { - if (e instanceof TimeoutException) { - logger.log( - Level.WARNING, "Timeout while waiting for pending tasks to complete.", e); - } - }) - .onErrorComplete() - .doFinally( + // Completion-ordered cleanup shared with the disposal path; see finalizeInvocation. + Action cleanup = + runOnce( () -> { - for (BatchProcessor processor : getBatchProcessors()) { - processor.close(); - // Fold after close() so rows dropped during the final drain are also counted. - foldDropStats(processor); + // Publish the closing gate FIRST: creators observing it refuse admission (or + // self-close after publication), so the drains below plus the creator-side + // rechecks cover every interleaving from permit acquisition to publication. + closing = true; + for (InvocationLifecycle lifecycle : lifecycles.values()) { + lifecycle.markFinalized(); + } + lifecycles.clear(); + // Drain every registered writer lease: constructed writers dispatch to the closer + // now; writers still mid-construction dispatch when their creator attaches them + // (attachWriter returns closeRequested). + for (WriterLease lease : liveLeases) { + StreamWriter leasedWriter = lease.requestClose(); + if (leasedWriter != null) { + submitWriterClose(leasedWriter, lease); + } + } + // Identity-remove each published processor while closing it: a creator racing + // publication re-checks the closing gate and self-closes if it still owns the + // mapping — exactly one party wins remove(id, processor), never neither. A blind + // clear() could silently drop a processor published after this iteration. + for (Map.Entry entry : batchProcessors.entrySet()) { + if (batchProcessors.remove(entry.getKey(), entry.getValue())) { + // Fold via the teardown callback; each drain consumes only the REMAINING + // shared budget, so N processors cannot take N timeouts. + entry.getValue().closeAndFold(this::foldStats, closeDeadline); + } } for (TraceManager traceManager : getTraceManagers()) { traceManager.clearStack(); } - clearBatchProcessors(); clearTraceManagers(); if (writeClient != null) { @@ -368,13 +807,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(); @@ -382,9 +822,25 @@ Completable close() { } else { offloadExecutor.shutdownNow(); } + // Detached writer closes drain without interruption: active closes finish on + // their daemon workers, and the unstarted queue transfers to a bounded reclaim + // owner so no writer loses its cleanup owner. + writerCloseExecutor.shutdown(); + remainingMillis = + java.time.Duration.between(java.time.Instant.now(), closeDeadline).toMillis(); + if (remainingMillis <= 0 + || !writerCloseExecutor.awaitTermination(remainingMillis, MILLISECONDS)) { + // Deadline expired with closes still pending. Do NOT shutdownNow(): that would + // interrupt ACTIVE closes mid-join, leaving partially-closed writers with no + // retry. Instead drain the unstarted queue to a bounded reclaim owner; active + // closes run to natural completion on their daemon worker threads. + drainQueuedWriterCloses(); + } } catch (InterruptedException e) { executor.shutdownNow(); offloadExecutor.shutdownNow(); + writerCloseExecutor.shutdown(); + drainQueuedWriterCloses(); Thread.currentThread().interrupt(); } @@ -396,5 +852,17 @@ Completable close() { logger.log(Level.WARNING, "Failed to close GCS offloader", e); } }); + return tasksState + .timeout(config.shutdownTimeout().toMillis(), MILLISECONDS) + .doOnError( + e -> { + if (e instanceof TimeoutException) { + logger.log( + Level.WARNING, "Timeout while waiting for pending tasks to complete.", e); + } + }) + .onErrorComplete() + .andThen(Completable.fromAction(cleanup)) + .doFinally(cleanup); } } diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/TraceManager.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/TraceManager.java index e455fc509..cb44102ab 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/TraceManager.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/TraceManager.java @@ -19,61 +19,88 @@ import com.google.adk.agents.InvocationContext; import com.google.auto.value.AutoValue; import com.google.errorprone.annotations.CanIgnoreReturnValue; -import io.opentelemetry.api.GlobalOpenTelemetry; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanContext; -import io.opentelemetry.api.trace.Tracer; -import io.opentelemetry.context.Context; import java.time.Duration; import java.time.Instant; +import java.util.ArrayList; +import java.util.Deque; import java.util.Iterator; +import java.util.List; import java.util.Objects; import java.util.Optional; import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.atomic.AtomicReference; -import java.util.function.Predicate; import java.util.logging.Logger; +import org.jspecify.annotations.Nullable; /** - * Manages OpenTelemetry-style trace and span context using InvocationContext callback data. + * Manages the BQAA-internal execution tree of span IDs for one invocation. * - *

Uses a stack of SpanRecord objects to keep span, ID, ownership, and timing in sync. + *

No OpenTelemetry spans are created: records are ID-only, so a host with an SDK exporter + * configured never receives a duplicate plugin-owned span tree next to ADK's framework spans. + * 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. + * + *

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()); static final String DEFAULT_ROOT_AGENT_NAME = "_bq_analytics_root_agent_name"; + private static final String ROOT_BRANCH = ""; - private final ConcurrentLinkedDeque records = new ConcurrentLinkedDeque<>(); - private String rootAgentName = DEFAULT_ROOT_AGENT_NAME; - private String activeInvocationId = "_bq_analytics_active_invocation_id"; - - private final Tracer tracer; + // Span records keyed by ADK branch string ("" for the invocation root / unbranched flows). + private final ConcurrentHashMap> stacksByBranch = + new ConcurrentHashMap<>(); + private volatile String rootAgentName = DEFAULT_ROOT_AGENT_NAME; + private volatile String activeInvocationId = "_bq_analytics_active_invocation_id"; + // Trace ID inherited from the ambient OpenTelemetry span at invocation-root seeding time; null + // when no ambient context existed (getTraceId then falls back to the invocation ID). + private volatile @Nullable String traceId; - TraceManager() { - this(GlobalOpenTelemetry.getTracer("google.adk.plugins.bigquery_agent_analytics")); - } - - TraceManager(Tracer tracer) { - this.tracer = tracer; - } + TraceManager() {} @AutoValue abstract static class SpanRecord { - abstract Span span(); - abstract String spanId(); - abstract boolean ownsSpan(); + /** Span kind ("invocation", "agent:NAME", "llm_request", "tool") for ownership-checked pops. */ + abstract String kind(); + + /** + * Identity of the operation this span belongs to (the tool's function-call ID), or null for + * spans whose kind executes sequentially within a branch. ADK runs an event's function calls + * concurrently by default, all under the same branch, so tool spans must be popped by operation + * identity rather than stack position. + */ + abstract @Nullable String operationId(); + + /** The enclosing span at push time, so concurrent siblings do not corrupt parent linkage. */ + abstract @Nullable String parentSpanId(); abstract Instant startTime(); abstract AtomicReference firstTokenTime(); - static SpanRecord create(Span span, String spanId, boolean ownsSpan, Instant startTime) { + static SpanRecord create( + String spanId, + String kind, + @Nullable String operationId, + @Nullable String parentSpanId, + Instant startTime) { return new AutoValue_TraceManager_SpanRecord( - span, spanId, ownsSpan, startTime, new AtomicReference<>()); + spanId, kind, operationId, parentSpanId, startTime, new AtomicReference<>()); } } @@ -81,10 +108,13 @@ static SpanRecord create(Span span, String spanId, boolean ownsSpan, Instant sta abstract static class RecordData { abstract String spanId(); + abstract Optional parentSpanId(); + abstract Duration duration(); - static RecordData create(String spanId, Duration duration) { - return new AutoValue_TraceManager_RecordData(spanId, duration); + static RecordData create(String spanId, @Nullable String parentSpanId, Duration duration) { + return new AutoValue_TraceManager_RecordData( + spanId, Optional.ofNullable(parentSpanId), duration); } } @@ -94,7 +124,7 @@ abstract static class SpanIds { abstract Optional parentSpanId(); - static SpanIds create(String spanId, String parentSpanId) { + static SpanIds create(@Nullable String spanId, @Nullable String parentSpanId) { return new AutoValue_TraceManager_SpanIds( Optional.ofNullable(spanId), Optional.ofNullable(parentSpanId)); } @@ -128,11 +158,9 @@ public void initTraceIfNeeded(InvocationContext context) { } public String getTraceId(InvocationContext context) { - if (!records.isEmpty()) { - Span currentSpan = records.peekLast().span(); - if (currentSpan.getSpanContext().isValid()) { - return currentSpan.getSpanContext().getTraceId(); - } + String tid = this.traceId; + if (tid != null) { + return tid; } // Fallback to the ambient span. SpanContext ambient = Span.current().getSpanContext(); @@ -143,49 +171,111 @@ public String getTraceId(InvocationContext context) { return context.invocationId(); } - @CanIgnoreReturnValue - public String pushSpan(String spanName) { - Context parentContext = Context.current(); - if (!records.isEmpty()) { - Span parentSpan = records.peekLast().span(); - if (parentSpan.getSpanContext().isValid()) { - parentContext = parentContext.with(parentSpan); - } + private static String newSpanId() { + // Aligns with the OpenTelemetry span ID format (16 hex chars). + return UUID.randomUUID().toString().replace("-", "").substring(0, 16); + } + + private static String branchKey(InvocationContext context) { + try { + return context.branch().orElse(ROOT_BRANCH); + } catch (RuntimeException e) { + return ROOT_BRANCH; } + } - Span span = tracer.spanBuilder(spanName).setParent(parentContext).startSpan(); - String spanIdStr; - if (span.getSpanContext().isValid()) { - spanIdStr = span.getSpanContext().getSpanId(); - } else { - // This span id aligns with the OpenTelemetry Span ID format. - spanIdStr = UUID.randomUUID().toString().replace("-", "").substring(0, 16); + private Deque stackFor(String branch) { + return stacksByBranch.computeIfAbsent(branch, b -> new ConcurrentLinkedDeque<>()); + } + + /** + * Returns the stacks from the given branch up its ancestor chain to the root branch, most + * specific first. A branch "p.a" resolves to ["p.a", "p", ""] (existing stacks only). + */ + private List> branchChain(String branch) { + List> chain = new ArrayList<>(); + String key = branch; + while (true) { + Deque stack = stacksByBranch.get(key); + if (stack != null) { + chain.add(stack); + } + if (key.isEmpty()) { + break; + } + int lastDot = key.lastIndexOf('.'); + key = lastDot >= 0 ? key.substring(0, lastDot) : ROOT_BRANCH; } + return chain; + } - SpanRecord record = SpanRecord.create(span, spanIdStr, true, Instant.now()); - records.add(record); - return spanIdStr; + /** Pushes an ID-only span record onto the calling branch's stack. No OTel span is created. */ + @CanIgnoreReturnValue + public String pushSpan(InvocationContext context, String spanName) { + return pushSpanRecord(context, spanName, null).spanId(); } + /** + * Pushes an ID-only span record with an optional operation identity (the tool's function-call + * ID). The parent span is resolved and stored at push time; when an operation identity is given, + * concurrent sibling records of the same kind are skipped so a second tool starting while the + * first is still running parents to the enclosing agent span, not to its sibling. + */ @CanIgnoreReturnValue - public String attachCurrentSpan() { - Span span = Span.current(); - String spanIdStr; - if (span.getSpanContext().isValid()) { - spanIdStr = span.getSpanContext().getSpanId(); - } else { - spanIdStr = UUID.randomUUID().toString().replace("-", "").substring(0, 16); + SpanRecord pushSpanRecord( + InvocationContext context, String spanName, @Nullable String operationId) { + String branch = branchKey(context); + String parentSpanId = findParentSpanId(branch, operationId == null ? null : spanName); + SpanRecord record = + SpanRecord.create(newSpanId(), spanName, operationId, parentSpanId, Instant.now()); + stackFor(branch).addLast(record); + return record; + } + + /** + * Newest record in the branch chain to serve as a new span's parent. When {@code + * skipConcurrentKind} is set, records of that kind carrying an operation identity are skipped: + * they are concurrent siblings of the span being pushed, not its ancestors. + */ + private @Nullable String findParentSpanId(String branch, @Nullable String skipConcurrentKind) { + for (Deque stack : branchChain(branch)) { + Iterator descending = stack.descendingIterator(); + while (descending.hasNext()) { + SpanRecord record = descending.next(); + if (skipConcurrentKind != null + && record.kind().equals(skipConcurrentKind) + && record.operationId() != null) { + continue; + } + return record.spanId(); + } } + return null; + } - SpanRecord record = SpanRecord.create(span, spanIdStr, false, Instant.now()); - records.add(record); - return spanIdStr; + /** + * Records the ambient OpenTelemetry span's IDs as the invocation root without creating or owning + * any span, so plugin-emitted rows correlate with the host's existing tracing. + */ + @CanIgnoreReturnValue + public String attachCurrentSpan(InvocationContext context) { + SpanContext ambient = Span.current().getSpanContext(); + String spanId; + if (ambient.isValid()) { + spanId = ambient.getSpanId(); + this.traceId = ambient.getTraceId(); + } else { + spanId = newSpanId(); + } + stackFor(branchKey(context)) + .addLast(SpanRecord.create(spanId, "invocation", null, null, Instant.now())); + return spanId; } public void ensureInvocationSpan(InvocationContext context) { String currentInv = context.invocationId(); - if (!records.isEmpty()) { + if (hasAnyRecords()) { if (currentInv.equals(activeInvocationId)) { return; } @@ -194,74 +284,166 @@ public void ensureInvocationSpan(InvocationContext context) { } activeInvocationId = currentInv; + // Reset the inherited trace ID so a new invocation without ambient context does not reuse the + // previous invocation's trace ID (attachCurrentSpan re-captures it when ambient is valid). + this.traceId = null; - Span ambient = Span.current(); - if (ambient.getSpanContext().isValid()) { - attachCurrentSpan(); + if (Span.current().getSpanContext().isValid()) { + attachCurrentSpan(context); } else { - pushSpan("invocation"); + pushSpan(context, "invocation"); } } + private boolean hasAnyRecords() { + for (Deque stack : stacksByBranch.values()) { + if (!stack.isEmpty()) { + return true; + } + } + return false; + } + + /** + * Pops the calling branch's top span record if its kind matches {@code expectedKindPrefix}. + * + *

The branch scoping prevents a concurrently completing {@code ParallelAgent} branch from + * popping another branch's span; the kind check prevents a mismatched pop (e.g. an error callback + * firing without its corresponding push) from corrupting the stack. + */ @CanIgnoreReturnValue - public Optional popSpan() { - if (records.isEmpty()) { + public Optional popSpan(InvocationContext context, String expectedKindPrefix) { + return popSpan(context, expectedKindPrefix, null); + } + + /** + * Pops the calling branch's matching span record. + * + *

With an {@code operationId}, the record is located by kind AND operation identity + * (newest-first) rather than stack position: ADK executes an event's function calls concurrently + * by default within one branch, so a completion must remove its own record even when a sibling + * tool's record sits above it. Without an {@code operationId}, only the branch's top record is + * popped, and only when its kind matches. + */ + @CanIgnoreReturnValue + public Optional popSpan( + InvocationContext context, String expectedKindPrefix, @Nullable String operationId) { + Deque stack = stacksByBranch.get(branchKey(context)); + if (stack == null || stack.isEmpty()) { return Optional.empty(); } - SpanRecord record = records.pollLast(); - if (record == null) { + if (operationId != null) { + Iterator descending = stack.descendingIterator(); + while (descending.hasNext()) { + SpanRecord record = descending.next(); + if (record.kind().startsWith(expectedKindPrefix) + && operationId.equals(record.operationId())) { + descending.remove(); + return Optional.of( + RecordData.create( + record.spanId(), + record.parentSpanId(), + Duration.between(record.startTime(), Instant.now()))); + } + } + logger.fine( + "No span with kind prefix '" + + expectedKindPrefix + + "' and operation ID '" + + operationId + + "' to pop."); return Optional.empty(); } - Duration duration = Duration.between(record.startTime(), Instant.now()); - if (record.ownsSpan()) { - record.span().end(); + SpanRecord top = stack.peekLast(); + if (top == null) { + return Optional.empty(); } - return Optional.of(RecordData.create(record.spanId(), duration)); + if (!top.kind().startsWith(expectedKindPrefix)) { + logger.fine( + "Not popping span of kind '" + + top.kind() + + "': expected kind prefix '" + + expectedKindPrefix + + "'."); + return Optional.empty(); + } + SpanRecord record = stack.pollLast(); + if (record == null) { + return Optional.empty(); + } + return Optional.of( + RecordData.create( + record.spanId(), + record.parentSpanId(), + Duration.between(record.startTime(), Instant.now()))); } public void clearStack() { - for (SpanRecord record : records) { - if (record.ownsSpan()) { - record.span().end(); - } - } - records.clear(); + // Records are ID-only; there are no OTel spans to end. + stacksByBranch.clear(); } - public SpanIds getCurrentSpanAndParent() { - if (records.isEmpty()) { + public SpanIds getCurrentSpanAndParent(InvocationContext context) { + List> chain = branchChain(branchKey(context)); + + SpanRecord current = null; + int currentChainIndex = -1; + for (int i = 0; i < chain.size(); i++) { + SpanRecord top = chain.get(i).peekLast(); + if (top != null) { + current = top; + currentChainIndex = i; + break; + } + } + if (current == null) { return SpanIds.create(null, null); } - String spanId = records.peekLast().spanId(); - String parentId = - findRecord(records.descendingIterator(), record -> !record.spanId().equals(spanId)) - .map(SpanRecord::spanId) - .orElse(null); - return SpanIds.create(spanId, parentId); - } - - public Optional getCurrentSpanId() { - if (records.isEmpty()) { - return Optional.empty(); + // Parent: the record below the current one in its own stack, else the nearest non-empty + // ancestor branch's top (a branch's first span parents to the invocation root). + SpanRecord parent = null; + Iterator descending = chain.get(currentChainIndex).descendingIterator(); + if (descending.hasNext()) { + descending.next(); // Skip the current record. + if (descending.hasNext()) { + parent = descending.next(); + } } - return Optional.of(records.peekLast().spanId()); + if (parent == null) { + for (int i = currentChainIndex + 1; i < chain.size(); i++) { + SpanRecord top = chain.get(i).peekLast(); + if (top != null) { + parent = top; + break; + } + } + } + return SpanIds.create(current.spanId(), parent == null ? null : parent.spanId()); } - private Optional findRecord( - Iterator iterator, Predicate predicate) { - while (iterator.hasNext()) { - SpanRecord record = iterator.next(); - if (predicate.test(record)) { - return Optional.of(record); + public Optional getCurrentSpanId(InvocationContext context) { + for (Deque stack : branchChain(branchKey(context))) { + SpanRecord top = stack.peekLast(); + if (top != null) { + return Optional.of(top.spanId()); } } return Optional.empty(); } private Optional findSpanRecord(String spanId) { - // Search from newest to oldest for efficiency. - return findRecord(records.descendingIterator(), record -> record.spanId().equals(spanId)); + for (Deque stack : stacksByBranch.values()) { + // Search from newest to oldest for efficiency. + Iterator iterator = stack.descendingIterator(); + while (iterator.hasNext()) { + SpanRecord record = iterator.next(); + if (record.spanId().equals(spanId)) { + return Optional.of(record); + } + } + } + return Optional.empty(); } public void recordFirstToken(String spanId) { diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/BatchProcessorTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/BatchProcessorTest.java index 2355c4def..85ca38d7e 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/BatchProcessorTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/BatchProcessorTest.java @@ -19,13 +19,16 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; import com.google.cloud.bigquery.storage.v1.AppendRowsResponse; import com.google.cloud.bigquery.storage.v1.RowError; import com.google.cloud.bigquery.storage.v1.StreamWriter; @@ -36,8 +39,11 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; import java.util.logging.Handler; import java.util.logging.Level; import java.util.logging.LogRecord; @@ -73,10 +79,33 @@ public class BatchProcessorTest { private Schema schema; private Handler mockHandler; + private ExecutorService closePool; + private java.util.function.Consumer writerCloser; + @Before public void setUp() { executor = Executors.newScheduledThreadPool(1); - batchProcessor = new BatchProcessor(mockWriter, 10, Duration.ofMinutes(1), 100, executor); + closePool = Executors.newCachedThreadPool(); + // Mirror production: the plugin-owned closer detaches writer closes off the caller thread. + writerCloser = + w -> + closePool.execute( + () -> { + try { + w.close(); + } catch (RuntimeException ignored) { + // Best-effort close in tests. + } + }); + batchProcessor = + new BatchProcessor( + mockWriter, + 10, + Duration.ofMinutes(1), + 100, + executor, + Duration.ofSeconds(10), + writerCloser); schema = BigQuerySchema.getArrowSchema(); when(mockWriter.append(any(ArrowRecordBatch.class))) @@ -286,7 +315,15 @@ public void flush_handlesGenericExceptionDuringAppend() throws Exception { @Test public void append_triggersFlushWhenBatchSizeReached() { ScheduledExecutorService mockExecutor = mock(ScheduledExecutorService.class); - BatchProcessor bp = new BatchProcessor(mockWriter, 2, Duration.ofMinutes(1), 10, mockExecutor); + BatchProcessor bp = + new BatchProcessor( + mockWriter, + 2, + Duration.ofMinutes(1), + 10, + mockExecutor, + Duration.ofSeconds(10), + writerCloser); Map row = new HashMap<>(); bp.append(row); @@ -359,13 +396,244 @@ public void flush_handlesAllocationFailure() throws Exception { @Test public void close_flushesAndClosesResources() throws Exception { try (BatchProcessor bp = - new BatchProcessor(mockWriter, 10, Duration.ofMinutes(1), 100, executor)) { + new BatchProcessor( + mockWriter, + 10, + Duration.ofMinutes(1), + 100, + executor, + Duration.ofSeconds(10), + writerCloser)) { Map row = new HashMap<>(); row.put("event_type", "CLOSE_EVENT"); bp.append(row); } verify(mockWriter).append(any(ArrowRecordBatch.class)); - verify(mockWriter).close(); + verify(mockWriter, org.mockito.Mockito.timeout(2000)).close(); + } + + @Test + public void close_cancelsPeriodicFlushTask() { + ScheduledExecutorService mockExecutor = mock(ScheduledExecutorService.class); + ScheduledFuture mockFuture = mock(ScheduledFuture.class); + doReturn(mockFuture) + .when(mockExecutor) + .scheduleWithFixedDelay(any(Runnable.class), anyLong(), anyLong(), any(TimeUnit.class)); + BatchProcessor bp = + new BatchProcessor( + mockWriter, + 10, + Duration.ofMinutes(1), + 10, + mockExecutor, + Duration.ofSeconds(1), + writerCloser); + bp.start(); + + bp.close(); + + // A completed invocation must not leave its periodic flush task scheduled (retaining the + // closed processor and writer) until plugin-wide shutdown. + verify(mockFuture).cancel(false); + } + + @Test + public void close_isIdempotent() { + BatchProcessor bp = + new BatchProcessor( + mockWriter, + 10, + Duration.ofMinutes(1), + 10, + executor, + Duration.ofSeconds(1), + writerCloser); + + bp.close(); + bp.close(); + + verify(mockWriter, org.mockito.Mockito.timeout(2000).times(1)).close(); + } + + @Test + public void append_afterClose_dropsWithAccounting() { + BatchProcessor bp = + new BatchProcessor( + mockWriter, + 10, + Duration.ofMinutes(1), + 10, + executor, + Duration.ofSeconds(1), + writerCloser); + bp.close(); + + Map row = new HashMap<>(); + row.put("event_type", "LATE_EVENT"); + bp.append(row); + + assertEquals(1L, (long) bp.getDropStats().get("after_close")); + assertEquals(0, bp.queue.size()); + } + + @Test + public void flush_neverCompletingAppend_isBoundedByShutdownTimeout() { + when(mockWriter.append(any(ArrowRecordBatch.class))).thenReturn(SettableApiFuture.create()); + BatchProcessor bp = + new BatchProcessor( + mockWriter, + 1, + Duration.ofMinutes(1), + 10, + executor, + Duration.ofMillis(200), + writerCloser); + Map row = new HashMap<>(); + row.put("event_type", "STUCK_EVENT"); + bp.queue.offer(row); + + long start = System.nanoTime(); + bp.flush(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue( + "flush must be bounded by the append deadline, took " + elapsedMs + "ms", + elapsedMs < 5_000); + assertEquals(1L, (long) bp.getDropStats().get("append_error")); + } + + @Test + public void close_waitsForInFlightFlushBeforeTeardown() throws Exception { + // A REAL blocked append: the writer call itself parks for 300ms, so the flush thread holds + // the flush mutex while blocked in Storage Write, with the queue already drained. + java.util.concurrent.CountDownLatch appendStarted = new java.util.concurrent.CountDownLatch(1); + org.mockito.Mockito.doAnswer( + invocation -> { + appendStarted.countDown(); + Thread.sleep(300); + return ApiFutures.immediateFuture(AppendRowsResponse.getDefaultInstance()); + }) + .when(mockWriter) + .append(any(ArrowRecordBatch.class)); + BatchProcessor bp = + new BatchProcessor( + mockWriter, + 1, + Duration.ofMinutes(1), + 10, + executor, + Duration.ofSeconds(5), + writerCloser); + + Map row = new HashMap<>(); + row.put("event_type", "IN_FLIGHT_EVENT"); + bp.queue.offer(row); + Thread inFlight = new Thread(bp::flush); + inFlight.start(); + assertTrue(appendStarted.await(2, java.util.concurrent.TimeUnit.SECONDS)); + + long start = System.nanoTime(); + bp.close(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + inFlight.join(2000); + + // close() must wait for the in-flight flush to release the mutex before tearing down the + // writer and Arrow resources underneath it, and still complete within its bound. + assertTrue( + "close should have waited for the in-flight flush, took " + elapsedMs + "ms", + elapsedMs >= 100); + assertTrue("close should be bounded, took " + elapsedMs + "ms", elapsedMs < 5_000); + verify(mockWriter, org.mockito.Mockito.timeout(2000)).close(); + } + + @Test + public void close_deadlineExpired_defersTeardownAndStatsToInFlightFlush() throws Exception { + // The writer call blocks for LONGER than the close deadline, forcing the deferred-ownership + // path: close() returns at its bound without tearing down, and the in-flight flush performs + // the teardown and delivers the final stats (including its own append failure) afterward. + java.util.concurrent.CountDownLatch appendStarted = new java.util.concurrent.CountDownLatch(1); + org.mockito.Mockito.doAnswer( + invocation -> { + appendStarted.countDown(); + Thread.sleep(1200); + throw new RuntimeException("append failed after close deadline"); + }) + .when(mockWriter) + .append(any(ArrowRecordBatch.class)); + BatchProcessor bp = + new BatchProcessor( + mockWriter, + 1, + Duration.ofMinutes(1), + 10, + executor, + Duration.ofMillis(300), + writerCloser); + + Map row = new HashMap<>(); + row.put("event_type", "STUCK_EVENT"); + bp.queue.offer(row); + Thread inFlight = new Thread(bp::flush); + inFlight.start(); + assertTrue(appendStarted.await(2, java.util.concurrent.TimeUnit.SECONDS)); + + java.util.concurrent.CountDownLatch statsDelivered = new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.atomic.AtomicReference< + com.google.common.collect.ImmutableMap> + finalStats = new java.util.concurrent.atomic.AtomicReference<>(); + long start = System.nanoTime(); + bp.closeAndFold( + stats -> { + finalStats.set(stats); + statsDelivered.countDown(); + }); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + // close() is bounded by its deadline even though the flush is still blocked. + assertTrue("close should be bounded, took " + elapsedMs + "ms", elapsedMs < 3_000); + // Ownership transferred: the in-flight flush finishes, tears down, and delivers the final + // snapshot INCLUDING the append failure it recorded after close() had already returned. + assertTrue( + "final stats should be delivered by the deferred flush", + statsDelivered.await(5, java.util.concurrent.TimeUnit.SECONDS)); + inFlight.join(2000); + assertEquals(1L, (long) finalStats.get().get("append_error")); + verify(mockWriter, org.mockito.Mockito.timeout(2000)).close(); + } + + @Test + public void close_blockingWriterClose_doesNotBlockTeardown() throws Exception { + // The real StreamWriter.close() can block for minutes (thread join + client/pool waits); the + // public close() must not inherit that, while the writer still gets closed eventually. + java.util.concurrent.CountDownLatch writerClosed = new java.util.concurrent.CountDownLatch(1); + org.mockito.Mockito.doAnswer( + invocation -> { + Thread.sleep(1500); + writerClosed.countDown(); + return null; + }) + .when(mockWriter) + .close(); + BatchProcessor bp = + new BatchProcessor( + mockWriter, + 10, + Duration.ofMinutes(1), + 10, + executor, + Duration.ofMillis(300), + writerCloser); + + long start = System.nanoTime(); + bp.close(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue( + "close() must not block on StreamWriter.close(), took " + elapsedMs + "ms", + elapsedMs < 1_000); + assertTrue( + "the writer must still be closed eventually", + writerClosed.await(5, java.util.concurrent.TimeUnit.SECONDS)); } } diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java index 8ff58d473..b4dfc2130 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPluginTest.java @@ -74,6 +74,7 @@ import com.google.genai.types.Content; import com.google.genai.types.CustomMetadata; import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionResponse; import com.google.genai.types.GenerateContentResponse; import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.Part; @@ -85,6 +86,7 @@ import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; import io.reactivex.rxjava3.core.Flowable; import java.time.Duration; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Objects; @@ -211,12 +213,14 @@ public void onUserMessageCallback_ensuresInvocationSpan() throws Exception { Content content = Content.builder().build(); // Verify initial state - assertTrue(state.getTraceManager("invocation_id").getCurrentSpanId().isEmpty()); + assertTrue( + state.getTraceManager("invocation_id").getCurrentSpanId(mockInvocationContext).isEmpty()); plugin.onUserMessageCallback(mockInvocationContext, content).blockingSubscribe(); // Verify that ensureInvocationSpan was called and created a span - assertTrue(state.getTraceManager("invocation_id").getCurrentSpanId().isPresent()); + assertTrue( + state.getTraceManager("invocation_id").getCurrentSpanId(mockInvocationContext).isPresent()); } @Test @@ -230,12 +234,14 @@ public void beforeRunCallback_appendsToWriter() throws Exception { @Test public void beforeRunCallback_ensuresInvocationSpan() throws Exception { // Verify initial state - assertTrue(state.getTraceManager("invocation_id").getCurrentSpanId().isEmpty()); + assertTrue( + state.getTraceManager("invocation_id").getCurrentSpanId(mockInvocationContext).isEmpty()); plugin.beforeRunCallback(mockInvocationContext).blockingSubscribe(); // Verify that ensureInvocationSpan was called and created a span - assertTrue(state.getTraceManager("invocation_id").getCurrentSpanId().isPresent()); + assertTrue( + state.getTraceManager("invocation_id").getCurrentSpanId(mockInvocationContext).isPresent()); } @Test @@ -418,7 +424,8 @@ public void afterAgentCallback_stampsInternalExecutionTreeSpanIds() throws Excep .blockingSubscribe(); plugin.beforeAgentCallback(fakeAgent, callbackContext).blockingSubscribe(); - TraceManager.SpanIds current = state.getTraceManager("invocation_id").getCurrentSpanAndParent(); + TraceManager.SpanIds current = + state.getTraceManager("invocation_id").getCurrentSpanAndParent(mockInvocationContext); String agentSpanId = current.spanId().orElseThrow(); String invocationSpanId = current.parentSpanId().orElseThrow(); @@ -566,7 +573,7 @@ public void logEvent_populatesTraceDetails() throws Exception { Span mockSpan = Span.wrap(mockSpanContext); try (Scope scope = mockSpan.makeCurrent()) { - state.getTraceManager("invocation_id").attachCurrentSpan(); + state.getTraceManager("invocation_id").attachCurrentSpan(mockInvocationContext); Content content = Content.builder().build(); plugin.onUserMessageCallback(mockInvocationContext, content).blockingSubscribe(); @@ -919,7 +926,7 @@ public void onModelErrorCallback_populatesCorrectFields() throws Exception { LlmRequest.Builder mockLlmRequestBuilder = mock(LlmRequest.Builder.class); Throwable error = new RuntimeException("model error message"); - state.getTraceManager("invocation_id").pushSpan("llm_request"); + state.getTraceManager("invocation_id").pushSpan(mockInvocationContext, "llm_request"); plugin .onModelErrorCallback(mockCallbackContext, mockLlmRequestBuilder, error) .blockingSubscribe(); @@ -945,7 +952,8 @@ public void onModelErrorCallback_stampsPoppedSpanId() throws Exception { when(mockCallbackContext.invocationContext()).thenReturn(mockInvocationContext); LlmRequest.Builder mockLlmRequestBuilder = mock(LlmRequest.Builder.class); - String llmSpanId = state.getTraceManager("invocation_id").pushSpan("llm_request"); + String llmSpanId = + state.getTraceManager("invocation_id").pushSpan(mockInvocationContext, "llm_request"); plugin .onModelErrorCallback( mockCallbackContext, mockLlmRequestBuilder, new RuntimeException("boom")) @@ -989,8 +997,8 @@ public void afterModelCallback_populatesCorrectFields() throws Exception { tracer.spanBuilder("ambient").setParent(Context.current().with(parentSpan)).startSpan(); // Set valid ambient span context try (Scope scope = ambientSpan.makeCurrent()) { - state.getTraceManager("invocation_id").pushSpan("parent_request"); - state.getTraceManager("invocation_id").pushSpan("llm_request"); + state.getTraceManager("invocation_id").pushSpan(mockInvocationContext, "parent_request"); + state.getTraceManager("invocation_id").pushSpan(mockInvocationContext, "llm_request"); plugin.afterModelCallback(mockCallbackContext, adkResponse).blockingSubscribe(); } finally { ambientSpan.end(); @@ -1029,11 +1037,22 @@ public void afterToolCallback_populatesCorrectFields() throws Exception { ImmutableMap toolArgs = ImmutableMap.of("arg1", "value1"); ImmutableMap result = ImmutableMap.of("res1", "value2"); - state.getTraceManager("invocation_id").pushSpan("tool_request"); + // 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); + plugin.beforeToolCallback(mockTool, toolArgs, mockToolContext).blockingSubscribe(); plugin.afterToolCallback(mockTool, toolArgs, mockToolContext, result).blockingSubscribe(); - Map 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 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"); @@ -1059,7 +1078,7 @@ public AgentOrigin toolOrigin() { AgentTool a2aTool = AgentTool.create(a2aAgent); - state.getTraceManager("invocation_id").pushSpan("tool_request"); + state.getTraceManager("invocation_id").pushSpan(mockInvocationContext, "tool_request"); plugin .afterToolCallback(a2aTool, ImmutableMap.of(), mockToolContext, ImmutableMap.of()) .blockingSubscribe(); @@ -1082,14 +1101,15 @@ public void afterToolCallback_stampsPoppedToolSpanId() throws Exception { plugin .onUserMessageCallback(mockInvocationContext, Content.builder().build()) .blockingSubscribe(); - String toolSpanId = state.getTraceManager("invocation_id").pushSpan("tool"); + String toolSpanId = + state.getTraceManager("invocation_id").pushSpan(mockInvocationContext, "tool"); // After the tool span is pushed, the enclosing span is the invocation span; afterTool must // stamp // it as the row's parent_span_id once the tool span is popped. String invocationSpanId = state .getTraceManager("invocation_id") - .getCurrentSpanAndParent() + .getCurrentSpanAndParent(mockInvocationContext) .parentSpanId() .orElseThrow(); @@ -1805,4 +1825,297 @@ protected Flowable runLiveImpl(InvocationContext invocationContext) { return Flowable.empty(); } } + + private Map> drainRowsByEventType() { + CompletableFuture.allOf( + state + .getPendingTasksForInvocation("invocation_id") + .toArray(new CompletableFuture[0])) + .join(); + Map> rowsByType = new HashMap<>(); + Map row; + while ((row = state.getBatchProcessor("invocation_id").queue.poll()) != null) { + rowsByType.put((String) row.get("event_type"), row); + } + return rowsByType; + } + + @Test + public void logEvent_redactsSensitiveKeysAtFinalAttributesBoundary() throws Exception { + Event event = + Event.builder() + .author("agent_author") + .actions( + EventActions.builder() + .stateDelta( + ImmutableMap.of( + "access_token", + "super-secret", + "nested", + ImmutableMap.of("api_key", "k-123"), + "safe", + "visible")) + .build()) + .build(); + + plugin.onEventCallback(mockInvocationContext, event).blockingSubscribe(); + + Map row = state.getBatchProcessor("invocation_id").queue.poll(); + assertNotNull("Row not found in queue", row); + assertEquals("STATE_DELTA", row.get("event_type")); + JsonNode attributes = (JsonNode) row.get("attributes"); + // state_delta enters attributes directly (not via the content formatter); the final + // output-boundary pass must still redact sensitive keys, including nested ones. + assertEquals("[REDACTED]", attributes.get("state_delta").get("access_token").asText()); + assertEquals("[REDACTED]", attributes.get("state_delta").get("nested").get("api_key").asText()); + assertEquals("visible", attributes.get("state_delta").get("safe").asText()); + } + + @Test + public void onEventCallback_hitlFunctionCall_emitsRequestNotCompleted() throws Exception { + Event event = + Event.builder() + .author("agent_author") + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .name("adk_request_confirmation") + .id("fc-1") + .args(ImmutableMap.of("prompt", "approve?")) + .build()) + .build())) + .build(); + + plugin.onEventCallback(mockInvocationContext, event).blockingSubscribe(); + Map> rows = drainRowsByEventType(); + + // The synthetic function CALL is the HITL request (the pause side), not a completion. + assertTrue( + "Expected HITL_CONFIRMATION_REQUEST, got: " + rows.keySet(), + rows.containsKey("HITL_CONFIRMATION_REQUEST")); + assertFalse(rows.containsKey("HITL_CONFIRMATION_REQUEST_COMPLETED")); + } + + @Test + public void onEventCallback_longRunningHitlCall_emitsPairedToolPaused() throws Exception { + Event event = + Event.builder() + .author("agent_author") + .longRunningToolIds(ImmutableSet.of("fc-1")) + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .name("adk_request_credential") + .id("fc-1") + .args(ImmutableMap.of("scope", "email")) + .build()) + .build())) + .build(); + + plugin.onEventCallback(mockInvocationContext, event).blockingSubscribe(); + Map> rows = drainRowsByEventType(); + + assertTrue(rows.containsKey("HITL_CREDENTIAL_REQUEST")); + Map paused = rows.get("TOOL_PAUSED"); + assertNotNull("TOOL_PAUSED row not found, got: " + rows.keySet(), paused); + JsonNode attributes = (JsonNode) paused.get("attributes"); + assertEquals("hitl_credential", attributes.get("pause_kind").asText()); + assertEquals("fc-1", attributes.get("function_call_id").asText()); + } + + @Test + public void onEventCallback_longRunningOrdinaryCall_emitsToolPausedWithToolKind() + throws Exception { + Event event = + Event.builder() + .author("agent_author") + .longRunningToolIds(ImmutableSet.of("fc-2")) + .content( + Content.fromParts( + Part.builder() + .functionCall( + FunctionCall.builder() + .name("my_long_tool") + .id("fc-2") + .args(ImmutableMap.of("job", "batch-7")) + .build()) + .build())) + .build(); + + plugin.onEventCallback(mockInvocationContext, event).blockingSubscribe(); + Map> rows = drainRowsByEventType(); + + Map paused = rows.get("TOOL_PAUSED"); + assertNotNull("TOOL_PAUSED row not found, got: " + rows.keySet(), paused); + JsonNode attributes = (JsonNode) paused.get("attributes"); + assertEquals("tool", attributes.get("pause_kind").asText()); + assertEquals("fc-2", attributes.get("function_call_id").asText()); + // An ordinary long-running call is not a HITL request. + assertFalse(rows.keySet().stream().anyMatch(k -> k.startsWith("HITL_"))); + } + + @Test + public void onUserMessageCallback_hitlFunctionResponse_emitsCompleted() throws Exception { + Content userMessage = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .name("adk_request_input") + .id("fc-3") + .response(ImmutableMap.of("value", "user typed this")) + .build()) + .build()); + + plugin.onUserMessageCallback(mockInvocationContext, userMessage).blockingSubscribe(); + Map> rows = drainRowsByEventType(); + + assertTrue(rows.containsKey("USER_MESSAGE_RECEIVED")); + // The resumed HITL input arrives as a FunctionResponse and completes the HITL pair; it must + // not also emit TOOL_COMPLETED. + Map completed = rows.get("HITL_INPUT_REQUEST_COMPLETED"); + assertNotNull("Expected HITL_INPUT_REQUEST_COMPLETED, got: " + rows.keySet(), completed); + assertFalse(rows.containsKey("TOOL_COMPLETED")); + // The completion carries the pause pair keys so it joins its HITL_*_REQUEST / TOOL_PAUSED + // rows even when multiple HITL requests share an invocation. + JsonNode completedAttributes = (JsonNode) completed.get("attributes"); + assertEquals("hitl_input", completedAttributes.get("pause_kind").asText()); + assertEquals("fc-3", completedAttributes.get("function_call_id").asText()); + } + + @Test + public void onUserMessageCallback_nonHitlFunctionResponse_emitsToolCompletedWithPairKeys() + throws Exception { + Content userMessage = + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .name("my_long_tool") + .id("fc-4") + .response(ImmutableMap.of("status", "done")) + .build()) + .build()); + + plugin.onUserMessageCallback(mockInvocationContext, userMessage).blockingSubscribe(); + Map> rows = drainRowsByEventType(); + + // A non-HITL FunctionResponse in a user message is the resume side of a paused long-running + // tool; it emits TOOL_COMPLETED carrying the pause pair keys for the BigQuery join. + Map completed = rows.get("TOOL_COMPLETED"); + assertNotNull("TOOL_COMPLETED row not found, got: " + rows.keySet(), completed); + JsonNode attributes = (JsonNode) completed.get("attributes"); + assertEquals("tool", attributes.get("pause_kind").asText()); + assertEquals("fc-4", attributes.get("function_call_id").asText()); + JsonNode content = (JsonNode) completed.get("content"); + assertEquals("my_long_tool", content.get("tool").asText()); + assertNotNull(content.get("result")); + } + + @Test + public void onEventCallback_hitlFunctionResponse_completionCarriesPairKeys() throws Exception { + Event event = + Event.builder() + .author("agent_author") + .content( + Content.fromParts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .name("adk_request_confirmation") + .id("fc-7") + .response(ImmutableMap.of("confirmed", true)) + .build()) + .build())) + .build(); + + plugin.onEventCallback(mockInvocationContext, event).blockingSubscribe(); + Map> rows = drainRowsByEventType(); + + Map completed = rows.get("HITL_CONFIRMATION_REQUEST_COMPLETED"); + assertNotNull( + "HITL_CONFIRMATION_REQUEST_COMPLETED row not found, got: " + rows.keySet(), completed); + 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 startingSpanByTool = new HashMap<>(); + Map completedSpanByTool = new HashMap<>(); + Map 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 + public void logEvent_sessionState_redactedBeforeTruncationFallback() throws Exception { + // One unserializable session-state value must not stringify the whole state map (which would + // put the sibling secret beyond the reach of key redaction); state is redacted BEFORE + // truncation, per leaf. + mockInvocationContext.session().state().put("api_key", "super-secret"); + mockInvocationContext.session().state().put("bad", new Object()); + mockInvocationContext.session().state().put("ok", "visible"); + + plugin.beforeRunCallback(mockInvocationContext).blockingSubscribe(); + + Map row = state.getBatchProcessor("invocation_id").queue.poll(); + assertNotNull("Row not found in queue", row); + JsonNode stateNode = ((JsonNode) row.get("attributes")).get("session_metadata").get("state"); + assertTrue("session state must remain structured, not stringified", stateNode.isObject()); + assertEquals("[REDACTED]", stateNode.get("api_key").asText()); + assertEquals("[UNSERIALIZABLE]", stateNode.get("bad").asText()); + assertEquals("visible", stateNode.get("ok").asText()); + } } diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/JsonFormatterTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/JsonFormatterTest.java index 597312b79..e46603ec9 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/JsonFormatterTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/JsonFormatterTest.java @@ -414,4 +414,56 @@ public void smartTruncate_atDepthBoundary_arrayNotTruncated() { "A nested-array structure exactly at the depth boundary must not be truncated", result.isTruncated()); } + + @Test + public void redactTree_unserializableValue_failsClosedPerLeaf() { + // An unserializable object anywhere in the attributes tree must not route the WHOLE map + // through a textual fallback (which would expose sibling secrets as plain text). + java.util.Map attributes = new java.util.HashMap<>(); + attributes.put("api_key", "secret-key"); + attributes.put("bad", new Object()); // Jackson cannot serialize a plain Object. + attributes.put("ok", "visible"); + + com.fasterxml.jackson.databind.JsonNode node = JsonFormatter.redactTree(attributes); + + assertTrue(node.isObject()); + assertEquals(JsonFormatter.REDACTED_MESSAGE, node.get("api_key").asText()); + assertEquals(JsonFormatter.UNSERIALIZABLE_MESSAGE, node.get("bad").asText()); + assertEquals("visible", node.get("ok").asText()); + } + + @Test + public void redactTree_redactsNestedContainersAndLists() { + java.util.Map attributes = + java.util.Map.of( + "custom_tags", + java.util.Map.of("password", "hunter2", "team", "analytics"), + "entries", + java.util.List.of(java.util.Map.of("refresh_token", "tok", "name", "a"))); + + com.fasterxml.jackson.databind.JsonNode node = JsonFormatter.redactTree(attributes); + + assertEquals(JsonFormatter.REDACTED_MESSAGE, node.get("custom_tags").get("password").asText()); + assertEquals("analytics", node.get("custom_tags").get("team").asText()); + assertEquals( + JsonFormatter.REDACTED_MESSAGE, node.get("entries").get(0).get("refresh_token").asText()); + assertEquals("a", node.get("entries").get(0).get("name").asText()); + } + + @Test + public void redactTree_redactsInsideConvertedPojoLeaves() { + // A leaf that Jackson converts into an object (e.g. a POJO) still gets key redaction. + java.util.Map attributes = + java.util.Map.of( + "node", + JsonFormatter.mapper + .createObjectNode() + .put("client_secret", "s3cret") + .put("plain", "ok")); + + com.fasterxml.jackson.databind.JsonNode node = JsonFormatter.redactTree(attributes); + + assertEquals(JsonFormatter.REDACTED_MESSAGE, node.get("node").get("client_secret").asText()); + assertEquals("ok", node.get("node").get("plain").asText()); + } } diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/PluginStateTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/PluginStateTest.java index 761a60e5e..924aa6723 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/PluginStateTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/PluginStateTest.java @@ -126,9 +126,20 @@ public void addPendingTask_removedTaskOnCompletion() { public void ensureInvocationCompleted_foldsClosedProcessorDropStats() throws IOException { String invocationId = "inv-fold"; BatchProcessor closedProcessor = mock(BatchProcessor.class); - when(closedProcessor.getDropStats()) - .thenReturn( - ImmutableMap.of("queue_full", 5L, "append_error", 3L, "serialization_error", 2L)); + ImmutableMap closedStats = + ImmutableMap.of("queue_full", 5L, "append_error", 3L, "serialization_error", 2L); + // closeAndFold delivers the final snapshot via its callback at teardown completion. + org.mockito.Mockito.doAnswer( + invocation -> { + @SuppressWarnings("unchecked") + java.util.function.Consumer> consumer = + (java.util.function.Consumer>) + invocation.getArgument(0); + consumer.accept(closedStats); + return null; + }) + .when(closedProcessor) + .closeAndFold(org.mockito.ArgumentMatchers.any(), org.mockito.ArgumentMatchers.any()); // Completing an invocation removes and closes its processor; each drop counter must be folded // into the plugin-level totals so it survives after the per-invocation processor is gone. @@ -325,4 +336,531 @@ protected GcsOffloader getGcsOffloader(BigQueryLoggerConfig config) { verify(mockOffloader).close(); } + + @Test + public void appendRow_writerCreationFails_countsDropAndAllowsRetry() throws IOException { + TestPluginState failingState = + new TestPluginState(config) { + @Override + protected StreamWriter createWriter() { + throw new IllegalStateException("writer construction failed"); + } + }; + + failingState.appendRow( + failingState.getLifecycle("inv-writer-fail"), + "inv-writer-fail", + ImmutableMap.of("event_type", "LLM_REQUEST")); + + assertEquals(1L, (long) failingState.getDropStats().get("writer_create_error")); + // The processor mapping must not be populated on failure, so a later event retries + // construction instead of being permanently broken. + assertTrue(failingState.getBatchProcessors().isEmpty()); + } + + @Test + public void appendRow_afterFinalize_dropsWithoutRecreatingProcessor() { + String invocationId = "inv-late"; + PluginState.InvocationLifecycle lifecycle = pluginState.getLifecycle(invocationId); + pluginState.markProcessed(invocationId); + + // A parse/offload continuation completing after the invocation was finalized must not + // recreate a BatchProcessor that nothing will ever close. + pluginState.appendRow(lifecycle, invocationId, ImmutableMap.of("event_type", "LLM_REQUEST")); + + assertTrue(pluginState.getBatchProcessors().isEmpty()); + assertEquals(1L, (long) pluginState.getDropStats().get("late_after_finalize")); + } + + @Test + public void appendRow_finalizedToken_dropsEvenAfterTombstoneEviction() throws Exception { + String invocationId = "inv-evicted"; + // The continuation captures its lifecycle token at logEvent time, while the invocation is + // active. + PluginState.InvocationLifecycle lifecycle = pluginState.getLifecycle(invocationId); + + pluginState.ensureInvocationCompleted(invocationId).blockingAwait(); + + // Simulate processed-cache eviction (size/TTL): invalidate the tombstone entirely, so the + // bounded cache can no longer gate the late continuation. + Field cacheField = PluginState.class.getDeclaredField("processedInvocations"); + cacheField.setAccessible(true); + ((com.google.common.cache.Cache) cacheField.get(pluginState)).invalidateAll(); + assertTrue(!pluginState.isProcessed(invocationId)); + + // The captured token is durable: the late continuation must still be dropped and must not + // resurrect a processor (writer, allocator, periodic task) for the finalized invocation. + pluginState.appendRow(lifecycle, invocationId, ImmutableMap.of("event_type", "LLM_REQUEST")); + + assertTrue(pluginState.getBatchProcessors().isEmpty()); + assertEquals(1L, (long) pluginState.getDropStats().get("late_after_finalize")); + } + + @Test + public void manyCompletedInvocations_leaveNoRetainedProcessorsOrTraceManagers() { + for (int i = 0; i < 100; i++) { + String invocationId = "inv-" + i; + var unusedProcessor = pluginState.getBatchProcessor(invocationId); + var unusedTraceManager = pluginState.getTraceManager(invocationId); + pluginState.ensureInvocationCompleted(invocationId).blockingAwait(); + } + + 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 row1 = new java.util.HashMap<>(); + row1.put("event_type", "A"); + java.util.Map 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; + + // Tight bound: the old per-phase restart behavior took ~2x (>=1000ms with a 500ms timeout); + // one shared absolute deadline finishes in ~one timeout plus scheduling slack. + assertTrue( + "multi-batch drain must fit one shutdownTimeout bound, took " + elapsedMs + "ms", + elapsedMs < 900); + } + + @Test + public void ensureInvocationCompleted_pendingTaskPlusDrain_shareOneDeadline() throws IOException { + // A stuck pending task consumes the budget; the processor drain must NOT get a fresh one. + config = config.toBuilder().shutdownTimeout(Duration.ofMillis(500)).build(); + TestPluginState slowState = + new TestPluginState(config) { + @Override + protected StreamWriter createWriter() { + StreamWriter writer = mock(StreamWriter.class); + when(writer.append(any(ArrowRecordBatch.class))) + .thenReturn(com.google.api.core.SettableApiFuture.create()); + return writer; + } + }; + String invocationId = "inv-task-plus-drain"; + slowState.addPendingTask(invocationId, new CompletableFuture<>()); // never completes + BatchProcessor processor = slowState.getBatchProcessor(invocationId); + java.util.Map row = new java.util.HashMap<>(); + row.put("event_type", "A"); + processor.queue.offer(row); + + long start = System.nanoTime(); + slowState.ensureInvocationCompleted(invocationId).blockingAwait(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + assertTrue( + "pending-task wait plus drain must share one shutdownTimeout, took " + elapsedMs + "ms", + elapsedMs < 900); + } + + @Test + public void close_multipleStuckProcessors_shareOneDeadline() throws IOException { + // N processors with never-completing appends must not take N sequential timeouts. + config = config.toBuilder().shutdownTimeout(Duration.ofMillis(500)).build(); + TestPluginState slowState = + new TestPluginState(config) { + @Override + protected StreamWriter createWriter() { + StreamWriter writer = mock(StreamWriter.class); + when(writer.append(any(ArrowRecordBatch.class))) + .thenReturn(com.google.api.core.SettableApiFuture.create()); + return writer; + } + }; + for (int i = 0; i < 2; i++) { + BatchProcessor processor = slowState.getBatchProcessor("inv-close-" + i); + java.util.Map row = new java.util.HashMap<>(); + row.put("event_type", "A"); + processor.queue.offer(row); + } + + long start = System.nanoTime(); + slowState.close().blockingAwait(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + + // Old behavior was ~one deadline PER processor (>=1000ms for two 500ms drains) before + // executor waits; the shared absolute deadline finishes in ~one timeout plus slack. + assertTrue( + "multi-processor close must share one shutdownTimeout, took " + elapsedMs + "ms", + elapsedMs < 900); + } + + @Test + public void ensureInvocationCompleted_doesNotCompleteBeforeCleanupFinishes() throws Exception { + // RxJava's doFinally notifies the downstream BEFORE running its action; finalization must be + // completion-ordered so callers cannot observe success while cleanup is still running. + String invocationId = "inv-ordered"; + java.util.concurrent.CountDownLatch cleanupStarted = new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.CountDownLatch cleanupRelease = new java.util.concurrent.CountDownLatch(1); + BatchProcessor blockingProcessor = mock(BatchProcessor.class); + org.mockito.Mockito.doAnswer( + invocation -> { + cleanupStarted.countDown(); + cleanupRelease.await(5, TimeUnit.SECONDS); + return null; + }) + .when(blockingProcessor) + .closeAndFold(org.mockito.ArgumentMatchers.any(), org.mockito.ArgumentMatchers.any()); + TestPluginState orderedState = + new TestPluginState(config) { + @Override + protected BatchProcessor removeProcessor(String id) { + return id.equals(invocationId) ? blockingProcessor : super.removeProcessor(id); + } + }; + CompletableFuture pending = new CompletableFuture<>(); + orderedState.addPendingTask(invocationId, pending); + + java.util.concurrent.atomic.AtomicBoolean observedComplete = + new java.util.concurrent.atomic.AtomicBoolean(false); + var unused = + orderedState + .ensureInvocationCompleted(invocationId) + .subscribe(() -> observedComplete.set(true)); + + // Complete the pending task ASYNCHRONOUSLY so the chain advances on another thread and + // blocks inside the (latched) cleanup. + Thread completer = new Thread(() -> pending.complete(null)); + completer.start(); + assertTrue(cleanupStarted.await(2, TimeUnit.SECONDS)); + + // Cleanup is running but blocked: the returned Completable must NOT have completed. + Thread.sleep(100); + assertTrue( + "completion must not be observable before cleanup finishes", !observedComplete.get()); + + cleanupRelease.countDown(); + completer.join(2000); + long deadline = Instant.now().plusMillis(2000).toEpochMilli(); + while (!observedComplete.get() && Instant.now().toEpochMilli() < deadline) { + Thread.sleep(10); + } + assertTrue("completion must be observable after cleanup finishes", observedComplete.get()); + } + + @Test + public void manyInvocationsWithBlockedWriterClose_boundedCloserThreads() throws Exception { + // A Storage outage makes every StreamWriter.close() block. Detached closes must run on the + // plugin-owned BOUNDED service: invocation throughput must not translate into raw-thread + // growth (one blocked closer per completed invocation would exhaust native threads). + java.util.concurrent.CountDownLatch closeRelease = new java.util.concurrent.CountDownLatch(1); + config = config.toBuilder().shutdownTimeout(Duration.ofMillis(300)).build(); + TestPluginState blockedState = + new TestPluginState(config) { + @Override + protected StreamWriter createWriter() { + StreamWriter writer = mock(StreamWriter.class); + when(writer.append(any(ArrowRecordBatch.class))) + .thenReturn(ApiFutures.immediateFuture(AppendRowsResponse.getDefaultInstance())); + org.mockito.Mockito.doAnswer( + invocation -> { + closeRelease.await(10, TimeUnit.SECONDS); + return null; + }) + .when(writer) + .close(); + return writer; + } + }; + + // Other plugin instances in this JVM (from sibling tests) may have idle closer threads; + // assert on the DELTA this instance produces across 25 blocked-close invocations. + long closerThreadsBefore = + Thread.getAllStackTraces().keySet().stream() + .filter(t -> t.getName().startsWith("bq-analytics-writer-close-")) + .count(); + + for (int i = 0; i < 25; i++) { + String invocationId = "inv-blocked-close-" + i; + var unusedProcessor = blockedState.getBatchProcessor(invocationId); + blockedState.ensureInvocationCompleted(invocationId).blockingAwait(); + } + + long closerThreadsAfter = + Thread.getAllStackTraces().keySet().stream() + .filter(t -> t.getName().startsWith("bq-analytics-writer-close-")) + .count(); + long delta = closerThreadsAfter - closerThreadsBefore; + assertTrue( + "closer thread growth must be bounded by the pool size, grew by " + delta, delta <= 2); + // All processors are gone despite the blocked closes. + assertTrue(blockedState.getBatchProcessors().isEmpty()); + + closeRelease.countDown(); + // Plugin shutdown completes within its bound even with a close backlog. + blockedState.close().test().awaitDone(5, SECONDS).assertComplete(); + } + + @Test + public void writerPermitCap_boundsLiveWritersAndPreservesCleanupOwnership() throws Exception { + // Every StreamWriter owns an internal client and a NON-DAEMON append thread; the permit cap + // must refuse new writers (with accounting) once closes back up, and every writer that WAS + // constructed must be closed exactly once when the backlog drains. + java.util.concurrent.CountDownLatch closeRelease = new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.atomic.AtomicInteger writersCreated = + new java.util.concurrent.atomic.AtomicInteger(); + java.util.concurrent.atomic.AtomicInteger writersClosed = + new java.util.concurrent.atomic.AtomicInteger(); + config = config.toBuilder().shutdownTimeout(Duration.ofMillis(200)).build(); + TestPluginState cappedState = + new TestPluginState(config) { + @Override + protected StreamWriter createWriter() { + writersCreated.incrementAndGet(); + StreamWriter writer = mock(StreamWriter.class); + when(writer.append(any(ArrowRecordBatch.class))) + .thenReturn(ApiFutures.immediateFuture(AppendRowsResponse.getDefaultInstance())); + org.mockito.Mockito.doAnswer( + invocation -> { + closeRelease.await(20, TimeUnit.SECONDS); + writersClosed.incrementAndGet(); + return null; + }) + .when(writer) + .close(); + return writer; + } + }; + + // Exhaust the permit cap: every finalized invocation's writer close is blocked, so permits + // are never returned. + for (int i = 0; i < PluginState.MAX_LIVE_WRITERS; i++) { + String invocationId = "inv-permit-" + i; + cappedState.appendRow( + cappedState.getLifecycle(invocationId), + invocationId, + ImmutableMap.of("event_type", "LLM_REQUEST")); + cappedState.ensureInvocationCompleted(invocationId).blockingAwait(); + } + assertEquals(PluginState.MAX_LIVE_WRITERS, writersCreated.get()); + + // One more invocation: refused BEFORE construction, with accounting — no new writer exists + // that could lose its cleanup owner. + cappedState.appendRow( + cappedState.getLifecycle("inv-over-cap"), + "inv-over-cap", + ImmutableMap.of("event_type", "LLM_REQUEST")); + assertEquals(PluginState.MAX_LIVE_WRITERS, writersCreated.get()); + assertEquals(1L, (long) cappedState.getDropStats().get("writer_permit_exhausted")); + + // Drain the backlog: every constructed writer is closed exactly once. + closeRelease.countDown(); + long deadline = Instant.now().plusMillis(10_000).toEpochMilli(); + while (writersClosed.get() < PluginState.MAX_LIVE_WRITERS + && Instant.now().toEpochMilli() < deadline) { + Thread.sleep(20); + } + assertEquals(PluginState.MAX_LIVE_WRITERS, writersClosed.get()); + } + + @Test + public void close_pastDeadline_queuedWriterClosesRetainCleanupOwnership() throws Exception { + // Plugin close() past its deadline drains the closer's unstarted queue to a bounded reclaim + // owner WITHOUT interrupting active closes. No writer may lose its cleanup owner. + java.util.concurrent.CountDownLatch closeRelease = new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.atomic.AtomicInteger writersClosed = + new java.util.concurrent.atomic.AtomicInteger(); + config = config.toBuilder().shutdownTimeout(Duration.ofMillis(300)).build(); + TestPluginState blockedState = + new TestPluginState(config) { + @Override + protected StreamWriter createWriter() { + StreamWriter writer = mock(StreamWriter.class); + when(writer.append(any(ArrowRecordBatch.class))) + .thenReturn(ApiFutures.immediateFuture(AppendRowsResponse.getDefaultInstance())); + org.mockito.Mockito.doAnswer( + invocation -> { + closeRelease.await(20, TimeUnit.SECONDS); + writersClosed.incrementAndGet(); + return null; + }) + .when(writer) + .close(); + return writer; + } + }; + + // 5 finalized invocations: 2 closes become active (and block), 3 sit queued. + int writers = 5; + for (int i = 0; i < writers; i++) { + String invocationId = "inv-owned-" + i; + var unusedProcessor = blockedState.getBatchProcessor(invocationId); + blockedState.ensureInvocationCompleted(invocationId).blockingAwait(); + } + + // Plugin shutdown times out on the blocked closers and drains the unstarted queue to the + // bounded reclaim owner, leaving the two active closes uninterrupted. + blockedState.close().test().awaitDone(5, SECONDS).assertComplete(); + + // Release: active closes finish naturally AND the drained queue runs via the reclaim owner. + closeRelease.countDown(); + long deadline = Instant.now().plusMillis(10_000).toEpochMilli(); + while (writersClosed.get() < writers && Instant.now().toEpochMilli() < deadline) { + Thread.sleep(20); + } + assertEquals( + "every constructed writer must be closed exactly once", writers, writersClosed.get()); + } + + @Test + public void close_racingWriterConstruction_writerClosedExactlyOnce() throws Exception { + // Plugin close() can run while a creator holds a permit and is still inside createWriter(). + // The lease registered before construction must ensure the writer — constructed AFTER the + // close drained the leases — is still closed exactly once and never published. + java.util.concurrent.CountDownLatch constructionStarted = + new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.CountDownLatch constructionRelease = + new java.util.concurrent.CountDownLatch(1); + java.util.concurrent.atomic.AtomicInteger writersClosed = + new java.util.concurrent.atomic.AtomicInteger(); + config = config.toBuilder().shutdownTimeout(Duration.ofMillis(300)).build(); + TestPluginState racingState = + new TestPluginState(config) { + @Override + protected StreamWriter createWriter() { + constructionStarted.countDown(); + try { + constructionRelease.await(10, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + StreamWriter writer = mock(StreamWriter.class); + org.mockito.Mockito.doAnswer( + invocation -> { + writersClosed.incrementAndGet(); + return null; + }) + .when(writer) + .close(); + return writer; + } + }; + + String invocationId = "inv-racing"; + Thread creator = + new Thread( + () -> + racingState.appendRow( + racingState.getLifecycle(invocationId), + invocationId, + ImmutableMap.of("event_type", "LLM_REQUEST"))); + creator.start(); + assertTrue(constructionStarted.await(2, TimeUnit.SECONDS)); + + // Plugin close runs while construction is blocked; it must stay bounded. + long start = System.nanoTime(); + racingState.close().test().awaitDone(5, SECONDS).assertComplete(); + long elapsedMs = (System.nanoTime() - start) / 1_000_000; + assertTrue("plugin close must stay bounded, took " + elapsedMs + "ms", elapsedMs < 3_000); + + // Release construction: the creator observes the drained lease and dispatches the close + // itself instead of publishing. + constructionRelease.countDown(); + creator.join(3000); + long deadline = Instant.now().plusMillis(3000).toEpochMilli(); + while (writersClosed.get() < 1 && Instant.now().toEpochMilli() < deadline) { + Thread.sleep(10); + } + assertEquals("the racing writer must be closed exactly once", 1, writersClosed.get()); + assertTrue( + "no processor may be published after close", racingState.getBatchProcessors().isEmpty()); + } + + @Test + public void startupRejection_writerStillClosedExactlyOnce() throws Exception { + // p.start() fails when the shared scheduler has concurrently shut down. The already + // constructed writer must be routed to the detached closer through its lease, not abandoned + // with a directly released permit. + java.util.concurrent.atomic.AtomicInteger writersClosed = + new java.util.concurrent.atomic.AtomicInteger(); + TestPluginState rejectingState = + new TestPluginState(config) { + @Override + protected StreamWriter createWriter() { + StreamWriter writer = mock(StreamWriter.class); + org.mockito.Mockito.doAnswer( + invocation -> { + writersClosed.incrementAndGet(); + return null; + }) + .when(writer) + .close(); + return writer; + } + }; + // Force start() to throw RejectedExecutionException. + rejectingState.getExecutor().shutdownNow(); + + String invocationId = "inv-start-reject"; + rejectingState.appendRow( + rejectingState.getLifecycle(invocationId), + invocationId, + ImmutableMap.of("event_type", "LLM_REQUEST")); + + assertEquals(1L, (long) rejectingState.getDropStats().get("writer_create_error")); + assertTrue(rejectingState.getBatchProcessors().isEmpty()); + long deadline = Instant.now().plusMillis(3000).toEpochMilli(); + while (writersClosed.get() < 1 && Instant.now().toEpochMilli() < deadline) { + Thread.sleep(10); + } + assertEquals( + "the writer from the failed startup must be closed exactly once", 1, writersClosed.get()); + } } diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/TraceManagerTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/TraceManagerTest.java index e56f68fe4..ba76aa56b 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/TraceManagerTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/TraceManagerTest.java @@ -36,6 +36,7 @@ import java.util.Map; import java.util.Optional; import java.util.concurrent.ConcurrentHashMap; +import org.jspecify.annotations.Nullable; import org.junit.Before; import org.junit.Rule; import org.junit.Test; @@ -55,9 +56,6 @@ public class TraceManagerTest { public void setUp() { tracer = openTelemetryRule.getOpenTelemetry().getTracer("test"); callbackData = new ConcurrentHashMap<>(); - mockContext = mock(InvocationContext.class); - when(mockContext.callbackContextData()).thenReturn(callbackData); - when(mockContext.invocationId()).thenReturn("test-invocation-id"); mockAgent = new BaseAgent("test-agent", "desc", null, null, null) { @Override @@ -70,51 +68,140 @@ protected Flowable runLiveImpl(InvocationContext invocationContext) { return Flowable.empty(); } }; - when(mockContext.agent()).thenReturn(mockAgent); - traceManager = new TraceManager(tracer); + mockContext = branchContext("test-invocation-id", null); + traceManager = new TraceManager(); + } + + private InvocationContext branchContext(String invocationId, @Nullable String branch) { + InvocationContext ctx = mock(InvocationContext.class); + when(ctx.callbackContextData()).thenReturn(callbackData); + when(ctx.invocationId()).thenReturn(invocationId); + when(ctx.branch()).thenReturn(Optional.ofNullable(branch)); + when(ctx.agent()).thenReturn(mockAgent); + return ctx; } @Test public void pushSpan_createsValidSpanId() { - String spanId = traceManager.pushSpan("test-span"); + String spanId = traceManager.pushSpan(mockContext, "test-span"); assertNotNull(spanId); assertTrue(spanId.length() >= 16); } @Test public void pushSpan_maintainsParentChildRelationship() { - String parentId = traceManager.pushSpan("parent"); - String childId = traceManager.pushSpan("child"); + String parentId = traceManager.pushSpan(mockContext, "parent"); + String childId = traceManager.pushSpan(mockContext, "child"); - TraceManager.SpanIds ids = traceManager.getCurrentSpanAndParent(); + TraceManager.SpanIds ids = traceManager.getCurrentSpanAndParent(mockContext); assertEquals(childId, ids.spanId().orElse(null)); assertEquals(parentId, ids.parentSpanId().orElse(null)); } @Test public void popSpan_removesFromStack() { - String parentId = traceManager.pushSpan("parent"); - traceManager.pushSpan("child"); + String parentId = traceManager.pushSpan(mockContext, "parent"); + traceManager.pushSpan(mockContext, "child"); - Optional popped = traceManager.popSpan(); + Optional popped = traceManager.popSpan(mockContext, ""); assertTrue(popped.isPresent()); assertFalse(popped.get().duration().isNegative()); - String currentId = traceManager.getCurrentSpanId().orElse(null); + String currentId = traceManager.getCurrentSpanId(mockContext).orElse(null); assertEquals(parentId, currentId); - TraceManager.SpanIds ids = traceManager.getCurrentSpanAndParent(); + TraceManager.SpanIds ids = traceManager.getCurrentSpanAndParent(mockContext); assertEquals(parentId, ids.spanId().orElse(null)); assertFalse(ids.parentSpanId().isPresent()); } + @Test + public void popSpan_kindMismatch_doesNotPop() { + traceManager.pushSpan(mockContext, "agent:root"); + + // An error callback firing without its matching push must not pop an unrelated record. + assertFalse(traceManager.popSpan(mockContext, "llm_request").isPresent()); + + Optional popped = traceManager.popSpan(mockContext, "agent:"); + assertTrue(popped.isPresent()); + } + + @Test + public void parallelBranches_outOfOrderCompletion_preserveSpanOwnership() { + // Same invocation ID across all branches, exactly like ParallelAgent's merged sub-agents. + InvocationContext root = branchContext("inv", null); + traceManager.ensureInvocationSpan(root); + String rootSpan = traceManager.getCurrentSpanId(root).orElse(null); + assertNotNull(rootSpan); + + InvocationContext branchA = branchContext("inv", "par.agentA"); + InvocationContext branchB = branchContext("inv", "par.agentB"); + // Branch A starts first, branch B second. + String spanA = traceManager.pushSpan(branchA, "agent:agentA"); + String spanB = traceManager.pushSpan(branchB, "agent:agentB"); + + // Rows logged concurrently in each branch see their own span, parented to the invocation root. + TraceManager.SpanIds idsA = traceManager.getCurrentSpanAndParent(branchA); + TraceManager.SpanIds idsB = traceManager.getCurrentSpanAndParent(branchB); + assertEquals(spanA, idsA.spanId().orElse(null)); + assertEquals(rootSpan, idsA.parentSpanId().orElse(null)); + assertEquals(spanB, idsB.spanId().orElse(null)); + assertEquals(rootSpan, idsB.parentSpanId().orElse(null)); + + // Branch A completes FIRST (out of order relative to a global LIFO): it must pop its OWN span, + // not branch B's. + Optional poppedA = traceManager.popSpan(branchA, "agent:"); + assertTrue(poppedA.isPresent()); + assertEquals(spanA, poppedA.get().spanId()); + + // Branch B still owns its span and pops it on its own completion. + assertEquals(spanB, traceManager.getCurrentSpanId(branchB).orElse(null)); + Optional poppedB = traceManager.popSpan(branchB, "agent:"); + assertTrue(poppedB.isPresent()); + assertEquals(spanB, poppedB.get().spanId()); + + // The invocation root span remains for INVOCATION_COMPLETED. + assertEquals(rootSpan, traceManager.getCurrentSpanId(root).orElse(null)); + } + + @Test + public void nestedSpansWithinBranch_parentWithinOwnStackFirst() { + InvocationContext root = branchContext("inv", null); + traceManager.ensureInvocationSpan(root); + + InvocationContext branchA = branchContext("inv", "par.agentA"); + String agentSpan = traceManager.pushSpan(branchA, "agent:agentA"); + String llmSpan = traceManager.pushSpan(branchA, "llm_request"); + + TraceManager.SpanIds ids = traceManager.getCurrentSpanAndParent(branchA); + assertEquals(llmSpan, ids.spanId().orElse(null)); + assertEquals(agentSpan, ids.parentSpanId().orElse(null)); + } + + @Test + public void spanLifecycle_exportsZeroPluginOwnedSpans() { + // Full push/pop lifecycle across kinds: the manager must never create real OTel spans, so a + // host with an SDK exporter configured receives no duplicate plugin-owned span tree. + traceManager.ensureInvocationSpan(mockContext); + traceManager.pushSpan(mockContext, "agent:test-agent"); + traceManager.pushSpan(mockContext, "llm_request"); + traceManager.popSpan(mockContext, "llm_request"); + traceManager.pushSpan(mockContext, "tool"); + traceManager.popSpan(mockContext, "tool"); + traceManager.popSpan(mockContext, "agent:"); + traceManager.popSpan(mockContext, "invocation"); + traceManager.clearStack(); + + assertTrue("BQAA must not export plugin-owned spans", openTelemetryRule.getSpans().isEmpty()); + } + @Test public void ensureInvocationSpan_isIdempotent() { traceManager.ensureInvocationSpan(mockContext); - String id1 = traceManager.getCurrentSpanId().orElse(null); + String id1 = traceManager.getCurrentSpanId(mockContext).orElse(null); traceManager.ensureInvocationSpan(mockContext); - String id2 = traceManager.getCurrentSpanId().orElse(null); + String id2 = traceManager.getCurrentSpanId(mockContext).orElse(null); assertEquals(id1, id2); } @@ -127,31 +214,48 @@ public void ensureInvocationSpan_clearsStaleRecords() { } finally { ambientSpan.end(); } - String id1 = traceManager.getCurrentSpanId().orElse(null); + String id1 = traceManager.getCurrentSpanId(mockContext).orElse(null); // Create a new context with same callback data but different invocation ID - InvocationContext mockContext2 = mock(InvocationContext.class); - when(mockContext2.callbackContextData()).thenReturn(callbackData); - when(mockContext2.invocationId()).thenReturn("new-invocation-id"); - when(mockContext2.agent()).thenReturn(mockAgent); + InvocationContext mockContext2 = branchContext("new-invocation-id", null); Span ambientSpan2 = tracer.spanBuilder("ambient2").startSpan(); try (Scope scope = ambientSpan2.makeCurrent()) { traceManager.ensureInvocationSpan(mockContext2); } finally { ambientSpan2.end(); } - String id2 = traceManager.getCurrentSpanId().orElse(null); + String id2 = traceManager.getCurrentSpanId(mockContext2).orElse(null); assertNotEquals(id1, id2); // Should only have 1 record now - TraceManager.SpanIds ids = traceManager.getCurrentSpanAndParent(); + TraceManager.SpanIds ids = traceManager.getCurrentSpanAndParent(mockContext2); assertFalse(ids.parentSpanId().isPresent()); } + @Test + public void ensureInvocationSpan_newInvocationWithoutAmbient_doesNotReuseOldTraceId() { + Span ambientSpan = tracer.spanBuilder("ambient").startSpan(); + try (Scope scope = ambientSpan.makeCurrent()) { + traceManager.ensureInvocationSpan(mockContext); + } finally { + ambientSpan.end(); + } + String firstTraceId = traceManager.getTraceId(mockContext); + assertEquals(ambientSpan.getSpanContext().getTraceId(), firstTraceId); + + // Second invocation seeds with no ambient context: it must fall back to its own invocation ID + // rather than reusing the previous invocation's inherited trace ID. + InvocationContext mockContext2 = branchContext("new-invocation-id", null); + try (Scope ignored = Context.root().makeCurrent()) { + traceManager.ensureInvocationSpan(mockContext2); + assertEquals("new-invocation-id", traceManager.getTraceId(mockContext2)); + } + } + @Test public void attachCurrentSpan_usesAmbientSpan() { Span ambientSpan = tracer.spanBuilder("ambient").startSpan(); try (Scope scope = ambientSpan.makeCurrent()) { - String attachedId = traceManager.attachCurrentSpan(); + String attachedId = traceManager.attachCurrentSpan(mockContext); String expectedId = ambientSpan.getSpanContext().getSpanId(); assertEquals(expectedId, attachedId); } finally { @@ -161,7 +265,7 @@ public void attachCurrentSpan_usesAmbientSpan() { @Test public void getTraceId_returnsCurrentTraceId() { - traceManager.pushSpan("test"); + traceManager.pushSpan(mockContext, "test"); String traceId = traceManager.getTraceId(mockContext); assertNotNull(traceId); if (traceId.equals("test-invocation-id")) { @@ -195,22 +299,23 @@ public void getTraceId_returnsAmbientTraceId_whenRecordsIsEmpty_butAmbientIsPres @Test public void attachCurrentSpan_worksWithoutAmbientSpan() { - // Ensure no ambient span - String attachedId = traceManager.attachCurrentSpan(); - assertNotNull(attachedId); - assertEquals(16, attachedId.length()); + try (Scope ignored = Context.root().makeCurrent()) { + String attachedId = traceManager.attachCurrentSpan(mockContext); + assertNotNull(attachedId); + assertEquals(16, attachedId.length()); - // Verify it's in records - assertEquals(attachedId, traceManager.getCurrentSpanId().orElse(null)); + // Verify it's in records + assertEquals(attachedId, traceManager.getCurrentSpanId(mockContext).orElse(null)); + } } @Test - public void getTraceId_fallsBackToInvocationId_whenRecordSpanIsInvalid() { + public void getTraceId_fallsBackToInvocationId_whenNoAmbientContext() { // Pin to the root context so a span leaked by another test sharing this JVM does not make the - // attached record (or the ambient fallback) valid; attachCurrentSpan with no ambient context - // then creates an invalid span record, exercising the invocation-id fallback. + // ambient fallback valid; attachCurrentSpan with no ambient context records a locally + // generated span ID and no trace ID, exercising the invocation-id fallback. try (Scope ignored = Context.root().makeCurrent()) { - traceManager.attachCurrentSpan(); + traceManager.attachCurrentSpan(mockContext); String traceId = traceManager.getTraceId(mockContext); assertEquals("test-invocation-id", traceId); @@ -219,14 +324,14 @@ public void getTraceId_fallsBackToInvocationId_whenRecordSpanIsInvalid() { @Test public void popSpan_returnsEmpty_whenRecordsIsEmpty() { - Optional popped = traceManager.popSpan(); + Optional popped = traceManager.popSpan(mockContext, ""); assertFalse(popped.isPresent()); } @Test public void clearStack_doesNothing_whenRecordsIsEmpty() { traceManager.clearStack(); - assertTrue(traceManager.getCurrentSpanAndParent().spanId().isEmpty()); + assertTrue(traceManager.getCurrentSpanAndParent(mockContext).spanId().isEmpty()); } @Test @@ -259,4 +364,44 @@ public void initTrace_nullRootAgent_keepsSentinelWithoutThrowing() { assertEquals(TraceManager.DEFAULT_ROOT_AGENT_NAME, traceManager.getRootAgentName()); } + + @Test + public void concurrentToolsInOneBranch_popByOperationIdentity() { + // ADK executes an event's function calls concurrently by default, all under the SAME branch: + // branch + kind alone cannot discriminate, so tool records carry an operation identity. + InvocationContext ctx = branchContext("inv", null); + traceManager.ensureInvocationSpan(ctx); + String agentSpan = traceManager.pushSpan(ctx, "agent:worker"); + + TraceManager.SpanRecord toolA = traceManager.pushSpanRecord(ctx, "tool", "fc-a"); + TraceManager.SpanRecord toolB = traceManager.pushSpanRecord(ctx, "tool", "fc-b"); + + // Both concurrent tools parent to the enclosing agent span, not to each other. + assertEquals(agentSpan, toolA.parentSpanId()); + assertEquals(agentSpan, toolB.parentSpanId()); + + // Tool A completes FIRST even though tool B sits above it on the stack: identity pop must + // remove A's record, not B's. + Optional poppedA = traceManager.popSpan(ctx, "tool", "fc-a"); + assertTrue(poppedA.isPresent()); + assertEquals(toolA.spanId(), poppedA.get().spanId()); + assertEquals(agentSpan, poppedA.get().parentSpanId().orElse(null)); + + Optional poppedB = traceManager.popSpan(ctx, "tool", "fc-b"); + assertTrue(poppedB.isPresent()); + assertEquals(toolB.spanId(), poppedB.get().spanId()); + assertEquals(agentSpan, poppedB.get().parentSpanId().orElse(null)); + + // The agent span remains for AGENT_COMPLETED. + assertEquals(agentSpan, traceManager.getCurrentSpanId(ctx).orElse(null)); + } + + @Test + public void popSpan_unknownOperationId_popsNothing() { + InvocationContext ctx = branchContext("inv", null); + traceManager.pushSpanRecord(ctx, "tool", "fc-a"); + + assertFalse(traceManager.popSpan(ctx, "tool", "fc-unknown").isPresent()); + assertTrue(traceManager.popSpan(ctx, "tool", "fc-a").isPresent()); + } }