feat: add android device control tool via WebSocket
Implement a unified architecture where the Go backend sends tool_request messages over WebSocket to the Android app, which processes them using Android APIs (PackageManager, AccessibilityService, Context) and returns tool_response messages. This avoids Termux `am` command limitations and eliminates shell injection risks. 10 actions: list_apps, app_info, launch_app, current_activity, tap, swipe, text, keyevent, broadcast, intent. Only active in assistant mode (client_type="assistant"), disabled by default. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
9693f9fc7c
commit
0b474b4a4d
16 changed files with 910 additions and 10 deletions
|
|
@ -75,10 +75,12 @@ class AssistantService : LifecycleService(), SavedStateRegistryOwner {
|
|||
private val httpClient: HttpClient by inject()
|
||||
private val ttsSettingsRepo: TtsSettingsRepository by inject()
|
||||
private val screenshotSource: ScreenshotSource by inject()
|
||||
private val deviceController: DeviceController by inject()
|
||||
|
||||
private lateinit var serviceScope: CoroutineScope
|
||||
private lateinit var connection: AssistantConnection
|
||||
private lateinit var assistantManager: AssistantManager
|
||||
private lateinit var toolRequestHandler: ToolRequestHandler
|
||||
private lateinit var ttsWrapper: TextToSpeechWrapper
|
||||
private lateinit var sttWrapper: SpeechRecognizerWrapper
|
||||
private lateinit var cameraCaptureManager: CameraCaptureManager
|
||||
|
|
@ -112,6 +114,20 @@ class AssistantService : LifecycleService(), SavedStateRegistryOwner {
|
|||
|
||||
connection = AssistantConnectionImpl(httpClient)
|
||||
|
||||
toolRequestHandler = ToolRequestHandler(
|
||||
context = applicationContext,
|
||||
deviceController = deviceController,
|
||||
onAccessibilityNeeded = { showAccessibilityGuide = true }
|
||||
)
|
||||
(connection as AssistantConnectionImpl).onToolRequest = { request ->
|
||||
val response = toolRequestHandler.handle(request)
|
||||
if (response.success) {
|
||||
response.result ?: ""
|
||||
} else {
|
||||
response.error ?: "unknown error"
|
||||
}
|
||||
}
|
||||
|
||||
sttWrapper = SpeechRecognizerWrapper(this)
|
||||
ttsWrapper = TextToSpeechWrapper(this, ttsSettingsRepo.ttsConfig)
|
||||
cameraCaptureManager = CameraCaptureManager(this)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,89 @@
|
|||
package io.picoclaw.android.assistant
|
||||
|
||||
import android.accessibilityservice.AccessibilityService
|
||||
import android.accessibilityservice.GestureDescription
|
||||
import android.graphics.Path
|
||||
import android.os.Bundle
|
||||
import android.view.accessibility.AccessibilityNodeInfo
|
||||
import kotlinx.coroutines.suspendCancellableCoroutine
|
||||
import kotlin.coroutines.resume
|
||||
|
||||
class DeviceController {
|
||||
|
||||
@Volatile
|
||||
private var service: AccessibilityService? = null
|
||||
|
||||
val isAvailable: Boolean get() = service != null
|
||||
|
||||
fun setService(s: AccessibilityService) {
|
||||
service = s
|
||||
}
|
||||
|
||||
fun clearService() {
|
||||
service = null
|
||||
}
|
||||
|
||||
suspend fun tap(x: Float, y: Float): Boolean {
|
||||
val svc = service ?: return false
|
||||
val path = Path().apply { moveTo(x, y) }
|
||||
val stroke = GestureDescription.StrokeDescription(path, 0, 100)
|
||||
val gesture = GestureDescription.Builder().addStroke(stroke).build()
|
||||
return dispatchGesture(svc, gesture)
|
||||
}
|
||||
|
||||
suspend fun swipe(x1: Float, y1: Float, x2: Float, y2: Float, durationMs: Long = 300): Boolean {
|
||||
val svc = service ?: return false
|
||||
val path = Path().apply {
|
||||
moveTo(x1, y1)
|
||||
lineTo(x2, y2)
|
||||
}
|
||||
val stroke = GestureDescription.StrokeDescription(path, 0, durationMs)
|
||||
val gesture = GestureDescription.Builder().addStroke(stroke).build()
|
||||
return dispatchGesture(svc, gesture)
|
||||
}
|
||||
|
||||
fun pressBack(): Boolean {
|
||||
return service?.performGlobalAction(AccessibilityService.GLOBAL_ACTION_BACK) ?: false
|
||||
}
|
||||
|
||||
fun pressHome(): Boolean {
|
||||
return service?.performGlobalAction(AccessibilityService.GLOBAL_ACTION_HOME) ?: false
|
||||
}
|
||||
|
||||
fun pressRecents(): Boolean {
|
||||
return service?.performGlobalAction(AccessibilityService.GLOBAL_ACTION_RECENTS) ?: false
|
||||
}
|
||||
|
||||
fun getCurrentPackage(): String? {
|
||||
return service?.rootInActiveWindow?.packageName?.toString()
|
||||
}
|
||||
|
||||
suspend fun inputText(text: String): Boolean {
|
||||
val svc = service ?: return false
|
||||
val focusedNode = svc.rootInActiveWindow?.findFocus(AccessibilityNodeInfo.FOCUS_INPUT)
|
||||
?: return false
|
||||
val args = Bundle().apply {
|
||||
putCharSequence(AccessibilityNodeInfo.ACTION_ARGUMENT_SET_TEXT_CHARSEQUENCE, text)
|
||||
}
|
||||
return focusedNode.performAction(AccessibilityNodeInfo.ACTION_SET_TEXT, args)
|
||||
}
|
||||
|
||||
private suspend fun dispatchGesture(
|
||||
svc: AccessibilityService,
|
||||
gesture: GestureDescription
|
||||
): Boolean = suspendCancellableCoroutine { cont ->
|
||||
svc.dispatchGesture(
|
||||
gesture,
|
||||
object : AccessibilityService.GestureResultCallback() {
|
||||
override fun onCompleted(gestureDescription: GestureDescription?) {
|
||||
cont.resume(true)
|
||||
}
|
||||
|
||||
override fun onCancelled(gestureDescription: GestureDescription?) {
|
||||
cont.resume(false)
|
||||
}
|
||||
},
|
||||
null
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
@ -7,19 +7,22 @@ import org.koin.android.ext.android.inject
|
|||
class PicoClawAccessibilityService : AccessibilityService() {
|
||||
|
||||
private val screenshotSource: AccessibilityScreenshotSource by inject()
|
||||
private val deviceController: DeviceController by inject()
|
||||
|
||||
override fun onServiceConnected() {
|
||||
super.onServiceConnected()
|
||||
screenshotSource.setService(this)
|
||||
deviceController.setService(this)
|
||||
}
|
||||
|
||||
override fun onDestroy() {
|
||||
screenshotSource.clearService()
|
||||
deviceController.clearService()
|
||||
super.onDestroy()
|
||||
}
|
||||
|
||||
override fun onAccessibilityEvent(event: AccessibilityEvent?) {
|
||||
// No-op: used only for screenshot capture
|
||||
// No-op
|
||||
}
|
||||
|
||||
override fun onInterrupt() {
|
||||
|
|
|
|||
|
|
@ -0,0 +1,260 @@
|
|||
package io.picoclaw.android.assistant
|
||||
|
||||
import android.content.Context
|
||||
import android.content.Intent
|
||||
import android.content.pm.ApplicationInfo
|
||||
import android.content.pm.PackageManager
|
||||
import android.net.Uri
|
||||
import android.os.Build
|
||||
import android.util.Log
|
||||
import io.picoclaw.android.core.data.remote.dto.ToolRequest
|
||||
import io.picoclaw.android.core.data.remote.dto.ToolResponse
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.JsonPrimitive
|
||||
import kotlinx.serialization.json.contentOrNull
|
||||
import kotlinx.serialization.json.doubleOrNull
|
||||
import kotlinx.serialization.json.jsonPrimitive
|
||||
import kotlinx.serialization.json.longOrNull
|
||||
|
||||
class ToolRequestHandler(
|
||||
private val context: Context,
|
||||
private val deviceController: DeviceController,
|
||||
private val onAccessibilityNeeded: () -> Unit
|
||||
) {
|
||||
|
||||
suspend fun handle(request: ToolRequest): ToolResponse {
|
||||
return try {
|
||||
when (request.action) {
|
||||
"list_apps" -> handleListApps(request)
|
||||
"app_info" -> handleAppInfo(request)
|
||||
"launch_app" -> handleLaunchApp(request)
|
||||
"current_activity" -> handleCurrentActivity(request)
|
||||
"tap" -> handleTap(request)
|
||||
"swipe" -> handleSwipe(request)
|
||||
"text" -> handleText(request)
|
||||
"keyevent" -> handleKeyEvent(request)
|
||||
"broadcast" -> handleBroadcast(request)
|
||||
"intent" -> handleIntent(request)
|
||||
else -> ToolResponse(
|
||||
requestId = request.requestId,
|
||||
success = false,
|
||||
error = "Unknown action: ${request.action}"
|
||||
)
|
||||
}
|
||||
} catch (e: Exception) {
|
||||
Log.e(TAG, "Error handling tool request: ${request.action}", e)
|
||||
ToolResponse(
|
||||
requestId = request.requestId,
|
||||
success = false,
|
||||
error = "Error: ${e.message}"
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
private fun requireAccessibility(request: ToolRequest): ToolResponse? {
|
||||
if (!deviceController.isAvailable) {
|
||||
onAccessibilityNeeded()
|
||||
return ToolResponse(
|
||||
requestId = request.requestId,
|
||||
success = false,
|
||||
error = "accessibility_required"
|
||||
)
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
private fun handleListApps(request: ToolRequest): ToolResponse {
|
||||
val pm = context.packageManager
|
||||
val apps = pm.getInstalledApplications(PackageManager.GET_META_DATA)
|
||||
.filter { pm.getLaunchIntentForPackage(it.packageName) != null }
|
||||
.map { app ->
|
||||
val label = pm.getApplicationLabel(app).toString()
|
||||
"${label} (${app.packageName})"
|
||||
}
|
||||
.sorted()
|
||||
|
||||
return ToolResponse(
|
||||
requestId = request.requestId,
|
||||
success = true,
|
||||
result = "Installed apps (${apps.size}):\n${apps.joinToString("\n")}"
|
||||
)
|
||||
}
|
||||
|
||||
private fun handleAppInfo(request: ToolRequest): ToolResponse {
|
||||
val packageName = request.params?.get("package_name")?.jsonPrimitive?.contentOrNull
|
||||
?: return ToolResponse(request.requestId, false, error = "package_name required")
|
||||
|
||||
val pm = context.packageManager
|
||||
return try {
|
||||
val info = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.TIRAMISU) {
|
||||
pm.getPackageInfo(packageName, PackageManager.PackageInfoFlags.of(0))
|
||||
} else {
|
||||
@Suppress("DEPRECATION")
|
||||
pm.getPackageInfo(packageName, 0)
|
||||
}
|
||||
val appInfo = info.applicationInfo
|
||||
val label = appInfo?.let { pm.getApplicationLabel(it).toString() } ?: packageName
|
||||
val isSystem = appInfo?.flags?.and(ApplicationInfo.FLAG_SYSTEM) != 0
|
||||
|
||||
val sb = StringBuilder()
|
||||
sb.appendLine("App: $label")
|
||||
sb.appendLine("Package: $packageName")
|
||||
sb.appendLine("Version: ${info.versionName ?: "unknown"}")
|
||||
sb.appendLine("System app: $isSystem")
|
||||
if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.P) {
|
||||
sb.appendLine("Version code: ${info.longVersionCode}")
|
||||
}
|
||||
|
||||
ToolResponse(request.requestId, true, result = sb.toString())
|
||||
} catch (e: PackageManager.NameNotFoundException) {
|
||||
ToolResponse(request.requestId, false, error = "Package not found: $packageName")
|
||||
}
|
||||
}
|
||||
|
||||
private fun handleLaunchApp(request: ToolRequest): ToolResponse {
|
||||
val packageName = request.params?.get("package_name")?.jsonPrimitive?.contentOrNull
|
||||
?: return ToolResponse(request.requestId, false, error = "package_name required")
|
||||
|
||||
val intent = context.packageManager.getLaunchIntentForPackage(packageName)
|
||||
?: return ToolResponse(request.requestId, false, error = "No launch intent for $packageName")
|
||||
|
||||
intent.addFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
|
||||
context.startActivity(intent)
|
||||
return ToolResponse(request.requestId, true, result = "Launched $packageName")
|
||||
}
|
||||
|
||||
private suspend fun handleCurrentActivity(request: ToolRequest): ToolResponse {
|
||||
requireAccessibility(request)?.let { return it }
|
||||
|
||||
val pkg = deviceController.getCurrentPackage()
|
||||
?: return ToolResponse(request.requestId, false, error = "Could not get current activity")
|
||||
|
||||
return ToolResponse(request.requestId, true, result = "Current package: $pkg")
|
||||
}
|
||||
|
||||
private suspend fun handleTap(request: ToolRequest): ToolResponse {
|
||||
requireAccessibility(request)?.let { return it }
|
||||
|
||||
val x = request.params?.get("x")?.jsonPrimitive?.doubleOrNull?.toFloat()
|
||||
?: return ToolResponse(request.requestId, false, error = "x coordinate required")
|
||||
val y = request.params?.get("y")?.jsonPrimitive?.doubleOrNull?.toFloat()
|
||||
?: return ToolResponse(request.requestId, false, error = "y coordinate required")
|
||||
|
||||
val success = deviceController.tap(x, y)
|
||||
return ToolResponse(
|
||||
request.requestId, success,
|
||||
result = if (success) "Tapped at ($x, $y)" else null,
|
||||
error = if (!success) "Tap failed" else null
|
||||
)
|
||||
}
|
||||
|
||||
private suspend fun handleSwipe(request: ToolRequest): ToolResponse {
|
||||
requireAccessibility(request)?.let { return it }
|
||||
|
||||
val x = request.params?.get("x")?.jsonPrimitive?.doubleOrNull?.toFloat()
|
||||
?: return ToolResponse(request.requestId, false, error = "x coordinate required")
|
||||
val y = request.params?.get("y")?.jsonPrimitive?.doubleOrNull?.toFloat()
|
||||
?: return ToolResponse(request.requestId, false, error = "y coordinate required")
|
||||
val x2 = request.params?.get("x2")?.jsonPrimitive?.doubleOrNull?.toFloat()
|
||||
?: return ToolResponse(request.requestId, false, error = "x2 coordinate required")
|
||||
val y2 = request.params?.get("y2")?.jsonPrimitive?.doubleOrNull?.toFloat()
|
||||
?: return ToolResponse(request.requestId, false, error = "y2 coordinate required")
|
||||
val durationMs = request.params?.get("duration_ms")?.jsonPrimitive?.longOrNull ?: 300L
|
||||
|
||||
val success = deviceController.swipe(x, y, x2, y2, durationMs)
|
||||
return ToolResponse(
|
||||
request.requestId, success,
|
||||
result = if (success) "Swiped from ($x,$y) to ($x2,$y2)" else null,
|
||||
error = if (!success) "Swipe failed" else null
|
||||
)
|
||||
}
|
||||
|
||||
private suspend fun handleText(request: ToolRequest): ToolResponse {
|
||||
requireAccessibility(request)?.let { return it }
|
||||
|
||||
val text = request.params?.get("text")?.jsonPrimitive?.contentOrNull
|
||||
?: return ToolResponse(request.requestId, false, error = "text required")
|
||||
|
||||
val success = deviceController.inputText(text)
|
||||
return ToolResponse(
|
||||
request.requestId, success,
|
||||
result = if (success) "Text input: $text" else null,
|
||||
error = if (!success) "Text input failed (no focused input field?)" else null
|
||||
)
|
||||
}
|
||||
|
||||
private fun handleKeyEvent(request: ToolRequest): ToolResponse {
|
||||
if (!deviceController.isAvailable) {
|
||||
onAccessibilityNeeded()
|
||||
return ToolResponse(request.requestId, false, error = "accessibility_required")
|
||||
}
|
||||
|
||||
val key = request.params?.get("key")?.jsonPrimitive?.contentOrNull
|
||||
?: return ToolResponse(request.requestId, false, error = "key required")
|
||||
|
||||
val success = when (key) {
|
||||
"back" -> deviceController.pressBack()
|
||||
"home" -> deviceController.pressHome()
|
||||
"recents" -> deviceController.pressRecents()
|
||||
else -> return ToolResponse(request.requestId, false, error = "Unknown key: $key")
|
||||
}
|
||||
return ToolResponse(
|
||||
request.requestId, success,
|
||||
result = if (success) "Key pressed: $key" else null,
|
||||
error = if (!success) "Key event failed" else null
|
||||
)
|
||||
}
|
||||
|
||||
private fun handleBroadcast(request: ToolRequest): ToolResponse {
|
||||
val action = request.params?.get("intent_action")?.jsonPrimitive?.contentOrNull
|
||||
?: return ToolResponse(request.requestId, false, error = "intent_action required")
|
||||
|
||||
val intent = Intent(action)
|
||||
applyExtras(intent, request.params)
|
||||
context.sendBroadcast(intent)
|
||||
return ToolResponse(request.requestId, true, result = "Broadcast sent: $action")
|
||||
}
|
||||
|
||||
private fun handleIntent(request: ToolRequest): ToolResponse {
|
||||
val action = request.params?.get("intent_action")?.jsonPrimitive?.contentOrNull
|
||||
?: return ToolResponse(request.requestId, false, error = "intent_action required")
|
||||
|
||||
val intent = Intent(action).addFlags(Intent.FLAG_ACTIVITY_NEW_TASK)
|
||||
|
||||
request.params?.get("intent_data")?.jsonPrimitive?.contentOrNull?.let {
|
||||
intent.data = Uri.parse(it)
|
||||
}
|
||||
request.params?.get("intent_package")?.jsonPrimitive?.contentOrNull?.let {
|
||||
intent.setPackage(it)
|
||||
}
|
||||
request.params?.get("intent_type")?.jsonPrimitive?.contentOrNull?.let {
|
||||
intent.type = it
|
||||
}
|
||||
applyExtras(intent, request.params)
|
||||
|
||||
return try {
|
||||
context.startActivity(intent)
|
||||
ToolResponse(request.requestId, true, result = "Intent started: $action")
|
||||
} catch (e: Exception) {
|
||||
ToolResponse(request.requestId, false, error = "Failed to start intent: ${e.message}")
|
||||
}
|
||||
}
|
||||
|
||||
private fun applyExtras(intent: Intent, params: JsonObject?) {
|
||||
val extras = params?.get("intent_extras") as? JsonObject ?: return
|
||||
for ((key, value) in extras) {
|
||||
val prim = value as? JsonPrimitive ?: continue
|
||||
when {
|
||||
prim.isString -> intent.putExtra(key, prim.content)
|
||||
prim.content.toBooleanStrictOrNull() != null ->
|
||||
intent.putExtra(key, prim.content.toBooleanStrict())
|
||||
prim.longOrNull != null -> intent.putExtra(key, prim.longOrNull!!)
|
||||
prim.doubleOrNull != null -> intent.putExtra(key, prim.doubleOrNull!!)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
companion object {
|
||||
private const val TAG = "ToolRequestHandler"
|
||||
}
|
||||
}
|
||||
|
|
@ -6,6 +6,7 @@ import io.ktor.client.engine.okhttp.OkHttp
|
|||
import io.ktor.client.plugins.websocket.WebSockets
|
||||
import io.picoclaw.android.core.data.local.AppDatabase
|
||||
import io.picoclaw.android.assistant.AccessibilityScreenshotSource
|
||||
import io.picoclaw.android.assistant.DeviceController
|
||||
import io.picoclaw.android.core.data.local.ImageFileStorage
|
||||
import io.picoclaw.android.feature.chat.voice.ScreenshotSource
|
||||
import io.picoclaw.android.core.data.remote.WebSocketClient
|
||||
|
|
@ -96,6 +97,9 @@ val appModule = module {
|
|||
single { AccessibilityScreenshotSource() }
|
||||
single<ScreenshotSource> { get<AccessibilityScreenshotSource>() }
|
||||
|
||||
// Device Controller (for Android tool)
|
||||
single { DeviceController() }
|
||||
|
||||
// Voice
|
||||
factory { SpeechRecognizerWrapper(androidContext()) }
|
||||
single { TextToSpeechWrapper(androidContext(), get<TtsSettingsRepository>().ttsConfig) }
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<accessibility-service xmlns:android="http://schemas.android.com/apk/res/android"
|
||||
android:canTakeScreenshot="true"
|
||||
android:canRetrieveWindowContent="false"
|
||||
android:canRetrieveWindowContent="true"
|
||||
android:canPerformGestures="true"
|
||||
android:accessibilityFlags="flagReportViewIds"
|
||||
android:description="@string/accessibility_service_description" />
|
||||
|
|
|
|||
|
|
@ -0,0 +1,12 @@
|
|||
package io.picoclaw.android.core.data.remote.dto
|
||||
|
||||
import kotlinx.serialization.SerialName
|
||||
import kotlinx.serialization.Serializable
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
|
||||
@Serializable
|
||||
data class ToolRequest(
|
||||
@SerialName("request_id") val requestId: String,
|
||||
val action: String,
|
||||
val params: JsonObject? = null
|
||||
)
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
package io.picoclaw.android.core.data.remote.dto
|
||||
|
||||
import kotlinx.serialization.SerialName
|
||||
import kotlinx.serialization.Serializable
|
||||
|
||||
@Serializable
|
||||
data class ToolResponse(
|
||||
@SerialName("request_id") val requestId: String,
|
||||
val success: Boolean,
|
||||
val result: String? = null,
|
||||
val error: String? = null
|
||||
)
|
||||
|
|
@ -8,5 +8,7 @@ data class WsIncoming(
|
|||
val content: String,
|
||||
@SerialName("sender_id") val senderId: String? = null,
|
||||
val images: List<String>? = null,
|
||||
@SerialName("input_mode") val inputMode: String? = null
|
||||
@SerialName("input_mode") val inputMode: String? = null,
|
||||
val type: String? = null,
|
||||
@SerialName("request_id") val requestId: String? = null
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
package io.picoclaw.android.core.data.repository
|
||||
|
||||
import android.util.Log
|
||||
import io.ktor.client.HttpClient
|
||||
import io.picoclaw.android.core.data.remote.WebSocketClient
|
||||
import io.picoclaw.android.core.data.remote.dto.ToolRequest
|
||||
import io.picoclaw.android.core.data.remote.dto.WsIncoming
|
||||
import io.picoclaw.android.core.domain.model.AssistantMessage
|
||||
import io.picoclaw.android.core.domain.model.ConnectionState
|
||||
|
|
@ -17,8 +19,11 @@ import kotlinx.coroutines.flow.StateFlow
|
|||
import kotlinx.coroutines.flow.asSharedFlow
|
||||
import kotlinx.coroutines.flow.asStateFlow
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.serialization.json.Json
|
||||
import java.util.UUID
|
||||
|
||||
typealias ToolRequestCallback = suspend (ToolRequest) -> String
|
||||
|
||||
class AssistantConnectionImpl(
|
||||
private val httpClient: HttpClient
|
||||
) : AssistantConnection {
|
||||
|
|
@ -26,6 +31,7 @@ class AssistantConnectionImpl(
|
|||
private val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)
|
||||
private val clientId = UUID.randomUUID().toString()
|
||||
private val wsClient = WebSocketClient(httpClient, scope, clientId, "assistant")
|
||||
private val json = Json { ignoreUnknownKeys = true }
|
||||
|
||||
private val _messages = MutableSharedFlow<AssistantMessage>(extraBufferCapacity = 64)
|
||||
override val messages: SharedFlow<AssistantMessage> = _messages.asSharedFlow()
|
||||
|
|
@ -35,12 +41,15 @@ class AssistantConnectionImpl(
|
|||
|
||||
override val connectionState: StateFlow<ConnectionState> = wsClient.connectionState
|
||||
|
||||
var onToolRequest: ToolRequestCallback? = null
|
||||
|
||||
init {
|
||||
scope.launch {
|
||||
wsClient.incomingMessages.collect { dto ->
|
||||
when (dto.type) {
|
||||
"status" -> _statusText.value = dto.content
|
||||
"status_end" -> _statusText.value = null
|
||||
"tool_request" -> handleToolRequest(dto.content)
|
||||
else -> {
|
||||
_statusText.value = null
|
||||
_messages.emit(AssistantMessage(content = dto.content, type = dto.type))
|
||||
|
|
@ -50,6 +59,29 @@ class AssistantConnectionImpl(
|
|||
}
|
||||
}
|
||||
|
||||
private fun handleToolRequest(content: String) {
|
||||
scope.launch {
|
||||
try {
|
||||
val request = json.decodeFromString<ToolRequest>(content)
|
||||
val callback = onToolRequest
|
||||
val resultContent = if (callback != null) {
|
||||
callback(request)
|
||||
} else {
|
||||
"error: tool request handler not configured"
|
||||
}
|
||||
|
||||
val response = WsIncoming(
|
||||
content = resultContent,
|
||||
type = "tool_response",
|
||||
requestId = request.requestId
|
||||
)
|
||||
wsClient.send(response)
|
||||
} catch (e: Exception) {
|
||||
Log.e(TAG, "Failed to handle tool request", e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override fun connect(wsUrl: String) {
|
||||
wsClient.wsUrl = wsUrl
|
||||
wsClient.connect()
|
||||
|
|
@ -68,4 +100,8 @@ class AssistantConnectionImpl(
|
|||
)
|
||||
wsClient.send(dto)
|
||||
}
|
||||
|
||||
companion object {
|
||||
private const val TAG = "AssistantConnectionImpl"
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -107,6 +107,21 @@ func createToolRegistry(workspace string, restrict bool, cfg *config.Config, msg
|
|||
registry.Register(tools.NewSPITool())
|
||||
}
|
||||
|
||||
// Android device control tool (assistant mode only)
|
||||
if cfg.Tools.Android.Enabled {
|
||||
androidTool := tools.NewAndroidTool()
|
||||
androidTool.SetSendCallback(func(channel, chatID, content, msgType string) error {
|
||||
msgBus.PublishOutbound(bus.OutboundMessage{
|
||||
Channel: channel,
|
||||
ChatID: chatID,
|
||||
Content: content,
|
||||
Type: msgType,
|
||||
})
|
||||
return nil
|
||||
})
|
||||
registry.Register(androidTool)
|
||||
}
|
||||
|
||||
// Message tool - available to both agent and subagent
|
||||
// Subagent uses it to communicate directly with user
|
||||
messageTool := tools.NewMessageTool()
|
||||
|
|
@ -510,7 +525,7 @@ func (al *AgentLoop) runAgentLoop(ctx context.Context, opts processOptions) (str
|
|||
}
|
||||
|
||||
// 1. Update tool contexts
|
||||
al.updateToolContexts(opts.Channel, opts.ChatID)
|
||||
al.updateToolContexts(opts.Channel, opts.ChatID, opts.Metadata)
|
||||
|
||||
// 2. Build messages (skip history for heartbeat)
|
||||
var history []providers.Message
|
||||
|
|
@ -959,7 +974,7 @@ func (al *AgentLoop) runLLMIteration(ctx context.Context, messages []providers.M
|
|||
}
|
||||
|
||||
// updateToolContexts updates the context for tools that need channel/chatID info.
|
||||
func (al *AgentLoop) updateToolContexts(channel, chatID string) {
|
||||
func (al *AgentLoop) updateToolContexts(channel, chatID string, metadata map[string]string) {
|
||||
// Use ContextualTool interface instead of type assertions
|
||||
if tool, ok := al.tools.Get("message"); ok {
|
||||
if mt, ok := tool.(tools.ContextualTool); ok {
|
||||
|
|
@ -976,6 +991,14 @@ func (al *AgentLoop) updateToolContexts(channel, chatID string) {
|
|||
st.SetContext(channel, chatID)
|
||||
}
|
||||
}
|
||||
if tool, ok := al.tools.Get("android"); ok {
|
||||
if at, ok := tool.(*tools.AndroidTool); ok {
|
||||
at.SetContext(channel, chatID)
|
||||
if metadata != nil {
|
||||
at.SetClientType(metadata["client_type"])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// maybeSummarize triggers summarization if the session history exceeds thresholds.
|
||||
|
|
|
|||
|
|
@ -56,6 +56,8 @@ func statusLabel(toolName string, args map[string]interface{}) string {
|
|||
return fmt.Sprintf("サブタスク実行中...(%s)", truncLabel(l, 20))
|
||||
}
|
||||
return "サブタスク実行中..."
|
||||
case "android":
|
||||
return androidStatusLabel(args)
|
||||
case "mcp":
|
||||
return mcpStatusLabel(args)
|
||||
case "i2c":
|
||||
|
|
@ -116,6 +118,42 @@ func cronStatusLabel(args map[string]interface{}) string {
|
|||
}
|
||||
}
|
||||
|
||||
func androidStatusLabel(args map[string]interface{}) string {
|
||||
switch strArg(args, "action") {
|
||||
case "list_apps":
|
||||
return "アプリ一覧取得中..."
|
||||
case "app_info":
|
||||
if p := strArg(args, "package_name"); p != "" {
|
||||
return fmt.Sprintf("アプリ情報取得中...(%s)", truncLabel(p, 25))
|
||||
}
|
||||
return "アプリ情報取得中..."
|
||||
case "launch_app":
|
||||
if p := strArg(args, "package_name"); p != "" {
|
||||
return fmt.Sprintf("アプリ起動中...(%s)", truncLabel(p, 25))
|
||||
}
|
||||
return "アプリ起動中..."
|
||||
case "current_activity":
|
||||
return "画面情報取得中..."
|
||||
case "tap":
|
||||
return "タップ中..."
|
||||
case "swipe":
|
||||
return "スワイプ中..."
|
||||
case "text":
|
||||
return "テキスト入力中..."
|
||||
case "keyevent":
|
||||
if k := strArg(args, "key"); k != "" {
|
||||
return fmt.Sprintf("キー操作中...(%s)", k)
|
||||
}
|
||||
return "キー操作中..."
|
||||
case "broadcast":
|
||||
return "ブロードキャスト送信中..."
|
||||
case "intent":
|
||||
return "インテント送信中..."
|
||||
default:
|
||||
return "デバイス操作中..."
|
||||
}
|
||||
}
|
||||
|
||||
func mcpStatusLabel(args map[string]interface{}) string {
|
||||
switch strArg(args, "action") {
|
||||
case "mcp_list":
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import (
|
|||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
)
|
||||
|
||||
// wsIncoming is the JSON message sent from APK to picoclaw.
|
||||
|
|
@ -21,6 +22,8 @@ type wsIncoming struct {
|
|||
SenderID string `json:"sender_id,omitempty"`
|
||||
Images []string `json:"images,omitempty"`
|
||||
InputMode string `json:"input_mode,omitempty"`
|
||||
Type string `json:"type,omitempty"` // "tool_response" for device tool responses
|
||||
RequestID string `json:"request_id,omitempty"` // correlates with tool_request
|
||||
}
|
||||
|
||||
// wsOutgoing is the JSON message sent from picoclaw to APK.
|
||||
|
|
@ -283,6 +286,12 @@ func (c *WebSocketChannel) readPump(conn *websocket.Conn, clientID, chatID, clie
|
|||
continue
|
||||
}
|
||||
|
||||
// Intercept tool_response messages: deliver to ResponseWaiter, skip inbound bus.
|
||||
if incoming.Type == "tool_response" && incoming.RequestID != "" {
|
||||
tools.DeviceResponseWaiter.Deliver(incoming.RequestID, incoming.Content)
|
||||
continue
|
||||
}
|
||||
|
||||
// Use sender_id from message if provided, otherwise use clientID.
|
||||
senderID := clientID
|
||||
if incoming.SenderID != "" {
|
||||
|
|
|
|||
|
|
@ -253,11 +253,16 @@ type MCPServerConfig struct {
|
|||
IdleTimeout int `json:"idle_timeout,omitempty"` // seconds, default 300
|
||||
}
|
||||
|
||||
type AndroidToolsConfig struct {
|
||||
Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_ANDROID_ENABLED"`
|
||||
}
|
||||
|
||||
type ToolsConfig struct {
|
||||
Web WebToolsConfig `json:"web"`
|
||||
Exec ExecToolsConfig `json:"exec"`
|
||||
I2C I2CToolsConfig `json:"i2c"`
|
||||
SPI SPIToolsConfig `json:"spi"`
|
||||
Android AndroidToolsConfig `json:"android"`
|
||||
MCP map[string]MCPServerConfig `json:"mcp,omitempty"`
|
||||
}
|
||||
|
||||
|
|
@ -374,6 +379,9 @@ func DefaultConfig() *Config {
|
|||
SPI: SPIToolsConfig{
|
||||
Enabled: false,
|
||||
},
|
||||
Android: AndroidToolsConfig{
|
||||
Enabled: false,
|
||||
},
|
||||
Web: WebToolsConfig{
|
||||
Brave: BraveConfig{
|
||||
Enabled: false,
|
||||
|
|
|
|||
334
pkg/tools/android.go
Normal file
334
pkg/tools/android.go
Normal file
|
|
@ -0,0 +1,334 @@
|
|||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const androidToolTimeout = 15 * time.Second
|
||||
|
||||
var (
|
||||
packageNameRe = regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9_]*(\.[a-zA-Z][a-zA-Z0-9_]*)*$`)
|
||||
intentActionRe = regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9_.]*$`)
|
||||
)
|
||||
|
||||
// SendCallbackWithType is like SendCallback but includes a message type field.
|
||||
type SendCallbackWithType func(channel, chatID, content, msgType string) error
|
||||
|
||||
// toolRequest is the JSON payload sent to the Android device via WebSocket.
|
||||
type toolRequest struct {
|
||||
RequestID string `json:"request_id"`
|
||||
Action string `json:"action"`
|
||||
Params map[string]interface{} `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
type AndroidTool struct {
|
||||
sendCallback SendCallbackWithType
|
||||
channel string
|
||||
chatID string
|
||||
clientType string
|
||||
}
|
||||
|
||||
func NewAndroidTool() *AndroidTool {
|
||||
return &AndroidTool{}
|
||||
}
|
||||
|
||||
func (t *AndroidTool) Name() string { return "android" }
|
||||
|
||||
func (t *AndroidTool) Description() string {
|
||||
return `Control the Android device. Available actions:
|
||||
- list_apps: List installed apps
|
||||
- app_info: Get app details (requires package_name)
|
||||
- launch_app: Launch an app (requires package_name)
|
||||
- current_activity: Get the currently active app/window
|
||||
- tap: Tap a screen coordinate (requires x, y)
|
||||
- swipe: Swipe between coordinates (requires x, y, x2, y2; optional duration_ms)
|
||||
- text: Input text into the focused field (requires text)
|
||||
- keyevent: Press a key (requires key: back/home/recents)
|
||||
- broadcast: Send a broadcast intent (requires intent_action; optional intent_extras)
|
||||
- intent: Start an activity via intent (requires intent_action; optional intent_data, intent_package, intent_type, intent_extras)
|
||||
|
||||
Only available in assistant mode.`
|
||||
}
|
||||
|
||||
func (t *AndroidTool) Parameters() map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"type": "object",
|
||||
"properties": map[string]interface{}{
|
||||
"action": map[string]interface{}{
|
||||
"type": "string",
|
||||
"enum": []string{
|
||||
"list_apps", "app_info", "launch_app", "current_activity",
|
||||
"tap", "swipe", "text", "keyevent",
|
||||
"broadcast", "intent",
|
||||
},
|
||||
"description": "The device action to perform",
|
||||
},
|
||||
"package_name": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Android package name (for app_info, launch_app)",
|
||||
},
|
||||
"x": map[string]interface{}{
|
||||
"type": "number",
|
||||
"description": "X coordinate (for tap, swipe start)",
|
||||
},
|
||||
"y": map[string]interface{}{
|
||||
"type": "number",
|
||||
"description": "Y coordinate (for tap, swipe start)",
|
||||
},
|
||||
"x2": map[string]interface{}{
|
||||
"type": "number",
|
||||
"description": "End X coordinate (for swipe)",
|
||||
},
|
||||
"y2": map[string]interface{}{
|
||||
"type": "number",
|
||||
"description": "End Y coordinate (for swipe)",
|
||||
},
|
||||
"duration_ms": map[string]interface{}{
|
||||
"type": "integer",
|
||||
"description": "Swipe duration in milliseconds (default 300)",
|
||||
},
|
||||
"text": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Text to input (for text action)",
|
||||
},
|
||||
"key": map[string]interface{}{
|
||||
"type": "string",
|
||||
"enum": []string{"back", "home", "recents"},
|
||||
"description": "Key to press (for keyevent action)",
|
||||
},
|
||||
"intent_action": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Intent action string (for broadcast, intent)",
|
||||
},
|
||||
"intent_data": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Intent data URI (for intent)",
|
||||
},
|
||||
"intent_package": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "Target package for intent (for intent)",
|
||||
},
|
||||
"intent_type": map[string]interface{}{
|
||||
"type": "string",
|
||||
"description": "MIME type for intent (for intent)",
|
||||
},
|
||||
"intent_extras": map[string]interface{}{
|
||||
"type": "object",
|
||||
"description": "Extra key-value pairs for broadcast/intent",
|
||||
},
|
||||
},
|
||||
"required": []string{"action"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *AndroidTool) SetContext(channel, chatID string) {
|
||||
t.channel = channel
|
||||
t.chatID = chatID
|
||||
}
|
||||
|
||||
func (t *AndroidTool) SetClientType(clientType string) {
|
||||
t.clientType = clientType
|
||||
}
|
||||
|
||||
func (t *AndroidTool) SetSendCallback(cb SendCallbackWithType) {
|
||||
t.sendCallback = cb
|
||||
}
|
||||
|
||||
func (t *AndroidTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
|
||||
if t.clientType == "" {
|
||||
return ErrorResult("android tool: not available (no assistant session)")
|
||||
}
|
||||
if t.clientType != "assistant" {
|
||||
return ErrorResult("android tool is only available in assistant mode")
|
||||
}
|
||||
if t.sendCallback == nil {
|
||||
return ErrorResult("android tool: send callback not configured")
|
||||
}
|
||||
if t.channel == "" || t.chatID == "" {
|
||||
return ErrorResult("android tool: no active channel context")
|
||||
}
|
||||
|
||||
action, _ := args["action"].(string)
|
||||
if action == "" {
|
||||
return ErrorResult("action is required")
|
||||
}
|
||||
|
||||
params, err := t.validateAndBuildParams(action, args)
|
||||
if err != nil {
|
||||
return ErrorResult(err.Error())
|
||||
}
|
||||
|
||||
return t.sendAndWait(ctx, action, params)
|
||||
}
|
||||
|
||||
func (t *AndroidTool) validateAndBuildParams(action string, args map[string]interface{}) (map[string]interface{}, error) {
|
||||
params := make(map[string]interface{})
|
||||
|
||||
switch action {
|
||||
case "list_apps":
|
||||
// No params needed
|
||||
|
||||
case "app_info", "launch_app":
|
||||
pkg, _ := args["package_name"].(string)
|
||||
if pkg == "" {
|
||||
return nil, fmt.Errorf("%s requires package_name", action)
|
||||
}
|
||||
if !packageNameRe.MatchString(pkg) {
|
||||
return nil, fmt.Errorf("invalid package_name: %s", pkg)
|
||||
}
|
||||
params["package_name"] = pkg
|
||||
|
||||
case "current_activity":
|
||||
// No params needed
|
||||
|
||||
case "tap":
|
||||
x, xOk := toFloat64(args["x"])
|
||||
y, yOk := toFloat64(args["y"])
|
||||
if !xOk || !yOk {
|
||||
return nil, fmt.Errorf("tap requires x and y coordinates")
|
||||
}
|
||||
params["x"] = x
|
||||
params["y"] = y
|
||||
|
||||
case "swipe":
|
||||
x, xOk := toFloat64(args["x"])
|
||||
y, yOk := toFloat64(args["y"])
|
||||
x2, x2Ok := toFloat64(args["x2"])
|
||||
y2, y2Ok := toFloat64(args["y2"])
|
||||
if !xOk || !yOk || !x2Ok || !y2Ok {
|
||||
return nil, fmt.Errorf("swipe requires x, y, x2, y2 coordinates")
|
||||
}
|
||||
params["x"] = x
|
||||
params["y"] = y
|
||||
params["x2"] = x2
|
||||
params["y2"] = y2
|
||||
if dur, ok := toFloat64(args["duration_ms"]); ok {
|
||||
params["duration_ms"] = int(dur)
|
||||
}
|
||||
|
||||
case "text":
|
||||
text, _ := args["text"].(string)
|
||||
if text == "" {
|
||||
return nil, fmt.Errorf("text action requires text parameter")
|
||||
}
|
||||
params["text"] = text
|
||||
|
||||
case "keyevent":
|
||||
key, _ := args["key"].(string)
|
||||
if key == "" {
|
||||
return nil, fmt.Errorf("keyevent requires key parameter")
|
||||
}
|
||||
switch key {
|
||||
case "back", "home", "recents":
|
||||
// valid
|
||||
default:
|
||||
return nil, fmt.Errorf("invalid key: %s (must be back, home, or recents)", key)
|
||||
}
|
||||
params["key"] = key
|
||||
|
||||
case "broadcast":
|
||||
intentAction, _ := args["intent_action"].(string)
|
||||
if intentAction == "" {
|
||||
return nil, fmt.Errorf("broadcast requires intent_action")
|
||||
}
|
||||
if !intentActionRe.MatchString(intentAction) {
|
||||
return nil, fmt.Errorf("invalid intent_action: %s", intentAction)
|
||||
}
|
||||
params["intent_action"] = intentAction
|
||||
if extras, ok := args["intent_extras"].(map[string]interface{}); ok {
|
||||
params["intent_extras"] = extras
|
||||
}
|
||||
|
||||
case "intent":
|
||||
intentAction, _ := args["intent_action"].(string)
|
||||
if intentAction == "" {
|
||||
return nil, fmt.Errorf("intent requires intent_action")
|
||||
}
|
||||
if !intentActionRe.MatchString(intentAction) {
|
||||
return nil, fmt.Errorf("invalid intent_action: %s", intentAction)
|
||||
}
|
||||
params["intent_action"] = intentAction
|
||||
if data, ok := args["intent_data"].(string); ok && data != "" {
|
||||
params["intent_data"] = data
|
||||
}
|
||||
if pkg, ok := args["intent_package"].(string); ok && pkg != "" {
|
||||
if !packageNameRe.MatchString(pkg) {
|
||||
return nil, fmt.Errorf("invalid intent_package: %s", pkg)
|
||||
}
|
||||
params["intent_package"] = pkg
|
||||
}
|
||||
if mimeType, ok := args["intent_type"].(string); ok && mimeType != "" {
|
||||
params["intent_type"] = mimeType
|
||||
}
|
||||
if extras, ok := args["intent_extras"].(map[string]interface{}); ok {
|
||||
params["intent_extras"] = extras
|
||||
}
|
||||
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown action: %s", action)
|
||||
}
|
||||
|
||||
return params, nil
|
||||
}
|
||||
|
||||
func (t *AndroidTool) sendAndWait(ctx context.Context, action string, params map[string]interface{}) *ToolResult {
|
||||
requestID := uuid.New().String()
|
||||
|
||||
req := toolRequest{
|
||||
RequestID: requestID,
|
||||
Action: action,
|
||||
Params: params,
|
||||
}
|
||||
|
||||
reqJSON, err := json.Marshal(req)
|
||||
if err != nil {
|
||||
return ErrorResult(fmt.Sprintf("failed to marshal tool request: %v", err))
|
||||
}
|
||||
|
||||
// Register waiter before sending to avoid race
|
||||
respCh := DeviceResponseWaiter.Register(requestID)
|
||||
|
||||
if err := t.sendCallback(t.channel, t.chatID, string(reqJSON), "tool_request"); err != nil {
|
||||
DeviceResponseWaiter.Cleanup(requestID)
|
||||
return ErrorResult(fmt.Sprintf("failed to send tool request: %v", err))
|
||||
}
|
||||
|
||||
// Wait for response with timeout
|
||||
select {
|
||||
case content := <-respCh:
|
||||
// Check if the response indicates accessibility_required
|
||||
if strings.HasPrefix(content, "accessibility_required") {
|
||||
return &ToolResult{
|
||||
ForUser: "この機能にはユーザー補助の設定が必要です",
|
||||
ForLLM: "accessibility_required: The accessibility service is not enabled. The settings dialog has been shown to the user. Do not retry automatically - wait for the user to enable the service and try again.",
|
||||
}
|
||||
}
|
||||
return SilentResult(content)
|
||||
case <-time.After(androidToolTimeout):
|
||||
DeviceResponseWaiter.Cleanup(requestID)
|
||||
return ErrorResult("android tool request timed out (15s)")
|
||||
case <-ctx.Done():
|
||||
DeviceResponseWaiter.Cleanup(requestID)
|
||||
return ErrorResult("android tool request cancelled")
|
||||
}
|
||||
}
|
||||
|
||||
// toFloat64 extracts a float64 from an interface{} (handles both float64 and int from JSON).
|
||||
func toFloat64(v interface{}) (float64, bool) {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n, true
|
||||
case int:
|
||||
return float64(n), true
|
||||
case int64:
|
||||
return float64(n), true
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
52
pkg/tools/response_waiter.go
Normal file
52
pkg/tools/response_waiter.go
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
package tools
|
||||
|
||||
import "sync"
|
||||
|
||||
// DeviceResponseWaiter is a package-level shared instance used to synchronize
|
||||
// tool_request/tool_response pairs between the Android tool and WebSocket channel.
|
||||
var DeviceResponseWaiter = NewResponseWaiter()
|
||||
|
||||
// ResponseWaiter manages pending request/response synchronization.
|
||||
// Each request registers a channel by ID; the response is delivered when it arrives.
|
||||
type ResponseWaiter struct {
|
||||
pending map[string]chan string
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func NewResponseWaiter() *ResponseWaiter {
|
||||
return &ResponseWaiter{
|
||||
pending: make(map[string]chan string),
|
||||
}
|
||||
}
|
||||
|
||||
// Register creates a buffered channel for the given request ID and returns it.
|
||||
// The caller should select on this channel with a timeout.
|
||||
func (w *ResponseWaiter) Register(id string) chan string {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
ch := make(chan string, 1)
|
||||
w.pending[id] = ch
|
||||
return ch
|
||||
}
|
||||
|
||||
// Deliver sends the response content to the waiting channel for the given ID.
|
||||
// If no waiter is registered for the ID, the delivery is silently dropped.
|
||||
func (w *ResponseWaiter) Deliver(id, content string) {
|
||||
w.mu.Lock()
|
||||
ch, ok := w.pending[id]
|
||||
if ok {
|
||||
delete(w.pending, id)
|
||||
}
|
||||
w.mu.Unlock()
|
||||
if ok {
|
||||
ch <- content
|
||||
}
|
||||
}
|
||||
|
||||
// Cleanup removes the pending channel for the given ID without delivering.
|
||||
// Used on timeout to prevent memory leaks.
|
||||
func (w *ResponseWaiter) Cleanup(id string) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
delete(w.pending, id)
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue