1

CustomKeyboard/View/KBFunctionView.m

@@ -21,12 +21,10 @@
 #import "KBStreamTextView.h"  // 流式文本视图
 #import "KBStreamOverlayView.h" // 带关闭按钮的流式层
 #import "KBFunctionTagListView.h"
-#import "KBStreamFetcher.h"   // 网络流封装
+#import "WJXEventSource.h"
 #import "KBTagItemModel.h"
 #import <MJExtension/MJExtension.h>
 
-static NSTimeInterval const kKBStreamIdleTimeout = 5.0;
-
 @interface KBFunctionView () <KBFunctionBarViewDelegate, KBStreamOverlayViewDelegate, KBFunctionTagListViewDelegate>
 // UI
 @property (nonatomic, strong) KBFunctionBarView *barViewInternal;
@@ -42,12 +40,12 @@ static NSTimeInterval const kKBStreamIdleTimeout = 5.0;
 @property (nonatomic, strong, nullable) KBStreamOverlayView *streamOverlay;
 
 // 网络流式（封装）
-@property (nonatomic, strong, nullable) KBStreamFetcher *streamFetcher;
+@property (nonatomic, strong, nullable) WJXEventSource *eventSource;
 @property (nonatomic, assign) BOOL streamHasOutput; // 是否已输出过正文（首段去首个 \t 用）
 @property (nonatomic, strong, nullable) NSNumber *loadingTagIndex; // 当前显示loading的标签index
 @property (nonatomic, copy, nullable) NSString *loadingTagTitle;
-@property (nonatomic, strong, nullable) NSTimer *streamIdleTimer;
+@property (nonatomic, assign) BOOL eventSourceDidReceiveDone;
-@property (nonatomic, assign) BOOL streamTimeoutTriggered;
+@property (nonatomic, copy, nullable) NSString *eventSourceSplitPrefix;
 
 // Data
 //@property (nonatomic, strong) NSArray<NSString *> *itemsInternal;
@@ -272,10 +270,7 @@ static NSTimeInterval const kKBStreamIdleTimeout = 5.0;
 }
 
 
-#pragma mark - Network Streaming (POST)
+#pragma mark - Network Streaming (WJXEventSource)
-
-// 后端测试地址（内网）。若被 ATS 拦截，请在扩展 target 的 Info.plist 中允许 HTTP 或添加例外域。
-//static NSString * const kKBStreamDemoURL = @"http://192.168.1.144:7529/api/demo/talk";
 
 - (void)kb_startNetworkStreamingWithSeed:(NSString *)seedTitle {
     [self kb_stopNetworkStreaming];
@@ -285,7 +280,6 @@ static NSTimeInterval const kKBStreamIdleTimeout = 5.0;
     NSURL *url = [NSURL URLWithString:apiUrl];
     if (!url) { return; }
 
-    // 组装 POST Body
     NSInteger characterId = 0;
     if (self.loadingTagIndex != nil) {
         NSInteger idx = self.loadingTagIndex.integerValue;
@@ -294,13 +288,13 @@ static NSTimeInterval const kKBStreamIdleTimeout = 5.0;
             characterId = model.characterId;
         }
     }
-    NSString *message = seedTitle.length > 0 ? seedTitle : @"";
     NSInteger resolvedCharacterId = (characterId > 0) ? characterId : 75;
-    message = [NSString stringWithFormat:@"%@-%u",message,arc4random() % 10000];
+    NSString *message = seedTitle.length > 0 ? seedTitle : @"aliqua non cupidatat";
     NSDictionary *payload = @{
         @"characterId": @(resolvedCharacterId),
-        @"message": message.length > 0 ? message : @"aliqua non cupidatat"
+        @"message": @"dolore ea cillum"
     };
+    NSLog(@"[KBFunction] request payload: %@", payload);
     NSError *bodyError = nil;
     NSData *bodyData = [NSJSONSerialization dataWithJSONObject:payload options:0 error:&bodyError];
     if (bodyError || bodyData.length == 0) {
@@ -308,63 +302,178 @@ static NSTimeInterval const kKBStreamIdleTimeout = 5.0;
         return;
     }
 
