@roj-ai/sdk 0.1.23 → 0.1.25

package/src/core/plugin-hooks.integration.test.ts

@@ -4,7 +4,9 @@ import { MockLLMProvider } from '~/core/llm/mock.js'
 import { definePlugin } from '~/core/plugins/plugin-builder.js'
 import { ToolCallId } from '~/core/tools/schema.js'
 import { toolEvents } from '~/core/tools/state.js'
+import { Ok } from '~/lib/utils/result.js'
 import { createTestPreset, TestHarness } from '~/testing/index.js'
+import z from 'zod/v4'
 
 describe('core: plugin hooks', () => {
 	// =========================================================================
@@ -373,6 +375,80 @@ describe('core: plugin hooks', () => {
 
 			await harness.shutdown()
 		})
+
+		describe('beforeMethod', () => {
+			it('returning { action: "deny" } → method call rejected and handler not run', async () => {
+				let handlerCalls = 0
+				const seenMethods: string[] = []
+
+				const gatedPlugin = definePlugin('gated')
+					.method('blocked', {
+						input: z.object({}),
+						output: z.object({ ok: z.boolean() }),
+						handler: async () => {
+							handlerCalls++
+							return Ok({ ok: true })
+						},
+					})
+					.sessionHook('beforeMethod', async (ctx) => {
+						seenMethods.push(ctx.method)
+						if (ctx.method === 'gated.blocked') {
+							return { action: 'deny', reason: 'blocked by test' }
+						}
+						return null
+					})
+					.build()
+
+				const harness = new TestHarness({
+					presets: [createTestPreset()],
+					llmProvider: MockLLMProvider.withFixedResponse({ content: 'Ok', toolCalls: [] }),
+					systemPlugins: [gatedPlugin],
+				})
+
+				const session = await harness.createSession('test')
+				const result = await session.callPluginMethod('gated.blocked', {})
+
+				expect(result.ok).toBe(false)
+				if (!result.ok) {
+					expect(result.error.type).toBe('method_denied')
+					expect(result.error.message).toBe('blocked by test')
+				}
+				expect(handlerCalls).toBe(0)
+				expect(seenMethods).toContain('gated.blocked')
+
+				await harness.shutdown()
+			})
+
+			it('returning null → method call proceeds to the handler', async () => {
+				let handlerCalls = 0
+
+				const gatedPlugin = definePlugin('gated')
+					.method('allowed', {
+						input: z.object({}),
+						output: z.object({ ok: z.boolean() }),
+						handler: async () => {
+							handlerCalls++
+							return Ok({ ok: true })
+						},
+					})
+					.sessionHook('beforeMethod', async () => null)
+					.build()
+
+				const harness = new TestHarness({
+					presets: [createTestPreset()],
+					llmProvider: MockLLMProvider.withFixedResponse({ content: 'Ok', toolCalls: [] }),
+					systemPlugins: [gatedPlugin],
+				})
+
+				const session = await harness.createSession('test')
+				const result = await session.callPluginMethod('gated.allowed', {})
+
+				expect(result.ok).toBe(true)
+				expect(handlerCalls).toBe(1)
+
+				await harness.shutdown()
+			})
+		})
 	})
 
 	// =========================================================================
@@ -425,4 +501,68 @@ describe('core: plugin hooks', () => {
 			await harness.shutdown()
 		})
 	})
+
+	// =========================================================================
+	// beforeDequeue gate
+	// =========================================================================
+
+	describe('beforeDequeue', () => {
+		it('holding a source → messages do not wake the agent, and are re-delivered (not dropped) once released', async () => {
+			let holding = true
+
+			const gatePlugin = definePlugin('budget-gate')
+				.beforeDequeue((ctx) => {
+					if (holding && ctx.sourcePlugin === 'user-chat') {
+						return { action: 'hold' }
+					}
+					return null
+				})
+				.build()
+
+			const harness = new TestHarness({
+				presets: [createTestPreset()],
+				llmProvider: MockLLMProvider.withFixedResponse({ content: 'Ok', toolCalls: [] }),
+				systemPlugins: [gatePlugin],
+			})
+
+			const session = await harness.createSession('test')
+
+			// Held: the user message is queued but must not trigger an inference.
+			await session.sendMessage('first')
+			await new Promise((r) => setTimeout(r, 50))
+			expect(harness.llmProvider.getLastRequest()).toBeUndefined()
+
+			// Released: a new message wakes the agent; the held 'first' must be
+			// re-delivered alongside 'second' (held, not dropped).
+			holding = false
+			await session.sendAndWaitForIdle('second')
+
+			const req = harness.llmProvider.getLastRequest()
+			expect(req).toBeDefined()
+			const serialized = JSON.stringify(req!.messages)
+			expect(serialized).toContain('first')
+			expect(serialized).toContain('second')
+
+			await harness.shutdown()
+		})
+
+		it('returning null → messages flow through normally', async () => {
+			const gatePlugin = definePlugin('budget-gate')
+				.beforeDequeue(() => null)
+				.build()
+
+			const harness = new TestHarness({
+				presets: [createTestPreset()],
+				llmProvider: MockLLMProvider.withFixedResponse({ content: 'Ok', toolCalls: [] }),
+				systemPlugins: [gatePlugin],
+			})
+
+			const session = await harness.createSession('test')
+			await session.sendAndWaitForIdle('hello')
+
+			expect(harness.llmProvider.getLastRequest()).toBeDefined()
+
+			await harness.shutdown()
+		})
+	})
 })

