Skip to content

Commit e200fbd

Browse files
maxisbeymaxisbey
committed
Replace magic strings with constants for task status and metadata
- Add TASK_STATUS_* constants to types.py for task status values - Use RELATED_TASK_METADATA_KEY constant instead of hardcoded string - Use RelatedTaskMetadata model instead of raw dict with 'taskId' - Update all task status usages to use constants
1 parent 9c535a7 commit e200fbd

5 files changed

Lines changed: 30 additions & 16 deletions

File tree

‎src/mcp/server/experimental/task_context.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@
1919
from mcp.shared.experimental.tasks.resolver import Resolver
2020
from mcp.shared.experimental.tasks.store import TaskStore
2121
from mcp.types import (
22+
TASK_STATUS_INPUT_REQUIRED,
23+
TASK_STATUS_WORKING,
2224
ClientCapabilities,
2325
CreateMessageResult,
2426
ElicitationCapability,
@@ -236,7 +238,7 @@ async def elicit(
236238
raise RuntimeError("handler is required for elicit(). Pass handler= to ServerTaskContext.")
237239

238240
# Update status to input_required
239-
await self._store.update_task(self.task_id, status="input_required")
241+
await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
240242

241243
# Build the request using session's helper
242244
request = self._session._build_elicit_request( # pyright: ignore[reportPrivateUsage]
@@ -262,10 +264,10 @@ async def elicit(
262264
try:
263265
# Wait for response (routed back via TaskResultHandler)
264266
response_data = await resolver.wait()
265-
await self._store.update_task(self.task_id, status="working")
267+
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
266268
return ElicitResult.model_validate(response_data)
267269
except anyio.get_cancelled_exc_class():
268-
await self._store.update_task(self.task_id, status="working")
270+
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
269271
raise
270272

271273
async def create_message(
@@ -313,7 +315,7 @@ async def create_message(
313315
raise RuntimeError("handler is required for create_message(). Pass handler= to ServerTaskContext.")
314316

315317
# Update status to input_required
316-
await self._store.update_task(self.task_id, status="input_required")
318+
await self._store.update_task(self.task_id, status=TASK_STATUS_INPUT_REQUIRED)
317319

318320
# Build the request using session's helper
319321
request = self._session._build_create_message_request( # pyright: ignore[reportPrivateUsage]
@@ -345,8 +347,8 @@ async def create_message(
345347
try:
346348
# Wait for response (routed back via TaskResultHandler)
347349
response_data = await resolver.wait()
348-
await self._store.update_task(self.task_id, status="working")
350+
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
349351
return CreateMessageResult.model_validate(response_data)
350352
except anyio.get_cancelled_exc_class():
351-
await self._store.update_task(self.task_id, status="working")
353+
await self._store.update_task(self.task_id, status=TASK_STATUS_WORKING)
352354
raise

‎src/mcp/server/experimental/task_result_handler.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
from mcp.server.session import ServerSession
1919
from mcp.shared.exceptions import McpError
20-
from mcp.shared.experimental.tasks.helpers import is_terminal
20+
from mcp.shared.experimental.tasks.helpers import RELATED_TASK_METADATA_KEY, is_terminal
2121
from mcp.shared.experimental.tasks.message_queue import TaskMessageQueue
2222
from mcp.shared.experimental.tasks.resolver import Resolver
2323
from mcp.shared.experimental.tasks.store import TaskStore
@@ -28,6 +28,7 @@
2828
GetTaskPayloadRequest,
2929
GetTaskPayloadResult,
3030
JSONRPCMessage,
31+
RelatedTaskMetadata,
3132
RequestId,
3233
)
3334

@@ -126,9 +127,9 @@ async def handle(
126127
result = await self._store.get_result(task_id)
127128
# GetTaskPayloadResult is a Result with extra="allow"
128129
# The stored result contains the actual payload data
129-
# Per spec: tasks/result MUST include _meta.io.modelcontextprotocol/related-task
130-
# with taskId, as the result structure itself does not contain the task ID
131-
related_task_meta: dict[str, Any] = {"io.modelcontextprotocol/related-task": {"taskId": task_id}}
130+
# Per spec: tasks/result MUST include _meta with related-task metadata
131+
related_task = RelatedTaskMetadata(taskId=task_id)
132+
related_task_meta: dict[str, Any] = {RELATED_TASK_METADATA_KEY: related_task.model_dump(by_alias=True)}
132133
if result is not None:
133134
# Copy result fields and add required metadata
134135
result_data = result.model_dump(by_alias=True)

‎src/mcp/shared/experimental/tasks/context.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
"""
88

99
from mcp.shared.experimental.tasks.store import TaskStore
10-
from mcp.types import Result, Task
10+
from mcp.types import TASK_STATUS_COMPLETED, TASK_STATUS_FAILED, Result, Task
1111

1212

1313
class TaskContext:
@@ -84,7 +84,7 @@ async def complete(self, result: Result) -> None:
8484
await self._store.store_result(self.task_id, result)
8585
self._task = await self._store.update_task(
8686
self.task_id,
87-
status="completed",
87+
status=TASK_STATUS_COMPLETED,
8888
)
8989

9090
async def fail(self, error: str) -> None:
@@ -96,6 +96,6 @@ async def fail(self, error: str) -> None:
9696
"""
9797
self._task = await self._store.update_task(
9898
self.task_id,
99-
status="failed",
99+
status=TASK_STATUS_FAILED,
100100
status_message=error,
101101
)

‎src/mcp/shared/experimental/tasks/helpers.py‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,10 @@
1515
from mcp.shared.experimental.tasks.store import TaskStore
1616
from mcp.types import (
1717
INVALID_PARAMS,
18+
TASK_STATUS_CANCELLED,
19+
TASK_STATUS_COMPLETED,
20+
TASK_STATUS_FAILED,
21+
TASK_STATUS_WORKING,
1822
CancelTaskResult,
1923
ErrorData,
2024
Task,
@@ -43,7 +47,7 @@ def is_terminal(status: TaskStatus) -> bool:
4347
Returns:
4448
True if the status is terminal (completed, failed, or cancelled)
4549
"""
46-
return status in ("completed", "failed", "cancelled")
50+
return status in (TASK_STATUS_COMPLETED, TASK_STATUS_FAILED, TASK_STATUS_CANCELLED)
4751

4852

4953
async def cancel_task(
@@ -94,7 +98,7 @@ async def handle_cancel(request: CancelTaskRequest) -> CancelTaskResult:
9498
)
9599

96100
# Update task to cancelled status
97-
cancelled_task = await store.update_task(task_id, status="cancelled")
101+
cancelled_task = await store.update_task(task_id, status=TASK_STATUS_CANCELLED)
98102
return CancelTaskResult(**cancelled_task.model_dump())
99103

100104

@@ -122,7 +126,7 @@ def create_task_state(
122126
now = datetime.now(timezone.utc)
123127
return Task(
124128
taskId=task_id or generate_task_id(),
125-
status="working",
129+
status=TASK_STATUS_WORKING,
126130
createdAt=now,
127131
lastUpdatedAt=now,
128132
ttl=metadata.ttl,

‎src/mcp/types.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -519,6 +519,13 @@ class ServerCapabilities(BaseModel):
519519

520520
TaskStatus = Literal["working", "input_required", "completed", "failed", "cancelled"]
521521

522+
# Task status constants
523+
TASK_STATUS_WORKING: Final[Literal["working"]] = "working"
524+
TASK_STATUS_INPUT_REQUIRED: Final[Literal["input_required"]] = "input_required"
525+
TASK_STATUS_COMPLETED: Final[Literal["completed"]] = "completed"
526+
TASK_STATUS_FAILED: Final[Literal["failed"]] = "failed"
527+
TASK_STATUS_CANCELLED: Final[Literal["cancelled"]] = "cancelled"
528+
522529

523530
class RelatedTaskMetadata(BaseModel):
524531
"""

0 commit comments

Comments
 (0)