From 82a23809a40e99f061d93884bebd825a4a95a2d0 Mon Sep 17 00:00:00 2001
From: Christo Todorov <chroxify@gmail.com>
Date: Fri, 15 May 2026 20:58:19 +0200
Subject: [PATCH] feat: allow custom tools to return structured results

---
 packages/runtime/src/client-compat.ts      |  3 ++
 packages/runtime/src/index.ts              |  3 ++
 packages/runtime/src/session.ts            | 34 +++++++++++++++++---
 packages/runtime/src/types.ts              | 13 ++++++--
 packages/runtime/test/custom-tools.test.ts | 37 ++++++++++++++++++++++
 5 files changed, 82 insertions(+), 8 deletions(-)
 create mode 100644 packages/runtime/test/custom-tools.test.ts

diff --git a/packages/runtime/src/client-compat.ts b/packages/runtime/src/client-compat.ts
index 7d72cec5..392fac3a 100644
--- a/packages/runtime/src/client-compat.ts
+++ b/packages/runtime/src/client-compat.ts
@@ -39,5 +39,8 @@ export type {
 	TaskOptions,
 	ThinkingLevel,
 	ToolDef,
+	ToolExecuteResult,
 	ToolParameters,
+	ToolResult,
+	ToolResultContent,
 } from './types.ts';
diff --git a/packages/runtime/src/index.ts b/packages/runtime/src/index.ts
index f766263d..8b79df7a 100644
--- a/packages/runtime/src/index.ts
+++ b/packages/runtime/src/index.ts
@@ -30,7 +30,10 @@ export type {
 	AgentConfig,
 	ModelConfig,
 	ToolDef,
+	ToolExecuteResult,
 	ToolParameters,
+	ToolResult,
+	ToolResultContent,
 	ThinkingLevel,
 	ProviderSettings,
 } from './types.ts';
diff --git a/packages/runtime/src/session.ts b/packages/runtime/src/session.ts
index dd9fa780..841ae95e 100644
--- a/packages/runtime/src/session.ts
+++ b/packages/runtime/src/session.ts
@@ -73,11 +73,38 @@ import type {
 	TaskOptions,
 	ThinkingLevel,
 	ToolDef,
+	ToolExecuteResult,
+	ToolResult,
 } from './types.ts';
 import { addUsage, emptyUsage, fromProviderUsage } from './usage.ts';
 
 const MAX_TASK_DEPTH = 4;
 
+function isToolResult(result: ToolExecuteResult): result is ToolResult {
+	return typeof result === 'object' && result !== null && Array.isArray(result.content);
+}
+
+export function toAgentToolResult(
+	toolName: string,
+	result: ToolExecuteResult,
+): AgentToolResult<any> {
+	if (typeof result === 'string') {
+		return {
+			content: [{ type: 'text' as const, text: result }],
+			details: { customTool: toolName },
+		};
+	}
+
+	if (!isToolResult(result)) {
+		throw new Error(
+			`[flue] Custom tool "${toolName}" returned an invalid result. ` +
+				'Return a string or an object with a content array.',
+		);
+	}
+
+	return result;
+}
+
 export interface CreateTaskSessionOptions {
 	parentSession: string;
 	taskId: string;
@@ -852,11 +879,8 @@ export class Session implements FlueSession {
 				parameters: toolDef.parameters as any,
 				async execute(_toolCallId: string, params: unknown, signal?: AbortSignal) {
 					if (signal?.aborted) throw new Error('Operation aborted');
-					const resultText = await toolDef.execute(params as Record<string, any>, signal);
-					return {
-						content: [{ type: 'text' as const, text: resultText }],
-						details: { customTool: toolDef.name },
-					};
+					const result = await toolDef.execute(params as Record<string, any>, signal);
+					return toAgentToolResult(toolDef.name, result);
 				},
 			}),
 		);
diff --git a/packages/runtime/src/types.ts b/packages/runtime/src/types.ts
index fba632d3..e3dec61d 100644
--- a/packages/runtime/src/types.ts
+++ b/packages/runtime/src/types.ts
@@ -1,4 +1,8 @@
-import type { AgentMessage, ThinkingLevel } from '@earendil-works/pi-agent-core';
+import type {
+	AgentMessage,
+	AgentToolResult,
+	ThinkingLevel,
+} from '@earendil-works/pi-agent-core';
 import type { ImageContent, Model, TSchema } from '@earendil-works/pi-ai';
 import type * as v from 'valibot';
 
@@ -47,6 +51,9 @@ export interface Role {
 // ─── Custom Tools ───────────────────────────────────────────────────────────
 
 export type ToolParameters = TSchema | Record<string, unknown>;
+export type ToolResultContent = AgentToolResult<any>['content'][number];
+export type ToolResult<TDetails = any> = AgentToolResult<TDetails>;
+export type ToolExecuteResult<TDetails = any> = string | ToolResult<TDetails>;
 
 /**
  * Custom tool passed to init(), prompt(), skill(), or task(). init() tools are
@@ -61,8 +68,8 @@ export interface ToolDef<TParams extends ToolParameters = ToolParameters> {
 	description: string;
 	/** JSON Schema-compatible parameter schema. */
 	parameters: TParams;
-	/** Returns a string result sent back to the LLM. Thrown errors become tool errors. */
-	execute: (args: Record<string, any>, signal?: AbortSignal) => Promise<string>;
+	/** Returns a result sent back to the LLM. Thrown errors become tool errors. */
+	execute: (args: Record<string, any>, signal?: AbortSignal) => Promise<ToolExecuteResult>;
 }
 
 // ─── File Stat ──────────────────────────────────────────────────────────────
diff --git a/packages/runtime/test/custom-tools.test.ts b/packages/runtime/test/custom-tools.test.ts
new file mode 100644
index 00000000..498c6985
--- /dev/null
+++ b/packages/runtime/test/custom-tools.test.ts
@@ -0,0 +1,37 @@
+import { describe, expect, it } from 'vitest';
+import { toAgentToolResult } from '../src/session.ts';
+
+describe('custom tool results', () => {
+	it('wraps string results as text content', () => {
+		expect(toAgentToolResult('lookup', 'hello')).toEqual({
+			content: [{ type: 'text', text: 'hello' }],
+			details: { customTool: 'lookup' },
+		});
+	});
+
+	it('preserves structured content and details', () => {
+		const result = toAgentToolResult('screenshot', {
+			content: [
+				{ type: 'text', text: 'screenshot captured' },
+				{ type: 'image', data: 'abc123', mimeType: 'image/png' },
+			],
+			details: { nodeId: 'page:page' },
+		});
+
+		expect(result).toEqual({
+			content: [
+				{ type: 'text', text: 'screenshot captured' },
+				{ type: 'image', data: 'abc123', mimeType: 'image/png' },
+			],
+			details: { nodeId: 'page:page' },
+		});
+	});
+
+	it('rejects invalid structured results', () => {
+		expect(() =>
+			toAgentToolResult('metadata', {
+				details: { nodeId: 'page:page' },
+			} as never),
+		).toThrow('Custom tool "metadata" returned an invalid result');
+	});
+});