Merge pull request #62 from maia-iyer/slack_agent_token_exchange

✨ Implement token exchange in Slack researcher agent

Author: mrsabath

a2a/slack_researcher/README.md

@@ -24,11 +24,19 @@ The Slack Researcher Agent is designed to perform research tasks across Slack ch
 | JWKS_URL | Endpoint to obtain JWKS for token validation. Enables token validation | No | - |
 | ISSUER | Expected `iss` value of incoming bearer tokens | No | - |
 | AUDIENCE | Expected `aud` value of incoming bearer tokens | No | - |
+| TOKEN_URL | Endpoint to perform token exchange. Required for token exchange. | No | - |
+| CLIENT_ID | Client ID to authenticate to auth server with. Required for token exchange. | No | - |
+| CLIENT_SECRET | Client secret to authenticate to auth server with. Required for token exchange. | No | - |
+| TARGET_AUDIENCE | Audience of token exchanged token | No | - |
+| TARGET_SCOPES | Requested scopes of token exchanged token | No | - |
 
 > **Note on Authorization configuration**
 > By default, no token validation is performed. To enable token validation, set `JWKS_URL`. 
 > If `ISSUER` is additionally set, the `iss` claim will be checked to equal this value. 
 > If `AUDIENCE` is additionally set, the `aud` claim will be checked to equal this value. 
+> If all of `TOKEN_URL`, `CLIENT_ID`, and `CLIENT_SECRET` are set in addition, token exchange will be performed using Bearer tokens from incoming requests, to send to the MCP endpoint. 
+> In addition to `TOKEN_URL`, `CLIENT_ID`, `CLIENT_SECRET`, which trigger token exchange, `TARGET_AUDIENCE` and `TARGET_SCOPES` can be optionally configured as the values of `audience` and `scope` in the token exchange request, respectively. 
+
 
 ## Running in Kagenti
 When deploying in the Kagenti UI - You will need to attach 3 environments to the agent deployment:

a2a/slack_researcher/a2a_agent.py

@@ -25,7 +25,7 @@
 from slack_researcher.config import settings, Settings
 from slack_researcher.event import Event
 from slack_researcher.main import SlackAgent
