Collect Foundry results returned as HTTP 202

changelog/395.added.md

@@ -0,0 +1 @@
+The Azure AI Foundry `TabPFNClassifier` and `TabPFNRegressor` gained `collect_timeout_s`, bounding how long `predict()` waits for a result the endpoint is still computing (default two hours). On expiry the fit continues server-side, so predicting again collects it.

changelog/395.fixed.md

@@ -0,0 +1 @@
+Thinking mode against an Azure AI Foundry endpoint now returns a result instead of failing. A fit that takes longer than the platform's request window is handed back to be collected, and the client collects it automatically rather than mistaking the hand-off for a prediction. Collecting re-sends only the fit's id and the test rows, so the training data is not uploaded again on each attempt.

src/tabpfn_client/foundry/estimator.py

@@ -11,10 +11,17 @@
 
 This client sends requests as `application/json` only (Foundry also
 accepts `multipart/form-data`, but we don't use it here).
+
+A fit that outlives the platform's request window answers HTTP 202 naming
+the fit rather than the prediction; `predict*` collects it by re-sending
+that id until the result is ready, so a long Thinking fit completes in one
+call from the caller's point of view. Only the id and the test rows go on
+the wire for those follow-ups, not the training data again.
 """
 
 from __future__ import annotations
 
+import time
 from typing import Any, Dict, Literal, Optional, Union, cast
 
 import httpx
@@ -97,6 +104,28 @@ def _raise_for_status(resp: httpx.Response) -> None:
     )
 
 
+#: Fallback pause between collect attempts when the endpoint sends no
+#: `Retry-After`. Short enough to feel responsive, long enough that polling a
+#: multi-minute fit does not hammer the endpoint.
+_DEFAULT_RETRY_AFTER_S = 10.0
+
+
+def _retry_after_s(resp: httpx.Response) -> float:
+    """How long the endpoint asked us to wait, in seconds.
+
+    Only the delta-seconds form is honoured; an HTTP-date, or anything
+    unparseable, falls back to the default rather than failing a collect that
+    is otherwise progressing.
+    """
+    raw = resp.headers.get("Retry-After")
+    if raw is None:
+        return _DEFAULT_RETRY_AFTER_S
+    try:
+        return max(0.0, float(raw))
+    except ValueError:
+        return _DEFAULT_RETRY_AFTER_S
+
+
 def _to_jsonable(X: Any) -> list:
     """Coerce numpy / pandas inputs to plain Python lists for JSON."""
     if isinstance(X, pd.DataFrame):
@@ -201,6 +230,7 @@ def __init__(
         use_kv_cache: bool = False,
         fit_mode: Optional[FitModeLiteral] = None,
         timeout_s: float = 300.0,
+        collect_timeout_s: float = 7200.0,
     ):
         self.endpoint_url = endpoint_url
         self.api_key = api_key
@@ -221,6 +251,7 @@ def __init__(
         self.use_kv_cache = use_kv_cache
         self.fit_mode = fit_mode
         self.timeout_s = timeout_s
+        self.collect_timeout_s = collect_timeout_s
         self._validate_args()
 
     def _validate_args(self) -> None:
@@ -265,6 +296,13 @@ def _validate_args(self) -> None:
                 f"limit); got {self.thinking_timeout_s!r}."
             )
 
+        if self.collect_timeout_s < 0:
+            raise ValueError(
+                f"collect_timeout_s must be >= 0 (0 does not wait at all: a "
+                f"result that is not ready raises immediately); got "
+                f"{self.collect_timeout_s!r}."
+            )
+
     @property
     def _thinking_active(self) -> bool:
         return self.thinking_mode or self.thinking_effort is not None
@@ -385,16 +423,52 @@ def _invoke(
             cached_model_id=self._cached_model_id if self._cache_active else None,
             thinking_block=self._build_thinking_block(),
         )
-        resp = self._http_client().post(
-            self.endpoint_url,
-            json=body,
-            headers=self._headers(),
-        )
-        _raise_for_status(resp)
-        payload = resp.json()
-        if self._cache_active:
-            self._cached_model_id = payload.get("model_id") or self._cached_model_id
-        return payload
+        deadline = time.monotonic() + self.collect_timeout_s
+        while True:
+            resp = self._http_client().post(
+                self.endpoint_url,
+                json=body,
+                headers=self._headers(),
+            )
+            _raise_for_status(resp)
+            payload = resp.json()
+            if self._cache_active:
+                self._cached_model_id = payload.get("model_id") or self._cached_model_id
+
+            if resp.status_code != 202:
+                return payload
+
+            # 202 means the work outlived the request window and is still
+            # running server-side. It names the fit, so the follow-up carries
+            # that id and the test rows alone — re-sending the training data
+            # would cost the whole payload again on every attempt.
+            model_id = payload.get("model_id")
+            if model_id is None:
+                raise FoundryEndpointError(
+                    "Endpoint answered HTTP 202 without a model_id, so the "
+                    "result cannot be collected.",
+                    request=resp.request,
+                    response=resp,
+                    error_code=payload.get("error_code"),
+                    trace_id=payload.get("trace_id"),
+                )
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                raise TimeoutError(
+                    f"Endpoint still working after collect_timeout_s="
+                    f"{self.collect_timeout_s:g}s. The fit continues server-side; "
+                    f"predict again with the same data to collect it, or raise "
+                    f"collect_timeout_s."
+                )
+            body = _build_request_body(
+                task=self._task,
+                tabpfn_config=self._build_tabpfn_config(),
+                predict_params=params,
+                X_test=X_test,
+                cached_model_id=model_id,
+                thinking_block=self._build_thinking_block(),
+            )
+            time.sleep(min(_retry_after_s(resp), remaining))
 
 
 class TabPFNClassifier(_FoundryBase, ClassifierMixin):

tests/unit/test_foundry_collect.py

@@ -0,0 +1,139 @@
+#  Copyright (c) Prior Labs GmbH 2026.
+#  Licensed under the Apache License, Version 2.0
+"""Collecting a result the endpoint could not return in one request.
+
+A fit that outlives the platform's request window answers HTTP 202 naming the
+fit instead of the prediction. `predict*` collects it by re-sending that id,
+so the caller still makes a single call — and the follow-ups carry only the id
+and the test rows, never the training data again, which at a large context is
+the whole payload on every attempt.
+"""
+
+import json
+
+import httpx
+import numpy as np
+import pytest
+import respx
+
+from tabpfn_client.foundry import FoundryEndpointError, TabPFNClassifier
+from tabpfn_client.foundry.estimator import _DEFAULT_RETRY_AFTER_S, _retry_after_s
+
+
+URL = "https://example.inference.ml.azure.com/invocations"
+KEY = "test-key"
+
+
+def _fitted(**kwargs):
+    """A thinking classifier with training data attached, plus that data."""
+    clf = TabPFNClassifier(
+        endpoint_url=URL, api_key=KEY, thinking_effort="medium", **kwargs
+    )
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(8, 3))
+    y = (X[:, 0] > 0).astype(int)
+    clf.fit(X, y)
+    return clf, X
+
+
+def _in_progress(model_id="mid-1", retry_after="0"):
+    return httpx.Response(
+        202,
+        headers={"Retry-After": retry_after},
+        json={
+            "message": "The thinking fit for this dataset is still running.",
+            "error_code": "FIT_IN_PROGRESS",
+            "model_id": model_id,
+        },
+    )
+
+
+def _done(model_id="mid-1"):
+    return httpx.Response(
+        200, json={"prediction": [0, 1], "metadata": {}, "model_id": model_id}
+    )
+
+
+def _body(call):
+    return json.loads(call.request.content)
+
+
+class TestCollect:
+    @respx.mock
+    def test_collects_the_result_after_a_202(self):
+        route = respx.post(URL).mock(side_effect=[_in_progress(), _done()])
+        clf, X = _fitted()
+
+        assert list(clf.predict(X[:2])) == [0, 1]
+        assert route.call_count == 2
+
+    @respx.mock
+    def test_resend_drops_the_training_data(self):
+        """The reason the id is worth returning at all."""
+        route = respx.post(URL).mock(side_effect=[_in_progress(), _done()])
+        clf, X = _fitted()
+        clf.predict(X[:2])
+
+        first, second = _body(route.calls[0]), _body(route.calls[1])
+        assert "x_train" in first and "y_train" in first
+        assert "x_train" not in second and "y_train" not in second
+        assert second["context"] == {"model_id": "mid-1"}
+        assert "x_test" in second
+
+    @respx.mock
+    def test_several_202s_are_all_collected(self):
+        route = respx.post(URL).mock(
+            side_effect=[_in_progress(), _in_progress(), _in_progress(), _done()]
+        )
+        clf, X = _fitted()
+
+        assert list(clf.predict(X[:2])) == [0, 1]
+        assert route.call_count == 4
+
+    @respx.mock
+    def test_202_without_a_model_id_is_reported(self):
+        """Nothing to collect with, so say so rather than looping forever."""
+        respx.post(URL).mock(
+            return_value=httpx.Response(
+                202, json={"message": "still running", "error_code": "FIT_IN_PROGRESS"}
+            )
+        )
+        clf, X = _fitted()
+
+        with pytest.raises(FoundryEndpointError, match="without a model_id"):
+            clf.predict(X[:2])
+
+    @respx.mock
+    def test_gives_up_once_the_collect_budget_is_spent(self):
+        respx.post(URL).mock(return_value=_in_progress())
+        clf, X = _fitted(collect_timeout_s=0)
+
+        with pytest.raises(TimeoutError, match="collect_timeout_s"):
+            clf.predict(X[:2])
+
+    def test_rejects_a_negative_collect_budget(self):
+        with pytest.raises(ValueError, match="collect_timeout_s"):
+            TabPFNClassifier(endpoint_url=URL, api_key=KEY, collect_timeout_s=-1)
+
+
+class TestRetryAfter:
+    @staticmethod
+    def _resp(**kwargs):
+        return httpx.Response(202, request=httpx.Request("POST", URL), **kwargs)
+
+    def test_defaults_when_the_endpoint_says_nothing(self):
+        assert _retry_after_s(self._resp()) == _DEFAULT_RETRY_AFTER_S
+
+    def test_honours_delta_seconds(self):
+        assert _retry_after_s(self._resp(headers={"Retry-After": "3"})) == 3.0
+
+    def test_falls_back_on_an_http_date(self):
+        # Only the delta-seconds form is honoured; an unparseable value must not
+        # fail a collect that is otherwise progressing.
+        stamp = "Wed, 21 Oct 2026 07:28:00 GMT"
+        assert _retry_after_s(self._resp(headers={"Retry-After": stamp})) == (
+            _DEFAULT_RETRY_AFTER_S
+        )
+
+    def test_clamps_a_negative_delta(self):
+        assert _retry_after_s(self._resp(headers={"Retry-After": "-5"})) == 0.0