feat: handle usageUpdate events and populate UsageMetadata in AgentResponse

.flattened-pom.xml

@@ -4,7 +4,7 @@
   <modelVersion>4.0.0</modelVersion>
   <groupId>io.github.glaforge.antigravity</groupId>
   <artifactId>antigravity-sdk-parent</artifactId>
-  <version>0.2.14-SNAPSHOT</version>
+  <version>0.2.16-SNAPSHOT</version>
   <packaging>pom</packaging>
   <name>Antigravity Java SDK</name>
   <description>A Java SDK for building agents using the Antigravity localharness.</description>
@@ -40,7 +40,7 @@
     <protobuf.version>3.25.3</protobuf.version>
     <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
     <jackson.version>2.17.1</jackson.version>
-    <revision>0.2.14-SNAPSHOT</revision>
+    <revision>0.2.16-SNAPSHOT</revision>
   </properties>
   <dependencyManagement>
     <dependencies>

antigravity-sdk-harness/.flattened-pom.xml

@@ -5,10 +5,10 @@
   <parent>
     <groupId>io.github.glaforge.antigravity</groupId>
     <artifactId>antigravity-sdk-parent</artifactId>
-    <version>0.2.14-SNAPSHOT</version>
+    <version>0.2.16-SNAPSHOT</version>
   </parent>
   <artifactId>antigravity-sdk-harness</artifactId>
-  <version>0.2.14-SNAPSHOT</version>
+  <version>0.2.16-SNAPSHOT</version>
   <name>Antigravity Java SDK - Harness Binaries</name>
   <description>Platform-specific native localharness Go binaries for offline deployment.</description>
   <url>https://github.com/glaforge/antigravity-java-sdk</url>

antigravity-sdk-protocol/.flattened-pom.xml

@@ -5,10 +5,10 @@
   <parent>
     <groupId>io.github.glaforge.antigravity</groupId>
     <artifactId>antigravity-sdk-parent</artifactId>
-    <version>0.2.14-SNAPSHOT</version>
+    <version>0.2.16-SNAPSHOT</version>
   </parent>
   <artifactId>antigravity-sdk-protocol</artifactId>
-  <version>0.2.14-SNAPSHOT</version>
+  <version>0.2.16-SNAPSHOT</version>
   <name>Antigravity Java SDK - Protocol</name>
   <description>The Protocol Buffers definitions and generated classes for Antigravity Java SDK.</description>
   <url>https://github.com/glaforge/antigravity-java-sdk</url>

antigravity-sdk-wrapper/.flattened-pom.xml

@@ -5,10 +5,10 @@
   <parent>
     <groupId>io.github.glaforge.antigravity</groupId>
     <artifactId>antigravity-sdk-parent</artifactId>
-    <version>0.2.14-SNAPSHOT</version>
+    <version>0.2.16-SNAPSHOT</version>
   </parent>
   <artifactId>antigravity-sdk-wrapper</artifactId>
-  <version>0.2.14-SNAPSHOT</version>
+  <version>0.2.16-SNAPSHOT</version>
   <name>Antigravity Java SDK - Wrapper</name>
   <description>The main Java wrapper and client for the Antigravity localharness.</description>
   <url>https://github.com/glaforge/antigravity-java-sdk</url>

antigravity-sdk-wrapper/src/main/java/io/github/glaforge/antigravity/Agent.java

@@ -64,6 +64,8 @@
 import java.util.function.Consumer;
 import java.util.List;
 import java.util.Set;
