Pass mcp config and auth headers in streaming_query too

Author: rawagner

src/app/endpoints/query.py

@@ -194,20 +194,23 @@ def retrieve_response(
                 "Authorization": f"Bearer {token}",
             }
 
+    extra_headers = {
+        "X-LlamaStack-Provider-Data": json.dumps(
+            {
+                "mcp_headers": mcp_headers,
+            }
+        ),
+    }
+
+    agent.extra_headers = extra_headers
+
     vector_db_ids = [vector_db.identifier for vector_db in client.vector_dbs.list()]
     response = agent.create_turn(
         messages=[UserMessage(role="user", content=query_request.query)],
         session_id=conversation_id,
         documents=query_request.get_documents(),
         stream=False,
         toolgroups=get_rag_toolgroups(vector_db_ids),
-        extra_headers={
-            "X-LlamaStack-Provider-Data": json.dumps(
-                {
-                    "mcp_headers": mcp_headers,
-                }
-            ),
-        },
     )
 
     return str(response.output_message.content), conversation_id  # type: ignore[union-attr]

src/app/endpoints/streaming_query.py

@@ -30,6 +30,7 @@
     validate_attachments_metadata,
 )
 
+
 logger = logging.getLogger("app.endpoints.handlers")
 router = APIRouter(tags=["streaming_query"])
 
@@ -57,7 +58,7 @@ async def get_agent(
         model=model_id,
         instructions=system_prompt,
         input_shields=available_shields if available_shields else [],
-        tools=[],  # mcp config ?
+        tools=[mcp.name for mcp in configuration.mcp_servers],
         enable_session_persistence=True,
     )
     conversation_id = await agent.create_session(get_suid())
@@ -165,7 +166,9 @@ async def streaming_query_endpoint_handler(
     logger.info("LLama stack config: %s", llama_stack_config)
     client = await get_async_llama_stack_client(llama_stack_config)
     model_id = select_model_id(await client.models.list(), query_request)
-    response, conversation_id = await retrieve_response(client, model_id, query_request)
+    response, conversation_id = await retrieve_response(
+        client, model_id, query_request, auth
+    )
 
     async def response_generator(turn_response: Any) -> AsyncIterator[str]:
         """Generate SSE formatted streaming response."""
@@ -204,7 +207,10 @@ async def response_generator(turn_response: Any) -> AsyncIterator[str]:
 
 
 async def retrieve_response(
-    client: AsyncLlamaStackClient, model_id: str, query_request: QueryRequest
+    client: AsyncLlamaStackClient,
+    model_id: str,
+    query_request: QueryRequest,
+    token: str,
 ) -> tuple[Any, str]:
     """Retrieve response from LLMs and agents."""
     available_shields = [shield.identifier for shield in await client.shields.list()]
@@ -234,6 +240,21 @@ async def retrieve_response(
         query_request.conversation_id,
     )
 
+    mcp_headers = {}
+    if token:
+        for mcp_server in configuration.mcp_servers:
+            mcp_headers[mcp_server.url] = {
+                "Authorization": f"Bearer {token}",
+            }
+
+    agent.extra_headers = {
+        "X-LlamaStack-Provider-Data": json.dumps(
+            {
+                "mcp_headers": mcp_headers,
+            }
+        ),
+    }
+
     logger.debug("Session ID: %s", conversation_id)
     vector_db_ids = [
         vector_db.identifier for vector_db in await client.vector_dbs.list()

tests/unit/app/endpoints/test_query.py

@@ -295,7 +295,6 @@ def test_retrieve_response_vector_db_available(mocker):
         documents=[],
         stream=False,
         toolgroups=get_rag_toolgroups(["VectorDB-1"]),
-        extra_headers={"X-LlamaStack-Provider-Data": '{"mcp_headers": {}}'},
     )
 
 
@@ -331,7 +330,6 @@ def test_retrieve_response_no_available_shields(mocker):
         documents=[],
         stream=False,
         toolgroups=None,
-        extra_headers={"X-LlamaStack-Provider-Data": '{"mcp_headers": {}}'},
     )
 
 
