Honor |=> captures on composite THEN / IF / LOOP nodes (#163)

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: pinodeca <32303022+pinodeca@users.noreply.github.com>

Author: Copilot

src/orchestrations/execute_function_graph.rs

@@ -332,6 +332,19 @@ async fn execute_sql_node(
     Ok(result)
 }
 
+fn store_named_result(
+    ctx: &OrchestrationContext,
+    node: &FunctionNode,
+    result: &str,
+    results: &mut HashMap<String, String>,
+    node_label: &str,
+) {
+    if let Some(name) = &node.result_name {
+        ctx.trace_info(format!("Storing {node_label} result as ${name}"));
+        results.insert(name.clone(), result.to_string());
+    }
+}
+
 async fn execute_then_node(
     ctx: &OrchestrationContext,
     graph: &FunctionGraph,
@@ -364,6 +377,8 @@ async fn execute_then_node(
     ))
     .await?;
 
+    store_named_result(ctx, node, &right_result, results, "THEN");
+
     Ok(right_result)
 }
 
@@ -474,6 +489,7 @@ async fn execute_loop_node(
         ctx.trace_info(format!(
             "Loop terminated by break with value: {break_value}"
         ));
+        store_named_result(ctx, node, &break_value, results, "LOOP");
         return Ok(break_value);
     }
 
@@ -499,6 +515,7 @@ async fn execute_loop_node(
 
                 if !should_continue {
                     ctx.trace_info("Loop condition false, exiting loop");
+                    store_named_result(ctx, node, &body_result, results, "LOOP");
                     return Ok(body_result);
                 }
             }
