Files
Ed Zynda 49f8b485be feat(extensions): add OnLLMUsage, SetState, enriched AgentEndEvent (#53) (#54)
* feat(extensions): add OnLLMUsage, SetState, enriched AgentEndEvent (#53)

Three additive primitives to the extension API:

- OnLLMUsage event: per-LLM-call token + cost deltas attributed to the
  specific model/provider used for each round-trip. Derived from the SDK
  StepFinishEvent in the extension bridge. Enables accurate budget
  enforcement between calls instead of only at turn boundaries.

- ctx.SetState / GetState / DeleteState / ListState: session-scoped,
  last-write-wins key-value store backed by a sidecar file
  (<session>.ext-state.json) outside the conversation tree. Reads are
  O(1), writes don't grow the JSONL, and the store is not duplicated on
  fork. State is preserved across hot-reloads.

- Enriched AgentEndEvent: ToolCallCount, ToolNames, LLMCallCount, token
  deltas (input/output/cache-read/cache-write), CostDelta, and
  DurationMs populated by a per-turn aggregator. Existing handlers
  reading only Response/StopReason are unaffected.

Includes unit tests for the state store, LLMUsage registration,
enriched AgentEndEvent, turn aggregator, llmUsageMeta, and sidecar path
derivation. Adds examples/extensions/usage-budget.go demoing all three
primitives together. Documents the additions in README, the docs site
(extensions overview, capabilities, examples), and the kit-extensions
and kit-sdk skill guides.

Fixes #53

* fix(extensions): address review feedback on state store and llmUsageMeta

- Serialize SetState/DeleteState saver invocations through a new saverMu
  so overlapping atomic-rename writes can no longer race on the shared
  .tmp file and persist an older snapshot after a newer one.
- LoadStateFromFile now clears the in-memory store when the sidecar is
  missing or empty, matching the documented "replace … with its
  contents" contract. This makes session-switching safe by preventing
  keys from a prior session leaking into a new one. Tests updated to
  cover both the missing-file and empty-file cases.
- llmUsageMeta now detects Anthropic OAuth credentials and returns
  Cost=0, matching the comment and the existing usage_tracker behavior
  for OAuth users. Mirrors the OAuth detection already used in
  cmd/extension_context.go.
- Document the single-in-flight-turn assumption baked into the
  per-turn aggregator with a clear migration path (per-turn ID) for if
  concurrent turns ever become a supported use case.

* fix(extensions): release saverMu on panic in state store

Extract a runSaver helper that locks saverMu and defers Unlock before
invoking the persistence callback. Without the deferred Unlock, a panic
inside the saver (e.g. disk full mid-write) would leave saverMu held
forever and deadlock the next SetState/DeleteState. Both SetState and
DeleteState now route through the helper. New TestRunner_State_Saver
PanicReleasesSaverMu reproduces the deadlock window with a 2s deadline
and proves the mutex is released after a panic.
2026-06-09 16:18:10 +03:00
..

Kit Extension Examples

A collection of example extensions demonstrating various Kit capabilities. These can be installed individually or as a complete collection.

Installation

Install all examples

kit install github.com/mark3labs/kit/examples/extensions

Install with interactive selection

kit install github.com/mark3labs/kit/examples/extensions --select

Install locally in your project

kit install github.com/mark3labs/kit/examples/extensions --local

Extension Index

Core Concepts

Extension Description Key API
minimal.go Minimal viable extension Basic Init() function
plan-mode.go Restrict agent to read-only tools OnBeforeAgentStart, SetActiveTools
tool-logger.go Log all tool calls to file OnToolCall, OnToolResult
notify.go Display notifications PrintInfo, PrintBlock

UI & Widgets

Extension Description Key API
widget-status.go Persistent status widget SetWidget, RemoveWidget
header-footer-demo.go Custom header/footer SetHeader, SetFooter
overlay-demo.go Modal overlay dialogs ShowOverlay
compact-notify.go Compact mode notifications PrintBlock
branded-output.go Custom styled output PrintBlock with colors

Input & Editor

Extension Description Key API
custom-editor-demo.go Custom key handling SetEditor, EditorKeyAction
pirate.go Transform user input OnInput, InputResult
interactive-shell.go Custom command input Slash commands with prompts
inline-bash.go Execute bash inline Input handling, exec

Session & Context

Extension Description Key API
context-inject.go Inject context into prompts OnContextPrepare
bookmark.go Bookmark messages AppendEntry, GetEntries
project-rules.go Project-specific rules Session data, file reading
protected-paths.go Block dangerous operations OnToolCall with blocking
permission-gate.go Confirm destructive actions OnToolCall with confirmation
usage-budget.go Soft cost cap + per-turn report OnLLMUsage, SetState/GetState, enriched AgentEndEvent

Tools & Commands

Extension Description Key API
auto-commit.go Auto-commit changes Custom tool, git operations
summarize.go Summarize conversation Custom tool with parameters
confirm-destructive.go Confirm destructive commands OnToolCall blocking
lsp-diagnostics.go LSP integration Complex extension, external process

Subagents & Background Tasks

Extension Description Key API
kit-kit.go Spawn Kit as subagent Subagent spawning
subagent-test.go Test subagent functionality SpawnSubagent
subagent-widget.go Widget with subagent updates Goroutines + widgets
dev-reload.go Hot reload extensions ReloadExtensions

Integrations

Extension Description Key API
kit-telegram/ Telegram relay for remote monitoring & control RegisterCommand, OnAgentStart/End, SetStatus, SendMessage

Themes

Extension Description Key API
neon-theme.go Register and switch custom themes RegisterTheme, SetTheme

Rendering

Extension Description Key API
tool-renderer-demo.go Custom tool output styling RegisterToolRenderer
prompt-demo.go Interactive prompts PromptSelect, PromptConfirm

Extension Details

minimal.go

The bare minimum extension showing the required structure:

  • Package main
  • Import kit/ext
  • Export Init(api ext.API) function

plan-mode.go

A complete example demonstrating:

  • Slash command (/plan)
  • Keyboard shortcut (ctrl+alt+p)
  • Option registration
  • Status bar indicators
  • System prompt injection
  • Tool filtering

widget-status.go

Shows how to create persistent UI elements:

  • Create widgets with SetWidget
  • Update content dynamically
  • Remove when done
  • Handle session lifecycle

context-inject.go

Advanced context manipulation:

  • Read project files
  • Inject into LLM context
  • Filter messages
  • Use negative indices for ephemeral content

lsp-diagnostics.go

Complex real-world example:

  • Multi-file extension
  • External process management (LSP server)
  • File watching
  • Diagnostics aggregation

kit-telegram/

Full-featured Telegram integration:

  • Slash command with subcommands and tab completion
  • Interactive guided setup flow with prompts
  • Background long-polling goroutine
  • Progress message rendering edited in place
  • Message queue with edit-before-dispatch
  • Remote command handling from Telegram
  • Status bar and widget updates
  • Config persistence with atomic writes

Multi-File Extension Example

The kit-kit-agents/ directory demonstrates the multi-file pattern:

kit-kit-agents/
├── main.go          # Entry point with Init()
├── agent.go         # Agent configuration
├── manager.go       # Agent lifecycle management
└── README.md        # Documentation

When the repo is installed, all files in subdirectories with main.go are loaded as separate extensions.

Testing & Validation

After installing, test the extensions:

# List all loaded extensions
kit extensions list

# Validate all extensions
kit extensions validate

# Run with a specific extension
kit -e ~/.local/share/kit/git/github.com/mark3labs/kit/examples/extensions/plan-mode.go

Creating Your Own

  1. Copy minimal.go as a starting point
  2. Modify the Init() function to register your handlers
  3. Use the other examples for reference on specific APIs
  4. Test with kit -e your-extension.go
  5. Share by pushing to a git repository!

Update

To get the latest examples:

kit install github.com/mark3labs/kit/examples/extensions --update

See Also