@agent-native/core 0.111.4 → 0.113.0

package/corpus/templates/brain/actions/list-projects.ts

@@ -0,0 +1,51 @@
+import { defineAction } from "@agent-native/core";
+import { getRequestUserEmail } from "@agent-native/core/server/request-context";
+import { getUserSetting } from "@agent-native/core/settings";
+import { accessFilter } from "@agent-native/core/sharing";
+import { inArray } from "drizzle-orm";
+import { z } from "zod";
+
+import { getDb, schema } from "../server/db/index.js";
+
+const DEFAULT_PROJECT_SETTING = "brain-default-project";
+
+export default defineAction({
+  description: "List accessible Brain projects and their source membership.",
+  schema: z.object({}),
+  http: { method: "GET" },
+  readOnly: true,
+  publicAgent: { expose: true, readOnly: true, requiresAuth: true },
+  run: async () => {
+    const db = getDb();
+    const projects = await db
+      .select()
+      .from(schema.brainProjects)
+      .where(accessFilter(schema.brainProjects, schema.brainProjectShares));
+    const projectIds = projects.map((project) => project.id);
+    const links = projectIds.length
+      ? await db
+          .select()
+          .from(schema.brainProjectSources)
+          .where(inArray(schema.brainProjectSources.projectId, projectIds))
+      : [];
+    const userEmail = getRequestUserEmail();
+    const setting = userEmail
+      ? ((await getUserSetting(userEmail, DEFAULT_PROJECT_SETTING)) as {
+          projectId?: string;
+        } | null)
+      : null;
+    return {
+      projects: projects.map((project) => ({
+        id: project.id,
+        title: project.title,
+        description: project.description,
+        visibility: project.visibility,
+        sourceIds: links
+          .filter((link) => link.projectId === project.id)
+          .map((link) => link.sourceId),
+        isDefault: setting?.projectId === project.id,
+        updatedAt: project.updatedAt,
+      })),
+    };
+  },
+});

package/corpus/templates/brain/actions/list-proposals.ts

@@ -1,26 +1,102 @@
 import { defineAction } from "@agent-native/core";
-import { accessFilter } from "@agent-native/core/sharing";
-import { and, desc, eq } from "drizzle-orm";
+import { accessFilter, assertAccess } from "@agent-native/core/sharing";
+import { and, desc, eq, inArray, isNull, or, sql } from "drizzle-orm";
 import { z } from "zod";
 
 import { getDb, schema } from "../server/db/index.js";
