mirror of
https://github.com/mark3labs/kit.git
synced 2026-06-14 03:30:26 +00:00
Ed Zynda
49f8b485be
* 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.
72 lines
2.3 KiB
Go
72 lines
2.3 KiB
Go
package extensions
|
|
|
|
import "testing"
|
|
|
|
func TestAllEventTypes_Count(t *testing.T) {
|
|
all := AllEventTypes()
|
|
if len(all) != 33 {
|
|
t.Fatalf("expected 33 event types, got %d", len(all))
|
|
}
|
|
}
|
|
|
|
func TestAllEventTypes_NoDuplicates(t *testing.T) {
|
|
seen := make(map[EventType]bool)
|
|
for _, et := range AllEventTypes() {
|
|
if seen[et] {
|
|
t.Fatalf("duplicate event type: %s", et)
|
|
}
|
|
seen[et] = true
|
|
}
|
|
}
|
|
|
|
func TestEventType_IsValid(t *testing.T) {
|
|
for _, et := range AllEventTypes() {
|
|
if !et.IsValid() {
|
|
t.Errorf("expected %s to be valid", et)
|
|
}
|
|
}
|
|
|
|
invalid := EventType("nonexistent_event")
|
|
if invalid.IsValid() {
|
|
t.Error("expected 'nonexistent_event' to be invalid")
|
|
}
|
|
}
|
|
|
|
func TestEventType_TypeMethod(t *testing.T) {
|
|
tests := []struct {
|
|
event Event
|
|
want EventType
|
|
}{
|
|
{ToolCallEvent{ToolName: "test"}, ToolCall},
|
|
{ToolCallInputStartEvent{ToolCallID: "x", ToolName: "test"}, ToolCallInputStart},
|
|
{ToolCallInputDeltaEvent{ToolCallID: "x", Delta: "{"}, ToolCallInputDelta},
|
|
{ToolCallInputEndEvent{ToolCallID: "x"}, ToolCallInputEnd},
|
|
{ToolExecutionStartEvent{ToolName: "test"}, ToolExecutionStart},
|
|
{ToolExecutionEndEvent{ToolName: "test"}, ToolExecutionEnd},
|
|
{ToolResultEvent{ToolName: "test"}, ToolResult},
|
|
{InputEvent{Text: "hello"}, Input},
|
|
{BeforeAgentStartEvent{Prompt: "test"}, BeforeAgentStart},
|
|
{AgentStartEvent{Prompt: "test"}, AgentStart},
|
|
{AgentEndEvent{Response: "done"}, AgentEnd},
|
|
{MessageStartEvent{}, MessageStart},
|
|
{MessageUpdateEvent{Chunk: "hi"}, MessageUpdate},
|
|
{MessageEndEvent{Content: "done"}, MessageEnd},
|
|
{SessionStartEvent{SessionID: "abc"}, SessionStart},
|
|
{SessionShutdownEvent{}, SessionShutdown},
|
|
{ModelChangeEvent{NewModel: "a/b"}, ModelChange},
|
|
{ContextPrepareEvent{Messages: []ContextMessage{{Index: 0, Role: "user", Content: "hi"}}}, ContextPrepare},
|
|
{BeforeForkEvent{TargetID: "abc"}, BeforeFork},
|
|
{BeforeSessionSwitchEvent{Reason: "new"}, BeforeSessionSwitch},
|
|
{BeforeCompactEvent{EstimatedTokens: 1000}, BeforeCompact},
|
|
{SubagentStartEvent{ToolCallID: "x", Task: "t"}, SubagentStart},
|
|
{SubagentChunkEvent{ToolCallID: "x", ChunkType: "text"}, SubagentChunk},
|
|
{SubagentEndEvent{ToolCallID: "x"}, SubagentEnd},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
if got := tt.event.Type(); got != tt.want {
|
|
t.Errorf("event %T.Type() = %s, want %s", tt.event, got, tt.want)
|
|
}
|
|
}
|
|
}
|