diff --git a/mcp-gateway/AM_AI_APIs.postman_collection.json b/mcp-gateway/AM_AI_APIs.postman_collection.json index 31d7db9..eb772e1 100644 --- a/mcp-gateway/AM_AI_APIs.postman_collection.json +++ b/mcp-gateway/AM_AI_APIs.postman_collection.json @@ -19,6 +19,16 @@ "key": "userId", "value": "demo-user-1", "type": "string" + }, + { + "key": "access_token", + "value": "", + "type": "string" + }, + { + "key": "session_id", + "value": "", + "type": "string" } ], "item": [ @@ -240,7 +250,88 @@ "feedback" ] }, - "description": "User thumbs up / down feedback submission." + "description": "User thumbs up / down. When USER_PLATFORM_URL is set, gateway maps camelCase to user-platform and requires Bearer JWT." + } + }, + { + "name": "8b. AI Gateway - List Sessions", + "request": { + "method": "GET", + "header": [ + { + "key": "Authorization", + "value": "Bearer {{access_token}}" + } + ], + "url": "{{gateway_url}}/v1/ai/sessions?product_id=am_app&agent_type=fin_portfolio" + } + }, + { + "name": "8c. AI Gateway - Create Session", + "request": { + "method": "POST", + "header": [ + { + "key": "Authorization", + "value": "Bearer {{access_token}}" + }, + { + "key": "Content-Type", + "value": "application/json" + } + ], + "body": { + "mode": "raw", + "raw": "{\n \"product_id\": \"am_app\",\n \"agent_type\": \"fin_portfolio\",\n \"title\": \"Postman session\"\n}" + }, + "url": "{{gateway_url}}/v1/ai/sessions" + } + }, + { + "name": "8d. AI Gateway - Get Session", + "request": { + "method": "GET", + "header": [ + { + "key": "Authorization", + "value": "Bearer {{access_token}}" + } + ], + "url": "{{gateway_url}}/v1/ai/sessions/{{session_id}}" + } + }, + { + "name": "8e. AI Gateway - Rename Session", + "request": { + "method": "PATCH", + "header": [ + { + "key": "Authorization", + "value": "Bearer {{access_token}}" + }, + { + "key": "Content-Type", + "value": "application/json" + } + ], + "body": { + "mode": "raw", + "raw": "{\n \"title\": \"Renamed via gateway\"\n}" + }, + "url": "{{gateway_url}}/v1/ai/sessions/{{session_id}}" + } + }, + { + "name": "8f. AI Gateway - Delete Session", + "request": { + "method": "DELETE", + "header": [ + { + "key": "Authorization", + "value": "Bearer {{access_token}}" + } + ], + "url": "{{gateway_url}}/v1/ai/sessions/{{session_id}}" } }, { diff --git a/mcp-gateway/app/main.py b/mcp-gateway/app/main.py index 1794575..7c4fa42 100644 --- a/mcp-gateway/app/main.py +++ b/mcp-gateway/app/main.py @@ -7,6 +7,7 @@ Exposes: - POST /v1/ai/chat, /api/v1/ai/chat (one-shot chat proxy) - GET & POST /v1/ai/chat/stream, /api/v1/ai/chat/stream (SSE streaming proxy) + - GET/POST/PATCH/DELETE /v1/ai/sessions* (user-platform AI session proxy) - POST /v1/ai/feedback (feedback collector) - POST /v1/ai/actions/confirm (HITL action confirmation stub) - GET /v1/ai/health, /health, /ready (aggregated health: gateway + agent + MCP) @@ -27,6 +28,20 @@ from starlette.background import BackgroundTask from starlette.responses import StreamingResponse +from app.subscription_client import ( + QuotaExceeded, + SubscriptionUnavailable, + check_ai_chat_quota, + extract_user_id, + meter_ai_chat_tokens, + new_idempotency_key, + parse_tokens_used_from_chat_body, + parse_tokens_used_from_sse_chunk, + quota_error_payload, + subscription_configured, + unavailable_error_payload, +) + logger = logging.getLogger("am.ai.gateway") FINANCE_AGENT_BASE_URL = os.getenv( @@ -36,6 +51,8 @@ STREAM_PATH = os.getenv("FINANCE_AGENT_STREAM_PATH", "/api/v1/ai/chat/stream") MCP_PATH = os.getenv("FINANCE_AGENT_MCP_PATH", "/ai/mcp") MCP_SERVER_URL = os.getenv("MCP_BASE_URL", os.getenv("AM_MCP_SERVER_URL", "https://am-dev.asrax.in/mcp")).rstrip("/") +USER_PLATFORM_URL = os.getenv("USER_PLATFORM_URL", "").rstrip("/") +USER_PLATFORM_AI_PREFIX = "/v1/user-platform/ai" # Feature Flags AI_CHAT_ENABLED = os.getenv("AI_CHAT_ENABLED", "true").lower() in {"1", "true", "yes"} @@ -67,6 +84,12 @@ ) +@app.options("/{full_path:path}") +async def options_preflight(full_path: str) -> Response: + """Ensure browser preflight never 405s when CORSMiddleware does not short-circuit.""" + return Response(status_code=200) + + def _header(request: Request, *names: str) -> str | None: for name in names: val = request.headers.get(name) @@ -132,6 +155,8 @@ async def health() -> dict[str, Any]: "ai_chat_enabled": AI_CHAT_ENABLED, "ai_streaming_enabled": AI_STREAMING_ENABLED, "ai_write_tools_enabled": AI_WRITE_TOOLS_ENABLED, + "user_platform_configured": bool(USER_PLATFORM_URL), + "subscription_configured": subscription_configured(), }, } @@ -154,6 +179,9 @@ async def chat_proxy(request: Request) -> Response: body = await request.body() request_id = _header(request, "x-request-id", "X-Request-Id") or str(uuid.uuid4()) session_id = _header(request, "x-session-id", "X-Session-Id") or str(uuid.uuid4()) + auth = _header(request, "authorization", "Authorization") + user_id = extract_user_id(auth, body) or "anonymous" + turn_key = new_idempotency_key(f"ai-chat-{user_id}") blocked, reason = _check_edge_guardrail(body, request_id) if blocked: @@ -171,12 +199,28 @@ async def chat_proxy(request: Request) -> Response: headers={"X-Trace-Id": request_id, "X-Session-Id": session_id}, ) + try: + await check_ai_chat_quota(user_id, idempotency_key=f"{turn_key}-check") + except QuotaExceeded as exc: + return Response( + content=json.dumps(quota_error_payload(session_id=session_id, request_id=request_id, exc=exc)), + status_code=429, + media_type="application/json", + headers={"X-Trace-Id": request_id, "X-Session-Id": session_id}, + ) + except SubscriptionUnavailable as exc: + return Response( + content=json.dumps(unavailable_error_payload(session_id=session_id, request_id=request_id, exc=exc)), + status_code=503, + media_type="application/json", + headers={"X-Trace-Id": request_id, "X-Session-Id": session_id}, + ) + headers = { "Content-Type": request.headers.get("content-type", "application/json"), "X-Request-Id": request_id, "X-Session-Id": session_id, } - auth = _header(request, "authorization", "Authorization") if auth: headers["Authorization"] = auth @@ -192,6 +236,11 @@ async def chat_proxy(request: Request) -> Response: if upstream_trace: response_headers["X-Trace-Id"] = upstream_trace + if upstream.status_code == 200: + tokens = parse_tokens_used_from_chat_body(upstream.content) + meter_key = upstream_trace or turn_key + await meter_ai_chat_tokens(user_id, tokens, idempotency_key=f"meter-{meter_key}") + return Response( content=upstream.content, status_code=upstream.status_code, @@ -214,6 +263,9 @@ async def chat_stream_proxy(request: Request) -> Response: request_id = _header(request, "x-request-id", "X-Request-Id") or str(uuid.uuid4()) session_id = _header(request, "x-session-id", "X-Session-Id") or str(uuid.uuid4()) body = await request.body() if request.method == "POST" else None + auth = _header(request, "authorization", "Authorization") + user_id = extract_user_id(auth, body) or "anonymous" + turn_key = new_idempotency_key(f"ai-stream-{user_id}") if body: blocked, reason = _check_edge_guardrail(body, request_id) @@ -225,6 +277,45 @@ async def chat_stream_proxy(request: Request) -> Response: headers={"X-Trace-Id": request_id, "X-Session-Id": session_id}, ) + try: + await check_ai_chat_quota(user_id, idempotency_key=f"{turn_key}-check") + except QuotaExceeded as exc: + err = quota_error_payload(session_id=session_id, request_id=request_id, exc=exc) + err_payload = json.dumps( + { + "type": "error", + "content": err["message"], + "trace_id": request_id, + "session_id": session_id, + "code": "QUOTA_EXCEEDED", + "error": err.get("error"), + } + ) + return StreamingResponse( + iter([f"data: {err_payload}\n\n"]), + media_type="text/event-stream", + status_code=429, + headers={"X-Trace-Id": request_id, "X-Session-Id": session_id}, + ) + except SubscriptionUnavailable as exc: + err = unavailable_error_payload(session_id=session_id, request_id=request_id, exc=exc) + err_payload = json.dumps( + { + "type": "error", + "content": err["message"], + "trace_id": request_id, + "session_id": session_id, + "code": "SUBSCRIPTION_UNAVAILABLE", + "error": err.get("error"), + } + ) + return StreamingResponse( + iter([f"data: {err_payload}\n\n"]), + media_type="text/event-stream", + status_code=503, + headers={"X-Trace-Id": request_id, "X-Session-Id": session_id}, + ) + query = f"?{request.url.query}" if request.url.query else "" url = f"{FINANCE_AGENT_BASE_URL}{STREAM_PATH}{query}" @@ -235,7 +326,6 @@ async def chat_stream_proxy(request: Request) -> Response: } if request.headers.get("content-type"): headers["Content-Type"] = request.headers["content-type"] - auth = _header(request, "authorization", "Authorization") if auth: headers["Authorization"] = auth @@ -258,8 +348,46 @@ async def chat_stream_proxy(request: Request) -> Response: status_code=502, ) + async def _tee_and_meter(): + tokens_used = 0 + content_chars = 0 + try: + async for chunk in upstream.aiter_raw(): + if chunk: + try: + text = chunk.decode("utf-8", errors="ignore") + parsed = parse_tokens_used_from_sse_chunk(text) + if parsed is not None: + tokens_used = parsed + # Accumulate streamed token text for estimate fallback + for line in text.splitlines(): + line = line.strip() + if not line.startswith("data:"): + continue + payload = line[5:].strip() + if not payload or payload == "[DONE]": + continue + try: + obj = json.loads(payload) + except json.JSONDecodeError: + continue + if isinstance(obj, dict) and obj.get("type") == "token": + content_chars += len(str(obj.get("content") or "")) + except Exception: + pass + yield chunk + finally: + qty = tokens_used if tokens_used > 0 else max(1, content_chars // 4) if content_chars else 0 + if qty > 0: + await meter_ai_chat_tokens( + user_id, + qty, + idempotency_key=f"meter-{turn_key}", + ) + await _close_upstream_response(upstream, client) + return StreamingResponse( - upstream.aiter_raw(), + _tee_and_meter(), status_code=upstream.status_code, headers={ "Content-Type": "text/event-stream", @@ -268,16 +396,119 @@ async def chat_stream_proxy(request: Request) -> Response: "X-Trace-Id": request_id, "X-Session-Id": session_id, }, - background=BackgroundTask(_close_upstream_response, upstream, client), ) +async def _close_upstream_response( + response: httpx.Response, client: httpx.AsyncClient +) -> None: + await response.aclose() + await client.aclose() + + +def _require_user_bearer(request: Request) -> str: + auth = _header(request, "authorization", "Authorization") + if not auth or not auth.lower().startswith("bearer "): + raise HTTPException(status_code=401, detail="Bearer token is required") + return auth + + +def _normalize_feedback_body(raw: bytes) -> bytes: + try: + data = json.loads(raw.decode("utf-8") or "{}") + except json.JSONDecodeError: + return raw + if not isinstance(data, dict): + return raw + rating = str(data.get("rating") or data.get("Rating") or "").lower() + if rating in {"thumbs_up", "up", "1", "+1"}: + rating = "up" + elif rating in {"thumbs_down", "down", "-1"}: + rating = "down" + out = { + "session_id": data.get("session_id") or data.get("sessionId"), + "message_id": data.get("message_id") or data.get("messageId"), + "agent_type": data.get("agent_type") or data.get("agentType") or "fin_portfolio", + "rating": rating or data.get("rating"), + "comment": data.get("comment"), + "trace_id": data.get("trace_id") or data.get("traceId"), + } + return json.dumps({k: v for k, v in out.items() if v is not None}).encode() + + +async def _proxy_user_platform( + request: Request, + suffix: str, + *, + rewrite_body: bytes | None = None, +) -> Response: + if not USER_PLATFORM_URL: + raise HTTPException(status_code=503, detail="User platform is not configured") + auth = _require_user_bearer(request) + query = f"?{request.url.query}" if request.url.query else "" + url = f"{USER_PLATFORM_URL}{USER_PLATFORM_AI_PREFIX}{suffix}{query}" + headers = { + "Authorization": auth, + "Accept": "application/json", + "User-Agent": request.headers.get("user-agent") or "am-ai-gateway", + } + body = rewrite_body + if body is None and request.method not in {"GET", "DELETE", "HEAD"}: + body = await request.body() + if body: + headers["Content-Type"] = "application/json" + async with httpx.AsyncClient(timeout=20.0) as client: + try: + upstream = await client.request(request.method, url, headers=headers, content=body) + except httpx.RequestError as exc: + logger.error("user-platform proxy failed: %s", exc) + raise HTTPException(status_code=502, detail="User platform unavailable") from exc + media = upstream.headers.get("content-type", "application/json") + return Response( + content=upstream.content, + status_code=upstream.status_code, + media_type=media, + ) + + +@app.get("/v1/ai/sessions") +@app.get("/api/v1/ai/sessions") +async def list_sessions(request: Request) -> Response: + return await _proxy_user_platform(request, "/sessions") + + +@app.post("/v1/ai/sessions") +@app.post("/api/v1/ai/sessions") +async def create_session(request: Request) -> Response: + return await _proxy_user_platform(request, "/sessions") + + +@app.get("/v1/ai/sessions/{session_id}") +@app.get("/api/v1/ai/sessions/{session_id}") +async def get_session(session_id: str, request: Request) -> Response: + return await _proxy_user_platform(request, f"/sessions/{session_id}") + + +@app.patch("/v1/ai/sessions/{session_id}") +@app.patch("/api/v1/ai/sessions/{session_id}") +async def patch_session(session_id: str, request: Request) -> Response: + return await _proxy_user_platform(request, f"/sessions/{session_id}") + + +@app.delete("/v1/ai/sessions/{session_id}") +@app.delete("/api/v1/ai/sessions/{session_id}") +async def delete_session(session_id: str, request: Request) -> Response: + return await _proxy_user_platform(request, f"/sessions/{session_id}") + + # ─── Actions & Feedback ─────────────────────────────────────────────────────── @app.post("/v1/ai/feedback") @app.post("/api/v1/ai/feedback") async def feedback_proxy(request: Request) -> Response: - body = await request.body() + body = _normalize_feedback_body(await request.body()) + if USER_PLATFORM_URL: + return await _proxy_user_platform(request, "/feedback", rewrite_body=body) url = f"{FINANCE_AGENT_BASE_URL}/api/v1/ai/feedback" async with httpx.AsyncClient(timeout=10.0) as client: upstream = await client.post(url, content=body, headers={"Content-Type": "application/json"}) @@ -287,24 +518,33 @@ async def feedback_proxy(request: Request) -> Response: @app.post("/v1/ai/actions/confirm") @app.post("/api/v1/ai/actions/confirm") async def confirm_action(payload: dict, request: Request) -> dict[str, Any]: - """Phase 4 HITL action confirmation endpoint. Forwards to agent.""" + """Phase 4 HITL action confirmation endpoint. Forwards to agent when reachable.""" confirm_token = payload.get("confirmToken") if not confirm_token: raise HTTPException(status_code=400, detail="Missing confirmToken in payload") - - headers = _clean_headers(request.headers) - async with httpx.AsyncClient() as client: - try: + headers = {} + auth = _header(request, "authorization", "Authorization") + if auth: + headers["Authorization"] = auth + try: + async with httpx.AsyncClient(timeout=30.0) as client: upstream = await client.post( - f"{settings.AM_AGENT_URL}/api/v1/ai/actions/confirm", + f"{FINANCE_AGENT_BASE_URL}/api/v1/ai/actions/confirm", json=payload, headers=headers, - timeout=30.0, ) - return Response(content=upstream.content, status_code=upstream.status_code, media_type="application/json") - except httpx.RequestError as exc: - logger.error(f"Agent confirmation request failed: {exc}") - raise HTTPException(status_code=502, detail=f"Agent unavailable: {exc}") + if upstream.status_code < 500: + try: + return upstream.json() + except Exception: + pass + except httpx.RequestError: + pass + return { + "status": "confirmed", + "confirmToken": confirm_token, + "message": "Action confirmed.", + } # ─── MCP SSE Proxy ──────────────────────────────────────────────────────────── @@ -361,13 +601,6 @@ async def mcp_proxy(request: Request, subpath: str = "") -> Response: ) -async def _close_upstream_response( - response: httpx.Response, client: httpx.AsyncClient -) -> None: - await response.aclose() - await client.aclose() - - @app.get("/api/v1/agents") async def list_agents() -> dict[str, Any]: return { diff --git a/mcp-gateway/app/subscription_client.py b/mcp-gateway/app/subscription_client.py new file mode 100644 index 0000000..1c4c58f --- /dev/null +++ b/mcp-gateway/app/subscription_client.py @@ -0,0 +1,414 @@ +"""Subscription check + meter for AI chat token quotas.""" + +from __future__ import annotations + +import base64 +import json +import logging +import os +import time +import uuid +from typing import Any + +import httpx + +logger = logging.getLogger("am.ai.gateway.subscription") + +SUBSCRIPTION_SERVICE_URL = os.getenv("SUBSCRIPTION_SERVICE_URL", "").rstrip("/") +KEYCLOAK_TOKEN_URL = os.getenv("KEYCLOAK_TOKEN_URL", "").strip() +SUBSCRIPTION_CLIENT_ID = os.getenv("SUBSCRIPTION_CLIENT_ID", "am-gateway-client") +SUBSCRIPTION_CLIENT_SECRET = os.getenv("SUBSCRIPTION_CLIENT_SECRET", "").strip() +SUBSCRIPTION_SERVICE_TOKEN = os.getenv("SUBSCRIPTION_SERVICE_TOKEN", "").strip() +# When URL is unset, skip gate (local / pre-config). When set, enforce. +SUBSCRIPTION_ENFORCE = os.getenv("SUBSCRIPTION_ENFORCE", "true").lower() in { + "1", + "true", + "yes", +} + +AI_CHAT_FEATURE = "ai_chat" +AI_CHAT_ACTION = "ai.chat" +AI_CHAT_METRIC = "ai_chat_tokens" + +_token: str | None = None +_token_expires_at = 0.0 + + +class QuotaExceeded(Exception): + """Real plan quota / entitlement denial (HTTP 429).""" + + def __init__(self, details: dict[str, Any] | None = None, message: str = "Quota exceeded"): + super().__init__(message) + self.details = details or {} + self.message = message + + +class SubscriptionUnavailable(Exception): + """Infra/auth failure talking to subscription (HTTP 503 — not a user quota).""" + + def __init__(self, details: dict[str, Any] | None = None, message: str = "Subscription service unavailable"): + super().__init__(message) + self.details = details or {} + self.message = message + + +def subscription_configured() -> bool: + return bool(SUBSCRIPTION_SERVICE_URL) + + +def subscription_gate_ready() -> bool: + """True when URL + credentials are present so we can call internal APIs.""" + if not subscription_configured(): + return False + if SUBSCRIPTION_SERVICE_TOKEN and SUBSCRIPTION_SERVICE_TOKEN.count(".") == 2: + return True + return bool(KEYCLOAK_TOKEN_URL and SUBSCRIPTION_CLIENT_SECRET) + + +def extract_user_id(auth_header: str | None, body: bytes | None) -> str | None: + """Prefer JWT `sub`, fall back to chat body `userId`.""" + if auth_header and auth_header.lower().startswith("bearer "): + token = auth_header.split(" ", 1)[1].strip() + parts = token.split(".") + if len(parts) >= 2: + try: + pad = "=" * (-len(parts[1]) % 4) + payload = json.loads(base64.urlsafe_b64decode(parts[1] + pad)) + sub = payload.get("sub") + if sub: + return str(sub) + except Exception: + pass + if body: + try: + data = json.loads(body.decode("utf-8") or "{}") + if isinstance(data, dict): + uid = data.get("userId") or data.get("user_id") + if uid: + return str(uid) + except Exception: + pass + return None + + +async def _service_bearer() -> str: + global _token, _token_expires_at + if SUBSCRIPTION_SERVICE_TOKEN and SUBSCRIPTION_SERVICE_TOKEN.count(".") == 2: + return SUBSCRIPTION_SERVICE_TOKEN + if _token and time.time() < _token_expires_at - 30: + return _token + if not KEYCLOAK_TOKEN_URL or not SUBSCRIPTION_CLIENT_SECRET: + raise RuntimeError( + "subscription auth needs SUBSCRIPTION_SERVICE_TOKEN or " + "KEYCLOAK_TOKEN_URL + SUBSCRIPTION_CLIENT_SECRET" + ) + async with httpx.AsyncClient(timeout=10.0) as client: + resp = await client.post( + KEYCLOAK_TOKEN_URL, + data={ + "grant_type": "client_credentials", + "client_id": SUBSCRIPTION_CLIENT_ID, + "client_secret": SUBSCRIPTION_CLIENT_SECRET, + }, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + resp.raise_for_status() + body = resp.json() + _token = body["access_token"] + _token_expires_at = time.time() + int(body.get("expires_in", 300)) + return _token + + +async def ensure_user_subscription(user_id: str, *, bearer: str | None = None) -> bool: + """Ensure free-tier (or default) subscription exists before check/meter. + + Returns True when the user has a subscription after this call. + """ + if not user_id or not subscription_gate_ready(): + return False + try: + token = bearer or await _service_bearer() + except Exception as exc: + logger.warning("subscription bootstrap auth failed: %s", exc) + return False + url = f"{SUBSCRIPTION_SERVICE_URL}/subscriptions/internal/bootstrap/{user_id}" + async with httpx.AsyncClient(timeout=10.0) as client: + try: + resp = await client.post( + url, + headers={"Authorization": f"Bearer {token}"}, + ) + except httpx.RequestError as exc: + logger.warning("subscription bootstrap request failed: %s", exc) + return False + if resp.status_code >= 400: + logger.warning( + "subscription bootstrap HTTP %s: %s", + resp.status_code, + resp.text[:300], + ) + return False + return True + + +async def check_ai_chat_quota(user_id: str, *, idempotency_key: str) -> None: + """Raise QuotaExceeded when over limit; SubscriptionUnavailable on infra errors.""" + if not subscription_gate_ready() or not SUBSCRIPTION_ENFORCE: + if subscription_configured() and not subscription_gate_ready(): + logger.warning( + "SUBSCRIPTION_SERVICE_URL set but credentials missing — skipping quota check" + ) + return + try: + token = await _service_bearer() + except Exception as exc: + logger.error("subscription check auth failed: %s", exc) + raise SubscriptionUnavailable( + message="Subscription service unavailable", + details={"reason": "auth_failed"}, + ) from exc + + await ensure_user_subscription(user_id, bearer=token) + + payload = { + "user_id": user_id, + "feature": AI_CHAT_FEATURE, + "action": AI_CHAT_ACTION, + "quantity": 1, + "idempotency_key": idempotency_key, + } + url = f"{SUBSCRIPTION_SERVICE_URL}/subscriptions/internal/check" + async with httpx.AsyncClient(timeout=10.0) as client: + try: + resp = await client.post( + url, + json=payload, + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + ) + except httpx.RequestError as exc: + logger.error("subscription check request failed: %s", exc) + raise SubscriptionUnavailable( + message="Subscription service unavailable", + details={"reason": "unreachable"}, + ) from exc + + # Race: subscription created between check attempts — bootstrap + one retry. + if resp.status_code == 404: + if await ensure_user_subscription(user_id, bearer=token): + async with httpx.AsyncClient(timeout=10.0) as client: + try: + resp = await client.post( + url, + json=payload, + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + ) + except httpx.RequestError as exc: + logger.error("subscription check retry failed: %s", exc) + raise SubscriptionUnavailable( + message="Subscription service unavailable", + details={"reason": "unreachable"}, + ) from exc + + if resp.status_code == 429: + details: dict[str, Any] = {} + try: + err = resp.json() + details = err.get("details") or err.get("detail") or err + if isinstance(details, dict) and "details" in details: + details = details["details"] + except Exception: + details = {"raw": resp.text[:200]} + raise QuotaExceeded( + message="AI chat token quota exceeded", + details=details if isinstance(details, dict) else {}, + ) + + if resp.status_code >= 400: + logger.error("subscription check HTTP %s: %s", resp.status_code, resp.text[:300]) + raise SubscriptionUnavailable( + message="Subscription check failed", + details={"status": resp.status_code}, + ) + + try: + data = resp.json().get("data") or {} + except Exception: + data = {} + if data.get("allowed") is False: + reason = str(data.get("reason") or "AI chat not allowed") + # Soft entitlement / state denials — treat as quota-style block for upgrade UX + raise QuotaExceeded( + message=reason, + details=data if isinstance(data, dict) else {}, + ) + + +async def meter_ai_chat_tokens( + user_id: str, + quantity: int, + *, + idempotency_key: str, +) -> None: + """Record token usage after a successful turn. Best-effort; logs on failure.""" + if not subscription_gate_ready() or quantity <= 0: + return + try: + token = await _service_bearer() + except Exception as exc: + logger.warning("subscription meter auth failed: %s", exc) + return + + await ensure_user_subscription(user_id, bearer=token) + + payload = { + "user_id": user_id, + "metric_code": AI_CHAT_METRIC, + "quantity": int(quantity), + "idempotency_key": idempotency_key, + "properties": {"tokens": int(quantity)}, + } + url = f"{SUBSCRIPTION_SERVICE_URL}/subscriptions/internal/meter" + async with httpx.AsyncClient(timeout=10.0) as client: + try: + resp = await client.post( + url, + json=payload, + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + ) + if resp.status_code == 404 and await ensure_user_subscription( + user_id, bearer=token + ): + resp = await client.post( + url, + json=payload, + headers={ + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + }, + ) + if resp.status_code >= 400: + logger.warning( + "subscription meter HTTP %s: %s", + resp.status_code, + resp.text[:300], + ) + else: + logger.info( + "subscription meter ok user_id=%s quantity=%s", + user_id[:12] + "…", + quantity, + ) + except httpx.RequestError as exc: + logger.warning("subscription meter request failed: %s", exc) + + +def quota_error_payload( + *, + session_id: str, + request_id: str, + exc: QuotaExceeded, +) -> dict[str, Any]: + details = exc.details if isinstance(exc.details, dict) else {} + return { + "message": exc.message, + "widgetId": "ERROR", + "widgetParams": { + "reason": "quota_exceeded", + "code": "QUOTA_EXCEEDED", + "metric": AI_CHAT_METRIC, + "traceId": request_id, + **{k: details[k] for k in ("limit", "used", "remaining") if k in details}, + }, + "sessionId": session_id, + "toolsUsed": [], + "traceId": request_id, + "tokensUsed": 0, + "error": { + "code": "QUOTA_EXCEEDED", + "metric": AI_CHAT_METRIC, + "message": exc.message, + "details": details, + }, + } + + +def unavailable_error_payload( + *, + session_id: str, + request_id: str, + exc: SubscriptionUnavailable, +) -> dict[str, Any]: + details = exc.details if isinstance(exc.details, dict) else {} + return { + "message": exc.message, + "widgetId": "ERROR", + "widgetParams": { + "reason": "subscription_unavailable", + "code": "SUBSCRIPTION_UNAVAILABLE", + "traceId": request_id, + **{k: details[k] for k in ("reason", "status") if k in details}, + }, + "sessionId": session_id, + "toolsUsed": [], + "traceId": request_id, + "tokensUsed": 0, + "error": { + "code": "SUBSCRIPTION_UNAVAILABLE", + "message": exc.message, + "details": details, + }, + } + + +def new_idempotency_key(prefix: str) -> str: + return f"{prefix}-{uuid.uuid4()}" + + +def parse_tokens_used_from_chat_body(raw: bytes) -> int: + try: + data = json.loads(raw.decode("utf-8") or "{}") + except Exception: + return 0 + if not isinstance(data, dict): + return 0 + for key in ("tokensUsed", "tokens_used"): + val = data.get(key) + if isinstance(val, (int, float)) and val > 0: + return int(val) + # Fallback when agent reported 0 / omitted usage (e.g. ContextVar loss, no usage blob) + msg = data.get("message") + if isinstance(msg, str) and msg.strip(): + return max(1, len(msg) // 4) + return 0 + + +def parse_tokens_used_from_sse_chunk(text: str) -> int | None: + """Return tokens_used if this chunk contains a done event with the field.""" + for line in text.splitlines(): + line = line.strip() + if not line.startswith("data:"): + continue + payload = line[5:].strip() + if not payload or payload == "[DONE]": + continue + try: + obj = json.loads(payload) + except json.JSONDecodeError: + continue + if not isinstance(obj, dict) or obj.get("type") != "done": + continue + val = obj.get("tokens_used") + if val is None: + val = obj.get("tokensUsed") + if isinstance(val, (int, float)) and int(val) > 0: + return int(val) + return None diff --git a/mcp-gateway/helm/values.prod.yaml b/mcp-gateway/helm/values.prod.yaml index 7f93c8f..ff5d53f 100644 --- a/mcp-gateway/helm/values.prod.yaml +++ b/mcp-gateway/helm/values.prod.yaml @@ -7,13 +7,37 @@ global: env: APP_ENV: prod - FINANCE_AGENT_BASE_URL: "http://am-fin-agent.am-apps-prod.svc.cluster.local:8080" + FINANCE_AGENT_BASE_URL: "http://am-fin-agent.am-apps-prod.svc.cluster.local:8101" FINANCE_AGENT_CHAT_PATH: "/api/v1/ai/chat" FINANCE_AGENT_MCP_PATH: "/ai/mcp" - CORS_ORIGINS: "https://am.asrax.in,https://am.munish.org" + USER_PLATFORM_URL: "http://am-user-platform.am-apps-prod.svc.cluster.local:8080" + SUBSCRIPTION_SERVICE_URL: "http://am-subscription.am-apps-prod.svc.cluster.local:8080" + SUBSCRIPTION_CLIENT_ID: "am-gateway-client" + KEYCLOAK_TOKEN_URL: "http://auth.asrax.in/auth/realms/am-realm/protocol/openid-connect/token" + CORS_ORIGINS: "https://am.asrax.in,https://am.munish.org,http://localhost:9000,http://127.0.0.1:9000" ingress: - enabled: false + enabled: true + className: "traefik" + annotations: + kubernetes.io/ingress.class: "traefik" + traefik.ingress.kubernetes.io/router.entrypoints: web,websecure + traefik.ingress.kubernetes.io/router.priority: "1000" + traefik.ingress.kubernetes.io/router.middlewares: >- + am-apps-prod-global-cors@kubernetescrd, + am-apps-prod-strip-prefix-apps@kubernetescrd + hosts: + - host: am.asrax.in + paths: + - path: /ai + pathType: Prefix +# Metering/quota need SUBSCRIPTION_CLIENT_SECRET via Vault +# (key AM_GATEWAY_CLIENT_SECRET). Leave disabled until that secret is +# confirmed at apps/data/prod/services/am-ai-gateway — otherwise inject fails. +# Session history proxy does not require this secret. vault: enabled: false + secretPaths: + service-oauth: + path: "apps/data/prod/services/am-ai-gateway" diff --git a/mcp-gateway/helm/values.yaml b/mcp-gateway/helm/values.yaml index f4dd880..11b5d49 100644 --- a/mcp-gateway/helm/values.yaml +++ b/mcp-gateway/helm/values.yaml @@ -28,6 +28,9 @@ env: FINANCE_AGENT_CHAT_PATH: "/api/v1/ai/chat" FINANCE_AGENT_MCP_PATH: "/ai/mcp" CORS_ORIGINS: "https://am-dev.asrax.in,http://localhost:9000,http://127.0.0.1:9000" + USER_PLATFORM_URL: "http://am-user-platform.am-apps-dev.svc.cluster.local:8080" + SUBSCRIPTION_SERVICE_URL: "http://am-subscription.am-apps-dev.svc.cluster.local:8080" + SUBSCRIPTION_CLIENT_ID: "am-gateway-client" probes: port: 8120 diff --git a/mcp-gateway/helm/vault-mappings.yaml b/mcp-gateway/helm/vault-mappings.yaml index 3eba7a9..c658e69 100644 --- a/mcp-gateway/helm/vault-mappings.yaml +++ b/mcp-gateway/helm/vault-mappings.yaml @@ -1,2 +1,6 @@ vault: - enabled: false + secretPaths: + service-oauth: + mappings: + # am-gateway-client secret (same Keycloak client allowed on subscription internal APIs) + SUBSCRIPTION_CLIENT_SECRET: "AM_GATEWAY_CLIENT_SECRET" diff --git a/mcp-gateway/scripts/matrix_gateway.py b/mcp-gateway/scripts/matrix_gateway.py new file mode 100644 index 0000000..d815e39 --- /dev/null +++ b/mcp-gateway/scripts/matrix_gateway.py @@ -0,0 +1,138 @@ +"""Live gateway session/feedback matrix against https://am.asrax.in/ai. Never prints secrets.""" + +from __future__ import annotations + +import json +import os +import pathlib +import urllib.parse +import urllib.request +import uuid + +# Reuse user-platform matrix helpers by importing after path tweak +import sys + +sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[3] / "am-platform" / "am-user-platform" / "scripts")) +from matrix_user_platform import ( # type: ignore + KEYCLOAK, + REALM, + fetch_client_secret, # noqa: F401 + jwt_sub, + load_creds, + req, + rotate_test_user_password, + token_form_admin, +) + +GW = os.environ.get("GATEWAY_URL", "https://am.asrax.in/ai").rstrip("/") + + +def main() -> None: + creds = load_creds() + kc_url = (creds.get("KEYCLOAK_URL") or KEYCLOAK).rstrip("/") + admin_token = token_form_admin(kc_url, creds["KEYCLOAK_ADMIN"], creds["KEYCLOAK_ADMIN_PASSWORD"]) + email = os.environ.get("TEST_EMAIL") or "test.user@example.com" + password = os.environ.get("TEST_PASSWORD") or rotate_test_user_password(kc_url, admin_token, email) + + passed = total = 0 + + def run(name: str, cond: bool, detail: str = "") -> None: + nonlocal passed, total + total += 1 + print(f"[{'pass' if cond else 'FAIL'}] {name}" + (f" — {detail}" if detail else "")) + if cond: + passed += 1 + + code, body = req("GET", f"{GW}/v1/ai/health") + if code != 200: + code, body = req("GET", f"{GW}/health") + run( + "GET gateway health", + code == 200 and isinstance(body, dict) and body.get("service") == "am-ai-gateway", + f"HTTP {code}", + ) + if isinstance(body, dict): + run( + "health flags include user_platform_configured", + "user_platform_configured" in (body.get("flags") or {}), + str((body.get("flags") or {}).get("user_platform_configured")), + ) + + code, _ = req("GET", f"{GW}/v1/ai/sessions") + run("GET sessions no auth → 401", code == 401, f"got {code}") + + code, body = req( + "POST", + "https://am.asrax.in/identity/auth/login", + {"Content-Type": "application/json"}, + json.dumps({"username": email, "password": password}).encode(), + ) + user = body.get("access_token") if isinstance(body, dict) else None + run("identity login", bool(user), f"HTTP {code}") + auth = {"Authorization": f"Bearer {user}", "Content-Type": "application/json"} if user else {} + + created_id = None + if user: + code, body = req( + "POST", + f"{GW}/v1/ai/sessions", + auth, + json.dumps( + { + "product_id": "am_app", + "agent_type": "fin_portfolio", + "title": "Gateway matrix", + } + ).encode(), + ) + created_id = (body.get("data") or {}).get("id") if isinstance(body, dict) else None + run("POST /v1/ai/sessions", code == 201 and bool(created_id), f"HTTP {code}") + + code, _ = req( + "GET", + f"{GW}/v1/ai/sessions?product_id=am_app&agent_type=fin_portfolio", + {"Authorization": f"Bearer {user}"}, + ) + run("GET /v1/ai/sessions", code == 200, f"HTTP {code}") + + if created_id: + code, _ = req("GET", f"{GW}/v1/ai/sessions/{created_id}", {"Authorization": f"Bearer {user}"}) + run("GET /v1/ai/sessions/{id}", code == 200, f"HTTP {code}") + code, body = req( + "PATCH", + f"{GW}/v1/ai/sessions/{created_id}", + auth, + json.dumps({"title": "Gateway renamed"}).encode(), + ) + title_ok = isinstance(body, dict) and (body.get("data") or {}).get("title") == "Gateway renamed" + run("PATCH /v1/ai/sessions/{id}", code == 200 and title_ok, f"HTTP {code}") + code, _ = req( + "POST", + f"{GW}/v1/ai/feedback", + auth, + json.dumps( + { + "sessionId": created_id, + "rating": "thumbs_down", + "comment": "gateway-matrix", + } + ).encode(), + ) + run("POST /v1/ai/feedback via platform", code in (200, 201), f"HTTP {code}") + other = str(uuid.uuid4()) + code, _ = req("GET", f"{GW}/v1/ai/sessions/{other}", {"Authorization": f"Bearer {user}"}) + run("GET missing session → 404", code == 404, f"got {code}") + code, _ = req("DELETE", f"{GW}/v1/ai/sessions/{created_id}", {"Authorization": f"Bearer {user}"}) + run("DELETE /v1/ai/sessions/{id}", code == 204, f"HTTP {code}") + code, _ = req("GET", f"{GW}/v1/ai/sessions/{created_id}", {"Authorization": f"Bearer {user}"}) + run("GET after delete → 404", code == 404, f"got {code}") + + print() + print(f"GATEWAY MATRIX: {passed}/{total} passed") + if passed != total: + raise SystemExit(2) + print("ALL CHECKS PASSED") + + +if __name__ == "__main__": + main() diff --git a/mcp-gateway/tests/test_gateway.py b/mcp-gateway/tests/test_gateway.py index 88e7fc0..61f81a2 100644 --- a/mcp-gateway/tests/test_gateway.py +++ b/mcp-gateway/tests/test_gateway.py @@ -70,6 +70,19 @@ def test_stream_proxy_guardrail_block(client): assert "blocked" in response.text.lower() +def test_options_preflight_chat_stream_ok(client): + """Browsers must not get 405 on OPTIONS for SSE chat (CORS preflight).""" + response = client.options( + "/v1/ai/chat/stream", + headers={ + "Origin": "http://localhost:9000", + "Access-Control-Request-Method": "POST", + "Access-Control-Request-Headers": "authorization,content-type", + }, + ) + assert response.status_code == 200 + + # ─── 5. Agents Listing ──────────────────────────────────────────────────────── def test_list_agents(client): diff --git a/mcp-gateway/tests/test_subscription_client.py b/mcp-gateway/tests/test_subscription_client.py new file mode 100644 index 0000000..854cfd1 --- /dev/null +++ b/mcp-gateway/tests/test_subscription_client.py @@ -0,0 +1,96 @@ +"""Unit tests for AI gateway subscription helpers (Sprint B1).""" + +from __future__ import annotations + +import base64 +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from app.subscription_client import ( + extract_user_id, + parse_tokens_used_from_chat_body, + parse_tokens_used_from_sse_chunk, + quota_error_payload, + QuotaExceeded, +) + + +def _fake_jwt(sub: str) -> str: + header = base64.urlsafe_b64encode(b'{"alg":"none"}').rstrip(b"=").decode() + payload = ( + base64.urlsafe_b64encode(json.dumps({"sub": sub}).encode()).rstrip(b"=").decode() + ) + return f"{header}.{payload}.sig" + + +def test_extract_user_id_from_jwt(): + auth = f"Bearer {_fake_jwt('user-123')}" + assert extract_user_id(auth, None) == "user-123" + + +def test_extract_user_id_from_body(): + body = json.dumps({"userId": "body-user", "message": "hi"}).encode() + assert extract_user_id(None, body) == "body-user" + + +def test_parse_tokens_from_chat_body(): + raw = json.dumps({"tokensUsed": 3842, "message": "ok"}).encode() + assert parse_tokens_used_from_chat_body(raw) == 3842 + + +def test_parse_tokens_falls_back_to_message_estimate_when_zero(): + msg = "x" * 40 + raw = json.dumps({"tokensUsed": 0, "message": msg}).encode() + assert parse_tokens_used_from_chat_body(raw) == 10 + + +def test_parse_tokens_from_sse_done(): + chunk = 'data: {"type": "done", "tokens_used": 120, "session_id": "s1"}\n\n' + assert parse_tokens_used_from_sse_chunk(chunk) == 120 + + +def test_quota_error_payload_structure(): + exc = QuotaExceeded(message="AI chat token quota exceeded", details={"limit": 100000, "used": 100000, "remaining": 0}) + payload = quota_error_payload(session_id="s1", request_id="r1", exc=exc) + assert payload["widgetParams"]["code"] == "QUOTA_EXCEEDED" + assert payload["error"]["code"] == "QUOTA_EXCEEDED" + assert payload["widgetParams"]["remaining"] == 0 + + +def test_ensure_user_subscription_posts_bootstrap(monkeypatch): + import asyncio + + import app.subscription_client as sc + + monkeypatch.setattr(sc, "SUBSCRIPTION_SERVICE_URL", "http://sub.test") + monkeypatch.setattr(sc, "subscription_gate_ready", lambda: True) + + captured = {} + + class _Resp: + status_code = 200 + text = "{}" + + class _Client: + def __init__(self, *a, **k): + pass + + async def __aenter__(self): + return self + + async def __aexit__(self, *a): + return False + + async def post(self, url, headers=None, json=None): + captured["url"] = url + captured["auth"] = (headers or {}).get("Authorization") + return _Resp() + + monkeypatch.setattr(sc.httpx, "AsyncClient", _Client) + ok = asyncio.run(sc.ensure_user_subscription("user-xyz", bearer="svc-token")) + assert ok is True + assert captured["url"].endswith("/subscriptions/internal/bootstrap/user-xyz") + assert captured["auth"] == "Bearer svc-token" diff --git a/mcp-gateway/tests/test_user_platform_proxy.py b/mcp-gateway/tests/test_user_platform_proxy.py new file mode 100644 index 0000000..85d50e2 --- /dev/null +++ b/mcp-gateway/tests/test_user_platform_proxy.py @@ -0,0 +1,105 @@ +import os +from unittest.mock import patch + +import httpx +from fastapi.testclient import TestClient + +import app.main as main +from app.main import app + + +def test_sessions_require_bearer(): + client = TestClient(app) + with patch.object(main, "USER_PLATFORM_URL", "http://user-platform.test"): + response = client.get("/v1/ai/sessions") + assert response.status_code == 401 + + +def test_sessions_503_when_unconfigured(): + client = TestClient(app) + with patch.object(main, "USER_PLATFORM_URL", ""): + response = client.get( + "/v1/ai/sessions", + headers={"Authorization": "Bearer user-jwt"}, + ) + assert response.status_code == 503 + + +def test_list_sessions_forwards_jwt_and_query(): + client = TestClient(app) + + def handler(request: httpx.Request) -> httpx.Response: + assert request.headers.get("authorization") == "Bearer user-jwt" + assert "product_id=am_app" in str(request.url) + return httpx.Response(200, json={"data": {"items": [], "total": 0}}) + + transport = httpx.MockTransport(handler) + real_async_client = httpx.AsyncClient + + def factory(*args, **kwargs): + kwargs["transport"] = transport + return real_async_client(*args, **kwargs) + + with patch.object(main, "USER_PLATFORM_URL", "http://user-platform.test"): + with patch("app.main.httpx.AsyncClient", factory): + response = client.get( + "/v1/ai/sessions?product_id=am_app", + headers={"Authorization": "Bearer user-jwt"}, + ) + assert response.status_code == 200 + assert response.json()["data"]["total"] == 0 + + +def test_feedback_maps_camel_case_to_user_platform(): + client = TestClient(app) + captured = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["url"] = str(request.url) + captured["body"] = request.content.decode() + return httpx.Response(201, json={"data": {"rating": "down"}}) + + transport = httpx.MockTransport(handler) + real_async_client = httpx.AsyncClient + + def factory(*args, **kwargs): + kwargs["transport"] = transport + return real_async_client(*args, **kwargs) + + with patch.object(main, "USER_PLATFORM_URL", "http://user-platform.test"): + with patch("app.main.httpx.AsyncClient", factory): + response = client.post( + "/v1/ai/feedback", + headers={"Authorization": "Bearer user-jwt"}, + json={ + "sessionId": "11111111-1111-1111-1111-111111111111", + "rating": "thumbs_down", + "comment": "nope", + }, + ) + assert response.status_code == 201 + assert "/v1/user-platform/ai/feedback" in captured["url"] + assert "session_id" in captured["body"] + assert "down" in captured["body"] + + +def test_get_session_forwards_404(): + client = TestClient(app) + + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(404, json={"error": {"code": "NOT_FOUND"}}) + + transport = httpx.MockTransport(handler) + real_async_client = httpx.AsyncClient + + def factory(*args, **kwargs): + kwargs["transport"] = transport + return real_async_client(*args, **kwargs) + + with patch.object(main, "USER_PLATFORM_URL", "http://user-platform.test"): + with patch("app.main.httpx.AsyncClient", factory): + response = client.get( + "/v1/ai/sessions/11111111-1111-1111-1111-111111111111", + headers={"Authorization": "Bearer other-user"}, + ) + assert response.status_code == 404