citadel-predict-mcp 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- citadel_predict_mcp-0.1.0/.gitignore +36 -0
- citadel_predict_mcp-0.1.0/PKG-INFO +171 -0
- citadel_predict_mcp-0.1.0/README.md +144 -0
- citadel_predict_mcp-0.1.0/pyproject.toml +59 -0
- citadel_predict_mcp-0.1.0/src/citadel_predict_mcp/__init__.py +8 -0
- citadel_predict_mcp-0.1.0/src/citadel_predict_mcp/server.py +198 -0
- citadel_predict_mcp-0.1.0/tests/test_configure_mcp.py +215 -0
- citadel_predict_mcp-0.1.0/tests/test_server.py +250 -0
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
# Virtual Environments
|
|
2
|
+
.venv/
|
|
3
|
+
**/.venv/
|
|
4
|
+
|
|
5
|
+
# Byte-compiled / cache files
|
|
6
|
+
__pycache__/
|
|
7
|
+
**/*.pyc
|
|
8
|
+
**/*.pyo
|
|
9
|
+
**/*.pyd
|
|
10
|
+
|
|
11
|
+
# Testing & Linting Caches
|
|
12
|
+
.pytest_cache/
|
|
13
|
+
.ruff_cache/
|
|
14
|
+
.mypy_cache/
|
|
15
|
+
|
|
16
|
+
# Node & Frontend
|
|
17
|
+
node_modules/
|
|
18
|
+
**/node_modules/
|
|
19
|
+
frontend/.next/
|
|
20
|
+
frontend/out/
|
|
21
|
+
|
|
22
|
+
# Environment Variables & Secrets
|
|
23
|
+
.env
|
|
24
|
+
.env.*
|
|
25
|
+
!.env.example
|
|
26
|
+
|
|
27
|
+
# IDE & Agents
|
|
28
|
+
.agents/
|
|
29
|
+
.gemini/
|
|
30
|
+
.vscode/
|
|
31
|
+
.idea/
|
|
32
|
+
|
|
33
|
+
# Distribution & Build
|
|
34
|
+
build/
|
|
35
|
+
dist/
|
|
36
|
+
*.egg-info/
|
|
@@ -0,0 +1,171 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: citadel-predict-mcp
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Model Context Protocol (MCP) server for Citadel Predict pre-execution AI agent cost estimation
|
|
5
|
+
Author: Citadel Predict Team
|
|
6
|
+
License: MIT
|
|
7
|
+
Keywords: agent-cost,citadel-predict,claude,claude-code,claude-desktop,llm-budget,mcp,model-context-protocol,token-estimation
|
|
8
|
+
Classifier: Development Status :: 4 - Beta
|
|
9
|
+
Classifier: Intended Audience :: Developers
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
17
|
+
Requires-Python: >=3.9
|
|
18
|
+
Requires-Dist: citadel-predict>=0.1.0
|
|
19
|
+
Requires-Dist: mcp>=1.0.0
|
|
20
|
+
Provides-Extra: dev
|
|
21
|
+
Requires-Dist: mypy>=1.10.0; extra == 'dev'
|
|
22
|
+
Requires-Dist: pytest-asyncio>=0.23.0; extra == 'dev'
|
|
23
|
+
Requires-Dist: pytest-cov>=4.1.0; extra == 'dev'
|
|
24
|
+
Requires-Dist: pytest>=8.0.0; extra == 'dev'
|
|
25
|
+
Requires-Dist: ruff>=0.4.0; extra == 'dev'
|
|
26
|
+
Description-Content-Type: text/markdown
|
|
27
|
+
|
|
28
|
+
# citadel-predict-mcp
|
|
29
|
+
|
|
30
|
+
[](https://pypi.org/project/citadel-predict-mcp/)
|
|
31
|
+
[](https://opensource.org/licenses/MIT)
|
|
32
|
+
|
|
33
|
+
**Model Context Protocol (MCP) server for Citadel Predict pre-execution AI agent cost estimation.**
|
|
34
|
+
|
|
35
|
+
`citadel-predict-mcp` connects your hosted [Citadel Predict API](https://github.com/Athullvr/Citadel) directly into **Claude Desktop** and **Claude Code** via a local standard I/O (stdio) MCP server.
|
|
36
|
+
|
|
37
|
+
With this server configured, Claude can natively estimate token usage ranges ($low, expected, high$) and flag out-of-distribution risks for autonomous workflows **before running them**—without requiring manual CLI execution.
|
|
38
|
+
|
|
39
|
+
---
|
|
40
|
+
|
|
41
|
+
## Features
|
|
42
|
+
|
|
43
|
+
- ⚡ **Native Claude Tool Calling**: Claude automatically decides when to call `estimate_agent_cost` when planning or dispatching tasks.
|
|
44
|
+
- 🔒 **Zero Network Exposure**: Runs strictly as a local `stdio` subprocess spawned by Claude Desktop / Claude Code.
|
|
45
|
+
- 🎯 **Pre-Execution Guardrails**: Predicts token consumption bounds before multi-step tools or reasoning loops execute.
|
|
46
|
+
- 🛡️ **User-Friendly Error Handling**: Catches authentication, rate-limiting, and validation issues, presenting clear, actionable suggestions to Claude rather than raw stack traces.
|
|
47
|
+
|
|
48
|
+
---
|
|
49
|
+
|
|
50
|
+
## Installation
|
|
51
|
+
|
|
52
|
+
Install the package via `pip`:
|
|
53
|
+
|
|
54
|
+
```bash
|
|
55
|
+
pip install citadel-predict-mcp
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
*(For local development from the repository root: `pip install -e packages/citadel-predict-mcp`)*
|
|
59
|
+
|
|
60
|
+
Verifying installation:
|
|
61
|
+
```bash
|
|
62
|
+
citadel-predict-mcp --help
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
---
|
|
66
|
+
|
|
67
|
+
## Quickstart: One-Command Auto-Configuration (Recommended)
|
|
68
|
+
|
|
69
|
+
To automatically configure **both** Claude Desktop and Claude Code on Windows, macOS, or Linux with zero manual JSON editing:
|
|
70
|
+
|
|
71
|
+
```bash
|
|
72
|
+
# Run from repository root with your active Python environment:
|
|
73
|
+
python scripts/configure_mcp.py
|
|
74
|
+
|
|
75
|
+
# Or pass parameters non-interactively:
|
|
76
|
+
python scripts/configure_mcp.py --api-key cp_live_your_key_here --api-url http://localhost:8000
|
|
77
|
+
```
|
|
78
|
+
|
|
79
|
+
The script automatically detects your active virtual environment, locates the platform-specific executable (`citadel-predict-mcp.exe` on Windows or `citadel-predict-mcp` on macOS/Linux), and cleanly merges the server definition into both `claude_desktop_config.json` and `.claude/settings.json` while preserving all existing configurations.
|
|
80
|
+
|
|
81
|
+
---
|
|
82
|
+
|
|
83
|
+
## Manual Configuration (Alternative)
|
|
84
|
+
|
|
85
|
+
If you prefer to configure manually:
|
|
86
|
+
|
|
87
|
+
### 1. Claude Desktop (`claude_desktop_config.json`)
|
|
88
|
+
|
|
89
|
+
| Operating System | Exact Configuration File Path |
|
|
90
|
+
| :--- | :--- |
|
|
91
|
+
| **macOS** | `~/Library/Application Support/Claude/claude_desktop_config.json` |
|
|
92
|
+
| **Windows** | `%APPDATA%\Claude\claude_desktop_config.json` |
|
|
93
|
+
| **Linux** | `~/.config/Claude/claude_desktop_config.json` |
|
|
94
|
+
|
|
95
|
+
Add `citadel-predict` under `mcpServers`:
|
|
96
|
+
|
|
97
|
+
```json
|
|
98
|
+
{
|
|
99
|
+
"mcpServers": {
|
|
100
|
+
"citadel-predict": {
|
|
101
|
+
"command": "citadel-predict-mcp",
|
|
102
|
+
"env": {
|
|
103
|
+
"CITADEL_API_KEY": "cp_live_your_api_key_here",
|
|
104
|
+
"CITADEL_API_URL": "http://localhost:8000"
|
|
105
|
+
}
|
|
106
|
+
}
|
|
107
|
+
}
|
|
108
|
+
}
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
### 2. Claude Code (`.claude/settings.json`)
|
|
112
|
+
|
|
113
|
+
```bash
|
|
114
|
+
claude mcp add citadel-predict -- citadel-predict-mcp
|
|
115
|
+
```
|
|
116
|
+
|
|
117
|
+
---
|
|
118
|
+
|
|
119
|
+
## Available Tools
|
|
120
|
+
|
|
121
|
+
### `estimate_agent_cost`
|
|
122
|
+
|
|
123
|
+
**Description**:
|
|
124
|
+
> *Estimate token cost and usage range for an AI agent task BEFORE running it. Use this when the user is about to execute, dispatch, or run a multi-step agent task and cost/budget matters.*
|
|
125
|
+
|
|
126
|
+
**Input Parameters**:
|
|
127
|
+
- `task_text` (*string, required*): The natural language description of the agent task (1 to 4000 characters).
|
|
128
|
+
- `tools` (*array of strings, optional*): List of tool names available to the agent (e.g. `["web_search", "draft_document"]`).
|
|
129
|
+
- `num_tools` (*integer, optional*): Tool count if specific tool names are not listed.
|
|
130
|
+
- `model_id` (*string, optional, default: `"claude-sonnet"`*): Model calibration profile to evaluate against.
|
|
131
|
+
|
|
132
|
+
**Sample Return Payload**:
|
|
133
|
+
```json
|
|
134
|
+
{
|
|
135
|
+
"success": true,
|
|
136
|
+
"model_id": "claude-sonnet",
|
|
137
|
+
"expected_tokens": 3200,
|
|
138
|
+
"low_tokens": 1500,
|
|
139
|
+
"high_tokens": 5800,
|
|
140
|
+
"out_of_distribution": false,
|
|
141
|
+
"ood_reasons": [],
|
|
142
|
+
"confidence": "normal",
|
|
143
|
+
"driving_factors": ["task_length", "tools_count"],
|
|
144
|
+
"summary": "Expected: 3,200 tokens (Range: 1,500 – 5,800)"
|
|
145
|
+
}
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
---
|
|
149
|
+
|
|
150
|
+
## Manual Test Checklist for Testers
|
|
151
|
+
|
|
152
|
+
Follow this 5-minute checklist to verify your MCP setup:
|
|
153
|
+
|
|
154
|
+
1. **Installation**:
|
|
155
|
+
- [ ] Run `citadel-predict-mcp --help` in your terminal to verify the command is accessible on your PATH.
|
|
156
|
+
2. **Configuration**:
|
|
157
|
+
- [ ] Add the JSON entry to `claude_desktop_config.json` with your active Citadel API key.
|
|
158
|
+
3. **Restart**:
|
|
159
|
+
- [ ] Fully quit and reopen Claude Desktop.
|
|
160
|
+
4. **Invocation Test**:
|
|
161
|
+
- [ ] Send the following prompt in a new Claude Desktop chat:
|
|
162
|
+
> *"I am planning to have an agent research competitor pricing across 5 company sites and compile a markdown report. Estimate the token cost and budget range before we start."*
|
|
163
|
+
5. **Verification**:
|
|
164
|
+
- [ ] Confirm Claude invokes the `estimate_agent_cost` tool (indicated by a tool-call widget in the conversation).
|
|
165
|
+
- [ ] Confirm Claude receives the token ranges (`expected_tokens`, `low_tokens`, `high_tokens`) and presents a natural language summary with the budget estimate to you.
|
|
166
|
+
|
|
167
|
+
---
|
|
168
|
+
|
|
169
|
+
## License
|
|
170
|
+
|
|
171
|
+
MIT
|
|
@@ -0,0 +1,144 @@
|
|
|
1
|
+
# citadel-predict-mcp
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/citadel-predict-mcp/)
|
|
4
|
+
[](https://opensource.org/licenses/MIT)
|
|
5
|
+
|
|
6
|
+
**Model Context Protocol (MCP) server for Citadel Predict pre-execution AI agent cost estimation.**
|
|
7
|
+
|
|
8
|
+
`citadel-predict-mcp` connects your hosted [Citadel Predict API](https://github.com/Athullvr/Citadel) directly into **Claude Desktop** and **Claude Code** via a local standard I/O (stdio) MCP server.
|
|
9
|
+
|
|
10
|
+
With this server configured, Claude can natively estimate token usage ranges ($low, expected, high$) and flag out-of-distribution risks for autonomous workflows **before running them**—without requiring manual CLI execution.
|
|
11
|
+
|
|
12
|
+
---
|
|
13
|
+
|
|
14
|
+
## Features
|
|
15
|
+
|
|
16
|
+
- ⚡ **Native Claude Tool Calling**: Claude automatically decides when to call `estimate_agent_cost` when planning or dispatching tasks.
|
|
17
|
+
- 🔒 **Zero Network Exposure**: Runs strictly as a local `stdio` subprocess spawned by Claude Desktop / Claude Code.
|
|
18
|
+
- 🎯 **Pre-Execution Guardrails**: Predicts token consumption bounds before multi-step tools or reasoning loops execute.
|
|
19
|
+
- 🛡️ **User-Friendly Error Handling**: Catches authentication, rate-limiting, and validation issues, presenting clear, actionable suggestions to Claude rather than raw stack traces.
|
|
20
|
+
|
|
21
|
+
---
|
|
22
|
+
|
|
23
|
+
## Installation
|
|
24
|
+
|
|
25
|
+
Install the package via `pip`:
|
|
26
|
+
|
|
27
|
+
```bash
|
|
28
|
+
pip install citadel-predict-mcp
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
*(For local development from the repository root: `pip install -e packages/citadel-predict-mcp`)*
|
|
32
|
+
|
|
33
|
+
Verifying installation:
|
|
34
|
+
```bash
|
|
35
|
+
citadel-predict-mcp --help
|
|
36
|
+
```
|
|
37
|
+
|
|
38
|
+
---
|
|
39
|
+
|
|
40
|
+
## Quickstart: One-Command Auto-Configuration (Recommended)
|
|
41
|
+
|
|
42
|
+
To automatically configure **both** Claude Desktop and Claude Code on Windows, macOS, or Linux with zero manual JSON editing:
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
# Run from repository root with your active Python environment:
|
|
46
|
+
python scripts/configure_mcp.py
|
|
47
|
+
|
|
48
|
+
# Or pass parameters non-interactively:
|
|
49
|
+
python scripts/configure_mcp.py --api-key cp_live_your_key_here --api-url http://localhost:8000
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
The script automatically detects your active virtual environment, locates the platform-specific executable (`citadel-predict-mcp.exe` on Windows or `citadel-predict-mcp` on macOS/Linux), and cleanly merges the server definition into both `claude_desktop_config.json` and `.claude/settings.json` while preserving all existing configurations.
|
|
53
|
+
|
|
54
|
+
---
|
|
55
|
+
|
|
56
|
+
## Manual Configuration (Alternative)
|
|
57
|
+
|
|
58
|
+
If you prefer to configure manually:
|
|
59
|
+
|
|
60
|
+
### 1. Claude Desktop (`claude_desktop_config.json`)
|
|
61
|
+
|
|
62
|
+
| Operating System | Exact Configuration File Path |
|
|
63
|
+
| :--- | :--- |
|
|
64
|
+
| **macOS** | `~/Library/Application Support/Claude/claude_desktop_config.json` |
|
|
65
|
+
| **Windows** | `%APPDATA%\Claude\claude_desktop_config.json` |
|
|
66
|
+
| **Linux** | `~/.config/Claude/claude_desktop_config.json` |
|
|
67
|
+
|
|
68
|
+
Add `citadel-predict` under `mcpServers`:
|
|
69
|
+
|
|
70
|
+
```json
|
|
71
|
+
{
|
|
72
|
+
"mcpServers": {
|
|
73
|
+
"citadel-predict": {
|
|
74
|
+
"command": "citadel-predict-mcp",
|
|
75
|
+
"env": {
|
|
76
|
+
"CITADEL_API_KEY": "cp_live_your_api_key_here",
|
|
77
|
+
"CITADEL_API_URL": "http://localhost:8000"
|
|
78
|
+
}
|
|
79
|
+
}
|
|
80
|
+
}
|
|
81
|
+
}
|
|
82
|
+
```
|
|
83
|
+
|
|
84
|
+
### 2. Claude Code (`.claude/settings.json`)
|
|
85
|
+
|
|
86
|
+
```bash
|
|
87
|
+
claude mcp add citadel-predict -- citadel-predict-mcp
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
---
|
|
91
|
+
|
|
92
|
+
## Available Tools
|
|
93
|
+
|
|
94
|
+
### `estimate_agent_cost`
|
|
95
|
+
|
|
96
|
+
**Description**:
|
|
97
|
+
> *Estimate token cost and usage range for an AI agent task BEFORE running it. Use this when the user is about to execute, dispatch, or run a multi-step agent task and cost/budget matters.*
|
|
98
|
+
|
|
99
|
+
**Input Parameters**:
|
|
100
|
+
- `task_text` (*string, required*): The natural language description of the agent task (1 to 4000 characters).
|
|
101
|
+
- `tools` (*array of strings, optional*): List of tool names available to the agent (e.g. `["web_search", "draft_document"]`).
|
|
102
|
+
- `num_tools` (*integer, optional*): Tool count if specific tool names are not listed.
|
|
103
|
+
- `model_id` (*string, optional, default: `"claude-sonnet"`*): Model calibration profile to evaluate against.
|
|
104
|
+
|
|
105
|
+
**Sample Return Payload**:
|
|
106
|
+
```json
|
|
107
|
+
{
|
|
108
|
+
"success": true,
|
|
109
|
+
"model_id": "claude-sonnet",
|
|
110
|
+
"expected_tokens": 3200,
|
|
111
|
+
"low_tokens": 1500,
|
|
112
|
+
"high_tokens": 5800,
|
|
113
|
+
"out_of_distribution": false,
|
|
114
|
+
"ood_reasons": [],
|
|
115
|
+
"confidence": "normal",
|
|
116
|
+
"driving_factors": ["task_length", "tools_count"],
|
|
117
|
+
"summary": "Expected: 3,200 tokens (Range: 1,500 – 5,800)"
|
|
118
|
+
}
|
|
119
|
+
```
|
|
120
|
+
|
|
121
|
+
---
|
|
122
|
+
|
|
123
|
+
## Manual Test Checklist for Testers
|
|
124
|
+
|
|
125
|
+
Follow this 5-minute checklist to verify your MCP setup:
|
|
126
|
+
|
|
127
|
+
1. **Installation**:
|
|
128
|
+
- [ ] Run `citadel-predict-mcp --help` in your terminal to verify the command is accessible on your PATH.
|
|
129
|
+
2. **Configuration**:
|
|
130
|
+
- [ ] Add the JSON entry to `claude_desktop_config.json` with your active Citadel API key.
|
|
131
|
+
3. **Restart**:
|
|
132
|
+
- [ ] Fully quit and reopen Claude Desktop.
|
|
133
|
+
4. **Invocation Test**:
|
|
134
|
+
- [ ] Send the following prompt in a new Claude Desktop chat:
|
|
135
|
+
> *"I am planning to have an agent research competitor pricing across 5 company sites and compile a markdown report. Estimate the token cost and budget range before we start."*
|
|
136
|
+
5. **Verification**:
|
|
137
|
+
- [ ] Confirm Claude invokes the `estimate_agent_cost` tool (indicated by a tool-call widget in the conversation).
|
|
138
|
+
- [ ] Confirm Claude receives the token ranges (`expected_tokens`, `low_tokens`, `high_tokens`) and presents a natural language summary with the budget estimate to you.
|
|
139
|
+
|
|
140
|
+
---
|
|
141
|
+
|
|
142
|
+
## License
|
|
143
|
+
|
|
144
|
+
MIT
|
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "citadel-predict-mcp"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Model Context Protocol (MCP) server for Citadel Predict pre-execution AI agent cost estimation"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [
|
|
13
|
+
{ name = "Citadel Predict Team" }
|
|
14
|
+
]
|
|
15
|
+
keywords = [
|
|
16
|
+
"mcp",
|
|
17
|
+
"model-context-protocol",
|
|
18
|
+
"claude",
|
|
19
|
+
"claude-desktop",
|
|
20
|
+
"claude-code",
|
|
21
|
+
"citadel-predict",
|
|
22
|
+
"token-estimation",
|
|
23
|
+
"agent-cost",
|
|
24
|
+
"llm-budget",
|
|
25
|
+
]
|
|
26
|
+
classifiers = [
|
|
27
|
+
"Development Status :: 4 - Beta",
|
|
28
|
+
"Intended Audience :: Developers",
|
|
29
|
+
"License :: OSI Approved :: MIT License",
|
|
30
|
+
"Programming Language :: Python :: 3",
|
|
31
|
+
"Programming Language :: Python :: 3.9",
|
|
32
|
+
"Programming Language :: Python :: 3.10",
|
|
33
|
+
"Programming Language :: Python :: 3.11",
|
|
34
|
+
"Programming Language :: Python :: 3.12",
|
|
35
|
+
"Topic :: Software Development :: Libraries :: Python Modules",
|
|
36
|
+
]
|
|
37
|
+
dependencies = [
|
|
38
|
+
"citadel-predict>=0.1.0",
|
|
39
|
+
"mcp>=1.0.0",
|
|
40
|
+
]
|
|
41
|
+
|
|
42
|
+
[project.optional-dependencies]
|
|
43
|
+
dev = [
|
|
44
|
+
"pytest>=8.0.0",
|
|
45
|
+
"pytest-cov>=4.1.0",
|
|
46
|
+
"pytest-asyncio>=0.23.0",
|
|
47
|
+
"ruff>=0.4.0",
|
|
48
|
+
"mypy>=1.10.0",
|
|
49
|
+
]
|
|
50
|
+
|
|
51
|
+
[project.scripts]
|
|
52
|
+
citadel-predict-mcp = "citadel_predict_mcp.server:main"
|
|
53
|
+
|
|
54
|
+
[tool.hatch.build.targets.wheel]
|
|
55
|
+
packages = ["src/citadel_predict_mcp"]
|
|
56
|
+
|
|
57
|
+
[tool.pytest.ini_options]
|
|
58
|
+
testpaths = ["tests"]
|
|
59
|
+
pythonpath = ["src"]
|
|
@@ -0,0 +1,198 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Model Context Protocol (MCP) server for Citadel Predict.
|
|
3
|
+
|
|
4
|
+
Exposes pre-execution token and cost prediction tools for AI agent workflows
|
|
5
|
+
over stdio transport for Claude Desktop and Claude Code.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import sys
|
|
9
|
+
from typing import Any, Optional
|
|
10
|
+
|
|
11
|
+
from citadel_predict import (
|
|
12
|
+
CitadelAuthError,
|
|
13
|
+
CitadelBadRequestError,
|
|
14
|
+
CitadelError,
|
|
15
|
+
CitadelNetworkError,
|
|
16
|
+
CitadelRateLimitError,
|
|
17
|
+
CitadelServerError,
|
|
18
|
+
CitadelValidationError,
|
|
19
|
+
predict_cost,
|
|
20
|
+
)
|
|
21
|
+
|
|
22
|
+
try:
|
|
23
|
+
from mcp.server.mcpserver import MCPServer
|
|
24
|
+
except ImportError: # pragma: no cover
|
|
25
|
+
from mcp.server.fastmcp import FastMCP as MCPServer # type: ignore[no-redef]
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def create_server() -> MCPServer:
|
|
29
|
+
"""
|
|
30
|
+
Factory function to instantiate and configure the Citadel Predict MCP server.
|
|
31
|
+
"""
|
|
32
|
+
app = MCPServer(
|
|
33
|
+
name="citadel-predict",
|
|
34
|
+
instructions=(
|
|
35
|
+
"Citadel Predict MCP Server provides pre-execution token budget "
|
|
36
|
+
"and cost estimation for AI agent tasks before running them."
|
|
37
|
+
),
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
@app.tool(
|
|
41
|
+
name="estimate_agent_cost",
|
|
42
|
+
description=(
|
|
43
|
+
"Estimate token cost and usage range for an AI agent task BEFORE running it. "
|
|
44
|
+
"Use this when the user is about to execute, dispatch, or run a multi-step agent "
|
|
45
|
+
"task and cost/budget matters."
|
|
46
|
+
),
|
|
47
|
+
)
|
|
48
|
+
def estimate_agent_cost(
|
|
49
|
+
task_text: str,
|
|
50
|
+
tools: Optional[list[str]] = None,
|
|
51
|
+
num_tools: Optional[int] = None,
|
|
52
|
+
model_id: str = "claude-sonnet",
|
|
53
|
+
) -> dict[str, Any]:
|
|
54
|
+
"""
|
|
55
|
+
Estimate token usage and cost bounds for an AI agent task.
|
|
56
|
+
|
|
57
|
+
Args:
|
|
58
|
+
task_text: Natural language task description (1-4000 characters).
|
|
59
|
+
tools: Optional list of tool names available to the agent (e.g. ['web_search', 'draft_doc']).
|
|
60
|
+
num_tools: Optional tool count if tool names are unspecified.
|
|
61
|
+
model_id: Model calibration identifier (default: 'claude-sonnet').
|
|
62
|
+
|
|
63
|
+
Returns:
|
|
64
|
+
Dictionary containing token estimates (expected_tokens, low_tokens, high_tokens),
|
|
65
|
+
out-of-distribution flags, and diagnostic reasons.
|
|
66
|
+
"""
|
|
67
|
+
try:
|
|
68
|
+
result = predict_cost(
|
|
69
|
+
task_text=task_text,
|
|
70
|
+
tools=tools,
|
|
71
|
+
num_tools=num_tools,
|
|
72
|
+
model_id=model_id,
|
|
73
|
+
)
|
|
74
|
+
return {
|
|
75
|
+
"success": True,
|
|
76
|
+
"model_id": result.get("model_id", model_id),
|
|
77
|
+
"expected_tokens": result.get("expected_tokens"),
|
|
78
|
+
"low_tokens": result.get("low_tokens"),
|
|
79
|
+
"high_tokens": result.get("high_tokens"),
|
|
80
|
+
"out_of_distribution": result.get("out_of_distribution", False),
|
|
81
|
+
"ood_reasons": result.get("ood_reasons", []),
|
|
82
|
+
"confidence": result.get("confidence", "normal"),
|
|
83
|
+
"driving_factors": result.get("driving_factors", []),
|
|
84
|
+
"summary": (
|
|
85
|
+
f"Expected: {result.get('expected_tokens', 0):,} tokens "
|
|
86
|
+
f"(Range: {result.get('low_tokens', 0):,} – {result.get('high_tokens', 0):,})"
|
|
87
|
+
),
|
|
88
|
+
}
|
|
89
|
+
except CitadelAuthError as exc:
|
|
90
|
+
return {
|
|
91
|
+
"success": False,
|
|
92
|
+
"error_type": "AuthenticationError",
|
|
93
|
+
"message": (
|
|
94
|
+
"Authentication failed: Missing or invalid Citadel API key. "
|
|
95
|
+
"Ensure CITADEL_API_KEY environment variable is configured in your "
|
|
96
|
+
"Claude Desktop or Claude Code configuration, or in ~/.citadel/config.toml."
|
|
97
|
+
),
|
|
98
|
+
"details": str(exc),
|
|
99
|
+
}
|
|
100
|
+
except CitadelRateLimitError as exc:
|
|
101
|
+
retry_note = f" (retry after {exc.retry_after}s)" if exc.retry_after else ""
|
|
102
|
+
return {
|
|
103
|
+
"success": False,
|
|
104
|
+
"error_type": "RateLimitError",
|
|
105
|
+
"message": (
|
|
106
|
+
f"Citadel Predict API rate limit exceeded{retry_note}. "
|
|
107
|
+
"Please wait a moment before sending more prediction requests."
|
|
108
|
+
),
|
|
109
|
+
"retry_after": exc.retry_after,
|
|
110
|
+
}
|
|
111
|
+
except CitadelValidationError as exc:
|
|
112
|
+
return {
|
|
113
|
+
"success": False,
|
|
114
|
+
"error_type": "ValidationError",
|
|
115
|
+
"message": (
|
|
116
|
+
f"Validation failed for task or tool parameters: {exc.message}. "
|
|
117
|
+
"Ensure task_text is between 1 and 4000 characters."
|
|
118
|
+
),
|
|
119
|
+
"details": exc.message,
|
|
120
|
+
}
|
|
121
|
+
except CitadelBadRequestError as exc:
|
|
122
|
+
return {
|
|
123
|
+
"success": False,
|
|
124
|
+
"error_type": "BadRequestError",
|
|
125
|
+
"message": f"Bad request rejected by Citadel Predict API: {exc.message}",
|
|
126
|
+
"details": exc.message,
|
|
127
|
+
}
|
|
128
|
+
except CitadelServerError as exc:
|
|
129
|
+
return {
|
|
130
|
+
"success": False,
|
|
131
|
+
"error_type": "ServerError",
|
|
132
|
+
"message": (
|
|
133
|
+
f"Citadel Predict server error ({exc.status_code}): {exc.message}. "
|
|
134
|
+
"Please try again shortly."
|
|
135
|
+
),
|
|
136
|
+
"details": exc.message,
|
|
137
|
+
}
|
|
138
|
+
except CitadelNetworkError as exc:
|
|
139
|
+
return {
|
|
140
|
+
"success": False,
|
|
141
|
+
"error_type": "NetworkError",
|
|
142
|
+
"message": (
|
|
143
|
+
"Unable to reach the Citadel Predict API. "
|
|
144
|
+
"Please verify your internet connection or API endpoint status."
|
|
145
|
+
),
|
|
146
|
+
"details": str(exc),
|
|
147
|
+
}
|
|
148
|
+
except CitadelError as exc:
|
|
149
|
+
return {
|
|
150
|
+
"success": False,
|
|
151
|
+
"error_type": "CitadelError",
|
|
152
|
+
"message": f"Citadel Predict request error: {exc.message}",
|
|
153
|
+
"details": str(exc),
|
|
154
|
+
}
|
|
155
|
+
except Exception as exc:
|
|
156
|
+
return {
|
|
157
|
+
"success": False,
|
|
158
|
+
"error_type": "UnexpectedError",
|
|
159
|
+
"message": f"Unexpected error during cost estimation: {exc}",
|
|
160
|
+
"details": str(exc),
|
|
161
|
+
}
|
|
162
|
+
|
|
163
|
+
return app
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
server = create_server()
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
def main() -> None:
|
|
170
|
+
"""Entrypoint for the citadel-predict-mcp CLI command (runs over stdio transport)."""
|
|
171
|
+
if len(sys.argv) > 1 and sys.argv[1] in ("-h", "--help", "-v", "--version"):
|
|
172
|
+
if sys.argv[1] in ("-v", "--version"):
|
|
173
|
+
from . import __version__
|
|
174
|
+
|
|
175
|
+
print(f"citadel-predict-mcp v{__version__}")
|
|
176
|
+
return
|
|
177
|
+
|
|
178
|
+
print("Citadel Predict MCP Server (stdio transport)")
|
|
179
|
+
print("\nUsage:")
|
|
180
|
+
print(" citadel-predict-mcp Run as an MCP stdio server (used by Claude Desktop/Code)")
|
|
181
|
+
print(" citadel-predict-mcp --help Show this help message")
|
|
182
|
+
print(" citadel-predict-mcp --version Show version")
|
|
183
|
+
print("\nConfiguration for Claude Desktop (claude_desktop_config.json):")
|
|
184
|
+
print(" {")
|
|
185
|
+
print(' "mcpServers": {')
|
|
186
|
+
print(' "citadel-predict": {')
|
|
187
|
+
print(' "command": "citadel-predict-mcp",')
|
|
188
|
+
print(' "env": { "CITADEL_API_KEY": "cp_live_..." }')
|
|
189
|
+
print(" }")
|
|
190
|
+
print(" }")
|
|
191
|
+
print(" }")
|
|
192
|
+
return
|
|
193
|
+
|
|
194
|
+
server.run(transport="stdio")
|
|
195
|
+
|
|
196
|
+
|
|
197
|
+
if __name__ == "__main__":
|
|
198
|
+
main()
|
|
@@ -0,0 +1,215 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Unit and integration tests for scripts/configure_mcp.py.
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
import os
|
|
7
|
+
import subprocess
|
|
8
|
+
import sys
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
|
|
11
|
+
import pytest
|
|
12
|
+
|
|
13
|
+
# Import configure_mcp directly
|
|
14
|
+
REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
|
|
15
|
+
sys.path.insert(0, str(REPO_ROOT / "scripts"))
|
|
16
|
+
from configure_mcp import (
|
|
17
|
+
build_server_config,
|
|
18
|
+
find_mcp_executable,
|
|
19
|
+
get_claude_desktop_config_path,
|
|
20
|
+
run_configuration,
|
|
21
|
+
update_config_file,
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def test_build_server_config():
|
|
26
|
+
"""Verify single source of truth configuration builder."""
|
|
27
|
+
cfg = build_server_config(
|
|
28
|
+
command_path="/path/to/citadel-predict-mcp",
|
|
29
|
+
api_key="cp_live_test123",
|
|
30
|
+
api_url="https://api.citadel.dev",
|
|
31
|
+
)
|
|
32
|
+
assert cfg["command"] == "/path/to/citadel-predict-mcp"
|
|
33
|
+
assert cfg["env"]["CITADEL_API_KEY"] == "cp_live_test123"
|
|
34
|
+
assert cfg["env"]["CITADEL_API_URL"] == "https://api.citadel.dev"
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_build_server_config_no_key():
|
|
38
|
+
"""Verify configuration builder when no API key is specified."""
|
|
39
|
+
cfg = build_server_config(
|
|
40
|
+
command_path="/path/to/citadel-predict-mcp",
|
|
41
|
+
api_key=None,
|
|
42
|
+
api_url="http://localhost:8000",
|
|
43
|
+
)
|
|
44
|
+
assert cfg["command"] == "/path/to/citadel-predict-mcp"
|
|
45
|
+
assert cfg["env"] == {"CITADEL_API_URL": "http://localhost:8000"}
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_platform_desktop_paths():
|
|
49
|
+
"""Verify desktop config paths resolve correctly for Windows, macOS, and Linux."""
|
|
50
|
+
win_path = get_claude_desktop_config_path("Windows")
|
|
51
|
+
assert win_path.name == "claude_desktop_config.json"
|
|
52
|
+
assert "Claude" in str(win_path)
|
|
53
|
+
|
|
54
|
+
mac_path = get_claude_desktop_config_path("Darwin")
|
|
55
|
+
assert mac_path == Path.home() / "Library" / "Application Support" / "Claude" / "claude_desktop_config.json"
|
|
56
|
+
|
|
57
|
+
linux_path = get_claude_desktop_config_path("Linux")
|
|
58
|
+
assert linux_path.name == "claude_desktop_config.json"
|
|
59
|
+
assert ".config" in str(linux_path) or "Claude" in str(linux_path)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def test_executable_detection():
|
|
63
|
+
"""Verify find_mcp_executable finds a real, existing executable on disk."""
|
|
64
|
+
exe = find_mcp_executable()
|
|
65
|
+
assert isinstance(exe, str)
|
|
66
|
+
assert Path(exe).exists()
|
|
67
|
+
assert Path(exe).is_file()
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def test_run_configuration_and_valid_json(tmp_path):
|
|
71
|
+
"""
|
|
72
|
+
Asserts:
|
|
73
|
+
(a) Resulting JSON is valid in both Claude Desktop and Claude Code targets.
|
|
74
|
+
(b) 'command' path actually exists on disk.
|
|
75
|
+
"""
|
|
76
|
+
desktop_target = tmp_path / "desktop" / "claude_desktop_config.json"
|
|
77
|
+
code_target = tmp_path / "code" / ".claude" / "settings.json"
|
|
78
|
+
|
|
79
|
+
res = run_configuration(
|
|
80
|
+
api_key="cp_live_secret_key",
|
|
81
|
+
api_url="http://localhost:8000",
|
|
82
|
+
desktop_config_path=desktop_target,
|
|
83
|
+
code_config_path=code_target,
|
|
84
|
+
)
|
|
85
|
+
|
|
86
|
+
# Check files exist
|
|
87
|
+
assert desktop_target.exists()
|
|
88
|
+
assert code_target.exists()
|
|
89
|
+
|
|
90
|
+
# (a) JSON is valid
|
|
91
|
+
desktop_data = json.loads(desktop_target.read_text(encoding="utf-8"))
|
|
92
|
+
code_data = json.loads(code_target.read_text(encoding="utf-8"))
|
|
93
|
+
|
|
94
|
+
assert "mcpServers" in desktop_data
|
|
95
|
+
assert "citadel-predict" in desktop_data["mcpServers"]
|
|
96
|
+
assert "mcpServers" in code_data
|
|
97
|
+
assert "citadel-predict" in code_data["mcpServers"]
|
|
98
|
+
|
|
99
|
+
# (b) Command path exists on disk
|
|
100
|
+
cmd = desktop_data["mcpServers"]["citadel-predict"]["command"]
|
|
101
|
+
assert Path(cmd).exists()
|
|
102
|
+
assert Path(cmd).is_file()
|
|
103
|
+
assert desktop_data["mcpServers"]["citadel-predict"]["env"]["CITADEL_API_KEY"] == "cp_live_secret_key"
|
|
104
|
+
assert desktop_data["mcpServers"]["citadel-predict"] == code_data["mcpServers"]["citadel-predict"]
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def test_run_configuration_idempotency(tmp_path):
|
|
108
|
+
"""
|
|
109
|
+
Asserts:
|
|
110
|
+
(c) Running configure_mcp twice produces no diff on the second run.
|
|
111
|
+
"""
|
|
112
|
+
desktop_target = tmp_path / "desktop.json"
|
|
113
|
+
code_target = tmp_path / "code.json"
|
|
114
|
+
|
|
115
|
+
res1 = run_configuration(
|
|
116
|
+
api_key="cp_live_test",
|
|
117
|
+
api_url="http://localhost:8000",
|
|
118
|
+
desktop_config_path=desktop_target,
|
|
119
|
+
code_config_path=code_target,
|
|
120
|
+
)
|
|
121
|
+
assert res1["targets"]["claude_desktop"]["modified"] is True
|
|
122
|
+
assert res1["targets"]["claude_code"]["modified"] is True
|
|
123
|
+
|
|
124
|
+
content1_desktop = desktop_target.read_text(encoding="utf-8")
|
|
125
|
+
content1_code = code_target.read_text(encoding="utf-8")
|
|
126
|
+
|
|
127
|
+
# Second run with exact same parameters
|
|
128
|
+
res2 = run_configuration(
|
|
129
|
+
api_key="cp_live_test",
|
|
130
|
+
api_url="http://localhost:8000",
|
|
131
|
+
desktop_config_path=desktop_target,
|
|
132
|
+
code_config_path=code_target,
|
|
133
|
+
)
|
|
134
|
+
assert res2["targets"]["claude_desktop"]["modified"] is False
|
|
135
|
+
assert res2["targets"]["claude_code"]["modified"] is False
|
|
136
|
+
|
|
137
|
+
content2_desktop = desktop_target.read_text(encoding="utf-8")
|
|
138
|
+
content2_code = code_target.read_text(encoding="utf-8")
|
|
139
|
+
|
|
140
|
+
# Content must be identical (zero diff)
|
|
141
|
+
assert content1_desktop == content2_desktop
|
|
142
|
+
assert content1_code == content2_code
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
def test_preserves_existing_unrelated_configs(tmp_path):
|
|
146
|
+
"""
|
|
147
|
+
Asserts:
|
|
148
|
+
(d) Existing unrelated mcpServers entries and other top-level keys are preserved.
|
|
149
|
+
"""
|
|
150
|
+
desktop_target = tmp_path / "claude_desktop_config.json"
|
|
151
|
+
existing_content = {
|
|
152
|
+
"theme": "dark",
|
|
153
|
+
"mcpServers": {
|
|
154
|
+
"github": {
|
|
155
|
+
"command": "npx",
|
|
156
|
+
"args": ["-y", "@modelcontextprotocol/server-github"],
|
|
157
|
+
},
|
|
158
|
+
"filesystem": {
|
|
159
|
+
"command": "npx",
|
|
160
|
+
"args": ["-y", "@modelcontextprotocol/server-filesystem", "/Users/example/Desktop"],
|
|
161
|
+
},
|
|
162
|
+
},
|
|
163
|
+
}
|
|
164
|
+
desktop_target.write_text(json.dumps(existing_content, indent=2), encoding="utf-8")
|
|
165
|
+
|
|
166
|
+
run_configuration(
|
|
167
|
+
api_key="cp_live_citadel",
|
|
168
|
+
api_url="https://citadel.example.com",
|
|
169
|
+
desktop_config_path=desktop_target,
|
|
170
|
+
skip_code=True,
|
|
171
|
+
)
|
|
172
|
+
|
|
173
|
+
updated_data = json.loads(desktop_target.read_text(encoding="utf-8"))
|
|
174
|
+
|
|
175
|
+
# Top level keys preserved
|
|
176
|
+
assert updated_data.get("theme") == "dark"
|
|
177
|
+
|
|
178
|
+
# Unrelated servers preserved intact
|
|
179
|
+
assert "github" in updated_data["mcpServers"]
|
|
180
|
+
assert updated_data["mcpServers"]["github"]["command"] == "npx"
|
|
181
|
+
assert "filesystem" in updated_data["mcpServers"]
|
|
182
|
+
|
|
183
|
+
# Citadel Predict added/updated
|
|
184
|
+
assert "citadel-predict" in updated_data["mcpServers"]
|
|
185
|
+
assert updated_data["mcpServers"]["citadel-predict"]["env"]["CITADEL_API_KEY"] == "cp_live_citadel"
|
|
186
|
+
assert updated_data["mcpServers"]["citadel-predict"]["env"]["CITADEL_API_URL"] == "https://citadel.example.com"
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
def test_configure_mcp_cli_subprocess(tmp_path):
|
|
190
|
+
"""Verify running the CLI as a subprocess works end-to-end with flags."""
|
|
191
|
+
desktop_target = tmp_path / "cli_desktop.json"
|
|
192
|
+
code_target = tmp_path / "cli_code.json"
|
|
193
|
+
|
|
194
|
+
script_path = REPO_ROOT / "scripts" / "configure_mcp.py"
|
|
195
|
+
|
|
196
|
+
cmd = [
|
|
197
|
+
sys.executable,
|
|
198
|
+
str(script_path),
|
|
199
|
+
"--non-interactive",
|
|
200
|
+
"--api-key",
|
|
201
|
+
"cp_live_cli_test",
|
|
202
|
+
"--api-url",
|
|
203
|
+
"http://127.0.0.1:8000",
|
|
204
|
+
"--desktop-config",
|
|
205
|
+
str(desktop_target),
|
|
206
|
+
"--code-config",
|
|
207
|
+
str(code_target),
|
|
208
|
+
]
|
|
209
|
+
|
|
210
|
+
proc = subprocess.run(cmd, capture_output=True, text=True)
|
|
211
|
+
assert proc.returncode == 0
|
|
212
|
+
assert "Configuration completed successfully." in proc.stdout
|
|
213
|
+
|
|
214
|
+
data = json.loads(desktop_target.read_text(encoding="utf-8"))
|
|
215
|
+
assert data["mcpServers"]["citadel-predict"]["env"]["CITADEL_API_KEY"] == "cp_live_cli_test"
|
|
@@ -0,0 +1,250 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Unit tests for the Citadel Predict MCP server.
|
|
3
|
+
"""
|
|
4
|
+
|
|
5
|
+
import asyncio
|
|
6
|
+
import json
|
|
7
|
+
from unittest.mock import MagicMock, patch
|
|
8
|
+
|
|
9
|
+
import pytest
|
|
10
|
+
from citadel_predict.errors import (
|
|
11
|
+
CitadelAuthError,
|
|
12
|
+
CitadelBadRequestError,
|
|
13
|
+
CitadelError,
|
|
14
|
+
CitadelNetworkError,
|
|
15
|
+
CitadelRateLimitError,
|
|
16
|
+
CitadelServerError,
|
|
17
|
+
CitadelValidationError,
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
from citadel_predict_mcp.server import create_server, main
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@pytest.fixture
|
|
24
|
+
def mcp_server():
|
|
25
|
+
return create_server()
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def test_tool_registration(mcp_server):
|
|
29
|
+
async def _test():
|
|
30
|
+
tools = await mcp_server.list_tools()
|
|
31
|
+
tool_names = [t.name for t in tools]
|
|
32
|
+
assert "estimate_agent_cost" in tool_names
|
|
33
|
+
|
|
34
|
+
tool_def = next(t for t in tools if t.name == "estimate_agent_cost")
|
|
35
|
+
assert "BEFORE running it" in tool_def.description
|
|
36
|
+
assert "token cost and usage range" in tool_def.description
|
|
37
|
+
|
|
38
|
+
asyncio.run(_test())
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def test_estimate_agent_cost_success(mcp_server):
|
|
42
|
+
async def _test():
|
|
43
|
+
mock_prediction = {
|
|
44
|
+
"model_id": "claude-sonnet",
|
|
45
|
+
"expected_tokens": 3200,
|
|
46
|
+
"low_tokens": 1500,
|
|
47
|
+
"high_tokens": 5800,
|
|
48
|
+
"out_of_distribution": False,
|
|
49
|
+
"ood_reasons": [],
|
|
50
|
+
"confidence": "high",
|
|
51
|
+
"driving_factors": ["task_complexity", "tools_count"],
|
|
52
|
+
}
|
|
53
|
+
|
|
54
|
+
with patch("citadel_predict_mcp.server.predict_cost", return_value=mock_prediction) as mock_predict:
|
|
55
|
+
result = await mcp_server.call_tool(
|
|
56
|
+
"estimate_agent_cost",
|
|
57
|
+
{
|
|
58
|
+
"task_text": "Audit repository and draft report",
|
|
59
|
+
"tools": ["list_files", "draft_document"],
|
|
60
|
+
"num_tools": 2,
|
|
61
|
+
"model_id": "claude-sonnet",
|
|
62
|
+
},
|
|
63
|
+
)
|
|
64
|
+
|
|
65
|
+
assert not result.is_error
|
|
66
|
+
content_text = result.content[0].text
|
|
67
|
+
data = json.loads(content_text)
|
|
68
|
+
|
|
69
|
+
assert data["success"] is True
|
|
70
|
+
assert data["expected_tokens"] == 3200
|
|
71
|
+
assert data["low_tokens"] == 1500
|
|
72
|
+
assert data["high_tokens"] == 5800
|
|
73
|
+
assert data["out_of_distribution"] is False
|
|
74
|
+
assert "Expected: 3,200 tokens" in data["summary"]
|
|
75
|
+
|
|
76
|
+
mock_predict.assert_called_once_with(
|
|
77
|
+
task_text="Audit repository and draft report",
|
|
78
|
+
tools=["list_files", "draft_document"],
|
|
79
|
+
num_tools=2,
|
|
80
|
+
model_id="claude-sonnet",
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
asyncio.run(_test())
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def test_auth_error_translation(mcp_server):
|
|
87
|
+
async def _test():
|
|
88
|
+
with patch(
|
|
89
|
+
"citadel_predict_mcp.server.predict_cost",
|
|
90
|
+
side_effect=CitadelAuthError("Invalid API key"),
|
|
91
|
+
):
|
|
92
|
+
result = await mcp_server.call_tool(
|
|
93
|
+
"estimate_agent_cost",
|
|
94
|
+
{"task_text": "Analyze data"},
|
|
95
|
+
)
|
|
96
|
+
data = json.loads(result.content[0].text)
|
|
97
|
+
assert data["success"] is False
|
|
98
|
+
assert data["error_type"] == "AuthenticationError"
|
|
99
|
+
assert "CITADEL_API_KEY" in data["message"]
|
|
100
|
+
|
|
101
|
+
asyncio.run(_test())
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def test_rate_limit_error_translation(mcp_server):
|
|
105
|
+
async def _test():
|
|
106
|
+
with patch(
|
|
107
|
+
"citadel_predict_mcp.server.predict_cost",
|
|
108
|
+
side_effect=CitadelRateLimitError("Rate limited", retry_after=45),
|
|
109
|
+
):
|
|
110
|
+
result = await mcp_server.call_tool(
|
|
111
|
+
"estimate_agent_cost",
|
|
112
|
+
{"task_text": "Analyze data"},
|
|
113
|
+
)
|
|
114
|
+
data = json.loads(result.content[0].text)
|
|
115
|
+
assert data["success"] is False
|
|
116
|
+
assert data["error_type"] == "RateLimitError"
|
|
117
|
+
assert data["retry_after"] == 45
|
|
118
|
+
assert "retry after 45s" in data["message"]
|
|
119
|
+
|
|
120
|
+
asyncio.run(_test())
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
def test_validation_error_translation(mcp_server):
|
|
124
|
+
async def _test():
|
|
125
|
+
with patch(
|
|
126
|
+
"citadel_predict_mcp.server.predict_cost",
|
|
127
|
+
side_effect=CitadelValidationError("task_text exceeds maximum length"),
|
|
128
|
+
):
|
|
129
|
+
result = await mcp_server.call_tool(
|
|
130
|
+
"estimate_agent_cost",
|
|
131
|
+
{"task_text": "A" * 5000},
|
|
132
|
+
)
|
|
133
|
+
data = json.loads(result.content[0].text)
|
|
134
|
+
assert data["success"] is False
|
|
135
|
+
assert data["error_type"] == "ValidationError"
|
|
136
|
+
assert "Validation failed" in data["message"]
|
|
137
|
+
|
|
138
|
+
asyncio.run(_test())
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
def test_bad_request_error_translation(mcp_server):
|
|
142
|
+
async def _test():
|
|
143
|
+
with patch(
|
|
144
|
+
"citadel_predict_mcp.server.predict_cost",
|
|
145
|
+
side_effect=CitadelBadRequestError("Unsupported model id"),
|
|
146
|
+
):
|
|
147
|
+
result = await mcp_server.call_tool(
|
|
148
|
+
"estimate_agent_cost",
|
|
149
|
+
{"task_text": "Task", "model_id": "unsupported-model"},
|
|
150
|
+
)
|
|
151
|
+
data = json.loads(result.content[0].text)
|
|
152
|
+
assert data["success"] is False
|
|
153
|
+
assert data["error_type"] == "BadRequestError"
|
|
154
|
+
assert "Unsupported model id" in data["message"]
|
|
155
|
+
|
|
156
|
+
asyncio.run(_test())
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
def test_server_error_translation(mcp_server):
|
|
160
|
+
async def _test():
|
|
161
|
+
with patch(
|
|
162
|
+
"citadel_predict_mcp.server.predict_cost",
|
|
163
|
+
side_effect=CitadelServerError("Internal database failure", status_code=500),
|
|
164
|
+
):
|
|
165
|
+
result = await mcp_server.call_tool(
|
|
166
|
+
"estimate_agent_cost",
|
|
167
|
+
{"task_text": "Task"},
|
|
168
|
+
)
|
|
169
|
+
data = json.loads(result.content[0].text)
|
|
170
|
+
assert data["success"] is False
|
|
171
|
+
assert data["error_type"] == "ServerError"
|
|
172
|
+
assert "server error (500)" in data["message"]
|
|
173
|
+
|
|
174
|
+
asyncio.run(_test())
|
|
175
|
+
|
|
176
|
+
|
|
177
|
+
def test_network_error_translation(mcp_server):
|
|
178
|
+
async def _test():
|
|
179
|
+
with patch(
|
|
180
|
+
"citadel_predict_mcp.server.predict_cost",
|
|
181
|
+
side_effect=CitadelNetworkError("ConnectTimeout"),
|
|
182
|
+
):
|
|
183
|
+
result = await mcp_server.call_tool(
|
|
184
|
+
"estimate_agent_cost",
|
|
185
|
+
{"task_text": "Task"},
|
|
186
|
+
)
|
|
187
|
+
data = json.loads(result.content[0].text)
|
|
188
|
+
assert data["success"] is False
|
|
189
|
+
assert data["error_type"] == "NetworkError"
|
|
190
|
+
assert "Unable to reach the Citadel Predict API" in data["message"]
|
|
191
|
+
|
|
192
|
+
asyncio.run(_test())
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
def test_generic_citadel_error_translation(mcp_server):
|
|
196
|
+
async def _test():
|
|
197
|
+
with patch(
|
|
198
|
+
"citadel_predict_mcp.server.predict_cost",
|
|
199
|
+
side_effect=CitadelError("Generic client failure"),
|
|
200
|
+
):
|
|
201
|
+
result = await mcp_server.call_tool(
|
|
202
|
+
"estimate_agent_cost",
|
|
203
|
+
{"task_text": "Task"},
|
|
204
|
+
)
|
|
205
|
+
data = json.loads(result.content[0].text)
|
|
206
|
+
assert data["success"] is False
|
|
207
|
+
assert data["error_type"] == "CitadelError"
|
|
208
|
+
assert "Generic client failure" in data["message"]
|
|
209
|
+
|
|
210
|
+
asyncio.run(_test())
|
|
211
|
+
|
|
212
|
+
|
|
213
|
+
def test_unexpected_exception_handling(mcp_server):
|
|
214
|
+
async def _test():
|
|
215
|
+
with patch(
|
|
216
|
+
"citadel_predict_mcp.server.predict_cost",
|
|
217
|
+
side_effect=RuntimeError("Unexpected OS crash"),
|
|
218
|
+
):
|
|
219
|
+
result = await mcp_server.call_tool(
|
|
220
|
+
"estimate_agent_cost",
|
|
221
|
+
{"task_text": "Task"},
|
|
222
|
+
)
|
|
223
|
+
data = json.loads(result.content[0].text)
|
|
224
|
+
assert data["success"] is False
|
|
225
|
+
assert data["error_type"] == "UnexpectedError"
|
|
226
|
+
assert "Unexpected OS crash" in data["details"]
|
|
227
|
+
|
|
228
|
+
asyncio.run(_test())
|
|
229
|
+
|
|
230
|
+
|
|
231
|
+
def test_main_entrypoint():
|
|
232
|
+
with patch("sys.argv", ["citadel-predict-mcp"]), patch("citadel_predict_mcp.server.server.run") as mock_run:
|
|
233
|
+
main()
|
|
234
|
+
mock_run.assert_called_once_with(transport="stdio")
|
|
235
|
+
|
|
236
|
+
|
|
237
|
+
def test_main_help(capsys):
|
|
238
|
+
with patch("sys.argv", ["citadel-predict-mcp", "--help"]):
|
|
239
|
+
main()
|
|
240
|
+
captured = capsys.readouterr()
|
|
241
|
+
assert "Citadel Predict MCP Server" in captured.out
|
|
242
|
+
assert "claude_desktop_config.json" in captured.out
|
|
243
|
+
|
|
244
|
+
|
|
245
|
+
def test_main_version(capsys):
|
|
246
|
+
with patch("sys.argv", ["citadel-predict-mcp", "--version"]):
|
|
247
|
+
main()
|
|
248
|
+
captured = capsys.readouterr()
|
|
249
|
+
assert "citadel-predict-mcp v0.1.0" in captured.out
|
|
250
|
+
|