Files
DevFlow/src-tauri/src/state/llm_concurrency.rs
T
lxy 2187ffe173 新增: spawn_agent 子代理派发与级联取消
- 纯内存子会话隔离上下文,一轮回传结论报告
- 白名单钳制低风险并恒剔 spawn_agent 防递归
- 子代理策略下高风险自动拒绝,不挂审批防卡死
- 轮数令牌挂钟三重熔断
- 父停止与删会话级联取消子代理
- AiAgentFinished 事件挂父会话
2026-10-01 19:35:21 +08:00

437 lines
25 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
use std::collections::HashMap;
use std::sync::atomic::AtomicUsize;
use std::sync::Arc;
use tokio::sync::{Mutex, Semaphore};
// ============================================================
// LLM 调用并发控制(双层 Semaphore + 可选 per-provider 层)
// ============================================================
/// LLM 调用并发控制 — 全局 LLM 调用限流 global + 单对话内 per_conv 双层 Semaphore + 可选 per-provider 层
///
/// 限流对象:所有真实 LLM 调用(主循环 stream_llm / 标题生成 / 知识提炼)。
/// 不限流本地工具执行(tools.execute)——本地操作无外部成本、不受 RPM 约束。
///
/// ## global(permits 默认 3,「LLM 调用并发限流」原义)
/// F-09 batch5 曾把 global 改「并发会话数上限」(loop 入口持整 loop),用户决策修正「不设并发会话上限」后
/// 已删除 loop 入口 acquire。global 回归原义:由各单次 LLM 调用点(stream_llm 重试循环 / 标题 / 压缩 /
/// 提炼 / 项目分析扫描)各自 `acquire_global()` 拿 permit、调用结束 Drop 释放,防 provider 429。
/// 多对话 loop 并发不限数(用户接受 token 暴增)。
///
/// ## per_conv 改 HashMap<conv_id, Semaphore>
/// 原应用级单信号量(AiSession 单例 + generating 互斥下退化)改为
/// `HashMap<String, Arc<Semaphore>>`:每对话一份(permits=2,主循环 + 标题 + 提炼各自限流)。
/// `run_agentic_loop` 入口 `acquire_per_conv(conv_id)` 拿 1 permit 持整个 loop 生命周期(含工具执行/审批
/// 等待/重试),防单对话内并发 LLM 调用失控。这是单对话内限流,非会话数限制,符合用户「不设上限」。
/// `acquire_per_conv(conv_id)`:lock HashMap → 无则建(permits=2)→ clone Arc → 释放 lock → acquire_owned。
/// conv 退出清理(`release_conv(conv_id)`):conv 删除时 remove 条目(防 HashMap 无限增长)。
///
/// 运行时调整:tokio Semaphore 的 permits 数构造时固定、不可增减,
/// 故 `global` 用 `Arc<Mutex<Arc<Semaphore>>>` 双层包装——替换内层 Arc 即重建 Semaphore。
/// 已持有旧 permit 的任务不受影响(permit 绑定旧 Semaphore,软收敛),
/// 新请求 lock 后克隆到最新 Arc、自动走新限制。旧 Semaphore 随最后 permit 释放而 drop。
/// per_conv HashMap 内每条 Arc<Semaphore> 不需运行时改 permits(对话内并发上限 2 固定),故无重建需求。
///
/// ## per-provider 层(可选)
/// `per_provider` 为 HashMap<provider_id, Arc<Semaphore>>。调用方(agentic loop)经
/// `try_acquire_for_provider(pid)` 取额外 permit,防单 provider 被打满(限流 429)。
/// **单 provider 场景**:若未调 `set_provider_caps`,HashMap 为空,
/// `try_acquire_for_provider` 返 `NotConfigured`(无限流,行为同 F-01 前)。零变化保证。
/// **三态语义(G4.1,2026-08-05)**:`NotConfigured`(未配置无限流) / `Acquired(permit)`
/// (拿到) / `Exhausted`(占满,非阻塞跳下一 candidate)。修复旧阻塞 `acquire_for_provider`
/// 在主 provider 信号量占满时整条 fallback 链卡死(降级失效)。
/// 全局容量 = min(sum(各 provider 上限), global_cap):由调用方在配置时约束
/// (set_provider_caps 传 min(sum, global_cap)),非运行时强约束。
///
/// ## Phase3 预留: per_sub_flow 层(批2-B,占位未接入)
/// `per_sub_flow` 为 HashMap<sub_flow_id, Arc<Semaphore>>,用于 Phase3 单对话并行多轮的
/// **子流并发上限**(单对话内并行 spawn 多条子流时,防子流数失控)。本批仅加层 + 占位,
/// **不接入调用点**(接入待 Phase3 子流 spawn 落地)。
///
/// **key 约定(文档化,非强约束)**:`sub_flow_id` 采用 `"{conv_id}::{sub_id}"` 格式,
/// 唯一标识某对话下的某条子流,便于 release_sub_flow 精确清理。
///
/// **三层语义**:
/// - per_sub_flow = 单对话内**子流**并发上限(每条子流一份 Semaphore,permits 默认 3,预留)
/// - per_conv = 单对话内**总**并发上限(permits 默认 2)
/// - global = 全应用**总**并发上限(permits 默认 3)
///
/// **理想配额(用户配置建议,非硬限)**: per_sub_permits ≤ per_conv_permits ≤ global_permits。
/// 实际限流由各层 Semaphore 自然保证:global Semaphore 硬限全应用并发不超 global permits
/// (跨对话 sum(per_conv) 可 > global,但 global acquire_owned 自然阻塞,不超卖)。
///
/// **acquire 语义(非阻塞)**:`acquire_per_sub_flow` 用 `try_acquire_owned`(非阻塞),
/// 耗尽返 None 降级串行——子流并行不阻塞主 loop(Phase3 子流 spawn 失败不拖垮整对话)。
#[derive(Clone)]
pub struct LlmConcurrency {
/// 全局 LLM 调用并发上限(permits 默认 3):F-09 batch5 修正后回归原义,由各单次 LLM 调用点
/// (stream_llm/标题/压缩/提炼/项目分析)各自 acquire/drop 防 429,不再由 loop 入口持整 loop。
global: Arc<Mutex<Arc<Semaphore>>>,
/// 单对话内并发上限(F-09 B 批5):HashMap<conv_id, Semaphore>,每对话 permits=2。
/// acquire 时按 conv_id 取/建;release_conv 在 loop 结束 + 无 pending 时 remove。
per_conv: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
/// 每对话内并发上限(permits 默认 2,构造时传入 new(3, 2) 的第二参)。
/// F-09 B 批5:AtomicUsize 支持运行时热改(ai_set_concurrency_config 的 per_conv_limit)。
/// 热改后**已建对话**的旧 Semaphore 不变(permits 构造时固定),**新建对话**用新值;
/// 为使热改立即全量生效,set_per_conv 同时清空 HashMap 强制重建(软收敛:旧 permit 随 Drop 释放)。
per_conv_permits: Arc<AtomicUsize>,
/// per-provider 信号量表。空 = 无 per-provider 限流(单 provider 路径零变化)。
/// Arc<Mutex<HashMap>>:运行时增删 provider 配置时替换/插入,acquire 时 clone Arc。
per_provider: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
/// Phase3 预留(批2-B): per-sub_flow 信号量表。key = sub_flow_id(约定 "{conv_id}::{sub_id}")。
/// 空 = 无子流并发限流(Phase3 未接入时零变化)。acquire 用 try_acquire_owned(非阻塞),
/// 耗尽返 None 降级串行,子流并行不阻塞主 loop。
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
per_sub_flow: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
/// Phase3 预留(批2-B): 每子流并发上限(permits 默认 3,内部初始化,预留)。
/// AtomicUsize 支持运行时热改;set_per_sub_permits 同时清空 HashMap(软收敛,对齐 set_per_conv)。
/// new(global, per_conv) 签名保持向后兼容,per_sub_permits 内部 AtomicUsize::new(3)。
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
per_sub_permits: Arc<AtomicUsize>,
}
/// F-260614-04(2026-08-05 G4.1 裁决): per-provider 限流 acquire 的**显式三态**结果。
///
/// **None 双关陷阱(本枚举的动机)**:旧 `acquire_for_provider` 用 `Option` 返回值,
/// `None` 既表示「未配置无限流」又可能被 naive 调用方误读为「耗尽」——耗尽语义下
/// naive 用 None 表耗尽会静默绕过限流打爆 429。三态把两种语义显式拆分:
///
/// - `NotConfigured`:未配置(per_provider 表无该 provider → 无限流,permits 跳过不落表)。
/// - `Acquired(permit)`:成功拿到 permit(绑 Semaphore,Drop 自动释放)。
/// - `Exhausted`:信号量占满(非阻塞 try_acquire 失败),调用方跳下一 candidate。
pub enum ProviderAcquire {
/// 未配置(per_provider 表无该 provider → 无限流,permits 跳过不落表)。
/// 与耗尽语义完全区分:调用方无 permit 直接 proceed(单 provider 零变化)。
NotConfigured,
/// 成功拿到 permit(绑 Semaphore,Drop 自动释放)。
Acquired(tokio::sync::OwnedSemaphorePermit),
/// 该 provider 信号量占满(try_acquire_owned 失败,非阻塞)。
/// 调用方跳下一 candidate(候选链降级),不阻塞等待。
Exhausted,
}
impl LlmConcurrency {
pub fn new(global: usize, per_conv: usize) -> Self {
Self {
global: Arc::new(Mutex::new(Arc::new(Semaphore::new(global)))),
per_conv: Arc::new(Mutex::new(HashMap::new())),
per_conv_permits: Arc::new(AtomicUsize::new(per_conv)),
per_provider: Arc::new(Mutex::new(HashMap::new())),
// Phase3 预留(批2-B): per_sub_flow 默认 permits=3,内部初始化。
// 签名保持 new(global, per_conv) 向后兼容;Phase3 接入后用户可经
// set_per_sub_permits 热改(配置化待后续批,本批仅占位)。
per_sub_flow: Arc::new(Mutex::new(HashMap::new())),
per_sub_permits: Arc::new(AtomicUsize::new(3)),
}
}
/// 取全局 LLM 调用并发 permit:由各单次 LLM 调用点(stream_llm/标题/压缩/提炼/项目分析)
/// 各自调用、调用结束 Drop 释放。F-09 batch5 修正后不再由 run_agentic_loop 入口持整 loop。
/// 重建后(set_global)新请求自动走最新 Semaphore。
pub async fn acquire_global(&self) -> tokio::sync::OwnedSemaphorePermit {
let sema = self.global.lock().await.clone();
sema.acquire_owned().await.expect("llm global semaphore closed")
}
/// 取全局 LLM 调用并发 permit 的**非阻塞**版本(对齐 try_acquire_per_conv 语义)。
///
/// 用途:spawn_agent 子代理循环入口一次性取 global + per_conv(child) 双份额,
/// 任一耗尽即 failed(子代理不阻塞等待,防占满父 loop 的执行槽)。
/// 锁持有短(不含 try_acquire)。
pub async fn try_acquire_global(&self) -> Option<tokio::sync::OwnedSemaphorePermit> {
let sema = self.global.lock().await.clone();
// try_acquire_owned 非阻塞:Err(NoPermits) 返 None,不阻塞调用方。
// Semaphore 不会 close(无 close() 调用),Err 仅可能为 NoPermits。
sema.try_acquire_owned().ok()
}
/// 取单对话内并发 permit(F-09 B 批5):按 conv_id 取/建 Semaphore,permits=当前 per_conv_permits(默认 2)。
/// lock HashMap → 无则建 → clone Arc → 释放 lock → acquire_owned。锁持有短(不含 await acquire)。
pub async fn acquire_per_conv(&self, conv_id: &str) -> tokio::sync::OwnedSemaphorePermit {
let sema = {
let mut map = self.per_conv.lock().await;
let permits = self.per_conv_permits.load(std::sync::atomic::Ordering::SeqCst);
map.entry(conv_id.to_string())
.or_insert_with(|| Arc::new(Semaphore::new(permits)))
.clone()
};
sema.acquire_owned().await.expect("llm per_conv semaphore closed")
}
/// 取单对话内并发 permit 的**非阻塞**版本(对齐 `acquire_per_sub_flow` 语义)。
///
/// 与 `acquire_per_conv` 同样按 conv_id 取/建 Semaphore(permits=当前 per_conv_permits),
/// 区别仅在最后一步用 `try_acquire_owned`(非阻塞):耗尽返 `None`,调用方降级,不阻塞。
///
/// **用途**:压缩(`compress_via_llm`)在 agentic loop 内被调用,loop 入口已持 1 个 per_conv
/// permit;当 `set_per_conv(1)` 后 per_conv Semaphore 仅 1 permit,loop 入口占满,压缩再
/// `acquire_per_conv`(阻塞)将永久等待 → `is_compressing` 永真 → 卡死。改用本非阻塞版本:
/// 拿不到(None)调用方提前返回降级(下轮再试),非崩溃。
///
/// - 返 `Some(permit)`:成功获取,permit 随 Drop 自动释放。
/// - 返 `None`:本对话并发已耗尽,调用方降级(非错误)。锁持有短(不含 try_acquire)。
pub async fn try_acquire_per_conv(
&self,
conv_id: &str,
) -> Option<tokio::sync::OwnedSemaphorePermit> {
let sema = {
let mut map = self.per_conv.lock().await;
let permits = self.per_conv_permits.load(std::sync::atomic::Ordering::SeqCst);
map.entry(conv_id.to_string())
.or_insert_with(|| Arc::new(Semaphore::new(permits)))
.clone()
};
// try_acquire_owned 非阻塞:Err(NoPermits) 返 None 降级,不阻塞调用方。
// 对齐 acquire_global/per_conv 的 expect 前提:Semaphore 不会 close(无 close() 调用)。
sema.try_acquire_owned().ok()
}
/// F-09 B 批5:conv 退出清理。loop 结束 + 无 pending 审批时 remove 该 conv 的 Semaphore 条目。
/// **时机由调用方判断**(agentic loop 退出点):仅在确信无后续 acquire 时调用,否则误删会致
/// 该 conv 下次 acquire 重建 Semaphore(限流计数清零,非致命,但语义偏离)。
/// remove 后已持 permit 不受影响(permit 绑旧 Arc,随 Drop 释放),仅阻止新条目累积。
pub async fn release_conv(&self, conv_id: &str) {
self.per_conv.lock().await.remove(conv_id);
}
/// 重建全局 Semaphore(软收敛:旧 permit 不回收,待其释放后新限制完全生效)
pub async fn set_global(&self, permits: usize) {
*self.global.lock().await = Arc::new(Semaphore::new(permits));
}
/// 重建单对话 Semaphore(F-09 B 批5 后 per_conv 为 HashMap)。
/// 行为:更新 per_conv_permits(AtomicUsize) + 清空 HashMap(软收敛:旧 permit 随 Drop 释放,
/// 新对话 acquire 用新 permits 值重建 Semaphore)。已建对话若仍在跑,旧 Semaphore 不变;
/// 下次该 conv 新 acquire 时因 HashMap 已清空会重建为新 permits。
pub async fn set_per_conv(&self, permits: usize) {
self.per_conv_permits
.store(permits, std::sync::atomic::Ordering::SeqCst);
self.per_conv.lock().await.clear();
}
/// 读当前 per_conv permits 生效值(对齐 set_per_conv 的写,补足读侧 getter)。
///
/// 用途:get_app_config AI 工具查「当前 LLM 并发配置」时返给 LLM(global 无对称 getter:
/// Semaphore permits 封装在 Arc<Mutex<Arc<Semaphore>>> 内层无读取接口,故此处只返 per_conv,
/// global 不返避免误导)。热改(set_per_conv)后立即反映,无需重启。
pub fn current_per_conv_permits(&self) -> usize {
self.per_conv_permits.load(std::sync::atomic::Ordering::SeqCst)
}
// ============================================================
// Phase3 预留(批2-B): per_sub_flow 层 — 占位未接入调用点
// ============================================================
/// Phase3 预留(批2-B): 取单子流内并发 permit(非阻塞)。
///
/// 按 `sub_flow_id` 取/建 Semaphore(permits=当前 per_sub_permits,默认 3)。
/// lock HashMap → 无则建(or_insert_with, permits 取 per_sub_permits 当前值)
/// → clone Arc → 释放 lock → **try_acquire_owned**(非阻塞)。
///
/// **非阻塞语义(关键)**:用 try_acquire_owned 而非 acquire_owned——耗尽返 None
/// 调用方降级串行处理,子流并行不阻塞主 loop(Phase3 子流 spawn 失败不拖垮整对话)。
///
/// - 返 `Ok(Some(permit))`:成功获取,permit 随 Drop 自动释放。
/// - 返 `Ok(None)`:子流并发已耗尽,调用方降级串行(非错误)。
///
/// **key 约定**:`sub_flow_id` 采用 `"{conv_id}::{sub_id}"` 格式(文档化,调用方组装)。
/// 锁持有短(不含 try_acquire),对齐 acquire_per_conv 模式。
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
pub async fn acquire_per_sub_flow(
&self,
sub_flow_id: &str,
) -> Option<tokio::sync::OwnedSemaphorePermit> {
let sema = {
let mut map = self.per_sub_flow.lock().await;
let permits = self.per_sub_permits.load(std::sync::atomic::Ordering::SeqCst);
map.entry(sub_flow_id.to_string())
.or_insert_with(|| Arc::new(Semaphore::new(permits)))
.clone()
};
// try_acquire_owned 非阻塞:Err(NoPermits) 返 None 降级串行,不阻塞主 loop。
// 对齐 acquire_global/per_conv 的 expect 前提:Semaphore 不会 close(无 close() 调用)。
sema.try_acquire_owned().ok()
}
/// Phase3 预留(批2-B): 子流完成清理。remove 该 sub_flow 的 Semaphore 条目(防 HashMap 无限增长)。
///
/// **对齐 release_conv 模式**:remove 后已持 permit 不受影响(permit 绑旧 Arc,随 Drop 释放),
/// 仅阻止新条目累积。时机由调用方判断(Phase3 子流退出点)。
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
pub async fn release_sub_flow(&self, sub_flow_id: &str) {
self.per_sub_flow.lock().await.remove(sub_flow_id);
}
/// Phase3 预留(批2-B): 热改 per_sub_flow permits 上限。
///
/// 行为(对齐 set_per_conv):更新 per_sub_permits(AtomicUsize) + 清空 HashMap(软收敛:
/// 旧 permit 随 Drop 释放,新子流 acquire 用新 permits 值重建 Semaphore)。已建子流若仍在跑,
/// 旧 Semaphore 不变;下次该 sub_flow 新 acquire 时因 HashMap 已清空会重建为新 permits。
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
pub async fn set_per_sub_permits(&self, permits: usize) {
self.per_sub_permits
.store(permits, std::sync::atomic::Ordering::SeqCst);
self.per_sub_flow.lock().await.clear();
}
/// F-260614-04(2026-08-05 G4.1 裁决): per-provider 并发 permit 的**非阻塞三态**版本。
///
/// **背景**:旧 `acquire_for_provider` 用阻塞 `acquire_owned()`,主 provider 信号量占满时
/// 整条 fallback 候选链卡死(降级失效)。compress/title 已改 try_acquire,候选链是唯一漏网。
/// 本方法用 `try_acquire_owned`(非阻塞),三态返回见 [`ProviderAcquire`]:
///
/// - `NotConfigured`:per_provider 表无该 provider(未配置 → 无限流,单 provider 零变化)。
/// - `Acquired(permit)`:成功拿到,permit 随 Drop 自动释放。
/// - `Exhausted`:信号量占满,调用方跳下一 candidate(候选链降级),不阻塞。
///
/// **仅 per-provider 层非阻塞**:global/per_conv 仍由 `acquire_global`/`acquire_per_conv`
/// 阻塞限流(调用方各自持有),不受本方法影响。
///
/// 锁持有短(不含 try_acquire),对齐 `acquire_per_sub_flow` 模式。
///
/// **调用方用法**(agentic 候选循环):
/// ```ignore
/// match llm_concurrency.try_acquire_for_provider(&candidate.id).await {
/// ProviderAcquire::NotConfigured => { /* 无 permit 直接 proceed */ }
/// ProviderAcquire::Acquired(permit) => { /* 持 permit stream,Drop 自动释放 */ }
/// ProviderAcquire::Exhausted => { /* 跳下一 candidate */ }
/// }
/// ```
pub async fn try_acquire_for_provider(
&self,
provider_id: &str,
) -> ProviderAcquire {
let sema = {
let map = self.per_provider.lock().await;
map.get(provider_id).cloned()
};
match sema {
// 无配置:无限流(permits 跳过不落表),显式区分于耗尽。
None => ProviderAcquire::NotConfigured,
Some(sema) => match sema.try_acquire_owned() {
// 成功拿到:permit 绑 Semaphore,Drop 自动释放。
Ok(permit) => ProviderAcquire::Acquired(permit),
// 占满:非阻塞返 Exhausted,调用方跳下一 candidate(降级)。
// try_acquire_owned 不会 Err 于 closed(Semaphore 无 close() 调用)。
Err(_) => ProviderAcquire::Exhausted,
},
}
}
/// 批量设置 per-provider 并发上限(替换整表)。
///
/// 调用方(Settings 配置热改 / 启动初始化)传入 `{ provider_id: permits }` map,
/// 替换整张 per_provider 表(软收敛:已持 permit 不回收)。全局容量约束 min(sum, global_cap)
/// 由调用方在构造 map 时应用(本函数不做强约束,只落表)。
///
/// 传空 map → 清空 per_provider 表(所有 provider 回退到无 per-provider 限流)。
pub async fn set_provider_caps(&self, caps: HashMap<String, usize>) {
let mut map = self.per_provider.lock().await;
map.clear();
for (pid, permits) in caps {
// permits=0 等同无配置(Semaphore::new(0) 永远 acquire 不到 → 死锁),
// 故 permits=0 跳过(不落表 → try_acquire_for_provider 返 NotConfigured → 无限流)。
if permits > 0 {
map.insert(pid, Arc::new(Semaphore::new(permits)));
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
// ============================================================
// F-260614-04(2026-08-05 G4.1): try_acquire_for_provider 三态测试
// NotConfigured / Acquired / Exhausted 显式区分(修 None 双关:None='未配置无限流',
// naive 用 None 表耗尽会静默绕过限流)。仅验证 per-provider 层三态语义,
// global/per_conv 仍阻塞限流(不经本方法)。
// ============================================================
/// 未配置(空 per_provider 表)→ NotConfigured:无限流,单 provider 零变化。
#[tokio::test]
async fn try_acquire_for_provider_not_configured() {
let lc = LlmConcurrency::new(3, 2);
let r = lc.try_acquire_for_provider("p1").await;
assert!(matches!(r, ProviderAcquire::NotConfigured));
}
/// 配置且有空闲 permit → Acquired(permit 绑 Semaphore,Drop 自动释放)。
#[tokio::test]
async fn try_acquire_for_provider_acquired() {
let lc = LlmConcurrency::new(3, 2);
let mut caps = HashMap::new();
caps.insert("p1".to_string(), 2usize);
lc.set_provider_caps(caps).await;
let r = lc.try_acquire_for_provider("p1").await;
let permit = match r {
ProviderAcquire::Acquired(p) => p,
other => panic!("配置后应返回 Acquired,实际 {:?}", discriminant(&other)),
};
// Drop 释放后仍可再拿(permit 绑旧 Semaphore Arc,set 后计数守恒)。
drop(permit);
let r2 = lc.try_acquire_for_provider("p1").await;
assert!(matches!(r2, ProviderAcquire::Acquired(_)));
}
/// 配置但信号量占满 → Exhausted(非阻塞,不等待)。
#[tokio::test]
async fn try_acquire_for_provider_exhausted() {
let lc = LlmConcurrency::new(3, 2);
let mut caps = HashMap::new();
caps.insert("p1".to_string(), 1usize); // 仅 1 permit
lc.set_provider_caps(caps).await;
// 第一次拿到(permit 仍持于作用域内,不 Drop)。
let first = lc.try_acquire_for_provider("p1").await;
assert!(matches!(first, ProviderAcquire::Acquired(_)));
// 第二次耗尽:信号量满,非阻塞返 Exhausted。
let second = lc.try_acquire_for_provider("p1").await;
assert!(matches!(second, ProviderAcquire::Exhausted));
// Drop 首 permit 后恢复可再拿。
drop(first);
let third = lc.try_acquire_for_provider("p1").await;
assert!(matches!(third, ProviderAcquire::Acquired(_)));
}
/// 配置后清空(传空 map)→ 回退 NotConfigured(permits 跳过不落表)。
#[tokio::test]
async fn try_acquire_for_provider_not_configured_after_caps_cleared() {
let lc = LlmConcurrency::new(3, 2);
let mut caps = HashMap::new();
caps.insert("p1".to_string(), 1usize);
lc.set_provider_caps(caps).await;
lc.set_provider_caps(HashMap::new()).await; // 清空整表
let r = lc.try_acquire_for_provider("p1").await;
assert!(matches!(r, ProviderAcquire::NotConfigured));
}
/// 配置了 p1,未配置 p2 → p1 Acquired / p2 NotConfigured(按 provider 区分,不互相污染)。
#[tokio::test]
async fn try_acquire_for_provider_distinct_providers() {
let lc = LlmConcurrency::new(3, 2);
let mut caps = HashMap::new();
caps.insert("p1".to_string(), 1usize);
lc.set_provider_caps(caps).await;
// p1 配置 → Acquired;p2 未配置 → NotConfigured。
assert!(matches!(lc.try_acquire_for_provider("p1").await, ProviderAcquire::Acquired(_)));
assert!(matches!(lc.try_acquire_for_provider("p2").await, ProviderAcquire::NotConfigured));
// p1 仅 1 permit 已满 → p2 不受影响。
assert!(matches!(lc.try_acquire_for_provider("p2").await, ProviderAcquire::NotConfigured));
}
/// 辅助:仅用于 panic 信息里给枚举变体一个 Debug 标签(无需求时保持简单)。
fn discriminant(r: &ProviderAcquire) -> &'static str {
match r {
ProviderAcquire::NotConfigured => "NotConfigured",
ProviderAcquire::Acquired(_) => "Acquired",
ProviderAcquire::Exhausted => "Exhausted",
}
}
}