feat: Implement agent graph definitions

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

@@ -0,0 +1,206 @@
+import { LDContext } from '@launchdarkly/js-server-sdk-common';
+
+import { AgentGraphDefinition } from '../src/api/graph/AgentGraphDefinition';
+import { LDAIClientImpl } from '../src/LDAIClientImpl';
+import { LDClientMin } from '../src/LDClientMin';
+
+// ---------------------------------------------------------------------------
+// Mocks
+// ---------------------------------------------------------------------------
+
+const mockTrack = jest.fn();
+const mockVariation = jest.fn();
+const mockDebug = jest.fn();
+
+const mockLdClient: LDClientMin = {
+  track: mockTrack,
+  variation: mockVariation,
+  logger: {
+    debug: mockDebug,
+    info: jest.fn(),
+    warn: jest.fn(),
+    error: jest.fn(),
+  },
+};
+
+const testContext: LDContext = { kind: 'user', key: 'test-user' };
+
+const makeClient = () => new LDAIClientImpl(mockLdClient);
+
+beforeEach(() => {
+  jest.clearAllMocks();
+});
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function makeGraphFlagValue(
+  root: string,
+  edges: Record<string, Array<{ key: string }>> = {},
+  variationKey = 'v1',
+  version = 1,
+) {
+  return { _ldMeta: { variationKey, version }, root, edges };
+}
+
+function makeAgentFlagValue(key: string, enabled = true) {
+  return {
+    _ldMeta: { variationKey: `${key}-v1`, enabled, version: 1, mode: 'agent' },
+    instructions: `Instructions for ${key}`,
+  };
+}
+
+// ---------------------------------------------------------------------------
+// agentGraph – disabled / validation failures
+// ---------------------------------------------------------------------------
+
+it('returns { enabled: false } when graph flag has no root', async () => {
+  const client = makeClient();
+  mockVariation.mockResolvedValueOnce({ root: '' }); // no root
+  const result = await client.agentGraph('my-graph', testContext);
+  expect(result.enabled).toBe(false);
+});
+
+it('logs debug when graph has no root', async () => {
+  const client = makeClient();
+  mockVariation.mockResolvedValueOnce({ root: '' });
+  await client.agentGraph('my-graph', testContext);
+  expect(mockDebug).toHaveBeenCalledWith(expect.stringContaining('not fetchable'));
+});
+
+it('returns { enabled: false } when a node is unconnected (not reachable from root)', async () => {
+  const client = makeClient();
+  // Graph says root → child, but "orphan" appears in edges with no path from root
+  const graphValue = makeGraphFlagValue('root', {
+    root: [{ key: 'child' }],
+    orphan: [{ key: 'other' }], // orphan is not reachable from root
+  });
+  mockVariation.mockResolvedValueOnce(graphValue);
+  const result = await client.agentGraph('my-graph', testContext);
+  expect(result.enabled).toBe(false);
+  expect(mockDebug).toHaveBeenCalledWith(expect.stringContaining('unconnected node'));
+});
+
+it('returns { enabled: false } when the graph contains a cycle', async () => {
+  const client = makeClient();
+  // A → B → A forms a cycle
+  const graphValue = makeGraphFlagValue('a', {
+    a: [{ key: 'b' }],
+    b: [{ key: 'a' }],
+  });
+  mockVariation.mockResolvedValueOnce(graphValue);
+  const result = await client.agentGraph('my-graph', testContext);
+  expect(result.enabled).toBe(false);
+  expect(mockDebug).toHaveBeenCalledWith(expect.stringContaining('cycle'));
+});
+
+it('returns { enabled: false } when a child agent config is disabled', async () => {
+  const client = makeClient();
+  const graphValue = makeGraphFlagValue('root', { root: [{ key: 'child' }] });
+  mockVariation
+    .mockResolvedValueOnce(graphValue) // graph flag
+    .mockResolvedValueOnce(makeAgentFlagValue('root', true)) // root agent config
+    .mockResolvedValueOnce(makeAgentFlagValue('child', false)); // child is disabled
+  const result = await client.agentGraph('my-graph', testContext);
+  expect(result.enabled).toBe(false);
+  expect(mockDebug).toHaveBeenCalledWith(expect.stringContaining('not enabled'));
+});
+
+// ---------------------------------------------------------------------------
+// agentGraph – success path
+// ---------------------------------------------------------------------------
+
+it('returns { enabled: true, graph } for a valid graph with a single node', async () => {
+  const client = makeClient();
+  const graphValue = makeGraphFlagValue('solo-agent');
+  mockVariation
+    .mockResolvedValueOnce(graphValue)
+    .mockResolvedValueOnce(makeAgentFlagValue('solo-agent'));
+  const result = await client.agentGraph('my-graph', testContext);
+  expect(result.enabled).toBe(true);
+  if (result.enabled) {
+    expect(result.graph).toBeInstanceOf(AgentGraphDefinition);
+    expect(result.graph.rootNode().getKey()).toBe('solo-agent');
+  }
+});
+
+it('returns a valid AgentGraphDefinition with correct nodes for multi-node graph', async () => {
+  const client = makeClient();
+  const graphValue = makeGraphFlagValue('root', {
+    root: [{ key: 'child-a' }, { key: 'child-b' }],
+    'child-a': [{ key: 'leaf' }],
+  });
+  // variation is called for: graph flag + root + child-a + child-b + leaf (order may vary)
+  mockVariation
+    .mockResolvedValueOnce(graphValue)
+    .mockResolvedValue(makeAgentFlagValue('agent', true)); // all agent configs succeed
+
+  const result = await client.agentGraph('my-graph', testContext);
+  expect(result.enabled).toBe(true);
+  if (result.enabled) {
+    const { graph } = result;
+    expect(graph.rootNode().getKey()).toBe('root');
+    expect(
+      graph
+        .getChildNodes('root')
+        .map((n) => n.getKey())
+        .sort(),
+    ).toEqual(['child-a', 'child-b']);
+    expect(
+      graph
+        .terminalNodes()
+        .map((n) => n.getKey())
+        .sort(),
+    ).toEqual(['child-b', 'leaf']);
+  }
+});
+
+it('tracks usage event when agentGraph is called', async () => {
+  const client = makeClient();
+  mockVariation.mockResolvedValue({ root: '' });
+  await client.agentGraph('my-graph', testContext);
+  expect(mockTrack).toHaveBeenCalledWith('$ld:ai:usage:agent-graph', testContext, 'my-graph', 1);
+});
+
+it('createTracker on returned graph produces a tracker with correct graphKey', async () => {
+  const client = makeClient();
+  const graphValue = makeGraphFlagValue('root', {}, 'varKey', 3);
+  mockVariation
+    .mockResolvedValueOnce(graphValue)
+    .mockResolvedValueOnce(makeAgentFlagValue('root', true));
+
+  const result = await client.agentGraph('graph-key', testContext);
+  expect(result.enabled).toBe(true);
+  if (result.enabled) {
+    const tracker = result.graph.createTracker();
+    expect(tracker.getTrackData().graphKey).toBe('graph-key');
+    expect(tracker.getTrackData().version).toBe(3);
+    expect(tracker.getTrackData().variationKey).toBe('varKey');
+  }
+});
+
+// ---------------------------------------------------------------------------
+// createGraphTracker
+// ---------------------------------------------------------------------------
+
+it('createGraphTracker reconstructs a tracker from a resumption token', async () => {
+  const client = makeClient();
+  const graphValue = makeGraphFlagValue('root', {}, 'v99', 7);
+  mockVariation
+    .mockResolvedValueOnce(graphValue)
+    .mockResolvedValueOnce(makeAgentFlagValue('root', true));
+
+  const result = await client.agentGraph('g-key', testContext);
+  expect(result.enabled).toBe(true);
+  if (result.enabled) {
+    const originalTracker = result.graph.createTracker();
+    const token = originalTracker.resumptionToken;
+
+    const reconstructed = client.createGraphTracker(token, testContext);
+    expect(reconstructed.getTrackData().graphKey).toBe('g-key');
+    expect(reconstructed.getTrackData().version).toBe(7);
+    expect(reconstructed.getTrackData().variationKey).toBe('v99');
+    expect(reconstructed.getTrackData().runId).toBe(originalTracker.getTrackData().runId);
+  }
+});

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

