fix(amber): pass unstamped boundary states through LoopEnd instead of consuming them

amber/src/main/python/core/runnables/main_loop.py

@@ -374,9 +374,23 @@ def _process_state_frame(self, frame: StateFrame) -> None:
             return
 
         if isinstance(executor, LoopEndOperator):
-            # Matching LoopEnd (in_counter == 0): it will consume this state
-            # and jump back. Remember which LoopStart to jump to (it rides
-            # the envelope) for complete()/_jump_to_loop_start.
+            if not frame.loop_start_id:
+                # An UNstamped counter-0 state at a LoopEnd is not the loop's
+                # own boundary state -- it was produced by a loop-body
+                # operator's produce_state_on_start/finish (a public API on
+                # both engine sides), which emits with the "no loop" envelope.
+                # A real loop state is always stamped: the matching LoopStart
+                # stamps its own id on every iteration's output state.
+                # Forward it downstream unchanged, skipping the operator, like
+                # any default pass-through: consuming it would clobber the
+                # captured back-jump id with "" and hand run_update a State
+                # with no `table` payload (#discussion_r3648708075).
+                self._emit_and_save_state(state, in_counter, frame.loop_start_id)
+                self._check_and_process_control()
+                return
+            # Matching LoopEnd (in_counter == 0, stamped): it will consume
+            # this state and jump back. Remember which LoopStart to jump to
+            # (it rides the envelope) for complete()/_jump_to_loop_start.
             self._loop_start_id = frame.loop_start_id
 
         self.context.state_processing_manager.current_input_state = state

amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndChannelHandler.scala

