Skip to content

Commit c604721

Browse files
localdenlocalden
committed
Fix cross-session task leakage
1 parent 0da5d3f commit c604721

2 files changed

Lines changed: 147 additions & 14 deletions

File tree

‎packages/core/src/experimental/tasks/stores/inMemory.ts‎

Lines changed: 42 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ interface StoredTask {
1111
task: Task;
1212
request: Request;
1313
requestId: RequestId;
14+
sessionId?: string;
1415
result?: Result;
1516
}
1617

@@ -30,7 +31,7 @@ export class InMemoryTaskStore implements TaskStore {
3031
return crypto.randomUUID().replaceAll('-', '');
3132
}
3233

33-
async createTask(taskParams: CreateTaskOptions, requestId: RequestId, request: Request, _sessionId?: string): Promise<Task> {
34+
async createTask(taskParams: CreateTaskOptions, requestId: RequestId, request: Request, sessionId?: string): Promise<Task> {
3435
// Generate a unique task ID
3536
const taskId = this.generateTaskId();
3637

@@ -55,7 +56,8 @@ export class InMemoryTaskStore implements TaskStore {
5556
this.tasks.set(taskId, {
5657
task,
5758
request,
58-
requestId
59+
requestId,
60+
sessionId
5961
});
6062

6163
// Schedule cleanup if ttl is specified
@@ -72,13 +74,30 @@ export class InMemoryTaskStore implements TaskStore {
7274
return task;
7375
}
7476

75-
async getTask(taskId: string, _sessionId?: string): Promise<Task | null> {
77+
/**
78+
* Retrieves a stored task, enforcing session ownership when a sessionId is provided.
79+
* Returns undefined if the task does not exist or belongs to a different session.
80+
*/
81+
private getStoredTask(taskId: string, sessionId?: string): StoredTask | undefined {
7682
const stored = this.tasks.get(taskId);
83+
if (!stored) {
84+
return undefined;
85+
}
86+
// Enforce session isolation: if a sessionId is provided and the task
87+
// was created with a sessionId, they must match.
88+
if (sessionId !== undefined && stored.sessionId !== undefined && stored.sessionId !== sessionId) {
89+
return undefined;
90+
}
91+
return stored;
92+
}
93+
94+
async getTask(taskId: string, sessionId?: string): Promise<Task | null> {
95+
const stored = this.getStoredTask(taskId, sessionId);
7796
return stored ? { ...stored.task } : null;
7897
}
7998

80-
async storeTaskResult(taskId: string, status: 'completed' | 'failed', result: Result, _sessionId?: string): Promise<void> {
81-
const stored = this.tasks.get(taskId);
99+
async storeTaskResult(taskId: string, status: 'completed' | 'failed', result: Result, sessionId?: string): Promise<void> {
100+
const stored = this.getStoredTask(taskId, sessionId);
82101
if (!stored) {
83102
throw new Error(`Task with ID ${taskId} not found`);
84103
}
@@ -110,8 +129,8 @@ export class InMemoryTaskStore implements TaskStore {
110129
}
111130
}
112131

113-
async getTaskResult(taskId: string, _sessionId?: string): Promise<Result> {
114-
const stored = this.tasks.get(taskId);
132+
async getTaskResult(taskId: string, sessionId?: string): Promise<Result> {
133+
const stored = this.getStoredTask(taskId, sessionId);
115134
if (!stored) {
116135
throw new Error(`Task with ID ${taskId} not found`);
117136
}
@@ -123,8 +142,8 @@ export class InMemoryTaskStore implements TaskStore {
123142
return stored.result;
124143
}
125144

126-
async updateTaskStatus(taskId: string, status: Task['status'], statusMessage?: string, _sessionId?: string): Promise<void> {
127-
const stored = this.tasks.get(taskId);
145+
async updateTaskStatus(taskId: string, status: Task['status'], statusMessage?: string, sessionId?: string): Promise<void> {
146+
const stored = this.getStoredTask(taskId, sessionId);
128147
if (!stored) {
129148
throw new Error(`Task with ID ${taskId} not found`);
130149
}
@@ -159,13 +178,22 @@ export class InMemoryTaskStore implements TaskStore {
159178
}
160179
}
161180

162-
async listTasks(cursor?: string, _sessionId?: string): Promise<{ tasks: Task[]; nextCursor?: string }> {
181+
async listTasks(cursor?: string, sessionId?: string): Promise<{ tasks: Task[]; nextCursor?: string }> {
163182
const PAGE_SIZE = 10;
164-
const allTaskIds = [...this.tasks.keys()];
183+
184+
// Filter tasks by session ownership before pagination
185+
const filteredTaskIds = [...this.tasks.entries()]
186+
.filter(([, stored]) => {
187+
if (sessionId === undefined || stored.sessionId === undefined) {
188+
return true;
189+
}
190+
return stored.sessionId === sessionId;
191+
})
192+
.map(([taskId]) => taskId);
165193

166194
let startIndex = 0;
167195
if (cursor) {
168-
const cursorIndex = allTaskIds.indexOf(cursor);
196+
const cursorIndex = filteredTaskIds.indexOf(cursor);
169197
if (cursorIndex === -1) {
170198
// Invalid cursor - throw error
171199
throw new Error(`Invalid cursor: ${cursor}`);
@@ -174,13 +202,13 @@ export class InMemoryTaskStore implements TaskStore {
174202
}
175203
}
176204

177-
const pageTaskIds = allTaskIds.slice(startIndex, startIndex + PAGE_SIZE);
205+
const pageTaskIds = filteredTaskIds.slice(startIndex, startIndex + PAGE_SIZE);
178206
const tasks = pageTaskIds.map(taskId => {
179207
const stored = this.tasks.get(taskId)!;
180208
return { ...stored.task };
181209
});
182210

183-
const nextCursor = startIndex + PAGE_SIZE < allTaskIds.length ? pageTaskIds.at(-1) : undefined;
211+
const nextCursor = startIndex + PAGE_SIZE < filteredTaskIds.length ? pageTaskIds.at(-1) : undefined;
184212

185213
return { tasks, nextCursor };
186214
}

‎packages/core/test/experimental/inMemory.test.ts‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -647,6 +647,111 @@ describe('InMemoryTaskStore', () => {
647647
});
648648
});
649649

650+
describe('session isolation', () => {
651+
const baseRequest: Request = { method: 'tools/call', params: { name: 'demo' } };
652+
653+
it('should not allow session-b to list tasks created by session-a', async () => {
654+
await store.createTask({}, 1, baseRequest, 'session-a');
655+
await store.createTask({}, 2, baseRequest, 'session-a');
656+
657+
const result = await store.listTasks(undefined, 'session-b');
658+
expect(result.tasks).toHaveLength(0);
659+
});
660+
661+
it('should not allow session-b to read a task created by session-a', async () => {
662+
const task = await store.createTask({}, 1, baseRequest, 'session-a');
663+
664+
const result = await store.getTask(task.taskId, 'session-b');
665+
expect(result).toBeNull();
666+
});
667+
668+
it('should not allow session-b to update a task created by session-a', async () => {
669+
const task = await store.createTask({}, 1, baseRequest, 'session-a');
670+
671+
await expect(
672+
store.updateTaskStatus(task.taskId, 'cancelled', undefined, 'session-b')
673+
).rejects.toThrow('not found');
674+
});
675+
676+
it('should not allow session-b to store a result on session-a task', async () => {
677+
const task = await store.createTask({}, 1, baseRequest, 'session-a');
678+
679+
await expect(
680+
store.storeTaskResult(task.taskId, 'completed', { content: [] }, 'session-b')
681+
).rejects.toThrow('not found');
682+
});
683+
684+
it('should not allow session-b to get the result of session-a task', async () => {
685+
const task = await store.createTask({}, 1, baseRequest, 'session-a');
686+
await store.storeTaskResult(task.taskId, 'completed', { content: [{ type: 'text', text: 'secret' }] }, 'session-a');
687+
688+
await expect(
689+
store.getTaskResult(task.taskId, 'session-b')
690+
).rejects.toThrow('not found');
691+
});
692+
693+
it('should allow the owning session to access its own tasks', async () => {
694+
const task = await store.createTask({}, 1, baseRequest, 'session-a');
695+
696+
const retrieved = await store.getTask(task.taskId, 'session-a');
697+
expect(retrieved).toBeDefined();
698+
expect(retrieved?.taskId).toBe(task.taskId);
699+
});
700+
701+
it('should list only tasks belonging to the requesting session', async () => {
702+
await store.createTask({}, 1, baseRequest, 'session-a');
703+
await store.createTask({}, 2, baseRequest, 'session-b');
704+
await store.createTask({}, 3, baseRequest, 'session-a');
705+
706+
const resultA = await store.listTasks(undefined, 'session-a');
707+
expect(resultA.tasks).toHaveLength(2);
708+
709+
const resultB = await store.listTasks(undefined, 'session-b');
710+
expect(resultB.tasks).toHaveLength(1);
711+
});
712+
713+
it('should allow access when no sessionId is provided (backward compatibility)', async () => {
714+
const task = await store.createTask({}, 1, baseRequest, 'session-a');
715+
716+
// No sessionId on read = no filtering
717+
const retrieved = await store.getTask(task.taskId);
718+
expect(retrieved).toBeDefined();
719+
});
720+
721+
it('should allow access when task was created without sessionId', async () => {
722+
const task = await store.createTask({}, 1, baseRequest);
723+
724+
// Any sessionId on read should still see the task
725+
const retrieved = await store.getTask(task.taskId, 'session-b');
726+
expect(retrieved).toBeDefined();
727+
});
728+
729+
it('should paginate correctly within a session', async () => {
730+
// Create 15 tasks for session-a, 5 for session-b
731+
for (let i = 1; i <= 15; i++) {
732+
await store.createTask({}, i, baseRequest, 'session-a');
733+
}
734+
for (let i = 16; i <= 20; i++) {
735+
await store.createTask({}, i, baseRequest, 'session-b');
736+
}
737+
738+
// First page for session-a should have 10
739+
const page1 = await store.listTasks(undefined, 'session-a');
740+
expect(page1.tasks).toHaveLength(10);
741+
expect(page1.nextCursor).toBeDefined();
742+
743+
// Second page for session-a should have 5
744+
const page2 = await store.listTasks(page1.nextCursor, 'session-a');
745+
expect(page2.tasks).toHaveLength(5);
746+
expect(page2.nextCursor).toBeUndefined();
747+
748+
// session-b should only see its 5
749+
const resultB = await store.listTasks(undefined, 'session-b');
750+
expect(resultB.tasks).toHaveLength(5);
751+
expect(resultB.nextCursor).toBeUndefined();
752+
});
753+
});
754+
650755
describe('cleanup', () => {
651756
it('should clear all timers and tasks', async () => {
652757
await store.createTask({ ttl: 1000 }, 1, {

0 commit comments

Comments
 (0)