@strav/brain 1.0.0-alpha.9 → 1.0.1

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 OpenAIBrainDriver(
+ *   '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

@@ -0,0 +1,108 @@
+/**
+ * `resolveMcpTools` — connects to each `MCPServer`, discovers its
+ * tools, and surfaces them as framework `Tool`s the standard agentic
+ * loop already knows how to invoke.
+ *
+ * Honors the per-server config in `MCPServer.tools`:
+ *   - `enabled: false` → server is skipped entirely.
+ *   - `allowedTools` → only those tool names are exposed.
+ *
+ * Naming: discovered tools are namespaced as `<server>__<tool>` to
+ * keep names unique when multiple servers expose overlapping names.
+ * The `Tool.execute` then routes back to the correct server. The
+ * underscore separator (not `.` or `/`) is chosen because OpenAI's
+ * tool-name regex rejects `.` and `/`.
+ *
+ * Lifecycle: this helper returns `{ tools, close }`. `close` runs all
+ * client `close()` calls in parallel — providers must call it in a
+ * `finally` to avoid leaking transports.
+ */
+
+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[]
+  close(): Promise<void>
+}
+
+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 = '__'
+
+export async function resolveMcpTools(
+  servers: readonly MCPServer[],
+  options: ResolveMcpToolsOptions = {},
+): 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.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
+
+    const descriptors = await client.listTools()
+    for (const descriptor of descriptors) {
+      if (allowedSet && !allowedSet.has(descriptor.name)) continue
+      tools.push(buildTool(server.name, client, descriptor))
+    }
+  }
+
+  return {
+    tools,
+    // 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()))
+        },
+  }
+}
+
+function buildTool(
+  serverName: string,
+  client: MCPClient,
+  descriptor: { name: string; description: string; inputSchema: Record<string, unknown> },
+): Tool {
+  return {
+    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,
+        ctx.signal !== undefined ? { signal: ctx.signal } : {},
+      )
+      if (result.isError) {
+        return `MCP tool error: ${result.content}`
+      }
+      return result.content
+    },
+  }
+}

package/src/mcp_server.ts

@@ -0,0 +1,63 @@
+/**
+ * `MCPServer` — declarative configuration for a remote MCP server
+ * that Anthropic's backend invokes on behalf of the model.
+ *
+ * V1 leverages Anthropic's server-side MCP support: apps declare
+ * server URLs (with optional bearer auth) on the request; the
+ * backend connects to them, discovers their tools, surfaces them to
+ * the model, and runs the tool calls itself. The agentic loop here
+ * doesn't intercept MCP tool calls — they appear in the response as
+ * `MCPToolUseBlock` / `MCPToolResultBlock` content blocks for
+ * observability and audit-trail rendering.
+ *
+ * For OpenAI / Gemini / DeepSeek providers (later slices), a local
+ * MCP client implementation will live alongside this to translate
+ * MCP-discovered tools into `Tool` records and let the framework run
+ * the loop. The V1 contract stays the same; the per-provider
+ * implementation differs.
+ *
+ * `allowedTools` opts into a subset of the server's exposed tools —
+ * useful for narrowing surface area when the MCP server exposes more
+ * capabilities than the agent should be able to invoke. `enabled`
+ * defaults to `true`; set to `false` to declare the server without
+ * routing model calls to it (rare, but handy for temporary
+ * disablement without re-deploying config).
+ */
+
+export interface MCPServerToolConfig {
+  /** Whitelist of tool names the agent can call. Omit for "all tools the server exposes." */
+  allowedTools?: readonly string[]
+  /** Default `true`. Set `false` to declare-but-disable. */
+  enabled?: boolean
+}
+
+export interface MCPServer {
+  /** Server identifier — used in MCPToolUseBlock.serverName + logging. */
+  name: string
+  /** HTTPS URL of the MCP server. */
+  url: string
+  /**
+   * 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/output_schema.ts

@@ -0,0 +1,72 @@
+/**
+ * `OutputSchema<T>` — declarative description of a structured-output
+ * target. Apps pass one of these to `BrainManager.generate(...)` and
+ * get back a `GenerateResult<T>` whose `value` is the parsed JSON
+ * the model produced.
+ *
+ * Schema shape:
+ *   - `name` — short identifier; OpenAI requires it on
+ *     `response_format.json_schema.name`. Other providers ignore it
+ *     but apps should still set something stable and meaningful for
+ *     logs / telemetry.
+ *   - `description` — optional hint the model sees (some providers
+ *     surface it next to the schema; others embed it in the prompt
+ *     translation).
+ *   - `jsonSchema` — plain JSON Schema (draft 2020-12 compatible).
+ *     The framework deliberately doesn't depend on Zod; apps that
+ *     want Zod use `zod-to-json-schema` (or similar) at the call
+ *     site and feed the result here.
+ *   - `parse` — optional runtime validator/parser. When set, the
+ *     framework runs every model response through it before
+ *     returning. Apps that just want type-level inference (and
+ *     trust the model + provider validation) can omit it; the
+ *     return type is then `T` by type assertion only.
+ *
+ * Why no built-in Zod dep:
+ *   Same reason `Tool.inputSchema` is plain JSON Schema — Strav
+ *   doesn't pin apps to a schema library. A thin `@strav/brain-zod`
+ *   helper that produces `OutputSchema<T>` from a Zod schema is
+ *   straightforward to ship later without touching this file.
+ */
+
+export interface OutputSchema<T = unknown> {
+  /** Short identifier — provider-specific; some surface it, others ignore. */
+  name: string
+  /** Optional hint shown to the model alongside the schema. */
+  description?: string
+  /** JSON Schema describing the expected shape. */
+  jsonSchema: Record<string, unknown>
+  /** Optional runtime parser/validator. Apps that use Zod plug it in here. */
+  parse?(value: unknown): T
+}
+
+/**
+ * Shared helper used by every provider's `generate` implementation:
+ * parse the raw JSON response, run the optional `parse` hook, and
+ * wrap parse failures in `BrainError` with the raw text on `.context`
+ * so apps can inspect what came back when validation rejects.
+ */
+import { BrainError } from './brain_error.ts'
+
+export function parseGenerated<T>(text: string, schema: OutputSchema<T>): T {
+  let parsed: unknown
+  try {
+    parsed = JSON.parse(text)
+  } catch (cause) {
+    throw new BrainError(
+      `BrainProvider.generate: response was not valid JSON for schema "${schema.name}".`,
+      { context: { schema: schema.name, text }, cause },
+    )
+  }
+  if (schema.parse) {
+    try {
+      return schema.parse(parsed)
+    } catch (cause) {
+      throw new BrainError(
+        `BrainProvider.generate: response failed schema.parse for "${schema.name}".`,
+        { context: { schema: schema.name, text }, cause },
+      )
+    }
+  }
+  return parsed as T
+}

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 './schemas/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
+}