Skip to content

Commit 19dd641

Browse files
committed
refactor(rust): 拆 ai_refine_task 步 1 — 并发 worker 池
1 parent b3aa4a9 commit 19dd641

2 files changed

Lines changed: 282 additions & 252 deletions

File tree

‎app/src-tauri/src/ai_refine_task.rs‎

Lines changed: 6 additions & 252 deletions
Original file line numberDiff line numberDiff line change
@@ -9,18 +9,13 @@
99
//! 线程内再开 worker——网络调用阻塞 worker 线程,互不干扰;
1010
//! AiClient 为 Clone+Send,跨线程共享安全)。
1111
12-
use std::collections::VecDeque;
13-
use std::sync::mpsc;
14-
use std::sync::{Arc, Mutex};
15-
1612
use base64::Engine;
1713
use tauri::Emitter;
1814

1915
use crate::ai_chat::{AiTurn, trajectory_to_json};
2016
use crate::ai_client::AiClient;
2117
use crate::ai_mock::AiMockAdapter;
2218
use crate::ai_note_refine::AiNoteRefineAdapter;
23-
use crate::ai_refine_protocol::AiRefineRequest;
2419
use crate::ai_strategy::ResolvedDims;
2520
use crate::ai_task::{slice_note, AiTaskFailure, AiTaskState, SLICE_MAX_CHARS};
2621
use crate::commands::AppState;
@@ -31,6 +26,12 @@ use crate::note_filter::PurifyEnv;
3126
use crate::outline::{detect_outline_smart, OutlineConfig};
3227
use crate::video_profile::ProfileKind;
3328

29+
/// 并发切片精修 worker 池子模块(拆分目的与锁窗口/时序口径见 workers.rs 模块头)
30+
#[path = "ai_refine_task/workers.rs"]
31+
mod workers;
32+
/// 子模块再导出:既有 `crate::ai_refine_task::{…}` 导入路径保持不变。
33+
pub(crate) use self::workers::{refine_slices_concurrent, RefineCtx};
34+
3435
/// 切片并发上限(REQ-145:并发 2-3——配额并发安全由 command 层启动前按
3536
/// 预估片数一次性消耗保证,此处 worker 数不超切片数)。v0.17.0:pub(crate)
3637
/// ——笔记级精修任务共用同一上限。
@@ -66,8 +67,6 @@ struct RefineStreamPayload {
6667
fn emit_refine_stream(st: &AppState, task_id: u64, frame: RefineStreamFrame) {
6768
let _ = st.app.emit("ai:refine-stream", RefineStreamPayload { task_id, frame });
6869
}
69-
/// 单片失败重试次数(幂等片——同片重跑不产生副作用)。
70-
const SLICE_RETRY: usize = 1;
7170

7271
/// 精修任务审计记录(F1:REQ-140 轨迹可见化——summary 不含原文,隐私红线)。
7372
/// v0.17.0:summary_ctx 泛化("session=1"/"note=3"——会话级/笔记级共用)。
@@ -359,46 +358,6 @@ fn run_refine_task_inner(
359358
))
360359
}
361360

