WIP: feat(extensions): support custom model providers

.golangci.yml

@@ -1,6 +1,12 @@
 version: "2"
 
 linters:
+  exclusions:
+    rules:
+      # Keep the third-party interpreter diff limited to the concurrency fix.
+      # Kit-owned integration and regression tests remain fully linted.
+      - path: internal/yaegi/
+        linters: [errcheck, govet, modernize, staticcheck, unused]
   enable:
     - modernize
     # staticcheck and unused catch bugs and dead code that went undetected

README.md

@@ -13,6 +13,10 @@
 
 # KIT (Knowledge Inference Tool)
 
+## Extension provider guide
+
+See [callback-based model providers](www/pages/extensions/providers.md) for the offline examples and current implementation status.
+
 ## SDK guides
 
 - [Decision API](https://go-kit.dev/sdk/decisions): typed probabilistic questions, not chat.

cmd/models.go

@@ -4,12 +4,16 @@ import (
 	"fmt"
 	"sort"
 
+	"github.com/mark3labs/kit/internal/extensions"
+	"github.com/mark3labs/kit/internal/kitsetup"
 	"github.com/mark3labs/kit/internal/models"
 	kit "github.com/mark3labs/kit/pkg/kit"
 	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
 )
 
 var modelsAllFlag bool
+var modelsExtensionsFlag bool
 
 var modelsCmd = &cobra.Command{
 	Use:   "models [provider]",
@@ -36,10 +40,41 @@ Examples:
 
 func init() {
 	modelsCmd.Flags().BoolVar(&modelsAllFlag, "all", false, "show all providers in the database, not just fantasy-compatible ones")
+	modelsCmd.Flags().BoolVar(&modelsExtensionsFlag, "extensions", false, "execute trusted extension Init code and include its declared models")
 	rootCmd.AddCommand(modelsCmd)
 }
 
 func runModels(_ *cobra.Command, args []string) error {
+	if modelsExtensionsFlag && !viper.GetBool("no-extensions") {
+		loaded, err := extensions.LoadExtensionsScoped(viper.GetStringSlice("extension"), bareFlag)
+		if err != nil {
+			return err
+		}
+		if _, err := kitsetup.PrepareProviderFactories(loaded, nil); err != nil {
+			return err
+		}
+		entries := kitsetup.PrepareProviderResolver(loaded).List()
+		if len(args) == 1 {
+			if items, ok := entries[args[0]]; ok {
+				for _, id := range sortedModelIDs(items) {
+					fmt.Println(id)
+				}
+				return nil
+			}
+		}
+		if len(args) == 0 {
+			providers := make([]string, 0, len(entries))
+			for provider := range entries {
+				providers = append(providers, provider)
+			}
+			sort.Strings(providers)
+			for _, provider := range providers {
+				for _, id := range sortedModelIDs(entries[provider]) {
+					fmt.Printf("%s/%s\n", provider, id)
+				}
+			}
+		}
+	}
 	if len(args) == 1 {
 		return printProvider(args[0])
 	}

cmd/root.go

@@ -1437,6 +1437,8 @@ func runInteractiveModeBubbleTea(_ context.Context, deps runModeDeps) error {
 		SkillItems:               snap.skillItems,
 		GetSkillItems:            prov.getSkillItems,
 		ExtensionItems:           snap.extensionItems,
+		GetModelEntries:          prov.getModelEntries,
+		GetThinkingLevels:        prov.getThinkingLevels,
 		GetExtensionItems:        prov.getExtensionItems,
 		StartupExtensionMessages: snap.startupExtensionMessages,
 		GetWidgets:               prov.getWidgets,

cmd/run_mode.go

@@ -73,6 +73,8 @@ type startupSnapshot struct {
 // uiProviders holds read-only getters that the TUI polls to render extension
 // chrome and to refresh lists after hot-reload. None of these mutate state.
 type uiProviders struct {
+	getModelEntries      func() []ui.ModelEntry
+	getThinkingLevels    func() []string
 	getPromptTemplates   func() []*prompts.PromptTemplate
 	getSkillItems        func() []ui.SkillItem
 	getExtensionItems    func() []ui.ExtensionItem
@@ -281,8 +283,12 @@ func newRunApp(k *kit.Kit, cli *ui.CLI, mcpConfig *config.Config, modelName stri
 	if cli != nil {
 		usageTracker = cli.GetUsageTracker()
 	} else {
-		usageTracker = ui.CreateUsageTracker(viper.GetString("model"), k.IsOAuth())
-		ui.UpdateUsageTrackerForMode(usageTracker, viper.GetString("model"), k.CurrentModelMode(), k.IsOAuth())
+		if info := k.GetModelInfo(); info != nil {
+			usageTracker = ui.NewUsageTracker(info, kit.GetCurrentProvider(k.GetCurrentModel()), 80, k.IsOAuth())
+		}
+		if k.CurrentModelMode() != "" {
+			ui.UpdateUsageTrackerForMode(usageTracker, viper.GetString("model"), k.CurrentModelMode(), k.IsOAuth())
+		}
 	}
 	if usageTracker != nil {
 		appOpts.UsageTracker = usageTracker
@@ -500,6 +506,14 @@ func mcpResourceReaderForUI(k *kit.Kit) ui.MCPResourceReader {
 // buildUIProviders wires the read-only getters the TUI polls.
 func buildUIProviders(k *kit.Kit) uiProviders {
 	return uiProviders{
+		getThinkingLevels: k.SupportedThinkingLevels,
+		getModelEntries: func() []ui.ModelEntry {
+			entries := []ui.ModelEntry{}
+			for _, item := range k.GetAvailableModels() {
+				entries = append(entries, ui.ModelEntry{Provider: item.Provider, ModelID: item.ModelID, Name: item.Name, ContextLimit: item.ContextLimit, Reasoning: item.Reasoning, Available: modelEntryAvailable(k, item.Provider, item.ModelID), AvailabilityUnknown: kit.GetProviderInfo(item.Provider) == nil})
+			}
+			return entries
+		},
 		getPromptTemplates: func() []*prompts.PromptTemplate { return loadPromptTemplates(true) },
 		getSkillItems:      func() []ui.SkillItem { return reloadSkillItems(k) },
 		// getExtensionItems re-collects the loaded extension list, used by the
@@ -573,7 +587,9 @@ func buildUIActions(k *kit.Kit, appInstance *app.App, usageTracker *ui.UsageTrac
 			// and m.modelName directly after setModel returns. Extension model
 			// changes instead use the deferred NotifyModelChanged notification.
 			// Update usage tracker with new model info for correct token counting.
-			ui.UpdateUsageTrackerForModel(usageTracker, modelString, k.IsOAuth())
+			if info := k.GetModelInfo(); info != nil && usageTracker != nil {
+				usageTracker.UpdateModelInfo(info, kit.GetCurrentProvider(modelString), k.IsOAuth())
+			}
 			return nil
 		},
 		// saveProviderAPIKey backs the /connect command. Keys go to the
@@ -717,3 +733,12 @@ func startContentWatcher(ctx context.Context, appInstance *app.App) func() {
 	go contentWatcher.Start(ctx)
 	return func() { _ = contentWatcher.Close() }
 }
+
+func modelEntryAvailable(k *kit.Kit, provider, model string) bool {
+	if kit.GetProviderInfo(provider) != nil {
+		return kit.ValidateEnvironment(provider, "") == nil
+	}
+	// Extension auth callbacks are not invoked from the UI event loop. Until
+	// asynchronous status refresh completes, only no-auth providers are enabled.
+	return k.CheckModelAvailable(provider + "/" + model)
+}

examples/extensions/fake-provider.go

@@ -0,0 +1,41 @@
+//go:build ignore
+
+package main
+
+import "kit/ext"
+
+// Helpers precede registration because Yaegi does not reliably resolve later
+// named helpers referenced from function literals.
+func fakeCreate(req ext.ProviderCreateRequest) ext.ProviderCreateResult {
+	return ext.ProviderCreateResult{Runtime: ext.ProviderRuntime{
+		Stream: func(req ext.ProviderRequest, emit func(ext.ProviderEvent) bool) ext.ProviderResult {
+			select {
+			case <-req.Done:
+				return ext.ProviderResult{Error: &ext.ProviderError{Code: "canceled"}}
+			default:
+			}
+			if !emit(ext.ProviderEvent{Type: ext.ProviderTextStart, ID: "greeting"}) {
+				return ext.ProviderResult{}
+			}
+			if !emit(ext.ProviderEvent{Type: ext.ProviderTextDelta, ID: "greeting", Delta: "Hello from a fake provider."}) {
+				return ext.ProviderResult{}
+			}
+			if !emit(ext.ProviderEvent{Type: ext.ProviderTextEnd, ID: "greeting"}) {
+				return ext.ProviderResult{}
+			}
+			emit(ext.ProviderEvent{Type: ext.ProviderFinish, FinishReason: ext.ProviderFinishStop})
+			return ext.ProviderResult{}
+		},
+		Close: func() error { return nil },
+	}}
+}
+
+// Init declares a deterministic offline provider. No API key is required.
+// kit --bare -e examples/extensions/fake-provider.go --model fake/echo --no-session --no-core-tools --quiet "hello"
+func Init(api ext.API) {
+	api.RegisterProvider(ext.ProviderDef{
+		ID: "fake", Name: "Fake provider", Auth: ext.ProviderAuth{Mode: ext.ProviderAuthNone},
+		Models: []ext.ProviderModel{{ID: "echo", Name: "Deterministic greeting", Capabilities: ext.ProviderCapabilities{Streaming: true, Tools: true, InputModalities: []string{"text"}}}},
+		Create: fakeCreate,
+	})
+}

examples/extensions/fake-tool-provider.go

@@ -0,0 +1,52 @@
+//go:build ignore
+
+package main
+
+import "kit/ext"
+
+func toolCreate(req ext.ProviderCreateRequest) ext.ProviderCreateResult {
+	return ext.ProviderCreateResult{Runtime: ext.ProviderRuntime{Stream: func(req ext.ProviderRequest, emit func(ext.ProviderEvent) bool) ext.ProviderResult {
+		// History, not a process-global counter, determines the next response.
+		resultSeen := false
+		for _, message := range req.Messages {
+			for _, part := range message.Parts {
+				if part.Type == ext.ProviderPartToolResult && part.CallID == "add-call" && part.Text == "5" && !part.IsError {
+					resultSeen = true
+				}
+			}
+		}
+		if resultSeen {
+			if !emit(ext.ProviderEvent{Type: ext.ProviderTextStart, ID: "answer"}) {
+				return ext.ProviderResult{}
+			}
+			if !emit(ext.ProviderEvent{Type: ext.ProviderTextDelta, ID: "answer", Delta: "5"}) {
+				return ext.ProviderResult{}
+			}
+			if !emit(ext.ProviderEvent{Type: ext.ProviderTextEnd, ID: "answer"}) {
+				return ext.ProviderResult{}
+			}
+			emit(ext.ProviderEvent{Type: ext.ProviderFinish, FinishReason: ext.ProviderFinishStop})
+		} else {
+			if !emit(ext.ProviderEvent{Type: ext.ProviderToolStart, ID: "add-call", ToolName: "add"}) {
+				return ext.ProviderResult{}
+			}
+			if !emit(ext.ProviderEvent{Type: ext.ProviderToolDelta, ID: "add-call", Delta: `{"a":2,`}) {
+				return ext.ProviderResult{}
+			}
+			if !emit(ext.ProviderEvent{Type: ext.ProviderToolDelta, ID: "add-call", Delta: `"b":3}`}) {
+				return ext.ProviderResult{}
+			}
+			if !emit(ext.ProviderEvent{Type: ext.ProviderToolEnd, ID: "add-call"}) {
+				return ext.ProviderResult{}
+			}
+			emit(ext.ProviderEvent{Type: ext.ProviderFinish, FinishReason: ext.ProviderFinishToolCalls})
+		}
+		return ext.ProviderResult{}
+	}}}
+}
+
+func Init(api ext.API) {
+	api.RegisterProvider(ext.ProviderDef{ID: "fake-tool", Name: "History-driven tool fixture", Create: toolCreate,
+		Models: []ext.ProviderModel{{ID: "add", Name: "Add", Capabilities: ext.ProviderCapabilities{Streaming: true, Tools: true}}},
+	})
+}

internal/agent/agent.go

@@ -1474,6 +1474,7 @@ func (a *Agent) SetModel(ctx context.Context, config *models.ProviderConfig) err
 	a.modelConfig = config
 	a.providerErr = nil
 	if a.currentModel != nil {
+		a.currentModel.resolver.Store(config.Resolver)
 		a.currentModel.set(config.ModelString)
 	}
 
@@ -1577,7 +1578,10 @@ func (a *Agent) GetMaxTokens() int {
 // If MCP tools are still loading in the background, Close waits for them
 // to finish before closing connections to avoid resource leaks. The wait
 // is bounded so a hung MCP transport can never block shutdown forever.
-func (a *Agent) Close() error {
+func (a *Agent) Close() error { return a.CloseContext(context.Background()) }
+
+// CloseContext bounds cooperative provider cleanup by the caller's deadline.
+func (a *Agent) CloseContext(ctx context.Context) error {
 	// Wait for background MCP loading to finish before closing connections.
 	if a.mcpReady != nil {
 		select {
@@ -1591,7 +1595,13 @@ func (a *Agent) Close() error {
 		toolErr = a.toolManager.Close()
 	}
 	if a.providerCloser != nil {
-		if err := a.providerCloser.Close(); err != nil && toolErr == nil {
+		var err error
+		if closer, ok := a.providerCloser.(interface{ CloseContext(context.Context) error }); ok {
+			err = closer.CloseContext(ctx)
+		} else {
+			err = a.providerCloser.Close()
+		}
+		if err != nil && toolErr == nil {
 			toolErr = err
 		}
 	}

internal/agent/image_gate.go

@@ -68,12 +68,14 @@ func imageUnsupportedText(resp fantasy.ToolResponse, model string) string {
 // activeModel holds the model string of the active model. Tools read it while
 // they run, possibly in parallel with SetModel, so it is atomic.
 type activeModel struct {
-	v atomic.Pointer[string]
+	resolver atomic.Pointer[models.Resolver]
+	v        atomic.Pointer[string]
 }
 
 func newActiveModel(cfg *models.ProviderConfig) *activeModel {
 	m := &activeModel{}
 	if cfg != nil {
+		m.resolver.Store(cfg.Resolver)
 		m.set(cfg.ModelString)
 	}
 	return m
@@ -87,6 +89,10 @@ func (m *activeModel) imageInput() (bool, string) {
 	if p == nil || *p == "" {
 		return true, ""
 	}
+	provider, _, _ := models.ParseModelString(*p)
+	if m.resolver.Load().Owns(provider) {
+		return true, *p
+	} // Preserve input for explicit adapter validation.
 	return modelImageInput(*p), *p
 }
 

internal/agent/retry.go

@@ -110,6 +110,10 @@ func retryableRequestError(err error) bool {
 	if err == nil || err == context.DeadlineExceeded || errors.Is(err, context.Canceled) {
 		return false
 	}
+	var extensionError interface{ ExtensionRetryable() bool }
+	if errors.As(err, &extensionError) {
+		return extensionError.ExtensionRetryable()
+	}
 	streamMessage := retryableStreamMessage(err.Error())
 	if pe, ok := errors.AsType[*fantasy.ProviderError](err); ok {
 		if pe.AuthError || pe.StatusCode == http.StatusUnauthorized || pe.StatusCode == http.StatusForbidden {

internal/agent/stream_turn.go

@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"errors"
 	"fmt"
+	"slices"
 	"strings"
 	"sync"
 
@@ -391,6 +392,52 @@ func (t *streamTurn) onToolResult(tr fantasy.ToolResultContent) error {
 }
 
 func (t *streamTurn) onStepFinish(step fantasy.StepResult) error {
+	// Host-executed tool results do not inherit model metadata in the model
+	// library. Carry the matching extension origin so persistence keeps IDs.
+	if metadata := message.ExtensionOptions(fantasy.ProviderOptions(step.ProviderMetadata)); metadata != nil {
+		for i := range step.Messages {
+			positions := map[string]int{}
+			for index, id := range metadata.Order {
+				positions[id] = index
+			}
+			slices.SortStableFunc(step.Messages[i].Content, func(a, b fantasy.MessagePart) int {
+				left, right := message.ExtensionOptions(a.Options()), message.ExtensionOptions(b.Options())
+				if left == nil || right == nil {
+					return 0
+				}
+				return positions[left.BlockID] - positions[right.BlockID]
+			})
+			step.Messages[i].ProviderOptions = message.CloneOptions(step.Messages[i].ProviderOptions)
+			if step.Messages[i].ProviderOptions == nil {
+				step.Messages[i].ProviderOptions = fantasy.ProviderOptions{}
+			}
+			step.Messages[i].ProviderOptions["kit-extension"] = metadata
+		}
+	}
+	origins := map[string]*message.ExtensionMetadata{}
+	for _, msg := range step.Messages {
+		for _, part := range msg.Content {
+			if call, ok := part.(fantasy.ToolCallPart); ok {
+				if metadata := message.ExtensionOptions(call.ProviderOptions); metadata != nil {
+					origins[call.ToolCallID] = metadata
+				}
+			}
+		}
+	}
+	for i, msg := range step.Messages {
+		for j, part := range msg.Content {
+			if result, ok := part.(fantasy.ToolResultPart); ok {
+				if metadata := origins[result.ToolCallID]; metadata != nil {
+					result.ProviderOptions = message.CloneOptions(result.ProviderOptions)
+					if result.ProviderOptions == nil {
+						result.ProviderOptions = fantasy.ProviderOptions{}
+					}
+					result.ProviderOptions["kit-extension"] = &message.ExtensionMetadata{Provider: metadata.Provider}
+					step.Messages[i].Content[j] = result
+				}
+			}
+		}
+	}
 	// Accumulate messages from completed steps so they can be
 	// persisted even if a later step is cancelled.
 	t.completedStepMessages = append(t.completedStepMessages, step.Messages...)