Proposal: Replace OverallState with the generic CustomState

[robocanic]

State 泛型化Proposal

提案背景

当前 spring-ai-alibaba-graph-core 中的状态管理完全依赖 OverAllState，它是一个基于 Map<String, Object>的弱类型容器。用户在节点中读写状态时必须使用字符串 key 和手动类型转换（如 state.value("messages", List.class)），这带来了以下问题：

1. 缺乏类型安全 ：编译期无法检测 key 拼写错误或类型不匹配，运行时才会暴露

2. 可读性差 ：每个节点无法知道某一个key是在哪被写入，在哪被读取

3. IDE 支持弱 ：无法利用自动补全、重构、跳转定义等 IDE 能力

用户期望能够定义一个普通的 Java 实体类（POJO）作为 Graph 的状态类型，通过字段声明和注解来描述状态结构和更新策略，同时保持与现有 OverAllState API 的完全向后兼容。

效果对比

改造前：

// 1. 定义状态策略
KeyStrategyFactory keyStrategyFactory = () -> Map.of(
    "messages", new AppendStrategy(),
    "currentStep", new ReplaceStrategy(),
    "input", new ReplaceStrategy()
);

// 2. 构建图
StateGraph stateGraph = new StateGraph(keyStrategyFactory);

// 3. 添加节点——必须用字符串 key 和手动类型转换
stateGraph.addNode("chatNode", (OverAllState state) -> {
    String input = state.value("input", "");
    List<Message> messages = state.value("messages", List.class).orElse(new ArrayList<>());
    
    // 调用 LLM...
    Message response = chatClient.call(input);
    messages.add(response);
    
    // 返回 Map 作为部分更新
    return Map.of(
        "messages", List.of(response),
        "currentStep", "completed"
    );
});

// 4. 添加边
stateGraph.addEdge(START, "chatNode");
stateGraph.addEdge("chatNode", END);

// 5. 编译并执行
CompiledGraph compiled = stateGraph.compile();
compiled.stream(Map.of("input", "你好")).subscribe(output -> {
    System.out.println(output);
});

改造后：

// 0. 接口仅做标记
public interface GraphState extends Serializable {
}

// 1. 定义状态类——POJO + 注解，一目了然
public class ChatState implements GraphState {
    
    @StateField(strategy = AppendStrategy.class, streaming = true)
    private List<Message> messages = new ArrayList<>();
    
    private String currentStep;  // 未标注，默认 ReplaceStrategy
    
    private String input;        // 未标注，默认 ReplaceStrategy
    
    transient Logger logger = LoggerFactory.getLogger(ChatState.class);  // 不参与序列化
    
    // getter/setter
    public List<Message> getMessages() { return messages; }
    public void setMessages(List<Message> messages) { this.messages = messages; }
    public String getCurrentStep() { return currentStep; }
    public void setCurrentStep(String step) { this.currentStep = step; }
    public String getInput() { return input; }
    public void setInput(String input) { this.input = input; }
}

// 2. 构建图——泛型参数指定状态类型，不再需要 KeyStrategyFactory
StateGraph<ChatState> stateGraph = new StateGraph<>(ChatState.class);

// 3. 添加节点——直接访问强类型字段，IDE 自动补全、重构、跳转全部可用
stateGraph.addNode("chatNode", (ChatState state) -> {
    String input = state.getInput();                    // 类型安全，IDE 自动补全
    List<Message> messages = state.getMessages();       // 无需类型转换
    
    // 调用 LLM...
    Message response = chatClient.call(input);
    state.getMessages().add(response);                  // 直接操作字段
    state.setCurrentStep("completed");
    
    return state;  // 返回强类型状态对象
});

// 4. 添加边（不变）
stateGraph.addEdge(START, "chatNode");
stateGraph.addEdge("chatNode", END);

// 5. 编译并执行
CompiledGraph<ChatState> compiled = stateGraph.compile();
compiled.stream(Map.of("input", "你好")).subscribe(output -> {
    System.out.println(output);
});

变更范围

GraphState

引入 GraphState 标记接口作为所有状态类型的基础契约，要求实现 Serializable，不包含任何方法。

public interface GraphState extends Serializable {
}

@StateField注解

引入 @StateField 可选注解

1. keyStrategy：允许用户在 POJO 字段上声明更新策略（KeyStrategy），未标注的字段默认使用 ReplaceStrategy，不想序列化的字段使用 Java 内置的 transient 关键字。

