State Management

/**
 * Base class for agent state with map-based data storage
 * @param initData Initial state data as key-value map
 */
class AgentState {
    AgentState(Map<String, Object> initData);

    /**
     * Returns unmodifiable view of state data
     * @return Immutable map of current state
     */
    Map<String, Object> data();

    /**
     * Retrieves typed value by key
     * @param key State key to retrieve
     * @return Optional containing value if present
     */
    <T> Optional<T> value(String key);
}

// Define custom state class
class MyAgentState extends AgentState {
    public MyAgentState(Map<String, Object> initData) {
        super(initData);
    }
}

// Create and use state
Map<String, Object> initialData = Map.of(
    "messages", new ArrayList<String>(),
    "score", 0,
    "active", true
);

MyAgentState state = new MyAgentState(initialData);

// Access state values
Optional<Integer> score = state.value("score");
Optional<List<String>> messages = state.value("messages");
Optional<Boolean> active = state.value("active");

// Get raw data
Map<String, Object> allData = state.data();

/**
 * Updates state with partial state using channels for merge logic
 * @param state Current state map
 * @param partialState State updates to apply
 * @param channels Channel definitions for merge behavior
 * @return New state map with updates applied
 * @throws NullPointerException if state is null
 */
static Map<String, Object> updateState(Map<String, Object> state, Map<String, Object> partialState, Map<String, Channel<?>> channels);

/**
 * Updates AgentState with partial state using channels
 * @param state Current AgentState instance
 * @param partialState State updates to apply
 * @param channels Channel definitions for merge behavior
 * @return New state map with updates applied
 * @throws NullPointerException if state is null
 */
static Map<String, Object> updateState(AgentState state, Map<String, Object> partialState, Map<String, Channel<?>> channels);

// Simple state update
Map<String, Object> currentState = Map.of("counter", 5, "name", "Alice");
Map<String, Object> updates = Map.of("counter", 10, "active", true);
Map<String, Object> newState = AgentState.updateState(currentState, updates, Map.of());
// Result: {counter=10, name=Alice, active=true}

// Update with channels for merge behavior
Map<String, Channel<?>> channels = Map.of(
    "messages", Channels.appender(ArrayList::new),
    "counter", (key, current, new_val) -> (Integer) current + (Integer) new_val
);

Map<String, Object> stateWithMessages = Map.of(
    "messages", new ArrayList<>(List.of("Hello")),
    "counter", 5
);

Map<String, Object> updatesWithMessages = Map.of(
    "messages", List.of("World"),
    "counter", 3
);

Map<String, Object> merged = AgentState.updateState(stateWithMessages, updatesWithMessages, channels);
// messages: ["Hello", "World"], counter: 8

// Constants for state control
static final Object MARK_FOR_REMOVAL;  // Removes key from state
static final Object MARK_FOR_RESET;    // Resets key to null/default

// Remove state keys
Map<String, Object> updates = Map.of(
    "temporary_data", AgentState.MARK_FOR_REMOVAL,
    "cache", AgentState.MARK_FOR_RESET,
    "new_value", "updated"
);

Map<String, Object> cleaned = AgentState.updateState(currentState, updates, Map.of());
// temporary_data is completely removed, cache is set to null

/**
 * Interface defining state value update behavior
 * @param <T> Type of values handled by this channel
 */
interface Channel<T> {
    /**
     * Updates a state value using channel-specific merge logic
     * @param key State key being updated
     * @param currentValue Current value (may be null)
     * @param newValue New value to merge
     * @return Updated value after applying channel logic
     */
    Object update(String key, Object currentValue, Object newValue);
}

// Utility class for creating standard channels
class Channels {
    /**
     * Creates appender channel for collections
     * @param factory Supplier creating empty collection instances
     * @return Channel that appends new values to collections
     */
    static <C extends Collection<Object>> Channel<C> appender(Supplier<C> factory);
}

