crabmate 0.5.0

Rust AI agent: OpenAI-compatible chat/completions, function calling, HTTP serve, ops CLI
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207

use crate::cm_sse_protocol::StreamEndReason;
use tokio::sync::mpsc;
use tokio::time::{Duration, timeout};

use crate::sse::SseStreamHub;
use crate::types::Message;

use super::ChatJobQueue;
use super::stream_finish::{
    emit_missing_final_response_fallback_if_needed, emit_stream_draining, emit_stream_ended_once,
    sse_payload_has_final_response_timeline,
};

#[tokio::test]
async fn queue_accepts_config_bounds() {
    let q = ChatJobQueue::new(2, 4);
    assert_eq!(q.max_concurrent(), 2);
    assert_eq!(q.max_pending(), 4);
}

#[tokio::test]
async fn queue_stream_cancel_flag_roundtrip() {
    let q = ChatJobQueue::new(1, 1);
    let flag = q.register_stream_cancel(42);
    assert!(q.request_stream_cancel(42));
    assert!(flag.load(std::sync::atomic::Ordering::SeqCst));
    q.unregister_stream_cancel(42);
    assert!(!q.request_stream_cancel(42));
}

// `final_response` 兜底回归覆盖：
// 1) 缺失时应补发（fallback_emits_final_response_when_missing）
// 2) 已存在且稍后可见时不应误补（fallback_skips_when_final_response_arrives_with_small_delay）
// 3) 同一 job 二次触发时保持幂等（fallback_is_idempotent_for_same_job）
#[tokio::test]
async fn fallback_emits_final_response_when_missing() {
    let hub = SseStreamHub::new();
    let job_id = 42_u64;
    hub.register_job(job_id);
    let (tx, mut rx) = mpsc::channel::<String>(8);
    let messages = vec![Message::assistant_only("最终总结内容")];

    emit_missing_final_response_fallback_if_needed(&hub, &tx, job_id, &messages).await;

    // 跳过 reasoning_end / text_start，读取 timeline + answer_phase
    let _reasoning_end = timeout(Duration::from_millis(200), rx.recv())
        .await
        .expect("recv reasoning_end")
        .expect("reasoning_end payload");
    let _text_start = timeout(Duration::from_millis(200), rx.recv())
        .await
        .expect("recv text_start")
        .expect("text_start payload");
    let first = timeout(Duration::from_millis(200), rx.recv())
        .await
        .expect("recv final_response frame")
        .expect("final_response payload");
    assert!(sse_payload_has_final_response_timeline(&first));

    let second = timeout(Duration::from_millis(200), rx.recv())
        .await
        .expect("recv answer_phase frame")
        .expect("answer_phase payload");
    assert!(
        second.contains("\"customType\":\"assistant_answer_phase\""),
        "second frame should be answer_phase CUSTOM: {second}"
    );
}

#[tokio::test]
async fn fallback_skips_when_final_response_arrives_with_small_delay() {
    let hub = SseStreamHub::new();
    let job_id = 43_u64;
    hub.register_job(job_id);
    let (tx, mut rx) = mpsc::channel::<String>(8);
    let messages = vec![Message::assistant_only("最终总结内容")];

    let delayed_payload = crate::sse::encode_message(crate::sse::SsePayload::TimelineLog {
        log: crate::sse::protocol::TimelineLogBody {
            kind: "final_response".to_string(),
            title: "已存在总结".to_string(),
            detail: None,
        },
    });
    let hub_for_task = hub.clone();
    tokio::spawn(async move {
        tokio::time::sleep(Duration::from_millis(10)).await;
        let _ = hub_for_task.publish(job_id, delayed_payload);
    });

    emit_missing_final_response_fallback_if_needed(&hub, &tx, job_id, &messages).await;

    let no_frame = timeout(Duration::from_millis(120), rx.recv()).await;
    assert!(no_frame.is_err(), "已有 final_response 时不应再补发");
}

