Harden referenced thread reads

Rodrigo Brechard committed 4 months ago

Commit dba8c741 parent e51be694 files · +189 −29Browse files

Unified · Split

4 files

apps/server/src/mcp/toolkits/threadReference/handlers.test.ts +139 −24
apps/server/src/mcp/toolkits/threadReference/handlers.ts +28 −1
packages/shared/src/composerInlineTokens.test.ts +13 −0
packages/shared/src/composerInlineTokens.ts +9 −4

apps/server/src/mcp/toolkits/threadReference/handlers.test.ts+139−24

@@ -1,5 +1,6 @@
1	1	import { expect, it } from "@effect/vitest";
2	2	import {
	3	+ EnvironmentId,
3	4	MessageId,
4	5	ProjectId,
5	6	ProviderInstanceId,
@@ -14,7 +15,11 @@ import { McpSchema, McpServer } from "effect/unstable/ai"; · 8 unmodified lines
14	15
15	16	import * as McpInvocationContext from "../../McpInvocationContext.ts";
16	17	import { ProjectionSnapshotQuery } from "../../../orchestration/Services/ProjectionSnapshotQuery.ts";
17		−import { buildThreadReferencePage, ThreadReferenceToolkitHandlersLive } from "./handlers.ts";
	18	+import {
	19	+ buildThreadReferencePage,
	20	+ hasUserThreadReference,
	21	+ ThreadReferenceToolkitHandlersLive,
	22	+} from "./handlers.ts";
18	23	import { ThreadReferenceToolkit } from "./tools.ts";
19	24
20	25	const TestLayer = McpServer.toolkit(ThreadReferenceToolkit).pipe(
@@ -22,9 +27,13 @@ const TestLayer = McpServer.toolkit(ThreadReferenceToolkit).pipe( · 1 unmodified line
22	27	Layer.provideMerge(McpServer.McpServer.layer),
23	28	);
24	29
25		−const makeMessage = (id: string, text: string): OrchestrationMessage => ({
	30	+const makeMessage = (
	31	+ id: string,
	32	+ text: string,
	33	+ role: OrchestrationMessage["role"] = "user",
	34	+): OrchestrationMessage => ({
26	35	id: MessageId.make(id),
27		− role: "user",
	36	+ role,
28	37	text,
29	38	turnId: null,
30	39	streaming: false,
@@ -32,12 +41,53 @@ const makeMessage = (id: string, text: string): OrchestrationMessage => ({ · 1 unmodified line
32	41	updatedAt: "2026-01-01T00:00:00.000Z",
33	42	});
34	43
	44	+const currentThreadId = ThreadId.make("thread-current");
	45	+const environmentId = EnvironmentId.make("environment-1");
35	46	const thread = {
36	47	id: ThreadId.make("thread-referenced"),
37	48	projectId: ProjectId.make("project-1"),
38	49	title: "Referenced work",
39	50	messages: [makeMessage("message-1", "abcdef"), makeMessage("message-2", "ghijkl")],
40	51	} as unknown as OrchestrationThread;
	52	+const currentThread = {
	53	+ ...thread,
	54	+ id: currentThreadId,
	55	+ messages: [
	56	+ makeMessage(
	57	+ "message-reference",
	58	+ `[Referenced work](t3-thread:///${environmentId}/${thread.id})`,
	59	+ ),
	60	+ ],
	61	+} as OrchestrationThread;
	62	+const invocation = {
	63	+ environmentId,
	64	+ threadId: currentThreadId,
	65	+ providerSessionId: "provider-session-1",
	66	+ providerInstanceId: ProviderInstanceId.make("codex"),
	67	+ capabilities: new Set(["thread-reference"] as const),
	68	+ issuedAt: 1,
	69	+ expiresAt: Number.MAX_SAFE_INTEGER,
	70	+};
	71	+const mcpServerClient = {
	72	+ clientId: 1,
	73	+ initializePayload: {
	74	+ protocolVersion: "2025-03-26",
	75	+ capabilities: {},
	76	+ clientInfo: { name: "test", version: "1" },
	77	+ },
	78	+ getClient: Effect.die("unused"),
	79	+};
	80	+const projectionQueryFor = (sourceThread: OrchestrationThread) =>
	81	+ ({
	82	+ getThreadDetailById: (threadId: ThreadId) =>
	83	+ Effect.succeed(
	84	+ threadId === sourceThread.id
	85	+ ? Option.some(sourceThread)
	86	+ : threadId === thread.id
	87	+ ? Option.some(thread)
	88	+ : Option.none(),
	89	+ ),
	90	+ }) as never;
41	91
42	92	it("paginates within a message without losing transcript text", () => {
43	93	const first = buildThreadReferencePage(thread, {
@@ -67,6 +117,67 @@ it("paginates within a message without losing transcript text", () => { · 23 unmodified lines
67	117	expect(longSecond.nextCursor).toBeNull();
68	118	});
69	119
	120	+it("continues pagination across an empty message at a page boundary", () => {
	121	+ const emptyMessageThread = {
	122	+ ...thread,
	123	+ messages: [
	124	+ makeMessage("message-full", "a".repeat(1_000)),
	125	+ makeMessage("message-empty", ""),
	126	+ makeMessage("message-tail", "tail"),
	127	+ ],
	128	+ };
	129	+ const first = buildThreadReferencePage(emptyMessageThread, {
	130	+ threadId: thread.id,
	131	+ maxChars: 1_000,
	132	+ });
	133	+ if ("_tag" in first) throw new Error("unexpected cursor error");
	134	+ expect(first.nextCursor).toBe("1:0");
	135	+
	136	+ const second = buildThreadReferencePage(emptyMessageThread, {
	137	+ threadId: thread.id,
	138	+ cursor: first.nextCursor!,
	139	+ maxChars: 1_000,
	140	+ });
	141	+ if ("_tag" in second) throw new Error("unexpected cursor error");
	142	+ expect(second.messages.map(({ text }) => text)).toEqual(["", "tail"]);
	143	+ expect(second.nextCursor).toBeNull();
	144	+});
	145	+
	146	+it("only authorizes user-supplied references from the invoking environment", () => {
	147	+ expect(hasUserThreadReference(currentThread, environmentId, thread.id)).toBe(true);
	148	+ expect(
	149	+ hasUserThreadReference(
	150	+ {
	151	+ ...currentThread,
	152	+ messages: [
	153	+ makeMessage(
	154	+ "message-assistant-reference",
	155	+ `[Referenced work](t3-thread:///${environmentId}/${thread.id})`,
	156	+ "assistant",
	157	+ ),
	158	+ ],
	159	+ },
	160	+ environmentId,
	161	+ thread.id,
	162	+ ),
	163	+ ).toBe(false);
	164	+ expect(
	165	+ hasUserThreadReference(
	166	+ {
	167	+ ...currentThread,
	168	+ messages: [
	169	+ makeMessage(
	170	+ "message-other-environment",
	171	+ `[Referenced work](t3-thread:///environment-2/${thread.id})`,
	172	+ ),
	173	+ ],
	174	+ },
	175	+ environmentId,
	176	+ thread.id,
	177	+ ),
	178	+ ).toBe(false);
	179	+});
	180	+
70	181	it.effect("reads a referenced thread through the MCP toolkit", () =>
71	182	Effect.gen(function* () {
72	183	const server = yield* McpServer.McpServer;
@@ -81,27 +192,31 @@ it.effect("reads a referenced thread through the MCP toolkit", () => · 8 unmodified lines
81	192	totalMessages: 2,
82	193	});
83	194	}).pipe(
84		− Effect.provideService(McpInvocationContext.McpInvocationContext, {
85		− environmentId: "environment-1" as never,
86		− threadId: ThreadId.make("thread-current"),
87		− providerSessionId: "provider-session-1",
88		− providerInstanceId: ProviderInstanceId.make("codex"),
89		− capabilities: new Set(["thread-reference"] as const),
90		− issuedAt: 1,
91		− expiresAt: Number.MAX_SAFE_INTEGER,
92		− }),
93		− Effect.provideService(McpSchema.McpServerClient, {
94		− clientId: 1,
95		− initializePayload: {
96		− protocolVersion: "2025-03-26",
97		− capabilities: {},
98		− clientInfo: { name: "test", version: "1" },
99		− },
100		− getClient: Effect.die("unused"),
101		− }),
102		− Effect.provideService(ProjectionSnapshotQuery, {
103		− getThreadDetailById: () => Effect.succeed(Option.some(thread)),
104		− } as never),
	195	+ Effect.provideService(McpInvocationContext.McpInvocationContext, invocation),
	196	+ Effect.provideService(McpSchema.McpServerClient, mcpServerClient),
	197	+ Effect.provideService(ProjectionSnapshotQuery, projectionQueryFor(currentThread)),
	198	+ Effect.provide(TestLayer),
	199	+ ),
	200	+);
	201	+
	202	+it.effect("rejects a thread that was not referenced by the invoking user", () =>
	203	+ Effect.gen(function* () {
	204	+ const server = yield* McpServer.McpServer;
	205	+ const result = yield* server.callTool({
	206	+ name: "thread_read",
	207	+ arguments: { threadId: thread.id },
	208	+ });
	209	+ expect(result.isError).toBe(true);
	210	+ }).pipe(
	211	+ Effect.provideService(McpInvocationContext.McpInvocationContext, invocation),
	212	+ Effect.provideService(McpSchema.McpServerClient, mcpServerClient),
	213	+ Effect.provideService(
	214	+ ProjectionSnapshotQuery,
	215	+ projectionQueryFor({
	216	+ ...currentThread,
	217	+ messages: [makeMessage("message-without-reference", "No thread reference here.")],
	218	+ }),
	219	+ ),
105	220	Effect.provide(TestLayer),
106	221	),
107	222	);

apps/server/src/mcp/toolkits/threadReference/handlers.ts+28−1

@@ -6,8 +6,11 @@ import { · 5 unmodified lines
6	6	type ThreadReferenceReadResult,
7	7	type ThreadReferenceTranscriptMessage,
8	8	ThreadReferenceUnavailableError,
	9	+ type EnvironmentId,
9	10	type OrchestrationThread,
	11	+ type ThreadId,
10	12	} from "@t3tools/contracts";
	13	+import { collectComposerThreadReferences } from "@t3tools/shared/composerInlineTokens";
11	14	import * as Effect from "effect/Effect";
12	15	import * as Option from "effect/Option";
13	16
@@ -35,7 +38,7 @@ function decodeCursor( · 21 unmodified lines
35	38	!Number.isSafeInteger(textOffset) \|\|
36	39	messageIndex < 0 \|\|
37	40	textOffset < 0 \|\|
38		− (!validEndCursor && (!message \|\| textOffset >= message.text.length))
	41	+ (!validEndCursor && (!message \|\| textOffset > message.text.length))
39	42	) {
40	43	return null;
41	44	}
@@ -96,6 +99,21 @@ export function buildThreadReferencePage( · 54 unmodified lines
96	99	};
97	100	}
98	101
	102	+export function hasUserThreadReference(
	103	+ sourceThread: OrchestrationThread,
	104	+ environmentId: EnvironmentId,
	105	+ targetThreadId: ThreadId,
	106	+): boolean {
	107	+ return sourceThread.messages.some(
	108	+ (message) =>
	109	+ message.role === "user" &&
	110	+ collectComposerThreadReferences(message.text).some(
	111	+ (reference) =>
	112	+ reference.environmentId === environmentId && reference.threadId === targetThreadId,
	113	+ ),
	114	+ );
	115	+}
	116	+
99	117	const threadRead = Effect.fn("ThreadReferenceToolkit.threadRead")(function* (
100	118	input: ThreadReferenceReadInput,
101	119	) {
@@ -104,6 +122,15 @@ const threadRead = Effect.fn("ThreadReferenceToolkit.threadRead")(function* ( · 2 unmodified lines
104	122	return yield* new ThreadReferenceUnavailableError({ threadId: input.threadId });
105	123	}
106	124	const snapshotQuery = yield* ProjectionSnapshotQuery;
	125	+ const sourceThreadOption = yield* snapshotQuery
	126	+ .getThreadDetailById(invocation.threadId)
	127	+ .pipe(Effect.mapError(() => new ThreadReferenceUnavailableError({ threadId: input.threadId })));
	128	+ if (
	129	+ Option.isNone(sourceThreadOption) \|\|
	130	+ !hasUserThreadReference(sourceThreadOption.value, invocation.environmentId, input.threadId)
	131	+ ) {
	132	+ return yield* new ThreadReferenceUnavailableError({ threadId: input.threadId });
	133	+ }
107	134	const threadOption = yield* snapshotQuery
108	135	.getThreadDetailById(input.threadId)
109	136	.pipe(Effect.mapError(() => new ThreadReferenceUnavailableError({ threadId: input.threadId })));

packages/shared/src/composerInlineTokens.test.ts+13−0

@@ -96,4 +96,17 @@ describe("collectComposerInlineTokens", () => { · 95 unmodified lines
96	96	it("ignores malformed thread links", () => {
97	97	expect(collectComposerInlineTokens("Read [chat](t3-thread:///missing) first")).toEqual([]);
98	98	});
	99	+
	100	+ it("collects a complete thread reference at the end of the prompt", () => {
	101	+ expect(
	102	+ collectComposerInlineTokens("Read [chat](t3-thread:///environment-1/thread-1)"),
	103	+ ).toMatchObject([
	104	+ {
	105	+ type: "thread",
	106	+ environmentId: "environment-1",
	107	+ threadId: "thread-1",
	108	+ title: "chat",
	109	+ },
	110	+ ]);
	111	+ });
99	112	});

packages/shared/src/composerInlineTokens.ts+9−4

@@ -30,15 +30,20 @@ export interface CollectComposerInlineTokensOptions { · 29 unmodified lines
30	30	readonly preserveTrailingFrom?: ReadonlyArray<ComposerInlineToken>;
31	31	}
32	32
	33	+type ComposerThreadInlineToken = Extract<ComposerInlineToken, { type: "thread" }>;
	34	+
33	35	const SKILL_TOKEN_REGEX = /(^\|\s)\$([a-zA-Z][a-zA-Z0-9:_-]*)(?=\s)/g;
34	36	const MENTION_TOKEN_REGEX = /(^\|\s)@(?:"((?:\\.\|[^"\\])*)"\|([^\s@"]+))(?=\s)/g;
35	37	const FILE_LINK_TOKEN_REGEX = /(^\|\s)\[((?:\\.\|[^\]\\])*)\]$([^)\s]+)$(?=\s)/g;
36		−const THREAD_LINK_TOKEN_REGEX = /(^\|\s)\[((?:\\.\|[^\]\\])*)\]$(t3-thread:\/\/\/[^)\s]+)$(?=\s)/g;
	38	+const THREAD_LINK_TOKEN_REGEX =
	39	+ /(^\|\s)\[((?:\\.\|[^\]\\])*)\]$(t3-thread:\/\/\/[^)\s]+)$(?=\s\|$)/g;
37	40	const URI_SCHEME_REGEX = /^[A-Za-z][A-Za-z0-9+.-]*:/;
38	41	const WINDOWS_DRIVE_PATH_REGEX = /^[A-Za-z]:[\\/]/;
39	42
40		−function collectThreadTokens(text: string): ComposerInlineToken[] {
41		− const matches: ComposerInlineToken[] = [];
	43	+export function collectComposerThreadReferences(
	44	+ text: string,
	45	+): ReadonlyArray<ComposerThreadInlineToken> {
	46	+ const matches: ComposerThreadInlineToken[] = [];
42	47
43	48	for (const match of text.matchAll(THREAD_LINK_TOKEN_REGEX)) {
44	49	const fullMatch = match[0];
@@ -122,7 +127,7 @@ export function collectComposerInlineTokens( · 77 unmodified lines
122	127	text: string,
123	128	options: CollectComposerInlineTokensOptions = {},
124	129	): ReadonlyArray<ComposerInlineToken> {
125		− const matches = [...collectMentionTokens(text), ...collectThreadTokens(text)];
	130	+ const matches = [...collectMentionTokens(text), ...collectComposerThreadReferences(text)];
126	131
127	132	for (const match of text.matchAll(SKILL_TOKEN_REGEX)) {
128	133	const fullMatch = match[0];

Harden referenced thread reads

Keyboard shortcuts

General

Issues and pull requests

Issues, pull requests and Inbox

Repository code