AaCompletionClient.kt

package afield.runtime.completion

import afield.app.diagnostics.CodeUseTelemetryService
import afield.data.table.div
import afield.data.table.get
import afield.data.table.minus
import afield.data.table.plus
import afield.data.table.rem
import afield.data.table.times
import afield.document.graph.box
import afield.document.graph.intro
import afield.document.graph.introImports
import afield.runtime.execution.KotlinExecution
import afield.runtime.execution.plus
import afield.runtime.kotlin.KotlinEngineConfig
import afield.runtime.kotlin.KotlinInterface2
import afield.runtime.lifecycle.plus
import afield.runtime.lifecycle.rem
import afield.runtime.lifecycle.then
import afield.ui.notation.long
import afield.ui.notation.with
import afield.ui.output.logs

import afield.app.protocol.put
import afield.app.protocol.send
import afield.app.settings.PersonalLayer
import afield.audio.synth.*
import afield.document.graph.preamble
import afield.document.workspace.metadata
import afield.graphics.vulkan.*
import afield.media.stream.*
import afield.media.video.*
import afield.runtime.kotlin.CompletionResult
import afield.runtime.kotlin.CompletionTokenBounds
import afield.runtime.lifecycle.map
import afield.ui.output.templates.filter
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json
import java.io.File
import java.net.URI
import java.net.http.HttpClient
import java.net.http.HttpRequest
import java.net.http.HttpResponse
import java.nio.charset.StandardCharsets
import java.time.Duration
import kotlin.script.experimental.api.SourceCodeCompletionVariant

    @Serializable
    private data class Req(
        val snippet: String,
        val cursor: Int,
        val preamble: String = "",
        /** File-scope declarations supplied by a synthetic language/runtime profile. */
        val declarations: String = "",
        val imports: List<String> = emptyList(),
        /** The REPL's own imports (Field's [intro]) — what a box gets for free, which a headless
         *  sidecar cannot see. Supersedes the sidecar's built-in fallback list. */
        val replImports: List<String> = emptyList(),
        val implicitReceivers: List<String> = emptyList(),
        val classpathRoots: List<String> = emptyList(),
        /** True for a KotlinJS box → the sidecar filters JVM-only stdlib members (no JS equivalent). */
        val jsTarget: Boolean = false,
        /**
         * How many candidates to ship. RANKING IS CLIENT-SIDE (`fieldSuggest.js` orders by fuzzy match,
         * category, recency and context), so a low limit here does not trim the tail of a ranked list — it
         * deletes rows the ranker never sees, from a list the sidecar sorts BY NAME. At 200 that meant `s.`
         * on a String stopped at `orEmpty(`: no `substring`, no `split`, no `trim`, no `toInt`, no
         * `uppercase` in the panel, and 31 of the caret sweep's misses were exactly this (their note reads
         * "200 candidates"). The candidates are already computed when the cap applies, so raising it costs
         * serialization only — measured: `String.` 200 → 299 rows, 43KB → 63KB, 13ms either way.
         *
         * 1200 bounds a pathological receiver rather than budgeting a normal one: the richest real receiver
         * measured (`Vec3`, whose generated swizzles alone are ~978 extensions) yields 337 rows after
         * de-duplication, and the sidecar's own ceiling is 2000.
         */
        val limit: Int = 1200,
    )

    @Serializable
    private data class Cand(
        val text: String,
        val displayText: String,
        val tail: String,
        val icon: String,
        val receiverType: String = "",
        val receiverQualifiedName: String = "",
        /** The candidate's type fits the type expected at the caret (AA's `expectedType`). */
        val expectedTypeMatch: Boolean = false,
        /** See [DECLARED_IN]. */
        val declaredIn: String = "",
        /** See [DECLARATION_PROXIMITY]. */
        val proximity: String = "",
    )

    @Serializable
    private data class Resp(
        val success: Boolean = false,
        val backend: String = "",
        val message: String = "",
        val candidates: List<Cand> = emptyList(),
    )

    private val json = Json { encodeDefaults = true; ignoreUnknownKeys = true }

    private val httpClient: HttpClient by lazy {
        HttpClient.newBuilder().connectTimeout(Duration.ofMillis(1500)).build()
    }

    fun complete(
        snippet: String,
        cursor: Int,
        preamble: String,
        implicitReceivers: List<String>,
        declarations: String = "",
        timeoutMs: Long = 4000,
        // Restrict the analysis classpath (e.g. stdlib-only for KotlinJS boxes). Null = Field's full
        // classpath (the normal rich-completion universe).
        classpathRootsOverride: List<String>? = null,
        // True for a KotlinJS box (compiled to JS, no JVM): the sidecar drops JVM-only stdlib members
        // (`stream`/`spliterator`/`toArray`/…). Default false → regular Kotlin completion is unchanged.
        jsTarget: Boolean = false,
        // The REPL imports the caller's box gets for free (`afield.document.graph.introImports` for a JVM box). Left
        // empty by the KotlinJS callers, whose scope is the stdlib jar and their own declarations —
        // sending them Field's `afield.*` imports would name types their classpath does not contain.
        replImports: List<String> = emptyList(),
    ): Outcome {
        endpoint() ?: return Outcome(null, "no-sidecar-url")
        val roots = classpathRootsOverride?.takeIf { it.isNotEmpty() } ?: classpathRoots
        if (roots.isEmpty()) return Outcome(null, "no-classpath")
        val uri = endpoint()!!
        return runCatching {
            val body = json.encodeToString(
                Req(
                    snippet = snippet,
                    cursor = cursor.coerceIn(0, snippet.length),
                    preamble = preamble,
                    declarations = declarations,
                    replImports = replImports,
                    implicitReceivers = implicitReceivers,
                    classpathRoots = roots,
                    jsTarget = jsTarget,
                )
            )
            val request = HttpRequest.newBuilder(uri)
                .timeout(Duration.ofMillis(timeoutMs.coerceIn(200, 30_000)))
                .header("Content-Type", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(body, StandardCharsets.UTF_8))
                .build()
            val response = httpClient.send(request, HttpResponse.BodyHandlers.ofString(StandardCharsets.UTF_8))
            if (response.statusCode() != 200) return Outcome(null, "http-${response.statusCode()}")
            val resp = json.decodeFromString<Resp>(response.body())
            if (!resp.success) return Outcome(null, "server-fail: ${resp.message.take(80)}")
            // Safety net for AA-as-default: an empty AA result (e.g. a preamble/receiver form the
            // headless session can't model, or the session still warming up) falls through to the
            // compiler-grade K1 delegate, so AA is never worse than K1 in coverage.
            if (resp.candidates.isEmpty()) return Outcome(null, "empty${resp.message.takeIf { it.isNotBlank() }?.let { ": $it" } ?: ""}")

            val prefix = Regex("[A-Za-z0-9_]*$").find(snippet.substring(0, cursor.coerceIn(0, snippet.length)))?.value ?: ""
            val bounds = CompletionTokenBounds(cursor - prefix.length, cursor)
            val metadata = resp.candidates.map {
                SourceCodeCompletionVariant(text = it.text, displayText = it.displayText, tail = it.tail, icon = it.icon).apply {
                    // Mirror the K1 delegate's extraInfo so the server's doc-enrichment pipeline
                    // (member KDoc via `receiverType.member`, AI-draft docs, receiver docs) attaches to
                    // AA candidates too — otherwise dot completions come back doc-less.
                    val info = buildMap {
                        if (it.receiverType.isNotBlank() || it.receiverQualifiedName.isNotBlank()) {
                            put("receiverType", it.receiverType)
                            put("receiverQualifiedName", it.receiverQualifiedName)
                        }
                        // Read back by `KotlinExecution.complete` as the type-fit ranking boost — the
                        // only channel from the sidecar to the server's ranking, since a
                        // SourceCodeCompletionVariant has no field for it.
                        if (it.expectedTypeMatch) put(EXPECTED_TYPE_MATCH, "true")
                        if (it.declaredIn.isNotBlank()) put(DECLARED_IN, it.declaredIn)
                        if (it.proximity.isNotBlank()) put(DECLARATION_PROXIMITY, it.proximity)
                    }
                    if (info.isNotEmpty()) extraInfo = info
                }
            }
            Outcome(
                CompletionResult.Success(
                    matches = metadata.map { it.text },
                    bounds = bounds,
                    metadata = metadata,
                    text = snippet,
                    cursor = cursor,
                ),
                // The note travels on a SUCCESSFUL response too, and dropping it here is how a silent
                // truncation stayed invisible: the sidecar reported "receiver-index truncated at 800 of
                // 973 candidates for kotlin/String" on every `String.` completion and nobody ever saw it,
                // because only the failure paths above read `resp.message`. KI2 logs this detail on the
                // `completion backend=AA …` line, so a capped answer now says so where it is being served.
                "${metadata.size} candidates" + (resp.message.takeIf { it.isNotBlank() }?.let { "; $it" } ?: ""),
            )
        }.getOrElse {
            Outcome(null, "error ${it.javaClass.simpleName}: ${it.message}")
        }
    }