#[tokio::test]
async fn fallback_is_idempotent_for_same_job() {
    let hub = SseStreamHub::new();
    let job_id = 44_u64;
    hub.register_job(job_id);
    let (tx, mut rx) = mpsc::channel::<String>(8);
    let messages = vec![Message::assistant_only("最终总结内容")];

    emit_missing_final_response_fallback_if_needed(&hub, &tx, job_id, &messages).await;
    // 跳过 reasoning_end / text_start
    let _r = rx.recv().await;
    let _t = rx.recv().await;
    let first = timeout(Duration::from_millis(200), rx.recv())
        .await
        .expect("recv first frame (timeline)")
        .expect("first frame payload");
    let second = timeout(Duration::from_millis(200), rx.recv())
        .await
        .expect("recv second frame (answer_phase)")
        .expect("second frame payload");
    assert!(sse_payload_has_final_response_timeline(&first));
    assert!(
        second.contains("\"customType\":\"assistant_answer_phase\""),
        "second frame should be answer_phase CUSTOM: {second}"
    );

    // 首次 fallback 通过桥接后，hub 已可见 final_response；二次调用应无输出。
    hub.publish(job_id, first);
    // eventual 内部有 5×20ms 轮询，等待足够久确保 publish 被读到
    tokio::time::sleep(Duration::from_millis(200)).await;
    // 二次调用应直接返回 false，不向 tx 写入任何数据
    let emitted_again =
        emit_missing_final_response_fallback_if_needed(&hub, &tx, job_id, &messages).await;
    assert!(
        !emitted_again,
        "二次 fallback 应返回 false（hub 中已有 final_response）"
    );
}

#[tokio::test]
async fn fallback_skips_when_turn_has_no_new_assistant() {
    let hub = SseStreamHub::new();
    let job_id = 45_u64;
    hub.register_job(job_id);
    let (tx, mut rx) = mpsc::channel::<String>(8);
    let messages = vec![
        Message::system_only("sys"),
        Message::user_only("上一轮提问"),
        Message::assistant_only("上一轮回答"),
        Message::user_only("本轮提问"),
    ];

    emit_missing_final_response_fallback_if_needed(&hub, &tx, job_id, &messages).await;

    let no_frame = timeout(Duration::from_millis(120), rx.recv()).await;
    assert!(
        no_frame.is_err(),
        "本轮无 assistant 输出时不应复用上一轮回答做 fallback"
    );
}

#[tokio::test]
async fn conversation_saved_emits_before_stream_ended() {
    let (tx, mut rx) = mpsc::channel::<String>(8);
    emit_stream_draining(&tx, 99).await;
    let saved = crate::sse::encode_message(crate::sse::SsePayload::ConversationSaved {
        saved: crate::sse::ConversationSavedBody {
            revision: 7,
            tiktoken_prompt_tokens: None,
        },
    });
    let _ = crate::sse::send_string_logged(
        &tx,
        saved,
        "chat_job_queue::tests conversation_saved_before_end",
    )
    .await;
    let mut sent = false;
    emit_stream_ended_once(
        &tx,
        99,
        StreamEndReason::Completed,
        &mut sent,
        "chat_job_queue::tests stream_ended_last",
        None,
    )
    .await;

    let first = timeout(Duration::from_millis(200), rx.recv())
        .await
        .expect("recv first")
        .expect("first payload");
    let second = timeout(Duration::from_millis(200), rx.recv())
        .await
        .expect("recv second")
        .expect("second payload");
    let third = timeout(Duration::from_millis(200), rx.recv())
        .await
        .expect("recv third")
        .expect("third payload");

    assert!(
        first.contains("\"customType\":\"stream_draining\"")
            || first.contains("\"stream_draining\"")
    );
    assert!(
        second.contains("\"customType\":\"conversation_saved\"")
            || second.contains("\"conversation_saved\"")
    );
    assert!(third.contains("\"type\":\"RUN_FINISHED\"") || third.contains("\"stream_ended\""));
}