+import java.util.Map;
+import java.util.LinkedHashMap;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.nio.file.Path;
@@ -92,6 +94,9 @@ public class Agent implements AutoCloseable, TriggerContext {
 	private StringBuilder currentText;
 	private StringBuilder currentThoughts;
 	private UsageMetadata currentUsage;
+	private UsageMetadata cumulativeUsage = new UsageMetadata(0, 0, 0, 0, 0);
+	private UsageMetadata turnStartUsage;
+	private final ConcurrentMap<String, UsageMetadata> trajectoryUsages = new ConcurrentHashMap<>();
 	private final List<Policy> policies;
 	private boolean hasStructuredOutput;
 	private StringBuilder wsBuffer = new StringBuilder();
@@ -116,7 +121,35 @@ public SandboxStatus getSandboxStatus() {
 	 * @return the usage metadata
 	 */
 	public UsageMetadata getUsageMetadata() {
-		return currentUsage;
+		if (currentUsage != null) {
+			return currentUsage;
+		}
+		if (cumulativeUsage != null && turnStartUsage != null) {
+			UsageMetadata diff = cumulativeUsage.subtract(turnStartUsage);
+			if (diff.totalTokenCount() > 0) {
+				return diff;
+			}
+		}
+		return cumulativeUsage;
+	}
+
+	/**
+	 * Returns the cumulative token usage across all turns in this session.
+	 *
+	 * @return the total cumulative usage metadata
+	 */
+	public UsageMetadata getTotalUsage() {
+		return cumulativeUsage;
+	}
+
+	/**
+	 * Returns a map of trajectory ID to cumulative token usage for subagents and
+	 * main agent.
+	 *
+	 * @return an unmodifiable map of trajectory usage metadata
+	 */
+	public Map<String, UsageMetadata> getTrajectoryUsages() {
+		return Collections.unmodifiableMap(new LinkedHashMap<>(trajectoryUsages));
 	}
 
 	/**
@@ -982,6 +1015,7 @@ public CompletableFuture<AgentResponse> chatStream(List<AgentInput> inputs, Cons
 		this.currentText = new StringBuilder();
 		this.currentThoughts = new StringBuilder();
 		this.hasStructuredOutput = false;
+		this.turnStartUsage = this.cumulativeUsage;
 		this.currentUsage = null;
 
 		StringBuilder combinedText = new StringBuilder();
@@ -1165,6 +1199,76 @@ void handleIncomingMessage(WebSocket webSocket, String message) {
 					this.sandboxStatus = new SandboxStatus(available, reason);
 					warnIfSandboxUnavailable(this.sandboxStatus);
 				}
+
+				if (initResp.has("cumulativeUsage") || initResp.has("cumulative_usage")) {
+					JsonNode cumNode = initResp.has("cumulativeUsage")
+							? initResp.get("cumulativeUsage")
+							: initResp.get("cumulative_usage");
+					UsageMetadata parsed = parseUsageMetadata(cumNode);
+					if (parsed != null) {
+						this.cumulativeUsage = parsed;
+						this.turnStartUsage = parsed;
+					}
+				}
+
+				JsonNode trajUsageNode = initResp.has("trajectoryUsage")
+						? initResp.get("trajectoryUsage")
+						: initResp.get("trajectory_usage");
+				if (trajUsageNode != null && trajUsageNode.isArray()) {
+					for (JsonNode entry : trajUsageNode) {
+						String trajId = entry.has("trajectoryId")
+								? entry.get("trajectoryId").asText()
+								: entry.path("trajectory_id").asText();
+						JsonNode uNode = entry.has("usage") ? entry.get("usage") : null;
+						if (trajId != null && !trajId.isEmpty() && uNode != null) {
+							UsageMetadata parsed = parseUsageMetadata(uNode);
+							if (parsed != null) {
+								trajectoryUsages.put(trajId, parsed);
+							}
+						}
+					}
+				}
+			}
+
+			if (payload.has("usageUpdate") || payload.has("usage_update")) {
+				JsonNode usageUpdate = payload.has("usageUpdate")
+						? payload.get("usageUpdate")
+						: payload.get("usage_update");
+
+				if (usageUpdate.has("total")) {
+					UsageMetadata newTotal = parseUsageMetadata(usageUpdate.get("total"));
+					if (newTotal != null) {
+						this.cumulativeUsage = newTotal;
+						this.currentUsage = (turnStartUsage != null) ? newTotal.subtract(turnStartUsage) : newTotal;
+					}
+				}
+
+				JsonNode agentsNode = usageUpdate.has("agents") ? usageUpdate.get("agents") : null;
+				if (agentsNode != null && agentsNode.isArray()) {
+					for (JsonNode entry : agentsNode) {
+						String trajId = entry.has("trajectoryId")
+								? entry.get("trajectoryId").asText()
+								: entry.path("trajectory_id").asText();
+						JsonNode uNode = entry.has("usage") ? entry.get("usage") : null;
+						if (trajId != null && !trajId.isEmpty() && uNode != null) {
+							UsageMetadata parsed = parseUsageMetadata(uNode);
+							if (parsed != null) {
+								trajectoryUsages.put(trajId, parsed);
+							}
+						}
+					}
+				}
+			}
+
+			if (payload.has("usageMetadata") || payload.has("usage_metadata")) {
+				JsonNode topUsage = payload.has("usageMetadata")
+						? payload.get("usageMetadata")
+						: payload.get("usage_metadata");
+				UsageMetadata parsed = parseUsageMetadata(topUsage);
+				if (parsed != null) {
+					this.cumulativeUsage = parsed;
+					this.currentUsage = (turnStartUsage != null) ? parsed.subtract(turnStartUsage) : parsed;
+				}
 			}
 
 			if (payload.has("stepUpdate")) {
@@ -1223,20 +1327,14 @@ void handleIncomingMessage(WebSocket webSocket, String message) {
 					}
 				}
 
-				if (stepUpdate.has("usageMetadata")) {
-					JsonNode usage = stepUpdate.get("usageMetadata");
-					List<ModalityTokenCount> promptDetails = parseModalityDetails(usage.path("promptTokensDetails"));
-					List<ModalityTokenCount> cacheDetails = parseModalityDetails(usage.path("cacheTokensDetails"));
-					List<ModalityTokenCount> candidateDetails = parseModalityDetails(
-							usage.path("candidatesTokensDetails"));
-					List<ModalityTokenCount> toolUseDetails = parseModalityDetails(
-							usage.path("toolUsePromptTokensDetails"));
-					String serviceTier = usage.has("serviceTier") ? usage.get("serviceTier").asText() : null;
-
-					currentUsage = new UsageMetadata(usage.path("promptTokenCount").asInt(),
-							usage.path("cachedContentTokenCount").asInt(), usage.path("candidatesTokenCount").asInt(),
-							usage.path("thoughtsTokenCount").asInt(), usage.path("totalTokenCount").asInt(),
-							serviceTier, promptDetails, cacheDetails, candidateDetails, toolUseDetails);
+				if (stepUpdate.has("usageMetadata") || stepUpdate.has("usage_metadata")) {
+					JsonNode usage = stepUpdate.has("usageMetadata")
+							? stepUpdate.get("usageMetadata")
+							: stepUpdate.get("usage_metadata");
+					UsageMetadata parsed = parseUsageMetadata(usage);
+					if (parsed != null) {
+						this.currentUsage = parsed;
+					}
 				}
 
 				if (stepUpdate.has("state") && "STATE_ERROR".equals(stepUpdate.path("state").asText())) {
@@ -1532,9 +1630,18 @@ else if (req.has("customToolCall")) {
 								currentToolCallsPublisher.closeExceptionally(new AgentCancelledException());
 							}
 						} else {
+							UsageMetadata finalUsage = currentUsage;
+							if (finalUsage == null && cumulativeUsage != null) {
+								UsageMetadata diff = (turnStartUsage != null)
+										? cumulativeUsage.subtract(turnStartUsage)
+										: cumulativeUsage;
+								if (diff.totalTokenCount() > 0) {
+									finalUsage = diff;
+								}
+							}
 							currentChatFuture
 									.complete(new AgentResponse(currentText != null ? currentText.toString() : "",
-											currentThoughts != null ? currentThoughts.toString() : "", currentUsage));
+											currentThoughts != null ? currentThoughts.toString() : "", finalUsage));
 							if (currentThoughtsPublisher != null) {
 								currentThoughtsPublisher.close();
 							}
@@ -1741,6 +1848,62 @@ private static String resolveGeminiApiKey() {
 		return propKey != null ? propKey : (envKey != null ? envKey : localEnvKey);
 	}
 
+	/**
+	 * Parses a {@link UsageMetadata} object from a Jackson {@link JsonNode},
+	 * supporting both camelCase and snake_case field variants.
+	 *
+	 * @param usage
+	 *            the JSON node representing UsageMetadata
+	 * @return the parsed UsageMetadata record, or null if node is null/missing
+	 */
+	static UsageMetadata parseUsageMetadata(JsonNode usage) {
+		if (usage == null || usage.isMissingNode() || usage.isNull()) {
+			return null;
+		}
+		int promptTokenCount = getIntField(usage, "promptTokenCount", "prompt_token_count");
+		int cachedContentTokenCount = getIntField(usage, "cachedContentTokenCount", "cached_content_token_count");
+		int candidatesTokenCount = getIntField(usage, "candidatesTokenCount", "candidates_token_count");
+		int thoughtsTokenCount = getIntField(usage, "thoughtsTokenCount", "thoughts_token_count");
+		int totalTokenCount = getIntField(usage, "totalTokenCount", "total_token_count");
+
+		String serviceTier = null;
+		if (usage.has("serviceTier")) {
+			serviceTier = usage.get("serviceTier").asText();
+		} else if (usage.has("service_tier")) {
+			serviceTier = usage.get("service_tier").asText();
+		}
+
+		JsonNode promptDetailsNode = usage.has("promptTokensDetails")
+				? usage.get("promptTokensDetails")
+				: usage.get("prompt_tokens_details");
+		JsonNode cacheDetailsNode = usage.has("cacheTokensDetails")
+				? usage.get("cacheTokensDetails")
+				: usage.get("cache_tokens_details");
+		JsonNode candidateDetailsNode = usage.has("candidatesTokensDetails")
+				? usage.get("candidatesTokensDetails")
+				: usage.get("candidates_tokens_details");
+		JsonNode toolUseDetailsNode = usage.has("toolUsePromptTokensDetails")
+				? usage.get("toolUsePromptTokensDetails")
+				: usage.get("tool_use_prompt_tokens_details");
+
+		List<ModalityTokenCount> promptDetails = parseModalityDetails(promptDetailsNode);
+		List<ModalityTokenCount> cacheDetails = parseModalityDetails(cacheDetailsNode);
+		List<ModalityTokenCount> candidateDetails = parseModalityDetails(candidateDetailsNode);
+		List<ModalityTokenCount> toolUseDetails = parseModalityDetails(toolUseDetailsNode);
+
+		return new UsageMetadata(promptTokenCount, cachedContentTokenCount, candidatesTokenCount, thoughtsTokenCount,
+				totalTokenCount, serviceTier, promptDetails, cacheDetails, candidateDetails, toolUseDetails);
+	}
+
+	private static int getIntField(JsonNode node, String camelCaseName, String snakeCaseName) {
+		if (node.has(camelCaseName)) {
+			return node.get(camelCaseName).asInt(0);
+		} else if (node.has(snakeCaseName)) {
+			return node.get(snakeCaseName).asInt(0);
+		}
+		return 0;
+	}
+
 	/**
 	 * Parses a JSON array of modality token details into a list of
 	 * {@link ModalityTokenCount} records.
@@ -1749,14 +1912,14 @@ private static String resolveGeminiApiKey() {
 	 *            the JSON node containing the modality details array
 	 * @return an unmodifiable list of ModalityTokenCount records
 	 */
-	private static List<ModalityTokenCount> parseModalityDetails(JsonNode node) {
+	static List<ModalityTokenCount> parseModalityDetails(JsonNode node) {
 		if (node == null || !node.isArray()) {
 			return List.of();
 		}
 		List<ModalityTokenCount> list = new ArrayList<>();
 		for (JsonNode item : node) {
 			String modStr = item.path("modality").asText("");
-			long count = item.path("tokenCount").asLong();
+			long count = item.has("tokenCount") ? item.get("tokenCount").asLong(0) : item.path("token_count").asLong(0);
 			Modality mod = Modality.fromString(modStr);
 			list.add(new ModalityTokenCount(mod, count));
 		}

antigravity-sdk-wrapper/src/main/java/io/github/glaforge/antigravity/AgentResponse.java

@@ -37,6 +37,15 @@ public record AgentResponse(String text, String thoughts, UsageMetadata usageMet
 		thoughts = thoughts != null ? thoughts : "";
 	}
 
+	/**
+	 * Convenience alias for {@link #usageMetadata()}.
+	 *
+	 * @return the usage metadata
+	 */
+	public UsageMetadata usage() {
+		return usageMetadata;
+	}
+
 	/**
 	 * Parses the text response as JSON and maps it to the specified class.
 	 *

antigravity-sdk-wrapper/src/main/java/io/github/glaforge/antigravity/UsageMetadata.java

@@ -74,9 +74,12 @@ public UsageMetadata(int promptTokenCount, int cachedContentTokenCount, int cand
 	 * @return a new UsageMetadata with summed token counts and merged service tier
 	 */
 	public UsageMetadata add(UsageMetadata other) {
-		if (other == null) {
+		if (other == null || other.totalTokenCount() == 0) {
 			return this;
 		}
+		if (this.totalTokenCount == 0) {
+			return other;
+		}
 		String mergedTier;
 		if (this.serviceTier == null) {
 			mergedTier = other.serviceTier();
@@ -111,7 +114,7 @@ public UsageMetadata plus(UsageMetadata other) {
 	 * @return a new UsageMetadata with subtracted token counts
 	 */
 	public UsageMetadata subtract(UsageMetadata other) {
-		if (other == null) {
+		if (other == null || other.totalTokenCount() == 0) {
 			return this;
 		}
 		String mergedTier = this.serviceTier != null ? this.serviceTier : other.serviceTier();

antigravity-sdk-wrapper/src/test/java/io/github/glaforge/antigravity/ObservabilityTest.java

@@ -38,14 +38,17 @@ public void testUsageObservability() throws Exception {
 				await().atMost(90, TimeUnit.SECONDS).until(future::isDone);
 				AgentResponse response = future.get();
 
-				// Verify usage metadata is populated if available
+				// Verify usage metadata is populated
 				UsageMetadata usage = response.usageMetadata();
-				if (usage != null) {
-					System.out.println("Metadata: " + usage);
-					assertTrue(usage.promptTokenCount() >= 0, "Prompt tokens should be >= 0");
-				} else {
-					System.out.println("Metadata was not returned by the harness in this test run.");
-				}
+				assertNotNull(usage, "UsageMetadata should not be null in response");
+				assertEquals(usage, response.usage(), "response.usage() should match response.usageMetadata()");
+				assertTrue(usage.promptTokenCount() > 0, "Prompt tokens should be > 0");
+				assertTrue(usage.totalTokenCount() > 0, "Total tokens should be > 0");
+
+				// Verify Agent getters
+				assertNotNull(agent.getUsageMetadata(), "agent.getUsageMetadata() should not be null");
+				assertNotNull(agent.getTotalUsage(), "agent.getTotalUsage() should not be null");
+				assertTrue(agent.getTotalUsage().totalTokenCount() > 0, "Cumulative total tokens should be > 0");
 			}
 		});
 	}