Coverage for gco_mcp / resources / tasks.py: 100.00%

1"""Task status resources (tasks://gco/...) for the GCO MCP server. 

2 

3Reads through the MCP tasks extension (SEP-2663, the ``fastmcp_tasks`` 

4package in FastMCP 4) to surface the status of a long-running tool 

5invocation as JSON. The extension's own ``tasks/get`` handler is reused so 

6this resource reports exactly what a protocol-native ``tasks/get`` request 

7would return — status, timestamps, poll interval, and the inlined result or 

8error for finished tasks. Returns a graceful error stub when the FastMCP 

9build in use doesn't ship the tasks extension. 

10""" 

11 

12from __future__ import annotations 

13 

14import json 

15import re 

16from typing import Any 

17 

18# Task IDs are client-controlled strings (FastMCP forwards whatever the 

19# client passed). Restrict to a generous alphanumeric+ punctuation set 

20# so a malformed URI expansion can't sneak shell metacharacters into 

21# downstream lookups. 

22_TASK_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:-]{0,255}$") 

23 

24 

25async def _lookup_task_state(task_id: str) -> dict[str, Any] | None: 

26 """Look up a task's current state through the tasks extension. 

27 

28 Delegates to ``fastmcp_tasks.handlers.tasks_get`` — the same handler 

29 that serves protocol ``tasks/get`` requests — so the resource view and 

30 the wire view can never disagree. Returns ``None`` when the extension 

31 is unavailable on this build (the caller turns that into a graceful 

32 "not available" JSON) and raises ``LookupError`` when the extension is 

33 present but knows nothing about ``task_id``. 

34 """ 

35 try: 

36 from fastmcp_tasks.handlers import tasks_get 

37 from server import mcp as _mcp 

38 except ImportError: 

39 return None 

40 

41 try: 

42 record = await tasks_get(_mcp, task_id) 

43 except Exception as exc: 

44 # The handler raises the protocol not-found error for unknown or 

45 # expired task ids (and for a docket that has not started yet). 

46 raise LookupError(str(exc)) from exc 

47 return _coerce_to_dict(record) 

48 

49 

50def _coerce_to_dict(record: object) -> dict[str, Any]: 

51 """Best-effort conversion of an opaque task record to a JSON-friendly dict.""" 

52 if isinstance(record, dict): 

53 return record 

54 for attr in ("model_dump", "dict", "to_dict", "_asdict"): 

55 method = getattr(record, attr, None) 

56 if callable(method): 

57 try: 

58 payload = method() 

59 except Exception: # noqa: BLE001 

60 continue 

61 if isinstance(payload, dict): 

62 return payload 

63 if hasattr(record, "__dict__"): 

64 return {k: v for k, v in vars(record).items() if not k.startswith("_")} 

65 return {"value": str(record)} 

66 

67 

68async def _task_resource(task_id: str) -> str: 

69 """Return the current status of ``task_id`` as JSON.""" 

70 if not _TASK_ID_RE.match(task_id): 

71 return json.dumps({"error": "invalid task_id", "value": task_id}) 

72 try: 

73 state = await _lookup_task_state(task_id) 

74 except LookupError as exc: 

75 return json.dumps( 

76 { 

77 "error": "task not found", 

78 "detail": str(exc)[:200], 

79 "task_id": task_id, 

80 } 

81 ) 

82 if state is None: 

83 return json.dumps( 

84 { 

85 "error": "task protocol not available", 

86 "detail": ( 

87 "this build of FastMCP does not ship the tasks extension " 

88 "(fastmcp_tasks) this resource handler reads through" 

89 ), 

90 "task_id": task_id, 

91 } 

92 ) 

93 return json.dumps({"task_id": task_id, "state": state}, indent=2, default=str) 

94 

95 

96def register(mcp_instance: Any) -> None: 

97 """Register the task-status resource against the shared MCP server.""" 

98 mcp_instance.resource("tasks://gco/{task_id}")(_task_resource)