AntigravityTextGeneration.ts

import {
  type ModelSelection,
  type ProviderSetupError,
  TextGenerationError,
} from "@t3tools/contracts";
import { sanitizeBranchFragment, sanitizeFeatureBranchName } from "@t3tools/shared/git";
import { extractJsonObject } from "@t3tools/shared/schemaJson";
import * as Deferred from "effect/Deferred";
import * as Effect from "effect/Effect";
import * as Exit from "effect/Exit";
import * as FileSystem from "effect/FileSystem";
import * as Option from "effect/Option";
import * as Path from "effect/Path";
import * as Ref from "effect/Ref";
import * as Schema from "effect/Schema";
import * as Scope from "effect/Scope";
import * as Stream from "effect/Stream";
import { type AcpError, AcpRequestError } from "effect-acp/errors";

import { applyAntigravityAcpModelSelection } from "../provider/acp/AntigravityAcpSupport.ts";
import { removeAntigravitySessionFiles } from "../provider/acp/AntigravitySessionFiles.ts";
import type { AcpSessionRuntime } from "../provider/acp/AcpSessionRuntime.ts";
import type * as TextGeneration from "./TextGeneration.ts";
import {
  buildBranchNamePrompt,
  buildCommitMessagePrompt,
  buildPrContentPrompt,
  buildPullRequestRankingPrompt,
  buildThreadTitlePrompt,
} from "./TextGenerationPrompts.ts";
import {
  clampPullRequestRankings,
  sanitizeCommitSubject,
  sanitizePrTitle,
  sanitizeThreadTitle,
} from "./TextGenerationUtils.ts";

const ANTIGRAVITY_TIMEOUT_MS = 180_000;
const MAX_OUTPUT_CHARS = 128_000;
const isTextGenerationError = Schema.is(TextGenerationError);
const isNativeSessionId = Schema.is(Schema.String.check(Schema.isUUID(4)));
const Configuration = Schema.Record(Schema.String, Schema.Unknown);
const decodeConfiguration = Schema.decodeEffect(Schema.fromJsonString(Configuration));
const decodeConfigurationObject = Schema.decodeUnknownEffect(Configuration);

type AntigravityTextRuntime = Pick<
  AcpSessionRuntime["Service"],
  | "start"
  | "setMode"
  | "getConfigOptions"
  | "getEvents"
  | "setModel"
  | "prompt"
  | "cancel"
  | "handleSessionUpdate"
  | "handleRequestPermission"
  | "handleElicitation"
  | "handleReadTextFile"
  | "handleWriteTextFile"
  | "handleCreateTerminal"
  | "handleTerminalOutput"
  | "handleTerminalWaitForExit"
  | "handleTerminalKill"
  | "handleTerminalRelease"
  | "handleUnknownExtRequest"
>;

export interface AntigravityTextGenerationOptions {
  readonly profileDirectory: string;
  /** Model the provider default alias selects, when the account offers it. */
  readonly defaultModel?: Effect.Effect<string | undefined>;
  /** Uses the instance's personal Google login, with no injected MCP servers or client tools. */
  readonly makeRuntime: (
    cwd: string,
  ) => Effect.Effect<AntigravityTextRuntime, AcpError | ProviderSetupError, Scope.Scope>;
  /** Registers the whole helper so sign-out can stop it before clearing credentials. */
  readonly withProcess: <A, E, R>(
    stop: Effect.Effect<void>,
    task: Effect.Effect<A, E, R>,
  ) => Effect.Effect<A, E | ProviderSetupError, R | Scope.Scope>;
}

