biosamples-mcp/src/middleware/cache_middleware.py at 97f4433e850c54b616f77d4bb0a6efdf0daddcd2 · EBIBioSamples/biosamples-mcp · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import hashlib
import json

from core.logger import get_logger
from middleware.base.tool_middleware import ToolMiddleware, NextHandler
from orchestrator.request_context import RequestContext

logger = get_logger(__name__)

class CacheMiddleware(ToolMiddleware):
    def __init__(
        self,
        redis_client,
        ttl_seconds: int = 300,
    ):
        self.redis = redis_client
        self.ttl_seconds = ttl_seconds

    async def process(
        self,
        context: RequestContext,
        next_handler: NextHandler,
    ):
        logger.info("RedisCacheMiddleware")
        non_cacheable = {
            "biosamples.submit_sample",
            "biosamples.prepare_submission",
        }

        if context.tool_name in non_cacheable:
            return await next_handler(context)

        if self.redis is None:
            return await next_handler(context)

        cache_key = self._cache_key(context)

        logger.info(
            "Cache check started",
            extra={
                "extra_fields": {
                    "event": "cache_check_started",
                    "tool": context.tool_name,
                    "requestId": context.request_id,
                }
            },
        )

        try:
            cached = await self.redis.get(cache_key)
        except Exception as error:
            logger.warning(
                "Redis cache read failed. Proceeding without cache.",
                extra={
                    "extra_fields": {
                        "event": "cache_read_failed",
                        "tool": context.tool_name,
                        "requestId": context.request_id,
                        "error": str(error),
                    }
                },
            )
            cached = None

        if cached:
            logger.info(
                "Cache hit",
                extra={
                    "extra_fields": {
                        "event": "cache_hit",
                        "tool": context.tool_name,
                        "requestId": context.request_id,
                    }
                },
            )
            return json.loads(cached)

        response = await next_handler(context)

        try:
            await self.redis.setex(
                cache_key,
                self.ttl_seconds,
                json.dumps(response, default=str),
            )
        except Exception as error:
            logger.warning(
                "Redis cache write failed. Returning live response.",
                extra={
                    "extra_fields": {
                        "event": "cache_write_failed",
                        "tool": context.tool_name,
                        "requestId": context.request_id,
                        "error": str(error),
                    }
                },
            )

        logger.info(
            "Cache miss",
            extra={
                "extra_fields": {
                    "event": "cache_miss",
                    "tool": context.tool_name,
                    "cache": {
                        "hit": False,
                        "type": "redis",
                        "ttlSeconds": self.ttl_seconds,
                    },
                }
            },
        )
        return response

    def _cache_key(self, context: RequestContext) -> str:
        raw = json.dumps(
            {
                "tool": context.tool_name,
                "payload": context.payload,
            },
            sort_keys=True,
            default=str,
        )

        digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()

        return f"tool-cache:{context.tool_name}:{digest}"