Accept the temporary gateway JWT from the portfolio launch URL

Author: Asaif Ali

streamlit_ui/api_client.py

@@ -18,12 +18,16 @@ def __init__(self, status_code: Optional[int], detail: str):
 
 
 class AgentServiceClient:
-    def __init__(self, base_url: str = DEFAULT_BASE_URL, token: str = "streamlit-user"):
+    def __init__(self, base_url: str = DEFAULT_BASE_URL, token: str = "streamlit-user", gateway_token: str = ""):
         self.base_url = base_url.rstrip("/")
         self.token = token or "streamlit-user"
+        self.gateway_token = (gateway_token or "").strip()
 
     def _headers(self) -> Dict[str, str]:
-        return {"Authorization": f"Bearer {self.token}"}
+        headers = {"Authorization": f"Bearer {self.token}"}
+        if self.gateway_token:
+            headers["X-LLM-Gateway-Token"] = self.gateway_token
+        return headers
 
     def _request(self, method: str, path: str, timeout: int = DEFAULT_TIMEOUT, **kwargs) -> Any:
         url = f"{self.base_url}{path}"

streamlit_ui/app.py

@@ -175,6 +175,16 @@
 # --------------------------------------------------------------------------
 if "base_url" not in st.session_state:
     st.session_state.base_url = DEFAULT_BASE_URL
+if "llm_gateway_session_token" not in st.session_state:
+    st.session_state.llm_gateway_session_token = ""
+
+portfolio_token = str(st.query_params.get("portfolio_llm_session", "")).strip()
+if portfolio_token:
+    st.session_state.llm_gateway_session_token = portfolio_token
+    try:
+        del st.query_params["portfolio_llm_session"]
+    except Exception:
+        pass
 if "token" not in st.session_state:
     st.session_state.token = "streamlit-user"
 if "active_task_id" not in st.session_state:
@@ -190,7 +200,7 @@
 
 
 def get_client() -> AgentServiceClient:
-    return AgentServiceClient(base_url=st.session_state.base_url, token=st.session_state.token)
+    return AgentServiceClient(base_url=st.session_state.base_url, token=st.session_state.token, gateway_token=st.session_state.get("llm_gateway_session_token", ""))
 
 
 def try_lottie(url: str, height: int = 160):