Skip to content
161 changes: 107 additions & 54 deletions pkgs/sdk/server-ai/src/Graph/AgentGraphDefinition.cs

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are a lot of new dictionary calls, toDictionary, as well as several contains and order by. Do you have any performance requirements wrt throughput or memory allocation delays? May be worth a quick benchmark to see if your expected cases are well behaved.

I'm guessing this would apply to the other PRs for the other SDKs if they are following similar implementations.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No hard requirements here — traversal runs once per graph invocation and every callback makes an LLM call, so we're spending microseconds against seconds. Benchmarked to confirm: a 20-node graph is ~65 µs / 60 KB and a 50-node fan-out ~150 µs, with degradation only past a few hundred nodes, well beyond expected graph sizes.

@mattrmc1 mattrmc1 Aug 26, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Took a look at the space concerns too. The old approach was actually worse since it carried context level to level. Now, if you have a fan out graph, each node gets initialContext plus ONLY its transitive ancestors. The old way gave each node all context from levels above as well as sibling context depending on order of BFS traversal. Assuming the the graph is a diamond shape, the bloated context on the exit node is the same either way.

        a
       / \
      b   c
      |   |
      d   e
       \ /
        f
node old context new context
a {} {}
b {a} {a}
c {a, b} {a}
d {a, b, c} {a, b}
e {a, b, c, d} {a, c}
f {a, b, c, d, e} {a, b, c, d, e}

Original file line number Diff line number Diff line change
Expand Up @@ -94,12 +94,13 @@ public IReadOnlyList<AgentGraphNode> TerminalNodes() =>
public AiGraphTracker CreateTracker() => _createTracker();

/// <summary>
/// Performs a breadth-first traversal of the graph starting from the root node.
/// For each visited node, <paramref name="fn"/> is called with the node and the
/// accumulated context dictionary. The return value of <paramref name="fn"/> is
/// stored in the context under the node's key and passed to subsequent calls.
/// Cycle-safe: each node is visited at most once.
/// Visits each reachable node in topological order (predecessors first, root first).
/// Ties break by discovery order. Cycle-safe.
/// </summary>
/// <remarks>
/// <paramref name="fn"/> receives <paramref name="initialContext"/> plus results
/// from that node's reachable predecessors only.
/// </remarks>
public void Traverse(
Func<AgentGraphNode, Dictionary<string, object>, object> fn,
Dictionary<string, object> initialContext = null)
Expand All @@ -108,91 +109,143 @@ public void Traverse(
if (root == null) return;

var context = initialContext ?? new Dictionary<string, object>();
var (reachable, order) = ReachableAndDiscovery(root.Key);

var indeg = reachable.ToDictionary(k => k, _ => 0);
foreach (var k in reachable)
{
foreach (var e in GetNode(k).Edges)
{
if (reachable.Contains(e.Key)) indeg[e.Key]++;
}
}
indeg[root.Key] = 0;

var visited = new HashSet<string>();
var queue = new Queue<AgentGraphNode>();
queue.Enqueue(root);
visited.Add(root.Key);
var results = new Dictionary<string, object>();
var ancestors = new Dictionary<string, HashSet<string>>();

while (queue.Count > 0)
Dictionary<string, object> Scoped(HashSet<string> deps)
{
var node = queue.Dequeue();
var result = fn(node, context);
context[node.Key] = result;
var c = new Dictionary<string, object>(context);
foreach (var k in deps) c[k] = results[k];
return c;
}

foreach (var child in GetChildNodes(node.Key))
while (visited.Count < reachable.Count)
{
var next = order.FirstOrDefault(k => !visited.Contains(k) && indeg[k] == 0)
?? order.Where(k => !visited.Contains(k)).OrderBy(k => indeg[k]).First();

var anc = new HashSet<string>();
foreach (var parent in GetParentNodes(next))
{
if (visited.Add(child.Key))
{
queue.Enqueue(child);
}
if (!visited.Contains(parent.Key)) continue;
anc.Add(parent.Key);
anc.UnionWith(ancestors[parent.Key]);
}
ancestors[next] = anc;
visited.Add(next);

results[next] = fn(GetNode(next), Scoped(anc));
foreach (var e in GetNode(next).Edges)
{
if (reachable.Contains(e.Key)) indeg[e.Key]--;
}
}
}