// Appender channel for list-like state values
class AppenderChannel<C extends Collection<Object>> implements Channel<C> {
    AppenderChannel(Supplier<C> factory);
    Object update(String key, Object currentValue, Object newValue);
}

import org.bsc.langgraph4j.state.Channels;

// Define schema with channels
Map<String, Channel<?>> schema = Map.of(
    // Append new messages to existing list
    "messages", Channels.appender(ArrayList::new),

    // Custom counter channel that adds values
    "counter", (key, current, newVal) -> {
        int currentInt = (current != null) ? (Integer) current : 0;
        int newInt = (Integer) newVal;
        return currentInt + newInt;
    },

    // Last-writer-wins channel (default behavior)
    "status", (key, current, newVal) -> newVal
);

// Use schema in StateGraph
StateGraph<MyAgentState> workflow = new StateGraph<>(schema, MyAgentState::new);

/**
 * Factory interface for creating agent state instances
 * @param <State> Type of agent state to create
 */
@FunctionalInterface
interface AgentStateFactory<State extends AgentState> extends Function<Map<String, Object>, State> {
    /**
     * Creates state instance from data map
     * @param data Initial state data
     * @return New agent state instance
     */
    State apply(Map<String, Object> data);

    /**
     * Creates initial data from schema channels
     * @param channels Channel schema defining state structure
     * @return Map with initial values based on schema
     */
    default Map<String, Object> initialDataFromSchema(Map<String, Channel<?>> channels) {
        // Default implementation returns empty map
        return Map.of();
    }
}

// Simple factory using constructor reference
AgentStateFactory<MyAgentState> factory = MyAgentState::new;

// Custom factory with default values
AgentStateFactory<MyAgentState> factoryWithDefaults = (data) -> {
    Map<String, Object> defaultData = Map.of(
        "messages", new ArrayList<String>(),
        "score", 0,
        "created", System.currentTimeMillis()
    );

    // Merge provided data with defaults
    Map<String, Object> merged = new HashMap<>(defaultData);
    merged.putAll(data);

    return new MyAgentState(merged);
};

// Factory with schema-based initialization
AgentStateFactory<MyAgentState> schemaFactory = new AgentStateFactory<MyAgentState>() {
    @Override
    public MyAgentState apply(Map<String, Object> data) {
        return new MyAgentState(data);
    }

    @Override
    public Map<String, Object> initialDataFromSchema(Map<String, Channel<?>> channels) {
        Map<String, Object> initial = new HashMap<>();

        // Initialize collections for appender channels
        for (Map.Entry<String, Channel<?>> entry : channels.entrySet()) {
            if (entry.getValue() instanceof AppenderChannel) {
                initial.put(entry.getKey(), new ArrayList<>());
            }
        }

        return initial;
    }
};

/**
 * Interface for state serialization and cloning
 * @param <State> Type of agent state
 */
interface StateSerializer<State extends AgentState> {
    /**
     * Get the state factory used by this serializer
     * @return State factory for creating instances
     */
    AgentStateFactory<State> stateFactory();

    /**
     * Clone state object for immutability and persistence
     * @param data State data to clone
     * @return Cloned state instance
     * @throws IOException, ClassNotFoundException, InstantiationException, IllegalAccessException
     */
    State cloneObject(Map<String, Object> data) throws IOException, ClassNotFoundException, InstantiationException, IllegalAccessException;
}

// Java object stream serialization
class ObjectStreamStateSerializer<State extends AgentState> implements StateSerializer<State> {
    ObjectStreamStateSerializer(AgentStateFactory<State> stateFactory);
    AgentStateFactory<State> stateFactory();
    State cloneObject(Map<String, Object> data) throws IOException, ClassNotFoundException;
}

// JSON-based serialization with Jackson
class JacksonStateSerializer<State extends AgentState> extends PlainTextStateSerializer<State> {
    JacksonStateSerializer(AgentStateFactory<State> stateFactory);
    JacksonStateSerializer(AgentStateFactory<State> stateFactory, ObjectMapper objectMapper);
}

