diff --git a/src/examples/server/simpleStreamableHttp.ts b/src/examples/server/simpleStreamableHttp.ts
index e3b754fa6..bd8bd1872 100644
--- a/src/examples/server/simpleStreamableHttp.ts
+++ b/src/examples/server/simpleStreamableHttp.ts
@@ -472,16 +472,16 @@ const getServer = () => {
             }
         },
         {
-            async createTask({ duration }, { taskStore, taskRequestedTtl }) {
+            async createTask({ duration }, { task }) {
                 // Create the task
-                const task = await taskStore.createTask({
-                    ttl: taskRequestedTtl
+                const newTask = await task.store.createTask({
+                    ttl: task.requestedTtl
                 });
 
                 // Simulate out-of-band work
                 (async () => {
                     await new Promise(resolve => setTimeout(resolve, duration));
-                    await taskStore.storeTaskResult(task.taskId, 'completed', {
+                    await task.store.storeTaskResult(newTask.taskId, 'completed', {
                         content: [
                             {
                                 type: 'text',
@@ -493,14 +493,14 @@ const getServer = () => {
 
                 // Return CreateTaskResult with the created task
                 return {
-                    task
+                    task: newTask
                 };
             },
-            async getTask(_args, { taskId, taskStore }) {
-                return await taskStore.getTask(taskId);
+            async getTask(_args, { task }) {
+                return await task.store.getTask(task.id);
             },
-            async getTaskResult(_args, { taskId, taskStore }) {
-                const result = await taskStore.getTaskResult(taskId);
+            async getTaskResult(_args, { task }) {
+                const result = await task.store.getTaskResult(task.id);
                 return result as CallToolResult;
             }
         }
diff --git a/src/experimental/tasks/interfaces.ts b/src/experimental/tasks/interfaces.ts
index 88e53028b..31a528dc4 100644
--- a/src/experimental/tasks/interfaces.ts
+++ b/src/experimental/tasks/interfaces.ts
@@ -19,7 +19,7 @@ import {
     Request
 } from '../../types.js';
 import { CreateTaskResult } from './types.js';
-import type { RequestHandlerExtra, RequestTaskStore } from '../../shared/protocol.js';
+import type { RequestHandlerExtra, TaskContext } from '../../shared/protocol.js';
 import type { ZodRawShapeCompat, AnySchema, ShapeOutput } from '../../server/zod-compat.js';
 
 // ============================================================================
@@ -27,11 +27,11 @@ import type { ZodRawShapeCompat, AnySchema, ShapeOutput } from '../../server/zod
 // ============================================================================
 
 /**
- * Extended handler extra with task store for task creation.
+ * Extended handler extra with task context for task creation.
  * @experimental
  */
 export interface CreateTaskRequestHandlerExtra extends RequestHandlerExtra<ServerRequest, ServerNotification> {
-    taskStore: RequestTaskStore;
+    task: TaskContext;
 }
 
 /**
@@ -39,8 +39,7 @@ export interface CreateTaskRequestHandlerExtra extends RequestHandlerExtra<Serve
  * @experimental
  */
 export interface TaskRequestHandlerExtra extends RequestHandlerExtra<ServerRequest, ServerNotification> {
-    taskId: string;
-    taskStore: RequestTaskStore;
+    task: TaskContext & { id: string };
 }
 
 /**
diff --git a/src/experimental/tasks/mcp-server.ts b/src/experimental/tasks/mcp-server.ts
index 506f3d72b..c553da7de 100644
--- a/src/experimental/tasks/mcp-server.ts
+++ b/src/experimental/tasks/mcp-server.ts
@@ -56,15 +56,15 @@ export class ExperimentalMcpServerTasks {
      *   execution: { taskSupport: 'required' }
      * }, {
      *   createTask: async (args, extra) => {
-     *     const task = await extra.taskStore.createTask({ ttl: 300000 });
-     *     startBackgroundWork(task.taskId, args);
-     *     return { task };
+     *     const newTask = await extra.task.store.createTask({ ttl: 300000 });
+     *     startBackgroundWork(newTask.taskId, args);
+     *     return { task: newTask };
      *   },
      *   getTask: async (args, extra) => {
-     *     return extra.taskStore.getTask(extra.taskId);
+     *     return extra.task.store.getTask(extra.task.id);
      *   },
      *   getTaskResult: async (args, extra) => {
-     *     return extra.taskStore.getTaskResult(extra.taskId);
+     *     return extra.task.store.getTaskResult(extra.task.id);
      *   }
      * });
      * ```
diff --git a/src/server/mcp.ts b/src/server/mcp.ts
index 7e61b4364..e9ff058b5 100644
--- a/src/server/mcp.ts
+++ b/src/server/mcp.ts
@@ -330,10 +330,10 @@ export class McpServer {
         const isTaskHandler = 'createTask' in handler;
 
         if (isTaskHandler) {
-            if (!extra.taskStore) {
+            if (!extra.task?.store) {
                 throw new Error('No task store provided.');
             }
-            const taskExtra = { ...extra, taskStore: extra.taskStore };
+            const taskExtra = { ...extra, task: extra.task };
 
             if (tool.inputSchema) {
                 const typedHandler = handler as ToolTaskHandler<ZodRawShapeCompat>;
@@ -365,14 +365,14 @@ export class McpServer {
         request: RequestT,
         extra: RequestHandlerExtra<ServerRequest, ServerNotification>
     ): Promise<CallToolResult> {
-        if (!extra.taskStore) {
+        if (!extra.task?.store) {
             throw new Error('No task store provided for task-capable tool.');
         }
 
         // Validate input and create task
         const args = await this.validateToolInput(tool, request.params.arguments, request.params.name);
         const handler = tool.handler as ToolTaskHandler<ZodRawShapeCompat | undefined>;
-        const taskExtra = { ...extra, taskStore: extra.taskStore };
+        const taskExtra = { ...extra, task: extra.task };
 
         const createTaskResult: CreateTaskResult = args // undefined only if tool.inputSchema is undefined
             ? await Promise.resolve((handler as ToolTaskHandler<ZodRawShapeCompat>).createTask(args, taskExtra))
@@ -386,7 +386,7 @@ export class McpServer {
 
         while (task.status !== 'completed' && task.status !== 'failed' && task.status !== 'cancelled') {
             await new Promise(resolve => setTimeout(resolve, pollInterval));
-            const updatedTask = await extra.taskStore.getTask(taskId);
+            const updatedTask = await extra.task.store.getTask(taskId);
             if (!updatedTask) {
                 throw new McpError(ErrorCode.InternalError, `Task ${taskId} not found during polling`);
             }
@@ -394,7 +394,7 @@ export class McpServer {
         }
 
         // Return the final result
-        return (await extra.taskStore.getTaskResult(taskId)) as CallToolResult;
+        return (await extra.task.store.getTaskResult(taskId)) as CallToolResult;
     }
 
     private _completionHandlerInitialized = false;
diff --git a/src/shared/protocol.ts b/src/shared/protocol.ts
index aa242a647..20fac2c0a 100644
--- a/src/shared/protocol.ts
+++ b/src/shared/protocol.ts
@@ -231,6 +231,18 @@ export interface RequestTaskStore {
     listTasks(cursor?: string): Promise<{ tasks: Task[]; nextCursor?: string }>;
 }
 
+/**
+ * Context for task-related operations in request handlers.
+ */
+export interface TaskContext {
+    /** The related task identifier (present when operating on existing task) */
+    id?: string;
+    /** Task store for managing task state */
+    store: RequestTaskStore;
+    /** Requested TTL in milliseconds (from client's task creation params) */
+    requestedTtl?: number;
+}
+
 /**
  * Extra data given to request handlers.
  */
@@ -261,11 +273,11 @@ export type RequestHandlerExtra<SendRequestT extends Request, SendNotificationT
      */
     requestId: RequestId;
 
-    taskId?: string;
-
-    taskStore?: RequestTaskStore;
-
-    taskRequestedTtl?: number | null;
+    /**
+     * Task context for task-related operations.
+     * Present when the server has task storage enabled.
+     */
+    task?: TaskContext;
 
     /**
      * The original HTTP request.
@@ -745,9 +757,13 @@ export abstract class Protocol<SendRequestT extends Request, SendNotificationT e
             authInfo: extra?.authInfo,
             requestId: request.id,
             requestInfo: extra?.requestInfo,
-            taskId: relatedTaskId,
-            taskStore: taskStore,
-            taskRequestedTtl: taskCreationParams?.ttl,
+            task: taskStore
+                ? {
+                      id: relatedTaskId,
+                      store: taskStore,
+                      requestedTtl: taskCreationParams?.ttl
+                  }
+                : undefined,
             closeSSEStream: extra?.closeSSEStream,
             closeStandaloneSSEStream: extra?.closeStandaloneSSEStream
         };
diff --git a/src/types.ts b/src/types.ts
index 6bec5190c..864e3f9b4 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -35,10 +35,9 @@ export const CursorSchema = z.string();
  */
 export const TaskCreationParamsSchema = z.looseObject({
     /**
-     * Time in milliseconds to keep task results available after completion.
-     * If null, the task has unlimited lifetime until manually cleaned up.
+     * Requested duration in milliseconds to retain task from creation.
      */
-    ttl: z.union([z.number(), z.null()]).optional(),
+    ttl: z.number().optional(),
 
     /**
      * Time in milliseconds to wait between task status requests.
diff --git a/test/client/index.test.ts b/test/client/index.test.ts
index 9735eb2ba..433b71239 100644
--- a/test/client/index.test.ts
+++ b/test/client/index.test.ts
@@ -2375,26 +2375,26 @@ describe('Task-based execution', () => {
                 },
                 {
                     async createTask(_args, extra) {
-                        const task = await extra.taskStore.createTask({
-                            ttl: extra.taskRequestedTtl
+                        const task = await extra.task!.store.createTask({
+                            ttl: extra.task?.requestedTtl
                         });
 
                         const result = {
                             content: [{ type: 'text', text: 'Tool executed successfully!' }]
                         };
-                        await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                        await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
 
                         return { task };
                     },
                     async getTask(_args, extra) {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task!.store.getTask(extra.task!.id);
                         if (!task) {
-                            throw new Error(`Task ${extra.taskId} not found`);
+                            throw new Error(`Task ${extra.task!.id} not found`);
                         }
                         return task;
                     },
                     async getTaskResult(_args, extra) {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task!.store.getTaskResult(extra.task!.id);
                         return result as { content: Array<{ type: 'text'; text: string }> };
                     }
                 }
@@ -2451,26 +2451,26 @@ describe('Task-based execution', () => {
                 },
                 {
                     async createTask(_args, extra) {
-                        const task = await extra.taskStore.createTask({
-                            ttl: extra.taskRequestedTtl
+                        const task = await extra.task!.store.createTask({
+                            ttl: extra.task?.requestedTtl
                         });
 
                         const result = {
                             content: [{ type: 'text', text: 'Success!' }]
                         };
-                        await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                        await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
 
                         return { task };
                     },
                     async getTask(_args, extra) {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task!.store.getTask(extra.task!.id);
                         if (!task) {
-                            throw new Error(`Task ${extra.taskId} not found`);
+                            throw new Error(`Task ${extra.task!.id} not found`);
                         }
                         return task;
                     },
                     async getTaskResult(_args, extra) {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task!.store.getTaskResult(extra.task!.id);
                         return result as { content: Array<{ type: 'text'; text: string }> };
                     }
                 }
@@ -2528,26 +2528,26 @@ describe('Task-based execution', () => {
                 },
                 {
                     async createTask(_args, extra) {
-                        const task = await extra.taskStore.createTask({
-                            ttl: extra.taskRequestedTtl
+                        const task = await extra.task!.store.createTask({
+                            ttl: extra.task?.requestedTtl
                         });
 
                         const result = {
                             content: [{ type: 'text', text: 'Result data!' }]
                         };
-                        await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                        await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
 
                         return { task };
                     },
                     async getTask(_args, extra) {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task!.store.getTask(extra.task!.id);
                         if (!task) {
-                            throw new Error(`Task ${extra.taskId} not found`);
+                            throw new Error(`Task ${extra.task!.id} not found`);
                         }
                         return task;
                     },
                     async getTaskResult(_args, extra) {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task!.store.getTaskResult(extra.task!.id);
                         return result as { content: Array<{ type: 'text'; text: string }> };
                     }
                 }
@@ -2609,26 +2609,26 @@ describe('Task-based execution', () => {
                 },
                 {
                     async createTask(_args, extra) {
-                        const task = await extra.taskStore.createTask({
-                            ttl: extra.taskRequestedTtl
+                        const task = await extra.task!.store.createTask({
+                            ttl: extra.task?.requestedTtl
                         });
 
                         const result = {
                             content: [{ type: 'text', text: 'Success!' }]
                         };
-                        await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                        await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
 
                         return { task };
                     },
                     async getTask(_args, extra) {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task!.store.getTask(extra.task!.id);
                         if (!task) {
-                            throw new Error(`Task ${extra.taskId} not found`);
+                            throw new Error(`Task ${extra.task!.id} not found`);
                         }
                         return task;
                     },
                     async getTaskResult(_args, extra) {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task!.store.getTaskResult(extra.task!.id);
                         return result as { content: Array<{ type: 'text'; text: string }> };
                     }
                 }
@@ -2712,11 +2712,11 @@ describe('Task-based execution', () => {
                 };
 
                 // Check if task creation is requested
-                if (request.params.task && extra.taskStore) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                if (request.params.task && extra.task!.store) {
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
-                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                    await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     // Return CreateTaskResult when task creation is requested
                     return { task };
                 }
@@ -2805,11 +2805,11 @@ describe('Task-based execution', () => {
                 };
 
                 // Check if task creation is requested
-                if (request.params.task && extra.taskStore) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                if (request.params.task && extra.task!.store) {
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
-                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                    await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     // Return CreateTaskResult when task creation is requested
                     return { task };
                 }
@@ -2897,11 +2897,11 @@ describe('Task-based execution', () => {
                 };
 
                 // Check if task creation is requested
-                if (request.params.task && extra.taskStore) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                if (request.params.task && extra.task!.store) {
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
-                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                    await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     // Return CreateTaskResult when task creation is requested
                     return { task };
                 }
@@ -2988,11 +2988,11 @@ describe('Task-based execution', () => {
                 };
 
                 // Check if task creation is requested
-                if (request.params.task && extra.taskStore) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                if (request.params.task && extra.task!.store) {
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
-                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                    await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     // Return CreateTaskResult when task creation is requested
                     return { task };
                 }
@@ -3094,26 +3094,26 @@ describe('Task-based execution', () => {
             },
             {
                 async createTask({ id }, extra) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
 
                     const result = {
                         content: [{ type: 'text', text: `Result for ${id || 'unknown'}` }]
                     };
-                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                    await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
 
                     return { task };
                 },
                 async getTask(_args, extra) {
-                    const task = await extra.taskStore.getTask(extra.taskId);
+                    const task = await extra.task!.store.getTask(extra.task!.id);
                     if (!task) {
-                        throw new Error(`Task ${extra.taskId} not found`);
+                        throw new Error(`Task ${extra.task!.id} not found`);
                     }
                     return task;
                 },
                 async getTaskResult(_args, extra) {
-                    const result = await extra.taskStore.getTaskResult(extra.taskId);
+                    const result = await extra.task!.store.getTaskResult(extra.task!.id);
                     return result as { content: Array<{ type: 'text'; text: string }> };
                 }
             }
@@ -3362,26 +3362,26 @@ test('should respect server task capabilities', async () => {
         },
         {
             async createTask(_args, extra) {
-                const task = await extra.taskStore.createTask({
-                    ttl: extra.taskRequestedTtl
+                const task = await extra.task!.store.createTask({
+                    ttl: extra.task?.requestedTtl
                 });
 
                 const result = {
                     content: [{ type: 'text', text: 'Success!' }]
                 };
-                await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
 
                 return { task };
             },
             async getTask(_args, extra) {
-                const task = await extra.taskStore.getTask(extra.taskId);
+                const task = await extra.task!.store.getTask(extra.task!.id);
                 if (!task) {
-                    throw new Error(`Task ${extra.taskId} not found`);
+                    throw new Error(`Task ${extra.task!.id} not found`);
                 }
                 return task;
             },
             async getTaskResult(_args, extra) {
-                const result = await extra.taskStore.getTaskResult(extra.taskId);
+                const result = await extra.task!.store.getTaskResult(extra.task!.id);
                 return result as { content: Array<{ type: 'text'; text: string }> };
             }
         }
diff --git a/test/experimental/tasks/stores/in-memory.test.ts b/test/experimental/tasks/stores/in-memory.test.ts
index ceef6c6d8..fd6bddc93 100644
--- a/test/experimental/tasks/stores/in-memory.test.ts
+++ b/test/experimental/tasks/stores/in-memory.test.ts
@@ -487,17 +487,16 @@ describe('InMemoryTaskStore', () => {
             expect(task).toBeNull();
         });
 
-        it('should support null TTL for unlimited lifetime', async () => {
-            // Test that null TTL means unlimited lifetime
-            const taskParams: TaskCreationParams = {
-                ttl: null
-            };
+        it('should support omitted TTL for unlimited lifetime', async () => {
+            // Test that omitting TTL means unlimited lifetime (server returns null)
+            // Per spec: clients omit ttl to let server decide, server returns null for unlimited
+            const taskParams: TaskCreationParams = {};
             const createdTask = await store.createTask(taskParams, 2222, {
                 method: 'tools/call',
                 params: {}
             });
 
-            // The returned task should have null TTL
+            // The returned task should have null TTL (unlimited)
             expect(createdTask.ttl).toBeNull();
 
             // Task should not be cleaned up even after a long time
diff --git a/test/experimental/tasks/task.test.ts b/test/experimental/tasks/task.test.ts
index 37e3938d2..de613a325 100644
--- a/test/experimental/tasks/task.test.ts
+++ b/test/experimental/tasks/task.test.ts
@@ -1,6 +1,7 @@
 import { describe, it, expect } from 'vitest';
 import { isTerminal } from '../../../src/experimental/tasks/interfaces.js';
 import type { Task } from '../../../src/types.js';
+import { TaskCreationParamsSchema } from '../../../src/types.js';
 
 describe('Task utility functions', () => {
     describe('isTerminal', () => {
@@ -115,3 +116,30 @@ describe('Task Schema Validation', () => {
         });
     });
 });
+
+describe('TaskCreationParams Schema Validation', () => {
+    it('should accept ttl as a number', () => {
+        const result = TaskCreationParamsSchema.safeParse({ ttl: 60000 });
+        expect(result.success).toBe(true);
+    });
+
+    it('should accept missing ttl (optional)', () => {
+        const result = TaskCreationParamsSchema.safeParse({});
+        expect(result.success).toBe(true);
+    });
+
+    it('should reject null ttl (not allowed in request, only response)', () => {
+        const result = TaskCreationParamsSchema.safeParse({ ttl: null });
+        expect(result.success).toBe(false);
+    });
+
+    it('should accept pollInterval as a number', () => {
+        const result = TaskCreationParamsSchema.safeParse({ pollInterval: 1000 });
+        expect(result.success).toBe(true);
+    });
+
+    it('should accept both ttl and pollInterval', () => {
+        const result = TaskCreationParamsSchema.safeParse({ ttl: 60000, pollInterval: 1000 });
+        expect(result.success).toBe(true);
+    });
+});
diff --git a/test/integration-tests/taskLifecycle.test.ts b/test/integration-tests/taskLifecycle.test.ts
index 629a61b66..35b71bfd5 100644
--- a/test/integration-tests/taskLifecycle.test.ts
+++ b/test/integration-tests/taskLifecycle.test.ts
@@ -64,7 +64,7 @@ describe('Task Lifecycle Integration Tests', () => {
             },
             {
                 async createTask({ duration, shouldFail }, extra) {
-                    const task = await extra.taskStore.createTask({
+                    const task = await extra.task.store.createTask({
                         ttl: 60000,
                         pollInterval: 100
                     });
@@ -75,12 +75,12 @@ describe('Task Lifecycle Integration Tests', () => {
 
                         try {
                             if (shouldFail) {
-                                await extra.taskStore.storeTaskResult(task.taskId, 'failed', {
+                                await extra.task.store.storeTaskResult(task.taskId, 'failed', {
                                     content: [{ type: 'text', text: 'Task failed as requested' }],
                                     isError: true
                                 });
                             } else {
-                                await extra.taskStore.storeTaskResult(task.taskId, 'completed', {
+                                await extra.task.store.storeTaskResult(task.taskId, 'completed', {
                                     content: [{ type: 'text', text: `Completed after ${duration}ms` }]
                                 });
                             }
@@ -92,14 +92,14 @@ describe('Task Lifecycle Integration Tests', () => {
                     return { task };
                 },
                 async getTask(_args, extra) {
-                    const task = await extra.taskStore.getTask(extra.taskId);
+                    const task = await extra.task.store.getTask(extra.task.id);
                     if (!task) {
-                        throw new Error(`Task ${extra.taskId} not found`);
+                        throw new Error(`Task ${extra.task.id} not found`);
                     }
                     return task;
                 },
                 async getTaskResult(_args, extra) {
-                    const result = await extra.taskStore.getTaskResult(extra.taskId);
+                    const result = await extra.task.store.getTaskResult(extra.task.id);
                     return result as { content: Array<{ type: 'text'; text: string }> };
                 }
             }
@@ -117,7 +117,7 @@ describe('Task Lifecycle Integration Tests', () => {
             },
             {
                 async createTask({ userName }, extra) {
-                    const task = await extra.taskStore.createTask({
+                    const task = await extra.task.store.createTask({
                         ttl: 60000,
                         pollInterval: 100
                     });
@@ -153,7 +153,7 @@ describe('Task Lifecycle Integration Tests', () => {
                                     ? elicitationResult.content.userName
                                     : 'Unknown';
                             try {
-                                await extra.taskStore.storeTaskResult(task.taskId, 'completed', {
+                                await extra.task.store.storeTaskResult(task.taskId, 'completed', {
                                     content: [{ type: 'text', text: `Hello, ${name}!` }]
                                 });
                             } catch {
@@ -162,7 +162,7 @@ describe('Task Lifecycle Integration Tests', () => {
                         } else {
                             // Complete immediately if userName was provided
                             try {
-                                await extra.taskStore.storeTaskResult(task.taskId, 'completed', {
+                                await extra.task.store.storeTaskResult(task.taskId, 'completed', {
                                     content: [{ type: 'text', text: `Hello, ${userName}!` }]
                                 });
                             } catch {
@@ -174,14 +174,14 @@ describe('Task Lifecycle Integration Tests', () => {
                     return { task };
                 },
                 async getTask(_args, extra) {
-                    const task = await extra.taskStore.getTask(extra.taskId);
+                    const task = await extra.task.store.getTask(extra.task.id);
                     if (!task) {
-                        throw new Error(`Task ${extra.taskId} not found`);
+                        throw new Error(`Task ${extra.task.id} not found`);
                     }
                     return task;
                 },
                 async getTaskResult(_args, extra) {
-                    const result = await extra.taskStore.getTaskResult(extra.taskId);
+                    const result = await extra.task.store.getTaskResult(extra.task.id);
                     return result as { content: Array<{ type: 'text'; text: string }> };
                 }
             }
@@ -421,7 +421,7 @@ describe('Task Lifecycle Integration Tests', () => {
                 },
                 {
                     async createTask({ requestCount }, extra) {
-                        const task = await extra.taskStore.createTask({
+                        const task = await extra.task.store.createTask({
                             ttl: 60000,
                             pollInterval: 100
                         });
@@ -460,7 +460,7 @@ describe('Task Lifecycle Integration Tests', () => {
 
                             // Complete with all responses
                             try {
-                                await extra.taskStore.storeTaskResult(task.taskId, 'completed', {
+                                await extra.task.store.storeTaskResult(task.taskId, 'completed', {
                                     content: [{ type: 'text', text: `Received responses: ${responses.join(', ')}` }]
                                 });
                             } catch {
@@ -471,14 +471,14 @@ describe('Task Lifecycle Integration Tests', () => {
                         return { task };
                     },
                     async getTask(_args, extra) {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) {
-                            throw new Error(`Task ${extra.taskId} not found`);
+                            throw new Error(`Task ${extra.task.id} not found`);
                         }
                         return task;
                     },
                     async getTaskResult(_args, extra) {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task.store.getTaskResult(extra.task.id);
                         return result as { content: Array<{ type: 'text'; text: string }> };
                     }
                 }
@@ -910,7 +910,7 @@ describe('Task Lifecycle Integration Tests', () => {
                 },
                 {
                     async createTask({ messageCount }, extra) {
-                        const task = await extra.taskStore.createTask({
+                        const task = await extra.task.store.createTask({
                             ttl: 60000,
                             pollInterval: 100
                         });
@@ -960,14 +960,14 @@ describe('Task Lifecycle Integration Tests', () => {
                         return { task };
                     },
                     async getTask(_args, extra) {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) {
-                            throw new Error(`Task ${extra.taskId} not found`);
+                            throw new Error(`Task ${extra.task.id} not found`);
                         }
                         return task;
                     },
                     async getTaskResult(_args, extra) {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task.store.getTaskResult(extra.task.id);
                         return result as { content: Array<{ type: 'text'; text: string }> };
                     }
                 }
@@ -1108,7 +1108,7 @@ describe('Task Lifecycle Integration Tests', () => {
                 },
                 {
                     async createTask({ messageCount, delayBetweenMessages }, extra) {
-                        const task = await extra.taskStore.createTask({
+                        const task = await extra.task.store.createTask({
                             ttl: 60000,
                             pollInterval: 100
                         });
@@ -1154,7 +1154,7 @@ describe('Task Lifecycle Integration Tests', () => {
 
                                 // Complete with all responses
                                 try {
-                                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', {
+                                    await extra.task.store.storeTaskResult(task.taskId, 'completed', {
                                         content: [{ type: 'text', text: `Received all responses: ${responses.join(', ')}` }]
                                     });
                                 } catch {
@@ -1163,7 +1163,7 @@ describe('Task Lifecycle Integration Tests', () => {
                             } catch (error) {
                                 // Handle errors
                                 try {
-                                    await extra.taskStore.storeTaskResult(task.taskId, 'failed', {
+                                    await extra.task.store.storeTaskResult(task.taskId, 'failed', {
                                         content: [{ type: 'text', text: `Error: ${error}` }],
                                         isError: true
                                     });
@@ -1176,14 +1176,14 @@ describe('Task Lifecycle Integration Tests', () => {
                         return { task };
                     },
                     async getTask(_args, extra) {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) {
-                            throw new Error(`Task ${extra.taskId} not found`);
+                            throw new Error(`Task ${extra.task.id} not found`);
                         }
                         return task;
                     },
                     async getTaskResult(_args, extra) {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task.store.getTaskResult(extra.task.id);
                         return result as { content: Array<{ type: 'text'; text: string }> };
                     }
                 }
@@ -1324,7 +1324,7 @@ describe('Task Lifecycle Integration Tests', () => {
                 },
                 {
                     async createTask({ messageCount }, extra) {
-                        const task = await extra.taskStore.createTask({
+                        const task = await extra.task.store.createTask({
                             ttl: 60000,
                             pollInterval: 100
                         });
@@ -1363,7 +1363,7 @@ describe('Task Lifecycle Integration Tests', () => {
 
                                 // Complete the task after all messages are queued
                                 try {
-                                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', {
+                                    await extra.task.store.storeTaskResult(task.taskId, 'completed', {
                                         content: [{ type: 'text', text: 'Task completed quickly' }]
                                     });
                                 } catch {
@@ -1372,7 +1372,7 @@ describe('Task Lifecycle Integration Tests', () => {
                             } catch (error) {
                                 // Handle errors
                                 try {
-                                    await extra.taskStore.storeTaskResult(task.taskId, 'failed', {
+                                    await extra.task.store.storeTaskResult(task.taskId, 'failed', {
                                         content: [{ type: 'text', text: `Error: ${error}` }],
                                         isError: true
                                     });
@@ -1385,14 +1385,14 @@ describe('Task Lifecycle Integration Tests', () => {
                         return { task };
                     },
                     async getTask(_args, extra) {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) {
-                            throw new Error(`Task ${extra.taskId} not found`);
+                            throw new Error(`Task ${extra.task.id} not found`);
                         }
                         return task;
                     },
                     async getTaskResult(_args, extra) {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task.store.getTaskResult(extra.task.id);
                         return result as { content: Array<{ type: 'text'; text: string }> };
                     }
                 }
diff --git a/test/server/index.test.ts b/test/server/index.test.ts
index e434e57fc..7210b26d1 100644
--- a/test/server/index.test.ts
+++ b/test/server/index.test.ts
@@ -2249,8 +2249,8 @@ describe('Task-based execution', () => {
             },
             {
                 async createTask(_args, extra) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
 
                     // Simulate some async work
@@ -2259,20 +2259,20 @@ describe('Task-based execution', () => {
                         const result = {
                             content: [{ type: 'text', text: 'Tool executed successfully!' }]
                         };
-                        await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                        await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     })();
 
                     return { task };
                 },
                 async getTask(_args, extra) {
-                    const task = await extra.taskStore.getTask(extra.taskId);
+                    const task = await extra.task!.store.getTask(extra.task!.id);
                     if (!task) {
-                        throw new Error(`Task ${extra.taskId} not found`);
+                        throw new Error(`Task ${extra.task!.id} not found`);
                     }
                     return task;
                 },
                 async getTaskResult(_args, extra) {
-                    const result = await extra.taskStore.getTaskResult(extra.taskId);
+                    const result = await extra.task!.store.getTaskResult(extra.task!.id);
                     return result as { content: Array<{ type: 'text'; text: string }> };
                 }
             }
@@ -2444,9 +2444,9 @@ describe('Task-based execution', () => {
             let taskId: string | undefined;
 
             // Check if task creation is requested
-            if (request.params.task && extra.taskStore) {
-                const createdTask = await extra.taskStore.createTask({
-                    ttl: extra.taskRequestedTtl
+            if (request.params.task && extra.task!.store) {
+                const createdTask = await extra.task!.store.createTask({
+                    ttl: extra.task?.requestedTtl
                 });
                 taskId = createdTask.taskId;
             }
@@ -2471,8 +2471,8 @@ describe('Task-based execution', () => {
             },
             {
                 async createTask(_args, extra) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
 
                     // Perform async work that makes a nested request
@@ -2504,20 +2504,20 @@ describe('Task-based execution', () => {
                                 }
                             ]
                         };
-                        await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                        await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     })();
 
                     return { task };
                 },
                 async getTask(_args, extra) {
-                    const task = await extra.taskStore.getTask(extra.taskId);
+                    const task = await extra.task!.store.getTask(extra.task!.id);
                     if (!task) {
-                        throw new Error(`Task ${extra.taskId} not found`);
+                        throw new Error(`Task ${extra.task!.id} not found`);
                     }
                     return task;
                 },
                 async getTaskResult(_args, extra) {
-                    const result = await extra.taskStore.getTaskResult(extra.taskId);
+                    const result = await extra.task!.store.getTaskResult(extra.task!.id);
                     return result as { content: Array<{ type: 'text'; text: string }> };
                 }
             }
@@ -2601,11 +2601,11 @@ describe('Task-based execution', () => {
                 };
 
                 // Check if task creation is requested
-                if (request.params.task && extra.taskStore) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                if (request.params.task && extra.task!.store) {
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
-                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                    await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     // Return CreateTaskResult when task creation is requested
                     return { task };
                 }
@@ -2682,11 +2682,11 @@ describe('Task-based execution', () => {
                 };
 
                 // Check if task creation is requested
-                if (request.params.task && extra.taskStore) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                if (request.params.task && extra.task!.store) {
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
-                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                    await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     // Return CreateTaskResult when task creation is requested
                     return { task };
                 }
@@ -2761,11 +2761,11 @@ describe('Task-based execution', () => {
                 };
 
                 // Check if task creation is requested
-                if (request.params.task && extra.taskStore) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                if (request.params.task && extra.task!.store) {
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
-                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                    await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     // Return CreateTaskResult when task creation is requested
                     return { task };
                 }
@@ -2842,11 +2842,11 @@ describe('Task-based execution', () => {
                 };
 
                 // Check if task creation is requested
-                if (request.params.task && extra.taskStore) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                if (request.params.task && extra.task!.store) {
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
-                    await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                    await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     // Return CreateTaskResult when task creation is requested
                     return { task };
                 }
@@ -2950,8 +2950,8 @@ describe('Task-based execution', () => {
             },
             {
                 async createTask({ delay, taskNum }, extra) {
-                    const task = await extra.taskStore.createTask({
-                        ttl: extra.taskRequestedTtl
+                    const task = await extra.task!.store.createTask({
+                        ttl: extra.task?.requestedTtl
                     });
 
                     // Simulate async work
@@ -2960,20 +2960,20 @@ describe('Task-based execution', () => {
                         const result = {
                             content: [{ type: 'text', text: `Completed task ${taskNum || 'unknown'}` }]
                         };
-                        await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+                        await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
                     })();
 
                     return { task };
                 },
                 async getTask(_args, extra) {
-                    const task = await extra.taskStore.getTask(extra.taskId);
+                    const task = await extra.task!.store.getTask(extra.task!.id);
                     if (!task) {
-                        throw new Error(`Task ${extra.taskId} not found`);
+                        throw new Error(`Task ${extra.task!.id} not found`);
                     }
                     return task;
                 },
                 async getTaskResult(_args, extra) {
-                    const result = await extra.taskStore.getTaskResult(extra.taskId);
+                    const result = await extra.task!.store.getTaskResult(extra.task!.id);
                     return result as { content: Array<{ type: 'text'; text: string }> };
                 }
             }
@@ -3185,11 +3185,11 @@ test('should respect client task capabilities', async () => {
         };
 
         // Check if task creation is requested
-        if (request.params.task && extra.taskStore) {
-            const task = await extra.taskStore.createTask({
-                ttl: extra.taskRequestedTtl
+        if (request.params.task && extra.task!.store) {
+            const task = await extra.task!.store.createTask({
+                ttl: extra.task?.requestedTtl
             });
-            await extra.taskStore.storeTaskResult(task.taskId, 'completed', result);
+            await extra.task!.store.storeTaskResult(task.taskId, 'completed', result);
             // Return CreateTaskResult when task creation is requested
             return { task };
         }
diff --git a/test/server/mcp.test.ts b/test/server/mcp.test.ts
index f6c2124e1..be3f5aecc 100644
--- a/test/server/mcp.test.ts
+++ b/test/server/mcp.test.ts
@@ -1903,16 +1903,16 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                 },
                 {
                     createTask: async (_args, extra) => {
-                        const task = await extra.taskStore.createTask({ ttl: 60000 });
+                        const task = await extra.task.store.createTask({ ttl: 60000 });
                         return { task };
                     },
                     getTask: async (_args, extra) => {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) throw new Error('Task not found');
                         return task;
                     },
                     getTaskResult: async (_args, extra) => {
-                        return (await extra.taskStore.getTaskResult(extra.taskId)) as CallToolResult;
+                        return (await extra.task.store.getTaskResult(extra.task.id)) as CallToolResult;
                     }
                 }
             );
@@ -1972,16 +1972,16 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                 },
                 {
                     createTask: async (_args, extra) => {
-                        const task = await extra.taskStore.createTask({ ttl: 60000 });
+                        const task = await extra.task.store.createTask({ ttl: 60000 });
                         return { task };
                     },
                     getTask: async (_args, extra) => {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) throw new Error('Task not found');
                         return task;
                     },
                     getTaskResult: async (_args, extra) => {
-                        return (await extra.taskStore.getTaskResult(extra.taskId)) as CallToolResult;
+                        return (await extra.task.store.getTaskResult(extra.task.id)) as CallToolResult;
                     }
                 }
             );
@@ -6302,10 +6302,10 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                 },
                 {
                     createTask: async ({ input }, extra) => {
-                        const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
+                        const task = await extra.task.store.createTask({ ttl: 60000, pollInterval: 100 });
 
                         // Capture taskStore for use in setTimeout
-                        const store = extra.taskStore;
+                        const store = extra.task.store;
 
                         // Simulate async work
                         setTimeout(async () => {
@@ -6317,14 +6317,14 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                         return { task };
                     },
                     getTask: async (_args, extra) => {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) {
                             throw new Error('Task not found');
                         }
                         return task;
                     },
                     getTaskResult: async (_input, extra) => {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task.store.getTaskResult(extra.task.id);
                         return result as CallToolResult;
                     }
                 }
@@ -6407,10 +6407,10 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                 },
                 {
                     createTask: async ({ value }, extra) => {
-                        const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
+                        const task = await extra.task.store.createTask({ ttl: 60000, pollInterval: 100 });
 
                         // Capture taskStore for use in setTimeout
-                        const store = extra.taskStore;
+                        const store = extra.task.store;
 
                         // Simulate async work
                         setTimeout(async () => {
@@ -6423,14 +6423,14 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                         return { task };
                     },
                     getTask: async (_args, extra) => {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) {
                             throw new Error('Task not found');
                         }
                         return task;
                     },
                     getTaskResult: async (_value, extra) => {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task.store.getTaskResult(extra.task.id);
                         return result as CallToolResult;
                     }
                 }
@@ -6515,10 +6515,10 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                 },
                 {
                     createTask: async ({ data }, extra) => {
-                        const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
+                        const task = await extra.task.store.createTask({ ttl: 60000, pollInterval: 100 });
 
                         // Capture taskStore for use in setTimeout
-                        const store = extra.taskStore;
+                        const store = extra.task.store;
 
                         // Simulate async work
                         setTimeout(async () => {
@@ -6531,14 +6531,14 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                         return { task };
                     },
                     getTask: async (_args, extra) => {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) {
                             throw new Error('Task not found');
                         }
                         return task;
                     },
                     getTaskResult: async (_data, extra) => {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task.store.getTaskResult(extra.task.id);
                         return result as CallToolResult;
                     }
                 }
@@ -6632,10 +6632,10 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                 },
                 {
                     createTask: async extra => {
-                        const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
+                        const task = await extra.task.store.createTask({ ttl: 60000, pollInterval: 100 });
 
                         // Capture taskStore for use in setTimeout
-                        const store = extra.taskStore;
+                        const store = extra.task.store;
 
                         // Simulate async failure
                         setTimeout(async () => {
@@ -6649,14 +6649,14 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                         return { task };
                     },
                     getTask: async extra => {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) {
                             throw new Error('Task not found');
                         }
                         return task;
                     },
                     getTaskResult: async extra => {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task.store.getTaskResult(extra.task.id);
                         return result as CallToolResult;
                     }
                 }
@@ -6738,10 +6738,10 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                 },
                 {
                     createTask: async extra => {
-                        const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
+                        const task = await extra.task.store.createTask({ ttl: 60000, pollInterval: 100 });
 
                         // Capture taskStore for use in setTimeout
-                        const store = extra.taskStore;
+                        const store = extra.task.store;
 
                         // Simulate async cancellation
                         setTimeout(async () => {
@@ -6752,14 +6752,14 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                         return { task };
                     },
                     getTask: async extra => {
-                        const task = await extra.taskStore.getTask(extra.taskId);
+                        const task = await extra.task.store.getTask(extra.task.id);
                         if (!task) {
                             throw new Error('Task not found');
                         }
                         return task;
                     },
                     getTaskResult: async extra => {
-                        const result = await extra.taskStore.getTaskResult(extra.taskId);
+                        const result = await extra.task.store.getTaskResult(extra.task.id);
                         return result as CallToolResult;
                     }
                 }
@@ -6825,18 +6825,18 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
                     },
                     {
                         createTask: async (_args, extra) => {
-                            const task = await extra.taskStore.createTask({ ttl: 60000, pollInterval: 100 });
+                            const task = await extra.task.store.createTask({ ttl: 60000, pollInterval: 100 });
                             return { task };
                         },
                         getTask: async (_args, extra) => {
-                            const task = await extra.taskStore.getTask(extra.taskId);
+                            const task = await extra.task.store.getTask(extra.task.id);
                             if (!task) {
                                 throw new Error('Task not found');
                             }
                             return task;
                         },
                         getTaskResult: async (_args, extra) => {
-                            const result = await extra.taskStore.getTaskResult(extra.taskId);
+                            const result = await extra.task.store.getTaskResult(extra.task.id);
                             return result as CallToolResult;
                         }
                     }
diff --git a/test/shared/protocol.test.ts b/test/shared/protocol.test.ts
index 886dcbb21..02b64e37d 100644
--- a/test/shared/protocol.test.ts
+++ b/test/shared/protocol.test.ts
@@ -2480,12 +2480,12 @@ describe('Progress notification support for tasks', () => {
 
         // Set up a request handler that will complete the task
         protocol.setRequestHandler(CallToolRequestSchema, async (request, extra) => {
-            if (extra.taskStore) {
-                const task = await extra.taskStore.createTask({ ttl: 60000 });
+            if (extra.task?.store) {
+                const task = await extra.task?.store.createTask({ ttl: 60000 });
 
                 // Simulate async work then complete the task
                 setTimeout(async () => {
-                    await extra.taskStore!.storeTaskResult(task.taskId, 'completed', {
+                    await extra.task?.store!.storeTaskResult(task.taskId, 'completed', {
                         content: [{ type: 'text', text: 'Done' }]
                     });
                 }, 50);