This commit is contained in:
vasily committed
1 parent
4db3f45164
commit
d25c04cdd2
175 files changed
+56243
-15
No files matched your search
@@ -0,0 +1,374 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from mcp import Client
|
||||
from mcp.server import MCPServer
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
from gewu.domain.contracts import SourceRevisionRecord
|
||||
from gewu.transports.stdio_connector import (
|
||||
SelectedFile,
|
||||
UploadStateStore,
|
||||
_local_source_id,
|
||||
_local_source_revision_id,
|
||||
_read_selected_file,
|
||||
_remote_factory,
|
||||
_selected_media_type,
|
||||
create_connector_server,
|
||||
upload_selected_file,
|
||||
)
|
||||
|
||||
|
||||
def test_connector_uploads_allowlisted_bytes_without_exposing_them_to_local_tools(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
async def exercise() -> None:
|
||||
source = tmp_path / "selected.md"
|
||||
payload = b"# Private source\n" + b"Exact source bytes stay in the connector.\n" * 20_000
|
||||
source.write_bytes(payload)
|
||||
backend = MCPServer("test GEWU backend")
|
||||
chunks: list[bytes] = []
|
||||
interrupted = False
|
||||
extraction_interrupted = False
|
||||
expected_digest = hashlib.sha256(payload).hexdigest()
|
||||
|
||||
@backend.tool(name="source_upload_begin", structured_output=False)
|
||||
def begin(
|
||||
project_id: str,
|
||||
filename: str,
|
||||
media_type: str,
|
||||
byte_length: int,
|
||||
expected_sha256: str,
|
||||
) -> str:
|
||||
assert (project_id, filename, media_type, byte_length, expected_sha256) == (
|
||||
"dome-performance",
|
||||
"selected.md",
|
||||
"text/markdown",
|
||||
len(payload),
|
||||
expected_digest,
|
||||
)
|
||||
return json.dumps({"upload_id": "opaque-upload-id"})
|
||||
|
||||
@backend.tool(name="source_upload_chunk", structured_output=False)
|
||||
def chunk(upload_id: str, offset: int, content_base64: str, chunk_sha256: str) -> str:
|
||||
nonlocal interrupted
|
||||
decoded = base64.b64decode(content_base64, validate=True)
|
||||
assert upload_id == "opaque-upload-id"
|
||||
assert hashlib.sha256(decoded).hexdigest() == chunk_sha256
|
||||
assert offset == sum(map(len, chunks))
|
||||
if offset > 0 and not interrupted:
|
||||
interrupted = True
|
||||
raise ValueError("simulated interrupted transfer")
|
||||
chunks.append(decoded)
|
||||
return json.dumps({"upload_id": upload_id})
|
||||
|
||||
@backend.tool(name="source_upload_status", structured_output=False)
|
||||
def status(upload_id: str) -> str:
|
||||
assert upload_id == "opaque-upload-id"
|
||||
return json.dumps({"verified_offsets": [0] if chunks else []})
|
||||
|
||||
@backend.tool(name="source_upload_finalize", structured_output=False)
|
||||
def finalize(upload_id: str) -> str:
|
||||
assert upload_id == "opaque-upload-id"
|
||||
return json.dumps(
|
||||
{"sha256": hashlib.sha256(b"".join(chunks)).hexdigest(), "status": "staged"}
|
||||
)
|
||||
|
||||
@backend.tool(name="source_extract", structured_output=False)
|
||||
def extract(project_id: str, artifact_sha256: str, source_revision_id: str) -> str:
|
||||
nonlocal extraction_interrupted
|
||||
assert project_id == "dome-performance"
|
||||
assert artifact_sha256 == expected_digest
|
||||
assert source_revision_id == "source-r1"
|
||||
if not extraction_interrupted:
|
||||
extraction_interrupted = True
|
||||
raise ValueError("simulated extraction interruption")
|
||||
return json.dumps({"extraction_id": "e" * 64, "blocks": []})
|
||||
|
||||
@asynccontextmanager
|
||||
async def remote_factory() -> AsyncIterator[Client]:
|
||||
async with Client(backend) as remote:
|
||||
yield remote
|
||||
|
||||
state_file = tmp_path / "connector-state.json"
|
||||
selected_files = (SelectedFile("research-note", source),)
|
||||
|
||||
async def connector_call() -> CallToolResult:
|
||||
connector_after_restart = create_connector_server(
|
||||
selected_files, remote_factory, UploadStateStore(state_file)
|
||||
)
|
||||
async with Client(connector_after_restart) as local_after_restart:
|
||||
return await local_after_restart.call_tool(
|
||||
"upload_selected_file",
|
||||
{
|
||||
"file_id": "research-note",
|
||||
"project_id": "dome-performance",
|
||||
"source_revision_id": "source-r1",
|
||||
},
|
||||
)
|
||||
|
||||
connector = create_connector_server(
|
||||
selected_files, remote_factory, UploadStateStore(state_file)
|
||||
)
|
||||
async with Client(connector) as local:
|
||||
tools = await local.list_tools()
|
||||
assert {tool.name for tool in tools.tools} == {
|
||||
"list_selected_files",
|
||||
"upload_selected_file",
|
||||
}
|
||||
listed = await local.call_tool("list_selected_files", {})
|
||||
assert not listed.is_error
|
||||
assert isinstance(listed.content[0], TextContent)
|
||||
assert "research-note" in listed.content[0].text
|
||||
assert str(source) not in listed.content[0].text
|
||||
denied = await local.call_tool(
|
||||
"upload_selected_file",
|
||||
{
|
||||
"file_id": "arbitrary-path",
|
||||
"project_id": "dome-performance",
|
||||
"source_revision_id": "source-r1",
|
||||
},
|
||||
)
|
||||
assert denied.is_error
|
||||
interrupted_result = await connector_call()
|
||||
assert interrupted_result.is_error
|
||||
assert state_file.exists()
|
||||
finalize_interrupted_result = await connector_call()
|
||||
assert finalize_interrupted_result.is_error
|
||||
assert json.loads(state_file.read_text())
|
||||
restarted_result = await connector_call()
|
||||
assert not restarted_result.is_error
|
||||
assert isinstance(restarted_result.content[0], TextContent)
|
||||
receipt = json.loads(restarted_result.content[0].text)
|
||||
assert receipt["sha256"] == expected_digest
|
||||
assert receipt["byte_length"] == len(payload)
|
||||
assert payload[:80].decode() not in restarted_result.content[0].text
|
||||
source_candidate = SourceRevisionRecord.model_validate_json(
|
||||
json.dumps(receipt["source_candidate"]), strict=True
|
||||
)
|
||||
assert source_candidate.archive_status == "archived"
|
||||
assert source_candidate.capture_method == "local-file"
|
||||
assert source_candidate.source_id.startswith("local:")
|
||||
assert source_candidate.revision_id == "source-r1"
|
||||
assert source_candidate.artifact_sha256 == expected_digest
|
||||
assert str(source) not in restarted_result.content[0].text
|
||||
assert b"".join(chunks) == payload
|
||||
state = json.loads(state_file.read_text())
|
||||
assert len(state) == 1
|
||||
assert "upload_id" not in next(iter(state.values()))
|
||||
duplicate_result = await connector_call()
|
||||
assert not duplicate_result.is_error
|
||||
assert json.loads(duplicate_result.content[0].text) == receipt
|
||||
|
||||
asyncio.run(exercise())
|
||||
|
||||
|
||||
def test_connector_rejects_symlinked_selected_paths(tmp_path: Path) -> None:
|
||||
outside = tmp_path / "outside.md"
|
||||
outside.write_text("private")
|
||||
link = tmp_path / "selected.md"
|
||||
link.symlink_to(outside)
|
||||
with pytest.raises(ValueError, match="symbolic links"):
|
||||
_read_selected_file(SelectedFile("selected", link))
|
||||
|
||||
|
||||
def test_local_file_identity_is_path_free_and_revisions_follow_content() -> None:
|
||||
from datetime import UTC, datetime
|
||||
|
||||
captured_at = datetime(2026, 10, 13, tzinfo=UTC)
|
||||
source_id = _local_source_id("research-note")
|
||||
first = _local_source_revision_id(source_id, hashlib.sha256(b"first").hexdigest(), captured_at)
|
||||
second = _local_source_revision_id(
|
||||
source_id, hashlib.sha256(b"second").hexdigest(), captured_at
|
||||
)
|
||||
assert source_id.startswith("local:")
|
||||
assert "/" not in source_id
|
||||
assert first != second
|
||||
assert first.startswith(f"{source_id}:r:")
|
||||
assert len(first) <= 128
|
||||
|
||||
|
||||
def test_selected_file_media_types_are_explicit() -> None:
|
||||
assert _selected_media_type(Path("source.pdf")) == "application/pdf"
|
||||
assert _selected_media_type(Path("source.md")) == "text/markdown"
|
||||
assert _selected_media_type(Path("source.html")) == "text/html"
|
||||
assert _selected_media_type(Path("source.txt")) == "text/plain"
|
||||
assert _selected_media_type(Path("source.vtt")) == "text/vtt"
|
||||
assert _selected_media_type(Path("source.srt")) == "application/x-subrip"
|
||||
with pytest.raises(ValueError, match="not supported"):
|
||||
_selected_media_type(Path("source.exe"))
|
||||
|
||||
|
||||
def test_vtt_connector_candidate_is_marked_as_transcript(tmp_path: Path) -> None:
|
||||
async def exercise() -> None:
|
||||
source = tmp_path / "interview.vtt"
|
||||
content = b"WEBVTT\n\n00:00:01.000 --> 00:00:02.000\nExact words.\n"
|
||||
source.write_bytes(content)
|
||||
digest = hashlib.sha256(content).hexdigest()
|
||||
|
||||
class Remote:
|
||||
async def call_tool(self, name: str, arguments: dict[str, object]) -> object:
|
||||
response: dict[str, object]
|
||||
if name == "source_upload_begin":
|
||||
response = {"upload_id": "transcript-upload"}
|
||||
elif name == "source_upload_chunk":
|
||||
response = {"upload_id": "transcript-upload"}
|
||||
elif name == "source_upload_finalize":
|
||||
response = {"sha256": digest}
|
||||
elif name == "source_extract":
|
||||
response = {"status": "complete", "extraction_id": "e" * 64, "blocks": []}
|
||||
else:
|
||||
raise AssertionError(f"Unexpected remote tool: {name}")
|
||||
return type(
|
||||
"ToolResult",
|
||||
(),
|
||||
{
|
||||
"is_error": False,
|
||||
"content": [TextContent(type="text", text=json.dumps(response))],
|
||||
},
|
||||
)()
|
||||
|
||||
receipt = await upload_selected_file(
|
||||
Remote(), SelectedFile("interview", source), project_id="dome-performance"
|
||||
)
|
||||
candidate = SourceRevisionRecord.model_validate_json(
|
||||
json.dumps(receipt["source_candidate"]), strict=True
|
||||
)
|
||||
assert candidate.capture_method == "transcript"
|
||||
|
||||
asyncio.run(exercise())
|
||||
|
||||
|
||||
def test_local_connector_returns_archived_candidate_when_pdf_extraction_is_unavailable(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
async def exercise() -> None:
|
||||
source = tmp_path / "scanned.pdf"
|
||||
content = b"%PDF-1.4 fixture"
|
||||
source.write_bytes(content)
|
||||
digest = hashlib.sha256(content).hexdigest()
|
||||
calls: list[str] = []
|
||||
|
||||
class Remote:
|
||||
async def call_tool(self, name: str, arguments: dict[str, object]) -> object:
|
||||
calls.append(name)
|
||||
response: dict[str, object]
|
||||
if name == "source_upload_begin":
|
||||
response = {"upload_id": "active-upload"}
|
||||
elif name == "source_upload_chunk":
|
||||
assert arguments["chunk_sha256"] == digest
|
||||
response = {"upload_id": "active-upload"}
|
||||
elif name == "source_upload_finalize":
|
||||
response = {"sha256": digest}
|
||||
elif name == "source_extract":
|
||||
response = {
|
||||
"status": "unavailable",
|
||||
"error_code": "PDF_NO_TEXT",
|
||||
"blocks": [],
|
||||
}
|
||||
else:
|
||||
raise AssertionError(f"Unexpected remote tool: {name}")
|
||||
return type(
|
||||
"ToolResult",
|
||||
(),
|
||||
{
|
||||
"is_error": False,
|
||||
"content": [TextContent(type="text", text=json.dumps(response))],
|
||||
},
|
||||
)()
|
||||
|
||||
state = UploadStateStore(tmp_path / "state.json")
|
||||
first = await upload_selected_file(
|
||||
Remote(),
|
||||
SelectedFile("scanned-paper", source),
|
||||
project_id="dome-performance",
|
||||
state_store=state,
|
||||
)
|
||||
second = await upload_selected_file(
|
||||
Remote(),
|
||||
SelectedFile("scanned-paper", source),
|
||||
project_id="dome-performance",
|
||||
state_store=state,
|
||||
)
|
||||
|
||||
candidate = SourceRevisionRecord.model_validate_json(
|
||||
json.dumps(first["source_candidate"]), strict=True
|
||||
)
|
||||
assert candidate.archive_status == "archived"
|
||||
assert candidate.artifact_sha256 == digest
|
||||
assert len(candidate.revision_id) <= 128
|
||||
assert first["extraction_status"] == "unavailable"
|
||||
assert first["extraction_error_code"] == "PDF_NO_TEXT"
|
||||
assert second["source_candidate"] == first["source_candidate"]
|
||||
assert calls.count("source_upload_begin") == 1
|
||||
assert calls.count("source_extract") == 2
|
||||
state_record = next(iter(json.loads((tmp_path / "state.json").read_text()).values()))
|
||||
assert "completed_receipt" not in state_record
|
||||
assert str(source) not in json.dumps(first)
|
||||
|
||||
asyncio.run(exercise())
|
||||
|
||||
|
||||
def test_connector_rejects_cleartext_remote_service_urls() -> None:
|
||||
with pytest.raises(ValueError, match="HTTPS"):
|
||||
_remote_factory("http://gewu.example/mcp", "gwu_test")
|
||||
_remote_factory("https://gewu.example/mcp", "gwu_test")
|
||||
_remote_factory("http://localhost:8000/mcp", "gwu_test")
|
||||
|
||||
|
||||
def test_remote_factory_verifies_service_certificate_before_authentication(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
import gewu.transports.stdio_connector as connector_module
|
||||
|
||||
observed: dict[str, object] = {}
|
||||
|
||||
class HttpClient:
|
||||
def __init__(self, **kwargs: object) -> None:
|
||||
observed.update(kwargs)
|
||||
|
||||
async def __aenter__(self) -> object:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: object) -> None:
|
||||
return None
|
||||
|
||||
class RemoteClient:
|
||||
def __init__(self, transport: object) -> None:
|
||||
self.transport = transport
|
||||
|
||||
async def __aenter__(self) -> object:
|
||||
await self.transport.__aenter__()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *args: object) -> None:
|
||||
await self.transport.__aexit__(*args)
|
||||
return None
|
||||
|
||||
@asynccontextmanager
|
||||
async def transport_factory(url: str, *, http_client: object):
|
||||
observed["url"] = url
|
||||
observed["http_client"] = http_client
|
||||
yield object()
|
||||
|
||||
monkeypatch.setattr(connector_module.httpx2, "AsyncClient", HttpClient)
|
||||
monkeypatch.setattr(connector_module, "Client", RemoteClient)
|
||||
monkeypatch.setattr(connector_module, "streamable_http_client", transport_factory)
|
||||
|
||||
async def exercise() -> None:
|
||||
async with _remote_factory("https://gewu.example/mcp", "secret-token")() as remote:
|
||||
assert isinstance(remote, RemoteClient)
|
||||
|
||||
asyncio.run(exercise())
|
||||
assert observed["verify"] is True
|
||||
assert observed["headers"] == {"Authorization": "Bearer secret-token"}
|
||||
assert observed["url"] == "https://gewu.example/mcp"
|
||||
assert "secret-token" not in str(observed["url"])
|
||||
Reference in new issue
Block a user