fix(client): harden artifacts.get and follow Modal 303 long requests

tests/test_artifacts.py

@@ -1,6 +1,10 @@
 """Tests for artifacts operations."""
 
+import io
+import warnings
+
 import pytest
+from PIL import Image
 
 from vlmrun.client.types import ArtifactListResponse
 
@@ -83,3 +87,75 @@ def test_list_artifacts_rejects_both_ids(mock_client):
             session_id="sess-123",
             execution_id="exec-456",
         )
+
+
+@pytest.mark.parametrize(
+    "raw,expected_type,expected_id",
+    [
+        ("img_a1b2c3", "img", "img_a1b2c3"),
+        ("img_4c129a.jpg", "img", "img_4c129a"),
+        ("vid_4d0e56.mp4", "vid", "vid_4d0e56"),
+        ("doc_abcdef.PDF", "doc", "doc_abcdef"),
+        ("IMG_ABCDEF.PNG", "img", "img_abcdef"),
+    ],
+)
+def test_normalize_object_id(raw, expected_type, expected_id):
+    from vlmrun.client.artifacts import normalize_object_id
+
+    obj_type, normalized = normalize_object_id(raw)
+    assert obj_type == expected_type
+    assert normalized == expected_id
+
+
+def test_normalize_object_id_invalid():
+    from vlmrun.client.artifacts import normalize_object_id
+
+    with pytest.raises(ValueError):
+        normalize_object_id("not-an-id")
+
+
+def test_artifacts_get_strips_extension_and_accepts_octet_stream(monkeypatch, tmp_path):
+    """artifacts.get should strip .jpg and not crash on application/octet-stream."""
+    from vlmrun.client import artifacts as artifacts_mod
+    from vlmrun.client.artifacts import Artifacts
+
+    class FakeClient:
+        api_key = "test"
+        base_url = "https://agent.vlm.run/v1"
+        max_retries = 1
+
+    buf = io.BytesIO()
+    Image.new("RGB", (2, 2), color=(0, 0, 255)).save(buf, format="PNG")
+    png_bytes = buf.getvalue()
+
+    def fake_request(self, method, url, params=None, raw_response=False, **kwargs):
+        assert params["object_id"] == "img_4c129a"
+        return png_bytes, 200, {"Content-Type": "application/octet-stream"}
+
+    monkeypatch.setattr(artifacts_mod.APIRequestor, "request", fake_request)
+    monkeypatch.setattr(artifacts_mod, "VLMRUN_ARTIFACTS_CACHE_DIR", tmp_path)
+
+    art = Artifacts(FakeClient())
+    with warnings.catch_warnings(record=True):
+        warnings.simplefilter("always")
+        image = art.get(
+            object_id="img_4c129a.jpg",
+            session_id="550e8400-e29b-41d4-a716-446655440000",
+        )
+    assert isinstance(image, Image.Image)
+    assert image.size == (2, 2)
+
+
+def test_long_request_pending_and_result_helpers():
+    from vlmrun.client.long_request import is_pending_status, is_result_status
+
+    assert is_pending_status(204)
+    assert is_pending_status(303)
+    assert is_pending_status(500)
+    assert is_pending_status(0)
+    assert not is_pending_status(400)
+
+    assert is_result_status(201, b'{"id":"x"}')
+    assert is_result_status(200, b'{"id":"x"}')
+    assert not is_result_status(201, b"")
+    assert not is_result_status(204, b"")

vlmrun/client/agent.py

@@ -26,33 +26,124 @@
 # via `extra_body`.
 _VLM_EXTRA_KEYS: frozenset[str] = frozenset({"skills", "toolsets", "models"})
 
