Skip to content

Commit 28960bd

Browse files
committed
feat(gatekeeper): enhance cost estimation and usage tracking
- Introduced a new pricing module to compute costs based on token usage across different providers. - Added functionality to extract usage statistics from responses for OpenAI, Anthropic, Gemini, and OpenRouter. - Updated GatekeeperCompletion and GatekeeperStats models to include token counts and cost sources. - Updated documentation to clarify cost estimation logic and configuration options. - Added tests to ensure accurate cost calculations and usage extraction across various clients.
1 parent 85c373a commit 28960bd

20 files changed

Lines changed: 727 additions & 44 deletions

docs/config-reference.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ These are used when `LINUX_MCP_TOOLSET` is set to `run_script` or `both`.
6363
| `--gatekeeper.reasoning_effort`<br>`LINUX_MCP_GATEKEEPER__REASONING_EFFORT` | _(model specific)_ | Reasoning effort (`none`, `minimal`, `low`, `medium`, `high`, `xhigh`). Not all values are supported for all models. |
6464
| `--gatekeeper.structured_output`<br>`LINUX_MCP_GATEKEEPER__STRUCTURED_OUTPUT` | `True` | Whether to use structured JSON output from the model |
6565
| `--gatekeeper.temperature`<br>`LINUX_MCP_GATEKEEPER__TEMPERATURE` | 0.0 | Temperature to use for the model |
66-
| `--gatekeeper.cost`<br>`LINUX_MCP_GATEKEEPER__COST` | _(none)_ | Gatekeeper cost for accounting (`input$/token:output$/token`) |
66+
| `--gatekeeper.cost`<br>`LINUX_MCP_GATEKEEPER__COST` | _(none)_ | Optional per-token cost override for gatekeeper accounting (`input$/token:output$/token`). When unset, cost is estimated from [models.dev](https://models.dev) pricing (with a vendored fallback), except OpenRouter which uses API-reported `usage.cost` when available. Local OpenAI-compatible inference (e.g. llama.cpp on localhost) is reported as `$0`. |
6767
| `--gatekeeper.openai.base_url`<br>`LINUX_MCP_GATEKEEPER__OPENAI__BASE_URL` / `OPENAI_API_BASE` | `https://api.openai.com/v1` | OpenAI provider: API base URL |
6868
| `--gatekeeper.openai.template_kwargs`<br>`LINUX_MCP_GATEKEEPER__OPENAI__TEMPLATE_KWARGS` | _(none)_ | OpenAI provider: extra chat-template arguments (e.g. llama.cpp `enable_thinking`), sent as `chat_template_kwargs` |
6969
| `--gatekeeper.openrouter.base_url`<br>`LINUX_MCP_GATEKEEPER__OPENROUTER__BASE_URL` | `https://openrouter.ai/api/v1` | OpenRouter provider: API base URL |

docs/guarded-command-execution.md

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,17 @@ LINUX_MCP_GATEKEEPER__VERTEX_AI__PROJECT=my-gcp-project
145145
GOOGLE_APPLICATION_CREDENTIALS=<path-to-service-account.json>
146146
```
147147

148+
**Gatekeeper cost estimation**
149+
150+
Eval runs and gatekeeper stats report an estimated dollar cost per check. Resolution order:
151+
152+
1. **API-reported cost** — OpenRouter `usage.cost` when present.
153+
2. **Config override**`LINUX_MCP_GATEKEEPER__COST` as `input$/token:output$/token` (useful for Vertex MaaS eval models).
154+
3. **models.dev** — runtime pricing lookup with a vendored snapshot fallback when the network is unavailable.
155+
4. **Local inference** — OpenAI-compatible providers pointed at `localhost` / `127.0.0.1` (e.g. llama.cpp) report `$0`.
156+
5. **Fallback** — hardcoded rates for known eval models, then a conservative cloud default.
157+
158+
Totals are **estimates** except when OpenRouter returns API-reported cost. Token counts come from each provider's usage metadata in the completion response.
148159

149160
**Configure your client's tool policy**
150161

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -205,6 +205,7 @@ order-by-type = false
205205
# Include the bundled mcp app html file
206206
[tool.hatch.build.targets.wheel.force-include]
207207
"mcp-app/dist" = "linux_mcp_server/ui_resources"
208+
"src/linux_mcp_server/gatekeeper/data" = "linux_mcp_server/gatekeeper/data"
208209

209210
# Force hatch to pick up the bundled app that is git-ignored
210211
[tool.hatch.build]
Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
#!/usr/bin/env python3
2+
"""Refresh the vendored models.dev pricing snapshot used by the gatekeeper."""
3+
4+
import json
5+
import sys
6+
import urllib.request
7+
8+
from dataclasses import dataclass
9+
from pathlib import Path
10+
from typing import Any
11+
12+
13+
REPO_ROOT = Path(__file__).resolve().parents[1]
14+
OUTPUT_PATH = REPO_ROOT / "src/linux_mcp_server/gatekeeper/data/models_dev_fallback.json"
15+
MODELS_DEV_URL = "https://models.dev/api.json"
16+
17+
# Provider -> model IDs used by eval/gatekeeper/standard-evals.sh
18+
WANTED: dict[str, list[str]] = {
19+
"anthropic": ["claude-sonnet-4-6", "claude-opus-4-6", "claude-haiku-4-5"],
20+
"openai": ["gpt-5.4", "gpt-oss-20b", "gpt-oss-120b"],
21+
"google": ["gemini-2.0-flash", "gemini-3.1-pro-preview"],
22+
"openrouter": [
23+
"openai/gpt-oss-120b",
24+
"anthropic/claude-sonnet-4-6",
25+
"google/gemma-4-26b-a4b-it",
26+
"ibm-granite/granite-4.0-h-small",
27+
"qwen/qwen3.5-35b-a3b",
28+
],
29+
}
30+
31+
32+
@dataclass(frozen=True)
33+
class TokenCostPerMillion:
34+
"""USD per million tokens, matching models.dev's cost.input / cost.output fields."""
35+
36+
input: float
37+
output: float
38+
39+
def to_dict(self) -> dict[str, float]:
40+
return {"input": self.input, "output": self.output}
41+
42+
43+
@dataclass(frozen=True)
44+
class ModelPricing:
45+
"""Pricing structure for a single model. Contains the cost per million tokens."""
46+
47+
cost: TokenCostPerMillion
48+
49+
def to_dict(self) -> dict[str, Any]:
50+
return {"cost": self.cost.to_dict()}
51+
52+
@classmethod
53+
def from_models_dev_entry(cls, entry: object) -> "ModelPricing | None":
54+
if not isinstance(entry, dict):
55+
return None
56+
cost = entry.get("cost")
57+
if not isinstance(cost, dict):
58+
return None
59+
input_mtok = cost.get("input")
60+
output_mtok = cost.get("output")
61+
if not isinstance(input_mtok, (int, float)) or not isinstance(output_mtok, (int, float)):
62+
return None
63+
return cls(cost=TokenCostPerMillion(input=float(input_mtok), output=float(output_mtok)))
64+
65+
66+
@dataclass
67+
class ProviderPricing:
68+
"""Pricing structure for a single provider. Maps model ID to its pricing."""
69+
70+
models: dict[str, ModelPricing]
71+
72+
def to_dict(self) -> dict[str, Any]:
73+
return {"models": {model_id: pricing.to_dict() for model_id, pricing in self.models.items()}}
74+
75+
76+
@dataclass
77+
class PricingSnapshot:
78+
"""Snapshot of models.dev pricing for the gatekeeper. Maps provider name to a pricing structure."""
79+
80+
providers: dict[str, ProviderPricing]
81+
82+
def to_dict(self) -> dict[str, Any]:
83+
return {provider: pricing.to_dict() for provider, pricing in self.providers.items()}
84+
85+
86+
def _models_dev_catalog(data: dict[str, Any], provider: str) -> dict[str, Any]:
87+
"""Extract the models catalog for a given provider from the models.dev data."""
88+
provider_data = data.get(provider, {})
89+
if not isinstance(provider_data, dict):
90+
return {}
91+
models = provider_data.get("models", {})
92+
return models if isinstance(models, dict) else {}
93+
94+
95+
def build_snapshot(data: dict[str, Any]) -> tuple[PricingSnapshot, list[str]]:
96+
"""Build a pricing snapshot from the models.dev data."""
97+
providers: dict[str, ProviderPricing] = {}
98+
missing: list[str] = []
99+
100+
for provider, model_ids in WANTED.items():
101+
src_models = _models_dev_catalog(data, provider)
102+
picked: dict[str, ModelPricing] = {}
103+
for model_id in model_ids:
104+
pricing = ModelPricing.from_models_dev_entry(src_models.get(model_id))
105+
if pricing is None:
106+
missing.append(f"{provider}/{model_id}")
107+
else:
108+
picked[model_id] = pricing
109+
if picked:
110+
providers[provider] = ProviderPricing(models=picked)
111+
112+
return PricingSnapshot(providers=providers), missing
113+
114+
115+
def main() -> int:
116+
with urllib.request.urlopen(MODELS_DEV_URL, timeout=30) as response:
117+
data = json.load(response)
118+
119+
snapshot, missing = build_snapshot(data)
120+
121+
OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
122+
OUTPUT_PATH.write_text(json.dumps(snapshot.to_dict(), indent=2) + "\n", encoding="utf-8")
123+
print(f"Wrote {OUTPUT_PATH}")
124+
125+
if missing:
126+
print("Missing pricing for:", ", ".join(missing), file=sys.stderr)
127+
return 1
128+
return 0
129+
130+
131+
if __name__ == "__main__":
132+
raise SystemExit(main())

src/linux_mcp_server/gatekeeper/anthropic_client.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from linux_mcp_server.gatekeeper.http_utils import normalize_model_id
1111
from linux_mcp_server.gatekeeper.http_utils import post_json
1212
from linux_mcp_server.gatekeeper.schema import anthropic_output_config
13+
from linux_mcp_server.gatekeeper.usage import extract_anthropic_usage
1314
from linux_mcp_server.models import GatekeeperCompletion
1415

1516

@@ -78,4 +79,9 @@ def complete_anthropic(prompt: str, *, max_tokens: int, timeout: int = DEFAULT_T
7879
body=build_messages_body(prompt, include_model=True, max_tokens=max_tokens),
7980
timeout=timeout,
8081
)
81-
return GatekeeperCompletion(text=extract_messages_text(response))
82+
usage = extract_anthropic_usage(response)
83+
return GatekeeperCompletion(
84+
text=extract_messages_text(response),
85+
prompt_tokens=usage.input_tokens,
86+
completion_tokens=usage.output_tokens,
87+
)

src/linux_mcp_server/gatekeeper/check_run_script.py

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77

88
from linux_mcp_server.config import CONFIG
99
from linux_mcp_server.gatekeeper.llm import complete_gatekeeper
10+
from linux_mcp_server.gatekeeper.pricing import compute_cost
11+
from linux_mcp_server.gatekeeper.pricing import CostSource
1012
from linux_mcp_server.utils import StrEnum
1113

1214

@@ -184,6 +186,7 @@ class GatekeeperStats(BaseModel):
184186
prompt_tokens: int = 0
185187
completion_tokens: int = 0
186188
cost: float = 0
189+
cost_source: CostSource | None = None
187190
latency: float = 0
188191

189192

@@ -193,15 +196,6 @@ def __init__(self, message: str, *, stats: GatekeeperStats | None = None):
193196
self.stats = stats
194197

195198

196-
def _compute_cost(prompt_tokens: int, completion_tokens: int, *, usage_cost: float | None = None) -> float:
197-
if usage_cost is not None:
198-
return usage_cost
199-
if CONFIG.gatekeeper.cost is None:
200-
return 0.0
201-
input_cost, output_cost = CONFIG.gatekeeper.cost
202-
return prompt_tokens * input_cost + completion_tokens * output_cost
203-
204-
205199
async def check_run_script_with_stats(
206200
description: str, script_type: str, script: str, *, readonly: bool
207201
) -> tuple[GatekeeperResult, GatekeeperStats]:
@@ -234,14 +228,17 @@ async def check_run_script_with_stats(
234228
except asyncio.TimeoutError:
235229
raise GatekeeperException("Timeout calling gatekeeper model") from None
236230

231+
cost, cost_source = compute_cost(
232+
completion.prompt_tokens,
233+
completion.completion_tokens,
234+
usage_cost=completion.usage_cost,
235+
)
236+
237237
stats = GatekeeperStats(
238238
prompt_tokens=completion.prompt_tokens,
239239
completion_tokens=completion.completion_tokens,
240-
cost=_compute_cost(
241-
completion.prompt_tokens,
242-
completion.completion_tokens,
243-
usage_cost=completion.usage_cost,
244-
),
240+
cost=cost,
241+
cost_source=cost_source,
245242
latency=time.perf_counter() - time_before,
246243
)
247244

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
{
2+
"anthropic": {
3+
"models": {
4+
"claude-sonnet-4-6": {
5+
"cost": {
6+
"input": 3.0,
7+
"output": 15.0
8+
}
9+
},
10+
"claude-opus-4-6": {
11+
"cost": {
12+
"input": 5.0,
13+
"output": 25.0
14+
}
15+
},
16+
"claude-haiku-4-5": {
17+
"cost": {
18+
"input": 1.0,
19+
"output": 5.0
20+
}
21+
}
22+
}
23+
},
24+
"openai": {
25+
"models": {
26+
"gpt-5.4": {
27+
"cost": {
28+
"input": 2.5,
29+
"output": 15.0
30+
}
31+
},
32+
"gpt-oss-20b": {
33+
"cost": {
34+
"input": 0.07,
35+
"output": 0.3
36+
}
37+
},
38+
"gpt-oss-120b": {
39+
"cost": {
40+
"input": 0.15,
41+
"output": 0.6
42+
}
43+
}
44+
}
45+
},
46+
"google": {
47+
"models": {
48+
"gemini-2.0-flash": {
49+
"cost": {
50+
"input": 0.1,
51+
"output": 0.4
52+
}
53+
},
54+
"gemini-3.1-pro-preview": {
55+
"cost": {
56+
"input": 2.0,
57+
"output": 12.0
58+
}
59+
}
60+
}
61+
},
62+
"openrouter": {
63+
"models": {
64+
"openai/gpt-oss-120b": {
65+
"cost": {
66+
"input": 0.09,
67+
"output": 0.45
68+
}
69+
},
70+
"anthropic/claude-sonnet-4-6": {
71+
"cost": {
72+
"input": 3.0,
73+
"output": 15.0
74+
}
75+
},
76+
"google/gemma-4-26b-a4b-it": {
77+
"cost": {
78+
"input": 0.15,
79+
"output": 0.6
80+
}
81+
},
82+
"ibm-granite/granite-4.0-h-small": {
83+
"cost": {
84+
"input": 0.1,
85+
"output": 0.4
86+
}
87+
},
88+
"qwen/qwen3.5-35b-a3b": {
89+
"cost": {
90+
"input": 0.2,
91+
"output": 0.8
92+
}
93+
}
94+
}
95+
}
96+
}

src/linux_mcp_server/gatekeeper/gemini_client.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
from linux_mcp_server.gatekeeper.http_utils import normalize_model_id
1111
from linux_mcp_server.gatekeeper.http_utils import post_json
1212
from linux_mcp_server.gatekeeper.schema import gemini_generation_config
13+
from linux_mcp_server.gatekeeper.usage import extract_gemini_usage
1314
from linux_mcp_server.models import GatekeeperCompletion
1415

1516

@@ -73,4 +74,9 @@ def complete_gemini(prompt: str, *, max_tokens: int, timeout: int = DEFAULT_TIME
7374
body=build_gemini_body(prompt, max_tokens=max_tokens),
7475
timeout=timeout,
7576
)
76-
return GatekeeperCompletion(text=extract_gemini_text(response))
77+
usage = extract_gemini_usage(response)
78+
return GatekeeperCompletion(
79+
text=extract_gemini_text(response),
80+
prompt_tokens=usage.input_tokens,
81+
completion_tokens=usage.output_tokens,
82+
)

src/linux_mcp_server/gatekeeper/openai_client.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
from linux_mcp_server.gatekeeper.http_utils import post_json
1414
from linux_mcp_server.gatekeeper.schema import openai_response_format
1515
from linux_mcp_server.gatekeeper.schema import openai_text_format
16+
from linux_mcp_server.gatekeeper.usage import extract_openai_chat_completions_usage
17+
from linux_mcp_server.gatekeeper.usage import extract_openai_responses_usage
1618
from linux_mcp_server.models import GatekeeperCompletion
1719

1820

@@ -138,7 +140,12 @@ def complete_openai(prompt: str, *, max_tokens: int, timeout: int = DEFAULT_TIME
138140
body=_build_responses_body(prompt, max_tokens=max_tokens),
139141
timeout=timeout,
140142
)
141-
return GatekeeperCompletion(text=_extract_responses_text(response))
143+
usage = extract_openai_responses_usage(response)
144+
return GatekeeperCompletion(
145+
text=_extract_responses_text(response),
146+
prompt_tokens=usage.input_tokens,
147+
completion_tokens=usage.output_tokens,
148+
)
142149
except GatekeeperHTTPError as exc:
143150
if exc.status_code not in {404, 405}:
144151
raise
@@ -150,4 +157,9 @@ def complete_openai(prompt: str, *, max_tokens: int, timeout: int = DEFAULT_TIME
150157
body=build_chat_completions_body(prompt, max_tokens=max_tokens),
151158
timeout=timeout,
152159
)
153-
return GatekeeperCompletion(text=extract_chat_completions_text(response))
160+
usage = extract_openai_chat_completions_usage(response)
161+
return GatekeeperCompletion(
162+
text=extract_chat_completions_text(response),
163+
prompt_tokens=usage.input_tokens,
164+
completion_tokens=usage.output_tokens,
165+
)

0 commit comments

Comments
 (0)