diff --git a/.gitignore b/.gitignore index 1d6c013..195dde0 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,7 @@ blacklist.txt secret_key.txt history.sqlite3* scheduler.sqlite3* +agent_tasks.sqlite3* scheduler_secret.txt *.lock cloud_staging/ diff --git a/account_manager.py b/account_manager.py index 1629175..ea5ad4d 100644 --- a/account_manager.py +++ b/account_manager.py @@ -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 diff --git a/tests/test_agent_tasks.py b/tests/test_agent_tasks.py new file mode 100644 index 0000000..cb40e58 --- /dev/null +++ b/tests/test_agent_tasks.py @@ -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() diff --git a/utils/agent_task_routes.py b/utils/agent_task_routes.py new file mode 100644 index 0000000..feee6b6 --- /dev/null +++ b/utils/agent_task_routes.py @@ -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) diff --git a/utils/agent_tasks.py b/utils/agent_tasks.py new file mode 100644 index 0000000..54dbd23 --- /dev/null +++ b/utils/agent_tasks.py @@ -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= 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)=?", + (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)) diff --git a/utils/jwt_auth.py b/utils/jwt_auth.py index 73157cb..d96c284 100644 --- a/utils/jwt_auth.py +++ b/utils/jwt_auth.py @@ -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")