mcp-client-kit 0.2.0__tar.gz → 0.3.0__tar.gz

{mcp_client_kit-0.2.0 → mcp_client_kit-0.3.0}/CHANGELOG.md

@@ -1,6 +1,40 @@
 # Changelog
 
-## [Unreleased] — 0.3.0
+## [Unreleased] — 0.4.0
+
+## [0.3.0] — 2026-07-14
+
+### Fixed
+
+- **`mcpgen login` — OAuth token exchange failed with `400 invalid_request`** on authorization
+  servers that strictly enforce RFC 6749 §2.3 (`"Client must not use multiple authentication
+  methods"`). Dynamic client registration omitted `token_endpoint_auth_method`, so the server
+  defaulted us to `client_secret_basic` and issued a `client_secret`; the MCP SDK then sent an
+  `Authorization: Basic` header *and* `client_id` in the form body — two client authentication
+  methods in one request. mcpgen now registers as a public client
+  (`token_endpoint_auth_method: "none"`), which is the correct posture for a distributed CLI
+  (RFC 8252 §8.4) and is fully secured by the PKCE (S256) the SDK already performs on every
+  authorization code grant. No security regression: the discarded `client_secret` was issued
+  per-install by dynamic registration and stored in the same `credentials.json` as the refresh
+  token, so anyone who could read the secret could already read the token — it never added a
+  layer of defense.
+
+- **Discriminated-union codegen mistyped `string` discriminators with numeric-looking
+  enum values as `int`.** Discriminator type was inferred solely by sniffing whether
+  shape-spec variant keys looked numeric (e.g. `"1"`, `"2"`), so a genuinely
+  `type: "string"` discriminator with those enum values got a wrong `int` `Literal`
+  in the generated `@overload` stubs. Discriminator type is now resolved from an
+  explicit `input_overrides` entry, then the tool's declared `inputSchema` type,
+  and only falls back to key-sniffing when neither is available.
+
+- **`mcpgen call`/`probe` crashed on `resource_link` content items** — real MCP SDK
+  `ResourceLink.uri` is a pydantic `AnyUrl`, not `str`; `json.dumps` on the summarized
+  content raised `TypeError`. The summary now coerces `uri` to `str`.
+
+- **`mcpgen probe` reported inflated payload sizes for non-ASCII content** — size was
+  measured as the character count of `json.dumps`'s default `ensure_ascii=True`
+  output, which escapes every non-ASCII character to a `\uXXXX` sequence. Size is
+  now measured as real UTF-8 byte length.
 
 ## [0.2.0] — 2026-06-20
 

