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 68d3f4c6d..8278e0ad6 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 @@ -28,6 +28,7 @@ import com.google.cloud.bigquery.storage.v1.Exceptions.AppendSerializationError; import com.google.cloud.bigquery.storage.v1.StreamWriter; import com.google.common.annotations.VisibleForTesting; +import com.google.common.collect.ImmutableMap; import java.time.Duration; import java.time.Instant; import java.util.ArrayList; @@ -37,6 +38,7 @@ import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.arrow.memory.BufferAllocator; @@ -68,6 +70,11 @@ class BatchProcessor implements AutoCloseable { private final Schema arrowSchema; private final VectorSchemaRoot root; + // Drop accounting so hosts can programmatically detect lost analytics rows. + private final AtomicLong droppedQueueFull = new AtomicLong(); + private final AtomicLong droppedAppendError = new AtomicLong(); + private final AtomicLong droppedSerializationError = new AtomicLong(); + public BatchProcessor( StreamWriter writer, int batchSize, @@ -105,6 +112,7 @@ public void start() { public void append(Map row) { if (!queue.offer(row)) { + droppedQueueFull.incrementAndGet(); logger.warning("BigQuery event queue is full, dropping event."); return; } @@ -139,6 +147,7 @@ public void flush() { try (ArrowRecordBatch recordBatch = new VectorUnloader(root).getRecordBatch()) { AppendRowsResponse result = writer.append(recordBatch).get(); if (result.hasError()) { + droppedAppendError.addAndGet(batch.size()); logger.severe("BigQuery append error: " + result.getError().getMessage()); for (var error : result.getRowErrorsList()) { logger.severe( @@ -153,6 +162,7 @@ public void flush() { Thread.currentThread().interrupt(); } if (e.getCause() instanceof AppendSerializationError ase) { + droppedSerializationError.addAndGet(batch.size()); logger.log( Level.SEVERE, "Failed to write batch to BigQuery due to serialization error", ase); Map rowIndexToErrorMessage = ase.getRowIndexToErrorMessage(); @@ -167,6 +177,7 @@ public void flush() { "AppendSerializationError occurred, but no row-specific errors were provided."); } } else { + droppedAppendError.addAndGet(batch.size()); logger.log(Level.SEVERE, "Failed to write batch to BigQuery", e); } } finally { @@ -247,6 +258,17 @@ private void populateVector(FieldVector vector, int index, Object value) { } } + /** + * Returns a snapshot of dropped-row counters keyed by reason ({@code queue_full}, {@code + * append_error}, {@code serialization_error}). Non-zero values indicate lost analytics rows. + */ + ImmutableMap getDropStats() { + return ImmutableMap.of( + "queue_full", droppedQueueFull.get(), + "append_error", droppedAppendError.get(), + "serialization_error", droppedSerializationError.get()); + } + @Override public void close() { while (this.queue != null && !this.queue.isEmpty()) { 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 5775d6d4d..ba1ccda18 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 @@ -51,9 +51,11 @@ import com.google.cloud.bigquery.Table; import com.google.cloud.bigquery.TableId; import com.google.cloud.bigquery.TableInfo; +import com.google.cloud.bigquery.TimePartitioning; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; import com.google.genai.types.CustomMetadata; import com.google.genai.types.Part; @@ -65,10 +67,12 @@ import java.util.ArrayList; import java.util.HashMap; import java.util.List; +import java.util.Locale; import java.util.Map; import java.util.Optional; import java.util.Set; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; import java.util.logging.Level; import java.util.logging.Logger; import org.jspecify.annotations.Nullable; @@ -97,6 +101,10 @@ public class BigQueryAgentAnalyticsPlugin extends BasePlugin { private final Object tableEnsuredLock = new Object(); private final PluginState state; private volatile boolean tableEnsured = false; + // Set only on the public construction paths (see registerShutdownHook); null for the + // package-private test constructor. Deregistered in close() so an explicit close does not leave + // the hook pinning this plugin for the JVM's lifetime. + private @Nullable Thread shutdownHook; public BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config) throws IOException { this(config, createBigQuery(config)); @@ -105,6 +113,9 @@ public BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config) throws IOExcept public BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config, BigQuery bigQuery) throws IOException { this(config, bigQuery, new PluginState(config)); + // Register on the public construction paths only (not the package-private test constructor), + // so a host that never calls close() still gets a best-effort drain at JVM exit. + registerShutdownHook(); } BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config, BigQuery bigQuery, PluginState state) { @@ -114,6 +125,43 @@ public BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config, BigQuery bigQue this.state = state; } + private void registerShutdownHook() { + shutdownHook = + new Thread( + () -> { + try { + boolean unused = + state + .close() + .blockingAwait(config.shutdownTimeout().toMillis(), TimeUnit.MILLISECONDS); + } catch (RuntimeException e) { + logger.log(Level.WARNING, "Error draining BQAA analytics on JVM shutdown", e); + } + }, + "bq-analytics-shutdown"); + Runtime.getRuntime().addShutdownHook(shutdownHook); + } + + private void removeShutdownHook() { + if (shutdownHook == null) { + return; + } + try { + Runtime.getRuntime().removeShutdownHook(shutdownHook); + } catch (IllegalStateException e) { + // The JVM is already shutting down; the hook is running (or has run) and cannot be removed. + } + } + + /** + * Returns aggregated dropped-row counters keyed by reason ({@code queue_full}, {@code + * append_error}, {@code serialization_error}). Non-zero values indicate analytics rows that never + * reached BigQuery. + */ + public ImmutableMap getDropStats() { + return state.getDropStats(); + } + private static BigQuery createBigQuery(BigQueryLoggerConfig config) throws IOException { BigQueryOptions.Builder builder = BigQueryOptions.newBuilder(); builder.setHeaderProvider( @@ -133,24 +181,36 @@ private void ensureTableExistsOnce() { if (!tableEnsured) { synchronized (tableEnsuredLock) { if (!tableEnsured) { - // Table creation is expensive, so we only do it once per plugin instance. - tableEnsured = true; - ensureTableExists(bigQuery, config); + // Only mark the table as ensured after a successful setup, so a transient first-run + // failure (auth blip, missing dataset, quota) is retried on subsequent events instead + // of permanently disabling table creation/upgrade for this plugin instance. + if (ensureTableExists(bigQuery, config)) { + tableEnsured = true; + } } } } } - private void ensureTableExists(BigQuery bigQuery, BigQueryLoggerConfig config) { + /** Returns true if the events table is present (created or already existed) and ready. */ + private boolean ensureTableExists(BigQuery bigQuery, BigQueryLoggerConfig config) { TableId tableId = TableId.of(config.projectId(), config.datasetId(), config.tableName()); Schema schema = BigQuerySchema.getEventsSchema(); + boolean tableReady = false; try { Table table = bigQuery.getTable(tableId); - logger.info("BigQuery table: " + tableId); + logger.fine("BigQuery table: " + tableId); if (table == null) { logger.info("Creating BigQuery table: " + tableId); StandardTableDefinition.Builder tableDefinitionBuilder = - StandardTableDefinition.newBuilder().setSchema(schema); + StandardTableDefinition.newBuilder() + .setSchema(schema) + // Day-partition on the event timestamp for cost/pruning parity with the Python + // plugin. Time-filtered analytics queries prune partitions instead of full scans. + .setTimePartitioning( + TimePartitioning.newBuilder(TimePartitioning.Type.DAY) + .setField("timestamp") + .build()); if (!config.clusteringFields().isEmpty()) { tableDefinitionBuilder.setClustering( Clustering.newBuilder().setFields(config.clusteringFields()).build()); @@ -161,9 +221,24 @@ private void ensureTableExists(BigQuery bigQuery, BigQueryLoggerConfig config) { ImmutableMap.of( BigQuerySchema.SCHEMA_VERSION_LABEL_KEY, BigQuerySchema.SCHEMA_VERSION)) .build(); - bigQuery.create(tableInfo); + try { + bigQuery.create(tableInfo); + } catch (BigQueryException e) { + // Another writer may have created the table concurrently; treat that as success. + String msg = e.getMessage(); + if (msg != null && msg.toLowerCase(Locale.ROOT).contains("already exists")) { + logger.info("BigQuery table already exists (concurrent create): " + tableId); + } else { + throw e; + } + } + tableReady = true; } else if (config.autoSchemaUpgrade()) { - maybeUpgradeSchema(bigQuery, table); + // Only treat the table as ready if the schema upgrade actually succeeded, so a failed + // upgrade is retried on a later event instead of being masked by tableEnsured=true. + tableReady = maybeUpgradeSchema(bigQuery, table); + } else { + tableReady = true; } } catch (BigQueryException e) { processBigQueryException(e, "Failed to check or create/upgrade BigQuery table: " + tableId); @@ -178,6 +253,7 @@ private void ensureTableExists(BigQuery bigQuery, BigQueryLoggerConfig config) { } catch (RuntimeException e) { logger.log(Level.WARNING, "Failed to create/update BigQuery views for table: " + tableId, e); } + return tableReady; } private void processBigQueryException(BigQueryException e, String logMessage) { @@ -234,7 +310,7 @@ private Completable logEvent( Map row = new HashMap<>(); row.put("timestamp", Instant.now()); row.put("event_type", eventType); - row.put("agent", invocationContext.agent().name()); + row.put("agent", resolveAgentName(invocationContext, eventData)); row.put("session_id", invocationContext.session().id()); row.put("invocation_id", invocationContext.invocationId()); row.put("user_id", invocationContext.userId()); @@ -295,6 +371,36 @@ private Completable logEvent( return Completable.complete(); } + /** + * Resolves the agent name defensively. Workflow-driven callbacks may have no current agent; fall + * back to the event author (via {@code EventData.fallbackAgentName}, mirroring the Python plugin) + * and finally to a sentinel, rather than letting an NPE drop the row. + */ + private static String resolveAgentName( + InvocationContext invocationContext, Optional eventData) { + BaseAgent agent = null; + try { + // Only the agent() lookup can throw; keep the guarded region narrow so a genuinely present + // agent still yields its (validated, non-null) name rather than being swallowed. + agent = invocationContext.agent(); + } catch (RuntimeException e) { + // Fall through to the author/sentinel fallback below. + } + if (agent != null) { + return agent.name(); + } + return eventData.flatMap(EventData::fallbackAgentName).orElse("unknown"); + } + + @CanIgnoreReturnValue + private static EventData.Builder withFallbackAgent( + EventData.Builder builder, @Nullable String author) { + if (author != null && !author.isEmpty()) { + builder.setFallbackAgentName(author); + } + return builder; + } + private ResolvedTraceIds getResolvedTraceIds( InvocationContext invocationContext, Optional eventData) { TraceManager traceManager = state.getTraceManager(invocationContext.invocationId()); @@ -302,8 +408,11 @@ private ResolvedTraceIds getResolvedTraceIds( eventData .flatMap(EventData::traceIdOverride) .orElseGet(() -> traceManager.getTraceId(invocationContext)); - Optional ambientSpanIds = traceManager.getAmbientSpanAndParent(); - SpanIds spanIds = ambientSpanIds.orElse(traceManager.getCurrentSpanAndParent()); + // span_id / parent_span_id must reference the BQAA internal execution tree (spans that are + // 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(); return new ResolvedTraceIds( traceId, @@ -329,6 +438,9 @@ private Map getAttributes( EventData eventData, InvocationContext invocationContext) { Map attributes = new HashMap<>(eventData.extraAttributes()); TraceManager traceManager = state.getTraceManager(invocationContext.invocationId()); + // Populate the root agent name from the invocation context if it has not been set yet, so + // attributes.root_agent_name is a real name rather than the sentinel default. + traceManager.initTraceIfNeeded(invocationContext); attributes.put("root_agent_name", traceManager.getRootAgentName()); eventData.model().ifPresent(m -> attributes.put("model", m)); eventData.modelVersion().ifPresent(mv -> attributes.put("model_version", mv)); @@ -370,7 +482,9 @@ private Map getAttributes( @Override public Completable close() { - return state.close(); + // Deregister the JVM shutdown hook first: an explicit close() supersedes the best-effort drain + // at exit, and leaving the hook registered would pin this plugin for the JVM's lifetime. + return Completable.fromRunnable(this::removeShutdownHook).andThen(state.close()); } @VisibleForTesting @@ -393,16 +507,13 @@ private Optional getCompletedEventData(InvocationContext invocationCo EventData.Builder eventDataBuilder = EventData.builder(); eventDataBuilder.setTraceIdOverride(traceId); eventDataBuilder.setLatency(popped.get().duration()); - // Only override span IDs when no ambient OTel span exists. - // Keep STARTING/COMPLETED pairs consistent. - if (!traceManager.hasAmbientSpan()) { - if (parentSpanId.isPresent()) { - eventDataBuilder.setParentSpanIdOverride(parentSpanId.get()); - } - if (popped.get().spanId() != null) { - eventDataBuilder.setSpanIdOverride(popped.get().spanId()); - } + // Always record the internal execution-tree span so the STARTING/COMPLETED pair stays + // internally joinable and parent_span_id references a logged row, regardless of ambient OTel. + if (parentSpanId.isPresent()) { + eventDataBuilder.setParentSpanIdOverride(parentSpanId.get()); } + // RecordData.spanId() is always populated by the trace manager, so record it unconditionally. + eventDataBuilder.setSpanIdOverride(popped.get().spanId()); return Optional.of(eventDataBuilder.build()); } @@ -446,19 +557,26 @@ public Maybe onEventCallback(InvocationContext invocationContext, Event e if (state.isProcessed(invocationContext.invocationId())) { return Maybe.empty(); } - EventData.Builder eventDataBuilder = - EventData.builder() - .setExtraAttributes( - ImmutableMap.builder() - .put("state_delta", event.actions().stateDelta()) - .put("author", event.author()) - .buildOrThrow()); - Completable logCompletable = - logEvent( - "STATE_DELTA", - invocationContext, - event.content().orElse(null), - Optional.of(eventDataBuilder.build())); + // Only emit STATE_DELTA when there is an actual state change, matching the Python plugin + // (which does not write a STATE_DELTA row for events with an empty state delta). + Completable logCompletable = Completable.complete(); + if (!event.actions().stateDelta().isEmpty()) { + EventData.Builder eventDataBuilder = + withFallbackAgent( + EventData.builder() + .setExtraAttributes( + ImmutableMap.builder() + .put("state_delta", event.actions().stateDelta()) + .put("author", event.author()) + .buildOrThrow()), + event.author()); + logCompletable = + logEvent( + "STATE_DELTA", + invocationContext, + event.content().orElse(null), + Optional.of(eventDataBuilder.build())); + } if (event.content().isPresent() && event.content().get().parts().isPresent()) { for (Part part : event.content().get().parts().get()) { @@ -539,7 +657,11 @@ public Maybe onEventCallback(InvocationContext invocationContext, Event e invocationContext, contentObject, a2aTruncated.isTruncated() || contentTruncated, - Optional.of(EventData.builder().setExtraAttributes(extraAttributes).build()))); + Optional.of( + withFallbackAgent( + EventData.builder().setExtraAttributes(extraAttributes), + event.author()) + .build()))); } } @@ -574,7 +696,11 @@ public Maybe onEventCallback(InvocationContext invocationContext, Event e invocationContext, visibleContent, false, - Optional.of(EventData.builder().setExtraAttributes(extraAttributes).build()))); + Optional.of( + withFallbackAgent( + EventData.builder().setExtraAttributes(extraAttributes), + event.author()) + .build()))); } } @@ -757,8 +883,9 @@ public Maybe afterModelCallback( } } - boolean hasAmbient = traceManager.hasAmbientSpan(); - boolean useOverride = isPopped && !hasAmbient; + // Always record the internal execution-tree span for the final response so parent_span_id + // references a logged row, regardless of any ambient OpenTelemetry span. + boolean useOverride = isPopped; EventData.Builder eventDataBuilder = EventData.builder(); if (!duration.isZero()) { @@ -806,19 +933,17 @@ public Maybe onModelErrorCallback( SpanIds spanIds = traceManager.getCurrentSpanAndParent(); String parentSpanId = spanIds.spanId().orElse(null); - boolean hasAmbient = traceManager.hasAmbientSpan(); EventData.Builder eventDataBuilder = EventData.builder().setStatus("ERROR").setErrorMessage(error.getMessage()); if (popped.isPresent()) { eventDataBuilder.setLatency(popped.get().duration()); } - if (!hasAmbient) { - if (spanId != null) { - eventDataBuilder.setSpanIdOverride(spanId); - } - if (parentSpanId != null) { - eventDataBuilder.setParentSpanIdOverride(parentSpanId); - } + // Always record the internal execution-tree span so parent_span_id references a logged row. + if (spanId != null) { + eventDataBuilder.setSpanIdOverride(spanId); + } + if (parentSpanId != null) { + eventDataBuilder.setParentSpanIdOverride(parentSpanId); } return logEvent("LLM_ERROR", invocationContext, null, Optional.of(eventDataBuilder.build())) .andThen(Maybe.empty()); @@ -863,16 +988,14 @@ public Maybe> afterToolCallback( getToolOrigin(tool)); SpanIds spanIds = traceManager.getCurrentSpanAndParent(); - boolean hasAmbient = traceManager.hasAmbientSpan(); EventData.Builder eventDataBuilder = EventData.builder(); if (popped.isPresent()) { eventDataBuilder.setLatency(popped.get().duration()); } - if (!hasAmbient) { - popped.ifPresent(p -> eventDataBuilder.setSpanIdOverride(p.spanId())); - spanIds.spanId().ifPresent(eventDataBuilder::setParentSpanIdOverride); - } + // Always record the internal execution-tree span so parent_span_id references a logged row. + popped.ifPresent(p -> eventDataBuilder.setSpanIdOverride(p.spanId())); + spanIds.spanId().ifPresent(eventDataBuilder::setParentSpanIdOverride); return logEvent( "TOOL_COMPLETED", @@ -905,17 +1028,15 @@ public Maybe> onToolErrorCallback( .buildOrThrow(); SpanIds spanIds = traceManager.getCurrentSpanAndParent(); - boolean hasAmbient = traceManager.hasAmbientSpan(); EventData.Builder eventDataBuilder = EventData.builder().setStatus("ERROR").setErrorMessage(error.getMessage()); if (popped.isPresent()) { eventDataBuilder.setLatency(popped.get().duration()); } - if (!hasAmbient) { - popped.ifPresent(p -> eventDataBuilder.setSpanIdOverride(p.spanId())); - spanIds.spanId().ifPresent(eventDataBuilder::setParentSpanIdOverride); - } + // Always record the internal execution-tree span so parent_span_id references a logged row. + popped.ifPresent(p -> eventDataBuilder.setSpanIdOverride(p.spanId())); + spanIds.spanId().ifPresent(eventDataBuilder::setParentSpanIdOverride); return logEvent( "TOOL_ERROR", diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java index 92a35b7d7..ec05c2616 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java @@ -222,7 +222,23 @@ public abstract Builder contentFormatter( @CanIgnoreReturnValue public abstract Builder credentials(Credentials credentials); - public abstract BigQueryLoggerConfig build(); + abstract BigQueryLoggerConfig autoBuild(); + + public BigQueryLoggerConfig build() { + BigQueryLoggerConfig config = autoBuild(); + if (config.batchSize() <= 0) { + throw new IllegalArgumentException("batchSize must be positive, got " + config.batchSize()); + } + if (config.queueMaxSize() <= 0) { + throw new IllegalArgumentException( + "queueMaxSize must be positive, got " + config.queueMaxSize()); + } + if (config.maxContentLength() <= 0) { + throw new IllegalArgumentException( + "maxContentLength must be positive, got " + config.maxContentLength()); + } + return config; + } } /** Retry configuration for BigQuery writes. */ 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 9d8be582a..3439606e1 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 @@ -16,8 +16,6 @@ package com.google.adk.plugins.agentanalytics; -import static com.google.adk.plugins.agentanalytics.BigQuerySchema.SCHEMA_VERSION; -import static com.google.adk.plugins.agentanalytics.BigQuerySchema.SCHEMA_VERSION_LABEL_KEY; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.ImmutableMap.toImmutableMap; import static java.util.stream.Collectors.toCollection; @@ -31,6 +29,7 @@ import com.google.cloud.bigquery.Schema; import com.google.cloud.bigquery.StandardSQLTypeName; import com.google.cloud.bigquery.Table; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -42,6 +41,7 @@ import java.util.Optional; import java.util.logging.Level; import java.util.logging.Logger; +import java.util.regex.Pattern; /** Utility for managing BigQuery schema upgrades and analytics views. */ final class BigQueryUtils { @@ -179,8 +179,27 @@ static String getVersionHeaderValue() { return FRAMEWORK_PREFIX + "/" + Version.JAVA_ADK_VERSION; } + private static final Pattern SAFE_IDENTIFIER = Pattern.compile("[A-Za-z0-9_\\-]+"); + + @VisibleForTesting + static boolean isSafeIdentifier(String id) { + return id != null && SAFE_IDENTIFIER.matcher(id).matches(); + } + /** Creates and/or replaces the analytics views in BigQuery. */ static void createAnalyticsViews(BigQuery bigQuery, BigQueryLoggerConfig config) { + // View DDL is assembled by string interpolation; refuse to build it if any operator-supplied + // identifier contains characters (backticks, quotes, dots, semicolons) that could break or + // redirect the statement. + if (!isSafeIdentifier(config.projectId()) + || !isSafeIdentifier(config.datasetId()) + || !isSafeIdentifier(config.tableName()) + || !isSafeIdentifier(config.viewPrefix())) { + logger.warning( + "Skipping analytics view creation: project/dataset/table/viewPrefix contains characters" + + " that are unsafe to interpolate into DDL."); + return; + } for (Map.Entry> entry : EVENT_VIEW_DEFS.entrySet()) { String eventType = entry.getKey(); ImmutableList extraCols = entry.getValue(); @@ -215,66 +234,66 @@ static void createAnalyticsViews(BigQuery bigQuery, BigQueryLoggerConfig config) } } - /** Adds missing columns to an existing table if the schema version has changed. */ - static void maybeUpgradeSchema(BigQuery bigQuery, Table existingTable) { - String storedVersion = - Optional.ofNullable(existingTable.getLabels()) - .map(labels -> labels.get(SCHEMA_VERSION_LABEL_KEY)) - .orElse(""); - - if (storedVersion.equals(SCHEMA_VERSION)) { - return; - } - + /** + * Adds missing columns to an existing table if the actual schema is behind the desired schema. + */ + static boolean maybeUpgradeSchema(BigQuery bigQuery, Table existingTable) { + // Always diff the actual table schema against the desired schema rather than trusting the + // stored version label alone: a table stamped with the current label can still be missing + // columns (e.g. it was created by an older build), and those must be reconciled. SchemaDiff diff = schemaFieldsMatch( existingTable.getDefinition().getSchema().getFields(), BigQuerySchema.getEventsSchema().getFields()); - if (!diff.newTopLevelFields().isEmpty() || !diff.updatedRecordFields().isEmpty()) { - ImmutableMap updatedFields = - diff.updatedRecordFields().stream().collect(toImmutableMap(Field::getName, f -> f)); - ImmutableSet updatedNames = updatedFields.keySet(); - - List mergedFields = new ArrayList<>(); - for (Field f : existingTable.getDefinition().getSchema().getFields()) { - if (updatedNames.contains(f.getName())) { - mergedFields.add(updatedFields.get(f.getName())); - } else { - mergedFields.add(f); - } - } - mergedFields.addAll(diff.newTopLevelFields()); + if (diff.newTopLevelFields().isEmpty() && diff.updatedRecordFields().isEmpty()) { + // Nothing to reconcile; the table already satisfies the desired schema. + return true; + } - logger.info( - String.format( - "Auto-upgrading table %s: new columns %s, updated RECORD fields %s", - existingTable.getTableId(), - diff.newTopLevelFields().stream().map(Field::getName).collect(toImmutableList()), - diff.updatedRecordFields().stream() - .map(Field::getName) - .collect(toCollection(ArrayList::new)))); + ImmutableMap updatedFields = + diff.updatedRecordFields().stream().collect(toImmutableMap(Field::getName, f -> f)); + ImmutableSet updatedNames = updatedFields.keySet(); - try { - Map labels = - new HashMap<>(Optional.ofNullable(existingTable.getLabels()).orElse(ImmutableMap.of())); - labels.put(BigQuerySchema.SCHEMA_VERSION_LABEL_KEY, BigQuerySchema.SCHEMA_VERSION); - - Table updatedTable = - existingTable.toBuilder() - .setDefinition( - existingTable.getDefinition().toBuilder() - .setSchema(Schema.of(mergedFields)) - .build()) - .setLabels(labels) - .build(); - - var unused = bigQuery.update(updatedTable); - } catch (BigQueryException e) { - logger.log( - Level.WARNING, "Schema auto-upgrade failed for " + existingTable.getTableId(), e); + List mergedFields = new ArrayList<>(); + for (Field f : existingTable.getDefinition().getSchema().getFields()) { + if (updatedNames.contains(f.getName())) { + mergedFields.add(updatedFields.get(f.getName())); + } else { + mergedFields.add(f); } } + mergedFields.addAll(diff.newTopLevelFields()); + + logger.info( + String.format( + "Auto-upgrading table %s: new columns %s, updated RECORD fields %s", + existingTable.getTableId(), + diff.newTopLevelFields().stream().map(Field::getName).collect(toImmutableList()), + diff.updatedRecordFields().stream() + .map(Field::getName) + .collect(toCollection(ArrayList::new)))); + + try { + Map labels = + new HashMap<>(Optional.ofNullable(existingTable.getLabels()).orElse(ImmutableMap.of())); + labels.put(BigQuerySchema.SCHEMA_VERSION_LABEL_KEY, BigQuerySchema.SCHEMA_VERSION); + + Table updatedTable = + existingTable.toBuilder() + .setDefinition( + existingTable.getDefinition().toBuilder() + .setSchema(Schema.of(mergedFields)) + .build()) + .setLabels(labels) + .build(); + + var unused = bigQuery.update(updatedTable); + return true; + } catch (BigQueryException e) { + logger.log(Level.WARNING, "Schema auto-upgrade failed for " + existingTable.getTableId(), e); + return false; + } } private static SchemaDiff schemaFieldsMatch(FieldList existing, FieldList desired) { @@ -292,6 +311,9 @@ private static SchemaDiff schemaFieldsMatch(FieldList existing, FieldList desire } else if (desiredField.getType().getStandardType().equals(StandardSQLTypeName.STRUCT) && existingField.getType().getStandardType().equals(StandardSQLTypeName.STRUCT) && desiredField.getSubFields() != null) { + // Mode drift on the STRUCT column itself (e.g. NULLABLE vs REPEATED) is just as + // un-upgradeable as on a scalar; check it before recursing into subfields. + warnOnIncompatibleDrift(existingField, desiredField); SchemaDiff subDiff = schemaFieldsMatch(existingField.getSubFields(), desiredField.getSubFields()); @@ -314,11 +336,44 @@ private static SchemaDiff schemaFieldsMatch(FieldList existing, FieldList desire .setType(StandardSQLTypeName.STRUCT, FieldList.of(mergedSub)) .build()); } + } else { + warnOnIncompatibleDrift(existingField, desiredField); } } return new SchemaDiff(ImmutableList.copyOf(newFields), ImmutableList.copyOf(updatedRecords)); } + // Additive auto-upgrade cannot reconcile a type or mode change on an existing column + // (including nested non-STRUCT fields, since schemaFieldsMatch recurses into STRUCTs). Surface + // it instead of silently ignoring it, since it otherwise appears later as opaque Storage + // Write append failures. + private static void warnOnIncompatibleDrift(Field existingField, Field desiredField) { + boolean typeDrift = + !desiredField.getType().getStandardType().equals(existingField.getType().getStandardType()); + boolean modeDrift = !modesEqual(existingField.getMode(), desiredField.getMode()); + if (typeDrift || modeDrift) { + logger.warning( + String.format( + "Incompatible schema drift on column '%s': table has %s/%s but the plugin expects" + + " %s/%s. This cannot be auto-upgraded; writes may fail until the column is" + + " fixed manually.", + desiredField.getName(), + existingField.getType().getStandardType(), + normalizeMode(existingField.getMode()), + desiredField.getType().getStandardType(), + normalizeMode(desiredField.getMode()))); + } + } + + // BigQuery leaves Field.getMode() null to mean NULLABLE; normalize before comparing. + private static Field.Mode normalizeMode(Field.Mode mode) { + return mode == null ? Field.Mode.NULLABLE : mode; + } + + private static boolean modesEqual(Field.Mode a, Field.Mode b) { + return normalizeMode(a) == normalizeMode(b); + } + private record SchemaDiff( ImmutableList newTopLevelFields, ImmutableList updatedRecordFields) {} diff --git a/core/src/main/java/com/google/adk/plugins/agentanalytics/EventData.java b/core/src/main/java/com/google/adk/plugins/agentanalytics/EventData.java index 8fd95a070..41d0ee312 100644 --- a/core/src/main/java/com/google/adk/plugins/agentanalytics/EventData.java +++ b/core/src/main/java/com/google/adk/plugins/agentanalytics/EventData.java @@ -1,3 +1,19 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + package com.google.adk.plugins.agentanalytics; import com.google.auto.value.AutoValue; @@ -6,7 +22,7 @@ import java.util.Map; import java.util.Optional; -/** Typed container for structured fields passed to _log_event. */ +/** Typed container for structured fields passed to the plugin's event-logging method. */ @AutoValue abstract class EventData { abstract Optional spanIdOverride(); @@ -31,6 +47,10 @@ abstract class EventData { abstract Optional traceIdOverride(); + // Fallback name for the `agent` column when the InvocationContext has no current agent (e.g. + // workflow-driven callbacks). Mirrors the Python plugin's fallback to Event.author. + abstract Optional fallbackAgentName(); + static Builder builder() { return new AutoValue_EventData.Builder().setStatus("OK").setExtraAttributes(ImmutableMap.of()); } @@ -59,6 +79,8 @@ abstract static class Builder { abstract Builder setTraceIdOverride(String value); + abstract Builder setFallbackAgentName(String value); + abstract EventData build(); } } 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 58049a9ef..922e03c02 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 @@ -24,7 +24,9 @@ import com.fasterxml.jackson.databind.node.ObjectNode; import com.google.auto.value.AutoValue; import com.google.common.base.Utf8; +import com.google.common.collect.ImmutableSet; import java.util.IdentityHashMap; +import java.util.Locale; import java.util.Map; import java.util.Set; import java.util.logging.Logger; @@ -36,6 +38,23 @@ final class JsonFormatter { static final ObjectMapper mapper = new ObjectMapper().findAndRegisterModules(); static final String TRUNCATION_SUFFIX = "...[truncated]"; static final String CYCLE_DETECTED_MESSAGE = "[cycle detected]"; + static final String MAX_DEPTH_MESSAGE = "[max depth exceeded]"; + static final String REDACTED_MESSAGE = "[REDACTED]"; + // Guard against unbounded recursion on deeply nested (non-cyclic) payloads. + static final int MAX_TRUNCATE_DEPTH = 200; + + // Keys whose values are redacted before logging. Mirrors the Python BQAA plugin's + // _SENSITIVE_KEYS (OAuth tokens / secrets); matching is case-insensitive, plus any + // key prefixed with "temp:" (ADK temporary session state). + private static final ImmutableSet SENSITIVE_KEYS = + ImmutableSet.of( + "client_secret", "access_token", "refresh_token", "id_token", "api_key", "password"); + private static final String TEMP_KEY_PREFIX = "temp:"; + + private static boolean isSensitiveKey(String key) { + String lower = key.toLowerCase(Locale.ROOT); + return SENSITIVE_KEYS.contains(lower) || lower.startsWith(TEMP_KEY_PREFIX); + } @AutoValue abstract static class TruncationResult { @@ -55,12 +74,14 @@ static TruncationResult smartTruncate(Object obj, int maxLength) { } try { if (obj instanceof JsonNode jsonNode) { - return recursiveSmartTruncate(jsonNode, maxLength, newSetFromMap(new IdentityHashMap<>())); + return recursiveSmartTruncate( + jsonNode, maxLength, newSetFromMap(new IdentityHashMap<>()), 0); } return recursiveSmartTruncate( - mapper.valueToTree(obj), maxLength, newSetFromMap(new IdentityHashMap<>())); + mapper.valueToTree(obj), maxLength, newSetFromMap(new IdentityHashMap<>()), 0); } catch (IllegalArgumentException e) { - // Fallback for types that mapper can't handle directly as a tree + // Fallback for types that mapper can't handle directly as a tree. + logger.fine("smartTruncate falling back to string conversion: " + e.getMessage()); return truncateWithStatus(safeToString(obj), maxLength); } } @@ -87,7 +108,10 @@ static String safeToString(Object obj) { } private static TruncationResult recursiveSmartTruncate( - JsonNode node, int maxLength, Set visited) { + JsonNode node, int maxLength, Set visited, int depth) { + if (depth > MAX_TRUNCATE_DEPTH) { + return TruncationResult.create(mapper.valueToTree(MAX_DEPTH_MESSAGE), true); + } if (node.isContainerNode()) { if (visited.contains(node)) { return TruncationResult.create(mapper.valueToTree(CYCLE_DETECTED_MESSAGE), true); @@ -106,7 +130,14 @@ private static TruncationResult recursiveSmartTruncate( ObjectNode newNode = mapper.createObjectNode(); Set> properties = node.properties(); for (Map.Entry entry : properties) { - TruncationResult res = recursiveSmartTruncate(entry.getValue(), maxLength, visited); + // Redact sensitive values without descending into them. Per parity with the + // Python plugin, redaction does not set the is_truncated flag. + if (isSensitiveKey(entry.getKey())) { + newNode.set(entry.getKey(), mapper.valueToTree(REDACTED_MESSAGE)); + continue; + } + TruncationResult res = + recursiveSmartTruncate(entry.getValue(), maxLength, visited, depth + 1); newNode.set(entry.getKey(), res.node()); isTruncated = isTruncated || res.isTruncated(); } @@ -114,7 +145,7 @@ private static TruncationResult recursiveSmartTruncate( } else if (node.isArray()) { ArrayNode newNode = mapper.createArrayNode(); for (JsonNode element : node) { - TruncationResult res = recursiveSmartTruncate(element, maxLength, visited); + TruncationResult res = recursiveSmartTruncate(element, maxLength, visited, depth + 1); newNode.add(res.node()); isTruncated = isTruncated || res.isTruncated(); } 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 d1826ec5e..afb3a1ea3 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 @@ -1,3 +1,19 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + package com.google.adk.plugins.agentanalytics; import static com.google.adk.plugins.agentanalytics.BigQueryUtils.getVersionHeaderValue; @@ -64,6 +80,11 @@ class PluginState { private final Parser parser; private final ConcurrentHashMap>> pendingTasks = new ConcurrentHashMap<>(); + // 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(); PluginState(BigQueryLoggerConfig config) throws IOException { this.config = config; @@ -106,7 +127,7 @@ ScheduledExecutorService getExecutor() { boolean isProcessed(String invocationId) { boolean isProcessed = processedInvocations.getIfPresent(invocationId) != null; if (isProcessed) { - logger.info("Invocation ID: " + invocationId + " already processed"); + logger.fine("Invocation ID: " + invocationId + " already processed"); } return isProcessed; } @@ -142,6 +163,9 @@ protected StreamWriter createWriter() { .setTraceId(BigQueryUtils.getVersionHeaderValue() + ":" + UUID.randomUUID()) .setRetrySettings(retrySettings) .setWriterSchema(BigQuerySchema.getArrowSchema()) + // Route Storage Write append RPCs to the dataset's region. Without this, appends to any + // location other than the US multi-region can fail with stream-not-found errors. + .setLocation(config.location()) .build(); } catch (Exception e) { throw new VerifyException("Failed to create StreamWriter for " + streamName, e); @@ -222,6 +246,11 @@ protected Set> getPendingTasksForInvocation(String invoc return pendingTasks.computeIfAbsent(invocationId, k -> ConcurrentHashMap.newKeySet()); } + // Relies on reference (identity) equality of CompletableFuture: the exact same future instance is + // added here and removed on completion, which is well-defined under the default Object.equals / + // hashCode. The set must remain concurrent (ConcurrentHashMap.newKeySet), so a JDK + // IdentityHashMap-based set is not an option. + @SuppressWarnings("CollectionUndefinedEquality") void addPendingTask(String invocationId, CompletableFuture task) { Set> tasks = getPendingTasksForInvocation(invocationId); tasks.add(task); @@ -236,7 +265,7 @@ Completable ensureInvocationCompleted(String invocationId) { Completable.fromCompletionStage( CompletableFuture.allOf(tasks.toArray(new CompletableFuture[0]))); } - logger.info("Waiting for pending tasks to complete for invocation ID: " + invocationId); + logger.fine("Waiting for pending tasks to complete for invocation ID: " + invocationId); return tasksState .timeout(config.shutdownTimeout().toMillis(), MILLISECONDS) .doOnError( @@ -258,16 +287,45 @@ Completable ensureInvocationCompleted(String invocationId) { if (processor != null) { processor.flush(); processor.close(); + // Fold after close() so rows dropped during the final drain are also counted. + foldDropStats(processor); } TraceManager traceManager = removeTraceManager(invocationId); if (traceManager != null) { traceManager.clearStack(); } - logger.info("Removing pending tasks for invocation ID: " + invocationId); + logger.fine("Removing pending tasks for invocation ID: " + invocationId); pendingTasks.remove(invocationId); }); } + private void foldDropStats(BatchProcessor processor) { + ImmutableMap stats = processor.getDropStats(); + droppedQueueFull.addAndGet(stats.getOrDefault("queue_full", 0L)); + droppedAppendError.addAndGet(stats.getOrDefault("append_error", 0L)); + droppedSerializationError.addAndGet(stats.getOrDefault("serialization_error", 0L)); + } + + /** + * Aggregated dropped-row counters across closed and still-live BatchProcessors. Non-zero values + * indicate analytics rows that never reached BigQuery. + */ + ImmutableMap getDropStats() { + long queueFull = droppedQueueFull.get(); + long appendError = droppedAppendError.get(); + long serializationError = droppedSerializationError.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); + } + return ImmutableMap.of( + "queue_full", queueFull, + "append_error", appendError, + "serialization_error", serializationError); + } + Completable close() { ImmutableList> tasks = pendingTasks.values().stream().flatMap(Set::stream).collect(toImmutableList()); @@ -291,6 +349,8 @@ Completable close() { () -> { for (BatchProcessor processor : getBatchProcessors()) { processor.close(); + // Fold after close() so rows dropped during the final drain are also counted. + foldDropStats(processor); } for (TraceManager traceManager : getTraceManagers()) { traceManager.clearStack(); 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 a02ea00b4..e455fc509 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 @@ -24,10 +24,10 @@ import io.opentelemetry.api.trace.SpanContext; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.context.Context; -import io.opentelemetry.sdk.trace.ReadableSpan; import java.time.Duration; import java.time.Instant; import java.util.Iterator; +import java.util.Objects; import java.util.Optional; import java.util.UUID; import java.util.concurrent.ConcurrentLinkedDeque; @@ -43,8 +43,10 @@ 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 final ConcurrentLinkedDeque records = new ConcurrentLinkedDeque<>(); - private String rootAgentName = "_bq_analytics_root_agent_name"; + private String rootAgentName = DEFAULT_ROOT_AGENT_NAME; private String activeInvocationId = "_bq_analytics_active_invocation_id"; private final Tracer tracer; @@ -103,8 +105,26 @@ public String getRootAgentName() { } public void initTrace(InvocationContext context) { - String rootAgentName = context.agent().rootAgent().name(); - this.rootAgentName = rootAgentName; + var rootAgent = context.agent().rootAgent(); + if (rootAgent != null && rootAgent.name() != null) { + this.rootAgentName = rootAgent.name(); + } + } + + /** + * Sets the root agent name from the invocation context if it is still the sentinel default. + * Null-safe: workflow-driven callbacks with no current agent leave the sentinel in place for a + * later event to resolve. + */ + public void initTraceIfNeeded(InvocationContext context) { + if (!Objects.equals(rootAgentName, DEFAULT_ROOT_AGENT_NAME)) { + return; + } + try { + initTrace(context); + } catch (RuntimeException e) { + // Leave the sentinel; a subsequent event may be able to resolve the root agent. + } } public String getTraceId(InvocationContext context) { @@ -123,10 +143,6 @@ public String getTraceId(InvocationContext context) { return context.invocationId(); } - public boolean hasAmbientSpan() { - return Span.current().getSpanContext().isValid(); - } - @CanIgnoreReturnValue public String pushSpan(String spanName) { Context parentContext = Context.current(); @@ -173,7 +189,7 @@ public void ensureInvocationSpan(InvocationContext context) { if (currentInv.equals(activeInvocationId)) { return; } - logger.info("Clearing stale span records from previous invocation."); + logger.fine("Clearing stale span records from previous invocation."); clearStack(); } @@ -225,22 +241,6 @@ public SpanIds getCurrentSpanAndParent() { return SpanIds.create(spanId, parentId); } - Optional getAmbientSpanAndParent() { - Span ambient = Span.current(); - if (!ambient.getSpanContext().isValid()) { - return Optional.empty(); - } - String spanId = ambient.getSpanContext().getSpanId(); - String parentSpanId = null; - if (ambient instanceof ReadableSpan readableSpan) { - SpanContext parentCtx = readableSpan.getParentSpanContext(); - if (parentCtx != null && parentCtx.isValid()) { - parentSpanId = parentCtx.getSpanId(); - } - } - return Optional.of(SpanIds.create(spanId, parentSpanId)); - } - public Optional getCurrentSpanId() { if (records.isEmpty()) { return Optional.empty(); 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 4f4350d1a..2355c4def 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 @@ -263,6 +263,8 @@ public void flush_handlesBigQueryErrorResponse() throws Exception { batchProcessor.flush(); verify(mockWriter).append(any(ArrowRecordBatch.class)); + // A BigQuery error response must count the batch as dropped under "append_error". + assertEquals(1L, (long) batchProcessor.getDropStats().get("append_error")); } @Test @@ -277,6 +279,8 @@ public void flush_handlesGenericExceptionDuringAppend() throws Exception { batchProcessor.flush(); verify(mockWriter).append(any(ArrowRecordBatch.class)); + // A generic append exception must count the batch as dropped under "append_error". + assertEquals(1L, (long) batchProcessor.getDropStats().get("append_error")); } @Test 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 a4def5bf7..8ff58d473 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 @@ -27,8 +27,11 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyMap; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.atLeast; import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -51,6 +54,7 @@ import com.google.auth.Credentials; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQueryOptions; +import com.google.cloud.bigquery.Field.Mode; import com.google.cloud.bigquery.FieldList; import com.google.cloud.bigquery.QueryJobConfiguration; import com.google.cloud.bigquery.StandardSQLTypeName; @@ -389,6 +393,51 @@ public void ensureTableExists_calledOnlyOnce() throws Exception { verify(mockBigQuery).getTable(any(TableId.class)); } + @Test + public void ensureTableExists_retriesAfterFailure() throws Exception { + when(mockBigQuery.getTable(any(TableId.class))) + .thenThrow(new RuntimeException("Table check failed")); + Content content = Content.builder().build(); + + plugin.onUserMessageCallback(mockInvocationContext, content).blockingSubscribe(); + plugin.onUserMessageCallback(mockInvocationContext, content).blockingSubscribe(); + + // A failed bootstrap must leave the table un-ensured so it is retried on the next event, rather + // than being masked as ready. With retry, getTable is invoked once per event. + verify(mockBigQuery, times(2)).getTable(any(TableId.class)); + } + + @Test + public void afterAgentCallback_stampsInternalExecutionTreeSpanIds() throws Exception { + CallbackContext callbackContext = mock(CallbackContext.class); + when(callbackContext.invocationContext()).thenReturn(mockInvocationContext); + + // Establish the invocation-level span, then push a child agent span. + plugin + .onUserMessageCallback(mockInvocationContext, Content.builder().build()) + .blockingSubscribe(); + plugin.beforeAgentCallback(fakeAgent, callbackContext).blockingSubscribe(); + + TraceManager.SpanIds current = state.getTraceManager("invocation_id").getCurrentSpanAndParent(); + String agentSpanId = current.spanId().orElseThrow(); + String invocationSpanId = current.parentSpanId().orElseThrow(); + + // Completing the agent pops the agent span and must stamp the row from the internal execution + // tree: span_id = the popped agent span, parent_span_id = the enclosing invocation span. + plugin.afterAgentCallback(fakeAgent, callbackContext).blockingSubscribe(); + + Map completedRow = null; + Map row; + while ((row = state.getBatchProcessor("invocation_id").queue.poll()) != null) { + if (Objects.equals(row.get("event_type"), "AGENT_COMPLETED")) { + completedRow = row; + } + } + assertNotNull("AGENT_COMPLETED row not found", completedRow); + assertEquals(agentSpanId, completedRow.get("span_id")); + assertEquals(invocationSpanId, completedRow.get("parent_span_id")); + } + @Test public void arrowSchema_handlesNestedFields() { Schema schema = BigQuerySchema.getArrowSchema(); @@ -545,6 +594,7 @@ public void onEventCallback_populatesCorrectFields() throws Exception { Event event = Event.builder() .author("agent_author") + .actions(EventActions.builder().stateDelta(ImmutableMap.of("key", "new_value")).build()) .content(Content.fromParts(Part.fromText("event content"))) .build(); @@ -556,10 +606,94 @@ public void onEventCallback_populatesCorrectFields() throws Exception { assertEquals("agent_name", row.get("agent")); ObjectNode attributes = (ObjectNode) row.get("attributes"); assertEquals("agent_author", attributes.get("author").asText()); + assertEquals("new_value", attributes.get("state_delta").get("key").asText()); assertTrue(row.get("content").toString().contains("event content")); assertEquals(false, row.get("is_truncated")); } + @Test + public void onEventCallback_noCurrentAgent_fallsBackToEventAuthor() throws Exception { + // Workflow-driven callbacks may have no current agent; the "agent" column must fall back to the + // event author rather than the "unknown" sentinel. + when(mockInvocationContext.agent()).thenReturn(null); + Event event = + Event.builder() + .author("agent_author") + .actions(EventActions.builder().stateDelta(ImmutableMap.of("key", "new_value")).build()) + .content(Content.fromParts(Part.fromText("event content"))) + .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")); + assertEquals("agent_author", row.get("agent")); + } + + @Test + public void onEventCallback_emptyAuthorNoCurrentAgent_fallsBackToUnknownSentinel() + throws Exception { + // An empty author is not a usable fallback: withFallbackAgent guards on + // `author != null && !author.isEmpty()`, so fallbackAgentName stays unset and resolveAgentName + // yields the "unknown" sentinel rather than an empty agent name. Pins the `&&` against a + // `||`-mutation (go/mutation-testing), which would stamp "" as the agent for empty-author + // events. + when(mockInvocationContext.agent()).thenReturn(null); + Event event = + Event.builder() + .author("") + .actions(EventActions.builder().stateDelta(ImmutableMap.of("key", "new_value")).build()) + .content(Content.fromParts(Part.fromText("event content"))) + .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")); + assertEquals("unknown", row.get("agent")); + } + + @Test + public void onEventCallback_nullAuthorNoCurrentAgent_fallsBackToUnknownSentinelWithoutNpe() + throws Exception { + // A null author must be short-circuited by the `author != null` half of withFallbackAgent's + // guard so `author.isEmpty()` is never dereferenced. Exercised via the AGENT_RESPONSE path + // (whose extraAttributes tolerate a null author, unlike the STATE_DELTA map). Pins the `&&` + // against a `||`-mutation (go/mutation-testing), which would NPE on null authors. + when(mockInvocationContext.agent()).thenReturn(null); + Event event = + Event.builder() + .id("evt-id") + .content(Content.fromParts(Part.fromText("agent final answer"))) + .build(); + assertNull("Precondition: author must be null for this test", event.author()); + + plugin.onEventCallback(mockInvocationContext, event).blockingSubscribe(); + CompletableFuture.allOf( + state + .getPendingTasksForInvocation("invocation_id") + .toArray(new CompletableFuture[0])) + .join(); + + Map row = state.getBatchProcessor("invocation_id").queue.poll(); + assertNotNull("AGENT_RESPONSE row not found in queue", row); + assertEquals("AGENT_RESPONSE", row.get("event_type")); + assertEquals("unknown", row.get("agent")); + } + + @Test + public void onEventCallback_emptyStateDelta_doesNotEmitStateDelta() throws Exception { + Event event = Event.builder().author("agent_author").build(); + + plugin.onEventCallback(mockInvocationContext, event).blockingSubscribe(); + + assertNull( + "No STATE_DELTA row should be emitted for an empty state delta", + state.getBatchProcessor("invocation_id").queue.poll()); + } + @Test public void onEventCallback_withA2AMetadata_emitsA2AInteraction() throws Exception { Event event = @@ -581,10 +715,6 @@ public void onEventCallback_withA2AMetadata_emitsA2AInteraction() throws Excepti .toArray(new CompletableFuture[0])) .join(); - Map stateDeltaRow = state.getBatchProcessor("invocation_id").queue.poll(); - assertNotNull(stateDeltaRow); - assertEquals("STATE_DELTA", stateDeltaRow.get("event_type")); - Map a2aRow = state.getBatchProcessor("invocation_id").queue.poll(); assertNotNull("A2A_INTERACTION row not found in queue", a2aRow); assertEquals("A2A_INTERACTION", a2aRow.get("event_type")); @@ -624,10 +754,6 @@ public void onEventCallback_agentResponse_emitsAgentResponse() throws Exception .toArray(new CompletableFuture[0])) .join(); - Map stateDeltaRow = state.getBatchProcessor("invocation_id").queue.poll(); - assertNotNull(stateDeltaRow); - assertEquals("STATE_DELTA", stateDeltaRow.get("event_type")); - Map agentResponseRow = state.getBatchProcessor("invocation_id").queue.poll(); assertNotNull("AGENT_RESPONSE row not found in queue", agentResponseRow); assertEquals("AGENT_RESPONSE", agentResponseRow.get("event_type")); @@ -670,10 +796,6 @@ public void onEventCallback_skipSummarizationAndFunctionCall_doesNotEmitAgentRes .toArray(new CompletableFuture[0])) .join(); - Map stateDeltaRow = state.getBatchProcessor("invocation_id").queue.poll(); - assertNotNull(stateDeltaRow); - assertEquals("STATE_DELTA", stateDeltaRow.get("event_type")); - Map nextRow = state.getBatchProcessor("invocation_id").queue.poll(); assertNull("No AGENT_RESPONSE row should be emitted", nextRow); } @@ -698,10 +820,6 @@ public void onEventCallback_longRunningToolIdsPresent_doesNotEmitAgentResponse() .toArray(new CompletableFuture[0])) .join(); - Map stateDeltaRow = state.getBatchProcessor("invocation_id").queue.poll(); - assertNotNull(stateDeltaRow); - assertEquals("STATE_DELTA", stateDeltaRow.get("event_type")); - Map nextRow = state.getBatchProcessor("invocation_id").queue.poll(); assertNull("No AGENT_RESPONSE row should be emitted", nextRow); } @@ -725,9 +843,6 @@ public void onEventCallback_withA2ARequestOnlyMetadata_emitsA2AInteraction() thr .toArray(new CompletableFuture[0])) .join(); - Map stateDeltaRow = state.getBatchProcessor("invocation_id").queue.poll(); - assertNotNull(stateDeltaRow); - Map a2aRow = state.getBatchProcessor("invocation_id").queue.poll(); assertNotNull("A2A_INTERACTION row not found in queue", a2aRow); assertEquals("A2A_INTERACTION", a2aRow.get("event_type")); @@ -781,10 +896,6 @@ protected StreamWriter createWriter() { .toArray(new CompletableFuture[0])) .join(); - // Consume STATE_DELTA - Map stateDeltaRow = customState.getBatchProcessor("invocation_id").queue.poll(); - assertNotNull(stateDeltaRow); - // Get AGENT_RESPONSE Map agentResponseRow = customState.getBatchProcessor("invocation_id").queue.poll(); @@ -828,6 +939,25 @@ public void onModelErrorCallback_populatesCorrectFields() throws Exception { row.containsKey("is_truncated")); } + @Test + public void onModelErrorCallback_stampsPoppedSpanId() throws Exception { + CallbackContext mockCallbackContext = mock(CallbackContext.class); + when(mockCallbackContext.invocationContext()).thenReturn(mockInvocationContext); + LlmRequest.Builder mockLlmRequestBuilder = mock(LlmRequest.Builder.class); + + String llmSpanId = state.getTraceManager("invocation_id").pushSpan("llm_request"); + plugin + .onModelErrorCallback( + mockCallbackContext, mockLlmRequestBuilder, new RuntimeException("boom")) + .blockingSubscribe(); + + Map row = state.getBatchProcessor("invocation_id").queue.poll(); + assertNotNull("Row not found in queue", row); + assertEquals("LLM_ERROR", row.get("event_type")); + // The error row's span_id must come from the popped internal span, not the post-pop stack. + assertEquals(llmSpanId, row.get("span_id")); + } + @Test public void afterModelCallback_populatesCorrectFields() throws Exception { CallbackContext mockCallbackContext = mock(CallbackContext.class); @@ -940,6 +1070,47 @@ public AgentOrigin toolOrigin() { assertEquals("A2A", contentMap.get("tool_origin").asText()); } + @Test + public void afterToolCallback_stampsPoppedToolSpanId() throws Exception { + ToolContext mockToolContext = mock(ToolContext.class); + when(mockToolContext.invocationContext()).thenReturn(mockInvocationContext); + BaseTool mockTool = mock(BaseTool.class); + when(mockTool.name()).thenReturn("test_tool"); + + // Establish the invocation span first (so afterTool's ensureInvocationSpan keeps the stack), + // then push the tool span that afterTool must pop and stamp onto the row. + plugin + .onUserMessageCallback(mockInvocationContext, Content.builder().build()) + .blockingSubscribe(); + String toolSpanId = state.getTraceManager("invocation_id").pushSpan("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() + .parentSpanId() + .orElseThrow(); + + plugin + .afterToolCallback(mockTool, ImmutableMap.of(), mockToolContext, ImmutableMap.of("r", "v")) + .blockingSubscribe(); + + Map completedRow = null; + Map row; + while ((row = state.getBatchProcessor("invocation_id").queue.poll()) != null) { + if (Objects.equals(row.get("event_type"), "TOOL_COMPLETED")) { + completedRow = row; + } + } + assertNotNull("TOOL_COMPLETED row not found", completedRow); + // span_id must be the popped tool span, not the enclosing invocation span left on the stack. + assertEquals(toolSpanId, completedRow.get("span_id")); + // parent_span_id must reference the enclosing invocation span from the post-pop stack top. + assertEquals(invocationSpanId, completedRow.get("parent_span_id")); + } + @Test public void logEvent_includesSessionMetadata_whenEnabled() throws Exception { // Config default has logSessionMetadata(true) @@ -1109,8 +1280,10 @@ public void maybeUpgradeSchema_addsNewTopLevelField() throws Exception { when(mockTableBuilder.setLabels(anyMap())).thenReturn(mockTableBuilder); when(mockTableBuilder.build()).thenReturn(mockTable); - BigQueryUtils.maybeUpgradeSchema(mockBigQuery, mockTable); + boolean upgraded = BigQueryUtils.maybeUpgradeSchema(mockBigQuery, mockTable); + // A successful upgrade must report the table as ready. + assertTrue(upgraded); ArgumentCaptor definitionCaptor = ArgumentCaptor.forClass(StandardTableDefinition.class); verify(mockTableBuilder).setDefinition(definitionCaptor.capture()); @@ -1161,7 +1334,7 @@ public void maybeUpgradeSchema_addsNewNestedField() throws Exception { when(mockTableBuilder.setLabels(anyMap())).thenReturn(mockTableBuilder); when(mockTableBuilder.build()).thenReturn(mockTable); - BigQueryUtils.maybeUpgradeSchema(mockBigQuery, mockTable); + var unused = BigQueryUtils.maybeUpgradeSchema(mockBigQuery, mockTable); ArgumentCaptor definitionCaptor = ArgumentCaptor.forClass(StandardTableDefinition.class); @@ -1173,6 +1346,191 @@ public void maybeUpgradeSchema_addsNewNestedField() throws Exception { verify(mockBigQuery).update(any(Table.class)); } + @Test + public void maybeUpgradeSchema_warnsOnStructModeDrift() throws Exception { + Table mockTable = mock(Table.class); + when(mockTable.getTableId()).thenReturn(TableId.of("project", "dataset", "table")); + when(mockTable.getLabels()).thenReturn(ImmutableMap.of()); + + // Existing table has 'content_parts' as a NULLABLE STRUCT instead of the expected REPEATED + ImmutableList initialFields = + BigQuerySchema.getEventsSchema().getFields().stream() + .map( + f -> + f.getName().equals("content_parts") + ? f.toBuilder().setMode(Mode.NULLABLE).build() + : f) + .collect(toImmutableList()); + + StandardTableDefinition tableDefinition = + StandardTableDefinition.newBuilder() + .setSchema(com.google.cloud.bigquery.Schema.of(initialFields)) + .build(); + when(mockTable.getDefinition()).thenReturn(tableDefinition); + + Logger logger = Logger.getLogger(BigQueryUtils.class.getName()); + Handler mockLogHandler = mock(Handler.class); + logger.addHandler(mockLogHandler); + try { + var unused = BigQueryUtils.maybeUpgradeSchema(mockBigQuery, mockTable); + } finally { + logger.removeHandler(mockLogHandler); + } + + ArgumentCaptor captor = ArgumentCaptor.forClass(LogRecord.class); + verify(mockLogHandler, atLeastOnce()).publish(captor.capture()); + assertTrue( + "Should have warned about STRUCT mode drift on content_parts", + captor.getAllValues().stream() + .anyMatch( + record -> + Objects.equals(record.getLevel(), Level.WARNING) + && record + .getMessage() + .contains("Incompatible schema drift on column 'content_parts'"))); + + // Mode drift alone is not auto-upgradeable, so no table update should be attempted. + verify(mockBigQuery, never()).update(any(Table.class)); + } + + @Test + public void maybeUpgradeSchema_noChanges_returnsTrueWithoutUpdateOrDriftWarning() + throws Exception { + Table mockTable = mock(Table.class); + when(mockTable.getTableId()).thenReturn(TableId.of("project", "dataset", "table")); + StandardTableDefinition tableDefinition = + StandardTableDefinition.newBuilder().setSchema(BigQuerySchema.getEventsSchema()).build(); + when(mockTable.getDefinition()).thenReturn(tableDefinition); + + Logger logger = Logger.getLogger(BigQueryUtils.class.getName()); + Handler mockLogHandler = mock(Handler.class); + logger.addHandler(mockLogHandler); + boolean upgraded; + try { + upgraded = BigQueryUtils.maybeUpgradeSchema(mockBigQuery, mockTable); + } finally { + logger.removeHandler(mockLogHandler); + } + + // When the existing schema already matches, the table is ready and no update is attempted. + assertTrue(upgraded); + verify(mockBigQuery, never()).update(any(Table.class)); + + // Every matching field has equal modes, so no incompatible-drift warning must be emitted. + ArgumentCaptor captor = ArgumentCaptor.forClass(LogRecord.class); + verify(mockLogHandler, atLeast(0)).publish(captor.capture()); + assertFalse( + "No drift warning should be logged when the schema already matches", + captor.getAllValues().stream() + .anyMatch(record -> record.getMessage().contains("Incompatible schema drift"))); + } + + @Test + public void maybeUpgradeSchema_treatsNullModeAsNullable_noDriftWarning() throws Exception { + Table mockTable = mock(Table.class); + when(mockTable.getTableId()).thenReturn(TableId.of("project", "dataset", "table")); + when(mockTable.getLabels()).thenReturn(ImmutableMap.of()); + + // BigQuery reports getMode() == null for NULLABLE columns. Represent the NULLABLE 'event_type' + // column with an unset (null) mode so its comparison against the NULLABLE desired field + // exercises normalizeMode's null -> NULLABLE path. Build the field WITHOUT setMode: the OSS + // BigQuery client stores toBuilder().setMode(null) as an empty mode string and throws on + // getMode(), whereas an unset mode is genuinely null. Every other field matches exactly. + com.google.cloud.bigquery.Field nullModeEventType = + com.google.cloud.bigquery.Field.newBuilder("event_type", StandardSQLTypeName.STRING) + .build(); + ImmutableList initialFields = + BigQuerySchema.getEventsSchema().getFields().stream() + .map(f -> f.getName().equals("event_type") ? nullModeEventType : f) + .collect(toImmutableList()); + StandardTableDefinition tableDefinition = + StandardTableDefinition.newBuilder() + .setSchema(com.google.cloud.bigquery.Schema.of(initialFields)) + .build(); + when(mockTable.getDefinition()).thenReturn(tableDefinition); + + Logger logger = Logger.getLogger(BigQueryUtils.class.getName()); + Handler mockLogHandler = mock(Handler.class); + logger.addHandler(mockLogHandler); + boolean upgraded; + try { + upgraded = BigQueryUtils.maybeUpgradeSchema(mockBigQuery, mockTable); + } finally { + logger.removeHandler(mockLogHandler); + } + + // A null (unset) mode is semantically NULLABLE, so the column already matches: the table is + // ready, no update is attempted, and no incompatible-drift warning must be emitted for it. + assertTrue(upgraded); + verify(mockBigQuery, never()).update(any(Table.class)); + + ArgumentCaptor captor = ArgumentCaptor.forClass(LogRecord.class); + verify(mockLogHandler, atLeast(0)).publish(captor.capture()); + assertFalse( + "A null mode must be normalized to NULLABLE, so no drift warning should be logged for" + + " 'event_type'", + captor.getAllValues().stream() + .anyMatch( + record -> + record + .getMessage() + .contains("Incompatible schema drift on column 'event_type'"))); + } + + @Test + public void maybeUpgradeSchema_warnsOnTypeDrift() throws Exception { + Table mockTable = mock(Table.class); + when(mockTable.getTableId()).thenReturn(TableId.of("project", "dataset", "table")); + when(mockTable.getLabels()).thenReturn(ImmutableMap.of()); + + // Existing 'timestamp' column has type STRING instead of the expected TIMESTAMP (same mode), so + // only the type-drift branch should fire. + ImmutableList initialFields = + BigQuerySchema.getEventsSchema().getFields().stream() + .map( + f -> + f.getName().equals("timestamp") + ? f.toBuilder().setType(StandardSQLTypeName.STRING).build() + : f) + .collect(toImmutableList()); + StandardTableDefinition tableDefinition = + StandardTableDefinition.newBuilder() + .setSchema(com.google.cloud.bigquery.Schema.of(initialFields)) + .build(); + when(mockTable.getDefinition()).thenReturn(tableDefinition); + + Logger logger = Logger.getLogger(BigQueryUtils.class.getName()); + Handler mockLogHandler = mock(Handler.class); + logger.addHandler(mockLogHandler); + try { + var unused = BigQueryUtils.maybeUpgradeSchema(mockBigQuery, mockTable); + } finally { + logger.removeHandler(mockLogHandler); + } + + ArgumentCaptor captor = ArgumentCaptor.forClass(LogRecord.class); + verify(mockLogHandler, atLeastOnce()).publish(captor.capture()); + assertTrue( + "Should have warned about type drift on the timestamp column", + captor.getAllValues().stream() + .anyMatch( + record -> + Objects.equals(record.getLevel(), Level.WARNING) + && record + .getMessage() + .contains("Incompatible schema drift on column 'timestamp'"))); + } + + @Test + public void isSafeIdentifier_nullIsRejectedWithoutThrowing() throws Exception { + // A null identifier must be rejected by the explicit guard. Without it, Pattern.matcher(null) + // would throw an NPE instead of returning false, so the DDL-safety check must short-circuit. + assertFalse(BigQueryUtils.isSafeIdentifier(null)); + // Sanity: well-formed identifiers pass and unsafe characters are rejected. + assertTrue(BigQueryUtils.isSafeIdentifier("project_123-abc")); + assertFalse(BigQueryUtils.isSafeIdentifier("bad;drop table")); + } + @Test public void createAnalyticsViews_executesQueries() throws Exception { BigQueryUtils.createAnalyticsViews(mockBigQuery, config); diff --git a/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfigTest.java b/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfigTest.java new file mode 100644 index 000000000..2a4844627 --- /dev/null +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfigTest.java @@ -0,0 +1,58 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.plugins.agentanalytics; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class BigQueryLoggerConfigTest { + + private static BigQueryLoggerConfig.Builder validBuilder() { + return BigQueryLoggerConfig.builder().projectId("test-project"); + } + + @Test + public void build_validConfig_succeeds() { + BigQueryLoggerConfig config = validBuilder().build(); + assertEquals("test-project", config.projectId()); + assertEquals(1, config.batchSize()); + assertEquals(10000, config.queueMaxSize()); + } + + @Test + public void build_nonPositiveBatchSize_throws() { + BigQueryLoggerConfig.Builder builder = validBuilder().batchSize(0); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void build_nonPositiveQueueMaxSize_throws() { + BigQueryLoggerConfig.Builder builder = validBuilder().queueMaxSize(0); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void build_nonPositiveMaxContentLength_throws() { + BigQueryLoggerConfig.Builder builder = validBuilder().maxContentLength(-1); + assertThrows(IllegalArgumentException.class, builder::build); + } +} 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 3ade94093..597312b79 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 @@ -39,8 +39,11 @@ import com.google.genai.types.FunctionCall; import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.Part; +import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.CompletableFuture; import org.junit.Test; import org.junit.runner.RunWith; @@ -272,4 +275,143 @@ public void smartTruncate_withCycle_detectsCycle() { assertTrue(result.isTruncated()); assertEquals("[cycle detected]", result.node().get("child").asText()); } + + @Test + public void smartTruncate_redactsSensitiveTopLevelKeys() { + ImmutableMap map = + ImmutableMap.of("api_key", "sk-secret", "password", "hunter2", "keep", "value"); + JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(map, 5000); + + JsonNode node = result.node(); + assertEquals("[REDACTED]", node.get("api_key").asText()); + assertEquals("[REDACTED]", node.get("password").asText()); + assertEquals("value", node.get("keep").asText()); + // Redaction must not flip the truncation flag. + assertFalse(result.isTruncated()); + } + + @Test + public void smartTruncate_redactsCaseInsensitiveAndTempPrefixKeys() { + ImmutableMap map = + ImmutableMap.of("Access_Token", "abc", "temp:scratch", "xyz", "keep", "ok"); + JsonNode node = JsonFormatter.smartTruncate(map, 5000).node(); + + assertEquals("[REDACTED]", node.get("Access_Token").asText()); + assertEquals("[REDACTED]", node.get("temp:scratch").asText()); + assertEquals("ok", node.get("keep").asText()); + } + + @Test + public void smartTruncate_redactsNestedSensitiveKeys() { + ImmutableMap map = + ImmutableMap.of("outer", ImmutableMap.of("client_secret", "s", "ok", "v")); + JsonNode node = JsonFormatter.smartTruncate(map, 5000).node(); + + assertEquals("[REDACTED]", node.get("outer").get("client_secret").asText()); + assertEquals("v", node.get("outer").get("ok").asText()); + } + + @Test + public void smartTruncate_depthGuard_replacesDeepSubtreeWithSentinel() { + Map root = new HashMap<>(); + Map cur = root; + for (int i = 0; i < 300; i++) { + Map next = new HashMap<>(); + cur.put("child", next); + cur = next; + } + + JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(root, 5000); + assertTrue(result.isTruncated()); + + JsonNode node = result.node(); + boolean foundSentinel = false; + for (int i = 0; i < 400; i++) { + JsonNode child = node.get("child"); + if (child == null) { + break; + } + if (child.isTextual() && child.asText().equals(JsonFormatter.MAX_DEPTH_MESSAGE)) { + foundSentinel = true; + break; + } + node = child; + } + assertTrue("Expected the max-depth sentinel in the deep chain", foundSentinel); + } + + @Test + public void smartTruncate_depthGuard_appliesToNestedArrays() { + // Deeply nested arrays must hit the same depth guard as objects; the array recursion must pass + // an increasing depth (not a reset/negated one) for the guard to ever fire. + List root = new ArrayList<>(); + List cur = root; + for (int i = 0; i < 300; i++) { + List next = new ArrayList<>(); + cur.add(next); + cur = next; + } + + JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(root, 5000); + + assertTrue("Deeply nested arrays should be truncated by the depth guard", result.isTruncated()); + } + + @Test + public void smartTruncate_atDepthBoundary_mapNotTruncated() { + // Exactly MAX_TRUNCATE_DEPTH levels must NOT be truncated. This pins the initial recursion + // depth + // to 0 (a seed of 1 would truncate at this boundary). + Map root = new HashMap<>(); + Map cur = root; + for (int i = 0; i < JsonFormatter.MAX_TRUNCATE_DEPTH; i++) { + Map next = new HashMap<>(); + cur.put("child", next); + cur = next; + } + + JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(root, 5000); + + assertFalse( + "A structure exactly at the depth boundary must not be truncated", result.isTruncated()); + } + + @Test + public void smartTruncate_atDepthBoundary_jsonNodeNotTruncated() { + // Same boundary check for the JsonNode input path (separate seed site in smartTruncate). + ObjectMapper mapper = new ObjectMapper(); + ObjectNode root = mapper.createObjectNode(); + ObjectNode cur = root; + for (int i = 0; i < JsonFormatter.MAX_TRUNCATE_DEPTH; i++) { + ObjectNode next = mapper.createObjectNode(); + cur.set("child", next); + cur = next; + } + + JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(root, 5000); + + assertFalse( + "A JsonNode structure exactly at the depth boundary must not be truncated", + result.isTruncated()); + } + + @Test + public void smartTruncate_atDepthBoundary_arrayNotTruncated() { + // Array analog of the map boundary check: exactly MAX_TRUNCATE_DEPTH levels of nested arrays + // must NOT be truncated. This pins the array-branch recursion to a +1 depth step; a +2 step + // would push the innermost element past the guard and truncate at this boundary. + List root = new ArrayList<>(); + List cur = root; + for (int i = 0; i < JsonFormatter.MAX_TRUNCATE_DEPTH; i++) { + List next = new ArrayList<>(); + cur.add(next); + cur = next; + } + + JsonFormatter.TruncationResult result = JsonFormatter.smartTruncate(root, 5000); + + assertFalse( + "A nested-array structure exactly at the depth boundary must not be truncated", + result.isTruncated()); + } } 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 e55b5a634..761a60e5e 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 @@ -17,6 +17,7 @@ package com.google.adk.plugins.agentanalytics; import static java.util.concurrent.TimeUnit.SECONDS; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -29,6 +30,7 @@ import com.google.cloud.bigquery.storage.v1.AppendRowsResponse; import com.google.cloud.bigquery.storage.v1.BigQueryWriteClient; import com.google.cloud.bigquery.storage.v1.StreamWriter; +import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.Uninterruptibles; import java.io.IOException; import java.lang.reflect.Field; @@ -120,6 +122,32 @@ public void addPendingTask_removedTaskOnCompletion() { // No specific log to check now, but we verify it completes without error. } + @Test + 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)); + + // 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. + TestPluginState stateWithClosedProcessor = + new TestPluginState(config) { + @Override + protected BatchProcessor removeProcessor(String id) { + return id.equals(invocationId) ? closedProcessor : super.removeProcessor(id); + } + }; + + stateWithClosedProcessor.ensureInvocationCompleted(invocationId).blockingAwait(); + + ImmutableMap stats = stateWithClosedProcessor.getDropStats(); + assertEquals(5L, (long) stats.get("queue_full")); + assertEquals(3L, (long) stats.get("append_error")); + assertEquals(2L, (long) stats.get("serialization_error")); + } + @Test public void ensureInvocationCompleted_noTasks_succeeds() { String invocationId = "testInvocation"; 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 b7ce7823f..e56f68fe4 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 @@ -29,6 +29,7 @@ import com.google.adk.events.Event; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.Tracer; +import io.opentelemetry.context.Context; import io.opentelemetry.context.Scope; import io.opentelemetry.sdk.testing.junit4.OpenTelemetryRule; import io.reactivex.rxjava3.core.Flowable; @@ -172,10 +173,10 @@ public void getTraceId_returnsCurrentTraceId() { @Test public void getTraceId_returnsInvocationId_whenRecordsIsEmpty() { - String traceId = traceManager.getTraceId(mockContext); - if (traceManager.hasAmbientSpan()) { - assertTrue(traceId.matches("[0-9a-f]{32}")); - } else { + // Pin to the root context so an ambient span leaked by another test sharing this JVM cannot + // make Span.current() valid and divert getTraceId away from the invocation-id fallback. + try (Scope ignored = Context.root().makeCurrent()) { + String traceId = traceManager.getTraceId(mockContext); assertEquals("test-invocation-id", traceId); } } @@ -205,13 +206,13 @@ public void attachCurrentSpan_worksWithoutAmbientSpan() { @Test public void getTraceId_fallsBackToInvocationId_whenRecordSpanIsInvalid() { - // attachCurrentSpan when no ambient context exists creates an invalid span record - traceManager.attachCurrentSpan(); + // 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. + try (Scope ignored = Context.root().makeCurrent()) { + traceManager.attachCurrentSpan(); - String traceId = traceManager.getTraceId(mockContext); - if (traceManager.hasAmbientSpan()) { - assertTrue(traceId.matches("[0-9a-f]{32}")); - } else { + String traceId = traceManager.getTraceId(mockContext); assertEquals("test-invocation-id", traceId); } } @@ -227,4 +228,35 @@ public void clearStack_doesNothing_whenRecordsIsEmpty() { traceManager.clearStack(); assertTrue(traceManager.getCurrentSpanAndParent().spanId().isEmpty()); } + + @Test + public void initTraceIfNeeded_setsRootAgentNameFromContext() { + assertEquals(TraceManager.DEFAULT_ROOT_AGENT_NAME, traceManager.getRootAgentName()); + traceManager.initTraceIfNeeded(mockContext); + assertEquals("test-agent", traceManager.getRootAgentName()); + } + + @Test + public void initTraceIfNeeded_nullAgent_keepsSentinel() { + InvocationContext ctx = mock(InvocationContext.class); + when(ctx.agent()).thenReturn(null); + traceManager.initTraceIfNeeded(ctx); + assertEquals(TraceManager.DEFAULT_ROOT_AGENT_NAME, traceManager.getRootAgentName()); + } + + @Test + public void initTrace_nullRootAgent_keepsSentinelWithoutThrowing() { + // rootAgent() may be null (e.g. workflow-driven callbacks with no resolved root); initTrace + // must + // guard on it directly rather than NPE. Called directly (not via initTraceIfNeeded, whose + // try/catch would otherwise mask a missing rootAgent null-check). + BaseAgent agentWithNullRoot = mock(BaseAgent.class); + when(agentWithNullRoot.rootAgent()).thenReturn(null); + InvocationContext ctx = mock(InvocationContext.class); + when(ctx.agent()).thenReturn(agentWithNullRoot); + + traceManager.initTrace(ctx); + + assertEquals(TraceManager.DEFAULT_ROOT_AGENT_NAME, traceManager.getRootAgentName()); + } }