422 lines
13 KiB
Python
422 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from tht.internal.litellm_completion import handle_request
|
|
|
|
HARNESS_ROOT = Path(__file__).resolve().parents[1]
|
|
HELPER_MODULE = "tht.internal.litellm_completion"
|
|
|
|
|
|
def _run_raw_helper(
|
|
tmp_path: Path, raw_request: str, fake_litellm: str
|
|
) -> subprocess.CompletedProcess[str]:
|
|
provider_path = tmp_path / "provider"
|
|
provider_path.mkdir()
|
|
(provider_path / "litellm.py").write_text(fake_litellm, encoding="utf-8")
|
|
|
|
env = os.environ.copy()
|
|
env["PYTHONPATH"] = os.pathsep.join((str(provider_path), str(HARNESS_ROOT)))
|
|
return subprocess.run(
|
|
[sys.executable, "-m", HELPER_MODULE],
|
|
input=raw_request,
|
|
text=True,
|
|
capture_output=True,
|
|
cwd=HARNESS_ROOT,
|
|
env=env,
|
|
timeout=10,
|
|
check=False,
|
|
)
|
|
|
|
|
|
def _run_helper(tmp_path: Path, request: object, fake_litellm: str) -> subprocess.CompletedProcess[str]:
|
|
return _run_raw_helper(tmp_path, json.dumps(request), fake_litellm)
|
|
|
|
|
|
def _valid_request() -> dict[str, object]:
|
|
return {
|
|
"model": "openai/gpt-4.1-mini",
|
|
"api_key": "sk-secret",
|
|
"messages": [{"role": "user", "content": "private prompt"}],
|
|
}
|
|
|
|
|
|
def test_subprocess_emits_one_pristine_success_line_and_silences_provider_output(tmp_path: Path):
|
|
secret = "sk-contract-secret"
|
|
prompt = "private patient prompt"
|
|
request = {
|
|
"model": "openai/gpt-4.1-mini",
|
|
"api_key": secret,
|
|
"messages": [
|
|
{"role": "system", "content": "Return one description."},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
"api_base": "https://models.example.test/v1",
|
|
"api_version": "2026-08-01-preview",
|
|
}
|
|
fake_litellm = f'''
|
|
import sys
|
|
|
|
suppress_debug_info = False
|
|
set_verbose = True
|
|
turn_off_message_logging = False
|
|
log_raw_request_response = True
|
|
|
|
print("provider import noise")
|
|
print("provider import diagnostic", file=sys.stderr)
|
|
|
|
def completion(**kwargs):
|
|
assert suppress_debug_info is True
|
|
assert set_verbose is False
|
|
assert turn_off_message_logging is True
|
|
assert log_raw_request_response is False
|
|
assert kwargs == {{
|
|
"model": "openai/gpt-4.1-mini",
|
|
"api_key": {secret!r},
|
|
"messages": [
|
|
{{"role": "system", "content": "Return one description."}},
|
|
{{"role": "user", "content": {prompt!r}}},
|
|
],
|
|
"api_base": "https://models.example.test/v1",
|
|
"api_version": "2026-08-01-preview",
|
|
"num_retries": 1,
|
|
"stream": False,
|
|
}}
|
|
print(repr(kwargs))
|
|
print(repr(kwargs), file=sys.stderr)
|
|
return {{"choices": [{{"message": {{"content": "Descrizione italiana"}}}}]}}
|
|
'''
|
|
|
|
result = _run_helper(tmp_path, request, fake_litellm)
|
|
|
|
assert result.returncode == 0
|
|
assert result.stdout == '{"ok":true,"content":"Descrizione italiana"}\n'
|
|
assert result.stderr == ""
|
|
assert secret not in result.stdout + result.stderr
|
|
assert prompt not in result.stdout + result.stderr
|
|
|
|
|
|
def test_subprocess_rejects_an_unknown_request_key_before_loading_provider(tmp_path: Path):
|
|
request = {**_valid_request(), "unexpected": True}
|
|
|
|
result = _run_helper(tmp_path, request, 'raise AssertionError("provider was loaded")')
|
|
|
|
assert result.returncode == 0
|
|
assert result.stdout == '{"ok":false,"error":"invalid_request"}\n'
|
|
assert result.stderr == "completion helper: invalid request\n"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"raw_request",
|
|
[
|
|
"{not-json",
|
|
"[]",
|
|
json.dumps({"model": "openai/gpt-4.1-mini", "api_key": "sk-secret"}),
|
|
json.dumps({**_valid_request(), "model": "bad model"}),
|
|
json.dumps({**_valid_request(), "model": "a" * 257}),
|
|
json.dumps({**_valid_request(), "api_key": " "}),
|
|
json.dumps({**_valid_request(), "api_key": "x" * (16 * 1024 + 1)}),
|
|
json.dumps(
|
|
{
|
|
**_valid_request(),
|
|
"messages": [{"role": "assistant", "content": "not allowed"}],
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
**_valid_request(),
|
|
"messages": [{"role": "user", "content": "ok", "name": "extra"}],
|
|
}
|
|
),
|
|
json.dumps({**_valid_request(), "messages": []}),
|
|
json.dumps(
|
|
{
|
|
**_valid_request(),
|
|
"messages": [{"role": "user", "content": "ok"}] * 33,
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
**_valid_request(),
|
|
"messages": [{"role": "user"}],
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
**_valid_request(),
|
|
"messages": [{"role": "user", "content": 1}],
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
**_valid_request(),
|
|
"messages": [{"role": "user", "content": "x" * (64 * 1024 + 1)}],
|
|
}
|
|
),
|
|
json.dumps(
|
|
{
|
|
**_valid_request(),
|
|
"messages": [
|
|
{"role": "user", "content": "x" * (48 * 1024)},
|
|
{"role": "system", "content": "y" * (48 * 1024)},
|
|
{"role": "user", "content": "z" * (48 * 1024)},
|
|
],
|
|
}
|
|
),
|
|
json.dumps({**_valid_request(), "api_base": "file:///tmp/provider"}),
|
|
json.dumps({**_valid_request(), "api_base": "https://user@example.test/v1"}),
|
|
json.dumps({**_valid_request(), "api_base": "https://example.test/v1?key=value"}),
|
|
json.dumps({**_valid_request(), "api_version": "2026 preview"}),
|
|
json.dumps({**_valid_request(), "disable_thinking": False}),
|
|
(
|
|
'{"model":"openai/first","model":"openai/second","api_key":"sk-secret",'
|
|
'"messages":[{"role":"user","content":"private prompt"}]}'
|
|
),
|
|
(
|
|
'{"model":"openai/gpt-4.1-mini","api_key":"sk-secret",'
|
|
'"messages":[{"role":"user","content":"private prompt"}],"number":'
|
|
+ "1" * 5000
|
|
+ "}"
|
|
),
|
|
json.dumps(
|
|
{
|
|
**_valid_request(),
|
|
"messages": [{"role": "user", "content": "x" * (256 * 1024)}],
|
|
}
|
|
),
|
|
],
|
|
ids=[
|
|
"malformed-json",
|
|
"non-object",
|
|
"missing-required-key",
|
|
"invalid-model",
|
|
"long-model",
|
|
"empty-secret",
|
|
"long-secret",
|
|
"invalid-role",
|
|
"message-extra-key",
|
|
"empty-messages",
|
|
"too-many-messages",
|
|
"message-missing-key",
|
|
"non-string-content",
|
|
"long-message",
|
|
"aggregate-content-too-long",
|
|
"invalid-url-scheme",
|
|
"url-credentials",
|
|
"url-query",
|
|
"invalid-api-version",
|
|
"invalid-disable-thinking",
|
|
"duplicate-key",
|
|
"oversized-number",
|
|
"oversized-stdin",
|
|
],
|
|
)
|
|
def test_subprocess_rejects_malformed_or_oversized_requests_without_loading_provider(
|
|
tmp_path: Path, raw_request: str
|
|
):
|
|
result = _run_raw_helper(
|
|
tmp_path,
|
|
raw_request,
|
|
'raise AssertionError("provider was loaded with private prompt and sk-secret")',
|
|
)
|
|
|
|
assert result.returncode == 0
|
|
assert result.stdout == '{"ok":false,"error":"invalid_request"}\n'
|
|
assert result.stderr == "completion helper: invalid request\n"
|
|
assert "private prompt" not in result.stdout + result.stderr
|
|
assert "sk-secret" not in result.stdout + result.stderr
|
|
|
|
|
|
def test_subprocess_normalizes_provider_failure_without_leaking_provider_output(tmp_path: Path):
|
|
secret = "sk-provider-canary"
|
|
prompt = "provider prompt canary"
|
|
request = {
|
|
**_valid_request(),
|
|
"api_key": secret,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
}
|
|
fake_litellm = '''
|
|
import os
|
|
import sys
|
|
|
|
def completion(**kwargs):
|
|
leaked = repr(kwargs).encode("utf-8")
|
|
print(repr(kwargs))
|
|
print(repr(kwargs), file=sys.stderr)
|
|
os.write(1, leaked + b"\\n")
|
|
os.write(2, leaked + b"\\n")
|
|
os.write(1, b"x" * (1024 * 1024))
|
|
os.write(2, b"x" * (1024 * 1024))
|
|
raise RuntimeError(repr(kwargs))
|
|
'''
|
|
|
|
result = _run_helper(tmp_path, request, fake_litellm)
|
|
|
|
assert result.returncode == 0
|
|
assert result.stdout == '{"ok":false,"error":"provider_failure"}\n'
|
|
assert result.stderr == "completion helper: provider failure\n"
|
|
assert len(result.stderr.encode("utf-8")) <= 256
|
|
assert secret not in result.stdout + result.stderr
|
|
assert prompt not in result.stdout + result.stderr
|
|
|
|
|
|
def test_subprocess_normalizes_an_invalid_response_to_one_safe_line(tmp_path: Path):
|
|
secret = "sk-invalid-response"
|
|
prompt = "invalid response prompt"
|
|
request = {
|
|
**_valid_request(),
|
|
"api_key": secret,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
}
|
|
fake_litellm = '''
|
|
import os
|
|
|
|
def completion(**kwargs):
|
|
os.write(1, repr(kwargs).encode("utf-8"))
|
|
os.write(2, repr(kwargs).encode("utf-8"))
|
|
return {"choices": [{"message": {"content": ""}}]}
|
|
'''
|
|
|
|
result = _run_helper(tmp_path, request, fake_litellm)
|
|
|
|
assert result.returncode == 0
|
|
assert result.stdout == '{"ok":false,"error":"invalid_response"}\n'
|
|
assert result.stderr == "completion helper: invalid response\n"
|
|
assert secret not in result.stdout + result.stderr
|
|
assert prompt not in result.stdout + result.stderr
|
|
|
|
|
|
def test_subprocess_escapes_multiline_content_without_adding_output_frames(tmp_path: Path):
|
|
content = 'Prima riga\nSeconda "riga" — fine'
|
|
fake_litellm = f'''
|
|
def completion(**kwargs):
|
|
return {{"choices": [{{"message": {{"content": {content!r}}}}}]}}
|
|
'''
|
|
|
|
result = _run_helper(tmp_path, _valid_request(), fake_litellm)
|
|
|
|
assert result.returncode == 0
|
|
assert result.stdout == json.dumps(
|
|
{"ok": True, "content": content}, ensure_ascii=False, separators=(",", ":")
|
|
) + "\n"
|
|
assert result.stdout.count("\n") == 1
|
|
assert result.stderr == ""
|
|
|
|
|
|
def test_injected_completion_configures_one_retry_without_fallback_or_secret_export():
|
|
secret = "sk-injected-canary"
|
|
request = {
|
|
**_valid_request(),
|
|
"api_key": secret,
|
|
"messages": [
|
|
{"role": "system", "content": "System instruction"},
|
|
{"role": "user", "content": "Private prompt"},
|
|
],
|
|
}
|
|
calls: list[dict[str, object]] = []
|
|
|
|
def completion(**kwargs: object) -> object:
|
|
assert secret not in os.environ.values()
|
|
calls.append(kwargs)
|
|
return SimpleNamespace(
|
|
choices=[SimpleNamespace(message=SimpleNamespace(content="Validated content"))]
|
|
)
|
|
|
|
result = handle_request(request, completion=completion)
|
|
|
|
assert result == {"ok": True, "content": "Validated content"}
|
|
assert calls == [{
|
|
"model": "openai/gpt-4.1-mini",
|
|
"api_key": secret,
|
|
"messages": [
|
|
{"role": "system", "content": "System instruction"},
|
|
{"role": "user", "content": "Private prompt"},
|
|
],
|
|
"num_retries": 1,
|
|
"stream": False,
|
|
}]
|
|
|
|
|
|
def test_injected_completion_uses_non_secret_client_placeholder_for_keyless_endpoint():
|
|
calls: list[dict[str, object]] = []
|
|
|
|
def completion(**kwargs: object) -> object:
|
|
calls.append(kwargs)
|
|
return {"choices": [{"message": {"content": "Descrizione Qwen"}}]}
|
|
|
|
request = {
|
|
"model": "openai/qwen3.6-35b-a3b",
|
|
"messages": [{"role": "user", "content": "Invented metadata"}],
|
|
"api_base": "https://models.internal.example/v1",
|
|
"disable_thinking": True,
|
|
}
|
|
|
|
assert handle_request(request, completion=completion) == {
|
|
"ok": True,
|
|
"content": "Descrizione Qwen",
|
|
}
|
|
assert calls == [{
|
|
"model": "openai/qwen3.6-35b-a3b",
|
|
"api_key": "not-required",
|
|
"messages": [{"role": "user", "content": "Invented metadata"}],
|
|
"num_retries": 1,
|
|
"stream": False,
|
|
"api_base": "https://models.internal.example/v1",
|
|
"extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
|
|
}]
|
|
|
|
|
|
def test_injected_provider_failure_is_not_retried():
|
|
attempts = 0
|
|
|
|
def completion(**kwargs: object) -> object:
|
|
nonlocal attempts
|
|
attempts += 1
|
|
raise RuntimeError(repr(kwargs))
|
|
|
|
result = handle_request(_valid_request(), completion=completion)
|
|
|
|
assert result == {"ok": False, "error": "provider_failure"}
|
|
assert attempts == 1
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"response",
|
|
[
|
|
None,
|
|
{},
|
|
{"choices": []},
|
|
{
|
|
"choices": [
|
|
{"message": {"content": "first"}},
|
|
{"message": {"content": "second"}},
|
|
]
|
|
},
|
|
{"choices": [{"message": {}}]},
|
|
{"choices": [{"message": {"content": None}}]},
|
|
{"choices": [{"message": {"content": " "}}]},
|
|
{"choices": [{"message": {"content": "x" * (64 * 1024 + 1)}}]},
|
|
],
|
|
ids=[
|
|
"none",
|
|
"missing-choices",
|
|
"empty-choices",
|
|
"ambiguous-choices",
|
|
"missing-content",
|
|
"non-string-content",
|
|
"empty-content",
|
|
"oversized-content",
|
|
],
|
|
)
|
|
def test_injected_completion_rejects_invalid_response_shapes(response: object):
|
|
result = handle_request(_valid_request(), completion=lambda **_: response)
|
|
|
|
assert result == {"ok": False, "error": "invalid_response"}
|