@strav/brain 1.0.0-alpha.16 → 1.0.0-alpha.18

package/src/mcp/oauth.ts

@@ -0,0 +1,227 @@
+/**
+ * OAuth support for the local MCP client.
+ *
+ * Most real-world MCP servers (Linear, Notion, GitHub, Asana,
+ * Atlassian) are OAuth-protected. The local client at
+ * `@strav/brain/mcp` already supports static bearer tokens via
+ * `MCPServer.authorizationToken` — fine for self-hosted servers,
+ * useless against any commercial server. This module closes the
+ * gap.
+ *
+ * The flow apps see:
+ *
+ * ```ts
+ * const store = new MemoryOAuthStore()
+ * const linear: MCPServer = {
+ *   name: 'linear',
+ *   url: 'https://mcp.linear.app',
+ *   oauth: {
+ *     redirectUri: 'https://myapp.com/mcp/linear/callback',
+ *     scope: 'read',
+ *     store,
+ *   },
+ * }
+ *
+ * try {
+ *   const client = new MCPClient(linear)
+ *   await client.connect()
+ * } catch (err) {
+ *   if (err instanceof MCPAuthRequiredError) {
+ *     // Redirect the user to err.authorizationUrl and remember
+ *     // who they were (so the callback handler can rebuild the
+ *     // store with the right per-user state).
+ *   }
+ * }
+ *
+ * // Later, in the callback handler:
+ * const client = new MCPClient(linear)
+ * await client.completeAuthorization(req.query.code)
+ * // The store now has tokens; subsequent connect()s succeed.
+ * ```
+ *
+ * The framework is server-side and headless — it can't redirect
+ * the user inline. So instead of blocking on `connect()`, we
+ * surface `MCPAuthRequiredError` carrying the authorization URL.
+ * Apps redirect the user themselves, then call
+ * `MCPClient.completeAuthorization(code)` from their callback
+ * route.
+ *
+ * Multi-tenancy: build a fresh `MCPOAuthStore` per `(user, server)`
+ * with the user id baked into the storage keys. The store
+ * interface is intentionally per-server (no `userId` arg) so apps
+ * pick the boundary that matches their data model.
+ */
+
+import type {
+  OAuthClientInformation,
+  OAuthClientInformationFull,
+  OAuthClientMetadata,
+  OAuthTokens,
+} from '@modelcontextprotocol/sdk/shared/auth.js'
+import type { OAuthClientProvider } from '@modelcontextprotocol/sdk/client/auth.js'
+import { BrainError } from '../brain_error.ts'
+
+/**
+ * Persistence contract for one MCP server's OAuth state.
+ *
+ * Methods may be sync or async — implementations are free to call
+ * a DB, a Redis cache, a file, or hold the state in memory. The
+ * framework awaits whichever shape returns.
+ *
+ * Per-(server) only by design — apps that need per-(user, server)
+ * scoping construct a fresh store per request with the user id
+ * baked into the underlying storage keys.
+ */
+export interface MCPOAuthStore {
+  /** Load the dynamic-client-registration record, or undefined if not registered. */
+  clientInformation():
+    | OAuthClientInformation
+    | undefined
+    | Promise<OAuthClientInformation | undefined>
+  /** Persist the dynamic-client-registration record after registration succeeds. */
+  saveClientInformation(info: OAuthClientInformationFull): void | Promise<void>
+  /** Load the active token set, or undefined if the user hasn't authorized. */
+  tokens(): OAuthTokens | undefined | Promise<OAuthTokens | undefined>
+  /** Persist tokens after authorization completes or a refresh succeeds. */
+  saveTokens(tokens: OAuthTokens): void | Promise<void>
+  /** Load the PKCE code verifier saved during the authorize step. */
+  codeVerifier(): string | Promise<string>
+  /** Persist the PKCE code verifier before redirecting to authorize. */
+  saveCodeVerifier(verifier: string): void | Promise<void>
+}
+
+/** Per-server OAuth configuration on `MCPServer.oauth`. */
+export interface MCPOAuthConfig {
+  /** Where the user comes back after authorizing. Must match a registered redirect URI on the OAuth server. */
+  redirectUri: string
+  /** Optional OAuth scopes to request. Some servers require specific scopes. */
+  scope?: string
+  /** Per-server token + client-info storage. */
+  store: MCPOAuthStore
+  /** Optional client metadata for dynamic client registration. Defaults to a minimal sane shape. */
+  clientMetadata?: Partial<OAuthClientMetadata>
+}
+
+/**
+ * `MemoryOAuthStore` — in-memory `MCPOAuthStore` implementation.
+ *
+ * Fine for tests and single-process dev. Production apps with
+ * multiple processes or restarts persist to a DB / Redis / KV.
+ */
+export class MemoryOAuthStore implements MCPOAuthStore {
+  private _clientInfo: OAuthClientInformationFull | undefined
+  private _tokens: OAuthTokens | undefined
+  private _verifier: string | undefined
+
+  clientInformation(): OAuthClientInformation | undefined {
+    return this._clientInfo
+  }
+  saveClientInformation(info: OAuthClientInformationFull): void {
+    this._clientInfo = info
+  }
+  tokens(): OAuthTokens | undefined {
+    return this._tokens
+  }
+  saveTokens(tokens: OAuthTokens): void {
+    this._tokens = tokens
+  }
+  codeVerifier(): string {
+    if (this._verifier === undefined) {
+      throw new BrainError(
+        'MemoryOAuthStore.codeVerifier(): no PKCE verifier saved. The authorization flow must call saveCodeVerifier before requesting the verifier.',
+      )
+    }
+    return this._verifier
+  }
+  saveCodeVerifier(verifier: string): void {
+    this._verifier = verifier
+  }
+}
+
+/**
+ * Thrown when an MCP server requires the user to authorize before
+ * the framework can connect. Apps catch this on `MCPClient.connect()`,
+ * redirect the user to `authorizationUrl`, and on the OAuth callback
+ * route call `MCPClient.completeAuthorization(code)` to finish the
+ * flow.
+ *
+ * `BrainError` subclass so the typed-exception handler renders it
+ * cleanly through the standard pathway.
+ */
+export class MCPAuthRequiredError extends BrainError {
+  /** URL the user should be redirected to in order to authorize. */
+  readonly authorizationUrl: string
+
+  constructor(serverName: string, authorizationUrl: string) {
+    super(
+      `MCPClient(${serverName}): authorization required. Redirect the user to authorizationUrl and call completeAuthorization(code) from your callback route.`,
+      {
+        context: { server: serverName, authorizationUrl },
+      },
+    )
+    this.authorizationUrl = authorizationUrl
+  }
+}
+
+/**
+ * Default client metadata used when dynamic client registration
+ * fires. Apps override per-server via
+ * `MCPOAuthConfig.clientMetadata`.
+ */
+function defaultClientMetadata(redirectUri: string): OAuthClientMetadata {
+  return {
+    redirect_uris: [redirectUri],
+    token_endpoint_auth_method: 'none',
+    grant_types: ['authorization_code', 'refresh_token'],
+    response_types: ['code'],
+    client_name: 'strav-brain',
+  }
+}
+
+/**
+ * Internal — implements the SDK's `OAuthClientProvider` against an
+ * `MCPOAuthStore`. Holds the auth URL captured from the SDK so
+ * `MCPClient.connect()` can surface it on `MCPAuthRequiredError`.
+ */
+export class StoreBackedOAuthProvider implements OAuthClientProvider {
+  /** Captured auth URL — populated by the SDK when authorization is needed. */
+  capturedAuthorizationUrl: URL | undefined
+
+  constructor(
+    private readonly config: MCPOAuthConfig,
+  ) {}
+
+  get redirectUrl(): string {
+    return this.config.redirectUri
+  }
+
+  get clientMetadata(): OAuthClientMetadata {
+    return {
+      ...defaultClientMetadata(this.config.redirectUri),
+      ...(this.config.scope !== undefined ? { scope: this.config.scope } : {}),
+      ...(this.config.clientMetadata ?? {}),
+    }
+  }
+
+  async clientInformation(): Promise<OAuthClientInformation | undefined> {
+    return await this.config.store.clientInformation()
+  }
+  async saveClientInformation(info: OAuthClientInformationFull): Promise<void> {
+    await this.config.store.saveClientInformation(info)
+  }
+  async tokens(): Promise<OAuthTokens | undefined> {
+    return await this.config.store.tokens()
+  }
+  async saveTokens(tokens: OAuthTokens): Promise<void> {
+    await this.config.store.saveTokens(tokens)
+  }
+  async redirectToAuthorization(authorizationUrl: URL): Promise<void> {
+    this.capturedAuthorizationUrl = authorizationUrl
+  }
+  async saveCodeVerifier(verifier: string): Promise<void> {
+    await this.config.store.saveCodeVerifier(verifier)
+  }
+  async codeVerifier(): Promise<string> {
+    return await this.config.store.codeVerifier()
+  }
+}

