Skip to content

Feat/spend tracking#3

Open
gettingstuffdone wants to merge 4 commits into
endemics:mainfrom
gettingstuffdone:feat/spend-tracking
Open

Feat/spend tracking#3
gettingstuffdone wants to merge 4 commits into
endemics:mainfrom
gettingstuffdone:feat/spend-tracking

Conversation

@gettingstuffdone

Copy link
Copy Markdown

Description:
This PR implements end-to-end query cost tracking and budget enforcement between the Go gateway and Python server:

  • Gateway Cost Recording:
    • Updated the Proxy to capture response bodies using a bodyRecorder.
    • Supports recording spend from the X-Limnos-Cost-USD header (preferred) or by scraping est. $X.XX from the Markdown response
      (fallback).
    • Fixed a path routing bug where the /mcp prefix wasn't being correctly stripped before proxying to workers.
  • Server Cost Reporting:
    • Added FastAPI middleware to automatically inject the X-Limnos-Cost-USD header into HTTP responses.
    • Implemented a ContextVar (current_query_cost) to safely communicate the estimated cost from the tool logic to the middleware
      across async boundaries.
  • Testing:
    • Added server/tests/test_query_helpers.py to verify the cost ContextVar behavior.
    • Updated gateway/internal/mcp/proxy_test.go to support the new authenticated proxy signature.
  • Dependency: This PR is based on feat/server-refactor.

Note: Since these branches are sequential, this PR will show the changes from the refactor as well unless feat/server-refactor is merged first.

@endemics endemics left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

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

Thanks!

@endemics

Copy link
Copy Markdown
Owner

There is a small linting issue with the golang code in the PR, would you be able to have a look and fix it? Thanks

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