Skip to content

refactor: restructure Graph class to utilize CompiledGraph#13

Merged
cantemizyurek merged 2 commits intomainfrom
focused-maxwell
Jan 25, 2026
Merged

refactor: restructure Graph class to utilize CompiledGraph#13
cantemizyurek merged 2 commits intomainfrom
focused-maxwell

Conversation

@cantemizyurek
Copy link
Owner

@cantemizyurek cantemizyurek commented Jan 25, 2026

Summary by CodeRabbit

  • New Features

    • Added a compiled graph execution mode with suspense support, checkpoint persistence, and lifecycle hooks.
    • Exposed a public SuspenseError and a CompiledGraph for executing compiled graphs.
    • Graphs can be compiled with runtime options (storage, onStart/onFinish) before execution.
  • Breaking Changes

    • Graph constructor and factory no longer accept execution options; use compile(...) to configure execution.
  • Tests

    • Tests updated to compile graphs before executing.

✏️ Tip: You can customize this high-level summary in your review settings.

…ion, enhancing modularity and maintainability
@changeset-bot
Copy link

changeset-bot bot commented Jan 25, 2026

🦋 Changeset detected

Latest commit: 7c468f8

The changes in this PR will be included in the next version bump.

This PR includes changesets to release 1 package
Name Type
ai-sdk-graph Minor

Not sure what this means? Click here to learn what changesets are.

Click here if you're a maintainer who wants to add another changeset to this PR

@coderabbitai
Copy link

coderabbitai bot commented Jan 25, 2026

📝 Walkthrough

Walkthrough

This change separates graph definition from execution by introducing CompiledGraph, which handles execution, suspense, checkpointing, state management, events, and subgraph execution; Graph becomes a lightweight builder with a new compile() method to produce an executable CompiledGraph.

Changes

Cohort / File(s) Summary
CompiledGraph Execution Engine
src/compiled-graph.ts
New module adding SuspenseError and CompiledGraph with node/edge/subgraph registries, storage (default InMemory), event emitter, state manager, checkpoint persistence, lifecycle hooks, suspense strategies (return/throw), and subgraph child execution.
Graph Refactoring
src/graph.ts
Removed runtime execution internals from Graph; constructor is now parameterless; added compile() method returning CompiledGraph; moved SuspenseError export to compiled-graph; built-in nodes inlined.
Public API Export
src/index.ts
Re-exported ./compiled-graph to expose CompiledGraph and SuspenseError on the public API.
Type Additions
src/types.ts
Added GraphSDK.CompileOptions<State, NodeKeys> with optional storage, onStart, and onFinish to configure compiled execution.
Tests Updated for Compilation Flow
tests/graph.test.ts
Tests changed to call g.compile(...) and execute the resulting CompiledGraph (storage and lifecycle hooks supplied via compile options); execution-related test flows updated accordingly.
Docs / Changeset
.changeset/separate-graph-runner.md
New changeset describing the breaking change: graph() no longer accepts options; use graph().compile(options) for execution configuration.

Sequence Diagrams

sequenceDiagram
    participant User
    participant Graph
    participant Compiled as CompiledGraph
    participant Storage
    participant State as StateManager
    participant Node

    User->>Graph: define graph
    User->>Graph: compile(options)
    Graph->>Compiled: new CompiledGraph(registries, options)
    Compiled->>Storage: init (InMemoryStorage default)
    Compiled-->>User: return CompiledGraph

    User->>Compiled: execute/run
    Compiled->>Compiled: executeInternal (checkpoint check)
    alt checkpoint exists
        Compiled->>Storage: restore state & pending nodes
    else no checkpoint
        Compiled->>State: init start state
    end
    Compiled->>Compiled: call onStart
    loop execution loop
        Compiled->>Node: execute(state, suspendFn)
        alt node completes
            Node->>State: apply updates
            Node-->>Compiled: result
        else node suspends
            Node-->>Compiled: throw SuspenseError(data)
            Compiled->>Storage: persist checkpoint (state, nodes)
            Compiled-->>User: emit/return suspense
        end
    end
    Compiled->>Compiled: call onFinish
    Compiled-->>User: return final state
Loading
sequenceDiagram
    participant Parent as ParentCompiledGraph
    participant SubNode as SubgraphNode
    participant Child as ChildCompiledGraph
    participant Storage as ChildStorage

    Parent->>SubNode: invoke subgraph node
    SubNode->>Child: create child CompiledGraph (registries, options)
    Child->>Storage: init child storage
    Child->>Child: execute child nodes
    alt child suspends
        Child->>Storage: persist child checkpoint
        Child-->>SubNode: propagate SuspenseError(data)
    else child completes
        Child-->>SubNode: return child final state
        SubNode->>Parent: map outputs into parent state
    end
Loading

Estimated Code Review Effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Poem

🐰 I stitched the graph from tip to tail,