@@ -43,6 +43,12 @@ trait EndChannelHandler {
     try {
       val outputState = dp.executor.produceStateOnFinish(portId.id)
       if (outputState.isDefined) {
+        // Deliberate "no loop" envelope defaults (loopCounter = 0,
+        // loopStartId = ""): this is operator-ORIGINATED boundary state, not
+        // a forwarded loop state, so it carries no LoopStart stamp. The
+        // Python LoopEnd runtime keys on that missing stamp to pass such
+        // states through instead of consuming them (see
+        // main_loop._process_state_frame).
         dp.outputManager.emitState(outputState.get)
       }
       dp.outputManager.outputIterator.setTupleOutput(

amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartChannelHandler.scala

@@ -42,6 +42,12 @@ trait StartChannelHandler {
     try {
       val outputState = dp.executor.produceStateOnStart(portId.id)
       if (outputState.isDefined) {
+        // Deliberate "no loop" envelope defaults (loopCounter = 0,
+        // loopStartId = ""): this is operator-ORIGINATED boundary state, not
+        // a forwarded loop state, so it carries no LoopStart stamp. The
+        // Python LoopEnd runtime keys on that missing stamp to pass such
+        // states through instead of consuming them (see
+        // main_loop._process_state_frame).
         dp.outputManager.emitState(outputState.get)
       }
     } catch safely {

amber/src/test/integration/org/apache/texera/amber/engine/e2e/LoopIntegrationSpec.scala

@@ -44,6 +44,7 @@ import org.apache.texera.amber.operator.LogicalOp
 import org.apache.texera.amber.operator.limit.LimitOpDesc
 import org.apache.texera.amber.operator.loop.{LoopEndOpDesc, LoopStartOpDesc}
 import org.apache.texera.amber.operator.sleep.SleepOpDesc
+import org.apache.texera.amber.operator.udf.python.PythonUDFOpDescV2
 import org.apache.texera.amber.operator.source.scan.text.TextInputSourceOpDesc
 import org.apache.texera.amber.tags.IntegrationTest
 import org.apache.texera.workflow.LogicalLink
@@ -184,6 +185,29 @@ class LoopIntegrationSpec
   // any loop logic runs. Production workers launch with a real classpath,
   // so this is a harness limitation, not an engine one.
 
+  /** A pass-through Python UDF that ALSO emits its own boundary state via the
+    * public `produce_state_on_finish` API -- the state arrives downstream
+    * with the "no loop" envelope (counter 0, no LoopStart stamp).
+    */
+  private def statefulPythonUDF(): PythonUDFOpDescV2 = {
+    val op = new PythonUDFOpDescV2()
+    op.code = """from pytexera import *
+                |
+                |class ProcessTupleOperator(UDFOperatorV2):
+                |
+                |    @overrides
+                |    def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]:
+                |        yield tuple_
+                |
+                |    @overrides
+                |    def produce_state_on_finish(self, port: int) -> Optional[State]:
+                |        return State({"note": "from-body-op"})
+                |""".stripMargin
+    op.workers = 1
+    op.retainInputColumns = true
+    op
+  }
+
   private def link(from: LogicalOp, to: LogicalOp): LogicalLink =
     LogicalLink(from.operatorIdentifier, PortIdentity(), to.operatorIdentifier, PortIdentity())
 
@@ -331,4 +355,31 @@ class LoopIntegrationSpec
     )
   }
 
+  it should "run a loop whose body operator emits its own boundary state" in {
+    // TextInput -> LoopStart -> stateful Python UDF -> LoopEnd.
+    //
+    // The UDF emits boundary state via produce_state_on_finish (a public API
+    // on both engine sides), which reaches the LoopEnd with the "no loop"
+    // envelope (counter 0, no LoopStart stamp) AFTER the forwarded loop
+    // state. The LoopEnd must pass that unstamped state through instead of
+    // consuming it: consuming would clobber the captured back-jump id with ""
+    // ("no loop-back state URI configured for LoopStart ''") and hand
+    // run_update a State with no `table` payload (KeyError). Regression test
+    // for #discussion_r3648708075 on #6661.
+    val src = textInput("1\n2\n3")
+    val start = loopStart("i = 0", "table.iloc[i]")
+    val mid = statefulPythonUDF()
+    val end = loopEnd("i += 1", "i < len(table)")
+    val materialized = runAndGetMaterializedRowCounts(
+      List(src, start, mid, end),
+      List(link(src, start), link(start, mid), link(mid, end))
+    )
+    val endRows = materialized.getOrElse(end.operatorIdentifier, -1L)
+    assert(
+      endRows == 3,
+      s"LoopEnd must accumulate all 3 iterations with a state-emitting " +
+        s"operator in the loop body: expected 3, got $endRows (all: $materialized)"
+    )
+  }
+
 }

amber/src/test/python/core/runnables/test_main_loop.py

@@ -2495,6 +2495,44 @@ def test_loopend_consume_invokes_operator_at_counter_zero(
         assert "loop_start_id" not in passed_to_operator
         assert "loop_counter" not in passed_to_operator
 
+    def test_loopend_forwards_unstamped_state_without_consuming(
+        self, main_loop, monkeypatch
+    ):
+        # Reviewer feedback (#discussion_r3648708075): a loop-body operator
+        # that emits its own boundary state (produce_state_on_start/finish --
+        # a public API on both engine sides) sends it with the "no loop"
+        # envelope (counter 0, id ""). That state is NOT the loop's own
+        # boundary state: consuming it would clobber the captured back-jump id
+        # with "" and hand run_update a State with no `table` payload
+        # (KeyError). A real loop state is always stamped -- the matching
+        # LoopStart stamps its own id on every iteration's output -- so an
+        # UNstamped counter-0 frame at a LoopEnd must be forwarded downstream
+        # unchanged, skipping the operator, like any default pass-through.
+        main_loop.context.executor_manager.executor = _FalseLoopEnd()
+        emitted, switched, reset_calls = self._capture_state_emit(
+            main_loop, monkeypatch
+        )
+        # The loop's own state was already consumed and its id captured.
+        main_loop._loop_start_id = "loop-start-1"
+
+        main_loop._process_state_frame(
+            StateFrame(
+                State({"note": "from-body-op"}),
+                loop_counter=0,
+                loop_start_id="",
+            )
+        )
+
+        assert switched == [], "unstamped state must not invoke the operator"
+        assert reset_calls == [], "unstamped state must not reset output"
+        assert emitted == [(State({"note": "from-body-op"}), 0, "")], (
+            "unstamped state must forward downstream with its envelope "
+            f"unchanged; emitted: {emitted}"
+        )
+        assert main_loop._loop_start_id == "loop-start-1", (
+            "an unstamped state must not clobber the captured back-jump id"
+        )
+
     # ------------------------------------------------------------------ #
     # _jump_to_loop_start
     #