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.
Example new syntax:
- Parentheses-Free Exceptions (PEP 758) <!-- coderabbit -->
## extra fields
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
try:
config = load_workflow_config(path)
except (FileNotFoundError, json.JSONDecodeError, ValidationError):
except FileNotFoundError, json.JSONDecodeError, ValidationError:
return 30.0
target = config.client.target
if isinstance(target, RpcHttpTargetConfig):
+6 -1
View File
@@ -12,4 +12,9 @@ from wf_sources_mcp.adapters import (
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
__all__ = ["AdapterLookupRef", "LegacyAdapterRef", "SourceAdapterRef", "require_adapter"]
__all__ = [
"AdapterLookupRef",
"LegacyAdapterRef",
"SourceAdapterRef",
"require_adapter",
]
+25 -5
View File
@@ -8,9 +8,11 @@ from mcp.types import (
CallToolResult,
ClientNotification,
ClientRequest,
GetPromptResult,
ListPromptsResult,
ListResourcesResult,
ListToolsResult,
ReadResourceResult,
)
from pydantic import AnyUrl
@@ -21,33 +23,41 @@ if TYPE_CHECKING:
DiscoveredPrompt,
DiscoveredResource,
DiscoveredTool,
)
)
from wf_sources_mcp.sdk.protocols import ToolCallResult
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_resources(self) -> ListResourcesResult: ...
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(
self,
name: str,
arguments: dict[str, str] | None = None,
/,
) -> Any: ...
) -> GetPromptResult: ...
# BaseSession stuff. not even complete signature, thats crazy
async def send_request(
self,
request: ClientRequest,
result_type: type[ClientResult],
) -> Any: ...
) -> ClientResult: ...
async def send_notification(self, notification: ClientNotification) -> None: ...
@@ -59,6 +69,16 @@ class McpClientSession(Protocol):
) -> 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)
class McpSourceClient:
"""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)),
env={
str(key): str(value)
for key, value in cast(
"dict[str, object]", env_raw
).items()
for key, value in cast("dict[str, object]", env_raw).items()
},
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]
headers={
str(key): str(value)
for key, value in cast(
"dict[str, object]", headers_raw
).items()
for key, value in cast("dict[str, object]", headers_raw).items()
},
)
raise ValueError(
+3 -1
View File
@@ -36,7 +36,9 @@ async def discover_connection_capabilities(
resources = await _list_optional_capabilities(
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)
return DiscoveredConnectionCapabilities(
tools=tools,
+2 -6
View File
@@ -117,9 +117,7 @@ def test_artifact_delete_confirmed_succeeds(monkeypatch) -> None:
lambda _ctx: _Context(handlers=handlers),
)
result = CliRunner().invoke(
app, ["artifact", "delete", "echo", "1", "--confirm"]
)
result = CliRunner().invoke(app, ["artifact", "delete", "echo", "1", "--confirm"])
assert result.exit_code == 0, 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),
)
result = CliRunner().invoke(
app, ["artifact", "delete", "echo", "1", "--confirm"]
)
result = CliRunner().invoke(app, ["artifact", "delete", "echo", "1", "--confirm"])
assert result.exit_code == 0, 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:
@@ -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
payload = json.loads(result.output)
-1
View File
@@ -13,7 +13,6 @@ from wf_mcp.auth import (
neutral_auth_from_mcp,
)
from wf_mcp.models import AuthRecord as McpAuthRecord
from wf_mcp.models import ConnectionConfig
from wf_mcp.storage import FileStore
from wf_sources_mcp.connections import McpSourceConnection
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"))
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):
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.paths import GraphSourcePath, LocalPath, StatePath
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
@@ -186,7 +186,9 @@ class FakeAdapter:
auth,
) -> dict[str, Any]:
return {
"server": getattr(connection, "provider", getattr(connection, "server", None)),
"server": getattr(
connection, "provider", getattr(connection, "server", None)
),
"account": connection.account,
"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_mcp.broker import WfMcpService
from wf_mcp.capabilities import DiscoveredTool
from wf_mcp.models import AuthRecord, ConnectionConfig
from wf_mcp.sdk import ToolCallResult
from wf_mcp.storage import FileStore
from wf_mcp.workflow_surface import WorkflowSurfaceHandlers
@@ -104,7 +103,11 @@ class ContentOnlyOutputAdapter:
connection,
auth,
) -> 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(
self,
+10 -2
View File
@@ -104,7 +104,11 @@ def test_combined_catalog_sorts_entries_and_serializes_payload() -> None:
first = snapshot_from_specs(
"zeta.default",
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)],
metadata={"order": "second"},
fetched_at_epoch_ms=2,
@@ -113,7 +117,11 @@ def test_combined_catalog_sorts_entries_and_serializes_payload() -> None:
second = snapshot_from_specs(
"alpha.default",
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)],
metadata={"order": "first"},
fetched_at_epoch_ms=1,
@@ -6,6 +6,7 @@ from dataclasses import dataclass
from typing import Any
import pytest
from pydantic import AnyHttpUrl
from wf_sources_mcp.auth import AuthRecord
from wf_sources_mcp.client.transport import open_mcp_session
@@ -35,9 +36,7 @@ async def _fake_streamable_http_client(
@asynccontextmanager
async def _fake_client_session(
read: Any, write: Any
) -> AsyncIterator[_FakeSession]:
async def _fake_client_session(read: Any, write: Any) -> AsyncIterator[_FakeSession]:
yield _FakeSession()
@@ -70,7 +69,7 @@ def _http_connection(
provider="test",
account="server",
transport=HttpSourceTransport(
url=url,
url=AnyHttpUrl(url),
headers=headers or {},
),
)
@@ -90,6 +89,7 @@ async def test_stdio_session_initializes_before_yielding() -> None:
connection = _stdio_connection()
async with open_mcp_session(connection, None) as session:
assert isinstance(session, _FakeSession)
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]
async with open_mcp_session(connection, auth) as session:
assert isinstance(session, _FakeSession)
assert session.initialized is True
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]
async with open_mcp_session(connection, None) as session:
assert isinstance(session, _FakeSession)
assert session.initialized is True
assert captured_params[0].cwd == "/workspace"
@@ -152,6 +154,7 @@ async def test_http_session_initializes_before_yielding() -> None:
connection = _http_connection()
async with open_mcp_session(connection, None) as session:
assert isinstance(session, _FakeSession)
assert session.initialized is True
@@ -184,6 +187,7 @@ async def test_http_auth_headers_passed_to_client(
monkeypatch.setattr(mod, "httpx", _PatchedHttpx())
async with open_mcp_session(connection, auth) as session:
assert isinstance(session, _FakeSession)
assert session.initialized is True
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:
transport = HttpSourceTransport.model_validate(
{"url": "http://127.0.0.1:8000/mcp"}
)
transport = HttpSourceTransport.model_validate({"url": "http://127.0.0.1:8000/mcp"})
assert transport.kind == "http"
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")
async def test_discover_connection_capabilities_collects_all_capability_families() -> None:
async def test_discover_connection_capabilities_collects_all_capability_families() -> (
None
):
adapter = _Adapter()
connection = _connection()
@@ -172,7 +174,9 @@ async def test_discover_connection_capabilities_collects_all_capability_families
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(
connection=_connection(),
auth=None,
@@ -184,7 +188,9 @@ async def test_discover_connection_capabilities_treats_missing_optional_families
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"):
await discover_connection_capabilities(
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:
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"):
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__):
if module_info.name in {"__init__", "base"}:
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):
if value.__module__ == module.__name__:
assert "_call" not in value.__dict__