-import { serializeProposal } from "../server/lib/brain.js";
+import { listAccessibleAudienceIds } from "../server/lib/audiences.js";
+import { parseJson, serializeProposal } from "../server/lib/brain.js";
 
 export default defineAction({
   description: "List Brain knowledge proposals requiring review.",
   schema: z.object({
-    status: z.enum(["pending", "approved", "rejected"]).default("pending"),
+    status: z
+      .enum(["pending", "approved", "rejected", "quarantine"])
+      .default("pending"),
     limit: z.coerce.number().int().min(1).max(100).default(50),
   }),
   http: { method: "GET" },
   readOnly: true,
   run: async ({ status, limit }) => {
+    if (status === "quarantine") {
+      const accessibleSources = await getDb()
+        .select({ id: schema.brainSources.id })
+        .from(schema.brainSources)
+        .where(accessFilter(schema.brainSources, schema.brainSourceShares));
+      const adminSourceIds = (
+        await Promise.all(
+          accessibleSources.map(async ({ id }) => {
+            try {
+              await assertAccess("brain-source", id, "admin");
+              return id;
+            } catch {
+              return null;
+            }
+          }),
+        )
+      ).filter((sourceId): sourceId is string => Boolean(sourceId));
+      if (!adminSourceIds.length) return { count: 0, proposals: [] };
+      const rows = await getDb()
+        .select({
+          id: schema.brainSensitivityEvents.id,
+          sourceId: schema.brainSensitivityEvents.sourceId,
+          sourceName: schema.brainSources.title,
+          disposition: schema.brainSensitivityEvents.disposition,
+          categoriesJson: schema.brainSensitivityEvents.categoriesJson,
+          confidenceBand: schema.brainSensitivityEvents.confidenceBand,
+          expiresAt: schema.brainSensitivityEvents.expiresAt,
+          createdAt: schema.brainSensitivityEvents.createdAt,
+          updatedAt: schema.brainSensitivityEvents.updatedAt,
+        })
+        .from(schema.brainSensitivityEvents)
+        .innerJoin(
+          schema.brainSources,
+          eq(schema.brainSources.id, schema.brainSensitivityEvents.sourceId),
+        )
+        .where(
+          and(
+            eq(schema.brainSensitivityEvents.disposition, "quarantined"),
+            inArray(schema.brainSensitivityEvents.sourceId, adminSourceIds),
+          ),
+        )
+        .orderBy(desc(schema.brainSensitivityEvents.createdAt))
+        .limit(limit);
+      const proposals = rows.map((row) => ({
+        id: row.id,
+        title: "Privacy quarantine",
+        sourceId: row.sourceId,
+        sourceName: row.sourceName,
+        reason: [
+          row.confidenceBand,
+          ...parseJson<string[]>(row.categoriesJson, []),
+          row.expiresAt ? `expires ${row.expiresAt}` : "metadata-only",
+        ].join(" · "),
+        status: "quarantine" as const,
+        createdAt: row.createdAt,
+        updatedAt: row.updatedAt,
+      }));
+      return { count: proposals.length, proposals };
+    }
+    const audienceIds = await listAccessibleAudienceIds();
     const rows = await getDb()
       .select()
       .from(schema.brainProposals)
       .where(
         and(
           accessFilter(schema.brainProposals, schema.brainProposalShares),
+          or(
+            isNull(schema.brainProposals.captureId),
+            audienceIds.length
+              ? and(
+                  sql`exists (
+                    select 1 from ${schema.brainSources}
+                    where ${schema.brainSources.id} = ${schema.brainProposals.sourceId}
+                      and ${accessFilter(schema.brainSources, schema.brainSourceShares)}
+                  )`,
+                  inArray(schema.brainProposals.audienceId, audienceIds),
+                )
+              : undefined,
+          ),
           eq(schema.brainProposals.status, status),
         ),
       )

package/corpus/templates/brain/actions/list-sources.ts

@@ -1,20 +1,34 @@
 import { defineAction } from "@agent-native/core";
 import { accessFilter } from "@agent-native/core/sharing";
-import { and, count, desc, eq, ne } from "drizzle-orm";
+import { and, countDistinct, desc, eq, inArray, ne } from "drizzle-orm";
 import { z } from "zod";
 
 import { getDb, schema } from "../server/db/index.js";
 import { nextBrainSourceSyncAt } from "../server/jobs/sync-sources.js";
+import { listAccessibleAudienceIds } from "../server/lib/audiences.js";
 import { parseJson, serializeSource } from "../server/lib/brain.js";
 import { sourceProviderSchema } from "./_schemas.js";
 
 async function sourceRecordCount(sourceId: string): Promise<number> {
+  const audienceIds = await listAccessibleAudienceIds([sourceId]);
+  if (!audienceIds.length) return 0;
+
   // Count only — never load the heavy `content` blob of every capture row
   // just to take `.length`.
   const [row] = await getDb()
-    .select({ value: count() })
+    .select({ value: countDistinct(schema.brainRawCaptures.id) })
     .from(schema.brainRawCaptures)
-    .where(eq(schema.brainRawCaptures.sourceId, sourceId));
+    .innerJoin(
+      schema.brainCaptureAudiences,
+      eq(schema.brainCaptureAudiences.captureId, schema.brainRawCaptures.id),
+    )
+    .where(
+      and(
+        eq(schema.brainRawCaptures.sourceId, sourceId),
+        eq(schema.brainRawCaptures.sensitivityDisposition, "allowed"),
+        inArray(schema.brainCaptureAudiences.audienceId, audienceIds),
+      ),
+    );
   return row?.value ?? 0;
 }
 

package/corpus/templates/brain/actions/list-staged-datasets.ts

@@ -1,40 +1,15 @@
 /**
  * Thin Brain re-export of staged dataset listing, pre-bound to appId="brain".
  */
-import { defineAction } from "@agent-native/core";
-import { listStagedDatasets } from "@agent-native/core/provider-api/staged-datasets-store";
-import { getCredentialContext } from "@agent-native/core/server/request-context";
+import { createListStagedDatasetsAction } from "@agent-native/core/provider-api/actions/staged-datasets";
 import { z } from "zod";
 
 import { BRAIN_APP_ID } from "../server/lib/provider-api.js";
 
-export default defineAction({
+export default createListStagedDatasetsAction({
   description:
     "List staged datasets stored by provider-api-request (stageAs) for the current user. Returns dataset ids, names, row counts, columns, and sizes. Use dataset ids with query-staged-dataset to aggregate, or with delete-staged-dataset to free scratch storage.",
   schema: z.object({}),
   http: { method: "GET" },
-  readOnly: true,
-  run: async () => {
-    const ctx = getCredentialContext();
-    if (!ctx) {
-      throw new Error("No authenticated context for list-staged-datasets.");
-    }
-
-    const datasets = await listStagedDatasets({
-      appId: BRAIN_APP_ID,
-      ownerEmail: ctx.userEmail,
-    });
-
-    return {
-      datasets: datasets.map((d) => ({
-        id: d.id,
-        name: d.name,
-        rowCount: d.rowCount,
-        columns: d.columns,
-        byteSize: d.byteSize,
-        updatedAt: new Date(d.updatedAt).toISOString(),
-      })),
-      total: datasets.length,
-    };
-  },
+  appId: BRAIN_APP_ID,
 });

package/corpus/templates/brain/actions/manage-project.ts

@@ -0,0 +1,118 @@
+import { defineAction } from "@agent-native/core";
+import {
+  getRequestOrgId,
+  getRequestUserEmail,
+} from "@agent-native/core/server/request-context";
+import { putUserSetting } from "@agent-native/core/settings";
+import { accessFilter, assertAccess } from "@agent-native/core/sharing";
+import { and, eq, inArray } from "drizzle-orm";
+import { z } from "zod";
+
+import { getDb, schema } from "../server/db/index.js";
+import { nanoid, nowIso } from "../server/lib/brain.js";
+
+const DEFAULT_PROJECT_SETTING = "brain-default-project";
+
+export default defineAction({
+  description:
+    "Create, update, remove, or select a Brain project that scopes related knowledge sources.",
+  schema: z.object({
+    operation: z.enum(["create", "update", "delete", "set-default"]),
+    projectId: z.string().optional(),
+    title: z.string().min(1).optional(),
+    description: z.string().optional(),
+    sourceIds: z.array(z.string().min(1)).optional(),
+  }),
+  run: async (args) => {
+    const userEmail = getRequestUserEmail();
+    if (!userEmail) throw new Error("Not authenticated");
+    const db = getDb();
+    const now = nowIso();
+
+    if (args.operation === "create") {
+      if (!args.title) throw new Error("title is required");
+      const id = args.projectId ?? nanoid();
+      await db.insert(schema.brainProjects).values({
+        id,
+        title: args.title,
+        description: args.description ?? "",
+        ownerEmail: userEmail,
+        orgId: getRequestOrgId() ?? null,
+        visibility: "private",
+        createdAt: now,
+        updatedAt: now,
+      });
+      await replaceProjectSources(id, args.sourceIds ?? []);
+      return { projectId: id, operation: args.operation };
+    }
+
+    if (!args.projectId) throw new Error("projectId is required");
+    if (args.operation === "set-default") {
+      await assertAccess("brain-project", args.projectId, "viewer");
+      await putUserSetting(userEmail, DEFAULT_PROJECT_SETTING, {
+        projectId: args.projectId,
+      });
+      return { projectId: args.projectId, default: true };
+    }
+
+    await assertAccess("brain-project", args.projectId, "editor");
+    if (args.operation === "delete") {
+      await db
+        .delete(schema.brainProjectSources)
+        .where(eq(schema.brainProjectSources.projectId, args.projectId));
+      await db
+        .delete(schema.brainProjectShares)
+        .where(eq(schema.brainProjectShares.resourceId, args.projectId));
+      await db
+        .delete(schema.brainProjects)
+        .where(eq(schema.brainProjects.id, args.projectId));
+      return { projectId: args.projectId, deleted: true };
+    }
+
+    await db
+      .update(schema.brainProjects)
+      .set({
+        ...(args.title ? { title: args.title } : {}),
+        ...(args.description !== undefined
+          ? { description: args.description }
+          : {}),
+        updatedAt: now,
+      })
+      .where(eq(schema.brainProjects.id, args.projectId));
+    if (args.sourceIds) {
+      await replaceProjectSources(args.projectId, args.sourceIds);
+    }
+    return { projectId: args.projectId, operation: args.operation };
+  },
+});
+
+async function replaceProjectSources(projectId: string, sourceIds: string[]) {
+  const db = getDb();
+  const unique = Array.from(new Set(sourceIds));
+  if (unique.length) {
+    const accessible = await db
+      .select({ id: schema.brainSources.id })
+      .from(schema.brainSources)
+      .where(
+        and(
+          accessFilter(schema.brainSources, schema.brainSourceShares),
+          inArray(schema.brainSources.id, unique),
+        ),
+      );
+    if (accessible.length !== unique.length) {
+      throw new Error("One or more Brain sources were not found");
+    }
+  }
+  await db
+    .delete(schema.brainProjectSources)
+    .where(eq(schema.brainProjectSources.projectId, projectId));
+  if (!unique.length) return;
+  await db.insert(schema.brainProjectSources).values(
+    unique.map((sourceId) => ({
+      id: nanoid(),
+      projectId,
+      sourceId,
+      createdAt: nowIso(),
+    })),
+  );
+}

package/corpus/templates/brain/actions/mark-capture-distilled.ts

@@ -1,5 +1,6 @@
 import { defineAction } from "@agent-native/core";
-import { and, eq, inArray } from "drizzle-orm";
+import { getDbExec } from "@agent-native/core/db";
+import { and, desc, eq } from "drizzle-orm";
 import { z } from "zod";
 
 import { getDb, schema } from "../server/db/index.js";
@@ -15,15 +16,76 @@ import {
 
 export default defineAction({
   description: "Mark a raw Brain capture as distilled or ignored.",
-  schema: z.object({
-    captureId: z.string().min(1),
-    status: z.enum(["distilled", "ignored"]).default("distilled"),
-  }),
-  run: async ({ captureId, status }) => {
+  schema: z
+    .object({
+      captureId: z.string().min(1),
+      status: z.enum(["distilled", "ignored"]).default("distilled"),
+      queueId: z.string().min(1).optional(),
+      claimToken: z.string().min(1).optional(),
+    })
+    .superRefine((value, ctx) => {
+      if (Boolean(value.queueId) !== Boolean(value.claimToken)) {
+        ctx.addIssue({
+          code: z.ZodIssueCode.custom,
+          message: "queueId and claimToken must be supplied together.",
+        });
+      }
+    }),
+  run: async ({ captureId, status, queueId, claimToken }) => {
     const access = await getAccessibleCapture(captureId);
     if (!access) throw new Error(`No access to capture ${captureId}`);
     const db = getDb();
     const now = nowIso();
+
+    if (queueId && claimToken) {
+      const result = await getDbExec().execute({
+        sql: `UPDATE brain_ingest_queue
+          SET status = ?, error = NULL, lease_token = NULL,
+              lease_expires_at = NULL, updated_at = ?
+          WHERE id = ? AND capture_id = ? AND operation = 'distill'
+            AND status = 'processing' AND lease_token = ?`,
+        args: ["done", now, queueId, captureId, claimToken],
+      });
+      if (result.rowsAffected === 0) {
+        throw new Error(
+          "The distillation claim is no longer active; do not complete a newer worker's queue item.",
+        );
+      }
+    } else {
+      const [processing] = await db
+        .select()
+        .from(schema.brainIngestQueue)
+        .where(
+          and(
+            eq(schema.brainIngestQueue.captureId, captureId),
+            eq(schema.brainIngestQueue.operation, "distill"),
+            eq(schema.brainIngestQueue.status, "processing"),
+          ),
+        )
+        .orderBy(desc(schema.brainIngestQueue.updatedAt))
+        .limit(1);
+      if (processing) {
+        throw new Error(
+          "This capture has an active distillation claim. Supply its queueId and claimToken to complete it.",
+        );
+      }
+      await db
+        .update(schema.brainIngestQueue)
+        .set({
+          status: "done",
+          error: null,
+          leaseToken: null,
+          leaseExpiresAt: null,
+          updatedAt: now,
+        })
+        .where(
+          and(
+            eq(schema.brainIngestQueue.captureId, captureId),
+            eq(schema.brainIngestQueue.operation, "distill"),
+            eq(schema.brainIngestQueue.status, "queued"),
+          ),
+        );
+    }
     await db
       .update(schema.brainRawCaptures)
       .set({
@@ -32,20 +94,6 @@ export default defineAction({
         updatedAt: now,
       })
       .where(eq(schema.brainRawCaptures.id, captureId));
-    await db
-      .update(schema.brainIngestQueue)
-      .set({
-        status: "done",
-        error: null,
-        updatedAt: now,
-      })
-      .where(
-        and(
-          eq(schema.brainIngestQueue.captureId, captureId),
-          eq(schema.brainIngestQueue.operation, "distill"),
-          inArray(schema.brainIngestQueue.status, ["queued", "processing"]),
-        ),
-      );
     const updated = await getAccessibleCapture(captureId);
     if (!updated) return { capture: null };
     const capture = serializeCapture(updated.capture);

package/corpus/templates/brain/actions/navigate.ts

@@ -28,6 +28,7 @@ export default defineAction({
     proposalId: z.string().optional(),
     extensionId: z.string().optional(),
     query: z.string().optional(),
+    projectId: z.string().optional(),
     type: z.enum(["knowledge", "capture", "source", "all"]).optional(),
     provider: z.string().optional(),
     status: z.string().optional(),

package/corpus/templates/brain/actions/provider-api-catalog.ts

@@ -1,14 +1,14 @@
-import { defineAction } from "@agent-native/core";
+import { createProviderApiCatalogAction } from "@agent-native/core/provider-api/actions/provider-api";
 import { z } from "zod";
 
 import {
   BRAIN_PROVIDER_API_IDS,
-  listProviderApiCatalog,
+  getBrainProviderApiRuntime,
 } from "../server/lib/provider-api.js";
 
 const ProviderSchema = z.enum(BRAIN_PROVIDER_API_IDS);
 
-export default defineAction({
+export default createProviderApiCatalogAction(getBrainProviderApiRuntime(), {
   description:
     "List raw HTTP API capabilities for configured/shared providers that Brain can use for ad hoc analysis. Use before provider-api-request when source readers or canned actions are too narrow. Returns base URLs, auth style, credential key names, docs/spec URLs, placeholders, and examples; never returns secret values.",
   schema: z.object({
@@ -17,12 +17,6 @@ export default defineAction({
     ),
   }),
   http: { method: "GET" },
-  readOnly: true,
-  run: async ({ provider }) => {
-    return {
-      providers: await listProviderApiCatalog(provider),
-      guidance:
-        "Brain source sync and retrieval actions are convenience readers, not capability limits. When a question needs an endpoint, object type, filter, pagination mode, or API version that a source action does not model, inspect docs/spec URLs here and call provider-api-request with the exact provider API method/path/query/body. For broad searches, joins, classification, or absence claims, stage or save the full bounded corpus and reduce it with query-staged-dataset or run-code.",
-    };
-  },
+  guidance:
+    "Brain source sync and retrieval actions are convenience readers, not capability limits. When a question needs an endpoint, object type, filter, pagination mode, or API version that a source action does not model, inspect docs/spec URLs here and call provider-api-request with the exact provider API method/path/query/body. For broad searches, joins, classification, or absence claims, stage or save the full bounded corpus and reduce it with query-staged-dataset or run-code.",
 });

package/corpus/templates/brain/actions/provider-api-docs.ts

@@ -1,9 +1,9 @@
-import { defineAction } from "@agent-native/core";
+import { createProviderApiDocsAction } from "@agent-native/core/provider-api/actions/provider-api";
 import { z } from "zod";
 
 import {
   BRAIN_PROVIDER_API_IDS,
-  fetchProviderApiDocs,
+  getBrainProviderApiRuntime,
 } from "../server/lib/provider-api.js";
 
 const ProviderSchema = z.enum(BRAIN_PROVIDER_API_IDS);
@@ -23,7 +23,7 @@ const WebContentSearchSchema = z.object({
   caseSensitive: BooleanFromQuerySchema.optional(),
 });
 
-export default defineAction({
+export default createProviderApiDocsAction(getBrainProviderApiRuntime(), {
   description:
     "Inspect provider API docs/spec metadata, or fetch a public provider docs/spec/changelog URL. Use this before arbitrary provider-api-request calls when the exact endpoint, filter operator, payload shape, pagination, or API version is uncertain.",
   schema: z.object({
@@ -69,6 +69,4 @@ export default defineAction({
     ),
   }),
   http: { method: "GET" },
-  readOnly: true,
-  run: async (args) => fetchProviderApiDocs(args),
 });

package/corpus/templates/brain/actions/provider-api-request.ts

@@ -1,12 +1,9 @@
-import { defineAction } from "@agent-native/core";
-import { stagingExecuteRequest } from "@agent-native/core/provider-api/staging";
-import { getCredentialContext } from "@agent-native/core/server/request-context";
+import { createProviderApiRequestAction } from "@agent-native/core/provider-api/actions/provider-api";
 import { z } from "zod";
 
 import {
   BRAIN_APP_ID,
   BRAIN_PROVIDER_API_IDS,
-  executeProviderApiRequest,
   getBrainProviderApiRuntime,
 } from "../server/lib/provider-api.js";
 
@@ -68,7 +65,7 @@ const PaginationSchema = z
   })
   .optional();
 
-export default defineAction({
+export default createProviderApiRequestAction(getBrainProviderApiRuntime(), {
   description:
     "Make an arbitrary authenticated HTTP request to a configured or shared provider API for Brain analysis. Use this as the flexible escape hatch when Brain source readers cannot express the needed endpoint, filters, pagination, payload, or API version. The request is constrained to the provider host, uses configured credentials automatically, blocks private/internal URLs, and redacts secrets from responses.",
   schema: z.object({
@@ -198,36 +195,5 @@ export default defineAction({
   }),
   http: false,
   toolCallable: false,
-  run: async (args) => {
-    if (args.stageAs) {
-      const ctx = getCredentialContext();
-      if (!ctx) {
-        throw new Error("No authenticated context for provider API staging.");
-      }
-      const providerRuntime = getBrainProviderApiRuntime();
-      return stagingExecuteRequest(
-        {
-          provider: args.provider,
-          method: args.method,
-          path: args.path,
-          query: args.query,
-          headers: args.headers,
-          body: args.body,
-          auth: args.auth,
-          connectionId: args.connectionId,
-          accountId: args.accountId,
-          timeoutMs: args.timeoutMs,
-          maxBytes: args.maxBytes,
-          stageAs: args.stageAs,
-          itemsPath: args.itemsPath,
-          pagination: args.pagination,
-        },
-        (reqArgs) => providerRuntime.executeRequest(reqArgs),
-        { appId: BRAIN_APP_ID, ownerEmail: ctx.userEmail },
-      );
-    }
-    return executeProviderApiRequest(
-      args as unknown as Parameters<typeof executeProviderApiRequest>[0],
-    );
-  },
+  appId: BRAIN_APP_ID,
 });

package/corpus/templates/brain/actions/query-staged-dataset.ts

@@ -1,13 +1,7 @@
 /**
  * Thin Brain re-export of the staged dataset query helper, pre-bound to appId="brain".
  */
-import { defineAction } from "@agent-native/core";
-import { runAggregateQuery } from "@agent-native/core/provider-api/staged-datasets-aggregate";
-import {
-  getStagedDatasetMeta,
-  getStagedDatasetRows,
-} from "@agent-native/core/provider-api/staged-datasets-store";
-import { getCredentialContext } from "@agent-native/core/server/request-context";
+import { createQueryStagedDatasetAction } from "@agent-native/core/provider-api/actions/staged-datasets";
 import { z } from "zod";
 
 import { BRAIN_APP_ID } from "../server/lib/provider-api.js";
@@ -40,7 +34,7 @@ const AggregateFieldSchema = z.object({
     .describe("Output column name. Default: {op}_{column}."),
 });
 
-export default defineAction({
+export default createQueryStagedDatasetAction({
   description:
     "Run a filter/aggregate/project query over a staged dataset previously written by provider-api-request (stageAs). Use after staging provider source records, messages, transcripts, documents, issues, events, or search results to count, group, filter, or project rows without re-fetching provider APIs.",
   schema: z.object({
@@ -86,44 +80,5 @@ export default defineAction({
       .describe("Maximum rows to return (default all, max 10000)."),
   }),
   http: false,
-  readOnly: true,
-  run: async (args) => {
-    const ctx = getCredentialContext();
-    if (!ctx) {
-      throw new Error("No authenticated context for query-staged-dataset.");
-    }
-
-    const meta = await getStagedDatasetMeta({
-      id: args.datasetId,
-      appId: BRAIN_APP_ID,
-      ownerEmail: ctx.userEmail,
-    });
-    if (!meta) {
-      throw new Error(
-        `Dataset ${args.datasetId} not found (or belongs to a different owner/app).`,
-      );
-    }
-
-    const rows = await getStagedDatasetRows({
-      id: args.datasetId,
-      appId: BRAIN_APP_ID,
-      ownerEmail: ctx.userEmail,
-    });
-
-    const result = runAggregateQuery(rows, {
-      where: args.where,
-      groupBy: args.groupBy,
-      aggregate: args.aggregate,
-      select: args.select,
-      orderBy: args.orderBy,
-      orderDir: args.orderDir,
-      limit: args.limit,
-    });
-
-    return {
-      dataset: { id: meta.id, name: meta.name, totalRows: meta.rowCount },
-      rowCount: result.length,
-      rows: result,
-    };
-  },
+  appId: BRAIN_APP_ID,
 });