@prismatic-io/lux 0.0.2-preview.4 → 0.0.2-preview.40

package/src/drivers/grok-build/index.ts

@@ -0,0 +1,331 @@
+import { realpath } from "node:fs/promises";
+import { agentEnvironment } from "../../core/agent-environment.js";
+import { durationMs } from "../../core/duration.js";
+import {
+  type AgentDriver,
+  type Answer,
+  type Artifact,
+  type DriverEventMap,
+  defineDriver,
+  discoverCliVersion,
+  type ReadyState,
+  type StartContext,
+  selectModelAndReasoningEffort,
+  TypedEmitter,
+} from "../../core/index.js";
+import { AcpEvents } from "../shared/acp-events.js";
+import { permissionAnswer, permissionRequestText, records } from "../shared/acp-interaction.js";
+import { AcpTransport, asRecord, type RpcId, type RpcRecord } from "../shared/acp-transport.js";
+import { walkArtifacts } from "../shared/artifacts.js";
+import {
+  buildGrokBuildArgs,
+  type GrokBuildDriverConfig,
+  GrokBuildDriverConfigSchema,
+} from "./config.js";
+
+class GrokBuildDriver implements AgentDriver {
+  private readonly emitter = new TypedEmitter<DriverEventMap>();
+  private readonly events: AcpEvents;
+  private transport: AcpTransport | undefined;
+  private sessionId: string | undefined;
+  private cwd: string | undefined;
+  private artifactsDir: string | undefined;
+  private ready = false;
+  private runPromise: Promise<void> | undefined;
+  private terminal = false;
+  private idleTimer: NodeJS.Timeout | undefined;
+  private interruptCount = 0;
+  private readonly interrupts = new Map<
+    string,
+    { id: RpcId; params: RpcRecord; timer: NodeJS.Timeout }
+  >();
+
+  private readonly config: GrokBuildDriverConfig;
+
+  constructor(config: GrokBuildDriverConfig) {
+    this.config = config;
+    this.events = new AcpEvents("grok-build", config.maxToolCalls, config.maxQueuedBytes);
+  }
+
+  on<K extends keyof DriverEventMap>(
+    event: K,
+    handler: (e: DriverEventMap[K]) => void | Promise<void>,
+  ): void {
+    this.emitter.on(event, handler);
+  }
+  off<K extends keyof DriverEventMap>(
+    event: K,
+    handler: (e: DriverEventMap[K]) => void | Promise<void>,
+  ): void {
+    this.emitter.off(event, handler);
+  }
+
+  async start(ctx: StartContext): Promise<ReadyState> {
+    if (this.transport || this.terminal)
+      throw new Error("Grok.build driver cannot be started twice or after close");
+    this.cwd = await realpath(this.config.cwd ?? ctx.artifactsDir);
+    this.artifactsDir = await realpath(ctx.artifactsDir);
+    const cliVersion = await discoverCliVersion(this.config.command);
+    const transport = new AcpTransport({
+      label: "Grok.build ACP",
+      command: this.config.command,
+      args: buildGrokBuildArgs(this.config),
+      cwd: this.cwd,
+      env: agentEnvironment(process.env, this.config.inheritEnv, this.config.env),
+      signal: ctx.abortSignal,
+      maxLineBytes: this.config.maxLineBytes,
+      maxQueuedBytes: this.config.maxQueuedBytes,
+      onMessage: (message) => this.message(message),
+      onFailure: (error) => this.fail(error),
+      onActivity: () => this.bumpIdleTimer(),
+    });
+    this.transport = transport;
+    try {
+      await transport.spawned();
+      const initialized = await this.request("initialize", {
+        protocolVersion: 1,
+        clientInfo: { name: "lux", version: "1" },
+        clientCapabilities: { fs: { readTextFile: false, writeTextFile: false }, terminal: false },
+      });
+      if (initialized.protocolVersion !== 1)
+        throw new Error("Grok.build ACP returned an unsupported protocol version");
+      if (asRecord(initialized._meta)?.grokShell !== true)
+        throw new Error("Expected Grok.build ACP; check the configured command points to grok");
+      const session = await this.request("session/new", {
+        cwd: this.cwd,
+        mcpServers: [],
+        _meta: { autoMode: this.config.permissionMode === "auto", yoloMode: false },
+      });
+      if (typeof session.sessionId !== "string")
+        throw new Error("Grok.build ACP omitted sessionId");
+      this.sessionId = session.sessionId;
+      let nativeConfig = session.configOptions;
+      for (const [configId, value] of [
+        ["model", this.config.model],
+        ["reasoning_effort", this.config.reasoningEffort],
+      ] as const) {
+        const updated = await this.request("session/set_config_option", {
+          sessionId: this.sessionId,
+          configId,
+          value,
+        });
+        nativeConfig = updated.configOptions;
+        if (records(nativeConfig).find((option) => option.id === configId)?.currentValue !== value)
+          throw new Error(`Grok.build ACP did not apply requested ${configId} '${value}'`);
+      }
+      this.ready = true;
+      this.bumpIdleTimer();
+      this.runPromise = this.run(ctx.prompt).catch((error: unknown) => this.fail(error));
+      return {
+        agentId: "grok-build",
+        model: this.config.model,
+        ...(typeof asRecord(initialized._meta)?.agentVersion === "string"
+          ? { agentVersion: String(asRecord(initialized._meta)?.agentVersion) }
+          : {}),
+        detail: {
+          sessionId: this.sessionId,
+          transport: "acp",
+          protocolVersion: 1,
+          ...(cliVersion ? { cliVersion } : {}),
+          permissionMode: this.config.permissionMode,
+          sandbox: this.config.sandbox,
+          trustProject: this.config.trustProject,
+          reasoningEffort: this.config.reasoningEffort,
+          modelResolution: "reported",
+          nativeConfig: nativeConfig ?? [],
+          usageAvailability: "unknown",
+          isolation: "none",
+        },
+      };
+    } catch (error) {
+      await this.close();
+      throw error;
+    }
+  }
+
+  private request(method: string, params: RpcRecord): Promise<RpcRecord> {
+    if (!this.transport) throw new Error("Grok.build ACP has not started");
+    return this.transport.request(method, params, durationMs(this.config.startupTimeout));
+  }
+
+  private async run(prompt: string): Promise<void> {
+    const result = await this.transport?.request("session/prompt", {
+      sessionId: this.sessionId,
+      prompt: [{ type: "text", text: prompt }],
+    });
+    await this.transport?.drained();
+    if (this.terminal) return;
+    if (result?.stopReason !== "end_turn")
+      throw new Error(`Grok.build stopped with ${String(result?.stopReason ?? "unknown reason")}`);
+    if (this.interrupts.size > 0)
+      throw new Error("Grok.build ended with pending permission requests");
+    this.terminal = true;
+    clearTimeout(this.idleTimer);
+    await this.emitter.emit("done", {
+      exitReason: "done",
+      summary: {
+        sessionId: this.sessionId,
+        stopReason: result.stopReason,
+        usageAvailability: "unknown",
+      },
+    });
+  }
+
+  private async message(message: RpcRecord): Promise<void> {
+    if (this.terminal) return;
+    const method = String(message.method);
+    const params = asRecord(message.params) ?? {};
+    const id = message.id;
+    if (typeof id === "number" || typeof id === "string") {
+      if (params.sessionId !== this.sessionId) {
+        this.transport?.write({ id, error: { code: -32602, message: "Unknown session" } });
+        return;
+      }
+      if (method !== "session/request_permission") {
+        this.transport?.write({
+          id,
+          error: { code: -32601, message: `Unsupported method: ${method}` },
+        });
+        return;
+      }
+      if (++this.interruptCount > this.config.maxInterrupts) {
+        this.transport?.write({ id, result: { outcome: { outcome: "cancelled" } } });
+        throw new Error("Grok.build exceeded maxInterrupts");
+      }
+      const interruptId = `grok-build-${id}`;
+      if (this.interrupts.has(interruptId))
+        throw new Error("Grok.build reused a pending request ID");
+      const timer = setTimeout(() => {
+        this.transport?.write({ id, result: { outcome: { outcome: "cancelled" } } });
+        this.interrupts.delete(interruptId);
+        this.fail(
+          new Error(
+            `Grok.build permission review timed out after ${durationMs(this.config.permissionTimeout)}ms`,
+          ),
+        );
+      }, durationMs(this.config.permissionTimeout));
+      this.interrupts.set(interruptId, { id, params, timer });
+      clearTimeout(this.idleTimer);
+      void this.emitter
+        .emit("interrupt", {
+          kind: "approve",
+          id: interruptId,
+          request: `${permissionRequestText(params)}\n\nOperation details:\n${JSON.stringify(params.toolCall)}`,
+          context: params,
+        })
+        .catch((error: unknown) => this.fail(error));
+      return;
+    }
+    if (method === "session/update") {
+      if (params.sessionId !== this.sessionId) return;
+      const update = asRecord(params.update);
+      if (update)
+        for (const event of this.events.update(update)) await this.emitter.emit("progress", event);
+    }
+  }
+
+  async respond(answer: Answer): Promise<void> {
+    if (this.terminal) throw new Error("Grok.build is closed");
+    const pending = this.interrupts.get(answer.id);
+    if (!pending) throw new Error(`Grok.build has no pending interrupt '${answer.id}'`);
+    const result = permissionAnswer(pending.params, answer);
+    this.transport?.write({ id: pending.id, result });
+    if (asRecord(result.outcome)?.outcome === "cancelled") {
+      this.fail(new Error("Grok.build offered no matching one-operation permission option"));
+      return;
+    }
+    clearTimeout(pending.timer);
+    this.interrupts.delete(answer.id);
+    this.bumpIdleTimer();
+  }
+
+  private bumpIdleTimer(): void {
+    clearTimeout(this.idleTimer);
+    if (!this.ready || this.terminal || this.interrupts.size > 0) return;
+    this.idleTimer = setTimeout(
+      () =>
+        this.fail(
+          new Error(`Grok.build emitted no event for ${durationMs(this.config.idleTimeout)}ms`),
+          "idle-timeout",
+        ),
+      durationMs(this.config.idleTimeout),
+    );
+  }
+
+  private fail(error: unknown, exitReason: "error" | "idle-timeout" = "error"): void {
+    if (this.terminal) return;
+    this.terminal = true;
+    clearTimeout(this.idleTimer);
+    for (const pending of this.interrupts.values()) clearTimeout(pending.timer);
+    if (this.ready)
+      void this.transport?.drained().then(() =>
+        this.emitter.emit("error", {
+          exitReason,
+          reason: error instanceof Error ? error.message : String(error),
+        }),
+      );
+    void this.transport?.close();
+  }
+
+  async quiesce(): Promise<void> {
+    clearTimeout(this.idleTimer);
+    const cancelling = !this.terminal && this.sessionId !== undefined;
+    this.terminal = true;
+    for (const pending of this.interrupts.values()) clearTimeout(pending.timer);
+    if (cancelling) {
+      for (const pending of this.interrupts.values())
+        this.transport?.write({ id: pending.id, result: { outcome: { outcome: "cancelled" } } });
+      this.transport?.write({ method: "session/cancel", params: { sessionId: this.sessionId } });
+      let timer: NodeJS.Timeout | undefined;
+      try {
+        await Promise.race([
+          this.runPromise,
+          new Promise<void>((resolve) => {
+            timer = setTimeout(resolve, 200);
+          }),
+        ]);
+      } finally {
+        clearTimeout(timer);
+      }
+    }
+    for (const pending of this.interrupts.values()) clearTimeout(pending.timer);
+    this.interrupts.clear();
+    await this.transport?.close();
+  }
+
+  async collect(): Promise<Artifact[]> {
+    await this.quiesce();
+    if (!this.cwd) return [];
+    const artifacts = await walkArtifacts(this.cwd);
+    return this.cwd === this.artifactsDir
+      ? artifacts
+      : artifacts.map((artifact) => ({ ...artifact, root: this.cwd as string }));
+  }
+
+  async close(): Promise<void> {
+    await this.quiesce();
+    this.emitter.removeAllListeners();
+  }
+}
+
+export const grokBuildDriver = defineDriver({
+  name: "grok-build",
+  configSchema: GrokBuildDriverConfigSchema,
+  runtime: (config) => ({
+    commands: [config.command],
+    capabilities: {
+      interactive: true,
+      "tool-calls": true,
+      "tool-results": true,
+      artifacts: true,
+      usage: false,
+      cost: false,
+      isolation: false,
+    },
+  }),
+  selectHarness: selectModelAndReasoningEffort,
+  create: (config) => new GrokBuildDriver(config),
+});
+
+export type { GrokBuildDriverConfig };
+export { buildGrokBuildArgs, GrokBuildDriverConfigSchema };

