Skip to content

Commit eae8a7b

Browse files
carlkesselmanclaude
andcommitted
Add background task system for long-running operations
Implements a multi-user safe background task manager for operations that may exceed MCP connection timeouts (e.g., catalog cloning). New tools: - clone_catalog_async: Start clone in background, returns task_id - get_task_status: Check progress and retrieve results - list_tasks: List user's tasks with optional filtering - cancel_task: Cancel pending/running tasks Key features: - Tasks isolated by user (from credentials) - Thread-safe with RLock for all state access - Progress tracking with step/percent/message - Auto-cleanup of old completed tasks - Extensible TaskType enum for future operations - Graceful shutdown via atexit handler Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent ed58058 commit eae8a7b

File tree

4 files changed

+941
-0
lines changed

4 files changed

+941
-0
lines changed

src/deriva_ml_mcp/server.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from deriva_ml_mcp.resources import register_resources
1616
from deriva_ml_mcp.tools import (
1717
register_annotation_tools,
18+
register_background_task_tools,
1819
register_catalog_tools,
1920
register_data_tools,
2021
register_dataset_tools,
@@ -506,6 +507,40 @@
506507
## Before Calling Tools
507508
508509
**Always verify required parameters before calling any tool.** Check the tool's description and parameter schema to understand which parameters are required vs optional. Never assume a parameter is optional - verify first.
510+
511+
## Background Tasks for Long-Running Operations
512+
513+
Some operations like catalog cloning can take many minutes. Use the async versions
514+
of these tools to avoid timeout issues:
515+
516+
**Starting a long-running operation:**
517+
```python
518+
# Instead of clone_catalog (which may timeout), use:
519+
clone_catalog_async("www.facebase.org", "1",
520+
root_rid="3-HXMC",
521+
dest_hostname="localhost",
522+
alias="my-clone")
523+
# Returns immediately with: {"task_id": "abc123", "status": "started", ...}
524+
```
525+
526+
**Checking progress:**
527+
```python
528+
get_task_status("abc123")
529+
# Returns: {"status": "running", "progress": {"percent_complete": 45.0, ...}}
530+
```
531+
532+
**When complete:**
533+
```python
534+
get_task_status("abc123")
535+
# Returns: {"status": "completed", "result": {...full clone result...}}
536+
```
537+
538+
**Managing tasks:**
539+
- `list_tasks()` - See all your tasks
540+
- `list_tasks(status="running")` - Filter by status
541+
- `cancel_task("abc123")` - Cancel a running task
542+
543+
Tasks are isolated per user - you can only see and manage your own tasks.
509544
""",
510545
)
511546

@@ -517,6 +552,7 @@ def register_all_tools(mcp_server: FastMCP, conn_manager: ConnectionManager) ->
517552
"""Register all DerivaML tools, resources, and prompts with the MCP server."""
518553
# Register tools
519554
register_annotation_tools(mcp_server, conn_manager)
555+
register_background_task_tools(mcp_server, conn_manager)
520556
register_catalog_tools(mcp_server, conn_manager)
521557
register_dataset_tools(mcp_server, conn_manager)
522558
register_vocabulary_tools(mcp_server, conn_manager)
@@ -541,6 +577,21 @@ def register_all_tools(mcp_server: FastMCP, conn_manager: ConnectionManager) ->
541577

542578
def main() -> None:
543579
"""Run the DerivaML MCP server."""
580+
import atexit
581+
582+
from deriva_ml_mcp.tasks import get_task_manager
583+
584+
# Register shutdown handler for background task manager
585+
def shutdown_task_manager() -> None:
586+
logger.info("Shutting down background task manager")
587+
try:
588+
task_manager = get_task_manager()
589+
task_manager.shutdown(wait=False) # Don't block on pending tasks
590+
except Exception as e:
591+
logger.warning(f"Error shutting down task manager: {e}")
592+
593+
atexit.register(shutdown_task_manager)
594+
544595
logger.info("Starting DerivaML MCP server")
545596
mcp.run(transport="stdio")
546597

0 commit comments

Comments
 (0)