content: """End-to-end integration tests for tasks functionality.
  
  These tests demonstrate the full task lifecycle:
  1. Client sends task-augmented request (tools/call with task metadata)
  2. Server creates task and returns CreateTaskResult immediately
  3. Background work executes (using task_execution context manager)
  4. Client polls with tasks/get
  5. Client retrieves result with tasks/result
  """
  
  from dataclasses import dataclass, field
  from typing import Any
  
  import anyio
  import pytest
  from anyio import Event
  from anyio.abc import TaskGroup
  
  from mcp.client.session import ClientSession
  from mcp.server import Server
  from mcp.server.lowlevel import NotificationOptions
  from mcp.server.models import InitializationOptions
  from mcp.server.session import ServerSession
  from mcp.shared.experimental.tasks.helpers import task_execution
  from mcp.shared.experimental.tasks.in_memory_task_store import InMemoryTaskStore
  from mcp.shared.message import SessionMessage
  from mcp.shared.session import RequestResponder
  from mcp.types import (
      TASK_REQUIRED,
      CallToolRequest,
      CallToolRequestParams,
      CallToolResult,
      ClientRequest,
      ClientResult,
      CreateTaskResult,
      GetTaskPayloadRequest,
      GetTaskPayloadRequestParams,
      GetTaskPayloadResult,
      GetTaskRequest,
      GetTaskRequestParams,
      GetTaskResult,
      ListTasksRequest,
      ListTasksResult,
      ServerNotification,
      ServerRequest,
      TaskMetadata,
      TextContent,
      Tool,
      ToolExecution,
  )
  
  
  @dataclass
  class AppContext:
      """Application context passed via lifespan_context."""
  
      task_group: TaskGroup
      store: InMemoryTaskStore
      # Events to signal when tasks complete (for testing without sleeps)
      task_done_events: dict[str, Event] = field(default_factory=lambda: {})
  
  
  @pytest.mark.anyio
  async def test_task_lifecycle_with_task_execution() -> None:
      """
      Test the complete task lifecycle using the task_execution pattern.
  
      This demonstrates the recommended way to implement task-augmented tools:
      1. Create task in store
      2. Spawn work using task_execution() context manager
      3. Return CreateTaskResult immediately
      4. Work executes in background, auto-fails on exception
      """
      # Note: We bypass the normal lifespan mechanism and pass context directly to _handle_message
      server: Server[AppContext, Any] = Server("test-tasks")  # type: ignore[assignment]
      store = InMemoryTaskStore()
  
      @server.list_tools()
      async def list_tools():
          return [
              Tool(
                  name="process_data",
                  description="Process data asynchronously",
                  inputSchema={
                      "type": "object",
                      "properties": {"input": {"type": "string"}},
                  },
                  execution=ToolExecution(taskSupport=TASK_REQUIRED),
              )
          ]
  
      @server.call_tool()
      async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
          ctx = server.request_context
          app = ctx.lifespan_context
          if name == "process_data" and ctx.experimental.is_task:
              # 1. Create task in store
              task_metadata = ctx.experimental.task_metadata
              assert task_metadata is not None
              task = await app.store.create_task(task_metadata)
  
              # 2. Create event to signal completion (for testing)
              done_event = Event()
              app.task_done_events[task.taskId] = done_event
  
              # 3. Define work function using task_execution for safety
              async def do_work():
                  async with task_execution(task.taskId, app.store) as task_ctx:
                      await task_ctx.update_status("Processing input...")
                      # Simulate work
                      input_value = arguments.get("input", "")
                      result_text = f"Processed: {input_value.upper()}"
                      await task_ctx.complete(CallToolResult(content=[TextContent(type="text", text=result_text)]))
                  # Signal completion
                  done_event.set()
  
              # 4. Spawn work in task group (from lifespan_context)
              app.task_group.start_soon(do_work)
  
              # 5. Return CreateTaskResult immediately
              return CreateTaskResult(task=task)
  
          raise NotImplementedError
  
      # Register task query handlers (delegate to store)
      @server.experimental.get_task()
      async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
          app = server.request_context.lifespan_context
          task = await app.store.get_task(request.params.taskId)
          assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
          return GetTaskResult(
              taskId=task.taskId,
              status=task.status,
              statusMessage=task.statusMessage,
              createdAt=task.createdAt,
              lastUpdatedAt=task.lastUpdatedAt,
              ttl=task.ttl,
              pollInterval=task.pollInterval,
          )
  
      @server.experimental.get_task_result()
      async def handle_get_task_result(
          request: GetTaskPayloadRequest,
      ) -> GetTaskPayloadResult:
          app = server.request_context.lifespan_context
          result = await app.store.get_result(request.params.taskId)
          assert result is not None, f"Test setup error: result for {request.params.taskId} should exist"
          assert isinstance(result, CallToolResult)
          # Return as GetTaskPayloadResult (which accepts extra fields)
          return GetTaskPayloadResult(**result.model_dump())
  
      @server.experimental.list_tasks()
      async def handle_list_tasks(request: ListTasksRequest) -> ListTasksResult:
          raise NotImplementedError
  
      # Set up client-server communication
      server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
      client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
  
      async def message_handler(
          message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
      ) -> None: ...  # pragma: no cover
  
      async def run_server(app_context: AppContext):
          async with ServerSession(
              client_to_server_receive,
              server_to_client_send,
              InitializationOptions(
                  server_name="test-server",
                  server_version="1.0.0",
                  capabilities=server.get_capabilities(
                      notification_options=NotificationOptions(),
                      experimental_capabilities={},
                  ),
              ),
          ) as server_session:
              async for message in server_session.incoming_messages:
                  await server._handle_message(message, server_session, app_context, raise_exceptions=False)
  
      async with anyio.create_task_group() as tg:
          # Create app context with task group and store
          app_context = AppContext(task_group=tg, store=store)
          tg.start_soon(run_server, app_context)
  
          async with ClientSession(
              server_to_client_receive,
              client_to_server_send,
              message_handler=message_handler,
          ) as client_session:
              await client_session.initialize()
  
              # === Step 1: Send task-augmented tool call ===
              create_result = await client_session.send_request(
                  ClientRequest(
                      CallToolRequest(
                          params=CallToolRequestParams(
                              name="process_data",
                              arguments={"input": "hello world"},
                              task=TaskMetadata(ttl=60000),
                          ),
                      )
                  ),
                  CreateTaskResult,
              )
  
              assert isinstance(create_result, CreateTaskResult)
              assert create_result.task.status == "working"
              task_id = create_result.task.taskId
  
              # === Step 2: Wait for task to complete ===
              await app_context.task_done_events[task_id].wait()
  
              task_status = await client_session.send_request(
                  ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId=task_id))),
                  GetTaskResult,
              )
  
              assert task_status.taskId == task_id
              assert task_status.status == "completed"
  
              # === Step 3: Retrieve the actual result ===
              task_result = await client_session.send_request(
                  ClientRequest(GetTaskPayloadRequest(params=GetTaskPayloadRequestParams(taskId=task_id))),
                  CallToolResult,
              )
  
              assert len(task_result.content) == 1
              content = task_result.content[0]
              assert isinstance(content, TextContent)
              assert content.text == "Processed: HELLO WORLD"
  
              tg.cancel_scope.cancel()
  
  
  @pytest.mark.anyio
  async def test_task_auto_fails_on_exception() -> None:
      """Test that task_execution automatically fails the task on unhandled exception."""
      # Note: We bypass the normal lifespan mechanism and pass context directly to _handle_message
      server: Server[AppContext, Any] = Server("test-tasks-failure")  # type: ignore[assignment]
      store = InMemoryTaskStore()
  
      @server.list_tools()
      async def list_tools():
          return [
              Tool(
                  name="failing_task",
                  description="A task that fails",
                  inputSchema={"type": "object", "properties": {}},
              )
          ]
  
      @server.call_tool()
      async def handle_call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent] | CreateTaskResult:
          ctx = server.request_context
          app = ctx.lifespan_context
          if name == "failing_task" and ctx.experimental.is_task:
              task_metadata = ctx.experimental.task_metadata
              assert task_metadata is not None
              task = await app.store.create_task(task_metadata)
  
              # Create event to signal completion (for testing)
              done_event = Event()
              app.task_done_events[task.taskId] = done_event
  
              async def do_failing_work():
                  async with task_execution(task.taskId, app.store) as task_ctx:
                      await task_ctx.update_status("About to fail...")
                      raise RuntimeError("Something went wrong!")
                      # Note: complete() is never called, but task_execution
                      # will automatically call fail() due to the exception
                  # This line is reached because task_execution suppresses the exception
                  done_event.set()
  
              app.task_group.start_soon(do_failing_work)
              return CreateTaskResult(task=task)
  
          raise NotImplementedError
  
      @server.experimental.get_task()
      async def handle_get_task(request: GetTaskRequest) -> GetTaskResult:
          app = server.request_context.lifespan_context
          task = await app.store.get_task(request.params.taskId)
          assert task is not None, f"Test setup error: task {request.params.taskId} should exist"
          return GetTaskResult(
              taskId=task.taskId,
              status=task.status,
              statusMessage=task.statusMessage,
              createdAt=task.createdAt,
              lastUpdatedAt=task.lastUpdatedAt,
              ttl=task.ttl,
              pollInterval=task.pollInterval,
          )
  
      # Set up streams
      server_to_client_send, server_to_client_receive = anyio.create_memory_object_stream[SessionMessage](10)
      client_to_server_send, client_to_server_receive = anyio.create_memory_object_stream[SessionMessage](10)
  
      async def message_handler(
          message: RequestResponder[ServerRequest, ClientResult] | ServerNotification | Exception,
      ) -> None: ...  # pragma: no cover
  
      async def run_server(app_context: AppContext):
          async with ServerSession(
              client_to_server_receive,
              server_to_client_send,
              InitializationOptions(
                  server_name="test-server",
                  server_version="1.0.0",
                  capabilities=server.get_capabilities(
                      notification_options=NotificationOptions(),
                      experimental_capabilities={},
                  ),
              ),
          ) as server_session:
              async for message in server_session.incoming_messages:
                  await server._handle_message(message, server_session, app_context, raise_exceptions=False)
  
      async with anyio.create_task_group() as tg:
          app_context = AppContext(task_group=tg, store=store)
          tg.start_soon(run_server, app_context)
  
          async with ClientSession(
              server_to_client_receive,
              client_to_server_send,
              message_handler=message_handler,
          ) as client_session:
              await client_session.initialize()
  
              # Send task request
              create_result = await client_session.send_request(
                  ClientRequest(
                      CallToolRequest(
                          params=CallToolRequestParams(
                              name="failing_task",
                              arguments={},
                              task=TaskMetadata(ttl=60000),
                          ),
                      )
                  ),
                  CreateTaskResult,
              )
  
              task_id = create_result.task.taskId
  
              # Wait for task to complete (even though it fails)
              await app_context.task_done_events[task_id].wait()
  
              # Check that task was auto-failed
              task_status = await client_session.send_request(
                  ClientRequest(GetTaskRequest(params=GetTaskRequestParams(taskId=task_id))),
                  GetTaskResult,
              )
  
              assert task_status.status == "failed"
              assert task_status.statusMessage == "Something went wrong!"
  
              tg.cancel_scope.cancel()
  