362-
/// 并发精修上下文(参数聚合——clippy too_many_arguments 修复;
363-
/// v0.17.0:字段 pub(crate)——笔记级精修任务共用)。
364-
pub(crate) struct RefineCtx<'a> {
365-
pub(crate) slices: &'a [String],
366-
pub(crate) chapters: &'a [String],
367-
pub(crate) glossary: &'a [String],
368-
pub(crate) profile: &'a str,
369-
pub(crate) adapter: &'a AiNoteRefineAdapter,
370-
pub(crate) mock_adapter: &'a AiMockAdapter,
371-
pub(crate) mock: bool,
372-
pub(crate) workers: usize,
373-
pub(crate) st: &'a AppState,
374-
pub(crate) task_id: u64,
375-
/// v0.12.0 M5:屏卡图 base64 data URI 列表(vision_refine_enabled 开启且
376-
/// 会话有归档图时非空;空 → 精修纯文本,现有行为零变化)
377-
pub(crate) vision_images: &'a [String],
378-
/// v0.17.0:策略解析结果(command 层 resolve——每片提示词一致)
379-
pub(crate) dims: &'a ResolvedDims,
380-
}
381-
382-
/// 片间摘要(F3 v2:前/后片首尾 N 字——提示词衔接上下文,防片间断裂)。
383-
/// 纯函数可单测:取片开头/结尾 SUMMARY_MAX_CHARS 字符(截断到字符边界)。
384-
pub(crate) fn slice_summary(text: &str, head: bool) -> Option<String> {
385-
let s = text.trim();
386-
if s.is_empty() {
387-
return None;
388-
}
389-
let max = crate::ai_refine_protocol::SUMMARY_MAX_CHARS;
390-
let out: String = if head {
391-
s.chars().take(max).collect()
392-
} else {
393-
s.chars().rev().take(max).collect::<String>().chars().rev().collect()
394-
};
395-
if out.chars().count() < s.chars().count() {
396-
Some(format!("{}…", out))
397-
} else {
398-
Some(out)
399-
}
400-
}
401-
402361
/// v0.12.0 M5:装载会话屏卡图(≤1280px 控 token → WebP/PNG base64 data URI)。
403362
///
404363
/// @ai-context: 仅 vision_refine_enabled 开启时调用;目录缺失/解码失败 → 跳过
@@ -463,208 +422,3 @@ fn load_session_vision_images(data_dir: &std::path::Path, session_id: i64) -> Ve
463422
}
464423
out
465424
}
466-
467-
/// 并发切片精修(纯编排):worker 池从共享队列取片 → 单片重试 → 收集。
468-
///
469-
/// @ai-context: 返回 (各片 markdown(保序,失败片跳过), 失败片数, 轨迹(每片
470-
/// 成功调用的提示词/回答——REQ-230;vision 调用只记图数占位,
471-
/// base64 不入轨迹库))。
472-
/// 单片失败不 panic、不中断其他片——部分成功语义(REQ-145)。
473-
/// 进度经 set_task 上报(finished = 已完成的片数,含失败片——
474-
/// 前端进度条推进不受单片失败影响)。
475-
/// @ai-context: F3 v2:请求携带 slice_index/slice_total/prev_summary/
476-
/// next_summary——模型知道自己是第几片、前后片衔接什么
477-
/// (防章节标题重复/内容断裂)。
478-
pub(crate) fn refine_slices_concurrent(ctx: RefineCtx<'_>) -> (Vec<String>, usize, Vec<AiTurn>) {
479-
let total = ctx.slices.len();
480-
if total == 0 {
481-
return (Vec::new(), 0, Vec::new());
482-
}
483-
let queue: Arc<Mutex<VecDeque<usize>>> = Arc::new(Mutex::new((0..total).collect()));
484-
let (tx, rx) = mpsc::channel::<(usize, Option<String>)>();
485-
let workers = ctx.workers.max(1).min(total);
486-
// 轨迹收集槽(REQ-230):worker 并发写,Mutex 保护;顺序由 turn 字段定
487-
let turns: Arc<Mutex<Vec<AiTurn>>> = Arc::new(Mutex::new(Vec::with_capacity(total)));
488-
// 请求一次构建(Arc 共享——worker 只读,避免每 worker 重复克隆切片)
489-
let reqs: Arc<Vec<AiRefineRequest>> = Arc::new(
490-
ctx.slices
491-
.iter()
492-
.enumerate()
493-
.map(|(i, s)| AiRefineRequest {
494-
content: s.clone(),
495-
profile: ctx.profile.to_string(),
496-
glossary: ctx.glossary.to_vec(),
497-
chapters: ctx.chapters.to_vec(),
498-
// F3 v2:片间上下文(前片结尾 / 后片开头摘要——衔接用)
499-
slice_index: i + 1,
500-
slice_total: total,
501-
prev_summary: if i > 0 { slice_summary(&ctx.slices[i - 1], false) } else { None },
502-
next_summary: ctx.slices.get(i + 1).and_then(|n| slice_summary(n, true)),
503-
})
504-
.collect(),
505-
);
506-
std::thread::scope(|scope| {
507-
for _ in 0..workers {
508-
let queue = queue.clone();
509-
let tx = tx.clone();
510-
let reqs = reqs.clone();
511-
// worker 捕获 ctx 字段(adapter/mock 只读共享;task_id 复制)
512-
let adapter = ctx.adapter;
513-
let mock_adapter = ctx.mock_adapter;
514-
let mock = ctx.mock;
515-
let task_id = ctx.task_id;
516-
let st = ctx.st; // REQ-290①:流式 Delta 帧推送(worker 线程内 emit)
517-
let profile = ctx.profile; // 轨迹 system 提示词构建用(模板分组同请求)
518-
let vision_images = ctx.vision_images;
519-
let dims = ctx.dims; // 策略解析结果(worker 只读共享)
520-
let turns = turns.clone();
521-
scope.spawn(move || loop {
522-
let idx = {
523-
let mut q = queue.lock().unwrap_or_else(|e| e.into_inner());
524-
q.pop_front()
525-
};
526-
let Some(idx) = idx else { break };
527-
let req = &reqs[idx];
528-
let mut outcome: Option<String> = None;
529-
for attempt in 0..=SLICE_RETRY {
530-
// REQ-290(v0.19.6)埋点先行:单片耗时归因——任务级 elapsed_ms
531-
// 已有(db_ai_tasks),此处补片级时间(含流式整包耗时口径)。
532-
let started = std::time::Instant::now();
533-
// REQ-290①(v0.19.7):首拍(attempt 0)优先流式(NDJSON 逐节,
534-
// 解析一节推一节 Delta——打字机正文);流式不可用/模型未遵守
535-
// 逐节约定 → 同拍回退非流式(与旧版逐字节一致);重试拍走非
536-
// 流式(幂等语义下不重复推流)。
537-
let stream_enabled = !mock
538-
&& attempt == 0
539-
&& std::env::var("REFINE_STREAM_NDJSON")
540-
.map(|v| v != "0")
541-
.unwrap_or(true);
542-
let resp = if mock {
543-
Ok(mock_adapter.refine(req))
544-
} else if stream_enabled {
545-
match adapter
546-
.refine_stream_ndjson(
547-
req,
548-
ctx.vision_images,
549-
Some(dims),
550-
|sec| {
551-
emit_refine_stream(
552-
st,
553-
task_id,
554-
RefineStreamFrame::Delta {
555-
slice_index: idx + 1,
556-
text: crate::ai_refine_protocol::render_sections(
557-
std::slice::from_ref(&sec),
558-
),
559-
},
560-
);
561-
},
562-
)
563-
.map_err(AiTaskFailure::from)
564-
{
565-
Ok(r) => Ok(r),
566-
Err(e) => {
567-
eprintln!(
568-
"[refine-task] task={} 片 {} 流式不可用,同拍回退非流式: {}",
569-
task_id, idx + 1, e.message()
570-
);
571-
if ctx.vision_images.is_empty() {
572-
adapter.refine(req, Some(dims)).map_err(AiTaskFailure::from)
573-
} else {
574-
adapter
575-
.refine_vision(req, ctx.vision_images, Some(dims))
576-
.map_err(AiTaskFailure::from)
577-
}
578-
}
579-
}
580-
} else if ctx.vision_images.is_empty() {
581-
adapter.refine(req, Some(dims)).map_err(AiTaskFailure::from)
582-
} else {
583-
// v0.12.0 M5:开启画面理解 → 屏卡图随切片请求送 vision-exp
584-
adapter
585-
.refine_vision(req, ctx.vision_images, Some(dims))
586-
.map_err(AiTaskFailure::from)
587-
};
588-
let elapsed_ms = started.elapsed().as_millis();
589-
match resp {
590-
Ok(r) => {
591-
eprintln!(
592-
"[refine-task] task={} 片 {}/{} attempt={} ok elapsed_ms={}",
593-
task_id, idx + 1, total, attempt + 1, elapsed_ms
594-
);
595-
// REQ-230:成功片记录轨迹(提示词/回答全文——任务对话视图)
596-
turns
597-
.lock()
598-
.unwrap_or_else(|e| e.into_inner())
599-
.push(AiTurn {
600-
turn: idx + 1,
601-
// 观察 2026-09-05-3(如实化):system=策略基座;流式
602-
// NDJSON 后缀/预算引导段为运行期动态段不入轨迹,
603-
// 取证以 response 全文为准
604-
system: adapter.prompt.build_system(profile, Some(dims)),
605-
user: turn_user_text(req, vision_images),
606-
response: serde_json::to_string(&r).unwrap_or_default(),
607-
});
608-
outcome = Some(r.to_markdown());
609-
break;
610-
}
611-
Err(e) if attempt < SLICE_RETRY => {
612-
eprintln!("[refine-task] task={} 片 {} 第{}次失败 elapsed_ms={},重试: {}", task_id, idx + 1, attempt + 1, elapsed_ms, e.message());
613-
}
614-
Err(e) => {
615-
eprintln!("[refine-task] task={} 片 {} 重试后仍失败 elapsed_ms={}(保留已成功片): {}", task_id, idx + 1, elapsed_ms, e.message());
616-
break;
617-
}
618-
}
619-
}
620-
let _ = tx.send((idx, outcome));
621-
});
622-
}
623-
drop(tx); // 所有 worker 结束后关闭通道(scope 内最后一个持有者)
624-
});
625-
// 收集(按消息携带的真实切片下标落位——channel 到达序 ≠ 切片序);
626-
// 进度按已收片数上报(含失败片——进度条推进不受单片失败影响)
627-
let received: Vec<(usize, Option<String>)> = rx.iter().collect();
628-
let mut by_index: Vec<Option<String>> = vec![None; total];
629-
for (pos, (slice_idx, out)) in received.iter().enumerate() {
630-
by_index[*slice_idx] = out.clone();
631-
set_task(
632-
ctx.st,
633-
ctx.task_id,
634-
AiTaskState::Running { finished_slices: pos + 1, total_slices: total },
635-
);
636-
// REQ-247:进度帧(片完成推进——前端进度行)
637-
emit_refine_stream(ctx.st, ctx.task_id, RefineStreamFrame::Progress {
638-
slice_index: pos + 1,
639-
slice_total: total,
640-
});
641-
// REQ-247:片级解析流帧(片完成即推——validate 已过;失败诚实提示)
642-
match out {
643-
Some(md) => emit_refine_stream(ctx.st, ctx.task_id, RefineStreamFrame::BlockDone {
644-
slice_index: slice_idx + 1,
645-
markdown: md.clone(),
646-
}),
647-
None => emit_refine_stream(ctx.st, ctx.task_id, RefineStreamFrame::SliceFailed {
648-
slice_index: slice_idx + 1,
649-
reason: "重试后仍失败(保留已成功片)".to_string(),
650-
}),
651-
}
652-
}
653-
let failed = by_index.iter().filter(|o| o.is_none()).count();
654-
let markdowns: Vec<String> = by_index.into_iter().flatten().collect();
655-
let turn_out = turns.lock().unwrap_or_else(|e| e.into_inner()).clone();
656-
(markdowns, failed, turn_out)
657-
}
658-
659-
/// 轨迹 user 文本:请求 JSON(vision 附加张数占位——base64 不入轨迹库,
660-
/// 原图在本机会话图库;REQ-230 提示词/回答可见 + 存储可控)。
661-
pub(crate) fn turn_user_text(req: &AiRefineRequest, images: &[String]) -> String {
662-
let mut s = serde_json::to_string(req).unwrap_or_default();
663-
if !images.is_empty() {
664-
s.push_str(&format!(
665-
"\n\n[附带画面图 {} 张——原始图在本机会话图库,轨迹不存 base64]",
666-
images.len()
667-
));
668-
}
669-
s
670-
}

0 commit comments

Comments
 (0)