koog-compose

2.1.0indexed

Single-block agent runtime orchestrating LLM conversation, native device tools and phases; streams tokens to Compose UI, enforces guardrails with confirmations, audit logging, circuit breakers and shared state.

AndroidJVMNative·BrianMwas/koog-compose

Stars

—

Used by

dependents

Health

/ 100

koog-compose

A Kotlin Multiplatform runtime for building on-device AI agents.

koog-compose lets you write a single koogCompose { } block that manages your LLM conversation, runs device tools (GPS, alarms, screen blocking), handles multi-step flows, and streams tokens to your Compose UI — with no server required.

Built on JetBrains Koog.

Production-ready security: All tool calls enforce guardrails (rate limits, allowlists, confirmations) with full audit logging. Circuit breakers prevent cascading failures. Thread-safe for parallel tool execution.

The core idea

Most AI integrations treat the LLM as a text box. koog-compose treats it as an orchestrator.

The user says "I'm going for a run." The agent starts a background GPS tracker, checks the weather, estimates duration from their history, and schedules a WorkManager task that survives the app being closed. When they say "I'm back," it stops everything, calculates pace, and responds conversationally. No buttons. No forms. No app-switching. The conversation is the UI.

This works because koog-compose bridges two things that usually live in separate worlds: the LLM conversation loop and the device's native APIs. KoogStateStore<S> is the shared state that flows from a tool result straight into your Compose UI via StateFlow. Device tools (location, alarms, screen time, camera) are first-class citizens of the agent graph, not afterthoughts bolted onto a chat widget.

Quick start

1. Install

dependencies {
    implementation("io.github.brianmwas.koog_compose:koog-compose-core:2.0.0")

    // Optional modules — add what you need
    implementation("io.github.brianmwas.koog_compose:koog-compose-ui:2.0.0")           // Material 3 chat components
    implementation("io.github.brianmwas.koog_compose:koog-compose-device:2.0.0")       // GPS, alarms, WorkManager (Android)
    implementation("io.github.brianmwas.koog_compose:koog-compose-mediapipe:2.0.0")    // On-device models (Gemma 4, Apple FMs)
    implementation("io.github.brianmwas.koog_compose:koog-compose-session-room:2.0.0") 
    implementation()      
}

2. Install the on-device bridge (Android only)

If you're using provider { onDevice(...) }, register the runtime bridge once at startup:

import io.github.koogcompose.provider.ondevice.installOnDeviceProviderSupport

fun initAi() {
    installOnDeviceProviderSupport()  // Application.onCreate() or main()
}

On iOS this happens automatically — iOSApp.init() installs the Apple Foundation Models bridge on launch.

3. Define your state

Everything flows through a single typed state object. Device tools write to it; your Compose UI reads from it.

@Serializable
data class RunState(
    val userName: String,
    val isRunning: Boolean = false,
    val distanceKm: Double = 0.0,
    val durationMs: Long = 0,
    val pace: String? = null,
)

4. Build the agent

val runCoach = koogCompose<RunState> {
    provider {
        onDevice(modelPath = "/data/models/gemma-4-E2B.litertlm") {
            onUnavailable { anthropic(apiKey = BuildConfig.KEY) }
        }
    }

    initialState { RunState(userName = ) }

    phases {
        phase(, initial = ) {
            instructions {  }
            tool(StartRunTimerTool(stateStore))
        }

        phase() {
            instructions {  }
            tool(BackgroundTimerTool())         
            tool(LocationTrackerTool(stateStore))
        }

        phase() {
            instructions {  }
            tool(StopTimerTool(stateStore))
            tool(CalculatePaceTool(stateStore))
        }
    }

    config {
        retry { maxAttempts = ; initialDelayMs =  }
        stuckDetection { threshold = ; fallbackMessage =  }
    }
}

5. Wire it to a ViewModel (or use Compose directly)

Option A: ViewModel + Compose (traditional)

