feat: context-aware functions and generic context typing

CHANGELOG.md

@@ -5,6 +5,23 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+
+- **Context-aware functions** (`addContextFunction`): register functions that receive a shallow-frozen snapshot of the evaluation context as their first parameter, enabling auth/permission/personalization patterns without threading context through expression arguments. Pure and context-aware functions share a single namespace; `isContextFunction(name)` introspects the kind.
+- **Generic context typing** (`bonsai<TCtx>()`): the factory is now generic over context type, with end-to-end type safety through `evaluate`, `evaluateSync`, `compile`, and `addContextFunction`. Backward compatible: defaults to `Record<string, unknown>` when unspecified.
+- New exported types: `ContextFunctionFn`, generic `BonsaiPlugin<TCtx>`, generic `CompiledExpression<TCtx>`, generic `BonsaiInstance<TCtx>`.
+- Internal `Bindings` snapshot consolidates transforms / functions / context-functions into a single cached object passed to the evaluator.
+
+### Resolves
+
+- #33: context access from registered functions.
+
+### Credits
+
+- Thanks to @jaenyf for raising #33 and contributing PR #34.
+
 ## [0.3.0] - 2026-03-21
 
 ### Added

README.md

@@ -125,6 +125,49 @@ expr.addFunction('lookupTier', async (userId) => {
 await expr.evaluate('lookupTier(userId) == "pro"', { userId: 'u_123' })
 ```
 
+### Context-aware functions
+
+Register functions that read the evaluation context directly. The function
+receives a shallow-frozen snapshot of the context as its first parameter, so
+you can keep expressions terse and let the function pull what it needs:
+
+```ts
+import { bonsai } from 'bonsai-js'
+
+interface AppContext {
+  currentUserId: string
+  perms: readonly string[]
+}
+
+const app = bonsai<AppContext>()
+
+app.addContextFunction('lookupCurrentUserTier', async (ctx) => {
+  const row = await db.users.findById(ctx.currentUserId)
+  return row?.tier ?? 'free'
+})
+
+app.addContextFunction('hasPermission', (ctx, action) =>
+  ctx.perms.includes(String(action)))
+
+await app.evaluate(
+  'lookupCurrentUserTier() == "pro" && hasPermission("admin")',
+  { currentUserId: 'u_123', perms: ['admin', 'write'] },
+)
+```
+
+The instance is generic over the context type (`bonsai<AppContext>()`), giving
+you end-to-end type safety: `ctx` is typed inside the function, and the call
+site is type-checked against the same shape. If your context type has required
+fields, TypeScript also requires you to pass the `context` argument to
+`evaluate`, `evaluateSync`, and compiled-expression evaluation. The frozen
+snapshot guarantees functions cannot mutate shared evaluation state.
+
+Pure functions (`addFunction`) and context-aware functions (`addContextFunction`)
+share a single namespace. Registering the same name with either method
+overwrites the prior registration. Use `isContextFunction(name)` for
+introspection when needed. Plugins typed against a minimal context requirement
+can still be applied to instances with a wider context shape.
+
 ### Editor validation
 
 ```ts
@@ -361,30 +404,38 @@ evaluateExpression<number>('x * 2', { x: 21 }) // 42
 ## Instance Methods
 
 ```ts
-interface BonsaiInstance {
-  use(plugin: BonsaiPlugin): this
+type EvaluationContextArgs<TCtx extends Record<string, unknown> = Record<string, unknown>> =
+  {} extends TCtx ? [context?: TCtx] : [context: TCtx]
+
+interface BonsaiInstance<TCtx extends Record<string, unknown> = Record<string, unknown>> {
+  use(plugin: BonsaiPlugin<TCtx>): this
   addTransform(name: string, fn: TransformFn): this
   addFunction(name: string, fn: FunctionFn): this
+  addContextFunction(name: string, fn: ContextFunctionFn<TCtx>): this
   removeTransform(name: string): boolean
   removeFunction(name: string): boolean
   hasTransform(name: string): boolean
   hasFunction(name: string): boolean
+  isContextFunction(name: string): boolean
   listTransforms(): string[]
   listFunctions(): string[]
   clearCache(): void
-  compile(expression: string): CompiledExpression
-  evaluate<T = unknown>(expression: string, context?: Record<string, unknown>): Promise<T>
-  evaluateSync<T = unknown>(expression: string, context?: Record<string, unknown>): T
+  compile(expression: string): CompiledExpression<TCtx>
+  evaluate<T = unknown>(expression: string, ...args: EvaluationContextArgs<TCtx>): Promise<T>
+  evaluateSync<T = unknown>(expression: string, ...args: EvaluationContextArgs<TCtx>): T
   validate(expression: string): ValidationResult
 }
 ```
 
 Method notes:
 
 - `use()` runs a plugin immediately and returns the same instance.
-- `addTransform()` and `addFunction()` overwrite any existing registration with the same name.
-- `listTransforms()` and `listFunctions()` return the currently registered names.
+- `addTransform()`, `addFunction()`, and `addContextFunction()` overwrite any existing registration with the same name. Pure and context-aware functions share a single namespace.
+- `addContextFunction()` registers a function that receives a shallow-frozen snapshot of the evaluation context as its first argument. See [Context-aware functions](#context-aware-functions).
+- `isContextFunction()` returns `true` if the named function was registered via `addContextFunction()`.
+- `listTransforms()` and `listFunctions()` return the currently registered names. `listFunctions()` includes both pure and context-aware functions.
 - `clearCache()` clears the internal AST cache and compiled-expression cache. It does not remove registered transforms/functions.
+- Pass a context type generic to `bonsai<MyContext>()` for end-to-end type safety: `evaluate`, `evaluateSync`, `addContextFunction`, `compile`, and `use` all propagate the type. If `MyContext` has required fields, TypeScript requires a context argument when evaluating.
 
 ## Extending the Runtime
 

benchmarks/context-functions.bench.ts

@@ -0,0 +1,60 @@
+import { describe, bench } from 'vitest'
+import { bonsai } from '../src/index.js'
+
+/**
+ * Validate that:
+ *   1. Pure function calls are unaffected by the context-functions feature
+ *      (we want zero regression on the hot path for users who do not adopt it).
+ *   2. Context function calls have predictable cost (one shallow-frozen copy
+ *      per top-level evaluation, amortised across multiple calls).
+ *   3. Lazy frozen-context caching genuinely avoids per-call freeze cost.
+ */
+
+const pureExpr = bonsai()
+pureExpr.addFunction('id', (x) => x)
+
+const ctxExpr = bonsai<{ value: number }>()
+ctxExpr.addContextFunction('current', (ctx) => ctx.value)
+
+const ctxExprMany = bonsai<{ value: number }>()
+ctxExprMany.addContextFunction('current', (ctx) => ctx.value)
+
+const compiledPure = pureExpr.compile('id(1)')
+const compiledCtx = ctxExpr.compile('current()')
+// Six calls in one expression. With lazy-cached freeze this should pay the
+// freeze cost once, then five "free" calls.
+const compiledCtxMany = ctxExprMany.compile('current() + current() + current() + current() + current() + current()')
+
+describe('context-functions: zero-regression on pure path', () => {
+  bench('pure function call: id(1)', () => {
+    compiledPure.evaluateSync({})
+  })
+})
+
+describe('context-functions: single call', () => {
+  bench('context function call: current()', () => {
+    compiledCtx.evaluateSync({ value: 42 })
+  })
+})
+
+describe('context-functions: amortised freeze cost', () => {
+  bench('six context calls in one expression (amortised freeze)', () => {
+    compiledCtxMany.evaluateSync({ value: 42 })
+  })
+})
+
+describe('context-functions: context realism', () => {
+  const richExpr = bonsai<{ userId: string; perms: readonly string[]; tenantId: string; trace: { reqId: string } }>()
+  richExpr.addContextFunction('hasPermission', (ctx, action) =>
+    ctx.perms.includes(action as string))
+  const compiledRich = richExpr.compile('hasPermission("write")')
+
+  bench('hasPermission("write") with 4-field context', () => {
+    compiledRich.evaluateSync({
+      userId: 'u_1',
+      perms: ['read', 'write'],
+      tenantId: 't_1',
+      trace: { reqId: 'r_1' },
+    })
+  })
+})

benchmarks/overhead.bench.ts

@@ -24,7 +24,7 @@ describe('overhead: literal 42', () => {
     compiled.evaluateSync({})
   })
   bench('raw evaluate (no cache lookup)', () => {
-    evaluate(ast, {}, transforms, functions, new ExecutionContext(policy))
+    evaluate(ast, {}, { transforms, functions, contextFunctions: {} }, new ExecutionContext(policy))
   })
   bench('baseline: just return 42', () => {
     const node = ast as { value: number }
@@ -41,7 +41,7 @@ describe('overhead: property access user.name', () => {
     compiled.evaluateSync(context)
   })
   bench('raw evaluate (no cache lookup)', () => {
-    evaluate(ast, context, {}, {}, new ExecutionContext(policy))
+    evaluate(ast, context, { transforms: {}, functions: {}, contextFunctions: {} }, new ExecutionContext(policy))
   })
 })
 
@@ -54,6 +54,6 @@ describe('overhead: comparison user.age >= 18 && user.verified', () => {
     compiled.evaluateSync(context)
   })
   bench('raw evaluate (no cache lookup)', () => {
-    evaluate(ast, context, {}, {}, new ExecutionContext(policy))
+    evaluate(ast, context, { transforms: {}, functions: {}, contextFunctions: {} }, new ExecutionContext(policy))
   })
 })

scripts/check-package.ts

@@ -4,8 +4,8 @@ import { tmpdir } from 'node:os'
 import { dirname, join, resolve } from 'node:path'
 import { fileURLToPath } from 'node:url'
 
-const __dirname = dirname(fileURLToPath(import.meta.url))
-const root = resolve(__dirname, '..')
+const scriptDir = dirname(fileURLToPath(import.meta.url))
+const root = resolve(scriptDir, '..')
 
 function writeLine(message: string): void {
   process.stdout.write(`${message}\n`)

src/autocomplete/index.ts

@@ -29,7 +29,7 @@ export interface AutocompleteInstance {
 }
 
 // Valid identifier pattern for transform name validation
-const VALID_IDENTIFIER = /^[a-zA-Z_$][\w$]*$/
+const VALID_IDENTIFIER = /^[a-zA-Z_$][\w$]*$/u
 
 /** Check if an error is an expected Bonsai error (syntax, security, type, or reference). */
 function isExpectedError(err: unknown): boolean {
@@ -456,7 +456,7 @@ function inferPipeInputType(
   onError?: ErrorHandler,
 ): InferredTypeName | undefined {
   const before = expression.slice(0, cursor)
-  const pipeMatch = before.match(/^(.*)\|>\s*\w*\s*$/s)
+  const pipeMatch = before.match(/^(.*)\|>\s*\w*\s*$/su)
   if (!pipeMatch) return undefined
 
   const exprBefore = pipeMatch[1].trim()

src/autocomplete/tokenizer.ts

@@ -96,7 +96,7 @@ export function isCursorInsideString(expression: string, cursor: number): boolea
   return inSingle || inDouble || (inTemplate && templateDepth === 0)
 }
 
-const TOKEN_RE = /\?\.|\.\.\.|\|>|\?\?|&&|\|\||[!=<>]=?|[+\-*/%]|\*\*|[.(){}[\],?:]|"(?:[^"\\]|\\.)*"?|'(?:[^'\\]|\\.)*'?|`(?:[^`\\$]|\\.|\$(?!\{))*`?|\d[\d_]*(?:\.\d[\d_]*)?(?:[eE][+-]?\d+)?|[a-zA-Z_$][\w$]*/g
+const TOKEN_RE = /\?\.|\.\.\.|\|>|\?\?|&&|\|\||[!=<>]=?|[+\-*/%]|\*\*|[.(){}[\],?:]|"(?:[^"\\]|\\.)*"?|'(?:[^'\\]|\\.)*'?|`(?:[^`\\$]|\\.|\$(?!\{))*`?|\d[\d_]*(?:\.\d[\d_]*)?(?:[eE][+-]?\d+)?|[a-zA-Z_$][\w$]*/gu
 
 function regexScan(expression: string, cursor: number): NonEofToken[] {
   const tokens: NonEofToken[] = []
@@ -121,8 +121,8 @@ function classifyToken(value: string, start: number, end: number): NonEofToken |
   if (value === 'true' || value === 'false') return { type: 'Boolean', value, start, end }
   if (value === 'null') return { type: 'Null', value, start, end }
   if (value === 'undefined') return { type: 'Undefined', value, start, end }
-  if (/^[a-zA-Z_$]/.test(value)) return { type: 'Identifier', value, start, end }
-  if (/^\d/.test(value)) return { type: 'Number', value, start, end }
+  if (/^[a-zA-Z_$]/u.test(value)) return { type: 'Identifier', value, start, end }
+  if (/^\d/u.test(value)) return { type: 'Number', value, start, end }
   if (value.startsWith('"') || value.startsWith("'")) return { type: 'String', value, start, end }
   if (value.startsWith('`')) return { type: 'TemplateLiteral', value, start, end }
   if (value.length === 1 && '(){}[],:?.'.includes(value)) return { type: 'Punctuation', value: value as PunctuationValue, start, end }

src/eval-ops.ts

@@ -3,6 +3,7 @@ import type { ExecutionContext } from './execution-context.js'
 import type {
   ASTNode,
   BinaryExpressionOperator,
+  ContextFunctionFn,
   TransformFn,
   FunctionFn,
   UnaryOperator,
@@ -121,6 +122,34 @@ export function resolveFunction(name: string, functions: Record<string, Function
   return functions[name]
 }
 
+/**
+ * Resolve a callable by name across the pure-function and context-function
+ * registries. Context functions are checked first because they share a
+ * namespace and registrations are last-write-wins regardless of kind.
+ *
+ * Returns a discriminated result so call sites can decide how to invoke
+ * (pure functions take only call args; context functions receive a frozen
+ * context first, then call args).
+ */
+export type ResolvedCallable =
+  | { kind: 'pure'; fn: FunctionFn }
+  | { kind: 'context'; fn: ContextFunctionFn }
+
+export function resolveCallable(
+  name: string,
+  functions: Record<string, FunctionFn>,
+  contextFunctions: Record<string, ContextFunctionFn>,
+): ResolvedCallable {
+  if (Object.hasOwn(contextFunctions, name)) {
+    return { kind: 'context', fn: contextFunctions[name] }
+  }
+  if (Object.hasOwn(functions, name)) {
+    return { kind: 'pure', fn: functions[name] }
+  }
+  const suggestion = suggest(name, [...Object.keys(functions), ...Object.keys(contextFunctions)])
+  throw new BonsaiReferenceError('function', name, suggestion)
+}
+
 export function getIdentifierName(node: ASTNode, message = 'Expected identifier'): string {
   if (node.type !== 'Identifier') {
     throw new Error(message)

src/evaluator-async.ts

@@ -1,15 +1,16 @@
-import type { ASTNode, FunctionFn, ObjectProperty, TransformFn } from './types.js'
+import type { ASTNode, ObjectProperty } from './types.js'
+import type { Bindings } from './plugins.js'
 import type { ExecutionContext } from './execution-context.js'
 import { attachLocation } from './errors.js'
-import { type EvalEnv } from './evaluator.js'
+import { type EvalEnv, getFrozenContext } from './evaluator.js'
 import {
   accessMember,
   applyBinaryOp,
   applyUnaryOp,
   expandSpreadValue,
   getIdentifierName,
   getObjectLiteralKeyName,
-  resolveFunction,
+  resolveCallable,
   resolveTransform,
   validateMethodArgs,
   validateMethodCall,
@@ -20,12 +21,18 @@ type AsyncEvalEnv = EvalEnv
 export async function evaluateAsync(
   node: ASTNode,
   context: Record<string, unknown>,
-  transforms: Record<string, TransformFn>,
-  functions: Record<string, FunctionFn>,
+  bindings: Bindings,
   guard: ExecutionContext,
   source?: string,
 ): Promise<unknown> {
-  const env: AsyncEvalEnv = { ctx: context, tr: transforms, fn: functions, g: guard, s: source }
+  const env: AsyncEvalEnv = {
+    ctx: context,
+    tr: bindings.transforms,
+    fn: bindings.functions,
+    cfn: bindings.contextFunctions,
+    g: guard,
+    s: source,
+  }
   const result = await evalNodeAsync(node, env)
   guard.checkTimeout()
   return result
@@ -214,12 +221,14 @@ async function evalCallExpressionAsync(
 
   if (node.callee.type === 'Identifier') {
     try {
-      const func = resolveFunction(node.callee.name, fn)
+      const resolved = resolveCallable(node.callee.name, fn, env.cfn)
       const args: unknown[] = []
       for (const arg of node.args) {
         await pushCallArgumentAsync(args, arg, env)
       }
-      const result = await func(...args)
+      const result = resolved.kind === 'context'
+        ? await resolved.fn(getFrozenContext(env), ...args)
+        : await resolved.fn(...args)
       g.checkTimeout()
       return result
     } catch (e) {