package/src/drivers/mcp/index.ts

@@ -1,12 +1,13 @@
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import { getDefaultEnvironment } from "@modelcontextprotocol/sdk/client/stdio.js";
 import { z } from "zod";
+import { DurationSchema, durationMs as parseDuration } from "../../core/duration.js";
 import {
   type AgentDriver,
   type Answer,
   type Artifact,
   type DriverEventMap,
-  type DriverPlugin,
+  defineDriver,
   type ListedTool,
   MCP_INIT_KIND,
   MCP_TOOLS_LIST_KIND,
@@ -72,9 +73,9 @@ export const McpProbeDriverConfigSchema = z
     /** The steps to run. Defaults to a single `tools/list` capture. */
     probe: z.array(ProbeStepSchema).default([{ op: "list-tools" }]),
     /** Timeout for the initialize handshake. */
-    startupTimeoutMs: z.number().positive().default(30_000),
+    startupTimeout: DurationSchema.default("30 seconds"),
     /** Timeout for each `tools/list` / `tools/call` request. */
-    stepTimeoutMs: z.number().positive().default(60_000),
+    stepTimeout: DurationSchema.default("1 minute"),
     /** Safety cap for a paginated `tools/list` operation. */
     maxListPages: z.number().int().positive().default(100),
   })