{mcp_client_kit-0.2.0 → mcp_client_kit-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mcp-client-kit
-Version: 0.2.0
+Version: 0.3.0
 Summary: mcp-client-kit — generate typed Python wrappers for any MCP server (codegen skill + CLI).
 Project-URL: Homepage, https://github.com/svd/mcp-client-kit
 Project-URL: Source, https://github.com/svd/mcp-client-kit

{mcp_client_kit-0.2.0 → mcp_client_kit-0.3.0}/mcpgen/_bridge.py

@@ -58,7 +58,7 @@ from urllib.parse import parse_qs, urlparse
 
 import httpx
 from mcp import ClientSession
-from mcp.client.auth import OAuthClientProvider, TokenStorage
+from mcp.client.auth import OAuthClientProvider, OAuthRegistrationError, TokenStorage
 from mcp.client.stdio import StdioServerParameters, stdio_client
 from mcp.client.streamable_http import streamable_http_client
 from mcp.shared._httpx_utils import create_mcp_http_client
@@ -330,6 +330,44 @@ def _resolve_client_name(server_name: str) -> str:
     return _client_names_cache.get(server_name) or f"mcpgen ({server_name})"
 
 
+def _client_metadata(server_name: str, callback_uri: str, client_name: str | None = None) -> OAuthClientMetadata:
+    """Dynamic-registration metadata (RFC 7591). Shared by both OAuth entry points.
+
+    `token_endpoint_auth_method="none"` is load-bearing. Omit it and the AS applies
+    the RFC 7591 §2 default of `client_secret_basic` and issues a `client_secret`;
+    the SDK then sends an `Authorization: Basic` header *and* `client_id` in the form
+    body — two client authentication methods in one request, which servers that
+    enforce RFC 6749 §2.3 reject with `400 invalid_request`. Registering as a public
+    client (RFC 8252 §8.4) is the correct posture for a distributed CLI anyway, and
+    the SDK's PKCE (S256, unconditional) is what actually secures the flow.
+    """
+    return OAuthClientMetadata(
+        client_name=client_name or _resolve_client_name(server_name),
+        redirect_uris=[callback_uri],  # type: ignore[list-item]  # Pydantic coerces str→AnyUrl
+        grant_types=["authorization_code", "refresh_token"],
+        response_types=["code"],
+        token_endpoint_auth_method="none",
+    )
+
+
+def _explain_registration_error(exc: OAuthRegistrationError) -> OAuthRegistrationError:
+    """Annotate `invalid_client_metadata` with its most likely cause.
+
+    An AS that does not support public clients MUST reject our registration
+    (RFC 7591 §2) rather than downgrade it. No such server is known, so rather than
+    carry a speculative override flag we make the failure legible if one turns up.
+    """
+    if "invalid_client_metadata" not in str(exc):
+        return exc
+    return OAuthRegistrationError(
+        f"{exc}\n\n"
+        "Likely cause: mcpgen registers as a public client "
+        "(token_endpoint_auth_method=none), and this authorization server appears to "
+        "require a client_secret. Please report this at "
+        "https://github.com/svd/mcp-client-kit/issues with the message above."
+    )
+
+
 # Treat a cached token as expired this many seconds before its real expiry.
 _MARGIN = 120
 
@@ -782,6 +820,12 @@ async def _http_session(
     storage = FileTokenStorage(server_name, backend=resolve_cred_backend(cred_backend))
     await _pre_flight_refresh(server_name, storage)
 
+    # Unlike login(), this does NOT clear a stale confidential-client registration,
+    # and must not: a pre-fix credential still refreshes fine here, because
+    # _pre_flight_refresh sends client_id + client_secret in the body — that is
+    # client_secret_post, a *single* auth method, so it never hit the double-auth
+    # bug. Once the refresh token dies, _no_browser raises ReauthenticationRequired
+    # → "Run: mcpgen login" → which does clear and re-register as a public client.
     data = storage._load()
     redirect_uris = data.get(server_name, {}).get("client_info", {}).get("redirect_uris", [])
     callback_uri = redirect_uris[0] if redirect_uris else "http://localhost:0/callback"
@@ -794,12 +838,7 @@ async def _http_session(
 
     provider = OAuthClientProvider(
         server_url=server_url,
-        client_metadata=OAuthClientMetadata(
-            client_name=client_name or _resolve_client_name(server_name),
-            redirect_uris=[callback_uri],  # type: ignore[list-item]  # Pydantic coerces str→AnyUrl
-            grant_types=["authorization_code", "refresh_token"],
-            response_types=["code"],
-        ),
+        client_metadata=_client_metadata(server_name, callback_uri, client_name),
         storage=storage,
         redirect_handler=_no_browser,
         callback_handler=_no_callback,
@@ -951,20 +990,61 @@ class McpBridgeCaller:
             env=self._env,
         ) as s:
             result = await s.call_tool(tool, arguments)
-            content = [{"type": item.type, "text": getattr(item, "text", "")} for item in result.content]
+            content = [_summarize_content_item(item) for item in result.content]
             return parse(content)
 
 
+def _summarize_content_item(item: Any) -> dict:
+    """Reduce one MCP content block to a compact, JSON-safe dict.
+
+    ``text`` blocks keep today's shape ({"type", "text"}). ``image`` / ``resource`` /
+    ``resource_link`` blocks carry no useful `.text` — capturing only `.text` (as the
+    old code did) silently degrades them to an empty string, hiding the payload from
+    shape inference entirely. Record presence/metadata instead (never the raw
+    base64/blob bytes, to keep shape-specs small).
+    """
+    item_type = getattr(item, "type", None)
+    if item_type == "image":
+        return {
+            "type": "image",
+            "mimeType": getattr(item, "mimeType", None),
+            "has_data": bool(getattr(item, "data", None)),
+        }
+    if item_type == "resource":
+        resource = getattr(item, "resource", None)
+        return {
+            "type": "resource",
+            "mimeType": getattr(resource, "mimeType", None),
+            "has_text": bool(getattr(resource, "text", None)),
+            "has_blob": bool(getattr(resource, "blob", None)),
+        }
+    if item_type == "resource_link":
+        uri = getattr(item, "uri", None)
+        return {
+            "type": "resource_link",
+            "uri": str(uri) if uri is not None else None,
+            "name": getattr(item, "name", None),
+        }
+    return {"type": item_type, "text": getattr(item, "text", "")}
+
+
 def parse(content_items: list) -> Any:
     """Extract and JSON-parse the text payload from an MCP tool result.
 
     Falls back to ``ast.literal_eval`` for Python-repr payloads (e.g. servers
     that return single-quoted dicts like sqlite), then to a plain string as a
     last resort so callers can still inspect the response.
+
+    Non-text content blocks (image / resource / resource_link) carry no `.text`
+    to parse — return the block's own summary dict as-is rather than falling
+    through to an empty string, so the shape-spec records something real.
     """
     if not content_items:
         raise ValueError("MCP tool result has empty content")
     item = content_items[0]
+    item_type = item.get("type") if isinstance(item, dict) else getattr(item, "type", None)
+    if item_type in ("image", "resource", "resource_link"):
+        return item if isinstance(item, dict) else _summarize_content_item(item)
     text = item.get("text", "") if isinstance(item, dict) else getattr(item, "text", "")
     try:
         return json.loads(text)
@@ -1067,35 +1147,33 @@ async def login(
 
         provider = OAuthClientProvider(
             server_url=server_url,
-            client_metadata=OAuthClientMetadata(
-                client_name=client_name or _resolve_client_name(server_name),
-                redirect_uris=[callback_uri],  # type: ignore[list-item]  # Pydantic coerces str→AnyUrl
-                grant_types=["authorization_code", "refresh_token"],
-                response_types=["code"],
-            ),
+            client_metadata=_client_metadata(server_name, callback_uri, client_name),
             storage=storage,
             redirect_handler=redirect_handler,
             callback_handler=callback_handler,
         )
 
         try:
-            async with _open_http(server_url, auth=provider) as (read, write, _):
-                async with ClientSession(read, write) as s:
-                    await s.initialize()
-
-                    # Persist the token endpoint for later pre-flight refresh. Do this
-                    # independently of any tool call — not every server exposes `whoami`.
-                    if provider.context.oauth_metadata is not None:
-                        endpoint_url = str(provider.context.oauth_metadata.token_endpoint)
-                        creds_data = storage._load()
-                        creds_data.setdefault(server_name, {})["token_endpoint"] = endpoint_url
-                        storage._save(creds_data)
-
-                    # Confirm the authenticated session works with a server-agnostic
-                    # call. `list_tools` is part of the MCP protocol — every server
-                    # supports it, unlike any specific tool name.
-                    tools = await s.list_tools()
-                    print(f"Login OK ({server_name}); {len(tools.tools)} tool(s) available")
+            try:
+                async with _open_http(server_url, auth=provider) as (read, write, _):
+                    async with ClientSession(read, write) as s:
+                        await s.initialize()
+
+                        # Persist the token endpoint for later pre-flight refresh. Do this
+                        # independently of any tool call — not every server exposes `whoami`.
+                        if provider.context.oauth_metadata is not None:
+                            endpoint_url = str(provider.context.oauth_metadata.token_endpoint)
+                            creds_data = storage._load()
+                            creds_data.setdefault(server_name, {})["token_endpoint"] = endpoint_url
+                            storage._save(creds_data)
+
+                        # Confirm the authenticated session works with a server-agnostic
+                        # call. `list_tools` is part of the MCP protocol — every server
+                        # supports it, unlike any specific tool name.
+                        tools = await s.list_tools()
+                        print(f"Login OK ({server_name}); {len(tools.tools)} tool(s) available")
+            except OAuthRegistrationError as e:
+                raise _explain_registration_error(e) from e
         finally:
             if not callback_future.done():
                 callback_future.set_result((None, None))

{mcp_client_kit-0.2.0 → mcp_client_kit-0.3.0}/mcpgen/cli.py

@@ -48,11 +48,14 @@ async def _list_tools(
         result = await s.list_tools()
     tools = []
     for t in result.tools:
+        raw_annotations = getattr(t, "annotations", None)
+        annotations = raw_annotations.model_dump(exclude_none=True) if raw_annotations else None
         tools.append(
             {
                 "name": t.name,
                 "description": t.description,
                 "inputSchema": t.inputSchema,
+                "annotations": annotations,
             }
         )
     return tools
@@ -70,7 +73,8 @@ async def _probe(
     config_path: str | None = None,
     cred_backend: str | None = None,
     env: dict[str, str] | None = None,
-) -> Any:
+) -> tuple[Any, int]:
+    """Return `(observed_shape, observed_byte_size)` for one live probe call."""
     caller = _bridge.McpBridgeCaller(
         cmd=cmd,
         url=url,
@@ -81,7 +85,8 @@ async def _probe(
         env=env,
     )
     raw = await caller.call(server, tool, args)
-    return codegen.summarize_shape(raw)
+    size = len(json.dumps(raw, default=str, ensure_ascii=False).encode("utf-8"))
+    return codegen.summarize_shape(raw), size
 
 
 async def _call(
@@ -283,7 +288,7 @@ def _cmd_codegen(ns: argparse.Namespace) -> int:
         args = json.loads(ns.probe_args) if ns.probe_args else {}
         print(f"[codegen] probing {ns.probe}({args}) …", file=sys.stderr)
         try:
-            shape = asyncio.run(_probe(ns.server, ns.probe, args, cmd=cmd, **conn))
+            shape, _size = asyncio.run(_probe(ns.server, ns.probe, args, cmd=cmd, **conn))
         except (FileNotFoundError, ValueError) as exc:
             print(f"[codegen] error: {exc}", file=sys.stderr)
             return 1
@@ -327,17 +332,19 @@ def _cmd_probe(ns: argparse.Namespace) -> int:
         env=_parse_env(ns),
     )
     shapes = []
+    sizes = []
     for i, args in enumerate(args_list):
         print(f"[probe]   [{i + 1}/{n}] args={args}", file=sys.stderr)
         # one session per probe (prototype); pooling is out of scope
         try:
-            shape = asyncio.run(_probe(ns.server, ns.tool, args, cmd=cmd, **conn))
+            shape, size = asyncio.run(_probe(ns.server, ns.tool, args, cmd=cmd, **conn))
         except (FileNotFoundError, ValueError) as exc:
             print(f"[probe] error: {exc}", file=sys.stderr)
             return 1
         shapes.append(shape)
+        sizes.append(size)
 
-    skeleton = codegen.probe_skeleton(ns.tool, args_list, shapes)
+    skeleton = codegen.probe_skeleton(ns.tool, args_list, shapes, observed_bytes=sizes)
     out = json.dumps(skeleton, indent=2)
     if ns.emit_shape:
         target = Path(ns.emit_shape)
@@ -470,6 +477,9 @@ def _cmd_merge(ns: argparse.Namespace) -> int:
                 pa = entry.get("probed_args")
                 if pa:  # non-empty dict or non-empty list
                     verify_map[tool_name] = pa
+    # Prune entries for tools no longer part of the merged shapes (e.g. the
+    # server dropped the tool) — otherwise dead entries persist forever.
+    verify_map = {k: v for k, v in verify_map.items() if k in merged}
     if verify_map:
         _atomic_write_text(verify_target, json.dumps(verify_map, indent=2) + "\n")
         print(
@@ -477,6 +487,10 @@ def _cmd_merge(ns: argparse.Namespace) -> int:
             " — ⚠  raw args/PII, git-ignored (verify sidecar)",
             file=sys.stderr,
         )
+    elif verify_target.is_file():
+        # Every entry was pruned as stale — don't leave dead content on disk.
+        verify_target.unlink()
+        print(f"[merge] removed {verify_target} (all entries stale)", file=sys.stderr)
 
     if not ns.keep_parts:
         shutil.rmtree(parts_d)
@@ -576,13 +590,17 @@ def _cmd_list(ns: argparse.Namespace) -> int:
     except (FileNotFoundError, ValueError) as exc:
         print(f"[list] error: {exc}", file=sys.stderr)
         return 1
-    if getattr(ns, "schema", False):
-        out = [
-            {"name": t["name"], "description": t.get("description") or "", "inputSchema": t.get("inputSchema") or {}}
-            for t in tools
-        ]
-    else:
-        out = [{"name": t["name"], "description": t.get("description") or ""} for t in tools]
+
+    def _entry(t: dict) -> dict:
+        entry = {"name": t["name"], "description": t.get("description") or ""}
+        if getattr(ns, "schema", False):
+            entry["inputSchema"] = t.get("inputSchema") or {}
+        annotations = t.get("annotations")
+        if annotations is not None:
+            entry["annotations"] = annotations
+        return entry
+
+    out = [_entry(t) for t in tools]
     sys.stdout.write(json.dumps(out, indent=2) + "\n")
 
     candidates = codegen.detect_discriminators(tools)

{mcp_client_kit-0.2.0 → mcp_client_kit-0.3.0}/mcpgen/codegen.py

@@ -116,6 +116,26 @@ def _str_literal(s: str) -> str:
     return repr(s)
 
 
+_LARGE_PAYLOAD_THRESHOLD_BYTES = 64 * 1024
+
+
+def _size_note(shape: dict | None) -> str | None:
+    """Docstring note when a probe recorded an unusually large response size."""
+    n = (shape or {}).get("_observed_bytes")
+    if not isinstance(n, (int, float)) or n < _LARGE_PAYLOAD_THRESHOLD_BYTES:
+        return None
+    kb = round(n / 1024)
+    return f"Returns a large payload (~{kb} KB observed) — prefer narrower queries where available."
+
+
+def _with_size_note(desc: str | None, shape: dict | None) -> str | None:
+    note = _size_note(shape)
+    if not note:
+        return desc
+    desc = (desc or "").strip()
+    return f"{desc}\n\n{note}" if desc else note
+
+
 def _docstring(text: str | None, indent: str) -> str:
     text = (text or "").strip()
     if not text:
@@ -215,9 +235,9 @@ _DIG_LIST = '''def _dig_list(obj: Any, path: tuple[str, ...]) -> list:
     return cur'''
 
 
-def _docstring_with_args(tool: dict, ordered: list[tuple[str, dict]], indent: str) -> str:
+def _docstring_with_args(tool: dict, ordered: list[tuple[str, dict]], indent: str, shape: dict | None = None) -> str:
     """Build docstring with description + Args section from inputSchema properties."""
-    desc = (tool.get("description") or "").strip()
+    desc = _with_size_note((tool.get("description") or "").strip(), shape) or ""
     props_lines = []
     for pname, pschema in ordered:
         py = sanitize(pname)
@@ -260,6 +280,108 @@ def _docstring_with_args(tool: dict, ordered: list[tuple[str, dict]], indent: st
     return f'{indent}"""{joined}\n{indent}"""'
 
 
+def _variants_are_numeric(variants: dict) -> bool:
+    """True iff every discriminator variant key parses as an int.
+
+    Fallback only — used when the discriminator's declared schema type can't be
+    determined (see `_discriminator_type`). Sniffing keys alone is unsound:
+    a `type: "string"` discriminator can have numeric-looking enum values
+    (e.g. `"1"`, `"2"`), which would be misclassified as int.
+    """
+    if not variants:
+        return True
+    for k in variants:
+        try:
+            int(k)
+        except ValueError:
+            return False
+    return True
+
+
+def _discriminator_type(disc: str, props: dict, overrides: dict, variants: dict) -> str:
+    """The discriminator argument's Python type — `"int"` or `"str"` — schema/override first.
+
+    Priority: an explicit `input_overrides` entry for the discriminator, then the
+    tool's declared `inputSchema` type for that property (a JSON Schema `type` may
+    also be a list, e.g. `["string", "null"]` — if the list names both `"integer"`
+    and `"string"`, `"integer"` wins), then (only if neither is available) a
+    best-effort sniff of the shape-spec's variant keys.
+
+    `float` is deliberately unsupported: `typing.Literal` cannot hold float values,
+    so a float discriminator can't be expressed as `@overload` stubs. An explicit
+    `"float"` override raises; a schema `type: "number"` (which permits floats)
+    falls back to key-sniffing rather than assuming int, so real float variant
+    keys (e.g. `"1.5"`) resolve to `"str"` instead of crashing on `int()`.
+
+    Any other explicit override (not `int`/`str`/`float`) also raises — a
+    discriminator's `Literal` type can only ever be `int` or `str`, so silently
+    ignoring an unrecognized override (e.g. a typo like `"strr"`) would drop
+    the user's intent without warning.
+    """
+    override = overrides.get(disc)
+    if override == "int":
+        return "int"
+    if override == "str":
+        return "str"
+    if override is not None:
+        raise ValueError(
+            f"discriminator {disc!r} has unsupported input_overrides type {override!r} — "
+            "only 'int' or 'str' are valid for a discriminator (Literal can't hold floats "
+            "or other types)"
+        )
+    pschema = props.get(disc)
+    if isinstance(pschema, dict):
+        t = pschema.get("type")
+        types = t if isinstance(t, list) else [t]
+        if "integer" in types:
+            return "int"
+        if "string" in types:
+            return "str"
+        if "number" in types:
+            return "int" if _variants_are_numeric(variants) else "str"
+    return "int" if _variants_are_numeric(variants) else "str"
+
+
+def _variant_key(k: str, disc: str, numeric: bool) -> int | str:
+    """Coerce a variant key to the discriminator's resolved type, erroring with context on mismatch."""
+    if not numeric:
+        return k
+    try:
+        return int(k)
+    except ValueError:
+        raise ValueError(
+            f"discriminator {disc!r} resolved to int, but variant key {k!r} isn't a valid int — "
+            "fix the shape-spec's input_overrides/variant keys or the tool's inputSchema type"
+        ) from None
+
+
+def _get_overrides(shape: dict) -> dict:
+    """`input_overrides` sidecar entry, validated as a dict.
+
+    A hand-edited shape-spec can set this to the wrong JSON type (e.g. a list,
+    from a typo); failing here with a clear message beats a confusing
+    `AttributeError` from `.get()` deep inside signature rendering.
+    """
+    overrides = shape.get("input_overrides") or {}
+    if not isinstance(overrides, dict):
+        raise ValueError(f"shape-spec input_overrides must be a dict, got {type(overrides).__name__}: {overrides!r}")
+    return overrides
+
+
+def _discriminator_context(tool: dict, shape: dict) -> tuple[str, bool]:
+    """Resolve (discriminator name, is-numeric) for a discriminated-union shape-spec.
+
+    Shared by `_render_overloaded` (signature/Literal generation) and
+    `render_module` (return-model registration) so the two call sites can't
+    drift on discriminator-type resolution.
+    """
+    disc: str = shape["discriminator"]
+    variants: dict = shape["variants"]
+    props = (tool.get("inputSchema") or {}).get("properties") or {}
+    numeric = _discriminator_type(disc, props, _get_overrides(shape), variants) == "int"
+    return disc, numeric
+
+
 def _render_overloaded(
     tool: dict,
     fn: str,
@@ -271,14 +393,16 @@ def _render_overloaded(
     schema: dict | None = None,
 ) -> str:
     """Emit @overload stubs (one per discriminator value) + impl signature."""
-    disc: str = shape["discriminator"]
     variants: dict = shape["variants"]
     unwrap: list = shape.get("unwrap") or []
-    overrides: dict = shape.get("input_overrides") or {}
+    overrides: dict = _get_overrides(shape)
     container: str | None = shape.get("return_container")
 
+    disc, numeric = _discriminator_context(tool, shape)
+    disc_type = "int" if numeric else "str"
+
     ordered = sorted(props.items(), key=lambda kv: kv[0] not in required)
-    sorted_variants = sorted(variants.items(), key=lambda kv: int(kv[0]))
+    coerced_variants = sorted(((_variant_key(k, disc, numeric), v) for k, v in variants.items()), key=lambda kv: kv[0])
 
     def _build_params(disc_type: str) -> list[str]:
         out = []
@@ -298,22 +422,23 @@ def _render_overloaded(
         all_p = ["caller: McpCaller", "*", *params] if params else ["caller: McpCaller"]
         return ", ".join(all_p)
 
-    variant_models = [(int(k), v.get("return_model", "Any")) for k, v in sorted_variants if v.get("return_model")]
+    variant_models = [(val, v.get("return_model", "Any")) for val, v in coerced_variants if v.get("return_model")]
     union_ret = " | ".join(m for _, m in variant_models) if variant_models else "Any"
 
     blocks: list[str] = []
 
     for val, model in variant_models:
         ret = f"list[{model}]" if container == "list" else model
-        params = _build_params(f"Literal[{val}]")
+        literal = str(val) if numeric else repr(val)
+        params = _build_params(f"Literal[{literal}]")
         blocks.append(f"@overload\nasync def {fn}({_sig(params)}) -> {ret}: ...")
 
     impl_ret = f"list[{union_ret}]" if container == "list" else union_ret
-    impl_lines = [f"async def {fn}({_sig(_build_params('int'))}) -> {impl_ret}:"]
+    impl_lines = [f"async def {fn}({_sig(_build_params(disc_type))}) -> {impl_ret}:"]
     if embed_schema:
-        impl_lines.append(_docstring_with_args(tool, ordered, "    "))
+        impl_lines.append(_docstring_with_args(tool, ordered, "    ", shape=shape))
     else:
-        impl_lines.append(_docstring(tool.get("description"), "    "))
+        impl_lines.append(_docstring(_with_size_note(tool.get("description"), shape), "    "))
 
     body_args = [(sanitize(pname), pname, pname in required or pname == disc) for pname, _ in ordered]
     req_pairs = [(py, pname) for py, pname, is_req in body_args if is_req]
@@ -369,7 +494,7 @@ def render_tool(tool: dict, shape: dict | None = None, embed_schema: bool = Fals
 
     unwrap: list = shape.get("unwrap") or []
     return_model: str | None = shape.get("return_model")
-    overrides: dict = shape.get("input_overrides") or {}
+    overrides: dict = _get_overrides(shape)
     # return_container="list" => the unwrapped value is a LIST of return_model
     # records (e.g. query_acme's data.results). Annotation/cast become
     # list[Model] and the body digs via _dig_list. Default (None) = dict/scalar.
@@ -400,9 +525,9 @@ def render_tool(tool: dict, shape: dict | None = None, embed_schema: bool = Fals
 
     lines = [f"async def {fn}({sig}) -> {ret_ann}:"]
     if embed_schema:
-        lines.append(_docstring_with_args(tool, ordered, "    "))
+        lines.append(_docstring_with_args(tool, ordered, "    ", shape=shape))
     else:
-        lines.append(_docstring(tool.get("description"), "    "))
+        lines.append(_docstring(_with_size_note(tool.get("description"), shape), "    "))
 
     req_pairs = [(py, pname) for py, pname, is_req in body_args if is_req]
     opt_pairs = [(py, pname) for py, pname, is_req in body_args if not is_req]
@@ -532,7 +657,9 @@ def render_module(
             if not sp:
                 continue
             if sp.get("discriminator") and sp.get("variants"):
-                for _, variant in sorted(sp["variants"].items(), key=lambda kv: int(kv[0])):
+                variants = sp["variants"]
+                disc, numeric = _discriminator_context(tool, sp)
+                for _, variant in sorted(variants.items(), key=lambda kv: _variant_key(kv[0], disc, numeric)):
                     if variant.get("return_model"):
                         _append_model(variant["return_model"], variant.get("fields") or {})
             elif sp.get("return_model"):
@@ -742,11 +869,17 @@ def merge_skeletons(skeletons: list[dict]) -> dict:
     return out
 
 
-def probe_skeleton(tool: str, args_list: list[dict], shapes: list[Any]) -> dict:
+def probe_skeleton(
+    tool: str, args_list: list[dict], shapes: list[Any], observed_bytes: list[int] | None = None
+) -> dict:
     """Build a shape-spec skeleton from N probe calls.
 
-    args_list  — the N arg-dicts passed to the tool (one per probe).
-    shapes     — the N summarize_shape() results from those calls.
+    args_list      — the N arg-dicts passed to the tool (one per probe).
+    shapes         — the N summarize_shape() results from those calls.
+    observed_bytes — optional raw response sizes (one per probe); the largest
+                     is recorded as `_observed_shape` so the generated docstring
+                     can warn callers about outsized payloads. Omitted entirely
+                     when not passed (byte-stable with existing shapes.json).
 
     When args_list has a single entry, probed_args is that dict (byte-stable
     with current single-probe output). With multiple entries, probed_args is
@@ -759,14 +892,15 @@ def probe_skeleton(tool: str, args_list: list[dict], shapes: list[Any]) -> dict:
         {k: _merge_scalar({v}) for k, v in merged.items() if isinstance(v, str)} if isinstance(merged, dict) else {}
     )
     probed_args: Any = args_list[0] if len(args_list) == 1 else list(args_list)
-    return {
-        tool: {
-            "unwrap": [],
-            "return_model": None,
-            "input_overrides": {},
-            "fields": fields,
-            "source": "live",
-            "probed_args": probed_args,
-            "_observed_shape": merged,
-        }
+    entry: dict[str, Any] = {
+        "unwrap": [],
+        "return_model": None,
+        "input_overrides": {},
+        "fields": fields,
+        "source": "live",
+        "probed_args": probed_args,
+        "_observed_shape": merged,
     }
+    if observed_bytes:
+        entry["_observed_bytes"] = max(observed_bytes)
+    return {tool: entry}

{mcp_client_kit-0.2.0 → mcp_client_kit-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mcp-client-kit"
-version = "0.2.0"
+version = "0.3.0"
 description = "mcp-client-kit — generate typed Python wrappers for any MCP server (codegen skill + CLI)."
 readme = "README.md"
 license = { text = "MIT" }