fix: align agent graph tracking (tokens, tracker, edges, logger) with SDK spec

Author: mattrmc1

lib/sdk/server-ai/src/main/java/com/launchdarkly/sdk/server/ai/AIGraphTracker.java

@@ -6,7 +6,6 @@
 import com.launchdarkly.sdk.LDValue;
 import com.launchdarkly.sdk.ObjectBuilder;
 import com.launchdarkly.sdk.server.ai.datamodel.LDAITrackingTypes.TokenUsage;
-import com.launchdarkly.sdk.server.ai.internal.Loggers;
 import com.launchdarkly.sdk.server.ai.internal.ResumptionTokens;
 import com.launchdarkly.sdk.server.interfaces.LDClientInterface;
 
@@ -83,32 +82,15 @@ public final class AIGraphTracker {
     this.resumptionToken = ResumptionTokens.encodeGraph(runId, graphKey, variationKey, version);
   }
 
-  /**
-   * Reconstructs a graph tracker from a resumption token, preserving the original run identity.
-   *
-   * @param token the resumption token produced by {@link #getResumptionToken()}
-   * @param client the LaunchDarkly client; must not be {@code null}
-   * @param context the evaluation context; must not be {@code null}
-   * @return a new tracker with the decoded run identity
-   * @throws IllegalArgumentException if the token is malformed
-   */
-  public static AIGraphTracker fromResumptionToken(
-      String token, LDClientInterface client, LDContext context) {
-    return fromResumptionToken(token, client, context, Loggers.defaultLogger());
-  }
-
   /**
    * Reconstructs a graph tracker from a resumption token, preserving the original run identity,
    * and logging through the supplied logger.
-   *
-   * @param token the resumption token produced by {@link #getResumptionToken()}
-   * @param client the LaunchDarkly client; must not be {@code null}
-   * @param context the evaluation context; must not be {@code null}
-   * @param logger the logger to use for at-most-once warnings; must not be {@code null}
-   * @return a new tracker with the decoded run identity
-   * @throws IllegalArgumentException if the token is malformed
+   * <p>
+   * This method is package-private. External callers should use
+   * {@link LDAIClient#createGraphTracker(String, LDContext)} instead, which correctly pipes the
+   * configured logger through from the top.
    */
