feat: use websocket in agent cli (#143)

apps/agent-test-cli/package.json

@@ -7,5 +7,9 @@
 		"format": "oxfmt --write .",
 		"start": "bun run src/agent-test-cli.ts",
 		"typecheck": "bun tsc --noEmit"
+	},
+	"dependencies": {
+		"@freya/agent-protocol": "workspace:*",
+		"@nym.sh/jrpc": "^0.1.0"
 	}
 }

apps/agent-test-cli/src/agent-test-cli.ts

@@ -1,3 +1,13 @@
+import type {
+	AgentClientApi,
+	AgentEvent,
+	AgentServerApi,
+	SendMessageResult,
+} from "@freya/agent-protocol"
+import type { JrpcChannel, JrpcMessage, JsonRpcMessage } from "@nym.sh/jrpc"
+
+import { JsonRpcClient, JsonRpcServer } from "@nym.sh/jrpc"
+
 type JsonObject = Record<string, unknown>
 
 interface AuthUser {
@@ -15,10 +25,6 @@ interface AuthSession {
 	}
 }
 
-interface QueryResponse {
-	message: string
-}
-
 interface QueryToolDefinition {
 	name: string
 	label: string
@@ -60,6 +66,219 @@ class CookieJar {
 	}
 }
 
+class AgentWebSocketSession implements AgentClientApi {
+	private readonly channel: WebSocketJrpcChannel
+	private readonly client: JsonRpcClient<AgentServerApi>
+	private readonly server: JsonRpcServer<AgentClientApi>
+	private conversationId: string | undefined
+	private responseHadText = false
+
+	private constructor(channel: WebSocketJrpcChannel) {
+		this.channel = channel
+		this.client = new JsonRpcClient<AgentServerApi>(channel)
+		this.server = new JsonRpcServer<AgentClientApi>(
+			{
+				notify: this.notify.bind(this),
+			},
+			channel,
+		)
+	}
+
+	static async connect(backendUrl: string, cookies: CookieJar): Promise<AgentWebSocketSession> {
+		const channel = new WebSocketJrpcChannel(agentWebSocketUrl(backendUrl), cookies.header())
+		const session = new AgentWebSocketSession(channel)
+
+		try {
+			await channel.waitUntilOpen()
+			void session.server.start().catch((err: unknown) => {
+				if (!channel.isClosed()) {
+					console.error(`\nWebSocket JSON-RPC failed: ${formatError(err)}\n`)
+				}
+			})
+			await session.client.call("ping")
+		} catch (err) {
+			channel.close()
+			throw err
+		}
+
+		return session
+	}
+
+	async ask(message: string): Promise<void> {
+		this.responseHadText = false
+
+		const result = await this.sendMessage(message)
+		if (result.conversationId) {
+			this.conversationId = result.conversationId
+		}
+
+		if (!this.responseHadText) {
+			console.log(`\nagent> ${result.message || "(no message)"}`)
+		}
+		console.log("")
+	}
+
+	notify(event: AgentEvent): void {
+		switch (event.type) {
+			case "conversation_started":
+				this.conversationId = event.conversationId
+				break
+			case "message_created":
+				this.printMessage(event.text)
+				break
+			case "tool_started":
+				console.log(`\ntool> ${event.toolName} started`)
+				break
+			case "tool_finished":
+				console.log(`tool> ${event.toolName} ${event.ok ? "finished" : "failed"}`)
+				break
+			case "message_finished":
+				break
+			case "message_failed":
+				console.log(`\nagent! ${event.error}`)
+				break
+		}
+	}
+
+	describeConversation(): string {
+		return this.conversationId ? `Conversation: ${this.conversationId}` : "No conversation yet."
+	}
+
+	close(): void {
+		this.channel.close()
+	}
+
+	private async sendMessage(message: string): Promise<SendMessageResult> {
+		return this.client.call("sendMessage", message)
+	}
+
+	private printMessage(text: string): void {
+		if (text === "") return
+
+		console.log(`\nagent> ${text}`)
+		this.responseHadText = true
+	}
+}
+
+class WebSocketJrpcChannel implements JrpcChannel {
+	private readonly ws: WebSocket
+	private readonly opened: Promise<void>
+	private closed = false
+	private openedOnce = false
+	private queue: JrpcMessage[] = []
+	private waiters: Array<(result: IteratorResult<JrpcMessage, void>) => void> = []
+
+	constructor(url: string, cookieHeader?: string) {
+		this.ws = new WebSocket(url, createWebSocketOptions(cookieHeader))
+		this.opened = new Promise((resolve, reject) => {
+			this.ws.onopen = () => {
+				this.openedOnce = true
+				resolve()
+			}
+			this.ws.onerror = () => {
+				if (!this.openedOnce) {
+					reject(new Error(`Could not connect to ${url}`))
+				}
+			}
+			this.ws.onclose = (event) => {
+				if (!this.openedOnce) {
+					reject(new Error(formatWebSocketClose(url, event)))
+				}
+				this.close()
+			}
+		})
+		this.ws.onmessage = (event) => {
+			this.receive(event.data)
+		}
+	}
+
+	waitUntilOpen(): Promise<void> {
+		return this.opened
+	}
+
+	isClosed(): boolean {
+		return this.closed
+	}
+
+	async send(msg: JsonRpcMessage): Promise<void> {
+		await this.opened
+		if (this.closed || this.ws.readyState !== WebSocket.OPEN) {
+			throw new Error("JSON-RPC WebSocket channel is closed")
+		}
+
+		this.ws.send(JSON.stringify(msg))
+	}
+
+	async next(): Promise<IteratorResult<JrpcMessage, void>> {
+		const msg = this.queue.shift()
+		if (msg) {
+			return { done: false, value: msg }
+		}
+
+		if (this.closed) {
+			return { done: true, value: undefined }
+		}
+
+		return new Promise((resolve) => {
+			this.waiters.push(resolve)
+		})
+	}
+
+	async return(): Promise<IteratorResult<JrpcMessage, void>> {
+		this.close()
+		return { done: true, value: undefined }
+	}
+
+	async throw(error?: unknown): Promise<IteratorResult<JrpcMessage, void>> {
+		this.close()
+		throw error
+	}
+
+	async [Symbol.asyncDispose](): Promise<void> {
+		await this.return()
+	}
+
+	close(): void {
+		if (this.closed) return
+
+		this.closed = true
+		for (const resolve of this.waiters.splice(0)) {
+			resolve({ done: true, value: undefined })
+		}
+
+		if (this.ws.readyState === WebSocket.CONNECTING || this.ws.readyState === WebSocket.OPEN) {
+			this.ws.close()
+		}
+	}
+
+	[Symbol.asyncIterator](): AsyncGenerator<JrpcMessage, void, unknown> {
+		return this
+	}
+
+	private receive(message: unknown): void {
+		const parsed = parseJrpcMessage(message)
+		if (!parsed) {
+			this.ws.close(1003, "Invalid JSON-RPC message")
+			this.close()
+			return
+		}
+
+		this.push(parsed)
+	}
+
+	private push(msg: JrpcMessage): void {
+		if (this.closed) return
+
+		const resolve = this.waiters.shift()
+		if (resolve) {
+			resolve({ done: false, value: msg })
+			return
+		}
+
+		this.queue.push(msg)
+	}
+}
+
 async function main(): Promise<void> {
 	if (wantsHelp()) {
 		printUsage()
@@ -111,73 +330,72 @@ async function runChatLoop(
 	cookies: CookieJar,
 	session: AuthSession,
 ): Promise<void> {
+	const agent = await AgentWebSocketSession.connect(backendUrl, cookies)
+	console.log("Connected to /api/agent/ws")
 	printHelp()
 
-	for (;;) {
-		const input = askOptional("you> ")?.trim()
-		if (!input) continue
+	try {
+		for (;;) {
+			const input = askOptional("you> ")?.trim()
+			if (!input) continue
 
-		if (input === "/quit" || input === "/exit") {
-			console.log("Bye.")
-			return
-		}
+			if (input === "/quit" || input === "/exit") {
+				console.log("Bye.")
+				return
+			}
 
-		if (input === "/help") {
-			printHelp()
-			continue
-		}
+			if (input === "/help") {
+				printHelp()
+				continue
+			}
 
-		if (input === "/session") {
-			console.log(`${session.user.name || session.user.email} (${session.user.id})`)
-			continue
-		}
+			if (input === "/session") {
+				console.log(`${session.user.name || session.user.email} (${session.user.id})`)
+				continue
+			}
 
-		if (input === "/tools") {
-			await runCliCommand(() => listQueryTools(backendUrl, cookies))
-			continue
-		}
+			if (input === "/conversation") {
+				console.log(agent.describeConversation())
+				continue
+			}
 
-		if (input.startsWith("/tool ")) {
-			await runCliCommand(() => executeQueryTool(backendUrl, cookies, input.slice("/tool ".length)))
-			continue
-		}
+			if (input === "/tools") {
+				await runCliCommand(() => listQueryTools(backendUrl, cookies))
+				continue
+			}
 
-		if (input.startsWith("/actions ")) {
-			await runCliCommand(() =>
-				listSourceActions(backendUrl, cookies, input.slice("/actions ".length)),
-			)
-			continue
-		}
+			if (input.startsWith("/tool ")) {
+				await runCliCommand(() =>
+					executeQueryTool(backendUrl, cookies, input.slice("/tool ".length)),
+				)
+				continue
+			}
 
-		if (input.startsWith("/action ")) {
-			await runCliCommand(() =>
-				executeSourceAction(backendUrl, cookies, input.slice("/action ".length)),
-			)
-			continue
-		}
+			if (input.startsWith("/actions ")) {
+				await runCliCommand(() =>
+					listSourceActions(backendUrl, cookies, input.slice("/actions ".length)),
+				)
+				continue
+			}
 
-		try {
-			await askAgent(backendUrl, cookies, input)
-		} catch (err) {
-			console.error(`\n${formatError(err)}\n`)
+			if (input.startsWith("/action ")) {
+				await runCliCommand(() =>
+					executeSourceAction(backendUrl, cookies, input.slice("/action ".length)),
+				)
+				continue
+			}
+
+			try {
+				await agent.ask(input)
+			} catch (err) {
+				console.error(`\n${formatError(err)}\n`)
+			}
 		}
+	} finally {
+		agent.close()
 	}
 }
 
-async function askAgent(backendUrl: string, cookies: CookieJar, message: string): Promise<void> {
-	const data = await requestJson(backendUrl, cookies, "/api/agent", {
-		method: "POST",
-		body: { message },
-	})
-
-	if (!isQueryResponse(data)) {
-		throw new Error("Query returned an unexpected response shape")
-	}
-
-	console.log(`\nagent> ${data.message || "(no message)"}`)
-	console.log("")
-}
-
 async function runCliCommand(command: () => Promise<void>): Promise<void> {
 	try {
 		await command()
@@ -327,7 +545,7 @@ async function requestJson(
 
 function printIntro(): void {
 	console.log("FREYA agent test CLI")
-	console.log("Connect to a backend, sign in, then send test messages to /api/agent.\n")
+	console.log("Connect to a backend, sign in, then send test messages to /api/agent/ws.\n")
 }
 
 function printUsage(): void {
@@ -348,6 +566,7 @@ function printHelp(): void {
 	console.log("  /tool        Execute an agent debug tool with JSON params")
 	console.log("  /actions     List source actions: /actions <source-id>")
 	console.log("  /action      Execute source action: /action <source-id> <action-id> <json-params>")
+	console.log("  /conversation  Show the current websocket conversation")
 	console.log("  /session     Show the signed-in user")
 	console.log("  /help        Show commands")
 	console.log("  /quit        Exit\n")
@@ -417,6 +636,33 @@ function normalizeBackendUrl(value: string): string {
 	}
 }
 
+function agentWebSocketUrl(backendUrl: string): string {
+	const url = new URL(backendUrl)
+	url.protocol = url.protocol === "https:" ? "wss:" : "ws:"
+	url.pathname = "/api/agent/ws"
+	url.search = ""
+	url.hash = ""
+	return url.toString()
+}
+
+function createWebSocketOptions(cookieHeader?: string): Bun.WebSocketOptions | undefined {
+	if (!cookieHeader) return undefined
+
+	return {
+		headers: {
+			Cookie: cookieHeader,
+		},
+	}
+}
+
+function formatWebSocketClose(
+	url: string,
+	event: { code: number; reason: string; wasClean: boolean },
+): string {
+	const reason = event.reason ? `: ${event.reason}` : ""
+	return `Could not connect to ${url} (${event.code}${reason})`
+}
+
 function formatPromptLabel(label: string, defaultValue?: string): string {
 	return defaultValue ? `${label} (${defaultValue}): ` : `${label}: `
 }
@@ -511,6 +757,25 @@ function splitSetCookieHeader(header: string): string[] {
 	return parts.filter(Boolean)
 }
 
+function parseJrpcMessage(message: unknown): JrpcMessage | null {
+	const text = webSocketMessageText(message)
+	if (!text) 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 Uint8Array) return Buffer.from(message).toString("utf8")
+	if (message instanceof ArrayBuffer) return Buffer.from(message).toString("utf8")
+	return null
+}
+
 async function readResponseError(response: Response, path: string): Promise<string> {
 	const text = await response.text()
 	if (response.status === 404 && path === "/api/agent") {
@@ -548,11 +813,6 @@ function isAuthSession(value: unknown): value is AuthSession {
 	)
 }
 
-function isQueryResponse(value: unknown): value is QueryResponse {
-	if (!isJsonObject(value)) return false
-	return typeof value.message === "string"
-}
-
 function isQueryToolsResponse(value: unknown): value is QueryToolsResponse {
 	if (!isJsonObject(value) || !Array.isArray(value.tools)) return false
 	return value.tools.every(isQueryToolDefinition)
@@ -585,6 +845,33 @@ function isSourceActionDefinition(value: unknown): value is { id: string; descri
 	)
 }
 
+function isJrpcMessage(value: unknown): value is JrpcMessage {
+	if (!isJsonObject(value) || value.jsonrpc !== "2.0" || typeof value.id !== "number") {
+		return false
+	}
+
+	if ("method" in value) {
+		return (
+			typeof value.method === "string" &&
+			(value.params === undefined || Array.isArray(value.params))
+		)
+	}
+
+	if ("result" in value) {
+		return true
+	}
+
+	if ("error" in value) {
+		return isJsonRpcErrorObject(value.error)
+	}
+
+	return false
+}
+
+function isJsonRpcErrorObject(value: unknown): boolean {
+	return isJsonObject(value) && typeof value.code === "number" && typeof value.message === "string"
+}
+
 function isJsonObject(value: unknown): value is JsonObject {
 	return typeof value === "object" && value !== null && !Array.isArray(value)
 }

apps/freya-backend/package.json

@@ -16,6 +16,7 @@
 	},
 	"dependencies": {
 		"@earendil-works/pi-coding-agent": "^0.79.1",
+		"@freya/agent-protocol": "workspace:*",
 		"@freya/core": "workspace:*",
 		"@freya/source-caldav": "workspace:*",
 		"@freya/source-google-calendar": "workspace:*",
@@ -25,6 +26,7 @@
 		"@freya/source-tfl": "workspace:*",
 		"@freya/source-weatherkit": "workspace:*",
 		"@freya/source-web-search": "workspace:*",
+		"@nym.sh/jrpc": "^0.1.0",
 		"@openrouter/sdk": "^0.9.11",
 		"arktype": "^2.1.29",
 		"better-auth": "^1",

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

@@ -0,0 +1,149 @@
+import type { AgentEvent } from "@freya/agent-protocol"
+
+import { describe, expect, test } from "bun:test"
+
+import type {
+	QueryAgent,
+	QueryAgentAsk,
+	QueryAgentEvent,
+	QueryAgentEventListener,
+	QueryAgentStreamEvent,
+} from "./query-agent.ts"
+import type { AgentResponseStreamItem } from "./streaming.ts"
+
+import { streamAgentResponse } from "./streaming.ts"
+
+class FakeQueryAgent implements QueryAgent {
+	readonly inputs: QueryAgentAsk[] = []
+	private readonly events: QueryAgentStreamEvent[]
+
+	constructor(events: QueryAgentStreamEvent[]) {
+		this.events = events
+	}
+
+	async *ask(input: QueryAgentAsk): AsyncIterable<QueryAgentStreamEvent> {
+		this.inputs.push(input)
+		for (const event of this.events) {
+			yield event
+		}
+	}
+
+	addEventListener<T extends QueryAgentEvent>(
+		_type: T,
+		_listener: QueryAgentEventListener<T>,
+	): () => void {
+		return () => {}
+	}
+
+	dispose(): void {}
+}
+
+describe("streamAgentResponse", () => {
+	test("emits one message event per completed newline", async () => {
+		const agent = new FakeQueryAgent([
+			{ type: "conversation", conversationId: "conversation-1" },
+			{ type: "text_delta", text: "First message\nSec" },
+			{ type: "text_delta", text: "ond message\nThird message" },
+			{ type: "done" },
+		])
+
+		const { events, result } = 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" },
+			{ type: "message_created", text: "Second message" },
+			{ type: "message_created", text: "Third message" },
+			{ type: "message_finished" },
+		])
+	})
+
+	test("preserves whitespace without emitting empty message events", async () => {
+		const agent = new FakeQueryAgent([
+			{ type: "conversation", conversationId: "conversation-1" },
+			{ type: "text_delta", text: "  const value = 1  \n\n  return value" },
+			{ type: "done" },
+		])
+
+		const { events, result } = 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  " },
+			{ type: "message_created", text: "  return value" },
+			{ type: "message_finished" },
+		])
+	})
+
+	test("emits tool and failure events", async () => {
+		const agent = new FakeQueryAgent([
+			{ type: "conversation", conversationId: "conversation-1" },
+			{ type: "text_delta", text: "I'll check" },
+			{ type: "tool_start", toolName: "calendar" },
+			{ type: "tool_end", toolName: "calendar", ok: false },
+			{ type: "text_delta", text: "That failed" },
+			{ type: "error", message: "model unavailable" },
+		])
+		const stream = streamAgentResponse({
+			agent,
+			input: { message: "hello" },
+		})
+		const events: AgentEvent[] = []
+
+		await expect(collectStreamAgentResponse(stream, events)).rejects.toThrow("model unavailable")
+
+		expect(events).toEqual([
+			{ type: "conversation_started", conversationId: "conversation-1" },
+			{ type: "message_created", text: "I'll check" },
+			{ type: "tool_started", toolName: "calendar" },
+			{ type: "tool_finished", toolName: "calendar", ok: false },
+			{ type: "message_created", text: "That failed" },
+			{ type: "message_failed", error: "model unavailable" },
+		])
+	})
+})
+
+async function collectStreamAgentResponse(
+	stream: AsyncIterable<AgentResponseStreamItem>,
+	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
+		}
+	}
+
+	if (!result) {
+		throw new Error("Expected stream result")
+	}
+
+	return { events, result }
+}

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

@@ -0,0 +1,125 @@
+import type { AgentEvent, SendMessageResult } from "@freya/agent-protocol"
+
+import type { QueryAgent, QueryAgentAsk } from "./query-agent.ts"
+
+export type AgentResponseStreamItem =
+	| { type: "event"; event: AgentEvent }
+	| { type: "result"; result: SendMessageResult }
+
+export async function* streamAgentResponse({
+	agent,
+	input,
+}: {
+	agent: QueryAgent
+	input: QueryAgentAsk
+}): AsyncGenerator<AgentResponseStreamItem, void, void> {
+	let message = ""
+	let conversationId: string | null = null
+	const splitter = new AgentMessageSplitter()
+
+	function messageEvent(text: string): AgentResponseStreamItem | null {
+		if (text.trim() === "") return null
+
+		return { type: "event", event: { type: "message_created", text } }
+	}
+
+	function flushPendingMessage(): AgentResponseStreamItem | null {
+		const text = splitter.flush()
+		if (text === null) return null
+
+		return messageEvent(text)
+	}
+
+	for await (const event of agent.ask(input)) {
+		switch (event.type) {
+			case "conversation":
+				conversationId = event.conversationId
+				yield { type: "event", event: { type: "conversation_started", conversationId } }
+				break
+
+			case "text_delta":
+				message += event.text
+				for (const line of splitter.push(event.text)) {
+					const item = messageEvent(line)
+					if (item) yield item
+				}
+				break
+
+			case "tool_start":
+				{
+					const item = flushPendingMessage()
+					if (item) yield item
+				}
+				yield { type: "event", event: { type: "tool_started", toolName: event.toolName } }
+				break
+
+			case "tool_end":
+				{
+					const item = flushPendingMessage()
+					if (item) yield item
+				}
+				yield {
+					type: "event",
+					event: {
+						type: "tool_finished",
+						toolName: event.toolName,
+						ok: event.ok,
+					},
+				}
+				break
+
+			case "error":
+				{
+					const item = flushPendingMessage()
+					if (item) yield item
+				}
+				yield { type: "event", event: { type: "message_failed", error: event.message } }
+				throw new Error(event.message)
+
+			case "done":
+				{
+					const item = flushPendingMessage()
+					if (item) yield item
+				}
+				const result = createResult(message, conversationId)
+				yield { type: "event", event: { type: "message_finished" } }
+				yield { type: "result", result }
+				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 }
+}
+
+class AgentMessageSplitter {
+	private pending = ""
+
+	push(text: string): string[] {
+		this.pending += text
+
+		const lines = this.pending.split(/\r?\n/)
+		this.pending = lines.pop() ?? ""
+
+		return lines
+	}
+
+	flush(): string | null {
+		if (this.pending === "") return null
+
+		const text = this.pending
+		this.pending = ""
+		return text
+	}
+}

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

@@ -0,0 +1,68 @@
+import { describe, expect, test } from "bun:test"
+import { Hono } from "hono"
+
+import type { UserSessionManager } from "../session/index.ts"
+
+import { registerAgentWebSocketHandlers } from "./ws.ts"
+
+describe("agent websocket handler", () => {
+	test("rejects disallowed browser origins before authenticating", async () => {
+		let sessionChecked = false
+		const app = new Hono()
+
+		registerAgentWebSocketHandlers(app, {
+			sessionManager: {} as UserSessionManager,
+			corsMiddleware: async (c, next) => {
+				const origin = c.req.header("origin")
+				if (origin && origin !== "https://app.freya.test") {
+					return c.text("Forbidden", 403)
+				}
+
+				await next()
+			},
+			authSessionMiddleware: async (c) => {
+				sessionChecked = true
+				return c.json({ error: "Unauthorized" }, 401)
+			},
+		})
+
+		const res = await app.fetch(
+			new Request("https://api.freya.test/api/agent/ws", {
+				headers: {
+					origin: "https://evil.test",
+					upgrade: "websocket",
+				},
+			}),
+		)
+
+		expect(res.status).toBe(403)
+		expect(sessionChecked).toBe(false)
+	})
+
+	test("allows requests without an origin header", async () => {
+		let sessionChecked = false
+		const app = new Hono()
+
+		registerAgentWebSocketHandlers(app, {
+			sessionManager: {} as UserSessionManager,
+			corsMiddleware: async (_c, next) => {
+				await next()
+			},
+			authSessionMiddleware: async (c) => {
+				sessionChecked = true
+				return c.json({ error: "Unauthorized" }, 401)
+			},
+		})
+
+		const res = await app.fetch(
+			new Request("https://api.freya.test/api/agent/ws", {
+				headers: {
+					upgrade: "websocket",
+				},
+			}),
+		)
+
+		expect(res.status).toBe(401)
+		expect(sessionChecked).toBe(true)
+	})
+})

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

@@ -0,0 +1,287 @@
+import type { AgentClientApi, AgentServerApi, SendMessageResult } 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 { 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"
+
+interface AgentWebSocketHandlerDeps {
+	sessionManager: UserSessionManager
+	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,
+): void {
+	app.get(
+		"/api/agent/ws",
+		corsMiddleware,
+		authSessionMiddleware,
+		upgradeWebSocket((c) => {
+			const user = c.get("user")
+			if (!user) {
+				throw new Error("Authenticated WebSocket user missing")
+			}
+
+			const channel = new HonoWebSocketJrpcChannel()
+			const connection = new AgentRpcConnection({
+				channel,
+				sessionManager,
+				userId: user.id,
+			})
+
+			return {
+				onOpen(_event, ws) {
+					channel.attach(ws)
+					void connection.start().catch((err: unknown) => {
+						console.error("[query] Agent WebSocket JSON-RPC failed:", errorMessage(err))
+						ws.close(1011, "Agent RPC connection failed")
+					})
+				},
+
+				onMessage(event) {
+					channel.receive(event.data)
+				},
+
+				onClose() {
+					channel.close()
+				},
+			}
+		}),
+	)
+}
+
+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 userId: string
+
+	constructor({
+		channel,
+		sessionManager,
+		userId,
+	}: {
+		channel: JrpcChannel
+		sessionManager: UserSessionManager
+		userId: string
+	}) {
+		this.sessionManager = sessionManager
+		this.userId = userId
+		this.client = new JsonRpcClient<AgentClientApi>(channel)
+		this.server = new JsonRpcServer<AgentServerApi>(
+			{
+				sendMessage: this.sendMessage.bind(this),
+				ping: this.ping.bind(this),
+			},
+			channel,
+		)
+	}
+
+	start(): Promise<void> {
+		return this.server.start()
+	}
+
+	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
+
+		try {
+			return await run
+		} finally {
+			if (this.activeMessage === run) {
+				this.activeMessage = null
+			}
+		}
+	}
+
+	ping(): "pong" {
+		return "pong"
+	}
+
+	private async runMessage(input: ValidSendMessageInput): Promise<SendMessageResult> {
+		const session = await this.sessionManager.getOrCreate(this.userId)
+		let result: SendMessageResult | null = null
+
+		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
+			}
+		}
+
+		if (!result) {
+			throw new Error("Agent response stream ended without a result")
+		}
+
+		return result
+	}
+}
+
+class HonoWebSocketJrpcChannel implements JrpcChannel {
+	private closed = false
+	private queue: JrpcMessage[] = []
+	private waiters: Array<(result: IteratorResult<JrpcMessage, void>) => void> = []
+	private ws: WSContext | null = null
+
+	attach(ws: WSContext): void {
+		this.ws = ws
+	}
+
+	async send(msg: JsonRpcMessage): Promise<void> {
+		if (this.closed || !this.ws) {
+			throw new Error("JSON-RPC WebSocket channel is closed")
+		}
+
+		this.ws.send(JSON.stringify(msg))
+	}
+
+	receive(message: unknown): void {
+		const parsed = parseJrpcMessage(message)
+		if (!parsed) {
+			this.ws?.close(1003, "Invalid JSON-RPC message")
+			return
+		}
+
+		this.push(parsed)
+	}
+
+	async next(): Promise<IteratorResult<JrpcMessage, void>> {
+		const msg = this.queue.shift()
+		if (msg) {
+			return { done: false, value: msg }
+		}
+
+		if (this.closed) {
+			return { done: true, value: undefined }
+		}
+
+		return new Promise((resolve) => {
+			this.waiters.push(resolve)
+		})
+	}
+
+	async return(): Promise<IteratorResult<JrpcMessage, void>> {
+		this.close()
+		this.ws?.close()
+		return { done: true, value: undefined }
+	}
+
+	async throw(error?: unknown): Promise<IteratorResult<JrpcMessage, void>> {
+		this.close()
+		throw error
+	}
+
+	async [Symbol.asyncDispose](): Promise<void> {
+		await this.return()
+	}
+
+	close(): void {
+		if (this.closed) return
+
+		this.closed = true
+		for (const resolve of this.waiters.splice(0)) {
+			resolve({ done: true, value: undefined })
+		}
+	}
+
+	[Symbol.asyncIterator](): AsyncGenerator<JrpcMessage, void, unknown> {
+		return this
+	}
+
+	private push(msg: JrpcMessage): void {
+		if (this.closed) return
+
+		const resolve = this.waiters.shift()
+		if (resolve) {
+			resolve({ done: false, value: msg })
+			return
+		}
+
+		this.queue.push(msg)
+	}
+}
+
+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)
+}

