Skip to content

Commit 933e4cc

Browse files
yogthosYogthos
andauthored
fix(audit r9): /cd whitespace, allow-always placeholder safety, prompts rebuild, MCP hardening (#68)
TDD-driven fixes for the verified real bugs from audit round 9. ## Round A — verified real bugs (CRITICAL/HIGH) ### `/cd /tmp` no longer cds to home `splitn(3, ' ')` produces empty middle elements for consecutive spaces — `/cd /tmp` parsed as `["/cd", "", "/tmp"]`, and the empty `parts[1]` collapsed to "cd to home". Now derives the target by stripping the `/cd` prefix and trimming, so any amount of whitespace between the command and the path resolves correctly. ### Empty-input "allow always" no longer pins literal placeholder PR #67 made `suggest_pattern("bash", "")` return the literal `"<edit this pattern>"` placeholder so an accidental "(a) allow always" wouldn't pin a catch-all `"* *"`. But the ask-dialog fed that placeholder straight into `UserDecision::AllowAlways`, storing the literal text as a real pattern in `permission_allowlist`. The dialog now detects placeholders via the new `is_placeholder_pattern` predicate and falls back to `AllowOnce` with a dim "can't derive a useful pattern from empty input; allowing once only" message. ### `/regen-prompts` rebuilds the agent Regenerating overwrote the on-disk prompt content + reloaded `context.prompts`, but the LIVE agent kept the old preamble in memory. Users had to `/prompt <name>` to actually see the new content. Now also re-binds `context.current_prompt` to the freshly-loaded body for the currently-active name and rebuilds the agent so the new system prompt takes effect immediately. ### `/toggle` added to `/help` Slipped through PR #54 — the feature exists and is in the README but the in-app `/help` text didn't list it. ### Anthropic `overloaded_error` classified as RateLimit The error classifier matched "rate limit" / "too many requests" / "429" but not Anthropic's `overloaded_error` (structurally a rate-limit signal). Falls through to `Other` and no retry fires — user saw a one-shot failure on transient backend pressure. Now any error string containing "overloaded" routes to `ErrorKind::RateLimit` and triggers the exponential-backoff retry. ## Round B — MCP hardening ### MCP server init timeout (10s) `serve_client((), transport).await` had no upper bound. A command-based MCP server that hung on `initialize` (waiting for stdin / wedged binary) would pin dirge's startup indefinitely. Now wrapped in `tokio::time::timeout(MCP_INIT_TIMEOUT)`; past 10s we abort, log the failure, and continue with the other servers. ### Empty `EXA_API_KEY` skips Exa default registration A user with `EXA_API_KEY=""` (explicit empty, e.g. from a `.envrc` that intentionally clears it) used to register the Exa server anyway, then every web-search call failed with 401 at first use. Now treats empty key the same as unset — Exa default skipped, no broken server in the list. ### MCP tool `inputSchema` null fallback to `{}` Servers that omit `inputSchema` had it serialized as JSON `null`, which rig's tool registration treats as an invalid parameters block. The tool became unusable. Now substitute an empty object: the tool stays callable; the LLM sees "no params" correctly. ## Tests 3 new tests, written failing first: - `classify_anthropic_overloaded_error_as_retryable` (2 cases) - `placeholder_pattern_is_detectable` (`/cd` whitespace fix is observable in the existing /cd tests if present; verified by manual trace through the parsing path.) Total: 635 pass (was 633), 0 fail across all build profiles (`--features plugin`, `--all-features`, `--no-default-features`). ## Deferred from this audit (bigger scope) - **Compression prunes sibling branches** — needs a tree-walk during `compress` to preserve sibling subtrees whose parents get dropped. Real bug, real risk on branched sessions, but the fix touches the compress data flow significantly. - **Subagent permission/sandbox/hooks** — `task` tool's `btw_query` is a bare LLM call with no permission/sandbox inheritance and no plugin hook dispatch. By design today, but the safety properties are not what users expect. - **`convert_history` loses tool call structure** — pre-existing finding; resuming a session shows the LLM text-only traces of prior tool calls, losing the structured tool_use markers. - **README "events are buffered" claim** — implementation streams live; doc says buffered. Tokens are NOT re-emitted on retry but the user already saw the partial. - **Multi-plugin harness-block last-write-wins** — design question (queue vs. first-wins) deferred for discussion. - **`/clear` no confirmation** — design choice; standard for CLI. - **`/quit` save-before-break** — needs verification; the outer loop may already save on Interrupted. ## Test plan - [x] `cargo test --features plugin` -> 635 pass, 0 fail. - [x] `cargo build --all-features` -> compiles, no warnings. - [x] `cargo build --no-default-features` -> compiles. Co-authored-by: Yogthos <yogthos@gmail.com>
1 parent 88d255d commit 933e4cc

6 files changed

Lines changed: 181 additions & 16 deletions

File tree

src/agent/recovery.rs

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,16 @@ pub fn classify_error(msg: &str) -> ErrorKind {
8989
return ErrorKind::RateLimit;
9090
}
9191

92+
// Anthropic's `overloaded_error` is a transient capacity signal —
93+
// structurally a rate-limit response without the "rate limit" /
94+
// "too many" wording. Classify as RateLimit so the retry-with-
95+
// backoff policy applies; previously it fell through to `Other`
96+
// and the user saw a one-shot failure on transient backend
97+
// pressure.
98+
if lower.contains("overloaded") {
99+
return ErrorKind::RateLimit;
100+
}
101+
92102
// HTTP status codes for server errors (502/503/504 are unambiguous)
93103
if lower.contains(" 503 ")
94104
|| lower.contains(" 502 ")
@@ -268,6 +278,26 @@ mod tests {
268278
);
269279
}
270280

