feat: use async task polling for image generation

Author: pengqun

src/modelscope_mcp_server/constants.py

@@ -11,3 +11,6 @@
 # Default timeout for requests
 DEFAULT_API_TIMEOUT_SECONDS = 5
 DEFAULT_IMAGE_GENERATION_TIMEOUT_SECONDS = 300
+
+# Task polling interval (seconds)
+DEFAULT_TASK_POLL_INTERVAL_SECONDS = 5

src/modelscope_mcp_server/settings.py

@@ -9,6 +9,7 @@
     DEFAULT_IMAGE_TO_IMAGE_MODEL,
     DEFAULT_MODELSCOPE_API_INFERENCE_DOMAIN,
     DEFAULT_MODELSCOPE_DOMAIN,
+    DEFAULT_TASK_POLL_INTERVAL_SECONDS,
     DEFAULT_TEXT_TO_IMAGE_MODEL,
 )
 
@@ -57,6 +58,12 @@ class Settings(BaseSettings):
         description="Default timeout for image generation requests",
     )
 
+    # Task polling
+    task_poll_interval_seconds: int = Field(
+        default=DEFAULT_TASK_POLL_INTERVAL_SECONDS,
+        description="Polling interval in seconds when waiting for async tasks",
+    )
+
     # Logging settings
     log_level: str = Field(default="INFO", description="Logging level")
 

src/modelscope_mcp_server/tools/aigc.py

