Skip to content

fix(concurrency): fix issue for data leakage across sessions #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
junxiant merged 1 commit into from
Sep 4, 2026
Merged

fix(concurrency): fix issue for data leakage across sessions #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 18 additions & 2 deletions src/agent/orchestrator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations
import re
import threading
from typing import Optional

from src.core.models import ResponseEnvelope, ActionEnum, Citation
Expand All @@ -16,8 +17,23 @@ def __init__(self, retriever: Optional[HybridRetriever] = None, llm: Optional[LL
self.retriever = retriever or HybridRetriever()
self.llm = llm or LLMClient()
self.session_manager = SessionStateManager
self.last_retrieved_chunks = []
self.last_raw_envelope: Optional[ResponseEnvelope] = None
self._local = threading.local()

@property
def last_retrieved_chunks(self) -> list:
return getattr(self._local, "last_retrieved_chunks", [])

@last_retrieved_chunks.setter
def last_retrieved_chunks(self, val: list) -> None:
self._local.last_retrieved_chunks = val

@property
def last_raw_envelope(self) -> Optional[ResponseEnvelope]:
return getattr(self._local, "last_raw_envelope", None)

@last_raw_envelope.setter
def last_raw_envelope(self, val: Optional[ResponseEnvelope]) -> None:
self._local.last_raw_envelope = val

def process_turn(self, session_id: str, user_message: str) -> ResponseEnvelope:
self.last_retrieved_chunks = []
Expand Down
17 changes: 11 additions & 6 deletions src/rag/llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import json
import os
import re
import threading
import time

from openai import OpenAI
Expand Down Expand Up @@ -196,6 +197,7 @@ def __init__(self):
# "openrouter" is the legacy config name for live mode.
self.mode = "live" if LLM_MODE == "openrouter" else LLM_MODE
self.client: OpenAI | None = None
self._rate_limit_lock = threading.Lock()
self.last_call_time = 0.0
self.rate_limit_delay = LLM_RATE_LIMIT_DELAY
self.fallback_models = OPENROUTER_FALLBACK_MODELS
Expand Down Expand Up @@ -303,15 +305,18 @@ def _call_with_fallback(self, messages: list[dict[str, str]]) -> tuple[ResponseE
logger.error("No OpenRouter client available for live call.")
return None, ""

elapsed = time.time() - self.last_call_time
if elapsed < self.rate_limit_delay:
sleep_duration = self.rate_limit_delay - elapsed
logger.debug(f"Rate limit buffer: sleeping for {sleep_duration:.2f}s before LLM call")
time.sleep(sleep_duration)
with self._rate_limit_lock:
elapsed = time.time() - self.last_call_time
if elapsed < self.rate_limit_delay:
sleep_duration = self.rate_limit_delay - elapsed
logger.debug(f"Rate limit buffer: sleeping for {sleep_duration:.2f}s before LLM call")
time.sleep(sleep_duration)
self.last_call_time = time.time()

for model_candidate in self.fallback_models:
try:
self.last_call_time = time.time()
with self._rate_limit_lock:
self.last_call_time = time.time()
resp = self.client.chat.completions.create(
model=model_candidate,
messages=messages,
Expand Down
168 changes: 98 additions & 70 deletions src/state/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
from __future__ import annotations
import json
import os
import sqlite3
import threading
import time
from contextlib import closing
from contextlib import closing, contextmanager
from pathlib import Path
from typing import Optional, Any, List, Union

Expand All @@ -17,63 +19,93 @@
class SessionStateManager:
_db_path: Path = SQLITE_DB_PATH
_initialized: bool = False
_init_lock: threading.Lock = threading.Lock()
_pg_pool: Optional[Any] = None
_use_postgres: bool = False

@classmethod
def is_postgres(cls) -> bool:
return DB_BACKEND == "postgres" and bool(DATABASE_URL)
return DB_BACKEND == "postgres" and bool(DATABASE_URL) and cls._use_postgres

@classmethod
@contextmanager
def _get_pg_conn(cls):
if cls._pg_pool is None:
raise RuntimeError("PostgreSQL connection pool is not initialized")
conn = cls._pg_pool.getconn()
try:
yield conn
finally:
cls._pg_pool.putconn(conn)

@classmethod
def _init_db_once(cls) -> None:
if cls._initialized:
return
if cls.is_postgres():
try:
import psycopg2
with closing(psycopg2.connect(DATABASE_URL)) as conn:
with conn:
with conn.cursor() as cur:
cur.execute("""
CREATE TABLE IF NOT EXISTS sessions (
session_id VARCHAR(255) PRIMARY KEY,
identified_model VARCHAR(255),
attempted_steps TEXT,
pending_confirmation VARCHAR(255),
dialogue_window TEXT,
created_at DOUBLE PRECISION,
updated_at DOUBLE PRECISION,
reported_issue TEXT,
confirmed_facts TEXT,
turns_count INTEGER DEFAULT 0,
is_escalated INTEGER DEFAULT 0,
is_resolved INTEGER DEFAULT 0
);
""")
cls._initialized = True
logger.info(f"Initialized PostgreSQL session storage: {DATABASE_URL}")
with cls._init_lock:
if cls._initialized:
return
except Exception as e:
logger.warning(f"Failed to connect to PostgreSQL ({e}). Falling back to SQLite at {cls._db_path}")

cls._db_path.parent.mkdir(parents=True, exist_ok=True)
with closing(sqlite3.connect(str(cls._db_path), check_same_thread=False)) as conn:
with conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS sessions (
session_id TEXT PRIMARY KEY,
identified_model TEXT,
attempted_steps TEXT,
pending_confirmation TEXT,
dialogue_window TEXT,
created_at REAL,
updated_at REAL,
reported_issue TEXT,
confirmed_facts TEXT,
turns_count INTEGER DEFAULT 0,
is_escalated INTEGER DEFAULT 0,
is_resolved INTEGER DEFAULT 0
if DB_BACKEND == "postgres" and bool(DATABASE_URL):
try:
from psycopg2 import pool
cls._pg_pool = pool.ThreadedConnectionPool(
minconn=1,
maxconn=int(os.getenv("DB_POOL_MAX", "10")),
dsn=DATABASE_URL,
)
""")
cls._initialized = True
with cls._get_pg_conn() as conn:
with conn:
with conn.cursor() as cur:
cur.execute("""
CREATE TABLE IF NOT EXISTS sessions (
session_id VARCHAR(255) PRIMARY KEY,
identified_model VARCHAR(255),
attempted_steps TEXT,
pending_confirmation VARCHAR(255),
dialogue_window TEXT,
created_at DOUBLE PRECISION,
updated_at DOUBLE PRECISION,
reported_issue TEXT,
confirmed_facts TEXT,
turns_count INTEGER DEFAULT 0,
is_escalated INTEGER DEFAULT 0,
is_resolved INTEGER DEFAULT 0
);
""")
cls._use_postgres = True
cls._initialized = True
logger.info(f"Initialized PostgreSQL session storage with ThreadedConnectionPool: {DATABASE_URL}")
return
except Exception as e:
cls._use_postgres = False
if cls._pg_pool:
try:
cls._pg_pool.closeall()
except Exception:
pass
cls._pg_pool = None
logger.warning(f"Failed to connect to PostgreSQL ({e}). Falling back to SQLite at {cls._db_path}")

cls._db_path.parent.mkdir(parents=True, exist_ok=True)
with closing(sqlite3.connect(str(cls._db_path), timeout=30.0, check_same_thread=False)) as conn:
with conn:
conn.execute("""
CREATE TABLE IF NOT EXISTS sessions (
session_id TEXT PRIMARY KEY,
identified_model TEXT,
attempted_steps TEXT,
pending_confirmation TEXT,
dialogue_window TEXT,
created_at REAL,
updated_at REAL,
reported_issue TEXT,
confirmed_facts TEXT,
turns_count INTEGER DEFAULT 0,
is_escalated INTEGER DEFAULT 0,
is_resolved INTEGER DEFAULT 0
)
""")
cls._initialized = True

@classmethod
def _row_to_state(cls, row: Union[sqlite3.Row, dict]) -> SessionState:
Expand Down Expand Up @@ -102,23 +134,22 @@ def get_or_create(cls, session_id: str) -> SessionState:
cls._init_db_once()
if cls.is_postgres():
try:
import psycopg2
from psycopg2.extras import RealDictCursor
with closing(psycopg2.connect(DATABASE_URL)) as conn:
with cls._get_pg_conn() as conn:
with conn.cursor(cursor_factory=RealDictCursor) as cur:
cur.execute("SELECT * FROM sessions WHERE session_id = %s", (session_id,))
row = cur.fetchone()
if row is not None:
return cls._row_to_state(row)
except Exception as e:
logger.error(f"Error fetching session from PostgreSQL: {e}")
else:
with closing(sqlite3.connect(str(cls._db_path), check_same_thread=False)) as conn:
conn.row_factory = sqlite3.Row
cursor = conn.execute("SELECT * FROM sessions WHERE session_id = ?", (session_id,))
row = cursor.fetchone()
if row is not None:
return cls._row_to_state(row)
logger.error(f"Error fetching session from PostgreSQL: {e}. Falling back to SQLite.")

with closing(sqlite3.connect(str(cls._db_path), timeout=30.0, check_same_thread=False)) as conn:
conn.row_factory = sqlite3.Row
cursor = conn.execute("SELECT * FROM sessions WHERE session_id = ?", (session_id,))
row = cursor.fetchone()
if row is not None:
return cls._row_to_state(row)

new_state = SessionState(session_id=session_id)
cls.update_session(new_state)
Expand Down Expand Up @@ -228,8 +259,7 @@ def update_session(cls, session: SessionState) -> None:

if cls.is_postgres():
try:
import psycopg2
with closing(psycopg2.connect(DATABASE_URL)) as conn:
with cls._get_pg_conn() as conn:
with conn:
with conn.cursor() as cur:
cur.execute("""
Expand All @@ -252,9 +282,9 @@ def update_session(cls, session: SessionState) -> None:
""", params)
return
except Exception as e:
logger.error(f"Error updating session in PostgreSQL: {e}")
logger.error(f"Error updating session in PostgreSQL: {e}. Falling back to SQLite.")

with closing(sqlite3.connect(str(cls._db_path), check_same_thread=False)) as conn:
with closing(sqlite3.connect(str(cls._db_path), timeout=30.0, check_same_thread=False)) as conn:
with conn:
conn.execute("""
INSERT INTO sessions (
Expand All @@ -280,16 +310,15 @@ def clear_session(cls, session_id: str) -> None:
cls._init_db_once()
if cls.is_postgres():
try:
import psycopg2
with closing(psycopg2.connect(DATABASE_URL)) as conn:
with cls._get_pg_conn() as conn:
with conn:
with conn.cursor() as cur:
cur.execute("DELETE FROM sessions WHERE session_id = %s", (session_id,))
return
except Exception as e:
logger.error(f"Error deleting session from PostgreSQL: {e}")
logger.error(f"Error deleting session from PostgreSQL: {e}. Falling back to SQLite.")

with closing(sqlite3.connect(str(cls._db_path), check_same_thread=False)) as conn:
with closing(sqlite3.connect(str(cls._db_path), timeout=30.0, check_same_thread=False)) as conn:
with conn:
conn.execute("DELETE FROM sessions WHERE session_id = ?", (session_id,))

Expand All @@ -298,16 +327,15 @@ def reset_all(cls) -> None:
cls._init_db_once()
if cls.is_postgres():
try:
import psycopg2
with closing(psycopg2.connect(DATABASE_URL)) as conn:
with cls._get_pg_conn() as conn:
with conn:
with conn.cursor() as cur:
cur.execute("DELETE FROM sessions")
return
except Exception as e:
logger.error(f"Error resetting sessions in PostgreSQL: {e}")
logger.error(f"Error resetting sessions in PostgreSQL: {e}. Falling back to SQLite.")

with closing(sqlite3.connect(str(cls._db_path), check_same_thread=False)) as conn:
with closing(sqlite3.connect(str(cls._db_path), timeout=30.0, check_same_thread=False)) as conn:
with conn:
conn.execute("DELETE FROM sessions")

Expand Down
Loading