@@ -372,7 +370,6 @@ def __init__(self, identifier):
         documents=[],
         stream=False,
         toolgroups=None,
-        extra_headers={"X-LlamaStack-Provider-Data": '{"mcp_headers": {}}'},
     )
 
 
@@ -416,7 +413,6 @@ def __init__(self, identifier):
         documents=[],
         stream=False,
         toolgroups=None,
-        extra_headers={"X-LlamaStack-Provider-Data": '{"mcp_headers": {}}'},
     )
 
 
@@ -465,7 +461,6 @@ def test_retrieve_response_with_one_attachment(mocker):
             },
         ],
         toolgroups=None,
-        extra_headers={"X-LlamaStack-Provider-Data": '{"mcp_headers": {}}'},
     )
 
 
@@ -523,7 +518,6 @@ def test_retrieve_response_with_two_attachments(mocker):
             },
         ],
         toolgroups=None,
-        extra_headers={"X-LlamaStack-Provider-Data": '{"mcp_headers": {}}'},
     )
 
 
@@ -573,23 +567,28 @@ def test_retrieve_response_with_mcp_servers(mocker):
         None,  # conversation_id
     )
 
-    # Check that the agent's create_turn was called with MCP headers
-    mock_agent.create_turn.assert_called_once()
-    call_args = mock_agent.create_turn.call_args
-
-    extra_headers_data = json.loads(
-        call_args[1]["extra_headers"]["X-LlamaStack-Provider-Data"]
-    )
-    mcp_headers = extra_headers_data["mcp_headers"]
+    # Check that the agent's extra_headers property was set correctly
+    expected_extra_headers = {
+        "X-LlamaStack-Provider-Data": json.dumps(
+            {
+                "mcp_headers": {
+                    "http://localhost:3000": {"Authorization": "Bearer test_token_123"},
+                    "https://git.example.com/mcp": {
+                        "Authorization": "Bearer test_token_123"
+                    },
+                }
+            }
+        )
+    }
+    assert mock_agent.extra_headers == expected_extra_headers
 
-    assert "http://localhost:3000" in mcp_headers
-    assert (
-        mcp_headers["http://localhost:3000"]["Authorization"] == "Bearer test_token_123"
-    )
-    assert "https://git.example.com/mcp" in mcp_headers
-    assert (
-        mcp_headers["https://git.example.com/mcp"]["Authorization"]
-        == "Bearer test_token_123"
+    # Check that create_turn was called with the correct parameters
+    mock_agent.create_turn.assert_called_once_with(
+        messages=[UserMessage(role="user", content="What is OpenStack?")],
+        session_id="fake_session_id",
+        documents=[],
+        stream=False,
+        toolgroups=None,
     )
 
 
@@ -632,15 +631,20 @@ def test_retrieve_response_with_mcp_servers_empty_token(mocker):
         None,  # conversation_id
     )
 
-    # Check that the agent's create_turn was called with empty MCP headers
-    mock_agent.create_turn.assert_called_once()
-    call_args = mock_agent.create_turn.call_args
+    # Check that the agent's extra_headers property was set correctly (empty mcp_headers)
+    expected_extra_headers = {
+        "X-LlamaStack-Provider-Data": json.dumps({"mcp_headers": {}})
+    }
+    assert mock_agent.extra_headers == expected_extra_headers
 
-    extra_headers_data = json.loads(
-        call_args[1]["extra_headers"]["X-LlamaStack-Provider-Data"]
+    # Check that create_turn was called with the correct parameters
+    mock_agent.create_turn.assert_called_once_with(
+        messages=[UserMessage(role="user", content="What is OpenStack?")],
+        session_id="fake_session_id",
+        documents=[],
+        stream=False,
+        toolgroups=None,
     )
-    mcp_headers = extra_headers_data["mcp_headers"]
-    assert len(mcp_headers) == 0
 
 
 def test_construct_transcripts_path(setup_configuration, mocker):