class RunViewModel(context: KoogComposeContext<RunState>, executor: PromptExecutor) : ViewModel() {
    val session = phaseSession(context, executor) {
        sessionId = "run_brian"
        scope = viewModelScope
    }

    val responseStream = session.responseStream  // Flow<String> — tokens as they arrive
    val runState       = session.appState        // StateFlow<RunState>
}

Option B: Pure Compose (recommended for new code)

@Composable
fun RunScreen(definition: KoogDefinition<RunState> = koogCompose { ... }) {
    val session = rememberPhaseSession(definition) {
        sessionId = "run_brian"
    }

    val responseStream = session.responseStream
    val runState by session.appState.collectAsState()
}

6. Render it

With ViewModel:

@Composable
fun RunScreen(viewModel: RunViewModel = viewModel()) {
    val chatState = rememberChatState(viewModel.session)
    val runState  by viewModel.runState.collectAsState()

    if (runState.isRunning) {
        Text("Running — ${runState.distanceKm} km")
    }

    Scaffold(bottomBar = { ChatInputBar(chatState) }) { padding ->
        ChatMessageList(chatState, modifier = Modifier.padding(padding))
    }
}

Pure Compose (no ViewModel):

@Composable
fun RunScreen(definition: KoogDefinition<RunState> = koogCompose { ... }) {
    val session = rememberPhaseSession(definition) {
        sessionId = "run_brian"
    }
    val runState by session.appState.collectAsState()

    if (runState.isRunning) {
        Text("Running — ${runState.distanceKm} km")
    }

    Scaffold(bottomBar = { ChatInputBar(rememberChatState(session)) }) { padding ->
        ChatMessageList(rememberChatState(session), modifier = Modifier.padding(padding))
    }
}

How it works

Phases

A phase is a named stage in your conversation flow. Each one has its own system instructions and tool access. The LLM transitions between phases automatically — no manual routing code.

ready ──► running ──► finished ──► END

For more complex flows, phases can contain ordered subphases (sequential steps invisible to the router) and parallel branches (concurrent tool execution).

phase("finish_run") {
    subphase("stop_timer") {
        instructions { "Stop the run timer and record final duration." }
        tool(StopTimerTool(stateStore))
    }
    subphase("calculate_stats") {
        instructions { "Calculate distance and pace from the GPS trace." }
        tool(CalculatePaceTool(stateStore))
    }
    subphase("save_run") {
        instructions { "Save the run to storage." }
        tool(SaveRunTool(stateStore))
    }
    onCondition("run saved", "summary")
}

Branches inside parallel { } run concurrently using Koog's nodeExecuteMultipleTools(parallelTools = true):

phase("gather_context", initial = true) {
    parallel {
        branch("weather")  { tool(WeatherTool(stateStore)) }
        branch("location") { tool(GeocoderTool(stateStore)) }
        branch("history")  { tool(RunHistoryTool(stateStore)) }
    }
    onCondition("context ready", "plan")
}

Shared state

KoogStateStore<S> connects tools to your UI without globals or manual wiring:

Tool executes
  → stateStore.update { it.copy(distanceKm = 3.2) }
      → StateFlow<RunState> emits
          → Compose UI recomposes automatically

Writing a tool

class (
      stateStore: KoogStateStore<RunState>
) : StatefulTool<RunState>() {
      name            = 
      description     = 
      permissionLevel = PermissionLevel.SENSITIVE

      : ToolResult {
         location = getCurrentLocation()
        stateStore.update {
            it.copy(gpsTrace = it.gpsTrace + location)
        }
         ToolResult.Success()
    }
}

Every tool call goes through a pipeline before execute() is reached:

LLM args → validateArgs() → GuardrailEnforcer → [SENSITIVE/CRITICAL: confirmation UI] → execute()
                                                   ↓
                                         SAFE: skipped, runs silently

validateArgs() — block malformed or unexpected args before they cause runtime errors
GuardrailEnforcer — rate limits and action allowlists per tool

Streaming

