Implementation:Microsoft Autogen Studio Builder Store

Sources	Microsoft_Autogen
Domains	State Management, Zustand Store, Graph State, Undo/Redo
Last Updated	2026-02-11 17:00 GMT

Overview

Description

The Studio_Builder_Store module implements the Zustand state management store for the AutoGen Studio team builder. This store manages the complete graph state including nodes, edges, selection, history, and synchronization between the visual graph and JSON team configurations. It provides actions for adding, updating, and removing nodes and edges, as well as undo/redo functionality.

Key features include:

Graph State Management - Maintains nodes and edges for the React Flow canvas
History Management - Implements undo/redo with a maximum of 50 history states
Component Integration - Handles adding models, tools, agents, and teams to the graph
JSON Synchronization - Converts between visual graph and JSON team configurations
Layout Management - Integrates with layout utilities for node positioning
Workbench Normalization - Handles both single and array workbench formats

The store uses Zustand for reactive state management and integrates with layout storage for persisting user-positioned nodes.

Usage

The store is used throughout the team builder to:

Add new components via drag-and-drop
Update node configurations
Remove nodes and their connections
Sync visual changes back to JSON format
Load team configurations from JSON
Track and restore history states
Manage selected node state

Code Reference

Source Location

python/packages/autogen-studio/frontend/src/components/views/teambuilder/builder/store.tsx

Repository: https://github.com/microsoft/autogen

Full path: /tmp/kapso_repo_2mr4n2g4/python/packages/autogen-studio/frontend/src/components/views/teambuilder/builder/store.tsx

Signature

// Main store interface
export interface TeamBuilderState {
  // State
  nodes: CustomNode[];
  edges: CustomEdge[];
  selectedNodeId: string | null;
  history: Array<{ nodes: CustomNode[]; edges: CustomEdge[] }>;
  currentHistoryIndex: number;
  originalComponent: Component<TeamConfig> | null;
  teamId: string | null;

  // Actions
  addNode: (
    position: Position,
    component: Component<ComponentConfig>,
    targetNodeId: string
  ) => void;
  updateNode: (nodeId: string, updates: Partial<NodeData>) => void;
  removeNode: (nodeId: string) => void;
  addEdge: (edge: CustomEdge) => void;
  removeEdge: (edgeId: string) => void;
  setSelectedNode: (nodeId: string | null) => void;
  setNodeUserPositioned: (nodeId: string, position: Position) => void;
  undo: () => void;
  redo: () => void;
  syncToJson: () => Component<TeamConfig> | null;
  loadFromJson: (
    config: Component<TeamConfig>,
    isInitialLoad?: boolean,
    teamId?: string
  ) => GraphState;
  layoutNodes: () => void;
  resetHistory: () => void;
  addToHistory: () => void;
}

// Store hook
export const useTeamBuilderStore: () => TeamBuilderState

Import

import { useTeamBuilderStore } from './store';
import type { TeamBuilderState } from './store';

I/O Contract

Inputs

addNode Parameters:

position - {x, y} coordinates for the new node
component - Component<ComponentConfig> to add to the graph
targetNodeId - ID of node to connect to (empty string for standalone)

updateNode Parameters:

nodeId - ID of node to update
updates - Partial<NodeData> containing fields to update

removeNode Parameters:

nodeId - ID of node to remove (and its connected edges)

loadFromJson Parameters:

config - Component<TeamConfig> team configuration
isInitialLoad - Boolean indicating if this is the first load (default: false)
teamId - Optional team ID for layout storage lookup

setNodeUserPositioned Parameters:

nodeId - ID of node that was manually positioned
position - {x, y} coordinates where user placed the node

Outputs

State Fields:

nodes - CustomNode[] array of all graph nodes
edges - CustomEdge[] array of all connections
selectedNodeId - string | null currently selected node ID
history - Array of previous graph states (max 50)
currentHistoryIndex - Index in history array
originalComponent - Component<TeamConfig> | null original team config
teamId - string | null current team identifier

syncToJson Returns:

Component<TeamConfig> | null - Current team configuration in JSON format

loadFromJson Returns:

GraphState - Object with nodes and edges arrays

Action Side Effects:

addNode - Updates nodes, edges, and history
updateNode - Updates node data and related team participants
removeNode - Removes node, edges, and updates team participants
undo/redo - Restores previous/next history state
layoutNodes - Recalculates node positions

Usage Examples

Example 1: Using the Store in Components

import { useTeamBuilderStore } from './store';

