@roj-ai/sdk 0.0.2

package/src/plugins/session-state/plugin.ts

@@ -0,0 +1,179 @@
+/**
+ * Session State Plugin
+ *
+ * Provides a typed, persistent state object per session that can be updated by:
+ * - Agent (via tool)
+ * - Client/SPA (via RPC method)
+ * - Backend/system (via RPC method with system caller context)
+ *
+ * State changes are broadcast as notifications and can trigger agent dequeue.
+ * A preset-defined validate function controls write permissions per caller.
+ */
+import z from 'zod/v4'
+import type { CallerContext } from '~/core/plugins/plugin-builder.js'
+import { definePlugin } from '~/core/plugins/plugin-builder.js'
+import { createEventsFactory } from '~/core/events/types.js'
+import { getEntryAgentId } from '~/core/sessions/state.js'
+import { Err, Ok } from '~/lib/utils/result.js'
+import { ValidationErrors } from '~/core/errors.js'
+
+export const sessionStateEvents = createEventsFactory({
+	events: {
+		session_state_updated: z.object({
+			state: z.record(z.string(), z.unknown()),
+			callerSource: z.enum(['agent', 'client', 'system']),
+		}),
+		session_state_external_consumed: z.object({}),
+	},
+})
+
+export interface SessionStatePluginConfig {
+	schema: z.ZodType
+	initial: Record<string, unknown>
+	validate?: (current: Record<string, unknown>, proposed: Record<string, unknown>, caller: CallerContext) => true | string
+}
+
+interface SessionStatePluginState {
+	state: Record<string, unknown>
+	pendingExternalUpdate: boolean
+}
+
+export const sessionStatePlugin = definePlugin('sessionState')
+	.pluginConfig<SessionStatePluginConfig>()
+	.isSessionEnabled(({ pluginConfig }) => pluginConfig !== undefined)
+	.events([sessionStateEvents])
+	.state({
+		key: 'sessionState',
+		initial: (): SessionStatePluginState => ({ state: {}, pendingExternalUpdate: false }),
+		reduce: (current, event, _sessionState, pluginConfig) => {
+			switch (event.type) {
+				case 'session_state_updated':
+					return {
+						state: event.state as Record<string, unknown>,
+						pendingExternalUpdate: event.callerSource !== 'agent',
+					}
+				case 'session_state_external_consumed':
+					return { ...current, pendingExternalUpdate: false }
+				default:
+					// Initialize with config defaults if state is empty and config is available
+					if (pluginConfig && Object.keys(current.state).length === 0) {
+						return { ...current, state: (pluginConfig as SessionStatePluginConfig).initial }
+					}
+					return current
+			}
+		},
+	})
+	.notification('sessionStateChanged', {
+		schema: z.object({
+			sessionId: z.string(),
+			state: z.record(z.string(), z.unknown()),
+		}),
+	})
+	.method('get', {
+		input: z.object({ sessionId: z.string() }),
+		output: z.object({ state: z.record(z.string(), z.unknown()) }),
+		handler: async (ctx) => {
+			const currentState = ctx.pluginState.state
+			const config = ctx.pluginConfig
+			const state = Object.keys(currentState).length === 0 && config ? config.initial : currentState
+			return Ok({ state })
+		},
+	})
+	.method('update', {
+		input: z.object({
+			sessionId: z.string(),
+			updates: z.record(z.string(), z.unknown()),
+		}),
+		output: z.object({ state: z.record(z.string(), z.unknown()) }),
+		handler: async (ctx, input) => {
+			const config = ctx.pluginConfig
+			if (!config) return Err(ValidationErrors.invalid('Session state plugin not configured'))
+
+			const currentState = Object.keys(ctx.pluginState.state).length === 0
+				? config.initial
+				: ctx.pluginState.state
+
+			const proposed = { ...currentState, ...input.updates }
+
+			// Validate against schema
+			const parsed = config.schema.safeParse(proposed)
+			if (!parsed.success) {
+				return Err(ValidationErrors.invalid(`Schema validation failed: ${parsed.error.message}`))
+			}
+
+			// Run custom validate function
+			if (config.validate) {
+				const result = config.validate(currentState, proposed, ctx.caller)
+				if (result !== true) {
+					return Err(ValidationErrors.invalid(result))
+				}
+			}
+
+			// Emit event
+			await ctx.emitEvent(sessionStateEvents.create('session_state_updated', {
+				state: proposed,
+				callerSource: ctx.caller.source,
+			}))
+
+			// Notify connected clients
+			ctx.notify('sessionStateChanged', {
+				sessionId: String(ctx.sessionId),
+				state: proposed,
+			})
+
+			// Schedule agent if external update
+			if (ctx.caller.source !== 'agent') {
+				const entryAgentId = getEntryAgentId(ctx.sessionState)
+				if (entryAgentId) {
+					ctx.scheduleAgent(entryAgentId)
+				}
+			}
+
+			return Ok({ state: proposed })
+		},
+	})
+	.tool('update_session_state', {
+		description: 'Update the session state. Use this to signal status changes (e.g., file ready for download, phase transitions).',
+		input: z.object({
+			updates: z.record(z.string(), z.unknown()).describe('Key-value pairs to merge into the current session state'),
+		}),
+		execute: async (input, _toolCtx, ctx) => {
+			const result = await ctx.self.update({
+				sessionId: String(ctx.sessionId),
+				updates: input.updates,
+			})
+
+			if (typeof result === 'object' && result !== null && 'ok' in result && !(result as { ok: boolean }).ok) {
+				const err = result as { ok: false; error: { message: string } }
+				return { ok: false, error: { message: err.error.message, recoverable: true } }
+			}
+
+			return { ok: true, value: [{ type: 'text' as const, text: 'Session state updated.' }] }
+		},
+	})
+	.dequeue({
+		hasPendingMessages: (ctx) => ctx.pluginState.pendingExternalUpdate,
+		getPendingMessages: (ctx) => {
+			if (!ctx.pluginState.pendingExternalUpdate) return null
+			const stateStr = JSON.stringify(ctx.pluginState.state, null, 2)
+			return {
+				messages: [{
+					role: 'user' as const,
+					content: `<system>Session state was updated externally. Current state:\n${stateStr}</system>`,
+				}],
+				token: true,
+			}
+		},
+		markConsumed: async (ctx) => {
+			await ctx.emitEvent(sessionStateEvents.create('session_state_external_consumed', {}))
+		},
+	})
+	.systemPrompt((ctx) => {
+		const config = ctx.pluginConfig
+		if (!config) return null
+		const state = Object.keys(ctx.pluginState.state).length === 0
+			? config.initial
+			: ctx.pluginState.state
+		return `## Session State\n\nCurrent session state:\n\`\`\`json\n${JSON.stringify(state, null, 2)}\n\`\`\`\n\nUse the \`update_session_state\` tool to update state values.`
+	})
+	.build()