apps/freya-backend/src/auth/session-middleware.ts

@@ -53,16 +53,6 @@ export function createRequireSession(auth: Auth): AuthSessionMiddleware {
 	}
 }
 
-/**
- * Creates a function to get session from headers. Useful for WebSocket upgrade validation.
- */
-export function createGetSessionFromHeaders(auth: Auth) {
-	return async (headers: Headers): Promise<{ user: AuthUser; session: AuthSession } | null> => {
-		const session = await auth.api.getSession({ headers })
-		return session
-	}
-}
-
 /**
  * Dev/test middleware that injects a fake user and session.
  * Pass userId to simulate an authenticated request, or omit to get 401.

apps/freya-backend/src/server.ts

@@ -1,15 +1,16 @@
 import { Hono } from "hono"
 import { cors } from "hono/cors"
+import { createMiddleware } from "hono/factory"
 
 import { registerAdminHttpHandlers } from "./admin/http.ts"
 import { createQueryDebugTools } from "./agent/debug-tools.ts"
 import { registerAgentHttpHandlers, registerDebugAgentHttpHandlers } from "./agent/http.ts"
+import { agentWebSocket, registerAgentWebSocketHandlers } from "./agent/ws.ts"
 import { createRequireAdmin } from "./auth/admin-middleware.ts"
 import { registerAuthHandlers } from "./auth/http.ts"
 import { createAuth } from "./auth/index.ts"
 import { createRequireSession } from "./auth/session-middleware.ts"
 import { CalDavSourceProvider } from "./caldav/provider.ts"
-import { registerConversationsHttpHandlers } from "./conversations/http.ts"
 import { createDatabase } from "./db/index.ts"
 import { registerFeedHttpHandlers } from "./engine/http.ts"
 import { createFeedEnhancer } from "./enhancement/enhance-feed.ts"
@@ -82,6 +83,15 @@ function main() {
 		return allowedOrigins.includes(origin) ? origin : undefined
 	}
 
+	const agentWebSocketCorsMiddleware = createMiddleware(async (c, next) => {
+		const origin = c.req.header("origin")
+		if (origin && resolveOrigin(origin) === undefined) {
+			return c.text("Forbidden", 403)
+		}
+
+		await next()
+	})
+
 	app.use(
 		"/api/auth/*",
 		cors({
@@ -109,7 +119,6 @@ function main() {
 
 	registerAuthHandlers(app, auth)
 
-	registerConversationsHttpHandlers(app, { db, authSessionMiddleware })
 	registerFeedHttpHandlers(app, {
 		sessionManager,
 		authSessionMiddleware,
@@ -129,6 +138,12 @@ function main() {
 	}
 	registerAdminHttpHandlers(app, { sessionManager, adminMiddleware, db })
 
+	registerAgentWebSocketHandlers(app, {
+		sessionManager,
+		authSessionMiddleware,
+		corsMiddleware: agentWebSocketCorsMiddleware,
+	})
+
 	process.on("SIGTERM", async () => {
 		sessionManager.dispose()
 		await closeDb()
@@ -144,4 +159,5 @@ export default {
 	port: 3000,
 	hostname: "0.0.0.0",
 	fetch: app.fetch,
+	websocket: agentWebSocket,
 }

bun.lock

@@ -49,12 +49,17 @@
     "apps/agent-test-cli": {
       "name": "@freya/agent-test-cli",
       "version": "0.0.0",
+      "dependencies": {
+        "@freya/agent-protocol": "workspace:*",
+        "@nym.sh/jrpc": "^0.1.0",
+      },
     },
     "apps/freya-backend": {
       "name": "@freya/backend",
       "version": "0.0.0",
       "dependencies": {
         "@earendil-works/pi-coding-agent": "^0.79.1",
+        "@freya/agent-protocol": "workspace:*",
         "@freya/core": "workspace:*",
         "@freya/source-caldav": "workspace:*",
         "@freya/source-google-calendar": "workspace:*",
@@ -64,6 +69,7 @@
         "@freya/source-tfl": "workspace:*",
         "@freya/source-weatherkit": "workspace:*",
         "@freya/source-web-search": "workspace:*",
+        "@nym.sh/jrpc": "^0.1.0",
         "@openrouter/sdk": "^0.9.11",
         "arktype": "^2.1.29",
         "better-auth": "^1",
@@ -153,6 +159,10 @@
         "vite-tsconfig-paths": "^5.1.4",
       },
     },
+    "packages/freya-agent-protocol": {
+      "name": "@freya/agent-protocol",
+      "version": "0.0.0",
+    },
     "packages/freya-components": {
       "name": "@freya/components",
       "version": "0.0.0",
@@ -749,6 +759,8 @@
 
     "@formatjs/intl-localematcher": ["@formatjs/intl-localematcher@0.6.2", "", { "dependencies": { "tslib": "^2.8.0" } }, "sha512-XOMO2Hupl0wdd172Y06h6kLpBz6Dv+J4okPLl4LPtzbr8f66WbIoy4ev98EBuZ6ZK4h5ydTN6XneT4QVpD7cdA=="],
 
+    "@freya/agent-protocol": ["@freya/agent-protocol@workspace:packages/freya-agent-protocol"],
+
     "@freya/agent-test-cli": ["@freya/agent-test-cli@workspace:apps/agent-test-cli"],
 
     "@freya/backend": ["@freya/backend@workspace:apps/freya-backend"],
@@ -901,6 +913,8 @@
 
     "@nolyfill/is-core-module": ["@nolyfill/is-core-module@1.0.39", "", {}, "sha512-nn5ozdjYQpUCZlWGuxcJY/KpxkWQs4DcbMCmKojjyrYDEAGy4Ce19NN4v5MduafTwJlbKc99UA8YhSVqq9yPZA=="],
 
+    "@nym.sh/jrpc": ["@nym.sh/jrpc@0.1.0", "", {}, "sha512-qH+vqKojPrX4RkW67U2R4J98uWHxZOwYxX2J5GLZcfm/yjklCcN5zTfDNLfgAa9jAoOFVscC3DFWhvdZOmN3fA=="],
+
     "@nym.sh/jrx": ["@nym.sh/jrx@0.2.0", "", { "peerDependencies": { "@json-render/core": ">=0.10.0" } }, "sha512-jd7Z1Q6T21366MtSUnwCFiu6Yl1AdNc9s5m6HxeUg265P+0enZCiyyxOuHsFwvpUcSEs/2DVBsqfMptdca44lA=="],
 
     "@oclif/core": ["@oclif/core@4.8.4", "", { "dependencies": { "ansi-escapes": "^4.3.2", "ansis": "^3.17.0", "clean-stack": "^3.0.1", "cli-spinners": "^2.9.2", "debug": "^4.4.3", "ejs": "^3.1.10", "get-package-type": "^0.1.0", "indent-string": "^4.0.0", "is-wsl": "^2.2.0", "lilconfig": "^3.1.3", "minimatch": "^10.2.4", "semver": "^7.7.3", "string-width": "^4.2.3", "supports-color": "^8", "tinyglobby": "^0.2.14", "widest-line": "^3.1.0", "wordwrap": "^1.0.0", "wrap-ansi": "^7.0.0" } }, "sha512-UTAqwXJJyRvLBvosL+1uPZYSpr8lEHgUb/EVGbPXo5WZqUIBHfJ0sR2bkBEsrj00/ar4IegKxx4YK0wn2c8SQg=="],

packages/freya-agent-protocol/package.json

@@ -0,0 +1,10 @@
+{
+	"name": "@freya/agent-protocol",
+	"version": "0.0.0",
+	"type": "module",
+	"main": "src/index.ts",
+	"types": "src/index.ts",
+	"scripts": {
+		"test": "bun test ./src"
+	}
+}

packages/freya-agent-protocol/src/index.test.ts

@@ -0,0 +1,20 @@
+import { describe, expect, test } from "bun:test"
+
+import type { AgentEvent, AgentServerApi } from "./index"
+
+describe("agent protocol", () => {
+	test("defines server methods and agent events", () => {
+		const server: AgentServerApi = {
+			async sendMessage(message) {
+				return { message, conversationId: "conversation-1" }
+			},
+			ping() {
+				return "pong"
+			},
+		}
+		const event: AgentEvent = { type: "message_finished" }
+
+		expect(server.ping()).toBe("pong")
+		expect(event.type).toBe("message_finished")
+	})
+})

packages/freya-agent-protocol/src/index.ts

@@ -0,0 +1,21 @@
+export interface SendMessageResult {
+	message: string
+	conversationId: string
+}
+
+export type AgentEvent =
+	| { type: "conversation_started"; conversationId: string }
+	| { type: "message_created"; text: string }
+	| { type: "tool_started"; toolName: string }
+	| { type: "tool_finished"; toolName: string; ok: boolean }
+	| { type: "message_finished" }
+	| { type: "message_failed"; error: string }
+
+export interface AgentServerApi {
+	sendMessage(message: string): Promise<SendMessageResult>
+	ping(): "pong"
+}
+
+export interface AgentClientApi {
+	notify(event: AgentEvent): void
+}

packages/freya-agent-protocol/tsconfig.json

@@ -0,0 +1,4 @@
+{
+	"extends": "../../tsconfig.json",
+	"include": ["src"]
+}