Fix: Improve robustness of location enrichment feature

app/actions.tsx

@@ -12,6 +12,7 @@ import { Spinner } from '@/components/ui/spinner'
 import { Section } from '@/components/section'
 import { FollowupPanel } from '@/components/followup-panel'
 import { inquire, researcher, taskManager, querySuggestor } from '@/lib/agents'
+import { geojsonEnricher } from '@/lib/agents/geojson-enricher'
 // 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'
@@ -25,11 +26,10 @@ import { CopilotDisplay } from '@/components/copilot-display'
 import RetrieveSection from '@/components/retrieve-section'
 import { VideoSearchSection } from '@/components/video-search-section'
 import { MapQueryHandler } from '@/components/map/map-query-handler' // Add this import
-
-// Define the type for related queries
-type RelatedQueries = {
-  items: { query: string }[]
-}
+import { LocationResponseHandler } from '@/components/map/location-response-handler'
+import { isLocationResponse } from '@/lib/utils/type-guards'
+import { PartialRelated } from '@/lib/schema/related'
+import { LocationResponse } from '@/lib/types/custom'
 
 // Removed mcp parameter from submit, as geospatialTool now handles its client.
 async function submit(formData?: FormData, skip?: boolean) {
@@ -222,7 +222,7 @@ async function submit(formData?: FormData, skip?: boolean) {
           {
             id: nanoid(),
             role: 'assistant',
-            content: `inquiry: ${inquiry?.question}`
+            content: `inquiry: ${(inquiry as any)?.question}`
           }
         ]
       })
