feat: Add async caching with background task management

[snowmead]

Summary

Transforms the cache_crate tool into an async operation that returns immediately with a task ID, enabling non-blocking caching of large crates with real-time progress monitoring and cancellation support.

Motivation

Large crates can take several minutes to cache (download + generate docs + index). The synchronous implementation blocked the MCP server during this time, preventing other operations and providing no visibility into progress. This makes the experience poor for AI agents and users working with large crates.

Changes

New Modules

• task_manager.rs: Core task management system

  • TaskManager: Thread-safe manager using Arc<RwLock<HashMap>>
  • CachingTask: Task metadata with status, stage, timestamps, cancellation token
  • TaskStatus enum: Pending, InProgress, Completed, Failed, Cancelled
  • CachingStage enum: Downloading, GeneratingDocs, Indexing, Completed

• task_formatter.rs: Rich markdown formatting for LLM consumption

  • AI-agent-optimized output with clear visual hierarchy
  • Embedded action commands for easy copy-paste
  • Grouped task listings by status
  • Human-readable durations and progress info

Modified Components

• cache_crate tool: Now spawns background tokio task and returns immediately
• CacheTools: Integrates TaskManager, new cache_operations method
• service.rs: New cache_operations tool with rich usage examples

Dependencies

• Added tokio-util 0.7 for CancellationToken
• Already had uuid for task IDs

API Design

Starting a Cache Operation

cache_crate({crate_name: "tokio", source_type: "cratesio", version: "1.35.0"})

Returns markdown with task ID and monitoring commands.

Monitoring Operations

// List all tasks
cache_operations({})

// Filter by status
cache_operations({status_filter: "in_progress"})

// Check specific task
cache_operations({task_id: "abc-123-def"})

// Cancel task
cache_operations({task_id: "abc-123-def", cancel: true})

// Clear completed/failed
cache_operations({clear: true})

Key Features

✅ Always Async - Consistent behavior, all operations return task ID
✅ Memory Only - No disk persistence, tasks cleared on restart
✅ No Auto-Cleanup - Tasks remain until explicitly cleared or server restart
✅ Cancellation Support - Can abort long-running operations
✅ Thread-Safe - Uses RwLock for concurrent access
✅ Progress Tracking - Track stage and status for each task
✅ LLM-Optimized Output - Markdown with embedded commands and clear hierarchy

Benefits

1. Non-Blocking: Server remains responsive during large crate caching
2. Visibility: Real-time progress updates with stage and elapsed time
3. Control: Cancel operations that are taking too long
4. AI-Friendly: Rich markdown output optimized for LLM parsing
5. Scalable: Multiple concurrent caching operations supported

Testing

• ✅ Compiles without errors or warnings
• ✅ Type-safe with strong enum types
• ✅ Follows Rust idioms with proper documentation
• ✅ Thread-safe with Arc/RwLock patterns
• Integration testing recommended for full workflow validation

Example Output

Starting Cache Operation

# Caching Started

**Crate**: tokio-1.35.0
**Source**: cratesio
**Task ID**: `abc-123-def`

The caching operation is running in the background.

**To check status**: `cache_operations({task_id: "abc-123-def"})`
**To cancel**: `cache_operations({task_id: "abc-123-def", cancel: true})`

Task Listing

# Caching Operations

## Summary
- **Total Operations**: 3
- **In Progress**: 1
- **Completed**: 2

## In Progress (1)

### Task: `abc-123-def`
**Crate**: tokio-1.35.0
**Source**: cratesio
**Status**: IN PROGRESS
**Stage**: Generating documentation
**Elapsed**: 2m 34s

**Actions**:
- Cancel: `cache_operations({task_id: "abc-123-def", cancel: true})`

Future Enhancements

• Add progress percentage tracking
• Implement rate limiting for concurrent operations
• Add webhook/callback support for completion notifications
• Persist task history to disk (optional)

🤖 Generated with Claude Code

[github-actions]

Claude finished @snowmead's task —— View job

---

Code Review Complete ✅

Review Checklist

• Gather context and understand the changes
• Review new task_manager.rs module
• Review new task_formatter.rs module
• Review modifications to cache/tools.rs
• Review modifications to service.rs
• Check thread safety and concurrency handling
• Evaluate error handling and edge cases
• Assess performance considerations
• Review security aspects
• Check test coverage

Overall Assessment

This PR successfully transforms the synchronous cache_crate tool into an async operation with background task management. The implementation is well-structured, thread-safe, and provides excellent UX with rich markdown output optimized for LLM consumption.

✅ Strengths

1. Architecture & Design

• Clean separation of concerns with dedicated modules for task management and formatting
• Consistent use of the async/await pattern throughout
• Strong type safety with comprehensive enum types for status and stages
• Thread-safe implementation using Arc<RwLock<HashMap>>

2. User Experience

• Rich markdown output with embedded commands for easy copy-paste
• Clear visual hierarchy optimized for LLM parsing
• Comprehensive progress tracking with elapsed time and current stage
• Intuitive API with flexible filtering and management options

3. Code Quality

• Well-documented code with clear module and function documentation
• Idiomatic Rust patterns and proper use of standard library types
• Consistent error handling patterns
• Clean parameter validation in cache_crate method

🔍 Areas for Improvement

1. Performance Considerations

Issue: Using RwLock for the entire task HashMap could become a bottleneck under high concurrency.

In task_manager.rs:187:

tasks: Arc<RwLock<HashMap<TaskId, CachingTask>>>,

Recommendation: Consider using DashMap for better concurrent performance:

use dashmap::DashMap;
// ...
tasks: Arc<DashMap<TaskId, CachingTask>>,

This would allow concurrent reads/writes to different keys without blocking.

2. Error Handling & Edge Cases

Issue: The elapsed_secs() method could panic if system time goes backwards.

In task_manager.rs:141-147:

pub fn elapsed_secs(&self) -> u64 {
    let end_time = self.completed_at.unwrap_or_else(SystemTime::now);
    end_time
        .duration_since(self.started_at)
        .unwrap_or_default()  // Good: handles error case
        .as_secs()
}

This is actually handled correctly with unwrap_or_default(), but consider logging when this edge case occurs.

3. Resource Management

Issue: No limit on concurrent background tasks could lead to resource exhaustion.

Recommendation: Add a semaphore to limit concurrent operations:

use tokio::sync::Semaphore;

pub struct TaskManager {
    tasks: Arc<RwLock<HashMap<TaskId, CachingTask>>>,
    max_concurrent: Arc<Semaphore>,  // e.g., permit 10 concurrent operations
}

4. Task Lifecycle

Issue: Tasks remain in memory indefinitely until explicitly cleared.

Recommendation: Consider adding:

• Optional TTL for completed tasks
• Maximum task count with FIFO eviction
• Automatic cleanup of old completed tasks after a configurable duration

5. Progress Granularity

Issue: No percentage tracking within stages (mentioned in Future Enhancements).

In task_formatter.rs:222-234, the stage context is static text. Consider adding progress percentage:

pub struct CachingStage {
    stage_type: CachingStageType,
    progress_percent: Option<u8>,  // 0-100
}

🔒 Security Review

✅ No Critical Security Issues Found

1. Cancellation Tokens: Properly implemented using tokio_util::CancellationToken
2. Input Validation: Parameters are validated before task creation
3. Resource Cleanup: Cancellation properly handled with cleanup mentioned
4. No Sensitive Data Exposure: Task IDs use UUIDs, no sensitive info in output

Minor Security Consideration

Consider rate limiting task creation to prevent potential DoS through rapid task spawning:

// Add to TaskManager
last_task_times: Arc<RwLock<VecDeque<SystemTime>>>,

// Check before creating task
fn check_rate_limit(&self) -> bool {
    // Allow max 10 tasks per second
}

🧪 Test Coverage

⚠️ Missing Test Coverage

No unit or integration tests found for the new async task management functionality. This is a significant gap that should be addressed.

Recommended Test Cases:

1. Task Lifecycle Tests:

   • Task creation and state transitions
   • Concurrent task operations
   • Task cancellation during different stages

2. Error Scenarios:

   • Cancelling already completed tasks
   • Clearing in-progress tasks
   • Invalid task ID handling

3. Concurrency Tests:

   • Multiple simultaneous cache operations
   • Race conditions in status updates
   • Thread safety validation

4. Formatting Tests:

   • Markdown output validation
   • Duration formatting edge cases
   • Task list sorting and grouping

🐛 Potential Bugs

1. Race Condition in Task Spawning

In tools.rs:573-635, the task status is updated after spawning but outside the spawn block. There's a small window where the task exists but isn't marked as InProgress:

tokio::spawn(async move {
    // Update status to in progress
    task_manager
        .update_status(&task_id, TaskStatus::InProgress)  // This happens after spawn
        .await;

Fix: Consider setting status before spawning or make it atomic with task creation.

2. Version Handling for Local Sources

In tools.rs:544, local sources use "auto" as placeholder:

let version = params.version.clone().unwrap_or_else(|| "auto".to_string());

This might cause issues if the actual version differs. Consider fetching the real version from Cargo.toml before creating the task.

📚 Documentation Suggestions

1. Add examples in the tool descriptions showing the markdown output format
2. Document the task lifecycle state machine
3. Add notes about memory-only persistence and server restart implications
4. Include performance characteristics (expected timing for different crate sizes)

✨ Recommendations

1. High Priority:

   • Add comprehensive test coverage
   • Implement concurrent task limiting
   • Add basic rate limiting

2. Medium Priority:

   • Consider using DashMap for better concurrent performance
   • Add automatic cleanup for old completed tasks
   • Implement progress percentage tracking

3. Low Priority:

   • Add metrics/telemetry for monitoring
   • Support for task priority/queuing
   • Webhook notifications for task completion

Conclusion

This is a well-implemented feature that significantly improves the user experience for caching operations. The async design is sound, the code quality is high, and the markdown formatting for LLM consumption is excellent. The main areas for improvement are test coverage, resource management, and some performance optimizations. With the addition of tests and minor enhancements to resource limiting, this PR would be ready for production use.

Verdict: Approve with suggestions ✅

The code is production-ready in terms of functionality and safety, but would benefit from test coverage and the resource management improvements mentioned above.

[github-actions]

Claude finished @snowmead's task —— View job

---

Code Review Complete ✅

Review Checklist

• Gather context and understand the changes
• Review new task_manager.rs module
• Review new task_formatter.rs module
• Review modifications to cache/tools.rs
• Review modifications to service.rs
• Check thread safety and concurrency handling
• Evaluate error handling and edge cases
• Assess performance considerations
• Review security aspects
• Check test coverage

Overall Assessment

This PR successfully transforms the synchronous cache_crate tool into an async operation with background task management. The implementation is well-structured, thread-safe, and provides excellent UX with rich markdown output optimized for LLM consumption.

✅ Strengths

1. Architecture & Design

• Clean separation of concerns with dedicated modules for task management and formatting
• Consistent use of the async/await pattern throughout
• Strong type safety with comprehensive enum types for status and stages
• Thread-safe implementation using Arc<RwLock<HashMap>>

2. User Experience

• Rich markdown output with embedded commands for easy copy-paste
• Clear visual hierarchy optimized for LLM parsing
• Comprehensive progress tracking with elapsed time and current stage
• Intuitive API with flexible filtering and management options

3. Code Quality

• Well-documented code with clear module and function documentation
• Idiomatic Rust patterns and proper use of standard library types
• Consistent error handling patterns
• Clean parameter validation in cache_crate method

🔍 Areas for Improvement

1. Performance Considerations

Issue: Using RwLock for the entire task HashMap could become a bottleneck under high concurrency.

In task_manager.rs:187:

tasks: Arc<RwLock<HashMap<TaskId, CachingTask>>>,

Recommendation: Consider using DashMap for better concurrent performance:

use dashmap::DashMap;
// ...
tasks: Arc<DashMap<TaskId, CachingTask>>,

This would allow concurrent reads/writes to different keys without blocking.

2. Error Handling & Edge Cases

Good Practice Found: The elapsed_secs() method properly handles system time edge cases:

In task_manager.rs:141-147:

pub fn elapsed_secs(&self) -> u64 {
    let end_time = self.completed_at.unwrap_or_else(SystemTime::now);
    end_time
        .duration_since(self.started_at)
        .unwrap_or_default()  // Good: handles error case
        .as_secs()
}

This correctly uses unwrap_or_default() to handle the case where system time goes backwards.

3. Resource Management

Issue: No limit on concurrent background tasks could lead to resource exhaustion.

Recommendation: Add a semaphore to limit concurrent operations:

use tokio::sync::Semaphore;

pub struct TaskManager {
    tasks: Arc<RwLock<HashMap<TaskId, CachingTask>>>,
    max_concurrent: Arc<Semaphore>,  // e.g., permit 10 concurrent operations
}

4. Task Lifecycle

Issue: Tasks remain in memory indefinitely until explicitly cleared.

Recommendation: Consider adding:

• Optional TTL for completed tasks
• Maximum task count with FIFO eviction
• Automatic cleanup of old completed tasks after a configurable duration

5. Progress Granularity

Issue: No percentage tracking within stages (mentioned in Future Enhancements).

In task_formatter.rs:222-234, the stage context is static text. Consider adding progress percentage:

pub struct CachingStage {
    stage_type: CachingStageType,
    progress_percent: Option<u8>,  // 0-100
}

🔒 Security Review

✅ No Critical Security Issues Found

1. Cancellation Tokens: Properly implemented using tokio_util::CancellationToken
2. Input Validation: Parameters are validated before task creation
3. Resource Cleanup: Cancellation properly handled
4. No Sensitive Data Exposure: Task IDs use UUIDs, no sensitive info in output

Minor Security Consideration

Consider rate limiting task creation to prevent potential DoS through rapid task spawning:

// Add to TaskManager
last_task_times: Arc<RwLock<VecDeque<SystemTime>>>,

// Check before creating task
fn check_rate_limit(&self) -> bool {
    // Allow max 10 tasks per second
}

🧪 Test Coverage

✅ Good Test Adaptation

The integration tests have been properly updated to handle the new async behavior:

• Tests now wait for task completion using wait_for_task_completion helper
• Proper parsing of CacheTaskStartedOutput for task IDs
• Validation of task status through cache_operations polling

Positive: The tests correctly handle the async nature of caching operations and validate task lifecycle.

Recommended Additional Test Cases

1. Concurrent Task Tests:

   • Multiple simultaneous cache operations
   • Race condition testing for status updates

2. Cancellation Tests:

   • Cancel during different stages
   • Cancel already completed tasks

3. Resource Limit Tests:

   • Many concurrent tasks
   • Memory usage under load

🐛 Minor Issues

1. Race Condition Fixed

Good fix in tools.rs:652-655 - setting status to InProgress before returning prevents race conditions:

// Update status to InProgress before returning (fixes race condition bug #1)
self.task_manager
    .update_status(&task.task_id, TaskStatus::InProgress)
    .await;

2. Local Version Resolution

Excellent improvement in tools.rs:613-622 - resolving version synchronously before task creation ensures accurate metadata:

// Resolve version synchronously before creating task (fixes bug #2)
let (version, auto_detected) = match Self::resolve_local_version(
    &path,
    params.version.as_deref(),
) {
    Ok(result) => result,
    Err(error_msg) => {
        return format!("# Error\n\n{}", error_msg);
    }
};

📚 Documentation Quality

The code is well-documented with:

• Clear module-level documentation
• Comprehensive function docs
• Helpful comments for complex logic
• Rich markdown output with usage examples

✨ Recommendations

High Priority:

1. Add concurrent task limiting - Implement semaphore-based limiting
2. Add basic rate limiting - Prevent rapid task spawning DoS

Medium Priority:

1. Consider DashMap for better concurrent performance
2. Add automatic cleanup for old completed tasks
3. Implement progress percentage tracking within stages

Low Priority:

1. Add metrics/telemetry for monitoring
2. Support task priority/queuing
3. Webhook notifications for task completion

Conclusion

This is a well-implemented feature that significantly improves the user experience for caching operations. The async design is sound, the code quality is high, and the markdown formatting for LLM consumption is excellent. The main areas for improvement are resource management and some performance optimizations.

The code demonstrates good Rust practices:

• Proper use of Arc<RwLock> for shared state
• Clean async/await patterns
• Strong type safety with enums
• Good error handling

Verdict: Approve with suggestions ✅

The implementation is production-ready with the following strengths:

• Thread-safe task management
• Proper cancellation support
• Rich user feedback
• Good test coverage adaptation

The suggested improvements would enhance scalability and resource management but are not blockers for merging.

[github-actions]

Claude finished @snowmead's task —— View job

---

PR Review: Async Caching with Background Task Management

Review Checklist

• Gather context and understand the changes
• Review new task_manager.rs module
• Review new task_formatter.rs module
• Review modifications to cache/tools.rs
• Review modifications to service.rs
• Check thread safety and concurrency handling
• Evaluate error handling and edge cases
• Assess performance considerations
• Review security aspects
• Check test coverage

Overall Assessment

This PR successfully transforms the synchronous cache_crate tool into an async operation with background task management. The implementation shows significant improvements over the previous review, addressing most of the concerns raised earlier. The code is well-structured, thread-safe, and provides excellent UX with rich markdown output.

✅ Strengths

1. Excellent Performance Optimizations

• DashMap Implementation: Already using DashMap (line 6 of task_manager.rs) instead of RwLock<HashMap>, which provides excellent concurrent performance with lock-free reads and minimal contention
• Efficient task tracking: Each task ID maps to its own entry, allowing parallel operations without blocking

2. Robust Error Handling

• System time edge cases: Properly handled with unwrap_or_default() in elapsed_secs() (lines 159-162)
• Version resolution: Smart synchronous resolution before task creation (lines 614-622 in tools.rs)
• Comprehensive validation: Thorough parameter validation with clear error messages

3. Fixed Race Conditions

• Status update ordering: Correctly sets InProgress status before returning (lines 652-655)
• Atomic operations: DashMap ensures thread-safe concurrent access

4. Excellent Documentation & UX

• Rich markdown output with embedded commands for easy copy-paste
• Clear visual hierarchy with status indicators (✓, ✗)
• Step-based progress tracking with descriptions
• LLM-optimized formatting for AI agent consumption

5. Good Test Adaptation

• Integration tests properly handle async behavior
• wait_for_task_completion helper polls correctly
• Tests validate task lifecycle and status transitions
• Proper timeout handling for different operation sizes

🔍 Areas for Improvement

1. Resource Management

While the implementation is solid, consider adding resource limits to prevent exhaustion:

// In TaskManager
use tokio::sync::Semaphore;

pub struct TaskManager {
    tasks: Arc<DashMap<TaskId, CachingTask>>,
    max_concurrent: Arc<Semaphore>, // Limit to e.g., 10 concurrent operations
}

// Before spawning task in cache_crate
let permit = self.task_manager.acquire_permit().await?;
tokio::spawn(async move {
    let _permit = permit; // Dropped when task completes
    // ... existing task logic
});

2. Task Lifecycle Management

Tasks remain in memory indefinitely. Consider:

• Optional TTL for completed tasks
• Maximum task count with FIFO eviction
• Configurable auto-cleanup after duration

impl TaskManager {
    pub async fn cleanup_old_tasks(&self, max_age: Duration) {
        let now = SystemTime::now();
        self.tasks.retain(|_, task| {
            task.is_terminal() && 
            task.completed_at
                .and_then(|t| now.duration_since(t).ok())
                .map(|d| d < max_age)
                .unwrap_or(true)
        });
    }
}

3. Progress Granularity Enhancement

Current implementation has step-based tracking. Consider adding percentage within steps:

pub struct CachingTask {
    // ... existing fields
    pub step_progress_percent: Option<u8>, // 0-100 within current step
}

4. Rate Limiting

Add protection against rapid task creation:

use std::collections::VecDeque;

pub struct TaskManager {
    // ... existing fields
    recent_tasks: Arc<RwLock<VecDeque<SystemTime>>>,
    max_tasks_per_second: usize,
}

impl TaskManager {
    pub async fn check_rate_limit(&self) -> Result<(), String> {
        let mut recent = self.recent_tasks.write().await;
        let now = SystemTime::now();
        
        // Remove entries older than 1 second
        while let Some(front) = recent.front() {
            if now.duration_since(*front).unwrap_or_default() > Duration::from_secs(1) {
                recent.pop_front();
            } else {
                break;
            }
        }
        
        if recent.len() >= self.max_tasks_per_second {
            return Err("Rate limit exceeded".to_string());
        }
        
        recent.push_back(now);
        Ok(())
    }
}

🔒 Security Review

✅ No Security Issues Found

• Proper input validation before task creation
• Safe path expansion with error handling
• No sensitive data in task output
• Cancellation tokens properly implemented
• UUID-based task IDs prevent enumeration

📊 Performance Analysis

Positive Aspects

• DashMap provides excellent concurrent performance
• Async/await used appropriately without blocking
• Efficient markdown formatting without excessive allocations
• Smart version resolution before spawning tasks

Recommendations

• Consider using tokio::time::interval for periodic cleanup
• Add metrics/telemetry for monitoring task performance
• Consider task priority queue for heavy load scenarios

🐛 Minor Issues

1. Potential Memory Leak in Long-Running Services

Without automatic cleanup, terminated tasks accumulate over time. While not critical, this could become an issue in long-running deployments.

Recommendation: Add a background cleanup task or implement LRU eviction.

2. Missing Cancellation Cleanup

When a task is cancelled, partial downloads might remain. The comment mentions cleanup but implementation isn't visible.

Verify: Ensure CrateCache::cleanup_partial_download() is called on cancellation.

🧪 Test Coverage

✅ Good Coverage

• Tests properly adapted for async behavior
• Task completion waiting implemented correctly
• Different timeout values for different crate sizes
• Binary-only and workspace detection handled

Suggested Additional Tests

1. Concurrent operation limits (when implemented)
2. Rate limiting behavior (when implemented)
3. Memory usage under load
4. Cancellation during different stages
5. Task cleanup and eviction

📚 Code Quality

Excellent Practices

• Clear module separation and responsibilities
• Comprehensive documentation
• Idiomatic Rust patterns
• Proper use of type system with enums
• Good error messages for users

Minor Suggestions

1. Consider extracting magic numbers as constants (e.g., poll interval, timeouts)
2. Add #[must_use] to functions returning important values
3. Consider using thiserror for error types instead of strings

Conclusion

This is a high-quality implementation that successfully achieves its goals. The async design is sound, thread safety is properly handled with DashMap, and the UX is excellent with rich markdown output. The code demonstrates good Rust practices and proper async patterns.

Verdict: Approve ✅

The implementation is production-ready as-is. The suggested improvements are optimizations that can be added in future PRs:

Future Enhancements (Priority Order)

1. High: Add concurrent task limiting with Semaphore
2. High: Implement automatic cleanup for old tasks
3. Medium: Add rate limiting for task creation
4. Medium: Implement progress percentage within steps
5. Low: Add metrics and monitoring
6. Low: Task priority queue for scheduling

Great work on addressing the previous review feedback, especially the DashMap implementation and race condition fixes! The code is clean, well-documented, and ready for production use.

[snowmead]

@claude fix the two failing ci workflows

[github-actions]

Claude finished @snowmead's task —— View job

---

Fixed CI Workflows ✅

Completed Tasks

• Identify failing CI workflows
• Analyze error logs to determine root causes
• Implement fixes for the failures
• Commit and push the changes

Fixes Applied

I've identified and fixed two issues that were causing CI failures:

1. Clippy warning - Unused import in task_manager.rs:

   • Removed unused use std::collections::HashMap; import (line 7)
   • The code uses DashMap instead, so the HashMap import was unnecessary

2. Formatting issue in integration_tests.rs:

   • Fixed line formatting at lines 1418-1420 where .await was improperly formatted
   • Changed to proper multi-line formatting that cargo fmt expects

Changes Committed

commit e8c87e4
fix: Remove unused import and fix formatting issues for CI

- Remove unused HashMap import from task_manager.rs
- Fix await formatting in integration_tests.rs line 1418-1420

The changes have been pushed to the PR branch. The CI workflows should now pass! 🎉