// JSON-based serialization with Gson
class GsonStateSerializer<State extends AgentState> extends PlainTextStateSerializer<State> {
    GsonStateSerializer(AgentStateFactory<State> stateFactory);
    GsonStateSerializer(AgentStateFactory<State> stateFactory, Gson gson);
}

// Object stream serializer
StateSerializer<MyAgentState> objectSerializer =
    new ObjectStreamStateSerializer<>(MyAgentState::new);

// Jackson JSON serializer
ObjectMapper mapper = new ObjectMapper();
StateSerializer<MyAgentState> jacksonSerializer =
    new JacksonStateSerializer<>(MyAgentState::new, mapper);

// Gson JSON serializer
Gson gson = new GsonBuilder().create();
StateSerializer<MyAgentState> gsonSerializer =
    new GsonStateSerializer<>(MyAgentState::new, gson);

// Use in StateGraph
StateGraph<MyAgentState> workflow = new StateGraph<>(schema, jacksonSerializer);

/**
 * Immutable snapshot of graph state at specific execution point
 */
class StateSnapshot<State extends AgentState> implements NodeOutput<State> {
    /**
     * Creates state snapshot from checkpoint and configuration
     * @param checkpoint Checkpoint data
     * @param config Runtime configuration
     * @param stateFactory Factory for creating state instances
     * @return State snapshot instance
     */
    static <State extends AgentState> StateSnapshot<State> of(Checkpoint checkpoint, RunnableConfig config, AgentStateFactory<State> stateFactory);

    // NodeOutput interface methods
    String node();
    State state();

    // Snapshot-specific methods
    String getCheckpointId();
    String getNodeId();
    String getNextNodeId();
    RunnableConfig getConfig();
}

// Access snapshot information
StateSnapshot<MyAgentState> snapshot = app.getState(config);

System.out.println("Execution at node: " + snapshot.getNodeId());
System.out.println("Next node will be: " + snapshot.getNextNodeId());
System.out.println("Checkpoint ID: " + snapshot.getCheckpointId());

// Access state data
MyAgentState currentState = snapshot.state();
Map<String, Object> stateData = currentState.data();

// Get execution context
RunnableConfig executionConfig = snapshot.getConfig();

interface Reducer<T> {
    T reduce(T current, T newValue);
}

// Custom channel with reducer
Channel<Integer> sumChannel = (key, current, newVal) -> {
    int currentInt = (current != null) ? (Integer) current : 0;
    return currentInt + (Integer) newVal;
};

Channel<Set<String>> setUnionChannel = (key, current, newVal) -> {
    Set<String> currentSet = (current != null) ? (Set<String>) current : new HashSet<>();
    Set<String> result = new HashSet<>(currentSet);
    if (newVal instanceof Collection) {
        result.addAll((Collection<String>) newVal);
    } else {
        result.add((String) newVal);
    }
    return result;
};

// Conditional update channel
Channel<String> conditionalStatus = (key, current, newVal) -> {
    String newStatus = (String) newVal;
    String currentStatus = (String) current;

    // Only update if new status has higher priority
    Map<String, Integer> priorities = Map.of(
        "pending", 1,
        "processing", 2,
        "completed", 3,
        "failed", 4
    );

    int currentPriority = priorities.getOrDefault(currentStatus, 0);
    int newPriority = priorities.getOrDefault(newStatus, 0);

    return (newPriority >= currentPriority) ? newStatus : currentStatus;
};

// Validating channel
Channel<Integer> validatedCounter = (key, current, newVal) -> {
    Integer newValue = (Integer) newVal;
    if (newValue < 0) {
        throw new IllegalArgumentException("Counter cannot be negative");
    }
    return newValue;
};

// Range-bound channel
Channel<Double> boundedValue = (key, current, newVal) -> {
    Double value = (Double) newVal;
    return Math.max(0.0, Math.min(100.0, value)); // Clamp to [0, 100]
};