diff --git a/.gitignore b/.gitignore index f29e32849..2244e713d 100644 --- a/.gitignore +++ b/.gitignore @@ -25,3 +25,4 @@ internal .isaac/ .codex-tmp/ +__pycache__/ diff --git a/docs/docs/plugins/omnigent.md b/docs/docs/plugins/omnigent.md new file mode 100644 index 000000000..d1e4f7d66 --- /dev/null +++ b/docs/docs/plugins/omnigent.md @@ -0,0 +1,117 @@ +# Omnigent + + +:::warning Beta plugin +This plugin is currently **beta**. APIs may change between minor releases. Import from `@databricks/appkit/beta`. See [Plugin Stability Tiers](./stability.md). +::: + + +The `omnigent` plugin gives your app's users [Omnigent](https://docs.databricks.com/aws/en/omnigent/) agent sessions: Claude Agent SDK, Pi and OpenAI Agents SDK, on the workspace's managed Omnigent. + +- **Sessions belong to the signed-in user.** Every call carries the user's token, so a session is theirs, appears in `/omnigent`, and is private to them until they share it. +- **Each session runs on a Databricks Sandbox**, so agent-written code runs in an isolated, Databricks-managed box with the user's own permissions. +- **Model calls go through Unity Gateway**, with its governance, limits and usage tracking. +- **Tools are Databricks-managed MCP servers**, called as the user, so Unity Catalog permissions apply. + +Omnigent owns the agent loop: sessions, history, approvals. Use the [`agents`](./agents.md) plugin when you want AppKit to own the loop against one serving endpoint. Use `omnigent` when you want the harnesses themselves. + +## Install + +```ts +import { createApp, server } from "@databricks/appkit"; +import { omnigent } from "@databricks/appkit/beta"; + +await createApp({ + plugins: [ + server(), + omnigent({ + instructions: "You help the team explore order data.", + tools: [ + { type: "sql" }, + { type: "genie", spaceId: "01f0..." }, + { type: "functions", catalog: "main", schema: "orders" }, + ], + }), + ], +}); +``` + +## Workspace setup + +- **Previews.** Enable the **Omnigent** and **Databricks Sandbox** previews in the workspace. Omnigent needs a region that supports Unity Gateway; Sandbox hosts are in select AWS regions. +- **The `all-apis` scope on the app.** Managed Omnigent only accepts tokens with the `all-apis` scope, which an app's user authorization cannot request itself. An account admin adds it to the app's OAuth integration once, keeping the app's existing scopes: + + ```bash + databricks apps get -o json # oauth2_app_integration_id, effective_user_api_scopes + databricks account custom-app-integration update \ + --json '{"scopes": ["", "all-apis"]}' + ``` + + The update replaces the whole scope list. New scopes can take up to an hour to reach users' tokens. A bundle deploy, or a change to the app's user API scopes, sets the scopes back to the declared list, so repeat the update after either. + +`GET /api/omnigent/status` reports whether the user can reach managed Omnigent, and why not. + +## Tools + +`tools` lists Databricks-managed MCP servers every session can call: + +| Source | Tools | +|---|---| +| `{ type: "sql" }` | Databricks SQL: run SQL statements | +| `{ type: "genie", spaceId }` | One Genie space | +| `{ type: "functions", catalog, schema, function? }` | Unity Catalog functions, for business logic as tools | +| `{ type: "ai-search", catalog, schema, index? }` | AI Search over an index | +| `{ type: "service", service }` | A Unity Gateway MCP Service, such as `system.ai.web_search` or one you registered | + +Tool calls run inside the session's Sandbox with the user's own credential, so a user only reaches the data and services they can already use. The user needs the usual grants: `CAN USE` on the warehouse, access to the Genie space, `EXECUTE` on functions and MCP Services. + +Tools defined in app code with `defineTool()` are not available: a Sandbox's credential cannot call back into a Databricks App. Put business logic in Unity Catalog functions instead. + +## Modes + +- `auto` runs every tool. +- `ask` pauses on tools that change data (and on Omnigent's own OS tools) until the user approves. +- `read` denies tools that change data. + +Databricks SQL's `execute_sql` counts as a change; its `execute_sql_read_only` does not. List other tools that change data in `writeTools`. A pending approval appears on the session as `pending_elicitations`; resolve it with `POST /api/omnigent/sessions/:id/elicitations/:eid` and `{"action": "accept" | "decline" | "cancel"}`. + +## Configuration + +| Option | Default | | +|---|---|---| +| `instructions` | a generic assistant prompt | System instructions for every session | +| `tools` | none | Databricks-managed MCP servers, see [Tools](#tools) | +| `writeTools` | none | Further tool names that change data | +| `defaultMode` | `ask` | `auto`, `ask` or `read` | +| `harnesses` | `claude-sdk`, `pi`, `openai-agents` | Harnesses users may pick | +| `defaultHarness`, `defaultModel` | first available | | + +Each harness is offered the Unity Gateway `system.ai` models whose API it speaks: Claude Agent SDK the Anthropic Messages API, Pi the Codex Responses API, OpenAI Agents SDK the OpenAI Responses API. + +## Routes + +All routes are under `/api/omnigent` and act as the signed-in user. Sessions are limited to the ones this app created. + +| Route | | +|---|---| +| `GET /status` | Whether managed Omnigent and Sandbox hosts are usable | +| `GET /harnesses` | Harnesses available to the user, each with the models it can drive | +| `GET /sessions`, `POST /sessions` | List, create (`harness`, `model`, `mode`, `message`, `title`) | +| `GET /sessions/:id`, `DELETE /sessions/:id` | Snapshot (items, status, mode, pending approvals), delete | +| `POST /sessions/:id/messages` | `{ "text": "..." }` | +| `GET /sessions/:id/stream` | Server-sent events (Omnigent's session stream, which replays from `Last-Event-ID`) | +| `POST /sessions/:id/interrupt` | Stop the current turn | +| `PUT /sessions/:id/mode` | `{ "mode": "ask" }` | +| `POST /sessions/:id/elicitations/:eid` | Resolve a pending approval | + +Creating a session returns once its Sandbox host is up, about 30 seconds on a cold start. + +## Local development + +With `NODE_ENV=development` and no forwarded token, the plugin calls Omnigent as the developer, with the token of the app's local Databricks CLI profile. + +## Limits + +- **Beta, in Omnigent's regions.** The plugin needs the workspace's managed Omnigent; it does not run Omnigent itself. +- **Only gateway harnesses.** Harnesses that sign in with their vendor (Cursor, Copilot, Gemini and others) are not offered. Codex is not offered yet: on a Sandbox host, its brokered sign-in rejects the Codex CLI version the Sandbox ships. +- **Broad app scope.** The `all-apis` scope lets the app call any API the user can, on the user's behalf. diff --git a/packages/appkit/src/plugins/beta-exports.generated.ts b/packages/appkit/src/plugins/beta-exports.generated.ts index 7e556ebd9..39f5b39d3 100644 --- a/packages/appkit/src/plugins/beta-exports.generated.ts +++ b/packages/appkit/src/plugins/beta-exports.generated.ts @@ -7,3 +7,4 @@ export { agents } from "./agents"; export { aiSearch } from "./ai-search"; +export { omnigent } from "./omnigent"; diff --git a/packages/appkit/src/plugins/omnigent/bundle.ts b/packages/appkit/src/plugins/omnigent/bundle.ts new file mode 100644 index 000000000..d503c838e --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/bundle.ts @@ -0,0 +1,92 @@ +import { gzipSync } from "node:zlib"; + +/** A stdio MCP server the session's host starts. */ +export interface McpServerSpec { + name: string; + description?: string; + transport: "stdio"; + command: string; + args: string[]; + /** Seconds. */ + timeout?: number; +} + +interface BundleSpec { + name: string; + description?: string; + harness: string; + model: string; + instructions: string; + mcpServers?: McpServerSpec[]; +} + +/** + * An Omnigent agent bundle: `config.yaml` (the agent spec), `AGENTS.md` (the + * instructions) and one `tools/mcp/.yaml` per MCP server, as a tar.gz. + * The YAML files are written as JSON, which is valid YAML. + */ +export function buildBundle(spec: BundleSpec): Buffer { + const config = { + spec_version: 1, + name: spec.name, + description: spec.description ?? spec.name, + executor: { + type: "omnigent", + model: spec.model, + config: { harness: spec.harness }, + }, + instructions: "AGENTS.md", + }; + const files: Record = { + "config.yaml": `${JSON.stringify(config, null, 2)}\n`, + "AGENTS.md": spec.instructions, + }; + for (const s of spec.mcpServers ?? []) { + if (!/^[a-z0-9][a-z0-9_-]*$/i.test(s.name)) { + throw new Error(`Invalid MCP server name: ${s.name}`); + } + if (files[`tools/mcp/${s.name}.yaml`]) { + throw new Error(`Duplicate MCP server name: ${s.name}`); + } + files[`tools/mcp/${s.name}.yaml`] = `${JSON.stringify( + { + name: s.name, + description: s.description ?? s.name, + transport: s.transport, + command: s.command, + args: s.args, + timeout: s.timeout ?? 300, + }, + null, + 2, + )}\n`; + } + return gzipSync(tar(files)); +} + +/** Minimal ustar writer: regular files only, names under 100 bytes. */ +function tar(files: Record): Buffer { + const blocks: Buffer[] = []; + const mtime = Math.floor(Date.now() / 1000); + for (const [name, content] of Object.entries(files)) { + if (Buffer.byteLength(name) >= 100) + throw new Error(`tar: name too long: ${name}`); + const data = Buffer.isBuffer(content) ? content : Buffer.from(content); + const header = Buffer.alloc(512); + header.write(name, 0, 100); + header.write("0000644\0", 100); + header.write("0000000\0", 108); + header.write("0000000\0", 116); + header.write(`${data.length.toString(8).padStart(11, "0")}\0`, 124); + header.write(`${mtime.toString(8).padStart(11, "0")}\0`, 136); + header.write(" ", 148); // checksum placeholder + header.write("0", 156); + header.write("ustar\u000000", 257); + let sum = 0; + for (const b of header) sum += b; + header.write(`${sum.toString(8).padStart(6, "0")}\0 `, 148); + blocks.push(header, data, Buffer.alloc((512 - (data.length % 512)) % 512)); + } + blocks.push(Buffer.alloc(1024)); + return Buffer.concat(blocks); +} diff --git a/packages/appkit/src/plugins/omnigent/client.ts b/packages/appkit/src/plugins/omnigent/client.ts new file mode 100644 index 000000000..56be2cf26 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/client.ts @@ -0,0 +1,57 @@ +/** + * Calls the workspace's managed Omnigent server (`/api/2.0/omnigent`) with + * the signed-in user's token, so sessions, Sandbox hosts and model calls are + * the user's own. + */ +export class ManagedOmnigent { + constructor(private readonly host: () => Promise) {} + + async request( + token: string, + method: string, + path: string, + init: { + json?: unknown; + body?: RequestInit["body"]; + headers?: Record; + signal?: AbortSignal; + } = {}, + ): Promise { + const r = await fetch(`${await this.host()}/api/2.0/omnigent${path}`, { + method, + headers: { + Authorization: `Bearer ${token}`, + ...(init.json !== undefined + ? { "content-type": "application/json" } + : {}), + ...init.headers, + }, + body: init.json !== undefined ? JSON.stringify(init.json) : init.body, + signal: init.signal, + redirect: "manual", + }); + if (r.status === 403) { + const text = await r.clone().text(); + if (text.includes("required scopes: all-apis")) { + throw new OmnigentScopeError(); + } + } + return r; + } +} + +/** + * Managed Omnigent only accepts tokens with the `all-apis` scope, which a + * Databricks App's user authorization cannot request itself. + */ +export class OmnigentScopeError extends Error { + readonly status = 403; + constructor() { + super( + "Omnigent needs this app's user token to carry the all-apis scope. " + + "An account admin adds it to the app's OAuth integration: " + + "databricks account custom-app-integration update " + + "with the app's current scopes plus all-apis. It can take up to an hour to reach users' tokens.", + ); + } +} diff --git a/packages/appkit/src/plugins/omnigent/harnesses.ts b/packages/appkit/src/plugins/omnigent/harnesses.ts new file mode 100644 index 000000000..7a6dc673d --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/harnesses.ts @@ -0,0 +1,78 @@ +interface HarnessInfo { + label: string; + /** + * Unity Gateway API types it calls (a UC model service's + * `supported_api_types`); a model must support one of them. + */ + apis: string[]; + /** Default models, best first, when the workspace has them. */ + preferred: string[]; +} + +/** + * Harnesses whose model calls a Databricks Sandbox routes through Unity + * Gateway. The other harnesses Omnigent supports sign in with their vendor, + * so they are not offered. Codex is left out for now: on a Sandbox host its + * brokered sign-in rejects the Codex CLI version the Sandbox ships. + */ +const GATEWAY_HARNESSES: Readonly> = { + "claude-sdk": { + label: "Claude Agent SDK", + apis: ["anthropic/v1/messages"], + preferred: ["system.ai.claude-sonnet-5", "system.ai.claude-opus-5-5"], + }, + pi: { + label: "Pi", + // Pi reaches Databricks models through the Codex responses API; on the + // chat-completions path Gemini fails with an empty 400 (tested). + apis: ["codex/v1/responses"], + preferred: ["system.ai.gpt-5-5"], + }, + "openai-agents": { + label: "OpenAI Agents SDK", + apis: ["openai/v1/responses"], + preferred: ["system.ai.gpt-5-5"], + }, +}; + +export const DEFAULT_HARNESSES = Object.keys(GATEWAY_HARNESSES); + +/** A Unity Gateway model, e.g. `system.ai.claude-sonnet-5`. */ +export interface GatewayModel { + name: string; + /** The model service's `supported_api_types`. */ + apiTypes: string[]; +} + +/** A harness the user can pick, with the models it can drive. */ +export interface HarnessOffer { + id: string; + label: string; + models: string[]; + defaultModel?: string; +} + +/** + * What a user can pick: the allowed gateway harnesses, each with the models + * whose API it speaks. + */ +export function offerHarnesses(opts: { + allowed: string[]; + models: GatewayModel[]; +}): HarnessOffer[] { + const out: HarnessOffer[] = []; + for (const id of opts.allowed) { + const info = GATEWAY_HARNESSES[id]; + if (!info) continue; + const models = opts.models + .filter((m) => m.apiTypes.some((t) => info.apis.includes(t))) + .map((m) => m.name); + out.push({ + id, + label: info.label, + models, + defaultModel: info.preferred.find((m) => models.includes(m)) ?? models[0], + }); + } + return out; +} diff --git a/packages/appkit/src/plugins/omnigent/index.ts b/packages/appkit/src/plugins/omnigent/index.ts new file mode 100644 index 000000000..24ab26045 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/index.ts @@ -0,0 +1,2 @@ +export * from "./omnigent"; +export * from "./types"; diff --git a/packages/appkit/src/plugins/omnigent/manifest.json b/packages/appkit/src/plugins/omnigent/manifest.json new file mode 100644 index 000000000..be2c79b9e --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/manifest.json @@ -0,0 +1,63 @@ +{ + "$schema": "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", + "name": "omnigent", + "displayName": "Omnigent Plugin", + "description": "Omnigent agent sessions for the app's users on the workspace's managed Omnigent: Claude Agent SDK, Pi and OpenAI Agents SDK on Databricks Sandboxes, with models through Unity Gateway", + "stability": "beta", + "resources": { + "required": [], + "optional": [] + }, + "scaffolding": { + "rules": { + "must": [ + "Enable the Omnigent and Databricks Sandbox previews in the workspace", + "Account admin: add all-apis to the app's OAuth integration (databricks account custom-app-integration update)", + "Re-run that update after a bundle deploy or a user_api_scopes change, which both drop all-apis" + ] + } + }, + "config": { + "schema": { + "type": "object", + "properties": { + "instructions": { + "type": "string", + "description": "System instructions for every session" + }, + "tools": { + "type": "array", + "items": { + "type": "object" + }, + "description": "Databricks-managed MCP servers sessions can call, as the user" + }, + "writeTools": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Further tool names that change data" + }, + "defaultMode": { + "type": "string", + "enum": ["auto", "ask", "read"], + "default": "ask" + }, + "harnesses": { + "type": "array", + "items": { + "type": "string" + }, + "description": "Harnesses users may pick" + }, + "defaultHarness": { + "type": "string" + }, + "defaultModel": { + "type": "string" + } + } + } + } +} diff --git a/packages/appkit/src/plugins/omnigent/omnigent.ts b/packages/appkit/src/plugins/omnigent/omnigent.ts new file mode 100644 index 000000000..993e0de1a --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/omnigent.ts @@ -0,0 +1,603 @@ +import { Readable } from "node:stream"; + +import type express from "express"; +import type { IAppRouter } from "shared"; + +import { ServiceContext } from "../../context"; +import { createLogger } from "../../logging"; +import { Plugin, toPlugin } from "../../plugin"; +import { defineManifest } from "../../registry"; +import { buildBundle } from "./bundle"; +import { ManagedOmnigent, OmnigentScopeError } from "./client"; +import { + DEFAULT_HARNESSES, + type GatewayModel, + offerHarnesses, + type HarnessOffer, +} from "./harnesses"; +import manifest from "./manifest.json"; +import { + MODE_POLICY_PREFIX, + modeFromPolicies, + modePolicies, + OMNIGENT_MODES, + type OmnigentMode, +} from "./policies"; +import { toolServer, writeToolNames } from "./tools"; +import type { + CreateSessionRequest, + IOmnigentConfig, + OmnigentStatus, +} from "./types"; + +const logger = createLogger("omnigent"); + +const MODELS_TTL_MS = 5 * 60_000; +const ELICITATION_ACTIONS = new Set(["accept", "decline", "cancel"]); + +/** + * Omnigent sessions for the app's users, on the workspace's managed + * Omnigent server. + * + * Every call carries the signed-in user's token, so each session belongs to + * that user, runs on a Databricks Sandbox host, and calls models through + * Unity Gateway. The plugin builds the agent (instructions, harness, model) + * and relays the conversation; Omnigent owns the session and its history, + * which also appear in `/omnigent`. + */ +export class OmnigentPlugin extends Plugin { + static manifest = defineManifest<"omnigent">(manifest); + + protected static description = + "Omnigent sessions (Claude Agent SDK, Pi, OpenAI Agents SDK) on Databricks Sandboxes"; + declare protected config: IOmnigentConfig; + + private readonly omnigent = new ManagedOmnigent(() => this.workspaceHost()); + private models?: { at: number; list: GatewayModel[] }; + + // ---------------------------------------------------------------- helpers + + private async workspaceHost(): Promise { + return (await ServiceContext.get().client.getHost()).replace(/\/+$/, ""); + } + + /** Sessions this app created carry its name as a label and agent name. */ + private appName(): string { + return (process.env.DATABRICKS_APP_NAME || "appkit") + .toLowerCase() + .replace(/[^a-z0-9-]/g, "-"); + } + + /** + * The signed-in user's token. In local development there is no forwarded + * token, so the developer's own CLI identity is used. + */ + private async userToken(req: express.Request): Promise { + const forwarded = req.header("x-forwarded-access-token")?.trim(); + if (forwarded) return forwarded; + if (process.env.NODE_ENV !== "development") return null; + const headers = new Headers(); + await ServiceContext.get().client.authenticate(headers); + return headers.get("Authorization")?.replace(/^Bearer /, "") ?? null; + } + + /** + * Chat models the user can call through Unity Gateway: the UC `system.ai` + * model services, with the API types each one supports. + */ + private async servedModels(token: string): Promise { + if (this.models && Date.now() - this.models.at < MODELS_TTL_MS) + return this.models.list; + const host = await this.workspaceHost(); + const list: GatewayModel[] = []; + let page: string | undefined; + do { + const q = new URLSearchParams({ + parent: "schemas/system.ai", + max_results: "100", + }); + if (page) q.set("page_token", page); + const r = await fetch( + `${host}/api/2.1/unity-catalog/model-services?${q}`, + { headers: { Authorization: `Bearer ${token}` } }, + ); + if (!r.ok) { + logger.warn("could not list UC model services: %d", r.status); + break; + } + const body = (await r.json()) as { + model_services?: Array<{ + name: string; + supported_api_types?: string[]; + }>; + next_page_token?: string; + }; + for (const svc of body.model_services ?? []) { + const apiTypes = svc.supported_api_types ?? []; + if ( + !apiTypes.some((t) => /chat\/completions|responses|messages/.test(t)) + ) + continue; + list.push({ + name: svc.name.replace(/^model-services\//, ""), + apiTypes, + }); + } + page = body.next_page_token || undefined; + } while (page); + list.sort((a, b) => a.name.localeCompare(b.name)); + if (list.length) this.models = { at: Date.now(), list }; + return list; + } + + private async offers(token: string): Promise { + return offerHarnesses({ + allowed: this.config.harnesses ?? DEFAULT_HARNESSES, + models: await this.servedModels(token), + }); + } + + private instructions(email: string | undefined): string { + const base = + this.config.instructions ?? + "You are an assistant inside this Databricks app."; + return email ? `${base}\n\nThe signed-in user is ${email}.` : base; + } + + private bundle( + host: string, + email: string | undefined, + harness: string, + model: string, + ): Buffer { + return buildBundle({ + name: this.appName(), + description: "AppKit app agent", + harness, + model, + instructions: this.instructions(email), + mcpServers: (this.config.tools ?? []).map((t) => toolServer(t, host)), + }); + } + + private async setMode( + token: string, + sid: string, + mode: OmnigentMode, + ): Promise { + const have = (await ( + await this.omnigent.request(token, "GET", `/v1/sessions/${sid}/policies`) + ).json()) as { data?: Array<{ id: string; name?: string }> }; + for (const p of have.data ?? []) { + if (String(p.name ?? "").startsWith(MODE_POLICY_PREFIX)) { + await this.omnigent.request( + token, + "DELETE", + `/v1/sessions/${sid}/policies/${p.id}`, + ); + } + } + const writes = writeToolNames( + this.config.tools ?? [], + this.config.writeTools, + ); + for (const p of modePolicies(mode, writes)) { + const r = await this.omnigent.request( + token, + "POST", + `/v1/sessions/${sid}/policies`, + { json: p }, + ); + if (!r.ok) + throw new Error( + `Could not set ${mode} mode: ${r.status} ${await r.text()}`, + ); + } + } + + /** The session, or a 404 unless it belongs to this app (Omnigent already scopes it to the user). */ + private async owned( + token: string, + sid: string, + res: express.Response, + items = false, + ): Promise | null> { + const r = await this.omnigent.request( + token, + "GET", + `/v1/sessions/${encodeURIComponent(sid)}${items ? "?include_items=true" : ""}`, + ); + const snap = r.ok ? ((await r.json()) as Record) : null; + const labels = (snap?.labels ?? {}) as Record; + if (!snap || labels.app !== this.appName()) { + res.status(404).json({ error: "Session not found" }); + return null; + } + return snap; + } + + private static async relay( + res: express.Response, + r: Response, + ): Promise { + res.status(r.status); + const ct = r.headers.get("content-type"); + if (ct) res.setHeader("content-type", ct); + res.send(Buffer.from(await r.arrayBuffer())); + } + + private handler( + fn: ( + req: express.Request, + res: express.Response, + token: string, + ) => Promise, + ) { + return async (req: express.Request, res: express.Response) => { + let token: string | null; + try { + token = await this.userToken(req); + } catch { + token = null; + } + if (!token) { + res.status(401).json({ error: "Sign in to the app to use its agent." }); + return; + } + try { + await fn(req, res, token); + } catch (err) { + if (err instanceof OmnigentScopeError) { + res.status(err.status).json({ error: err.message }); + return; + } + logger.error("omnigent route failed: %O", err); + if (!res.headersSent) { + res + .status(502) + .json({ error: err instanceof Error ? err.message : String(err) }); + } + } + }; + } + + // ----------------------------------------------------------------- routes + + injectRoutes(router: IAppRouter) { + this.route(router, { + name: "status", + method: "get", + path: "/status", + handler: this.handler(async (_req, res, token) => { + let status: OmnigentStatus; + try { + const r = await this.omnigent.request(token, "GET", "/v1/info"); + const info = r.ok + ? ((await r.json()) as { + server_version?: string; + managed_sandboxes_enabled?: boolean; + }) + : null; + status = info + ? { + ready: Boolean(info.managed_sandboxes_enabled), + version: info.server_version, + sandboxes: Boolean(info.managed_sandboxes_enabled), + ...(info.managed_sandboxes_enabled + ? {} + : { + error: + "Databricks Sandbox is not enabled in this workspace.", + }), + } + : { + ready: false, + error: `Omnigent is not available in this workspace (HTTP ${r.status}). Enable the Omnigent preview.`, + }; + } catch (err) { + status = { + ready: false, + error: err instanceof Error ? err.message : String(err), + }; + } + res.json(status); + }), + }); + + this.route(router, { + name: "harnesses", + method: "get", + path: "/harnesses", + handler: this.handler(async (_req, res, token) => { + const harnesses = await this.offers(token); + res.json({ + harnesses, + modes: OMNIGENT_MODES, + defaultMode: this.config.defaultMode ?? "ask", + defaultHarness: + harnesses.find((h) => h.id === this.config.defaultHarness)?.id ?? + harnesses[0]?.id, + }); + }), + }); + + this.route(router, { + name: "listSessions", + method: "get", + path: "/sessions", + handler: this.handler(async (_req, res, token) => { + const q = new URLSearchParams({ + agent_name: this.appName(), + limit: "100", + }); + await OmnigentPlugin.relay( + res, + await this.omnigent.request(token, "GET", `/v1/sessions?${q}`), + ); + }), + }); + + this.route(router, { + name: "createSession", + method: "post", + path: "/sessions", + handler: this.handler(async (req, res, token) => { + const body = (req.body ?? {}) as CreateSessionRequest; + const offers = await this.offers(token); + const harness = + offers.find( + (h) => h.id === (body.harness ?? this.config.defaultHarness), + ) ?? offers[0]; + if (!harness) { + res + .status(409) + .json({ error: "No harness is available on this app." }); + return; + } + if (body.harness && body.harness !== harness.id) { + res + .status(400) + .json({ error: `Harness not available: ${body.harness}` }); + return; + } + const model = + body.model ?? this.config.defaultModel ?? harness.defaultModel; + if (!model || !harness.models.includes(model)) { + res.status(400).json({ + error: `Model not available for ${harness.id}: ${model ?? "(none)"}`, + }); + return; + } + const mode = body.mode ?? this.config.defaultMode ?? "ask"; + if (!OMNIGENT_MODES.includes(mode)) { + res.status(400).json({ error: `Unknown mode: ${mode}` }); + return; + } + const host = await this.workspaceHost(); + const form = new FormData(); + form.set( + "metadata", + JSON.stringify({ + title: + body.title ?? + body.message?.split("\n")[0]?.slice(0, 80) ?? + "New session", + labels: { app: this.appName() }, + host_type: "managed", + }), + ); + form.set( + "bundle", + new Blob( + [ + new Uint8Array( + this.bundle( + host, + req.header("x-forwarded-email"), + harness.id, + model, + ), + ), + ], + { type: "application/gzip" }, + ), + "bundle.tar.gz", + ); + // Returns once the Sandbox host is up (about 30 s on a cold start). + const created = await this.omnigent.request( + token, + "POST", + "/v1/sessions", + { body: form }, + ); + if (!created.ok) { + await OmnigentPlugin.relay(res, created); + return; + } + const { session_id: sid } = (await created.json()) as { + session_id: string; + }; + await this.setMode(token, sid, mode); + if (body.message) { + const sent = await this.omnigent.request( + token, + "POST", + `/v1/sessions/${sid}/events`, + { json: userMessage(body.message) }, + ); + if (!sent.ok) { + await OmnigentPlugin.relay(res, sent); + return; + } + } + res + .status(201) + .json({ session_id: sid, harness: harness.id, model, mode }); + }), + }); + + this.route(router, { + name: "getSession", + method: "get", + path: "/sessions/:id", + handler: this.handler(async (req, res, token) => { + const snap = await this.owned(token, req.params.id, res, true); + if (!snap) return; + const pol = (await ( + await this.omnigent.request( + token, + "GET", + `/v1/sessions/${encodeURIComponent(req.params.id)}/policies`, + ) + ).json()) as { data?: Array<{ name?: string }> }; + res.json({ + ...snap, + mode: modeFromPolicies((pol.data ?? []).map((p) => String(p.name))), + }); + }), + }); + + this.route(router, { + name: "deleteSession", + method: "delete", + path: "/sessions/:id", + handler: this.handler(async (req, res, token) => { + if (!(await this.owned(token, req.params.id, res))) return; + await OmnigentPlugin.relay( + res, + await this.omnigent.request( + token, + "DELETE", + `/v1/sessions/${encodeURIComponent(req.params.id)}`, + ), + ); + }), + }); + + this.route(router, { + name: "sendMessage", + method: "post", + path: "/sessions/:id/messages", + handler: this.handler(async (req, res, token) => { + const text = (req.body as { text?: string })?.text; + if (!text) { + res.status(400).json({ error: "text is required" }); + return; + } + if (!(await this.owned(token, req.params.id, res))) return; + await OmnigentPlugin.relay( + res, + await this.omnigent.request( + token, + "POST", + `/v1/sessions/${encodeURIComponent(req.params.id)}/events`, + { json: userMessage(text) }, + ), + ); + }), + }); + + this.route(router, { + name: "interrupt", + method: "post", + path: "/sessions/:id/interrupt", + handler: this.handler(async (req, res, token) => { + if (!(await this.owned(token, req.params.id, res))) return; + await OmnigentPlugin.relay( + res, + await this.omnigent.request( + token, + "POST", + `/v1/sessions/${encodeURIComponent(req.params.id)}/events`, + { json: { type: "interrupt", data: {} } }, + ), + ); + }), + }); + + this.route(router, { + name: "setMode", + method: "put", + path: "/sessions/:id/mode", + handler: this.handler(async (req, res, token) => { + const mode = (req.body as { mode?: OmnigentMode })?.mode; + if (!mode || !OMNIGENT_MODES.includes(mode)) { + res.status(400).json({ error: "mode must be auto, ask or read" }); + return; + } + if (!(await this.owned(token, req.params.id, res))) return; + await this.setMode(token, encodeURIComponent(req.params.id), mode); + res.json({ mode }); + }), + }); + + this.route(router, { + name: "resolveElicitation", + method: "post", + path: "/sessions/:id/elicitations/:eid", + handler: this.handler(async (req, res, token) => { + const action = (req.body as { action?: string })?.action ?? ""; + if (!ELICITATION_ACTIONS.has(action)) { + res + .status(400) + .json({ error: "action must be accept, decline or cancel" }); + return; + } + if (!(await this.owned(token, req.params.id, res))) return; + await OmnigentPlugin.relay( + res, + await this.omnigent.request( + token, + "POST", + `/v1/sessions/${encodeURIComponent(req.params.id)}/elicitations/${encodeURIComponent(req.params.eid)}/resolve`, + { json: { action } }, + ), + ); + }), + }); + + this.route(router, { + name: "stream", + method: "get", + path: "/sessions/:id/stream", + handler: this.handler(async (req, res, token) => { + if (!(await this.owned(token, req.params.id, res))) return; + const ac = new AbortController(); + req.on("close", () => ac.abort()); + const headers: Record = { Accept: "text/event-stream" }; + const last = req.header("last-event-id"); + if (last) headers["Last-Event-ID"] = last; + const up = await this.omnigent.request( + token, + "GET", + `/v1/sessions/${encodeURIComponent(req.params.id)}/stream`, + { headers, signal: ac.signal }, + ); + if (!up.ok || !up.body) { + await OmnigentPlugin.relay(res, up); + return; + } + res.status(200); + res.setHeader("content-type", "text/event-stream; charset=utf-8"); + res.setHeader("cache-control", "no-cache"); + res.setHeader("x-accel-buffering", "no"); + res.flushHeaders(); + Readable.fromWeb(up.body as import("node:stream/web").ReadableStream) + .on("error", () => res.end()) + .pipe(res); + }), + }); + } +} + +function userMessage(text: string) { + return { + type: "message", + data: { role: "user", content: [{ type: "input_text", text }] }, + }; +} + +/** + * @internal + */ +export const omnigent = toPlugin(OmnigentPlugin); diff --git a/packages/appkit/src/plugins/omnigent/policies.ts b/packages/appkit/src/plugins/omnigent/policies.ts new file mode 100644 index 000000000..4866768ca --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/policies.ts @@ -0,0 +1,89 @@ +/** + * Session modes, compiled to Omnigent session policies. + * + * - `auto`: no gate. + * - `ask`: tools marked as writes pause for the user's approval, and so do + * Omnigent's own OS tools. + * - `read`: writes are denied. + * + * Omnigent enforces these on MCP tool calls itself, whatever the harness. + */ +export type OmnigentMode = "auto" | "ask" | "read"; + +export const OMNIGENT_MODES: readonly OmnigentMode[] = ["auto", "ask", "read"]; + +/** Name prefix of the policies a mode owns, so a mode change replaces only them. */ +export const MODE_POLICY_PREFIX = "appkit-mode"; + +interface OmnigentPolicy { + name: string; + type: "python"; + handler: string; + factory_params?: Record; +} + +function celGate( + writeTools: string[], + result: "ASK" | "DENY", + reason: string, +): string { + if (writeTools.length === 0) return `{"result": "ALLOW"}`; + // MCP tool names reach the policy prefixed with the server name (e.g. + // "mcp__omnigent__app__save_note"), so match on the suffix. + const match = writeTools + .map((t) => `event.data.name.endsWith(${JSON.stringify(t)})`) + .join(" || "); + return ( + `event.type == "tool_call" && (${match}) ? ` + + `{"result": "${result}", "reason": ${JSON.stringify(reason)}} : {"result": "ALLOW"}` + ); +} + +export function modePolicies( + mode: OmnigentMode, + writeTools: string[], +): OmnigentPolicy[] { + if (mode === "ask") { + return [ + { + name: `${MODE_POLICY_PREFIX}-ask`, + type: "python", + handler: "omnigent.policies.builtins.cel.cel_policy", + factory_params: { + expression: celGate(writeTools, "ASK", "Changes app data"), + }, + }, + { + name: `${MODE_POLICY_PREFIX}-os`, + type: "python", + handler: "omnigent.policies.builtins.safety.ask_on_os_tools", + }, + ]; + } + if (mode === "read") { + return [ + { + name: `${MODE_POLICY_PREFIX}-read`, + type: "python", + handler: "omnigent.policies.builtins.cel.cel_policy", + factory_params: { + expression: celGate(writeTools, "DENY", "Read-only session"), + }, + }, + { + name: `${MODE_POLICY_PREFIX}-os`, + type: "python", + handler: "omnigent.policies.builtins.orchestration.read_only_os", + }, + ]; + } + return []; +} + +/** Recovers a session's mode from the names of its policies. */ +export function modeFromPolicies(names: Iterable): OmnigentMode { + const set = new Set(names); + if (set.has(`${MODE_POLICY_PREFIX}-ask`)) return "ask"; + if (set.has(`${MODE_POLICY_PREFIX}-read`)) return "read"; + return "auto"; +} diff --git a/packages/appkit/src/plugins/omnigent/python/mcp_proxy.py b/packages/appkit/src/plugins/omnigent/python/mcp_proxy.py new file mode 100644 index 000000000..bb4ca64d9 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/python/mcp_proxy.py @@ -0,0 +1,46 @@ +"""stdio MCP server that forwards to a Databricks-managed MCP endpoint, for an Omnigent session on a Databricks Sandbox. + + python3 -c + +A Sandbox's Databricks credential is the session owner's own (its default ~/.databrickscfg profile), and Omnigent does +not attach it to HTTP MCP servers. This proxy runs inside the Sandbox, so every tool call reaches the managed MCP +endpoint as the session owner, with their Unity Catalog permissions. Standard library plus the Databricks SDK, which +the Sandbox has. +""" +import json +import sys +import urllib.error +import urllib.request + +from databricks.sdk.core import Config + +URL = sys.argv[1] +_config = Config() + + +def forward(message): + headers = {**_config.authenticate(), "Content-Type": "application/json", "Accept": "application/json, text/event-stream"} + request = urllib.request.Request(URL, data=json.dumps(message).encode(), method="POST", headers=headers) + with urllib.request.urlopen(request, timeout=300) as response: + raw = response.read().decode() + if "text/event-stream" in (response.headers.get("content-type") or ""): + data = [line[5:].strip() for line in raw.splitlines() if line.startswith("data:")] + return json.loads(data[-1]) if data else None + return json.loads(raw) if raw.strip() else None + + +for line in sys.stdin: + try: + message = json.loads(line) + except ValueError: + continue + try: + reply = forward(message) + except urllib.error.HTTPError as error: + reply = {"jsonrpc": "2.0", "id": message.get("id"), + "error": {"code": -32000, "message": f"{error.code} {error.read().decode()[:300]}"}} + except Exception as error: # surfaced to the agent as the tool's error + reply = {"jsonrpc": "2.0", "id": message.get("id"), "error": {"code": -32000, "message": repr(error)[:300]}} + if message.get("id") is not None and reply is not None: + sys.stdout.write(json.dumps(reply) + "\n") + sys.stdout.flush() diff --git a/packages/appkit/src/plugins/omnigent/tests/plugin.test.ts b/packages/appkit/src/plugins/omnigent/tests/plugin.test.ts new file mode 100644 index 000000000..f8efc7d51 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/tests/plugin.test.ts @@ -0,0 +1,277 @@ +import { + createMockRequest, + createMockResponse, + createMockRouter, +} from "@tools/test-helpers"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { ServiceContext } from "../../../context/service-context"; +import { OmnigentPlugin } from "../omnigent"; + +const HOST = "https://ws.example.com"; + +type Route = ( + init: RequestInit, + url: URL, +) => { status?: number; body?: unknown }; + +function plugin(config = {}) { + const p = new OmnigentPlugin(config); + const { router, getHandler } = createMockRouter(); + p.injectRoutes(router); + return { p, getHandler }; +} + +async function call( + handler: (req: unknown, res: unknown) => Promise, + req: object, +) { + const res = createMockResponse(); + await handler(createMockRequest(req), res); + return res; +} + +const user = { token: "user-tok", userId: "u1", email: "alice@example.com" }; + +describe("OmnigentPlugin routes", () => { + let routes: Record; + let calls: Array<{ + method: string; + url: string; + auth: string | null; + body: unknown; + }>; + + beforeEach(() => { + process.env.DATABRICKS_APP_NAME = "orders-app"; + vi.spyOn(ServiceContext, "get").mockReturnValue({ + client: { + getHost: async () => HOST, + authenticate: async (h: Headers) => + h.set("Authorization", "Bearer dev-tok"), + }, + } as never); + routes = {}; + calls = []; + vi.stubGlobal( + "fetch", + vi.fn(async (input: string, init: RequestInit = {}) => { + const url = new URL(input); + const method = init.method ?? "GET"; + const headers = new Headers(init.headers); + calls.push({ + method, + url: url.pathname + url.search, + auth: headers.get("Authorization"), + body: + typeof init.body === "string" ? JSON.parse(init.body) : init.body, + }); + const route = routes[`${method} ${url.pathname}`]; + const out = route ? route(init, url) : { status: 404, body: {} }; + return new Response(JSON.stringify(out.body ?? {}), { + status: out.status ?? 200, + headers: { "content-type": "application/json" }, + }); + }), + ); + }); + + afterEach(() => { + delete process.env.DATABRICKS_APP_NAME; + vi.unstubAllGlobals(); + vi.restoreAllMocks(); + }); + + test("registers the session API", () => { + const { getHandler } = plugin(); + for (const [m, path] of [ + ["GET", "/status"], + ["GET", "/harnesses"], + ["GET", "/sessions"], + ["POST", "/sessions"], + ["GET", "/sessions/:id"], + ["DELETE", "/sessions/:id"], + ["POST", "/sessions/:id/messages"], + ["POST", "/sessions/:id/interrupt"], + ["PUT", "/sessions/:id/mode"], + ["POST", "/sessions/:id/elicitations/:eid"], + ["GET", "/sessions/:id/stream"], + ]) { + expect(getHandler(m, path), `${m} ${path}`).toBeTypeOf("function"); + } + }); + + test("needs a forwarded user token outside development", async () => { + const { getHandler } = plugin(); + const res = await call(getHandler("GET", "/sessions"), {}); + expect(res.status).toHaveBeenCalledWith(401); + expect(calls).toEqual([]); + }); + + test("calls managed Omnigent with the user's token and lists only this app's sessions", async () => { + routes["GET /api/2.0/omnigent/v1/sessions"] = () => ({ + body: { data: [] }, + }); + const { getHandler } = plugin(); + await call(getHandler("GET", "/sessions"), { obo: user }); + expect(calls).toEqual([ + { + method: "GET", + url: "/api/2.0/omnigent/v1/sessions?agent_name=orders-app&limit=100", + auth: "Bearer user-tok", + body: undefined, + }, + ]); + }); + + test("explains the all-apis setup step when Omnigent rejects the token's scope", async () => { + routes["GET /api/2.0/omnigent/v1/sessions"] = () => ({ + status: 403, + body: { message: "Invalid scope, required scopes: all-apis" }, + }); + const { getHandler } = plugin(); + const res = await call(getHandler("GET", "/sessions"), { obo: user }); + expect(res.status).toHaveBeenCalledWith(403); + expect(res.json.mock.calls[0][0].error).toMatch( + /custom-app-integration update/, + ); + }); + + test("hides sessions another app created", async () => { + routes["GET /api/2.0/omnigent/v1/sessions/s1"] = () => ({ + body: { id: "s1", labels: { app: "other-app" } }, + }); + const { getHandler } = plugin(); + const res = await call(getHandler("POST", "/sessions/:id/messages"), { + obo: user, + params: { id: "s1" }, + body: { text: "hi" }, + }); + expect(res.status).toHaveBeenCalledWith(404); + expect(calls.map((c) => c.method)).toEqual(["GET"]); + }); + + test("creates a managed session with the app's agent and sends the first message", async () => { + routes["GET /api/2.1/unity-catalog/model-services"] = () => ({ + body: { + model_services: [ + { + name: "system.ai.claude-sonnet-5", + supported_api_types: ["anthropic/v1/messages"], + }, + ], + }, + }); + let metadata: Record = {}; + routes["POST /api/2.0/omnigent/v1/sessions"] = (init) => { + metadata = JSON.parse(String((init.body as FormData).get("metadata"))); + return { status: 201, body: { session_id: "s9" } }; + }; + routes["POST /api/2.0/omnigent/v1/sessions/s9/events"] = () => ({ + status: 202, + body: { queued: true }, + }); + routes["GET /api/2.0/omnigent/v1/sessions/s9/policies"] = () => ({ + body: { data: [] }, + }); + const policies: unknown[] = []; + routes["POST /api/2.0/omnigent/v1/sessions/s9/policies"] = (init) => { + policies.push(JSON.parse(String(init.body))); + return { body: {} }; + }; + const { getHandler } = plugin({ + instructions: "Help with orders.", + tools: [{ type: "sql" }], + }); + const res = await call(getHandler("POST", "/sessions"), { + obo: user, + body: { harness: "claude-sdk", message: "How many orders?" }, + }); + expect(res.status).toHaveBeenCalledWith(201); + expect(res.json).toHaveBeenCalledWith({ + session_id: "s9", + harness: "claude-sdk", + model: "system.ai.claude-sonnet-5", + mode: "ask", + }); + // ask mode: Databricks SQL's execute_sql pauses for approval + expect(policies.map((p) => (p as { name: string }).name)).toEqual([ + "appkit-mode-ask", + "appkit-mode-os", + ]); + expect(JSON.stringify(policies[0])).toContain("execute_sql"); + expect(metadata).toMatchObject({ + host_type: "managed", + labels: { app: "orders-app" }, + title: "How many orders?", + }); + expect(calls.at(-1)?.body).toMatchObject({ + type: "message", + data: { role: "user" }, + }); + expect(calls.every((c) => c.auth === "Bearer user-tok")).toBe(true); + }); + + test("rejects a model the harness cannot drive", async () => { + routes["GET /api/2.1/unity-catalog/model-services"] = () => ({ + body: { + model_services: [ + { + name: "system.ai.gpt-5-5", + supported_api_types: ["codex/v1/responses"], + }, + ], + }, + }); + const { getHandler } = plugin(); + const res = await call(getHandler("POST", "/sessions"), { + obo: user, + body: { harness: "claude-sdk", model: "system.ai.gpt-5-5" }, + }); + expect(res.status).toHaveBeenCalledWith(400); + }); + + test("resolves an approval only with accept, decline or cancel", async () => { + routes["GET /api/2.0/omnigent/v1/sessions/s1"] = () => ({ + body: { id: "s1", labels: { app: "orders-app" } }, + }); + routes["POST /api/2.0/omnigent/v1/sessions/s1/elicitations/e1/resolve"] = + () => ({ + body: { ok: true }, + }); + const { getHandler } = plugin(); + const handler = getHandler("POST", "/sessions/:id/elicitations/:eid"); + let res = await call(handler, { + obo: user, + params: { id: "s1", eid: "e1" }, + body: { action: "approve" }, + }); + expect(res.status).toHaveBeenCalledWith(400); + res = await call(handler, { + obo: user, + params: { id: "s1", eid: "e1" }, + body: { action: "accept" }, + }); + expect(res.status).toHaveBeenCalledWith(200); + expect(calls.at(-1)).toMatchObject({ + method: "POST", + url: "/api/2.0/omnigent/v1/sessions/s1/elicitations/e1/resolve", + body: { action: "accept" }, + }); + }); + + test("status reports when Sandbox hosts are not enabled", async () => { + routes["GET /api/2.0/omnigent/v1/info"] = () => ({ + body: { server_version: "0.18.0", managed_sandboxes_enabled: false }, + }); + const { getHandler } = plugin(); + const res = await call(getHandler("GET", "/status"), { obo: user }); + expect(res.json).toHaveBeenCalledWith( + expect.objectContaining({ + ready: false, + version: "0.18.0", + sandboxes: false, + }), + ); + }); +}); diff --git a/packages/appkit/src/plugins/omnigent/tests/units.test.ts b/packages/appkit/src/plugins/omnigent/tests/units.test.ts new file mode 100644 index 000000000..b7d2af277 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/tests/units.test.ts @@ -0,0 +1,229 @@ +import { gunzipSync } from "node:zlib"; + +import { describe, expect, test } from "vitest"; + +import { buildBundle } from "../bundle"; +import { offerHarnesses } from "../harnesses"; +import { modeFromPolicies, modePolicies } from "../policies"; +import { toolServer, writeToolNames } from "../tools"; + +function untar(buf: Buffer): Record { + const out: Record = {}; + let off = 0; + while (off + 512 <= buf.length) { + const name = buf + .subarray(off, off + 100) + .toString() + .replace(/\0.*$/s, ""); + if (!name) break; + const size = Number.parseInt( + buf + .subarray(off + 124, off + 136) + .toString() + .replace(/\0.*$/s, ""), + 8, + ); + out[name] = buf.subarray(off + 512, off + 512 + size).toString(); + off += 512 + Math.ceil(size / 512) * 512; + } + return out; +} + +describe("buildBundle", () => { + test("writes the spec and the instructions", () => { + const files = untar( + gunzipSync( + buildBundle({ + name: "orders-app", + harness: "pi", + model: "system.ai.gpt-5-5", + instructions: "Be terse.", + }), + ), + ); + expect(Object.keys(files).sort()).toEqual(["AGENTS.md", "config.yaml"]); + expect(JSON.parse(files["config.yaml"]).executor).toEqual({ + type: "omnigent", + model: "system.ai.gpt-5-5", + config: { harness: "pi" }, + }); + expect(files["AGENTS.md"]).toBe("Be terse."); + }); +}); + +describe("offerHarnesses", () => { + const catalog = [ + { + name: "system.ai.claude-sonnet-5", + apiTypes: ["mlflow/v1/chat/completions", "anthropic/v1/messages"], + }, + { + name: "system.ai.gpt-5-5", + apiTypes: [ + "mlflow/v1/chat/completions", + "openai/v1/responses", + "codex/v1/responses", + ], + }, + { + name: "system.ai.gemini-3-5-flash", + apiTypes: ["mlflow/v1/chat/completions", "gemini/v1/generateContent"], + }, + { + name: "system.ai.glm-5-3", + apiTypes: [ + "mlflow/v1/chat/completions", + "anthropic/v1/messages", + "codex/v1/responses", + ], + }, + { + name: "system.ai.grok-4-7", + apiTypes: ["mlflow/v1/chat/completions", "openai/v1/responses"], + }, + ]; + + test("matches UC models to harnesses by the API each harness calls", () => { + const offers = offerHarnesses({ + allowed: ["claude-sdk", "pi", "openai-agents"], + models: catalog, + }); + const by = Object.fromEntries(offers.map((o) => [o.id, o])); + expect(by["claude-sdk"].models).toEqual([ + "system.ai.claude-sonnet-5", + "system.ai.glm-5-3", + ]); + expect(by.pi.models).toEqual(["system.ai.gpt-5-5", "system.ai.glm-5-3"]); + expect(by["openai-agents"].models).toEqual([ + "system.ai.gpt-5-5", + "system.ai.grok-4-7", + ]); + expect(by["claude-sdk"].defaultModel).toBe("system.ai.claude-sonnet-5"); + expect(by.pi.defaultModel).toBe("system.ai.gpt-5-5"); + }); + + test("offers only gateway harnesses, in the allowed order", () => { + const offers = offerHarnesses({ + allowed: ["pi", "cursor", "claude-sdk"], + models: catalog, + }); + expect(offers.map((o) => o.id)).toEqual(["pi", "claude-sdk"]); + }); +}); + +describe("tool sources", () => { + const host = "https://ws.example.com"; + + test("each source becomes a stdio proxy to its managed MCP URL", () => { + const urls = [ + { type: "sql" }, + { type: "genie", spaceId: "01f0abcdef12345678" }, + { type: "functions", catalog: "main", schema: "tools" }, + { + type: "functions", + catalog: "main", + schema: "tools", + function: "price", + }, + { type: "ai-search", catalog: "main", schema: "docs", index: "faq" }, + { type: "service", service: "system.ai.web_search" }, + ].map((s) => { + const spec = toolServer(s as never, host); + expect(spec).toMatchObject({ transport: "stdio", command: "python3" }); + expect(spec.args[0]).toBe("-c"); + expect(spec.args[1]).toContain("databricks.sdk.core"); + return [spec.name, spec.args[2]]; + }); + expect(urls).toEqual([ + ["sql", `${host}/api/2.0/mcp/sql`], + ["genie_12345678", `${host}/api/2.0/mcp/genie/01f0abcdef12345678`], + ["functions_tools", `${host}/api/2.0/mcp/functions/main/tools`], + [ + "functions_tools_price", + `${host}/api/2.0/mcp/functions/main/tools/price`, + ], + ["search_faq", `${host}/api/2.0/mcp/ai-search/main/docs/faq`], + [ + "system_ai_web_search", + `${host}/ai-gateway/mcp-services/system.ai.web_search`, + ], + ]); + }); + + test("names can be set and are made MCP-safe", () => { + expect(toolServer({ type: "sql", name: "Orders DB" }, host).name).toBe( + "orders_db", + ); + }); + + test("Databricks SQL's execute_sql is a write by default; apps add their own", () => { + expect(writeToolNames([{ type: "sql" }], ["send_message"])).toEqual([ + "execute_sql", + "send_message", + ]); + expect(writeToolNames([{ type: "genie", spaceId: "x" }])).toEqual([]); + }); + + test("the bundle carries one tools/mcp file per server", () => { + const files = untar( + gunzipSync( + buildBundle({ + name: "orders-app", + harness: "pi", + model: "system.ai.gpt-5-5", + instructions: "", + mcpServers: [toolServer({ type: "sql" }, host)], + }), + ), + ); + expect(JSON.parse(files["tools/mcp/sql.yaml"])).toMatchObject({ + name: "sql", + transport: "stdio", + command: "python3", + }); + expect(() => + buildBundle({ + name: "a", + harness: "pi", + model: "m", + instructions: "", + mcpServers: [ + toolServer({ type: "sql" }, host), + toolServer({ type: "sql" }, host), + ], + }), + ).toThrow(/Duplicate MCP server name/); + }); +}); + +describe("policies", () => { + test("ask gates write tools and OS tools", () => { + const p = modePolicies("ask", ["execute_sql"]); + expect(p.map((x) => x.name)).toEqual(["appkit-mode-ask", "appkit-mode-os"]); + expect(String(p[0].factory_params?.expression)).toContain( + 'endsWith("execute_sql")', + ); + expect(String(p[0].factory_params?.expression)).toContain('"ASK"'); + }); + + test("read denies writes; auto has no policies", () => { + expect( + String(modePolicies("read", ["w"])[0].factory_params?.expression), + ).toContain('"DENY"'); + expect(modePolicies("auto", ["w"])).toEqual([]); + }); + + test("no write tools means the gate allows everything", () => { + expect(modePolicies("ask", [])[0].factory_params?.expression).toBe( + '{"result": "ALLOW"}', + ); + }); + + test("mode round-trips through policy names", () => { + for (const m of ["ask", "read", "auto"] as const) { + expect(modeFromPolicies(modePolicies(m, ["w"]).map((p) => p.name))).toBe( + m, + ); + } + }); +}); diff --git a/packages/appkit/src/plugins/omnigent/tools.ts b/packages/appkit/src/plugins/omnigent/tools.ts new file mode 100644 index 000000000..47b920331 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/tools.ts @@ -0,0 +1,115 @@ +import fs from "node:fs"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; + +import type { McpServerSpec } from "./bundle"; + +/** + * A Databricks-managed MCP server the app's sessions can call. Calls run as + * the session's user, so Unity Catalog permissions apply. + */ +export type OmnigentToolSource = + /** Databricks SQL: run SQL statements on a warehouse. */ + | { type: "sql"; name?: string } + /** One Genie space. */ + | { type: "genie"; spaceId: string; name?: string } + /** Unity Catalog functions in a schema, or one function. */ + | { + type: "functions"; + catalog: string; + schema: string; + function?: string; + name?: string; + } + /** AI Search over an index, or every index in a schema. */ + | { + type: "ai-search"; + catalog: string; + schema: string; + index?: string; + name?: string; + } + /** A Unity Gateway MCP Service, e.g. `system.ai.web_search` or `main.tools.jira`. */ + | { type: "service"; service: string; name?: string }; + +/** Tools that change data by default: `ask` pauses on them and `read` denies them. */ +const DEFAULT_WRITE_TOOLS: Partial< + Record +> = { + sql: ["execute_sql"], +}; + +const PROXY = fs.readFileSync( + path.join( + path.dirname(fileURLToPath(import.meta.url)), + "python", + "mcp_proxy.py", + ), + "utf8", +); + +function urlPath(source: OmnigentToolSource): string { + const seg = encodeURIComponent; + switch (source.type) { + case "sql": + return "/api/2.0/mcp/sql"; + case "genie": + return `/api/2.0/mcp/genie/${seg(source.spaceId)}`; + case "functions": + return `/api/2.0/mcp/functions/${seg(source.catalog)}/${seg(source.schema)}${source.function ? `/${seg(source.function)}` : ""}`; + case "ai-search": + return `/api/2.0/mcp/ai-search/${seg(source.catalog)}/${seg(source.schema)}${source.index ? `/${seg(source.index)}` : ""}`; + case "service": + return `/ai-gateway/mcp-services/${seg(source.service)}`; + } +} + +function defaultName(source: OmnigentToolSource): string { + switch (source.type) { + case "sql": + return "sql"; + case "genie": + return `genie_${source.spaceId.slice(-8)}`; + case "functions": + return `functions_${source.schema}${source.function ? `_${source.function}` : ""}`; + case "ai-search": + return `search_${source.index ?? source.schema}`; + case "service": + // The full name: a bare service name can collide with an Omnigent built-in tool (`web_search`). + return source.service.replace(/\./g, "_"); + } +} + +/** + * The bundle entry for a tool source: a stdio MCP server that runs inside the + * session's Sandbox and forwards to the managed MCP URL as the session's user. + * (A Sandbox's credential is not attached to HTTP MCP servers.) + */ +export function toolServer( + source: OmnigentToolSource, + workspaceHost: string, +): McpServerSpec { + const name = (source.name ?? defaultName(source)) + .toLowerCase() + .replace(/[^a-z0-9_-]/g, "_"); + return { + name, + description: `Databricks ${source.type} tools, run as the signed-in user`, + transport: "stdio", + command: "python3", + args: ["-c", PROXY, `${workspaceHost}${urlPath(source)}`], + }; +} + +/** Tool names that change data across the sources, plus any the app lists. */ +export function writeToolNames( + sources: OmnigentToolSource[], + extra: string[] = [], +): string[] { + return [ + ...new Set([ + ...sources.flatMap((s) => DEFAULT_WRITE_TOOLS[s.type] ?? []), + ...extra, + ]), + ]; +} diff --git a/packages/appkit/src/plugins/omnigent/types.ts b/packages/appkit/src/plugins/omnigent/types.ts new file mode 100644 index 000000000..99eae4a02 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/types.ts @@ -0,0 +1,53 @@ +import type { BasePluginConfig } from "shared"; + +import type { OmnigentMode } from "./policies"; +import type { OmnigentToolSource } from "./tools"; + +export type { OmnigentMode } from "./policies"; +export type { OmnigentToolSource } from "./tools"; + +export interface IOmnigentConfig extends BasePluginConfig { + /** System instructions for every session (the bundle's `AGENTS.md`). */ + instructions?: string; + /** + * Databricks-managed MCP servers every session can call, as the session's + * user: Databricks SQL, Genie spaces, Unity Catalog functions, AI Search, + * Unity Gateway MCP Services. + */ + tools?: OmnigentToolSource[]; + /** + * Tool names that change data, on top of the built-in ones (Databricks + * SQL's `execute_sql`). `ask` mode pauses on them and `read` mode denies + * them. + */ + writeTools?: string[]; + /** Mode new sessions start in. Default `ask`. */ + defaultMode?: OmnigentMode; + /** + * Harnesses users may pick, out of the ones whose model calls go through + * Unity Gateway. Default: all of them. + */ + harnesses?: string[]; + /** Default harness and model for new sessions. */ + defaultHarness?: string; + defaultModel?: string; +} + +export interface CreateSessionRequest { + harness?: string; + model?: string; + mode?: OmnigentMode; + message?: string; + title?: string; +} + +/** From `GET /status`: whether the workspace's managed Omnigent is usable by this user. */ +export interface OmnigentStatus { + ready: boolean; + /** Why it is not ready, e.g. the app's token lacks the `all-apis` scope. */ + error?: string; + /** Managed Omnigent server version. */ + version?: string; + /** Whether Databricks Sandbox hosts are available. */ + sandboxes?: boolean; +} diff --git a/packages/appkit/tsdown.config.ts b/packages/appkit/tsdown.config.ts index bb17c29ba..8474de357 100644 --- a/packages/appkit/tsdown.config.ts +++ b/packages/appkit/tsdown.config.ts @@ -43,6 +43,11 @@ export default defineConfig([ to: "dist/plugins/server/remote-tunnel", flatten: true, }, + { + from: "src/plugins/omnigent/python/*.py", + to: "dist/plugins/omnigent/python", + flatten: true, + }, ], }, ]); diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index db53b96b7..fe3f78234 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -459,6 +459,26 @@ } } }, + "omnigent": { + "name": "omnigent", + "displayName": "Omnigent Plugin", + "description": "Omnigent agent sessions for the app's users on the workspace's managed Omnigent: Claude Agent SDK, Pi and OpenAI Agents SDK on Databricks Sandboxes, with models through Unity Gateway", + "package": "@databricks/appkit", + "resources": { + "required": [], + "optional": [] + }, + "stability": "beta", + "scaffolding": { + "rules": { + "must": [ + "Enable the Omnigent and Databricks Sandbox previews in the workspace", + "Account admin: add all-apis to the app's OAuth integration (databricks account custom-app-integration update)", + "Re-run that update after a bundle deploy or a user_api_scopes change, which both drop all-apis" + ] + } + } + }, "server": { "name": "server", "displayName": "Server Plugin",