160 lines
4.7 KiB
Python
160 lines
4.7 KiB
Python
"""SharedWorkspace - Agent 间共享工作空间
|
||
|
||
基于 Redis 的共享状态存储，支持读写、订阅、锁操作。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import time
|
||
from typing import Any
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class SharedWorkspace:
|
||
"""Agent 间共享工作空间
|
||
|
||
基于 Redis 的共享状态存储，支持：
|
||
- write/read: 读写共享数据
|
||
- lock/unlock: 分布式锁
|
||
- 版本控制：每次写入递增版本号
|
||
"""
|
||
|
||
def __init__(self, redis_client: Any = None, prefix: str = "workspace"):
|
||
"""
|
||
Args:
|
||
redis_client: aioredis.Redis 实例，None 时使用内存字典
|
||
prefix: Redis key 前缀
|
||
"""
|
||
self._redis = redis_client
|
||
self._prefix = prefix
|
||
self._local_store: dict[str, dict[str, Any]] = {}
|
||
self._locks: dict[str, str] = {} # key -> lock_owner
|
||
|
||
def _make_key(self, key: str) -> str:
|
||
return f"{self._prefix}:{key}"
|
||
|
||
async def write(
|
||
self, key: str, value: Any, agent_id: str, ttl: int | None = None
|
||
) -> int:
|
||
"""写入共享数据
|
||
|
||
Args:
|
||
key: 数据键
|
||
value: 数据值
|
||
agent_id: 写入者 ID
|
||
ttl: 过期时间（秒），None 表示不过期
|
||
|
||
Returns:
|
||
版本号
|
||
"""
|
||
entry = {
|
||
"value": value,
|
||
"agent_id": agent_id,
|
||
"version": await self._get_version(key) + 1,
|
||
"timestamp": time.time(),
|
||
}
|
||
|
||
if self._redis:
|
||
redis_key = self._make_key(key)
|
||
data = json.dumps(entry, default=str)
|
||
if ttl:
|
||
await self._redis.setex(redis_key, ttl, data)
|
||
else:
|
||
await self._redis.set(redis_key, data)
|
||
else:
|
||
self._local_store[key] = entry
|
||
|
||
return entry["version"]
|
||
|
||
async def read(self, key: str) -> dict[str, Any] | None:
|
||
"""读取共享数据
|
||
|
||
Returns:
|
||
{"value": ..., "agent_id": ..., "version": ..., "timestamp": ...} 或 None
|
||
"""
|
||
if self._redis:
|
||
redis_key = self._make_key(key)
|
||
data = await self._redis.get(redis_key)
|
||
if data is None:
|
||
return None
|
||
return json.loads(data)
|
||
else:
|
||
return self._local_store.get(key)
|
||
|
||
async def delete(self, key: str) -> bool:
|
||
"""删除共享数据"""
|
||
if self._redis:
|
||
redis_key = self._make_key(key)
|
||
result = await self._redis.delete(redis_key)
|
||
return result > 0
|
||
else:
|
||
return self._local_store.pop(key, None) is not None
|
||
|
||
async def lock(self, key: str, agent_id: str, timeout: float = 30.0) -> bool:
|
||
"""获取分布式锁
|
||
|
||
Args:
|
||
key: 要锁定的数据键
|
||
agent_id: 请求锁的 Agent ID
|
||
timeout: 锁超时时间（秒）
|
||
|
||
Returns:
|
||
是否成功获取锁
|
||
"""
|
||
lock_key = f"{self._prefix}:lock:{key}"
|
||
|
||
if self._redis:
|
||
# Redis SET with NX (only if not exists) and EX (expiry)
|
||
result = await self._redis.set(lock_key, agent_id, nx=True, ex=int(timeout))
|
||
return result is not None
|
||
else:
|
||
if key in self._locks:
|
||
return False
|
||
self._locks[key] = agent_id
|
||
return True
|
||
|
||
async def unlock(self, key: str, agent_id: str) -> bool:
|
||
"""释放分布式锁
|
||
|
||
只有锁的持有者才能释放锁。
|
||
"""
|
||
lock_key = f"{self._prefix}:lock:{key}"
|
||
|
||
if self._redis:
|
||
current_owner = await self._redis.get(lock_key)
|
||
if current_owner and current_owner.decode() == agent_id:
|
||
await self._redis.delete(lock_key)
|
||
return True
|
||
return False
|
||
else:
|
||
if self._locks.get(key) == agent_id:
|
||
del self._locks[key]
|
||
return True
|
||
return False
|
||
|
||
async def _get_version(self, key: str) -> int:
|
||
"""获取当前版本号"""
|
||
data = await self.read(key)
|
||
if data is None:
|
||
return 0
|
||
return data.get("version", 0)
|
||
|
||
async def list_keys(self) -> list[str]:
|
||
"""列出所有键"""
|
||
if self._redis:
|
||
pattern = f"{self._prefix}:*"
|
||
keys = []
|
||
async for key in self._redis.scan_iter(match=pattern):
|
||
# Strip prefix
|
||
k = key.decode() if isinstance(key, bytes) else key
|
||
k = k[len(self._prefix) + 1:] # Remove "prefix:"
|
||
# Skip lock keys
|
||
if not k.startswith("lock:"):
|
||
keys.append(k)
|
||
return keys
|
||
else:
|
||
return list(self._local_store.keys())
|