@@ -639,15 +656,19 @@ async fn execute_if_node(
     ctx.trace_info(format!("Condition evaluated to: {is_true}"));
 
     if is_true {
-        Box::pin(execute_function_node_with_vars(
+        let result = Box::pin(execute_function_node_with_vars(
             ctx, graph, then_id, results, exec_ctx,
         ))
-        .await
+        .await?;
+        store_named_result(ctx, node, &result, results, "IF");
+        Ok(result)
     } else {
-        Box::pin(execute_function_node_with_vars(
+        let result = Box::pin(execute_function_node_with_vars(
             ctx, graph, else_id, results, exec_ctx,
         ))
-        .await
+        .await?;
+        store_named_result(ctx, node, &result, results, "IF");
+        Ok(result)
     }
 }
 

tests/e2e/sql/02_conditionals.sql

@@ -299,5 +299,46 @@ END $$;
 DROP TABLE _test_state2;
 DROP TABLE test_if_rows_dot;
 
+-- Test 4: named result on the IF node itself is accessible downstream
+DROP TABLE IF EXISTS test_if_named;
+CREATE TABLE test_if_named (id SERIAL, chosen INT);
+
+CREATE TEMP TABLE _test_if_named (instance_id TEXT);
+
+INSERT INTO _test_if_named SELECT df.start(
+    df.if(
+        'SELECT true',
+        'SELECT 41 AS chosen',
+        'SELECT 99 AS chosen'
+    ) |=> 'decision'
+    ~> $$INSERT INTO test_if_named (chosen) VALUES ($decision.chosen)$$,
+    'test-if-named-result'
+);
+
+DO $$
+DECLARE
+    inst_id TEXT;
+    status TEXT;
+    chosen_val INT;
+BEGIN
+    SELECT instance_id INTO inst_id FROM _test_if_named;
+    SELECT df.wait_for_completion(inst_id, 10) INTO status;
+
+    IF status != 'completed' THEN
+        RAISE EXCEPTION 'TEST FAILED [if-named]: status = %', status;
+    END IF;
+
+    SELECT chosen INTO chosen_val FROM test_if_named ORDER BY id DESC LIMIT 1;
+
+    IF chosen_val != 41 THEN
+        RAISE EXCEPTION 'TEST FAILED [if-named]: expected 41, got %', chosen_val;
+    END IF;
+
+    RAISE NOTICE 'PASSED: named IF result stored';
+END $$;
+
+DROP TABLE _test_if_named;
+DROP TABLE test_if_named;
+
 RESET SESSION AUTHORIZATION;
 SELECT 'TEST PASSED' AS result;

tests/e2e/sql/03_loops.sql

@@ -201,6 +201,44 @@ END $$;
 DROP TABLE _test3_state;
 DROP TABLE test_break_log;
 
+-- Test 4: named result on the LOOP node itself is accessible downstream
+DROP TABLE IF EXISTS test_loop_named;
+CREATE TABLE test_loop_named (id SERIAL, status TEXT);
+
+CREATE TEMP TABLE _test4_state AS
+SELECT df.start(
+    (df.loop(df.break('{"status": "done"}')) |=> 'loop_result')
+    ~> $$INSERT INTO test_loop_named (status) VALUES ($loop_result::jsonb->>'status')$$,
+    'test-loop-named-result'
+) AS instance_id;
+
+DO $$
+DECLARE
+    v_instance_id TEXT;
+    v_status TEXT;
+    v_loop_status TEXT;
+BEGIN
+    SELECT instance_id INTO v_instance_id FROM _test4_state;
+    RAISE NOTICE 'Test 4 - named loop result: instance %', v_instance_id;
+
+    SELECT df.wait_for_completion(v_instance_id, 20) INTO v_status;
+
+    IF v_status != 'completed' THEN
+        RAISE EXCEPTION 'TEST FAILED [loop-named]: expected Completed, got %', v_status;
+    END IF;
+
+    SELECT status INTO v_loop_status FROM test_loop_named ORDER BY id DESC LIMIT 1;
+
+    IF v_loop_status != 'done' THEN
+        RAISE EXCEPTION 'TEST FAILED [loop-named]: expected done, got %', v_loop_status;
+    END IF;
+
+    RAISE NOTICE 'PASSED: named LOOP result stored';
+END $$;
+
+DROP TABLE _test4_state;
+DROP TABLE test_loop_named;
+
 -- === Test: running_status_during_loop ===
 -- Verify that df.status() reports 'running' while a loop is actively executing
 -- (regression test for: loops reporting 'pending' instead of 'running')

tests/e2e/sql/07_signals.sql

@@ -224,6 +224,57 @@ BEGIN
 END $$;
 
 DROP TABLE _test_signal_text;
+DELETE FROM signal_test_log;
+
+-- Test 5: |=> applied after a THEN composite still captures the final result
+CREATE TEMP TABLE _test_signal_then_named (instance_id TEXT);
+
+INSERT INTO _test_signal_then_named SELECT df.start(
+    'INSERT INTO signal_test_log (msg) VALUES (''waiting-for-decision'')'
+    ~> df.wait_for_signal('test_approval_then', 60) |=> 'decision'
+    ~> df.if(
+        'SELECT ($decision::jsonb->''data''->>''approved'')::boolean',
+        $$INSERT INTO signal_test_log (msg, data) VALUES ('approved-after-then', $decision::jsonb)$$,
+        $$INSERT INTO signal_test_log (msg, data) VALUES ('rejected-after-then', $decision::jsonb)$$
+    ),
+    'test-signal-then-named'
+);
+
+SELECT pg_sleep(1);
+
+DO $$
+DECLARE
+    inst_id TEXT;
+BEGIN
+    SELECT instance_id INTO inst_id FROM _test_signal_then_named;
+    PERFORM df.signal(inst_id, 'test_approval_then', '{"approved": true}');
+END $$;
+
+DO $$
+DECLARE
+    inst_id TEXT;
+    status TEXT;
+BEGIN
+    SELECT instance_id INTO inst_id FROM _test_signal_then_named;
+    SELECT df.wait_for_completion(inst_id, 10) INTO status;
+
+    IF status != 'completed' THEN
+        RAISE EXCEPTION 'TEST FAILED: composite THEN capture status = %', status;
+    END IF;
+
+    IF NOT EXISTS (
+        SELECT 1
+        FROM signal_test_log
+        WHERE msg = 'approved-after-then'
+          AND (data->'data'->>'approved')::boolean = true
+    ) THEN
+        RAISE EXCEPTION 'TEST FAILED: THEN composite capture did not substitute $decision';
+    END IF;
+
+    RAISE NOTICE 'TEST PASSED: composite THEN capture';
+END $$;
+
+DROP TABLE _test_signal_then_named;
 DROP TABLE signal_test_log;
 
 RESET SESSION AUTHORIZATION;