@@ -296,6 +296,29 @@ async function submit(formData?: FormData, skip?: boolean) {
     }
 
     if (!errorOccurred) {
+      let locationResponse: LocationResponse;
+      try {
+        // Create a timeout promise that rejects after 30 seconds.
+        const timeoutPromise = new Promise((_, reject) =>
+          setTimeout(() => reject(new Error('GeoJSON enrichment timed out after 30 seconds')), 30000)
+        );
+
+        // Race the agent call against the timeout.
+        locationResponse = (await Promise.race([
+          geojsonEnricher(answer),
+          timeoutPromise
+        ])) as LocationResponse;
+      } catch (e) {
+        console.error("Error during geojson enrichment:", e);
+        // Fallback to a response without location data, ensuring consistent structure.
+        locationResponse = {
+          type: 'tool',
+          text: answer,
+          geojson: null,
+          map_commands: null,
+        };
+      }
+
       const relatedQueries = await querySuggestor(uiStream, messages)
       uiStream.append(
         <Section title="Follow-up">
@@ -312,9 +335,16 @@ async function submit(formData?: FormData, skip?: boolean) {
           {
             id: groupeId,
             role: 'assistant',
-            content: answer,
+            content: locationResponse.text,
             type: 'response'
           },
+          {
+            id: nanoid(),
+            role: 'tool',
+            name: 'geojsonEnrichment',
+            content: JSON.stringify(locationResponse),
+            type: 'tool',
+          },
           {
             id: groupeId,
             role: 'assistant',
@@ -518,7 +548,7 @@ export const getUIStateFromAIState = (aiState: AIState): UIState => {
                 )
               }
             case 'related':
-              const relatedQueries = createStreamableValue<RelatedQueries>()
+              const relatedQueries = createStreamableValue<PartialRelated>()
               relatedQueries.done(JSON.parse(content as string))
               return {
                 id,
@@ -556,6 +586,25 @@ export const getUIStateFromAIState = (aiState: AIState): UIState => {
               }
             }
 
+            if (name === 'geojsonEnrichment') {
+              // Use the type guard for robust runtime validation.
+              if (isLocationResponse(toolOutput)) {
+                return {
+                  id,
+                  component: (
+                    <LocationResponseHandler locationResponse={toolOutput} />
+                  ),
+                  isCollapsed: false,
+                };
+              } else {
+                // The type guard already logs the specific validation error.
+                console.warn(
+                  'Skipping render for geojsonEnrichment due to invalid data structure.'
+                );
+                return { id, component: null };
+              }
+            }
+
             const searchResults = createStreamableValue()
             searchResults.done(JSON.stringify(toolOutput))
             switch (name) {

components/map/location-response-handler.tsx

@@ -0,0 +1,28 @@
+'use client';
+
+import { useEffect } from 'react';
+import { useMapData } from './map-data-context';
+import { LocationResponse } from '@/lib/types/custom';
+
+interface LocationResponseHandlerProps {
+  locationResponse: LocationResponse;
+}
+
+export const LocationResponseHandler: React.FC<LocationResponseHandlerProps> = ({ locationResponse }) => {
+  const { setMapData } = useMapData();
+
+  useEffect(() => {
+    if (locationResponse) {
+      const { geojson, map_commands } = locationResponse;
+      console.log('LocationResponseHandler: Received data', locationResponse);
+      setMapData(prevData => ({
+        ...prevData,
+        geojson: geojson,
+        mapCommands: map_commands,
+      }));
+    }
+  }, [locationResponse, setMapData]);
+
+  // This component handles logic and does not render any UI.
+  return null;
+};

components/map/map-data-context.tsx

@@ -2,6 +2,10 @@
 
 import React, { createContext, useContext, useState, ReactNode } from 'react';
 import { LngLatLike } from 'mapbox-gl'; // Import LngLatLike
+import {
+  GeoJSONFeatureCollection,
+  MapCommand
+} from '@/lib/types/custom';
 
 // Define the shape of the map data you want to share
 export interface MapData {
@@ -14,6 +18,8 @@ export interface MapData {
     measurement: string;
     geometry: any;
   }>;
+  geojson?: GeoJSONFeatureCollection | null;
+  mapCommands?: MapCommand[] | null;
 }
 
 interface MapDataContextType {

components/map/mapbox-map.tsx

@@ -524,6 +524,87 @@ export const Mapbox: React.FC<{ position?: { latitude: number; longitude: number
     // }
   }, [mapData.targetPosition, mapData.mapFeature, updateMapPosition]);
 
+  // Effect to handle GeoJSON data updates
+  useEffect(() => {
+    if (!map.current) return;
+
+    const mapInstance = map.current;
+    const source = mapInstance.getSource('geojson-data');
+
+    // If GeoJSON data is present, add or update the source and layers
+    if (mapData.geojson) {
+      if (source) {
+        (source as mapboxgl.GeoJSONSource).setData(mapData.geojson as any);
+      } else {
+        mapInstance.addSource('geojson-data', {
+          type: 'geojson',
+          data: mapData.geojson as any,
+        });
+
+        // Add layer for points
+        mapInstance.addLayer({
+          id: 'geojson-points',
+          type: 'circle',
+          source: 'geojson-data',
+          paint: {
+            'circle-radius': 8,
+            'circle-color': '#007cbf',
+            'circle-stroke-width': 2,
+            'circle-stroke-color': '#ffffff',
+          },
+          filter: ['==', '$type', 'Point'],
+        });
+
+        // Add layer for lines
+        mapInstance.addLayer({
+          id: 'geojson-lines',
+          type: 'line',
+          source: 'geojson-data',
+          paint: {
+            'line-color': '#ff4500',
+            'line-width': 3,
+          },
+          filter: ['==', '$type', 'LineString'],
+        });
+      }
+    } else {
+      // If no GeoJSON data, remove layers and source if they exist
+      if (mapInstance.getLayer('geojson-points')) mapInstance.removeLayer('geojson-points');
+      if (mapInstance.getLayer('geojson-lines')) mapInstance.removeLayer('geojson-lines');
+      if (source) mapInstance.removeSource('geojson-data');
+    }
+  }, [mapData.geojson]);
+
+  // Effect to execute map commands
+  useEffect(() => {
+    if (!map.current || !mapData.mapCommands || mapData.mapCommands.length === 0) return;
+
+    const mapInstance = map.current;
+
+    mapData.mapCommands.forEach(command => {
+      switch (command.command) {
+        case 'flyTo':
+          mapInstance.flyTo(command.params);
+          break;
+        case 'easeTo':
+          mapInstance.easeTo(command.params);
+          break;
+        case 'fitBounds':
+          // Destructure params to pass 'bounds' and 'options' separately,
+          // ensuring the function is called with the correct signature.
+          const { bounds, options } = command.params;
+          mapInstance.fitBounds(bounds, options || {});
+          break;
+        default:
+          console.warn(`Unknown map command: ${command.command}`);
+      }
+    });
+
+    // Clear commands after execution to prevent re-triggering
+    setMapData(prev => ({ ...prev, mapCommands: null }));
+
+  }, [mapData.mapCommands, setMapData]);
+
   // Long-press handlers
   const handleMouseDown = useCallback(() => {
     // Only activate long press if not in real-time mode (as that mode has its own interactions)

lib/agents/geojson-enricher.tsx

@@ -0,0 +1,92 @@
+import { CoreMessage, LanguageModel, streamText } from 'ai';
+import { getModel } from '../utils';
+import { LocationResponse } from '../types/custom';
+import { isLocationResponse } from '../utils/type-guards';
+
+// A specialized prompt instructing the LLM to parse a textual response
+// and extract structured GeoJSON data and map commands.
+const GEOJSON_ENRICHMENT_PROMPT = `
+You are an AI assistant specializing in geospatial data extraction.
+The final output MUST be a single JSON object that strictly follows the LocationResponse interface.
+Return ONLY the raw JSON object with no surrounding markdown, code fences, or any additional text or explanation.
+
+Your task is to process a given text and extract the following information:
+
+1.  "type": This field must always be the string 'tool'.
+2.  "text": The original textual response that should be displayed to the user.
+3.  "geojson": A valid GeoJSON FeatureCollection representing any locations, addresses, coordinates, or routes mentioned in the text.
+4.  "map_commands": A list of map camera commands to control the map view, such as flying to a location.
+
+Rules for GeoJSON:
+- Convert all found locations into appropriate GeoJSON features (Point, LineString).
+- Use the correct coordinate format: [Longitude, Latitude] in WGS84.
+- Include meaningful properties for each feature (e.g., "name", "description").
+- If no geographic data can be extracted, set "geojson" to null.
+
+Rules for Map Commands:
+- Identify actions in thetext that imply map movements (e.g., "fly to," "center on," "zoom to").
+- Create a list of command objects, for example: { "command": "flyTo", "params": { "center": [-71.05633, 42.356823], "zoom": 15 } }.
+- If no map commands can be inferred, set "map_commands" to null.
+
+Here is the text to process:
+`;
+
+/**
+ * An asynchronous agent that enriches a textual response with GeoJSON data and map commands.
+ * @param researcherResponse The text generated by the researcher agent.
+ * @returns A promise that resolves to a LocationResponse object.
+ */
+export async function geojsonEnricher(
+  researcherResponse: string
+): Promise<LocationResponse> {
+  const model = getModel() as LanguageModel;
+  const messages: CoreMessage[] = [
+    {
+      role: 'user',
+      content: `${GEOJSON_ENRICHMENT_PROMPT}\n\n${researcherResponse}`,
+    },
+  ];
+
+  try {
+    // Await the streaming text promise to get the full response.
+    const { text } = await streamText({
+      model,
+      messages,
+      maxTokens: 4096, // Increased maxTokens for potentially larger GeoJSON payloads.
+    });
+    let responseText = await text;
+
+    // Strip any surrounding markdown code fences.
+    const jsonMatch = responseText.match(/```(json)?\n([\s\S]*?)\n```/);
+    if (jsonMatch && jsonMatch[2]) {
+      responseText = jsonMatch[2].trim();
+    }
+
+    // Parse the cleaned text into a JSON object.
+    const parsedJson = JSON.parse(responseText);
+
+    // Validate the parsed object against the LocationResponse interface.
+    if (isLocationResponse(parsedJson)) {
+      return parsedJson;
+    } else {
+      // If validation fails, log the error and the invalid object.
+      console.error('Validation failed: The parsed object does not match the LocationResponse interface.', parsedJson);
+      // Return a default response to avoid crashing the application.
+      return {
+        type: 'tool',
+        text: researcherResponse,
+        geojson: null,
+        map_commands: null,
+      };
+    }
+  } catch (error) {
+    console.error('Error enriching response with GeoJSON:', error);
+    // If parsing or any other part of the process fails, return a default response.
+    return {
+      type: 'tool',
+      text: researcherResponse,
+      geojson: null,
+      map_commands: null,
+    };
+  }
+}

lib/types/custom.ts

@@ -0,0 +1,34 @@
+// Defines the structure for a map command, like 'flyTo' or 'easeTo'.
+export interface MapCommand {
+  command: 'flyTo' | 'easeTo' | 'fitBounds'; // Add other valid map commands as needed
+  params: any; // Parameters for the command, e.g., { center: [lon, lat], zoom: 10 }
+}
+
+// Defines the structure for the geometry part of a GeoJSON feature.
+export interface GeoJSONGeometry {
+  type: 'Point' | 'LineString' | 'Polygon'; // Can be extended with other GeoJSON geometry types
+  coordinates: number[] | number[][] | number[][][];
+}
+
+// Defines a single feature in a GeoJSON FeatureCollection.
+export interface GeoJSONFeature {
+  type: 'Feature';
+  geometry: GeoJSONGeometry;
+  properties: {
+    [key: string]: any; // Features can have any number of properties
+  };
+}
+
+// Defines the structure for a GeoJSON FeatureCollection.
+export interface GeoJSONFeatureCollection {
+  type: 'FeatureCollection';
+  features: GeoJSONFeature[];
+}
+
+// Defines the structured response that includes textual data, GeoJSON, and map commands.
+export interface LocationResponse {
+  type: 'tool';
+  text: string;
+  geojson: GeoJSONFeatureCollection | null;
+  map_commands?: MapCommand[] | null;
+}