Compare commits

..
Author SHA1 Message Date
Quanzheng Long d26d4050c8 simp 2026-03-17 16:33:47 -07:00
Quanzheng Long 8f0f4a1be7 rm-po3 2026-03-17 16:23:08 -07:00
Quanzheng Long e13004da77 stream-draft 2026-03-17 16:12:16 -07:00
Quanzheng Long 5a9264d124 v2 2026-03-16 16:53:09 -07:00
Quanzheng Long efa86a6b14 fix 2026-03-16 16:52:05 -07:00
Quanzheng Long f9866186d9 makefile 2026-03-16 16:46:33 -07:00
Quanzheng Long bfdd7deb60 makefile 2026-03-16 16:41:14 -07:00
Quanzheng Long 09c8bb7c1c revert 2026-03-16 16:34:17 -07:00
Quanzheng Long 84d59adcd8 publish 2026-03-16 16:31:02 -07:00
Quanzheng Long 841ebf0c77 publish 2026-03-16 16:18:25 -07:00
35 changed files with 1527 additions and 853 deletions
+6
View File
@@ -0,0 +1,6 @@
PYTHON_VERSION ?= 3.13
.PHONY: run-basic
run-basic:
uv run --python $(PYTHON_VERSION) python examples/basic_run.py
+14
View File
@@ -0,0 +1,14 @@
# advanced-graph-examples
Examples for the published `saf-python-sdk` package.
## Quick start
```bash
cd advanced-graph-examples
uv sync
uv run python examples/basic_run.py
```
> Note: current published wheel is for Python 3.13 on macOS arm64.
@@ -0,0 +1,36 @@
from __future__ import annotations
import asyncio
from typing import TypedDict
from saf_python_sdk import Command, Send
from saf_python_sdk.advanced_graph import AdvancedStateGraph
class MyState(TypedDict):
count: int
logs: list[str]
async def start_node(state: MyState) -> Command:
state["logs"].append("start")
return Command(update=state, goto=Send("finish_node", "hello"))
async def finish_node(input: str, state: MyState) -> Command:
state["logs"].append(f"finish:{input}")
state["count"] += 1
return Command(update=state)
async def main() -> None:
graph = AdvancedStateGraph(MyState)
graph.add_entry_node(start_node)
graph.add_finish_node(finish_node)
result = await graph.compile().ainvoke({"count": 0, "logs": []})
print(result)
if __name__ == "__main__":
asyncio.run(main())
+12
View File
@@ -0,0 +1,12 @@
[project]
name = "advanced-graph-examples"
version = "0.1.0"
requires-python = ">=3.13,<3.14"
dependencies = [
"saf-python-sdk>=0.1.1",
]
[tool.uv]
dev-dependencies = []
+25
View File
@@ -0,0 +1,25 @@
version = 1
revision = 3
requires-python = "==3.13.*"
[[package]]
name = "advanced-graph-examples"
version = "0.1.0"
source = { virtual = "." }
dependencies = [
{ name = "saf-python-sdk" },
]
[package.metadata]
requires-dist = [{ name = "saf-python-sdk", specifier = ">=0.1.1" }]
[package.metadata.requires-dev]
dev = []
[[package]]
name = "saf-python-sdk"
version = "0.1.2"
source = { registry = "https://pypi.org/simple" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/c5/43/2a86cdc8a1fbb185b9ce4642f0e3676acc168d48ed2e68d4ead6ad57e95e/saf_python_sdk-0.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:9035ab6ede9884b5588925c47d5b8be2c1c107becf5d6f88e0ebb39f4e8fac3b", size = 517053, upload-time = "2026-03-16T23:49:25.661Z" },
]
+28 -1
View File
@@ -111,6 +111,10 @@ func (c *Context) PublishToChannel(channel string, value any) error {
return c.engine.Publish(channel, value)
}
func (c *Context) SendCustomStreamEvent(value any) error {
return c.engine.SendCustomStreamEvent(value)
}
type Handler[StateT any] struct {
engine *RustEngine
done chan resultOrErr[StateT]
@@ -130,7 +134,29 @@ func (h *Handler[StateT]) WaitForResult() (StateT, error) {
return res.state, res.err
}
func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT) (*Handler[StateT], error) {
func (h *Handler[StateT]) ReceiveStream() (any, error) {
event, hasEvent, err := h.engine.ReceiveStream()
if err != nil {
return nil, err
}
if !hasEvent {
return nil, nil
}
return event, nil
}
func (h *Handler[StateT]) CloseStream() error {
return h.engine.CloseStream()
}
func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT, streamMode ...string) (*Handler[StateT], error) {
resolvedStreamMode := ""
if len(streamMode) > 1 {
return nil, fmt.Errorf("start accepts at most one stream mode")
}
if len(streamMode) == 1 {
resolvedStreamMode = streamMode[0]
}
engine := NewRustEngine()
for _, ch := range g.asyncChannels {
if err := engine.AddAsyncChannel(ch); err != nil {
@@ -147,6 +173,7 @@ func (g *CompiledGraph[StateT]) Start(initialInput any, initialState StateT) (*H
rawState, err := engine.RunGraph(
g.entryPoint,
g.finishPoint,
resolvedStreamMode,
initialState,
initialInput,
func(node string, nodeInput any, fallbackState map[string]any) (Command, error) {
+62
View File
@@ -128,6 +128,59 @@ func (e *RustEngine) AddAsyncChannel(channel string) error {
return parseRustStatus(resp)
}
func (e *RustEngine) StartStream(streamMode string) error {
var cmode *C.char
if streamMode != "" {
cmode = C.CString(streamMode)
defer C.free(unsafe.Pointer(cmode))
}
resp := C.rc_start_stream(e.ptr, cmode)
return parseRustStatus(resp)
}
func (e *RustEngine) ReceiveStream() (any, bool, error) {
resp := C.rc_receive_stream_json(e.ptr)
defer C.rc_string_free(resp)
raw := C.GoString(resp)
var status struct {
OK bool `json:"ok"`
Error string `json:"error"`
HasEvent bool `json:"has_event"`
Event json.RawMessage `json:"event"`
}
if err := json.Unmarshal([]byte(raw), &status); err != nil {
return nil, false, fmt.Errorf("decode rust stream response: %w", err)
}
if !status.OK {
return nil, false, fmt.Errorf("rust stream failed: %s", status.Error)
}
if !status.HasEvent {
return nil, false, nil
}
var event any
if err := json.Unmarshal(status.Event, &event); err != nil {
return nil, false, fmt.Errorf("decode stream event: %w", err)
}
return coerceJSONValue(event), true, nil
}
func (e *RustEngine) SendCustomStreamEvent(value any) error {
payload, err := json.Marshal(value)
if err != nil {
return fmt.Errorf("marshal stream event: %w", err)
}
cval := C.CString(string(payload))
defer C.free(unsafe.Pointer(cval))
resp := C.rc_send_custom_stream_event(e.ptr, cval)
return parseRustStatus(resp)
}
func (e *RustEngine) CloseStream() error {
resp := C.rc_close_stream(e.ptr)
return parseRustStatus(resp)
}
func (e *RustEngine) Publish(channel string, value any) error {
payload, err := json.Marshal(value)
if err != nil {
@@ -173,6 +226,7 @@ func (e *RustEngine) WaitAnyOf(cond AnyOfCondition) (WaitEvent, error) {
func (e *RustEngine) RunGraph(
entryPoint string,
finishPoint string,
streamMode string,
initialState any,
initialInput any,
exec func(node string, nodeInput any, state map[string]any) (Command, error),
@@ -189,10 +243,17 @@ func (e *RustEngine) RunGraph(
cfinish := C.CString(finishPoint)
cinitial := C.CString(string(initialJSON))
cinitialInput := C.CString(string(initialInputJSON))
var cstreamMode *C.char
if streamMode != "" {
cstreamMode = C.CString(streamMode)
}
defer C.free(unsafe.Pointer(centry))
defer C.free(unsafe.Pointer(cfinish))
defer C.free(unsafe.Pointer(cinitial))
defer C.free(unsafe.Pointer(cinitialInput))
if cstreamMode != nil {
defer C.free(unsafe.Pointer(cstreamMode))
}
callbackID := registerRunGraphCallbackCtx(&runGraphCallbackCtx{exec: exec})
defer unregisterRunGraphCallbackCtx(callbackID)
@@ -203,6 +264,7 @@ func (e *RustEngine) RunGraph(
cfinish,
cinitial,
cinitialInput,
cstreamMode,
C.ulong(callbackID),
(C.rc_node_callback_t)(C.goNodeCallback),
)
+102
View File
@@ -0,0 +1,102 @@
package tests
import (
"strings"
"testing"
"time"
ag "github.com/langchain-ai/langgraph/langgraph-go/advancedgraph"
)
type streamState struct {
Done bool `json:"done"`
}
type streamWorkflow struct{}
func (w *streamWorkflow) startNode(ctx *ag.Context, _ any, state streamState) (ag.Command, error) {
if err := ctx.SendCustomStreamEvent(map[string]any{"step": "start", "value": 1}); err != nil {
return ag.Command{}, err
}
time.Sleep(80 * time.Millisecond)
if err := ctx.SendCustomStreamEvent(map[string]any{"step": "start", "value": 2}); err != nil {
return ag.Command{}, err
}
return ag.Command{
Update: state,
Goto: []ag.Send{
{Node: w.finishNode},
},
}, nil
}
func (w *streamWorkflow) finishNode(ctx *ag.Context, _ any, state streamState) (ag.Command, error) {
state.Done = true
return ag.Command{Update: state}, nil
}
func TestCustomStreamReceiveAndClose(t *testing.T) {
workflow := &streamWorkflow{}
graph := ag.NewAdvancedStateGraph[streamState]()
graph.AddEntryNode(workflow.startNode)
graph.AddFinishNode(workflow.finishNode)
handler, err := graph.Compile().Start(nil, streamState{Done: false}, "custom")
if err != nil {
t.Fatalf("start failed: %v", err)
}
event, err := handler.ReceiveStream()
if err != nil {
t.Fatalf("receive stream failed: %v", err)
}
if event == nil {
t.Fatalf("expected first stream event, got nil")
}
eventMap, ok := event.(map[string]any)
if !ok {
t.Fatalf("unexpected event type: %T", event)
}
if eventMap["step"] != "start" {
t.Fatalf("unexpected stream event payload: %#v", eventMap)
}
if err := handler.CloseStream(); err != nil {
t.Fatalf("close stream failed: %v", err)
}
closedEvent, err := handler.ReceiveStream()
if err != nil {
t.Fatalf("receive stream after close failed: %v", err)
}
if closedEvent != nil {
t.Fatalf("expected nil stream event after close, got %#v", closedEvent)
}
result, err := handler.WaitForResult()
if err != nil {
t.Fatalf("result failed: %v", err)
}
if !result.Done {
t.Fatalf("expected final state done=true, got %#v", result)
}
}
func TestOnlyCustomStreamModeSupported(t *testing.T) {
workflow := &streamWorkflow{}
graph := ag.NewAdvancedStateGraph[streamState]()
graph.AddEntryNode(workflow.startNode)
graph.AddFinishNode(workflow.finishNode)
handler, err := graph.Compile().Start(nil, streamState{Done: false}, "values")
if err != nil {
t.Fatalf("start failed: %v", err)
}
_, runErr := handler.WaitForResult()
if runErr == nil {
t.Fatalf("expected run error for unsupported stream mode")
}
if !strings.Contains(runErr.Error(), "only `custom` is supported") {
t.Fatalf("unexpected error: %v", runErr)
}
}
-124
View File
@@ -2,12 +2,6 @@
# It is not intended for manual editing.
version = 4
[[package]]
name = "autocfg"
version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8"
[[package]]
name = "bitflags"
version = "2.11.0"
@@ -20,21 +14,6 @@ version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
[[package]]
name = "heck"
version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
[[package]]
name = "indoc"
version = "2.0.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706"
dependencies = [
"rustversion",
]
[[package]]
name = "itoa"
version = "1.0.17"
@@ -47,7 +26,6 @@ version = "0.1.0"
dependencies = [
"libc",
"parking_lot",
"pyo3",
"serde",
"serde_json",
"tokio",
@@ -74,21 +52,6 @@ version = "2.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
[[package]]
name = "memoffset"
version = "0.9.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a"
dependencies = [
"autocfg",
]
[[package]]
name = "once_cell"
version = "1.21.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
[[package]]
name = "parking_lot"
version = "0.12.5"
@@ -118,12 +81,6 @@ version = "0.2.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
[[package]]
name = "portable-atomic"
version = "1.13.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49"
[[package]]
name = "proc-macro2"
version = "1.0.106"
@@ -133,69 +90,6 @@ dependencies = [
"unicode-ident",
]
[[package]]
name = "pyo3"
version = "0.23.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7778bffd85cf38175ac1f545509665d0b9b92a198ca7941f131f85f7a4f9a872"
dependencies = [
"cfg-if",
"indoc",
"libc",
"memoffset",
"once_cell",
"portable-atomic",
"pyo3-build-config",
"pyo3-ffi",
"pyo3-macros",
"unindent",
]
[[package]]
name = "pyo3-build-config"
version = "0.23.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94f6cbe86ef3bf18998d9df6e0f3fc1050a8c5efa409bf712e661a4366e010fb"
dependencies = [
"once_cell",
"target-lexicon",
]
[[package]]
name = "pyo3-ffi"
version = "0.23.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e9f1b4c431c0bb1c8fb0a338709859eed0d030ff6daa34368d3b152a63dfdd8d"
dependencies = [
"libc",
"pyo3-build-config",
]
[[package]]
name = "pyo3-macros"
version = "0.23.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fbc2201328f63c4710f68abdf653c89d8dbc2858b88c5d88b0ff38a75288a9da"
dependencies = [
"proc-macro2",
"pyo3-macros-backend",
"quote",
"syn",
]
[[package]]
name = "pyo3-macros-backend"
version = "0.23.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fca6726ad0f3da9c9de093d6f116a93c1a38e417ed73bf138472cf4064f72028"
dependencies = [
"heck",
"proc-macro2",
"pyo3-build-config",
"quote",
"syn",
]
[[package]]
name = "quote"
version = "1.0.45"
@@ -214,12 +108,6 @@ dependencies = [
"bitflags",
]
[[package]]
name = "rustversion"
version = "1.0.22"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d"
[[package]]
name = "scopeguard"
version = "1.2.0"
@@ -286,12 +174,6 @@ dependencies = [
"unicode-ident",
]
[[package]]
name = "target-lexicon"
version = "0.12.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1"
[[package]]
name = "tokio"
version = "1.50.0"
@@ -319,12 +201,6 @@ version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "unindent"
version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3"
[[package]]
name = "windows-link"
version = "0.2.1"
-5
View File
@@ -7,12 +7,7 @@ edition = "2021"
name = "langgraph_rust_core"
crate-type = ["cdylib", "rlib"]
[features]
default = ["python-bindings"]
python-bindings = ["dep:pyo3"]
[dependencies]
pyo3 = { version = "0.23.5", features = ["extension-module"], optional = true }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
parking_lot = "0.12"
+5
View File
@@ -19,12 +19,17 @@ void rc_engine_free(Engine* ptr);
char* rc_add_async_channel(Engine* ptr, const char* channel);
char* rc_publish_json(Engine* ptr, const char* channel, const char* value_json);
char* rc_wait_any_of_json(Engine* ptr, const char* any_of_json);
char* rc_start_stream(Engine* ptr, const char* stream_mode);
char* rc_receive_stream_json(Engine* ptr);
char* rc_send_custom_stream_event(Engine* ptr, const char* value_json);
char* rc_close_stream(Engine* ptr);
char* rc_run_graph_json(
Engine* ptr,
const char* entry_point,
const char* finish_point,
const char* initial_state_json,
const char* initial_input_json,
const char* stream_mode,
unsigned long user_data,
rc_node_callback_t callback
);
+281 -7
View File
@@ -10,7 +10,7 @@ use std::sync::OnceLock;
use std::thread;
use std::time::{Duration, Instant};
use tokio::runtime::Runtime;
use tokio::sync::Notify;
use tokio::sync::{mpsc as tokio_mpsc, Mutex as AsyncMutex, Notify};
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "kind")]
@@ -62,6 +62,31 @@ pub enum NodeOutcome<U, A> {
Suspended { wait: WaitRequest },
}
#[derive(Debug, Deserialize)]
struct CallbackSendPayloadJson {
node: String,
#[serde(default)]
arg: Value,
}
#[derive(Debug, Deserialize)]
struct CallbackNodeExecResultJsonWire {
update: Option<Value>,
#[serde(default)]
sends: Vec<CallbackSendPayloadJson>,
}
#[derive(Debug, Deserialize)]
struct CallbackEnvelopeIn {
ok: bool,
#[serde(default)]
payload: Option<CallbackNodeExecResultJsonWire>,
#[serde(default)]
suspend: Option<WaitRequest>,
#[serde(default)]
error: Option<String>,
}
type Task = Box<dyn FnOnce() + Send + 'static>;
fn debug_enabled() -> bool {
@@ -279,10 +304,33 @@ pub fn merge_json_update(state: &mut Value, update: Option<Value>) {
}
}
enum SchedulerEventJson {
Node(Result<NodeExecutionJson, String>),
Resume {
node: String,
arg: Value,
event: WaitEvent,
},
WaitError(String),
}
struct NodeExecutionJson {
node: String,
arg: Value,
outcome: NodeOutcome<Value, Value>,
}
#[derive(Clone, Default)]
pub struct Engine {
channels: Arc<StdMutex<HashMap<String, VecDeque<serde_json::Value>>>>,
channel_notify: Arc<Notify>,
stream: Arc<StdMutex<Option<StreamChannel>>>,
}
#[derive(Clone)]
struct StreamChannel {
sender: tokio_mpsc::UnboundedSender<serde_json::Value>,
receiver: Arc<AsyncMutex<tokio_mpsc::UnboundedReceiver<serde_json::Value>>>,
}
impl Engine {
@@ -297,6 +345,60 @@ impl Engine {
channels.entry(name.to_owned()).or_default();
}
pub fn start_stream(&self, stream_mode: Option<&str>) -> Result<(), String> {
let mode = stream_mode.and_then(|m| {
let trimmed = m.trim();
if trimmed.is_empty() {
None
} else {
Some(trimmed)
}
});
let mut stream = self.stream.lock().expect("stream mutex poisoned");
if let Some(mode) = mode {
if mode != "custom" {
return Err(format!(
"unsupported stream_mode `{mode}`, only `custom` is supported"
));
}
let (sender, receiver) = tokio_mpsc::unbounded_channel();
*stream = Some(StreamChannel {
sender,
receiver: Arc::new(AsyncMutex::new(receiver)),
});
} else {
*stream = None;
}
Ok(())
}
pub fn close_stream(&self) {
let mut stream = self.stream.lock().expect("stream mutex poisoned");
*stream = None;
}
pub fn send_custom_stream_event(&self, value: serde_json::Value) {
let sender = {
let stream = self.stream.lock().expect("stream mutex poisoned");
stream.as_ref().map(|s| s.sender.clone())
};
if let Some(tx) = sender {
let _ = tx.send(value);
}
}
pub async fn receive_stream_async(&self) -> Option<serde_json::Value> {
let receiver = {
let stream = self.stream.lock().expect("stream mutex poisoned");
stream.as_ref().map(|s| Arc::clone(&s.receiver))
};
let Some(receiver) = receiver else {
return None;
};
let mut guard = receiver.lock().await;
guard.recv().await
}
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().expect("channels mutex poisoned");
@@ -339,7 +441,10 @@ impl Engine {
}
}
pub async fn wait_for_any_of_async(&self, any_of: &AnyOfCondition) -> Result<WaitEvent, String> {
pub async fn wait_for_any_of_async(
&self,
any_of: &AnyOfCondition,
) -> Result<WaitEvent, String> {
debug_log(&format!(
"Engine::wait_for_any_of_async(conditions={})",
any_of.conditions.len()
@@ -398,11 +503,7 @@ impl Engine {
run_loop_block_on(self.wait_for_any_of_async(any_of))
}
fn try_take_channel_event(
&self,
channel: &str,
n: usize,
) -> Result<Option<WaitEvent>, String> {
fn try_take_channel_event(&self, channel: &str, n: usize) -> Result<Option<WaitEvent>, String> {
let mut channels = self.channels.lock().expect("channels mutex poisoned");
let queue = channels
.get_mut(channel)
@@ -431,3 +532,176 @@ impl Engine {
}))
}
}
fn spawn_json_node_task<F>(
node: String,
arg: Value,
state_snapshot: Value,
tx: tokio_mpsc::UnboundedSender<SchedulerEventJson>,
callback: Arc<F>,
) -> Result<(), String>
where
F: Fn(String, Value, Value) -> Result<NodeOutcome<Value, Value>, String>
+ Send
+ Sync
+ 'static,
{
node_pool_execute(move || {
let node_for_result = node.clone();
let arg_for_result = arg.clone();
let result = callback(node.clone(), arg, state_snapshot).map(|payload| NodeExecutionJson {
node: node_for_result,
arg: arg_for_result,
outcome: payload,
});
let _ = tx.send(SchedulerEventJson::Node(result));
})
}
pub async fn run_graph_json_with_callback<F>(
entry_point: String,
finish_point: String,
initial_state: Value,
initial_input: Value,
engine: Engine,
callback: F,
) -> Result<Value, String>
where
F: Fn(String, Value, Value) -> Result<NodeOutcome<Value, Value>, String>
+ Send
+ Sync
+ 'static,
{
let callback = Arc::new(callback);
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<SchedulerEventJson>();
let state = Arc::new(StdMutex::new(initial_state));
let tx_for_spawn = tx.clone();
let state_for_spawn = Arc::clone(&state);
let state_for_merge = Arc::clone(&state);
let mut active: usize = 1;
let mut waiting: usize = 0;
spawn_json_node_task(
entry_point,
initial_input,
state_for_spawn
.lock()
.expect("state mutex poisoned")
.clone(),
tx_for_spawn.clone(),
Arc::clone(&callback),
)?;
while active > 0 || waiting > 0 {
let evt = rx
.recv()
.await
.ok_or_else(|| "scheduler event channel closed".to_string())?;
match evt {
SchedulerEventJson::Node(result) => {
active = active.saturating_sub(1);
let exec = result?;
match exec.outcome {
NodeOutcome::Completed(node_result) => {
let mut guard = state_for_merge.lock().expect("state mutex poisoned");
merge_json_update(&mut guard, node_result.update);
drop(guard);
if exec.node == finish_point {
break;
}
for send in node_result.sends {
active += 1;
let snapshot = state_for_spawn
.lock()
.expect("state mutex poisoned")
.clone();
spawn_json_node_task(
send.node,
send.arg,
snapshot,
tx_for_spawn.clone(),
Arc::clone(&callback),
)?;
}
}
NodeOutcome::Suspended { wait } => {
waiting += 1;
let tx_wait = tx_for_spawn.clone();
let node = exec.node;
let arg = exec.arg;
let engine_for_wait = engine.clone();
tokio::spawn(async move {
match engine_for_wait.wait_request_async(&wait).await {
Ok(event) => {
let _ = tx_wait.send(SchedulerEventJson::Resume {
node,
arg,
event,
});
}
Err(e) => {
let _ = tx_wait.send(SchedulerEventJson::WaitError(e));
}
}
});
}
}
}
SchedulerEventJson::Resume { node, arg, event } => {
waiting = waiting.saturating_sub(1);
active += 1;
let snapshot = state_for_spawn
.lock()
.expect("state mutex poisoned")
.clone();
spawn_json_node_task(
node,
wrap_resume_arg(arg, event),
snapshot,
tx_for_spawn.clone(),
Arc::clone(&callback),
)?;
}
SchedulerEventJson::WaitError(e) => return Err(e),
}
}
let final_state = state.lock().expect("state mutex poisoned").clone();
Ok(final_state)
}
fn wrap_resume_arg(arg: Value, event: WaitEvent) -> Value {
serde_json::json!({
"__lg_resume_arg__": arg,
"__lg_resume_event__": event,
})
}
pub fn parse_callback_envelope_json(
raw: &str,
node_name: &str,
) -> Result<NodeOutcome<Value, Value>, String> {
let parsed: CallbackEnvelopeIn = serde_json::from_str(raw)
.map_err(|e| format!("decode callback envelope for `{node_name}` failed: {e}"))?;
if !parsed.ok {
return Err(parsed
.error
.unwrap_or_else(|| format!("callback reported error for `{node_name}`")));
}
if let Some(wait) = parsed.suspend {
return Ok(NodeOutcome::Suspended { wait });
}
let payload = parsed
.payload
.ok_or_else(|| format!("callback payload missing for `{node_name}`"))?;
let sends = payload
.sends
.into_iter()
.map(|s| SendPayload {
node: s.node,
arg: s.arg,
})
.collect();
Ok(NodeOutcome::Completed(NodeExecResult {
update: payload.update,
sends,
}))
}
-2
View File
@@ -1,4 +1,2 @@
mod engine;
mod lib_c;
#[cfg(feature = "python-bindings")]
mod lib_py;
+133 -245
View File
@@ -1,39 +1,11 @@
use crate::engine::{
merge_json_update, node_pool_execute, run_loop_block_on, run_loop_spawn, AnyOfCondition,
Engine, NodeExecResult, NodeOutcome, SendPayload, WaitEvent, WaitRequest,
parse_callback_envelope_json, run_graph_json_with_callback, run_loop_block_on, run_loop_spawn,
AnyOfCondition, Engine, NodeOutcome,
};
use serde::Deserialize;
use serde_json::Value;
use std::ffi::{CStr, CString};
use std::os::raw::c_char;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use tokio::sync::mpsc as tokio_mpsc;
#[derive(Debug, Deserialize)]
struct SendPayloadJson {
node: String,
#[serde(default)]
arg: Value,
}
#[derive(Debug, Deserialize)]
struct NodeExecResultJsonWire {
update: Option<Value>,
#[serde(default)]
sends: Vec<SendPayloadJson>,
}
#[derive(Debug, Deserialize)]
struct CallbackEnvelopeIn {
ok: bool,
#[serde(default)]
payload: Option<NodeExecResultJsonWire>,
#[serde(default)]
suspend: Option<WaitRequest>,
#[serde(default)]
error: Option<String>,
}
type CNodeCallback = unsafe extern "C" fn(
user_data: libc::c_ulong,
@@ -42,9 +14,6 @@ type CNodeCallback = unsafe extern "C" fn(
state_json: *mut c_char,
) -> *mut c_char;
#[derive(Clone, Copy)]
struct CUserData(libc::c_ulong);
fn cstr_to_str<'a>(ptr: *const c_char) -> Result<&'a str, String> {
if ptr.is_null() {
return Err("Received null pointer".to_string());
@@ -63,213 +32,6 @@ fn into_c_ptr(s: String) -> *mut c_char {
}
}
fn parse_c_callback_result(
raw: String,
node_name: &str,
) -> Result<NodeOutcome<Value, Value>, String> {
let parsed: CallbackEnvelopeIn = serde_json::from_str(&raw)
.map_err(|e| format!("decode callback envelope for `{node_name}` failed: {e}"))?;
if !parsed.ok {
return Err(parsed
.error
.unwrap_or_else(|| format!("callback reported error for `{node_name}`")));
}
if let Some(wait) = parsed.suspend {
return Ok(NodeOutcome::Suspended { wait });
}
let payload = parsed
.payload
.ok_or_else(|| format!("callback payload missing for `{node_name}`"))?;
let sends = payload
.sends
.into_iter()
.map(|s| SendPayload {
node: s.node,
arg: s.arg,
})
.collect();
Ok(NodeOutcome::Completed(NodeExecResult {
update: payload.update,
sends,
}))
}
enum SchedulerEventJson {
Node(Result<NodeExecutionJson, String>),
Resume {
node: String,
arg: Value,
event: WaitEvent,
},
WaitError(String),
}
struct NodeExecutionJson {
node: String,
arg: Value,
outcome: NodeOutcome<Value, Value>,
}
fn spawn_json_node_task(
node: String,
arg: Value,
state_snapshot: Value,
tx: tokio_mpsc::UnboundedSender<SchedulerEventJson>,
user_data_bits: libc::c_ulong,
callback: CNodeCallback,
) -> Result<(), String> {
node_pool_execute(move || {
let node_for_result = node.clone();
let arg_for_result = arg.clone();
let result = (|| -> Result<NodeExecutionJson, String> {
let node_c =
CString::new(node.clone()).map_err(|e| format!("invalid node name: {e}"))?;
let arg_json = serde_json::to_string(&arg)
.map_err(|e| format!("serialize arg for `{node}` failed: {e}"))?;
let state_json = serde_json::to_string(&state_snapshot)
.map_err(|e| format!("serialize state for `{node}` failed: {e}"))?;
let arg_c =
CString::new(arg_json).map_err(|e| format!("invalid arg JSON bytes: {e}"))?;
let state_c =
CString::new(state_json).map_err(|e| format!("invalid state JSON bytes: {e}"))?;
let out_ptr = unsafe {
callback(
user_data_bits,
node_c.as_ptr() as *mut c_char,
arg_c.as_ptr() as *mut c_char,
state_c.as_ptr() as *mut c_char,
)
};
if out_ptr.is_null() {
return Err(format!("callback returned null for `{node}`"));
}
let out_raw = unsafe { CStr::from_ptr(out_ptr) }
.to_string_lossy()
.into_owned();
unsafe {
libc::free(out_ptr.cast());
}
let payload = parse_c_callback_result(out_raw, &node)?;
Ok(NodeExecutionJson {
node: node_for_result,
arg: arg_for_result,
outcome: payload,
})
})();
let _ = tx.send(SchedulerEventJson::Node(result));
})
}
async fn run_graph_scheduler_json(
entry_point: String,
finish_point: String,
initial_state: Value,
initial_input: Value,
engine: Engine,
user_data: CUserData,
callback: CNodeCallback,
) -> Result<Value, String> {
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<SchedulerEventJson>();
let state = Arc::new(Mutex::new(initial_state));
let user_data_bits = user_data.0;
let tx_for_spawn = tx.clone();
let state_for_spawn = Arc::clone(&state);
let state_for_merge = Arc::clone(&state);
let mut active: usize = 1;
let mut waiting: usize = 0;
spawn_json_node_task(
entry_point,
initial_input,
state_for_spawn
.lock()
.expect("state mutex poisoned")
.clone(),
tx_for_spawn.clone(),
user_data_bits,
callback,
)?;
while active > 0 || waiting > 0 {
let evt = rx
.recv()
.await
.ok_or_else(|| "scheduler event channel closed".to_string())?;
match evt {
SchedulerEventJson::Node(result) => {
active = active.saturating_sub(1);
let exec = result?;
match exec.outcome {
NodeOutcome::Completed(node_result) => {
let mut guard = state_for_merge.lock().expect("state mutex poisoned");
merge_json_update(&mut guard, node_result.update);
drop(guard);
if exec.node == finish_point {
break;
}
for send in node_result.sends {
active += 1;
let snapshot = state_for_spawn
.lock()
.expect("state mutex poisoned")
.clone();
spawn_json_node_task(
send.node,
send.arg,
snapshot,
tx_for_spawn.clone(),
user_data_bits,
callback,
)?;
}
}
NodeOutcome::Suspended { wait } => {
waiting += 1;
let tx_wait = tx_for_spawn.clone();
let node = exec.node;
let arg = exec.arg;
let engine_for_wait = engine.clone();
tokio::spawn(async move {
match engine_for_wait.wait_request_async(&wait).await {
Ok(event) => {
let _ = tx_wait.send(SchedulerEventJson::Resume { node, arg, event });
}
Err(e) => {
let _ = tx_wait.send(SchedulerEventJson::WaitError(e));
}
}
});
}
}
}
SchedulerEventJson::Resume { node, arg, event } => {
waiting = waiting.saturating_sub(1);
active += 1;
let snapshot = state_for_spawn
.lock()
.expect("state mutex poisoned")
.clone();
spawn_json_node_task(
node,
wrap_resume_arg(arg, event),
snapshot,
tx_for_spawn.clone(),
user_data_bits,
callback,
)?;
}
SchedulerEventJson::WaitError(e) => return Err(e),
}
}
let final_state = state.lock().expect("state mutex poisoned").clone();
Ok(final_state)
}
fn wrap_resume_arg(arg: Value, event: WaitEvent) -> Value {
serde_json::json!({
"__lg_resume_arg__": arg,
"__lg_resume_event__": event,
})
}
#[no_mangle]
pub extern "C" fn rc_engine_new() -> *mut Engine {
Box::into_raw(Box::new(Engine::new()))
@@ -382,6 +144,86 @@ pub unsafe extern "C" fn rc_wait_any_of_json(
}
}
#[no_mangle]
/// # Safety
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
/// `stream_mode` must be null or a valid null-terminated UTF-8 string pointer.
pub unsafe extern "C" fn rc_start_stream(
ptr: *mut Engine,
stream_mode: *const c_char,
) -> *mut c_char {
if ptr.is_null() {
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
}
let mode = if stream_mode.is_null() {
None
} else {
match cstr_to_str(stream_mode) {
Ok(v) => Some(v),
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
}
};
match (*ptr).start_stream(mode) {
Ok(()) => into_c_ptr("{\"ok\":true}".to_string()),
Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
}
}
#[no_mangle]
/// # Safety
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
pub unsafe extern "C" fn rc_receive_stream_json(ptr: *mut Engine) -> *mut c_char {
if ptr.is_null() {
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
}
let event = run_loop_block_on((*ptr).receive_stream_async());
match event {
Some(value) => match serde_json::to_string(&value) {
Ok(s) => into_c_ptr(format!("{{\"ok\":true,\"has_event\":true,\"event\":{s}}}")),
Err(e) => into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
},
None => into_c_ptr("{\"ok\":true,\"has_event\":false}".to_string()),
}
}
#[no_mangle]
/// # Safety
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
/// `value_json` must be a valid null-terminated UTF-8 string pointer.
pub unsafe extern "C" fn rc_send_custom_stream_event(
ptr: *mut Engine,
value_json: *const c_char,
) -> *mut c_char {
if ptr.is_null() {
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
}
let value_json = match cstr_to_str(value_json) {
Ok(v) => v,
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
};
let value: Value = match serde_json::from_str(value_json) {
Ok(v) => v,
Err(e) => {
return into_c_ptr(format!(
"{{\"ok\":false,\"error\":\"invalid JSON value: {e}\"}}"
))
}
};
(*ptr).send_custom_stream_event(value);
into_c_ptr("{\"ok\":true}".to_string())
}
#[no_mangle]
/// # Safety
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
pub unsafe extern "C" fn rc_close_stream(ptr: *mut Engine) -> *mut c_char {
if ptr.is_null() {
return into_c_ptr("{\"ok\":false,\"error\":\"null engine pointer\"}".to_string());
}
(*ptr).close_stream();
into_c_ptr("{\"ok\":true}".to_string())
}
#[no_mangle]
/// # Safety
/// `ptr` must be a valid engine pointer from `rc_engine_new`.
@@ -393,6 +235,7 @@ pub unsafe extern "C" fn rc_run_graph_json(
finish_point: *const c_char,
initial_state_json: *const c_char,
initial_input_json: *const c_char,
stream_mode: *const c_char,
user_data: libc::c_ulong,
callback: Option<CNodeCallback>,
) -> *mut c_char {
@@ -434,21 +277,66 @@ pub unsafe extern "C" fn rc_run_graph_json(
))
}
};
let stream_mode = if stream_mode.is_null() {
None
} else {
match cstr_to_str(stream_mode) {
Ok(v) => Some(v.to_string()),
Err(e) => return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}")),
}
};
if let Err(e) = (*ptr).start_stream(stream_mode.as_deref()) {
return into_c_ptr(format!("{{\"ok\":false,\"error\":\"{e}\"}}"));
}
let (tx, rx) = mpsc::channel::<Result<Value, String>>();
let user_data = CUserData(user_data);
let user_data_bits = user_data;
let run_engine = (*ptr).clone();
let submit = run_loop_spawn(async move {
let out = run_graph_scheduler_json(
let callback_wrapper = move |node: String,
arg: Value,
state_snapshot: Value|
-> Result<NodeOutcome<Value, Value>, String> {
let node_c =
CString::new(node.clone()).map_err(|e| format!("invalid node name: {e}"))?;
let arg_json = serde_json::to_string(&arg)
.map_err(|e| format!("serialize arg for `{node}` failed: {e}"))?;
let state_json = serde_json::to_string(&state_snapshot)
.map_err(|e| format!("serialize state for `{node}` failed: {e}"))?;
let arg_c =
CString::new(arg_json).map_err(|e| format!("invalid arg JSON bytes: {e}"))?;
let state_c =
CString::new(state_json).map_err(|e| format!("invalid state JSON bytes: {e}"))?;
let out_ptr = unsafe {
callback(
user_data_bits,
node_c.as_ptr() as *mut c_char,
arg_c.as_ptr() as *mut c_char,
state_c.as_ptr() as *mut c_char,
)
};
if out_ptr.is_null() {
return Err(format!("callback returned null for `{node}`"));
}
let out_raw = unsafe { CStr::from_ptr(out_ptr) }
.to_string_lossy()
.into_owned();
unsafe {
libc::free(out_ptr.cast());
}
parse_callback_envelope_json(&out_raw, &node)
};
let out = run_graph_json_with_callback(
entry_point,
finish_point,
initial_state,
initial_input,
run_engine,
user_data,
callback,
run_engine.clone(),
callback_wrapper,
)
.await;
run_engine.close_stream();
let _ = tx.send(out);
});
if let Err(e) = submit {
-423
View File
@@ -1,423 +0,0 @@
#[cfg(feature = "python-bindings")]
use crate::engine::{
node_pool_execute, run_loop_block_on, run_loop_spawn, AnyOfCondition, Engine, NodeExecResult,
NodeOutcome, SendPayload, WaitCondition, WaitEvent, WaitRequest,
};
#[cfg(feature = "python-bindings")]
use pyo3::exceptions::PyValueError;
#[cfg(feature = "python-bindings")]
use pyo3::prelude::*;
#[cfg(feature = "python-bindings")]
use pyo3::types::PyAny;
#[cfg(feature = "python-bindings")]
use pyo3::types::{PyDict, PyList, PyTuple};
#[cfg(feature = "python-bindings")]
use serde_json::Value;
#[cfg(feature = "python-bindings")]
use std::sync::mpsc;
#[cfg(feature = "python-bindings")]
use std::sync::Arc;
#[cfg(feature = "python-bindings")]
use tokio::sync::mpsc as tokio_mpsc;
#[cfg(feature = "python-bindings")]
#[pyclass]
struct PyRustEngine {
inner: Engine,
}
#[cfg(feature = "python-bindings")]
#[pymethods]
impl PyRustEngine {
#[new]
fn new() -> Self {
Self {
inner: Engine::new(),
}
}
fn add_async_channel(&self, name: &str) {
self.inner.add_async_channel(name);
}
fn publish_json(&self, channel: &str, value_json: &str) -> PyResult<()> {
let value: Value = serde_json::from_str(value_json)
.map_err(|e| PyValueError::new_err(format!("Invalid JSON value: {e}")))?;
self.inner
.publish_json(channel, value)
.map_err(PyValueError::new_err)
}
fn publish_obj(&self, py: Python<'_>, channel: &str, value: Py<PyAny>) -> PyResult<()> {
let value_json = py_obj_to_json_string(py, &value.bind(py))?;
let parsed: Value = serde_json::from_str(&value_json)
.map_err(|e| PyValueError::new_err(format!("Invalid Python JSON value: {e}")))?;
self.inner
.publish_json(channel, parsed)
.map_err(PyValueError::new_err)
}
fn wait_any_of_json(&self, any_of_json: &str) -> PyResult<String> {
let any_of: AnyOfCondition = serde_json::from_str(any_of_json)
.map_err(|e| PyValueError::new_err(format!("Invalid any_of JSON: {e}")))?;
let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of))
.map_err(PyValueError::new_err)?;
serde_json::to_string(&event)
.map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))
}
fn wait_channel(&self, py: Python<'_>, channel: &str, n: usize) -> PyResult<Py<PyAny>> {
let cond = WaitCondition::Channel {
channel: channel.to_string(),
n,
};
let event =
run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?;
let event_json = serde_json::to_string(&event)
.map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?;
json_string_to_py_obj(py, &event_json)
}
fn wait_timer(&self, py: Python<'_>, seconds: f64) -> PyResult<Py<PyAny>> {
let cond = WaitCondition::Timer { seconds };
let event =
run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?;
let event_json = serde_json::to_string(&event)
.map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?;
json_string_to_py_obj(py, &event_json)
}
fn wait_any_of_obj(&self, py: Python<'_>, any_of_payload: Py<PyAny>) -> PyResult<Py<PyAny>> {
let payload_json = py_obj_to_json_string(py, &any_of_payload.bind(py))?;
let any_of: AnyOfCondition = serde_json::from_str(&payload_json)
.map_err(|e| PyValueError::new_err(format!("Invalid any_of payload: {e}")))?;
let event = run_loop_block_on(self.inner.wait_for_any_of_async(&any_of))
.map_err(PyValueError::new_err)?;
let event_json = serde_json::to_string(&event)
.map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))?;
json_string_to_py_obj(py, &event_json)
}
fn wait_condition_json(&self, cond_json: &str) -> PyResult<String> {
let cond: WaitCondition = serde_json::from_str(cond_json)
.map_err(|e| PyValueError::new_err(format!("Invalid condition JSON: {e}")))?;
let event =
run_loop_block_on(self.inner.wait_for_async(&cond)).map_err(PyValueError::new_err)?;
serde_json::to_string(&event)
.map_err(|e| PyValueError::new_err(format!("Serialize event failed: {e}")))
}
fn run_graph_py(
&self,
py: Python<'_>,
entry_point: &str,
finish_point: &str,
initial_state: Py<PyAny>,
callback: Py<PyAny>,
) -> PyResult<Py<PyAny>> {
let state = Arc::new(initial_state);
let callback = Arc::new(callback);
let entry_point = entry_point.to_string();
let finish_point = finish_point.to_string();
let (done_tx, done_rx) = mpsc::channel::<Result<(), String>>();
let state_for_run = Arc::clone(&state);
let engine = self.inner.clone();
run_loop_spawn(async move {
let run_result =
run_graph_scheduler(entry_point, finish_point, callback, state_for_run, engine)
.await;
let _ = done_tx.send(run_result);
})
.map_err(PyValueError::new_err)?;
let run_result = py
.allow_threads(move || done_rx.recv())
.map_err(|e| PyValueError::new_err(format!("run-loop recv failed: {e}")))?;
run_result.map_err(PyValueError::new_err)?;
Ok((*state).clone_ref(py))
}
}
#[cfg(feature = "python-bindings")]
enum SchedulerEventPy {
Node(Result<NodeExecutionPy, String>),
Resume {
node: String,
arg: Py<PyAny>,
event: WaitEvent,
},
WaitError(String),
}
struct NodeExecutionPy {
node: String,
arg: Py<PyAny>,
outcome: NodeOutcome<Py<PyAny>, Py<PyAny>>,
}
#[cfg(feature = "python-bindings")]
async fn run_graph_scheduler(
entry_point: String,
finish_point: String,
callback: Arc<Py<PyAny>>,
state: Arc<Py<PyAny>>,
engine: Engine,
) -> Result<(), String> {
let (tx, mut rx) = tokio_mpsc::unbounded_channel::<SchedulerEventPy>();
let initial_arg = Python::with_gil(|py| (*state).clone_ref(py));
let callback_for_spawn = Arc::clone(&callback);
let state_for_spawn = Arc::clone(&state);
let tx_for_spawn = tx.clone();
let state_for_merge = Arc::clone(&state);
let mut active: usize = 1;
let mut waiting: usize = 0;
spawn_node_task(
entry_point,
initial_arg,
tx_for_spawn.clone(),
Arc::clone(&callback_for_spawn),
Arc::clone(&state_for_spawn),
)?;
while active > 0 || waiting > 0 {
let event = rx
.recv()
.await
.ok_or_else(|| "scheduler event channel closed".to_string())?;
match event {
SchedulerEventPy::Node(result) => {
active = active.saturating_sub(1);
let exec = result?;
match exec.outcome {
NodeOutcome::Completed(node_result) => {
Python::with_gil(|py| -> Result<(), String> {
if let Some(update) = node_result.update {
apply_update_to_state(py, state_for_merge.as_ref(), &update)
.map_err(|e| {
format!("state merge failed for `{}`: {e}", exec.node)
})?;
}
Ok(())
})?;
if exec.node == finish_point {
break;
}
for send in node_result.sends {
active += 1;
spawn_node_task(
send.node,
send.arg,
tx_for_spawn.clone(),
Arc::clone(&callback_for_spawn),
Arc::clone(&state_for_spawn),
)?;
}
}
NodeOutcome::Suspended { wait } => {
waiting += 1;
let tx_wait = tx_for_spawn.clone();
let node = exec.node;
let arg = exec.arg;
let engine_for_wait = engine.clone();
tokio::spawn(async move {
let outcome = engine_for_wait.wait_request_async(&wait).await;
match outcome {
Ok(event) => {
let _ = tx_wait.send(SchedulerEventPy::Resume { node, arg, event });
}
Err(e) => {
let _ = tx_wait.send(SchedulerEventPy::WaitError(e));
}
}
});
}
}
}
SchedulerEventPy::Resume { node, arg, event } => {
waiting = waiting.saturating_sub(1);
active += 1;
let resume_arg = wrap_resume_arg(&arg, &event)?;
spawn_node_task(
node,
resume_arg,
tx_for_spawn.clone(),
Arc::clone(&callback_for_spawn),
Arc::clone(&state_for_spawn),
)?;
}
SchedulerEventPy::WaitError(e) => return Err(e),
}
}
Ok(())
}
#[cfg(feature = "python-bindings")]
fn spawn_node_task(
node: String,
arg: Py<PyAny>,
tx: tokio_mpsc::UnboundedSender<SchedulerEventPy>,
callback: Arc<Py<PyAny>>,
state_for_task: Arc<Py<PyAny>>,
) -> Result<(), String> {
node_pool_execute(move || {
let node_for_result = node.clone();
let arg_for_result = Python::with_gil(|py| arg.clone_ref(py));
let outcome = Python::with_gil(
|py| -> Result<NodeExecutionPy, String> {
let callback_bound = callback.as_ref().bind(py);
let payload_obj = callback_bound
.call1((node.as_str(), arg, (*state_for_task).clone_ref(py)))
.map_err(|e| format!("callback failed for node `{node}`: {e}"))?;
let payload = parse_node_outcome(py, &payload_obj)
.map_err(|e| format!("invalid callback payload for `{node}`: {e}"))?;
Ok(NodeExecutionPy {
node: node_for_result,
arg: arg_for_result,
outcome: payload,
})
},
);
let _ = tx.send(SchedulerEventPy::Node(outcome));
})
}
#[cfg(feature = "python-bindings")]
fn parse_node_outcome(
py: Python<'_>,
payload_obj: &Bound<'_, PyAny>,
) -> Result<NodeOutcome<Py<PyAny>, Py<PyAny>>, String> {
let payload_dict = payload_obj
.downcast::<PyDict>()
.map_err(|_| "payload must be a dict".to_string())?;
let suspended_item = payload_dict
.get_item("suspend")
.map_err(|e| format!("failed to read suspend: {e}"))?;
if let Some(wait_obj) = suspended_item {
let wait_json = py_obj_to_json_string(py, &wait_obj)
.map_err(|e| format!("failed to encode suspend payload: {e}"))?;
let wait: WaitRequest =
serde_json::from_str(&wait_json).map_err(|e| format!("invalid suspend payload: {e}"))?;
return Ok(NodeOutcome::Suspended { wait });
}
let update_item = payload_dict
.get_item("update")
.map_err(|e| format!("failed to read update: {e}"))?;
let update = match update_item {
Some(v) if !v.is_none() => Some(v.unbind()),
_ => None,
};
let sends_obj = payload_dict
.get_item("sends")
.map_err(|e| format!("failed to read sends: {e}"))?
.ok_or_else(|| "missing sends".to_string())?;
let sends_list = sends_obj
.downcast::<PyList>()
.map_err(|_| "sends must be a list".to_string())?;
let mut sends = Vec::with_capacity(sends_list.len());
for item in sends_list.iter() {
let send_dict = item
.downcast::<PyDict>()
.map_err(|_| "send item must be a dict".to_string())?;
let node_obj = send_dict
.get_item("node")
.map_err(|e| format!("failed to read send.node: {e}"))?
.ok_or_else(|| "send.node is required".to_string())?;
let node = node_obj
.extract::<String>()
.map_err(|e| format!("send.node must be string: {e}"))?;
let arg: Py<PyAny> = match send_dict.get_item("arg") {
Ok(Some(v)) => v.unbind(),
Ok(None) => Python::with_gil(|py| py.None()),
Err(e) => return Err(format!("failed to read send.arg: {e}")),
};
sends.push(SendPayload { node, arg });
}
Ok(NodeOutcome::Completed(NodeExecResult { update, sends }))
}
#[cfg(feature = "python-bindings")]
fn wrap_resume_arg(arg: &Py<PyAny>, event: &WaitEvent) -> Result<Py<PyAny>, String> {
Python::with_gil(|py| -> Result<Py<PyAny>, String> {
let wrapper = PyDict::new(py);
wrapper
.set_item("__lg_resume_arg__", arg.clone_ref(py))
.map_err(|e| format!("failed to set resume arg: {e}"))?;
let event_json =
serde_json::to_string(event).map_err(|e| format!("failed to encode wait event: {e}"))?;
let event_obj =
json_string_to_py_obj(py, &event_json).map_err(|e| format!("failed to parse event: {e}"))?;
wrapper
.set_item("__lg_resume_event__", event_obj.bind(py))
.map_err(|e| format!("failed to set resume event: {e}"))?;
Ok(wrapper.unbind().into_any())
})
}
#[cfg(feature = "python-bindings")]
fn apply_update_to_state(py: Python<'_>, state: &Py<PyAny>, update: &Py<PyAny>) -> PyResult<()> {
let state_obj = state.bind(py);
let update_obj = update.bind(py);
if update_obj.is_none() {
return Ok(());
}
if state_obj.is_instance_of::<PyDict>() && update_obj.is_instance_of::<PyDict>() {
let state_dict = state_obj.downcast::<PyDict>()?;
let update_dict = update_obj.downcast::<PyDict>()?;
state_dict.call_method1("update", (update_dict,))?;
return Ok(());
}
if let Ok(tuple_like) = update_obj.downcast::<PyList>() {
apply_pair_updates(state_obj, tuple_like)?;
return Ok(());
}
if let Ok(tuple_like) = update_obj.downcast::<PyTuple>() {
let list = PyList::new(py, tuple_like)?;
apply_pair_updates(state_obj, &list)?;
}
Ok(())
}
#[cfg(feature = "python-bindings")]
fn apply_pair_updates(state_obj: &Bound<'_, PyAny>, entries: &Bound<'_, PyList>) -> PyResult<()> {
if !state_obj.is_instance_of::<PyDict>() {
return Ok(());
}
let state_dict = state_obj.downcast::<PyDict>()?;
for entry in entries.iter() {
if let Ok(pair) = entry.downcast::<PyTuple>() {
if pair.len() == 2 {
let key_obj = pair.get_item(0)?;
if let Ok(key) = key_obj.extract::<String>() {
let value_obj = pair.get_item(1)?;
state_dict.set_item(key, value_obj)?;
}
}
}
}
Ok(())
}
#[cfg(feature = "python-bindings")]
fn py_obj_to_json_string(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult<String> {
let json_mod = py.import("json")?;
let dumped = json_mod.call_method1("dumps", (obj,))?;
dumped.extract::<String>()
}
#[cfg(feature = "python-bindings")]
fn json_string_to_py_obj(py: Python<'_>, value: &str) -> PyResult<Py<PyAny>> {
let json_mod = py.import("json")?;
let loaded = json_mod.call_method1("loads", (value,))?;
Ok(loaded.unbind())
}
#[cfg(feature = "python-bindings")]
#[pymodule]
fn langgraph_rust_core(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyRustEngine>()?;
Ok(())
}
+214
View File
@@ -0,0 +1,214 @@
# This file is automatically @generated by Cargo.
# It is not intended for manual editing.
version = 4
[[package]]
name = "bitflags"
version = "2.11.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af"
[[package]]
name = "cfg-if"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801"
[[package]]
name = "itoa"
version = "1.0.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"
[[package]]
name = "langgraph_rust_core"
version = "0.1.0"
dependencies = [
"libc",
"parking_lot",
"serde",
"serde_json",
"tokio",
]
[[package]]
name = "libc"
version = "0.2.183"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d"
[[package]]
name = "lock_api"
version = "0.4.14"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965"
dependencies = [
"scopeguard",
]
[[package]]
name = "memchr"
version = "2.8.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79"
[[package]]
name = "parking_lot"
version = "0.12.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a"
dependencies = [
"lock_api",
"parking_lot_core",
]
[[package]]
name = "parking_lot_core"
version = "0.9.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1"
dependencies = [
"cfg-if",
"libc",
"redox_syscall",
"smallvec",
"windows-link",
]
[[package]]
name = "pin-project-lite"
version = "0.2.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
[[package]]
name = "proc-macro2"
version = "1.0.106"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934"
dependencies = [
"unicode-ident",
]
[[package]]
name = "quote"
version = "1.0.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
dependencies = [
"proc-macro2",
]
[[package]]
name = "redox_syscall"
version = "0.5.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d"
dependencies = [
"bitflags",
]
[[package]]
name = "scopeguard"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49"
[[package]]
name = "serde"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e"
dependencies = [
"serde_core",
"serde_derive",
]
[[package]]
name = "serde_core"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad"
dependencies = [
"serde_derive",
]
[[package]]
name = "serde_derive"
version = "1.0.228"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "serde_json"
version = "1.0.149"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86"
dependencies = [
"itoa",
"memchr",
"serde",
"serde_core",
"zmij",
]
[[package]]
name = "smallvec"
version = "1.15.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
[[package]]
name = "syn"
version = "2.0.117"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99"
dependencies = [
"proc-macro2",
"quote",
"unicode-ident",
]
[[package]]
name = "tokio"
version = "1.50.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d"
dependencies = [
"pin-project-lite",
"tokio-macros",
]
[[package]]
name = "tokio-macros"
version = "2.6.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c"
dependencies = [
"proc-macro2",
"quote",
"syn",
]
[[package]]
name = "unicode-ident"
version = "1.0.24"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75"
[[package]]
name = "windows-link"
version = "0.2.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5"
[[package]]
name = "zmij"
version = "1.0.21"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa"
+17
View File
@@ -0,0 +1,17 @@
[package]
name = "langgraph_rust_core"
version = "0.1.0"
edition = "2021"
[lib]
name = "langgraph_rust_core"
path = "../rust-core/src/lib.rs"
crate-type = ["cdylib", "rlib"]
[dependencies]
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
parking_lot = "0.12"
libc = "0.2"
tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time"] }
+33
View File
@@ -0,0 +1,33 @@
PYPI_REPOSITORY ?= pypi
PYPI_TOKEN ?=
PYTHON_VERSION ?= 3.13
TEST_BASE = uv run --python $(PYTHON_VERSION) --with pytest --with anyio --with typing_extensions --with pydantic --with langgraph pytest -q -s
.PHONY: publish-to-pypi-saf-python-sdk all-tests test_primitives test_sub_agents test_update_elision test_run_pool_size test-benchmark
publish-to-pypi-saf-python-sdk:
@test -n "$(PYPI_TOKEN)" || (echo "PYPI_TOKEN is required"; exit 1)
cd "$(CURDIR)" && \
MATURIN_PYPI_TOKEN="$(PYPI_TOKEN)" \
uvx maturin publish --repository $(PYPI_REPOSITORY) --non-interactive --skip-existing --no-sdist
all-tests:
$(TEST_BASE) tests/advanced-graph/test_*.py
test_primitives:
$(TEST_BASE) tests/advanced-graph/test_primitives.py
test_sub_agents:
$(TEST_BASE) tests/advanced-graph/test_sub_agents.py
test_update_elision:
$(TEST_BASE) tests/advanced-graph/test_update_elision.py
test_run_pool_size:
$(TEST_BASE) tests/advanced-graph/test_run_pool_size.py
test-benchmark:
# Increase per-process file descriptor limit for high-concurrency benchmark runs.
ulimit -n 65536 || true; \
uv run --python $(PYTHON_VERSION) --with anyio --with typing_extensions --with pydantic --with langgraph python tests/advanced-graph/benchmark_stategraph_vs_advancedgraph.py
+16
View File
@@ -0,0 +1,16 @@
# saf-python-sdk
Standalone Python SDK for the `advanced_graph` runtime backed by the Rust engine.
This package intentionally contains only:
- `saf_python_sdk.advanced_graph` (Python API)
- Rust core engine via C bindings (`ctypes`)
It does not package the original `langgraph` `stategraph` stack.
## Moved assets
- Original advanced graph design doc: `evolve-extend-langgraph.md`
- Original advanced graph tests: `tests/advanced-graph/`
@@ -32,11 +32,10 @@ Sub-agents also cannot simply be modeled as subgraphs, because subgraphs today e
The closest workaround today is double texting, but it has a fundamental flaw: when a new audio input arrives, the previous one is interrupted and canceled rather than being allowed to gracefully complete. The workflow code itself should have the control to decide whether to stop running.
LangGraph is, at its core, a general-purpose workflow engine. Although we focus primarily on agent development, none of the primitives it offers are exclusive to agents or dedicated solely to agentic use cases. Conversely, there is nothing that a general-purpose workflow engine provides that we can safely assume agent development will _never_ need.
LangGraph is, at its core, a general-purpose workflow engine. Although we focus primarily on agent development, none of the primitives it offers are exclusive to agents or dedicated solely to agentic use cases. Conversely, there is nothing that a general-purpose workflow engine provides that we can safely assume agent development will _never_ need.
The difference is probably only priority. For example, durable timer where a step can sleep for hours, days or months before resuming. Traditional workflow engines — those built for general microservice orchestration(which doesn't need streaming) -- they may need durable timers. In the agent development world today, most agents are still relatively simple. There are not yet many scenarios that require a step to wait for hours or days before proceeding.
## Deriving What's Needed from First Principles
Before jumping to solutions, it is worth stepping back and asking a fundamental question: what is an orchestration engine, and what do users expect it to provide?
@@ -47,7 +46,7 @@ Starting from the simple. A developer could write a simple `main` function — a
But if that machine crashes, you probably do not want the process to start over from scratch. You want it to resume from the last step that completed successfully. And if a step fails, you might want it to retry automatically before giving up.
LangGraph handles this case very well.
LangGraph handles this case very well.
There is an important constraint worth calling out explicitly: LangGraph requires the developer to organize their code into **nodes**, which serve as the boundaries at which checkpoints can be taken. This is a constraint shared by every workflow engine — it is simply not feasible to persist a checkpoint after every single line of arbitrary code.
@@ -59,9 +58,9 @@ This is precisely why LangGraph's superstep restriction feels awkward in practic
Multiple threads and processes do, however, need to coordinate with each other. In concurrent programming, channels are an essential primitive precisely because they provide a safe, structured way for threads to communicate and synchronize without relying on shared mutable memory — avoiding data races and deadlocks. In some cases, threads may use locking for coordination, but the preferred approach is message passing through channels.
NOTE: "channel" is overloaded term here as it's also an internal term within current LangGraph pregel algorithm.
NOTE: "channel" is overloaded term here as it's also an internal term within current LangGraph pregel algorithm.
LangGraph already has a mechanism that is closely related: `interrupt`. A run can be interrupted, and then another run can resume it. If we look at this through the lens of channels, `interrupt` is essentially a **channel with size 0** — a synchronous rendezvous point where one side blocks until the other side is ready.
LangGraph already has a mechanism that is closely related: `interrupt`. A run can be interrupted, and then another run can resume it. If we look at this through the lens of channels, `interrupt` is essentially a **channel with size 0** — a synchronous rendezvous point where one side blocks until the other side is ready.
The natural extension:
@@ -97,7 +96,6 @@ Branch 2: a → b2 → b22 → ...
Each branch advances at its own pace. `b1` finishing triggers `b11` immediately, without waiting for `b2`.
No API change is needed from the user's perspective — the graph definition stays the same. The change is in the execution semantics: the engine no longer forces all parallel nodes to synchronize at each step boundary. Each branch is checkpointed independently, so if `b1 → b11` completes while `b2` is still running, `b11`'s result is already persisted.
This is necessary for the next one -- Light-weight Interrupt: Only Block the Current Node. Because we want to let other nodes continue to run while a node is waiting on something.
@@ -118,13 +116,12 @@ The `wait_for` call takes a channel name and optionally a count `N`, meaning "wa
See [test_sub_agents.py](../libs/langgraph/tests/advanced-graph/test_sub_agents.py)
### P2: likely needed
#### subGraph redesign
#### durable timers
#### more flexiable waiting conditions on interrupts
#### locking on state fields
### P3: future needed or nice to have
#### RPC
#### RPC
+26
View File
@@ -0,0 +1,26 @@
[build-system]
requires = ["maturin>=1.7,<2"]
build-backend = "maturin"
[project]
name = "saf-python-sdk"
version = "0.1.2"
description = "Standalone advanced graph runtime powered by Rust engine"
readme = "README.md"
requires-python = ">=3.10"
license = "MIT"
authors = [{ name = "LangGraph Contributors" }]
classifiers = [
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3 :: Only",
"Programming Language :: Rust",
"License :: OSI Approved :: MIT License",
"Operating System :: OS Independent",
]
[tool.maturin]
python-source = "python"
module-name = "saf_python_sdk.langgraph_rust_core"
bindings = "pyo3"
features = ["python-bindings"]
@@ -0,0 +1,4 @@
from .types import Command, Send
__all__ = ["Command", "Send"]
@@ -1,4 +1,4 @@
from langgraph.advanced_graph.state import (
from .state import (
AdvancedStateGraph,
AnyOfCondition,
ChannelCondition,
@@ -11,15 +11,16 @@ from langgraph.advanced_graph.state import (
timer_condition,
)
__all__ = (
__all__ = [
"AdvancedStateGraph",
"AnyOfCondition",
"ChannelCondition",
"Context",
"CompiledGraphEngine",
"Context",
"GraphRunHandler",
"ChannelCondition",
"TimerCondition",
"any_of",
"AnyOfCondition",
"channel_condition",
"timer_condition",
)
"any_of",
]
@@ -2,18 +2,18 @@ from __future__ import annotations
import atexit
import asyncio
import os
import inspect
import os
import threading
from concurrent.futures import ThreadPoolExecutor
from collections.abc import Callable, Coroutine, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from datetime import timedelta
from typing import Any, Generic, TypeVar, cast
from langgraph_rust_core import PyRustEngine # type: ignore[import-untyped]
from saf_python_sdk.rust_core_cffi import PyRustEngine
from langgraph.types import Command, Send
from saf_python_sdk.types import Command, Send
StateT = TypeVar("StateT")
@@ -53,7 +53,7 @@ def _advanced_graph_executor() -> ThreadPoolExecutor:
worker_count = max(worker_count, 1)
_EXECUTOR = ThreadPoolExecutor(
max_workers=worker_count,
thread_name_prefix="langgraph-advanced-py",
thread_name_prefix="saf-advanced-py",
)
atexit.register(_shutdown_advanced_graph_executor)
return _EXECUTOR
@@ -162,12 +162,15 @@ class CompiledGraphEngine(Generic[StateT]):
handler = await self.astart(initial_state)
return await handler
async def astart(self, initial_state: StateT) -> GraphRunHandler[StateT]:
async def astart(
self, initial_state: StateT, *, stream_mode: str | None = None
) -> GraphRunHandler[StateT]:
run = _GraphEngineRun(
nodes=self._nodes,
async_channel_specs=self._async_channels,
entry_point=self._entry_point,
finish_point=self._finish_point,
stream_mode=stream_mode,
)
task = asyncio.create_task(run.run(initial_state))
return GraphRunHandler(run=run, task=task)
@@ -191,6 +194,9 @@ class Context:
async def apublish_to_channel(self, channel: str, value: Any) -> None:
await self._run.publish(channel, value)
def send_custom_stream_event(self, value: Any) -> None:
self._run.send_custom_stream_event(value)
class GraphRunHandler(Generic[StateT]):
"""Handle for an active in-memory run."""
@@ -204,6 +210,20 @@ class GraphRunHandler(Generic[StateT]):
raise RuntimeError("Run has already completed")
await self._run.publish(channel, value)
async def receive_stream(self) -> Any | None:
while True:
loop = asyncio.get_running_loop()
event = await loop.run_in_executor(
_advanced_graph_executor(),
self._run.receive_stream_sync,
)
if event is not None or self._task.done():
return event
await asyncio.sleep(0.005)
def close_stream(self) -> None:
self._run.close_stream_sync()
async def aresult(self) -> StateT:
return await self._task
@@ -219,10 +239,12 @@ class _GraphEngineRun:
async_channel_specs: dict[str, _ChannelSpec],
entry_point: str,
finish_point: str | None,
stream_mode: str | None,
) -> None:
self._nodes = nodes
self._entry_point = entry_point
self._finish_point = finish_point
self._stream_mode = stream_mode
self._rust_engine = PyRustEngine()
for name in async_channel_specs:
self._rust_engine.add_async_channel(name)
@@ -242,6 +264,7 @@ class _GraphEngineRun:
finish_point,
initial_state,
self._execute_node_for_rust,
self._stream_mode,
)
self._state = result_obj
return cast(StateT, self._state)
@@ -305,6 +328,15 @@ class _GraphEngineRun:
def _publish_sync(self, channel: str, value: Any) -> None:
self._rust_engine.publish_obj(channel, value)
def send_custom_stream_event(self, value: Any) -> None:
self._rust_engine.send_custom_stream_event_obj(value)
def receive_stream_sync(self) -> Any | None:
return self._rust_engine.receive_stream_obj()
def close_stream_sync(self) -> None:
self._rust_engine.close_stream()
def _execute_node_for_rust(
self, node_name: str, node_input: Any, state: Any
) -> dict[str, Any]:
@@ -350,14 +382,10 @@ class _GraphEngineRun:
return event
def _run_awaitable_in_worker(self, awaitable: Coroutine[Any, Any, Any]) -> Any:
loop = cast(
asyncio.AbstractEventLoop | None,
getattr(self._local, "worker_loop", None),
)
if loop is None or loop.is_closed():
loop = asyncio.new_event_loop()
self._local.worker_loop = loop
return loop.run_until_complete(awaitable)
# Create and close a dedicated loop per execution to avoid
# interpreter-shutdown warnings from lingering thread-local loops.
return asyncio.run(awaitable)
def _normalize_result_to_sends(result: Any, *, default_input: Any) -> list[Send]:
if result is None:
@@ -542,3 +570,4 @@ def _invoke_node(node: Callable[..., Any], ctx: Context, node_input: Any, state:
return node(node_input, state)
return node(ctx, node_input, state)
@@ -0,0 +1 @@
@@ -0,0 +1,347 @@
from __future__ import annotations
import ctypes
import dataclasses
import json
import subprocess
from copy import deepcopy
from pathlib import Path
import threading
from typing import Any, Callable, get_args, get_origin
class PyRustEngine:
def __init__(self) -> None:
self._lib = _load_rust_lib()
self._engine = self._lib.rc_engine_new()
if not self._engine:
raise RuntimeError("failed to create rust engine")
def __del__(self) -> None:
engine = getattr(self, "_engine", None)
if engine:
self._lib.rc_engine_free(engine)
self._engine = None
def add_async_channel(self, name: str) -> None:
self._call_status(self._lib.rc_add_async_channel, name.encode())
def publish_obj(self, channel: str, value: Any) -> None:
payload = json.dumps(value, ensure_ascii=False).encode()
self._call_status(self._lib.rc_publish_json, channel.encode(), payload)
def wait_any_of_obj(self, any_of_payload: Any) -> Any:
payload = json.dumps(any_of_payload, ensure_ascii=False).encode()
raw = self._consume_json_ptr(self._lib.rc_wait_any_of_json(self._engine, payload))
if not raw.get("ok"):
raise ValueError(raw.get("error", "rust wait_any_of failed"))
return raw["event"]
def wait_channel(self, channel: str, n: int) -> Any:
return self.wait_any_of_obj({"conditions": [{"kind": "channel", "channel": channel, "n": n}]})
def wait_timer(self, seconds: float) -> Any:
return self.wait_any_of_obj({"conditions": [{"kind": "timer", "seconds": seconds}]})
def start_stream(self, stream_mode: str | None) -> None:
encoded = stream_mode.encode() if stream_mode is not None else None
self._call_status(self._lib.rc_start_stream, encoded)
def receive_stream_obj(self) -> Any | None:
raw = self._consume_json_ptr(self._lib.rc_receive_stream_json(self._engine))
if not raw.get("ok"):
raise ValueError(raw.get("error", "rust receive_stream failed"))
if not raw.get("has_event", False):
return None
return raw.get("event")
def send_custom_stream_event_obj(self, value: Any) -> None:
payload = json.dumps(value, ensure_ascii=False).encode()
self._call_status(self._lib.rc_send_custom_stream_event, payload)
def close_stream(self) -> None:
self._call_status(self._lib.rc_close_stream)
def run_graph_py(
self,
entry_point: str,
finish_point: str,
initial_state: Any,
callback: Callable[[str, Any, Any], dict[str, Any]],
stream_mode: str | None = None,
) -> Any:
shared_state = initial_state
shared_state_lock = threading.Lock()
state_type = type(initial_state)
use_shared_state = dataclasses.is_dataclass(initial_state)
initial_state_json = json.dumps(_to_jsonable(shared_state), ensure_ascii=False).encode()
callback_c = _make_node_callback(
callback,
state_type,
use_shared_state,
shared_state,
shared_state_lock,
)
stream_mode_encoded = stream_mode.encode() if stream_mode is not None else None
out = self._consume_json_ptr(
self._lib.rc_run_graph_json(
self._engine,
entry_point.encode(),
finish_point.encode(),
initial_state_json,
initial_state_json,
stream_mode_encoded,
ctypes.c_ulong(0),
callback_c,
)
)
if not out.get("ok"):
raise ValueError(out.get("error", "rust run_graph failed"))
if use_shared_state:
return shared_state
return _coerce_for_type(out["state"], state_type)
def _call_status(self, func: Any, *args: Any) -> None:
raw = self._consume_json_ptr(func(self._engine, *args))
if not raw.get("ok"):
raise ValueError(raw.get("error", "rust call failed"))
def _consume_json_ptr(self, ptr: ctypes.c_void_p) -> dict[str, Any]:
if not ptr:
raise RuntimeError("rust returned null string pointer")
try:
text = ctypes.cast(ptr, ctypes.c_char_p).value
if text is None:
raise RuntimeError("rust returned empty string pointer")
return json.loads(text.decode())
finally:
self._lib.rc_string_free(ptr)
def _make_node_callback(
callback: Callable[[str, Any, Any], dict[str, Any]],
state_type: type[Any],
use_shared_state: bool,
shared_state: Any,
shared_state_lock: threading.Lock,
) -> ctypes.CFUNCTYPE: # type: ignore[type-arg]
cb_type = ctypes.CFUNCTYPE(
ctypes.c_void_p,
ctypes.c_ulong,
ctypes.c_char_p,
ctypes.c_char_p,
ctypes.c_char_p,
)
libc = ctypes.CDLL(None)
libc.malloc.argtypes = [ctypes.c_size_t]
libc.malloc.restype = ctypes.c_void_p
@cb_type
def _callback(
_user_data: int,
node_ptr: bytes,
arg_ptr: bytes,
state_ptr: bytes,
) -> ctypes.c_void_p:
try:
node = node_ptr.decode()
arg = json.loads(arg_ptr.decode())
if use_shared_state:
with shared_state_lock:
before = deepcopy(_to_jsonable(shared_state))
result = callback(node, arg, shared_state)
state_after_call = shared_state
else:
state_raw = json.loads(state_ptr.decode())
state_snapshot = _coerce_for_type(state_raw, state_type)
before = deepcopy(_to_jsonable(state_snapshot))
result = callback(node, arg, state_snapshot)
state_after_call = state_snapshot
if "suspend" in result:
envelope = {"ok": True, "suspend": result["suspend"]}
else:
update = result.get("update")
if update is not None:
update_json = _to_jsonable(update)
if update_json == before:
update = None
update_json = None
else:
if use_shared_state:
_apply_update_to_state(shared_state, update)
after = _to_jsonable(state_after_call)
if update is None and after != before:
update_json = after
elif update is None:
update_json = None
else:
update_json = _to_jsonable(update)
sends = []
for item in result.get("sends", []):
if not isinstance(item, dict):
continue
sends.append(
{
"node": item.get("node"),
"arg": _to_jsonable(item.get("arg")),
}
)
envelope = {
"ok": True,
"payload": {
"update": update_json,
"sends": sends,
},
}
except Exception as exc: # noqa: BLE001
envelope = {"ok": False, "error": f"python callback failed: {exc}"}
return _malloc_c_string(json.dumps(envelope, ensure_ascii=False).encode(), libc)
return _callback
def _malloc_c_string(payload: bytes, libc: Any) -> ctypes.c_void_p:
size = len(payload) + 1
ptr = libc.malloc(size)
if not ptr:
return ctypes.c_void_p(0)
ctypes.memmove(ptr, payload, len(payload))
ctypes.memset(ctypes.c_void_p(ptr + len(payload)), 0, 1)
return ptr
def _load_rust_lib() -> ctypes.CDLL:
env = Path.cwd()
root = _find_repo_root(env)
rust_core = root / "rust-core"
lib_path = _resolve_lib_path(rust_core)
if not lib_path.exists():
subprocess.run(["cargo", "build"], cwd=rust_core, check=True)
lib = ctypes.CDLL(str(lib_path))
_configure_signatures(lib)
return lib
def _find_repo_root(start: Path) -> Path:
current = start.resolve()
for candidate in [current, *current.parents]:
if (candidate / "rust-core").exists() and (candidate / "saf-python-sdk").exists():
return candidate
here = Path(__file__).resolve()
return here.parents[3]
def _resolve_lib_path(rust_core: Path) -> Path:
if (rust_core / "target" / "debug" / "liblanggraph_rust_core.dylib").exists():
return rust_core / "target" / "debug" / "liblanggraph_rust_core.dylib"
if (rust_core / "target" / "debug" / "liblanggraph_rust_core.so").exists():
return rust_core / "target" / "debug" / "liblanggraph_rust_core.so"
if (rust_core / "target" / "debug" / "langgraph_rust_core.dll").exists():
return rust_core / "target" / "debug" / "langgraph_rust_core.dll"
return rust_core / "target" / "debug" / "liblanggraph_rust_core.dylib"
def _configure_signatures(lib: ctypes.CDLL) -> None:
cb_type = ctypes.CFUNCTYPE(
ctypes.c_void_p,
ctypes.c_ulong,
ctypes.c_char_p,
ctypes.c_char_p,
ctypes.c_char_p,
)
lib.rc_engine_new.argtypes = []
lib.rc_engine_new.restype = ctypes.c_void_p
lib.rc_engine_free.argtypes = [ctypes.c_void_p]
lib.rc_engine_free.restype = None
lib.rc_string_free.argtypes = [ctypes.c_void_p]
lib.rc_string_free.restype = None
lib.rc_add_async_channel.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
lib.rc_add_async_channel.restype = ctypes.c_void_p
lib.rc_publish_json.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p]
lib.rc_publish_json.restype = ctypes.c_void_p
lib.rc_wait_any_of_json.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
lib.rc_wait_any_of_json.restype = ctypes.c_void_p
lib.rc_start_stream.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
lib.rc_start_stream.restype = ctypes.c_void_p
lib.rc_receive_stream_json.argtypes = [ctypes.c_void_p]
lib.rc_receive_stream_json.restype = ctypes.c_void_p
lib.rc_send_custom_stream_event.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
lib.rc_send_custom_stream_event.restype = ctypes.c_void_p
lib.rc_close_stream.argtypes = [ctypes.c_void_p]
lib.rc_close_stream.restype = ctypes.c_void_p
lib.rc_run_graph_json.argtypes = [
ctypes.c_void_p,
ctypes.c_char_p,
ctypes.c_char_p,
ctypes.c_char_p,
ctypes.c_char_p,
ctypes.c_char_p,
ctypes.c_ulong,
cb_type,
]
lib.rc_run_graph_json.restype = ctypes.c_void_p
def _to_jsonable(value: Any) -> Any:
if value is None or isinstance(value, (str, int, float, bool)):
return value
if dataclasses.is_dataclass(value):
return {field.name: _to_jsonable(getattr(value, field.name)) for field in dataclasses.fields(value)}
model_dump = getattr(value, "model_dump", None)
if callable(model_dump):
return _to_jsonable(model_dump())
if isinstance(value, dict):
return {str(k): _to_jsonable(v) for k, v in value.items()}
if isinstance(value, (list, tuple, set)):
return [_to_jsonable(v) for v in value]
return value
def _coerce_for_type(value: Any, typ: Any) -> Any:
if value is None:
return None
origin = get_origin(typ)
args = get_args(typ)
if origin is not None:
if origin in (list, tuple, set):
item_type = args[0] if args else Any
items = [_coerce_for_type(v, item_type) for v in value]
if origin is tuple:
return tuple(items)
if origin is set:
return set(items)
return items
if origin is dict:
value_type = args[1] if len(args) == 2 else Any
return {k: _coerce_for_type(v, value_type) for k, v in value.items()}
if isinstance(typ, type):
if dataclasses.is_dataclass(typ):
kwargs = {}
for field in dataclasses.fields(typ):
kwargs[field.name] = _coerce_for_type(value.get(field.name), field.type)
return typ(**kwargs)
model_validate = getattr(typ, "model_validate", None)
if callable(model_validate):
return model_validate(value)
return value
def _apply_update_to_state(state: Any, update: Any) -> None:
if update is None:
return
if isinstance(state, dict):
if isinstance(update, dict):
state.update(update)
return
if dataclasses.is_dataclass(update):
state.update(_to_jsonable(update))
return
return
if dataclasses.is_dataclass(state):
if dataclasses.is_dataclass(update):
for field in dataclasses.fields(state):
setattr(state, field.name, getattr(update, field.name))
return
if isinstance(update, dict):
for key, val in update.items():
setattr(state, key, val)
@@ -0,0 +1,19 @@
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Generic, TypeVar
N = TypeVar("N")
@dataclass
class Send(Generic[N]):
node: N
arg: Any = None
@dataclass
class Command:
update: Any = None
goto: Any = field(default=None)
@@ -9,13 +9,13 @@ from typing import Any
os.environ["LANGGRAPH_RUN_POOL_SIZE"] = "10"
os.environ["LANGGRAPH_NODE_POOL_SIZE"] = "1000"
from langgraph.advanced_graph import (
from saf_python_sdk.advanced_graph import (
AdvancedStateGraph,
channel_condition,
timer_condition,
)
from saf_python_sdk.types import Command, Send
from langgraph.graph import END, START, StateGraph
from langgraph.types import Command, Send
RUNS = 100
MIDDLE_COUNT = 10
@@ -47,6 +47,7 @@ def build_advanced_parallel() -> Any:
graph.add_entry_node(start_node)
for i in range(MIDDLE_COUNT):
async def middle_node(ctx: Any, state: dict[str, Any], idx: int = i) -> None:
_ = idx
_ = state
@@ -71,6 +72,7 @@ def build_advanced_sequential() -> Any:
return out
graph.add_entry_node(start_node)
def make_middle(target: str):
async def middle_node(ctx: Any, state: dict[str, Any]) -> Command:
_ = state
@@ -99,6 +101,7 @@ def build_stategraph_parallel() -> Any:
graph.add_node("start_node", start_node)
for i in range(MIDDLE_COUNT):
async def middle_node(state: dict[str, Any], idx: int = i) -> None:
_ = idx
_ = state
@@ -128,6 +131,7 @@ def build_stategraph_sequential() -> Any:
graph.add_node("start_node", start_node)
for i in range(MIDDLE_COUNT):
async def middle_node(state: dict[str, Any], idx: int = i) -> None:
_ = idx
_ = state
@@ -164,6 +168,7 @@ def build_advanced_parallel_blocking() -> Any:
graph.add_entry_node(start_node)
for i in range(MIDDLE_COUNT):
async def middle_blocking(
ctx: Any, state: dict[str, Any], idx: int = i
) -> None:
@@ -223,6 +228,7 @@ def build_stategraph_parallel_blocking() -> Any:
graph.add_node("start_node", start_node)
for i in range(MIDDLE_COUNT):
async def middle_blocking(state: dict[str, Any], idx: int = i) -> None:
_ = idx
_ = state
@@ -252,6 +258,7 @@ def build_stategraph_sequential_blocking() -> Any:
graph.add_node("start_node", start_node)
for i in range(MIDDLE_COUNT):
async def middle_blocking_seq(state: dict[str, Any], idx: int = i) -> None:
_ = idx
_ = state
@@ -304,3 +311,4 @@ async def main() -> None:
if __name__ == "__main__":
asyncio.run(main())
@@ -1,8 +1,8 @@
import pytest
from typing_extensions import TypedDict
from langgraph.advanced_graph import AdvancedStateGraph, Context
from langgraph.types import Command, Send
from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context
from saf_python_sdk.types import Command, Send
pytestmark = pytest.mark.anyio
@@ -64,3 +64,4 @@ async def test_run_ends_without_finish_node() -> None:
assert result["counter"] == 8
assert result["done"] == "stopped"
assert result["logs"] == ["start", "middle:from_start"]
@@ -8,8 +8,8 @@ def test_run_pool_size_one_still_allows_parallel_runs() -> None:
import asyncio
import time
from typing_extensions import TypedDict
from langgraph.advanced_graph import AdvancedStateGraph, Context, timer_condition
from langgraph.types import Command, Send
from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context, timer_condition
from saf_python_sdk.types import Command, Send
class RunState(TypedDict):
@@ -53,3 +53,4 @@ asyncio.run(main())
)
elapsed = float(completed.stdout.strip().splitlines()[-1])
assert elapsed < 0.35, completed.stdout
@@ -0,0 +1,56 @@
import asyncio
import pytest
from typing_extensions import TypedDict
from saf_python_sdk.advanced_graph import AdvancedStateGraph, Context
from saf_python_sdk.types import Command, Send
pytestmark = pytest.mark.anyio
class StreamState(TypedDict):
done: bool
async def test_custom_stream_receive_and_close() -> None:
graph: AdvancedStateGraph[StreamState] = AdvancedStateGraph(StreamState)
async def start_node(ctx: Context, state: StreamState) -> Command:
ctx.send_custom_stream_event({"step": "start", "value": 1})
await asyncio.sleep(0.08)
ctx.send_custom_stream_event({"step": "start", "value": 2})
return Command(update=state, goto=Send("finish_node", None))
async def finish_node(state: StreamState) -> dict[str, bool]:
return {"done": True}
graph.add_entry_node(start_node)
graph.add_finish_node(finish_node)
handler = await graph.compile().astart({"done": False}, stream_mode="custom")
event = await handler.receive_stream()
assert isinstance(event, dict)
assert event["step"] == "start"
assert event["value"] == 1
handler.close_stream()
assert await handler.receive_stream() is None
result = await handler.aresult()
assert result["done"] is True
async def test_only_custom_stream_mode_supported() -> None:
graph: AdvancedStateGraph[StreamState] = AdvancedStateGraph(StreamState)
async def start_node(ctx: Context, state: StreamState) -> Command:
ctx.send_custom_stream_event({"hello": "world"})
return Command(update=state)
graph.add_entry_node(start_node)
handler = await graph.compile().astart({"done": False}, stream_mode="values")
with pytest.raises(Exception, match="only `custom` is supported"):
await handler.aresult()
@@ -1,11 +1,12 @@
import asyncio
import json
from dataclasses import dataclass
from typing import Any, Literal
import pytest
from typing_extensions import TypedDict
from langgraph.advanced_graph import (
from saf_python_sdk.advanced_graph import (
AdvancedStateGraph,
Context,
any_of,
@@ -14,7 +15,7 @@ from langgraph.advanced_graph import (
)
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.types import Command, Send
from saf_python_sdk.types import Command, Send
pytestmark = pytest.mark.anyio
@@ -212,6 +213,6 @@ async def test_async_sub_graph() -> None:
order_food_idx = output.index("order_food: order submitted")
assert first_sub_idx < second_sub_idx < order_food_idx
assert llm._idx == len(llm.responses)
import json
print(json.dumps(result, ensure_ascii=False, indent=2))
@@ -1,12 +1,12 @@
import asyncio
from dataclasses import dataclass
import pytest
from pydantic import BaseModel
from typing_extensions import TypedDict
import pytest
from langgraph.advanced_graph import AdvancedStateGraph, CompiledGraphEngine
from langgraph.types import Command, Send
from saf_python_sdk.advanced_graph import AdvancedStateGraph, CompiledGraphEngine
from saf_python_sdk.types import Command, Send
pytestmark = pytest.mark.anyio
@@ -46,7 +46,6 @@ def _initial_state() -> UpdateElisionState:
)
async def test_noop_slow_update_does_not_override_fast_update() -> None:
graph: AdvancedStateGraph[UpdateElisionState] = AdvancedStateGraph(UpdateElisionState)
@@ -65,7 +64,6 @@ async def test_noop_slow_update_does_not_override_fast_update() -> None:
async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
await asyncio.sleep(0.1)
# Returns the same values as the initial snapshot.
return state
graph.add_entry_node(start_node)
@@ -101,7 +99,6 @@ async def test_changed_slow_update_overrides_fast_update() -> None:
async def slow_node(state: UpdateElisionState) -> UpdateElisionState:
await asyncio.sleep(0.1)
# Slow node makes real changes for all field types.
state.x = 2
state.dc.value = 2
state.model.value = 2
@@ -124,3 +121,4 @@ async def test_changed_slow_update_overrides_fast_update() -> None:
assert result.td == {"flag": False, "n": 2}
assert result.obj == {"n": 2}
assert result.items == [0, 1, 2]
+8
View File
@@ -0,0 +1,8 @@
version = 1
revision = 3
requires-python = ">=3.10"
[[package]]
name = "saf-python-sdk"
version = "0.1.2"
source = { editable = "." }