Add durable independent Agent tasks and delegated child authorization
This commit is contained in:
magicwang1111
@@ -9,6 +9,7 @@ blacklist.txt
|
||||
secret_key.txt
|
||||
history.sqlite3*
|
||||
scheduler.sqlite3*
|
||||
agent_tasks.sqlite3*
|
||||
scheduler_secret.txt
|
||||
*.lock
|
||||
cloud_staging/
|
||||
|
||||
@@ -15,6 +15,8 @@ from .utils.frontend_reload import install_frontend_reload
|
||||
from .utils.api_recovery import ApiRecoveryManager
|
||||
from .utils.history_store import HistoryStore
|
||||
from .utils.task_sessions import TaskSessions
|
||||
from .utils.agent_tasks import AgentTaskStore
|
||||
from .utils.agent_task_routes import AgentTaskRoutes
|
||||
|
||||
instance = PromptServer.instance
|
||||
app = instance.app
|
||||
@@ -98,6 +100,25 @@ def current_task_environment():
|
||||
|
||||
|
||||
instance.account_manager_task_environment = current_task_environment
|
||||
if scheduler:
|
||||
agent_store = AgentTaskStore(
|
||||
os.path.join(os.path.dirname(__file__), config.get('agent_tasks_db', 'agent_tasks.sqlite3')),
|
||||
concurrency=config.get('agent_concurrency', 4),
|
||||
per_user=config.get('agent_concurrency_per_user', 2),
|
||||
pending_per_user=config.get('agent_pending_per_user', 20),
|
||||
)
|
||||
agent_routes = AgentTaskRoutes(agent_store, jwt_auth.task_sessions, internal_signer, scheduler, instance,
|
||||
HistoryStore(HISTORY_FILE))
|
||||
agent_routes.install(routes)
|
||||
jwt_auth.agent_tasks = agent_routes
|
||||
|
||||
def submit_agent_task(node_id, job):
|
||||
context = get_executing_context()
|
||||
if context is None:
|
||||
raise ValueError('独立 Agent 提交需要已认证的执行上下文。')
|
||||
return agent_routes.create_from_node(context.prompt_id, node_id, job)
|
||||
|
||||
instance.account_manager_submit_agent = submit_agent_task
|
||||
access_control.sync_comfy_users()
|
||||
|
||||
api_recovery = None
|
||||
|
||||
@@ -0,0 +1,169 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
from aiohttp import web
|
||||
from aiohttp.test_utils import TestClient, TestServer
|
||||
|
||||
|
||||
def module(name):
|
||||
spec = importlib.util.spec_from_file_location(name, Path(__file__).resolve().parents[1] / 'utils' / (name + '.py'))
|
||||
result = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(result)
|
||||
return result
|
||||
|
||||
|
||||
Store = module('agent_tasks').AgentTaskStore
|
||||
Routes = module('agent_task_routes').AgentTaskRoutes
|
||||
|
||||
|
||||
class StoreTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.temp.cleanup)
|
||||
self.store = Store(Path(self.temp.name) / 'tasks.db', concurrency=4, per_user=2)
|
||||
|
||||
def create(self, owner='alice', key=None):
|
||||
return self.store.create(owner, key or str(time.time_ns()), {'job_dir': self.temp.name})
|
||||
|
||||
def test_concurrent_claims_enforce_global_and_account_limits(self):
|
||||
for _ in range(6):
|
||||
self.create()
|
||||
self.create('bob')
|
||||
with ThreadPoolExecutor(max_workers=12) as pool:
|
||||
claims = [task for task in pool.map(lambda _: self.store.claim(), range(12)) if task]
|
||||
self.assertEqual(len(claims), 4)
|
||||
self.assertEqual(len({task['id'] for task in claims}), 4)
|
||||
self.assertEqual(sum(task['owner'] == 'alice' for task in claims), 2)
|
||||
|
||||
def test_submission_is_idempotent_and_pending_is_bounded(self):
|
||||
first = self.create(key='one')
|
||||
self.assertEqual(self.create(key='one')['id'], first['id'])
|
||||
self.store.pending_per_user = 1
|
||||
with self.assertRaises(ValueError):
|
||||
self.create(key='two')
|
||||
|
||||
def test_expired_worker_is_fenced_and_not_reexecuted(self):
|
||||
self.create()
|
||||
task = self.store.claim()
|
||||
with closing(self.store.connect()) as db, db:
|
||||
db.execute('UPDATE agent_tasks SET lease_until=0')
|
||||
self.assertIsNone(self.store.claim())
|
||||
self.assertEqual(self.store.get(task['id'])['state'], 'recovery_required')
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.update(task['id'], task['lease'], state='completed')
|
||||
|
||||
def test_delivery_checkpoint_reclaimed_without_rerunning_agent(self):
|
||||
self.create()
|
||||
task = self.store.claim()
|
||||
self.store.update(task['id'], task['lease'], checkpoint={'text': 'done'}, state='delivering')
|
||||
with closing(self.store.connect()) as db, db:
|
||||
db.execute('UPDATE agent_tasks SET lease_until=0')
|
||||
again = self.store.claim()
|
||||
self.assertEqual(again['checkpoint'], {'text': 'done'})
|
||||
self.assertNotEqual(again['lease'], task['lease'])
|
||||
|
||||
def test_cancel_completion_race_cannot_report_success(self):
|
||||
self.create()
|
||||
task = self.store.claim()
|
||||
self.store.cancel(task['id'], task['owner'])
|
||||
self.assertEqual(self.store.update(task['id'], task['lease'], state='completed'), 'cancelled')
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.reserve_child(task['id'], 'video', 'hash')
|
||||
|
||||
def test_event_replay_and_child_idempotency(self):
|
||||
self.create()
|
||||
task = self.store.claim()
|
||||
self.store.update(task['id'], task['lease'], events=[{'message': 'a'}, {'message': 'b'}])
|
||||
events = self.store.events(task['id'])
|
||||
self.assertEqual(self.store.events(task['id'], events[0]['seq']), events[1:])
|
||||
child, fresh = self.store.reserve_child(task['id'], 'video', 'hash')
|
||||
again, fresh_again = self.store.reserve_child(task['id'], 'video', 'hash')
|
||||
self.assertTrue(fresh)
|
||||
self.assertFalse(fresh_again)
|
||||
self.assertEqual(child['prompt_id'], again['prompt_id'])
|
||||
with self.assertRaises(ValueError):
|
||||
self.store.reserve_child(task['id'], 'video', 'different')
|
||||
|
||||
|
||||
class HTTPTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncSetUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.store = Store(Path(self.temp.name) / 'tasks.db')
|
||||
self.created = self.store.create('alice', 'one', {'job_dir': self.temp.name})
|
||||
self.task = self.store.claim()
|
||||
self.claims = {'id': 'alice', 'am_agent': self.task['id'], 'am_lease': self.task['lease']}
|
||||
self.scheduler = Mock()
|
||||
self.scheduler.get_job.return_value = None
|
||||
self.submitted = []
|
||||
sessions = SimpleNamespace(jwt_auth=Mock(), users_db=Mock())
|
||||
self.routes = Routes(self.store, sessions, Mock(verify_request=AsyncMock(return_value=False)),
|
||||
self.scheduler, SimpleNamespace(port=8180))
|
||||
|
||||
@web.middleware
|
||||
async def auth(request, handler):
|
||||
request['user_id'] = request.headers.get('X-Test-Owner', 'alice')
|
||||
if request.path in {'/prompt', '/history/foreign', '/view', '/users'}:
|
||||
return await self.routes.delegated(request, self.claims, handler)
|
||||
return await handler(request)
|
||||
|
||||
async def submit(request):
|
||||
payload = await request.json()
|
||||
self.submitted.append(payload)
|
||||
return web.json_response({'prompt_id': payload['prompt_id']})
|
||||
|
||||
app = web.Application(middlewares=[auth])
|
||||
routes = web.RouteTableDef()
|
||||
self.routes.install(routes)
|
||||
app.add_routes(routes)
|
||||
app.router.add_post('/prompt', submit)
|
||||
for path in ['/history/foreign', '/view', '/users']:
|
||||
app.router.add_get(path, lambda _: web.json_response({}))
|
||||
self.client = TestClient(TestServer(app))
|
||||
await self.client.start_server()
|
||||
|
||||
async def asyncTearDown(self):
|
||||
await self.client.close()
|
||||
self.temp.cleanup()
|
||||
|
||||
async def test_full_video_pool_still_accepts_and_deduplicates_child(self):
|
||||
self.scheduler.idle_worker_ports.return_value = []
|
||||
payload = {'prompt': {'1': {'class_type': 'Video', 'inputs': {}}}, 'client_id': 'video-1'}
|
||||
first = await self.client.post('/prompt', json=payload)
|
||||
second = await self.client.post('/prompt', json=payload)
|
||||
self.assertEqual(first.status, 200)
|
||||
self.assertEqual(await first.json(), await second.json())
|
||||
self.assertEqual(len(self.submitted), 1)
|
||||
self.scheduler.idle_worker_ports.assert_not_called()
|
||||
|
||||
async def test_owner_and_task_scope_are_enforced(self):
|
||||
response = await self.client.get('/llm/agent-tasks/' + self.task['id'], headers={'X-Test-Owner': 'bob'})
|
||||
self.assertEqual(response.status, 404)
|
||||
for path in ['/history/foreign', '/view?filename=foreign.mp4', '/users']:
|
||||
self.assertEqual((await self.client.get(path)).status, 403)
|
||||
|
||||
async def test_old_lease_cannot_submit_after_cancel(self):
|
||||
self.store.cancel(self.task['id'], 'alice')
|
||||
response = await self.client.post('/prompt', json={'prompt': {}, 'client_id': 'video'})
|
||||
self.assertEqual(response.status, 401)
|
||||
|
||||
async def test_internal_endpoint_requires_signature(self):
|
||||
response = await self.client.post(Routes.INTERNAL, json={'action': 'claim'})
|
||||
self.assertEqual(response.status, 403)
|
||||
|
||||
async def test_result_file_cannot_escape_task_directory(self):
|
||||
self.store.update(self.task['id'], self.task['lease'], state='completed',
|
||||
result={'artifacts': [{'path': str(Path(self.temp.name).parent / 'secret')}]})
|
||||
response = await self.client.get(f"/llm/agent-tasks/{self.task['id']}/files/0")
|
||||
self.assertEqual(response.status, 404)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@@ -0,0 +1,210 @@
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from aiohttp import ClientError, ClientSession, ClientTimeout, web
|
||||
|
||||
|
||||
class AgentTaskRoutes:
|
||||
INTERNAL = '/account-manager/agent-internal'
|
||||
PUBLIC = '/llm/agent-tasks'
|
||||
|
||||
def __init__(self, store, sessions, signer, scheduler, server, history=None):
|
||||
self.store = store
|
||||
self.sessions = sessions
|
||||
self.signer = signer
|
||||
self.scheduler = scheduler
|
||||
self.server = server
|
||||
self.history = history
|
||||
|
||||
def create_from_node(self, prompt_id, node_id, job):
|
||||
owner = self.sessions._local_owner(prompt_id)
|
||||
if not owner or not self.scheduler:
|
||||
raise ValueError('独立 Agent 需要启用 Account Manager 分布式队列并取得当前任务账号。')
|
||||
_, user = self.sessions.users_db.get_user(user_id=owner)
|
||||
if not user:
|
||||
raise ValueError('当前任务账号不存在。')
|
||||
job['identity'] = {'external_user_id': owner, 'external_username': user['username']}
|
||||
task = self.store.create(owner, f'{prompt_id}:{node_id}', job)
|
||||
return task['id']
|
||||
|
||||
def token(self, task):
|
||||
_, user = self.sessions.users_db.get_user(user_id=task['owner'])
|
||||
if not user:
|
||||
raise ValueError('Agent task owner no longer exists')
|
||||
token = self.sessions.jwt_auth.create_access_token(
|
||||
{'id': task['owner'], 'username': user['username'], 'am_agent': task['id'], 'am_lease': task['lease']},
|
||||
expire_minutes=24 * 60,
|
||||
)
|
||||
return {'COMFY_TASK_TOKEN': token, 'COMFY_TASK_BASE_URL': f'http://127.0.0.1:{self.server.port}'}
|
||||
|
||||
def public_task(self, task):
|
||||
children = []
|
||||
for child in self.store.children(task['id']):
|
||||
job = self.scheduler.get_job(child['prompt_id']) if self.scheduler else None
|
||||
children.append({'prompt_id': child['prompt_id'], 'state': job['status'] if job else 'submission_unknown'})
|
||||
state = task['state']
|
||||
if state in {'running', 'delivering', 'cancel_requested'} and task['lease_until'] < time.time():
|
||||
state = 'recovery_required'
|
||||
result = task['result']
|
||||
if result:
|
||||
result = dict(result)
|
||||
result['artifacts'] = [
|
||||
{k: v for k, v in item.items() if k not in {'path'}}
|
||||
| {'url': f"{self.PUBLIC}/{task['id']}/files/{index}"}
|
||||
for index, item in enumerate(result.get('artifacts', []))
|
||||
]
|
||||
return {'id': task['id'], 'state': state, 'created': task['created'], 'updated': task['updated'],
|
||||
'recoverable': state == 'recovery_required' and bool(task['checkpoint'] or children),
|
||||
'error': task['error'], 'children': children, 'result': result}
|
||||
|
||||
def owned(self, request):
|
||||
task = self.store.get(request.match_info['task_id'])
|
||||
if not task or task['owner'] != request.get('user_id'):
|
||||
raise web.HTTPNotFound()
|
||||
return task
|
||||
|
||||
async def listing(self, request):
|
||||
return web.json_response([self.public_task(task) for task in self.store.list(request['user_id'])])
|
||||
|
||||
async def detail(self, request):
|
||||
task = self.owned(request)
|
||||
try:
|
||||
after = max(0, int(request.query.get('after', '0')))
|
||||
except ValueError:
|
||||
raise web.HTTPBadRequest(reason='Invalid event sequence')
|
||||
return web.json_response(self.public_task(task) | {'events': self.store.events(task['id'], after)})
|
||||
|
||||
async def cancel(self, request):
|
||||
task = self.owned(request)
|
||||
self.store.cancel(task['id'], task['owner'])
|
||||
# Queued children can be removed immediately. Running paid jobs remain tracked.
|
||||
for child in self.store.children(task['id']):
|
||||
self.scheduler.cancel(child['prompt_id'], owner_id=task['owner'])
|
||||
current = self.scheduler.get_job(child['prompt_id'])
|
||||
if current and current['status'] == 'running' and current.get('worker_port'):
|
||||
try:
|
||||
token = self.sessions.jwt_auth.get_token_from_request(request)
|
||||
async with ClientSession(timeout=ClientTimeout(total=10)) as client:
|
||||
async with client.post(f"http://127.0.0.1:{current['worker_port']}/interrupt",
|
||||
json={'prompt_id': child['prompt_id']},
|
||||
headers={'Authorization': 'Bearer ' + token}) as response:
|
||||
await response.read()
|
||||
except (ClientError, TimeoutError):
|
||||
# Keep the actual child status visible if cancellation could not be delivered.
|
||||
pass
|
||||
return web.json_response(self.public_task(self.store.get(task['id'])))
|
||||
|
||||
async def recover(self, request):
|
||||
task = self.owned(request)
|
||||
try:
|
||||
self.store.recover(task['id'], task['owner'])
|
||||
except ValueError as exc:
|
||||
raise web.HTTPConflict(reason=str(exc))
|
||||
return web.json_response(self.public_task(self.store.get(task['id'])))
|
||||
|
||||
async def file(self, request):
|
||||
task = self.owned(request)
|
||||
try:
|
||||
index = int(request.match_info['index'])
|
||||
if index < 0:
|
||||
raise ValueError()
|
||||
record = task['result']['artifacts'][index]
|
||||
path = Path(record['path']).resolve()
|
||||
except (ValueError, IndexError, TypeError, KeyError):
|
||||
raise web.HTTPNotFound()
|
||||
if not path.is_relative_to(Path(task['job']['job_dir']).resolve()) or not path.is_file():
|
||||
raise web.HTTPNotFound()
|
||||
return web.FileResponse(path)
|
||||
|
||||
async def internal(self, request):
|
||||
body = await request.read()
|
||||
if not self.signer or not await self.signer.verify_request(request, body):
|
||||
raise web.HTTPForbidden(reason='Invalid internal signature')
|
||||
try:
|
||||
data = json.loads(body)
|
||||
if data['action'] == 'claim':
|
||||
task = self.store.claim()
|
||||
if task:
|
||||
task['environment'] = self.token(task)
|
||||
task['job']['children'] = self.store.children(task['id'])
|
||||
return web.json_response(task)
|
||||
if data['action'] == 'update':
|
||||
state = self.store.update(data['id'], data['lease'], events=data.get('events', []),
|
||||
checkpoint=data.get('checkpoint'), state=data.get('state'),
|
||||
result=data.get('result'), error=data.get('error', ''))
|
||||
return web.json_response({'state': state})
|
||||
except (ValueError, KeyError) as exc:
|
||||
raise web.HTTPConflict(reason=str(exc))
|
||||
raise web.HTTPBadRequest(reason='Unknown Agent action')
|
||||
|
||||
async def delegated(self, request, claims, handler):
|
||||
task = self.store.get(claims['am_agent'])
|
||||
if (not task or task['owner'] != claims['id'] or task['lease'] != claims.get('am_lease')
|
||||
or task['lease_until'] < time.time() or task['state'] not in {'running', 'delivering'}):
|
||||
raise web.HTTPUnauthorized(reason='Agent task is no longer active')
|
||||
path = request.path.removeprefix('/api')
|
||||
if request.method == 'GET' and path.startswith('/history/'):
|
||||
if path.split('/')[-1] not in {c['prompt_id'] for c in self.store.children(task['id'])}:
|
||||
raise web.HTTPForbidden(reason='History does not belong to this Agent')
|
||||
return await handler(request)
|
||||
if request.method == 'GET' and path == '/view':
|
||||
wanted = {k: request.query.get(k, 'output' if k == 'type' else '')
|
||||
for k in ('filename', 'subfolder', 'type')}
|
||||
allowed = False
|
||||
for child in self.store.children(task['id']):
|
||||
entries = self.history.query(owner_id=task['owner'], prompt_id=child['prompt_id']) if self.history else {}
|
||||
for entry in entries.values():
|
||||
for output in entry.get('outputs', {}).values():
|
||||
for value in output.values():
|
||||
if isinstance(value, list):
|
||||
allowed = allowed or any(isinstance(item, dict) and all(
|
||||
item.get(k, 'output' if k == 'type' else '') == v for k, v in wanted.items()
|
||||
) for item in value)
|
||||
if not allowed:
|
||||
raise web.HTTPForbidden(reason='Asset does not belong to an Agent child')
|
||||
return await handler(request)
|
||||
if request.method == 'GET' and path == '/object_info':
|
||||
return await handler(request)
|
||||
if request.method == 'POST' and path == '/upload/image':
|
||||
return await handler(request)
|
||||
if request.method != 'POST' or path != '/prompt':
|
||||
raise web.HTTPForbidden(reason='Agent task cannot access this endpoint')
|
||||
payload = await request.json()
|
||||
workflow = payload.get('prompt', {})
|
||||
if any('Agent' in str(node.get('class_type', '')) and 'ComfyUI-LLM' in str(node.get('class_type', ''))
|
||||
for node in workflow.values()):
|
||||
raise web.HTTPBadRequest(reason='Nested Agent workflows are not supported')
|
||||
key = str(payload.get('client_id') or '')
|
||||
if not key:
|
||||
raise web.HTTPBadRequest(reason='Agent child submission requires a client_id')
|
||||
fingerprint = hashlib.sha256(json.dumps(workflow, sort_keys=True).encode()).hexdigest()
|
||||
try:
|
||||
child, fresh = self.store.reserve_child(task['id'], key, fingerprint)
|
||||
except ValueError as exc:
|
||||
raise web.HTTPConflict(reason=str(exc))
|
||||
if not fresh:
|
||||
if child['response']:
|
||||
saved = json.loads(child['response'])
|
||||
return web.json_response(saved['body'], status=saved['status'])
|
||||
existing = self.scheduler.get_job(child['prompt_id'])
|
||||
if existing:
|
||||
return web.json_response({'prompt_id': child['prompt_id']})
|
||||
raise web.HTTPConflict(reason='Submission outcome unknown; do not regenerate automatically')
|
||||
payload['prompt_id'] = child['prompt_id']
|
||||
# aiohttp Request.json reads the cached body; keep downstream prompt validation intact.
|
||||
request._read_bytes = json.dumps(payload).encode()
|
||||
response = await handler(request)
|
||||
self.store.child_response(child['prompt_id'], {'status': response.status, 'body': json.loads(response.body)})
|
||||
if self.store.get(task['id'])['state'] != 'running':
|
||||
self.scheduler.cancel(child['prompt_id'], owner_id=task['owner'])
|
||||
return response
|
||||
|
||||
def install(self, routes):
|
||||
routes.get(self.PUBLIC)(self.listing)
|
||||
routes.get(self.PUBLIC + '/{task_id}')(self.detail)
|
||||
routes.post(self.PUBLIC + '/{task_id}/cancel')(self.cancel)
|
||||
routes.post(self.PUBLIC + '/{task_id}/recover')(self.recover)
|
||||
routes.get(self.PUBLIC + '/{task_id}/files/{index}')(self.file)
|
||||
routes.post(self.INTERNAL)(self.internal)
|
||||
@@ -0,0 +1,176 @@
|
||||
"""Durable Agent tasks, independent of the ComfyUI execution queue."""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import time
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
from uuid import uuid4
|
||||
|
||||
|
||||
TERMINAL = {"completed", "failed", "cancelled", "recovery_required"}
|
||||
ACTIVE = {"running", "delivering", "cancel_requested"}
|
||||
|
||||
|
||||
class AgentTaskStore:
|
||||
def __init__(self, path, concurrency=4, per_user=2, pending_per_user=20):
|
||||
self.path = str(path)
|
||||
self.concurrency = max(1, int(concurrency))
|
||||
self.per_user = max(1, int(per_user))
|
||||
self.pending_per_user = max(1, int(pending_per_user))
|
||||
Path(path).parent.mkdir(parents=True, exist_ok=True)
|
||||
with closing(self.connect()) as db, db:
|
||||
db.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS agent_tasks (
|
||||
id TEXT PRIMARY KEY, owner TEXT NOT NULL, submission_key TEXT NOT NULL,
|
||||
job TEXT NOT NULL, state TEXT NOT NULL, created REAL NOT NULL,
|
||||
updated REAL NOT NULL, lease TEXT, lease_until REAL NOT NULL DEFAULT 0,
|
||||
checkpoint TEXT, result TEXT, error TEXT NOT NULL DEFAULT '',
|
||||
UNIQUE(owner, submission_key));
|
||||
CREATE TABLE IF NOT EXISTS agent_events (
|
||||
seq INTEGER PRIMARY KEY AUTOINCREMENT, task_id TEXT NOT NULL,
|
||||
event TEXT NOT NULL);
|
||||
CREATE INDEX IF NOT EXISTS agent_events_task ON agent_events(task_id, seq);
|
||||
CREATE TABLE IF NOT EXISTS agent_children (
|
||||
task_id TEXT NOT NULL, request_key TEXT NOT NULL, fingerprint TEXT NOT NULL,
|
||||
prompt_id TEXT NOT NULL UNIQUE, response TEXT,
|
||||
PRIMARY KEY(task_id, request_key));
|
||||
""")
|
||||
|
||||
def connect(self):
|
||||
db = sqlite3.connect(self.path, timeout=30)
|
||||
db.row_factory = sqlite3.Row
|
||||
db.execute("PRAGMA journal_mode=WAL")
|
||||
return db
|
||||
|
||||
@staticmethod
|
||||
def decode(row):
|
||||
if row is None:
|
||||
return None
|
||||
result = dict(row)
|
||||
for key in ("job", "checkpoint", "result"):
|
||||
result[key] = json.loads(result[key]) if result[key] else None
|
||||
return result
|
||||
|
||||
def create(self, owner, submission_key, job):
|
||||
now = time.time()
|
||||
with closing(self.connect()) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
old = db.execute("SELECT * FROM agent_tasks WHERE owner=? AND submission_key=?",
|
||||
(owner, submission_key)).fetchone()
|
||||
if old:
|
||||
return self.decode(old)
|
||||
count = db.execute("SELECT COUNT(*) FROM agent_tasks WHERE owner=? AND state NOT IN "
|
||||
"('completed','failed','cancelled','recovery_required')", (owner,)).fetchone()[0]
|
||||
if count >= self.pending_per_user:
|
||||
raise ValueError("独立 Agent 待处理任务已达账户上限，请等待已有任务完成。")
|
||||
task_id = uuid4().hex
|
||||
db.execute("INSERT INTO agent_tasks(id,owner,submission_key,job,state,created,updated) "
|
||||
"VALUES(?,?,?,?,'queued',?,?)", (task_id, owner, submission_key, json.dumps(job), now, now))
|
||||
return self.decode(db.execute("SELECT * FROM agent_tasks WHERE id=?", (task_id,)).fetchone())
|
||||
|
||||
def get(self, task_id):
|
||||
with closing(self.connect()) as db:
|
||||
return self.decode(db.execute("SELECT * FROM agent_tasks WHERE id=?", (task_id,)).fetchone())
|
||||
|
||||
def list(self, owner):
|
||||
with closing(self.connect()) as db:
|
||||
return [self.decode(row) for row in db.execute(
|
||||
"SELECT * FROM agent_tasks WHERE owner=? ORDER BY created DESC LIMIT 100", (owner,))]
|
||||
|
||||
def expire(self, db, now):
|
||||
# Never replay an interrupted model/tool operation: it may already have charged money.
|
||||
db.execute("UPDATE agent_tasks SET state=CASE WHEN state='delivering' THEN 'queued' "
|
||||
"ELSE 'recovery_required' END, lease=NULL, updated=?, "
|
||||
"error='执行服务失联；已保留产物和子任务，未自动重新生成。' "
|
||||
"WHERE state IN ('running','delivering','cancel_requested') AND lease_until<?", (now, now))
|
||||
|
||||
def claim(self):
|
||||
now = time.time()
|
||||
with closing(self.connect()) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
self.expire(db, now)
|
||||
active = db.execute("SELECT COUNT(*) FROM agent_tasks WHERE lease IS NOT NULL").fetchone()[0]
|
||||
if active >= self.concurrency:
|
||||
return None
|
||||
row = db.execute("SELECT * FROM agent_tasks t WHERE state='queued' AND "
|
||||
"(SELECT COUNT(*) FROM agent_tasks a WHERE a.owner=t.owner AND a.lease IS NOT NULL)<? "
|
||||
"ORDER BY created LIMIT 1", (self.per_user,)).fetchone()
|
||||
if not row:
|
||||
return None
|
||||
lease = uuid4().hex
|
||||
state = 'delivering' if row['checkpoint'] else 'running'
|
||||
db.execute("UPDATE agent_tasks SET state=?,lease=?,lease_until=?,updated=?,error='' WHERE id=?",
|
||||
(state, lease, now + 90, now, row['id']))
|
||||
return self.decode(db.execute("SELECT * FROM agent_tasks WHERE id=?", (row['id'],)).fetchone())
|
||||
|
||||
def update(self, task_id, lease, *, events=(), checkpoint=None, state=None, result=None, error=""):
|
||||
now = time.time()
|
||||
with closing(self.connect()) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
row = db.execute("SELECT * FROM agent_tasks WHERE id=? AND lease=? AND lease_until>=?",
|
||||
(task_id, lease, now)).fetchone()
|
||||
if not row:
|
||||
raise ValueError("Agent 执行租约已失效。")
|
||||
if state is not None and state not in TERMINAL | {'delivering'}:
|
||||
raise ValueError("Invalid Agent state")
|
||||
target = state or row['state']
|
||||
if row['state'] == 'cancel_requested' and target == 'completed':
|
||||
target = 'cancelled'
|
||||
for event in events:
|
||||
db.execute("INSERT INTO agent_events(task_id,event) VALUES(?,?)", (task_id, json.dumps(event)))
|
||||
db.execute("UPDATE agent_tasks SET state=?,updated=?,lease_until=?,lease=?, "
|
||||
"checkpoint=COALESCE(?,checkpoint),result=COALESCE(?,result),error=? WHERE id=?",
|
||||
(target, now, now + 90, None if target in TERMINAL else lease,
|
||||
json.dumps(checkpoint) if checkpoint is not None else None,
|
||||
json.dumps(result) if result is not None else None, error[:1000], task_id))
|
||||
return target
|
||||
|
||||
def cancel(self, task_id, owner):
|
||||
with closing(self.connect()) as db, db:
|
||||
db.execute("UPDATE agent_tasks SET state=CASE WHEN state='queued' THEN 'cancelled' "
|
||||
"ELSE 'cancel_requested' END,updated=? WHERE id=? AND owner=? "
|
||||
"AND state IN ('queued','running','delivering')", (time.time(), task_id, owner))
|
||||
|
||||
def recover(self, task_id, owner):
|
||||
with closing(self.connect()) as db, db:
|
||||
db.execute('BEGIN IMMEDIATE')
|
||||
self.expire(db, time.time())
|
||||
row = db.execute('SELECT * FROM agent_tasks WHERE id=? AND owner=?', (task_id, owner)).fetchone()
|
||||
if not row or row['state'] != 'recovery_required':
|
||||
raise ValueError('Only interrupted tasks can be recovered')
|
||||
job = json.loads(row['job'])
|
||||
if not row['checkpoint']:
|
||||
if not db.execute('SELECT 1 FROM agent_children WHERE task_id=?', (task_id,)).fetchone():
|
||||
raise ValueError('No saved video task or delivery checkpoint; inspect the interrupted Agent manually')
|
||||
job['recovery_only'] = True
|
||||
db.execute("UPDATE agent_tasks SET state='queued',lease=NULL,error='',job=?,updated=? WHERE id=?",
|
||||
(json.dumps(job), time.time(), task_id))
|
||||
|
||||
def events(self, task_id, after=0):
|
||||
with closing(self.connect()) as db:
|
||||
return [{"seq": row['seq'], **json.loads(row['event'])} for row in db.execute(
|
||||
"SELECT * FROM agent_events WHERE task_id=? AND seq>? ORDER BY seq LIMIT 200", (task_id, after))]
|
||||
|
||||
def children(self, task_id):
|
||||
with closing(self.connect()) as db:
|
||||
return [dict(row) for row in db.execute("SELECT * FROM agent_children WHERE task_id=?", (task_id,))]
|
||||
|
||||
def reserve_child(self, task_id, key, fingerprint):
|
||||
with closing(self.connect()) as db, db:
|
||||
db.execute("BEGIN IMMEDIATE")
|
||||
row = db.execute("SELECT state,lease_until FROM agent_tasks WHERE id=?", (task_id,)).fetchone()
|
||||
if not row or row['state'] != 'running' or row['lease_until'] < time.time():
|
||||
raise ValueError("Agent task cannot submit new children")
|
||||
old = db.execute("SELECT * FROM agent_children WHERE task_id=? AND request_key=?", (task_id, key)).fetchone()
|
||||
if old:
|
||||
if old['fingerprint'] != fingerprint:
|
||||
raise ValueError("Submission key reused with different workflow")
|
||||
return dict(old), False
|
||||
prompt_id = str(uuid4())
|
||||
db.execute("INSERT INTO agent_children VALUES(?,?,?,?,NULL)", (task_id, key, fingerprint, prompt_id))
|
||||
return {'prompt_id': prompt_id, 'response': None}, True
|
||||
|
||||
def child_response(self, prompt_id, response):
|
||||
with closing(self.connect()) as db, db:
|
||||
db.execute("UPDATE agent_children SET response=? WHERE prompt_id=?", (json.dumps(response), prompt_id))
|
||||
@@ -24,6 +24,7 @@ class JWTAuth:
|
||||
self.expire_minutes = expire_minutes
|
||||
self.algorithm = algorithm
|
||||
self.task_sessions = None
|
||||
self.agent_tasks = None
|
||||
|
||||
self.__secret_key = secret_key
|
||||
|
||||
@@ -101,6 +102,10 @@ class JWTAuth:
|
||||
request, "/logout", message="Unexpected error"
|
||||
)
|
||||
|
||||
if "am_agent" in user:
|
||||
if self.agent_tasks is None:
|
||||
raise web.HTTPUnauthorized(reason="Independent Agent tasks are unavailable")
|
||||
return await self.agent_tasks.delegated(request, user, handler)
|
||||
if "am_task" in user:
|
||||
if self.task_sessions is None:
|
||||
raise web.HTTPUnauthorized(reason="Task sessions are unavailable")
|
||||
|
||||
Reference in New Issue
Block a user