chore: Move graphKey to LDAIConfigTracker constructor (#1279)

Author: jsonbailey

packages/sdk/server-ai/__tests__/LDAIClientImpl.test.ts

@@ -139,7 +139,14 @@ describe('config evaluation', () => {
     const evaluateSpy = jest.spyOn(client as any, '_evaluate');
     const result = await client.agentConfig(key, testContext, defaultValue, variables);
 
-    expect(evaluateSpy).toHaveBeenCalledWith(key, testContext, defaultValue, 'agent', variables);
+    expect(evaluateSpy).toHaveBeenCalledWith(
+      key,
+      testContext,
+      defaultValue,
+      'agent',
+      variables,
+      undefined,
+    );
     expect(result.instructions).toBe(
       'You are a helpful assistant. Your name is John and your score is 42',
     );
@@ -464,7 +471,14 @@ describe('agentConfig method', () => {
       key,
       1,
     );
-    expect(evaluateSpy).toHaveBeenCalledWith(key, testContext, defaultValue, 'agent', variables);
+    expect(evaluateSpy).toHaveBeenCalledWith(
+      key,
+      testContext,
+      defaultValue,
+      'agent',
+      variables,
+      undefined,
+    );
     expect(result).toBe(mockConfig);
     evaluateSpy.mockRestore();
   });

packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts

@@ -933,7 +933,7 @@ describe('trackToolCall', () => {
     );
   });
 
-  it('includes graphKey when provided', () => {
+  it('includes graphKey when set on constructor', () => {
     const tracker = new LDAIConfigTrackerImpl(
       mockLdClient,
       testRunId,
@@ -943,9 +943,10 @@ describe('trackToolCall', () => {
       modelName,
       providerName,
       testContext,
+      'my-graph',
     );
 
-    tracker.trackToolCall('my-tool', 'my-graph');
+    tracker.trackToolCall('my-tool');
 
     expect(mockTrack).toHaveBeenCalledWith(
       '$ld:ai:tool_call',
@@ -993,8 +994,8 @@ describe('trackToolCalls', () => {
   });
 });
 
-describe('graphKey parameter support', () => {
-  it('includes graphKey in trackDuration event', () => {
+describe('graphKey constructor support', () => {
+  it('includes graphKey in trackDuration event when set on constructor', () => {
     const tracker = new LDAIConfigTrackerImpl(
       mockLdClient,
       testRunId,
@@ -1004,9 +1005,10 @@ describe('graphKey parameter support', () => {
       modelName,
       providerName,
       testContext,
+      'my-graph',
     );
 
-    tracker.trackDuration(1000, 'my-graph');
+    tracker.trackDuration(1000);
 
     expect(mockTrack).toHaveBeenCalledWith(
       '$ld:ai:duration:total',
@@ -1016,7 +1018,7 @@ describe('graphKey parameter support', () => {
     );
   });
 
-  it('includes graphKey in trackSuccess event', () => {
+  it('includes graphKey in trackSuccess event when set on constructor', () => {
     const tracker = new LDAIConfigTrackerImpl(
       mockLdClient,
       testRunId,
@@ -1026,9 +1028,10 @@ describe('graphKey parameter support', () => {
       modelName,
       providerName,
       testContext,
+      'my-graph',
     );
 
-    tracker.trackSuccess('my-graph');
+    tracker.trackSuccess();
 
     expect(mockTrack).toHaveBeenCalledWith(
       '$ld:ai:generation:success',
@@ -1038,7 +1041,7 @@ describe('graphKey parameter support', () => {
     );
   });
 
-  it('does not include graphKey when not provided', () => {
+  it('does not include graphKey when not set on constructor', () => {
     const tracker = new LDAIConfigTrackerImpl(
       mockLdClient,
       testRunId,
@@ -1059,6 +1062,41 @@ describe('graphKey parameter support', () => {
       1,
     );
   });
+
+  it('includes graphKey in getTrackData when set on constructor', () => {
+    const tracker = new LDAIConfigTrackerImpl(
+      mockLdClient,
+      testRunId,
+      configKey,
+      variationKey,
+      version,
+      modelName,
+      providerName,
+      testContext,
+      'my-graph',
+    );
+
+    expect(tracker.getTrackData()).toEqual({
+      ...getExpectedTrackData(),
+      graphKey: 'my-graph',
+    });
+  });
+
+  it('does not include graphKey in getTrackData when not set', () => {
+    const tracker = new LDAIConfigTrackerImpl(
+      mockLdClient,
+      testRunId,
+      configKey,
+      variationKey,
+      version,
+      modelName,
+      providerName,
+      testContext,
+    );
+
+    expect(tracker.getTrackData()).toEqual(getExpectedTrackData());
+    expect('graphKey' in tracker.getTrackData()).toBe(false);
+  });
 });
 
 describe('at-most-once semantics', () => {
@@ -1311,4 +1349,96 @@ describe('fromResumptionToken', () => {
       1,
     );
   });
+
+  it('includes graphKey in resumption token when set on constructor', () => {
+    const tracker = new LDAIConfigTrackerImpl(
+      mockLdClient,
+      testRunId,
+      configKey,
+      variationKey,
+      version,
+      modelName,
+      providerName,
+      testContext,
+      'my-graph',
+    );
+
+    const token = tracker.resumptionToken;
+    const decoded = JSON.parse(Buffer.from(token, 'base64url').toString('utf8'));
+
+    expect(decoded).toEqual({
+      runId: testRunId,
+      configKey,
+      variationKey,
+      version,
+      graphKey: 'my-graph',
+    });
+  });
+
+  it('does not include graphKey in resumption token when not set', () => {
+    const tracker = new LDAIConfigTrackerImpl(
+      mockLdClient,
+      testRunId,
+      configKey,
+      variationKey,
+      version,
+      modelName,
+      providerName,
+      testContext,
+    );
+
+    const token = tracker.resumptionToken;
+    const decoded = JSON.parse(Buffer.from(token, 'base64url').toString('utf8'));
+
+    expect(decoded).toEqual({
+      runId: testRunId,
+      configKey,
+      variationKey,
+      version,
+    });
+    expect('graphKey' in decoded).toBe(false);
+  });
+
+  it('reconstructs tracker with graphKey from resumption token', () => {
+    const original = new LDAIConfigTrackerImpl(
+      mockLdClient,
+      testRunId,
+      configKey,
+      variationKey,
+      version,
+      modelName,
+      providerName,
+      testContext,
+      'my-graph',
+    );
+
+    const reconstructed = LDAIConfigTrackerImpl.fromResumptionToken(
+      original.resumptionToken,
+      mockLdClient,
+      testContext,
+    );
+
+    expect(reconstructed.getTrackData().graphKey).toBe('my-graph');
+  });
+
+  it('reconstructed tracker without graphKey does not include graphKey in track data', () => {
+    const original = new LDAIConfigTrackerImpl(
+      mockLdClient,
+      testRunId,
+      configKey,
+      variationKey,
+      version,
+      modelName,
+      providerName,
+      testContext,
+    );
+
+    const reconstructed = LDAIConfigTrackerImpl.fromResumptionToken(
+      original.resumptionToken,
+      mockLdClient,
+      testContext,
+    );
+
+    expect('graphKey' in reconstructed.getTrackData()).toBe(false);
+  });
 });

packages/sdk/server-ai/src/LDAIClientImpl.ts

@@ -74,6 +74,7 @@ export class LDAIClientImpl implements LDAIClient {
     defaultValue: LDAIConfigDefaultKind,
     mode: LDAIConfigMode,
     variables?: Record<string, unknown>,
+    graphKey?: string,
   ): Promise<LDAIConfigKind> {
     const ldFlagValue = LDAIConfigUtils.toFlagValue(defaultValue, mode);
 
@@ -101,6 +102,7 @@ export class LDAIClientImpl implements LDAIClient {
         value.model?.name ?? '',
         value.provider?.name ?? '',
         context,
+        graphKey,
       );
 
     const config = LDAIConfigUtils.fromFlagValue(key, value, trackerFactory);
@@ -217,21 +219,25 @@ export class LDAIClientImpl implements LDAIClient {
     return this._judgeConfig(key, context, defaultValue ?? disabledAIConfig, variables);
   }
 
+  private async _agentConfig(
+    key: string,
+    context: LDContext,
+    defaultValue: LDAIAgentConfigDefault,
+    variables?: Record<string, unknown>,
+    graphKey?: string,
+  ): Promise<LDAIAgentConfig> {
+    const config = await this._evaluate(key, context, defaultValue, 'agent', variables, graphKey);
+    return config as LDAIAgentConfig;
+  }
+
   async agentConfig(
     key: string,
     context: LDContext,
     defaultValue?: LDAIAgentConfigDefault,
     variables?: Record<string, unknown>,
   ): Promise<LDAIAgentConfig> {
     this._ldClient.track(TRACK_USAGE_AGENT_CONFIG, context, key, 1);
-    const config = await this._evaluate(
-      key,
-      context,
-      defaultValue ?? disabledAIConfig,
-      'agent',
-      variables,
-    );
-    return config as LDAIAgentConfig;
+    return this._agentConfig(key, context, defaultValue ?? disabledAIConfig, variables);
   }
 
   /**
@@ -261,14 +267,13 @@ export class LDAIClientImpl implements LDAIClient {
 
     await Promise.all(
       agentConfigs.map(async (config) => {
-        const agent = await this._evaluate(
+        const agent = await this._agentConfig(
           config.key,
           context,
           config.defaultValue ?? disabledAIConfig,
-          'agent',
           config.variables,
         );
-        agents[config.key as T[number]['key']] = agent as LDAIAgentConfig;
+        agents[config.key as T[number]['key']] = agent;
       }),
     );