jetty-mcp-server 1.0.0 → 1.0.3

package/src/tools.ts

@@ -8,6 +8,60 @@ function jsonResult(data: unknown) {
   };
 }
 
+const PROVIDER_KEYS = [
+  "OPENAI_API_KEY",
+  "ANTHROPIC_API_KEY",
+  "GEMINI_API_KEY",
+  "REPLICATE_API_TOKEN",
+];
+
+async function detectTrialEligibility(
+  client: JettyClient,
+  collection: string
+): Promise<{ useTrialKeys: boolean; trialInfo: Record<string, unknown> }> {
+  try {
+    const trialStatus = (await client.getTrialStatus(collection)) as Record<
+      string,
+      unknown
+    >;
+
+    if (!trialStatus.active) {
+      return {
+        useTrialKeys: false,
+        trialInfo: { trial_active: false },
+      };
+    }
+
+    const env = (await client.getCollectionEnvironment(collection)) as Record<
+      string,
+      unknown
+    >;
+    const envKeys = Object.keys(env);
+    const hasProviderKeys = PROVIDER_KEYS.some((k) => envKeys.includes(k));
+
+    if (hasProviderKeys) {
+      return {
+        useTrialKeys: false,
+        trialInfo: {
+          trial_active: true,
+          using_trial_keys: false,
+          reason: "Collection has its own provider keys",
+        },
+      };
+    }
+
+    return {
+      useTrialKeys: true,
+      trialInfo: {
+        trial_active: true,
+        using_trial_keys: true,
+      },
+    };
+  } catch {
+    return { useTrialKeys: false, trialInfo: {} };
+  }
+}
+
 export function registerTools(server: McpServer, client: JettyClient) {
   // ── Collections ──────────────────────────────────────────────
 
@@ -87,11 +141,31 @@ export function registerTools(server: McpServer, client: JettyClient) {
     }
   );
 