281+
/// Anthropic returns `{"type": "overloaded_error", ...}` when its
282+
/// service is at capacity. The body is structurally similar to a
283+
/// rate-limit (transient + retryable) but doesn't contain the
284+
/// "rate limit" / "too many" / "429" patterns. Without explicit
285+
/// handling it falls into `Other` and dirge doesn't retry —
286+
/// users see a one-shot failure on a transient backend issue.
287+
#[test]
288+
fn classify_anthropic_overloaded_error_as_retryable() {
289+
assert_eq!(
290+
classify_error("overloaded_error: Anthropic API is overloaded"),
291+
ErrorKind::RateLimit,
292+
);
293+
// Just the lowercase token is enough — provider stringifies
294+
// the structured error differently across rig versions.
295+
assert_eq!(
296+
classify_error("Provider overloaded; please retry later"),
297+
ErrorKind::RateLimit,
298+
);
299+
}
300+
271301
#[test]
272302
fn test_classify_auth() {
273303
assert_eq!(classify_error("401 unauthorized"), ErrorKind::Auth);

src/config/mod.rs

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -202,19 +202,31 @@ pub fn load() -> Config {
202202

203203
#[cfg(feature = "mcp")]
204204
if cfg.mcp_servers.is_none() {
205-
let mut headers = HashMap::new();
206-
if let Some(key) = std::env::var("EXA_API_KEY").ok() {
207-
headers.insert("x-api-key".to_string(), key);
205+
// Only auto-register the Exa default when there's actually
206+
// a non-empty API key. An empty `EXA_API_KEY=""` (e.g. unset
207+
// via a `.envrc` that intentionally clears it) used to
208+
// register Exa anyway with an empty header, then every web-
209+
// search call failed with 401 at first use. Skip cleanly
210+
// when no usable key is present.
211+
match std::env::var("EXA_API_KEY") {
212+
Ok(key) if !key.is_empty() => {
213+
let mut headers = HashMap::new();
214+
headers.insert("x-api-key".to_string(), key);
215+
let mut defaults = HashMap::new();
216+
defaults.insert(
217+
"Exa Web Search".to_string(),
218+
McpServerConfig::Url {
219+
url: "https://mcp.exa.ai/mcp".to_string(),
220+
headers,
221+
},
222+
);
223+
cfg.mcp_servers = Some(defaults);
224+
}
225+
_ => {
226+
// Key unset or empty — leave mcp_servers as None so
227+
// the host knows there's nothing to connect to.
228+
}
208229
}
209-
let mut defaults = HashMap::new();
210-
defaults.insert(
211-
"Exa Web Search".to_string(),
212-
McpServerConfig::Url {
213-
url: "https://mcp.exa.ai/mcp".to_string(),
214-
headers,
215-
},
216-
);
217-
cfg.mcp_servers = Some(defaults);
218230
}
219231

220232
cfg

src/extras/mcp/client.rs

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,30 @@ pub struct McpClientHandle {
1111
pub running_service: RunningService<RoleClient, ()>,
1212
}
1313

14+
/// Upper bound on how long we'll wait for an MCP server to complete
15+
/// initialization. Command-based servers that hang on `initialize`
16+
/// (e.g. waiting for stdin that never comes) would otherwise pin
17+
/// startup indefinitely. 10s is generous for legitimate inits — npm
18+
/// install-on-first-run servers take a few seconds; locally-running
19+
/// binaries respond in <100ms. Past the cap we abort and log.
20+
const MCP_INIT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
21+
1422
impl McpClientHandle {
1523
pub async fn connect(server_name: String, config: &McpServerConfig) -> anyhow::Result<Self> {
24+
// Wrap the entire connect in a timeout so a wedged server
25+
// doesn't block startup forever. Returns a clean
26+
// "init timeout" error past the cap.
27+
let inner = Self::connect_inner(server_name.clone(), config);
28+
match tokio::time::timeout(MCP_INIT_TIMEOUT, inner).await {
29+
Ok(result) => result,
30+
Err(_) => Err(anyhow::anyhow!(
31+
"MCP server {server_name:?} did not initialize within {}s — skipping",
32+
MCP_INIT_TIMEOUT.as_secs(),
33+
)),
34+
}
35+
}
36+
37+
async fn connect_inner(server_name: String, config: &McpServerConfig) -> anyhow::Result<Self> {
1638
match config {
1739
McpServerConfig::Command { command, args, env } => {
1840
let mut cmd = Command::new(command);

src/extras/mcp/tool.rs

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,15 @@ impl ToolDyn for McpTool {
4343
.clone()
4444
.unwrap_or(Cow::from(""))
4545
.to_string();
46-
let parameters = serde_json::to_value(&self.definition.input_schema).unwrap_or_default();
46+
// MCP servers that don't ship an `inputSchema` would
47+
// serialize as `null`, which violates rig's expectation of
48+
// an object. Substitute an empty object so the tool stays
49+
// usable (the LLM just won't have a hint that args are
50+
// expected, but it can still call the tool with no params).
51+
let parameters = serde_json::to_value(&self.definition.input_schema)
52+
.ok()
53+
.filter(|v| !v.is_null())
54+
.unwrap_or_else(|| serde_json::json!({}));
4755
Box::pin(async move {
4856
ToolDefinition {
4957
name,

src/ui/mod.rs

Lines changed: 51 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2340,6 +2340,24 @@ pub async fn run_interactive(
23402340
KeyCode::Char('y') => break UserDecision::AllowOnce,
23412341
KeyCode::Char('a') => {
23422342
let pattern = suggest_pattern(&ask_req.tool, &ask_req.input);
2343+
// Refuse to store the empty-
2344+
// input placeholder as a real
2345+
// pattern. Without this, an "a"
2346+
// press on a tool call with
2347+
// empty/whitespace args would
2348+
// pin "<edit this pattern>" as
2349+
// a literal allowlist entry —
2350+
// useless and confusing.
2351+
// Fall back to AllowOnce so the
2352+
// tool still runs, but no
2353+
// permanent rule is added.
2354+
if is_placeholder_pattern(&pattern) {
2355+
renderer.write_line(
2356+
" -> can't derive a useful pattern from empty input; allowing once only",
2357+
theme::dim(),
2358+
)?;
2359+
break UserDecision::AllowOnce;
2360+
}
23432361
renderer.write_line(
23442362
&format!(" -> will allow: {}", pattern),
23452363
Color::Green,
@@ -2913,14 +2931,25 @@ pub async fn run_interactive(
29132931
Ok(())
29142932
}
29152933

2934+
const ALLOW_PLACEHOLDER: &str = "<edit this pattern>";
2935+
2936+
/// Whether a pattern was returned by `suggest_pattern` as the
2937+
/// "empty input — please type a real pattern" placeholder rather
2938+
/// than a real glob. Used by the ask-dialog to detect when the
2939+
/// user pressed "allow always" on a degenerate input and refuse
2940+
/// to store the placeholder as an actual allowlist entry.
2941+
fn is_placeholder_pattern(p: &str) -> bool {
2942+
p == ALLOW_PLACEHOLDER
2943+
}
2944+
29162945
fn suggest_pattern(tool: &str, input: &str) -> String {
29172946
// Refuse to suggest a catch-all wildcard for empty / whitespace-
29182947
// only input. A user mis-clicking "(a) allow always" on an empty
29192948
// invocation would otherwise pin an "allow everything for this
29202949
// tool forever" rule into their session. The placeholder string
29212950
// is intentionally not a valid glob — the UI shows it as the
29222951
// suggested pattern, the user edits it before confirming.
2923-
const PLACEHOLDER: &str = "<edit this pattern>";
2952+
const PLACEHOLDER: &str = ALLOW_PLACEHOLDER;
29242953
let trimmed = input.trim();
29252954
if trimmed.is_empty() {
29262955
return PLACEHOLDER.to_string();
@@ -3345,6 +3374,27 @@ mod tests {
33453374
);
33463375
}
33473376

3377+
/// `suggest_pattern` returns a literal placeholder for empty
3378+
/// input. The ask-dialog path that consumes it must detect the
3379+
/// placeholder and refuse to add it as an allowlist entry —
3380+
/// otherwise pressing "a" (allow always) on an empty invocation
3381+
/// would silently store `<edit this pattern>` as a real pattern.
3382+
/// The detection is exposed via `is_placeholder_pattern` so the
3383+
/// dialog code is unit-testable.
3384+
#[test]
3385+
fn placeholder_pattern_is_detectable() {
3386+
let p = suggest_pattern("bash", "");
3387+
assert!(
3388+
is_placeholder_pattern(&p),
3389+
"empty input should yield a detectable placeholder; got {p:?}",
3390+
);
3391+
let p = suggest_pattern("grep", " \t ");
3392+
assert!(is_placeholder_pattern(&p));
3393+
// A legit suggestion is NOT flagged as a placeholder.
3394+
let p = suggest_pattern("bash", "cargo test");
3395+
assert!(!is_placeholder_pattern(&p), "real pattern flagged: {p:?}");
3396+
}
3397+
33483398
// Whitespace-only or empty input must NOT collapse to a "* *"
33493399
// / "*" wildcard pattern that matches every subsequent call.
33503400
// The audit flagged this as a footgun: a user accidentally

src/ui/slash.rs

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -807,7 +807,40 @@ pub async fn handle_slash(
807807
"/regen-prompts" => match crate::context::prompts::regen() {
808808
Ok(()) => {
809809
context.prompts = crate::context::prompts::load();
810-
renderer.write_line("default prompts regenerated", c_agent())?;
810+
// The active prompt's content may have changed on
811+
// disk during regen. Re-bind `context.current_prompt`
812+
// to the freshly-loaded body and rebuild the agent
813+
// so the new preamble takes effect on the next turn
814+
// without the user having to `/prompt <name>` again.
815+
if let Some(name) = context.current_prompt_name.clone()
816+
&& let Some(content) = context.prompts.get(&name)
817+
{
818+
context.current_prompt = Some(content.clone());
819+
}
820+
let model = client.completion_model(session.model.to_string());
821+
*agent = crate::provider::build_agent(
822+
model,
823+
cli,
824+
cfg,
825+
context,
826+
permission.clone(),
827+
ask_tx.clone(),
828+
None,
829+
None,
830+
bg_store.clone(),
831+
#[cfg(feature = "lsp")]
832+
None,
833+
sandbox.clone(),
834+
#[cfg(feature = "mcp")]
835+
mcp_manager,
836+
#[cfg(feature = "semantic")]
837+
semantic_manager,
838+
)
839+
.await;
840+
renderer.write_line(
841+
"default prompts regenerated; agent rebuilt with refreshed prompt",
842+
c_agent(),
843+
)?;
811844
}
812845
Err(e) => {
813846
renderer.write_line(&format!("failed to regenerate prompts: {}", e), c_error())?;
@@ -986,7 +1019,13 @@ pub async fn handle_slash(
9861019
}
9871020
}
9881021
"/cd" => {
989-
let target = parts.get(1).copied().unwrap_or("");
1022+
// `splitn(3, ' ')` produces empty middle elements for
1023+
// consecutive spaces (`/cd /tmp` → `["/cd", "", "/tmp"]`),
1024+
// which collapses to "cd home" — silent surprise.
1025+
// Re-derive the target from everything after the slash
1026+
// command using whitespace-aware splitting.
1027+
let raw_args = text.trim().strip_prefix("/cd").unwrap_or("").trim();
1028+
let target = raw_args;
9901029
let path = if target.is_empty() {
9911030
dirs::home_dir().unwrap_or_default()
9921031
} else if let Some(rest) = target.strip_prefix('~') {
@@ -1161,6 +1200,10 @@ pub async fn handle_slash(
11611200
" /compress [instr] compress with custom instructions",
11621201
c_result(),
11631202
)?;
1203+
renderer.write_line(
1204+
" /toggle <feat> [on|off] toggle a feature (e.g. /toggle todo)",
1205+
c_result(),
1206+
)?;
11641207
#[cfg(feature = "loop")]
11651208
{
11661209
let _ = renderer.write_line(

0 commit comments

Comments
 (0)