@@ -49,6 +49,11 @@ use a3s_code_core::hooks::{
4949 HookMatcher as RustHookMatcher , HookResponse as RustHookResponse ,
5050} ;
5151use 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+ } ;
5257use 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