[appkit plugin] [5/x] Add chat server plugin to @databricks/appkit-agent by hubertzub-db · Pull Request #387 · databricks/databricks-ai-bridge

hubertzub-db · 2026-03-20T13:58:26Z

🥞 Stacked PR

Summary

Adds a ChatPlugin to @databricks/appkit-agent, exported from the @databricks/appkit-agent/chat subpath. The plugin provides a complete chat backend: streaming responses from agent or Databricks model serving endpoints, session management, optional PostgreSQL persistence, feedback via MLflow Traces, and stream resumption.

Adapted from app-templates/e2e-chatbot-app-next](https://github.com/databricks/app-templates/tree/main/e2e-chatbot-app-next), with @databricks/appkit as a peer dependency and a standalone auth module replacing internal imports.

Usage

import { createApp } from "@databricks/appkit";
import { agent, chat } from "@databricks/appkit-agent";
import { Pool } from "pg";

const app = createApp({
  plugins: [
    agent({ /* ... */ }),
    chat({
      pool: new Pool(),       // optional — omit for ephemeral (no history)
      autoMigrate: true,      // creates ai_chatbot schema/tables on startup
      backend: "agent",       // route chat through local agent plugin
    }),
  ],
});

Demo of the chat with persistence and tools (the UI runs from existing `e2e-chatbot-app-next`)

Screen.Recording.2026-03-25.at.11.04.21.mov

What's included

ChatPlugin class extending AppKit's Plugin<ChatConfig>, registered as "chat". Exposes a chat() factory via toPlugin() and wires 13 Express routes under /api/chat/.
The plugin code is mostly copied from https://github.com/databricks/app-templates/tree/main/e2e-chatbot-app-next/server
Exposed API uses https://ai-sdk.dev/ (like the original)
Migrations are handled by drizzle-kit and enabled using autoMigrate flag

hubertzub-db · 2026-03-24T15:00:11Z

integrations/appkit-agent/src/chat-plugin/auth.ts

@@ -0,0 +1,299 @@
+/**


note: moved from https://github.com/databricks/app-templates/tree/main/e2e-chatbot-app-next/packages/auth

hubertzub-db · 2026-03-24T15:00:53Z

integrations/appkit-agent/src/chat-plugin/chat-acl.ts

@@ -0,0 +1,28 @@
+import type { ChatRow } from "./schema";


note: moved from https://github.com/databricks/app-templates/blob/main/e2e-chatbot-app-next/packages/core/src/chat-acl.ts

hubertzub-db · 2026-03-24T15:01:35Z

integrations/appkit-agent/src/chat-plugin/message-meta.ts

@@ -0,0 +1,52 @@
+const MAX_ENTRIES = 10_000;
+
+function setBounded<K, V>(map: Map<K, V>, key: K, value: V): void {


note: moved from https://github.com/databricks/app-templates/blob/main/e2e-chatbot-app-next/server/src/lib/message-meta-store.ts

hubertzub-db · 2026-03-24T15:06:09Z

integrations/appkit-agent/src/chat-plugin/persistence.ts

+
+export type { ChatRow, MessageRow };
+
+export function createDb(pool: Pool) {


note: reimplemented from https://github.com/databricks/app-templates/tree/main/e2e-chatbot-app-next/packages/db

bbqiu

lgtm with a few nits!

bbqiu · 2026-03-31T08:27:37Z

integrations/appkit-agent/src/chat-plugin/auth.ts

+ * Consolidated Databricks authentication module.
+ *
+ * Supported auth methods (checked in priority order):
+ *   1. PAT            — DATABRICKS_TOKEN env var


nit: can we make pat the lowest priority?
oauth u2m, oauth sp, then pat

could we also add a test to codify this priority order in the future

reordered and added tests

bbqiu · 2026-03-31T08:41:18Z

integrations/appkit-agent/src/chat-plugin/index.ts

+    this.registerEndpoint("chat", `/api/${this.name}`);
+  }
+
+  private resolveBackend(req: express.Request): {


nit: can we add tests for this?

bbqiu · 2026-03-31T08:44:41Z

integrations/appkit-agent/src/chat-plugin/provider.ts

+}
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+type LanguageModel = any;


should this be the type from the ai-sdk?

correct, fixing now

bbqiu · 2026-03-31T08:47:14Z

integrations/appkit-agent/src/chat-plugin/session.ts

+ * Resolve session for the request.
+ *
+ * Resolution order:
+ * 1. Custom getSession callback (if provided in plugin config)


can we add some tests to codify this ordering? priority makes sense / is the same as the existing templat

bbqiu · 2026-03-31T08:51:07Z

integrations/appkit-agent/src/chat-plugin/index.ts

+    );
+    this.registerEndpoint("deleteChat", `/api/${this.name}/:id`);
+
+    // ── POST / (main chat handler) ───────────────────────────────────────


generally for this file, can we add tests for expected behavior of each route?

i think the post is probably the most important to test (ex. fallback behvior, what to do with MCP approvals / denials etc)

let me add some tests with simulated requests/responses

bbqiu · 2026-03-31T08:54:49Z

integrations/appkit-agent/src/chat-plugin/types.ts

+  /** Enable feedback feature (thumbs up/down). Defaults to !!process.env.MLFLOW_EXPERIMENT_ID. */
+  feedbackEnabled?: boolean;
+  /** Auto-create the ai_chatbot schema and tables on startup. Defaults to false. */
+  autoMigrate?: boolean;


should this default to true? in what case would we not want this to run

good point, changed to true as default and retested - seems to be idempotent and works well

Move agent plugin source code into src/agent-plugin/ subfolder for better organization when hosting multiple plugins.

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

Add per-package pnpm-lock.yaml and CI/release workflows with explicit pnpm version and cache-dependency-path.

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

This was referenced Mar 20, 2026

[appkit plugin] [2/x] Add @databricks/appkit-agent package #384

Open

[appkit plugin] [3/x] Add appkit-agent documentation #385

Open

[appkit plugin] [4/x] Add @databricks/chat-ui package #386

Open

hubertzub-db marked this pull request as draft March 20, 2026 13:59

hubertzub-db force-pushed the stack/agent-plugin-p5 branch 2 times, most recently from 2d61336 to fa5fffc Compare March 24, 2026 10:46

hubertzub-db mentioned this pull request Mar 24, 2026

[appkit plugin] [6/x] Add bundled chat UI to @databricks/appkit-agent #391

Open

hubertzub-db force-pushed the stack/agent-plugin-p5 branch 5 times, most recently from bceae2d to 2be31ff Compare March 24, 2026 12:11

hubertzub-db mentioned this pull request Mar 24, 2026

[appkit plugin] [7/x] Add default-ui playground for @databricks/appkit-agent #392

Open

hubertzub-db force-pushed the stack/agent-plugin-p5 branch from 2be31ff to 59b4915 Compare March 24, 2026 14:58

hubertzub-db commented Mar 24, 2026

View reviewed changes

hubertzub-db force-pushed the stack/agent-plugin-p5 branch from 1d83c9e to 61bf6ac Compare March 24, 2026 15:33

hubertzub-db changed the title ~~Add @databricks/appkit-chat-server package~~ [appkit plugin] [5/x] Add chat server plugin to @databricks/appkit-agent Mar 24, 2026

hubertzub-db requested a review from bbqiu March 24, 2026 16:02

hubertzub-db marked this pull request as ready for review March 25, 2026 08:06

hubertzub-db force-pushed the stack/agent-plugin-p5 branch 3 times, most recently from c469eda to c4d0cb8 Compare March 25, 2026 10:38

hubertzub-db mentioned this pull request Mar 25, 2026

[appkit plugin] [7/x] Add custom UI elements (+playground) to @databricks/appkit-agent #395

Open

3 tasks

bbqiu approved these changes Mar 31, 2026

View reviewed changes

bbqiu reviewed Mar 31, 2026

View reviewed changes

hubertzub-db added 2 commits April 1, 2026 11:00

Agent plugin

c7b0c54

Move agent plugin source code into src/agent-plugin/ subfolder for better organization when hosting multiple plugins.

CR feedback

eabd0b6

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

hubertzub-db added 9 commits April 1, 2026 11:00

Responses

c463ebf

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

docs

d9af352

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

CR feedback

96a56f3

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

simple chatui

adc9176

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

Standalone pnpm for chat-ui

64e1d95

Add per-package pnpm-lock.yaml and CI/release workflows with explicit pnpm version and cache-dependency-path.

Move to appkit-agent

ce40d32

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

Chat plugin

9a32b44

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

bugfixes

2e23f68

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

Backend option

e9154b6

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

hubertzub-db force-pushed the stack/agent-plugin-p5 branch from c4d0cb8 to b8cc404 Compare April 1, 2026 13:26

Cr feedback

638d1ee

Signed-off-by: Hubert Zub <hubert.zub@databricks.com>

hubertzub-db force-pushed the stack/agent-plugin-p5 branch from b8cc404 to 638d1ee Compare April 1, 2026 13:40

		@@ -0,0 +1,52 @@
		const MAX_ENTRIES = 10_000;

		function setBounded<K, V>(map: Map<K, V>, key: K, value: V): void {


		export type { ChatRow, MessageRow };

		export function createDb(pool: Pool) {

Conversation

hubertzub-db commented Mar 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🥞 Stacked PR

Summary

Usage

Demo of the chat with persistence and tools (the UI runs from existing e2e-chatbot-app-next)

What's included

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bbqiu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

hubertzub-db commented Mar 20, 2026 •

edited

Loading

Demo of the chat with persistence and tools (the UI runs from existing `e2e-chatbot-app-next`)