Skip to content

Commit dc00dbc

Browse files
committed
fix(runtime): wire hooks/cache/retry, session locking, MCP timeouts, cron injection
Build the hook executor from settings and enable prompt caching and a retry policy at runtime init instead of hardcoding them off. Keep an authoritative conversation snapshot so mid-query reads do not see an empty shell, and route permission asks through per-request oneshots. Own session grants in memory under a per-session file lock instead of re-reading the file. Bound MCP stdio/ws requests with a 30s timeout and drain pending requests when a server dies. Drain fired cron prompts on an idle tick and inject them as user turns.
1 parent 130a58b commit dc00dbc

18 files changed

Lines changed: 878 additions & 110 deletions

File tree

crates/agents/src/prompt/git_context.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -267,7 +267,6 @@ mod tests {
267267
.output()
268268
.unwrap();
269269

270-
// Create a commit
271270
fs::write(path.join("file.txt"), "content").unwrap();
272271
Command::new("git")
273272
.args(["add", "."])

crates/agents/src/runtime.rs

Lines changed: 468 additions & 37 deletions
Large diffs are not rendered by default.

crates/agents/src/summarizer.rs

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -453,11 +453,7 @@ fn extract_topic(msg: &Message) -> Option<String> {
453453

454454
/// Truncate a line to max length, adding "..." if truncated.
455455
fn truncate_line(s: &str, max: usize) -> String {
456-
if s.len() <= max {
457-
s.to_string()
458-
} else {
459-
format!("{}...", &s[..max.saturating_sub(3)])
460-
}
456+
crab_utils::text::truncate_chars_within(s, max, "...")
461457
}
462458

463459
#[cfg(test)]

crates/agents/src/teams/worker_pool.rs

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -356,18 +356,25 @@ impl WorkerPool {
356356

357357
/// Truncate a string to `max_len` chars, appending "…" if truncated.
358358
fn truncate_preview(s: &str, max_len: usize) -> String {
359-
if s.len() <= max_len {
360-
s.to_string()
361-
} else {
362-
format!("{}…", &s[..max_len])
363-
}
359+
crab_utils::text::truncate_chars(s, max_len, "…")
364360
}
365361

366362
#[cfg(test)]
367363
mod tests {
368364
use super::*;
369365
use crab_session::Conversation;
370366

367+
#[test]
368+
fn truncate_preview_multibyte_does_not_panic() {
369+
// The previous byte-slice implementation panicked when `max_len`
370+
// landed inside a multi-byte codepoint.
371+
let prompt = "实现一个跨平台的终端用户界面组件".repeat(10);
372+
let out = truncate_preview(&prompt, 80);
373+
assert_eq!(out.chars().count(), 81);
374+
assert!(out.ends_with('…'));
375+
assert_eq!(truncate_preview("短", 80), "短");
376+
}
377+
371378
#[test]
372379
fn coordinator_creation() {
373380
let coord = WorkerPool::new("main".into(), "Main Agent".into());

crates/agents/tests/integration_m7b.rs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ use crab_tools::builtin::glob::GLOB_TOOL_NAME;
2525
use crab_tools::builtin::grep::GREP_TOOL_NAME;
2626
use crab_tools::builtin::notebook::NOTEBOOK_EDIT_TOOL_NAME;
2727
use crab_tools::builtin::read::READ_TOOL_NAME;
28-
use crab_tools::builtin::remote_trigger::REMOTE_TRIGGER_TOOL_NAME;
2928
use crab_tools::builtin::task::{
3029
TASK_CREATE_TOOL_NAME, TASK_GET_TOOL_NAME, TASK_LIST_TOOL_NAME, TASK_OUTPUT_TOOL_NAME,
3130
TASK_STOP_TOOL_NAME, TASK_UPDATE_TOOL_NAME,
@@ -483,7 +482,6 @@ fn default_registry_includes_all_builtin_tools() {
483482
CRON_CREATE_TOOL_NAME,
484483
CRON_DELETE_TOOL_NAME,
485484
CRON_LIST_TOOL_NAME,
486-
REMOTE_TRIGGER_TOOL_NAME,
487485
];
488486

489487
for name in &expected {

crates/cli/src/agent.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -468,6 +468,9 @@ pub async fn run(cli: &Cli, resume_session_id: Option<String>) -> anyhow::Result
468468
backend,
469469
skill_dirs,
470470
mcp_servers: settings.mcp_servers.clone(),
471+
hooks: settings.hooks.clone(),
472+
disable_hooks: cli.bare || settings.disable_all_hooks.unwrap_or(false),
473+
cache_enabled: true,
471474
settings_warnings,
472475
prefill: cli.prefill.clone(),
473476
open_resume_picker,

crates/mcp/src/elicitation.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -196,8 +196,12 @@ mod tests {
196196
request_id: None,
197197
is_confirmation: false,
198198
};
199+
// `approve(Null)` still carries `Some(data)`, so it passes the
200+
// presence check but fails schema validation against an object type.
199201
let resp = ElicitationResponse::approve(serde_json::Value::Null);
200-
// Data is Null which is still Some; let's test with confirm
202+
let err = validate_response(&req, &resp).unwrap_err();
203+
assert!(err.contains("does not conform"));
204+
201205
let resp2 = ElicitationResponse {
202206
approved: true,
203207
data: None,

crates/mcp/src/transport/stdio.rs

Lines changed: 74 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use std::collections::HashMap;
22
use std::future::Future;
33
use std::pin::Pin;
44
use std::sync::Arc;
5+
use std::time::Duration;
56

67
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
78
use tokio::process::{Child, Command};
@@ -10,6 +11,10 @@ use tokio::sync::{Mutex, oneshot};
1011
use crate::protocol::{JsonRpcNotification, JsonRpcRequest, JsonRpcResponse};
1112
use crate::transport::Transport;
1213

14+
/// Per-request timeout — a request that gets no response within this window
15+
/// fails rather than hanging forever (e.g. the server died mid-request).
16+
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
17+
1318
/// Stdin/stdout transport for MCP servers launched as child processes.
1419
///
1520
/// The MCP stdio protocol frames messages as `Content-Length: N\r\n\r\n{json}`
@@ -72,18 +77,23 @@ impl StdioTransport {
7277
loop {
7378
match read_message(&mut reader).await {
7479
Ok(Some(data)) => {
75-
// Try to parse as a response (has "id" field)
80+
// A response carries an "id"; a notification does not.
7681
if let Ok(resp) = serde_json::from_str::<JsonRpcResponse>(&data) {
7782
let mut map = pending_clone.lock().await;
7883
if let Some(tx) = map.remove(&resp.id) {
7984
let _ = tx.send(resp);
8085
}
86+
} else if let Ok(notif) = serde_json::from_str::<JsonRpcNotification>(&data)
87+
{
88+
// No consumer surface for stdio server notifications
89+
// yet; log so they're observable but not lost silently.
90+
tracing::debug!(
91+
method = notif.method,
92+
"received MCP server notification (no consumer wired)"
93+
);
8194
}
82-
// Notifications from server are silently dropped for now.
83-
// TODO: emit server notifications through a channel.
8495
}
8596
Ok(None) => {
86-
// EOF — server process closed stdout
8797
tracing::debug!("MCP server stdout closed");
8898
break;
8999
}
@@ -93,6 +103,10 @@ impl StdioTransport {
93103
}
94104
}
95105
}
106+
// Reader is exiting (EOF or error): drop every pending sender so
107+
// in-flight `rx.await` callers get an error promptly instead of
108+
// hanging forever waiting on a dead server.
109+
pending_clone.lock().await.clear();
96110
});
97111

98112
Ok(Self {
@@ -142,10 +156,20 @@ impl Transport for StdioTransport {
142156
tracing::debug!(method = %req.method, id, "sending MCP request");
143157
self.write_message(&json).await?;
144158

145-
// Wait for the response from the reader task.
146-
rx.await.map_err(|_| {
147-
crab_core::Error::Other("MCP server closed connection before responding".into())
148-
})
159+
// Wait for the response, bounded by REQUEST_TIMEOUT so a dead or
160+
// wedged server can't hang the caller indefinitely.
161+
match tokio::time::timeout(REQUEST_TIMEOUT, rx).await {
162+
Ok(Ok(resp)) => Ok(resp),
163+
Ok(Err(_)) => Err(crab_core::Error::Other(
164+
"MCP server closed connection before responding".into(),
165+
)),
166+
Err(_) => {
167+
self.pending.lock().await.remove(&id);
168+
Err(crab_core::Error::Other(format!(
169+
"MCP server did not respond within {REQUEST_TIMEOUT:?}"
170+
)))
171+
}
172+
}
149173
})
150174
}
151175

@@ -258,4 +282,46 @@ mod tests {
258282
let msg = read_message(&mut reader).await.unwrap().unwrap();
259283
assert_eq!(msg, "{}");
260284
}
285+
286+
#[test]
287+
fn request_timeout_is_bounded() {
288+
assert_eq!(REQUEST_TIMEOUT, Duration::from_secs(30));
289+
}
290+
291+
/// A server that exits immediately (closing stdout) must cause an in-flight
292+
/// request to fail promptly rather than hang forever. The reader task hits
293+
/// EOF, drains `pending`, and the dropped sender resolves `rx.await` with an
294+
/// error well before the 30s request timeout.
295+
#[tokio::test]
296+
async fn send_errors_when_server_dies_before_responding() {
297+
let (command, args) = if cfg!(windows) {
298+
(
299+
"cmd".to_string(),
300+
vec!["/C".to_string(), "exit".to_string()],
301+
)
302+
} else {
303+
(
304+
"sh".to_string(),
305+
vec!["-c".to_string(), "exit 0".to_string()],
306+
)
307+
};
308+
309+
let transport = StdioTransport::spawn(&command, &args, None)
310+
.await
311+
.expect("spawn short-lived server");
312+
313+
let req = JsonRpcRequest::new("initialize", None);
314+
// Bound the whole call so a regression (a real hang) fails the test
315+
// instead of stalling the suite.
316+
let result = tokio::time::timeout(Duration::from_secs(5), transport.send(req)).await;
317+
318+
assert!(
319+
result.is_ok(),
320+
"send hung instead of erroring on dead server"
321+
);
322+
assert!(
323+
result.unwrap().is_err(),
324+
"send should error when the server dies before responding"
325+
);
326+
}
261327
}

crates/mcp/src/transport/ws.rs

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,10 @@ use crate::transport::Transport;
1919
/// Connection timeout for the initial WebSocket handshake.
2020
const CONNECT_TIMEOUT: Duration = Duration::from_secs(30);
2121

22+
/// Per-request timeout — a request that gets no response within this window
23+
/// fails rather than hanging forever (e.g. the server died mid-request).
24+
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
25+
2226
/// Receiver for server-initiated JSON-RPC notifications.
2327
///
2428
/// Returned by [`WsTransport::take_notifications`] exactly once per
@@ -230,12 +234,21 @@ impl Transport for WsTransport {
230234
tracing::debug!(method = %req.method, id, url = %self.url, "sending WS request");
231235
self.send_text(&json).await?;
232236

233-
// Wait for the response from the reader task.
234-
rx.await.map_err(|_| {
235-
crab_core::Error::Other(
237+
// Wait for the response, bounded by REQUEST_TIMEOUT so a dead or
238+
// wedged server can't hang the caller indefinitely.
239+
match tokio::time::timeout(REQUEST_TIMEOUT, rx).await {
240+
Ok(Ok(resp)) => Ok(resp),
241+
Ok(Err(_)) => Err(crab_core::Error::Other(
236242
"WebSocket connection closed before response received".into(),
237-
)
238-
})
243+
)),
244+
Err(_) => {
245+
self.pending.lock().await.remove(&id);
246+
Err(crab_core::Error::Other(format!(
247+
"MCP server at {} did not respond within {REQUEST_TIMEOUT:?}",
248+
self.url
249+
)))
250+
}
251+
}
239252
})
240253
}
241254

crates/session/src/context.rs

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,12 @@ impl Default for ContextManager {
2525

2626
impl ContextManager {
2727
/// Check usage and return the appropriate action.
28+
///
29+
/// Occupancy is taken from [`Conversation::effective_token_estimate`],
30+
/// which prefers the real `input_tokens` from the latest API response over
31+
/// the char-based heuristic.
2832
pub fn check(&self, conversation: &Conversation) -> ContextAction {
29-
let used = conversation.estimated_tokens();
33+
let used = conversation.effective_token_estimate();
3034
let limit = conversation.context_window;
3135
if limit == 0 {
3236
return ContextAction::Ok;
@@ -176,6 +180,58 @@ mod tests {
176180
);
177181
}
178182

183+
#[test]
184+
fn check_uses_real_input_tokens_over_char_estimate() {
185+
use crab_core::model::TokenUsage;
186+
let cm = ContextManager {
187+
warn_threshold_percent: 30,
188+
upgrade_threshold_percent: 50,
189+
compact_threshold_percent: 80,
190+
};
191+
// Short content → tiny char estimate, but the API reported a large real
192+
// prompt size (e.g. CJK that chars/4 underestimates). The compaction
193+
// decision must honor the real token count.
194+
let mut conv = Conversation::new("s".into(), String::new(), 1_000);
195+
conv.push_user("短");
196+
assert!(
197+
matches!(cm.check(&conv), ContextAction::Ok),
198+
"char estimate alone is below all thresholds"
199+
);
200+
conv.record_usage(TokenUsage {
201+
input_tokens: 850,
202+
..TokenUsage::default()
203+
});
204+
assert!(
205+
matches!(cm.check(&conv), ContextAction::NeedsCompaction { .. }),
206+
"real input_tokens (850/1000) must trigger compaction"
207+
);
208+
}
209+
210+
#[test]
211+
fn check_adds_post_usage_messages_to_real_tokens() {
212+
use crab_core::model::TokenUsage;
213+
let cm = ContextManager {
214+
warn_threshold_percent: 30,
215+
upgrade_threshold_percent: 50,
216+
compact_threshold_percent: 80,
217+
};
218+
let mut conv = Conversation::new("s".into(), String::new(), 1_000);
219+
conv.push_user("first");
220+
conv.record_usage(TokenUsage {
221+
input_tokens: 400,
222+
..TokenUsage::default()
223+
});
224+
// 400 real tokens = 40% → Warning band only.
225+
assert!(matches!(cm.check(&conv), ContextAction::Warning { .. }));
226+
// Append a large message (~500 tokens) not yet seen by any API call.
227+
conv.push_user("x".repeat(2000));
228+
// 400 + ~500 = ~900 → above compact threshold.
229+
assert!(
230+
matches!(cm.check(&conv), ContextAction::NeedsCompaction { .. }),
231+
"appended messages must be folded into the real-token estimate"
232+
);
233+
}
234+
179235
#[test]
180236
fn context_action_equality() {
181237
assert_eq!(ContextAction::Ok, ContextAction::Ok);

0 commit comments

Comments
 (0)