-    self.streamHasOutput = NO; // 重置首段处理标记
+    NSMutableURLRequest *request = [NSMutableURLRequest requestWithURL:url cachePolicy:NSURLRequestReloadIgnoringCacheData timeoutInterval:60];
-    __weak typeof(self) weakSelf = self;
+    request.HTTPMethod = @"POST";
-    KBStreamFetcher *fetcher = [KBStreamFetcher fetcherWithURL:url];
+    [request setValue:@"application/json" forHTTPHeaderField:@"Content-Type"];
-    fetcher.httpMethod = @"POST";
-    fetcher.httpBody = bodyData;
-    NSMutableDictionary *headers = [@{ @"Content-Type": @"application/json" } mutableCopy];
     NSString *token = KBAuthManager.shared.current.accessToken ?: @"";
     if (token.length > 0) {
-        headers[@"auth-token"] = token;
+        [request setValue:token forHTTPHeaderField:@"auth-token"];
     }
-    fetcher.extraHeaders = headers;
+    request.HTTPBody = bodyData;
-    // 由本类统一做 /t->\t 与首段去 \t，fetcher 只负责增量与协议解析
+
-    fetcher.disableCompression = YES;
+    self.streamHasOutput = NO;
-    fetcher.acceptEventStream = NO;  // 响应头若是 SSE 仍会自动解析
+    self.eventSourceSplitPrefix = nil;
-    // 将 \t 与首段去 \t 的处理下沉到 Fetcher，避免 UI 抖动
+    self.eventSourceDidReceiveDone = NO;
-    fetcher.treatSlashTAsTab = YES;
+
-    fetcher.trimLeadingTabOnce = YES;
+    __weak typeof(self) weakSelf = self;
-    fetcher.flushInterval = 0.05; // 更接近“立刻显示”的节奏
+    WJXEventSource *source = [[WJXEventSource alloc] initWithRquest:request];
-    fetcher.onChunk = ^(NSString *chunk) {
+    source.ignoreRetryAction = YES;
+    [source addListener:^(WJXEvent * _Nonnull event) {
         __strong typeof(weakSelf) self = weakSelf; if (!self) return;
-        [self kb_appendChunkToStreamView:chunk];
+        [self kb_handleEventSourceMessage:event];
-    };
+    } forEvent:WJXEventNameMessage queue:NSOperationQueue.mainQueue];
-    fetcher.onFinish = ^(NSError * _Nullable error) {
+    [source addListener:^(WJXEvent * _Nonnull event) {
         __strong typeof(weakSelf) self = weakSelf; if (!self) return;
-        [self kb_invalidateStreamIdleTimer];
+        [self kb_handleEventSourceError:event.error];
-        if (self.streamTimeoutTriggered) {
+    } forEvent:WJXEventNameError queue:NSOperationQueue.mainQueue];
-            self.streamTimeoutTriggered = NO;
+    self.eventSource = source;
-            return;
+    [self.eventSource open];
-        }
-        // 若还未出现任何数据且仍在 loading，则取消小菊花
-        if (!self.streamHasOutput && self.loadingTagIndex) {
-            [self.tagListView setLoading:NO atIndex:self.loadingTagIndex.integerValue];
-            self.loadingTagIndex = nil; self.loadingTagTitle = nil;
-        }
-        BOOL isCancelledError = (error && error.code == NSURLErrorCancelled);
-        if (error && !isCancelledError) { [KBHUD showInfo:KBLocalized(@"拉取失败")]; }
-        if (self.streamOverlay) { [self.streamOverlay finish]; }
-        self.streamTimeoutTriggered = NO;
-    };
-    self.streamFetcher = fetcher;
-    self.streamTimeoutTriggered = NO;
-    [self kb_restartStreamIdleTimer];
-    [self.streamFetcher start];
 }
 
 - (void)kb_stopNetworkStreaming {
-    [self.streamFetcher cancel];
+    [self.eventSource close];
-    self.streamFetcher = nil;
+    self.eventSource = nil;
+    self.eventSourceSplitPrefix = nil;
+    self.eventSourceDidReceiveDone = NO;
     self.streamHasOutput = NO;
-    [self kb_invalidateStreamIdleTimer];
+}
-    self.streamTimeoutTriggered = NO;
+
+- (void)kb_handleEventSourceMessage:(WJXEvent *)event {
+    if (event.data.length == 0) { return; }
+    NSLog(@"[KBStream] SSE raw payload: %@", event.data);
+    NSData *jsonData = [event.data dataUsingEncoding:NSUTF8StringEncoding];
+    if (!jsonData) { return; }
+    NSError *error = nil;
+    NSDictionary *payload = [NSJSONSerialization JSONObjectWithData:jsonData options:0 error:&error];
+    if (error || ![payload isKindOfClass:[NSDictionary class]]) { return; }
+    NSString *type = payload[@"type"];
+    if (![type isKindOfClass:[NSString class]]) { return; }
+
+    if ([type isEqualToString:@"llm_chunk"]) {
+        NSString *chunk = [self kb_normalizedLLMChunkString:payload[@"data"]];
+        if (chunk.length > 0) {
+            [self kb_appendChunkToStreamView:chunk];
+        }
+        return;
+    }
+    if ([type isEqualToString:@"search_result"]) {
+        NSString *text = [self kb_formattedSearchResultString:payload[@"data"]];
+        if (text.length > 0) {
+            [self kb_appendChunkToStreamView:text];
+        }
+        return;
+    }
+    if ([type isEqualToString:@"done"]) {
+        self.eventSourceDidReceiveDone = YES;
+        [self kb_finishEventSourceWithError:nil];
+        return;
+    }
+}
+
+- (void)kb_handleEventSourceError:(NSError * _Nullable)error {
+    if (self.eventSourceDidReceiveDone) { return; }
+    [self kb_finishEventSourceWithError:error];
+}
+
+- (void)kb_finishEventSourceWithError:(NSError * _Nullable)error {
+    [self.eventSource close];
+    self.eventSource = nil;
+    if (!self.streamHasOutput && self.loadingTagIndex) {
+        [self.tagListView setLoading:NO atIndex:self.loadingTagIndex.integerValue];
+        self.loadingTagIndex = nil;
+        self.loadingTagTitle = nil;
+    }
+    BOOL shouldShowError = (error != nil);
+    if (shouldShowError) {
+        [KBHUD showInfo:error.localizedDescription ?: KBLocalized(@"拉取失败")];
+    }
+    if (self.streamOverlay) {
+        [self.streamOverlay finish];
+    }
+    self.eventSourceSplitPrefix = nil;
+    self.eventSourceDidReceiveDone = NO;
+}
+
+#pragma mark - Event Parsing
+
+- (NSString *)kb_normalizedLLMChunkString:(id)dataValue {
+    if (![dataValue isKindOfClass:[NSString class]]) { return @""; }
+    NSString *text = (NSString *)dataValue;
+    if (self.eventSourceSplitPrefix.length > 0) {
+        text = [self.eventSourceSplitPrefix stringByAppendingString:text ?: @""];
+        self.eventSourceSplitPrefix = nil;
+    }
+    text = [text stringByReplacingOccurrencesOfString:@"\r\n\t" withString:@"\t"];
+    text = [text stringByReplacingOccurrencesOfString:@"\n\t" withString:@"\t"];
+    text = [text stringByReplacingOccurrencesOfString:@"\r\t" withString:@"\t"];
+    while (text.length > 0) {
+        unichar c0 = [text characterAtIndex:0];
+        if (c0 == '\n' || c0 == '\r') {
+            text = [text substringFromIndex:1];
+            continue;
+        }
+        break;
+    }
+    text = [text stringByReplacingOccurrencesOfString:@"/t" withString:@"\t"];
+    text = [text stringByReplacingOccurrencesOfString:@"<SPLIT>" withString:@"\t"];
+    NSString *suffix = [self kb_pendingSplitSuffixForString:text];
+    if (suffix.length > 0) {
+        self.eventSourceSplitPrefix = suffix;
+        text = [text substringToIndex:text.length - suffix.length];
+    }
+    text = [text stringByReplacingOccurrencesOfString:@"\t " withString:@"\t"];
+    return text;
+}
+
+- (NSString *)kb_formattedSearchResultString:(id)dataValue {
+    if (![dataValue isKindOfClass:[NSArray class]]) { return @""; }
+    NSArray *list = (NSArray *)dataValue;
+    NSMutableArray<NSString *> *segments = [NSMutableArray array];
+    [list enumerateObjectsUsingBlock:^(id  _Nonnull obj, NSUInteger idx, BOOL * _Nonnull stop) {
+        NSString *payload = nil;
+        if ([obj isKindOfClass:[NSDictionary class]]) {
+            id val = obj[@"payload"];
+            if ([val isKindOfClass:[NSString class]]) {
+                payload = (NSString *)val;
+            }
+        } else if ([obj isKindOfClass:[NSString class]]) {
+            payload = (NSString *)obj;
+        }
+        payload = [payload stringByTrimmingCharactersInSet:[NSCharacterSet whitespaceAndNewlineCharacterSet]];
+        if (payload.length > 0) {
+            NSString *line = [NSString stringWithFormat:@"%lu. %@", (unsigned long)(segments.count + 1), payload];
+            [segments addObject:line];
+        }
+    }];
+    if (segments.count == 0) { return @""; }
+    NSString *title = KBLocalized(@"Search result");
+    NSMutableString *result = [NSMutableString stringWithFormat:@"\t%@:", title.length > 0 ? title : @"Search result"];
+    for (NSString *line in segments) {
+        [result appendFormat:@"\t%@", line];
+    }
+    return result;
+}
+
+- (NSString *)kb_pendingSplitSuffixForString:(NSString *)text {
+    static NSString * const token = @"<SPLIT>";
+    if (text.length == 0) { return @""; }
+    NSUInteger tokenLen = token.length;
+    if (tokenLen <= 1) { return @""; }
+    NSUInteger maxLen = MIN(tokenLen - 1, text.length);
+    for (NSUInteger len = maxLen; len > 0; len--) {
+        NSString *suffix = [text substringFromIndex:text.length - len];
+        NSString *prefix = [token substringToIndex:len];
+        if ([suffix isEqualToString:prefix]) {
+            return suffix;
+        }
+    }
+    return @"";
 }
 
 #pragma mark - Helpers
 
 /// 统一处理需要输出到 KBStreamTextView 的分片：
