mcp/services/gateway_session_store.py

import json
from datetime import datetime, timezone

from utils.security import hash_gateway_session_id


class GatewaySessionStore:
    def __init__(self, client, prefix='fms:mcp:gateway:', ttl_seconds=0):
        self.client = client
        self.prefix = str(prefix or 'fms:mcp:gateway:')
        self.ttl_seconds = int(ttl_seconds or 0)

    def key_for(self, gateway_session_id):
        return self.prefix.rstrip(':') + ':session:' + hash_gateway_session_id(gateway_session_id)

    def save(self, gateway_session_id, session):
        payload = dict(session or {})
        payload['gateway_session_id_hash'] = hash_gateway_session_id(gateway_session_id)
        set_kwargs = {}
        if self.ttl_seconds > 0:
            set_kwargs['ex'] = self.ttl_seconds
        self.client.set(
            self.key_for(gateway_session_id),
            json.dumps(payload, ensure_ascii=False),
            **set_kwargs,
        )
        return payload

    def get(self, gateway_session_id):
        raw = self.client.get(self.key_for(gateway_session_id))
        if raw is None or raw == '':
            return None
        if isinstance(raw, bytes):
            raw = raw.decode('utf-8')
        return json.loads(raw)

    def touch_session(self, gateway_session_id):
        session = self.get(gateway_session_id)
        if not session:
            return None
        session['last_access_time'] = datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z')
        return self.save(gateway_session_id, session)

    def delete(self, gateway_session_id):
        return self.client.delete(self.key_for(gateway_session_id))