From 92898f486b2c3f67ceeee6707b2b79374be4eff4 Mon Sep 17 00:00:00 2001
From: Mario Zechner <badlogicgames@gmail.com>
Date: Sun, 28 Dec 2025 11:39:46 +0100
Subject: [PATCH] Allow async streamFn for dynamic proxy settings

- StreamFn type now allows returning Promise
- agent-loop awaits streamFn result
- createStreamFn takes getProxyUrl callback, reads settings on each call
---
 packages/agent/src/agent-loop.ts         |  2 +-
 packages/agent/src/types.ts              |  5 ++++-
 packages/web-ui/example/src/main.ts      | 11 +++++------
 packages/web-ui/src/utils/proxy-utils.ts |  9 ++++++---
 4 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts
index 4927a62a..0fd0d6a0 100644
--- a/packages/agent/src/agent-loop.ts
+++ b/packages/agent/src/agent-loop.ts
@@ -211,7 +211,7 @@ async function streamAssistantResponse(
 	const resolvedApiKey =
 		(config.getApiKey ? await config.getApiKey(config.model.provider) : undefined) || config.apiKey;
 
-	const response = streamFunction(config.model, llmContext, {
+	const response = await streamFunction(config.model, llmContext, {
 		...config,
 		apiKey: resolvedApiKey,
 		signal,
diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts
index 1c9c7f2d..e8af618e 100644
--- a/packages/agent/src/types.ts
+++ b/packages/agent/src/types.ts
@@ -11,7 +11,10 @@ import type {
 } from "@mariozechner/pi-ai";
 import type { Static, TSchema } from "@sinclair/typebox";
 
-export type StreamFn = typeof streamSimple;
+/** Stream function - can return sync or Promise for async config lookup */
+export type StreamFn = (
+	...args: Parameters<typeof streamSimple>
+) => ReturnType<typeof streamSimple> | Promise<ReturnType<typeof streamSimple>>;
 
 /**
  * Configuration for the agent loop.
diff --git a/packages/web-ui/example/src/main.ts b/packages/web-ui/example/src/main.ts
index 2a87c9c5..3705f9b0 100644
--- a/packages/web-ui/example/src/main.ts
+++ b/packages/web-ui/example/src/main.ts
@@ -161,10 +161,6 @@ const createAgent = async (initialState?: Partial<AgentState>) => {
 		agentUnsubscribe();
 	}
 
-	// Read proxy settings for streamFn
-	const proxyEnabled = await storage.settings.get<boolean>("proxy.enabled");
-	const proxyUrl = proxyEnabled ? (await storage.settings.get<string>("proxy.url")) || undefined : undefined;
-
 	agent = new Agent({
 		initialState: initialState || {
 			systemPrompt: `You are a helpful AI assistant with access to various tools.
@@ -186,8 +182,11 @@ Feel free to use these tools when needed to provide accurate and helpful respons
 			const key = await storage.providerKeys.get(provider);
 			return key ?? undefined;
 		},
-		// Use streamFn with CORS proxy support
-		streamFn: createStreamFn(proxyUrl),
+		// Use streamFn with CORS proxy support (reads settings on each call)
+		streamFn: createStreamFn(async () => {
+			const enabled = await storage.settings.get<boolean>("proxy.enabled");
+			return enabled ? (await storage.settings.get<string>("proxy.url")) || undefined : undefined;
+		}),
 	});
 
 	agentUnsubscribe = agent.subscribe((event: any) => {
diff --git a/packages/web-ui/src/utils/proxy-utils.ts b/packages/web-ui/src/utils/proxy-utils.ts
index 0792809a..27816b18 100644
--- a/packages/web-ui/src/utils/proxy-utils.ts
+++ b/packages/web-ui/src/utils/proxy-utils.ts
@@ -114,13 +114,16 @@ export function isCorsError(error: unknown): boolean {
 
 /**
  * Create a streamFn that applies CORS proxy when needed.
+ * Reads proxy settings from storage on each call.
  *
- * @param proxyUrl - CORS proxy URL, or undefined to disable
+ * @param getProxyUrl - Async function to get current proxy URL (or undefined if disabled)
  * @returns A streamFn compatible with Agent's streamFn option
  */
-export function createStreamFn(proxyUrl?: string) {
-	return (model: Model<any>, context: Context, options?: SimpleStreamOptions) => {
+export function createStreamFn(getProxyUrl: () => Promise<string | undefined>) {
+	return async (model: Model<any>, context: Context, options?: SimpleStreamOptions) => {
 		const apiKey = options?.apiKey;
+		const proxyUrl = await getProxyUrl();
+
 		if (!apiKey || !proxyUrl) {
 			return streamSimple(model, context, options);
 		}