fix: update API token last_used and log usage stats by shoummu1 · Pull Request #2711 · IBM/mcp-context-forge

shoummu1 · 2026-02-05T13:10:59Z

📌 Summary

API tokens "Last Used" timestamp was always showing "Never" and "Usage Stats" showed no data, even after tokens were actively used. This PR fixes both issues by updating the last_used timestamp during JWT validation and implementing middleware to log all API token usage.

🔗 Related Issue

Closes #2572

🐞 Root Cause

Two root causes identified:

last_used not updating: API tokens created via /tokens are JWT tokens. When used, they pass JWT validation directly (auth.py:612) and the auth flow never reaches _lookup_api_token_sync() (auth.py:836) where last_used was updated. The database lookup was designed as a fallback for non-JWT tokens, but all API tokens are JWTs.
Usage Stats empty: The log_token_usage() method in TokenCatalogService was defined but never called from production code.

💡 Fix Description

1. Last Used Timestamp Tracking

Added _update_api_token_last_used_sync(jti: str) helper function in auth.py (line 325)
Calls this function during JWT validation in _set_auth_method_from_payload():
- For standard API tokens with auth_provider == "api_token" (line 548)
- For legacy API tokens via DB lookup fallback (line 564)
Uses asyncio.to_thread() for non-blocking database updates
Thread-safe implementation with fresh_db_session()

2. Usage Statistics Logging

Implemented TokenUsageMiddleware in mcpgateway/middleware/token_usage_middleware.py
Middleware logs all API token requests after completion
Captures: endpoint, method, status code, response time, IP address, user agent
Only activates when request.state.auth_method == "api_token"
Registered in main.py after AuthContextMiddleware to ensure auth context is available
Falls back to JWT decoding if JTI not in request.state

3. Test Coverage

Added unit tests for _update_api_token_last_used_sync() function
Added integration tests for API token auth flow with last_used updates
Added comprehensive middleware tests (8+ test cases)
Tests cover: standard tokens, legacy tokens, error handling, fallback paths

🧪 Verification

Check	Command	Status
Lint suite	`make lint`	✅
Unit tests	`make test`	✅
Coverage ≥ 90 %	`make coverage`	✅
Manual regression no longer fails	See steps below	✅

📐 MCP Compliance (if relevant)

Matches current MCP spec
No breaking change to MCP clients

✅ Checklist

Code formatted (make black isort pre-commit)
No secrets/credentials committed

crivetimihai · 2026-02-10T10:58:08Z

Thanks for fixing the last_used tracking, @shoummu1! The root cause analysis is spot-on — API tokens being JWTs means they skip the DB lookup path where last_used was being updated. The test coverage is thorough and the error handling is properly non-blocking.

My main concern is performance at scale:

DB write per request: _update_api_token_last_used_sync runs a SELECT + UPDATE + COMMIT on every API token request. Consider rate-limiting the update (e.g., only update if last update was >N minutes ago, using an in-memory cache).
No feature flag: The TokenUsageMiddleware is always registered. Other middleware in the codebase (observability) has enable/disable configuration — this should too.
Session management: The middleware uses SessionLocal() directly instead of fresh_db_session() context manager, which could leak sessions on edge-case exceptions. Please align with the codebase pattern.
BaseHTTPMiddleware: Starlette's maintainers recommend raw ASGI middleware for production — BaseHTTPMiddleware buffers entire response bodies. Worth considering for a middleware that runs on every request.

crivetimihai

Thanks for this PR, @shoummu1! The separation of last_used update from usage logging and the rate-limiting approach are good ideas. A few issues:

1. Unbounded in-memory cache (auth.py:143-168)
_update_api_token_last_used_sync._cache is a plain dict that grows with every unique JTI. In deployments with many API tokens, this leaks memory indefinitely. Consider using functools.lru_cache or a bounded dict.

2. await on synchronous DB inside ASGI middleware (token_usage_middleware.py)
log_token_usage is async def but performs synchronous SQLAlchemy db.add()/db.commit() — this blocks the event loop. In auth.py you correctly used asyncio.to_thread() for the sync path; the middleware should do the same.