2. fieldName：指定序列化时使用的字段名，默认为空字符串（表示使用 Java 字段名）。当 POJO 字段名与 Map<String, Object> 中的 key 不一致时（例如与旧版 OverAllState 的 key 对齐），可通过此属性指定映射名。

3. streaming：标记该字段是否承载流式数据，默认为 false。当 NodeActionResult 携带 Flux<?> 时，框架通过此标记确定流式数据对应的状态字段 key。每个状态类中最多只能有一个字段标记 streaming = true。

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.FIELD)
public @interface StateField {
    Class<? extends KeyStrategy> strategy() default ReplaceStrategy.class;
    String fieldName() default "";
    boolean streaming() default false;
}

StateFieldScanner

引入 StateFieldScanner 工具类，通过反射扫描状态类的字段及 @StateField 注解，自动提取 KeyStrategy 映射

核心API泛型化

将 StateGraph、CompiledGraph、NodeAction、AsyncNodeAction、AsyncNodeActionWithConfig 等核心 API 泛型化，引入类型参数 <S extends GraphState>

@FunctionalInterface
public interface NodeAction<S extends GraphState> {
    NodeActionResult<S> apply(S state) throws Exception;
}

public class NodeActionResult<S extends GraphState> {
    private final S state;
    private final Flux<?> streamingFlux;  // 可选，为 null 表示非流式

    private NodeActionResult(S state, Flux<?> streamingFlux) {
        this.state = state;
        this.streamingFlux = streamingFlux;
    }

    // 非流式场景：仅返回状态
    public static <S extends GraphState> NodeActionResult<S> of(S state) {
        return new NodeActionResult<>(state, null);
    }

    // 流式场景：返回状态 + Flux
    public static <S extends GraphState> NodeActionResult<S> ofStreaming(S state, Flux<?> streamingFlux) {
        return new NodeActionResult<>(state, streamingFlux);
    }

    public S state() { return state; }
    public Flux<?> streamingFlux() { return streamingFlux; }
    public boolean hasStreamingFlux() { return streamingFlux != null; }
}

用户原有的NodeAction的兼容方式

提供 OverallStateNodeAction 函数式接口，保留旧版返回 Map<String, Object> 的 Action 签名，并通过内置的 wrap 静态方法将其包装为新版泛型 NodeAction<OverAllState>

@FunctionalInterface
public interface OverallStateNodeAction {
    Map<String, Object> apply(OverAllState state) throws Exception;
    
    // 便捷包装方法
    static NodeAction<OverAllState> wrap(OverallStateNodeAction nodeAction) {
        return new NodeAction<OverallState>(state -> {
            Map<String, Object> partialUpdate = nodeAction.apply(state);
            return NodeActionResult.of(new OverAllState(partialUpdate));
        });
    }
}

序列化

泛型化 StateSerializer，使其支持任意 GraphState 子类型的序列化

节点流式输出

原先用户如果需要一个节点里面返回流式输出，只需要在apply()方法中返回的Map value中使用Flux流就行。graph框架侧会判断类型是否为flux，如果是flux，则会走flux的流式逻辑。

@Override
	public Map<String, Object> apply(OverAllState state) {
		// Create prompt with user message
		UserMessage message = new UserMessage((String) state.value(OverAllState.DEFAULT_INPUT_KEY).get());
		Flux<ChatResponse> stream = chatModel.stream(new Prompt(message));
		return Map.of("messages", stream);
	}

State泛型改造后，NodeAction返回的结果为NodeActionResult, 框架侧通过判断streamingFlux==null来判断是否走流式的逻辑。

流式逻辑中除了将message 包装成StreamingOutput之外，还需要将flux流式输出的完整结果再更新到对应的字段里面。StateField注解中提供了**streaming=true**（约束：整个State中，只能有一个字段带上streaming = true）可以找到这个对应的字段名，进而将flux流式输出后的完整结果再更新到对应的字段值中。

public class MyState implements GraphState {
    @StateField(streaming = true)
    private String messages;
    
    private String currentStep;
    // getters/setters...
}

NodeAction<MyState> action = (MyState state) -> {
    Flux<ChatResponse> chatFlux = chatModel.stream(prompt);
    state.setCurrentStep("streaming");
    return NodeActionResult.ofStreaming(state, chatFlux);
};

其他

agent-framework

对用户暴露的API改造小，主要是各种自定义Node和Edge的改造

spring-ai-alibaba-starter-builtin-nodes(break change)

所有的nodes都需要适配新的NodeAction定义

[robocanic]

1. 流式输出支持任意类型