Skip to content

feat(lens): track worker spend through virtual keys - #43989

Merged
moe-berri merged 7 commits into
mainfrom
litellm_lens_reliability_followup
Oct 1, 2026
Merged

moe-berri merged 7 commits into
mainfrom
litellm_lens_reliability_followup

Conversation

@moe-berri

@moe-berri moe-berri commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor

TLDR

Problem this solves:

  • Lens analysis lacks attributable virtual-key spend and permissions
  • Competing worker claims can duplicate scans using shared credentials

How it solves it:

  • Select or create a billing key during worker setup
  • Reuse normal proxy authorization, limits, logging, and inference
  • Require a successful database update before returning a claim

Intentional product change: existing workers must select a billing key in Set up analysis before resuming; unbilled inference no longer runs

User Flow

Before: Lens analysis has no associated virtual key

  1. Open /ui/lens/ and select Set up analysis
  2. Generate a Docker command and start the worker
  3. Run an investigation and open Virtual Keys
  4. No key shows the worker's analysis spend

Before worker setup

After: analysis is attributed to the worker's selected key

  1. Open /ui/lens/ and select Set up analysis
  2. Select an existing key or click Create worker key, then start the worker
  3. Run an investigation and open Virtual Keys
  4. The selected key shows spend and its permissions and limits apply

After worker setup

Screenshots / Proof of Fix

Local proxy with a real gateway-backed model. The ordinary virtual key is selected by its ID; only the separate worker credential leaves the proxy. Representative request: POST /engine/worker/{lens_id}/{run_id}/model with {"prompt":"Return exactly this JSON: {\"observations\": []}","purpose":"extract"}

Before (2dc58e9)

  1. Create a virtual key and register a worker, supplying its key ID
  2. Call the worker model endpoint: HTTP 200, Lens cost $0.000918, virtual-key spend $0

After (2e1b0fa)

  1. Create a virtual key and register a worker, supplying analysis_key_id
  2. Call the worker model endpoint with a real model: HTTP 200, Lens cost $0.000858, virtual-key spend $0.000858
  3. Block the key: HTTP 400; exhaust its budget: HTTP 402; restrict its model: HTTP 403; delete it: HTTP 401
  4. Race eight claims through the real proxy: exactly one receives the scan
  5. Change the worker's billing key: spend lands on the replacement; its RPM limit returns HTTP 429; worker revocation returns HTTP 401
  6. Enable prompt storage: ordinary requests retain their prompts, Lens analysis prompts and responses remain redacted; allowlisted workers succeed and other addresses fail
  7. Fill the Docker worker's 64 KiB temporary mount: the scan fails clearly, files are cleaned, and the next scan completes

Implementation and validation

The large authentication diff extracts existing virtual-key validation into a shared function, preserving its check order. Normal inference and Lens both call it; a bearer containing a stored key hash is still rejected. Four existing source-inspection tests now follow the extracted function with their assertions unchanged

101 Lens tests, all three live HTTP billing/privacy cases, and a fresh real paid model call pass at the latest tip, with exact Lens/key cost agreement. Local make check passes. Earlier validation also passed 197 focused Python tests, 11 dashboard tests, and three live HTTP billing/privacy cases. All branch-required CI checks, Codecov, and bot reviews now pass. Six previous jobs timed out during dependency installation; the Lens protocol-ordering test failure is fixed without changing its assertions. The published amd64 worker passed CI and is pinned by verified digest

The IP regression fails against the previous implementation and passes with the fix: an untrusted worker cannot impersonate an allowed address through forwarded headers, trusted forwarding still works, and a forged prefix cannot override the configured hop count. Resolution reuses the existing trusted-proxy settings

The prior database shard passed all 929 tests, but its coverage uploader used a signing-key URL that now returns 404. The Lens coverage step now pins the official Codecov action v7.1.1 and CLI v11.3.1, retaining signature and checksum validation. Coverage includes shared authorization code with repository-relative filenames; no thresholds or test assertions were relaxed. The updated uploader succeeded in CI with a verified GPG signature and queued the complete report for processing

Current-head Codecov passes at 83.66% against an 81.73% target. Greptile is 5/5 and Bugbot reports no issues. Veria confirms all previously flagged issues are resolved and reports no open security concerns

The Python CodeQL job exceeds the 2 GiB result limit in LogInjection.ql. The same failure exists on main, so no security query is disabled here. The unchanged MCP discovery test timed out once; its targeted CI rerun now passes. All ten branch-required CI checks and Codecov pass. The OAuth job awaits approval for the e2e-changed environment

