feat: accept MCP authorization headers for MCP toolgroups

tests/integration/tool_runtime/test_mcp.py

@@ -0,0 +1,189 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import json
+import socket
+import threading
+import time
+
+import httpx
+import mcp.types as types
+import pytest
+import uvicorn
+from llama_stack_client import Agent
+from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.sse import SseServerTransport
+from starlette.applications import Starlette
+from starlette.exceptions import HTTPException
+from starlette.responses import Response
+from starlette.routing import Mount, Route
+
+AUTH_TOKEN = "test-token"
+
+
+@pytest.fixture(scope="module")
+def mcp_server():
+    server = FastMCP("FastMCP Test Server")
+
+    @server.tool()
+    async def greet_everyone(
+        url: str, ctx: Context
+    ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
+        return [types.TextContent(type="text", text="Hello, world!")]
+
+    sse = SseServerTransport("/messages/")
+
+    async def handle_sse(request):
+        auth_header = request.headers.get("Authorization")
+        auth_token = None
+        if auth_header and auth_header.startswith("Bearer "):
+            auth_token = auth_header.split(" ")[1]
+
+        if auth_token != AUTH_TOKEN:
+            raise HTTPException(status_code=401, detail="Unauthorized")
+
+        async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
+            await server._mcp_server.run(
+                streams[0],
+                streams[1],
+                server._mcp_server.create_initialization_options(),
+            )
+            return Response()
+
+    app = Starlette(
+        routes=[
+            Route("/sse", endpoint=handle_sse),
+            Mount("/messages/", app=sse.handle_post_message),
+        ],
+    )
+
+    def get_open_port():
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.bind(("", 0))
+            return sock.getsockname()[1]
+
+    port = get_open_port()
+
+    config = uvicorn.Config(app, host="0.0.0.0", port=port)
+    server_instance = uvicorn.Server(config)
+    app.state.uvicorn_server = server_instance
+
+    def run_server():
+        server_instance.run()
+
+    # Start the server in a new thread
+    server_thread = threading.Thread(target=run_server, daemon=True)
+    server_thread.start()
+
+    # Polling until the server is ready
+    timeout = 10
+    start_time = time.time()
+
+    while time.time() - start_time < timeout:
+        try:
+            response = httpx.get(f"http://localhost:{port}/sse")
+            if response.status_code == 401:
+                break
+        except httpx.RequestError:
+            pass
+        time.sleep(0.1)
+
+    yield port
+
+    # Tell server to exit
+    server_instance.should_exit = True
+    server_thread.join(timeout=5)
+
+
+def test_mcp_invocation(llama_stack_client, mcp_server):
+    port = mcp_server
+    test_toolgroup_id = "remote::mcptest"
+
+    # registering itself should fail since it requires listing tools
+    with pytest.raises(Exception, match="Unauthorized"):
+        llama_stack_client.toolgroups.register(
+            toolgroup_id=test_toolgroup_id,
+            provider_id="model-context-protocol",
+            mcp_endpoint=dict(uri=f"http://localhost:{port}/sse"),
+        )
+
+    provider_data = {
+        "mcp_headers": {
+            f"http://localhost:{port}/sse": [
+                f"Authorization: Bearer {AUTH_TOKEN}",
+            ],
+        },
+    }
+    auth_headers = {
+        "X-LlamaStack-Provider-Data": json.dumps(provider_data),
+    }
+    llama_stack_client.toolgroups.register(
+        toolgroup_id=test_toolgroup_id,
+        provider_id="model-context-protocol",
+        mcp_endpoint=dict(uri=f"http://localhost:{port}/sse"),
+        extra_headers=auth_headers,
+    )
+    response = llama_stack_client.tools.list(
+        toolgroup_id=test_toolgroup_id,
+        extra_headers=auth_headers,
+    )
+    assert len(response) == 1
+    assert response[0].identifier == "greet_everyone"
+    assert response[0].type == "tool"
+    assert len(response[0].parameters) == 1
+    p = response[0].parameters[0]
+    assert p.name == "url"
+    assert p.parameter_type == "string"
+    assert p.required
+
+    response = llama_stack_client.tool_runtime.invoke_tool(
+        tool_name=response[0].identifier,
+        kwargs=dict(url="https://www.google.com"),
+        extra_headers=auth_headers,
+    )
+    content = response.content
+    assert len(content) == 1
+    assert content[0].type == "text"
+    assert content[0].text == "Hello, world!"
+
+    models = llama_stack_client.models.list()
+    model_id = models[0].identifier
+    print(f"Using model: {model_id}")
+    agent = Agent(
+        client=llama_stack_client,
+        model=model_id,
+        instructions="You are a helpful assistant.",
+        tools=[test_toolgroup_id],
+    )
+    session_id = agent.create_session("test-session")
+    response = agent.create_turn(
+        session_id=session_id,
+        messages=[
+            {
+                "role": "user",
+                "content": "Yo. Use tools.",
+            }
+        ],
+        stream=False,
+        extra_headers=auth_headers,
+    )
+    steps = response.steps
+    first = steps[0]
+    assert first.step_type == "inference"
+    assert len(first.api_model_response.tool_calls) == 1
+    tool_call = first.api_model_response.tool_calls[0]
+    assert tool_call.tool_name == "greet_everyone"
+
+    second = steps[1]
+    assert second.step_type == "tool_execution"
+    tool_response_content = second.tool_responses[0].content
+    assert len(tool_response_content) == 1
+    assert tool_response_content[0].type == "text"
+    assert tool_response_content[0].text == "Hello, world!"
+
+    third = steps[2]
+    assert third.step_type == "inference"
+    assert len(third.api_model_response.tool_calls) == 0