-/// - 目前网络层（KBStreamFetcher）已做 “/t->\t、首段去一个 \t、段间去一个空格”
+/// - 已将 `<SPLIT>` 转换为 `\t` 并去掉多余换行
 /// - 这里仅负责附加到视图与标记首段状态，避免 UI 抖动
 - (void)kb_appendChunkToStreamView:(NSString *)chunk {
     if (chunk.length == 0) return;
@@ -379,42 +488,6 @@ static NSTimeInterval const kKBStreamIdleTimeout = 5.0;
     if (!self.streamOverlay) return;
     [self.streamOverlay appendChunk:chunk];
     self.streamHasOutput = YES;
-    [self kb_restartStreamIdleTimer];
-}
-
-- (void)kb_restartStreamIdleTimer {
-    [self.streamIdleTimer invalidate];
-    if (!self.streamFetcher) {
-        self.streamIdleTimer = nil;
-        return;
-    }
-    __weak typeof(self) weakSelf = self;
-    self.streamIdleTimer = [NSTimer scheduledTimerWithTimeInterval:kKBStreamIdleTimeout repeats:NO block:^(NSTimer * _Nonnull timer) {
-        __strong typeof(weakSelf) self = weakSelf;
-        [self kb_handleStreamIdleTimeout];
-    }];
-}
-
-- (void)kb_invalidateStreamIdleTimer {
-    [self.streamIdleTimer invalidate];
-    self.streamIdleTimer = nil;
-}
-
-- (void)kb_handleStreamIdleTimeout {
-    [self kb_invalidateStreamIdleTimer];
-    if (!self.streamFetcher) { return; }
-    self.streamTimeoutTriggered = YES;
-    [self.streamFetcher cancel];
-    self.streamFetcher = nil;
-    if (self.loadingTagIndex) {
-        [self.tagListView setLoading:NO atIndex:self.loadingTagIndex.integerValue];
-        self.loadingTagIndex = nil;
-        self.loadingTagTitle = nil;
-    }
-    if (self.streamOverlay) {
-        [self.streamOverlay finish];
-    }
-    [KBHUD showInfo:KBLocalized(@"Response timeout")];
 }
 
 /// 统一更新左侧粘贴按钮的展示：