| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | import json |
| 5 | import time |
| 6 | import uuid |
| 7 | from dataclasses import dataclass, field, replace |
| 8 | from typing import Any, Literal, TYPE_CHECKING |
| 9 | |
| 10 | from helpers import extract_tools |
| 11 | from helpers.defer import DeferredTask, THREAD_BACKGROUND |
| 12 | from helpers.extension import call_extensions_async |
| 13 | from helpers.print_style import PrintStyle |
| 14 | |
| 15 | if TYPE_CHECKING: |
| 16 | from agent import Agent, AgentConfig, AgentContext |
| 17 | from helpers.log import LogItem |
| 18 | |
| 19 | |
| 20 | PARALLEL_JOBS_KEY = "_parallel_jobs" |
| 21 | PARALLEL_WORKER_PARENT_CONTEXT_KEY = "_parallel_parent_context_id" |
| 22 | PARALLEL_WORKER_JOB_KEY = "_parallel_job_id" |
| 23 | PARALLEL_WORKER_KIND_KEY = "_parallel_worker_kind" |
| 24 | |
| 25 | CHILD_PARENT_CONTEXT_ID_KEY = "parent_context_id" |
| 26 | CHILD_PARENT_AGENT_NUMBER_KEY = "parent_agent_number" |
| 27 | CHILD_PARENT_CONTEXT_KIND_KEY = "parent_context_kind" |
| 28 | CHILD_PARENT_CONTEXT_LABEL_KEY = "parent_context_label" |
| 29 | CHILD_PARALLEL_JOB_ID_KEY = "parallel_job_id" |
| 30 | CHILD_PARALLEL_TOOL_NAME_KEY = "parallel_tool_name" |
| 31 | |
| 32 | DEFAULT_MAX_CALLS = 8 |
| 33 | DEFAULT_TIMEOUT_SECONDS = 300 |
| 34 | POLL_INTERVAL_SECONDS = 0.5 |
| 35 | DISALLOWED_PARALLEL_TOOLS = {"document_query", "response"} |
| 36 | |
| 37 | TERMINAL_STATES = {"success", "error", "cancelled", "timeout"} |
| 38 | JobState = Literal["pending", "running", "success", "error", "cancelled", "timeout"] |
| 39 | JobKind = Literal["tool", "subordinate"] |
| 40 | |
| 41 | |
| 42 | @dataclass |
| 43 | class NormalizedToolCall: |
| 44 | index: int |
| 45 | tool_name: str |
| 46 | tool_args: dict[str, Any] |
| 47 | |
| 48 | |
| 49 | @dataclass |
| 50 | class ParallelJob: |
| 51 | id: str |
| 52 | parent_context_id: str |
| 53 | index: int |
| 54 | tool_name: str |
| 55 | tool_args: dict[str, Any] |
| 56 | kind: JobKind |
| 57 | parent_agent: "Agent | None" = field(default=None, repr=False) |
| 58 | state: JobState = "pending" |
| 59 | created_at: float = field(default_factory=time.time) |
| 60 | started_at: float | None = None |
| 61 | completed_at: float | None = None |
| 62 | result: str | None = None |
| 63 | error: str | None = None |
| 64 | worker_context_id: str | None = None |
| 65 | log_id: str = field(default_factory=lambda: str(uuid.uuid4())) |
| 66 | log_item: "LogItem | None" = field(default=None, repr=False) |
| 67 | deferred_task: DeferredTask | None = field(default=None, repr=False) |
| 68 | parent_history: list[tuple[Any, int]] = field(default_factory=list, repr=False) |
| 69 | |
| 70 | def elapsed(self) -> float: |
| 71 | end = self.completed_at or time.time() |
| 72 | start = self.started_at or self.created_at |
| 73 | return max(0.0, end - start) |
| 74 | |
| 75 | |
| 76 | def extract_tool_calls(args: dict[str, Any]) -> Any: |
| 77 | for key in ("tool_calls", "calls", "items"): |
| 78 | if key in args: |
| 79 | return args.get(key) |
| 80 | return None |
| 81 | |
| 82 | |
| 83 | def normalize_parallel_tool_calls(raw_calls: Any) -> list[NormalizedToolCall]: |
| 84 | if isinstance(raw_calls, str): |
| 85 | try: |
| 86 | raw_calls = json.loads(raw_calls) |
| 87 | except json.JSONDecodeError as exc: |
| 88 | raise ValueError( |
| 89 | "`tool_calls` must be an array of normal tool-call objects." |
| 90 | ) from exc |
| 91 | if not isinstance(raw_calls, list): |
| 92 | raise ValueError("`tool_calls` must be an array of normal tool-call objects.") |
| 93 | if not raw_calls: |
| 94 | raise ValueError("`tool_calls` must contain at least one tool call.") |
| 95 | if len(raw_calls) > DEFAULT_MAX_CALLS: |
| 96 | raise ValueError(f"`tool_calls` supports at most {DEFAULT_MAX_CALLS} items.") |
| 97 | |
| 98 | calls: list[NormalizedToolCall] = [] |
| 99 | for index, raw_call in enumerate(raw_calls): |
| 100 | try: |
| 101 | tool_name, tool_args = extract_tools.normalize_tool_request(raw_call) |
| 102 | except ValueError as exc: |
| 103 | raise ValueError(f"tool_calls[{index}] is not a valid tool call: {exc}") from exc |
| 104 | |
| 105 | if tool_name == "parallel": |
| 106 | raise ValueError("`parallel` cannot be nested inside another `parallel` call.") |
| 107 | if tool_name in DISALLOWED_PARALLEL_TOOLS: |
| 108 | raise ValueError( |
| 109 | f"`{tool_name}` cannot be used inside `parallel`; call it sequentially." |
| 110 | ) |
| 111 | |
| 112 | calls.append( |
| 113 | NormalizedToolCall( |
| 114 | index=index, |
| 115 | tool_name=tool_name, |
| 116 | tool_args=dict(tool_args), |
| 117 | ) |
| 118 | ) |
| 119 | return calls |
| 120 | |
| 121 | |
| 122 | def normalize_job_ids(raw_job_ids: Any) -> list[str]: |
| 123 | if raw_job_ids is None: |
| 124 | return [] |
| 125 | if isinstance(raw_job_ids, str): |
| 126 | return [raw_job_ids] |
| 127 | if isinstance(raw_job_ids, list): |
| 128 | return [str(item) for item in raw_job_ids if str(item).strip()] |
| 129 | raise ValueError("`job_ids` must be a string or an array of strings.") |
| 130 | |
| 131 | |
| 132 | def coerce_bool(value: Any, default: bool) -> bool: |
| 133 | if value is None: |
| 134 | return default |
| 135 | if isinstance(value, bool): |
| 136 | return value |
| 137 | if isinstance(value, str): |
| 138 | normalized = value.strip().lower() |
| 139 | if normalized in {"1", "true", "yes", "on"}: |
| 140 | return True |
| 141 | if normalized in {"0", "false", "no", "off"}: |
| 142 | return False |
| 143 | return bool(value) |
| 144 | |
| 145 | |
| 146 | def coerce_timeout(value: Any) -> int: |
| 147 | if value in (None, ""): |
| 148 | return DEFAULT_TIMEOUT_SECONDS |
| 149 | try: |
| 150 | timeout = int(value) |
| 151 | except (TypeError, ValueError) as exc: |
| 152 | raise ValueError(f"`timeout` must be an integer number of seconds, got {value!r}.") from exc |
| 153 | if timeout <= 0: |
| 154 | raise ValueError("`timeout` must be greater than 0.") |
| 155 | return timeout |
| 156 | |
| 157 | |
| 158 | def _parallel_worker_kind(agent: "Agent | None") -> JobKind | None: |
| 159 | context = getattr(agent, "context", None) |
| 160 | if not context: |
| 161 | return None |
| 162 | kind = context.get_data(PARALLEL_WORKER_KIND_KEY) |
| 163 | if kind in {"tool", "subordinate"}: |
| 164 | return kind |
| 165 | if context.get_data(PARALLEL_WORKER_JOB_KEY): |
| 166 | return "tool" |
| 167 | return None |
| 168 | |
| 169 | |
| 170 | def is_parallel_worker(agent: "Agent | None") -> bool: |
| 171 | return _parallel_worker_kind(agent) == "tool" |
| 172 | |
| 173 | |
| 174 | def queue_parallel_parent_history( |
| 175 | agent: "Agent", |
| 176 | *, |
| 177 | content: Any, |
| 178 | tokens: int = 0, |
| 179 | ) -> bool: |
| 180 | if not is_parallel_worker(agent): |
| 181 | return False |
| 182 | context = agent.context |
| 183 | parent_context_id = str( |
| 184 | context.get_data(PARALLEL_WORKER_PARENT_CONTEXT_KEY) or "" |
| 185 | ) |
| 186 | job_id = str(context.get_data(PARALLEL_WORKER_JOB_KEY) or "") |
| 187 | job = _get_job(parent_context_id, job_id) |
| 188 | if not job or job.kind != "tool": |
| 189 | return False |
| 190 | job.parent_history.append((content, tokens)) |
| 191 | return True |
| 192 | |
| 193 | |
| 194 | def _jobs_for_context(context: "AgentContext") -> dict[str, ParallelJob]: |
| 195 | jobs = context.get_data(PARALLEL_JOBS_KEY) |
| 196 | if not isinstance(jobs, dict): |
| 197 | jobs = {} |
| 198 | context.set_data(PARALLEL_JOBS_KEY, jobs) |
| 199 | return jobs |
| 200 | |
| 201 | |
| 202 | def _get_job(parent_context_id: str, job_id: str) -> ParallelJob | None: |
| 203 | from agent import AgentContext |
| 204 | |
| 205 | context = AgentContext.get(parent_context_id) |
| 206 | if not context: |
| 207 | return None |
| 208 | job = _jobs_for_context(context).get(job_id) |
| 209 | return job if isinstance(job, ParallelJob) else None |
| 210 | |
| 211 | |
| 212 | def _new_job_id(tool_name: str) -> str: |
| 213 | prefix = "".join(ch for ch in tool_name if ch.isalnum())[:12] or "job" |
| 214 | return f"{prefix}-{uuid.uuid4().hex[:8]}" |
| 215 | |
| 216 | |
| 217 | async def start_parallel_jobs( |
| 218 | agent: "Agent", |
| 219 | calls: list[NormalizedToolCall], |
| 220 | ) -> list[ParallelJob]: |
| 221 | jobs: list[ParallelJob] = [] |
| 222 | context = agent.context |
| 223 | job_store = _jobs_for_context(context) |
| 224 | |
| 225 | for call in calls: |
| 226 | kind: JobKind = "subordinate" if call.tool_name == "call_subordinate" else "tool" |
| 227 | job = ParallelJob( |
| 228 | id=_new_job_id(call.tool_name), |
| 229 | parent_context_id=context.id, |
| 230 | index=call.index, |
| 231 | tool_name=call.tool_name, |
| 232 | tool_args=call.tool_args, |
| 233 | kind=kind, |
| 234 | parent_agent=agent, |
| 235 | ) |
| 236 | job_store[job.id] = job |
| 237 | jobs.append(job) |
| 238 | _log_parallel_child_started(agent, job) |
| 239 | |
| 240 | try: |
| 241 | job.state = "running" |
| 242 | job.started_at = time.time() |
| 243 | task = DeferredTask(thread_name=THREAD_BACKGROUND) |
| 244 | job.deferred_task = task |
| 245 | if _parallel_worker_kind(agent) == "subordinate" and context.task: |
| 246 | context.task.add_child_task(task) |
| 247 | task.start_task(_run_parallel_job, context.id, job.id) |
| 248 | except Exception as exc: |
| 249 | _finish_job(job, "error", error=str(exc)) |
| 250 | |
| 251 | return jobs |
| 252 | |
| 253 | |
| 254 | async def await_parallel_jobs( |
| 255 | agent: "Agent", |
| 256 | job_ids: list[str], |
| 257 | timeout: int = DEFAULT_TIMEOUT_SECONDS, |
| 258 | *, |
| 259 | collect: bool = True, |
| 260 | wait: bool = True, |
| 261 | ) -> list[dict[str, Any]]: |
| 262 | if not job_ids: |
| 263 | raise ValueError("No `job_ids` were provided to await.") |
| 264 | |
| 265 | deadline = time.time() + timeout |
| 266 | wait_timed_out_job_ids: set[str] = set() |
| 267 | while True: |
| 268 | await refresh_parallel_jobs(agent) |
| 269 | jobs = [_jobs_for_context(agent.context).get(job_id) for job_id in job_ids] |
| 270 | missing = [job_id for job_id, job in zip(job_ids, jobs) if job is None] |
| 271 | if missing: |
| 272 | raise ValueError(f"Unknown parallel job id(s): {', '.join(missing)}") |
| 273 | |
| 274 | active = [job for job in jobs if job and job.state not in TERMINAL_STATES] |
| 275 | if not wait or not active: |
| 276 | break |
| 277 | |
| 278 | if time.time() >= deadline: |
| 279 | wait_timed_out_job_ids = {job.id for job in active} |
| 280 | break |
| 281 | |
| 282 | await asyncio.sleep(POLL_INTERVAL_SECONDS) |
| 283 | |
| 284 | snapshots = [] |
| 285 | for job_id in job_ids: |
| 286 | job = _jobs_for_context(agent.context).get(job_id) |
| 287 | if job: |
| 288 | snapshot = _job_snapshot(job, include_result=True) |
| 289 | if job.id in wait_timed_out_job_ids and job.state not in TERMINAL_STATES: |
| 290 | snapshot["wait_timed_out"] = True |
| 291 | snapshots.append(snapshot) |
| 292 | |
| 293 | if collect: |
| 294 | await collect_parallel_jobs(agent, job_ids) |
| 295 | |
| 296 | return snapshots |
| 297 | |
| 298 | |
| 299 | async def cancel_parallel_jobs(agent: "Agent", job_ids: list[str]) -> list[dict[str, Any]]: |
| 300 | if not job_ids: |
| 301 | raise ValueError("No `job_ids` were provided to cancel.") |
| 302 | |
| 303 | await refresh_parallel_jobs(agent) |
| 304 | snapshots = [] |
| 305 | for job_id in job_ids: |
| 306 | job = _jobs_for_context(agent.context).get(job_id) |
| 307 | if not job: |
| 308 | raise ValueError(f"Unknown parallel job id: {job_id}") |
| 309 | await _cancel_job(job) |
| 310 | snapshots.append(_job_snapshot(job, include_result=True)) |
| 311 | await cleanup_parallel_job(agent, job) |
| 312 | _jobs_for_context(agent.context).pop(job_id, None) |
| 313 | return snapshots |
| 314 | |
| 315 | |
| 316 | async def refresh_parallel_jobs(agent: "Agent") -> list[ParallelJob]: |
| 317 | jobs = list(_jobs_for_context(agent.context).values()) |
| 318 | for job in jobs: |
| 319 | if job.state in TERMINAL_STATES: |
| 320 | continue |
| 321 | task = job.deferred_task |
| 322 | if not task: |
| 323 | continue |
| 324 | if task.is_ready(): |
| 325 | try: |
| 326 | await task.result() |
| 327 | except asyncio.CancelledError: |
| 328 | _finish_job(job, "cancelled", error="Parallel job was cancelled.") |
| 329 | except Exception as exc: |
| 330 | _finish_job(job, "error", error=str(exc)) |
| 331 | elif task.is_alive(): |
| 332 | job.state = "running" |
| 333 | if job.started_at is None: |
| 334 | job.started_at = time.time() |
| 335 | return jobs |
| 336 | |
| 337 | |
| 338 | async def cleanup_parallel_job(agent: "Agent", job: ParallelJob) -> None: |
| 339 | if job.deferred_task and job.deferred_task.is_alive(): |
| 340 | job.deferred_task.kill() |
| 341 | if job.kind == "tool": |
| 342 | await _remove_context(job.worker_context_id) |
| 343 | |
| 344 | |
| 345 | async def collect_parallel_jobs( |
| 346 | agent: "Agent", |
| 347 | job_ids: list[str], |
| 348 | *, |
| 349 | promote_parent_history: bool = False, |
| 350 | ) -> None: |
| 351 | jobs = _jobs_for_context(agent.context) |
| 352 | for job_id in dict.fromkeys(job_ids): |
| 353 | job = jobs.get(job_id) |
| 354 | if not job or job.state not in TERMINAL_STATES: |
| 355 | continue |
| 356 | if promote_parent_history: |
| 357 | for content, tokens in job.parent_history: |
| 358 | agent.hist_add_message(False, content=content, tokens=tokens) |
| 359 | job.parent_history.clear() |
| 360 | await cleanup_parallel_job(agent, job) |
| 361 | jobs.pop(job_id, None) |
| 362 | |
| 363 | |
| 364 | async def build_parallel_jobs_extras(agent: "Agent") -> str: |
| 365 | await refresh_parallel_jobs(agent) |
| 366 | jobs = [ |
| 367 | job |
| 368 | for job in _jobs_for_context(agent.context).values() |
| 369 | if isinstance(job, ParallelJob) and job.state not in {"cancelled", "timeout"} |
| 370 | ] |
| 371 | if not jobs: |
| 372 | return "" |
| 373 | |
| 374 | active = [job for job in jobs if job.state not in TERMINAL_STATES] |
| 375 | ready = [job for job in jobs if job.state in TERMINAL_STATES] |
| 376 | if not active and not ready: |
| 377 | return "" |
| 378 | |
| 379 | lines = ["parallel jobs:"] |
| 380 | if active: |
| 381 | lines.append("running:") |
| 382 | for job in active: |
| 383 | lines.append( |
| 384 | f"- {job.id}: {job.tool_name} [{job.state}], running for {job.elapsed():.1f}s" |
| 385 | ) |
| 386 | if ready: |
| 387 | lines.append("ready to collect with `parallel` and `job_ids`:") |
| 388 | for job in ready: |
| 389 | lines.append( |
| 390 | f"- {job.id}: {job.tool_name} [{job.state}], duration {job.elapsed():.1f}s" |
| 391 | ) |
| 392 | lines.append("call the `parallel` tool with `job_ids` to await/collect results or `action: \"cancel\"` to cancel.") |
| 393 | return "\n".join(lines) |
| 394 | |
| 395 | |
| 396 | def format_started_jobs(jobs: list[ParallelJob]) -> str: |
| 397 | payload = { |
| 398 | "status": "started", |
| 399 | "jobs": [_job_snapshot(job, include_result=False) for job in jobs], |
| 400 | "instruction": "Use the parallel tool with job_ids to await or cancel these background jobs.", |
| 401 | } |
| 402 | return json.dumps(payload, ensure_ascii=False, separators=(",", ":")) |
| 403 | |
| 404 | |
| 405 | def format_parallel_results(results: list[dict[str, Any]]) -> str: |
| 406 | states = [result.get("state") for result in results] |
| 407 | has_active_jobs = any(state not in TERMINAL_STATES for state in states) |
| 408 | wait_timed_out = any(result.get("wait_timed_out") for result in results) |
| 409 | if has_active_jobs: |
| 410 | status = "waiting" if wait_timed_out else "running" |
| 411 | elif states and all(state == "success" for state in states): |
| 412 | status = "success" |
| 413 | elif states and all(state == "cancelled" for state in states): |
| 414 | status = "cancelled" |
| 415 | elif any(state == "success" for state in states): |
| 416 | status = "partial" |
| 417 | else: |
| 418 | status = "error" |
| 419 | |
| 420 | payload = { |
| 421 | "status": status, |
| 422 | "jobs": results, |
| 423 | } |
| 424 | if wait_timed_out: |
| 425 | payload["wait_timeout"] = True |
| 426 | if has_active_jobs: |
| 427 | payload["instruction"] = ( |
| 428 | "Some jobs are still running. Call `parallel` with `action: \"await\"` " |
| 429 | "and the listed `job_ids` to wait again, or `action: \"cancel\"` to stop them." |
| 430 | ) |
| 431 | return json.dumps(payload, ensure_ascii=False, separators=(",", ":")) |
| 432 | |
| 433 | |
| 434 | async def _run_parallel_job(parent_context_id: str, job_id: str) -> None: |
| 435 | job = _get_job(parent_context_id, job_id) |
| 436 | if not job: |
| 437 | return |
| 438 | try: |
| 439 | if job.kind == "subordinate": |
| 440 | result = await _run_subordinate_context_job(parent_context_id, job) |
| 441 | else: |
| 442 | result = await _run_direct_tool_job(parent_context_id, job) |
| 443 | _finish_job(job, "success", result=result) |
| 444 | except asyncio.CancelledError: |
| 445 | _finish_job(job, "cancelled", error="Parallel job was cancelled.") |
| 446 | raise |
| 447 | except Exception as exc: |
| 448 | _finish_job(job, "error", error=str(exc)) |
| 449 | PrintStyle.error(f"Parallel job {job.id} failed: {exc}") |
| 450 | |
| 451 | |
| 452 | async def _run_subordinate_context_job(parent_context_id: str, job: ParallelJob) -> str: |
| 453 | from agent import AgentContext |
| 454 | from helpers.tool_policy import ensure_tool_allowed |
| 455 | from tools.call_subordinate import get_or_create_subordinate, run_subordinate |
| 456 | |
| 457 | parent_context = AgentContext.get(parent_context_id) |
| 458 | if not parent_context: |
| 459 | raise ValueError("Parent context not found.") |
| 460 | parent_agent = job.parent_agent or parent_context.agent0 |
| 461 | ensure_tool_allowed(parent_agent, "call_subordinate") |
| 462 | |
| 463 | args = job.tool_args |
| 464 | message = str(args.get("message") or "").strip() |
| 465 | if not message: |
| 466 | raise ValueError("call_subordinate requires `tool_args.message`.") |
| 467 | |
| 468 | context_id = str(args.get("context_id") or args.get("agent_id") or "").strip() |
| 469 | reset = args.get("reset", False) |
| 470 | slot = ( |
| 471 | job.id |
| 472 | if coerce_bool(reset, False) and not context_id |
| 473 | else "default" |
| 474 | ) |
| 475 | attachments = args.get("attachments") if isinstance(args.get("attachments"), list) else [] |
| 476 | subordinate = get_or_create_subordinate( |
| 477 | parent_agent, |
| 478 | profile=str(args.get("profile") or args.get("agent_profile") or ""), |
| 479 | reset=reset, |
| 480 | context_id=context_id, |
| 481 | name=str(args.get("name") or ""), |
| 482 | message=message, |
| 483 | slot=slot, |
| 484 | ) |
| 485 | worker_context = subordinate.context |
| 486 | job.worker_context_id = worker_context.id |
| 487 | if job.deferred_task: |
| 488 | worker_context.task = job.deferred_task |
| 489 | |
| 490 | worker_context.set_data(PARALLEL_WORKER_PARENT_CONTEXT_KEY, parent_context.id) |
| 491 | worker_context.set_data(PARALLEL_WORKER_JOB_KEY, job.id) |
| 492 | worker_context.set_data(PARALLEL_WORKER_KIND_KEY, job.kind) |
| 493 | worker_context.set_output_data(CHILD_PARALLEL_JOB_ID_KEY, job.id) |
| 494 | worker_context.set_output_data(CHILD_PARALLEL_TOOL_NAME_KEY, job.tool_name) |
| 495 | return await run_subordinate(parent_agent, subordinate, message, attachments) |
| 496 | |
| 497 | |
| 498 | async def _run_direct_tool_job(parent_context_id: str, job: ParallelJob) -> str: |
| 499 | from agent import AgentContext, AgentContextType, LoopData |
| 500 | |
| 501 | parent_context = AgentContext.get(parent_context_id) |
| 502 | if not parent_context: |
| 503 | raise ValueError("Parent context not found.") |
| 504 | |
| 505 | worker_context: AgentContext | None = None |
| 506 | try: |
| 507 | worker_context = AgentContext( |
| 508 | config=_clone_config(parent_context.config), |
| 509 | name=f"parallel:{job.tool_name}", |
| 510 | type=AgentContextType.BACKGROUND, |
| 511 | ) |
| 512 | worker_context.set_data(PARALLEL_WORKER_PARENT_CONTEXT_KEY, parent_context_id) |
| 513 | worker_context.set_data(PARALLEL_WORKER_JOB_KEY, job.id) |
| 514 | worker_context.set_data(PARALLEL_WORKER_KIND_KEY, job.kind) |
| 515 | worker_context.set_data( |
| 516 | "chat_model_override", |
| 517 | parent_context.get_data("chat_model_override"), |
| 518 | ) |
| 519 | job.worker_context_id = worker_context.id |
| 520 | _copy_project(parent_context, worker_context) |
| 521 | |
| 522 | worker_agent = worker_context.agent0 |
| 523 | worker_agent.last_user_message = parent_context.agent0.last_user_message |
| 524 | worker_agent.loop_data = LoopData() |
| 525 | return await execute_tool_call( |
| 526 | worker_agent, |
| 527 | job.tool_name, |
| 528 | job.tool_args, |
| 529 | log_item=job.log_item, |
| 530 | ) |
| 531 | finally: |
| 532 | if worker_context: |
| 533 | await _remove_context(worker_context.id) |
| 534 | |
| 535 | |
| 536 | async def execute_tool_call( |
| 537 | agent: "Agent", |
| 538 | tool_name: str, |
| 539 | tool_args: dict[str, Any], |
| 540 | *, |
| 541 | log_item: "LogItem | None" = None, |
| 542 | ) -> str: |
| 543 | if tool_name == "parallel": |
| 544 | raise ValueError("`parallel` cannot be nested inside a parallel worker.") |
| 545 | |
| 546 | tool = _resolve_parallel_tool(agent, tool_name, tool_args, strict=True) |
| 547 | if not tool: |
| 548 | raise ValueError(f"Tool '{tool_name}' not found or could not be initialized.") |
| 549 | |
| 550 | original_get_log_object = None |
| 551 | if log_item is not None: |
| 552 | original_get_log_object = tool.get_log_object |
| 553 | tool.get_log_object = lambda: log_item |
| 554 | |
| 555 | agent.loop_data.current_tool = tool |
| 556 | try: |
| 557 | await agent.handle_intervention() |
| 558 | await tool.before_execution(**tool_args) |
| 559 | await agent.handle_intervention() |
| 560 | await call_extensions_async( |
| 561 | "tool_execute_before", |
| 562 | agent, |
| 563 | tool_args=tool_args or {}, |
| 564 | tool_name=tool_name, |
| 565 | ) |
| 566 | response = await tool.execute(**tool_args) |
| 567 | await agent.handle_intervention() |
| 568 | await call_extensions_async( |
| 569 | "tool_execute_after", |
| 570 | agent, |
| 571 | response=response, |
| 572 | tool_name=tool_name, |
| 573 | ) |
| 574 | await tool.after_execution(response) |
| 575 | await agent.handle_intervention() |
| 576 | return response.message |
| 577 | finally: |
| 578 | if original_get_log_object is not None: |
| 579 | tool.get_log_object = original_get_log_object |
| 580 | agent.loop_data.current_tool = None |
| 581 | |
| 582 | |
| 583 | def _resolve_parallel_tool( |
| 584 | agent: "Agent", |
| 585 | tool_name: str, |
| 586 | tool_args: dict[str, Any], |
| 587 | *, |
| 588 | strict: bool = False, |
| 589 | ): |
| 590 | message = json.dumps({"tool_name": tool_name, "tool_args": tool_args}) |
| 591 | |
| 592 | tool = None |
| 593 | try: |
| 594 | import helpers.mcp_handler as mcp_helper |
| 595 | |
| 596 | tool = mcp_helper.MCPConfig.get_instance().get_tool(agent, tool_name) |
| 597 | except ImportError: |
| 598 | tool = None |
| 599 | except Exception as exc: |
| 600 | if strict: |
| 601 | raise |
| 602 | PrintStyle.warning(f"Failed to initialize MCP tool '{tool_name}' for parallel job: {exc}") |
| 603 | |
| 604 | if not tool: |
| 605 | get_tool = getattr(agent, "get_tool", None) |
| 606 | if not callable(get_tool): |
| 607 | if strict: |
| 608 | raise ValueError(f"Tool '{tool_name}' not found or could not be initialized.") |
| 609 | return None |
| 610 | try: |
| 611 | tool = get_tool( |
| 612 | name=tool_name, |
| 613 | method=None, |
| 614 | args=tool_args, |
| 615 | message=message, |
| 616 | loop_data=getattr(agent, "loop_data", None), |
| 617 | ) |
| 618 | except Exception as exc: |
| 619 | if strict: |
| 620 | raise |
| 621 | PrintStyle.warning(f"Failed to initialize tool '{tool_name}' for parallel job: {exc}") |
| 622 | tool = None |
| 623 | |
| 624 | if not tool: |
| 625 | return None |
| 626 | |
| 627 | try: |
| 628 | tool.args = dict(tool_args) |
| 629 | except Exception: |
| 630 | pass |
| 631 | |
| 632 | return tool |
| 633 | |
| 634 | |
| 635 | async def _cancel_job( |
| 636 | job: ParallelJob, |
| 637 | *, |
| 638 | state: JobState = "cancelled", |
| 639 | message: str = "Parallel job was cancelled.", |
| 640 | ) -> None: |
| 641 | if job.deferred_task and job.deferred_task.is_alive(): |
| 642 | job.deferred_task.kill() |
| 643 | _finish_job(job, state, error=message) |
| 644 | |
| 645 | |
| 646 | def _finish_job( |
| 647 | job: ParallelJob, |
| 648 | state: JobState, |
| 649 | *, |
| 650 | result: str | None = None, |
| 651 | error: str | None = None, |
| 652 | ) -> None: |
| 653 | job.state = state |
| 654 | job.completed_at = time.time() |
| 655 | if job.started_at is None: |
| 656 | job.started_at = job.created_at |
| 657 | if result is not None: |
| 658 | job.result = result |
| 659 | if error is not None: |
| 660 | job.error = error |
| 661 | _update_parallel_child_log(job) |
| 662 | |
| 663 | |
| 664 | async def _remove_context(context_id: str | None) -> None: |
| 665 | if not context_id: |
| 666 | return |
| 667 | from agent import AgentContext |
| 668 | from helpers import persist_chat |
| 669 | |
| 670 | context = AgentContext.get(context_id) |
| 671 | if context: |
| 672 | try: |
| 673 | context.reset() |
| 674 | except Exception: |
| 675 | pass |
| 676 | AgentContext.remove(context_id) |
| 677 | persist_chat.remove_chat(context_id) |
| 678 | |
| 679 | |
| 680 | def _log_parallel_child_started(agent: "Agent", job: ParallelJob) -> None: |
| 681 | if job.kind == "subordinate": |
| 682 | job.log_item = agent.context.log.log( |
| 683 | type="subagent", |
| 684 | heading=f"icon://communication {agent.agent_name}: Calling Subordinate Agent", |
| 685 | content="", |
| 686 | kvps=job.tool_args, |
| 687 | id=job.log_id, |
| 688 | ) |
| 689 | return |
| 690 | |
| 691 | tool = _resolve_parallel_tool(agent, job.tool_name, job.tool_args) |
| 692 | if tool is not None: |
| 693 | try: |
| 694 | job.log_item = tool.get_log_object() |
| 695 | if job.log_item is not None: |
| 696 | return |
| 697 | except Exception as exc: |
| 698 | PrintStyle.warning( |
| 699 | f"Failed to derive parallel child log for {job.tool_name}: {exc}" |
| 700 | ) |
| 701 | |
| 702 | heading = f"icon://construction {agent.agent_name}: Using tool '{job.tool_name}'" |
| 703 | job.log_item = agent.context.log.log( |
| 704 | type="tool", |
| 705 | heading=heading, |
| 706 | content="", |
| 707 | kvps=job.tool_args, |
| 708 | id=job.log_id, |
| 709 | _tool_name=job.tool_name, |
| 710 | ) |
| 711 | |
| 712 | |
| 713 | def _update_parallel_child_log(job: ParallelJob) -> None: |
| 714 | if not job.log_item: |
| 715 | return |
| 716 | if job.state == "success": |
| 717 | content = job.result if job.result else "(completed without textual output)" |
| 718 | else: |
| 719 | content = f"Error: {job.error or job.state}" |
| 720 | try: |
| 721 | job.log_item.update(content=content) |
| 722 | except Exception: |
| 723 | pass |
| 724 | |
| 725 | |
| 726 | def _job_snapshot(job: ParallelJob, *, include_result: bool) -> dict[str, Any]: |
| 727 | data: dict[str, Any] = { |
| 728 | "job_id": job.id, |
| 729 | "tool_name": job.tool_name, |
| 730 | "state": job.state, |
| 731 | "duration_seconds": round(job.elapsed(), 3), |
| 732 | } |
| 733 | if job.worker_context_id: |
| 734 | data["context_id"] = job.worker_context_id |
| 735 | if include_result: |
| 736 | if job.result is not None: |
| 737 | data["result"] = job.result |
| 738 | if job.error is not None: |
| 739 | data["error"] = job.error |
| 740 | return data |
| 741 | |
| 742 | |
| 743 | def _clone_config(config: "AgentConfig") -> "AgentConfig": |
| 744 | try: |
| 745 | return replace( |
| 746 | config, |
| 747 | knowledge_subdirs=list(config.knowledge_subdirs), |
| 748 | additional=dict(config.additional), |
| 749 | ) |
| 750 | except Exception: |
| 751 | return config |
| 752 | |
| 753 | |
| 754 | def _copy_project(parent_context: "AgentContext", worker_context: "AgentContext") -> None: |
| 755 | try: |
| 756 | from helpers import projects |
| 757 | |
| 758 | project_name = projects.get_context_project_name(parent_context) |
| 759 | if project_name: |
| 760 | projects.activate_project(worker_context.id, project_name, mark_dirty=False) |
| 761 | except Exception: |
| 762 | pass |