Develop

config/dev.exs

@@ -1,4 +1,6 @@
 import Config
 
 # Development-specific configuration
-config :logger, level: :debug
+# Set to :error to avoid log messages messing up the TUI
+# Change to :debug or :info when debugging without the TUI
+config :logger, level: :error

lib/jido_code/agents/llm_agent.ex

@@ -46,6 +46,7 @@ defmodule JidoCode.Agents.LLMAgent do
   alias Jido.AI.Model.Registry.Adapter, as: RegistryAdapter
   alias Jido.AI.Prompt
   alias JidoCode.Config
+  alias JidoCode.Language
   alias JidoCode.PubSubTopics
   alias JidoCode.Session.ProcessRegistry
   alias JidoCode.Session.State, as: SessionState
@@ -58,7 +59,8 @@ defmodule JidoCode.Agents.LLMAgent do
 
   # System prompt should NOT include user input to prevent prompt injection attacks.
   # User messages are passed separately to the AI agent via chat_response/3.
-  @system_prompt """
+  # The base prompt is extended with language-specific instructions at runtime.
+  @base_system_prompt """
   You are JidoCode, an expert coding assistant running in a terminal interface.
 
   Your capabilities:
@@ -76,6 +78,9 @@ defmodule JidoCode.Agents.LLMAgent do
   - Acknowledge limitations when you're uncertain
   """
 
+  # For backwards compatibility and non-session contexts
+  @system_prompt @base_system_prompt
+
   # ============================================================================
   # Client API
   # ============================================================================
@@ -506,7 +511,8 @@ defmodule JidoCode.Agents.LLMAgent do
               ai_pid: ai_pid,
               config: config,
               session_id: actual_session_id,
-              topic: build_topic(actual_session_id)
+              topic: build_topic(actual_session_id),
+              is_processing: false
             }
 
             {:ok, state}
@@ -539,6 +545,12 @@ defmodule JidoCode.Agents.LLMAgent do
     {:noreply, state}
   end
 
+  @impl true
+  def handle_info(:stream_complete, state) do
+    # Reset processing state when stream completes (success or failure)
+    {:noreply, %{state | is_processing: false}}
+  end
+
   @impl true
   def handle_call({:chat, message}, from, state) do
     # ARCH-1 Fix: Use Task.Supervisor for monitored async tasks
@@ -576,8 +588,12 @@ defmodule JidoCode.Agents.LLMAgent do
 
   @impl true
   def handle_call(:get_status, _from, state) do
+    # ready is false when processing or when agent is not alive
+    agent_alive = is_pid(state.ai_pid) and Process.alive?(state.ai_pid)
+    ready = agent_alive and not state.is_processing
+
     status = %{
-      ready: is_pid(state.ai_pid) and Process.alive?(state.ai_pid),
+      ready: ready,
       config: state.config,
       session_id: state.session_id,
       topic: state.topic
@@ -627,23 +643,38 @@ defmodule JidoCode.Agents.LLMAgent do
     topic = state.topic
     config = state.config
     session_id = state.session_id
+    agent_pid = self()
 
     # ARCH-1 Fix: Use Task.Supervisor for monitored async streaming
     Task.Supervisor.start_child(JidoCode.TaskSupervisor, fn ->
+      # Trap exits to prevent ReqLLM's internal cleanup tasks from crashing us
+      Process.flag(:trap_exit, true)
+
       try do
         do_chat_stream_with_timeout(config, message, topic, timeout, session_id)
+        # Notify agent that streaming is complete
+        send(agent_pid, :stream_complete)
       catch
         :exit, {:timeout, _} ->
           Logger.warning("Stream timed out after #{timeout}ms")
           broadcast_stream_error(topic, :timeout)
+          send(agent_pid, :stream_complete)
 
         kind, reason ->
           Logger.error("Stream failed: #{kind} - #{inspect(reason)}")
           broadcast_stream_error(topic, {kind, reason})
+          send(agent_pid, :stream_complete)
+      end
+
+      # Drain any EXIT messages from ReqLLM cleanup tasks
+      receive do
+        {:EXIT, _pid, _reason} -> :ok
+      after
+        100 -> :ok
       end
     end)
 
-    {:noreply, state}
+    {:noreply, %{state | is_processing: true}}
   end
 
   @impl true
@@ -800,11 +831,11 @@ defmodule JidoCode.Agents.LLMAgent do
     Phoenix.PubSub.broadcast(@pubsub, topic, {:stream_chunk, session_id, chunk})
   end
 
-  defp broadcast_stream_end(topic, full_content, session_id) do
+  defp broadcast_stream_end(topic, full_content, session_id, metadata) do
     # Finalize message in Session.State (skip if session_id is PID string)
     end_session_streaming(session_id)
-    # Also broadcast for TUI (include session_id for routing)
-    Phoenix.PubSub.broadcast(@pubsub, topic, {:stream_end, session_id, full_content})
+    # Also broadcast for TUI (include session_id, content, and metadata for routing)
+    Phoenix.PubSub.broadcast(@pubsub, topic, {:stream_end, session_id, full_content, metadata})
   end
 
   defp broadcast_stream_error(topic, reason) do
@@ -851,11 +882,14 @@ defmodule JidoCode.Agents.LLMAgent do
   end
 
   defp execute_stream(model, message, topic, session_id) do
+    # Build dynamic system prompt with language-specific instructions
+    system_prompt = build_system_prompt(session_id)
+
     # Build prompt with system message and user message
     prompt =
       Prompt.new(%{
         messages: [
-          %{role: :system, content: @system_prompt, engine: :none},
+          %{role: :system, content: system_prompt, engine: :none},
           %{role: :user, content: message, engine: :none}
         ]
       })
@@ -884,33 +918,120 @@ defmodule JidoCode.Agents.LLMAgent do
     end
   end
 
-  defp process_stream(stream, topic, session_id) do
-    # Accumulate full content while streaming chunks
+  defp process_stream(stream_response, topic, session_id) do
+    # Extract inner stream and metadata_task from StreamResponse
+    {actual_stream, metadata_task} =
+      case stream_response do
+        %ReqLLM.StreamResponse{stream: inner, metadata_task: task} when inner != nil ->
+          {inner, task}
+
+        %ReqLLM.StreamResponse{metadata_task: task} ->
+          Logger.warning("LLMAgent: StreamResponse has nil stream field")
+          {[], task}
+
+        other ->
+          # Not a StreamResponse, just a raw stream
+          {other, nil}
+      end
+
+    # Accumulate full content while streaming chunks (catch handles ReqLLM process cleanup race conditions)
     full_content =
-      Enum.reduce_while(stream, "", fn chunk, acc ->
-        case extract_chunk_content(chunk) do
-          {:ok, content} ->
-            broadcast_stream_chunk(topic, content, session_id)
-            {:cont, acc <> content}
-
-          {:finish, content} ->
-            # Last chunk with finish_reason
-            if content != "" do
-              broadcast_stream_chunk(topic, content, session_id)
-            end
-
-            {:halt, acc <> content}
-        end
-      end)
+      try do
+        Enum.reduce_while(actual_stream, "", fn chunk, acc ->
+          case extract_chunk_content(chunk) do
+            {:ok, content} ->
+              # Only broadcast non-empty content
+              if content != "" do
+                broadcast_stream_chunk(topic, content, session_id)
+              end
+
+              {:cont, acc <> content}
+
+            {:finish, content} ->
+              if content != "" do
+                broadcast_stream_chunk(topic, content, session_id)
+              end
+
+              {:halt, acc <> content}
+          end
+        end)
+      rescue
+        e in Protocol.UndefinedError ->
+          Logger.error("LLMAgent: Protocol error during enumeration: #{inspect(e)}")
+          broadcast_stream_error(topic, e)
+          ""
+
+        e ->
+          Logger.error("LLMAgent: Error during enumeration: #{inspect(e)}")
+          broadcast_stream_error(topic, e)
+          ""
+      catch
+        :exit, {:noproc, _} ->
+          # StreamServer died - this is expected at end of stream due to ReqLLM cleanup
+          Logger.debug("LLMAgent: Stream server terminated (normal cleanup)")
+          ""
+
+        :exit, reason ->
+          Logger.warning("LLMAgent: Stream exited: #{inspect(reason)}")
+          ""
+      end
+
+    # Await metadata from the stream (usage info, finish_reason, etc.)
+    # This keeps the StreamServer alive until metadata is collected
+    metadata = await_stream_metadata(metadata_task)
+
+    # Log usage information if available
+    if metadata[:usage] do
+      Logger.info("LLMAgent: Token usage - #{inspect(metadata[:usage])}")
+    end
 
     # Broadcast stream completion and finalize in Session.State
-    broadcast_stream_end(topic, full_content, session_id)
+    broadcast_stream_end(topic, full_content, session_id, metadata)
   rescue
     error ->
       Logger.error("Stream processing error: #{inspect(error)}")
       broadcast_stream_error(topic, error)
   end
 
+  # Await metadata from the stream's metadata task
+  # Returns empty map if task is nil or fails
+  defp await_stream_metadata(nil), do: %{}
+
+  defp await_stream_metadata(task) do
+    try do
+      # Await with reasonable timeout (10 seconds)
+      Task.await(task, 10_000)
+    catch
+      :exit, {:timeout, _} ->
+        Logger.warning("LLMAgent: Metadata task timed out")
+        %{}
+
+      :exit, reason ->
+        Logger.debug("LLMAgent: Metadata task exited: #{inspect(reason)}")
+        %{}
+    end
+  end
+
+  # ReqLLM.StreamChunk format - content type with text field
+  defp extract_chunk_content(%ReqLLM.StreamChunk{type: :content, text: text}) do
+    {:ok, text || ""}
+  end
+
+  # ReqLLM.StreamChunk format - meta type signals end of stream
+  defp extract_chunk_content(%ReqLLM.StreamChunk{type: :meta, metadata: metadata}) do
+    if Map.get(metadata, :finish_reason) do
+      {:finish, ""}
+    else
+      {:ok, ""}
+    end
+  end
+
+  # ReqLLM.StreamChunk format - thinking/tool_call types (skip for now)
+  defp extract_chunk_content(%ReqLLM.StreamChunk{type: _type}) do
+    {:ok, ""}
+  end
+
+  # Legacy format - content with finish_reason
   defp extract_chunk_content(%{content: content, finish_reason: nil}) do
     {:ok, content || ""}
   end
@@ -919,6 +1040,7 @@ defmodule JidoCode.Agents.LLMAgent do
     {:finish, content || ""}
   end
 
+  # Legacy format - delta content
   defp extract_chunk_content(%{delta: %{content: content}} = chunk) do
     finish_reason = Map.get(chunk, :finish_reason)
 
@@ -934,8 +1056,8 @@ defmodule JidoCode.Agents.LLMAgent do
   end
 
   defp extract_chunk_content(chunk) do
-    # Unknown chunk format - try to extract content
-    content = Map.get(chunk, :content) || Map.get(chunk, "content") || ""
+    # Unknown chunk format - try to extract content or text
+    content = Map.get(chunk, :content) || Map.get(chunk, :text) || Map.get(chunk, "content") || ""
     {:ok, content}
   end
 
@@ -1022,6 +1144,40 @@ defmodule JidoCode.Agents.LLMAgent do
     not String.starts_with?(session_id, "#PID<")
   end
 
+  # ============================================================================
+  # System Prompt Building
+  # ============================================================================
+
+  # Build the system prompt with optional language-specific instructions
+  defp build_system_prompt(session_id) when is_binary(session_id) do
+    if is_valid_session_id?(session_id) do
+      case SessionState.get_state(session_id) do
+        {:ok, %{session: session}} when not is_nil(session.language) ->
+          add_language_instruction(@base_system_prompt, session.language)
+
+        _ ->
+          @base_system_prompt
+      end
+    else
+      @base_system_prompt
+    end
+  end
+
+  defp build_system_prompt(_), do: @base_system_prompt
+
+  # Add language-specific instruction to the system prompt
+  defp add_language_instruction(base_prompt, language) do
+    lang_name = Language.display_name(language)
+
+    language_instruction = """
+
+    Language Context:
+    This project uses #{lang_name}. When providing code snippets, write them in #{lang_name} unless a different language is specifically requested.
+    """
+
+    base_prompt <> language_instruction
+  end
+
   # ============================================================================
   # Validation Functions
   # ============================================================================