package/src/core/plugins/hook-types.ts

@@ -41,6 +41,26 @@ export type OnErrorResult =
 
 export type OnPauseResult = null
 
+/**
+ * Result of a `beforeMethod` session hook — lets a plugin veto a plugin-method
+ * call before its handler runs (e.g. a budget guard blocking new user messages
+ * while allowing internal agent-to-agent calls).
+ */
+export type BeforeMethodResult =
+	| null
+	| { action: 'deny'; reason: string }
+
+/**
+ * Result of a `beforeDequeue` hook — lets a plugin HOLD (not deliver) the
+ * pending messages of a given source plugin for the current cycle. Held
+ * messages stay queued (not consumed) and are re-delivered on a later cycle —
+ * e.g. a budget guard holding new user/upload input while letting inter-agent
+ * mailbox traffic through so running agents can finish.
+ */
+export type BeforeDequeueResult =
+	| null
+	| { action: 'hold' }
+
 export type HandlerName =
 	| 'onStart'
 	| 'beforeInference'

package/src/core/plugins/plugin-builder.ts

@@ -26,7 +26,9 @@ import type { ToolCall } from '../tools/schema.js'
 import type {
 	AfterInferenceResult,
 	AfterToolCallResult,
+	BeforeDequeueResult,
 	BeforeInferenceResult,
+	BeforeMethodResult,
 	BeforeToolCallResult,
 	OnCompleteResult,
 	OnErrorResult,
@@ -237,6 +239,7 @@ type HookMap<TCtx> = {
 type SessionHookMap<TCtx> = {
 	onSessionReady: (ctx: TCtx) => Promise<void>
 	onSessionClose: (ctx: TCtx) => Promise<void>
+	beforeMethod: (ctx: TCtx & { method: string; input: unknown; agentId?: AgentId }) => Promise<BeforeMethodResult>
 }
 
 // ============================================================================
@@ -311,6 +314,7 @@ type ErasedAgentHookMap = {
 type ErasedSessionHookMap = {
 	onSessionReady: (ctx: BaseSessionHookContext) => Promise<void>
 	onSessionClose: (ctx: BaseSessionHookContext) => Promise<void>
+	beforeMethod: (ctx: BaseSessionHookContext & { method: string; input: unknown; agentId?: AgentId }) => Promise<BeforeMethodResult>
 }
 
 // ============================================================================
@@ -351,6 +355,8 @@ export interface ConfiguredPlugin {
 		getPendingMessages: (ctx: BasePluginHookContext) => PluginPendingMessages | null
 		markConsumed: (ctx: BasePluginHookContext, token: unknown) => Promise<void>
 	}
+	/** Gate run before delivering another plugin's dequeued messages — return 'hold' to keep them queued this cycle. */
+	beforeDequeue?: (ctx: BasePluginHookContext & { sourcePlugin: string }) => BeforeDequeueResult
 	slice?: StateSlice
 	isEnabled?: (ctx: { pluginConfig: unknown; pluginAgentConfig: unknown; agentConfig: AgentConfig }) => boolean
 	isSessionEnabled?: (ctx: { pluginConfig: unknown }) => boolean
@@ -443,6 +449,7 @@ interface BuilderConfig {
 		getPendingMessages: (ctx: BasePluginHookContext) => PluginPendingMessages | null
 		markConsumed: (ctx: BasePluginHookContext, token: unknown) => Promise<void>
 	} | undefined
+	beforeDequeueFn: ((ctx: BasePluginHookContext & { sourcePlugin: string }) => BeforeDequeueResult) | undefined
 }
 
 // ============================================================================
@@ -486,6 +493,7 @@ export class PluginBuilder<
 			sessionHooks: {},
 			isEnabledFn: undefined,
 			dequeueHook: undefined,
+			beforeDequeueFn: undefined,
 		}
 	}
 
@@ -748,6 +756,20 @@ export class PluginBuilder<
 		return this
 	}
 
