Skip to content

fix: harden disabled mode, config checks and error capture #1727

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
+477 −24
Open

fix: harden disabled mode, config checks and error capture #1727

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 55 additions & 19 deletions core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryAgentAnalyticsPlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -136,14 +138,18 @@ 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;
this.state = state;
}

private void registerShutdownHook() {
if (!config.enabled()) {
return;
}
shutdownHook =
new Thread(
() -> {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -565,6 +580,16 @@ PluginState getState() {
return state;
}

@VisibleForTesting
@Nullable BigQuery getBigQuery() {
return bigQuery;
}

@VisibleForTesting
@Nullable Thread getShutdownHook() {
return shutdownHook;
}

private Optional<EventData> getCompletedEventData(
InvocationContext invocationContext, String expectedKindPrefix) {
TraceManager traceManager = state.getTraceManager(invocationContext.invocationId());
Expand Down Expand Up @@ -595,7 +620,7 @@ private Optional<EventData> getCompletedEventData(
@Override
public Maybe<Content> 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);
Expand Down Expand Up @@ -663,7 +688,7 @@ public Maybe<Content> onUserMessageCallback(

@Override
public Maybe<Event> 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
Expand Down Expand Up @@ -851,7 +876,7 @@ public Maybe<Event> onEventCallback(InvocationContext invocationContext, Event e

@Override
public Maybe<Content> beforeRunCallback(InvocationContext invocationContext) {
if (state.isProcessed(invocationContext.invocationId())) {
if (!config.enabled() || state.isProcessed(invocationContext.invocationId())) {
return Maybe.empty();
}
state.getTraceManager(invocationContext.invocationId()).ensureInvocationSpan(invocationContext);
Expand All @@ -861,6 +886,9 @@ public Maybe<Content> beforeRunCallback(InvocationContext invocationContext) {

@Override
public Completable afterRunCallback(InvocationContext invocationContext) {
if (!config.enabled()) {
return Completable.complete();
}
return logEvent(
"INVOCATION_COMPLETED",
invocationContext,
Expand All @@ -871,7 +899,8 @@ public Completable afterRunCallback(InvocationContext invocationContext) {

@Override
public Maybe<Content> beforeAgentCallback(BaseAgent agent, CallbackContext callbackContext) {
if (state.isProcessed(callbackContext.invocationContext().invocationId())) {
if (!config.enabled()
|| state.isProcessed(callbackContext.invocationContext().invocationId())) {
return Maybe.empty();
}
state
Expand All @@ -883,6 +912,10 @@ public Maybe<Content> beforeAgentCallback(BaseAgent agent, CallbackContext callb

@Override
public Maybe<Content> afterAgentCallback(BaseAgent agent, CallbackContext callbackContext) {
if (!config.enabled()
|| state.isProcessed(callbackContext.invocationContext().invocationId())) {
return Maybe.empty();
}
return logEvent(
"AGENT_COMPLETED",
callbackContext.invocationContext(),
Expand All @@ -902,7 +935,8 @@ public Maybe<Content> afterAgentCallback(BaseAgent agent, CallbackContext callba
@Override
public Maybe<LlmResponse> beforeModelCallback(
CallbackContext callbackContext, LlmRequest.Builder llmRequest) {
if (state.isProcessed(callbackContext.invocationContext().invocationId())) {
if (!config.enabled()
|| state.isProcessed(callbackContext.invocationContext().invocationId())) {
return Maybe.empty();
}
Map<String, Object> attributes = new HashMap<>();
Expand Down Expand Up @@ -971,7 +1005,8 @@ public Maybe<LlmResponse> beforeModelCallback(
@Override
public Maybe<LlmResponse> afterModelCallback(
CallbackContext callbackContext, LlmResponse llmResponse) {
if (state.isProcessed(callbackContext.invocationContext().invocationId())) {
if (!config.enabled()
|| state.isProcessed(callbackContext.invocationContext().invocationId())) {
return Maybe.empty();
}
TraceManager traceManager =
Expand Down Expand Up @@ -1063,7 +1098,8 @@ public Maybe<LlmResponse> afterModelCallback(
@Override
public Maybe<LlmResponse> 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 =
Expand All @@ -1076,7 +1112,7 @@ public Maybe<LlmResponse> 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());
}
Expand All @@ -1094,7 +1130,7 @@ public Maybe<LlmResponse> onModelErrorCallback(
@Override
public Maybe<Map<String, Object>> beforeToolCallback(
BaseTool tool, Map<String, Object> toolArgs, ToolContext toolContext) {
if (state.isProcessed(toolContext.invocationContext().invocationId())) {
if (!config.enabled() || state.isProcessed(toolContext.invocationContext().invocationId())) {
return Maybe.empty();
}
ImmutableMap<String, Object> contentMap =
Expand Down Expand Up @@ -1125,7 +1161,7 @@ public Maybe<Map<String, Object>> afterToolCallback(
Map<String, Object> toolArgs,
ToolContext toolContext,
Map<String, Object> result) {
if (state.isProcessed(toolContext.invocationContext().invocationId())) {
if (!config.enabled() || state.isProcessed(toolContext.invocationContext().invocationId())) {
return Maybe.empty();
}
state
Expand Down Expand Up @@ -1167,7 +1203,7 @@ public Maybe<Map<String, Object>> afterToolCallback(
@Override
public Maybe<Map<String, Object>> onToolErrorCallback(
BaseTool tool, Map<String, Object> toolArgs, ToolContext toolContext, Throwable error) {
if (state.isProcessed(toolContext.invocationContext().invocationId())) {
if (!config.enabled() || state.isProcessed(toolContext.invocationContext().invocationId())) {
return Maybe.empty();
}
state
Expand All @@ -1187,7 +1223,7 @@ public Maybe<Map<String, Object>> 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());
}
Expand Down
55 changes: 54 additions & 1 deletion core/src/main/java/com/google/adk/plugins/agentanalytics/BigQueryLoggerConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down Expand Up @@ -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);
}
}
}
}
}
9 changes: 5 additions & 4 deletions core/src/main/java/com/google/adk/plugins/agentanalytics/PluginState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, BatchProcessor> batchProcessors =
Expand All @@ -96,7 +96,7 @@ class PluginState {
private final ConcurrentHashMap<String, TraceManager> traceManagers = new ConcurrentHashMap<>();
// Cache of invocation ID to Boolean indicating invocation ID has been processed.
private final Cache<String, Boolean> processedInvocations;
private final GcsOffloader offloader;
private final @Nullable GcsOffloader offloader;
private final Parser parser;
private final ConcurrentHashMap<String, Set<CompletableFuture<Void>>> pendingTasks =
new ConcurrentHashMap<>();
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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())
Expand Down
Loading
Loading