feat(pkg/kit): expose runtime add/remove of native Go tools (#74)

README.md

@@ -917,6 +917,34 @@ mutators and readers (`GetSkills`, `GetContextFiles`) are safe to call
 concurrently from multiple goroutines. See the [SDK overview docs](/sdk/overview#runtime-skills-and-context-files)
 for the full reference.
 
+### Runtime Native Tools
+
+Native Go tools can also be added and removed on a live host, mirroring the
+runtime MCP-server and skill APIs. This is useful for progressive disclosure
+(dynamically loading a domain toolset only when the model asks for it) and for
+multi-tenant hosts that need to swap tool catalogs without rebuilding the host.
+
+```go
+// Add tools that persist for the session.
+host.AddTools(crmTools...)
+
+// Drop a domain when it is no longer needed.
+if err := host.RemoveTools("crm_search_contacts", "crm_create_deal"); err != nil {
+    log.Printf("remove tools: %v", err)
+}
+
+// Replace the entire extra-tool set wholesale.
+host.SetExtraTools(activeTools...)
+
+// Read back the currently registered extra tools.
+extra := host.GetExtraTools()
+```
+
+`AddTools` replaces existing extra tools by name (last-write-wins) and appends
+new ones. `RemoveTools` is atomic: if any supplied name is not currently
+registered, it returns an error and leaves the tool set unchanged. Core tools
+and MCP tools are unaffected. Mutations take effect on the next LLM step.
+
 ## Advanced Usage
 
 ### Subagent Pattern

pkg/kit/kit.go

@@ -8,6 +8,7 @@ import (
 	"log"
 	"os"
 	"path/filepath"
+	"sort"
 	"strings"
 	"sync"
 	"time"
@@ -117,10 +118,17 @@ type Kit struct {
 	steerCh       chan agent.SteerMessage
 	leftoverSteer []agent.SteerMessage // unconsumed steer messages from the last turn
 
-	// promptOptsMu serializes per-call PromptOptions overrides that mutate
-	// shared agent state (model, thinking level, provider creds, extra tools)
-	// so the apply/restore window of one call never races another.
-	promptOptsMu sync.Mutex
+	// promptOptsMu protects shared agent state that can be mutated at runtime
+	// (model, thinking level, provider creds, extra tools). It serializes
+	// writers so the apply/restore window of one call never races another, while
+	// allowing concurrent readers of the extra-tool set.
+	promptOptsMu sync.RWMutex
+
+	// runtimeExtraTools holds native tools added via AddTools / SetExtraTools /
+	// RemoveTools and via Options.ExtraTools at construction. Extension tools
+	// are kept separately on the extension runner and recomposed with this
+	// slice when either side changes.
+	runtimeExtraTools []Tool
 }
 
 // Subscribe registers an EventListener that will be called for every lifecycle
@@ -157,6 +165,139 @@ func (m *Kit) GetToolsForSubagent() []Tool {
 	return tools
 }
 
+// AddTools additively registers native Go tools on the live host. Added tools
+// persist for the session and become visible to the model on the next turn.
+// If a provided tool shares a name with a tool that is already in the
+// extra-tool set, the new tool replaces the previous one (last-write-wins).
+// Core tools and MCP tools are not affected.
+//
+// AddTools is safe to call while the agent is idle. If a turn is in progress
+// ([Kit.IsGenerating] returns true), the change takes effect starting from the
+// next LLM step.
+func (m *Kit) AddTools(tools ...Tool) {
+	m.promptOptsMu.Lock()
+	defer m.promptOptsMu.Unlock()
+
+	cur := m.runtimeExtraTools
+	if len(cur) == 0 && len(tools) == 0 {
+		return
+	}
+
+	replacements := make(map[string]Tool, len(tools))
+	for _, t := range tools {
+		replacements[t.Info().Name] = t
+	}
+
+	seen := make(map[string]struct{}, len(cur)+len(tools))
+	merged := make([]Tool, 0, len(cur)+len(tools))
+	for _, t := range cur {
+		if r, ok := replacements[t.Info().Name]; ok {
+			merged = append(merged, r)
+		} else {
+			merged = append(merged, t)
+		}
+		seen[t.Info().Name] = struct{}{}
+	}
+	for _, t := range tools {
+		if _, ok := seen[t.Info().Name]; !ok {
+			merged = append(merged, replacements[t.Info().Name])
+			seen[t.Info().Name] = struct{}{}
+		}
+	}
+
+	m.runtimeExtraTools = merged
+	m.recomposeExtraTools()
+}
+
+// RemoveTools removes previously-added native Go tools by name. Core tools and
+// MCP tools are unaffected. If any of the supplied names is not currently in
+// the extra-tool set, an error is returned listing the missing names and no
+// tools are removed.
+//
+// RemoveTools is safe to call while the agent is idle. If a turn is in
+// progress, the tools are removed at the next LLM step.
+func (m *Kit) RemoveTools(names ...string) error {
+	m.promptOptsMu.Lock()
+	defer m.promptOptsMu.Unlock()
+
+	cur := m.runtimeExtraTools
+
+	drop := make(map[string]struct{}, len(names))
+	for _, n := range names {
+		drop[n] = struct{}{}
+	}
+
+	missing := make(map[string]struct{}, len(names))
+	for n := range drop {
+		missing[n] = struct{}{}
+	}
+
+	kept := make([]Tool, 0, len(cur))
+	for _, t := range cur {
+		if _, ok := drop[t.Info().Name]; ok {
+			delete(missing, t.Info().Name)
+			continue
+		}
+		kept = append(kept, t)
+	}
+
+	if len(missing) > 0 {
+		list := make([]string, 0, len(missing))
+		for n := range missing {
+			list = append(list, n)
+		}
+		sort.Strings(list)
+		return fmt.Errorf("tool(s) not found: %s", strings.Join(list, ", "))
+	}
+
+	m.runtimeExtraTools = kept
+	m.recomposeExtraTools()
+	return nil
+}
+
+// SetExtraTools replaces the entire native extra-tool set in one call. Core
+// tools and MCP tools are unaffected. Pass an empty slice to clear all
+// extra tools.
+//
+// SetExtraTools is safe to call while the agent is idle. If a turn is in
+// progress, the change takes effect starting from the next LLM step.
+func (m *Kit) SetExtraTools(tools ...Tool) {
+	m.promptOptsMu.Lock()
+	defer m.promptOptsMu.Unlock()
+	m.runtimeExtraTools = append([]Tool(nil), tools...)
+	m.recomposeExtraTools()
+}
+
+// GetExtraTools returns a snapshot of the native extra tools that were added
+// via AddTools / SetExtraTools / RemoveTools or Options.ExtraTools. Extension
+// tools are not included. The returned slice is a copy; modifying it does not
+// affect the tools registered on the host.
+func (m *Kit) GetExtraTools() []Tool {
+	m.promptOptsMu.RLock()
+	defer m.promptOptsMu.RUnlock()
+	if len(m.runtimeExtraTools) == 0 {
+		return nil
+	}
+	out := make([]Tool, len(m.runtimeExtraTools))
+	copy(out, m.runtimeExtraTools)
+	return out
+}
+
+// recomposeExtraTools rebuilds the agent's extra-tool list from extension
+// tools plus runtime native tools. Callers must hold promptOptsMu.
+func (m *Kit) recomposeExtraTools() {
+	var combined []Tool
+	if m.extRunner != nil {
+		extTools := extensions.ExtensionToolsAsLLMTools(m.extRunner.RegisteredTools(), m.extRunner)
+		if len(extTools) > 0 {
+			combined = make([]Tool, 0, len(extTools)+len(m.runtimeExtraTools))
+			combined = append(combined, extTools...)
+		}
+	}
+	combined = append(combined, m.runtimeExtraTools...)
+	m.agent.SetExtraTools(combined)
+}
+
 // GetLoadingMessage returns the agent's startup info message (e.g. GPU
 // fallback info), or empty string if none.
 func (m *Kit) GetLoadingMessage() string {
@@ -774,8 +915,9 @@ func (m *Kit) ReloadExtensions() error {
 
 	// Update extension tools on the agent so the LLM sees changes.
 	if m.agent != nil {
-		extTools := extensions.ExtensionToolsAsLLMTools(m.extRunner.RegisteredTools(), m.extRunner)
-		m.agent.SetExtraTools(extTools)
+		m.promptOptsMu.Lock()
+		m.recomposeExtraTools()
+		m.promptOptsMu.Unlock()
 	}
 
 	// Re-set context and emit SessionStart.
@@ -1691,8 +1833,13 @@ func New(ctx context.Context, opts *Options) (*Kit, error) {
 		contextPrepare:        contextPrepare,
 		beforeCompact:         beforeCompact,
 		prepareStep:           prepareStep,
+		runtimeExtraTools:     append([]Tool(nil), extraTools...),
 	}
 
+	// Ensure the agent's extra-tool list reflects the current extension tools
+	// plus the runtime native tools captured above.
+	k.recomposeExtraTools()
+
 	// Point the activate_skill provider closure at the live Kit instance so it
 	// resolves skills mutated after construction.
 	skillToolKit = k

pkg/kit/tools_test.go

@@ -2,6 +2,9 @@ package kit_test
 
 import (
 	"context"
+	"fmt"
+	"slices"
+	"sync"
 	"testing"
 
 	kit "github.com/mark3labs/kit/pkg/kit"
@@ -263,3 +266,152 @@ func TestNewParallelTool_BinaryImageResponse(t *testing.T) {
 		t.Errorf("ToolResponse.Type = %q, want %q", resp.Type, "image")
 	}
 }
+
+func makeTestTool(name string) kit.Tool {
+	return kit.NewTool(name, "test tool "+name,
+		func(ctx context.Context, input struct{}) (kit.ToolOutput, error) {
+			return kit.TextResult("ok"), nil
+		},
+	)
+}
+
+func sortedStrings(s []string) []string {
+	c := make([]string, len(s))
+	copy(c, s)
+	slices.Sort(c)
+	return c
+}
+
+// TestRuntimeToolMutation verifies that native Go tools can be added, removed,
+// replaced, and enumerated on a live Kit host.
+func TestRuntimeToolMutation(t *testing.T) {
+	t.Setenv("OPENAI_API_KEY", "sk-test")
+	ctx := context.Background()
+
+	host, err := kit.New(ctx, &kit.Options{
+		Model:            "openai/gpt-4o-mini",
+		Quiet:            true,
+		NoSession:        true,
+		NoExtensions:     true,
+		DisableCoreTools: true,
+		SkipConfig:       true,
+		NoSkills:         true,
+		NoContextFiles:   true,
+	})
+	if err != nil {
+		t.Fatalf("Failed to create Kit: %v", err)
+	}
+	defer func() { _ = host.Close() }()
+
+	if got := host.GetExtraTools(); len(got) != 0 {
+		t.Fatalf("expected no extra tools initially, got %d", len(got))
+	}
+	if got := host.GetToolNames(); len(got) != 0 {
+		t.Fatalf("expected no tools initially, got %d: %v", len(got), got)
+	}
+
+	// Add two tools.
+	host.AddTools(makeTestTool("alpha"), makeTestTool("beta"))
+	if got := sortedStrings(host.GetToolNames()); !slices.Equal(got, []string{"alpha", "beta"}) {
+		t.Errorf("after AddTools: expected [alpha beta], got %v", got)
+	}
+	if got := host.GetExtraTools(); len(got) != 2 {
+		t.Errorf("expected 2 extra tools, got %d", len(got))
+	}
+
+	// Adding a duplicate by name replaces it in place without growing the set.
+	host.AddTools(makeTestTool("alpha"))
+	if got := sortedStrings(host.GetToolNames()); !slices.Equal(got, []string{"alpha", "beta"}) {
+		t.Errorf("after duplicate AddTools: expected [alpha beta], got %v", got)
+	}
+	if got := host.GetExtraTools(); len(got) != 2 {
+		t.Errorf("expected 2 extra tools after duplicate add, got %d", len(got))
+	}
+
+	// Removing a missing name returns an error and leaves tools untouched.
+	if err := host.RemoveTools("missing"); err == nil {
+		t.Error("RemoveTools(missing) expected error, got nil")
+	}
+	if got := sortedStrings(host.GetToolNames()); !slices.Equal(got, []string{"alpha", "beta"}) {
+		t.Errorf("after failed RemoveTools: expected [alpha beta], got %v", got)
+	}
+
+	// Remove one tool.
+	if err := host.RemoveTools("alpha"); err != nil {
+		t.Errorf("RemoveTools(alpha) unexpected error: %v", err)
+	}
+	if got := host.GetToolNames(); !slices.Equal(got, []string{"beta"}) {
+		t.Errorf("after RemoveTools(alpha): expected [beta], got %v", got)
+	}
+
+	// Replace the whole extra-tool set.
+	host.SetExtraTools(makeTestTool("x"), makeTestTool("y"), makeTestTool("z"))
+	if got := sortedStrings(host.GetToolNames()); !slices.Equal(got, []string{"x", "y", "z"}) {
+		t.Errorf("after SetExtraTools: expected [x y z], got %v", got)
+	}
+
+	// Remove multiple tools, one missing should fail atomically.
+	if err := host.RemoveTools("x", "notfound"); err == nil {
+		t.Error("RemoveTools(x, notfound) expected error, got nil")
+	}
+	if got := sortedStrings(host.GetToolNames()); !slices.Equal(got, []string{"x", "y", "z"}) {
+		t.Errorf("after failed multi-RemoveTools: expected [x y z], got %v", got)
+	}
+
+	// Remove remaining tools.
+	if err := host.RemoveTools("x", "y", "z"); err != nil {
+		t.Errorf("RemoveTools(x, y, z) unexpected error: %v", err)
+	}
+	if got := host.GetToolNames(); len(got) != 0 {
+		t.Errorf("after removing all tools: expected none, got %v", got)
+	}
+}
+
+// TestRuntimeToolMutationConcurrent exercises the runtime tool mutators from
+// multiple goroutines to verify locking and snapshotting are race-free.
+func TestRuntimeToolMutationConcurrent(t *testing.T) {
+	t.Setenv("OPENAI_API_KEY", "sk-test")
+	ctx := context.Background()
+
+	host, err := kit.New(ctx, &kit.Options{
+		Model:            "openai/gpt-4o-mini",
+		Quiet:            true,
+		NoSession:        true,
+		NoExtensions:     true,
+		DisableCoreTools: true,
+		SkipConfig:       true,
+		NoSkills:         true,
+		NoContextFiles:   true,
+	})
+	if err != nil {
+		t.Fatalf("Failed to create Kit: %v", err)
+	}
+	defer func() { _ = host.Close() }()
+
+	const n = 20
+	var wg sync.WaitGroup
+	wg.Add(n)
+	for i := range n {
+		go func(i int) {
+			defer wg.Done()
+			name := fmt.Sprintf("tool-%d", i)
+			host.AddTools(makeTestTool(name))
+			_ = host.GetExtraTools()
+			if i%2 == 0 {
+				_ = host.RemoveTools(name)
+			}
+		}(i)
+	}
+	wg.Wait()
+
+	// Even-numbered tools were removed; only odd-numbered tools remain.
+	got := sortedStrings(host.GetToolNames())
+	want := make([]string, 0, n/2)
+	for i := 1; i < n; i += 2 {
+		want = append(want, fmt.Sprintf("tool-%d", i))
+	}
+	slices.Sort(want)
+	if !slices.Equal(got, want) {
+		t.Errorf("expected odd-numbered tools after concurrent mutation\n got: %v\nwant: %v", got, want)
+	}
+}