Skip to content
Source
Sign inSign up

McpSessionRegistry.ts

247 lines9.3 KB

1import { ProviderInstanceId, ThreadId } from "@t3tools/contracts";
2import * as Clock from "effect/Clock";
3import * as Context from "effect/Context";
4import * as Digest from "effect/crypto/Digest";
5import * as SecureRandom from "effect/crypto/SecureRandom";
6import * as Effect from "effect/Effect";
7import * as Layer from "effect/Layer";
8import * as SynchronizedRef from "effect/SynchronizedRef";
9import { HttpServer } from "effect/http";
10import * as NetAddress from "effect/net/NetAddress";
11​
12import * as ServerEnvironment from "../environment/ServerEnvironment.ts";
13import * as McpInvocationContext from "./McpInvocationContext.ts";
14import * as McpProviderSession from "@t3tools/provider-core/server/mcpSession";
15​
16export interface McpCredentialRequest {
17 readonly threadId: ThreadId;
18 readonly providerInstanceId: ProviderInstanceId;
19 /** False omits the preview capability wherever the token is honored. Defaults to full access. */
20 readonly browserToolsAvailable?: boolean;
21 readonly capabilities?: ReadonlySet<McpInvocationContext.McpCapability>;
22 readonly subagent?: boolean;
24​
25export interface McpIssuedCredential {
26 readonly config: McpProviderSession.McpProviderSessionConfig;
28​
29export interface McpSessionRegistryShape {
30 readonly issue: (request: McpCredentialRequest) => Effect.Effect<McpIssuedCredential>;
31 readonly resolve: (
32 rawToken: string,
33 ) => Effect.Effect<McpInvocationContext.McpThreadInvocationScope | undefined>;
34 readonly touch: (threadId: ThreadId) => Effect.Effect<void>;
35 readonly revokeProviderSession: (providerSessionId: string) => Effect.Effect<void>;
36 readonly revokeThread: (threadId: ThreadId) => Effect.Effect<void>;
37 readonly revokeAll: Effect.Effect<void>;
39​
40export class McpSessionRegistry extends Context.Service<
41 McpSessionRegistry,
42 McpSessionRegistryShape
43>()("t3/mcp/McpSessionRegistry") {}
44​
45interface AliveMark {
46 readonly wallMs: number;
47 readonly monotonicMs: number;
49​
50interface CredentialRecord {
51 readonly tokenHash: string;
52 readonly scope: McpInvocationContext.McpThreadInvocationScope;
53 readonly lastAlive: AliveMark;
54 /** Set when a request finds the token idle; only `touch` clears it, whatever the clocks do. */
55 readonly expired: boolean;
57​
58interface RegistryState {
59 readonly records: ReadonlyMap<string, CredentialRecord>;
61​
62export interface McpSessionRegistryOptions {
63 readonly livenessWindowMs?: number;
64 /** Drives both the wall and monotonic readings. */
65 readonly now?: () => number;
67​
68/**
69 * `/mcp` is outside environment auth, so idle tokens stop authenticating.
70 * Keep expired records until explicit revocation: a trusted provider turn can
71 * refresh the token its subprocess still holds; MCP requests cannot.
72 */
73const DEFAULT_LIVENESS_WINDOW_MS = 24 * 60 * 60 * 1_000;
74​
75const bytesToHex = (bytes: Uint8Array): string =>
76 Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join("");
77​
78const tokenFromBytes = (bytes: Uint8Array): string => Buffer.from(bytes).toString("base64url");
79​
80// The larger elapsed reading wins: stepping the wall clock back cannot extend a token,
81// and wall time that passes while the host sleeps still counts.
82const idleMs = (since: AliveMark, now: AliveMark): number =>
83 Math.max(now.wallMs - since.wallMs, now.monotonicMs - since.monotonicMs);
84​
85const getHttpMcpEndpointHost = (address: NetAddress.IpAddress): string =>
86 NetAddress.isUnspecified(address)
87 ? "127.0.0.1"
88 : NetAddress.formatUrlHostString(NetAddress.formatIp(address));
89​
90const makeWithOptions = Effect.fn("McpSessionRegistry.make")(function* (
91 options: McpSessionRegistryOptions = {},
92) {
93 const random = yield* SecureRandom.SecureRandom;
94 const digest = yield* Digest.Digest;
95 const environment = yield* ServerEnvironment.ServerEnvironment;
96 const environmentId = yield* environment.getEnvironmentId;
97 const httpServer = yield* HttpServer.HttpServer;
98 const state = yield* SynchronizedRef.make<RegistryState>({ records: new Map() });
99 const { now } = options;
100 const currentAliveMark: Effect.Effect<AliveMark> = now
101 ? Effect.sync(() => {
102 const timestamp = now();
103 return { wallMs: timestamp, monotonicMs: timestamp };
104 })
105 : Effect.all({
106 wallMs: Clock.currentTimeMillis,
107 monotonicMs: Clock.monotonicTimeNanos.pipe(
108 Effect.map((nanos) => Number(nanos / 1_000_000n)),
109 ),
110 });
111 const livenessWindowMs = options.livenessWindowMs ?? DEFAULT_LIVENESS_WINDOW_MS;
112 const endpoint = NetAddress.isInetAddress(httpServer.address)
113 ? `http://${getHttpMcpEndpointHost(httpServer.address.address)}:${httpServer.address.port}/mcp`
114 : "http://127.0.0.1/mcp";
115​
116 const hashToken = (token: string) =>
117 digest
118 .digest(Digest.sha256, new TextEncoder().encode(token))
119 .pipe(Effect.map(bytesToHex), Effect.orDie);
120​
121 const issue: McpSessionRegistryShape["issue"] = Effect.fn("McpSessionRegistry.issue")(
122 function* (request) {
123 const issuedAlive = yield* currentAliveMark;
124 const providerSessionId = yield* random.uuidV4();
125 const rawToken = yield* random.randomBytes(32).pipe(Effect.map(tokenFromBytes), Effect.orDie);
126 const tokenHash = yield* hashToken(rawToken);
127 const browserToolsAvailable = request.browserToolsAvailable ?? true;
128 const scope: McpInvocationContext.McpThreadInvocationScope = {
129 environmentId,
130 requestNamespace: providerSessionId,
131 thread: {
132 threadId: ThreadId.make(request.threadId),
133 providerSessionId,
134 providerInstanceId: ProviderInstanceId.make(request.providerInstanceId),
135 subagent: request.subagent === true,
136 },
137 client: undefined,
138 capabilities: new Set<McpInvocationContext.McpCapability>([
139 "orchestration",
140 "worktree",
141 "pull-requests",
142 ...(request.capabilities ?? (browserToolsAvailable ? (["preview"] as const) : [])),
143 ]),
144 issuedAt: issuedAlive.wallMs,
145 };
146 yield* SynchronizedRef.update(state, ({ records }) => {
147 const next = new Map(records);
148 next.set(tokenHash, { tokenHash, scope, lastAlive: issuedAlive, expired: false });
149 return { records: next };
150 });
151 return {
152 config: {
153 environmentId,
154 threadId: scope.thread.threadId,
155 providerSessionId,
156 providerInstanceId: scope.thread.providerInstanceId,
157 endpoint,
158 authorizationHeader: `Bearer ${rawToken}`,
159 browserToolsAvailable: scope.capabilities.has("preview"),
160 capabilities: scope.capabilities,
161 },
162 };
163 },
164 );
165​
166 const resolve: McpSessionRegistryShape["resolve"] = Effect.fn("McpSessionRegistry.resolve")(
167 function* (rawToken) {
168 if (rawToken.length === 0) return undefined;
169 const tokenHash = yield* hashToken(rawToken);
170 const alive = yield* currentAliveMark;
171 return yield* SynchronizedRef.modify(state, ({ records }) => {
172 const record = records.get(tokenHash);
173 if (!record || record.expired) return [undefined, { records }] as const;
174 if (idleMs(record.lastAlive, alive) > livenessWindowMs) {
175 const next = new Map(records);
176 next.set(tokenHash, { ...record, expired: true });
177 return [undefined, { records: next }] as const;
178 }
179 const next = new Map(records);
180 next.set(tokenHash, { ...record, lastAlive: alive });
181 return [record.scope, { records: next }] as const;
182 });
183 },
184 );
185​
186 const touch: McpSessionRegistryShape["touch"] = Effect.fn("McpSessionRegistry.touch")(
187 function* (threadId) {
188 const alive = yield* currentAliveMark;
189 yield* SynchronizedRef.update(state, ({ records }) => {
190 const next = new Map(records);
191 for (const [tokenHash, record] of records) {
192 if (record.scope.thread.threadId === threadId) {
193 next.set(tokenHash, { ...record, lastAlive: alive, expired: false });
194 }
195 }
196 return { records: next };
197 });
198 },
199 );
200​
201 const revokeWhere = (predicate: (record: CredentialRecord) => boolean) =>
202 SynchronizedRef.update(state, ({ records }) => ({
203 records: new Map(Array.from(records).filter(([, record]) => !predicate(record))),
204 }));
205​
206 return McpSessionRegistry.of({
207 issue,
208 resolve,
209 touch,
210 revokeProviderSession: Effect.fn("McpSessionRegistry.revokeProviderSession")(
211 function* (providerSessionId) {
212 yield* revokeWhere((record) => record.scope.thread.providerSessionId === providerSessionId);
213 },
214 ),
215 revokeThread: Effect.fn("McpSessionRegistry.revokeThread")(function* (threadId) {
216 yield* revokeWhere((record) => record.scope.thread.threadId === threadId);
217 }),
218 revokeAll: SynchronizedRef.set(state, { records: new Map() }),
219 });
220});
221​
222let activeMcpSessionRegistry: McpSessionRegistryShape | undefined;
223​
224const make = Effect.acquireRelease(
225 makeWithOptions().pipe(
226 Effect.tap((registry) =>
227 Effect.sync(() => {
228 activeMcpSessionRegistry = registry;
229 }),
230 ),
231 ),
232 (registry) =>
233 Effect.sync(() => {
234 if (activeMcpSessionRegistry === registry) {
235 activeMcpSessionRegistry = undefined;
236 }
237 }),
238);
239​
240export const layer = Layer.effect(McpSessionRegistry, make);
241​
242export const touchActiveMcpThread = (threadId: ThreadId): Effect.Effect<void> =>
243 activeMcpSessionRegistry ? activeMcpSessionRegistry.touch(threadId) : Effect.void;
244​
245export const __testing = {
246 make: makeWithOptions,
247};