1import {
2 DEFAULT_TEXT_GENERATION_REASONING_EFFORT,
3 MUSE_DEFAULT_MODEL,
4 type ModelSelection,
5 type ServerProviderModel,
6 TextGenerationError,
7} from "@t3tools/contracts";
8import type { MuseSettings } from "../settings.ts";
9import { sanitizeBranchFragment, sanitizeFeatureBranchName } from "@t3tools/shared/git";
10import { getModelSelectionStringOptionValue } from "@t3tools/shared/model";
11import { extractJsonObject } from "@t3tools/shared/schemaJson";
12import * as Effect from "effect/Effect";
13import * as FileSystem from "effect/FileSystem";
14import * as Option from "effect/Option";
15import * as Schema from "effect/Schema";
16
17import { createMuseSdkHost, createMuseSdkHostEffect, type MuseSdkHost } from "./sdk.ts";
18import { museModelCapabilities, resolveMuseReasoningEffort } from "./modelCatalog.ts";
19import type { ProviderTextGeneration } from "@t3tools/provider-core/server/textGeneration";
20import { investigationUnsupported } from "@t3tools/provider-core/server/textGenerationOperations";
21import {
22 buildBranchNamePrompt,
23 buildCommitMessagePrompt,
24 buildPrContentPrompt,
25 buildThreadTitlePrompt,
26} from "@t3tools/provider-core/server/textGenerationPrompts";
27import {
28 sanitizeCommitSubject,
29 sanitizePrTitle,
30 sanitizeThreadTitle,
31 toJsonSchemaObject,
32} from "@t3tools/provider-core/server/textGenerationUtils";
33
34const SessionStarted = Schema.Struct({ session: Schema.Struct({ sessionId: Schema.String }) });
35const ItemNotification = Schema.Struct({
36 sessionId: Schema.String,
37 item: Schema.Struct({
38 itemId: Schema.String,
39 kind: Schema.String,
40 revision: Schema.Int,
41 turnId: Schema.optionalKey(Schema.String),
42 text: Schema.optionalKey(Schema.String),
43 retracted: Schema.optionalKey(Schema.Boolean),
44 truncated: Schema.optionalKey(Schema.Boolean),
45 }),
46});
47const TurnCompleted = Schema.Struct({
48 sessionId: Schema.String,
49 turnId: Schema.String,
50 terminal: Schema.String,
51 error: Schema.optionalKey(Schema.Struct({ message: Schema.String })),
52});
53const ApprovalRequested = Schema.Struct({
54 sessionId: Schema.String,
55 approvalId: Schema.String,
56 currentRequirementId: Schema.Unknown,
57 availableChoices: Schema.Array(
58 Schema.Struct({ choiceId: Schema.String, decision: Schema.String }),
59 ),
60});
61const UserInputRequested = Schema.Struct({
62 sessionId: Schema.String,
63 userInputId: Schema.String,
64});
65const decodeSession = Schema.decodeUnknownSync(SessionStarted);
66const decodeItem = Schema.decodeUnknownSync(ItemNotification);
67const decodeCompleted = Schema.decodeUnknownSync(TurnCompleted);
68const decodeApproval = Schema.decodeUnknownSync(ApprovalRequested);
69const decodeUserInput = Schema.decodeUnknownSync(UserInputRequested);
70const encodeJson = Schema.encodeEffect(Schema.fromJsonString(Schema.Unknown));
71const decodeJson = Schema.decodeEffect(Schema.fromJsonString(Schema.Unknown));
72
73type Operation = keyof ProviderTextGeneration;
74
75async function generateMuseText(
76 host: MuseSdkHost,
77 workspaceRoot: string,
78 prompt: string,
79 modelId: string | undefined,
80 reasoningEffort: string | undefined,
81 signal: AbortSignal,
82) {
83 let sessionId: string | undefined;
84 const turnId = host.connection.mintCommandId();
85 const messages = new Map<string, (typeof ItemNotification.Type)["item"]>();
86 let resolveCompletion: (text: string) => void = () => {};
87 let rejectCompletion: (error: unknown) => void = () => {};
88 const completion = new Promise<string>((resolve, reject) => {
89 resolveCompletion = resolve;
90 rejectCompletion = reject;
91 });
92 // Notifications may settle the turn before the turn/start acknowledgement arrives.
93 void completion.catch(() => {});
94 const rejectInteractive = async (method: string, params: Record<string, unknown>) => {
95 if (method === "approval/requested") {
96 const approval = decodeApproval(params);
97 if (approval.sessionId !== sessionId) return;
98 rejectCompletion(new Error("Muse requested an interactive action during text generation."));
99 const deny = approval.availableChoices.find((choice) => choice.decision === "denied");
100 if (deny) {
101 await host.connection.command("approval/decide", {
102 sessionId,
103 approvalId: approval.approvalId,
104 requirementId: approval.currentRequirementId,
105 choiceId: deny.choiceId,
106 });
107 }
108 } else {
109 const request = decodeUserInput(params);
110 if (request.sessionId !== sessionId) return;
111 rejectCompletion(new Error("Muse requested an interactive action during text generation."));
112 await host.connection.command("userInput/cancel", {
113 sessionId,
114 userInputId: request.userInputId,
115 reason: "T3 Code text generation cannot request user input.",
116 });
117 }
118 };
119 host.connection.onNotification((notification) => {
120 try {
121 if (notification.method === "item/completed" || notification.method === "item/updated") {
122 const { item, sessionId: itemSessionId } = decodeItem(notification.params);
123 if (itemSessionId !== sessionId || item.turnId !== turnId || item.kind !== "agentMessage")
124 return;
125 const previous = messages.get(item.itemId);
126 if (!previous || item.revision > previous.revision) messages.set(item.itemId, item);
127 } else if (notification.method === "turn/completed") {
128 const completed = decodeCompleted(notification.params);
129 if (completed.sessionId !== sessionId || completed.turnId !== turnId) return;
130 if (completed.terminal === "completed") {
131 // A turn may include commentary before its final structured answer.
132 const answer = [...messages.values()].findLast((item) => !item.retracted);
133 if (answer?.truncated) {
134 rejectCompletion(new Error("Muse text generation returned a truncated response."));
135 } else {
136 resolveCompletion(answer?.text ?? "");
137 }
138 } else {
139 rejectCompletion(
140 new Error(completed.error?.message ?? `Muse turn ${completed.terminal}.`),
141 );
142 }
143 } else if (
144 notification.method === "approval/requested" ||
145 notification.method === "userInput/requested"
146 ) {
147 void rejectInteractive(notification.method, notification.params ?? {}).catch(
148 rejectCompletion,
149 );
150 } else if (notification.method === "view/gap") {
151 rejectCompletion(new Error("Muse text generation lost part of its response."));
152 }
153 } catch (error) {
154 rejectCompletion(error);
155 }
156 });
157 host.connection.onProtocolError(rejectCompletion);
158 host.connection.onServerRequest(async () => {
159 const error = new Error("T3 Code text generation cannot approve interactive requests.");
160 rejectCompletion(error);
161 throw error;
162 });
163 void host.connection.closed.then(() => {
164 rejectCompletion(new Error("Muse connection closed before text generation completed."));
165 });
166 void host.exited.then(() => {
167 rejectCompletion(new Error("Muse exited before text generation completed."));
168 });
169 const onAbort = () => rejectCompletion(signal.reason);
170 signal.addEventListener("abort", onAbort, { once: true });
171 try {
172 signal.throwIfAborted();
173 const started = decodeSession(
174 await host.connection.command("session/start", {
175 workspaceRoot,
176 providerId: "meta",
177 ...(modelId ? { modelId } : {}),
178 approvalMode: "denyUnmatched",
179 }),
180 );
181 sessionId = started.session.sessionId;
182 const acknowledged = host.connection.command(
183 "turn/start",
184 {
185 sessionId,
186 input: [{ type: "text", text: prompt }],
187 ...(reasoningEffort !== undefined ? { reasoningEffort } : {}),
188 },
189 { commandId: turnId },
190 );
191 return await Promise.race([acknowledged.then(() => completion), completion]);
192 } finally {
193 signal.removeEventListener("abort", onAbort);
194 }
195}
196
197export const makeMuseTextGeneration = Effect.fn("makeMuseTextGeneration")(function* (
198 settings: MuseSettings,
199 options: {
200 readonly environment?: NodeJS.ProcessEnv;
201 readonly modelCatalog?: Effect.Effect<ReadonlyArray<ServerProviderModel>>;
202 /** Test hook. */
203 readonly createHost?: typeof createMuseSdkHost;
204 } = {},
205) {
206 const { environment, modelCatalog = Effect.succeed([]), createHost } = options;
207 const fileSystem = yield* FileSystem.FileSystem;
208 const runMuseJson = Effect.fn("runMuseJson")(function* <S extends Schema.Top>(input: {
209 operation: Operation;
210 prompt: string;
211 outputSchema: S;
212 modelSelection: ModelSelection;
213 }): Effect.fn.Return<S["Type"], TextGenerationError, S["DecodingServices"]> {
214 const { operation } = input;
215 if (!settings.enabled) {
216 return yield* new TextGenerationError({ operation, detail: "Muse Code is disabled." });
217 }
218 const selectedModel = input.modelSelection.model.trim();
219 const models = yield* modelCatalog;
220 // "default" is the catalog's default model, as in chat. Without one, Muse picks.
221 const catalogModel =
222 selectedModel && selectedModel !== MUSE_DEFAULT_MODEL
223 ? models.find((model) => model.slug === selectedModel)
224 : models.find((model) => model.isDefault && !model.isCustom);
225 const modelId =
226 selectedModel && selectedModel !== MUSE_DEFAULT_MODEL ? selectedModel : catalogModel?.slug;
227 const reasoningEffort = resolveMuseReasoningEffort(
228 catalogModel?.capabilities ?? museModelCapabilities(),
229 getModelSelectionStringOptionValue(input.modelSelection, "reasoningEffort") ??
230 DEFAULT_TEXT_GENERATION_REASONING_EFFORT,
231 );
232 const jsonSchema = yield* encodeJson(toJsonSchemaObject(input.outputSchema)).pipe(
233 Effect.mapError(
234 (cause) => new TextGenerationError({ operation, detail: "Invalid output schema.", cause }),
235 ),
236 );
237 const response = yield* Effect.gen(function* () {
238 // Metadata requests already include their context and do not need checkout configuration.
239 const cwd = yield* fileSystem.makeTempDirectoryScoped({ prefix: "t3code-muse-text-" });
240 const host = yield* Effect.acquireRelease(
241 createMuseSdkHostEffect(
242 {
243 binaryPath: settings.binaryPath,
244 cwd,
245 ...(environment ? { environment } : {}),
246 readOnly: true,
247 // Muse derives SDK turn events from its session log, including completion.
248 sessionLogging: true,
249 },
250 createHost,
251 ),
252 (host: MuseSdkHost) => Effect.promise(() => host.close()),
253 { interruptible: true },
254 );
255 return yield* Effect.tryPromise((signal) =>
256 generateMuseText(
257 host,
258 cwd,
259 `${input.prompt}\n\nReturn only a JSON object matching this schema, with no markdown fences. Do not use tools or ask questions.\n${jsonSchema}`,
260 modelId,
261 reasoningEffort,
262 signal,
263 ),
264 );
265 }).pipe(
266 Effect.scoped,
267 Effect.mapError(
268 (cause) =>
269 new TextGenerationError({
270 operation,
271 detail:
272 "Muse Code text generation failed. Check Muse login and availability on this T3 server host.",
273 cause,
274 }),
275 ),
276 Effect.timeoutOption(180_000),
277 Effect.flatMap(
278 Option.match({
279 onNone: () =>
280 Effect.fail(new TextGenerationError({ operation, detail: "Muse request timed out." })),
281 onSome: Effect.succeed,
282 }),
283 ),
284 );
285 const decodeOutput = Schema.decodeEffect(input.outputSchema);
286 return yield* decodeJson(extractJsonObject(response)).pipe(
287 Effect.flatMap(decodeOutput),
288 Effect.mapError(
289 (cause) =>
290 new TextGenerationError({
291 operation,
292 detail: "Muse returned invalid structured output.",
293 cause,
294 }),
295 ),
296 );
297 });
298
299 const generateCommitMessage: ProviderTextGeneration["generateCommitMessage"] = Effect.fn(
300 "MuseTextGeneration.generateCommitMessage",
301 )(function* (input) {
302 const generated = yield* runMuseJson({
303 operation: "generateCommitMessage",
304 ...buildCommitMessagePrompt({ ...input, includeBranch: input.includeBranch === true }),
305 modelSelection: input.modelSelection,
306 });
307 return {
308 subject: sanitizeCommitSubject(generated.subject),
309 body: generated.body.trim(),
310 ...("branch" in generated && typeof generated.branch === "string"
311 ? { branch: sanitizeFeatureBranchName(generated.branch) }
312 : {}),
313 };
314 });
315 const generatePrContent: ProviderTextGeneration["generatePrContent"] = Effect.fn(
316 "MuseTextGeneration.generatePrContent",
317 )(function* (input) {
318 const generated = yield* runMuseJson({
319 operation: "generatePrContent",
320 ...buildPrContentPrompt(input),
321 modelSelection: input.modelSelection,
322 });
323 return { title: sanitizePrTitle(generated.title), body: generated.body.trim() };
324 });
325 const generateBranchName: ProviderTextGeneration["generateBranchName"] = Effect.fn(
326 "MuseTextGeneration.generateBranchName",
327 )(function* (input) {
328 const generated = yield* runMuseJson({
329 operation: "generateBranchName",
330 ...buildBranchNamePrompt(input),
331 modelSelection: input.modelSelection,
332 });
333 return { branch: sanitizeBranchFragment(generated.branch) };
334 });
335 const generateThreadTitle: ProviderTextGeneration["generateThreadTitle"] = Effect.fn(
336 "MuseTextGeneration.generateThreadTitle",
337 )(function* (input) {
338 const generated = yield* runMuseJson({
339 operation: "generateThreadTitle",
340 ...buildThreadTitlePrompt(input),
341 modelSelection: input.modelSelection,
342 });
343 return { title: sanitizeThreadTitle(generated.title) };
344 });
345 // Muse runs in an empty read-only directory and cannot be pointed at a checkout to read.
346 const generateReviewGuide: ProviderTextGeneration["generateReviewGuide"] = () =>
347 Effect.fail(investigationUnsupported("generateReviewGuide", "Muse Code"));
348 return {
349 generateCommitMessage,
350 generatePrContent,
351 generateBranchName,
352 generateThreadTitle,
353 generateReviewGuide,
354 } satisfies ProviderTextGeneration;
355});