responseStream emits tokens as they arrive from the model. Reset accumulation on each new turn using turnId:

val displayText by remember {
    viewModel.turnId.flatMapLatest { _ ->
        viewModel.responseStream.runningFold("") { acc, token -> acc + token }
    }
}.collectAsState(initial = "")

Generative UI layout

The layout engine lets the agent drive what UI is shown through a small, declarative vocabulary instead of free-form code generation. The agent emits AgentLayoutDirectives — ShowComponent, , , , — into named, host-declared . Each directive runs through a validation pipeline () before it touches the live your Compose UI renders.

Every directive carries a correlationId, and the engine publishes a DirectiveOutcome the agent reads back on its next turn:

positionFallback — ShowComponent can request a relative Position.Before(ref) or Position.After(ref). If the referenced component isn't in the slot, the engine silently appends to the end instead of rejecting the directive. When that happens, the Accepted (or Rewritten) outcome carries , so the agent can detect that its requested ordering wasn't honored and correct course on the next turn:

processor.outcomes.collect { outcome ->
    if (outcome is DirectiveOutcome.Accepted && outcome.positionFallback) {
        // The Before/After reference was missing — component went to the end.
        // The agent can re-issue a ReorderComponents directive if ordering matters.
    }
}

Testing

koog-compose-testing swaps the live provider for a scripted FakePromptExecutor. You test real phase transitions and tool dispatch without hitting a model.

@Test
fun `"I'm back" transitions from running to finished`() {
    val session = testPhaseSession(context) {
        on("I'm back", phase = "running") {
            transitionTo("finished")
            callTool("StopTimer")
            callTool("CalculateStats")
            respondWith("Great run! 3.2 km in 18 minutes — 5:38 pace.")
        }
    }

    session.send("I'm back")

    assertPhase(session, "finished")
    assertToolCalled(session, )
    assertState(session) { assertFalse(it.isRunning) }
}

Run tests without an emulator:

./gradlew :koog-compose-core:desktopTest

Test assertions:

This DSL approach makes tests deterministic and fast — no network calls, no model inference, no flakiness.

On-device models & Privacy

koog-compose runs inference locally on the device by default — no API key, no network call, all data stays on-device.

provider {
    onDevice(modelPath = "/data/models/gemma-4-E2B.litertlm") {
        maxToolRounds(8)
        onUnavailable {
            // Fallback only if model is unavailable:
            // - File missing or corrupted
            // - Device hardware incompatible
            // ⚠️ This fallback sends data to Anthropic's servers
            anthropic(apiKey = BuildConfig.KEY)
        }
    }
}

Data flow & privacy:

Important: If you use onUnavailable { anthropic(...) } as a fallback, that provider will see:

Full conversation history (messages + tool results)
Tool names and arguments
Application context (phase name, session ID)

If this is unacceptable, use onUnavailable { throw UnsupportedOperationException(...) } instead — users will see the error, but no data leaves the device.

Platform	Backend	Scope
Android	LiteRT-LM with Gemma 4 (E2B / E4B)

On Android, koog-compose disables LiteRT-LM's automatic tool calling loop so Gemma 4's <tool_call> responses are routed through koog's own SecureTool pipeline — validation and guardrails stay active regardless of the model backend.

Multi-agent handoff

Define specialist agents and the orchestrator delegates to them automatically:

val focusAgent = koogAgent("focus") {
    instructions { "You are a focus session specialist." }
    phases { phase("active") { /* ... */ } }
}

val session = koogSession<Unit> {
    provider { ollama(model = "llama3.2") }
    main {
        phases {
            phase("root", initial = true) {
                handoff(focusAgent) {
                    "User asks about focus, productivity, or pomodoro"
                }
            }
        }
    }
    agents(focusAgent)
}

Observability & event tracking

Route structured lifecycle events to Firebase, Datadog, a local database, or any custom backend. Events capture every significant moment: session starts, phase transitions, tool calls, guardrails denying access, stuck detection, and failures.