package/src/mcp/pool.ts

@@ -0,0 +1,106 @@
+/**
+ * `MCPClientPool` — long-lived, per-server `MCPClient` cache.
+ *
+ * Default `resolveMcpTools` flow constructs a fresh `MCPClient` per
+ * call to `runTools` / `runWithTools` / etc., handshakes the
+ * Streamable HTTP transport, lists tools, executes them, then
+ * closes the transport in a `finally`. For one-shot calls that's
+ * fine. For long-running agent workers — chat servers, background
+ * job processors — the per-call handshake adds noticeable
+ * latency and burns connection slots upstream.
+ *
+ * The pool keeps one connected `MCPClient` per `(server.name,
+ * server.url)` pair for the lifetime of the pool. `borrow(server)`
+ * returns the pooled client (lazily creating + connecting on
+ * first use). When the pool is in play, `resolveMcpTools` skips
+ * the per-call `close()` — the pool owns the lifetime — so
+ * subsequent calls reuse the existing transport.
+ *
+ * Apps own the pool's lifetime. Construct one at app boot, hand it
+ * to every provider (or to `BrainProvider` if using the DI
+ * helper), and call `pool.close()` on shutdown.
+ *
+ * ```ts
+ * const pool = new MCPClientPool()
+ *
+ * const openai = new OpenAIProvider(
+ *   'openai',
+ *   { driver: 'openai', apiKey: ... },
+ *   { mcpPool: pool },
+ * )
+ *
+ * // ... many runTools calls later, on graceful shutdown:
+ * await pool.close()
+ * ```
+ *
+ * Concurrency: `borrow()` is synchronous; `MCPClient.connect()`
+ * itself dedupes concurrent calls. Two parallel `runTools` calls
+ * sharing the same pooled client both await one handshake.
+ *
+ * Re-auth: when a borrowed client throws `MCPAuthRequiredError`,
+ * the pool keeps the (still un-authorized) client. Apps call
+ * `pool.evict(server)` after running `completeAuthorization` on
+ * a fresh client so subsequent borrows see the renewed state —
+ * or just reuse the same client the app authorized via the
+ * standard `MCPClient.completeAuthorization` flow.
+ */
+
+import type { MCPServer } from '../mcp_server.ts'
+import { MCPClient } from './client.ts'
+
+/** Internal — factory injection for tests. Defaults to `new MCPClient(server)`. */
+export type MCPClientFactory = (server: MCPServer) => MCPClient
+
+export class MCPClientPool {
+  private readonly clients: Map<string, MCPClient> = new Map()
+  private readonly factory: MCPClientFactory
+
+  constructor(factory: MCPClientFactory = (s) => new MCPClient(s)) {
+    this.factory = factory
+  }
+
+  /**
+   * Return the pooled client for `server`, constructing + caching it on
+   * first call. The client is NOT eagerly connected — the first
+   * `listTools` / `callTool` invocation triggers `connect()` once.
+   */
+  borrow(server: MCPServer): MCPClient {
+    const key = poolKey(server)
+    const existing = this.clients.get(key)
+    if (existing) return existing
+    const client = this.factory(server)
+    this.clients.set(key, client)
+    return client
+  }
+
+  /**
+   * Drop the cached client for `server` and close its transport.
+   * Useful after the app re-authorizes an OAuth server, or after a
+   * transient failure where the connection state is suspect and a
+   * fresh handshake on next borrow is preferable.
+   */
+  async evict(server: MCPServer): Promise<void> {
+    const key = poolKey(server)
+    const client = this.clients.get(key)
+    if (!client) return
+    this.clients.delete(key)
+    await client.close()
+  }
+
+  /** Close every pooled client. Call on app shutdown. */
+  async close(): Promise<void> {
+    const all = [...this.clients.values()]
+    this.clients.clear()
+    await Promise.all(all.map((c) => c.close()))
+  }
+
+  /** Whether the pool currently holds a client for `server`. Used by tests. */
+  has(server: MCPServer): boolean {
+    return this.clients.has(poolKey(server))
+  }
+}
+
+/** Pool key: name + url, so two `MCPServer`s with the same name but different URLs don't collide. */
+function poolKey(server: MCPServer): string {
+  return `${server.name}|${server.url}`
+}