The complete Docker investigation also finished with one reviewed run and one saved finding in 11.53 seconds; Lens cost and the key-spend increase both measured $0.028827 at runtime revision bf6ab45

The follow-up also separates accepted patterns from issue identities so feedback cannot dismiss a distinct issue, and reports temporary-storage failures without abandoning the worker process

Pre-Submission checklist

  • Added meaningful regression tests
  • Focused tests pass locally
  • Required CI and coverage pass
  • Scope limited to Lens worker billing and reliability
  • Greptile confidence at least 4/5 and acceptable Veria/Bugbot reviews

Caveats

Severe

  • Existing workers need billing before new analysis; heartbeat/results remain available
  • Forwarded IP allowlists require configured trusted proxy ranges

Medium

  • Lease recovery can repeat analysis after worker failure
  • Real 1,000–2,000-agent scale validation remains a separate follow-up

Note

High Risk
Changes authentication, budget enforcement, and billing paths for Lens inference, plus IP allowlist behavior for workers—security- and spend-critical areas requiring careful rollout for existing workers without billing keys.

Overview
Lens workers must be tied to an analysis virtual key at registration (or via a new billing-key API). Model calls during scans no longer bypass the proxy router; they run through the same virtual-key authorization, budgets, rate limits, and spend logging as normal inference, with prompts/responses redacted in spend logs while cost still lands on the chosen key.

Proxy auth is refactored so shared validate_resolved_virtual_key and authorize_internal_virtual_key can authorize server-side Lens billing without treating a stored key hash as a bearer token. Worker model requests honor allowed_ips using existing trusted-proxy (mcp_trusted_proxy_ranges) settings.

Worker packaging and reliability move writable trace storage from a baked-in Docker volume to an explicit tmpfs /tmp (default 1 GiB) in generated commands and Compose; full-disk failures surface as failed scans with cleanup instead of silent truncation. Job claim uses optimistic updates that return no claim when the engine row did not change, reducing duplicate claims under concurrency. Finding merge identity now includes kind so issues and patterns with the same title stay separate.

The Lens dashboard Set up analysis flow adds key pick/create UI, treats workers without a billing key as not ready, and updates docs/CI (storage smoke test, pinned worker image digest, Codecov action for postgres coverage).

Reviewed by Cursor Bugbot for commit bfdf7fa. Bugbot is set up for automated code reviews on this repo. Configure here.

@moe-berri

Copy link
Copy Markdown
Contributor Author

@greptileai Please review the latest worker billing changes, especially shared virtual-key authorization, spend attribution, and concurrent job claims.

@moe-berri

Copy link
Copy Markdown
Contributor Author

@veria-ai Please review worker authentication, stored billing-key selection, authorization checks, revocation races, and concurrent scan ownership on this revision.

@moe-berri

Copy link
Copy Markdown
Contributor Author

bugbot run Please check the latest Lens worker changes for authorization bypasses, billing mistakes, duplicate claims, and storage recovery regressions.

@greptile-apps

greptile-apps Bot commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 5/5

[High risk] Adds billing tracking to worker analysis through API keys.

The PR appears safe to merge; no new actionable issue or outstanding previous finding was identified.

Summary

The PR assigns virtual keys to Lens workers, routes analysis through normal proxy authorization and spend tracking, and improves claim and temporary-storage recovery. Since the previous review, it connects the internal request to worker-disconnect handling and extends the privacy regression test to cover both disconnect settings.

Reviews (7) · Last reviewed commit: "fix(lens): yield to worker disconnects a..."

Comment thread litellm/proxy/engine/billing.py
Comment thread litellm/proxy/engine/endpoints.py Outdated

@cursor cursor Bot left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale Bugbot comment from a previous run.

Comment thread litellm/proxy/engine/billing.py
@veria-ai

veria-ai Bot commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor

PR overview

All previously flagged issues have been addressed. No open security concerns remain on this pull request.

Security review

No open security issues remain on this pull request.

Fixed/addressed: 2 · PR risk: 0/10

@codspeed

codspeed Bot commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor

Merging this PR will not alter performance

✅ 31 untouched benchmarks


Comparing litellm_lens_reliability_followup (bfdf7fa) with main (3a11192)1

Open in CodSpeed

Footnotes

  1. No successful run was found on main (4b06d04) during the generation of this report, so 3a11192 was used instead as the comparison base. There might be some changes unrelated to this pull request in this report. ↩

@moe-berri

Copy link
Copy Markdown
Contributor Author

@greptileai Recheck bf6ab45: worker network context is preserved, heartbeat and results remain accessible during billing migration, and tests cover both.

@moe-berri

Copy link
Copy Markdown
Contributor Author

