From 19c13ac971ca5324e21f410e4c4783375a5cdbda Mon Sep 17 00:00:00 2001 From: Zachary Lyon Date: Mon, 21 Sep 2026 21:21:43 -0700 Subject: [PATCH] feat(proxy): add TinyFish Agent API passthrough with per-step billing (#41099) * feat(proxy): add TinyFish Agent API passthrough with per-step billing * chore(ui): regenerate dashboard API types for /tinyfish passthrough * fix(proxy): satisfy strict lint budget for tinyfish passthrough * style: ruff format tinyfish passthrough handler * test(proxy): exercise tinyfish route through the app with a faked upstream * refactor(proxy): make cross-module tinyfish billing hooks public * fix(proxy): tolerate transient tinyfish poll failures instead of dropping the charge * fix(proxy): defer billing for disconnected tinyfish SSE runs to the background poller * Revert "fix(proxy): defer billing for disconnected tinyfish SSE runs to the background poller" This reverts commit ef0bcfb4a0b9189eda0f8f43619949c4d50accd2. * fix(proxy): bill tinyfish SSE runs via detached poller and only COMPLETED runs Disconnected run-sse clients previously left completed runs unbilled: the stream-end handler saw a still-RUNNING run and logged $0. The poller now spawns from the streaming path on the first run_id frame, outlives the disconnect, and writes the one spend row when the run turns terminal; the stream-end path only logs the $0 fallback for run_id-less streams. Costs now apply only to COMPLETED runs ($0 for FAILED/CANCELLED, matching the upstream invoice), spend rows carry the request's litellm_call_id (previously NULL request_ids collided and were silently dropped), and the GET /v1/runs listing is blocked so callers behind the shared key cannot discover each other's runs. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * fix(proxy): drop GET /v1/runs from the tinyfish allowlist error message Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * chore(proxy): sync openapi artifacts for tinyfish docstring, suppress LIT011 on flag write Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * style(proxy): ruff-format the sse poller flag write Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * style(proxy): keep the rebind-ok suppression on the flag write's own line Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * fix(proxy): harden tinyfish billing paths from review findings Skip failure dispatch when the SSE poller owns billing (a failure row collided with the poller's billed row on request_id and dropped the charge), late-spawn the poller for run_ids that arrive in unterminated frames instead of mispricing RUNNING runs at $0, thread litellm_params into poller-billed standard logging objects so SLO consumers see attribution, untype the run error field so upstream error-shape drift cannot void a billable run, normalize a schemeless TINYFISH_AGENT_API_BASE, extend the poll budget to cover queue wait (3600s) with ~60s outage tolerance, and log poller cancellation. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * chore(proxy): satisfy ratcheted BLE001/LIT002 budgets from main in tinyfish handler Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * style(proxy): drop stray blank line from merge resolution Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * fix(proxy): reject passthrough envelope controls on tinyfish route, raise blocking-run timeout The generic passthrough unwraps a caller-supplied custom_body as the forwarded request and honors a caller stream flag, so custom_body.use_vault bypassed the credentialed-run 403 and stream: true flipped a blocking run into the streaming pipeline. The route now 400s the envelope fields (custom_body, stream, query_params); streaming comes from the endpoint. Blocking runs also get a 1500s default timeout covering the upstream 1200s run cap, unless the operator configured pass_through_request_timeout. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * style(proxy): resolve operator timeout without a dict-literal default Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Pxz11TQKYvCXvKawNUa2g4 * test(passthrough): list the TinyFish route among protocol-constrained pass-through routes * chore(proxy): regenerate the lazy OpenAPI snapshot after merging main * chore(proxy): keep the lazy OpenAPI snapshot as CI's Python 3.12 renders it * fix(proxy): reject TinyFish POST bodies that are not a JSON object A form-encoded or text body carried stream and use_vault past both field gates, because the gates only saw fields the body parsed to as JSON. The route now checks the content type before reading the body and answers 400 for anything that is not a JSON object. * fix(tinyfish): reject submit paths with extra slashes so run-async always bills The allowlist dropped empty path segments, so POST /v1/automation/run-async/ was forwarded upstream while the billing dispatch only recognises the exact path and would have logged the submit at $0 without starting the poller. Any path with a trailing or doubled slash now returns 403 before forwarding. --------- Co-authored-by: Claude Fable 5 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- gateway/routes/allowlist.py | 1 + litellm/proxy/_lazy_features.py | 1 + litellm/proxy/_lazy_openapi_snapshot.json | 90 ++++ litellm/proxy/_types.py | 1 + .../llm_passthrough_endpoints.py | 139 ++++++ .../tinyfish_passthrough_logging_handler.py | 425 ++++++++++++++++++ .../pass_through_endpoints.py | 5 + .../streaming_handler.py | 52 +++ .../pass_through_endpoints/success_handler.py | 38 ++ .../pass_through_endpoints.py | 1 + .../types/passthrough_endpoints/tinyfish.py | 55 +++ .../test_pass_through_unit_tests.py | 1 + ...st_tinyfish_passthrough_logging_handler.py | 400 +++++++++++++++++ .../test_llm_pass_through_endpoints.py | 197 ++++++++ .../test_streaming_handler_interrupt.py | 220 +++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 120 +++++ 16 files changed, 1746 insertions(+) create mode 100644 litellm/proxy/pass_through_endpoints/llm_provider_handlers/tinyfish_passthrough_logging_handler.py create mode 100644 litellm/types/passthrough_endpoints/tinyfish.py create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_tinyfish_passthrough_logging_handler.py 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?: {