@@ -3,6 +3,7 @@
 Provides MCP tools for text-to-image generation, etc.
 """
 
+import time
 from typing import Annotated
 
 from fastmcp import FastMCP
@@ -77,7 +78,8 @@ async def generate_image(
         if not settings.is_api_token_configured():
             raise ValueError("API token is not set")
 
-        url = f"{settings.api_inference_domain}/v1/images/generations"
+        # Step 1: submit async generation task
+        submit_url = f"{settings.api_inference_domain}/v1/images/generations"
 
         payload = {
             "model": model,
@@ -87,21 +89,44 @@ async def generate_image(
         if generation_type == GenerationType.IMAGE_TO_IMAGE and image_url:
             payload["image_url"] = image_url
 
-        response = default_client.post(
-            url, json_data=payload, timeout=settings.default_image_generation_timeout_seconds
+        submit_response = default_client.post(
+            submit_url,
+            json_data=payload,
+            timeout=settings.default_image_generation_timeout_seconds,
+            headers={"X-ModelScope-Async-Mode": "true"},
         )
 
-        images_data = response.get("images", [])
-
-        if len(images_data) == 0:
-            raise Exception(f"No images found in response: {response}")
-
-        generated_image_url = images_data[0].get("url", "")
-        if len(generated_image_url) == 0:
-            raise Exception(f"No image URL found in response: {response}")
+        task_id = submit_response.get("task_id")
+        if not task_id:
+            raise Exception(f"No task_id found in response: {submit_response}")
+
+        # 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:
+            # timeout check
+            if time.time() - start_time > settings.default_image_generation_timeout_seconds:
+                raise TimeoutError("Image generation timed out - please try again later")
+
+            task_result = default_client.get(
+                task_url,
+                timeout=settings.default_api_timeout_seconds,
+                headers={"X-ModelScope-Task-Type": "image_generation"},
+            )
 
-        return ImageGenerationResult(
-            type=generation_type,
-            model=model,
-            image_url=generated_image_url,
-        )
+            status = task_result.get("task_status")
+            if status == "SUCCEED":
+                output_images = task_result.get("output_images") or []
+                if not output_images:
+                    raise Exception(f"No output images found in task result: {task_result}")
+                generated_image_url = output_images[0]
+                return ImageGenerationResult(
+                    type=generation_type,
+                    model=model,
+                    image_url=generated_image_url,
+                )
+            if status == "FAILED":
+                raise Exception("Image Generation Failed.")
+
+            logger.info(f"Image generation task {task_id} is {status}, waiting for next poll...")
+            time.sleep(settings.task_poll_interval_seconds)

tests/tools/test_generate_image.py

@@ -7,9 +7,18 @@
 
 
 async def test_text_to_image_generation_success(mcp_server, mocker):
-    """Test successful text-to-image generation."""
-    mock_response_data = {"images": [{"url": "https://example.com/generated_image.jpg"}]}
-    mock_post = mocker.patch("modelscope_mcp_server.client.default_client.post", return_value=mock_response_data)
+    """Test successful text-to-image generation with async polling."""
+    mock_post = mocker.patch(
+        "modelscope_mcp_server.client.default_client.post",
+        return_value={"task_id": "task-text-1"},
+    )
+    mock_get = mocker.patch(
+        "modelscope_mcp_server.client.default_client.get",
+        return_value={
+            "task_status": "SUCCEED",
+            "output_images": ["https://example.com/generated_image.jpg"],
+        },
+    )
 
     async with Client(mcp_server) as client:
         result = await client.call_tool(
@@ -32,12 +41,22 @@ async def test_text_to_image_generation_success(mcp_server, mocker):
         )
 
         mock_post.assert_called_once()
+        mock_get.assert_called()
 
 
 async def test_image_to_image_generation_success(mcp_server, mocker):
-    """Test successful image-to-image generation."""
-    mock_response_data = {"images": [{"url": "https://example.com/modified_image.jpg"}]}
-    mock_post = mocker.patch("modelscope_mcp_server.client.default_client.post", return_value=mock_response_data)
+    """Test successful image-to-image generation with async polling."""
+    mock_post = mocker.patch(
+        "modelscope_mcp_server.client.default_client.post",
+        return_value={"task_id": "task-image-1"},
+    )
+    mock_get = mocker.patch(
+        "modelscope_mcp_server.client.default_client.get",
+        return_value={
+            "task_status": "SUCCEED",
+            "output_images": ["https://example.com/modified_image.jpg"],
+        },
+    )
 
     async with Client(mcp_server) as client:
         result = await client.call_tool(
@@ -61,12 +80,22 @@ async def test_image_to_image_generation_success(mcp_server, mocker):
         )
 
         mock_post.assert_called_once()
+        mock_get.assert_called()
 
 
 async def test_generate_image_with_default_model(mcp_server, mocker):
-    """Test image generation with default model when no model is specified."""
-    mock_response_data = {"images": [{"url": "https://example.com/default_model_image.jpg"}]}
-    mocker.patch("modelscope_mcp_server.client.default_client.post", return_value=mock_response_data)
+    """Test image generation with default model when no model is specified (async)."""
+    mocker.patch(
+        "modelscope_mcp_server.client.default_client.post",
+        return_value={"task_id": "task-default-1"},
+    )
+    mocker.patch(
+        "modelscope_mcp_server.client.default_client.get",
+        return_value={
+            "task_status": "SUCCEED",
+            "output_images": ["https://example.com/default_model_image.jpg"],
+        },
+    )
 
     async with Client(mcp_server) as client:
         result = await client.call_tool(
@@ -141,10 +170,10 @@ async def test_generate_image_timeout_error(mcp_server, mocker):
 
 
 async def test_generate_image_malformed_response(mcp_server, mocker):
-    """Test handling of malformed API response."""
+    """Test handling of malformed API response on submit (missing task_id)."""
     malformed_response_data = {
         "result": "success",
-        # Missing 'images' field
+        # Missing 'task_id' field
     }
     mocker.patch("modelscope_mcp_server.client.default_client.post", return_value=malformed_response_data)
 
@@ -159,13 +188,22 @@ async def test_generate_image_malformed_response(mcp_server, mocker):
             )
 
         print(f"✅ Malformed response error handled correctly: {exc_info.value}")
-        assert "No images found in response" in str(exc_info.value)
+        assert "No task_id found in response" in str(exc_info.value)
 
 
 async def test_generate_image_request_parameters(mcp_server, mocker):
-    """Test that the correct parameters are sent in the request."""
-    mock_response_data = {"images": [{"url": "https://example.com/test_image.jpg"}]}
-    mock_post = mocker.patch("modelscope_mcp_server.client.default_client.post", return_value=mock_response_data)
+    """Test that the correct parameters are sent in the request (async)."""
+    mock_post = mocker.patch(
+        "modelscope_mcp_server.client.default_client.post",
+        return_value={"task_id": "task-param-1"},
+    )
+    mock_get = mocker.patch(
+        "modelscope_mcp_server.client.default_client.get",
+        return_value={
+            "task_status": "SUCCEED",
+            "output_images": ["https://example.com/test_image.jpg"],
+        },
+    )
 
     async with Client(mcp_server) as client:
         await client.call_tool(
@@ -177,7 +215,7 @@ async def test_generate_image_request_parameters(mcp_server, mocker):
             },
         )
 
-        # Verify the request was called with correct parameters
+        # Verify the submit request was called with correct parameters
         mock_post.assert_called_once()
         call_args = mock_post.call_args
 
@@ -194,4 +232,11 @@ async def test_generate_image_request_parameters(mcp_server, mocker):
         # Check timeout
         assert call_args.kwargs["timeout"] == 300
 
+        # Check headers include async mode
+        headers = call_args.kwargs.get("headers", {})
+        assert headers.get("X-ModelScope-Async-Mode") == "true"
+
+        # Verify polling was performed
+        mock_get.assert_called()
+
         print("✅ Request parameters verified correctly")