@veria-ai Please recheck bf6ab45. Analysis payloads are now redacted through existing privacy controls, with real stored-log and spend assertions.

@moe-berri

Copy link
Copy Markdown
Contributor Author

bugbot run Please review bf6ab45, including forwarded client context, prompt redaction, normal billing cost reconciliation, and existing-worker result submission.

@codecov

codecov Bot commented Oct 1, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 83.46457% with 42 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
litellm/proxy/auth/user_api_key_auth.py 82.06% 26 Missing ⚠️
litellm/proxy/engine/billing.py 83.67% 8 Missing ⚠️
litellm/proxy/engine/endpoints.py 77.77% 6 Missing ⚠️
litellm/proxy/engine/inference.py 88.23% 2 Missing ⚠️

📢 Thoughts on this report? Let us know!

Comment thread litellm/proxy/auth/user_api_key_auth.py
@moe-berri

Copy link
Copy Markdown
Contributor Author

@greptileai The latest tip adds a regression proving old workers can submit results. Please recheck the resolved network-context and billing-migration findings.

@moe-berri

Copy link
Copy Markdown
Contributor Author

@veria-ai The latest tip retains the privacy fix and adds worker-upgrade coverage. Please verify the stored-content disclosure is resolved.

@moe-berri

Copy link
Copy Markdown
Contributor Author

bugbot run Please review the latest tip; runtime fixes include billing privacy, client IP checks, and accurate cost attribution.

@cursor cursor Bot left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale Bugbot comment from a previous run.

Comment thread litellm/proxy/engine/billing.py Outdated

@cursor cursor Bot left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale Bugbot comment from a previous run.

@moe-berri

Copy link
Copy Markdown
Contributor Author

@greptileai Please review b999302. Worker IP validation and protocol ordering are fixed; the flagged membership lookup is unchanged from main.

@moe-berri

Copy link
Copy Markdown
Contributor Author

@veria-ai Please review b999302 for remaining security issues. Worker IP spoofing is blocked, and billing privacy regression coverage remains intact.

@moe-berri

Copy link
Copy Markdown
Contributor Author

bugbot run Please verify b999302 resolves forwarded-IP spoofing. Regression tests cover untrusted headers, trusted forwarding, and forged address chains.

Comment thread litellm/proxy/engine/billing.py

@cursor cursor Bot left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale Bugbot comment from a previous run.

@moe-berri

Copy link
Copy Markdown
Contributor Author

@greptileai Please review 2e1b0fa. Trusted proxy requirements are documented; accepting untrusted forwarded addresses would restore the confirmed allowlist bypass.

@moe-berri

Copy link
Copy Markdown
Contributor Author

@veria-ai Please review 2e1b0fa for remaining issues. Runtime retains trusted network validation and billing redaction; the latest commit documents deployment requirements.

@moe-berri

Copy link
Copy Markdown
Contributor Author

bugbot run Please review 2e1b0fa. Runtime is unchanged from your clean review; worker allowlist setup requirements are now documented.

@cursor cursor Bot left a comment •

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Stale Bugbot comment from a previous run.

@moe-berri

Copy link
Copy Markdown
Contributor Author

@veria-ai Please reassess discussion_r4152598001: inference.py applies inherit_message_logging_privacy(True), and test_worker_spend_logs_do_not_expose_investigation_content verifies stored rows exclude private prompts and responses while retaining spend.

@moe-berri
moe-berri marked this pull request as ready for review October 1, 2026 08:21
@moe-berri
moe-berri requested a review from a team October 1, 2026 08:21
Comment thread litellm/proxy/engine/billing.py Outdated
@moe-berri
moe-berri marked this pull request as draft October 1, 2026 16:12
@moe-berri

Copy link
Copy Markdown
Contributor Author

@greptileai

@moe-berri

Copy link
Copy Markdown
Contributor Author

@veria-ai

@moe-berri

Copy link
Copy Markdown
Contributor Author

bugbot run

@moe-berri
moe-berri marked this pull request as ready for review October 1, 2026 16:24

@cursor cursor Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

✅ Bugbot reviewed your changes and found no new issues!

Comment @cursor review or bugbot run to trigger another review on this PR

Reviewed by Cursor Bugbot for commit bfdf7fa. Configure here.

@moe-berri
moe-berri merged commit d9f7324 into main Oct 1, 2026
102 of 105 checks passed
@moe-berri
moe-berri deleted the litellm_lens_reliability_followup branch October 1, 2026 16:47

This branch is waiting to be deployed

1 waiting deployment
e2e-changed — bfdf7fa8 Waiting Oct 1, 2026 by moe-berri via oauth #2313
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants