374 lines
15 KiB
Python
374 lines
15 KiB
Python
|
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"])
|