import { ZodType, z } from 'zod';
import {
    CallToolRequestSchema,
    ClientCapabilities,
    ErrorCode,
    JSONRPCMessage,
    McpError,
    RELATED_TASK_META_KEY,
    RequestId,
    ServerCapabilities,
    Task,
    TaskCreationParams,
    type Request,
    type Notification,
    type Result
} from '../../src/types.js';
import { Protocol, mergeCapabilities } from '../../src/shared/protocol.js';
import { Transport, TransportSendOptions } from '../../src/shared/transport.js';
import { TaskStore, TaskMessageQueue, QueuedMessage, QueuedNotification, QueuedRequest } from '../../src/experimental/tasks/interfaces.js';
import { MockInstance, vi } from 'vitest';
import { JSONRPCResultResponse, JSONRPCRequest, JSONRPCErrorResponse } from '../../src/types.js';
import { ErrorMessage, ResponseMessage, toArrayAsync } from '../../src/shared/responseMessage.js';
import { InMemoryTaskMessageQueue } from '../../src/experimental/tasks/stores/in-memory.js';

// Type helper for accessing private/protected Protocol properties in tests
interface TestProtocol {
    _taskMessageQueue?: TaskMessageQueue;
    _requestResolvers: Map<RequestId, (response: JSONRPCResultResponse | Error) => void>;
    _responseHandlers: Map<RequestId, (response: JSONRPCResultResponse | Error) => void>;
    _taskProgressTokens: Map<string, number>;
    _clearTaskQueue: (taskId: string, sessionId?: string) => Promise<void>;
    requestTaskStore: (request: Request, authInfo: unknown) => TaskStore;
    // Protected task methods (exposed for testing)
    listTasks: (params?: { cursor?: string }) => Promise<{ tasks: Task[]; nextCursor?: string }>;
    cancelTask: (params: { taskId: string }) => Promise<Result>;
    requestStream: <T extends Result>(request: Request, schema: ZodType<T>, options?: unknown) => AsyncGenerator<ResponseMessage<T>>;
}

// Mock Transport class
class MockTransport implements Transport {
    onclose?: () => void;
    onerror?: (error: Error) => void;
    onmessage?: (message: unknown) => void;

    async start(): Promise<void> {}
    async close(): Promise<void> {
        this.onclose?.();
    }
    async send(_message: JSONRPCMessage, _options?: TransportSendOptions): Promise<void> {}
}

function createMockTaskStore(options?: {
    onStatus?: (status: Task['status']) => void;
    onList?: () => void;
}): TaskStore & { [K in keyof TaskStore]: MockInstance } {
    const tasks: Record<string, Task & { result?: Result }> = {};
    return {
        createTask: vi.fn((taskParams: TaskCreationParams, _1: RequestId, _2: Request) => {
            // Generate a unique task ID
            const taskId = `test-task-${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
            const createdAt = new Date().toISOString();
            const task = (tasks[taskId] = {
                taskId,
                status: 'working',
                ttl: taskParams.ttl ?? null,
                createdAt,
                lastUpdatedAt: createdAt,
                pollInterval: taskParams.pollInterval ?? 1000
            });
            options?.onStatus?.('working');
            return Promise.resolve(task);
        }),
        getTask: vi.fn((taskId: string) => {
            return Promise.resolve(tasks[taskId] ?? null);
        }),
        updateTaskStatus: vi.fn((taskId, status, statusMessage) => {
            const task = tasks[taskId];
            if (task) {
                task.status = status;
                task.statusMessage = statusMessage;
                options?.onStatus?.(task.status);
            }
            return Promise.resolve();
        }),
        storeTaskResult: vi.fn((taskId: string, status: 'completed' | 'failed', result: Result) => {
            const task = tasks[taskId];
            if (task) {
                task.status = status;
                task.result = result;
                options?.onStatus?.(status);
            }
            return Promise.resolve();
        }),
        getTaskResult: vi.fn((taskId: string) => {
            const task = tasks[taskId];
            if (task?.result) {
                return Promise.resolve(task.result);
            }
            throw new Error('Task result not found');
        }),
        listTasks: vi.fn(() => {
            const result = {
                tasks: Object.values(tasks)
            };
            options?.onList?.();
            return Promise.resolve(result);
        })
    };
}

function createLatch() {
    let latch = false;
    const waitForLatch = async () => {
        while (!latch) {
            await new Promise(resolve => setTimeout(resolve, 0));
        }
    };

    return {
        releaseLatch: () => {
            latch = true;
        },
        waitForLatch
    };
}

function assertErrorResponse(o: ResponseMessage<Result>): asserts o is ErrorMessage {
    expect(o.type).toBe('error');
}

function assertQueuedNotification(o?: QueuedMessage): asserts o is QueuedNotification {
    expect(o).toBeDefined();
    expect(o?.type).toBe('notification');
}

function assertQueuedRequest(o?: QueuedMessage): asserts o is QueuedRequest {
    expect(o).toBeDefined();
    expect(o?.type).toBe('request');
}

describe('protocol tests', () => {
    let protocol: Protocol<Request, Notification, Result>;
    let transport: MockTransport;
    let sendSpy: MockInstance;

    beforeEach(() => {
        transport = new MockTransport();
        sendSpy = vi.spyOn(transport, 'send');
        protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })();
    });

    test('should throw a timeout error if the request exceeds the timeout', async () => {
        await protocol.connect(transport);
        const request = { method: 'example', params: {} };
        try {
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            await protocol.request(request, mockSchema, {
                timeout: 0
            });
        } catch (error) {
            expect(error).toBeInstanceOf(McpError);
            if (error instanceof McpError) {
                expect(error.code).toBe(ErrorCode.RequestTimeout);
            }
        }
    });

    test('should invoke onclose when the connection is closed', async () => {
        const oncloseMock = vi.fn();
        protocol.onclose = oncloseMock;
        await protocol.connect(transport);
        await transport.close();
        expect(oncloseMock).toHaveBeenCalled();
    });

    test('should not overwrite existing hooks when connecting transports', async () => {
        const oncloseMock = vi.fn();
        const onerrorMock = vi.fn();
        const onmessageMock = vi.fn();
        transport.onclose = oncloseMock;
        transport.onerror = onerrorMock;
        transport.onmessage = onmessageMock;
        await protocol.connect(transport);
        transport.onclose();
        transport.onerror(new Error());
        transport.onmessage('');
        expect(oncloseMock).toHaveBeenCalled();
        expect(onerrorMock).toHaveBeenCalled();
        expect(onmessageMock).toHaveBeenCalled();
    });

    describe('_meta preservation with onprogress', () => {
        test('should preserve existing _meta when adding progressToken', async () => {
            await protocol.connect(transport);
            const request = {
                method: 'example',
                params: {
                    data: 'test',
                    _meta: {
                        customField: 'customValue',
                        anotherField: 123
                    }
                }
            };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            const onProgressMock = vi.fn();

            // Start request but don't await - we're testing the sent message
            void protocol
                .request(request, mockSchema, {
                    onprogress: onProgressMock
                })
                .catch(() => {
                    // May not complete, ignore error
                });

            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    method: 'example',
                    params: {
                        data: 'test',
                        _meta: {
                            customField: 'customValue',
                            anotherField: 123,
                            progressToken: expect.any(Number)
                        }
                    },
                    jsonrpc: '2.0',
                    id: expect.any(Number)
                }),
                expect.any(Object)
            );
        });

        test('should create _meta with progressToken when no _meta exists', async () => {
            await protocol.connect(transport);
            const request = {
                method: 'example',
                params: {
                    data: 'test'
                }
            };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            const onProgressMock = vi.fn();

            // Start request but don't await - we're testing the sent message
            void protocol
                .request(request, mockSchema, {
                    onprogress: onProgressMock
                })
                .catch(() => {
                    // May not complete, ignore error
                });

            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    method: 'example',
                    params: {
                        data: 'test',
                        _meta: {
                            progressToken: expect.any(Number)
                        }
                    },
                    jsonrpc: '2.0',
                    id: expect.any(Number)
                }),
                expect.any(Object)
            );
        });

        test('should not modify _meta when onprogress is not provided', async () => {
            await protocol.connect(transport);
            const request = {
                method: 'example',
                params: {
                    data: 'test',
                    _meta: {
                        customField: 'customValue'
                    }
                }
            };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });

            // Start request but don't await - we're testing the sent message
            void protocol.request(request, mockSchema).catch(() => {
                // May not complete, ignore error
            });

            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    method: 'example',
                    params: {
                        data: 'test',
                        _meta: {
                            customField: 'customValue'
                        }
                    },
                    jsonrpc: '2.0',
                    id: expect.any(Number)
                }),
                expect.any(Object)
            );
        });

        test('should handle params being undefined with onprogress', async () => {
            await protocol.connect(transport);
            const request = {
                method: 'example'
            };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            const onProgressMock = vi.fn();

            // Start request but don't await - we're testing the sent message
            void protocol
                .request(request, mockSchema, {
                    onprogress: onProgressMock
                })
                .catch(() => {
                    // May not complete, ignore error
                });

            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    method: 'example',
                    params: {
                        _meta: {
                            progressToken: expect.any(Number)
                        }
                    },
                    jsonrpc: '2.0',
                    id: expect.any(Number)
                }),
                expect.any(Object)
            );
        });
    });

    describe('progress notification timeout behavior', () => {
        beforeEach(() => {
            vi.useFakeTimers();
        });
        afterEach(() => {
            vi.useRealTimers();
        });

        test('should not reset timeout when resetTimeoutOnProgress is false', async () => {
            await protocol.connect(transport);
            const request = { method: 'example', params: {} };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            const onProgressMock = vi.fn();
            const requestPromise = protocol.request(request, mockSchema, {
                timeout: 1000,
                resetTimeoutOnProgress: false,
                onprogress: onProgressMock
            });

            vi.advanceTimersByTime(800);

            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/progress',
                    params: {
                        progressToken: 0,
                        progress: 50,
                        total: 100
                    }
                });
            }
            await Promise.resolve();

            expect(onProgressMock).toHaveBeenCalledWith({
                progress: 50,
                total: 100
            });

            vi.advanceTimersByTime(201);

            await expect(requestPromise).rejects.toThrow('Request timed out');
        });

        test('should reset timeout when progress notification is received', async () => {
            await protocol.connect(transport);
            const request = { method: 'example', params: {} };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            const onProgressMock = vi.fn();
            const requestPromise = protocol.request(request, mockSchema, {
                timeout: 1000,
                resetTimeoutOnProgress: true,
                onprogress: onProgressMock
            });
            vi.advanceTimersByTime(800);
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/progress',
                    params: {
                        progressToken: 0,
                        progress: 50,
                        total: 100
                    }
                });
            }
            await Promise.resolve();
            expect(onProgressMock).toHaveBeenCalledWith({
                progress: 50,
                total: 100
            });
            vi.advanceTimersByTime(800);
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    id: 0,
                    result: { result: 'success' }
                });
            }
            await Promise.resolve();
            await expect(requestPromise).resolves.toEqual({ result: 'success' });
        });

        test('should respect maxTotalTimeout', async () => {
            await protocol.connect(transport);
            const request = { method: 'example', params: {} };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            const onProgressMock = vi.fn();
            const requestPromise = protocol.request(request, mockSchema, {
                timeout: 1000,
                maxTotalTimeout: 150,
                resetTimeoutOnProgress: true,
                onprogress: onProgressMock
            });

            // First progress notification should work
            vi.advanceTimersByTime(80);
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/progress',
                    params: {
                        progressToken: 0,
                        progress: 50,
                        total: 100
                    }
                });
            }
            await Promise.resolve();
            expect(onProgressMock).toHaveBeenCalledWith({
                progress: 50,
                total: 100
            });
            vi.advanceTimersByTime(80);
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/progress',
                    params: {
                        progressToken: 0,
                        progress: 75,
                        total: 100
                    }
                });
            }
            await expect(requestPromise).rejects.toThrow('Maximum total timeout exceeded');
            expect(onProgressMock).toHaveBeenCalledTimes(1);
        });

        test('should timeout if no progress received within timeout period', async () => {
            await protocol.connect(transport);
            const request = { method: 'example', params: {} };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            const requestPromise = protocol.request(request, mockSchema, {
                timeout: 100,
                resetTimeoutOnProgress: true
            });
            vi.advanceTimersByTime(101);
            await expect(requestPromise).rejects.toThrow('Request timed out');
        });

        test('should handle multiple progress notifications correctly', async () => {
            await protocol.connect(transport);
            const request = { method: 'example', params: {} };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            const onProgressMock = vi.fn();
            const requestPromise = protocol.request(request, mockSchema, {
                timeout: 1000,
                resetTimeoutOnProgress: true,
                onprogress: onProgressMock
            });

            // Simulate multiple progress updates
            for (let i = 1; i <= 3; i++) {
                vi.advanceTimersByTime(800);
                if (transport.onmessage) {
                    transport.onmessage({
                        jsonrpc: '2.0',
                        method: 'notifications/progress',
                        params: {
                            progressToken: 0,
                            progress: i * 25,
                            total: 100
                        }
                    });
                }
                await Promise.resolve();
                expect(onProgressMock).toHaveBeenNthCalledWith(i, {
                    progress: i * 25,
                    total: 100
                });
            }
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    id: 0,
                    result: { result: 'success' }
                });
            }
            await Promise.resolve();
            await expect(requestPromise).resolves.toEqual({ result: 'success' });
        });

        test('should handle progress notifications with message field', async () => {
            await protocol.connect(transport);
            const request = { method: 'example', params: {} };
            const mockSchema: ZodType<{ result: string }> = z.object({
                result: z.string()
            });
            const onProgressMock = vi.fn();

            const requestPromise = protocol.request(request, mockSchema, {
                timeout: 1000,
                onprogress: onProgressMock
            });

            vi.advanceTimersByTime(200);

            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/progress',
                    params: {
                        progressToken: 0,
                        progress: 25,
                        total: 100,
                        message: 'Initializing process...'
                    }
                });
            }
            await Promise.resolve();

            expect(onProgressMock).toHaveBeenCalledWith({
                progress: 25,
                total: 100,
                message: 'Initializing process...'
            });

            vi.advanceTimersByTime(200);

            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/progress',
                    params: {
                        progressToken: 0,
                        progress: 75,
                        total: 100,
                        message: 'Processing data...'
                    }
                });
            }
            await Promise.resolve();

            expect(onProgressMock).toHaveBeenCalledWith({
                progress: 75,
                total: 100,
                message: 'Processing data...'
            });

            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    id: 0,
                    result: { result: 'success' }
                });
            }
            await Promise.resolve();
            await expect(requestPromise).resolves.toEqual({ result: 'success' });
        });
    });

    describe('Debounced Notifications', () => {
        // We need to flush the microtask queue to test the debouncing logic.
        // This helper function does that.
        const flushMicrotasks = () => new Promise(resolve => setImmediate(resolve));

        it('should NOT debounce a notification that has parameters', async () => {
            // ARRANGE
            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ debouncedNotificationMethods: ['test/debounced_with_params'] });
            await protocol.connect(transport);

            // ACT
            // These notifications are configured for debouncing but contain params, so they should be sent immediately.
            await protocol.notification({ method: 'test/debounced_with_params', params: { data: 1 } });
            await protocol.notification({ method: 'test/debounced_with_params', params: { data: 2 } });

            // ASSERT
            // Both should have been sent immediately to avoid data loss.
            expect(sendSpy).toHaveBeenCalledTimes(2);
            expect(sendSpy).toHaveBeenCalledWith(expect.objectContaining({ params: { data: 1 } }), undefined);
            expect(sendSpy).toHaveBeenCalledWith(expect.objectContaining({ params: { data: 2 } }), undefined);
        });

        it('should NOT debounce a notification that has a relatedRequestId', async () => {
            // ARRANGE
            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ debouncedNotificationMethods: ['test/debounced_with_options'] });
            await protocol.connect(transport);

            // ACT
            await protocol.notification({ method: 'test/debounced_with_options' }, { relatedRequestId: 'req-1' });
            await protocol.notification({ method: 'test/debounced_with_options' }, { relatedRequestId: 'req-2' });

            // ASSERT
            expect(sendSpy).toHaveBeenCalledTimes(2);
            expect(sendSpy).toHaveBeenCalledWith(expect.any(Object), { relatedRequestId: 'req-1' });
            expect(sendSpy).toHaveBeenCalledWith(expect.any(Object), { relatedRequestId: 'req-2' });
        });

        it('should clear pending debounced notifications on connection close', async () => {
            // ARRANGE
            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ debouncedNotificationMethods: ['test/debounced'] });
            await protocol.connect(transport);

            // ACT
            // Schedule a notification but don't flush the microtask queue.
            protocol.notification({ method: 'test/debounced' });

            // Close the connection. This should clear the pending set.
            await protocol.close();

            // Now, flush the microtask queue.
            await flushMicrotasks();

            // ASSERT
            // The send should never have happened because the transport was cleared.
            expect(sendSpy).not.toHaveBeenCalled();
        });

        it('should debounce multiple synchronous calls when params property is omitted', async () => {
            // ARRANGE
            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ debouncedNotificationMethods: ['test/debounced'] });
            await protocol.connect(transport);

            // ACT
            // This is the more idiomatic way to write a notification with no params.
            protocol.notification({ method: 'test/debounced' });
            protocol.notification({ method: 'test/debounced' });
            protocol.notification({ method: 'test/debounced' });

            expect(sendSpy).not.toHaveBeenCalled();
            await flushMicrotasks();

            // ASSERT
            expect(sendSpy).toHaveBeenCalledTimes(1);
            // The final sent object might not even have the `params` key, which is fine.
            // We can check that it was called and that the params are "falsy".
            const sentNotification = sendSpy.mock.calls[0][0];
            expect(sentNotification.method).toBe('test/debounced');
            expect(sentNotification.params).toBeUndefined();
        });

        it('should debounce calls when params is explicitly undefined', async () => {
            // ARRANGE
            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ debouncedNotificationMethods: ['test/debounced'] });
            await protocol.connect(transport);

            // ACT
            protocol.notification({ method: 'test/debounced', params: undefined });
            protocol.notification({ method: 'test/debounced', params: undefined });
            await flushMicrotasks();

            // ASSERT
            expect(sendSpy).toHaveBeenCalledTimes(1);
            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    method: 'test/debounced',
                    params: undefined
                }),
                undefined
            );
        });

        it('should send non-debounced notifications immediately and multiple times', async () => {
            // ARRANGE
            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ debouncedNotificationMethods: ['test/debounced'] }); // Configure for a different method
            await protocol.connect(transport);

            // ACT
            // Call a non-debounced notification method multiple times.
            await protocol.notification({ method: 'test/immediate' });
            await protocol.notification({ method: 'test/immediate' });

            // ASSERT
            // Since this method is not in the debounce list, it should be sent every time.
            expect(sendSpy).toHaveBeenCalledTimes(2);
        });

        it('should not debounce any notifications if the option is not provided', async () => {
            // ARRANGE
            // Use the default protocol from beforeEach, which has no debounce options.
            await protocol.connect(transport);

            // ACT
            await protocol.notification({ method: 'any/method' });
            await protocol.notification({ method: 'any/method' });

            // ASSERT
            // Without the config, behavior should be immediate sending.
            expect(sendSpy).toHaveBeenCalledTimes(2);
        });

        it('should handle sequential batches of debounced notifications correctly', async () => {
            // ARRANGE
            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ debouncedNotificationMethods: ['test/debounced'] });
            await protocol.connect(transport);

            // ACT (Batch 1)
            protocol.notification({ method: 'test/debounced' });
            protocol.notification({ method: 'test/debounced' });
            await flushMicrotasks();

            // ASSERT (Batch 1)
            expect(sendSpy).toHaveBeenCalledTimes(1);

            // ACT (Batch 2)
            // After the first batch has been sent, a new batch should be possible.
            protocol.notification({ method: 'test/debounced' });
            protocol.notification({ method: 'test/debounced' });
            await flushMicrotasks();

            // ASSERT (Batch 2)
            // The total number of sends should now be 2.
            expect(sendSpy).toHaveBeenCalledTimes(2);
        });
    });
});

describe('InMemoryTaskMessageQueue', () => {
    let queue: TaskMessageQueue;
    const taskId = 'test-task-id';

    beforeEach(() => {
        queue = new InMemoryTaskMessageQueue();
    });

    describe('enqueue/dequeue maintains FIFO order', () => {
        it('should maintain FIFO order for multiple messages', async () => {
            const msg1 = {
                type: 'notification' as const,
                message: { jsonrpc: '2.0' as const, method: 'test1' },
                timestamp: 1
            };
            const msg2 = {
                type: 'request' as const,
                message: { jsonrpc: '2.0' as const, id: 1, method: 'test2' },
                timestamp: 2
            };
            const msg3 = {
                type: 'notification' as const,
                message: { jsonrpc: '2.0' as const, method: 'test3' },
                timestamp: 3
            };

            await queue.enqueue(taskId, msg1);
            await queue.enqueue(taskId, msg2);
            await queue.enqueue(taskId, msg3);

            expect(await queue.dequeue(taskId)).toEqual(msg1);
            expect(await queue.dequeue(taskId)).toEqual(msg2);
            expect(await queue.dequeue(taskId)).toEqual(msg3);
        });

        it('should return undefined when dequeuing from empty queue', async () => {
            expect(await queue.dequeue(taskId)).toBeUndefined();
        });
    });

    describe('dequeueAll operation', () => {
        it('should return all messages in FIFO order', async () => {
            const msg1 = {
                type: 'notification' as const,
                message: { jsonrpc: '2.0' as const, method: 'test1' },
                timestamp: 1
            };
            const msg2 = {
                type: 'request' as const,
                message: { jsonrpc: '2.0' as const, id: 1, method: 'test2' },
                timestamp: 2
            };
            const msg3 = {
                type: 'notification' as const,
                message: { jsonrpc: '2.0' as const, method: 'test3' },
                timestamp: 3
            };

            await queue.enqueue(taskId, msg1);
            await queue.enqueue(taskId, msg2);
            await queue.enqueue(taskId, msg3);

            const allMessages = await queue.dequeueAll(taskId);

            expect(allMessages).toEqual([msg1, msg2, msg3]);
        });

        it('should return empty array for empty queue', async () => {
            const allMessages = await queue.dequeueAll(taskId);
            expect(allMessages).toEqual([]);
        });

        it('should clear queue after dequeueAll', async () => {
            await queue.enqueue(taskId, {
                type: 'notification' as const,
                message: { jsonrpc: '2.0' as const, method: 'test1' },
                timestamp: 1
            });
            await queue.enqueue(taskId, {
                type: 'notification' as const,
                message: { jsonrpc: '2.0' as const, method: 'test2' },
                timestamp: 2
            });

            await queue.dequeueAll(taskId);

            expect(await queue.dequeue(taskId)).toBeUndefined();
        });
    });
});

describe('mergeCapabilities', () => {
    it('should merge client capabilities', () => {
        const base: ClientCapabilities = {
            sampling: {},
            roots: {
                listChanged: true
            }
        };

        const additional: ClientCapabilities = {
            experimental: {
                feature: {
                    featureFlag: true
                }
            },
            elicitation: {},
            roots: {
                listChanged: true
            }
        };

        const merged = mergeCapabilities(base, additional);
        expect(merged).toEqual({
            sampling: {},
            elicitation: {},
            roots: {
                listChanged: true
            },
            experimental: {
                feature: {
                    featureFlag: true
                }
            }
        });
    });

    it('should merge server capabilities', () => {
        const base: ServerCapabilities = {
            logging: {},
            prompts: {
                listChanged: true
            }
        };

        const additional: ServerCapabilities = {
            resources: {
                subscribe: true
            },
            prompts: {
                listChanged: true
            }
        };

        const merged = mergeCapabilities(base, additional);
        expect(merged).toEqual({
            logging: {},
            prompts: {
                listChanged: true
            },
            resources: {
                subscribe: true
            }
        });
    });

    it('should override existing values with additional values', () => {
        const base: ServerCapabilities = {
            prompts: {
                listChanged: false
            }
        };

        const additional: ServerCapabilities = {
            prompts: {
                listChanged: true
            }
        };

        const merged = mergeCapabilities(base, additional);
        expect(merged.prompts!.listChanged).toBe(true);
    });

    it('should handle empty objects', () => {
        const base = {};
        const additional = {};
        const merged = mergeCapabilities(base, additional);
        expect(merged).toEqual({});
    });
});

describe('Task-based execution', () => {
    let protocol: Protocol<Request, Notification, Result>;
    let transport: MockTransport;
    let sendSpy: MockInstance;

    beforeEach(() => {
        transport = new MockTransport();
        sendSpy = vi.spyOn(transport, 'send');
        protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })({ taskStore: createMockTaskStore(), taskMessageQueue: new InMemoryTaskMessageQueue() });
    });

    describe('request with task metadata', () => {
        it('should include task parameters at top level', async () => {
            await protocol.connect(transport);

            const request = {
                method: 'tools/call',
                params: { name: 'test-tool' }
            };

            const resultSchema = z.object({
                content: z.array(z.object({ type: z.literal('text'), text: z.string() }))
            });

            void protocol
                .request(request, resultSchema, {
                    task: {
                        ttl: 30000,
                        pollInterval: 1000
                    }
                })
                .catch(() => {
                    // May not complete, ignore error
                });

            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    method: 'tools/call',
                    params: {
                        name: 'test-tool',
                        task: {
                            ttl: 30000,
                            pollInterval: 1000
                        }
                    }
                }),
                expect.any(Object)
            );
        });

        it('should preserve existing _meta and add task parameters at top level', async () => {
            await protocol.connect(transport);

            const request = {
                method: 'tools/call',
                params: {
                    name: 'test-tool',
                    _meta: {
                        customField: 'customValue'
                    }
                }
            };

            const resultSchema = z.object({
                content: z.array(z.object({ type: z.literal('text'), text: z.string() }))
            });

            void protocol
                .request(request, resultSchema, {
                    task: {
                        ttl: 60000
                    }
                })
                .catch(() => {
                    // May not complete, ignore error
                });

            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    params: {
                        name: 'test-tool',
                        _meta: {
                            customField: 'customValue'
                        },
                        task: {
                            ttl: 60000
                        }
                    }
                }),
                expect.any(Object)
            );
        });

        it('should return Promise for task-augmented request', async () => {
            await protocol.connect(transport);

            const request = {
                method: 'tools/call',
                params: { name: 'test-tool' }
            };

            const resultSchema = z.object({
                content: z.array(z.object({ type: z.literal('text'), text: z.string() }))
            });

            const resultPromise = protocol.request(request, resultSchema, {
                task: {
                    ttl: 30000
                }
            });

            expect(resultPromise).toBeDefined();
            expect(resultPromise).toBeInstanceOf(Promise);
        });
    });

    describe('relatedTask metadata', () => {
        it('should inject relatedTask metadata into _meta field', async () => {
            await protocol.connect(transport);

            const request = {
                method: 'notifications/message',
                params: { data: 'test' }
            };

            const resultSchema = z.object({});

            // Start the request (don't await completion, just let it send)
            void protocol
                .request(request, resultSchema, {
                    relatedTask: {
                        taskId: 'parent-task-123'
                    }
                })
                .catch(() => {
                    // May not complete, ignore error
                });

            // Wait a bit for the request to be queued
            await new Promise(resolve => setTimeout(resolve, 10));

            // Requests with relatedTask should be queued, not sent via transport
            // This prevents duplicate delivery for bidirectional transports
            expect(sendSpy).not.toHaveBeenCalled();

            // Verify the message was queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();
        });

        it('should work with notification method', async () => {
            await protocol.connect(transport);

            await protocol.notification(
                {
                    method: 'notifications/message',
                    params: { level: 'info', data: 'test message' }
                },
                {
                    relatedTask: {
                        taskId: 'parent-task-456'
                    }
                }
            );

            // Notifications with relatedTask should be queued, not sent via transport
            // This prevents duplicate delivery for bidirectional transports
            expect(sendSpy).not.toHaveBeenCalled();

            // Verify the message was queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            const queuedMessage = await queue!.dequeue('parent-task-456');
            assertQueuedNotification(queuedMessage);
            expect(queuedMessage.message.method).toBe('notifications/message');
            expect(queuedMessage.message.params!._meta![RELATED_TASK_META_KEY]).toEqual({ taskId: 'parent-task-456' });
        });
    });

    describe('task metadata combination', () => {
        it('should combine task, relatedTask, and progress metadata', async () => {
            await protocol.connect(transport);

            const request = {
                method: 'tools/call',
                params: { name: 'test-tool' }
            };

            const resultSchema = z.object({
                content: z.array(z.object({ type: z.literal('text'), text: z.string() }))
            });

            // Start the request (don't await completion, just let it send)
            void protocol
                .request(request, resultSchema, {
                    task: {
                        ttl: 60000,
                        pollInterval: 1000
                    },
                    relatedTask: {
                        taskId: 'parent-task'
                    },
                    onprogress: vi.fn()
                })
                .catch(() => {
                    // May not complete, ignore error
                });

            // Wait a bit for the request to be queued
            await new Promise(resolve => setTimeout(resolve, 10));

            // Requests with relatedTask should be queued, not sent via transport
            // This prevents duplicate delivery for bidirectional transports
            expect(sendSpy).not.toHaveBeenCalled();

            // Verify the message was queued with all metadata combined
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            const queuedMessage = await queue!.dequeue('parent-task');
            assertQueuedRequest(queuedMessage);
            expect(queuedMessage.message.params).toMatchObject({
                name: 'test-tool',
                task: {
                    ttl: 60000,
                    pollInterval: 1000
                },
                _meta: {
                    [RELATED_TASK_META_KEY]: {
                        taskId: 'parent-task'
                    },
                    progressToken: expect.any(Number)
                }
            });
        });
    });

    describe('task status transitions', () => {
        it('should be handled by tool implementors, not protocol layer', () => {
            // Task status management is now the responsibility of tool implementors
            expect(true).toBe(true);
        });

        it('should handle requests with task creation parameters in top-level task field', async () => {
            // This test documents that task creation parameters are now in the top-level task field
            // rather than in _meta, and that task management is handled by tool implementors
            const mockTaskStore = createMockTaskStore();

            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });

            await protocol.connect(transport);

            protocol.setRequestHandler(CallToolRequestSchema, async request => {
                // Tool implementor can access task creation parameters from request.params.task
                expect(request.params.task).toEqual({
                    ttl: 60000,
                    pollInterval: 1000
                });
                return { result: 'success' };
            });

            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 1,
                method: 'tools/call',
                params: {
                    name: 'test',
                    arguments: {},
                    task: {
                        ttl: 60000,
                        pollInterval: 1000
                    }
                }
            });

            // Wait for the request to be processed
            await new Promise(resolve => setTimeout(resolve, 10));
        });
    });

    describe('listTasks', () => {
        it('should handle tasks/list requests and return tasks from TaskStore', async () => {
            const listedTasks = createLatch();
            const mockTaskStore = createMockTaskStore({
                onList: () => listedTasks.releaseLatch()
            });
            const task1 = await mockTaskStore.createTask(
                {
                    pollInterval: 500
                },
                1,
                {
                    method: 'test/method',
                    params: {}
                }
            );
            // Manually set status to completed for this test
            await mockTaskStore.updateTaskStatus(task1.taskId, 'completed');

            const task2 = await mockTaskStore.createTask(
                {
                    ttl: 60000,
                    pollInterval: 1000
                },
                2,
                {
                    method: 'test/method',
                    params: {}
                }
            );

            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });

            await protocol.connect(transport);

            // Simulate receiving a tasks/list request
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 3,
                method: 'tasks/list',
                params: {}
            });

            await listedTasks.waitForLatch();

            expect(mockTaskStore.listTasks).toHaveBeenCalledWith(undefined, undefined);
            const sentMessage = sendSpy.mock.calls[0][0];
            expect(sentMessage.jsonrpc).toBe('2.0');
            expect(sentMessage.id).toBe(3);
            expect(sentMessage.result.tasks).toEqual([
                {
                    taskId: task1.taskId,
                    status: 'completed',
                    ttl: null,
                    createdAt: expect.any(String),
                    lastUpdatedAt: expect.any(String),
                    pollInterval: 500
                },
                {
                    taskId: task2.taskId,
                    status: 'working',
                    ttl: 60000,
                    createdAt: expect.any(String),
                    lastUpdatedAt: expect.any(String),
                    pollInterval: 1000
                }
            ]);
            expect(sentMessage.result._meta).toEqual({});
        });

        it('should handle tasks/list requests with cursor for pagination', async () => {
            const listedTasks = createLatch();
            const mockTaskStore = createMockTaskStore({
                onList: () => listedTasks.releaseLatch()
            });
            const task3 = await mockTaskStore.createTask(
                {
                    pollInterval: 500
                },
                1,
                {
                    method: 'test/method',
                    params: {}
                }
            );

            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });

            await protocol.connect(transport);

            // Simulate receiving a tasks/list request with cursor
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 2,
                method: 'tasks/list',
                params: {
                    cursor: 'task-2'
                }
            });

            await listedTasks.waitForLatch();

            expect(mockTaskStore.listTasks).toHaveBeenCalledWith('task-2', undefined);
            const sentMessage = sendSpy.mock.calls[0][0];
            expect(sentMessage.jsonrpc).toBe('2.0');
            expect(sentMessage.id).toBe(2);
            expect(sentMessage.result.tasks).toEqual([
                {
                    taskId: task3.taskId,
                    status: 'working',
                    ttl: null,
                    createdAt: expect.any(String),
                    lastUpdatedAt: expect.any(String),
                    pollInterval: 500
                }
            ]);
            expect(sentMessage.result.nextCursor).toBeUndefined();
            expect(sentMessage.result._meta).toEqual({});
        });

        it('should handle tasks/list requests with empty results', async () => {
            const listedTasks = createLatch();
            const mockTaskStore = createMockTaskStore({
                onList: () => listedTasks.releaseLatch()
            });

            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });

            await protocol.connect(transport);

            // Simulate receiving a tasks/list request
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 3,
                method: 'tasks/list',
                params: {}
            });

            await listedTasks.waitForLatch();

            expect(mockTaskStore.listTasks).toHaveBeenCalledWith(undefined, undefined);
            const sentMessage = sendSpy.mock.calls[0][0];
            expect(sentMessage.jsonrpc).toBe('2.0');
            expect(sentMessage.id).toBe(3);
            expect(sentMessage.result.tasks).toEqual([]);
            expect(sentMessage.result.nextCursor).toBeUndefined();
            expect(sentMessage.result._meta).toEqual({});
        });

        it('should return error for invalid cursor', async () => {
            const mockTaskStore = createMockTaskStore();
            mockTaskStore.listTasks.mockRejectedValue(new Error('Invalid cursor: bad-cursor'));

            protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });

            await protocol.connect(transport);

            // Simulate receiving a tasks/list request with invalid cursor
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 4,
                method: 'tasks/list',
                params: {
                    cursor: 'bad-cursor'
                }
            });

            await new Promise(resolve => setTimeout(resolve, 10));

            expect(mockTaskStore.listTasks).toHaveBeenCalledWith('bad-cursor', undefined);
            const sentMessage = sendSpy.mock.calls[0][0];
            expect(sentMessage.jsonrpc).toBe('2.0');
            expect(sentMessage.id).toBe(4);
            expect(sentMessage.error).toBeDefined();
            expect(sentMessage.error.code).toBe(-32602); // InvalidParams error code
            expect(sentMessage.error.message).toContain('Failed to list tasks');
            expect(sentMessage.error.message).toContain('Invalid cursor');
        });

        it('should call listTasks method from client side', async () => {
            await protocol.connect(transport);

            const listTasksPromise = (protocol as unknown as TestProtocol).listTasks();

            // Simulate server response
            setTimeout(() => {
                transport.onmessage?.({
                    jsonrpc: '2.0',
                    id: sendSpy.mock.calls[0][0].id,
                    result: {
                        tasks: [
                            {
                                taskId: 'task-1',
                                status: 'completed',
                                ttl: null,
                                createdAt: '2024-01-01T00:00:00Z',
                                lastUpdatedAt: '2024-01-01T00:00:00Z',
                                pollInterval: 500
                            }
                        ],
                        nextCursor: undefined,
                        _meta: {}
                    }
                });
            }, 10);

            const result = await listTasksPromise;

            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    method: 'tasks/list',
                    params: undefined
                }),
                expect.any(Object)
            );
            expect(result.tasks).toHaveLength(1);
            expect(result.tasks[0].taskId).toBe('task-1');
        });

        it('should call listTasks with cursor from client side', async () => {
            await protocol.connect(transport);

            const listTasksPromise = (protocol as unknown as TestProtocol).listTasks({ cursor: 'task-10' });

            // Simulate server response
            setTimeout(() => {
                transport.onmessage?.({
                    jsonrpc: '2.0',
                    id: sendSpy.mock.calls[0][0].id,
                    result: {
                        tasks: [
                            {
                                taskId: 'task-11',
                                status: 'working',
                                ttl: 30000,
                                createdAt: '2024-01-01T00:00:00Z',
                                lastUpdatedAt: '2024-01-01T00:00:00Z',
                                pollInterval: 1000
                            }
                        ],
                        nextCursor: 'task-11',
                        _meta: {}
                    }
                });
            }, 10);

            const result = await listTasksPromise;

            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    method: 'tasks/list',
                    params: {
                        cursor: 'task-10'
                    }
                }),
                expect.any(Object)
            );
            expect(result.tasks).toHaveLength(1);
            expect(result.tasks[0].taskId).toBe('task-11');
            expect(result.nextCursor).toBe('task-11');
        });
    });

    describe('cancelTask', () => {
        it('should handle tasks/cancel requests and update task status to cancelled', async () => {
            const taskDeleted = createLatch();
            const mockTaskStore = createMockTaskStore();
            const task = await mockTaskStore.createTask({}, 1, {
                method: 'test/method',
                params: {}
            });

            mockTaskStore.getTask.mockResolvedValue(task);
            mockTaskStore.updateTaskStatus.mockImplementation(async (taskId: string, status: string) => {
                if (taskId === task.taskId && status === 'cancelled') {
                    taskDeleted.releaseLatch();
                    return;
                }
                throw new Error('Task not found');
            });

            const serverProtocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });
            const serverTransport = new MockTransport();
            const sendSpy = vi.spyOn(serverTransport, 'send');

            await serverProtocol.connect(serverTransport);

            serverTransport.onmessage?.({
                jsonrpc: '2.0',
                id: 5,
                method: 'tasks/cancel',
                params: {
                    taskId: task.taskId
                }
            });

            await taskDeleted.waitForLatch();

            expect(mockTaskStore.getTask).toHaveBeenCalledWith(task.taskId, undefined);
            expect(mockTaskStore.updateTaskStatus).toHaveBeenCalledWith(
                task.taskId,
                'cancelled',
                'Client cancelled task execution.',
                undefined
            );
            const sentMessage = sendSpy.mock.calls[0][0] as unknown as JSONRPCResultResponse;
            expect(sentMessage.jsonrpc).toBe('2.0');
            expect(sentMessage.id).toBe(5);
            expect(sentMessage.result._meta).toBeDefined();
        });

        it('should return error with code -32602 when task does not exist', async () => {
            const taskDeleted = createLatch();
            const mockTaskStore = createMockTaskStore();

            mockTaskStore.getTask.mockResolvedValue(null);

            const serverProtocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });
            const serverTransport = new MockTransport();
            const sendSpy = vi.spyOn(serverTransport, 'send');

            await serverProtocol.connect(serverTransport);

            serverTransport.onmessage?.({
                jsonrpc: '2.0',
                id: 6,
                method: 'tasks/cancel',
                params: {
                    taskId: 'non-existent'
                }
            });

            // Wait a bit for the async handler to complete
            await new Promise(resolve => setTimeout(resolve, 10));
            taskDeleted.releaseLatch();

            expect(mockTaskStore.getTask).toHaveBeenCalledWith('non-existent', undefined);
            const sentMessage = sendSpy.mock.calls[0][0] as unknown as JSONRPCErrorResponse;
            expect(sentMessage.jsonrpc).toBe('2.0');
            expect(sentMessage.id).toBe(6);
            expect(sentMessage.error).toBeDefined();
            expect(sentMessage.error.code).toBe(-32602); // InvalidParams error code
            expect(sentMessage.error.message).toContain('Task not found');
        });

        it('should return error with code -32602 when trying to cancel a task in terminal status', async () => {
            const mockTaskStore = createMockTaskStore();
            const completedTask = await mockTaskStore.createTask({}, 1, {
                method: 'test/method',
                params: {}
            });
            // Set task to completed status
            await mockTaskStore.updateTaskStatus(completedTask.taskId, 'completed');
            completedTask.status = 'completed';

            // Reset the mock so we can check it's not called during cancellation
            mockTaskStore.updateTaskStatus.mockClear();
            mockTaskStore.getTask.mockResolvedValue(completedTask);

            const serverProtocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });
            const serverTransport = new MockTransport();
            const sendSpy = vi.spyOn(serverTransport, 'send');

            await serverProtocol.connect(serverTransport);

            serverTransport.onmessage?.({
                jsonrpc: '2.0',
                id: 7,
                method: 'tasks/cancel',
                params: {
                    taskId: completedTask.taskId
                }
            });

            // Wait a bit for the async handler to complete
            await new Promise(resolve => setTimeout(resolve, 10));

            expect(mockTaskStore.getTask).toHaveBeenCalledWith(completedTask.taskId, undefined);
            expect(mockTaskStore.updateTaskStatus).not.toHaveBeenCalled();
            const sentMessage = sendSpy.mock.calls[0][0] as unknown as JSONRPCErrorResponse;
            expect(sentMessage.jsonrpc).toBe('2.0');
            expect(sentMessage.id).toBe(7);
            expect(sentMessage.error).toBeDefined();
            expect(sentMessage.error.code).toBe(-32602); // InvalidParams error code
            expect(sentMessage.error.message).toContain('Cannot cancel task in terminal status');
        });

        it('should call cancelTask method from client side', async () => {
            await protocol.connect(transport);

            const deleteTaskPromise = (protocol as unknown as TestProtocol).cancelTask({ taskId: 'task-to-delete' });

            // Simulate server response - per MCP spec, CancelTaskResult is Result & Task
            setTimeout(() => {
                transport.onmessage?.({
                    jsonrpc: '2.0',
                    id: sendSpy.mock.calls[0][0].id,
                    result: {
                        _meta: {},
                        taskId: 'task-to-delete',
                        status: 'cancelled',
                        ttl: 60000,
                        createdAt: new Date().toISOString(),
                        lastUpdatedAt: new Date().toISOString()
                    }
                });
            }, 0);

            const result = await deleteTaskPromise;

            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    method: 'tasks/cancel',
                    params: {
                        taskId: 'task-to-delete'
                    }
                }),
                expect.any(Object)
            );
            expect(result._meta).toBeDefined();
            expect(result.taskId).toBe('task-to-delete');
            expect(result.status).toBe('cancelled');
        });
    });

    describe('task status notifications', () => {
        it('should call getTask after updateTaskStatus to enable notification sending', async () => {
            const mockTaskStore = createMockTaskStore();

            // Create a task first
            const task = await mockTaskStore.createTask({}, 1, {
                method: 'test/method',
                params: {}
            });

            const serverProtocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });
            const serverTransport = new MockTransport();

            await serverProtocol.connect(serverTransport);

            // Simulate cancelling the task
            serverTransport.onmessage?.({
                jsonrpc: '2.0',
                id: 2,
                method: 'tasks/cancel',
                params: {
                    taskId: task.taskId
                }
            });

            // Wait for async processing
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify that updateTaskStatus was called
            expect(mockTaskStore.updateTaskStatus).toHaveBeenCalledWith(
                task.taskId,
                'cancelled',
                'Client cancelled task execution.',
                undefined
            );

            // Verify that getTask was called after updateTaskStatus
            // This is done by the RequestTaskStore wrapper to get the updated task for the notification
            const getTaskCalls = mockTaskStore.getTask.mock.calls;
            const lastGetTaskCall = getTaskCalls[getTaskCalls.length - 1];
            expect(lastGetTaskCall[0]).toBe(task.taskId);
        });
    });

    describe('task metadata handling', () => {
        it('should NOT include related-task metadata in tasks/get response', async () => {
            const mockTaskStore = createMockTaskStore();

            // Create a task first
            const task = await mockTaskStore.createTask({}, 1, {
                method: 'test/method',
                params: {}
            });

            const serverProtocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });
            const serverTransport = new MockTransport();
            const sendSpy = vi.spyOn(serverTransport, 'send');

            await serverProtocol.connect(serverTransport);

            // Request task status
            serverTransport.onmessage?.({
                jsonrpc: '2.0',
                id: 2,
                method: 'tasks/get',
                params: {
                    taskId: task.taskId
                }
            });

            // Wait for async processing
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify response does NOT include related-task metadata
            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    result: expect.objectContaining({
                        taskId: task.taskId,
                        status: 'working'
                    })
                })
            );

            // Verify _meta is not present or doesn't contain RELATED_TASK_META_KEY
            const response = sendSpy.mock.calls[0][0] as { result?: { _meta?: Record<string, unknown> } };
            expect(response.result?._meta?.[RELATED_TASK_META_KEY]).toBeUndefined();
        });

        it('should NOT include related-task metadata in tasks/list response', async () => {
            const mockTaskStore = createMockTaskStore();

            // Create a task first
            await mockTaskStore.createTask({}, 1, {
                method: 'test/method',
                params: {}
            });

            const serverProtocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });
            const serverTransport = new MockTransport();
            const sendSpy = vi.spyOn(serverTransport, 'send');

            await serverProtocol.connect(serverTransport);

            // Request task list
            serverTransport.onmessage?.({
                jsonrpc: '2.0',
                id: 2,
                method: 'tasks/list',
                params: {}
            });

            // Wait for async processing
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify response does NOT include related-task metadata
            const response = sendSpy.mock.calls[0][0] as { result?: { _meta?: Record<string, unknown> } };
            expect(response.result?._meta).toEqual({});
        });

        it('should NOT include related-task metadata in tasks/cancel response', async () => {
            const mockTaskStore = createMockTaskStore();

            // Create a task first
            const task = await mockTaskStore.createTask({}, 1, {
                method: 'test/method',
                params: {}
            });

            const serverProtocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });
            const serverTransport = new MockTransport();
            const sendSpy = vi.spyOn(serverTransport, 'send');

            await serverProtocol.connect(serverTransport);

            // Cancel the task
            serverTransport.onmessage?.({
                jsonrpc: '2.0',
                id: 2,
                method: 'tasks/cancel',
                params: {
                    taskId: task.taskId
                }
            });

            // Wait for async processing
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify response does NOT include related-task metadata
            const response = sendSpy.mock.calls[0][0] as { result?: { _meta?: Record<string, unknown> } };
            expect(response.result?._meta).toEqual({});
        });

        it('should include related-task metadata in tasks/result response', async () => {
            const mockTaskStore = createMockTaskStore();

            // Create a task and complete it
            const task = await mockTaskStore.createTask({}, 1, {
                method: 'test/method',
                params: {}
            });

            const testResult = {
                content: [{ type: 'text', text: 'test result' }]
            };

            await mockTaskStore.storeTaskResult(task.taskId, 'completed', testResult);

            const serverProtocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });
            const serverTransport = new MockTransport();
            const sendSpy = vi.spyOn(serverTransport, 'send');

            await serverProtocol.connect(serverTransport);

            // Request task result
            serverTransport.onmessage?.({
                jsonrpc: '2.0',
                id: 2,
                method: 'tasks/result',
                params: {
                    taskId: task.taskId
                }
            });

            // Wait for async processing
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify response DOES include related-task metadata
            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    result: expect.objectContaining({
                        content: testResult.content,
                        _meta: expect.objectContaining({
                            [RELATED_TASK_META_KEY]: {
                                taskId: task.taskId
                            }
                        })
                    })
                })
            );
        });

        it('should propagate related-task metadata to handler sendRequest and sendNotification', async () => {
            const mockTaskStore = createMockTaskStore();

            const serverProtocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });

            const serverTransport = new MockTransport();
            const sendSpy = vi.spyOn(serverTransport, 'send');

            await serverProtocol.connect(serverTransport);

            // Set up a handler that uses sendRequest and sendNotification
            serverProtocol.setRequestHandler(CallToolRequestSchema, async (_request, extra) => {
                // Send a notification using the extra.sendNotification
                await extra.sendNotification({
                    method: 'notifications/message',
                    params: { level: 'info', data: 'test' }
                });

                return {
                    content: [{ type: 'text', text: 'done' }]
                };
            });

            // Send a request with related-task metadata
            let handlerPromise: Promise<void> | undefined;
            const originalOnMessage = serverTransport.onmessage;

            serverTransport.onmessage = message => {
                handlerPromise = Promise.resolve(originalOnMessage?.(message));
                return handlerPromise;
            };

            serverTransport.onmessage({
                jsonrpc: '2.0',
                id: 1,
                method: 'tools/call',
                params: {
                    name: 'test-tool',
                    _meta: {
                        [RELATED_TASK_META_KEY]: {
                            taskId: 'parent-task-123'
                        }
                    }
                }
            });

            // Wait for handler to complete
            if (handlerPromise) {
                await handlerPromise;
            }
            await new Promise(resolve => setTimeout(resolve, 100));

            // Verify the notification was QUEUED (not sent via transport)
            // Messages with relatedTask metadata should be queued for delivery via tasks/result
            // to prevent duplicate delivery for bidirectional transports
            const queue = (serverProtocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            const queuedMessage = await queue!.dequeue('parent-task-123');
            assertQueuedNotification(queuedMessage);
            expect(queuedMessage.message.method).toBe('notifications/message');
            expect(queuedMessage.message.params!._meta![RELATED_TASK_META_KEY]).toEqual({
                taskId: 'parent-task-123'
            });

            // Verify the notification was NOT sent via transport (should be queued instead)
            const notificationCalls = sendSpy.mock.calls.filter(call => 'method' in call[0] && call[0].method === 'notifications/message');
            expect(notificationCalls).toHaveLength(0);
        });
    });
});

describe('Request Cancellation vs Task Cancellation', () => {
    let protocol: Protocol<Request, Notification, Result>;
    let transport: MockTransport;
    let taskStore: TaskStore;

    beforeEach(() => {
        transport = new MockTransport();
        taskStore = createMockTaskStore();
        protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })({ taskStore });
    });

    describe('notifications/cancelled behavior', () => {
        test('should abort request handler when notifications/cancelled is received', async () => {
            await protocol.connect(transport);

            // Set up a request handler that checks if it was aborted
            let wasAborted = false;
            const TestRequestSchema = z.object({
                method: z.literal('test/longRunning'),
                params: z.optional(z.record(z.unknown()))
            });
            protocol.setRequestHandler(TestRequestSchema, async (_request, extra) => {
                // Simulate a long-running operation
                await new Promise(resolve => setTimeout(resolve, 100));
                wasAborted = extra.signal.aborted;
                return { _meta: {} } as Result;
            });

            // Simulate an incoming request
            const requestId = 123;
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    id: requestId,
                    method: 'test/longRunning',
                    params: {}
                });
            }

            // Wait a bit for the handler to start
            await new Promise(resolve => setTimeout(resolve, 10));

            // Send cancellation notification
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/cancelled',
                    params: {
                        requestId: requestId,
                        reason: 'User cancelled'
                    }
                });
            }

            // Wait for the handler to complete
            await new Promise(resolve => setTimeout(resolve, 150));

            // Verify the request was aborted
            expect(wasAborted).toBe(true);
        });

        test('should NOT automatically cancel associated tasks when notifications/cancelled is received', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await taskStore.createTask({ ttl: 60000 }, 'req-1', {
                method: 'test/method',
                params: {}
            });

            // Send cancellation notification for the request
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/cancelled',
                    params: {
                        requestId: 'req-1',
                        reason: 'User cancelled'
                    }
                });
            }

            // Wait a bit
            await new Promise(resolve => setTimeout(resolve, 10));

            // Verify the task status was NOT changed to cancelled
            const updatedTask = await taskStore.getTask(task.taskId);
            expect(updatedTask?.status).toBe('working');
            expect(taskStore.updateTaskStatus).not.toHaveBeenCalledWith(task.taskId, 'cancelled', expect.any(String));
        });
    });

    describe('tasks/cancel behavior', () => {
        test('should cancel task independently of request cancellation', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await taskStore.createTask({ ttl: 60000 }, 'req-1', {
                method: 'test/method',
                params: {}
            });

            // Cancel the task using tasks/cancel
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    id: 999,
                    method: 'tasks/cancel',
                    params: {
                        taskId: task.taskId
                    }
                });
            }

            // Wait for the handler to complete
            await new Promise(resolve => setTimeout(resolve, 10));

            // Verify the task was cancelled
            expect(taskStore.updateTaskStatus).toHaveBeenCalledWith(
                task.taskId,
                'cancelled',
                'Client cancelled task execution.',
                undefined
            );
        });

        test('should reject cancellation of terminal tasks', async () => {
            await protocol.connect(transport);
            const sendSpy = vi.spyOn(transport, 'send');

            // Create a task and mark it as completed
            const task = await taskStore.createTask({ ttl: 60000 }, 'req-1', {
                method: 'test/method',
                params: {}
            });
            await taskStore.updateTaskStatus(task.taskId, 'completed');

            // Try to cancel the completed task
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    id: 999,
                    method: 'tasks/cancel',
                    params: {
                        taskId: task.taskId
                    }
                });
            }

            // Wait for the handler to complete
            await new Promise(resolve => setTimeout(resolve, 10));

            // Verify an error was sent
            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    jsonrpc: '2.0',
                    id: 999,
                    error: expect.objectContaining({
                        code: ErrorCode.InvalidParams,
                        message: expect.stringContaining('Cannot cancel task in terminal status')
                    })
                })
            );
        });

        test('should return error when task not found', async () => {
            await protocol.connect(transport);
            const sendSpy = vi.spyOn(transport, 'send');

            // Try to cancel a non-existent task
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    id: 999,
                    method: 'tasks/cancel',
                    params: {
                        taskId: 'non-existent-task'
                    }
                });
            }

            // Wait for the handler to complete
            await new Promise(resolve => setTimeout(resolve, 10));

            // Verify an error was sent
            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    jsonrpc: '2.0',
                    id: 999,
                    error: expect.objectContaining({
                        code: ErrorCode.InvalidParams,
                        message: expect.stringContaining('Task not found')
                    })
                })
            );
        });
    });

    describe('separation of concerns', () => {
        test('should allow request cancellation without affecting task', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await taskStore.createTask({ ttl: 60000 }, 'req-1', {
                method: 'test/method',
                params: {}
            });

            // Cancel the request (not the task)
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/cancelled',
                    params: {
                        requestId: 'req-1',
                        reason: 'User cancelled request'
                    }
                });
            }

            await new Promise(resolve => setTimeout(resolve, 10));

            // Verify task is still working
            const updatedTask = await taskStore.getTask(task.taskId);
            expect(updatedTask?.status).toBe('working');
        });

        test('should allow task cancellation without affecting request', async () => {
            await protocol.connect(transport);

            // Set up a request handler
            let requestCompleted = false;
            const TestMethodSchema = z.object({
                method: z.literal('test/method'),
                params: z.optional(z.record(z.unknown()))
            });
            protocol.setRequestHandler(TestMethodSchema, async () => {
                await new Promise(resolve => setTimeout(resolve, 50));
                requestCompleted = true;
                return { _meta: {} } as Result;
            });

            // Create a task
            const task = await taskStore.createTask({ ttl: 60000 }, 'req-1', {
                method: 'test/method',
                params: {}
            });

            // Start a request
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    id: 123,
                    method: 'test/method',
                    params: {}
                });
            }

            // Cancel the task (not the request)
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    id: 999,
                    method: 'tasks/cancel',
                    params: {
                        taskId: task.taskId
                    }
                });
            }

            // Wait for request to complete
            await new Promise(resolve => setTimeout(resolve, 100));

            // Verify request completed normally
            expect(requestCompleted).toBe(true);

            // Verify task was cancelled
            expect(taskStore.updateTaskStatus).toHaveBeenCalledWith(
                task.taskId,
                'cancelled',
                'Client cancelled task execution.',
                undefined
            );
        });
    });
});

describe('Progress notification support for tasks', () => {
    let protocol: Protocol<Request, Notification, Result>;
    let transport: MockTransport;
    let sendSpy: MockInstance;

    beforeEach(() => {
        transport = new MockTransport();
        sendSpy = vi.spyOn(transport, 'send');
        protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })();
    });

    it('should maintain progress token association after CreateTaskResult is returned', async () => {
        const taskStore = createMockTaskStore();
        const protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })({ taskStore });

        const transport = new MockTransport();
        const sendSpy = vi.spyOn(transport, 'send');
        await protocol.connect(transport);

        const progressCallback = vi.fn();
        const request = {
            method: 'tools/call',
            params: { name: 'test-tool' }
        };

        const resultSchema = z.object({
            task: z.object({
                taskId: z.string(),
                status: z.string(),
                ttl: z.number().nullable(),
                createdAt: z.string()
            })
        });

        // Start a task-augmented request with progress callback
        void protocol
            .request(request, resultSchema, {
                task: { ttl: 60000 },
                onprogress: progressCallback
            })
            .catch(() => {
                // May not complete, ignore error
            });

        // Wait a bit for the request to be sent
        await new Promise(resolve => setTimeout(resolve, 10));

        // Get the message ID from the sent request
        const sentRequest = sendSpy.mock.calls[0][0] as { id: number; params: { _meta: { progressToken: number } } };
        const messageId = sentRequest.id;
        const progressToken = sentRequest.params._meta.progressToken;

        expect(progressToken).toBe(messageId);

        // Simulate CreateTaskResult response
        const taskId = 'test-task-123';
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                id: messageId,
                result: {
                    task: {
                        taskId,
                        status: 'working',
                        ttl: 60000,
                        createdAt: new Date().toISOString()
                    }
                }
            });
        }

        // Wait for response to be processed
        await Promise.resolve();
        await Promise.resolve();

        // Send a progress notification - should still work after CreateTaskResult
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                method: 'notifications/progress',
                params: {
                    progressToken,
                    progress: 50,
                    total: 100
                }
            });
        }

        // Wait for notification to be processed
        await Promise.resolve();

        // Verify progress callback was invoked
        expect(progressCallback).toHaveBeenCalledWith({
            progress: 50,
            total: 100
        });
    });

    it('should stop progress notifications when task reaches terminal status (completed)', async () => {
        const taskStore = createMockTaskStore();
        const protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })({ taskStore });

        const transport = new MockTransport();
        const sendSpy = vi.spyOn(transport, 'send');
        await protocol.connect(transport);

        // 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 });

                // Simulate async work then complete the task
                setTimeout(async () => {
                    await extra.taskStore!.storeTaskResult(task.taskId, 'completed', {
                        content: [{ type: 'text', text: 'Done' }]
                    });
                }, 50);

                return { task };
            }
            return { content: [] };
        });

        const progressCallback = vi.fn();
        const request = {
            method: 'tools/call',
            params: { name: 'test-tool' }
        };

        const resultSchema = z.object({
            task: z.object({
                taskId: z.string(),
                status: z.string(),
                ttl: z.number().nullable(),
                createdAt: z.string()
            })
        });

        // Start a task-augmented request with progress callback
        void protocol
            .request(request, resultSchema, {
                task: { ttl: 60000 },
                onprogress: progressCallback
            })
            .catch(() => {
                // May not complete, ignore error
            });

        // Wait a bit for the request to be sent
        await new Promise(resolve => setTimeout(resolve, 10));

        const sentRequest = sendSpy.mock.calls[0][0] as { id: number; params: { _meta: { progressToken: number } } };
        const messageId = sentRequest.id;
        const progressToken = sentRequest.params._meta.progressToken;

        // Create a task in the mock store first so it exists when we try to get it later
        const createdTask = await taskStore.createTask({ ttl: 60000 }, messageId, request);
        const taskId = createdTask.taskId;

        // Simulate CreateTaskResult response
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                id: messageId,
                result: {
                    task: createdTask
                }
            });
        }

        await Promise.resolve();
        await Promise.resolve();

        // Progress notification should work while task is working
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                method: 'notifications/progress',
                params: {
                    progressToken,
                    progress: 50,
                    total: 100
                }
            });
        }

        await Promise.resolve();

        expect(progressCallback).toHaveBeenCalledTimes(1);

        // Verify the task-progress association was created
        const taskProgressTokens = (protocol as unknown as TestProtocol)._taskProgressTokens as Map<string, number>;
        expect(taskProgressTokens.has(taskId)).toBe(true);
        expect(taskProgressTokens.get(taskId)).toBe(progressToken);

        // Simulate task completion by calling through the protocol's task store
        // This will trigger the cleanup logic
        const mockRequest = { jsonrpc: '2.0' as const, id: 999, method: 'test', params: {} };
        const requestTaskStore = (protocol as unknown as TestProtocol).requestTaskStore(mockRequest, undefined);
        await requestTaskStore.storeTaskResult(taskId, 'completed', { content: [] });

        // Wait for all async operations including notification sending to complete
        await new Promise(resolve => setTimeout(resolve, 50));

        // Verify the association was cleaned up
        expect(taskProgressTokens.has(taskId)).toBe(false);

        // Try to send progress notification after task completion - should be ignored
        progressCallback.mockClear();
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                method: 'notifications/progress',
                params: {
                    progressToken,
                    progress: 100,
                    total: 100
                }
            });
        }

        await Promise.resolve();

        // Progress callback should NOT be invoked after task completion
        expect(progressCallback).not.toHaveBeenCalled();
    });

    it('should stop progress notifications when task reaches terminal status (failed)', async () => {
        const taskStore = createMockTaskStore();
        const protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })({ taskStore });

        const transport = new MockTransport();
        const sendSpy = vi.spyOn(transport, 'send');
        await protocol.connect(transport);

        const progressCallback = vi.fn();
        const request = {
            method: 'tools/call',
            params: { name: 'test-tool' }
        };

        const resultSchema = z.object({
            task: z.object({
                taskId: z.string(),
                status: z.string(),
                ttl: z.number().nullable(),
                createdAt: z.string()
            })
        });

        void protocol.request(request, resultSchema, {
            task: { ttl: 60000 },
            onprogress: progressCallback
        });

        const sentRequest = sendSpy.mock.calls[0][0] as { id: number; params: { _meta: { progressToken: number } } };
        const messageId = sentRequest.id;
        const progressToken = sentRequest.params._meta.progressToken;

        // Simulate CreateTaskResult response
        const taskId = 'test-task-456';
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                id: messageId,
                result: {
                    task: {
                        taskId,
                        status: 'working',
                        ttl: 60000,
                        createdAt: new Date().toISOString()
                    }
                }
            });
        }

        await new Promise(resolve => setTimeout(resolve, 10));

        // Simulate task failure via storeTaskResult
        await taskStore.storeTaskResult(taskId, 'failed', {
            content: [],
            isError: true
        });

        // Manually trigger the status notification
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                method: 'notifications/tasks/status',
                params: {
                    taskId,
                    status: 'failed',
                    ttl: 60000,
                    createdAt: new Date().toISOString(),
                    lastUpdatedAt: new Date().toISOString(),
                    statusMessage: 'Task failed'
                }
            });
        }

        await new Promise(resolve => setTimeout(resolve, 10));

        // Try to send progress notification after task failure - should be ignored
        progressCallback.mockClear();
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                method: 'notifications/progress',
                params: {
                    progressToken,
                    progress: 75,
                    total: 100
                }
            });
        }

        expect(progressCallback).not.toHaveBeenCalled();
    });

    it('should stop progress notifications when task is cancelled', async () => {
        const taskStore = createMockTaskStore();
        const protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })({ taskStore });

        const transport = new MockTransport();
        const sendSpy = vi.spyOn(transport, 'send');
        await protocol.connect(transport);

        const progressCallback = vi.fn();
        const request = {
            method: 'tools/call',
            params: { name: 'test-tool' }
        };

        const resultSchema = z.object({
            task: z.object({
                taskId: z.string(),
                status: z.string(),
                ttl: z.number().nullable(),
                createdAt: z.string()
            })
        });

        void protocol.request(request, resultSchema, {
            task: { ttl: 60000 },
            onprogress: progressCallback
        });

        const sentRequest = sendSpy.mock.calls[0][0] as { id: number; params: { _meta: { progressToken: number } } };
        const messageId = sentRequest.id;
        const progressToken = sentRequest.params._meta.progressToken;

        // Simulate CreateTaskResult response
        const taskId = 'test-task-789';
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                id: messageId,
                result: {
                    task: {
                        taskId,
                        status: 'working',
                        ttl: 60000,
                        createdAt: new Date().toISOString()
                    }
                }
            });
        }

        await new Promise(resolve => setTimeout(resolve, 10));

        // Simulate task cancellation via updateTaskStatus
        await taskStore.updateTaskStatus(taskId, 'cancelled', 'User cancelled');

        // Manually trigger the status notification
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                method: 'notifications/tasks/status',
                params: {
                    taskId,
                    status: 'cancelled',
                    ttl: 60000,
                    createdAt: new Date().toISOString(),
                    lastUpdatedAt: new Date().toISOString(),
                    statusMessage: 'User cancelled'
                }
            });
        }

        await new Promise(resolve => setTimeout(resolve, 10));

        // Try to send progress notification after cancellation - should be ignored
        progressCallback.mockClear();
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                method: 'notifications/progress',
                params: {
                    progressToken,
                    progress: 25,
                    total: 100
                }
            });
        }

        expect(progressCallback).not.toHaveBeenCalled();
    });

    it('should use the same progressToken throughout task lifetime', async () => {
        const taskStore = createMockTaskStore();
        const protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })({ taskStore });

        const transport = new MockTransport();
        const sendSpy = vi.spyOn(transport, 'send');
        await protocol.connect(transport);

        const progressCallback = vi.fn();
        const request = {
            method: 'tools/call',
            params: { name: 'test-tool' }
        };

        const resultSchema = z.object({
            task: z.object({
                taskId: z.string(),
                status: z.string(),
                ttl: z.number().nullable(),
                createdAt: z.string()
            })
        });

        void protocol.request(request, resultSchema, {
            task: { ttl: 60000 },
            onprogress: progressCallback
        });

        const sentRequest = sendSpy.mock.calls[0][0] as { id: number; params: { _meta: { progressToken: number } } };
        const messageId = sentRequest.id;
        const progressToken = sentRequest.params._meta.progressToken;

        // Simulate CreateTaskResult response
        const taskId = 'test-task-consistency';
        if (transport.onmessage) {
            transport.onmessage({
                jsonrpc: '2.0',
                id: messageId,
                result: {
                    task: {
                        taskId,
                        status: 'working',
                        ttl: 60000,
                        createdAt: new Date().toISOString()
                    }
                }
            });
        }

        await Promise.resolve();
        await Promise.resolve();

        // Send multiple progress notifications with the same token
        const progressUpdates = [
            { progress: 25, total: 100 },
            { progress: 50, total: 100 },
            { progress: 75, total: 100 }
        ];

        for (const update of progressUpdates) {
            if (transport.onmessage) {
                transport.onmessage({
                    jsonrpc: '2.0',
                    method: 'notifications/progress',
                    params: {
                        progressToken, // Same token for all notifications
                        ...update
                    }
                });
            }
            await Promise.resolve();
        }

        // Verify all progress notifications were received with the same token
        expect(progressCallback).toHaveBeenCalledTimes(3);
        expect(progressCallback).toHaveBeenNthCalledWith(1, { progress: 25, total: 100 });
        expect(progressCallback).toHaveBeenNthCalledWith(2, { progress: 50, total: 100 });
        expect(progressCallback).toHaveBeenNthCalledWith(3, { progress: 75, total: 100 });
    });

    it('should maintain progressToken throughout task lifetime', async () => {
        await protocol.connect(transport);

        const request = {
            method: 'tools/call',
            params: { name: 'long-running-tool' }
        };

        const resultSchema = z.object({
            content: z.array(z.object({ type: z.literal('text'), text: z.string() }))
        });

        const onProgressMock = vi.fn();

        void protocol.request(request, resultSchema, {
            task: {
                ttl: 60000
            },
            onprogress: onProgressMock
        });

        const sentMessage = sendSpy.mock.calls[0][0];
        expect(sentMessage.params._meta.progressToken).toBeDefined();
    });

    it('should support progress notifications with task-augmented requests', async () => {
        await protocol.connect(transport);

        const request = {
            method: 'tools/call',
            params: { name: 'test-tool' }
        };

        const resultSchema = z.object({
            content: z.array(z.object({ type: z.literal('text'), text: z.string() }))
        });

        const onProgressMock = vi.fn();

        void protocol.request(request, resultSchema, {
            task: {
                ttl: 30000
            },
            onprogress: onProgressMock
        });

        const sentMessage = sendSpy.mock.calls[0][0];
        const progressToken = sentMessage.params._meta.progressToken;

        // Simulate progress notification
        transport.onmessage?.({
            jsonrpc: '2.0',
            method: 'notifications/progress',
            params: {
                progressToken,
                progress: 50,
                total: 100,
                message: 'Processing...'
            }
        });

        await new Promise(resolve => setTimeout(resolve, 10));

        expect(onProgressMock).toHaveBeenCalledWith({
            progress: 50,
            total: 100,
            message: 'Processing...'
        });
    });

    it('should continue progress notifications after CreateTaskResult', async () => {
        await protocol.connect(transport);

        const request = {
            method: 'tools/call',
            params: { name: 'test-tool' }
        };

        const resultSchema = z.object({
            task: z.object({
                taskId: z.string(),
                status: z.string(),
                ttl: z.number().nullable(),
                createdAt: z.string()
            })
        });

        const onProgressMock = vi.fn();

        void protocol.request(request, resultSchema, {
            task: {
                ttl: 30000
            },
            onprogress: onProgressMock
        });

        const sentMessage = sendSpy.mock.calls[0][0];
        const progressToken = sentMessage.params._meta.progressToken;

        // Simulate CreateTaskResult response
        setTimeout(() => {
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: sentMessage.id,
                result: {
                    task: {
                        taskId: 'task-123',
                        status: 'working',
                        ttl: 30000,
                        createdAt: new Date().toISOString()
                    }
                }
            });
        }, 5);

        // Progress notifications should still work
        setTimeout(() => {
            transport.onmessage?.({
                jsonrpc: '2.0',
                method: 'notifications/progress',
                params: {
                    progressToken,
                    progress: 75,
                    total: 100
                }
            });
        }, 10);

        await new Promise(resolve => setTimeout(resolve, 20));

        expect(onProgressMock).toHaveBeenCalledWith({
            progress: 75,
            total: 100
        });
    });
});

describe('Capability negotiation for tasks', () => {
    it('should use empty objects for capability fields', () => {
        const serverCapabilities = {
            tasks: {
                list: {},
                cancel: {},
                requests: {
                    tools: {
                        call: {}
                    }
                }
            }
        };

        expect(serverCapabilities.tasks.list).toEqual({});
        expect(serverCapabilities.tasks.cancel).toEqual({});
        expect(serverCapabilities.tasks.requests.tools.call).toEqual({});
    });

    it('should include list and cancel in server capabilities', () => {
        const serverCapabilities = {
            tasks: {
                list: {},
                cancel: {}
            }
        };

        expect('list' in serverCapabilities.tasks).toBe(true);
        expect('cancel' in serverCapabilities.tasks).toBe(true);
    });

    it('should include list and cancel in client capabilities', () => {
        const clientCapabilities = {
            tasks: {
                list: {},
                cancel: {}
            }
        };

        expect('list' in clientCapabilities.tasks).toBe(true);
        expect('cancel' in clientCapabilities.tasks).toBe(true);
    });
});

describe('Message interception for task-related notifications', () => {
    it('should queue notifications with io.modelcontextprotocol/related-task metadata', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });

        await server.connect(transport);

        // Create a task first
        const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });

        // Send a notification with related task metadata
        await server.notification(
            {
                method: 'notifications/message',
                params: { level: 'info', data: 'test message' }
            },
            {
                relatedTask: { taskId: task.taskId }
            }
        );

        // Access the private queue to verify the message was queued
        const queue = (server as unknown as TestProtocol)._taskMessageQueue;
        expect(queue).toBeDefined();

        const queuedMessage = await queue!.dequeue(task.taskId);
        assertQueuedNotification(queuedMessage);
        expect(queuedMessage.message.method).toBe('notifications/message');
        expect(queuedMessage.message.params!._meta![RELATED_TASK_META_KEY]).toEqual({ taskId: task.taskId });
    });

    it('should not queue notifications without related-task metadata', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });

        await server.connect(transport);

        // Send a notification without related task metadata
        await server.notification({
            method: 'notifications/message',
            params: { level: 'info', data: 'test message' }
        });

        // Verify message was not queued (notification without metadata goes through transport)
        // We can't directly check the queue, but we know it wasn't queued because
        // notifications without relatedTask metadata are sent via transport, not queued
    });

    // Test removed: _taskResultWaiters was removed in favor of polling-based task updates
    // The functionality is still tested through integration tests that verify message queuing works

    it('should propagate queue overflow errors without failing the task', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue(), maxTaskQueueSize: 100 });

        await server.connect(transport);

        // Create a task
        const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });

        // Fill the queue to max capacity (100 messages)
        for (let i = 0; i < 100; i++) {
            await server.notification(
                {
                    method: 'notifications/message',
                    params: { level: 'info', data: `message ${i}` }
                },
                {
                    relatedTask: { taskId: task.taskId }
                }
            );
        }

        // Try to add one more message - should throw an error
        await expect(
            server.notification(
                {
                    method: 'notifications/message',
                    params: { level: 'info', data: 'overflow message' }
                },
                {
                    relatedTask: { taskId: task.taskId }
                }
            )
        ).rejects.toThrow('overflow');

        // Verify the task was NOT automatically failed by the Protocol
        // (implementations can choose to fail tasks on overflow if they want)
        expect(taskStore.updateTaskStatus).not.toHaveBeenCalledWith(task.taskId, 'failed', expect.anything(), expect.anything());
    });

    it('should extract task ID correctly from metadata', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });

        await server.connect(transport);

        const taskId = 'custom-task-id-123';

        // Send a notification with custom task ID
        await server.notification(
            {
                method: 'notifications/message',
                params: { level: 'info', data: 'test message' }
            },
            {
                relatedTask: { taskId }
            }
        );

        // Verify the message was queued under the correct task ID
        const queue = (server as unknown as TestProtocol)._taskMessageQueue;
        expect(queue).toBeDefined();
        const queuedMessage = await queue!.dequeue(taskId);
        expect(queuedMessage).toBeDefined();
    });

    it('should preserve message order when queuing multiple notifications', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });

        await server.connect(transport);

        // Create a task
        const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });

        // Send multiple notifications
        for (let i = 0; i < 5; i++) {
            await server.notification(
                {
                    method: 'notifications/message',
                    params: { level: 'info', data: `message ${i}` }
                },
                {
                    relatedTask: { taskId: task.taskId }
                }
            );
        }

        // Verify messages are in FIFO order
        const queue = (server as unknown as TestProtocol)._taskMessageQueue;
        expect(queue).toBeDefined();

        for (let i = 0; i < 5; i++) {
            const queuedMessage = await queue!.dequeue(task.taskId);
            assertQueuedNotification(queuedMessage);
            expect(queuedMessage.message.params!.data).toBe(`message ${i}`);
        }
    });
});

describe('Message interception for task-related requests', () => {
    it('should queue requests with io.modelcontextprotocol/related-task metadata', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });

        await server.connect(transport);

        // Create a task first
        const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });

        // Send a request with related task metadata (don't await - we're testing queuing)
        const requestPromise = server.request(
            {
                method: 'ping',
                params: {}
            },
            z.object({}),
            {
                relatedTask: { taskId: task.taskId }
            }
        );

        // Access the private queue to verify the message was queued
        const queue = (server as unknown as TestProtocol)._taskMessageQueue;
        expect(queue).toBeDefined();

        const queuedMessage = await queue!.dequeue(task.taskId);
        assertQueuedRequest(queuedMessage);
        expect(queuedMessage.message.method).toBe('ping');
        expect(queuedMessage.message.params!._meta![RELATED_TASK_META_KEY]).toEqual({ taskId: task.taskId });

        // Verify resolver is stored in _requestResolvers map (not in the message)
        const requestId = (queuedMessage!.message as JSONRPCRequest).id as RequestId;
        const resolvers = (server as unknown as TestProtocol)._requestResolvers;
        expect(resolvers.has(requestId)).toBe(true);

        // Clean up - send a response to prevent hanging promise
        transport.onmessage?.({
            jsonrpc: '2.0',
            id: requestId,
            result: {}
        });

        await requestPromise;
    });

    it('should not queue requests without related-task metadata', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });

        await server.connect(transport);

        // Send a request without related task metadata
        const requestPromise = server.request(
            {
                method: 'ping',
                params: {}
            },
            z.object({})
        );

        // Verify queue exists (but we don't track size in the new API)
        const queue = (server as unknown as TestProtocol)._taskMessageQueue;
        expect(queue).toBeDefined();

        // Clean up - send a response
        transport.onmessage?.({
            jsonrpc: '2.0',
            id: 0,
            result: {}
        });

        await requestPromise;
    });

    // Test removed: _taskResultWaiters was removed in favor of polling-based task updates
    // The functionality is still tested through integration tests that verify message queuing works

    it('should store request resolver for response routing', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });

        await server.connect(transport);

        // Create a task
        const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });

        // Send a request with related task metadata
        const requestPromise = server.request(
            {
                method: 'ping',
                params: {}
            },
            z.object({}),
            {
                relatedTask: { taskId: task.taskId }
            }
        );

        // Verify the resolver was stored
        const resolvers = (server as unknown as TestProtocol)._requestResolvers;
        expect(resolvers.size).toBe(1);

        // Get the request ID from the queue
        const queue = (server as unknown as TestProtocol)._taskMessageQueue;
        const queuedMessage = await queue!.dequeue(task.taskId);
        const requestId = (queuedMessage!.message as JSONRPCRequest).id as RequestId;

        expect(resolvers.has(requestId)).toBe(true);

        // Send a response to trigger resolver
        transport.onmessage?.({
            jsonrpc: '2.0',
            id: requestId,
            result: {}
        });

        await requestPromise;

        // Verify resolver was cleaned up after response
        expect(resolvers.has(requestId)).toBe(false);
    });

    it('should route responses to side-channeled requests', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const queue = new InMemoryTaskMessageQueue();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: queue });

        await server.connect(transport);

        // Create a task
        const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });

        // Send a request with related task metadata
        const requestPromise = server.request(
            {
                method: 'ping',
                params: {}
            },
            z.object({ message: z.string() }),
            {
                relatedTask: { taskId: task.taskId }
            }
        );

        // Get the request ID from the queue
        const queuedMessage = await queue.dequeue(task.taskId);
        const requestId = (queuedMessage!.message as JSONRPCRequest).id as RequestId;

        // Enqueue a response message to the queue (simulating client sending response back)
        await queue.enqueue(task.taskId, {
            type: 'response',
            message: {
                jsonrpc: '2.0',
                id: requestId,
                result: { message: 'pong' }
            },
            timestamp: Date.now()
        });

        // Simulate a client calling tasks/result which will process the response
        // This is done by creating a mock request handler that will trigger the GetTaskPayloadRequest handler
        const mockRequestId = 999;
        transport.onmessage?.({
            jsonrpc: '2.0',
            id: mockRequestId,
            method: 'tasks/result',
            params: { taskId: task.taskId }
        });

        // Wait for the response to be processed
        await new Promise(resolve => setTimeout(resolve, 50));

        // Mark task as completed
        await taskStore.updateTaskStatus(task.taskId, 'completed');
        await taskStore.storeTaskResult(task.taskId, 'completed', { _meta: {} });

        // Verify the response was routed correctly
        const result = await requestPromise;
        expect(result).toEqual({ message: 'pong' });
    });

    it('should log error when resolver is missing for side-channeled request', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });

        const errors: Error[] = [];
        server.onerror = (error: Error) => {
            errors.push(error);
        };

        await server.connect(transport);

        // Create a task
        const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });

        // Send a request with related task metadata
        void server.request(
            {
                method: 'ping',
                params: {}
            },
            z.object({ message: z.string() }),
            {
                relatedTask: { taskId: task.taskId }
            }
        );

        // Get the request ID from the queue
        const queue = (server as unknown as TestProtocol)._taskMessageQueue;
        const queuedMessage = await queue!.dequeue(task.taskId);
        const requestId = (queuedMessage!.message as JSONRPCRequest).id as RequestId;

        // Manually delete the resolver to simulate missing resolver
        (server as unknown as TestProtocol)._requestResolvers.delete(requestId);

        // Enqueue a response message - this should trigger the error logging when processed
        await queue!.enqueue(task.taskId, {
            type: 'response',
            message: {
                jsonrpc: '2.0',
                id: requestId,
                result: { message: 'pong' }
            },
            timestamp: Date.now()
        });

        // Simulate a client calling tasks/result which will process the response
        const mockRequestId = 888;
        transport.onmessage?.({
            jsonrpc: '2.0',
            id: mockRequestId,
            method: 'tasks/result',
            params: { taskId: task.taskId }
        });

        // Wait for the response to be processed
        await new Promise(resolve => setTimeout(resolve, 50));

        // Mark task as completed
        await taskStore.updateTaskStatus(task.taskId, 'completed');
        await taskStore.storeTaskResult(task.taskId, 'completed', { _meta: {} });

        // Wait a bit more for error to be logged
        await new Promise(resolve => setTimeout(resolve, 50));

        // Verify error was logged
        expect(errors.length).toBeGreaterThanOrEqual(1);
        expect(errors.some(e => e.message.includes('Response handler missing for request'))).toBe(true);
    });

    it('should propagate queue overflow errors for requests without failing the task', async () => {
        const taskStore = createMockTaskStore();
        const transport = new MockTransport();
        const server = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({ taskStore, taskMessageQueue: new InMemoryTaskMessageQueue(), maxTaskQueueSize: 100 });

        await server.connect(transport);

        // Create a task
        const task = await taskStore.createTask({ ttl: 60000 }, 'test-request-1', { method: 'tools/call', params: {} });

        // Fill the queue to max capacity (100 messages)
        const promises: Promise<unknown>[] = [];
        for (let i = 0; i < 100; i++) {
            const promise = server
                .request(
                    {
                        method: 'ping',
                        params: {}
                    },
                    z.object({}),
                    {
                        relatedTask: { taskId: task.taskId }
                    }
                )
                .catch(() => {
                    // Requests will remain pending until task completes or fails
                });
            promises.push(promise);
        }

        // Try to add one more request - should throw an error
        await expect(
            server.request(
                {
                    method: 'ping',
                    params: {}
                },
                z.object({}),
                {
                    relatedTask: { taskId: task.taskId }
                }
            )
        ).rejects.toThrow('overflow');

        // Verify the task was NOT automatically failed by the Protocol
        // (implementations can choose to fail tasks on overflow if they want)
        expect(taskStore.updateTaskStatus).not.toHaveBeenCalledWith(task.taskId, 'failed', expect.anything(), expect.anything());
    });
});

describe('Message Interception', () => {
    let protocol: Protocol<Request, Notification, Result>;
    let transport: MockTransport;
    let mockTaskStore: TaskStore & { [K in keyof TaskStore]: MockInstance };

    beforeEach(() => {
        transport = new MockTransport();
        mockTaskStore = createMockTaskStore();
        protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })({ taskStore: mockTaskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });
    });

    describe('messages with relatedTask metadata are queued', () => {
        it('should queue notifications with relatedTask metadata', async () => {
            await protocol.connect(transport);

            // Send a notification with relatedTask metadata
            await protocol.notification(
                {
                    method: 'notifications/message',
                    params: { level: 'info', data: 'test message' }
                },
                {
                    relatedTask: {
                        taskId: 'task-123'
                    }
                }
            );

            // Access the private _taskMessageQueue to verify the message was queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            const queuedMessage = await queue!.dequeue('task-123');
            assertQueuedNotification(queuedMessage);
            expect(queuedMessage!.message.method).toBe('notifications/message');
        });

        it('should queue requests with relatedTask metadata', async () => {
            await protocol.connect(transport);

            const mockSchema = z.object({ result: z.string() });

            // Send a request with relatedTask metadata
            const requestPromise = protocol.request(
                {
                    method: 'test/request',
                    params: { data: 'test' }
                },
                mockSchema,
                {
                    relatedTask: {
                        taskId: 'task-456'
                    }
                }
            );

            // Access the private _taskMessageQueue to verify the message was queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            const queuedMessage = await queue!.dequeue('task-456');
            assertQueuedRequest(queuedMessage);
            expect(queuedMessage.message.method).toBe('test/request');

            // Verify resolver is stored in _requestResolvers map (not in the message)
            const requestId = queuedMessage.message.id as RequestId;
            const resolvers = (protocol as unknown as TestProtocol)._requestResolvers;
            expect(resolvers.has(requestId)).toBe(true);

            // Clean up the pending request
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: requestId,
                result: { result: 'success' }
            });
            await requestPromise;
        });
    });

    describe('server queues responses/errors for task-related requests', () => {
        it('should queue response when handling a request with relatedTask metadata', async () => {
            await protocol.connect(transport);

            // Set up a request handler that returns a result
            const TestRequestSchema = z.object({
                method: z.literal('test/taskRequest'),
                params: z
                    .object({
                        _meta: z.optional(z.record(z.unknown()))
                    })
                    .passthrough()
            });

            protocol.setRequestHandler(TestRequestSchema, async () => {
                return { content: 'test result' } as Result;
            });

            // Simulate an incoming request with relatedTask metadata
            const requestId = 456;
            const taskId = 'task-response-test';
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: requestId,
                method: 'test/taskRequest',
                params: {
                    _meta: {
                        'io.modelcontextprotocol/related-task': { taskId }
                    }
                }
            });

            // Wait for the handler to complete
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify the response was queued instead of sent directly
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            const queuedMessage = await queue!.dequeue(taskId);
            expect(queuedMessage).toBeDefined();
            expect(queuedMessage!.type).toBe('response');
            if (queuedMessage!.type === 'response') {
                expect(queuedMessage!.message.id).toBe(requestId);
                expect(queuedMessage!.message.result).toEqual({ content: 'test result' });
            }
        });

        it('should queue error when handling a request with relatedTask metadata that throws', async () => {
            await protocol.connect(transport);

            // Set up a request handler that throws an error
            const TestRequestSchema = z.object({
                method: z.literal('test/taskRequestError'),
                params: z
                    .object({
                        _meta: z.optional(z.record(z.unknown()))
                    })
                    .passthrough()
            });

            protocol.setRequestHandler(TestRequestSchema, async () => {
                throw new McpError(ErrorCode.InternalError, 'Test error message');
            });

            // Simulate an incoming request with relatedTask metadata
            const requestId = 789;
            const taskId = 'task-error-test';
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: requestId,
                method: 'test/taskRequestError',
                params: {
                    _meta: {
                        'io.modelcontextprotocol/related-task': { taskId }
                    }
                }
            });

            // Wait for the handler to complete
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify the error was queued instead of sent directly
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            const queuedMessage = await queue!.dequeue(taskId);
            expect(queuedMessage).toBeDefined();
            expect(queuedMessage!.type).toBe('error');
            if (queuedMessage!.type === 'error') {
                expect(queuedMessage!.message.id).toBe(requestId);
                expect(queuedMessage!.message.error.code).toBe(ErrorCode.InternalError);
                expect(queuedMessage!.message.error.message).toContain('Test error message');
            }
        });

        it('should queue MethodNotFound error for unknown method with relatedTask metadata', async () => {
            await protocol.connect(transport);

            // Simulate an incoming request for unknown method with relatedTask metadata
            const requestId = 101;
            const taskId = 'task-not-found-test';
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: requestId,
                method: 'unknown/method',
                params: {
                    _meta: {
                        'io.modelcontextprotocol/related-task': { taskId }
                    }
                }
            });

            // Wait for processing
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify the error was queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            const queuedMessage = await queue!.dequeue(taskId);
            expect(queuedMessage).toBeDefined();
            expect(queuedMessage!.type).toBe('error');
            if (queuedMessage!.type === 'error') {
                expect(queuedMessage!.message.id).toBe(requestId);
                expect(queuedMessage!.message.error.code).toBe(ErrorCode.MethodNotFound);
            }
        });

        it('should send response normally when request has no relatedTask metadata', async () => {
            await protocol.connect(transport);
            const sendSpy = vi.spyOn(transport, 'send');

            // Set up a request handler
            const TestRequestSchema = z.object({
                method: z.literal('test/normalRequest'),
                params: z.optional(z.record(z.unknown()))
            });

            protocol.setRequestHandler(TestRequestSchema, async () => {
                return { content: 'normal result' } as Result;
            });

            // Simulate an incoming request WITHOUT relatedTask metadata
            const requestId = 202;
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: requestId,
                method: 'test/normalRequest',
                params: {}
            });

            // Wait for the handler to complete
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify the response was sent through transport, not queued
            expect(sendSpy).toHaveBeenCalledWith(
                expect.objectContaining({
                    jsonrpc: '2.0',
                    id: requestId,
                    result: { content: 'normal result' }
                })
            );
        });
    });

    describe('messages without metadata bypass the queue', () => {
        it('should not queue notifications without relatedTask metadata', async () => {
            await protocol.connect(transport);

            // Send a notification without relatedTask metadata
            await protocol.notification({
                method: 'notifications/message',
                params: { level: 'info', data: 'test message' }
            });

            // Access the private _taskMessageQueue to verify no messages were queued
            // Since we can't check if queues exist without messages, we verify that
            // attempting to dequeue returns undefined (no messages queued)
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();
        });

        it('should not queue requests without relatedTask metadata', async () => {
            await protocol.connect(transport);

            const mockSchema = z.object({ result: z.string() });
            const sendSpy = vi.spyOn(transport, 'send');

            // Send a request without relatedTask metadata
            const requestPromise = protocol.request(
                {
                    method: 'test/request',
                    params: { data: 'test' }
                },
                mockSchema
            );

            // Access the private _taskMessageQueue to verify no messages were queued
            // Since we can't check if queues exist without messages, we verify that
            // attempting to dequeue returns undefined (no messages queued)
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Clean up the pending request
            const requestId = (sendSpy.mock.calls[0][0] as JSONRPCResultResponse).id;
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: requestId,
                result: { result: 'success' }
            });
            await requestPromise;
        });
    });

    describe('task ID extraction from metadata', () => {
        it('should extract correct task ID from relatedTask metadata for notifications', async () => {
            await protocol.connect(transport);

            const taskId = 'extracted-task-789';

            // Send a notification with relatedTask metadata
            await protocol.notification(
                {
                    method: 'notifications/message',
                    params: { data: 'test' }
                },
                {
                    relatedTask: {
                        taskId: taskId
                    }
                }
            );

            // Verify the message was queued under the correct task ID
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Verify a message was queued for this task
            const queuedMessage = await queue!.dequeue(taskId);
            assertQueuedNotification(queuedMessage);
            expect(queuedMessage.message.method).toBe('notifications/message');
        });

        it('should extract correct task ID from relatedTask metadata for requests', async () => {
            await protocol.connect(transport);

            const taskId = 'extracted-task-999';
            const mockSchema = z.object({ result: z.string() });

            // Send a request with relatedTask metadata
            const requestPromise = protocol.request(
                {
                    method: 'test/request',
                    params: { data: 'test' }
                },
                mockSchema,
                {
                    relatedTask: {
                        taskId: taskId
                    }
                }
            );

            // Verify the message was queued under the correct task ID
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Clean up the pending request
            const queuedMessage = await queue!.dequeue(taskId);
            assertQueuedRequest(queuedMessage);
            expect(queuedMessage.message.method).toBe('test/request');
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: queuedMessage.message.id,
                result: { result: 'success' }
            });
            await requestPromise;
        });

        it('should handle multiple messages for different task IDs', async () => {
            await protocol.connect(transport);

            // Send messages for different tasks
            await protocol.notification({ method: 'test1', params: {} }, { relatedTask: { taskId: 'task-A' } });
            await protocol.notification({ method: 'test2', params: {} }, { relatedTask: { taskId: 'task-B' } });
            await protocol.notification({ method: 'test3', params: {} }, { relatedTask: { taskId: 'task-A' } });

            // Verify messages are queued under correct task IDs
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Verify two messages for task-A
            const msg1A = await queue!.dequeue('task-A');
            const msg2A = await queue!.dequeue('task-A');
            const msg3A = await queue!.dequeue('task-A'); // Should be undefined
            expect(msg1A).toBeDefined();
            expect(msg2A).toBeDefined();
            expect(msg3A).toBeUndefined();

            // Verify one message for task-B
            const msg1B = await queue!.dequeue('task-B');
            const msg2B = await queue!.dequeue('task-B'); // Should be undefined
            expect(msg1B).toBeDefined();
            expect(msg2B).toBeUndefined();
        });
    });

    describe('queue creation on first message', () => {
        it('should queue messages for a task', async () => {
            await protocol.connect(transport);

            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Send first message for a task
            await protocol.notification({ method: 'test', params: {} }, { relatedTask: { taskId: 'new-task' } });

            // Verify message was queued
            const msg = await queue!.dequeue('new-task');
            assertQueuedNotification(msg);
            expect(msg.message.method).toBe('test');
        });

        it('should queue multiple messages for the same task', async () => {
            await protocol.connect(transport);

            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Send first message
            await protocol.notification({ method: 'test1', params: {} }, { relatedTask: { taskId: 'reuse-task' } });

            // Send second message
            await protocol.notification({ method: 'test2', params: {} }, { relatedTask: { taskId: 'reuse-task' } });

            // Verify both messages were queued in order
            const msg1 = await queue!.dequeue('reuse-task');
            const msg2 = await queue!.dequeue('reuse-task');
            assertQueuedNotification(msg1);
            expect(msg1.message.method).toBe('test1');
            assertQueuedNotification(msg2);
            expect(msg2.message.method).toBe('test2');
        });

        it('should queue messages for different tasks separately', async () => {
            await protocol.connect(transport);

            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Send messages for different tasks
            await protocol.notification({ method: 'test1', params: {} }, { relatedTask: { taskId: 'task-1' } });
            await protocol.notification({ method: 'test2', params: {} }, { relatedTask: { taskId: 'task-2' } });

            // Verify messages are queued separately
            const msg1 = await queue!.dequeue('task-1');
            const msg2 = await queue!.dequeue('task-2');
            assertQueuedNotification(msg1);
            expect(msg1?.message.method).toBe('test1');
            assertQueuedNotification(msg2);
            expect(msg2?.message.method).toBe('test2');
        });
    });

    describe('metadata preservation in queued messages', () => {
        it('should preserve relatedTask metadata in queued notification', async () => {
            await protocol.connect(transport);

            const relatedTask = { taskId: 'task-meta-123' };

            await protocol.notification(
                {
                    method: 'test/notification',
                    params: { data: 'test' }
                },
                { relatedTask }
            );

            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            const queuedMessage = await queue!.dequeue('task-meta-123');

            // Verify the metadata is preserved in the queued message
            expect(queuedMessage).toBeDefined();
            assertQueuedNotification(queuedMessage);
            expect(queuedMessage.message.params!._meta).toBeDefined();
            expect(queuedMessage.message.params!._meta![RELATED_TASK_META_KEY]).toEqual(relatedTask);
        });

        it('should preserve relatedTask metadata in queued request', async () => {
            await protocol.connect(transport);

            const relatedTask = { taskId: 'task-meta-456' };
            const mockSchema = z.object({ result: z.string() });

            const requestPromise = protocol.request(
                {
                    method: 'test/request',
                    params: { data: 'test' }
                },
                mockSchema,
                { relatedTask }
            );

            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            const queuedMessage = await queue!.dequeue('task-meta-456');

            // Verify the metadata is preserved in the queued message
            expect(queuedMessage).toBeDefined();
            assertQueuedRequest(queuedMessage);
            expect(queuedMessage.message.params!._meta).toBeDefined();
            expect(queuedMessage.message.params!._meta![RELATED_TASK_META_KEY]).toEqual(relatedTask);

            // Clean up
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: (queuedMessage!.message as JSONRPCRequest).id,
                result: { result: 'success' }
            });
            await requestPromise;
        });

        it('should preserve existing _meta fields when adding relatedTask', async () => {
            await protocol.connect(transport);

            await protocol.notification(
                {
                    method: 'test/notification',
                    params: {
                        data: 'test',
                        _meta: {
                            customField: 'customValue',
                            anotherField: 123
                        }
                    }
                },
                {
                    relatedTask: { taskId: 'task-preserve-meta' }
                }
            );

            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            const queuedMessage = await queue!.dequeue('task-preserve-meta');

            // Verify both existing and new metadata are preserved
            expect(queuedMessage).toBeDefined();
            assertQueuedNotification(queuedMessage);
            expect(queuedMessage.message.params!._meta!.customField).toBe('customValue');
            expect(queuedMessage.message.params!._meta!.anotherField).toBe(123);
            expect(queuedMessage.message.params!._meta![RELATED_TASK_META_KEY]).toEqual({
                taskId: 'task-preserve-meta'
            });
        });
    });
});

describe('Queue lifecycle management', () => {
    let protocol: Protocol<Request, Notification, Result>;
    let transport: MockTransport;
    let mockTaskStore: TaskStore & { [K in keyof TaskStore]: MockInstance };

    beforeEach(() => {
        transport = new MockTransport();
        mockTaskStore = createMockTaskStore();
        protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })({ taskStore: mockTaskStore, taskMessageQueue: new InMemoryTaskMessageQueue() });
    });

    describe('queue cleanup on task completion', () => {
        it('should clear queue when task reaches completed status', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await mockTaskStore.createTask({}, 1, { method: 'test', params: {} });
            const taskId = task.taskId;

            // Queue some messages for the task
            await protocol.notification({ method: 'test/notification', params: { data: 'test1' } }, { relatedTask: { taskId } });
            await protocol.notification({ method: 'test/notification', params: { data: 'test2' } }, { relatedTask: { taskId } });

            // Verify messages are queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Verify messages can be dequeued
            const msg1 = await queue!.dequeue(taskId);
            const msg2 = await queue!.dequeue(taskId);
            expect(msg1).toBeDefined();
            expect(msg2).toBeDefined();

            // Directly call the cleanup method (simulating what happens when task reaches terminal status)
            (protocol as unknown as TestProtocol)._clearTaskQueue(taskId);

            // After cleanup, no more messages should be available
            const msg3 = await queue!.dequeue(taskId);
            expect(msg3).toBeUndefined();
        });

        it('should clear queue after delivering messages on tasks/result for completed task', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await mockTaskStore.createTask({}, 1, { method: 'test', params: {} });
            const taskId = task.taskId;

            // Queue a message
            await protocol.notification({ method: 'test/notification', params: { data: 'test' } }, { relatedTask: { taskId } });

            // Mark task as completed
            const completedTask = { ...task, status: 'completed' as const };
            mockTaskStore.getTask.mockResolvedValue(completedTask);
            mockTaskStore.getTaskResult.mockResolvedValue({ content: [{ type: 'text', text: 'done' }] });

            // Simulate tasks/result request
            const resultPromise = new Promise(resolve => {
                transport.onmessage?.({
                    jsonrpc: '2.0',
                    id: 100,
                    method: 'tasks/result',
                    params: { taskId }
                });
                setTimeout(resolve, 50);
            });

            await resultPromise;

            // Verify queue is cleared after delivery (no messages available)
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            const msg = await queue!.dequeue(taskId);
            expect(msg).toBeUndefined();
        });
    });

    describe('queue cleanup on task cancellation', () => {
        it('should clear queue when task is cancelled', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await mockTaskStore.createTask({}, 1, { method: 'test', params: {} });
            const taskId = task.taskId;

            // Queue some messages
            await protocol.notification({ method: 'test/notification', params: { data: 'test1' } }, { relatedTask: { taskId } });

            // Verify message is queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            const msg1 = await queue!.dequeue(taskId);
            expect(msg1).toBeDefined();

            // Re-queue the message for cancellation test
            await protocol.notification({ method: 'test/notification', params: { data: 'test1' } }, { relatedTask: { taskId } });

            // Mock task as non-terminal
            mockTaskStore.getTask.mockResolvedValue(task);

            // Cancel the task
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 200,
                method: 'tasks/cancel',
                params: { taskId }
            });

            // Wait for cancellation to process
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify queue is cleared (no messages available)
            const msg2 = await queue!.dequeue(taskId);
            expect(msg2).toBeUndefined();
        });

        it('should reject pending request resolvers when task is cancelled', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await mockTaskStore.createTask({}, 1, { method: 'test', params: {} });
            const taskId = task.taskId;

            // Queue a request (catch rejection to avoid unhandled promise rejection)
            const requestPromise = protocol
                .request({ method: 'test/request', params: { data: 'test' } }, z.object({ result: z.string() }), {
                    relatedTask: { taskId }
                })
                .catch(err => err);

            // Verify request is queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Mock task as non-terminal
            mockTaskStore.getTask.mockResolvedValue(task);

            // Cancel the task
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 201,
                method: 'tasks/cancel',
                params: { taskId }
            });

            // Wait for cancellation to process
            await new Promise(resolve => setTimeout(resolve, 50));

            // Verify the request promise is rejected
            const result = await requestPromise;
            expect(result).toBeInstanceOf(McpError);
            expect(result.message).toContain('Task cancelled or completed');

            // Verify queue is cleared (no messages available)
            const msg = await queue!.dequeue(taskId);
            expect(msg).toBeUndefined();
        });
    });

    describe('queue cleanup on task failure', () => {
        it('should clear queue when task reaches failed status', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await mockTaskStore.createTask({}, 1, { method: 'test', params: {} });
            const taskId = task.taskId;

            // Queue some messages
            await protocol.notification({ method: 'test/notification', params: { data: 'test1' } }, { relatedTask: { taskId } });
            await protocol.notification({ method: 'test/notification', params: { data: 'test2' } }, { relatedTask: { taskId } });

            // Verify messages are queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Verify messages can be dequeued
            const msg1 = await queue!.dequeue(taskId);
            const msg2 = await queue!.dequeue(taskId);
            expect(msg1).toBeDefined();
            expect(msg2).toBeDefined();

            // Directly call the cleanup method (simulating what happens when task reaches terminal status)
            (protocol as unknown as TestProtocol)._clearTaskQueue(taskId);

            // After cleanup, no more messages should be available
            const msg3 = await queue!.dequeue(taskId);
            expect(msg3).toBeUndefined();
        });

        it('should reject pending request resolvers when task fails', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await mockTaskStore.createTask({}, 1, { method: 'test', params: {} });
            const taskId = task.taskId;

            // Queue a request (catch the rejection to avoid unhandled promise rejection)
            const requestPromise = protocol
                .request({ method: 'test/request', params: { data: 'test' } }, z.object({ result: z.string() }), {
                    relatedTask: { taskId }
                })
                .catch(err => err);

            // Verify request is queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Directly call the cleanup method (simulating what happens when task reaches terminal status)
            (protocol as unknown as TestProtocol)._clearTaskQueue(taskId);

            // Verify the request promise is rejected
            const result = await requestPromise;
            expect(result).toBeInstanceOf(McpError);
            expect(result.message).toContain('Task cancelled or completed');

            // Verify queue is cleared (no messages available)
            const msg = await queue!.dequeue(taskId);
            expect(msg).toBeUndefined();
        });
    });

    describe('resolver rejection on cleanup', () => {
        it('should reject all pending request resolvers when queue is cleared', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await mockTaskStore.createTask({}, 1, { method: 'test', params: {} });
            const taskId = task.taskId;

            // Queue multiple requests (catch rejections to avoid unhandled promise rejections)
            const request1Promise = protocol
                .request({ method: 'test/request1', params: { data: 'test1' } }, z.object({ result: z.string() }), {
                    relatedTask: { taskId }
                })
                .catch(err => err);

            const request2Promise = protocol
                .request({ method: 'test/request2', params: { data: 'test2' } }, z.object({ result: z.string() }), {
                    relatedTask: { taskId }
                })
                .catch(err => err);

            const request3Promise = protocol
                .request({ method: 'test/request3', params: { data: 'test3' } }, z.object({ result: z.string() }), {
                    relatedTask: { taskId }
                })
                .catch(err => err);

            // Verify requests are queued
            const queue = (protocol as unknown as TestProtocol)._taskMessageQueue;
            expect(queue).toBeDefined();

            // Directly call the cleanup method (simulating what happens when task reaches terminal status)
            (protocol as unknown as TestProtocol)._clearTaskQueue(taskId);

            // Verify all request promises are rejected
            const result1 = await request1Promise;
            const result2 = await request2Promise;
            const result3 = await request3Promise;

            expect(result1).toBeInstanceOf(McpError);
            expect(result1.message).toContain('Task cancelled or completed');
            expect(result2).toBeInstanceOf(McpError);
            expect(result2.message).toContain('Task cancelled or completed');
            expect(result3).toBeInstanceOf(McpError);
            expect(result3.message).toContain('Task cancelled or completed');

            // Verify queue is cleared (no messages available)
            const msg = await queue!.dequeue(taskId);
            expect(msg).toBeUndefined();
        });

        it('should clean up resolver mappings when rejecting requests', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await mockTaskStore.createTask({}, 1, { method: 'test', params: {} });
            const taskId = task.taskId;

            // Queue a request (catch rejection to avoid unhandled promise rejection)
            const requestPromise = protocol
                .request({ method: 'test/request', params: { data: 'test' } }, z.object({ result: z.string() }), {
                    relatedTask: { taskId }
                })
                .catch(err => err);

            // Get the request ID that was sent
            const requestResolvers = (protocol as unknown as TestProtocol)._requestResolvers;
            const initialResolverCount = requestResolvers.size;
            expect(initialResolverCount).toBeGreaterThan(0);

            // Complete the task (triggers cleanup)
            const completedTask = { ...task, status: 'completed' as const };
            mockTaskStore.getTask.mockResolvedValue(completedTask);

            // Directly call the cleanup method (simulating what happens when task reaches terminal status)
            (protocol as unknown as TestProtocol)._clearTaskQueue(taskId);

            // Verify request promise is rejected
            const result = await requestPromise;
            expect(result).toBeInstanceOf(McpError);
            expect(result.message).toContain('Task cancelled or completed');

            // Verify resolver mapping is cleaned up
            // The resolver should be removed from the map
            expect(requestResolvers.size).toBeLessThan(initialResolverCount);
        });
    });
});

describe('requestStream() method', () => {
    const CallToolResultSchema = z.object({
        content: z.array(z.object({ type: z.string(), text: z.string() })),
        _meta: z.object({}).optional()
    });

    test('should yield result immediately for non-task requests', async () => {
        const transport = new MockTransport();
        const protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })();
        await protocol.connect(transport);

        // Start the request stream
        const streamPromise = (async () => {
            const messages = [];
            const stream = (protocol as unknown as TestProtocol).requestStream(
                { method: 'tools/call', params: { name: 'test', arguments: {} } },
                CallToolResultSchema
            );
            for await (const message of stream) {
                messages.push(message);
            }
            return messages;
        })();

        // Simulate server response
        await new Promise(resolve => setTimeout(resolve, 10));
        transport.onmessage?.({
            jsonrpc: '2.0',
            id: 0,
            result: {
                content: [{ type: 'text', text: 'test result' }],
                _meta: {}
            }
        });

        const messages = await streamPromise;

        // Should yield exactly one result message
        expect(messages).toHaveLength(1);
        expect(messages[0].type).toBe('result');
        expect(messages[0]).toHaveProperty('result');
    });

    test('should yield error message on request failure', async () => {
        const transport = new MockTransport();
        const protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })();
        await protocol.connect(transport);

        // Start the request stream
        const streamPromise = (async () => {
            const messages = [];
            const stream = (protocol as unknown as TestProtocol).requestStream(
                { method: 'tools/call', params: { name: 'test', arguments: {} } },
                CallToolResultSchema
            );
            for await (const message of stream) {
                messages.push(message);
            }
            return messages;
        })();

        // Simulate server error response
        await new Promise(resolve => setTimeout(resolve, 10));
        transport.onmessage?.({
            jsonrpc: '2.0',
            id: 0,
            error: {
                code: ErrorCode.InternalError,
                message: 'Test error'
            }
        });

        const messages = await streamPromise;

        // Should yield exactly one error message
        expect(messages).toHaveLength(1);
        expect(messages[0].type).toBe('error');
        expect(messages[0]).toHaveProperty('error');
        if (messages[0].type === 'error') {
            expect(messages[0].error.message).toContain('Test error');
        }
    });

    test('should handle cancellation via AbortSignal', async () => {
        const transport = new MockTransport();
        const protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })();
        await protocol.connect(transport);

        const abortController = new AbortController();

        // Abort immediately before starting the stream
        abortController.abort('User cancelled');

        // Start the request stream with already-aborted signal
        const messages = [];
        const stream = (protocol as unknown as TestProtocol).requestStream(
            { method: 'tools/call', params: { name: 'test', arguments: {} } },
            CallToolResultSchema,
            {
                signal: abortController.signal
            }
        );
        for await (const message of stream) {
            messages.push(message);
        }

        // Should yield error message about cancellation
        expect(messages).toHaveLength(1);
        expect(messages[0].type).toBe('error');
        if (messages[0].type === 'error') {
            expect(messages[0].error.message).toContain('cancelled');
        }
    });

    describe('Error responses', () => {
        test('should yield error as terminal message for server error response', async () => {
            const transport = new MockTransport();
            const protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })();
            await protocol.connect(transport);

            const messagesPromise = toArrayAsync(
                (protocol as unknown as TestProtocol).requestStream(
                    { method: 'tools/call', params: { name: 'test', arguments: {} } },
                    CallToolResultSchema
                )
            );

            // Simulate server error response
            await new Promise(resolve => setTimeout(resolve, 10));
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 0,
                error: {
                    code: ErrorCode.InternalError,
                    message: 'Server error'
                }
            });

            // Collect messages
            const messages = await messagesPromise;

            // Verify error is terminal and last message
            expect(messages.length).toBeGreaterThan(0);
            const lastMessage = messages[messages.length - 1];
            assertErrorResponse(lastMessage);
            expect(lastMessage.error).toBeDefined();
            expect(lastMessage.error.message).toContain('Server error');
        });

        test('should yield error as terminal message for timeout', async () => {
            vi.useFakeTimers();
            try {
                const transport = new MockTransport();
                const protocol = new (class extends Protocol<Request, Notification, Result> {
                    protected assertCapabilityForMethod(): void {}
                    protected assertNotificationCapability(): void {}
                    protected assertRequestHandlerCapability(): void {}
                    protected assertTaskCapability(): void {}
                    protected assertTaskHandlerCapability(): void {}
                })();
                await protocol.connect(transport);

                const messagesPromise = toArrayAsync(
                    (protocol as unknown as TestProtocol).requestStream(
                        { method: 'tools/call', params: { name: 'test', arguments: {} } },
                        CallToolResultSchema,
                        {
                            timeout: 100
                        }
                    )
                );

                // Advance time to trigger timeout
                await vi.advanceTimersByTimeAsync(101);

                // Collect messages
                const messages = await messagesPromise;

                // Verify error is terminal and last message
                expect(messages.length).toBeGreaterThan(0);
                const lastMessage = messages[messages.length - 1];
                assertErrorResponse(lastMessage);
                expect(lastMessage.error).toBeDefined();
                expect(lastMessage.error.code).toBe(ErrorCode.RequestTimeout);
            } finally {
                vi.useRealTimers();
            }
        });

        test('should yield error as terminal message for cancellation', async () => {
            const transport = new MockTransport();
            const protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })();
            await protocol.connect(transport);

            const abortController = new AbortController();
            abortController.abort('User cancelled');

            // Collect messages
            const messages = await toArrayAsync(
                (protocol as unknown as TestProtocol).requestStream(
                    { method: 'tools/call', params: { name: 'test', arguments: {} } },
                    CallToolResultSchema,
                    {
                        signal: abortController.signal
                    }
                )
            );

            // Verify error is terminal and last message
            expect(messages.length).toBeGreaterThan(0);
            const lastMessage = messages[messages.length - 1];
            assertErrorResponse(lastMessage);
            expect(lastMessage.error).toBeDefined();
            expect(lastMessage.error.message).toContain('cancelled');
        });

        test('should not yield any messages after error message', async () => {
            const transport = new MockTransport();
            const protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })();
            await protocol.connect(transport);

            const messagesPromise = toArrayAsync(
                (protocol as unknown as TestProtocol).requestStream(
                    { method: 'tools/call', params: { name: 'test', arguments: {} } },
                    CallToolResultSchema
                )
            );

            // Simulate server error response
            await new Promise(resolve => setTimeout(resolve, 10));
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 0,
                error: {
                    code: ErrorCode.InternalError,
                    message: 'Test error'
                }
            });

            // Collect messages
            const messages = await messagesPromise;

            // Verify only one message (the error) was yielded
            expect(messages).toHaveLength(1);
            expect(messages[0].type).toBe('error');

            // Try to send another message (should be ignored)
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 0,
                result: {
                    content: [{ type: 'text', text: 'should not appear' }]
                }
            });

            await new Promise(resolve => setTimeout(resolve, 10));

            // Verify no additional messages were yielded
            expect(messages).toHaveLength(1);
        });

        test('should yield error as terminal message for task failure', async () => {
            const transport = new MockTransport();
            const mockTaskStore = createMockTaskStore();
            const protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })({ taskStore: mockTaskStore });
            await protocol.connect(transport);

            const messagesPromise = toArrayAsync(
                (protocol as unknown as TestProtocol).requestStream(
                    { method: 'tools/call', params: { name: 'test', arguments: {} } },
                    CallToolResultSchema
                )
            );

            // Simulate task creation response
            await new Promise(resolve => setTimeout(resolve, 10));
            const taskId = 'test-task-123';
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 0,
                result: {
                    _meta: {
                        task: {
                            taskId,
                            status: 'working',
                            createdAt: new Date().toISOString(),
                            pollInterval: 100
                        }
                    }
                }
            });

            // Wait for task creation to be processed
            await new Promise(resolve => setTimeout(resolve, 20));

            // Update task to failed status
            const failedTask = {
                taskId,
                status: 'failed' as const,
                createdAt: new Date().toISOString(),
                pollInterval: 100,
                ttl: null,
                statusMessage: 'Task failed'
            };
            mockTaskStore.getTask.mockResolvedValue(failedTask);

            // Collect messages
            const messages = await messagesPromise;

            // Verify error is terminal and last message
            expect(messages.length).toBeGreaterThan(0);
            const lastMessage = messages[messages.length - 1];
            assertErrorResponse(lastMessage);
            expect(lastMessage.error).toBeDefined();
        });

        test('should yield error as terminal message for network error', async () => {
            const transport = new MockTransport();
            const protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })();
            await protocol.connect(transport);

            // Override send to simulate network error
            transport.send = vi.fn().mockRejectedValue(new Error('Network error'));

            const messages = await toArrayAsync(
                (protocol as unknown as TestProtocol).requestStream(
                    { method: 'tools/call', params: { name: 'test', arguments: {} } },
                    CallToolResultSchema
                )
            );

            // Verify error is terminal and last message
            expect(messages.length).toBeGreaterThan(0);
            const lastMessage = messages[messages.length - 1];
            assertErrorResponse(lastMessage);
            expect(lastMessage.error).toBeDefined();
        });

        test('should ensure error is always the final message', async () => {
            const transport = new MockTransport();
            const protocol = new (class extends Protocol<Request, Notification, Result> {
                protected assertCapabilityForMethod(): void {}
                protected assertNotificationCapability(): void {}
                protected assertRequestHandlerCapability(): void {}
                protected assertTaskCapability(): void {}
                protected assertTaskHandlerCapability(): void {}
            })();
            await protocol.connect(transport);

            const messagesPromise = toArrayAsync(
                (protocol as unknown as TestProtocol).requestStream(
                    { method: 'tools/call', params: { name: 'test', arguments: {} } },
                    CallToolResultSchema
                )
            );

            // Simulate server error response
            await new Promise(resolve => setTimeout(resolve, 10));
            transport.onmessage?.({
                jsonrpc: '2.0',
                id: 0,
                error: {
                    code: ErrorCode.InternalError,
                    message: 'Test error'
                }
            });

            // Collect messages
            const messages = await messagesPromise;

            // Verify error is the last message
            expect(messages.length).toBeGreaterThan(0);
            const lastMessage = messages[messages.length - 1];
            expect(lastMessage.type).toBe('error');

            // Verify all messages before the last are not terminal
            for (let i = 0; i < messages.length - 1; i++) {
                expect(messages[i].type).not.toBe('error');
                expect(messages[i].type).not.toBe('result');
            }
        });
    });
});

describe('Error handling for missing resolvers', () => {
    let protocol: Protocol<Request, Notification, Result>;
    let transport: MockTransport;
    let taskStore: TaskStore & { [K in keyof TaskStore]: MockInstance };
    let taskMessageQueue: TaskMessageQueue;
    let errorHandler: MockInstance;

    beforeEach(() => {
        taskStore = createMockTaskStore();
        taskMessageQueue = new InMemoryTaskMessageQueue();
        errorHandler = vi.fn();

        protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(_method: string): void {}
            protected assertNotificationCapability(_method: string): void {}
            protected assertRequestHandlerCapability(_method: string): void {}
            protected assertTaskCapability(_method: string): void {}
            protected assertTaskHandlerCapability(_method: string): void {}
        })({
            taskStore,
            taskMessageQueue,
            defaultTaskPollInterval: 100
        });

        // @ts-expect-error deliberately overriding error handler with mock
        protocol.onerror = errorHandler;
        transport = new MockTransport();
    });

    describe('Response routing with missing resolvers', () => {
        it('should log error for unknown request ID without throwing', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });

            // Enqueue a response message without a corresponding resolver
            await taskMessageQueue.enqueue(task.taskId, {
                type: 'response',
                message: {
                    jsonrpc: '2.0',
                    id: 999, // Non-existent request ID
                    result: { content: [] }
                },
                timestamp: Date.now()
            });

            // Set up the GetTaskPayloadRequest handler to process the message
            const testProtocol = protocol as unknown as TestProtocol;

            // Simulate dequeuing and processing the response
            const queuedMessage = await taskMessageQueue.dequeue(task.taskId);
            expect(queuedMessage).toBeDefined();
            expect(queuedMessage?.type).toBe('response');

            // Manually trigger the response handling logic
            if (queuedMessage && queuedMessage.type === 'response') {
                const responseMessage = queuedMessage.message as JSONRPCResultResponse;
                const requestId = responseMessage.id as RequestId;
                const resolver = testProtocol._requestResolvers.get(requestId);

                if (!resolver) {
                    // This simulates what happens in the actual handler
                    protocol.onerror?.(new Error(`Response handler missing for request ${requestId}`));
                }
            }

            // Verify error was logged
            expect(errorHandler).toHaveBeenCalledWith(
                expect.objectContaining({
                    message: expect.stringContaining('Response handler missing for request 999')
                })
            );
        });

        it('should continue processing after missing resolver error', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });

            // Enqueue a response with missing resolver, then a valid notification
            await taskMessageQueue.enqueue(task.taskId, {
                type: 'response',
                message: {
                    jsonrpc: '2.0',
                    id: 999,
                    result: { content: [] }
                },
                timestamp: Date.now()
            });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'notification',
                message: {
                    jsonrpc: '2.0',
                    method: 'notifications/progress',
                    params: { progress: 50, total: 100 }
                },
                timestamp: Date.now()
            });

            // Process first message (response with missing resolver)
            const msg1 = await taskMessageQueue.dequeue(task.taskId);
            expect(msg1?.type).toBe('response');

            // Process second message (should work fine)
            const msg2 = await taskMessageQueue.dequeue(task.taskId);
            expect(msg2?.type).toBe('notification');
            expect(msg2?.message).toMatchObject({
                method: 'notifications/progress'
            });
        });
    });

    describe('Task cancellation with missing resolvers', () => {
        it('should log error when resolver is missing during cleanup', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });

            // Enqueue a request without storing a resolver
            await taskMessageQueue.enqueue(task.taskId, {
                type: 'request',
                message: {
                    jsonrpc: '2.0',
                    id: 42,
                    method: 'tools/call',
                    params: { name: 'test-tool', arguments: {} }
                },
                timestamp: Date.now()
            });

            // Clear the task queue (simulating cancellation)
            const testProtocol = protocol as unknown as TestProtocol;
            await testProtocol._clearTaskQueue(task.taskId);

            // Verify error was logged for missing resolver
            expect(errorHandler).toHaveBeenCalledWith(
                expect.objectContaining({
                    message: expect.stringContaining('Resolver missing for request 42')
                })
            );
        });

        it('should handle cleanup gracefully when resolver exists', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });

            const requestId = 42;
            const resolverMock = vi.fn();

            // Store a resolver
            const testProtocol = protocol as unknown as TestProtocol;
            testProtocol._requestResolvers.set(requestId, resolverMock);

            // Enqueue a request
            await taskMessageQueue.enqueue(task.taskId, {
                type: 'request',
                message: {
                    jsonrpc: '2.0',
                    id: requestId,
                    method: 'tools/call',
                    params: { name: 'test-tool', arguments: {} }
                },
                timestamp: Date.now()
            });

            // Clear the task queue
            await testProtocol._clearTaskQueue(task.taskId);

            // Verify resolver was called with cancellation error
            expect(resolverMock).toHaveBeenCalledWith(expect.any(McpError));

            // Verify the error has the correct properties
            const calledError = resolverMock.mock.calls[0][0];
            expect(calledError.code).toBe(ErrorCode.InternalError);
            expect(calledError.message).toContain('Task cancelled or completed');

            // Verify resolver was removed
            expect(testProtocol._requestResolvers.has(requestId)).toBe(false);
        });

        it('should handle mixed messages during cleanup', async () => {
            await protocol.connect(transport);

            // Create a task
            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });

            const testProtocol = protocol as unknown as TestProtocol;

            // Enqueue multiple messages: request with resolver, request without, notification
            const requestId1 = 42;
            const resolverMock = vi.fn();
            testProtocol._requestResolvers.set(requestId1, resolverMock);

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'request',
                message: {
                    jsonrpc: '2.0',
                    id: requestId1,
                    method: 'tools/call',
                    params: { name: 'test-tool', arguments: {} }
                },
                timestamp: Date.now()
            });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'request',
                message: {
                    jsonrpc: '2.0',
                    id: 43, // No resolver for this one
                    method: 'tools/call',
                    params: { name: 'test-tool', arguments: {} }
                },
                timestamp: Date.now()
            });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'notification',
                message: {
                    jsonrpc: '2.0',
                    method: 'notifications/progress',
                    params: { progress: 50, total: 100 }
                },
                timestamp: Date.now()
            });

            // Clear the task queue
            await testProtocol._clearTaskQueue(task.taskId);

            // Verify resolver was called for first request
            expect(resolverMock).toHaveBeenCalledWith(expect.any(McpError));

            // Verify the error has the correct properties
            const calledError = resolverMock.mock.calls[0][0];
            expect(calledError.code).toBe(ErrorCode.InternalError);
            expect(calledError.message).toContain('Task cancelled or completed');

            // Verify error was logged for second request
            expect(errorHandler).toHaveBeenCalledWith(
                expect.objectContaining({
                    message: expect.stringContaining('Resolver missing for request 43')
                })
            );

            // Verify queue is empty
            const remaining = await taskMessageQueue.dequeue(task.taskId);
            expect(remaining).toBeUndefined();
        });
    });

    describe('Side-channeled request error handling', () => {
        it('should log error when response handler is missing for side-channeled request', async () => {
            await protocol.connect(transport);

            const testProtocol = protocol as unknown as TestProtocol;
            const messageId = 123;

            // Create a response resolver without a corresponding response handler
            const responseResolver = (response: JSONRPCResultResponse | Error) => {
                const handler = testProtocol._responseHandlers.get(messageId);
                if (handler) {
                    handler(response);
                } else {
                    protocol.onerror?.(new Error(`Response handler missing for side-channeled request ${messageId}`));
                }
            };

            // Simulate the resolver being called without a handler
            const mockResponse: JSONRPCResultResponse = {
                jsonrpc: '2.0',
                id: messageId,
                result: { content: [] }
            };

            responseResolver(mockResponse);

            // Verify error was logged
            expect(errorHandler).toHaveBeenCalledWith(
                expect.objectContaining({
                    message: expect.stringContaining('Response handler missing for side-channeled request 123')
                })
            );
        });
    });

    describe('Error handling does not throw exceptions', () => {
        it('should not throw when processing response with missing resolver', async () => {
            await protocol.connect(transport);

            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'response',
                message: {
                    jsonrpc: '2.0',
                    id: 999,
                    result: { content: [] }
                },
                timestamp: Date.now()
            });

            // This should not throw
            const processMessage = async () => {
                const msg = await taskMessageQueue.dequeue(task.taskId);
                if (msg && msg.type === 'response') {
                    const testProtocol = protocol as unknown as TestProtocol;
                    const responseMessage = msg.message as JSONRPCResultResponse;
                    const requestId = responseMessage.id as RequestId;
                    const resolver = testProtocol._requestResolvers.get(requestId);
                    if (!resolver) {
                        protocol.onerror?.(new Error(`Response handler missing for request ${requestId}`));
                    }
                }
            };

            await expect(processMessage()).resolves.not.toThrow();
        });

        it('should not throw during task cleanup with missing resolvers', async () => {
            await protocol.connect(transport);

            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'request',
                message: {
                    jsonrpc: '2.0',
                    id: 42,
                    method: 'tools/call',
                    params: { name: 'test-tool', arguments: {} }
                },
                timestamp: Date.now()
            });

            const testProtocol = protocol as unknown as TestProtocol;

            // This should not throw
            await expect(testProtocol._clearTaskQueue(task.taskId)).resolves.not.toThrow();
        });
    });

    describe('Error message routing', () => {
        it('should route error messages to resolvers correctly', async () => {
            await protocol.connect(transport);

            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });
            const requestId = 42;
            const resolverMock = vi.fn();

            // Store a resolver
            const testProtocol = protocol as unknown as TestProtocol;
            testProtocol._requestResolvers.set(requestId, resolverMock);

            // Enqueue an error message
            await taskMessageQueue.enqueue(task.taskId, {
                type: 'error',
                message: {
                    jsonrpc: '2.0',
                    id: requestId,
                    error: {
                        code: ErrorCode.InvalidRequest,
                        message: 'Invalid request parameters'
                    }
                },
                timestamp: Date.now()
            });

            // Simulate dequeuing and processing the error
            const queuedMessage = await taskMessageQueue.dequeue(task.taskId);
            expect(queuedMessage).toBeDefined();
            expect(queuedMessage?.type).toBe('error');

            // Manually trigger the error handling logic
            if (queuedMessage && queuedMessage.type === 'error') {
                const errorMessage = queuedMessage.message as JSONRPCErrorResponse;
                const reqId = errorMessage.id as RequestId;
                const resolver = testProtocol._requestResolvers.get(reqId);

                if (resolver) {
                    testProtocol._requestResolvers.delete(reqId);
                    const error = new McpError(errorMessage.error.code, errorMessage.error.message, errorMessage.error.data);
                    resolver(error);
                }
            }

            // Verify resolver was called with McpError
            expect(resolverMock).toHaveBeenCalledWith(expect.any(McpError));
            const calledError = resolverMock.mock.calls[0][0];
            expect(calledError.code).toBe(ErrorCode.InvalidRequest);
            expect(calledError.message).toContain('Invalid request parameters');

            // Verify resolver was removed from map
            expect(testProtocol._requestResolvers.has(requestId)).toBe(false);
        });

        it('should log error for unknown request ID in error messages', async () => {
            await protocol.connect(transport);

            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });

            // Enqueue an error message without a corresponding resolver
            await taskMessageQueue.enqueue(task.taskId, {
                type: 'error',
                message: {
                    jsonrpc: '2.0',
                    id: 999,
                    error: {
                        code: ErrorCode.InternalError,
                        message: 'Something went wrong'
                    }
                },
                timestamp: Date.now()
            });

            // Simulate dequeuing and processing the error
            const queuedMessage = await taskMessageQueue.dequeue(task.taskId);
            expect(queuedMessage).toBeDefined();
            expect(queuedMessage?.type).toBe('error');

            // Manually trigger the error handling logic
            if (queuedMessage && queuedMessage.type === 'error') {
                const testProtocol = protocol as unknown as TestProtocol;
                const errorMessage = queuedMessage.message as JSONRPCErrorResponse;
                const requestId = errorMessage.id as RequestId;
                const resolver = testProtocol._requestResolvers.get(requestId);

                if (!resolver) {
                    protocol.onerror?.(new Error(`Error handler missing for request ${requestId}`));
                }
            }

            // Verify error was logged
            expect(errorHandler).toHaveBeenCalledWith(
                expect.objectContaining({
                    message: expect.stringContaining('Error handler missing for request 999')
                })
            );
        });

        it('should handle error messages with data field', async () => {
            await protocol.connect(transport);

            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });
            const requestId = 42;
            const resolverMock = vi.fn();

            // Store a resolver
            const testProtocol = protocol as unknown as TestProtocol;
            testProtocol._requestResolvers.set(requestId, resolverMock);

            // Enqueue an error message with data field
            await taskMessageQueue.enqueue(task.taskId, {
                type: 'error',
                message: {
                    jsonrpc: '2.0',
                    id: requestId,
                    error: {
                        code: ErrorCode.InvalidParams,
                        message: 'Validation failed',
                        data: { field: 'userName', reason: 'required' }
                    }
                },
                timestamp: Date.now()
            });

            // Simulate dequeuing and processing the error
            const queuedMessage = await taskMessageQueue.dequeue(task.taskId);

            if (queuedMessage && queuedMessage.type === 'error') {
                const errorMessage = queuedMessage.message as JSONRPCErrorResponse;
                const reqId = errorMessage.id as RequestId;
                const resolver = testProtocol._requestResolvers.get(reqId);

                if (resolver) {
                    testProtocol._requestResolvers.delete(reqId);
                    const error = new McpError(errorMessage.error.code, errorMessage.error.message, errorMessage.error.data);
                    resolver(error);
                }
            }

            // Verify resolver was called with McpError including data
            expect(resolverMock).toHaveBeenCalledWith(expect.any(McpError));
            const calledError = resolverMock.mock.calls[0][0];
            expect(calledError.code).toBe(ErrorCode.InvalidParams);
            expect(calledError.message).toContain('Validation failed');
            expect(calledError.data).toEqual({ field: 'userName', reason: 'required' });
        });

        it('should not throw when processing error with missing resolver', async () => {
            await protocol.connect(transport);

            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'error',
                message: {
                    jsonrpc: '2.0',
                    id: 999,
                    error: {
                        code: ErrorCode.InternalError,
                        message: 'Error occurred'
                    }
                },
                timestamp: Date.now()
            });

            // This should not throw
            const processMessage = async () => {
                const msg = await taskMessageQueue.dequeue(task.taskId);
                if (msg && msg.type === 'error') {
                    const testProtocol = protocol as unknown as TestProtocol;
                    const errorMessage = msg.message as JSONRPCErrorResponse;
                    const requestId = errorMessage.id as RequestId;
                    const resolver = testProtocol._requestResolvers.get(requestId);
                    if (!resolver) {
                        protocol.onerror?.(new Error(`Error handler missing for request ${requestId}`));
                    }
                }
            };

            await expect(processMessage()).resolves.not.toThrow();
        });
    });

    describe('Response and error message routing integration', () => {
        it('should handle mixed response and error messages in queue', async () => {
            await protocol.connect(transport);

            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });
            const testProtocol = protocol as unknown as TestProtocol;

            // Set up resolvers for multiple requests
            const resolver1 = vi.fn();
            const resolver2 = vi.fn();
            const resolver3 = vi.fn();

            testProtocol._requestResolvers.set(1, resolver1);
            testProtocol._requestResolvers.set(2, resolver2);
            testProtocol._requestResolvers.set(3, resolver3);

            // Enqueue mixed messages: response, error, response
            await taskMessageQueue.enqueue(task.taskId, {
                type: 'response',
                message: {
                    jsonrpc: '2.0',
                    id: 1,
                    result: { content: [{ type: 'text', text: 'Success' }] }
                },
                timestamp: Date.now()
            });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'error',
                message: {
                    jsonrpc: '2.0',
                    id: 2,
                    error: {
                        code: ErrorCode.InvalidRequest,
                        message: 'Request failed'
                    }
                },
                timestamp: Date.now()
            });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'response',
                message: {
                    jsonrpc: '2.0',
                    id: 3,
                    result: { content: [{ type: 'text', text: 'Another success' }] }
                },
                timestamp: Date.now()
            });

            // Process all messages
            let msg;
            while ((msg = await taskMessageQueue.dequeue(task.taskId))) {
                if (msg.type === 'response') {
                    const responseMessage = msg.message as JSONRPCResultResponse;
                    const requestId = responseMessage.id as RequestId;
                    const resolver = testProtocol._requestResolvers.get(requestId);
                    if (resolver) {
                        testProtocol._requestResolvers.delete(requestId);
                        resolver(responseMessage);
                    }
                } else if (msg.type === 'error') {
                    const errorMessage = msg.message as JSONRPCErrorResponse;
                    const requestId = errorMessage.id as RequestId;
                    const resolver = testProtocol._requestResolvers.get(requestId);
                    if (resolver) {
                        testProtocol._requestResolvers.delete(requestId);
                        const error = new McpError(errorMessage.error.code, errorMessage.error.message, errorMessage.error.data);
                        resolver(error);
                    }
                }
            }

            // Verify all resolvers were called correctly
            expect(resolver1).toHaveBeenCalledWith(expect.objectContaining({ id: 1 }));
            expect(resolver2).toHaveBeenCalledWith(expect.any(McpError));
            expect(resolver3).toHaveBeenCalledWith(expect.objectContaining({ id: 3 }));

            // Verify error has correct properties
            const error = resolver2.mock.calls[0][0];
            expect(error.code).toBe(ErrorCode.InvalidRequest);
            expect(error.message).toContain('Request failed');

            // Verify all resolvers were removed
            expect(testProtocol._requestResolvers.size).toBe(0);
        });

        it('should maintain FIFO order when processing responses and errors', async () => {
            await protocol.connect(transport);

            const task = await taskStore.createTask({ ttl: 60000 }, 1, { method: 'test', params: {} });
            const testProtocol = protocol as unknown as TestProtocol;

            const callOrder: number[] = [];
            const resolver1 = vi.fn(() => callOrder.push(1));
            const resolver2 = vi.fn(() => callOrder.push(2));
            const resolver3 = vi.fn(() => callOrder.push(3));

            testProtocol._requestResolvers.set(1, resolver1);
            testProtocol._requestResolvers.set(2, resolver2);
            testProtocol._requestResolvers.set(3, resolver3);

            // Enqueue in specific order
            await taskMessageQueue.enqueue(task.taskId, {
                type: 'response',
                message: { jsonrpc: '2.0', id: 1, result: {} },
                timestamp: 1000
            });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'error',
                message: {
                    jsonrpc: '2.0',
                    id: 2,
                    error: { code: -32600, message: 'Error' }
                },
                timestamp: 2000
            });

            await taskMessageQueue.enqueue(task.taskId, {
                type: 'response',
                message: { jsonrpc: '2.0', id: 3, result: {} },
                timestamp: 3000
            });

            // Process all messages
            let msg;
            while ((msg = await taskMessageQueue.dequeue(task.taskId))) {
                if (msg.type === 'response') {
                    const responseMessage = msg.message as JSONRPCResultResponse;
                    const requestId = responseMessage.id as RequestId;
                    const resolver = testProtocol._requestResolvers.get(requestId);
                    if (resolver) {
                        testProtocol._requestResolvers.delete(requestId);
                        resolver(responseMessage);
                    }
                } else if (msg.type === 'error') {
                    const errorMessage = msg.message as JSONRPCErrorResponse;
                    const requestId = errorMessage.id as RequestId;
                    const resolver = testProtocol._requestResolvers.get(requestId);
                    if (resolver) {
                        testProtocol._requestResolvers.delete(requestId);
                        const error = new McpError(errorMessage.error.code, errorMessage.error.message, errorMessage.error.data);
                        resolver(error);
                    }
                }
            }

            // Verify FIFO order was maintained
            expect(callOrder).toEqual([1, 2, 3]);
        });
    });
});

describe('_onclose cleanup', () => {
    let protocol: Protocol<Request, Notification, Result>;
    let transport: MockTransport;
    let sendSpy: MockInstance;

    beforeEach(() => {
        vi.useFakeTimers();
        transport = new MockTransport();
        sendSpy = vi.spyOn(transport, 'send');
        protocol = new (class extends Protocol<Request, Notification, Result> {
            protected assertCapabilityForMethod(): void {}
            protected assertNotificationCapability(): void {}
            protected assertRequestHandlerCapability(): void {}
            protected assertTaskCapability(): void {}
            protected assertTaskHandlerCapability(): void {}
        })();
    });

    afterEach(() => {
        vi.useRealTimers();
    });

    test('should clear pending timeouts in _onclose to prevent spurious cancellation after reconnect', async () => {
        await protocol.connect(transport);

        // Start a request with a long timeout
        const request = { method: 'example', params: {} };
        const mockSchema = z.object({ result: z.string() });

        const requestPromise = protocol
            .request(request, mockSchema, {
                timeout: 60000
            })
            .catch(() => {
                /* expected ConnectionClosed error */
            });

        // Verify the request was sent
        expect(sendSpy).toHaveBeenCalled();

        // Spy on clearTimeout to verify it gets called during close
        const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');

        // Close the transport - this should clear timeouts
        await transport.close();

        // Verify clearTimeout was called (at least once for our timeout)
        expect(clearTimeoutSpy).toHaveBeenCalled();

        // Now reconnect with a new transport
        const transport2 = new MockTransport();
        const sendSpy2 = vi.spyOn(transport2, 'send');
        await protocol.connect(transport2);

        // Advance past the original timeout - if not cleared, this would fire the callback
        await vi.advanceTimersByTimeAsync(60000);

        // Verify no spurious cancellation notification was sent to the new transport
        const cancellationCalls = sendSpy2.mock.calls.filter(call => {
            const msg = call[0] as Record<string, unknown>;
            return msg.method === 'notifications/cancelled';
        });
        expect(cancellationCalls).toHaveLength(0);

        await transport2.close();
        await requestPromise;
        clearTimeoutSpy.mockRestore();
    });

    test('should not let stale .finally() delete a new connections abort controller after reconnect', async () => {
        await protocol.connect(transport);

        const TestRequestSchema = z.object({
            method: z.literal('test/longRunning'),
            params: z.optional(z.record(z.unknown()))
        });

        // Set up a handler with a deferred resolution we control
        let resolveHandler!: () => void;
        const handlerStarted = new Promise<void>(resolve => {
            protocol.setRequestHandler(TestRequestSchema, async () => {
                resolve(); // signal that handler has started
                // Wait for explicit resolution
                await new Promise<void>(r => {
                    resolveHandler = r;
                });
                return { _meta: {} } as Result;
            });
        });

        // Simulate an incoming request with id=1 on the first connection
        const requestId = 1;
        transport.onmessage!({
            jsonrpc: '2.0',
            id: requestId,
            method: 'test/longRunning',
            params: {}
        });

        // Wait for handler to start
        await handlerStarted;

        // Close the connection (aborts the controller and clears the map)
        await transport.close();

        // Reconnect with a new transport
        const transport2 = new MockTransport();
        await protocol.connect(transport2);

        // Set up a new handler for the second connection that we can verify cancellation on
        let wasAborted = false;
        let resolveHandler2!: () => void;
        const handler2Started = new Promise<void>(resolve => {
            protocol.setRequestHandler(TestRequestSchema, async (_request, extra) => {
                resolve();
                await new Promise<void>(r => {
                    resolveHandler2 = r;
                });
                wasAborted = extra.signal.aborted;
                return { _meta: {} } as Result;
            });
        });

        // Simulate same request id=1 on the new connection
        transport2.onmessage!({
            jsonrpc: '2.0',
            id: requestId,
            method: 'test/longRunning',
            params: {}
        });

        await handler2Started;

        // Resolve the OLD handler so its .finally() runs
        resolveHandler();
        // Flush microtasks so .finally() executes
        await vi.advanceTimersByTimeAsync(0);

        // Send cancellation for request id=1 on the new connection.
        // If the old .finally() incorrectly deleted the new controller, this won't work.
        transport2.onmessage!({
            jsonrpc: '2.0',
            method: 'notifications/cancelled',
            params: {
                requestId: requestId,
                reason: 'test cancel'
            }
        });

        // Resolve handler2 so it can check the abort signal
        resolveHandler2();
        await vi.advanceTimersByTimeAsync(0);

        expect(wasAborted).toBe(true);

        await transport2.close();
    });
});