Skip to content

Commit 973dae6

Browse files
committed
fix: apply rustfmt formatting and remove unused import
1 parent 45f10c5 commit 973dae6

3 files changed

Lines changed: 90 additions & 37 deletions

File tree

src/cortex-engine/src/tools/mod.rs

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,8 @@ pub use handlers::*;
4848
pub use registry::{PluginTool, ToolRegistry};
4949
pub use response_store::{
5050
CLEANUP_INTERVAL, DEFAULT_TTL, MAX_STORE_SIZE, StoreInfo, StoreStats, StoredResponse,
51-
ToolResponseStore, ToolResponseStoreConfig, create_shared_store, create_shared_store_with_config,
51+
ToolResponseStore, ToolResponseStoreConfig, create_shared_store,
52+
create_shared_store_with_config,
5253
};
5354
pub use router::ToolRouter;
5455
pub use spec::{ToolCall, ToolDefinition, ToolHandler, ToolResult};

src/cortex-engine/src/tools/response_store.rs

Lines changed: 50 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@ use std::sync::Arc;
1212
use std::time::{Duration, Instant};
1313

1414
use tokio::sync::RwLock;
15-
use tracing::{debug, warn};
15+
use tracing::debug;
1616

1717
use crate::tools::spec::ToolResult;
1818

