From 0bc2968a5a048efdb96152cf4bf233b4229eab0d Mon Sep 17 00:00:00 2001 From: Noah Solomon Date: Wed, 30 Sep 2026 00:05:01 +1000 Subject: [PATCH 1/4] feat(appkit): add the omnigent plugin (1/4) Embeds Omnigent, the open-source agent meta-harness, in an AppKit app: Claude Agent SDK, Pi and OpenAI Agents SDK sessions, with model calls going through Unity AI Gateway as the app's service principal. - One Omnigent server on loopback in header-auth mode (random header name) and one host per active user, supervised with restart backoff and idle shutdown; a parent-death watchdog so no child outlives Node. - A loopback model gateway: children hold a placeholder token, the gateway forwards only model-serving paths with the real token. - Harnesses are offered the UC system.ai models whose API types they speak, plus the workspace's chat serving endpoints. - Session routes under /api/omnigent, scoped to the signed-in user. - docs/docs/plugins/omnigent.md, extended by the later PRs. Co-authored-by: Isaac Signed-off-by: Noah Solomon --- .gitignore | 1 + docs/docs/plugins/omnigent.md | 69 +++ .../src/plugins/beta-exports.generated.ts | 1 + .../appkit/src/plugins/omnigent/bundle.ts | 64 +++ .../appkit/src/plugins/omnigent/harnesses.ts | 83 +++ packages/appkit/src/plugins/omnigent/index.ts | 2 + .../appkit/src/plugins/omnigent/manifest.json | 44 ++ .../appkit/src/plugins/omnigent/omnigent.ts | 540 ++++++++++++++++++ .../src/plugins/omnigent/runtime/env.ts | 144 +++++ .../src/plugins/omnigent/runtime/gateway.ts | 153 +++++ .../omnigent/runtime/python/watchdog.py | 43 ++ .../src/plugins/omnigent/runtime/runtime.ts | 373 ++++++++++++ .../src/plugins/omnigent/runtime/tls.ts | 66 +++ .../plugins/omnigent/tests/gateway.test.ts | 104 ++++ .../src/plugins/omnigent/tests/plugin.test.ts | 73 +++ .../src/plugins/omnigent/tests/units.test.ts | 222 +++++++ packages/appkit/src/plugins/omnigent/types.ts | 47 ++ packages/appkit/tsdown.config.ts | 5 + template/appkit.plugins.json | 20 + 19 files changed, 2054 insertions(+) create mode 100644 docs/docs/plugins/omnigent.md create mode 100644 packages/appkit/src/plugins/omnigent/bundle.ts create mode 100644 packages/appkit/src/plugins/omnigent/harnesses.ts create mode 100644 packages/appkit/src/plugins/omnigent/index.ts create mode 100644 packages/appkit/src/plugins/omnigent/manifest.json create mode 100644 packages/appkit/src/plugins/omnigent/omnigent.ts create mode 100644 packages/appkit/src/plugins/omnigent/runtime/env.ts create mode 100644 packages/appkit/src/plugins/omnigent/runtime/gateway.ts create mode 100644 packages/appkit/src/plugins/omnigent/runtime/python/watchdog.py create mode 100644 packages/appkit/src/plugins/omnigent/runtime/runtime.ts create mode 100644 packages/appkit/src/plugins/omnigent/runtime/tls.ts create mode 100644 packages/appkit/src/plugins/omnigent/tests/gateway.test.ts create mode 100644 packages/appkit/src/plugins/omnigent/tests/plugin.test.ts create mode 100644 packages/appkit/src/plugins/omnigent/tests/units.test.ts create mode 100644 packages/appkit/src/plugins/omnigent/types.ts diff --git a/.gitignore b/.gitignore index b70e0c2e8..b9a39ae83 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,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..09823ba9a --- /dev/null +++ b/docs/docs/plugins/omnigent.md @@ -0,0 +1,69 @@ +# 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 embeds [Omnigent](https://github.com/omnigent-ai/omnigent), the open-source agent meta-harness, in a Databricks AppKit app. Your users get Claude Agent SDK, Pi and OpenAI Agents SDK sessions inside the app, with model calls going through Unity AI Gateway as the app's service principal. + +Omnigent owns the agent loop: sessions, history, approvals. The plugin runs it safely inside a Databricks App. 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." }), + ], +}); +``` + +The app also needs: + +- **Python.** Add a `pyproject.toml` with `requires-python = ">=3.12,<3.13"` and `omnigent[databricks]`, plus a `uv.lock`. Databricks Apps then builds a Python 3.12 venv next to `npm install`. A plain `requirements.txt` gets Python 3.11, which Omnigent does not support. The lockfile must reference public PyPI. +- **Harness CLIs.** Add them to your `package.json`, so you pin their versions: `@anthropic-ai/claude-code` and `@earendil-works/pi-coding-agent`. The plugin only offers harnesses whose CLI is installed. The Claude Agent SDK's bundled CLI is too old for current Claude models, so add `@anthropic-ai/claude-code` even though the SDK ships one. +- **Start with `node`.** In `app.yaml`, use `command: ["node", "server.js"]`, not `npm run start`. npm does not pass SIGTERM on, so the plugin never gets to stop its child processes. As a backstop, each child also exits when Node does, so a redeploy still succeeds. + +## Configuration + +| Option | Default | | +|---|---|---| +| `instructions` | a generic assistant prompt | System instructions for every session | +| `harnesses` | `claude-sdk`, `pi`, `openai-agents` | Harnesses users may pick | +| `defaultHarness`, `defaultModel` | first available | | +| `runtime.python` | `./.venv/bin/python` | | +| `runtime.home` | `$TMPDIR/appkit-omnigent` | The plugin's state directory | +| `runtime.hostIdleMs` | 30 min | Stop a user's Omnigent host after this long idle | + +Each harness is offered the Unity AI 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 scoped to the signed-in user. + +| Route | | +|---|---| +| `GET /harnesses` | Harnesses available to the user, each with the gateway models it can drive | +| `GET /sessions`, `POST /sessions` | List, create (`harness`, `model`, `message`, `title`) | +| `GET /sessions/:id`, `DELETE /sessions/:id` | Snapshot (items, status), 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 | +| `GET /status` | Runtime health: hosts, gateway counters | + +## How it works + +- **One Omnigent server, one host per user.** The server listens on loopback in Omnigent's header-auth mode, under a random header name only the plugin knows. Each active user gets their own host process, registered as that user; Omnigent only runs a session on its owner's host. An idle host stops after `runtime.hostIdleMs`. +- **No credential in any child process.** The harnesses' Databricks profile points at a model gateway inside the plugin, on loopback, with a placeholder token. The gateway forwards only model-serving paths to the workspace and adds the service principal's token. Every other path returns 403. + +## Limits + +- **Sessions live in the container.** Omnigent keeps them in SQLite under the plugin's state directory, so they survive a process restart but not a redeploy. +- **Only gateway harnesses.** Harnesses that sign in with their vendor (Cursor, Copilot, Gemini and others) bypass Unity AI Gateway, so they are not offered. Native terminal harnesses (`claude-native`, `pi-native`) are not offered yet: they need tmux and a terminal client. +- **Shared gateway limits.** All users' model traffic uses the service principal's gateway limits. 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..b82cb64db --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/bundle.ts @@ -0,0 +1,64 @@ +import { gzipSync } from "node:zlib"; + +interface BundleSpec { + name: string; + description?: string; + harness: string; + model: string; + /** Databricks profile the harness uses for model calls. */ + profile: string; + instructions: string; +} + +/** + * An Omnigent agent bundle: `config.yaml` (the agent spec) and `AGENTS.md` + * (the instructions), as a tar.gz. The YAML is 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, + profile: spec.profile, + config: { harness: spec.harness }, + auth: { type: "databricks", profile: spec.profile }, + }, + instructions: "AGENTS.md", + }; + const files: Record = { + "config.yaml": `${JSON.stringify(config, null, 2)}\n`, + "AGENTS.md": spec.instructions, + }; + 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/harnesses.ts b/packages/appkit/src/plugins/omnigent/harnesses.ts new file mode 100644 index 000000000..3a726f90c --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/harnesses.ts @@ -0,0 +1,83 @@ +import type { OmnigentHarness } from "./types"; + +interface HarnessInfo { + label: string; + /** + * Unity AI 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 Omnigent routes through a Databricks profile + * (`AuthModel.OMNIGENT_CREDENTIAL`), so they go through Unity AI Gateway. + * The other harnesses Omnigent supports sign in with their vendor, so they + * are not offered. + */ +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 AI Gateway model, e.g. `system.ai.claude-sonnet-5`. */ +export interface GatewayModel { + name: string; + /** The model service's `supported_api_types`. */ + apiTypes: string[]; +} + +export interface HarnessOffer extends OmnigentHarness { + models: string[]; + defaultModel?: string; +} + +/** + * What a user can pick: the allowed gateway harnesses their host reports as + * usable, each with the models whose API it speaks. + */ +export function offerHarnesses(opts: { + allowed: string[]; + configured: Record; + models: GatewayModel[]; +}): HarnessOffer[] { + const out: HarnessOffer[] = []; + for (const id of opts.allowed) { + const info = GATEWAY_HARNESSES[id]; + const ready = opts.configured[id]; + if (!info) continue; + // "needs-auth" is advisory for gateway harnesses: their auth comes from the profile. + if (ready === undefined || ready === false || ready === "binary-missing") + 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, + ready, + 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..7e2a18198 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/manifest.json @@ -0,0 +1,44 @@ +{ + "$schema": "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", + "name": "omnigent", + "displayName": "Omnigent Plugin", + "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway", + "stability": "beta", + "resources": { + "required": [], + "optional": [] + }, + "scaffolding": { + "rules": { + "must": [ + "Add pyproject.toml (python >=3.12,<3.13) with omnigent[databricks] and a uv.lock that uses public PyPI", + "Add harness CLIs to package.json: @anthropic-ai/claude-code, @earendil-works/pi-coding-agent", + "Start the app with node, not npm run, so SIGTERM reaches the plugin and it can stop its child processes" + ] + } + }, + "config": { + "schema": { + "type": "object", + "properties": { + "instructions": { + "type": "string", + "description": "System instructions for every session" + }, + "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..77c230052 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/omnigent.ts @@ -0,0 +1,540 @@ +import os from "node:os"; +import path from "node:path"; +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 { + DEFAULT_HARNESSES, + type GatewayModel, + offerHarnesses, + type HarnessOffer, +} from "./harnesses"; +import manifest from "./manifest.json"; +import { GATEWAY_PROFILE } from "./runtime/env"; +import { OmnigentRuntime } from "./runtime/runtime"; +import type { CreateSessionRequest, IOmnigentConfig } from "./types"; + +const logger = createLogger("omnigent"); + +const MODELS_TTL_MS = 5 * 60_000; + +/** + * Omnigent, embedded in the app. + * + * Runs an Omnigent server and one host per active user as supervised child + * processes, so every harness Omnigent can route through Unity AI Gateway is + * available to the app's users. Model calls are billed to the app's service + * principal through a loopback gateway (no child holds a real credential). + */ +export class OmnigentPlugin extends Plugin { + static manifest = defineManifest<"omnigent">(manifest); + + protected static description = + "Embedded Omnigent: Claude Agent SDK, Pi and OpenAI Agents SDK, through Unity AI Gateway"; + declare protected config: IOmnigentConfig; + + private runtime?: OmnigentRuntime; + private models?: { at: number; list: GatewayModel[] }; + + async setup(): Promise { + const ctx = ServiceContext.get(); + const upstream = (await ctx.client.config.getHost()) + .toString() + .replace(/\/+$/, ""); + const appDir = process.cwd(); + this.runtime = new OmnigentRuntime({ + root: + this.config.runtime?.home ?? path.join(os.tmpdir(), "appkit-omnigent"), + appDir, + python: + this.config.runtime?.python ?? + path.join(appDir, ".venv", "bin", "python"), + upstream, + token: () => this.serviceToken(), + hostIdleMs: this.config.runtime?.hostIdleMs ?? 30 * 60_000, + }); + // Do not block app start on the server. + await this.runtime.start(); + this.runtime + .whenReady() + .catch((err) => logger.error("Omnigent did not start: %O", err)); + } + + async shutdown(): Promise { + this.streamManager.abortAll(); + await this.runtime?.stop(); + } + + // ---------------------------------------------------------------- helpers + + private rt(): OmnigentRuntime { + if (!this.runtime) throw new Error("Omnigent runtime is not started"); + return this.runtime; + } + + private async serviceToken(): Promise { + const headers = new Headers(); + await ServiceContext.get().client.config.authenticate(headers); + const auth = headers.get("Authorization") ?? ""; + if (!auth.startsWith("Bearer ")) + throw new Error("service principal has no bearer token"); + return auth.slice("Bearer ".length); + } + + /** + * Chat models the app can call through Unity AI Gateway: the UC `system.ai` + * model services, with the API types each one supports. + */ + private async servedModels(): Promise { + if (this.models && Date.now() - this.models.at < MODELS_TTL_MS) + return this.models.list; + const host = (await ServiceContext.get().client.config.getHost()) + .toString() + .replace(/\/+$/, ""); + const headers = { Authorization: `Bearer ${await this.serviceToken()}` }; + + const catalog: GatewayModel[] = []; + try { + 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 }, + ); + if (!r.ok) 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; + catalog.push({ + name: svc.name.replace(/^model-services\//, ""), + apiTypes, + }); + } + page = body.next_page_token || undefined; + } while (page); + } catch (err) { + logger.warn("could not list UC model services: %s", String(err)); + } + + const list = catalog.sort((a, b) => a.name.localeCompare(b.name)); + this.models = { at: Date.now(), list }; + return list; + } + + private allowedHarnesses(): string[] { + return this.config.harnesses ?? DEFAULT_HARNESSES; + } + + private async offers(user: string): Promise { + const rt = this.rt(); + const host = await rt.ensureHost(user); + const r = await rt.request("GET", "/v1/hosts", user); + const { hosts = [] } = (await r.json()) as { + hosts?: Array<{ + host_id: string; + configured_harnesses?: Record; + }>; + }; + const mine = hosts.find((h) => h.host_id === host.hostId); + return offerHarnesses({ + allowed: this.allowedHarnesses(), + configured: mine?.configured_harnesses ?? {}, + models: await this.servedModels(), + }); + } + + 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( + email: string | undefined, + harness: string, + model: string, + ): Buffer { + return buildBundle({ + name: "appkit", + description: "AppKit app agent", + harness, + model, + profile: GATEWAY_PROFILE, + instructions: this.instructions(email), + }); + } + + /** + * Waits until the session's runner has finished initializing. A turn posted + * while it initializes is delivered twice (initialization replays the history + * and so does the post); the duplicate is rejected with HTTP 204 and the + * session is marked failed. Initialization publishes the session's REPL + * terminal as a resource just before it finishes, so wait for that. + */ + private async runnerReady( + user: string, + sid: string, + ms: number, + ): Promise { + const end = Date.now() + ms; + while (Date.now() < end) { + const r = await this.rt().request( + "GET", + `/v1/sessions/${sid}?include_items=true`, + user, + ); + if (r.ok) { + const snap = (await r.json()) as { + runner_online?: boolean; + items?: Array<{ type?: string }>; + }; + if ( + snap.runner_online && + (snap.items ?? []).some((i) => i.type === "resource_event") + ) { + await new Promise((resolve) => setTimeout(resolve, 300)); + return; + } + } + await new Promise((resolve) => setTimeout(resolve, 250)); + } + } + + /** 404s unless the session exists and belongs to the caller (Omnigent scopes by user). */ + private async owned( + user: string, + sid: string, + res: express.Response, + ): Promise | null> { + const r = await this.rt().request( + "GET", + `/v1/sessions/${encodeURIComponent(sid)}?include_items=true`, + user, + ); + if (!r.ok) { + res + .status(r.status === 403 ? 404 : r.status) + .json({ error: "Session not found" }); + return null; + } + return (await r.json()) as Record; + } + + 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, + user: string, + ) => Promise, + ) { + return async (req: express.Request, res: express.Response) => { + let user: string; + try { + user = this.resolveUserId(req); + } catch { + res.status(401).json({ error: "Sign in to the app to use its agent." }); + return; + } + try { + await fn(req, res, user); + } catch (err) { + logger.error("omnigent route failed: %O", err); + if (!res.headersSent) { + res + .status(503) + .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) => { + res.json(this.rt().status()); + }), + }); + + this.route(router, { + name: "harnesses", + method: "get", + path: "/harnesses", + handler: this.handler(async (_req, res, user) => { + const harnesses = await this.offers(user); + res.json({ + harnesses, + 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, user) => { + await OmnigentPlugin.relay( + res, + await this.rt().request("GET", "/v1/sessions?limit=100", user), + ); + }), + }); + + this.route(router, { + name: "createSession", + method: "post", + path: "/sessions", + handler: this.handler(async (req, res, user) => { + const body = (req.body ?? {}) as CreateSessionRequest; + const offers = await this.offers(user); + 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 rt = this.rt(); + const host = await rt.ensureHost(user); + 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.name }, + host_type: "external", + host_id: host.hostId, + workspace: rt.layout.userWork(host.userId), + }), + ); + form.set( + "bundle", + new Blob( + [ + new Uint8Array( + this.bundle(req.header("x-forwarded-email"), harness.id, model), + ), + ], + { + type: "application/gzip", + }, + ), + "bundle.tar.gz", + ); + const created = await rt.request("POST", "/v1/sessions", user, { + body: form, + }); + if (!created.ok) { + await OmnigentPlugin.relay(res, created); + return; + } + const { session_id: sid } = (await created.json()) as { + session_id: string; + }; + if (body.message) { + // A turn delivered before the session's runner is up can be rejected (HTTP 204); wait briefly. + await this.runnerReady(user, sid, 5_000); + const sent = await rt.request( + "POST", + `/v1/sessions/${sid}/events`, + user, + { + json: userMessage(body.message), + }, + ); + if (!sent.ok) { + await OmnigentPlugin.relay(res, sent); + return; + } + } + res.status(201).json({ session_id: sid, harness: harness.id, model }); + }), + }); + + this.route(router, { + name: "getSession", + method: "get", + path: "/sessions/:id", + handler: this.handler(async (req, res, user) => { + const snap = await this.owned(user, req.params.id, res); + if (!snap) return; + res.json(snap); + }), + }); + + this.route(router, { + name: "deleteSession", + method: "delete", + path: "/sessions/:id", + handler: this.handler(async (req, res, user) => { + await OmnigentPlugin.relay( + res, + await this.rt().request( + "DELETE", + `/v1/sessions/${encodeURIComponent(req.params.id)}`, + user, + ), + ); + }), + }); + + this.route(router, { + name: "sendMessage", + method: "post", + path: "/sessions/:id/messages", + handler: this.handler(async (req, res, user) => { + const text = (req.body as { text?: string })?.text; + if (!text) { + res.status(400).json({ error: "text is required" }); + return; + } + await this.rt().ensureHost(user); // a host idled out or a fresh container + await OmnigentPlugin.relay( + res, + await this.rt().request( + "POST", + `/v1/sessions/${encodeURIComponent(req.params.id)}/events`, + user, + { + json: userMessage(text), + }, + ), + ); + }), + }); + + this.route(router, { + name: "interrupt", + method: "post", + path: "/sessions/:id/interrupt", + handler: this.handler(async (req, res, user) => { + await OmnigentPlugin.relay( + res, + await this.rt().request( + "POST", + `/v1/sessions/${encodeURIComponent(req.params.id)}/events`, + user, + { + json: { type: "interrupt", data: {} }, + }, + ), + ); + }), + }); + + this.route(router, { + name: "stream", + method: "get", + path: "/sessions/:id/stream", + handler: this.handler(async (req, res, user) => { + 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.rt().request( + "GET", + `/v1/sessions/${encodeURIComponent(req.params.id)}/stream`, + user, + { + 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); + }), + }); + } + + exports() { + return { + /** Runtime health: server, hosts, gateway counters. */ + status: () => this.rt().status(), + }; + } +} + +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/runtime/env.ts b/packages/appkit/src/plugins/omnigent/runtime/env.ts new file mode 100644 index 000000000..bdb1b31d0 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/runtime/env.ts @@ -0,0 +1,144 @@ +import { createHash } from "node:crypto"; +import fs from "node:fs"; +import path from "node:path"; + +/** The Databricks profile the children use; it points at the model gateway. */ +export const GATEWAY_PROFILE = "appkit-omnigent"; + +/** + * Directory layout under the runtime root: + * + * server/ Omnigent server data + * users//home one user's $HOME: Databricks profile, Omnigent config + * users//data that user's host data + * work/ that user's workspace (sessions' cwd) + * bin/ public: CA bundle + * tls/ gateway key + * logs/ + */ +export class RuntimeLayout { + constructor(readonly root: string) {} + + static userId(email: string): string { + return createHash("sha256") + .update(email.toLowerCase()) + .digest("hex") + .slice(0, 16); + } + + get serverDir() { + return path.join(this.root, "server"); + } + get bin() { + return path.join(this.root, "bin"); + } + get tls() { + return path.join(this.root, "tls"); + } + get logs() { + return path.join(this.root, "logs"); + } + userHome(id: string) { + return path.join(this.root, "users", id, "home"); + } + userData(id: string) { + return path.join(this.root, "users", id, "data"); + } + userWork(id: string) { + return path.join(this.root, "work", id); + } + + create(): void { + for (const d of [this.root, this.serverDir, this.bin, this.logs]) { + fs.mkdirSync(d, { recursive: true }); + } + } +} + +interface GatewayEndpoints { + httpsUrl: string; + placeholder: string; +} + +/** + * Writes a user's `$HOME` config: a `.databrickscfg` whose profile points at + * the gateway with the placeholder token (no real credential), and an + * Omnigent `config.yaml` that authenticates through that profile. + */ +export function writeUserConfig(home: string, gw: GatewayEndpoints): void { + fs.mkdirSync(path.join(home, ".config", "omnigent"), { recursive: true }); + const cfg = [ + `[${GATEWAY_PROFILE}]`, + `host = ${gw.httpsUrl}`, + `token = ${gw.placeholder}`, + "auth_type = pat", + "", + ].join("\n"); + fs.writeFileSync(path.join(home, ".databrickscfg"), cfg, { mode: 0o600 }); + const omnigentConfig = [ + "auth:", + " type: databricks", + ` profile: ${GATEWAY_PROFILE}`, + "", + ].join("\n"); + fs.writeFileSync( + path.join(home, ".config", "omnigent", "config.yaml"), + omnigentConfig, + ); +} + +/** Environment variables a child must never inherit. */ +const STRIPPED = [ + /^DATABRICKS_(CLIENT_ID|CLIENT_SECRET|TOKEN|HOST|CONFIG_PROFILE|CONFIG_FILE)$/, + /^(PGPASSWORD|OMNIGENT_KEY_SECRET)$/, + /^(CLAUDE|CLAUDECODE|UVICORN|npm_)/, + /^NODE_OPTIONS$/, +]; + +interface ChildEnvOptions { + home: string; + dataDir: string; + gateway: GatewayEndpoints; + caBundle: string; + caCert: string; + pathPrepend: string[]; + extra?: Record; +} + +/** + * The environment of an Omnigent child: the app's environment minus every + * credential, with the gateway as its Databricks host and its trust stores + * extended with the gateway's certificate. + */ +export function childEnv( + base: NodeJS.ProcessEnv, + opts: ChildEnvOptions, +): Record { + const env: Record = {}; + for (const [k, v] of Object.entries(base)) { + if (v === undefined || STRIPPED.some((re) => re.test(k))) continue; + env[k] = v; + } + Object.assign(env, { + HOME: opts.home, + OMNIGENT_DATA_DIR: opts.dataDir, + OMNIGENT_CONFIG_HOME: path.join(opts.home, ".config", "omnigent"), + DATABRICKS_CONFIG_FILE: path.join(opts.home, ".databrickscfg"), + DATABRICKS_HOST: opts.gateway.httpsUrl, + SSL_CERT_FILE: opts.caBundle, + REQUESTS_CA_BUNDLE: opts.caBundle, + NODE_EXTRA_CA_CERTS: opts.caCert, + OMNIGENT_SKIP_UPDATE_CHECK: "1", + OMNIGENT_UPDATE_CHECK: "0", + NO_COLOR: "1", + PYTHONUNBUFFERED: "1", + PATH: [...opts.pathPrepend, base.PATH ?? ""] + .filter(Boolean) + .join(path.delimiter), + }); + for (const [k, v] of Object.entries(opts.extra ?? {})) { + if (v === undefined) delete env[k]; + else env[k] = v; + } + return env; +} diff --git a/packages/appkit/src/plugins/omnigent/runtime/gateway.ts b/packages/appkit/src/plugins/omnigent/runtime/gateway.ts new file mode 100644 index 000000000..7adda8be8 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/runtime/gateway.ts @@ -0,0 +1,153 @@ +import { randomBytes } from "node:crypto"; +import type http from "node:http"; +import https from "node:https"; +import type { AddressInfo } from "node:net"; + +import { createLogger } from "../../../logging"; + +const logger = createLogger("omnigent:gateway"); + +/** + * Paths a harness may reach through the gateway: model calls and the read-only + * model listings Omnigent's catalog uses. Everything else is refused, so a + * process that finds the placeholder token can call models and nothing more. + */ +const DEFAULT_GATEWAY_ALLOW: readonly RegExp[] = [ + /^\/ai-gateway\//, + /^\/serving-endpoints\//, + /^\/api\/2\.0\/serving-endpoints(\/|\?|$)/, + /^\/api\/2\.1\/unity-catalog\/model-services(\/|\?|$)/, +]; + +interface GatewayOptions { + /** The real workspace, e.g. `https://my-ws.cloud.databricks.com`. */ + upstream: string; + /** Returns a bearer token for the identity model calls are billed to. */ + token: () => Promise; + /** TLS material for the https listener (Omnigent's token helper only serves https hosts). */ + tls: { key: Buffer; cert: Buffer }; + /** Loopback port; 0 picks a free one. */ + port?: number; + allow?: readonly RegExp[]; +} + +interface GatewayStats { + forwarded: number; + denied: number; +} + +/** + * A loopback stand-in for the workspace that Omnigent's children use as their + * Databricks host. They hold only a placeholder token; the gateway forwards + * allowed paths to the real workspace with the real token. No child ever sees + * a credential that works anywhere else. + */ +export class ModelGateway { + readonly placeholder = `appkit-omnigent-${randomBytes(18).toString("hex")}`; + readonly stats: GatewayStats = { forwarded: 0, denied: 0 }; + private readonly allow: readonly RegExp[]; + private readonly upstream: URL; + private server?: https.Server; + private port = 0; + + constructor(private readonly opts: GatewayOptions) { + this.allow = opts.allow ?? DEFAULT_GATEWAY_ALLOW; + this.upstream = new URL(opts.upstream); + } + + /** The gateway's URL (https: Omnigent's token helper only serves https hosts). */ + get httpsUrl(): string { + return `https://127.0.0.1:${this.port}`; + } + + async start(): Promise { + const handler = (req: http.IncomingMessage, res: http.ServerResponse) => { + void this.handle(req, res); + }; + this.server = https.createServer(this.opts.tls, handler); + this.port = await listen(this.server, this.opts.port ?? 0); + logger.debug("model gateway on %s", this.httpsUrl); + } + + async stop(): Promise { + const s = this.server; + if (!s) return; + await new Promise((resolve) => { + s.close(() => resolve()); + s.closeAllConnections(); + }); + this.server = undefined; + } + + /** Whether a request may be forwarded. Exposed for tests. */ + permits(path: string, authorization: string | undefined): boolean { + if (authorization && authorization !== `Bearer ${this.placeholder}`) + return false; + return this.allow.some((re) => re.test(path)); + } + + private async handle(req: http.IncomingMessage, res: http.ServerResponse) { + const path = req.url ?? "/"; + // Host metadata: answer "none" so the SDK stays on plain token auth. + if (path.startsWith("/.well-known/")) { + res.writeHead(404, { "content-type": "application/json" }).end("{}"); + return; + } + if (!this.permits(path, req.headers.authorization)) { + this.stats.denied++; + res.writeHead(403, { "content-type": "application/json" }).end( + JSON.stringify({ + error_code: "PERMISSION_DENIED", + message: + "This path is not available through the app's model gateway.", + }), + ); + req.resume(); + return; + } + let token: string; + try { + token = await this.opts.token(); + } catch (err) { + logger.error("model gateway could not get a token: %O", err); + res.writeHead(502).end(); + req.resume(); + return; + } + const headers = { + ...req.headers, + host: this.upstream.host, + authorization: `Bearer ${token}`, + }; + const upstream = https.request( + { + hostname: this.upstream.hostname, + port: this.upstream.port || 443, + path, + method: req.method, + headers, + }, + (up) => { + this.stats.forwarded++; + res.writeHead(up.statusCode ?? 502, up.headers); + up.pipe(res); + }, + ); + upstream.on("error", (err) => { + logger.warn("model gateway upstream error: %s", err.message); + if (!res.headersSent) res.writeHead(502); + res.end(); + }); + req.pipe(upstream); + } +} + +function listen(server: https.Server, port: number): Promise { + return new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(port, "127.0.0.1", () => { + server.off("error", reject); + resolve((server.address() as AddressInfo).port); + }); + }); +} diff --git a/packages/appkit/src/plugins/omnigent/runtime/python/watchdog.py b/packages/appkit/src/plugins/omnigent/runtime/python/watchdog.py new file mode 100644 index 000000000..b84d22821 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/runtime/python/watchdog.py @@ -0,0 +1,43 @@ +"""Run a Python entrypoint and take its whole process group down if the Node parent goes away. + + python watchdog.py [args...] + +The AppKit omnigent plugin starts each Omnigent process in its own process group. If Node dies without running its +shutdown hook (a crash, SIGKILL, or a platform stop that never reached it), this thread notices the parent change and +kills the group, so no Omnigent process outlives the app and blocks the next deployment. +""" +import os +import runpy +import signal +import sys +import threading +import time + +PARENT = os.getppid() + + +def _watch(): + while True: + if os.getppid() != PARENT: + try: + os.killpg(os.getpgid(0), signal.SIGTERM) + time.sleep(3) + os.killpg(os.getpgid(0), signal.SIGKILL) + finally: + os._exit(1) + time.sleep(1) + + +def main(): + threading.Thread(target=_watch, daemon=True, name="parent-watchdog").start() + target = sys.argv[1] + sys.argv = sys.argv[1:] + if target.endswith(".py"): + sys.path.insert(0, os.path.dirname(os.path.abspath(target))) + runpy.run_path(target, run_name="__main__") + else: + runpy.run_module(target, run_name="__main__", alter_sys=True) + + +if __name__ == "__main__": + main() diff --git a/packages/appkit/src/plugins/omnigent/runtime/runtime.ts b/packages/appkit/src/plugins/omnigent/runtime/runtime.ts new file mode 100644 index 000000000..fbc4ee209 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/runtime/runtime.ts @@ -0,0 +1,373 @@ +import { type ChildProcess, spawn } from "node:child_process"; +import { createHash, randomBytes } from "node:crypto"; +import fs from "node:fs"; +import net from "node:net"; +import path from "node:path"; +import { fileURLToPath } from "node:url"; + +import { createLogger } from "../../../logging"; +import type { OmnigentStatus } from "../types"; +import { RuntimeLayout, childEnv, writeUserConfig } from "./env"; +import { ModelGateway } from "./gateway"; +import { ensureGatewayTls, type GatewayTls } from "./tls"; + +const logger = createLogger("omnigent:runtime"); + +const PYTHON_DIR = path.join( + path.dirname(fileURLToPath(import.meta.url)), + "python", +); + +interface RuntimeOptions { + root: string; + appDir: string; + python: string; + /** The real workspace URL. */ + upstream: string; + /** Bearer for model calls (the app's service principal). */ + token: () => Promise; + hostIdleMs: number; +} + +interface HostEntry { + proc: ChildProcess; + hostId: string; + userId: string; + ready: Promise; + /** Set when the host never came online; the entry is then replaced, not reused. */ + failed?: boolean; + lastUsed: number; +} + +/** + * Omnigent embedded in the app: one server on loopback in header-auth mode + * (under a random header name) and one host per active user, supervised as + * child processes. + */ +export class OmnigentRuntime { + readonly layout: RuntimeLayout; + /** Header the server trusts for identity. Random, so no other process can guess it. */ + readonly authHeader = `x-appkit-omnigent-user-${randomBytes(12).toString("hex")}`; + private tls!: GatewayTls; + private gateway!: ModelGateway; + private server?: ChildProcess; + private serverReady?: Promise; + private readonly hosts = new Map(); + private port = 0; + private stopping = false; + private reaper?: NodeJS.Timeout; + private lastError?: string; + private restartDelayMs = 3000; + private version?: string; + + constructor(private readonly opts: RuntimeOptions) { + this.layout = new RuntimeLayout(opts.root); + } + + get url(): string { + return `http://127.0.0.1:${this.port}`; + } + + async start(): Promise { + this.layout.create(); + this.tls = ensureGatewayTls(this.layout.tls, this.layout.bin); + this.gateway = new ModelGateway({ + upstream: this.opts.upstream, + token: this.opts.token, + tls: { key: this.tls.key, cert: this.tls.cert }, + }); + await this.gateway.start(); + + this.port = await freePort(); + this.startServer(); + this.reaper = setInterval(() => this.reapIdleHosts(), 60_000); + this.reaper.unref(); + } + + /** Resolves once the server answers `/health`. */ + whenReady(): Promise { + return ( + this.serverReady ?? + Promise.reject(new Error("Omnigent runtime not started")) + ); + } + + private startServer(): void { + const args = [ + this.opts.python, + path.join(PYTHON_DIR, "watchdog.py"), + "omnigent", + "server", + "--host", + "127.0.0.1", + "-p", + String(this.port), + "--no-open", + ]; + const home = path.join(this.layout.serverDir, "home"); + fs.mkdirSync(home, { recursive: true }); + writeUserConfig(home, this.gateway); + const env = childEnv(process.env, { + home, + dataDir: path.join(this.layout.serverDir, "data"), + gateway: this.gateway, + caBundle: this.tls.bundleFile, + caCert: this.tls.certFile, + pathPrepend: [path.dirname(this.opts.python)], + extra: { + OMNIGENT_AUTH_PROVIDER: "header", + OMNIGENT_AUTH_HEADER: this.authHeader, + }, + }); + this.server = this.spawn("server", args, env, this.layout.serverDir); + this.server.once("exit", (code, signal) => { + if (this.stopping) return; + this.lastError = `Omnigent server exited (${code ?? signal}); restarting`; + logger.warn("%s. Log: %s", this.lastError, this.logFile("server")); + const delay = this.restartDelayMs; + this.restartDelayMs = Math.min(delay * 2, 60_000); + setTimeout(() => !this.stopping && this.startServer(), delay).unref(); + }); + this.serverReady = this.waitFor( + async () => (await fetch(`${this.url}/health`)).ok, + 10 * 60_000, + "Omnigent server", + ).then(async () => { + this.lastError = undefined; + this.restartDelayMs = 3000; + try { + const info = await this.request("GET", "/v1/info", "appkit-runtime"); + this.version = ( + (await info.json()) as { server_version?: string } + ).server_version; + } catch {} + logger.info("Omnigent server ready on %s", this.url); + }); + this.serverReady.catch((e) => { + this.lastError = String(e); + }); + } + + /** Starts (or reuses) the host that runs this user's sessions. */ + async ensureHost(email: string): Promise { + await this.whenReady(); + const current = this.hosts.get(email); + if (current?.failed) { + // Alive but never registered: take it down and start a fresh host. + killGroup(current.proc, "SIGKILL"); + this.hosts.delete(email); + } else if ( + current && + current.proc.exitCode === null && + current.proc.signalCode === null + ) { + current.lastUsed = Date.now(); + return current.ready; + } + const userId = RuntimeLayout.userId(email); + // Stable per app and user, so a thread from before a restart re-binds to its host. + const hostId = createHash("sha256") + .update( + `${this.opts.upstream}/apps/${process.env.DATABRICKS_APP_NAME ?? "appkit"}/${email.toLowerCase()}`, + ) + .digest("hex") + .slice(0, 32); + const home = this.layout.userHome(userId); + const work = this.layout.userWork(userId); + for (const d of [home, this.layout.userData(userId), work]) + fs.mkdirSync(d, { recursive: true }); + writeUserConfig(home, this.gateway); + const env = childEnv(process.env, { + home, + dataDir: this.layout.userData(userId), + gateway: this.gateway, + caBundle: this.tls.bundleFile, + caCert: this.tls.certFile, + pathPrepend: [ + path.dirname(this.opts.python), + path.join(this.opts.appDir, "node_modules", ".bin"), + ], + extra: { + OMNIGENT_HOST_ID: hostId, + OMNIGENT_HOST_NAME: `appkit-${userId}`, + OMNIGENT_DATABRICKS_EXTRA_HEADERS: JSON.stringify({ + [this.authHeader]: email, + }), + }, + }); + const proc = this.spawn( + `host-${userId}`, + [ + this.opts.python, + path.join(PYTHON_DIR, "watchdog.py"), + "omnigent", + "host", + "--server", + this.url, + ], + env, + ); + const entry: HostEntry = { + proc, + hostId, + userId, + lastUsed: Date.now(), + ready: undefined as never, + }; + entry.ready = this.waitFor( + async () => { + const r = await this.request("GET", "/v1/hosts", email); + const { hosts = [] } = (await r.json()) as { + hosts?: Array<{ host_id: string; status: string }>; + }; + return hosts.some((h) => h.host_id === hostId && h.status === "online") + ? entry + : undefined; + }, + 120_000, + "Omnigent host", + ); + entry.ready.catch(() => { + entry.failed = true; + }); + this.hosts.set(email, entry); + proc.once("exit", (code, signal) => { + if (!this.stopping && code !== 0 && signal !== "SIGTERM") { + logger.warn( + "Omnigent host for %s exited (%s). Log: %s", + userId, + code ?? signal, + this.logFile(`host-${userId}`), + ); + } + if (this.hosts.get(email) === entry) this.hosts.delete(email); + }); + return entry.ready; + } + + /** A request to the Omnigent server as `email`. */ + request( + method: string, + p: string, + email: string, + init: { + json?: unknown; + body?: RequestInit["body"]; + headers?: Record; + signal?: AbortSignal; + } = {}, + ): Promise { + return fetch(this.url + p, { + method, + headers: { + [this.authHeader]: email, + ...(init.json !== undefined + ? { "content-type": "application/json" } + : {}), + ...init.headers, + }, + body: init.json !== undefined ? JSON.stringify(init.json) : init.body, + signal: init.signal, + }); + } + + status(): OmnigentStatus { + return { + ready: Boolean( + this.server && this.lastError === undefined && this.version, + ), + error: this.lastError, + version: this.version, + hosts: this.hosts.size, + gateway: this.gateway?.stats ?? { forwarded: 0, denied: 0 }, + }; + } + + async stop(): Promise { + this.stopping = true; + if (this.reaper) clearInterval(this.reaper); + const procs = [ + this.server, + ...[...this.hosts.values()].map((h) => h.proc), + ].filter((p): p is ChildProcess => Boolean(p?.pid) && p?.exitCode === null); + for (const p of procs) killGroup(p, "SIGTERM"); + await Promise.race([ + Promise.all(procs.map((p) => new Promise((r) => p.once("exit", r)))), + new Promise((r) => setTimeout(r, 4000)), + ]); + for (const p of procs) killGroup(p, "SIGKILL"); + await this.gateway?.stop(); + } + + private reapIdleHosts(): void { + const cutoff = Date.now() - this.opts.hostIdleMs; + for (const [email, h] of this.hosts) { + if (h.lastUsed < cutoff) { + logger.debug("stopping idle host for %s", h.userId); + killGroup(h.proc, "SIGTERM"); + this.hosts.delete(email); + } + } + } + + private logFile(name: string): string { + return path.join(this.layout.logs, `${name}.log`); + } + + private spawn( + name: string, + argv: string[], + env: Record, + cwd: string = this.layout.root, + ): ChildProcess { + const log = fs.openSync(this.logFile(name), "a"); + const child = spawn(argv[0], argv.slice(1), { + env, + cwd, + stdio: ["ignore", log, log], + detached: true, // own process group, so stop() can take the whole tree down + }); + fs.closeSync(log); + child.on("error", (err) => { + this.lastError = `${name}: ${err.message}`; + logger.error("failed to start %s: %s", name, err.message); + }); + return child; + } + + private async waitFor( + fn: () => Promise, + ms: number, + what: string, + ): Promise { + const end = Date.now() + ms; + while (Date.now() < end && !this.stopping) { + try { + const v = await fn(); + if (v) return v; + } catch {} + await new Promise((r) => setTimeout(r, 1000)); + } + throw new Error( + `${what} did not come up within ${Math.round(ms / 1000)} s`, + ); + } +} + +function killGroup(p: ChildProcess, signal: NodeJS.Signals): void { + if (!p.pid) return; + try { + process.kill(-p.pid, signal); + } catch {} +} + +function freePort(): Promise { + return new Promise((resolve, reject) => { + const s = net.createServer(); + s.once("error", reject); + s.listen(0, "127.0.0.1", () => { + const { port } = s.address() as net.AddressInfo; + s.close(() => resolve(port)); + }); + }); +} diff --git a/packages/appkit/src/plugins/omnigent/runtime/tls.ts b/packages/appkit/src/plugins/omnigent/runtime/tls.ts new file mode 100644 index 000000000..3fe9e9718 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/runtime/tls.ts @@ -0,0 +1,66 @@ +import { execFileSync } from "node:child_process"; +import fs from "node:fs"; +import path from "node:path"; + +export interface GatewayTls { + key: Buffer; + cert: Buffer; + /** The certificate alone (for `NODE_EXTRA_CA_CERTS`). */ + certFile: string; + /** System roots plus the certificate (for `SSL_CERT_FILE` / `REQUESTS_CA_BUNDLE`). */ + bundleFile: string; +} + +const SYSTEM_ROOTS = [ + "/etc/ssl/certs/ca-certificates.crt", + "/etc/pki/tls/certs/ca-bundle.crt", + "/etc/ssl/cert.pem", +]; + +/** + * A self-signed certificate for `127.0.0.1`, made once per runtime directory + * with the system `openssl`, and a CA bundle that trusts it alongside the + * system roots. The key stays in `privateDir`; the certificate and + * bundle go to `publicDir`, which the harnesses read. + */ +export function ensureGatewayTls( + privateDir: string, + publicDir: string, +): GatewayTls { + fs.mkdirSync(privateDir, { recursive: true, mode: 0o700 }); + fs.mkdirSync(publicDir, { recursive: true }); + const keyFile = path.join(privateDir, "key.pem"); + const certFile = path.join(publicDir, "gateway-cert.pem"); + if (!fs.existsSync(certFile) || !fs.existsSync(keyFile)) { + execFileSync( + "openssl", + [ + "req", + "-x509", + "-newkey", + "rsa:2048", + "-nodes", + "-days", + "30", + "-subj", + "/CN=127.0.0.1", + "-addext", + "subjectAltName=IP:127.0.0.1", + "-keyout", + keyFile, + "-out", + certFile, + ], + { stdio: "ignore" }, + ); + fs.chmodSync(keyFile, 0o600); + } + const cert = fs.readFileSync(certFile); + const roots = SYSTEM_ROOTS.find((f) => fs.existsSync(f)); + const bundleFile = path.join(publicDir, "ca-bundle.pem"); + fs.writeFileSync( + bundleFile, + `${roots ? fs.readFileSync(roots, "utf8") : ""}\n${cert.toString("utf8")}`, + ); + return { key: fs.readFileSync(keyFile), cert, certFile, bundleFile }; +} diff --git a/packages/appkit/src/plugins/omnigent/tests/gateway.test.ts b/packages/appkit/src/plugins/omnigent/tests/gateway.test.ts new file mode 100644 index 000000000..3d4817d0d --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/tests/gateway.test.ts @@ -0,0 +1,104 @@ +import fs from "node:fs"; +import https from "node:https"; +import os from "node:os"; +import path from "node:path"; + +import { afterAll, beforeAll, describe, expect, test } from "vitest"; + +import { ModelGateway } from "../runtime/gateway"; +import { ensureGatewayTls } from "../runtime/tls"; + +function get( + url: string, + ca: Buffer, + headers: Record = {}, +): Promise<{ status: number; body: string }> { + return new Promise((resolve, reject) => { + https + .get(url, { headers, ca }, (res) => { + let body = ""; + res.on("data", (c) => (body += c)); + res.on("end", () => resolve({ status: res.statusCode ?? 0, body })); + }) + .on("error", reject); + }); +} + +describe("ModelGateway", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "omni-gw-")); + let gw: ModelGateway; + let cert: Buffer; + let tokenCalls = 0; + + beforeAll(async () => { + const tls = ensureGatewayTls( + path.join(root, "tls"), + path.join(root, "bin"), + ); + cert = tls.cert; + gw = new ModelGateway({ + upstream: "https://127.0.0.1:9", // never reached in these tests + token: async () => { + tokenCalls++; + return "real-token"; + }, + tls, + }); + await gw.start(); + }); + + afterAll(async () => { + await gw.stop(); + fs.rmSync(root, { recursive: true, force: true }); + }); + + test("TLS material: key private, cert and bundle public", () => { + expect(fs.existsSync(path.join(root, "tls", "key.pem"))).toBe(true); + expect( + fs.readFileSync(path.join(root, "bin", "ca-bundle.pem"), "utf8"), + ).toContain("BEGIN CERTIFICATE"); + expect( + fs.readdirSync(path.join(root, "bin")).some((f) => f.includes("key")), + ).toBe(false); + }); + + test("permits model paths with the placeholder or no auth, nothing else", () => { + const ph = `Bearer ${gw.placeholder}`; + expect(gw.permits("/ai-gateway/anthropic/v1/messages", ph)).toBe(true); + expect( + gw.permits( + "/serving-endpoints/databricks-gpt-5-5/invocations", + undefined, + ), + ).toBe(true); + expect(gw.permits("/api/2.0/serving-endpoints", ph)).toBe(true); + expect( + gw.permits("/api/2.1/unity-catalog/model-services?parent=x", ph), + ).toBe(true); + expect(gw.permits("/api/2.0/clusters/list", ph)).toBe(false); + expect(gw.permits("/api/2.0/sql/statements", ph)).toBe(false); + expect(gw.permits("/api/2.1/unity-catalog/tables", ph)).toBe(false); + expect( + gw.permits("/ai-gateway/anthropic/v1/messages", "Bearer something-else"), + ).toBe(false); + }); + + test("refuses non-model paths with 403 and never mints a token for them", async () => { + const before = tokenCalls; + const r = await get(`${gw.httpsUrl}/api/2.0/clusters/list`, cert, { + authorization: `Bearer ${gw.placeholder}`, + }); + expect(r.status).toBe(403); + expect(tokenCalls).toBe(before); + expect(gw.stats.denied).toBeGreaterThan(0); + }); + + test("answers host metadata with 404 so SDKs stay on token auth", async () => { + const r = await get(`${gw.httpsUrl}/.well-known/databricks-config`, cert); + expect(r.status).toBe(404); + }); + + test("listens on loopback only", () => { + expect(gw.httpsUrl).toMatch(/^https:\/\/127\.0\.0\.1:\d+$/); + }); +}); 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..797eb0365 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/tests/plugin.test.ts @@ -0,0 +1,73 @@ +import { + createMockRequest, + createMockResponse, + createMockRouter, + mockServiceContext, + setupDatabricksEnv, +} from "@tools/test-helpers"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +import { ServiceContext } from "../../../context/service-context"; +import { OmnigentPlugin } from "../omnigent"; + +function plugin() { + const p = new OmnigentPlugin({}); + 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; +} + +describe("OmnigentPlugin routes", () => { + let ctx: Awaited>; + + beforeEach(async () => { + setupDatabricksEnv(); + ServiceContext.reset(); + ctx = await mockServiceContext({}); + }); + + afterEach(() => { + ctx?.restore(); + 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"], + ["GET", "/sessions/:id/stream"], + ]) { + expect(getHandler(m, path), `${m} ${path}`).toBeTypeOf("function"); + } + }); + + test("session routes need a forwarded identity", async () => { + const { getHandler } = plugin(); + const res = await call(getHandler("GET", "/sessions"), {}); + expect(res.status).toHaveBeenCalledWith(401); + }); + + test("a signed-in request before the runtime starts gets 503, not a crash", async () => { + const { getHandler } = plugin(); + const res = await call(getHandler("GET", "/sessions"), { + obo: { token: "t", userId: "u", email: "a@example.com" }, + }); + expect(res.status).toHaveBeenCalledWith(503); + }); +}); 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..25f379a84 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/tests/units.test.ts @@ -0,0 +1,222 @@ +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; +import { gunzipSync } from "node:zlib"; + +import { afterEach, describe, expect, test } from "vitest"; + +import { buildBundle } from "../bundle"; +import { offerHarnesses } from "../harnesses"; +import { childEnv, GATEWAY_PROFILE, writeUserConfig } from "../runtime/env"; + +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: "appkit", + harness: "pi", + model: "system.ai.gpt-5-5", + profile: GATEWAY_PROFILE, + instructions: "Be terse.", + }), + ), + ); + expect(Object.keys(files).sort()).toEqual(["AGENTS.md", "config.yaml"]); + const cfg = JSON.parse(files["config.yaml"]); + expect(cfg.executor).toMatchObject({ + model: "system.ai.gpt-5-5", + profile: "appkit-omnigent", + config: { harness: "pi" }, + auth: { type: "databricks", profile: "appkit-omnigent" }, + }); + expect(files["AGENTS.md"]).toBe("Be terse."); + }); +}); + +describe("child environment", () => { + const dirs: string[] = []; + afterEach(() => { + for (const d of dirs.splice(0)) + fs.rmSync(d, { recursive: true, force: true }); + }); + + const gw = { + httpsUrl: "https://127.0.0.1:1", + placeholder: "appkit-omnigent-ph", + }; + + test("no credential reaches a child", () => { + const env = childEnv( + { + PATH: "/usr/bin", + DATABRICKS_CLIENT_ID: "id", + DATABRICKS_CLIENT_SECRET: "secret", + DATABRICKS_TOKEN: "dapi", + DATABRICKS_CONFIG_PROFILE: "prod", + PGPASSWORD: "pw", + CLAUDECODE: "1", + NODE_OPTIONS: "--require x", + KEEP_ME: "yes", + }, + { + home: "/h", + dataDir: "/d", + gateway: gw, + caBundle: "/b", + caCert: "/c", + pathPrepend: ["/venv/bin"], + }, + ); + const leaked = Object.values(env).filter((v) => /secret|dapi|pw/.test(v)); + expect(leaked).toEqual([]); + for (const k of [ + "DATABRICKS_CLIENT_ID", + "DATABRICKS_CONFIG_PROFILE", + "CLAUDECODE", + "NODE_OPTIONS", + ]) { + expect(env[k]).toBeUndefined(); + } + expect(env).toMatchObject({ + DATABRICKS_HOST: gw.httpsUrl, + HOME: "/h", + KEEP_ME: "yes", + SSL_CERT_FILE: "/b", + }); + expect(env.PATH).toBe(`/venv/bin${path.delimiter}/usr/bin`); + }); + + test("extra can add and remove variables", () => { + const env = childEnv( + { A: "1" }, + { + home: "/h", + dataDir: "/d", + gateway: gw, + caBundle: "", + caCert: "", + pathPrepend: [], + extra: { A: undefined, B: "2" }, + }, + ); + expect(env.A).toBeUndefined(); + expect(env.B).toBe("2"); + }); + + test("user config holds only the placeholder", () => { + const home = fs.mkdtempSync(path.join(os.tmpdir(), "omni-home-")); + dirs.push(home); + writeUserConfig(home, gw); + const cfg = fs.readFileSync(path.join(home, ".databrickscfg"), "utf8"); + expect(cfg).toContain(`host = ${gw.httpsUrl}`); + expect(cfg.match(/token = (.*)/g)).toEqual([`token = ${gw.placeholder}`]); + expect( + (fs.statSync(path.join(home, ".databrickscfg")).mode & 0o777).toString(8), + ).toBe("600"); + const omni = fs.readFileSync( + path.join(home, ".config", "omnigent", "config.yaml"), + "utf8", + ); + expect(omni).toContain(`profile: ${GATEWAY_PROFILE}`); + }); +}); + +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"], + }, + ]; + const all = { + "claude-sdk": true, + pi: "needs-auth", + "openai-agents": true, + } as const; + + test("matches UC models to harnesses by the API each harness calls", () => { + const offers = offerHarnesses({ + allowed: Object.keys(all), + configured: all, + 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["openai-agents"].models).toEqual([ + "system.ai.gpt-5-5", + "system.ai.grok-4-7", + ]); + // Pi: the Codex responses API only, and never Claude + expect(by.pi.models).toEqual(["system.ai.gpt-5-5", "system.ai.glm-5-3"]); + 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", () => { + const offers = offerHarnesses({ + allowed: ["claude-sdk", "cursor"], + configured: { "claude-sdk": true, cursor: true }, + models: catalog, + }); + expect(offers.map((o) => o.id)).toEqual(["claude-sdk"]); + }); + + test("drops harnesses whose binary is missing", () => { + const offers = offerHarnesses({ + allowed: ["claude-sdk", "pi"], + configured: { "claude-sdk": true, pi: "binary-missing" }, + models: catalog, + }); + expect(offers.map((o) => o.id)).toEqual(["claude-sdk"]); + }); +}); diff --git a/packages/appkit/src/plugins/omnigent/types.ts b/packages/appkit/src/plugins/omnigent/types.ts new file mode 100644 index 000000000..e871f649b --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/types.ts @@ -0,0 +1,47 @@ +import type { BasePluginConfig } from "shared"; + +export interface IOmnigentConfig extends BasePluginConfig { + /** System instructions for every session (the bundle's `AGENTS.md`). */ + instructions?: string; + /** + * Harnesses users may pick, out of the ones that route model calls through + * Unity AI Gateway. Default: all of them whose CLI is installed. + */ + harnesses?: string[]; + /** Default harness and model for new sessions. */ + defaultHarness?: string; + defaultModel?: string; + runtime?: { + /** Python with `omnigent[databricks]` installed. Default: `./.venv/bin/python`. */ + python?: string; + /** Private state directory. Default: `$TMPDIR/appkit-omnigent`. */ + home?: string; + /** Stop a user's host after this long without requests. Default 30 min. */ + hostIdleMs?: number; + }; +} + +export interface OmnigentHarness { + id: string; + label: string; + ready: boolean | string; +} + +export interface CreateSessionRequest { + harness?: string; + model?: string; + message?: string; + title?: string; +} + +/** Runtime health, from `GET /status` and the plugin's `status()` export. */ +export interface OmnigentStatus { + ready: boolean; + error?: string; + /** Omnigent server version. */ + version?: string; + /** Running per-user hosts. */ + hosts: number; + /** Model gateway counters. */ + gateway: { forwarded: number; denied: number }; +} diff --git a/packages/appkit/tsdown.config.ts b/packages/appkit/tsdown.config.ts index bb17c29ba..b390305e7 100644 --- a/packages/appkit/tsdown.config.ts +++ b/packages/appkit/tsdown.config.ts @@ -38,6 +38,11 @@ export default defineConfig([ }, tsconfig: "./tsconfig.json", copy: [ + { + from: "src/plugins/omnigent/runtime/python/*.py", + to: "dist/plugins/omnigent/runtime/python", + flatten: true, + }, { from: "src/plugins/server/remote-tunnel/*.html", to: "dist/plugins/server/remote-tunnel", diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index 61a3e0b12..07ffbdc6a 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -342,6 +342,26 @@ } } }, + "omnigent": { + "name": "omnigent", + "displayName": "Omnigent Plugin", + "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway", + "package": "@databricks/appkit", + "resources": { + "required": [], + "optional": [] + }, + "stability": "beta", + "scaffolding": { + "rules": { + "must": [ + "Add pyproject.toml (python >=3.12,<3.13) with omnigent[databricks] and a uv.lock that uses public PyPI", + "Add harness CLIs to package.json: @anthropic-ai/claude-code, @earendil-works/pi-coding-agent", + "Start the app with node, not npm run, so SIGTERM reaches the plugin and it can stop its child processes" + ] + } + } + }, "server": { "name": "server", "displayName": "Server Plugin", From 27a8d0252283395c56c4e19aebf61591c7327dfe Mon Sep 17 00:00:00 2001 From: Noah Solomon Date: Wed, 30 Sep 2026 00:05:05 +1000 Subject: [PATCH 2/4] feat(appkit): offer Codex in the omnigent plugin, sandboxed with bubblewrap (2/4) Codex has its own shell, so it runs through a bubblewrap wrapper: own PID namespace, the plugin's state directory masked, only the user's home and work directory bound in, and /proc of the app's processes hidden. The wrapper is named codex and put first on PATH, because Omnigent resolves harness CLIs by name. Without bubblewrap, Codex is not offered. Codex only trusts public TLS roots, so the model gateway also listens on plain http for it. Co-authored-by: Isaac Signed-off-by: Noah Solomon --- docs/docs/plugins/omnigent.md | 21 +++-- .../appkit/src/plugins/omnigent/harnesses.ts | 17 +++- .../appkit/src/plugins/omnigent/manifest.json | 4 +- .../appkit/src/plugins/omnigent/omnigent.ts | 9 ++- .../src/plugins/omnigent/runtime/codex.ts | 59 ++++++++++++++ .../src/plugins/omnigent/runtime/env.ts | 32 +++++--- .../src/plugins/omnigent/runtime/gateway.ts | 51 +++++++----- .../src/plugins/omnigent/runtime/runtime.ts | 39 ++++++++++ .../src/plugins/omnigent/runtime/sandbox.ts | 78 +++++++++++++++++++ .../src/plugins/omnigent/runtime/tls.ts | 4 +- .../plugins/omnigent/tests/gateway.test.ts | 14 ++-- .../src/plugins/omnigent/tests/units.test.ts | 65 +++++++++++++--- packages/appkit/src/plugins/omnigent/types.ts | 4 + template/appkit.plugins.json | 4 +- 14 files changed, 335 insertions(+), 66 deletions(-) create mode 100644 packages/appkit/src/plugins/omnigent/runtime/codex.ts create mode 100644 packages/appkit/src/plugins/omnigent/runtime/sandbox.ts diff --git a/docs/docs/plugins/omnigent.md b/docs/docs/plugins/omnigent.md index 09823ba9a..defc68959 100644 --- a/docs/docs/plugins/omnigent.md +++ b/docs/docs/plugins/omnigent.md @@ -6,7 +6,7 @@ This plugin is currently **beta**. APIs may change between minor releases. Impor ::: -The `omnigent` plugin embeds [Omnigent](https://github.com/omnigent-ai/omnigent), the open-source agent meta-harness, in a Databricks AppKit app. Your users get Claude Agent SDK, Pi and OpenAI Agents SDK sessions inside the app, with model calls going through Unity AI Gateway as the app's service principal. +The `omnigent` plugin embeds [Omnigent](https://github.com/omnigent-ai/omnigent), the open-source agent meta-harness, in a Databricks AppKit app. Your users get Claude Agent SDK, Codex, Pi and OpenAI Agents SDK sessions inside the app, with model calls going through Unity AI Gateway as the app's service principal. Omnigent owns the agent loop: sessions, history, approvals. The plugin runs it safely inside a Databricks App. 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. @@ -27,7 +27,7 @@ await createApp({ The app also needs: - **Python.** Add a `pyproject.toml` with `requires-python = ">=3.12,<3.13"` and `omnigent[databricks]`, plus a `uv.lock`. Databricks Apps then builds a Python 3.12 venv next to `npm install`. A plain `requirements.txt` gets Python 3.11, which Omnigent does not support. The lockfile must reference public PyPI. -- **Harness CLIs.** Add them to your `package.json`, so you pin their versions: `@anthropic-ai/claude-code` and `@earendil-works/pi-coding-agent`. The plugin only offers harnesses whose CLI is installed. The Claude Agent SDK's bundled CLI is too old for current Claude models, so add `@anthropic-ai/claude-code` even though the SDK ships one. +- **Harness CLIs.** Add them to your `package.json`, so you pin their versions: `@anthropic-ai/claude-code`, `@openai/codex` and `@earendil-works/pi-coding-agent`. The plugin only offers harnesses whose CLI is installed. The Claude Agent SDK's bundled CLI is too old for current Claude models, so add `@anthropic-ai/claude-code` even though the SDK ships one. - **Start with `node`.** In `app.yaml`, use `command: ["node", "server.js"]`, not `npm run start`. npm does not pass SIGTERM on, so the plugin never gets to stop its child processes. As a backstop, each child also exits when Node does, so a redeploy still succeeds. ## Configuration @@ -35,13 +35,13 @@ The app also needs: | Option | Default | | |---|---|---| | `instructions` | a generic assistant prompt | System instructions for every session | -| `harnesses` | `claude-sdk`, `pi`, `openai-agents` | Harnesses users may pick | +| `harnesses` | `claude-sdk`, `codex`, `pi`, `openai-agents` | Harnesses users may pick | | `defaultHarness`, `defaultModel` | first available | | | `runtime.python` | `./.venv/bin/python` | | | `runtime.home` | `$TMPDIR/appkit-omnigent` | The plugin's state directory | | `runtime.hostIdleMs` | 30 min | Stop a user's Omnigent host after this long idle | -Each harness is offered the Unity AI 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. +Each harness is offered the Unity AI Gateway `system.ai` models whose API it speaks: Claude Agent SDK the Anthropic Messages API, Codex and Pi the Codex Responses API, OpenAI Agents SDK the OpenAI Responses API. ## Routes @@ -55,15 +55,22 @@ All routes are under `/api/omnigent` and scoped to the signed-in user. | `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 | -| `GET /status` | Runtime health: hosts, gateway counters | +| `GET /status` | Runtime health: sandbox, hosts, gateway counters | ## How it works - **One Omnigent server, one host per user.** The server listens on loopback in Omnigent's header-auth mode, under a random header name only the plugin knows. Each active user gets their own host process, registered as that user; Omnigent only runs a session on its owner's host. An idle host stops after `runtime.hostIdleMs`. -- **No credential in any child process.** The harnesses' Databricks profile points at a model gateway inside the plugin, on loopback, with a placeholder token. The gateway forwards only model-serving paths to the workspace and adds the service principal's token. Every other path returns 403. +- **No credential in any child process.** The harnesses' Databricks profile points at a model gateway inside the plugin, on loopback, with a placeholder token. The gateway forwards only model-serving paths to the workspace and adds the service principal's token. Every other path returns 403. Codex only trusts public TLS roots, so it reaches the gateway over plain http on loopback; the others use https with a certificate the plugin generates. +- **Codex runs sandboxed.** Codex has its own shell. Its binary runs through a bubblewrap wrapper that gives it: + - its own PID namespace + - the plugin's state directory masked + - only the user's home and work directory bound in + - the `/proc` entries of the app's own processes hidden + + Without bubblewrap on `PATH` (or bundled with Codex), Codex is not offered. ## Limits - **Sessions live in the container.** Omnigent keeps them in SQLite under the plugin's state directory, so they survive a process restart but not a redeploy. -- **Only gateway harnesses.** Harnesses that sign in with their vendor (Cursor, Copilot, Gemini and others) bypass Unity AI Gateway, so they are not offered. Native terminal harnesses (`claude-native`, `pi-native`) are not offered yet: they need tmux and a terminal client. +- **Only gateway harnesses.** Harnesses that sign in with their vendor (Cursor, Copilot, Gemini and others) bypass Unity AI Gateway, so they are not offered. Native terminal harnesses (`claude-native`, `codex-native`, `pi-native`) are not offered yet: they need tmux and a terminal client. - **Shared gateway limits.** All users' model traffic uses the service principal's gateway limits. diff --git a/packages/appkit/src/plugins/omnigent/harnesses.ts b/packages/appkit/src/plugins/omnigent/harnesses.ts index 3a726f90c..096ed11de 100644 --- a/packages/appkit/src/plugins/omnigent/harnesses.ts +++ b/packages/appkit/src/plugins/omnigent/harnesses.ts @@ -2,6 +2,8 @@ import type { OmnigentHarness } from "./types"; interface HarnessInfo { label: string; + /** Runs its own shell (so it needs the sandbox). */ + shell: boolean; /** * Unity AI Gateway API types it calls (a UC model service's * `supported_api_types`); a model must support one of them. @@ -20,11 +22,19 @@ interface HarnessInfo { const GATEWAY_HARNESSES: Readonly> = { "claude-sdk": { label: "Claude Agent SDK", + shell: false, apis: ["anthropic/v1/messages"], preferred: ["system.ai.claude-sonnet-5", "system.ai.claude-opus-5-5"], }, + codex: { + label: "Codex", + shell: true, + apis: ["codex/v1/responses"], + preferred: ["system.ai.gpt-5-5"], + }, pi: { label: "Pi", + shell: false, // 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"], @@ -32,6 +42,7 @@ const GATEWAY_HARNESSES: Readonly> = { }, "openai-agents": { label: "OpenAI Agents SDK", + shell: false, apis: ["openai/v1/responses"], preferred: ["system.ai.gpt-5-5"], }, @@ -53,11 +64,13 @@ export interface HarnessOffer extends OmnigentHarness { /** * What a user can pick: the allowed gateway harnesses their host reports as - * usable, each with the models whose API it speaks. + * usable, minus shell harnesses when they cannot be sandboxed, each with the + * models whose API it speaks. */ export function offerHarnesses(opts: { allowed: string[]; configured: Record; + shellAllowed: boolean; models: GatewayModel[]; }): HarnessOffer[] { const out: HarnessOffer[] = []; @@ -68,12 +81,14 @@ export function offerHarnesses(opts: { // "needs-auth" is advisory for gateway harnesses: their auth comes from the profile. if (ready === undefined || ready === false || ready === "binary-missing") continue; + if (info.shell && !opts.shellAllowed) 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, + shell: info.shell, ready, models, defaultModel: info.preferred.find((m) => models.includes(m)) ?? models[0], diff --git a/packages/appkit/src/plugins/omnigent/manifest.json b/packages/appkit/src/plugins/omnigent/manifest.json index 7e2a18198..ea13a9126 100644 --- a/packages/appkit/src/plugins/omnigent/manifest.json +++ b/packages/appkit/src/plugins/omnigent/manifest.json @@ -2,7 +2,7 @@ "$schema": "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", "name": "omnigent", "displayName": "Omnigent Plugin", - "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway", + "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Codex, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway", "stability": "beta", "resources": { "required": [], @@ -12,7 +12,7 @@ "rules": { "must": [ "Add pyproject.toml (python >=3.12,<3.13) with omnigent[databricks] and a uv.lock that uses public PyPI", - "Add harness CLIs to package.json: @anthropic-ai/claude-code, @earendil-works/pi-coding-agent", + "Add harness CLIs to package.json: @anthropic-ai/claude-code, @openai/codex, @earendil-works/pi-coding-agent", "Start the app with node, not npm run, so SIGTERM reaches the plugin and it can stop its child processes" ] } diff --git a/packages/appkit/src/plugins/omnigent/omnigent.ts b/packages/appkit/src/plugins/omnigent/omnigent.ts index 77c230052..e5838805c 100644 --- a/packages/appkit/src/plugins/omnigent/omnigent.ts +++ b/packages/appkit/src/plugins/omnigent/omnigent.ts @@ -17,7 +17,7 @@ import { type HarnessOffer, } from "./harnesses"; import manifest from "./manifest.json"; -import { GATEWAY_PROFILE } from "./runtime/env"; +import { profileFor } from "./runtime/env"; import { OmnigentRuntime } from "./runtime/runtime"; import type { CreateSessionRequest, IOmnigentConfig } from "./types"; @@ -37,7 +37,7 @@ export class OmnigentPlugin extends Plugin { static manifest = defineManifest<"omnigent">(manifest); protected static description = - "Embedded Omnigent: Claude Agent SDK, Pi and OpenAI Agents SDK, through Unity AI Gateway"; + "Embedded Omnigent: Claude Agent SDK, Codex, Pi and more, through Unity AI Gateway"; declare protected config: IOmnigentConfig; private runtime?: OmnigentRuntime; @@ -163,6 +163,7 @@ export class OmnigentPlugin extends Plugin { return offerHarnesses({ allowed: this.allowedHarnesses(), configured: mine?.configured_harnesses ?? {}, + shellAllowed: rt.shellHarnessesAllowed, models: await this.servedModels(), }); } @@ -184,7 +185,7 @@ export class OmnigentPlugin extends Plugin { description: "AppKit app agent", harness, model, - profile: GATEWAY_PROFILE, + profile: profileFor(harness), instructions: this.instructions(email), }); } @@ -521,7 +522,7 @@ export class OmnigentPlugin extends Plugin { exports() { return { - /** Runtime health: server, hosts, gateway counters. */ + /** Runtime health: server, sandbox, hosts, gateway counters. */ status: () => this.rt().status(), }; } diff --git a/packages/appkit/src/plugins/omnigent/runtime/codex.ts b/packages/appkit/src/plugins/omnigent/runtime/codex.ts new file mode 100644 index 000000000..34d6e2bce --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/runtime/codex.ts @@ -0,0 +1,59 @@ +import fs from "node:fs"; +import path from "node:path"; + +interface CodexTools { + /** The Codex binary from the app's `@openai/codex` package. */ + bin?: string; + /** A usable `bwrap`, to sandbox it. */ + bwrap?: string; +} + +function readdir(dir: string): string[] { + try { + return fs.readdirSync(dir); + } catch { + return []; + } +} + +function onPath(name: string): string | undefined { + for (const dir of (process.env.PATH ?? "").split(path.delimiter)) { + const p = path.join(dir, name); + try { + fs.accessSync(p, fs.constants.X_OK); + if (dir) return p; + } catch {} + } + return undefined; +} + +/** Finds the Codex binary (to wrap it) and bubblewrap (PATH, else the one Codex's package ships). */ +export function findCodex(appDir: string): CodexTools { + const nm = path.join(appDir, "node_modules", "@openai"); + let bin: string | undefined; + // vendor//bin/codex (0.15x+) or vendor//codex/codex (older). + for (const pkg of readdir(nm).filter((d) => d.startsWith("codex-"))) { + for (const triple of readdir(path.join(nm, pkg, "vendor"))) { + const base = path.join(nm, pkg, "vendor", triple); + bin = + [ + path.join(base, "bin", "codex"), + path.join(base, "codex", "codex"), + ].find((p) => fs.existsSync(p)) ?? bin; + } + } + let bwrap = onPath("bwrap"); + if (!bwrap && bin) { + const shipped = path.resolve( + path.dirname(bin), + "..", + "codex-resources", + "bwrap", + ); + if (fs.existsSync(shipped)) { + fs.chmodSync(shipped, 0o755); + bwrap = shipped; + } + } + return { bin, bwrap }; +} diff --git a/packages/appkit/src/plugins/omnigent/runtime/env.ts b/packages/appkit/src/plugins/omnigent/runtime/env.ts index bdb1b31d0..36190596d 100644 --- a/packages/appkit/src/plugins/omnigent/runtime/env.ts +++ b/packages/appkit/src/plugins/omnigent/runtime/env.ts @@ -2,18 +2,24 @@ import { createHash } from "node:crypto"; import fs from "node:fs"; import path from "node:path"; -/** The Databricks profile the children use; it points at the model gateway. */ -export const GATEWAY_PROFILE = "appkit-omnigent"; +/** Databricks profile names the children use; both point at the model gateway. */ +const PROFILE_HTTPS = "appkit-omnigent"; +const PROFILE_HTTP = "appkit-omnigent-http"; + +/** Harnesses whose TLS stack only trusts public roots use the plain-http listener. */ +export function profileFor(harness: string): string { + return harness.startsWith("codex") ? PROFILE_HTTP : PROFILE_HTTPS; +} /** * Directory layout under the runtime root: * - * server/ Omnigent server data + * server/ Omnigent server data (masked from harnesses) * users//home one user's $HOME: Databricks profile, Omnigent config * users//data that user's host data * work/ that user's workspace (sessions' cwd) - * bin/ public: CA bundle - * tls/ gateway key + * bin/ public: sandbox wrappers, CA bundle + * tls/ gateway key (masked) * logs/ */ export class RuntimeLayout { @@ -57,28 +63,34 @@ export class RuntimeLayout { interface GatewayEndpoints { httpsUrl: string; + httpUrl: string; placeholder: string; } /** - * Writes a user's `$HOME` config: a `.databrickscfg` whose profile points at - * the gateway with the placeholder token (no real credential), and an - * Omnigent `config.yaml` that authenticates through that profile. + * Writes a user's `$HOME` config: a `.databrickscfg` whose two profiles point + * at the gateway with the placeholder token (no real credential), and an + * Omnigent `config.yaml` that authenticates through the https one. */ export function writeUserConfig(home: string, gw: GatewayEndpoints): void { fs.mkdirSync(path.join(home, ".config", "omnigent"), { recursive: true }); const cfg = [ - `[${GATEWAY_PROFILE}]`, + `[${PROFILE_HTTPS}]`, `host = ${gw.httpsUrl}`, `token = ${gw.placeholder}`, "auth_type = pat", "", + `[${PROFILE_HTTP}]`, + `host = ${gw.httpUrl}`, + `token = ${gw.placeholder}`, + "auth_type = pat", + "", ].join("\n"); fs.writeFileSync(path.join(home, ".databrickscfg"), cfg, { mode: 0o600 }); const omnigentConfig = [ "auth:", " type: databricks", - ` profile: ${GATEWAY_PROFILE}`, + ` profile: ${PROFILE_HTTPS}`, "", ].join("\n"); fs.writeFileSync( diff --git a/packages/appkit/src/plugins/omnigent/runtime/gateway.ts b/packages/appkit/src/plugins/omnigent/runtime/gateway.ts index 7adda8be8..7760dd071 100644 --- a/packages/appkit/src/plugins/omnigent/runtime/gateway.ts +++ b/packages/appkit/src/plugins/omnigent/runtime/gateway.ts @@ -1,5 +1,5 @@ import { randomBytes } from "node:crypto"; -import type http from "node:http"; +import http from "node:http"; import https from "node:https"; import type { AddressInfo } from "node:net"; @@ -26,8 +26,9 @@ interface GatewayOptions { token: () => Promise; /** TLS material for the https listener (Omnigent's token helper only serves https hosts). */ tls: { key: Buffer; cert: Buffer }; - /** Loopback port; 0 picks a free one. */ - port?: number; + /** Loopback ports; 0 picks a free one. */ + httpsPort?: number; + httpPort?: number; allow?: readonly RegExp[]; } @@ -47,36 +48,47 @@ export class ModelGateway { readonly stats: GatewayStats = { forwarded: 0, denied: 0 }; private readonly allow: readonly RegExp[]; private readonly upstream: URL; - private server?: https.Server; - private port = 0; + private servers: Array = []; + private ports = { https: 0, http: 0 }; constructor(private readonly opts: GatewayOptions) { this.allow = opts.allow ?? DEFAULT_GATEWAY_ALLOW; this.upstream = new URL(opts.upstream); } - /** The gateway's URL (https: Omnigent's token helper only serves https hosts). */ + /** https host for harnesses whose token helper requires https. */ get httpsUrl(): string { - return `https://127.0.0.1:${this.port}`; + return `https://127.0.0.1:${this.ports.https}`; + } + + /** http host for harnesses that only trust public roots (Codex). */ + get httpUrl(): string { + return `http://127.0.0.1:${this.ports.http}`; } async start(): Promise { const handler = (req: http.IncomingMessage, res: http.ServerResponse) => { void this.handle(req, res); }; - this.server = https.createServer(this.opts.tls, handler); - this.port = await listen(this.server, this.opts.port ?? 0); - logger.debug("model gateway on %s", this.httpsUrl); + const tlsServer = https.createServer(this.opts.tls, handler); + const plainServer = http.createServer(handler); + this.ports.https = await listen(tlsServer, this.opts.httpsPort ?? 0); + this.ports.http = await listen(plainServer, this.opts.httpPort ?? 0); + this.servers = [tlsServer, plainServer]; + logger.debug("model gateway on %s and %s", this.httpsUrl, this.httpUrl); } async stop(): Promise { - const s = this.server; - if (!s) return; - await new Promise((resolve) => { - s.close(() => resolve()); - s.closeAllConnections(); - }); - this.server = undefined; + await Promise.all( + this.servers.map( + (s) => + new Promise((resolve) => { + s.close(() => resolve()); + if ("closeAllConnections" in s) s.closeAllConnections(); + }), + ), + ); + this.servers = []; } /** Whether a request may be forwarded. Exposed for tests. */ @@ -142,7 +154,10 @@ export class ModelGateway { } } -function listen(server: https.Server, port: number): Promise { +function listen( + server: http.Server | https.Server, + port: number, +): Promise { return new Promise((resolve, reject) => { server.once("error", reject); server.listen(port, "127.0.0.1", () => { diff --git a/packages/appkit/src/plugins/omnigent/runtime/runtime.ts b/packages/appkit/src/plugins/omnigent/runtime/runtime.ts index fbc4ee209..4c2c3a2a9 100644 --- a/packages/appkit/src/plugins/omnigent/runtime/runtime.ts +++ b/packages/appkit/src/plugins/omnigent/runtime/runtime.ts @@ -7,8 +7,10 @@ import { fileURLToPath } from "node:url"; import { createLogger } from "../../../logging"; import type { OmnigentStatus } from "../types"; +import { findCodex } from "./codex"; import { RuntimeLayout, childEnv, writeUserConfig } from "./env"; import { ModelGateway } from "./gateway"; +import { HarnessSandbox } from "./sandbox"; import { ensureGatewayTls, type GatewayTls } from "./tls"; const logger = createLogger("omnigent:runtime"); @@ -50,6 +52,9 @@ export class OmnigentRuntime { readonly authHeader = `x-appkit-omnigent-user-${randomBytes(12).toString("hex")}`; private tls!: GatewayTls; private gateway!: ModelGateway; + private sandbox?: HarnessSandbox; + /** Points Omnigent at the sandboxed Codex wrapper. */ + private harnessEnv: Record = {}; private server?: ChildProcess; private serverReady?: Promise; private readonly hosts = new Map(); @@ -68,6 +73,11 @@ export class OmnigentRuntime { return `http://127.0.0.1:${this.port}`; } + /** Harnesses with their own shell (Codex) run only inside the sandbox. */ + get shellHarnessesAllowed(): boolean { + return Boolean(this.sandbox); + } + async start(): Promise { this.layout.create(); this.tls = ensureGatewayTls(this.layout.tls, this.layout.bin); @@ -78,6 +88,20 @@ export class OmnigentRuntime { }); await this.gateway.start(); + const codex = findCodex(this.opts.appDir); + if (codex.bin && codex.bwrap) { + this.sandbox = new HarnessSandbox(this.layout.root, codex.bwrap); + this.harnessEnv = { + OMNIGENT_CODEX_PATH: this.sandbox.wrap("codex", codex.bin), + }; + // Omnigent's own REPL terminal sandbox looks for `bwrap` on PATH. + const shimBwrap = path.join(this.sandbox.shimDir, "bwrap"); + fs.copyFileSync(codex.bwrap, shimBwrap); + fs.chmodSync(shimBwrap, 0o755); + } else if (codex.bin) { + logger.warn("bubblewrap not found: Codex is not offered"); + } + this.port = await freePort(); this.startServer(); this.reaper = setInterval(() => this.reapIdleHosts(), 60_000); @@ -120,6 +144,7 @@ export class OmnigentRuntime { }, }); this.server = this.spawn("server", args, env, this.layout.serverDir); + this.protect(); this.server.once("exit", (code, signal) => { if (this.stopping) return; this.lastError = `Omnigent server exited (${code ?? signal}); restarting`; @@ -184,10 +209,13 @@ export class OmnigentRuntime { caBundle: this.tls.bundleFile, caCert: this.tls.certFile, pathPrepend: [ + // Sandbox wrappers first: Omnigent may resolve a harness CLI by name from PATH. + ...(this.sandbox ? [this.sandbox.shimDir] : []), path.dirname(this.opts.python), path.join(this.opts.appDir, "node_modules", ".bin"), ], extra: { + ...this.harnessEnv, OMNIGENT_HOST_ID: hostId, OMNIGENT_HOST_NAME: `appkit-${userId}`, OMNIGENT_DATABRICKS_EXTRA_HEADERS: JSON.stringify({ @@ -231,6 +259,7 @@ export class OmnigentRuntime { entry.failed = true; }); this.hosts.set(email, entry); + this.protect(); proc.once("exit", (code, signal) => { if (!this.stopping && code !== 0 && signal !== "SIGTERM") { logger.warn( @@ -241,6 +270,7 @@ export class OmnigentRuntime { ); } if (this.hosts.get(email) === entry) this.hosts.delete(email); + this.protect(); }); return entry.ready; } @@ -278,6 +308,7 @@ export class OmnigentRuntime { ), error: this.lastError, version: this.version, + sandbox: Boolean(this.sandbox), hosts: this.hosts.size, gateway: this.gateway?.stats ?? { forwarded: 0, denied: 0 }, }; @@ -314,6 +345,14 @@ export class OmnigentRuntime { return path.join(this.layout.logs, `${name}.log`); } + private protect(): void { + this.sandbox?.protect([ + process.pid, + this.server?.pid, + ...[...this.hosts.values()].map((h) => h.proc.pid), + ]); + } + private spawn( name: string, argv: string[], diff --git a/packages/appkit/src/plugins/omnigent/runtime/sandbox.ts b/packages/appkit/src/plugins/omnigent/runtime/sandbox.ts new file mode 100644 index 000000000..45c532243 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/runtime/sandbox.ts @@ -0,0 +1,78 @@ +import fs from "node:fs"; +import path from "node:path"; + +/** + * Harness binaries with their own shell run through a bubblewrap wrapper, so + * the shell cannot read the app's state or other processes' memory: + * + * - own PID, IPC and UTS namespaces; + * - the runtime directory masked with a tmpfs, with only this user's home and + * work directory (and the public `bin` directory) bound back in; + * - `/proc/` of every protected process (Node, the Omnigent server, the + * hosts) overmounted, since Apps containers cannot mount a fresh procfs. + * + * The protected list is read at launch, so a process started later is not + * masked from a harness that is already running; the plugin keeps secrets out + * of every child's environment so such a process has nothing to leak. + */ +export class HarnessSandbox { + private readonly pidFile: string; + + constructor( + private readonly root: string, + private readonly bwrap: string, + ) { + this.pidFile = path.join(root, "protected.pids"); + } + + get publicBin(): string { + return path.join(this.root, "bin"); + } + + /** Records the processes sandboxed harnesses must not see. */ + protect(pids: Array): void { + const list = [ + ...new Set(pids.filter((p): p is number => typeof p === "number")), + ]; + fs.writeFileSync(this.pidFile, `${list.join("\n")}\n`); + } + + /** Directory of wrappers named after the real CLIs; must come first on the hosts' PATH. */ + get shimDir(): string { + return path.join(this.publicBin, "sandboxed"); + } + + /** + * Writes a wrapper named like the CLI it wraps (`codex`) and returns its + * path. Omnigent may launch a harness by `OMNIGENT_*_PATH` or by name from + * `PATH`, so both must resolve to the wrapper: put {@link shimDir} first on + * `PATH`. + */ + wrap(name: string, real: string): string { + fs.mkdirSync(this.shimDir, { recursive: true }); + const wrapper = path.join(this.shimDir, name); + fs.writeFileSync(wrapper, this.script(real), { mode: 0o755 }); + return wrapper; + } + + script(real: string): string { + const q = shellQuote; + return `#!/bin/sh +# Generated by the AppKit omnigent plugin. $HOME is /users//home. +U="$(dirname "$HOME")" +W=${q(path.join(this.root, "work"))}/"$(basename "$U")" +MASK="" +for p in $(cat ${q(this.pidFile)} 2>/dev/null); do + [ -d "/proc/$p" ] && MASK="$MASK --tmpfs /proc/$p" +done +exec ${q(this.bwrap)} --die-with-parent --unshare-pid --unshare-ipc --unshare-uts \\ + --dev-bind / / --tmpfs ${q(this.root)} $MASK \\ + --bind "$U" "$U" --bind "$W" "$W" --ro-bind ${q(this.publicBin)} ${q(this.publicBin)} \\ + -- ${q(real)} "$@" +`; + } +} + +function shellQuote(s: string): string { + return `'${s.replace(/'/g, `'\\''`)}'`; +} diff --git a/packages/appkit/src/plugins/omnigent/runtime/tls.ts b/packages/appkit/src/plugins/omnigent/runtime/tls.ts index 3fe9e9718..33cb8a91a 100644 --- a/packages/appkit/src/plugins/omnigent/runtime/tls.ts +++ b/packages/appkit/src/plugins/omnigent/runtime/tls.ts @@ -20,8 +20,8 @@ const SYSTEM_ROOTS = [ /** * A self-signed certificate for `127.0.0.1`, made once per runtime directory * with the system `openssl`, and a CA bundle that trusts it alongside the - * system roots. The key stays in `privateDir`; the certificate and - * bundle go to `publicDir`, which the harnesses read. + * system roots. The key stays in `privateDir` (masked from sandboxed + * harnesses); the certificate and bundle go to `publicDir`, which they read. */ export function ensureGatewayTls( privateDir: string, diff --git a/packages/appkit/src/plugins/omnigent/tests/gateway.test.ts b/packages/appkit/src/plugins/omnigent/tests/gateway.test.ts index 3d4817d0d..c0cdb5a9a 100644 --- a/packages/appkit/src/plugins/omnigent/tests/gateway.test.ts +++ b/packages/appkit/src/plugins/omnigent/tests/gateway.test.ts @@ -1,5 +1,5 @@ import fs from "node:fs"; -import https from "node:https"; +import http from "node:http"; import os from "node:os"; import path from "node:path"; @@ -10,12 +10,11 @@ import { ensureGatewayTls } from "../runtime/tls"; function get( url: string, - ca: Buffer, headers: Record = {}, ): Promise<{ status: number; body: string }> { return new Promise((resolve, reject) => { - https - .get(url, { headers, ca }, (res) => { + http + .get(url, { headers }, (res) => { let body = ""; res.on("data", (c) => (body += c)); res.on("end", () => resolve({ status: res.statusCode ?? 0, body })); @@ -27,7 +26,6 @@ function get( describe("ModelGateway", () => { const root = fs.mkdtempSync(path.join(os.tmpdir(), "omni-gw-")); let gw: ModelGateway; - let cert: Buffer; let tokenCalls = 0; beforeAll(async () => { @@ -35,7 +33,6 @@ describe("ModelGateway", () => { path.join(root, "tls"), path.join(root, "bin"), ); - cert = tls.cert; gw = new ModelGateway({ upstream: "https://127.0.0.1:9", // never reached in these tests token: async () => { @@ -85,7 +82,7 @@ describe("ModelGateway", () => { test("refuses non-model paths with 403 and never mints a token for them", async () => { const before = tokenCalls; - const r = await get(`${gw.httpsUrl}/api/2.0/clusters/list`, cert, { + const r = await get(`${gw.httpUrl}/api/2.0/clusters/list`, { authorization: `Bearer ${gw.placeholder}`, }); expect(r.status).toBe(403); @@ -94,11 +91,12 @@ describe("ModelGateway", () => { }); test("answers host metadata with 404 so SDKs stay on token auth", async () => { - const r = await get(`${gw.httpsUrl}/.well-known/databricks-config`, cert); + const r = await get(`${gw.httpUrl}/.well-known/databricks-config`); expect(r.status).toBe(404); }); test("listens on loopback only", () => { expect(gw.httpsUrl).toMatch(/^https:\/\/127\.0\.0\.1:\d+$/); + expect(gw.httpUrl).toMatch(/^http:\/\/127\.0\.0\.1:\d+$/); }); }); diff --git a/packages/appkit/src/plugins/omnigent/tests/units.test.ts b/packages/appkit/src/plugins/omnigent/tests/units.test.ts index 25f379a84..8fb82e2c2 100644 --- a/packages/appkit/src/plugins/omnigent/tests/units.test.ts +++ b/packages/appkit/src/plugins/omnigent/tests/units.test.ts @@ -7,7 +7,8 @@ import { afterEach, describe, expect, test } from "vitest"; import { buildBundle } from "../bundle"; import { offerHarnesses } from "../harnesses"; -import { childEnv, GATEWAY_PROFILE, writeUserConfig } from "../runtime/env"; +import { childEnv, profileFor, writeUserConfig } from "../runtime/env"; +import { HarnessSandbox } from "../runtime/sandbox"; function untar(buf: Buffer): Record { const out: Record = {}; @@ -37,9 +38,9 @@ describe("buildBundle", () => { gunzipSync( buildBundle({ name: "appkit", - harness: "pi", + harness: "codex", model: "system.ai.gpt-5-5", - profile: GATEWAY_PROFILE, + profile: profileFor("codex"), instructions: "Be terse.", }), ), @@ -48,12 +49,18 @@ describe("buildBundle", () => { const cfg = JSON.parse(files["config.yaml"]); expect(cfg.executor).toMatchObject({ model: "system.ai.gpt-5-5", - profile: "appkit-omnigent", - config: { harness: "pi" }, - auth: { type: "databricks", profile: "appkit-omnigent" }, + profile: "appkit-omnigent-http", + config: { harness: "codex" }, + auth: { type: "databricks", profile: "appkit-omnigent-http" }, }); expect(files["AGENTS.md"]).toBe("Be terse."); }); + + test("Claude and Pi use the https profile, Codex the http one", () => { + expect(profileFor("claude-sdk")).toBe("appkit-omnigent"); + expect(profileFor("pi")).toBe("appkit-omnigent"); + expect(profileFor("codex")).toBe("appkit-omnigent-http"); + }); }); describe("child environment", () => { @@ -65,6 +72,7 @@ describe("child environment", () => { const gw = { httpsUrl: "https://127.0.0.1:1", + httpUrl: "http://127.0.0.1:2", placeholder: "appkit-omnigent-ph", }; @@ -132,7 +140,11 @@ describe("child environment", () => { writeUserConfig(home, gw); const cfg = fs.readFileSync(path.join(home, ".databrickscfg"), "utf8"); expect(cfg).toContain(`host = ${gw.httpsUrl}`); - expect(cfg.match(/token = (.*)/g)).toEqual([`token = ${gw.placeholder}`]); + expect(cfg).toContain(`host = ${gw.httpUrl}`); + expect(cfg.match(/token = (.*)/g)).toEqual([ + `token = ${gw.placeholder}`, + `token = ${gw.placeholder}`, + ]); expect( (fs.statSync(path.join(home, ".databrickscfg")).mode & 0o777).toString(8), ).toBe("600"); @@ -140,7 +152,7 @@ describe("child environment", () => { path.join(home, ".config", "omnigent", "config.yaml"), "utf8", ); - expect(omni).toContain(`profile: ${GATEWAY_PROFILE}`); + expect(omni).toContain("profile: appkit-omnigent"); }); }); @@ -177,6 +189,7 @@ describe("offerHarnesses", () => { ]; const all = { "claude-sdk": true, + codex: true, pi: "needs-auth", "openai-agents": true, } as const; @@ -185,6 +198,7 @@ describe("offerHarnesses", () => { const offers = offerHarnesses({ allowed: Object.keys(all), configured: all, + shellAllowed: true, models: catalog, }); const by = Object.fromEntries(offers.map((o) => [o.id, o])); @@ -192,6 +206,7 @@ describe("offerHarnesses", () => { "system.ai.claude-sonnet-5", "system.ai.glm-5-3", ]); + expect(by.codex.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", @@ -199,24 +214,50 @@ describe("offerHarnesses", () => { // Pi: the Codex responses API only, and never Claude expect(by.pi.models).toEqual(["system.ai.gpt-5-5", "system.ai.glm-5-3"]); expect(by["claude-sdk"].defaultModel).toBe("system.ai.claude-sonnet-5"); - expect(by.pi.defaultModel).toBe("system.ai.gpt-5-5"); + expect(by.codex.defaultModel).toBe("system.ai.gpt-5-5"); }); test("offers only gateway harnesses", () => { const offers = offerHarnesses({ allowed: ["claude-sdk", "cursor"], configured: { "claude-sdk": true, cursor: true }, + shellAllowed: true, models: catalog, }); expect(offers.map((o) => o.id)).toEqual(["claude-sdk"]); }); - test("drops harnesses whose binary is missing", () => { + test("drops missing binaries and, without a sandbox, shell harnesses", () => { const offers = offerHarnesses({ - allowed: ["claude-sdk", "pi"], - configured: { "claude-sdk": true, pi: "binary-missing" }, + allowed: ["claude-sdk", "codex", "pi"], + configured: { "claude-sdk": true, codex: true, pi: "binary-missing" }, + shellAllowed: false, models: catalog, }); expect(offers.map((o) => o.id)).toEqual(["claude-sdk"]); }); }); + +describe("HarnessSandbox", () => { + test("wrapper masks the runtime root and protected pids, and binds only the user's dirs", () => { + const root = fs.mkdtempSync(path.join(os.tmpdir(), "omni-sbx-")); + try { + const sbx = new HarnessSandbox(root, "/usr/bin/bwrap"); + sbx.protect([1234, undefined, 1234, 99]); + expect(fs.readFileSync(path.join(root, "protected.pids"), "utf8")).toBe( + "1234\n99\n", + ); + const wrapper = sbx.wrap("codex", "/app/node_modules/codex"); + // named like the CLI, in the shim dir that goes first on PATH + expect(wrapper).toBe(path.join(sbx.shimDir, "codex")); + const script = fs.readFileSync(wrapper, "utf8"); + expect(script).toContain(`--tmpfs '${root}'`); + expect(script).toContain("--unshare-pid"); + expect(script).toContain("--tmpfs /proc/$p"); + expect(script).toContain('--bind "$U" "$U" --bind "$W" "$W"'); + expect(script).toContain("-- '/app/node_modules/codex' \"$@\""); + } finally { + fs.rmSync(root, { recursive: true, force: true }); + } + }); +}); diff --git a/packages/appkit/src/plugins/omnigent/types.ts b/packages/appkit/src/plugins/omnigent/types.ts index e871f649b..0cecdae82 100644 --- a/packages/appkit/src/plugins/omnigent/types.ts +++ b/packages/appkit/src/plugins/omnigent/types.ts @@ -24,6 +24,8 @@ export interface IOmnigentConfig extends BasePluginConfig { export interface OmnigentHarness { id: string; label: string; + /** Has its own shell, so runs sandboxed. */ + shell: boolean; ready: boolean | string; } @@ -40,6 +42,8 @@ export interface OmnigentStatus { error?: string; /** Omnigent server version. */ version?: string; + /** Harnesses with their own shell run in bubblewrap. */ + sandbox: boolean; /** Running per-user hosts. */ hosts: number; /** Model gateway counters. */ diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index 07ffbdc6a..af1e98032 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -345,7 +345,7 @@ "omnigent": { "name": "omnigent", "displayName": "Omnigent Plugin", - "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway", + "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Codex, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway", "package": "@databricks/appkit", "resources": { "required": [], @@ -356,7 +356,7 @@ "rules": { "must": [ "Add pyproject.toml (python >=3.12,<3.13) with omnigent[databricks] and a uv.lock that uses public PyPI", - "Add harness CLIs to package.json: @anthropic-ai/claude-code, @earendil-works/pi-coding-agent", + "Add harness CLIs to package.json: @anthropic-ai/claude-code, @openai/codex, @earendil-works/pi-coding-agent", "Start the app with node, not npm run, so SIGTERM reaches the plugin and it can stop its child processes" ] } From d88aa927a34b4b7ea121e4f7a6fa5f3291ede3b8 Mon Sep 17 00:00:00 2001 From: Noah Solomon Date: Wed, 30 Sep 2026 00:05:09 +1000 Subject: [PATCH 3/4] feat(appkit): serve app tools to omnigent harnesses, with modes and approvals (3/4) - The app's defineTool() entries (and optionally other plugins' tools) are served to every harness over a stateless MCP endpoint and run in the session user's execution context. - The endpoint is guarded by HMAC session keys keyed on the app's OAuth client secret (Apps proxy traffic also arrives from loopback). - auto / ask / read modes compile to Omnigent CEL policies; writes pause for approval in ask and are denied in read. Co-authored-by: Isaac Signed-off-by: Noah Solomon --- docs/docs/plugins/omnigent.md | 53 +++- .../appkit/src/plugins/omnigent/bundle.ts | 33 ++- packages/appkit/src/plugins/omnigent/keys.ts | 69 +++++ .../appkit/src/plugins/omnigent/manifest.json | 10 +- .../appkit/src/plugins/omnigent/mcp-server.ts | 127 +++++++++ .../appkit/src/plugins/omnigent/omnigent.ts | 251 +++++++++++++++++- .../appkit/src/plugins/omnigent/policies.ts | 89 +++++++ .../src/plugins/omnigent/tests/plugin.test.ts | 142 +++++++++- .../src/plugins/omnigent/tests/units.test.ts | 212 ++++++++++++++- .../appkit/src/plugins/omnigent/tokens.ts | 42 +++ packages/appkit/src/plugins/omnigent/types.ts | 20 ++ template/appkit.plugins.json | 5 +- 12 files changed, 1023 insertions(+), 30 deletions(-) create mode 100644 packages/appkit/src/plugins/omnigent/keys.ts create mode 100644 packages/appkit/src/plugins/omnigent/mcp-server.ts create mode 100644 packages/appkit/src/plugins/omnigent/policies.ts create mode 100644 packages/appkit/src/plugins/omnigent/tokens.ts diff --git a/docs/docs/plugins/omnigent.md b/docs/docs/plugins/omnigent.md index defc68959..9ce3efbbc 100644 --- a/docs/docs/plugins/omnigent.md +++ b/docs/docs/plugins/omnigent.md @@ -6,20 +6,43 @@ This plugin is currently **beta**. APIs may change between minor releases. Impor ::: -The `omnigent` plugin embeds [Omnigent](https://github.com/omnigent-ai/omnigent), the open-source agent meta-harness, in a Databricks AppKit app. Your users get Claude Agent SDK, Codex, Pi and OpenAI Agents SDK sessions inside the app, with model calls going through Unity AI Gateway as the app's service principal. +The `omnigent` plugin embeds [Omnigent](https://github.com/omnigent-ai/omnigent), the open-source agent meta-harness, in a Databricks AppKit app. Your users get Claude Agent SDK, Codex, Pi and OpenAI Agents SDK sessions inside the app: + +- **Model calls go through Unity AI Gateway** as the app's service principal. +- **The app's tools run as the signed-in user**, so Unity Catalog grants apply. Omnigent owns the agent loop: sessions, history, approvals. The plugin runs it safely inside a Databricks App. 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"; +import { createApp, getExecutionContext, server } from "@databricks/appkit"; +import { defineTool, omnigent } from "@databricks/appkit/beta"; +import { z } from "zod"; await createApp({ plugins: [ server(), - omnigent({ instructions: "You help the team explore order data." }), + omnigent({ + instructions: "You help the team explore order data.", + tools: { + count_orders: defineTool({ + description: "Count this month's orders", + schema: z.object({}), + annotations: { effect: "read" }, + execute: async () => { + // Runs as the session's user. + const { client } = getExecutionContext(); + return client.statementExecution.executeStatement({ + warehouse_id: process.env.DATABRICKS_WAREHOUSE_ID!, + statement: + "SELECT count(*) FROM main.shop.orders WHERE date_trunc('MONTH', order_date) = date_trunc('MONTH', current_date())", + wait_timeout: "30s", + }); + }, + }), + }, + }), ], }); ``` @@ -29,12 +52,16 @@ The app also needs: - **Python.** Add a `pyproject.toml` with `requires-python = ">=3.12,<3.13"` and `omnigent[databricks]`, plus a `uv.lock`. Databricks Apps then builds a Python 3.12 venv next to `npm install`. A plain `requirements.txt` gets Python 3.11, which Omnigent does not support. The lockfile must reference public PyPI. - **Harness CLIs.** Add them to your `package.json`, so you pin their versions: `@anthropic-ai/claude-code`, `@openai/codex` and `@earendil-works/pi-coding-agent`. The plugin only offers harnesses whose CLI is installed. The Claude Agent SDK's bundled CLI is too old for current Claude models, so add `@anthropic-ai/claude-code` even though the SDK ships one. - **Start with `node`.** In `app.yaml`, use `command: ["node", "server.js"]`, not `npm run start`. npm does not pass SIGTERM on, so the plugin never gets to stop its child processes. As a backstop, each child also exits when Node does, so a redeploy still succeeds. +- **User API scopes.** Grant the scopes your tools need (for example `sql`), so tools can run as the user. ## Configuration | Option | Default | | |---|---|---| | `instructions` | a generic assistant prompt | System instructions for every session | +| `tools` | none | `defineTool()` entries, served to every harness over MCP | +| `pluginTools` | none | `true`, or a list of plugin names, to also serve other plugins' tools (for example `genie`) | +| `defaultMode` | `ask` | `auto`, `ask` or `read` | | `harnesses` | `claude-sdk`, `codex`, `pi`, `openai-agents` | Harnesses users may pick | | `defaultHarness`, `defaultModel` | first available | | | `runtime.python` | `./.venv/bin/python` | | @@ -43,6 +70,16 @@ The app also needs: Each harness is offered the Unity AI Gateway `system.ai` models whose API it speaks: Claude Agent SDK the Anthropic Messages API, Codex and Pi the Codex Responses API, OpenAI Agents SDK the OpenAI Responses API. +## Modes + +Tools whose `annotations.effect` is `write`, `update` or `destructive` are writes: + +- `auto` runs everything. +- `ask` pauses on writes (and on Omnigent's own OS tools) until the user approves. +- `read` denies writes. + +Omnigent enforces the mode on every harness. 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"}`. + ## Routes All routes are under `/api/omnigent` and scoped to the signed-in user. @@ -50,17 +87,20 @@ All routes are under `/api/omnigent` and scoped to the signed-in user. | Route | | |---|---| | `GET /harnesses` | Harnesses available to the user, each with the gateway models it can drive | -| `GET /sessions`, `POST /sessions` | List, create (`harness`, `model`, `message`, `title`) | -| `GET /sessions/:id`, `DELETE /sessions/:id` | Snapshot (items, status), delete | +| `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 | | `GET /status` | Runtime health: sandbox, hosts, gateway counters | ## How it works - **One Omnigent server, one host per user.** The server listens on loopback in Omnigent's header-auth mode, under a random header name only the plugin knows. Each active user gets their own host process, registered as that user; Omnigent only runs a session on its owner's host. An idle host stops after `runtime.hostIdleMs`. - **No credential in any child process.** The harnesses' Databricks profile points at a model gateway inside the plugin, on loopback, with a placeholder token. The gateway forwards only model-serving paths to the workspace and adds the service principal's token. Every other path returns 403. Codex only trusts public TLS roots, so it reaches the gateway over plain http on loopback; the others use https with a certificate the plugin generates. +- **Tools over MCP.** Each session's bundle points an MCP server at the app's own `/api/omnigent/mcp`, with a key that is an HMAC over a nonce and the user. The plugin verifies the key, then runs the tool inside that user's AppKit execution context, using their latest forwarded token. Keys survive restarts because nothing is stored. - **Codex runs sandboxed.** Codex has its own shell. Its binary runs through a bubblewrap wrapper that gives it: - its own PID namespace - the plugin's state directory masked @@ -73,4 +113,5 @@ All routes are under `/api/omnigent` and scoped to the signed-in user. - **Sessions live in the container.** Omnigent keeps them in SQLite under the plugin's state directory, so they survive a process restart but not a redeploy. - **Only gateway harnesses.** Harnesses that sign in with their vendor (Cursor, Copilot, Gemini and others) bypass Unity AI Gateway, so they are not offered. Native terminal harnesses (`claude-native`, `codex-native`, `pi-native`) are not offered yet: they need tmux and a terminal client. +- **Tool calls need a recent visit.** A tool call runs on the user's most recent forwarded token. After a restart, a user's tools work again once they make a request to the app. - **Shared gateway limits.** All users' model traffic uses the service principal's gateway limits. diff --git a/packages/appkit/src/plugins/omnigent/bundle.ts b/packages/appkit/src/plugins/omnigent/bundle.ts index b82cb64db..ea749a381 100644 --- a/packages/appkit/src/plugins/omnigent/bundle.ts +++ b/packages/appkit/src/plugins/omnigent/bundle.ts @@ -1,5 +1,14 @@ import { gzipSync } from "node:zlib"; +interface McpServerSpec { + name: string; + description?: string; + url: string; + headers?: Record; + /** Seconds. */ + timeout?: number; +} + interface BundleSpec { name: string; description?: string; @@ -8,12 +17,13 @@ interface BundleSpec { /** Databricks profile the harness uses for model calls. */ profile: string; instructions: string; + mcpServers: McpServerSpec[]; } /** - * An Omnigent agent bundle: `config.yaml` (the agent spec) and `AGENTS.md` - * (the instructions), as a tar.gz. The YAML is written as JSON, which is - * valid YAML. + * 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 = { @@ -33,6 +43,23 @@ export function buildBundle(spec: BundleSpec): Buffer { "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}`); + } + files[`tools/mcp/${s.name}.yaml`] = `${JSON.stringify( + { + name: s.name, + description: s.description ?? s.name, + transport: "http", + url: s.url, + timeout: s.timeout ?? 600, + headers: s.headers ?? {}, + }, + null, + 2, + )}\n`; + } return gzipSync(tar(files)); } diff --git a/packages/appkit/src/plugins/omnigent/keys.ts b/packages/appkit/src/plugins/omnigent/keys.ts new file mode 100644 index 000000000..28393703f --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/keys.ts @@ -0,0 +1,69 @@ +import { createHmac, randomBytes, timingSafeEqual } from "node:crypto"; + +/** + * Credentials a session's agent bundle carries to call the app's MCP endpoint. + * + * The key is `HMAC(secret, nonce + "\n" + user)`, so the plugin can verify a + * tool call after a restart without storing anything: it recomputes the HMAC + * from the headers. A key only ever stands for the user it was minted for. + */ +interface SessionKey { + user: string; + nonce: string; + key: string; +} + +export const MCP_USER_HEADER = "x-omnigent-user"; +const MCP_NONCE_HEADER = "x-omnigent-nonce"; +export const MCP_KEY_HEADER = "x-omnigent-key"; + +export class SessionKeys { + private readonly secret: string; + + /** + * @param secret Stable across restarts (the app's OAuth client secret in + * Apps). Without one, keys are only valid for the life of the process. + */ + constructor(secret?: string) { + this.secret = secret || randomBytes(32).toString("hex"); + } + + mint(user: string): SessionKey { + const nonce = randomBytes(18).toString("base64url"); + return { user, nonce, key: this.sign(nonce, user) }; + } + + /** Returns the user a request's key stands for, or `null`. */ + verify( + headers: Record, + ): string | null { + const user = first(headers[MCP_USER_HEADER]); + const nonce = first(headers[MCP_NONCE_HEADER]); + const key = first(headers[MCP_KEY_HEADER]); + if (!user || !nonce || !key) return null; + const expected = Buffer.from(this.sign(nonce, user)); + const given = Buffer.from(key); + if (given.length !== expected.length || !timingSafeEqual(given, expected)) { + return null; + } + return user; + } + + headers(k: SessionKey): Record { + return { + [MCP_USER_HEADER]: k.user, + [MCP_NONCE_HEADER]: k.nonce, + [MCP_KEY_HEADER]: k.key, + }; + } + + private sign(nonce: string, user: string): string { + return createHmac("sha256", this.secret) + .update(`${nonce}\n${user.toLowerCase()}`) + .digest("base64url"); + } +} + +function first(v: string | string[] | undefined): string | undefined { + return Array.isArray(v) ? v[0] : v; +} diff --git a/packages/appkit/src/plugins/omnigent/manifest.json b/packages/appkit/src/plugins/omnigent/manifest.json index ea13a9126..765cbc270 100644 --- a/packages/appkit/src/plugins/omnigent/manifest.json +++ b/packages/appkit/src/plugins/omnigent/manifest.json @@ -2,7 +2,7 @@ "$schema": "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", "name": "omnigent", "displayName": "Omnigent Plugin", - "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Codex, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway", + "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Codex, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway, with the app's tools run as the signed-in user", "stability": "beta", "resources": { "required": [], @@ -13,7 +13,8 @@ "must": [ "Add pyproject.toml (python >=3.12,<3.13) with omnigent[databricks] and a uv.lock that uses public PyPI", "Add harness CLIs to package.json: @anthropic-ai/claude-code, @openai/codex, @earendil-works/pi-coding-agent", - "Start the app with node, not npm run, so SIGTERM reaches the plugin and it can stop its child processes" + "Start the app with node, not npm run, so SIGTERM reaches the plugin and it can stop its child processes", + "Grant the user API scopes the app's tools need (e.g. sql) so tools run as the signed-in user" ] } }, @@ -25,6 +26,11 @@ "type": "string", "description": "System instructions for every session" }, + "defaultMode": { + "type": "string", + "enum": ["auto", "ask", "read"], + "default": "ask" + }, "harnesses": { "type": "array", "items": { diff --git a/packages/appkit/src/plugins/omnigent/mcp-server.ts b/packages/appkit/src/plugins/omnigent/mcp-server.ts new file mode 100644 index 000000000..d71d8484a --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/mcp-server.ts @@ -0,0 +1,127 @@ +import type { AgentToolDefinition, ToolEffect } from "shared"; + +/** One tool served over MCP, with where it came from. */ +export interface McpTool { + /** MCP-safe name (`[A-Za-z0-9_-]{1,64}`). */ + name: string; + definition: AgentToolDefinition; + run: (args: unknown, signal?: AbortSignal) => Promise; +} + +const WRITE_EFFECTS: ReadonlySet = new Set([ + "write", + "update", + "destructive", +]); + +export function isWriteTool(def: AgentToolDefinition): boolean { + const a = def.annotations; + if (!a) return false; + if (a.effect) return WRITE_EFFECTS.has(a.effect); + return a.destructive === true || a.readOnly === false; +} + +/** Maps an AppKit tool name (may contain dots) to an MCP-safe one. */ +export function mcpToolName(name: string, prefix?: string): string { + const raw = prefix ? `${prefix}__${name}` : name; + return raw.replace(/[^A-Za-z0-9_-]/g, "_").slice(0, 64); +} + +interface JsonRpcRequest { + jsonrpc?: string; + id?: string | number | null; + method?: string; + params?: Record; +} + +type JsonRpcReply = Record; + +const MCP_PROTOCOL_VERSION = "2025-06-18"; + +/** + * A stateless MCP server over plain JSON-RPC (streamable HTTP with JSON + * responses): `initialize`, `ping`, `tools/list`, `tools/call`. The caller has + * already authenticated the session key and set up the user context. + */ +export class McpToolServer { + constructor( + private readonly tools: () => McpTool[], + private readonly info = { name: "appkit", version: "1" }, + ) {} + + /** Handles one JSON-RPC message; returns `null` when nothing needs a reply. */ + async handle( + body: unknown, + signal?: AbortSignal, + ): Promise { + const m = body as JsonRpcRequest; + if (!m || typeof m !== "object" || typeof m.method !== "string") { + return rpcError(null, -32600, "Invalid request"); + } + const isNotification = m.id === undefined || m.id === null; + const reply = (result: unknown) => ({ jsonrpc: "2.0", id: m.id, result }); + switch (m.method) { + case "initialize": + return reply({ + protocolVersion: + typeof m.params?.protocolVersion === "string" + ? m.params.protocolVersion + : MCP_PROTOCOL_VERSION, + capabilities: { tools: { listChanged: false } }, + serverInfo: this.info, + }); + case "ping": + return reply({}); + case "tools/list": + return reply({ + tools: this.tools().map((t) => ({ + name: t.name, + description: t.definition.description, + inputSchema: t.definition.parameters, + annotations: { + readOnlyHint: !isWriteTool(t.definition), + destructiveHint: + t.definition.annotations?.effect === "destructive", + }, + })), + }); + case "tools/call": { + const name = String(m.params?.name ?? ""); + const tool = this.tools().find((t) => t.name === name); + if (!tool) + return rpcError(m.id ?? null, -32602, `Unknown tool: ${name}`); + try { + const out = await tool.run(m.params?.arguments ?? {}, signal); + return reply({ content: [{ type: "text", text: asText(out) }] }); + } catch (err) { + // Tool failures go back to the model as a tool result, not a protocol error. + return reply({ + content: [ + { + type: "text", + text: err instanceof Error ? err.message : String(err), + }, + ], + isError: true, + }); + } + } + default: + if (isNotification) return null; + return rpcError(m.id ?? null, -32601, `Method not found: ${m.method}`); + } + } +} + +function asText(v: unknown): string { + if (typeof v === "string") return v; + try { + return JSON.stringify(v); + } catch { + return String(v); + } +} + +function rpcError(id: unknown, code: number, message: string): JsonRpcReply { + return { jsonrpc: "2.0", id, error: { code, message } }; +} diff --git a/packages/appkit/src/plugins/omnigent/omnigent.ts b/packages/appkit/src/plugins/omnigent/omnigent.ts index e5838805c..e1e6dde88 100644 --- a/packages/appkit/src/plugins/omnigent/omnigent.ts +++ b/packages/appkit/src/plugins/omnigent/omnigent.ts @@ -5,7 +5,11 @@ import { Readable } from "node:stream"; import type express from "express"; import type { IAppRouter } from "shared"; -import { ServiceContext } from "../../context"; +import { runInUserContext, ServiceContext } from "../../context"; +import { + executeFromRegistry, + toolsFromRegistry, +} from "../../core/agent/tools/define-tool"; import { createLogger } from "../../logging"; import { Plugin, toPlugin } from "../../plugin"; import { defineManifest } from "../../registry"; @@ -16,13 +20,29 @@ import { offerHarnesses, type HarnessOffer, } from "./harnesses"; +import { SessionKeys } from "./keys"; import manifest from "./manifest.json"; +import { + isWriteTool, + type McpTool, + McpToolServer, + mcpToolName, +} from "./mcp-server"; +import { + MODE_POLICY_PREFIX, + modeFromPolicies, + modePolicies, + OMNIGENT_MODES, + type OmnigentMode, +} from "./policies"; import { profileFor } from "./runtime/env"; import { OmnigentRuntime } from "./runtime/runtime"; +import { UserTokens } from "./tokens"; import type { CreateSessionRequest, IOmnigentConfig } from "./types"; const logger = createLogger("omnigent"); +const ELICITATION_ACTIONS = new Set(["accept", "decline", "cancel"]); const MODELS_TTL_MS = 5 * 60_000; /** @@ -31,7 +51,8 @@ const MODELS_TTL_MS = 5 * 60_000; * Runs an Omnigent server and one host per active user as supervised child * processes, so every harness Omnigent can route through Unity AI Gateway is * available to the app's users. Model calls are billed to the app's service - * principal through a loopback gateway (no child holds a real credential). + * principal through a loopback gateway (no child holds a real credential); + * the app's tools reach the harnesses over MCP and run as the session's user. */ export class OmnigentPlugin extends Plugin { static manifest = defineManifest<"omnigent">(manifest); @@ -41,6 +62,11 @@ export class OmnigentPlugin extends Plugin { declare protected config: IOmnigentConfig; private runtime?: OmnigentRuntime; + private readonly tokens = new UserTokens(); + private readonly keys = new SessionKeys( + process.env.DATABRICKS_CLIENT_SECRET || process.env.OMNIGENT_KEY_SECRET, + ); + private readonly mcp = new McpToolServer(() => this.mcpTools()); private models?: { at: number; list: GatewayModel[] }; async setup(): Promise { @@ -168,18 +194,77 @@ export class OmnigentPlugin extends Plugin { }); } + /** The app's tools plus opted-in plugins' tools, under MCP-safe names. */ + private mcpTools(): McpTool[] { + const tools: McpTool[] = []; + const registry = this.config.tools ?? {}; + for (const def of toolsFromRegistry(registry)) { + tools.push({ + name: mcpToolName(def.name), + definition: def, + run: (args, signal) => + executeFromRegistry(registry, def.name, args, signal), + }); + } + const wanted = this.config.pluginTools; + if (wanted && this.context) { + for (const { name, provider } of this.context.getToolProviders()) { + if ( + name === this.name || + (Array.isArray(wanted) && !wanted.includes(name)) + ) + continue; + for (const def of provider.getAgentTools()) { + tools.push({ + name: mcpToolName(def.name, name), + definition: def, + run: (args, signal) => + provider.executeAgentTool(def.name, args, signal), + }); + } + } + } + return tools; + } + + private writeToolNames(): string[] { + return this.mcpTools() + .filter((t) => isWriteTool(t.definition)) + .map((t) => t.name); + } + + /** The port the app actually listens on (dev servers fall back to a free one). */ + private appPort(): number { + const server = this.context?.getPlugins().get("server") as + | { getServer?: () => { address(): unknown } | null } + | undefined; + try { + const addr = server?.getServer?.()?.address(); + if (addr && typeof addr === "object" && "port" in addr) + return Number((addr as { port: number }).port); + } catch { + // server not listening yet + } + return Number(process.env.DATABRICKS_APP_PORT || process.env.PORT || 8000); + } + 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; + "You are an assistant inside this Databricks app. Use the app's tools to answer from its data."; + return email + ? `${base}\n\nThe signed-in user is ${email}. Tool calls run with their permissions.` + : base; } private bundle( + user: string, email: string | undefined, harness: string, model: string, ): Buffer { + const key = this.keys.mint(user); + const port = this.appPort(); return buildBundle({ name: "appkit", description: "AppKit app agent", @@ -187,9 +272,49 @@ export class OmnigentPlugin extends Plugin { model, profile: profileFor(harness), instructions: this.instructions(email), + mcpServers: [ + { + name: "app", + description: "This app's tools, run as the signed-in user", + url: `http://127.0.0.1:${port}/api/${this.name}/mcp`, + headers: this.keys.headers(key), + timeout: 600, + }, + ], }); } + private async setMode( + user: string, + sid: string, + mode: OmnigentMode, + ): Promise { + const rt = this.rt(); + const have = (await ( + await rt.request("GET", `/v1/sessions/${sid}/policies`, user) + ).json()) as { + data?: Array<{ id: string; name?: string }>; + }; + for (const p of have.data ?? []) { + if (String(p.name ?? "").startsWith(MODE_POLICY_PREFIX)) { + await rt.request( + "DELETE", + `/v1/sessions/${sid}/policies/${p.id}`, + user, + ); + } + } + for (const p of modePolicies(mode, this.writeToolNames())) { + const r = await rt.request("POST", `/v1/sessions/${sid}/policies`, user, { + json: p, + }); + if (!r.ok) + throw new Error( + `Could not set ${mode} mode: ${r.status} ${await r.text()}`, + ); + } + } + /** * Waits until the session's runner has finished initializing. A turn posted * while it initializes is delivered twice (initialization replays the history @@ -267,6 +392,7 @@ export class OmnigentPlugin extends Plugin { let user: string; try { user = this.resolveUserId(req); + this.tokens.remember(req, user); } catch { res.status(401).json({ error: "Sign in to the app to use its agent." }); return; @@ -304,6 +430,8 @@ export class OmnigentPlugin extends Plugin { const harnesses = await this.offers(user); res.json({ harnesses, + modes: OMNIGENT_MODES, + defaultMode: this.config.defaultMode ?? "ask", defaultHarness: harnesses.find((h) => h.id === this.config.defaultHarness)?.id ?? harnesses[0]?.id, @@ -354,6 +482,11 @@ export class OmnigentPlugin extends Plugin { }); return; } + const mode = body.mode ?? this.config.defaultMode ?? "ask"; + if (!OMNIGENT_MODES.includes(mode)) { + res.status(400).json({ error: `Unknown mode: ${mode}` }); + return; + } const rt = this.rt(); const host = await rt.ensureHost(user); const form = new FormData(); @@ -364,7 +497,7 @@ export class OmnigentPlugin extends Plugin { body.title ?? body.message?.split("\n")[0]?.slice(0, 80) ?? "New session", - labels: { app: this.name }, + labels: { app: this.name, mode }, host_type: "external", host_id: host.hostId, workspace: rt.layout.userWork(host.userId), @@ -375,7 +508,12 @@ export class OmnigentPlugin extends Plugin { new Blob( [ new Uint8Array( - this.bundle(req.header("x-forwarded-email"), harness.id, model), + this.bundle( + user, + req.header("x-forwarded-email"), + harness.id, + model, + ), ), ], { @@ -394,6 +532,7 @@ export class OmnigentPlugin extends Plugin { const { session_id: sid } = (await created.json()) as { session_id: string; }; + await this.setMode(user, sid, mode); if (body.message) { // A turn delivered before the session's runner is up can be rejected (HTTP 204); wait briefly. await this.runnerReady(user, sid, 5_000); @@ -410,7 +549,9 @@ export class OmnigentPlugin extends Plugin { return; } } - res.status(201).json({ session_id: sid, harness: harness.id, model }); + res + .status(201) + .json({ session_id: sid, harness: harness.id, model, mode }); }), }); @@ -421,7 +562,17 @@ export class OmnigentPlugin extends Plugin { handler: this.handler(async (req, res, user) => { const snap = await this.owned(user, req.params.id, res); if (!snap) return; - res.json(snap); + const pol = (await ( + await this.rt().request( + "GET", + `/v1/sessions/${req.params.id}/policies`, + user, + ) + ).json()) as { data?: Array<{ name?: string }> }; + res.json({ + ...snap, + mode: modeFromPolicies((pol.data ?? []).map((p) => String(p.name))), + }); }), }); @@ -485,6 +636,47 @@ export class OmnigentPlugin extends Plugin { }), }); + this.route(router, { + name: "resolveElicitation", + method: "post", + path: "/sessions/:id/elicitations/:eid", + handler: this.handler(async (req, res, user) => { + 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; + } + await this.rt().ensureHost(user); + await OmnigentPlugin.relay( + res, + await this.rt().request( + "POST", + `/v1/sessions/${encodeURIComponent(req.params.id)}/elicitations/${encodeURIComponent(req.params.eid)}/resolve`, + user, + { json: { action } }, + ), + ); + }), + }); + + this.route(router, { + name: "setMode", + method: "put", + path: "/sessions/:id/mode", + handler: this.handler(async (req, res, user) => { + 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(user, req.params.id, res))) return; + await this.setMode(user, req.params.id, mode); + res.json({ mode }); + }), + }); + this.route(router, { name: "stream", method: "get", @@ -518,6 +710,49 @@ export class OmnigentPlugin extends Plugin { .pipe(res); }), }); + + // Tool calls from the harnesses (loopback). The session key is the only guard: + // Apps proxy traffic also arrives from 127.0.0.1. + this.route(router, { + name: "mcp", + method: "post", + path: "/mcp", + handler: async (req: express.Request, res: express.Response) => { + const user = this.keys.verify(req.headers); + if (!user) { + res.status(401).json({ + jsonrpc: "2.0", + id: null, + error: { code: -32001, message: "Invalid session key" }, + }); + return; + } + const ac = new AbortController(); + res.on("close", () => ac.abort()); + let reply: unknown; + try { + const userContext = this.tokens.userContext(user); + const run = () => this.mcp.handle(req.body, ac.signal); + reply = userContext + ? await runInUserContext(userContext, run) + : await run(); + } catch (err) { + reply = { + jsonrpc: "2.0", + id: (req.body as { id?: unknown })?.id ?? null, + error: { + code: -32002, + message: err instanceof Error ? err.message : String(err), + }, + }; + } + if (reply === null) { + res.status(202).end(); + return; + } + res.json(reply); + }, + }); } exports() { 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/tests/plugin.test.ts b/packages/appkit/src/plugins/omnigent/tests/plugin.test.ts index 797eb0365..11e8879d3 100644 --- a/packages/appkit/src/plugins/omnigent/tests/plugin.test.ts +++ b/packages/appkit/src/plugins/omnigent/tests/plugin.test.ts @@ -6,15 +6,46 @@ import { setupDatabricksEnv, } from "@tools/test-helpers"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; +import { z } from "zod"; +import { + getExecutionContext, + isInUserContext, +} from "../../../context/execution-context"; import { ServiceContext } from "../../../context/service-context"; +import { defineTool } from "../../../core/agent/tools/define-tool"; +import { SessionKeys } from "../keys"; import { OmnigentPlugin } from "../omnigent"; +const SECRET = "test-key-secret"; + function plugin() { - const p = new OmnigentPlugin({}); + const seen: Array<{ inUserContext: boolean; client: unknown }> = []; + const p = new OmnigentPlugin({ + tools: { + whoami: defineTool({ + description: "Who is calling", + schema: z.object({}), + annotations: { effect: "read" }, + execute: async () => { + seen.push({ + inUserContext: isInUserContext(), + client: getExecutionContext().client, + }); + return { ok: true }; + }, + }), + "notes.save": defineTool({ + description: "Save", + schema: z.object({ text: z.string() }), + annotations: { effect: "write" }, + execute: async () => ({ saved: true }), + }), + }, + }); const { router, getHandler } = createMockRouter(); p.injectRoutes(router); - return { p, getHandler }; + return { p, getHandler, seen }; } async function call( @@ -32,15 +63,18 @@ describe("OmnigentPlugin routes", () => { beforeEach(async () => { setupDatabricksEnv(); ServiceContext.reset(); + process.env.OMNIGENT_KEY_SECRET = SECRET; + delete process.env.DATABRICKS_CLIENT_SECRET; ctx = await mockServiceContext({}); }); afterEach(() => { ctx?.restore(); + delete process.env.OMNIGENT_KEY_SECRET; vi.restoreAllMocks(); }); - test("registers the session API", () => { + test("registers the session API and the MCP endpoint", () => { const { getHandler } = plugin(); for (const [m, path] of [ ["GET", "/status"], @@ -51,23 +85,117 @@ describe("OmnigentPlugin routes", () => { ["DELETE", "/sessions/:id"], ["POST", "/sessions/:id/messages"], ["POST", "/sessions/:id/interrupt"], + ["POST", "/sessions/:id/elicitations/:eid"], + ["PUT", "/sessions/:id/mode"], ["GET", "/sessions/:id/stream"], + ["POST", "/mcp"], ]) { expect(getHandler(m, path), `${m} ${path}`).toBeTypeOf("function"); } }); + test("a signed-in request before the runtime starts gets 503, not a crash", async () => { + const { getHandler } = plugin(); + const res = await call(getHandler("GET", "/sessions"), { + obo: { token: "t", userId: "u", email: "a@example.com" }, + }); + expect(res.status).toHaveBeenCalledWith(503); + }); + test("session routes need a forwarded identity", async () => { const { getHandler } = plugin(); const res = await call(getHandler("GET", "/sessions"), {}); expect(res.status).toHaveBeenCalledWith(401); }); - test("a signed-in request before the runtime starts gets 503, not a crash", async () => { + test("MCP rejects calls without a valid session key", async () => { const { getHandler } = plugin(); - const res = await call(getHandler("GET", "/sessions"), { - obo: { token: "t", userId: "u", email: "a@example.com" }, + const body = { jsonrpc: "2.0", id: 1, method: "tools/list" }; + let res = await call(getHandler("POST", "/mcp"), { body }); + expect(res.status).toHaveBeenCalledWith(401); + const forged = new SessionKeys("some-other-secret").mint("alice-id"); + res = await call(getHandler("POST", "/mcp"), { + body, + headers: new SessionKeys().headers(forged), }); - expect(res.status).toHaveBeenCalledWith(503); + expect(res.status).toHaveBeenCalledWith(401); + }); + + test("MCP runs tools as the key's user, with their latest forwarded token", async () => { + const { getHandler, seen } = plugin(); + const keys = new SessionKeys(SECRET); + const headers = keys.headers(keys.mint("alice-id")); + + // Before alice has made a request, there is no token to act with. + let res = await call(getHandler("POST", "/mcp"), { + body: { + jsonrpc: "2.0", + id: 1, + method: "tools/call", + params: { name: "whoami" }, + }, + headers, + }); + expect(res.json.mock.calls[0][0].error.message).toMatch( + /Open the app again/, + ); + + // A request from alice records her identity (the route itself fails: no runtime). + await call(getHandler("GET", "/sessions"), { + obo: { + token: "alice-token", + userId: "alice-id", + email: "alice@example.com", + }, + }); + + res = await call(getHandler("POST", "/mcp"), { + body: { + jsonrpc: "2.0", + id: 2, + method: "tools/call", + params: { name: "whoami", arguments: {} }, + }, + headers, + }); + const reply = res.json.mock.calls[0][0]; + expect(reply.result.content[0].text).toBe('{"ok":true}'); + expect(seen).toHaveLength(1); + expect(seen[0].inUserContext).toBe(true); + expect(ctx.createUserContextSpy).toHaveBeenCalledWith( + "alice-token", + "alice-id", + undefined, + "alice@example.com", + ); + }); + + test("tools/list serves MCP-safe names", async () => { + const { getHandler } = plugin(); + const keys = new SessionKeys(SECRET); + await call(getHandler("GET", "/sessions"), { + obo: { token: "t", userId: "u", email: "bob@example.com" }, + }); + const res = await call(getHandler("POST", "/mcp"), { + body: { jsonrpc: "2.0", id: 1, method: "tools/list" }, + headers: keys.headers(keys.mint("u")), + }); + const names = res.json.mock.calls[0][0].result.tools.map( + (t: { name: string }) => t.name, + ); + expect(names).toEqual(["whoami", "notes_save"]); + }); + + test("notifications get 202 with no body", async () => { + const { getHandler } = plugin(); + const keys = new SessionKeys(SECRET); + await call(getHandler("GET", "/sessions"), { + obo: { token: "t", userId: "u", email: "c@example.com" }, + }); + const res = await call(getHandler("POST", "/mcp"), { + body: { jsonrpc: "2.0", method: "notifications/initialized" }, + headers: keys.headers(keys.mint("u")), + }); + expect(res.status).toHaveBeenCalledWith(202); }); }); diff --git a/packages/appkit/src/plugins/omnigent/tests/units.test.ts b/packages/appkit/src/plugins/omnigent/tests/units.test.ts index 8fb82e2c2..524bd2413 100644 --- a/packages/appkit/src/plugins/omnigent/tests/units.test.ts +++ b/packages/appkit/src/plugins/omnigent/tests/units.test.ts @@ -4,9 +4,17 @@ import path from "node:path"; import { gunzipSync } from "node:zlib"; import { afterEach, describe, expect, test } from "vitest"; +import { z } from "zod"; +import { + defineTool, + toolsFromRegistry, +} from "../../../core/agent/tools/define-tool"; import { buildBundle } from "../bundle"; import { offerHarnesses } from "../harnesses"; +import { MCP_KEY_HEADER, MCP_USER_HEADER, SessionKeys } from "../keys"; +import { isWriteTool, McpToolServer, mcpToolName } from "../mcp-server"; +import { modeFromPolicies, modePolicies } from "../policies"; import { childEnv, profileFor, writeUserConfig } from "../runtime/env"; import { HarnessSandbox } from "../runtime/sandbox"; @@ -32,8 +40,65 @@ function untar(buf: Buffer): Record { return out; } +describe("SessionKeys", () => { + test("a minted key verifies for its user only", () => { + const keys = new SessionKeys("secret"); + const k = keys.mint("Alice@Example.com"); + expect(keys.verify(keys.headers(k))).toBe("Alice@Example.com"); + expect( + keys.verify({ ...keys.headers(k), [MCP_USER_HEADER]: "bob@example.com" }), + ).toBeNull(); + expect( + keys.verify({ ...keys.headers(k), [MCP_KEY_HEADER]: "x".repeat(43) }), + ).toBeNull(); + expect(keys.verify({})).toBeNull(); + }); + + test("keys survive a restart when the secret is stable", () => { + const k = new SessionKeys("stable").mint("a@b.c"); + expect( + new SessionKeys("stable").verify(new SessionKeys("stable").headers(k)), + ).toBe("a@b.c"); + expect( + new SessionKeys("other").verify(new SessionKeys("stable").headers(k)), + ).toBeNull(); + }); +}); + +describe("policies", () => { + test("ask gates write tools and OS tools", () => { + const p = modePolicies("ask", ["save_note"]); + expect(p.map((x) => x.name)).toEqual(["appkit-mode-ask", "appkit-mode-os"]); + expect(String(p[0].factory_params?.expression)).toContain( + 'endsWith("save_note")', + ); + 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, + ); + } + }); +}); + describe("buildBundle", () => { - test("writes the spec and the instructions", () => { + test("writes spec, instructions and one file per MCP server", () => { const files = untar( gunzipSync( buildBundle({ @@ -42,10 +107,21 @@ describe("buildBundle", () => { model: "system.ai.gpt-5-5", profile: profileFor("codex"), instructions: "Be terse.", + mcpServers: [ + { + name: "app", + url: "http://127.0.0.1:8000/api/omnigent/mcp", + headers: { a: "b" }, + }, + ], }), ), ); - expect(Object.keys(files).sort()).toEqual(["AGENTS.md", "config.yaml"]); + expect(Object.keys(files).sort()).toEqual([ + "AGENTS.md", + "config.yaml", + "tools/mcp/app.yaml", + ]); const cfg = JSON.parse(files["config.yaml"]); expect(cfg.executor).toMatchObject({ model: "system.ai.gpt-5-5", @@ -53,9 +129,26 @@ describe("buildBundle", () => { config: { harness: "codex" }, auth: { type: "databricks", profile: "appkit-omnigent-http" }, }); + expect(JSON.parse(files["tools/mcp/app.yaml"])).toMatchObject({ + transport: "http", + headers: { a: "b" }, + }); expect(files["AGENTS.md"]).toBe("Be terse."); }); + test("rejects unsafe MCP server names", () => { + expect(() => + buildBundle({ + name: "a", + harness: "pi", + model: "m", + profile: "p", + instructions: "", + mcpServers: [{ name: "../x", url: "u" }], + }), + ).toThrow(/Invalid MCP server name/); + }); + test("Claude and Pi use the https profile, Codex the http one", () => { expect(profileFor("claude-sdk")).toBe("appkit-omnigent"); expect(profileFor("pi")).toBe("appkit-omnigent"); @@ -63,6 +156,121 @@ describe("buildBundle", () => { }); }); +describe("McpToolServer", () => { + const registry = { + "notes.save": defineTool({ + description: "Save a note", + schema: z.object({ text: z.string() }), + annotations: { effect: "write" }, + execute: async ({ text }) => ({ saved: text }), + }), + boom: defineTool({ + description: "Fails", + schema: z.object({}), + execute: async () => { + throw new Error("nope"); + }, + }), + }; + const server = new McpToolServer(() => + toolsFromRegistry(registry).map((definition) => ({ + name: mcpToolName(definition.name), + definition, + run: async (args) => + (registry as Record unknown }>)[ + definition.name + ].execute(args as never), + })), + ); + + test("initialize echoes the client's protocol version", async () => { + const r = (await server.handle({ + jsonrpc: "2.0", + id: 1, + method: "initialize", + params: { protocolVersion: "2025-11-25" }, + })) as any; + expect(r.result.protocolVersion).toBe("2025-11-25"); + expect(r.result.capabilities.tools).toBeDefined(); + }); + + test("tools/list uses MCP-safe names and marks writes", async () => { + const r = (await server.handle({ + jsonrpc: "2.0", + id: 2, + method: "tools/list", + })) as any; + const names = r.result.tools.map((t: { name: string }) => t.name); + expect(names).toEqual(["notes_save", "boom"]); + expect(r.result.tools[0].annotations.readOnlyHint).toBe(false); + expect(r.result.tools[1].annotations.readOnlyHint).toBe(true); + }); + + test("tools/call returns results and tool errors as isError", async () => { + const ok = (await server.handle({ + jsonrpc: "2.0", + id: 3, + method: "tools/call", + params: { name: "notes_save", arguments: { text: "hi" } }, + })) as any; + expect(JSON.parse(ok.result.content[0].text)).toEqual({ saved: "hi" }); + const bad = (await server.handle({ + jsonrpc: "2.0", + id: 4, + method: "tools/call", + params: { name: "boom" }, + })) as any; + expect(bad.result.isError).toBe(true); + expect(bad.result.content[0].text).toBe("nope"); + }); + + test("notifications get no reply; unknown methods and tools are errors", async () => { + expect( + await server.handle({ + jsonrpc: "2.0", + method: "notifications/initialized", + }), + ).toBeNull(); + expect( + ((await server.handle({ jsonrpc: "2.0", id: 5, method: "nope" })) as any) + .error.code, + ).toBe(-32601); + expect( + ( + (await server.handle({ + jsonrpc: "2.0", + id: 6, + method: "tools/call", + params: { name: "x" }, + })) as any + ).error.code, + ).toBe(-32602); + }); + + test("rejects JSON-RPC batches (not part of MCP)", async () => { + const r = (await server.handle([ + { jsonrpc: "2.0", id: 7, method: "ping" }, + ])) as any; + expect(r.error.code).toBe(-32600); + }); + + test("isWriteTool honours effect and the deprecated flags", () => { + const def = (annotations: object) => + ({ name: "t", description: "", parameters: {}, annotations }) as any; + expect(isWriteTool(def({ effect: "destructive" }))).toBe(true); + expect(isWriteTool(def({ effect: "read" }))).toBe(false); + expect(isWriteTool(def({ destructive: true }))).toBe(true); + expect(isWriteTool(def({}))).toBe(false); + }); + + test("mcpToolName prefixes plugin tools and stays within 64 chars", () => { + expect(mcpToolName("default.sendMessage", "genie")).toBe( + "genie__default_sendMessage", + ); + expect(mcpToolName("x".repeat(100))).toHaveLength(64); + }); +}); + describe("child environment", () => { const dirs: string[] = []; afterEach(() => { diff --git a/packages/appkit/src/plugins/omnigent/tokens.ts b/packages/appkit/src/plugins/omnigent/tokens.ts new file mode 100644 index 000000000..8c2520d47 --- /dev/null +++ b/packages/appkit/src/plugins/omnigent/tokens.ts @@ -0,0 +1,42 @@ +import type express from "express"; + +import { ServiceContext } from "../../context/service-context"; +import type { UserContext } from "../../context/user-context"; + +/** + * The latest forwarded token of each user, in memory only. + * + * Tool calls arrive from a harness, not from the user's browser, so they carry + * no forwarded token. Each request the user makes to the app refreshes theirs + * here, and a tool call runs with it. After a restart the map is empty until + * the user's next request; a tool call before that fails with a clear error. + */ +export class UserTokens { + private readonly byUser = new Map< + string, + { token: string; email?: string } + >(); + + /** Remembers the request's forwarded token for `user`. */ + remember(req: express.Request, user: string): void { + const token = req.header("x-forwarded-access-token")?.trim(); + if (token) + this.byUser.set(user, { token, email: req.header("x-forwarded-email") }); + } + + /** + * An AppKit user context for running a tool as `user`. In development + * without a forwarded token, `undefined`: the tool runs as the app, like + * `asUser()` does there. + */ + userContext(user: string): UserContext | undefined { + const t = this.byUser.get(user); + if (!t) { + if (process.env.NODE_ENV === "development") return undefined; + throw new Error( + "No current sign-in for this session's user. Open the app again to continue.", + ); + } + return ServiceContext.createUserContext(t.token, user, undefined, t.email); + } +} diff --git a/packages/appkit/src/plugins/omnigent/types.ts b/packages/appkit/src/plugins/omnigent/types.ts index 0cecdae82..292d41769 100644 --- a/packages/appkit/src/plugins/omnigent/types.ts +++ b/packages/appkit/src/plugins/omnigent/types.ts @@ -1,8 +1,27 @@ import type { BasePluginConfig } from "shared"; +import type { ToolRegistry } from "../../core/agent/tools/define-tool"; +import type { OmnigentMode } from "./policies"; + +export type { OmnigentMode } from "./policies"; + export interface IOmnigentConfig extends BasePluginConfig { /** System instructions for every session (the bundle's `AGENTS.md`). */ instructions?: string; + /** + * The app's own tools, served to every harness over the plugin's MCP + * endpoint and run as the session's user. Same `defineTool()` entries the + * agents plugin takes. Tools with `annotations.effect` of `write`, `update` + * or `destructive` are gated by `ask` and `read` modes. + */ + tools?: ToolRegistry; + /** + * Also serve the tools of other registered AppKit plugins (e.g. `genie`). + * `true` for all tool providers, or a list of plugin names. Default: none. + */ + pluginTools?: boolean | string[]; + /** Mode new sessions start in. Default `ask`. */ + defaultMode?: OmnigentMode; /** * Harnesses users may pick, out of the ones that route model calls through * Unity AI Gateway. Default: all of them whose CLI is installed. @@ -32,6 +51,7 @@ export interface OmnigentHarness { export interface CreateSessionRequest { harness?: string; model?: string; + mode?: OmnigentMode; message?: string; title?: string; } diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index af1e98032..740bb962c 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -345,7 +345,7 @@ "omnigent": { "name": "omnigent", "displayName": "Omnigent Plugin", - "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Codex, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway", + "description": "Embedded Omnigent agent runtime: Claude Agent SDK, Codex, Pi and OpenAI Agents SDK harnesses through Unity AI Gateway, with the app's tools run as the signed-in user", "package": "@databricks/appkit", "resources": { "required": [], @@ -357,7 +357,8 @@ "must": [ "Add pyproject.toml (python >=3.12,<3.13) with omnigent[databricks] and a uv.lock that uses public PyPI", "Add harness CLIs to package.json: @anthropic-ai/claude-code, @openai/codex, @earendil-works/pi-coding-agent", - "Start the app with node, not npm run, so SIGTERM reaches the plugin and it can stop its child processes" + "Start the app with node, not npm run, so SIGTERM reaches the plugin and it can stop its child processes", + "Grant the user API scopes the app's tools need (e.g. sql) so tools run as the signed-in user" ] } } From ac48f7e1688bcea0ded905c366e91d80d6060f8a Mon Sep 17 00:00:00 2001 From: Noah Solomon Date: Wed, 30 Sep 2026 00:05:13 +1000 Subject: [PATCH 4/4] feat(appkit-ui): add omnigent hooks, playground page and docs (4/4) - useOmnigentSession, useOmnigentHarnesses and useOmnigentSessions in @databricks/appkit-ui/react/beta (headless, like useAgentChat). - A dev-playground /omnigent page, registered only with OMNIGENT_ENABLED=1. - A React section in docs/docs/plugins/omnigent.md. Co-authored-by: Isaac Signed-off-by: Noah Solomon --- apps/dev-playground/client/src/lib/nav.ts | 8 + .../client/src/routeTree.gen.ts | 21 + .../client/src/routes/omnigent.route.tsx | 311 ++++++++++ apps/dev-playground/server/index.ts | 40 +- docs/docs/plugins/omnigent.md | 34 ++ packages/appkit-ui/src/react/beta.ts | 20 + .../__tests__/use-omnigent-session.test.ts | 334 +++++++++++ .../src/react/hooks/use-omnigent-session.ts | 536 ++++++++++++++++++ 8 files changed, 1303 insertions(+), 1 deletion(-) create mode 100644 apps/dev-playground/client/src/routes/omnigent.route.tsx create mode 100644 packages/appkit-ui/src/react/hooks/__tests__/use-omnigent-session.test.ts create mode 100644 packages/appkit-ui/src/react/hooks/use-omnigent-session.ts diff --git a/apps/dev-playground/client/src/lib/nav.ts b/apps/dev-playground/client/src/lib/nav.ts index 5e6eda2a9..c446a89d5 100644 --- a/apps/dev-playground/client/src/lib/nav.ts +++ b/apps/dev-playground/client/src/lib/nav.ts @@ -18,6 +18,7 @@ import { SigmaIcon, Wand2Icon, ZapIcon, + TerminalIcon, } from "lucide-react"; /** @@ -123,6 +124,13 @@ export const NAV_GROUPS: ReadonlyArray = [ "Chat agent over Databricks Model Serving with tools auto-discovered from AppKit plugins.", icon: BotIcon, }, + { + to: "/omnigent", + label: "Omnigent", + description: + "Claude Agent SDK, Codex and Pi sessions through Unity AI Gateway, with the app's tools run as you.", + icon: TerminalIcon, + }, { to: "/genie", label: "Genie", diff --git a/apps/dev-playground/client/src/routeTree.gen.ts b/apps/dev-playground/client/src/routeTree.gen.ts index 8658c1952..f1896ed1c 100644 --- a/apps/dev-playground/client/src/routeTree.gen.ts +++ b/apps/dev-playground/client/src/routeTree.gen.ts @@ -18,6 +18,7 @@ import { Route as ServingRouteRouteImport } from './routes/serving.route' import { Route as ReconnectRouteRouteImport } from './routes/reconnect.route' import { Route as QueryDedupRouteRouteImport } from './routes/query-dedup.route' import { Route as PolicyMatrixRouteRouteImport } from './routes/policy-matrix.route' +import { Route as OmnigentRouteRouteImport } from './routes/omnigent.route' import { Route as MetricViewsRouteRouteImport } from './routes/metric-views.route' import { Route as LakebaseRouteRouteImport } from './routes/lakebase.route' import { Route as JobsRouteRouteImport } from './routes/jobs.route' @@ -77,6 +78,11 @@ const PolicyMatrixRouteRoute = PolicyMatrixRouteRouteImport.update({ path: '/policy-matrix', getParentRoute: () => rootRouteImport, } as any) +const OmnigentRouteRoute = OmnigentRouteRouteImport.update({ + id: '/omnigent', + path: '/omnigent', + getParentRoute: () => rootRouteImport, +} as any) const MetricViewsRouteRoute = MetricViewsRouteRouteImport.update({ id: '/metric-views', path: '/metric-views', @@ -157,6 +163,7 @@ export interface FileRoutesByFullPath { '/jobs': typeof JobsRouteRoute '/lakebase': typeof LakebaseRouteRoute '/metric-views': typeof MetricViewsRouteRoute + '/omnigent': typeof OmnigentRouteRoute '/policy-matrix': typeof PolicyMatrixRouteRoute '/query-dedup': typeof QueryDedupRouteRoute '/reconnect': typeof ReconnectRouteRoute @@ -181,6 +188,7 @@ export interface FileRoutesByTo { '/jobs': typeof JobsRouteRoute '/lakebase': typeof LakebaseRouteRoute '/metric-views': typeof MetricViewsRouteRoute + '/omnigent': typeof OmnigentRouteRoute '/policy-matrix': typeof PolicyMatrixRouteRoute '/query-dedup': typeof QueryDedupRouteRoute '/reconnect': typeof ReconnectRouteRoute @@ -206,6 +214,7 @@ export interface FileRoutesById { '/jobs': typeof JobsRouteRoute '/lakebase': typeof LakebaseRouteRoute '/metric-views': typeof MetricViewsRouteRoute + '/omnigent': typeof OmnigentRouteRoute '/policy-matrix': typeof PolicyMatrixRouteRoute '/query-dedup': typeof QueryDedupRouteRoute '/reconnect': typeof ReconnectRouteRoute @@ -232,6 +241,7 @@ export interface FileRouteTypes { | '/jobs' | '/lakebase' | '/metric-views' + | '/omnigent' | '/policy-matrix' | '/query-dedup' | '/reconnect' @@ -256,6 +266,7 @@ export interface FileRouteTypes { | '/jobs' | '/lakebase' | '/metric-views' + | '/omnigent' | '/policy-matrix' | '/query-dedup' | '/reconnect' @@ -280,6 +291,7 @@ export interface FileRouteTypes { | '/jobs' | '/lakebase' | '/metric-views' + | '/omnigent' | '/policy-matrix' | '/query-dedup' | '/reconnect' @@ -305,6 +317,7 @@ export interface RootRouteChildren { JobsRouteRoute: typeof JobsRouteRoute LakebaseRouteRoute: typeof LakebaseRouteRoute MetricViewsRouteRoute: typeof MetricViewsRouteRoute + OmnigentRouteRoute: typeof OmnigentRouteRoute PolicyMatrixRouteRoute: typeof PolicyMatrixRouteRoute QueryDedupRouteRoute: typeof QueryDedupRouteRoute ReconnectRouteRoute: typeof ReconnectRouteRoute @@ -381,6 +394,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof PolicyMatrixRouteRouteImport parentRoute: typeof rootRouteImport } + '/omnigent': { + id: '/omnigent' + path: '/omnigent' + fullPath: '/omnigent' + preLoaderRoute: typeof OmnigentRouteRouteImport + parentRoute: typeof rootRouteImport + } '/metric-views': { id: '/metric-views' path: '/metric-views' @@ -489,6 +509,7 @@ const rootRouteChildren: RootRouteChildren = { JobsRouteRoute: JobsRouteRoute, LakebaseRouteRoute: LakebaseRouteRoute, MetricViewsRouteRoute: MetricViewsRouteRoute, + OmnigentRouteRoute: OmnigentRouteRoute, PolicyMatrixRouteRoute: PolicyMatrixRouteRoute, QueryDedupRouteRoute: QueryDedupRouteRoute, ReconnectRouteRoute: ReconnectRouteRoute, diff --git a/apps/dev-playground/client/src/routes/omnigent.route.tsx b/apps/dev-playground/client/src/routes/omnigent.route.tsx new file mode 100644 index 000000000..c302c1747 --- /dev/null +++ b/apps/dev-playground/client/src/routes/omnigent.route.tsx @@ -0,0 +1,311 @@ +import { + Badge, + Button, + Card, + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, + Textarea, +} from "@databricks/appkit-ui/react"; +import { + type OmnigentItem, + type OmnigentMode, + useOmnigentHarnesses, + useOmnigentSession, + useOmnigentSessions, +} from "@databricks/appkit-ui/react/beta"; +import { createFileRoute } from "@tanstack/react-router"; +import { useEffect, useState } from "react"; + +export const Route = createFileRoute("/omnigent")({ + component: OmnigentRoute, +}); + +const MODE_LABELS: Record = { + auto: "Auto", + ask: "Ask before changes", + read: "Read only", +}; + +function text(item: OmnigentItem): string { + return (item.data?.content ?? []) + .map((c) => c.text ?? "") + .join(" ") + .trim(); +} + +function ItemView({ item }: { item: OmnigentItem }) { + const d = item.data ?? {}; + if (item.type === "message") { + const mine = d.role === "user"; + return ( +
+
+ {text(item)} +
+
+ ); + } + if (item.type === "function_call") { + return ( +
+ → {d.name}({String(d.arguments ?? "")}) +
+ ); + } + if (item.type === "function_call_output") { + return ( + 
+        {String(d.output ?? "")}
+
+ ); + } + if (item.type === "error") { + return ( +
+ {String(d.message ?? "Error")} +
+ ); + } + return null; +} + +function OmnigentRoute() { + const options = useOmnigentHarnesses(); + const threads = useOmnigentSessions(); + const [openId, setOpenId] = useState(null); + const s = useOmnigentSession({ sessionId: openId }); + + const [harness, setHarness] = useState(""); + const [model, setModel] = useState(""); + const [mode, setMode] = useState("ask"); + const [draftText, setDraftText] = useState(""); + + const current = options.harnesses.find((h) => h.id === harness); + + // Defaults once the options load; follow an opened session's settings. + useEffect(() => { + if (!harness && options.defaultHarness) setHarness(options.defaultHarness); + if (!openId) setMode(options.defaultMode); + }, [options.defaultHarness, options.defaultMode, harness, openId]); + useEffect(() => { + if (current && !current.models.includes(model)) + setModel(current.defaultModel ?? current.models[0] ?? ""); + }, [current, model]); + useEffect(() => { + if (s.session?.harness) setHarness(s.session.harness); + if (s.session?.mode) setMode(s.session.mode); + }, [s.session?.harness, s.session?.mode]); + + const submit = async () => { + const message = draftText.trim(); + if (!message) return; + setDraftText(""); + if (s.sessionId) { + await s.send(message); + } else { + const id = await s.start({ harness, model, mode, message }); + setOpenId(id); + void threads.reload(); + } + }; + + const changeMode = async (next: OmnigentMode) => { + setMode(next); + if (s.sessionId) await s.setMode(next); + }; + + return ( +
+
+
+
+

+ Omnigent +

+

+ Claude Agent SDK, Codex, Pi and OpenAI Agents SDK sessions, with + model calls through Unity AI Gateway and the app's tools run as + you. +

+
+ + {options.error && ( + + The omnigent plugin is not available: {options.error}. Start the + playground with OMNIGENT_ENABLED=1 (and a Python 3.12 venv with + omnigent installed). + + )} + +
+ + + {threads.sessions.map((t) => ( + + ))} + + + +
+ + + + {s.status && ( + + {s.status} + + )} +
+ +
+ {s.items.map((it, i) => ( + + ))} + {s.draft && ( +
+ {s.draft} +
+ )} + {s.approvals.map((a) => ( + +
+ Approval needed: {a.params?.message} +
+ {a.params?.content_preview && ( + 
+                        {a.params.content_preview}
+
+ )} +
+ + +
+ + ))} + {s.error && ( +
{s.error}
+ )} +
+ +
+