@@ -211,8 +212,8 @@ class McpProbeDriverImpl implements AgentDriver {
 
     try {
       await client.connect(transport, {
-        timeout: this.config.startupTimeoutMs,
-        maxTotalTimeout: this.config.startupTimeoutMs,
+        timeout: parseDuration(this.config.startupTimeout),
+        maxTotalTimeout: parseDuration(this.config.startupTimeout),
         signal: ctx.abortSignal,
       });
     } catch (err) {
@@ -261,8 +262,8 @@ class McpProbeDriverImpl implements AgentDriver {
               throw new Error(`tools/list exceeded ${this.config.maxListPages} pages`);
             }
             const result = await client.listTools(cursor ? { cursor } : undefined, {
-              timeout: this.config.stepTimeoutMs,
-              maxTotalTimeout: this.config.stepTimeoutMs,
+              timeout: parseDuration(this.config.stepTimeout),
+              maxTotalTimeout: parseDuration(this.config.stepTimeout),
               ...(this.abortSignal ? { signal: this.abortSignal } : {}),
             });
             tools.push(...(result.tools as ListedTool[]));
@@ -285,8 +286,8 @@ class McpProbeDriverImpl implements AgentDriver {
             { name: step.tool, arguments: step.args },
             undefined,
             {
-              timeout: this.config.stepTimeoutMs,
-              maxTotalTimeout: this.config.stepTimeoutMs,
+              timeout: parseDuration(this.config.stepTimeout),
+              maxTotalTimeout: parseDuration(this.config.stepTimeout),
               ...(this.abortSignal ? { signal: this.abortSignal } : {}),
             },
           );
@@ -344,8 +345,21 @@ class McpProbeDriverImpl implements AgentDriver {
   }
 }
 
-export const mcpProbeDriver: DriverPlugin<McpProbeDriverConfig> = {
+export const mcpProbeDriver = defineDriver({
   name: "mcp-probe",
   configSchema: McpProbeDriverConfigSchema,
+  runtime: (config) => ({
+    commands: [config.command],
+    capabilities: {
+      interactive: false,
+      "tool-calls": true,
+      "tool-results": true,
+      artifacts: true,
+      subagents: false,
+      usage: false,
+      cost: false,
+      isolation: false,
+    },
+  }),
   create: (config) => new McpProbeDriverImpl(config),
-};
+});

package/src/drivers/shared/acp-events.ts

@@ -0,0 +1,94 @@
+import { type Progress, toolResultProgress } from "../../core/index.js";
+import { records } from "./acp-interaction.js";
+import { asRecord, type RpcRecord } from "./acp-transport.js";
+
+type ToolState = { value: RpcRecord; bytes: number; name: string };
+
+/** Normalizes ACP `session/update` notifications into Lux progress events. */
+export class AcpEvents {
+  private seq = 0;
+  private retainedBytes = 0;
+  private readonly tools = new Map<string, ToolState>();
+  private readonly completed = new Set<string>();
+
+  private readonly prefix: string;
+  private readonly maxToolCalls: number;
+  private readonly maxBytes: number;
+
+  constructor(prefix: string, maxToolCalls: number, maxBytes: number) {
+    this.prefix = prefix;
+    this.maxToolCalls = maxToolCalls;
+    this.maxBytes = maxBytes;
+  }
+
+  update(update: RpcRecord): Progress[] {
+    const ts = Date.now();
+    if (
+      update.sessionUpdate === "agent_message_chunk" ||
+      update.sessionUpdate === "agent_thought_chunk"
+    ) {
+      const content = asRecord(update.content);
+      if (content?.type !== "text" || typeof content.text !== "string") return [];
+      return [
+        {
+          id: `${this.prefix}-message-${++this.seq}`,
+          kind: update.sessionUpdate === "agent_message_chunk" ? "agent-message" : "reasoning",
+          payload: { text: content.text },
+          ts,
+        },
+      ];
+    }
+    if (update.sessionUpdate !== "tool_call" && update.sessionUpdate !== "tool_call_update") {
+      return [
+        {
+          id: `${this.prefix}-update-${++this.seq}`,
+          kind: `${this.prefix}-update`,
+          payload: update,
+          ts,
+        },
+      ];
+    }
+    const id = update.toolCallId;
+    if (typeof id !== "string")
+      throw new Error(`${this.prefix} ACP tool update omitted toolCallId`);
+    if (this.completed.has(id)) return [];
+    const old = this.tools.get(id);
+    if (!old && this.tools.size + this.completed.size >= this.maxToolCalls)
+      throw new Error(`${this.prefix} ACP exceeded maxToolCalls`);
+    const value = { ...old?.value, ...update };
+    const bytes = Buffer.byteLength(JSON.stringify(value));
+    this.retainedBytes += bytes - (old?.bytes ?? 0);
+    if (this.retainedBytes > this.maxBytes)
+      throw new Error(`${this.prefix} ACP tool state exceeded maxQueuedBytes`);
+    const name = old?.name ?? String(value.title ?? value.kind ?? `${this.prefix}-tool`);
+    this.tools.set(id, { value, bytes, name });
+    const events: Progress[] = old
+      ? []
+      : [
+          {
+            id: `tc-${id}`,
+            kind: "tool-call",
+            payload: { name, input: value.rawInput ?? {}, toolUseId: id },
+            ts,
+          },
+        ];
+    if (value.status === "completed" || value.status === "failed") {
+      const content = value.rawOutput ?? value.content ?? [];
+      const text = records(value.content)
+        .map((block) => asRecord(block.content)?.text)
+        .filter((entry) => typeof entry === "string")
+        .join("\n");
+      events.push(
+        toolResultProgress(
+          ++this.seq,
+          { name, toolUseId: id, content, isError: value.status === "failed", text },
+          ts,
+        ),
+      );
+      this.retainedBytes -= bytes;
+      this.tools.delete(id);
+      this.completed.add(id);
+    }
+    return events;
+  }
+}

package/src/drivers/shared/acp-interaction.ts

@@ -0,0 +1,27 @@
+import type { Answer } from "../../core/index.js";
+import { asRecord, type RpcRecord } from "./acp-transport.js";
+
+export const records = (value: unknown): RpcRecord[] =>
+  Array.isArray(value)
+    ? value.flatMap((entry): RpcRecord[] => {
+        const record = asRecord(entry);
+        return record ? [record] : [];
+      })
+    : [];
+
+/** A Lux approval applies to this operation only, never a remembered permission. */
+export const permissionAnswer = (params: RpcRecord, answer: Answer): RpcRecord => {
+  if (answer.kind !== "approve") throw new Error("ACP permission requires an approval answer");
+  const option = records(params.options).find(
+    (entry) => entry.kind === (answer.approved ? "allow_once" : "reject_once"),
+  );
+  return {
+    outcome:
+      option && typeof option.optionId === "string"
+        ? { outcome: "selected", optionId: option.optionId }
+        : { outcome: "cancelled" },
+  };
+};
+
+export const permissionRequestText = (params: RpcRecord): string =>
+  String(asRecord(params.toolCall)?.title ?? "Approve this tool call?");

package/src/drivers/shared/acp-session.ts

@@ -0,0 +1,40 @@
+import { records } from "./acp-interaction.js";
+import { asRecord, type RpcRecord } from "./acp-transport.js";
+
+export const acpModel = (session: RpcRecord): unknown =>
+  records(session.configOptions).find(
+    (option) => option.category === "model" || option.id === "model",
+  )?.currentValue ?? asRecord(session.models)?.currentModelId;
+
+const hasOption = (options: unknown, value: string): boolean =>
+  records(options).some((option) => option.value === value || hasOption(option.options, value));
+
+export const selectAcpModel = async (
+  label: string,
+  session: RpcRecord,
+  model: string,
+  request: (method: string, params: RpcRecord) => Promise<RpcRecord>,
+): Promise<RpcRecord> => {
+  let selected = session;
+  const selector = records(session.configOptions).find(
+    (option) => option.category === "model" || option.id === "model",
+  );
+  if (
+    acpModel(session) !== model &&
+    typeof selector?.id === "string" &&
+    hasOption(selector.options, model)
+  ) {
+    const updated = await request("session/set_config_option", {
+      sessionId: session.sessionId,
+      configId: selector.id,
+      value: model,
+    });
+    selected = { ...session, ...updated };
+  }
+  const reported = acpModel(selected);
+  if (typeof reported === "string" && reported !== model)
+    throw new Error(
+      `${label} ACP selected '${reported}' instead of requested '${model}'; use an exact model ID advertised by the ACP session`,
+    );
+  return selected;
+};