diff --git a/crates/tw-gateway/src/plugin/bridge.rs b/crates/tw-gateway/src/plugin/bridge.rs index 40f0823d..b03e3509 100644 --- a/crates/tw-gateway/src/plugin/bridge.rs +++ b/crates/tw-gateway/src/plugin/bridge.rs @@ -38,6 +38,11 @@ pub struct Bridge { /// `values` 的下标，按值的字节序排：一段尾巴是不是哪个值的开头，二分查一次（见 /// [`Bridge::hold_from`]） by_bytes: Vec, + /// `values` 的下标，按占位符的字节序排：插件交回来的一段结尾是不是哪个占位符的开头， + /// 二分查一次（见 [`Bridge::unfinished_from`]） + by_placeholder: Vec, + /// 最长的占位符有多长：结尾扣住的不会比它长 + longest_placeholder: usize, } impl Bridge { @@ -48,6 +53,8 @@ impl Bridge { values: Vec::new(), lengths: Lengths::default(), by_bytes: Vec::new(), + by_placeholder: Vec::new(), + longest_placeholder: 0, } } @@ -106,8 +113,12 @@ impl Bridge { values.sort_by(|a, b| b.0.len().cmp(&a.0.len()).then_with(|| a.0.cmp(&b.0))); let mut by_bytes: Vec = (0..values.len()).collect(); by_bytes.sort_unstable_by(|&a, &b| values[a].0.cmp(&values[b].0)); + let mut by_placeholder: Vec = (0..values.len()).collect(); + by_placeholder.sort_unstable_by(|&a, &b| values[a].1.cmp(&values[b].1)); + self.longest_placeholder = values.iter().map(|(_, p)| p.len()).max().unwrap_or(0); self.lengths = Lengths::of(&values); self.by_bytes = by_bytes; + self.by_placeholder = by_placeholder; self.values = values; } @@ -265,6 +276,31 @@ impl Bridge { o.len() > tail.len() && o.starts_with(tail) }) } + + /// 流式插件交回来的一段，从哪个字节起是账里某个占位符**没写完**的开头：插件从占位符 + /// 中间切开它交出的文字时，前后两截各自都换不回原值，结尾这半个要扣住，和它交出的下一段 + /// 拼起来再换。 + /// + /// 返回 `s.len()` 是没有要扣的。只认账里的占位符（同 [`tw_guard::redact::stream`]）： + /// `<<` 在 C++ 代码、heredoc 里到处都是，不能见一个扣一个。扣住的因此不会比最长的占位符 + /// 长；切点在占位符开头那个 ASCII 字节上，总在字符边界上。 + pub fn unfinished_from(&self, s: &str) -> usize { + let b = s.as_bytes(); + (b.len().saturating_sub(self.longest_placeholder)..b.len()) + .find(|&at| s.is_char_boundary(at) && self.starts_a_placeholder(&b[at..])) + .unwrap_or(b.len()) + } + + /// `tail` 是不是账里某个占位符的**真**前缀（同 [`Bridge::starts_a_value`]） + fn starts_a_placeholder(&self, tail: &[u8]) -> bool { + let first = self + .by_placeholder + .partition_point(|&v| self.values[v].1.as_bytes() < tail); + self.by_placeholder[first..].iter().take(2).any(|&v| { + let p = self.values[v].1.as_bytes(); + p.len() > tail.len() && p.starts_with(tail) + }) + } } /// 账里的值按长度分组，在一段文字里一遍找出它们出现的每一处（见 [`Bridge::hide_known`]）。 @@ -481,6 +517,32 @@ mod tests { assert_eq!(b.hold_from("什么都不像"), "什么都不像".len()); } + #[test] + fn only_the_unfinished_start_of_a_placeholder_in_the_ledger_is_held() { + let mut b = Bridge::new(rules()); + b.learn( + format!("{{\"messages\":[{{\"content\":\"{KEY}\"}}]}}").as_bytes(), + crate::guard::TEST_SHAPE, + ); + for (s, held) in [ + ("用 <", "<"), + ("用 <", "<"), + ("<<>", ""), + // 账里没有的占位符、C++ 的 `<<` 不扣 + ("用 < Bridge { use tw_guard::redact::rules::{Hit, Rule}; diff --git a/crates/tw-gateway/src/plugin/reply/mod.rs b/crates/tw-gateway/src/plugin/reply/mod.rs index 09994fb6..c34d2af8 100644 --- a/crates/tw-gateway/src/plugin/reply/mod.rs +++ b/crates/tw-gateway/src/plugin/reply/mod.rs @@ -20,7 +20,8 @@ //! //! - **文字**按块交：整块模式攒齐一块再交一次，交回来的才发给客户端；流式模式每段 //! 增量交一次，交回什么现在就发什么（空串是先扣着），块结束时调 `onReplyTextEnd` -//! 把扣着的补上。几个插件串起来，前一个交出的是后一个收到的。 +//! 把扣着的补上。交回来的结尾是半个占位符的，那半个等它交出下一段拼齐了再换回原值。 +//! 几个插件串起来，前一个交出的是后一个收到的。 //! - **工具调用**攒到完整再交：不变、换成别的（一个或几个）、去掉。换出来的按客户端 //! 的格式写成完整的调用，后面块的序号跟着挪。 //! - 推理内容不交给插件，原样发。 @@ -112,6 +113,9 @@ struct StageLane { /// 流式模式：上一次交出非空的东西之后收到的原文。插件半路出错又是跳过时，它 /// 扣着的那些字从这里补回去，模型说过的话不丢 pending: String, + /// 流式模式：插件交回来的结尾是半个占位符（它从占位符中间切开了交出的文字），先扣着， + /// 和它交出的下一段拼起来再换回原值（见 [`Bridge::unfinished_from`]） + unrevealed: String, } #[derive(Default, Clone, Copy)] @@ -399,7 +403,11 @@ impl Chain { p: String, ) -> Result, GatewayError> { if !self.live(i) { - return Ok(Some(p)); + // 在别的块里出错被拿掉的：这一块里它攒着的原样接在这一段前面放出来 + return Ok(Some(match self.stages[i].lanes.remove(&lane) { + Some(l) => format!("{}{}{}{p}", l.unrevealed, l.pending, l.buf), + None => p, + })); } if self.stages[i].mode == ReplyMode::Block { self.stages[i] @@ -433,8 +441,12 @@ impl Chain { let hidden = self.bridge.hide(&send); let given = hidden.clone(); self.stages[i].counts.text_calls += 1; + let carrying = self.stages[i] + .lanes + .get(&lane) + .is_some_and(|l| !l.unrevealed.is_empty()); match self.invoke(i, move |r| r.on_text(&given)).await { - Ok(None) => { + Ok(None) if !carrying => { self.stages[i] .lanes .entry(lane) @@ -443,29 +455,49 @@ impl Chain { .clear(); Ok(Some(send)) } - Ok(Some(s)) => { + // 没改就是交回了交给它的那样：前一段结尾扣着半个占位符时，接在它后面一起换 + Ok(s) => { + let s = s.unwrap_or_else(|| hidden.clone()); if s != hidden { self.stages[i].counts.text_changed += 1; } - let out = self.bridge.reveal(&s); + let out = self.reveal(i, lane, &s, false); let l = self.stages[i].lanes.entry(lane).or_default(); - if out.is_empty() { + // 交出了东西（哪怕整段都还扣在半个占位符里）：收到的原文有了着落 + if s.is_empty() { l.pending.push_str(&send); - Ok(None) } else { l.pending.clear(); - Ok(Some(out)) } + Ok((!out.is_empty()).then_some(out)) } Err(e) => { self.fail(i, e)?; - // 跳过：它扣着的、这一段、还没交给它的尾巴，原样往下走 + // 跳过：它交出的半个占位符、它扣着的、这一段、还没交给它的尾巴，原样往下走 let l = self.stages[i].lanes.remove(&lane).unwrap_or_default(); - Ok(Some(format!("{}{send}{}", l.pending, l.buf))) + Ok(Some(format!( + "{}{}{send}{}", + l.unrevealed, l.pending, l.buf + ))) } } } + /// 流式插件在这一块里交回来的一段，换回原值。结尾是半个占位符的先扣着，和它交出的 + /// 下一段拼起来再换。`end`：块结束了，扣着的一起放出去，凑不完整的原样放 + fn reveal(&mut self, i: usize, lane: Lane, s: &str, end: bool) -> String { + let l = self.stages[i].lanes.entry(lane).or_default(); + let mut joined = std::mem::take(&mut l.unrevealed); + joined.push_str(s); + let cut = if end { + joined.len() + } else { + self.bridge.unfinished_from(&joined) + }; + l.unrevealed = joined.split_off(cut); + self.bridge.reveal(&joined) + } + /// 一块文字的一段增量。返回此刻该发给客户端的（可能是空的：插件扣着） pub async fn text(&mut self, lane: Lane, piece: &str) -> Result { let mut pieces = vec![piece.to_string()]; @@ -506,14 +538,14 @@ impl Chain { if !self.live(i) { // 半路被拿掉的：它攒着的原样放出来 if let Some(l) = self.stages[i].lanes.remove(&lane) { - out.push(format!("{}{}", l.pending, l.buf)); + out.push(format!("{}{}{}", l.unrevealed, l.pending, l.buf)); } carry = out; continue; } - let l = self.stages[i].lanes.remove(&lane).unwrap_or_default(); match self.stages[i].mode { ReplyMode::Block => { + let l = self.stages[i].lanes.remove(&lane).unwrap_or_default(); if !l.buf.is_empty() { let whole = l.buf; let hidden = self.bridge.hide(&whole); @@ -535,12 +567,19 @@ impl Chain { } } ReplyMode::Stream => { - if !l.buf.is_empty() { + // 这一块的账留到最后：扣着的半个占位符、出错时要补回去的原文都在上面 + let buf = self.stages[i] + .lanes + .get_mut(&lane) + .map(|l| std::mem::take(&mut l.buf)) + .unwrap_or_default(); + if !buf.is_empty() { // 扣着的尾巴到头了：不会再长成别的，交出去 - if let Some(o) = self.stream_call(i, lane, l.buf).await? { + if let Some(o) = self.stream_call(i, lane, buf).await? { out.push(o); } } + let mut last = String::new(); if self.live(i) && self.stages[i].text_end { match self.invoke(i, |r| r.on_text_end()).await { Ok(None) => {} @@ -548,19 +587,17 @@ impl Chain { if !s.is_empty() { self.stages[i].counts.text_changed += 1; } - out.push(self.bridge.reveal(&s)); + last = s; } Err(e) => { - let pending = self.stages[i] - .lanes - .remove(&lane) - .map(|l| l.pending) - .unwrap_or_default(); + let l = self.stages[i].lanes.remove(&lane).unwrap_or_default(); self.fail(i, e)?; - out.push(pending); + out.push(format!("{}{}", l.unrevealed, l.pending)); } } } + // 扣着的半个占位符接上最后这一段换回去，凑不完整的原样放出去 + out.push(self.reveal(i, lane, &last, true)); self.stages[i].lanes.remove(&lane); } } diff --git a/crates/tw-gateway/src/plugin/reply/tests/mod.rs b/crates/tw-gateway/src/plugin/reply/tests/mod.rs index 746d2127..7094708b 100644 --- a/crates/tw-gateway/src/plugin/reply/tests/mod.rs +++ b/crates/tw-gateway/src/plugin/reply/tests/mod.rs @@ -452,6 +452,269 @@ async fn a_failing_plugin_cuts_under_reject_and_steps_aside_under_skip() { assert_eq!(anthropic_blocks(&out)[1].2, "HELLO 世界"); } +// ───────────────────────────────────────────────────────── 切开的占位符 + +/// 流式：每段只交出除最后 `n` 个字以外的，扣着的块结束时放出来。交出去的切口会落在 +/// 占位符中间 +fn hold_last(n: usize) -> Double { + Double::new("hold-last") + .permit(&[Permission::ReplyText]) + .mode(ReplyMode::Stream) + .on_reply(true, true, false, move |_| { + let buf = Arc::new(Mutex::new(String::new())); + let (a, b) = (buf.clone(), buf); + Ok(Box::new(Closures { + text: Box::new(move |t| { + let mut buf = a.lock().unwrap(); + buf.push_str(t); + let cut = buf.char_indices().rev().nth(n - 1).map_or(0, |(i, _)| i); + Invocation::ok(Some(buf.drain(..cut).collect())) + }), + end: Box::new(move || { + Invocation::ok(Some(std::mem::take(&mut *b.lock().unwrap()))) + }), + tool: Box::new(|_| Invocation::ok(ToolCallOutcome::Unchanged)), + })) + }) +} + +/// 流式：记下每次收到的，原样放过 +fn spy(seen: Arc>>) -> Double { + Double::new("spy") + .permit(&[Permission::ReplyText]) + .mode(ReplyMode::Stream) + .on_reply(true, false, false, move |_| { + let seen = seen.clone(); + Ok(Box::new(Closures { + text: Box::new(move |t| { + seen.lock().unwrap().push(t.to_string()); + Invocation::ok(None) + }), + end: Box::new(|| Invocation::ok(None)), + tool: Box::new(|_| Invocation::ok(ToolCallOutcome::Unchanged)), + })) + }) +} + +/// 只有一块文字的 Anthropic 流，按这几段增量发 +fn anthropic_text(deltas: &[&str]) -> String { + let mut frames = vec![ev( + "content_block_start", + json!({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}), + )]; + for d in deltas { + frames.push(ev( + "content_block_delta", + json!({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text": d}}), + )); + } + frames.push(ev( + "content_block_stop", + json!({"type":"content_block_stop","index":0}), + )); + frames.concat() +} + +/// 请求里有 `KEY`：回答里出现时插件看到的是 `<>` +fn key_request() -> Value { + json!({ "messages": [{ "role": "user", "content": format!("key {KEY}") }] }) +} + +#[tokio::test] +async fn a_placeholder_a_stream_plugin_splits_is_restored_from_both_halves() { + let want = format!("use {KEY} now"); + let (head, tail) = (format!("use {}", &KEY[..12]), format!("{} now", &KEY[12..])); + let splits: [Vec<&str>; 3] = [vec![&want], vec!["use ", KEY, " now"], vec![&head, &tail]]; + for deltas in &splits { + let input = anthropic_text(deltas); + for plugins in [vec![hold_last(5)], vec![hold_last(5), hold_last(7)]] { + let n = plugins.len(); + for step in [3, 4096] { + let mut s = Stream::new( + chain_of(Dialect::Anthropic, plugins.clone(), &key_request()).await, + Framing::Sse, + ); + let (out, err) = run(&mut s, &input, step).await; + assert!(err.is_none()); + assert_eq!( + anthropic_blocks(&out)[0].2, + want, + "{deltas:?}, {n} plugins, step {step}" + ); + } + } + } +} + +#[tokio::test] +async fn the_next_plugin_in_the_chain_sees_whole_placeholders() { + let seen = Arc::new(Mutex::new(Vec::new())); + let input = anthropic_text(&["use ", KEY, " now"]); + let mut s = Stream::new( + chain_of( + Dialect::Anthropic, + vec![hold_last(5), spy(seen.clone())], + &key_request(), + ) + .await, + Framing::Sse, + ); + let (out, _) = run(&mut s, &input, 4096).await; + assert_eq!(anthropic_blocks(&out)[0].2, format!("use {KEY} now")); + let seen = seen.lock().unwrap().clone(); + assert_eq!(seen.concat(), "use <> now"); + // 每一段里的占位符都是整的 + for t in &seen { + assert_eq!( + t.matches("<<").count(), + t.matches("<>").count(), + "{seen:?}" + ); + } +} + +#[tokio::test] +async fn text_given_back_when_a_plugin_is_skipped_comes_out_restored() { + let want = format!("use {KEY} now"); + let input = anthropic_text(&["use ", KEY, " now"]); + // 后一个插件头一段就出错：前一个切开交出的照样拼回原值 + let boom = Double::new("boom") + .permit(&[Permission::ReplyText]) + .mode(ReplyMode::Stream) + .on_reply(true, false, false, |_| { + Ok(Box::new(Closures { + text: Box::new(|_| Invocation::err(RunError::CpuLimit)), + end: Box::new(|| Invocation::ok(None)), + tool: Box::new(|_| Invocation::ok(ToolCallOutcome::Unchanged)), + })) + }); + let mut chain = chain_of(Dialect::Anthropic, vec![hold_last(5), boom], &key_request()).await; + chain.stages[1].on_error = OnError::Skip; + let mut s = Stream::new(chain, Framing::Sse); + let (out, err) = run(&mut s, &input, 4096).await; + assert!(err.is_none()); + assert_eq!(anthropic_blocks(&out)[0].2, want); +} + +#[tokio::test] +async fn a_stream_plugin_failing_at_the_end_of_a_block_gives_back_what_it_held() { + let want = format!("use {KEY} now"); + let input = anthropic_text(&["use ", KEY, " now"]); + // 一直扣着、到块结束时才出错：扣着的那些原文补回来 + let late = Double::new("late") + .permit(&[Permission::ReplyText]) + .mode(ReplyMode::Stream) + .on_reply(true, true, false, |_| { + Ok(Box::new(Closures { + text: Box::new(|_| Invocation::ok(Some(String::new()))), + end: Box::new(|| Invocation::err(RunError::CpuLimit)), + tool: Box::new(|_| Invocation::ok(ToolCallOutcome::Unchanged)), + })) + }); + let mut chain = chain_of(Dialect::Anthropic, vec![late], &key_request()).await; + chain.stages[0].on_error = OnError::Skip; + let mut s = Stream::new(chain, Framing::Sse); + let (out, err) = run(&mut s, &input, 4096).await; + assert!(err.is_none()); + assert_eq!(anthropic_blocks(&out)[0].2, want); +} + +#[tokio::test] +async fn a_plugin_skipped_in_one_block_gives_back_another_open_blocks_text_in_order() { + // 两块文字同时开着、增量交替来。插件一直扣着，在 0 号块出错跳过：1 号块里它扣着的 + // （连同为了不交出半截密钥先扣着的尾巴）排在 1 号块后来的段前面，一个字不丢 + let fails_on_a2 = Double::new("fails-on-a2") + .permit(&[Permission::ReplyText]) + .mode(ReplyMode::Stream) + .on_reply(true, false, false, |_| { + Ok(Box::new(Closures { + text: Box::new(|t| { + if t.contains("A2") { + Invocation::err(RunError::CpuLimit) + } else { + Invocation::ok(Some(String::new())) + } + }), + end: Box::new(|| Invocation::ok(None)), + tool: Box::new(|_| Invocation::ok(ToolCallOutcome::Unchanged)), + })) + }); + let start = |i: u64| { + ev( + "content_block_start", + json!({"type":"content_block_start","index":i,"content_block":{"type":"text","text":""}}), + ) + }; + let delta = |i: u64, t: &str| { + ev( + "content_block_delta", + json!({"type":"content_block_delta","index":i,"delta":{"type":"text_delta","text":t}}), + ) + }; + let stop = |i: u64| { + ev( + "content_block_stop", + json!({"type":"content_block_stop","index":i}), + ) + }; + let input = [ + start(0), + start(1), + delta(1, "B1 "), + delta(0, "A1 "), + delta(1, &format!("B2 {}", &KEY[..12])), + delta(0, "A2 "), + delta(1, &format!("{} B3", &KEY[12..])), + stop(0), + stop(1), + ] + .concat(); + let mut chain = chain_of(Dialect::Anthropic, vec![fails_on_a2], &key_request()).await; + chain.stages[0].on_error = OnError::Skip; + let mut s = Stream::new(chain, Framing::Sse); + let (out, err) = run(&mut s, &input, 4096).await; + assert!(err.is_none()); + let blocks = anthropic_blocks(&out); + assert_eq!(blocks[0].2, "A1 A2 "); + assert_eq!(blocks[1].2, format!("B1 B2 {KEY} B3")); +} + +#[tokio::test] +async fn whole_bodies_restore_a_placeholder_a_stream_plugin_splits() { + let t = format!("use {KEY} now"); + let cases = [ + ( + Dialect::Anthropic, + json!({"id":"m","type":"message","role":"assistant", + "content":[{"type":"text","text":t}],"stop_reason":"end_turn"}), + ), + ( + Dialect::Chat, + json!({"id":"c","object":"chat.completion","choices":[{"index":0, + "message":{"role":"assistant","content":t},"finish_reason":"stop"}]}), + ), + ( + Dialect::Responses, + json!({"id":"r","object":"response","status":"completed","output":[ + {"type":"message","id":"msg","role":"assistant","content":[{"type":"output_text","text":t}]} + ]}), + ), + ( + Dialect::Gemini, + json!({"candidates":[{"content":{"role":"model","parts":[{"text":t}]},"finishReason":"STOP"}]}), + ), + ]; + for (d, body) in cases { + let mut chain = chain_of(d, vec![hold_last(5), hold_last(7)], &key_request()).await; + let out = whole(&mut chain, body.to_string().as_bytes()) + .await + .unwrap(); + let text = String::from_utf8(out).unwrap(); + assert!(text.contains(&t), "{d:?}: {text}"); + assert!(!text.contains("TW_SECRET"), "{d:?}: {text}"); + } +} + // ───────────────────────────────────────────────────────── Chat fn chat_stream() -> String {