@agentunion/fastaun 0.5.27 → 0.5.28

package/dist/transport.js

@@ -14,6 +14,8 @@ import { AUNError, AuthError, CertificateRevokedError, ClientSignatureError, Con
 import { trafficLogContext, withLogContext } from './logger.js';
 import { isJsonObject } from './types.js';
 import { parseJsonStrict, validateOrgRequestJSON } from './org-protocol.js';
+import { RiskController } from './risk-control.js';
+import { messageClassOf } from './client/control-messages.js';
 const MAX_WS_PAYLOAD_SIZE = 1_000_000;
 const MAX_RPC_INFLIGHT = 16;
 const MAX_BACKGROUND_RPC_INFLIGHT = 8;
@@ -572,11 +574,17 @@ export class RPCTransport {
     _drainingRpcQueue = false;
     _onForegroundDrained;
     _onRpcResponse;
-    _sendChain = Promise.resolve();
+    _sendQueue = [];
+    _sending = false;
+    _riskDeliveries = 0;
     _idCounter = 0;
     // Gateway 在 RPC envelope 注入 _meta 字段（与 result 同级），由 client 层 observer 接收。
     // observer 抛异常会被吞掉，不影响 RPC result 返回。
     _metaObserver = null;
+    _requestAdmission = null;
+    _requestCheck = null;
+    _riskControlHandler = null;
+    _classMessageHandler = null;
     _protocolMetaObserver = false;
     // Trace 模式：off / log / diag
     _traceMode = 'off';
@@ -829,6 +837,16 @@ export class RPCTransport {
         }
         if (this._workerGeneration === null || message.generation !== this._workerGeneration)
             return;
+        if (message.event === '_transport.risk_control') {
+            this._dispatchRiskControl(message.payload, () => message.generation === this._workerGeneration);
+            return;
+        }
+        if (message.event === '_transport.class_message') {
+            const value = message.payload;
+            if (value && typeof value.event === 'string')
+                this._routeClassMessage(value.event, value.payload);
+            return;
+        }
         if (message.event === '_transport.observer') {
             const item = message.payload;
             const observer = item?.kind?.includes('trace') ? this._traceObserver : this._metaObserver;
@@ -997,6 +1015,66 @@ export class RPCTransport {
         this._metaObserver = observer;
         this._protocolMetaObserver = false;
     }
+    setRequestAdmission(admission) { this._requestAdmission = admission; }
+    setRequestCheck(check) { this._requestCheck = check; }
+    setClassMessageHandler(handler) { this._classMessageHandler = handler; }
+    _routeClassMessage(event, payload) {
+        if (!['message.received', 'group.message_created'].includes(event) || !isJsonObject(payload))
+            return false;
+        const value = payload;
+        const messageClass = messageClassOf(value);
+        if (value.delivery_class === 'normal' && messageClass !== 'sdk_control')
+            return false;
+        if (messageClass === 'normal' && !(event === 'group.message_created' && (value.storage === 'volatile' || value.volatile === true)))
+            return false;
+        if (!isMainThread && Boolean(workerData?.aunTransportWorker) && parentPort) {
+            parentPort.postMessage({ type: 'event', generation: this._connectionGeneration, event: '_transport.class_message', payload: { event, payload: value } });
+        }
+        else {
+            try {
+                this._classMessageHandler?.(event, value);
+            }
+            catch { /* 控制正文保持内部隔离 */ }
+        }
+        return true;
+    }
+    async setRiskSnapshot(state) {
+        if (this._worker)
+            await this._workerCommand({ type: 'risk_snapshot', riskState: state });
+    }
+    setRiskControlHandler(handler) {
+        this._riskControlHandler = handler;
+    }
+    _finishRiskDelivery() {
+        this._riskDeliveries = Math.max(0, this._riskDeliveries - 1);
+        if (this._riskDeliveries === 0)
+            this._startSendQueue();
+    }
+    _receiveRiskControl(payload) {
+        if (!isJsonObject(payload))
+            return;
+        if (!isMainThread && Boolean(workerData?.aunTransportWorker) && parentPort) {
+            this._riskDeliveries += 1;
+            parentPort.postMessage({ type: 'event', generation: this._connectionGeneration, event: '_transport.risk_control', payload });
+        }
+        else {
+            const generation = this._connectionGeneration;
+            this._dispatchRiskControl(payload, () => generation === this._connectionGeneration);
+        }
+    }
+    _dispatchRiskControl(payload, current) {
+        const handler = this._riskControlHandler;
+        this._riskDeliveries += 1;
+        const value = structuredClone(payload);
+        void Promise.resolve().then(() => {
+            if (current() && handler && this._riskControlHandler === handler && isJsonObject(value))
+                return handler(value);
+        }).catch(() => { }).finally(() => {
+            this._finishRiskDelivery();
+            if (current() && this._worker)
+                this._worker.postMessage({ type: 'risk_release', generation: this._workerGeneration ?? undefined });
+        });
+    }
     setProtocolMetaObserver(observer) {
         this._metaObserver = observer;
         this._protocolMetaObserver = true;
@@ -1453,6 +1531,8 @@ export class RPCTransport {
         this._handshakeAbort = null;
         abortHandshake?.();
         this._closed = true;
+        this._riskDeliveries = 0;
+        this._startSendQueue();
         this._connectionGeneration += 1;
         this._stopWebSocketHeartbeat();
         // 取消所有待处理的 RPC
@@ -1548,6 +1628,20 @@ export class RPCTransport {
         return Object.assign(result, { cancel: request.cancel });
     }
     _call(method, params, timeout, trace, background = false, onSent, sentCheckpoint) {
+        if (!this._requestAdmission)
+            return this._callAdmitted(method, params, timeout, trace, background, onSent, sentCheckpoint);
+        let release;
+        try {
+            release = this._requestAdmission?.(`rpc:${method}`);
+            const request = this._callAdmitted(method, params, timeout, trace, background, onSent, sentCheckpoint);
+            return Object.assign(request.finally(() => release?.()), { cancel: request.cancel });
+        }
+        catch (error) {
+            release?.();
+            return Promise.reject(error);
+        }
+    }
+    _callAdmitted(method, params, timeout, trace, background = false, onSent, sentCheckpoint) {
         if (this._worker) {
             const requestContext = trafficLogContext('outbound', params ?? {});
             const requestLog = withLogContext(this._logger, requestContext);
@@ -1680,6 +1774,7 @@ export class RPCTransport {
                 tStart,
                 timeoutMs: effectiveTimeout,
                 background: backgroundRpc,
+                priority: this._messagePriority(localParams.message_class),
                 logger: requestLog,
                 onSent,
             });
@@ -1693,6 +1788,7 @@ export class RPCTransport {
                     tStart,
                     timeoutMs: effectiveTimeout,
                     background: true,
+                    priority: this._messagePriority(localParams.message_class),
                     logger: requestLog,
                     onSent,
                 });
@@ -1703,6 +1799,15 @@ export class RPCTransport {
     }
     /** 发送 JSON-RPC 2.0 Notification，不分配 id，也不等待响应。 */
     async notify(method, params) {
+        const release = this._requestAdmission?.(`rpc:${String(method).replace(/^(notification|event)\//, '')}`);
+        try {
+            await this._notifyAdmitted(method, params);
+        }
+        finally {
+            release?.();
+        }
+    }
+    async _notifyAdmitted(method, params) {
         if (this._worker) {
             await this._workerCommand({ type: 'notify', method, params });
             return;
@@ -1733,15 +1838,43 @@ export class RPCTransport {
             throw new ValidationError('payload is too large');
         }
         const notifyLog = withLogContext(this._logger, trafficLogContext('outbound', params ?? {}));
-        await this._sendText(payload, `notification ${normalizedMethod}`);
+        await this._sendText(payload, `notification ${normalizedMethod}`, () => { this._requestCheck?.(`rpc:${normalizedMethod.replace(/^(notification|event)\//, '')}`); return true; }, this._messagePriority(params?.message_class));
         notifyLog.debug(`notification sent: method=${normalizedMethod}, size=${payloadSize}`);
     }
-    _enqueueSend(task) {
-        const run = this._sendChain.then(task, task);
-        this._sendChain = run.catch(() => { });
-        return run;
+    _messagePriority(messageClass) { return messageClass === 'sdk_control' ? 2 : messageClass === 'control' ? 1 : 0; }
+    _isRiskControlMethod(method) { return ['gateway.risk_ack', 'meta.ping', 'system.ping', 'ping', 'auth.disconnect', 'gateway.risk_status'].includes(method); }
+    _enqueueSend(task, priority = 0) {
+        const result = new Promise((resolve, reject) => { this._sendQueue.push({ task, priority, resolve, reject }); });
+        this._startSendQueue();
+        return result;
+    }
+    _startSendQueue() {
+        if (this._sending)
+            return;
+        this._sending = true;
+        void Promise.resolve().then(async () => {
+            try {
+                while (this._sendQueue.length) {
+                    this._sendQueue.sort((left, right) => right.priority - left.priority);
+                    const index = this._riskDeliveries > 0 ? this._sendQueue.findIndex(entry => entry.priority === 3) : 0;
+                    if (index < 0)
+                        break;
+                    const entry = this._sendQueue.splice(index, 1)[0];
+                    try {
+                        await entry.task();
+                        entry.resolve();
+                    }
+                    catch (error) {
+                        entry.reject(error);
+                    }
+                }
+            }
+            finally {
+                this._sending = false;
+            }
+        });
     }
-    _sendText(payload, context, beforeSend) {
+    _sendText(payload, context, beforeSend, priority = 0) {
         return this._enqueueSend(async () => {
             if (beforeSend && !beforeSend()) {
                 throw new ConnectionError(`send cancelled before write: ${context}`);
@@ -1767,7 +1900,7 @@ export class RPCTransport {
                     this._sendFailureError(context, err, ws).then(reject, reject);
                 }
             });
-        });
+        }, priority);
     }
     async _sendFailureError(context, err, ws) {
         const original = err instanceof Error ? err.message : String(err);
@@ -1839,19 +1972,29 @@ export class RPCTransport {
         this._drainingRpcQueue = true;
         let foregroundDrained = false;
         try {
-            while (!this._closed && this._ws && this._pending.size < MAX_RPC_INFLIGHT) {
+            while (!this._closed && this._ws) {
+                const controlOnly = this._pending.size >= MAX_RPC_INFLIGHT;
+                if (controlOnly && (!this._rpcQueue.some(entry => this._isRiskControlMethod(entry.method)) || Array.from(this._pendingMeta.values()).filter(entry => this._isRiskControlMethod(entry.method)).length >= 4))
+                    break;
                 let entry;
-                if (this._rpcQueue.length > 0) {
-                    entry = this._rpcQueue.shift();
-                    if (this._rpcQueue.length === 0)
-                        foregroundDrained = true;
-                }
-                else if (this._backgroundRpcQueue.length > 0 && this._pendingBackground.size < MAX_BACKGROUND_RPC_INFLIGHT) {
-                    entry = this._backgroundRpcQueue.shift();
+                let selectedQueue;
+                let selectedIndex = -1;
+                for (const queue of [this._rpcQueue, ...(this._pendingBackground.size < MAX_BACKGROUND_RPC_INFLIGHT ? [this._backgroundRpcQueue] : [])]) {
+                    for (let index = 0; index < queue.length; index++) {
+                        if (controlOnly && !this._isRiskControlMethod(queue[index].method))
+                            continue;
+                        if (!entry || queue[index].priority > entry.priority) {
+                            entry = queue[index];
+                            selectedQueue = queue;
+                            selectedIndex = index;
+                        }
+                    }
                 }
-                else {
+                if (!entry || !selectedQueue)
                     break;
-                }
+                selectedQueue.splice(selectedIndex, 1);
+                if (selectedQueue === this._rpcQueue && this._rpcQueue.length === 0)
+                    foregroundDrained = true;
                 const elapsed = Date.now() - entry.tStart;
                 if (elapsed >= entry.timeoutMs) {
                     clearTimeout(entry.pending.timer);
@@ -1870,7 +2013,12 @@ export class RPCTransport {
                 if (entry.background) {
                     this._pendingBackground.add(entry.rpcId);
                 }
-                this._sendText(entry.payload, `rpc ${entry.method}`, () => this._pending.get(entry.rpcId) === entry.pending).then(() => {
+                this._sendText(entry.payload, `rpc ${entry.method}`, () => {
+                    if (this._pending.get(entry.rpcId) !== entry.pending)
+                        return false;
+                    this._requestCheck?.(`rpc:${entry.method}`);
+                    return true;
+                }, this._isRiskControlMethod(entry.method) ? 3 : entry.priority).then(() => {
                     entry.logger.debug(`RPC request sent: method=${entry.method}, id=${entry.rpcId}, background=${entry.background}`);
                     return entry.onSent?.();
                 }).catch((err) => {
@@ -1878,7 +2026,7 @@ export class RPCTransport {
                         return;
                     this._removeRpc(entry.rpcId, entry.pending);
                     entry.logger.error(`RPC send failed: method=${entry.method}, id=${entry.rpcId}, error=${err instanceof Error ? err.message : String(err)}`);
-                    entry.pending.reject(err instanceof ConnectionError
+                    entry.pending.reject(err instanceof ConnectionError || err instanceof StateError
                         ? err
                         : new ConnectionError(`failed to send rpc ${entry.method}: ${err instanceof Error ? err.message : String(err)}`));
                     this._drainRpcQueue();
@@ -2073,6 +2221,14 @@ export class RPCTransport {
         this._logger.debug(`_routeMessage enter: has_id=${hasId}, method=${methodStr || '<none>'}`);
         // RPC 响应（有 id 字段）
         if ('id' in message) {
+            const error = message.error;
+            if (isJsonObject(error) && isJsonObject(error.data) && isJsonObject(error.data.risk_state)) {
+                const states = error.data.risk_state.states;
+                if (Array.isArray(states))
+                    for (const state of states)
+                        this._receiveRiskControl(state);
+                delete error.data.risk_state;
+            }
             const rpcId = String(message.id);
             const pending = this._pending.get(rpcId);
             if (pending) {
@@ -2086,6 +2242,10 @@ export class RPCTransport {
             return;
         }
         const method = String(message.method ?? '');
+        if (method === 'event/gateway.risk_control' || method === 'gateway.risk_control') {
+            this._receiveRiskControl(message.params);
+            return;
+        }
         // Challenge 消息
         if (method === 'challenge') {
             this._challenge = message;
@@ -2127,6 +2287,8 @@ export class RPCTransport {
                     eventLog.info(`[trace=${String(traceObj.trace_id ?? '')}] event_recv event=${sdkEvent}`);
                 }
             }
+            if (this._routeClassMessage(sdkEvent, params))
+                return;
             // 发布为 _raw.{event}，由 AUNClient 处理后再发布用户可见的事件
             if (sdkEvent.startsWith('app.') || sdkEvent.startsWith('org.member.') || sdkEvent.startsWith('org.lifecycle.')) {
                 this._dispatcher.enqueue(sdkEvent, params);
@@ -2238,6 +2400,8 @@ if (!isMainThread && Boolean(workerData?.aunTransportWorker) && parentPort) {
         onDisconnect: () => { },
     });
     workerActor = actor;
+    const riskGate = new RiskController();
+    actor.setRequestCheck((request) => riskGate.check(request));
     actor._onForegroundDrained = () => parentPort.postMessage({
         type: 'foreground_drained', generation: bridgeGeneration(),
     });
@@ -2260,6 +2424,11 @@ if (!isMainThread && Boolean(workerData?.aunTransportWorker) && parentPort) {
     parentPort.on('message', async (command) => {
         const id = command.id;
         try {
+            if (command.type === 'risk_release') {
+                if (command.generation === bridgeGeneration())
+                    actor['_finishRiskDelivery']();
+                return;
+            }
             if (command.type === 'cancel') {
                 const pending = active.get(String(command.targetId ?? ''));
                 pending?.cancel?.();
@@ -2267,6 +2436,16 @@ if (!isMainThread && Boolean(workerData?.aunTransportWorker) && parentPort) {
             }
             let result;
             switch (command.type) {
+                case 'risk_snapshot': {
+                    const state = command.riskState;
+                    const target = state.target;
+                    const [keyId, key] = Object.entries(state.signing_keys)[0] ?? [];
+                    riskGate.configure({ aid: String(target.aid), device_id: String(target.device_id ?? ''), slot_id: String(target.slot_id ?? ''), connection_id: String(target.connection_id ?? '') }, { key_id: keyId, public_key: Buffer.from(key).toString('base64url'), algorithm: 'ECDSA-P256-SHA256' }, Number(state.server_time_ms));
+                    await riskGate.restore(state);
+                    riskGate.close();
+                    result = null;
+                    break;
+                }
                 case 'connect': {
                     dnsFallback = command.dnsFallback ?? null;
                     const request = actor.connect(String(command.url ?? ''));