diff --git a/gateway/routes/allowlist.py b/gateway/routes/allowlist.py index 01ba9da3364..fe58e2dd58c 100644 --- a/gateway/routes/allowlist.py +++ b/gateway/routes/allowlist.py @@ -108,6 +108,7 @@ GATEWAY_PATH_PREFIXES: tuple[str, ...] = ( "/cursor/", "/milvus/", "/openai_passthrough/", + "/tinyfish/", # Dynamic provider / toolset passthrough (path templates) "/{provider}/", "/toolset/", diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index c53f7625550..5be87a8bf4d 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -211,6 +211,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( "/nvidia_nim/", "/openai/", "/openai_passthrough/", + "/tinyfish/", "/transcribe", "/typesafe/", "/openrouter/", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 3985feff08a..03122f25870 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -21163,6 +21163,96 @@ ] } }, + "/tinyfish/{endpoint}": { + "get": { + "description": "Pass-through for the TinyFish Agent API (goal-based web automation).\n\nForwarded endpoints:\n- POST /v1/automation/run \u2014 run to completion (blocking)\n- POST /v1/automation/run-async \u2014 submit a run, poll GET /v1/runs/{id} for the result\n- POST /v1/automation/run-sse \u2014 run with SSE progress events\n- GET /v1/runs/{id} \u2014 run status / result\n- POST /v1/runs/{id}/cancel \u2014 cancel a run\n\nEvery other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs\nlisting, which would let any caller discover other callers' run ids) returns 403: all\nproxy callers share one upstream key.\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. TINYFISH_API_KEY environment variable\n\n[Docs](https://docs.litellm.ai/docs/pass_through/tinyfish)", + "operationId": "tinyfish_proxy_route_tinyfish__endpoint__get", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Tinyfish Proxy Route", + "tags": [ + "llm_passthrough" + ] + }, + "post": { + "description": "Pass-through for the TinyFish Agent API (goal-based web automation).\n\nForwarded endpoints:\n- POST /v1/automation/run \u2014 run to completion (blocking)\n- POST /v1/automation/run-async \u2014 submit a run, poll GET /v1/runs/{id} for the result\n- POST /v1/automation/run-sse \u2014 run with SSE progress events\n- GET /v1/runs/{id} \u2014 run status / result\n- POST /v1/runs/{id}/cancel \u2014 cancel a run\n\nEvery other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs\nlisting, which would let any caller discover other callers' run ids) returns 403: all\nproxy callers share one upstream key.\n\nCredential lookup order:\n1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through)\n2. TINYFISH_API_KEY environment variable\n\n[Docs](https://docs.litellm.ai/docs/pass_through/tinyfish)", + "operationId": "tinyfish_proxy_route_tinyfish__endpoint__post", + "parameters": [ + { + "in": "path", + "name": "endpoint", + "required": true, + "schema": { + "title": "Endpoint", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Tinyfish Proxy Route", + "tags": [ + "llm_passthrough" + ] + } + }, "/transcribe": { "post": { "description": "AWS-SDK-shaped pass-through for Amazon Transcribe: point the SDK's `endpoint_url`\nat `/transcribe` and the operation is read from the `X-Amz-Target` header, per the\nAWS JSON 1.1 protocol.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/transcribe)", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 2a6b15e4097..8b2c81fea77 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -491,6 +491,7 @@ class LiteLLMRoutes(enum.Enum): "/openai_passthrough", "/assemblyai", "/eu.assemblyai", + "/tinyfish", "/vllm", "/mistral", "/typesafe", diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 74caa1050bb..b1960b9a046 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -14,6 +14,7 @@ import json import os import posixpath import re +import sys from collections.abc import AsyncGenerator, Callable, Mapping, Sequence from dataclasses import dataclass from functools import partial @@ -101,6 +102,12 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, ) +from litellm.types.passthrough_endpoints.tinyfish import ( + TINYFISH_AUTHENTICATED_RUN_FIELDS, + TINYFISH_PASSTHROUGH_TIMEOUT_SECONDS, + TINYFISH_REJECTED_ENVELOPE_FIELDS, + is_allowed_tinyfish_endpoint, +) from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials from litellm.types.router import LiteLLMParamsTypedDict from litellm.types.utils import LlmProviders @@ -3357,6 +3364,138 @@ async def cursor_proxy_route( return received_value +TINYFISH_JSON_OBJECT_BODY_DETAIL: Final = ( + "TinyFish requests must be a JSON object body sent with Content-Type: application/json." +) + + +async def _tinyfish_json_object_field_names(request: Request) -> frozenset[str] | None: + content_type: Final = request.headers.get("content-type", "") + if content_type and not is_json_content_type(content_type): + return None + raw_body: Final = await request.body() + if not raw_body: + return frozenset() + try: + parsed: Final[object] = json.loads(raw_body) # any-ok: json.loads -> Any + except (json.JSONDecodeError, UnicodeDecodeError): + return None + return frozenset(parsed) if isinstance(parsed, dict) else None + + +def _tinyfish_route_timeout() -> float | None: + # only raise the 600s default to cover legal 1200s runs; an operator's configured timeout still wins + proxy_server: Final = sys.modules.get("litellm.proxy.proxy_server") + operator_settings: Final = getattr(proxy_server, "general_settings", None) + operator_timeout: Final = ( + operator_settings.get("pass_through_request_timeout") if isinstance(operator_settings, Mapping) else None + ) + return None if operator_timeout is not None else TINYFISH_PASSTHROUGH_TIMEOUT_SECONDS + + +@router.api_route( + "/tinyfish/{endpoint:path}", + methods=["GET", "POST"], # mutable-ok: fastapi api_route requires List[str] + tags=["TinyFish Pass-through", "pass-through"], # mutable-ok: fastapi api_route requires a list +) +async def tinyfish_proxy_route( + endpoint: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection +) -> Response: + """ + Pass-through for the TinyFish Agent API (goal-based web automation). + + Forwarded endpoints: + - POST /v1/automation/run — run to completion (blocking) + - POST /v1/automation/run-async — submit a run, poll GET /v1/runs/{id} for the result + - POST /v1/automation/run-sse — run with SSE progress events + - GET /v1/runs/{id} — run status / result + - POST /v1/runs/{id}/cancel — cancel a run + + Every other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs + listing, which would let any caller discover other callers' run ids) returns 403: all + proxy callers share one upstream key. + + Credential lookup order: + 1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through) + 2. TINYFISH_API_KEY environment variable + + [Docs](https://docs.litellm.ai/docs/pass_through/tinyfish) + """ + from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + resolve_tinyfish_agent_api_base, + ) + + raw_endpoint_path: Final = httpx.URL(endpoint).path + encoded_endpoint: Final = raw_endpoint_path if raw_endpoint_path.startswith("/") else f"/{raw_endpoint_path}" + + if not is_allowed_tinyfish_endpoint(request.method, encoded_endpoint): + raise HTTPException( + status_code=403, + detail=f"{request.method} {encoded_endpoint} is not an allowed TinyFish Agent passthrough endpoint. " + "Allowed: POST /v1/automation/run, POST /v1/automation/run-async, POST /v1/automation/run-sse, " + "GET /v1/runs/{id}, POST /v1/runs/{id}/cancel.", + ) + + if request.method == "POST": + body_fields: Final = await _tinyfish_json_object_field_names(request) + if body_fields is None: + raise HTTPException(status_code=400, detail=TINYFISH_JSON_OBJECT_BODY_DETAIL) + envelope_fields: Final = tuple(sorted(body_fields & TINYFISH_REJECTED_ENVELOPE_FIELDS)) + if envelope_fields: + raise HTTPException( + status_code=400, + detail=f"Request fields [{', '.join(envelope_fields)}] are LiteLLM pass-through envelope controls " + "and are not accepted on the TinyFish route. Send the native TinyFish request body; streaming is " + "determined by the endpoint.", + ) + blocked_fields: Final = tuple(sorted(body_fields & TINYFISH_AUTHENTICATED_RUN_FIELDS)) + if ( + blocked_fields + and encoded_endpoint.startswith("/v1/automation/") + and str_to_bool(os.getenv("TINYFISH_ALLOW_AUTHENTICATED_RUNS")) is not True + ): + raise HTTPException( + status_code=403, + detail=f"Request fields [{', '.join(blocked_fields)}] run with the shared TinyFish account's saved " + "credentials and are disabled on this proxy. Ask the proxy admin to set " + "TINYFISH_ALLOW_AUTHENTICATED_RUNS=true to allow them.", + ) + + tinyfish_api_key: Final = passthrough_endpoint_router.get_credentials( + custom_llm_provider="tinyfish", + region_name=None, + ) + if tinyfish_api_key is None: + raise HTTPException( + status_code=401, + detail="TinyFish API key not found. Set the TINYFISH_API_KEY environment variable or add a " + "deployment with use_in_pass_through: true.", + ) + + base_url: Final = httpx.URL(resolve_tinyfish_agent_api_base()) + updated_url: Final = base_url.copy_with( + path=HttpPassThroughEndpointHelpers.join_base_and_endpoint_path(base_url, encoded_endpoint) + ) + + endpoint_func: Final = create_pass_through_route( + endpoint=endpoint, + target=str(updated_url), + custom_headers=MappingProxyType({"X-API-Key": tinyfish_api_key}), + custom_llm_provider="tinyfish", + timeout=_tinyfish_route_timeout(), + ) + received_value: Final = await endpoint_func( + request, + fastapi_response, + user_api_key_dict, + ) + + return received_value + + VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON: Final = ( "Vertex AI auth failed: set a use_in_pass_through vertex model, default_vertex_config, or DEFAULT_VERTEXAI_* env" ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py new file mode 100644 index 00000000000..a6c3cb669a6 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py @@ -0,0 +1,425 @@ +import asyncio +import json +import os +import time +import urllib.parse +from collections.abc import Mapping, Sequence +from datetime import datetime +from types import MappingProxyType +from typing import Final, NamedTuple +from urllib.parse import urlparse + +import httpx +from pydantic import TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import ( + get_standard_logging_object_payload, +) +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, +) +from litellm.proxy._types import PassThroughEndpointLoggingTypedDict +from litellm.types.llms.custom_http import httpxSpecialProvider +from litellm.types.passthrough_endpoints.tinyfish import ( + TINYFISH_AGENT_DEFAULT_API_BASE, + TINYFISH_DEFAULT_COST_PER_STEP, + TINYFISH_MAX_CONSECUTIVE_POLL_FAILURES, + TINYFISH_MAX_POLLING_SECONDS, + TINYFISH_MODEL_NAME, + TINYFISH_POLLING_INTERVAL_SECONDS, + TINYFISH_TERMINAL_RUN_STATUSES, + TinyfishRun, +) +from litellm.types.utils import StandardPassThroughResponseObject + +_RUN_ADAPTER: Final = TypeAdapter(TinyfishRun) + +_EMPTY_KWARGS: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _TinyfishLoggingPayload(NamedTuple): + result: StandardPassThroughResponseObject + kwargs: Mapping[str, object] + + def as_handler_result(self) -> PassThroughEndpointLoggingTypedDict: + handler_result: Final[PassThroughEndpointLoggingTypedDict] = { + "result": self.result, + "kwargs": {**self.kwargs}, + } + return handler_result + + +# asyncio tasks are weakly referenced by the loop; hold them until done or they can vanish mid-poll +_BACKGROUND_BILLING_TASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: task registry + + +def _register_billing_task(task: "asyncio.Task[None]") -> None: + _BACKGROUND_BILLING_TASKS.add(task) + task.add_done_callback(_BACKGROUND_BILLING_TASKS.discard) + task.add_done_callback(_warn_if_cancelled) + + +def _warn_if_cancelled(task: "asyncio.Task[None]") -> None: + # CancelledError bypasses the poller's exception handler, so shutdown-time charge loss must be logged here + if task.cancelled(): + verbose_proxy_logger.warning("TinyFish passthrough: billing poller cancelled mid-poll; the run may go unbilled") + + +_SSE_POLLER_SPAWNED_KEY: Final = "tinyfish_sse_poller_spawned" + + +def mark_sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> None: + logging_obj.model_call_details[_SSE_POLLER_SPAWNED_KEY] = True # rebind-ok: request-scoped scratch dict + + +def sse_poller_spawned(logging_obj: LiteLLMLoggingObj) -> bool: + return logging_obj.model_call_details.get(_SSE_POLLER_SPAWNED_KEY) is True + + +def run_id_from_sse_frames(frames: bytes) -> str | None: + return _run_id_from_sse_chunks(frames.decode("utf-8", errors="replace").splitlines()) + + +def resolve_tinyfish_agent_api_base() -> str: + raw: Final = (os.getenv("TINYFISH_AGENT_API_BASE") or TINYFISH_AGENT_DEFAULT_API_BASE).rstrip("/") + # a schemeless override would silently break both routing and billing (urlparse hostname becomes None) + return raw if "://" in raw else f"https://{raw}" + + +def resolve_tinyfish_cost_per_step() -> float: + raw: Final = os.getenv("TINYFISH_COST_PER_STEP") + if raw is None: + return TINYFISH_DEFAULT_COST_PER_STEP + try: + return float(raw) + except ValueError: + verbose_proxy_logger.warning( + "TINYFISH_COST_PER_STEP=%r is not a number; using the default rate %s", + raw, + TINYFISH_DEFAULT_COST_PER_STEP, + ) + return TINYFISH_DEFAULT_COST_PER_STEP + + +def is_tinyfish_agent_url(url: str) -> bool: + hostname: Final = urlparse(url).hostname + return hostname is not None and hostname == urlparse(resolve_tinyfish_agent_api_base()).hostname + + +def _parse_run(payload: object) -> TinyfishRun | None: + try: + return _RUN_ADAPTER.validate_python(payload) + except ValidationError as e: + verbose_proxy_logger.warning("TinyFish passthrough: unexpected run object shape: %s", e) + return None + + +def _run_cost(run: TinyfishRun | None) -> float | None: + if run is None: + return None + # TinyFish only invoices COMPLETED runs, so FAILED/CANCELLED runs must charge the team $0 + if run.get("status") != "COMPLETED": + return None + num_of_steps: Final = run.get("num_of_steps") + if num_of_steps is None: + return None + return num_of_steps * resolve_tinyfish_cost_per_step() + + +class TinyFishPassthroughLoggingHandler: + @staticmethod + def should_log_request(request_method: str, url_route: str) -> bool: + """Only run submissions are billed; GET /v1/runs* polling and cancels never write spend rows.""" + return request_method == "POST" and "/v1/automation/" in urlparse(url_route).path + + @staticmethod + def is_run_async_route(url_route: str) -> bool: + return urlparse(url_route).path.endswith("/v1/automation/run-async") + + @staticmethod + def tinyfish_passthrough_handler( + httpx_response: httpx.Response, + response_body: Mapping[str, object] | None, + logging_obj: LiteLLMLoggingObj, + url_route: str, + result: str, + start_time: datetime, + end_time: datetime, + cache_hit: bool, + request_body: Mapping[str, object], + **kwargs: object, # kwargs-ok: the passthrough logging dispatch forwards shared logging kwargs to every handler + ) -> PassThroughEndpointLoggingTypedDict: + """Bill a blocking POST /v1/automation/run: the response is the terminal run object.""" + try: + run: Final = _parse_run(response_body) if response_body is not None else None + handler_payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload( + run=run, + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + kwargs=kwargs, + ).as_handler_result() + except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request + verbose_proxy_logger.exception("Error in TinyFish passthrough logging handler: %s", e) + fallback_payload: Final[PassThroughEndpointLoggingTypedDict] = { + "result": StandardPassThroughResponseObject(response=result), + "kwargs": kwargs, + } + return fallback_payload + return handler_payload + + @staticmethod + def start_async_run_billing( + response_body: Mapping[str, object] | None, + logging_obj: LiteLLMLoggingObj, + result: str, + start_time: datetime, + cache_hit: bool, + **kwargs: object, # kwargs-ok: shared logging kwargs, replayed into _handle_logging when the run finishes + ) -> None: + """Bill POST /v1/automation/run-async once, when the polled run turns terminal.""" + submitted: Final = _parse_run(response_body) if response_body is not None else None + run_id: Final = submitted.get("run_id") if submitted is not None else None + if not run_id: + verbose_proxy_logger.warning( + "TinyFish passthrough: run-async response carried no run_id; logging the request without cost" + ) + task: Final = asyncio.create_task( + TinyFishPassthroughLoggingHandler._poll_and_log( + run_id=run_id, + logging_obj=logging_obj, + result=result, + start_time=start_time, + cache_hit=cache_hit, + kwargs=kwargs, + ) + ) + _register_billing_task(task) + + @staticmethod + def start_sse_run_billing( + run_id: str, + litellm_logging_obj: LiteLLMLoggingObj, + start_time: datetime, + client: AsyncHTTPHandler | None = None, + ) -> None: + """Bill POST /v1/automation/run-sse once via a detached poller that outlives client disconnects.""" + mark_sse_poller_spawned(litellm_logging_obj) + task: Final = asyncio.create_task( + TinyFishPassthroughLoggingHandler._poll_and_log( + run_id=run_id, + logging_obj=litellm_logging_obj, + result="", + start_time=start_time, + cache_hit=litellm_logging_obj.model_call_details.get("cache_hit") is True, + kwargs=_EMPTY_KWARGS, + client=client, + ) + ) + _register_billing_task(task) + + @staticmethod + async def _poll_and_log( + run_id: str | None, + logging_obj: LiteLLMLoggingObj, + result: str, + start_time: datetime, + cache_hit: bool, + kwargs: Mapping[str, object], + client: AsyncHTTPHandler | None = None, + ) -> None: + from ..pass_through_endpoints import pass_through_endpoint_logging + + try: + run: Final = ( + await TinyFishPassthroughLoggingHandler._poll_until_terminal(run_id, client) if run_id else None + ) + run_end_time: Final = datetime.now() # noqa: DTZ005 # naive to match the start_time stamped by pass_through_request + payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload( + run=run, + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=run_end_time, + kwargs=kwargs, + ) + await pass_through_endpoint_logging._handle_logging( # pyright: ignore[reportPrivateUsage] # shared passthrough logging dispatcher, same access as the assemblyai handler + logging_obj=logging_obj, + standard_logging_response_object=payload.result, + result=result, + start_time=start_time, + end_time=run_end_time, + cache_hit=cache_hit, + **payload.kwargs, + ) + except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request + verbose_proxy_logger.exception("[Non blocking logging error] TinyFish run-async billing failed: %s", e) + + @staticmethod + async def _poll_until_terminal( + run_id: str, + client: AsyncHTTPHandler | None = None, + poll_interval_seconds: float = TINYFISH_POLLING_INTERVAL_SECONDS, + ) -> TinyfishRun | None: + deadline: Final = time.monotonic() + TINYFISH_MAX_POLLING_SECONDS + last_run: TinyfishRun | None = None # rebind-ok: poll-loop state + consecutive_failures = 0 # rebind-ok: poll-loop state + while time.monotonic() < deadline: + run = await TinyFishPassthroughLoggingHandler._fetch_run(run_id, client) + if run is None: + # a single transient poll failure must not drop the run's charge + consecutive_failures += 1 + if consecutive_failures >= TINYFISH_MAX_CONSECUTIVE_POLL_FAILURES: + verbose_proxy_logger.warning( + "TinyFish passthrough: giving up on run %s after %s consecutive poll failures; " + "logging the request without cost", + run_id, + consecutive_failures, + ) + return last_run + else: + consecutive_failures = 0 + last_run = run + if (run.get("status") or "") in TINYFISH_TERMINAL_RUN_STATUSES: + return run + await asyncio.sleep(poll_interval_seconds) + verbose_proxy_logger.warning( + "TinyFish passthrough: run %s not terminal after %ss; logging the request without cost", + run_id, + TINYFISH_MAX_POLLING_SECONDS, + ) + return last_run + + @staticmethod + async def _fetch_run(run_id: str, client: AsyncHTTPHandler | None = None) -> TinyfishRun | None: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + passthrough_endpoint_router, + ) + + api_key: Final = passthrough_endpoint_router.get_credentials(custom_llm_provider="tinyfish", region_name=None) + if api_key is None: + verbose_proxy_logger.warning("TinyFish passthrough: no API key available to poll run %s", run_id) + return None + if any(c in run_id for c in ("/", "\\", "#", "?")) or ".." in run_id: + verbose_proxy_logger.warning("TinyFish passthrough: invalid run_id %r", run_id) + return None + safe_run_id: Final = urllib.parse.quote(run_id, safe="") + resolved_client: Final = client or get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint, + params={"timeout": 30.0}, # mutable-ok: get_async_httpx_client takes a plain dict of client params + ) + try: + # screenshots=none keeps the poll payload small (no per-step screenshot URLs needed) + response: Final = await resolved_client.get( + f"{resolve_tinyfish_agent_api_base()}/v1/runs/{safe_run_id}?screenshots=none", + headers={"X-API-Key": api_key}, # mutable-ok: httpx headers= takes a plain dict + ) + if not (200 <= response.status_code < 300): + verbose_proxy_logger.warning( + "TinyFish passthrough: GET /v1/runs/%s returned %s", safe_run_id, response.status_code + ) + return None + payload: Final[object] = response.json() # any-ok: httpx Response.json() -> Any + return _parse_run(payload) + except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request + verbose_proxy_logger.warning("[Non blocking logging error] TinyFish run fetch failed: %s", e) + return None + + @staticmethod + async def handle_logging_tinyfish_collected_chunks( + litellm_logging_obj: LiteLLMLoggingObj, + url_route: str, + start_time: datetime, + all_chunks: Sequence[str], + end_time: datetime, + client: AsyncHTTPHandler | None = None, + ) -> PassThroughEndpointLoggingTypedDict: + """Fallback for run-sse streams with no poller: logs the request, pricing via one GET if a run_id parses.""" + try: + run_id: Final = _run_id_from_sse_chunks(all_chunks) + if run_id is None: + verbose_proxy_logger.warning( + "TinyFish passthrough: no run_id in SSE stream; logging the request without cost" + ) + run: Final = await TinyFishPassthroughLoggingHandler._fetch_run(run_id, client) if run_id else None + payload: Final = TinyFishPassthroughLoggingHandler._build_logging_payload( + run=run, + logging_obj=litellm_logging_obj, + result="", + start_time=start_time, + end_time=end_time, + kwargs=_EMPTY_KWARGS, + ).as_handler_result() + except Exception as e: # noqa: BLE001 # billing/logging must never break the relayed request + verbose_proxy_logger.exception("Error in TinyFish SSE passthrough logging handler: %s", e) + fallback_payload: Final[PassThroughEndpointLoggingTypedDict] = { + "result": StandardPassThroughResponseObject(response=""), + "kwargs": {}, + } + return fallback_payload + return payload + + @staticmethod + def _build_logging_payload( + run: TinyfishRun | None, + logging_obj: LiteLLMLoggingObj, + result: str, + start_time: datetime, + end_time: datetime, + kwargs: Mapping[str, object], + ) -> _TinyfishLoggingPayload: + response_cost: Final = _run_cost(run) + updated_kwargs: Final = { # mutable-ok: the logging pipeline requires a plain kwargs dict + **kwargs, + "model": TINYFISH_MODEL_NAME, + "custom_llm_provider": "tinyfish", + "response_cost": response_cost, + # spend rows key on this as request_id; without it every poller-billed row is a NULL-key collision + "litellm_call_id": logging_obj.litellm_call_id, + # the poller paths pass no request kwargs, so SLO attribution (key hash, team, tags) needs the stored params + "litellm_params": kwargs.get("litellm_params") + or logging_obj.model_call_details.get("litellm_params") + or {}, # mutable-ok: the logging pipeline requires a plain kwargs dict + } + logging_obj.model_call_details.update( + model=TINYFISH_MODEL_NAME, + custom_llm_provider="tinyfish", + response_cost=response_cost, + ) + + logged_response: Final = StandardPassThroughResponseObject( + response=json.dumps(run) if run is not None else result + ) + standard_logging_object: Final = get_standard_logging_object_payload( + kwargs=updated_kwargs, + init_response_obj=logged_response, + start_time=start_time, + end_time=end_time, + logging_obj=logging_obj, + status="success", + ) + return _TinyfishLoggingPayload( + result=logged_response, + kwargs=MappingProxyType({**updated_kwargs, "standard_logging_object": standard_logging_object}), + ) + + +def _run_id_from_sse_chunks(all_chunks: Sequence[str]) -> str | None: + for line in all_chunks: + if not line.startswith("data:"): + continue + try: + event_payload: object = json.loads(line[5:].strip()) # any-ok: json.loads -> Any + except json.JSONDecodeError: + continue + event = _parse_run(event_payload) + if event is None: + continue + run_id = event.get("run_id") + if run_id: + return run_id + return None diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 8902e599788..c2874ac948f 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -113,6 +113,9 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( ) from litellm.types.utils import TRUSTED_CALLBACK_VARS_FIELD, Usage +from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + is_tinyfish_agent_url, +) from .streaming_handler import PassThroughStreamingHandler from .success_handler import PassThroughEndpointLogging from .upstream_usage_headers import ( @@ -383,6 +386,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): or (parsed_url.hostname and "openai.com" in parsed_url.hostname) ): return EndpointType.OPENAI + elif is_tinyfish_agent_url(url): + return EndpointType.TINYFISH return EndpointType.GENERIC @staticmethod diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index fae26b5a72d..e1f13f2bee0 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -27,6 +27,11 @@ from .llm_provider_handlers.gemini_passthrough_logging_handler import ( from .llm_provider_handlers.openai_passthrough_logging_handler import ( OpenAIPassthroughLoggingHandler, ) +from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + TinyFishPassthroughLoggingHandler, + run_id_from_sse_frames, + sse_poller_spawned, +) from .llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) @@ -70,6 +75,9 @@ class PassThroughStreamingHandler: exception: Exception, stream_context: PassThroughStreamContext | None = None, ) -> None: + # the tinyfish poller writes the one authoritative row; a failure row here would collide on its request_id + if endpoint_type == EndpointType.TINYFISH and sse_poller_spawned(litellm_logging_obj): + return await asyncify(PassThroughStreamingHandler._record_partial_usage_for_failure)( litellm_logging_obj=litellm_logging_obj, endpoint_type=endpoint_type, @@ -179,12 +187,25 @@ class PassThroughStreamingHandler: ) ) ) + # TinyFish SSE bills via a detached poller spawned on the first run_id frame, so disconnects can't lose the charge + tinyfish_scan_active = endpoint_type == EndpointType.TINYFISH # rebind-ok: scan stops once the poller spawns + tinyfish_pending = b"" # rebind-ok: SSE frame reassembly buffer across transport chunks try: if not cost_injection_active: # Hot path: just buffer for end-of-stream logging and forward. async for chunk in response.aiter_bytes(): raw_bytes.append(chunk) PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj) + if tinyfish_scan_active: + complete_frames, tinyfish_pending = split_complete_sse_frames(tinyfish_pending + chunk) + run_id = run_id_from_sse_frames(complete_frames) if b"run_id" in complete_frames else None + if run_id: + TinyFishPassthroughLoggingHandler.start_sse_run_billing( + run_id=run_id, + litellm_logging_obj=litellm_logging_obj, + start_time=start_time, + ) + tinyfish_scan_active = False yield chunk else: # ``cost_injection_active`` already requires ``model_name`` to @@ -294,6 +315,37 @@ class PassThroughStreamingHandler: and not _is_provider_error_chunk(complete_frames) ) try: + # TinyFish billing is owned by the detached poller; the $0 fallback below is only for streams with no run_id + if endpoint_type == EndpointType.TINYFISH: + if sse_poller_spawned(litellm_logging_obj): + return + late_run_id: Final = run_id_from_sse_frames(b"".join(raw_bytes)) + if late_run_id: + # the run_id arrived in an unterminated frame; poll to terminal instead of mispricing a RUNNING run + TinyFishPassthroughLoggingHandler.start_sse_run_billing( + run_id=late_run_id, + litellm_logging_obj=litellm_logging_obj, + start_time=start_time, + ) + return + tinyfish_payload: Final = ( + await TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks( + litellm_logging_obj=litellm_logging_obj, + url_route=url_route, + start_time=start_time, + all_chunks=PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes), + end_time=end_time, + ) + ) + await litellm_logging_obj.dispatch_success_handlers( + result=tinyfish_payload["result"], + start_time=start_time, + end_time=end_time, + cache_hit=litellm_logging_obj.model_call_details.get("cache_hit") is True, + prefer_async_handlers=True, + **tinyfish_payload["kwargs"], + ) + return ( standard_logging_response_object, kwargs, diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index d1e4da2e47c..6bba879b6c1 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -35,6 +35,10 @@ from .llm_provider_handlers.fal_ai_passthrough_logging_handler import ( from .llm_provider_handlers.gemini_passthrough_logging_handler import ( GeminiPassthroughLoggingHandler, ) +from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + TinyFishPassthroughLoggingHandler, + is_tinyfish_agent_url, +) from .llm_provider_handlers.transcribe_passthrough_logging_handler import ( TRANSCRIBE_CUSTOM_LLM_PROVIDER, PassThroughLogDispatch, @@ -281,6 +285,22 @@ class PassThroughEndpointLogging: ) standard_logging_response_object = comprehend_medical_handler_result["result"] # rebind-ok: elif-chain kwargs = comprehend_medical_handler_result["kwargs"] # rebind-ok: elif-chain contract + elif self.is_tinyfish_route(url_route, custom_llm_provider): + tinyfish_handler_result: Final = TinyFishPassthroughLoggingHandler.tinyfish_passthrough_handler( + httpx_response=httpx_response, + response_body=response_body if isinstance(response_body, dict) else None, + logging_obj=logging_obj, + url_route=url_route, + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + request_body=request_body, + **kwargs, + ) + standard_logging_response_object = tinyfish_handler_result["result"] # rebind-ok: elif-chain + kwargs = tinyfish_handler_result["kwargs"] # rebind-ok: elif-chain contract + elif self.is_azure_speech_route(custom_llm_provider): from .llm_provider_handlers.azure_speech_passthrough_logging_handler import ( AzureSpeechPassthroughLoggingHandler, @@ -336,6 +356,7 @@ class PassThroughEndpointLogging: ) standard_logging_response_object = typesafe_handler_result["result"] kwargs = typesafe_handler_result["kwargs"] + elif self.is_vertex_ai_live_route(url_route): from .llm_provider_handlers.vertex_ai_live_passthrough_logging_handler import ( VertexAILivePassthroughLoggingHandler, @@ -405,6 +426,20 @@ class PassThroughEndpointLogging: ): standard_logging_response_object: PassThroughEndpointLoggingResultValues | None = None logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload + if self.is_tinyfish_route(url_route, custom_llm_provider): + # polls and cancels never write spend rows; run-async bills once from the background poller + if not TinyFishPassthroughLoggingHandler.should_log_request(httpx_response.request.method, url_route): + return + if TinyFishPassthroughLoggingHandler.is_run_async_route(url_route): + TinyFishPassthroughLoggingHandler.start_async_run_billing( + response_body=response_body if isinstance(response_body, dict) else None, + logging_obj=logging_obj, + result=result, + start_time=start_time, + cache_hit=cache_hit, + **kwargs, + ) + return if self.is_assemblyai_route(url_route) and not self.is_azure_speech_route(custom_llm_provider): if AssemblyAIPassthroughLoggingHandler._should_log_request(httpx_response.request.method) is not True: return @@ -512,6 +547,9 @@ class PassThroughEndpointLogging: def is_comprehend_medical_route(self, custom_llm_provider: str | None) -> bool: return custom_llm_provider == "comprehendmedical" + def is_tinyfish_route(self, url_route: str, custom_llm_provider: str | None) -> bool: + return custom_llm_provider == "tinyfish" or is_tinyfish_agent_url(url_route) + def is_azure_speech_route(self, custom_llm_provider: str | None) -> bool: return custom_llm_provider == AZURE_SPEECH_CUSTOM_LLM_PROVIDER diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index e47acf9d68b..619001a5791 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -28,6 +28,7 @@ class EndpointType(str, Enum): GEMINI = "gemini" ANTHROPIC = "anthropic" OPENAI = "openai" + TINYFISH = "tinyfish" GENERIC = "generic" diff --git a/litellm/types/passthrough_endpoints/tinyfish.py b/litellm/types/passthrough_endpoints/tinyfish.py new file mode 100644 index 00000000000..365eaef77a6 --- /dev/null +++ b/litellm/types/passthrough_endpoints/tinyfish.py @@ -0,0 +1,55 @@ +from typing import Final + +from typing_extensions import ReadOnly, TypedDict + +TINYFISH_AGENT_DEFAULT_API_BASE: Final = "https://agent.tinyfish.ai" +TINYFISH_AGENT_DOCS_URL: Final = "https://docs.tinyfish.ai/agent-api" +# TinyFish's published Agent API rate (USD per run step); override with env TINYFISH_COST_PER_STEP +TINYFISH_DEFAULT_COST_PER_STEP: Final = 0.016 +TINYFISH_MODEL_NAME: Final = "tinyfish/automation-run" +TINYFISH_POLLING_INTERVAL_SECONDS: Final = 5.0 +# the Agent API caps runs at 1200s but queue wait extends wall time, so billing polls with generous headroom +TINYFISH_MAX_POLLING_SECONDS: Final = 3600.0 +# at the 5s interval this tolerates a ~60s upstream outage before abandoning the charge +TINYFISH_MAX_CONSECUTIVE_POLL_FAILURES: Final = 12 + +TINYFISH_TERMINAL_RUN_STATUSES: Final = frozenset({"COMPLETED", "FAILED", "CANCELLED"}) + +# these fields use the shared account's saved logins/vault, so they 403 unless TINYFISH_ALLOW_AUTHENTICATED_RUNS=true +TINYFISH_AUTHENTICATED_RUN_FIELDS: Final = frozenset({"use_profile", "profile_id", "use_vault", "credential_item_ids"}) + +# litellm's pass-through envelope controls; rejected here or custom_body smuggles past the field gate and +# a caller stream flag flips the billing mode away from what the endpoint dictates +TINYFISH_REJECTED_ENVELOPE_FIELDS: Final = frozenset({"custom_body", "stream", "query_params"}) + +# covers the upstream 1200s max run duration plus response headroom for blocking runs +TINYFISH_PASSTHROUGH_TIMEOUT_SECONDS: Final = 1500.0 + +_RUN_SUBMIT_PATHS: Final = frozenset( + {("v1", "automation", "run"), ("v1", "automation", "run-async"), ("v1", "automation", "run-sse")} +) + + +class TinyfishRun(TypedDict, total=False): + """Run objects are null-heavy until terminal, so every field must tolerate None.""" + + run_id: ReadOnly[str | None] + status: ReadOnly[str | None] + num_of_steps: ReadOnly[int | None] + result: ReadOnly[object] + # left untyped on purpose: a strict error shape would fail whole-run validation on upstream drift and drop the charge + error: ReadOnly[object] + type: ReadOnly[str | None] + + +def is_allowed_tinyfish_endpoint(method: str, path: str) -> bool: + """The host also serves vault/wallet/profile management under the same key, so only run endpoints forward.""" + segments: Final = tuple(path.split("/")[1:]) + if not path.startswith("/") or any(segment in ("", ".", "..") for segment in segments): + return False + if method == "POST" and segments in _RUN_SUBMIT_PATHS: + return True + # no GET /v1/runs listing: run ids are unguessable, so blocking the list keeps teams out of each other's runs + if method == "GET" and len(segments) == 3 and segments[:2] == ("v1", "runs"): + return True + return method == "POST" and len(segments) == 4 and segments[:2] == ("v1", "runs") and segments[3] == "cancel" diff --git a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py index 1d4e13474a7..6c57e59f7e3 100644 --- a/tests/pass_through_unit_tests/test_pass_through_unit_tests.py +++ b/tests/pass_through_unit_tests/test_pass_through_unit_tests.py @@ -413,6 +413,7 @@ PROTOCOL_CONSTRAINED_PASS_THROUGH_ROUTES = { "/comprehendmedical/{operation}": {"POST"}, "/transcribe": {"POST"}, "/transcribe/{operation}": {"POST"}, + "/tinyfish/{endpoint:path}": {"GET", "POST"}, } diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py new file mode 100644 index 00000000000..1b83c5140ca --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py @@ -0,0 +1,400 @@ +import asyncio +import json +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + _BACKGROUND_BILLING_TASKS, + TinyFishPassthroughLoggingHandler, + is_tinyfish_agent_url, + resolve_tinyfish_agent_api_base, + resolve_tinyfish_cost_per_step, + run_id_from_sse_frames, + sse_poller_spawned, +) +from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, +) +from litellm.types.passthrough_endpoints.tinyfish import is_allowed_tinyfish_endpoint + +RUN_URL = "https://agent.tinyfish.ai/v1/automation/run" +RUN_ASYNC_URL = "https://agent.tinyfish.ai/v1/automation/run-async" + + +def _make_logging_obj() -> MagicMock: + logging_obj = MagicMock() + logging_obj.litellm_call_id = "test-call-id" + logging_obj.model_call_details = {} + return logging_obj + + +def _make_response(method: str, url: str, body: dict) -> httpx.Response: + request = httpx.Request(method, url) + return httpx.Response(200, request=request, text=json.dumps(body)) + + +class _FakeClient: + """Payload items are dicts served with status_code, or (status, dict) tuples for scripted failures.""" + + def __init__(self, payloads: list, status_code: int = 200): + self.payloads = payloads + self.status_code = status_code + self.requested_urls: list[str] = [] + + async def get(self, url: str, headers: dict) -> httpx.Response: + self.requested_urls.append(url) + item = self.payloads[min(len(self.requested_urls) - 1, len(self.payloads) - 1)] + status, payload = item if isinstance(item, tuple) else (self.status_code, item) + return httpx.Response(status, text=json.dumps(payload), request=httpx.Request("GET", url)) + + +@pytest.fixture +def tinyfish_env(monkeypatch): + monkeypatch.setenv("TINYFISH_API_KEY", "sk-tf-test") + monkeypatch.delenv("TINYFISH_COST_PER_STEP", raising=False) + monkeypatch.delenv("TINYFISH_AGENT_API_BASE", raising=False) + + +class TestCostResolution: + def test_default_rate(self, tinyfish_env): + assert resolve_tinyfish_cost_per_step() == pytest.approx(0.016) + + def test_env_override(self, tinyfish_env, monkeypatch): + monkeypatch.setenv("TINYFISH_COST_PER_STEP", "0.02") + assert resolve_tinyfish_cost_per_step() == pytest.approx(0.02) + + def test_invalid_env_falls_back_to_default(self, tinyfish_env, monkeypatch): + monkeypatch.setenv("TINYFISH_COST_PER_STEP", "free") + assert resolve_tinyfish_cost_per_step() == pytest.approx(0.016) + + +class TestBillingGate: + @pytest.mark.parametrize( + "method,url,expected", + [ + ("POST", RUN_URL, True), + ("POST", RUN_ASYNC_URL, True), + ("POST", "https://agent.tinyfish.ai/v1/automation/run-sse", True), + ("GET", "https://agent.tinyfish.ai/v1/runs", False), + ("GET", "https://agent.tinyfish.ai/v1/runs/run-123?screenshots=none", False), + ("POST", "https://agent.tinyfish.ai/v1/runs/run-123/cancel", False), + ], + ) + def test_only_run_submissions_are_billed(self, method, url, expected): + assert TinyFishPassthroughLoggingHandler.should_log_request(method, url) is expected + + def test_polling_writes_no_spend_row(self, tinyfish_env): + logging_obj = _make_logging_obj() + logging_obj.dispatch_success_handlers = AsyncMock() + poll_url = "https://agent.tinyfish.ai/v1/runs/run-123" + + asyncio.run( + PassThroughEndpointLogging().pass_through_async_success_handler( + httpx_response=_make_response("GET", poll_url, {"run_id": "run-123", "status": "RUNNING"}), + response_body={"run_id": "run-123", "status": "RUNNING"}, + logging_obj=logging_obj, + url_route=poll_url, + result="", + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={}, + passthrough_logging_payload={"url": poll_url}, + custom_llm_provider="tinyfish", + ) + ) + + logging_obj.dispatch_success_handlers.assert_not_awaited() + + +class TestBlockingRunBilling: + def _handle(self, response_body: dict, logging_obj: MagicMock): + return TinyFishPassthroughLoggingHandler.tinyfish_passthrough_handler( + httpx_response=_make_response("POST", RUN_URL, response_body), + response_body=response_body, + logging_obj=logging_obj, + url_route=RUN_URL, + result=json.dumps(response_body), + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + request_body={"url": "https://scrapeme.live/shop", "goal": "extract products"}, + ) + + def test_bills_steps_times_rate(self, tinyfish_env): + logging_obj = _make_logging_obj() + run = {"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 3, "result": {"products": []}} + + handler_result = self._handle(run, logging_obj) + + assert handler_result["kwargs"]["model"] == "tinyfish/automation-run" + assert handler_result["kwargs"]["custom_llm_provider"] == "tinyfish" + assert handler_result["kwargs"]["response_cost"] == pytest.approx(0.048) + assert "standard_logging_object" in handler_result["kwargs"] + assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.048) + + def test_env_rate_override_applies(self, tinyfish_env, monkeypatch): + monkeypatch.setenv("TINYFISH_COST_PER_STEP", "0.5") + run = {"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 2} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] == pytest.approx(1.0) + + def test_failed_run_logs_without_cost(self, tinyfish_env): + run = {"run_id": "run-1", "status": "FAILED", "num_of_steps": 2, "error": {"code": "AGENT_FAILURE"}} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] is None + + def test_cancelled_run_logs_without_cost(self, tinyfish_env): + run = {"run_id": "run-1", "status": "CANCELLED", "num_of_steps": 2} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] is None + + def test_null_steps_logs_without_cost(self, tinyfish_env): + run = {"run_id": "run-1", "status": "RUNNING", "num_of_steps": None} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] is None + + def test_unexpected_error_shape_still_bills(self, tinyfish_env): + run = {"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 2, "error": {"retry_after": "5s"}} + + handler_result = self._handle(run, _make_logging_obj()) + + assert handler_result["kwargs"]["response_cost"] == pytest.approx(0.032) + + +class TestRunAsyncBilling: + def test_poll_and_log_bills_once_terminal(self, tinyfish_env): + logging_obj = _make_logging_obj() + logging_obj.dispatch_success_handlers = AsyncMock() + fake_client = _FakeClient( + payloads=[{"run_id": "run-9", "status": "COMPLETED", "num_of_steps": 4, "result": "ok"}] + ) + + asyncio.run( + TinyFishPassthroughLoggingHandler._poll_and_log( + run_id="run-9", + logging_obj=logging_obj, + result="", + start_time=datetime.now(), + cache_hit=False, + kwargs={}, + client=fake_client, + ) + ) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + awaited_kwargs = logging_obj.dispatch_success_handlers.await_args.kwargs + assert awaited_kwargs["response_cost"] == pytest.approx(0.064) + assert awaited_kwargs["model"] == "tinyfish/automation-run" + assert fake_client.requested_urls == ["https://agent.tinyfish.ai/v1/runs/run-9?screenshots=none"] + assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.064) + + def test_transient_poll_failure_keeps_polling(self, tinyfish_env): + fake_client = _FakeClient( + payloads=[(500, {}), {"run_id": "run-9", "status": "COMPLETED", "num_of_steps": 3}] + ) + + run = asyncio.run( + TinyFishPassthroughLoggingHandler._poll_until_terminal("run-9", fake_client, poll_interval_seconds=0.0) + ) + + assert run is not None + assert run["num_of_steps"] == 3 + assert len(fake_client.requested_urls) == 2 + + def test_gives_up_after_consecutive_poll_failures(self, tinyfish_env): + fake_client = _FakeClient(payloads=[(500, {})]) + + run = asyncio.run( + TinyFishPassthroughLoggingHandler._poll_until_terminal("run-9", fake_client, poll_interval_seconds=0.0) + ) + + assert run is None + assert len(fake_client.requested_urls) == 12 + + def test_traversal_run_id_is_rejected(self, tinyfish_env): + fake_client = _FakeClient(payloads=[{}]) + + run = asyncio.run(TinyFishPassthroughLoggingHandler._fetch_run("../vault/items", fake_client)) + + assert run is None + assert fake_client.requested_urls == [] + + def test_upstream_error_status_returns_none(self, tinyfish_env): + fake_client = _FakeClient(payloads=[{"error": {"code": "NOT_FOUND"}}], status_code=404) + + run = asyncio.run(TinyFishPassthroughLoggingHandler._fetch_run("run-1", fake_client)) + + assert run is None + + +class TestRunCostStatusGate: + def test_poller_bills_zero_for_terminal_failed_run(self, tinyfish_env): + logging_obj = _make_logging_obj() + logging_obj.dispatch_success_handlers = AsyncMock() + fake_client = _FakeClient(payloads=[{"run_id": "run-9", "status": "FAILED", "num_of_steps": 4}]) + + asyncio.run( + TinyFishPassthroughLoggingHandler._poll_and_log( + run_id="run-9", + logging_obj=logging_obj, + result="", + start_time=datetime.now(), + cache_hit=False, + kwargs={}, + client=fake_client, + ) + ) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + assert logging_obj.dispatch_success_handlers.await_args.kwargs["response_cost"] is None + assert len(fake_client.requested_urls) == 1 + + +class TestRunIdFromSseFrames: + def test_finds_run_id_in_first_frame(self): + frames = b'data: {"run_id": "run-7", "event": "INITIALIZED"}\n\ndata: {"run_id": "run-7", "event": "ACTION"}\n\n' + assert run_id_from_sse_frames(frames) == "run-7" + + def test_skips_frames_without_run_id(self): + frames = b': keepalive\n\ndata: not-json\n\ndata: {"event": "HEARTBEAT"}\n\n' + assert run_id_from_sse_frames(frames) is None + + +class TestStartSseRunBilling: + def test_spawns_detached_poller_that_bills_once(self, tinyfish_env): + logging_obj = _make_logging_obj() + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.model_call_details["litellm_params"] = {"metadata": {"user_api_key_hash": "hash-team-a"}} + fake_client = _FakeClient( + payloads=[{"run_id": "run-7", "status": "COMPLETED", "num_of_steps": 5, "result": "done"}] + ) + + async def _run() -> None: + tasks_before = set(_BACKGROUND_BILLING_TASKS) + TinyFishPassthroughLoggingHandler.start_sse_run_billing( + run_id="run-7", + litellm_logging_obj=logging_obj, + start_time=datetime.now(), + client=fake_client, + ) + assert sse_poller_spawned(logging_obj) + await asyncio.gather(*(_BACKGROUND_BILLING_TASKS - tasks_before)) + + asyncio.run(_run()) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + awaited_kwargs = logging_obj.dispatch_success_handlers.await_args.kwargs + assert awaited_kwargs["response_cost"] == pytest.approx(0.08) + # a missing call id makes every poller row a NULL request_id primary-key collision + assert awaited_kwargs["standard_logging_object"]["id"] == "test-call-id" + # SLO consumers (Prometheus, Langfuse) must see the caller's attribution despite the empty poller kwargs + assert awaited_kwargs["standard_logging_object"]["metadata"]["user_api_key_hash"] == "hash-team-a" + assert fake_client.requested_urls == ["https://agent.tinyfish.ai/v1/runs/run-7?screenshots=none"] + + def test_flag_defaults_to_not_spawned(self): + assert not sse_poller_spawned(_make_logging_obj()) + def test_collected_chunks_price_via_run_fetch(self, tinyfish_env): + logging_obj = _make_logging_obj() + chunks = [ + 'data: {"type": "STARTED", "run_id": "run-7", "status": "RUNNING"}', + 'data: {"type": "PROGRESS", "run_id": "run-7"}', + 'data: {"type": "COMPLETE", "run_id": "run-7", "status": "COMPLETED", "result": "done"}', + ] + fake_client = _FakeClient( + payloads=[{"run_id": "run-7", "status": "COMPLETED", "num_of_steps": 5, "result": "done"}] + ) + + payload = asyncio.run( + TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks( + litellm_logging_obj=logging_obj, + url_route="https://agent.tinyfish.ai/v1/automation/run-sse", + start_time=datetime.now(), + all_chunks=chunks, + end_time=datetime.now(), + client=fake_client, + ) + ) + + assert payload["kwargs"]["response_cost"] == pytest.approx(0.08) + assert payload["kwargs"]["model"] == "tinyfish/automation-run" + assert fake_client.requested_urls == ["https://agent.tinyfish.ai/v1/runs/run-7?screenshots=none"] + + def test_stream_without_run_id_logs_without_cost(self, tinyfish_env): + fake_client = _FakeClient(payloads=[{}]) + + payload = asyncio.run( + TinyFishPassthroughLoggingHandler.handle_logging_tinyfish_collected_chunks( + litellm_logging_obj=_make_logging_obj(), + url_route="https://agent.tinyfish.ai/v1/automation/run-sse", + start_time=datetime.now(), + all_chunks=["data: not-json", ": keepalive"], + end_time=datetime.now(), + client=fake_client, + ) + ) + + assert payload["kwargs"]["response_cost"] is None + assert fake_client.requested_urls == [] + + +class TestRouteDetection: + def test_provider_tag_claims_route(self): + assert PassThroughEndpointLogging().is_tinyfish_route("https://example.com/x", "tinyfish") + + def test_agent_host_claims_route(self): + assert PassThroughEndpointLogging().is_tinyfish_route("https://agent.tinyfish.ai/v1/runs", None) + + def test_other_providers_do_not_claim(self): + assert not PassThroughEndpointLogging().is_tinyfish_route("https://api.openai.com/v1", "openai") + + def test_env_base_override_claims_route(self, monkeypatch): + monkeypatch.setenv("TINYFISH_AGENT_API_BASE", "https://agent.staging.tinyfish.ai") + assert is_tinyfish_agent_url("https://agent.staging.tinyfish.ai/v1/runs/x") + assert not is_tinyfish_agent_url("https://agent.tinyfish.ai/v1/runs/x") + + def test_schemeless_env_base_is_normalized(self, monkeypatch): + monkeypatch.setenv("TINYFISH_AGENT_API_BASE", "agent.staging.tinyfish.ai") + assert resolve_tinyfish_agent_api_base() == "https://agent.staging.tinyfish.ai" + assert is_tinyfish_agent_url("https://agent.staging.tinyfish.ai/v1/runs/x") + + +class TestEndpointAllowlist: + @pytest.mark.parametrize( + "method,path,expected", + [ + ("POST", "/v1/automation/run", True), + ("POST", "/v1/automation/run-async", True), + ("POST", "/v1/automation/run-sse", True), + ("GET", "/v1/runs", False), + ("GET", "/v1/runs/run-abc-123", True), + ("POST", "/v1/runs/run-abc-123/cancel", True), + ("GET", "/v1/vault/items", False), + ("GET", "/v1/wallet", False), + ("POST", "/v1/browser-profiles", False), + ("DELETE", "/v1/runs/run-abc-123", False), + ("GET", "/v1/automation/run", False), + ("POST", "/v1/runs", False), + ("GET", "/v1/runs/..", False), + ("POST", "/v1/runs/../automation/run/cancel", False), + ("POST", "/v1/automation/run/", False), + ("POST", "/v1/automation/run-async/", False), + ("POST", "/v1/automation/run-sse/", False), + ("POST", "/v1//automation/run-async", False), + ("GET", "/v1/runs/run-abc-123/", False), + ("POST", "/v1/runs/run-abc-123/cancel/", False), + ], + ) + def test_allowlist(self, method, path, expected): + assert is_allowed_tinyfish_endpoint(method, path) is expected diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index dcc3bd2b690..fd81fcc8e72 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7371,3 +7371,200 @@ class TestOpenRouterPassthroughRoute: ) assert create_route.call_args.kwargs["target"] == f"{expected_root}/{endpoint}" + + +class TestTinyFishProxyRoute: + """Tests for the TinyFish Agent pass-through route, faking the upstream HTTP boundary.""" + + RUN_BODY = {"url": "https://scrapeme.live/shop", "goal": "Extract the first 2 product names. Return JSON."} + + @pytest.fixture + def tinyfish_client(self, monkeypatch: pytest.MonkeyPatch) -> Iterator[TestClient]: + from litellm.proxy.proxy_server import app + + monkeypatch.setenv("TINYFISH_API_KEY", "sk-tf-upstream") + monkeypatch.delenv("TINYFISH_AGENT_API_BASE", raising=False) + monkeypatch.delenv("TINYFISH_ALLOW_AUTHENTICATED_RUNS", raising=False) + monkeypatch.delenv("SERVER_ROOT_PATH", raising=False) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, lambda: UserAPIKeyAuth(api_key="sk-virtual")) + yield TestClient(app) + + def test_forwards_run_with_server_key_not_callers(self, tinyfish_client: TestClient) -> None: + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 2}) + ) + response = tinyfish_client.post( + "/tinyfish/v1/automation/run", json=self.RUN_BODY, headers={"X-API-Key": "sk-callers-virtual-key"} + ) + + assert (response.status_code, response.json()["run_id"]) == (200, "run-1") + assert route.calls.last.request.headers["x-api-key"] == "sk-tf-upstream" + + @pytest.mark.parametrize( + "method,path", + [ + ("GET", "/tinyfish/v1/vault/items"), + ("GET", "/tinyfish/v1/wallet"), + ("POST", "/tinyfish/v1/browser-profiles"), + ("GET", "/tinyfish/v1/automation/run"), + ("GET", "/tinyfish/v1/runs"), + ], + ) + def test_blocks_endpoints_outside_allowlist(self, tinyfish_client: TestClient, method: str, path: str) -> None: + with respx.mock: + response = tinyfish_client.request(method, path) + + assert response.status_code == 403 + assert "not an allowed TinyFish Agent passthrough endpoint" in response.json()["detail"] + + @pytest.mark.parametrize( + "path", + [ + "/tinyfish/v1/automation/run/", + "/tinyfish/v1/automation/run-async/", + "/tinyfish/v1/automation/run-sse/", + "/tinyfish/v1//automation/run-async", + ], + ) + def test_submit_paths_with_extra_slashes_are_rejected_before_forwarding( + self, tinyfish_client: TestClient, path: str + ) -> None: + with respx.mock(assert_all_called=False) as upstream: + upstream.post(url__regex=r"https://agent\.tinyfish\.ai/.*").mock( + return_value=httpx.Response(200, json={"run_id": "run-slash", "status": "PENDING"}) + ) + response = tinyfish_client.post(path, json=self.RUN_BODY) + + assert response.status_code == 403 + assert "not an allowed TinyFish Agent passthrough endpoint" in response.json()["detail"] + assert upstream.calls.call_count == 0 + + def test_rejects_authenticated_run_fields_by_default(self, tinyfish_client: TestClient) -> None: + with respx.mock: + response = tinyfish_client.post("/tinyfish/v1/automation/run", json={**self.RUN_BODY, "use_vault": True}) + + assert response.status_code == 403 + assert "use_vault" in response.json()["detail"] + + def test_env_opt_in_allows_authenticated_run_fields( + self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("TINYFISH_ALLOW_AUTHENTICATED_RUNS", "true") + + with respx.mock(assert_all_called=True) as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-2", "status": "COMPLETED", "num_of_steps": 1}) + ) + response = tinyfish_client.post("/tinyfish/v1/automation/run", json={**self.RUN_BODY, "use_vault": True}) + + assert response.status_code == 200 + assert json.loads(route.calls.last.request.content)["use_vault"] is True + + def test_returns_401_on_missing_api_key( + self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("TINYFISH_API_KEY") + + with respx.mock: + response = tinyfish_client.get("/tinyfish/v1/runs/run-123") + + assert response.status_code == 401 + assert "TINYFISH_API_KEY" in response.json()["detail"] + + def test_env_base_override_changes_target( + self, tinyfish_client: TestClient, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("TINYFISH_AGENT_API_BASE", "https://agent.staging.tinyfish.ai") + + with respx.mock(assert_all_called=True) as upstream: + upstream.get("https://agent.staging.tinyfish.ai/v1/runs/run-123").mock( + return_value=httpx.Response(200, json={"run_id": "run-123", "status": "RUNNING"}) + ) + response = tinyfish_client.get("/tinyfish/v1/runs/run-123") + + assert (response.status_code, response.json()["status"]) == (200, "RUNNING") + + @pytest.mark.parametrize( + "body", + [ + {"custom_body": {"url": "https://scrapeme.live/shop", "goal": "g", "use_vault": True}}, + {"url": "https://scrapeme.live/shop", "goal": "g", "stream": True}, + {"url": "https://scrapeme.live/shop", "goal": "g", "query_params": {"x": "1"}}, + ], + ) + def test_rejects_passthrough_envelope_controls(self, tinyfish_client: TestClient, body: dict) -> None: + """custom_body smuggled vault fields past the 403 gate and a stream flag flipped the + billing mode, because the generic passthrough honors both from the caller's body.""" + with respx.mock as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 1}) + ) + response = tinyfish_client.post("/tinyfish/v1/automation/run", json=body) + + assert response.status_code == 400 + assert "envelope" in response.json()["detail"] + assert not route.called + + def test_rejects_envelope_stream_on_cancel(self, tinyfish_client: TestClient) -> None: + with respx.mock as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/runs/run-1/cancel").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "CANCELLED"}) + ) + response = tinyfish_client.post("/tinyfish/v1/runs/run-1/cancel", json={"stream": True}) + + assert response.status_code == 400 + assert not route.called + + @pytest.mark.parametrize( + "content,content_type", + [ + ("url=https%3A%2F%2Fscrapeme.live%2Fshop&goal=g&stream=true", "application/x-www-form-urlencoded"), + ("url=https%3A%2F%2Fscrapeme.live%2Fshop&goal=g&use_vault=true", "application/x-www-form-urlencoded"), + ('{"url": "https://scrapeme.live/shop", "goal": "g", "use_vault": true}', "text/plain"), + ('[{"url": "https://scrapeme.live/shop", "goal": "g", "stream": true}]', "application/json"), + ], + ) + def test_rejects_bodies_that_are_not_json_objects( + self, tinyfish_client: TestClient, content: str, content_type: str + ) -> None: + """A form-encoded body carried stream and use_vault past both field gates, because + the gates only saw fields the body parsed to as JSON.""" + with respx.mock as upstream: + route = upstream.post("https://agent.tinyfish.ai/v1/automation/run").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "COMPLETED", "num_of_steps": 1}) + ) + response = tinyfish_client.post( + "/tinyfish/v1/automation/run", content=content, headers={"Content-Type": content_type} + ) + + assert response.status_code == 400 + assert "JSON object" in response.json()["detail"] + assert not route.called + + def test_cancel_without_body_forwards(self, tinyfish_client: TestClient) -> None: + with respx.mock(assert_all_called=True) as upstream: + upstream.post("https://agent.tinyfish.ai/v1/runs/run-1/cancel").mock( + return_value=httpx.Response(200, json={"run_id": "run-1", "status": "CANCELLED"}) + ) + response = tinyfish_client.post("/tinyfish/v1/runs/run-1/cancel") + + assert (response.status_code, response.json()["status"]) == (200, "CANCELLED") + + +class TestTinyFishRouteTimeout: + def test_default_covers_upstream_run_cap(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _tinyfish_route_timeout + + monkeypatch.setattr(proxy_server, "general_settings", {}, raising=False) + assert _tinyfish_route_timeout() == 1500.0 + + def test_operator_configured_timeout_wins(self, monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import _tinyfish_route_timeout + + monkeypatch.setattr(proxy_server, "general_settings", {"pass_through_request_timeout": 30}, raising=False) + assert _tinyfish_route_timeout() is None diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 88b82349c83..00a3606bbf5 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -7,8 +7,14 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx import litellm +import litellm.proxy.pass_through_endpoints.llm_provider_handlers.tinyfish_passthrough_logging_handler as tinyfish_handler_module +from litellm.proxy.pass_through_endpoints.llm_provider_handlers.tinyfish_passthrough_logging_handler import ( + mark_sse_poller_spawned, + sse_poller_spawned, +) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -840,6 +846,220 @@ async def test_chunk_processor_bills_partial_google_usage_on_mid_stream_exceptio assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) +class TestTinyFishStreamBilling: + """SSE billing is owned by the detached poller spawned on the first run_id frame; it must + survive gen.aclose() (client disconnect) and the stream-end path must not double-bill.""" + + RUNS_URL = "https://agent.tinyfish.ai/v1/runs/run-sse-1?screenshots=none" + SSE_ROUTE = "https://agent.tinyfish.ai/v1/automation/run-sse" + + @pytest.fixture + def tinyfish_env(self, monkeypatch): + monkeypatch.setenv("TINYFISH_API_KEY", "sk-tf-test") + monkeypatch.delenv("TINYFISH_COST_PER_STEP", raising=False) + monkeypatch.delenv("TINYFISH_AGENT_API_BASE", raising=False) + # aiohttp transport bypasses respx; force plain httpx and drop any cached aiohttp client + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + def _tinyfish_logging_obj(self): + obj = _unarmed_logging_obj() + obj.model_call_details = {} + obj.dispatch_success_handlers = AsyncMock() + return obj + + def _spawned_since(self, tasks_before): + return tinyfish_handler_module._BACKGROUND_BILLING_TASKS - tasks_before + + @pytest.mark.asyncio + async def test_run_id_split_across_chunks_spawns_one_poller(self, tinyfish_env): + chunks = [ + b'data: {"run_id": "run-s', + b'se-1", "event": "INITIALIZED"}\n\n', + b'data: {"run_id": "run-sse-1", "event": "COMPLETE"}\n\n', + ] + logging_obj = self._tinyfish_logging_obj() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + with respx.mock(assert_all_called=True) as upstream: + upstream.get(self.RUNS_URL).respond( + json={"run_id": "run-sse-1", "status": "COMPLETED", "num_of_steps": 2, "result": "ok"} + ) + received = [] + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=_make_streaming_response(chunks), + request_body={"url": "https://scrapeme.live/shop", "goal": "extract"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + route_streaming_logging=AsyncMock(), + ): + received.append(chunk) + + assert received == chunks + assert sse_poller_spawned(logging_obj) + spawned = self._spawned_since(tasks_before) + assert len(spawned) == 1 + await asyncio.gather(*spawned) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + assert logging_obj.dispatch_success_handlers.await_args.kwargs["response_cost"] == pytest.approx(0.032) + + @pytest.mark.asyncio + async def test_poller_survives_client_disconnect_and_bills(self, tinyfish_env): + chunks = [ + b'data: {"run_id": "run-sse-1", "event": "INITIALIZED"}\n\n', + b'data: {"run_id": "run-sse-1", "event": "ACTION"}\n\n', + ] + logging_obj = self._tinyfish_logging_obj() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + with respx.mock(assert_all_called=True) as upstream: + upstream.get(self.RUNS_URL).respond( + json={"run_id": "run-sse-1", "status": "COMPLETED", "num_of_steps": 2, "result": "ok"} + ) + gen = PassThroughStreamingHandler.chunk_processor( + response=_make_streaming_response(chunks), + request_body={"url": "https://scrapeme.live/shop", "goal": "extract"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + route_streaming_logging=AsyncMock(), + ) + await gen.__anext__() + spawned = self._spawned_since(tasks_before) + assert len(spawned) == 1 + await gen.aclose() + + task = next(iter(spawned)) + assert not task.cancelled() + await task + + logging_obj.dispatch_success_handlers.assert_awaited_once() + assert logging_obj.dispatch_success_handlers.await_args.kwargs["response_cost"] == pytest.approx(0.032) + + @pytest.mark.asyncio + async def test_stream_without_run_id_spawns_nothing(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + async for _ in PassThroughStreamingHandler.chunk_processor( + response=_make_streaming_response([b": keepalive\n\n", b'data: {"event": "HEARTBEAT"}\n\n']), + request_body={}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + route_streaming_logging=AsyncMock(), + ): + pass + + assert not sse_poller_spawned(logging_obj) + assert self._spawned_since(tasks_before) == set() + + @pytest.mark.asyncio + async def test_stream_end_skips_dispatch_when_poller_owns_billing(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + mark_sse_poller_spawned(logging_obj) + + await PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + request_body={}, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + raw_bytes=[b'data: {"run_id": "run-sse-1", "event": "COMPLETE"}\n\n'], + end_time=datetime.now(), + ) + + logging_obj.dispatch_success_handlers.assert_not_awaited() + + @pytest.mark.asyncio + async def test_stream_end_fallback_still_logs_when_no_poller_spawned(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + + await PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + request_body={}, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + raw_bytes=[b": keepalive\n\n"], + end_time=datetime.now(), + ) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + assert logging_obj.dispatch_success_handlers.await_args.kwargs["response_cost"] is None + + @pytest.mark.asyncio + async def test_upstream_error_after_spawn_skips_failure_dispatch(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + logging_obj.dispatch_failure_handlers = MagicMock() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + async def _aiter_bytes(): + yield b'data: {"run_id": "run-sse-1", "event": "INITIALIZED"}\n\n' + raise httpx.ReadTimeout("upstream died") + + response = MagicMock(spec=httpx.Response) + response.status_code = 200 + response.aiter_bytes = _aiter_bytes + + async def _consume(): + async for _ in PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + route_streaming_logging=AsyncMock(), + ): + pass + + with pytest.raises(httpx.ReadTimeout): + await _consume() + + spawned = self._spawned_since(tasks_before) + assert len(spawned) == 1 + # the poller owns the single row; a failure dispatch would collide on its request_id + logging_obj.dispatch_failure_handlers.assert_not_called() + for task in spawned: + task.cancel() + + @pytest.mark.asyncio + async def test_unterminated_run_id_frame_late_spawns_poller(self, tinyfish_env): + logging_obj = self._tinyfish_logging_obj() + tasks_before = set(tinyfish_handler_module._BACKGROUND_BILLING_TASKS) + + await PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route=self.SSE_ROUTE, + request_body={}, + endpoint_type=EndpointType.TINYFISH, + start_time=datetime.now(), + raw_bytes=[b'data: {"run_id": "run-sse-1", "event": "INITIALIZED"}'], + end_time=datetime.now(), + ) + + spawned = self._spawned_since(tasks_before) + assert len(spawned) == 1 + assert sse_poller_spawned(logging_obj) + # the poller polls to terminal instead of the old single fetch that mispriced a RUNNING run at $0 + logging_obj.dispatch_success_handlers.assert_not_awaited() + for task in spawned: + task.cancel() + + @pytest.mark.asyncio @pytest.mark.parametrize( "deferred_dispatch_armed", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index bd0de11e9fb..6a184d0765d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16600,6 +16600,64 @@ export interface paths { patch?: never; trace?: never; }; + "/tinyfish/{endpoint}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Tinyfish Proxy Route + * @description Pass-through for the TinyFish Agent API (goal-based web automation). + * + * Forwarded endpoints: + * - POST /v1/automation/run — run to completion (blocking) + * - POST /v1/automation/run-async — submit a run, poll GET /v1/runs/{id} for the result + * - POST /v1/automation/run-sse — run with SSE progress events + * - GET /v1/runs/{id} — run status / result + * - POST /v1/runs/{id}/cancel — cancel a run + * + * Every other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs + * listing, which would let any caller discover other callers' run ids) returns 403: all + * proxy callers share one upstream key. + * + * Credential lookup order: + * 1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through) + * 2. TINYFISH_API_KEY environment variable + * + * [Docs](https://docs.litellm.ai/docs/pass_through/tinyfish) + */ + get: operations["tinyfish_proxy_route_tinyfish__endpoint__get"]; + put?: never; + /** + * Tinyfish Proxy Route + * @description Pass-through for the TinyFish Agent API (goal-based web automation). + * + * Forwarded endpoints: + * - POST /v1/automation/run — run to completion (blocking) + * - POST /v1/automation/run-async — submit a run, poll GET /v1/runs/{id} for the result + * - POST /v1/automation/run-sse — run with SSE progress events + * - GET /v1/runs/{id} — run status / result + * - POST /v1/runs/{id}/cancel — cancel a run + * + * Every other Agent API endpoint (vault, wallet, browser profiles, and the GET /v1/runs + * listing, which would let any caller discover other callers' run ids) returns 403: all + * proxy callers share one upstream key. + * + * Credential lookup order: + * 1. passthrough_endpoint_router (config.yaml deployments with use_in_pass_through) + * 2. TINYFISH_API_KEY environment variable + * + * [Docs](https://docs.litellm.ai/docs/pass_through/tinyfish) + */ + post: operations["tinyfish_proxy_route_tinyfish__endpoint__post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/token": { parameters: { query?: never; @@ -63022,6 +63080,68 @@ export interface operations { }; }; }; + tinyfish_proxy_route_tinyfish__endpoint__get: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + tinyfish_proxy_route_tinyfish__endpoint__post: { + parameters: { + query?: never; + header?: never; + path: { + endpoint: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; token_endpoint_token_post: { parameters: { query?: {