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 6982c7df0..9510991f8 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 @@ -115,17 +115,19 @@ public class BigQueryAgentAnalyticsPlugin extends BasePlugin { "hitl_input"); private final BigQueryLoggerConfig config; - private final BigQuery bigQuery; + private final @Nullable BigQuery bigQuery; 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. + // Set only on enabled public construction paths (see registerShutdownHook); null for disabled + // configs (enabled(false)) and 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)); + // A disabled plugin resolves no credentials and builds no BigQuery client. + this(config, config.enabled() ? createBigQuery(config) : null, new PluginState(config)); + registerShutdownHook(); } public BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config, BigQuery bigQuery) @@ -136,7 +138,8 @@ public BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config, BigQuery bigQue registerShutdownHook(); } - BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config, BigQuery bigQuery, PluginState state) { + BigQueryAgentAnalyticsPlugin( + BigQueryLoggerConfig config, @Nullable BigQuery bigQuery, PluginState state) { super("bigquery_agent_analytics"); this.config = config; this.bigQuery = bigQuery; @@ -144,6 +147,9 @@ public BigQueryAgentAnalyticsPlugin(BigQueryLoggerConfig config, BigQuery bigQue } private void registerShutdownHook() { + if (!config.enabled()) { + return; + } shutdownHook = new Thread( () -> { @@ -196,6 +202,9 @@ private static BigQuery createBigQuery(BigQueryLoggerConfig config) throws IOExc } private void ensureTableExistsOnce() { + if (bigQuery == null) { + return; + } if (!tableEnsured) { synchronized (tableEnsuredLock) { if (!tableEnsured) { @@ -275,13 +284,19 @@ private boolean ensureTableExists(BigQuery bigQuery, BigQueryLoggerConfig config } private void processBigQueryException(BigQueryException e, String logMessage) { - if (e.getMessage().contains("invalid_grant")) { + String message = e.getMessage(); + if (message != null && message.contains("invalid_grant")) { logger.log(Level.SEVERE, "Failed to authenticate with BigQuery.", e); } else { logger.log(Level.WARNING, logMessage, e); } } + private static String formatErrorMessage(Throwable error) { + String message = error.getMessage(); + return (message != null && !message.isEmpty()) ? message : error.getClass().getName(); + } + private Completable logEvent( String eventType, InvocationContext invocationContext, @@ -565,6 +580,16 @@ PluginState getState() { return state; } + @VisibleForTesting + @Nullable BigQuery getBigQuery() { + return bigQuery; + } + + @VisibleForTesting + @Nullable Thread getShutdownHook() { + return shutdownHook; + } + private Optional getCompletedEventData( InvocationContext invocationContext, String expectedKindPrefix) { TraceManager traceManager = state.getTraceManager(invocationContext.invocationId()); @@ -595,7 +620,7 @@ private Optional getCompletedEventData( @Override public Maybe onUserMessageCallback( InvocationContext invocationContext, Content userMessage) { - if (state.isProcessed(invocationContext.invocationId())) { + if (!config.enabled() || state.isProcessed(invocationContext.invocationId())) { return Maybe.empty(); } state.getTraceManager(invocationContext.invocationId()).ensureInvocationSpan(invocationContext); @@ -663,7 +688,7 @@ public Maybe onUserMessageCallback( @Override public Maybe onEventCallback(InvocationContext invocationContext, Event event) { - if (state.isProcessed(invocationContext.invocationId())) { + if (!config.enabled() || state.isProcessed(invocationContext.invocationId())) { return Maybe.empty(); } // Only emit STATE_DELTA when there is an actual state change, matching the Python plugin @@ -851,7 +876,7 @@ public Maybe onEventCallback(InvocationContext invocationContext, Event e @Override public Maybe beforeRunCallback(InvocationContext invocationContext) { - if (state.isProcessed(invocationContext.invocationId())) { + if (!config.enabled() || state.isProcessed(invocationContext.invocationId())) { return Maybe.empty(); } state.getTraceManager(invocationContext.invocationId()).ensureInvocationSpan(invocationContext); @@ -861,6 +886,9 @@ public Maybe beforeRunCallback(InvocationContext invocationContext) { @Override public Completable afterRunCallback(InvocationContext invocationContext) { + if (!config.enabled()) { + return Completable.complete(); + } return logEvent( "INVOCATION_COMPLETED", invocationContext, @@ -871,7 +899,8 @@ public Completable afterRunCallback(InvocationContext invocationContext) { @Override public Maybe beforeAgentCallback(BaseAgent agent, CallbackContext callbackContext) { - if (state.isProcessed(callbackContext.invocationContext().invocationId())) { + if (!config.enabled() + || state.isProcessed(callbackContext.invocationContext().invocationId())) { return Maybe.empty(); } state @@ -883,6 +912,10 @@ public Maybe beforeAgentCallback(BaseAgent agent, CallbackContext callb @Override public Maybe afterAgentCallback(BaseAgent agent, CallbackContext callbackContext) { + if (!config.enabled() + || state.isProcessed(callbackContext.invocationContext().invocationId())) { + return Maybe.empty(); + } return logEvent( "AGENT_COMPLETED", callbackContext.invocationContext(), @@ -902,7 +935,8 @@ public Maybe afterAgentCallback(BaseAgent agent, CallbackContext callba @Override public Maybe beforeModelCallback( CallbackContext callbackContext, LlmRequest.Builder llmRequest) { - if (state.isProcessed(callbackContext.invocationContext().invocationId())) { + if (!config.enabled() + || state.isProcessed(callbackContext.invocationContext().invocationId())) { return Maybe.empty(); } Map attributes = new HashMap<>(); @@ -971,7 +1005,8 @@ public Maybe beforeModelCallback( @Override public Maybe afterModelCallback( CallbackContext callbackContext, LlmResponse llmResponse) { - if (state.isProcessed(callbackContext.invocationContext().invocationId())) { + if (!config.enabled() + || state.isProcessed(callbackContext.invocationContext().invocationId())) { return Maybe.empty(); } TraceManager traceManager = @@ -1063,7 +1098,8 @@ public Maybe afterModelCallback( @Override public Maybe onModelErrorCallback( CallbackContext callbackContext, LlmRequest.Builder llmRequest, Throwable error) { - if (state.isProcessed(callbackContext.invocationContext().invocationId())) { + if (!config.enabled() + || state.isProcessed(callbackContext.invocationContext().invocationId())) { return Maybe.empty(); } TraceManager traceManager = @@ -1076,7 +1112,7 @@ public Maybe onModelErrorCallback( String parentSpanId = spanIds.spanId().orElse(null); EventData.Builder eventDataBuilder = - EventData.builder().setStatus("ERROR").setErrorMessage(error.getMessage()); + EventData.builder().setStatus("ERROR").setErrorMessage(formatErrorMessage(error)); if (popped.isPresent()) { eventDataBuilder.setLatency(popped.get().duration()); } @@ -1094,7 +1130,7 @@ public Maybe onModelErrorCallback( @Override public Maybe> beforeToolCallback( BaseTool tool, Map toolArgs, ToolContext toolContext) { - if (state.isProcessed(toolContext.invocationContext().invocationId())) { + if (!config.enabled() || state.isProcessed(toolContext.invocationContext().invocationId())) { return Maybe.empty(); } ImmutableMap contentMap = @@ -1125,7 +1161,7 @@ public Maybe> afterToolCallback( Map toolArgs, ToolContext toolContext, Map result) { - if (state.isProcessed(toolContext.invocationContext().invocationId())) { + if (!config.enabled() || state.isProcessed(toolContext.invocationContext().invocationId())) { return Maybe.empty(); } state @@ -1167,7 +1203,7 @@ public Maybe> afterToolCallback( @Override public Maybe> onToolErrorCallback( BaseTool tool, Map toolArgs, ToolContext toolContext, Throwable error) { - if (state.isProcessed(toolContext.invocationContext().invocationId())) { + if (!config.enabled() || state.isProcessed(toolContext.invocationContext().invocationId())) { return Maybe.empty(); } state @@ -1187,7 +1223,7 @@ public Maybe> onToolErrorCallback( .buildOrThrow(); EventData.Builder eventDataBuilder = - EventData.builder().setStatus("ERROR").setErrorMessage(error.getMessage()); + EventData.builder().setStatus("ERROR").setErrorMessage(formatErrorMessage(error)); if (popped.isPresent()) { eventDataBuilder.setLatency(popped.get().duration()); } 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 301ead886..8c9219406 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 @@ -236,13 +236,34 @@ public BigQueryLoggerConfig build() { throw new IllegalArgumentException( "maxContentLength must be positive, got " + config.maxContentLength()); } + if (toMillisOrThrow(config.batchFlushInterval(), "batchFlushInterval") < 1) { + throw new IllegalArgumentException( + "batchFlushInterval must be at least 1 ms, got " + config.batchFlushInterval()); + } + if (toMillisOrThrow(config.shutdownTimeout(), "shutdownTimeout") < 1) { + throw new IllegalArgumentException( + "shutdownTimeout must be at least 1 ms, got " + config.shutdownTimeout()); + } return config; } } + private static long toMillisOrThrow(Duration duration, String fieldName) { + try { + return duration.toMillis(); + } catch (ArithmeticException e) { + throw new IllegalArgumentException( + fieldName + " overflows millisecond range, got " + duration, e); + } + } + /** Retry configuration for BigQuery writes. */ @AutoValue public abstract static class RetryConfig { + /** + * Maximum number of retry attempts after the initial write attempt ({@code 0} disables + * retries). + */ public abstract int maxRetries(); public abstract Duration initialDelay(); @@ -274,7 +295,39 @@ public abstract static class Builder { @CanIgnoreReturnValue public abstract Builder maxDelay(Duration maxDelay); - public abstract RetryConfig build(); + abstract RetryConfig autoBuild(); + + public RetryConfig build() { + RetryConfig retry = autoBuild(); + if (retry.maxRetries() < 0) { + throw new IllegalArgumentException( + "maxRetries must be non-negative, got " + retry.maxRetries()); + } + // Write retries back off in whole milliseconds, so (0, 1 ms) would silently become 0. + validateRetryDelay(retry.initialDelay(), "initialDelay"); + validateRetryDelay(retry.maxDelay(), "maxDelay"); + if (retry.maxDelay().compareTo(retry.initialDelay()) < 0) { + throw new IllegalArgumentException( + "maxDelay (" + + retry.maxDelay() + + ") must be >= initialDelay (" + + retry.initialDelay() + + ")"); + } + if (!Double.isFinite(retry.multiplier()) || retry.multiplier() < 1.0) { + throw new IllegalArgumentException( + "multiplier must be finite and >= 1.0, got " + retry.multiplier()); + } + return retry; + } + + private static void validateRetryDelay(Duration delay, String fieldName) { + long millis = toMillisOrThrow(delay, fieldName); + if (delay.isNegative() || (!delay.isZero() && millis < 1)) { + throw new IllegalArgumentException( + fieldName + " must be 0 or at least 1 ms, got " + delay); + } + } } } } 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 d60c8d911..1f712047d 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 @@ -87,7 +87,7 @@ class PluginState { private final ScheduledExecutorService executor; private final ExecutorService offloadExecutor; private final ThreadPoolExecutor writerCloseExecutor; - private final BigQueryWriteClient writeClient; + private final @Nullable BigQueryWriteClient writeClient; private static final AtomicLong threadCounter = new AtomicLong(0); // Map of invocation ID to BatchProcessor. private final ConcurrentHashMap batchProcessors = @@ -96,7 +96,7 @@ class PluginState { private final ConcurrentHashMap traceManagers = new ConcurrentHashMap<>(); // Cache of invocation ID to Boolean indicating invocation ID has been processed. private final Cache processedInvocations; - private final GcsOffloader offloader; + private final @Nullable GcsOffloader offloader; private final Parser parser; private final ConcurrentHashMap>> pendingTasks = new ConcurrentHashMap<>(); @@ -229,13 +229,13 @@ synchronized boolean attachWriter(StreamWriter writer) { new ThreadPoolExecutor.AbortPolicy()); this.writerCloseExecutor.allowCoreThreadTimeOut(true); // One write client per plugin instance, shared by all invocations. - this.writeClient = createWriteClient(config); + this.writeClient = config.enabled() ? createWriteClient(config) : null; this.processedInvocations = CacheBuilder.newBuilder() .maximumSize(10000) .expireAfterWrite(java.time.Duration.ofMinutes(10)) .build(); - this.offloader = getGcsOffloader(config); + this.offloader = config.enabled() ? getGcsOffloader(config) : null; this.parser = new Parser( offloader, @@ -285,6 +285,7 @@ protected BigQueryWriteClient createWriteClient(BigQueryLoggerConfig config) thr protected StreamWriter createWriter() { BigQueryLoggerConfig.RetryConfig retryConfig = config.retryConfig(); + // StreamWriter treats maxAttempts as a retry count; 0 disables retries. RetrySettings retrySettings = RetrySettings.newBuilder() .setMaxAttempts(retryConfig.maxRetries()) 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 fd7f3355f..93cbf4bb6 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 @@ -28,6 +28,7 @@ import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyMap; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.atLeast; import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.mock; @@ -54,6 +55,7 @@ import com.google.api.core.ApiFutures; import com.google.auth.Credentials; import com.google.cloud.bigquery.BigQuery; +import com.google.cloud.bigquery.BigQueryException; import com.google.cloud.bigquery.BigQueryOptions; import com.google.cloud.bigquery.Field.Mode; import com.google.cloud.bigquery.FieldList; @@ -2285,4 +2287,220 @@ public void logEvent_sessionState_redactedBeforeTruncationFallback() throws Exce assertEquals("[UNSERIALIZABLE]", stateNode.get("bad").asText()); assertEquals("visible", stateNode.get("ok").asText()); } + + @Test + public void publicConstructor_disabledConfig_callbacksAndCloseAreNoOps() throws Exception { + // No credentials are set: an enabled plugin would resolve ADC and build clients here. + BigQueryLoggerConfig disabledConfig = + BigQueryLoggerConfig.builder() + .enabled(false) + .projectId("project") + .datasetId("dataset") + .build(); + BigQueryAgentAnalyticsPlugin disabledPlugin = new BigQueryAgentAnalyticsPlugin(disabledConfig); + assertNull(disabledPlugin.getBigQuery()); + assertNull(disabledPlugin.getShutdownHook()); + PluginState disabledState = disabledPlugin.getState(); + CallbackContext callbackContext = mock(CallbackContext.class); + when(callbackContext.invocationContext()).thenReturn(mockInvocationContext); + ToolContext toolContext = mock(ToolContext.class); + when(toolContext.invocationContext()).thenReturn(mockInvocationContext); + BaseTool tool = mock(BaseTool.class); + when(tool.name()).thenReturn("test_tool"); + Content message = Content.fromParts(Part.fromText("hello")); + Event event = Event.builder().author("agent_name").content(message).build(); + RuntimeException error = new RuntimeException("boom"); + + disabledPlugin.onUserMessageCallback(mockInvocationContext, message).test().assertResult(); + disabledPlugin.onEventCallback(mockInvocationContext, event).test().assertResult(); + disabledPlugin.beforeRunCallback(mockInvocationContext).test().assertResult(); + disabledPlugin.beforeAgentCallback(fakeAgent, callbackContext).test().assertResult(); + disabledPlugin.beforeModelCallback(callbackContext, LlmRequest.builder()).test().assertResult(); + disabledPlugin + .afterModelCallback(callbackContext, LlmResponse.builder().build()) + .test() + .assertResult(); + disabledPlugin + .onModelErrorCallback(callbackContext, LlmRequest.builder(), error) + .test() + .assertResult(); + disabledPlugin.beforeToolCallback(tool, ImmutableMap.of(), toolContext).test().assertResult(); + disabledPlugin + .afterToolCallback(tool, ImmutableMap.of(), toolContext, ImmutableMap.of()) + .test() + .assertResult(); + disabledPlugin + .onToolErrorCallback(tool, ImmutableMap.of(), toolContext, error) + .test() + .assertResult(); + disabledPlugin.afterAgentCallback(fakeAgent, callbackContext).test().assertResult(); + // afterRun finalization would drop trace managers, so check for spans before it runs. + assertTrue(disabledState.getTraceManagers().isEmpty()); + disabledPlugin.afterRunCallback(mockInvocationContext).test().assertResult(); + + assertFalse(disabledState.isProcessed("invocation_id")); + assertTrue(disabledState.getBatchProcessors().isEmpty()); + assertTrue(disabledPlugin.close().blockingAwait(5_000, MILLISECONDS)); + } + + @Test + public void publicConstructor_enabledConfig_registersShutdownHookUntilClose() throws Exception { + // A mock's null metrics type would NPE in gax's ClientContext header builder. + when(config.credentials().getMetricsCredentialType()).thenCallRealMethod(); + BigQueryAgentAnalyticsPlugin enabledPlugin = new BigQueryAgentAnalyticsPlugin(config); + Thread shutdownHook = enabledPlugin.getShutdownHook(); + assertTrue(enabledPlugin.close().blockingAwait(5_000, MILLISECONDS)); + + assertNotNull(enabledPlugin.getBigQuery()); + assertNotNull(shutdownHook); + // close() deregistered the hook, so removing it again finds nothing. + assertFalse(Runtime.getRuntime().removeShutdownHook(shutdownHook)); + } + + @Test + public void onModelErrorCallback_nullErrorMessage_recordsExceptionClassName() throws Exception { + CallbackContext mockCallbackContext = mock(CallbackContext.class); + when(mockCallbackContext.invocationContext()).thenReturn(mockInvocationContext); + state.getTraceManager("invocation_id").pushSpan(mockInvocationContext, "llm_request"); + + // An explicitly constructed NPE has a null message (helpful NPE text is JVM-raised only). + plugin + .onModelErrorCallback( + mockCallbackContext, mock(LlmRequest.Builder.class), new NullPointerException()) + .blockingSubscribe(); + + Map row = state.getBatchProcessor("invocation_id").queue.poll(); + assertNotNull("Row not found in queue", row); + assertEquals("LLM_ERROR", row.get("event_type")); + assertEquals("java.lang.NullPointerException", row.get("error_message")); + } + + @Test + public void onToolErrorCallback_nullErrorMessage_recordsExceptionClassName() throws Exception { + ToolContext mockToolContext = mock(ToolContext.class); + when(mockToolContext.invocationContext()).thenReturn(mockInvocationContext); + BaseTool mockTool = mock(BaseTool.class); + when(mockTool.name()).thenReturn("failing_tool"); + + plugin + .onToolErrorCallback( + mockTool, ImmutableMap.of(), mockToolContext, new NullPointerException()) + .blockingSubscribe(); + + Map errorRow = drainRowsByEventType().get("TOOL_ERROR"); + assertNotNull("TOOL_ERROR row not found", errorRow); + assertEquals("java.lang.NullPointerException", errorRow.get("error_message")); + } + + @Test + public void onToolErrorCallback_emptyErrorMessage_recordsExceptionClassName() throws Exception { + ToolContext mockToolContext = mock(ToolContext.class); + when(mockToolContext.invocationContext()).thenReturn(mockInvocationContext); + BaseTool mockTool = mock(BaseTool.class); + when(mockTool.name()).thenReturn("failing_tool"); + + plugin + .onToolErrorCallback( + mockTool, ImmutableMap.of(), mockToolContext, new IllegalStateException("")) + .blockingSubscribe(); + + Map errorRow = drainRowsByEventType().get("TOOL_ERROR"); + assertNotNull("TOOL_ERROR row not found", errorRow); + assertEquals("java.lang.IllegalStateException", errorRow.get("error_message")); + } + + @Test + public void onModelErrorCallback_emptyErrorMessage_recordsExceptionClassName() throws Exception { + CallbackContext mockCallbackContext = mock(CallbackContext.class); + when(mockCallbackContext.invocationContext()).thenReturn(mockInvocationContext); + state.getTraceManager("invocation_id").pushSpan(mockInvocationContext, "llm_request"); + + plugin + .onModelErrorCallback( + mockCallbackContext, mock(LlmRequest.Builder.class), new IllegalStateException("")) + .blockingSubscribe(); + + Map row = state.getBatchProcessor("invocation_id").queue.poll(); + assertNotNull("Row not found in queue", row); + assertEquals("LLM_ERROR", row.get("event_type")); + assertEquals("java.lang.IllegalStateException", row.get("error_message")); + } + + @Test + public void afterAgentCallback_processedInvocation_doesNotRecreateTraceManager() { + CallbackContext callbackContext = mock(CallbackContext.class); + when(callbackContext.invocationContext()).thenReturn(mockInvocationContext); + // Finalization marks the invocation processed and removes its trace manager. + state.markProcessed("invocation_id"); + + plugin.afterAgentCallback(fakeAgent, callbackContext).test().assertResult(); + + assertTrue(state.getTraceManagers().isEmpty()); + } + + @Test + public void afterAgentCallback_alreadyProcessed_skipsSpanPopAndLogging() throws Exception { + CallbackContext callbackContext = mock(CallbackContext.class); + when(callbackContext.invocationContext()).thenReturn(mockInvocationContext); + + plugin + .onUserMessageCallback(mockInvocationContext, Content.builder().build()) + .blockingSubscribe(); + plugin.beforeAgentCallback(fakeAgent, callbackContext).blockingSubscribe(); + state.getBatchProcessor("invocation_id").queue.clear(); + String agentSpanId = + state + .getTraceManager("invocation_id") + .getCurrentSpanAndParent(mockInvocationContext) + .spanId() + .orElseThrow(); + state.markProcessed("invocation_id"); + + plugin.afterAgentCallback(fakeAgent, callbackContext).test().assertResult(); + + assertTrue(state.getBatchProcessor("invocation_id").queue.isEmpty()); + assertEquals( + agentSpanId, + state + .getTraceManager("invocation_id") + .getCurrentSpanAndParent(mockInvocationContext) + .spanId() + .orElseThrow()); + } + + @Test + public void onUserMessageCallback_tableCheckErrorWithoutMessage_stillCompletes() + throws Exception { + when(mockBigQuery.getTable(any(TableId.class))).thenThrow(new BigQueryException(500, null)); + + plugin + .onUserMessageCallback(mockInvocationContext, Content.builder().build()) + .test() + .assertResult(); + + verify(mockBigQuery).getTable(any(TableId.class)); + } + + @Test + public void onUserMessageCallback_withoutBigQueryClient_skipsTableSetupAndQueuesRow() + throws Exception { + BigQueryAgentAnalyticsPlugin pluginWithoutClient = + new BigQueryAgentAnalyticsPlugin(config, /* bigQuery= */ null, state); + Logger logger = Logger.getLogger(BigQueryAgentAnalyticsPlugin.class.getName()); + Handler logHandler = mock(Handler.class); + logger.addHandler(logHandler); + try { + pluginWithoutClient + .onUserMessageCallback(mockInvocationContext, Content.fromParts(Part.fromText("hi"))) + .test() + .assertResult(); + } finally { + logger.removeHandler(logHandler); + } + + assertNotNull(drainRowsByEventType().get("USER_MESSAGE_RECEIVED")); + // A table check against the missing client would log a WARNING on every event. + verify(logHandler, never()) + .publish(argThat(record -> record.getLevel().intValue() >= Level.WARNING.intValue())); + } } 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 index 58e1c5565..612554204 100644 --- a/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfigTest.java +++ b/core/src/test/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfigTest.java @@ -20,6 +20,7 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertThrows; +import java.time.Duration; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -68,4 +69,128 @@ public void build_nonPositiveMaxContentLength_throws() { BigQueryLoggerConfig.Builder builder = validBuilder().maxContentLength(-1); assertThrows(IllegalArgumentException.class, builder::build); } + + @Test + public void build_zeroBatchFlushInterval_throws() { + BigQueryLoggerConfig.Builder builder = validBuilder().batchFlushInterval(Duration.ZERO); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void build_subMillisecondBatchFlushInterval_throws() { + BigQueryLoggerConfig.Builder builder = + validBuilder().batchFlushInterval(Duration.ofNanos(500_000)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void build_negativeShutdownTimeout_throws() { + BigQueryLoggerConfig.Builder builder = validBuilder().shutdownTimeout(Duration.ofSeconds(-1)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void build_overflowingBatchFlushInterval_throwsIllegalArgumentException() { + BigQueryLoggerConfig.Builder builder = + validBuilder().batchFlushInterval(Duration.ofSeconds(Long.MAX_VALUE)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void build_overflowingShutdownTimeout_throwsIllegalArgumentException() { + BigQueryLoggerConfig.Builder builder = + validBuilder().shutdownTimeout(Duration.ofSeconds(Long.MAX_VALUE)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_negativeMaxRetries_throws() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder().maxRetries(-1); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_negativeInitialDelay_throws() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder().initialDelay(Duration.ofMillis(-1)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_subMillisecondInitialDelay_throws() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder().initialDelay(Duration.ofNanos(500_000)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_negativeMaxDelay_throws() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder() + .initialDelay(Duration.ZERO) + .maxDelay(Duration.ofMillis(-1)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_subMillisecondMaxDelay_throws() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder() + .initialDelay(Duration.ZERO) + .maxDelay(Duration.ofNanos(500_000)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_overflowingDelay_throwsIllegalArgumentException() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder() + .initialDelay(Duration.ofSeconds(Long.MAX_VALUE)) + .maxDelay(Duration.ofSeconds(Long.MAX_VALUE)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_maxDelayBelowInitialDelay_throws() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder() + .initialDelay(Duration.ofSeconds(2)) + .maxDelay(Duration.ofSeconds(1)); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_multiplierBelowOne_throws() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder().multiplier(0.5); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_nanMultiplier_throws() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder().multiplier(Double.NaN); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_infiniteMultiplier_throws() { + BigQueryLoggerConfig.RetryConfig.Builder builder = + BigQueryLoggerConfig.RetryConfig.builder().multiplier(Double.POSITIVE_INFINITY); + assertThrows(IllegalArgumentException.class, builder::build); + } + + @Test + public void retryConfigBuild_zeroDelaysAndUnitMultiplier_succeeds() { + BigQueryLoggerConfig.RetryConfig retry = + BigQueryLoggerConfig.RetryConfig.builder() + .initialDelay(Duration.ZERO) + .maxDelay(Duration.ZERO) + .multiplier(1.0) + .build(); + assertThat(retry.initialDelay()).isEqualTo(Duration.ZERO); + assertThat(retry.maxDelay()).isEqualTo(Duration.ZERO); + assertThat(retry.multiplier()).isEqualTo(1.0); + } } 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 3b3474392..59f6087ed 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 @@ -944,4 +944,24 @@ public void getDropStats_foldsAfterCloseFromLiveProcessors() { // already-removed ones. assertEquals(1L, (long) stats.get("after_close")); } + + @Test + public void constructor_disabledConfig_createsNoWriteClientOrOffloader() throws IOException { + BigQueryLoggerConfig disabledConfig = + config.toBuilder().enabled(false).gcsBucketName("test-bucket").build(); + PluginState disabledState = + new PluginState(disabledConfig) { + @Override + protected BigQueryWriteClient createWriteClient(BigQueryLoggerConfig config) { + throw new AssertionError("A disabled plugin must not create a write client"); + } + + @Override + protected GcsOffloader getGcsOffloader(BigQueryLoggerConfig config) { + throw new AssertionError("A disabled plugin must not create a GCS offloader"); + } + }; + + assertTrue(disabledState.close().blockingAwait(5, SECONDS)); + } }