diff --git a/stackone_ai/toolset.py b/stackone_ai/toolset.py index fef9198..d6006d5 100644 --- a/stackone_ai/toolset.py +++ b/stackone_ai/toolset.py @@ -7,6 +7,7 @@ import json import logging import os +import re import threading from collections.abc import Coroutine, Sequence from dataclasses import dataclass @@ -92,6 +93,14 @@ class ExecuteToolsConfig(TypedDict, total=False): } _USER_AGENT = f"stackone-ai-python/{_SDK_VERSION}" +# Param-style pinned on the /mcp tool-listing URL. The MCP schema and the RPC-execution unwrap +# (_split_envelope_params) must agree on this, so it is pinned rather than following the server +# default — the server default is free to change without breaking the SDK. +_MCP_PARAM_STYLE = "flat_prefixed" + +# Matches a flat_prefixed envelope key: `_` (e.g. `path_id`, `query_limit`). +_FLAT_ENVELOPE_KEY_PATTERN = re.compile(r"^(path|query|body|headers)_(.+)$") + # --- Internal tool_search + tool_execute --- @@ -444,24 +453,17 @@ def execute( ) -> dict[str, Any]: parsed_arguments = self._parse_arguments(arguments) - body_payload = self._extract_record(parsed_arguments.pop("body", None)) - headers_payload = self._extract_record(parsed_arguments.pop("headers", None)) - path_payload = self._extract_record(parsed_arguments.pop("path", None)) - query_payload = self._extract_record(parsed_arguments.pop("query", None)) - - rpc_body: dict[str, Any] = dict(body_payload or {}) - for key, value in parsed_arguments.items(): - rpc_body[key] = value + envelope = self._split_envelope_params(parsed_arguments) payload: dict[str, Any] = { "action": self.name, - "body": rpc_body, - "headers": self._build_action_headers(headers_payload), + "body": envelope["body"], + "headers": self._build_action_headers(envelope["headers"] or None), } - if path_payload: - payload["path"] = path_payload - if query_payload: - payload["query"] = query_payload + if envelope["path"]: + payload["path"] = envelope["path"] + if envelope["query"]: + payload["query"] = envelope["query"] return super().execute(payload, options=options) @@ -477,10 +479,28 @@ def _parse_arguments(self, arguments: str | dict[str, Any] | None) -> dict[str, return dict(parsed) @staticmethod - def _extract_record(value: Any) -> dict[str, Any] | None: - if isinstance(value, dict): - return dict(value) - return None + def _split_envelope_params(params: dict[str, Any]) -> dict[str, dict[str, Any]]: + """Split LLM-supplied tool arguments into the RPC envelope (path/query/headers/body). + + Tools are listed with ``?param-style=flat_prefixed``, so keys arrive as + ``_`` (for example ``path_id``, ``query_limit``). The prefix carries + the parameter location, so the split needs no per-action schema. A bare dict-valued + ``path``/``query``/``headers``/``body`` key is still bucketed for clients holding a + cached nested schema, and any other key falls through to the body. + """ + buckets: dict[str, dict[str, Any]] = {"path": {}, "query": {}, "headers": {}, "body": {}} + for key, value in params.items(): + match = _FLAT_ENVELOPE_KEY_PATTERN.match(key) + if match: + location, field = match.group(1), match.group(2) + buckets[location].setdefault(field, value) + continue + if key in ("path", "query", "headers", "body") and isinstance(value, dict): + for field, field_value in value.items(): + buckets[key].setdefault(field, field_value) + continue + buckets["body"][key] = value + return buckets def _build_action_headers(self, additional_headers: dict[str, Any] | None) -> dict[str, str]: headers: dict[str, str] = {} @@ -1240,7 +1260,7 @@ def fetch_tools( if cached is not None: return cached - endpoint = f"{self.base_url.rstrip('/')}/mcp" + endpoint = f"{self.base_url.rstrip('/')}/mcp?param-style={_MCP_PARAM_STYLE}" def _fetch_for_account(account: str | None) -> list[StackOneTool]: headers = self._build_mcp_headers(account) diff --git a/tests/test_fetch_tools.py b/tests/test_fetch_tools.py index 17fc8f1..b39fc39 100644 --- a/tests/test_fetch_tools.py +++ b/tests/test_fetch_tools.py @@ -662,3 +662,21 @@ def counting_init(self, tools, hybrid_alpha=None): toolset.search_tools("bar") assert build_count["count"] == 2 + + +class TestMcpParamStylePinning: + """The /mcp listing URL must pin param-style so the schema matches the RPC unwrap.""" + + def test_fetch_tools_pins_flat_prefixed_param_style(self, monkeypatch): + captured: dict[str, str] = {} + + def fake_fetch(endpoint: str, headers: dict[str, str]) -> list[_McpToolDefinition]: + captured["endpoint"] = endpoint + return [] + + monkeypatch.setattr("stackone_ai.toolset._fetch_mcp_tools", fake_fetch) + + toolset = StackOneToolSet(api_key="test-key", base_url="https://api.example.com") + toolset.fetch_tools(account_ids=["acc1"]) + + assert captured["endpoint"] == "https://api.example.com/mcp?param-style=flat_prefixed" diff --git a/tests/test_tool_calling.py b/tests/test_tool_calling.py index 08e04f0..783c67a 100644 --- a/tests/test_tool_calling.py +++ b/tests/test_tool_calling.py @@ -327,16 +327,33 @@ def test_parse_arguments_non_dict(self, rpc_tool): with pytest.raises(ValueError, match="Tool arguments must be a JSON object"): rpc_tool._parse_arguments("[1, 2, 3]") - def test_extract_record_with_dict(self, rpc_tool): - """Test _extract_record with dict input""" - result = rpc_tool._extract_record({"key": "value"}) - assert result == {"key": "value"} - - def test_extract_record_with_non_dict(self, rpc_tool): - """Test _extract_record with non-dict input""" - assert rpc_tool._extract_record("string") is None - assert rpc_tool._extract_record(123) is None - assert rpc_tool._extract_record(None) is None + def test_split_envelope_params_routes_flat_prefixed_keys(self, rpc_tool): + """flat_prefixed keys are bucketed into the RPC envelope by their location prefix""" + actual = rpc_tool._split_envelope_params( + { + "path_id": "123", + "query_limit": 10, + "headers_x-custom": "value", + "body_name": "test", + } + ) + assert actual["path"] == {"id": "123"} + assert actual["query"] == {"limit": 10} + assert actual["headers"] == {"x-custom": "value"} + assert actual["body"] == {"name": "test"} + + def test_split_envelope_params_buckets_nested_and_unprefixed_keys(self, rpc_tool): + """Bare nested envelopes are accepted and unprefixed keys fall through to the body""" + actual = rpc_tool._split_envelope_params( + { + "body": {"nested": "value"}, + "path": {"id": "1"}, + "extra": "x", + } + ) + assert actual["path"] == {"id": "1"} + assert actual["query"] == {} + assert actual["body"] == {"nested": "value", "extra": "x"} class TestBinaryDownloadResponse: