From 83198aab67333773bf5d54a8281789f0424417dc Mon Sep 17 00:00:00 2001
From: Mattia Cerutti <mattiacerutti04@gmail.com>
Date: Sat, 3 Oct 2026 13:20:02 +0100
Subject: [PATCH] feat(durable): expose thinking budgets and websocket connect
 timeout

---
 packages/durable/src/harness/types.ts            |  5 +++++
 packages/durable/test/harness-generation.test.ts | 11 ++++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/packages/durable/src/harness/types.ts b/packages/durable/src/harness/types.ts
index 941a3daa4..691f09ab8 100644
--- a/packages/durable/src/harness/types.ts
+++ b/packages/durable/src/harness/types.ts
@@ -6,6 +6,7 @@ import type {
 	Models,
 	ModelThinkingLevel,
 	Static,
+	ThinkingBudgets,
 	Tool,
 	ToolCall,
 	ToolResultMessage,
@@ -359,6 +360,10 @@ export type ConversationCreateOptions = {
 export type ConversationStreamOptions = {
 	transport?: Transport;
 	timeoutMs?: number;
+	/** WebSocket connection/open handshake timeout; stream idleness uses `timeoutMs`. */
+	websocketConnectTimeoutMs?: number;
+	/** Custom token budgets for thinking levels on token-based providers. */
+	thinkingBudgets?: ThinkingBudgets;
 	/** Provider/SDK retries inside one request attempt. */
 	maxRetries?: number;
 	maxRetryDelayMs?: number;
diff --git a/packages/durable/test/harness-generation.test.ts b/packages/durable/test/harness-generation.test.ts
index e792ebd06..2fb667083 100644
--- a/packages/durable/test/harness-generation.test.ts
+++ b/packages/durable/test/harness-generation.test.ts
@@ -362,7 +362,12 @@ describe("generation", () => {
 				return fauxAssistantMessage("b");
 			},
 		]);
-		setup.settings.stream = { timeoutMs: 1234, headers: { "x-test": "1" } };
+		setup.settings.stream = {
+			timeoutMs: 1234,
+			headers: { "x-test": "1" },
+			websocketConnectTimeoutMs: 5678,
+			thinkingBudgets: { minimal: 128, low: 512, medium: 2048, high: 8192 },
+		};
 		const { harness, root } = await openChat(new MemoryStorage(), setup);
 		await root.configure({ thinkingLevel: "high" }, context);
 		harness.resume();
@@ -377,11 +382,15 @@ describe("generation", () => {
 			headers: { "x-test": "1" },
 			reasoning: "high",
 			sessionId,
+			websocketConnectTimeoutMs: 5678,
+			thinkingBudgets: { minimal: 128, low: 512, medium: 2048, high: 8192 },
 		});
 		expect(seen[0]!.signal).toBeInstanceOf(AbortSignal);
 		expect(seen[1]!.reasoning).toBeUndefined();
 		expect(seen[1]).toMatchObject({ timeoutMs: 99, sessionId });
 		expect(seen[1]!.headers).toBeUndefined();
+		expect(seen[1]!.websocketConnectTimeoutMs).toBeUndefined();
+		expect(seen[1]!.thinkingBudgets).toBeUndefined();
 		await harness.close(context);
 	});