-from slack_researcher.auth import on_auth_error, BearerAuthBackend
+from slack_researcher.auth import on_auth_error, BearerAuthBackend, auth_headers
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.DEBUG, stream=sys.stdout, format='%(levelname)s: %(message)s')
@@ -138,7 +138,7 @@ async def execute(self, context: RequestContext, event_queue: EventQueue):
         Returns:
             None
         """
-        user_token = context.call_context.user.user_name
+        user_token = context.call_context.user._user.access_token
         user_input = [context.get_user_input()]
         task = context.current_task
         if not task:
@@ -164,9 +164,11 @@ async def execute(self, context: RequestContext, event_queue: EventQueue):
             if settings.MCP_URL:
                 logging.info("Connecting to MCP server at %s", settings.MCP_URL)
 
-                headers={}
-                if user_token:
-                    headers={"Authorization": f"Bearer {user_token}"}
+                headers = await auth_headers(
+                    user_token, 
+                    target_audience=settings.TARGET_AUDIENCE, 
+                    target_scopes=settings.TARGET_SCOPES
+                )
 
                 async with streamablehttp_client(
                     url=settings.MCP_URL,

a2a/slack_researcher/slack_researcher/auth.py

@@ -97,8 +97,75 @@ async def authenticate(self, conn):
             logger.debug("Token successfully validated.")
 
             # return user
-            user = SimpleUser(token)
-            return AuthCredentials(["authenticated"]), user
+            user = AgentUser(token=token, claims=claims)
+            return AuthCredentials(user.scopes()), user
         except AuthlibBaseError as e:
             logger.error(f"Token validation failed: {e}")
-            raise AuthenticationError(f"Invalid token: {e}, status_code=401")
+            raise AuthenticationError(f"Invalid token: {e}, status_code=401")
+
+class AgentUser(SimpleUser):
+    def __init__(self, token, claims) -> None:
+        super().__init__(username=claims.get("sub"))
+        self.access_token = token
+        self.claims = claims
+
+    def scopes(self) -> list[str]:
+        scope = self.claims.get("scope", "")
+        return scope.split()
+
+class TokenExchanger:
+    def __init__(self):
+        if None in [settings.TOKEN_URL, settings.CLIENT_ID, settings.CLIENT_SECRET]:
+            raise Exception("One of TOKEN_URL, CLIENT_ID, CLIENT_SECRET env vars not set - token exchange will not be performed")
+        self.token_url = settings.TOKEN_URL
+        self.client_id = settings.CLIENT_ID
+        self.client_secret = settings.CLIENT_SECRET
+
+    async def exchange(self, subject_token: str, audience: str = None, scope: str = None) -> str:
+        # headers
+        headers = {'Content-Type': 'application/x-www-form-urlencoded'}
+        # data
+        data = {
+            'grant_type': 'urn:ietf:params:oauth:grant-type:token-exchange',
+            'subject_token_type': 'urn:ietf:params:oauth:token-type:access_token',
+            'requested_token_type': 'urn:ietf:params:oauth:token-type:access_token',
+            'client_id': self.client_id,
+            'client_secret': self.client_secret,
+            'subject_token': subject_token,
+        }
+        if not audience is None:
+            data['audience'] = audience
+        if not scope is None:
+            data['scope'] = scope
+        # make token endpoint call
+        logger.debug('Performing token exchange')
+        async with httpx.AsyncClient() as client:
+            try: 
+                response = await client.post(self.token_url, data=data, headers=headers)
+                response.raise_for_status() # raise exception if Http status error
+                token_data = response.json()
+                if "access_token" in token_data:
+                    new_token = token_data["access_token"]
+                    logger.debug(f"Successful token exchange. Using token: {new_token}")
+                    return new_token
+                logger.error("Token exchange failed.")
+                raise AuthenticationError("Token exchange failed. Identity provider response did not include 'access_token'")
+            except httpx.HTTPStatusError as e:
+                logger.error(f"Token exchange failed with status {e.response.status_code}: {e}")
+                raise AuthenticationError("Token endpoint call failed.")
+
+async def auth_headers(access_token, target_audience = None, target_scopes = None):
+    headers = {}
+    if not access_token:
+        return headers
+    try:
+        token_exchanger = TokenExchanger()
+        access_token = await token_exchanger.exchange(access_token, audience=target_audience, scope=target_scopes)
+    except AuthenticationError as e:
+        logging.error(f"Error performing token exchange - returning empty headers: {e}")
+        return headers # 
+    except Exception as e:
+        logging.debug(f"Error creating token exchanger - will passthrough token")
+
+    headers["Authorization"] = f"Bearer {access_token}"
+    return headers

a2a/slack_researcher/slack_researcher/config.py

@@ -36,9 +36,9 @@ class Settings(BaseSettings):
     MCP_URL: str = Field(os.getenv("MCP_URL", "http://slack-tool:8000"), description="Endpoint for an option MCP server")
     SERVICE_PORT: int = Field(os.getenv("SERVICE_URL", 8000), description="Port on which the service will run.")
 
-    # auth variables
+    # auth variables for token validation
     ISSUER: Optional[str] = Field(
-        os.getenv("ISSUER", None), 
+        os.getenv("ISSUER", None),
         description="The issuer for incoming JWT tokens"
     )
     JWKS_URI: Optional[str] = Field(
@@ -50,6 +50,28 @@ class Settings(BaseSettings):
         description="Expected audience value during resource validation"
     )
 
+    # auth variables for token exchange
+    TOKEN_URL: Optional[str] = Field(
+        os.getenv("TOKEN_URL", None),
+        description="Token endpoint to obtain new access tokens"
+    )
+    CLIENT_ID: Optional[str] = Field(
+        os.getenv("CLIENT_ID", None),
+        description="Client ID to authenticate to OAuth server"
+    )
+    CLIENT_SECRET: Optional[str] = Field(
+        os.getenv("CLIENT_SECRET", None),
+        description="Client secret to authenticate to OAuth server"
+    )
+    TARGET_AUDIENCE: Optional[str] = Field(
+        os.getenv("TARGET_AUDIENCE", None),
+        description="Target audience to request during token exchange"
+    )
+    TARGET_SCOPES: Optional[str] = Field(
+        os.getenv("TARGET_SCOPES", None),
+        description="Target scopes to request during token exchange"
+    )
+
     class Config:
         env_file = ".env"
         env_file_encoding = "utf-8"