| import asyncio |
| from collections import OrderedDict |
|
|
| from a2a.server.context import ServerCallContext |
| from a2a.server.owner_resolver import resolve_user_scope |
| from a2a.server.tasks import InMemoryTaskStore |
| from a2a.types import Task, TaskState |
|
|
| TERMINAL_STATES = { |
| TaskState.TASK_STATE_COMPLETED, |
| TaskState.TASK_STATE_CANCELED, |
| TaskState.TASK_STATE_FAILED, |
| TaskState.TASK_STATE_REJECTED, |
| } |
|
|
|
|
| class GameTaskStore(InMemoryTaskStore): |
| def __init__(self, *, max_completed: int = 64): |
| if max_completed < 1: |
| raise ValueError("已结束任务的保留数量必须大于零") |
| super().__init__() |
| self.max_completed = max_completed |
| self.completed: OrderedDict[tuple[str, str], ServerCallContext] = OrderedDict() |
| self.retention_lock = asyncio.Lock() |
|
|
| async def save(self, task: Task, context: ServerCallContext) -> None: |
| async with self.retention_lock: |
| await super().save(task, context) |
| key = (resolve_user_scope(context), task.id) |
| self.completed.pop(key, None) |
| if task.status.state in TERMINAL_STATES: |
| self.completed[key] = context |
| |
| while len(self.completed) > self.max_completed: |
| (_, task_id), owner_context = self.completed.popitem(last=False) |
| await super().delete(task_id, owner_context) |
|
|
| async def delete(self, task_id: str, context: ServerCallContext) -> None: |
| async with self.retention_lock: |
| await super().delete(task_id, context) |
| self.completed.pop((resolve_user_scope(context), task_id), None) |
|
|
