apache · fat-catTW · Jul 6, 2026
@@ -177,7 +177,7 @@ def test_ti_state_transitions(sdk_client, task_instance_id):
     Test task instance state transition to terminal state.
     """
     console.print("[yellow]Testing state transition: RUNNING → FAILED...")
-    sdk_client.task_instances.finish(
+    sdk_client.task_instances.set_terminal_state(
         id=task_instance_id, state=TerminalStateNonSuccess.FAILED, when=utcnow(), rendered_map_index="-1"
     )
 

@@ -265,11 +265,12 @@ def start(self, id: uuid.UUID, pid: int, when: datetime) -> TIRunContext:
             raise
         return TIRunContext.model_validate_json(resp.read())
 
-    def finish(self, id: uuid.UUID, state: TerminalStateNonSuccess, when: datetime, rendered_map_index):
-        """Tell the API server that this TI has reached a terminal state."""
+    def set_terminal_state(
+        self, id: uuid.UUID, state: TerminalStateNonSuccess, when: datetime, rendered_map_index
+    ):
+        """Tell the API server that this TI has reached a non-success terminal state."""
         if state == TaskInstanceState.SUCCESS:
             raise ValueError("Logic error. SUCCESS state should call the `succeed` function instead")
-        # TODO: handle the naming better. finish sounds wrong as "even" deferred is essentially finishing.
         body = TITerminalStatePayload(
             end_date=when, state=TerminalStateNonSuccess(state), rendered_map_index=rendered_map_index
         )

@@ -1316,9 +1316,9 @@ class ActivitySubprocess(WatchedSubprocess):
     # this attribute stays set and the dispatcher in
     # `update_task_state_if_needed` re-issues the matching API call on
     # subprocess exit — re-attempting the original transition rather than
-    # falling back to `finish()`, which doesn't accept SUCCESS / DEFERRED /
-    # SERVER_TERMINATED on the server side. Cleared (and `_terminal_state`
-    # set) only after the API call returns successfully.
+    # falling back to the generic non-success terminal-state update, which
+    # does not accept SUCCESS / DEFERRED / SERVER_TERMINATED. Cleared
+    # (and `_terminal_state` set) only after the API call returns successfully.
     _pending_terminal_state_msg: (
         SucceedTask | RetryTask | DeferTask | RescheduleTask | AwaitInputTask | None
     ) = attrs.field(default=None, init=False)
@@ -1446,20 +1446,19 @@ def update_task_state_if_needed(self):
         # the original request. Re-issue the matching dedicated API call so
         # the server learns the terminal state we couldn't deliver earlier.
         # Without this recovery, a transient API failure during the direct
-        # call would leave the TI stuck RUNNING on the server — `finish()`
-        # cannot substitute because the server-side `finish` endpoint does
-        # not accept SUCCESS / DEFERRED / SERVER_TERMINATED transitions.
+        # call would leave the TI stuck RUNNING on the server. The generic
+        # non-success terminal-state update cannot substitute because its
+        # payload does not accept SUCCESS / DEFERRED / SERVER_TERMINATED.
         if self._pending_terminal_state_msg is not None:
             self._replay_pending_terminal_state_msg()
             return
 
-        # If the process has finished a non-directly-patched state (e.g.
-        # FAILED, UP_FOR_RETRY without RetryTask), `finish()` is the
-        # dedicated endpoint for those transitions. For states already in
-        # STATES_SENT_DIRECTLY whose direct API call succeeded, no further
-        # action is needed.
+        # If the process has finished in a non-success terminal state that was
+        # not sent directly, report it with `set_terminal_state()`. For states
+        # already in STATES_SENT_DIRECTLY whose direct API call succeeded, no
+        # further action is needed.
         if self.final_state not in STATES_SENT_DIRECTLY:
-            self.client.task_instances.finish(
+            self.client.task_instances.set_terminal_state(
                 id=self.id,
                 state=self.final_state,
                 when=datetime.now(tz=timezone.utc),
@@ -1695,7 +1694,7 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id:
         dump_opts: dict[str, bool] = {}
         if isinstance(msg, TaskState):
             # No direct API call here — the recovery path in
-            # `update_task_state_if_needed` will call `finish()` for
+            # `update_task_state_if_needed` will call `set_terminal_state()` for
             # non-direct states (FAILED, etc.) once the subprocess exits.
             self._terminal_state = msg.state
             self._task_end_time_monotonic = time.monotonic()

@@ -405,8 +405,8 @@ def handle_request(request: httpx.Request) -> httpx.Response:
     @pytest.mark.parametrize(
         "state", [state for state in TerminalTIState if state != TerminalTIState.SUCCESS]
     )
-    def test_task_instance_finish(self, state):
-        # Simulate a successful response from the server that finishes (moved to terminal state) a task
+    def test_task_instance_set_terminal_state(self, state):
+        # Simulate a successful response from the server that sets a non-success terminal state for a task
         ti_id = uuid6.uuid7()
 
         def handle_request(request: httpx.Request) -> httpx.Response:
@@ -421,7 +421,7 @@ def handle_request(request: httpx.Request) -> httpx.Response:
             return httpx.Response(status_code=400, json={"detail": "Bad Request"})
 
         client = make_client(transport=httpx.MockTransport(handle_request))
-        client.task_instances.finish(
+        client.task_instances.set_terminal_state(
             ti_id, state=state, when="2024-10-31T12:00:00Z", rendered_map_index="test"
         )
 

@@ -3233,9 +3233,9 @@ def test_update_task_state_replays_pending_terminal_state_call(
         self, watched_subprocess, mocker, msg, api_method, expected_state
     ):
         """If a direct terminal-state API call was attempted and raised, the
