feat: add maximum polling attempts for async tasks

Author: pengqun

src/modelscope_mcp_server/constants.py

@@ -14,3 +14,6 @@
 
 # Task polling interval (seconds)
 DEFAULT_TASK_POLL_INTERVAL_SECONDS = 5
+
+# Maximum number of polling attempts for async tasks
+DEFAULT_MAX_POLL_ATTEMPTS = 60  # 60 attempts * 5 seconds = 5 minutes max

src/modelscope_mcp_server/settings.py

@@ -7,6 +7,7 @@
     DEFAULT_API_TIMEOUT_SECONDS,
     DEFAULT_IMAGE_GENERATION_TIMEOUT_SECONDS,
     DEFAULT_IMAGE_TO_IMAGE_MODEL,
+    DEFAULT_MAX_POLL_ATTEMPTS,
     DEFAULT_MODELSCOPE_API_INFERENCE_DOMAIN,
     DEFAULT_MODELSCOPE_DOMAIN,
     DEFAULT_TASK_POLL_INTERVAL_SECONDS,
@@ -63,6 +64,10 @@ class Settings(BaseSettings):
         default=DEFAULT_TASK_POLL_INTERVAL_SECONDS,
         description="Polling interval in seconds when waiting for async tasks",
     )
+    max_poll_attempts: int = Field(
+        default=DEFAULT_MAX_POLL_ATTEMPTS,
+        description="Maximum number of polling attempts for async tasks",
+    )
 
     # Logging settings
     log_level: str = Field(default="INFO", description="Logging level")

src/modelscope_mcp_server/tools/aigc.py

@@ -105,11 +105,15 @@ async def generate_image(
         # Step 2: poll task result until succeed/failed or timeout
         start_time = time.time()
         task_url = f"{settings.api_inference_domain}/v1/tasks/{task_id}"
-        while True:
+        attempt_count = 0
+
+        while attempt_count < settings.max_poll_attempts:
             # timeout check
             if time.time() - start_time > settings.default_image_generation_timeout_seconds:
                 raise TimeoutError("Image generation timed out - please try again later")
 
+            attempt_count += 1
+
             task_result = default_client.get(
                 task_url,
                 timeout=settings.default_api_timeout_seconds,
@@ -138,5 +142,13 @@ async def generate_image(
                     f"Image generation failed: {error_msg}. Task ID: {task_id}, Request ID: {request_id}"
                 )
 
-            logger.info(f"Image generation task {task_id} is {status}, waiting for next poll...")
+            logger.info(
+                f"Image generation task {task_id} is {status}, waiting for next poll... "
+                f"(attempt {attempt_count}/{settings.max_poll_attempts})"
+            )
             time.sleep(settings.task_poll_interval_seconds)
+
+        # If we exit the loop without success or failure, max attempts exceeded
+        raise TimeoutError(
+            f"Image generation exceeded maximum polling attempts ({settings.max_poll_attempts}). Task ID: {task_id}"
+        )