function TeamBuilderCanvas() {
  // Select state
  const nodes = useTeamBuilderStore((state) => state.nodes);
  const edges = useTeamBuilderStore((state) => state.edges);
  const selectedNodeId = useTeamBuilderStore((state) => state.selectedNodeId);

  // Select actions
  const addNode = useTeamBuilderStore((state) => state.addNode);
  const updateNode = useTeamBuilderStore((state) => state.updateNode);
  const setSelectedNode = useTeamBuilderStore((state) => state.setSelectedNode);

  return (
    <ReactFlow
      nodes={nodes}
      edges={edges}
      onNodeClick={(_, node) => setSelectedNode(node.id)}
    />
  );
}

Example 2: Adding Components to Graph

import { useTeamBuilderStore } from './store';
import { Component, AgentConfig } from '../../../types/datamodel';

function DragDropHandler() {
  const addNode = useTeamBuilderStore((state) => state.addNode);

  const handleDrop = (
    position: { x: number; y: number },
    component: Component<AgentConfig>,
    targetNodeId: string
  ) => {
    // Add agent to team
    addNode(position, component, targetNodeId);
  };

  // Add model to agent
  const handleModelDrop = (modelComponent, agentNodeId) => {
    addNode({ x: 0, y: 0 }, modelComponent, agentNodeId);
  };

  // Add tool to agent (creates/updates workbench)
  const handleToolDrop = (toolComponent, agentNodeId) => {
    addNode({ x: 0, y: 0 }, toolComponent, agentNodeId);
  };
}

Example 3: Loading and Syncing Team Configurations

import { useTeamBuilderStore } from './store';
import { Component, TeamConfig } from '../../../types/datamodel';

function TeamLoader() {
  const loadFromJson = useTeamBuilderStore((state) => state.loadFromJson);
  const syncToJson = useTeamBuilderStore((state) => state.syncToJson);

  // Load team from JSON
  const loadTeam = (teamConfig: Component<TeamConfig>, teamId: string) => {
    const graphState = loadFromJson(teamConfig, true, teamId);
    console.log('Loaded team with', graphState.nodes.length, 'nodes');
  };

  // Save current state to JSON
  const saveTeam = () => {
    const teamConfig = syncToJson();
    if (teamConfig) {
      // Send to API or save locally
      console.log('Team config:', JSON.stringify(teamConfig, null, 2));
    }
  };

  return (
    <div>
      <button onClick={saveTeam}>Save Team</button>
    </div>
  );
}

Example 4: Undo/Redo Functionality

import { useTeamBuilderStore } from './store';

function UndoRedoControls() {
  const undo = useTeamBuilderStore((state) => state.undo);
  const redo = useTeamBuilderStore((state) => state.redo);
  const history = useTeamBuilderStore((state) => state.history);
  const currentIndex = useTeamBuilderStore((state) => state.currentHistoryIndex);

  const canUndo = currentIndex > 0;
  const canRedo = currentIndex < history.length - 1;

  // Keyboard shortcuts
  useEffect(() => {
    const handleKeyPress = (e: KeyboardEvent) => {
      if ((e.metaKey || e.ctrlKey) && e.key === 'z') {
        e.preventDefault();
        if (e.shiftKey) {
          redo();
        } else {
          undo();
        }
      }
    };

    window.addEventListener('keydown', handleKeyPress);
    return () => window.removeEventListener('keydown', handleKeyPress);
  }, [undo, redo]);

  return (
    <div className="flex gap-2">
      <button onClick={undo} disabled={!canUndo}>
        Undo (Ctrl+Z)
      </button>
      <button onClick={redo} disabled={!canRedo}>
        Redo (Ctrl+Shift+Z)
      </button>
    </div>
  );
}

Example 5: Updating Node Configuration

import { useTeamBuilderStore } from './store';
import { Component, AgentConfig } from '../../../types/datamodel';

function AgentEditor({ nodeId }: { nodeId: string }) {
  const updateNode = useTeamBuilderStore((state) => state.updateNode);
  const nodes = useTeamBuilderStore((state) => state.nodes);

  const node = nodes.find(n => n.id === nodeId);
  const component = node?.data.component as Component<AgentConfig>;

  const handleNameChange = (newName: string) => {
    const updatedComponent = {
      ...component,
      config: {
        ...component.config,
        name: newName,
      },
    };

    updateNode(nodeId, { component: updatedComponent });
  };

  return (
    <input
      value={component?.config.name || ''}
      onChange={(e) => handleNameChange(e.target.value)}
    />
  );
}

Related Pages

Studio_Builder_Nodes - React components for rendering graph nodes
Studio_Builder_Utils - Utility functions for graph layout and conversion
Studio_TeamBuilder_Sidebar - Sidebar for team management
Types_Datamodel - Core type definitions
Types_Guards - Type guard functions

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment