1717from mcp .shared .experimental .tasks .store import TaskStore
1818from mcp .types import Result , Task , TaskMetadata , TaskStatus
1919
20+ CLEANUP_INTERVAL_SECONDS = 1.0
21+
2022
2123@dataclass
2224class StoredTask :
23- """Internal storage representation of a task."""
24-
2525 task : Task
2626 result : Result | None = None
27- # Time when this task should be removed (None = never)
2827 expires_at : datetime | None = field (default = None )
2928
3029
@@ -49,21 +48,26 @@ def __init__(self, page_size: int = 10) -> None:
4948 self ._tasks : dict [str , StoredTask ] = {}
5049 self ._page_size = page_size
5150 self ._update_events : dict [str , anyio .Event ] = {}
51+ self ._last_cleanup : datetime | None = None
5252
5353 def _calculate_expiry (self , ttl_ms : int | None ) -> datetime | None :
54- """Calculate expiry time from TTL in milliseconds."""
5554 if ttl_ms is None :
5655 return None
5756 return datetime .now (timezone .utc ) + timedelta (milliseconds = ttl_ms )
5857
5958 def _is_expired (self , stored : StoredTask ) -> bool :
60- """Check if a task has expired."""
6159 if stored .expires_at is None :
6260 return False
6361 return datetime .now (timezone .utc ) >= stored .expires_at
6462
6563 def _cleanup_expired (self ) -> None :
66- """Remove all expired tasks. Called lazily during access operations."""
64+ now = datetime .now (timezone .utc )
65+ if self ._last_cleanup is not None :
66+ elapsed = (now - self ._last_cleanup ).total_seconds ()
67+ if elapsed < CLEANUP_INTERVAL_SECONDS :
68+ return
69+
70+ self ._last_cleanup = now
6771 expired_ids = [task_id for task_id , stored in self ._tasks .items () if self ._is_expired (stored )]
6872 for task_id in expired_ids :
6973 del self ._tasks [task_id ]
@@ -73,34 +77,21 @@ async def create_task(
7377 metadata : TaskMetadata ,
7478 task_id : str | None = None ,
7579 ) -> Task :
76- """Create a new task with the given metadata."""
77- # Cleanup expired tasks on access
7880 self ._cleanup_expired ()
79-
8081 task = create_task_state (metadata , task_id )
8182
8283 if task .taskId in self ._tasks :
8384 raise ValueError (f"Task with ID { task .taskId } already exists" )
8485
85- stored = StoredTask (
86- task = task ,
87- expires_at = self ._calculate_expiry (metadata .ttl ),
88- )
86+ stored = StoredTask (task = task , expires_at = self ._calculate_expiry (metadata .ttl ))
8987 self ._tasks [task .taskId ] = stored
90-
91- # Return a copy to prevent external modification
9288 return Task (** task .model_dump ())
9389
9490 async def get_task (self , task_id : str ) -> Task | None :
95- """Get a task by ID."""
96- # Cleanup expired tasks on access
9791 self ._cleanup_expired ()
98-
9992 stored = self ._tasks .get (task_id )
10093 if stored is None :
10194 return None
102-
103- # Return a copy to prevent external modification
10495 return Task (** stored .task .model_dump ())
10596
10697 async def update_task (
@@ -109,12 +100,10 @@ async def update_task(
109100 status : TaskStatus | None = None ,
110101 status_message : str | None = None ,
111102 ) -> Task :
112- """Update a task's status and/or message."""
113103 stored = self ._tasks .get (task_id )
114104 if stored is None :
115105 raise ValueError (f"Task with ID { task_id } not found" )
116106
117- # Per spec: Terminal states MUST NOT transition to any other status
118107 if status is not None and status != stored .task .status and is_terminal (stored .task .status ):
119108 raise ValueError (f"Cannot transition from terminal status '{ stored .task .status } '" )
120109
@@ -126,94 +115,69 @@ async def update_task(
126115 if status_message is not None :
127116 stored .task .statusMessage = status_message
128117
129- # Update lastUpdatedAt on any change
130118 stored .task .lastUpdatedAt = datetime .now (timezone .utc )
131119
132- # If task is now terminal and has TTL, reset expiry timer
133120 if status is not None and is_terminal (status ) and stored .task .ttl is not None :
134121 stored .expires_at = self ._calculate_expiry (stored .task .ttl )
135122
136- # Notify waiters if status changed
137123 if status_changed :
138124 await self .notify_update (task_id )
139125
140126 return Task (** stored .task .model_dump ())
141127
142128 async def store_result (self , task_id : str , result : Result ) -> None :
143- """Store the result for a task."""
144129 stored = self ._tasks .get (task_id )
145130 if stored is None :
146131 raise ValueError (f"Task with ID { task_id } not found" )
147-
148132 stored .result = result
149133
150134 async def get_result (self , task_id : str ) -> Result | None :
151- """Get the stored result for a task."""
152135 stored = self ._tasks .get (task_id )
153- if stored is None :
154- return None
155-
156- return stored .result
136+ return stored .result if stored else None
157137
158138 async def list_tasks (
159139 self ,
160140 cursor : str | None = None ,
161141 ) -> tuple [list [Task ], str | None ]:
162- """List tasks with pagination."""
163- # Cleanup expired tasks on access
164142 self ._cleanup_expired ()
165-
166143 all_task_ids = list (self ._tasks .keys ())
167144
168145 start_index = 0
169146 if cursor is not None :
170147 try :
171- cursor_index = all_task_ids .index (cursor )
172- start_index = cursor_index + 1
148+ start_index = all_task_ids .index (cursor ) + 1
173149 except ValueError :
174150 raise ValueError (f"Invalid cursor: { cursor } " )
175151
176152 page_task_ids = all_task_ids [start_index : start_index + self ._page_size ]
177153 tasks = [Task (** self ._tasks [tid ].task .model_dump ()) for tid in page_task_ids ]
178154
179- # Determine next cursor
180155 next_cursor = None
181156 if start_index + self ._page_size < len (all_task_ids ) and page_task_ids :
182157 next_cursor = page_task_ids [- 1 ]
183158
184159 return tasks , next_cursor
185160
186161 async def delete_task (self , task_id : str ) -> bool :
187- """Delete a task."""
188162 if task_id not in self ._tasks :
189163 return False
190-
191164 del self ._tasks [task_id ]
192165 return True
193166
194167 async def wait_for_update (self , task_id : str ) -> None :
195- """Wait until the task status changes."""
196168 if task_id not in self ._tasks :
197169 raise ValueError (f"Task with ID { task_id } not found" )
198-
199- # Create a fresh event for waiting (anyio.Event can't be cleared)
200170 self ._update_events [task_id ] = anyio .Event ()
201- event = self ._update_events [task_id ]
202- await event .wait ()
171+ await self ._update_events [task_id ].wait ()
203172
204173 async def notify_update (self , task_id : str ) -> None :
205- """Signal that a task has been updated."""
206174 if task_id in self ._update_events :
207175 self ._update_events [task_id ].set ()
208176
209- # --- Testing/debugging helpers ---
210-
211177 def cleanup (self ) -> None :
212- """Cleanup all tasks (useful for testing or graceful shutdown)."""
213178 self ._tasks .clear ()
214179 self ._update_events .clear ()
215180
216181 def get_all_tasks (self ) -> list [Task ]:
217- """Get all tasks (useful for debugging). Returns copies to prevent modification."""
218182 self ._cleanup_expired ()
219183 return [Task (** stored .task .model_dump ()) for stored in self ._tasks .values ()]
0 commit comments