mirror of
https://github.com/mark3labs/kit.git
synced 2026-06-14 03:30:26 +00:00
Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| e8e99b19a8 | |||
| ef072f6e59 | |||
| 49f8b485be | |||
| febdc530e1 | |||
| e610bdd2d0 |
@@ -228,6 +228,10 @@ kit auth login [provider] --set-default # Set provider's default model as syste
|
||||
kit auth logout [provider] # Remove credentials for provider
|
||||
kit auth status # Check authentication status
|
||||
|
||||
# GitHub Copilot login (experimental; requires active Copilot subscription)
|
||||
kit auth login copilot
|
||||
kit --model copilot/gpt-5.5 "Hello"
|
||||
|
||||
# Model database
|
||||
kit models [provider] # List available models (optionally filter by provider)
|
||||
kit models --all # Show all providers (not just LLM-compatible)
|
||||
@@ -308,12 +312,15 @@ kit -e examples/extensions/minimal.go
|
||||
|
||||
### Extension Capabilities
|
||||
|
||||
**Lifecycle Events**: OnSessionStart, OnSessionShutdown, OnBeforeAgentStart, OnAgentStart, OnAgentEnd, OnToolCall, OnToolCallInputStart, OnToolCallInputDelta, OnToolCallInputEnd, OnToolExecutionStart, OnToolOutput, OnToolExecutionEnd, OnToolResult, OnInput, OnMessageStart, OnMessageUpdate, OnMessageEnd, OnModelChange, OnContextPrepare, OnBeforeFork, OnBeforeSessionSwitch, OnBeforeCompact, OnCustomEvent, OnSubagentStart, OnSubagentChunk, OnSubagentEnd
|
||||
**Lifecycle Events**: OnSessionStart, OnSessionShutdown, OnBeforeAgentStart, OnAgentStart, OnAgentEnd, OnLLMUsage, OnToolCall, OnToolCallInputStart, OnToolCallInputDelta, OnToolCallInputEnd, OnToolExecutionStart, OnToolOutput, OnToolExecutionEnd, OnToolResult, OnInput, OnMessageStart, OnMessageUpdate, OnMessageEnd, OnModelChange, OnContextPrepare, OnBeforeFork, OnBeforeSessionSwitch, OnBeforeCompact, OnCustomEvent, OnSubagentStart, OnSubagentChunk, OnSubagentEnd
|
||||
|
||||
`OnAgentEnd` carries per-turn aggregates (`ToolCallCount`, `ToolNames`, `LLMCallCount`, `InputTokensDelta`, `OutputTokensDelta`, `CostDelta`, `DurationMs`) so observers don't need to maintain parallel bookkeeping. `OnLLMUsage` fires after each LLM provider call with token + cost deltas attributed to that specific call/model — use it for accurate budget enforcement *between* calls instead of waiting for the turn to finish.
|
||||
|
||||
**Custom Components**:
|
||||
- **Tools**: Add new tools the LLM can invoke
|
||||
- **Commands**: Register slash commands (e.g., `/mycommand`)
|
||||
- **Options**: Register configurable extension options
|
||||
- **Session State**: Last-write-wins key-value store via `ctx.SetState` / `GetState` / `DeleteState` / `ListState`, persisted to a per-session sidecar file outside the conversation tree
|
||||
- **Widgets**: Persistent status displays above/below input
|
||||
- **Headers/Footers**: Persistent content above/below the conversation
|
||||
- **Status Bar**: Custom status bar entries
|
||||
@@ -369,6 +376,7 @@ See the `examples/extensions/` directory:
|
||||
- [`tool-logger.go`](examples/extensions/tool-logger.go) - Log all tool calls
|
||||
- [`neon-theme.go`](examples/extensions/neon-theme.go) - Custom theme registration and switching
|
||||
- [`tool-renderer-demo.go`](examples/extensions/tool-renderer-demo.go) - Custom tool call rendering
|
||||
- [`usage-budget.go`](examples/extensions/usage-budget.go) - Per-call usage callback (`OnLLMUsage`), session state, and enriched `OnAgentEnd` per-turn report
|
||||
- [`widget-status.go`](examples/extensions/widget-status.go) - Persistent status widgets
|
||||
|
||||
Also see [`.kit/extensions/go-edit-lint.go`](.kit/extensions/go-edit-lint.go) (in this repo) for a project-local extension example that runs gopls and golangci-lint on Go file edits.
|
||||
@@ -949,6 +957,7 @@ npm/ - NPM package wrapper for distribution
|
||||
|
||||
- **Anthropic** - Claude models (native, prompt caching, OAuth)
|
||||
- **OpenAI** - GPT models
|
||||
- **Copilot** - GitHub Copilot models (`copilot`, requires active Copilot subscription)
|
||||
- **Google** - Gemini models
|
||||
- **Ollama** - Local models
|
||||
- **Azure OpenAI** - Azure-hosted OpenAI
|
||||
|
||||
+157
-4
@@ -31,10 +31,12 @@ using OAuth flows. Stored credentials take precedence over environment variables
|
||||
Available providers:
|
||||
- anthropic: Anthropic Claude API (OAuth)
|
||||
- openai: OpenAI API (OAuth and API key)
|
||||
- copilot: GitHub Copilot (GitHub device login)
|
||||
|
||||
Examples:
|
||||
kit auth login anthropic
|
||||
kit auth login openai
|
||||
kit auth login copilot
|
||||
kit auth logout anthropic
|
||||
kit auth status`,
|
||||
}
|
||||
@@ -54,6 +56,7 @@ environment variables when making API calls.
|
||||
Available providers:
|
||||
- anthropic: Anthropic Claude API (OAuth)
|
||||
- openai: OpenAI ChatGPT Plus/Pro (Codex OAuth)
|
||||
- copilot: GitHub Copilot (GitHub device login, experimental)
|
||||
|
||||
Flags:
|
||||
--set-default Set this provider's default model as the system default
|
||||
@@ -61,7 +64,8 @@ Flags:
|
||||
Examples:
|
||||
kit auth login anthropic
|
||||
kit auth login openai
|
||||
kit auth login openai --set-default`,
|
||||
kit auth login copilot
|
||||
kit auth login copilot --set-default`,
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: runAuthLogin,
|
||||
}
|
||||
@@ -80,10 +84,12 @@ You will need to use environment variables or command-line flags for authenticat
|
||||
Available providers:
|
||||
- anthropic: Anthropic Claude API
|
||||
- openai: OpenAI API
|
||||
- copilot: GitHub Copilot
|
||||
|
||||
Example:
|
||||
kit auth logout anthropic
|
||||
kit auth logout openai`,
|
||||
kit auth logout openai
|
||||
kit auth logout copilot`,
|
||||
Args: cobra.ExactArgs(1),
|
||||
RunE: runAuthLogout,
|
||||
}
|
||||
@@ -113,6 +119,7 @@ var (
|
||||
var defaultModels = map[string]string{
|
||||
"anthropic": "anthropic/claude-sonnet-4-5-20250929",
|
||||
"openai": "openai/gpt-5.4",
|
||||
"copilot": "copilot/gpt-5.5",
|
||||
}
|
||||
|
||||
// setDefaultModelIfRequested sets the default model for the given provider
|
||||
@@ -143,6 +150,7 @@ func init() {
|
||||
authLoginCmd.Flags().BoolVar(&loginSetDefault, "set-default", false, "Set this provider's default model as the system default after login")
|
||||
}
|
||||
|
||||
// runAuthLogin dispatches OAuth login to the selected provider.
|
||||
func runAuthLogin(cmd *cobra.Command, args []string) error {
|
||||
provider := strings.ToLower(args[0])
|
||||
|
||||
@@ -151,8 +159,10 @@ func runAuthLogin(cmd *cobra.Command, args []string) error {
|
||||
return loginAnthropic()
|
||||
case "openai":
|
||||
return loginOpenAI()
|
||||
case "copilot":
|
||||
return loginCopilot(cmd.Context())
|
||||
default:
|
||||
return fmt.Errorf("unsupported provider: %s. Available providers: anthropic, openai", provider)
|
||||
return fmt.Errorf("unsupported provider: %s. Available providers: anthropic, openai, copilot", provider)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -164,8 +174,10 @@ func runAuthLogout(cmd *cobra.Command, args []string) error {
|
||||
return logoutAnthropic()
|
||||
case "openai":
|
||||
return logoutOpenAI()
|
||||
case "copilot":
|
||||
return logoutCopilot()
|
||||
default:
|
||||
return fmt.Errorf("unsupported provider: %s. Available providers: anthropic, openai", provider)
|
||||
return fmt.Errorf("unsupported provider: %s. Available providers: anthropic, openai, copilot", provider)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -244,9 +256,31 @@ func runAuthStatus(cmd *cobra.Command, args []string) error {
|
||||
}
|
||||
}
|
||||
|
||||
// Check GitHub Copilot credentials
|
||||
fmt.Print("\nGitHub Copilot: ")
|
||||
if hasCopilotCreds, err := cm.HasCopilotCredentials(); err != nil {
|
||||
fmt.Printf("Error checking credentials: %v\n", err)
|
||||
} else if hasCopilotCreds {
|
||||
if creds, err := cm.GetCopilotCredentials(); err != nil {
|
||||
fmt.Printf("Error reading credentials: %v\n", err)
|
||||
} else {
|
||||
status := "✓ Authenticated"
|
||||
if creds.IsExpired() {
|
||||
status = "⚠️ Token expired (will refresh automatically)"
|
||||
} else if creds.NeedsRefresh() {
|
||||
status = "⚠️ Token expires soon (will refresh automatically)"
|
||||
}
|
||||
|
||||
fmt.Printf("%s (GitHub OAuth, stored %s)\n", status, creds.CreatedAt.Format("2006-01-02 15:04:05"))
|
||||
}
|
||||
} else {
|
||||
fmt.Println("✗ Not authenticated")
|
||||
}
|
||||
|
||||
fmt.Println("\nTo authenticate with a provider:")
|
||||
fmt.Println(" kit auth login anthropic")
|
||||
fmt.Println(" kit auth login openai")
|
||||
fmt.Println(" kit auth login copilot")
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -517,6 +551,85 @@ func loginOpenAI() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// loginCopilot authenticates GitHub Copilot using GitHub device flow.
|
||||
func loginCopilot(ctx context.Context) error {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
|
||||
cm, err := kit.NewCredentialManager()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to initialize credential manager: %w", err)
|
||||
}
|
||||
|
||||
if hasAuth, err := cm.HasCopilotCredentials(); err == nil && hasAuth {
|
||||
var reauth bool
|
||||
err := huh.NewConfirm().
|
||||
Title("You are already authenticated with GitHub Copilot").
|
||||
Description("Do you want to re-authenticate?").
|
||||
Affirmative("Yes").
|
||||
Negative("No").
|
||||
Value(&reauth).
|
||||
Run()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to prompt for re-authentication: %w", err)
|
||||
}
|
||||
if !reauth {
|
||||
fmt.Println("Authentication cancelled.")
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
client := auth.NewCopilotOAuthClient()
|
||||
|
||||
fmt.Println("🔐 Starting GitHub Copilot authentication...")
|
||||
fmt.Println("This uses GitHub device login and requires an active GitHub Copilot subscription.")
|
||||
fmt.Println("Experimental: this uses VS Code Copilot Chat client identifiers.")
|
||||
fmt.Println()
|
||||
|
||||
deviceCode, err := client.StartDeviceFlow(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to start GitHub device login: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("📱 Open this page and enter the code:")
|
||||
fmt.Printf("\n%s\n\n", deviceCode.VerificationURI)
|
||||
fmt.Printf("Code: %s\n\n", deviceCode.UserCode)
|
||||
auth.TryOpenBrowser(deviceCode.VerificationURI)
|
||||
|
||||
fmt.Println("Waiting for GitHub authorization...")
|
||||
githubToken, err := client.PollDeviceToken(ctx, deviceCode)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to complete GitHub device login: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("\n🔄 Exchanging GitHub token for Copilot access token...")
|
||||
creds, err := client.ExchangeGitHubToken(ctx, githubToken)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get GitHub Copilot token: %w", err)
|
||||
}
|
||||
|
||||
if err := cm.SetCopilotOAuthCredentials(creds); err != nil {
|
||||
return fmt.Errorf("failed to store credentials: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("✅ Successfully authenticated with GitHub Copilot!")
|
||||
fmt.Printf("📁 Credentials stored in: %s\n", cm.GetCredentialsPath())
|
||||
fmt.Println("\n🎉 Your GitHub Copilot credentials will now be used for copilot/* models.")
|
||||
fmt.Println("💡 You can check your authentication status with: kit auth status")
|
||||
|
||||
if err := setDefaultModelIfRequested("copilot"); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if !loginSetDefault {
|
||||
fmt.Println("\n💡 To set Copilot as your default model, run:")
|
||||
fmt.Println(" kit auth login copilot --set-default")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// callbackServer holds the HTTP server and channel for receiving the OAuth callback
|
||||
type callbackServer struct {
|
||||
Server *http.Server
|
||||
@@ -635,3 +748,43 @@ func logoutOpenAI() error {
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func logoutCopilot() error {
|
||||
cm, err := kit.NewCredentialManager()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to initialize credential manager: %w", err)
|
||||
}
|
||||
|
||||
hasAuth, err := cm.HasCopilotCredentials()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to check authentication status: %w", err)
|
||||
}
|
||||
|
||||
if !hasAuth {
|
||||
fmt.Println("You are not currently authenticated with GitHub Copilot.")
|
||||
return nil
|
||||
}
|
||||
|
||||
var confirm bool
|
||||
err = huh.NewConfirm().
|
||||
Title("Remove GitHub Copilot credentials").
|
||||
Description("Are you sure you want to remove your stored credentials?").
|
||||
Affirmative("Yes").
|
||||
Negative("No").
|
||||
Value(&confirm).
|
||||
Run()
|
||||
if err != nil || !confirm {
|
||||
fmt.Println("Logout cancelled.")
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := cm.RemoveCopilotCredentials(); err != nil {
|
||||
return fmt.Errorf("failed to remove credentials: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("✓ Successfully logged out from GitHub Copilot!")
|
||||
fmt.Println("You will need to authenticate again with 'kit auth login copilot'.")
|
||||
fmt.Println("Tip: this removes local credentials only. Revoke the GitHub OAuth grant at https://github.com/settings/applications")
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
+264
-429
@@ -4,13 +4,11 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/spf13/viper"
|
||||
"golang.org/x/term"
|
||||
|
||||
"github.com/mark3labs/kit/internal/app"
|
||||
"github.com/mark3labs/kit/internal/auth"
|
||||
"github.com/mark3labs/kit/internal/extbridge"
|
||||
"github.com/mark3labs/kit/internal/extensions"
|
||||
"github.com/mark3labs/kit/internal/models"
|
||||
@@ -35,439 +33,276 @@ type extensionContextDeps struct {
|
||||
// the three print routes appropriately for their phase (startup buffering
|
||||
// vs. live runtime routing).
|
||||
//
|
||||
// This consolidates two near-identical 400-line literal expressions that
|
||||
// previously appeared inline in runNormalMode.
|
||||
// The headless half (data access, state, options, tree navigation, skills,
|
||||
// templates, model resolution, subagents) comes from extbridge.BaseContext;
|
||||
// this function overlays the TUI-specific fields and overrides SetModel /
|
||||
// ReloadExtensions with TUI-aware versions.
|
||||
func buildInteractiveExtensionContext(deps extensionContextDeps) extensions.Context {
|
||||
kitInstance := deps.kitInstance
|
||||
appInstance := deps.appInstance
|
||||
usageTracker := deps.usageTracker
|
||||
ctx := deps.ctx
|
||||
|
||||
return extensions.Context{
|
||||
CWD: deps.cwd,
|
||||
Model: deps.modelName,
|
||||
Interactive: deps.interactive,
|
||||
PrintBlock: func(opts extensions.PrintBlockOpts) {
|
||||
appInstance.PrintBlockFromExtension(opts)
|
||||
},
|
||||
SendMessage: func(text string) { appInstance.Run(text) },
|
||||
CancelAndSend: func(text string) { appInstance.InterruptAndSend(text) },
|
||||
Abort: func() { appInstance.Abort() },
|
||||
IsIdle: func() bool { return !appInstance.IsBusy() },
|
||||
Compact: func(cfg extensions.CompactConfig) error {
|
||||
return appInstance.CompactAsync(cfg.CustomInstructions, cfg.OnComplete, cfg.OnError)
|
||||
},
|
||||
SendMultimodalMessage: func(text string, files []extensions.FilePart) {
|
||||
parts := make([]kit.LLMFilePart, len(files))
|
||||
for i, f := range files {
|
||||
parts[i] = kit.LLMFilePart{
|
||||
Filename: f.Filename,
|
||||
Data: f.Data,
|
||||
MediaType: f.MediaType,
|
||||
}
|
||||
}
|
||||
appInstance.RunWithFiles(text, parts)
|
||||
},
|
||||
GetSessionUsage: func() extensions.SessionUsage {
|
||||
if usageTracker == nil {
|
||||
return extensions.SessionUsage{}
|
||||
}
|
||||
stats := usageTracker.GetSessionStats()
|
||||
return extensions.SessionUsage{
|
||||
TotalInputTokens: stats.TotalInputTokens,
|
||||
TotalOutputTokens: stats.TotalOutputTokens,
|
||||
TotalCacheReadTokens: stats.TotalCacheReadTokens,
|
||||
TotalCacheWriteTokens: stats.TotalCacheWriteTokens,
|
||||
TotalCost: stats.TotalCost,
|
||||
RequestCount: stats.RequestCount,
|
||||
}
|
||||
},
|
||||
Exit: func() { appInstance.QuitFromExtension() },
|
||||
SetWidget: func(config extensions.WidgetConfig) {
|
||||
kitInstance.Extensions().SetWidget(config)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
RemoveWidget: func(id string) {
|
||||
kitInstance.Extensions().RemoveWidget(id)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
SetHeader: func(config extensions.HeaderFooterConfig) {
|
||||
kitInstance.Extensions().SetHeader(config)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
RemoveHeader: func() {
|
||||
kitInstance.Extensions().RemoveHeader()
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
SetFooter: func(config extensions.HeaderFooterConfig) {
|
||||
kitInstance.Extensions().SetFooter(config)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
RemoveFooter: func() {
|
||||
kitInstance.Extensions().RemoveFooter()
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
PromptSelect: func(config extensions.PromptSelectConfig) extensions.PromptSelectResult {
|
||||
ch := make(chan app.PromptResponse, 1)
|
||||
appInstance.SendPromptRequest(app.PromptRequestEvent{
|
||||
PromptType: "select",
|
||||
Message: config.Message,
|
||||
Options: config.Options,
|
||||
ResponseCh: ch,
|
||||
})
|
||||
resp := <-ch
|
||||
if resp.Cancelled {
|
||||
return extensions.PromptSelectResult{Cancelled: true}
|
||||
}
|
||||
return extensions.PromptSelectResult{Value: resp.Value, Index: resp.Index}
|
||||
},
|
||||
PromptConfirm: func(config extensions.PromptConfirmConfig) extensions.PromptConfirmResult {
|
||||
ch := make(chan app.PromptResponse, 1)
|
||||
def := "false"
|
||||
if config.DefaultValue {
|
||||
def = "true"
|
||||
}
|
||||
appInstance.SendPromptRequest(app.PromptRequestEvent{
|
||||
PromptType: "confirm",
|
||||
Message: config.Message,
|
||||
Default: def,
|
||||
ResponseCh: ch,
|
||||
})
|
||||
resp := <-ch
|
||||
if resp.Cancelled {
|
||||
return extensions.PromptConfirmResult{Cancelled: true}
|
||||
}
|
||||
return extensions.PromptConfirmResult{Value: resp.Confirmed}
|
||||
},
|
||||
PromptInput: func(config extensions.PromptInputConfig) extensions.PromptInputResult {
|
||||
ch := make(chan app.PromptResponse, 1)
|
||||
appInstance.SendPromptRequest(app.PromptRequestEvent{
|
||||
PromptType: "input",
|
||||
Message: config.Message,
|
||||
Placeholder: config.Placeholder,
|
||||
Default: config.Default,
|
||||
ResponseCh: ch,
|
||||
})
|
||||
resp := <-ch
|
||||
if resp.Cancelled {
|
||||
return extensions.PromptInputResult{Cancelled: true}
|
||||
}
|
||||
return extensions.PromptInputResult{Value: resp.Value}
|
||||
},
|
||||
SetUIVisibility: func(v extensions.UIVisibility) {
|
||||
kitInstance.Extensions().SetUIVisibility(v)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
GetContextStats: func() extensions.ContextStats {
|
||||
s := kitInstance.GetContextStats()
|
||||
return extensions.ContextStats{
|
||||
EstimatedTokens: s.EstimatedTokens,
|
||||
ContextLimit: s.ContextLimit,
|
||||
UsagePercent: s.UsagePercent,
|
||||
MessageCount: s.MessageCount,
|
||||
}
|
||||
},
|
||||
SetEditor: func(config extensions.EditorConfig) {
|
||||
kitInstance.Extensions().SetEditor(config)
|
||||
// Always use a goroutine for NotifyWidgetUpdate: prog.Send()
|
||||
// deadlocks if called synchronously from inside BubbleTea's
|
||||
// Update() handler. All call sites use go-routines uniformly.
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
ResetEditor: func() {
|
||||
kitInstance.Extensions().ResetEditor()
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
GetMessages: func() []extensions.SessionMessage {
|
||||
return kitInstance.Extensions().GetSessionMessages()
|
||||
},
|
||||
GetSessionPath: func() string {
|
||||
return kitInstance.GetSessionPath()
|
||||
},
|
||||
AppendEntry: func(entryType string, data string) (string, error) {
|
||||
return kitInstance.Extensions().AppendEntry(entryType, data)
|
||||
},
|
||||
GetEntries: func(entryType string) []extensions.ExtensionEntry {
|
||||
return kitInstance.Extensions().GetEntries(entryType)
|
||||
},
|
||||
SetEditorText: func(text string) {
|
||||
appInstance.SetEditorTextFromExtension(text)
|
||||
},
|
||||
SetStatus: func(key string, text string, priority int) {
|
||||
kitInstance.Extensions().SetStatus(extensions.StatusBarEntry{
|
||||
Key: key,
|
||||
Text: text,
|
||||
Priority: priority,
|
||||
})
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
RemoveStatus: func(key string) {
|
||||
kitInstance.Extensions().RemoveStatus(key)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
},
|
||||
GetOption: func(name string) string {
|
||||
return kitInstance.Extensions().GetOption(name)
|
||||
},
|
||||
SetOption: func(name string, value string) {
|
||||
kitInstance.Extensions().SetOption(name, value)
|
||||
},
|
||||
SetModel: func(modelString string) error {
|
||||
// Capture previous model for the ModelChange event.
|
||||
previousModel := kitInstance.Extensions().GetContext().Model
|
||||
err := kitInstance.SetModel(context.Background(), modelString)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Notify TUI so it updates model in status bar.
|
||||
p, m, _ := models.ParseModelString(modelString)
|
||||
appInstance.NotifyModelChanged(p, m)
|
||||
// Update the context's Model field so handlers see it.
|
||||
kitInstance.Extensions().UpdateContextModel(modelString)
|
||||
// Fire OnModelChange event to extensions.
|
||||
kitInstance.Extensions().EmitModelChange(modelString, previousModel, "extension")
|
||||
// Update usage tracker with new model info for correct token counting.
|
||||
if usageTracker != nil {
|
||||
newProvider, newModel, _ := models.ParseModelString(modelString)
|
||||
if newProvider != "unknown" && newModel != "unknown" && newProvider != "ollama" {
|
||||
registry := models.GetGlobalRegistry()
|
||||
if modelInfo := registry.LookupModel(newProvider, newModel); modelInfo != nil {
|
||||
// Check OAuth status for Anthropic models
|
||||
isOAuth := false
|
||||
if newProvider == "anthropic" {
|
||||
_, source, err := auth.GetAnthropicAPIKey(viper.GetString("provider-api-key"))
|
||||
if err == nil && strings.HasPrefix(source, "stored OAuth") {
|
||||
isOAuth = true
|
||||
}
|
||||
}
|
||||
usageTracker.UpdateModelInfo(modelInfo, newProvider, isOAuth)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
},
|
||||
GetAvailableModels: func() []extensions.ModelInfoEntry {
|
||||
return kitInstance.GetAvailableModels()
|
||||
},
|
||||
EmitCustomEvent: func(name string, data string) {
|
||||
kitInstance.Extensions().EmitCustomEvent(name, data)
|
||||
},
|
||||
Complete: func(req extensions.CompleteRequest) (extensions.CompleteResponse, error) {
|
||||
return kitInstance.ExecuteCompletion(context.Background(), req)
|
||||
},
|
||||
SuspendTUI: func(callback func()) error {
|
||||
return appInstance.SuspendTUI(callback)
|
||||
},
|
||||
RenderMessage: func(rendererName, content string) {
|
||||
renderer := kitInstance.Extensions().GetMessageRenderer(rendererName)
|
||||
if renderer == nil || renderer.Render == nil {
|
||||
appInstance.PrintFromExtension("", content)
|
||||
return
|
||||
}
|
||||
w, _, _ := term.GetSize(int(os.Stdout.Fd()))
|
||||
if w == 0 {
|
||||
w = 80
|
||||
}
|
||||
rendered := renderer.Render(content, w)
|
||||
appInstance.PrintFromExtension("", rendered)
|
||||
},
|
||||
ReloadExtensions: func() error {
|
||||
err := kitInstance.Extensions().Reload()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Notify TUI that widgets/status/commands may have changed.
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
return nil
|
||||
},
|
||||
GetAllTools: func() []extensions.ToolInfo {
|
||||
return kitInstance.Extensions().GetToolInfos()
|
||||
},
|
||||
SetActiveTools: func(names []string) {
|
||||
kitInstance.Extensions().SetActiveTools(names)
|
||||
},
|
||||
RegisterTheme: func(name string, config extensions.ThemeColorConfig) {
|
||||
tc := func(c extensions.ThemeColor) [2]string { return [2]string{c.Light, c.Dark} }
|
||||
ui.RegisterThemeFromConfig(name,
|
||||
tc(config.Primary), tc(config.Secondary),
|
||||
tc(config.Success), tc(config.Warning),
|
||||
tc(config.Error), tc(config.Info),
|
||||
tc(config.Text), tc(config.Muted),
|
||||
tc(config.VeryMuted), tc(config.Background),
|
||||
tc(config.Border), tc(config.MutedBorder),
|
||||
tc(config.System), tc(config.Tool),
|
||||
tc(config.Accent), tc(config.Highlight),
|
||||
tc(config.MdHeading), tc(config.MdLink),
|
||||
tc(config.MdKeyword), tc(config.MdString),
|
||||
tc(config.MdNumber), tc(config.MdComment),
|
||||
)
|
||||
},
|
||||
SetTheme: func(name string) error {
|
||||
return ui.ApplyTheme(name)
|
||||
},
|
||||
ListThemes: func() []string {
|
||||
return ui.ListThemes()
|
||||
},
|
||||
ShowOverlay: func(config extensions.OverlayConfig) extensions.OverlayResult {
|
||||
ch := make(chan app.OverlayResponse, 1)
|
||||
appInstance.SendOverlayRequest(app.OverlayRequestEvent{
|
||||
Title: config.Title,
|
||||
Content: config.Content.Text,
|
||||
Markdown: config.Content.Markdown,
|
||||
BorderColor: config.Style.BorderColor,
|
||||
Background: config.Style.Background,
|
||||
Width: config.Width,
|
||||
MaxHeight: config.MaxHeight,
|
||||
Anchor: string(config.Anchor),
|
||||
Actions: config.Actions,
|
||||
ResponseCh: ch,
|
||||
})
|
||||
resp := <-ch
|
||||
if resp.Cancelled {
|
||||
return extensions.OverlayResult{Cancelled: true, Index: -1}
|
||||
}
|
||||
return extensions.OverlayResult{
|
||||
Action: resp.Action,
|
||||
Index: resp.Index,
|
||||
}
|
||||
},
|
||||
SpawnSubagent: func(config extensions.SubagentConfig) (*extensions.SubagentHandle, *extensions.SubagentResult, error) {
|
||||
return extbridge.SpawnSubagent(ctx, kitInstance, config)
|
||||
},
|
||||
// -------------------------------------------------------------------
|
||||
// Tree Navigation API
|
||||
// -------------------------------------------------------------------
|
||||
GetTreeNode: func(entryID string) *extensions.TreeNode {
|
||||
node := kitInstance.GetTreeNode(entryID)
|
||||
if node == nil {
|
||||
return nil
|
||||
}
|
||||
return &extensions.TreeNode{
|
||||
ID: node.ID,
|
||||
ParentID: node.ParentID,
|
||||
Type: node.Type,
|
||||
Role: node.Role,
|
||||
Content: node.Content,
|
||||
Model: node.Model,
|
||||
Provider: node.Provider,
|
||||
Timestamp: node.Timestamp,
|
||||
Children: node.Children,
|
||||
}
|
||||
},
|
||||
GetCurrentBranch: func() []extensions.TreeNode {
|
||||
nodes := kitInstance.GetCurrentBranch()
|
||||
result := make([]extensions.TreeNode, len(nodes))
|
||||
for i, n := range nodes {
|
||||
result[i] = extensions.TreeNode{
|
||||
ID: n.ID,
|
||||
ParentID: n.ParentID,
|
||||
Type: n.Type,
|
||||
Role: n.Role,
|
||||
Content: n.Content,
|
||||
Model: n.Model,
|
||||
Provider: n.Provider,
|
||||
Timestamp: n.Timestamp,
|
||||
Children: n.Children,
|
||||
}
|
||||
}
|
||||
return result
|
||||
},
|
||||
GetChildren: func(parentID string) []string {
|
||||
return kitInstance.GetChildren(parentID)
|
||||
},
|
||||
NavigateTo: func(entryID string) extensions.TreeNavigationResult {
|
||||
err := kitInstance.NavigateTo(entryID)
|
||||
if err != nil {
|
||||
return extensions.TreeNavigationResult{Success: false, Error: err.Error()}
|
||||
}
|
||||
return extensions.TreeNavigationResult{Success: true}
|
||||
},
|
||||
SummarizeBranch: func(fromID, toID string) string {
|
||||
summary, _ := kitInstance.SummarizeBranch(fromID, toID)
|
||||
return summary
|
||||
},
|
||||
CollapseBranch: func(fromID, toID, summary string) extensions.TreeNavigationResult {
|
||||
err := kitInstance.CollapseBranch(fromID, toID, summary)
|
||||
if err != nil {
|
||||
return extensions.TreeNavigationResult{Success: false, Error: err.Error()}
|
||||
}
|
||||
return extensions.TreeNavigationResult{Success: true}
|
||||
},
|
||||
ec := extbridge.BaseContext(deps.ctx, kitInstance)
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// Skill Loading API
|
||||
// -------------------------------------------------------------------
|
||||
LoadSkill: func(path string) (*extensions.Skill, string) {
|
||||
s, err := kitInstance.LoadSkillForExtension(path)
|
||||
return s, err
|
||||
},
|
||||
LoadSkillsFromDir: func(dir string) extensions.SkillLoadResult {
|
||||
return kitInstance.LoadSkillsFromDirForExtension(dir)
|
||||
},
|
||||
DiscoverSkills: func() extensions.SkillLoadResult {
|
||||
skills := kitInstance.DiscoverSkillsForExtension()
|
||||
return extensions.SkillLoadResult{Skills: skills}
|
||||
},
|
||||
InjectSkillAsContext: func(skillName string) string {
|
||||
skills := kitInstance.DiscoverSkillsForExtension()
|
||||
for _, s := range skills {
|
||||
if s.Name == skillName {
|
||||
appInstance.Run(fmt.Sprintf("<skill name=%q>\n%s\n</skill>", s.Name, s.Content))
|
||||
return ""
|
||||
}
|
||||
}
|
||||
return fmt.Sprintf("skill not found: %s", skillName)
|
||||
},
|
||||
InjectRawSkillAsContext: func(path string) string {
|
||||
s, err := kitInstance.LoadSkillForExtension(path)
|
||||
if err != "" {
|
||||
return err
|
||||
}
|
||||
appInstance.Run(fmt.Sprintf("<skill name=%q>\n%s\n</skill>", s.Name, s.Content))
|
||||
return ""
|
||||
},
|
||||
GetAvailableSkills: func() []extensions.Skill {
|
||||
return kitInstance.DiscoverSkillsForExtension()
|
||||
},
|
||||
ec.CWD = deps.cwd
|
||||
ec.Model = deps.modelName
|
||||
ec.Interactive = deps.interactive
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// Template Parsing API
|
||||
// -------------------------------------------------------------------
|
||||
ParseTemplate: func(name, content string) extensions.PromptTemplate {
|
||||
return kit.ParseTemplate(name, content)
|
||||
},
|
||||
RenderTemplate: func(tpl extensions.PromptTemplate, vars map[string]string) string {
|
||||
return kit.RenderTemplate(tpl, vars)
|
||||
},
|
||||
ParseArguments: func(input string, pattern extensions.ArgumentPattern) extensions.ParseResult {
|
||||
return kit.ParseArguments(input, pattern)
|
||||
},
|
||||
SimpleParseArguments: func(input string, count int) []string {
|
||||
return kit.SimpleParseArguments(input, count)
|
||||
},
|
||||
EvaluateModelConditional: func(condition string) bool {
|
||||
return kit.EvaluateModelConditional(kitInstance.Extensions().GetContext().Model, condition)
|
||||
},
|
||||
RenderWithModelConditionals: func(content string) string {
|
||||
return kit.RenderWithModelConditionals(content, kitInstance.Extensions().GetContext().Model)
|
||||
},
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// Model Resolution API
|
||||
// -------------------------------------------------------------------
|
||||
ResolveModelChain: func(preferences []string) extensions.ModelResolutionResult {
|
||||
return kit.ResolveModelChain(preferences)
|
||||
},
|
||||
GetModelCapabilities: func(model string) (extensions.ModelCapabilities, string) {
|
||||
return kit.GetModelCapabilities(model)
|
||||
},
|
||||
CheckModelAvailable: func(model string) bool {
|
||||
return kit.CheckModelAvailable(model)
|
||||
},
|
||||
GetCurrentProvider: func() string {
|
||||
return kit.GetCurrentProvider(kitInstance.Extensions().GetContext().Model)
|
||||
},
|
||||
GetCurrentModelID: func() string {
|
||||
return kit.GetCurrentModelID(kitInstance.Extensions().GetContext().Model)
|
||||
},
|
||||
ec.PrintBlock = func(opts extensions.PrintBlockOpts) {
|
||||
appInstance.PrintBlockFromExtension(opts)
|
||||
}
|
||||
ec.SendMessage = func(text string) { appInstance.Run(text) }
|
||||
ec.CancelAndSend = func(text string) { appInstance.InterruptAndSend(text) }
|
||||
ec.Abort = func() { appInstance.Abort() }
|
||||
ec.IsIdle = func() bool { return !appInstance.IsBusy() }
|
||||
ec.Compact = func(cfg extensions.CompactConfig) error {
|
||||
return appInstance.CompactAsync(cfg.CustomInstructions, cfg.OnComplete, cfg.OnError)
|
||||
}
|
||||
ec.SendMultimodalMessage = func(text string, files []extensions.FilePart) {
|
||||
parts := make([]kit.LLMFilePart, len(files))
|
||||
for i, f := range files {
|
||||
parts[i] = kit.LLMFilePart{
|
||||
Filename: f.Filename,
|
||||
Data: f.Data,
|
||||
MediaType: f.MediaType,
|
||||
}
|
||||
}
|
||||
appInstance.RunWithFiles(text, parts)
|
||||
}
|
||||
ec.GetSessionUsage = func() extensions.SessionUsage {
|
||||
if usageTracker == nil {
|
||||
return extensions.SessionUsage{}
|
||||
}
|
||||
stats := usageTracker.GetSessionStats()
|
||||
return extensions.SessionUsage{
|
||||
TotalInputTokens: stats.TotalInputTokens,
|
||||
TotalOutputTokens: stats.TotalOutputTokens,
|
||||
TotalCacheReadTokens: stats.TotalCacheReadTokens,
|
||||
TotalCacheWriteTokens: stats.TotalCacheWriteTokens,
|
||||
TotalCost: stats.TotalCost,
|
||||
RequestCount: stats.RequestCount,
|
||||
}
|
||||
}
|
||||
ec.Exit = func() { appInstance.QuitFromExtension() }
|
||||
|
||||
// TUI widgets/chrome — mutate runner state, then notify the TUI.
|
||||
// Always use a goroutine for NotifyWidgetUpdate: prog.Send() deadlocks
|
||||
// if called synchronously from inside BubbleTea's Update() handler.
|
||||
// All call sites use go-routines uniformly.
|
||||
ec.SetWidget = func(config extensions.WidgetConfig) {
|
||||
kitInstance.Extensions().SetWidget(config)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.RemoveWidget = func(id string) {
|
||||
kitInstance.Extensions().RemoveWidget(id)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.SetHeader = func(config extensions.HeaderFooterConfig) {
|
||||
kitInstance.Extensions().SetHeader(config)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.RemoveHeader = func() {
|
||||
kitInstance.Extensions().RemoveHeader()
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.SetFooter = func(config extensions.HeaderFooterConfig) {
|
||||
kitInstance.Extensions().SetFooter(config)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.RemoveFooter = func() {
|
||||
kitInstance.Extensions().RemoveFooter()
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.SetUIVisibility = func(v extensions.UIVisibility) {
|
||||
kitInstance.Extensions().SetUIVisibility(v)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.SetEditor = func(config extensions.EditorConfig) {
|
||||
kitInstance.Extensions().SetEditor(config)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.ResetEditor = func() {
|
||||
kitInstance.Extensions().ResetEditor()
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.SetEditorText = func(text string) {
|
||||
appInstance.SetEditorTextFromExtension(text)
|
||||
}
|
||||
ec.SetStatus = func(key string, text string, priority int) {
|
||||
kitInstance.Extensions().SetStatus(extensions.StatusBarEntry{
|
||||
Key: key,
|
||||
Text: text,
|
||||
Priority: priority,
|
||||
})
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
ec.RemoveStatus = func(key string) {
|
||||
kitInstance.Extensions().RemoveStatus(key)
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
}
|
||||
|
||||
// Interactive prompts — channel-based round trips through the TUI.
|
||||
ec.PromptSelect = func(config extensions.PromptSelectConfig) extensions.PromptSelectResult {
|
||||
ch := make(chan app.PromptResponse, 1)
|
||||
appInstance.SendPromptRequest(app.PromptRequestEvent{
|
||||
PromptType: "select",
|
||||
Message: config.Message,
|
||||
Options: config.Options,
|
||||
ResponseCh: ch,
|
||||
})
|
||||
resp := <-ch
|
||||
if resp.Cancelled {
|
||||
return extensions.PromptSelectResult{Cancelled: true}
|
||||
}
|
||||
return extensions.PromptSelectResult{Value: resp.Value, Index: resp.Index}
|
||||
}
|
||||
ec.PromptConfirm = func(config extensions.PromptConfirmConfig) extensions.PromptConfirmResult {
|
||||
ch := make(chan app.PromptResponse, 1)
|
||||
def := "false"
|
||||
if config.DefaultValue {
|
||||
def = "true"
|
||||
}
|
||||
appInstance.SendPromptRequest(app.PromptRequestEvent{
|
||||
PromptType: "confirm",
|
||||
Message: config.Message,
|
||||
Default: def,
|
||||
ResponseCh: ch,
|
||||
})
|
||||
resp := <-ch
|
||||
if resp.Cancelled {
|
||||
return extensions.PromptConfirmResult{Cancelled: true}
|
||||
}
|
||||
return extensions.PromptConfirmResult{Value: resp.Confirmed}
|
||||
}
|
||||
ec.PromptInput = func(config extensions.PromptInputConfig) extensions.PromptInputResult {
|
||||
ch := make(chan app.PromptResponse, 1)
|
||||
appInstance.SendPromptRequest(app.PromptRequestEvent{
|
||||
PromptType: "input",
|
||||
Message: config.Message,
|
||||
Placeholder: config.Placeholder,
|
||||
Default: config.Default,
|
||||
ResponseCh: ch,
|
||||
})
|
||||
resp := <-ch
|
||||
if resp.Cancelled {
|
||||
return extensions.PromptInputResult{Cancelled: true}
|
||||
}
|
||||
return extensions.PromptInputResult{Value: resp.Value}
|
||||
}
|
||||
ec.ShowOverlay = func(config extensions.OverlayConfig) extensions.OverlayResult {
|
||||
ch := make(chan app.OverlayResponse, 1)
|
||||
appInstance.SendOverlayRequest(app.OverlayRequestEvent{
|
||||
Title: config.Title,
|
||||
Content: config.Content.Text,
|
||||
Markdown: config.Content.Markdown,
|
||||
BorderColor: config.Style.BorderColor,
|
||||
Background: config.Style.Background,
|
||||
Width: config.Width,
|
||||
MaxHeight: config.MaxHeight,
|
||||
Anchor: string(config.Anchor),
|
||||
Actions: config.Actions,
|
||||
ResponseCh: ch,
|
||||
})
|
||||
resp := <-ch
|
||||
if resp.Cancelled {
|
||||
return extensions.OverlayResult{Cancelled: true, Index: -1}
|
||||
}
|
||||
return extensions.OverlayResult{
|
||||
Action: resp.Action,
|
||||
Index: resp.Index,
|
||||
}
|
||||
}
|
||||
ec.SuspendTUI = func(callback func()) error {
|
||||
return appInstance.SuspendTUI(callback)
|
||||
}
|
||||
|
||||
// TUI-aware model switch: also notifies the TUI status bar and
|
||||
// refreshes the usage tracker for correct token counting.
|
||||
ec.SetModel = func(modelString string) error {
|
||||
// Capture previous model for the ModelChange event.
|
||||
previousModel := kitInstance.Extensions().GetContext().Model
|
||||
err := kitInstance.SetModel(context.Background(), modelString)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Notify TUI so it updates model in status bar.
|
||||
p, m, _ := models.ParseModelString(modelString)
|
||||
appInstance.NotifyModelChanged(p, m)
|
||||
// Update the context's Model field so handlers see it.
|
||||
kitInstance.Extensions().UpdateContextModel(modelString)
|
||||
// Fire OnModelChange event to extensions.
|
||||
kitInstance.Extensions().EmitModelChange(modelString, previousModel, "extension")
|
||||
// Update usage tracker with new model info for correct token counting.
|
||||
ui.UpdateUsageTrackerForModel(usageTracker, modelString, viper.GetString("provider-api-key"))
|
||||
return nil
|
||||
}
|
||||
|
||||
ec.RenderMessage = func(rendererName, content string) {
|
||||
renderer := kitInstance.Extensions().GetMessageRenderer(rendererName)
|
||||
if renderer == nil || renderer.Render == nil {
|
||||
appInstance.PrintFromExtension("", content)
|
||||
return
|
||||
}
|
||||
w, _, _ := term.GetSize(int(os.Stdout.Fd()))
|
||||
if w == 0 {
|
||||
w = 80
|
||||
}
|
||||
rendered := renderer.Render(content, w)
|
||||
appInstance.PrintFromExtension("", rendered)
|
||||
}
|
||||
ec.ReloadExtensions = func() error {
|
||||
err := kitInstance.Extensions().Reload()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Notify TUI that widgets/status/commands may have changed.
|
||||
go appInstance.NotifyWidgetUpdate()
|
||||
return nil
|
||||
}
|
||||
|
||||
// Theme management (TUI only).
|
||||
ec.RegisterTheme = func(name string, config extensions.ThemeColorConfig) {
|
||||
tc := func(c extensions.ThemeColor) [2]string { return [2]string{c.Light, c.Dark} }
|
||||
ui.RegisterThemeFromConfig(name,
|
||||
tc(config.Primary), tc(config.Secondary),
|
||||
tc(config.Success), tc(config.Warning),
|
||||
tc(config.Error), tc(config.Info),
|
||||
tc(config.Text), tc(config.Muted),
|
||||
tc(config.VeryMuted), tc(config.Background),
|
||||
tc(config.Border), tc(config.MutedBorder),
|
||||
tc(config.System), tc(config.Tool),
|
||||
tc(config.Accent), tc(config.Highlight),
|
||||
tc(config.MdHeading), tc(config.MdLink),
|
||||
tc(config.MdKeyword), tc(config.MdString),
|
||||
tc(config.MdNumber), tc(config.MdComment),
|
||||
)
|
||||
}
|
||||
ec.SetTheme = func(name string) error {
|
||||
return ui.ApplyTheme(name)
|
||||
}
|
||||
ec.ListThemes = func() []string {
|
||||
return ui.ListThemes()
|
||||
}
|
||||
|
||||
// Skill context-injection (drives a new agent turn through the TUI).
|
||||
ec.InjectSkillAsContext = func(skillName string) string {
|
||||
skills := kitInstance.DiscoverSkillsForExtension()
|
||||
for _, s := range skills {
|
||||
if s.Name == skillName {
|
||||
appInstance.Run(fmt.Sprintf("<skill name=%q>\n%s\n</skill>", s.Name, s.Content))
|
||||
return ""
|
||||
}
|
||||
}
|
||||
return fmt.Sprintf("skill not found: %s", skillName)
|
||||
}
|
||||
ec.InjectRawSkillAsContext = func(path string) string {
|
||||
s, err := kitInstance.LoadSkillForExtension(path)
|
||||
if err != "" {
|
||||
return err
|
||||
}
|
||||
appInstance.Run(fmt.Sprintf("<skill name=%q>\n%s\n</skill>", s.Name, s.Content))
|
||||
return ""
|
||||
}
|
||||
|
||||
return ec
|
||||
}
|
||||
|
||||
+63
-40
@@ -12,7 +12,6 @@ import (
|
||||
|
||||
tea "charm.land/bubbletea/v2"
|
||||
"github.com/mark3labs/kit/internal/app"
|
||||
"github.com/mark3labs/kit/internal/auth"
|
||||
"github.com/mark3labs/kit/internal/config"
|
||||
"github.com/mark3labs/kit/internal/extensions"
|
||||
"github.com/mark3labs/kit/internal/models"
|
||||
@@ -677,8 +676,8 @@ func globalShortcutsProviderForUI(k *kit.Kit) func() map[string]func() {
|
||||
}
|
||||
}
|
||||
|
||||
func runNormalMode(ctx context.Context) error {
|
||||
// Validate flag combinations
|
||||
// validateModeFlags rejects invalid flag combinations for the root command.
|
||||
func validateModeFlags() error {
|
||||
if quietFlag && positionalPrompt == "" {
|
||||
return fmt.Errorf("--quiet requires a prompt (e.g. kit \"your question\" --quiet)")
|
||||
}
|
||||
@@ -691,21 +690,14 @@ func runNormalMode(ctx context.Context) error {
|
||||
if noExitFlag && positionalPrompt == "" {
|
||||
return fmt.Errorf("--no-exit requires a prompt (e.g. kit \"your question\" --no-exit)")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Set up logging
|
||||
if debugMode {
|
||||
log.SetFlags(log.LstdFlags | log.Lshortfile)
|
||||
}
|
||||
|
||||
// Update debug mode from viper
|
||||
if viper.GetBool("debug") && !debugMode {
|
||||
debugMode = viper.GetBool("debug")
|
||||
log.SetFlags(log.LstdFlags | log.Lshortfile)
|
||||
}
|
||||
|
||||
// Restore persisted model preference when no explicit --model flag or
|
||||
// config file model is set. Precedence: CLI flag > config file > saved
|
||||
// preference > built-in default. This mirrors how themes are persisted.
|
||||
// restorePersistedPreferences applies saved model / thinking-level
|
||||
// preferences into viper when neither a CLI flag nor a config-file value
|
||||
// takes precedence. Precedence: CLI flag > config file > saved preference >
|
||||
// built-in default. This mirrors how themes are persisted.
|
||||
func restorePersistedPreferences() {
|
||||
// Skip custom/* models unless --provider-url is also provided, since the
|
||||
// custom provider requires a URL that was only valid for the previous session.
|
||||
if !modelFlagChanged && !viper.InConfig("model") {
|
||||
@@ -724,6 +716,15 @@ func runNormalMode(ctx context.Context) error {
|
||||
viper.Set("thinking-level", pref)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// applyProviderURLRouting rewrites the model in viper when --provider-url
|
||||
// is set, routing requests through the "custom" (OpenAI-compatible)
|
||||
// provider. Must run after restorePersistedPreferences.
|
||||
func applyProviderURLRouting() {
|
||||
if viper.GetString("provider-url") == "" {
|
||||
return
|
||||
}
|
||||
|
||||
// When --provider-url is set but no explicit --model was provided,
|
||||
// default to "custom/custom" so the user doesn't need to remember a
|
||||
@@ -731,18 +732,53 @@ func runNormalMode(ctx context.Context) error {
|
||||
// This intentionally overrides saved preferences but respects config-file
|
||||
// models — if you specify a model in ~/.kit.yml, it will be used with
|
||||
// custom/custom's provider routing.
|
||||
if viper.GetString("provider-url") != "" && !modelFlagChanged && !viper.InConfig("model") {
|
||||
if !modelFlagChanged && !viper.InConfig("model") {
|
||||
viper.Set("model", "custom/custom")
|
||||
}
|
||||
|
||||
// When --provider-url is set with an explicit --model that lacks a provider
|
||||
// prefix (no "/"), auto-prefix with "custom/" for OpenAI-compatible endpoints.
|
||||
if viper.GetString("provider-url") != "" && modelFlagChanged {
|
||||
// When --provider-url is set with an explicit --model, route through the
|
||||
// "custom" provider (OpenAI-compatible wire). This honors the user's
|
||||
// intent: passing a custom URL means "use THIS endpoint", not "speak
|
||||
// the Google/Anthropic/etc. wire protocol against this endpoint".
|
||||
//
|
||||
// Any provider prefix on the model is stripped so a model name that
|
||||
// happens to collide with a known provider (e.g. `google/gemma-4-12b`
|
||||
// served by LM Studio) still resolves correctly. If you genuinely need
|
||||
// to point a non-OpenAI wire (Anthropic, Google, ...) at a proxy URL,
|
||||
// use the explicit `custom/<name>` form to opt out of the rewrite by
|
||||
// configuring the proxy as that provider in your config file instead.
|
||||
if modelFlagChanged {
|
||||
model := viper.GetString("model")
|
||||
if model != "" && !strings.Contains(model, "/") {
|
||||
viper.Set("model", "custom/"+model)
|
||||
if model != "" {
|
||||
name := model
|
||||
if _, after, ok := strings.Cut(model, "/"); ok {
|
||||
name = after
|
||||
}
|
||||
if !strings.HasPrefix(model, "custom/") {
|
||||
viper.Set("model", "custom/"+name)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func runNormalMode(ctx context.Context) error {
|
||||
if err := validateModeFlags(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Set up logging
|
||||
if debugMode {
|
||||
log.SetFlags(log.LstdFlags | log.Lshortfile)
|
||||
}
|
||||
|
||||
// Update debug mode from viper
|
||||
if viper.GetBool("debug") && !debugMode {
|
||||
debugMode = viper.GetBool("debug")
|
||||
log.SetFlags(log.LstdFlags | log.Lshortfile)
|
||||
}
|
||||
|
||||
restorePersistedPreferences()
|
||||
applyProviderURLRouting()
|
||||
|
||||
// Load MCP configuration.
|
||||
mcpConfig, err := config.LoadAndValidateConfig()
|
||||
@@ -916,6 +952,9 @@ func runNormalMode(ctx context.Context) error {
|
||||
startupExtensionMessages = append(startupExtensionMessages, text)
|
||||
}
|
||||
kitInstance.Extensions().SetContext(extCtx)
|
||||
if err := kitInstance.Extensions().InitStatePersistence(); err != nil {
|
||||
log.Printf("WARN extension state init failed: %v", err)
|
||||
}
|
||||
kitInstance.Extensions().EmitSessionStart()
|
||||
|
||||
// Restore normal print functions for runtime use.
|
||||
@@ -1146,23 +1185,7 @@ func runNormalMode(ctx context.Context) error {
|
||||
// NotifyModelChanged calls prog.Send() which deadlocks. The UI layer
|
||||
// updates m.providerName and m.modelName directly after setModel returns.
|
||||
// Update usage tracker with new model info for correct token counting.
|
||||
if usageTracker != nil {
|
||||
newProvider, newModel, _ := models.ParseModelString(modelString)
|
||||
if newProvider != "unknown" && newModel != "unknown" && newProvider != "ollama" {
|
||||
registry := models.GetGlobalRegistry()
|
||||
if modelInfo := registry.LookupModel(newProvider, newModel); modelInfo != nil {
|
||||
// Check OAuth status for Anthropic models
|
||||
isOAuth := false
|
||||
if newProvider == "anthropic" {
|
||||
_, source, err := auth.GetAnthropicAPIKey(viper.GetString("provider-api-key"))
|
||||
if err == nil && strings.HasPrefix(source, "stored OAuth") {
|
||||
isOAuth = true
|
||||
}
|
||||
}
|
||||
usageTracker.UpdateModelInfo(modelInfo, newProvider, isOAuth)
|
||||
}
|
||||
}
|
||||
}
|
||||
ui.UpdateUsageTrackerForModel(usageTracker, modelString, viper.GetString("provider-api-key"))
|
||||
return nil
|
||||
}
|
||||
emitModelChangeForUI := func(newModel, previousModel, source string) {
|
||||
|
||||
@@ -58,6 +58,7 @@ kit install github.com/mark3labs/kit/examples/extensions --local
|
||||
| `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
|
||||
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
//go:build ignore
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
|
||||
"kit/ext"
|
||||
)
|
||||
|
||||
// Init demonstrates the three primitives added in issue #53:
|
||||
//
|
||||
// 1. api.OnLLMUsage(...) — per-LLM-call usage callback with token + cost
|
||||
// deltas. Use this for budget enforcement that reacts between calls
|
||||
// within a single agent turn, rather than only at turn boundaries.
|
||||
//
|
||||
// 2. ctx.SetState / ctx.GetState / ctx.DeleteState / ctx.ListState —
|
||||
// last-write-wins, session-scoped key-value store backed by a sidecar
|
||||
// file. Use this for snapshot state (current value of X) instead of
|
||||
// ctx.AppendEntry, which is append-only and bloats branch reads.
|
||||
//
|
||||
// 3. ext.AgentEndEvent.ToolCallCount / .ToolNames / .LLMCallCount /
|
||||
// .InputTokensDelta / .OutputTokensDelta / .CostDelta / .DurationMs —
|
||||
// per-turn aggregates so observer extensions don't need to maintain
|
||||
// parallel bookkeeping.
|
||||
//
|
||||
// Together these support a simple soft-budget cap: warn when the
|
||||
// cumulative cost in this session exceeds a threshold, and print a
|
||||
// per-turn report on AgentEnd.
|
||||
//
|
||||
// Usage: kit -e examples/extensions/usage-budget.go
|
||||
func Init(api ext.API) {
|
||||
const warnAtKey = "usage-budget:warn-at-usd"
|
||||
|
||||
// 1. Print per-LLM-call usage with provider, model, and cost.
|
||||
api.OnLLMUsage(func(e ext.LLMUsageEvent, ctx ext.Context) {
|
||||
ctx.Print(fmt.Sprintf(
|
||||
"[usage] step=%d %s/%s tokens=↑%d ↓%d cache=↑%d/↓%d cost=$%.4f (%s)",
|
||||
e.StepNumber, e.Provider, e.Model,
|
||||
e.InputTokens, e.OutputTokens,
|
||||
e.CacheWriteTokens, e.CacheReadTokens,
|
||||
e.Cost, e.FinishReason,
|
||||
))
|
||||
|
||||
// 2. Persist running total in last-write-wins state.
|
||||
current := 0.0
|
||||
if raw, ok := ctx.GetState("usage-budget:total-cost"); ok {
|
||||
current, _ = strconv.ParseFloat(raw, 64)
|
||||
}
|
||||
current += e.Cost
|
||||
ctx.SetState("usage-budget:total-cost", strconv.FormatFloat(current, 'f', 6, 64))
|
||||
|
||||
// Soft warn-at threshold (configurable via state).
|
||||
warnAt := 0.50
|
||||
if raw, ok := ctx.GetState(warnAtKey); ok {
|
||||
if v, err := strconv.ParseFloat(raw, 64); err == nil {
|
||||
warnAt = v
|
||||
}
|
||||
}
|
||||
if current > warnAt {
|
||||
ctx.PrintError(fmt.Sprintf(
|
||||
"[usage] session cost $%.4f exceeds soft cap $%.2f",
|
||||
current, warnAt,
|
||||
))
|
||||
}
|
||||
})
|
||||
|
||||
// 3. Print a per-turn summary using the enriched AgentEndEvent.
|
||||
api.OnAgentEnd(func(e ext.AgentEndEvent, ctx ext.Context) {
|
||||
ctx.Print(fmt.Sprintf(
|
||||
"[turn] stop=%s tools=%d llm-calls=%d tokens=↑%d ↓%d cost=$%.4f duration=%dms",
|
||||
e.StopReason, e.ToolCallCount, e.LLMCallCount,
|
||||
e.InputTokensDelta, e.OutputTokensDelta, e.CostDelta, e.DurationMs,
|
||||
))
|
||||
if len(e.ToolNames) > 0 {
|
||||
ctx.Print(fmt.Sprintf("[turn] tool order: %v", e.ToolNames))
|
||||
}
|
||||
})
|
||||
|
||||
// Bootstrap default soft cap once per session.
|
||||
api.OnSessionStart(func(e ext.SessionStartEvent, ctx ext.Context) {
|
||||
if _, ok := ctx.GetState(warnAtKey); !ok {
|
||||
ctx.SetState(warnAtKey, "0.50")
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -73,111 +73,70 @@ func (r *sessionRegistry) create(ctx context.Context, cwd string) (*acpSession,
|
||||
|
||||
// Wire extension context with headless implementations so extensions
|
||||
// work in ACP mode. TUI-dependent features (widgets, prompts, editor)
|
||||
// become no-ops or return cancelled; all data/model/tool APIs work
|
||||
// identically to interactive mode.
|
||||
// become no-ops or return cancelled; all data/model/tool APIs come from
|
||||
// extbridge.BaseContext and work identically to interactive mode.
|
||||
if kitInstance.Extensions().HasExtensions() {
|
||||
kitInstance.Extensions().SetContext(extensions.Context{
|
||||
SessionID: sessionID,
|
||||
CWD: cwd,
|
||||
Model: kitInstance.GetModelString(),
|
||||
Interactive: false,
|
||||
// Use a background context for subagent spawns: the create() ctx is
|
||||
// request-scoped and may be cancelled before extensions spawn anything.
|
||||
ec := extbridge.BaseContext(context.Background(), kitInstance)
|
||||
|
||||
// Output — route through structured logger.
|
||||
Print: func(text string) { log.Debug("extension: print", "text", text) },
|
||||
PrintInfo: func(text string) { log.Info("extension: info", "text", text) },
|
||||
PrintError: func(text string) { log.Error("extension: error", "text", text) },
|
||||
PrintBlock: func(opts extensions.PrintBlockOpts) {
|
||||
log.Info("extension: block", "subtitle", opts.Subtitle, "text", opts.Text)
|
||||
},
|
||||
ec.SessionID = sessionID
|
||||
ec.CWD = cwd
|
||||
ec.Model = kitInstance.GetModelString()
|
||||
ec.Interactive = false
|
||||
|
||||
// Message injection — no-ops for now; ACP clients drive prompts.
|
||||
SendMessage: func(string) {},
|
||||
CancelAndSend: func(string) {},
|
||||
Exit: func() {},
|
||||
// Output — route through structured logger.
|
||||
ec.Print = func(text string) { log.Debug("extension: print", "text", text) }
|
||||
ec.PrintInfo = func(text string) { log.Info("extension: info", "text", text) }
|
||||
ec.PrintError = func(text string) { log.Error("extension: error", "text", text) }
|
||||
ec.PrintBlock = func(opts extensions.PrintBlockOpts) {
|
||||
log.Info("extension: block", "subtitle", opts.Subtitle, "text", opts.Text)
|
||||
}
|
||||
|
||||
// TUI widgets/chrome — silent no-ops (no TUI in ACP).
|
||||
SetWidget: func(extensions.WidgetConfig) {},
|
||||
RemoveWidget: func(string) {},
|
||||
SetHeader: func(extensions.HeaderFooterConfig) {},
|
||||
RemoveHeader: func() {},
|
||||
SetFooter: func(extensions.HeaderFooterConfig) {},
|
||||
RemoveFooter: func() {},
|
||||
SetEditor: func(extensions.EditorConfig) {},
|
||||
ResetEditor: func() {},
|
||||
SetEditorText: func(string) {},
|
||||
SetUIVisibility: func(extensions.UIVisibility) {},
|
||||
SetStatus: func(string, string, int) {},
|
||||
RemoveStatus: func(string) {},
|
||||
// Message injection — no-ops for now; ACP clients drive prompts.
|
||||
ec.SendMessage = func(string) {}
|
||||
ec.CancelAndSend = func(string) {}
|
||||
ec.Exit = func() {}
|
||||
|
||||
// Interactive prompts — return cancelled (no user to prompt).
|
||||
PromptSelect: func(extensions.PromptSelectConfig) extensions.PromptSelectResult {
|
||||
return extensions.PromptSelectResult{Cancelled: true}
|
||||
},
|
||||
PromptConfirm: func(extensions.PromptConfirmConfig) extensions.PromptConfirmResult {
|
||||
return extensions.PromptConfirmResult{Cancelled: true}
|
||||
},
|
||||
PromptInput: func(extensions.PromptInputConfig) extensions.PromptInputResult {
|
||||
return extensions.PromptInputResult{Cancelled: true}
|
||||
},
|
||||
ShowOverlay: func(extensions.OverlayConfig) extensions.OverlayResult {
|
||||
return extensions.OverlayResult{Cancelled: true, Index: -1}
|
||||
},
|
||||
SuspendTUI: func(callback func()) error { callback(); return nil },
|
||||
// TUI widgets/chrome — silent no-ops (no TUI in ACP).
|
||||
ec.SetWidget = func(extensions.WidgetConfig) {}
|
||||
ec.RemoveWidget = func(string) {}
|
||||
ec.SetHeader = func(extensions.HeaderFooterConfig) {}
|
||||
ec.RemoveHeader = func() {}
|
||||
ec.SetFooter = func(extensions.HeaderFooterConfig) {}
|
||||
ec.RemoveFooter = func() {}
|
||||
ec.SetEditor = func(extensions.EditorConfig) {}
|
||||
ec.ResetEditor = func() {}
|
||||
ec.SetEditorText = func(string) {}
|
||||
ec.SetUIVisibility = func(extensions.UIVisibility) {}
|
||||
ec.SetStatus = func(string, string, int) {}
|
||||
ec.RemoveStatus = func(string) {}
|
||||
|
||||
// Data access — delegate to Kit instance.
|
||||
GetContextStats: func() extensions.ContextStats {
|
||||
s := kitInstance.GetContextStats()
|
||||
return extensions.ContextStats{
|
||||
EstimatedTokens: s.EstimatedTokens,
|
||||
ContextLimit: s.ContextLimit,
|
||||
UsagePercent: s.UsagePercent,
|
||||
MessageCount: s.MessageCount,
|
||||
}
|
||||
},
|
||||
GetMessages: func() []extensions.SessionMessage { return kitInstance.Extensions().GetSessionMessages() },
|
||||
GetSessionPath: func() string { return kitInstance.GetSessionPath() },
|
||||
AppendEntry: func(entryType, data string) (string, error) {
|
||||
return kitInstance.Extensions().AppendEntry(entryType, data)
|
||||
},
|
||||
GetEntries: func(entryType string) []extensions.ExtensionEntry {
|
||||
return kitInstance.Extensions().GetEntries(entryType)
|
||||
},
|
||||
// Interactive prompts — return cancelled (no user to prompt).
|
||||
ec.PromptSelect = func(extensions.PromptSelectConfig) extensions.PromptSelectResult {
|
||||
return extensions.PromptSelectResult{Cancelled: true}
|
||||
}
|
||||
ec.PromptConfirm = func(extensions.PromptConfirmConfig) extensions.PromptConfirmResult {
|
||||
return extensions.PromptConfirmResult{Cancelled: true}
|
||||
}
|
||||
ec.PromptInput = func(extensions.PromptInputConfig) extensions.PromptInputResult {
|
||||
return extensions.PromptInputResult{Cancelled: true}
|
||||
}
|
||||
ec.ShowOverlay = func(extensions.OverlayConfig) extensions.OverlayResult {
|
||||
return extensions.OverlayResult{Cancelled: true, Index: -1}
|
||||
}
|
||||
ec.SuspendTUI = func(callback func()) error { callback(); return nil }
|
||||
|
||||
// Options, model, and tool management.
|
||||
GetOption: func(name string) string { return kitInstance.Extensions().GetOption(name) },
|
||||
SetOption: func(name, value string) { kitInstance.Extensions().SetOption(name, value) },
|
||||
SetModel: func(modelString string) error {
|
||||
previousModel := kitInstance.Extensions().GetContext().Model
|
||||
if err := kitInstance.SetModel(context.Background(), modelString); err != nil {
|
||||
return err
|
||||
}
|
||||
kitInstance.Extensions().UpdateContextModel(modelString)
|
||||
kitInstance.Extensions().EmitModelChange(modelString, previousModel, "extension")
|
||||
return nil
|
||||
},
|
||||
GetAvailableModels: func() []extensions.ModelInfoEntry { return kitInstance.GetAvailableModels() },
|
||||
EmitCustomEvent: func(name, data string) { kitInstance.Extensions().EmitCustomEvent(name, data) },
|
||||
GetAllTools: func() []extensions.ToolInfo { return kitInstance.Extensions().GetToolInfos() },
|
||||
SetActiveTools: func(names []string) { kitInstance.Extensions().SetActiveTools(names) },
|
||||
// Render — fall back to logging.
|
||||
ec.RenderMessage = func(name, content string) {
|
||||
renderer := kitInstance.Extensions().GetMessageRenderer(name)
|
||||
if renderer != nil && renderer.Render != nil {
|
||||
content = renderer.Render(content, 80)
|
||||
}
|
||||
log.Info("extension: message", "renderer", name, "content", content)
|
||||
}
|
||||
|
||||
// LLM completions and subagents.
|
||||
Complete: func(req extensions.CompleteRequest) (extensions.CompleteResponse, error) {
|
||||
return kitInstance.ExecuteCompletion(context.Background(), req)
|
||||
},
|
||||
SpawnSubagent: func(config extensions.SubagentConfig) (*extensions.SubagentHandle, *extensions.SubagentResult, error) {
|
||||
return extbridge.SpawnSubagent(context.Background(), kitInstance, config)
|
||||
},
|
||||
|
||||
// Render — fall back to logging.
|
||||
RenderMessage: func(name, content string) {
|
||||
renderer := kitInstance.Extensions().GetMessageRenderer(name)
|
||||
if renderer != nil && renderer.Render != nil {
|
||||
content = renderer.Render(content, 80)
|
||||
}
|
||||
log.Info("extension: message", "renderer", name, "content", content)
|
||||
},
|
||||
ReloadExtensions: func() error { return kitInstance.Extensions().Reload() },
|
||||
})
|
||||
kitInstance.Extensions().SetContext(ec)
|
||||
kitInstance.Extensions().EmitSessionStart()
|
||||
}
|
||||
|
||||
|
||||
@@ -13,11 +13,6 @@ type MessageStore struct {
|
||||
messages []kit.LLMMessage
|
||||
}
|
||||
|
||||
// NewMessageStore creates an empty MessageStore.
|
||||
func NewMessageStore() *MessageStore {
|
||||
return &MessageStore{}
|
||||
}
|
||||
|
||||
// NewMessageStoreWithMessages creates a MessageStore pre-populated with the
|
||||
// given messages. This is used when loading an existing session at startup.
|
||||
func NewMessageStoreWithMessages(msgs []kit.LLMMessage) *MessageStore {
|
||||
|
||||
@@ -29,7 +29,7 @@ func textOf(msg kit.LLMMessage) string {
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestNewMessageStore_empty(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
if s == nil {
|
||||
t.Fatal("expected non-nil store")
|
||||
}
|
||||
@@ -72,7 +72,7 @@ func TestNewMessageStoreWithMessages_isolatesInput(t *testing.T) {
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestAdd_appendsMessage(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
s.Add(makeTextMsg("user", "first"))
|
||||
s.Add(makeTextMsg("assistant", "second"))
|
||||
|
||||
@@ -82,7 +82,7 @@ func TestAdd_appendsMessage(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAdd_preservesOrder(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
texts := []string{"a", "b", "c"}
|
||||
for _, t2 := range texts {
|
||||
s.Add(makeTextMsg("user", t2))
|
||||
@@ -100,7 +100,7 @@ func TestAdd_preservesOrder(t *testing.T) {
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestReplace_swapsHistory(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
s.Add(makeTextMsg("user", "old"))
|
||||
|
||||
replacement := []kit.LLMMessage{
|
||||
@@ -120,7 +120,7 @@ func TestReplace_swapsHistory(t *testing.T) {
|
||||
|
||||
// Replace must deep-copy the incoming slice.
|
||||
func TestReplace_isolatesInput(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
replacement := []kit.LLMMessage{makeTextMsg("user", "original")}
|
||||
s.Replace(replacement)
|
||||
|
||||
@@ -137,7 +137,7 @@ func TestReplace_isolatesInput(t *testing.T) {
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestGetAll_returnsCopy(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
s.Add(makeTextMsg("user", "hello"))
|
||||
|
||||
got := s.GetAll()
|
||||
@@ -151,7 +151,7 @@ func TestGetAll_returnsCopy(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestGetAll_emptyStore(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
got := s.GetAll()
|
||||
if len(got) != 0 {
|
||||
t.Fatalf("expected empty slice, got %d elements", len(got))
|
||||
@@ -163,7 +163,7 @@ func TestGetAll_emptyStore(t *testing.T) {
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestClear_removesAllMessages(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
s.Add(makeTextMsg("user", "a"))
|
||||
s.Add(makeTextMsg("user", "b"))
|
||||
s.Clear()
|
||||
@@ -174,7 +174,7 @@ func TestClear_removesAllMessages(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestClear_allowsSubsequentAdds(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
s.Add(makeTextMsg("user", "before"))
|
||||
s.Clear()
|
||||
s.Add(makeTextMsg("user", "after"))
|
||||
@@ -193,7 +193,7 @@ func TestClear_allowsSubsequentAdds(t *testing.T) {
|
||||
// --------------------------------------------------------------------------
|
||||
|
||||
func TestConcurrentAccess(t *testing.T) {
|
||||
s := NewMessageStore()
|
||||
s := NewMessageStoreWithMessages(nil)
|
||||
done := make(chan struct{})
|
||||
|
||||
// Writer goroutine.
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
@@ -9,11 +10,11 @@ import (
|
||||
"time"
|
||||
)
|
||||
|
||||
// CredentialStore holds all stored credentials for various providers.
|
||||
// Currently supports Anthropic and OpenAI credentials with both OAuth and API key authentication methods.
|
||||
// CredentialStore holds stored credentials for Anthropic, OpenAI, and GitHub Copilot.
|
||||
type CredentialStore struct {
|
||||
Anthropic *AnthropicCredentials `json:"anthropic,omitempty"`
|
||||
OpenAI *OpenAICredentials `json:"openai,omitempty"`
|
||||
Copilot *CopilotCredentials `json:"copilot,omitempty"`
|
||||
}
|
||||
|
||||
// AnthropicCredentials holds Anthropic API credentials supporting both OAuth
|
||||
@@ -43,6 +44,16 @@ type OpenAICredentials struct {
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// CopilotCredentials holds GitHub OAuth credentials and the short-lived
|
||||
// GitHub Copilot API token derived from them.
|
||||
type CopilotCredentials struct {
|
||||
Type string `json:"type"` // "oauth"
|
||||
GitHubToken string `json:"github_token,omitempty"` // GitHub device-flow OAuth token
|
||||
CopilotAccessToken string `json:"copilot_access_token,omitempty"` // Short-lived Copilot API token
|
||||
ExpiresAt int64 `json:"expires_at,omitempty"` // Copilot token expiry
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
}
|
||||
|
||||
// oauthTokenExpired reports whether an OAuth token with the given type and
|
||||
// expiry unix timestamp is past its expiry. Returns false for API key
|
||||
// credentials or when no expiry is set.
|
||||
@@ -91,6 +102,16 @@ func (c *OpenAICredentials) NeedsRefresh() bool {
|
||||
return oauthTokenNeedsRefresh(c.Type, c.ExpiresAt)
|
||||
}
|
||||
|
||||
// IsExpired checks if the Copilot API token is expired.
|
||||
func (c *CopilotCredentials) IsExpired() bool {
|
||||
return oauthTokenExpired(c.Type, c.ExpiresAt)
|
||||
}
|
||||
|
||||
// NeedsRefresh reports whether the Copilot API token should be renewed.
|
||||
func (c *CopilotCredentials) NeedsRefresh() bool {
|
||||
return oauthTokenNeedsRefresh(c.Type, c.ExpiresAt)
|
||||
}
|
||||
|
||||
// CredentialManager handles secure storage and retrieval of authentication credentials.
|
||||
// It manages a JSON file stored in the user's config directory with appropriate
|
||||
// file permissions for security.
|
||||
@@ -222,7 +243,7 @@ func (cm *CredentialManager) RemoveAnthropicCredentials() error {
|
||||
store.Anthropic = nil
|
||||
|
||||
// If store is empty, remove the file entirely
|
||||
if store.Anthropic == nil {
|
||||
if store.Anthropic == nil && store.OpenAI == nil && store.Copilot == nil {
|
||||
if err := os.Remove(cm.credentialsPath); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("failed to remove credentials file: %w", err)
|
||||
}
|
||||
@@ -279,7 +300,7 @@ func (cm *CredentialManager) RemoveOpenAICredentials() error {
|
||||
store.OpenAI = nil
|
||||
|
||||
// If store is empty, remove the file entirely
|
||||
if store.Anthropic == nil && store.OpenAI == nil {
|
||||
if store.Anthropic == nil && store.OpenAI == nil && store.Copilot == nil {
|
||||
if err := os.Remove(cm.credentialsPath); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("failed to remove credentials file: %w", err)
|
||||
}
|
||||
@@ -289,6 +310,104 @@ func (cm *CredentialManager) RemoveOpenAICredentials() error {
|
||||
return cm.SaveCredentials(store)
|
||||
}
|
||||
|
||||
// GetCopilotCredentials retrieves stored GitHub Copilot credentials.
|
||||
func (cm *CredentialManager) GetCopilotCredentials() (*CopilotCredentials, error) {
|
||||
store, err := cm.LoadCredentials()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return store.Copilot, nil
|
||||
}
|
||||
|
||||
// RemoveCopilotCredentials removes stored GitHub Copilot credentials.
|
||||
func (cm *CredentialManager) RemoveCopilotCredentials() error {
|
||||
store, err := cm.LoadCredentials()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
store.Copilot = nil
|
||||
|
||||
if store.Anthropic == nil && store.OpenAI == nil && store.Copilot == nil {
|
||||
if err := os.Remove(cm.credentialsPath); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("failed to remove credentials file: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
return cm.SaveCredentials(store)
|
||||
}
|
||||
|
||||
// HasCopilotCredentials checks if valid GitHub Copilot credentials are stored.
|
||||
func (cm *CredentialManager) HasCopilotCredentials() (bool, error) {
|
||||
creds, err := cm.GetCopilotCredentials()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
if creds == nil {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
return creds.Type == "oauth" && creds.GitHubToken != "", nil
|
||||
}
|
||||
|
||||
// SetCopilotOAuthCredentials stores GitHub Copilot OAuth credentials.
|
||||
func (cm *CredentialManager) SetCopilotOAuthCredentials(creds *CopilotCredentials) error {
|
||||
store, err := cm.LoadCredentials()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
store.Copilot = creds
|
||||
return cm.SaveCredentials(store)
|
||||
}
|
||||
|
||||
// GetValidCopilotAccessToken returns a fresh Copilot API token, renewing it
|
||||
// with the stored GitHub OAuth token when needed.
|
||||
func (cm *CredentialManager) GetValidCopilotAccessToken() (string, error) {
|
||||
return cm.GetValidCopilotAccessTokenContext(context.Background())
|
||||
}
|
||||
|
||||
// GetValidCopilotAccessTokenContext returns a fresh Copilot API token, renewing
|
||||
// it with the stored GitHub OAuth token when needed.
|
||||
func (cm *CredentialManager) GetValidCopilotAccessTokenContext(ctx context.Context) (string, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
|
||||
creds, err := cm.GetCopilotCredentials()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if creds == nil {
|
||||
return "", fmt.Errorf("no Copilot credentials found")
|
||||
}
|
||||
if creds.Type != "oauth" {
|
||||
return "", fmt.Errorf("unknown credential type: %s", creds.Type)
|
||||
}
|
||||
if creds.GitHubToken == "" {
|
||||
return "", fmt.Errorf("GitHub OAuth token missing from Copilot credentials")
|
||||
}
|
||||
|
||||
if creds.CopilotAccessToken == "" || creds.NeedsRefresh() {
|
||||
client := NewCopilotOAuthClient()
|
||||
newCreds, err := client.RefreshCopilotToken(ctx, creds.GitHubToken)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to refresh Copilot token: %w", err)
|
||||
}
|
||||
newCreds.CreatedAt = creds.CreatedAt
|
||||
|
||||
if err := cm.SetCopilotOAuthCredentials(newCreds); err != nil {
|
||||
return "", fmt.Errorf("failed to save refreshed Copilot token: %w", err)
|
||||
}
|
||||
|
||||
return newCreds.CopilotAccessToken, nil
|
||||
}
|
||||
|
||||
return creds.CopilotAccessToken, nil
|
||||
}
|
||||
|
||||
// HasOpenAICredentials checks if valid OpenAI credentials are stored.
|
||||
// Returns true if either a non-empty OAuth access token or API key is present,
|
||||
// false otherwise. Returns an error if credentials cannot be loaded.
|
||||
@@ -394,6 +513,20 @@ func validateAnthropicAPIKey(apiKey string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// CredentialSourceOAuth is the source description returned by
|
||||
// GetAnthropicAPIKey when the key resolves to stored OAuth credentials.
|
||||
// Consumers should compare against this constant (or use IsAnthropicOAuth)
|
||||
// rather than matching the string literal.
|
||||
const CredentialSourceOAuth = "stored OAuth credentials"
|
||||
|
||||
// IsAnthropicOAuth reports whether the active Anthropic credential resolves
|
||||
// to a stored OAuth token (in which case the user is not billed per-token).
|
||||
// flagValue is the --provider-api-key flag value (may be empty).
|
||||
func IsAnthropicOAuth(flagValue string) bool {
|
||||
_, source, err := GetAnthropicAPIKey(flagValue)
|
||||
return err == nil && source == CredentialSourceOAuth
|
||||
}
|
||||
|
||||
// GetAnthropicAPIKey retrieves an Anthropic API key from multiple sources in priority order:
|
||||
// 1. Command-line flag value (highest priority)
|
||||
// 2. Stored credentials (OAuth or API key)
|
||||
@@ -416,7 +549,7 @@ func GetAnthropicAPIKey(flagValue string) (string, string, error) {
|
||||
if err != nil {
|
||||
return "", "", fmt.Errorf("failed to get valid OAuth token: %w", err)
|
||||
}
|
||||
return token, "stored OAuth credentials", nil
|
||||
return token, CredentialSourceOAuth, nil
|
||||
} else if creds.Type == "api_key" && creds.APIKey != "" {
|
||||
return creds.APIKey, "stored API key", nil
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCredentialManager(t *testing.T) {
|
||||
@@ -215,6 +216,7 @@ func TestCredentialStorePersistence(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
|
||||
defer func() { _ = os.RemoveAll(tempDir) }()
|
||||
|
||||
credentialsPath := filepath.Join(tempDir, "credentials.json")
|
||||
@@ -252,3 +254,98 @@ func TestCredentialStorePersistence(t *testing.T) {
|
||||
t.Errorf("Expected file permissions 0600, got %v", info.Mode().Perm())
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopilotCredentials(t *testing.T) {
|
||||
tempDir, err := os.MkdirTemp("", "kit-auth-test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer func() { _ = os.RemoveAll(tempDir) }()
|
||||
|
||||
cm := &CredentialManager{
|
||||
credentialsPath: filepath.Join(tempDir, "credentials.json"),
|
||||
}
|
||||
|
||||
creds := &CopilotCredentials{
|
||||
Type: "oauth",
|
||||
GitHubToken: "github-token",
|
||||
CopilotAccessToken: "copilot-token",
|
||||
ExpiresAt: time.Now().Add(time.Hour).Unix(),
|
||||
CreatedAt: time.Now(),
|
||||
}
|
||||
|
||||
if err := cm.SetCopilotOAuthCredentials(creds); err != nil {
|
||||
t.Fatalf("SetCopilotOAuthCredentials failed: %v", err)
|
||||
}
|
||||
|
||||
hasAuth, err := cm.HasCopilotCredentials()
|
||||
if err != nil {
|
||||
t.Fatalf("HasCopilotCredentials failed: %v", err)
|
||||
}
|
||||
if !hasAuth {
|
||||
t.Fatal("Expected Copilot credentials")
|
||||
}
|
||||
|
||||
token, err := cm.GetValidCopilotAccessToken()
|
||||
if err != nil {
|
||||
t.Fatalf("GetValidCopilotAccessToken failed: %v", err)
|
||||
}
|
||||
if token != creds.CopilotAccessToken {
|
||||
t.Fatalf("Expected Copilot token %q, got %q", creds.CopilotAccessToken, token)
|
||||
}
|
||||
|
||||
if err := cm.RemoveCopilotCredentials(); err != nil {
|
||||
t.Fatalf("RemoveCopilotCredentials failed: %v", err)
|
||||
}
|
||||
hasAuth, err = cm.HasCopilotCredentials()
|
||||
if err != nil {
|
||||
t.Fatalf("HasCopilotCredentials after removal failed: %v", err)
|
||||
}
|
||||
if hasAuth {
|
||||
t.Fatal("Expected no Copilot credentials after removal")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRemoveCredentialsPreservesOtherProviders(t *testing.T) {
|
||||
tempDir, err := os.MkdirTemp("", "kit-auth-test")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
defer func() { _ = os.RemoveAll(tempDir) }()
|
||||
|
||||
cm := &CredentialManager{
|
||||
credentialsPath: filepath.Join(tempDir, "credentials.json"),
|
||||
}
|
||||
|
||||
if err := cm.SetOpenAIOAuthCredentials(&OpenAICredentials{
|
||||
Type: "oauth",
|
||||
AccessToken: "openai-token",
|
||||
RefreshToken: "refresh-token",
|
||||
ExpiresAt: time.Now().Add(time.Hour).Unix(),
|
||||
AccountID: "account",
|
||||
CreatedAt: time.Now(),
|
||||
}); err != nil {
|
||||
t.Fatalf("SetOpenAIOAuthCredentials failed: %v", err)
|
||||
}
|
||||
if err := cm.SetCopilotOAuthCredentials(&CopilotCredentials{
|
||||
Type: "oauth",
|
||||
GitHubToken: "github-token",
|
||||
CopilotAccessToken: "copilot-token",
|
||||
ExpiresAt: time.Now().Add(time.Hour).Unix(),
|
||||
CreatedAt: time.Now(),
|
||||
}); err != nil {
|
||||
t.Fatalf("SetCopilotOAuthCredentials failed: %v", err)
|
||||
}
|
||||
|
||||
if err := cm.RemoveCopilotCredentials(); err != nil {
|
||||
t.Fatalf("RemoveCopilotCredentials failed: %v", err)
|
||||
}
|
||||
|
||||
hasOpenAI, err := cm.HasOpenAICredentials()
|
||||
if err != nil {
|
||||
t.Fatalf("HasOpenAICredentials failed: %v", err)
|
||||
}
|
||||
if !hasOpenAI {
|
||||
t.Fatal("Expected OpenAI credentials to remain after removing Copilot credentials")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
@@ -211,6 +212,262 @@ type OpenAIOAuthClient struct {
|
||||
Scopes string
|
||||
}
|
||||
|
||||
// CopilotOAuthClient handles GitHub device-flow OAuth and exchanges the
|
||||
// GitHub token for a short-lived GitHub Copilot API token.
|
||||
//
|
||||
// The GitHub token comes from GitHub's OAuth device flow. It is then presented
|
||||
// to GitHub's internal Copilot token endpoint, which returns the bearer token
|
||||
// used by api.githubcopilot.com.
|
||||
type CopilotOAuthClient struct {
|
||||
ClientID string
|
||||
DeviceURL string
|
||||
TokenURL string
|
||||
CopilotURL string
|
||||
Scopes string
|
||||
PollTimeout time.Duration
|
||||
ClientTimeout time.Duration
|
||||
}
|
||||
|
||||
// CopilotDeviceCode contains data returned by GitHub's device-code endpoint.
|
||||
type CopilotDeviceCode struct {
|
||||
DeviceCode string `json:"device_code"`
|
||||
UserCode string `json:"user_code"`
|
||||
VerificationURI string `json:"verification_uri"`
|
||||
ExpiresIn int `json:"expires_in"`
|
||||
Interval int `json:"interval"`
|
||||
}
|
||||
|
||||
// NewCopilotOAuthClient creates a GitHub Copilot OAuth client.
|
||||
func NewCopilotOAuthClient() *CopilotOAuthClient {
|
||||
return &CopilotOAuthClient{
|
||||
ClientID: "Iv1.b507a08c87ecfe98",
|
||||
DeviceURL: "https://github.com/login/device/code",
|
||||
TokenURL: "https://github.com/login/oauth/access_token",
|
||||
CopilotURL: "https://api.github.com/copilot_internal/v2/token",
|
||||
Scopes: "read:user",
|
||||
PollTimeout: 15 * time.Minute,
|
||||
ClientTimeout: 30 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// StartDeviceFlow requests a GitHub device code for browser login.
|
||||
//
|
||||
// The returned user code and verification URI are displayed by loginCopilot.
|
||||
// GitHub's response may omit interval, so this method normalizes it to the
|
||||
// documented five-second default.
|
||||
func (c *CopilotOAuthClient) StartDeviceFlow(ctx context.Context) (*CopilotDeviceCode, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
|
||||
data := url.Values{
|
||||
"client_id": {c.ClientID},
|
||||
"scope": {c.Scopes},
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", c.DeviceURL, strings.NewReader(data.Encode()))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create device-code request: %w", err)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
|
||||
resp, err := (&http.Client{Timeout: c.ClientTimeout}).Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to request device code: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("device-code request failed with status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var code CopilotDeviceCode
|
||||
if err := json.NewDecoder(resp.Body).Decode(&code); err != nil {
|
||||
return nil, fmt.Errorf("failed to decode device-code response: %w", err)
|
||||
}
|
||||
if code.DeviceCode == "" || code.UserCode == "" || code.VerificationURI == "" {
|
||||
return nil, fmt.Errorf("device-code response missing required fields")
|
||||
}
|
||||
if code.Interval <= 0 {
|
||||
code.Interval = 5
|
||||
}
|
||||
return &code, nil
|
||||
}
|
||||
|
||||
// PollDeviceToken waits until the user authorizes the device code and returns
|
||||
// the resulting GitHub OAuth token.
|
||||
//
|
||||
// It follows GitHub's device-flow polling contract: authorization_pending keeps
|
||||
// polling, slow_down increases the interval, and polling stops at the earlier of
|
||||
// the client timeout or the device-code expiry.
|
||||
func (c *CopilotOAuthClient) PollDeviceToken(ctx context.Context, deviceCode *CopilotDeviceCode) (string, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
|
||||
if deviceCode == nil || deviceCode.DeviceCode == "" {
|
||||
return "", fmt.Errorf("device code missing")
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(c.PollTimeout)
|
||||
if deviceCode.ExpiresIn > 0 {
|
||||
expiresAt := time.Now().Add(time.Duration(deviceCode.ExpiresIn) * time.Second)
|
||||
if expiresAt.Before(deadline) {
|
||||
deadline = expiresAt
|
||||
}
|
||||
}
|
||||
|
||||
interval := time.Duration(deviceCode.Interval) * time.Second
|
||||
if interval <= 0 {
|
||||
interval = 5 * time.Second
|
||||
}
|
||||
|
||||
for time.Now().Before(deadline) {
|
||||
wait := interval
|
||||
if remaining := time.Until(deadline); remaining < wait {
|
||||
wait = remaining
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return "", ctx.Err()
|
||||
case <-time.After(wait):
|
||||
}
|
||||
|
||||
data := url.Values{
|
||||
"client_id": {c.ClientID},
|
||||
"device_code": {deviceCode.DeviceCode},
|
||||
"grant_type": {"urn:ietf:params:oauth:grant-type:device_code"},
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", c.TokenURL, strings.NewReader(data.Encode()))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to create device-token request: %w", err)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
|
||||
|
||||
resp, err := (&http.Client{Timeout: c.ClientTimeout}).Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to poll device token: %w", err)
|
||||
}
|
||||
|
||||
var tokenResp struct {
|
||||
AccessToken string `json:"access_token"`
|
||||
Error string `json:"error"`
|
||||
Description string `json:"error_description"`
|
||||
}
|
||||
decodeErr := json.NewDecoder(resp.Body).Decode(&tokenResp)
|
||||
_ = resp.Body.Close()
|
||||
if decodeErr != nil {
|
||||
return "", fmt.Errorf("failed to decode device-token response: %w", decodeErr)
|
||||
}
|
||||
|
||||
if tokenResp.AccessToken != "" {
|
||||
return tokenResp.AccessToken, nil
|
||||
}
|
||||
|
||||
switch tokenResp.Error {
|
||||
case "authorization_pending":
|
||||
continue
|
||||
case "slow_down":
|
||||
interval += 5 * time.Second
|
||||
continue
|
||||
case "expired_token":
|
||||
return "", fmt.Errorf("device code expired; restart login")
|
||||
case "access_denied":
|
||||
return "", fmt.Errorf("github login denied")
|
||||
case "":
|
||||
return "", fmt.Errorf("device-token request failed with status %d", resp.StatusCode)
|
||||
default:
|
||||
if tokenResp.Description != "" {
|
||||
return "", fmt.Errorf("device-token request failed: %s: %s", tokenResp.Error, tokenResp.Description)
|
||||
}
|
||||
return "", fmt.Errorf("device-token request failed: %s", tokenResp.Error)
|
||||
}
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("timed out waiting for github device authorization")
|
||||
}
|
||||
|
||||
// ExchangeGitHubToken converts a GitHub OAuth token into a Copilot API token.
|
||||
// It is a semantic wrapper over RefreshCopilotToken used by the login flow.
|
||||
func (c *CopilotOAuthClient) ExchangeGitHubToken(ctx context.Context, githubToken string) (*CopilotCredentials, error) {
|
||||
return c.RefreshCopilotToken(ctx, githubToken)
|
||||
}
|
||||
|
||||
// RefreshCopilotToken obtains a fresh short-lived Copilot token from GitHub.
|
||||
//
|
||||
// GitHub may return expires_at as either a Unix timestamp or RFC3339 string.
|
||||
// parseCopilotExpiry handles both forms and falls back to a conservative
|
||||
// 20-minute lifetime when the field is absent or unrecognized.
|
||||
func (c *CopilotOAuthClient) RefreshCopilotToken(ctx context.Context, githubToken string) (*CopilotCredentials, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", c.CopilotURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create copilot token request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "token "+githubToken)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("User-Agent", "kit")
|
||||
req.Header.Set("X-GitHub-Api-Version", "2022-11-28")
|
||||
|
||||
resp, err := (&http.Client{Timeout: c.ClientTimeout}).Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to request copilot token: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return nil, fmt.Errorf("copilot token request failed with status %d: %s", resp.StatusCode, string(body))
|
||||
}
|
||||
|
||||
var tokenResp struct {
|
||||
Token string `json:"token"`
|
||||
ExpiresAt any `json:"expires_at"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to decode copilot token response: %w", err)
|
||||
}
|
||||
if tokenResp.Token == "" {
|
||||
return nil, fmt.Errorf("copilot token response missing token")
|
||||
}
|
||||
|
||||
expiresAt := parseCopilotExpiry(tokenResp.ExpiresAt)
|
||||
if expiresAt == 0 {
|
||||
expiresAt = time.Now().Add(20 * time.Minute).Unix()
|
||||
}
|
||||
|
||||
return &CopilotCredentials{
|
||||
Type: "oauth",
|
||||
GitHubToken: githubToken,
|
||||
CopilotAccessToken: tokenResp.Token,
|
||||
ExpiresAt: expiresAt,
|
||||
CreatedAt: time.Now(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// parseCopilotExpiry normalizes GitHub's expires_at variants to a Unix second.
|
||||
func parseCopilotExpiry(value any) int64 {
|
||||
switch v := value.(type) {
|
||||
case float64:
|
||||
return int64(v)
|
||||
case string:
|
||||
if parsed, err := strconv.ParseInt(v, 10, 64); err == nil {
|
||||
return parsed
|
||||
}
|
||||
if parsed, err := time.Parse(time.RFC3339, v); err == nil {
|
||||
return parsed.Unix()
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// NewOpenAIOAuthClient creates a new OAuth client configured for OpenAI Codex OAuth.
|
||||
// This uses the public client ID for CLI applications with PKCE for security.
|
||||
func NewOpenAIOAuthClient() *OpenAIOAuthClient {
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCopilotStartDeviceFlow(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
t.Fatalf("expected POST, got %s", r.Method)
|
||||
}
|
||||
if err := r.ParseForm(); err != nil {
|
||||
t.Fatalf("ParseForm failed: %v", err)
|
||||
}
|
||||
if r.Form.Get("client_id") != "client-id" {
|
||||
t.Fatalf("expected client id, got %q", r.Form.Get("client_id"))
|
||||
}
|
||||
if r.Form.Get("scope") != "read:user" {
|
||||
t.Fatalf("expected scope, got %q", r.Form.Get("scope"))
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"device_code": "device-code",
|
||||
"user_code": "USER-CODE",
|
||||
"verification_uri": "https://github.com/login/device",
|
||||
"expires_in": 600,
|
||||
"interval": 1,
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewCopilotOAuthClient()
|
||||
client.ClientID = "client-id"
|
||||
client.DeviceURL = server.URL
|
||||
|
||||
code, err := client.StartDeviceFlow(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("StartDeviceFlow failed: %v", err)
|
||||
}
|
||||
if code.DeviceCode != "device-code" || code.UserCode != "USER-CODE" || code.Interval != 1 {
|
||||
t.Fatalf("unexpected device code: %#v", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopilotPollDeviceToken(t *testing.T) {
|
||||
polls := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
polls++
|
||||
if r.Method != http.MethodPost {
|
||||
t.Fatalf("expected POST, got %s", r.Method)
|
||||
}
|
||||
if err := r.ParseForm(); err != nil {
|
||||
t.Fatalf("ParseForm failed: %v", err)
|
||||
}
|
||||
if r.Form.Get("grant_type") != "urn:ietf:params:oauth:grant-type:device_code" {
|
||||
t.Fatalf("unexpected grant type: %q", r.Form.Get("grant_type"))
|
||||
}
|
||||
if polls == 1 {
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"error": "authorization_pending"})
|
||||
return
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"access_token": "github-token"})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewCopilotOAuthClient()
|
||||
client.ClientID = "client-id"
|
||||
client.TokenURL = server.URL
|
||||
client.PollTimeout = 5 * time.Second
|
||||
client.ClientTimeout = time.Second
|
||||
|
||||
token, err := client.PollDeviceToken(context.Background(), &CopilotDeviceCode{
|
||||
DeviceCode: "device-code",
|
||||
ExpiresIn: 10,
|
||||
Interval: 1,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("PollDeviceToken failed: %v", err)
|
||||
}
|
||||
if token != "github-token" {
|
||||
t.Fatalf("expected github-token, got %q", token)
|
||||
}
|
||||
if polls != 2 {
|
||||
t.Fatalf("expected 2 polls, got %d", polls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopilotRefreshToken(t *testing.T) {
|
||||
expiresAt := time.Now().Add(time.Hour).Unix()
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
t.Fatalf("expected GET, got %s", r.Method)
|
||||
}
|
||||
if r.Header.Get("Authorization") != "token github-token" {
|
||||
t.Fatalf("unexpected authorization header: %q", r.Header.Get("Authorization"))
|
||||
}
|
||||
if r.Header.Get("User-Agent") != "kit" {
|
||||
t.Fatalf("unexpected user agent: %q", r.Header.Get("User-Agent"))
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"token": "copilot-token",
|
||||
"expires_at": expiresAt,
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewCopilotOAuthClient()
|
||||
client.CopilotURL = server.URL
|
||||
|
||||
creds, err := client.RefreshCopilotToken(context.Background(), "github-token")
|
||||
if err != nil {
|
||||
t.Fatalf("RefreshCopilotToken failed: %v", err)
|
||||
}
|
||||
if creds.GitHubToken != "github-token" || creds.CopilotAccessToken != "copilot-token" {
|
||||
t.Fatalf("unexpected credentials: %#v", creds)
|
||||
}
|
||||
if creds.ExpiresAt != expiresAt {
|
||||
t.Fatalf("expected expires_at %d, got %d", expiresAt, creds.ExpiresAt)
|
||||
}
|
||||
}
|
||||
@@ -56,9 +56,3 @@ func (e *EnvSubstituter) SubstituteEnvVars(content string) (string, error) {
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// HasEnvVars checks if content contains environment variable patterns (${env://...}).
|
||||
// This is useful for determining if substitution is needed before processing.
|
||||
func HasEnvVars(content string) bool {
|
||||
return envVarPattern.MatchString(content)
|
||||
}
|
||||
|
||||
@@ -187,41 +187,3 @@ func TestEnvSubstituter_SubstituteEnvVars(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHasEnvVars(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
content string
|
||||
expected bool
|
||||
}{
|
||||
{
|
||||
name: "has env vars",
|
||||
content: `{"token": "${env://GITHUB_TOKEN}"}`,
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "has env vars with default",
|
||||
content: `{"debug": "${env://DEBUG:-false}"}`,
|
||||
expected: true,
|
||||
},
|
||||
{
|
||||
name: "no env vars",
|
||||
content: `{"name": "${username}", "normal": "value"}`,
|
||||
expected: false,
|
||||
},
|
||||
{
|
||||
name: "empty content",
|
||||
content: "",
|
||||
expected: false,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := HasEnvVars(tt.content)
|
||||
if result != tt.expected {
|
||||
t.Errorf("Expected %v, got %v", tt.expected, result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -59,12 +59,6 @@ func passwordPromptFromContext(ctx context.Context) PasswordPromptCallback {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ContextWithSudoPassword returns a new context with the sudo password set.
|
||||
// When present, the bash tool will use sudo -S to pipe this password to sudo commands.
|
||||
func ContextWithSudoPassword(ctx context.Context, password string) context.Context {
|
||||
return context.WithValue(ctx, sudoPasswordKey, password)
|
||||
}
|
||||
|
||||
// sudoPasswordFromContext retrieves the sudo password from context.
|
||||
func sudoPasswordFromContext(ctx context.Context) string {
|
||||
if pw, ok := ctx.Value(sudoPasswordKey).(string); ok {
|
||||
|
||||
@@ -183,7 +183,7 @@ func TestRewriteSudoForStdin(t *testing.T) {
|
||||
|
||||
func TestSudoPasswordFromContext(t *testing.T) {
|
||||
// Test with password in context
|
||||
ctx := ContextWithSudoPassword(context.Background(), "secret123")
|
||||
ctx := context.WithValue(context.Background(), sudoPasswordKey, "secret123")
|
||||
pw := sudoPasswordFromContext(ctx)
|
||||
if pw != "secret123" {
|
||||
t.Errorf("expected password 'secret123', got %q", pw)
|
||||
|
||||
@@ -0,0 +1,234 @@
|
||||
package extbridge
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/mark3labs/kit/internal/extensions"
|
||||
kit "github.com/mark3labs/kit/pkg/kit"
|
||||
)
|
||||
|
||||
// BaseContext returns an extensions.Context populated with the headless,
|
||||
// TUI-independent delegation fields: data access, state, options,
|
||||
// model/tool management, completions, subagents, tree navigation, skills,
|
||||
// template parsing, and model resolution.
|
||||
//
|
||||
// Callers overlay their UI-specific fields (print routes, widgets, prompts,
|
||||
// editor, TUI-aware SetModel/ReloadExtensions, etc.) on the returned value:
|
||||
// cmd/extension_context.go for the interactive TUI and
|
||||
// internal/acpserver/session.go for headless ACP mode. Keeping the shared
|
||||
// half here means a new data-access Context field only has to be wired once.
|
||||
//
|
||||
// ctx is used for subagent spawns; pass a long-lived context (not a
|
||||
// per-request one) so later spawns aren't cancelled prematurely.
|
||||
func BaseContext(ctx context.Context, kitInstance *kit.Kit) extensions.Context {
|
||||
return extensions.Context{
|
||||
// -------------------------------------------------------------------
|
||||
// Data access
|
||||
// -------------------------------------------------------------------
|
||||
GetContextStats: func() extensions.ContextStats {
|
||||
s := kitInstance.GetContextStats()
|
||||
return extensions.ContextStats{
|
||||
EstimatedTokens: s.EstimatedTokens,
|
||||
ContextLimit: s.ContextLimit,
|
||||
UsagePercent: s.UsagePercent,
|
||||
MessageCount: s.MessageCount,
|
||||
}
|
||||
},
|
||||
GetMessages: func() []extensions.SessionMessage {
|
||||
return kitInstance.Extensions().GetSessionMessages()
|
||||
},
|
||||
GetSessionPath: func() string {
|
||||
return kitInstance.GetSessionPath()
|
||||
},
|
||||
AppendEntry: func(entryType string, data string) (string, error) {
|
||||
return kitInstance.Extensions().AppendEntry(entryType, data)
|
||||
},
|
||||
GetEntries: func(entryType string) []extensions.ExtensionEntry {
|
||||
return kitInstance.Extensions().GetEntries(entryType)
|
||||
},
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// Extension state
|
||||
// -------------------------------------------------------------------
|
||||
SetState: func(key string, value string) {
|
||||
kitInstance.Extensions().SetState(key, value)
|
||||
},
|
||||
GetState: func(key string) (string, bool) {
|
||||
return kitInstance.Extensions().GetState(key)
|
||||
},
|
||||
DeleteState: func(key string) {
|
||||
kitInstance.Extensions().DeleteState(key)
|
||||
},
|
||||
ListState: func() []string {
|
||||
return kitInstance.Extensions().ListState()
|
||||
},
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// Options, model, and tool management
|
||||
// -------------------------------------------------------------------
|
||||
GetOption: func(name string) string {
|
||||
return kitInstance.Extensions().GetOption(name)
|
||||
},
|
||||
SetOption: func(name string, value string) {
|
||||
kitInstance.Extensions().SetOption(name, value)
|
||||
},
|
||||
// Headless model switch. The interactive TUI overrides this with a
|
||||
// version that also notifies the TUI and refreshes the usage tracker.
|
||||
SetModel: func(modelString string) error {
|
||||
previousModel := kitInstance.Extensions().GetContext().Model
|
||||
if err := kitInstance.SetModel(context.Background(), modelString); err != nil {
|
||||
return err
|
||||
}
|
||||
kitInstance.Extensions().UpdateContextModel(modelString)
|
||||
kitInstance.Extensions().EmitModelChange(modelString, previousModel, "extension")
|
||||
return nil
|
||||
},
|
||||
GetAvailableModels: func() []extensions.ModelInfoEntry {
|
||||
return kitInstance.GetAvailableModels()
|
||||
},
|
||||
EmitCustomEvent: func(name string, data string) {
|
||||
kitInstance.Extensions().EmitCustomEvent(name, data)
|
||||
},
|
||||
GetAllTools: func() []extensions.ToolInfo {
|
||||
return kitInstance.Extensions().GetToolInfos()
|
||||
},
|
||||
SetActiveTools: func(names []string) {
|
||||
kitInstance.Extensions().SetActiveTools(names)
|
||||
},
|
||||
// Headless reload. The interactive TUI overrides this to also
|
||||
// refresh widgets/status/commands.
|
||||
ReloadExtensions: func() error {
|
||||
return kitInstance.Extensions().Reload()
|
||||
},
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// LLM completions and subagents
|
||||
// -------------------------------------------------------------------
|
||||
Complete: func(req extensions.CompleteRequest) (extensions.CompleteResponse, error) {
|
||||
return kitInstance.ExecuteCompletion(context.Background(), req)
|
||||
},
|
||||
SpawnSubagent: func(config extensions.SubagentConfig) (*extensions.SubagentHandle, *extensions.SubagentResult, error) {
|
||||
return SpawnSubagent(ctx, kitInstance, config)
|
||||
},
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// Tree Navigation API
|
||||
// -------------------------------------------------------------------
|
||||
GetTreeNode: func(entryID string) *extensions.TreeNode {
|
||||
node := kitInstance.GetTreeNode(entryID)
|
||||
if node == nil {
|
||||
return nil
|
||||
}
|
||||
return &extensions.TreeNode{
|
||||
ID: node.ID,
|
||||
ParentID: node.ParentID,
|
||||
Type: node.Type,
|
||||
Role: node.Role,
|
||||
Content: node.Content,
|
||||
Model: node.Model,
|
||||
Provider: node.Provider,
|
||||
Timestamp: node.Timestamp,
|
||||
Children: node.Children,
|
||||
}
|
||||
},
|
||||
GetCurrentBranch: func() []extensions.TreeNode {
|
||||
nodes := kitInstance.GetCurrentBranch()
|
||||
result := make([]extensions.TreeNode, len(nodes))
|
||||
for i, n := range nodes {
|
||||
result[i] = extensions.TreeNode{
|
||||
ID: n.ID,
|
||||
ParentID: n.ParentID,
|
||||
Type: n.Type,
|
||||
Role: n.Role,
|
||||
Content: n.Content,
|
||||
Model: n.Model,
|
||||
Provider: n.Provider,
|
||||
Timestamp: n.Timestamp,
|
||||
Children: n.Children,
|
||||
}
|
||||
}
|
||||
return result
|
||||
},
|
||||
GetChildren: func(parentID string) []string {
|
||||
return kitInstance.GetChildren(parentID)
|
||||
},
|
||||
NavigateTo: func(entryID string) extensions.TreeNavigationResult {
|
||||
err := kitInstance.NavigateTo(entryID)
|
||||
if err != nil {
|
||||
return extensions.TreeNavigationResult{Success: false, Error: err.Error()}
|
||||
}
|
||||
return extensions.TreeNavigationResult{Success: true}
|
||||
},
|
||||
SummarizeBranch: func(fromID, toID string) string {
|
||||
summary, _ := kitInstance.SummarizeBranch(fromID, toID)
|
||||
return summary
|
||||
},
|
||||
CollapseBranch: func(fromID, toID, summary string) extensions.TreeNavigationResult {
|
||||
err := kitInstance.CollapseBranch(fromID, toID, summary)
|
||||
if err != nil {
|
||||
return extensions.TreeNavigationResult{Success: false, Error: err.Error()}
|
||||
}
|
||||
return extensions.TreeNavigationResult{Success: true}
|
||||
},
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// Skill Loading API (context-injection variants are TUI-specific and
|
||||
// wired by the interactive overlay)
|
||||
// -------------------------------------------------------------------
|
||||
LoadSkill: func(path string) (*extensions.Skill, string) {
|
||||
s, err := kitInstance.LoadSkillForExtension(path)
|
||||
return s, err
|
||||
},
|
||||
LoadSkillsFromDir: func(dir string) extensions.SkillLoadResult {
|
||||
return kitInstance.LoadSkillsFromDirForExtension(dir)
|
||||
},
|
||||
DiscoverSkills: func() extensions.SkillLoadResult {
|
||||
skills := kitInstance.DiscoverSkillsForExtension()
|
||||
return extensions.SkillLoadResult{Skills: skills}
|
||||
},
|
||||
GetAvailableSkills: func() []extensions.Skill {
|
||||
return kitInstance.DiscoverSkillsForExtension()
|
||||
},
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// Template Parsing API
|
||||
// -------------------------------------------------------------------
|
||||
ParseTemplate: func(name, content string) extensions.PromptTemplate {
|
||||
return kit.ParseTemplate(name, content)
|
||||
},
|
||||
RenderTemplate: func(tpl extensions.PromptTemplate, vars map[string]string) string {
|
||||
return kit.RenderTemplate(tpl, vars)
|
||||
},
|
||||
ParseArguments: func(input string, pattern extensions.ArgumentPattern) extensions.ParseResult {
|
||||
return kit.ParseArguments(input, pattern)
|
||||
},
|
||||
SimpleParseArguments: func(input string, count int) []string {
|
||||
return kit.SimpleParseArguments(input, count)
|
||||
},
|
||||
EvaluateModelConditional: func(condition string) bool {
|
||||
return kit.EvaluateModelConditional(kitInstance.Extensions().GetContext().Model, condition)
|
||||
},
|
||||
RenderWithModelConditionals: func(content string) string {
|
||||
return kit.RenderWithModelConditionals(content, kitInstance.Extensions().GetContext().Model)
|
||||
},
|
||||
|
||||
// -------------------------------------------------------------------
|
||||
// Model Resolution API
|
||||
// -------------------------------------------------------------------
|
||||
ResolveModelChain: func(preferences []string) extensions.ModelResolutionResult {
|
||||
return kit.ResolveModelChain(preferences)
|
||||
},
|
||||
GetModelCapabilities: func(model string) (extensions.ModelCapabilities, string) {
|
||||
return kit.GetModelCapabilities(model)
|
||||
},
|
||||
CheckModelAvailable: func(model string) bool {
|
||||
return kit.CheckModelAvailable(model)
|
||||
},
|
||||
GetCurrentProvider: func() string {
|
||||
return kit.GetCurrentProvider(kitInstance.Extensions().GetContext().Model)
|
||||
},
|
||||
GetCurrentModelID: func() string {
|
||||
return kit.GetCurrentModelID(kitInstance.Extensions().GetContext().Model)
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -66,6 +66,7 @@ func SpawnSubagent(ctx context.Context, k *kit.Kit, cfg extensions.SubagentConfi
|
||||
SystemPrompt: cfg.SystemPrompt,
|
||||
Timeout: cfg.Timeout,
|
||||
NoSession: cfg.NoSession,
|
||||
Tools: k.GetToolsForSubagent(),
|
||||
}
|
||||
if cfg.OnEvent != nil {
|
||||
sdkCfg.OnEvent = func(e kit.Event) {
|
||||
|
||||
+135
-1
@@ -341,6 +341,13 @@ type Context struct {
|
||||
// The data survives across session restarts and can be retrieved via
|
||||
// GetEntries. Use entryType to namespace your data (e.g. "myext:state").
|
||||
//
|
||||
// AppendEntry is append-only and lives in the conversation tree, which
|
||||
// makes it the right tool for audit logs and event histories. For
|
||||
// last-write-wins snapshot state — "what's the current value of X?" —
|
||||
// prefer SetState / GetState instead. Those primitives store data in a
|
||||
// sidecar file outside the conversation tree, are O(1) to read/write,
|
||||
// and do not bloat branch reads or duplicate on fork.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// data, _ := json.Marshal(myState)
|
||||
@@ -360,6 +367,45 @@ type Context struct {
|
||||
// }
|
||||
GetEntries func(entryType string) []ExtensionEntry
|
||||
|
||||
// SetState stores a key-value pair in session-scoped, last-write-wins
|
||||
// extension state. Unlike AppendEntry the value is kept in a sidecar
|
||||
// file outside the conversation tree, so:
|
||||
// - reads are O(1) (no branch walk)
|
||||
// - writes don't bloat the session JSONL
|
||||
// - state is not duplicated on fork (branches share the sidecar)
|
||||
// - state is invisible to the LLM
|
||||
//
|
||||
// Use SetState for snapshot state ("current value of X"); use
|
||||
// AppendEntry for audit logs and event histories. Namespace keys with
|
||||
// your extension name to avoid collisions (e.g. "myext:budget-cap").
|
||||
//
|
||||
// State persists for the lifetime of the session. For ephemeral or
|
||||
// in-memory sessions the state lives only in memory.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// ctx.SetState("myext:budget-cap", "10.00")
|
||||
SetState func(key string, value string)
|
||||
|
||||
// GetState returns the value previously stored via SetState. The bool
|
||||
// is false when the key was never written. Returns ("", false) when
|
||||
// state is unavailable.
|
||||
//
|
||||
// Example:
|
||||
//
|
||||
// if cap, ok := ctx.GetState("myext:budget-cap"); ok {
|
||||
// fmt.Println("current cap:", cap)
|
||||
// }
|
||||
GetState func(key string) (string, bool)
|
||||
|
||||
// DeleteState removes a key from session-scoped extension state.
|
||||
// No-op when the key is missing.
|
||||
DeleteState func(key string)
|
||||
|
||||
// ListState returns all keys currently stored in session-scoped
|
||||
// extension state, in unspecified order.
|
||||
ListState func() []string
|
||||
|
||||
// SetEditorText sets the text content of the input editor. This can
|
||||
// be used to pre-fill the editor with suggested text (e.g. extracted
|
||||
// questions, handoff prompts). The cursor is moved to the end.
|
||||
@@ -1102,6 +1148,7 @@ type API struct {
|
||||
onError func(func(ErrorEvent, Context))
|
||||
onRetry func(func(RetryEvent, Context))
|
||||
onPrepareStep func(func(PrepareStepEvent, Context) *PrepareStepResult)
|
||||
onLLMUsage func(func(LLMUsageEvent, Context))
|
||||
}
|
||||
|
||||
// OnToolCall registers a handler that fires before a tool executes.
|
||||
@@ -1359,6 +1406,19 @@ func (a *API) OnPrepareStep(handler func(PrepareStepEvent, Context) *PrepareStep
|
||||
a.onPrepareStep(handler)
|
||||
}
|
||||
|
||||
// OnLLMUsage registers a handler that fires after each LLM provider call
|
||||
// with the token and cost deltas for that single call. Use this for
|
||||
// per-call usage attribution, real-time budget enforcement, and cost
|
||||
// dashboards that need to react between calls within a single agent turn.
|
||||
//
|
||||
// Handlers receive an LLMUsageEvent describing the call's input/output
|
||||
// tokens, cache tokens, computed cost, model, and provider. A single agent
|
||||
// turn typically fires multiple LLMUsageEvents (one per tool-loop
|
||||
// iteration).
|
||||
func (a *API) OnLLMUsage(handler func(LLMUsageEvent, Context)) {
|
||||
a.onLLMUsage(handler)
|
||||
}
|
||||
|
||||
// RegisterToolRenderer registers a custom renderer for a specific tool's
|
||||
// display in the TUI. The renderer controls the header (parameter summary)
|
||||
// and/or body (result display) of the tool's output block. If multiple
|
||||
@@ -2091,10 +2151,47 @@ type AgentStartEvent struct {
|
||||
|
||||
func (e AgentStartEvent) Type() EventType { return AgentStart }
|
||||
|
||||
// AgentEndEvent fires when the agent finishes responding.
|
||||
// AgentEndEvent fires when the agent finishes responding. In addition to the
|
||||
// final response and stop reason, the event carries per-turn aggregates so
|
||||
// observer-style extensions don't have to maintain parallel bookkeeping in
|
||||
// OnToolResult / OnStepFinish handlers.
|
||||
type AgentEndEvent struct {
|
||||
Response string
|
||||
StopReason string // "completed", "cancelled", "error"
|
||||
|
||||
// ToolCallCount is the total number of tool invocations observed during
|
||||
// this turn (sum across all steps).
|
||||
ToolCallCount int
|
||||
|
||||
// ToolNames lists the tool names invoked during this turn, in call order.
|
||||
// Duplicates are preserved (e.g. two bash calls produce ["bash", "bash"]).
|
||||
ToolNames []string
|
||||
|
||||
// LLMCallCount is the number of LLM round-trips (tool-loop iterations)
|
||||
// performed during this turn. Always >= 1 for a successful turn.
|
||||
LLMCallCount int
|
||||
|
||||
// InputTokensDelta is the sum of input tokens consumed during this turn
|
||||
// across every LLM call (including cache-hit input tokens).
|
||||
InputTokensDelta int
|
||||
|
||||
// OutputTokensDelta is the sum of output tokens generated during this turn.
|
||||
OutputTokensDelta int
|
||||
|
||||
// CacheReadTokensDelta is the sum of cache-read tokens during this turn.
|
||||
CacheReadTokensDelta int
|
||||
|
||||
// CacheWriteTokensDelta is the sum of cache-write tokens during this turn.
|
||||
CacheWriteTokensDelta int
|
||||
|
||||
// CostDelta is the total cost in USD attributable to this turn. Computed
|
||||
// from per-step usage and current model pricing. Zero when pricing is
|
||||
// unknown or OAuth credentials are in use.
|
||||
CostDelta float64
|
||||
|
||||
// DurationMs is the elapsed wall-clock time from AgentStart to AgentEnd,
|
||||
// in milliseconds.
|
||||
DurationMs int64
|
||||
}
|
||||
|
||||
func (e AgentEndEvent) Type() EventType { return AgentEnd }
|
||||
@@ -2403,6 +2500,43 @@ type PrepareStepResult struct {
|
||||
|
||||
func (PrepareStepResult) isResult() {}
|
||||
|
||||
// LLMUsageEvent fires after each LLM provider call with the per-call token
|
||||
// and cost deltas. Use this for accurate budget tracking, cost dashboards,
|
||||
// and any logic that needs to react between LLM calls within a single agent
|
||||
// turn (rather than only at turn boundaries).
|
||||
//
|
||||
// A single agent turn typically produces multiple LLMUsageEvents (one per
|
||||
// tool-loop iteration). The Model and Provider fields reflect the model used
|
||||
// for that specific call, which may differ from earlier calls if the
|
||||
// extension switched models mid-turn via ctx.SetModel().
|
||||
type LLMUsageEvent struct {
|
||||
// InputTokens is the number of input tokens for this call.
|
||||
InputTokens int
|
||||
// OutputTokens is the number of output tokens generated by this call.
|
||||
OutputTokens int
|
||||
// CacheReadTokens is the number of cache-hit input tokens (provider-specific).
|
||||
CacheReadTokens int
|
||||
// CacheWriteTokens is the number of cache-write tokens.
|
||||
CacheWriteTokens int
|
||||
// Cost is the USD cost of this call computed from the model's per-token
|
||||
// pricing. Zero when pricing is unknown or OAuth credentials are in use.
|
||||
Cost float64
|
||||
// Model is the model identifier used for this call (e.g. "claude-sonnet-4-5-20250929").
|
||||
Model string
|
||||
// Provider is the provider identifier (e.g. "anthropic", "openai").
|
||||
Provider string
|
||||
// RequestID is an optional correlation id for the underlying provider
|
||||
// call. May be empty when the provider does not surface one.
|
||||
RequestID string
|
||||
// StepNumber is the zero-based step index within the current agent turn.
|
||||
StepNumber int
|
||||
// FinishReason mirrors the provider's finish reason for this call
|
||||
// (e.g. "stop", "tool_calls", "length"). May be empty.
|
||||
FinishReason string
|
||||
}
|
||||
|
||||
func (e LLMUsageEvent) Type() EventType { return LLMUsage }
|
||||
|
||||
// ThemeColor is an adaptive color pair with light and dark hex values.
|
||||
// Either field may be empty to inherit from the default theme.
|
||||
type ThemeColor struct {
|
||||
|
||||
@@ -125,6 +125,11 @@ const (
|
||||
// after steering messages are injected and before messages are sent
|
||||
// to the LLM. Handlers can replace the context window for this step.
|
||||
PrepareStep EventType = "prepare_step"
|
||||
|
||||
// LLMUsage fires after each LLM provider call with the token and cost
|
||||
// deltas for that single call. Extensions use it to attribute usage to
|
||||
// specific calls/models and to drive budget enforcement between calls.
|
||||
LLMUsage EventType = "llm_usage"
|
||||
)
|
||||
|
||||
// AllEventTypes returns every supported event type.
|
||||
@@ -139,7 +144,7 @@ func AllEventTypes() []EventType {
|
||||
BeforeFork, BeforeSessionSwitch, BeforeCompact,
|
||||
SubagentStart, SubagentChunk, SubagentEnd,
|
||||
StepStart, StepFinish, ReasoningStart, Warnings, Source, Error, Retry,
|
||||
PrepareStep,
|
||||
PrepareStep, LLMUsage,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -4,8 +4,8 @@ import "testing"
|
||||
|
||||
func TestAllEventTypes_Count(t *testing.T) {
|
||||
all := AllEventTypes()
|
||||
if len(all) != 32 {
|
||||
t.Fatalf("expected 32 event types, got %d", len(all))
|
||||
if len(all) != 33 {
|
||||
t.Fatalf("expected 33 event types, got %d", len(all))
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,119 @@
|
||||
package extensions
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestRunner_EmitLLMUsage(t *testing.T) {
|
||||
var got LLMUsageEvent
|
||||
var called bool
|
||||
ext := makeHandlerExt("llmusage.go", map[EventType][]HandlerFunc{
|
||||
LLMUsage: {
|
||||
func(e Event, c Context) Result {
|
||||
got = e.(LLMUsageEvent)
|
||||
called = true
|
||||
return nil
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
r := makeRunner(ext)
|
||||
_, err := r.Emit(LLMUsageEvent{
|
||||
InputTokens: 100,
|
||||
OutputTokens: 50,
|
||||
Cost: 0.0012,
|
||||
Model: "claude-sonnet-4-5-20250929",
|
||||
Provider: "anthropic",
|
||||
StepNumber: 2,
|
||||
FinishReason: "tool_calls",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("emit: %v", err)
|
||||
}
|
||||
if !called {
|
||||
t.Fatal("expected LLMUsage handler to be called")
|
||||
}
|
||||
if got.InputTokens != 100 || got.OutputTokens != 50 {
|
||||
t.Errorf("token fields not propagated: %+v", got)
|
||||
}
|
||||
if got.Cost != 0.0012 {
|
||||
t.Errorf("cost not propagated, got %v", got.Cost)
|
||||
}
|
||||
if got.Model != "claude-sonnet-4-5-20250929" || got.Provider != "anthropic" {
|
||||
t.Errorf("model/provider not propagated: %+v", got)
|
||||
}
|
||||
if got.StepNumber != 2 || got.FinishReason != "tool_calls" {
|
||||
t.Errorf("step/finish reason not propagated: %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunner_LLMUsageRegisteredViaTestAPI(t *testing.T) {
|
||||
// Verify NewTestAPI wires up onLLMUsage so the extension can call
|
||||
// api.OnLLMUsage during Init.
|
||||
ext := &LoadedExtension{Handlers: make(map[EventType][]HandlerFunc)}
|
||||
api := NewTestAPI(ext)
|
||||
|
||||
var calls int
|
||||
api.OnLLMUsage(func(e LLMUsageEvent, c Context) {
|
||||
calls++
|
||||
})
|
||||
|
||||
if len(ext.Handlers[LLMUsage]) != 1 {
|
||||
t.Fatalf("expected 1 LLMUsage handler registered, got %d", len(ext.Handlers[LLMUsage]))
|
||||
}
|
||||
|
||||
r := makeRunner(*ext)
|
||||
_, _ = r.Emit(LLMUsageEvent{InputTokens: 1})
|
||||
if calls != 1 {
|
||||
t.Errorf("expected handler called once, got %d", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAgentEndEvent_EnrichedFields(t *testing.T) {
|
||||
// Verify the enriched event carries through Emit without mangling.
|
||||
var got AgentEndEvent
|
||||
ext := makeHandlerExt("end.go", map[EventType][]HandlerFunc{
|
||||
AgentEnd: {
|
||||
func(e Event, c Context) Result {
|
||||
got = e.(AgentEndEvent)
|
||||
return nil
|
||||
},
|
||||
},
|
||||
})
|
||||
r := makeRunner(ext)
|
||||
_, err := r.Emit(AgentEndEvent{
|
||||
Response: "done",
|
||||
StopReason: "completed",
|
||||
ToolCallCount: 3,
|
||||
ToolNames: []string{"bash", "read", "bash"},
|
||||
LLMCallCount: 4,
|
||||
InputTokensDelta: 1500,
|
||||
OutputTokensDelta: 400,
|
||||
CacheReadTokensDelta: 200,
|
||||
CacheWriteTokensDelta: 100,
|
||||
CostDelta: 0.0123,
|
||||
DurationMs: 2500,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("emit: %v", err)
|
||||
}
|
||||
if got.ToolCallCount != 3 {
|
||||
t.Errorf("ToolCallCount: got %d want 3", got.ToolCallCount)
|
||||
}
|
||||
if len(got.ToolNames) != 3 || got.ToolNames[0] != "bash" || got.ToolNames[2] != "bash" {
|
||||
t.Errorf("ToolNames: %v", got.ToolNames)
|
||||
}
|
||||
if got.LLMCallCount != 4 {
|
||||
t.Errorf("LLMCallCount: got %d want 4", got.LLMCallCount)
|
||||
}
|
||||
if got.InputTokensDelta != 1500 || got.OutputTokensDelta != 400 {
|
||||
t.Errorf("token deltas: %+v", got)
|
||||
}
|
||||
if got.CacheReadTokensDelta != 200 || got.CacheWriteTokensDelta != 100 {
|
||||
t.Errorf("cache deltas: %+v", got)
|
||||
}
|
||||
if got.CostDelta != 0.0123 {
|
||||
t.Errorf("CostDelta: got %v", got.CostDelta)
|
||||
}
|
||||
if got.DurationMs != 2500 {
|
||||
t.Errorf("DurationMs: got %d", got.DurationMs)
|
||||
}
|
||||
}
|
||||
@@ -669,6 +669,12 @@ func loadSingleExtension(path string) (*LoadedExtension, error) {
|
||||
return *r
|
||||
})
|
||||
},
|
||||
onLLMUsage: func(h func(LLMUsageEvent, Context)) {
|
||||
reg(LLMUsage, func(e Event, c Context) Result {
|
||||
h(e.(LLMUsageEvent), c)
|
||||
return nil
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
// Call Init — the extension registers its handlers, tools, commands.
|
||||
|
||||
@@ -2,9 +2,12 @@ package extensions
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log"
|
||||
"maps"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strconv"
|
||||
@@ -99,6 +102,10 @@ type Runner struct {
|
||||
customEventSubs map[string][]func(string) // inter-extension event bus
|
||||
optionOverrides map[string]string // runtime option overrides
|
||||
configStore *viper.Viper // per-instance config store (nil = global)
|
||||
state map[string]string // session-scoped extension state (last-write-wins)
|
||||
stateMu sync.RWMutex // guards state independently of mu
|
||||
saverMu sync.Mutex // serializes stateSaver invocations so atomic-rename writes don't interleave
|
||||
stateSaver func() // optional persistence hook invoked after each state mutation
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
@@ -264,6 +271,18 @@ func normalizeContext(ctx Context) Context {
|
||||
if ctx.GetEntries == nil {
|
||||
ctx.GetEntries = func(string) []ExtensionEntry { return nil }
|
||||
}
|
||||
if ctx.SetState == nil {
|
||||
ctx.SetState = func(string, string) {}
|
||||
}
|
||||
if ctx.GetState == nil {
|
||||
ctx.GetState = func(string) (string, bool) { return "", false }
|
||||
}
|
||||
if ctx.DeleteState == nil {
|
||||
ctx.DeleteState = func(string) {}
|
||||
}
|
||||
if ctx.ListState == nil {
|
||||
ctx.ListState = func() []string { return nil }
|
||||
}
|
||||
if ctx.GetOption == nil {
|
||||
ctx.GetOption = func(string) string { return "" }
|
||||
}
|
||||
@@ -745,6 +764,168 @@ func (r *Runner) GetMessageRenderer(name string) *MessageRendererConfig {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension state store (session-scoped, last-write-wins)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// SetState records a key-value pair in the runner's session-scoped extension
|
||||
// state store. The store is in-memory; callers wire SetStateSaver to persist
|
||||
// changes to a sidecar file. Thread-safe.
|
||||
//
|
||||
// When a saver is installed, concurrent SetState/DeleteState invocations are
|
||||
// serialized through saverMu so that overlapping snapshot-and-rename writes
|
||||
// cannot interleave (which would otherwise race on the shared tmp file and
|
||||
// risk persisting an older snapshot after a newer one).
|
||||
func (r *Runner) SetState(key, value string) {
|
||||
r.stateMu.Lock()
|
||||
if r.state == nil {
|
||||
r.state = make(map[string]string)
|
||||
}
|
||||
r.state[key] = value
|
||||
saver := r.stateSaver
|
||||
r.stateMu.Unlock()
|
||||
r.runSaver(saver)
|
||||
}
|
||||
|
||||
// GetState returns the value previously stored via SetState, plus a bool
|
||||
// indicating whether the key was present. Thread-safe.
|
||||
func (r *Runner) GetState(key string) (string, bool) {
|
||||
r.stateMu.RLock()
|
||||
defer r.stateMu.RUnlock()
|
||||
v, ok := r.state[key]
|
||||
return v, ok
|
||||
}
|
||||
|
||||
// DeleteState removes a key from the state store. No-op if the key is
|
||||
// missing. Thread-safe. Saver invocations are serialized via saverMu — see
|
||||
// SetState for the rationale.
|
||||
func (r *Runner) DeleteState(key string) {
|
||||
r.stateMu.Lock()
|
||||
_, existed := r.state[key]
|
||||
if existed {
|
||||
delete(r.state, key)
|
||||
}
|
||||
saver := r.stateSaver
|
||||
r.stateMu.Unlock()
|
||||
if !existed {
|
||||
return
|
||||
}
|
||||
r.runSaver(saver)
|
||||
}
|
||||
|
||||
// runSaver invokes the optional persistence callback under saverMu so
|
||||
// concurrent SetState/DeleteState writers cannot race on the shared tmp
|
||||
// file used by SaveStateToFile's atomic rename. The deferred Unlock
|
||||
// guarantees saverMu is released even if the saver panics.
|
||||
func (r *Runner) runSaver(saver func()) {
|
||||
if saver == nil {
|
||||
return
|
||||
}
|
||||
r.saverMu.Lock()
|
||||
defer r.saverMu.Unlock()
|
||||
saver()
|
||||
}
|
||||
|
||||
// ListState returns all keys currently in the state store, in unspecified
|
||||
// order. Thread-safe.
|
||||
func (r *Runner) ListState() []string {
|
||||
r.stateMu.RLock()
|
||||
defer r.stateMu.RUnlock()
|
||||
if len(r.state) == 0 {
|
||||
return nil
|
||||
}
|
||||
keys := make([]string, 0, len(r.state))
|
||||
for k := range r.state {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
// SetStateSaver installs an optional persistence hook invoked after each
|
||||
// mutation to the state store (SetState / DeleteState / LoadStateFromFile).
|
||||
// Pass nil to disable persistence. Thread-safe.
|
||||
func (r *Runner) SetStateSaver(saver func()) {
|
||||
r.stateMu.Lock()
|
||||
defer r.stateMu.Unlock()
|
||||
r.stateSaver = saver
|
||||
}
|
||||
|
||||
// SnapshotState returns a copy of the current state store as a
|
||||
// fresh map. Useful for persisting to disk without holding the lock.
|
||||
// Thread-safe.
|
||||
func (r *Runner) SnapshotState() map[string]string {
|
||||
r.stateMu.RLock()
|
||||
defer r.stateMu.RUnlock()
|
||||
if len(r.state) == 0 {
|
||||
return nil
|
||||
}
|
||||
copyMap := make(map[string]string, len(r.state))
|
||||
maps.Copy(copyMap, r.state)
|
||||
return copyMap
|
||||
}
|
||||
|
||||
// LoadStateFromFile reads a JSON map from path and replaces the in-memory
|
||||
// state store with its contents. Missing or empty files are treated as
|
||||
// "no prior state": the in-memory store is replaced with an empty map so
|
||||
// callers can safely switch sessions without leaking keys from a prior
|
||||
// session into a new one. Malformed JSON returns the parse error without
|
||||
// touching the existing store. Thread-safe.
|
||||
func (r *Runner) LoadStateFromFile(path string) error {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
r.stateMu.Lock()
|
||||
r.state = map[string]string{}
|
||||
r.stateMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("reading extension state: %w", err)
|
||||
}
|
||||
if len(data) == 0 {
|
||||
r.stateMu.Lock()
|
||||
r.state = map[string]string{}
|
||||
r.stateMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
var loaded map[string]string
|
||||
if err := json.Unmarshal(data, &loaded); err != nil {
|
||||
return fmt.Errorf("parsing extension state: %w", err)
|
||||
}
|
||||
r.stateMu.Lock()
|
||||
r.state = loaded
|
||||
r.stateMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// SaveStateToFile writes the current state store to path as JSON, creating
|
||||
// parent directories as needed. An empty store writes an empty object so
|
||||
// that consumers can distinguish "loaded but empty" from "never saved".
|
||||
// Writes are atomic via a tmp-file-and-rename sequence. Thread-safe.
|
||||
func (r *Runner) SaveStateToFile(path string) error {
|
||||
snap := r.SnapshotState()
|
||||
if snap == nil {
|
||||
snap = map[string]string{}
|
||||
}
|
||||
data, err := json.MarshalIndent(snap, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshalling extension state: %w", err)
|
||||
}
|
||||
if dir := filepath.Dir(path); dir != "." && dir != "" {
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return fmt.Errorf("creating state directory: %w", err)
|
||||
}
|
||||
}
|
||||
tmp := path + ".tmp"
|
||||
if err := os.WriteFile(tmp, data, 0o644); err != nil {
|
||||
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)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Hot-reload
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -768,7 +949,9 @@ func (r *Runner) Reload(exts []LoadedExtension) {
|
||||
r.uiVisibility = nil
|
||||
r.disabledTools = nil
|
||||
r.customEventSubs = nil
|
||||
// optionOverrides are intentionally preserved.
|
||||
// optionOverrides and state are intentionally preserved across reloads:
|
||||
// they represent user/session intent (not extension code) and would be
|
||||
// surprising to lose on a hot-reload.
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,262 @@
|
||||
package extensions
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestRunner_State_BasicSetGetDelete(t *testing.T) {
|
||||
r := NewRunner(nil)
|
||||
|
||||
if _, ok := r.GetState("missing"); ok {
|
||||
t.Fatal("expected GetState to return ok=false for missing key")
|
||||
}
|
||||
|
||||
r.SetState("a", "1")
|
||||
r.SetState("b", "2")
|
||||
r.SetState("a", "3") // last-write-wins
|
||||
|
||||
if v, ok := r.GetState("a"); !ok || v != "3" {
|
||||
t.Errorf("expected GetState(a)=(3,true), got (%q,%v)", v, ok)
|
||||
}
|
||||
if v, ok := r.GetState("b"); !ok || v != "2" {
|
||||
t.Errorf("expected GetState(b)=(2,true), got (%q,%v)", v, ok)
|
||||
}
|
||||
|
||||
keys := r.ListState()
|
||||
if len(keys) != 2 {
|
||||
t.Errorf("expected 2 keys, got %d (%v)", len(keys), keys)
|
||||
}
|
||||
|
||||
r.DeleteState("a")
|
||||
if _, ok := r.GetState("a"); ok {
|
||||
t.Error("expected key a to be gone after DeleteState")
|
||||
}
|
||||
if len(r.ListState()) != 1 {
|
||||
t.Errorf("expected 1 key after delete, got %v", r.ListState())
|
||||
}
|
||||
|
||||
// Deleting missing key is a no-op.
|
||||
r.DeleteState("never-there")
|
||||
}
|
||||
|
||||
func TestRunner_State_SaverFires(t *testing.T) {
|
||||
r := NewRunner(nil)
|
||||
var calls int
|
||||
var mu sync.Mutex
|
||||
r.SetStateSaver(func() {
|
||||
mu.Lock()
|
||||
calls++
|
||||
mu.Unlock()
|
||||
})
|
||||
|
||||
r.SetState("a", "1")
|
||||
r.SetState("a", "2")
|
||||
r.DeleteState("a")
|
||||
r.DeleteState("a") // missing → no save
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if calls != 3 {
|
||||
t.Errorf("expected saver to fire 3 times (2 sets + 1 delete), got %d", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunner_State_SaveAndLoadRoundTrip(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "ext-state.json")
|
||||
|
||||
r1 := NewRunner(nil)
|
||||
r1.SetState("k1", "v1")
|
||||
r1.SetState("k2", `{"json":"value"}`)
|
||||
if err := r1.SaveStateToFile(path); err != nil {
|
||||
t.Fatalf("SaveStateToFile: %v", err)
|
||||
}
|
||||
|
||||
// Verify file contains JSON map.
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("reading saved file: %v", err)
|
||||
}
|
||||
var parsed map[string]string
|
||||
if err := json.Unmarshal(data, &parsed); err != nil {
|
||||
t.Fatalf("unmarshalling: %v", err)
|
||||
}
|
||||
if parsed["k1"] != "v1" || parsed["k2"] != `{"json":"value"}` {
|
||||
t.Errorf("unexpected file contents: %v", parsed)
|
||||
}
|
||||
|
||||
r2 := NewRunner(nil)
|
||||
if err := r2.LoadStateFromFile(path); err != nil {
|
||||
t.Fatalf("LoadStateFromFile: %v", err)
|
||||
}
|
||||
if v, ok := r2.GetState("k1"); !ok || v != "v1" {
|
||||
t.Errorf("expected k1=v1 after load, got (%q,%v)", v, ok)
|
||||
}
|
||||
if v, ok := r2.GetState("k2"); !ok || v != `{"json":"value"}` {
|
||||
t.Errorf("expected k2 to round-trip, got %q", v)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunner_State_LoadMissingFileClearsState(t *testing.T) {
|
||||
// LoadStateFromFile is documented to "replace the in-memory state store
|
||||
// with its contents"; for a missing file that means clearing the store.
|
||||
// This is what makes session-switching safe: a new session that has not
|
||||
// yet written a sidecar must not inherit keys from a prior session.
|
||||
r := NewRunner(nil)
|
||||
r.SetState("a", "1")
|
||||
if err := r.LoadStateFromFile(filepath.Join(t.TempDir(), "does-not-exist.json")); err != nil {
|
||||
t.Errorf("expected nil error for missing file, got %v", err)
|
||||
}
|
||||
if _, ok := r.GetState("a"); ok {
|
||||
t.Error("expected pre-existing state to be cleared when target file is missing")
|
||||
}
|
||||
if keys := r.ListState(); keys != nil {
|
||||
t.Errorf("expected ListState() to be nil after clearing, got %v", keys)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunner_State_LoadEmptyFileClearsState(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "empty.json")
|
||||
if err := os.WriteFile(path, nil, 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
r := NewRunner(nil)
|
||||
r.SetState("a", "1")
|
||||
if err := r.LoadStateFromFile(path); err != nil {
|
||||
t.Errorf("expected nil error for empty file, got %v", err)
|
||||
}
|
||||
if _, ok := r.GetState("a"); ok {
|
||||
t.Error("expected pre-existing state to be cleared when target file is empty")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunner_State_LoadMalformedFileError(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "bad.json")
|
||||
if err := os.WriteFile(path, []byte("{not json"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
r := NewRunner(nil)
|
||||
if err := r.LoadStateFromFile(path); err == nil {
|
||||
t.Error("expected error loading malformed JSON, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunner_State_PersistenceViaSaver(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "ext-state.json")
|
||||
|
||||
r := NewRunner(nil)
|
||||
r.SetStateSaver(func() {
|
||||
_ = r.SaveStateToFile(path)
|
||||
})
|
||||
r.SetState("hello", "world")
|
||||
|
||||
// File should exist with the value already.
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("reading saved file: %v", err)
|
||||
}
|
||||
var parsed map[string]string
|
||||
if err := json.Unmarshal(data, &parsed); err != nil {
|
||||
t.Fatalf("unmarshalling: %v", err)
|
||||
}
|
||||
if parsed["hello"] != "world" {
|
||||
t.Errorf("expected file to contain hello=world, got %v", parsed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunner_State_ConcurrentSet(t *testing.T) {
|
||||
r := NewRunner(nil)
|
||||
var wg sync.WaitGroup
|
||||
const goroutines = 16
|
||||
const iterations = 100
|
||||
wg.Add(goroutines)
|
||||
for range goroutines {
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for range iterations {
|
||||
r.SetState("k", "v")
|
||||
_, _ = r.GetState("k")
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if v, ok := r.GetState("k"); !ok || v != "v" {
|
||||
t.Errorf("expected k=v after concurrent writes, got (%q,%v)", v, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunner_State_ContextNoOpsWhenUnset(t *testing.T) {
|
||||
// Verify normalizeContext installs safe no-ops for SetState/GetState/etc.
|
||||
// when not provided by the caller.
|
||||
ext := makeHandlerExt("state.go", map[EventType][]HandlerFunc{
|
||||
SessionStart: {
|
||||
func(e Event, c Context) Result {
|
||||
// All four state functions should be non-nil and safe to call.
|
||||
c.SetState("a", "b")
|
||||
if v, ok := c.GetState("a"); ok || v != "" {
|
||||
t.Errorf("no-op GetState should return (\"\", false); got (%q,%v)", v, ok)
|
||||
}
|
||||
c.DeleteState("a")
|
||||
if keys := c.ListState(); keys != nil {
|
||||
t.Errorf("no-op ListState should return nil; got %v", keys)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
})
|
||||
r := makeRunner(ext)
|
||||
// SetContext with empty Context to exercise normalizeContext defaults.
|
||||
r.SetContext(Context{})
|
||||
_, err := r.Emit(SessionStartEvent{})
|
||||
if err != nil {
|
||||
t.Fatalf("emit: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunner_State_SaverPanicReleasesSaverMu(t *testing.T) {
|
||||
// If the saver callback panics (e.g. disk full mid-write), runSaver
|
||||
// must still release saverMu so subsequent SetState/DeleteState calls
|
||||
// can make progress. Without `defer Unlock()` the lock would be
|
||||
// permanently held and the next write would deadlock.
|
||||
r := NewRunner(nil)
|
||||
var calls int
|
||||
r.SetStateSaver(func() {
|
||||
calls++
|
||||
if calls == 1 {
|
||||
panic("simulated disk-write failure")
|
||||
}
|
||||
})
|
||||
|
||||
// First call panics. Recover, then verify a follow-up call still works
|
||||
// without blocking (proving saverMu was released).
|
||||
func() {
|
||||
defer func() {
|
||||
if rec := recover(); rec == nil {
|
||||
t.Fatal("expected panic from first saver invocation")
|
||||
}
|
||||
}()
|
||||
r.SetState("a", "1")
|
||||
}()
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
r.SetState("b", "2") // would deadlock if saverMu were still held
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("SetState after saver panic blocked — saverMu was not released")
|
||||
}
|
||||
if calls != 2 {
|
||||
t.Errorf("expected saver to fire twice (panic + recovery write), got %d", calls)
|
||||
}
|
||||
}
|
||||
@@ -183,6 +183,7 @@ func Symbols() interp.Exports {
|
||||
"RetryEvent": reflect.ValueOf((*RetryEvent)(nil)),
|
||||
"PrepareStepEvent": reflect.ValueOf((*PrepareStepEvent)(nil)),
|
||||
"PrepareStepResult": reflect.ValueOf((*PrepareStepResult)(nil)),
|
||||
"LLMUsageEvent": reflect.ValueOf((*LLMUsageEvent)(nil)),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -189,5 +189,11 @@ func NewTestAPI(ext *LoadedExtension) API {
|
||||
return nil
|
||||
})
|
||||
},
|
||||
onLLMUsage: func(h func(LLMUsageEvent, Context)) {
|
||||
reg(LLMUsage, func(e Event, c Context) Result {
|
||||
h(e.(LLMUsageEvent), c)
|
||||
return nil
|
||||
})
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
package extensions
|
||||
|
||||
// ToolKind constants classify what a tool does, enabling UIs to render
|
||||
// appropriate visualizations (e.g. diff view for edit tools, command+output
|
||||
// for execute tools) and file trackers to identify which results contain
|
||||
// modifications.
|
||||
//
|
||||
// This is the single source of truth for tool-kind classification; the
|
||||
// pkg/kit SDK re-exports these constants.
|
||||
const (
|
||||
ToolKindExecute = "execute" // Shell execution (bash)
|
||||
ToolKindEdit = "edit" // File modification (edit, write)
|
||||
ToolKindRead = "read" // File reading (read, ls)
|
||||
ToolKindSearch = "search" // Content/file search (grep, find)
|
||||
ToolKindSubagent = "agent" // Subagent spawning (subagent)
|
||||
)
|
||||
|
||||
// coreToolKinds maps built-in tool names to their kind classification.
|
||||
// MCP and extension tools without an entry default to ToolKindExecute.
|
||||
var coreToolKinds = map[string]string{
|
||||
"bash": ToolKindExecute,
|
||||
"edit": ToolKindEdit,
|
||||
"write": ToolKindEdit,
|
||||
"read": ToolKindRead,
|
||||
"ls": ToolKindRead,
|
||||
"grep": ToolKindSearch,
|
||||
"find": ToolKindSearch,
|
||||
"subagent": ToolKindSubagent,
|
||||
}
|
||||
|
||||
// ToolKindFor returns the ToolKind for a given tool name, defaulting to
|
||||
// ToolKindExecute for unknown tools (including MCP tools).
|
||||
func ToolKindFor(toolName string) string {
|
||||
if kind, ok := coreToolKinds[toolName]; ok {
|
||||
return kind
|
||||
}
|
||||
return ToolKindExecute
|
||||
}
|
||||
@@ -40,27 +40,6 @@ func ExtensionToolsAsLLMTools(defs []ToolDef, runner *Runner) []fantasy.AgentToo
|
||||
return tools
|
||||
}
|
||||
|
||||
// coreToolKinds maps built-in tool names to their kind classification.
|
||||
var coreToolKinds = map[string]string{
|
||||
"bash": "execute",
|
||||
"edit": "edit",
|
||||
"write": "edit",
|
||||
"read": "read",
|
||||
"ls": "read",
|
||||
"grep": "search",
|
||||
"find": "search",
|
||||
"subagent": "agent",
|
||||
}
|
||||
|
||||
// toolKindFor returns the ToolKind for a given tool name, defaulting to
|
||||
// "execute" for unknown tools (including MCP tools).
|
||||
func toolKindFor(toolName string) string {
|
||||
if kind, ok := coreToolKinds[toolName]; ok {
|
||||
return kind
|
||||
}
|
||||
return "execute"
|
||||
}
|
||||
|
||||
// parseToolArgsJSON attempts to parse JSON-encoded tool args into a map.
|
||||
// Returns nil on failure (non-fatal convenience parsing).
|
||||
func parseToolArgsJSON(input string) map[string]any {
|
||||
@@ -93,7 +72,7 @@ func (w *wrappedTool) Run(ctx context.Context, call fantasy.ToolCall) (fantasy.T
|
||||
fmt.Sprintf("Error: tool %q is currently disabled", toolName)), nil
|
||||
}
|
||||
|
||||
kind := toolKindFor(toolName)
|
||||
kind := ToolKindFor(toolName)
|
||||
|
||||
// 1. Emit ToolCall — extensions can block execution.
|
||||
if w.runner.HasHandlers(ToolCall) {
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
package models
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestCopilotProviderAliasUsesCatalog(t *testing.T) {
|
||||
registry := NewModelsRegistry()
|
||||
|
||||
models, err := registry.GetModelsForProvider("copilot")
|
||||
if err != nil {
|
||||
t.Fatalf("GetModelsForProvider(copilot) failed: %v", err)
|
||||
}
|
||||
if len(models) == 0 {
|
||||
t.Fatal("expected copilot alias to return github-copilot catalog models")
|
||||
}
|
||||
if registry.LookupModel("copilot", "gpt-5.5") == nil {
|
||||
t.Fatal("expected copilot/gpt-5.5 to resolve through github-copilot catalog")
|
||||
}
|
||||
if registry.GetProviderInfo("copilot") == nil {
|
||||
t.Fatal("expected copilot alias to return github-copilot provider info")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopilotRejectsNonGPTModels(t *testing.T) {
|
||||
_, err := CreateProvider(t.Context(), &ProviderConfig{ModelString: "copilot/claude-sonnet-4.6"})
|
||||
if err == nil {
|
||||
t.Fatal("expected non-GPT Copilot model to be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopilotHTTPClientCachesToken(t *testing.T) {
|
||||
client := createCopilotHTTPClient("cached-token", time.Now().Add(time.Hour).Unix(), false)
|
||||
transport, ok := client.Transport.(*copilotTransport)
|
||||
if !ok {
|
||||
t.Fatal("expected *copilotTransport")
|
||||
}
|
||||
|
||||
token := transport.cachedToken(t.Context())
|
||||
if token != "cached-token" {
|
||||
t.Fatalf("expected cached token, got %q", token)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCopilotTransportHeaders(t *testing.T) {
|
||||
req, err := http.NewRequest(http.MethodGet, "https://example.com", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewRequest failed: %v", err)
|
||||
}
|
||||
|
||||
transport := &copilotTransport{
|
||||
base: roundTripFunc(func(req *http.Request) (*http.Response, error) {
|
||||
if req.Header.Get("Authorization") != "Bearer cached-token" {
|
||||
t.Fatalf("unexpected Authorization header: %q", req.Header.Get("Authorization"))
|
||||
}
|
||||
if req.Header.Get("Copilot-Integration-Id") != copilotIntegrationID {
|
||||
t.Fatalf("unexpected Copilot-Integration-Id header: %q", req.Header.Get("Copilot-Integration-Id"))
|
||||
}
|
||||
if req.Header.Get("Editor-Version") != copilotEditorVersion {
|
||||
t.Fatalf("unexpected Editor-Version header: %q", req.Header.Get("Editor-Version"))
|
||||
}
|
||||
if req.Header.Get("User-Agent") != copilotUserAgent {
|
||||
t.Fatalf("unexpected User-Agent header: %q", req.Header.Get("User-Agent"))
|
||||
}
|
||||
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody}, nil
|
||||
}),
|
||||
token: "cached-token",
|
||||
expiresAt: time.Now().Add(time.Hour).Unix(),
|
||||
}
|
||||
|
||||
resp, err := transport.RoundTrip(req)
|
||||
if err != nil {
|
||||
t.Fatalf("RoundTrip failed: %v", err)
|
||||
}
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return f(req)
|
||||
}
|
||||
@@ -69,13 +69,6 @@ func modelConfigToModelInfo(modelID string, cfg CustomModelConfig) ModelInfo {
|
||||
return info
|
||||
}
|
||||
|
||||
// LoadModelSettingsFromConfig loads per-model generation parameter overrides
|
||||
// from the process-global viper store. Keys are "provider/model" strings.
|
||||
// Returns nil if no model settings are configured.
|
||||
func LoadModelSettingsFromConfig() map[string]*GenerationParams {
|
||||
return LoadModelSettingsFrom(viper.GetViper())
|
||||
}
|
||||
|
||||
// LoadModelSettingsFrom loads per-model generation parameter overrides from the
|
||||
// supplied per-instance store. When v is nil the process-global store is used.
|
||||
// Keys are "provider/model" strings. Returns nil if no model settings are
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
@@ -33,6 +34,24 @@ import (
|
||||
const (
|
||||
// ClaudeCodePrompt is the required system prompt for OAuth authentication.
|
||||
ClaudeCodePrompt = "You are Claude Code, Anthropic's official CLI for Claude."
|
||||
|
||||
// copilotProviderID is the canonical models.dev provider key. The CLI also
|
||||
// accepts the shorter "copilot" alias for user-facing model strings.
|
||||
copilotProviderID = "github-copilot"
|
||||
// copilotAliasProviderID is the short provider prefix accepted by kit.
|
||||
copilotAliasProviderID = "copilot"
|
||||
// copilotBaseURL is the fallback API URL if the model catalog has no API URL.
|
||||
copilotBaseURL = "https://api.githubcopilot.com"
|
||||
|
||||
// GitHub Copilot currently expects VS Code Copilot Chat client identifiers.
|
||||
// Keep these centralized so they are easy to audit and update when GitHub
|
||||
// changes accepted client metadata.
|
||||
copilotIntegrationID = "vscode-chat"
|
||||
copilotEditorVersion = "vscode/1.104.1"
|
||||
copilotEditorPluginVersion = "copilot-chat/0.31.0"
|
||||
copilotUserAgent = "GitHubCopilotChat/0.31.0"
|
||||
copilotOpenAIIntent = "conversation-agent"
|
||||
copilotGitHubAPIVersion = "2026-01-09"
|
||||
)
|
||||
|
||||
// resolveModelAlias resolves model aliases to their full names using the registry
|
||||
@@ -215,6 +234,20 @@ func ParseModelString(modelString string) (provider, model string, err error) {
|
||||
return "", "", fmt.Errorf("invalid model format %q: expected provider/model (e.g. anthropic/claude-sonnet-4-5)", modelString)
|
||||
}
|
||||
|
||||
// isCopilotProvider reports whether provider is the canonical catalog key or
|
||||
// the user-facing shorthand alias.
|
||||
func isCopilotProvider(provider string) bool {
|
||||
return provider == copilotAliasProviderID || provider == copilotProviderID
|
||||
}
|
||||
|
||||
// catalogProviderID maps supported provider aliases to their models.dev keys.
|
||||
func catalogProviderID(provider string) string {
|
||||
if isCopilotProvider(provider) {
|
||||
return copilotProviderID
|
||||
}
|
||||
return provider
|
||||
}
|
||||
|
||||
// CreateProvider creates a fantasy LanguageModel based on the provider configuration.
|
||||
// Model metadata is looked up from the models.dev database for cost tracking and
|
||||
// capability detection, but unknown models are passed through to the provider
|
||||
@@ -238,17 +271,30 @@ func CreateProvider(ctx context.Context, config *ProviderConfig) (*ProviderResul
|
||||
}
|
||||
|
||||
registry := GetGlobalRegistry()
|
||||
lookupProvider := catalogProviderID(provider)
|
||||
|
||||
// Look up model metadata (advisory, not blocking).
|
||||
// Look up model metadata (advisory for most providers, strict for Copilot).
|
||||
// When the model is known we validate config limits and print
|
||||
// suggestions on likely typos; when unknown we let the provider
|
||||
// API be the authority.
|
||||
modelInfo := registry.LookupModel(provider, modelName)
|
||||
if modelInfo == nil && provider != "ollama" && config.ProviderURL == "" {
|
||||
// API be the authority except for Copilot, whose non-GPT catalog entries
|
||||
// require unsupported wire protocols.
|
||||
modelInfo := registry.LookupModel(lookupProvider, modelName)
|
||||
if isCopilotProvider(provider) {
|
||||
providerInfo := registry.GetProviderInfo(copilotProviderID)
|
||||
if providerInfo == nil {
|
||||
return nil, fmt.Errorf("unsupported provider: %s (not found in model database)", copilotProviderID)
|
||||
}
|
||||
if modelInfo == nil {
|
||||
if suggestions := registry.SuggestModels(copilotProviderID, modelName); len(suggestions) > 0 {
|
||||
return nil, fmt.Errorf("model %q not found for provider %s. Did you mean one of: %s", modelName, copilotProviderID, strings.Join(suggestions, ", "))
|
||||
}
|
||||
return nil, fmt.Errorf("model %q not found for provider %s", modelName, copilotProviderID)
|
||||
}
|
||||
} else if modelInfo == nil && provider != "ollama" && config.ProviderURL == "" {
|
||||
// Model not in database — warn with suggestions but don't block.
|
||||
if suggestions := registry.SuggestModels(provider, modelName); len(suggestions) > 0 {
|
||||
if suggestions := registry.SuggestModels(lookupProvider, modelName); len(suggestions) > 0 {
|
||||
fmt.Fprintf(os.Stderr, "Warning: model %q not found in model database for provider %s. Similar models: %s\n",
|
||||
modelName, provider, strings.Join(suggestions, ", "))
|
||||
modelName, lookupProvider, strings.Join(suggestions, ", "))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -282,6 +328,8 @@ func CreateProvider(ctx context.Context, config *ProviderConfig) (*ProviderResul
|
||||
result, createErr = createAnthropicProvider(ctx, config, modelName)
|
||||
case "openai":
|
||||
result, createErr = createOpenAIProvider(ctx, config, modelName)
|
||||
case "copilot", "github-copilot":
|
||||
result, createErr = createCopilotProvider(ctx, config, modelName)
|
||||
case "google", "gemini":
|
||||
result, createErr = createGoogleProvider(ctx, config, modelName)
|
||||
case "ollama":
|
||||
@@ -884,7 +932,7 @@ func createAnthropicProvider(ctx context.Context, config *ProviderConfig, modelN
|
||||
}
|
||||
|
||||
// Handle OAuth vs API key authentication
|
||||
if strings.HasPrefix(source, "stored OAuth") {
|
||||
if source == auth.CredentialSourceOAuth {
|
||||
httpClient := createOAuthHTTPClient(apiKey, config.TLSSkipVerify)
|
||||
opts = append(opts, anthropic.WithHTTPClient(httpClient))
|
||||
// Note: For OAuth, the API key is set as a placeholder; the transport handles auth
|
||||
@@ -1023,6 +1071,72 @@ func createOpenAIProvider(ctx context.Context, config *ProviderConfig, modelName
|
||||
return &ProviderResult{Model: model, ProviderOptions: providerOpts}, nil
|
||||
}
|
||||
|
||||
// createCopilotProvider builds a GitHub Copilot provider through fantasy's
|
||||
// OpenAI-compatible provider. The catalog key is github-copilot, but the public
|
||||
// model prefix may be either copilot/ or github-copilot/.
|
||||
//
|
||||
// Only gpt-* Copilot models are enabled here. The catalog also lists Claude and
|
||||
// Gemini Copilot models, but those require different wire protocols and must be
|
||||
// routed explicitly before they can be safely accepted.
|
||||
func createCopilotProvider(ctx context.Context, config *ProviderConfig, modelName string) (*ProviderResult, error) {
|
||||
if !strings.HasPrefix(modelName, "gpt-") {
|
||||
return nil, fmt.Errorf("GitHub Copilot model %q is not supported yet: only gpt-* models use the OpenAI-compatible protocol", modelName)
|
||||
}
|
||||
|
||||
cm, err := auth.NewCredentialManager()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to initialize credential manager: %w", err)
|
||||
}
|
||||
|
||||
token, err := cm.GetValidCopilotAccessTokenContext(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("GitHub Copilot credentials not available. Use 'kit auth login copilot': %w", err)
|
||||
}
|
||||
|
||||
expiresAt := int64(0)
|
||||
if creds, err := cm.GetCopilotCredentials(); err == nil && creds != nil && creds.CopilotAccessToken == token {
|
||||
expiresAt = creds.ExpiresAt
|
||||
}
|
||||
|
||||
baseURL := copilotBaseURL
|
||||
if providerInfo := GetGlobalRegistry().GetProviderInfo(copilotProviderID); providerInfo != nil && providerInfo.API != "" {
|
||||
baseURL = providerInfo.API
|
||||
}
|
||||
if config.ProviderURL != "" {
|
||||
baseURL = config.ProviderURL
|
||||
}
|
||||
|
||||
opts := []openai.Option{
|
||||
openai.WithName(copilotAliasProviderID),
|
||||
openai.WithBaseURL(baseURL),
|
||||
openai.WithAPIKey(token),
|
||||
openai.WithHTTPClient(createCopilotHTTPClient(token, expiresAt, config.TLSSkipVerify)),
|
||||
openai.WithUseResponsesAPI(),
|
||||
openai.WithResponsesAPIFunc(copilotUsesResponsesAPI),
|
||||
openai.WithObjectMode(fantasy.ObjectModeTool),
|
||||
}
|
||||
|
||||
provider, err := openai.New(opts...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create GitHub Copilot provider: %w", err)
|
||||
}
|
||||
|
||||
model, err := provider.LanguageModel(ctx, modelName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create GitHub Copilot model: %w", err)
|
||||
}
|
||||
|
||||
providerOpts := buildOpenAIProviderOptions(config, modelName)
|
||||
|
||||
return &ProviderResult{Model: model, ProviderOptions: providerOpts}, nil
|
||||
}
|
||||
|
||||
// copilotUsesResponsesAPI selects the OpenAI Responses API for Copilot models
|
||||
// known to support it. Non-gpt models are rejected before provider creation.
|
||||
func copilotUsesResponsesAPI(modelID string) bool {
|
||||
return strings.HasPrefix(modelID, "gpt-5")
|
||||
}
|
||||
|
||||
// createOpenAICodexProvider creates a provider for ChatGPT/Codex OAuth tokens.
|
||||
// Uses the chatgpt.com/backend-api/codex endpoint with special headers.
|
||||
func createOpenAICodexProvider(ctx context.Context, config *ProviderConfig, modelName, token, accountID string) (*ProviderResult, error) {
|
||||
@@ -1152,6 +1266,87 @@ func (t *codexTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
return t.base.RoundTrip(newReq)
|
||||
}
|
||||
|
||||
// createCopilotHTTPClient returns an HTTP client that injects Copilot-specific
|
||||
// authorization and client metadata headers. The token and expiry are cached in
|
||||
// the transport so streaming requests do not hit credentials.json on every
|
||||
// RoundTrip; the credential manager is consulted only near expiry.
|
||||
func createCopilotHTTPClient(token string, expiresAt int64, skipVerify bool) *http.Client {
|
||||
var base http.RoundTripper
|
||||
if skipVerify {
|
||||
base = &http.Transport{
|
||||
TLSClientConfig: &tls.Config{
|
||||
InsecureSkipVerify: true,
|
||||
},
|
||||
}
|
||||
} else {
|
||||
base = http.DefaultTransport
|
||||
}
|
||||
|
||||
return &http.Client{
|
||||
Transport: &copilotTransport{
|
||||
base: base,
|
||||
token: token,
|
||||
expiresAt: expiresAt,
|
||||
},
|
||||
Timeout: 120 * time.Second,
|
||||
}
|
||||
}
|
||||
|
||||
// copilotTransport decorates requests for api.githubcopilot.com.
|
||||
//
|
||||
// It owns a cached Copilot access token. When the token is still valid, the hot
|
||||
// path is in-memory only. Near expiry it refreshes through CredentialManager,
|
||||
// which updates both the cache here and credentials.json.
|
||||
type copilotTransport struct {
|
||||
base http.RoundTripper
|
||||
token string
|
||||
expiresAt int64
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func (t *copilotTransport) RoundTrip(req *http.Request) (*http.Response, error) {
|
||||
token := t.cachedToken(req.Context())
|
||||
|
||||
newReq := req.Clone(req.Context())
|
||||
newReq.Header.Set("Authorization", "Bearer "+token)
|
||||
newReq.Header.Set("Copilot-Integration-Id", copilotIntegrationID)
|
||||
newReq.Header.Set("Editor-Version", copilotEditorVersion)
|
||||
newReq.Header.Set("Editor-Plugin-Version", copilotEditorPluginVersion)
|
||||
newReq.Header.Set("Openai-Intent", copilotOpenAIIntent)
|
||||
newReq.Header.Set("User-Agent", copilotUserAgent)
|
||||
newReq.Header.Set("X-GitHub-Api-Version", copilotGitHubAPIVersion)
|
||||
|
||||
return t.base.RoundTrip(newReq)
|
||||
}
|
||||
|
||||
// cachedToken returns the cached token unless it is within the five-minute
|
||||
// refresh window. Refresh errors fall back to the last token so the request can
|
||||
// surface any authoritative auth failure from the Copilot API.
|
||||
func (t *copilotTransport) cachedToken(ctx context.Context) string {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
|
||||
if t.expiresAt == 0 || time.Now().Unix() < t.expiresAt-300 {
|
||||
return t.token
|
||||
}
|
||||
|
||||
cm, err := auth.NewCredentialManager()
|
||||
if err != nil {
|
||||
return t.token
|
||||
}
|
||||
|
||||
fresh, err := cm.GetValidCopilotAccessTokenContext(ctx)
|
||||
if err != nil || fresh == "" {
|
||||
return t.token
|
||||
}
|
||||
|
||||
t.token = fresh
|
||||
if creds, err := cm.GetCopilotCredentials(); err == nil && creds != nil && creds.CopilotAccessToken == fresh {
|
||||
t.expiresAt = creds.ExpiresAt
|
||||
}
|
||||
return t.token
|
||||
}
|
||||
|
||||
func createGoogleProvider(ctx context.Context, config *ProviderConfig, modelName string) (*ProviderResult, error) {
|
||||
apiKey := firstNonEmpty(
|
||||
config.ProviderAPIKey,
|
||||
|
||||
@@ -246,6 +246,7 @@ func loadEmbeddedProviders() map[string]modelsDBProvider {
|
||||
// doesn't track yet. Callers should treat a nil return as "unknown model"
|
||||
// and continue with sensible defaults.
|
||||
func (r *ModelsRegistry) LookupModel(provider, modelID string) *ModelInfo {
|
||||
provider = catalogProviderID(provider)
|
||||
providerInfo, exists := r.providers[provider]
|
||||
if !exists {
|
||||
return nil
|
||||
@@ -273,6 +274,7 @@ func LookupModelForSettings(modelString string) *ModelInfo {
|
||||
|
||||
// getRequiredEnvVars returns the required environment variables for a provider.
|
||||
func (r *ModelsRegistry) getRequiredEnvVars(provider string) ([]string, error) {
|
||||
provider = catalogProviderID(provider)
|
||||
providerInfo, exists := r.providers[provider]
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("unsupported provider: %s", provider)
|
||||
@@ -287,6 +289,7 @@ func (r *ModelsRegistry) getRequiredEnvVars(provider string) ([]string, error) {
|
||||
// variables. Returns nil for providers not in the registry (unknown
|
||||
// providers are assumed to handle auth themselves or via --provider-api-key).
|
||||
func (r *ModelsRegistry) ValidateEnvironment(provider string, apiKey string) error {
|
||||
provider = catalogProviderID(provider)
|
||||
if apiKey != "" {
|
||||
return nil
|
||||
}
|
||||
@@ -311,6 +314,15 @@ func (r *ModelsRegistry) ValidateEnvironment(provider string, apiKey string) err
|
||||
}
|
||||
}
|
||||
|
||||
// For GitHub Copilot, check stored GitHub OAuth credentials.
|
||||
if provider == copilotProviderID {
|
||||
if cm, err := auth.NewCredentialManager(); err == nil {
|
||||
if has, _ := cm.HasCopilotCredentials(); has {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
envVars, err := r.getRequiredEnvVars(provider)
|
||||
if err != nil {
|
||||
// Unknown provider — nothing to validate
|
||||
@@ -350,6 +362,7 @@ func (r *ModelsRegistry) ValidateEnvironment(provider string, apiKey string) err
|
||||
|
||||
// SuggestModels returns similar model names when an invalid model is provided.
|
||||
func (r *ModelsRegistry) SuggestModels(provider, invalidModel string) []string {
|
||||
provider = catalogProviderID(provider)
|
||||
providerInfo, exists := r.providers[provider]
|
||||
if !exists {
|
||||
return nil
|
||||
@@ -415,6 +428,7 @@ func isProviderLLMSupported(providerID string, info *ProviderInfo) bool {
|
||||
|
||||
// GetModelsForProvider returns all models for a specific provider.
|
||||
func (r *ModelsRegistry) GetModelsForProvider(provider string) (map[string]ModelInfo, error) {
|
||||
provider = catalogProviderID(provider)
|
||||
providerInfo, exists := r.providers[provider]
|
||||
if !exists {
|
||||
return nil, fmt.Errorf("unsupported provider: %s", provider)
|
||||
@@ -425,6 +439,7 @@ func (r *ModelsRegistry) GetModelsForProvider(provider string) (map[string]Model
|
||||
|
||||
// GetProviderInfo returns the full provider info, or nil if not found.
|
||||
func (r *ModelsRegistry) GetProviderInfo(provider string) *ProviderInfo {
|
||||
provider = catalogProviderID(provider)
|
||||
info, exists := r.providers[provider]
|
||||
if !exists {
|
||||
return nil
|
||||
|
||||
@@ -70,7 +70,8 @@ func ParseTemplate(path string) (*PromptTemplate, error) {
|
||||
}
|
||||
|
||||
// ParseCommandArgs splits a command line into arguments respecting quotes.
|
||||
// It handles single quotes, double quotes, and backslash escaping.
|
||||
// It handles single quotes, double quotes, backslash escaping, and splits on
|
||||
// spaces and tabs.
|
||||
func ParseCommandArgs(input string) []string {
|
||||
var args []string
|
||||
var current strings.Builder
|
||||
@@ -78,7 +79,7 @@ func ParseCommandArgs(input string) []string {
|
||||
inDoubleQuote := false
|
||||
escaped := false
|
||||
|
||||
for i, r := range input {
|
||||
for _, r := range input {
|
||||
if escaped {
|
||||
current.WriteRune(r)
|
||||
escaped = false
|
||||
@@ -101,7 +102,7 @@ func ParseCommandArgs(input string) []string {
|
||||
continue
|
||||
}
|
||||
|
||||
if r == ' ' && !inSingleQuote && !inDoubleQuote {
|
||||
if (r == ' ' || r == '\t') && !inSingleQuote && !inDoubleQuote {
|
||||
if current.Len() > 0 {
|
||||
args = append(args, current.String())
|
||||
current.Reset()
|
||||
@@ -110,7 +111,6 @@ func ParseCommandArgs(input string) []string {
|
||||
}
|
||||
|
||||
current.WriteRune(r)
|
||||
_ = i // silence unused warning when we need position later
|
||||
}
|
||||
|
||||
if current.Len() > 0 {
|
||||
@@ -325,8 +325,3 @@ func (t *PromptTemplate) Expand(argsInput string) string {
|
||||
args := ParseCommandArgs(argsInput)
|
||||
return SubstituteArgs(t.Content, args)
|
||||
}
|
||||
|
||||
// ExpandWithArgs substitutes the provided arguments into the template content.
|
||||
func (t *PromptTemplate) ExpandWithArgs(args []string) string {
|
||||
return SubstituteArgs(t.Content, args)
|
||||
}
|
||||
|
||||
@@ -18,8 +18,11 @@ type PromptTemplate struct {
|
||||
Variables []string
|
||||
}
|
||||
|
||||
// variableRe matches {{variable_name}} placeholders.
|
||||
var variableRe = regexp.MustCompile(`\{\{(\w+)\}\}`)
|
||||
// variableRe matches {{variable_name}} placeholders, tolerating surrounding
|
||||
// whitespace inside the braces (e.g. {{ name }}). This is the canonical
|
||||
// template grammar shared by skill prompts and the extension template API
|
||||
// (pkg/kit ParseTemplate/RenderTemplate delegate here).
|
||||
var variableRe = regexp.MustCompile(`\{\{\s*(\w+)\s*\}\}`)
|
||||
|
||||
// NewPromptTemplate creates a PromptTemplate, automatically extracting
|
||||
// variable names from {{...}} placeholders in content.
|
||||
@@ -50,11 +53,13 @@ func LoadPromptTemplate(path string) (*PromptTemplate, error) {
|
||||
// Expand replaces all {{variable}} placeholders with values from the
|
||||
// provided map. Missing variables are left as-is (no error).
|
||||
func (t *PromptTemplate) Expand(values map[string]string) string {
|
||||
result := t.Content
|
||||
for k, v := range values {
|
||||
result = strings.ReplaceAll(result, "{{"+k+"}}", v)
|
||||
}
|
||||
return result
|
||||
return variableRe.ReplaceAllStringFunc(t.Content, func(m string) string {
|
||||
name := variableRe.FindStringSubmatch(m)[1]
|
||||
if v, ok := values[name]; ok {
|
||||
return v
|
||||
}
|
||||
return m
|
||||
})
|
||||
}
|
||||
|
||||
// ExpandStrict replaces all {{variable}} placeholders and returns an error
|
||||
|
||||
+60
-57
@@ -641,30 +641,16 @@ func (m *MCPToolManager) ExecuteTool(ctx context.Context, prefixedName, inputJSO
|
||||
Request: mcp.Request{Method: "tools/call"},
|
||||
Params: callParams,
|
||||
}
|
||||
result, callErr := conn.client.CallTool(ctx, callRequest)
|
||||
if callErr != nil {
|
||||
if m.connectionPool.oauthFlow != nil && IsOAuthError(callErr) {
|
||||
if flowErr := m.connectionPool.oauthFlow.RunAuthFlow(ctx, mapping.serverName, callErr); flowErr != nil {
|
||||
return nil, fmt.Errorf("OAuth re-authorization failed for tool %s: %w", mapping.originalName, flowErr)
|
||||
}
|
||||
result, callErr = conn.client.CallTool(ctx, callRequest)
|
||||
if callErr != nil {
|
||||
m.connectionPool.HandleConnectionError(mapping.serverName, callErr)
|
||||
return nil, fmt.Errorf("failed to call mcp tool after re-auth: %w", callErr)
|
||||
}
|
||||
} else {
|
||||
m.connectionPool.HandleConnectionError(mapping.serverName, callErr)
|
||||
return nil, fmt.Errorf("failed to call mcp tool: %w", callErr)
|
||||
}
|
||||
var result *mcp.CallToolResult
|
||||
err := m.withOAuthRetry(ctx, mapping.serverName, mapping.originalName, func() error {
|
||||
var callErr error
|
||||
result, callErr = conn.client.CallTool(ctx, callRequest)
|
||||
return callErr
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
marshaledResult, mErr := json.Marshal(result)
|
||||
if mErr != nil {
|
||||
return nil, fmt.Errorf("failed to marshal mcp tool result: %w", mErr)
|
||||
}
|
||||
return &MCPToolResult{
|
||||
Content: string(marshaledResult),
|
||||
IsError: result.IsError,
|
||||
}, nil
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
// Task-augmented path. Bypass the upstream CallTool helper because its
|
||||
@@ -683,40 +669,25 @@ func (m *MCPToolManager) ExecuteTool(ctx context.Context, prefixedName, inputJSO
|
||||
m.connectionPool.HandleConnectionError(mapping.serverName, callErr)
|
||||
return nil, fmt.Errorf("failed to call mcp tool: %w", callErr)
|
||||
}
|
||||
marshaledResult, mErr := json.Marshal(result)
|
||||
if mErr != nil {
|
||||
return nil, fmt.Errorf("failed to marshal mcp tool result: %w", mErr)
|
||||
}
|
||||
return &MCPToolResult{Content: string(marshaledResult), IsError: result.IsError}, nil
|
||||
return marshalToolResult(result)
|
||||
}
|
||||
|
||||
callResult, taskResult, callErr := callToolWithTask(ctx, rawClient, callParams)
|
||||
if callErr != nil {
|
||||
if m.connectionPool.oauthFlow != nil && IsOAuthError(callErr) {
|
||||
if flowErr := m.connectionPool.oauthFlow.RunAuthFlow(ctx, mapping.serverName, callErr); flowErr != nil {
|
||||
return nil, fmt.Errorf("OAuth re-authorization failed for tool %s: %w", mapping.originalName, flowErr)
|
||||
}
|
||||
callResult, taskResult, callErr = callToolWithTask(ctx, rawClient, callParams)
|
||||
if callErr != nil {
|
||||
m.connectionPool.HandleConnectionError(mapping.serverName, callErr)
|
||||
return nil, fmt.Errorf("failed to call mcp tool after re-auth: %w", callErr)
|
||||
}
|
||||
} else {
|
||||
m.connectionPool.HandleConnectionError(mapping.serverName, callErr)
|
||||
return nil, fmt.Errorf("failed to call mcp tool: %w", callErr)
|
||||
}
|
||||
var (
|
||||
callResult *mcp.CallToolResult
|
||||
taskResult *mcp.CreateTaskResult
|
||||
)
|
||||
err = m.withOAuthRetry(ctx, mapping.serverName, mapping.originalName, func() error {
|
||||
var callErr error
|
||||
callResult, taskResult, callErr = callToolWithTask(ctx, rawClient, callParams)
|
||||
return callErr
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Server chose to answer synchronously — same shape as the no-task path.
|
||||
if callResult != nil {
|
||||
marshaledResult, mErr := json.Marshal(callResult)
|
||||
if mErr != nil {
|
||||
return nil, fmt.Errorf("failed to marshal mcp tool result: %w", mErr)
|
||||
}
|
||||
return &MCPToolResult{
|
||||
Content: string(marshaledResult),
|
||||
IsError: callResult.IsError,
|
||||
}, nil
|
||||
return marshalToolResult(callResult)
|
||||
}
|
||||
|
||||
// Asynchronous task path: poll until terminal, then return the result.
|
||||
@@ -732,18 +703,50 @@ func (m *MCPToolManager) ExecuteTool(ctx context.Context, prefixedName, inputJSO
|
||||
}
|
||||
|
||||
// Adapt TaskResultResult → CallToolResult for downstream JSON shape parity.
|
||||
adapted := &mcp.CallToolResult{
|
||||
return marshalToolResult(&mcp.CallToolResult{
|
||||
Content: final.Content,
|
||||
StructuredContent: final.StructuredContent,
|
||||
IsError: final.IsError,
|
||||
})
|
||||
}
|
||||
|
||||
// withOAuthRetry runs call once; when it fails with an OAuth error and an
|
||||
// OAuth flow is configured, it re-authorizes the server and retries once.
|
||||
// Connection failures are reported to the pool and wrapped uniformly. This
|
||||
// consolidates the retry/error chain shared by the synchronous and
|
||||
// task-augmented tool-call paths.
|
||||
func (m *MCPToolManager) withOAuthRetry(ctx context.Context, serverName, toolName string, call func() error) error {
|
||||
callErr := call()
|
||||
if callErr == nil {
|
||||
return nil
|
||||
}
|
||||
marshaledResult, mErr := json.Marshal(adapted)
|
||||
if mErr != nil {
|
||||
return nil, fmt.Errorf("failed to marshal mcp tool result: %w", mErr)
|
||||
if m.connectionPool.oauthFlow != nil && IsOAuthError(callErr) {
|
||||
if flowErr := m.connectionPool.oauthFlow.RunAuthFlow(ctx, serverName, callErr); flowErr != nil {
|
||||
return fmt.Errorf("OAuth re-authorization failed for tool %s: %w", toolName, flowErr)
|
||||
}
|
||||
if callErr = call(); callErr != nil {
|
||||
m.connectionPool.HandleConnectionError(serverName, callErr)
|
||||
return fmt.Errorf("failed to call mcp tool after re-auth: %w", callErr)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
m.connectionPool.HandleConnectionError(serverName, callErr)
|
||||
return fmt.Errorf("failed to call mcp tool: %w", callErr)
|
||||
}
|
||||
|
||||
// marshalToolResult converts an MCP CallToolResult into the JSON-encoded
|
||||
// MCPToolResult shape returned to the agent.
|
||||
func marshalToolResult(result *mcp.CallToolResult) (*MCPToolResult, error) {
|
||||
if result == nil {
|
||||
return nil, errors.New("mcp tool call returned nil result")
|
||||
}
|
||||
marshaled, err := json.Marshal(result)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal mcp tool result: %w", err)
|
||||
}
|
||||
return &MCPToolResult{
|
||||
Content: string(marshaledResult),
|
||||
IsError: final.IsError,
|
||||
Content: string(marshaled),
|
||||
IsError: result.IsError,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
+29
-35
@@ -2,7 +2,6 @@ package ui
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/mark3labs/kit/internal/auth"
|
||||
"github.com/mark3labs/kit/internal/models"
|
||||
@@ -44,28 +43,39 @@ func parseModelName(modelString string) (provider, model string) {
|
||||
// ollama or unrecognised models). This is used by the interactive TUI path
|
||||
// which doesn't go through SetupCLI.
|
||||
func CreateUsageTracker(modelString, providerAPIKey string) *UsageTracker {
|
||||
provider, model := parseModelName(modelString)
|
||||
if provider == "unknown" || model == "unknown" || provider == "ollama" {
|
||||
return nil
|
||||
}
|
||||
|
||||
registry := models.GetGlobalRegistry()
|
||||
modelInfo := registry.LookupModel(provider, model)
|
||||
modelInfo, provider := lookupTrackableModel(modelString)
|
||||
if modelInfo == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
isOAuth := false
|
||||
if provider == "anthropic" {
|
||||
_, source, err := auth.GetAnthropicAPIKey(providerAPIKey)
|
||||
if err == nil && strings.HasPrefix(source, "stored OAuth") {
|
||||
isOAuth = true
|
||||
}
|
||||
}
|
||||
|
||||
isOAuth := provider == "anthropic" && auth.IsAnthropicOAuth(providerAPIKey)
|
||||
return NewUsageTracker(modelInfo, provider, 80, isOAuth)
|
||||
}
|
||||
|
||||
// UpdateUsageTrackerForModel refreshes an existing tracker after a model
|
||||
// switch so token counting and cost reporting use the new model's metadata.
|
||||
// No-op for a nil tracker or untrackable models (unknown/ollama).
|
||||
func UpdateUsageTrackerForModel(t *UsageTracker, modelString, providerAPIKey string) {
|
||||
if t == nil {
|
||||
return
|
||||
}
|
||||
modelInfo, provider := lookupTrackableModel(modelString)
|
||||
if modelInfo == nil {
|
||||
return
|
||||
}
|
||||
isOAuth := provider == "anthropic" && auth.IsAnthropicOAuth(providerAPIKey)
|
||||
t.UpdateModelInfo(modelInfo, provider, isOAuth)
|
||||
}
|
||||
|
||||
// lookupTrackableModel resolves a model string to registry metadata, returning
|
||||
// nil for models without usage tracking support (unknown or ollama models).
|
||||
func lookupTrackableModel(modelString string) (*models.ModelInfo, string) {
|
||||
provider, model := parseModelName(modelString)
|
||||
if provider == "unknown" || model == "unknown" || provider == "ollama" {
|
||||
return nil, provider
|
||||
}
|
||||
return models.GetGlobalRegistry().LookupModel(provider, model), provider
|
||||
}
|
||||
|
||||
// SetupCLI creates, configures, and initializes a CLI instance with the provided
|
||||
// options. It sets up model display, usage tracking for supported providers, and
|
||||
// shows initial loading information. Returns nil in quiet mode or an initialized
|
||||
@@ -89,24 +99,8 @@ func SetupCLI(opts *CLISetupOptions) (*CLI, error) {
|
||||
}
|
||||
|
||||
// Set up usage tracking for supported providers
|
||||
if provider != "unknown" && model != "unknown" {
|
||||
// Skip usage tracking for ollama as it's not in models.dev
|
||||
if provider != "ollama" {
|
||||
registry := models.GetGlobalRegistry()
|
||||
if modelInfo := registry.LookupModel(provider, model); modelInfo != nil {
|
||||
// Check if OAuth credentials are being used for Anthropic models
|
||||
isOAuth := false
|
||||
if provider == "anthropic" {
|
||||
_, source, err := auth.GetAnthropicAPIKey(opts.ProviderAPIKey)
|
||||
if err == nil && strings.HasPrefix(source, "stored OAuth") {
|
||||
isOAuth = true
|
||||
}
|
||||
}
|
||||
|
||||
usageTracker := NewUsageTracker(modelInfo, provider, 80, isOAuth) // Will be updated with actual width
|
||||
cli.SetUsageTracker(usageTracker)
|
||||
}
|
||||
}
|
||||
if usageTracker := CreateUsageTracker(opts.ModelString, opts.ProviderAPIKey); usageTracker != nil {
|
||||
cli.SetUsageTracker(usageTracker)
|
||||
}
|
||||
|
||||
// Display model info (the system message block provides its own spacing).
|
||||
|
||||
+84
-116
@@ -1208,53 +1208,7 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||||
m.modelSelector = nil
|
||||
m.state = stateInput
|
||||
if m.setModel != nil {
|
||||
previousModel := m.providerName + "/" + m.modelName
|
||||
|
||||
// Check if thinking level needs adjustment for the new model.
|
||||
// Some models (e.g., OpenAI gpt-5.4) don't support "minimal" and require "none".
|
||||
if m.thinkingLevel != "" && m.thinkingLevel != "off" {
|
||||
parts := strings.SplitN(msg.ModelString, "/", 2)
|
||||
if len(parts) == 2 {
|
||||
modelName := parts[1]
|
||||
currentLevel := models.ParseThinkingLevel(m.thinkingLevel)
|
||||
if !models.IsValidThinkingLevelForModel(currentLevel, modelName) {
|
||||
fallback := models.SuggestThinkingLevelFallback(currentLevel, modelName)
|
||||
if fallback != models.ThinkingOff {
|
||||
m.printSystemMessage(fmt.Sprintf(
|
||||
"Note: Model %s doesn't support '%s' thinking level. Adjusted to '%s'.",
|
||||
modelName, currentLevel, fallback,
|
||||
))
|
||||
m.thinkingLevel = string(fallback)
|
||||
if m.setThinkingLevel != nil {
|
||||
_ = m.setThinkingLevel(string(fallback))
|
||||
}
|
||||
go func() { _ = prefs.SaveThinkingLevelPreference(string(fallback)) }()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err := m.setModel(msg.ModelString); err != nil {
|
||||
m.printSystemMessage(fmt.Sprintf("Failed to switch model: %v", err))
|
||||
} else {
|
||||
// Update display state directly — we cannot use
|
||||
// NotifyModelChanged (prog.Send) from inside Update()
|
||||
// without deadlocking BubbleTea.
|
||||
parts := strings.SplitN(msg.ModelString, "/", 2)
|
||||
if len(parts) == 2 {
|
||||
m.providerName = parts[0]
|
||||
m.modelName = parts[1]
|
||||
}
|
||||
m.printSystemMessage(fmt.Sprintf("Switched to %s", msg.ModelString))
|
||||
// Persist model selection for next launch.
|
||||
go func() { _ = prefs.SaveModelPreference(msg.ModelString) }()
|
||||
if m.emitModelChange != nil {
|
||||
emit := m.emitModelChange
|
||||
newModel := msg.ModelString
|
||||
prev := previousModel
|
||||
go emit(newModel, prev, "user")
|
||||
}
|
||||
}
|
||||
m.switchModel(msg.ModelString)
|
||||
}
|
||||
return m, tea.Batch(cmds...)
|
||||
|
||||
@@ -4211,11 +4165,31 @@ func (m *AppModel) handleModelCommand(args string) tea.Cmd {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Direct model switch with the provided model string.
|
||||
m.switchModel(args)
|
||||
return nil
|
||||
}
|
||||
|
||||
// switchModel performs a direct model switch, shared by the model selector
|
||||
// overlay and the /model slash command: it adjusts the thinking level when
|
||||
// the new model doesn't support the current one, calls the setModel
|
||||
// callback, updates display state, persists preferences, and emits the
|
||||
// ModelChange extension event.
|
||||
//
|
||||
// Display state is updated directly — we cannot use NotifyModelChanged
|
||||
// (prog.Send) from inside Update() without deadlocking BubbleTea.
|
||||
func (m *AppModel) switchModel(modelString string) {
|
||||
if m.setModel == nil {
|
||||
m.printSystemMessage("Model switching is not available.")
|
||||
return
|
||||
}
|
||||
|
||||
previousModel := m.providerName + "/" + m.modelName
|
||||
|
||||
// Check if thinking level needs adjustment for the new model.
|
||||
// Some models (e.g., OpenAI gpt-5.4) don't support "minimal" and require "none".
|
||||
if m.thinkingLevel != "" && m.thinkingLevel != "off" {
|
||||
parts := strings.SplitN(args, "/", 2)
|
||||
if len(parts) == 2 {
|
||||
if parts := strings.SplitN(modelString, "/", 2); len(parts) == 2 {
|
||||
modelName := parts[1]
|
||||
currentLevel := models.ParseThinkingLevel(m.thinkingLevel)
|
||||
if !models.IsValidThinkingLevelForModel(currentLevel, modelName) {
|
||||
@@ -4235,32 +4209,26 @@ func (m *AppModel) handleModelCommand(args string) tea.Cmd {
|
||||
}
|
||||
}
|
||||
|
||||
// Direct model switch with the provided model string.
|
||||
previousModel := m.providerName + "/" + m.modelName
|
||||
if err := m.setModel(args); err != nil {
|
||||
if err := m.setModel(modelString); err != nil {
|
||||
m.printSystemMessage(fmt.Sprintf("Failed to switch model: %v", err))
|
||||
return nil
|
||||
return
|
||||
}
|
||||
|
||||
// Update display state directly (cannot use prog.Send from Update).
|
||||
parts := strings.SplitN(args, "/", 2)
|
||||
if len(parts) == 2 {
|
||||
if parts := strings.SplitN(modelString, "/", 2); len(parts) == 2 {
|
||||
m.providerName = parts[0]
|
||||
m.modelName = parts[1]
|
||||
}
|
||||
|
||||
if m.emitModelChange != nil {
|
||||
emit := m.emitModelChange
|
||||
prev := previousModel
|
||||
newModel := args
|
||||
go emit(newModel, prev, "user")
|
||||
}
|
||||
m.printSystemMessage(fmt.Sprintf("Switched to %s", modelString))
|
||||
|
||||
// Persist model selection for next launch.
|
||||
go func() { _ = prefs.SaveModelPreference(args) }()
|
||||
go func() { _ = prefs.SaveModelPreference(modelString) }()
|
||||
|
||||
m.printSystemMessage(fmt.Sprintf("Switched to %s", args))
|
||||
return nil
|
||||
if m.emitModelChange != nil {
|
||||
emit := m.emitModelChange
|
||||
go emit(modelString, previousModel, "user")
|
||||
}
|
||||
}
|
||||
|
||||
// --------------------------------------------------------------------------
|
||||
@@ -4827,61 +4795,11 @@ func (m *AppModel) handleShareCommand() tea.Cmd {
|
||||
return r
|
||||
}, name)
|
||||
|
||||
tmpFile, err := os.CreateTemp("", fmt.Sprintf("kit-%s-*.jsonl", name))
|
||||
tmpPath, err := buildShareFile(name, data, sysPromptJSON)
|
||||
if err != nil {
|
||||
m.printSystemMessage(fmt.Sprintf("Failed to create temp file: %v", err))
|
||||
m.printSystemMessage(fmt.Sprintf("Failed to share session: %v", err))
|
||||
return nil
|
||||
}
|
||||
tmpPath := tmpFile.Name()
|
||||
|
||||
// Write the session data with the system prompt entry inserted after the header.
|
||||
// The header is the first line, so we write:
|
||||
// 1. First line (header) from original data
|
||||
// 2. System prompt entry
|
||||
// 3. Remaining lines from original data
|
||||
lines := strings.Split(string(data), "\n")
|
||||
if len(lines) > 0 && lines[len(lines)-1] == "" {
|
||||
lines = lines[:len(lines)-1] // Remove trailing empty line
|
||||
}
|
||||
|
||||
if len(lines) > 0 {
|
||||
// Write header (first line)
|
||||
if _, err := tmpFile.WriteString(lines[0] + "\n"); err != nil {
|
||||
_ = tmpFile.Close()
|
||||
_ = os.Remove(tmpPath)
|
||||
m.printSystemMessage(fmt.Sprintf("Failed to write temp file: %v", err))
|
||||
return nil
|
||||
}
|
||||
|
||||
// Write system prompt entry
|
||||
if _, err := tmpFile.Write(sysPromptJSON); err != nil {
|
||||
_ = tmpFile.Close()
|
||||
_ = os.Remove(tmpPath)
|
||||
m.printSystemMessage(fmt.Sprintf("Failed to write system prompt: %v", err))
|
||||
return nil
|
||||
}
|
||||
if _, err := tmpFile.WriteString("\n"); err != nil {
|
||||
_ = tmpFile.Close()
|
||||
_ = os.Remove(tmpPath)
|
||||
m.printSystemMessage(fmt.Sprintf("Failed to write temp file: %v", err))
|
||||
return nil
|
||||
}
|
||||
|
||||
// Write remaining lines
|
||||
for i := 1; i < len(lines); i++ {
|
||||
if lines[i] == "" {
|
||||
continue // Skip empty lines
|
||||
}
|
||||
if _, err := tmpFile.WriteString(lines[i] + "\n"); err != nil {
|
||||
_ = tmpFile.Close()
|
||||
_ = os.Remove(tmpPath)
|
||||
m.printSystemMessage(fmt.Sprintf("Failed to write temp file: %v", err))
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_ = tmpFile.Close()
|
||||
|
||||
m.printSystemMessage("Uploading session to GitHub Gist...")
|
||||
|
||||
@@ -4907,6 +4825,56 @@ func (m *AppModel) handleShareCommand() tea.Cmd {
|
||||
}
|
||||
}
|
||||
|
||||
// buildShareFile assembles a temp JSONL file containing the session data
|
||||
// with the system-prompt entry inserted after the header line. On success
|
||||
// the caller owns the returned file and must remove it when done; on error
|
||||
// any partially-written temp file has already been cleaned up.
|
||||
func buildShareFile(name string, data, sysPromptJSON []byte) (tmpPath string, err error) {
|
||||
tmpFile, err := os.CreateTemp("", fmt.Sprintf("kit-%s-*.jsonl", name))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("create temp file: %w", err)
|
||||
}
|
||||
tmpPath = tmpFile.Name()
|
||||
defer func() {
|
||||
_ = tmpFile.Close()
|
||||
if err != nil {
|
||||
_ = os.Remove(tmpPath)
|
||||
}
|
||||
}()
|
||||
|
||||
// Write the session data with the system prompt entry inserted after the
|
||||
// header. The header is the first line, so we write:
|
||||
// 1. First line (header) from original data
|
||||
// 2. System prompt entry
|
||||
// 3. Remaining lines from original data
|
||||
lines := strings.Split(string(data), "\n")
|
||||
if len(lines) > 0 && lines[len(lines)-1] == "" {
|
||||
lines = lines[:len(lines)-1] // Remove trailing empty line
|
||||
}
|
||||
if len(lines) == 0 {
|
||||
return tmpPath, nil
|
||||
}
|
||||
|
||||
if _, err = tmpFile.WriteString(lines[0] + "\n"); err != nil {
|
||||
return "", fmt.Errorf("write temp file: %w", err)
|
||||
}
|
||||
if _, err = tmpFile.Write(sysPromptJSON); err != nil {
|
||||
return "", fmt.Errorf("write system prompt: %w", err)
|
||||
}
|
||||
if _, err = tmpFile.WriteString("\n"); err != nil {
|
||||
return "", fmt.Errorf("write temp file: %w", err)
|
||||
}
|
||||
for i := 1; i < len(lines); i++ {
|
||||
if lines[i] == "" {
|
||||
continue // Skip empty lines
|
||||
}
|
||||
if _, err = tmpFile.WriteString(lines[i] + "\n"); err != nil {
|
||||
return "", fmt.Errorf("write temp file: %w", err)
|
||||
}
|
||||
}
|
||||
return tmpPath, nil
|
||||
}
|
||||
|
||||
// handleImportCommand imports a session from a JSONL file.
|
||||
// Usage: /import path.jsonl
|
||||
func (m *AppModel) handleImportCommand(args string) tea.Cmd {
|
||||
|
||||
+2
-2
@@ -11,12 +11,12 @@ import (
|
||||
// treeManagerAdapter adapts TreeManager to SessionManager interface.
|
||||
// This is unexported - users don't interact with it directly.
|
||||
type treeManagerAdapter struct {
|
||||
inner *session.TreeManager
|
||||
inner *TreeManager
|
||||
}
|
||||
|
||||
// NewTreeManagerAdapter creates an adapter (exported for use in New function).
|
||||
// This is used by the SDK when no custom SessionManager is provided.
|
||||
func NewTreeManagerAdapter(tm *session.TreeManager) SessionManager {
|
||||
func NewTreeManagerAdapter(tm *TreeManager) SessionManager {
|
||||
return &treeManagerAdapter{inner: tm}
|
||||
}
|
||||
|
||||
|
||||
@@ -17,6 +17,9 @@ type AnthropicCredentials = auth.AnthropicCredentials
|
||||
// and API key authentication methods.
|
||||
type OpenAICredentials = auth.OpenAICredentials
|
||||
|
||||
// CopilotCredentials holds GitHub OAuth and Copilot API credentials.
|
||||
type CopilotCredentials = auth.CopilotCredentials
|
||||
|
||||
// CredentialStore holds all stored credentials for various providers.
|
||||
type CredentialStore = auth.CredentialStore
|
||||
|
||||
@@ -65,6 +68,37 @@ func HasOpenAICredentials() bool {
|
||||
return has
|
||||
}
|
||||
|
||||
// HasCopilotCredentials checks if valid GitHub Copilot credentials are stored.
|
||||
func HasCopilotCredentials() bool {
|
||||
cm, err := auth.NewCredentialManager()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
has, err := cm.HasCopilotCredentials()
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return has
|
||||
}
|
||||
|
||||
// GetCopilotCredentials retrieves stored GitHub Copilot credentials.
|
||||
func GetCopilotCredentials() (*CopilotCredentials, error) {
|
||||
cm, err := auth.NewCredentialManager()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return cm.GetCopilotCredentials()
|
||||
}
|
||||
|
||||
// GetValidCopilotAccessToken returns a fresh GitHub Copilot access token.
|
||||
func GetValidCopilotAccessToken() (string, error) {
|
||||
cm, err := auth.NewCredentialManager()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return cm.GetValidCopilotAccessToken()
|
||||
}
|
||||
|
||||
// GetOpenAIAPIKey resolves the OpenAI API key using the standard
|
||||
// resolution order: stored credentials -> OPENAI_API_KEY env var.
|
||||
// Returns an empty string if no key is found.
|
||||
|
||||
+11
-22
@@ -3,6 +3,8 @@ package kit
|
||||
import (
|
||||
"encoding/json"
|
||||
"sync"
|
||||
|
||||
"github.com/mark3labs/kit/internal/extensions"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -103,34 +105,21 @@ type Event interface {
|
||||
// appropriate visualizations (e.g. diff view for edit tools, command+output
|
||||
// for execute tools) and file trackers to identify which results contain
|
||||
// modifications.
|
||||
//
|
||||
// These constants re-export the canonical classification used by extension
|
||||
// events, so SDK events and extension events always agree.
|
||||
const (
|
||||
ToolKindExecute = "execute" // Shell execution (bash)
|
||||
ToolKindEdit = "edit" // File modification (edit, write)
|
||||
ToolKindRead = "read" // File reading (read, ls)
|
||||
ToolKindSearch = "search" // Content/file search (grep, find)
|
||||
ToolKindSubagent = "agent" // Subagent spawning (subagent)
|
||||
ToolKindExecute = extensions.ToolKindExecute // Shell execution (bash)
|
||||
ToolKindEdit = extensions.ToolKindEdit // File modification (edit, write)
|
||||
ToolKindRead = extensions.ToolKindRead // File reading (read, ls)
|
||||
ToolKindSearch = extensions.ToolKindSearch // Content/file search (grep, find)
|
||||
ToolKindSubagent = extensions.ToolKindSubagent // Subagent spawning (subagent)
|
||||
)
|
||||
|
||||
// coreToolKinds maps built-in tool names to their kind. MCP and extension
|
||||
// tools without an entry default to ToolKindExecute.
|
||||
var coreToolKinds = map[string]string{
|
||||
"bash": ToolKindExecute,
|
||||
"edit": ToolKindEdit,
|
||||
"write": ToolKindEdit,
|
||||
"read": ToolKindRead,
|
||||
"ls": ToolKindRead,
|
||||
"grep": ToolKindSearch,
|
||||
"find": ToolKindSearch,
|
||||
"subagent": ToolKindSubagent,
|
||||
}
|
||||
|
||||
// toolKindFor returns the ToolKind for a given tool name, defaulting to
|
||||
// ToolKindExecute for unknown tools.
|
||||
func toolKindFor(toolName string) string {
|
||||
if kind, ok := coreToolKinds[toolName]; ok {
|
||||
return kind
|
||||
}
|
||||
return ToolKindExecute
|
||||
return extensions.ToolKindFor(toolName)
|
||||
}
|
||||
|
||||
// parseToolArgs attempts to parse a JSON-encoded tool args string into a map.
|
||||
|
||||
@@ -2,6 +2,8 @@ package kit
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"strings"
|
||||
|
||||
"github.com/mark3labs/kit/internal/extensions"
|
||||
"github.com/mark3labs/kit/internal/message"
|
||||
@@ -96,6 +98,23 @@ type ExtensionAPI interface {
|
||||
AppendEntry(extType, data string) (string, error)
|
||||
GetEntries(extType string) []ExtensionEntry
|
||||
|
||||
// Session-scoped extension state (last-write-wins key-value store).
|
||||
// Backed by an in-memory map and (optionally) a sidecar file per session;
|
||||
// state lives outside the conversation tree and is not visible to the LLM.
|
||||
SetState(key, value string)
|
||||
GetState(key string) (string, bool)
|
||||
DeleteState(key string)
|
||||
ListState() []string
|
||||
|
||||
// InitStatePersistence loads any existing state from the per-session
|
||||
// sidecar file and installs a saver hook so that subsequent SetState /
|
||||
// DeleteState mutations are flushed to disk. Safe to call multiple times;
|
||||
// repeat calls simply reload and reinstall the saver.
|
||||
//
|
||||
// For ephemeral or in-memory sessions (no session file path), the call
|
||||
// is a no-op and state remains in memory for the lifetime of the runner.
|
||||
InitStatePersistence() error
|
||||
|
||||
// Status bar
|
||||
SetStatus(entry ExtensionStatusBarEntry)
|
||||
RemoveStatus(key string)
|
||||
@@ -332,6 +351,67 @@ func (e *extensionAPI) AppendEntry(extType, data string) (string, error) {
|
||||
return e.kit.session.AppendExtensionData(extType, data)
|
||||
}
|
||||
|
||||
func (e *extensionAPI) SetState(key, value string) {
|
||||
if e.kit.extRunner != nil {
|
||||
e.kit.extRunner.SetState(key, value)
|
||||
}
|
||||
}
|
||||
|
||||
func (e *extensionAPI) GetState(key string) (string, bool) {
|
||||
if e.kit.extRunner == nil {
|
||||
return "", false
|
||||
}
|
||||
return e.kit.extRunner.GetState(key)
|
||||
}
|
||||
|
||||
func (e *extensionAPI) DeleteState(key string) {
|
||||
if e.kit.extRunner != nil {
|
||||
e.kit.extRunner.DeleteState(key)
|
||||
}
|
||||
}
|
||||
|
||||
func (e *extensionAPI) ListState() []string {
|
||||
if e.kit.extRunner == nil {
|
||||
return nil
|
||||
}
|
||||
return e.kit.extRunner.ListState()
|
||||
}
|
||||
|
||||
func (e *extensionAPI) InitStatePersistence() error {
|
||||
if e.kit.extRunner == nil {
|
||||
return nil
|
||||
}
|
||||
path := extStateSidecarPath(e.kit.GetSessionPath())
|
||||
if path == "" {
|
||||
// Ephemeral or in-memory session; no on-disk state.
|
||||
e.kit.extRunner.SetStateSaver(nil)
|
||||
return nil
|
||||
}
|
||||
if err := e.kit.extRunner.LoadStateFromFile(path); err != nil {
|
||||
return err
|
||||
}
|
||||
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)
|
||||
}
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// extStateSidecarPath returns the path to the per-session extension state
|
||||
// sidecar file derived from the session's JSONL path. Returns empty for
|
||||
// ephemeral / in-memory sessions where no JSONL is being written.
|
||||
func extStateSidecarPath(sessionPath string) string {
|
||||
if sessionPath == "" {
|
||||
return ""
|
||||
}
|
||||
if trimmed, ok := strings.CutSuffix(sessionPath, ".jsonl"); ok {
|
||||
return trimmed + ".ext-state.json"
|
||||
}
|
||||
return sessionPath + ".ext-state.json"
|
||||
}
|
||||
|
||||
func (e *extensionAPI) GetEntries(extType string) []ExtensionEntry {
|
||||
if e.kit.session == nil {
|
||||
return nil
|
||||
|
||||
@@ -3,8 +3,11 @@ package kit
|
||||
import (
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/mark3labs/kit/internal/auth"
|
||||
"github.com/mark3labs/kit/internal/extensions"
|
||||
"github.com/mark3labs/kit/internal/models"
|
||||
)
|
||||
|
||||
// bridgeExtensions registers extension event handlers as SDK hooks and
|
||||
@@ -19,6 +22,30 @@ import (
|
||||
// wrapper (internal/extensions/wrapper.go) which composes underneath the SDK
|
||||
// hook wrapper.
|
||||
func (m *Kit) bridgeExtensions(runner *extensions.Runner) {
|
||||
// Per-turn aggregator: collects tool/LLM/usage signals between AgentStart
|
||||
// and AgentEnd so the enriched AgentEndEvent can be populated without
|
||||
// requiring extensions to maintain parallel bookkeeping.
|
||||
//
|
||||
// NOTE: this aggregator assumes a single in-flight turn per *Kit instance,
|
||||
// which is the current contract — runTurn does not serialize callers and
|
||||
// the SDK's TurnStartEvent/TurnEndEvent do not carry a turn ID, so two
|
||||
// concurrent Prompt() calls on the same *Kit would clobber the counters.
|
||||
// All current callers (TUI app layer, CLI runner, SDK examples) serialize
|
||||
// turns above this layer. If concurrent turns become a supported use case,
|
||||
// extend TurnStartEvent/TurnEndEvent with a turn ID and key this map per
|
||||
// turn instead.
|
||||
turnAgg := &turnAggregator{kit: m}
|
||||
m.Subscribe(func(e Event) {
|
||||
switch ev := e.(type) {
|
||||
case TurnStartEvent:
|
||||
turnAgg.start()
|
||||
case ToolResultEvent:
|
||||
turnAgg.recordTool(ev.ToolName)
|
||||
case StepFinishEvent:
|
||||
turnAgg.recordStep(ev.Usage)
|
||||
}
|
||||
})
|
||||
|
||||
// --- Interception hooks ---
|
||||
|
||||
// Extension Input → BeforeTurn hook (high priority, runs first).
|
||||
@@ -109,9 +136,19 @@ func (m *Kit) bridgeExtensions(runner *extensions.Runner) {
|
||||
} else if stopReason == "" {
|
||||
stopReason = "completed"
|
||||
}
|
||||
agg := turnAgg.consume()
|
||||
_, _ = runner.Emit(extensions.AgentEndEvent{
|
||||
Response: response,
|
||||
StopReason: stopReason,
|
||||
Response: response,
|
||||
StopReason: stopReason,
|
||||
ToolCallCount: agg.toolCallCount,
|
||||
ToolNames: agg.toolNames,
|
||||
LLMCallCount: agg.llmCallCount,
|
||||
InputTokensDelta: agg.inputTokens,
|
||||
OutputTokensDelta: agg.outputTokens,
|
||||
CacheReadTokensDelta: agg.cacheReadTokens,
|
||||
CacheWriteTokensDelta: agg.cacheWriteTokens,
|
||||
CostDelta: agg.cost,
|
||||
DurationMs: agg.durationMs(),
|
||||
})
|
||||
}
|
||||
})
|
||||
@@ -302,6 +339,32 @@ func (m *Kit) bridgeExtensions(runner *extensions.Runner) {
|
||||
}
|
||||
})
|
||||
|
||||
// LLMUsage: derive per-call usage from StepFinish. Each step corresponds
|
||||
// to one LLM provider call, so the step's usage is the per-call delta.
|
||||
// Cost is computed from the current model's pricing (zero when unknown
|
||||
// or OAuth credentials are in use). RequestID is left empty until the
|
||||
// SDK surfaces a correlation id from the underlying provider.
|
||||
if runner.HasHandlers(extensions.LLMUsage) {
|
||||
m.Subscribe(func(e Event) {
|
||||
ev, ok := e.(StepFinishEvent)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
provider, modelID, cost := llmUsageMeta(m, ev.Usage)
|
||||
_, _ = runner.Emit(extensions.LLMUsageEvent{
|
||||
InputTokens: int(ev.Usage.InputTokens),
|
||||
OutputTokens: int(ev.Usage.OutputTokens),
|
||||
CacheReadTokens: int(ev.Usage.CacheReadTokens),
|
||||
CacheWriteTokens: int(ev.Usage.CacheCreationTokens),
|
||||
Cost: cost,
|
||||
Model: modelID,
|
||||
Provider: provider,
|
||||
StepNumber: ev.StepNumber,
|
||||
FinishReason: ev.FinishReason,
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
bridgeObserve(m, runner, extensions.ReasoningStart, func(ev ReasoningStartEvent) extensions.Event {
|
||||
return extensions.ReasoningStartEvent{ID: ev.ID}
|
||||
})
|
||||
@@ -363,6 +426,167 @@ func bridgeObserve[In Event](m *Kit, runner *extensions.Runner, kind extensions.
|
||||
})
|
||||
}
|
||||
|
||||
// turnAggregator collects per-turn signals (tool calls, LLM round-trips, token
|
||||
// usage, wall-clock duration) so that the enriched AgentEndEvent can be
|
||||
// populated without requiring extensions to maintain parallel bookkeeping.
|
||||
//
|
||||
// The aggregator resets on each TurnStartEvent and is consumed (snapshotted +
|
||||
// reset) on TurnEndEvent. All access is serialized via a mutex because the
|
||||
// underlying event bus may fan handlers across goroutines in the future.
|
||||
type turnAggregator struct {
|
||||
mu sync.Mutex
|
||||
started time.Time
|
||||
ended time.Time
|
||||
toolCallCount int
|
||||
toolNames []string
|
||||
llmCallCount int
|
||||
inputTokens int
|
||||
outputTokens int
|
||||
cacheReadTokens int
|
||||
cacheWriteTokens int
|
||||
cost float64
|
||||
kit *Kit
|
||||
}
|
||||
|
||||
type turnSnapshot struct {
|
||||
started time.Time
|
||||
ended time.Time
|
||||
toolCallCount int
|
||||
toolNames []string
|
||||
llmCallCount int
|
||||
inputTokens int
|
||||
outputTokens int
|
||||
cacheReadTokens int
|
||||
cacheWriteTokens int
|
||||
cost float64
|
||||
}
|
||||
|
||||
func (s turnSnapshot) durationMs() int64 {
|
||||
if s.started.IsZero() {
|
||||
return 0
|
||||
}
|
||||
end := s.ended
|
||||
if end.IsZero() {
|
||||
end = time.Now()
|
||||
}
|
||||
return end.Sub(s.started).Milliseconds()
|
||||
}
|
||||
|
||||
// start resets all counters and records the turn's start time. Called from
|
||||
// the TurnStartEvent subscriber.
|
||||
func (a *turnAggregator) start() {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.started = time.Now()
|
||||
a.ended = time.Time{}
|
||||
a.toolCallCount = 0
|
||||
a.toolNames = nil
|
||||
a.llmCallCount = 0
|
||||
a.inputTokens = 0
|
||||
a.outputTokens = 0
|
||||
a.cacheReadTokens = 0
|
||||
a.cacheWriteTokens = 0
|
||||
a.cost = 0
|
||||
}
|
||||
|
||||
func (a *turnAggregator) recordTool(name string) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.toolCallCount++
|
||||
if name != "" {
|
||||
a.toolNames = append(a.toolNames, name)
|
||||
}
|
||||
}
|
||||
|
||||
func (a *turnAggregator) recordStep(usage LLMUsage) {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.llmCallCount++
|
||||
a.inputTokens += int(usage.InputTokens)
|
||||
a.outputTokens += int(usage.OutputTokens)
|
||||
a.cacheReadTokens += int(usage.CacheReadTokens)
|
||||
a.cacheWriteTokens += int(usage.CacheCreationTokens)
|
||||
if a.kit != nil {
|
||||
_, _, c := llmUsageMeta(a.kit, usage)
|
||||
a.cost += c
|
||||
}
|
||||
}
|
||||
|
||||
// consume returns a snapshot of the current turn and marks it ended.
|
||||
// Subsequent start() calls clear the snapshot.
|
||||
func (a *turnAggregator) consume() turnSnapshot {
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
a.ended = time.Now()
|
||||
names := a.toolNames
|
||||
if len(names) > 0 {
|
||||
copied := make([]string, len(names))
|
||||
copy(copied, names)
|
||||
names = copied
|
||||
}
|
||||
return turnSnapshot{
|
||||
started: a.started,
|
||||
ended: a.ended,
|
||||
toolCallCount: a.toolCallCount,
|
||||
toolNames: names,
|
||||
llmCallCount: a.llmCallCount,
|
||||
inputTokens: a.inputTokens,
|
||||
outputTokens: a.outputTokens,
|
||||
cacheReadTokens: a.cacheReadTokens,
|
||||
cacheWriteTokens: a.cacheWriteTokens,
|
||||
cost: a.cost,
|
||||
}
|
||||
}
|
||||
|
||||
// llmUsageMeta returns the current provider, model id, and computed cost for
|
||||
// the given usage values using the Kit instance's active model. Cost is zero
|
||||
// in any of the following cases:
|
||||
// - the *Kit pointer is nil or has no active model;
|
||||
// - the model is not in the registry (custom fine-tunes, unknown providers);
|
||||
// - the model has no pricing fields set;
|
||||
// - the active credential is an Anthropic OAuth token (matches the
|
||||
// existing usage_tracker behavior of suppressing cost for OAuth users).
|
||||
func llmUsageMeta(m *Kit, usage LLMUsage) (provider, modelID string, cost float64) {
|
||||
if m == nil {
|
||||
return "", "", 0
|
||||
}
|
||||
modelString := m.GetModelString()
|
||||
if modelString == "" {
|
||||
return "", "", 0
|
||||
}
|
||||
p, id, err := models.ParseModelString(modelString)
|
||||
if err != nil {
|
||||
return "", "", 0
|
||||
}
|
||||
provider, modelID = p, id
|
||||
info := models.GetGlobalRegistry().LookupModel(provider, modelID)
|
||||
if info == nil {
|
||||
return provider, modelID, 0
|
||||
}
|
||||
if isAnthropicOAuth(m, provider) {
|
||||
return provider, modelID, 0
|
||||
}
|
||||
cost = float64(usage.InputTokens) * info.Cost.Input / 1_000_000
|
||||
cost += float64(usage.OutputTokens) * info.Cost.Output / 1_000_000
|
||||
if info.Cost.CacheRead != nil {
|
||||
cost += float64(usage.CacheReadTokens) * (*info.Cost.CacheRead) / 1_000_000
|
||||
}
|
||||
if info.Cost.CacheWrite != nil {
|
||||
cost += float64(usage.CacheCreationTokens) * (*info.Cost.CacheWrite) / 1_000_000
|
||||
}
|
||||
return provider, modelID, cost
|
||||
}
|
||||
|
||||
// isAnthropicOAuth reports whether the current Anthropic credential resolves
|
||||
// to a stored OAuth token (in which case the user is not billed per-token),
|
||||
// so OnLLMUsage cost reporting agrees with ctx.GetSessionUsage().
|
||||
func isAnthropicOAuth(m *Kit, provider string) bool {
|
||||
if m == nil || provider != "anthropic" {
|
||||
return false
|
||||
}
|
||||
return auth.IsAnthropicOAuth(m.v.GetString("provider-api-key"))
|
||||
}
|
||||
|
||||
// llmToContextMessages converts a slice of LLM messages to extension
|
||||
// ContextMessage values, extracting plain text from each message.
|
||||
func llmToContextMessages(msgs []LLMMessage) []extensions.ContextMessage {
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
package kit
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestTurnAggregator_BasicLifecycle exercises the per-turn aggregator:
|
||||
// start → record several tools and steps → consume → snapshot should reflect
|
||||
// the accumulated counts and zero out for the next turn.
|
||||
func TestTurnAggregator_BasicLifecycle(t *testing.T) {
|
||||
agg := &turnAggregator{}
|
||||
|
||||
agg.start()
|
||||
agg.recordTool("bash")
|
||||
agg.recordTool("read")
|
||||
agg.recordTool("bash")
|
||||
agg.recordStep(LLMUsage{
|
||||
InputTokens: 100,
|
||||
OutputTokens: 50,
|
||||
CacheReadTokens: 10,
|
||||
CacheCreationTokens: 5,
|
||||
})
|
||||
agg.recordStep(LLMUsage{
|
||||
InputTokens: 200,
|
||||
OutputTokens: 75,
|
||||
})
|
||||
|
||||
snap := agg.consume()
|
||||
if snap.toolCallCount != 3 {
|
||||
t.Errorf("toolCallCount: got %d want 3", snap.toolCallCount)
|
||||
}
|
||||
wantNames := []string{"bash", "read", "bash"}
|
||||
if len(snap.toolNames) != len(wantNames) {
|
||||
t.Fatalf("toolNames length: got %d want %d", len(snap.toolNames), len(wantNames))
|
||||
}
|
||||
for i, n := range wantNames {
|
||||
if snap.toolNames[i] != n {
|
||||
t.Errorf("toolNames[%d]: got %q want %q", i, snap.toolNames[i], n)
|
||||
}
|
||||
}
|
||||
if snap.llmCallCount != 2 {
|
||||
t.Errorf("llmCallCount: got %d want 2", snap.llmCallCount)
|
||||
}
|
||||
if snap.inputTokens != 300 {
|
||||
t.Errorf("inputTokens: got %d want 300", snap.inputTokens)
|
||||
}
|
||||
if snap.outputTokens != 125 {
|
||||
t.Errorf("outputTokens: got %d want 125", snap.outputTokens)
|
||||
}
|
||||
if snap.cacheReadTokens != 10 {
|
||||
t.Errorf("cacheReadTokens: got %d want 10", snap.cacheReadTokens)
|
||||
}
|
||||
if snap.cacheWriteTokens != 5 {
|
||||
t.Errorf("cacheWriteTokens: got %d want 5", snap.cacheWriteTokens)
|
||||
}
|
||||
if snap.durationMs() < 0 {
|
||||
t.Errorf("durationMs should not be negative, got %d", snap.durationMs())
|
||||
}
|
||||
}
|
||||
|
||||
func TestTurnAggregator_StartResetsCounters(t *testing.T) {
|
||||
agg := &turnAggregator{}
|
||||
agg.start()
|
||||
agg.recordTool("bash")
|
||||
agg.recordStep(LLMUsage{InputTokens: 50})
|
||||
|
||||
// Begin a new turn — previous counters should be cleared.
|
||||
agg.start()
|
||||
snap := agg.consume()
|
||||
|
||||
if snap.toolCallCount != 0 || snap.llmCallCount != 0 || snap.inputTokens != 0 {
|
||||
t.Errorf("expected counters zeroed after start(), got %+v", snap)
|
||||
}
|
||||
if snap.toolNames != nil {
|
||||
t.Errorf("expected toolNames=nil after start(), got %v", snap.toolNames)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTurnAggregator_DurationMs verifies the snapshot computes a positive
|
||||
// duration when consume() runs after start().
|
||||
func TestTurnAggregator_DurationMs(t *testing.T) {
|
||||
agg := &turnAggregator{}
|
||||
agg.start()
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
snap := agg.consume()
|
||||
if snap.durationMs() < 1 {
|
||||
t.Errorf("expected positive duration, got %d", snap.durationMs())
|
||||
}
|
||||
}
|
||||
|
||||
// TestTurnAggregator_ZeroStartSafe ensures a snapshot taken without a prior
|
||||
// start() doesn't crash and reports zero duration.
|
||||
func TestTurnAggregator_ZeroStartSafe(t *testing.T) {
|
||||
agg := &turnAggregator{}
|
||||
snap := agg.consume()
|
||||
if snap.durationMs() != 0 {
|
||||
t.Errorf("expected zero duration for unstarted aggregator, got %d", snap.durationMs())
|
||||
}
|
||||
}
|
||||
|
||||
// TestLLMUsageMeta_NilKit verifies the helper degrades gracefully when given
|
||||
// a nil Kit instance (zero values, no panic).
|
||||
func TestLLMUsageMeta_NilKit(t *testing.T) {
|
||||
provider, modelID, cost := llmUsageMeta(nil, LLMUsage{InputTokens: 100})
|
||||
if provider != "" || modelID != "" || cost != 0 {
|
||||
t.Errorf("expected zero values for nil kit, got (%q,%q,%v)", provider, modelID, cost)
|
||||
}
|
||||
}
|
||||
|
||||
// TestIsAnthropicOAuth_NonAnthropic verifies the helper short-circuits for any
|
||||
// provider other than "anthropic" without touching the credential store.
|
||||
func TestIsAnthropicOAuth_NonAnthropic(t *testing.T) {
|
||||
for _, provider := range []string{"openai", "google", "openrouter", ""} {
|
||||
if isAnthropicOAuth(nil, provider) {
|
||||
t.Errorf("isAnthropicOAuth(nil, %q) = true, want false", provider)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtStateSidecarPath(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
in string
|
||||
want string
|
||||
}{
|
||||
{"empty", "", ""},
|
||||
{"jsonl", "/tmp/sessions/abc.jsonl", "/tmp/sessions/abc.ext-state.json"},
|
||||
{"jsonl with subdir", "/a/b/c.jsonl", "/a/b/c.ext-state.json"},
|
||||
{"no extension", "/tmp/session-blob", "/tmp/session-blob.ext-state.json"},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := extStateSidecarPath(tc.in)
|
||||
if got != tc.want {
|
||||
t.Errorf("extStateSidecarPath(%q): got %q want %q", tc.in, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+23
-3
@@ -138,6 +138,19 @@ func (m *Kit) GetToolNames() []string {
|
||||
return names
|
||||
}
|
||||
|
||||
// GetToolsForSubagent like GetTools but eliminates subagent tool
|
||||
// to avoid infinite recursion.
|
||||
func (m *Kit) GetToolsForSubagent() []Tool {
|
||||
var tools []Tool
|
||||
for _, t := range m.agent.GetTools() {
|
||||
if t.Info().Name == "subagent" {
|
||||
continue
|
||||
}
|
||||
tools = append(tools, t)
|
||||
}
|
||||
return tools
|
||||
}
|
||||
|
||||
// GetLoadingMessage returns the agent's startup info message (e.g. GPU
|
||||
// fallback info), or empty string if none.
|
||||
func (m *Kit) GetLoadingMessage() string {
|
||||
@@ -1147,7 +1160,7 @@ type CLIOptions struct {
|
||||
// - Continue: resume most recent session for SessionDir (or cwd)
|
||||
// - SessionPath: open a specific JSONL session file
|
||||
// - default: create a new tree session for SessionDir (or cwd)
|
||||
func InitTreeSession(opts *Options) (*session.TreeManager, error) {
|
||||
func InitTreeSession(opts *Options) (*TreeManager, error) {
|
||||
if opts == nil {
|
||||
opts = &Options{}
|
||||
}
|
||||
@@ -1814,8 +1827,14 @@ type SubagentConfig struct {
|
||||
// Empty string uses a minimal default prompt.
|
||||
SystemPrompt string
|
||||
|
||||
// Tools overrides the tool set. If nil, SubagentTools() is used (all
|
||||
// core tools except subagent, preventing infinite recursion).
|
||||
// Tools overrides the tool set available to the subagent.
|
||||
// If nil and the subagent is created via the SDK (Kit.Subagent()), the
|
||||
// static SubagentTools() set (all core tools except "subagent") is used.
|
||||
// When spawned internally by the agent loop, the parent's active tools
|
||||
// minus "subagent" are used instead (see GetToolsForSubagent()).
|
||||
// Pass m.GetToolsForSubagent() explicitly to opt into inheritance from
|
||||
// SDK call sites.
|
||||
// (The subagent tool is dropped to prevent infinite recursion.)
|
||||
Tools []Tool
|
||||
|
||||
// NoSession, when true, uses an in-memory ephemeral session. When false
|
||||
@@ -2076,6 +2095,7 @@ func (m *Kit) generate(ctx context.Context, messages []fantasy.Message) (*agent.
|
||||
SystemPrompt: systemPrompt,
|
||||
Timeout: timeout,
|
||||
OnEvent: onEvent,
|
||||
Tools: m.GetToolsForSubagent(),
|
||||
})
|
||||
m.cleanupSubagentListeners(toolCallID)
|
||||
if result == nil {
|
||||
|
||||
+27
-71
@@ -7,45 +7,36 @@ import (
|
||||
|
||||
"github.com/mark3labs/kit/internal/extensions"
|
||||
"github.com/mark3labs/kit/internal/models"
|
||||
"github.com/mark3labs/kit/internal/prompts"
|
||||
"github.com/mark3labs/kit/internal/skills"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Template Parsing Bridge for Extensions (Phase 3)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// varRegex matches {{variable}} placeholders in templates.
|
||||
var varRegex = regexp.MustCompile(`\{\{\s*([a-zA-Z_][a-zA-Z0-9_]*)\s*\}\}`)
|
||||
|
||||
// ParseTemplate extracts {{variables}} from template content.
|
||||
// ParseTemplate extracts {{variables}} from template content. The template
|
||||
// grammar is shared with skill prompt templates, so a template parses
|
||||
// identically regardless of which API loads it.
|
||||
func ParseTemplate(name, content string) extensions.PromptTemplate {
|
||||
matches := varRegex.FindAllStringSubmatch(content, -1)
|
||||
vars := make([]string, 0, len(matches))
|
||||
seen := make(map[string]bool)
|
||||
for _, m := range matches {
|
||||
if len(m) > 1 && !seen[m[1]] {
|
||||
seen[m[1]] = true
|
||||
vars = append(vars, m[1])
|
||||
}
|
||||
tpl := skills.NewPromptTemplate(name, content)
|
||||
vars := tpl.Variables
|
||||
if vars == nil {
|
||||
vars = []string{}
|
||||
}
|
||||
return extensions.PromptTemplate{
|
||||
Name: name,
|
||||
Content: content,
|
||||
Name: tpl.Name,
|
||||
Content: tpl.Content,
|
||||
Variables: vars,
|
||||
}
|
||||
}
|
||||
|
||||
// RenderTemplate substitutes variables into template content.
|
||||
// Handles {{name}} and {{ name }} (any whitespace) placeholders.
|
||||
// Handles {{name}} and {{ name }} (any whitespace) placeholders; missing
|
||||
// variables are left as-is.
|
||||
func RenderTemplate(tpl extensions.PromptTemplate, vars map[string]string) string {
|
||||
return varRegex.ReplaceAllStringFunc(tpl.Content, func(m string) string {
|
||||
sub := varRegex.FindStringSubmatch(m)
|
||||
if len(sub) > 1 {
|
||||
if v, ok := vars[sub[1]]; ok {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return m
|
||||
})
|
||||
t := skills.PromptTemplate{Content: tpl.Content}
|
||||
return t.Expand(vars)
|
||||
}
|
||||
|
||||
// ParseArguments parses command-line style arguments.
|
||||
@@ -183,44 +174,12 @@ func SimpleParseArguments(input string, count int) []string {
|
||||
return result
|
||||
}
|
||||
|
||||
// parseFields splits input respecting quoted strings.
|
||||
// parseFields splits input into arguments respecting quoted strings and
|
||||
// backslash escaping. It delegates to the canonical tokenizer in
|
||||
// internal/prompts so extension argument parsing and builtin prompt-template
|
||||
// parsing agree on grammar.
|
||||
func parseFields(input string) []string {
|
||||
var fields []string
|
||||
var current strings.Builder
|
||||
inQuote := false
|
||||
quoteChar := rune(0)
|
||||
|
||||
for _, r := range input {
|
||||
switch r {
|
||||
case '"', '\'':
|
||||
if !inQuote {
|
||||
inQuote = true
|
||||
quoteChar = r
|
||||
} else if r == quoteChar {
|
||||
inQuote = false
|
||||
quoteChar = 0
|
||||
} else {
|
||||
current.WriteRune(r)
|
||||
}
|
||||
case ' ', '\t':
|
||||
if inQuote {
|
||||
current.WriteRune(r)
|
||||
} else {
|
||||
if current.Len() > 0 {
|
||||
fields = append(fields, current.String())
|
||||
current.Reset()
|
||||
}
|
||||
}
|
||||
default:
|
||||
current.WriteRune(r)
|
||||
}
|
||||
}
|
||||
|
||||
if current.Len() > 0 {
|
||||
fields = append(fields, current.String())
|
||||
}
|
||||
|
||||
return fields
|
||||
return prompts.ParseCommandArgs(input)
|
||||
}
|
||||
|
||||
// EvaluateModelConditional checks if condition matches current model.
|
||||
@@ -417,21 +376,18 @@ func MatchModelGlob(model, pattern string) bool {
|
||||
}
|
||||
|
||||
// ExtractProviderFromPath extracts provider from a path-like model string.
|
||||
//
|
||||
// Deprecated: Use GetCurrentProvider instead.
|
||||
func ExtractProviderFromPath(model string) string {
|
||||
parts := strings.Split(model, "/")
|
||||
if len(parts) >= 2 {
|
||||
return parts[0]
|
||||
}
|
||||
return ""
|
||||
return GetCurrentProvider(model)
|
||||
}
|
||||
|
||||
// ExtractModelFromPath extracts model ID from a path-like model string.
|
||||
//
|
||||
// Deprecated: Use RemoveProviderFromModel instead, which correctly handles
|
||||
// model IDs containing "/" (e.g. "openrouter/meta/llama").
|
||||
func ExtractModelFromPath(model string) string {
|
||||
parts := strings.Split(model, "/")
|
||||
if len(parts) >= 2 {
|
||||
return parts[1]
|
||||
}
|
||||
return model
|
||||
return RemoveProviderFromModel(model)
|
||||
}
|
||||
|
||||
// IsBareModelID checks if a string is a bare model ID (no provider).
|
||||
|
||||
@@ -88,7 +88,8 @@ api.OnAgentStart(func(e ext.AgentStartEvent, ctx ext.Context) {
|
||||
// e.Prompt string
|
||||
})
|
||||
|
||||
// Agent finished responding.
|
||||
// Agent finished responding. Carries per-turn aggregates so observer-style
|
||||
// extensions don't need to maintain parallel bookkeeping.
|
||||
api.OnAgentEnd(func(e ext.AgentEndEvent, ctx ext.Context) {
|
||||
// e.Response string
|
||||
// e.StopReason string — "error" (on failure), "completed" (when LLM returns
|
||||
@@ -96,6 +97,33 @@ api.OnAgentEnd(func(e ext.AgentEndEvent, ctx ext.Context) {
|
||||
// (e.g. "stop", "length" (max output tokens hit), "tool-calls", "content-filter").
|
||||
// To detect errors, check e.StopReason == "error".
|
||||
// Do NOT compare against "completed" for success — instead check != "error".
|
||||
//
|
||||
// Per-turn aggregates (computed by Kit's runtime):
|
||||
// e.ToolCallCount int — total tool invocations this turn
|
||||
// e.ToolNames []string — tool names in call order (duplicates preserved)
|
||||
// e.LLMCallCount int — LLM round-trips / tool-loop iterations
|
||||
// e.InputTokensDelta int — sum of input tokens across LLM calls this turn
|
||||
// e.OutputTokensDelta int
|
||||
// e.CacheReadTokensDelta int
|
||||
// e.CacheWriteTokensDelta int
|
||||
// e.CostDelta float64 — USD cost (zero when pricing unknown / OAuth)
|
||||
// e.DurationMs int64 — wall-clock duration AgentStart→AgentEnd
|
||||
})
|
||||
|
||||
// Per-LLM-call usage — fires after each provider round-trip with token + cost
|
||||
// deltas attributed to that specific call. A single turn typically produces
|
||||
// multiple LLMUsageEvents (one per tool-loop iteration). Use this for accurate
|
||||
// budget enforcement that needs to react between calls instead of waiting
|
||||
// for the turn to finish.
|
||||
api.OnLLMUsage(func(e ext.LLMUsageEvent, ctx ext.Context) {
|
||||
// e.InputTokens, e.OutputTokens int
|
||||
// e.CacheReadTokens, e.CacheWriteTokens int
|
||||
// e.Cost float64 — USD; zero when pricing unknown / OAuth
|
||||
// e.Model, e.Provider string — model used for THIS call
|
||||
// (may differ across calls if SetModel was called)
|
||||
// e.StepNumber int — zero-based step index in this turn
|
||||
// e.FinishReason string — "stop" / "tool_calls" / "length" / ...
|
||||
// e.RequestID string — optional provider correlation id (may be empty)
|
||||
})
|
||||
```
|
||||
|
||||
@@ -528,11 +556,38 @@ stats := ctx.GetContextStats() // .EstimatedTokens, .ContextLimit, .UsagePer
|
||||
msgs := ctx.GetMessages() // []ext.SessionMessage on current branch
|
||||
path := ctx.GetSessionPath() // file path of session JSONL
|
||||
|
||||
// Persist custom data in the session tree:
|
||||
// Append-only log in the session tree (fork-aware, walked on every branch read):
|
||||
id, err := ctx.AppendEntry("my-type", "data string")
|
||||
entries := ctx.GetEntries("my-type") // []ext.ExtensionEntry{ID, EntryType, Data, Timestamp}
|
||||
```
|
||||
|
||||
### Session State (last-write-wins)
|
||||
|
||||
Key-value store scoped to the session, persisted to a sidecar file
|
||||
(`<session>.ext-state.json`) outside the conversation tree. Reads are O(1)
|
||||
(no branch walk), writes don't grow the JSONL, and the store is not
|
||||
duplicated on fork. State is invisible to the LLM and survives session
|
||||
resume. For ephemeral / in-memory sessions, state lives only in memory.
|
||||
|
||||
```go
|
||||
ctx.SetState("myext:budget-cap", "10.00") // last write wins
|
||||
val, ok := ctx.GetState("myext:budget-cap") // (string, bool)
|
||||
ctx.DeleteState("myext:budget-cap") // no-op if missing
|
||||
keys := ctx.ListState() // []string, unspecified order
|
||||
```
|
||||
|
||||
**When to use which:**
|
||||
|
||||
| Need | Use |
|
||||
|------|-----|
|
||||
| Snapshot state ("current value of X") | `SetState` / `GetState` |
|
||||
| Audit log / event history | `AppendEntry` / `GetEntries` |
|
||||
| One-shot per-turn signal | enriched `AgentEndEvent` fields |
|
||||
| Per-LLM-call observation | `OnLLMUsage` event |
|
||||
|
||||
Namespace keys with your extension name (e.g. `"myext:budget-cap"`) to avoid
|
||||
collisions across extensions.
|
||||
|
||||
### Model Management
|
||||
|
||||
```go
|
||||
|
||||
@@ -1104,6 +1104,19 @@ if extAPI.HasExtensions() {
|
||||
tools := extAPI.GetToolInfos()
|
||||
extAPI.SetActiveTools([]string{"bash", "read"})
|
||||
|
||||
// Session-scoped extension state (last-write-wins key-value store).
|
||||
// Backed by an in-memory map and a per-session sidecar file
|
||||
// (<session>.ext-state.json) outside the conversation tree.
|
||||
extAPI.SetState("myext:budget-cap", "10.00")
|
||||
val, ok := extAPI.GetState("myext:budget-cap")
|
||||
extAPI.DeleteState("myext:budget-cap")
|
||||
keys := extAPI.ListState()
|
||||
|
||||
// Load any existing state from the sidecar and install a saver hook so
|
||||
// subsequent SetState/DeleteState mutations are flushed atomically.
|
||||
// No-op for ephemeral / in-memory sessions. Safe to call multiple times.
|
||||
_ = extAPI.InitStatePersistence()
|
||||
|
||||
// Events
|
||||
extAPI.EmitSessionStart()
|
||||
extAPI.EmitModelChange("new/model", "old/model", "extension")
|
||||
|
||||
@@ -7,7 +7,7 @@ description: All extension capabilities — lifecycle events, tools, commands, w
|
||||
|
||||
## Lifecycle events
|
||||
|
||||
Extensions can hook into 26 lifecycle events:
|
||||
Extensions can hook into 27 lifecycle events:
|
||||
|
||||
| Event | Description |
|
||||
|-------|-------------|
|
||||
@@ -15,7 +15,8 @@ Extensions can hook into 26 lifecycle events:
|
||||
| `OnSessionShutdown` | Session ending |
|
||||
| `OnBeforeAgentStart` | Before the agent loop begins |
|
||||
| `OnAgentStart` | Agent loop started |
|
||||
| `OnAgentEnd` | Agent loop completed |
|
||||
| `OnAgentEnd` | Agent loop completed (carries per-turn aggregates: tool counts, token deltas, cost, duration) |
|
||||
| `OnLLMUsage` | Per-LLM-call token + cost delta (fires once per provider round-trip) |
|
||||
| `OnToolCall` | Tool call requested by the model |
|
||||
| `OnToolCallInputStart` | LLM began generating tool call arguments (tool name known, args streaming) |
|
||||
| `OnToolCallInputDelta` | Streamed JSON fragment of tool call arguments |
|
||||
@@ -45,11 +46,52 @@ api.OnToolCall(func(event ext.ToolCallEvent, ctx ext.Context) {
|
||||
ctx.PrintInfo("Calling tool: " + event.Name)
|
||||
})
|
||||
|
||||
api.OnAgentEnd(func(_ ext.AgentEndEvent, ctx ext.Context) {
|
||||
ctx.PrintInfo("Agent finished")
|
||||
api.OnAgentEnd(func(e ext.AgentEndEvent, ctx ext.Context) {
|
||||
// Per-turn aggregates populated by Kit's runtime — no parallel
|
||||
// bookkeeping required in the handler.
|
||||
ctx.PrintInfo(fmt.Sprintf(
|
||||
"Turn finished: %d tool calls (%v), %d LLM round-trips, $%.4f, %dms",
|
||||
e.ToolCallCount, e.ToolNames, e.LLMCallCount, e.CostDelta, e.DurationMs,
|
||||
))
|
||||
})
|
||||
|
||||
// Per-LLM-call usage — fires multiple times per turn (once per round-trip).
|
||||
// Use for accurate budget enforcement between calls.
|
||||
api.OnLLMUsage(func(e ext.LLMUsageEvent, ctx ext.Context) {
|
||||
ctx.PrintInfo(fmt.Sprintf(
|
||||
"%s/%s step=%d tokens=↑%d ↓%d cost=$%.4f (%s)",
|
||||
e.Provider, e.Model, e.StepNumber,
|
||||
e.InputTokens, e.OutputTokens, e.Cost, e.FinishReason,
|
||||
))
|
||||
})
|
||||
```
|
||||
|
||||
**`AgentEndEvent` fields** (in addition to `Response` and `StopReason`):
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `ToolCallCount` | `int` | Total tool invocations during the turn |
|
||||
| `ToolNames` | `[]string` | Tool names in call order (duplicates preserved) |
|
||||
| `LLMCallCount` | `int` | LLM round-trips / tool-loop iterations |
|
||||
| `InputTokensDelta` | `int` | Sum of input tokens across all LLM calls this turn |
|
||||
| `OutputTokensDelta` | `int` | Sum of output tokens across all LLM calls this turn |
|
||||
| `CacheReadTokensDelta` | `int` | Sum of cache-read tokens this turn |
|
||||
| `CacheWriteTokensDelta` | `int` | Sum of cache-write tokens this turn |
|
||||
| `CostDelta` | `float64` | Cost in USD (zero when pricing is unknown or OAuth credentials) |
|
||||
| `DurationMs` | `int64` | Wall-clock time from `AgentStart` to `AgentEnd` |
|
||||
|
||||
**`LLMUsageEvent` fields**:
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `InputTokens` / `OutputTokens` | `int` | Per-call token deltas |
|
||||
| `CacheReadTokens` / `CacheWriteTokens` | `int` | Per-call cache token deltas |
|
||||
| `Cost` | `float64` | Per-call USD cost (zero when pricing unknown) |
|
||||
| `Model` / `Provider` | `string` | Model used for this specific call — may differ from earlier calls if `ctx.SetModel` was called mid-turn |
|
||||
| `StepNumber` | `int` | Zero-based step index within the turn |
|
||||
| `FinishReason` | `string` | Provider finish reason for this call (`"stop"`, `"tool_calls"`, `"length"`, ...) |
|
||||
| `RequestID` | `string` | Optional provider correlation id (may be empty) |
|
||||
|
||||
## Tools
|
||||
|
||||
Register custom tools that the LLM can invoke:
|
||||
@@ -338,6 +380,36 @@ api.OnCustomEvent("my-extension:data-ready", func(data any, ctx ext.Context) {
|
||||
})
|
||||
```
|
||||
|
||||
## Session state
|
||||
|
||||
Last-write-wins key-value store, scoped to the current session and persisted to a sidecar file (`<session>.ext-state.json`) outside the conversation tree:
|
||||
|
||||
```go
|
||||
ctx.SetState("myext:budget-cap", "10.00")
|
||||
|
||||
if cap, ok := ctx.GetState("myext:budget-cap"); ok {
|
||||
// ...
|
||||
}
|
||||
|
||||
ctx.DeleteState("myext:budget-cap")
|
||||
keys := ctx.ListState() // []string, unspecified order
|
||||
```
|
||||
|
||||
Reads are O(1) (no branch walk), writes don't grow the session JSONL, and the store is not duplicated when the conversation forks. State is invisible to the LLM and survives session resume.
|
||||
|
||||
### When to use which persistence primitive
|
||||
|
||||
| Need | Use | Why |
|
||||
|------|-----|-----|
|
||||
| Snapshot state ("current value of X") | `SetState` / `GetState` | O(1) reads, sidecar file, last-write-wins |
|
||||
| Audit log / event history | `AppendEntry` / `GetEntries` | Append-only, lives in conversation tree, fork-aware |
|
||||
| One-shot per-turn signal | Enriched `AgentEndEvent` fields | No persistence needed; runtime tracks it for you |
|
||||
| Per-LLM-call observation | `OnLLMUsage` event | Already attributed to model/provider/step |
|
||||
|
||||
Using `AppendEntry` for snapshot state has a cost: it's O(branch_length) to read, fsyncs into the JSONL on every write, and the entry list duplicates on every fork. Prefer `SetState` for "what's the current value of X?"-style data.
|
||||
|
||||
For ephemeral / in-memory sessions (no JSONL path) the state lives only in memory for the lifetime of the runner.
|
||||
|
||||
## Bridged SDK APIs
|
||||
|
||||
Extensions can access powerful internal SDK capabilities that enable advanced features like conversation tree navigation, dynamic skill loading, template parsing, and model resolution.
|
||||
|
||||
@@ -50,6 +50,7 @@ Kit ships with a rich set of example extensions in the `examples/extensions/` di
|
||||
| [`context-inject.go`](https://github.com/mark3labs/kit/blob/master/examples/extensions/context-inject.go) | Inject context into conversations |
|
||||
| [`summarize.go`](https://github.com/mark3labs/kit/blob/master/examples/extensions/summarize.go) | Conversation summarization |
|
||||
| [`lsp-diagnostics.go`](https://github.com/mark3labs/kit/blob/master/examples/extensions/lsp-diagnostics.go) | LSP diagnostic integration |
|
||||
| [`usage-budget.go`](https://github.com/mark3labs/kit/blob/master/examples/extensions/usage-budget.go) | Per-call usage callback (`OnLLMUsage`), session state (`SetState`/`GetState`), and enriched `OnAgentEnd` per-turn report |
|
||||
|
||||
## Bridged SDK APIs
|
||||
|
||||
|
||||
@@ -65,7 +65,8 @@ Passed to event handlers, the `Context` object provides runtime access to Kit's
|
||||
- **Model** — `ctx.SetModel(...)`, `ctx.GetAvailableModels()`
|
||||
- **Tools** — `ctx.GetAllTools()`, `ctx.SetActiveTools(...)`
|
||||
- **Context stats** — `ctx.GetContextStats()`
|
||||
- **Session data** — `ctx.AppendEntry(...)`, `ctx.GetEntries(...)`
|
||||
- **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(...)`
|
||||
|
||||
+19
-3
@@ -13,6 +13,7 @@ Kit supports a wide range of LLM providers through a unified `provider/model` st
|
||||
|----------|--------|-------------|
|
||||
| **Anthropic** | `anthropic/` | Claude models (native, prompt caching, OAuth) |
|
||||
| **OpenAI** | `openai/` | GPT models |
|
||||
| **GitHub Copilot** | `copilot/` | Copilot models through GitHub device login (experimental) |
|
||||
| **Google** | `google/` or `gemini/` | Gemini models |
|
||||
| **Ollama** | `ollama/` | Local models |
|
||||
| **Azure OpenAI** | `azure/` | Azure-hosted OpenAI |
|
||||
@@ -29,6 +30,7 @@ Kit supports a wide range of LLM providers through a unified `provider/model` st
|
||||
provider/model # Standard format
|
||||
anthropic/claude-sonnet-latest
|
||||
openai/gpt-4o
|
||||
copilot/gpt-5.5
|
||||
ollama/llama3
|
||||
google/gemini-2.5-flash
|
||||
```
|
||||
@@ -117,14 +119,19 @@ kit --provider-api-key "sk-..." --model openai/gpt-4o
|
||||
|
||||
### OAuth
|
||||
|
||||
For providers that support OAuth (e.g., Anthropic):
|
||||
For providers that support OAuth:
|
||||
|
||||
```bash
|
||||
kit auth login anthropic # Start OAuth flow
|
||||
kit auth login anthropic # Anthropic OAuth
|
||||
kit auth login openai # ChatGPT/Codex OAuth
|
||||
kit auth login copilot # GitHub Copilot device login (experimental)
|
||||
kit auth status # Check authentication status
|
||||
kit auth logout anthropic # Remove credentials
|
||||
kit auth logout copilot # Remove credentials
|
||||
```
|
||||
|
||||
The experimental `copilot/` provider requires an active GitHub Copilot subscription
|
||||
and uses GitHub device login; no OpenAI account or OpenAI API key is required.
|
||||
|
||||
### Custom provider URL
|
||||
|
||||
For self-hosted or proxy endpoints:
|
||||
@@ -133,6 +140,15 @@ For self-hosted or proxy endpoints:
|
||||
kit --provider-url "https://my-proxy.example.com/v1" --model openai/gpt-4o
|
||||
```
|
||||
|
||||
When `--provider-url` is set with an explicit `--model`, Kit routes through the
|
||||
`custom` (OpenAI-compatible) wire and strips any provider prefix from the model
|
||||
name. So `openai/gpt-4o`, `google/gemma-4-12b`, and bare `gpt-4o` all resolve
|
||||
to the same endpoint — Kit treats `--provider-url` as authoritative about *where*
|
||||
to send the request, and the model string as just the upstream model id.
|
||||
|
||||
This avoids name collisions when a local server (LM Studio, Ollama, vLLM, ...)
|
||||
happens to expose a model whose name matches a known cloud provider.
|
||||
|
||||
When `--provider-url` is provided without `--model`, Kit automatically defaults to `custom/custom`:
|
||||
|
||||
```bash
|
||||
|
||||
Reference in New Issue
Block a user