chore: tighten mcp source typing

This commit is contained in:
lda
2026-06-09 07:03:36 +07:00 Verified
parent 19d65be77c
commit 5ebe1451c7
19 changed files with 103 additions and 43 deletions
+4
View File
@@ -4,6 +4,10 @@
Python baseline is 3.14 (`requires-python = ">=3.14"`). Python 3.14 syntax is allowed; do not "fix" valid new syntax just because it looks unusual. Python baseline is 3.14 (`requires-python = ">=3.14"`). Python 3.14 syntax is allowed; do not "fix" valid new syntax just because it looks unusual.
Example new syntax:
- Parentheses-Free Exceptions (PEP 758) <!-- coderabbit -->
## extra fields ## extra fields
prefer asserts actual['field'] == expected['field'] over assert actual == expected unless we know better (eg. no extra fields allowed) prefer asserts actual['field'] == expected['field'] over assert actual == expected unless we know better (eg. no extra fields allowed)
+1 -1
View File
@@ -282,7 +282,7 @@ def _rpc_timeout_from_optional_config(
return override return override
try: try:
config = load_workflow_config(path) config = load_workflow_config(path)
except (FileNotFoundError, json.JSONDecodeError, ValidationError): except FileNotFoundError, json.JSONDecodeError, ValidationError:
return 30.0 return 30.0
target = config.client.target target = config.client.target
if isinstance(target, RpcHttpTargetConfig): if isinstance(target, RpcHttpTargetConfig):
+6 -1
View File
@@ -12,4 +12,9 @@ from wf_sources_mcp.adapters import (
require_adapter, require_adapter,
) )
__all__ = ["AdapterLookupRef", "LegacyAdapterRef", "SourceAdapterRef", "require_adapter"] __all__ = [
"AdapterLookupRef",
"LegacyAdapterRef",
"SourceAdapterRef",
"require_adapter",
]
+6 -1
View File
@@ -59,4 +59,9 @@ def require_adapter(
return adapter return adapter
__all__ = ["AdapterLookupRef", "LegacyAdapterRef", "SourceAdapterRef", "require_adapter"] __all__ = [
"AdapterLookupRef",
"LegacyAdapterRef",
"SourceAdapterRef",
"require_adapter",
]
+24 -4
View File
@@ -8,9 +8,11 @@ from mcp.types import (
CallToolResult, CallToolResult,
ClientNotification, ClientNotification,
ClientRequest, ClientRequest,
GetPromptResult,
ListPromptsResult, ListPromptsResult,
ListResourcesResult, ListResourcesResult,
ListToolsResult, ListToolsResult,
ReadResourceResult,
) )
from pydantic import AnyUrl from pydantic import AnyUrl
@@ -26,28 +28,36 @@ if TYPE_CHECKING:
class McpClientSession(Protocol): class McpClientSession(Protocol):
"""Subset of MCP SDK ClientSession operations used by source clients.""" """Subset of MCP SDK ClientSession operations used by source clients.
This protocol targets the low-level MCP SDK ``ClientSession`` shape, not
``fastmcp.client.Client``. FastMCP exposes higher-level convenience methods
with different return types; if we use it here later, wrap it in an adapter
instead of pretending it satisfies this session protocol.
"""
# ClientSession stuff. we dont even use their Pagination system...
async def list_tools(self) -> ListToolsResult: ... async def list_tools(self) -> ListToolsResult: ...
async def list_resources(self) -> ListResourcesResult: ... async def list_resources(self) -> ListResourcesResult: ...
async def list_prompts(self) -> ListPromptsResult: ... async def list_prompts(self) -> ListPromptsResult: ...
async def read_resource(self, uri: AnyUrl) -> Any: ... async def read_resource(self, uri: AnyUrl) -> ReadResourceResult: ...
async def get_prompt( async def get_prompt(
self, self,
name: str, name: str,
arguments: dict[str, str] | None = None, arguments: dict[str, str] | None = None,
/, /,
) -> Any: ... ) -> GetPromptResult: ...
# BaseSession stuff. not even complete signature, thats crazy
async def send_request( async def send_request(
self, self,
request: ClientRequest, request: ClientRequest,
result_type: type[ClientResult], result_type: type[ClientResult],
) -> Any: ... ) -> ClientResult: ...
async def send_notification(self, notification: ClientNotification) -> None: ... async def send_notification(self, notification: ClientNotification) -> None: ...
@@ -59,6 +69,16 @@ class McpClientSession(Protocol):
) -> CallToolResult: ... ) -> CallToolResult: ...
if TYPE_CHECKING:
from mcp.client.session import ClientSession as SdkClientSession
def _typecheck_sdk_client_session(
session: SdkClientSession,
) -> McpClientSession:
"""Static-only guard: MCP SDK ClientSession must satisfy our subset."""
return session
@dataclass(slots=True) @dataclass(slots=True)
class McpSourceClient: class McpSourceClient:
"""Operation facade over an initialized MCP SDK ClientSession. """Operation facade over an initialized MCP SDK ClientSession.
+2 -6
View File
@@ -131,9 +131,7 @@ def _transport_from_connection_metadata(
args=tuple(str(arg) for arg in cast("tuple[object, ...]", args_raw)), args=tuple(str(arg) for arg in cast("tuple[object, ...]", args_raw)),
env={ env={
str(key): str(value) str(key): str(value)
for key, value in cast( for key, value in cast("dict[str, object]", env_raw).items()
"dict[str, object]", env_raw
).items()
}, },
cwd=( cwd=(
str(connection.metadata["cwd"]) str(connection.metadata["cwd"])
@@ -148,9 +146,7 @@ def _transport_from_connection_metadata(
url=url if isinstance(url, str) else str(url), # type: ignore[arg-type] url=url if isinstance(url, str) else str(url), # type: ignore[arg-type]
headers={ headers={
str(key): str(value) str(key): str(value)
for key, value in cast( for key, value in cast("dict[str, object]", headers_raw).items()
"dict[str, object]", headers_raw
).items()
}, },
) )
raise ValueError( raise ValueError(
+3 -1
View File
@@ -36,7 +36,9 @@ async def discover_connection_capabilities(
resources = await _list_optional_capabilities( resources = await _list_optional_capabilities(
lambda: adapter.list_resources(connection, auth) lambda: adapter.list_resources(connection, auth)
) )
prompts = await _list_optional_capabilities(lambda: adapter.list_prompts(connection, auth)) prompts = await _list_optional_capabilities(
lambda: adapter.list_prompts(connection, auth)
)
metadata = await adapter.get_connection_metadata(connection, auth) metadata = await adapter.get_connection_metadata(connection, auth)
return DiscoveredConnectionCapabilities( return DiscoveredConnectionCapabilities(
tools=tools, tools=tools,
+2 -6
View File
@@ -117,9 +117,7 @@ def test_artifact_delete_confirmed_succeeds(monkeypatch) -> None:
lambda _ctx: _Context(handlers=handlers), lambda _ctx: _Context(handlers=handlers),
) )
result = CliRunner().invoke( result = CliRunner().invoke(app, ["artifact", "delete", "echo", "1", "--confirm"])
app, ["artifact", "delete", "echo", "1", "--confirm"]
)
assert result.exit_code == 0, result.output assert result.exit_code == 0, result.output
payload = json.loads(result.output) payload = json.loads(result.output)
@@ -136,9 +134,7 @@ def test_artifact_delete_blocked_returns_blocker_ids(monkeypatch) -> None:
lambda _ctx: _BlockedContext(handlers=handlers), lambda _ctx: _BlockedContext(handlers=handlers),
) )
result = CliRunner().invoke( result = CliRunner().invoke(app, ["artifact", "delete", "echo", "1", "--confirm"])
app, ["artifact", "delete", "echo", "1", "--confirm"]
)
assert result.exit_code == 0, result.output assert result.exit_code == 0, result.output
payload = json.loads(result.output) payload = json.loads(result.output)
+6 -2
View File
@@ -290,7 +290,9 @@ def _patch_rpc_client_to_server(monkeypatch, server) -> None:
), ),
) )
monkeypatch.setattr(cli_context, "rpc_client_from_target", fake_rpc_client_from_target) monkeypatch.setattr(
cli_context, "rpc_client_from_target", fake_rpc_client_from_target
)
def test_wf_cap_commands_use_rpc_url_override(monkeypatch, tmp_path) -> None: def test_wf_cap_commands_use_rpc_url_override(monkeypatch, tmp_path) -> None:
@@ -742,7 +744,9 @@ def test_wf_draft_delete_succeeds_with_confirm(monkeypatch, tmp_path) -> None:
], ],
) )
result = runner.invoke(app, [*base_args, "draft", "delete", "delete-me", "--confirm"]) result = runner.invoke(
app, [*base_args, "draft", "delete", "delete-me", "--confirm"]
)
assert result.exit_code == 0, result.output assert result.exit_code == 0, result.output
payload = json.loads(result.output) payload = json.loads(result.output)
-1
View File
@@ -13,7 +13,6 @@ from wf_mcp.auth import (
neutral_auth_from_mcp, neutral_auth_from_mcp,
) )
from wf_mcp.models import AuthRecord as McpAuthRecord from wf_mcp.models import AuthRecord as McpAuthRecord
from wf_mcp.models import ConnectionConfig
from wf_mcp.storage import FileStore from wf_mcp.storage import FileStore
from wf_sources_mcp.connections import McpSourceConnection from wf_sources_mcp.connections import McpSourceConnection
from wf_sources_mcp.transports import StdioSourceTransport from wf_sources_mcp.transports import StdioSourceTransport
+5 -1
View File
@@ -40,7 +40,11 @@ class _ToolsOnlyAdapter:
raise McpError(ErrorData(code=-32601, message="Method not found")) raise McpError(ErrorData(code=-32601, message="Method not found"))
async def get_connection_metadata(self, connection, auth): async def get_connection_metadata(self, connection, auth):
return {"server": getattr(connection, "provider", getattr(connection, "server", None))} return {
"server": getattr(
connection, "provider", getattr(connection, "server", None)
)
}
async def read_resource(self, connection, auth, uri): async def read_resource(self, connection, auth, uri):
raise NotImplementedError raise NotImplementedError
+4 -2
View File
@@ -12,7 +12,7 @@ from wf_core import RuntimeContext
from wf_core.models.steps import InputPathBinding, OutputBinding from wf_core.models.steps import InputPathBinding, OutputBinding
from wf_core.paths import GraphSourcePath, LocalPath, StatePath from wf_core.paths import GraphSourcePath, LocalPath, StatePath
from wf_mcp.capabilities import DiscoveredPrompt, DiscoveredResource, DiscoveredTool from wf_mcp.capabilities import DiscoveredPrompt, DiscoveredResource, DiscoveredTool
from wf_mcp.models import AuthRecord, ConnectionConfig from wf_mcp.models import ConnectionConfig
from wf_mcp.sdk import ToolCallResult from wf_mcp.sdk import ToolCallResult
@@ -186,7 +186,9 @@ class FakeAdapter:
auth, auth,
) -> dict[str, Any]: ) -> dict[str, Any]:
return { return {
"server": getattr(connection, "provider", getattr(connection, "server", None)), "server": getattr(
connection, "provider", getattr(connection, "server", None)
),
"account": connection.account, "account": connection.account,
"auth_scheme": auth.scheme if auth is not None else None, "auth_scheme": auth.scheme if auth is not None else None,
} }
+5 -2
View File
@@ -14,7 +14,6 @@ from wf_artifacts import (
from wf_authoring import node, reducer from wf_authoring import node, reducer
from wf_mcp.broker import WfMcpService from wf_mcp.broker import WfMcpService
from wf_mcp.capabilities import DiscoveredTool from wf_mcp.capabilities import DiscoveredTool
from wf_mcp.models import AuthRecord, ConnectionConfig
from wf_mcp.sdk import ToolCallResult from wf_mcp.sdk import ToolCallResult
from wf_mcp.storage import FileStore from wf_mcp.storage import FileStore
from wf_mcp.workflow_surface import WorkflowSurfaceHandlers from wf_mcp.workflow_surface import WorkflowSurfaceHandlers
@@ -104,7 +103,11 @@ class ContentOnlyOutputAdapter:
connection, connection,
auth, auth,
) -> dict[str, Any]: ) -> dict[str, Any]:
return {"server": getattr(connection, "provider", getattr(connection, "server", None))} return {
"server": getattr(
connection, "provider", getattr(connection, "server", None)
)
}
async def call_tool( async def call_tool(
self, self,
+10 -2
View File
@@ -104,7 +104,11 @@ def test_combined_catalog_sorts_entries_and_serializes_payload() -> None:
first = snapshot_from_specs( first = snapshot_from_specs(
"zeta.default", "zeta.default",
specs={"echo": _echo_spec()}, specs={"echo": _echo_spec()},
resources=[DiscoveredResource(uri="zeta://guide", name="guide", title=None, description=None)], resources=[
DiscoveredResource(
uri="zeta://guide", name="guide", title=None, description=None
)
],
prompts=[DiscoveredPrompt(name="prompt", title=None, description=None)], prompts=[DiscoveredPrompt(name="prompt", title=None, description=None)],
metadata={"order": "second"}, metadata={"order": "second"},
fetched_at_epoch_ms=2, fetched_at_epoch_ms=2,
@@ -113,7 +117,11 @@ def test_combined_catalog_sorts_entries_and_serializes_payload() -> None:
second = snapshot_from_specs( second = snapshot_from_specs(
"alpha.default", "alpha.default",
specs={"echo": _echo_spec()}, specs={"echo": _echo_spec()},
resources=[DiscoveredResource(uri="alpha://guide", name="guide", title=None, description=None)], resources=[
DiscoveredResource(
uri="alpha://guide", name="guide", title=None, description=None
)
],
prompts=[DiscoveredPrompt(name="prompt", title=None, description=None)], prompts=[DiscoveredPrompt(name="prompt", title=None, description=None)],
metadata={"order": "first"}, metadata={"order": "first"},
fetched_at_epoch_ms=1, fetched_at_epoch_ms=1,
@@ -6,6 +6,7 @@ from dataclasses import dataclass
from typing import Any from typing import Any
import pytest import pytest
from pydantic import AnyHttpUrl
from wf_sources_mcp.auth import AuthRecord from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.client.transport import open_mcp_session from wf_sources_mcp.client.transport import open_mcp_session
@@ -35,9 +36,7 @@ async def _fake_streamable_http_client(
@asynccontextmanager @asynccontextmanager
async def _fake_client_session( async def _fake_client_session(read: Any, write: Any) -> AsyncIterator[_FakeSession]:
read: Any, write: Any
) -> AsyncIterator[_FakeSession]:
yield _FakeSession() yield _FakeSession()
@@ -70,7 +69,7 @@ def _http_connection(
provider="test", provider="test",
account="server", account="server",
transport=HttpSourceTransport( transport=HttpSourceTransport(
url=url, url=AnyHttpUrl(url),
headers=headers or {}, headers=headers or {},
), ),
) )
@@ -90,6 +89,7 @@ async def test_stdio_session_initializes_before_yielding() -> None:
connection = _stdio_connection() connection = _stdio_connection()
async with open_mcp_session(connection, None) as session: async with open_mcp_session(connection, None) as session:
assert isinstance(session, _FakeSession)
assert session.initialized is True assert session.initialized is True
@@ -116,6 +116,7 @@ async def test_stdio_env_merges_transport_and_auth_wins_on_duplicate() -> None:
mod.stdio_client = _capturing_stdio_client # type: ignore[assignment] mod.stdio_client = _capturing_stdio_client # type: ignore[assignment]
async with open_mcp_session(connection, auth) as session: async with open_mcp_session(connection, auth) as session:
assert isinstance(session, _FakeSession)
assert session.initialized is True assert session.initialized is True
params = captured_params[0] params = captured_params[0]
@@ -142,6 +143,7 @@ async def test_stdio_cwd_propagated_to_server_parameters() -> None:
mod.stdio_client = _capturing_stdio_client # type: ignore[assignment] mod.stdio_client = _capturing_stdio_client # type: ignore[assignment]
async with open_mcp_session(connection, None) as session: async with open_mcp_session(connection, None) as session:
assert isinstance(session, _FakeSession)
assert session.initialized is True assert session.initialized is True
assert captured_params[0].cwd == "/workspace" assert captured_params[0].cwd == "/workspace"
@@ -152,6 +154,7 @@ async def test_http_session_initializes_before_yielding() -> None:
connection = _http_connection() connection = _http_connection()
async with open_mcp_session(connection, None) as session: async with open_mcp_session(connection, None) as session:
assert isinstance(session, _FakeSession)
assert session.initialized is True assert session.initialized is True
@@ -184,6 +187,7 @@ async def test_http_auth_headers_passed_to_client(
monkeypatch.setattr(mod, "httpx", _PatchedHttpx()) monkeypatch.setattr(mod, "httpx", _PatchedHttpx())
async with open_mcp_session(connection, auth) as session: async with open_mcp_session(connection, auth) as session:
assert isinstance(session, _FakeSession)
assert session.initialized is True assert session.initialized is True
assert len(captured_clients) == 1 assert len(captured_clients) == 1
+1 -3
View File
@@ -48,9 +48,7 @@ def test_stdio_source_transport_is_typed() -> None:
def test_http_source_transport_is_typed() -> None: def test_http_source_transport_is_typed() -> None:
transport = HttpSourceTransport.model_validate( transport = HttpSourceTransport.model_validate({"url": "http://127.0.0.1:8000/mcp"})
{"url": "http://127.0.0.1:8000/mcp"}
)
assert transport.kind == "http" assert transport.kind == "http"
assert str(transport.url) == "http://127.0.0.1:8000/mcp" assert str(transport.url) == "http://127.0.0.1:8000/mcp"
+9 -3
View File
@@ -155,7 +155,9 @@ class _BrokenResourceAdapter(_Adapter):
raise RuntimeError("resource listing broke") raise RuntimeError("resource listing broke")
async def test_discover_connection_capabilities_collects_all_capability_families() -> None: async def test_discover_connection_capabilities_collects_all_capability_families() -> (
None
):
adapter = _Adapter() adapter = _Adapter()
connection = _connection() connection = _connection()
@@ -172,7 +174,9 @@ async def test_discover_connection_capabilities_collects_all_capability_families
assert adapter.seen_connections == [connection] assert adapter.seen_connections == [connection]
async def test_discover_connection_capabilities_treats_missing_optional_families_as_empty() -> None: async def test_discover_connection_capabilities_treats_missing_optional_families_as_empty() -> (
None
):
capabilities = await discover_connection_capabilities( capabilities = await discover_connection_capabilities(
connection=_connection(), connection=_connection(),
auth=None, auth=None,
@@ -184,7 +188,9 @@ async def test_discover_connection_capabilities_treats_missing_optional_families
assert capabilities.prompts == [] assert capabilities.prompts == []
async def test_discover_connection_capabilities_reraises_non_method_not_found_errors() -> None: async def test_discover_connection_capabilities_reraises_non_method_not_found_errors() -> (
None
):
with pytest.raises(RuntimeError, match="resource listing broke"): with pytest.raises(RuntimeError, match="resource listing broke"):
await discover_connection_capabilities( await discover_connection_capabilities(
connection=_connection(), connection=_connection(),
+3 -1
View File
@@ -165,7 +165,9 @@ def test_connection_config_to_registry_entry_accepts_flat_http_metadata() -> Non
def test_connection_config_to_registry_entry_requires_transport_metadata() -> None: def test_connection_config_to_registry_entry_requires_transport_metadata() -> None:
connection = _LegacyConnectionLike(id="github.work", server="github", account="work") connection = _LegacyConnectionLike(
id="github.work", server="github", account="work"
)
with pytest.raises(ValueError, match="requires metadata.transport"): with pytest.raises(ValueError, match="requires metadata.transport"):
connection_config_to_registry_entry(connection) connection_config_to_registry_entry(connection)
@@ -40,7 +40,9 @@ def test_rpc_client_mixins_share_one_call_contract() -> None:
for module_info in pkgutil.iter_modules(rpc_client_package.__path__): for module_info in pkgutil.iter_modules(rpc_client_package.__path__):
if module_info.name in {"__init__", "base"}: if module_info.name in {"__init__", "base"}:
continue continue
module = importlib.import_module(f"wf_transport_rpc_http.client.{module_info.name}") module = importlib.import_module(
f"wf_transport_rpc_http.client.{module_info.name}"
)
for _name, value in inspect.getmembers(module, inspect.isclass): for _name, value in inspect.getmembers(module, inspect.isclass):
if value.__module__ == module.__name__: if value.__module__ == module.__name__:
assert "_call" not in value.__dict__ assert "_call" not in value.__dict__