mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
Merge branch 'litellm_internal_staging' into fix/async-stream-finalization
This commit is contained in:
Dixing Xu
commit
339fbe83a5
24 changed files with 1095 additions and 166 deletions
10
.github/workflows/test-litellm-ui-lint.yml
vendored
10
.github/workflows/test-litellm-ui-lint.yml
vendored
|
|
@ -29,7 +29,15 @@ jobs:
|
|||
id: changed
|
||||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
# base.sha is the base branch tip from when the PR was opened, while
|
||||
# actions/checkout leaves HEAD on a merge of the PR into the *current*
|
||||
# base tip. "$BASE_SHA"...HEAD therefore spans every base-branch commit
|
||||
# landed since, so a PR that touches no UI file still gets linted
|
||||
# against hundreds of other people's files. Diff the PR head against its
|
||||
# own merge base instead, which is exactly what this PR changed.
|
||||
merge_base=$(git merge-base "$BASE_SHA" "$HEAD_SHA")
|
||||
: > "$RUNNER_TEMP/prettier_files.txt"
|
||||
: > "$RUNNER_TEMP/eslint_files.txt"
|
||||
while IFS= read -r f; do
|
||||
|
|
@ -41,7 +49,7 @@ jobs:
|
|||
*.json | *.css | *.scss | *.md | *.mdx | *.yml | *.yaml | *.html)
|
||||
printf '%s\n' "$f" >> "$RUNNER_TEMP/prettier_files.txt" ;;
|
||||
esac
|
||||
done < <(git diff --name-only --diff-filter=ACMR --relative "$BASE_SHA"...HEAD -- .)
|
||||
done < <(git diff --name-only --diff-filter=ACMR --relative "$merge_base" "$HEAD_SHA" -- .)
|
||||
if [ -s "$RUNNER_TEMP/prettier_files.txt" ] || [ -s "$RUNNER_TEMP/eslint_files.txt" ]; then
|
||||
echo "has_files=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
|
|
|
|||
|
|
@ -316,26 +316,34 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
where_clause: Dict[str, Any] = {"file_purpose": "batch", **owner_filter}
|
||||
|
||||
if after:
|
||||
where_clause["id"] = {"gt": after}
|
||||
cursor_row = (
|
||||
await self.prisma_client.db.litellm_managedobjecttable.find_first(
|
||||
where={**where_clause, "unified_object_id": after}
|
||||
)
|
||||
)
|
||||
if cursor_row is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Invalid 'after' cursor: no batch found with id '{after}'.",
|
||||
)
|
||||
|
||||
fetch_limit = limit or 20
|
||||
if target_model_names:
|
||||
# Oversample so post-fetch model-name filtering still has enough rows.
|
||||
fetch_limit = max(fetch_limit * 3, 100)
|
||||
page_size = limit or 20
|
||||
cursor_args: Dict[str, Any] = (
|
||||
{"cursor": {"unified_object_id": after}, "skip": 1} if after else {}
|
||||
)
|
||||
|
||||
batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
where=where_clause,
|
||||
take=fetch_limit,
|
||||
order={"created_at": "desc"},
|
||||
take=page_size + 1,
|
||||
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
|
||||
**cursor_args,
|
||||
)
|
||||
|
||||
batch_objects: List[LiteLLMBatch] = []
|
||||
for batch in batches:
|
||||
try:
|
||||
# Stop once we have enough after filtering
|
||||
if len(batch_objects) >= (limit or 20):
|
||||
break
|
||||
has_more = len(batches) > page_size
|
||||
|
||||
batch_objects: List[LiteLLMBatch] = []
|
||||
for batch in batches[:page_size]:
|
||||
try:
|
||||
batch_data = (
|
||||
json.loads(batch.file_object)
|
||||
if isinstance(batch.file_object, str)
|
||||
|
|
@ -351,9 +359,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
continue
|
||||
|
||||
return build_list_page(
|
||||
batch_objects, has_more=len(batch_objects) == (limit or 20)
|
||||
)
|
||||
return build_list_page(batch_objects, has_more=has_more)
|
||||
|
||||
async def get_user_created_file_ids(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str]
|
||||
|
|
|
|||
|
|
@ -2727,7 +2727,7 @@ if MCP_AVAILABLE:
|
|||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"User not allowed to call this tool. Allowed MCP servers: {allowed_mcp_servers}",
|
||||
detail="User not allowed to call this tool.",
|
||||
)
|
||||
|
||||
standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = _get_standard_logging_mcp_tool_call(
|
||||
|
|
|
|||
|
|
@ -163,6 +163,15 @@ class MCPServer(BaseModel):
|
|||
allow_elicitation: bool = False
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"MCPServer(server_id={self.server_id!r}, name={self.name!r}, "
|
||||
f"transport={self.transport!r}, auth_type={self.auth_type!r})"
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.__repr__()
|
||||
|
||||
@property
|
||||
def has_client_credentials(self) -> bool:
|
||||
"""True if this server should use the OAuth2 client_credentials (M2M) flow.
|
||||
|
|
|
|||
|
|
@ -32,10 +32,17 @@ from e2e_config import unique_marker
|
|||
from e2e_http import Result, UnknownApiError, unwrap
|
||||
from lifecycle import ResourceManager
|
||||
|
||||
# No api_key: litellm resolves ANTHROPIC_API_KEY from the proxy's own environment
|
||||
# for this provider, which is what the agent-owner flow relies on. Pinning
|
||||
# "os.environ/ANTHROPIC_API_KEY" here (#34512) was worse than redundant. The a2a
|
||||
# bridge forwards litellm_params to acompletion() without expanding "os.environ/"
|
||||
# indirection, so that literal string was sent upstream as x-api-key and every
|
||||
# message/send failed with `AnthropicException - invalid x-api-key`. Omitting the
|
||||
# key lets the normal provider resolution apply. Verified against a live proxy:
|
||||
# omitted -> 200, "os.environ/..." -> 500 invalid x-api-key, literal key -> 200.
|
||||
BRIDGE = A2ABridgeParams(
|
||||
custom_llm_provider="anthropic",
|
||||
model="claude-haiku-4-5",
|
||||
api_key="os.environ/ANTHROPIC_API_KEY",
|
||||
)
|
||||
|
||||
MOVEHOME_AGENT_CARD_URL = "https://movehome.org/.well-known/agent.json"
|
||||
|
|
|
|||
|
|
@ -524,6 +524,17 @@ class TestOpenAIFiles:
|
|||
"llm.files.openai.list.nonstream.works",
|
||||
exercised_on=["files"],
|
||||
)
|
||||
@pytest.mark.skip(
|
||||
reason=(
|
||||
"LIT-4820 (https://linear.app/litellm-ai/issue/LIT-4820): GET /v1/files omits "
|
||||
"newly uploaded files. The upload succeeds and "
|
||||
"GET /v1/files/{id} returns the file, but it never appears in the listing; the "
|
||||
"returned set is stable with its newest entry ~10h old, on both the managed "
|
||||
"(/v1/files?model=) and provider-scoped (/openai/v1/files) routes. Skipped rather "
|
||||
"than weakened because the assertion below is the correct contract. Remove this "
|
||||
"marker when LIT-4820 is fixed; do not relax the assertion to make it pass."
|
||||
)
|
||||
)
|
||||
def test_uploaded_file_appears_in_list(
|
||||
self, client: BatchClient, resources: ResourceManager, batch_deployments: None
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -58,7 +58,7 @@
|
|||
- {id: mgmt.budget.delete.persists, module: mgmt, tier: P1, surface: api, assertions: [persists], source: "budget_management_endpoints.py:280", rationale: "Clears limits"}
|
||||
- {id: mgmt.budget.list.happy_path, module: mgmt, tier: P1, surface: api, assertions: [happy_path], source: "budget_management_endpoints.py:215", rationale: "Budget enumeration"}
|
||||
- {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"}
|
||||
- {id: mgmt.cache_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cache_settings_endpoints.py", rationale: "Cache config (smoke)"}
|
||||
- {id: mgmt.cache_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cache_settings_endpoints.py", rationale: "Cache config (smoke). Deliberately uncovered: the previous test read the live settings and wrote them back, which proves nothing (identical values in, so a no-op POST still passes) while being able to break the deployment. /cache/settings persists what it receives and that row outranks YAML cache_params, re-applied on a timer, so a write that omits ssl or redis_startup_nodes turns a TLS cluster into a plaintext standalone node and every later Redis call hangs. That took out 60 of 72 tests on 2026-07-25. GET cannot round-trip it either: it resolves the stored row overlaid with REDIS_* env and never reads YAML, so on a fresh deploy it cannot see YAML ssl to echo back. A safe test needs an isolated proxy, or LIT-4816 fixed so a partial write cannot downgrade transport. Do not re-add a read-then-write-back test against a shared proxy."}
|
||||
- {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"}
|
||||
- {id: mgmt.router_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "router_settings_endpoints.py", rationale: "Router config (smoke)"}
|
||||
- {id: mgmt.jwt_key_mapping.new.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "jwt_key_mapping_endpoints.py", rationale: "JWT->key mapping (smoke)"}
|
||||
|
|
|
|||
|
|
@ -77,10 +77,16 @@ REQUEST_TIMEOUT = float(os.environ.get("E2E_REQUEST_TIMEOUT", "60"))
|
|||
|
||||
EXPECT_RUST = os.environ.get("E2E_EXPECT_RUST", "").strip().lower() in ("1", "true", "yes")
|
||||
|
||||
LOAD_USERS = int(os.environ.get("E2E_LOAD_USERS", "750"))
|
||||
LOAD_SPAWN_RATE = float(os.environ.get("E2E_LOAD_SPAWN_RATE", "50"))
|
||||
# Deliberately modest concurrency. The suite shares its proxy with every other
|
||||
# suite in the run, and 750 users at spawn rate 50 saturated the request path hard
|
||||
# enough to distort latency-sensitive neighbours (and to spend real provider money
|
||||
# fast). The SLO is scaled with the user count to keep the same per-user throughput
|
||||
# expectation (~0.47 RPS/user), so this still catches a request-path regression
|
||||
# rather than becoming a formality. Raise all four via env for a real load run.
|
||||
LOAD_USERS = int(os.environ.get("E2E_LOAD_USERS", "200"))
|
||||
LOAD_SPAWN_RATE = float(os.environ.get("E2E_LOAD_SPAWN_RATE", "20"))
|
||||
LOAD_DURATION_SECONDS = float(os.environ.get("E2E_LOAD_DURATION_SECONDS", "60"))
|
||||
LOAD_MIN_RPS = float(os.environ.get("E2E_LOAD_MIN_RPS", "355"))
|
||||
LOAD_MIN_RPS = float(os.environ.get("E2E_LOAD_MIN_RPS", "90"))
|
||||
LOAD_MAX_FAILURE_RATIO = float(os.environ.get("E2E_LOAD_MAX_FAILURE_RATIO", "0.01"))
|
||||
|
||||
WEEKLY_ANOMALY_OPT_IN_ENV = "E2E_WEEKLY_ANOMALY"
|
||||
|
|
|
|||
|
|
@ -140,7 +140,17 @@ class GuardrailsClient:
|
|||
*,
|
||||
identifier: str,
|
||||
version: str,
|
||||
default_on: bool = False,
|
||||
) -> str:
|
||||
"""Register a Bedrock guardrail, opted out of `default_on` by default.
|
||||
|
||||
`default_on=True` applies the guardrail to every request the proxy serves,
|
||||
not just this test's. When the upstream ApplyGuardrail call fails (a missing
|
||||
bedrock:ApplyGuardrail permission answers 403), that failure is returned to
|
||||
unrelated traffic as `403 Bedrock guardrail request failed`, so one guardrail
|
||||
test takes out whatever else is running. Callers select the guardrail
|
||||
per-request instead, which keeps the blast radius to the test that wants it.
|
||||
"""
|
||||
return unwrap(
|
||||
self.proxy.transport.post(
|
||||
"/guardrails",
|
||||
|
|
@ -150,7 +160,7 @@ class GuardrailsClient:
|
|||
guardrail_name=name,
|
||||
litellm_params=BedrockGuardrailParamsBody(
|
||||
mode="pre_call",
|
||||
default_on=True,
|
||||
default_on=default_on,
|
||||
guardrailIdentifier=identifier,
|
||||
guardrailVersion=version,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -47,7 +47,10 @@ class TestBedrockGuardrail:
|
|||
)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
result = client.chat(scoped_key, MODEL, BLOCKED_PROMPT)
|
||||
# Selected per request rather than registered default_on, so an upstream
|
||||
# ApplyGuardrail failure surfaces here instead of 403ing every other suite
|
||||
# running against this proxy.
|
||||
result = client.chat(scoped_key, MODEL, BLOCKED_PROMPT, guardrails=[name])
|
||||
|
||||
match result:
|
||||
case UnknownApiError(status_code=status, body=body):
|
||||
|
|
|
|||
|
|
@ -105,7 +105,12 @@ class AnthropicMessageBody(BaseModel):
|
|||
class OpenAIChatBody(BaseModel):
|
||||
model: str
|
||||
messages: list[ChatMessage]
|
||||
max_tokens: int = 64
|
||||
# Passthrough sends this body to OpenAI untranslated, so it has to satisfy
|
||||
# OpenAI's current contract directly: newer models reject `max_tokens` with
|
||||
# "Unsupported parameter: 'max_tokens' is not supported with this model. Use
|
||||
# 'max_completion_tokens' instead." litellm's drop_params/translation does not
|
||||
# apply on this route.
|
||||
max_completion_tokens: int = 64
|
||||
|
||||
|
||||
class VllmChatBody(BaseModel):
|
||||
|
|
@ -198,14 +203,14 @@ class PassthroughClient:
|
|||
)
|
||||
|
||||
def openai_chat(
|
||||
self, key: str, model: str, text: str, *, max_tokens: int = 64
|
||||
self, key: str, model: str, text: str, *, max_completion_tokens: int = 64
|
||||
) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/openai/v1/chat/completions",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=OpenAIChatBody(
|
||||
model=model,
|
||||
max_tokens=max_tokens,
|
||||
max_completion_tokens=max_completion_tokens,
|
||||
messages=[ChatMessage(role="user", content=text)],
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -535,13 +535,25 @@ class TestOpenAIChatCompletions:
|
|||
ChatBody(
|
||||
model=model,
|
||||
messages=[
|
||||
# The prompt and effort have to make the model actually think,
|
||||
# otherwise this asserts something the model is free not to do:
|
||||
# at reasoning_effort="low" a one-step arithmetic question comes
|
||||
# back with reasoning_tokens=0, which is correct behavior and not
|
||||
# a reporting bug. Verified against the live model:
|
||||
# low + "60 miles in 1.5 hours" -> reasoning_tokens=0,
|
||||
# high + the prompt below -> reasoning_tokens=90.
|
||||
# What is under test is that litellm surfaces the field, so the
|
||||
# request has to be one where the field is populated.
|
||||
ChatMessage(
|
||||
role="user",
|
||||
content="A train travels 60 miles in 1.5 hours. What is its average speed in mph?",
|
||||
content=(
|
||||
"Prove that the sum of two odd integers is even, then find the "
|
||||
"smallest prime p greater than 100 such that p+2 is also prime."
|
||||
),
|
||||
)
|
||||
],
|
||||
reasoning_effort="low",
|
||||
max_tokens=2048,
|
||||
reasoning_effort="high",
|
||||
max_tokens=3000,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,9 +3,13 @@
|
|||
One method per registry cell, each asserting the real contract against a live
|
||||
proxy: read-only inventory routes return their documented shape, stateless
|
||||
validators compute their verdict from the request, and the write routes persist
|
||||
so a read-back reflects the change. The two routes that mutate global proxy state
|
||||
(cache settings and router settings, both driven from the admin UI) are exercised
|
||||
with a benign, self-restoring change so a shared proxy is left as it was found.
|
||||
so a read-back reflects the change. Router settings, which mutate global proxy
|
||||
state, are exercised with a benign, self-restoring change so a shared proxy is left
|
||||
as it was found.
|
||||
|
||||
Cache settings are deliberately not covered here; see the rationale on
|
||||
mgmt.cache_settings.update.happy_path in coverage_registry/mgmt.yaml before adding
|
||||
a test for that route.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -124,34 +128,6 @@ class ComplianceResponse(BaseModel):
|
|||
checks: list[ComplianceCheck]
|
||||
|
||||
|
||||
# ---- cache settings --------------------------------------------------------
|
||||
|
||||
|
||||
class CacheSettingsValue(BaseModel):
|
||||
type: str
|
||||
host: str = ""
|
||||
port: str = ""
|
||||
|
||||
|
||||
class CacheSettingsUpdateBody(BaseModel):
|
||||
cache_settings: CacheSettingsValue
|
||||
|
||||
|
||||
class CacheCurrentValues(BaseModel):
|
||||
type: str | None = None
|
||||
host: str | None = None
|
||||
port: str | None = None
|
||||
|
||||
|
||||
class CacheGetResponse(BaseModel):
|
||||
current_values: CacheCurrentValues
|
||||
|
||||
|
||||
class CacheUpdateResponse(BaseModel):
|
||||
status: str
|
||||
settings: CacheSettingsValue
|
||||
|
||||
|
||||
# ---- fallback management ---------------------------------------------------
|
||||
|
||||
|
||||
|
|
@ -369,70 +345,6 @@ class TestComplianceRoutes:
|
|||
)
|
||||
|
||||
|
||||
class TestCacheSettings:
|
||||
@pytest.mark.covers("mgmt.cache_settings.update.happy_path")
|
||||
def test_update_persists_cache_backend_to_get(
|
||||
self, client: ManagementClient, resources: ResourceManager
|
||||
) -> None:
|
||||
"""Exercise the update route without changing global state: capture the live
|
||||
cache backend and write exactly that back, so the config the proxy ends on is
|
||||
byte-for-byte the one it started with. A teardown restore of the same captured
|
||||
settings is the safety net if the body fails partway. The update route is only
|
||||
meaningful against a configured cache, so an unconfigured proxy fails loudly
|
||||
here rather than being silently switched to redis."""
|
||||
before = self._read_settings(client)
|
||||
assert before.type is not None, (
|
||||
"GET /cache/settings reported no cache type; refusing to invent one and mutate the shared proxy"
|
||||
)
|
||||
captured = CacheSettingsValue(type=before.type, host=before.host or "", port=before.port or "")
|
||||
resources.defer(lambda: self._write_settings(client, captured))
|
||||
|
||||
updated = unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/cache/settings",
|
||||
headers=client.proxy.transport.master,
|
||||
json=CacheSettingsUpdateBody(cache_settings=captured),
|
||||
response_type=CacheUpdateResponse,
|
||||
)
|
||||
)
|
||||
assert updated.status == "success", f"/cache/settings update status {updated.status!r}, expected 'success'"
|
||||
assert updated.settings.type == captured.type, (
|
||||
f"/cache/settings echoed type {updated.settings.type!r}, wrote {captured.type!r}"
|
||||
)
|
||||
|
||||
def reflected() -> CacheCurrentValues | None:
|
||||
current = self._read_settings(client)
|
||||
return current if current.type == captured.type else None
|
||||
|
||||
after = _poll(client, reflected, f"/cache/settings never reported type {captured.type!r} after the update")
|
||||
assert after.host == captured.host and after.port == captured.port, (
|
||||
f"/cache/settings persisted host/port {after.host!r}/{after.port!r}, "
|
||||
f"wrote {captured.host!r}/{captured.port!r}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _read_settings(client: ManagementClient) -> CacheCurrentValues:
|
||||
return unwrap(
|
||||
client.proxy.transport.get(
|
||||
"/cache/settings",
|
||||
headers=client.proxy.transport.master,
|
||||
params=NoBody(),
|
||||
response_type=CacheGetResponse,
|
||||
)
|
||||
).current_values
|
||||
|
||||
@staticmethod
|
||||
def _write_settings(client: ManagementClient, settings: CacheSettingsValue) -> None:
|
||||
_ = unwrap(
|
||||
client.proxy.transport.post(
|
||||
"/cache/settings",
|
||||
headers=client.proxy.transport.master,
|
||||
json=CacheSettingsUpdateBody(cache_settings=settings),
|
||||
response_type=CacheUpdateResponse,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class TestFallbackManagement:
|
||||
@pytest.mark.covers("mgmt.fallback_management.update.happy_path")
|
||||
def test_create_persists_and_is_read_back(self, client: ManagementClient, resources: ResourceManager) -> None:
|
||||
|
|
|
|||
|
|
@ -1771,8 +1771,8 @@ async def test_list_batches_from_managed_objects_table():
|
|||
# Should filter by user_id (created_by)
|
||||
prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
|
||||
where={"file_purpose": "batch", "created_by": "test-user"},
|
||||
take=10,
|
||||
order={"created_at": "desc"},
|
||||
take=11,
|
||||
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1801,8 +1801,8 @@ async def test_list_batches_from_managed_objects_table_empty_list():
|
|||
# Default take is 20 when no limit is provided
|
||||
prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
|
||||
where={"file_purpose": "batch", "created_by": "test-user"},
|
||||
take=20,
|
||||
order={"created_at": "desc"},
|
||||
take=21,
|
||||
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1918,8 +1918,8 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by():
|
|||
assert result_user1["data"][0].id == "unified-batch-user1"
|
||||
prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with(
|
||||
where={"file_purpose": "batch", "created_by": "user1"},
|
||||
take=10,
|
||||
order={"created_at": "desc"},
|
||||
take=11,
|
||||
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
|
||||
)
|
||||
|
||||
# Query with user2's API key - should only return user2's batch
|
||||
|
|
@ -1933,11 +1933,505 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by():
|
|||
assert result_user2["data"][0].id == "unified-batch-user2"
|
||||
prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with(
|
||||
where={"file_purpose": "batch", "created_by": "user2"},
|
||||
take=10,
|
||||
order={"created_at": "desc"},
|
||||
take=11,
|
||||
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_batches_pagination_uses_unified_object_id_cursor():
|
||||
"""Regression for LIT-4678.
|
||||
|
||||
The ``after`` cursor a client sends back is a batch's ``unified_object_id``
|
||||
(that is what is returned as ``.id`` / ``last_id``). Paginating must use a
|
||||
Prisma cursor on the unique ``unified_object_id`` column, not a
|
||||
``where id > after`` filter against the random-uuid primary key.
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
prisma_client = AsyncMock()
|
||||
prisma_client.db.litellm_managedobjecttable.find_first.return_value = MagicMock()
|
||||
prisma_client.db.litellm_managedobjecttable.find_many.return_value = []
|
||||
|
||||
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
|
||||
DualCache(), prisma_client=prisma_client
|
||||
)
|
||||
|
||||
await proxy_managed_files.list_user_batches(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="test-user"),
|
||||
limit=5,
|
||||
after="unified-batch-id-7",
|
||||
)
|
||||
|
||||
prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
|
||||
where={"file_purpose": "batch", "created_by": "test-user"},
|
||||
take=6,
|
||||
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
|
||||
cursor={"unified_object_id": "unified-batch-id-7"},
|
||||
skip=1,
|
||||
)
|
||||
|
||||
_, call_kwargs = prisma_client.db.litellm_managedobjecttable.find_many.call_args
|
||||
assert "id" not in call_kwargs["where"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_batches_pagination_walks_all_pages_without_loops_or_gaps():
|
||||
"""Regression for LIT-4678.
|
||||
|
||||
Simulates the managed-objects table (random-uuid ``id`` primary key,
|
||||
base64 ``unified_object_id``, reverse-chronological ``created_at``) and
|
||||
walks every page the way a client would, feeding ``last_id`` back as
|
||||
``after``. With the old ``where id > after`` cursor this loops and drops
|
||||
batches; the fixed cursor returns each batch exactly once, newest first.
|
||||
"""
|
||||
import uuid as _uuid
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
def _unified_id(i: int) -> str:
|
||||
raw = f"litellm_proxy;model_id:gpt-4o-batch;llm_batch_id:batch_{i:03d}"
|
||||
return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
|
||||
|
||||
total = 10
|
||||
rows = []
|
||||
for i in range(total):
|
||||
row = MagicMock()
|
||||
row.id = str(_uuid.uuid4())
|
||||
row.unified_object_id = _unified_id(i)
|
||||
row.created_at = 1_000_000 + i
|
||||
row.file_object = json.dumps(
|
||||
{
|
||||
"id": f"batch_provider_{i:03d}",
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"status": "completed",
|
||||
"created_at": 1_000_000 + i,
|
||||
"input_file_id": f"file-input-{i:03d}",
|
||||
"request_counts": {"total": 1, "completed": 1, "failed": 0},
|
||||
}
|
||||
)
|
||||
rows.append(row)
|
||||
|
||||
async def fake_find_many(where, take, order, cursor=None, skip=0):
|
||||
result = list(rows)
|
||||
id_filter = where.get("id")
|
||||
if isinstance(id_filter, dict) and "gt" in id_filter:
|
||||
result = [r for r in result if r.id > id_filter["gt"]]
|
||||
order_keys = order if isinstance(order, list) else [order]
|
||||
for clause in reversed(order_keys):
|
||||
(order_field, direction), = clause.items()
|
||||
result.sort(
|
||||
key=lambda r: getattr(r, order_field), reverse=(direction == "desc")
|
||||
)
|
||||
if cursor is not None:
|
||||
(cur_field, cur_val), = cursor.items()
|
||||
idx = next(
|
||||
(i for i, r in enumerate(result) if getattr(r, cur_field) == cur_val),
|
||||
None,
|
||||
)
|
||||
if idx is None:
|
||||
return []
|
||||
result = result[idx + skip:]
|
||||
return result[:take]
|
||||
|
||||
async def fake_find_first(where):
|
||||
return next(
|
||||
(r for r in rows if r.unified_object_id == where.get("unified_object_id")),
|
||||
None,
|
||||
)
|
||||
|
||||
prisma_client = AsyncMock()
|
||||
prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
side_effect=fake_find_many
|
||||
)
|
||||
prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
|
||||
side_effect=fake_find_first
|
||||
)
|
||||
|
||||
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
|
||||
DualCache(), prisma_client=prisma_client
|
||||
)
|
||||
user = UserAPIKeyAuth(user_id="test-user")
|
||||
|
||||
seen: list = []
|
||||
after = None
|
||||
for _ in range(total + 5):
|
||||
resp = await proxy_managed_files.list_user_batches(
|
||||
user_api_key_dict=user, limit=3, after=after
|
||||
)
|
||||
page_ids = [b.id for b in resp["data"]]
|
||||
seen.extend(page_ids)
|
||||
if not resp["has_more"]:
|
||||
break
|
||||
assert page_ids, "has_more was true but the page was empty"
|
||||
assert resp["last_id"] != after, "cursor did not advance (pagination loop)"
|
||||
after = resp["last_id"]
|
||||
|
||||
expected = [_unified_id(i) for i in reversed(range(total))]
|
||||
assert seen == expected
|
||||
assert len(seen) == len(set(seen))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_batches_pagination_stable_when_created_at_ties():
|
||||
"""Regression for LIT-4678.
|
||||
|
||||
Cursor pagination is only well-defined when the ``order`` fully determines
|
||||
row order. If listing ordered by non-unique ``created_at`` alone, batches
|
||||
sharing a timestamp come back in an arbitrary order that can shift between
|
||||
page requests, so a cursor row's neighbours change and batches get skipped
|
||||
or duplicated. Listing must add the unique ``unified_object_id`` as a
|
||||
tie-breaker so the order is total and pagination is stable.
|
||||
"""
|
||||
import itertools
|
||||
import uuid as _uuid
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
def _unified_id(i: int) -> str:
|
||||
raw = f"litellm_proxy;model_id:gpt-4o-batch;llm_batch_id:batch_{i:03d}"
|
||||
return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
|
||||
|
||||
total = 6
|
||||
shared_created_at = 1_000_000
|
||||
rows = []
|
||||
for i in range(total):
|
||||
row = MagicMock()
|
||||
row.id = str(_uuid.uuid4())
|
||||
row.unified_object_id = _unified_id(i)
|
||||
row.created_at = shared_created_at
|
||||
row.file_object = json.dumps(
|
||||
{
|
||||
"id": f"batch_provider_{i:03d}",
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"status": "completed",
|
||||
"created_at": shared_created_at,
|
||||
"input_file_id": f"file-input-{i:03d}",
|
||||
"request_counts": {"total": 1, "completed": 1, "failed": 0},
|
||||
}
|
||||
)
|
||||
rows.append(row)
|
||||
|
||||
call_counter = itertools.count()
|
||||
|
||||
async def fake_find_many(where, take, order, cursor=None, skip=0):
|
||||
call = next(call_counter)
|
||||
order_keys = order if isinstance(order, list) else [order]
|
||||
fields = [next(iter(clause)) for clause in order_keys]
|
||||
result = list(rows)
|
||||
for clause in reversed(order_keys):
|
||||
field, direction = next(iter(clause.items()))
|
||||
result.sort(
|
||||
key=lambda r: getattr(r, field), reverse=(direction == "desc")
|
||||
)
|
||||
|
||||
def order_key(r):
|
||||
return tuple(getattr(r, f) for f in fields)
|
||||
|
||||
stabilized = []
|
||||
i = 0
|
||||
while i < len(result):
|
||||
j = i
|
||||
while j < len(result) and order_key(result[j]) == order_key(result[i]):
|
||||
j += 1
|
||||
group = result[i:j]
|
||||
if len(group) > 1:
|
||||
rot = call % len(group)
|
||||
group = group[rot:] + group[:rot]
|
||||
stabilized.extend(group)
|
||||
i = j
|
||||
result = stabilized
|
||||
|
||||
if cursor is not None:
|
||||
cur_field, cur_val = next(iter(cursor.items()))
|
||||
idx = next(
|
||||
(k for k, r in enumerate(result) if getattr(r, cur_field) == cur_val),
|
||||
None,
|
||||
)
|
||||
if idx is None:
|
||||
return []
|
||||
result = result[idx + skip:]
|
||||
return result[:take]
|
||||
|
||||
async def fake_find_first(where):
|
||||
return next(
|
||||
(r for r in rows if r.unified_object_id == where.get("unified_object_id")),
|
||||
None,
|
||||
)
|
||||
|
||||
prisma_client = AsyncMock()
|
||||
prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
side_effect=fake_find_many
|
||||
)
|
||||
prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
|
||||
side_effect=fake_find_first
|
||||
)
|
||||
|
||||
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
|
||||
DualCache(), prisma_client=prisma_client
|
||||
)
|
||||
user = UserAPIKeyAuth(user_id="test-user")
|
||||
|
||||
seen: list = []
|
||||
after = None
|
||||
for _ in range(total + 5):
|
||||
resp = await proxy_managed_files.list_user_batches(
|
||||
user_api_key_dict=user, limit=2, after=after
|
||||
)
|
||||
page_ids = [b.id for b in resp["data"]]
|
||||
seen.extend(page_ids)
|
||||
if not resp["has_more"]:
|
||||
break
|
||||
assert page_ids, "has_more was true but the page was empty"
|
||||
assert resp["last_id"] != after, "cursor did not advance (pagination loop)"
|
||||
after = resp["last_id"]
|
||||
|
||||
assert sorted(seen) == sorted(_unified_id(i) for i in range(total))
|
||||
assert len(seen) == len(set(seen)), "a tied batch was returned more than once"
|
||||
|
||||
|
||||
def _managed_batch_row(index, file_object=None):
|
||||
row = MagicMock()
|
||||
row.id = f"pk-{index:03d}"
|
||||
raw = f"litellm_proxy;model_id:gpt-4o-batch;llm_batch_id:batch_{index:03d}"
|
||||
row.unified_object_id = base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
|
||||
row.created_at = 1_000_000 + index
|
||||
row.file_object = (
|
||||
file_object
|
||||
if file_object is not None
|
||||
else json.dumps(
|
||||
{
|
||||
"id": f"batch_provider_{index:03d}",
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"completion_window": "24h",
|
||||
"status": "completed",
|
||||
"created_at": 1_000_000 + index,
|
||||
"input_file_id": f"file-input-{index:03d}",
|
||||
"request_counts": {"total": 1, "completed": 1, "failed": 0},
|
||||
}
|
||||
)
|
||||
)
|
||||
return row
|
||||
|
||||
|
||||
def _fake_managed_object_table(rows):
|
||||
async def find_many(where, take, order, cursor=None, skip=0):
|
||||
result = sorted(
|
||||
rows, key=lambda r: (r.created_at, r.unified_object_id), reverse=True
|
||||
)
|
||||
if cursor is not None:
|
||||
(cur_field, cur_val), = cursor.items()
|
||||
idx = next(
|
||||
(i for i, r in enumerate(result) if getattr(r, cur_field) == cur_val),
|
||||
None,
|
||||
)
|
||||
if idx is None:
|
||||
return []
|
||||
result = result[idx + skip:]
|
||||
return result[:take]
|
||||
|
||||
async def find_first(where):
|
||||
return next(
|
||||
(r for r in rows if r.unified_object_id == where.get("unified_object_id")),
|
||||
None,
|
||||
)
|
||||
|
||||
prisma_client = AsyncMock()
|
||||
prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
|
||||
side_effect=find_many
|
||||
)
|
||||
prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
|
||||
side_effect=find_first
|
||||
)
|
||||
return prisma_client
|
||||
|
||||
|
||||
async def _walk_batch_pages(proxy_managed_files, user, limit, max_pages=20):
|
||||
pages = []
|
||||
after = None
|
||||
for _ in range(max_pages):
|
||||
resp = await proxy_managed_files.list_user_batches(
|
||||
user_api_key_dict=user, limit=limit, after=after
|
||||
)
|
||||
pages.append(resp)
|
||||
if not resp["has_more"]:
|
||||
break
|
||||
assert resp["last_id"] is not None, "has_more was true but there is no cursor"
|
||||
after = resp["last_id"]
|
||||
return pages
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_batches_rejects_unknown_after_cursor():
|
||||
"""An ``after`` that does not resolve to a batch the caller can see is a
|
||||
client error, not an empty page.
|
||||
|
||||
Returning ``[]`` for an unresolvable cursor is indistinguishable from
|
||||
"you have reached the end of the list", so a client walking pages with a
|
||||
stale or malformed cursor silently sees a truncated batch list instead of
|
||||
an error it can act on.
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
prisma_client = AsyncMock()
|
||||
prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
|
||||
DualCache(), prisma_client=prisma_client
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await proxy_managed_files.list_user_batches(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="test-user"),
|
||||
limit=3,
|
||||
after="does-not-exist-xyz",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "does-not-exist-xyz" in str(exc_info.value.detail)
|
||||
prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called()
|
||||
prisma_client.db.litellm_managedobjecttable.find_first.assert_called_once_with(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"created_by": "test-user",
|
||||
"unified_object_id": "does-not-exist-xyz",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_batches_treats_empty_after_as_no_cursor():
|
||||
"""``?after=`` means "start from the beginning", as it always has.
|
||||
|
||||
Only a cursor the client actually sent is validated, so an SDK that always
|
||||
emits the query parameter does not get a 400 on its first page.
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
rows = [_managed_batch_row(i) for i in range(2)]
|
||||
prisma_client = _fake_managed_object_table(rows)
|
||||
|
||||
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
|
||||
DualCache(), prisma_client=prisma_client
|
||||
)
|
||||
|
||||
page = await proxy_managed_files.list_user_batches(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=2, after=""
|
||||
)
|
||||
|
||||
assert [batch.id for batch in page["data"]] == [
|
||||
rows[1].unified_object_id,
|
||||
rows[0].unified_object_id,
|
||||
]
|
||||
prisma_client.db.litellm_managedobjecttable.find_first.assert_not_called()
|
||||
_, call_kwargs = prisma_client.db.litellm_managedobjecttable.find_many.call_args
|
||||
assert "cursor" not in call_kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_batches_rejects_after_cursor_owned_by_another_user():
|
||||
"""The cursor lookup must be scoped to the rows the caller can list.
|
||||
|
||||
A Prisma cursor resolves by unique column regardless of the ``where``
|
||||
filter, so an unscoped cursor would let one user anchor their page window
|
||||
to another user's batch and learn when it was created.
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
other_users_batch = _managed_batch_row(0)
|
||||
|
||||
async def find_first(where):
|
||||
if where.get("created_by") != "user-b":
|
||||
return None
|
||||
return other_users_batch
|
||||
|
||||
prisma_client = AsyncMock()
|
||||
prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
|
||||
side_effect=find_first
|
||||
)
|
||||
prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[])
|
||||
|
||||
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
|
||||
DualCache(), prisma_client=prisma_client
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await proxy_managed_files.list_user_batches(
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id="user-a"),
|
||||
limit=3,
|
||||
after=other_users_batch.unified_object_id,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_batches_has_more_false_on_exactly_full_final_page():
|
||||
"""``has_more`` must mean "another row exists", not "this page is full".
|
||||
|
||||
With a batch count that is an exact multiple of ``limit``, reporting
|
||||
``has_more`` off page fullness makes every client fetch one extra empty
|
||||
page before it can stop.
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
rows = [_managed_batch_row(i) for i in range(4)]
|
||||
prisma_client = _fake_managed_object_table(rows)
|
||||
|
||||
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
|
||||
DualCache(), prisma_client=prisma_client
|
||||
)
|
||||
|
||||
pages = await _walk_batch_pages(
|
||||
proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=2
|
||||
)
|
||||
|
||||
seen = [batch.id for page in pages for batch in page["data"]]
|
||||
assert seen == [r.unified_object_id for r in reversed(rows)]
|
||||
assert [page["has_more"] for page in pages] == [True, False]
|
||||
assert prisma_client.db.litellm_managedobjecttable.find_many.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_batches_unparseable_row_does_not_truncate_pagination():
|
||||
"""A row that fails to parse must not end pagination early.
|
||||
|
||||
Skipping a corrupt row shortens the page, so deriving ``has_more`` from
|
||||
the number of returned batches reports "no more results" while older
|
||||
batches are still unread, silently hiding them from the caller.
|
||||
"""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
rows = [_managed_batch_row(i) for i in range(4)]
|
||||
rows[2].file_object = "{ not valid json"
|
||||
prisma_client = _fake_managed_object_table(rows)
|
||||
|
||||
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
|
||||
DualCache(), prisma_client=prisma_client
|
||||
)
|
||||
|
||||
pages = await _walk_batch_pages(
|
||||
proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=2
|
||||
)
|
||||
|
||||
seen = [batch.id for page in pages for batch in page["data"]]
|
||||
assert seen == [rows[3].unified_object_id, rows[1].unified_object_id, rows[0].unified_object_id]
|
||||
assert len(seen) == len(set(seen))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_return_unified_file_id_includes_expires_at():
|
||||
from litellm.types.llms.openai import OpenAIFileObject
|
||||
|
|
@ -2312,8 +2806,8 @@ async def test_list_batches_only_returns_user_own_batches():
|
|||
# Verify the database query filtered by user_id
|
||||
prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
|
||||
where={"file_purpose": "batch", "created_by": "user_a_id"},
|
||||
take=10,
|
||||
order={"created_at": "desc"},
|
||||
take=11,
|
||||
order=[{"created_at": "desc"}, {"unified_object_id": "desc"}],
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3355,6 +3355,121 @@ async def test_call_mcp_tool_user_unauthorized_access():
|
|||
assert "User not allowed to call this tool" in exc_info.value.detail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_mcp_tool_unauthorized_403_does_not_leak_server_credentials():
|
||||
"""Regression for LIT-4703 / GH #29936.
|
||||
|
||||
Calling a tool on a server the key is not scoped to must 403 with a bare
|
||||
message. The prior code interpolated the caller's allowed ``List[MCPServer]``
|
||||
config objects into the 403 detail, dumping every upstream credential
|
||||
(authentication_token, client_secret, AWS keys, private keys, env,
|
||||
static_headers) in cleartext to the caller.
|
||||
"""
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.server import call_mcp_tool
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
mock_user_auth = UserAPIKeyAuth(
|
||||
api_key="test-key",
|
||||
user_id="test-user",
|
||||
team_id="team-basic",
|
||||
object_permission_id="key-permission-123",
|
||||
)
|
||||
|
||||
secret_fields = {
|
||||
"authentication_token": "sk-LEAK-authtok",
|
||||
"client_id": "LEAK-clientid",
|
||||
"client_secret": "sk-LEAK-clientsecret",
|
||||
"aws_access_key_id": "LEAK-akid",
|
||||
"aws_secret_access_key": "LEAK-awssecret",
|
||||
"aws_session_token": "LEAK-awssess",
|
||||
"client_private_key": "LEAK-privkey",
|
||||
}
|
||||
allowed_server_obj = MCPServer(
|
||||
server_id="allowed_server",
|
||||
name="allowed_server",
|
||||
server_name="allowed_server",
|
||||
alias="allowed_server",
|
||||
transport="http",
|
||||
auth_type=MCPAuth.bearer_token,
|
||||
env={"UPSTREAM_API_KEY": "LEAK-env"},
|
||||
static_headers={"X-Upstream-Auth": "LEAK-header"},
|
||||
**secret_fields,
|
||||
)
|
||||
|
||||
def mock_get_server_by_id(server_id):
|
||||
if server_id == "allowed_server":
|
||||
return allowed_server_obj
|
||||
return None
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.MCPRequestHandler.get_allowed_mcp_servers",
|
||||
AsyncMock(return_value=["allowed_server"]),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager.get_mcp_server_by_id",
|
||||
side_effect=mock_get_server_by_id,
|
||||
),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await call_mcp_tool(
|
||||
name="restricted_server-send_email",
|
||||
arguments={"to": "test@example.com"},
|
||||
user_api_key_auth=mock_user_auth,
|
||||
mcp_auth_header="Bearer test_token",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
assert exc_info.value.detail == "User not allowed to call this tool."
|
||||
all_secret_values = list(secret_fields.values()) + ["LEAK-env", "LEAK-header"]
|
||||
detail_text = str(exc_info.value.detail)
|
||||
leaked = [value for value in all_secret_values if value in detail_text]
|
||||
assert leaked == [], f"403 body leaked upstream credentials: {leaked}"
|
||||
|
||||
|
||||
def test_mcpserver_repr_and_str_mask_credentials():
|
||||
"""Regression for LIT-4703 / GH #29936.
|
||||
|
||||
``MCPServer.__repr__``/``__str__`` must never render credential fields, so a
|
||||
stray f-string, log line, or list interpolation cannot leak them. Only
|
||||
display is masked; ``model_dump`` serialization is unchanged.
|
||||
"""
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
secret_fields = {
|
||||
"authentication_token": "sk-SENTINEL-authtok",
|
||||
"client_id": "SENTINEL-clientid",
|
||||
"client_secret": "sk-SENTINEL-clientsecret",
|
||||
"aws_access_key_id": "SENTINEL-akid",
|
||||
"aws_secret_access_key": "SENTINEL-awssecret",
|
||||
"aws_session_token": "SENTINEL-awssess",
|
||||
"client_private_key": "SENTINEL-privkey",
|
||||
"client_private_key_id": "SENTINEL-privkeyid",
|
||||
}
|
||||
server = MCPServer(
|
||||
server_id="srv1",
|
||||
name="srv1",
|
||||
transport="http",
|
||||
auth_type=MCPAuth.bearer_token,
|
||||
env={"UPSTREAM_API_KEY": "SENTINEL-env"},
|
||||
static_headers={"X-Upstream-Auth": "SENTINEL-header"},
|
||||
env_vars=[{"name": "K", "value": "SENTINEL-envvar"}],
|
||||
**secret_fields,
|
||||
)
|
||||
|
||||
all_secrets = list(secret_fields.values()) + ["SENTINEL-env", "SENTINEL-header", "SENTINEL-envvar"]
|
||||
for text in (repr(server), str(server), repr([server]), f"{server}", f"{[server]}"):
|
||||
leaked = [secret for secret in all_secrets if secret in text]
|
||||
assert leaked == [], f"MCPServer rendering leaked credentials {leaked} in {text!r}"
|
||||
|
||||
assert "srv1" in repr(server)
|
||||
assert server.model_dump()["authentication_token"] == "sk-SENTINEL-authtok"
|
||||
assert server.model_dump()["client_secret"] == "sk-SENTINEL-clientsecret"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_tools_filters_by_key_team_permissions():
|
||||
"""Test that list_tools filters tools based on key/team mcp_tool_permissions"""
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ vi.mock("@/components/shared/charts", () => ({
|
|||
AreaChart: () => <div />,
|
||||
DonutChart: () => <div />,
|
||||
BarChart: () => <div />,
|
||||
CustomLegend: () => <div />,
|
||||
DEFAULT_COLOR_CYCLE: ["emerald"],
|
||||
}));
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import { render } from "@testing-library/react";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
import type { ToolSpendResponse } from "@/components/networking";
|
||||
|
||||
import type { DailyData, SpendMetrics } from "@/components/UsagePage/types";
|
||||
|
|
@ -25,6 +26,9 @@ vi.mock("@/components/shared/charts", () => ({
|
|||
BarChart: ({ data, categories }: { data: unknown; categories: string[] }) => (
|
||||
<div data-testid="bar-chart" data-categories={categories.join(",")} data-series={JSON.stringify(data)} />
|
||||
),
|
||||
CustomLegend: ({ categories }: { categories: readonly string[] }) => (
|
||||
<div data-testid="chart-legend">{categories.join(",")}</div>
|
||||
),
|
||||
DEFAULT_COLOR_CYCLE: ["emerald", "blue", "violet", "amber"],
|
||||
}));
|
||||
|
||||
|
|
@ -58,13 +62,20 @@ const day = (date: string, metrics: Partial<SpendMetrics>): DailyData => ({
|
|||
},
|
||||
});
|
||||
|
||||
const renderWith = (results: DailyData[], toolSpend = emptyToolSpend) => {
|
||||
interface RenderOptions {
|
||||
toolSpend?: ToolSpendResponse;
|
||||
from?: Date;
|
||||
to?: Date;
|
||||
}
|
||||
|
||||
const renderWith = (results: DailyData[], options: RenderOptions = {}) => {
|
||||
const { toolSpend = emptyToolSpend, from = new Date(2026, 6, 1), to = new Date(2026, 6, 14) } = options;
|
||||
mockGetToolSpend.mockResolvedValue(toolSpend);
|
||||
return render(
|
||||
<UsageTab
|
||||
accessToken="test-token"
|
||||
activity={{
|
||||
dateValue: { from: new Date("2026-07-01"), to: new Date("2026-07-14") },
|
||||
dateValue: { from, to },
|
||||
onDateChange: vi.fn(),
|
||||
results,
|
||||
loading: false,
|
||||
|
|
@ -74,7 +85,13 @@ const renderWith = (results: DailyData[], toolSpend = emptyToolSpend) => {
|
|||
);
|
||||
};
|
||||
|
||||
const readSeries = (element: HTMLElement) => JSON.parse(element.getAttribute("data-series") ?? "[]");
|
||||
|
||||
describe("UsageTab", () => {
|
||||
beforeEach(() => {
|
||||
mockGetToolSpend.mockReset();
|
||||
});
|
||||
|
||||
it("sums compression and caching dollars across days into the summary cards", () => {
|
||||
const { getByText } = renderWith([
|
||||
day("2026-07-12", {
|
||||
|
|
@ -95,16 +112,88 @@ describe("UsageTab", () => {
|
|||
expect(getByText("140,000 tokens compressed")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("builds a per-day time series and per-driver donut from the daily rows", () => {
|
||||
const { getByTestId } = renderWith([
|
||||
day("2026-07-12", { compression_savings_spend: 0.04, prompt_caching_savings_spend: 0.006 }),
|
||||
day("2026-07-13", { compression_savings_spend: 0.1, prompt_caching_savings_spend: 0.01 }),
|
||||
]);
|
||||
const twoDays = () => [
|
||||
day("2026-07-12", { compression_savings_spend: 0.04, prompt_caching_savings_spend: 0.006 }),
|
||||
day("2026-07-13", { compression_savings_spend: 0.1, prompt_caching_savings_spend: 0.01 }),
|
||||
];
|
||||
|
||||
const series = JSON.parse(getByTestId("area-chart").getAttribute("data-series") ?? "[]");
|
||||
it("opens on a running total anchored at $0 at the start of the range", () => {
|
||||
const { getByTestId } = renderWith(twoDays());
|
||||
|
||||
// Cumulative prepends a synthetic $0 point at the range start (Jul 1) so the
|
||||
// line rises from zero rather than floating; the daily running totals follow.
|
||||
const series = readSeries(getByTestId("area-chart"));
|
||||
expect(series).toHaveLength(3);
|
||||
expect(series[0]).toMatchObject({ date: "Jul 1", Compression: 0, "Prompt caching": 0 });
|
||||
expect(series[1]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 });
|
||||
expect(series[2].Compression).toBeCloseTo(0.14, 5);
|
||||
expect(series[2]["Prompt caching"]).toBeCloseTo(0.016, 5);
|
||||
});
|
||||
|
||||
it("rises from $0 to the day's cumulative total for a single-day range", () => {
|
||||
// The original complaint: a one-day range plotted a single floating dot. The
|
||||
// synthetic start anchor gives the line a zero origin to climb from.
|
||||
const oneDay = new Date(2026, 6, 24);
|
||||
const { getByTestId } = renderWith(
|
||||
[day("2026-07-24", { compression_savings_spend: 0.2, prompt_caching_savings_spend: 0.05 })],
|
||||
{ from: oneDay, to: oneDay },
|
||||
);
|
||||
|
||||
const series = readSeries(getByTestId("area-chart"));
|
||||
expect(series).toHaveLength(2);
|
||||
expect(series[0]).toMatchObject({ date: "Jul 24", Compression: 0, "Prompt caching": 0 });
|
||||
expect(series[1]).toMatchObject({ date: "Jul 24", Compression: 0.2, "Prompt caching": 0.05 });
|
||||
});
|
||||
|
||||
it("plots the daily series oldest first even though the rollup arrives newest first", async () => {
|
||||
// The daily activity endpoint returns days newest first; the chart must
|
||||
// still read left to right in time, and the running total must climb toward
|
||||
// the newest day, not fall away from it.
|
||||
const newestFirst = [
|
||||
day("2026-07-13", { prompt_caching_savings_spend: 0.1 }),
|
||||
day("2026-07-12", { prompt_caching_savings_spend: 0.04 }),
|
||||
];
|
||||
const { getByTestId, getByRole } = renderWith(newestFirst);
|
||||
|
||||
// The $0 anchor leads, then the days climb oldest to newest.
|
||||
const cumulative = readSeries(getByTestId("area-chart"));
|
||||
expect(cumulative.map((p: { date: string }) => p.date)).toEqual(["Jul 1", "Jul 12", "Jul 13"]);
|
||||
expect(cumulative[1]["Prompt caching"]).toBeCloseTo(0.04, 5);
|
||||
expect(cumulative[2]["Prompt caching"]).toBeCloseTo(0.14, 5);
|
||||
expect(cumulative[2]["Prompt caching"]).toBeGreaterThan(cumulative[1]["Prompt caching"]);
|
||||
|
||||
await userEvent.click(getByRole("tab", { name: "Per day" }));
|
||||
const perDay = readSeries(getByTestId("bar-chart"));
|
||||
expect(perDay.map((p: { date: string }) => p.date)).toEqual(["Jul 12", "Jul 13"]);
|
||||
});
|
||||
|
||||
it("draws bars of the raw per-interval readings on the other tab", async () => {
|
||||
const { getByRole, getByTestId, queryByTestId } = renderWith(twoDays());
|
||||
|
||||
// Cumulative opens on the area line.
|
||||
expect(getByTestId("area-chart")).toBeInTheDocument();
|
||||
|
||||
await userEvent.click(getByRole("tab", { name: "Per day" }));
|
||||
|
||||
// Per day switches to a bar chart of the unaccumulated daily savings, with no
|
||||
// synthetic anchor prepended.
|
||||
expect(queryByTestId("area-chart")).toBeNull();
|
||||
const series = readSeries(getByTestId("bar-chart"));
|
||||
expect(series).toHaveLength(2);
|
||||
expect(series[0]).toMatchObject({ Compression: 0.04, "Prompt caching": 0.006 });
|
||||
expect(series[1]).toMatchObject({ Compression: 0.1, "Prompt caching": 0.01 });
|
||||
});
|
||||
|
||||
it("says what the line means and over what range", async () => {
|
||||
const { getByText, getByRole } = renderWith(twoDays());
|
||||
|
||||
expect(getByText("Running total saved · Jul 1 – Jul 14")).toBeInTheDocument();
|
||||
await userEvent.click(getByRole("tab", { name: "Per day" }));
|
||||
expect(getByText("Saved per day · Jul 1 – Jul 14")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("builds the per-driver donut from the range totals, not the running total", () => {
|
||||
const { getByTestId } = renderWith(twoDays());
|
||||
|
||||
const slices = JSON.parse(getByTestId("donut-chart").getAttribute("data-slices") ?? "[]");
|
||||
expect(slices).toEqual([
|
||||
|
|
@ -131,7 +220,7 @@ describe("UsageTab", () => {
|
|||
start_date: "2026-07-12",
|
||||
end_date: "2026-07-12",
|
||||
};
|
||||
const { findAllByTestId } = renderWith([day("2026-07-12", {})], toolSpend);
|
||||
const { findAllByTestId } = renderWith([day("2026-07-12", {})], { toolSpend });
|
||||
|
||||
const bars = await findAllByTestId("bar-chart");
|
||||
const series = JSON.parse(bars[0].getAttribute("data-series") ?? "[]");
|
||||
|
|
@ -146,7 +235,7 @@ describe("UsageTab", () => {
|
|||
start_date: "2026-07-05",
|
||||
end_date: "2026-07-14",
|
||||
};
|
||||
const { findByText } = renderWith([day("2026-07-12", {})], toolSpend);
|
||||
const { findByText } = renderWith([day("2026-07-12", {})], { toolSpend });
|
||||
|
||||
expect(await findByText(/capped at 30 days before the end of the selected range/)).toBeInTheDocument();
|
||||
});
|
||||
|
|
@ -159,7 +248,7 @@ describe("UsageTab", () => {
|
|||
start_date: "2026-07-01",
|
||||
end_date: "2026-07-14",
|
||||
};
|
||||
const { findAllByTestId, queryByText } = renderWith([day("2026-07-12", {})], toolSpend);
|
||||
const { findAllByTestId, queryByText } = renderWith([day("2026-07-12", {})], { toolSpend });
|
||||
|
||||
await findAllByTestId("bar-chart");
|
||||
expect(queryByText(/capped at 30 days before the end of the selected range/)).not.toBeInTheDocument();
|
||||
|
|
|
|||
|
|
@ -3,14 +3,27 @@
|
|||
import React, { useEffect, useMemo, useState } from "react";
|
||||
import { Info } from "lucide-react";
|
||||
|
||||
import { AreaChart, BarChart, DonutChart, DEFAULT_COLOR_CYCLE } from "@/components/shared/charts";
|
||||
import { AreaChart, BarChart, CustomLegend, DonutChart, DEFAULT_COLOR_CYCLE } from "@/components/shared/charts";
|
||||
import AdvancedDatePicker from "@/components/shared/advanced_date_picker";
|
||||
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
|
||||
import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import { getToolSpend, ToolSpendResponse } from "@/components/networking";
|
||||
import { SpendMetrics } from "@/components/UsagePage/types";
|
||||
import { formatNumberWithCommas } from "@/utils/dataUtils";
|
||||
import { buildDailyToolSeries, topToolsBySpend, usd } from "./costOptimizationUtils";
|
||||
import {
|
||||
buildDailyToolSeries,
|
||||
formatRangeLabel,
|
||||
localIsoDay,
|
||||
MAX_POINTS_WITH_DOTS,
|
||||
SAVINGS_SERIES,
|
||||
SavingsAccumulation,
|
||||
SavingsPoint,
|
||||
toCumulative,
|
||||
topToolsBySpend,
|
||||
usd,
|
||||
withStartAnchor,
|
||||
} from "./costOptimizationUtils";
|
||||
import { DailyActivityRange } from "./useDailyActivityRange";
|
||||
|
||||
interface UsageTabProps {
|
||||
|
|
@ -26,6 +39,8 @@ const EMPTY_TOOL_SPEND: ToolSpendResponse = {
|
|||
end_date: null,
|
||||
};
|
||||
|
||||
const SAVINGS_COLORS = ["emerald", "blue"] as const;
|
||||
|
||||
const shortDate = (iso: string): string =>
|
||||
new Date(`${iso}T00:00:00`).toLocaleDateString("en-US", { month: "short", day: "numeric" });
|
||||
|
||||
|
|
@ -95,16 +110,41 @@ const UsageTab: React.FC<UsageTabProps> = ({ accessToken, activity }) => {
|
|||
const savedTokensTotal = useMemo(() => results.reduce((sum, d) => sum + savedTokensOf(d.metrics), 0), [results]);
|
||||
const totalSaved = compressionTotal + cachingTotal;
|
||||
|
||||
const overTime = useMemo(
|
||||
const [accumulation, setAccumulation] = useState<SavingsAccumulation>("cumulative");
|
||||
|
||||
// The daily rollup arrives newest first; sort on the raw ISO date so the axis
|
||||
// reads oldest to newest and the running total accumulates forward in time
|
||||
// rather than backward. Sort here, before shortDate() drops the year and makes
|
||||
// the labels unsortable.
|
||||
const perInterval = useMemo<SavingsPoint[]>(
|
||||
() =>
|
||||
results.map((d) => ({
|
||||
date: shortDate(d.date),
|
||||
Compression: compressionOf(d.metrics),
|
||||
"Prompt caching": cachingOf(d.metrics),
|
||||
})),
|
||||
[...results]
|
||||
.sort((a, b) => a.date.localeCompare(b.date))
|
||||
.map((d) => ({
|
||||
date: shortDate(d.date),
|
||||
Compression: compressionOf(d.metrics),
|
||||
"Prompt caching": cachingOf(d.metrics),
|
||||
})),
|
||||
[results],
|
||||
);
|
||||
|
||||
// Cumulative anchors on a synthetic $0 point at the range start so a short
|
||||
// range (down to a single day) rises from zero instead of floating as one dot.
|
||||
const overTime = useMemo(() => {
|
||||
if (accumulation !== "cumulative") return perInterval;
|
||||
const startLabel = startTime ? shortDate(localIsoDay(startTime)) : "";
|
||||
return withStartAnchor(toCumulative(perInterval), startLabel);
|
||||
}, [accumulation, perInterval, startTime]);
|
||||
|
||||
const intervalLabel = "Per day";
|
||||
const rangeLabel = formatRangeLabel(startTime ?? undefined, endTime ?? undefined);
|
||||
const savingsSubtitle = [
|
||||
accumulation === "cumulative" ? "Running total saved" : `Saved ${intervalLabel.toLowerCase()}`,
|
||||
rangeLabel,
|
||||
]
|
||||
.filter(Boolean)
|
||||
.join(" \u00b7 ");
|
||||
|
||||
const byDriver = useMemo(
|
||||
() =>
|
||||
[
|
||||
|
|
@ -159,16 +199,44 @@ const UsageTab: React.FC<UsageTabProps> = ({ accessToken, activity }) => {
|
|||
<div className="grid grid-cols-1 gap-6 lg:grid-cols-3">
|
||||
<Card className="lg:col-span-2">
|
||||
<CardHeader>
|
||||
<CardTitle>Savings over time</CardTitle>
|
||||
<div className="flex flex-wrap items-start justify-between gap-3">
|
||||
<div>
|
||||
<CardTitle>Savings</CardTitle>
|
||||
<p className="text-sm text-muted-foreground">{savingsSubtitle}</p>
|
||||
</div>
|
||||
<div className="flex items-center gap-4">
|
||||
<CustomLegend categories={SAVINGS_SERIES} colors={SAVINGS_COLORS} />
|
||||
<Tabs value={accumulation} onValueChange={(value) => setAccumulation(value as SavingsAccumulation)}>
|
||||
<TabsList>
|
||||
<TabsTrigger value="cumulative">Cumulative</TabsTrigger>
|
||||
<TabsTrigger value="per-interval">{intervalLabel}</TabsTrigger>
|
||||
</TabsList>
|
||||
</Tabs>
|
||||
</div>
|
||||
</div>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
<AreaChart
|
||||
data={overTime}
|
||||
index="date"
|
||||
categories={["Compression", "Prompt caching"]}
|
||||
colors={["emerald", "blue"]}
|
||||
valueFormatter={usd}
|
||||
/>
|
||||
{accumulation === "cumulative" ? (
|
||||
<AreaChart
|
||||
data={overTime}
|
||||
index="date"
|
||||
categories={SAVINGS_SERIES}
|
||||
colors={SAVINGS_COLORS}
|
||||
valueFormatter={usd}
|
||||
showLegend={false}
|
||||
showDots={overTime.length <= MAX_POINTS_WITH_DOTS}
|
||||
/>
|
||||
) : (
|
||||
<BarChart
|
||||
data={overTime}
|
||||
index="date"
|
||||
categories={SAVINGS_SERIES}
|
||||
colors={SAVINGS_COLORS}
|
||||
stack
|
||||
valueFormatter={usd}
|
||||
showLegend={false}
|
||||
/>
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
<Card>
|
||||
|
|
|
|||
|
|
@ -2,7 +2,16 @@ import { describe, expect, it } from "vitest";
|
|||
|
||||
import type { DailyData, SpendMetrics } from "@/components/UsagePage/types";
|
||||
import type { ToolSpendDailyEntry, ToolSpendEntry } from "@/components/networking";
|
||||
import { buildDailyToolSeries, computeCacheLeakage, isAnthropicModel, topToolsBySpend } from "./costOptimizationUtils";
|
||||
import {
|
||||
buildDailyToolSeries,
|
||||
computeCacheLeakage,
|
||||
formatRangeLabel,
|
||||
isAnthropicModel,
|
||||
localIsoDay,
|
||||
toCumulative,
|
||||
topToolsBySpend,
|
||||
withStartAnchor,
|
||||
} from "./costOptimizationUtils";
|
||||
|
||||
const metrics = (overrides: Partial<SpendMetrics>): SpendMetrics => ({
|
||||
spend: 0,
|
||||
|
|
@ -221,3 +230,79 @@ describe("topToolsBySpend", () => {
|
|||
expect(topToolsBySpend(byTool, 2).map((t) => t.tool_name)).toEqual(["b", "c"]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("localIsoDay", () => {
|
||||
it("reads the date off the viewer's clock rather than shifting it to UTC", () => {
|
||||
expect(localIsoDay(new Date(2026, 6, 23, 23, 30))).toBe("2026-07-23");
|
||||
expect(localIsoDay(new Date(2026, 0, 5, 0, 30))).toBe("2026-01-05");
|
||||
});
|
||||
});
|
||||
|
||||
describe("toCumulative", () => {
|
||||
const point = (date: string, compression: number, caching: number) => ({
|
||||
date,
|
||||
Compression: compression,
|
||||
"Prompt caching": caching,
|
||||
});
|
||||
|
||||
it("turns each reading into everything saved up to that point", () => {
|
||||
const running = toCumulative([point("Jul 1", 1, 10), point("Jul 2", 2, 20), point("Jul 3", 3, 30)]);
|
||||
expect(running.map((p) => p.Compression)).toEqual([1, 3, 6]);
|
||||
expect(running.map((p) => p["Prompt caching"])).toEqual([10, 30, 60]);
|
||||
});
|
||||
|
||||
it("accumulates each driver on its own, so one flat series cannot lift the other", () => {
|
||||
const running = toCumulative([point("Jul 1", 0, 5), point("Jul 2", 0, 5)]);
|
||||
expect(running.map((p) => p.Compression)).toEqual([0, 0]);
|
||||
expect(running.map((p) => p["Prompt caching"])).toEqual([5, 10]);
|
||||
});
|
||||
|
||||
it("never falls, even across a quiet interval", () => {
|
||||
const running = toCumulative([point("Jul 1", 4, 0), point("Jul 2", 0, 0), point("Jul 3", 1, 0)]);
|
||||
expect(running.map((p) => p.Compression)).toEqual([4, 4, 5]);
|
||||
});
|
||||
|
||||
it("keeps the labels and length of the readings it was given", () => {
|
||||
const running = toCumulative([point("9am", 1, 1), point("10am", 1, 1)]);
|
||||
expect(running.map((p) => p.date)).toEqual(["9am", "10am"]);
|
||||
expect(toCumulative([])).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("withStartAnchor", () => {
|
||||
const point = (date: string, compression: number, caching: number) => ({
|
||||
date,
|
||||
Compression: compression,
|
||||
"Prompt caching": caching,
|
||||
});
|
||||
|
||||
it("lifts a single-day cumulative off a floating dot by prepending a $0 origin", () => {
|
||||
const anchored = withStartAnchor([point("Jul 24", 12, 30)], "Jul 24");
|
||||
expect(anchored).toEqual([point("Jul 24", 0, 0), point("Jul 24", 12, 30)]);
|
||||
});
|
||||
|
||||
it("starts the range at zero without disturbing the running totals that follow", () => {
|
||||
const anchored = withStartAnchor([point("Jul 16", 5, 1), point("Jul 17", 9, 4)], "Jul 16");
|
||||
expect(anchored.map((p) => p.Compression)).toEqual([0, 5, 9]);
|
||||
expect(anchored.map((p) => p["Prompt caching"])).toEqual([0, 1, 4]);
|
||||
});
|
||||
|
||||
it("leaves an empty series alone so the chart's own no-data state can show", () => {
|
||||
expect(withStartAnchor([], "Jul 24")).toEqual([]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("formatRangeLabel", () => {
|
||||
it("reads as a range across days", () => {
|
||||
expect(formatRangeLabel(new Date(2026, 6, 16), new Date(2026, 6, 23))).toBe("Jul 16 \u2013 Jul 23");
|
||||
});
|
||||
|
||||
it("collapses to one date when both ends are the same day", () => {
|
||||
expect(formatRangeLabel(new Date(2026, 6, 23), new Date(2026, 6, 23))).toBe("Jul 23");
|
||||
});
|
||||
|
||||
it("is empty until both ends are picked", () => {
|
||||
expect(formatRangeLabel(undefined, new Date(2026, 6, 23))).toBe("");
|
||||
expect(formatRangeLabel(new Date(2026, 6, 23), undefined)).toBe("");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -149,3 +149,63 @@ const seedPoint = (date: string, toolNames: readonly string[]): DailyToolSpendPo
|
|||
|
||||
export const topToolsBySpend = (byTool: readonly ToolSpendEntry[], limit = 8): ToolSpendEntry[] =>
|
||||
[...byTool].sort((a, b) => b.spend - a.spend).slice(0, limit);
|
||||
|
||||
export const localIsoDay = (d: Date): string =>
|
||||
`${d.getFullYear()}-${String(d.getMonth() + 1).padStart(2, "0")}-${String(d.getDate()).padStart(2, "0")}`;
|
||||
|
||||
export type SavingsAccumulation = "cumulative" | "per-interval";
|
||||
|
||||
// A type alias, not an interface: only aliases get the implicit index signature
|
||||
// that the chart wrappers' `Record<string, unknown>` datum bound requires.
|
||||
export type SavingsPoint = {
|
||||
date: string;
|
||||
Compression: number;
|
||||
"Prompt caching": number;
|
||||
};
|
||||
|
||||
export const SAVINGS_SERIES = ["Compression", "Prompt caching"] as const;
|
||||
|
||||
/**
|
||||
* Running total of each series across the selected window. The total restarts
|
||||
* at the beginning of the range rather than carrying in earlier spend, which is
|
||||
* what "running total saved, <range>" claims on the card.
|
||||
*/
|
||||
export const toCumulative = (points: readonly SavingsPoint[]): SavingsPoint[] =>
|
||||
points.reduce<SavingsPoint[]>((acc, point) => {
|
||||
const previous = acc[acc.length - 1];
|
||||
return [
|
||||
...acc,
|
||||
{
|
||||
date: point.date,
|
||||
Compression: (previous?.Compression ?? 0) + point.Compression,
|
||||
"Prompt caching": (previous?.["Prompt caching"] ?? 0) + point["Prompt caching"],
|
||||
},
|
||||
];
|
||||
}, []);
|
||||
|
||||
/**
|
||||
* Prepend a synthetic $0 point at the start of the range so the cumulative line
|
||||
* rises from zero instead of floating as a single dot. The daily rollup only
|
||||
* resolves whole days, so a one-day range would otherwise be one point; with the
|
||||
* anchor it reads as "start of range $0 climbing to the range's running total".
|
||||
* An empty series is left untouched so the chart's own "No data" state shows.
|
||||
*/
|
||||
export const withStartAnchor = (cumulative: readonly SavingsPoint[], startLabel: string): SavingsPoint[] =>
|
||||
cumulative.length === 0
|
||||
? [...cumulative]
|
||||
: [{ date: startLabel, Compression: 0, "Prompt caching": 0 }, ...cumulative];
|
||||
|
||||
/** "Jul 16 – Jul 23", collapsing to a single date when the range is one day. */
|
||||
export const formatRangeLabel = (from: Date | undefined, to: Date | undefined): string => {
|
||||
if (!from || !to) return "";
|
||||
const short = (d: Date) => d.toLocaleDateString("en-US", { month: "short", day: "numeric" });
|
||||
const start = short(from);
|
||||
const end = short(to);
|
||||
return start === end ? start : `${start} – ${end}`;
|
||||
};
|
||||
|
||||
/**
|
||||
* Dots mark each reading, as in the design. Past this many readings they crowd
|
||||
* into a solid band and stop being readable, so the line carries it alone.
|
||||
*/
|
||||
export const MAX_POINTS_WITH_DOTS = 31;
|
||||
|
|
|
|||
|
|
@ -40,4 +40,12 @@ describe("AreaChart", () => {
|
|||
expect(area.getAttribute("fill")).toMatch(/^url\(#fill-/);
|
||||
}
|
||||
});
|
||||
|
||||
it("marks each reading with a dot only when asked", () => {
|
||||
const withoutDots = render(<AreaChart data={data} index="date" categories={["tokens"]} />);
|
||||
expect(withoutDots.container.querySelectorAll("circle.recharts-dot")).toHaveLength(0);
|
||||
|
||||
const withDots = render(<AreaChart data={data} index="date" categories={["tokens"]} showDots />);
|
||||
expect(withDots.container.querySelectorAll("circle.recharts-dot").length).toBeGreaterThanOrEqual(data.length);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ export type AreaChartProps<TDatum extends Record<string, unknown>> = {
|
|||
showLegend?: boolean;
|
||||
showGridLines?: boolean;
|
||||
showTooltip?: boolean;
|
||||
showDots?: boolean;
|
||||
customTooltip?: ChartTooltipComponent;
|
||||
className?: string;
|
||||
style?: React.CSSProperties;
|
||||
|
|
@ -32,6 +33,7 @@ export function AreaChart<TDatum extends Record<string, unknown>>({
|
|||
showLegend = true,
|
||||
showGridLines = true,
|
||||
showTooltip = true,
|
||||
showDots = false,
|
||||
customTooltip,
|
||||
className,
|
||||
style,
|
||||
|
|
@ -94,7 +96,7 @@ export function AreaChart<TDatum extends Record<string, unknown>>({
|
|||
strokeWidth={2}
|
||||
fill={`url(#fill-${gradientId}-${i})`}
|
||||
fillOpacity={1}
|
||||
dot={false}
|
||||
dot={showDots ? { r: 3.5, strokeWidth: 2, stroke: fills[i], fill: "var(--background, #fff)" } : false}
|
||||
isAnimationActive={false}
|
||||
/>
|
||||
))}
|
||||
|
|
|
|||
|
|
@ -20,6 +20,13 @@ describe("BarChart", () => {
|
|||
expect(fills).toEqual(new Set(["var(--color-green-500, #22c55e)", "var(--color-red-500, #ef4444)"]));
|
||||
});
|
||||
|
||||
it("renders the No data placeholder instead of a chart when data is empty", () => {
|
||||
const { container, getByText } = render(<BarChart data={[]} index="date" categories={["passed"]} />);
|
||||
|
||||
expect(getByText("No data")).toBeTruthy();
|
||||
expect(container.querySelector('[data-slot="chart"]')).toBeNull();
|
||||
});
|
||||
|
||||
it("falls back to the tremor default color cycle when no colors are passed", () => {
|
||||
const { container } = render(<BarChart data={data} index="date" categories={["passed", "blocked"]} />);
|
||||
|
||||
|
|
|
|||
|
|
@ -46,6 +46,17 @@ export function BarChart<TDatum extends Record<string, unknown>>({
|
|||
className,
|
||||
style,
|
||||
}: BarChartProps<TDatum>) {
|
||||
if (data.length === 0) {
|
||||
return (
|
||||
<div
|
||||
className={cn("flex h-80 w-full items-center justify-center rounded-lg border border-dashed", className)}
|
||||
style={style}
|
||||
>
|
||||
<p className="text-sm text-muted-foreground">No data</p>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
const fills = categoryFills(categories.length, colors);
|
||||
const config: ChartConfig = Object.fromEntries(categories.map((category) => [category, { label: category }]));
|
||||
const vertical = layout === "vertical";
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue