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..81b595fb 100644
--- a/ai-assistant/build.gradle.kts
+++ b/ai-assistant/build.gradle.kts
@@ -22,6 +22,7 @@ android {
buildFeatures {
viewBinding = true
+ buildConfig = true
}
buildTypes {
@@ -45,6 +46,10 @@ android {
}
}
+ testOptions {
+ unitTests.isReturnDefaultValues = true
+ }
+
packaging {
resources {
excludes += setOf(
@@ -77,10 +82,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 0f0e60ea..72fecfbb 100644
--- a/ai-assistant/src/main/AndroidManifest.xml
+++ b/ai-assistant/src/main/AndroidManifest.xml
@@ -1,6 +1,5 @@
-
+
()
+ private val animatingHolders = mutableSetOf()
companion object {
private const val VIEW_TYPE_DEFAULT = 0
@@ -219,7 +220,7 @@ class ChatAdapter(
hideGeneratingDots(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)
}
@@ -227,7 +228,7 @@ class ChatAdapter(
// tag has to follow the role it currently has.
wireTooltip(holder.btnRetry, AiAssistantPlugin.TOOLTIP_TAG_MESSAGE_OPEN_SETTINGS)
} else {
- holder.btnRetry.text = "Retry"
+ holder.btnRetry.text = holder.btnRetry.context.getString(R.string.action_retry)
holder.btnRetry.setOnClickListener {
onMessageAction(ACTION_RETRY, message)
}
@@ -257,7 +258,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 {
@@ -268,14 +269,14 @@ class ChatAdapter(
}
/**
- * Starts — or restarts — the "..." animation, cancelling any step already queued for [holder]
- * so repeated binds of one recycled row cannot stack loops. The step is posted on the dots
- * view, not a bare main-looper Handler, so [hideGeneratingDots] can cancel it.
+ * Starts the "..." animation, or leaves an already-running one alone: restarting on every
+ * streaming rebind would reset the loop to "." and it would never visibly advance. The step is
+ * posted on the dots view, not a bare main-looper Handler, so [hideGeneratingDots] can cancel it.
*
* @param holder the row whose dots should animate
*/
private fun animateGeneratingDots(holder: DefaultMessageViewHolder) {
- hideGeneratingDots(holder)
+ if (holder.generatingDotsStep != null) return
holder.generatingDots.visibility = View.VISIBLE
val dotStates = arrayOf(".", "..", "...")
var currentIndex = 0
@@ -292,6 +293,7 @@ class ChatAdapter(
}
}
holder.generatingDotsStep = step
+ animatingHolders.add(holder)
holder.generatingDots.post(step)
}
@@ -304,15 +306,29 @@ class ChatAdapter(
private fun hideGeneratingDots(holder: DefaultMessageViewHolder) {
holder.generatingDotsStep?.let { holder.generatingDots.removeCallbacks(it) }
holder.generatingDotsStep = null
+ animatingHolders.remove(holder)
holder.generatingDots.visibility = View.GONE
}
- /** Stops the dots animation of a row leaving the screen, so its step can't outlive the view. */
+ /**
+ * Stop every live "…" animation. Call from the host fragment's `onDestroyView`:
+ * a message still streaming when the tab closes never reaches a terminal status
+ * and its holder is never recycled, so nothing else cancels its Runnable.
+ */
+ fun stopAllAnimations() {
+ animatingHolders.toList().forEach { hideGeneratingDots(it) }
+ }
+
override fun onViewRecycled(holder: RecyclerView.ViewHolder) {
super.onViewRecycled(holder)
if (holder is DefaultMessageViewHolder) hideGeneratingDots(holder)
}
+ override fun onDetachedFromRecyclerView(recyclerView: RecyclerView) {
+ super.onDetachedFromRecyclerView(recyclerView)
+ stopAllAnimations()
+ }
+
private fun createPreview(rawText: String): String {
val cleanedText = rawText
.replace(Regex("`{1,3}|\\*{1,2}|_"), "")
@@ -383,7 +399,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 145deb02..c1c93d61 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
@@ -80,9 +80,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()
}
}
}
@@ -247,12 +247,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
}
@@ -271,7 +271,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
}
@@ -282,19 +282,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)
}
}
}
@@ -574,7 +574,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 { _, _ ->
@@ -590,8 +590,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 9dc94021..06bffb10 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
@@ -16,6 +16,7 @@ import com.google.android.material.chip.Chip
import com.google.android.material.snackbar.Snackbar
import com.itsaky.androidide.plugins.PluginContext
import com.itsaky.androidide.plugins.aiassistant.AiAssistantPlugin
+import com.itsaky.androidide.plugins.aiassistant.BuildConfig
import com.itsaky.androidide.plugins.aiassistant.R
import com.itsaky.androidide.plugins.aiassistant.adapters.ChatAdapter
import com.itsaky.androidide.plugins.aiassistant.databinding.FragmentChatBinding
@@ -61,21 +62,6 @@ class ChatFragment : Fragment() {
}
}
- companion object {
- // Test prompt injection (for E2E testing via broadcast receiver)
- @Volatile
- private var pendingTestPrompt: String? = null
-
- fun injectTestPrompt(prompt: String) {
- pendingTestPrompt = prompt
- }
-
- fun getPendingTestPrompt(): String? {
- return pendingTestPrompt?.also { pendingTestPrompt = null }
- }
- }
-
-
/**
* Route inflation through the host so the plugin's views resolve against a Context whose
* Configuration tracks the IDE's day/night setting — this is what lets values-night/ colors
@@ -99,6 +85,9 @@ class ChatFragment : Fragment() {
}
override fun onDestroyView() {
+ if (::chatAdapter.isInitialized) {
+ chatAdapter.stopAllAnimations()
+ }
super.onDestroyView()
viewModel.stopProcessing()
_binding = null
@@ -136,8 +125,12 @@ class ChatFragment : Fragment() {
/**
* Check for test prompt from broadcast receiver and auto-send if present.
* Uses SharedPreferences set by TestBroadcastReceiver for reliable communication.
+ *
+ * Debug builds only. This path auto-drives the agent — which owns file-mutating
+ * tools — without any user gesture, so it must not exist in a released plugin.
*/
private fun injectPendingTestPrompt() {
+ if (!BuildConfig.DEBUG) return
try {
// Check SharedPreferences for pending test prompt (set by TestBroadcastReceiver)
val context = requireContext()
@@ -320,33 +313,33 @@ 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"
+ binding.sendButton.text = getString(R.string.btn_stop)
viewModel.startStateTimer(state)
}
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..627874e9
--- /dev/null
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/AgentLoop.kt
@@ -0,0 +1,236 @@
+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 the transcript so far. Receives the turns
+ * structurally rather than pre-flattened, so a backend that speaks a real chat format can
+ * emit one turn per message; flattening callers can use [renderTranscript].
+ * @param executeTools runs a batch of tool calls.
+ * @param events UI/state callbacks.
+ * @return the run [Result].
+ */
+ suspend fun run(
+ history: MutableList,
+ generate: suspend (turns: List) -> String,
+ executeTools: suspend (List) -> List,
+ events: Events = object : Events {},
+ ): Result {
+ var turn = 0
+ var previousSignature: String? = null
+ var consecutiveRepeats = 0
+ var previousBatchSucceeded = false
+ while (turn < maxIterations) {
+ turn++
+
+ val text = generate(history.toList())
+ 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)
+ }
+ }
+
+ val signature = signatureOf(realCalls)
+ if (signature == previousSignature) {
+ if (previousBatchSucceeded) {
+ return Result(turn, StopReason.COMPLETED)
+ }
+ consecutiveRepeats++
+ if (consecutiveRepeats >= maxConsecutiveRepeats) {
+ events.onRepeatedToolCalls(turn)
+ return Result(turn, StopReason.REPEATED)
+ }
+ } else {
+ consecutiveRepeats = 0
+ }
+ previousSignature = signature
+
+ val results = executeTools(realCalls)
+ previousBatchSucceeded = results.isNotEmpty() && results.all { it.success }
+ 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).
+ *
+ * For backends whose transport carries only a single string. A backend that renders real
+ * conversation turns must be given the [ChatMessage] list instead — flattening a multi-turn
+ * run into one string leaves the assistant's tool calls and the tool results sitting inside
+ * whatever single turn the backend wraps this in.
+ *
+ * @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.
+ *
+ * Each result is wrapped in `` tags. Chat-tuned models are trained to read
+ * tool output inside that delimiter, and it reads as a single token once the backend
+ * tokenizes with special tokens enabled. Handed the same content as bare prose, a small
+ * model tends not to register that the call already ran and re-issues it, which the
+ * [maxConsecutiveRepeats] guard then has to abort.
+ *
+ * @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()
+ 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("\n")
+ .append("[").append(name).append("] ").append(truncate(body)).append("\n")
+ .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/ToolApprovalManager.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolApprovalManager.kt
index 6fa89098..caf6cef8 100644
--- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolApprovalManager.kt
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/ToolApprovalManager.kt
@@ -2,6 +2,9 @@ package com.itsaky.androidide.plugins.aiassistant.tool
import android.util.Log
import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.withTimeoutOrNull
/**
@@ -31,8 +34,16 @@ class ToolApprovalManager {
// Pending approval request
private var pendingApproval: CompletableDeferred? = null
- private var currentApprovalRequest: ApprovalRequest? = null
-
+
+ private val _currentApprovalRequest = MutableStateFlow(null)
+
+ /**
+ * The pending approval request, or null when none is outstanding. A flow rather than a
+ * plain field so the UI is woken when a request appears, instead of polling for one.
+ */
+ val currentApprovalRequest: StateFlow = _currentApprovalRequest.asStateFlow()
+
+
/**
* Check if a tool needs approval and request it if needed.
* @return ApprovalResponse with approved status and optional denial message
@@ -59,7 +70,7 @@ class ToolApprovalManager {
description = handler.description
)
- currentApprovalRequest = request
+ _currentApprovalRequest.value = request
pendingApproval = CompletableDeferred()
Log.d(TAG, "Requesting approval for $toolName (timeout: ${APPROVAL_TIMEOUT_MS}ms)")
@@ -69,7 +80,7 @@ class ToolApprovalManager {
pendingApproval!!.await()
}
- currentApprovalRequest = null
+ _currentApprovalRequest.value = null
pendingApproval = null
// Handle timeout or decision
@@ -101,19 +112,15 @@ class ToolApprovalManager {
}
}
- /**
- * Get the current pending approval request, if any.
- */
- fun getCurrentApprovalRequest(): ApprovalRequest? {
- return currentApprovalRequest
- }
-
/**
* Submit user's approval decision.
*/
fun submitApproval(result: ApprovalResult) {
if (pendingApproval?.isActive == true) {
pendingApproval?.complete(result)
+ // Clear here as well as in ensureApproved(): completing the deferred only resumes
+ // that coroutine on the next dispatch, and the dialog must dismiss immediately.
+ _currentApprovalRequest.value = null
Log.d(TAG, "Approval decision submitted: $result")
}
}
@@ -133,7 +140,7 @@ class ToolApprovalManager {
* Check if there's a pending approval request.
*/
fun hasPendingApproval(): Boolean {
- return currentApprovalRequest != null && pendingApproval?.isActive == true
+ return _currentApprovalRequest.value != null && pendingApproval?.isActive == true
}
/**
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/AddDependencyHandler.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/AddDependencyHandler.kt
index 46d71f03..e48f4d1e 100644
--- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/AddDependencyHandler.kt
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/AddDependencyHandler.kt
@@ -16,6 +16,8 @@ class AddDependencyHandler(
override val description = "Add a Maven dependency to the project build file"
override val requiresApproval = true
+ override val pathArgs = listOf("build_file")
+
override suspend fun execute(args: Map): ToolResult {
val dependency = args["dependency"]?.toString()?.trim()
if (dependency.isNullOrBlank()) {
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 "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/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 ("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/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/tool/handlers/UpdateFileHandler.kt b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/UpdateFileHandler.kt
index 92e5e23c..ae59c6b5 100644
--- a/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/UpdateFileHandler.kt
+++ b/ai-assistant/src/main/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/handlers/UpdateFileHandler.kt
@@ -32,11 +32,14 @@ class UpdateFileHandler(
} else if (!file.isFile) {
ToolResult.failure("Path is not a file: $filePath")
} else {
- // Backup existing content
- val backup = file.readText()
-
- // Write new content
- file.writeText(content)
+ val previous = file.readText()
+ try {
+ file.writeText(content)
+ } catch (e: Exception) {
+ runCatching { file.writeText(previous) }
+ .onFailure { Log.e("UpdateFileHandler", "Could not restore $filePath after a failed write", it) }
+ throw e
+ }
ToolResult.success(
message = "Updated file: $filePath (${content.length} characters)",
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..e0ebd8f6 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,10 +149,26 @@ class ChatViewModel(
private val approvalManager = ToolApprovalManager()
private val toolRouter: ToolRouter
private val executor: Executor
+ private val agentLoop = AgentLoop(terminalTool = RESPOND_TOOL)
val toolExecutionTracker = ToolExecutionTracker()
- private val _pendingApprovalRequest = MutableStateFlow(null)
- val pendingApprovalRequest: StateFlow = _pendingApprovalRequest.asStateFlow()
+ /** 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
+
+ /** The current run's most recent fully-successful tool batch, or null; reset per run. */
+ @Volatile
+ private var lastSucceededCalls: List? = null
+
+ /** The tool awaiting approval, straight from [approvalManager] — no polling in between. */
+ val pendingApprovalRequest: StateFlow = approvalManager.currentApprovalRequest
private var contextFiles = listOf()
@@ -152,17 +205,6 @@ class ChatViewModel(
toolRouter = ToolRouter(handlers)
executor = Executor(toolRouter, approvalManager, toolExecutionTracker)
-
- // Monitor approval requests
- viewModelScope.launch {
- while (true) {
- delay(100) // Poll every 100ms
- val request = approvalManager.getCurrentApprovalRequest()
- if (request != _pendingApprovalRequest.value) {
- _pendingApprovalRequest.value = request
- }
- }
- }
}
fun initializeStorage(context: android.content.Context) {
@@ -196,7 +238,6 @@ class ChatViewModel(
*/
fun submitApproval(result: ApprovalResult) {
approvalManager.submitApproval(result)
- _pendingApprovalRequest.value = null
}
/**
@@ -301,6 +342,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 a block — 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 +369,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 +384,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 +409,53 @@ class ChatViewModel(
}
/**
- * Parse tool calls from text using multi-strategy extraction.
- * Tries: XML tags → JSON objects → Implicit actions
- */
- private fun parseToolCalls(text: String): List {
- return ToolCallExtractor.extractToolCalls(text)
- }
-
- /**
- * Execute tool calls and add results to chat.
+ * 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 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
+
+ lastSucceededCalls = toolCalls.takeIf { results.isNotEmpty() && results.all { r -> r.success } }
+
// 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 +555,296 @@ 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 tracking.
+ lastToolFailedThisRun = false
+ lastSucceededCalls = null
+ 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
- }
+ try {
+ agentLoop.run(
+ history = history,
+ generate = { turns ->
+ withContext(Dispatchers.Main) {
+ _agentState.value = AgentState.Processing(str(R.string.msg_generating))
}
-
- // 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
+ runModelTurn(llmService, turns, 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
+ 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)
+ }
+ }
+ }
- // Remove the agent message
+ /**
+ * Runs one streaming model turn: creates an agent bubble, streams tokens into it,
+ * and suspends until completion. Throws on backend error.
+ *
+ * Sends [turns] structurally to the local backend, which renders one chat turn per message.
+ * Gemini's transport carries a single string, so it keeps the flattened transcript.
+ *
+ * @param llmService the inference service.
+ * @param turns the conversation so far; the last entry is the current user 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,
+ turns: List,
+ 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)
+ }
+
+ val streamCallback = 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 realCalls = toolCalls.filterNot { it.name == RESPOND_TOOL }
+ if (realCalls.isNotEmpty() && realCalls == lastSucceededCalls) {
+ viewModelScope.launch(Dispatchers.Main) {
_messages.value = _messages.value.filter { it.id != agentMessageId }
+ }
+ deferred.complete(response.text)
+ return
+ }
- // 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)
+ 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
+ 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))
+ }
+ }
+
+ try {
+ if (currentBackendId == "gemini") {
+ llmService.generateStreaming(agentLoop.renderTranscript(turns), config, streamCallback)
+ } else {
+ llmService.generateStreamingWithTools(
+ turns.lastOrNull()?.content.orEmpty(),
+ turns.dropLast(1),
+ config,
+ emptyList(),
+ object : LlmInferenceService.ToolStreamCallback {
+ override fun onToken(token: String) = streamCallback.onToken(token)
+ override fun onToolCall(request: LlmInferenceService.ToolCallRequest) = Unit
+ override fun onComplete(response: LlmInferenceService.LlmResponse) =
+ streamCallback.onComplete(response)
+ override fun onError(error: String) = streamCallback.onError(error)
+ }
)
- _messages.value = _messages.value + errorMessage
}
+ } 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 +852,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 +871,7 @@ class ChatViewModel(
_sessions.value = _sessions.value + newSession
_currentSessionId.value = newSession.id
_messages.value = emptyList()
+ _history.value = emptyList()
}
/**
@@ -702,6 +884,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 +905,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 +948,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 a669d8c3..50ef9b78 100644
--- a/ai-assistant/src/main/res/values/strings.xml
+++ b/ai-assistant/src/main/res/values/strings.xml
@@ -123,4 +123,28 @@
Toggle All
Open a project first — context files are picked from the project you have open.
Project directory not found
+
+
+ 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…
+
+
+ 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..3ef3542c
--- /dev/null
+++ b/ai-assistant/src/test/kotlin/com/itsaky/androidide/plugins/aiassistant/tool/AgentLoopTest.kt
@@ -0,0 +1,370 @@
+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 the turns it was given and returns scripted model responses in order. */
+ private class ScriptedModel(private val responses: List) {
+ val turns = mutableListOf>()
+ var calls = 0
+ suspend fun generate(history: List): String {
+ turns += history
+ 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 turn must receive the fed-back tool results (the whole point), and receive
+ // them as their own USER turn rather than folded into the preceding one.
+ val secondTurnInput = model.turns[1]
+ assertEquals(3, secondTurnInput.size)
+ assertEquals(Role.ASSISTANT, secondTurnInput[1].role)
+ assertEquals(Role.USER, secondTurnInput[2].role)
+ assertTrue(
+ "2nd turn should include tool results",
+ secondTurnInput[2].content.contains("") &&
+ secondTurnInput[2].content.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(""))
+ }
+
+ @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 givenAModelThatRepeatsAToolCallThatSucceeded_whenTheLoopRuns_thenItEndsWithoutRunningItAgain() = 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
+
+ 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 }
+ }
+ )
+
+ // The work already succeeded, so a re-request means "done" — not an error.
+ assertTrue(result.completed)
+ assertEquals(AgentLoop.StopReason.COMPLETED, result.reason)
+ assertEquals(2, result.turns)
+ assertEquals(1, toolBatches) // the side effect runs exactly once
+ assertEquals(-1, repeatedTurns) // no "kept requesting the same action" warning
+ }
+
+ @Test
+ fun givenAModelThatRepeatsAToolCallThatFailed_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
+
+ // A failed batch keeps the retry tolerance: one repeat allowed, the second aborts.
+ val result = AgentLoop(maxIterations = 8).run(
+ history = history,
+ generate = model::generate,
+ executeTools = { toolBatches++; listOf(ToolResult.failure("nope")) },
+ 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_whenAFailedCallIsRepeated_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.failure("nope")) }
+ )
+
+ 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("") }
+ 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("") }
+ 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 "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/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 "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/../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 "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/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("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/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("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/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..d25bb977 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..86381579 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..3089b2e6 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;
@@ -302,6 +312,12 @@ Java_android_llama_cpp_LLamaAndroid_new_1context(JNIEnv *env, jobject, jlong jmo
return 0;
}
+ // A fresh context has an empty KV cache, so the prefix record must start empty too.
+ {
+ std::lock_guard lock(g_globals_mutex);
+ g_cached_tokens.clear();
+ }
+
return reinterpret_cast(context);
}
@@ -309,6 +325,11 @@ extern "C"
JNIEXPORT void JNICALL
Java_android_llama_cpp_LLamaAndroid_free_1context(JNIEnv *, jobject, jlong context) {
llama_free(reinterpret_cast(context));
+
+ // g_cached_tokens outlives the context it describes; left stale, the next completion_init()
+ // reuses a prefix this now-empty cache doesn't have and decodes from a truncated context.
+ std::lock_guard lock(g_globals_mutex);
+ g_cached_tokens.clear();
}
extern "C"
@@ -554,6 +575,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 +698,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 +853,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 f0056540..efa1008c 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
@@ -31,7 +31,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)
@@ -454,6 +454,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..a676a171 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
@@ -122,7 +122,11 @@ class LlmInferenceServiceImpl : LlmInferenceService {
override fun onComplete(response: LlmResponse) = callback.onComplete(response)
override fun onError(error: String) = callback.onError(error)
}
- backend.generateStreaming(prompt, config, streamCallback)
+ if (backend is LocalLlmBackend) {
+ backend.generateStreamingWithHistory(history, prompt, config, streamCallback)
+ } else {
+ backend.generateStreaming(prompt, config, streamCallback)
+ }
return
}
@@ -168,5 +172,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 3e83461f..61bc01f0 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,25 +6,35 @@ 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.isActive
+import kotlinx.coroutines.ensureActive
+import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.launch
-import kotlinx.coroutines.runBlocking
-import kotlinx.coroutines.withTimeoutOrNull
+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"
+
/**
* Belt-and-braces guard: `<|im_end|>` is an EOG control token, so the native loop
* normally stops on it by itself. This only matters if a model emits it as plain
@@ -33,22 +43,30 @@ class LocalLlmBackend(private val context: PluginContext) : LlmBackend {
private val CHAT_STOP = listOf("<|im_end|>")
}
- private val llama by lazy { LLamaAndroid.instance() }
+ private val llamaLazy = lazy { LLamaAndroid.instance() }
+ private val llama by llamaLazy
private val scope = CoroutineScope(Dispatchers.IO)
/**
- * Separate from [scope] because [close] cancels [scope] and then has to run the unload —
- * work submitted to a cancelled scope never starts. Cancelled once its teardown job finishes;
- * a `var` because that is terminal and a second [close] must re-create it to do any work.
+ * Owns the teardown coroutine in [close]. Separate from [scope] because that one is cancelled
+ * first (to stop in-flight generation) and so could not run the unload itself; a
+ * [SupervisorJob] keeps the two teardown steps independent. [close] cancels this scope once
+ * the work completes, so nothing outlives the plugin.
*/
- @Volatile private var teardownScope = CoroutineScope(Dispatchers.IO)
+ private val cleanupScope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
+
+ /** Single-flight guard serializing whole generations on the shared native context. */
+ private val generationMutex = Mutex()
- /** The [close] teardown, retained so [awaitClose] can join it. */
- @Volatile private var teardownJob: Job? = null
+ @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"
@@ -67,10 +85,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`.
@@ -276,39 +321,52 @@ 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)
-
- val startTime = System.currentTimeMillis()
-
- // Collect all tokens
- val responseBuilder = StringBuilder()
- var tokenCount = 0
-
- llama.send(
- message = fullPrompt,
- formatChat = true,
- stop = CHAT_STOP,
- clearCache = false
- ).collect { token ->
- responseBuilder.append(token)
- tokenCount++
- }
+ // 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)
+
+ 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 = true,
+ stop = CHAT_STOP,
+ clearCache = false
+ ).collect { token ->
+ // Honor cancellation so Stop frees the run loop early.
+ ensureActive()
+ responseBuilder.append(token)
+ tokenCount++
+ }
- val responseText = responseBuilder.toString()
- context.logger.info("Generated response: ${responseText.take(50)}... ($tokenCount tokens)")
+ val responseText = responseBuilder.toString()
+ context.logger.info("Generated response: ${responseText.take(50)}... ($tokenCount tokens)")
- future.complete(LlmResponse.success(responseText, tokenCount, System.currentTimeMillis() - startTime))
+ 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) {
@@ -323,7 +381,37 @@ class LocalLlmBackend(private val context: PluginContext) : LlmBackend {
override fun generateStreaming(prompt: String, config: LlmConfig, callback: StreamCallback) {
context.logger.info("LocalLlmBackend.generateStreaming() called")
+ streamGeneration(buildPrompt(config.systemPrompt, prompt), config, callback)
+ }
+
+ /**
+ * Streams a reply for a multi-turn conversation, rendering each earlier turn as its own
+ * ChatML turn.
+ *
+ * @param history earlier turns, oldest first, excluding [prompt]
+ * @param prompt the current user turn
+ * @param config sampling settings; [LlmConfig.systemPrompt] becomes the leading system turn
+ * @param callback receives tokens, completion, and errors
+ */
+ fun generateStreamingWithHistory(
+ history: List,
+ prompt: String,
+ config: LlmConfig,
+ callback: StreamCallback
+ ) {
+ context.logger.info("LocalLlmBackend.generateStreamingWithHistory() called with ${history.size} messages")
+ streamGeneration(buildPrompt(config.systemPrompt, prompt, history), config, callback)
+ }
+ /**
+ * Streams a generation over an already-formatted prompt, serialized against every other
+ * generation on the shared native context.
+ *
+ * @param fullPrompt the complete prompt, as built by [buildPrompt]
+ * @param config sampling settings for this request
+ * @param callback receives tokens, completion, and errors
+ */
+ private fun streamGeneration(fullPrompt: String, config: LlmConfig, callback: StreamCallback) {
// Check if model is configured
val prefs = try {
val aiAssistantContext = SharedServices.get(PluginContext::class.java)
@@ -339,38 +427,51 @@ 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 = buildPrompt(config.systemPrompt, prompt)
-
- val startTime = System.currentTimeMillis()
- var tokenCount = 0
- val responseBuilder = StringBuilder()
-
- llama.send(
- message = fullPrompt,
- formatChat = true,
- stop = CHAT_STOP,
- 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)
+
+ 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 = true,
+ stop = CHAT_STOP,
+ // Keep the KV cache so the native layer reuses the common prefix (system prompt).
+ 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) {
@@ -381,6 +482,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,
@@ -407,46 +519,29 @@ class LocalLlmBackend(private val context: PluginContext) : LlmBackend {
* on Main. Cancel generation, then unload on a background thread, then stop
* the Llm-RunLoop thread so it doesn't outlive the plugin.
*
- * Teardown runs on [teardownScope] rather than a throwaway `CoroutineScope(...)` so the
- * work has an owner: the returned [Job] is retained in [teardownJob], letting a caller
- * observe or await it via [awaitClose] instead of dispose() returning while native work is
- * still in flight with no handle to it, and is cancelled once that job completes.
+ * The teardown runs in [cleanupScope] rather than a floating `CoroutineScope(...)`: the scope
+ * is owned by this object and cancelled as soon as the work finishes, so there is no orphan
+ * job left behind. It cannot be joined — dispose() may be on the main thread and unload()
+ * blocks on the native run loop — so deterministic teardown is the strongest guarantee here.
*/
fun close() {
scope.cancel()
- // A prior close() cancelled the scope on completion, and a dead scope never starts work.
- val teardown = teardownScope.takeIf { it.isActive }
- ?: CoroutineScope(Dispatchers.IO).also { teardownScope = it }
- teardownJob = teardown.launch {
+ val cleanup = cleanupScope.launch {
+ if (!llamaLazy.isInitialized()) {
+ return@launch
+ }
try {
unloadModelInternal()
- } catch (e: Exception) {
- context.logger.error("Error unloading model during close()", e)
+ } catch (t: Throwable) {
+ context.logger.error("Error unloading model during close()", t)
} finally {
- // The Llm-RunLoop executor thread exists even if no model was
- // ever loaded; shut it down unconditionally so the plugin's
- // classloader can be collected after unload.
- llama.shutdown()
+ try {
+ llama.shutdown()
+ } catch (t: Throwable) {
+ context.logger.error("Error shutting down Llm-RunLoop during close()", t)
+ }
}
}
- // Cancel the captured scope, not the field: a later close() may have replaced it.
- teardownJob?.invokeOnCompletion { teardown.cancel() }
- }
-
- /**
- * Block until the [close] teardown finishes, at most [timeoutMs].
- *
- * Intended for tests and for a host that wants unload to have completed before it drops the
- * plugin's classloader. Never call from the main thread — that is the deadlock [close] exists
- * to avoid.
- *
- * @param timeoutMs how long to wait before giving up
- * @return true if teardown finished (or never started), false if it was still running
- */
- fun awaitClose(timeoutMs: Long = 10_000): Boolean {
- val job = teardownJob ?: return true
- return runBlocking {
- withTimeoutOrNull(timeoutMs) { job.join() } != null
- }
+ cleanup.invokeOnCompletion { cleanupScope.cancel() }
}
}