package/src/mcp/resolve_mcp_tools.ts

@@ -21,6 +21,7 @@
 import type { MCPServer } from '../mcp_server.ts'
 import type { Tool, ToolContext } from '../tool.ts'
 import { MCPClient } from './client.ts'
+import type { MCPClientPool } from './pool.ts'
 
 export interface ResolvedMcpTools {
   tools: Tool[]
@@ -30,6 +31,14 @@ export interface ResolvedMcpTools {
 export interface ResolveMcpToolsOptions {
   /** Override the client factory — tests inject mock clients per server here. */
   clientFactory?(server: MCPServer): MCPClient
+  /**
+   * When set, clients are borrowed from the pool instead of being
+   * constructed fresh per call, and the returned `close` becomes a
+   * no-op — the pool owns the lifetime, and apps call
+   * `pool.close()` on shutdown. Mutually beneficial with
+   * `clientFactory` (tests pass a factory to the pool itself).
+   */
+  pool?: MCPClientPool
 }
 
 const NAME_SEPARATOR = '__'
@@ -40,13 +49,16 @@ export async function resolveMcpTools(
 ): Promise<ResolvedMcpTools> {
   const clients: MCPClient[] = []
   const tools: Tool[] = []
+  const pooled = options.pool !== undefined
 
   for (const server of servers) {
     if (server.tools?.enabled === false) continue
-    const client = options.clientFactory
-      ? options.clientFactory(server)
-      : new MCPClient(server)
-    clients.push(client)
+    const client = options.pool
+      ? options.pool.borrow(server)
+      : options.clientFactory
+        ? options.clientFactory(server)
+        : new MCPClient(server)
+    if (!pooled) clients.push(client)
 
     const allowed = server.tools?.allowedTools
     const allowedSet = allowed ? new Set(allowed) : null
@@ -60,9 +72,15 @@ export async function resolveMcpTools(
 
   return {
     tools,
-    close: async () => {
-      await Promise.all(clients.map((c) => c.close()))
-    },
+    // Pooled clients live across calls — `close` becomes a no-op
+    // and the pool owns the lifetime. Non-pooled clients close
+    // here so each `runWithTools` invocation cleans up its own
+    // transports.
+    close: pooled
+      ? async () => {}
+      : async () => {
+          await Promise.all(clients.map((c) => c.close()))
+        },
   }
 }
 
@@ -75,8 +93,12 @@ function buildTool(
     name: `${serverName}${NAME_SEPARATOR}${descriptor.name}`,
     description: descriptor.description,
     inputSchema: descriptor.inputSchema,
-    async execute(input: unknown, _ctx: ToolContext): Promise<string> {
-      const result = await client.callTool(descriptor.name, input)
+    async execute(input: unknown, ctx: ToolContext): Promise<string> {
+      const result = await client.callTool(
+        descriptor.name,
+        input,
+        ctx.signal !== undefined ? { signal: ctx.signal } : {},
+      )
       if (result.isError) {
         return `MCP tool error: ${result.content}`
       }

package/src/mcp_server.ts

@@ -40,8 +40,24 @@ export interface MCPServer {
    * Optional bearer token. Apps source from env vars / secrets
    * managers — never hardcode. The framework forwards this verbatim
    * to the provider's `authorization_token` field.
+   *
+   * Mutually exclusive with `oauth`. Use `authorizationToken` for
+   * self-hosted servers where you control the token; use `oauth`
+   * for commercial servers (Linear, Notion, GitHub, ...).
    */
   authorizationToken?: string
+  /**
+   * OAuth configuration for servers that require it. When set, the
+   * local MCP client (`@strav/brain/mcp`) drives the
+   * authorization-code-with-PKCE flow against the server's OAuth
+   * endpoints, storing tokens via the supplied `store`. The
+   * Anthropic server-side path doesn't use this — Anthropic's
+   * connector handles its own auth.
+   *
+   * Mutually exclusive with `authorizationToken`. See the OAuth
+   * section of `docs/brain/guides/mcp.md`.
+   */
+  oauth?: import('./mcp/oauth.ts').MCPOAuthConfig
   /** Per-server tool config (allowlist / enable flag). */
   tools?: MCPServerToolConfig
 }

package/src/persistence/brain_message.ts

@@ -0,0 +1,34 @@
+/**
+ * `BrainMessage` — the typed row of `brain_message`. One per turn.
+ *
+ * `content` mirrors `Message.content` — string for plain text or
+ * `ContentBlock[]` when the turn carries structured blocks
+ * (tool_use, tool_result, image, compaction, ...). JSONB hydration
+ * is automatic.
+ *
+ * Assistant turns carry `model` / `usage` / `stop_reason` /
+ * `response_id`; user turns leave them NULL. The repository's
+ * `appendTurn` helper writes the right shape per role.
+ */
+
+import { Model } from '@strav/database'
+import type { ChatUsage, ContentBlock } from '../types.ts'
+import { brainMessageSchema } from './schema/brain_message_schema.ts'
+
+export type BrainMessageRole = 'user' | 'assistant'
+
+export class BrainMessage extends Model {
+  static override readonly schema = brainMessageSchema
+
+  id!: string
+  tenant_id!: string
+  thread_id!: string
+  turn_index!: number
+  role!: BrainMessageRole
+  content!: string | ContentBlock[]
+  model!: string | null
+  usage!: ChatUsage | null
+  stop_reason!: string | null
+  response_id!: string | null
+  created_at!: Date
+}

package/src/persistence/brain_message_repository.ts

@@ -0,0 +1,106 @@
+/**
+ * `BrainMessageRepository` — data-access object for `BrainMessage`.
+ *
+ * Append-only by design: messages get inserted and never updated.
+ * `appendTurn` handles the next-`turn_index` lookup + INSERT in a
+ * single round-trip via `INSERT ... SELECT COALESCE(MAX, -1) + 1`
+ * so concurrent appends on the same thread don't race.
+ *
+ * Reads:
+ *   - `loadForThread(threadId, opts?)` — paginated history,
+ *     ordered by `turn_index ASC`.
+ *   - `countForThread(threadId)` — total turn count, useful for
+ *     pagination UIs.
+ */
+
+// biome-ignore lint/style/useImportType: PostgresDatabase needs to be a value import for @inject().
+import { PostgresDatabase, quoteIdent, Repository, type RepositoryScope } from '@strav/database'
+// biome-ignore lint/style/useImportType: EventBus value import for @inject().
+import { EventBus, inject, ulid } from '@strav/kernel'
+import type { ChatUsage, ContentBlock } from '../types.ts'
+import { BrainMessage, type BrainMessageRole } from './brain_message.ts'
+import { brainMessageSchema } from './schema/brain_message_schema.ts'
+
+export interface AppendTurnInput {
+  threadId: string
+  role: BrainMessageRole
+  content: string | ContentBlock[]
+  model?: string
+  usage?: ChatUsage
+  stopReason?: string
+  responseId?: string
+}
+
+export interface LoadMessagesOptions {
+  /** Pagination — defaults to no limit (full history). */
+  limit?: number
+  offset?: number
+}
+
+@inject()
+export class BrainMessageRepository extends Repository<BrainMessage> {
+  static override readonly schema = brainMessageSchema
+  static override readonly model = BrainMessage
+
+  // biome-ignore lint/complexity/noUselessConstructor: explicit constructor for @inject() metadata emission.
+  constructor(db: PostgresDatabase, events: EventBus) {
+    super(db, events)
+  }
+
+  /**
+   * Insert a new turn at the next `turn_index` for the thread. The
+   * `turn_index` is computed in-SQL so two concurrent appends
+   * don't collide — the unique `(thread_id, turn_index)` index on
+   * the table catches any race that slips through.
+   *
+   * Lifecycle: routes through `create()` so `brain_message.created`
+   * events fire. The `turn_index` is filled in by the SELECT side
+   * of an explicit INSERT here rather than `create()` because the
+   * value isn't known client-side.
+   */
+  async appendTurn(input: AppendTurnInput, opts?: RepositoryScope): Promise<BrainMessage> {
+    const table = quoteIdent(brainMessageSchema.name)
+    const sql = `
+      INSERT INTO ${table}
+        ("id", "thread_id", "turn_index", "role", "content",
+         "model", "usage", "stop_reason", "response_id", "created_at")
+      SELECT
+        $1, $2,
+        COALESCE((SELECT MAX("turn_index") FROM ${table} WHERE "thread_id" = $2), -1) + 1,
+        $3, $4::jsonb, $5, $6::jsonb, $7, $8, NOW()
+      RETURNING *
+    `
+    const params = [
+      ulid(),
+      input.threadId,
+      input.role,
+      JSON.stringify(input.content),
+      input.model ?? null,
+      input.usage !== undefined ? JSON.stringify(input.usage) : null,
+      input.stopReason ?? null,
+      input.responseId ?? null,
+    ]
+    const rows = await this.executor(opts).query<Record<string, unknown>>(sql, params)
+    if (rows.length === 0) {
+      throw new Error('BrainMessageRepository.appendTurn: INSERT returned no rows.')
+    }
+    return this.hydrate(rows[0]!)
+  }
+
+  /** Load every turn for a thread, ordered by `turn_index ASC`. */
+  async loadForThread(
+    threadId: string,
+    opts: LoadMessagesOptions = {},
+  ): Promise<BrainMessage[]> {
+    let q = this.query().where('thread_id', threadId).orderBy('turn_index', 'asc')
+    if (opts.limit !== undefined) q = q.limit(opts.limit)
+    if (opts.offset !== undefined) q = q.offset(opts.offset)
+    return q.get()
+  }
+
+  /** Total turn count for a thread — useful for pagination UIs. */
+  async countForThread(threadId: string): Promise<number> {
+    return this.count({ thread_id: threadId } as Partial<BrainMessage>)
+  }
+}
+