feat: add agent response scheduler

apps/freya-backend/src/agent/job.ts

@@ -0,0 +1,145 @@
+import type { AgentEvent } from "@freya/agent-protocol"
+
+import {
+	AssistantMessagePayload,
+	ConversationEntryKind,
+	UserMessagePayload,
+	ToolCallPayload,
+	ToolResultPayload,
+} from "@freya/core"
+import { type } from "arktype"
+
+import type { ConversationStorage } from "../conversations/storage"
+import type { Job } from "../lib/job"
+import type { JobExecutor } from "../lib/worker"
+import type { NotificationCentral } from "../notification/notification-central"
+import type { UserSessionManager } from "../session"
+
+import { ConversationResponseStateStatus } from "../db/schema"
+import { streamAgentResponse } from "./streaming"
+
+export interface AgentResponseJobPayload {
+	conversationId: string
+}
+
+interface AgentResponseWorkerConfig {
+	conversationStorage: ConversationStorage
+	userSessionManager: UserSessionManager
+	notificationCentral: NotificationCentral
+}
+
+export class AgentResponseJobExecutor implements JobExecutor<AgentResponseJobPayload> {
+	private conversationStorage: ConversationStorage
+	private userSessionManager: UserSessionManager
+	private notificationCentral: NotificationCentral
+
+	constructor({
+		conversationStorage,
+		userSessionManager,
+		notificationCentral,
+	}: AgentResponseWorkerConfig) {
+		this.conversationStorage = conversationStorage
+		this.userSessionManager = userSessionManager
+		this.notificationCentral = notificationCentral
+	}
+
+	async execute(job: Job<AgentResponseJobPayload>): Promise<void> {
+		const conversation = await this.conversationStorage.findConversation(job.payload.conversationId)
+		if (!conversation) {
+			return
+		}
+
+		const claimed = await this.conversationStorage.claimPendingConversationResponseState(
+			job.payload.conversationId,
+		)
+		if (!claimed) {
+			// conversation response state not found or already claimed
+			return
+		}
+
+		const pendingEntries = await this.conversationStorage.listPendingUserConversationEntries(
+			conversation.userId,
+			conversation.id,
+		)
+		if (pendingEntries.length === 0) {
+			await this.conversationStorage.clearConversationResponseState(job.payload.conversationId)
+			return
+		}
+
+		const message = pendingEntries.reduce((acc, entry) => {
+			const payload = UserMessagePayload(entry.payload)
+			if (payload instanceof type.errors) {
+				return acc
+			}
+			return (
+				acc + "\n" + payload.parts.reduce((msg, p) => (p.type === "text" ? msg + p.text : msg), "")
+			)
+		}, "")
+
+		const session = await this.userSessionManager.getOrCreate(conversation.userId)
+
+		try {
+			for await (const event of streamAgentResponse({
+				agent: session.agent,
+				input: { message, signal: job.signal },
+			})) {
+				if (job.signal.aborted) {
+					break
+				}
+
+				await this.recordAgentEvent(event, conversation.id)
+				await this.notificationCentral.notifyUser(conversation.userId, {
+					kind: "agent",
+					payload: event,
+				})
+			}
+
+			// if job is aborted, stop everything immediately, including clean up.
+			// the aborter is assumed responsibility on how to proceed.
+			if (!job.signal.aborted) {
+				await this.conversationStorage.clearConversationResponseState(job.payload.conversationId)
+			}
+		} catch (err) {
+			console.error("[agent job executor] error streaming agent response:", err)
+			if (!job.signal.aborted) {
+				await this.conversationStorage.markResponseStateStatus(
+					[job.payload.conversationId],
+					ConversationResponseStateStatus.Failed,
+				)
+			}
+		}
+	}
+
+	private async recordAgentEvent(event: AgentEvent, conversationId: string) {
+		switch (event.type) {
+			case "message_created":
+				await this.conversationStorage.appendEntry(conversationId, {
+					kind: ConversationEntryKind.AssistantMessage,
+					payload: {
+						role: "assistant",
+						parts: [{ type: "text", text: event.text }],
+					} satisfies AssistantMessagePayload,
+				})
+				break
+
+			case "tool_started":
+				await this.conversationStorage.appendEntry(conversationId, {
+					kind: ConversationEntryKind.ToolCall,
+					payload: {
+						toolName: event.toolName,
+					} satisfies ToolCallPayload,
+				})
+				break
+
+			case "tool_finished":
+				await this.conversationStorage.appendEntry(conversationId, {
+					kind: ConversationEntryKind.ToolResult,
+					payload: {
+						toolName: event.toolName,
+						ok: event.ok,
+					} satisfies ToolResultPayload,
+				})
+				break
+		}
+	}
+}

