Restore a placeholder a stream-mode reply plugin splits in two

crates/tw-gateway/src/plugin/bridge.rs

@@ -38,6 +38,11 @@ pub struct Bridge {
     /// `values` 的下标，按值的字节序排：一段尾巴是不是哪个值的开头，二分查一次（见
     /// [`Bridge::hold_from`]）
     by_bytes: Vec<usize>,
+    /// `values` 的下标，按占位符的字节序排：插件交回来的一段结尾是不是哪个占位符的开头，
+    /// 二分查一次（见 [`Bridge::unfinished_from`]）
+    by_placeholder: Vec<usize>,
+    /// 最长的占位符有多长：结尾扣住的不会比它长
+    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<usize> = (0..values.len()).collect();
         by_bytes.sort_unstable_by(|&a, &b| values[a].0.cmp(&values[b].0));
+        let mut by_placeholder: Vec<usize> = (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 [
+            ("用 <<TW_SECRET_1", "<<TW_SECRET_1"),
+            ("用 <<TW_SECRET_1>", "<<TW_SECRET_1>"),
+            ("用 <", "<"),
+            ("<<<TW_SEC", "<<TW_SEC"),
+            ("用 <<TW_SECRET_1>>", ""),
+            // 账里没有的占位符、C++ 的 `<<` 不扣
+            ("用 <<TW_SECRET_9", ""),
+            ("std::cout << x", ""),
+        ] {
+            assert_eq!(&s[b.unfinished_from(s)..], held, "{s:?}");
+        }
+        assert_eq!(
+            Bridge::new(rules()).unfinished_from("<<TW_"),
+            5,
+            "空账什么都不扣"
+        );
+    }
+
     /// 一本账里装着这些值（按出现的先后发号）
     fn bridge_with(values: &[&str]) -> Bridge {
         use tw_guard::redact::rules::{Hit, Rule};

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<Option<String>, 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<String, GatewayError> {
         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,32 +567,37 @@ 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) => {}
                             Ok(Some(s)) => {
                                 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);
                 }
             }