Add AgentState handling and token management to WHIP client

Author: jason-shen

src/streamcore/__init__.py

@@ -1,5 +1,6 @@
 from .audio import CHANNELS, FRAME_SIZE, SAMPLE_RATE
 from .types import (
+    AgentState,
     ConnectionStatus,
     TranscriptEntry,
     DataChannelMessage,
@@ -13,6 +14,7 @@
     "Client",
     "Config",
     "EventHandler",
+    "AgentState",
     "ConnectionStatus",
     "TranscriptEntry",
     "DataChannelMessage",

src/streamcore/client.py

@@ -5,6 +5,7 @@
 import logging
 from threading import Lock
 
+import aiohttp
 import av
 import numpy as np
 from aiortc import (
@@ -17,6 +18,7 @@
 
 from .audio import SAMPLE_RATE, _SDKAudioTrack
 from .types import (
+    AgentState,
     Config,
     ConnectionStatus,
     DataChannelMessage,
@@ -104,6 +106,7 @@ def on_dc_message(message):
                         text=data.get("text", ""),
                         final=data.get("final", False),
                         message=data.get("message", ""),
+                        state=data.get("state", ""),
                     )
                     if self.events.on_data_channel_message:
                         self.events.on_data_channel_message(msg)
@@ -132,8 +135,21 @@ async def on_connection_state_change():
             # so localDescription.sdp already contains all candidates.
             offer_sdp = pc.localDescription.sdp
 
+            # Fetch a fresh token from the token endpoint if configured.
+            token = self.config.token
+            if self.config.token_url:
+                fetch_headers: dict[str, str] = {}
+                if self.config.api_key:
+                    fetch_headers["Authorization"] = f"Bearer {self.config.api_key}"
+                async with aiohttp.ClientSession() as http_session:
+                    async with http_session.post(self.config.token_url, headers=fetch_headers) as resp:
+                        if resp.status != 200:
+                            raise RuntimeError(f"Token request failed ({resp.status})")
+                        data = await resp.json()
+                        token = data["token"]
+
             # WHIP exchange.
-            result = await whip_offer(self.config.whip_endpoint, offer_sdp)
+            result = await whip_offer(self.config.whip_endpoint, offer_sdp, token)
             self._session_url = result.session_url
 
             answer = RTCSessionDescription(sdp=result.answer_sdp, type="answer")
@@ -152,7 +168,7 @@ async def disconnect(self) -> None:
         if self._sdk_track is not None:
             self._sdk_track.stop()
             self._sdk_track = None
-        await whip_delete(self._session_url)
+        await whip_delete(self._session_url, self.config.token)
         self._session_url = ""
         if self._blackhole:
             await self._blackhole.stop()
@@ -270,6 +286,14 @@ def _handle_data_channel_message(self, msg: DataChannelMessage) -> None:
                 if self.events.on_error:
                     self.events.on_error(RuntimeError(msg.message))
                 return
+
+            elif msg.type == "state":
+                if self.events.on_agent_state_change and msg.state:
+                    try:
+                        self.events.on_agent_state_change(AgentState(msg.state))
+                    except ValueError:
+                        logger.warning("Unknown agent state: %s", msg.state)
+                return
             else:
                 return
 

src/streamcore/types.py

@@ -20,19 +20,29 @@ class TranscriptEntry:
     partial: bool = False
 
 
+class AgentState(str, Enum):
+    LISTENING = "listening"
+    THINKING = "thinking"
+    SPEAKING = "speaking"
+
+
 @dataclass
 class DataChannelMessage:
-    type: str  # "transcript", "response", or "error"
+    type: str  # "transcript", "response", "error", or "state"
     text: str = ""
     final: bool = False
     message: str = ""  # for error type
+    state: str = ""  # for state type
 
 
 @dataclass
 class Config:
     """Configuration for a StreamCoreAIClient."""
 
     whip_endpoint: str = "http://localhost:8080/whip"
+    token: str = ""  # JWT token for authenticating with the WHIP endpoint
+    token_url: str = ""  # Token endpoint URL; if set, fetches a JWT before each connection (overrides token)
+    api_key: str = ""  # API key sent as Bearer header when fetching from token_url
     ice_servers: list[str] = field(
         default_factory=lambda: ["stun:stun.l.google.com:19302"]
     )
@@ -47,4 +57,5 @@ class EventHandler:
         None
     )
     on_error: Callable[[Exception], None] | None = None
+    on_agent_state_change: Callable[[AgentState], None] | None = None
     on_data_channel_message: Callable[[DataChannelMessage], None] | None = None

src/streamcore/whip.py

@@ -12,16 +12,20 @@ class WhipResult:
     session_url: str
 
 
-async def whip_offer(endpoint: str, offer_sdp: str) -> WhipResult:
+async def whip_offer(endpoint: str, offer_sdp: str, token: str = "") -> WhipResult:
     """Perform a WHIP signaling exchange per RFC 9725 §4.2.
 
     POST an SDP offer, receive a 201 Created with SDP answer and Location header.
     """
+    headers: dict[str, str] = {"Content-Type": "application/sdp"}
+    if token:
+        headers["Authorization"] = f"Bearer {token}"
+
     async with aiohttp.ClientSession() as session:
         async with session.post(
             endpoint,
             data=offer_sdp,
-            headers={"Content-Type": "application/sdp"},
+            headers=headers,
         ) as resp:
             if resp.status != 201:
                 body = await resp.text()
@@ -40,17 +44,20 @@ async def whip_offer(endpoint: str, offer_sdp: str) -> WhipResult:
             return WhipResult(answer_sdp=answer_sdp, session_url=session_url)
 
 
-async def whip_delete(session_url: str) -> None:
+async def whip_delete(session_url: str, token: str = "") -> None:
     """Terminate a WHIP session per RFC 9725 §4.2.
 
     Send HTTP DELETE to the WHIP session URL. Best-effort; errors are ignored.
     """
     if not session_url:
         return
     try:
+        headers: dict[str, str] = {}
+        if token:
+            headers["Authorization"] = f"Bearer {token}"
         timeout = aiohttp.ClientTimeout(total=3)
         async with aiohttp.ClientSession(timeout=timeout) as session:
-            async with session.delete(session_url):
+            async with session.delete(session_url, headers=headers):
                 pass
     except Exception:
         pass