apps/freya-backend/src/agent/pi-query-agent.ts

@@ -166,6 +166,16 @@ export class PiQueryAgent implements QueryAgent {
 			this.handlePiEvent(event, pushRunEvent)
 		})
 
+		input.signal?.addEventListener(
+			"abort",
+			async () => {
+				await session.abort()
+				close()
+				unsubscribe()
+			},
+			{ once: true },
+		)
+
 		session
 			.prompt(input.message)
 			.then(() => {

apps/freya-backend/src/agent/query-agent.ts

@@ -2,6 +2,7 @@ export interface QueryAgentAsk {
 	message: string
 	conversationId?: string
 	userMessageEntry?: QueryAgentConversationEntryRef
+	signal?: AbortSignal
 }
 
 export type QueryAgentStreamEvent =

apps/freya-backend/src/agent/reconciler.ts

@@ -0,0 +1,70 @@
+import type { ConversationStorage } from "../conversations/storage"
+import type { AgentWorkScheduler } from "./scheduler"
+
+interface AgentResponseReconcilerConfig {
+	storage: ConversationStorage
+	interval: number
+	scheduler: AgentWorkScheduler
+	signal: AbortSignal
+}
+
+export class AgentResponseReconciler {
+	private storage: ConversationStorage
+	private interval: number
+	private scheduler: AgentWorkScheduler
+	private signal: AbortSignal
+
+	private stopLoop: ReturnType<typeof setInterval> | null = null
+
+	constructor({ storage, interval, scheduler, signal }: AgentResponseReconcilerConfig) {
+		this.storage = storage
+		this.interval = interval
+		this.scheduler = scheduler
+		this.signal = signal
+	}
+
+	start() {
+		this.signal.throwIfAborted()
+
+		this.signal.addEventListener(
+			"abort",
+			() => {
+				if (this.stopLoop !== null) {
+					clearInterval(this.stopLoop)
+					this.stopLoop = null
+				}
+			},
+			{ once: true },
+		)
+
+		this.stopLoop = setInterval(this.reconcile.bind(this), this.interval)
+	}
+
+	private async reconcile() {
+		// enqueue pending responses
+		const pendingStates = await this.storage.listPendingResponseStates()
+		const now = new Date().getTime()
+		for (const state of pendingStates) {
+			if (state.maxWaitUntil.getTime() < now) {
+				this.scheduler.enqueueAgentResponse(state.conversationId)
+			}
+		}
+
+		// re-enqueue stuck responses
+		const runningStates = await this.storage.listRunningResponseStates()
+		const stuckIds: string[] = []
+		for (const state of runningStates) {
+			if (state.runningSince && Math.max(now - state.runningSince.getTime(), 0) > 5 * 1000 * 60) {
+				// if the response is running for more than 5 minutes
+				// we assume that its stuck and enqueue it for retry
+				stuckIds.push(state.conversationId)
+			}
+		}
+		if (stuckIds.length > 0) {
+			await this.storage.markResponseStateStatus(stuckIds, "pending")
+			for (const id of stuckIds) {
+				this.scheduler.enqueueAgentResponse(id)
+			}
+		}
+	}
+}

apps/freya-backend/src/agent/scheduler.ts

@@ -0,0 +1,160 @@
+import type { UserEvent } from "@freya/agent-protocol"
+
+import { ConversationEntryKind, UserMessagePayload } from "@freya/core"
+
+import type { ConversationStorage } from "../conversations/storage"
+import type { Job, JobRegistry } from "../lib/job"
+import type { AgentResponseJobPayload } from "./job"
+import { ConversationNotFoundError } from "../conversations/errors";
+import { ConversationResponseStateStatus } from "../db/schema";
+
+interface AgentMessageSchedulerConfig {
+	storage: ConversationStorage
+	maxWaitTime: number
+
+	/**
+	 * How long to wait before responding to the user.
+	 */
+	waitTIme: number
+
+	jobRegistry: JobRegistry<AgentResponseJobPayload>
+}
+
+/**
+ * Schedules and manages the flow of messages between the user and the query agent for a specific conversation.
+ */
+export class AgentWorkScheduler {
+	private conversationStorage: ConversationStorage
+	private jobRegistry: JobRegistry<AgentResponseJobPayload>
+
+	private timing: {
+		maxWaitTime: number
+		waitTime: number
+	}
+
+	private timers = new Map<string, ReturnType<typeof setTimeout>>()
+	private runningJobs = new Map<string, Job<AgentResponseJobPayload>>()
+
+	constructor(config: AgentMessageSchedulerConfig) {
+		this.conversationStorage = config.storage
+		this.jobRegistry = config.jobRegistry
+		this.timing = {
+			maxWaitTime: config.maxWaitTime,
+			waitTime: config.waitTIme,
+		}
+
+		this.jobRegistry.addEventListener("settled", this.eraseJob.bind(this))
+		this.jobRegistry.addEventListener("cancelled", this.eraseJob.bind(this))
+	}
+
+	async receiveMessage(conversationId: string, message: string) {
+		await this.conversationStorage.transaction(async (storage) => {
+			const now = new Date()
+
+			const entry = await storage.appendEntry(conversationId, {
+				kind: ConversationEntryKind.UserMessage,
+				payload: {
+					role: "user",
+					parts: [{ type: "text", text: message }],
+				} satisfies UserMessagePayload,
+			})
+
+			await storage.upsertConversationResponseState(conversationId, {
+				maxWaitUntil: new Date(now.getTime() + this.timing.maxWaitTime),
+				pendingSinceEntryId: entry.id,
+				status: "pending",
+			})
+
+			return entry
+		})
+		this.scheduleAgentResponse(conversationId, this.timing.waitTime)
+	}
+
+	async receiveUserEvent(conversationId: string, event: UserEvent) {
+		if (event.type === "typing") {
+			await this.delayAgentResponse(conversationId)
+		}
+	}
+
+	enqueueAgentResponse(conversationId: string): void {
+		const existing = this.timers.get(conversationId)
+		if (existing) {
+			clearTimeout(existing)
+			this.timers.delete(conversationId)
+		}
+
+		this.cancelCurrentJob(conversationId)
+
+		const job = this.jobRegistry.addJob({
+			payload: { conversationId },
+		})
+		this.runningJobs.set(conversationId, job)
+	}
+
+	private async delayAgentResponse(conversationId: string) {
+		this.cancelCurrentJob(conversationId);
+
+		try {
+			const ok = await this.conversationStorage.transaction(async (storage) => {
+				const state = await storage.findConversationResponseState(conversationId);
+				if (state && state.status !== ConversationResponseStateStatus.Failed) {
+					await storage.updateConversationResponseState(conversationId, {
+						status: ConversationResponseStateStatus.Pending,
+						// the agent response was cancelled, so its no longer running
+						// clear runningSince timestamp
+						runningSince: null,
+					})
+					return true
+				}
+				return false
+			})
+			if (ok) {
+				await this.scheduleAgentResponse(conversationId, this.timing.waitTime)
+			}
+		} catch (error) {
+			if (error instanceof ConversationNotFoundError) {
+				// the user is typing but there isn't a scheduled agent response yet
+				// which means the user is typing their first message after the agent has previously responded
+				// swallow the error
+			} else {
+				console.error("[agent response scheduler] error delaying agent response", error)
+			}
+			return
+		}
+	}
+
+	private async scheduleAgentResponse(conversationId: string, delay: number) {
+		const existing = this.timers.get(conversationId)
+		if (existing) {
+			clearTimeout(existing)
+		}
+
+		this.cancelCurrentJob(conversationId)
+
+		this.timers.set(
+			conversationId,
+			setTimeout(() => {
+				this.enqueueAgentResponse(conversationId)
+			}, delay),
+		)
+	}
+
+	/**
+	 * cancels the current job for agent response for the given conversation id
+	 * no-op if there is no active job for the conversation.
+	 */
+	private cancelCurrentJob(conversationId: string): void {
+		const job = this.runningJobs.get(conversationId)
+		if (!job) return
+
+		// If an active response is working on stale context, abort it so the next
+		// job can answer using the latest pending user messages.
+		this.jobRegistry.cancelJob(job)
+	}
+
+	private eraseJob(job: Job<AgentResponseJobPayload>) {
+		if (this.runningJobs.get(job.payload.conversationId) === job) {
+			this.runningJobs.delete(job.payload.conversationId)
+		}
+	}
+}

apps/freya-backend/src/agent/service.ts

@@ -0,0 +1,66 @@
+import type { UserEvent } from "@freya/agent-protocol"
+
+import type { ConversationStorage } from "../conversations/storage"
+import type { NotificationCentral } from "../notification/notification-central"
+import type { UserSessionManager } from "../session"
+
+import { JobRegistry } from "../lib/job"
+import { Worker } from "../lib/worker"
+import { AgentResponseJobExecutor, type AgentResponseJobPayload } from "./job"
+import { AgentResponseReconciler } from "./reconciler"
+import { AgentWorkScheduler } from "./scheduler"
+
+interface AgentServiceConfig {
+	storage: ConversationStorage
+	userSessionManager: UserSessionManager
+	notificationCentral: NotificationCentral
+	signal: AbortSignal
+}
+
+export class AgentService {
+	private readonly storage: ConversationStorage
+	private readonly scheduler: AgentWorkScheduler
+	private readonly reconciler: AgentResponseReconciler
+	private readonly worker: Worker<AgentResponseJobPayload>
+
+	private readonly jobRegistry = new JobRegistry<AgentResponseJobPayload>()
+
+	constructor({ storage, userSessionManager, notificationCentral, signal }: AgentServiceConfig) {
+		this.storage = storage
+		this.scheduler = new AgentWorkScheduler({
+			storage,
+			jobRegistry: this.jobRegistry,
+			waitTIme: 5 * 1000,
+			maxWaitTime: 5 * 1000 * 60,
+		})
+		this.reconciler = new AgentResponseReconciler({
+			signal,
+			storage: this.storage,
+			interval: 60 * 1000,
+			scheduler: this.scheduler,
+		})
+		this.worker = new Worker<AgentResponseJobPayload>({
+			signal,
+			concurrency: 10,
+			registry: this.jobRegistry,
+			runner: new AgentResponseJobExecutor({
+				conversationStorage: storage,
+				notificationCentral,
+				userSessionManager,
+			}),
+		})
+	}
+
+	start() {
+		this.worker.start()
+		this.reconciler.start()
+	}
+
+	async scheduleAgentResponse(conversationId: string, message: string) {
+		await this.scheduler.receiveMessage(conversationId, message)
+	}
+
+	async handleUserEvent(conversationId: string, event: UserEvent) {
+		await this.scheduler.receiveUserEvent(conversationId, event)
+	}
+}

apps/freya-backend/src/agent/streaming.test.ts

@@ -9,7 +9,6 @@ import type {
 	QueryAgentEventListener,
 	QueryAgentStreamEvent,
 } from "./query-agent.ts"
-import type { AgentResponseStreamItem } from "./streaming.ts"
 
 import { streamAgentResponse } from "./streaming.ts"
 
@@ -47,17 +46,13 @@ describe("streamAgentResponse", () => {
 			{ type: "done" },
 		])
 
-		const { events, result } = await collectStreamAgentResponse(
+		const events = await collectStreamAgentResponse(
 			streamAgentResponse({
 				agent,
 				input: { message: "hello" },
 			}),
 		)
 
-		expect(result).toEqual({
-			conversationId: "conversation-1",
-			message: "First message\nSecond message\nThird message",
-		})
 		expect(events).toEqual([
 			{ type: "conversation_started", conversationId: "conversation-1" },
 			{ type: "message_created", text: "First message" },
@@ -74,17 +69,13 @@ describe("streamAgentResponse", () => {
 			{ type: "done" },
 		])
 
-		const { events, result } = await collectStreamAgentResponse(
+		const events = await collectStreamAgentResponse(
 			streamAgentResponse({
 				agent,
 				input: { message: "hello" },
 			}),
 		)
 
-		expect(result).toEqual({
-			conversationId: "conversation-1",
-			message: "  const value = 1  \n\n  return value",
-		})
 		expect(events).toEqual([
 			{ type: "conversation_started", conversationId: "conversation-1" },
 			{ type: "message_created", text: "  const value = 1  " },
@@ -122,28 +113,12 @@ describe("streamAgentResponse", () => {
 })
 
 async function collectStreamAgentResponse(
-	stream: AsyncIterable<AgentResponseStreamItem>,
+	stream: AsyncIterable<AgentEvent>,
 	events: AgentEvent[] = [],
-): Promise<{
-	events: AgentEvent[]
-	result: { message: string; conversationId: string }
-}> {
-	let result: { message: string; conversationId: string } | null = null
-
-	for await (const item of stream) {
-		switch (item.type) {
-			case "event":
-				events.push(item.event)
-				break
-			case "result":
-				result = item.result
-				break
-		}
+): Promise<AgentEvent[]> {
+	for await (const event of stream) {
+		events.push(event)
 	}
 
-	if (!result) {
-		throw new Error("Expected stream result")
-	}
-
-	return { events, result }
+	return events
 }

apps/freya-backend/src/agent/streaming.ts

@@ -1,10 +1,8 @@
-import type { AgentEvent, SendMessageResult } from "@freya/agent-protocol"
+import type { AgentEvent } from "@freya/agent-protocol"
 
 import type { QueryAgent, QueryAgentAsk } from "./query-agent.ts"
 
-export type AgentResponseStreamItem =
-	| { type: "event"; event: AgentEvent }
-	| { type: "result"; result: SendMessageResult }
+export type AgentResponseStreamItem = { type: "event"; event: AgentEvent }
 
 export async function* streamAgentResponse({
 	agent,
@@ -12,18 +10,18 @@ export async function* streamAgentResponse({
 }: {
 	agent: QueryAgent
 	input: QueryAgentAsk
-}): AsyncGenerator<AgentResponseStreamItem, void, void> {
+}): AsyncGenerator<AgentEvent, void, void> {
 	let message = ""
 	let conversationId: string | null = null
 	const splitter = new AgentMessageSplitter()
 
-	function messageEvent(text: string): AgentResponseStreamItem | null {
+	function messageEvent(text: string): AgentEvent | null {
 		if (text.trim() === "") return null
 
-		return { type: "event", event: { type: "message_created", text } }
+		return { type: "message_created", text }
 	}
 
-	function flushPendingMessage(): AgentResponseStreamItem | null {
+	function flushPendingMessage(): AgentEvent | null {
 		const text = splitter.flush()
 		if (text === null) return null
 
@@ -31,10 +29,14 @@ export async function* streamAgentResponse({
 	}
 
 	for await (const event of agent.ask(input)) {
+		if (input.signal?.aborted) {
+			break
+		}
+
 		switch (event.type) {
 			case "conversation":
 				conversationId = event.conversationId
-				yield { type: "event", event: { type: "conversation_started", conversationId } }
+				yield { type: "conversation_started", conversationId }
 				break
 
 			case "text_delta":
@@ -50,7 +52,7 @@ export async function* streamAgentResponse({
 					const item = flushPendingMessage()
 					if (item) yield item
 				}
-				yield { type: "event", event: { type: "tool_started", toolName: event.toolName } }
+				yield { type: "tool_started", toolName: event.toolName }
 				break
 
 			case "tool_end":
@@ -59,12 +61,9 @@ export async function* streamAgentResponse({
 					if (item) yield item
 				}
 				yield {
-					type: "event",
-					event: {
-						type: "tool_finished",
-						toolName: event.toolName,
-						ok: event.ok,
-					},
+					type: "tool_finished",
+					toolName: event.toolName,
+					ok: event.ok,
 				}
 				break
 
@@ -73,7 +72,7 @@ export async function* streamAgentResponse({
 					const item = flushPendingMessage()
 					if (item) yield item
 				}
-				yield { type: "event", event: { type: "message_failed", error: event.message } }
+				yield { type: "message_failed", error: event.message }
 				throw new Error(event.message)
 
 			case "done":
@@ -81,26 +80,15 @@ export async function* streamAgentResponse({
 					const item = flushPendingMessage()
 					if (item) yield item
 				}
-				const result = createResult(message, conversationId)
-				yield { type: "event", event: { type: "message_finished" } }
-				yield { type: "result", result }
+				yield { type: "message_finished" }
 				return
 		}
 	}
 
 	const item = flushPendingMessage()
 	if (item) yield item
-	const result = createResult(message, conversationId)
-	yield { type: "event", event: { type: "message_finished" } }
-	yield { type: "result", result }
-}
 
-function createResult(message: string, conversationId: string | null): SendMessageResult {
-	if (!conversationId) {
-		throw new Error("Agent response stream ended without a conversation id")
-	}
-
-	return { message, conversationId }
+	yield { type: "message_finished" }
 }
 
 class AgentMessageSplitter {

apps/freya-backend/src/agent/ws.test.ts

@@ -1,8 +1,10 @@
 import { describe, expect, test } from "bun:test"
 import { Hono } from "hono"
 
-import type { UserSessionManager } from "../session/index.ts"
+import type { ConversationStorage } from "../conversations/storage.ts"
+import type { NotificationCentral } from "../notification/notification-central.ts"
 
+import type { AgentService } from "./service.ts"
 import { registerAgentWebSocketHandlers } from "./ws.ts"
 
 describe("agent websocket handler", () => {
@@ -11,7 +13,9 @@ describe("agent websocket handler", () => {
 		const app = new Hono()
 
 		registerAgentWebSocketHandlers(app, {
-			sessionManager: {} as UserSessionManager,
+			agentService: {} as AgentService,
+			storage: {} as ConversationStorage,
+			notificationCentral: {} as NotificationCentral,
 			corsMiddleware: async (c, next) => {
 				const origin = c.req.header("origin")
 				if (origin && origin !== "https://app.freya.test") {
@@ -44,7 +48,9 @@ describe("agent websocket handler", () => {
 		const app = new Hono()
 
 		registerAgentWebSocketHandlers(app, {
-			sessionManager: {} as UserSessionManager,
+			agentService: {} as AgentService,
+			storage: {} as ConversationStorage,
+			notificationCentral: {} as NotificationCentral,
 			corsMiddleware: async (_c, next) => {
 				await next()
 			},

apps/freya-backend/src/agent/ws.ts

@@ -1,53 +1,58 @@
-import type { AgentClientApi, AgentServerApi, SendMessageResult } from "@freya/agent-protocol"
+import type { AgentClientApi, AgentServerApi, UserEvent } from "@freya/agent-protocol"
 import type { JrpcChannel, JrpcMessage, JsonRpcMessage } from "@nym.sh/jrpc"
 import type { Hono, MiddlewareHandler } from "hono"
 import type { WSContext } from "hono/ws"
 
-import { JsonRpcClient, JsonRpcServer } from "@nym.sh/jrpc"
-import { type } from "arktype"
+import { JsonRpcClient, JsonRpcServer, deserializeJrpcMessage } from "@nym.sh/jrpc"
 import { upgradeWebSocket, websocket } from "hono/bun"
 
 import type { AuthSessionMiddleware } from "../auth/session-middleware.ts"
-import type { UserSessionManager } from "../session/index.ts"
-
-import { streamAgentResponse } from "./streaming.ts"
+import type { ConversationStorage } from "../conversations/storage.ts"
+import type {
+	NotificationCentral,
+	NotificationPayload,
+} from "../notification/notification-central.ts"
+import type { AgentService } from "./service.ts"
 
 interface AgentWebSocketHandlerDeps {
-	sessionManager: UserSessionManager
+	agentService: AgentService
+	storage: ConversationStorage
+	notificationCentral: NotificationCentral
 	authSessionMiddleware: AuthSessionMiddleware
 	corsMiddleware: MiddlewareHandler
 }
 
-interface ValidSendMessageInput {
-	message: string
-}
-
 export const agentWebSocket = websocket
 
-const SendMessageInputBody = type({
-	"+": "reject",
-	message: "string",
-})
-
 export function registerAgentWebSocketHandlers(
 	app: Hono,
-	{ sessionManager, authSessionMiddleware, corsMiddleware }: AgentWebSocketHandlerDeps,
+	{
+		agentService,
+		storage,
+		notificationCentral,
+		authSessionMiddleware,
+		corsMiddleware,
+	}: AgentWebSocketHandlerDeps,
 ): void {
 	app.get(
 		"/api/agent/ws",
 		corsMiddleware,
 		authSessionMiddleware,
-		upgradeWebSocket((c) => {
+		upgradeWebSocket(async (c) => {
 			const user = c.get("user")
 			if (!user) {
 				throw new Error("Authenticated WebSocket user missing")
 			}
 
+			const conversation = await storage.getOrCreateConversation(user.id)
+
 			const channel = new HonoWebSocketJrpcChannel()
 			const connection = new AgentRpcConnection({
 				channel,
-				sessionManager,
+				notificationCentral,
+				agentService,
 				userId: user.id,
+				conversationId: conversation.id,
 			})
 
 			return {
@@ -64,6 +69,7 @@ export function registerAgentWebSocketHandlers(
 				},
 
 				onClose() {
+					connection.close()
 					channel.close()
 				},
 			}
@@ -74,54 +80,52 @@ export function registerAgentWebSocketHandlers(
 class AgentRpcConnection implements AgentServerApi {
 	private readonly client: JsonRpcClient<AgentClientApi>
 	private readonly server: JsonRpcServer<AgentServerApi>
-	private activeMessage: Promise<SendMessageResult> | null = null
-	private readonly sessionManager: UserSessionManager
+	private readonly agentService: AgentService
+	private readonly notificationCentral: NotificationCentral
 	private readonly userId: string
+	private readonly conversationId: string
+
+	private cleanup: (() => void) | null = null
 
 	constructor({
+		agentService,
+		notificationCentral,
 		channel,
-		sessionManager,
 		userId,
+		conversationId,
 	}: {
+		agentService: AgentService
+		notificationCentral: NotificationCentral
 		channel: JrpcChannel
-		sessionManager: UserSessionManager
 		userId: string
+		conversationId: string
 	}) {
-		this.sessionManager = sessionManager
-		this.userId = userId
 		this.client = new JsonRpcClient<AgentClientApi>(channel)
+		this.agentService = agentService
+		this.notificationCentral = notificationCentral
+		this.userId = userId
+		this.conversationId = conversationId
 		this.server = new JsonRpcServer<AgentServerApi>(
 			{
 				sendMessage: this.sendMessage.bind(this),
+				notify: this.notify.bind(this),
 				ping: this.ping.bind(this),
 			},
 			channel,
 		)
 	}
 
-	start(): Promise<void> {
-		return this.server.start()
+	notify(event: UserEvent): void {
+		this.agentService.handleUserEvent(this.conversationId, event)
 	}
 
-	async sendMessage(message: string): Promise<SendMessageResult> {
-		const parsed = SendMessageInputBody({ message })
-		if (parsed instanceof type.errors) {
-			throw new Error(parsed.summary)
-		}
-
-		if (this.activeMessage) {
-			throw new Error("A message is already running")
-		}
-
-		const run = this.runMessage(parsed)
-		this.activeMessage = run
-
+	async sendMessage(message: string): Promise<boolean> {
 		try {
-			return await run
-		} finally {
-			if (this.activeMessage === run) {
-				this.activeMessage = null
-			}
+			await this.agentService.scheduleAgentResponse(this.conversationId, message)
+			return true
+		} catch (error) {
+			console.log("[agent rpc connection] error when scheduling agent response", error)
+			return false
 		}
 	}
 
@@ -129,26 +133,22 @@ class AgentRpcConnection implements AgentServerApi {
 		return "pong"
 	}
 
-	private async runMessage(input: ValidSendMessageInput): Promise<SendMessageResult> {
-		const session = await this.sessionManager.getOrCreate(this.userId)
-		let result: SendMessageResult | null = null
+	async start() {
+		this.cleanup = this.notificationCentral.registerListenerForUser(
+			this.userId,
+			this.onNotificationReceived.bind(this),
+		)
+		await this.server.start()
+	}
 
-		for await (const item of streamAgentResponse({ agent: session.agent, input })) {
-			switch (item.type) {
-				case "event":
-					await this.client.call("notify", item.event)
-					break
-				case "result":
-					result = item.result
-					break
-			}
+	close() {
+		this.cleanup?.()
+	}
+
+	private async onNotificationReceived(notification: NotificationPayload) {
+		if (notification.kind === "agent") {
+			await this.client.call("notify", notification.payload)
 		}
-
-		if (!result) {
-			throw new Error("Agent response stream ended without a result")
-		}
-
-		return result
 	}
 }
 
@@ -171,7 +171,11 @@ class HonoWebSocketJrpcChannel implements JrpcChannel {
 	}
 
 	receive(message: unknown): void {
-		const parsed = parseJrpcMessage(message)
+		if (typeof message !== "string") {
+			return
+		}
+
+		const parsed = deserializeJrpcMessage(message)
 		if (!parsed) {
 			this.ws?.close(1003, "Invalid JSON-RPC message")
 			return
@@ -236,52 +240,6 @@ class HonoWebSocketJrpcChannel implements JrpcChannel {
 	}
 }
 
-function parseJrpcMessage(message: unknown): JrpcMessage | null {
-	const text = webSocketMessageText(message)
-	if (text === null) return null
-
-	try {
-		const value: unknown = JSON.parse(text)
-		return isJrpcMessage(value) ? value : null
-	} catch {
-		return null
-	}
-}
-
-function webSocketMessageText(message: unknown): string | null {
-	if (typeof message === "string") return message
-	if (message instanceof ArrayBuffer) return Buffer.from(message).toString("utf8")
-	if (ArrayBuffer.isView(message)) {
-		return Buffer.from(message.buffer, message.byteOffset, message.byteLength).toString("utf8")
-	}
-
-	return null
-}
-
-function isJrpcMessage(value: unknown): value is JrpcMessage {
-	if (typeof value !== "object" || value === null) return false
-	if (!("jsonrpc" in value) || value.jsonrpc !== "2.0") return false
-
-	if ("method" in value) {
-		return "id" in value && typeof value.id === "number" && typeof value.method === "string"
-	}
-
-	if ("result" in value) {
-		return "id" in value && typeof value.id === "number"
-	}
-
-	if ("error" in value) {
-		return (
-			"id" in value &&
-			typeof value.id === "number" &&
-			typeof value.error === "object" &&
-			value.error !== null
-		)
-	}
-
-	return false
-}
-
 function errorMessage(error: unknown): string {
 	return error instanceof Error ? error.message : String(error)
 }