Skip to content

Commit 5f604d3

Browse files
committed
fix(gemini): 🐛 generate stable session ids based on prompt content
This change updates the session identification logic to be deterministic and stateless, addressing isolation issues in multi-user environments. - Replaces the static, instance-level `_session_id` with a dynamic generation method based on the SHA256 hash of the first user message. - Ensures consistent session IDs for the same conversation context, maintaining continuity across server restarts. - Decouples session identity from the provider instance lifecycle to prevent context bleeding between requests. - Aligns ID formatting with native client behavior using UUID-formatted hashes.
1 parent 1850a7f commit 5f604d3

File tree

1 file changed

+41
-5
lines changed

1 file changed

+41
-5
lines changed

src/rotator_library/providers/gemini_cli_provider.py

Lines changed: 41 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
from pathlib import Path
3636
import uuid
3737
import secrets
38+
import hashlib
3839
from datetime import datetime
3940

4041
lib_logger = logging.getLogger("rotator_library")
@@ -409,9 +410,6 @@ def __init__(self):
409410
self._learned_costs: Dict[str, Dict[str, float]] = {}
410411
self._learned_costs_loaded: bool = False
411412

412-
# Session ID for Code Assist API (persistent per provider instance)
413-
# Matches native gemini-cli behavior: single session UUID for conversation lifetime
414-
self._session_id: str = str(uuid.uuid4())
415413

416414
# =========================================================================
417415
# CREDENTIAL TIER LOOKUP (Provider-specific - uses cache)
@@ -476,6 +474,44 @@ def _generate_user_prompt_id(self) -> str:
476474
"""
477475
return secrets.token_hex(7) # 14 hex characters
478476

477+
def _generate_stable_session_id(self, contents: List[Dict[str, Any]]) -> str:
478+
"""
479+
Generate a stable session ID based on the first user message.
480+
481+
This ensures:
482+
- Same conversation = same session_id (even across server restarts)
483+
- Different conversations = different session_ids
484+
- Multi-user scenarios are properly isolated
485+
486+
Uses SHA256 hash of the first user message to create a deterministic
487+
UUID-formatted session ID. Falls back to random UUID if no user message.
488+
489+
This approach mirrors Antigravity's _generate_stable_session_id() but
490+
uses UUID format instead of the -{number} format to match native
491+
gemini-cli's crypto.randomUUID() output format.
492+
493+
Args:
494+
contents: List of message contents in Gemini format
495+
496+
Returns:
497+
UUID-formatted session ID string
498+
"""
499+
# Find first user message text
500+
for content in contents:
501+
if content.get("role") == "user":
502+
parts = content.get("parts", [])
503+
for part in parts:
504+
if isinstance(part, dict):
505+
text = part.get("text", "")
506+
if text:
507+
# SHA256 hash and use first 16 bytes to create UUID
508+
h = hashlib.sha256(text.encode("utf-8")).digest()
509+
# Format as UUID (8-4-4-4-12 hex chars)
510+
return f"{h[:4].hex()}-{h[4:6].hex()}-{h[6:8].hex()}-{h[8:10].hex()}-{h[10:16].hex()}"
511+
512+
# Fallback to random UUID if no user message found
513+
return str(uuid.uuid4())
514+
479515
def _get_gemini_cli_request_headers(self, model: str) -> Dict[str, str]:
480516
"""
481517
Build request headers matching native gemini-cli client.
@@ -1415,7 +1451,7 @@ async def do_call(attempt_model: str, is_fallback: bool = False):
14151451
"request": {
14161452
"contents": contents,
14171453
"generationConfig": gen_config,
1418-
"session_id": self._session_id,
1454+
"session_id": self._generate_stable_session_id(contents),
14191455
},
14201456
}
14211457

@@ -1751,7 +1787,7 @@ async def count_tokens(
17511787
"user_prompt_id": self._generate_user_prompt_id(),
17521788
"request": {
17531789
"contents": contents,
1754-
"session_id": self._session_id,
1790+
"session_id": self._generate_stable_session_id(contents),
17551791
},
17561792
}
17571793

0 commit comments

Comments
 (0)