+	/**
+	 * Gate the delivery of OTHER plugins' dequeued messages. Runs once per source
+	 * plugin that has pending messages, before they wake the agent; return
+	 * `{ action: 'hold' }` to keep that source's messages queued this cycle
+	 * (re-delivered later). `ctx.sourcePlugin` is the name of the plugin whose
+	 * messages are being considered.
+	 */
+	beforeDequeue(
+		fn: (ctx: PluginHookContext<TConfig, TMethods, TAgentConfig, TContext, TState, TNotifications, TDeps> & { sourcePlugin: string }) => BeforeDequeueResult,
+	): this {
+		this._cfg.beforeDequeueFn = fn as BuilderConfig['beforeDequeueFn']
+		return this
+	}
+
 	// --- Build ---
 
 	build(): PluginDefinition<TName, TConfig, TAgentConfig, TManagerMethods, TMethods> {
@@ -811,11 +833,13 @@ function buildConfiguredPlugin(cfg: BuilderConfig, pluginConfig: unknown): Confi
 		) as Partial<ErasedAgentHookMap>
 		: undefined
 
-	const sessionHookEntries = Object.entries(cfg.sessionHooks)
+	const sessionHookEntries = Object.entries(cfg.sessionHooks) as [string, (ctx: BaseSessionHookContext) => Promise<unknown>][]
 	const sessionHooks: Partial<ErasedSessionHookMap> | undefined = sessionHookEntries.length > 0
 		? Object.fromEntries(
 			sessionHookEntries.map(([name, fn]) => [
 				name,
+				// Session.ts provides the full context with extra fields (method, input,
+				// agentId for beforeMethod) — this wrapper just injects pluginConfig.
 				(ctx: BaseSessionHookContext) => fn({ ...ctx, pluginConfig }),
 			]),
 		) as Partial<ErasedSessionHookMap>
@@ -872,6 +896,11 @@ function buildConfiguredPlugin(cfg: BuilderConfig, pluginConfig: unknown): Confi
 		}
 		: undefined
 
+	const beforeDequeueFn = cfg.beforeDequeueFn
+	const wrappedBeforeDequeue = beforeDequeueFn
+		? (ctx: BasePluginHookContext & { sourcePlugin: string }) => beforeDequeueFn({ ...ctx, pluginConfig })
+		: undefined
+
 	// Wrap isEnabled to pass pluginConfig
 	const wrappedIsEnabled = cfg.isEnabledFn
 		? (ctx: { pluginConfig: unknown; pluginAgentConfig: unknown; agentConfig: AgentConfig }) => {
@@ -916,6 +945,7 @@ function buildConfiguredPlugin(cfg: BuilderConfig, pluginConfig: unknown): Confi
 		getStatus: wrappedStatus,
 		getSystemPrompt: wrappedSystemPrompt,
 		dequeue: wrappedDequeue,
+		beforeDequeue: wrappedBeforeDequeue,
 		slice,
 		isEnabled: wrappedIsEnabled,
 		isSessionEnabled: wrappedIsSessionEnabled,

package/src/core/sessions/session.ts

@@ -12,7 +12,7 @@ import { COMMUNICATOR_ROLE, ORCHESTRATOR_ROLE } from '~/core/agents/agent-roles.
 import { AgentId, generateAgentId } from '~/core/agents/schema.js'
 import type { AgentState } from '~/core/agents/state.js'
 import { agentEvents, getChildren } from '~/core/agents/state.js'
-import { AgentErrors, type DomainError, SessionErrors, ValidationErrors } from '~/core/errors.js'
+import { AgentErrors, type DomainError, MethodErrors, SessionErrors, ValidationErrors } from '~/core/errors.js'
 import { withSessionId } from '~/core/events/test-helpers.js'
 import type { DomainEvent } from '~/core/events/types.js'
 import type { LLMLogger } from '~/core/llm/logger.js'
@@ -478,6 +478,29 @@ export class Session {
 			return Err(ValidationErrors.invalid(`Invalid input for ${method}: ${parsed.error.message}`))
 		}
 
+		// beforeMethod gate — any plugin may veto this call before the handler runs
+		// (e.g. a budget guard blocking new user input). The hook receives `caller`,
+		// so a gate can allow internal AGENT_CALLER traffic (agent-to-agent) while
+		// denying external/user-originated calls. First deny wins.
+		for (const gatePlugin of this.plugins) {
+			if (!gatePlugin.sessionHooks?.beforeMethod) continue
+			try {
+				const gateCtx = {
+					...this.buildSessionHookContext(gatePlugin),
+					caller: caller ?? DEFAULT_CALLER,
+					method,
+					input: parsed.data,
+					agentId,
+				}
+				const gate = await gatePlugin.sessionHooks.beforeMethod(gateCtx)
+				if (gate?.action === 'deny') {
+					return Err(MethodErrors.denied(gate.reason))
+				}
+			} catch (err) {
+				this.logger.error(`Plugin '${gatePlugin.name}' beforeMethod hook failed`, err instanceof Error ? err : undefined, { method })
+			}
+		}
+
 		// Build MethodHandlerContext with plugin state, context, scheduleAgent, notify, and deps
 		const sessionContext = this.buildSessionContext()
 		const pluginState = plugin.slice ? plugin.slice.select(this.store.getState()) : undefined