mirror of https://github.com/mark3labs/kit.git synced 2026-06-13 19:20:06 +00:00

Files

T

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

3.0 KiB

Raw Blame History

title, description

title	description
Extension System	Overview of Kit's Go-based extension system.

Extension System

Extensions are Go source files interpreted at runtime via Yaegi. They can add custom tools, slash commands, widgets, keyboard shortcuts, and intercept lifecycle events — all without recompiling Kit.

Minimal extension

//go:build ignore

package main

import "kit/ext"

func Init(api ext.API) {
    api.OnSessionStart(func(_ ext.SessionStartEvent, ctx ext.Context) {
        ctx.SetFooter(ext.HeaderFooterConfig{
            Content: ext.WidgetContent{Text: "Custom Footer"},
        })
    })
}

Run it with:

kit -e examples/extensions/minimal.go

How extensions work

Kit discovers extension files from auto-discovery paths or explicit -e flags
Each .go file is loaded into a Yaegi interpreter with access to the kit/ext package
Kit calls the Init(api ext.API) function in each extension
The extension registers callbacks, tools, commands, and UI components via the api and ctx objects

Key concepts

The `API` object

Passed to Init(), the API object is used to register lifecycle event handlers and static components:

Lifecycle handlers — api.OnSessionStart(...), api.OnToolCall(...), etc.
Tools — api.RegisterTool(ext.ToolDef{...})
Commands — api.RegisterCommand(ext.CommandDef{...})
Shortcuts — api.RegisterShortcut(ext.ShortcutDef{...}, handler)
Tool renderers — api.RegisterToolRenderer(ext.ToolRenderConfig{...})
Message renderers — api.RegisterMessageRenderer(ext.MessageRendererConfig{...})
Options — api.RegisterOption(ext.OptionDef{...})

The `Context` object

Passed to event handlers, the Context object provides runtime access to Kit's state and UI:

Output — ctx.Print(...), ctx.PrintInfo(...), ctx.PrintError(...)
UI components — ctx.SetWidget(...), ctx.SetHeader(...), ctx.SetFooter(...), ctx.SetStatus(...)
Editor — ctx.SetEditor(...), ctx.ResetEditor()
Prompts — ctx.PromptSelect(...), ctx.PromptConfirm(...), ctx.PromptInput(...)
Overlays — ctx.ShowOverlay(...)
Messages — ctx.SendMessage(...), ctx.GetMessages()
Model — ctx.SetModel(...), ctx.GetAvailableModels()
Tools — ctx.GetAllTools(), ctx.SetActiveTools(...)
Context stats — ctx.GetContextStats()
Session data — ctx.AppendEntry(...), ctx.GetEntries(...) (append-only, in conversation tree)
Session state — ctx.SetState(...), ctx.GetState(...), ctx.DeleteState(...), ctx.ListState() (last-write-wins, sidecar file)
Subagents — ctx.SpawnSubagent(...)
LLM completion — ctx.Complete(...)
Custom events — ctx.EmitCustomEvent(...)

See Capabilities for full details on each component type, and Testing for writing tests for your extensions.

3.0 KiB Raw Blame History