@@ -21,11 +21,13 @@ import {
   LDMessage,
 } from './api/config';
 import { LDAIConfigFlagValue, LDAIConfigUtils } from './api/config/LDAIConfigUtils';
+import { AgentGraphDefinition, LDAgentGraphFlagValue, LDGraphTracker } from './api/graph';
 import { Judge } from './api/judge/Judge';
 import { LDAIClient } from './api/LDAIClient';
 import { AIProviderFactory, SupportedAIProvider } from './api/providers';
 import { LDAIConfigTrackerImpl } from './LDAIConfigTrackerImpl';
 import { LDClientMin } from './LDClientMin';
+import { LDGraphTrackerImpl } from './LDGraphTrackerImpl';
 import { aiSdkLanguage, aiSdkName, aiSdkVersion } from './sdkInfo';
 
 /**
@@ -38,6 +40,7 @@ const TRACK_USAGE_JUDGE_CONFIG = '$ld:ai:usage:judge-config';
 const TRACK_USAGE_CREATE_JUDGE = '$ld:ai:usage:create-judge';
 const TRACK_USAGE_AGENT_CONFIG = '$ld:ai:usage:agent-config';
 const TRACK_USAGE_AGENT_CONFIGS = '$ld:ai:usage:agent-configs';
+const TRACK_USAGE_AGENT_GRAPH = '$ld:ai:usage:agent-graph';
 
 const INIT_TRACK_CONTEXT: LDContext = {
   kind: 'ld_ai',
@@ -388,4 +391,148 @@ export class LDAIClientImpl implements LDAIClient {
   createTracker(token: string, context: LDContext): LDAIConfigTracker {
     return LDAIConfigTrackerImpl.fromResumptionToken(token, this._ldClient, context);
   }
+
+  async agentGraph(
+    graphKey: string,
+    context: LDContext,
+  ): Promise<{ enabled: false } | { enabled: true; graph: AgentGraphDefinition }> {
+    this._ldClient.track(TRACK_USAGE_AGENT_GRAPH, context, graphKey, 1);
+
+    const disabled = { enabled: false as const };
+
+    // Step 1: Fetch the graph flag value
+    const defaultGraphValue: LDAgentGraphFlagValue = { root: '' };
+    const graphFlagValue = (await this._ldClient.variation(
+      graphKey,
+      context,
+      defaultGraphValue,
+    )) as LDAgentGraphFlagValue;
+
+    // Step 2: Validate - graph must be fetchable (has a root)
+    if (!graphFlagValue.root) {
+      this._logger?.debug(`agentGraph: graph "${graphKey}" is not fetchable or has no root node.`);
+      return disabled;
+    }
+
+    // Step 3: Validate - collect all node keys and check connectivity from root
+    const allKeys = AgentGraphDefinition.collectAllKeys(graphFlagValue);
+    const reachableKeys = this._collectReachableKeys(graphFlagValue);
+
+    const unreachableKey = [...allKeys].find((key) => !reachableKeys.has(key));
+    if (unreachableKey) {
+      this._logger?.debug(
+        `agentGraph: graph "${graphKey}" has unconnected node "${unreachableKey}" that is not reachable from the root.`,
+      );
+      return disabled;
+    }
+
+    // Step 4: Validate - graph must be acyclic (cycles cause reverseTraverse to hang)
+    if (this._hasCycle(graphFlagValue)) {
+      this._logger?.debug(`agentGraph: graph "${graphKey}" contains a cycle and cannot be used.`);
+      return disabled;
+    }
+
+    // Step 5: Validate - fetch all child agent configs
+    const agentConfigs: Record<string, LDAIAgentConfig> = {};
+    const fetchResults = await Promise.all(
+      [...allKeys].map(async (key) => {
+        const config = await this._agentConfigInternal(key, context);
+        return { key, config };
+      }),
+    );
+
+    const disabledResult = fetchResults.find(({ config }) => !config.enabled);
+    if (disabledResult) {
+      this._logger?.debug(
+        `agentGraph: agent config "${disabledResult.key}" in graph "${graphKey}" is not enabled or could not be fetched.`,
+      );
+      return disabled;
+    }
+    fetchResults.forEach(({ key, config }) => {
+      agentConfigs[key] = config;
+    });
+
+    // Build the node map
+    const nodes = AgentGraphDefinition.buildNodes(graphFlagValue, agentConfigs);
+
+    // eslint-disable-next-line no-underscore-dangle
+    const variationKey = graphFlagValue._ldMeta?.variationKey;
+    // eslint-disable-next-line no-underscore-dangle
+    const version = graphFlagValue._ldMeta?.version ?? 1;
+    const ldClient = this._ldClient;
+
+    const graph = new AgentGraphDefinition(
+      graphFlagValue,
+      nodes,
+      () =>
+        new LDGraphTrackerImpl(ldClient, randomUUID(), graphKey, variationKey, version, context),
+    );
+
+    return { enabled: true, graph };
+  }
+
+  createGraphTracker(token: string, context: LDContext): LDGraphTracker {
+    return LDGraphTrackerImpl.fromResumptionToken(token, this._ldClient, context);
+  }
+
+  /**
+   * Fetches a single agent config without tracking usage (used internally by agentGraph).
+   */
+  private async _agentConfigInternal(key: string, context: LDContext): Promise<LDAIAgentConfig> {
+    const config = await this._evaluate(key, context, disabledAIConfig, 'agent');
+    return config as LDAIAgentConfig;
+  }
+
+  /**
+   * Returns true if the graph contains a cycle, using DFS with three-color marking.
+   */
+  private _hasCycle(graph: LDAgentGraphFlagValue): boolean {
+    const color: Record<string, 'gray' | 'black'> = {};
+
+    const dfs = (key: string): boolean => {
+      color[key] = 'gray';
+      const edges = graph.edges?.[key] ?? [];
+      const foundCycle = edges.some((edge) => {
+        if (color[edge.key] === 'gray') {
+          return true;
+        }
+        if (color[edge.key] === 'black') {
+          return false;
+        }
+        return dfs(edge.key);
+      });
+      color[key] = 'black';
+      return foundCycle;
+    };
+
+    const allKeys = AgentGraphDefinition.collectAllKeys(graph);
+    return [...allKeys].some((key) => {
+      if (!color[key]) {
+        return dfs(key);
+      }
+      return false;
+    });
+  }
+
+  /**
+   * Returns the set of all node keys reachable from the root via BFS.
+   */
+  private _collectReachableKeys(graph: LDAgentGraphFlagValue): Set<string> {
+    const visited = new Set<string>();
+    const queue: string[] = [graph.root];
+    visited.add(graph.root);
+
+    while (queue.length > 0) {
+      const key = queue.shift()!;
+      const edges = graph.edges?.[key] ?? [];
+      edges.forEach((edge) => {
+        if (!visited.has(edge.key)) {
+          visited.add(edge.key);
+          queue.push(edge.key);
+        }
+      });
+    }
+
+    return visited;
+  }
 }