timbal 2.1.7__tar.gz → 2.2.0__tar.gz

{timbal-2.1.7 → timbal-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: timbal
-Version: 2.1.7
+Version: 2.2.0
 Summary: Simple, performant, battle-tested framework for building reliable AI applications
 Project-URL: Source, https://github.com/timbal-ai/timbal
 Author-email: Timbal <team@timbal.ai>

{timbal-2.1.7 → timbal-2.2.0}/python/timbal/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '2.1.7'
-__version_tuple__ = version_tuple = (2, 1, 7)
+__version__ = version = '2.2.0'
+__version_tuple__ = version_tuple = (2, 2, 0)
 
 __commit_id__ = commit_id = None

{timbal-2.1.7 → timbal-2.2.0}/python/timbal/core/tool.py

@@ -11,6 +11,8 @@ except ImportError:
 
 from pydantic import BaseModel, Field, SkipValidation, computed_field, model_validator
 
+from ..errors import CredentialNotAvailable, PlatformError, ToolProxyUnavailable
+from ..platform.tool_proxy import execute_tool_proxy
 from ..utils import create_model_from_handler
 from .runnable import Runnable
 
@@ -119,6 +121,49 @@ class Tool(Runnable):
         self._is_async_gen = inspect_result["is_async_gen"]
         self._dependencies = inspect_result["dependencies"]
 
+    @override
+    async def _execute_handler(
+        self,
+        validated_input: dict[str, Any],
+        run_context: Any,
+        span: Any,
+        event_queue: Any = None,
+    ):
+        """Execute handler locally; on missing credentials, run via platform tool proxy.
+
+        If no proxy is reachable for this tool — the platform has no proxy for it
+        (HTTP 403, or 404/501 fallbacks), or there's no platform config at all
+        (e.g. running locally) — re-raise the original ``CredentialNotAvailable``
+        so the user gets the actionable "configure credentials locally" message
+        instead of an opaque platform error.
+
+        When platform auth is configured but org/subject is missing (e.g.
+        ``TIMBAL_API_KEY`` without ``TIMBAL_ORG_ID``), propagate
+        ``ToolProxyUnavailable`` so the user fixes platform setup rather than
+        local provider credentials.
+        """
+        try:
+            async for event, final_output, collector in super()._execute_handler(
+                validated_input, run_context, span, event_queue
+            ):
+                yield (event, final_output, collector)
+        except CredentialNotAvailable as cred_error:
+            try:
+                output = await execute_tool_proxy(self.name, validated_input)
+            except ToolProxyUnavailable:
+                if run_context.platform_config is None:
+                    # Local run without platform config — fall back to credential error.
+                    raise cred_error from None
+                raise
+            except PlatformError as proxy_error:
+                # 403 is what the platform returns when no proxy is available for this
+                # tool (no service-account credentials configured). 404/501 are kept as
+                # fallbacks. In all these cases, surface the actionable credential error.
+                if proxy_error.status_code in (403, 404, 501):
+                    raise cred_error from None
+                raise
+            yield (None, output, None)
+
     @override
     def nest(self, parent_path: str) -> None:
         """See base class."""

{timbal-2.1.7 → timbal-2.2.0}/python/timbal/errors.py

@@ -9,6 +9,38 @@ class APIKeyNotFoundError(TimbalError):
     """Error raised when an API key is not found."""
 
 
+class CredentialNotAvailable(APIKeyNotFoundError, ValueError):
+    """Raised when a tool credential can't be resolved from any source.
+
+    Subclasses ValueError for backward compat with tool code/tests. Carries
+    structured fields (provider_name, missing, env_vars) so callers (e.g. a
+    remote service-account executor) can react without parsing the message.
+    """
+
+    def __init__(
+        self,
+        provider_name: str,
+        *,
+        missing: list[str] | None = None,
+        env_vars: list[str] | None = None,
+        message: str | None = None,
+    ) -> None:
+        self.provider_name = provider_name
+        self.missing = list(missing or [])
+        self.env_vars = list(env_vars or [])
+        super().__init__(message or self._default_message())
+
+    def _default_message(self) -> str:
+        parts = [f"{self.provider_name} credentials not found."]
+        if self.missing:
+            parts.append(f"Missing: {', '.join(self.missing)}.")
+        hints = ["pass them on the tool", "configure an Integration"]
+        if self.env_vars:
+            hints.append(f"set the {', '.join(self.env_vars)} environment variable(s)")
+        parts.append("To fix: " + ", or ".join(hints) + ".")
+        return " ".join(parts)
+
+
 class EarlyExit(TimbalError):
     """Error raised when an early exit is requested.
 
@@ -133,7 +165,25 @@ class PDFProcessingError(TimbalError):
 
 
 class PlatformError(TimbalError):
-    """Error raised when a platform API call fails."""
+    """Error raised when a platform API call fails.
+
+    Carries the HTTP ``status_code`` (when the failure came from an HTTP
+    response) so callers can branch on it — e.g. treat 404/501 as "endpoint
+    not implemented" rather than a hard failure.
+    """
+
+    def __init__(self, message: str, *, status_code: int | None = None) -> None:
+        super().__init__(message)
+        self.status_code = status_code
+
+
+class ToolProxyUnavailable(TimbalError):
+    """Raised when the platform tool proxy can't be used for a tool execution.
+
+    Signals that no proxy is reachable for this run — e.g. no platform config
+    (running locally without TIMBAL_API_KEY / TIMBAL_ORG_ID). Callers fall back
+    to the original credential error so the user gets an actionable message.
+    """
 
 
 class FallbackExhausted(TimbalError):

timbal-2.2.0/python/timbal/platform/tool_proxy.py

@@ -0,0 +1,98 @@
+"""Platform tool proxy — remote tool execution without exporting credentials."""
+
+import os
+from typing import Any
+
+import structlog
+
+from .. import __version__
+from ..errors import ToolProxyUnavailable
+from ..state import get_call_id, get_or_create_run_context
+from .utils import _request
+
+logger = structlog.get_logger("timbal.platform.tool_proxy")
+
+_PROXY_PREFIX = "/proxies/v1/tools"
+
+
+def _subject_header_env_map() -> dict[str, str]:
+    """Map x-timbal-* header names to TIMBAL_* env vars (fallback when subject fields are unset)."""
+    return {
+        "x-timbal-app-id": "TIMBAL_APP_ID",
+        "x-timbal-project-id": "TIMBAL_PROJECT_ID",
+        "x-timbal-rev": "TIMBAL_REV",
+    }
+
+
+def build_tool_proxy_headers() -> dict[str, str]:
+    """Build per-request headers for tool proxy calls (run context + subject + version)."""
+    run_context = get_or_create_run_context()
+    headers: dict[str, str] = {
+        "x-timbal-run-id": run_context.id,
+        "x-timbal-version": __version__,
+    }
+
+    # Match llm_router: always attach call_id when the context var is set (tool runs always set it).
+    call_id = get_call_id()
+    if call_id is not None:
+        headers["x-timbal-call-id"] = call_id
+
+    platform_config = run_context.platform_config
+    if platform_config and platform_config.subject:
+        subject = platform_config.subject
+        if subject.app_id:
+            headers["x-timbal-app-id"] = subject.app_id
+        if subject.project_id:
+            headers["x-timbal-project-id"] = subject.project_id
+        if subject.rev:
+            headers["x-timbal-rev"] = subject.rev
+
+    for header_name, env_var in _subject_header_env_map().items():
+        if header_name not in headers:
+            value = os.getenv(env_var)
+            if value:
+                headers[header_name] = value
+
+    return headers
+
+
+def _tool_proxy_path(org_id: str, tool_name: str) -> str:
+    return f"orgs/{org_id}{_PROXY_PREFIX}/{tool_name}"
+
+
+async def execute_tool_proxy(tool_name: str, params: dict[str, Any]) -> Any:
+    """Execute a tool via the platform proxy (no credentials in the client runtime).
+
+    POST ``/orgs/{org_id}/proxies/v1/tools/{tool_name}`` with handler params as JSON body.
+    Returns the tool handler result JSON as-is (no response wrapping).
+    """
+    run_context = get_or_create_run_context()
+    platform_config = run_context.platform_config
+    if platform_config is None or platform_config.subject is None:
+        # The run context may have skipped platform_config resolution (e.g. tracing was
+        # explicitly disabled). Resolve it on demand — the proxy needs the org subject.
+        # force_refresh bypasses a stale cached default (a default call could have cached
+        # None before TIMBAL_API_KEY / TIMBAL_ORG_ID were set in this process).
+        from ..state.config_loader import resolve_platform_config
+
+        platform_config = resolve_platform_config(platform_config, force_refresh=True)
+        run_context.platform_config = platform_config
+
+    if platform_config is None or platform_config.subject is None:
+        raise ToolProxyUnavailable(
+            "Tool proxy requires platform_config with org subject. Set TIMBAL_API_KEY and TIMBAL_ORG_ID."
+        )
+
+    org_id = platform_config.subject.org_id
+    path = _tool_proxy_path(org_id, tool_name)
+    headers = build_tool_proxy_headers()
+
+    logger.debug("tool_proxy_execute", tool=tool_name, path=path)
+
+    response = await _request(
+        "POST",
+        path,
+        headers=headers,
+        json=params,
+    )
+    return response.json()

{timbal-2.1.7 → timbal-2.2.0}/python/timbal/platform/utils.py

@@ -134,7 +134,8 @@ async def _request(
                     f"\n"
                     f"  URL: {exc.request.url}\n"
                     f"  Status: {exc.response.status_code} {exc.response.reason_phrase}\n"
-                    f"  Response body: {error_body or None}"
+                    f"  Response body: {error_body or None}",
+                    status_code=exc.response.status_code,
                 ) from exc
             # Retry on 429, 5xx, or other errors
             if attempt == max_retries:
@@ -142,7 +143,8 @@ async def _request(
                     f"\n"
                     f"  URL: {exc.request.url}\n"
                     f"  Status: {exc.response.status_code} {exc.response.reason_phrase}\n"
-                    f"  Response body: {error_body or None}"
+                    f"  Response body: {error_body or None}",
+                    status_code=exc.response.status_code,
                 ) from exc
             wait_time = 0.1 * (2**attempt)  # Exponential backoff: 100ms, 200ms, 400ms
             if exc.response.status_code == 429:
@@ -236,7 +238,8 @@ async def _stream(
                     f"\n"
                     f"  URL: {exc.request.url}\n"
                     f"  Status: {exc.response.status_code} {exc.response.reason_phrase}\n"
-                    f"  Response body: {error_body or None}"
+                    f"  Response body: {error_body or None}",
+                    status_code=exc.response.status_code,
                 ) from exc
             # Retry on 429, 5xx, or other errors
             if attempt == max_retries:
@@ -244,7 +247,8 @@ async def _stream(
                     f"\n"
                     f"  URL: {exc.request.url}\n"
                     f"  Status: {exc.response.status_code} {exc.response.reason_phrase}\n"
-                    f"  Response body: {error_body or None}"
+                    f"  Response body: {error_body or None}",
+                    status_code=exc.response.status_code,
                 ) from exc
             wait_time = 0.1 * (2**attempt)
             if exc.response.status_code == 429:

{timbal-2.1.7 → timbal-2.2.0}/python/timbal/state/config.py

@@ -54,6 +54,8 @@ class PlatformSubject(BaseModel):
     """Organization identifier."""
     app_id: str | None = None
     """Application identifier. Either project or app must be specified."""
+    project_id: str | None = None
+    """Project identifier (platform project scope)."""
     rev: str | None = None
     """Git revision the run was executed against. Internal use."""
 

{timbal-2.1.7 → timbal-2.2.0}/python/timbal/state/config_loader.py

@@ -150,6 +150,7 @@ def resolve_platform_config(
     platform_config: PlatformConfig | None = None,
     profile: str | None = None,
     config_dir: Path | None = None,
+    force_refresh: bool = False,
 ) -> PlatformConfig | None:
     """Resolve platform configuration by merging sources with precedence.
 
@@ -170,11 +171,14 @@ def resolve_platform_config(
         platform_config: Existing config to fill in missing fields for.
         profile: Profile name for file config. Defaults to TIMBAL_PROFILE env var or "default".
         config_dir: Override the config directory (for testing).
+        force_refresh: Re-read env/file sources even on a default call, bypassing
+            (and refreshing) the cached default. Use when env credentials may have
+            been set after an earlier default resolution cached ``None``.
     """
     global _cached_default_config, _default_config_resolved
 
     is_default_call = platform_config is None and profile is None and config_dir is None
-    if is_default_call and _default_config_resolved:
+    if is_default_call and _default_config_resolved and not force_refresh:
         return _cached_default_config.model_copy() if _cached_default_config else None
 
     file_config = load_file_config(profile=profile, config_dir=config_dir)
@@ -241,17 +245,27 @@ def resolve_platform_config(
     if not app_id:
         app_id = os.getenv("TIMBAL_APP_ID")
 
+    project_id = None
+    if existing_subject:
+        project_id = existing_subject.project_id
+    if not project_id:
+        project_id = os.getenv("TIMBAL_PROJECT_ID")
+
     rev = existing_subject.rev if existing_subject else None
+    if not rev:
+        rev = os.getenv("TIMBAL_REV")
 
     platform_config.subject = PlatformSubject(
         org_id=org_id,
         app_id=app_id,
+        project_id=project_id,
         rev=rev,
     )
     logger.debug(
         "Resolved platform subject.",
         org_id=org_id,
         app_id=app_id,
+        project_id=project_id,
         rev=rev,
     )
 

timbal-2.2.0/python/timbal/tools/_creds.py

@@ -0,0 +1,82 @@
+"""Credential resolution for tools.
+
+Per credential, first match wins: explicit value on the tool → Integration
+credentials → env var. Raises CredentialNotAvailable when nothing resolves.
+"""
+
+import os
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from pydantic import SecretStr
+
+from ..errors import CredentialNotAvailable
+from ..platform.integrations import Integration
+
+__all__ = [
+    "resolve_api_key",
+]
+
+
+def _secret_value(value: Any) -> str:
+    if isinstance(value, SecretStr):
+        return value.get_secret_value()
+    return str(value)
+
+
+def _first_present(credentials: Any, keys: Sequence[str]) -> Any:
+    """First truthy value for `keys` in a credentials dict or object."""
+    if isinstance(credentials, Mapping):
+        for key in keys:
+            value = credentials.get(key)
+            if value:
+                return value
+        return None
+    for key in keys:
+        value = getattr(credentials, key, None)
+        if value:
+            return value
+    return None
+
+
+async def resolve_api_key(
+    *,
+    provider_name: str,
+    env_var: str | None = None,
+    integration: Any = None,
+    api_key: SecretStr | str | None = None,
+    tool: Any = None,
+    explicit_attr: str = "api_key",
+    integration_keys: Sequence[str] = ("api_key",),
+) -> str:
+    """Resolve a single provider secret.
+
+    Pass `tool=self` to pull `integration` / `explicit_attr` off it, or pass
+    `integration` / `api_key` explicitly. `integration_keys` are the keys tried
+    in the resolved integration credentials.
+    """
+    if tool is not None:
+        if integration is None:
+            integration = getattr(tool, "integration", None)
+        if api_key is None:
+            api_key = getattr(tool, explicit_attr, None)
+
+    if api_key is not None:
+        return _secret_value(api_key)
+
+    if isinstance(integration, Integration):
+        credentials = await integration.resolve()
+        value = _first_present(credentials, integration_keys)
+        if value is not None:
+            return str(value)
+
+    if env_var:
+        env_value = os.getenv(env_var)
+        if env_value:
+            return env_value
+
+    raise CredentialNotAvailable(
+        provider_name,
+        missing=[explicit_attr],
+        env_vars=[env_var] if env_var else [],
+    )

{timbal-2.1.7 → timbal-2.2.0}/python/timbal/tools/cala.py

@@ -4,24 +4,13 @@ from pydantic import Field, SecretStr
 
 from ..core.tool import Tool
 from ..platform.integrations import Integration
+from ._creds import resolve_api_key
 
 _DEFAULT_BASE_URL = "https://api.cala.ai/v1"
 # Back-compat alias for web_search imports
 _BASE_URL = _DEFAULT_BASE_URL
 
 
-async def _resolve_api_key(*, integration: Any = None, api_key: SecretStr | None = None) -> str:
-    """Resolve Cala API key from integration, explicit field, or env var."""
-    from ._creds import resolve_api_key
-
-    return await resolve_api_key(
-        env_var="CALA_API_KEY",
-        provider_name="Cala",
-        integration=integration,
-        api_key=api_key,
-    )
-
-
 def _normalize_base_url(base_url: str | None) -> str:
     return (base_url or _DEFAULT_BASE_URL).rstrip("/")
 
@@ -106,7 +95,12 @@ class CalaSearch(_CalaKnowledgeTool):
             explainability: bool = True,
             return_entities: bool = True,
         ) -> dict:
-            api_key = await _resolve_api_key(integration=self.integration, api_key=self.api_key)
+            api_key = await resolve_api_key(
+                env_var="CALA_API_KEY",
+                provider_name="Cala",
+                integration=self.integration,
+                api_key=self.api_key,
+            )
             return await _post_cala(
                 path="/knowledge/search",
                 api_key=api_key,
@@ -137,7 +131,12 @@ class CalaQuery(_CalaKnowledgeTool):
             ),
             return_entities: bool = True,
         ) -> dict:
-            api_key = await _resolve_api_key(integration=self.integration, api_key=self.api_key)
+            api_key = await resolve_api_key(
+                env_var="CALA_API_KEY",
+                provider_name="Cala",
+                integration=self.integration,
+                api_key=self.api_key,
+            )
             return await _post_cala(
                 path="/knowledge/query",
                 api_key=api_key,

{timbal-2.1.7 → timbal-2.2.0}/python/timbal/tools/elevenlabs.py

@@ -5,24 +5,11 @@ from pydantic import Field, SecretStr
 
 from ..core.tool import Tool
 from ..platform.integrations import Integration
+from ._creds import resolve_api_key
 
 _ELEVENLABS_BASE = "https://api.elevenlabs.io/v1"
 
 
-async def _resolve_api_key(tool: Any) -> str:
-    if isinstance(tool.integration, Integration):
-        credentials = await tool.integration.resolve()
-        return credentials["api_key"]
-    if tool.api_key is not None:
-        return tool.api_key.get_secret_value()
-    env_key = os.getenv("ELEVENLABS_API_KEY")
-    if env_key:
-        return env_key
-    raise ValueError(
-        "ElevenLabs API key not found. Set ELEVENLABS_API_KEY, pass api_key, or configure an integration."
-    )
-
-
 class ElevenLabsTextToSpeech(Tool):
     name: str = "elevenlabs_text_to_speech"
     description: str | None = (
@@ -49,7 +36,7 @@ class ElevenLabsTextToSpeech(Tool):
             style: float = Field(0.0, description="Style exaggeration (0.0–1.0, only for v2 models)."),
             use_speaker_boost: bool = Field(True, description="Improve speaker clarity (slight latency increase)."),
         ) -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import base64
 
             import httpx
@@ -99,7 +86,7 @@ class ElevenLabsListPhoneNumbers(Tool):
 
     def __init__(self, **kwargs: Any) -> None:
         async def _list_phone_numbers() -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import httpx
 
             async with httpx.AsyncClient(timeout=httpx.Timeout(120.0, connect=10.0)) as client:
@@ -140,7 +127,7 @@ class ElevenLabsMakeOutboundCall(Tool):
                 description="Destination phone number in E.164 format, e.g. '+14155552671'.",
             ),
         ) -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import httpx
 
             payload: dict[str, Any] = {
@@ -177,7 +164,7 @@ class ElevenLabsListAgents(Tool):
 
     def __init__(self, **kwargs: Any) -> None:
         async def _list_agents() -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import httpx
 
             async with httpx.AsyncClient(timeout=httpx.Timeout(120.0, connect=10.0)) as client:
@@ -205,7 +192,7 @@ class ElevenLabsGetModels(Tool):
 
     def __init__(self, **kwargs: Any) -> None:
         async def _get_models() -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import httpx
 
             async with httpx.AsyncClient(timeout=httpx.Timeout(120.0, connect=10.0)) as client:
@@ -238,7 +225,7 @@ class ElevenLabsGetVoicesWithDescriptions(Tool):
         async def _get_voices_with_descriptions(
             show_legacy: bool = Field(False, description="Include legacy pre-XTTS voices in the response."),
         ) -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import httpx
 
             async with httpx.AsyncClient(timeout=httpx.Timeout(120.0, connect=10.0)) as client:
@@ -292,7 +279,7 @@ class ElevenLabsListHistory(Tool):
             page_size: int = Field(100, description="Max items to return (default 100)."),
             source: str | None = Field(None, description="Filter by source: 'TTS', 'STS', or 'ConvAI'."),
         ) -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import httpx
 
             params: dict[str, Any] = {"page_size": page_size}
@@ -329,7 +316,7 @@ class ElevenLabsGetAudioFromHistoryItem(Tool):
         async def _get_audio_from_history_item(
             history_item_id: str = Field(..., description="The ID of the history item to retrieve audio for."),
         ) -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import base64
 
             import httpx
@@ -369,7 +356,7 @@ class ElevenLabsDownloadHistoryItems(Tool):
         async def _download_history_items(
             history_item_ids: list[str] = Field(..., description="List of history item IDs to download. Single ID returns audio; multiple IDs return a .zip."),
         ) -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import base64
             import io
             import zipfile
@@ -435,7 +422,7 @@ class ElevenLabsCreateAgent(Tool):
             max_duration_seconds: int = Field(600, description="Maximum call/conversation length in seconds."),
             turn_timeout: int = Field(20, description="Seconds of silence before the agent ends its turn."),
         ) -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import httpx
 
             payload: dict[str, Any] = {
@@ -505,7 +492,7 @@ class ElevenLabsAddVoice(Tool):
                 description="Optional metadata. Keys: language (BCP-47: 'ca', 'es', 'en', 'fr'), accent, gender, age. E.g. {'language': 'ca', 'accent': 'British'}.",
             ),
         ) -> Any:
-            api_key = await _resolve_api_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="ElevenLabs", env_var="ELEVENLABS_API_KEY")
             import base64
 
             import httpx

{timbal-2.1.7 → timbal-2.2.0}/python/timbal/tools/fal.py

@@ -1,4 +1,3 @@
-import os
 from typing import Annotated, Any
 from urllib.parse import quote, urlparse
 
@@ -6,28 +5,11 @@ from pydantic import Field, SecretStr
 
 from ..core.tool import Tool
 from ..platform.integrations import Integration
+from ._creds import resolve_api_key
 
 _QUEUE_BASE = "https://queue.fal.run"
 
 
-async def _resolve_fal_key(tool: Any) -> str:
-    """Resolve fal API key from integration, explicit field, or env var."""
-    if isinstance(tool.integration, Integration):
-        credentials = await tool.integration.resolve()
-        key = credentials.get("api_key")
-        if key:
-            return str(key)
-    if tool.api_key is not None:
-        return tool.api_key.get_secret_value()
-    env_key = os.getenv("FAL_KEY")
-    if env_key:
-        return env_key
-    raise ValueError(
-        "fal API key not found. Set FAL_KEY environment variable, "
-        "pass api_key in config, or configure an integration."
-    )
-
-
 def _fal_auth_headers(api_key: str) -> dict[str, str]:
     return {"Authorization": f"Key {api_key}", "Content-Type": "application/json"}
 
@@ -91,7 +73,7 @@ class FalQueueSubmit(Tool):
                 description="If set, fal POSTs the result to this URL when the job completes.",
             ),
         ) -> dict[str, Any]:
-            api_key = await _resolve_fal_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="fal", env_var="FAL_KEY")
             import httpx
 
             path = _encode_model_path(model_id)
@@ -149,7 +131,7 @@ class FalQueueStatus(Tool):
             request_id: str = Field(..., description="request_id returned by fal_queue_submit."),
             include_logs: bool = Field(False, description="If True, append ?logs=1 for runner log lines."),
         ) -> dict[str, Any]:
-            api_key = await _resolve_fal_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="fal", env_var="FAL_KEY")
             import httpx
 
             path = _encode_model_path(model_id)
@@ -197,7 +179,7 @@ class FalQueueResult(Tool):
             ),
             request_id: str = Field(..., description="request_id returned by fal_queue_submit."),
         ) -> dict[str, Any]:
-            api_key = await _resolve_fal_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="fal", env_var="FAL_KEY")
             import httpx
 
             path = _encode_model_path(model_id)
@@ -242,7 +224,7 @@ class FalQueueCancel(Tool):
             ),
             request_id: str = Field(..., description="request_id returned by fal_queue_submit."),
         ) -> dict[str, Any]:
-            api_key = await _resolve_fal_key(self)
+            api_key = await resolve_api_key(tool=self, provider_name="fal", env_var="FAL_KEY")
             import httpx
 
             path = _encode_model_path(model_id)