feat(extensions): add OnLLMUsage, SetState, enriched AgentEndEvent (#53)#54
Merged
Conversation
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
|
Connected to Huly®: KIT-55 |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds per-LLM-call usage events (OnLLMUsage), session-scoped key/value state APIs with optional per-session sidecar persistence, enriched AgentEndEvent per-turn aggregates, runtime turn aggregation and LLM cost/meta bridging, startup/context wiring, tests, examples, and documentation. ChangesExtension API Contracts
Runner state and persistence
Turn aggregation and LLM usage bridge
Tests and validation
Documentation and example
Sequence Diagram(s)sequenceDiagram
participant TurnStartEvent
participant turnAggregator
participant StepFinishEvent
participant Runner
participant StateSidecarFile
participant ExtensionHandler
TurnStartEvent->>turnAggregator: start()
StepFinishEvent->>turnAggregator: record step/tool usage
StepFinishEvent->>Runner: emit LLMUsageEvent
Runner->>ExtensionHandler: call OnLLMUsage handlers
Runner->>StateSidecarFile: stateSaver -> SaveStateToFile()
turnAggregator->>Runner: consume() -> AgentEndEvent enriched
Runner->>ExtensionHandler: call OnAgentEnd handlers
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
internal/extensions/events.go (1)
137-147:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winInclude
ToolOutputinAllEventTypes().
ToolOutputis declared as a supported event above, but this slice still skips it, soEventType.IsValid()returns false for a real lifecycle event.🐛 Proposed fix
return []EventType{ ToolCall, ToolCallInputStart, ToolCallInputDelta, ToolCallInputEnd, - ToolExecutionStart, ToolExecutionEnd, ToolResult, + ToolExecutionStart, ToolExecutionEnd, ToolOutput, ToolResult, Input, BeforeAgentStart, AgentStart, AgentEnd,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/extensions/events.go` around lines 137 - 147, AllEventTypes() currently returns a slice missing the ToolOutput EventType, causing EventType.IsValid() to reject real lifecycle events; update the slice returned by AllEventTypes (in internal/extensions/events.go) to include ToolOutput among the other EventType entries (e.g., alongside ToolResult/ToolExecutionEnd) so ToolOutput is recognized as a valid event.cmd/root.go (1)
1197-1204:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReinitialize extension state when switching sessions.
This callback swaps the tree session but leaves the extension state map and saver bound to the previous session. After a
/newor resume into a different JSONL,GetStatewill still read the old session's keys andSetStatewill keep writing the old sidecar.🐛 Proposed fix
func switchSessionForUI(path string) error { ts, err := kit.OpenTreeSession(path) if err != nil { return fmt.Errorf("failed to open session: %w", err) } kitInstance.SetTreeSession(ts) appInstance.SwitchTreeSession(ts) + if err := kitInstance.Extensions().InitStatePersistence(); err != nil { + return fmt.Errorf("reinitializing extension state: %w", err) + } return nil }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cmd/root.go` around lines 1197 - 1204, The switchSessionForUI callback currently replaces the TreeSession but leaves extension state and its saver bound to the old session; update switchSessionForUI (after kitInstance.SetTreeSession and appInstance.SwitchTreeSession) to reinitialize the extension state map and recreate/rebind the extension state saver for the newly opened session so GetState/SetState operate on the new session's keys/sidecar (use the TreeSession/ts identifier or path from kit.OpenTreeSession to create the new saver and an empty state map, replacing the previous extStateMap/extStateSaver used by GetState/SetState).
🧹 Nitpick comments (2)
pkg/kit/extension_api.go (1)
5-6: ⚡ Quick winUse structured logging here.
Please route this warning through
github.com/charmbracelet/loginstead oflog.Printfso it stays consistent with the repo's logging contract.As per coding guidelines,
**/*.go:Use github.com/charmbracelet/logfor structured logging.♻️ Proposed fix
-import ( - "fmt" - "log" - "strings" +import ( + "fmt" + "strings" + "github.com/charmbracelet/log" "github.com/mark3labs/kit/internal/extensions" "github.com/mark3labs/kit/internal/message" "github.com/mark3labs/kit/internal/session" ) @@ runner := e.kit.extRunner runner.SetStateSaver(func() { if err := runner.SaveStateToFile(path); err != nil { - log.Printf("WARN extension state save failed: path=%s err=%v", path, err) + log.Warn("extension state save failed", "path", path, "err", err) } })Also applies to: 394-397
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/kit/extension_api.go` around lines 5 - 6, Replace usages of the standard library logger with the project's structured logger: remove the "log" import and import "github.com/charmbracelet/log" instead, then replace log.Printf calls that emit warnings with the charmbracelet logger equivalents (e.g., log.Warn or log.Warnf) so messages are structured; update every occurrence in this file (including the warnings around the earlier occurrence and the ones at the other locations noted around lines 394-397) to use log.Warn/ log.Warnf and preserve the original message and formatting arguments.Source: Coding guidelines
cmd/root.go (1)
934-935: ⚡ Quick winPrefer structured logging for the startup state warning.
This new warning uses stdlib logging instead of the repository's structured logger.
As per coding guidelines,
**/*.go:Use github.com/charmbracelet/logfor structured logging.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cmd/root.go` around lines 934 - 935, The call to log.Printf in the kitInstance.Extensions().InitStatePersistence() error path uses stdlib logging; replace it with the repository's structured logger from github.com/charmbracelet/log (e.g., call logger.Warn or log.Warn with field-style context) and include the error as a field (error=err) and a clear message like "extension state init failed" so the InitStatePersistence() failure is recorded using structured logging instead of fmt-based Printf.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/extensions/runner.go`:
- Around line 773-783: Concurrent callers to Runner.SetState (and DeleteState)
can race because saver() is invoked after releasing stateMu, allowing
out-of-order snapshot writes to path+".tmp"; fix by serializing saver execution:
either call r.stateSaver() while still holding r.stateMu (move the saver
invocation inside the critical section in Runner.SetState and similarly in
DeleteState) or add a dedicated mutex (e.g., r.saverMu) and wrap saver() calls
with that mutex while ensuring the saver reads state under r.stateMu so
snapshots are consistent; update SetState, DeleteState and any other places
noted (lines ~797-807, 877-899) to use the same approach to prevent overlapping
saver runs and temp-file races.
- Around line 852-869: LoadStateFromFile currently returns early on missing file
or empty data and leaves r.state unchanged, causing stale keys to persist;
modify Runner.LoadStateFromFile so that when os.IsNotExist(err) is true or when
len(data)==0 it locks r.stateMu and replaces r.state with an empty map (e.g.,
map[string]string{}), then unlocks, ensuring the in-memory store is cleared for
a fresh session; keep the rest of the function behavior (JSON unmarshal path)
unchanged.
In `@pkg/kit/extensions_bridge.go`:
- Around line 332-356: The llmUsageMeta call is producing non-zero cost even for
OAuth-backed sessions; update the LLM usage path to detect OAuth credentials and
force Cost=0 when OAuth is in use. Either extend llmUsageMeta to accept the
session/credentials (or return an explicit oauth-backed boolean) and have
m.Subscribe handlers for StepFinishEvent use that to set Cost/Cos tDelta to 0
before emitting extensions.LLMUsageEvent, or check ev (e.g., ev.Session /
ev.Credentials) in the subscriber and override the computed cost to zero when
OAuth is present; apply the same change to the other LLM-usage emitter block
around the later section (the second occurrence referenced in the comment).
- Around line 24-37: turnAggregator is currently a single shared accumulator
(created as turnAgg := &turnAggregator{kit: m}) that is reset on TurnStartEvent
and snapshotted on TurnEndEvent, which allows interleaved/concurrent turns to
clobber each other's tool/step/usage counts; fix by making aggregation per-turn
(e.g., change turnAggregator to maintain a map[keyed by turn ID] and use the
event TurnStartEvent/TurnEndEvent IDs to start/stop/snapshot via
recordTool/recordStep into the per-turn entry) or alternatively enforce single
in-flight turn by adding mutual exclusion around Kit.runTurn/Prompt so events
cannot overlap. Also address the llmUsageMeta comment mismatch: either update
the comment on llmUsageMeta to accurately state when cost is zero (i.e., when
model/pricing info is missing/unparseable) or implement the intended OAuth
branch (detect OAuth credentials and set cost to zero) inside llmUsageMeta so
the behavior matches the comment. Ensure references to turnAggregator,
TurnStartEvent, TurnEndEvent, recordTool, recordStep, Kit.runTurn/Prompt, and
llmUsageMeta are used to locate the changes.
---
Outside diff comments:
In `@cmd/root.go`:
- Around line 1197-1204: The switchSessionForUI callback currently replaces the
TreeSession but leaves extension state and its saver bound to the old session;
update switchSessionForUI (after kitInstance.SetTreeSession and
appInstance.SwitchTreeSession) to reinitialize the extension state map and
recreate/rebind the extension state saver for the newly opened session so
GetState/SetState operate on the new session's keys/sidecar (use the
TreeSession/ts identifier or path from kit.OpenTreeSession to create the new
saver and an empty state map, replacing the previous extStateMap/extStateSaver
used by GetState/SetState).
In `@internal/extensions/events.go`:
- Around line 137-147: AllEventTypes() currently returns a slice missing the
ToolOutput EventType, causing EventType.IsValid() to reject real lifecycle
events; update the slice returned by AllEventTypes (in
internal/extensions/events.go) to include ToolOutput among the other EventType
entries (e.g., alongside ToolResult/ToolExecutionEnd) so ToolOutput is
recognized as a valid event.
---
Nitpick comments:
In `@cmd/root.go`:
- Around line 934-935: The call to log.Printf in the
kitInstance.Extensions().InitStatePersistence() error path uses stdlib logging;
replace it with the repository's structured logger from
github.com/charmbracelet/log (e.g., call logger.Warn or log.Warn with
field-style context) and include the error as a field (error=err) and a clear
message like "extension state init failed" so the InitStatePersistence() failure
is recorded using structured logging instead of fmt-based Printf.
In `@pkg/kit/extension_api.go`:
- Around line 5-6: Replace usages of the standard library logger with the
project's structured logger: remove the "log" import and import
"github.com/charmbracelet/log" instead, then replace log.Printf calls that emit
warnings with the charmbracelet logger equivalents (e.g., log.Warn or log.Warnf)
so messages are structured; update every occurrence in this file (including the
warnings around the earlier occurrence and the ones at the other locations noted
around lines 394-397) to use log.Warn/ log.Warnf and preserve the original
message and formatting arguments.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: c2291a81-b571-4208-9da0-99195b039ae6
📒 Files selected for processing (22)
README.mdcmd/extension_context.gocmd/root.goexamples/extensions/README.mdexamples/extensions/usage-budget.gointernal/extensions/api.gointernal/extensions/events.gointernal/extensions/events_test.gointernal/extensions/llmusage_test.gointernal/extensions/loader.gointernal/extensions/runner.gointernal/extensions/state_test.gointernal/extensions/symbols.gointernal/extensions/test_api.gopkg/kit/extension_api.gopkg/kit/extensions_bridge.gopkg/kit/extensions_bridge_test.goskills/kit-extensions/SKILL.mdskills/kit-sdk/SKILL.mdwww/pages/extensions/capabilities.mdwww/pages/extensions/examples.mdwww/pages/extensions/overview.md
- 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.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
internal/extensions/runner.go (1)
896-916:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
SaveStateToFilestill races for direct callers.The new
saverMuonly protects the callback path inSetState/DeleteState. Two goroutines can still callSaveStateToFile(path)concurrently, both write topath + ".tmp", and one rename/remove sequence can clobber the other's temp file or return a spurious error. That breaks the method's "Thread-safe" contract for persistence itself.🔧 Suggested direction
func (r *Runner) SaveStateToFile(path string) error { + r.saverMu.Lock() + defer r.saverMu.Unlock() + snap := r.SnapshotState() if snap == nil { snap = map[string]string{} } @@ - tmp := path + ".tmp" - if err := os.WriteFile(tmp, data, 0o644); err != nil { + tmpFile, err := os.CreateTemp(filepath.Dir(path), filepath.Base(path)+".*.tmp") + if err != nil { + return fmt.Errorf("creating temp state file: %w", err) + } + tmp := tmpFile.Name() + if _, err := tmpFile.Write(data); err != nil { + _ = tmpFile.Close() + _ = os.Remove(tmp) + return fmt.Errorf("writing extension state: %w", err) + } + if err := tmpFile.Close(); err != nil { + _ = os.Remove(tmp) + return fmt.Errorf("closing extension state temp file: %w", err) + } - return fmt.Errorf("writing extension state: %w", err) - } if err := os.Rename(tmp, path); err != nil { _ = os.Remove(tmp) return fmt.Errorf("renaming extension state: %w", err) }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/extensions/runner.go` around lines 896 - 916, SaveStateToFile can race when multiple goroutines write the same tmp path; fix by serializing persistence: acquire the existing r.saverMu (or a new file-specific mutex on Runner) at the start of SaveStateToFile and defer its unlock so only one caller writes/renames/removes the temp file at a time, and also switch to creating a unique temp file per call (e.g., using os.CreateTemp/os.TempFile in the target directory) before renaming to the final path to avoid tmp-name collisions; update SaveStateToFile to use these changes (referencing SaveStateToFile and r.saverMu).
🧹 Nitpick comments (1)
internal/extensions/state_test.go (1)
174-193: ⚡ Quick winExercise the saver path in the concurrency test.
This only stresses the in-memory map, and every goroutine writes the same
"v", so it will not catch a regression in the serializedSaveStateToFilepath that this PR is actually hardening. A saver-backed concurrent mutation test would lock that behavior in much better.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@internal/extensions/state_test.go` around lines 174 - 193, The concurrency test TestRunner_State_ConcurrentSet only exercises the in-memory map; modify it to exercise the saver-backed path by constructing the Runner with a file-backed saver (or the existing SaveStateToFile implementation) instead of NewRunner(nil). Create a temporary file or temp dir, initialize the saver used by SaveStateToFile, instantiate the runner with that saver (e.g., via NewRunnerWithSaver or by setting runner.saver/savePath before starting goroutines), then spawn goroutines that call Runner.SetState("k", value) with varying values (or at least simultaneous writes) so the saver path is exercised; clean up the temp file/dir after the test and assert final state as before. Ensure you reference TestRunner_State_ConcurrentSet, NewRunner / runner.saver / SaveStateToFile, and Runner.SetState when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@internal/extensions/runner.go`:
- Around line 787-790: The mutex saverMu is locked then unlocked around a call
to an arbitrary callback (stateSaver) in SetState and DeleteState, risking a
permanent lock if the callback panics; change the pattern to unlock with defer
immediately after locking (e.g., lock saverMu, defer unlock) so the mutex is
always released, or extract a small helper (runStateSaver) that acquires
saverMu, defers unlock, and invokes the stateSaver callback to avoid duplication
and ensure safety in both SetState and DeleteState.
---
Outside diff comments:
In `@internal/extensions/runner.go`:
- Around line 896-916: SaveStateToFile can race when multiple goroutines write
the same tmp path; fix by serializing persistence: acquire the existing
r.saverMu (or a new file-specific mutex on Runner) at the start of
SaveStateToFile and defer its unlock so only one caller writes/renames/removes
the temp file at a time, and also switch to creating a unique temp file per call
(e.g., using os.CreateTemp/os.TempFile in the target directory) before renaming
to the final path to avoid tmp-name collisions; update SaveStateToFile to use
these changes (referencing SaveStateToFile and r.saverMu).
---
Nitpick comments:
In `@internal/extensions/state_test.go`:
- Around line 174-193: The concurrency test TestRunner_State_ConcurrentSet only
exercises the in-memory map; modify it to exercise the saver-backed path by
constructing the Runner with a file-backed saver (or the existing
SaveStateToFile implementation) instead of NewRunner(nil). Create a temporary
file or temp dir, initialize the saver used by SaveStateToFile, instantiate the
runner with that saver (e.g., via NewRunnerWithSaver or by setting
runner.saver/savePath before starting goroutines), then spawn goroutines that
call Runner.SetState("k", value) with varying values (or at least simultaneous
writes) so the saver path is exercised; clean up the temp file/dir after the
test and assert final state as before. Ensure you reference
TestRunner_State_ConcurrentSet, NewRunner / runner.saver / SaveStateToFile, and
Runner.SetState when making the change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 3ea24ef0-459b-448d-b3c3-871b14fc0b29
📒 Files selected for processing (4)
internal/extensions/runner.gointernal/extensions/state_test.gopkg/kit/extensions_bridge.gopkg/kit/extensions_bridge_test.go
🚧 Files skipped from review as they are similar to previous changes (1)
- pkg/kit/extensions_bridge_test.go
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.
10 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds three independent, additive primitives to the extension API surface that issue #53 called out as missing, all delivered together because they share an implementation theme (give extensions accurate per-turn signals and a lightweight place to put state):
OnLLMUsageevent — fires after every LLM provider round-trip with per-call token + cost deltas attributed to the specific model/provider used. Derived from the existing SDKStepFinishEventin the extension bridge; cost is computed from the models registry (zero when pricing is unknown or OAuth credentials are in use). This unblocks budget enforcement that needs to react between calls within a turn, not 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 session JSONL, and the store is not duplicated when the conversation forks. State is invisible to the LLM, survives session resume, and is preserved across extension hot-reloads. Use this for snapshot state ("current value of X"); useAppendEntryfor audit logs.AgentEndEvent— addsToolCallCount,ToolNames,LLMCallCount,InputTokensDelta,OutputTokensDelta,CacheReadTokensDelta,CacheWriteTokensDelta,CostDelta, andDurationMs. Populated by a per-turn aggregator hooked intoTurnStartEvent/ToolResultEvent/StepFinishEvent. Existing handlers reading onlyResponse/StopReasonare unaffected.A new example extension
usage-budget.godemonstrates all three primitives together as a soft-cost-cap with a per-turn report.Fixes #53
Type of Change
Checklist
go vetandgolangci-lint runboth clean)go test -race ./...)npx tome buildinwww/)Additional Information
New files
internal/extensions/state_test.go— 7 tests for the state store (CRUD, saver hook, save/load round-trip, malformed JSON, missing file, concurrent writes, no-op defaults).internal/extensions/llmusage_test.go— 3 tests forLLMUsageEventregistration and enrichedAgentEndEventfield round-trip.pkg/kit/extensions_bridge_test.go— 6 tests forturnAggregator,llmUsageMetanil-safety, andextStateSidecarPath.examples/extensions/usage-budget.go— demo extension wiring all three primitives.Notable changes
internal/extensions/api.go— newLLMUsageEventstruct, four newContextfunction fields, enrichedAgentEndEvent, newOnLLMUsageregistrar.internal/extensions/runner.go— state store (SetState/GetState/DeleteState/ListState/SnapshotState/LoadStateFromFile/SaveStateToFile/SetStateSaver) withsync.RWMutexprotection. State preserved acrossReload.internal/extensions/{events,loader,test_api,symbols}.go— wiredLLMUsageevent type, handler registration, and Yaegi export.pkg/kit/extensions_bridge.go—turnAggregatorsubscribes to turn/tool/step events; emitsLLMUsageEventderived fromStepFinishEvent; populates enriched fields onAgentEndEvent.pkg/kit/extension_api.go— exposes the state methods onExtensionAPI. NewInitStatePersistence()loads any existing sidecar and installs a saver hook.cmd/{extension_context,root}.go— wires the fourContextcallbacks and callsInitStatePersistencebeforeEmitSessionStart.internal/extensions/events_test.go— bumps the expected count from 32 → 33.Backwards compatibility
Fully additive. All previously valid handler signatures continue to compile and run unchanged. Extensions that read only
Response/StopReasononAgentEndEventsee the same values they always have; the new fields default to zero values. The state store is opt-in — no behaviour change for extensions that don't callSetState.Documentation
README.md— addedOnLLMUsageto lifecycle list, added Session State bullet, addedusage-budget.goto examples list, added a paragraph explaining the enriched event.www/pages/extensions/{overview,capabilities,examples}.md— capability table, field tables for both new events, new Session state section with a "when to use which" comparison vsAppendEntry.skills/kit-extensions/SKILL.mdandskills/kit-sdk/SKILL.md— updated agent skill guides.Summary by CodeRabbit