aizz-server 1.0.0__py3-none-any.whl

aizz_server/auth_middleware.py

@@ -0,0 +1,73 @@
+"""Authenticate all HTTP routes, and cancel in-flight work after device revocation."""
+import asyncio
+from contextlib import suppress
+
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+
+from . import device_auth
+from . import codex_proxy
+
+
+class DeviceAuthMiddleware:
+    def __init__(self, app, db_path):
+        self.app = app
+        self.db_path = db_path
+
+    async def __call__(self, scope, receive, send):
+        if scope['type'] != 'http':
+            return await self.app(scope, receive, send)
+        request = Request(scope)
+        compatible = codex_proxy.is_v1_path(request.url.path)
+        public = (request.method, request.url.path) in {
+            ('GET', '/api/health'), ('POST', '/api/auth/pair'),
+        }
+        if compatible:
+            try:
+                codex_proxy.authorize(request, scope['app'].state.codex_proxy.settings)
+            except codex_proxy.ProxyError as exc:
+                response = await codex_proxy.proxy_error_handler(request, exc)
+                return await response(scope, receive, send)
+            return await self.app(scope, receive, send)
+        token = request.headers.get('authorization', '')
+        token = token[7:] if token.startswith('Bearer ') else ''
+        device = None if public else await asyncio.to_thread(device_auth.authenticate, self.db_path(), token)
+        if not public and device is None:
+            return await JSONResponse({'detail': '设备未授权或授权已撤销'}, status_code=401,
+                                      headers={'WWW-Authenticate': 'Bearer', 'Cache-Control': 'no-store'})(scope, receive, send)
+        scope.setdefault('state', {})['device'] = device
+        started = False
+        finished = False
+
+        async def secure_send(message):
+            nonlocal started, finished
+            if message['type'] == 'http.response.start':
+                started = True
+                message['headers'] = list(message.get('headers', [])) + [(b'cache-control', b'no-store')]
+            if message['type'] == 'http.response.body' and not message.get('more_body', False):
+                finished = True
+            await send(message)
+
+        if public:
+            return await self.app(scope, receive, secure_send)
+        task = asyncio.create_task(self.app(scope, receive, secure_send))
+        try:
+            while not task.done():
+                done, _ = await asyncio.wait({task}, timeout=1)
+                if done:
+                    break
+                if not await asyncio.to_thread(device_auth.authenticate, self.db_path(), token, touch=False):
+                    task.cancel()
+                    with suppress(asyncio.CancelledError):
+                        await task
+                    if not started:
+                        await JSONResponse({'detail': '设备授权已撤销'}, status_code=401)(scope, receive, secure_send)
+                    elif not finished:
+                        await secure_send({'type': 'http.response.body', 'body': b'', 'more_body': False})
+                    return
+            await task
+        finally:
+            if not task.done():
+                task.cancel()
+                with suppress(asyncio.CancelledError):
+                    await task

aizz_server/background_tasks.py