/// <summary>
/// Performs a reverse breadth-first traversal: starts from terminal nodes and
/// works upward toward the root. The root node is always processed last.
/// Cycle-safe: each node is visited at most once.
/// Visits each reachable node in reverse topological order (descendants first,
/// root last). Ties break by discovery order. Cycle-safe.
/// </summary>
/// <remarks>
/// <paramref name="fn"/> receives <paramref name="initialContext"/> plus results
/// from that node's reachable descendants only.
/// </remarks>
public void ReverseTraverse(
Func<AgentGraphNode, Dictionary<string, object>, object> fn,
Dictionary<string, object> initialContext = null)
{
if (_nodes.Count == 0) return;
var root = RootNode();
if (root == null) return;

var rootKey = root.Key;
var context = initialContext ?? new Dictionary<string, object>();
var (reachable, order) = ReachableAndDiscovery(rootKey);

var outdeg = reachable.ToDictionary(
k => k, k => GetNode(k).Edges.Count(e => reachable.Contains(e.Key)));
Comment thread
mattrmc1 marked this conversation as resolved.

var root = RootNode();
var visited = new HashSet<string>();
var queue = new Queue<AgentGraphNode>();
var results = new Dictionary<string, object>();
var descendants = new Dictionary<string, HashSet<string>>();

// Seed with terminal nodes (excluding root if it happens to be terminal and there are others)
foreach (var terminal in TerminalNodes())
Dictionary<string, object> Scoped(HashSet<string> deps)
{
if (root != null && terminal.Key == root.Key && _nodes.Count > 1)
{
continue;
}
if (visited.Add(terminal.Key))
{
queue.Enqueue(terminal);
}
var c = new Dictionary<string, object>(context);
foreach (var k in deps) c[k] = results[k];
return c;
}

while (queue.Count > 0)
bool NonRootRemaining() => reachable.Any(k => k != rootKey && !visited.Contains(k));
while (NonRootRemaining())
{
var node = queue.Dequeue();
var next = order.FirstOrDefault(k => k != rootKey && !visited.Contains(k) && outdeg[k] == 0)
?? order.Where(k => k != rootKey && !visited.Contains(k)).OrderBy(k => outdeg[k]).First();

// Defer root until the very end (unless it's the only node)
if (root != null && node.Key == root.Key && _nodes.Count > 1)
var desc = new HashSet<string>();
foreach (var e in GetNode(next).Edges)
{
continue;
if (!reachable.Contains(e.Key) || !visited.Contains(e.Key)) continue;
desc.Add(e.Key);
desc.UnionWith(descendants[e.Key]);
}
descendants[next] = desc;
visited.Add(next);

results[next] = fn(GetNode(next), Scoped(desc));
foreach (var parent in GetParentNodes(next))
{
if (parent.Key != rootKey && reachable.Contains(parent.Key))
outdeg[parent.Key]--;
}
}

// Root last
var rootDeps = new HashSet<string>(reachable.Where(k => k != rootKey));
results[rootKey] = fn(root, Scoped(rootDeps));
}

var result = fn(node, context);
context[node.Key] = result;
/// <summary>
/// Reachable nodes from <paramref name="rootKey"/> and BFS discovery order
/// (declared edge order). Used as a topological tie-break.
/// </summary>
private (HashSet<string> reachable, List<string> order) ReachableAndDiscovery(string rootKey)
{
var reachable = new HashSet<string>();
var order = new List<string>();
var queue = new Queue<string>();
reachable.Add(rootKey);
order.Add(rootKey);
queue.Enqueue(rootKey);

foreach (var parent in GetParentNodes(node.Key))
while (queue.Count > 0)
{
var key = queue.Dequeue();
var node = GetNode(key);
if (node == null) continue;
foreach (var edge in node.Edges)
{
if (root != null && parent.Key == root.Key)
{
// Don't enqueue root yet — process it last
continue;
}
if (visited.Add(parent.Key))
if (GetNode(edge.Key) != null && reachable.Add(edge.Key))
{
queue.Enqueue(parent);
order.Add(edge.Key);
queue.Enqueue(edge.Key);
}
}
}

// Process root last
if (root != null && _nodes.Count > 1 && visited.Count > 0)
{
var result = fn(root, context);
context[root.Key] = result;
}
return (reachable, order);
}

/// <summary>
Expand Down
Loading
Loading