@@ -145,7 +145,12 @@ impl ToolResponseStore {
145145
///
146146
/// If the store is at capacity, the oldest entry will be evicted.
147147
/// Returns `true` if an entry was evicted to make room.
148-
pub async fn store(&self, call_id: impl Into<String>, tool_name: impl Into<String>, result: ToolResult) -> bool {
148+
pub async fn store(
149+
&self,
150+
call_id: impl Into<String>,
151+
tool_name: impl Into<String>,
152+
result: ToolResult,
153+
) -> bool {
149154
let call_id = call_id.into();
150155
let tool_name = tool_name.into();
151156
let mut evicted = false;
@@ -185,7 +190,7 @@ impl ToolResponseStore {
185190
/// Marks the response as read but does not consume it.
186191
pub async fn get(&self, call_id: &str) -> Option<ToolResult> {
187192
let mut responses = self.responses.write().await;
188-
193+
189194
if let Some(response) = responses.get_mut(call_id) {
190195
response.read = true;
191196
let mut stats = self.stats.write().await;
@@ -202,7 +207,7 @@ impl ToolResponseStore {
202207
/// entries are cleaned up after being consumed (#5293).
203208
pub async fn take(&self, call_id: &str) -> Option<ToolResult> {
204209
let mut responses = self.responses.write().await;
205-
210+
206211
if let Some(response) = responses.remove(call_id) {
207212
let mut stats = self.stats.write().await;
208213
stats.takes += 1;
@@ -234,16 +239,16 @@ impl ToolResponseStore {
234239
let mut responses = self.responses.write().await;
235240
let ttl = self.config.ttl;
236241
let before = responses.len();
237-
242+
238243
responses.retain(|_, v| !v.is_expired(ttl));
239-
244+
240245
let removed = before - responses.len();
241246
if removed > 0 {
242247
debug!(removed, "Cleaned up expired responses");
243248
let mut stats = self.stats.write().await;
244249
stats.expired_cleanups += removed as u64;
245250
}
246-
251+
247252
removed
248253
}
249254

@@ -253,14 +258,14 @@ impl ToolResponseStore {
253258
pub async fn cleanup_read(&self) -> usize {
254259
let mut responses = self.responses.write().await;
255260
let before = responses.len();
256-
261+
257262
responses.retain(|_, v| !v.read);
258-
263+
259264
let removed = before - responses.len();
260265
if removed > 0 {
261266
debug!(removed, "Cleaned up read-but-not-consumed responses");
262267
}
263-
268+
264269
removed
265270
}
266271

@@ -278,7 +283,7 @@ impl ToolResponseStore {
278283
pub async fn info(&self) -> StoreInfo {
279284
let responses = self.responses.read().await;
280285
let stats = self.stats.read().await;
281-
286+
282287
StoreInfo {
283288
current_size: responses.len(),
284289
max_size: self.config.max_size,
@@ -411,14 +416,22 @@ mod tests {
411416
let store = ToolResponseStore::with_config(config);
412417

413418
// Fill to capacity
414-
store.store("call-1", "Read", ToolResult::success("1")).await;
415-
store.store("call-2", "Read", ToolResult::success("2")).await;
416-
store.store("call-3", "Read", ToolResult::success("3")).await;
419+
store
420+
.store("call-1", "Read", ToolResult::success("1"))
421+
.await;
422+
store
423+
.store("call-2", "Read", ToolResult::success("2"))
424+
.await;
425+
store
426+
.store("call-3", "Read", ToolResult::success("3"))
427+
.await;
417428

418429
assert_eq!(store.len().await, 3);
419430

420431
// Add one more, should evict oldest
421-
let evicted = store.store("call-4", "Read", ToolResult::success("4")).await;
432+
let evicted = store
433+
.store("call-4", "Read", ToolResult::success("4"))
434+
.await;
422435
assert!(evicted);
423436
assert_eq!(store.len().await, 3);
424437

@@ -429,11 +442,12 @@ mod tests {
429442

430443
#[tokio::test]
431444
async fn test_expired_cleanup() {
432-
let config = ToolResponseStoreConfig::default()
433-
.with_ttl(Duration::from_millis(50));
445+
let config = ToolResponseStoreConfig::default().with_ttl(Duration::from_millis(50));
434446
let store = ToolResponseStore::with_config(config);
435447

436-
store.store("call-1", "Read", ToolResult::success("1")).await;
448+
store
449+
.store("call-1", "Read", ToolResult::success("1"))
450+
.await;
437451
assert_eq!(store.len().await, 1);
438452

439453
// Wait for expiration
@@ -448,8 +462,12 @@ mod tests {
448462
async fn test_cleanup_read() {
449463
let store = ToolResponseStore::new();
450464

451-
store.store("call-1", "Read", ToolResult::success("1")).await;
452-
store.store("call-2", "Read", ToolResult::success("2")).await;
465+
store
466+
.store("call-1", "Read", ToolResult::success("1"))
467+
.await;
468+
store
469+
.store("call-2", "Read", ToolResult::success("2"))
470+
.await;
453471

454472
// Read one entry
455473
store.get("call-1").await;
@@ -466,7 +484,9 @@ mod tests {
466484
async fn test_stats() {
467485
let store = ToolResponseStore::new();
468486

469-
store.store("call-1", "Read", ToolResult::success("1")).await;
487+
store
488+
.store("call-1", "Read", ToolResult::success("1"))
489+
.await;
470490
store.get("call-1").await;
471491
store.take("call-1").await;
472492

@@ -489,8 +509,12 @@ mod tests {
489509
async fn test_clear() {
490510
let store = ToolResponseStore::new();
491511

492-
store.store("call-1", "Read", ToolResult::success("1")).await;
493-
store.store("call-2", "Read", ToolResult::success("2")).await;
512+
store
513+
.store("call-1", "Read", ToolResult::success("1"))
514+
.await;
515+
store
516+
.store("call-2", "Read", ToolResult::success("2"))
517+
.await;
494518

495519
assert_eq!(store.len().await, 2);
496520

@@ -502,7 +526,9 @@ mod tests {
502526
async fn test_info() {
503527
let store = ToolResponseStore::new();
504528

505-
store.store("call-1", "Read", ToolResult::success("1")).await;
529+
store
530+
.store("call-1", "Read", ToolResult::success("1"))
531+
.await;
506532

507533
let info = store.info().await;
508534
assert_eq!(info.current_size, 1);

src/cortex-engine/src/validation.rs

Lines changed: 38 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -280,9 +280,8 @@ fn normalize_command(cmd: &str) -> String {
280280
.enumerate()
281281
.map(|(idx, part)| {
282282
// Remove surrounding quotes (single and double)
283-
let unquoted = part
284-
.trim_matches(|c| c == '\'' || c == '"');
285-
283+
let unquoted = part.trim_matches(|c| c == '\'' || c == '"');
284+
286285
// For the first part (command), extract basename to handle path variants
287286
if idx == 0 {
288287
Path::new(unquoted)
@@ -738,10 +737,16 @@ mod tests {
738737

739738
// Extra whitespace should not bypass validation
740739
let result = validator.validate("rm -rf /");
741-
assert!(!result.valid, "Extra whitespace should not bypass blocked command");
740+
assert!(
741+
!result.valid,
742+
"Extra whitespace should not bypass blocked command"
743+
);
742744

743745
let result = validator.validate("rm -rf /");
744-
assert!(!result.valid, "Multiple spaces should not bypass blocked command");
746+
assert!(
747+
!result.valid,
748+
"Multiple spaces should not bypass blocked command"
749+
);
745750
}
746751

747752
#[test]
@@ -750,13 +755,22 @@ mod tests {
750755

751756
// Quoted commands should not bypass validation
752757
let result = validator.validate("'rm' -rf /");
753-
assert!(!result.valid, "Single quotes should not bypass blocked command");
758+
assert!(
759+
!result.valid,
760+
"Single quotes should not bypass blocked command"
761+
);
754762

755763
let result = validator.validate("\"rm\" -rf /");
756-
assert!(!result.valid, "Double quotes should not bypass blocked command");
764+
assert!(
765+
!result.valid,
766+
"Double quotes should not bypass blocked command"
767+
);
757768

758769
let result = validator.validate("'rm' '-rf' '/'");
759-
assert!(!result.valid, "Fully quoted command should not bypass blocked command");
770+
assert!(
771+
!result.valid,
772+
"Fully quoted command should not bypass blocked command"
773+
);
760774
}
761775

762776
#[test]
@@ -765,13 +779,19 @@ mod tests {
765779

766780
// Path variants should not bypass validation
767781
let result = validator.validate("/bin/rm -rf /");
768-
assert!(!result.valid, "Absolute path should not bypass blocked command");
782+
assert!(
783+
!result.valid,
784+
"Absolute path should not bypass blocked command"
785+
);
769786

770787
let result = validator.validate("/usr/bin/rm -rf /");
771788
assert!(!result.valid, "Full path should not bypass blocked command");
772789

773790
let result = validator.validate("./rm -rf /");
774-
assert!(!result.valid, "Relative path should not bypass blocked command");
791+
assert!(
792+
!result.valid,
793+
"Relative path should not bypass blocked command"
794+
);
775795
}
776796

777797
#[test]
@@ -780,10 +800,16 @@ mod tests {
780800

781801
// Combined bypass attempts
782802
let result = validator.validate("'/bin/rm' -rf /");
783-
assert!(!result.valid, "Combined path and whitespace should not bypass");
803+
assert!(
804+
!result.valid,
805+
"Combined path and whitespace should not bypass"
806+
);
784807

785808
let result = validator.validate("\"/usr/bin/rm\" '-rf' '/'");
786-
assert!(!result.valid, "Combined quotes, path, and whitespace should not bypass");
809+
assert!(
810+
!result.valid,
811+
"Combined quotes, path, and whitespace should not bypass"
812+
);
787813
}
788814

789815
#[test]

0 commit comments

Comments
 (0)