/** Global hooks and MCP servers can run before a helper can deny a tool request. */
export const isAntigravityTextGenerationAvailable = Effect.fn(
  "isAntigravityTextGenerationAvailable",
)(function* (profileDirectory: string) {
  const fs = yield* FileSystem.FileSystem;
  const path = yield* Path.Path;
  for (const name of ["hooks.json", "mcp_config.json"]) {
    const configurationPath = path.join(profileDirectory, "config", name);
    if (!(yield* fs.exists(configurationPath))) {
      continue;
    }
    const info = yield* fs.stat(configurationPath);
    if (info.type !== "File" || info.size > 64_000n) {
      return false;
    }
    const empty = yield* fs.readFileString(configurationPath).pipe(
      Effect.flatMap(decodeConfiguration),
      Effect.flatMap((configuration) =>
        decodeConfigurationObject(
          configuration[name === "hooks.json" ? "hooks" : "mcpServers"] ?? configuration,
        ),
      ),
      Effect.map((configuration) => Object.keys(configuration).length === 0),
      Effect.orElseSucceed(() => false),
    );
    if (!empty) return false;
  }
  return true;
});

/** Runs short-lived subscription helpers without giving them the user's workspace. */
export const makeAntigravityTextGeneration = Effect.fn("makeAntigravityTextGeneration")(function* (
  options: AntigravityTextGenerationOptions,
) {
  const fs = yield* FileSystem.FileSystem;
  const path = yield* Path.Path;
  const available = isAntigravityTextGenerationAvailable(options.profileDirectory).pipe(
    Effect.provideService(FileSystem.FileSystem, fs),
    Effect.provideService(Path.Path, path),
  );

  const runAntigravityJson = Effect.fn("AntigravityTextGeneration.runJson")(
    function* <S extends Schema.Top>(input: {
      readonly operation: keyof TextGeneration.TextGeneration["Service"];
      readonly prompt: string;
      readonly outputSchema: S;
      readonly modelSelection: ModelSelection;
    }) {
      const { operation } = input;
      const scope = yield* Scope.make();
      yield* Effect.addFinalizer((exit) => Scope.close(scope, exit));
      const helper = Effect.gen(function* () {
        if (!(yield* available)) {
          return yield* new TextGenerationError({
            operation,
            detail:
              "Antigravity text generation is unavailable for profiles with global hooks or MCP configuration. Select another system model.",
          });
        }

        const cwd = yield* fs.makeTempDirectoryScoped({ prefix: "t3-antigravity-text-" });
        let sessionId: string | undefined;
        yield* Effect.addFinalizer(() =>
          removeAntigravitySessionFiles({
            profileDirectory: options.profileDirectory,
            sessionId,
            cwd,
          }).pipe(
            Effect.provideService(FileSystem.FileSystem, fs),
            Effect.provideService(Path.Path, path),
          ),
        );

        const rawResult = yield* Effect.gen(function* () {
          const runtime = yield* options.makeRuntime(cwd);
          yield* runtime.getEvents().pipe(
            Stream.runForEach((event) =>
              event._tag === "EventStreamBarrier"
                ? Deferred.succeed(event.acknowledge, undefined).pipe(Effect.asVoid)
                : Effect.void,
            ),
            Effect.forkScoped,
          );
          const output = yield* Ref.make("");
          const rejected = yield* Deferred.make<never, TextGenerationError>();
          const reject = (detail: string) =>
            Deferred.fail(rejected, new TextGenerationError({ operation, detail })).pipe(
              Effect.asVoid,
            );
          const rejectToolRequest = () =>
            reject("Antigravity text generation requested a tool or user input.").pipe(
              Effect.andThen(
                Effect.fail(
                  new AcpRequestError({
                    code: -32601,
                    errorMessage: "Tools and user input are disabled for text generation.",
                  }),
                ),
              ),
            );

          yield* runtime.handleRequestPermission(() =>
            reject("Antigravity text generation requested a tool permission or user input.").pipe(
              Effect.as({ outcome: { outcome: "cancelled" as const } }),
            ),
          );
          yield* runtime.handleElicitation(() =>
            reject("Antigravity text generation requested user input.").pipe(
              Effect.as({ action: { action: "decline" as const } }),
            ),
          );
          yield* runtime.handleReadTextFile(rejectToolRequest);
          yield* runtime.handleWriteTextFile(rejectToolRequest);
          yield* runtime.handleCreateTerminal(rejectToolRequest);
          yield* runtime.handleTerminalOutput(rejectToolRequest);
          yield* runtime.handleTerminalWaitForExit(rejectToolRequest);
          yield* runtime.handleTerminalKill(rejectToolRequest);
          yield* runtime.handleTerminalRelease(rejectToolRequest);
          yield* runtime.handleUnknownExtRequest(rejectToolRequest);
          yield* runtime.handleSessionUpdate((notification) =>
            Effect.gen(function* () {
              const update = notification.update;
              if (
                update.sessionUpdate === "tool_call" ||
                update.sessionUpdate === "tool_call_update"
              ) {
                return yield* reject("Antigravity attempted tool work during text generation.");
              }
              if (
                notification.sessionId !== sessionId ||
                update.sessionUpdate !== "agent_message_chunk" ||
                update.content.type !== "text"
              ) {
                return;
              }
              const text = update.content.text;
              const exceeded = yield* Ref.modify(output, (current) =>
                current.length + text.length > MAX_OUTPUT_CHARS
                  ? [true, current]
                  : [false, current + text],
              );
              if (exceeded) {
                return yield* reject("Antigravity text generation exceeded the output limit.");
              }
            }),
          );

          return yield* Effect.gen(function* () {
            const started = yield* runtime.start();
            sessionId = started.sessionId;
            if (!isNativeSessionId(sessionId)) {
              return yield* new TextGenerationError({
                operation,
                detail: "Antigravity returned an invalid text helper session ID.",
              });
            }
            yield* runtime.setMode("default");
            yield* applyAntigravityAcpModelSelection({
              runtime,
              model: input.modelSelection.model,
              defaultModel: yield* options.defaultModel ?? Effect.succeed(undefined),
              mapError: (cause) =>
                new TextGenerationError({
                  operation,
                  detail: "Could not select the Antigravity model for text generation.",
                  cause,
                }),
            });

            const result = yield* runtime.prompt({
              prompt: [
                {
                  type: "text",
                  text: [
                    "Use only the input below. Do not use tools, read or write files, run commands, or ask questions.",
                    "Return only the requested JSON object.",
                    "",
                    input.prompt,
                  ].join("\n"),
                },
              ],
            });
            if (yield* Deferred.isDone(rejected)) {
              return yield* Deferred.await(rejected);
            }
            if (result.stopReason === "cancelled") {
              return yield* new TextGenerationError({
                operation,
                detail: "Antigravity text generation was cancelled.",
              });
            }
            return (yield* Ref.get(output)).trim();
          }).pipe(
            Effect.onInterrupt(() =>
              runtime.cancel.pipe(Effect.timeoutOption(2_000), Effect.ignore),
            ),
            Effect.raceFirst(Deferred.await(rejected)),
          );
        }).pipe(Effect.scoped);

        if ((yield* fs.readDirectory(cwd)).length > 0) {
          return yield* new TextGenerationError({
            operation,
            detail: "Antigravity wrote files during text generation.",
          });
        }
        if (!rawResult) {
          return yield* new TextGenerationError({
            operation,
            detail: "Antigravity returned empty text generation output.",
          });
        }
        const decodeOutput = Schema.decodeEffect(Schema.fromJsonString(input.outputSchema));
        return yield* decodeOutput(extractJsonObject(rawResult)).pipe(
          Effect.mapError(
            (cause) =>
              new TextGenerationError({
                operation,
                detail: "Antigravity returned invalid structured output.",
                cause,
              }),
          ),
        );
      }).pipe(
        Effect.scoped,
        Effect.timeoutOption(ANTIGRAVITY_TIMEOUT_MS),
        Effect.flatMap(
          Option.match({
            onNone: () =>
              Effect.fail(
                new TextGenerationError({
                  operation,
                  detail: "Antigravity text generation timed out.",
                }),
              ),
            onSome: Effect.succeed,
          }),
        ),
      );

      return yield* options
        .withProcess(Scope.close(scope, Exit.void), helper)
        .pipe(Effect.provideService(Scope.Scope, scope));
    },
    (effect, input) =>
      effect.pipe(
        Effect.mapError((cause) =>
          isTextGenerationError(cause)
            ? cause
            : new TextGenerationError({
                operation: input.operation,
                detail: "Antigravity text generation failed.",
                cause,
              }),
        ),
        Effect.scoped,
      ),
  );

  const generateCommitMessage: TextGeneration.TextGeneration["Service"]["generateCommitMessage"] =
    Effect.fn("AntigravityTextGeneration.generateCommitMessage")(function* (input) {
      const generated = yield* runAntigravityJson({
        operation: "generateCommitMessage",
        ...buildCommitMessagePrompt({
          branch: input.branch,
          stagedSummary: input.stagedSummary,
          stagedPatch: input.stagedPatch,
          includeBranch: input.includeBranch === true,
          policy: input.policy,
        }),
        modelSelection: input.modelSelection,
      });
      return {
        subject: sanitizeCommitSubject(generated.subject),
        body: generated.body.trim(),
        ...("branch" in generated && typeof generated.branch === "string"
          ? { branch: sanitizeFeatureBranchName(generated.branch) }
          : {}),
      };
    });

  const generatePrContent: TextGeneration.TextGeneration["Service"]["generatePrContent"] =
    Effect.fn("AntigravityTextGeneration.generatePrContent")(function* (input) {
      const generated = yield* runAntigravityJson({
        operation: "generatePrContent",
        ...buildPrContentPrompt({
          baseBranch: input.baseBranch,
          headBranch: input.headBranch,
          commitSummary: input.commitSummary,
          diffSummary: input.diffSummary,
          diffPatch: input.diffPatch,
          policy: input.policy,
          changeRequestTemplate: input.changeRequestTemplate,
        }),
        modelSelection: input.modelSelection,
      });
      return { title: sanitizePrTitle(generated.title), body: generated.body.trim() };
    });

  const generateBranchName: TextGeneration.TextGeneration["Service"]["generateBranchName"] =
    Effect.fn("AntigravityTextGeneration.generateBranchName")(function* (input) {
      const generated = yield* runAntigravityJson({
        operation: "generateBranchName",
        ...buildBranchNamePrompt({ message: input.message, attachments: input.attachments }),
        modelSelection: input.modelSelection,
      });
      return { branch: sanitizeBranchFragment(generated.branch) };
    });

  const generateThreadTitle: TextGeneration.TextGeneration["Service"]["generateThreadTitle"] =
    Effect.fn("AntigravityTextGeneration.generateThreadTitle")(function* (input) {
      const generated = yield* runAntigravityJson({
        operation: "generateThreadTitle",
        ...buildThreadTitlePrompt({
          message: input.message,
          previousTitle: input.previousTitle,
          attachments: input.attachments,
        }),
        modelSelection: input.modelSelection,
      });
      return { title: sanitizeThreadTitle(generated.title) };
    });

  const rankPullRequests: TextGeneration.TextGeneration["Service"]["rankPullRequests"] = Effect.fn(
    "AntigravityTextGeneration.rankPullRequests",
  )(function* (input) {
    const generated = yield* runAntigravityJson({
      operation: "rankPullRequests",
      ...buildPullRequestRankingPrompt({
        repository: input.repository,
        intoRepository: input.intoRepository,
        candidates: input.candidates,
      }),
      modelSelection: input.modelSelection,
    });
    return {
      rankings: clampPullRequestRankings(generated.rankings),
    } satisfies TextGeneration.PullRequestRankingResult;
  });

  return {
    generateCommitMessage,
    generatePrContent,
    generateBranchName,
    generateThreadTitle,
    rankPullRequests,
  } satisfies TextGeneration.TextGeneration["Service"];
});