File size: 1,714 Bytes
ad88834 | 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 | 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)
|