Draft: Move the Orbit usage-quota check into Rails

Sharmad Nachnolkarrequested to merge
snachnolkar/630306-orbit-quota-check-in-rails into master

What does this MR do and why?

Moves the Orbit (GitLab Knowledge Graph) usage-quota check out of the Rust orbit-server service and into GitLab Rails, ahead of the Orbit query endpoints.

Why: on self-managed, orbit-server runs on customer infrastructure and cannot hold CustomersDot admin credentials. CustomersDot authenticates self-managed callers by license checksum (X-License-Token). Rails already has that path via Gitlab::SubscriptionPortal::Client. orbit-server does not.

Key design decision

The Rails gate deliberately reproduces orbit-server's existing decision table instead of improving on it, so this MR is behaviour-neutral. Any change to what is actually enforced stays a separate, deliberate step. The gate:

  • Denies only on a CustomersDot 402.
  • Fails open on 403, on 422, on a missing namespace, on a missing user, and on transport errors.
  • Skips the check entirely on SaaS when there is no namespace to bill, matching what orbit-server does when root_namespace_id is absent.
  • Sends event_type: orbit_workflow_completion and feature_qualified_name: orbit_<source_type>, matching what CustomersDot already receives from orbit-server.

Note for reviewers: orbit-server fails open on a CustomersDot 403, whereas the tracking issue's plan said to block on it. This MR follows orbit-server so the move stays behaviour-neutral. Whether a 403 should deny is tracked as follow-up, together with 422 handling and the CustomersDot fail-close policy on self-managed.

Call sites gated

These are the only Rails paths that reach orbit-server's execute_query RPC, which is the only RPC orbit-server meters:

  • POST /orbit/query and POST /orbit/query/:name — always.
  • POST /orbit/agent/commands/:name and the MCP tools/call handler — only for the query_graph command, because the same helpers also serve schema and catalog commands that orbit-server never metered. The predicate is the one the adjacent orbit_query rate limiter already uses.

Error contract on denial

Matches exactly what Workhorse already writes when orbit-server denies, so clients see one contract either way:

  • REST: HTTP 402 with body {"code":"quota_exhausted","message":"GitLab credits exhausted","reason":"GITLAB_CREDITS_EXHAUSTED"}
  • MCP: HTTP 200 with a JSON-RPC result {"content":[{"type":"text","text":"GitLab credits exhausted"}],"isError":true}

Also included

  • An http_timeout option on Ai::UsageQuotaService, threaded to verify_usage_quota. Orbit passes 1 second. Without it the check inherits Gitlab::HTTP defaults of 10s open and 20s read with no circuit breaker. When no timeout is given, the portal call's argument list is unchanged, so Duo Chat, Duo Workflows, and DAP agent-status checks are untouched.
  • Self-managed instances without an online cloud license skip the CustomersDot call, since CustomersDot rejects them with a 401 and the request only buys an error log per query.

Feature flag

orbit_quota_check_in_rails, type beta, disabled by default.

Not in this MR

Removing the gate call in orbit-server, and any CustomersDot changes, are tracked separately. orbit-server's quota client itself is deliberately being kept, because the indexer and planned GB/volume billing have no per-request Rails path to hang a check on.

Files changed (3 commits)

  1. "Allow an explicit timeout on the CustomersDot quota call" — ee/app/services/ai/usage_quota_service.rb, ee/lib/gitlab/subscription_portal/clients/rest.rb, plus spec.
  2. "Add an Orbit quota gate that mirrors orbit-server" — new ee/lib/analytics/knowledge_graph/quota_gate.rb, new feature flag YAML, plus spec.
  3. "Enforce the Orbit quota at the query entry points" — ee/lib/api/orbit/data.rb, ee/lib/api/orbit/mcp.rb, plus request specs.

Test results

  • ee/spec/lib/analytics/knowledge_graph/quota_gate_spec.rb: 13 examples, 0 failures.
  • ee/spec/services/ai/usage_quota_service_spec.rb: 36 examples, 0 failures.
  • ee/spec/requests/api/orbit/data_spec.rb and ee/spec/requests/api/orbit/mcp_spec.rb: 297 examples, 0 failures, 26 pending (the pending ones are pre-existing skips in a shared granular-token example that does not apply to user-boundary routes).

References

Screenshots or screen recordings

No UI change. This MR is backend and API only.

How to set up and validate locally

  1. The flag is off by default. Enable it in the Rails console: Feature.enable(:orbit_quota_check_in_rails).
  2. Stub or point CustomersDot so /api/v1/consumers/resolve returns 402.
  3. Run a query against POST /api/v4/orbit/query with a valid token and confirm the 402 body: {"code":"quota_exhausted","message":"GitLab credits exhausted","reason":"GITLAB_CREDITS_EXHAUSTED"}.
  4. Confirm the same denial through MCP tools/call with invoke_command / query_graph returns the JSON-RPC isError envelope with HTTP 200.
  5. Confirm a tools/list call is not gated.
  6. Disable the flag and confirm the endpoints behave exactly as before.

MR acceptance checklist

This checklist encourages us to confirm any changes have been analyzed to reduce risks in quality, performance, reliability, security, and maintainability.

Merge request reports