config {
    eventSink = PrintlnEventSink          // dev: logs to console
    // or
    eventSink = FirebaseEventSink()       // prod: Firebase Analytics
    // or
    eventSink = NoOpEventSink             // tests: silent
}

Events emitted at runtime:

Implement a custom sink by extending EventSink:

Wire in development, production, and test builds separately:

// In your DI container or ViewModel factory
val eventSink = when {
    BuildConfig.DEBUG    -> PrintlnEventSink
    BuildConfig.FIREBASE -> FirebaseEventSink(analytics)
    else                 -> NoOpEventSink
}

config { eventSink = eventSink }

Events are emitted from within coroutines and the sink is safe to suspend — use emit(event) to write to databases, call remote APIs, or batch events without blocking the agent.

Resume from any external trigger

Jump to a specific phase from a push notification, deep link, or WorkManager callback:

// From a notification
session.resumeAt("notify_user", userMessage = "Your run is ready to view!")

// From a deep link — no user message, no history pollution
session.resumeAt("onboarding_flow")

Reusable templates

Define common phase patterns once and include them anywhere:

val researchSubphase = subphaseTemplate("research") {
    instructions { "Search and summarise relevant information." }
    tool(WebSearchTool(stateStore))
}

phase("respond") {
    include(researchSubphase)    // adds the "research" subphase
    subphase("compose_answer") { /* ... */ }
}

Persistence

Session creation (DSL-consistent)

Create sessions using the phaseSession() DSL builder for consistency with koogCompose { }:

// In ViewModel
class MyViewModel(context: KoogComposeContext<MyState>, executor: PromptExecutor) : ViewModel() {
    val session = phaseSession(context, executor) {
        sessionId = "my_session"
        scope = viewModelScope
        store = RoomSessionStore(db.sessionDao())
    }
}

Or in Compose:

@Composable
fun MyScreen(definition: KoogDefinition<MyState>) {
    val session = rememberPhaseSession(definition) {
        sessionId = "my_screen_session"
        store = RedisSessionStore()
    }
}

All parameters are optional; defaults are sensible:

sessionId — defaults to "default"
scope — defaults to

Session storage

Drop in Room-backed persistence by passing a custom store:

val session = phaseSession(context, executor) {
    sessionId = "run_brian"
    scope = viewModelScope
    store = RoomSessionStore(db.sessionDao())  // ← Room backend
}

Or implement SessionStore directly to use any backend (Redis, SQLite, custom).

State migration

When your app state evolves, increment the schema version and define upgrade paths. Migrations are chained — if a user skips versions, all intermediate steps run automatically:

How chaining works:

If stored version is 1 and current is 3: v1 → v2 runs, then v2 → v3 runs
Each step transforms the JSON once
All steps happen before the app sees the state

Quick migrations (no explicit handler needed):

Added fields with defaults — handled automatically
Removed fields — ignored automatically
Use ignoreUnknownKeys + coerceInputValues in serializer

Explicit migrations only for:

Renamed fields
Retyped fields (e.g., String → Int)
Complex transformations (e.g., splitting one field into many)

Session Creation Patterns

koog-compose is DSL-first. All three ways to create a session follow the same builder pattern for consistency:

1. Non-Compose (ViewModel, Services)

class MyViewModel(context: KoogComposeContext<MyState>, executor: PromptExecutor) : ViewModel() {
    val session = phaseSession(context, executor) {
        sessionId = "my_session"
        scope = viewModelScope
        store = RoomSessionStore(db.sessionDao())  // optional
    }
}

All parameters except context and executor are optional.

2. Compose (recommended for new code)

@Composable
fun MyScreen(definition: KoogDefinition<MyState> = koogCompose { ... }) {
    val session = rememberPhaseSession(definition) {
        sessionId = "my_screen_session"
    }
}

rememberPhaseSession() automatically:

Binds to the Compose lifecycle
Uses lifecycleScope (no need to pass scope)
Memoizes across recompositions