package/src/plugins/session-stats/index.ts

@@ -0,0 +1,3 @@
+export { sessionStatsPlugin } from './plugin.js'
+export type { SessionStatsState } from './plugin.js'
+export { selectSessionStats } from './plugin.js'

package/src/plugins/session-stats/plugin.ts

@@ -0,0 +1,120 @@
+import { agentEvents } from '~/core/agents/state.js'
+import { llmEvents } from '~/core/llm/state.js'
+import { definePlugin } from '~/core/plugins/plugin-builder.js'
+import { selectPluginState } from '~/core/sessions/reducer.js'
+import type { SessionState } from '~/core/sessions/state.js'
+import { toolEvents } from '~/core/tools/state.js'
+
+// ============================================================================
+// Session stats state
+// ============================================================================
+
+export interface ProviderStats {
+	llmCalls: number
+	totalTokens: number
+	promptTokens: number
+	completionTokens: number
+	totalCost: number
+}
+
+export interface SessionStatsState {
+	totalTokens: number
+	promptTokens: number
+	completionTokens: number
+	totalCost: number
+	llmCalls: number
+	llmErrors: number
+	toolCalls: number
+	toolErrors: number
+	agentCount: number
+	firstEventAt: number | null
+	lastEventAt: number | null
+	/** Per-provider breakdown of LLM usage */
+	byProvider: Record<string, ProviderStats>
+}
+
+const createInitialStats = (): SessionStatsState => ({
+	totalTokens: 0,
+	promptTokens: 0,
+	completionTokens: 0,
+	totalCost: 0,
+	llmCalls: 0,
+	llmErrors: 0,
+	toolCalls: 0,
+	toolErrors: 0,
+	agentCount: 0,
+	firstEventAt: null,
+	lastEventAt: null,
+	byProvider: {},
+})
+
+const PLUGIN_STATE_KEY = 'sessionStats'
+
+/**
+ * Select session stats from session state (for external consumers).
+ */
+export function selectSessionStats(sessionState: SessionState): SessionStatsState {
+	return selectPluginState<SessionStatsState>(sessionState, PLUGIN_STATE_KEY) ?? createInitialStats()
+}
+
+// ============================================================================
+// Plugin
+// ============================================================================
+
+export const sessionStatsPlugin = definePlugin('session-stats')
+	.events([agentEvents, llmEvents, toolEvents])
+	.state({
+		key: PLUGIN_STATE_KEY,
+		initial: createInitialStats,
+		reduce: (stats, event) => {
+			const now = event.timestamp
+			const withTimestamp = (partial: Partial<SessionStatsState>): SessionStatsState => ({
+				...stats,
+				...partial,
+				firstEventAt: stats.firstEventAt ?? now,
+				lastEventAt: now,
+			})
+
+			switch (event.type) {
+				case 'agent_spawned':
+					return withTimestamp({ agentCount: stats.agentCount + 1 })
+
+				case 'inference_completed': {
+					const provider = event.metrics.provider
+					const byProvider = provider
+						? {
+							...stats.byProvider,
+							[provider]: {
+								llmCalls: (stats.byProvider[provider]?.llmCalls ?? 0) + 1,
+								totalTokens: (stats.byProvider[provider]?.totalTokens ?? 0) + event.metrics.totalTokens,
+								promptTokens: (stats.byProvider[provider]?.promptTokens ?? 0) + event.metrics.promptTokens,
+								completionTokens: (stats.byProvider[provider]?.completionTokens ?? 0) + event.metrics.completionTokens,
+								totalCost: (stats.byProvider[provider]?.totalCost ?? 0) + (event.metrics.cost ?? 0),
+							},
+						}
+						: stats.byProvider
+					return withTimestamp({
+						llmCalls: stats.llmCalls + 1,
+						totalTokens: stats.totalTokens + event.metrics.totalTokens,
+						promptTokens: stats.promptTokens + event.metrics.promptTokens,
+						completionTokens: stats.completionTokens + event.metrics.completionTokens,
+						totalCost: stats.totalCost + (event.metrics.cost ?? 0),
+						byProvider,
+					})
+				}
+
+				case 'inference_failed':
+					return withTimestamp({ llmErrors: stats.llmErrors + 1 })
+
+				case 'tool_started':
+					return withTimestamp({ toolCalls: stats.toolCalls + 1 })
+
+				case 'tool_failed':
+					return withTimestamp({ toolErrors: stats.toolErrors + 1 })
+
+				default:
+					return stats
+			}
+		},
+	})
+	.build()