modelcontextprotocol/typescript-sdk
Watch
1
0
Fork
You've already forked typescript-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 1

Fix cross-session task leakage (#1486)

Browse source
This commit is contained in:
Den Delimarsky 2026-02-06 07:47:17 -08:00 • committed by GitHub
commit 65bbceab77
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 146 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

56
packages/core/src/experimental/tasks/stores/inMemory.ts
View file

@ -11,6 +11,7 @@ interface StoredTask {
task: Task;
request: Request;
requestId: RequestId;
sessionId?: string;
result?: Result;
}
@ -30,7 +31,7 @@ export class InMemoryTaskStore implements TaskStore {
return crypto.randomUUID().replaceAll('-', '');
}
async createTask(taskParams: CreateTaskOptions, requestId: RequestId, request: Request, _sessionId?: string): Promise<Task> {
async createTask(taskParams: CreateTaskOptions, requestId: RequestId, request: Request, sessionId?: string): Promise<Task> {
// Generate a unique task ID
const taskId = this.generateTaskId();
@ -55,7 +56,8 @@ export class InMemoryTaskStore implements TaskStore {
this.tasks.set(taskId, {
task,
request,
requestId
requestId,
sessionId
});
// Schedule cleanup if ttl is specified
@ -72,13 +74,30 @@ export class InMemoryTaskStore implements TaskStore {
return task;
}
async getTask(taskId: string, _sessionId?: string): Promise<Task | null> {
/**
* Retrieves a stored task, enforcing session ownership when a sessionId is provided.
* Returns undefined if the task does not exist or belongs to a different session.
*/
private getStoredTask(taskId: string, sessionId?: string): StoredTask | undefined {
const stored = this.tasks.get(taskId);
if (!stored) {
return undefined;
}
// Enforce session isolation: if a sessionId is provided and the task
// was created with a sessionId, they must match.
if (sessionId !== undefined && stored.sessionId !== undefined && stored.sessionId !== sessionId) {
return undefined;
}
return stored;
}
async getTask(taskId: string, sessionId?: string): Promise<Task | null> {
const stored = this.getStoredTask(taskId, sessionId);
return stored ? { ...stored.task } : null;
}
async storeTaskResult(taskId: string, status: 'completed' | 'failed', result: Result, _sessionId?: string): Promise<void> {
const stored = this.tasks.get(taskId);
async storeTaskResult(taskId: string, status: 'completed' | 'failed', result: Result, sessionId?: string): Promise<void> {
const stored = this.getStoredTask(taskId, sessionId);
if (!stored) {
throw new Error(`Task with ID ${taskId} not found`);
}
@ -110,8 +129,8 @@ export class InMemoryTaskStore implements TaskStore {
}
}
async getTaskResult(taskId: string, _sessionId?: string): Promise<Result> {
const stored = this.tasks.get(taskId);
async getTaskResult(taskId: string, sessionId?: string): Promise<Result> {
const stored = this.getStoredTask(taskId, sessionId);
if (!stored) {
throw new Error(`Task with ID ${taskId} not found`);
}
@ -123,8 +142,8 @@ export class InMemoryTaskStore implements TaskStore {
return stored.result;
}
async updateTaskStatus(taskId: string, status: Task['status'], statusMessage?: string, _sessionId?: string): Promise<void> {
const stored = this.tasks.get(taskId);
async updateTaskStatus(taskId: string, status: Task['status'], statusMessage?: string, sessionId?: string): Promise<void> {
const stored = this.getStoredTask(taskId, sessionId);
if (!stored) {
throw new Error(`Task with ID ${taskId} not found`);
}
@ -159,13 +178,22 @@ export class InMemoryTaskStore implements TaskStore {
}
}
async listTasks(cursor?: string, _sessionId?: string): Promise<{ tasks: Task[]; nextCursor?: string }> {
async listTasks(cursor?: string, sessionId?: string): Promise<{ tasks: Task[]; nextCursor?: string }> {
const PAGE_SIZE = 10;
const allTaskIds = [...this.tasks.keys()];
// Filter tasks by session ownership before pagination
const filteredTaskIds = [...this.tasks.entries()]
.filter(([, stored]) => {
if (sessionId === undefined || stored.sessionId === undefined) {
return true;
}
return stored.sessionId === sessionId;
})
.map(([taskId]) => taskId);
let startIndex = 0;
if (cursor) {
const cursorIndex = allTaskIds.indexOf(cursor);
const cursorIndex = filteredTaskIds.indexOf(cursor);
if (cursorIndex === -1) {
// Invalid cursor - throw error
throw new Error(`Invalid cursor: ${cursor}`);
@ -174,13 +202,13 @@ export class InMemoryTaskStore implements TaskStore {
}
}
const pageTaskIds = allTaskIds.slice(startIndex, startIndex + PAGE_SIZE);
const pageTaskIds = filteredTaskIds.slice(startIndex, startIndex + PAGE_SIZE);
const tasks = pageTaskIds.map(taskId => {
const stored = this.tasks.get(taskId)!;
return { ...stored.task };
});
const nextCursor = startIndex + PAGE_SIZE < allTaskIds.length ? pageTaskIds.at(-1) : undefined;
const nextCursor = startIndex + PAGE_SIZE < filteredTaskIds.length ? pageTaskIds.at(-1) : undefined;
return { tasks, nextCursor };
}

99
packages/core/test/experimental/inMemory.test.ts
View file

@ -647,6 +647,105 @@ describe('InMemoryTaskStore', () => {
});
});
describe('session isolation', () => {
const baseRequest: Request = { method: 'tools/call', params: { name: 'demo' } };
it('should not allow session-b to list tasks created by session-a', async () => {
await store.createTask({}, 1, baseRequest, 'session-a');
await store.createTask({}, 2, baseRequest, 'session-a');
const result = await store.listTasks(undefined, 'session-b');
expect(result.tasks).toHaveLength(0);
});
it('should not allow session-b to read a task created by session-a', async () => {
const task = await store.createTask({}, 1, baseRequest, 'session-a');
const result = await store.getTask(task.taskId, 'session-b');
expect(result).toBeNull();
});
it('should not allow session-b to update a task created by session-a', async () => {
const task = await store.createTask({}, 1, baseRequest, 'session-a');
await expect(store.updateTaskStatus(task.taskId, 'cancelled', undefined, 'session-b')).rejects.toThrow('not found');
});
it('should not allow session-b to store a result on session-a task', async () => {
const task = await store.createTask({}, 1, baseRequest, 'session-a');
await expect(store.storeTaskResult(task.taskId, 'completed', { content: [] }, 'session-b')).rejects.toThrow('not found');
});
it('should not allow session-b to get the result of session-a task', async () => {
const task = await store.createTask({}, 1, baseRequest, 'session-a');
await store.storeTaskResult(task.taskId, 'completed', { content: [{ type: 'text', text: 'secret' }] }, 'session-a');
await expect(store.getTaskResult(task.taskId, 'session-b')).rejects.toThrow('not found');
});
it('should allow the owning session to access its own tasks', async () => {
const task = await store.createTask({}, 1, baseRequest, 'session-a');
const retrieved = await store.getTask(task.taskId, 'session-a');
expect(retrieved).toBeDefined();
expect(retrieved?.taskId).toBe(task.taskId);
});
it('should list only tasks belonging to the requesting session', async () => {
await store.createTask({}, 1, baseRequest, 'session-a');
await store.createTask({}, 2, baseRequest, 'session-b');
await store.createTask({}, 3, baseRequest, 'session-a');
const resultA = await store.listTasks(undefined, 'session-a');
expect(resultA.tasks).toHaveLength(2);
const resultB = await store.listTasks(undefined, 'session-b');
expect(resultB.tasks).toHaveLength(1);
});
it('should allow access when no sessionId is provided (backward compatibility)', async () => {
const task = await store.createTask({}, 1, baseRequest, 'session-a');
// No sessionId on read = no filtering
const retrieved = await store.getTask(task.taskId);
expect(retrieved).toBeDefined();
});
it('should allow access when task was created without sessionId', async () => {
const task = await store.createTask({}, 1, baseRequest);
// Any sessionId on read should still see the task
const retrieved = await store.getTask(task.taskId, 'session-b');
expect(retrieved).toBeDefined();
});
it('should paginate correctly within a session', async () => {
// Create 15 tasks for session-a, 5 for session-b
for (let i = 1; i <= 15; i++) {
await store.createTask({}, i, baseRequest, 'session-a');
}
for (let i = 16; i <= 20; i++) {
await store.createTask({}, i, baseRequest, 'session-b');
}
// First page for session-a should have 10
const page1 = await store.listTasks(undefined, 'session-a');
expect(page1.tasks).toHaveLength(10);
expect(page1.nextCursor).toBeDefined();
// Second page for session-a should have 5
const page2 = await store.listTasks(page1.nextCursor, 'session-a');
expect(page2.tasks).toHaveLength(5);
expect(page2.nextCursor).toBeUndefined();
// session-b should only see its 5
const resultB = await store.listTasks(undefined, 'session-b');
expect(resultB.tasks).toHaveLength(5);
expect(resultB.nextCursor).toBeUndefined();
});
});
describe('cleanup', () => {
it('should clear all timers and tasks', async () => {
await store.createTask({ ttl: 1000 }, 1, {