From 612015b6d757e9928e789514b0c945cf681e2c98 Mon Sep 17 00:00:00 2001 From: chenchen Date: Sat, 11 Jul 2026 01:38:10 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=B7=A5=E4=BD=9C=E6=B5=81?= =?UTF-8?q?=E6=A3=80=E6=9F=A5=E7=82=B9=E5=BA=8F=E5=88=97=E5=8C=96=20(#193)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: ruanmingchen <“ruanm@chenchen”> --- internal/workflow/eino_compile_test.go | 85 ++++++++++++++++++++++++++ internal/workflow/state.go | 7 +++ 2 files changed, 92 insertions(+) diff --git a/internal/workflow/eino_compile_test.go b/internal/workflow/eino_compile_test.go index 06ec6edf..691bdbce 100644 --- a/internal/workflow/eino_compile_test.go +++ b/internal/workflow/eino_compile_test.go @@ -2,6 +2,7 @@ package workflow import ( "context" + "fmt" "path/filepath" "strings" "testing" @@ -9,6 +10,7 @@ import ( "cyberstrike-ai/internal/config" "cyberstrike-ai/internal/database" + "github.com/cloudwego/eino/compose" "go.uber.org/zap" ) @@ -235,6 +237,89 @@ func TestExecuteEinoGraph_linearStartOutput(t *testing.T) { } } +func TestExecuteEinoGraph_checkpointRestoresStartOutput(t *testing.T) { + ctx := context.Background() + checkpointStore, err := newFileCheckPointStore(t.TempDir()) + if err != nil { + t.Fatalf("new checkpoint store: %v", err) + } + state := newWorkflowLocalState(map[string]interface{}{"message": "ping"}, "run-checkpoint") + node := graphNode{ID: "start-1", Type: "start"} + wf := compose.NewWorkflow[WorkflowInput, WorkflowOutput]( + compose.WithGenLocalState(func(context.Context) *WorkflowLocalState { return state }), + ) + start := wf.AddLambdaNode("start-1", compose.InvokableLambda(func(_ context.Context, input WorkflowInput) (WorkflowNodeOutput, error) { + result := startOutputMap(node, input.Message, input.ConversationID, input.ProjectID) + state.NodeOutputs[node.ID] = result + state.NodeOutputs["condition-1"] = conditionOutputMap(graphNode{ID: "condition-1", Type: "condition"}, "{{inputs.message}} == ping", true) + state.NodeOutputs["tool-1"] = toolOutputMap(graphNode{ID: "tool-1", Type: "tool"}, "tool result", "lookup", map[string]any{"id": "1"}, "exec-1", false) + state.NodeOutputs["agent-1"] = agentOutputMap(graphNode{ID: "agent-1", Type: "agent"}, "agent result", "chat", []string{"exec-1"}) + state.NodeOutputs["hitl-1"] = hitlOutputMap(graphNode{ID: "hitl-1", Type: "hitl"}, "completed", "approved", "continue?", "reviewer", true) + state.NodeOutputs["output-1"] = outputNodeOutputMap(graphNode{ID: "output-1", Type: "output"}, "result", "ping") + state.NodeOutputs["end-1"] = endOutputMap(graphNode{ID: "end-1", Type: "end"}, "done") + state.LastOutput = result + state.Outputs["seed"] = "preserved" + return result, nil + })) + outputNode := wf.AddLambdaNode("out-1", compose.InvokableLambda(func(_ context.Context, input WorkflowNodeOutput) (WorkflowNodeOutput, error) { + return input, nil + })) + start.AddInput(compose.START) + outputNode.AddInput("start-1") + wf.End().AddInput("out-1", compose.ToField("out-1")) + runnable, err := wf.Compile(ctx, + compose.WithCheckPointStore(checkpointStore), + compose.WithInterruptAfterNodes([]string{"start-1"}), + ) + if err != nil { + t.Fatalf("compile: %v", err) + } + + _, err = runnable.Invoke(ctx, workflowInputFromMap(state.Inputs), compose.WithCheckPointID("run-checkpoint")) + info, ok := compose.ExtractInterruptInfo(err) + if !ok { + t.Fatalf("invoke error = %v, want checkpoint interrupt", err) + } + restored, ok := info.State.(*WorkflowLocalState) + if !ok { + t.Fatalf("checkpoint state = %T, want *WorkflowLocalState", info.State) + } + for nodeID, wantType := range map[string]string{ + "start-1": "StartOutput", + "condition-1": "ConditionOutput", + "tool-1": "ToolOutput", + "agent-1": "AgentOutput", + "hitl-1": "HITLOutput", + "output-1": "OutputNodeOutput", + "end-1": "NodeOutputEnvelope", + } { + if got := fmt.Sprintf("%T", restored.NodeOutputs[nodeID]["typed"]); got != "workflow."+wantType { + t.Fatalf("restored %s typed output = %s, want workflow.%s", nodeID, got, wantType) + } + } + if got := valueFromPath("previous.message", restored); got != "ping" { + t.Fatalf("restored previous.message = %v, want ping", got) + } + if got := valueFromPath("inputs.message", restored); got != "ping" { + t.Fatalf("restored inputs.message = %v, want ping", got) + } + if got := valueFromPath("outputs.seed", restored); got != "preserved" { + t.Fatalf("restored outputs.seed = %v, want preserved", got) + } + + result, err := runnable.Invoke(ctx, WorkflowInput{}, compose.WithCheckPointID("run-checkpoint")) + if err != nil { + t.Fatalf("resume checkpoint: %v", err) + } + output, ok := result["out-1"].(map[string]any) + if !ok { + t.Fatalf("resumed output type = %T, want map[string]any", result["out-1"]) + } + if got := output["output"]; got != "ping" { + t.Fatalf("resumed output = %v, want ping", got) + } +} + func TestExecuteEinoGraph_conditionBranch(t *testing.T) { ctx := context.Background() SetCheckpointDir(t.TempDir()) diff --git a/internal/workflow/state.go b/internal/workflow/state.go index 23f348f5..8b3a0376 100644 --- a/internal/workflow/state.go +++ b/internal/workflow/state.go @@ -11,6 +11,13 @@ import ( func init() { schema.RegisterName[*WorkflowLocalState]("_cyberstrike_workflow_local_state") + schema.RegisterName[NodeOutputEnvelope]("_cyberstrike_workflow_node_output_envelope") + schema.RegisterName[StartOutput]("_cyberstrike_workflow_start_output") + schema.RegisterName[ConditionOutput]("_cyberstrike_workflow_condition_output") + schema.RegisterName[ToolOutput]("_cyberstrike_workflow_tool_output") + schema.RegisterName[AgentOutput]("_cyberstrike_workflow_agent_output") + schema.RegisterName[HITLOutput]("_cyberstrike_workflow_hitl_output") + schema.RegisterName[OutputNodeOutput]("_cyberstrike_workflow_output_node_output") } // WorkflowLocalState is the Eino WithGenLocalState payload (checkpoint-serializable).