1import * as NodeServices from "@effect/platform-node/NodeServices";
2import { Connection, type DuplexTransport } from "@muse-code/sdk";
3import {
4 ProviderDriverKind,
5 ProviderInstanceId,
6 ProviderReplayEntry,
7 type ProviderReplayEntry as ProviderReplayEntryType,
8 type ProviderReplayTranscript,
9} from "@t3tools/contracts";
10import { MuseSettings } from "@t3tools/provider-muse/settings";
11import * as Effect from "effect/Effect";
12import * as FileSystem from "effect/FileSystem";
13import * as Layer from "effect/Layer";
14import * as Schema from "effect/Schema";
15
16import * as TestProviderHost from "@t3tools/provider-testing/TestProviderHost";
17import {
18 museInitializeParams,
19 museServeArgs,
20 museWorkspaceRoot,
21 type MuseSdkHost,
22 type MuseSdkHostOptions,
23} from "@t3tools/provider-muse/testing";
24import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator";
25import * as ProviderAdapterRegistry from "../ProviderAdapterRegistry.ts";
26import * as ProviderContinuationRequests from "@t3tools/provider-core/server/ProviderContinuationRequests";
27import type { OrchestratorV2ProviderReplayHarness } from "../testkit/ProviderReplayHarness.ts";
28import { materializeReplayTranscriptWorkspace } from "@t3tools/provider-testing/replayTranscript";
29import { makeMuseAdapterV2 } from "@t3tools/provider-muse/server";
30
31export const MUSE_PROVIDER_KIND = "muse";
32export const MUSE_MSP_REPLAY_PROTOCOL = "muse.msp-jsonl";
33const MUSE_REPLAY_ANY = "<any>";
34
35const MuseReplayTranscript = Schema.Struct({
36 provider: Schema.Literal(MUSE_PROVIDER_KIND),
37 protocol: Schema.Literal(MUSE_MSP_REPLAY_PROTOCOL),
38 version: Schema.String,
39 scenario: Schema.String,
40 metadata: Schema.Struct({
41 commandIds: Schema.Array(Schema.Array(Schema.String)),
42 }),
43 entries: Schema.Array(ProviderReplayEntry),
44});
45type MuseReplayTranscript = typeof MuseReplayTranscript.Type;
46const decodeMuseReplayTranscript = Schema.decodeUnknownEffect(MuseReplayTranscript);
47
48class MuseReplayTranscriptDecodeError extends Schema.TaggedError<MuseReplayTranscriptDecodeError>()(
49 "MuseReplayTranscriptDecodeError",
50 { scenario: Schema.optional(Schema.String), cause: Schema.Defect() },
51) {
52 override get message(): string {
53 return `Failed to decode Muse MSP replay transcript for scenario ${this.scenario ?? "<unknown>"}.`;
54 }
55}
56
57class MuseReplayMismatchError extends Schema.TaggedError<MuseReplayMismatchError>()(
58 "MuseReplayMismatchError",
59 {
60 scenario: Schema.String,
61 cursor: Schema.Number,
62 host: Schema.Number,
63 expected: Schema.Unknown,
64 actual: Schema.Unknown,
65 },
66) {
67 override get message(): string {
68 return `Muse replay frame mismatch at cursor ${this.cursor} from host ${this.host} in scenario ${this.scenario}.`;
69 }
70}
71
72class MuseReplayIncompleteError extends Schema.TaggedError<MuseReplayIncompleteError>()(
73 "MuseReplayIncompleteError",
74 {
75 scenario: Schema.String,
76 cursor: Schema.Number,
77 remaining: Schema.Number,
78 nextLabel: Schema.optional(Schema.String),
79 },
80) {
81 override get message(): string {
82 return `Muse replay ended with ${this.remaining} unconsumed entries at cursor ${this.cursor} (next ${this.nextLabel ?? "<none>"}) in scenario ${this.scenario}.`;
83 }
84}
85
86function isRecord(value: unknown): value is Record<string, unknown> {
87 return typeof value === "object" && value !== null && !Array.isArray(value);
88}
89
90function replayValueMatches(expected: unknown, actual: unknown): boolean {
91 if (expected === MUSE_REPLAY_ANY) return true;
92 if (Array.isArray(expected)) {
93 return (
94 Array.isArray(actual) &&
95 expected.length === actual.length &&
96 expected.every((entry, index) => replayValueMatches(entry, actual[index]))
97 );
98 }
99 if (isRecord(expected)) {
100 if (!isRecord(actual)) return false;
101 const keys = new Set([...Object.keys(expected), ...Object.keys(actual)]);
102 return [...keys].every((key) => replayValueMatches(expected[key], actual[key]));
103 }
104 return Object.is(expected, actual);
105}
106
107function clientRequestId(frame: unknown): number | undefined {
108 return isRecord(frame) && typeof frame.method === "string" && typeof frame.id === "number"
109 ? frame.id
110 : undefined;
111}
112
113function withRequestId(expected: unknown, actual: unknown): unknown {
114 const actualId = clientRequestId(actual);
115 return isRecord(expected) && clientRequestId(expected) !== undefined && actualId !== undefined
116 ? { ...expected, id: actualId }
117 : expected;
118}
119
120/** A write may not match past one of these: it would be a turn early. */
121function startsTurn(entry: ProviderReplayEntryType): boolean {
122 if (entry.type !== "expect_outbound" || !isRecord(entry.frame)) return false;
123 const method = entry.frame.method ?? entry.frame.type;
124 return method === "turn/start" || method === "session/compact" || method === "host_start";
125}
126
127function entryHost(entry: ProviderReplayEntryType): number {
128 const label = entry.type === "runtime_exit" ? undefined : entry.label;
129 const match = label === undefined ? null : /@h(\d+)$/u.exec(label);
130 return match === null ? 1 : Number(match[1]);
131}
132
133function entryLabel(entry: ProviderReplayEntryType | undefined): string | undefined {
134 if (entry === undefined) return undefined;
135 return entry.type === "runtime_exit" ? "runtime_exit" : (entry.label ?? entry.type);
136}
137
138export function museRecordLabel(
139 frame: unknown,
140 direction: "expect_outbound" | "emit_inbound",
141 requestMethods: ReadonlyMap<unknown, string>,
142 host: number,
143): string | undefined {
144 if (!isRecord(frame)) return undefined;
145 const method = typeof frame.method === "string" ? frame.method : undefined;
146 const base =
147 frame.type === "host_start"
148 ? "host_start"
149 : method !== undefined
150 ? direction === "emit_inbound" && "id" in frame
151 ? `request:${method}`
152 : method
153 : requestMethods.has(frame.id)
154 ? `${direction === "emit_inbound" ? "response" : "reply"}:${requestMethods.get(frame.id)}`
155 : undefined;
156 return base === undefined || host === 1 ? base : `${base}@h${host}`;
157}
158
159class LineQueue implements AsyncIterable<string> {
160 private readonly lines: Array<string> = [];
161 private waiting: ((result: IteratorResult<string>) => void) | undefined;
162 private ended = false;
163
164 push(line: string): void {
165 if (this.ended) return;
166 const waiting = this.waiting;
167 if (waiting) {
168 this.waiting = undefined;
169 waiting({ value: line, done: false });
170 } else this.lines.push(line);
171 }
172
173 end(): void {
174 this.ended = true;
175 this.waiting?.({ value: undefined, done: true });
176 this.waiting = undefined;
177 }
178
179 [Symbol.asyncIterator](): AsyncIterator<string> {
180 return {
181 next: () => {
182 const line = this.lines.shift();
183 if (line !== undefined) return Promise.resolve({ value: line, done: false });
184 if (this.ended) return Promise.resolve({ value: undefined, done: true });
185 return new Promise((resolve) => {
186 this.waiting = resolve;
187 });
188 },
189 };
190 }
191}
192
193interface MuseReplayHost {
194 readonly ordinal: number;
195 readonly incoming: LineQueue;
196 readonly ids: Map<number, number>;
197 readonly minted: Array<string>;
198 readonly exited: Promise<{ readonly code: number | null; readonly signal: null }>;
199 readonly exit: () => void;
200 open: boolean;
201}
202
203/** Adapter fibers can reorder stdin writes; hold matches within the turn until reached, and fail unmatched writes. */
204class MuseReplayController {
205 private cursor = 0;
206 private failure: unknown = null;
207 private readonly consumed = new Set<number>();
208 private readonly held: Array<{ readonly host: MuseReplayHost; readonly actual: unknown }> = [];
209 private readonly hosts: Array<MuseReplayHost> = [];
210 private readonly transcript: MuseReplayTranscript;
211
212 constructor(transcript: MuseReplayTranscript) {
213 this.transcript = transcript;
214 }
215
216 spawn(args: ReadonlyArray<string>): MuseReplayHost {
217 let exit = () => {};
218 const exited = new Promise<{ readonly code: number | null; readonly signal: null }>(
219 (resolve) => {
220 exit = () => resolve({ code: 0, signal: null });
221 },
222 );
223 const host: MuseReplayHost = {
224 ordinal: this.hosts.length + 1,
225 incoming: new LineQueue(),
226 ids: new Map(),
227 minted: [...(this.transcript.metadata.commandIds[this.hosts.length] ?? [])],
228 exited,
229 exit,
230 open: true,
231 };
232 this.hosts.push(host);
233 this.receive(host, { type: "host_start", args: [...args] });
234 return host;
235 }
236
237 mint(host: MuseReplayHost): string {
238 const id = host.minted.shift();
239 if (id === undefined) {
240 throw new Error(`Muse replay host ${host.ordinal} minted more ids than were recorded.`);
241 }
242 return id;
243 }
244
245 receive(host: MuseReplayHost, actual: unknown): void {
246 if (this.failure !== null) return;
247 const entries = this.transcript.entries;
248 let possible = false;
249 for (let index = this.cursor; index < entries.length && !possible; index += 1) {
250 const entry = entries[index]!;
251 const pending = entry.type === "expect_outbound" && !this.consumed.has(index);
252 possible =
253 pending &&
254 entryHost(entry) === host.ordinal &&
255 replayValueMatches(withRequestId(entry.frame, actual), actual);
256 if (!possible && pending && startsTurn(entry)) break;
257 }
258 if (!possible) {
259 this.fail(
260 new MuseReplayMismatchError({
261 scenario: this.transcript.scenario,
262 cursor: this.cursor,
263 host: host.ordinal,
264 expected: entries[this.cursor] ?? null,
265 actual,
266 }),
267 );
268 return;
269 }
270 this.held.push({ host, actual });
271 this.advance();
272 }
273
274 close(host: MuseReplayHost): void {
275 host.open = false;
276 host.incoming.end();
277 host.exit();
278 }
279
280 assertComplete(): void {
281 if (this.failure !== null) throw this.failure;
282 if (this.cursor !== this.transcript.entries.length || this.held.length > 0) {
283 const nextLabel = entryLabel(this.transcript.entries[this.cursor]);
284 throw new MuseReplayIncompleteError({
285 scenario: this.transcript.scenario,
286 cursor: this.cursor,
287 remaining: this.transcript.entries.length - this.cursor,
288 ...(nextLabel === undefined ? {} : { nextLabel }),
289 });
290 }
291 }
292
293 private advance(): void {
294 let progressed = true;
295 while (progressed && this.failure === null) {
296 progressed = this.drainInbound();
297 for (const [heldIndex, write] of this.held.entries()) {
298 if (this.consumePending(write.host, write.actual)) {
299 this.held.splice(heldIndex, 1);
300 progressed = true;
301 break;
302 }
303 }
304 }
305 }
306
307 private consumePending(host: MuseReplayHost, actual: unknown): boolean {
308 const entries = this.transcript.entries;
309 for (let index = this.cursor; index < entries.length; index += 1) {
310 const entry = entries[index]!;
311 if (entry.type !== "expect_outbound") return false;
312 if (this.consumed.has(index) || entryHost(entry) !== host.ordinal) continue;
313 if (!replayValueMatches(withRequestId(entry.frame, actual), actual)) continue;
314 const recordedId = clientRequestId(entry.frame);
315 const actualId = clientRequestId(actual);
316 if (recordedId !== undefined && actualId !== undefined) host.ids.set(recordedId, actualId);
317 this.consumed.add(index);
318 while (this.consumed.has(this.cursor)) this.cursor += 1;
319 return true;
320 }
321 return false;
322 }
323
324 private drainInbound(): boolean {
325 let emitted = false;
326 while (this.failure === null) {
327 const entry = this.transcript.entries[this.cursor];
328 if (entry?.type === "runtime_exit") {
329 this.cursor += 1;
330 emitted = true;
331 const host = this.hosts.at(-1);
332 if (host !== undefined) this.close(host);
333 continue;
334 }
335 if (entry?.type !== "emit_inbound") return emitted;
336 const host = this.hosts[entryHost(entry) - 1];
337 if (host === undefined) return emitted;
338 this.cursor += 1;
339 emitted = true;
340 // Output a host wrote after T3 released it was never observed.
341 if (!host.open) continue;
342 const frame = entry.frame;
343 // Responses carry the client's request id; server requests carry their own.
344 const reboundId =
345 isRecord(frame) && frame.method === undefined && typeof frame.id === "number"
346 ? host.ids.get(frame.id)
347 : undefined;
348 const rebound =
349 isRecord(frame) && reboundId !== undefined ? { ...frame, id: reboundId } : frame;
350 host.incoming.push(`${JSON.stringify(rebound)}\n`);
351 }
352 return emitted;
353 }
354
355 private fail(cause: unknown): void {
356 if (this.failure !== null) return;
357 this.failure = cause;
358 for (const host of this.hosts) host.incoming.end();
359 }
360}
361
362export async function connectMuseTransport(input: {
363 readonly transport: DuplexTransport;
364 readonly readOnly: boolean | undefined;
365 readonly mintCommandId?: () => string;
366 readonly exited: MuseSdkHost["exited"];
367 readonly close: () => Promise<void>;
368}): Promise<MuseSdkHost> {
369 const connection = new Connection(
370 input.transport,
371 input.mintCommandId ? { mintCommandId: input.mintCommandId } : {},
372 );
373 const result = await connection.request("initialize", museInitializeParams(input.readOnly));
374 connection.notify("initialized");
375 const granted = result.grantedCapabilities;
376 return {
377 connection,
378 initializeResult: {
379 grantedCapabilities: Array.isArray(granted)
380 ? granted.filter((capability) => typeof capability === "string")
381 : [],
382 },
383 exited: input.exited,
384 close: input.close,
385 };
386}
387
388function makeReplayCreateHost(controller: MuseReplayController) {
389 return async (options: MuseSdkHostOptions): Promise<MuseSdkHost> => {
390 const host = controller.spawn(museServeArgs(options));
391 return connectMuseTransport({
392 transport: {
393 incoming: host.incoming,
394 write: async (line) => {
395 for (const part of line.split("\n")) {
396 if (part.trim()) controller.receive(host, JSON.parse(part));
397 }
398 },
399 close: () => controller.close(host),
400 },
401 readOnly: options.readOnly,
402 mintCommandId: () => controller.mint(host),
403 exited: host.exited,
404 close: async () => controller.close(host),
405 });
406 };
407}
408
409const DEFAULT_MUSE_SETTINGS = Schema.decodeSync(MuseSettings)({ enabled: true });
410
411export function layer(input: {
412 readonly scenario: string;
413 readonly createHost: NonNullable<Parameters<typeof makeMuseAdapterV2>[0]["createHost"]>;
414 readonly environment?: NodeJS.ProcessEnv;
415}) {
416 return ProviderAdapterRegistry.layerFromAdaptersEffect(
417 Effect.gen(function* () {
418 return [
419 yield* makeMuseAdapterV2({
420 instanceId: ProviderInstanceId.make(MUSE_PROVIDER_KIND),
421 settings: DEFAULT_MUSE_SETTINGS,
422 environment: input.environment ?? {},
423 createHost: input.createHost,
424 continuationRequests: yield* ProviderContinuationRequests.ProviderContinuationRequests,
425 }),
426 ];
427 }),
428 ).pipe(
429 Layer.provide(
430 Layer.mergeAll(
431 TestProviderHost.layer().pipe(Layer.provide(NodeServices.layer)),
432 NodeServices.layer,
433 IdAllocator.layer,
434 ),
435 ),
436 );
437}
438
439/** Uses canonical workspace paths, with turn/start in museWorkspaceRoot form (verbatim on Windows). */
440export const materializeMuseReplayWorkspace = Effect.fn("materializeMuseReplayWorkspace")(
441 function* (transcript: ProviderReplayTranscript, workspace: string) {
442 const canonical = yield* FileSystem.FileSystem.pipe(
443 Effect.flatMap((fs) => fs.realPath(workspace)),
444 Effect.provide(NodeServices.layer),
445 );
446 const plain = materializeReplayTranscriptWorkspace(transcript, canonical);
447 const turnStarts = materializeReplayTranscriptWorkspace(
448 transcript,
449 yield* museWorkspaceRoot(canonical),
450 );
451 return {
452 ...plain,
453 entries: plain.entries.map((entry, index) =>
454 entry.type === "expect_outbound" &&
455 isRecord(entry.frame) &&
456 entry.frame.method === "turn/start"
457 ? turnStarts.entries[index]!
458 : entry,
459 ),
460 };
461 },
462);
463
464export const MuseOrchestratorReplayHarness: OrchestratorV2ProviderReplayHarness<
465 MuseReplayTranscript,
466 MuseReplayTranscriptDecodeError
467> = {
468 driver: ProviderDriverKind.make(MUSE_PROVIDER_KIND),
469 decodeTranscript: (transcript: ProviderReplayTranscript) =>
470 decodeMuseReplayTranscript(transcript).pipe(
471 Effect.mapError(
472 (cause) => new MuseReplayTranscriptDecodeError({ scenario: transcript.scenario, cause }),
473 ),
474 ),
475 makeProviderAdapterRegistryLayer: (transcript) =>
476 Layer.unwrap(
477 Effect.gen(function* () {
478 const controller = new MuseReplayController(transcript);
479 yield* Effect.addFinalizer(() => Effect.sync(() => controller.assertComplete()));
480 return layer({
481 scenario: transcript.scenario,
482 createHost: makeReplayCreateHost(controller),
483 });
484 }),
485 ),
486};