From 13847a32c0c7159d0e7b40ab207f3aa9c1465bac Mon Sep 17 00:00:00 2001 From: Quanzheng Long Date: Thu, 12 Mar 2026 22:41:32 -0700 Subject: [PATCH] refactor-go --- langgraph-go/advancedgraph/graph.go | 55 +++++- langgraph-go/tests/sub_agents_test.go | 230 +++++++++++++++----------- rust-core/src/engine.rs | 77 ++++++++- rust-core/src/lib_c.rs | 3 +- 4 files changed, 262 insertions(+), 103 deletions(-) diff --git a/langgraph-go/advancedgraph/graph.go b/langgraph-go/advancedgraph/graph.go index c38cd18d4..54a0113be 100644 --- a/langgraph-go/advancedgraph/graph.go +++ b/langgraph-go/advancedgraph/graph.go @@ -2,9 +2,12 @@ package advancedgraph import ( "fmt" + "reflect" + "runtime" + "strings" ) -type NodeFunc func(ctx *Context, arg any) (Command, error) +type NodeFunc func(ctx *Context, state map[string]any) (Command, error) type AdvancedStateGraph struct { nodes map[string]NodeFunc @@ -19,20 +22,25 @@ func NewAdvancedStateGraph() *AdvancedStateGraph { } } -func (g *AdvancedStateGraph) AddNode(name string, fn NodeFunc) { +func (g *AdvancedStateGraph) AddNode(fn NodeFunc) string { + name := NodeName(fn) + if _, exists := g.nodes[name]; exists { + panic(fmt.Sprintf("node `%s` already exists", name)) + } g.nodes[name] = fn + return name } func (g *AdvancedStateGraph) AddAsyncChannel(name string) { g.asyncChannels = append(g.asyncChannels, name) } -func (g *AdvancedStateGraph) SetEntryPoint(name string) { - g.entryPoint = name +func (g *AdvancedStateGraph) SetEntryNode(fn NodeFunc) { + g.entryPoint = NodeName(fn) } -func (g *AdvancedStateGraph) SetFinishPoint(name string) { - g.finishPoint = name +func (g *AdvancedStateGraph) SetFinishNode(fn NodeFunc) { + g.finishPoint = NodeName(fn) } func (g *AdvancedStateGraph) Compile() *CompiledGraph { @@ -100,12 +108,19 @@ func (g *CompiledGraph) Start(initialState map[string]any) (*Handler, error) { g.entryPoint, g.finishPoint, initialState, - func(node string, arg any, _ map[string]any) (Command, error) { + func(node string, arg any, fallbackState map[string]any) (Command, error) { fn, ok := g.nodes[node] if !ok { return Command{}, fmt.Errorf("unknown node `%s`", node) } - return fn(&Context{engine: engine}, arg) + stateArg, ok := arg.(map[string]any) + if !ok { + stateArg = fallbackState + } + if stateArg == nil { + return Command{}, fmt.Errorf("node `%s` expected map state argument", node) + } + return fn(&Context{engine: engine}, stateArg) }, ) handler.done <- resultOrErr{state: state, err: err} @@ -113,3 +128,27 @@ func (g *CompiledGraph) Start(initialState map[string]any) (*Handler, error) { }() return handler, nil } + +func NodeName(fn NodeFunc) string { + pc := reflect.ValueOf(fn).Pointer() + f := runtime.FuncForPC(pc) + if f == nil { + panic("cannot infer node name from nil function") + } + full := f.Name() + if strings.Contains(full, ".func") { + panic("anonymous functions are not allowed as nodes") + } + short := full + if i := strings.LastIndex(short, "/"); i >= 0 { + short = short[i+1:] + } + if i := strings.LastIndex(short, "."); i >= 0 { + short = short[i+1:] + } + short = strings.TrimSuffix(short, "-fm") + if short == "" || strings.Contains(short, "func") { + panic(fmt.Sprintf("cannot infer stable node name from `%s`", full)) + } + return short +} diff --git a/langgraph-go/tests/sub_agents_test.go b/langgraph-go/tests/sub_agents_test.go index badd06e22..58fe6fbdb 100644 --- a/langgraph-go/tests/sub_agents_test.go +++ b/langgraph-go/tests/sub_agents_test.go @@ -29,6 +29,128 @@ func (m *mockPlanner) invoke() []decision { return resp } +type lunchWorkflow struct { + planner *mockPlanner + names map[string]string +} + +func cloneState(state map[string]any) map[string]any { + out := make(map[string]any, len(state)+2) + for k, v := range state { + out[k] = v + } + out["output"] = append([]string(nil), outputSlice(state)...) + return out +} + +func outputSlice(state map[string]any) []string { + raw, ok := state["output"] + if !ok || raw == nil { + return []string{} + } + switch v := raw.(type) { + case []string: + return v + case []any: + out := make([]string, 0, len(v)) + for _, item := range v { + if s, ok := item.(string); ok { + out = append(out, s) + } + } + return out + default: + return []string{} + } +} + +func (w *lunchWorkflow) llmNode(ctx *ag.Context, state map[string]any) (ag.Command, error) { + decisions := w.planner.invoke() + sends := make([]ag.Send, 0, 4) + for _, d := range decisions { + if d.Type == "end" { + next := cloneState(state) + next["complete"] = d.Complete + return ag.Command{ + Goto: []ag.Send{ + {Node: w.names["order"], Arg: next}, + }, + }, nil + } + if d.Type == "sub_agent" { + next := cloneState(state) + next["sub_agent_input"] = d.SubAgent + sends = append(sends, ag.Send{Node: w.names["sub"], Arg: next}) + } + if d.Type == "tool" { + next := cloneState(state) + next["tool_input"] = d.Tool + sends = append(sends, ag.Send{Node: w.names["tool"], Arg: next}) + } + } + sends = append(sends, ag.Send{Node: w.names["wait"], Arg: cloneState(state)}) + return ag.Command{Goto: sends}, nil +} + +func (w *lunchWorkflow) waitNode(ctx *ag.Context, state map[string]any) (ag.Command, error) { + event, err := ctx.WaitFor( + ag.AnyOf( + ag.ChannelCondition{Channel: "tool_completion_channel", N: 1}, + ag.ChannelCondition{Channel: "subagent_completion_channel", N: 1}, + ag.ChannelCondition{Channel: "user_input_channel", N: 1}, + ag.TimerCondition{Seconds: 1}, + ), + ) + if err != nil { + return ag.Command{}, err + } + + output := outputSlice(state) + if event.Condition == "channel" { + payload := ag.DecodeString(event.Value) + switch event.Channel { + case "tool_completion_channel": + output = append(output, "tool: "+payload) + case "subagent_completion_channel": + output = append(output, "sub_agent: "+payload) + case "user_input_channel": + output = append(output, "user_input: "+payload) + } + state["output"] = output + return ag.Command{Goto: []ag.Send{{Node: w.names["llm"], Arg: cloneState(state)}}}, nil + } + + output = append(output, "timer: no updates yet") + state["output"] = output + return ag.Command{Goto: []ag.Send{{Node: w.names["wait"], Arg: cloneState(state)}}}, nil +} + +func (w *lunchWorkflow) toolNode(ctx *ag.Context, state map[string]any) (ag.Command, error) { + toolInput, _ := state["tool_input"].(string) + time.Sleep(100 * time.Millisecond) + err := ctx.PublishToChannel("tool_completion_channel", "tool completed for: "+toolInput) + return ag.Command{}, err +} + +func (w *lunchWorkflow) subAgentNode(ctx *ag.Context, state map[string]any) (ag.Command, error) { + subInput, _ := state["sub_agent_input"].(string) + time.Sleep(5 * time.Second) + err := ctx.PublishToChannel( + "subagent_completion_channel", + "research sub agent completed for: "+subInput, + ) + return ag.Command{}, err +} + +func (w *lunchWorkflow) orderFoodNode(ctx *ag.Context, state map[string]any) (ag.Command, error) { + complete, _ := state["complete"].(string) + output := outputSlice(state) + output = append(output, "order_food: "+complete) + state["output"] = output + state["done"] = complete + return ag.Command{Update: state}, nil +} + func TestSubAgentsEquivalentFlow(t *testing.T) { planner := &mockPlanner{ responses: [][]decision{ @@ -42,105 +164,24 @@ func TestSubAgentsEquivalentFlow(t *testing.T) { {{Type: "end", Complete: "order submitted"}}, }, } + workflow := &lunchWorkflow{ + planner: planner, + names: make(map[string]string), + } graph := ag.NewAdvancedStateGraph() graph.AddAsyncChannel("tool_completion_channel") graph.AddAsyncChannel("subagent_completion_channel") graph.AddAsyncChannel("user_input_channel") - graph.AddNode("llm_node", func(ctx *ag.Context, arg any) (ag.Command, error) { - state := arg.(map[string]any) - decisions := planner.invoke() - sends := make([]ag.Send, 0, 4) - for _, d := range decisions { - if d.Type == "end" { - return ag.Command{ - Goto: []ag.Send{ - { - Node: "order_food_node", - Arg: map[string]any{ - "state": state, - "complete": d.Complete, - }, - }, - }, - }, nil - } - if d.Type == "sub_agent" { - sends = append(sends, ag.Send{Node: "sub_agent_node", Arg: d.SubAgent}) - } - if d.Type == "tool" { - sends = append(sends, ag.Send{Node: "tool_node", Arg: d.Tool}) - } - } - sends = append(sends, ag.Send{Node: "wait_node", Arg: state}) - return ag.Command{Goto: sends}, nil - }) + workflow.names["llm"] = graph.AddNode(workflow.llmNode) + workflow.names["wait"] = graph.AddNode(workflow.waitNode) + workflow.names["tool"] = graph.AddNode(workflow.toolNode) + workflow.names["sub"] = graph.AddNode(workflow.subAgentNode) + workflow.names["order"] = graph.AddNode(workflow.orderFoodNode) - graph.AddNode("wait_node", func(ctx *ag.Context, arg any) (ag.Command, error) { - state := arg.(map[string]any) - event, err := ctx.WaitFor( - ag.AnyOf( - ag.ChannelCondition{Channel: "tool_completion_channel", N: 1}, - ag.ChannelCondition{Channel: "subagent_completion_channel", N: 1}, - ag.ChannelCondition{Channel: "user_input_channel", N: 1}, - ag.TimerCondition{Seconds: 1}, - ), - ) - if err != nil { - return ag.Command{}, err - } - - output := state["output"].([]string) - if event.Condition == "channel" { - payload := ag.DecodeString(event.Value) - switch event.Channel { - case "tool_completion_channel": - output = append(output, "tool: "+payload) - case "subagent_completion_channel": - output = append(output, "sub_agent: "+payload) - case "user_input_channel": - output = append(output, "user_input: "+payload) - } - state["output"] = output - return ag.Command{Goto: []ag.Send{{Node: "llm_node", Arg: state}}}, nil - } - - output = append(output, "timer: no updates yet") - state["output"] = output - return ag.Command{Goto: []ag.Send{{Node: "wait_node", Arg: state}}}, nil - }) - - graph.AddNode("tool_node", func(ctx *ag.Context, arg any) (ag.Command, error) { - toolInput := arg.(string) - time.Sleep(100 * time.Millisecond) - err := ctx.PublishToChannel("tool_completion_channel", "tool completed for: "+toolInput) - return ag.Command{}, err - }) - - graph.AddNode("sub_agent_node", func(ctx *ag.Context, arg any) (ag.Command, error) { - subInput := arg.(string) - time.Sleep(5 * time.Second) - err := ctx.PublishToChannel( - "subagent_completion_channel", - "research sub agent completed for: "+subInput, - ) - return ag.Command{}, err - }) - - graph.AddNode("order_food_node", func(ctx *ag.Context, arg any) (ag.Command, error) { - payload := arg.(map[string]any) - state := payload["state"].(map[string]any) - complete := payload["complete"].(string) - output := state["output"].([]string) - output = append(output, "order_food: "+complete) - state["output"] = output - state["done"] = complete - return ag.Command{Update: state}, nil - }) - - graph.SetEntryPoint("llm_node") - graph.SetFinishPoint("order_food_node") + graph.SetEntryNode(workflow.llmNode) + graph.SetFinishNode(workflow.orderFoodNode) handler, err := graph.Compile().Start( map[string]any{ @@ -163,7 +204,10 @@ func TestSubAgentsEquivalentFlow(t *testing.T) { t.Fatalf("result failed: %v", err) } - output := result["output"].([]string) + output := outputSlice(result) + if len(output) == 0 { + t.Fatalf("output is empty, full result=%#v", result) + } if result["done"] != "order submitted" { t.Fatalf("unexpected done: %v", result["done"]) } diff --git a/rust-core/src/engine.rs b/rust-core/src/engine.rs index c6314528e..6658df5af 100644 --- a/rust-core/src/engine.rs +++ b/rust-core/src/engine.rs @@ -2,6 +2,7 @@ use parking_lot::{Condvar, Mutex}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::{HashMap, VecDeque}; +use std::env; use std::sync::mpsc; use std::sync::Arc; use std::sync::Mutex as StdMutex; @@ -47,6 +48,28 @@ pub struct NodeExecResult { type Task = Box; +fn debug_enabled() -> bool { + static DEBUG: OnceLock = OnceLock::new(); + *DEBUG.get_or_init(|| { + matches!( + env::var("DEBUG") + .unwrap_or_default() + .trim() + .to_ascii_lowercase() + .as_str(), + "1" | "true" | "yes" | "on" + ) + }) +} + +fn debug_log(message: &str) { + if debug_enabled() { + let current = thread::current(); + let thread_name = current.name().unwrap_or("unnamed"); + println!("[advanced-graph][{thread_name}] {message}"); + } +} + struct ThreadPool { tx: mpsc::Sender, _workers: Vec>, @@ -85,6 +108,7 @@ impl ThreadPool { where F: FnOnce() + Send + 'static, { + debug_log("thread-pool execute() called"); self.tx .send(Box::new(task)) .map_err(|e| format!("thread-pool send failed: {e}")) @@ -95,6 +119,7 @@ pub fn node_pool_execute(task: F) -> Result<(), String> where F: FnOnce() + Send + 'static, { + debug_log("node_pool_execute() called"); static NODE_POOL: OnceLock = OnceLock::new(); let pool = NODE_POOL.get_or_init(|| { let size = thread::available_parallelism() @@ -109,6 +134,7 @@ pub fn run_loop_pool_execute(task: F) -> Result<(), String> where F: FnOnce() + Send + 'static, { + debug_log("run_loop_pool_execute() called"); static RUN_LOOP_POOL: OnceLock = OnceLock::new(); let pool = RUN_LOOP_POOL.get_or_init(|| ThreadPool::new(2, "langgraph-runloop")); pool.execute(task) @@ -126,42 +152,65 @@ where FSpawn: FnMut(String, A) -> Result<(), String>, FMerge: FnMut(&str, Option) -> Result<(), String>, { + debug_log("run_scheduler_loop() started"); let mut active: usize = 1; + debug_log("scheduling initial entry node"); spawn(entry_point, initial_arg)?; while active > 0 { + debug_log(&format!( + "scheduler waiting for node result (active={active})" + )); let item = rx .recv() .map_err(|e| format!("scheduler recv failed: {e}"))?; active = active.saturating_sub(1); let (node_name, node_result) = item.map_err(|e| format!("node execution failed: {e}"))?; + debug_log(&format!("scheduler received result from node={node_name}")); merge(&node_name, node_result.update)?; + debug_log(&format!("merged update from node={node_name}")); if node_name == finish_point { + debug_log("finish node reached, stopping scheduler loop"); break; } for send in node_result.sends { active += 1; + debug_log(&format!( + "scheduling next node={} (active={active})", + send.node + )); spawn(send.node, send.arg)?; } } + debug_log("run_scheduler_loop() finished"); Ok(()) } pub fn merge_json_update(state: &mut Value, update: Option) { + debug_log("merge_json_update() called"); let Some(update_value) = update else { + debug_log("merge_json_update(): no update payload"); return; }; match (&mut *state, update_value) { (Value::Object(state_obj), Value::Object(update_obj)) => { + debug_log(&format!( + "merge_json_update(): object merge with {} keys", + update_obj.len() + )); for (k, v) in update_obj { state_obj.insert(k, v); } } (Value::Object(state_obj), Value::Array(entries)) => { + debug_log(&format!( + "merge_json_update(): tuple-list merge with {} entries", + entries.len() + )); for entry in entries { if let Value::Array(pair) = entry { if pair.len() == 2 { @@ -172,7 +221,7 @@ pub fn merge_json_update(state: &mut Value, update: Option) { } } } - _ => {} + _ => debug_log("merge_json_update(): unsupported update shape, ignored"), } } @@ -184,15 +233,18 @@ pub struct Engine { impl Engine { pub fn new() -> Self { + debug_log("Engine::new()"); Self::default() } pub fn add_async_channel(&self, name: &str) { + debug_log(&format!("Engine::add_async_channel(name={name})")); let mut channels = self.channels.lock(); channels.entry(name.to_owned()).or_default(); } pub fn publish_json(&self, channel: &str, value: serde_json::Value) -> Result<(), String> { + debug_log(&format!("Engine::publish_json(channel={channel})")); let mut channels = self.channels.lock(); let queue = channels .get_mut(channel) @@ -204,6 +256,7 @@ impl Engine { } pub fn wait_for(&self, cond: &WaitCondition) -> Result { + debug_log(&format!("Engine::wait_for(cond={cond:?})")); match cond { WaitCondition::Channel { channel, n } => { if *n < 1 { @@ -215,6 +268,9 @@ impl Engine { .get_mut(channel) .ok_or_else(|| format!("Unknown channel `{channel}`"))?; if queue.len() >= *n { + debug_log(&format!( + "Engine::wait_for channel ready (channel={channel}, n={n})" + )); if *n == 1 { if let Some(value) = queue.pop_front() { return Ok(WaitEvent::Channel { @@ -235,6 +291,9 @@ impl Engine { }); } } + debug_log(&format!( + "Engine::wait_for waiting on channel condvar (channel={channel}, n={n})" + )); self.channel_notify.wait(&mut channels); } } @@ -243,12 +302,17 @@ impl Engine { return Err("timer condition must be > 0".to_string()); } std::thread::sleep(Duration::from_secs_f64(*seconds)); + debug_log(&format!("Engine::wait_for timer fired (seconds={seconds})")); Ok(WaitEvent::Timer { seconds: *seconds }) } } } pub fn wait_for_any_of(&self, any_of: &AnyOfCondition) -> Result { + debug_log(&format!( + "Engine::wait_for_any_of(conditions={})", + any_of.conditions.len() + )); if any_of.conditions.is_empty() { return Err("any_of requires at least one condition".to_string()); } @@ -275,6 +339,9 @@ impl Engine { .get_mut(channel) .ok_or_else(|| format!("Unknown channel `{channel}`"))?; if queue.len() >= *n { + debug_log(&format!( + "Engine::wait_for_any_of channel hit (channel={channel}, n={n})" + )); if *n == 1 { if let Some(value) = queue.pop_front() { return Ok(WaitEvent::Channel { @@ -302,11 +369,19 @@ impl Engine { let timeout = Duration::from_secs_f64(seconds); let elapsed = started.elapsed(); if elapsed >= timeout { + debug_log(&format!( + "Engine::wait_for_any_of timer hit (seconds={seconds})" + )); return Ok(WaitEvent::Timer { seconds }); } let remaining = timeout.saturating_sub(elapsed); + debug_log(&format!( + "Engine::wait_for_any_of waiting on condvar with timeout {:?}", + remaining + )); self.channel_notify.wait_for(&mut channels, remaining); } else { + debug_log("Engine::wait_for_any_of waiting on condvar without timeout"); self.channel_notify.wait(&mut channels); } } diff --git a/rust-core/src/lib_c.rs b/rust-core/src/lib_c.rs index c721402b0..9677d0b15 100644 --- a/rust-core/src/lib_c.rs +++ b/rust-core/src/lib_c.rs @@ -147,10 +147,11 @@ fn run_graph_scheduler_json( let tx_for_spawn = tx.clone(); let state_for_spawn = Arc::clone(&state); let state_for_merge = Arc::clone(&state); + let initial_arg = state.lock().expect("state mutex poisoned").clone(); run_scheduler_loop( entry_point, &finish_point, - state.lock().expect("state mutex poisoned").clone(), + initial_arg, move |node, arg| { let snapshot = state_for_spawn .lock()