+  // ── Trial Keys ────────────────────────────────────────────────
+
+  server.tool(
+    "get-trial-status",
+    "Get trial key status for a collection",
+    { collection: z.string().describe("Collection name") },
+    async ({ collection }) => {
+      return jsonResult(await client.getTrialStatus(collection));
+    }
+  );
+
+  server.tool(
+    "activate-trial",
+    "Activate trial keys for a collection",
+    { collection: z.string().describe("Collection name") },
+    async ({ collection }) => {
+      return jsonResult(await client.activateTrial(collection));
+    }
+  );
+
   // ── Run Workflows ────────────────────────────────────────────
 
   server.tool(
     "run-workflow",
-    "Run a workflow asynchronously (returns immediately with workflow_id)",
+    "Run a workflow asynchronously (returns immediately with workflow_id). Auto-detects trial key eligibility.",
     {
       collection: z.string().describe("Collection name"),
       task: z.string().describe("Task name"),
@@ -101,19 +175,23 @@ export function registerTools(server: McpServer, client: JettyClient) {
         .describe("Input parameters for the workflow"),
     },
     async ({ collection, task, init_params }) => {
-      return jsonResult(
-        await client.runWorkflow(
-          collection,
-          task,
-          init_params as Record<string, unknown>
-        )
+      const { useTrialKeys, trialInfo } = await detectTrialEligibility(
+        client,
+        collection
+      );
+      const result = await client.runWorkflow(
+        collection,
+        task,
+        init_params as Record<string, unknown>,
+        useTrialKeys
       );
+      return jsonResult({ ...trialInfo, result });
     }
   );
 
   server.tool(
     "run-workflow-sync",
-    "Run a workflow synchronously (blocks until completion, may take 30-60s)",
+    "Run a workflow synchronously (blocks until completion, may take 30-60s). Auto-detects trial key eligibility.",
     {
       collection: z.string().describe("Collection name"),
       task: z.string().describe("Task name"),
@@ -123,13 +201,17 @@ export function registerTools(server: McpServer, client: JettyClient) {
         .describe("Input parameters for the workflow"),
     },
     async ({ collection, task, init_params }) => {
-      return jsonResult(
-        await client.runWorkflowSync(
-          collection,
-          task,
-          init_params as Record<string, unknown>
-        )
+      const { useTrialKeys, trialInfo } = await detectTrialEligibility(
+        client,
+        collection
       );
+      const result = await client.runWorkflowSync(
+        collection,
+        task,
+        init_params as Record<string, unknown>,
+        useTrialKeys
+      );
+      return jsonResult({ ...trialInfo, result });
     }
   );
 
@@ -219,4 +301,285 @@ export function registerTools(server: McpServer, client: JettyClient) {
       return jsonResult(await client.getStepTemplate(name));
     }
   );
+
+  // ── Environment Variables ─────────────────────────────────────
+
+  server.tool(
+    "check-secrets",
+    "Check which environment variables a collection has configured vs. what a runbook needs",
+    {
+      collection: z.string().describe("Collection name"),
+      required_keys: z
+        .array(z.string())
+        .describe("Environment variable names the runbook requires"),
+    },
+    async ({ collection, required_keys }) => {
+      const col = (await client.getCollectionEnvironment(
+        collection
+      )) as Record<string, unknown>;
+      const envVars = (col.environment_variables || {}) as Record<
+        string,
+        unknown
+      >;
+      const configured = Object.keys(envVars);
+      const missing = required_keys.filter((k) => !configured.includes(k));
+      return jsonResult({
+        configured,
+        missing,
+        ready: missing.length === 0,
+      });
+    }
+  );
+
+  server.tool(
+    "set-environment-vars",
+    "Set environment variables on a collection (merge semantics, pass null to delete a key)",
+    {
+      collection: z.string().describe("Collection name"),
+      variables: z
+        .record(z.string().nullable())
+        .describe("Key-value pairs to set (null to delete)"),
+    },
+    async ({ collection, variables }) => {
+      return jsonResult(
+        await client.setEnvironmentVars(
+          collection,
+          variables as Record<string, string>
+        )
+      );
+    }
+  );
+
+  // ── Routines (Scheduled Runs) ─────────────────────────────────
+
+  const cadenceSchema = z
+    .object({
+      type: z
+        .enum(["manual", "hourly", "daily", "weekdays", "weekly"])
+        .describe(
+          "Cadence type. 'manual' creates a saved invocation preset (no schedule); the others register a Temporal schedule."
+        ),
+      hour_utc: z
+        .number()
+        .int()
+        .min(0)
+        .max(23)
+        .optional()
+        .describe(
+          "Hour of day in UTC (0-23). Required for daily/weekdays/weekly."
+        ),
+      minute_utc: z
+        .number()
+        .int()
+        .min(0)
+        .max(59)
+        .optional()
+        .describe("Minute of the hour in UTC (0-59). Defaults to 0."),
+      day_of_week: z
+        .enum(["mon", "tue", "wed", "thu", "fri", "sat", "sun"])
+        .optional()
+        .describe("Day of week. Required for weekly cadence."),
+    })
+    .describe(
+      "Cadence config. Examples: {type:'hourly', minute_utc:15}, {type:'daily', hour_utc:9}, {type:'weekly', day_of_week:'mon', hour_utc:9}."
+    );
+
+  server.tool(
+    "list-routines",
+    "List scheduled routines for a collection (optionally scoped to a task)",
+    {
+      collection: z.string().describe("Collection name"),
+      task: z
+        .string()
+        .optional()
+        .describe("Task name (omit to list all routines in the collection)"),
+    },
+    async ({ collection, task }) => {
+      return jsonResult(await client.listRoutines(collection, task));
+    }
+  );
+
+  server.tool(
+    "get-routine",
+    "Get a single routine including resolved next_run_at from Temporal",
+    {
+      collection: z.string().describe("Collection name"),
+      task: z.string().describe("Task name"),
+      name: z.string().describe("Routine name (unique per collection+task)"),
+    },
+    async ({ collection, task, name }) => {
+      return jsonResult(await client.getRoutine(collection, task, name));
+    }
+  );
+
+  server.tool(
+    "create-routine",
+    "Create a scheduled routine for an existing task. init_params_overrides keys must be a subset of task.workflow.init_params.",
+    {
+      collection: z.string().describe("Collection name"),
+      task: z.string().describe("Task name"),
+      name: z.string().describe("Routine name (slug, unique per task)"),
+      cadence: cadenceSchema,
+      init_params_overrides: z
+        .record(z.unknown())
+        .optional()
+        .describe(
+          "Init param overrides. Keys must be a subset of task.workflow.init_params; unknown keys return 400."
+        ),
+      secret_params: z
+        .record(z.unknown())
+        .optional()
+        .describe(
+          "Secrets injected at fire time (encrypted at rest, never logged)."
+        ),
+      paused: z
+        .boolean()
+        .optional()
+        .describe("Create the routine in paused state"),
+      webhook_url: z
+        .string()
+        .optional()
+        .describe("Optional webhook to notify on each run"),
+      webhook_secret: z
+        .string()
+        .optional()
+        .describe("Optional shared secret for webhook signing"),
+    },
+    async ({
+      collection,
+      task,
+      name,
+      cadence,
+      init_params_overrides,
+      secret_params,
+      paused,
+      webhook_url,
+      webhook_secret,
+    }) => {
+      const body: Record<string, unknown> = { name, cadence };
+      if (init_params_overrides !== undefined)
+        body.init_params_overrides = init_params_overrides;
+      if (secret_params !== undefined) body.secret_params = secret_params;
+      if (paused !== undefined) body.paused = paused;
+      if (webhook_url !== undefined) body.webhook_url = webhook_url;
+      if (webhook_secret !== undefined) body.webhook_secret = webhook_secret;
+      return jsonResult(await client.createRoutine(collection, task, body));
+    }
+  );
+
+  server.tool(
+    "update-routine",
+    "Patch a routine — provide any subset of cadence/overrides/secrets/paused/webhook fields",
+    {
+      collection: z.string().describe("Collection name"),
+      task: z.string().describe("Task name"),
+      name: z.string().describe("Routine name"),
+      cadence: cadenceSchema.optional(),
+      init_params_overrides: z
+        .record(z.unknown())
+        .optional()
+        .describe("Replace init param overrides"),
+      secret_params: z
+        .record(z.unknown())
+        .optional()
+        .describe("Replace secret params"),
+      paused: z.boolean().optional().describe("Pause/resume the routine"),
+      webhook_url: z.string().optional().describe("Update webhook URL"),
+      webhook_secret: z
+        .string()
+        .optional()
+        .describe("Update webhook signing secret"),
+    },
+    async ({
+      collection,
+      task,
+      name,
+      cadence,
+      init_params_overrides,
+      secret_params,
+      paused,
+      webhook_url,
+      webhook_secret,
+    }) => {
+      const patch: Record<string, unknown> = {};
+      if (cadence !== undefined) patch.cadence = cadence;
+      if (init_params_overrides !== undefined)
+        patch.init_params_overrides = init_params_overrides;
+      if (secret_params !== undefined) patch.secret_params = secret_params;
+      if (paused !== undefined) patch.paused = paused;
+      if (webhook_url !== undefined) patch.webhook_url = webhook_url;
+      if (webhook_secret !== undefined) patch.webhook_secret = webhook_secret;
+      return jsonResult(
+        await client.updateRoutine(collection, task, name, patch)
+      );
+    }
+  );
+
+  server.tool(
+    "delete-routine",
+    "Delete a routine and unregister its Temporal schedule",
+    {
+      collection: z.string().describe("Collection name"),
+      task: z.string().describe("Task name"),
+      name: z.string().describe("Routine name"),
+    },
+    async ({ collection, task, name }) => {
+      return jsonResult(await client.deleteRoutine(collection, task, name));
+    }
+  );
+
+  server.tool(
+    "pause-routine",
+    "Pause a routine (stops future fires until resumed)",
+    {
+      collection: z.string().describe("Collection name"),
+      task: z.string().describe("Task name"),
+      name: z.string().describe("Routine name"),
+    },
+    async ({ collection, task, name }) => {
+      return jsonResult(await client.pauseRoutine(collection, task, name));
+    }
+  );
+
+  server.tool(
+    "resume-routine",
+    "Resume a paused routine",
+    {
+      collection: z.string().describe("Collection name"),
+      task: z.string().describe("Task name"),
+      name: z.string().describe("Routine name"),
+    },
+    async ({ collection, task, name }) => {
+      return jsonResult(await client.resumeRoutine(collection, task, name));
+    }
+  );
+
+  server.tool(
+    "run-routine-now",
+    "Fire a routine immediately, bypassing the schedule. Returns workflow_id (parity with run-workflow).",
+    {
+      collection: z.string().describe("Collection name"),
+      task: z.string().describe("Task name"),
+      name: z.string().describe("Routine name"),
+    },
+    async ({ collection, task, name }) => {
+      return jsonResult(await client.runRoutineNow(collection, task, name));
+    }
+  );
+
+  server.tool(
+    "list-routine-runs",
+    "List recent trajectories triggered by this routine (filtered by triggered_by_routine_id)",
+    {
+      collection: z.string().describe("Collection name"),
+      task: z.string().describe("Task name"),
+      name: z.string().describe("Routine name"),
+      limit: z.number().optional().describe("Max results"),
+    },
+    async ({ collection, task, name, limit }) => {
+      return jsonResult(
+        await client.listRoutineRuns(collection, task, name, limit)
+      );
+    }
+  );
 }

package/tests/auth.test.mjs

@@ -0,0 +1,129 @@
+// Auth/token-resolution tests for JettyClient.
+//
+// Regression for the "MCP can't authenticate" bug: the plugin's .mcp.json used
+// to pass JETTY_API_TOKEN="" (empty), which both supplied no token AND shadowed
+// any inherited one. The server now (a) treats empty/whitespace env as unset and
+// (b) falls back to ~/.config/jetty/token written by `jetty login`.
+//
+// Imports from ../dist (build first; `npm test` runs the pretest build hook).
+
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { mkdtempSync, mkdirSync, writeFileSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { JettyClient } from "../dist/client.js";
+
+// Snapshot env we mutate so tests don't leak into each other.
+function withEnv(overrides, fn) {
+  const saved = {};
+  for (const k of ["JETTY_API_TOKEN", "JETTY_API_URL", "HOME"]) {
+    saved[k] = process.env[k];
+    if (overrides[k] === undefined) delete process.env[k];
+    else process.env[k] = overrides[k];
+  }
+  try {
+    return fn();
+  } finally {
+    for (const k of Object.keys(saved)) {
+      if (saved[k] === undefined) delete process.env[k];
+      else process.env[k] = saved[k];
+    }
+  }
+}
+
+// Build a fake HOME containing ~/.config/jetty/token (or not, if token === null).
+function fakeHome(token) {
+  const home = mkdtempSync(join(tmpdir(), "jetty-auth-test-"));
+  if (token !== null) {
+    mkdirSync(join(home, ".config", "jetty"), { recursive: true });
+    writeFileSync(join(home, ".config", "jetty", "token"), token);
+  }
+  return home;
+}
+
+// requireToken() is private; exercise it through a no-network request and assert
+// on the thrown message (or the lack of an auth-shaped one).
+async function tokenError(client) {
+  const orig = globalThis.fetch;
+  globalThis.fetch = async () => {
+    throw new Error("__network_reached__");
+  };
+  try {
+    await client.listCollections();
+    return null;
+  } catch (e) {
+    return e.message;
+  } finally {
+    globalThis.fetch = orig;
+  }
+}
+
+test("empty JETTY_API_TOKEN falls back to the token file", async () => {
+  const home = fakeHome("mlc_filetoken");
+  try {
+    await withEnv({ JETTY_API_TOKEN: "", HOME: home }, async () => {
+      const client = new JettyClient();
+      // Token file present → no auth error; the only failure is the mocked network.
+      const msg = await tokenError(client);
+      assert.equal(msg, "__network_reached__");
+    });
+  } finally {
+    rmSync(home, { recursive: true, force: true });
+  }
+});
+
+test("whitespace-only JETTY_API_TOKEN does not count as set", async () => {
+  const home = fakeHome("mlc_filetoken");
+  try {
+    await withEnv({ JETTY_API_TOKEN: "   ", HOME: home }, async () => {
+      const client = new JettyClient();
+      assert.equal(await tokenError(client), "__network_reached__");
+    });
+  } finally {
+    rmSync(home, { recursive: true, force: true });
+  }
+});
+
+test("env token takes precedence over the file", async () => {
+  const home = fakeHome("mlc_filetoken");
+  try {
+    await withEnv({ JETTY_API_TOKEN: "mlc_envtoken", HOME: home }, async () => {
+      const calls = [];
+      const orig = globalThis.fetch;
+      globalThis.fetch = async (url, init) => {
+        calls.push(init?.headers?.Authorization);
+        return {
+          ok: true,
+          status: 200,
+          headers: new Headers({ "content-type": "application/json" }),
+          json: async () => ({}),
+          text: async () => "{}",
+        };
+      };
+      try {
+        const client = new JettyClient();
+        await client.listCollections();
+        assert.equal(calls[0], "Bearer mlc_envtoken");
+      } finally {
+        globalThis.fetch = orig;
+      }
+    });
+  } finally {
+    rmSync(home, { recursive: true, force: true });
+  }
+});
+
+test("no env and no file yields a helpful error mentioning jetty login", async () => {
+  const home = fakeHome(null); // no token file
+  try {
+    await withEnv({ JETTY_API_TOKEN: undefined, HOME: home }, async () => {
+      const client = new JettyClient();
+      const msg = await tokenError(client);
+      assert.match(msg, /jetty login/);
+      assert.notEqual(msg, "__network_reached__"); // failed before any fetch
+    });
+  } finally {
+    rmSync(home, { recursive: true, force: true });
+  }
+});