refactor: extract build_session_key() as single source of truth

The session key construction logic was duplicated in 4 places
(session.py + 3 inline copies in run.py), which is exactly the
kind of drift that caused issue #349 in the first place.

Extracted build_session_key() as a public function in session.py.
SessionStore._generate_session_key() now delegates to it, and all
inline key construction in run.py has been replaced with calls to
the shared function. Tests updated to test the function directly.
This commit is contained in:
teknium1
2026-03-04 03:34:45 -08:00
parent bfd08d5648
commit 3db3d60368
3 changed files with 35 additions and 71 deletions

View File

@@ -118,6 +118,7 @@ from gateway.session import (
SessionContext,
build_session_context,
build_session_context_prompt,
build_session_key,
)
from gateway.delivery import DeliveryRouter, DeliveryTarget
from gateway.platforms.base import BasePlatformAdapter, MessageEvent, MessageType
@@ -637,12 +638,7 @@ class GatewayRunner:
# PRIORITY: If an agent is already running for this session, interrupt it
# immediately. This is before command parsing to minimize latency -- the
# user's "stop" message reaches the agent as fast as possible.
if source.chat_type != "dm":
_quick_key = f"agent:main:{source.platform.value}:{source.chat_type}:{source.chat_id}"
elif source.platform.value == "whatsapp" and source.chat_id:
_quick_key = f"agent:main:{source.platform.value}:dm:{source.chat_id}"
else:
_quick_key = f"agent:main:{source.platform.value}:dm"
_quick_key = build_session_key(source)
if _quick_key in self._running_agents:
running_agent = self._running_agents[_quick_key]
logger.debug("PRIORITY interrupt for session %s", _quick_key[:20])
@@ -720,12 +716,7 @@ class GatewayRunner:
logger.debug("Skill command check failed (non-fatal): %s", e)
# Check for pending exec approval responses
if source.chat_type != "dm":
session_key_preview = f"agent:main:{source.platform.value}:{source.chat_type}:{source.chat_id}"
elif source.platform and source.platform.value == "whatsapp" and source.chat_id:
session_key_preview = f"agent:main:{source.platform.value}:dm:{source.chat_id}"
else:
session_key_preview = f"agent:main:{source.platform.value}:dm"
session_key_preview = build_session_key(source)
if session_key_preview in self._pending_approvals:
user_text = event.text.strip().lower()
if user_text in ("yes", "y", "approve", "ok", "go", "do it"):
@@ -1362,12 +1353,7 @@ class GatewayRunner:
async def _handle_usage_command(self, event: MessageEvent) -> str:
"""Handle /usage command -- show token usage for the session's last agent run."""
source = event.source
if source.chat_type != "dm":
session_key = f"agent:main:{source.platform.value}:{source.chat_type}:{source.chat_id}"
elif source.platform.value == "whatsapp" and source.chat_id:
session_key = f"agent:main:{source.platform.value}:dm:{source.chat_id}"
else:
session_key = f"agent:main:{source.platform.value}:dm"
session_key = build_session_key(source)
agent = self._running_agents.get(session_key)
if agent and hasattr(agent, "session_total_tokens") and agent.session_api_calls > 0:

View File

@@ -281,6 +281,20 @@ class SessionEntry:
)
def build_session_key(source: SessionSource) -> str:
"""Build a deterministic session key from a message source.
This is the single source of truth for session key construction.
WhatsApp DMs include chat_id (multi-user), other DMs do not (single owner).
"""
platform = source.platform.value
if source.chat_type == "dm":
if platform == "whatsapp" and source.chat_id:
return f"agent:main:{platform}:dm:{source.chat_id}"
return f"agent:main:{platform}:dm"
return f"agent:main:{platform}:{source.chat_type}:{source.chat_id}"
class SessionStore:
"""
Manages session storage and retrieval.
@@ -337,16 +351,7 @@ class SessionStore:
def _generate_session_key(self, source: SessionSource) -> str:
"""Generate a session key from a source."""
platform = source.platform.value
if source.chat_type == "dm":
# WhatsApp DMs come from different people, each needs its own session.
# Other platforms (Telegram, Discord) have a single DM with the bot owner.
if platform == "whatsapp" and source.chat_id:
return f"agent:main:{platform}:dm:{source.chat_id}"
return f"agent:main:{platform}:dm"
else:
return f"agent:main:{platform}:{source.chat_type}:{source.chat_id}"
return build_session_key(source)
def _should_reset(self, entry: SessionEntry, source: SessionSource) -> bool:
"""