diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts
index dec35373..8d5d7706 100644
--- a/packages/agent/src/agent.ts
+++ b/packages/agent/src/agent.ts
@@ -82,8 +82,8 @@ export class Agent {
 	private transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]>;
 	private messageQueue: AgentMessage[] = [];
 	private queueMode: "all" | "one-at-a-time";
-	private streamFn: StreamFn;
-	private getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
+	public streamFn: StreamFn;
+	public getApiKey?: (provider: string) => Promise<string | undefined> | string | undefined;
 	private runningPrompt?: Promise<void>;
 	private resolveRunningPrompt?: () => void;
 
diff --git a/packages/web-ui/example/src/main.ts b/packages/web-ui/example/src/main.ts
index 3705f9b0..f0b5425e 100644
--- a/packages/web-ui/example/src/main.ts
+++ b/packages/web-ui/example/src/main.ts
@@ -8,7 +8,6 @@ import {
 	ChatPanel,
 	CustomProvidersStore,
 	createJavaScriptReplTool,
-	createStreamFn,
 	IndexedDBStorageBackend,
 	// PersistentStorageDialog, // TODO: Fix - currently broken
 	ProviderKeysStore,
@@ -177,16 +176,6 @@ Feel free to use these tools when needed to provide accurate and helpful respons
 		},
 		// Custom transformer: convert custom messages to LLM-compatible format
 		convertToLlm: customConvertToLlm,
-		// Get API keys from provider keys store
-		getApiKey: async (provider: string) => {
-			const key = await storage.providerKeys.get(provider);
-			return key ?? undefined;
-		},
-		// 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/components/AgentInterface.ts b/packages/web-ui/src/components/AgentInterface.ts
index 92b9352b..fd4872a5 100644
--- a/packages/web-ui/src/components/AgentInterface.ts
+++ b/packages/web-ui/src/components/AgentInterface.ts
@@ -1,4 +1,4 @@
-import type { ToolResultMessage, Usage } from "@mariozechner/pi-ai";
+import { streamSimple, type ToolResultMessage, type Usage } from "@mariozechner/pi-ai";
 import { html, LitElement } from "lit";
 import { customElement, property, query } from "lit/decorators.js";
 import { ModelSelector } from "../dialogs/ModelSelector.js";
@@ -12,6 +12,7 @@ import type { Agent, AgentEvent } from "@mariozechner/pi-agent-core";
 import type { Attachment } from "../utils/attachment-utils.js";
 import { formatUsage } from "../utils/format.js";
 import { i18n } from "../utils/i18n.js";
+import { createStreamFn } from "../utils/proxy-utils.js";
 import type { UserMessageWithAttachments } from "./Messages.js";
 import type { StreamingMessageContainer } from "./StreamingMessageContainer.js";
 
@@ -130,6 +131,23 @@ export class AgentInterface extends LitElement {
 			this._unsubscribeSession = undefined;
 		}
 		if (!this.session) return;
+
+		// Set default streamFn with proxy support if not already set
+		if (this.session.streamFn === streamSimple) {
+			this.session.streamFn = createStreamFn(async () => {
+				const enabled = await getAppStorage().settings.get<boolean>("proxy.enabled");
+				return enabled ? (await getAppStorage().settings.get<string>("proxy.url")) || undefined : undefined;
+			});
+		}
+
+		// Set default getApiKey if not already set
+		if (!this.session.getApiKey) {
+			this.session.getApiKey = async (provider: string) => {
+				const key = await getAppStorage().providerKeys.get(provider);
+				return key ?? undefined;
+			};
+		}
+
 		this._unsubscribeSession = this.session.subscribe(async (ev: AgentEvent) => {
 			switch (ev.type) {
 				case "message_start":