-  public static AIGraphTracker fromResumptionToken(
+  static AIGraphTracker fromResumptionToken(
       String token, LDClientInterface client, LDContext context, LDLogger logger) {
     ResumptionTokens.DecodedGraph d = ResumptionTokens.decodeGraph(token);
     int version = d.getVersion();
@@ -174,8 +156,7 @@ public void trackDuration(double durationMs) {
   /**
    * Records the total token usage for the graph invocation.
    * <p>
-   * At-most-once: subsequent calls are silently dropped. Calls where all counts are zero do not
-   * consume the at-most-once slot.
+   * At-most-once: subsequent calls are silently dropped.
    *
    * @param tokens the token usage; ignored if {@code null}
    */
@@ -184,17 +165,11 @@ public void trackTotalTokens(TokenUsage tokens) {
       logger.debug("Skipping trackTotalTokens: tokens was null.");
       return;
     }
-    boolean hasPositive = tokens.getTotal() > 0 || tokens.getInput() > 0 || tokens.getOutput() > 0;
-    if (!hasPositive) {
-      return;
-    }
     if (!tokensRecorded.compareAndSet(null, tokens)) {
       logger.warn("Skipping trackTotalTokens: token usage already recorded on this graph tracker.");
       return;
     }
-    if (tokens.getTotal() > 0) {
-      client.trackMetric(GRAPH_TOTAL_TOKENS, context, baseData().build(), tokens.getTotal());
-    }
+    client.trackMetric(GRAPH_TOTAL_TOKENS, context, baseData().build(), tokens.getTotal());
   }
 
   /**

lib/sdk/server-ai/src/main/java/com/launchdarkly/sdk/server/ai/AgentGraphDefinition.java

@@ -19,7 +19,9 @@
  * <p>
  * An {@code AgentGraphDefinition} is obtained from {@link LDAIClient#agentGraph}. When
  * {@link #isEnabled()} returns {@code false}, the graph definition was not fetchable or failed
- * validation; in that case all node collections are empty and traversal methods are no-ops.
+ * validation; in that case all node collections are empty and traversal methods are no-ops. Only
+ * {@link #getConfig()} and {@link #createTracker()} remain meaningful, so callers can still inspect
+ * the raw flag value and fire graph-level usage events for a disabled graph.
  * <p>
  * Traversal methods ({@link #traverse} and {@link #reverseTraverse}) are BFS-based and
  * cycle-safe: each node is visited at most once.
@@ -142,13 +144,14 @@ AgentGraphFlagValue getConfig() {
   /**
    * Creates a new {@link AIGraphTracker} for this graph invocation.
    * <p>
-   * Each call produces a fresh tracker with a new run ID. Returns {@code null} if the graph is
-   * disabled.
+   * Each call produces a fresh tracker with a new run ID. A tracker is returned even when the
+   * graph is disabled, so callers can still fire graph-level usage events (e.g. invocation
+   * failure) when the graph's configuration could not be resolved.
    *
-   * @return a new tracker, or {@code null} if disabled
+   * @return a new tracker, or {@code null} if no tracker factory was provided
    */
   public AIGraphTracker createTracker() {
-    if (!enabled || trackerFactory == null) {
+    if (trackerFactory == null) {
       return null;
     }
     return trackerFactory.get();

lib/sdk/server-ai/src/main/java/com/launchdarkly/sdk/server/ai/AgentGraphNode.java

@@ -1,5 +1,6 @@
 package com.launchdarkly.sdk.server.ai;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
@@ -18,7 +19,9 @@ public final class AgentGraphNode {
   AgentGraphNode(String key, AIAgentConfig config, List<GraphEdge> edges) {
     this.key = key;
     this.config = config;
-    this.edges = edges == null ? Collections.<GraphEdge>emptyList() : edges;
+    this.edges = edges == null
+        ? Collections.<GraphEdge>emptyList()
+        : Collections.unmodifiableList(new ArrayList<>(edges));
   }
 
   /**

lib/sdk/server-ai/src/main/java/com/launchdarkly/sdk/server/ai/LDAIClientImpl.java

@@ -368,6 +368,10 @@ private Supplier<LDAIConfigTracker> trackerFactory(
   @Override
   public AgentGraphDefinition agentGraph(
       String graphKey, LDContext context, Map<String, Object> variables) {
+    Objects.requireNonNull(graphKey, "graphKey");
+    if (graphKey.trim().isEmpty()) {
+      throw new IllegalArgumentException("graphKey must not be blank");
+    }
     client.trackMetric(TRACK_USAGE_AGENT_GRAPH, context, LDValue.of(graphKey), 1);
 
     LDValue flagValue = client.jsonValueVariation(graphKey, context, LDValue.ofNull());

lib/sdk/server-ai/src/test/java/com/launchdarkly/sdk/server/ai/AIGraphTrackerTest.java

@@ -176,22 +176,29 @@ public void trackTotalTokensIsAtMostOnce() {
   }
 
   @Test
-  public void trackTotalTokensAllZeroDoesNotBurnSlot() {
-    tracker.trackTotalTokens(new TokenUsage(0, 0, 0));
+  public void trackTotalTokensNullIsIgnored() {
+    tracker.trackTotalTokens(null);
     verify(client, never()).trackMetric(
         eq("$ld:ai:graph:total_tokens"), any(), any(), anyDouble());
-    // Slot not consumed — a subsequent non-zero call should fire
-    tracker.trackTotalTokens(new TokenUsage(5, 5, 0));
-    verify(client, times(1)).trackMetric(
-        eq("$ld:ai:graph:total_tokens"), any(), any(), anyDouble());
+    assertThat(debugs().stream().anyMatch(w -> w.contains("tokens was null")), is(true));
   }
 
   @Test
-  public void trackTotalTokensNullIsIgnored() {
-    tracker.trackTotalTokens(null);
-    verify(client, never()).trackMetric(
+  public void trackTotalTokensEmitsRawTotalIgnoringInputOutput() {
+    // total=0 with positive input/output: emits 0, not input+output
+    tracker.trackTotalTokens(new TokenUsage(0, 5, 3));
+    verify(client).trackMetric(
+        eq("$ld:ai:graph:total_tokens"), eq(CONTEXT), eq(baseExpectedData()), eq(0.0));
+  }
+
+  @Test
+  public void trackTotalTokensSlotConsumedEvenForZeroTotal() {
+    tracker.trackTotalTokens(new TokenUsage(0, 0, 0));
+    // slot already consumed; second call should be dropped with a warning
+    tracker.trackTotalTokens(new TokenUsage(10, 5, 5));
+    verify(client, times(1)).trackMetric(
         eq("$ld:ai:graph:total_tokens"), any(), any(), anyDouble());
-    assertThat(debugs().stream().anyMatch(w -> w.contains("tokens was null")), is(true));
+    assertThat(warnings().stream().anyMatch(w -> w.contains("token usage already recorded")), is(true));
   }
 
   // ---- trackPath ------------------------------------------------------------
@@ -366,14 +373,14 @@ public void getResumptionTokenIsNotNull() {
   @Test
   public void fromResumptionTokenRoundTrips() {
     String token = tracker.getResumptionToken();
-    AIGraphTracker reconstructed = AIGraphTracker.fromResumptionToken(token, client, CONTEXT);
+    AIGraphTracker reconstructed = AIGraphTracker.fromResumptionToken(token, client, CONTEXT, logger);
     assertThat(reconstructed.getResumptionToken(), is(token));
   }
 
   @Test
   public void fromResumptionTokenPreservesRunId() {
     String token = tracker.getResumptionToken();
-    AIGraphTracker reconstructed = AIGraphTracker.fromResumptionToken(token, client, CONTEXT);
+    AIGraphTracker reconstructed = AIGraphTracker.fromResumptionToken(token, client, CONTEXT, logger);
     // Verify same events are emitted by the reconstructed tracker
     reconstructed.trackInvocationSuccess();
     ArgumentCaptor<LDValue> captor = ArgumentCaptor.forClass(LDValue.class);
@@ -404,7 +411,7 @@ public void fromResumptionTokenWithLoggerRoutesWarningsThroughIt() {
   public void fromResumptionTokenPreservesVersionZero() {
     String token = com.launchdarkly.sdk.server.ai.internal.ResumptionTokens.encodeGraph(
         RUN_ID, GRAPH_KEY, null, 0);
-    AIGraphTracker reconstructed = AIGraphTracker.fromResumptionToken(token, client, CONTEXT);
+    AIGraphTracker reconstructed = AIGraphTracker.fromResumptionToken(token, client, CONTEXT, logger);
     reconstructed.trackInvocationSuccess();
     ArgumentCaptor<LDValue> captor = ArgumentCaptor.forClass(LDValue.class);
     verify(client).trackMetric(eq("$ld:ai:graph:invocation_success"), any(), captor.capture(), anyDouble());

lib/sdk/server-ai/src/test/java/com/launchdarkly/sdk/server/ai/AgentGraphDefinitionTest.java

@@ -63,7 +63,8 @@ private static AgentGraphFlagValue flagValue(String root, String[][] edges) {
 
   private static AIAgentConfig makeConfig(String key, boolean enabled) {
     return new AIAgentConfig(key, enabled, null, null, null, null, null,
-        () -> mock(com.launchdarkly.sdk.server.ai.LDAIConfigTracker.class));
+        () -> mock(com.launchdarkly.sdk.server.ai.LDAIConfigTracker.class),
+        Evaluator.noop());
   }
 
   private static Map<String, AIAgentConfig> configs(String... keys) {
@@ -264,11 +265,26 @@ public void isEnabledReflectsConstructorValue() {
 
   @Test
   public void createTrackerReturnsNullWhenDisabled() {
+    // A null factory is the only case that returns null (defensive guard).
     AgentGraphDefinition graph = new AgentGraphDefinition(
         AgentGraphFlagValue.disabled(), Collections.emptyMap(), false, null);
     assertThat(graph.createTracker(), is(nullValue()));
   }
 
+  @Test
+  public void createTrackerReturnsTrackerEvenWhenDisabled() {
+    // Disabled graphs still produce a tracker so callers can fire graph-level usage events
+    // (e.g. invocation failure) when the graph's configuration could not be resolved.
+    LDClientInterface client = mock(LDClientInterface.class);
+    AgentGraphDefinition graph = new AgentGraphDefinition(
+        AgentGraphFlagValue.disabled(), Collections.emptyMap(), false,
+        () -> new AIGraphTracker(client, "run-id", "graph-key", null, 1,
+            com.launchdarkly.sdk.LDContext.create("user"),
+            com.launchdarkly.logging.LDLogger.withAdapter(
+                com.launchdarkly.logging.Logs.none(), "")));
+    assertThat(graph.createTracker(), is(notNullValue()));
+  }
+
   @Test
   public void createTrackerReturnsTrackerWhenEnabled() {
     LDClientInterface client = mock(LDClientInterface.class);

lib/sdk/server-ai/src/test/java/com/launchdarkly/sdk/server/ai/LDAIClientImplTest.java

@@ -333,6 +333,16 @@ private static LDValue agentFlagValue(boolean enabled) {
         + "\"instructions\":\"test instructions\"}");
   }
 
+  @Test(expected = NullPointerException.class)
+  public void agentGraphThrowsOnNullGraphKey() {
+    ai.agentGraph(null, context, null);
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void agentGraphThrowsOnBlankGraphKey() {
+    ai.agentGraph("  ", context, null);
+  }
+
   @Test
   public void agentGraphFiresUsageEvent() {
     when(client.jsonValueVariation(eq("g"), any(), any()))