-
-def _patch_create(create_fn: Any) -> Any:
-    """Wrap an OpenAI ``create`` callable to accept VLM Run-specific kwargs.
+# Default wall-clock budget for following a Modal 303 long-request poll URL.
+_DEFAULT_LONG_REQUEST_TIMEOUT = 900.0
+
+
+def _pop_vlm_extra_kwargs(kwargs: dict[str, Any]) -> None:
+    """Move VLM-specific kwargs into ``extra_body`` in-place."""
+    vlm_kwargs = {k: kwargs.pop(k) for k in _VLM_EXTRA_KEYS if k in kwargs}
+    if vlm_kwargs:
+        kwargs["extra_body"] = {**(kwargs.get("extra_body") or {}), **vlm_kwargs}
+
+
+def _status_code_from_exc(exc: BaseException) -> int | None:
+    """Best-effort extract of an HTTP status code from an OpenAI/httpx error."""
+    status = getattr(exc, "status_code", None)
+    if isinstance(status, int):
+        return status
+    response = getattr(exc, "response", None)
+    if response is not None:
+        status = getattr(response, "status_code", None)
+        if isinstance(status, int):
+            return status
+    return None
+
+
+def _headers_from_exc(exc: BaseException) -> dict[str, Any]:
+    """Best-effort extract of response headers from an OpenAI/httpx error."""
+    response = getattr(exc, "response", None)
+    if response is None:
+        return {}
+    headers = getattr(response, "headers", None) or {}
+    try:
+        return dict(headers)
+    except Exception:
+        return {}
+
+
+def _parse_chat_completion(body: bytes) -> Any:
+    """Parse poll-response bytes into an OpenAI ChatCompletion object."""
+    import json
+
+    from openai.types.chat import ChatCompletion
+
+    payload = json.loads(body)
+    return ChatCompletion.model_validate(payload)
+
+
+def _follow_long_request_303(
+    exc: BaseException,
+    *,
+    api_key: str | None,
+    timeout: float,
+) -> Any:
+    """Poll the Modal ``Location`` URL from a 303 and return the ChatCompletion."""
+    from vlmrun.client.long_request import (
+        extract_location,
+        poll_location,
+    )
+
+    location = extract_location(_headers_from_exc(exc))
+    if not location:
+        raise exc
+
+    headers: dict[str, str] = {}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+
+    body, _status, _resp_headers = poll_location(
+        location,
+        headers=headers,
+        timeout=timeout,
+    )
+    return _parse_chat_completion(body)
+
+
+def _patch_create(create_fn: Any, *, api_key: str | None, timeout: float) -> Any:
+    """Wrap an OpenAI ``create`` callable for VLM kwargs + Modal 303 polling.
 
     ``skills``, ``toolsets``, and ``models`` are popped from ``kwargs`` and
-    merged into ``extra_body`` before the underlying call is made.
+    merged into ``extra_body`` before the underlying call is made. If the
+    gateway returns ``303 See Other`` (long request), the Location URL is
+    polled until the chat completion is ready.
     """
 
     @functools.wraps(create_fn)
     def _create(*args: Any, **kwargs: Any) -> Any:
-        vlm_kwargs = {k: kwargs.pop(k) for k in _VLM_EXTRA_KEYS if k in kwargs}
-        if vlm_kwargs:
-            kwargs["extra_body"] = {**(kwargs.get("extra_body") or {}), **vlm_kwargs}
-        return create_fn(*args, **kwargs)
+        _pop_vlm_extra_kwargs(kwargs)
+        try:
+            return create_fn(*args, **kwargs)
+        except Exception as exc:
+            if _status_code_from_exc(exc) != 303:
+                raise
+            return _follow_long_request_303(
+                exc,
+                api_key=api_key,
+                timeout=max(float(timeout or 0), _DEFAULT_LONG_REQUEST_TIMEOUT),
+            )
 
     return _create
 
 
-def _patch_async_create(create_fn: Any) -> Any:
+def _patch_async_create(create_fn: Any, *, api_key: str | None, timeout: float) -> Any:
     """Async variant of :func:`_patch_create`."""
+    import asyncio
 
     @functools.wraps(create_fn)
     async def _create(*args: Any, **kwargs: Any) -> Any:
-        vlm_kwargs = {k: kwargs.pop(k) for k in _VLM_EXTRA_KEYS if k in kwargs}
-        if vlm_kwargs:
-            kwargs["extra_body"] = {**(kwargs.get("extra_body") or {}), **vlm_kwargs}
-        return await create_fn(*args, **kwargs)
+        _pop_vlm_extra_kwargs(kwargs)
+        try:
+            return await create_fn(*args, **kwargs)
+        except Exception as exc:
+            if _status_code_from_exc(exc) != 303:
+                raise
+            return await asyncio.to_thread(
+                _follow_long_request_303,
+                exc,
+                api_key=api_key,
+                timeout=max(float(timeout or 0), _DEFAULT_LONG_REQUEST_TIMEOUT),
+            )
 
     return _create
 
@@ -325,7 +416,11 @@ def completions(self):
         )
 
         completions = openai_client.chat.completions
-        completions.create = _patch_create(completions.create)
+        completions.create = _patch_create(
+            completions.create,
+            api_key=self._client.api_key,
+            timeout=float(self._client.timeout or _DEFAULT_LONG_REQUEST_TIMEOUT),
+        )
         return completions
 
     @cached_property
@@ -379,5 +474,9 @@ async def main():
         )
 
         completions = async_openai_client.chat.completions
-        completions.create = _patch_async_create(completions.create)
+        completions.create = _patch_async_create(
+            completions.create,
+            api_key=self._client.api_key,
+            timeout=float(self._client.timeout or _DEFAULT_LONG_REQUEST_TIMEOUT),
+        )
         return completions