3. Triple commit per request
token_service.log_token_usage() calls self.db.commit(), then the middleware explicitly calls db.commit(), and then fresh_db_session.__exit__ auto-commits. Since the context manager auto-commits on success, the explicit commits are redundant.

4. Redis client never retries after failure (auth.py:75)
If _get_sync_redis_client fails, _SYNC_REDIS_CLIENT stays None permanently for the process lifetime. Redis becoming temporarily unavailable means the in-memory fallback is permanent — consider a retry mechanism or TTL on the failure state.

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

- Store JTI in request.state.jti for standard API tokens in _set_auth_method_from_payload, ensuring middleware can access it without re-decoding the JWT on cached/batched auth paths - Remove redundant db.commit() in TokenUsageMiddleware since log_token_usage() already commits the transaction - Add bounded eviction (max 1024 entries) to in-memory rate-limit cache in _update_api_token_last_used_sync to prevent unbounded growth - Add test coverage for cache eviction and JTI propagation - Fix formatting (missing space, unused import, black/isort) Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>

Two admin tests failed because .env sets MCPGATEWAY_UI_HIDE_SECTIONS which leaks into test execution, causing section-dependent code paths (resource loading, gRPC services) to be skipped. Patch settings to ensure no sections are hidden during these tests. Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>

crivetimihai · 2026-02-16T23:33:25Z

Review Changes

Rebased onto latest main (no conflicts) and added two fix commits after code review:

1. `fix: address review findings in token usage tracking`

Bug fix — missing JTI propagation (mcpgateway/auth.py)

Added request.state.jti = jti in _set_auth_method_from_payload for the standard api_token auth path. Without this, the token usage middleware had to re-decode the JWT on every request to extract the JTI, even though auth already decoded it.

Unbounded cache growth (mcpgateway/auth.py)

The in-memory rate-limit cache (_update_api_token_last_used_sync._cache) could grow without bound. Added eviction logic: when cache reaches 1024 entries, evict the oldest 50%.

Redundant double commit (mcpgateway/middleware/token_usage_middleware.py)

Removed db.commit() after token_service.log_token_usage() since log_token_usage() already commits internally. The double commit was harmless but unnecessary.

Test updates

test_auth.py: Added assertion that request.state.jti is set after API token auth; added cache eviction test (pre-fills 1024 entries, verifies eviction); added no-JTI graceful handling test; fixed formatting (=MagicMock() → = MagicMock()).
test_token_usage_middleware.py: Removed stale mock_db.commit.assert_called_once() assertions that expected the now-removed double commit.

2. `fix: patch hidden sections env leakage in admin UI tests`

Pre-existing test failures — two admin tests (test_admin_ui_with_service_failures, test_admin_ui_tuple_unwrap_filtering_and_grpc_success) were failing because .env sets MCPGATEWAY_UI_HIDE_SECTIONS with nearly all sections, which leaks into test execution and causes section-dependent code paths (resource loading, gRPC services) to be skipped entirely.

Fixed by patching settings.mcpgateway_ui_hide_sections to [] in both tests so they exercise the intended code paths regardless of local .env configuration.

crivetimihai · 2026-02-16T23:45:10Z

Verification Against Running Instance (localhost:8080)

Tested the feature end-to-end against the Docker Compose deployment (3 gateway replicas, PostgreSQL, Redis, nginx).

Setup

Created API token via POST /tokens — last_used: null initially
Token has auth_provider: api_token and token_use: api claims

Results

last_used tracking — Working correctly:

First API request updated last_used from null to 2026-02-16T23:38:48Z
Subsequent requests within 5-minute rate-limit window correctly skip DB writes (by design: token_last_used_update_interval_minutes: 5)
After the rate-limit window elapsed, last_used updated to 2026-02-16T23:44:05Z

Usage logging — Working correctly:

10 API token requests logged in token_usage_logs table
All logged with correct: token_jti, user_email, endpoint, method, status_code, response_time_ms, ip_address, timestamp

GET /tokens/{id}/usage endpoint returns accurate aggregated stats:

total_requests: 10
successful_requests: 10
blocked_requests: 0
success_rate: 1.0
avg_response_time: 26.3ms
top_endpoints: /gateways: 2, /prompts: 2, /resources: 2, /servers: 2, /tools: 2

Auth method propagation — Confirmed JTI flows through all auth paths (cached, standard) via request.state.jti, enabling the middleware to log without re-decoding the JWT.

Rate limiting — Rate-limit cache correctly prevents excessive DB writes for last_used updates while usage logging captures every request.

- Fix get_current_user request param type (Optional[object] → Request) so FastAPI auto-injects it, enabling last_used and usage logging - Add 30s backoff in _get_sync_redis_client after connection failures to prevent retry storms when Redis is down - Move in-memory cache from hasattr-based function attributes to module-level globals, eliminating initialization race condition - Store user_email in request state for DB-fallback opaque tokens so usage middleware can log them without JWT decode Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>

crivetimihai

Addressing review findings with 4 targeted fixes. See inline comments for details on each change.

crivetimihai · 2026-02-17T00:36:01Z

mcpgateway/auth.py

 async def get_current_user(
    credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
-    request: Optional[object] = None,
+    request: Request = None,  # type: ignore[assignment]


Fix 1 (High priority): Changed request: Optional[object] = None → request: Request = None.

FastAPI's dependency injection matches parameters by type using lenient_issubclass. With Optional[object], the type check Request ≠ object fails, so FastAPI never auto-injects the request — it was always None. This caused _set_auth_method_from_payload to bail out immediately at line 813, skipping last_used updates, auth_method/jti state propagation, and usage logging for all 9 routes using Depends(get_current_user).

With Request type, FastAPI recognizes and auto-injects it. Manual callers (RBAC middleware, AuthContextMiddleware) already pass request=request explicitly, so they're unaffected.

crivetimihai · 2026-02-17T00:36:08Z

mcpgateway/auth.py

+    if _SYNC_REDIS_CLIENT is not None or not (config_settings.redis_url and config_settings.redis_url.strip() and config_settings.cache_type == "redis"):
+        return _SYNC_REDIS_CLIENT
+
+    # Backoff after recent failure (30 seconds)


Fix 2 (Medium priority): Added 30-second backoff after Redis connection failures.

Previously, when Redis was down, every call to _get_sync_redis_client() would retry redis.from_url() + ping() with a 2-second socket timeout. Under load this creates a retry storm. Now, after a failure we record _SYNC_REDIS_FAILURE_TIME and skip retries for 30 seconds. On successful connection, the failure timestamp is cleared.

crivetimihai · 2026-02-17T00:36:17Z

mcpgateway/auth.py

+
+# Module-level in-memory cache for last_used rate-limiting (fallback when Redis unavailable)
+_LAST_USED_CACHE: dict = {}
+_LAST_USED_CACHE_LOCK = threading.Lock()


Fix 3 (Low priority): Moved in-memory cache from hasattr-based lazy function attributes to module-level _LAST_USED_CACHE / _LAST_USED_CACHE_LOCK.

The previous pattern used hasattr(_update_api_token_last_used_sync, "_cache") for lazy init. Two threads could both see hasattr return False — if thread A creates _cache but hasn't created _cache_lock yet, thread B could skip init (hasattr on _cache returns True) and try to acquire the non-existent lock. Module-level initialization eliminates this race entirely.

crivetimihai · 2026-02-17T00:36:26Z

mcpgateway/middleware/token_usage_middleware.py

+        jti = state.get("jti") if state else None
+        user = state.get("user") if state else None
+        user_email = getattr(user, "email", None) if user else None
+        if not user_email:


Fix 4 (Low priority): Added user_email fallback from request state for DB-fallback opaque tokens.

The DB-fallback API token path (auth.py ~line 1310) sets request.state.auth_method and request.state.jti but not request.state.user. Without a user object, the middleware couldn't extract user_email and the JWT decode fallback fails for opaque tokens — silently dropping usage logs.

Now auth.py stores request.state.user_email directly for the DB-fallback path, and the middleware checks it as a fallback when the user object is absent.

* fix: update API token last_used and log usage stats Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * flake8 fix Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * additional changes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * remove duplicate last_used update in token usage logging Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * additional test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * added exception handling Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * fix coverage Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * fix: optimize token tracking with rate-limiting and raw ASGI middleware Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * additional test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * fix: address review findings in token usage tracking - Store JTI in request.state.jti for standard API tokens in _set_auth_method_from_payload, ensuring middleware can access it without re-decoding the JWT on cached/batched auth paths - Remove redundant db.commit() in TokenUsageMiddleware since log_token_usage() already commits the transaction - Add bounded eviction (max 1024 entries) to in-memory rate-limit cache in _update_api_token_last_used_sync to prevent unbounded growth - Add test coverage for cache eviction and JTI propagation - Fix formatting (missing space, unused import, black/isort) Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: patch hidden sections env leakage in admin UI tests Two admin tests failed because .env sets MCPGATEWAY_UI_HIDE_SECTIONS which leaks into test execution, causing section-dependent code paths (resource loading, gRPC services) to be skipped. Patch settings to ensure no sections are hidden during these tests. Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: resolve DI, retry storm, race, and opaque token tracking issues - Fix get_current_user request param type (Optional[object] → Request) so FastAPI auto-injects it, enabling last_used and usage logging - Add 30s backoff in _get_sync_redis_client after connection failures to prevent retry storms when Redis is down - Move in-memory cache from hasattr-based function attributes to module-level globals, eliminating initialization race condition - Store user_email in request state for DB-fallback opaque tokens so usage middleware can log them without JWT decode Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> --------- Signed-off-by: Shoumi <shoumimukherjee@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> Co-authored-by: Mihai Criveti <crivetimihai@gmail.com> Signed-off-by: Vishu Bhatnagar <vishu.bhatnagar@ibm.com>

* fix: update API token last_used and log usage stats Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * flake8 fix Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * additional changes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * remove duplicate last_used update in token usage logging Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * additional test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * added exception handling Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * fix coverage Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * fix: optimize token tracking with rate-limiting and raw ASGI middleware Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * additional test fixes Signed-off-by: Shoumi <shoumimukherjee@gmail.com> * fix: address review findings in token usage tracking - Store JTI in request.state.jti for standard API tokens in _set_auth_method_from_payload, ensuring middleware can access it without re-decoding the JWT on cached/batched auth paths - Remove redundant db.commit() in TokenUsageMiddleware since log_token_usage() already commits the transaction - Add bounded eviction (max 1024 entries) to in-memory rate-limit cache in _update_api_token_last_used_sync to prevent unbounded growth - Add test coverage for cache eviction and JTI propagation - Fix formatting (missing space, unused import, black/isort) Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: patch hidden sections env leakage in admin UI tests Two admin tests failed because .env sets MCPGATEWAY_UI_HIDE_SECTIONS which leaks into test execution, causing section-dependent code paths (resource loading, gRPC services) to be skipped. Patch settings to ensure no sections are hidden during these tests. Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> * fix: resolve DI, retry storm, race, and opaque token tracking issues - Fix get_current_user request param type (Optional[object] → Request) so FastAPI auto-injects it, enabling last_used and usage logging - Add 30s backoff in _get_sync_redis_client after connection failures to prevent retry storms when Redis is down - Move in-memory cache from hasattr-based function attributes to module-level globals, eliminating initialization race condition - Store user_email in request state for DB-fallback opaque tokens so usage middleware can log them without JWT decode Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> --------- Signed-off-by: Shoumi <shoumimukherjee@gmail.com> Signed-off-by: Mihai Criveti <crivetimihai@gmail.com> Co-authored-by: Mihai Criveti <crivetimihai@gmail.com>

shoummu1 marked this pull request as ready for review February 5, 2026 15:30

shoummu1 requested review from crivetimihai, kevalmahajan and madhav165 as code owners February 5, 2026 15:30

shoummu1 force-pushed the fix/2572-api-token-usage-stats branch 2 times, most recently from 1fdaaae to 420c276 Compare February 6, 2026 08:59

crivetimihai self-assigned this Feb 6, 2026

crivetimihai added this to the Release 1.0.0-RC1 milestone Feb 7, 2026

shoummu1 force-pushed the fix/2572-api-token-usage-stats branch from 420c276 to 967656e Compare February 10, 2026 05:44

shoummu1 marked this pull request as draft February 10, 2026 06:01

shoummu1 marked this pull request as ready for review February 10, 2026 10:37

shoummu1 force-pushed the fix/2572-api-token-usage-stats branch 3 times, most recently from 7aa61ec to d59b551 Compare February 11, 2026 10:47

crivetimihai modified the milestones: Release 1.0.0-RC1, Release 1.0.0-GA Feb 11, 2026

shoummu1 force-pushed the fix/2572-api-token-usage-stats branch from e0600b4 to b450386 Compare February 12, 2026 08:24

crivetimihai reviewed Feb 13, 2026

View reviewed changes

crivetimihai modified the milestones: Release 1.0.0-GA, Release 1.0.0-RC1 Feb 16, 2026

shoummu1 added 9 commits February 16, 2026 22:41

fix: update API token last_used and log usage stats

9d13f62

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

test fixes

c8b55d6

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

flake8 fix

a9ba3ad

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

additional changes

dbc7ef2

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

remove duplicate last_used update in token usage logging

a76b344

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

test fixes

326ba3d

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

test fixes

528cadb

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

additional test fixes

0ceb99c

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

added exception handling

651a8db

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

shoummu1 and others added 7 commits February 16, 2026 22:41

fix coverage

8a4e5c6

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

fix: optimize token tracking with rate-limiting and raw ASGI middleware

a785f82

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

test fixes

9f52600

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

test fixes

801cd99

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

additional test fixes

47e4396

Signed-off-by: Shoumi <shoumimukherjee@gmail.com>

crivetimihai force-pushed the fix/2572-api-token-usage-stats branch from b450386 to 6a8aa8c Compare February 16, 2026 23:28

crivetimihai reviewed Feb 17, 2026

View reviewed changes

crivetimihai merged commit a235827 into main Feb 17, 2026
54 checks passed

crivetimihai deleted the fix/2572-api-token-usage-stats branch February 17, 2026 00:59

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: update API token last_used and log usage stats#2711

fix: update API token last_used and log usage stats#2711
crivetimihai merged 17 commits intomainfrom
fix/2572-api-token-usage-stats

shoummu1 commented Feb 5, 2026 •

edited

Loading

Uh oh!

crivetimihai commented Feb 10, 2026

Uh oh!

crivetimihai left a comment

Uh oh!

crivetimihai commented Feb 16, 2026

Uh oh!

crivetimihai commented Feb 16, 2026

Uh oh!

crivetimihai left a comment

Uh oh!

crivetimihai Feb 17, 2026

Uh oh!

crivetimihai Feb 17, 2026

Uh oh!

crivetimihai Feb 17, 2026

Uh oh!

crivetimihai Feb 17, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

shoummu1 commented Feb 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

📌 Summary

🔗 Related Issue

🐞 Root Cause

💡 Fix Description

1. Last Used Timestamp Tracking

2. Usage Statistics Logging

3. Test Coverage

🧪 Verification

📐 MCP Compliance (if relevant)

✅ Checklist

Uh oh!

crivetimihai commented Feb 10, 2026

Uh oh!

crivetimihai left a comment

Choose a reason for hiding this comment

Uh oh!

crivetimihai commented Feb 16, 2026

Review Changes

1. fix: address review findings in token usage tracking

2. fix: patch hidden sections env leakage in admin UI tests

Uh oh!

crivetimihai commented Feb 16, 2026

Verification Against Running Instance (localhost:8080)

Setup

Results

Uh oh!

crivetimihai left a comment

Choose a reason for hiding this comment

Uh oh!

crivetimihai Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

crivetimihai Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

crivetimihai Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

crivetimihai Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

shoummu1 commented Feb 5, 2026 •

edited

Loading

1. `fix: address review findings in token usage tracking`

2. `fix: patch hidden sections env leakage in admin UI tests`