Skip to content

Commit d8a0097

Browse files
committed
feat(orchestration): Python SDK parallel/resumable/pipeline (Phase 5c/5d)
Exposes the orchestration grammar to the Python SDK, matching Node: - Session.parallel(specs) and Session.parallel_resumable(specs, workflow_id): specs are dicts (snake_case keys); outcomes returned as dicts in input order. resumable raises if no session_store is configured. - Session.pipeline(items, stages): stages are callables stage(ctx={'previous','item'}) -> spec_dict | None; chains run with no inter-stage barrier and stop on None / failure. Returns one entry per item. FFI safety: methods run sync via py.allow_threads + block_on (no event-loop block). Stage callables are bridged via Python::with_gil from the worker runtime (same as the hook/budget bridges) with an exception-catch that fails closed to None — a raising stage stops only its own chain, never aborts. JSON round-trip (json.dumps/loads + serde) converts dicts <-> AgentStepSpec / StepOutcome. Per-stage output_schema not supported on pipeline stages (use parallel for schema-validated steps). Python SDK fmt + clippy --lib clean.
1 parent 60a6887 commit d8a0097

1 file changed

Lines changed: 187 additions & 0 deletions

File tree

sdk/python/src/lib.rs

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,11 @@ use a3s_code_core::hooks::{
4949
HookMatcher as RustHookMatcher, HookResponse as RustHookResponse,
5050
};
5151
use a3s_code_core::llm::Message as RustMessage;
52+
use a3s_code_core::orchestration::{
53+
execute_pipeline, execute_steps_parallel, execute_steps_parallel_resumable,
54+
AgentStepSpec as RustAgentStepSpec, PipelineStage as RustPipelineStage,
55+
StepOutcome as RustStepOutcome,
56+
};
5257
use a3s_code_core::permissions::{
5358
PermissionDecision as RustPermissionDecision, PermissionPolicy as RustPermissionPolicy,
5459
PermissionRule as RustPermissionRule,
@@ -1396,6 +1401,110 @@ impl PySession {
13961401
Ok(PyAgentResult::from(result))
13971402
}
13981403

1404+
/// Run `specs` as a fan-out of agent steps and return each step's outcome
1405+
/// (a dict) in input order. Each spec is a dict with snake_case keys:
1406+
/// `task_id`, `agent`, `description`, `prompt`, optional `max_steps`,
1407+
/// `parent_session_id`, `output_schema`. A failed step surfaces as
1408+
/// `success: False` without failing the batch.
1409+
fn parallel(&self, py: Python<'_>, specs: Vec<Bound<'_, PyAny>>) -> PyResult<Vec<PyObject>> {
1410+
let rust_specs = specs
1411+
.iter()
1412+
.map(|s| py_to_step_spec(py, s))
1413+
.collect::<PyResult<Vec<_>>>()?;
1414+
let session = self.inner.clone();
1415+
let outcomes = py.allow_threads(move || {
1416+
get_runtime().block_on(async move {
1417+
let executor = session.agent_executor();
1418+
execute_steps_parallel(executor, rust_specs, None).await
1419+
})
1420+
});
1421+
outcomes.iter().map(|o| step_outcome_to_py(py, o)).collect()
1422+
}
1423+
1424+
/// Like `parallel`, but resumable: progress is journaled under
1425+
/// `workflow_id` via the session's store, so an interrupted run skips
1426+
/// already-completed steps. Raises if no `session_store` is configured.
1427+
fn parallel_resumable(
1428+
&self,
1429+
py: Python<'_>,
1430+
specs: Vec<Bound<'_, PyAny>>,
1431+
workflow_id: String,
1432+
) -> PyResult<Vec<PyObject>> {
1433+
let rust_specs = specs
1434+
.iter()
1435+
.map(|s| py_to_step_spec(py, s))
1436+
.collect::<PyResult<Vec<_>>>()?;
1437+
let session = self.inner.clone();
1438+
let outcomes = py
1439+
.allow_threads(move || {
1440+
get_runtime().block_on(async move {
1441+
let Some(store) = session.session_store() else {
1442+
return Err("parallel_resumable requires a session_store on the session");
1443+
};
1444+
let executor = session.agent_executor();
1445+
Ok(execute_steps_parallel_resumable(
1446+
executor,
1447+
rust_specs,
1448+
&workflow_id,
1449+
store,
1450+
None,
1451+
)
1452+
.await)
1453+
})
1454+
})
1455+
.map_err(PyRuntimeError::new_err)?;
1456+
outcomes.iter().map(|o| step_outcome_to_py(py, o)).collect()
1457+
}
1458+
1459+
/// Run each item through a chain of `stages`, with no barrier between
1460+
/// stages. Each stage is a callable `stage(ctx) -> spec_dict | None`, where
1461+
/// `ctx = {"previous": <outcome dict or None>, "item": <item>}`. Return a
1462+
/// spec dict (snake_case keys) to run that step, or `None` to stop the
1463+
/// item's chain. A chain also stops when a step fails. Returns one entry
1464+
/// per item (the last outcome dict, or `None`), in input order.
1465+
///
1466+
/// A stage callable that raises is caught and treated as `None` (stops that
1467+
/// chain). Per-stage `output_schema` is not supported here — use `parallel`
1468+
/// for schema-validated steps.
1469+
fn pipeline(
1470+
&self,
1471+
py: Python<'_>,
1472+
items: Vec<Bound<'_, PyAny>>,
1473+
stages: Vec<Bound<'_, PyAny>>,
1474+
) -> PyResult<Vec<Option<PyObject>>> {
1475+
let rust_items = items
1476+
.iter()
1477+
.map(|i| py_to_json_value(py, i))
1478+
.collect::<PyResult<Vec<_>>>()?;
1479+
let rust_stages: Vec<RustPipelineStage<serde_json::Value>> = stages
1480+
.into_iter()
1481+
.map(|s| {
1482+
let stage = std::sync::Arc::new(PythonPipelineStage {
1483+
callback: s.unbind(),
1484+
});
1485+
let ps: RustPipelineStage<serde_json::Value> =
1486+
std::sync::Arc::new(move |prev, item| stage.invoke(prev, item));
1487+
ps
1488+
})
1489+
.collect();
1490+
1491+
let session = self.inner.clone();
1492+
let outcomes = py.allow_threads(move || {
1493+
get_runtime().block_on(async move {
1494+
let executor = session.agent_executor();
1495+
execute_pipeline(executor, rust_items, rust_stages, None).await
1496+
})
1497+
});
1498+
1499+
outcomes
1500+
.iter()
1501+
.map(|o| match o {
1502+
Some(outcome) => step_outcome_to_py(py, outcome).map(Some),
1503+
None => Ok(None),
1504+
})
1505+
.collect()
1506+
}
1507+
13991508
/// Send a prompt or request and get a streaming iterator of events.
14001509
///
14011510
/// When ``history`` is omitted, session history and verification evidence are
@@ -3205,6 +3314,84 @@ fn parse_py_hook_response(
32053314
Ok(RustHookResponse::continue_())
32063315
}
32073316

3317+
// ============================================================================
3318+
// Orchestration: Python <-> Rust conversion + pipeline-stage bridge
3319+
// ============================================================================
3320+
3321+
fn py_dumps(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult<String> {
3322+
let json_mod = py.import("json")?;
3323+
json_mod.call_method1("dumps", (obj,))?.extract()
3324+
}
3325+
3326+
/// Convert a Python spec dict into an `AgentStepSpec` (snake_case keys) via a
3327+
/// JSON round-trip.
3328+
fn py_to_step_spec(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult<RustAgentStepSpec> {
3329+
serde_json::from_str(&py_dumps(py, obj)?)
3330+
.map_err(|e| PyValueError::new_err(format!("invalid AgentStepSpec: {e}")))
3331+
}
3332+
3333+
/// Convert an arbitrary Python value into a `serde_json::Value`.
3334+
fn py_to_json_value(py: Python<'_>, obj: &Bound<'_, PyAny>) -> PyResult<serde_json::Value> {
3335+
serde_json::from_str(&py_dumps(py, obj)?)
3336+
.map_err(|e| PyValueError::new_err(format!("invalid JSON: {e}")))
3337+
}
3338+
3339+
/// Convert a `StepOutcome` into a Python dict.
3340+
fn step_outcome_to_py(py: Python<'_>, outcome: &RustStepOutcome) -> PyResult<PyObject> {
3341+
let json = serde_json::to_string(outcome)
3342+
.map_err(|e| PyRuntimeError::new_err(format!("serialize outcome: {e}")))?;
3343+
json_string_to_py(py, &json)
3344+
}
3345+
3346+
/// Bridges a Python pipeline-stage callable into a synchronous `PipelineStage`.
3347+
///
3348+
/// GIL safety: `pipeline()` releases the GIL via `py.allow_threads`, so
3349+
/// re-acquiring it here from a tokio worker thread does not deadlock (same as
3350+
/// the hook/budget bridges). A raised exception is caught and treated as
3351+
/// `None` (stop the chain).
3352+
struct PythonPipelineStage {
3353+
callback: pyo3::Py<pyo3::PyAny>,
3354+
}
3355+
3356+
impl PythonPipelineStage {
3357+
fn invoke(
3358+
&self,
3359+
prev: Option<&RustStepOutcome>,
3360+
item: &serde_json::Value,
3361+
) -> Option<RustAgentStepSpec> {
3362+
pyo3::Python::with_gil(|py| {
3363+
let result = (|| -> PyResult<Option<RustAgentStepSpec>> {
3364+
let json_mod = py.import("json")?;
3365+
let previous = match prev {
3366+
Some(o) => {
3367+
let s = serde_json::to_string(o)
3368+
.map_err(|e| PyValueError::new_err(e.to_string()))?;
3369+
json_mod.call_method1("loads", (s,))?
3370+
}
3371+
None => py.None().into_bound(py),
3372+
};
3373+
let item_str = serde_json::to_string(item)
3374+
.map_err(|e| PyValueError::new_err(e.to_string()))?;
3375+
let item_py = json_mod.call_method1("loads", (item_str,))?;
3376+
let ctx = PyDict::new(py);
3377+
ctx.set_item("previous", previous)?;
3378+
ctx.set_item("item", item_py)?;
3379+
let ret = self.callback.call1(py, (ctx,))?;
3380+
let bound = ret.bind(py);
3381+
if bound.is_none() {
3382+
return Ok(None);
3383+
}
3384+
let spec_json: String = json_mod.call_method1("dumps", (bound,))?.extract()?;
3385+
serde_json::from_str::<RustAgentStepSpec>(&spec_json)
3386+
.map(Some)
3387+
.map_err(|e| PyValueError::new_err(format!("invalid step spec: {e}")))
3388+
})();
3389+
// Fail-closed: any exception → stop this chain.
3390+
result.unwrap_or(None)
3391+
})
3392+
}
3393+
}
3394+
32083395
// ============================================================================
32093396
// Python BudgetGuard bridge
32103397
// ============================================================================

0 commit comments

Comments
 (0)