fix: handle_response_headers

Author: hanchchch

projects/fal_client/src/fal_client/_headers.py

@@ -1,6 +1,23 @@
 from __future__ import annotations
 
-from typing import Literal, Union, get_args
+from typing import Literal, Union, get_args, Optional, Any
+
+from httpx import Headers
+
+try:
+    from fal.ref import get_current_app
+except ImportError:
+
+    def get_current_app() -> Optional[Any]:
+        return None
+
+
+def _current_fal_app_request() -> Optional[Any]:
+    """Get the current request if we are running in a fal app."""
+    if (app := get_current_app()) is not None and app.current_request is not None:
+        return app.current_request
+    return None
+
 
 MIN_REQUEST_TIMEOUT_SECONDS = 1  # Minimum allowed request timeout in seconds
 
@@ -54,15 +71,13 @@ def add_priority_header(priority: Priority, headers: dict[str, str]) -> None:
     headers[QUEUE_PRIORITY_HEADER] = priority
 
 
-def add_forwarded_headers(request_headers, _headers: dict[str, str]) -> None:
-    if cdn_token := request_headers.get("x-fal-cdn-token"):
-        _headers["x-fal-forwarded-cdn-token"] = cdn_token
+def add_fal_app_context_headers(headers: dict[str, str]) -> None:
+    if request := _current_fal_app_request():
+        if cdn_token := request.headers.get("x-fal-cdn-token"):
+            headers["x-fal-cdn-token"] = cdn_token
 
-    forwarded_request_ids = []
-    if request_ids := request_headers.get("x-fal-forwarded-request-id"):
-        forwarded_request_ids.extend(request_ids.split(","))
-    if request_id := request_headers.get("x-fal-request-id"):
-        forwarded_request_ids.append(request_id)
 
-    if forwarded_request_ids:
-        _headers["x-fal-forwarded-request-id"] = ",".join(forwarded_request_ids)
+def handle_response_headers(response_headers: Headers) -> None:
+    if request := _current_fal_app_request():
+        if cdn_token := response_headers.get("x-fal-cdn-token"):
+            request.headers["x-fal-cdn-token"] = cdn_token

projects/fal_client/src/fal_client/client.py

@@ -49,9 +49,10 @@
     add_priority_header,
     add_timeout_header,
     add_hint_header,
+    add_fal_app_context_headers,
+    handle_response_headers,
     REQUEST_TIMEOUT_TYPE_HEADER,
     REQUEST_TIMEOUT_HEADER,
-    add_forwarded_headers,
 )
 
 if TYPE_CHECKING:
@@ -63,14 +64,6 @@
 if TYPE_CHECKING:
     from PIL import Image
 
-try:
-    from fal.ref import get_current_app
-except ImportError:
-
-    def get_current_app() -> Optional[Any]:
-        return None
-
-
 AnyJSON = Dict[str, Any]
 UploadRepositoryId = Literal["fal_v3", "cdn", "fal"]
 
@@ -1609,8 +1602,7 @@ async def run(
         if start_timeout is not None:
             add_timeout_header(start_timeout, _headers)
 
-        if (app := get_current_app()) is not None and app.current_request is not None:
-            add_forwarded_headers(app.current_request.headers, _headers)
+        add_fal_app_context_headers(_headers)
 
         response = await _async_maybe_retry_request(
             self._client,
@@ -1620,8 +1612,9 @@ async def run(
             timeout=timeout,
             headers=_headers,
         )
-
         _raise_for_status(response)
+        handle_response_headers(response.headers)
+
         return response.json()
 
     async def submit(
@@ -1665,8 +1658,7 @@ async def submit(
         if start_timeout is not None:
             add_timeout_header(start_timeout, _headers)
 
-        if (app := get_current_app()) is not None and app.current_request is not None:
-            add_forwarded_headers(app.current_request.headers, _headers)
+        add_fal_app_context_headers(_headers)
 
         response = await _async_maybe_retry_request(
             self._client,
@@ -1677,6 +1669,7 @@ async def submit(
             headers=_headers,
         )
         _raise_for_status(response)
+        handle_response_headers(response.headers)
 
         data = response.json()
         return AsyncRequestHandle(
@@ -2111,8 +2104,7 @@ def run(
         if start_timeout is not None:
             add_timeout_header(start_timeout, _headers)
 
-        if (app := get_current_app()) is not None and app.current_request is not None:
-            add_forwarded_headers(app.current_request.headers, _headers)
+        add_fal_app_context_headers(_headers)
 
         response = _maybe_retry_request(
             self._client,
@@ -2123,6 +2115,8 @@ def run(
             headers=_headers,
         )
         _raise_for_status(response)
+        handle_response_headers(response.headers)
+
         return response.json()
 
     def submit(
@@ -2163,8 +2157,7 @@ def submit(
         if start_timeout is not None:
             add_timeout_header(start_timeout, _headers)
 
-        if (app := get_current_app()) is not None and app.current_request is not None:
-            add_forwarded_headers(app.current_request.headers, _headers)
+        add_fal_app_context_headers(_headers)
 
         response = _maybe_retry_request(
             self._client,
@@ -2175,6 +2168,7 @@ def submit(
             headers=_headers,
         )
         _raise_for_status(response)
+        handle_response_headers(response.headers)
 
         data = response.json()
         return SyncRequestHandle(