@@ -0,0 +1,248 @@
+"""Durable, reconnectable subscriptions to independently owned generation tasks."""
+import asyncio
+import hashlib
+import logging
+from datetime import datetime, timezone
+import json
+import time
+import uuid
+from contextlib import suppress
+from contextvars import ContextVar
+
+from fastapi import HTTPException
+
+ACCEPTANCE = ContextVar('task_acceptance', default=None)
+
+TERMINAL = {'completed', 'failed', 'stopped', 'interrupted'}
+
+
+class TaskManager:
+    def __init__(self, db, stop, notify):
+        self.db, self.stop, self.notify = db, stop, notify
+        self.workers = {}
+        self.lock = asyncio.Lock()
+        self.controls = {}
+
+    def initialize(self):
+        self.lock = asyncio.Lock()
+        with self.db() as c:
+            c.executescript('''
+                CREATE TABLE IF NOT EXISTS generation_tasks (
+                    id TEXT PRIMARY KEY, session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
+                    device_id TEXT NOT NULL, request_id TEXT NOT NULL, request_hash TEXT NOT NULL,
+                    state TEXT NOT NULL, snapshot TEXT NOT NULL, created_at REAL NOT NULL,
+                    UNIQUE(device_id, request_id)
+                );
+                CREATE TABLE IF NOT EXISTS generation_events (
+                    task_id TEXT NOT NULL REFERENCES generation_tasks(id) ON DELETE CASCADE,
+                    seq INTEGER NOT NULL, event TEXT NOT NULL, created_at REAL NOT NULL,
+                    PRIMARY KEY(task_id,seq)
+                );
+                CREATE INDEX IF NOT EXISTS generation_session ON generation_tasks(session_id,created_at);
+            ''')
+            # A command may already have executed before a crash. Never replay it.
+            rows = c.execute("SELECT * FROM generation_tasks WHERE state NOT IN ('completed','failed','stopped','interrupted')").fetchall()
+            for row in rows:
+                snap = json.loads(row['snapshot'])
+                snap.update(state='interrupted', pending={}, error='服务端已重启，任务已中断，请确认已有操作后重新发送')
+                snap['content'] += '\n\n（服务端重启，任务已中断）'
+                saved = c.execute("SELECT id,content FROM messages WHERE session_id=? AND role='assistant' AND rowid>(SELECT rowid FROM messages WHERE id=?) ORDER BY rowid LIMIT 1", (row['session_id'], snap.get('user_message_id'))).fetchone()
+                message = saved['id'] if saved else str(uuid.uuid4())
+                snap['message_id'] = message
+                if saved:
+                    # The message transaction may have committed just before the process died.
+                    snap['content'] = saved['content'] + '\n\n（服务端重启，任务已中断）'
+                    c.execute('UPDATE messages SET content=?,mode=? WHERE id=?', (snap['content'], 'error', message))
+                else:
+                    c.execute('INSERT INTO messages VALUES (?,?,?,?,?,?)', (message, row['session_id'], 'assistant', snap['content'], 'error', datetime.now(timezone.utc).isoformat()))
+                c.execute('INSERT OR REPLACE INTO message_runs VALUES (?,?)', (message, json.dumps(snap['execution'], ensure_ascii=False)))
+                snap['seq'] += 1
+                c.execute('UPDATE generation_tasks SET state=?,snapshot=? WHERE id=?', ('interrupted', json.dumps(snap, ensure_ascii=False), row['id']))
+                c.execute('INSERT INTO generation_events VALUES (?,?,?,?)', (row['id'], snap['seq'], json.dumps({'type': 'snapshot', 'data': snap, 'seq': snap['seq']}, ensure_ascii=False), time.time()))
+            c.execute('DELETE FROM generation_events WHERE created_at < ?', (time.time() - 7 * 86400,))
+
+    def get(self, task_id):
+        with self.db() as c:
+            row = c.execute('SELECT * FROM generation_tasks WHERE id=?', (task_id,)).fetchone()
+        if row is None:
+            raise HTTPException(404, '任务不存在')
+        return json.loads(row['snapshot'])
+
+    def session(self, session_id):
+        with self.db() as c:
+            row = c.execute('SELECT id FROM generation_tasks WHERE session_id=? ORDER BY created_at DESC,rowid DESC LIMIT 1', (session_id,)).fetchone()
+        return self.get(row['id']) if row else None
+
+    def dismiss_plan(self, task_id):
+        with self.db() as c:
+            row = c.execute('SELECT snapshot FROM generation_tasks WHERE id=?', (task_id,)).fetchone()
+            if row is None:
+                raise HTTPException(404, '任务不存在')
+            snap = json.loads(row['snapshot'])
+            if snap['state'] != 'completed' or not snap.get('plan'):
+                raise HTTPException(409, '没有待确认的计划')
+            snap['plan_dismissed'] = True
+            c.execute('UPDATE generation_tasks SET snapshot=? WHERE id=?', (json.dumps(snap, ensure_ascii=False), task_id))
+        return snap
+
+    async def create(self, session_id, device_id, request_id, body, prepare):
+        if not 16 <= len(request_id) <= 128:
+            raise HTTPException(400, '请求 ID 长度必须为 16–128 字符')
+        fingerprint = hashlib.sha256(json.dumps({'session': session_id, 'body': body}, sort_keys=True).encode()).hexdigest()
+        async with self.lock:
+            with self.db() as c:
+                existing = c.execute('SELECT id,request_hash FROM generation_tasks WHERE device_id=? AND request_id=?', (device_id, request_id)).fetchone()
+            if existing:
+                if existing['request_hash'] != fingerprint:
+                    raise HTTPException(409, '请求 ID 已用于不同消息')
+                return self.get(existing['id'])
+            current = self.session(session_id)
+            if current and current['state'] not in TERMINAL:
+                raise HTTPException(409, '当前会话已有回复正在生成')
+            task_id = str(uuid.uuid4())
+            snap = dict(user_message_id=None, id=task_id, session_id=session_id, request_id=request_id, state='running', seq=0,
+                        content='', mode=body['mode'], plan=body.get('plan', False), plan_dismissed=False, execution={'duration_ms': 0, 'items': []},
+                        pending={}, message_id=None, error=None)
+            # The existing message transaction also writes the idempotency record.
+            # Either both are durable, or neither is accepted.
+            token = ACCEPTANCE.set((snap, device_id, fingerprint))
+            try:
+                response = await prepare()
+            finally:
+                ACCEPTANCE.reset(token)
+            self.workers[task_id] = asyncio.create_task(self.run(snap, device_id, response))
+            return snap.copy()
+
+    def accept_message(self, connection, role, message_id):
+        acceptance = ACCEPTANCE.get()
+        if role != 'user' or acceptance is None:
+            return
+        snap, device_id, fingerprint = acceptance
+        snap['user_message_id'] = message_id
+        snap['mode'] = connection.execute('SELECT mode FROM messages WHERE id=?', (message_id,)).fetchone()['mode']
+        connection.execute('INSERT INTO generation_tasks VALUES (?,?,?,?,?,?,?,?)',
+                           (snap['id'], snap['session_id'], device_id, snap['request_id'], fingerprint,
+                            snap['state'], json.dumps(snap, ensure_ascii=False), time.time()))
+
+    async def publish(self, snap, kind, data):
+        if kind == 'heartbeat':
+            return
+        if kind == 'delta':
+            snap['content'] += data
+        elif kind == 'execution':
+            items = snap['execution']['items']
+            items[:] = [item for item in items if item['id'] != data['id']] + [data]
+        elif kind == 'activity':
+            snap['execution']['items'].append({'id': str(uuid.uuid4()), 'type': 'notice', 'text': str(data)})
+        elif kind in ('approval', 'user_input'):
+            snap['pending'][data['token']] = {'type': kind, 'data': data}
+        elif kind == 'approval_resolved':
+            snap['pending'].pop(data, None)
+        elif kind == 'error':
+            snap.update(state='failed', error=str(data))
+        elif kind == 'interrupted':
+            snap.update(state='interrupted', error=str(data))
+        elif kind == 'done':
+            snap['state'] = 'stopped' if data.get('stopped') else 'completed'
+        if kind not in ('done', 'error', 'interrupted'):
+            snap['state'] = ('waiting_approval' if any(p['type'] == 'approval' for p in snap['pending'].values()) else
+                             'waiting_input' if snap['pending'] else 'running')
+        if snap['state'] in TERMINAL:
+            snap['pending'].clear()
+            with self.db() as c:
+                row = c.execute("SELECT id FROM messages WHERE session_id=? AND role='assistant' ORDER BY rowid DESC LIMIT 1", (snap['session_id'],)).fetchone()
+            snap['message_id'] = row['id'] if row else None
+            if kind == 'interrupted' and row:
+                with self.db() as c:
+                    c.execute('UPDATE messages SET content=?,mode=? WHERE id=?', (snap['content'] + '\n\n（服务端关闭，任务已中断）', 'error', row['id']))
+        snap['seq'] += 1
+        payload = {'type': 'snapshot' if kind == 'interrupted' else kind, 'data': snap if kind == 'interrupted' else data, 'seq': snap['seq']}
+        with self.db() as c:
+            if snap['state'] in TERMINAL:
+                c.execute('DELETE FROM generation_events WHERE created_at < ?', (time.time() - 7 * 86400,))
+            c.execute('UPDATE generation_tasks SET state=?,snapshot=? WHERE id=?', (snap['state'], json.dumps(snap, ensure_ascii=False), snap['id']))
+            c.execute('INSERT INTO generation_events VALUES (?,?,?,?)', (snap['id'], snap['seq'], json.dumps(payload, ensure_ascii=False), time.time()))
+        if kind in ('approval', 'user_input', 'error') or (kind == 'done' and not data.get('stopped')):
+            try:
+                await self.notify(snap, kind, data)
+            except Exception:
+                logging.getLogger('aizz.agent').exception('任务通知入队失败 task=%s', snap['id'])
+
+    async def resolved(self, session_id, token):
+        current = self.session(session_id)
+        if current and token in current['pending']:
+            # The runner owns its snapshot; mutate that object via a control event.
+            self.controls.setdefault(current['id'], []).append(token)
+
+    async def run(self, snap, device_id, response):
+        started = time.monotonic()
+        iterator = response.body_iterator
+        pending = None
+        try:
+            while True:
+                if pending is None:
+                    pending = asyncio.create_task(anext(iterator))
+                done, _ = await asyncio.wait({pending}, timeout=0.25)
+                for token in self.controls.pop(snap['id'], []):
+                    if token in snap['pending']:
+                        await self.publish(snap, 'approval_resolved', token)
+                with self.db() as c:
+                    active = c.execute('SELECT 1 FROM authorized_devices WHERE id=? AND revoked_at IS NULL', (device_id,)).fetchone()
+                if not active:
+                    await self.stop(snap['session_id'])
+                if not done:
+                    continue
+                try:
+                    raw = pending.result()
+                except StopAsyncIteration:
+                    break
+                pending = None
+                if isinstance(raw, bytes):
+                    raw = raw.decode()
+                item = json.loads(raw.removeprefix('data: ').strip())
+                snap['execution']['duration_ms'] = int((time.monotonic() - started) * 1000)
+                await self.publish(snap, item['type'], item['data'])
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            await self.stop(snap['session_id'])
+            await self.publish(snap, 'error', str(exc))
+        finally:
+            if pending:
+                pending.cancel()
+                await asyncio.gather(pending, return_exceptions=True)
+            with suppress(Exception):
+                await iterator.aclose()
+            self.workers.pop(snap['id'], None)
+
+    async def events(self, task_id, after):
+        while True:
+            snap = self.get(task_id)
+            with self.db() as c:
+                rows = c.execute('SELECT seq,event FROM generation_events WHERE task_id=? AND seq>? ORDER BY seq LIMIT 100', (task_id, after)).fetchall()
+            if rows and rows[0]['seq'] == after + 1:
+                for row in rows:
+                    after = row['seq']
+                    yield 'data: ' + row['event'] + '\n\n'
+                continue
+            if after != snap['seq'] or after == 0:
+                after = snap['seq']
+                yield 'data: ' + json.dumps({'type': 'snapshot', 'data': snap, 'seq': after}, ensure_ascii=False) + '\n\n'
+            if snap['state'] in TERMINAL:
+                return
+            yield ': heartbeat\n\n'
+            await asyncio.sleep(0.5)
+
+    async def shutdown(self):
+        running = list(self.workers)
+        for task_id in running:
+            await self.stop(self.get(task_id)['session_id'])
+        if self.workers:
+            await asyncio.gather(*list(self.workers.values()), return_exceptions=True)
+        for task_id in running:
+            try:
+                snap = self.get(task_id)
+                if snap['state'] == 'stopped':
+                    await self.publish(snap, 'interrupted', '服务端关闭，任务已中断，请确认已有操作后重新发送')
+            except HTTPException:
+                pass