then compiled it swift down the rabbit trail.
It pauses, saves, then springs anew—
subgraphs hop through, state follows too.
A little thunk and off we run—hooray!

🚥 Pre-merge checks | ✅ 2 | ❌ 1
❌ Failed checks (1 warning)
Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (2 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and specifically describes the main architectural change in the PR: restructuring the Graph class to utilize CompiledGraph.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing touches
  • 📝 Generate docstrings
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Post copyable unit tests in a comment
  • Commit unit tests in branch focused-maxwell

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

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

Actionable comments posted: 2

🤖 Fix all issues with AI agents
In `@src/compiled-graph.ts`:
- Around line 238-246: executeBatch currently runs executeSingleNode in parallel
via Promise.all which can cause race conditions because individual nodes call
update() to mutate context.state; either document this limitation or make state
mutations deterministic by changing executeBatch to (a) run nodes serially
(await each executeSingleNode in a for...of) so updates are applied in a defined
order, or (b) have executeSingleNode return a collected state-delta instead of
applying immediately and then merge/apply those deltas to context.state in a
deterministic order after all promises resolve; update references: executeBatch,
executeSingleNode, update(), and context.state.

In `@tests/graph.test.ts`:
- Line 289: The test uses an unnecessary cast "storage as any" when calling
g.compile(...); fix it by instantiating the InMemoryStorage with the correct
NodeKeys generic to match the graph (e.g., change new InMemoryStorage<{ value:
number }, string>() to new InMemoryStorage<{ value: number }, 'START' |
'END'>()) so g.compile accepts it without casting; update the InMemoryStorage
generic to the actual node identifiers for any custom graph, then remove the "as
any" cast from the runGraph(g.compile(...)) call.
🧹 Nitpick comments (2)
src/types.ts (1)

85-92: Consider consolidating duplicate type definitions.

CompileOptions is identical to GraphOptions. If they're expected to diverge in the future, this separation is fine. Otherwise, consider using a shared type or type alias to reduce duplication.

// Option 1: Type alias
export type CompileOptions<State, NodeKeys> = GraphOptions<State, NodeKeys>

// Option 2: Shared base (if divergence is expected)
interface BaseOptions<State, NodeKeys> {
  storage?: GraphStorage<State, NodeKeys>
  onFinish?: (args: { state: State }) => Promise<void> | void
  onStart?: (args: { state: State; writer: UIMessageStreamWriter }) => Promise<void> | void
}
export interface GraphOptions<State, NodeKeys> extends BaseOptions<State, NodeKeys> {}
export interface CompileOptions<State, NodeKeys> extends BaseOptions<State, NodeKeys> {}
src/compiled-graph.ts (1)

13-21: Add Object.setPrototypeOf for robust error inheritance.

When extending built-in classes like Error in TypeScript/ES6, the prototype chain may not be properly set in some transpilation targets (ES5). This can cause instanceof SuspenseError checks to fail.

♻️ Proposed fix
 export class SuspenseError extends Error {
   readonly data?: unknown

   constructor(data?: unknown) {
     super('Suspense')
     this.name = 'SuspenseError'
     this.data = data
+    Object.setPrototypeOf(this, SuspenseError.prototype)
   }
 }

Comment on lines +238 to +246
private async executeBatch(
nodes: GraphSDK.Node<State, NodeKeys>[],
context: GraphSDK.ExecutionContext<State, NodeKeys>
): Promise<Array<{ node: GraphSDK.Node<State, NodeKeys>; error: SuspenseError }>> {
const results = await Promise.all(
nodes.map((node) => this.executeSingleNode(node, context))
)
return results.filter((r): r is { node: GraphSDK.Node<State, NodeKeys>; error: SuspenseError } => r !== null)
}
Copy link

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟠 Major

Parallel state updates may cause race conditions.

executeBatch runs nodes in parallel via Promise.all. Each node can call update() which mutates context.state. If two parallel nodes update the same state property, the final value depends on execution order, potentially causing lost updates.

The tests "parallel nodes can update different state properties" pass because they update different properties. However, if two parallel nodes update the same property, the behavior is undefined.

Consider documenting this limitation or implementing synchronized state updates if concurrent writes to the same property should be supported.

🤖 Prompt for AI Agents
In `@src/compiled-graph.ts` around lines 238 - 246, executeBatch currently runs
executeSingleNode in parallel via Promise.all which can cause race conditions
because individual nodes call update() to mutate context.state; either document
this limitation or make state mutations deterministic by changing executeBatch
to (a) run nodes serially (await each executeSingleNode in a for...of) so
updates are applied in a defined order, or (b) have executeSingleNode return a
collected state-delta instead of applying immediately and then merge/apply those
deltas to context.state in a deterministic order after all promises resolve;
update references: executeBatch, executeSingleNode, update(), and context.state.

… separation of concerns and execution flexibility
@cantemizyurek cantemizyurek merged commit 227c948 into main Jan 25, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant