History

Harden referenced thread reads

Rodrigo Brechard committed 4 months ago

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

Unified · Split

4 files
  1. apps/server/src/mcp/toolkits/threadReference/handlers.test.ts +139 −24
  2. apps/server/src/mcp/toolkits/threadReference/handlers.ts +28 −1
  3. packages/shared/src/composerInlineTokens.test.ts +13 −0
  4. packages/shared/src/composerInlineTokens.ts +9 −4
apps/server/src/mcp/toolkits/threadReference/handlers.test.ts+139−24
@@ -1,5 +1,6 @@
1import { expect, it } from "@effect/vitest";1import { expect, it } from "@effect/vitest";
2import {2import {
3 EnvironmentId,
3 MessageId,4 MessageId,
4 ProjectId,5 ProjectId,
5 ProviderInstanceId,6 ProviderInstanceId,
@@ -14,7 +15,11 @@ import { McpSchema, McpServer } from "effect/unstable/ai"; · 8 unmodified lines
14 15
15import * as McpInvocationContext from "../../McpInvocationContext.ts";16import * as McpInvocationContext from "../../McpInvocationContext.ts";
16import { ProjectionSnapshotQuery } from "../../../orchestration/Services/ProjectionSnapshotQuery.ts";17import { ProjectionSnapshotQuery } from "../../../orchestration/Services/ProjectionSnapshotQuery.ts";
17import { buildThreadReferencePage, ThreadReferenceToolkitHandlersLive } from "./handlers.ts";18import {
19 buildThreadReferencePage,
20 hasUserThreadReference,
21 ThreadReferenceToolkitHandlersLive,
22} from "./handlers.ts";
18import { ThreadReferenceToolkit } from "./tools.ts";23import { ThreadReferenceToolkit } from "./tools.ts";
19 24
20const TestLayer = McpServer.toolkit(ThreadReferenceToolkit).pipe(25const TestLayer = McpServer.toolkit(ThreadReferenceToolkit).pipe(
@@ -22,9 +27,13 @@ const TestLayer = McpServer.toolkit(ThreadReferenceToolkit).pipe( · 1 unmodified line
22 Layer.provideMerge(McpServer.McpServer.layer),27 Layer.provideMerge(McpServer.McpServer.layer),
23);28);
24 29
25const makeMessage = (id: string, text: string): OrchestrationMessage => ({30const makeMessage = (
31 id: string,
32 text: string,
33 role: OrchestrationMessage["role"] = "user",
34): OrchestrationMessage => ({
26 id: MessageId.make(id),35 id: MessageId.make(id),
27 role: "user",36 role,
28 text,37 text,
29 turnId: null,38 turnId: null,
30 streaming: false,39 streaming: false,
@@ -32,12 +41,53 @@ const makeMessage = (id: string, text: string): OrchestrationMessage => ({ · 1 unmodified line
32 updatedAt: "2026-01-01T00:00:00.000Z",41 updatedAt: "2026-01-01T00:00:00.000Z",
33});42});
34 43
44const currentThreadId = ThreadId.make("thread-current");
45const environmentId = EnvironmentId.make("environment-1");
35const thread = {46const thread = {
36 id: ThreadId.make("thread-referenced"),47 id: ThreadId.make("thread-referenced"),
37 projectId: ProjectId.make("project-1"),48 projectId: ProjectId.make("project-1"),
38 title: "Referenced work",49 title: "Referenced work",
39 messages: [makeMessage("message-1", "abcdef"), makeMessage("message-2", "ghijkl")],50 messages: [makeMessage("message-1", "abcdef"), makeMessage("message-2", "ghijkl")],
40} as unknown as OrchestrationThread;51} as unknown as OrchestrationThread;
52const 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;
62const 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};
71const 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};
80const 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
42it("paginates within a message without losing transcript text", () => {92it("paginates within a message without losing transcript text", () => {
43 const first = buildThreadReferencePage(thread, {93 const first = buildThreadReferencePage(thread, {
@@ -67,6 +117,67 @@ it("paginates within a message without losing transcript text", () => { · 23 unmodified lines
67 expect(longSecond.nextCursor).toBeNull();117 expect(longSecond.nextCursor).toBeNull();
68});118});
69 119
120it("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
146it("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
70it.effect("reads a referenced thread through the MCP toolkit", () =>181it.effect("reads a referenced thread through the MCP toolkit", () =>
71 Effect.gen(function* () {182 Effect.gen(function* () {
72 const server = yield* McpServer.McpServer;183 const server = yield* McpServer.McpServer;
@@ -81,27 +192,31 @@ it.effect("reads a referenced thread through the MCP toolkit", () => · 8 unmodified lines
81 totalMessages: 2,192 totalMessages: 2,
82 });193 });
83 }).pipe(194 }).pipe(
84 Effect.provideService(McpInvocationContext.McpInvocationContext, {195 Effect.provideService(McpInvocationContext.McpInvocationContext, invocation),
85 environmentId: "environment-1" as never,196 Effect.provideService(McpSchema.McpServerClient, mcpServerClient),
86 threadId: ThreadId.make("thread-current"),197 Effect.provideService(ProjectionSnapshotQuery, projectionQueryFor(currentThread)),
87 providerSessionId: "provider-session-1",198 Effect.provide(TestLayer),
88 providerInstanceId: ProviderInstanceId.make("codex"),199 ),
89 capabilities: new Set(["thread-reference"] as const),200);
90 issuedAt: 1,201
91 expiresAt: Number.MAX_SAFE_INTEGER,202it.effect("rejects a thread that was not referenced by the invoking user", () =>
92 }),203 Effect.gen(function* () {
93 Effect.provideService(McpSchema.McpServerClient, {204 const server = yield* McpServer.McpServer;
94 clientId: 1,205 const result = yield* server.callTool({
95 initializePayload: {206 name: "thread_read",
96 protocolVersion: "2025-03-26",207 arguments: { threadId: thread.id },
97 capabilities: {},208 });
98 clientInfo: { name: "test", version: "1" },209 expect(result.isError).toBe(true);
99 },210 }).pipe(
100 getClient: Effect.die("unused"),211 Effect.provideService(McpInvocationContext.McpInvocationContext, invocation),
101 }),212 Effect.provideService(McpSchema.McpServerClient, mcpServerClient),
102 Effect.provideService(ProjectionSnapshotQuery, {213 Effect.provideService(
103 getThreadDetailById: () => Effect.succeed(Option.some(thread)),214 ProjectionSnapshotQuery,
104 } as never),215 projectionQueryFor({
216 ...currentThread,
217 messages: [makeMessage("message-without-reference", "No thread reference here.")],
218 }),
219 ),
105 Effect.provide(TestLayer),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 type ThreadReferenceReadResult,6 type ThreadReferenceReadResult,
7 type ThreadReferenceTranscriptMessage,7 type ThreadReferenceTranscriptMessage,
8 ThreadReferenceUnavailableError,8 ThreadReferenceUnavailableError,
9 type EnvironmentId,
9 type OrchestrationThread,10 type OrchestrationThread,
11 type ThreadId,
10} from "@t3tools/contracts";12} from "@t3tools/contracts";
13import { collectComposerThreadReferences } from "@t3tools/shared/composerInlineTokens";
11import * as Effect from "effect/Effect";14import * as Effect from "effect/Effect";
12import * as Option from "effect/Option";15import * as Option from "effect/Option";
13 16
@@ -35,7 +38,7 @@ function decodeCursor( · 21 unmodified lines
35 !Number.isSafeInteger(textOffset) ||38 !Number.isSafeInteger(textOffset) ||
36 messageIndex < 0 ||39 messageIndex < 0 ||
37 textOffset < 0 ||40 textOffset < 0 ||
38 (!validEndCursor && (!message || textOffset >= message.text.length))41 (!validEndCursor && (!message || textOffset > message.text.length))
39 ) {42 ) {
40 return null;43 return null;
41 }44 }
@@ -96,6 +99,21 @@ export function buildThreadReferencePage( · 54 unmodified lines
96 };99 };
97}100}
98 101
102export 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
99const threadRead = Effect.fn("ThreadReferenceToolkit.threadRead")(function* (117const threadRead = Effect.fn("ThreadReferenceToolkit.threadRead")(function* (
100 input: ThreadReferenceReadInput,118 input: ThreadReferenceReadInput,
101) {119) {
@@ -104,6 +122,15 @@ const threadRead = Effect.fn("ThreadReferenceToolkit.threadRead")(function* ( · 2 unmodified lines
104 return yield* new ThreadReferenceUnavailableError({ threadId: input.threadId });122 return yield* new ThreadReferenceUnavailableError({ threadId: input.threadId });
105 }123 }
106 const snapshotQuery = yield* ProjectionSnapshotQuery;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 const threadOption = yield* snapshotQuery134 const threadOption = yield* snapshotQuery
108 .getThreadDetailById(input.threadId)135 .getThreadDetailById(input.threadId)
109 .pipe(Effect.mapError(() => new ThreadReferenceUnavailableError({ threadId: input.threadId })));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 it("ignores malformed thread links", () => {96 it("ignores malformed thread links", () => {
97 expect(collectComposerInlineTokens("Read [chat](t3-thread:///missing) first")).toEqual([]);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 readonly preserveTrailingFrom?: ReadonlyArray<ComposerInlineToken>;30 readonly preserveTrailingFrom?: ReadonlyArray<ComposerInlineToken>;
31}31}
32 32
33type ComposerThreadInlineToken = Extract<ComposerInlineToken, { type: "thread" }>;
34
33const SKILL_TOKEN_REGEX = /(^|\s)\$([a-zA-Z][a-zA-Z0-9:_-]*)(?=\s)/g;35const SKILL_TOKEN_REGEX = /(^|\s)\$([a-zA-Z][a-zA-Z0-9:_-]*)(?=\s)/g;
34const MENTION_TOKEN_REGEX = /(^|\s)@(?:"((?:\\.|[^"\\])*)"|([^\s@"]+))(?=\s)/g;36const MENTION_TOKEN_REGEX = /(^|\s)@(?:"((?:\\.|[^"\\])*)"|([^\s@"]+))(?=\s)/g;
35const FILE_LINK_TOKEN_REGEX = /(^|\s)\[((?:\\.|[^\]\\])*)\]\(([^)\s]+)\)(?=\s)/g;37const FILE_LINK_TOKEN_REGEX = /(^|\s)\[((?:\\.|[^\]\\])*)\]\(([^)\s]+)\)(?=\s)/g;
36const THREAD_LINK_TOKEN_REGEX = /(^|\s)\[((?:\\.|[^\]\\])*)\]\((t3-thread:\/\/\/[^)\s]+)\)(?=\s)/g;38const THREAD_LINK_TOKEN_REGEX =
39 /(^|\s)\[((?:\\.|[^\]\\])*)\]\((t3-thread:\/\/\/[^)\s]+)\)(?=\s|$)/g;
37const URI_SCHEME_REGEX = /^[A-Za-z][A-Za-z0-9+.-]*:/;40const URI_SCHEME_REGEX = /^[A-Za-z][A-Za-z0-9+.-]*:/;
38const WINDOWS_DRIVE_PATH_REGEX = /^[A-Za-z]:[\\/]/;41const WINDOWS_DRIVE_PATH_REGEX = /^[A-Za-z]:[\\/]/;
39 42
40function collectThreadTokens(text: string): ComposerInlineToken[] {43export function collectComposerThreadReferences(
41 const matches: ComposerInlineToken[] = [];44 text: string,
45): ReadonlyArray<ComposerThreadInlineToken> {
46 const matches: ComposerThreadInlineToken[] = [];
42 47
43 for (const match of text.matchAll(THREAD_LINK_TOKEN_REGEX)) {48 for (const match of text.matchAll(THREAD_LINK_TOKEN_REGEX)) {
44 const fullMatch = match[0];49 const fullMatch = match[0];
@@ -122,7 +127,7 @@ export function collectComposerInlineTokens( · 77 unmodified lines
122 text: string,127 text: string,
123 options: CollectComposerInlineTokensOptions = {},128 options: CollectComposerInlineTokensOptions = {},
124): ReadonlyArray<ComposerInlineToken> {129): ReadonlyArray<ComposerInlineToken> {
125 const matches = [...collectMentionTokens(text), ...collectThreadTokens(text)];130 const matches = [...collectMentionTokens(text), ...collectComposerThreadReferences(text)];
126 131
127 for (const match of text.matchAll(SKILL_TOKEN_REGEX)) {132 for (const match of text.matchAll(SKILL_TOKEN_REGEX)) {
128 const fullMatch = match[0];133 const fullMatch = match[0];

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t