3. Bridge Pattern (when you already have a definition)

val definition = koogCompose<MyState> { ... }
val session = definition.createPhaseSession(executor, viewModelScope) {
    sessionId = "my_session"
}

All three patterns are equivalent; choose based on your UI framework.

Platform support

Modules

Error Handling & Resilience

koog-compose provides production-grade error recovery patterns to keep your agent running even when dependencies fail.

Recovery Hints

Tool failures now carry metadata to guide the agent's recovery strategy:

Recovery hint types:

Hint	Use Case

Circuit Breaker

Prevent cascading failures when an external service keeps failing:

val breaker = CircuitBreaker(failureThreshold = 5, cooldownMs = 60_000)
val tool = CircuitBreakerGuard(
    delegate = SavePhotoTool(stateStore),
    circuitBreaker = breaker
)

// After 5 failures: circuit opens, returns user-friendly message
// After 60s cooldown: circuit enters half-open (trial mode)
// On success: circuit closes, normal operation resumed

The circuit breaker counts both thrown exceptions and ToolResult.Failure as failures. ToolResult.Denied (policy/user denials) are not counted — they're not service failures.

States:

CLOSED (normal) → failures counted
OPEN (broken) → calls rejected immediately
HALF_OPEN (trial) → one success closes it, one failure reopens

Pass a sessionId and an EventSink to observe state transitions — the guard emits AgentEvent.CircuitBreakerOpened when it trips and AgentEvent.CircuitBreakerClosed when it recovers, so you can surface degraded-mode banners or alert your backend:

val tool = CircuitBreakerGuard(
    delegate = SavePhotoTool(stateStore),
    circuitBreaker = breaker,
    sessionId = session.id,
    eventSink = myEventSink,   // routes to Firebase / Datadog / logs
)

Both parameters are optional (default: no session id, NoOpEventSink), so existing call sites keep working unchanged.

Session Corruption Recovery

Sessions corrupted by storage errors are detected and recovered:

val store = RoomSessionStore(dao, serializer)

// Load with automatic recovery
val result = store.loadOrRecover(sessionId)
when (result) {
    is SessionLoadResult.Success -> {
        session = resumeSession(result.session)
    }
    is SessionLoadResult.Recovered -> {
        showMessage(result.reason)  // "Session corrupted, starting fresh"
        session = startNewSession()
    }
    is SessionLoadResult.NotFound -> { }
}

Retry with Backoff

Configure automatic retries in your session config:

config {
    retry {
        maxAttempts = 3
        initialDelayMs = 1_000
        backoffMultiplier = 2.0   // 1s → 2s → 4s
    }
}

Error Mapping for Users

Never show raw exceptions to users. Map internal errors to friendly messages:

// Inside your tool
catch (e: IOException) {
    val userMessage = when {
        e.isNetworkRelated() -> "Internet connection problem — trying again..."
        e.isStorageFull() -> "Your device is full — please free up space"
        else -> "Something went wrong — our team is aware"
    }
    ToolResult.Failure(userMessage, retryable = false)
}

Privacy

All data stays on the device by default. koog-compose does not transmit prompts, responses, tool args, or telemetry anywhere. You own the SessionStore. Audit logs stay in-memory only, with optional PII redaction:

config {
    auditLog { redactArgs = true }
}

Contributing

Bug reports and feature requests → GitHub Issues Questions → GitHub Discussions

Read CONTRIBUTING.md before opening a PR.

License

Copyright 2025-2026 Brian Mwangi

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    https://www.apache.org/licenses/LICENSE-2.0

Related libraries

Surfaced from shared tags and platforms — no rankings paid for.

kaluga★ 397

splendoOffers modular tools for cross-platform mobile app development, including features like MVVM architecture, location services, permissions management, Bluetooth communication, and more, using coroutines and `Flow` for efficient design.Shared: kotlin-flow, kotlin-coroutines, device

KSafe★ 301

