diff --git a/packages/ai/src/ai/modules/task/task_server.py b/packages/ai/src/ai/modules/task/task_server.py index edcf4a070..3ba9645bc 100644 --- a/packages/ai/src/ai/modules/task/task_server.py +++ b/packages/ai/src/ai/modules/task/task_server.py @@ -1314,7 +1314,13 @@ async def stop_task(self, token: str): # Only terminate tasks that were launched or executed directly if control.launch_type in (LAUNCH_TYPE.LAUNCH, LAUNCH_TYPE.EXECUTE): + # Stop the inner task first so its state-transition monitor can + # still find this control in _task_control while the busy-wait + # in task_engine.stop_task() polls for COMPLETED/CANCELLED. await control.task.stop_task() + # Free the slot so a subsequent use() with the same token does + # not race against a phantom registry entry. + self._task_control.pop(token, None) self.debug_message(f'Task "{control.id}" stopped on request') except Exception as e: diff --git a/packages/client-typescript/tests/RocketRideClient.test.ts b/packages/client-typescript/tests/RocketRideClient.test.ts index 0a59d5f1d..798e970a0 100644 --- a/packages/client-typescript/tests/RocketRideClient.test.ts +++ b/packages/client-typescript/tests/RocketRideClient.test.ts @@ -174,7 +174,7 @@ describe('RocketRideClient Integration Tests', () => { expect(Object.values(TASK_STATE)).toContain(status!.state); await client.terminate(result.token); - }, 90000); + }, TEST_CONFIG.timeout); it( 'should terminate a pipeline',