mcp/tests/test_public_server.py

import unittest

from public_server import PublicMcpHttpHandler, extract_client_ip
from utils.rate_limiter import SimpleRateLimiter


class FakeContext:
    def __init__(self, gateway_session_id):
        self.gateway_session_id = gateway_session_id

    def has_session(self):
        return bool(self.gateway_session_id)


class FakeParser:
    def parse(self, headers):
        return FakeContext(headers.get('X-Gateway-Session', ''))


class FakeGateway:
    def __init__(self):
        self.calls = []
        self.list_calls = []
        self.tool_result = {'code': 'MCP_0000', 'data': {'ok': True}}

    def registered_tool_names(self):
        return ('query_order', 'query_track')

    def list_tools(self, gateway_session_id):
        self.list_calls.append(gateway_session_id)
        return [{'name': 'query_track', 'description': 'query track', 'input_schema': {'type': 'object'}}]

    def call_tool(self, gateway_session_id, name, arguments=None, request_id='', client_ip=''):
        self.calls.append((gateway_session_id, name, arguments, request_id, client_ip))
        return self.tool_result


class PublicMcpHttpHandlerTest(unittest.TestCase):
    def test_constructor_uses_registered_names_without_loading_dynamic_list(self):
        gateway = FakeGateway()

        PublicMcpHttpHandler(gateway, context_parser=FakeParser())

        self.assertEqual([], gateway.list_calls)

    def test_handle_tools_list_passes_session_and_returns_dynamic_tools(self):
        gateway = FakeGateway()
        handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser())

        response = handler.handle_json_rpc(
            headers={'X-Gateway-Session': 'GWS_A'},
            message={
                'jsonrpc': '2.0',
                'id': 2,
                'method': 'tools/list',
                'params': {},
            },
            client_ip='10.0.0.5',
        )

        self.assertEqual(['GWS_A'], gateway.list_calls)
        self.assertEqual(
            ['query_track'],
            [tool['name'] for tool in response['result']['tools']],
        )
        self.assertIn('inputSchema', response['result']['tools'][0])

    def test_missing_session_on_tools_list_returns_protocol_error(self):
        gateway = FakeGateway()
        handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser())

        response = handler.handle_json_rpc(
            headers={},
            message={
                'jsonrpc': '2.0',
                'id': 3,
                'method': 'tools/list',
                'params': {},
            },
            client_ip='10.0.0.5',
        )

        self.assertEqual(-32000, response['error']['code'])
        self.assertIn('Workbuddy', response['error']['message'])
        self.assertEqual([], gateway.list_calls)

    def test_handle_tools_call_passes_gateway_session_to_public_gateway(self):
        gateway = FakeGateway()
        handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser())
        response = handler.handle_json_rpc(
            headers={'X-Gateway-Session': 'GWS_A'},
            message={
                'jsonrpc': '2.0',
                'id': 1,
                'method': 'tools/call',
                'params': {
                    'name': 'query_order',
                    'arguments': {'keyword': 'USC'},
                },
            },
            client_ip='10.0.0.5'
        )

        self.assertEqual(False, response['result']['isError'])
        self.assertEqual('GWS_A', gateway.calls[0][0])
        self.assertEqual('query_order', gateway.calls[0][1])
        self.assertEqual('10.0.0.5', gateway.calls[0][4])

    def test_backend_business_error_returns_tool_error_result(self):
        gateway = FakeGateway()
        gateway.tool_result = {
            'code': 'MCP_1401',
            'msg': 'invalid exact order query conditions',
            'data': [],
            'meta': {'request_id': 'rq_backend'},
        }
        handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser())

        response = handler.handle_json_rpc(
            headers={'X-Gateway-Session': 'GWS_A'},
            message={
                'jsonrpc': '2.0',
                'id': 7,
                'method': 'tools/call',
                'params': {
                    'name': 'query_order',
                    'arguments': {'keyword': 'USC'},
                },
            },
            client_ip='10.0.0.5',
        )

        self.assertNotIn('error', response)
        self.assertTrue(response['result']['isError'])
        self.assertIn('MCP_1401', response['result']['content'][0]['text'])
        self.assertEqual(
            {
                'code': 'MCP_1401',
                'msg': 'invalid exact order query conditions',
                'meta': {'request_id': 'rq_backend'},
            },
            response['result']['structuredContent'],
        )

    def test_missing_session_on_tool_call_returns_error_content(self):
        gateway = FakeGateway()
        handler = PublicMcpHttpHandler(gateway, context_parser=FakeParser())
        response = handler.handle_json_rpc(
            headers={},
            message={
                'jsonrpc': '2.0',
                'id': 1,
                'method': 'tools/call',
                'params': {'name': 'query_order', 'arguments': {'keyword': 'USC'}},
            },
        )

        self.assertTrue(response['result']['isError'])
        self.assertIn('这台设备的 Workbuddy 配置已失效，请重新生成配置', response['result']['content'][0]['text'])

    def test_extract_client_ip_ignores_spoofable_forwarded_for_header(self):
        client_ip = extract_client_ip(
            {'X-Forwarded-For': '203.0.113.9'},
            ('10.0.0.5', 54321),
        )

        self.assertEqual('10.0.0.5', client_ip)


