Coverage for mcp_bridge/tools/background_tasks.py: 33%

1"""

2 Background Task Manager for Stravinsky.

3

4 Provides mechanisms to spawn, monitor, and manage async sub-agents.

5 Tasks are persisted to .stravinsky/tasks.json.

6"""

7

8 import json

9 import subprocess

10 import sys

11 from dataclasses import asdict, dataclass

12 from datetime import datetime

13 from pathlib import Path

14 from typing import Any

15

16

17@dataclass

18 class BackgroundTask:

19 id: str

20 prompt: str

21 model: str

22 status: str # pending, running, completed, failed

23 created_at: str

24 started_at: str | None = None

25 completed_at: str | None = None

26 result: str | None = None

27 error: str | None = None

28 pid: int | None = None

29

30

31 class BackgroundManager:

32 def __init__(self, base_dir: str | None = None):

33 if base_dir:

34 self.base_dir = Path(base_dir)

35 else:

36 # Default to .stravinsky in the current working directory

37 self.base_dir = Path.cwd() / ".stravinsky"

38

39 self.tasks_dir = self.base_dir / "tasks"

40 self.state_file = self.base_dir / "tasks.json"

41

42 self.base_dir.mkdir(parents=True, exist_ok=True)

43 self.tasks_dir.mkdir(parents=True, exist_ok=True)

44

45 if not self.state_file.exists():

46 self._save_tasks({})

47

48 def _load_tasks(self) -> dict[str, Any]:

49 try:

50 with open(self.state_file) as f:

51 return json.load(f)

52 except (json.JSONDecodeError, FileNotFoundError):

53 return {}

54

55 def _save_tasks(self, tasks: dict[str, Any]):

56 with open(self.state_file, "w") as f:

57 json.dump(tasks, f, indent=2)

58

59 def create_task(self, prompt: str, model: str) -> str:

60 import uuid as uuid_module # Local import for MCP context

61 task_id = str(uuid_module.uuid4())[:8]

62 task = BackgroundTask(

63 id=task_id,

64 prompt=prompt,

65 model=model,

66 status="pending",

67 created_at=datetime.isoformat(datetime.now()),

68 )

69

70 tasks = self._load_tasks()

71 tasks[task_id] = asdict(task)

72 self._save_tasks(tasks)

73 return task_id

74

75 def update_task(self, task_id: str, **kwargs):

76 tasks = self._load_tasks()

77 if task_id in tasks:

78 tasks[task_id].update(kwargs)

79 self._save_tasks(tasks)

80

81 def get_task(self, task_id: str) -> dict[str, Any] | None:

82 tasks = self._load_tasks()

83 return tasks.get(task_id)

84

85 def list_tasks(self) -> list[dict[str, Any]]:

86 tasks = self._load_tasks()

87 return list(tasks.values())

88

89 def spawn(self, task_id: str):

90 """

91 Spawns a background process to execute the task.

92 In this implementation, it uses another instance of the MCP bridge

93 or a dedicated tool invoker script.

94 """

95 task = self.get_task(task_id)

96 if not task:

97 return

98

99 # We'll use a wrapper script that handles the model invocation and updates the status

100 # mcp_bridge/tools/task_runner.py (we will create this)

101

102 log_file = self.tasks_dir / f"{task_id}.log"

103

104 # Start the background process

105 cmd = [

106 sys.executable,

107 "-m", "mcp_bridge.tools.task_runner",

108 "--task-id", task_id,

109 "--base-dir", str(self.base_dir)

110 ]

111

112 try:

113 # Using Popen to run in background

114 process = subprocess.Popen(

115 cmd,

116 stdout=open(log_file, "w"),

117 stderr=subprocess.STDOUT,

118 start_new_session=True # Run in its own session so it doesn't die with the server

119 )

120

121 self.update_task(task_id, status="running", pid=process.pid, started_at=datetime.isoformat(datetime.now()))

122 except Exception as e:

123 self.update_task(task_id, status="failed", error=str(e))

124

125

126# Tool interface functions

127

128 async def task_spawn(prompt: str, model: str = "gemini-3-flash") -> str:

129 """Spawns a new background task."""

130 manager = BackgroundManager()

131 task_id = manager.create_task(prompt, model)

132 manager.spawn(task_id)

133 return f"Task spawned with ID: {task_id}. Use task_status('{task_id}') to check progress."

134

135

136 async def task_status(task_id: str) -> str:

137 """Checks the status of a background task."""

138 manager = BackgroundManager()

139 task = manager.get_task(task_id)

140 if not task:

141 return f"Task {task_id} not found."

142

143 status = task["status"]

144 if status == "completed":

145 return f"Task {task_id} COMPLETED:\n\n{task.get('result')}"

146 elif status == "failed":

147 return f"Task {task_id} FAILED:\n\n{task.get('error')}"

148 else:

149 return f"Task {task_id} is currently {status} (PID: {task.get('pid')})."

150

151

152 async def task_list() -> str:

153 """Lists all background tasks."""

154 manager = BackgroundManager()

155 tasks = manager.list_tasks()

156 if not tasks:

157 return "No background tasks found."

158

159 lines = ["Background Tasks:"]

160 for t in tasks:

161 lines.append(f"- [{t['id']}] {t['status']}: {t['prompt'][:50]}...")

162

163 return "\n".join(lines)