Skip to content

Commit fe4579d

Browse files
fix(mcp): finalize preview responses in caller
1 parent 56876a9 commit fe4579d

2 files changed

Lines changed: 155 additions & 33 deletions

File tree

‎apps/server/src/mcp/PreviewAutomationBroker.test.ts‎

Lines changed: 111 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
PreviewAutomationMalformedResponseError,
88
PreviewAutomationNoAvailableHostError,
99
PreviewAutomationTargetNotEditableError,
10+
PreviewAutomationTimeoutError,
1011
PreviewTabId,
1112
ProviderInstanceId,
1213
ThreadId,
@@ -19,6 +20,7 @@ import * as Deferred from "effect/Deferred";
1920
import * as Fiber from "effect/Fiber";
2021
import * as Result from "effect/Result";
2122
import * as Stream from "effect/Stream";
23+
import * as TestClock from "effect/testing/TestClock";
2224

2325
import * as PreviewAutomationBroker from "./PreviewAutomationBroker.ts";
2426

@@ -224,7 +226,115 @@ it.effect("blocks a closed in-flight tab without discarding another selected tab
224226
),
225227
);
226228

227-
it.effect("commits a response tab before delivery so a following close cannot be undone", () =>
229+
it.effect("allows an older response for another tab after closing the selected tab", () =>
230+
Effect.scoped(
231+
Effect.gen(function* () {
232+
const broker = yield* makeBroker;
233+
const selectedTabId = PreviewTabId.make("tab-selected-before-close");
234+
const otherTabId = PreviewTabId.make("tab-other-in-flight");
235+
const connected = yield* Deferred.make<void>();
236+
const otherRequestRouted = yield* Deferred.make<void>();
237+
const releaseOtherResponse = yield* Deferred.make<void>();
238+
const routedRequests: RoutedRequest[] = [];
239+
const events = yield* broker.connect(makeHost());
240+
yield* Stream.runForEach(events, (event) => {
241+
if (event.type === "connected") return Deferred.succeed(connected, undefined);
242+
const request = { ...event.request, connectionId: event.connectionId };
243+
routedRequests.push(request);
244+
return Effect.gen(function* () {
245+
if (request.tabId === otherTabId) {
246+
yield* Deferred.succeed(otherRequestRouted, undefined);
247+
yield* Deferred.await(releaseOtherResponse);
248+
}
249+
yield* broker.respond({
250+
clientId: "client-1",
251+
connectionId: request.connectionId,
252+
requestId: request.requestId,
253+
ok: true,
254+
result:
255+
request.operation === "open"
256+
? { available: true, tabId: selectedTabId }
257+
: request.tabId === otherTabId
258+
? { available: true, tabId: otherTabId }
259+
: { url: "http://localhost:3200" },
260+
});
261+
}).pipe(Effect.forkScoped, Effect.asVoid);
262+
}).pipe(Effect.forkScoped);
263+
yield* Deferred.await(connected);
264+
265+
yield* broker.invoke({ scope, operation: "open", input: {} });
266+
const other = yield* broker
267+
.invoke({ scope, operation: "status", input: {}, tabId: otherTabId })
268+
.pipe(Effect.forkScoped);
269+
yield* Deferred.await(otherRequestRouted);
270+
yield* broker.forgetClosedTab(scope, selectedTabId);
271+
yield* Deferred.succeed(releaseOtherResponse, undefined);
272+
yield* Fiber.join(other);
273+
yield* broker.invoke({ scope, operation: "snapshot", input: {} });
274+
275+
expect(routedRequests.at(-1)?.tabId).toBe(otherTabId);
276+
}),
277+
),
278+
);
279+
280+
it.effect("does not install a host response after the caller has timed out", () =>
281+
Effect.scoped(
282+
Effect.gen(function* () {
283+
const broker = yield* makeBroker;
284+
const lateTabId = PreviewTabId.make("tab-after-timeout");
285+
const connected = yield* Deferred.make<void>();
286+
const requestRouted = yield* Deferred.make<void>();
287+
const releaseResponse = yield* Deferred.make<void>();
288+
const responseHandled = yield* Deferred.make<void>();
289+
const routedRequests: RoutedRequest[] = [];
290+
const events = yield* broker.connect(makeHost());
291+
yield* Stream.runForEach(events, (event) => {
292+
if (event.type === "connected") return Deferred.succeed(connected, undefined);
293+
const request = { ...event.request, connectionId: event.connectionId };
294+
routedRequests.push(request);
295+
return Effect.gen(function* () {
296+
if (request.operation === "status") {
297+
yield* Deferred.succeed(requestRouted, undefined);
298+
yield* Deferred.await(releaseResponse);
299+
}
300+
yield* broker.respond({
301+
clientId: "client-1",
302+
connectionId: request.connectionId,
303+
requestId: request.requestId,
304+
ok: true,
305+
result:
306+
request.operation === "status"
307+
? { available: true, tabId: lateTabId }
308+
: { url: "http://localhost:3200" },
309+
});
310+
if (request.operation === "status") {
311+
yield* Deferred.succeed(responseHandled, undefined);
312+
}
313+
}).pipe(Effect.forkScoped, Effect.asVoid);
314+
}).pipe(Effect.forkScoped);
315+
yield* Deferred.await(connected);
316+
317+
const timedOut = yield* broker
318+
.invoke<{ readonly available: boolean }>({
319+
scope,
320+
operation: "status",
321+
input: {},
322+
timeoutMs: 1_000,
323+
})
324+
.pipe(Effect.flip, Effect.forkScoped);
325+
yield* Deferred.await(requestRouted);
326+
yield* TestClock.adjust("1 second");
327+
expect(yield* Fiber.join(timedOut)).toBeInstanceOf(PreviewAutomationTimeoutError);
328+
yield* Deferred.succeed(releaseResponse, undefined);
329+
yield* Deferred.await(responseHandled);
330+
yield* broker.invoke({ scope, operation: "snapshot", input: {} });
331+
332+
expect(routedRequests.at(-1)?.tabId).toBeUndefined();
333+
}),
334+
),
335+
);
336+
337+
it.effect("keeps a response pending until close can reject its returned tab", () =>
228338
Effect.scoped(
229339
Effect.gen(function* () {
230340
const broker = yield* makeBroker;

‎apps/server/src/mcp/PreviewAutomationBroker.ts‎

Lines changed: 44 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -430,21 +430,42 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() {
430430
const respond: PreviewAutomationBroker["Service"]["respond"] = Effect.fn(
431431
"PreviewAutomationBroker.respond",
432432
)(function* (response) {
433-
const pending = yield* SynchronizedRef.modify(state, (current) => {
434-
const entry = current.pending.get(response.requestId);
435-
if (
436-
!entry ||
437-
entry.context.clientId !== response.clientId ||
438-
entry.context.connectionId !== response.connectionId
439-
) {
440-
return [undefined, current] as const;
441-
}
442-
const nextPending = new Map(current.pending);
443-
nextPending.delete(response.requestId);
444-
let assignments: ReadonlyMap<string, HostAssignment> = current.assignments;
445-
if (response.ok) {
433+
const pending = yield* SynchronizedRef.get(state).pipe(
434+
Effect.map((current) => {
435+
const entry = current.pending.get(response.requestId);
436+
if (
437+
!entry ||
438+
entry.context.clientId !== response.clientId ||
439+
entry.context.connectionId !== response.connectionId
440+
) {
441+
return undefined;
442+
}
443+
return entry;
444+
}),
445+
);
446+
if (!pending) return;
447+
if (response.ok) {
448+
yield* Deferred.succeed(pending.deferred, response.result);
449+
} else {
450+
yield* Deferred.fail(
451+
pending.deferred,
452+
response.error
453+
? classifyResponseError(pending.context, response.error)
454+
: new PreviewAutomationMalformedResponseError(pending.context),
455+
);
456+
}
457+
});
458+
459+
const commitSuccessfulResponse = Effect.fn("PreviewAutomationBroker.commitSuccessfulResponse")(
460+
function* (requestId: string, result: unknown) {
461+
yield* SynchronizedRef.update(state, (current) => {
462+
const entry = current.pending.get(requestId);
463+
if (entry === undefined) return current;
464+
const pending = new Map(current.pending);
465+
pending.delete(requestId);
466+
let assignments: ReadonlyMap<string, HostAssignment> = current.assignments;
446467
const assignment = assignments.get(entry.assignmentKey);
447-
const responseTabId = readResultTabId(response.result);
468+
const responseTabId = readResultTabId(result);
448469
const resultTabId =
449470
responseTabId === undefined && entry.tabIdExplicit ? entry.context.tabId : responseTabId;
450471
if (
@@ -473,22 +494,11 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() {
473494
}
474495
assignments = nextAssignments;
475496
}
476-
}
477-
assignments = pruneClosedTabBarriers(assignments, nextPending, entry.assignmentKey);
478-
return [entry, { ...current, assignments, pending: nextPending }] as const;
479-
});
480-
if (!pending) return;
481-
if (response.ok) {
482-
yield* Deferred.succeed(pending.deferred, response.result);
483-
} else {
484-
yield* Deferred.fail(
485-
pending.deferred,
486-
response.error
487-
? classifyResponseError(pending.context, response.error)
488-
: new PreviewAutomationMalformedResponseError(pending.context),
489-
);
490-
}
491-
});
497+
assignments = pruneClosedTabBarriers(assignments, pending, entry.assignmentKey);
498+
return { ...current, assignments, pending };
499+
});
500+
},
501+
);
492502

493503
const forgetClosedTab: PreviewAutomationBroker["Service"]["forgetClosedTab"] = Effect.fn(
494504
"PreviewAutomationBroker.forgetClosedTab",
@@ -513,7 +523,6 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() {
513523
const { tabId: _tabId, ...withoutTabId } = assignment;
514524
assignments.set(assignmentKey, {
515525
...withoutTabId,
516-
tabSequence: current.requestSequence,
517526
...(closedTabSequences.size === 0 ? {} : { closedTabSequences }),
518527
});
519528
} else {
@@ -673,7 +682,10 @@ export const make = Effect.gen(function* PreviewAutomationBrokerMake() {
673682
onSome: (value) => Effect.succeed(value as A),
674683
});
675684
});
676-
const result = yield* awaitResponse().pipe(Effect.ensuring(removePending));
685+
const result = yield* awaitResponse().pipe(
686+
Effect.tap((response) => commitSuccessfulResponse(requestId, response)),
687+
Effect.ensuring(removePending),
688+
);
677689
return result;
678690
});
679691

0 commit comments

Comments
 (0)