feat: Implement Tool Coordinator for efficient multi-step tool execution

app/actions.tsx

@@ -12,7 +12,7 @@ import type { FeatureCollection } from 'geojson'
 import { Spinner } from '@/components/ui/spinner'
 import { Section } from '@/components/section'
 import { FollowupPanel } from '@/components/followup-panel'
-import { inquire, researcher, taskManager, querySuggestor, resolutionSearch } from '@/lib/agents'
+import { inquire, researcher, taskManager, querySuggestor, resolutionSearch, toolCoordinator, executeToolPlan, aggregateToolResults } from '@/lib/agents'
 // Removed import of useGeospatialToolMcp as it no longer exists and was incorrectly used here.
 // The geospatialTool (if used by agents like researcher) now manages its own MCP client.
 import { writer } from '@/lib/agents/writer'
@@ -125,6 +125,7 @@ async function submit(formData?: FormData, skip?: boolean) {
 
   const groupeId = nanoid()
   const useSpecificAPI = process.env.USE_SPECIFIC_API_FOR_WRITER === 'true'
+  const useToolCoordinator = process.env.USE_TOOL_COORDINATOR === 'true'
   const maxMessages = useSpecificAPI ? 5 : 10
   messages.splice(0, Math.max(messages.length - maxMessages, 0))
 
@@ -319,17 +320,55 @@ async function submit(formData?: FormData, skip?: boolean) {
     const streamText = createStreamableValue<string>()
     uiStream.update(<Spinner />)
 
+    let finalMessages = messages
+
+    if (useToolCoordinator) {
+      uiStream.update(<div><Spinner /> Planning tool execution...</div>)
+      try {
+        const plan = await toolCoordinator(messages)
+        uiStream.update(<div><Spinner /> Executing tool plan...</div>)
+        const results = await executeToolPlan(plan)
+        toolOutputs = results
+        const summary = aggregateToolResults(results, plan)
+
+        // Add the summary to the messages for the final synthesis agent
+        finalMessages = [
+          ...messages,
+          {
+            id: nanoid(),
+            role: 'tool',
+            content: summary,
+            type: 'tool_coordinator_summary'
+          } as any // Cast to any to satisfy CoreMessage type for custom type
+        ]
+
+        // Stream a message to the user about the tool execution completion
+        uiStream.append(
+          <BotMessage content="Tool execution complete. Synthesizing final answer..." />
+        )
+      } catch (e) {
+        console.error('Tool Coordinator failed:', e)
+        uiStream.append(
+          <BotMessage content="Tool Coordinator failed. Falling back to streaming researcher." />
+        )
+        // Fallback: continue with the original messages and let the researcher handle it
+        finalMessages = messages
+      }
+    }
+
     while (
       useSpecificAPI
         ? answer.length === 0
         : answer.length === 0 && !errorOccurred
     ) {
+      // If coordinator was used, pass finalMessages and disable tools for researcher
       const { fullResponse, hasError, toolResponses } = await researcher(
         currentSystemPrompt,
         uiStream,
         streamText,
-        messages,
-        useSpecificAPI
+        finalMessages,
+        useSpecificAPI,
+        !useToolCoordinator // Pass a flag to disable tools if coordinator was used
       )
       answer = fullResponse
       toolOutputs = toolResponses

lib/agents/index.tsx

@@ -3,3 +3,4 @@ export * from './inquire'
 export * from './query-suggestor'
 export * from './researcher'
 export * from './resolution-search'
+export * from './tool-coordinator'

lib/agents/researcher.tsx

@@ -78,7 +78,8 @@ export async function researcher(
   uiStream: ReturnType<typeof createStreamableUI>,
   streamText: ReturnType<typeof createStreamableValue<string>>,
   messages: CoreMessage[],
-  useSpecificModel?: boolean
+  useSpecificModel?: boolean,
+  useTools: boolean = true
 ) {
   let fullResponse = ''
   let hasError = false
@@ -101,7 +102,7 @@ export async function researcher(
     maxTokens: 4096,
     system: systemPromptToUse,
     messages,
-    tools: getTools({ uiStream, fullResponse }),
+    tools: useTools ? getTools({ uiStream, fullResponse }) : undefined,
   })
 
   uiStream.update(null) // remove spinner

lib/agents/tool-coordinator.tsx

@@ -0,0 +1,197 @@
+import { generateObject } from 'ai'
+import { z } from 'zod'
+import { Message } from 'ai/react'
+import { getTools } from '@/lib/agents/tools'
+import { ToolResultPart } from '@/lib/types'
+
+// ——————————————————————————————————————
+// Fallbacks if the original files don't exist yet
+// ——————————————————————————————————————
+
+let getModel: () => any
+let createStreamableUI: () => any
+
+try {
+  // Try the most common real locations first
+  const models = require('@/lib/models')
+  getModel = models.getModel || models.default || (() => null)
+} catch {
+  try {
+    const mod = require('@/lib/ai/models')
+    getModel = mod.getModel || mod.default
+  } catch {
+    getModel = () => {
+      throw new Error('getModel not available — check your @/lib/models setup')
+    }
+  }
+}
+
+try {
+  const streamable = require('@/lib/streamable')
+  createStreamableUI = streamable.createStreamableUI || streamable.default
+} catch {
+  try {
+    const s = require('@/lib/ui/streamable')
+    createStreamableUI = s.createStreamableUI
+  } catch {
+    // Minimal no-op version that won't break tool calling
+    createStreamableUI = () => ({
+      append: () => {},
+      update: () => {},
+      done: () => {},
+      value: null
+    })
+  }
+}
+
+// ——————————————————————————————————————
+// Schemas
+// ——————————————————————————————————————
+
+const toolStepSchema = z.object({
+  toolName: z.string(),
+  toolArgs: z.record(z.any()),
+  dependencyIndices: z.array(z.number()).optional(),
+  purpose: z.string()
+})
+
+const toolPlanSchema = z.object({
+  reasoning: z.string(),
+  steps: z.array(toolStepSchema)
+})
+
+export type ToolPlan = z.infer<typeof toolPlanSchema>
+export type ToolStep = z.infer<typeof toolStepSchema>
+
+// ——————————————————————————————————————
+// 1. Plan Generation
+// ——————————————————————————————————————
+
+export async function toolCoordinator(messages: Message[]): Promise<ToolPlan> {
+  const model = getModel()
+
+  const toolsObj = getTools({
+    uiStream: createStreamableUI(),
+    fullResponse: ''
+  })
+
+  const toolDescriptions = Object.values(toolsObj).map(tool => ({
+    name: tool.toolName,
+    description: tool.description,
+    parameters: tool.parameters
+  }))
+
+  const systemPrompt = `You are an expert Tool Coordinator. Create a precise multi-step plan using only these tools.
+
+Rules:
+- Use exact toolName from the list.
+- Use dependencyIndices (0-based) when a step needs prior results.
+- Output must be valid JSON matching the schema.
+
+Available Tools:
+${JSON.stringify(toolDescriptions, null, 2)}
+`
+
+  const { object } = await generateObject({
+    model,
+    system: systemPrompt,
+    messages,
+    schema: toolPlanSchema
+  })
+
+  return object
+}
+
+// ——————————————————————————————————————
+// 2. Execution
+// ——————————————————————————————————————
+
+interface ExecutionContext {
+  uiStream: any
+  fullResponse: string
+}
+
+export async function executeToolPlan(
+  plan: ToolPlan,
+  context: ExecutionContext
+): Promise<ToolResultPart[]> {
+  const { uiStream, fullResponse } = context
+
+  const toolsObj = getTools({ uiStream, fullResponse })
+  const toolMap = new Map(Object.values(toolsObj).map(t => [t.toolName, t]))
+
+  const results = new Map<number, any>()
+  const toolResults: ToolResultPart[] = []
+
+  const resolveDeps = (indices: number[] = []) =>
+    indices.map(i => {
+      if (!results.has(i)) throw new Error(`Dependency step ${i} missing`)
+      return results.get(i)
+    })
+
+  for (let i = 0; i < plan.steps.length; i++) {
+    const step = plan.steps[i]
+    const tool = toolMap.get(step.toolName)
+
+    let result: any = { error: `Tool "${step.toolName}" not found` }
+
+    try {
+      if (!tool) throw new Error(`Tool not found: ${step.toolName}`)
+
+      const deps = step.dependencyIndices ? resolveDeps(step.dependencyIndices) : []
+      const args = {
+        ...step.toolArgs,
+        ...(deps.length > 0 && { _dependencyResults: deps })
+      }
+
+      console.log(`[ToolCoordinator] Step ${i}: ${step.toolName}`)
+      result = await tool.execute(args)
+    } catch (err: any) {
+      const msg = err?.message || String(err)
+      console.error(`[ToolCoordinator] Step ${i} failed:`, msg)
+      result = { error: msg }
+    }
+
+    results.set(i, result)
+    toolResults.push({
+      toolName: step.toolName,
+      toolCallId: `coord-${i}`,
+      result
+    })
+  }
+
+  return toolResults
+}
+
+// ——————————————————————————————————————
+// 3. Aggregation
+// ——————————————————————————————————————
+
+export function aggregateToolResults(toolResults: ToolResultPart[], plan: ToolPlan): string {
+  let out = `# Tool Coordinator Results
+
+### Plan
+${plan.reasoning}
+
+### Steps
+`
+
+  toolResults.forEach((tr, i) => {
+    const step = plan.steps[i]
+    const hasError = tr.result && typeof tr.result === 'object' && 'error' in tr.result
+
+    out += `\n#### Step ${i + 1}: ${step.purpose} (\`${step.toolName}\`)`
+
+    if (hasError) {
+      out += `\n**Status:** Failed\n**Error:** ${tr.result.error}`
+    } else {
+      const json = JSON.stringify(tr.result, null, 2)
+      const truncated = json.length > 600 ? json.slice(0, 600) + '...' : json
+      out += `\n**Status:** Success\n**Result:**\n\`\`\`json\n${truncated}\n\`\`\``
+    }
+  })
+
+  out += `\n\n---\n**INSTRUCTION:** Write a natural, helpful final answer using only the information above. Do not mention tools, steps, or internal process.`
+
+  return out
+}

lib/agents/tools/geospatial.tsx

@@ -217,10 +217,21 @@ Uses the Mapbox Search Box Text Search API endpoint to power searching for and g
 
 ,
   parameters: geospatialQuerySchema,
-  execute: async (params: z.infer<typeof geospatialQuerySchema>) => {
-    const { queryType, includeMap = true } = params;
+  execute: async (params: z.infer<typeof geospatialQuerySchema> & { _dependencyResults?: any[] }) => {
+    const { queryType, includeMap = true, _dependencyResults } = params;
     console.log('[GeospatialTool] Execute called with:', params);
 
+    if (_dependencyResults && _dependencyResults.length > 0) {
+      console.log('[GeospatialTool] Processing dependency results:', _dependencyResults);
+      // Logic to process dependency results can be added here.
+      // For example, if a previous step was a search, the result might contain coordinates
+      // that can be used as input for a subsequent directions query.
+      // Since the full logic for dependency injection is complex and depends on the
+      // specific tool schema, we will log it for now and ensure the tool can handle it.
+      // The LLM planning step is responsible for generating the correct 'params'
+      // based on the dependency results. The tool only needs to be aware of them.
+    }
+
     const uiFeedbackStream = createStreamableValue<string>();
     uiStream.append(<BotMessage content={uiFeedbackStream.value} />);