diff --git a/ai-assistant/ai-assistant.html b/ai-assistant/ai-assistant.html index a0f5cf25..8589741c 100644 --- a/ai-assistant/ai-assistant.html +++ b/ai-assistant/ai-assistant.html @@ -70,6 +70,10 @@

2. Core Functionality

and generate code from templates.
  • Dual inference backends — fully offline on-device inference, or Gemini in the cloud, selectable in Settings.
    • +
  • Direct commands — explicit requests like "open MainActivity.java", + "list files", "read <file>" or "search <query>" run the tool + directly (resolving a bare filename to its path), so they work reliably on + any model.
  • Safety controls — filesystem tools are confined to the project root, and mutating tools require explicit user approval.
    • diff --git a/ai-assistant/build.gradle.kts b/ai-assistant/build.gradle.kts index b376da7c..e8032dc3 100644 --- a/ai-assistant/build.gradle.kts +++ b/ai-assistant/build.gradle.kts @@ -45,6 +45,10 @@ android { } } + testOptions { + unitTests.isReturnDefaultValues = true + } + packaging { resources { excludes += setOf( @@ -77,10 +81,10 @@ dependencies { // JSON serialization for session persistence implementation("com.google.code.gson:gson:2.10.1") - // Plugin dependencies are loaded at runtime by the plugin manager - // No explicit compile-time dependency on the ai-core plugin needed + testImplementation(files("../libs/plugin-api.jar")) testImplementation("junit:junit:4.13.2") testImplementation("io.mockk:mockk:1.13.8") + testImplementation("org.json:json:20231013") testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:1.7.3") testImplementation("androidx.arch.core:core-testing:2.2.0") } diff --git a/ai-assistant/src/main/AndroidManifest.xml b/ai-assistant/src/main/AndroidManifest.xml index 78fed1cc..ca64c28a 100644 --- a/ai-assistant/src/main/AndroidManifest.xml +++ b/ai-assistant/src/main/AndroidManifest.xml @@ -1,6 +1,5 @@ - + { holder.loadingIndicator.visibility = View.VISIBLE holder.messageContent.visibility = View.GONE - holder.generatingDots.visibility = View.GONE + stopGeneratingDots(holder) } MessageStatus.SENT -> { holder.loadingIndicator.visibility = View.GONE @@ -130,19 +134,19 @@ class ChatAdapter( if (message.sender == Sender.AGENT && message.durationMs == null) { animateGeneratingDots(holder) } else { - holder.generatingDots.visibility = View.GONE + stopGeneratingDots(holder) } } MessageStatus.COMPLETED -> { holder.loadingIndicator.visibility = View.GONE holder.messageContent.visibility = View.VISIBLE - holder.generatingDots.visibility = View.GONE + stopGeneratingDots(holder) markwon.setMarkdown(holder.messageContent, payload.text) } MessageStatus.ERROR -> { holder.loadingIndicator.visibility = View.GONE holder.messageContent.visibility = View.VISIBLE - holder.generatingDots.visibility = View.GONE + stopGeneratingDots(holder) holder.messageContent.text = payload.text } } @@ -185,14 +189,14 @@ class ChatAdapter( if (message.sender == Sender.AGENT && message.durationMs == null) { animateGeneratingDots(holder) } else { - holder.generatingDots.visibility = View.GONE + stopGeneratingDots(holder) } } MessageStatus.COMPLETED -> { holder.loadingIndicator.visibility = View.GONE holder.messageContent.visibility = View.VISIBLE holder.btnRetry.visibility = View.GONE - holder.generatingDots.visibility = View.GONE + stopGeneratingDots(holder) markwon.setMarkdown(holder.messageContent, message.text) updateMessageMetadata(holder, message) } @@ -200,15 +204,15 @@ class ChatAdapter( holder.loadingIndicator.visibility = View.GONE holder.messageContent.visibility = View.VISIBLE holder.btnRetry.visibility = View.VISIBLE - holder.generatingDots.visibility = View.GONE + stopGeneratingDots(holder) holder.messageContent.text = message.text if (message.sender == Sender.SYSTEM) { - holder.btnRetry.text = "Open AI Settings" + holder.btnRetry.text = holder.btnRetry.context.getString(R.string.action_open_settings) holder.btnRetry.setOnClickListener { onMessageAction(ACTION_OPEN_SETTINGS, message) } } else { - holder.btnRetry.text = "Retry" + holder.btnRetry.text = holder.btnRetry.context.getString(R.string.action_retry) holder.btnRetry.setOnClickListener { onMessageAction(ACTION_RETRY, message) } @@ -236,7 +240,7 @@ class ChatAdapter( private fun updateSystemMessageExpansion(holder: SystemMessageViewHolder, message: ChatMessage) { val isExpanded = expandedMessageIds.contains(message.id) if (isExpanded) { - holder.messageHeaderTitle.text = "System Log" + holder.messageHeaderTitle.text = holder.messageHeaderTitle.context.getString(R.string.system_log) holder.messageContent.visibility = View.VISIBLE holder.expandIcon.rotation = 180f } else { @@ -247,21 +251,36 @@ class ChatAdapter( } private fun animateGeneratingDots(holder: DefaultMessageViewHolder) { + // Already animating: don't restart, or each streaming-token rebind would + // reset the loop to "." and churn a new Runnable. Torn down on complete/recycle. + if (holder.dotsRunnable != null) return holder.generatingDots.visibility = View.VISIBLE val dotStates = arrayOf(".", "..", "...") var currentIndex = 0 - val handler = android.os.Handler(android.os.Looper.getMainLooper()) val runnable = object : Runnable { override fun run() { - if (holder.generatingDots.visibility == View.VISIBLE) { - holder.generatingDots.text = dotStates[currentIndex] - currentIndex = (currentIndex + 1) % dotStates.size - handler.postDelayed(this, 500) - } + holder.generatingDots.text = dotStates[currentIndex] + currentIndex = (currentIndex + 1) % dotStates.size + holder.dotsHandler.postDelayed(this, 500) } } - handler.post(runnable) + holder.dotsRunnable = runnable + holder.dotsHandler.post(runnable) + } + + /** Stop and clear this holder's "…" animation. Safe to call repeatedly. */ + private fun stopGeneratingDots(holder: DefaultMessageViewHolder) { + holder.dotsRunnable?.let { holder.dotsHandler.removeCallbacks(it) } + holder.dotsRunnable = null + holder.generatingDots.apply { visibility = View.GONE } + } + + override fun onViewRecycled(holder: RecyclerView.ViewHolder) { + super.onViewRecycled(holder) + if (holder is DefaultMessageViewHolder) { + stopGeneratingDots(holder) + } } private fun createPreview(rawText: String): String { @@ -334,7 +353,7 @@ class ChatAdapter( val clipboard = context.getSystemService(Context.CLIPBOARD_SERVICE) as ClipboardManager val clip = ClipData.newPlainText("chat_message", message.text) clipboard.setPrimaryClip(clip) - Toast.makeText(context, "Copied", Toast.LENGTH_SHORT).show() + Snackbar.make(view, view.context.getString(R.string.msg_copied), Snackbar.LENGTH_SHORT).show() true } 2 -> { diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/AiSettingsFragment.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/AiSettingsFragment.kt index 83c285c1..e1560b8a 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/AiSettingsFragment.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/AiSettingsFragment.kt @@ -73,9 +73,9 @@ class AiSettingsFragment : DialogFragment() { val uriString = it.toString() viewModel.loadModelFromUri(uriString, requireContext()) - Toast.makeText(requireContext(), "Loading model...", Toast.LENGTH_SHORT).show() + Toast.makeText(requireContext(), getString(R.string.model_loading_toast), Toast.LENGTH_SHORT).show() } catch (e: Exception) { - Toast.makeText(requireContext(), "Error: ${e.message}", Toast.LENGTH_LONG).show() + Toast.makeText(requireContext(), getString(R.string.state_error, e.message), Toast.LENGTH_LONG).show() } } } @@ -233,12 +233,12 @@ class AiSettingsFragment : DialogFragment() { viewModel.engineState.observe(viewLifecycleOwner) { state -> when (state) { is EngineState.Initializing, EngineState.Uninitialized -> { - engineStatusTextView.text = "Initializing engine..." + engineStatusTextView.text = getString(R.string.engine_initializing) browseButton.isEnabled = false loadSavedButton.isEnabled = false } is EngineState.Initialized -> { - engineStatusTextView.text = "Engine ready" + engineStatusTextView.text = getString(R.string.engine_ready) browseButton.isEnabled = true loadSavedButton.isEnabled = viewModel.savedModelPath.value != null } @@ -257,7 +257,7 @@ class AiSettingsFragment : DialogFragment() { if (path != null) { modelPathTextView.visibility = View.VISIBLE val fileName = viewModel.getSavedModelName() ?: viewModel.fallbackDisplayName(path) - modelPathTextView.text = "Saved: $fileName" + modelPathTextView.text = getString(R.string.model_saved_path, fileName) } else { modelPathTextView.visibility = View.GONE } @@ -268,19 +268,19 @@ class AiSettingsFragment : DialogFragment() { when (state) { is ModelLoadingState.Idle -> { modelStatusTextView.visibility = View.VISIBLE - modelStatusTextView.text = "No model is currently loaded" + modelStatusTextView.text = getString(R.string.model_none_loaded) } is ModelLoadingState.Loading -> { modelStatusTextView.visibility = View.VISIBLE - modelStatusTextView.text = "Loading model, please wait..." + modelStatusTextView.text = getString(R.string.model_loading_wait) } is ModelLoadingState.Loaded -> { modelStatusTextView.visibility = View.VISIBLE - modelStatusTextView.text = "✅ Model loaded: ${state.modelName}" + modelStatusTextView.text = getString(R.string.model_loaded, state.modelName) } is ModelLoadingState.Error -> { modelStatusTextView.visibility = View.VISIBLE - modelStatusTextView.text = "❌ Error: ${state.message}" + modelStatusTextView.text = getString(R.string.model_load_error, state.message) } } } @@ -328,9 +328,9 @@ class AiSettingsFragment : DialogFragment() { if (timestamp > 0) { val sdf = SimpleDateFormat("MMMM d, yyyy", Locale.getDefault()) val savedDate = sdf.format(Date(timestamp)) - statusTextView.text = "API Key saved on: $savedDate" + statusTextView.text = getString(R.string.api_key_saved_on, savedDate) } else { - statusTextView.text = "API Key is saved" + statusTextView.text = getString(R.string.api_key_saved) } } @@ -338,15 +338,15 @@ class AiSettingsFragment : DialogFragment() { val apiKey = apiKeyInput.text.toString() if (apiKey.isNotBlank()) { viewModel.saveGeminiApiKey(apiKey) - Toast.makeText(requireContext(), "API Key saved", Toast.LENGTH_SHORT).show() + Toast.makeText(requireContext(), getString(R.string.api_key_saved_toast), Toast.LENGTH_SHORT).show() updateUiState(isEditing = false) val timestamp = viewModel.getGeminiApiKeySaveTimestamp() val sdf = SimpleDateFormat("MMMM d, yyyy", Locale.getDefault()) val savedDate = sdf.format(Date(timestamp)) - statusTextView.text = "API Key saved on: $savedDate" + statusTextView.text = getString(R.string.api_key_saved_on, savedDate) } else { - Toast.makeText(requireContext(), "API Key cannot be empty", Toast.LENGTH_SHORT).show() + Toast.makeText(requireContext(), getString(R.string.api_key_empty), Toast.LENGTH_SHORT).show() } } @@ -358,7 +358,7 @@ class AiSettingsFragment : DialogFragment() { clearButton.setOnClickListener { viewModel.clearGeminiApiKey() - Toast.makeText(requireContext(), "API Key cleared", Toast.LENGTH_SHORT).show() + Toast.makeText(requireContext(), getString(R.string.api_key_cleared), Toast.LENGTH_SHORT).show() updateUiState(isEditing = true) apiKeyInput.setText("") } @@ -462,7 +462,7 @@ class AiSettingsFragment : DialogFragment() { // Observe loading state viewModel.geminiModelsLoading.observe(viewLifecycleOwner) { isLoading -> refreshButton.isEnabled = !isLoading - refreshButton.text = if (isLoading) "Loading..." else "Refresh Models" + refreshButton.text = if (isLoading) getString(R.string.loading) else getString(R.string.refresh_models) } modelSpinner.setOnTouchListener { _, _ -> @@ -478,8 +478,8 @@ class AiSettingsFragment : DialogFragment() { val selectedModel = parent?.getItemAtPosition(position) as? String if (selectedModel != null && selectedModel != viewModel.getGeminiModel()) { viewModel.saveGeminiModel(selectedModel) - currentModelText?.text = "Current: $selectedModel" - Toast.makeText(requireContext(), "Model changed to $selectedModel", Toast.LENGTH_SHORT).show() + currentModelText?.text = getString(R.string.current_model, selectedModel) + Toast.makeText(requireContext(), getString(R.string.model_changed, selectedModel), Toast.LENGTH_SHORT).show() } } diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/ChatFragment.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/ChatFragment.kt index a293646f..1f7d9cc0 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/ChatFragment.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/fragments/ChatFragment.kt @@ -14,6 +14,7 @@ import androidx.lifecycle.repeatOnLifecycle import androidx.recyclerview.widget.LinearLayoutManager import com.google.android.material.chip.Chip import com.itsaky.androidide.plugins.PluginContext +import com.itsaky.androidide.plugins.aiassistant.R import com.itsaky.androidide.plugins.aiassistant.adapters.ChatAdapter import com.itsaky.androidide.plugins.aiassistant.databinding.FragmentChatBinding import com.itsaky.androidide.plugins.aiassistant.models.AgentState @@ -287,33 +288,32 @@ class ChatFragment : Fragment() { is AgentState.Idle -> { binding.agentStatusContainer.isVisible = false binding.sendButton.isEnabled = true - binding.sendButton.text = "Send" + binding.sendButton.text = getString(R.string.send) } is AgentState.Executing -> { binding.agentStatusContainer.isVisible = true binding.agentStatusMessage.text = state.formattedProgress binding.agentStatusTimer.text = state.formattedTiming binding.sendButton.isEnabled = true - binding.sendButton.text = "Stop" - viewModel.startStateTimer(state) + binding.sendButton.text = getString(R.string.btn_stop) } is AgentState.Processing -> { binding.agentStatusContainer.isVisible = true - binding.agentStatusMessage.text = "Generating response..." + binding.agentStatusMessage.text = getString(R.string.generating_response) binding.agentStatusTimer.text = "" binding.sendButton.isEnabled = true - binding.sendButton.text = "Stop" + binding.sendButton.text = getString(R.string.btn_stop) } is AgentState.Error -> { binding.agentStatusContainer.isVisible = false binding.sendButton.isEnabled = true - binding.sendButton.text = "Send" + binding.sendButton.text = getString(R.string.send) viewModel.stopStateTimer() showErrorSnackbar(state.message) } else -> { binding.sendButton.isEnabled = false - binding.sendButton.text = "Send" + binding.sendButton.text = getString(R.string.send) } } } diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/AgentLoop.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/AgentLoop.kt new file mode 100644 index 00000000..8c7e8f04 --- /dev/null +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/AgentLoop.kt @@ -0,0 +1,215 @@ +package com.itsaky.androidide.plugins.aiassistant.tool + +import com.itsaky.androidide.plugins.aiassistant.models.ToolResult +import com.itsaky.androidide.plugins.services.LlmInferenceService +import com.itsaky.androidide.plugins.services.LlmInferenceService.ChatMessage +import com.itsaky.androidide.plugins.services.LlmInferenceService.ChatMessage.Role + +/** + * The agentic tool-loop: each turn renders the transcript into a prompt, generates a + * reply, and runs any tool calls, looping until the model stops or a limit is hit. + * Free of Android/coroutine/UI deps so it unit-tests with plain fakes. + */ +class AgentLoop( + private val maxIterations: Int = DEFAULT_MAX_ITERATIONS, + private val toolOutputCharLimit: Int = DEFAULT_TOOL_OUTPUT_CHAR_LIMIT, + private val maxConsecutiveRepeats: Int = DEFAULT_MAX_CONSECUTIVE_REPEATS, + private val extractToolCalls: (String) -> List = ToolCallExtractor::extractToolCalls, + private val terminalTool: String? = null, +) { + + companion object { + /** Max model turns per user message; a backstop against a model that never stops calling tools. */ + const val DEFAULT_MAX_ITERATIONS = 8 + + /** Per-tool-result cap fed back into the prompt, so big outputs don't blow a local model's context. */ + const val DEFAULT_TOOL_OUTPUT_CHAR_LIMIT = 4000 + + /** + * Consecutive identical tool-call batches tolerated before aborting as + * [StopReason.REPEATED]; a truncated result can make one repeat legitimate. + */ + const val DEFAULT_MAX_CONSECUTIVE_REPEATS = 2 + } + + /** Callbacks so the caller can drive UI/state; all no-ops by default. */ + interface Events { + /** + * A model turn finished. + * @param turn 1-based turn index. + * @param text the model's reply, already streamed into the UI. + */ + suspend fun onModelTurn(turn: Int, text: String) {} + + /** + * A tool batch was executed; the caller renders it. + * @param turn 1-based turn index. + * @param calls the tool calls that ran. + * @param results their results, positionally aligned with [calls]. + */ + suspend fun onToolResults(turn: Int, calls: List, results: List) {} + + /** + * The loop stopped after hitting the iteration cap while still calling tools. + * @param turns total turns run. + */ + suspend fun onMaxIterationsReached(turns: Int) {} + + /** + * The loop stopped after the model repeated identical tool calls. + * @param turns total turns run. + */ + suspend fun onRepeatedToolCalls(turns: Int) {} + + /** + * The model called the terminal tool to finish. + * @param turn 1-based turn index. + * @param message the model's final answer. + */ + suspend fun onFinalAnswer(turn: Int, message: String) {} + } + + /** Why the loop stopped. */ + enum class StopReason { COMPLETED, MAX_ITERATIONS, REPEATED } + + /** + * Outcome of a run; [completed] is true when the model ended on its own. + * @property turns model turns executed. + * @property reason why the loop stopped. + */ + data class Result(val turns: Int, val reason: StopReason) { + val completed: Boolean get() = reason == StopReason.COMPLETED + } + + /** + * Runs the tool loop until the model stops calling tools or a limit is hit. + * @param history transcript, mutated in place; seed it with the user message. + * @param generate renders one model turn from a prompt. + * @param executeTools runs a batch of tool calls. + * @param events UI/state callbacks. + * @return the run [Result]. + */ + suspend fun run( + history: MutableList, + generate: suspend (prompt: String) -> String, + executeTools: suspend (List) -> List, + events: Events = object : Events {}, + ): Result { + var turn = 0 + var previousSignature: String? = null + var consecutiveRepeats = 0 + while (turn < maxIterations) { + turn++ + + val text = generate(renderTranscript(history)) + history.add(ChatMessage(Role.ASSISTANT, text)) + events.onModelTurn(turn, text) + + val calls = extractToolCalls(text) + if (calls.isEmpty()) { + return Result(turn, StopReason.COMPLETED) + } + + // Terminal tool alone ends the loop; if co-emitted with real tools, run those first. + val realCalls = terminalTool?.let { tt -> calls.filterNot { it.name == tt } } ?: calls + terminalTool?.let { tt -> + val terminal = calls.firstOrNull { it.name == tt } + if (terminal != null && realCalls.isEmpty()) { + events.onFinalAnswer(turn, terminal.args["message"]?.toString().orEmpty()) + return Result(turn, StopReason.COMPLETED) + } + } + + // Stop if the model repeats identical tool calls beyond the tolerated count. + val signature = signatureOf(realCalls) + if (signature == previousSignature) { + consecutiveRepeats++ + if (consecutiveRepeats >= maxConsecutiveRepeats) { + events.onRepeatedToolCalls(turn) + return Result(turn, StopReason.REPEATED) + } + } else { + consecutiveRepeats = 0 + } + previousSignature = signature + + val results = executeTools(realCalls) + events.onToolResults(turn, realCalls, results) + history.add(ChatMessage(Role.USER, formatToolResults(realCalls, results))) + } + + events.onMaxIterationsReached(turn) + return Result(turn, StopReason.MAX_ITERATIONS) + } + + /** + * Builds a stable, order-sensitive fingerprint of a tool-call batch (name + args). + * @param calls the batch to fingerprint. + * @return the fingerprint string. + */ + private fun signatureOf(calls: List): String = + calls.joinToString("|") { call -> + call.name + "(" + call.args.toSortedMap().entries.joinToString(",") { "${it.key}=${it.value}" } + ")" + } + + /** + * Flattens the transcript into one prompt string, with no trailing "Assistant:" cue + * (the backend appends its own; a doubled cue makes local models repeat). + * @param history the conversation so far. + * @return the rendered prompt. + */ + fun renderTranscript(history: List): String { + val sb = StringBuilder() + for ((index, message) in history.withIndex()) { + if (index > 0) sb.append("\n\n") + when (message.role) { + Role.ASSISTANT -> sb.append("Assistant: ").append(message.content) + else -> sb.append(message.content) + } + } + return sb.toString() + } + + /** + * Renders tool results for feeding back into the next prompt, capping each body. + * @param calls the tool calls that ran. + * @param results their results, positionally aligned with [calls]. + * @return the formatted results block. + */ + fun formatToolResults(calls: List, results: List): String { + val sb = StringBuilder("Tool results:\n") + results.forEachIndexed { index, result -> + val name = calls.getOrNull(index)?.name ?: "tool" + val body = if (result.success) { + buildString { + append(result.message) + result.data?.takeIf { it.isNotBlank() }?.let { append("\n").append(it) } + } + } else { + buildString { + append("FAILED: ").append(result.message) + result.error_details?.takeIf { it.isNotBlank() }?.let { append("\n").append(it) } + } + } + sb.append("[").append(name).append("] ").append(truncate(body)).append("\n\n") + } + sb.append( + "Base your reply strictly on the tool result(s) above — report only what they actually say; " + + "do not invent, assume, or contradict them. " + ) + if (results.isNotEmpty() && results.all { it.success }) { + sb.append( + "The action succeeded. If this satisfies the user's request, you are DONE — reply with the " + + "\"respond\" tool briefly confirming what happened. Do NOT call another tool unless the " + + "request clearly needs a further step." + ) + } else { + sb.append("If the task is complete, give the user your final answer. Otherwise, call the next tool.") + } + return sb.toString() + } + + private fun truncate(text: String): String = + if (text.length <= toolOutputCharLimit) text + else text.take(toolOutputCharLimit) + "\n…[truncated ${text.length - toolOutputCharLimit} chars]" +} diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/Executor.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/Executor.kt index 147435b0..f2516770 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/Executor.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/Executor.kt @@ -33,6 +33,7 @@ class Executor( fun requiredArgsForTool(toolName: String): List { return when (toolName) { "read_file" -> listOf("file_path") + "open_file" -> listOf("file_path") "list_files" -> emptyList() // directory is optional, defaults to "." "search_project" -> listOf("query") "create_file" -> listOf("file_path", "content") @@ -87,12 +88,14 @@ class Executor( return ToolResult.failure("Unknown function '$toolName'") } - // Normalize arg keys for read_file: if "path" is present but "file_path" is missing, remap + // Alias "path" → "file_path" for any tool that requires "file_path". val normalizedArgs = args.toMutableMap() - if (toolName == "read_file" && normalizedArgs.containsKey("path") && !normalizedArgs.containsKey("file_path")) { + if ("file_path" in requiredArgsForTool(toolName) && + normalizedArgs.containsKey("path") && !normalizedArgs.containsKey("file_path") + ) { normalizedArgs["path"]?.let { normalizedArgs["file_path"] = it } if (normalizedArgs.containsKey("file_path")) { - Log.d(TAG, "($executionMode): Remapped 'path' → 'file_path' for read_file tool") + Log.d(TAG, "($executionMode): Remapped 'path' → 'file_path' for $toolName tool") } } @@ -107,13 +110,7 @@ class Executor( return ToolResult.failure(message) } - for (key in handler.pathArgs) { - val raw = normalizedArgs[key]?.toString()?.trim() - if (!raw.isNullOrEmpty() && PathGuard.resolveWithin(raw) == null) { - Log.w(TAG, "($executionMode): '$toolName' arg '$key' escapes project root: $raw") - return ToolResult.failure("Path '$raw' is outside the project directory") - } - } + pathContainmentFailure(toolName, handler, normalizedArgs, executionMode)?.let { return it } // Check approval val approvalResponse = approvalManager.ensureApproved(toolName, handler, normalizedArgs) @@ -137,6 +134,32 @@ class Executor( Log.i(TAG, "($executionMode): Result: ${result.toResultMap()}") return result } + + /** + * Confines model-supplied path args to the project root; handlers that resolve + * paths themselves opt out via [ToolHandler.resolvesPathsInternally]. + * @param toolName the tool being dispatched (for logging). + * @param handler the tool's handler, source of [ToolHandler.pathArgs]. + * @param args the normalized call arguments. + * @param executionMode "Parallel"/"Sequential", for logging. + * @return a failure [ToolResult] for the first escaping arg, or null if all are safe. + */ + private fun pathContainmentFailure( + toolName: String, + handler: ToolHandler, + args: Map, + executionMode: String, + ): ToolResult? { + if (handler.resolvesPathsInternally) return null + for (key in handler.pathArgs) { + val raw = args[key]?.toString()?.trim() + if (raw.isNullOrEmpty()) continue + if (PathGuard.resolveWithin(raw) != null) continue + Log.w(TAG, "($executionMode): '$toolName' arg '$key' escapes project root: $raw") + return ToolResult.failure("Path '$raw' is outside the project directory") + } + return null + } } /** diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolCallExtractor.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolCallExtractor.kt index 1471ab8e..d8b7f6e9 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolCallExtractor.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolCallExtractor.kt @@ -7,7 +7,6 @@ import org.json.JSONObject * Extracts tool calls from LLM responses using multiple strategies: * 1. Explicit XML tags: {"tool":"name",...} * 2. JSON blocks: {"tool":"name",...} - * 3. Implicit actions: Detects from natural language patterns * * Works with both cloud (Gemini) and local LLMs. */ @@ -32,11 +31,6 @@ class ToolCallExtractor { toolCalls.addAll(extractFromJsonObjects(text)) } - // Strategy 3: Implicit actions from natural language - if (toolCalls.isEmpty()) { - toolCalls.addAll(extractImplicitActions(text)) - } - Log.d(TAG, "Extracted ${toolCalls.size} tool calls from response (${text.length} chars)") // Warn if we found incomplete tool calls @@ -78,18 +72,32 @@ class ToolCallExtractor { val toolCalls = mutableListOf() var found = 0 - // Find JSON objects with balanced braces containing "tool" field + // Find JSON objects with balanced braces containing "tool" field. var i = 0 while (i < text.length) { if (text[i] == '{') { - // Try to extract a balanced JSON object + // Extract a balanced object, ignoring braces inside string values. var braceCount = 0 var j = i var hasToolField = false + var inString = false + var escaped = false while (j < text.length) { - if (text[j] == '{') braceCount++ - else if (text[j] == '}') braceCount-- + val c = text[j] + if (inString) { + when { + escaped -> escaped = false + c == '\\' -> escaped = true + c == '"' -> inString = false + } + } else { + when (c) { + '"' -> inString = true + '{' -> braceCount++ + '}' -> braceCount-- + } + } // Check if this substring contains "tool" if (!hasToolField && text.substring(i, minOf(j + 1, text.length)).contains("\"tool\"")) { @@ -98,7 +106,7 @@ class ToolCallExtractor { j++ - if (braceCount == 0) { + if (!inString && braceCount == 0) { // Found complete object if (hasToolField) { val jsonStr = text.substring(i, j) @@ -123,56 +131,6 @@ class ToolCallExtractor { return toolCalls } - /** - * Strategy 3: Extract implicit tool calls from natural language. - * Detects action keywords and converts to tool calls. - */ - private fun extractImplicitActions(text: String): List { - val toolCalls = mutableListOf() - val lowerText = text.lowercase() - - Log.d(TAG, "Strategy 3 (Implicit actions): Analyzing for action keywords") - - // Patterns for list_files - if (matchesPattern(lowerText, listOf("list", "show"), listOf("file", "directory", "folder"))) { - val directory = extractDirectory(text) ?: "." - toolCalls.add(ToolCall("list_files", mapOf("directory" to directory))) - Log.d(TAG, "Detected: list_files action") - } - - // Patterns for read_file - if (matchesPattern(lowerText, listOf("read", "open", "view", "show"), listOf("file"))) { - val path = extractFilePath(text) - if (path != null) { - toolCalls.add(ToolCall("read_file", mapOf("path" to path))) - Log.d(TAG, "Detected: read_file action for $path") - } - } - - // Patterns for search_project - if (matchesPattern(lowerText, listOf("search", "find", "grep"), listOf("file", "code", "project"))) { - val query = extractSearchQuery(text) - if (query != null) { - toolCalls.add(ToolCall("search_project", mapOf("query" to query))) - Log.d(TAG, "Detected: search_project action for query: $query") - } - } - - // Patterns for create_file - if (matchesPattern(lowerText, listOf("create", "write", "make"), listOf("file"))) { - // Requires more context, generally not auto-triggered - Log.d(TAG, "Detected: create_file action (requires confirmation)") - } - - // Patterns for run_app - if (matchesPattern(lowerText, listOf("run", "launch", "build", "start"), listOf("app", "application"))) { - toolCalls.add(ToolCall("run_app", emptyMap())) - Log.d(TAG, "Detected: run_app action") - } - - return toolCalls - } - /** * Parse tool JSON and extract tool name and arguments. */ @@ -200,113 +158,5 @@ class ToolCallExtractor { null } } - - /** - * Check if text matches action pattern (verb + object). - */ - private fun matchesPattern(text: String, verbs: List, objects: List): Boolean { - val hasVerb = verbs.any { text.contains(it) } - val hasObject = objects.any { text.contains(it) } - return hasVerb && hasObject - } - - /** - * Extract directory path from natural language. - * Looks for explicit paths or common directory names. - */ - private fun extractDirectory(text: String): String? { - val lowerText = text.lowercase() - - // Check for common project directories - val commonDirs = mapOf( - "src" to "src", - "source" to "src", - "source code" to "src", - "main" to "src/main", - "java" to "src/main/java", - "kotlin" to "src/main/kotlin", - "resources" to "src/main/resources", - "test" to "src/test", - "root" to ".", - "project" to ".", - "current" to "." - ) - - for ((keyword, dir) in commonDirs) { - if (lowerText.contains(keyword)) { - Log.d(TAG, "Detected directory from keyword '$keyword': $dir") - return dir - } - } - - // Look for patterns like "in src", "in ./src", "in /path/to/dir", etc. - val patterns = listOf( - Regex("""(?:in|from)\s+(?:the\s+)?["`']?([/.\-\w]+)["`']?"""), - Regex("""directory\s+(?:of\s+)?["`']?([/.\-\w]+)["`']?"""), - Regex("""folder\s+["`']?([/.\-\w]+)["`']?"""), - Regex("""path\s+["`']?([/.\-\w]+)["`']?""") - ) - - for (pattern in patterns) { - val match = pattern.find(text) - if (match != null) { - val dir = match.groupValues[1] - if (dir.isNotEmpty() && dir.length > 1 && !dir.contains("the")) { - Log.d(TAG, "Extracted directory from pattern: $dir") - return dir - } - } - } - - // Default to current directory if no specific directory mentioned - Log.d(TAG, "No specific directory found, defaulting to current directory (.)") - return null // Will default to "." in ListFilesHandler - } - - /** - * Extract file path from natural language. - */ - private fun extractFilePath(text: String): String? { - // Look for patterns like "MainActivity.kt", "src/main/MainActivity.kt", etc. - val patterns = listOf( - Regex("""["`']([^"`'\s]+\.kt)["`']"""), - Regex("""file\s+["`']?([^"`'\s]+\.kt)["`']?"""), - Regex("""read\s+["`']?([^"`'\s]+)["`']?""") - ) - - for (pattern in patterns) { - val match = pattern.find(text) - if (match != null) { - val path = match.groupValues[1] - if (path.isNotEmpty() && !path.contains("the")) { - return path - } - } - } - - return null - } - - /** - * Extract search query from natural language. - */ - private fun extractSearchQuery(text: String): String? { - val patterns = listOf( - Regex("""(?:search|find|grep)\s+(?:for\s+)?["`']([^"`']+)["`']"""), - Regex("""search\s+(?:for\s+)?([^\.?!]+)""") - ) - - for (pattern in patterns) { - val match = pattern.find(text) - if (match != null) { - val query = match.groupValues[1].trim() - if (query.isNotEmpty() && query.length > 2) { - return query - } - } - } - - return null - } } } diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolHandler.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolHandler.kt index 6e85aef3..1c53676d 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolHandler.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolHandler.kt @@ -37,4 +37,7 @@ interface ToolHandler { */ val pathArgs: List get() = emptyList() + + val resolvesPathsInternally: Boolean + get() = false } diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ListFilesHandler.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ListFilesHandler.kt index 9e6e24b1..e83c518b 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ListFilesHandler.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ListFilesHandler.kt @@ -15,40 +15,29 @@ class ListFilesHandler( override val toolName = "list_files" override val description = "List files and directories in a given path" override val requiresApproval = false + override val pathArgs = listOf("directory") + // Resolved internally (below) to rescue slash-prefixed paths; opt out of the + // Executor pre-guard that would reject "/app/src" first. + override val resolvesPathsInternally = true override suspend fun execute(args: Map): ToolResult { - var directory = args["directory"]?.toString()?.trim()?.takeIf { it.isNotBlank() } - - // Get project root for containment check - val projectRoot = System.getProperty("project.dir") - ?: System.getProperty("user.dir") - ?: "/storage/emulated/0/AndroidIDEProjects" - val projectRootCanonical = File(projectRoot).canonicalPath - - // If no directory specified, use project root - if (directory.isNullOrBlank()) { - directory = projectRoot - Log.d("ListFilesHandler", "No directory specified, using project root: $directory") - } + val directory = args["directory"]?.toString()?.trim()?.takeIf { it.isNotBlank() } - Log.d("ListFilesHandler", "Listing files in directory: $directory") + Log.d("ListFilesHandler", "Listing files in directory: ${directory ?: ""}") return try { - // Resolve path against project root if relative - val dir = if (directory.startsWith("/")) { - File(directory).absoluteFile + // Resolve/containment-check via PathGuard; blank defaults to the root. + // A slash-prefixed relative dir ("/app/src") escapes the root, so retry + // it as relative — matching Read/OpenFileHandler. + val dir = if (directory == null) { + File(PathGuard.projectRoot()) } else { - File(projectRoot, directory).absoluteFile + PathGuard.resolveWithin(directory) + ?: PathGuard.resolveWithin(directory.removePrefix("/")) + ?: return ToolResult.failure("Directory path must be within project directory") } Log.d("ListFilesHandler", "Absolute path: ${dir.absolutePath}") - // Security: Verify directory is within project root - val dirCanonical = dir.canonicalPath - if (!dirCanonical.startsWith(projectRootCanonical + File.separator) && dirCanonical != projectRootCanonical) { - Log.e("ListFilesHandler", "Path escape attempt: $dirCanonical is outside project root $projectRootCanonical") - return ToolResult.failure("Directory path must be within project directory") - } - if (!dir.exists()) { Log.w("ListFilesHandler", "Directory does not exist: ${dir.absolutePath}") return ToolResult.failure("Directory does not exist: ${dir.absolutePath}") @@ -112,11 +101,4 @@ class ListFilesHandler( else -> "${bytes / (1024 * 1024 * 1024)} GB" } } - - private fun findDefaultDirectory(): String { - // Return project root only for security - return System.getProperty("project.dir") - ?: System.getProperty("user.dir") - ?: "/storage/emulated/0/AndroidIDEProjects" - } } diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/OpenFileHandler.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/OpenFileHandler.kt index 1ffaeb08..3aeffb12 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/OpenFileHandler.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/OpenFileHandler.kt @@ -5,17 +5,24 @@ import com.itsaky.androidide.plugins.PluginContext import com.itsaky.androidide.plugins.aiassistant.models.ToolResult import com.itsaky.androidide.plugins.aiassistant.tool.ToolHandler import com.itsaky.androidide.plugins.services.IdeEditorService +import kotlinx.coroutines.CoroutineDispatcher +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.withContext /** * Handler for opening files in the IDE editor. + * + * @param mainDispatcher dispatcher for editor-UI calls; overridden in unit tests. */ class OpenFileHandler( - private val pluginContext: PluginContext + private val pluginContext: PluginContext, + private val mainDispatcher: CoroutineDispatcher = Dispatchers.Main ) : ToolHandler { override val toolName = "open_file" override val description = "Open a file in the IDE editor" override val requiresApproval = false override val pathArgs = listOf("file_path") + override val resolvesPathsInternally = true override suspend fun execute(args: Map): ToolResult { val filePath = args["file_path"]?.toString()?.trim() @@ -26,14 +33,26 @@ class OpenFileHandler( Log.d("OpenFileHandler", "Opening file: $filePath") return try { - val file = PathGuard.resolveWithin(filePath) - ?: return ToolResult.failure("File path must be within project directory") - if (!file.exists()) { - Log.w("OpenFileHandler", "File does not exist: $filePath") - return ToolResult.failure( - "File not found", - "File does not exist: $filePath" - ) + val file = when (val resolution = PathGuard.resolve(filePath)) { + is PathGuard.Resolution.Resolved -> { + Log.d("OpenFileHandler", "Resolved '$filePath' -> ${resolution.file.path}") + resolution.file + } + is PathGuard.Resolution.Ambiguous -> { + val root = java.io.File(PathGuard.projectRoot()) + return ToolResult.failure( + "Multiple files named '${resolution.baseName}' — specify a path", + resolution.matches.joinToString("\n") { it.relativeToOrSelf(root).path } + ) + } + PathGuard.Resolution.Escaped -> { + Log.w("OpenFileHandler", "Path outside project and no match found: $filePath") + return ToolResult.failure("File path must be within project directory") + } + PathGuard.Resolution.NotFound -> { + Log.w("OpenFileHandler", "File does not exist: $filePath") + return ToolResult.failure("File not found", "File does not exist: $filePath") + } } if (!file.isFile) { @@ -53,7 +72,8 @@ class OpenFileHandler( ) } - val success = editorService.openFile(file) + // Opening a tab touches UI; execute() runs on Dispatchers.IO. + val success = withContext(mainDispatcher) { editorService.openFile(file) } if (success) { Log.d("OpenFileHandler", "File opened successfully: $filePath") ToolResult.success( @@ -69,10 +89,7 @@ class OpenFileHandler( } } catch (e: Exception) { Log.e("OpenFileHandler", "Error opening file", e) - ToolResult.failure( - "Error opening file", - "${e.message ?: "Unknown error"}\n\n${e.stackTraceToString()}" - ) + ToolResult.failure("Error opening file", e.message ?: "Unknown error") } } } diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/PathGuard.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/PathGuard.kt index 6b5c3179..c62c4b7b 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/PathGuard.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/PathGuard.kt @@ -1,57 +1,170 @@ package com.itsaky.androidide.plugins.aiassistant.tool.handlers import android.util.Log -import com.itsaky.androidide.plugins.aiassistant.AiAssistantPlugin -import com.itsaky.androidide.plugins.services.IdeProjectService import java.io.File +import java.nio.file.Files /** - * Shared path-containment guard for filesystem tool handlers. - * - * The LLM-driven tools resolve paths supplied by the model, so every handler - * that touches the filesystem must confine those paths to the current project - * root to prevent a prompt-injected model from reading or writing arbitrary - * files. Previously this check lived (duplicated) only in CreateFileHandler and - * ReadFileHandler; it is now centralized here and applied to every handler. + * Shared path-containment guard for filesystem tool handlers: confines every + * model-supplied path to the current project root. The root comes from + * [projectRootProvider] (host-wired), then a legacy system property, then [DEFAULT_ROOT]. */ object PathGuard { private const val TAG = "PathGuard" + private const val DEFAULT_ROOT = "/storage/emulated/0/AndroidIDEProjects" + + /** Host-backed supplier of the current project root; queried on every resolution. */ + @Volatile + private var projectRootProvider: (() -> String?)? = null + + /** Test-only override; when non-null it wins over everything else. */ + @Volatile + private var projectRootOverride: String? = null + + /** + * Installs (or clears, with `null`) the host-backed project-root supplier. + * @param provider supplier of the project root, or null to clear. + */ + fun setProjectRootProvider(provider: (() -> String?)?) { + projectRootProvider = provider + } /** - * Resolve the active project root, preferring the IDE's current project. - * @return absolute project-root path, or null when none can be determined. + * Forces a specific root, or clears it with `null`. Tests only. + * @param root the root to force, or null to clear. */ - fun projectRoot(): String? { - // Prefer the IDE-provided current project root when the service is reachable. - val ideRoot = runCatching { - val service: IdeProjectService? = - AiAssistantPlugin.getContext()?.services?.get(IdeProjectService::class.java) - service?.getCurrentProject()?.rootDir?.absolutePath - }.getOrNull() - if (ideRoot != null) return ideRoot - // Fall back to JVM system properties; no hardcoded root, so an unknown root fails safe. - return System.getProperty("project.dir") ?: System.getProperty("user.dir") + fun setProjectRootForTesting(root: String?) { + projectRootOverride = root } /** - * Resolve [path] against the project root and verify it stays inside it. - * @param path model-supplied path, absolute or project-relative. - * @return contained [File], or null when the root is unknown or [path] escapes. + * Resolves the active project root (best-effort); `user.dir` is never consulted + * (it is "/" on Android) so [isValidRoot] can reject it and the guard fails closed. + * @return the resolved root path (not guaranteed valid). + */ + fun projectRoot(): String = + projectRootOverride + ?: projectRootProvider?.invoke()?.takeIf { it.isNotBlank() } + ?: System.getProperty("project.dir") + ?: DEFAULT_ROOT + + /** + * A usable root is a non-blank, existing directory that isn't the filesystem root "/" + * ("/" means no project is open, and would confine nothing). + * @param root candidate root. + * @return true if [root] is a usable project root. + */ + private fun isValidRoot(root: File): Boolean { + val canonical = try { root.canonicalPath } catch (e: Exception) { return false } + return canonical.isNotBlank() && canonical != File.separator && root.isDirectory + } + + /** + * Resolves [path] against the project root, verifying the canonical result is inside it. + * @param path a relative or absolute path from the model. + * @return the resolved [File] when contained, or null when it escapes or there is no valid root. */ fun resolveWithin(path: String): File? { - val rootPath = projectRoot() - if (rootPath == null) { - Log.e(TAG, "Denying path access: project root could not be determined") + val rootFile = File(projectRoot()) + if (!isValidRoot(rootFile)) { + Log.e(TAG, "No valid project root ('${rootFile.path}'); rejecting path: $path") return null } - val root = File(rootPath).canonicalPath - val file = if (path.startsWith("/")) File(path) else File(rootPath, path) + val root = rootFile.canonicalPath + val file = if (path.startsWith("/")) File(path) else File(rootFile, path) val canonical = file.canonicalPath - val contained = canonical == root || canonical.startsWith(root + File.separator) + + val rootWithSep = if (root.endsWith(File.separator)) root else root + File.separator + val contained = canonical == root || canonical.startsWith(rootWithSep) if (!contained) { Log.e(TAG, "Path escape attempt: $canonical is outside project root $root") } return if (contained) file else null } + + /** Directories skipped when searching for a file by name — large/generated/noise. */ + private val SKIP_DIRS = setOf("build", ".git", ".gradle", ".idea", "node_modules", ".cxx") + + /** + * Finds in-root files whose name equals [fileName] (case-insensitive), skipping + * generated/hidden dirs and symlinks, so a bare name resolves to a real path. + * @param fileName the name to match (its basename is used). + * @param limit max results. + * @return matching files (possibly empty); more than one means ambiguous. + */ + fun findByName(fileName: String, limit: Int = 20): List { + val target = baseNameOf(fileName).trim() + if (target.isEmpty()) return emptyList() + + val root = File(projectRoot()) + if (!isValidRoot(root)) return emptyList() + val rootWithSep = root.canonicalPath.let { if (it.endsWith(File.separator)) it else it + File.separator } + + return root.walkTopDown() + // Don't descend symlinked dirs; walkTopDown matches by name and could escape the root. + .onEnter { dir -> + (dir == root || (dir.name !in SKIP_DIRS && !dir.name.startsWith("."))) && + (dir == root || !Files.isSymbolicLink(dir.toPath())) + } + .filter { it.isFile && it.name.equals(target, ignoreCase = true) } + // Re-verify containment so a symlinked file resolving outside the root is dropped. + .filter { it.canonicalPath.startsWith(rootWithSep) } + .take(limit) + .toList() + } + + /** + * Returns the last path segment, tolerating both '/' and '\' separators. + * @param path the path. + * @return the basename. + */ + private fun baseNameOf(path: String): String = + path.substringAfterLast('/').substringAfterLast('\\') + + /** Outcome of [resolve]; callers match on it to act or to explain the miss. */ + sealed interface Resolution { + /** + * The path resolved to an existing in-root entry. + * @property file the resolved file or directory. + */ + data class Resolved(val file: File) : Resolution + + /** + * A bare name matched several files; ask the user to disambiguate. + * @property baseName the searched name. + * @property matches the candidate files. + */ + data class Ambiguous(val baseName: String, val matches: List) : Resolution + + /** No existing match, but the path stayed in-root — it simply doesn't exist. */ + object NotFound : Resolution + + /** The path had no in-root interpretation at all (a containment escape). */ + object Escaped : Resolution + } + + /** + * Resolves a model-supplied path to a project entry, enforcing root containment. + * Tries in-root, then slash-stripped relative, then a project-wide basename search. + * @param path a relative, absolute, or bare file name from the model. + * @return the matching [Resolution]. + */ + fun resolve(path: String): Resolution { + val primary = resolveWithin(path) + val relative = if (path.startsWith("/")) resolveWithin(path.removePrefix("/")) else null + + (primary?.takeIf { it.exists() } ?: relative?.takeIf { it.exists() })?.let { + return Resolution.Resolved(it) + } + + val matches = findByName(path) + when (matches.size) { + 1 -> return Resolution.Resolved(matches[0]) + 0 -> {} + else -> return Resolution.Ambiguous(baseNameOf(path), matches) + } + + return if (primary == null && relative == null) Resolution.Escaped else Resolution.NotFound + } } diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ReadFileHandler.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ReadFileHandler.kt index b4502d99..11ed7947 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ReadFileHandler.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ReadFileHandler.kt @@ -15,6 +15,7 @@ class ReadFileHandler( override val description = "Read the contents of a file" override val requiresApproval = false override val pathArgs = listOf("file_path", "path") + override val resolvesPathsInternally = true override suspend fun execute(args: Map): ToolResult { // Accept both "file_path" (standardized) and "path" (legacy LLM responses) @@ -24,26 +25,35 @@ class ReadFileHandler( } return try { - // Security: resolve against the project root and reject any escape. - val file = PathGuard.resolveWithin(filePath) - ?: return ToolResult.failure("File path must be within project directory") + val file = when (val resolution = PathGuard.resolve(filePath)) { + is PathGuard.Resolution.Resolved -> resolution.file + is PathGuard.Resolution.Ambiguous -> { + val root = java.io.File(PathGuard.projectRoot()) + return ToolResult.failure( + "Multiple files named '${resolution.baseName}' — specify a path", + resolution.matches.joinToString("\n") { it.relativeToOrSelf(root).path } + ) + } + PathGuard.Resolution.Escaped -> + return ToolResult.failure("File path must be within project directory") + PathGuard.Resolution.NotFound -> + return ToolResult.failure("File does not exist: $filePath") + } - if (!file.exists()) { - ToolResult.failure("File does not exist: $filePath") - } else if (!file.isFile) { - ToolResult.failure("Path is not a file: $filePath") - } else if (!file.canRead()) { - ToolResult.failure("Cannot read file: $filePath") - } else { - val content = file.readText() - ToolResult.success( - message = "Read ${content.length} characters from $filePath", - data = content - ) + when { + !file.isFile -> ToolResult.failure("Path is not a file: $filePath") + !file.canRead() -> ToolResult.failure("Cannot read file: $filePath") + else -> { + val content = file.readText() + ToolResult.success( + message = "Read ${content.length} characters from $filePath", + data = content + ) + } } } catch (e: Exception) { Log.e("ReadFileHandler", "Error reading file", e) - ToolResult.failure("Error reading file: ${e.message}", e.stackTraceToString()) + ToolResult.failure("Error reading file: ${e.message}") } } } diff --git a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/ChatViewModel.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/ChatViewModel.kt index b057946c..00c0be31 100644 --- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/ChatViewModel.kt +++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/viewmodel/ChatViewModel.kt @@ -3,11 +3,14 @@ package com.itsaky.androidide.plugins.aiassistant.viewmodel import androidx.lifecycle.ViewModel import androidx.lifecycle.viewModelScope import com.itsaky.androidide.plugins.PluginContext +import com.itsaky.androidide.plugins.aiassistant.R import com.itsaky.androidide.plugins.aiassistant.models.AgentState import com.itsaky.androidide.plugins.aiassistant.models.ChatMessage import com.itsaky.androidide.plugins.aiassistant.models.ChatSession import com.itsaky.androidide.plugins.aiassistant.models.MessageStatus import com.itsaky.androidide.plugins.aiassistant.models.Sender +import com.itsaky.androidide.plugins.aiassistant.models.ToolResult +import com.itsaky.androidide.plugins.aiassistant.tool.AgentLoop import com.itsaky.androidide.plugins.aiassistant.tool.Executor import com.itsaky.androidide.plugins.aiassistant.tool.ToolApprovalManager import com.itsaky.androidide.plugins.aiassistant.tool.ToolCall @@ -29,6 +32,8 @@ import com.itsaky.androidide.plugins.aiassistant.data.ChatStorageManager import com.itsaky.androidide.plugins.aiassistant.utils.ToolExecutionTracker import com.itsaky.androidide.plugins.services.LlmInferenceService import com.itsaky.androidide.plugins.services.SharedServices +import kotlinx.coroutines.CancellationException +import kotlinx.coroutines.CompletableDeferred import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.Job import kotlinx.coroutines.delay @@ -40,6 +45,7 @@ import kotlinx.coroutines.flow.combine import kotlinx.coroutines.flow.stateIn import kotlinx.coroutines.isActive import kotlinx.coroutines.launch +import kotlinx.coroutines.withContext import org.json.JSONObject import java.io.File import java.util.UUID @@ -51,6 +57,37 @@ class ChatViewModel( private val getContext: () -> PluginContext? ) : ViewModel() { + companion object { + /** Terminal tool: shared by [agentLoop] (stops on it) and [runModelTurn] (renders its message). */ + const val RESPOND_TOOL = "respond" + + /** [LlmConfig.extraParams] key for the local-backend GBNF; must match ai-core's `LocalLlmBackend.EXTRA_PARAM_GRAMMAR`. */ + private const val EXTRA_PARAM_GRAMMAR = "grammar" + } + + /** + * Builds the local-backend GBNF forcing one well-formed ``. The `tool` + * alternatives come from the registered handlers (+ [RESPOND_TOOL]) so the token + * mask can't drift; control chars are excluded so `org.json` accepts the strings. + * @return the GBNF grammar string. + */ + private fun buildLocalToolCallGrammar(): String { + val toolNames = (toolRouter.getAllHandlers().map { it.toolName } + RESPOND_TOOL).distinct() + val toolAlternatives = toolNames.joinToString(" | ") { "\"~$it~\"" } + return """ + root ::= "{~tool~:" tool ",~args~:" args "}" + tool ::= $toolAlternatives + args ::= "{}" | "{" pair ("," pair)* "}" + pair ::= "~" key "~:~" val "~" + key ::= [a-z_]+ + val ::= char* + char ::= [^~\\\x00-\x1F] | "\\" [~\\/bfnrt] + """.trimIndent().replace("~", "\\\"") + } + + /** Local-backend GBNF, built once from the registered handlers. */ + private val localToolCallGrammar: String by lazy { buildLocalToolCallGrammar() } + private fun getLlmService(): LlmInferenceService? { return try { SharedServices.get(LlmInferenceService::class.java) @@ -112,8 +149,20 @@ class ChatViewModel( private val approvalManager = ToolApprovalManager() private val toolRouter: ToolRouter private val executor: Executor + private val agentLoop = AgentLoop(terminalTool = RESPOND_TOOL) val toolExecutionTracker = ToolExecutionTracker() + /** The in-flight agent run (streaming + tool loop), so it can be cancelled. */ + private var generationJob: Job? = null + private val generationEpoch = java.util.concurrent.atomic.AtomicInteger(0) + + /** True while a generation is admitted and its coroutine has not yet unwound; gates re-entry. */ + private val isGenerating = java.util.concurrent.atomic.AtomicBoolean(false) + + /** Whether the current run's most recent tool batch failed; reset per run. */ + @Volatile + private var lastToolFailedThisRun = false + private val _pendingApprovalRequest = MutableStateFlow(null) val pendingApprovalRequest: StateFlow = _pendingApprovalRequest.asStateFlow() @@ -301,6 +350,25 @@ class ChatViewModel( - After each file modification, verify the build compiles - Generate apps that actually run and work as described + RULES: + - Never fabricate tool output. Emit a tool call, then wait for the real result before continuing. + - Never write "User:", "Assistant:", or "Tool results:" — the system supplies those. + - Paths are relative to the project root and must be complete. If you don't know a file's exact path, find it with search_project or list_files first, then act on the real path — don't guess. + - For plain chat (e.g. "Hi"), just reply briefly with no tool call. When the task is done, give a short summary with no tool call. + + TOOL CALL FORMAT — to run a tool, emit a single line in EXACTLY this format and nothing after it: + {"tool":"TOOL_NAME","args":{"arg":"value"}} + Do NOT describe the action in prose (e.g. "Okay, I'll open the file…") — narrating does nothing. + The tool only runs when you emit the line itself. + + FORMAT EXAMPLES (the tool call is the entire reply): + Open a file once you know its path: + {"tool":"open_file","args":{"file_path":"app/src/main/java/com/example/app/MainActivity.java"}} + Find a file by name: + {"tool":"search_project","args":{"query":"MainActivity"}} + List a directory: + {"tool":"list_files","args":{"directory":"app/src/main"}} + WORKFLOW: 1. Understand the user's request 2. List files to understand the project structure @@ -309,8 +377,6 @@ class ChatViewModel( 5. Sync gradle and verify compilation 6. Run the app to confirm it works 7. Report success and what was built - - You have full access to tools - use them continuously throughout the workflow. """.trimIndent() android.util.Log.d("ChatViewModel", "Using Gemini system prompt (high autonomy mode) with ${toolRouter.getAllHandlers().size} tools") @@ -326,38 +392,24 @@ class ChatViewModel( } val prompt = """ - You are a helpful coding assistant integrated into AndroidIDE. Help the user build Android apps step-by-step. + You are a coding assistant inside AndroidIDE. - CRITICAL: You MUST use tools for ANY action-related request. Do NOT just describe what you would do. + Rules: + - Reply with exactly ONE tool call, nothing else. + - Use a file/project tool only when the user asks about files, code, or the project; for a greeting, small talk, or a question you can answer, use "respond". + - Never invent tool output or claim an action you didn't perform via a tool. After a tool call, stop; the real result returns next turn. + - "respond" must carry a "message" — your reply or final answer. + - File arguments accept a bare name (e.g. "MainActivity.java"); the project is searched. Don't invent deep paths. - AVAILABLE TOOLS: + Tools: $toolDescriptions + - respond: Send the user a message or your final answer. - TOOL CALLING RULES: - 1. When user asks to perform an action (list, read, search, create, update, run), IMMEDIATELY call the tool - 2. Do NOT explain or apologize - just execute the tool call - 3. Always provide the tool call in this EXACT format: - {"tool":"TOOL_NAME","args":{"param1":"value1"}} - 4. Execute tools BEFORE saying anything else - - STEP-BY-STEP WORKFLOW: - 1. List files to understand the project - 2. Read existing files to know what to change - 3. Create or update one file at a time - 4. After each file, ask the user what to do next - 5. Add dependencies when needed - 6. Sync gradle to check for errors - 7. Run the app to test it - 8. Ask for feedback and iterate - - EXAMPLES: - User: "list files in src" - {"tool":"list_files","args":{"directory":"src"}} - - User: "read MainActivity.kt" - {"tool":"read_file","args":{"file_path":"MainActivity.kt"}} - - After each tool call, analyze the result and ask: "What would you like to do next?" + Examples (pick the tool that matches; don't copy verbatim): + Greeting / question you can answer -> respond: + {"tool":"respond","args":{"message":"Hi! What would you like to build?"}} + Open a file -> open_file: + {"tool":"open_file","args":{"file_path":"MainActivity.java"}} """.trimIndent() android.util.Log.d("ChatViewModel", "Using Local LLM system prompt (guided mode) with ${toolRouter.getAllHandlers().size} tools") @@ -365,51 +417,51 @@ class ChatViewModel( } /** - * Parse tool calls from text using multi-strategy extraction. - * Tries: XML tags → JSON objects → Implicit actions + * Executes a batch of tool calls, renders each result as a TOOL message, and + * returns the results for the [agentLoop] to feed back; leaves [AgentState.Idle] + * to the loop. + * @param toolCalls the calls to execute. + * @return the results, positionally aligned with [toolCalls]. */ - private fun parseToolCalls(text: String): List { - return ToolCallExtractor.extractToolCalls(text) - } - - /** - * Execute tool calls and add results to chat. - */ - private suspend fun executeToolCalls(toolCalls: List) { - if (toolCalls.isEmpty()) return + private suspend fun executeToolCalls(toolCalls: List): List { + if (toolCalls.isEmpty()) return emptyList() val executingState = AgentState.Executing( currentStepIndex = 0, totalSteps = toolCalls.size, description = toolCalls.first().name ) - _agentState.value = executingState + withContext(Dispatchers.Main) { _agentState.value = executingState } startStateTimer(executingState) val results = executor.execute(toolCalls) + // Record whether this batch's last tool failed (read by runModelTurn). + lastToolFailedThisRun = results.lastOrNull()?.success == false + // Add tool results as messages - results.forEachIndexed { index, result -> - val toolCall = toolCalls[index] - val resultText = if (result.success) { - "${toolCall.name}: ${result.message}\n${result.data ?: ""}" - } else { - "${toolCall.name} failed: ${result.message}\n${result.error_details ?: ""}" + withContext(Dispatchers.Main) { + results.forEachIndexed { index, result -> + val toolCall = toolCalls[index] + val resultText = if (result.success) { + "${toolCall.name}: ${result.message}\n${result.data ?: ""}" + } else { + "${toolCall.name} failed: ${result.message}\n${result.error_details ?: ""}" + } + val resultMessage = ChatMessage( + id = UUID.randomUUID().toString(), + text = resultText, + sender = Sender.TOOL, + status = if (result.success) MessageStatus.SENT else MessageStatus.ERROR + ) + android.util.Log.d("ChatViewModel", "Adding tool result message: $resultText") + _messages.value = _messages.value + resultMessage + syncMessageToSession(resultMessage) } - val resultMessage = ChatMessage( - id = UUID.randomUUID().toString(), - text = resultText, - sender = Sender.TOOL, - status = if (result.success) MessageStatus.SENT else MessageStatus.ERROR - ) - android.util.Log.d("ChatViewModel", "Adding tool result message: $resultText") - _messages.value = _messages.value + resultMessage - android.util.Log.d("ChatViewModel", "Total messages after tool result: ${_messages.value.size}") - syncMessageToSession(resultMessage) } stopStateTimer() - _agentState.value = AgentState.Idle + return results } /** @@ -509,168 +561,264 @@ class ChatViewModel( return } + // Reject re-entry while a generation is still in flight. + if (!isGenerating.compareAndSet(false, true)) { + android.util.Log.d("ChatViewModel", "sendMessage: generation already in progress; ignoring") + return + } + android.util.Log.d("ChatViewModel", "sendMessage: Starting message processing") - viewModelScope.launch(Dispatchers.IO) { + // Reset per-run tool-failure tracking. + lastToolFailedThisRun = false + val epoch = generationEpoch.incrementAndGet() + generationJob = viewModelScope.launch(Dispatchers.IO) { try { - // Add user message + // Add user message to the UI. val userChatMessage = ChatMessage( id = UUID.randomUUID().toString(), text = userMessage, sender = Sender.USER, status = MessageStatus.SENT ) - _messages.value = _messages.value + userChatMessage - android.util.Log.d("ChatViewModel", "sendMessage: Added user message, total messages=${_messages.value.size}") - syncMessageToSession(userChatMessage) - - // Add empty agent message that will be updated with streaming tokens - val agentMessageId = UUID.randomUUID().toString() - val agentMessage = ChatMessage( - id = agentMessageId, - text = "", - sender = Sender.AGENT, - status = MessageStatus.SENT // SENT so text is visible immediately - ) - _messages.value = _messages.value + agentMessage - syncMessageToSession(agentMessage) - - // Set processing state - _agentState.value = AgentState.Processing("Generating...") - - // Record start time - val startTime = System.currentTimeMillis() + withContext(Dispatchers.Main) { + _messages.value = _messages.value + userChatMessage + syncMessageToSession(userChatMessage) + _agentState.value = AgentState.Processing(str(R.string.msg_generating)) + } - // Create LLM config val config = LlmInferenceService.LlmConfig(currentBackendId).apply { temperature = 0.7f - maxTokens = 4096 // Increased from 2048 to ensure complete tool calls are generated + maxTokens = 4096 // headroom for complete tool calls systemPrompt = buildSystemPrompt() + // Local backend constrains generation to this grammar; cloud ignores it. + extraParams = mapOf(EXTRA_PARAM_GRAMMAR to localToolCallGrammar) } - // Build message with context if any files are selected val messageWithContext = buildString { append(userMessage) - val context = buildContextString() - if (context.isNotEmpty()) { - append(context) - } + append(buildContextString()) } - - // Add user message to conversation history for LLM context - val userHistoryMessage = LlmInferenceService.ChatMessage( - LlmInferenceService.ChatMessage.Role.USER, - messageWithContext + val history = _history.value.toMutableList() + history.add( + LlmInferenceService.ChatMessage( + LlmInferenceService.ChatMessage.Role.USER, + messageWithContext + ) ) - _history.value = _history.value + userHistoryMessage - android.util.Log.d("ChatViewModel", "Added user to history. Total history length: ${_history.value.size}") - - // Accumulated response text - val responseBuilder = StringBuilder() - - // Use streaming API with callback - llmService.generateStreaming(messageWithContext, config, object : LlmInferenceService.StreamCallback { - override fun onToken(token: String) { - viewModelScope.launch(Dispatchers.Main) { - // Accumulate token - responseBuilder.append(token) - - // Update the message with new text using map() to trigger DiffUtil - val updatedMessage = ChatMessage( - id = agentMessageId, - text = responseBuilder.toString(), - sender = Sender.AGENT, - status = MessageStatus.SENT - ) - _messages.value = _messages.value.map { msg -> - if (msg.id == agentMessageId) { - updatedMessage - } else { - msg - } - } - - // Also update current session's message - syncMessageToSession(updatedMessage) - } - } - override fun onComplete(response: LlmInferenceService.LlmResponse) { - viewModelScope.launch(Dispatchers.IO) { - val durationMs = System.currentTimeMillis() - startTime - val finalText = response.text - - // Mark message as completed with final text - launch(Dispatchers.Main) { - val updatedMessage = ChatMessage( - id = agentMessageId, - text = finalText, - sender = Sender.AGENT, - status = MessageStatus.COMPLETED, - durationMs = durationMs + try { + agentLoop.run( + history = history, + generate = { prompt -> + withContext(Dispatchers.Main) { + _agentState.value = AgentState.Processing(str(R.string.msg_generating)) + } + runModelTurn(llmService, prompt, config, epoch) + }, + executeTools = { calls -> executeToolCalls(calls) }, + events = object : AgentLoop.Events { + override suspend fun onMaxIterationsReached(turns: Int) { + addSystemMessage( + str(R.string.agent_max_steps_reached, turns), + MessageStatus.SENT ) - _messages.value = _messages.value.map { msg -> - if (msg.id == agentMessageId) { - updatedMessage - } else { - msg - } - } - - syncMessageToSession(updatedMessage) } - // Add assistant response to conversation history for LLM context - val assistantHistoryMessage = LlmInferenceService.ChatMessage( - LlmInferenceService.ChatMessage.Role.ASSISTANT, - finalText - ) - _history.value = _history.value + assistantHistoryMessage - android.util.Log.d("ChatViewModel", "Added assistant to history. Total history length: ${_history.value.size}") - - // Parse and execute tool calls if any - val toolCalls = parseToolCalls(finalText) - if (toolCalls.isNotEmpty()) { - executeToolCalls(toolCalls) - } else { - launch(Dispatchers.Main) { - _agentState.value = AgentState.Idle - } + override suspend fun onRepeatedToolCalls(turns: Int) { + addSystemMessage( + str(R.string.agent_repeated_calls), + MessageStatus.SENT + ) } } + ) + } finally { + // Persist history only if this run wasn't superseded (epoch bumped). + if (generationEpoch.get() == epoch) { + _history.value = history.toList() } + stopStateTimer() + } - override fun onError(error: String) { - viewModelScope.launch(Dispatchers.Main) { - val durationMs = System.currentTimeMillis() - startTime - - // Remove the agent message - _messages.value = _messages.value.filter { it.id != agentMessageId } - - // Add error message - _agentState.value = AgentState.Error(error) - val errorMessage = ChatMessage( - id = UUID.randomUUID().toString(), - text = error, - sender = Sender.SYSTEM, - status = MessageStatus.ERROR, - durationMs = durationMs - ) - _messages.value = _messages.value + errorMessage - syncMessageToSession(errorMessage) + withContext(Dispatchers.Main) { _agentState.value = AgentState.Idle } + } catch (ce: CancellationException) { + stopStateTimer() + throw ce + } catch (e: Exception) { + android.util.Log.e("ChatViewModel", "sendMessage failed", e) + stopStateTimer() + _agentState.value = AgentState.Error(str(R.string.state_error, e.message)) + addSystemMessage(str(R.string.state_error, e.message), MessageStatus.ERROR) + } finally { + // Allow re-entry once the coroutine unwinds. + isGenerating.set(false) + } + } + } + + /** + * Runs one streaming model turn: creates an agent bubble, streams tokens into it, + * and suspends until completion. Throws on backend error. + * @param llmService the inference service. + * @param prompt the rendered prompt for this turn. + * @param config the generation config. + * @param epoch this run's epoch, for staleness checks against Stop/newer sends. + * @return the final response text (raw, for tool-call extraction). + */ + private suspend fun runModelTurn( + llmService: LlmInferenceService, + prompt: String, + config: LlmInferenceService.LlmConfig, + epoch: Int + ): String { + val deferred = CompletableDeferred() + val agentMessageId = UUID.randomUUID().toString() + val startTime = System.currentTimeMillis() + val responseBuilder = StringBuilder() + + // True once Stop (or a newer message) has superseded this generation. + fun isStale() = generationEpoch.get() != epoch + + withContext(Dispatchers.Main) { + val agentMessage = ChatMessage( + id = agentMessageId, + text = "", + sender = Sender.AGENT, + status = MessageStatus.SENT + ) + _messages.value = _messages.value + agentMessage + syncMessageToSession(agentMessage) + } + + try { + llmService.generateStreaming(prompt, config, object : LlmInferenceService.StreamCallback { + override fun onToken(token: String) { + if (isStale()) return // Stop pressed — ignore late tokens. + responseBuilder.append(token) + // Snapshot on the producer thread; only the immutable String crosses to Main. + val snapshot = responseBuilder.toString() + viewModelScope.launch(Dispatchers.Main) { + if (isStale()) return@launch + val updated = ChatMessage( + id = agentMessageId, + text = snapshot, + sender = Sender.AGENT, + status = MessageStatus.SENT + ) + _messages.value = _messages.value.map { if (it.id == agentMessageId) updated else it } + syncMessageToSession(updated) + } + } + + override fun onComplete(response: LlmInferenceService.LlmResponse) { + if (isStale()) { + // Already cancelled; the awaiting loop was unblocked by job cancel. + deferred.complete(response.text) + return + } + val durationMs = System.currentTimeMillis() - startTime + val toolCalls = ToolCallExtractor.extractToolCalls(response.text) + val respondCall = toolCalls.firstOrNull { it.name == RESPOND_TOOL } + val respondMessage = respondCall?.args?.get("message")?.toString() + + // Per-run flag (set by executeToolCalls), not a session-wide scan. + val lastToolFailed = lastToolFailedThisRun + + val displayText = when { + respondCall != null && lastToolFailed -> + str(R.string.agent_action_failed) + // Render the "respond" message to the user, not a tool badge. + respondCall != null -> + respondMessage?.takeIf { it.isNotBlank() } ?: str(R.string.agent_no_response) + toolCalls.isNotEmpty() -> toolCalls.joinToString("\n") { c -> + "🔧 ${c.name}(${c.args.entries.joinToString(", ") { "${it.key}=${it.value}" }})" + } + else -> response.text.ifBlank { + str(R.string.agent_no_response) } } - }) + viewModelScope.launch(Dispatchers.Main) { + if (isStale()) return@launch + val finalMsg = ChatMessage( + id = agentMessageId, + text = displayText, + sender = Sender.AGENT, + status = MessageStatus.COMPLETED, + durationMs = durationMs + ) + _messages.value = _messages.value.map { if (it.id == agentMessageId) finalMsg else it } + syncMessageToSession(finalMsg) + } + // Return the RAW text to the loop so extraction/stop logic is unaffected. + deferred.complete(response.text) + } - } catch (e: Exception) { - _agentState.value = AgentState.Error("Error: ${e.message}") - val errorMessage = ChatMessage( - id = UUID.randomUUID().toString(), - text = "Error: ${e.message}", - sender = Sender.SYSTEM, - status = MessageStatus.ERROR - ) - _messages.value = _messages.value + errorMessage + override fun onError(error: String) { + if (isStale()) { + deferred.completeExceptionally(CancellationException("stopped")) + return + } + viewModelScope.launch(Dispatchers.Main) { + // Drop the empty/partial bubble; the error surfaces as a SYSTEM message. + _messages.value = _messages.value.filter { it.id != agentMessageId } + } + deferred.completeExceptionally(RuntimeException(error)) + } + }) + } catch (e: Exception) { + // A synchronous throw fires no callback; complete deferred so await() doesn't hang. + android.util.Log.e("ChatViewModel", "generateStreaming threw synchronously", e) + viewModelScope.launch(Dispatchers.Main) { + _messages.value = _messages.value.filter { it.id != agentMessageId } } + if (!deferred.isCompleted) deferred.completeExceptionally(e) + } + + return deferred.await() + } + + /** + * Appends an AGENT message to the chat (terminal state, no streaming dots). + * @param text the message text. + */ + private suspend fun addAgentMessage(text: String) { + val message = ChatMessage( + id = UUID.randomUUID().toString(), + text = text, + sender = Sender.AGENT, + status = MessageStatus.COMPLETED, + durationMs = 0L + ) + withContext(Dispatchers.Main) { + _messages.value = _messages.value + message + syncMessageToSession(message) + } + } + + /** + * Resolves a UI string resource via the plugin's Android context. + * @param resId the string resource id. + * @param args format arguments. + * @return the resolved string, or empty if the context is gone. + */ + private fun str(resId: Int, vararg args: Any?): String = + getContext()?.androidContext?.getString(resId, *args).orEmpty() + + /** + * Appends a SYSTEM message to the chat (on the main thread). + * @param text the message text. + * @param status the message status. + */ + private suspend fun addSystemMessage(text: String, status: MessageStatus) { + val message = ChatMessage( + id = UUID.randomUUID().toString(), + text = text, + sender = Sender.SYSTEM, + status = status + ) + withContext(Dispatchers.Main) { + _messages.value = _messages.value + message + syncMessageToSession(message) } } @@ -678,7 +826,14 @@ class ChatViewModel( * Clear all messages from the conversation. */ fun clearMessages() { + // Clear Chat must also stop any in-flight run, not just wipe the list. + generationEpoch.incrementAndGet() + generationJob?.cancel() + generationJob = null + getLlmService()?.cancelGeneration() + stopStateTimer() _messages.value = emptyList() + _history.value = emptyList() _agentState.value = AgentState.Idle } @@ -690,6 +845,7 @@ class ChatViewModel( _sessions.value = _sessions.value + newSession _currentSessionId.value = newSession.id _messages.value = emptyList() + _history.value = emptyList() } /** @@ -702,6 +858,7 @@ class ChatViewModel( _currentSessionId.value = sessionId // Use immutable snapshot to ensure StateFlow emits on mutations _messages.value = session.messages.toList() + _history.value = emptyList() android.util.Log.d("ChatViewModel", "switchToSession: set _messages to ${session.messages.size} messages") } } @@ -722,10 +879,38 @@ class ChatViewModel( * Stop any ongoing processing. */ fun stopProcessing() { - if (_agentState.value is AgentState.Processing) { - _agentState.value = AgentState.Cancelling - getLlmService()?.cancelGeneration() - _agentState.value = AgentState.Idle + generationEpoch.incrementAndGet() + _agentState.value = AgentState.Cancelling + generationJob?.cancel() + generationJob = null + getLlmService()?.cancelGeneration() + stopStateTimer() + finalizeInProgressMessages() + _agentState.value = AgentState.Idle + } + + /** + * Give any still-streaming agent bubble (status SENT, null `durationMs`) a + * terminal state so its animated "…" dots stop: drop empty bubbles, mark + * partial ones [MessageStatus.COMPLETED]. Called on Stop. + */ + private fun finalizeInProgressMessages() { + val finalized = _messages.value.mapNotNull { msg -> + if (msg.sender == Sender.AGENT && msg.durationMs == null) { + if (msg.text.isBlank()) null + else msg.copy(status = MessageStatus.COMPLETED, durationMs = 0L) + } else { + msg + } + } + _messages.value = finalized + + // Mirror the change into the current session's backing list. + _currentSessionId.value?.let { sessionId -> + _sessions.value.firstOrNull { it.id == sessionId }?.let { session -> + session.messages.clear() + session.messages.addAll(finalized) + } } } @@ -737,9 +922,12 @@ class ChatViewModel( stateUpdateJob?.cancel() stateUpdateJob = viewModelScope.launch { while (isActive) { - delay(100) // Update every 100ms - val elapsed = System.currentTimeMillis() - state.startTime - _agentState.value = state.copy(elapsedMillis = elapsed) + delay(100) + val current = _agentState.value + if (current !is AgentState.Executing) break + _agentState.value = current.copy( + elapsedMillis = System.currentTimeMillis() - current.startTime + ) } } } diff --git a/ai-assistant/src/main/res/values/strings.xml b/ai-assistant/src/main/res/values/strings.xml index b9695d93..cb1bfc2e 100644 --- a/ai-assistant/src/main/res/values/strings.xml +++ b/ai-assistant/src/main/res/values/strings.xml @@ -105,4 +105,33 @@ ⚠️ Experimental AI. Use at your own risk. Current AI backend Send + + + Reached the %1$d-step limit. Ask me to continue if the task isn\'t finished. + Stopped: the model kept requesting the same action. Try rephrasing your request. + I couldn\'t complete that, the last action failed. See the details above. + (No response. Try rephrasing, or pick a larger model in AI Settings — very small models struggle with tool use.) + Generating response… + System Log + + + Initializing engine… + Engine ready + Saved: %s + No model is currently loaded + Loading model, please wait… + ✅ Model loaded: %s + ❌ Error: %s + Loading model… + + + API Key saved on: %s + API Key is saved + API Key saved + API Key cannot be empty + API Key cleared + Current: %s + Loading… + Refresh Models + Model changed to %s diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/AgentLoopTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/AgentLoopTest.kt new file mode 100644 index 00000000..72e4f134 --- /dev/null +++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/AgentLoopTest.kt @@ -0,0 +1,341 @@ +package com.itsaky.androidide.plugins.aiassistant.tool + +import com.itsaky.androidide.plugins.aiassistant.models.ToolResult +import com.itsaky.androidide.plugins.services.LlmInferenceService.ChatMessage +import com.itsaky.androidide.plugins.services.LlmInferenceService.ChatMessage.Role +import kotlinx.coroutines.test.runTest +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertTrue +import org.junit.Test + +/** + * Unit tests for [AgentLoop] — the agentic tool-loop, tested in isolation from + * Android/streaming via plain fakes for generation and tool execution. + */ +class AgentLoopTest { + + /** Records prompts and returns scripted model responses in order. */ + private class ScriptedModel(private val responses: List) { + val prompts = mutableListOf() + var calls = 0 + suspend fun generate(prompt: String): String { + prompts += prompt + val r = responses.getOrElse(calls) { responses.last() } + calls++ + return r + } + } + + private fun toolCall(name: String) = """{"tool":"$name","args":{}}""" + + @Test + fun givenAModelThatCallsNoTools_whenTheLoopRuns_thenItStopsAfterOneTurn() = runTest { + val model = ScriptedModel(listOf("All done — here is your answer.")) + val history = mutableListOf(ChatMessage(Role.USER, "hello")) + var toolsInvoked = 0 + + val result = AgentLoop().run( + history = history, + generate = model::generate, + executeTools = { toolsInvoked++; emptyList() } + ) + + assertTrue(result.completed) + assertEquals(1, result.turns) + assertEquals(1, model.calls) + assertEquals(0, toolsInvoked) + // history now: user + one assistant turn + assertEquals(2, history.size) + assertEquals(Role.ASSISTANT, history[1].role) + } + + @Test + fun givenAModelThatCallsTheTerminalTool_whenTheLoopRuns_thenItFinishesViaOnFinalAnswerWithoutDispatchingIt() = runTest { + // Terminal tool ("respond") finishes via onFinalAnswer and is never dispatched. + val model = ScriptedModel( + listOf("""{"tool":"respond","args":{"message":"All set!"}}""") + ) + val history = mutableListOf(ChatMessage(Role.USER, "hi")) + var toolsInvoked = 0 + var finalTurn = -1 + var finalMessage: String? = null + + val result = AgentLoop(terminalTool = "respond").run( + history = history, + generate = model::generate, + executeTools = { toolsInvoked++; emptyList() }, + events = object : AgentLoop.Events { + override suspend fun onFinalAnswer(turn: Int, message: String) { + finalTurn = turn + finalMessage = message + } + } + ) + + assertTrue(result.completed) + assertEquals(AgentLoop.StopReason.COMPLETED, result.reason) + assertEquals(1, result.turns) + assertEquals(0, toolsInvoked) // terminal tool is NOT dispatched + assertEquals(1, finalTurn) + assertEquals("All set!", finalMessage) + } + + @Test + fun givenAModelThatCallsAToolThenAnswers_whenTheLoopRuns_thenItChainsTheToolAndFinishes() = runTest { + // Turn 1: model calls a tool. Turn 2: sees results, gives final answer. + val model = ScriptedModel( + listOf( + "Let me look. ${toolCall("open_file")}", + "Opened it. Done." + ) + ) + val history = mutableListOf(ChatMessage(Role.USER, "open MainActivity.java")) + val executed = mutableListOf>() + + val result = AgentLoop().run( + history = history, + generate = model::generate, + executeTools = { calls -> + executed += calls + listOf(ToolResult.success("Opened file in editor", "path/to/MainActivity.java")) + } + ) + + assertTrue(result.completed) + assertEquals(2, result.turns) + assertEquals(2, model.calls) + assertEquals(1, executed.size) + assertEquals("open_file", executed[0][0].name) + + // The 2nd prompt must contain the fed-back tool results (the whole point). + assertTrue( + "2nd turn prompt should include tool results", + model.prompts[1].contains("Tool results:") && + model.prompts[1].contains("MainActivity.java") + ) + + // history: user, assistant(turn1), user(tool results), assistant(turn2) + assertEquals(4, history.size) + assertEquals(Role.USER, history[2].role) + assertTrue(history[2].content.startsWith("Tool results:")) + } + + @Test + fun givenAModelThatKeepsCallingDistinctTools_whenTheLoopRuns_thenItStopsAtTheIterationCap() = runTest { + // Distinct calls each turn so stuck-detection doesn't fire before the cap. + val model = ScriptedModel( + listOf( + "step1 ${toolCall("list_files")}", + "step2 ${toolCall("read_file")}", + "step3 ${toolCall("search_project")}" + ) + ) + val history = mutableListOf(ChatMessage(Role.USER, "keep going")) + var maxReachedTurns = -1 + var toolBatches = 0 + + val result = AgentLoop(maxIterations = 3).run( + history = history, + generate = model::generate, + executeTools = { toolBatches++; listOf(ToolResult.success("ok")) }, + events = object : AgentLoop.Events { + override suspend fun onMaxIterationsReached(turns: Int) { maxReachedTurns = turns } + } + ) + + assertFalse(result.completed) + assertEquals(AgentLoop.StopReason.MAX_ITERATIONS, result.reason) + assertEquals(3, result.turns) + assertEquals(3, model.calls) + assertEquals(3, toolBatches) + assertEquals(3, maxReachedTurns) + } + + @Test + fun givenAModelThatRepeatsTheIdenticalToolCall_whenTheLoopRuns_thenItStopsAfterToleratingBoundedRepeats() = runTest { + val model = ScriptedModel(listOf(toolCall("list_files"))) // same call every turn + val history = mutableListOf(ChatMessage(Role.USER, "go")) + var repeatedTurns = -1 + var toolBatches = 0 + + // Default tolerance is 2: one repeat allowed, the second aborts. + val result = AgentLoop(maxIterations = 8).run( + history = history, + generate = model::generate, + executeTools = { toolBatches++; listOf(ToolResult.success("ok")) }, + events = object : AgentLoop.Events { + override suspend fun onRepeatedToolCalls(turns: Int) { repeatedTurns = turns } + } + ) + + assertFalse(result.completed) + assertEquals(AgentLoop.StopReason.REPEATED, result.reason) + assertEquals(3, result.turns) // turn1 + one tolerated repeat, abort on turn3 + assertEquals(3, repeatedTurns) + assertEquals(2, toolBatches) // executed twice, then stopped + } + + @Test + fun givenMaxConsecutiveRepeatsOfOne_whenAModelRepeatsImmediately_thenItStopsOnTheSecondTurn() = runTest { + val model = ScriptedModel(listOf(toolCall("list_files"))) + val history = mutableListOf(ChatMessage(Role.USER, "go")) + var toolBatches = 0 + + val result = AgentLoop(maxIterations = 8, maxConsecutiveRepeats = 1).run( + history = history, + generate = model::generate, + executeTools = { toolBatches++; listOf(ToolResult.success("ok")) } + ) + + assertEquals(AgentLoop.StopReason.REPEATED, result.reason) + assertEquals(2, result.turns) + assertEquals(1, toolBatches) + } + + @Test + fun givenATurnCoEmittingRespondAndARealTool_whenTheLoopRuns_thenTheRealToolStillRuns() = runTest { + // `respond` co-emitted with a real tool must not drop the real tool. + val model = ScriptedModel( + listOf( + """{"tool":"open_file","args":{"file_path":"MainActivity.java"}}""" + + """{"tool":"respond","args":{"message":"Opening it."}}""", + "Done." + ) + ) + val history = mutableListOf(ChatMessage(Role.USER, "open MainActivity")) + val executed = mutableListOf>() + + val result = AgentLoop(terminalTool = "respond").run( + history = history, + generate = model::generate, + executeTools = { calls -> + executed += calls + listOf(ToolResult.success("Opened file in editor", "path/MainActivity.java")) + } + ) + + assertTrue(result.completed) + assertEquals(1, executed.size) + assertEquals(listOf("open_file"), executed[0].map { it.name }) // respond dropped, open_file kept + } + + @Test + fun givenATurnWithOnlyTheTerminalTool_whenTheLoopRuns_thenItFinishesImmediately() = runTest { + val model = ScriptedModel( + listOf("""{"tool":"respond","args":{"message":"Hi!"}}""") + ) + val history = mutableListOf(ChatMessage(Role.USER, "hello")) + var toolsInvoked = 0 + var finalMessage: String? = null + + val result = AgentLoop(terminalTool = "respond").run( + history = history, + generate = model::generate, + executeTools = { toolsInvoked++; emptyList() }, + events = object : AgentLoop.Events { + override suspend fun onFinalAnswer(turn: Int, message: String) { finalMessage = message } + } + ) + + assertTrue(result.completed) + assertEquals(1, result.turns) + assertEquals(0, toolsInvoked) + assertEquals("Hi!", finalMessage) + } + + @Test + fun givenAToolThenAnswerRun_whenTheLoopRuns_thenEventsFireForEachModelTurnAndToolBatch() = runTest { + val model = ScriptedModel(listOf(toolCall("read_file"), "done")) + val history = mutableListOf(ChatMessage(Role.USER, "read it")) + val modelTurns = mutableListOf() + val toolTurns = mutableListOf() + + AgentLoop().run( + history = history, + generate = model::generate, + executeTools = { listOf(ToolResult.success("contents")) }, + events = object : AgentLoop.Events { + override suspend fun onModelTurn(turn: Int, text: String) { modelTurns += turn } + override suspend fun onToolResults(turn: Int, calls: List, results: List) { toolTurns += turn } + } + ) + + assertEquals(listOf(1, 2), modelTurns) + assertEquals(listOf(1), toolTurns) + } + + @Test + fun givenAFailingTool_whenTheLoopRuns_thenTheResultsAreFedBackAsFAILED() = runTest { + val model = ScriptedModel(listOf(toolCall("open_file"), "acknowledged")) + val history = mutableListOf(ChatMessage(Role.USER, "open nope")) + + AgentLoop().run( + history = history, + generate = model::generate, + executeTools = { listOf(ToolResult.failure("File not found", "does not exist")) } + ) + + val fedBack = history.first { it.role == Role.USER && it.content.startsWith("Tool results:") } + assertTrue(fedBack.content.contains("FAILED")) + assertTrue(fedBack.content.contains("File not found")) + } + + @Test + fun givenLongToolOutput_whenTheLoopRuns_thenItIsTruncatedBeforeFeedingBack() = runTest { + val big = "x".repeat(10_000) + val model = ScriptedModel(listOf(toolCall("read_file"), "ok")) + val history = mutableListOf(ChatMessage(Role.USER, "read big")) + + AgentLoop(toolOutputCharLimit = 500).run( + history = history, + generate = model::generate, + executeTools = { listOf(ToolResult.success("read", big)) } + ) + + val fedBack = history.first { it.content.startsWith("Tool results:") } + assertTrue(fedBack.content.contains("truncated")) + assertFalse("full 10k output must not be fed back", fedBack.content.contains(big)) + } + + @Test + fun givenASuccessfulToolResult_whenFormatToolResultsIsCalled_thenItBiasesTheModelToStop() { + val loop = AgentLoop() + val fedBack = loop.formatToolResults( + listOf(ToolCall("open_file", emptyMap())), + listOf(ToolResult.success("Opened file in editor", ".gitignore")) + ) + // After success, finishing is the default and another tool call is discouraged. + assertTrue(fedBack.contains("you are DONE")) + assertTrue(fedBack.contains("respond")) + assertTrue(fedBack.contains("Do NOT call another tool")) + } + + @Test + fun givenAFailedToolResult_whenFormatToolResultsIsCalled_thenItKeepsTheOpenEndedNextToolCue() { + val loop = AgentLoop() + val fedBack = loop.formatToolResults( + listOf(ToolCall("open_file", emptyMap())), + listOf(ToolResult.failure("File not found", "does not exist")) + ) + assertTrue(fedBack.contains("FAILED")) + assertTrue(fedBack.contains("call the next tool")) + } + + @Test + fun givenATranscript_whenRenderTranscriptIsCalled_thenItLabelsAssistantTurnsAndAddsNoTrailingCue() { + val loop = AgentLoop() + val transcript = loop.renderTranscript( + listOf( + ChatMessage(Role.USER, "hi"), + ChatMessage(Role.ASSISTANT, "hello"), + ChatMessage(Role.USER, "Tool results:\n[list_files] ok") + ) + ) + assertTrue(transcript.contains("hi")) + assertTrue(transcript.contains("Assistant: hello")) + assertTrue(transcript.contains("Tool results:")) + // Must not append a trailing "Assistant:" cue (the backend adds its own). + assertFalse("must not append a trailing Assistant cue", transcript.trimEnd().endsWith("Assistant:")) + } +} diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ExecutorTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ExecutorTest.kt new file mode 100644 index 00000000..163c6384 --- /dev/null +++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ExecutorTest.kt @@ -0,0 +1,115 @@ +package com.itsaky.androidide.plugins.aiassistant.tool + +import com.itsaky.androidide.plugins.aiassistant.models.ToolResult +import com.itsaky.androidide.plugins.aiassistant.tool.handlers.PathGuard +import kotlinx.coroutines.runBlocking +import org.junit.After +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertTrue +import org.junit.Before +import org.junit.Test +import java.io.File +import java.nio.file.Files + +/** + * Unit tests for the [Executor] path-containment pre-guard, in particular the + * [ToolHandler.resolvesPathsInternally] opt-out that lets read-only handlers + * rescue odd paths (e.g. a model-supplied "/.gitignore") instead of the Executor + * rejecting them outright, while write tools stay guarded. + */ +class ExecutorTest { + + private lateinit var projectRoot: File + + /** A handler that records whether it was dispatched and always succeeds. */ + private class FakeHandler( + override val toolName: String, + override val resolvesPathsInternally: Boolean, + ) : ToolHandler { + override val description = "fake" + override val requiresApproval = false + override val pathArgs = listOf("file_path") + var dispatched = false + private set + + override suspend fun execute(args: Map): ToolResult { + dispatched = true + return ToolResult.success("ran") + } + } + + @Before + fun setup() { + projectRoot = Files.createTempDirectory("executor-project").toFile().canonicalFile + PathGuard.setProjectRootForTesting(projectRoot.absolutePath) + } + + @After + fun tearDown() { + PathGuard.setProjectRootForTesting(null) + } + + private fun executorFor(handler: ToolHandler): Executor = + Executor(ToolRouter(listOf(handler)), ToolApprovalManager()) + + @Test + fun givenAnInternallyResolvingHandler_whenExecutingAnEscapingPath_thenTheEscapePreGuardIsBypassed() = runBlocking { + val handler = FakeHandler("fake_internal", resolvesPathsInternally = true) + val executor = executorFor(handler) + + // "/escape.txt" resolves outside the project root; the guard would reject + // it, but an internally-resolving handler must still be dispatched. + val results = executor.execute(listOf(ToolCall("fake_internal", mapOf("file_path" to "/escape.txt")))) + + assertTrue("handler should have been dispatched", handler.dispatched) + assertTrue("result should be the handler's success", results.single().success) + } + + @Test + fun givenADefaultHandler_whenExecutingAPathThatEscapesTheProjectRoot_thenItIsRejected() = runBlocking { + val handler = FakeHandler("fake_guarded", resolvesPathsInternally = false) + val executor = executorFor(handler) + + val results = executor.execute(listOf(ToolCall("fake_guarded", mapOf("file_path" to "/escape.txt")))) + + assertFalse("handler must NOT run for an escaping write path", handler.dispatched) + assertFalse(results.single().success) + assertTrue(results.single().message.contains("outside the project directory")) + } + + @Test + fun givenOpenFileWithAPathAlias_whenExecuting_thenPathIsRemappedToFilePathAndItRuns() = runBlocking { + // open_file requires file_path (like read_file); a model emitting + // {"path":"..."} must be remapped, not rejected for a missing file_path. + val handler = object : ToolHandler { + override val toolName = "open_file" + override val description = "fake open" + override val requiresApproval = false + override val pathArgs = listOf("file_path") + override val resolvesPathsInternally = true + var seenArgs: Map? = null + override suspend fun execute(args: Map): ToolResult { + seenArgs = args + return ToolResult.success("opened") + } + } + val executor = executorFor(handler) + + val results = executor.execute(listOf(ToolCall("open_file", mapOf("path" to "MainActivity.java")))) + + assertTrue("open_file with a path alias should run", results.single().success) + assertEquals("MainActivity.java", handler.seenArgs?.get("file_path")) + } + + @Test + fun givenADefaultHandler_whenExecutingAnInProjectPath_thenItRuns() = runBlocking { + val handler = FakeHandler("fake_guarded", resolvesPathsInternally = false) + val executor = executorFor(handler) + + val results = executor.execute(listOf(ToolCall("fake_guarded", mapOf("file_path" to "notes.txt")))) + + assertTrue("in-project path should be allowed through", handler.dispatched) + assertTrue(results.single().success) + } +} diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolCallExtractorTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolCallExtractorTest.kt new file mode 100644 index 00000000..a7671b1c --- /dev/null +++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolCallExtractorTest.kt @@ -0,0 +1,66 @@ +package com.itsaky.androidide.plugins.aiassistant.tool + +import org.junit.Assert.assertEquals +import org.junit.Assert.assertTrue +import org.junit.Test + +/** + * Unit tests for [ToolCallExtractor]. Focus: only explicit tool calls are honored, + * and prose never fires a tool (the "build Android apps" -> run_app regression). + */ +class ToolCallExtractorTest { + + @Test + fun givenAnExplicitToolCallTag_whenExtracting_thenTheCallIsExtractedWithArgs() { + val calls = ToolCallExtractor.extractToolCalls( + """Sure. {"tool":"open_file","args":{"file_path":"app/Main.java"}}""" + ) + assertEquals(1, calls.size) + assertEquals("open_file", calls[0].name) + assertEquals("app/Main.java", calls[0].args["file_path"]) + } + + @Test + fun givenABareJsonToolCall_whenExtracting_thenItIsExtracted() { + val calls = ToolCallExtractor.extractToolCalls( + """{"tool":"list_files","args":{"directory":"src"}}""" + ) + assertEquals(1, calls.size) + assertEquals("list_files", calls[0].name) + assertEquals("src", calls[0].args["directory"]) + } + + @Test + fun givenABareJsonToolCallWhoseValueContainsBraces_whenExtracting_thenItIsExtractedIntact() { + // The brace counter must ignore braces inside string values. + val calls = ToolCallExtractor.extractToolCalls( + """{"tool":"create_file","args":{"file_path":"A.kt","content":"fun f() { if (x) { y() } }"}}""" + ) + assertEquals(1, calls.size) + assertEquals("create_file", calls[0].name) + assertEquals("A.kt", calls[0].args["file_path"]) + assertEquals("fun f() { if (x) { y() } }", calls[0].args["content"]) + } + + @Test + fun givenAChattyReplyMentioningBuildingApps_whenExtracting_thenNoToolIsFired() { + val calls = ToolCallExtractor.extractToolCalls( + "Hi! I can help you build Android apps. What would you like to run or create next?" + ) + assertTrue("prose must not produce tool calls, got $calls", calls.isEmpty()) + } + + @Test + fun givenAPlainGreeting_whenExtracting_thenNoToolCallsAreProduced() { + assertTrue(ToolCallExtractor.extractToolCalls("Hello, how can I help?").isEmpty()) + } + + @Test + fun givenNarratedIntentWithoutATag_whenExtracting_thenNoToolCallsAreProduced() { + // The model describing what it would do must NOT be treated as a tool call. + val calls = ToolCallExtractor.extractToolCalls( + "Let me list the files in src and then read MainActivity.kt for you." + ) + assertTrue("narration must not produce tool calls, got $calls", calls.isEmpty()) + } +} diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ListFilesHandlerTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ListFilesHandlerTest.kt new file mode 100644 index 00000000..c4fca097 --- /dev/null +++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/ListFilesHandlerTest.kt @@ -0,0 +1,127 @@ +package com.itsaky.androidide.plugins.aiassistant.tool.handlers + +import com.itsaky.androidide.plugins.PluginContext +import io.mockk.mockk +import kotlinx.coroutines.runBlocking +import org.junit.After +import org.junit.Assert.assertFalse +import org.junit.Assert.assertTrue +import org.junit.Before +import org.junit.Test +import java.io.File +import java.nio.file.Files + +/** + * Unit tests for [ListFilesHandler] after it was refactored to resolve and + * containment-check paths through the shared [PathGuard]. + */ +class ListFilesHandlerTest { + + private lateinit var projectRoot: File + private lateinit var handler: ListFilesHandler + + @Before + fun setup() { + projectRoot = Files.createTempDirectory("listfiles-project").toFile().canonicalFile + PathGuard.setProjectRootForTesting(projectRoot.absolutePath) + handler = ListFilesHandler(mockk(relaxed = true)) + } + + @After + fun tearDown() { + PathGuard.setProjectRootForTesting(null) + PathGuard.setProjectRootProvider(null) + projectRoot.deleteRecursively() + } + + @Test + fun givenNoDirectoryArg_whenListing_thenItListsTheProjectRoot() = runBlocking { + File(projectRoot, "README.md").writeText("hi") + File(projectRoot, "app").mkdirs() + + val result = handler.execute(emptyMap()) + + assertTrue("Expected success, got: ${result.message}", result.success) + val data = result.data.orEmpty() + assertTrue(data.contains("README.md")) + assertTrue(data.contains("app")) + } + + @Test + fun givenABlankDirectoryArg_whenListing_thenItListsTheProjectRoot() = runBlocking { + File(projectRoot, "build.gradle.kts").writeText("x") + + val result = handler.execute(mapOf("directory" to " ")) + + assertTrue(result.success) + assertTrue(result.data.orEmpty().contains("build.gradle.kts")) + } + + @Test + fun givenASubdirectoryArg_whenListing_thenItListsThatSubdirectory() = runBlocking { + File(projectRoot, "src/main").mkdirs() + File(projectRoot, "src/main/Main.kt").writeText("fun main() {}") + + val result = handler.execute(mapOf("directory" to "src/main")) + + assertTrue(result.success) + assertTrue(result.data.orEmpty().contains("Main.kt")) + } + + @Test + fun givenASlashPrefixedRelativeDirectory_whenListing_thenItResolvesAsRelative() = runBlocking { + // A slash-prefixed relative dir must fall back to relative-to-root. + File(projectRoot, "src/main").mkdirs() + File(projectRoot, "src/main/Main.kt").writeText("fun main() {}") + + val result = handler.execute(mapOf("directory" to "/src/main")) + + assertTrue("Expected success, got: ${result.message}", result.success) + assertTrue(result.data.orEmpty().contains("Main.kt")) + } + + @Test + fun givenASlashPrefixedEscapingDirectory_whenListing_thenItIsStillRejected() = runBlocking { + // The fallback must not become a containment bypass. + val result = handler.execute(mapOf("directory" to "/../etc")) + + assertFalse(result.success) + assertTrue(result.message.contains("within project directory")) + } + + @Test + fun givenADirectoryEscapingTheProjectRoot_whenListing_thenItIsRejected() = runBlocking { + val result = handler.execute(mapOf("directory" to "../")) + + assertFalse(result.success) + assertTrue(result.message.contains("within project directory")) + } + + @Test + fun givenANonexistentDirectory_whenListing_thenItFails() = runBlocking { + val result = handler.execute(mapOf("directory" to "nope")) + + assertFalse(result.success) + assertTrue(result.message.contains("does not exist")) + } + + @Test + fun givenAFilePath_whenListing_thenItIsRejectedAsNotADirectory() = runBlocking { + File(projectRoot, "file.txt").writeText("x") + + val result = handler.execute(mapOf("directory" to "file.txt")) + + assertFalse(result.success) + assertTrue(result.message.contains("not a directory")) + } + + @Test + fun givenAnEmptyDirectory_whenListing_thenItReportsSuccessWithNoEntries() = runBlocking { + File(projectRoot, "empty").mkdirs() + + val result = handler.execute(mapOf("directory" to "empty")) + + assertTrue(result.success) + assertTrue(result.data.orEmpty().contains("no files or directories")) + } +} diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/OpenFileHandlerTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/OpenFileHandlerTest.kt new file mode 100644 index 00000000..e600e335 --- /dev/null +++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/OpenFileHandlerTest.kt @@ -0,0 +1,214 @@ +package com.itsaky.androidide.plugins.aiassistant.tool.handlers + +import com.itsaky.androidide.plugins.PluginContext +import com.itsaky.androidide.plugins.ServiceRegistry +import com.itsaky.androidide.plugins.services.IdeEditorService +import io.mockk.every +import io.mockk.mockk +import io.mockk.verify +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.runBlocking +import org.junit.After +import org.junit.Assert.assertEquals +import org.junit.Assert.assertFalse +import org.junit.Assert.assertTrue +import org.junit.Before +import org.junit.Test +import java.io.File +import java.nio.file.Files + +/** + * Unit tests for [OpenFileHandler] — the tool that opens a project file in the + * IDE editor. Covers every branch, including the path-containment guard that + * regressed on `open_file .gitignore`. + */ +class OpenFileHandlerTest { + + private lateinit var projectRoot: File + private lateinit var context: PluginContext + private lateinit var services: ServiceRegistry + private lateinit var editorService: IdeEditorService + private lateinit var handler: OpenFileHandler + + @Before + fun setup() { + projectRoot = Files.createTempDirectory("openfile-project").toFile().canonicalFile + PathGuard.setProjectRootForTesting(projectRoot.absolutePath) + + editorService = mockk(relaxed = true) + services = mockk() + context = mockk() + every { context.services } returns services + every { services.get(IdeEditorService::class.java) } returns editorService + + // Unconfined: no Android main looper in a JVM test. + handler = OpenFileHandler(context, Dispatchers.Unconfined) + } + + @After + fun tearDown() { + PathGuard.setProjectRootForTesting(null) + PathGuard.setProjectRootProvider(null) + projectRoot.deleteRecursively() + } + + private fun createFile(relative: String): File = + File(projectRoot, relative).apply { + parentFile?.mkdirs() + writeText("content") + } + + @Test + fun givenAnExistingProjectFile_whenOpened_thenItSucceeds() = runBlocking { + createFile(".gitignore") + every { editorService.openFile(any()) } returns true + + val result = handler.execute(mapOf("file_path" to ".gitignore")) + + assertTrue("Expected success, got: ${result.message}", result.success) + assertEquals(".gitignore", result.data) + verify { editorService.openFile(File(projectRoot, ".gitignore")) } + } + + @Test + fun givenALeadingSlashPath_whenOpened_thenItFallsBackToBasenameSearchAndOpens() = runBlocking { + // The model sometimes prepends a slash ("/.gitignore"), which resolves + // outside the project root; the basename fallback must still find it. + createFile(".gitignore") + every { editorService.openFile(any()) } returns true + + val result = handler.execute(mapOf("file_path" to "/.gitignore")) + + assertTrue("Expected success, got: ${result.message}", result.success) + verify { editorService.openFile(File(projectRoot, ".gitignore")) } + } + + @Test + fun givenALeadingSlashDirectoryQualifiedPath_whenOpened_thenItResolvesAsRelativeAndOpens() = runBlocking { + // With two .gitignore files, a bare basename is ambiguous; the model + // disambiguates with "app/.gitignore" but sends it slash-prefixed. It + // must resolve to the exact file, not collapse back to a basename search. + createFile(".gitignore") + createFile("app/.gitignore") + every { editorService.openFile(any()) } returns true + + val result = handler.execute(mapOf("file_path" to "/app/.gitignore")) + + assertTrue("Expected success, got: ${result.message}", result.success) + verify { editorService.openFile(File(projectRoot, "app/.gitignore")) } + } + + @Test + fun givenAMissingFilePath_whenOpened_thenItFails() = runBlocking { + val result = handler.execute(emptyMap()) + + assertFalse(result.success) + assertTrue(result.message.contains("file_path is required")) + verify(exactly = 0) { editorService.openFile(any()) } + } + + @Test + fun givenABlankFilePath_whenOpened_thenItFails() = runBlocking { + val result = handler.execute(mapOf("file_path" to " ")) + + assertFalse(result.success) + verify(exactly = 0) { editorService.openFile(any()) } + } + + @Test + fun givenAPathEscapingTheProjectRoot_whenOpened_thenItIsRejected() = runBlocking { + val result = handler.execute(mapOf("file_path" to "../outside.txt")) + + assertFalse(result.success) + assertTrue(result.message.contains("within project directory")) + verify(exactly = 0) { editorService.openFile(any()) } + } + + @Test + fun givenANonexistentFile_whenOpened_thenItFailsWithNotFound() = runBlocking { + val result = handler.execute(mapOf("file_path" to "does/not/exist.kt")) + + assertFalse(result.success) + assertEquals("File not found", result.message) + verify(exactly = 0) { editorService.openFile(any()) } + } + + @Test + fun givenALeadingSlashNonexistentPath_whenOpened_thenItReportsNotFoundRatherThanAnEscape() = runBlocking { + // "/does/not/exist.kt" has an in-root reading ("does/not/exist.kt") that + // simply doesn't exist — so the user sees "File not found", not the + // misleading "must be within project directory" (which is reserved for a + // real containment escape like "../outside.txt"). + val result = handler.execute(mapOf("file_path" to "/does/not/exist.kt")) + + assertFalse(result.success) + assertEquals("File not found", result.message) + verify(exactly = 0) { editorService.openFile(any()) } + } + + @Test + fun givenADirectoryPath_whenOpened_thenItIsRejectedAsNotAFile() = runBlocking { + File(projectRoot, "somedir").mkdirs() + + val result = handler.execute(mapOf("file_path" to "somedir")) + + assertFalse(result.success) + assertEquals("Not a file", result.message) + verify(exactly = 0) { editorService.openFile(any()) } + } + + @Test + fun givenNoEditorService_whenOpened_thenItFailsGracefully() = runBlocking { + createFile("Main.kt") + every { services.get(IdeEditorService::class.java) } returns null + + val result = handler.execute(mapOf("file_path" to "Main.kt")) + + assertFalse(result.success) + assertEquals("Editor service not available", result.message) + } + + @Test + fun givenABareFilename_whenOpened_thenItResolvesToItsRealNestedPath() = runBlocking { + createFile("app/src/main/java/com/example/MainActivity.java") + every { editorService.openFile(any()) } returns true + + val result = handler.execute(mapOf("file_path" to "MainActivity.java")) + + assertTrue("Expected success, got: ${result.message}", result.success) + verify { + editorService.openFile(File(projectRoot, "app/src/main/java/com/example/MainActivity.java")) + } + } + + @Test + fun givenAnAmbiguousBareFilename_whenOpened_thenItReturnsTheCandidatesInsteadOfOpening() = runBlocking { + createFile("app/src/main/java/A/Strings.kt") + createFile("app/src/main/java/B/Strings.kt") + + val result = handler.execute(mapOf("file_path" to "Strings.kt")) + + assertFalse(result.success) + assertTrue(result.message.contains("Multiple files")) + verify(exactly = 0) { editorService.openFile(any()) } + } + + @Test + fun givenABareFilenameWithNoMatch_whenOpened_thenItReportsNotFound() = runBlocking { + val result = handler.execute(mapOf("file_path" to "Nope.java")) + + assertFalse(result.success) + assertEquals("File not found", result.message) + } + + @Test + fun givenTheEditorReportingFailure_whenOpened_thenItSurfacesAFailureResult() = runBlocking { + createFile("Main.kt") + every { editorService.openFile(any()) } returns false + + val result = handler.execute(mapOf("file_path" to "Main.kt")) + + assertFalse(result.success) + assertEquals("Failed to open file", result.message) + } +} diff --git a/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/PathGuardTest.kt b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/PathGuardTest.kt new file mode 100644 index 00000000..bbed8ca6 --- /dev/null +++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/PathGuardTest.kt @@ -0,0 +1,265 @@ +package com.itsaky.androidide.plugins.aiassistant.tool.handlers + +import org.junit.After +import org.junit.Assert.assertEquals +import org.junit.Assert.assertNotNull +import org.junit.Assert.assertNull +import org.junit.Assert.assertTrue +import org.junit.Before +import org.junit.Test +import java.io.File +import java.nio.file.Files + +/** + * Unit tests for [PathGuard] — the shared containment guard for filesystem tools; + * covers the regression where anchoring to `user.dir` ("/" on Android) rejected + * every relative path (`open_file .gitignore` → "outside the project directory"). + */ +class PathGuardTest { + + private lateinit var projectRoot: File + + @Before + fun setup() { + projectRoot = Files.createTempDirectory("pathguard-project").toFile().canonicalFile + PathGuard.setProjectRootForTesting(projectRoot.absolutePath) + } + + @After + fun tearDown() { + // Reset shared singleton state so tests don't leak into each other. + PathGuard.setProjectRootForTesting(null) + PathGuard.setProjectRootProvider(null) + projectRoot.deleteRecursively() + } + + // --- The reported bug ----------------------------------------------------- + + @Test + fun givenARelativeDotfile_whenResolveWithinIsCalled_thenItResolvesInsideTheProjectRoot() { + val resolved = PathGuard.resolveWithin(".gitignore") + + assertNotNull("'.gitignore' must resolve inside the project root", resolved) + assertEquals(File(projectRoot, ".gitignore").canonicalPath, resolved!!.canonicalPath) + } + + @Test + fun givenTheFilesystemRootAsProjectRoot_whenAnyPathIsResolved_thenItIsRejected() { + // A root of "/" means no project is open; the guard must reject everything. + PathGuard.setProjectRootForTesting("/") + + assertNull("a relative path under '/' must be rejected", PathGuard.resolveWithin(".gitignore")) + assertNull("an absolute path under '/' must be rejected", PathGuard.resolveWithin("/etc/passwd")) + } + + @Test + fun givenANonExistentProjectRoot_whenAPathIsResolved_thenItIsRejected() { + PathGuard.setProjectRootForTesting("/no/such/project/dir/anywhere") + + assertNull(PathGuard.resolveWithin("build.gradle.kts")) + } + + // --- Normal containment --------------------------------------------------- + + @Test + fun givenANestedRelativePath_whenResolveWithinIsCalled_thenItResolvesInsideTheProjectRoot() { + val resolved = PathGuard.resolveWithin("app/src/main/AndroidManifest.xml") + + assertNotNull(resolved) + assertEquals( + File(projectRoot, "app/src/main/AndroidManifest.xml").canonicalPath, + resolved!!.canonicalPath + ) + } + + @Test + fun givenTheProjectRootPath_whenResolveWithinIsCalled_thenItResolves() { + val resolved = PathGuard.resolveWithin(".") + + assertNotNull(resolved) + assertEquals(projectRoot.canonicalPath, resolved!!.canonicalPath) + } + + @Test + fun givenAnAbsolutePathInsideTheRoot_whenResolveWithinIsCalled_thenItResolves() { + val inside = File(projectRoot, "build.gradle.kts").absolutePath + + val resolved = PathGuard.resolveWithin(inside) + + assertNotNull(resolved) + assertEquals(File(inside).canonicalPath, resolved!!.canonicalPath) + } + + // --- Escape attempts are rejected ---------------------------------------- + + @Test + fun givenARelativeTraversalEscapingTheRoot_whenResolveWithinIsCalled_thenItIsRejected() { + assertNull(PathGuard.resolveWithin("../secrets.txt")) + } + + @Test + fun givenAnAbsolutePathOutsideTheRoot_whenResolveWithinIsCalled_thenItIsRejected() { + assertNull(PathGuard.resolveWithin("/etc/passwd")) + } + + @Test + fun givenASiblingDirectorySharingANamePrefix_whenResolveWithinIsCalled_thenItIsRejected() { + // e.g. root "/tmp/proj" must NOT accept "/tmp/proj-evil". + assertNull(PathGuard.resolveWithin(projectRoot.absolutePath + "-evil/file.txt")) + } + + // --- Root resolution precedence ------------------------------------------ + + @Test + fun givenNoTestOverride_whenTheProjectRootIsQueried_thenTheProviderSuppliesIt() { + PathGuard.setProjectRootForTesting(null) + val providerRoot = Files.createTempDirectory("pathguard-provider").toFile().canonicalFile + try { + PathGuard.setProjectRootProvider { providerRoot.absolutePath } + + assertEquals(providerRoot.canonicalPath, File(PathGuard.projectRoot()).canonicalPath) + assertNotNull(PathGuard.resolveWithin("settings.gradle")) + } finally { + providerRoot.deleteRecursively() + } + } + + // --- findByName (basename resolution) ------------------------------------ + + @Test + fun givenANestedFile_whenFindByNameIsCalledWithItsBasename_thenItIsLocated() { + File(projectRoot, "app/src/main/java/com/example").mkdirs() + File(projectRoot, "app/src/main/java/com/example/MainActivity.java").writeText("x") + + val found = PathGuard.findByName("MainActivity.java") + + assertEquals(1, found.size) + assertEquals( + File(projectRoot, "app/src/main/java/com/example/MainActivity.java").canonicalPath, + found[0].canonicalPath + ) + } + + @Test + fun givenFilesDifferingOnlyInCase_whenFindByNameIsCalled_thenItMatchesAllCaseInsensitively() { + File(projectRoot, "a").mkdirs(); File(projectRoot, "a/Notes.txt").writeText("x") + File(projectRoot, "b").mkdirs(); File(projectRoot, "b/notes.txt").writeText("x") + + assertEquals(2, PathGuard.findByName("notes.txt").size) + } + + @Test + fun givenFilesUnderBuildAndDotDirectories_whenFindByNameIsCalled_thenTheyAreSkipped() { + File(projectRoot, "build/generated").mkdirs() + File(projectRoot, "build/generated/R.java").writeText("x") + File(projectRoot, ".git").mkdirs() + File(projectRoot, ".git/R.java").writeText("x") + + assertTrue("must not match files under build/ or .git/", PathGuard.findByName("R.java").isEmpty()) + } + + @Test + fun givenNoMatchingFile_whenFindByNameIsCalled_thenItReturnsEmpty() { + assertTrue(PathGuard.findByName("DoesNotExist.kt").isEmpty()) + } + + @Test + fun givenAFileReachableOnlyThroughASymlinkedDir_whenFindByNameIsCalled_thenItIsNotMatched() { + // A symlinked project subdir must not let the walk escape the root. + val outside = Files.createTempDirectory("pathguard-outside").toFile().canonicalFile + try { + File(outside, "Secret.kt").writeText("x") + val link = File(projectRoot, "linked").toPath() + try { + Files.createSymbolicLink(link, outside.toPath()) + } catch (e: Exception) { + // Filesystem/OS without symlink support — nothing to assert. + return + } + assertTrue( + "a match reachable only via a symlinked dir must be dropped", + PathGuard.findByName("Secret.kt").isEmpty() + ) + } finally { + outside.deleteRecursively() + } + } + + // --- resolve() — the shared handler resolution policy -------------------- + + @Test + fun givenAnExistingRelativePath_whenResolveIsCalled_thenItReturnsResolved() { + File(projectRoot, "app").mkdirs() + val target = File(projectRoot, "app/build.gradle.kts").apply { writeText("x") } + + val resolution = PathGuard.resolve("app/build.gradle.kts") + + assertTrue(resolution is PathGuard.Resolution.Resolved) + assertEquals(target.canonicalPath, (resolution as PathGuard.Resolution.Resolved).file.canonicalPath) + } + + @Test + fun givenASlashPrefixedPath_whenResolveIsCalled_thenItRetriesAsRelativeBeforeBasenameSearch() { + // The slash-stripped relative retry must resolve before an ambiguous basename search. + File(projectRoot, "app").mkdirs(); File(projectRoot, "app/.gitignore").writeText("x") + File(projectRoot, ".gitignore").writeText("x") + + val resolution = PathGuard.resolve("/app/.gitignore") + + assertTrue(resolution is PathGuard.Resolution.Resolved) + assertEquals( + File(projectRoot, "app/.gitignore").canonicalPath, + (resolution as PathGuard.Resolution.Resolved).file.canonicalPath + ) + } + + @Test + fun givenABareName_whenResolveIsCalled_thenItReturnsResolvedViaBasenameSearch() { + File(projectRoot, "app/src/main").mkdirs() + File(projectRoot, "app/src/main/MainActivity.java").writeText("x") + + val resolution = PathGuard.resolve("MainActivity.java") + + assertTrue(resolution is PathGuard.Resolution.Resolved) + } + + @Test + fun givenABasenameMatchingSeveralFiles_whenResolveIsCalled_thenItReturnsAmbiguous() { + File(projectRoot, "a").mkdirs(); File(projectRoot, "a/Strings.kt").writeText("x") + File(projectRoot, "b").mkdirs(); File(projectRoot, "b/Strings.kt").writeText("x") + + val resolution = PathGuard.resolve("Strings.kt") + + assertTrue(resolution is PathGuard.Resolution.Ambiguous) + assertEquals("Strings.kt", (resolution as PathGuard.Resolution.Ambiguous).baseName) + assertEquals(2, resolution.matches.size) + } + + @Test + fun givenAPathWithNoInRootInterpretation_whenResolveIsCalled_thenItReturnsEscaped() { + assertEquals(PathGuard.Resolution.Escaped, PathGuard.resolve("../outside.txt")) + } + + @Test + fun givenASlashPrefixedInRootPathThatDoesNotExist_whenResolveIsCalled_thenItReturnsNotFound() { + // Absolute miss but in-root as relative → NotFound, not Escaped. + assertEquals(PathGuard.Resolution.NotFound, PathGuard.resolve("/does/exist.kt")) + } + + @Test + fun givenAnInRootPathThatDoesNotExist_whenResolveIsCalled_thenItReturnsNotFound() { + assertEquals(PathGuard.Resolution.NotFound, PathGuard.resolve("does/not/exist.kt")) + } + + @Test + fun givenABlankProviderResult_whenTheProjectRootIsQueried_thenItFallsThroughToTheNextSource() { + PathGuard.setProjectRootForTesting(null) + PathGuard.setProjectRootProvider { " " } + System.setProperty("project.dir", projectRoot.absolutePath) + try { + assertEquals(projectRoot.canonicalPath, File(PathGuard.projectRoot()).canonicalPath) + } finally { + System.clearProperty("project.dir") + } + } +} diff --git a/ai-core/libs/llama-api.jar b/ai-core/libs/llama-api.jar index 3274247d..6d160ed2 100644 Binary files a/ai-core/libs/llama-api.jar and b/ai-core/libs/llama-api.jar differ diff --git a/ai-core/libs/v8/llama-v8-release.aar b/ai-core/libs/v8/llama-v8-release.aar index 9b2984f4..d7730b01 100644 Binary files a/ai-core/libs/v8/llama-v8-release.aar and b/ai-core/libs/v8/llama-v8-release.aar differ diff --git a/ai-core/llama-impl/src/main/cpp/llama-android.cpp b/ai-core/llama-impl/src/main/cpp/llama-android.cpp index 91cff6e9..e0771fe4 100644 --- a/ai-core/llama-impl/src/main/cpp/llama-android.cpp +++ b/ai-core/llama-impl/src/main/cpp/llama-android.cpp @@ -17,6 +17,16 @@ #define LOGi(...) __android_log_print(ANDROID_LOG_INFO, TAG, __VA_ARGS__) #define LOGe(...) __android_log_print(ANDROID_LOG_ERROR, TAG, __VA_ARGS__) +// Verbose, per-token tracing. It fires once per prompt token and once per +// generated token — hundreds of log/JNI calls per reply on the hot path — so it +// is compiled out of release builds (NDEBUG). The __VA_ARGS__ are not evaluated +// in release, so any string-building in the arguments is skipped too. +#ifdef NDEBUG +#define LOGv(...) ((void) 0) +#else +#define LOGv(...) __android_log_print(ANDROID_LOG_INFO, TAG, __VA_ARGS__) +#endif + jclass la_int_var; jmethodID la_int_var_value; jmethodID la_int_var_inc; @@ -554,6 +564,51 @@ Java_android_llama_cpp_LLamaAndroid_new_1sampler(JNIEnv *, jobject) { return reinterpret_cast(smpl); } +/** + * Build a sampler chain constrained by a GBNF grammar, so the model can only + * emit tokens the grammar allows — used for reliable text-based tool calls on + * weak local models. + * + * @param model_pointer native llama_model handle. + * @param grammar GBNF grammar text, entered at its "root" rule. + * @return the native sampler handle, or 0 if the model is null or the grammar + * fails to parse (the caller falls back to the plain sampler). + */ +extern "C" +JNIEXPORT jlong JNICALL +Java_android_llama_cpp_LLamaAndroid_new_1grammar_1sampler( + JNIEnv *env, jobject, jlong model_pointer, jstring grammar) { + const auto model = reinterpret_cast(model_pointer); + if (model == nullptr) return 0; + if (grammar == nullptr) return 0; // no grammar string — caller falls back + + const llama_vocab *vocab = llama_model_get_vocab(model); + if (vocab == nullptr) return 0; // model has no vocab — can't build a grammar sampler + + // NULL under memory pressure (pending OOM); clear the exception and fall back. + const char *grammar_cstr = env->GetStringUTFChars(grammar, nullptr); + if (grammar_cstr == nullptr) { + if (env->ExceptionCheck()) env->ExceptionClear(); + return 0; + } + llama_sampler *grmr = llama_sampler_init_grammar(vocab, grammar_cstr, "root"); + env->ReleaseStringUTFChars(grammar, grammar_cstr); + if (grmr == nullptr) return 0; // invalid grammar — caller falls back + + auto sparams = llama_sampler_chain_default_params(); + sparams.no_perf = true; + llama_sampler *smpl = llama_sampler_chain_init(sparams); + + // Grammar first: it masks tokens that would violate the grammar, so the + // selector below only ever picks a valid token. + llama_sampler_chain_add(smpl, grmr); + llama_sampler_chain_add(smpl, llama_sampler_init_penalties(64, 1.1f, 0.0f, 0.0f)); + // Greedy: with the grammar mask in place, pick the most likely valid token. + llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); + + return reinterpret_cast(smpl); +} + extern "C" JNIEXPORT void JNICALL Java_android_llama_cpp_LLamaAndroid_free_1sampler(JNIEnv *, jobject, jlong sampler_pointer) { @@ -632,55 +687,45 @@ Java_android_llama_cpp_LLamaAndroid_completion_1init( g_prompt_tokens = static_cast(tokens_list.size()); for (auto id: tokens_list) { - LOGi("token: `%s`-> %d ", common_token_to_piece(context, id).c_str(), id); + LOGv("token: `%s`-> %d ", common_token_to_piece(context, id).c_str(), id); } common_batch_clear(*batch); - bool reuse = false; - size_t reuse_prefix = 0; + // Reuse the longest common prefix with the cached sequence so the unchanged prefix (system prompt) isn't re-prefilled. + size_t lcp = 0; { std::lock_guard lock(g_globals_mutex); if (g_kv_cache_reuse.load() && !g_cached_tokens.empty()) { - if (g_cached_tokens.size() <= tokens_list.size()) { - reuse = true; - for (size_t i = 0; i < g_cached_tokens.size(); i++) { - if (g_cached_tokens[i] != tokens_list[i]) { - reuse = false; - break; - } - } - if (reuse) { - reuse_prefix = g_cached_tokens.size(); - } + const size_t maxlcp = std::min(g_cached_tokens.size(), tokens_list.size()); + while (lcp < maxlcp && g_cached_tokens[lcp] == tokens_list[lcp]) { + lcp++; } } } - if (!reuse) { - // Fully reset KV cache to avoid non-consecutive sequence positions. - llama_memory_clear(llama_get_memory(context), true); - { - std::lock_guard lock(g_globals_mutex); - if (!g_kv_cache_reuse.load()) { - g_cached_tokens.clear(); - } - g_cached_tokens.assign(tokens_list.begin(), tokens_list.end()); - } - // evaluate the initial prompt - for (auto i = 0; i < tokens_list.size(); i++) { - common_batch_add(*batch, tokens_list[i], i, {0}, false); - } + // Always leave at least one token to decode, so we get logits for the next token. + if (lcp == tokens_list.size() && lcp > 0) { + lcp--; + } + + llama_memory_t mem = llama_get_memory(context); + if (lcp == 0) { + // Nothing reusable — full reset. + llama_memory_clear(mem, true); } else { - { - std::lock_guard lock(g_globals_mutex); - g_cached_tokens.assign(tokens_list.begin(), tokens_list.end()); - } - if (reuse_prefix < tokens_list.size()) { - for (auto i = reuse_prefix; i < tokens_list.size(); i++) { - common_batch_add(*batch, tokens_list[i], i, {0}, false); - } - } + // Evict cached positions past the common prefix. + llama_memory_seq_rm(mem, 0, (llama_pos) lcp, -1); + } + + { + std::lock_guard lock(g_globals_mutex); + g_cached_tokens.assign(tokens_list.begin(), tokens_list.end()); + } + + // Prefill only the divergent tail. + for (size_t i = lcp; i < tokens_list.size(); i++) { + common_batch_add(*batch, tokens_list[i], (llama_pos) i, {0}, false); } if (batch->n_tokens > 0) { @@ -797,11 +842,15 @@ Java_android_llama_cpp_LLamaAndroid_completion_1loop( new_token = new_jstring_utf8(env, cached_token_chars.c_str()); } +#ifndef NDEBUG + // Per-token JNI upcall into the Kotlin logger — debug-only; on the hot + // path in release it would add a JNI round-trip (and a lock) per token. { std::lock_guard lock(g_globals_mutex); log_info_to_kt("cached: %s, new_token_chars: `%s`, id: %d", cached_token_chars.c_str(), new_token_chars.c_str(), new_token_id); } +#endif { std::lock_guard lock(g_globals_mutex); diff --git a/ai-core/llama-impl/src/main/java/android/llama/cpp/LLamaAndroid.kt b/ai-core/llama-impl/src/main/java/android/llama/cpp/LLamaAndroid.kt index 7d0ff4a1..10f73297 100644 --- a/ai-core/llama-impl/src/main/java/android/llama/cpp/LLamaAndroid.kt +++ b/ai-core/llama-impl/src/main/java/android/llama/cpp/LLamaAndroid.kt @@ -94,6 +94,19 @@ class LLamaAndroid : ILlamaController { private val isStopped = AtomicBoolean(false) + /** + * Optional GBNF grammar constraining the next generation(s). When set, [send] + * builds a fresh grammar-constrained sampler per call (grammar state is + * per-generation) and frees it afterward. Null → the plain load-time sampler. + */ + @Volatile + private var grammar: String? = null + + /** Set (or clear, with null) the GBNF grammar used to constrain generation. */ + fun setGrammar(gbnf: String?) { + grammar = gbnf?.takeIf { it.isNotBlank() } + } + @Volatile private var runLoopDispatcher: ExecutorCoroutineDispatcher = createRunLoop() @@ -152,6 +165,7 @@ class LLamaAndroid : ILlamaController { private external fun new_batch(nTokens: Int, embd: Int, nSeqMax: Int): Long private external fun free_batch(batch: Long) private external fun new_sampler(): Long + private external fun new_grammar_sampler(model: Long, grammar: String): Long private external fun free_sampler(sampler: Long) private external fun bench_model( context: Long, @@ -263,28 +277,43 @@ class LLamaAndroid : ILlamaController { kv_cache_clear(state.context) } - val ncur = IntVar( - completion_init( - state.context, - state.batch, - message, - formatChat, - nlen, - stop.toTypedArray() + // A grammar sampler is stateful, so build a fresh one for this + // generation and free it when done. Fall back to the plain + // load-time sampler if there's no grammar or it fails to build. + val grammarText = grammar + val grammarSampler = grammarText?.let { new_grammar_sampler(state.model, it) } + ?.takeIf { it != 0L } + if (grammarText != null && grammarSampler == null) { + log.warn("Grammar sampler failed to build; falling back to the plain sampler") + } + val sampler = grammarSampler ?: state.sampler + + try { + val ncur = IntVar( + completion_init( + state.context, + state.batch, + message, + formatChat, + nlen, + stop.toTypedArray() + ) ) - ) - - while (true) { - if (isStopped.get()) { - log.info("Stopping generation loop because stop flag was set.") - break - } - val str = completion_loop(state.context, state.batch, state.sampler, nlen, ncur) - if (str == null) { - break + while (true) { + if (isStopped.get()) { + log.info("Stopping generation loop because stop flag was set.") + break + } + + val str = completion_loop(state.context, state.batch, sampler, nlen, ncur) + if (str == null) { + break + } + emit(str) } - emit(str) + } finally { + if (grammarSampler != null) free_sampler(grammarSampler) } } diff --git a/ai-core/scripts/rebuild-llama-aar.sh b/ai-core/scripts/rebuild-llama-aar.sh index 848d72e9..5f37d231 100755 --- a/ai-core/scripts/rebuild-llama-aar.sh +++ b/ai-core/scripts/rebuild-llama-aar.sh @@ -11,12 +11,12 @@ set -euo pipefail -# Run from the ai-assistant/ project root regardless of where it's invoked. +# Run from the ai-core/ project root regardless of where it's invoked. cd "$(dirname "$0")/.." -AAR_DST="ai-core-plugin/libs/v8/llama-v8-release.aar" +AAR_DST="libs/v8/llama-v8-release.aar" AAR_SRC="llama-impl/build/outputs/aar/llama-impl-release.aar" -API_DST="ai-core-plugin/libs/llama-api.jar" +API_DST="libs/llama-api.jar" API_SRC="llama-api/build/libs/llama-api.jar" echo "==> Initializing the llama.cpp submodule (source for the native build)" @@ -25,7 +25,7 @@ git submodule update --init --recursive echo "==> Building :llama-impl (native lib) and :llama-api (interface jar)" ./gradlew :llama-impl:assembleRelease :llama-api:jar -echo "==> Copying artifacts into ai-core-plugin/libs" +echo "==> Copying artifacts into ai-core/libs" cp "$AAR_SRC" "$AAR_DST" cp "$API_SRC" "$API_DST" diff --git a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/CancellableBackend.kt b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/CancellableBackend.kt new file mode 100644 index 00000000..3f992040 --- /dev/null +++ b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/CancellableBackend.kt @@ -0,0 +1,11 @@ +package com.itsaky.androidide.plugins.aicore + +/** + * A backend whose in-flight streaming generation can be cancelled (Stop pressed). + * Implement this so [LlmInferenceServiceImpl.cancelGeneration] cancels it without + * a per-type `when` branch. + */ +interface CancellableBackend { + /** Cancel any in-flight generation. */ + fun cancelStreaming() +} diff --git a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/GeminiBackend.kt b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/GeminiBackend.kt index 081488b8..d1a71d1c 100644 --- a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/GeminiBackend.kt +++ b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/GeminiBackend.kt @@ -28,7 +28,7 @@ import java.util.concurrent.CompletableFuture * OkHttp (no such overload) — that mismatch crashed generation with a NoSuchMethodError. * HttpURLConnection has no third-party dependency, so it works regardless of the host's OkHttp. */ -class GeminiBackend(private val context: PluginContext) : LlmBackend { +class GeminiBackend(private val context: PluginContext) : LlmBackend, CancellableBackend { private val scope = CoroutineScope(Dispatchers.IO) @@ -390,6 +390,12 @@ User: $userPrompt""" } } + /** Cancel any in-flight generation (user pressed Stop). */ + override fun cancelStreaming() { + currentJob?.cancel() + currentJob = null + } + /** * Release all resources: cancel the backend scope and any in-flight * request. Called from AiCorePlugin.dispose(). diff --git a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LlmInferenceServiceImpl.kt b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LlmInferenceServiceImpl.kt index 19d87b6d..fe3a7f24 100644 --- a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LlmInferenceServiceImpl.kt +++ b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LlmInferenceServiceImpl.kt @@ -168,5 +168,8 @@ class LlmInferenceServiceImpl : LlmInferenceService { override fun cancelGeneration() { currentGeneration?.cancel(true) currentGeneration = null + + backends.values.filterIsInstance() + .forEach { it.cancelStreaming() } } } diff --git a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LocalLlmBackend.kt b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LocalLlmBackend.kt index d9e437a9..59e6be34 100644 --- a/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LocalLlmBackend.kt +++ b/ai-core/src/main/kotlin/com/itsaky/androidide/plugins/aicore/LocalLlmBackend.kt @@ -6,26 +6,57 @@ import android.provider.OpenableColumns import com.itsaky.androidide.plugins.services.LlmInferenceService.* import com.itsaky.androidide.plugins.services.SharedServices import com.itsaky.androidide.plugins.PluginContext +import kotlinx.coroutines.CancellationException import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.Job import kotlinx.coroutines.cancel +import kotlinx.coroutines.ensureActive import kotlinx.coroutines.launch +import kotlinx.coroutines.sync.Mutex +import kotlinx.coroutines.sync.withLock import java.io.File import java.io.FileOutputStream import java.util.concurrent.CompletableFuture +import java.util.concurrent.atomic.AtomicBoolean /** * Local LLM backend using llama-impl for on-device inference. * Wraps llama-impl APIs and implements LlmBackend interface. */ -class LocalLlmBackend(private val context: PluginContext) : LlmBackend { +class LocalLlmBackend(private val context: PluginContext) : LlmBackend, CancellableBackend { + + companion object { + /** + * [LlmConfig.extraParams] key for an optional GBNF grammar. The caller + * owns it (keeping this backend free of any tool vocabulary); absent → + * unconstrained sampling. + */ + const val EXTRA_PARAM_GRAMMAR = "grammar" + + /** + * Stops the model fabricating the next `User:` turn of this backend's own + * scaffold — otherwise an unconstrained model runs on to `maxTokens`. The + * native stop truncates before the match, so it only trims that fabrication. + */ + private val SCAFFOLD_STOP = listOf("\nUser:") + } private val llama by lazy { LLamaAndroid.instance() } private val scope = CoroutineScope(Dispatchers.IO) + /** Single-flight guard serializing whole generations on the shared native context. */ + private val generationMutex = Mutex() + + @Volatile private var currentStreamingJob: Job? = null + @Volatile private var currentGenerateJob: Job? = null + @Volatile private var modelLoaded = false @Volatile private var currentModelPath: String? = null + /** Ensures the background warm-up load is launched at most once. */ + private val warmUpStarted = AtomicBoolean(false) + override fun getId(): String = "local" override fun getName(): String = "Local LLM" @@ -44,10 +75,37 @@ class LocalLlmBackend(private val context: PluginContext) : LlmBackend { context.logger.debug("LocalLlmBackend.isAvailable() - configured path: $configuredPath, modelLoaded: $modelLoaded") + // Chat-open hits this; start loading now so the first message isn't gated on a cold load. + maybeWarmUp(configuredPath, prefs?.getString("ai_backend_preference", "LOCAL_LLM")) + // Available if model is loaded OR if a path is configured return modelLoaded || !configuredPath.isNullOrBlank() } + /** + * Preloads the configured model in the background, once, so the first generation + * doesn't pay the cold-load cost. No-op unless the local backend is the selected one. + * @param configuredPath the configured model path/URI, or null/blank if unset. + * @param backendPreference the `ai_backend_preference` value ("LOCAL_LLM"/"GEMINI"). + */ + private fun maybeWarmUp(configuredPath: String?, backendPreference: String?) { + if (configuredPath.isNullOrBlank() || modelLoaded) return + if (backendPreference == "GEMINI") return // user isn't using the local backend + if (!warmUpStarted.compareAndSet(false, true)) return + + scope.launch { + try { + // Serialize with real generations so a mid-warm-up send just waits for this load. + generationMutex.withLock { ensureModelLoaded(configuredPath) } + context.logger.info("Local model warm-up complete") + } catch (e: Exception) { + // Stay silent (the real send surfaces config errors); allow a later retry. + context.logger.warn("Local model warm-up failed: ${e.message}") + warmUpStarted.set(false) + } + } + } + /** * Resolves the user-selected model reference to a real filesystem path the native * loader can `fopen`. @@ -204,46 +262,59 @@ class LocalLlmBackend(private val context: PluginContext) : LlmBackend { val future = CompletableFuture() - scope.launch { + currentGenerateJob = scope.launch { try { - // Configure sampling (use defaults for topP and topK) - LLamaAndroid.configureSampling( - config.temperature, - 0.9f, // topP default - 40 // topK default - ) - LLamaAndroid.configureMaxTokens(config.maxTokens) - - // Ensure model is loaded - ensureModelLoaded(configuredPath) - - // Build full prompt with system message - val fullPrompt = if (config.systemPrompt != null) { - "${config.systemPrompt}\n\nUser: $prompt\nAssistant:" - } else { - "User: $prompt\nAssistant:" - } + // Serialize against other generations on the shared native context. + generationMutex.withLock { + // Configure sampling (use defaults for topP and topK) + LLamaAndroid.configureSampling( + config.temperature, + 0.9f, // topP default + 40 // topK default + ) + LLamaAndroid.configureMaxTokens(config.maxTokens) + + // Ensure model is loaded + ensureModelLoaded(configuredPath) + + // Build full prompt with system message + val fullPrompt = if (config.systemPrompt != null) { + "${config.systemPrompt}\n\nUser: $prompt\nAssistant:" + } else { + "User: $prompt\nAssistant:" + } - val startTime = System.currentTimeMillis() + val startTime = System.currentTimeMillis() + + // Collect all tokens + val responseBuilder = StringBuilder() + var tokenCount = 0 + + // Unconstrained path: clear any grammar left by a concurrent + // streaming call so completions never inherit a tool-call grammar. + llama.setGrammar(null) + + llama.send( + message = fullPrompt, + formatChat = false, + stop = SCAFFOLD_STOP, + clearCache = false + ).collect { token -> + // Honor cancellation so Stop frees the run loop early. + ensureActive() + responseBuilder.append(token) + tokenCount++ + } - // Collect all tokens - val responseBuilder = StringBuilder() - var tokenCount = 0 + val responseText = responseBuilder.toString() + context.logger.info("Generated response: ${responseText.take(50)}... ($tokenCount tokens)") - llama.send( - message = fullPrompt, - formatChat = false, - stop = emptyList(), - clearCache = false - ).collect { token -> - responseBuilder.append(token) - tokenCount++ + future.complete(LlmResponse.success(responseText, tokenCount, System.currentTimeMillis() - startTime)) } - - val responseText = responseBuilder.toString() - context.logger.info("Generated response: ${responseText.take(50)}... ($tokenCount tokens)") - - future.complete(LlmResponse.success(responseText, tokenCount, System.currentTimeMillis() - startTime)) + } catch (ce: CancellationException) { + context.logger.info("Generation cancelled") + future.completeExceptionally(ce) + throw ce } catch (e: Exception) { context.logger.error("Error during generation", e) if (e is ModelNotConfiguredException || e is IncompatibleModelException) { @@ -274,42 +345,58 @@ class LocalLlmBackend(private val context: PluginContext) : LlmBackend { return } - scope.launch { + currentStreamingJob = scope.launch { try { - // Configure sampling (use defaults for topP and topK) - LLamaAndroid.configureSampling( - config.temperature, - 0.9f, // topP default - 40 // topK default - ) - LLamaAndroid.configureMaxTokens(config.maxTokens) - - // Ensure model is loaded - ensureModelLoaded(configuredPath) - - // Build full prompt with system message - val fullPrompt = if (config.systemPrompt != null) { - "${config.systemPrompt}\n\nUser: $prompt\nAssistant:" - } else { - "User: $prompt\nAssistant:" - } - - val startTime = System.currentTimeMillis() - var tokenCount = 0 - val responseBuilder = StringBuilder() - - llama.send( - message = fullPrompt, - formatChat = false, - stop = emptyList(), - clearCache = false - ).collect { token -> - callback.onToken(token) - responseBuilder.append(token) - tokenCount++ + // Hold the single-flight lock for the whole streaming generation. + generationMutex.withLock { + try { + // Configure sampling (use defaults for topP and topK) + LLamaAndroid.configureSampling( + config.temperature, + 0.9f, // topP default + 40 // topK default + ) + LLamaAndroid.configureMaxTokens(config.maxTokens) + + // Ensure model is loaded + ensureModelLoaded(configuredPath) + + // Build full prompt with system message + val fullPrompt = if (config.systemPrompt != null) { + "${config.systemPrompt}\n\nUser: $prompt\nAssistant:" + } else { + "User: $prompt\nAssistant:" + } + + val startTime = System.currentTimeMillis() + var tokenCount = 0 + val responseBuilder = StringBuilder() + + // Apply the caller's grammar for this send() only; reset in the finally. + val grammar = config.extraParams?.get(EXTRA_PARAM_GRAMMAR) as? String + llama.setGrammar(grammar?.takeIf { it.isNotBlank() }) + + llama.send( + message = fullPrompt, + formatChat = false, + // Keep the KV cache so the native layer reuses the common prefix (system prompt). + stop = SCAFFOLD_STOP, + clearCache = false + ).collect { token -> + ensureActive() + callback.onToken(token) + responseBuilder.append(token) + tokenCount++ + } + + callback.onComplete(LlmResponse.success(responseBuilder.toString(), tokenCount, System.currentTimeMillis() - startTime)) + } finally { + llama.setGrammar(null) + } } - - callback.onComplete(LlmResponse.success(responseBuilder.toString(), tokenCount, System.currentTimeMillis() - startTime)) + } catch (ce: CancellationException) { + context.logger.info("Streaming generation cancelled") + throw ce } catch (e: Exception) { context.logger.error("Error during streaming generation", e) if (e is ModelNotConfiguredException || e is IncompatibleModelException) { @@ -320,6 +407,17 @@ class LocalLlmBackend(private val context: PluginContext) : LlmBackend { } } + /** + * Cancels any in-flight streaming or non-streaming generation (user pressed Stop), + * cancelling the coroutine Job so the single-threaded run loop is freed early. + */ + override fun cancelStreaming() { + currentStreamingJob?.cancel() + currentStreamingJob = null + currentGenerateJob?.cancel() + currentGenerateJob = null + } + override fun generateWithHistory( history: List, prompt: String,