Add tuple unpacking support to GeneratedToken and improve thread safety

Author: Mateusz

.gitignore

@@ -117,4 +117,5 @@ var/*.log
 var/*.md
 var/*.json
 var/*.jsonl
-var/*.py
+var/*.py
+pyright_output.json

dev/pyright_output.json

@@ -1,13 +1,13 @@
-{
-    "version": "1.1.407",
-    "time": "1767034635449",
-    "generalDiagnostics": [],
-    "summary": {
-        "filesAnalyzed": 1011,
-        "errorCount": 0,
-        "warningCount": 0,
-        "informationCount": 0,
-        "timeInSec": 41.452
-    }
-}
-
+{
+    "version": "1.1.407",
+    "time": "1767041279293",
+    "generalDiagnostics": [],
+    "summary": {
+        "filesAnalyzed": 1011,
+        "errorCount": 0,
+        "warningCount": 0,
+        "informationCount": 0,
+        "timeInSec": 18.36
+    }
+}
+

src/core/auth/sso/token_service.py

@@ -1,35 +1,40 @@
-"""
-Token service for SSO authentication.
-
-This module provides secure token generation, hashing, and verification
-using Argon2id with 2025-recommended security parameters.
-"""
-
-import base64
-import secrets
-
-import pydantic
-from argon2 import PasswordHasher
-from argon2.exceptions import InvalidHashError, VerificationError, VerifyMismatchError
-
-from src.core.auth.sso.exceptions import TokenError
-
-
-class GeneratedToken(pydantic.BaseModel):
-    """Result of token generation.
-
-    Contains both the plaintext token (for immediate use) and its hash
-    (for secure storage and verification).
-
-    Attributes:
-        plaintext: Base64url-encoded token (43+ characters) with 256-bit entropy
-        hash: Argon2id hash of the token for secure storage
-    """
-
-    plaintext: str
-    hash: str
-
-    model_config = {"frozen": True}
+"""
+Token service for SSO authentication.
+
+This module provides secure token generation, hashing, and verification
+using Argon2id with 2025-recommended security parameters.
+"""
+
+import base64
+import secrets
+
+import pydantic
+from argon2 import PasswordHasher
+from argon2.exceptions import InvalidHashError, VerificationError, VerifyMismatchError
+
+from src.core.auth.sso.exceptions import TokenError
+
+
+class GeneratedToken(pydantic.BaseModel):
+    """Result of token generation.
+
+    Contains both the plaintext token (for immediate use) and its hash
+    (for secure storage and verification).
+
+    Attributes:
+        plaintext: Base64url-encoded token (43+ characters) with 256-bit entropy
+        hash: Argon2id hash of the token for secure storage
+    """
+
+    plaintext: str
+    hash: str
+
+    model_config = {"frozen": True}
+
+    def __iter__(self):  # type: ignore[override]
+        """Allow tuple unpacking for backward compatibility."""
+        yield self.plaintext
+        yield self.hash
 
 
 class TokenService:
@@ -99,41 +104,41 @@ def __init__(
             salt_len=16,  # Always 16 bytes salt
         )
 
-    def generate_token(self) -> GeneratedToken:
-        """
-        Generate a new agent token with 256-bit entropy.
-
-        The token is generated using cryptographically secure random bytes
-        and encoded as base64url for Bearer token compatibility.
-
-        Returns:
-            GeneratedToken: Object containing both plaintext token and its hash
-                - plaintext: Base64url-encoded token (43+ characters)
-                - hash: Argon2id hash of token
-
-        Raises:
-            TokenError: If token generation or hashing fails
-        """
-        try:
-            # Generate 256 bits (32 bytes) of cryptographically secure random data
-            token_bytes = secrets.token_bytes(32)
-
-            # Encode as base64url (URL-safe, no padding)
-            plaintext_token = (
-                base64.urlsafe_b64encode(token_bytes).decode("ascii").rstrip("=")
-            )
-
-            # Hash the token using Argon2id
-            token_hash = self.hash_token(plaintext_token)
-
-            return GeneratedToken(plaintext=plaintext_token, hash=token_hash)
-
-        except Exception as e:
-            raise TokenError(
-                "Failed to generate token",
-                details={"error": str(e)},
-                original_error=e,
-            ) from e
+    def generate_token(self) -> GeneratedToken:
+        """
+        Generate a new agent token with 256-bit entropy.
+
+        The token is generated using cryptographically secure random bytes
+        and encoded as base64url for Bearer token compatibility.
+
+        Returns:
+            GeneratedToken: Object containing both plaintext token and its hash
+                - plaintext: Base64url-encoded token (43+ characters)
+                - hash: Argon2id hash of token
+
+        Raises:
+            TokenError: If token generation or hashing fails
+        """
+        try:
+            # Generate 256 bits (32 bytes) of cryptographically secure random data
+            token_bytes = secrets.token_bytes(32)
+
+            # Encode as base64url (URL-safe, no padding)
+            plaintext_token = (
+                base64.urlsafe_b64encode(token_bytes).decode("ascii").rstrip("=")
+            )
+
+            # Hash the token using Argon2id
+            token_hash = self.hash_token(plaintext_token)
+
+            return GeneratedToken(plaintext=plaintext_token, hash=token_hash)
+
+        except Exception as e:
+            raise TokenError(
+                "Failed to generate token",
+                details={"error": str(e)},
+                original_error=e,
+            ) from e
 
     def hash_token(self, token: str) -> str:
         """

src/core/ports/streaming_processors.py

@@ -196,7 +196,7 @@ def __init__(self, max_cached_sessions: int = 256) -> None:
         self._max_cached_sessions = max_cached_sessions
         self._session_order: list[str] = []  # Track LRU order
         self._logger = logging.getLogger(__name__)
-        self._lock = threading.Lock()
+        self._lock = threading.RLock()
 
     async def process(self, content: StreamingContent) -> StreamingContent:
         """Process streaming content and check for tool call loops.

tests/unit/test_token_service.py

@@ -6,20 +6,19 @@
 
 import pytest
 from src.core.auth.sso.exceptions import TokenError
-from src.core.auth.sso.token_service import TokenService
+from src.core.auth.sso.token_service import GeneratedToken, TokenService
 
 
 class TestTokenService:
     """Unit tests for TokenService."""
 
     def test_generate_token_returns_tuple(self) -> None:
-        """Test that generate_token returns a tuple of (token, hash)."""
+        """Test that generate_token returns a GeneratedToken that can be unpacked as a tuple."""
         service = TokenService(memory_cost=8192, time_cost=1, parallelism=1)
         result = service.generate_token()
 
-        assert isinstance(result, tuple)
-        assert len(result) == 2
-
+        assert isinstance(result, GeneratedToken)
+        # Verify it can be unpacked as a tuple for backward compatibility
         plaintext_token, token_hash = result
         assert isinstance(plaintext_token, str)
         assert isinstance(token_hash, str)
@@ -98,17 +97,17 @@ def test_hash_token_produces_different_hashes(self) -> None:
         assert service.verify_token(plaintext_token, hash2) is True
         assert service.verify_token(plaintext_token, hash3) is True
 
-    def test_generated_tokens_are_unique(self) -> None:
-        """Test that multiple generated tokens are unique."""
-        service = TokenService(memory_cost=8, time_cost=1, parallelism=1)
-
-        tokens = set()
-        for _ in range(20):
-            plaintext_token, _ = service.generate_token()
-            tokens.add(plaintext_token)
-
-        # All tokens should be unique
-        assert len(tokens) == 20
+    def test_generated_tokens_are_unique(self) -> None:
+        """Test that multiple generated tokens are unique."""
+        service = TokenService(memory_cost=8, time_cost=1, parallelism=1)
+
+        tokens = set()
+        for _ in range(20):
+            plaintext_token, _ = service.generate_token()
+            tokens.add(plaintext_token)
+
+        # All tokens should be unique
+        assert len(tokens) == 20
 
     def test_token_hash_does_not_contain_plaintext(self) -> None:
         """Test that token hash does not contain the plaintext token."""