diff --git a/README.md b/README.md index 06b8020f..9b371883 100644 --- a/README.md +++ b/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 diff --git a/pkg/kit/kit.go b/pkg/kit/kit.go index d8abbf98..a2353434 100644 --- a/pkg/kit/kit.go +++ b/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 diff --git a/pkg/kit/tools_test.go b/pkg/kit/tools_test.go index 25e78269..abbc61b0 100644 --- a/pkg/kit/tools_test.go +++ b/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) + } +} diff --git a/www/pages/sdk/options.md b/www/pages/sdk/options.md index 7f6e206d..8a28c056 100644 --- a/www/pages/sdk/options.md +++ b/www/pages/sdk/options.md @@ -55,7 +55,7 @@ host, err := kit.New(ctx, &kit.Options{ // Tools Tools: []kit.Tool{...}, // Replace default tool set entirely - ExtraTools: []kit.Tool{...}, // Add tools alongside defaults + ExtraTools: []kit.Tool{...}, // Add tools alongside defaults (mutate at runtime via host.AddTools/RemoveTools) DisableCoreTools: true, // Use no core tools (0 tools, for chat-only) CoreToolList: []string, // List of core tools to include, if empty (default) include all diff --git a/www/pages/sdk/overview.md b/www/pages/sdk/overview.md index 5def1580..56edb075 100644 --- a/www/pages/sdk/overview.md +++ b/www/pages/sdk/overview.md @@ -200,6 +200,8 @@ Binary data (images, audio, etc.) in `ToolOutput.Data` is automatically forwarde Use `kit.NewParallelTool` for tools that are safe to run concurrently. Use `kit.ToolCallIDFromContext(ctx)` to retrieve the LLM-assigned call ID for logging or tracing. +`Options.ExtraTools` fixes the native tool set at construction time. To add or remove native tools on a live host, see [Runtime native tools](#runtime-native-tools). + ### Schema-driven tools When the tool's input shape isn't known at compile time — tools sourced from @@ -351,6 +353,58 @@ host, _ := kit.New(ctx, &kit.Options{ n, _ := host.AddInProcessMCPServer(ctx, "docs", mcpSrv) ``` +## Runtime native tools + +`Options.Tools` / `Options.ExtraTools` freeze the native Go tool set at +construction time. For progressive disclosure (loading a domain's tools only +when the model asks for them) or multi-tenant hosts that swap tool catalogs per +request, mutate the native tool set on a live host — mirroring the runtime MCP +and skill APIs. No host rebuild, so session history, MCP connections, and the +system-prompt snapshot all survive. + +```go +weatherTool := kit.NewTool("get_weather", "Get current weather for a city", + func(ctx context.Context, input WeatherInput) (kit.ToolOutput, error) { + return kit.TextResult("72°F, sunny in " + input.City), nil + }, +) + +// Add tools that persist for the session (visible on the next turn). +host.AddTools(weatherTool) + +// Drop tools by name when a domain is no longer needed. +if err := host.RemoveTools("get_weather"); err != nil { + log.Printf("remove tools: %v", err) +} + +// Replace the entire native extra-tool set in one call. +host.SetExtraTools(activeToolsForUser...) + +// Inspect the current set (snapshot copy — safe to mutate). +extra := host.GetExtraTools() +``` + +Key points: + +- **Scope is `extraTools` only.** These methods manage the same slice as + `Options.ExtraTools`. Core tools, MCP tools, and extension-registered tools are + never touched, and `GetExtraTools` excludes extension tools. +- **Last-write-wins on name.** `AddTools` replaces any existing extra tool that + shares a `Info().Name`, then appends the rest. Duplicate names within a single + call also resolve to the last one provided. +- **`RemoveTools` is atomic.** If any supplied name is not currently registered, + it returns an error listing the missing names (deduped and sorted) and leaves + the tool set unchanged. +- **Next-step visibility.** Mutations apply from the next LLM step. If a turn is + in progress, the running step finishes with its existing tool set. +- **Composes with per-call tools.** [`PromptOptions.ExtraTools`](#per-call-overrides) + still layers on top for a single call and is reverted afterwards, snapshotting + around whatever persistent set is active. +- **Thread safety.** All four methods are safe to call concurrently; the + extra-tool state is guarded by an internal `RWMutex`. +- **Not session-persisted.** Native tool *definitions* are not serialized into + session state. Re-add them on session resume, just as with `Options.ExtraTools`. + ## Runtime skills and context files Kit auto-discovers skills and `AGENTS.md`-style context files during `New()`,