class RateLimitTest(unittest.TestCase):
    def _make_handler(self, max_requests=2):
        gateway = FakeGateway()
        limiter = SimpleRateLimiter(max_requests=max_requests, window_seconds=60)
        return PublicMcpHttpHandler(gateway, context_parser=FakeParser(), rate_limiter=limiter)

    def _tools_call_msg(self, tool_name='query_order'):
        return {
            'jsonrpc': '2.0', 'id': 1,
            'method': 'tools/call',
            'params': {'name': tool_name, 'arguments': {'keyword': 'test'}},
        }

    # --- tools/call per-tool独立限流 ---

    def test_known_tool_rate_limited_after_quota_exhausted(self):
        handler = self._make_handler(max_requests=1)
        handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.2.3.4')
        response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.2.3.4')
        self.assertIn('error', response)
        self.assertIn('Rate limit', response['error']['message'])

    def test_two_known_tools_have_independent_quotas(self):
        # query_order 限流不影响 query_track
        handler = self._make_handler(max_requests=1)
        handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('query_order'), client_ip='1.2.3.4')
        # query_order quota exhausted, query_track should still work
        response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('query_track'), client_ip='1.2.3.4')
        self.assertNotIn('error', response)

    def test_unknown_tool_name_uses_shared_session_bucket_not_new_bucket(self):
        # 未知工具名应归入 session_id 桶，同一会话的不同未知工具共享同一配额，不能无限新建桶
        handler = self._make_handler(max_requests=1)
        # 先用 session 桶打一次（用未知工具名触发）
        handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('fake_tool_1'), client_ip='1.2.3.4')
        # 再用同会话的另一个未知工具名，应该命中同一个 session 桶，被限流
        response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg('fake_tool_2'), client_ip='1.2.3.4')
        self.assertIn('error', response)
        self.assertIn('Rate limit', response['error']['message'])

    def test_different_sessions_have_independent_quotas(self):
        # 不同会话（不同员工）即使来自同一IP，也拥有独立的配额
        handler = self._make_handler(max_requests=1)
        handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.1.1.1')
        # 不同 session，即使同一 IP，quota 独立 → 应该允许
        response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_B'}, self._tools_call_msg(), client_ip='1.1.1.1')
        self.assertNotIn('error', response)

    def test_same_session_from_different_ips_shares_quota(self):
        # 同一会话从不同IP发来（如移动网络切换），仍共享同一会话配额
        handler = self._make_handler(max_requests=1)
        handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='1.1.1.1')
        # 同 session，不同 IP → 命中同一 session 桶，被限流
        response = handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, self._tools_call_msg(), client_ip='2.2.2.2')
        self.assertIn('error', response)
        self.assertIn('Rate limit', response['error']['message'])

    # --- initialize 不受限流，tools/list 有独立 bucket ---

    def test_initialize_is_never_rate_limited(self):
        # initialize 是握手协议，无论请求多少次都不应被限流
        handler = self._make_handler(max_requests=1)
        for i in range(5):
            response = handler.handle_json_rpc({}, {'jsonrpc': '2.0', 'id': i, 'method': 'initialize'}, client_ip='1.2.3.4')
            self.assertNotIn('error', response, f'initialize should never be rate limited (attempt {i})')
            self.assertIn('protocolVersion', response['result'])

    def test_tools_list_is_rate_limited(self):
        handler = self._make_handler(max_requests=1)
        headers = {'X-Gateway-Session': 'GWS_A'}
        handler.handle_json_rpc(headers, {'jsonrpc': '2.0', 'id': 1, 'method': 'tools/list'}, client_ip='1.2.3.4')
        response = handler.handle_json_rpc(headers, {'jsonrpc': '2.0', 'id': 2, 'method': 'tools/list'}, client_ip='1.2.3.4')
        self.assertIn('error', response)
        self.assertIn('Rate limit', response['error']['message'])

    def test_tools_list_quota_independent_from_tools_call_quota(self):
        # tools/list 和 tools/call 使用不同 bucket，互不影响
        handler = self._make_handler(max_requests=1)
        # 用掉 tools/list 的1个 slot
        handler.handle_json_rpc({'X-Gateway-Session': 'GWS_A'}, {'jsonrpc': '2.0', 'id': 1, 'method': 'tools/list'}, client_ip='1.2.3.4')
        # tools/call 应仍可正常执行（使用独立的 client_ip:tool_name bucket）
        response = handler.handle_json_rpc(
            {'X-Gateway-Session': 'GWS_A'},
            {'jsonrpc': '2.0', 'id': 2, 'method': 'tools/call',
             'params': {'name': 'query_order', 'arguments': {'keyword': 'test'}}},
            client_ip='1.2.3.4',
        )
        self.assertNotIn('error', response)

if __name__ == '__main__':
    unittest.main()