Skip to content

Commit 85c373a

Browse files
subpopsubpop
committed
fix: Restore GATEKEEPER_MAX_TOKENS
Passing it down into each provider's completion function.
1 parent 1d5aee7 commit 85c373a

15 files changed

Lines changed: 72 additions & 72 deletions

‎src/linux_mcp_server/gatekeeper/anthropic_client.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515

1616
ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages"
1717
ANTHROPIC_API_VERSION = "2023-06-01"
18-
ANTHROPIC_DEFAULT_MAX_TOKENS = 4096
1918

2019

2120
def _get_anthropic_api_key() -> str:
@@ -41,9 +40,9 @@ def _anthropic_thinking_block(reasoning_effort: ReasoningEffort | None) -> dict[
4140
return {"type": "enabled", "budget_tokens": budget}
4241

4342

44-
def build_messages_body(prompt: str, *, include_model: bool) -> dict[str, Any]:
43+
def build_messages_body(prompt: str, *, include_model: bool, max_tokens: int) -> dict[str, Any]:
4544
body: dict[str, Any] = {
46-
"max_tokens": ANTHROPIC_DEFAULT_MAX_TOKENS,
45+
"max_tokens": max_tokens,
4746
"messages": [{"role": "user", "content": prompt}],
4847
"temperature": CONFIG.gatekeeper.temperature,
4948
}
@@ -66,7 +65,7 @@ def extract_messages_text(response: dict[str, Any]) -> str:
6665
return ""
6766

6867

69-
def complete_anthropic(prompt: str, *, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> GatekeeperCompletion:
68+
def complete_anthropic(prompt: str, *, max_tokens: int, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> GatekeeperCompletion:
7069
headers = {
7170
"x-api-key": _get_anthropic_api_key(),
7271
"anthropic-version": ANTHROPIC_API_VERSION,
@@ -76,7 +75,7 @@ def complete_anthropic(prompt: str, *, timeout: int = DEFAULT_TIMEOUT_SECONDS) -
7675
provider="anthropic",
7776
url=ANTHROPIC_API_URL,
7877
headers=headers,
79-
body=build_messages_body(prompt, include_model=True),
78+
body=build_messages_body(prompt, include_model=True, max_tokens=max_tokens),
8079
timeout=timeout,
8180
)
8281
return GatekeeperCompletion(text=extract_messages_text(response))

‎src/linux_mcp_server/gatekeeper/check_run_script.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,8 @@ def get_model() -> str:
103103
"""
104104

105105

106+
# Maximum number of completion tokens (including reasoning)
107+
GATEKEEPER_MAX_TOKENS = 8000
106108
# Timeout (s)
107109
GATEKEEPER_TIMEOUT = 120
108110

@@ -226,7 +228,7 @@ async def check_run_script_with_stats(
226228
time_before = time.perf_counter()
227229
try:
228230
completion = await asyncio.wait_for(
229-
asyncio.to_thread(complete_gatekeeper, prompt),
231+
asyncio.to_thread(complete_gatekeeper, prompt, max_tokens=GATEKEEPER_MAX_TOKENS),
230232
timeout=GATEKEEPER_TIMEOUT,
231233
)
232234
except asyncio.TimeoutError:

‎src/linux_mcp_server/gatekeeper/gemini_client.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -36,10 +36,11 @@ def _gemini_thinking_level(reasoning_effort: ReasoningEffort | None) -> str | No
3636
return mapping.get(reasoning_effort)
3737

3838

39-
def build_gemini_body(prompt: str) -> dict[str, Any]:
39+
def build_gemini_body(prompt: str, *, max_tokens: int) -> dict[str, Any]:
4040
generation_config = gemini_generation_config(
4141
temperature=CONFIG.gatekeeper.temperature,
4242
structured_output=CONFIG.gatekeeper.structured_output,
43+
max_tokens=max_tokens,
4344
)
4445
thinking_level = _gemini_thinking_level(CONFIG.gatekeeper.reasoning_effort)
4546
if thinking_level is not None:
@@ -62,14 +63,14 @@ def extract_gemini_text(response: dict[str, Any]) -> str:
6263
return text.strip() if isinstance(text, str) else ""
6364

6465

65-
def complete_gemini(prompt: str, *, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> GatekeeperCompletion:
66+
def complete_gemini(prompt: str, *, max_tokens: int, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> GatekeeperCompletion:
6667
model = normalize_model_id(CONFIG.gatekeeper.model or "")
6768
api_key = _get_google_api_key()
6869
response = post_json(
6970
provider="gemini",
7071
url=f"{GOOGLE_AI_BASE_URL}/models/{model}:generateContent?key={api_key}",
7172
headers={"Content-Type": "application/json"},
72-
body=build_gemini_body(prompt),
73+
body=build_gemini_body(prompt, max_tokens=max_tokens),
7374
timeout=timeout,
7475
)
7576
return GatekeeperCompletion(text=extract_gemini_text(response))

‎src/linux_mcp_server/gatekeeper/llm.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -33,19 +33,19 @@ def resolve_provider() -> GatekeeperProvider:
3333
return _infer_provider_from_model(CONFIG.gatekeeper.model)
3434

3535

36-
def complete_gatekeeper(prompt: str) -> GatekeeperCompletion:
36+
def complete_gatekeeper(prompt: str, *, max_tokens: int) -> GatekeeperCompletion:
3737
provider = resolve_provider()
3838
match provider:
3939
case GatekeeperProvider.OPENAI:
40-
completion = complete_openai(prompt)
40+
completion = complete_openai(prompt, max_tokens=max_tokens)
4141
case GatekeeperProvider.ANTHROPIC:
42-
completion = complete_anthropic(prompt)
42+
completion = complete_anthropic(prompt, max_tokens=max_tokens)
4343
case GatekeeperProvider.GEMINI:
44-
completion = complete_gemini(prompt)
44+
completion = complete_gemini(prompt, max_tokens=max_tokens)
4545
case GatekeeperProvider.OPENROUTER:
46-
completion = complete_openrouter(prompt)
46+
completion = complete_openrouter(prompt, max_tokens=max_tokens)
4747
case GatekeeperProvider.VERTEX_AI:
48-
completion = complete_vertex_ai(prompt)
48+
completion = complete_vertex_ai(prompt, max_tokens=max_tokens)
4949
case _: # pragma: no cover
5050
raise ValueError(f"Unsupported gatekeeper provider: {provider}")
5151

‎src/linux_mcp_server/gatekeeper/openai_client.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,11 @@ def _apply_chat_completions_extras(body: dict[str, Any]) -> dict[str, Any]:
6161
return body
6262

6363

64-
def _build_responses_body(prompt: str) -> dict[str, Any]:
64+
def _build_responses_body(prompt: str, *, max_tokens: int) -> dict[str, Any]:
6565
body: dict[str, Any] = {
6666
"model": normalize_model_id(CONFIG.gatekeeper.model or ""),
6767
"input": prompt,
68+
"max_output_tokens": max_tokens,
6869
"temperature": CONFIG.gatekeeper.temperature,
6970
"store": False,
7071
}
@@ -76,10 +77,11 @@ def _build_responses_body(prompt: str) -> dict[str, Any]:
7677
return body
7778

7879

79-
def build_chat_completions_body(prompt: str) -> dict[str, Any]:
80+
def build_chat_completions_body(prompt: str, *, max_tokens: int) -> dict[str, Any]:
8081
body: dict[str, Any] = {
8182
"model": normalize_model_id(CONFIG.gatekeeper.model or ""),
8283
"messages": [{"role": "user", "content": prompt}],
84+
"max_completion_tokens": max_tokens,
8385
"temperature": CONFIG.gatekeeper.temperature,
8486
}
8587
if CONFIG.gatekeeper.structured_output:
@@ -119,7 +121,7 @@ def extract_chat_completions_text(response: dict[str, Any]) -> str:
119121
return (content or "").strip() if isinstance(content, str) else ""
120122

121123

122-
def complete_openai(prompt: str, *, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> GatekeeperCompletion:
124+
def complete_openai(prompt: str, *, max_tokens: int, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> GatekeeperCompletion:
123125
base_url = _get_openai_base_url()
124126
headers = {
125127
**_openai_auth_headers(),
@@ -133,7 +135,7 @@ def complete_openai(prompt: str, *, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> G
133135
provider="openai",
134136
url=f"{base_url}/responses",
135137
headers=headers,
136-
body=_build_responses_body(prompt),
138+
body=_build_responses_body(prompt, max_tokens=max_tokens),
137139
timeout=timeout,
138140
)
139141
return GatekeeperCompletion(text=_extract_responses_text(response))
@@ -145,7 +147,7 @@ def complete_openai(prompt: str, *, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> G
145147
provider="openai",
146148
url=f"{base_url}/chat/completions",
147149
headers=headers,
148-
body=build_chat_completions_body(prompt),
150+
body=build_chat_completions_body(prompt, max_tokens=max_tokens),
149151
timeout=timeout,
150152
)
151153
return GatekeeperCompletion(text=extract_chat_completions_text(response))

‎src/linux_mcp_server/gatekeeper/openrouter_client.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ def _openrouter_config() -> dict[str, Any]:
5454
}
5555

5656

57-
def _build_chat_completions_body(prompt: str) -> dict[str, Any]:
57+
def _build_chat_completions_body(prompt: str, *, max_tokens: int) -> dict[str, Any]:
5858
openrouter = _openrouter_config()
5959
provider: dict[str, Any] = {"require_parameters": True}
6060
if openrouter["quantization"]:
@@ -63,6 +63,7 @@ def _build_chat_completions_body(prompt: str) -> dict[str, Any]:
6363
body: dict[str, Any] = {
6464
"model": _normalize_openrouter_model_id(CONFIG.gatekeeper.model or ""),
6565
"messages": [{"role": "user", "content": prompt}],
66+
"max_tokens": max_tokens,
6667
"temperature": CONFIG.gatekeeper.temperature,
6768
"provider": provider,
6869
}
@@ -99,7 +100,9 @@ def _extract_usage(response: dict[str, Any]) -> tuple[int, int, float | None]:
99100
)
100101

101102

102-
def complete_openrouter(prompt: str, *, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> GatekeeperCompletion:
103+
def complete_openrouter(
104+
prompt: str, *, max_tokens: int, timeout: int = DEFAULT_TIMEOUT_SECONDS
105+
) -> GatekeeperCompletion:
103106
base_url = _get_openrouter_base_url()
104107
headers = {
105108
**_openrouter_auth_headers(),
@@ -109,7 +112,7 @@ def complete_openrouter(prompt: str, *, timeout: int = DEFAULT_TIMEOUT_SECONDS)
109112
provider="openrouter",
110113
url=f"{base_url}/chat/completions",
111114
headers=headers,
112-
body=_build_chat_completions_body(prompt),
115+
body=_build_chat_completions_body(prompt, max_tokens=max_tokens),
113116
timeout=timeout,
114117
)
115118
prompt_tokens, completion_tokens, usage_cost = _extract_usage(response)

‎src/linux_mcp_server/gatekeeper/schema.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,8 +99,8 @@ def anthropic_output_config() -> dict[str, Any]:
9999
}
100100

101101

102-
def gemini_generation_config(*, temperature: float, structured_output: bool) -> dict[str, Any]:
103-
config: dict[str, Any] = {"temperature": temperature}
102+
def gemini_generation_config(*, temperature: float, structured_output: bool, max_tokens: int) -> dict[str, Any]:
103+
config: dict[str, Any] = {"temperature": temperature, "maxOutputTokens": max_tokens}
104104
if structured_output:
105105
config["responseMimeType"] = "application/json"
106106
config["responseSchema"] = gemini_json_schema()

‎src/linux_mcp_server/gatekeeper/vertex_ai_client.py‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -61,9 +61,9 @@ def _gemini_vertex_url(model: str) -> str:
6161
return f"https://{host}/v1/projects/{project}/locations/{location}/publishers/google/models/{model}:generateContent"
6262

6363

64-
def _complete_anthropic_on_vertex(prompt: str, *, timeout: int) -> GatekeeperCompletion:
64+
def _complete_anthropic_on_vertex(prompt: str, *, max_tokens: int, timeout: int) -> GatekeeperCompletion:
6565
model = normalize_model_id(CONFIG.gatekeeper.model or "")
66-
body = build_messages_body(prompt, include_model=False)
66+
body = build_messages_body(prompt, include_model=False, max_tokens=max_tokens)
6767
body["anthropic_version"] = ANTHROPIC_VERTEX_VERSION
6868
response = post_json(
6969
provider="anthropic",
@@ -75,36 +75,36 @@ def _complete_anthropic_on_vertex(prompt: str, *, timeout: int) -> GatekeeperCom
7575
return GatekeeperCompletion(text=extract_messages_text(response))
7676

7777

78-
def _complete_gemini_on_vertex(prompt: str, *, timeout: int) -> GatekeeperCompletion:
78+
def _complete_gemini_on_vertex(prompt: str, *, max_tokens: int, timeout: int) -> GatekeeperCompletion:
7979
model = normalize_model_id(CONFIG.gatekeeper.model or "")
8080
response = post_json(
8181
provider="gemini",
8282
url=_gemini_vertex_url(model),
8383
headers={**_vertex_auth_headers(), "Content-Type": "application/json"},
84-
body=build_gemini_body(prompt),
84+
body=build_gemini_body(prompt, max_tokens=max_tokens),
8585
timeout=timeout,
8686
)
8787
return GatekeeperCompletion(text=extract_gemini_text(response))
8888

8989

90-
def _complete_openai_compatible_on_vertex(prompt: str, *, timeout: int) -> GatekeeperCompletion:
90+
def _complete_openai_compatible_on_vertex(prompt: str, *, max_tokens: int, timeout: int) -> GatekeeperCompletion:
9191
base_url = _get_vertex_openapi_base_url()
9292
response = post_json(
9393
provider="openai",
9494
url=f"{base_url}/chat/completions",
9595
headers={**_vertex_auth_headers(), "Content-Type": "application/json"},
96-
body=build_chat_completions_body(prompt),
96+
body=build_chat_completions_body(prompt, max_tokens=max_tokens),
9797
timeout=timeout,
9898
)
9999
return GatekeeperCompletion(text=extract_chat_completions_text(response))
100100

101101

102-
def complete_vertex_ai(prompt: str, *, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> GatekeeperCompletion:
102+
def complete_vertex_ai(prompt: str, *, max_tokens: int, timeout: int = DEFAULT_TIMEOUT_SECONDS) -> GatekeeperCompletion:
103103
model = CONFIG.gatekeeper.model or ""
104104
match _vertex_api_style(model):
105105
case "anthropic":
106-
return _complete_anthropic_on_vertex(prompt, timeout=timeout)
106+
return _complete_anthropic_on_vertex(prompt, max_tokens=max_tokens, timeout=timeout)
107107
case "gemini":
108-
return _complete_gemini_on_vertex(prompt, timeout=timeout)
108+
return _complete_gemini_on_vertex(prompt, max_tokens=max_tokens, timeout=timeout)
109109
case "openai_compatible":
110-
return _complete_openai_compatible_on_vertex(prompt, timeout=timeout)
110+
return _complete_openai_compatible_on_vertex(prompt, max_tokens=max_tokens, timeout=timeout)

‎tests/gatekeeper/test_anthropic_client.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ def test_complete_anthropic_direct(self, gatekeeper_config, mocker):
3232
return_value={"content": [{"type": "text", "text": '{"status": "OK", "detail": ""}'}]},
3333
)
3434

35-
result = anthropic_client.complete_anthropic("prompt")
35+
result = anthropic_client.complete_anthropic("prompt", max_tokens=8000)
3636

3737
assert result.text == '{"status": "OK", "detail": ""}'
3838
assert mock_post.call_args.kwargs["url"] == "https://api.anthropic.com/v1/messages"

‎tests/gatekeeper/test_check_run_script.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ def _completion(text: str) -> GatekeeperCompletion:
8585
return mocker.patch.object(
8686
check_run_script_module,
8787
"complete_gatekeeper",
88-
side_effect=lambda prompt: _completion('{"status": "OK", "detail": ""}'),
88+
side_effect=lambda prompt, **kwargs: _completion('{"status": "OK", "detail": ""}'),
8989
)
9090

9191
async def test_rejects_script_with_prompt_injection_attempts(self):
@@ -128,7 +128,7 @@ async def test_parse_errors(self, mocker, response_text):
128128
await check_run_script(description="test", script_type="bash", script="echo hi", readonly=True)
129129

130130
async def test_timeout(self, mocker):
131-
def slow_complete(_prompt: str) -> GatekeeperCompletion:
131+
def slow_complete(_prompt: str, **kwargs: object) -> GatekeeperCompletion:
132132
time.sleep(10)
133133
return GatekeeperCompletion(text='{"status": "OK"}')
134134

0 commit comments

Comments
 (0)