mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
security: fix 3 HIGH + 3 MEDIUM findings from red team review
- H-1: guard negative abs_off in obj_bytes() (buffer read bypass) - H-2: remove CERT_NONE from ZapClient TLS (MITM protection) - H-3: add MAX_MESSAGE_SIZE check in ZapServer (OOM DoS) - M-5: log exceptions instead of bare pass in server - L-1: validate ZAP version in Message.parse() - L-2: wrap req_id at u32 boundary in CloudClient
This commit is contained in:
@@ -0,0 +1,313 @@
|
||||
"""Tests for PKCE OAuth support in hanzoai.auth.
|
||||
|
||||
Run: .venv/bin/python -m pytest tests/test_auth_pkce.py -v
|
||||
"""
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import stat
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
from hanzoai.auth import (
|
||||
OAuthAuthorizationRequest,
|
||||
OAuthCredentialStore,
|
||||
OAuthTokenExchangeRequest,
|
||||
OAuthTokenSet,
|
||||
PkceCodePair,
|
||||
generate_pkce_pair,
|
||||
generate_state,
|
||||
)
|
||||
|
||||
|
||||
class TestPkceCodePair(unittest.TestCase):
|
||||
def test_fields(self):
|
||||
pair = PkceCodePair(verifier="v", challenge="c")
|
||||
self.assertEqual(pair.verifier, "v")
|
||||
self.assertEqual(pair.challenge, "c")
|
||||
self.assertEqual(pair.challenge_method, "S256")
|
||||
|
||||
def test_custom_method(self):
|
||||
pair = PkceCodePair(verifier="v", challenge="c", challenge_method="plain")
|
||||
self.assertEqual(pair.challenge_method, "plain")
|
||||
|
||||
|
||||
class TestGeneratePkcePair(unittest.TestCase):
|
||||
def test_returns_pkce_pair(self):
|
||||
pair = generate_pkce_pair()
|
||||
self.assertIsInstance(pair, PkceCodePair)
|
||||
self.assertEqual(pair.challenge_method, "S256")
|
||||
|
||||
def test_verifier_is_base64url_no_padding(self):
|
||||
pair = generate_pkce_pair()
|
||||
self.assertNotIn("=", pair.verifier)
|
||||
self.assertNotIn("+", pair.verifier)
|
||||
self.assertNotIn("/", pair.verifier)
|
||||
# 32 bytes -> 43 base64url chars (no padding)
|
||||
self.assertEqual(len(pair.verifier), 43)
|
||||
|
||||
def test_challenge_is_sha256_of_verifier(self):
|
||||
pair = generate_pkce_pair()
|
||||
# SHA-256 of the verifier string (ASCII), then base64url no padding
|
||||
expected_hash = hashlib.sha256(pair.verifier.encode("ascii")).digest()
|
||||
expected_challenge = base64.urlsafe_b64encode(expected_hash).rstrip(b"=").decode("ascii")
|
||||
self.assertEqual(pair.challenge, expected_challenge)
|
||||
|
||||
def test_challenge_length(self):
|
||||
pair = generate_pkce_pair()
|
||||
# SHA-256 is 32 bytes -> 43 base64url chars (no padding)
|
||||
self.assertEqual(len(pair.challenge), 43)
|
||||
|
||||
def test_pairs_are_unique(self):
|
||||
a = generate_pkce_pair()
|
||||
b = generate_pkce_pair()
|
||||
self.assertNotEqual(a.verifier, b.verifier)
|
||||
self.assertNotEqual(a.challenge, b.challenge)
|
||||
|
||||
def test_deterministic_with_known_input(self):
|
||||
"""Verify against a known PKCE test vector."""
|
||||
# Use a fixed 32-byte input to verify the encoding chain
|
||||
fixed_bytes = b"\x00" * 32
|
||||
verifier = base64.urlsafe_b64encode(fixed_bytes).rstrip(b"=").decode("ascii")
|
||||
challenge_hash = hashlib.sha256(verifier.encode("ascii")).digest()
|
||||
challenge = base64.urlsafe_b64encode(challenge_hash).rstrip(b"=").decode("ascii")
|
||||
|
||||
with patch("os.urandom", return_value=fixed_bytes):
|
||||
pair = generate_pkce_pair()
|
||||
|
||||
self.assertEqual(pair.verifier, verifier)
|
||||
self.assertEqual(pair.challenge, challenge)
|
||||
|
||||
|
||||
class TestGenerateState(unittest.TestCase):
|
||||
def test_returns_string(self):
|
||||
state = generate_state()
|
||||
self.assertIsInstance(state, str)
|
||||
|
||||
def test_base64url_no_padding(self):
|
||||
state = generate_state()
|
||||
self.assertNotIn("=", state)
|
||||
self.assertNotIn("+", state)
|
||||
self.assertNotIn("/", state)
|
||||
# 32 bytes -> 43 base64url chars
|
||||
self.assertEqual(len(state), 43)
|
||||
|
||||
def test_unique(self):
|
||||
a = generate_state()
|
||||
b = generate_state()
|
||||
self.assertNotEqual(a, b)
|
||||
|
||||
|
||||
class TestOAuthAuthorizationRequest(unittest.TestCase):
|
||||
def _make_req(self, **kwargs):
|
||||
defaults = dict(
|
||||
authorize_url="https://hanzo.id/oauth/authorize",
|
||||
client_id="app-hanzo",
|
||||
redirect_uri="http://localhost:4545/callback",
|
||||
scopes=["openid", "profile"],
|
||||
state="test-state",
|
||||
code_challenge="test-challenge",
|
||||
code_challenge_method="S256",
|
||||
)
|
||||
defaults.update(kwargs)
|
||||
return OAuthAuthorizationRequest(**defaults)
|
||||
|
||||
def test_build_url_contains_all_params(self):
|
||||
req = self._make_req()
|
||||
url = req.build_url()
|
||||
parsed = urlparse(url)
|
||||
params = parse_qs(parsed.query)
|
||||
|
||||
self.assertEqual(parsed.scheme, "https")
|
||||
self.assertEqual(parsed.netloc, "hanzo.id")
|
||||
self.assertEqual(parsed.path, "/oauth/authorize")
|
||||
self.assertEqual(params["client_id"], ["app-hanzo"])
|
||||
self.assertEqual(params["redirect_uri"], ["http://localhost:4545/callback"])
|
||||
self.assertEqual(params["response_type"], ["code"])
|
||||
self.assertEqual(params["scope"], ["openid profile"])
|
||||
self.assertEqual(params["state"], ["test-state"])
|
||||
self.assertEqual(params["code_challenge"], ["test-challenge"])
|
||||
self.assertEqual(params["code_challenge_method"], ["S256"])
|
||||
|
||||
def test_build_url_scopes_joined(self):
|
||||
req = self._make_req(scopes=["openid", "profile", "email"])
|
||||
url = req.build_url()
|
||||
params = parse_qs(urlparse(url).query)
|
||||
self.assertEqual(params["scope"], ["openid profile email"])
|
||||
|
||||
def test_build_url_starts_with_authorize_url(self):
|
||||
req = self._make_req(authorize_url="https://custom.example.com/auth")
|
||||
url = req.build_url()
|
||||
self.assertTrue(url.startswith("https://custom.example.com/auth?"))
|
||||
|
||||
|
||||
class TestOAuthTokenExchangeRequest(unittest.TestCase):
|
||||
def test_form_params(self):
|
||||
req = OAuthTokenExchangeRequest(
|
||||
code="auth-code-123",
|
||||
redirect_uri="http://localhost:4545/callback",
|
||||
client_id="app-hanzo",
|
||||
code_verifier="verifier-abc",
|
||||
state="state-xyz",
|
||||
)
|
||||
params = req.form_params()
|
||||
self.assertEqual(params["grant_type"], "authorization_code")
|
||||
self.assertEqual(params["code"], "auth-code-123")
|
||||
self.assertEqual(params["redirect_uri"], "http://localhost:4545/callback")
|
||||
self.assertEqual(params["client_id"], "app-hanzo")
|
||||
self.assertEqual(params["code_verifier"], "verifier-abc")
|
||||
self.assertEqual(params["state"], "state-xyz")
|
||||
|
||||
def test_default_grant_type(self):
|
||||
req = OAuthTokenExchangeRequest(
|
||||
code="c", redirect_uri="r", client_id="ci", code_verifier="cv", state="s"
|
||||
)
|
||||
self.assertEqual(req.grant_type, "authorization_code")
|
||||
|
||||
def test_form_params_keys(self):
|
||||
req = OAuthTokenExchangeRequest(
|
||||
code="c", redirect_uri="r", client_id="ci", code_verifier="cv", state="s"
|
||||
)
|
||||
params = req.form_params()
|
||||
expected_keys = {"grant_type", "code", "redirect_uri", "client_id", "code_verifier", "state"}
|
||||
self.assertEqual(set(params.keys()), expected_keys)
|
||||
|
||||
|
||||
class TestOAuthTokenSet(unittest.TestCase):
|
||||
def test_required_fields(self):
|
||||
ts = OAuthTokenSet(access_token="tok123", scopes=["openid"])
|
||||
self.assertEqual(ts.access_token, "tok123")
|
||||
self.assertEqual(ts.scopes, ["openid"])
|
||||
self.assertIsNone(ts.refresh_token)
|
||||
self.assertIsNone(ts.expires_at)
|
||||
|
||||
def test_all_fields(self):
|
||||
ts = OAuthTokenSet(
|
||||
access_token="tok",
|
||||
scopes=["openid"],
|
||||
refresh_token="ref",
|
||||
expires_at=1234567890,
|
||||
)
|
||||
self.assertEqual(ts.refresh_token, "ref")
|
||||
self.assertEqual(ts.expires_at, 1234567890)
|
||||
|
||||
|
||||
class TestOAuthCredentialStore(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.store = OAuthCredentialStore(config_home=self.tmpdir)
|
||||
|
||||
def test_credentials_path_default(self):
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
# Remove HANZO_CONFIG_HOME if set
|
||||
env = os.environ.copy()
|
||||
env.pop("HANZO_CONFIG_HOME", None)
|
||||
with patch.dict(os.environ, env, clear=True):
|
||||
store = OAuthCredentialStore()
|
||||
expected = Path.home() / ".hanzo" / "credentials.json"
|
||||
self.assertEqual(store.credentials_path(), expected)
|
||||
|
||||
def test_credentials_path_custom_config_home(self):
|
||||
self.assertEqual(
|
||||
self.store.credentials_path(),
|
||||
Path(self.tmpdir) / "credentials.json",
|
||||
)
|
||||
|
||||
def test_credentials_path_env_override(self):
|
||||
with patch.dict(os.environ, {"HANZO_CONFIG_HOME": "/custom/path"}):
|
||||
store = OAuthCredentialStore()
|
||||
self.assertEqual(store.credentials_path(), Path("/custom/path/credentials.json"))
|
||||
|
||||
def test_save_and_load_roundtrip(self):
|
||||
token_set = OAuthTokenSet(
|
||||
access_token="access123",
|
||||
refresh_token="refresh456",
|
||||
expires_at=9999999999,
|
||||
scopes=["openid", "profile"],
|
||||
)
|
||||
self.store.save(token_set)
|
||||
|
||||
loaded = self.store.load()
|
||||
self.assertIsNotNone(loaded)
|
||||
self.assertEqual(loaded.access_token, "access123")
|
||||
self.assertEqual(loaded.refresh_token, "refresh456")
|
||||
self.assertEqual(loaded.expires_at, 9999999999)
|
||||
self.assertEqual(loaded.scopes, ["openid", "profile"])
|
||||
|
||||
def test_save_preserves_other_keys(self):
|
||||
creds_path = self.store.credentials_path()
|
||||
creds_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(creds_path, "w") as f:
|
||||
json.dump({"api_key": "keep-me", "other": "data"}, f)
|
||||
|
||||
token_set = OAuthTokenSet(access_token="tok", scopes=["openid"])
|
||||
self.store.save(token_set)
|
||||
|
||||
with open(creds_path) as f:
|
||||
data = json.load(f)
|
||||
|
||||
self.assertEqual(data["api_key"], "keep-me")
|
||||
self.assertEqual(data["other"], "data")
|
||||
self.assertIn("oauth", data)
|
||||
|
||||
def test_load_returns_none_when_no_file(self):
|
||||
self.assertIsNone(self.store.load())
|
||||
|
||||
def test_load_returns_none_when_no_oauth_key(self):
|
||||
creds_path = self.store.credentials_path()
|
||||
creds_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(creds_path, "w") as f:
|
||||
json.dump({"api_key": "something"}, f)
|
||||
|
||||
self.assertIsNone(self.store.load())
|
||||
|
||||
def test_clear_removes_oauth_preserves_rest(self):
|
||||
creds_path = self.store.credentials_path()
|
||||
creds_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(creds_path, "w") as f:
|
||||
json.dump({"api_key": "keep", "oauth": {"access_token": "remove"}}, f)
|
||||
|
||||
self.store.clear()
|
||||
|
||||
with open(creds_path) as f:
|
||||
data = json.load(f)
|
||||
|
||||
self.assertEqual(data["api_key"], "keep")
|
||||
self.assertNotIn("oauth", data)
|
||||
|
||||
def test_clear_noop_when_no_file(self):
|
||||
self.store.clear() # should not raise
|
||||
|
||||
def test_save_atomic_no_tmp_left(self):
|
||||
"""After save, .tmp file should not exist."""
|
||||
token_set = OAuthTokenSet(access_token="tok", scopes=["openid"])
|
||||
self.store.save(token_set)
|
||||
|
||||
creds_path = self.store.credentials_path()
|
||||
tmp_path = creds_path.with_suffix(".json.tmp")
|
||||
self.assertTrue(creds_path.exists())
|
||||
self.assertFalse(tmp_path.exists())
|
||||
|
||||
def test_save_sets_restrictive_permissions(self):
|
||||
token_set = OAuthTokenSet(access_token="tok", scopes=["openid"])
|
||||
self.store.save(token_set)
|
||||
|
||||
creds_path = self.store.credentials_path()
|
||||
mode = os.stat(creds_path).st_mode
|
||||
self.assertEqual(stat.S_IMODE(mode), 0o600)
|
||||
|
||||
def test_save_creates_parent_dirs(self):
|
||||
nested = os.path.join(self.tmpdir, "deep", "nested")
|
||||
store = OAuthCredentialStore(config_home=nested)
|
||||
token_set = OAuthTokenSet(access_token="tok", scopes=["openid"])
|
||||
store.save(token_set)
|
||||
self.assertTrue(store.credentials_path().exists())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,410 @@
|
||||
"""Tests for hanzoai.config — hierarchical config discovery and merging."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from hanzoai.config import (
|
||||
ConfigEntry,
|
||||
ConfigError,
|
||||
ConfigLoader,
|
||||
ConfigSource,
|
||||
McpClaudeAiProxyConfig,
|
||||
McpOAuthConfig,
|
||||
McpRemoteConfig,
|
||||
McpSdkConfig,
|
||||
McpStdioConfig,
|
||||
McpTransport,
|
||||
McpWsConfig,
|
||||
OAuthConfig,
|
||||
RuntimeConfig,
|
||||
_deep_merge,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Deep merge
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDeepMerge:
|
||||
def test_flat_override(self):
|
||||
assert _deep_merge({"a": 1}, {"a": 2}) == {"a": 2}
|
||||
|
||||
def test_flat_union(self):
|
||||
assert _deep_merge({"a": 1}, {"b": 2}) == {"a": 1, "b": 2}
|
||||
|
||||
def test_nested_merge(self):
|
||||
base = {"a": {"x": 1, "y": 2}}
|
||||
over = {"a": {"y": 3, "z": 4}}
|
||||
assert _deep_merge(base, over) == {"a": {"x": 1, "y": 3, "z": 4}}
|
||||
|
||||
def test_nested_replace_non_dict(self):
|
||||
assert _deep_merge({"a": {"x": 1}}, {"a": 42}) == {"a": 42}
|
||||
|
||||
def test_empty(self):
|
||||
assert _deep_merge({}, {"a": 1}) == {"a": 1}
|
||||
assert _deep_merge({"a": 1}, {}) == {"a": 1}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ConfigLoader.discover
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDiscover:
|
||||
def test_returns_three_entries(self, tmp_path: Path):
|
||||
loader = ConfigLoader(cwd=tmp_path, config_home=tmp_path / ".hanzo")
|
||||
entries = loader.discover()
|
||||
assert len(entries) == 3
|
||||
assert entries[0].source == ConfigSource.User
|
||||
assert entries[1].source == ConfigSource.Project
|
||||
assert entries[2].source == ConfigSource.Local
|
||||
|
||||
def test_paths(self, tmp_path: Path):
|
||||
home = tmp_path / "home" / ".hanzo"
|
||||
cwd = tmp_path / "project"
|
||||
loader = ConfigLoader(cwd=cwd, config_home=home)
|
||||
entries = loader.discover()
|
||||
assert entries[0].path == home / "settings.json"
|
||||
assert entries[1].path == cwd / ".hanzo" / "settings.json"
|
||||
assert entries[2].path == cwd / ".hanzo" / "settings.local.json"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ConfigLoader.default_for
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestDefaultFor:
|
||||
def test_uses_env_var(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
custom = tmp_path / "custom"
|
||||
monkeypatch.setenv("HANZO_CONFIG_HOME", str(custom))
|
||||
loader = ConfigLoader.default_for(tmp_path)
|
||||
assert loader.config_home == custom
|
||||
|
||||
def test_falls_back_to_home(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.delenv("HANZO_CONFIG_HOME", raising=False)
|
||||
loader = ConfigLoader.default_for(tmp_path)
|
||||
assert loader.config_home == Path.home() / ".hanzo"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ConfigLoader.load — file handling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _write_json(path: Path, data: object) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(data), encoding="utf-8")
|
||||
|
||||
|
||||
class TestLoad:
|
||||
def test_empty_when_no_files(self, tmp_path: Path):
|
||||
loader = ConfigLoader(cwd=tmp_path, config_home=tmp_path / "h")
|
||||
cfg = loader.load()
|
||||
assert cfg.merged == {}
|
||||
assert cfg.loaded_entries == []
|
||||
|
||||
def test_single_user_file(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", {"theme": "dark"})
|
||||
cfg = ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
assert cfg.merged == {"theme": "dark"}
|
||||
assert len(cfg.loaded_entries) == 1
|
||||
assert cfg.loaded_entries[0].source == ConfigSource.User
|
||||
|
||||
def test_project_overrides_user(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", {"a": 1, "b": 2})
|
||||
_write_json(tmp_path / ".hanzo" / "settings.json", {"b": 99})
|
||||
cfg = ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
assert cfg.merged == {"a": 1, "b": 99}
|
||||
assert len(cfg.loaded_entries) == 2
|
||||
|
||||
def test_local_overrides_project(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", {"a": 1})
|
||||
_write_json(tmp_path / ".hanzo" / "settings.json", {"b": 2})
|
||||
_write_json(tmp_path / ".hanzo" / "settings.local.json", {"b": 3, "c": 4})
|
||||
cfg = ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
assert cfg.merged == {"a": 1, "b": 3, "c": 4}
|
||||
assert len(cfg.loaded_entries) == 3
|
||||
|
||||
def test_invalid_json_raises(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
p = home / "settings.json"
|
||||
p.parent.mkdir(parents=True, exist_ok=True)
|
||||
p.write_text("{bad json", encoding="utf-8")
|
||||
with pytest.raises(ConfigError, match="invalid JSON"):
|
||||
ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
|
||||
def test_non_object_top_level_raises(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", [1, 2, 3])
|
||||
with pytest.raises(ConfigError, match="top-level value must be a JSON object"):
|
||||
ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
|
||||
def test_mcp_servers_replaced_not_merged(self, tmp_path: Path):
|
||||
home = tmp_path / "h"
|
||||
_write_json(home / "settings.json", {
|
||||
"mcpServers": {
|
||||
"alpha": {"command": "a", "transport": "stdio"},
|
||||
"beta": {"command": "b", "transport": "stdio"},
|
||||
}
|
||||
})
|
||||
_write_json(tmp_path / ".hanzo" / "settings.json", {
|
||||
"mcpServers": {
|
||||
"gamma": {"command": "g", "transport": "stdio"},
|
||||
}
|
||||
})
|
||||
cfg = ConfigLoader(cwd=tmp_path, config_home=home).load()
|
||||
servers = cfg.mcp_servers()
|
||||
assert set(servers.keys()) == {"gamma"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RuntimeConfig accessors
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRuntimeConfig:
|
||||
def test_empty(self):
|
||||
cfg = RuntimeConfig.empty()
|
||||
assert cfg.merged == {}
|
||||
assert cfg.get("anything") is None
|
||||
|
||||
def test_get(self):
|
||||
cfg = RuntimeConfig(merged={"foo": "bar"})
|
||||
assert cfg.get("foo") == "bar"
|
||||
assert cfg.get("nope") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MCP server parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestMcpServers:
|
||||
def test_stdio(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"myserver": {
|
||||
"transport": "stdio",
|
||||
"command": "node",
|
||||
"args": ["server.js"],
|
||||
"env": {"PORT": "8080"},
|
||||
}
|
||||
}
|
||||
})
|
||||
servers = cfg.mcp_servers()
|
||||
s = servers["myserver"]
|
||||
assert isinstance(s, McpStdioConfig)
|
||||
assert s.command == "node"
|
||||
assert s.args == ["server.js"]
|
||||
assert s.env == {"PORT": "8080"}
|
||||
assert s.transport == McpTransport.Stdio
|
||||
|
||||
def test_stdio_default_transport(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"s": {"command": "echo"}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["s"]
|
||||
assert isinstance(s, McpStdioConfig)
|
||||
|
||||
def test_http_remote(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"remote": {
|
||||
"transport": "http",
|
||||
"url": "https://example.com/mcp",
|
||||
"headers": {"Authorization": "Bearer tok"},
|
||||
}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["remote"]
|
||||
assert isinstance(s, McpRemoteConfig)
|
||||
assert s.url == "https://example.com/mcp"
|
||||
assert s.transport == McpTransport.Http
|
||||
|
||||
def test_sse_remote(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"sse": {"transport": "sse", "url": "https://x.com/events"}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["sse"]
|
||||
assert isinstance(s, McpRemoteConfig)
|
||||
assert s.transport == McpTransport.Sse
|
||||
|
||||
def test_ws(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"ws": {"transport": "ws", "url": "wss://x.com/ws"}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["ws"]
|
||||
assert isinstance(s, McpWsConfig)
|
||||
assert s.transport == McpTransport.Ws
|
||||
assert s.url == "wss://x.com/ws"
|
||||
|
||||
def test_sdk(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"my-sdk": {"transport": "sdk", "name": "my-sdk-server"}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["my-sdk"]
|
||||
assert isinstance(s, McpSdkConfig)
|
||||
assert s.transport == McpTransport.Sdk
|
||||
assert s.name == "my-sdk-server"
|
||||
|
||||
def test_claudeai_proxy(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"proxy": {
|
||||
"transport": "claudeai-proxy",
|
||||
"url": "https://claude.ai/proxy",
|
||||
"id": "server-123",
|
||||
}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["proxy"]
|
||||
assert isinstance(s, McpClaudeAiProxyConfig)
|
||||
assert s.transport == McpTransport.ClaudeAiProxy
|
||||
assert s.url == "https://claude.ai/proxy"
|
||||
assert s.id == "server-123"
|
||||
|
||||
def test_remote_with_oauth(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {
|
||||
"authed": {
|
||||
"transport": "http",
|
||||
"url": "https://example.com/mcp",
|
||||
"oauth": {
|
||||
"clientId": "cid",
|
||||
"callbackPort": 9999,
|
||||
"authServerMetadataUrl": "https://auth.example.com/.well-known",
|
||||
"xaa": True,
|
||||
},
|
||||
}
|
||||
}
|
||||
})
|
||||
s = cfg.mcp_servers()["authed"]
|
||||
assert isinstance(s, McpRemoteConfig)
|
||||
assert s.oauth is not None
|
||||
assert s.oauth.client_id == "cid"
|
||||
assert s.oauth.callback_port == 9999
|
||||
assert s.oauth.xaa is True
|
||||
|
||||
def test_sdk_missing_name_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "sdk"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="requires 'name'"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_claudeai_proxy_missing_id_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "claudeai-proxy", "url": "https://x.com"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="requires 'id'"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_unknown_transport_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "grpc"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="unknown transport"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_stdio_missing_command_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "stdio"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="requires 'command'"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_http_missing_url_raises(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"mcpServers": {"bad": {"transport": "http"}}
|
||||
})
|
||||
with pytest.raises(ConfigError, match="requires 'url'"):
|
||||
cfg.mcp_servers()
|
||||
|
||||
def test_no_mcp_servers_returns_empty(self):
|
||||
assert RuntimeConfig.empty().mcp_servers() == {}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# OAuth parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestOAuth:
|
||||
def test_parses_camel_case(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"clientId": "my-id",
|
||||
"authorizeUrl": "https://hanzo.id/authorize",
|
||||
"tokenUrl": "https://hanzo.id/token",
|
||||
"scopes": ["read", "write"],
|
||||
"callbackPort": 8080,
|
||||
}
|
||||
})
|
||||
oauth = cfg.oauth()
|
||||
assert isinstance(oauth, OAuthConfig)
|
||||
assert oauth.client_id == "my-id"
|
||||
assert oauth.authorize_url == "https://hanzo.id/authorize"
|
||||
assert oauth.token_url == "https://hanzo.id/token"
|
||||
assert oauth.scopes == ["read", "write"]
|
||||
assert oauth.callback_port == 8080
|
||||
|
||||
def test_parses_snake_case(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"client_id": "cid",
|
||||
"authorize_url": "https://hanzo.id/authorize",
|
||||
"token_url": "https://hanzo.id/token",
|
||||
}
|
||||
})
|
||||
oauth = cfg.oauth()
|
||||
assert oauth is not None
|
||||
assert oauth.client_id == "cid"
|
||||
assert oauth.authorize_url == "https://hanzo.id/authorize"
|
||||
|
||||
def test_returns_none_when_missing(self):
|
||||
assert RuntimeConfig.empty().oauth() is None
|
||||
|
||||
def test_returns_none_when_no_client_id(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"authorizeUrl": "https://x.com/auth",
|
||||
"tokenUrl": "https://x.com/token",
|
||||
}
|
||||
})
|
||||
assert cfg.oauth() is None
|
||||
|
||||
def test_returns_none_when_no_authorize_url(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"clientId": "cid",
|
||||
"tokenUrl": "https://x.com/token",
|
||||
}
|
||||
})
|
||||
assert cfg.oauth() is None
|
||||
|
||||
def test_returns_none_when_no_token_url(self):
|
||||
cfg = RuntimeConfig(merged={
|
||||
"oauth": {
|
||||
"clientId": "cid",
|
||||
"authorizeUrl": "https://x.com/auth",
|
||||
}
|
||||
})
|
||||
assert cfg.oauth() is None
|
||||
@@ -0,0 +1,242 @@
|
||||
"""Tests for MCP client implementation."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import textwrap
|
||||
|
||||
import pytest
|
||||
|
||||
from hanzoai.mcp import (
|
||||
JsonRpcError,
|
||||
JsonRpcRequest,
|
||||
JsonRpcResponse,
|
||||
MCPClient,
|
||||
MCPClientError,
|
||||
MCPHttpClient,
|
||||
mcp_tool_name,
|
||||
normalize_mcp_name,
|
||||
)
|
||||
|
||||
|
||||
# --- Name normalization tests ---
|
||||
|
||||
|
||||
class TestNormalizeMcpName:
|
||||
def test_passthrough(self):
|
||||
assert normalize_mcp_name("hello") == "hello"
|
||||
|
||||
def test_hyphens_kept(self):
|
||||
# Rust parity: hyphens are valid chars and kept as-is
|
||||
assert normalize_mcp_name("my-tool") == "my-tool"
|
||||
|
||||
def test_dots_to_underscores(self):
|
||||
assert normalize_mcp_name("my.tool") == "my_tool"
|
||||
|
||||
def test_slashes_to_underscores(self):
|
||||
assert normalize_mcp_name("my/tool") == "my_tool"
|
||||
|
||||
def test_mixed_keeps_hyphens(self):
|
||||
assert normalize_mcp_name("my-tool.v2/action") == "my-tool_v2_action"
|
||||
|
||||
def test_non_claude_keeps_consecutive_underscores(self):
|
||||
# Non-claude.ai names do NOT collapse underscores
|
||||
assert normalize_mcp_name("my--tool") == "my--tool"
|
||||
|
||||
def test_non_claude_keeps_leading_trailing(self):
|
||||
# Non-claude.ai names do NOT strip leading/trailing
|
||||
assert normalize_mcp_name("-tool-") == "-tool-"
|
||||
|
||||
def test_non_claude_trailing_underscore(self):
|
||||
assert normalize_mcp_name("tool name!") == "tool_name_"
|
||||
|
||||
def test_github_dot_com(self):
|
||||
assert normalize_mcp_name("github.com") == "github_com"
|
||||
|
||||
def test_claude_ai_collapses_and_strips(self):
|
||||
# claude.ai names get underscore collapsing + strip
|
||||
assert normalize_mcp_name("claude.ai Example Server!!") == "claude_ai_Example_Server"
|
||||
|
||||
def test_claude_ai_simple(self):
|
||||
assert normalize_mcp_name("claude.ai My Tool") == "claude_ai_My_Tool"
|
||||
|
||||
|
||||
class TestMcpToolName:
|
||||
def test_basic(self):
|
||||
assert mcp_tool_name("server", "tool") == "mcp__server__tool"
|
||||
|
||||
def test_normalizes_both(self):
|
||||
# Hyphens kept, so my-server stays my-server
|
||||
assert mcp_tool_name("my-server", "my-tool") == "mcp__my-server__my-tool"
|
||||
|
||||
def test_claude_ai_server(self):
|
||||
assert mcp_tool_name("claude.ai Example Server", "weather tool") == "mcp__claude_ai_Example_Server__weather_tool"
|
||||
|
||||
|
||||
# --- Protocol types tests ---
|
||||
|
||||
|
||||
class TestJsonRpcRequest:
|
||||
def test_to_json(self):
|
||||
req = JsonRpcRequest(method="initialize", params={"a": 1}, id=1)
|
||||
d = req.to_dict()
|
||||
assert d["jsonrpc"] == "2.0"
|
||||
assert d["method"] == "initialize"
|
||||
assert d["params"] == {"a": 1}
|
||||
assert d["id"] == 1
|
||||
|
||||
def test_to_line(self):
|
||||
req = JsonRpcRequest(method="test", params={}, id=42)
|
||||
line = req.to_line()
|
||||
assert line.endswith(b"\n")
|
||||
parsed = json.loads(line)
|
||||
assert parsed["id"] == 42
|
||||
|
||||
|
||||
class TestJsonRpcResponse:
|
||||
def test_from_dict_success(self):
|
||||
resp = JsonRpcResponse.from_dict({"jsonrpc": "2.0", "result": {"ok": True}, "id": 1})
|
||||
assert resp.result == {"ok": True}
|
||||
assert resp.error is None
|
||||
assert resp.id == 1
|
||||
|
||||
def test_from_dict_error(self):
|
||||
resp = JsonRpcResponse.from_dict({
|
||||
"jsonrpc": "2.0",
|
||||
"error": {"code": -32600, "message": "Invalid Request"},
|
||||
"id": 1,
|
||||
})
|
||||
assert resp.result is None
|
||||
assert resp.error is not None
|
||||
assert resp.error.code == -32600
|
||||
assert resp.error.message == "Invalid Request"
|
||||
|
||||
|
||||
# --- Stdio MCPClient tests ---
|
||||
|
||||
# A minimal MCP server script that we spawn as a subprocess.
|
||||
# It reads JSON-RPC from stdin and responds on stdout.
|
||||
MOCK_SERVER_SCRIPT = textwrap.dedent("""\
|
||||
import json
|
||||
import sys
|
||||
|
||||
def respond(id, result):
|
||||
msg = json.dumps({"jsonrpc": "2.0", "result": result, "id": id})
|
||||
sys.stdout.write(msg + "\\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
def respond_error(id, code, message):
|
||||
msg = json.dumps({"jsonrpc": "2.0", "error": {"code": code, "message": message}, "id": id})
|
||||
sys.stdout.write(msg + "\\n")
|
||||
sys.stdout.flush()
|
||||
|
||||
for line in sys.stdin:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
req = json.loads(line)
|
||||
method = req["method"]
|
||||
rid = req["id"]
|
||||
|
||||
if method == "initialize":
|
||||
respond(rid, {
|
||||
"protocolVersion": "2024-11-05",
|
||||
"capabilities": {"tools": {}, "resources": {}},
|
||||
"serverInfo": {"name": "mock", "version": "0.1.0"},
|
||||
})
|
||||
elif method == "tools/list":
|
||||
respond(rid, {
|
||||
"tools": [
|
||||
{"name": "echo", "description": "echoes input", "inputSchema": {"type": "object"}},
|
||||
],
|
||||
})
|
||||
elif method == "tools/call":
|
||||
tool_name = req["params"]["name"]
|
||||
args = req["params"].get("arguments", {})
|
||||
if tool_name == "echo":
|
||||
respond(rid, {"content": [{"type": "text", "text": json.dumps(args)}]})
|
||||
elif tool_name == "fail":
|
||||
respond_error(rid, -32000, "tool failed")
|
||||
else:
|
||||
respond_error(rid, -32601, f"unknown tool: {tool_name}")
|
||||
elif method == "resources/list":
|
||||
respond(rid, {"resources": [{"uri": "file:///test.txt", "name": "test.txt"}]})
|
||||
else:
|
||||
respond_error(rid, -32601, f"unknown method: {method}")
|
||||
""")
|
||||
|
||||
|
||||
class TestMCPClientStdio:
|
||||
@pytest.fixture
|
||||
async def client(self, tmp_path):
|
||||
script = tmp_path / "mock_server.py"
|
||||
script.write_text(MOCK_SERVER_SCRIPT)
|
||||
c = MCPClient(server_command=[sys.executable, str(script)])
|
||||
await c.connect()
|
||||
yield c
|
||||
await c.disconnect()
|
||||
|
||||
async def test_connect_initializes(self, client):
|
||||
assert client.server_info is not None
|
||||
assert client.server_info["name"] == "mock"
|
||||
|
||||
async def test_list_tools(self, client):
|
||||
tools = await client.list_tools()
|
||||
assert len(tools) == 1
|
||||
assert tools[0]["name"] == "echo"
|
||||
|
||||
async def test_call_tool(self, client):
|
||||
result = await client.call_tool("echo", arguments={"msg": "hello"})
|
||||
assert result["content"][0]["text"] == '{"msg": "hello"}'
|
||||
|
||||
async def test_call_tool_error(self, client):
|
||||
with pytest.raises(MCPClientError, match="tool failed"):
|
||||
await client.call_tool("fail")
|
||||
|
||||
async def test_list_resources(self, client):
|
||||
resources = await client.list_resources()
|
||||
assert len(resources) == 1
|
||||
assert resources[0]["uri"] == "file:///test.txt"
|
||||
|
||||
async def test_context_manager(self, tmp_path):
|
||||
script = tmp_path / "mock_server.py"
|
||||
script.write_text(MOCK_SERVER_SCRIPT)
|
||||
async with MCPClient(server_command=[sys.executable, str(script)]) as c:
|
||||
tools = await c.list_tools()
|
||||
assert len(tools) == 1
|
||||
|
||||
async def test_request_ids_increment(self, client):
|
||||
# Each call increments the request id
|
||||
id_before = client._next_id
|
||||
await client.list_tools()
|
||||
assert client._next_id == id_before + 1
|
||||
|
||||
async def test_disconnect_kills_process(self, tmp_path):
|
||||
script = tmp_path / "mock_server.py"
|
||||
script.write_text(MOCK_SERVER_SCRIPT)
|
||||
c = MCPClient(server_command=[sys.executable, str(script)])
|
||||
await c.connect()
|
||||
proc = c._process
|
||||
await c.disconnect()
|
||||
assert c._process is None
|
||||
# Process should be terminated
|
||||
assert proc.returncode is not None
|
||||
|
||||
|
||||
class TestMCPClientNotConnected:
|
||||
async def test_list_tools_without_connect(self):
|
||||
c = MCPClient(server_command=["true"])
|
||||
with pytest.raises(MCPClientError, match="not connected"):
|
||||
await c.list_tools()
|
||||
|
||||
async def test_call_tool_without_connect(self):
|
||||
c = MCPClient(server_command=["true"])
|
||||
with pytest.raises(MCPClientError, match="not connected"):
|
||||
await c.call_tool("x")
|
||||
|
||||
|
||||
class TestMCPClientBadProcess:
|
||||
async def test_connect_to_nonexistent_binary(self):
|
||||
c = MCPClient(server_command=["__nonexistent_binary_xyz__"])
|
||||
with pytest.raises(MCPClientError):
|
||||
await c.connect()
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Tests for session management and compaction."""
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from hanzoai.session import (
|
||||
CompactionConfig,
|
||||
CompactionResult,
|
||||
ConversationMessage,
|
||||
MessageRole,
|
||||
Session,
|
||||
TextBlock,
|
||||
TokenUsage,
|
||||
ToolResultBlock,
|
||||
ToolUseBlock,
|
||||
compact_session,
|
||||
estimate_session_tokens,
|
||||
format_compact_summary,
|
||||
should_compact,
|
||||
)
|
||||
|
||||
|
||||
class TestTokenUsage:
|
||||
def test_total_tokens(self):
|
||||
u = TokenUsage(input_tokens=10, output_tokens=4, cache_creation_input_tokens=2, cache_read_input_tokens=1)
|
||||
assert u.total_tokens() == 17
|
||||
|
||||
def test_roundtrip(self):
|
||||
u = TokenUsage(input_tokens=5, output_tokens=3, cache_creation_input_tokens=1, cache_read_input_tokens=0)
|
||||
assert TokenUsage.from_dict(u.to_dict()) == u
|
||||
|
||||
|
||||
class TestContentBlocks:
|
||||
def test_text_block_roundtrip(self):
|
||||
b = TextBlock(text="hello")
|
||||
d = b.to_dict()
|
||||
assert d == {"type": "text", "text": "hello"}
|
||||
assert TextBlock.from_dict(d) == b
|
||||
|
||||
def test_tool_use_block_roundtrip(self):
|
||||
b = ToolUseBlock(id="t1", name="bash", input={"command": "ls"})
|
||||
d = b.to_dict()
|
||||
assert d["type"] == "tool_use"
|
||||
assert ToolUseBlock.from_dict(d) == b
|
||||
|
||||
def test_tool_result_block_roundtrip(self):
|
||||
b = ToolResultBlock(tool_use_id="t1", tool_name="bash", output="ok", is_error=False)
|
||||
d = b.to_dict()
|
||||
assert d["type"] == "tool_result"
|
||||
assert ToolResultBlock.from_dict(d) == b
|
||||
|
||||
|
||||
class TestConversationMessage:
|
||||
def test_user_text_factory(self):
|
||||
m = ConversationMessage.user_text("hello")
|
||||
assert m.role == MessageRole.User
|
||||
assert len(m.blocks) == 1
|
||||
assert isinstance(m.blocks[0], TextBlock)
|
||||
|
||||
def test_tool_result_factory(self):
|
||||
m = ConversationMessage.tool_result("t1", "bash", "output", is_error=True)
|
||||
assert m.role == MessageRole.Tool
|
||||
assert isinstance(m.blocks[0], ToolResultBlock)
|
||||
assert m.blocks[0].is_error
|
||||
|
||||
def test_assistant_with_usage(self):
|
||||
u = TokenUsage(input_tokens=10, output_tokens=4)
|
||||
m = ConversationMessage.assistant_with_usage([TextBlock(text="hi")], u)
|
||||
assert m.usage == u
|
||||
|
||||
def test_roundtrip(self):
|
||||
m = ConversationMessage.user_text("test")
|
||||
assert ConversationMessage.from_dict(m.to_dict()).role == m.role
|
||||
|
||||
|
||||
class TestSession:
|
||||
def test_empty_session(self):
|
||||
s = Session()
|
||||
assert s.version == 1
|
||||
assert s.messages == []
|
||||
|
||||
def test_roundtrip_dict(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("hello"))
|
||||
s.messages.append(ConversationMessage.assistant([TextBlock(text="hi")]))
|
||||
s.messages.append(ConversationMessage.tool_result("t1", "bash", "ok", False))
|
||||
|
||||
restored = Session.from_dict(s.to_dict())
|
||||
assert len(restored.messages) == 3
|
||||
assert restored.messages[0].role == MessageRole.User
|
||||
assert restored.messages[2].role == MessageRole.Tool
|
||||
|
||||
def test_save_and_load(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("hello"))
|
||||
s.messages.append(
|
||||
ConversationMessage.assistant_with_usage(
|
||||
[
|
||||
TextBlock(text="thinking"),
|
||||
ToolUseBlock(id="t1", name="bash", input={"cmd": "echo hi"}),
|
||||
],
|
||||
TokenUsage(input_tokens=10, output_tokens=4, cache_creation_input_tokens=1, cache_read_input_tokens=2),
|
||||
)
|
||||
)
|
||||
s.messages.append(ConversationMessage.tool_result("t1", "bash", "hi", False))
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".json", delete=False) as f:
|
||||
path = f.name
|
||||
|
||||
s.save(path)
|
||||
restored = Session.load(path)
|
||||
|
||||
assert restored == s
|
||||
assert restored.messages[1].usage.total_tokens() == 17
|
||||
assert restored.messages[2].role == MessageRole.Tool
|
||||
Path(path).unlink()
|
||||
|
||||
|
||||
class TestCompaction:
|
||||
def _big_session(self, n: int = 10) -> Session:
|
||||
s = Session()
|
||||
for i in range(n):
|
||||
s.messages.append(ConversationMessage.user_text("x " * 200))
|
||||
s.messages.append(ConversationMessage.assistant([TextBlock(text="y " * 200)]))
|
||||
return s
|
||||
|
||||
def test_small_session_unchanged(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("hello"))
|
||||
result = compact_session(s, CompactionConfig())
|
||||
assert result.removed_message_count == 0
|
||||
assert result.compacted_session is s
|
||||
|
||||
def test_compacts_old_messages(self):
|
||||
s = self._big_session(10)
|
||||
config = CompactionConfig(preserve_recent_messages=2, max_estimated_tokens=1)
|
||||
result = compact_session(s, config)
|
||||
|
||||
assert result.removed_message_count == 18
|
||||
assert result.compacted_session.messages[0].role == MessageRole.System
|
||||
assert len(result.compacted_session.messages) == 3 # system + 2 preserved
|
||||
assert result.summary != ""
|
||||
assert result.formatted_summary != ""
|
||||
|
||||
def test_should_compact_checks_both(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("short"))
|
||||
# Few messages, low tokens - should not compact
|
||||
assert not should_compact(s, CompactionConfig(preserve_recent_messages=0, max_estimated_tokens=100))
|
||||
|
||||
def test_estimate_tokens(self):
|
||||
s = Session()
|
||||
s.messages.append(ConversationMessage.user_text("a" * 400))
|
||||
tokens = estimate_session_tokens(s)
|
||||
assert tokens == 400 // 4 + 1 # 101
|
||||
|
||||
def test_compaction_reduces_tokens(self):
|
||||
s = self._big_session(10)
|
||||
config = CompactionConfig(preserve_recent_messages=2, max_estimated_tokens=1)
|
||||
result = compact_session(s, config)
|
||||
assert estimate_session_tokens(result.compacted_session) < estimate_session_tokens(s)
|
||||
|
||||
def test_summary_contains_tool_names(self):
|
||||
s = Session()
|
||||
for _ in range(6):
|
||||
s.messages.append(ConversationMessage.user_text("x " * 200))
|
||||
s.messages.append(
|
||||
ConversationMessage.assistant(
|
||||
[TextBlock(text="ok"), ToolUseBlock(id="t1", name="bash", input={"cmd": "ls"})]
|
||||
)
|
||||
)
|
||||
s.messages.append(ConversationMessage.tool_result("t1", "bash", "output " * 50, False))
|
||||
|
||||
config = CompactionConfig(preserve_recent_messages=2, max_estimated_tokens=1)
|
||||
result = compact_session(s, config)
|
||||
assert "bash" in result.summary
|
||||
|
||||
def test_summary_contains_file_candidates(self):
|
||||
s = Session()
|
||||
for _ in range(5):
|
||||
s.messages.append(ConversationMessage.user_text("Update src/main.rs and pkg/config.py next " + "x " * 200))
|
||||
s.messages.append(ConversationMessage.assistant([TextBlock(text="done " * 200)]))
|
||||
|
||||
config = CompactionConfig(preserve_recent_messages=2, max_estimated_tokens=1)
|
||||
result = compact_session(s, config)
|
||||
assert "src/main.rs" in result.summary
|
||||
|
||||
def test_format_compact_summary_strips_analysis(self):
|
||||
summary = "<analysis>scratch</analysis>\n<summary>Kept work</summary>"
|
||||
formatted = format_compact_summary(summary)
|
||||
assert "analysis" not in formatted
|
||||
assert "Kept work" in formatted
|
||||
|
||||
def test_format_compact_summary_no_tags(self):
|
||||
formatted = format_compact_summary("plain text summary")
|
||||
assert formatted == "plain text summary"
|
||||
@@ -0,0 +1,181 @@
|
||||
"""Tests for SSE chunking and decoding in hanzoai._streaming."""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Iterator, AsyncIterator
|
||||
|
||||
import pytest
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent / "pkg" / "hanzoai"))
|
||||
|
||||
from hanzoai._streaming import SSEDecoder, ServerSentEvent
|
||||
|
||||
|
||||
def make_sync_iterator(chunks: list[bytes]) -> Iterator[bytes]:
|
||||
yield from chunks
|
||||
|
||||
|
||||
async def make_async_iterator(chunks: list[bytes]) -> AsyncIterator[bytes]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
|
||||
class TestIterChunks:
|
||||
"""Test that _iter_chunks correctly splits SSE frames."""
|
||||
|
||||
def _collect_chunks(self, raw_chunks: list[bytes]) -> list[bytes]:
|
||||
decoder = SSEDecoder()
|
||||
return list(decoder._iter_chunks(make_sync_iterator(raw_chunks)))
|
||||
|
||||
def test_simple_frame_lf(self):
|
||||
"""Single frame terminated by \\n\\n."""
|
||||
raw = [b"data: hello\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\n\n"]
|
||||
|
||||
def test_simple_frame_crlf(self):
|
||||
"""Single frame terminated by \\r\\n\\r\\n."""
|
||||
raw = [b"data: hello\r\n\r\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
def test_simple_frame_cr(self):
|
||||
"""Single frame terminated by \\r\\r."""
|
||||
raw = [b"data: hello\r\r"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\r"]
|
||||
|
||||
def test_multiple_frames(self):
|
||||
"""Multiple frames in one byte chunk."""
|
||||
raw = [b"data: one\n\ndata: two\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: one\n\n", b"data: two\n\n"]
|
||||
|
||||
def test_split_across_chunks(self):
|
||||
"""Frame split across multiple byte chunks."""
|
||||
raw = [b"data: hel", b"lo\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\n\n"]
|
||||
|
||||
def test_separator_split_across_chunks_lf(self):
|
||||
"""\\n\\n separator split across chunk boundary."""
|
||||
raw = [b"data: hello\n", b"\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\n\n"]
|
||||
|
||||
def test_separator_split_across_chunks_crlf(self):
|
||||
"""\\r\\n\\r\\n separator split across chunk boundary at various points."""
|
||||
# Split between \r\n and \r\n
|
||||
raw = [b"data: hello\r\n", b"\r\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
def test_crlf_split_mid_separator(self):
|
||||
"""The critical bug case: \\r\\n\\r\\n split so \\r and \\n land in different chunks."""
|
||||
# This is the case splitlines(keepends=True) gets wrong:
|
||||
# chunk1 ends with \r, chunk2 starts with \n\r\n
|
||||
raw = [b"data: hello\r", b"\n\r\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
def test_crlf_split_three_ways(self):
|
||||
"""\\r\\n\\r\\n split across three chunks."""
|
||||
raw = [b"data: hello\r\n", b"\r", b"\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
def test_trailing_data_no_terminator(self):
|
||||
"""Data without a trailing separator is yielded on stream end."""
|
||||
raw = [b"data: hello"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello"]
|
||||
|
||||
def test_empty_input(self):
|
||||
"""No data yields nothing."""
|
||||
chunks = self._collect_chunks([])
|
||||
assert chunks == []
|
||||
|
||||
def test_only_separator(self):
|
||||
"""Just a separator yields it as a chunk."""
|
||||
raw = [b"\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"\n\n"]
|
||||
|
||||
def test_multiple_frames_split(self):
|
||||
"""Two frames, each split across chunks."""
|
||||
raw = [b"data: one\n", b"\ndata: tw", b"o\n\n"]
|
||||
chunks = self._collect_chunks(raw)
|
||||
assert chunks == [b"data: one\n\n", b"data: two\n\n"]
|
||||
|
||||
|
||||
class TestAiterChunks:
|
||||
"""Test that _aiter_chunks matches _iter_chunks behavior."""
|
||||
|
||||
async def _collect_chunks(self, raw_chunks: list[bytes]) -> list[bytes]:
|
||||
decoder = SSEDecoder()
|
||||
result = []
|
||||
async for chunk in decoder._aiter_chunks(make_async_iterator(raw_chunks)):
|
||||
result.append(chunk)
|
||||
return result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_simple_frame(self):
|
||||
chunks = await self._collect_chunks([b"data: hello\n\n"])
|
||||
assert chunks == [b"data: hello\n\n"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_crlf_split_mid_separator(self):
|
||||
"""The critical bug case for async path."""
|
||||
raw = [b"data: hello\r", b"\n\r\n"]
|
||||
chunks = await self._collect_chunks(raw)
|
||||
assert chunks == [b"data: hello\r\n\r\n"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_frames_split(self):
|
||||
raw = [b"data: one\n", b"\ndata: tw", b"o\n\n"]
|
||||
chunks = await self._collect_chunks(raw)
|
||||
assert chunks == [b"data: one\n\n", b"data: two\n\n"]
|
||||
|
||||
|
||||
class TestDoneAndPingSentinel:
|
||||
"""Test that [DONE] and event: ping are skipped."""
|
||||
|
||||
def _collect_events(self, raw_chunks: list[bytes]) -> list[ServerSentEvent]:
|
||||
decoder = SSEDecoder()
|
||||
return list(decoder.iter_bytes(make_sync_iterator(raw_chunks)))
|
||||
|
||||
def test_done_sentinel_skipped(self):
|
||||
"""data: [DONE] should not produce an event."""
|
||||
raw = [b"data: {\"text\": \"hi\"}\n\n", b"data: [DONE]\n\n"]
|
||||
events = self._collect_events(raw)
|
||||
assert len(events) == 1
|
||||
assert events[0].data == '{"text": "hi"}'
|
||||
|
||||
def test_ping_event_skipped(self):
|
||||
"""event: ping should not produce an event."""
|
||||
raw = [b"event: ping\ndata: \n\n", b"data: {\"text\": \"hi\"}\n\n"]
|
||||
events = self._collect_events(raw)
|
||||
assert len(events) == 1
|
||||
assert events[0].data == '{"text": "hi"}'
|
||||
|
||||
def test_normal_events_pass_through(self):
|
||||
"""Normal events should still work."""
|
||||
raw = [b"data: {\"a\": 1}\n\n", b"data: {\"b\": 2}\n\n"]
|
||||
events = self._collect_events(raw)
|
||||
assert len(events) == 2
|
||||
assert events[0].json() == {"a": 1}
|
||||
assert events[1].json() == {"b": 2}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_done_sentinel_skipped_async(self):
|
||||
"""data: [DONE] should not produce an event in async path."""
|
||||
decoder = SSEDecoder()
|
||||
raw = [b"data: {\"text\": \"hi\"}\n\n", b"data: [DONE]\n\n"]
|
||||
events = []
|
||||
async for sse in decoder.aiter_bytes(make_async_iterator(raw)):
|
||||
events.append(sse)
|
||||
assert len(events) == 1
|
||||
assert events[0].data == '{"text": "hi"}'
|
||||
Reference in New Issue
Block a user