ioannisaEffortlessly encrypts and persists data using hardware-backed security, offering one-code-path simplicity, seamless integration, and protection for sensitive information like OAuth tokens.Shared: state-management, kotlin-flow, kotlin-coroutines

LazyCardStack★ 92

HukumisterProvides a Tinder-like card stack interface with a LazyColumn-like API, supporting swipe gestures, advanced animations, and programmatic card swiping. Includes callback for swipe detection and method to return previous cards.Shared: state-management, kotlin-coroutines, compose-multiplatform

kilua★ 364

rjarosComposable web framework enabling modern web application development with declarative UI components, state management, server-side rendering, and support for WebAssembly and JavaScript targets.Shared: state-management, kotlin-flow, kotlin-coroutines

ag-ui★ 14.5k

ag-ui-protocolOpen, lightweight event-based protocol standardizing agent-to-user application interactions, with transport-agnostic middleware, streaming real-time chat, bi-directional state sync, generative UI, structured messages, and tool integrations.Shared: state-management, agents, agentic-ai

koog★ 4.4k

JetBrainsFramework designed for building AI agents with tool interaction, complex workflows, semantic search, and persistent memory. Offers modular architecture, real-time processing, and comprehensive tracing.Shared: llm, agents, agentic-ai

`Accepted`	Applied as-is.
`Rewritten`	Modified by policy before applying (e.g. evict-then-show on a Single slot).
`Rejected`	Dropped; `rejectedAt` names the pipeline stage that refused it.
`Coalesced`	Deduplicated against an in-flight directive with the same `correlationId`.

`assertPhase(session, "phase_name")`	Verify current phase
`assertToolCalled(session, "ToolName")`	Verify tool was invoked
`assertToolNotCalled(session, "ToolName")`	Verify tool was NOT invoked
`assertState(session) { block }`	Assert app state with lambda
`assertResponse(session, "text")`	Verify agent response contains text

On-device model available	All inference runs locally	✅ 100% on-device, no internet
Model file missing	Falls back to `onUnavailable` block	⚠️ Data sent to fallback provider (Anthropic, OpenAI, etc.)
User revokes permissions	Tool execution denied, conversation continues	✅ On-device, no network
Tool calls device APIs (GPS, camera)	Local, permission-gated	✅ On-device, gated by OS permissions

`SessionStarted`	First user message	Session analytics, trace IDs
`PhaseTransitioned`	LLM routes to a new phase	Funnel analysis, flow tracing
`ToolCalled`	Tool executes successfully	Usage metrics, feature adoption
`GuardrailDenied`	Tool blocked by rate limit, allowlist, or user refusal	Security/compliance audit, UX friction
`AgentStuck`	LLM repeats the same phase N times	Loop detection, fallback messaging
`TurnFailed`	Retry exhausted after N attempts	Error rates, provider reliability
`CircuitBreakerOpened`	A `CircuitBreakerGuard` trips OPEN after repeated failures	Degraded-mode banners, dependency alerting
`CircuitBreakerClosed`	A tripped breaker recovers to CLOSED	Recovery tracking, clearing degraded UI
`LLMRequested`	(Reserved for future use)	—

Core DSL & phases	✅	✅	✅
Subphases & parallel branches	✅	✅	✅
Token streaming	✅	✅	✅
Multi-agent handoff	✅	✅	✅
On-device model (LiteRT-LM)	✅	—	—
On-device model (Apple FMs)	—	✅	—
Provider fallback routing	✅	✅	✅
Compose UI components	✅	✅	—
Room session store	✅	✅	—
Device tools & WorkManager	✅	—	—

`koog-compose-core`	DSL, agent runtime, phase engine — required
`koog-compose-ui`	Material 3 chat UI components
`koog-compose-device`	Android device tools (GPS, alarms, WorkManager)
`koog-compose-mediapipe`	On-device model providers (LiteRT-LM, Apple FMs)
`koog-compose-testing`	Deterministic fake executor + test assertions
`koog-compose-session-room`	Room-backed session persistence