MCP server exposing ZoomInfo GTM data to Claude under a credit ceiling

[project]
name = "zoominfo-gtm-mcp"
version = "0.1.0"
description = "Read-only, credit-governed MCP server over the ZoomInfo GTM API for unattended agents"
readme = "README.md"
requires-python = ">=3.11"
license = { text = "MIT" }
dependencies = [
  "mcp>=1.2.0",
  "httpx>=0.27.0",
  "pydantic>=2.6.0",
]

[project.optional-dependencies]
dev = [
  "pytest>=8.0.0",
  "pytest-httpx>=0.30.0",
  "ruff>=0.4.0",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.wheel]
packages = ["src/zoominfo_gtm_mcp"]

[tool.ruff]
line-length = 100
target-version = "py311"

# mcp-server-zoominfo-gtm-revops

A read-only MCP server over the [ZoomInfo](https://www.zoominfo.com) GTM API with a credit governor in front of it. Five tools: two free searches, two budgeted enrichments, and a credit-status tool. Every credit-spending call is checked against a daily ceiling before it goes out, served from a local cache when the record is still inside the Records Under Management window, and written to an append-only audit log naming the run that spent the money.

> **STATUS: scaffold — not runtime-tested.** The code follows the official `mcp` Python SDK conventions. Endpoint paths, OAuth scopes, request shapes, rate-limit behaviour and the credit rules track the public ZoomInfo GTM API docs (docs.zoominfo.com) as of August 2026. The company enrich/search paths and the usage path are quoted directly from those docs; the contact paths follow the documented symmetry and should be confirmed against `https://docs.zoominfo.com/llms.txt` before you rely on them. It has not been executed against a live ZoomInfo tenant.

## Read this first: ZoomInfo ships an official hosted MCP server

ZoomInfo hosts its own MCP server at `https://mcp.zoominfo.com/mcp`. It authenticates over OAuth 2.0 in the browser, is included with every subscription at no extra cost, and exposes 19 tools — 16 data tools (company and contact search and enrich, intent, scoops, news, lookup, lookalikes, recommended contacts, audiences, GTM context) plus three agentic ones (Account Research, Contact Research, Update GTM Context). Admins enable it per user in the Admin Portal.

**For a person doing interactive research, that is the correct answer and this scaffold is wasted work.** Connect it and move on:

```bash
claude mcp add --transport http zoominfo https://mcp.zoominfo.com/mcp
```

Build this instead when one of the following is true.

**Your agent runs unattended or on a schedule.** This is ZoomInfo's own guidance, not a preference of ours: the hosted MCP server is documented as unsuitable for bulk exports, CRM write-back, and scheduled jobs, and scheduled pipelines are directed to the API instead. If your agent wakes up at 06:00 and enriches an account list without a human in the loop, the hosted server is the wrong instrument by the vendor's own description.

**You need a service identity, not a user identity.** The hosted server runs as the signed-in person, with that person's entitlements, enabled per user by an admin. A shared agent triggered from Slack or a job runner has no person to be. The client-credentials flow this scaffold uses gives it its own client id and its own scope set.

**You need a hard spend ceiling.** The hosted server has no per-run credit cap. Enrichment charges one bulk data credit per returned record, so an agent looping over an account list can spend real money before anyone notices — see the arithmetic below. `ZI_DAILY_CREDIT_LIMIT` is a number the agent cannot argue with.

**Your subscription runs on recurring monthly credits.** The hosted MCP server uses bulk data credits and does not work with recurring monthly credits. If that is your contract shape, the hosted server will not function for you at all and the API is the only route.

If none of those apply, use the hosted server.

## What a runaway costs

Enrichment charges one bulk data credit per record returned, up to 25 records per request, with no charge for no-match results or errors. Bulk credits are commonly quoted in the range of $0.60–$1.00 each at small volumes, falling toward roughly $0.20 at high volume — those are third-party resale figures, not a ZoomInfo rate card, and your contract governs.

An agent researching 200 accounts and pulling four contacts each is 800 new records. That is 32 enrich calls and 800 bulk data credits — on the order of $480–$800 for one unattended afternoon, using the third-party band above.

The 32 calls are nothing against the rate limits: even the smallest documented package, Builder, allows 5 requests/second, 10,800/hour and 129,600/day (Standard: 25/s, 54,000/hr, 648,000/day; Scaling: 35/s, 75,600/hr, 907,200/day). **The binding constraint on an enrichment agent is the credit pool, not throughput.** That is why this scaffold governs credits and merely reports rate limits.

## What it exposes

All five tools are reads. There is no write tool, no audience-mutation tool, and no GTM-config tool — the scopes in `ZI_SCOPES` do not request them.

- `zi_credit_status()` — free. Combines ZoomInfo's subscription counters (`GET /data/v1/users/usage`, returning `limitType` / `totalLimit` / `currentUsage` / `usageRemaining`) with this server's local ledger: spent today, ceiling, held by calls in flight, available, and per-tool spend over 7 days. The tool description tells the agent to call this before planning a batch. If ZoomInfo's usage endpoint fails, the tool degrades to the local ledger rather than erroring — the local ceiling is still enforced.
- `zi_search_companies(criteria, page_size=25, page_number=1, sort='-revenue')` — `POST /data/v1/companies/search`. Free: charges no credits and returned companies do not count against record limits, though each request counts against rate limits. Page size is clamped to 100.
- `zi_search_contacts(criteria, page_size=25, page_number=1)` — `POST /data/v1/contacts/search`. Free, same terms.
- `zi_enrich_companies(company_ids, output_fields?)` — `POST /data/v1/companies/enrich`. Costs credits. At most 25 ids.
- `zi_enrich_contacts(contact_ids, output_fields?)` — `POST /data/v1/contacts/enrich`. Costs credits. At most 25 ids. This is the tool that returns verified business email and direct dial, so it is both the expensive one and the one carrying personal data.

Search results are deliberately slimmed to ids plus a thin label. Search is free and enrichment is not, so the job of a search result here is to let the agent pick which ids are worth paying for — returning full search payloads invites the model to treat unverified search fields as enriched data.

## How the credit governor works

The cost of an enrich call is not knowable before the response is parsed: no-match records and records already under management are not charged. So the budget is enforced in two steps, in `src/zoominfo_gtm_mcp/budget.py`.

1. **Reserve the worst case.** Before the request, hold one credit per record not already in the cache — every input matching, every match new. If that worst case exceeds the remaining ceiling, the call is refused. Refusal is all-or-nothing: a partial reservation would let an agent enrich the first 8 of 25 accounts and report success, which reads as a complete answer and is not one.
2. **Settle against reality.** After the response, count records ZoomInfo returned as a match, write that to the durable ledger, and release the difference.

The refusal comes back as a **result**, not an exception — a JSON object with `refused: true`, the remaining allowance, and a next step. A model can read that and re-plan against what is left; a protocol error usually just ends the turn.

The cache is SQLite, keyed on ZoomInfo's own record id, TTL 365 days to match the 12-month Records Under Management window during which re-enrichment is free. A hit costs no credit *and* no HTTP request, which is what matters when an agent re-asks the same question four times in one session.

Reservations are in-process; the ledger is on disk. If the process dies between reserve and settle, the reservation dies with it and the ledger never records that spend — the next run's ceiling is then slightly generous rather than slightly strict. That is the safer direction to be wrong for a guard that could otherwise deadlock an agent against a phantom hold.

## Setup

### 1. Install

```bash
cd mcp-server-zoominfo-gtm-revops
python -m venv .venv
source .venv/bin/activate     # or .venv\Scripts\activate on Windows
pip install -e .
```

### 2. Create an API application in ZoomInfo

You need a ZoomInfo subscription with API access and bulk data credits enabled. In the ZoomInfo admin portal, create an API application and configure it for the **client credentials** grant, then copy its client id and secret. Scopes are configured on the application, not requested freely at token time — a token request for a scope the application does not hold is rejected.

Grant the application exactly two scopes for this server:

- `api:data:company` — search and enrich company data
- `api:data:contact` — search and enrich contact data

Do not add `api:gtm-config:manage`, `api:audience:manage`, or `api:gtm-data-model:manage`. Nothing here uses them, and a scope an agent holds is a scope an agent can be talked into using.

### 3. Environment variables

| Variable | Required | Default | Where the value comes from |
|---|---|---|---|
| `ZI_CLIENT_ID` | yes | — | The API application you created in step 2. |
| `ZI_CLIENT_SECRET` | yes | — | Same application. Shown once at creation; store it in your secret manager, not in the MCP config file. |
| `ZI_DAILY_CREDIT_LIMIT` | no | `250` | Your call. Set it to the number of new records you are willing to buy in one day. 250 credits is roughly $150–$250 at the third-party band above. |
| `ZI_STATE_PATH` | no | `./zi_state.db` | Path to the SQLite ledger + cache. Put it somewhere durable — deleting it resets today's spend counter to zero and empties the cache. |
| `ZI_AUDIT_LOG` | no | unset | Path to an append-only JSONL audit file. Unset means no audit trail is written, which is fine for a laptop and not fine for a shared agent. |
| `ZI_SCOPES` | no | `api:data:company api:data:contact` | Space-delimited. Must be a subset of what the application holds. |
| `ZI_RUN_ID` | no | random | Stamped on every ledger and audit row. Set it to your job's run id so spend attributes to the run that caused it. |
| `ZI_BASE_URL` | no | `https://api.zoominfo.com/gtm` | Override only if ZoomInfo gives you a different host. |

### 4. Register with Claude

Claude Code:

```bash
claude mcp add zoominfo-gtm -- python -m zoominfo_gtm_mcp.server
```

Claude Desktop — add to `claude_desktop_config.json`:

```json
{
  "mcpServers": {
    "zoominfo-gtm": {
      "command": "/absolute/path/to/.venv/bin/python",
      "args": ["-m", "zoominfo_gtm_mcp.server"],
      "env": {
        "ZI_CLIENT_ID": "your-client-id",
        "ZI_CLIENT_SECRET": "your-client-secret",
        "ZI_DAILY_CREDIT_LIMIT": "250",
        "ZI_STATE_PATH": "/absolute/path/to/zi_state.db",
        "ZI_AUDIT_LOG": "/absolute/path/to/zi_audit.jsonl"
      }
    }
  }
}
```

### 5. Sanity check

Run these three in order. They prove auth, the free path, and the governor without spending much.

1. **"What is my ZoomInfo credit status?"** — calls `zi_credit_status`. Proves the client credentials work and the ledger initialised. Costs nothing. If the subscription-usage block reports an error but the local ledger renders, your token works and the usage scope or endpoint does not — the server is still safe to run.
2. **"Find software companies in California with revenue over $50M."** — calls `zi_search_companies`. Costs nothing. Returns ids.
3. **"Enrich the first two of those."** — calls `zi_enrich_companies` with 2 ids. Should charge at most 2 credits. Run it a second time with the same ids: `credits_charged` must be `0` and `cache_hits` must be `2`. That is the cache doing its job. Then set `ZI_DAILY_CREDIT_LIMIT=1`, restart, and ask for three new ids — you should get the structured `refused: true` object rather than an enrichment.

## Security model

**The token is a service identity with two read scopes.** It cannot write to ZoomInfo, cannot modify audiences or GTM config, and cannot reach intent, news, or scoops without scopes this server does not request.

**Enriched contact data enters the conversation.** `zi_enrich_contacts` returns verified business email and direct dial. Every field it returns is visible to the model and lives in the transcript. Narrow `output_fields` to what the task actually needs — the field set does not change the credit cost, but it does change what personal data is exposed and how many tokens the answer burns. If contact PII cannot reach an LLM at all, no MCP server over a contact database is the right project.

**ZoomInfo's own connection guidance requires model training to be turned off** in your AI account or client settings before connecting. That applies here too.

**The audit log is the accountability record.** Each line carries timestamp, run id, tool, records requested, cache hits, credits charged, no-match count, and the running daily total. That is what a finance chargeback conversation needs and what the hosted server does not give you locally.

**The state file is security-relevant.** `zi_state.db` holds cached enrichment payloads — real contact records — in plaintext SQLite. Put it on an encrypted volume and treat deleting it as both a cache flush and a spend-counter reset.

## Known limits — work through these before production

1. **Not runtime-tested.** Verify every path against `https://docs.zoominfo.com/llms.txt` before trusting it, starting with the two contact paths, which are inferred from the documented company symmetry rather than quoted.
2. **`credits_charged` is an upper bound, not an invoice.** It counts matched records returned. Records already under management are returned without a new charge, so the true bill can be lower — never higher. Reconcile against ZoomInfo's own usage reporting before using these numbers for chargeback.
3. **The daily ceiling rolls at UTC midnight**, not local midnight. Change `_today()` in `budget.py` if your finance day differs.
4. **No cross-process reservation.** Two servers sharing one `ZI_STATE_PATH` see each other's *settled* spend but not each other's in-flight holds, so concurrent runs can jointly overshoot the ceiling by up to 25 records each. Give each agent its own state file, or move reservations into the database, before running several at once.
5. **No retry or backoff.** A 429 raises with the rejected bucket and `Retry-After` in the message and the reservation is released, but nothing retries. Add backoff — 1–5s exponential for a per-second rejection, the full `Retry-After` for hour or day rejections, which can exceed 700 seconds.
6. **Search criteria are passed through unvalidated.** The tool schemas take a free-form `criteria` object, so a malformed filter surfaces as a ZoomInfo 4xx rather than a local error. Pin the fields your team actually uses into the schema once you know them.
7. **The cache never invalidates on change.** A 364-day-old cached record is served as current. If a use case needs freshness inside the RUM window, add a `max_age_days` argument that bypasses the cache and knowingly spends a credit.

"""zoominfo_gtm_mcp — read-only, credit-governed MCP server over the ZoomInfo GTM API."""

__version__ = "0.1.0"

"""Credit ledger, record cache, and audit log for the ZoomInfo GTM MCP server.

ZoomInfo charges one bulk data credit per *record returned* by an enrich call, unless
that record is already under management. No-match results and errors are not charged.
That asymmetry is the whole reason this module exists: the cost of a call is not known
until the response is parsed, so a budget can only be enforced in two steps.

  1. reserve(n)   — before the call, hold the worst case (every input matches, every
                    match is new). Refuses the call if the worst case would breach the
                    ceiling. Pessimistic on purpose: an agent that discovers its budget
                    mid-batch has already spent the money.
  2. settle(...)  — after the call, replace the reservation with the count actually
                    charged, derived from matchStatus, and release the difference.

Reservations live in memory for the process; the ledger is durable. If the process dies
between reserve and settle, the reservation dies with it and the durable ledger simply
never records the spend — the next run's ceiling is then slightly generous rather than
slightly strict. That is the safer direction to be wrong for a guard that can otherwise
deadlock an agent against a phantom hold.

The cache is keyed on ZoomInfo's own record id with a 365-day TTL, matching the Records
Under Management window during which re-enrichment is free. A hit costs no credit *and*
no HTTP request, which is the difference that matters when an agent re-asks the same
question four times in one session.
"""

from __future__ import annotations

import json
import os
import sqlite3
import threading
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

# Records Under Management: ZoomInfo does not re-charge for a record already under
# management. The documented window is 12 months, so a longer local TTL would serve
# stale data while a shorter one throws away free freshness.
CACHE_TTL_SECONDS = 365 * 24 * 60 * 60


class BudgetExceeded(RuntimeError):
    """Raised before any HTTP call when the worst-case cost breaches the daily ceiling."""


@dataclass(frozen=True)
class BudgetState:
    daily_limit: int
    spent_today: int
    reserved: int

    @property
    def available(self) -> int:
        return max(0, self.daily_limit - self.spent_today - self.reserved)


class CreditGovernor:
    """Durable credit ledger + record cache + append-only audit log, backed by SQLite."""

    def __init__(self, db_path: str | Path, daily_limit: int, run_id: str, audit_path: str | Path | None = None):
        self.db_path = str(db_path)
        self.daily_limit = daily_limit
        self.run_id = run_id
        self.audit_path = Path(audit_path) if audit_path else None
        self._reserved = 0
        self._lock = threading.Lock()
        self._init_db()

    # ----- schema -----

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.db_path, timeout=10.0)
        # The ledger is read on every enrich call and written on every settle. WAL keeps
        # a concurrent reader from blocking the writer when two agents share a state file.
        conn.execute("PRAGMA journal_mode=WAL")
        return conn

    def _init_db(self) -> None:
        Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
        with self._connect() as conn:
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS credit_ledger (
                    id          INTEGER PRIMARY KEY AUTOINCREMENT,
                    ts          TEXT    NOT NULL,
                    day         TEXT    NOT NULL,
                    run_id      TEXT    NOT NULL,
                    tool        TEXT    NOT NULL,
                    requested   INTEGER NOT NULL,
                    cache_hits  INTEGER NOT NULL,
                    charged     INTEGER NOT NULL,
                    no_match    INTEGER NOT NULL
                )
                """
            )
            conn.execute("CREATE INDEX IF NOT EXISTS idx_ledger_day ON credit_ledger(day)")
            conn.execute(
                """
                CREATE TABLE IF NOT EXISTS record_cache (
                    entity     TEXT NOT NULL,
                    zi_id      TEXT NOT NULL,
                    payload    TEXT NOT NULL,
                    fetched_at REAL NOT NULL,
                    PRIMARY KEY (entity, zi_id)
                )
                """
            )

    # ----- budget -----

    def state(self) -> BudgetState:
        return BudgetState(self.daily_limit, self.spent_today(), self._reserved)

    def spent_today(self) -> int:
        with self._connect() as conn:
            row = conn.execute(
                "SELECT COALESCE(SUM(charged), 0) FROM credit_ledger WHERE day = ?",
                (_today(),),
            ).fetchone()
        return int(row[0])

    def reserve(self, worst_case: int) -> int:
        """Hold `worst_case` credits. Raises BudgetExceeded rather than partially reserving.

        Partial reservation would let an agent silently enrich the first 8 of 25 accounts
        and report success, which reads as a complete answer and is not one.
        """
        with self._lock:
            st = self.state()
            if worst_case > st.available:
                raise BudgetExceeded(
                    f"This call could charge up to {worst_case} bulk data credits and only "
                    f"{st.available} remain in today's ceiling of {st.daily_limit} "
                    f"({st.spent_today} already spent, {st.reserved} held by calls in flight). "
                    "Narrow the batch, raise ZI_DAILY_CREDIT_LIMIT deliberately, or wait for "
                    "the ceiling to roll over at UTC midnight."
                )
            self._reserved += worst_case
            return worst_case

    def settle(self, *, tool: str, reserved: int, requested: int, charged: int, cache_hits: int, no_match: int) -> None:
        """Release the reservation and record what was actually charged."""
        with self._lock:
            self._reserved = max(0, self._reserved - reserved)
        with self._connect() as conn:
            conn.execute(
                "INSERT INTO credit_ledger (ts, day, run_id, tool, requested, cache_hits, charged, no_match) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
                (_now_iso(), _today(), self.run_id, tool, requested, cache_hits, charged, no_match),
            )
        self._audit(
            {
                "ts": _now_iso(),
                "run_id": self.run_id,
                "tool": tool,
                "requested": requested,
                "cache_hits": cache_hits,
                "charged": charged,
                "no_match": no_match,
                "spent_today_after": self.spent_today(),
                "daily_limit": self.daily_limit,
            }
        )

    def release(self, reserved: int) -> None:
        """Drop a reservation without charging — used when the HTTP call itself fails."""
        with self._lock:
            self._reserved = max(0, self._reserved - reserved)

    def spend_by_tool(self, days: int = 7) -> list[dict[str, Any]]:
        with self._connect() as conn:
            rows = conn.execute(
                "SELECT day, tool, SUM(charged), SUM(cache_hits) FROM credit_ledger "
                "WHERE day >= date('now', ?) GROUP BY day, tool ORDER BY day DESC, tool",
                (f"-{int(days)} days",),
            ).fetchall()
        return [
            {"day": d, "tool": t, "credits_charged": int(c), "cache_hits": int(h)}
            for d, t, c, h in rows
        ]

    # ----- cache -----

    def cache_get(self, entity: str, zi_ids: Iterable[str]) -> dict[str, Any]:
        ids = [str(i) for i in zi_ids]
        if not ids:
            return {}
        cutoff = time.time() - CACHE_TTL_SECONDS
        placeholders = ",".join("?" for _ in ids)
        with self._connect() as conn:
            rows = conn.execute(
                f"SELECT zi_id, payload FROM record_cache "
                f"WHERE entity = ? AND fetched_at >= ? AND zi_id IN ({placeholders})",
                (entity, cutoff, *ids),
            ).fetchall()
        return {zid: json.loads(payload) for zid, payload in rows}

    def cache_put(self, entity: str, records: dict[str, Any]) -> None:
        if not records:
            return
        now = time.time()
        with self._connect() as conn:
            conn.executemany(
                "INSERT OR REPLACE INTO record_cache (entity, zi_id, payload, fetched_at) VALUES (?, ?, ?, ?)",
                [(entity, str(k), json.dumps(v), now) for k, v in records.items()],
            )

    # ----- audit -----

    def _audit(self, event: dict[str, Any]) -> None:
        if not self.audit_path:
            return
        self.audit_path.parent.mkdir(parents=True, exist_ok=True)
        with self.audit_path.open("a", encoding="utf-8") as fh:
            fh.write(json.dumps(event, ensure_ascii=False) + "\n")


def _today() -> str:
    # The ceiling rolls at UTC midnight, not local midnight, so a team spread across
    # time zones sees one shared boundary rather than an argument about whose day it is.
    return datetime.now(timezone.utc).date().isoformat()


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def governor_from_env(run_id: str) -> CreditGovernor:
    return CreditGovernor(
        db_path=os.environ.get("ZI_STATE_PATH", "./zi_state.db"),
        daily_limit=int(os.environ.get("ZI_DAILY_CREDIT_LIMIT", "250")),
        run_id=run_id,
        audit_path=os.environ.get("ZI_AUDIT_LOG"),
    )

"""zoominfo-gtm-mcp — a read-only, credit-governed MCP server over the ZoomInfo GTM API.

