forked from modelcontextprotocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathprogress.py
More file actions
40 lines (29 loc) · 1.03 KB
/
progress.py
File metadata and controls
40 lines (29 loc) · 1.03 KB
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
from contextlib import contextmanager
from dataclasses import dataclass, field
from pydantic import BaseModel
from mcp_python.shared.context import RequestContext
from mcp_python.shared.session import BaseSession
from mcp_python.types import ProgressToken
class Progress(BaseModel):
progress: float
total: float | None
@dataclass
class ProgressContext:
session: BaseSession
progress_token: ProgressToken
total: float | None
current: float = field(default=0.0, init=False)
async def progress(self, amount: float) -> None:
self.current += amount
await self.session.send_progress_notification(
self.progress_token, self.current, total=self.total
)
@contextmanager
def progress(ctx: RequestContext, total: float | None = None):
if ctx.meta is None or ctx.meta.progressToken is None:
raise ValueError("No progress token provided")
progress_ctx = ProgressContext(ctx.session, ctx.meta.progressToken, total)
try:
yield progress_ctx
finally:
pass