Skip to content

Commit 3b83949

Browse files
quangdang46claude
andcommitted
feat: implement YOLO classifier and MCP unified pipeline
Phase 3.4 YOLO Classifier (jcode-only): - Create src/yolo_classifier.rs with 2-stage LLM approach - Stage 1 (fast, 64 tokens): BLOCK/ALLOW decision - Stage 2 (thinking, 4096 tokens): CoT evidence gathering when blocked - Fail closed on errors/timeouts → falls back to interactive prompt - Circuit breaker: 3 consecutive denials → bypass YOLO until reset - Integrated into dcg_bridge::classify_with_mode() for Mode::Auto Phase 3.6 MCP Unified Pipeline: - Extend action_to_tool_call() to handle mcp__* actions - MCP tools mapped to ToolCall::read with [Read, Write, Spawn] effects - Three matching levels: mcp__server, mcp__server__tool, mcp__server__* Also updated src/lib.rs to export the yolo_classifier module. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent 5a062eb commit 3b83949

3 files changed

Lines changed: 436 additions & 5 deletions

File tree

src/dcg_bridge.rs

Lines changed: 47 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ use std::sync::{LazyLock, Mutex};
3737

3838
use dcg_core::{Decision, Effect, Engine, EngineConfig, Mode, Session, ToolCall};
3939

40+
pub use crate::yolo_classifier::YoloClassifier;
41+
4042
/// Globally configured permission mode. Set once during CLI startup, read
4143
/// from every `SafetySystem::classify` call.
4244
///
@@ -141,11 +143,33 @@ pub fn classify_with_mode(action: &str, mode: Mode) -> BridgeDecision {
141143
// `Engine::evaluate` because their pre-check semantics are well
142144
// defined without rule data.
143145
if matches!(mode, Mode::Default | Mode::Auto) {
144-
return if is_legacy_auto_allowed(&lower) {
145-
BridgeDecision::Allow
146-
} else {
147-
BridgeDecision::Prompt
148-
};
146+
// Legacy auto-allowed tools always allow in both Default and Auto.
147+
if is_legacy_auto_allowed(&lower) {
148+
return BridgeDecision::Allow;
149+
}
150+
151+
// For Mode::Auto, non-legacy tools go through YOLO classifier.
152+
if mode == Mode::Auto {
153+
let (tool, effects) = action_to_tool_call(&lower);
154+
let effect_strings: Vec<&str> = effects
155+
.iter()
156+
.map(|e| match e {
157+
Effect::Read => "Read",
158+
Effect::Write => "Write",
159+
Effect::Spawn => "Spawn",
160+
Effect::Fs => "Fs",
161+
Effect::Irreversible => "Irreversible",
162+
Effect::Network => "Network",
163+
Effect::CredentialAccess => "CredentialAccess",
164+
Effect::PrivilegeEscalation => "PrivilegeEscalation",
165+
})
166+
.collect();
167+
168+
let classifier = YoloClassifier::get_or_init();
169+
return classifier.evaluate(&lower, &format!("{:?}", tool), &effect_strings);
170+
}
171+
172+
return BridgeDecision::Prompt;
149173
}
150174

151175
let (tool, effects) = action_to_tool_call(&lower);
@@ -242,6 +266,24 @@ fn action_to_tool_call(action_lower: &str) -> (ToolCall, Vec<Effect>) {
242266
return (ToolCall::bash(""), vec![Spawn, Write, Irreversible]);
243267
}
244268

269+
// MCP tool actions: mcp__serverName__toolName
270+
// Three matching levels:
271+
// mcp__github → matches ALL tools from github server
272+
// mcp__github__* → wildcard, same as above
273+
// mcp__github__create_pull_request → exact tool
274+
if action_lower.starts_with("mcp__") {
275+
let parts: Vec<&str> = action_lower.split("__").collect();
276+
if parts.len() >= 2 {
277+
// MCP tools carry Read + Write + Spawn effects since they can
278+
// read/write data and spawn background processes.
279+
// Path is unknown at classify time — use placeholder.
280+
return (
281+
ToolCall::read(placeholder),
282+
vec![Read, Write, Spawn],
283+
);
284+
}
285+
}
286+
245287
// Conservative default for unknown / future tools. We still surface a
246288
// ToolCall::Bash so the engine treats it as command-shaped rather
247289
// than file-shaped.

src/lib.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
clippy::useless_conversion
77
)]
88

9+
pub mod yolo_classifier;
910
pub mod agent;
1011
pub mod ambient;
1112
pub mod ambient_runner;

0 commit comments

Comments
 (0)