Exposes five read tools: two free searches, two budgeted enrichments, and a credit-status
tool the agent is told to call before any batch. Every credit-spending call passes through
the governor in budget.py, which holds the worst-case cost, checks a daily ceiling, serves
records from a local cache when they are still inside the Records Under Management window,
and writes an audit row naming the run that spent the money.

This exists alongside ZoomInfo's own hosted MCP server at https://mcp.zoominfo.com/mcp,
which is included with a subscription, exposes 19 tools, and is the right answer for a
person doing interactive research. ZoomInfo's own guidance is that its hosted server is
not for scheduled jobs or bulk pipelines — those run through the API. This scaffold is for
exactly that case: an unattended agent, running on a service identity, against a hard
credit ceiling that a per-user OAuth grant cannot express.

STATUS: scaffold — not runtime-tested. Endpoint paths, scopes, request shapes, and the
credit rules track the public ZoomInfo GTM API docs (docs.zoominfo.com) as of August 2026.
The company paths and the usage path are quoted directly from those docs; the contact
paths follow the documented symmetry and should be confirmed against
https://docs.zoominfo.com/llms.txt before production use.

Run as: python -m zoominfo_gtm_mcp.server
"""

from __future__ import annotations

import json
import os
import time
import uuid
from typing import Any

import httpx
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import TextContent, Tool

from .budget import BudgetExceeded, governor_from_env

# ----- Configuration (read from env at startup) -----

ZI_CLIENT_ID = os.environ.get("ZI_CLIENT_ID")
ZI_CLIENT_SECRET = os.environ.get("ZI_CLIENT_SECRET")
ZI_BASE_URL = os.environ.get("ZI_BASE_URL", "https://api.zoominfo.com/gtm").rstrip("/")

# Client credentials, not the authorization-code flow the hosted server uses. A service
# identity is the point: an unattended agent should carry its own scope set, not borrow
# whichever human happened to authorize it last.
TOKEN_PATH = "/oauth/v1/token"
ZI_SCOPES = os.environ.get("ZI_SCOPES", "api:data:company api:data:contact")

# ZoomInfo caps enrich at 25 records per request and search at 100 per page. Both are the
# vendor's numbers, not ours; sending more is a 4xx, not a slow success.
MAX_ENRICH_BATCH = 25
MAX_SEARCH_PAGE = 100
DEFAULT_SEARCH_PAGE = 25

# Tokens come back with expires_in around 1000 seconds. Refreshing at 80% of the stated
# lifetime avoids the case where a token passes the local check and expires in flight.
TOKEN_REFRESH_MARGIN = 0.8

RUN_ID = os.environ.get("ZI_RUN_ID") or f"run-{uuid.uuid4().hex[:12]}"

_governor = None
_token: dict[str, Any] = {"access_token": None, "expires_at": 0.0}


def governor():
    global _governor
    if _governor is None:
        _governor = governor_from_env(RUN_ID)
    return _governor


def require_config() -> None:
    missing = [n for n, v in (("ZI_CLIENT_ID", ZI_CLIENT_ID), ("ZI_CLIENT_SECRET", ZI_CLIENT_SECRET)) if not v]
    if missing:
        raise RuntimeError(f"{' and '.join(missing)} env var(s) required")


# ----- Auth -----


async def access_token(client: httpx.AsyncClient) -> str:
    if _token["access_token"] and time.time() < _token["expires_at"]:
        return _token["access_token"]
    r = await client.post(
        f"{ZI_BASE_URL}{TOKEN_PATH}",
        auth=(ZI_CLIENT_ID, ZI_CLIENT_SECRET),
        data={"grant_type": "client_credentials", "scope": ZI_SCOPES},
        headers={"Content-Type": "application/x-www-form-urlencoded", "Accept": "application/json"},
    )
    if r.status_code in (400, 401):
        raise PermissionError(
            "ZoomInfo rejected the client credentials. Check ZI_CLIENT_ID/ZI_CLIENT_SECRET, and "
            "confirm the requested scopes are a subset of what the application is configured for "
            f"(requested: {ZI_SCOPES}). Scopes are set on the application, not on the token request."
        )
    r.raise_for_status()
    payload = r.json()
    _token["access_token"] = payload["access_token"]
    _token["expires_at"] = time.time() + float(payload.get("expires_in", 1000)) * TOKEN_REFRESH_MARGIN
    return _token["access_token"]


# ----- HTTP -----


async def zi_request(method: str, path: str, *, json_body: dict[str, Any] | None = None) -> dict[str, Any]:
    async with httpx.AsyncClient(timeout=45.0) as client:
        token = await access_token(client)
        r = await client.request(
            method,
            f"{ZI_BASE_URL}{path}",
            headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
            json=json_body,
        )
        _raise_for_zi(r)
        return r.json() if r.content else {}


def _raise_for_zi(r: httpx.Response) -> None:
    if r.status_code == 403:
        raise PermissionError(
            "ZoomInfo returned 403. The token is missing a scope this call needs: company "
            "search and enrich need api:data:company, contact calls need api:data:contact. "
            "Scopes are configured on the application in the ZoomInfo admin portal."
        )
    if r.status_code == 429:
        bucket = r.headers.get("X-RateLimit-Rejected-Bucket", "unknown")
        retry_after = r.headers.get("Retry-After", "unknown")
        raise RuntimeError(
            f"ZoomInfo returned 429; rejected bucket: {bucket}, Retry-After: {retry_after}s. "
            "Requests are evaluated against per-second, per-hour and per-day windows at once. "
            "A per-second rejection deserves 1-5s of exponential backoff; an hour or day "
            "rejection needs the full Retry-After, which can exceed 700 seconds. Do not retry "
            "an hour or day rejection immediately."
        )
    r.raise_for_status()


# ----- Server + tool registry -----

server = Server("zoominfo-gtm")


@server.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(
            name="zi_credit_status",
            description=(
                "Report bulk data credit position: ZoomInfo's own subscription counters "
                "(GET /data/v1/users/usage) plus this server's local ledger — spent today, "
                "daily ceiling, remaining, and per-tool spend for the last 7 days. Free; "
                "charges no credits. CALL THIS FIRST before planning any enrichment batch, "
                "and treat 'remaining' as the number of new records you may enrich today."
            ),
            inputSchema={"type": "object", "properties": {}},
        ),
        Tool(
            name="zi_search_companies",
            description=(
                "Search companies by firmographic criteria (POST /data/v1/companies/search). "
                "Read-only and free — it charges no credits and returned companies do not "
                "count against record limits, though each request counts against rate limits. "
                "Use search to narrow a target set, then enrich only the survivors. Returns "
                "ZoomInfo company ids you pass to zi_enrich_companies."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "criteria": {
                        "type": "object",
                        "description": (
                            "Company search attributes, e.g. {'industryKeywords': 'logistics', "
                            "'revenueMin': 50000000, 'employeeCountMin': 200, 'state': 'CA'}."
                        ),
                    },
                    "page_size": {"type": "integer", "default": DEFAULT_SEARCH_PAGE, "maximum": MAX_SEARCH_PAGE},
                    "page_number": {"type": "integer", "default": 1, "minimum": 1},
                    "sort": {
                        "type": "string",
                        "description": "name, employeeCount, or revenue. Prefix '-' for descending.",
                        "default": "-revenue",
                    },
                },
                "required": ["criteria"],
            },
        ),
        Tool(
            name="zi_search_contacts",
            description=(
                "Search contacts by role, seniority, function, or company "
                "(POST /data/v1/contacts/search). Read-only and free. Returns ZoomInfo contact "
                "ids for zi_enrich_contacts. Search results carry identifying fields only — "
                "verified emails and direct dials come from enrichment, which costs credits."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "criteria": {
                        "type": "object",
                        "description": (
                            "Contact search attributes, e.g. {'companyId': 344589814, "
                            "'managementLevel': 'VP Level Executives', 'department': 'Sales'}."
                        ),
                    },
                    "page_size": {"type": "integer", "default": DEFAULT_SEARCH_PAGE, "maximum": MAX_SEARCH_PAGE},
                    "page_number": {"type": "integer", "default": 1, "minimum": 1},
                },
                "required": ["criteria"],
            },
        ),
        Tool(
            name="zi_enrich_companies",
            description=(
                "Enrich up to 25 companies by ZoomInfo company id "
                "(POST /data/v1/companies/enrich). COSTS CREDITS: one bulk data credit per "
                "record returned, unless the record is already under management. No-match "
                "results and errors are not charged. Served from the local 365-day cache when "
                "possible, which costs nothing. Refuses the call outright if the worst-case "
                "cost would breach the daily ceiling — check zi_credit_status first."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "company_ids": {
                        "type": "array",
                        "items": {"type": "string"},
                        "maxItems": MAX_ENRICH_BATCH,
                        "description": "ZoomInfo company ids, at most 25 per call.",
                    },
                    "output_fields": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": (
                            "Fields to return, e.g. ['id','name','website','revenue','employeeCount']. "
                            "Narrow this — the field set does not change the credit cost but does "
                            "change how many tokens the answer burns."
                        ),
                    },
                },
                "required": ["company_ids"],
            },
        ),
        Tool(
            name="zi_enrich_contacts",
            description=(
                "Enrich up to 25 contacts by ZoomInfo contact id "
                "(POST /data/v1/contacts/enrich). COSTS CREDITS on the same terms as "
                "zi_enrich_companies: one bulk data credit per matched new record, nothing for "
                "no-match. This is the tool that returns verified business email and direct "
                "dial, so it is both the expensive one and the one carrying personal data. "
                "Cached for 365 days; budget-checked before the call."
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "contact_ids": {
                        "type": "array",
                        "items": {"type": "string"},
                        "maxItems": MAX_ENRICH_BATCH,
                        "description": "ZoomInfo contact ids, at most 25 per call.",
                    },
                    "output_fields": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Fields to return, e.g. ['id','firstName','lastName','jobTitle','email'].",
                    },
                },
                "required": ["contact_ids"],
            },
        ),
    ]


# ----- Tool dispatch -----


@server.call_tool()
async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
    require_config()

    if name == "zi_credit_status":
        return [TextContent(type="text", text=json.dumps(await _credit_status(), indent=2))]

    if name == "zi_search_companies":
        body = _search_body("CompanySearch", arguments)
        data = await zi_request("POST", "/data/v1/companies/search", json_body=body)
        return [TextContent(type="text", text=json.dumps(_slim_search(data, "companies"), indent=2))]

    if name == "zi_search_contacts":
        body = _search_body("ContactSearch", arguments)
        data = await zi_request("POST", "/data/v1/contacts/search", json_body=body)
        return [TextContent(type="text", text=json.dumps(_slim_search(data, "contacts"), indent=2))]

    if name == "zi_enrich_companies":
        return await _enrich(
            entity="companies",
            path="/data/v1/companies/enrich",
            payload_type="CompanyEnrich",
            match_key="matchCompanyInput",
            id_field="companyId",
            ids=arguments.get("company_ids") or [],
            output_fields=arguments.get("output_fields"),
            tool="zi_enrich_companies",
        )

    if name == "zi_enrich_contacts":
        return await _enrich(
            entity="contacts",
            path="/data/v1/contacts/enrich",
            payload_type="ContactEnrich",
            match_key="matchPersonInput",
            id_field="personId",
            ids=arguments.get("contact_ids") or [],
            output_fields=arguments.get("output_fields"),
            tool="zi_enrich_contacts",
        )

    raise ValueError(f"Unknown tool: {name}")


async def _credit_status() -> dict[str, Any]:
    gov = governor()
    st = gov.state()
    status: dict[str, Any] = {
        "local_ledger": {
            "run_id": RUN_ID,
            "daily_limit": st.daily_limit,
            "spent_today": st.spent_today,
            "held_by_calls_in_flight": st.reserved,
            "available_today": st.available,
        },
        "recent_spend_by_tool": gov.spend_by_tool(days=7),
    }
    # The subscription counters are authoritative for what ZoomInfo will actually allow;
    # the local ledger is authoritative for what this server will allow. They answer
    # different questions and a run can be blocked by either, so report both.
    try:
        usage = await zi_request("GET", "/data/v1/users/usage")
        status["zoominfo_subscription_usage"] = [
            {
                "limitType": u.get("limitType"),
                "description": u.get("description"),
                "totalLimit": u.get("totalLimit"),
                "currentUsage": u.get("currentUsage"),
                "usageRemaining": u.get("usageRemaining"),
            }
            for item in usage.get("data", [])
            for u in (item.get("attributes", {}) or {}).get("usage", [])
        ]
    except Exception as exc:  # noqa: BLE001 — status must degrade, not fail
        status["zoominfo_subscription_usage_error"] = (
            f"{type(exc).__name__}: {exc}. The local ceiling below is still enforced."
        )
    return status


async def _enrich(
    *,
    entity: str,
    path: str,
    payload_type: str,
    match_key: str,
    id_field: str,
    ids: list[str],
    output_fields: list[str] | None,
    tool: str,
) -> list[TextContent]:
    gov = governor()
    ids = [str(i) for i in ids if str(i).strip()]
    if not ids:
        raise ValueError(f"No ids supplied to {tool}.")
    if len(ids) > MAX_ENRICH_BATCH:
        raise ValueError(
            f"{len(ids)} ids supplied; ZoomInfo enriches at most {MAX_ENRICH_BATCH} records per "
            "request. Split the batch — and note that splitting does not reduce the credit cost, "
            "only the request size."
        )

    cached = gov.cache_get(entity, ids)
    to_fetch = [i for i in ids if i not in cached]

    if not to_fetch:
        gov.settle(tool=tool, reserved=0, requested=len(ids), charged=0, cache_hits=len(cached), no_match=0)
        return [
            TextContent(
                type="text",
                text=json.dumps(
                    {
                        "entity": entity,
                        "records": list(cached.values()),
                        "credits_charged": 0,
                        "cache_hits": len(cached),
                        "note": "Every record was served from the local cache inside the 365-day "
                        "Records Under Management window. No request was sent and no credit spent.",
                    },
                    indent=2,
                ),
            )
        ]

    try:
        reserved = gov.reserve(len(to_fetch))
    except BudgetExceeded as exc:
        # Returned as a result, not raised. A refusal the model can read lets it re-plan
        # against the remaining allowance — enrich the top 40, defer the rest — where a
        # protocol error usually just ends the turn.
        st = gov.state()
        return [
            TextContent(
                type="text",
                text=json.dumps(
                    {
                        "refused": True,
                        "reason": str(exc),
                        "requested_new_records": len(to_fetch),
                        "cache_hits_available_free": len(cached),
                        "available_today": st.available,
                        "next_step": "Re-call with at most available_today ids, or stop and tell "
                        "the human the ceiling was reached.",
                    },
                    indent=2,
                ),
            )
        ]
    body = {
        "data": {
            "type": payload_type,
            "attributes": {
                match_key: [{id_field: _coerce_id(i)} for i in to_fetch],
                **({"outputFields": output_fields} if output_fields else {}),
            },
        }
    }

    try:
        data = await zi_request("POST", path, json_body=body)
    except Exception:
        gov.release(reserved)
        raise

    fetched: dict[str, Any] = {}
    no_match = 0
    for rec in data.get("data", []):
        if rec.get("type") == "NoMatch" or (rec.get("meta", {}) or {}).get("matchStatus") == "NO_MATCH":
            no_match += 1
            continue
        rid = str(rec.get("id"))
        fetched[rid] = {"id": rid, **(rec.get("attributes", {}) or {})}

    gov.cache_put(entity, fetched)
    charged = len(fetched)
    gov.settle(
        tool=tool,
        reserved=reserved,
        requested=len(ids),
        charged=charged,
        cache_hits=len(cached),
        no_match=no_match,
    )

    st = gov.state()
    return [
        TextContent(
            type="text",
            text=json.dumps(
                {
                    "entity": entity,
                    "records": list(cached.values()) + list(fetched.values()),
                    "credits_charged": charged,
                    "cache_hits": len(cached),
                    "no_match": no_match,
                    "budget_after": {"spent_today": st.spent_today, "available_today": st.available},
                    "note": "credits_charged counts records ZoomInfo returned as a match. Records "
                    "already under management are returned without a new charge, so the true "
                    "invoice may be lower than this figure; it is never higher.",
                },
                indent=2,
            ),
        )
    ]


# ----- Request/response shaping -----


def _search_body(payload_type: str, arguments: dict[str, Any]) -> dict[str, Any]:
    page_size = max(1, min(int(arguments.get("page_size", DEFAULT_SEARCH_PAGE)), MAX_SEARCH_PAGE))
    attributes = dict(arguments.get("criteria") or {})
    if sort := arguments.get("sort"):
        attributes["sort"] = sort
    return {
        "data": {"type": payload_type, "attributes": attributes},
        "page": {"number": max(1, int(arguments.get("page_number", 1))), "size": page_size},
    }


def _slim_search(data: dict[str, Any], entity: str) -> dict[str, Any]:
    """Return ids and a thin label per hit, not the full record.

    Search is free and enrichment is not, so the only job of a search result here is to let
    the agent decide which ids are worth paying for. Returning the whole payload would
    invite the model to treat unverified search fields as enriched data.
    """
    rows = []
    for rec in data.get("data", []):
        attrs = rec.get("attributes", {}) or {}
        rows.append(
            {
                "id": rec.get("id"),
                "name": attrs.get("name") or " ".join(
                    x for x in (attrs.get("firstName"), attrs.get("lastName")) if x
                ),
                "jobTitle": attrs.get("jobTitle"),
                "company": attrs.get("companyName"),
                "website": attrs.get("website"),
            }
        )
        rows[-1] = {k: v for k, v in rows[-1].items() if v}
    meta = data.get("meta", {}) or {}
    return {
        "entity": entity,
        "returned": len(rows),
        "total_results": meta.get("totalResults"),
        "results": rows,
        "credits_charged": 0,
        "next_step": f"Pass the ids you actually want to zi_enrich_{entity}, at most 25 per call.",
    }


def _coerce_id(value: str) -> Any:
    # ZoomInfo ids are numeric in the documented request examples but arrive as strings
    # through MCP's JSON schema. Send an int when it is one; leave anything else alone.
    try:
        return int(value)
    except (TypeError, ValueError):
        return value


# ----- Entrypoint -----


async def main() -> None:
    require_config()
    governor()
    async with stdio_server() as (read, write):
        await server.run(read, write, server.create_initialization_options())


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())