-        recovery dispatcher must re-issue the dedicated endpoint (not
-        `finish()`, which the server-side endpoint refuses for SUCCESS /
-        DEFERRED / SERVER_TERMINATED). Covers all four message types.
+        recovery dispatcher must re-issue the dedicated method, not the
+        generic non-success terminal-state update, which refuses SUCCESS /
+        DEFERRED / SERVER_TERMINATED. Covers all four message types.
         """
         watched_subprocess, _ = watched_subprocess
         watched_subprocess._exit_code = 0
@@ -3244,9 +3244,9 @@ def test_update_task_state_replays_pending_terminal_state_call(
 
         watched_subprocess.update_task_state_if_needed()
 
-        # Recovery re-issues the dedicated endpoint, NOT finish().
+        # Recovery re-issues the dedicated method, not the generic terminal-state update.
         getattr(watched_subprocess.client.task_instances, api_method).assert_called_once()
-        watched_subprocess.client.task_instances.finish.assert_not_called()
+        watched_subprocess.client.task_instances.set_terminal_state.assert_not_called()
         assert watched_subprocess._terminal_state == expected_state
         assert watched_subprocess._pending_terminal_state_msg is None
 
@@ -3260,7 +3260,7 @@ def test_update_task_state_no_recovery_without_pending_msg(self, watched_subproc
 
         watched_subprocess.update_task_state_if_needed()
 
-        watched_subprocess.client.task_instances.finish.assert_not_called()
+        watched_subprocess.client.task_instances.set_terminal_state.assert_not_called()
         watched_subprocess.client.task_instances.succeed.assert_not_called()
 
 
@@ -3378,7 +3378,7 @@ def execute(self, context: Context):
         )
 
         # Patch the API client used by InProcessTestSupervisor to return a predictable TI context
-        fake_task_instances = mock.MagicMock(spec_set=["start", "finish"])
+        fake_task_instances = mock.MagicMock(spec_set=["start", "set_terminal_state"])
         fake_task_instances.start.return_value = make_ti_context()
         fake_client = mock.MagicMock(spec_set=["task_instances"])
         fake_client.task_instances = fake_task_instances