@agentunion/fastaun 0.5.26 → 0.5.28

package/dist/transport.d.ts

@@ -56,9 +56,15 @@ export declare class RPCTransport {
     private _drainingRpcQueue;
     _onForegroundDrained?: () => void;
     _onRpcResponse?: () => void;
-    private _sendChain;
+    private _sendQueue;
+    private _sending;
+    private _riskDeliveries;
     private _idCounter;
     private _metaObserver;
+    private _requestAdmission;
+    private _requestCheck;
+    private _riskControlHandler;
+    private _classMessageHandler;
     private _protocolMetaObserver;
     private _traceMode;
     private _traceObserver;
@@ -106,8 +112,18 @@ export declare class RPCTransport {
      * 不影响 RPC 结果返回。
      */
     setMetaObserver(observer: MetaObserver | null): void;
+    setRequestAdmission(admission: ((request: string) => () => void) | null): void;
+    setRequestCheck(check: (request: string) => void): void;
+    setClassMessageHandler(handler: (event: string, payload: JsonObject) => boolean): void;
+    private _routeClassMessage;
+    setRiskSnapshot(state: JsonObject): Promise<void>;
+    setRiskControlHandler(handler: ((payload: JsonObject) => void | Promise<void>) | null): void;
+    private _finishRiskDelivery;
+    private _receiveRiskControl;
+    private _dispatchRiskControl;
     setProtocolMetaObserver(observer: MetaObserver | null): void;
     setProtocolDisconnectCallback(callback: DisconnectCallback | null): void;
+    waitForPending(timeoutMs: number, signal?: AbortSignal): Promise<void>;
     getTraceMode(): string;
     /** 设置 trace 模式：off / log / diag */
     setTraceMode(mode: string): void;
@@ -145,11 +161,16 @@ export declare class RPCTransport {
     call(method: string, params?: RpcParams, timeout?: number, trace?: string, background?: boolean): Promise<RpcResult>;
     private _callWithSent;
     private _call;
+    private _callAdmitted;
     private _callLocal;
     /** 发送 JSON-RPC 2.0 Notification，不分配 id，也不等待响应。 */
     notify(method: string, params?: RpcParams | null): Promise<void>;
+    private _notifyAdmitted;
     private _notifyLocal;
+    private _messagePriority;
+    private _isRiskControlMethod;
     private _enqueueSend;
+    private _startSendQueue;
     private _sendText;
     private _sendFailureError;
     private _waitForCloseCodeIfClosing;

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;
@@ -858,7 +876,8 @@ export class RPCTransport {
                 this._dispatcher.enqueue(message.event, message.payload);
             }
             else {
-                this._dispatcher.enqueueProtocol(message.event, message.payload);
+                const migration = message.event === 'notification' && message.payload?.method === 'gateway.migrate';
+                this._dispatcher.enqueueProtocol(message.event, message.payload, migration ? () => message.generation === this._workerGeneration : undefined);
             }
         }
     }
@@ -996,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;
@@ -1006,6 +1085,21 @@ export class RPCTransport {
         this._onDisconnect = callback;
         this._protocolDisconnectCallback = callback !== null;
     }
+    async waitForPending(timeoutMs, signal) {
+        const deadline = Date.now() + Math.max(0, timeoutMs);
+        while (!signal?.aborted && Date.now() < deadline) {
+            const active = this._worker
+                ? [...this._workerPending.values()].some((entry) => entry.commandType === 'call' || entry.commandType === 'notify')
+                : this._pending.size > 0 || this._rpcQueue.length > 0 || this._backgroundRpcQueue.length > 0;
+            if (!active)
+                return;
+            await new Promise((resolve) => {
+                const finish = () => { clearTimeout(timer); signal?.removeEventListener('abort', finish); resolve(); };
+                const timer = setTimeout(finish, Math.min(10, deadline - Date.now()));
+                signal?.addEventListener('abort', finish, { once: true });
+            });
+        }
+    }
     getTraceMode() {
         return this._traceMode;
     }
@@ -1437,6 +1531,8 @@ export class RPCTransport {
         this._handshakeAbort = null;
         abortHandshake?.();
         this._closed = true;
+        this._riskDeliveries = 0;
+        this._startSendQueue();
         this._connectionGeneration += 1;
         this._stopWebSocketHeartbeat();
         // 取消所有待处理的 RPC
@@ -1532,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);
@@ -1664,6 +1774,7 @@ export class RPCTransport {
                 tStart,
                 timeoutMs: effectiveTimeout,
                 background: backgroundRpc,
+                priority: this._messagePriority(localParams.message_class),
                 logger: requestLog,
                 onSent,
             });
@@ -1677,6 +1788,7 @@ export class RPCTransport {
                     tStart,
                     timeoutMs: effectiveTimeout,
                     background: true,
+                    priority: this._messagePriority(localParams.message_class),
                     logger: requestLog,
                     onSent,
                 });
@@ -1687,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;
@@ -1717,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}`);
@@ -1751,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);
@@ -1823,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);
@@ -1854,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) => {
@@ -1862,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();
@@ -1964,7 +2128,9 @@ export class RPCTransport {
     _startWebSocketHeartbeat(ws) {
         this._stopWebSocketHeartbeat();
         this._heartbeatWs = ws;
+        let failures = 0;
         const onPong = () => {
+            failures = 0;
             if (this._pongTimeout !== null) {
                 clearTimeout(this._pongTimeout);
                 this._pongTimeout = null;
@@ -1972,6 +2138,7 @@ export class RPCTransport {
         };
         this._pongListener = onPong;
         ws.on('pong', onPong);
+        ws.on('message', onPong);
         const ping = () => {
             if (this._closed || this._ws !== ws || ws.readyState !== WebSocket.OPEN)
                 return;
@@ -1981,6 +2148,8 @@ export class RPCTransport {
                 this._pongTimeout = null;
                 if (this._closed || this._ws !== ws)
                     return;
+                if (++failures < 3)
+                    return;
                 this._logger.warn('WebSocket pong timeout; terminating half-open connection');
                 try {
                     ws.terminate?.();
@@ -2036,6 +2205,10 @@ export class RPCTransport {
                 target.removeListener('pong', this._pongListener);
             }
             catch { /* noop */ }
+            try {
+                target.removeListener('message', this._pongListener);
+            }
+            catch { /* noop */ }
         }
         this._pongListener = null;
         this._heartbeatWs = null;
@@ -2048,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) {
@@ -2061,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;
@@ -2102,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);
@@ -2115,7 +2302,8 @@ export class RPCTransport {
         if (this._metaObserver !== null && meta !== null && typeof meta === 'object' && !Array.isArray(meta)) {
             this._queueObserver('notification meta_observer', this._metaObserver, meta);
         }
-        this._dispatcher.enqueueProtocol('notification', message);
+        const generation = this._connectionGeneration;
+        this._dispatcher.enqueueProtocol('notification', message, method === 'gateway.migrate' ? () => generation === this._connectionGeneration : undefined);
     }
     /** 解码 WebSocket 消息为 JSON 对象 */
     _decodeMessage(raw) {
@@ -2212,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(),
     });
@@ -2234,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?.();
@@ -2241,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 ?? ''));