|
35 | 35 | from pathlib import Path |
36 | 36 | import uuid |
37 | 37 | import secrets |
| 38 | +import hashlib |
38 | 39 | from datetime import datetime |
39 | 40 |
|
40 | 41 | lib_logger = logging.getLogger("rotator_library") |
@@ -409,9 +410,6 @@ def __init__(self): |
409 | 410 | self._learned_costs: Dict[str, Dict[str, float]] = {} |
410 | 411 | self._learned_costs_loaded: bool = False |
411 | 412 |
|
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()) |
415 | 413 |
|
416 | 414 | # ========================================================================= |
417 | 415 | # CREDENTIAL TIER LOOKUP (Provider-specific - uses cache) |
@@ -476,6 +474,44 @@ def _generate_user_prompt_id(self) -> str: |
476 | 474 | """ |
477 | 475 | return secrets.token_hex(7) # 14 hex characters |
478 | 476 |
|
| 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 | + |
479 | 515 | def _get_gemini_cli_request_headers(self, model: str) -> Dict[str, str]: |
480 | 516 | """ |
481 | 517 | Build request headers matching native gemini-cli client. |
@@ -1415,7 +1451,7 @@ async def do_call(attempt_model: str, is_fallback: bool = False): |
1415 | 1451 | "request": { |
1416 | 1452 | "contents": contents, |
1417 | 1453 | "generationConfig": gen_config, |
1418 | | - "session_id": self._session_id, |
| 1454 | + "session_id": self._generate_stable_session_id(contents), |
1419 | 1455 | }, |
1420 | 1456 | } |
1421 | 1457 |
|
@@ -1751,7 +1787,7 @@ async def count_tokens( |
1751 | 1787 | "user_prompt_id": self._generate_user_prompt_id(), |
1752 | 1788 | "request": { |
1753 | 1789 | "contents": contents, |
1754 | | - "session_id": self._session_id, |
| 1790 | + "session_id": self._generate_stable_session_id(contents), |
1755 | 1791 | }, |
1756 | 1792 | } |
1757 | 1793 |
|
|
0 commit comments