Compare commits

...

57 Commits

Author SHA1 Message Date
Ed Zynda 9a662d440c fix(ui): reduce TUI visual noise and improve layout
- remove "You" label and icon from user messages, use borderless content block
- remove input title bar ("Enter your prompt...") and hint line
- increase textarea from 3 to 4 rows with top/bottom margin
- hide input hints permanently for a cleaner UI
- match separator colors (use theme.Border for both startup and input dividers)
- make startup separator full terminal width instead of hardcoded 80
- add /help for help hint and pipe separators to status bar
- add printCustomMessage/RenderCustomMessage for custom alert labels
- render /help output as markdown with "Help" alert label
- add Ctrl+V (paste image) to help message keys section
- fix reasoning text wrapping using ANSI-aware lipgloss.Style.Width
- export HighlightFileTokens for cross-package use
2026-04-22 11:41:09 +03:00
Ed Zynda 4ba9d6fab3 feat(events): mirror Fantasy tool input streaming callbacks as Kit events
- Add ToolCallStartEvent, ToolCallDeltaEvent, ToolCallEndEvent to SDK
- Wire Fantasy OnToolInputStart/Delta/End through agent to EventBus
- Add typed convenience subscribers: OnToolCallStart/Delta/End on Kit
- Bridge new events to TUI via ToolCallInputStart/Delta/End app events
- Extend extension system with OnToolCallInputStart/Delta/End handlers
- Add extension event types, API methods, loader wiring, Yaegi symbols
- Update docs: README, SDK skill, extensions skill, www/sdk, www/extensions

Closes #16
2026-04-21 23:28:13 +03:00
Ed Zynda aec0e7cc01 docs: document noOAuth MCP server config field
- Add noOAuth to MCP server fields table in www/pages/configuration.md
- Add pubmed example with noOAuth in README and www config docs
2026-04-21 22:44:27 +03:00
Ed Zynda bac04636bf feat(config): add noOAuth flag to skip OAuth on public MCP servers
- Add NoOAuth field to MCPServerConfig with JSON/YAML support
- Guard OAuth error handling and transport setup with the new flag
- Prevents failed dynamic client registration on servers like PubMed
  that do not support OAuth
2026-04-21 22:24:10 +03:00
Ed Zynda 5f851fd08e fix(ui): require double ctrl+c to quit, matching double-esc pattern
- First ctrl+c clears input and arms quit flag with 3s timeout
- Second ctrl+c within timeout window actually quits
- Show '⚠ Press Ctrl+C again to quit' warning after first press
- Empty input no longer quits immediately on single ctrl+c
- Prompt/overlay states: ctrl+c cancels dialog, re-dispatches to
  main handler for double-press tracking instead of quitting
- Update placeholder, help text, and tests to match new behavior
2026-04-21 22:05:13 +03:00
Ed Zynda f8371836d8 fix(cmd): fix character encoding in OAuth success page
Add charset=utf-8 to Content-Type header and use HTML entity
✓ instead of raw Unicode checkmark to prevent garbled
text display in browsers.

Fixes #9
2026-04-21 21:19:51 +03:00
Ed Zynda 74f00244be fix(ui): wrap reasoning blocks to terminal width to prevent clipping
- wrap thinking text in StreamComponent and render.ReasoningBlock
- plumb width through renderer and streaming item paths
- keeps style consistent with user/assistant blocks and avoids cut-off lines
2026-04-21 20:42:53 +03:00
Ed Zynda b5d7fd4f3e update docs 2026-04-21 20:33:32 +03:00
Ed Zynda 5857d40978 cleanup 2026-04-21 20:27:32 +03:00
Ed Zynda 3ff701054a fix(models): add gpt-5.4 reasoning level support with auto-adjustment
Adds 'none' thinking level to support OpenAI gpt-5.4 models which use
'reasoning_effort: none' instead of 'minimal'. Includes validation and
auto-adjustment when switching models with incompatible levels.

- Add ThinkingNone constant mapping to ReasoningEffortNone
- Add IsValidThinkingLevelForModel() with gpt-5.4 detection
- Add SuggestThinkingLevelFallback() for level migration
- Auto-adjust thinking level on model switch with user notification
- Update all docs to include 'none' in valid levels

Fixes #11
2026-04-21 20:19:00 +03:00
Ed Zynda c1dee3ceba feat(cmd): add --set-default flag and improve auth error messages
Add --set-default flag to 'kit auth login' to automatically set the
provider's default model after successful authentication. When no Anthropic
credentials exist but OpenAI credentials are detected, error messages
now suggest using OpenAI with the correct --model flag.

Fixes #9
2026-04-21 19:52:06 +03:00
Ed Zynda 2d9783a44d fix(ui): make ctrl+c clear input before quitting
Change Ctrl+C behavior to match other terminal AI tools (claude, codex, pi):
- First Ctrl+C clears the current input when text is present
- Second Ctrl+C (within 3 seconds) quits the application
- Ctrl+C on empty input quits immediately
- 3-second auto-reset timer clears the 'pressed once' state
- Flag also resets after message submission

Updates placeholder text and help message to reflect new behavior.

Fixes #13
2026-04-21 19:32:48 +03:00
Ed Zynda 88dd216e15 fix(session): prevent circular parent references in tree session
Add defensive validation to detect and prevent cycles in the session tree
parent chain that could occur after compaction or file corruption.

- Add tree_validation.go with cycle detection and parent chain validation
- Validate parent chain before appending messages (AppendMessage)
- Validate firstKeptEntryID exists in AppendCompaction
- Add depth limit and cycle detection to buildTreeNode to prevent infinite recursion
- Log diagnostics on session open to detect existing cycles
- Add tests for cycle detection and graceful handling
2026-04-21 16:24:38 +03:00
Ed Zynda 9e5806ade8 fix(subagent): remove biased model example from tool schema
- Remove vendor-specific model example that could bias LLM selection
- Add minimum recommended timeout guidance to subagent schema
2026-04-21 11:28:32 +03:00
Ed Zynda 50f586ec8f chore(models): update embedded model database from models.dev
Update internal/models/embedded_models.json with the latest snapshot
from https://models.dev/api.json.

- Providers: 111 → 115 (+4)
- Models: 4,191 → 4,259 (+68)
2026-04-21 10:38:23 +03:00
Ed Zynda 8a8e684dff docs(sdk): document MCPAuthHandler and OAuth opt-in behavior
Reflect the refactor that made MCPAuthHandler an explicit, opt-in
dependency for remote MCP OAuth. Four surfaces updated:

- README.md: new 'MCP OAuth (remote MCP servers)' subsection under the
  Go SDK section, outlining the three consumer patterns (nil / CLI /
  custom) and linking to the full options docs.
- pkg/kit/README.md: type cheat-sheet now lists MCPAuthHandler,
  DefaultMCPAuthHandler, and CLIMCPAuthHandler alongside the existing
  MCPTokenStore entries.
- skills/kit-sdk/SKILL.md: Options example annotated with nil-disables-
  OAuth semantics; new 'MCP OAuth Authorization' section precedes the
  existing token-storage section; re-exported types list expanded.
- www/pages/sdk/options.md: Options fields table gains MCPAuthHandler
  row; new top-level 'MCP OAuth Authorization' section with consumer
  matrix, CLI/custom/fully-custom code samples, and a warning callout
  about the OnAuthURL nil-hang footgun.
2026-04-17 15:30:10 +03:00
Ed Zynda 7ef99ac60f refactor(sdk): remove UX policy from MCP OAuth handler
Strip user-facing I/O out of the SDK's OAuth surface so library, daemon,
and web-app embedders are not surprised by port binds or browser opens.

- DefaultMCPAuthHandler no longer calls openBrowser; it exposes an
  OnAuthURL(serverName, authURL) hook and performs no presentation I/O.
- kit.New no longer auto-constructs a default handler when
  Options.MCPAuthHandler is nil. OAuth is opt-in; remote MCP servers
  requiring authorization fail with a clear error if no handler is set.
- CLIMCPAuthHandler owns the CLI policy (browser open + stderr prints)
  by wiring an OnAuthURL closure on the inner DefaultMCPAuthHandler.
- openBrowser is now unexported and colocated with its sole caller; no
  new exported helper is added to the SDK surface.

BREAKING CHANGE: SDK consumers relying on implicit OAuth with a nil
MCPAuthHandler must now pass kit.NewCLIMCPAuthHandler() (or a custom
implementation) explicitly. The kit CLI is unaffected — cmd/root.go
already constructs the handler explicitly.
2026-04-17 15:26:35 +03:00
Ed Zynda a67f514560 chore(models): refresh embedded models.dev database
- update internal/models/embedded_models.json from https://models.dev/api.json
- 110 → 111 providers, 4172 → 4191 models
2026-04-17 12:19:21 +03:00
Ed Zynda b6bb35cb71 Merge pull request #7 from mark3labs/feat/sdk-options-overrides
feat(sdk): expose generation and provider params on Options
2026-04-17 12:15:47 +03:00
Ed Zynda 4e82fac442 fix(fileutil): decouple TestDetectMediaType from system MIME db
TestDetectMediaType/.go fails on CI images (Ubuntu mime-support) where
/etc/mime.types registers '.go → text/x-go', because mime.TypeByExtension
reads those files at init. The test intended to exercise the 'unknown
extension falls through to text/plain' branch but used a real extension,
making the assertion environment-dependent.

Replace '.go' with '.kitsyntheticext', an invented extension that no
system MIME database registers. The fallback path is now exercised
deterministically on any host.
2026-04-17 12:13:28 +03:00
Ed Zynda 5ec2217b0f docs(sdk): document global viper state leakage in New and Options
The SDK applies Options by calling viper.Set on viper's process-global
store, which means two Kits constructed in the same process are not
isolated from each other: the second New overwrites the first's keys,
and downstream readers (SetModel, GetThinkingLevel, BuildProviderConfig)
observe the most recent value.

- Add a 'Global viper state warning' block to the Options godoc
  explaining the leak, the zero-value-does-not-clear gotcha, and
  pointing at viper.Reset() as the migration workaround.
- Add a matching warning to the New godoc so consumers discover the
  constraint from either entry point.
- Detach the viperInitMu godoc (previously lodged inside New's comment
  block) and clarify that the mutex only guards the construction
  window, not instance isolation.
- Add a TODO noting the proper fix: refactor to a per-call viper.New()
  instance so each Kit owns its own config store.
2026-04-17 12:09:13 +03:00
Ed Zynda 8a851723ba style(sdk): gofmt trailing newlines in kit_test.go 2026-04-17 12:07:54 +03:00
Ed Zynda 53b628c5f8 fix(sdk): map hyphenated config keys to KIT_* env vars
- InitConfig now installs a viper env key replacer so keys like
  "max-tokens" bind to KIT_MAX_TOKENS under AutomaticEnv; previously
  hyphenated keys silently missed their documented env overrides.
- Simplify TestNewPreservesIsSetSemantics: with SkipConfig: true no env
  bindings are registered, so the os.Getenv guard and upper() helper
  were dead weight. Remove both and drop the unused helper.
2026-04-17 12:07:29 +03:00
Ed Zynda e1c94cb362 fix(sdk): align SDK max-tokens floor with CLI default (4096 → 8192)
The SDK last-resort MaxTokens floor is applied in kit.New() when
Options.MaxTokens, KIT_MAX_TOKENS, .kit.yml, and per-model defaults
are all unset. It was 4096 (inherited from the old setSDKDefaults
viper default) while the CLI --max-tokens cobra default is 8192.

Bump the floor to 8192 so SDK and CLI callers start from the same
base value before rightSizeMaxTokens runs, then update README,
skills/kit-sdk/SKILL.md, and www/pages/{configuration,sdk/options}.md
to match.
2026-04-17 11:59:49 +03:00
Ed Zynda ecf95b52e1 fix(sdk): preserve IsSet semantics for generation param overrides
Previously setSDKDefaults() registered viper.SetDefault for max-tokens,
temperature, top-p, top-k, frequency/presence-penalty, and thinking-level.
viper.SetDefault makes IsSet() return true, which silently suppressed
per-model defaults (ApplyModelSettings) and automatic right-sizing
(rightSizeMaxTokens) for every SDK-created Kit — and for CLI runs too,
since cmd/root.go routes through kit.New. Effective max-tokens for
claude-sonnet-4-5 was pinned at 4096 instead of 32768.

- Drop SetDefault for all IsSet-sensitive keys; keep only model,
  system-prompt, stream, num-gpu-layers, main-gpu.
- Apply a 4096 max-tokens floor directly on the *models.ProviderConfig
  struct in kit.New() when nothing else resolved a value. Keeps
  viper.IsSet("max-tokens") == false so rightSizeMaxTokens and
  per-model maxTokens overrides still fire.
- Update Options.MaxTokens / ThinkingLevel godoc to describe the real
  precedence chain.
- Strengthen tests: add Temperature subtest; add
  TestNewPreservesIsSetSemantics regression covering all seven keys;
  split TestNewWithProviderOptions into three subtests including
  Options-beats-viper-state and ProviderURL propagation; add
  resetViper helper so subtests don't bleed state.
- Document the new SDK fields (MaxTokens, ThinkingLevel, Temperature,
  TopP, TopK, FrequencyPenalty, PresencePenalty, ProviderAPIKey,
  ProviderURL, TLSSkipVerify) in README, skills/kit-sdk, and the www
  configuration / sdk/options / sdk/overview pages, including a
  dedicated precedence table.
2026-04-17 11:50:45 +03:00
Ed Zynda 0641c92acc feat(sdk): expose generation and provider params on Options
Adds programmatic overrides on kit.Options for the model/provider knobs
that were previously only reachable through viper.Set() — letting SDK
consumers (web apps, services, embedded agents) configure kit fully
in-code without polluting global viper state or shipping .kit.yml.

Generation parameters:
  - MaxTokens         int      (max output tokens per response)
  - ThinkingLevel     string   (off/low/medium/high)
  - Temperature       *float32
  - TopP              *float32
  - TopK              *int32
  - FrequencyPenalty  *float32
  - PresencePenalty   *float32

Sampling params use pointer types so explicit 0 is distinguishable from
unset; nil leaves provider/per-model defaults in place.

Provider configuration:
  - ProviderAPIKey    string
  - ProviderURL       string
  - TLSSkipVerify     bool

Implementation just pushes Options values into viper inside New(),
so all existing downstream code (BuildProviderConfig, SetModel,
modelSettings lookups, runtime model switching) picks them up
uniformly without any new code paths. Tests added for MaxTokens,
ThinkingLevel, and ProviderAPIKey.
2026-04-17 11:24:00 +03:00
Ed Zynda 3bb20f5283 feat(models): surface and prevent silent max-tokens truncation
- Raise --max-tokens default from 4096 to 8192.
- Auto-raise MaxTokens toward the model's catalog Limit.Output (capped at
  32768) when the user hasn't set --max-tokens explicitly and no per-model
  modelSettings override applied. Prevents silent 4k/8k truncation on
  models that support 32k-262k output.
- Surface FinishReasonLength at turn end: the app now subscribes to
  TurnEndEvent and renders a system-message banner explaining the current
  cap, the model's known ceiling, and how to raise it. Previously the TUI
  swallowed 'length' stops, producing 'ghost' truncations.
- Export FinishReason* constants on pkg/kit (Stop, Length, ToolCalls,
  ContentFilter, Error, Other, Unknown) and fix stale comments that used
  Anthropic-style strings.
- Add Kit.MaxTokens() and Kit.MaxOutputLimit() SDK accessors, backed by
  Agent.GetMaxTokens() which correctly returns 0 for providers that
  suppress the param (e.g. Codex OAuth).
- Tests: rightSizeMaxTokens covers 7 paths (cap, raise, preserve,
  explicit flag, nil info, zero limit); handleTurnEnd covers length/
  non-length/nil-sendFn and the fallback message formatter.
- Docs: update configuration.md, cli/flags.md, and kit-extensions skill
  to reflect the new default and behavior.
2026-04-16 23:12:10 +03:00
Ed Zynda 633fa38b2b fix(ui): regenerate spinner frames on theme change
- UpdateTheme() only refreshed typography styles, leaving spinner
  frames rendered with the old theme's colors
- Now calls knightRiderFrames() to rebuild frames with the new
  theme's Primary, Muted, VeryMuted, and MutedBorder colors
2026-04-16 12:32:49 +03:00
Ed Zynda f905cee48c fix(ui): dynamically size slash command name column in popup
- Replace hardcoded nameWidth of 15 with dynamic calculation based on
  the longest command name in the filtered list
- Prevents truncation of longer names like /feature-request and
  /release-tagger that were cut off with ellipsis
- Cap name column to leave at least 20 chars for descriptions
- Add 1 char gap between name and description columns
2026-04-16 12:27:56 +03:00
Ed Zynda 182c10ea1a refactor(ui): improve keybinding ergonomics for terminal multiplexers
- Move thinking toggle from ctrl+t to leader chord (ctrl+x t) to avoid
  conflicts with tmux/zellij tab mode and terminal new-tab shortcuts
- Change scrollback jump from alt+home/alt+end to ctrl+home/ctrl+end
  for better compatibility across SSH and older tmux versions
- Remove ctrl+d as submit alias (enter suffices); avoids EOF convention
  confusion and accidental shell disconnects
- Remove ctrl+a from tree selector filter shortcuts to avoid conflict
  with the common tmux prefix remap (ctrl+o cycle still reaches all
  filter modes)
2026-04-16 12:21:37 +03:00
Ed Zynda fcaa52bf1c fix(extensions): serialize handler calls per-extension to prevent data races
- Add per-extension reentrant mutex to Runner that serializes handler
  invocations from concurrent goroutines (e.g. parallel subagent events)
  while allowing re-entrant calls (handler → EmitCustomEvent → handler)
- Fix subagent-monitor slice aliasing bug: submonEntries[:0] reuses the
  backing array, corrupting entries during in-place filtering
- Pass parent's loaded MCPConfig to child subagents in Kit.Subagent(),
  eliminating concurrent viper map access during parallel kit.New() calls
- Add Options.MCPConfig field so SDK consumers can also skip viper reads
- Add tests for concurrent emit, cross-extension concurrency, and
  re-entrant EmitCustomEvent
2026-04-16 12:11:10 +03:00
Ed Zynda 7e6455732c docs: update documentation for sudo password prompt feature
- README.md: mention interactive sudo password prompt in features
- skills/kit-sdk/SKILL.md: add PasswordPromptEvent to event types table
- www/pages/index.md: update features list with sudo prompt
- www/pages/development.md: update project structure description
- www/pages/sdk/callbacks.md: add complete event types table
2026-04-15 18:06:11 +03:00
Ed Zynda 71301a9035 feat: add interactive sudo password prompt for bash tool
Add core TUI support for handling sudo password prompts when executing
bash commands that require elevated privileges.

- Detect sudo commands and check if credentials are cached (sudo -n)
- Show modal password prompt with masked input (• characters) when needed
- Pipe password via stdin using sudo -S -p '' (no password in command string)
- Password flows through context callbacks, never stored in session history
- Add PasswordPromptHandler to agent and SDK event system
- Add password prompt overlay to TUI with 🔐 icon and hidden input
- Include tests for sudo command detection and rewriting

The password is never persisted to disk - it only exists in memory
during execution and is piped directly to sudo via stdin.
2026-04-15 17:33:03 +03:00
Ed Zynda 0974d37ab2 feat(sdk): support mcp-go in-process transport for MCP servers
- Add InProcessServer field to MCPServerConfig (json:"-", never serialized)
- Add "inprocess" transport type to config, validation, and connection pool
- Add createInProcessClient() using mcp-go client.NewInProcessClient()
- Add Kit.AddInProcessMCPServer() convenience method
- Add Options.InProcessMCPServers for init-time registration
- Export MCPServer type alias (= server.MCPServer) in pkg/kit/types.go
- Add 8 tests covering config, pool, tool manager, and edge cases
- Update SDK README, kit-sdk skill, and www docs
2026-04-15 16:29:07 +03:00
Ed Zynda 398e825df8 docs: update docs for recent features and API additions
- Add smart @ attachments (MIME detection, @mcp:server:uri syntax)
- Add MCP Prompts and Resources SDK APIs to skill and www docs
- Add $+ required variadic placeholder for prompt templates
- Add Ctrl+X e (external editor) and Ctrl+X s (steer) keyboard shortcuts
- Fix stale Ctrl+S references, now Ctrl+X s for mid-turn steering
- Add --frequency-penalty and --presence-penalty CLI flags
- Add per-model settings (modelSettings) to configuration docs
- Add NoExtensions, NoSkills, NoContextFiles, SessionManager,
  MCPTokenStoreFactory to SDK options docs
- Add bridge_demo.go to extension examples
- Add dynamic MCP servers, subagents to SDK overview
2026-04-15 16:02:49 +03:00
Ed Zynda 3c51c20be7 feat(mcp): handle embedded resources in prompt messages
- Extract all MCP content types in prompt expansion: ImageContent,
  AudioContent, EmbeddedResource (text and blob), and ResourceLink
- Add MCPFilePart type to carry decoded binary attachments through
  the tools → SDK → bridge → UI layers
- Inline text resources as fenced code blocks with URI annotation
- Decode image/audio/blob content from base64 into LLMFilePart
  attachments submitted via RunWithFiles
- Render ResourceLink as text annotation for the LLM
- Show attachment badges on user messages (e.g. '1 image(s) attached')
  matching the existing clipboard paste UI pattern
- Log warnings on base64 decode failures instead of silently dropping
2026-04-15 15:23:01 +03:00
Ed Zynda 25410af440 feat: add smart @ attachments with MIME detection and MCP resource support
Phase 1: Smart @ for local files
- ProcessFileAttachments now returns FileAttachmentResult with separate
  ProcessedText and FileParts fields instead of a plain string
- Binary files (images, audio, video, PDFs, etc.) detected via MIME type
  are extracted as multimodal FileParts instead of XML-wrapped text garbage
- detectMediaType() uses extension-based lookup then content sniffing
- isBinaryMediaType() classifies image/*, audio/*, video/*, and specific
  application types as binary
- @mcp:server:uri token format for referencing MCP resources in text
- All 4 submission paths (TUI submit, TUI steer, MCP prompt, CLI) updated
- App.RunOnceWithFiles/RunOnceResultWithFiles/RunOnceWithDisplayAndFiles
  added for non-interactive multimodal submission

Phase 2: MCP resources in @ autocomplete
- MCPToolManager gains loadServerResources(), GetResources(), ReadResource(),
  SubscribeResource(), UnsubscribeResource(), RefreshServerResources()
- MCPResource and MCPResourceContent types for resource metadata/content
- FileSuggestion extended with IsMCPResource, MCPServerName, MCPResourceURI
- InputComponent.SetMCPResourceProvider() wires resource suggestions into
  the @ popup alongside local files
- @ popup merges local file suggestions with MCP resource suggestions,
  sorted by fuzzy match score
- MCP resources display 'mcp:servername' in the popup description
- Selecting an MCP resource inserts @mcp:server:uri format
- ProcessFileAttachments resolves @mcp: tokens via MCPResourceReader callback
- Text resources are XML-wrapped as <resource>; binary resources become
  FileParts for multimodal submission
- Agent, Kit SDK, and cmd/root.go wired end-to-end

Phase 3: Resource subscriptions (foundation)
- SubscribeResource/UnsubscribeResource on MCPToolManager
- onResourcesChanged callback for live refresh (wired but not yet
  triggering UI refresh automatically)
- RefreshServerResources for manual resource list refresh
2026-04-15 13:01:36 +03:00
Ed Zynda 26c9f009f9 refactor: remove fantasy dependency name leaks from SDK surface
- Rename ExtensionToolsAsFantasy -> ExtensionToolsAsLLMTools
- Rename convertKitMessagesToFantasy -> convertToLLMMessages
- Delete GetFantasyProviders, ToFantasyMessages, FromFantasyMessage
- Replace direct fantasy type usage with kit.LLM* aliases in app tests
- Scrub fantasy references from godoc comments across pkg/kit and internal
2026-04-15 12:24:52 +03:00
Ed Zynda e068487ff7 style(ui): fix gofmt alignment in MCPPromptInfo struct 2026-04-15 11:50:33 +03:00
Ed Zynda 0ffb0ba788 refactor(tools): remove fantasy dependency from internal/tools
- Replace fantasy.AgentTool with plain MCPTool struct in MCPToolManager
- Move fantasy adapter from internal/tools to internal/agent as mcpAgentTool
- Add MCPToolManager.ExecuteTool() for framework-agnostic tool execution
- Remove dead fantasy.LanguageModel field from MCPConnectionPool
- Remove MCPToolManager.SetModel() (was only feeding the dead field)

internal/tools is now a pure MCP client library with no LLM framework
dependency. The fantasy-to-MCP bridging is confined to the agent layer
where it belongs.
2026-04-15 11:27:47 +03:00
Ed Zynda 65c6e9f797 refactor(models): decouple TUI progress from SDK dependency tree
- Remove direct internal/ui/progress import from internal/models/providers.go
- Add ProgressReaderFunc callback to ProviderConfig for dependency inversion
- Wire Bubble Tea progress reader via CLIOptions in cmd/root.go
- Add NewProgressReadCloser convenience wrapper in progress package
- SDK consumers (pkg/kit) no longer transitively pull bubbletea, lipgloss
  v2, or bubbles
- Update embedded_models.json from models.dev (110 providers, 4172 models)
2026-04-14 17:17:01 +03:00
Ed Zynda 68d798d2f4 feat(prompts): add $+ required variadic, skip code in placeholders
- Add internal/fences package for detecting markdown code regions
  (fenced blocks and inline code spans) with ReplaceOutside/StripCode
- SubstituteArgs, HasArgPlaceholders, RequiredArgs now skip $
  placeholders inside ``` fences and `inline` code spans
- ProcessFileAttachments skips @file tokens inside code regions
- Add $+ placeholder: expands like $@ but requires at least 1 argument
- Add RequiredArgs() method; expandPromptTemplate validates arg count
  and re-populates input on failure instead of submitting
- Update feature-request, file-issue, new-prompt to use $+
2026-04-14 13:22:10 +03:00
Ed Zynda eefd5565f8 feat(ui): populate input instead of auto-submitting prompts with args
- Add HasArgPlaceholders() method to PromptTemplate to detect , $@,
  $ARGUMENTS, etc. placeholders in template content
- Add HasArgs field to SlashCommand struct
- Set HasArgs when registering prompt templates as slash commands
- In fuzzy finder Enter handler, populate input with command + trailing
  space when HasArgs is true, letting the user type arguments naturally
- Fix potential index bug by capturing selectedCmd before resetting index
2026-04-14 12:46:12 +03:00
Ed Zynda 9d1b8a102e feat(ui): open external $EDITOR via ctrl+x e chord
- Add ctrl+x e leader key chord to open $VISUAL/$EDITOR in a temp file
  pre-populated with the current input text
- On save & quit, replace the input textarea with the edited content
- On error exit (e.g. :cq in vim), silently preserve original input
- Use charmbracelet/x/editor for editor command construction
- Use tea.ExecProcess to suspend/resume the TUI around the editor
- Update input hint text at all width breakpoints to show the shortcut
- Add ctrl+x e to /help output

Closes #5
2026-04-14 12:39:29 +03:00
Ed Zynda f57e045c69 feat(ui): highlight @file tokens in user messages with accent color
- Add highlightFileTokens() to style @file references with theme.Accent + bold
- Apply highlighting in UserBlock() after line wrapping, before herald rendering
- Support both unquoted (@path/to/file) and quoted (@"path with spaces") tokens
- Add tests for highlightFileTokens and UserBlock integration

Closes #6
2026-04-14 12:28:04 +03:00
Ed Zynda eb5da28a15 chore(deps): update all dependencies
- bump go directive to 1.26.2 (required by fantasy v0.17.2)
- fantasy v0.17.1 → v0.17.2
- bubbletea/v2 v2.0.2 → v2.0.5
- lipgloss/v2 v2.0.2 → v2.0.3
- mcp-go v0.47.1 → v0.48.0
- x/ansi v0.11.6 → v0.11.7
- x/term v0.41.0 → v0.42.0
- genai v1.52.1 → v1.54.0
- ultraviolet, x/crypto, x/net, x/text, x/exp, and other indirects
- allow MODEL env override in acp_smoke_test.py
2026-04-14 12:16:22 +03:00
Ed Zynda cd8e2a7654 feat(extensions): expand inline bash in editor for interactive mode
- Add editor interceptor via OnSessionStart so !{...} expansions
  appear in the user message block on screen
- Restrict OnInput handler to non-interactive sources (CLI, script,
  queue) to avoid double expansion
- Extract expand() helper and hoist regex to package level
- Update doc comments and examples
2026-04-14 11:56:41 +03:00
Ed Zynda 64da1caf41 docs(readme): add clickable links to extension examples
- Convert 31 extension example entries from plain code spans to
  Markdown links pointing to their source files
- Link the go-edit-lint.go and tool-logger_test.go references
2026-04-14 11:39:43 +03:00
Ed Zynda 7eaeafff8c fix(mcp): propagate OAuth config for runtime-added servers
- Store authHandler and tokenStoreFactory on Agent struct so
  AddMCPServer() can propagate them to new MCPToolManagers (#3)
- Add OAuthClientID, OAuthClientSecret, OAuthScopes fields to
  MCPServerConfig for servers without dynamic registration (#4)
- Pass OAuth fields from server config to transport OAuthConfig
  in both SSE and Streamable HTTP client creation paths
- Add GetAuthHandler() accessor to MCPToolManager
- Add tests for auth handler propagation and OAuth config parsing

Closes #3, closes #4
2026-04-11 15:24:47 +03:00
Ed Zynda 8ed8d23c73 docs(sdk): update kit-sdk skill with recent API additions
- Add NoSkills, NoExtensions, NoContextFiles options
- Add MCPTokenStoreFactory option and MCP OAuth types
- Add dynamic MCP server management (AddMCPServer, RemoveMCPServer,
  ListMCPServers, MCPServerStatus)
- Add per-model system prompts and generation parameters sections
- Add MCPToolsReady() to tool querying section
- Expand LLMUsage fields to include CacheCreationTokens/CacheReadTokens
- Update FinalUsage and ShouldCompact docs for cache-aware token counting
- Add MCP OAuth types to re-exported types reference
2026-04-11 12:09:51 +03:00
Ed Zynda 2de98d32be fix(ui): accurate context token tracking including cache tokens
- Include all token categories in context fill calculation:
  InputTokens + CacheReadTokens + CacheCreationTokens + OutputTokens
- With Anthropic/kimi prompt caching, InputTokens can be near-zero
  while CacheReadTokens holds the bulk of the context
- Include OutputTokens since assistant output becomes context next turn
- Remove max-only guard in SetContextTokens so context shrinks after
  compaction instead of staying stuck at the high-water mark
- Reset context tokens to 0 after compaction in both SDK and UI layers
- Use real API-reported token counts in ShouldCompact() instead of
  the chars/4 text heuristic which misses system prompts and tool defs
2026-04-10 17:05:47 +03:00
Ed Zynda 83127467c5 feat(sdk): add NoExtensions, NoSkills, and NoContextFiles options
- Add NoExtensions bool to Options, OR with viper fallback
- Add NoSkills bool to Options, guards all skill loading
- Add NoContextFiles bool to Options, skips AGENTS.md discovery
- SDK consumers can disable autoloading without touching viper
2026-04-09 17:07:31 +03:00
Ed Zynda e07c94f49d feat(mcp): add dynamic MCP server loading and unloading
- Add AddServer/RemoveServer to MCPToolManager for runtime server management
- Add RemoveConnection to MCPConnectionPool for per-server teardown
- Add AddMCPServer/RemoveMCPServer/ListMCPServers to Agent and SDK Kit
- Lazily create connection pool so AddServer works without prior LoadTools
- Wire onToolsChanged callback to trigger agent tool list rebuild
- Make MCPToolManager.Close nil-safe when pool was never initialized

Tests:
- Integration tests with real stdio MCP server (Python echo server)
- Agent-level tests using mock LLM model (no API key needed)
- Unit tests for error paths, callbacks, idempotency, nil safety
- SDK type surface tests
2026-04-09 13:54:11 +03:00
Ed Zynda b87146a284 feat(sdk): add MCPTokenStoreFactory for custom OAuth token storage
- Add MCPTokenStoreFactory option to kit.Options allowing SDK consumers
  to provide custom token storage backends for remote MCP servers
- Thread TokenStoreFactory through the full chain: kit.Options →
  kitsetup → agent → MCPToolManager → MCPConnectionPool
- Add createTokenStore() helper on connection pool that delegates to the
  factory or falls back to the default FileTokenStore
- Export MCPTokenStore, MCPToken, MCPTokenStoreFactory, and ErrMCPNoToken
  in pkg/kit/types.go following SDK naming conventions
- Default behavior (file-based storage) is preserved when factory is nil
2026-04-09 13:27:40 +03:00
Ed Zynda 186d9f7f44 fix(ui): route raw fmt.Print calls through proper renderers
- event_handler: route default extension print level through DisplayInfo
  instead of bare fmt.Println for consistent styling and timestamps
- factory: remove orphan fmt.Println("") before system messages; the
  renderer already manages its own spacing
- app: PrintFromExtension non-interactive fallback now respects level,
  writing errors/info to stderr with prefix to keep stdout clean
- app: PrintBlockFromExtension non-interactive fallback writes framed
  blocks to stderr instead of raw text to stdout
2026-04-09 13:00:23 +03:00
Ed Zynda 3a8ffc2104 feat(models): add per-model system prompt support
- Add systemPrompt field to GenerationParams and config structs
- On init, replace default system prompt with per-model prompt when
  user hasn't explicitly set one (via flag, config, or SDK option)
- On model switch, detect per-model prompt and compose it with
  AGENTS.md, skills, and date/cwd context
- Fix viper.IsSet bug: BindPFlag causes IsSet to return true for
  unset flags, so compare against defaultSystemPrompt instead
- Agent.SetModel now updates stored system prompt from config
- Export LoadModelSettingsFromConfig, LoadSystemPromptValue, and
  LookupModelForSettings for use by Kit.SetModel
- Add tests for prompt apply, precedence, file path, and
  modelSettings override
2026-04-09 12:35:00 +03:00
Ed Zynda e54570162e feat(models): add per-model generation parameter defaults
- Add modelSettings config section for attaching generation params
  (temperature, topP, topK, frequencyPenalty, presencePenalty,
  maxTokens, stopSequences, thinkingLevel) to any model by
  provider/model key
- Add params field to customModels definitions for inline defaults
- Change BuildProviderConfig and SetModel to use viper.IsSet so
  unset params remain nil, allowing model-level defaults to apply
- Wire ApplyModelSettings into CreateProvider with priority order:
  CLI flags > global config > modelSettings > customModels params
- Add GenerationParams to ModelInfo in the registry
- Update default config template with modelSettings and customModels
  params examples
2026-04-09 12:07:42 +03:00
115 changed files with 10962 additions and 1002 deletions
+6 -3
View File
@@ -13,6 +13,8 @@
// - No channels in maps (Yaegi panics on range over map[string]chan)
// - All ctx.* calls guarded with nil checks
// - Simple data structures only
// - The extension runner serializes handler calls per-extension, so
// concurrent subagent events cannot race on this shared state.
package main
import (
@@ -43,7 +45,8 @@ const (
)
// ---------------------------------------------------------------------------
// Package-level state - all simple types
// Package-level state — safe because the runner serializes all handler
// invocations for the same extension (per-extension reentrant mutex).
// ---------------------------------------------------------------------------
var (
@@ -282,8 +285,8 @@ func Init(api ext.API) {
submonPushWidget()
// Remove the entry immediately (no goroutine to avoid races)
newEntries := submonEntries[:0]
// Remove the entry — build a new slice to avoid aliasing bugs
newEntries := make([]*submonEntry, 0, len(submonEntries))
for _, en := range submonEntries {
if en.callID != e.ToolCallID {
newEntries = append(newEntries, en)
+2 -2
View File
@@ -2,7 +2,7 @@
description: Create a feature request using the GitHub template
---
Create a feature request for the Kit repository. The user wants to request: $@
Create a feature request for the Kit repository. The user wants to request: $+
## Feature Request Template
@@ -16,7 +16,7 @@ This prompt uses the `feature_request` GitHub template which requires:
## Steps
1. **Understand the request** from `$@`
1. **Understand the request** from `$+`
- What capability is missing?
- What would the ideal behavior look like?
+2 -2
View File
@@ -2,7 +2,7 @@
description: File a GitHub issue using the appropriate template
---
File a GitHub issue for the Kit repository. The user wants to create an issue about: $@
File a GitHub issue for the Kit repository. The user wants to create an issue about: $+
## Issue Templates Available
@@ -16,7 +16,7 @@ This repository has structured issue templates. You MUST use the appropriate tem
## Steps
1. **Determine the issue type** from `$@`:
1. **Determine the issue type** from `$+`:
- Bug → use `--template bug_report`
- Feature → use `--template feature_request`
- Documentation → use `--template documentation`
+6 -4
View File
@@ -2,7 +2,7 @@
description: Scaffold a new prompt template in .kit/prompts/
---
Create a new kit prompt template. The user wants a prompt that does: $@
Create a new kit prompt template. The user wants a prompt that does: $+
## What a prompt template is
@@ -23,19 +23,21 @@ $1 $2 etc. for positional arguments.
- **Filename** → slug: `commit-push.md` becomes `/commit-push`
- **Frontmatter**: only `description` is recognised; keep it under ~80 chars
- **Body**: plain markdown; the full text is submitted as the user's message when the template fires
- **Arguments**: `$@` expands to everything the user typed after the slash command name;
- **Arguments**: `$+` expands to everything the user typed after the slash command name
(requires at least one argument); `$@` is the same but allows zero arguments;
`$1`, `$2` for individual positional args; omit entirely if no arguments are needed
## Steps
1. **Understand the workflow** the user described in `$@` — ask a clarifying question if the intent is ambiguous
1. **Understand the workflow** the user described in `$+` — ask a clarifying question if the intent is ambiguous
2. **Choose a filename**: short, lowercase, hyphen-separated, descriptive (e.g. `code-review.md`)
3. **Write the description**: one sentence, imperative, fits in autocomplete
4. **Draft the body**:
- Open with a single sentence stating the goal
- Use `## Steps` for multi-step workflows; use plain prose for simple prompts
- Be specific: name commands, flags, and file paths where relevant
- End with `$@` on its own line if the user might want to pass context or a hint; omit if the prompt is self-contained
- End with `$+` on its own line if the user must pass context; use `$@` if arguments
are optional; omit if the prompt is self-contained
5. **Write the file** to `.kit/prompts/<slug>.md`
6. **Confirm** by showing the final file content and the slash command that activates it
-80
View File
@@ -1,80 +0,0 @@
# Autoscroll Fix - Final Summary
## Root Cause
The autoscroll was failing for streaming assistant messages due to a bug in how `GotoBottom()` calculated item heights.
### The Problem
1. **Reasoning blocks** (`StreamingMessageItem` with `role="reasoning"`) are **never cached** because they have live duration counters that update every render
2. The `Height()` method returns `0` when `cachedRender == ""`
3. `GotoBottom()` was calling:
```go
itemHeight := item.Height() // Returns 0 for reasoning
if itemHeight == 0 {
item.Render(s.width) // Renders but doesn't cache (reasoning)
itemHeight = item.Height() // Still returns 0!
}
```
4. This caused incorrect scroll position calculations, especially during reasoning → assistant transitions
## The Solution
Changed `GotoBottom()` and `AtBottom()` to calculate height **directly from the rendered string** instead of relying on the cached height:
```go
// OLD: item.Height() which checks cached render
itemHeight := item.Height()
if itemHeight == 0 {
item.Render(s.width)
itemHeight = item.Height() // Still might be 0!
}
// NEW: Calculate from rendered string directly
rendered := item.Render(s.width)
itemHeight := strings.Count(rendered, "\n") + 1
```
This works for **all** items regardless of whether they cache their render or not.
## Files Changed
### `internal/ui/scrolllist.go`
- **`GotoBottom()`**: Calculate height from rendered string (2 loops)
- **`AtBottom()`**: Calculate height from rendered string (1 loop)
### `internal/ui/model.go`
- **`appendStreamingChunk()`**: For existing messages, call `GotoBottom()` directly (iteratr pattern)
- **`refreshContent()`**: Simplified to only call `SetItems()` (removed redundant `GotoBottom()`)
- **Bash streaming handler**: Removed redundant `GotoBottom()` after `refreshContent()`
## Testing Results
✅ **Test prompt**: "explore this repo"
**Before fix**:
- Autoscroll stopped after reasoning block completed
- Viewport stuck showing end of reasoning ("Thought for 203ms")
- Assistant response streamed off-screen below
**After fix**:
- Autoscroll works throughout reasoning block
- Autoscroll continues during reasoning → assistant transition
- Viewport stays at bottom showing latest assistant content
- Final position shows end of response (build commands section)
## Behavior Verified
1. ✅ Streaming text auto-scrolls to bottom
2. ✅ Works across reasoning → assistant transition
3. ✅ Manual scroll up (PgUp) disables autoscroll
4. ✅ Scroll to bottom (Alt+End) re-enables autoscroll
5. ✅ Accurate positioning with no offset errors
## Performance Note
The fix calls `Render()` on all items during `GotoBottom()` calculations. This is acceptable because:
- `Render()` is already optimized with caching for non-reasoning items
- `GotoBottom()` is only called during content updates (not every frame)
- Reasoning blocks need to render anyway for live duration updates
- This matches iteratr's approach of ensuring items are rendered before height calculations
+122 -41
View File
@@ -18,7 +18,8 @@ A powerful, extensible AI coding agent CLI with multi-provider support, built-in
## Features
- **Multi-Provider LLM Support**: Anthropic, OpenAI, Google Gemini, Ollama, Azure OpenAI, AWS Bedrock, OpenRouter, and more
- **Built-in Core Tools**: bash, read, write, edit, grep, find, ls, subagent - no MCP overhead
- **Built-in Core Tools**: bash (with interactive sudo password prompt), read, write, edit, grep, find, ls, subagent - no MCP overhead
- **Smart @ Attachments**: Binary files auto-detected via MIME type, MCP resources via `@mcp:server:uri`
- **MCP Integration**: Connect external MCP servers for expanded capabilities
- **Extension System**: Write custom tools, commands, widgets, and UI modifications in Go
- **Theming**: 22 built-in color themes (KITT, Catppuccin, Dracula, Nord, etc.) with runtime switching, persistence, and custom theme files
@@ -125,8 +126,13 @@ model: anthropic/claude-sonnet-latest
max-tokens: 4096
temperature: 0.7
stream: true
thinking-level: off # off, none, minimal, low, medium, high
```
All of the above keys can also be set programmatically via the SDK
(`kit.Options.MaxTokens`, `Options.Temperature`, `Options.ThinkingLevel`, etc.)
without touching config files — see [SDK options](#with-options).
### Environment Variables
```bash
@@ -151,6 +157,11 @@ mcpServers:
search:
type: remote
url: "https://mcp.example.com/search"
pubmed:
type: remote
url: "https://pubmed.mcp.example.com"
noOAuth: true # skip OAuth for public servers that don't require auth
```
## CLI Reference
@@ -186,12 +197,14 @@ mcpServers:
--no-prompt-templates Disable prompt template loading
# Generation parameters
--max-tokens Maximum tokens in response (default: 4096)
--max-tokens Maximum tokens in response (default: 8192, auto-raised up to 32768 for models with larger known output limits)
--temperature Randomness 0.0-1.0 (default: 0.7)
--top-p Nucleus sampling 0.0-1.0 (default: 0.95)
--top-k Limit top K tokens (default: 40)
--stop-sequences Custom stop sequences (comma-separated)
--thinking-level Extended thinking level: off, minimal, low, medium, high (default: off)
--frequency-penalty Penalize frequent tokens 0.0-2.0 (default: 0.0)
--presence-penalty Penalize present tokens 0.0-2.0 (default: 0.0)
--thinking-level Extended thinking level: off, none, minimal, low, medium, high (default: off)
# System
--config Config file path (default: ~/.kit.yml)
@@ -203,9 +216,10 @@ mcpServers:
```bash
# Authentication (for OAuth-enabled providers)
kit auth login [provider] # Start OAuth flow (e.g., anthropic)
kit auth logout [provider] # Remove credentials for provider
kit auth status # Check authentication status
kit auth login [provider] # Start OAuth flow (e.g., anthropic)
kit auth login [provider] --set-default # Set provider's default model as system default
kit auth logout [provider] # Remove credentials for provider
kit auth status # Check authentication status
# Model database
kit models [provider] # List available models (optionally filter by provider)
@@ -287,7 +301,7 @@ kit -e examples/extensions/minimal.go
### Extension Capabilities
**Lifecycle Events**: OnSessionStart, OnSessionShutdown, OnBeforeAgentStart, OnAgentStart, OnAgentEnd, OnToolCall, OnToolExecutionStart, OnToolOutput, OnToolExecutionEnd, OnToolResult, OnInput, OnMessageStart, OnMessageUpdate, OnMessageEnd, OnModelChange, OnContextPrepare, OnBeforeFork, OnBeforeSessionSwitch, OnBeforeCompact, OnCustomEvent, OnSubagentStart, OnSubagentChunk, OnSubagentEnd
**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
**Custom Components**:
- **Tools**: Add new tools the LLM can invoke
@@ -317,39 +331,40 @@ kit -e examples/extensions/minimal.go
See the `examples/extensions/` directory:
- `minimal.go` - Clean UI with custom footer
- `auto-commit.go` - Auto-commit on shutdown
- `bookmark.go` - Bookmark conversations
- `branded-output.go` - Branded output rendering
- `compact-notify.go` - Notification on compaction
- `confirm-destructive.go` - Confirm destructive operations
- `context-inject.go` - Inject context into conversations
- `conversation-manager.go` - **NEW** Tree navigation, branch summarization, and fresh context loops
- `custom-editor-demo.go` - Vim-like modal editor
- `dev-reload.go` - Development live-reload
- `header-footer-demo.go` - Custom headers and footers
- `inline-bash.go` - Inline bash execution
- `interactive-shell.go` - Interactive shell integration
- `kit-kit.go` - Kit-in-Kit (sub-agent spawning)
- `lsp-diagnostics.go` - LSP diagnostic integration
- `notify.go` - Desktop notifications
- `overlay-demo.go` - Modal dialogs
- `permission-gate.go` - Permission gating for tools
- `pirate.go` - Pirate-themed personality
- `plan-mode.go` - Read-only planning mode
- `project-rules.go` - Project-specific rules
- `prompt-demo.go` - Interactive prompts (select/confirm/input)
- `prompt-templates.go` - **NEW** Frontmatter-driven templates with model switching and skill injection
- `protected-paths.go` - Path protection for sensitive files
- `subagent-widget.go` - Multi-agent orchestration with status widget
- `subagent-test.go` - Subagent testing utilities
- `summarize.go` - Conversation summarization
- `tool-logger.go` - Log all tool calls
- `neon-theme.go` - Custom theme registration and switching
- `tool-renderer-demo.go` - Custom tool call rendering
- `widget-status.go` - Persistent status widgets
- [`minimal.go`](examples/extensions/minimal.go) - Clean UI with custom footer
- [`auto-commit.go`](examples/extensions/auto-commit.go) - Auto-commit on shutdown
- [`bookmark.go`](examples/extensions/bookmark.go) - Bookmark conversations
- [`branded-output.go`](examples/extensions/branded-output.go) - Branded output rendering
- [`bridge-demo.go`](examples/extensions/bridge_demo.go) - Bridged SDK API demo (tree navigation, skills, templates, model resolution)
- [`compact-notify.go`](examples/extensions/compact-notify.go) - Notification on compaction
- [`confirm-destructive.go`](examples/extensions/confirm-destructive.go) - Confirm destructive operations
- [`context-inject.go`](examples/extensions/context-inject.go) - Inject context into conversations
- [`conversation-manager.go`](examples/extensions/conversation-manager.go) - **NEW** Tree navigation, branch summarization, and fresh context loops
- [`custom-editor-demo.go`](examples/extensions/custom-editor-demo.go) - Vim-like modal editor
- [`dev-reload.go`](examples/extensions/dev-reload.go) - Development live-reload
- [`header-footer-demo.go`](examples/extensions/header-footer-demo.go) - Custom headers and footers
- [`inline-bash.go`](examples/extensions/inline-bash.go) - Inline bash execution
- [`interactive-shell.go`](examples/extensions/interactive-shell.go) - Interactive shell integration
- [`kit-kit.go`](examples/extensions/kit-kit.go) - Kit-in-Kit (sub-agent spawning)
- [`lsp-diagnostics.go`](examples/extensions/lsp-diagnostics.go) - LSP diagnostic integration
- [`notify.go`](examples/extensions/notify.go) - Desktop notifications
- [`overlay-demo.go`](examples/extensions/overlay-demo.go) - Modal dialogs
- [`permission-gate.go`](examples/extensions/permission-gate.go) - Permission gating for tools
- [`pirate.go`](examples/extensions/pirate.go) - Pirate-themed personality
- [`plan-mode.go`](examples/extensions/plan-mode.go) - Read-only planning mode
- [`project-rules.go`](examples/extensions/project-rules.go) - Project-specific rules
- [`prompt-demo.go`](examples/extensions/prompt-demo.go) - Interactive prompts (select/confirm/input)
- [`prompt-templates.go`](examples/extensions/prompt-templates.go) - **NEW** Frontmatter-driven templates with model switching and skill injection
- [`protected-paths.go`](examples/extensions/protected-paths.go) - Path protection for sensitive files
- [`subagent-widget.go`](examples/extensions/subagent-widget.go) - Multi-agent orchestration with status widget
- [`subagent-test.go`](examples/extensions/subagent-test.go) - Subagent testing utilities
- [`summarize.go`](examples/extensions/summarize.go) - Conversation summarization
- [`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
- [`widget-status.go`](examples/extensions/widget-status.go) - Persistent status widgets
Also see `.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.
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.
### Loading Extensions
@@ -406,7 +421,7 @@ func TestMyExtension(t *testing.T) {
- `AssertPrinted()`, `AssertPrintedContains()` — Verify output
- `AssertToolRegistered()`, `AssertCommandRegistered()` — Verify registration
See `examples/extensions/tool-logger_test.go` for a complete example with 14 test cases covering tool calls, input handling, and session lifecycle.
See [`examples/extensions/tool-logger_test.go`](examples/extensions/tool-logger_test.go) for a complete example with 14 test cases covering tool calls, input handling, and session lifecycle.
### Prompt Templates
@@ -428,10 +443,13 @@ Focus on $1 specifically.
**Argument placeholders:**
- `$1`, `$2`, etc. — Individual arguments
- `$@` or `$ARGUMENTS` — All arguments
- `$@` or `$ARGUMENTS` — All arguments (zero or more)
- `$+` — All arguments (one or more required; error if none given)
- `${@:2}` — Arguments from position 2 onwards
- `${@:1:3}` — 3 arguments starting at position 1
Placeholders inside fenced code blocks (```) and inline code spans are ignored.
Disable templates with `--no-prompt-templates` or load a specific template with `--prompt-template <name>`.
## Session Management
@@ -480,6 +498,15 @@ During an interactive session, use these slash commands:
| `/fork` | Fork to new session from an earlier message |
| `/new` | Start a fresh session |
### Keyboard Shortcuts
| Shortcut | Description |
|----------|-------------|
| `Ctrl+X e` | Open `$VISUAL`/`$EDITOR` to compose or edit your prompt |
| `Ctrl+X s` | Steer — inject a system-level instruction mid-turn |
| `ESC ESC` | Cancel the current operation (tool call or streaming) |
| `` / `` | Navigate prompt history |
## Go SDK
Embed Kit in your Go applications:
@@ -525,6 +552,20 @@ host, err := kit.New(ctx, &kit.Options{
Streaming: true,
Quiet: true,
// Generation parameters (override env/config/per-model defaults)
MaxTokens: 16384, // 0 = auto-resolve (env → config → per-model → 8192 floor)
ThinkingLevel: "medium", // "off", "none", "minimal", "low", "medium", "high"
Temperature: ptr(float32(0.2)), // pointer so 0.0 != unset; nil = provider default
TopP: nil, // nil = leave provider/per-model default
TopK: nil,
FrequencyPenalty: nil,
PresencePenalty: nil,
// Provider configuration (override env/config without reaching into viper)
ProviderAPIKey: "sk-...", // "" = use config / provider env var
ProviderURL: "https://proxy.internal/v1", // "" = provider default
TLSSkipVerify: false, // only takes effect when true
// Session options
SessionPath: "./session.jsonl", // Open specific session
Continue: true, // Resume most recent session
@@ -545,6 +586,46 @@ host, err := kit.New(ctx, &kit.Options{
})
```
**Generation & provider fields** (added in v0.55+) let SDK consumers configure
Kit entirely in-code without `viper.Set()` workarounds or shipping a `.kit.yml`.
Precedence is `Options` > `KIT_*` env vars > `.kit.yml` > per-model defaults
(`modelSettings` / `customModels`) > provider-level defaults. Sampling params
are pointer types so explicit `0.0` is distinguishable from "leave alone"; a
non-zero `MaxTokens` suppresses automatic right-sizing the same way `--max-tokens`
does on the CLI.
### MCP OAuth (remote MCP servers)
When a remote MCP server returns 401, Kit runs the full OAuth flow (dynamic
client registration → PKCE → token exchange → persistence) but delegates the
user-facing step — showing the authorization URL and receiving the callback —
to an `MCPAuthHandler` that you pass explicitly via `Options.MCPAuthHandler`.
If nil, OAuth is disabled and the authorization-required error surfaces to the
caller; the SDK never auto-opens a browser or binds a localhost port.
```go
// CLI/TUI apps: opens the system browser + prints status to stderr.
authHandler, _ := kit.NewCLIMCPAuthHandler()
defer authHandler.Close()
host, _ := kit.New(ctx, &kit.Options{
MCPAuthHandler: authHandler,
})
// Custom UX: reuse the SDK's port + callback server, supply your own
// presentation via OnAuthURL (TUI modal, QR code, web redirect, etc.).
// h, _ := kit.NewDefaultMCPAuthHandler()
// h.OnAuthURL = func(server, authURL string) { myUI.Show(server, authURL) }
//
// Full control (web apps, daemons): implement kit.MCPAuthHandler yourself —
// no localhost binding, no side effects.
```
Tokens are persisted to `$XDG_CONFIG_HOME/.kit/mcp_tokens.json` by default; swap
in a custom `MCPTokenStoreFactory` for encrypted, DB-backed, or in-memory
storage. See the [SDK options docs](/sdk/options#mcp-oauth-authorization) for
the full matrix.
### Custom Tools
Create custom tools with automatic schema generation — no external dependencies needed:
+64 -4
View File
@@ -11,6 +11,7 @@ import (
"charm.land/huh/v2"
"github.com/mark3labs/kit/internal/auth"
"github.com/mark3labs/kit/internal/ui"
kit "github.com/mark3labs/kit/pkg/kit"
"github.com/spf13/cobra"
)
@@ -54,9 +55,13 @@ Available providers:
- anthropic: Anthropic Claude API (OAuth)
- openai: OpenAI ChatGPT Plus/Pro (Codex OAuth)
Example:
Flags:
--set-default Set this provider's default model as the system default
Examples:
kit auth login anthropic
kit auth login openai`,
kit auth login openai
kit auth login openai --set-default`,
Args: cobra.ExactArgs(1),
RunE: runAuthLogin,
}
@@ -99,10 +104,43 @@ Example:
RunE: runAuthStatus,
}
var (
loginSetDefault bool
)
// defaultModels maps providers to their recommended default models.
// These are used when --set-default flag is passed to auth login.
var defaultModels = map[string]string{
"anthropic": "anthropic/claude-sonnet-4-5-20250929",
"openai": "openai/gpt-5.4",
}
// setDefaultModelIfRequested sets the default model for the given provider
// if the --set-default flag was provided.
func setDefaultModelIfRequested(provider string) error {
if !loginSetDefault {
return nil
}
model, ok := defaultModels[provider]
if !ok {
return fmt.Errorf("no default model configured for provider: %s", provider)
}
if err := ui.SaveModelPreference(model); err != nil {
return fmt.Errorf("failed to save model preference: %w", err)
}
fmt.Printf("\n✓ Set default model to: %s\n", model)
return nil
}
func init() {
authCmd.AddCommand(authLoginCmd)
authCmd.AddCommand(authLogoutCmd)
authCmd.AddCommand(authStatusCmd)
authLoginCmd.Flags().BoolVar(&loginSetDefault, "set-default", false, "Set this provider's default model as the system default after login")
}
func runAuthLogin(cmd *cobra.Command, args []string) error {
@@ -288,6 +326,17 @@ func loginAnthropic() error {
fmt.Println("\n🎉 Your OAuth credentials will now be used for Anthropic API calls.")
fmt.Println("💡 You can check your authentication status with: kit auth status")
// Set default model if requested
if err := setDefaultModelIfRequested("anthropic"); err != nil {
return err
}
// Remind users how to set this as default if they didn't use --set-default
if !loginSetDefault {
fmt.Println("\n💡 To set Anthropic as your default model, run:")
fmt.Println(" kit auth login anthropic --set-default")
}
return nil
}
@@ -454,6 +503,17 @@ func loginOpenAI() error {
fmt.Println("\n🎉 Your OAuth credentials will now be used for OpenAI API calls.")
fmt.Println("💡 You can check your authentication status with: kit auth status")
// Set default model if requested
if err := setDefaultModelIfRequested("openai"); err != nil {
return err
}
// Remind users how to set this as default if they didn't use --set-default
if !loginSetDefault {
fmt.Println("\n💡 To set OpenAI as your default model, run:")
fmt.Println(" kit auth login openai --set-default")
}
return nil
}
@@ -504,13 +564,13 @@ func startOpenAICallbackServer(expectedState string) (*callbackServer, error) {
}
// Return success page
w.Header().Set("Content-Type", "text/html")
w.Header().Set("Content-Type", "text/html; charset=utf-8")
w.WriteHeader(http.StatusOK)
_, _ = fmt.Fprintf(w, `<!DOCTYPE html>
<html>
<head><title>Authentication Successful</title></head>
<body style="font-family: sans-serif; text-align: center; padding: 50px;">
<h1> Authentication Successful</h1>
<h1>&#10003; Authentication Successful</h1>
<p>You can close this window and return to the terminal.</p>
</body>
</html>`)
+140 -30
View File
@@ -19,6 +19,7 @@ import (
"github.com/mark3labs/kit/internal/prompts"
"github.com/mark3labs/kit/internal/ui"
"github.com/mark3labs/kit/internal/ui/commands"
"github.com/mark3labs/kit/internal/ui/progress"
"github.com/mark3labs/kit/internal/watcher"
kit "github.com/mark3labs/kit/pkg/kit"
"github.com/spf13/cobra"
@@ -33,13 +34,18 @@ var (
providerURL string
providerAPIKey string
debugMode bool
positionalPrompt string // set by processPositionalArgs from CLI positional args
quietFlag bool
jsonFlag bool
noExitFlag bool
maxSteps int
streamFlag bool // Enable streaming output
autoCompactFlag bool // Enable auto-compaction near context limit
positionalPrompt string // set by processPositionalArgs from CLI positional args
positionalFiles []ui.FilePart // binary @file parts from processPositionalArgs
// MCP resource callbacks, set in runNormalMode, consumed by runInteractiveModeBubbleTea.
mcpGetResources func() []ui.FileSuggestion
mcpResourceReader ui.MCPResourceReader
quietFlag bool
jsonFlag bool
noExitFlag bool
maxSteps int
streamFlag bool // Enable streaming output
autoCompactFlag bool // Enable auto-compaction near context limit
// Session management
sessionPath string
@@ -291,14 +297,14 @@ func init() {
flags.BoolVar(&noPromptTemplates, "no-prompt-templates", false, "disable prompt template discovery")
// Model generation parameters
flags.IntVar(&maxTokens, "max-tokens", 4096, "maximum number of tokens in the response")
flags.IntVar(&maxTokens, "max-tokens", 8192, "maximum number of output tokens per response (auto-raised up to 32768 for models with higher known output limits; see internal/models/embedded_models.json)")
flags.Float32Var(&temperature, "temperature", 0.7, "controls randomness in responses (0.0-1.0)")
flags.Float32Var(&topP, "top-p", 0.95, "controls diversity via nucleus sampling (0.0-1.0)")
flags.Int32Var(&topK, "top-k", 40, "controls diversity by limiting top K tokens to sample from")
flags.Float32Var(&frequencyPenalty, "frequency-penalty", 0.0, "penalizes tokens based on frequency of appearance (0.0-2.0)")
flags.Float32Var(&presencePenalty, "presence-penalty", 0.0, "penalizes tokens based on whether they have appeared (0.0-2.0)")
flags.StringSliceVar(&stopSequences, "stop-sequences", nil, "custom stop sequences (comma-separated)")
flags.StringVar(&thinkingLevel, "thinking-level", "off", "extended thinking level: off, minimal, low, medium, high")
flags.StringVar(&thinkingLevel, "thinking-level", "off", "extended thinking level: off, none, minimal, low, medium, high")
// Ollama-specific parameters
flags.Int32Var(&numGPU, "num-gpu-layers", -1, "number of model layers to offload to GPU for Ollama models (-1 for auto-detect)")
@@ -338,12 +344,14 @@ func init() {
}
// processPositionalArgs separates positional CLI arguments into @file
// attachments and prompt text. File content is read and prepended to
// positionalPrompt so the agent receives it. Positional args are the primary
// way to run non-interactive mode:
// attachments and prompt text. Text file content is read and prepended to
// positionalPrompt; binary files (images, audio) are stored in positionalFiles
// for multimodal submission. Positional args are the primary way to run
// non-interactive mode:
//
// kit "Explain this codebase"
// kit @code.ts @test.ts "Review these files"
// kit @screenshot.png "What's in this image?"
func processPositionalArgs(args []string) {
cwd, err := os.Getwd()
if err != nil {
@@ -362,14 +370,17 @@ func processPositionalArgs(args []string) {
}
// Build file content prefix from @file arguments.
// Text files are XML-wrapped inline; binary files become multimodal parts.
var fileContent strings.Builder
for _, token := range fileTokens {
expanded := ui.ProcessFileAttachments(token, cwd)
if expanded != token {
// File was resolved — add it.
fileContent.WriteString(expanded)
result := ui.ProcessFileAttachments(token, cwd)
if result.ProcessedText != token {
// Text file was resolved — add it.
fileContent.WriteString(result.ProcessedText)
fileContent.WriteString("\n\n")
}
// Collect binary file parts for multimodal submission.
positionalFiles = append(positionalFiles, result.FileParts...)
}
// Combine: positional prompt text is appended to any existing --prompt
@@ -753,10 +764,11 @@ func runNormalMode(ctx context.Context) error {
}
},
CLI: &kit.CLIOptions{
MCPConfig: mcpConfig,
ShowSpinner: true,
SpinnerFunc: spinnerFunc,
UseBufferedLogger: true,
MCPConfig: mcpConfig,
ShowSpinner: true,
SpinnerFunc: spinnerFunc,
UseBufferedLogger: true,
ProgressReaderFunc: progress.NewProgressReadCloser,
},
}
if resumeFlag {
@@ -1704,6 +1716,81 @@ func runNormalMode(ctx context.Context) error {
return kitInstance.GetMCPToolCount()
}
// Build MCP prompt provider callbacks for the TUI.
// Convert kit.MCPPrompt → ui.MCPPromptInfo for the UI layer.
convertMCPPromptsForUI := func() []ui.MCPPromptInfo {
prompts := kitInstance.ListMCPPrompts()
if len(prompts) == 0 {
return nil
}
result := make([]ui.MCPPromptInfo, len(prompts))
for i, p := range prompts {
args := make([]ui.MCPPromptArgInfo, len(p.Arguments))
for j, a := range p.Arguments {
args[j] = ui.MCPPromptArgInfo{
Name: a.Name,
Description: a.Description,
Required: a.Required,
}
}
result[i] = ui.MCPPromptInfo{
Name: p.Name,
Description: p.Description,
Arguments: args,
ServerName: p.ServerName,
}
}
return result
}
mcpPrompts := convertMCPPromptsForUI()
getMCPPrompts := func() []ui.MCPPromptInfo {
return convertMCPPromptsForUI()
}
expandMCPPrompt := func(serverName, promptName string, args map[string]string) (*ui.MCPPromptExpandResult, error) {
result, err := kitInstance.GetMCPPrompt(context.Background(), serverName, promptName, args)
if err != nil {
return nil, err
}
msgs := make([]ui.MCPPromptMessageInfo, len(result.Messages))
for i, m := range result.Messages {
msgs[i] = ui.MCPPromptMessageInfo{
Role: m.Role,
Content: m.Content,
FileParts: m.FileParts,
}
}
return &ui.MCPPromptExpandResult{Messages: msgs}, nil
}
// MCP resource callbacks for @ autocomplete and submit-time resolution.
getMCPResources := func() []ui.FileSuggestion {
resources := kitInstance.ListMCPResources()
suggestions := make([]ui.FileSuggestion, len(resources))
for i, r := range resources {
suggestions[i] = ui.FileSuggestion{
RelPath: r.Name,
IsMCPResource: true,
MCPServerName: r.ServerName,
MCPResourceURI: r.URI,
MCPMIMEType: r.MIMEType,
Score: 100, // default score, filtered later
}
}
return suggestions
}
mcpResourceReaderFn := func(serverName, uri string) (string, []byte, string, bool, error) {
content, err := kitInstance.ReadMCPResource(context.Background(), serverName, uri)
if err != nil {
return "", nil, "", false, err
}
return content.Text, content.BlobData, content.MIMEType, content.IsBlob, nil
}
// Store MCP resource callbacks at package level for consumption by
// runInteractiveModeBubbleTea and runNonInteractiveModeApp.
mcpGetResources = getMCPResources
mcpResourceReader = mcpResourceReaderFn
// Start a goroutine that waits for background MCP tool loading to
// complete and notifies the TUI so it can refresh tool names and counts.
if len(mcpConfig.MCPServers) > 0 {
@@ -1840,7 +1927,7 @@ func runNormalMode(ctx context.Context) error {
// Check if running in non-interactive mode
if positionalPrompt != "" {
return runNonInteractiveModeApp(ctx, appInstance, cli, positionalPrompt, quietFlag, jsonFlag, noExitFlag, modelName, parsedProvider, kitInstance.GetLoadingMessage(), serverNames, toolNames, mcpToolCount, extensionToolCount, usageTracker, extCommands, promptTemplates, contextPaths, skillItems, getPromptTemplates, getSkillItems, getToolNames, getMCPToolCount, getWidgets, getHeader, getFooter, getToolRenderer, getEditorInterceptor, getUIVisibility, getStatusBarEntries, emitBeforeFork, emitBeforeSessionSwitch, getGlobalShortcuts, getExtensionCommands, setModelForUI, emitModelChangeForUI, kitInstance.IsReasoningModel(), kitInstance.GetThinkingLevel(), setThinkingLevelForUI, switchSessionForUI, reloadExtensionsForUI)
return runNonInteractiveModeApp(ctx, appInstance, cli, positionalPrompt, quietFlag, jsonFlag, noExitFlag, modelName, parsedProvider, kitInstance.GetLoadingMessage(), serverNames, toolNames, mcpToolCount, extensionToolCount, usageTracker, extCommands, promptTemplates, contextPaths, skillItems, getPromptTemplates, getSkillItems, getToolNames, getMCPToolCount, mcpPrompts, getMCPPrompts, expandMCPPrompt, getWidgets, getHeader, getFooter, getToolRenderer, getEditorInterceptor, getUIVisibility, getStatusBarEntries, emitBeforeFork, emitBeforeSessionSwitch, getGlobalShortcuts, getExtensionCommands, setModelForUI, emitModelChangeForUI, kitInstance.IsReasoningModel(), kitInstance.GetThinkingLevel(), setThinkingLevelForUI, switchSessionForUI, reloadExtensionsForUI)
}
// Quiet mode is not allowed in interactive mode
@@ -1848,7 +1935,7 @@ func runNormalMode(ctx context.Context) error {
return fmt.Errorf("--quiet requires a prompt")
}
return runInteractiveModeBubbleTea(ctx, appInstance, modelName, parsedProvider, kitInstance.GetLoadingMessage(), serverNames, toolNames, mcpToolCount, extensionToolCount, usageTracker, extCommands, promptTemplates, contextPaths, skillItems, getPromptTemplates, getSkillItems, getToolNames, getMCPToolCount, getWidgets, getHeader, getFooter, getToolRenderer, getEditorInterceptor, getUIVisibility, getStatusBarEntries, emitBeforeFork, emitBeforeSessionSwitch, getGlobalShortcuts, getExtensionCommands, setModelForUI, emitModelChangeForUI, kitInstance.IsReasoningModel(), kitInstance.GetThinkingLevel(), setThinkingLevelForUI, switchSessionForUI, reloadExtensionsForUI, startupExtensionMessages)
return runInteractiveModeBubbleTea(ctx, appInstance, modelName, parsedProvider, kitInstance.GetLoadingMessage(), serverNames, toolNames, mcpToolCount, extensionToolCount, usageTracker, extCommands, promptTemplates, contextPaths, skillItems, getPromptTemplates, getSkillItems, getToolNames, getMCPToolCount, mcpPrompts, getMCPPrompts, expandMCPPrompt, getWidgets, getHeader, getFooter, getToolRenderer, getEditorInterceptor, getUIVisibility, getStatusBarEntries, emitBeforeFork, emitBeforeSessionSwitch, getGlobalShortcuts, getExtensionCommands, setModelForUI, emitModelChangeForUI, kitInstance.IsReasoningModel(), kitInstance.GetThinkingLevel(), setThinkingLevelForUI, switchSessionForUI, reloadExtensionsForUI, startupExtensionMessages)
}
// runNonInteractiveModeApp executes a single prompt via the app layer and exits,
@@ -1861,15 +1948,33 @@ func runNormalMode(ctx context.Context) error {
//
// When --no-exit is set, after the prompt completes the interactive BubbleTea
// TUI is started so the user can continue the conversation.
func runNonInteractiveModeApp(ctx context.Context, appInstance *app.App, cli *ui.CLI, prompt string, quiet, jsonOutput, noExit bool, modelName, providerName, loadingMessage string, serverNames, toolNames []string, mcpToolCount, extensionToolCount int, usageTracker *ui.UsageTracker, extCommands []commands.ExtensionCommand, promptTemplates []*prompts.PromptTemplate, contextPaths []string, skillItems []ui.SkillItem, getPromptTemplates func() []*prompts.PromptTemplate, getSkillItems func() []ui.SkillItem, getToolNames func() []string, getMCPToolCount func() int, getWidgets func(string) []ui.WidgetData, getHeader, getFooter func() *ui.WidgetData, getToolRenderer func(string) *ui.ToolRendererData, getEditorInterceptor func() *ui.EditorInterceptor, getUIVisibility func() *ui.UIVisibility, getStatusBarEntries func() []ui.StatusBarEntryData, emitBeforeFork func(string, bool, string) (bool, string), emitBeforeSessionSwitch func(string) (bool, string), getGlobalShortcuts func() map[string]func(), getExtensionCommands func() []commands.ExtensionCommand, setModel func(string) error, emitModelChange func(string, string, string), isReasoningModel bool, thinkingLevel string, setThinkingLevel func(string) error, switchSession func(string) error, reloadExtensions func() error) error {
func runNonInteractiveModeApp(ctx context.Context, appInstance *app.App, cli *ui.CLI, prompt string, quiet, jsonOutput, noExit bool, modelName, providerName, loadingMessage string, serverNames, toolNames []string, mcpToolCount, extensionToolCount int, usageTracker *ui.UsageTracker, extCommands []commands.ExtensionCommand, promptTemplates []*prompts.PromptTemplate, contextPaths []string, skillItems []ui.SkillItem, getPromptTemplates func() []*prompts.PromptTemplate, getSkillItems func() []ui.SkillItem, getToolNames func() []string, getMCPToolCount func() int, mcpPrompts []ui.MCPPromptInfo, getMCPPrompts func() []ui.MCPPromptInfo, expandMCPPrompt func(string, string, map[string]string) (*ui.MCPPromptExpandResult, error), getWidgets func(string) []ui.WidgetData, getHeader, getFooter func() *ui.WidgetData, getToolRenderer func(string) *ui.ToolRendererData, getEditorInterceptor func() *ui.EditorInterceptor, getUIVisibility func() *ui.UIVisibility, getStatusBarEntries func() []ui.StatusBarEntryData, emitBeforeFork func(string, bool, string) (bool, string), emitBeforeSessionSwitch func(string) (bool, string), getGlobalShortcuts func() map[string]func(), getExtensionCommands func() []commands.ExtensionCommand, setModel func(string) error, emitModelChange func(string, string, string), isReasoningModel bool, thinkingLevel string, setThinkingLevel func(string) error, switchSession func(string) error, reloadExtensions func() error) error {
// Expand @file references in the prompt before sending to the agent.
// Text files are XML-inlined; binary files are extracted as multimodal parts.
var fileParts []kit.LLMFilePart
if cwd, err := os.Getwd(); err == nil {
prompt = ui.ProcessFileAttachments(prompt, cwd)
result := ui.ProcessFileAttachments(prompt, cwd, mcpResourceReader)
prompt = result.ProcessedText
for _, fp := range result.FileParts {
fileParts = append(fileParts, kit.LLMFilePart{
Filename: fp.Filename,
Data: fp.Data,
MediaType: fp.MediaType,
})
}
}
// Also include binary files from processPositionalArgs (CLI @file args).
for _, fp := range positionalFiles {
fileParts = append(fileParts, kit.LLMFilePart{
Filename: fp.Filename,
Data: fp.Data,
MediaType: fp.MediaType,
})
}
if jsonOutput {
// JSON mode: no intermediate display, structured JSON output.
result, err := appInstance.RunOnceResult(ctx, prompt)
result, err := appInstance.RunOnceResultWithFiles(ctx, prompt, fileParts)
if err != nil {
writeJSONError(err)
return err
@@ -1881,7 +1986,7 @@ func runNonInteractiveModeApp(ctx context.Context, appInstance *app.App, cli *ui
fmt.Println(string(data))
} else if quiet {
// Quiet mode: no intermediate display, just print final response.
if err := appInstance.RunOnce(ctx, prompt); err != nil {
if err := appInstance.RunOnceWithFiles(ctx, prompt, fileParts); err != nil {
return err
}
} else if cli != nil {
@@ -1890,21 +1995,21 @@ func runNonInteractiveModeApp(ctx context.Context, appInstance *app.App, cli *ui
// Route events through the shared CLI event handler.
eventHandler := ui.NewCLIEventHandler(cli, modelName)
err := appInstance.RunOnceWithDisplay(ctx, prompt, eventHandler.Handle)
err := appInstance.RunOnceWithDisplayAndFiles(ctx, prompt, eventHandler.Handle, fileParts)
eventHandler.Cleanup()
if err != nil {
return err
}
} else {
// No CLI available (shouldn't happen in non-quiet mode, but be safe).
if err := appInstance.RunOnce(ctx, prompt); err != nil {
if err := appInstance.RunOnceWithFiles(ctx, prompt, fileParts); err != nil {
return err
}
}
// If --no-exit was requested, hand off to the interactive TUI.
if noExit {
return runInteractiveModeBubbleTea(ctx, appInstance, modelName, providerName, loadingMessage, serverNames, toolNames, mcpToolCount, extensionToolCount, usageTracker, extCommands, promptTemplates, contextPaths, skillItems, getPromptTemplates, getSkillItems, getToolNames, getMCPToolCount, getWidgets, getHeader, getFooter, getToolRenderer, getEditorInterceptor, getUIVisibility, getStatusBarEntries, emitBeforeFork, emitBeforeSessionSwitch, getGlobalShortcuts, getExtensionCommands, setModel, emitModelChange, isReasoningModel, thinkingLevel, setThinkingLevel, switchSession, reloadExtensions, nil)
return runInteractiveModeBubbleTea(ctx, appInstance, modelName, providerName, loadingMessage, serverNames, toolNames, mcpToolCount, extensionToolCount, usageTracker, extCommands, promptTemplates, contextPaths, skillItems, getPromptTemplates, getSkillItems, getToolNames, getMCPToolCount, mcpPrompts, getMCPPrompts, expandMCPPrompt, getWidgets, getHeader, getFooter, getToolRenderer, getEditorInterceptor, getUIVisibility, getStatusBarEntries, emitBeforeFork, emitBeforeSessionSwitch, getGlobalShortcuts, getExtensionCommands, setModel, emitModelChange, isReasoningModel, thinkingLevel, setThinkingLevel, switchSession, reloadExtensions, nil)
}
return nil
@@ -2002,7 +2107,7 @@ func writeJSONError(err error) {
// 4. Calls program.Run() which blocks until the user quits (Ctrl+C or /quit).
//
// SetupCLI is not used for interactive mode; the TUI (AppModel) handles its own rendering.
func runInteractiveModeBubbleTea(_ context.Context, appInstance *app.App, modelName, providerName, loadingMessage string, serverNames, toolNames []string, mcpToolCount, extensionToolCount int, usageTracker *ui.UsageTracker, extCommands []commands.ExtensionCommand, promptTemplates []*prompts.PromptTemplate, contextPaths []string, skillItems []ui.SkillItem, getPromptTemplates func() []*prompts.PromptTemplate, getSkillItems func() []ui.SkillItem, getToolNames func() []string, getMCPToolCount func() int, getWidgets func(string) []ui.WidgetData, getHeader, getFooter func() *ui.WidgetData, getToolRenderer func(string) *ui.ToolRendererData, getEditorInterceptor func() *ui.EditorInterceptor, getUIVisibility func() *ui.UIVisibility, getStatusBarEntries func() []ui.StatusBarEntryData, emitBeforeFork func(string, bool, string) (bool, string), emitBeforeSessionSwitch func(string) (bool, string), getGlobalShortcuts func() map[string]func(), getExtensionCommands func() []commands.ExtensionCommand, setModel func(string) error, emitModelChange func(string, string, string), isReasoningModel bool, thinkingLevel string, setThinkingLevel func(string) error, switchSession func(string) error, reloadExtensions func() error, startupExtensionMessages []string) error {
func runInteractiveModeBubbleTea(_ context.Context, appInstance *app.App, modelName, providerName, loadingMessage string, serverNames, toolNames []string, mcpToolCount, extensionToolCount int, usageTracker *ui.UsageTracker, extCommands []commands.ExtensionCommand, promptTemplates []*prompts.PromptTemplate, contextPaths []string, skillItems []ui.SkillItem, getPromptTemplates func() []*prompts.PromptTemplate, getSkillItems func() []ui.SkillItem, getToolNames func() []string, getMCPToolCount func() int, mcpPrompts []ui.MCPPromptInfo, getMCPPrompts func() []ui.MCPPromptInfo, expandMCPPrompt func(string, string, map[string]string) (*ui.MCPPromptExpandResult, error), getWidgets func(string) []ui.WidgetData, getHeader, getFooter func() *ui.WidgetData, getToolRenderer func(string) *ui.ToolRendererData, getEditorInterceptor func() *ui.EditorInterceptor, getUIVisibility func() *ui.UIVisibility, getStatusBarEntries func() []ui.StatusBarEntryData, emitBeforeFork func(string, bool, string) (bool, string), emitBeforeSessionSwitch func(string) (bool, string), getGlobalShortcuts func() map[string]func(), getExtensionCommands func() []commands.ExtensionCommand, setModel func(string) error, emitModelChange func(string, string, string), isReasoningModel bool, thinkingLevel string, setThinkingLevel func(string) error, switchSession func(string) error, reloadExtensions func() error, startupExtensionMessages []string) error {
// Redirect all log output (stdlib and charm) to a file so that log
// messages don't write to stderr and corrupt the TUI. Bubble Tea
// captures stdout for rendering; any stray stderr output from
@@ -2041,6 +2146,9 @@ func runInteractiveModeBubbleTea(_ context.Context, appInstance *app.App, modelN
ExtensionCommands: extCommands,
PromptTemplates: promptTemplates,
GetPromptTemplates: getPromptTemplates,
MCPPrompts: mcpPrompts,
GetMCPPrompts: getMCPPrompts,
ExpandMCPPrompt: expandMCPPrompt,
ContextPaths: contextPaths,
SkillItems: skillItems,
GetSkillItems: getSkillItems,
@@ -2064,6 +2172,8 @@ func runInteractiveModeBubbleTea(_ context.Context, appInstance *app.App, modelN
SwitchSession: switchSession,
ReloadExtensions: reloadExtensions,
ShowSessionPicker: resumeFlag,
GetMCPResources: mcpGetResources,
MCPResourceReader: mcpResourceReader,
})
program := tea.NewProgram(appModel)
+57 -19
View File
@@ -10,13 +10,21 @@ import (
"kit/ext"
)
// re matches !{...} with non-greedy content.
var re = regexp.MustCompile(`!\{([^}]+)\}`)
// Init expands inline bash expressions in user prompts before they reach the
// LLM. Text like !{git branch --show-current} is replaced with the command's
// stdout.
// LLM. Text like !{git rev-parse --abbrev-ref HEAD} is replaced with the
// command's stdout.
//
// In interactive mode the expansion happens at submit time via an editor
// interceptor, so the expanded text is also visible in the user message
// block on screen. In non-interactive mode (CLI, script, queue) the
// expansion happens via OnInput transform.
//
// Examples:
//
// "Fix the tests on !{git branch --show-current}"
// "Fix the tests on !{git rev-parse --abbrev-ref HEAD}"
// → "Fix the tests on main"
//
// "The current directory is !{pwd}"
@@ -24,29 +32,59 @@ import (
//
// Usage: kit -e examples/extensions/inline-bash.go
func Init(api ext.API) {
// Matches !{...} with non-greedy content.
re := regexp.MustCompile(`!\{([^}]+)\}`)
// ── Interactive mode: editor interceptor ──────────────────────────
// Intercept Enter / Ctrl+D so we can expand !{...} BEFORE the
// SubmitMsg is created. This ensures the expanded text appears in
// the user message block on screen as well as in the LLM prompt.
api.OnSessionStart(func(_ ext.SessionStartEvent, ctx ext.Context) {
if !ctx.Interactive {
return
}
ctx.SetEditor(ext.EditorConfig{
HandleKey: func(key string, currentText string) ext.EditorKeyAction {
if (key == "enter" || key == "ctrl+d") && re.MatchString(currentText) {
expanded := expand(currentText)
// Clear the textarea asynchronously — calling
// SetEditorText synchronously from inside Update()
// would deadlock the BubbleTea event loop.
go ctx.SetEditorText("")
return ext.EditorKeyAction{
Type: ext.EditorKeySubmit,
SubmitText: expanded,
}
}
return ext.EditorKeyAction{Type: ext.EditorKeyPassthrough}
},
})
})
// ── Non-interactive fallback: OnInput transform ──────────────────
// For CLI, script, and queue sources the editor interceptor is not
// active, so we fall back to OnInput which still rewrites the
// prompt text sent to the LLM.
api.OnInput(func(ev ext.InputEvent, ctx ext.Context) *ext.InputResult {
if !re.MatchString(ev.Text) {
if ev.Source == "interactive" || !re.MatchString(ev.Text) {
return nil
}
expanded := re.ReplaceAllStringFunc(ev.Text, func(match string) string {
// Extract the command between !{ and }.
cmd := re.FindStringSubmatch(match)[1]
cmd = strings.TrimSpace(cmd)
out, err := exec.Command("bash", "-c", cmd).Output()
if err != nil {
return match // keep original on error
}
return strings.TrimSpace(string(out))
})
return &ext.InputResult{
Action: "transform",
Text: expanded,
Text: expand(ev.Text),
}
})
}
// expand replaces every !{cmd} in text with the command's stdout.
// On error the original !{cmd} token is preserved.
func expand(text string) string {
return re.ReplaceAllStringFunc(text, func(match string) string {
cmd := re.FindStringSubmatch(match)[1]
cmd = strings.TrimSpace(cmd)
out, err := exec.Command("bash", "-c", cmd).Output()
if err != nil {
return match // keep original on error
}
return strings.TrimSpace(string(out))
})
}
@@ -130,6 +130,58 @@ func TestSubagentMonitor_MultipleSubagents(t *testing.T) {
time.Sleep(100 * time.Millisecond)
}
// TestSubagentMonitor_ConcurrentSubagents verifies no panics when multiple
// subagents emit events concurrently from different goroutines.
func TestSubagentMonitor_ConcurrentSubagents(t *testing.T) {
harness := test.New(t)
harness.LoadFile("../../.kit/extensions/subagent-monitor.go")
_, err := harness.Emit(extensions.SessionStartEvent{SessionID: "test-session"})
if err != nil {
t.Fatalf("SessionStart should not error: %v", err)
}
// Start 5 subagents concurrently
done := make(chan struct{}, 5)
for i := range 5 {
go func(idx int) {
defer func() { done <- struct{}{} }()
callID := fmt.Sprintf("concurrent-%d", idx)
task := fmt.Sprintf("concurrent task %d", idx)
_, _ = harness.Emit(extensions.SubagentStartEvent{
ToolCallID: callID,
Task: task,
})
// Emit many chunks rapidly
for j := range 20 {
_, _ = harness.Emit(extensions.SubagentChunkEvent{
ToolCallID: callID,
Task: task,
ChunkType: "text",
Content: fmt.Sprintf("agent %d chunk %d", idx, j),
})
}
_, _ = harness.Emit(extensions.SubagentEndEvent{
ToolCallID: callID,
Task: task,
Response: "done",
})
}(i)
}
// Wait for all goroutines
for range 5 {
<-done
}
// Allow any final processing
time.Sleep(200 * time.Millisecond)
}
// TestSubagentMonitor_SessionShutdown verifies shutdown doesn't panic
// even with nil ctx functions.
func TestSubagentMonitor_SessionShutdown(t *testing.T) {
+153
View File
@@ -0,0 +1,153 @@
//go:build ignore
// sudo-handler.go - Extension to handle sudo password prompts securely
//
// This extension intercepts bash commands containing "sudo" and:
// 1. Checks if sudo credentials are already cached (via sudo -n)
// 2. If not cached, prompts the user for their password (with masking)
// 3. Temporarily sets SUDO_PASSWORD environment variable for execution
// 4. The bash tool automatically uses sudo -S -p '' to pipe the password
//
// Usage: kit -e examples/extensions/sudo-handler.go
//
// Security notes:
// - Password is only stored in memory for the duration of the session
// - Password is never logged or displayed
// - Each session requires re-authentication (sudo -k is used)
// - The SUDO_PASSWORD env var is set only during tool execution
package main
import (
"encoding/json"
"os"
"strings"
"sync"
"kit/ext"
)
var (
// cachedPassword stores the sudo password for the session
cachedPassword string
// hasCachedPassword tracks if we have a valid cached password
hasCachedPassword bool
// mu protects cached password access
mu sync.RWMutex
)
// Init sets up the sudo handler extension
func Init(api ext.API) {
api.OnToolCall(func(tc ext.ToolCallEvent, ctx ext.Context) *ext.ToolCallResult {
if tc.ToolName != "bash" {
return nil
}
// Parse the command from tool input
var input struct {
Command string `json:"command"`
}
if err := json.Unmarshal([]byte(tc.Input), &input); err != nil {
return nil
}
// Check if command contains sudo
if !containsSudo(input.Command) {
return nil
}
// Check if we already have cached credentials
mu.RLock()
password := cachedPassword
hasCached := hasCachedPassword
mu.RUnlock()
if hasCached {
// Use cached password
os.Setenv("SUDO_PASSWORD", password)
return nil
}
// No cached password - prompt user
result := ctx.PromptInput(ext.PromptInputConfig{
Message: "🔐 Sudo password required for:\n " + truncateCommand(input.Command, 60),
Placeholder: "Enter your password",
})
if result.Cancelled {
return &ext.ToolCallResult{
Block: true,
Reason: "Sudo password prompt cancelled by user",
}
}
if result.Value == "" {
return &ext.ToolCallResult{
Block: true,
Reason: "No password provided",
}
}
// Cache the password for this session
mu.Lock()
cachedPassword = result.Value
hasCachedPassword = true
mu.Unlock()
// Set environment variable for the bash tool to use
os.Setenv("SUDO_PASSWORD", result.Value)
// Show confirmation (without revealing password)
ctx.PrintInfo("Sudo password cached for this session")
return nil
})
// Clear cached password when session ends
api.OnSessionShutdown(func(event ext.SessionShutdownEvent, ctx ext.Context) {
mu.Lock()
cachedPassword = ""
hasCachedPassword = false
mu.Unlock()
os.Unsetenv("SUDO_PASSWORD")
})
}
// containsSudo checks if the command contains sudo as a command (not in a string)
func containsSudo(command string) bool {
// Simple check for sudo as a word, not inside quotes or as part of another word
lower := strings.ToLower(command)
// Check for sudo at start or after separators
patterns := []string{
"sudo ",
"sudo\t",
";sudo ",
"&& sudo ",
"|| sudo ",
"| sudo ",
"$(sudo ",
"`sudo ",
}
for _, pattern := range patterns {
if strings.Contains(lower, pattern) {
return true
}
}
// Check if command starts with sudo
if strings.HasPrefix(lower, "sudo ") {
return true
}
return false
}
// truncateCommand truncates a long command for display
func truncateCommand(cmd string, maxLen int) string {
if len(cmd) <= maxLen {
return cmd
}
return cmd[:maxLen-3] + "..."
}
+19 -18
View File
@@ -1,31 +1,32 @@
module github.com/mark3labs/kit
go 1.26.1
go 1.26.2
require (
charm.land/bubbles/v2 v2.1.0
charm.land/bubbletea/v2 v2.0.2
charm.land/fantasy v0.17.1
charm.land/bubbletea/v2 v2.0.5
charm.land/fantasy v0.17.2
charm.land/huh/v2 v2.0.3
charm.land/lipgloss/v2 v2.0.2
charm.land/lipgloss/v2 v2.0.3
github.com/alecthomas/chroma/v2 v2.23.1
github.com/atotto/clipboard v0.1.4
github.com/aymanbagabas/go-udiff v0.4.1
github.com/charmbracelet/fang v1.0.0
github.com/charmbracelet/log v1.0.0
github.com/charmbracelet/openai-go v0.0.0-20260319145158-d0740cc34266
github.com/charmbracelet/ultraviolet v0.0.0-20260330092749-0f94982c930b
github.com/charmbracelet/ultraviolet v0.0.0-20260414011438-8c69ec811b1e
github.com/charmbracelet/x/editor v0.2.0
github.com/clipperhouse/displaywidth v0.11.0
github.com/clipperhouse/uax29/v2 v2.7.0
github.com/coder/acp-go-sdk v0.6.3
github.com/fsnotify/fsnotify v1.9.0
github.com/indaco/herald v0.13.0
github.com/indaco/herald-md v0.3.0
github.com/mark3labs/mcp-go v0.47.1
github.com/mark3labs/mcp-go v0.48.0
github.com/spf13/cobra v1.10.2
github.com/spf13/viper v1.21.0
github.com/traefik/yaegi v0.16.1
golang.org/x/term v0.41.0
golang.org/x/term v0.42.0
gopkg.in/yaml.v3 v3.0.1
)
@@ -58,9 +59,9 @@ require (
github.com/charmbracelet/harmonica v0.2.0 // indirect
github.com/charmbracelet/lipgloss v1.1.1-0.20250404203927-76690c660834 // indirect
github.com/charmbracelet/x/cellbuf v0.0.15 // indirect
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260406091427-a791e22d5143 // indirect
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260413165052-6921c759c913 // indirect
github.com/charmbracelet/x/exp/ordered v0.1.0 // indirect
github.com/charmbracelet/x/exp/slice v0.0.0-20260406091427-a791e22d5143 // indirect
github.com/charmbracelet/x/exp/slice v0.0.0-20260413165052-6921c759c913 // indirect
github.com/charmbracelet/x/exp/strings v0.1.0 // indirect
github.com/charmbracelet/x/json v0.2.0 // indirect
github.com/charmbracelet/x/termios v0.1.1 // indirect
@@ -81,10 +82,10 @@ require (
github.com/googleapis/enterprise-certificate-proxy v0.3.14 // indirect
github.com/googleapis/gax-go/v2 v2.21.0 // indirect
github.com/gorilla/websocket v1.5.3 // indirect
github.com/kaptinlin/go-i18n v0.3.1 // indirect
github.com/kaptinlin/go-i18n v0.4.0 // indirect
github.com/kaptinlin/jsonpointer v0.4.17 // indirect
github.com/kaptinlin/jsonschema v0.7.7 // indirect
github.com/kaptinlin/messageformat-go v0.4.19 // indirect
github.com/kaptinlin/messageformat-go v0.4.20 // indirect
github.com/mitchellh/hashstructure/v2 v2.0.2 // indirect
github.com/muesli/mango v0.2.0 // indirect
github.com/muesli/mango-cobra v1.3.0 // indirect
@@ -109,14 +110,14 @@ require (
go.opentelemetry.io/otel/metric v1.43.0 // indirect
go.opentelemetry.io/otel/trace v1.43.0 // indirect
go.yaml.in/yaml/v3 v3.0.4 // indirect
golang.org/x/crypto v0.49.0 // indirect
golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90 // indirect
golang.org/x/net v0.52.0 // indirect
golang.org/x/crypto v0.50.0 // indirect
golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f // indirect
golang.org/x/net v0.53.0 // indirect
golang.org/x/oauth2 v0.36.0 // indirect
golang.org/x/time v0.15.0 // indirect
google.golang.org/api v0.275.0 // indirect
google.golang.org/genai v1.52.1 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d // indirect
google.golang.org/genai v1.54.0 // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 // indirect
google.golang.org/grpc v1.80.0 // indirect
google.golang.org/protobuf v1.36.11 // indirect
gopkg.in/yaml.v2 v2.4.0 // indirect
@@ -124,7 +125,7 @@ require (
require (
github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect
github.com/charmbracelet/x/ansi v0.11.6
github.com/charmbracelet/x/ansi v0.11.7
github.com/charmbracelet/x/term v0.2.2 // indirect
github.com/inconshreveable/mousetrap v1.1.0 // indirect
github.com/lucasb-eyer/go-colorful v1.4.0 // indirect
@@ -136,5 +137,5 @@ require (
github.com/spf13/pflag v1.0.10 // indirect
golang.org/x/sync v0.20.0 // indirect
golang.org/x/sys v0.43.0 // indirect
golang.org/x/text v0.35.0
golang.org/x/text v0.36.0
)
+40 -38
View File
@@ -1,13 +1,13 @@
charm.land/bubbles/v2 v2.1.0 h1:YSnNh5cPYlYjPxRrzs5VEn3vwhtEn3jVGRBT3M7/I0g=
charm.land/bubbles/v2 v2.1.0/go.mod h1:l97h4hym2hvWBVfmJDtrEHHCtkIKeTEb3TTJ4ZOB3wY=
charm.land/bubbletea/v2 v2.0.2 h1:4CRtRnuZOdFDTWSff9r8QFt/9+z6Emubz3aDMnf/dx0=
charm.land/bubbletea/v2 v2.0.2/go.mod h1:3LRff2U4WIYXy7MTxfbAQ+AdfM3D8Xuvz2wbsOD9OHQ=
charm.land/fantasy v0.17.1 h1:SQzfnyJPDuQWt6e//KKmQmEEXdqHMC0IZz10XwkLcEM=
charm.land/fantasy v0.17.1/go.mod h1:FF5ALCCHETacHJPBqU42CtwMInYQ0ul52fdzIHQMbQk=
charm.land/bubbletea/v2 v2.0.5 h1:TQlLFqxo39AAHSVuOhJ5D3nH7O9Nk8JGinsfWQ4y1U4=
charm.land/bubbletea/v2 v2.0.5/go.mod h1:dvbsYZD+MHkdIZl+Z67D212hEvB+GII2tfH8f9SnoDw=
charm.land/fantasy v0.17.2 h1:ojTMufMxY/PVH7TzYUxht2SVkvD90iCTJfmPR6c8BR8=
charm.land/fantasy v0.17.2/go.mod h1:V9cCIUMZB9g3Bq40aKEY8xBNzDd48EdfHp2OMS0uzWs=
charm.land/huh/v2 v2.0.3 h1:2cJsMqEPwSywGHvdlKsJyQKPtSJLVnFKyFbsYZTlLkU=
charm.land/huh/v2 v2.0.3/go.mod h1:93eEveeeqn47MwiC3tf+2atZ2l7Is88rAtmZNZ8x9Wc=
charm.land/lipgloss/v2 v2.0.2 h1:xFolbF8JdpNkM2cEPTfXEcW1p6NRzOWTSamRfYEw8cs=
charm.land/lipgloss/v2 v2.0.2/go.mod h1:KjPle2Qd3YmvP1KL5OMHiHysGcNwq6u83MUjYkFvEkM=
charm.land/lipgloss/v2 v2.0.3 h1:yM2zJ4Cf5Y51b7RHIwioil4ApI/aypFXXVHSwlM6RzU=
charm.land/lipgloss/v2 v2.0.3/go.mod h1:7myLU9iG/3xluAWzpY/fSxYYHCgoKTie7laxk6ATwXA=
cloud.google.com/go v0.123.0 h1:2NAUJwPR47q+E35uaJeYoNhuNEM9kM8SjgRgdeOJUSE=
cloud.google.com/go v0.123.0/go.mod h1:xBoMV08QcqUGuPW65Qfm1o9Y4zKZBpGS+7bImXLTAZU=
cloud.google.com/go/auth v0.20.0 h1:kXTssoVb4azsVDoUiF8KvxAqrsQcQtB53DcSgta74CA=
@@ -86,24 +86,26 @@ github.com/charmbracelet/log v1.0.0 h1:HVVVMmfOorfj3BA9i8X8UL69Hoz9lI0PYwXfJvOdR
github.com/charmbracelet/log v1.0.0/go.mod h1:uYgY3SmLpwJWxmlrPwXvzVYujxis1vAKRV/0VQB7yWA=
github.com/charmbracelet/openai-go v0.0.0-20260319145158-d0740cc34266 h1:BW/sZtyd1JyYy0h5adMm3tzpNyL857LWjuTRET6OhpY=
github.com/charmbracelet/openai-go v0.0.0-20260319145158-d0740cc34266/go.mod h1:1DahUaExbUZx/jD+FNT2PKP4L9rLE5+ZBRuI8mZjd/E=
github.com/charmbracelet/ultraviolet v0.0.0-20260330092749-0f94982c930b h1:ASDO9RT6SNKTQN87jO2bRfxHFJq8cgeYdFzivY2gCeM=
github.com/charmbracelet/ultraviolet v0.0.0-20260330092749-0f94982c930b/go.mod h1:Vo8TffMf0q7Uho/n8e6XpBZvOWtd3g39yX+9P5rRutA=
github.com/charmbracelet/x/ansi v0.11.6 h1:GhV21SiDz/45W9AnV2R61xZMRri5NlLnl6CVF7ihZW8=
github.com/charmbracelet/x/ansi v0.11.6/go.mod h1:2JNYLgQUsyqaiLovhU2Rv/pb8r6ydXKS3NIttu3VGZQ=
github.com/charmbracelet/ultraviolet v0.0.0-20260414011438-8c69ec811b1e h1:O5hZFj55wZQWxMiRtQLa3uLKhZGZGS/j8M3OXinQlrw=
github.com/charmbracelet/ultraviolet v0.0.0-20260414011438-8c69ec811b1e/go.mod h1:bAAz7dh/FTYfC+oiHavL4mX1tOIBZ0ZwYjSi3qE6ivM=
github.com/charmbracelet/x/ansi v0.11.7 h1:kzv1kJvjg2S3r9KHo8hDdHFQLEqn4RBCb39dAYC84jI=
github.com/charmbracelet/x/ansi v0.11.7/go.mod h1:9qGpnAVYz+8ACONkZBUWPtL7lulP9No6p1epAihUZwQ=
github.com/charmbracelet/x/cellbuf v0.0.15 h1:ur3pZy0o6z/R7EylET877CBxaiE1Sp1GMxoFPAIztPI=
github.com/charmbracelet/x/cellbuf v0.0.15/go.mod h1:J1YVbR7MUuEGIFPCaaZ96KDl5NoS0DAWkskup+mOY+Q=
github.com/charmbracelet/x/conpty v0.1.1 h1:s1bUxjoi7EpqiXysVtC+a8RrvPPNcNvAjfi4jxsAuEs=
github.com/charmbracelet/x/conpty v0.1.1/go.mod h1:OmtR77VODEFbiTzGE9G1XiRJAga6011PIm4u5fTNZpk=
github.com/charmbracelet/x/editor v0.2.0 h1:7XLUKtaRaB8jN7bWU2p2UChiySyaAuIfYiIRg8gGWwk=
github.com/charmbracelet/x/editor v0.2.0/go.mod h1:p3oQ28TSL3YPd+GKJ1fHWcp+7bVGpedHpXmo0D6t1dY=
github.com/charmbracelet/x/errors v0.0.0-20240508181413-e8d8b6e2de86 h1:JSt3B+U9iqk37QUU2Rvb6DSBYRLtWqFqfxf8l5hOZUA=
github.com/charmbracelet/x/errors v0.0.0-20240508181413-e8d8b6e2de86/go.mod h1:2P0UgXMEa6TsToMSuFqKFQR+fZTO9CNGUNokkPatT/0=
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260406091427-a791e22d5143 h1:zmBor0ftFNqVFp9U59ZoEDRUCIYSGOGSIfGGkNZRufs=
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260406091427-a791e22d5143/go.mod h1:nsExn0DGyX0lh9LwLHTn2Gg+hafdzfSXnC+QmEJTZFY=
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260413165052-6921c759c913 h1:6F/6bu5nBLjodsvaU5xAszTaxtHrDU5UiJarpMPZj48=
github.com/charmbracelet/x/exp/charmtone v0.0.0-20260413165052-6921c759c913/go.mod h1:nsExn0DGyX0lh9LwLHTn2Gg+hafdzfSXnC+QmEJTZFY=
github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f h1:pk6gmGpCE7F3FcjaOEKYriCvpmIN4+6OS/RD0vm4uIA=
github.com/charmbracelet/x/exp/golden v0.0.0-20250806222409-83e3a29d542f/go.mod h1:IfZAMTHB6XkZSeXUqriemErjAWCCzT0LwjKFYCZyw0I=
github.com/charmbracelet/x/exp/ordered v0.1.0 h1:55/qLwjIh0gL0Vni+QAWk7T/qRVP6sBf+2agPBgnOFE=
github.com/charmbracelet/x/exp/ordered v0.1.0/go.mod h1:5UHwmG+is5THxMyCJHNPCn2/ecI07aKNrW+LcResjJ8=
github.com/charmbracelet/x/exp/slice v0.0.0-20260406091427-a791e22d5143 h1:aEppolah2k9c0LzKX2fk5ryuyQ0Lq8kCOjkvMw1b8o4=
github.com/charmbracelet/x/exp/slice v0.0.0-20260406091427-a791e22d5143/go.mod h1:vqEfX6xzqW1pKKZUUiFOKg0OQ7bCh54Q2vR/tserrRA=
github.com/charmbracelet/x/exp/slice v0.0.0-20260413165052-6921c759c913 h1:RiZFY92Ug9iz1CenzxSSQla2Z3WflsR7bIuXq40JlpU=
github.com/charmbracelet/x/exp/slice v0.0.0-20260413165052-6921c759c913/go.mod h1:vqEfX6xzqW1pKKZUUiFOKg0OQ7bCh54Q2vR/tserrRA=
github.com/charmbracelet/x/exp/strings v0.1.0 h1:i69S2XI7uG1u4NLGeJPSYU++Nmjvpo9nwd6aoEm7gkA=
github.com/charmbracelet/x/exp/strings v0.1.0/go.mod h1:/ehtMPNh9K4odGFkqYJKpIYyePhdp1hLBRvyY4bWkH8=
github.com/charmbracelet/x/json v0.2.0 h1:DqB+ZGx2h+Z+1s98HOuOyli+i97wsFQIxP2ZQANTPrQ=
@@ -185,14 +187,14 @@ github.com/indaco/herald v0.13.0 h1:+xVG9Fx5NpuWhwku/9IlRL6I009NnX4VUGKvlZHTRxU=
github.com/indaco/herald v0.13.0/go.mod h1:T5g1+XLYvpjouhzAGHnAHDCKizhESkoV6+QPZ3DhgWA=
github.com/indaco/herald-md v0.3.0 h1:hN1cKyrexPPM9PeHBsKuaWvIizSi/iYvM9yzRgtdb8M=
github.com/indaco/herald-md v0.3.0/go.mod h1:RUHVaDSG45ymJjKyxpDwBocLXrZo93FB4OeYMsw9B9s=
github.com/kaptinlin/go-i18n v0.3.1 h1:plXi3XQE1aYamFi8TU0K6actODmw2+5FSobmhTkfQ/0=
github.com/kaptinlin/go-i18n v0.3.1/go.mod h1:ZRoAHj7elWYamfbv7wev7Ajch6LOzjtBaq8nWe8HIVk=
github.com/kaptinlin/go-i18n v0.4.0 h1:i7L3U2yurg+xhokITtJ0k+mjHnXqkoyz8ju5Wb7W8Oc=
github.com/kaptinlin/go-i18n v0.4.0/go.mod h1:njA6x0+4MWGcLWT0KLrwekhRPmze1Hnstf2+VJFzwpM=
github.com/kaptinlin/jsonpointer v0.4.17 h1:mY9k8ciWncxbsECyaxKnR0MdmxamNdp2tLQkAKVrtSk=
github.com/kaptinlin/jsonpointer v0.4.17/go.mod h1:SsfsjqnHG5zuKo1DTBzk1VknaHlL4osHw+X9kZKukpU=
github.com/kaptinlin/jsonschema v0.7.7 h1:41BlQJ9dskH0oE5DSzBUrl/w4JQYIr6N6L0B5GNyDoM=
github.com/kaptinlin/jsonschema v0.7.7/go.mod h1:rKjWfyySHSxAD7Li2ctYkPlOu960igoKBvZ2ADRtd5Q=
github.com/kaptinlin/messageformat-go v0.4.19 h1:A5kuuZ1ybXDQ7kD1aoEWGAOemX7hLsMY0yolgSbgpRI=
github.com/kaptinlin/messageformat-go v0.4.19/go.mod h1:utSDTfiXTxl66OC5RIEuObLH7Ue3YjbA2X86SYMBYWg=
github.com/kaptinlin/messageformat-go v0.4.20 h1:a0ufTd5liiUubIGeGxpSTnNS8ZSrN4DV01/wGFmfzMs=
github.com/kaptinlin/messageformat-go v0.4.20/go.mod h1:FqdEPfQLkqVBX7OBRMPgYwUPvKYJohFD9Ok1BMzCfIo=
github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk=
github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
@@ -201,8 +203,8 @@ github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
github.com/lucasb-eyer/go-colorful v1.4.0 h1:UtrWVfLdarDgc44HcS7pYloGHJUjHV/4FwW4TvVgFr4=
github.com/lucasb-eyer/go-colorful v1.4.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
github.com/mark3labs/mcp-go v0.47.1 h1:A9sJJ20mscl/ssLYHjodfaoBmq6uuhMG7pAPNYaQymQ=
github.com/mark3labs/mcp-go v0.47.1/go.mod h1:JKTC7R2LLVagkEWK7Kwu7DbmA6iIvnNAod6yrHiQMag=
github.com/mark3labs/mcp-go v0.48.0 h1:o+MXuGW/HCeR2ny5LcAcZQn2bo6I2xaZMEHnpRG+dtw=
github.com/mark3labs/mcp-go v0.48.0/go.mod h1:JKTC7R2LLVagkEWK7Kwu7DbmA6iIvnNAod6yrHiQMag=
github.com/mattn/go-isatty v0.0.21 h1:xYae+lCNBP7QuW4PUnNG61ffM4hVIfm+zUzDuSzYLGs=
github.com/mattn/go-isatty v0.0.21/go.mod h1:ZXfXG4SQHsB/w3ZeOYbR0PrPwLy+n6xiMrJlRFqopa4=
github.com/mattn/go-runewidth v0.0.23 h1:7ykA0T0jkPpzSvMS5i9uoNn2Xy3R383f9HDx3RybWcw=
@@ -288,36 +290,36 @@ go.opentelemetry.io/otel/trace v1.43.0 h1:BkNrHpup+4k4w+ZZ86CZoHHEkohws8AY+WTX09
go.opentelemetry.io/otel/trace v1.43.0/go.mod h1:/QJhyVBUUswCphDVxq+8mld+AvhXZLhe+8WVFxiFff0=
go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc=
go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
golang.org/x/crypto v0.49.0 h1:+Ng2ULVvLHnJ/ZFEq4KdcDd/cfjrrjjNSXNzxg0Y4U4=
golang.org/x/crypto v0.49.0/go.mod h1:ErX4dUh2UM+CFYiXZRTcMpEcN8b/1gxEuv3nODoYtCA=
golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90 h1:jiDhWWeC7jfWqR9c/uplMOqJ0sbNlNWv0UkzE0vX1MA=
golang.org/x/exp v0.0.0-20260312153236-7ab1446f8b90/go.mod h1:xE1HEv6b+1SCZ5/uscMRjUBKtIxworgEcEi+/n9NQDQ=
golang.org/x/net v0.52.0 h1:He/TN1l0e4mmR3QqHMT2Xab3Aj3L9qjbhRm78/6jrW0=
golang.org/x/net v0.52.0/go.mod h1:R1MAz7uMZxVMualyPXb+VaqGSa3LIaUqk0eEt3w36Sw=
golang.org/x/crypto v0.50.0 h1:zO47/JPrL6vsNkINmLoo/PH1gcxpls50DNogFvB5ZGI=
golang.org/x/crypto v0.50.0/go.mod h1:3muZ7vA7PBCE6xgPX7nkzzjiUq87kRItoJQM1Yo8S+Q=
golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f h1:W3F4c+6OLc6H2lb//N1q4WpJkhzJCK5J6kUi1NTVXfM=
golang.org/x/exp v0.0.0-20260410095643-746e56fc9e2f/go.mod h1:J1xhfL/vlindoeF/aINzNzt2Bket5bjo9sdOYzOsU80=
golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA=
golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs=
golang.org/x/oauth2 v0.36.0 h1:peZ/1z27fi9hUOFCAZaHyrpWG5lwe0RJEEEeH0ThlIs=
golang.org/x/oauth2 v0.36.0/go.mod h1:YDBUJMTkDnJS+A4BP4eZBjCqtokkg1hODuPjwiGPO7Q=
golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
golang.org/x/sys v0.43.0 h1:Rlag2XtaFTxp19wS8MXlJwTvoh8ArU6ezoyFsMyCTNI=
golang.org/x/sys v0.43.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
golang.org/x/term v0.41.0 h1:QCgPso/Q3RTJx2Th4bDLqML4W6iJiaXFq2/ftQF13YU=
golang.org/x/term v0.41.0/go.mod h1:3pfBgksrReYfZ5lvYM0kSO0LIkAl4Yl2bXOkKP7Ec2A=
golang.org/x/text v0.35.0 h1:JOVx6vVDFokkpaq1AEptVzLTpDe9KGpj5tR4/X+ybL8=
golang.org/x/text v0.35.0/go.mod h1:khi/HExzZJ2pGnjenulevKNX1W67CUy0AsXcNubPGCA=
golang.org/x/term v0.42.0 h1:UiKe+zDFmJobeJ5ggPwOshJIVt6/Ft0rcfrXZDLWAWY=
golang.org/x/term v0.42.0/go.mod h1:Dq/D+snpsbazcBG5+F9Q1n2rXV8Ma+71xEjTRufARgY=
golang.org/x/text v0.36.0 h1:JfKh3XmcRPqZPKevfXVpI1wXPTqbkE5f7JA92a55Yxg=
golang.org/x/text v0.36.0/go.mod h1:NIdBknypM8iqVmPiuco0Dh6P5Jcdk8lJL0CUebqK164=
golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U=
golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno=
gonum.org/v1/gonum v0.17.0 h1:VbpOemQlsSMrYmn7T2OUvQ4dqxQXU+ouZFQsZOx50z4=
gonum.org/v1/gonum v0.17.0/go.mod h1:El3tOrEuMpv2UdMrbNlKEh9vd86bmQ6vqIcDwxEOc1E=
google.golang.org/api v0.275.0 h1:vfY5d9vFVJeWEZT65QDd9hbndr7FyZ2+6mIzGAh71NI=
google.golang.org/api v0.275.0/go.mod h1:Fnag/EWUPIcJXuIkP1pjoTgS5vdxlk3eeemL7Do6bvw=
google.golang.org/genai v1.52.1 h1:dYoljKtLDXMiBdVaClSJ/ZPwZ7j1N0lGjMhwOKOQUlk=
google.golang.org/genai v1.52.1/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk=
google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7 h1:XzmzkmB14QhVhgnawEVsOn6OFsnpyxNPRY9QV01dNB0=
google.golang.org/genproto v0.0.0-20260319201613-d00831a3d3e7/go.mod h1:L43LFes82YgSonw6iTXTxXUX1OlULt4AQtkik4ULL/I=
google.golang.org/genproto/googleapis/api v0.0.0-20260319201613-d00831a3d3e7 h1:41r6JMbpzBMen0R/4TZeeAmGXSJC7DftGINUodzTkPI=
google.golang.org/genproto/googleapis/api v0.0.0-20260319201613-d00831a3d3e7/go.mod h1:EIQZ5bFCfRQDV4MhRle7+OgjNtZ6P1PiZBgAKuxXu/Y=
google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d h1:wT2n40TBqFY6wiwazVK9/iTWbsQrgk5ZfCSVFLO9LQA=
google.golang.org/genproto/googleapis/rpc v0.0.0-20260406210006-6f92a3bedf2d/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
google.golang.org/genai v1.54.0 h1:ZQCa70WMTJDI11FdqWCzGvZ5PanpcpfoO6jl/lrSnGU=
google.golang.org/genai v1.54.0/go.mod h1:A3kkl0nyBjyFlNjgxIwKq70julKbIxpSxqKO5gw/gmk=
google.golang.org/genproto v0.0.0-20260406210006-6f92a3bedf2d h1:N1Ec54vZnIPd7MnxRiYLW+oY4fDR4BOS/LrssdD9+ek=
google.golang.org/genproto v0.0.0-20260406210006-6f92a3bedf2d/go.mod h1:c2hJ1grtnH0xUiEKGDGkjGNTJ1Hy2LrblyKOHF0sqRM=
google.golang.org/genproto/googleapis/api v0.0.0-20260406210006-6f92a3bedf2d h1:/aDRtSZJjyLQzm75d+a1wOJaqyKBMvIAfeQmoa3ORiI=
google.golang.org/genproto/googleapis/api v0.0.0-20260406210006-6f92a3bedf2d/go.mod h1:etfGUgejTiadZAUaEP14NP97xi1RGeawqkjDARA/UOs=
google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478 h1:RmoJA1ujG+/lRGNfUnOMfhCy5EipVMyvUE+KNbPbTlw=
google.golang.org/genproto/googleapis/rpc v0.0.0-20260414002931-afd174a4e478/go.mod h1:4Hqkh8ycfw05ld/3BWL7rJOSfebL2Q+DVDeRgYgxUU8=
google.golang.org/grpc v1.80.0 h1:Xr6m2WmWZLETvUNvIUmeD5OAagMw3FiKmMlTdViWsHM=
google.golang.org/grpc v1.80.0/go.mod h1:ho/dLnxwi3EDJA4Zghp7k2Ec1+c2jqup0bFkw07bwF4=
google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
+214 -13
View File
@@ -30,6 +30,11 @@ type AgentConfig struct {
// If nil, remote MCP servers that require OAuth will fail to connect.
AuthHandler tools.MCPAuthHandler
// TokenStoreFactory, if non-nil, creates a custom token store for each
// remote MCP server's OAuth tokens. When nil, the default file-based
// token store is used.
TokenStoreFactory tools.TokenStoreFactory
// CoreTools overrides the default core tool set. If empty, core.AllTools()
// is used. This allows SDK users to provide a custom tool set (e.g.
// CodingTools or tools with a custom WorkDir).
@@ -82,6 +87,19 @@ type ReasoningDeltaHandler func(delta string)
// Called when the last reasoning token has been processed, before text streaming starts.
type ReasoningCompleteHandler func()
// ToolCallStartHandler is a function type for handling the moment when the LLM
// begins generating tool call arguments. The tool name is known but the full
// argument JSON is still streaming.
type ToolCallStartHandler func(toolCallID, toolName string)
// ToolCallDeltaHandler is a function type for handling streamed fragments of
// tool call arguments as they arrive from the LLM.
type ToolCallDeltaHandler func(toolCallID, delta string)
// ToolCallEndHandler is a function type for handling the end of tool argument
// streaming, before the tool call is parsed and execution begins.
type ToolCallEndHandler func(toolCallID string)
// ToolOutputHandler is a function type for handling streaming tool output chunks.
// Used by tools like bash to stream output as it arrives rather than waiting
// for the command to complete. The isStderr flag indicates if the chunk
@@ -89,6 +107,12 @@ type ReasoningCompleteHandler func()
// Note: This is an alias for core.ToolOutputCallback to avoid import cycles.
type ToolOutputHandler = core.ToolOutputCallback
// PasswordPromptHandler is a function type for password prompts.
// Used by the bash tool when sudo requires a password. The handler receives
// a prompt message and returns the password and whether it was cancelled.
// Note: This is an alias for core.PasswordPromptCallback.
type PasswordPromptHandler = core.PasswordPromptCallback
// StepMessagesHandler is a function type for persisting messages after each
// complete step in a multi-step agent turn. The handler receives the messages
// produced by the step (typically an assistant message with tool calls followed
@@ -130,6 +154,12 @@ type Agent struct {
skipMaxOutputTokens bool
modelConfig *models.ProviderConfig
// authHandler and tokenStoreFactory are stored from AgentConfig so that
// AddMCPServer() can propagate them when creating a new MCPToolManager
// at runtime (i.e. when no MCP servers were configured at init time).
authHandler tools.MCPAuthHandler
tokenStoreFactory tools.TokenStoreFactory
// mcpReady is closed when background MCP tool loading completes (success
// or failure). nil when no MCP servers are configured.
mcpReady chan struct{}
@@ -226,16 +256,20 @@ func NewAgent(ctx context.Context, agentConfig *AgentConfig) (*Agent, error) {
providerOptions: providerResult.ProviderOptions,
skipMaxOutputTokens: providerResult.SkipMaxOutputTokens,
modelConfig: agentConfig.ModelConfig,
authHandler: agentConfig.AuthHandler,
tokenStoreFactory: agentConfig.TokenStoreFactory,
}
// Start MCP tool loading in the background if servers are configured.
// The mcpReady channel is closed when loading completes (success or failure).
if agentConfig.MCPConfig != nil && len(agentConfig.MCPConfig.MCPServers) > 0 {
toolManager := tools.NewMCPToolManager()
toolManager.SetModel(providerResult.Model)
if agentConfig.AuthHandler != nil {
toolManager.SetAuthHandler(agentConfig.AuthHandler)
}
if agentConfig.TokenStoreFactory != nil {
toolManager.SetTokenStoreFactory(agentConfig.TokenStoreFactory)
}
if agentConfig.DebugLogger != nil {
toolManager.SetDebugLogger(agentConfig.DebugLogger)
}
@@ -309,7 +343,7 @@ func (a *Agent) rebuildFantasyAgent() {
allTools := make([]fantasy.AgentTool, len(a.coreTools))
copy(allTools, a.coreTools)
if a.toolManager != nil {
allTools = append(allTools, a.toolManager.GetTools()...)
allTools = append(allTools, mcpToolsToAgentTools(a.toolManager.GetTools(), a.toolManager)...)
}
if len(a.extraTools) > 0 {
allTools = append(allTools, a.extraTools...)
@@ -390,7 +424,7 @@ func (a *Agent) GenerateWithLoop(ctx context.Context, messages []fantasy.Message
onResponse ResponseHandler, onToolCallContent ToolCallContentHandler,
) (*GenerateWithLoopResult, error) {
return a.GenerateWithLoopAndStreaming(ctx, messages, onToolCall, onToolExecution, onToolResult,
onResponse, onToolCallContent, nil, nil, nil, nil, nil, nil)
onResponse, onToolCallContent, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
}
// GenerateWithLoopAndStreaming processes messages using the agent with streaming and callbacks.
@@ -405,6 +439,10 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
onToolOutput ToolOutputHandler,
onStepMessages StepMessagesHandler,
onStepUsage StepUsageHandler,
onPasswordPrompt PasswordPromptHandler,
onToolCallStart ToolCallStartHandler,
onToolCallDelta ToolCallDeltaHandler,
onToolCallEnd ToolCallEndHandler,
) (*GenerateWithLoopResult, error) {
// Wait for background MCP tool loading to complete and rebuild the
@@ -417,6 +455,11 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
ctx = core.ContextWithToolOutputCallback(ctx, onToolOutput)
}
// Inject password prompt handler into context for use by bash tool.
if onPasswordPrompt != nil {
ctx = core.ContextWithPasswordPrompt(ctx, onPasswordPrompt)
}
// The agent requires the current user input as Prompt, with prior messages as history.
// Extract the last user message text and files as the prompt, and pass everything
// before it as Messages. Files (e.g. clipboard images) are passed via the Files
@@ -435,7 +478,8 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
// Stream is required to observe tool execution in real time. The non-streaming
// Generate path is reserved for the simple case with no callbacks at all.
hasCallbacks := onToolCall != nil || onToolExecution != nil || onToolResult != nil ||
onToolCallContent != nil || onStreamingResponse != nil || onReasoningDelta != nil
onToolCallContent != nil || onStreamingResponse != nil || onReasoningDelta != nil ||
onToolCallStart != nil || onToolCallDelta != nil || onToolCallEnd != nil
if a.streamingEnabled || hasCallbacks {
// Track completed step messages so we can return partial results
@@ -454,6 +498,35 @@ func (a *Agent) GenerateWithLoopAndStreaming(ctx context.Context, messages []fan
Files: files,
Messages: history,
// Tool input streaming callbacks — fire during tool argument generation
OnToolInputStart: func(id, toolName string) error {
if ctx.Err() != nil {
return ctx.Err()
}
if onToolCallStart != nil {
onToolCallStart(id, toolName)
}
return nil
},
OnToolInputDelta: func(id, delta string) error {
if ctx.Err() != nil {
return ctx.Err()
}
if onToolCallDelta != nil {
onToolCallDelta(id, delta)
}
return nil
},
OnToolInputEnd: func(id string) error {
if ctx.Err() != nil {
return ctx.Err()
}
if onToolCallEnd != nil {
onToolCallEnd(id)
}
return nil
},
// Reasoning/thinking streaming callback
OnReasoningDelta: func(id, delta string) error {
if ctx.Err() != nil {
@@ -792,7 +865,7 @@ func (a *Agent) GetTools() []fantasy.AgentTool {
allTools := make([]fantasy.AgentTool, len(a.coreTools))
copy(allTools, a.coreTools)
if a.toolManager != nil {
allTools = append(allTools, a.toolManager.GetTools()...)
allTools = append(allTools, mcpToolsToAgentTools(a.toolManager.GetTools(), a.toolManager)...)
}
if len(a.extraTools) > 0 {
allTools = append(allTools, a.extraTools...)
@@ -826,6 +899,64 @@ func (a *Agent) SetExtraTools(extraTools []fantasy.AgentTool) {
a.rebuildFantasyAgent()
}
// AddMCPServer connects to a new MCP server at runtime and makes its tools
// available to the agent. Returns the number of tools loaded.
// If the agent has no tool manager (no MCP servers were configured at init),
// one is created automatically.
func (a *Agent) AddMCPServer(ctx context.Context, name string, cfg config.MCPServerConfig) (int, error) {
// Ensure MCP tools from initial load are settled first.
a.ensureMCPTools()
if a.toolManager == nil {
a.toolManager = tools.NewMCPToolManager()
if a.authHandler != nil {
a.toolManager.SetAuthHandler(a.authHandler)
}
if a.tokenStoreFactory != nil {
a.toolManager.SetTokenStoreFactory(a.tokenStoreFactory)
}
a.toolManager.SetOnToolsChanged(func() {
a.rebuildFantasyAgent()
})
}
count, err := a.toolManager.AddServer(ctx, name, cfg)
if err != nil {
return 0, err
}
// AddServer's onToolsChanged callback triggers rebuildFantasyAgent,
// but only if it was wired. Ensure rebuild happens regardless.
a.rebuildFantasyAgent()
return count, nil
}
// RemoveMCPServer disconnects an MCP server and removes its tools from the agent.
func (a *Agent) RemoveMCPServer(name string) error {
if a.toolManager == nil {
return fmt.Errorf("no MCP servers loaded")
}
// Ensure MCP tools from initial load are settled first.
a.ensureMCPTools()
err := a.toolManager.RemoveServer(name)
if err != nil {
return err
}
// RemoveServer's onToolsChanged callback triggers rebuildFantasyAgent,
// but ensure rebuild happens regardless.
a.rebuildFantasyAgent()
return nil
}
// GetMCPToolManager returns the underlying MCP tool manager.
// Returns nil if no MCP servers have been configured.
func (a *Agent) GetMCPToolManager() *tools.MCPToolManager {
return a.toolManager
}
// GetLoadingMessage returns the loading message from provider creation.
func (a *Agent) GetLoadingMessage() string {
return a.loadingMessage
@@ -839,9 +970,61 @@ func (a *Agent) GetLoadedServerNames() []string {
return a.toolManager.GetLoadedServerNames()
}
// SetModel swaps the agent's LLM provider to a new model. The existing tools,
// system prompt, and configuration are preserved. The old provider is closed
// if it has a closer. Returns the previous model string for notification.
// GetMCPPrompts returns all prompts discovered from connected MCP servers.
// Returns nil if no MCP servers are configured or no prompts were found.
func (a *Agent) GetMCPPrompts() []tools.MCPPrompt {
if a.toolManager == nil {
return nil
}
return a.toolManager.GetPrompts()
}
// GetMCPPrompt retrieves and expands a specific prompt from an MCP server.
// This is a lazy call — the server is contacted each time.
func (a *Agent) GetMCPPrompt(ctx context.Context, serverName, promptName string, args map[string]string) (*tools.MCPPromptResult, error) {
if a.toolManager == nil {
return nil, fmt.Errorf("no MCP servers configured")
}
return a.toolManager.GetPrompt(ctx, serverName, promptName, args)
}
// GetMCPResources returns all resources discovered from connected MCP servers.
func (a *Agent) GetMCPResources() []tools.MCPResource {
if a.toolManager == nil {
return nil
}
return a.toolManager.GetResources()
}
// ReadMCPResource reads a specific resource from an MCP server by URI.
func (a *Agent) ReadMCPResource(ctx context.Context, serverName, uri string) (*tools.MCPResourceContent, error) {
if a.toolManager == nil {
return nil, fmt.Errorf("no MCP servers configured")
}
return a.toolManager.ReadResource(ctx, serverName, uri)
}
// SubscribeMCPResource subscribes to change notifications for a resource.
func (a *Agent) SubscribeMCPResource(ctx context.Context, serverName, uri string) error {
if a.toolManager == nil {
return fmt.Errorf("no MCP servers configured")
}
return a.toolManager.SubscribeResource(ctx, serverName, uri)
}
// UnsubscribeMCPResource cancels change notifications for a resource.
func (a *Agent) UnsubscribeMCPResource(ctx context.Context, serverName, uri string) error {
if a.toolManager == nil {
return fmt.Errorf("no MCP servers configured")
}
return a.toolManager.UnsubscribeResource(ctx, serverName, uri)
}
// SetModel swaps the agent's LLM provider to a new model. The existing tools
// and configuration are preserved. When the new model's ProviderConfig carries
// a system prompt (from per-model settings), it replaces the agent's stored
// prompt so the rebuilt fantasy agent uses it. The old provider is closed if
// it has a closer.
func (a *Agent) SetModel(ctx context.Context, config *models.ProviderConfig) error {
// Ensure MCP tools are loaded before rebuilding (SetModel may be called
// before the first LLM call).
@@ -856,11 +1039,6 @@ func (a *Agent) SetModel(ctx context.Context, config *models.ProviderConfig) err
_ = a.providerCloser.Close()
}
// Update model info on MCP tool manager.
if a.toolManager != nil {
a.toolManager.SetModel(providerResult.Model)
}
// Swap fields.
a.model = providerResult.Model
a.providerCloser = providerResult.Closer
@@ -868,6 +1046,13 @@ func (a *Agent) SetModel(ctx context.Context, config *models.ProviderConfig) err
a.skipMaxOutputTokens = providerResult.SkipMaxOutputTokens
a.modelConfig = config
// Update system prompt when the config carries one (from per-model
// settings or the global config). This allows model-specific system
// prompts to take effect on model switch.
if config.SystemPrompt != "" {
a.systemPrompt = config.SystemPrompt
}
// Update provider type.
if config.ModelString != "" {
if p, _, err := models.ParseModelString(config.ModelString); err == nil {
@@ -886,6 +1071,22 @@ func (a *Agent) GetModel() fantasy.LanguageModel {
return a.model
}
// GetMaxTokens returns the effective max output tokens the agent currently
// sends to the LLM provider, after per-model defaults, right-sizing, and any
// Anthropic thinking-budget adjustments. Returns 0 when no ModelConfig is
// attached (e.g. early init) or when the provider suppresses the parameter
// (e.g. Codex OAuth), which allows callers to differentiate "default" from
// "explicitly capped".
func (a *Agent) GetMaxTokens() int {
if a.skipMaxOutputTokens {
return 0
}
if a.modelConfig == nil {
return 0
}
return a.modelConfig.MaxTokens
}
// Close closes the agent and cleans up resources.
// If MCP tools are still loading in the background, Close waits for them
// to finish before closing connections to avoid resource leaks.
+302
View File
@@ -0,0 +1,302 @@
package agent
import (
"context"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
"charm.land/fantasy"
"github.com/mark3labs/kit/internal/config"
)
// mockModel is a minimal LanguageModel that satisfies the interface
// without making real API calls. Used to test tool management wiring.
type mockModel struct{}
func (m *mockModel) Generate(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) {
return &fantasy.Response{}, nil
}
func (m *mockModel) Stream(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) {
return nil, nil
}
func (m *mockModel) GenerateObject(_ context.Context, _ fantasy.ObjectCall) (*fantasy.ObjectResponse, error) {
return &fantasy.ObjectResponse{}, nil
}
func (m *mockModel) StreamObject(_ context.Context, _ fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) {
return nil, nil
}
func (m *mockModel) Provider() string { return "mock" }
func (m *mockModel) Model() string { return "mock-model" }
// testdataDir returns the absolute path to the tools testdata directory.
func testdataDir(t *testing.T) string {
t.Helper()
_, file, _, ok := runtime.Caller(0)
if !ok {
t.Fatal("cannot determine test file path")
}
return filepath.Join(filepath.Dir(file), "..", "tools", "testdata")
}
// echoServerConfig returns an MCPServerConfig for the test echo MCP server.
func echoServerConfig(t *testing.T) config.MCPServerConfig {
t.Helper()
script := filepath.Join(testdataDir(t), "echo_server.py")
if _, err := os.Stat(script); err != nil {
t.Skipf("echo_server.py not found: %v", err)
}
return config.MCPServerConfig{
Command: []string{"python3", script},
}
}
// mockAuthHandler is a minimal MCPAuthHandler for testing that auth handler
// propagation works without requiring a real OAuth server.
type mockAuthHandler struct {
redirectURI string
}
func (h *mockAuthHandler) RedirectURI() string { return h.redirectURI }
func (h *mockAuthHandler) HandleAuth(_ context.Context, _ string, _ string) (string, error) {
return "", nil
}
// newTestAgent creates a minimal Agent with a mock model and no core tools,
// suitable for testing MCP server management without an API key.
func newTestAgent() *Agent {
model := &mockModel{}
a := &Agent{
model: model,
coreTools: nil,
extraTools: nil,
maxSteps: 10,
systemPrompt: "test",
fantasyAgent: fantasy.NewAgent(model),
}
return a
}
func TestAgent_AddMCPServer(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
a := newTestAgent()
defer func() { _ = a.Close() }()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
// Initially no MCP tools.
if a.GetMCPToolCount() != 0 {
t.Fatalf("Expected 0 MCP tools initially, got %d", a.GetMCPToolCount())
}
// Add a server.
count, err := a.AddMCPServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("AddMCPServer failed: %v", err)
}
if count != 2 {
t.Errorf("Expected 2 tools, got %d", count)
}
// Verify tools are in the agent's tool list.
if a.GetMCPToolCount() != 2 {
t.Errorf("Expected 2 MCP tools, got %d", a.GetMCPToolCount())
}
allTools := a.GetTools()
toolNames := make(map[string]bool)
for _, tool := range allTools {
toolNames[tool.Info().Name] = true
}
if !toolNames["echo__echo"] {
t.Error("Expected tool 'echo__echo' in agent tools")
}
if !toolNames["echo__greet"] {
t.Error("Expected tool 'echo__greet' in agent tools")
}
// Verify loaded server names.
names := a.GetLoadedServerNames()
found := false
for _, n := range names {
if n == "echo" {
found = true
}
}
if !found {
t.Errorf("Expected 'echo' in loaded server names: %v", names)
}
}
func TestAgent_RemoveMCPServer(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
a := newTestAgent()
defer func() { _ = a.Close() }()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
// Add then remove.
_, err := a.AddMCPServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("AddMCPServer failed: %v", err)
}
err = a.RemoveMCPServer("echo")
if err != nil {
t.Fatalf("RemoveMCPServer failed: %v", err)
}
// Verify tools removed.
if a.GetMCPToolCount() != 0 {
t.Errorf("Expected 0 MCP tools after removal, got %d", a.GetMCPToolCount())
}
// Verify agent's tool list has no MCP tools.
for _, tool := range a.GetTools() {
if strings.Contains(tool.Info().Name, "echo__") {
t.Errorf("Found leftover tool after removal: %s", tool.Info().Name)
}
}
}
func TestAgent_RemoveMCPServer_NoToolManager(t *testing.T) {
a := newTestAgent()
defer func() { _ = a.Close() }()
err := a.RemoveMCPServer("nonexistent")
if err == nil {
t.Fatal("Expected error when no tool manager exists")
}
if !strings.Contains(err.Error(), "no MCP servers loaded") {
t.Errorf("Expected 'no MCP servers loaded' error, got: %v", err)
}
}
func TestAgent_AddMCPServer_CreatesToolManager(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
a := newTestAgent()
defer func() { _ = a.Close() }()
// Initially no tool manager.
if a.GetMCPToolManager() != nil {
t.Fatal("Expected nil tool manager initially")
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
_, err := a.AddMCPServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("AddMCPServer failed: %v", err)
}
// Tool manager should now exist.
if a.GetMCPToolManager() == nil {
t.Fatal("Expected tool manager to be created by AddMCPServer")
}
}
func TestAgent_AddRemoveAdd_MCP(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
a := newTestAgent()
defer func() { _ = a.Close() }()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
// Add → Remove → Add cycle.
_, err := a.AddMCPServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("First add failed: %v", err)
}
err = a.RemoveMCPServer("echo")
if err != nil {
t.Fatalf("Remove failed: %v", err)
}
count, err := a.AddMCPServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("Re-add failed: %v", err)
}
if count != 2 {
t.Errorf("Expected 2 tools on re-add, got %d", count)
}
if a.GetMCPToolCount() != 2 {
t.Errorf("Expected 2 MCP tools after re-add, got %d", a.GetMCPToolCount())
}
}
// TestAgent_AddMCPServer_InheritsAuthHandler verifies that AddMCPServer()
// propagates the agent's authHandler and tokenStoreFactory to a newly created
// MCPToolManager (fix for issue #3).
func TestAgent_AddMCPServer_InheritsAuthHandler(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
handler := &mockAuthHandler{redirectURI: "http://localhost:9999/oauth/callback"}
model := &mockModel{}
a := &Agent{
model: model,
coreTools: nil,
extraTools: nil,
maxSteps: 10,
systemPrompt: "test",
fantasyAgent: fantasy.NewAgent(model),
authHandler: handler,
tokenStoreFactory: nil, // nil is fine; we just test authHandler propagation
}
defer func() { _ = a.Close() }()
// Initially no tool manager.
if a.GetMCPToolManager() != nil {
t.Fatal("Expected nil tool manager initially")
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
_, err := a.AddMCPServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("AddMCPServer failed: %v", err)
}
// Tool manager should now exist and have the auth handler set.
tm := a.GetMCPToolManager()
if tm == nil {
t.Fatal("Expected tool manager to be created by AddMCPServer")
}
// Verify the auth handler was propagated by checking the field directly.
if tm.GetAuthHandler() == nil {
t.Fatal("Expected auth handler to be propagated to tool manager")
}
}
+5
View File
@@ -38,6 +38,10 @@ type AgentCreationOptions struct {
DebugLogger tools.DebugLogger // Optional debug logger
// AuthHandler handles OAuth authorization for remote MCP servers
AuthHandler tools.MCPAuthHandler
// TokenStoreFactory, if non-nil, creates a custom token store for each
// remote MCP server's OAuth tokens. When nil, the default file-based
// token store is used.
TokenStoreFactory tools.TokenStoreFactory
// CoreTools overrides the default core tool set. If empty, core.AllTools()
// is used.
CoreTools []fantasy.AgentTool
@@ -66,6 +70,7 @@ func CreateAgent(ctx context.Context, opts *AgentCreationOptions) (*Agent, error
StreamingEnabled: opts.StreamingEnabled,
DebugLogger: opts.DebugLogger,
AuthHandler: opts.AuthHandler,
TokenStoreFactory: opts.TokenStoreFactory,
CoreTools: opts.CoreTools,
DisableCoreTools: opts.DisableCoreTools,
ToolWrapper: opts.ToolWrapper,
+65
View File
@@ -0,0 +1,65 @@
package agent
import (
"context"
"fmt"
"charm.land/fantasy"
"github.com/mark3labs/kit/internal/tools"
)
// mcpAgentTool adapts an tools.MCPTool to the fantasy.AgentTool interface.
// This keeps the fantasy dependency confined to the agent layer — the tools
// package is a pure MCP client library with no LLM framework dependency.
type mcpAgentTool struct {
tool tools.MCPTool
manager *tools.MCPToolManager
providerOptions fantasy.ProviderOptions
}
// Info returns the fantasy tool info including name, description, and parameter schema.
func (t *mcpAgentTool) Info() fantasy.ToolInfo {
return fantasy.ToolInfo{
Name: t.tool.Name,
Description: t.tool.Description,
Parameters: t.tool.Parameters,
Required: t.tool.Required,
}
}
// Run executes the MCP tool by delegating to the MCPToolManager.
func (t *mcpAgentTool) Run(ctx context.Context, call fantasy.ToolCall) (fantasy.ToolResponse, error) {
result, err := t.manager.ExecuteTool(ctx, t.tool.Name, call.Input)
if err != nil {
return fantasy.ToolResponse{}, fmt.Errorf("mcp tool execution failed: %w", err)
}
if result.IsError {
return fantasy.NewTextErrorResponse(result.Content), nil
}
return fantasy.NewTextResponse(result.Content), nil
}
// ProviderOptions returns provider-specific options for this tool.
func (t *mcpAgentTool) ProviderOptions() fantasy.ProviderOptions {
return t.providerOptions
}
// SetProviderOptions sets provider-specific options for this tool.
func (t *mcpAgentTool) SetProviderOptions(opts fantasy.ProviderOptions) {
t.providerOptions = opts
}
// mcpToolsToAgentTools converts a slice of MCPTool to fantasy.AgentTool
// implementations that route execution through the MCPToolManager.
func mcpToolsToAgentTools(mcpTools []tools.MCPTool, manager *tools.MCPToolManager) []fantasy.AgentTool {
agentTools := make([]fantasy.AgentTool, len(mcpTools))
for i, t := range mcpTools {
agentTools[i] = &mcpAgentTool{
tool: t,
manager: manager,
}
}
return agentTools
}
+145 -21
View File
@@ -497,6 +497,12 @@ func (a *App) CompactAsync(customInstructions string, onComplete func(), onError
// response text to stdout. No intermediate events are emitted. Blocks until
// the step completes or ctx is cancelled.
func (a *App) RunOnce(ctx context.Context, prompt string) error {
return a.RunOnceWithFiles(ctx, prompt, nil)
}
// RunOnceWithFiles executes a single agent step synchronously with optional
// multimodal file attachments. Prints the response to stdout and returns.
func (a *App) RunOnceWithFiles(ctx context.Context, prompt string, files []kit.LLMFilePart) error {
stepCtx, cancel := context.WithCancel(ctx)
defer cancel()
@@ -504,7 +510,7 @@ func (a *App) RunOnce(ctx context.Context, prompt string) error {
a.cancelStep = cancel
a.mu.Unlock()
result, err := a.executeStep(stepCtx, prompt, nil, nil)
result, err := a.executeStep(stepCtx, prompt, nil, files)
if err != nil {
return err
}
@@ -519,6 +525,12 @@ func (a *App) RunOnce(ctx context.Context, prompt string) error {
// full TurnResult without printing anything. This is used by --json mode to
// capture structured output for serialization.
func (a *App) RunOnceResult(ctx context.Context, prompt string) (*kit.TurnResult, error) {
return a.RunOnceResultWithFiles(ctx, prompt, nil)
}
// RunOnceResultWithFiles executes a single agent step synchronously with
// optional multimodal file attachments and returns the full TurnResult.
func (a *App) RunOnceResultWithFiles(ctx context.Context, prompt string, files []kit.LLMFilePart) (*kit.TurnResult, error) {
stepCtx, cancel := context.WithCancel(ctx)
defer cancel()
@@ -526,7 +538,7 @@ func (a *App) RunOnceResult(ctx context.Context, prompt string) (*kit.TurnResult
a.cancelStep = cancel
a.mu.Unlock()
return a.executeStep(stepCtx, prompt, nil, nil)
return a.executeStep(stepCtx, prompt, nil, files)
}
// RunOnceWithDisplay executes a single agent step synchronously, sending
@@ -540,6 +552,12 @@ func (a *App) RunOnceResult(ctx context.Context, prompt string) (*kit.TurnResult
//
// Blocks until the step completes or ctx is cancelled.
func (a *App) RunOnceWithDisplay(ctx context.Context, prompt string, eventFn func(tea.Msg)) error {
return a.RunOnceWithDisplayAndFiles(ctx, prompt, eventFn, nil)
}
// RunOnceWithDisplayAndFiles executes a single agent step synchronously with
// optional multimodal file attachments, sending intermediate display events.
func (a *App) RunOnceWithDisplayAndFiles(ctx context.Context, prompt string, eventFn func(tea.Msg), files []kit.LLMFilePart) error {
stepCtx, cancel := context.WithCancel(ctx)
defer cancel()
@@ -547,7 +565,7 @@ func (a *App) RunOnceWithDisplay(ctx context.Context, prompt string, eventFn fun
a.cancelStep = cancel
a.mu.Unlock()
result, err := a.executeStep(stepCtx, prompt, eventFn, nil)
result, err := a.executeStep(stepCtx, prompt, eventFn, files)
if err != nil {
return err
}
@@ -870,6 +888,12 @@ func (a *App) subscribeSDKEvents(sendFn func(tea.Msg), stepUsageSeen *atomic.Boo
switch ev := e.(type) {
case kit.ToolCallEvent:
sendFn(ToolCallStartedEvent{ToolCallID: ev.ToolCallID, ToolName: ev.ToolName, ToolArgs: ev.ToolArgs})
case kit.ToolCallStartEvent:
sendFn(ToolCallInputStartEvent{ToolCallID: ev.ToolCallID, ToolName: ev.ToolName, ToolKind: ev.ToolKind})
case kit.ToolCallDeltaEvent:
sendFn(ToolCallInputDeltaEvent{ToolCallID: ev.ToolCallID, Delta: ev.Delta})
case kit.ToolCallEndEvent:
sendFn(ToolCallInputEndEvent{ToolCallID: ev.ToolCallID})
case kit.ToolExecutionStartEvent:
sendFn(ToolExecutionEvent{ToolCallID: ev.ToolCallID, ToolName: ev.ToolName, ToolArgs: ev.ToolArgs, IsStarting: true})
case kit.ToolExecutionEndEvent:
@@ -900,6 +924,22 @@ func (a *App) subscribeSDKEvents(sendFn func(tea.Msg), stepUsageSeen *atomic.Boo
sendFn(SteerConsumedEvent{})
case kit.StepUsageEvent:
a.recordStepUsage(ev, stepUsageSeen)
case kit.PasswordPromptEvent:
// Convert SDK PasswordPromptEvent to app PasswordPromptEvent
// The TUI will handle this and send the response back
responseCh := make(chan PasswordPromptResponse, 1)
sendFn(PasswordPromptEvent{
Prompt: ev.Prompt,
ResponseCh: responseCh,
})
// Wait for TUI response and forward to SDK
resp := <-responseCh
ev.ResponseCh <- kit.PasswordPromptResponse{
Password: resp.Password,
Cancelled: resp.Cancelled,
}
case kit.TurnEndEvent:
a.handleTurnEnd(ev, sendFn)
}
}))
@@ -910,6 +950,64 @@ func (a *App) subscribeSDKEvents(sendFn func(tea.Msg), stepUsageSeen *atomic.Boo
}
}
// handleTurnEnd inspects a turn's final StopReason and surfaces actionable
// feedback to the user when the turn ended in a state they can act on.
//
// Today the only surfaced case is FinishReasonLength — the model hit its
// configured max_output_tokens budget and the reply was truncated. Without
// this banner the TUI used to swallow the truncation silently, leading to
// "ghost" cut-offs with no indication of why.
//
// Separated from subscribeSDKEvents so tests can exercise it directly via a
// stubbed sendFn without standing up a full Kit.
func (a *App) handleTurnEnd(ev kit.TurnEndEvent, sendFn func(tea.Msg)) {
if sendFn == nil {
return
}
if ev.StopReason != kit.FinishReasonLength {
return
}
sendFn(ExtensionPrintEvent{
Level: "info",
Text: a.formatMaxTokensTruncatedMessage(),
})
}
// formatMaxTokensTruncatedMessage builds the user-facing explanation for a
// truncated turn. It reports the active max_output_tokens budget and, when
// known, the model's catalog output ceiling so the user can judge how much
// headroom is available.
func (a *App) formatMaxTokensTruncatedMessage() string {
k := a.opts.Kit
if k == nil {
// Extremely early / test-stub case: still emit a useful generic hint.
return "⚠ Response truncated: the model hit the configured max_output_tokens limit. " +
"Raise it with --max-tokens N, KIT_MAX_TOKENS=N, or per-model " +
"modelSettings[provider/model].maxTokens in config."
}
current := k.MaxTokens()
ceiling := k.MaxOutputLimit()
model := k.GetModelString()
msg := "⚠ Response truncated: "
if model != "" {
msg += fmt.Sprintf("%s hit the configured max_output_tokens limit", model)
} else {
msg += "the model hit the configured max_output_tokens limit"
}
if current > 0 {
msg += fmt.Sprintf(" (%d)", current)
}
msg += "."
if ceiling > 0 && current > 0 && ceiling > current {
msg += fmt.Sprintf(" This model supports up to %d output tokens.", ceiling)
}
msg += "\n\nRaise it with --max-tokens N, KIT_MAX_TOKENS=N, " +
"or per-model modelSettings[provider/model].maxTokens in your config. " +
"Re-run the last prompt after raising it to get the full response."
return msg
}
// QuitFromExtension triggers a graceful shutdown. In interactive mode it
// sends a tea.QuitMsg to the program so the TUI exits cleanly. In
// non-interactive mode it cancels the root context, stopping any in-flight
@@ -930,7 +1028,8 @@ func (a *App) QuitFromExtension() {
// controls styling: "" for plain text, "info" for a system message block,
// "error" for an error block. In interactive mode it sends an
// ExtensionPrintEvent through the program so the TUI can render it with the
// appropriate renderer. In non-interactive mode it falls back to stdout.
// appropriate renderer. In non-interactive mode it falls back to stderr with
// a level prefix so errors are distinguishable from plain output.
func (a *App) PrintFromExtension(level, text string) {
a.mu.Lock()
prog := a.program
@@ -939,8 +1038,16 @@ func (a *App) PrintFromExtension(level, text string) {
prog.Send(ExtensionPrintEvent{Text: text, Level: level})
return
}
// Non-interactive fallback: write directly to stdout.
fmt.Println(text)
// Non-interactive fallback: write to stderr with a level prefix so that
// errors and info messages are distinguishable from plain output.
switch level {
case "error":
fmt.Fprintf(os.Stderr, "[ERROR] %s\n", text)
case "info":
fmt.Fprintf(os.Stderr, "[INFO] %s\n", text)
default:
fmt.Println(text)
}
}
// SetEditorTextFromExtension sends an EditorTextSetEvent to the TUI to
@@ -1122,11 +1229,12 @@ func (a *App) PrintBlockFromExtension(opts extensions.PrintBlockOpts) {
})
return
}
// Non-interactive fallback.
// Non-interactive fallback: render a simple framed block to stderr so
// it is visually distinct from plain stdout output.
if opts.Subtitle != "" {
fmt.Printf("%s\n — %s\n", opts.Text, opts.Subtitle)
fmt.Fprintf(os.Stderr, "--- %s ---\n%s\n", opts.Subtitle, opts.Text)
} else {
fmt.Println(opts.Text)
fmt.Fprintf(os.Stderr, "---\n%s\n---\n", opts.Text)
}
}
@@ -1155,9 +1263,10 @@ func (a *App) recordStepUsage(ev kit.StepUsageEvent, stepUsageSeen *atomic.Bool)
int(ev.CacheWriteTokens),
)
// NOTE: We do NOT call SetContextTokens here. Context fill is set once
// at turn completion via updateUsageFromTurnResult using FinalUsage.InputTokens,
// which reflects the full accumulated context. Per-step context tokens would
// cause the display to jump around during multi-step tool calls.
// at turn completion via updateUsageFromTurnResult, which sums all token
// categories (Input + CacheRead + CacheCreate + Output) from FinalUsage.
// Per-step context tokens would cause the display to jump around during
// multi-step tool calls.
}
// updateUsageFromTurnResult records token usage from an SDK TurnResult into the
@@ -1221,15 +1330,30 @@ func (a *App) updateUsageFromTurnResult(result *kit.TurnResult, userPrompt strin
}
// --- Context window fill (drives the % bar) ---
// Use FinalUsage.InputTokens as the context window fill. The API's InputTokens
// already includes the full conversation history (system prompt + all previous
// messages + current user message). Adding OutputTokens would double-count since
// the output becomes part of the input for the next turn.
if result.FinalUsage != nil && result.FinalUsage.InputTokens > 0 {
if a.opts.Debug {
log.Printf("[DEBUG] updateUsageFromTurnResult: calling SetContextTokens=%d (FinalUsage.InputTokens)",
result.FinalUsage.InputTokens)
// Calculate context fill from the LAST API call's usage. The context
// window is filled by everything sent to and received from the model:
//
// InputTokens — non-cached input (may be small with prompt caching)
// CacheReadTokens — input tokens served from cache
// CacheCreationTokens — input tokens written to cache this call
// OutputTokens — assistant output (becomes input next turn)
//
// With Anthropic prompt caching, InputTokens can drop to near-zero while
// CacheReadTokens holds the bulk of the context. We must sum all four to
// get the true context window utilization.
//
// We use FinalUsage (last step only), NOT TotalUsage, because TotalUsage
// sums across all tool-calling steps — and each step re-sends the full
// conversation, so TotalUsage massively overstates the actual window fill.
if result.FinalUsage != nil {
u := result.FinalUsage
contextFill := int(u.InputTokens) + int(u.CacheReadTokens) + int(u.CacheCreationTokens) + int(u.OutputTokens)
if contextFill > 0 {
if a.opts.Debug {
log.Printf("[DEBUG] updateUsageFromTurnResult: SetContextTokens=%d (Input=%d + CacheRead=%d + CacheCreate=%d + Output=%d)",
contextFill, u.InputTokens, u.CacheReadTokens, u.CacheCreationTokens, u.OutputTokens)
}
a.opts.UsageTracker.SetContextTokens(contextFill)
}
a.opts.UsageTracker.SetContextTokens(int(result.FinalUsage.InputTokens))
}
}
+112 -13
View File
@@ -3,10 +3,12 @@ package app
import (
"context"
"errors"
"strings"
"sync"
"testing"
"time"
tea "charm.land/bubbletea/v2"
kit "github.com/mark3labs/kit/pkg/kit"
)
@@ -630,10 +632,12 @@ func TestUpdateUsageFromTurnResult_recordsWhenInputTokensZero(t *testing.T) {
}
}
// TestUpdateUsageFromTurnResult_contextTokensUsesInputOnly verifies that context
// window fill uses InputTokens only (not input+output). The API's InputTokens
// already includes the full conversation history; adding output would double-count.
func TestUpdateUsageFromTurnResult_contextTokensUsesInputOnly(t *testing.T) {
// TestUpdateUsageFromTurnResult_contextTokensUsesAllCategories verifies that
// context window fill uses all token categories from the final API call:
// InputTokens + CacheReadTokens + CacheCreationTokens + OutputTokens.
// With Anthropic prompt caching, InputTokens can be near-zero while
// CacheReadTokens holds the bulk of the context.
func TestUpdateUsageFromTurnResult_contextTokensUsesAllCategories(t *testing.T) {
usage := &usageUpdaterStub{}
app := New(Options{UsageTracker: usage}, nil)
defer app.Close()
@@ -641,22 +645,117 @@ func TestUpdateUsageFromTurnResult_contextTokensUsesInputOnly(t *testing.T) {
app.updateUsageFromTurnResult(&kit.TurnResult{
Response: "ok",
TotalUsage: &kit.LLMUsage{
InputTokens: 1000,
OutputTokens: 200,
InputTokens: 3,
OutputTokens: 5,
CacheReadTokens: 0,
CacheCreationTokens: 4317,
},
FinalUsage: &kit.LLMUsage{
InputTokens: 1000, // Full context including history
OutputTokens: 200,
InputTokens: 3, // Non-cached input (small with caching)
OutputTokens: 5, // Assistant output
CacheReadTokens: 0, // No cache reads on first call
CacheCreationTokens: 4317, // System prompt + tools written to cache
},
}, "prompt", false)
usage.mu.Lock()
defer usage.mu.Unlock()
// Context tokens should be InputTokens only (1000), not input+output (1200)
// because InputTokens already includes the full conversation history
if usage.contextCalls != 1 || usage.lastContextTokens != 1000 {
t.Fatalf("expected context tokens=1000 (InputTokens only), got calls=%d tokens=%d",
usage.contextCalls, usage.lastContextTokens)
// Context tokens should be Input + CacheRead + CacheCreate + Output = 4325
expected := 3 + 0 + 4317 + 5
if usage.contextCalls != 1 || usage.lastContextTokens != expected {
t.Fatalf("expected context tokens=%d (all categories), got calls=%d tokens=%d",
expected, usage.contextCalls, usage.lastContextTokens)
}
}
// TestHandleTurnEnd_LengthEmitsWarning verifies that when the SDK reports a
// FinishReasonLength (max_output_tokens hit), the app surfaces a user-visible
// ExtensionPrintEvent with Level="info" so the TUI can render a banner
// instead of silently showing a truncated reply.
func TestHandleTurnEnd_LengthEmitsWarning(t *testing.T) {
app := New(Options{}, nil)
defer app.Close()
var mu sync.Mutex
var received []tea.Msg
sendFn := func(m tea.Msg) {
mu.Lock()
defer mu.Unlock()
received = append(received, m)
}
app.handleTurnEnd(kit.TurnEndEvent{StopReason: kit.FinishReasonLength}, sendFn)
mu.Lock()
defer mu.Unlock()
if len(received) != 1 {
t.Fatalf("expected 1 event on length stop, got %d", len(received))
}
ev, ok := received[0].(ExtensionPrintEvent)
if !ok {
t.Fatalf("expected ExtensionPrintEvent, got %T", received[0])
}
if ev.Level != "info" {
t.Errorf("expected Level=info, got %q", ev.Level)
}
if ev.Text == "" {
t.Error("expected non-empty warning text")
}
if !strings.Contains(ev.Text, "max_output_tokens") {
t.Errorf("warning text should mention max_output_tokens, got: %s", ev.Text)
}
}
// TestHandleTurnEnd_NonLengthIgnored verifies that ordinary stop reasons
// (stop, tool-calls, error, unknown, "") do not produce a warning banner.
func TestHandleTurnEnd_NonLengthIgnored(t *testing.T) {
app := New(Options{}, nil)
defer app.Close()
reasons := []string{
kit.FinishReasonStop,
kit.FinishReasonToolCalls,
kit.FinishReasonError,
kit.FinishReasonContentFilter,
kit.FinishReasonOther,
kit.FinishReasonUnknown,
"",
}
for _, r := range reasons {
var called bool
app.handleTurnEnd(kit.TurnEndEvent{StopReason: r}, func(m tea.Msg) {
called = true
})
if called {
t.Errorf("stop reason %q unexpectedly emitted a warning", r)
}
}
}
// TestHandleTurnEnd_NilSendFn guards against panics when no TUI listener is
// attached (e.g. early init or headless teardown).
func TestHandleTurnEnd_NilSendFn(t *testing.T) {
app := New(Options{}, nil)
defer app.Close()
// Should not panic with a nil sendFn.
app.handleTurnEnd(kit.TurnEndEvent{StopReason: kit.FinishReasonLength}, nil)
}
// TestFormatMaxTokensTruncatedMessage_NoKit verifies the fallback message
// when Options.Kit is nil (test/stub path).
func TestFormatMaxTokensTruncatedMessage_NoKit(t *testing.T) {
app := New(Options{}, nil)
defer app.Close()
msg := app.formatMaxTokensTruncatedMessage()
if msg == "" {
t.Fatal("expected non-empty fallback message")
}
for _, needle := range []string{"max_output_tokens", "--max-tokens", "KIT_MAX_TOKENS", "modelSettings"} {
if !strings.Contains(msg, needle) {
t.Errorf("fallback message missing %q:\n%s", needle, msg)
}
}
}
+48
View File
@@ -32,6 +32,36 @@ type ToolCallStartedEvent struct {
ToolArgs string
}
// ToolCallInputStartEvent is sent when the LLM begins generating tool call
// arguments. The tool name is known but the full argument JSON is still being
// streamed. UIs can use this to show a "running" indicator immediately instead
// of waiting for the full argument JSON to finish streaming.
type ToolCallInputStartEvent struct {
// ToolCallID is the stable identifier for correlating tool lifecycle events.
ToolCallID string
// ToolName is the name of the tool being called.
ToolName string
// ToolKind classifies the tool: "execute", "edit", "read", "search", "agent".
ToolKind string
}
// ToolCallInputDeltaEvent is sent for each streamed fragment of tool call
// arguments as they arrive from the LLM. Useful for live-previewing content
// or showing a progress indicator with byte count.
type ToolCallInputDeltaEvent struct {
// ToolCallID is the stable identifier for correlating tool lifecycle events.
ToolCallID string
// Delta is a JSON fragment of tool call arguments.
Delta string
}
// ToolCallInputEndEvent is sent when tool argument streaming is complete,
// before the tool call is parsed and execution begins.
type ToolCallInputEndEvent struct {
// ToolCallID is the stable identifier for correlating tool lifecycle events.
ToolCallID string
}
// ToolExecutionEvent is sent when a tool starts or finishes executing.
// The IsStarting flag distinguishes between the start and end of execution.
type ToolExecutionEvent struct {
@@ -79,6 +109,24 @@ type ToolCallContentEvent struct {
Content string
}
// PasswordPromptEvent is sent when a sudo command needs a password.
// The TUI should display a password prompt overlay and send the result back.
type PasswordPromptEvent struct {
// Prompt is the message to display to the user.
Prompt string
// ResponseCh receives the password from the TUI.
// The TUI must send exactly one value.
ResponseCh chan<- PasswordPromptResponse
}
// PasswordPromptResponse carries the user's password input.
type PasswordPromptResponse struct {
// Password is the entered password.
Password string
// Cancelled is true if the user cancelled the prompt.
Cancelled bool
}
// ResponseCompleteEvent is sent when the LLM produces a final (non-streaming) response.
// In streaming mode, this may be empty if all content was delivered via StreamChunkEvents.
type ResponseCompleteEvent struct {
+3 -6
View File
@@ -3,24 +3,21 @@ package app
import (
"testing"
"charm.land/fantasy"
kit "github.com/mark3labs/kit/pkg/kit"
)
// makeTextMsg builds a minimal kit.LLMMessage using fantasy.NewUserMessage
// or constructing with the given role.
// makeTextMsg builds a minimal kit.LLMMessage with the given role and text.
func makeTextMsg(role, text string) kit.LLMMessage {
return kit.LLMMessage{
Role: kit.LLMMessageRole(role),
Content: []fantasy.MessagePart{fantasy.TextPart{Text: text}},
Content: []kit.LLMMessagePart{kit.LLMTextPart{Text: text}},
}
}
// textOf extracts the plain text from an LLMMessage for assertions.
func textOf(msg kit.LLMMessage) string {
for _, part := range msg.Content {
if tp, ok := part.(fantasy.TextPart); ok {
if tp, ok := part.(kit.LLMTextPart); ok {
return tp.Text
}
}
+4 -2
View File
@@ -21,8 +21,10 @@ type UsageUpdater interface {
// the provider does not return exact counts.
EstimateAndUpdateUsage(inputText, outputText string)
// SetContextTokens records the approximate current context window fill
// level. This should be the final API call's input+output tokens (from
// FinalResponse.Usage), NOT the aggregate TotalUsage.
// level. This should be the sum of ALL token categories from the last
// API call: InputTokens + CacheReadTokens + CacheCreationTokens +
// OutputTokens. With Anthropic prompt caching, InputTokens can be
// near-zero while CacheReadTokens holds the bulk of the context.
SetContextTokens(tokens int)
}
+8
View File
@@ -471,5 +471,13 @@ func GetAnthropicAPIKey(flagValue string) (string, string, error) {
return envKey, "ANTHROPIC_API_KEY environment variable", nil
}
// Check if OpenAI credentials exist to provide a helpful suggestion
if cm != nil {
hasOpenAI, _ := cm.HasOpenAICredentials()
if hasOpenAI {
return "", "", fmt.Errorf("no Anthropic API key found. Use 'kit auth login anthropic', set ANTHROPIC_API_KEY environment variable, or use --provider-api-key flag\n\nNote: OpenAI credentials were detected. To use OpenAI, run with --model openai/gpt-5.4 or set it as default:\n kit auth login openai --set-default")
}
}
return "", "", fmt.Errorf("no Anthropic API key found. Use 'kit auth login anthropic', set ANTHROPIC_API_KEY environment variable, or use --provider-api-key flag")
}
+113 -9
View File
@@ -22,6 +22,28 @@ type MCPServerConfig struct {
AllowedTools []string `json:"allowedTools,omitempty" yaml:"allowedTools,omitempty"`
ExcludedTools []string `json:"excludedTools,omitempty" yaml:"excludedTools,omitempty"`
// OAuth configuration for remote servers that don't support dynamic
// client registration (e.g. GitHub). When OAuthClientID is set, it is
// passed directly to the transport's OAuthConfig instead of relying on
// dynamic registration.
OAuthClientID string `json:"oauthClientId,omitempty" yaml:"oauthClientId,omitempty"`
OAuthClientSecret string `json:"oauthClientSecret,omitempty" yaml:"oauthClientSecret,omitempty"`
OAuthScopes []string `json:"oauthScopes,omitempty" yaml:"oauthScopes,omitempty"`
// NoOAuth disables OAuth transport configuration for this server, even
// when the connection pool has an auth handler. Use this for public MCP
// servers (e.g. PubMed) that don't require authentication. Without this
// flag, the pool would attach OAuth transport to every remote server,
// causing proactive dynamic-client-registration attempts that fail on
// servers that don't support it.
NoOAuth bool `json:"noOAuth,omitempty" yaml:"noOAuth,omitempty"`
// InProcessServer holds a live *server.MCPServer for in-process transport.
// When set (and Type is "inprocess"), the connection pool creates an
// in-process client instead of spawning a subprocess or making HTTP calls.
// This field is never serialized — it is only used programmatically via the SDK.
InProcessServer any `json:"-" yaml:"-"`
// Legacy fields for backward compatibility
Transport string `json:"transport,omitempty"`
Args []string `json:"args,omitempty"`
@@ -35,13 +57,17 @@ type MCPServerConfig struct {
func (s *MCPServerConfig) UnmarshalJSON(data []byte) error {
// First try to unmarshal as the new format
type newFormat struct {
Type string `json:"type"`
Command []string `json:"command,omitempty"`
Environment map[string]string `json:"environment,omitempty"`
URL string `json:"url,omitempty"`
Headers []string `json:"headers,omitempty"`
AllowedTools []string `json:"allowedTools,omitempty" yaml:"allowedTools,omitempty"`
ExcludedTools []string `json:"excludedTools,omitempty" yaml:"excludedTools,omitempty"`
Type string `json:"type"`
Command []string `json:"command,omitempty"`
Environment map[string]string `json:"environment,omitempty"`
URL string `json:"url,omitempty"`
Headers []string `json:"headers,omitempty"`
AllowedTools []string `json:"allowedTools,omitempty" yaml:"allowedTools,omitempty"`
ExcludedTools []string `json:"excludedTools,omitempty" yaml:"excludedTools,omitempty"`
OAuthClientID string `json:"oauthClientId,omitempty" yaml:"oauthClientId,omitempty"`
OAuthClientSecret string `json:"oauthClientSecret,omitempty" yaml:"oauthClientSecret,omitempty"`
OAuthScopes []string `json:"oauthScopes,omitempty" yaml:"oauthScopes,omitempty"`
NoOAuth bool `json:"noOAuth,omitempty" yaml:"noOAuth,omitempty"`
}
// Also try legacy format
@@ -66,6 +92,10 @@ func (s *MCPServerConfig) UnmarshalJSON(data []byte) error {
s.Headers = newConfig.Headers
s.AllowedTools = newConfig.AllowedTools
s.ExcludedTools = newConfig.ExcludedTools
s.OAuthClientID = newConfig.OAuthClientID
s.OAuthClientSecret = newConfig.OAuthClientSecret
s.OAuthScopes = newConfig.OAuthScopes
s.NoOAuth = newConfig.NoOAuth
return nil
}
@@ -157,6 +187,21 @@ type Theme struct {
Markdown MarkdownThemeConfig `json:"markdown,omitzero" yaml:"markdown,omitempty"`
}
// GenerationParams defines generation parameter defaults that can be attached
// to individual models. These act as model-level defaults — CLI flags and
// global config values take precedence when explicitly set.
type GenerationParams struct {
MaxTokens *int `json:"maxTokens,omitempty" yaml:"maxTokens,omitempty"`
Temperature *float32 `json:"temperature,omitempty" yaml:"temperature,omitempty"`
TopP *float32 `json:"topP,omitempty" yaml:"topP,omitempty"`
TopK *int32 `json:"topK,omitempty" yaml:"topK,omitempty"`
FrequencyPenalty *float32 `json:"frequencyPenalty,omitempty" yaml:"frequencyPenalty,omitempty"`
PresencePenalty *float32 `json:"presencePenalty,omitempty" yaml:"presencePenalty,omitempty"`
StopSequences []string `json:"stopSequences,omitempty" yaml:"stopSequences,omitempty"`
ThinkingLevel string `json:"thinkingLevel,omitempty" yaml:"thinkingLevel,omitempty"`
SystemPrompt string `json:"systemPrompt,omitempty" yaml:"systemPrompt,omitempty"`
}
// CustomModelConfig defines a custom model that can be used with custom/custom
// or other custom/ prefixed models. These models are loaded from the config file
// and merged into the custom provider in the model registry.
@@ -171,6 +216,11 @@ type CustomModelConfig struct {
Knowledge string `json:"knowledge,omitempty" yaml:"knowledge,omitempty"`
Cost CostConfig `json:"cost" yaml:"cost"`
Limit LimitConfig `json:"limit" yaml:"limit"`
// Generation parameter defaults for this model.
// These are applied when the user hasn't explicitly set the corresponding
// CLI flag or global config value.
Params GenerationParams `json:"params,omitzero" yaml:"params,omitempty"`
}
// CostConfig defines the pricing for a custom model.
@@ -219,6 +269,12 @@ type Config struct {
// Custom model definitions (under custom/ provider)
CustomModels map[string]CustomModelConfig `json:"customModels,omitempty" yaml:"customModels,omitempty"`
// Per-model generation parameter overrides. Keys are "provider/model" strings
// (e.g. "anthropic/claude-sonnet-4-5-20250929", "openai/gpt-4o"). These
// settings act as model-level defaults — CLI flags and global config values
// take precedence when explicitly set.
ModelSettings map[string]GenerationParams `json:"modelSettings,omitempty" yaml:"modelSettings,omitempty"`
}
// GetTransportType returns the transport type for the server config, mapping
@@ -237,11 +293,18 @@ func (s *MCPServerConfig) GetTransportType() string {
return "stdio"
case "remote":
return "streamable"
case "inprocess":
return "inprocess"
default:
return s.Type
}
}
// Programmatic in-process server detection.
if s.InProcessServer != nil {
return "inprocess"
}
// Backward compatibility: infer transport type
if len(s.Command) > 0 {
return "stdio"
@@ -272,8 +335,12 @@ func (c *Config) Validate() error {
if serverConfig.URL == "" {
return fmt.Errorf("server %s: url is required for %s transport", serverName, transport)
}
case "inprocess":
if serverConfig.InProcessServer == nil {
return fmt.Errorf("server %s: InProcessServer is required for inprocess transport", serverName)
}
default:
return fmt.Errorf("server %s: unsupported transport type '%s'. Supported types: stdio, sse, streamable", serverName, transport)
return fmt.Errorf("server %s: unsupported transport type '%s'. Supported types: stdio, sse, streamable, inprocess", serverName, transport)
}
}
return nil
@@ -367,7 +434,7 @@ mcpServers:
# debug: false # Enable debug logging
# system-prompt: "/path/to/system-prompt.txt" # System prompt text file
# Model generation parameters (all optional)
# Model generation parameters (all optional, apply globally to all models)
# max-tokens: 4096 # Maximum tokens in response
# temperature: 0.7 # Randomness (0.0-1.0)
# top-p: 0.95 # Nucleus sampling (0.0-1.0)
@@ -376,9 +443,46 @@ mcpServers:
# presence-penalty: 0.0 # Penalize present tokens (0.0-2.0)
# stop-sequences: ["Human:", "Assistant:"] # Custom stop sequences
# Per-model generation parameter overrides (apply to specific models)
# These act as model-level defaults — CLI flags and global settings above take precedence.
# Keys are "provider/model" strings matching the model you use.
# modelSettings:
# anthropic/claude-sonnet-4-5-20250929:
# temperature: 0.3
# maxTokens: 8192
# openai/gpt-4o:
# temperature: 0.7
# topP: 0.95
# topK: 40
# frequencyPenalty: 0.1
# presencePenalty: 0.1
# anthropic/claude-opus-4-6:
# thinkingLevel: "high"
# maxTokens: 16384
# systemPrompt: "You are a deep reasoning assistant." # or a file path
# API Configuration (can also use environment variables)
# provider-api-key: "your-api-key" # API key for OpenAI, Anthropic, or Google
# provider-url: "https://api.openai.com/v1" # Base URL for OpenAI, Anthropic, or Ollama
# Custom model definitions (under custom/ provider)
# customModels:
# my-local-llama:
# name: "Local Llama 3"
# baseUrl: "http://localhost:8080/v1"
# family: "llama"
# temperature: true
# cost:
# input: 0.0
# output: 0.0
# limit:
# context: 131072
# output: 8192
# params: # Generation parameter defaults for this model
# temperature: 0.8
# topP: 0.95
# topK: 40
# systemPrompt: "You are a helpful local assistant."
`
_, err = file.WriteString(content)
+85
View File
@@ -6,6 +6,8 @@ import (
"path/filepath"
"strings"
"testing"
"gopkg.in/yaml.v3"
)
func TestMCPServerConfig_NewFormat(t *testing.T) {
@@ -542,3 +544,86 @@ func TestEnsureConfigExistsWhenFileExists(t *testing.T) {
t.Error("Existing config file was modified when it shouldn't have been")
}
}
func TestMCPServerConfig_OAuthFields_JSON(t *testing.T) {
jsonData := `{
"type": "remote",
"url": "https://api.githubcopilot.com/mcp/",
"oauthClientId": "Ov23liXXXXXXXXXXXXXX",
"oauthClientSecret": "secret123",
"oauthScopes": ["read:user", "repo"]
}`
var cfg MCPServerConfig
err := json.Unmarshal([]byte(jsonData), &cfg)
if err != nil {
t.Fatalf("Failed to unmarshal: %v", err)
}
if cfg.Type != "remote" {
t.Errorf("Expected type 'remote', got %q", cfg.Type)
}
if cfg.URL != "https://api.githubcopilot.com/mcp/" {
t.Errorf("Expected URL, got %q", cfg.URL)
}
if cfg.OAuthClientID != "Ov23liXXXXXXXXXXXXXX" {
t.Errorf("Expected OAuthClientID 'Ov23liXXXXXXXXXXXXXX', got %q", cfg.OAuthClientID)
}
if cfg.OAuthClientSecret != "secret123" {
t.Errorf("Expected OAuthClientSecret 'secret123', got %q", cfg.OAuthClientSecret)
}
if len(cfg.OAuthScopes) != 2 || cfg.OAuthScopes[0] != "read:user" || cfg.OAuthScopes[1] != "repo" {
t.Errorf("Expected OAuthScopes [read:user, repo], got %v", cfg.OAuthScopes)
}
}
func TestMCPServerConfig_OAuthFields_YAML(t *testing.T) {
yamlData := `
type: remote
url: https://api.githubcopilot.com/mcp/
oauthClientId: "Ov23liXXXXXXXXXXXXXX"
oauthScopes:
- read:user
- repo
`
var cfg MCPServerConfig
err := yaml.Unmarshal([]byte(yamlData), &cfg)
if err != nil {
t.Fatalf("Failed to unmarshal YAML: %v", err)
}
if cfg.Type != "remote" {
t.Errorf("Expected type 'remote', got %q", cfg.Type)
}
if cfg.OAuthClientID != "Ov23liXXXXXXXXXXXXXX" {
t.Errorf("Expected OAuthClientID 'Ov23liXXXXXXXXXXXXXX', got %q", cfg.OAuthClientID)
}
if len(cfg.OAuthScopes) != 2 || cfg.OAuthScopes[0] != "read:user" || cfg.OAuthScopes[1] != "repo" {
t.Errorf("Expected OAuthScopes [read:user, repo], got %v", cfg.OAuthScopes)
}
}
func TestMCPServerConfig_OAuthFields_Omitted(t *testing.T) {
// Verify that omitting OAuth fields still works (backward compat).
jsonData := `{
"type": "remote",
"url": "https://example.com/mcp"
}`
var cfg MCPServerConfig
err := json.Unmarshal([]byte(jsonData), &cfg)
if err != nil {
t.Fatalf("Failed to unmarshal: %v", err)
}
if cfg.OAuthClientID != "" {
t.Errorf("Expected empty OAuthClientID, got %q", cfg.OAuthClientID)
}
if cfg.OAuthClientSecret != "" {
t.Errorf("Expected empty OAuthClientSecret, got %q", cfg.OAuthClientSecret)
}
if len(cfg.OAuthScopes) != 0 {
t.Errorf("Expected empty OAuthScopes, got %v", cfg.OAuthScopes)
}
}
+176 -6
View File
@@ -19,10 +19,18 @@ import (
// It receives tool call ID, tool name, output chunk, and whether it's stderr.
type ToolOutputCallback func(toolCallID, toolName, chunk string, isStderr bool)
// PasswordPromptCallback is the signature for password prompts.
// It receives a prompt message and returns the password and whether it was cancelled.
type PasswordPromptCallback func(prompt string) (password string, cancelled bool)
// contextKey is a custom type for context keys to avoid collisions.
type contextKey string
const toolOutputCallbackKey contextKey = "toolOutputCallback"
const (
toolOutputCallbackKey contextKey = "toolOutputCallback"
sudoPasswordKey contextKey = "sudoPassword"
passwordPromptKey contextKey = "passwordPrompt"
)
// ContextWithToolOutputCallback returns a new context with the tool output callback set.
func ContextWithToolOutputCallback(ctx context.Context, callback ToolOutputCallback) context.Context {
@@ -37,6 +45,34 @@ func toolOutputCallbackFromContext(ctx context.Context) ToolOutputCallback {
return nil
}
// ContextWithPasswordPrompt returns a new context with the password prompt callback set.
// This allows the TUI to show a modal password prompt when sudo needs a password.
func ContextWithPasswordPrompt(ctx context.Context, callback PasswordPromptCallback) context.Context {
return context.WithValue(ctx, passwordPromptKey, callback)
}
// passwordPromptFromContext retrieves the password prompt callback from context.
func passwordPromptFromContext(ctx context.Context) PasswordPromptCallback {
if cb, ok := ctx.Value(passwordPromptKey).(PasswordPromptCallback); ok {
return cb
}
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 {
return pw
}
return ""
}
const defaultBashTimeout = 120 * time.Second
const maxBashTimeout = 600 * time.Second
@@ -73,6 +109,66 @@ func NewBashTool(opts ...ToolOption) fantasy.AgentTool {
}
}
// sudoCommandRe matches sudo commands that need to be rewritten for -S mode.
// It matches "sudo" as a word boundary, optionally preceded by environment variables.
var sudoCommandRe = regexp.MustCompile(`(?i)(^|[&|;|]|\|\||&&)\s*(\w+=\S+\s+)?\bsudo\b`)
// truncateCommand truncates a long command for display.
func truncateCommand(cmd string, maxLen int) string {
if len(cmd) <= maxLen {
return cmd
}
return cmd[:maxLen-3] + "..."
}
// rewriteSudoForStdin rewrites sudo commands to use -S -p ” for stdin password input.
// It transforms: sudo cmd → sudo -S -p ” cmd
func rewriteSudoForStdin(command string) string {
// Find all matches and their positions
matches := sudoCommandRe.FindAllStringIndex(command, -1)
if matches == nil {
return command
}
// Build result from end to start to preserve indices
result := command
for i := len(matches) - 1; i >= 0; i-- {
match := matches[i]
start, end := match[0], match[1]
matchedText := result[start:end]
// Extract just the "sudo" part (after any prefix)
sudoIdx := strings.Index(strings.ToLower(matchedText), "sudo")
if sudoIdx == -1 {
continue
}
prefix := matchedText[:sudoIdx]
sudoPart := matchedText[sudoIdx:]
// Check if the text immediately after "sudo" in the result contains -S
afterSudo := result[end:]
if strings.HasPrefix(strings.TrimLeft(afterSudo, " \t"), "-S") {
// Already has -S flag, skip
continue
}
// Insert -S -p '' after "sudo"
newSudo := strings.Replace(sudoPart, "sudo", "sudo -S -p ''", 1)
result = result[:start] + prefix + newSudo + result[end:]
}
return result
}
// SudoPasswordRequiredResult is a special marker that indicates sudo needs a password.
// This is stored in tool response metadata to signal the TUI to prompt for password.
const SudoPasswordRequiredMetadata = `{"sudo_password_required":true}`
// IsSudoPasswordRequiredResult checks if a tool response indicates sudo password is needed.
func IsSudoPasswordRequiredResult(resp fantasy.ToolResponse) bool {
return resp.Metadata == SudoPasswordRequiredMetadata
}
func executeBash(ctx context.Context, call fantasy.ToolCall, workDir string) (fantasy.ToolResponse, error) {
var args bashArgs
if err := parseArgs(call.Input, &args); err != nil {
@@ -97,7 +193,47 @@ func executeBash(ctx context.Context, call fantasy.ToolCall, workDir string) (fa
cmdCtx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
cmd := exec.CommandContext(cmdCtx, "bash", "-c", args.Command)
// Check for sudo password in context or environment
sudoPassword := sudoPasswordFromContext(ctx)
if sudoPassword == "" {
sudoPassword = os.Getenv("SUDO_PASSWORD")
}
command := args.Command
// If command contains sudo and we don't have a password, check if sudo needs one
if sudoPassword == "" && sudoCommandRe.MatchString(command) {
// Check if sudo credentials are cached using sudo -n (non-interactive)
testCmd := exec.CommandContext(cmdCtx, "sudo", "-n", "true")
testCmd.Dir = workDir
if err := testCmd.Run(); err != nil {
// Sudo needs a password - try to prompt via callback
if promptCallback := passwordPromptFromContext(ctx); promptCallback != nil {
pw, cancelled := promptCallback("Sudo password required for: " + truncateCommand(args.Command, 60))
if cancelled {
return fantasy.NewTextErrorResponse("sudo password prompt cancelled"), nil
}
if pw == "" {
return fantasy.NewTextErrorResponse("no sudo password provided"), nil
}
sudoPassword = pw
command = rewriteSudoForStdin(command)
} else {
// No callback available - return error with helpful message
return fantasy.NewTextErrorResponse(
"This command requires sudo access. " +
"Please run 'sudo -v' in your terminal first to cache credentials, " +
"or set the SUDO_PASSWORD environment variable."), nil
}
}
// Credentials are cached or password was provided, proceed
}
// If we have a sudo password, rewrite the command to use sudo -S
if sudoPassword != "" && sudoCommandRe.MatchString(command) {
command = rewriteSudoForStdin(command)
}
cmd := exec.CommandContext(cmdCtx, "bash", "-c", command)
if workDir != "" {
cmd.Dir = workDir
}
@@ -115,18 +251,18 @@ func executeBash(ctx context.Context, call fantasy.ToolCall, workDir string) (fa
if outputCallback != nil {
// Streaming mode: use pipes to capture output as it arrives
return executeBashStreaming(cmdCtx, call, cmd, outputCallback)
return executeBashStreaming(cmdCtx, call, cmd, outputCallback, sudoPassword)
}
// Non-streaming mode: collect all output at once (original behavior)
return executeBashBuffered(cmdCtx, call, cmd)
return executeBashBuffered(cmdCtx, call, cmd, sudoPassword)
}
// executeBashBuffered collects all output before returning (original behavior).
// It uses explicit pipes (not cmd.Stdout) so that cmd.WaitDelay can forcibly
// close them when grandchild processes hold pipe handles open after the
// direct child exits.
func executeBashBuffered(cmdCtx context.Context, call fantasy.ToolCall, cmd *exec.Cmd) (fantasy.ToolResponse, error) {
func executeBashBuffered(cmdCtx context.Context, call fantasy.ToolCall, cmd *exec.Cmd, sudoPassword string) (fantasy.ToolResponse, error) {
stdoutPipe, err := cmd.StdoutPipe()
if err != nil {
return fantasy.NewTextErrorResponse("failed to create stdout pipe"), nil
@@ -136,10 +272,27 @@ func executeBashBuffered(cmdCtx context.Context, call fantasy.ToolCall, cmd *exe
return fantasy.NewTextErrorResponse("failed to create stderr pipe"), nil
}
// If we have a sudo password, create a stdin pipe and write the password
var stdinPipe io.WriteCloser
if sudoPassword != "" {
stdinPipe, err = cmd.StdinPipe()
if err != nil {
return fantasy.NewTextErrorResponse("failed to create stdin pipe"), nil
}
}
if err := cmd.Start(); err != nil {
return fantasy.NewTextErrorResponse(fmt.Sprintf("failed to start command: %v", err)), nil
}
// Write password to stdin if needed, then close stdin
if sudoPassword != "" && stdinPipe != nil {
go func() {
defer func() { _ = stdinPipe.Close() }()
_, _ = io.WriteString(stdinPipe, sudoPassword+"\n")
}()
}
// Read pipes concurrently
var wg sync.WaitGroup
var stdout, stderr strings.Builder
@@ -181,7 +334,7 @@ func executeBashBuffered(cmdCtx context.Context, call fantasy.ToolCall, cmd *exe
}
// executeBashStreaming streams output as it arrives via the callback.
func executeBashStreaming(cmdCtx context.Context, call fantasy.ToolCall, cmd *exec.Cmd, outputCallback ToolOutputCallback) (fantasy.ToolResponse, error) {
func executeBashStreaming(cmdCtx context.Context, call fantasy.ToolCall, cmd *exec.Cmd, outputCallback ToolOutputCallback, sudoPassword string) (fantasy.ToolResponse, error) {
stdoutPipe, err := cmd.StdoutPipe()
if err != nil {
return fantasy.NewTextErrorResponse("failed to create stdout pipe"), nil
@@ -191,11 +344,28 @@ func executeBashStreaming(cmdCtx context.Context, call fantasy.ToolCall, cmd *ex
return fantasy.NewTextErrorResponse("failed to create stderr pipe"), nil
}
// If we have a sudo password, create a stdin pipe
var stdinPipe io.WriteCloser
if sudoPassword != "" {
stdinPipe, err = cmd.StdinPipe()
if err != nil {
return fantasy.NewTextErrorResponse("failed to create stdin pipe"), nil
}
}
// Start command execution
if err := cmd.Start(); err != nil {
return fantasy.NewTextErrorResponse(fmt.Sprintf("failed to start command: %v", err)), nil
}
// Write password to stdin if needed, then close stdin
if sudoPassword != "" && stdinPipe != nil {
go func() {
defer func() { _ = stdinPipe.Close() }()
_, _ = io.WriteString(stdinPipe, sudoPassword+"\n")
}()
}
// Stream stdout and stderr concurrently
var wg sync.WaitGroup
var mu sync.Mutex
+69
View File
@@ -127,3 +127,72 @@ func TestBash_EmptyCommand(t *testing.T) {
t.Fatal("expected error for empty command")
}
}
func TestRewriteSudoForStdin(t *testing.T) {
tests := []struct {
name string
input string
expected string
}{
{
name: "simple sudo",
input: "sudo apt update",
expected: "sudo -S -p '' apt update",
},
{
name: "sudo with env var",
input: "DEBIAN_FRONTEND=noninteractive sudo apt update",
expected: "DEBIAN_FRONTEND=noninteractive sudo -S -p '' apt update",
},
{
name: "sudo in pipeline",
input: "echo test | sudo tee /etc/test.conf",
expected: "echo test | sudo -S -p '' tee /etc/test.conf",
},
{
name: "sudo after &&",
input: "apt update && sudo apt upgrade",
expected: "apt update && sudo -S -p '' apt upgrade",
},
{
name: "already has -S flag",
input: "sudo -S apt update",
expected: "sudo -S apt update",
},
{
name: "no sudo",
input: "apt update && apt upgrade",
expected: "apt update && apt upgrade",
},
{
name: "sudo in string (should not match)",
input: "echo 'use sudo carefully'",
expected: "echo 'use sudo carefully'",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := rewriteSudoForStdin(tt.input)
if result != tt.expected {
t.Errorf("rewriteSudoForStdin(%q) = %q, want %q", tt.input, result, tt.expected)
}
})
}
}
func TestSudoPasswordFromContext(t *testing.T) {
// Test with password in context
ctx := ContextWithSudoPassword(context.Background(), "secret123")
pw := sudoPasswordFromContext(ctx)
if pw != "secret123" {
t.Errorf("expected password 'secret123', got %q", pw)
}
// Test without password
ctx = context.Background()
pw = sudoPasswordFromContext(ctx)
if pw != "" {
t.Errorf("expected empty password, got %q", pw)
}
}
+2 -2
View File
@@ -86,7 +86,7 @@ Example use cases:
},
"model": map[string]any{
"type": "string",
"description": "Optional model override (e.g. 'anthropic/claude-haiku-3-5-20241022' for faster/cheaper tasks)",
"description": "Optional model override. Empty string uses the current model.",
},
"system_prompt": map[string]any{
"type": "string",
@@ -94,7 +94,7 @@ Example use cases:
},
"timeout_seconds": map[string]any{
"type": "number",
"description": "Maximum execution time in seconds (default: 300, max: 1800)",
"description": "Maximum execution time in seconds (default: 300, max: 1800, minimum recommended: 240)",
},
},
Required: []string{"task"},
+52 -1
View File
@@ -918,7 +918,7 @@ type ExtensionEntry struct {
type ContextMessage struct {
// Index is the position of this message in the original context array
// (0-based). When returning messages from a ContextPrepareResult,
// messages with Index >= 0 reuse the original fantasy.Message at that
// messages with Index >= 0 reuse the original LLM message at that
// position (preserving tool calls, reasoning, and other complex parts).
// Set Index to -1 for newly injected messages (created from Role + Content).
Index int
@@ -1063,6 +1063,9 @@ type PrintBlockOpts struct {
type API struct {
// Event-specific registration functions (wired by the loader).
onToolCall func(func(ToolCallEvent, Context) *ToolCallResult)
onToolCallInputStart func(func(ToolCallInputStartEvent, Context))
onToolCallInputDelta func(func(ToolCallInputDeltaEvent, Context))
onToolCallInputEnd func(func(ToolCallInputEndEvent, Context))
onToolExecStart func(func(ToolExecutionStartEvent, Context))
onToolExecEnd func(func(ToolExecutionEndEvent, Context))
onToolOutput func(func(ToolOutputEvent, Context))
@@ -1099,6 +1102,26 @@ func (a *API) OnToolCall(handler func(ToolCallEvent, Context) *ToolCallResult) {
a.onToolCall(handler)
}
// OnToolCallInputStart registers a handler that fires when the LLM begins
// generating tool call arguments. The tool name is known but the full
// argument JSON is still being streamed. Useful for showing a "running"
// indicator immediately without waiting for the full arguments.
func (a *API) OnToolCallInputStart(handler func(ToolCallInputStartEvent, Context)) {
a.onToolCallInputStart(handler)
}
// OnToolCallInputDelta registers a handler that fires for each streamed
// fragment of tool call arguments as they arrive from the LLM.
func (a *API) OnToolCallInputDelta(handler func(ToolCallInputDeltaEvent, Context)) {
a.onToolCallInputDelta(handler)
}
// OnToolCallInputEnd registers a handler that fires when tool argument
// streaming is complete, before the tool call is parsed and execution begins.
func (a *API) OnToolCallInputEnd(handler func(ToolCallInputEndEvent, Context)) {
a.onToolCallInputEnd(handler)
}
// OnToolExecutionStart registers a handler for tool execution start.
func (a *API) OnToolExecutionStart(handler func(ToolExecutionStartEvent, Context)) {
a.onToolExecStart(handler)
@@ -1890,6 +1913,34 @@ type ToolCallResult struct {
func (ToolCallResult) isResult() {}
// ToolCallInputStartEvent fires when the LLM begins generating tool call
// arguments. The tool name is known but the full argument JSON is still
// being streamed.
type ToolCallInputStartEvent struct {
ToolCallID string
ToolName string
ToolKind string // Tool classification: "execute", "edit", "read", "search", "agent"
}
func (e ToolCallInputStartEvent) Type() EventType { return ToolCallInputStart }
// ToolCallInputDeltaEvent fires for each streamed fragment of tool call
// arguments as they arrive from the LLM.
type ToolCallInputDeltaEvent struct {
ToolCallID string
Delta string // JSON fragment of tool arguments
}
func (e ToolCallInputDeltaEvent) Type() EventType { return ToolCallInputDelta }
// ToolCallInputEndEvent fires when tool argument streaming is complete,
// before the tool call is parsed and execution begins.
type ToolCallInputEndEvent struct {
ToolCallID string
}
func (e ToolCallInputEndEvent) Type() EventType { return ToolCallInputEnd }
// ToolExecutionStartEvent fires when a tool begins executing.
type ToolExecutionStartEvent struct {
ToolCallID string
+15 -1
View File
@@ -13,6 +13,19 @@ const (
// ToolCall fires before a tool executes. Handlers can block execution.
ToolCall EventType = "tool_call"
// ToolCallInputStart fires when the LLM begins generating tool call
// arguments. The tool name is known but the full argument JSON is still
// being streamed.
ToolCallInputStart EventType = "tool_call_input_start"
// ToolCallInputDelta fires for each streamed fragment of tool call
// arguments as they arrive from the LLM.
ToolCallInputDelta EventType = "tool_call_input_delta"
// ToolCallInputEnd fires when tool argument streaming is complete,
// before the tool call is parsed and execution begins.
ToolCallInputEnd EventType = "tool_call_input_end"
// ToolExecutionStart fires when a tool begins executing.
ToolExecutionStart EventType = "tool_execution_start"
@@ -88,7 +101,8 @@ const (
// AllEventTypes returns every supported event type.
func AllEventTypes() []EventType {
return []EventType{
ToolCall, ToolExecutionStart, ToolExecutionEnd, ToolResult,
ToolCall, ToolCallInputStart, ToolCallInputDelta, ToolCallInputEnd,
ToolExecutionStart, ToolExecutionEnd, ToolResult,
Input, BeforeAgentStart, AgentStart, AgentEnd,
MessageStart, MessageUpdate, MessageEnd,
SessionStart, SessionShutdown,
+5 -2
View File
@@ -4,8 +4,8 @@ import "testing"
func TestAllEventTypes_Count(t *testing.T) {
all := AllEventTypes()
if len(all) != 21 {
t.Fatalf("expected 21 event types, got %d", len(all))
if len(all) != 24 {
t.Fatalf("expected 24 event types, got %d", len(all))
}
}
@@ -38,6 +38,9 @@ func TestEventType_TypeMethod(t *testing.T) {
want EventType
}{
{ToolCallEvent{ToolName: "test"}, ToolCall},
{ToolCallInputStartEvent{ToolCallID: "x", ToolName: "test"}, ToolCallInputStart},
{ToolCallInputDeltaEvent{ToolCallID: "x", Delta: "{"}, ToolCallInputDelta},
{ToolCallInputEndEvent{ToolCallID: "x"}, ToolCallInputEnd},
{ToolExecutionStartEvent{ToolName: "test"}, ToolExecutionStart},
{ToolExecutionEndEvent{ToolName: "test"}, ToolExecutionEnd},
{ToolResultEvent{ToolName: "test"}, ToolResult},
+18
View File
@@ -429,6 +429,24 @@ func loadSingleExtension(path string) (*LoadedExtension, error) {
return *r
})
},
onToolCallInputStart: func(h func(ToolCallInputStartEvent, Context)) {
reg(ToolCallInputStart, func(e Event, c Context) Result {
h(e.(ToolCallInputStartEvent), c)
return nil
})
},
onToolCallInputDelta: func(h func(ToolCallInputDeltaEvent, Context)) {
reg(ToolCallInputDelta, func(e Event, c Context) Result {
h(e.(ToolCallInputDeltaEvent), c)
return nil
})
},
onToolCallInputEnd: func(h func(ToolCallInputEndEvent, Context)) {
reg(ToolCallInputEnd, func(e Event, c Context) Result {
h(e.(ToolCallInputEndEvent), c)
return nil
})
},
onToolExecStart: func(h func(ToolExecutionStartEvent, Context)) {
reg(ToolExecutionStart, func(e Event, c Context) Result {
h(e.(ToolExecutionStartEvent), c)
+92 -3
View File
@@ -1,21 +1,93 @@
package extensions
import (
"bytes"
"fmt"
"log"
"os"
"runtime"
"sort"
"strconv"
"strings"
"sync"
"github.com/spf13/viper"
)
// ---------------------------------------------------------------------------
// reentrantMu — a per-extension mutex that allows the same goroutine to
// re-enter (e.g. handler → ctx.EmitCustomEvent → handler in same extension).
// Different goroutines are serialized, preventing concurrent state mutation.
// ---------------------------------------------------------------------------
type reentrantMu struct {
mu sync.Mutex
cond *sync.Cond
owner int64 // goroutine ID that holds the lock, or 0
depth int // re-entrancy depth
}
// initReentrantMu initializes the reentrant mutex in-place. Must be called
// after the struct is at its final memory location (not before copying).
func (r *reentrantMu) init() {
r.cond = sync.NewCond(&r.mu)
}
// lock acquires the mutex. If the calling goroutine already holds it, the
// call succeeds immediately (re-entrant). Every call to lock must be paired
// with a call to unlock.
func (r *reentrantMu) lock() {
gid := goroutineID()
r.mu.Lock()
if r.owner == gid {
// Re-entrant: same goroutine already holds the lock.
r.depth++
r.mu.Unlock()
return
}
// Wait for the current owner to release.
for r.owner != 0 {
r.cond.Wait() // releases mu, blocks, re-acquires mu on wake
}
r.owner = gid
r.depth = 1
r.mu.Unlock()
}
// unlock releases the mutex (or decrements re-entrancy depth).
func (r *reentrantMu) unlock() {
r.mu.Lock()
r.depth--
if r.depth == 0 {
r.owner = 0
r.cond.Signal()
}
r.mu.Unlock()
}
// goroutineID extracts the current goroutine's ID from runtime.Stack output.
// This is a well-known technique used by Go testing infrastructure.
func goroutineID() int64 {
var buf [64]byte
n := runtime.Stack(buf[:], false)
// Stack output starts with "goroutine NNN ["
s := buf[:n]
s = s[len("goroutine "):]
s = s[:bytes.IndexByte(s, ' ')]
id, _ := strconv.ParseInt(string(s), 10, 64)
return id
}
// Runner manages loaded extensions and dispatches events to their handlers
// sequentially. Handlers execute in extension
// load order; for cancellable events the first blocking result wins.
//
// Each extension has a dedicated reentrant mutex so that handlers for the
// same extension are serialized (preventing data races on shared package-level
// state), while handlers for different extensions may execute concurrently.
type Runner struct {
extensions []LoadedExtension
extMu []reentrantMu // per-extension reentrant mutex, indexed by extension position
ctx Context
widgets map[string]WidgetConfig // keyed by widget ID
statusEntries map[string]StatusBarEntry // keyed by status key
@@ -52,7 +124,11 @@ type LoadedExtension struct {
// NewRunner creates a Runner from a set of loaded extensions.
func NewRunner(exts []LoadedExtension) *Runner {
return &Runner{extensions: exts}
mus := make([]reentrantMu, len(exts))
for i := range mus {
mus[i].init()
}
return &Runner{extensions: exts, extMu: mus}
}
// SetContext updates the runtime context (session ID, model, etc.) that is
@@ -367,6 +443,11 @@ func (r *Runner) Emit(event Event) (Result, error) {
for i := range r.extensions {
ext := &r.extensions[i]
handlers := ext.Handlers[event.Type()]
if len(handlers) == 0 {
continue
}
r.extMu[i].lock()
for _, handler := range handlers {
result, err := safeCall(handler, event, ctx)
if err != nil {
@@ -379,6 +460,7 @@ func (r *Runner) Emit(event Event) (Result, error) {
// Check for blocking/short-circuit results.
if isBlocking(result) {
r.extMu[i].unlock()
return result, nil
}
@@ -386,6 +468,7 @@ func (r *Runner) Emit(event Event) (Result, error) {
// the caller is responsible for applying the modifications.
accumulated = result
}
r.extMu[i].unlock()
}
return accumulated, nil
}
@@ -712,11 +795,17 @@ func (r *Runner) EmitCustomEvent(name, data string) {
// Extension-registered handlers first (in load order).
for i := range r.extensions {
for _, h := range r.extensions[i].CustomEventHandlers[name] {
extHandlers := r.extensions[i].CustomEventHandlers[name]
if len(extHandlers) == 0 {
continue
}
r.extMu[i].lock()
for _, h := range extHandlers {
safeInvoke(h)
}
r.extMu[i].unlock()
}
// Then dynamic subscriptions.
// Then dynamic subscriptions (not extension-scoped, no per-ext lock).
for _, h := range dynamicHandlers {
safeInvoke(h)
}
+140
View File
@@ -1,6 +1,7 @@
package extensions
import (
"sync"
"testing"
)
@@ -571,3 +572,142 @@ func TestRunner_ContextPrintNilSafe(t *testing.T) {
t.Fatalf("unexpected error: %v", err)
}
}
func TestRunner_ConcurrentEmitSameExtension(t *testing.T) {
// Verify that concurrent Emit calls for the same extension are serialized
// and don't cause data races on shared handler state.
var counter int
ext := makeHandlerExt("shared-state.go", map[EventType][]HandlerFunc{
SubagentStart: {
func(e Event, c Context) Result {
// Read-modify-write: racy without serialization.
v := counter
counter = v + 1
return nil
},
},
SubagentChunk: {
func(e Event, c Context) Result {
v := counter
counter = v + 1
return nil
},
},
})
r := makeRunner(ext)
var wg sync.WaitGroup
const goroutines = 20
const iterations = 50
wg.Add(goroutines)
for range goroutines {
go func() {
defer wg.Done()
for range iterations {
_, _ = r.Emit(SubagentStartEvent{ToolCallID: "x"})
_, _ = r.Emit(SubagentChunkEvent{ToolCallID: "x"})
}
}()
}
wg.Wait()
if counter != goroutines*iterations*2 {
t.Errorf("expected counter=%d, got %d (race detected)", goroutines*iterations*2, counter)
}
}
func TestRunner_ConcurrentEmitDifferentExtensions(t *testing.T) {
// Two extensions with independent state should not block each other
// and should both run correctly under concurrent Emit calls.
var counter1, counter2 int
ext1 := makeHandlerExt("ext1.go", map[EventType][]HandlerFunc{
SubagentStart: {
func(e Event, c Context) Result {
v := counter1
counter1 = v + 1
return nil
},
},
})
ext2 := makeHandlerExt("ext2.go", map[EventType][]HandlerFunc{
SubagentStart: {
func(e Event, c Context) Result {
v := counter2
counter2 = v + 1
return nil
},
},
})
r := makeRunner(ext1, ext2)
var wg sync.WaitGroup
const goroutines = 20
const iterations = 50
wg.Add(goroutines)
for range goroutines {
go func() {
defer wg.Done()
for range iterations {
_, _ = r.Emit(SubagentStartEvent{ToolCallID: "x"})
}
}()
}
wg.Wait()
expected := goroutines * iterations
if counter1 != expected {
t.Errorf("ext1 counter: expected %d, got %d", expected, counter1)
}
if counter2 != expected {
t.Errorf("ext2 counter: expected %d, got %d", expected, counter2)
}
}
func TestRunner_ReentrantEmitCustomEvent(t *testing.T) {
// Verify that a handler can call EmitCustomEvent (which dispatches to
// the same extension's custom event handlers) without deadlocking.
var order []string
ext := LoadedExtension{
Path: "reentrant.go",
Handlers: map[EventType][]HandlerFunc{
SessionStart: {
func(e Event, c Context) Result {
order = append(order, "session_start")
// This triggers EmitCustomEvent for the same extension
// via a direct runner call (simulating ctx.EmitCustomEvent).
return nil
},
},
},
CustomEventHandlers: map[string][]func(string){
"test-event": {
func(data string) {
order = append(order, "custom:"+data)
},
},
},
}
r := makeRunner(ext)
// Wire up the handler to call EmitCustomEvent re-entrantly.
ext.Handlers[SessionStart] = []HandlerFunc{
func(e Event, c Context) Result {
order = append(order, "session_start")
r.EmitCustomEvent("test-event", "hello")
return nil
},
}
r.extensions[0] = ext
// Rebuild mutexes after modifying extensions slice.
r.extMu = make([]reentrantMu, len(r.extensions))
for i := range r.extMu {
r.extMu[i].init()
}
_, err := r.Emit(SessionStartEvent{})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if len(order) != 2 || order[0] != "session_start" || order[1] != "custom:hello" {
t.Errorf("expected [session_start, custom:hello], got %v", order)
}
}
+3
View File
@@ -152,6 +152,9 @@ func Symbols() interp.Exports {
// Event structs
"ToolCallEvent": reflect.ValueOf((*ToolCallEvent)(nil)),
"ToolCallResult": reflect.ValueOf((*ToolCallResult)(nil)),
"ToolCallInputStartEvent": reflect.ValueOf((*ToolCallInputStartEvent)(nil)),
"ToolCallInputDeltaEvent": reflect.ValueOf((*ToolCallInputDeltaEvent)(nil)),
"ToolCallInputEndEvent": reflect.ValueOf((*ToolCallInputEndEvent)(nil)),
"ToolExecutionStartEvent": reflect.ValueOf((*ToolExecutionStartEvent)(nil)),
"ToolExecutionEndEvent": reflect.ValueOf((*ToolExecutionEndEvent)(nil)),
"ToolOutputEvent": reflect.ValueOf((*ToolOutputEvent)(nil)),
+5 -5
View File
@@ -28,11 +28,11 @@ func WrapToolsWithExtensions(tools []fantasy.AgentTool, runner *Runner) []fantas
return wrapped
}
// ExtensionToolsAsFantasy converts ToolDef values registered by extensions
// into fantasy.AgentTool implementations so the LLM can invoke them.
// ExtensionToolsAsLLMTools converts ToolDef values registered by extensions
// into LLM agent tool implementations so the LLM can invoke them.
// The runner is optional; if provided, ToolContext.OnProgress routes
// progress messages through the runner's Print function.
func ExtensionToolsAsFantasy(defs []ToolDef, runner *Runner) []fantasy.AgentTool {
func ExtensionToolsAsLLMTools(defs []ToolDef, runner *Runner) []fantasy.AgentTool {
tools := make([]fantasy.AgentTool, 0, len(defs))
for _, def := range defs {
tools = append(tools, &extensionTool{def: def, runner: runner})
@@ -154,7 +154,7 @@ func (w *wrappedTool) Run(ctx context.Context, call fantasy.ToolCall) (fantasy.T
}
// ---------------------------------------------------------------------------
// extensionTool — wraps a ToolDef into a fantasy.AgentTool
// extensionTool — wraps a ToolDef into an LLM agent tool
// ---------------------------------------------------------------------------
type extensionTool struct {
@@ -182,7 +182,7 @@ func (t *extensionTool) Info() fantasy.ToolInfo {
info.Parameters = props
} else {
// Schema doesn't have "properties" — use as-is (may be
// a flat property map already matching fantasy's format).
// a flat property map already matching the expected format).
info.Parameters = schema
}
// Extract required fields if present.
+8 -8
View File
@@ -192,7 +192,7 @@ func TestWrappedTool_ExecutionStartEnd(t *testing.T) {
}
}
func TestExtensionToolsAsFantasy(t *testing.T) {
func TestExtensionToolsAsLLMTools(t *testing.T) {
defs := []ToolDef{
{
Name: "greet",
@@ -202,7 +202,7 @@ func TestExtensionToolsAsFantasy(t *testing.T) {
},
}
tools := ExtensionToolsAsFantasy(defs, nil)
tools := ExtensionToolsAsLLMTools(defs, nil)
if len(tools) != 1 {
t.Fatalf("expected 1 tool, got %d", len(tools))
}
@@ -232,7 +232,7 @@ func TestExtensionTool_Error(t *testing.T) {
},
}
tools := ExtensionToolsAsFantasy(defs, nil)
tools := ExtensionToolsAsLLMTools(defs, nil)
resp, err := tools[0].Run(context.Background(), fantasy.ToolCall{Input: "x"})
if err == nil {
t.Error("expected error")
@@ -259,7 +259,7 @@ func TestExtensionTool_ExecuteWithContext(t *testing.T) {
}
// Without runner, OnProgress is a no-op.
tools := ExtensionToolsAsFantasy(defs, nil)
tools := ExtensionToolsAsLLMTools(defs, nil)
resp, err := tools[0].Run(context.Background(), fantasy.ToolCall{Input: "test"})
if err != nil {
t.Fatalf("unexpected error: %v", err)
@@ -285,7 +285,7 @@ func TestExtensionTool_ExecuteWithContext(t *testing.T) {
},
},
}
tools2 := ExtensionToolsAsFantasy(defs2, runner)
tools2 := ExtensionToolsAsLLMTools(defs2, runner)
_, err = tools2[0].Run(context.Background(), fantasy.ToolCall{Input: ""})
if err != nil {
t.Fatalf("unexpected error: %v", err)
@@ -306,7 +306,7 @@ func TestExtensionTool_ExecuteWithContextPriority(t *testing.T) {
},
},
}
tools := ExtensionToolsAsFantasy(defs, nil)
tools := ExtensionToolsAsLLMTools(defs, nil)
resp, err := tools[0].Run(context.Background(), fantasy.ToolCall{Input: ""})
if err != nil {
t.Fatalf("unexpected error: %v", err)
@@ -330,7 +330,7 @@ func TestExtensionTool_CancelledContext(t *testing.T) {
},
},
}
tools := ExtensionToolsAsFantasy(defs, nil)
tools := ExtensionToolsAsLLMTools(defs, nil)
_, _ = tools[0].Run(ctx, fantasy.ToolCall{Input: ""})
if !sawCancelled {
t.Error("expected IsCancelled=true for cancelled context")
@@ -339,7 +339,7 @@ func TestExtensionTool_CancelledContext(t *testing.T) {
func TestExtensionTool_ProviderOptions(t *testing.T) {
defs := []ToolDef{{Name: "test", Execute: func(string) (string, error) { return "", nil }}}
tools := ExtensionToolsAsFantasy(defs, nil)
tools := ExtensionToolsAsLLMTools(defs, nil)
// Initially nil.
opts := tools[0].ProviderOptions()
+248
View File
@@ -0,0 +1,248 @@
// Package fences provides utilities for detecting markdown code regions
// (fenced code blocks and inline code spans) and applying transformations
// only to text outside those regions.
//
// This prevents special tokens like $1, $@, or @file from being interpreted
// when they appear inside ``` fences, ~~~ fences, or `inline` code spans.
package fences
import "strings"
// Ranges returns byte ranges [start, end) of fenced code blocks in content.
// Recognises both backtick (```) and tilde (~~~) fences, with optional
// leading indentation (up to 3 spaces) and optional info strings.
// An unclosed fence extends to the end of content.
func Ranges(content string) [][2]int {
var result [][2]int
var inFence bool
var fenceChar byte
var fenceCount int
var fenceStart int
pos := 0
for pos < len(content) {
// Find the end of the current line.
lineEnd := strings.IndexByte(content[pos:], '\n')
var line string
var nextPos int
if lineEnd < 0 {
line = content[pos:]
nextPos = len(content)
} else {
line = content[pos : pos+lineEnd]
nextPos = pos + lineEnd + 1
}
trimmed := strings.TrimLeft(line, " ")
indent := len(line) - len(trimmed)
if !inFence {
if indent <= 3 {
if ch, n := parseFenceOpen(trimmed); n > 0 {
inFence = true
fenceChar = ch
fenceCount = n
fenceStart = pos
}
}
} else {
if indent <= 3 && isFenceClose(trimmed, fenceChar, fenceCount) {
result = append(result, [2]int{fenceStart, nextPos})
inFence = false
}
}
pos = nextPos
}
// Unclosed fence extends to end of content.
if inFence {
result = append(result, [2]int{fenceStart, len(content)})
}
return result
}
// ReplaceOutside applies fn to each text segment that is outside fenced code
// blocks and inline code spans, leaving code content unchanged. This is the
// primary entry point for callers that need to do regex replacement only on
// non-code text.
func ReplaceOutside(content string, fn func(string) string) string {
ranges := Ranges(content)
if len(ranges) == 0 {
return replaceOutsideInline(content, fn)
}
var b strings.Builder
b.Grow(len(content))
pos := 0
for _, r := range ranges {
if pos < r[0] {
// Within non-fenced segments, also skip inline code spans.
b.WriteString(replaceOutsideInline(content[pos:r[0]], fn))
}
// Preserve fenced content verbatim.
b.WriteString(content[r[0]:r[1]])
pos = r[1]
}
if pos < len(content) {
b.WriteString(replaceOutsideInline(content[pos:], fn))
}
return b.String()
}
// StripCode returns content with fenced code blocks and inline code spans
// removed. Useful for detection/matching where only non-code text matters.
func StripCode(content string) string {
// First strip fenced blocks.
stripped := StripFenced(content)
// Then strip inline code spans from what remains.
return stripInlineCode(stripped)
}
// StripFenced returns content with fenced code block regions removed.
// Useful for detection/matching where only non-fenced text matters.
// NOTE: this does NOT strip inline code spans; use StripCode for both.
func StripFenced(content string) string {
ranges := Ranges(content)
if len(ranges) == 0 {
return content
}
var b strings.Builder
b.Grow(len(content))
pos := 0
for _, r := range ranges {
b.WriteString(content[pos:r[0]])
pos = r[1]
}
b.WriteString(content[pos:])
return b.String()
}
// parseFenceOpen checks whether trimmed (leading spaces already removed)
// starts a fenced code block. Returns the fence character and count, or
// (0, 0) if it is not a fence opener.
func parseFenceOpen(trimmed string) (byte, int) {
if len(trimmed) == 0 {
return 0, 0
}
ch := trimmed[0]
if ch != '`' && ch != '~' {
return 0, 0
}
count := 0
for count < len(trimmed) && trimmed[count] == ch {
count++
}
if count < 3 {
return 0, 0
}
// Per CommonMark: backtick fences cannot have backticks in the info string.
if ch == '`' && strings.ContainsRune(trimmed[count:], '`') {
return 0, 0
}
return ch, count
}
// isFenceClose checks whether trimmed is a closing fence matching fenceChar
// with at least minCount characters. A closing fence line contains only the
// fence characters and optional trailing spaces.
func isFenceClose(trimmed string, fenceChar byte, minCount int) bool {
if len(trimmed) == 0 || trimmed[0] != fenceChar {
return false
}
count := 0
for count < len(trimmed) && trimmed[count] == fenceChar {
count++
}
if count < minCount {
return false
}
// Closing fence must contain only fence chars (and optional trailing spaces).
return strings.TrimRight(trimmed[count:], " ") == ""
}
// --------------------------------------------------------------------------
// Inline code span handling
// --------------------------------------------------------------------------
// inlineCodeRanges returns byte ranges [start, end) of inline code spans
// in segment. Per CommonMark, a code span opens with N backticks and closes
// with exactly N backticks.
func inlineCodeRanges(s string) [][2]int {
var result [][2]int
i := 0
for i < len(s) {
if s[i] != '`' {
i++
continue
}
// Count opening backticks.
start := i
n := 0
for i < len(s) && s[i] == '`' {
n++
i++
}
// Scan for a closing run of exactly n backticks.
for j := i; j < len(s); {
if s[j] != '`' {
j++
continue
}
m := 0
for j < len(s) && s[j] == '`' {
m++
j++
}
if m == n {
result = append(result, [2]int{start, j})
i = j
break
}
}
// If no closing run was found, i is already past the opening
// backticks so the outer loop advances naturally.
}
return result
}
// replaceOutsideInline applies fn only to text outside inline code spans.
func replaceOutsideInline(segment string, fn func(string) string) string {
ranges := inlineCodeRanges(segment)
if len(ranges) == 0 {
return fn(segment)
}
var b strings.Builder
b.Grow(len(segment))
pos := 0
for _, r := range ranges {
if pos < r[0] {
b.WriteString(fn(segment[pos:r[0]]))
}
b.WriteString(segment[r[0]:r[1]])
pos = r[1]
}
if pos < len(segment) {
b.WriteString(fn(segment[pos:]))
}
return b.String()
}
// stripInlineCode removes inline code spans from s.
func stripInlineCode(s string) string {
ranges := inlineCodeRanges(s)
if len(ranges) == 0 {
return s
}
var b strings.Builder
b.Grow(len(s))
pos := 0
for _, r := range ranges {
b.WriteString(s[pos:r[0]])
pos = r[1]
}
b.WriteString(s[pos:])
return b.String()
}
+313
View File
@@ -0,0 +1,313 @@
package fences
import (
"testing"
)
func TestRanges(t *testing.T) {
tests := []struct {
name string
content string
want [][2]int
}{
{
name: "no fences",
content: "hello world\nno code here",
want: nil,
},
{
name: "single backtick fence",
content: "before\n```\ncode\n```\nafter",
want: [][2]int{{7, 20}},
},
{
name: "single tilde fence",
content: "before\n~~~\ncode\n~~~\nafter",
want: [][2]int{{7, 20}},
},
{
name: "fence with info string",
content: "before\n```go\ncode\n```\nafter",
want: [][2]int{{7, 22}},
},
{
name: "multiple fences",
content: "a\n```\nx\n```\nb\n~~~\ny\n~~~\nc",
want: [][2]int{{2, 12}, {14, 24}},
},
{
name: "unclosed fence",
content: "before\n```\ncode\nmore code",
want: [][2]int{{7, 25}},
},
{
name: "longer closing fence",
content: "before\n```\ncode\n`````\nafter",
want: [][2]int{{7, 22}},
},
{
name: "shorter closing fence ignored",
content: "before\n`````\ncode\n```\nmore\n`````\nafter",
want: [][2]int{{7, 33}},
},
{
name: "indented fence up to 3 spaces",
content: "before\n ```\ncode\n ```\nafter",
want: [][2]int{{7, 26}},
},
{
name: "4 space indent is not a fence",
content: "before\n ```\ncode\n ```\nafter",
want: nil,
},
{
name: "backtick in info string rejects open",
// The ```foo`bar line is not a valid opener (backtick in info).
// The standalone ``` becomes an opener with no close.
content: "before\n```foo`bar\ncode\n```\nafter",
want: [][2]int{{23, 32}},
},
{
name: "empty content",
content: "",
want: nil,
},
{
name: "fence only",
content: "```\ncode\n```",
want: [][2]int{{0, 12}},
},
{
name: "fence at end without trailing newline",
content: "```\ncode\n```",
want: [][2]int{{0, 12}},
},
{
name: "tilde fence does not close with backticks",
content: "~~~\ncode\n```\nmore\n~~~\nafter",
want: [][2]int{{0, 22}},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := Ranges(tt.content)
if len(got) != len(tt.want) {
t.Fatalf("Ranges() = %v, want %v", got, tt.want)
}
for i := range got {
if got[i] != tt.want[i] {
t.Errorf("Ranges()[%d] = %v, want %v", i, got[i], tt.want[i])
}
}
})
}
}
func TestReplaceOutside(t *testing.T) {
upper := func(s string) string {
b := []byte(s)
for i, c := range b {
if c >= 'a' && c <= 'z' {
b[i] = c - 32
}
}
return string(b)
}
tests := []struct {
name string
content string
want string
}{
{
name: "no fences",
content: "hello world",
want: "HELLO WORLD",
},
{
name: "text around fence",
content: "before\n```\ncode\n```\nafter",
want: "BEFORE\n```\ncode\n```\nAFTER",
},
{
name: "multiple fences",
content: "aaa\n```\nxxx\n```\nbbb\n~~~\nyyy\n~~~\nccc",
want: "AAA\n```\nxxx\n```\nBBB\n~~~\nyyy\n~~~\nCCC",
},
{
name: "unclosed fence preserves code",
content: "before\n```\ncode",
want: "BEFORE\n```\ncode",
},
{
name: "only fenced content",
content: "```\ncode\n```",
want: "```\ncode\n```",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := ReplaceOutside(tt.content, upper)
if got != tt.want {
t.Errorf("ReplaceOutside() =\n%s\nwant:\n%s", got, tt.want)
}
})
}
}
func TestStripFenced(t *testing.T) {
tests := []struct {
name string
content string
want string
}{
{
name: "no fences",
content: "hello $1 world",
want: "hello $1 world",
},
{
name: "strips fenced code",
content: "before $1\n```\n$2 inside\n```\nafter $3",
want: "before $1\nafter $3",
},
{
name: "multiple fences",
content: "a\n```\nx\n```\nb\n~~~\ny\n~~~\nc",
want: "a\nb\nc",
},
{
name: "unclosed fence",
content: "before\n```\n$1 inside",
want: "before\n",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := StripFenced(tt.content)
if got != tt.want {
t.Errorf("StripFenced() = %q, want %q", got, tt.want)
}
})
}
}
func TestInlineCodeRanges(t *testing.T) {
tests := []struct {
name string
s string
want [][2]int
}{
{"no backticks", "hello world", nil},
{"single backtick span", "use `$1` here", [][2]int{{4, 8}}},
{"double backtick span", "use ``$1`` here", [][2]int{{4, 10}}},
{"multiple spans", "`$1` and `$2`", [][2]int{{0, 4}, {9, 13}}},
{"unmatched backtick", "use `$1 here", nil},
{"mismatched backtick counts", "use ``$1` here", nil},
{"empty inline content", "use `` `` here", [][2]int{{4, 9}}},
{"backticks inside double", "use ``foo`bar`` here", [][2]int{{4, 15}}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := inlineCodeRanges(tt.s)
if len(got) != len(tt.want) {
t.Fatalf("inlineCodeRanges() = %v, want %v", got, tt.want)
}
for i := range got {
if got[i] != tt.want[i] {
t.Errorf("inlineCodeRanges()[%d] = %v, want %v", i, got[i], tt.want[i])
}
}
})
}
}
func TestReplaceOutside_InlineCode(t *testing.T) {
upper := func(s string) string {
b := []byte(s)
for i, c := range b {
if c >= 'a' && c <= 'z' {
b[i] = c - 32
}
}
return string(b)
}
tests := []struct {
name string
content string
want string
}{
{
name: "inline code preserved",
content: "use `code` here",
want: "USE `code` HERE",
},
{
name: "double backtick inline code",
content: "use ``co`de`` here",
want: "USE ``co`de`` HERE",
},
{
name: "mixed fenced and inline",
content: "before `x` mid\n```\nfenced\n```\nafter `y` end",
want: "BEFORE `x` MID\n```\nfenced\n```\nAFTER `y` END",
},
{
name: "only inline code",
content: "`code`",
want: "`code`",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := ReplaceOutside(tt.content, upper)
if got != tt.want {
t.Errorf("ReplaceOutside() =\n%s\nwant:\n%s", got, tt.want)
}
})
}
}
func TestStripCode(t *testing.T) {
tests := []struct {
name string
content string
want string
}{
{
name: "no code",
content: "hello $1 world",
want: "hello $1 world",
},
{
name: "strips inline code",
content: "use `$1` and `$2` for positional args",
want: "use and for positional args",
},
{
name: "strips fenced and inline",
content: "before `$1`\n```\n$2 inside\n```\nafter",
want: "before \nafter",
},
{
name: "real world prompt template",
content: "Use $@ for all args.\n`$1`, `$2` for positional.\n```bash\necho $1\n```\n",
want: "Use $@ for all args.\n, for positional.\n",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := StripCode(tt.content)
if got != tt.want {
t.Errorf("StripCode() = %q, want %q", got, tt.want)
}
})
}
}
+45 -21
View File
@@ -65,6 +65,10 @@ type AgentSetupOptions struct {
// AuthHandler handles OAuth authorization for remote MCP servers.
// When set, remote transports are configured with OAuth support.
AuthHandler tools.MCPAuthHandler
// TokenStoreFactory, if non-nil, creates a custom token store for each
// remote MCP server's OAuth tokens. When nil, the default file-based
// token store is used.
TokenStoreFactory tools.TokenStoreFactory
// OnMCPServerLoaded, if non-nil, is called when each MCP server finishes
// loading (successfully or with error). Called from the background goroutine.
OnMCPServerLoaded func(serverName string, toolCount int, err error)
@@ -82,36 +86,55 @@ type AgentSetupResult struct {
// BuildProviderConfig creates a *models.ProviderConfig from the current viper
// state. All entry points (root, script, SDK) converge through this function.
//
// Generation parameter pointers (Temperature, TopP, etc.) are only set when
// the user has explicitly configured them via CLI flag, environment variable,
// or global config file. This allows per-model defaults from modelSettings
// and customModels to fill in unset parameters downstream.
func BuildProviderConfig() (*models.ProviderConfig, string, error) {
systemPrompt, err := config.LoadSystemPrompt(viper.GetString("system-prompt"))
if err != nil {
return nil, "", fmt.Errorf("failed to load system prompt: %w", err)
}
temperature := float32(viper.GetFloat64("temperature"))
topP := float32(viper.GetFloat64("top-p"))
topK := int32(viper.GetInt("top-k"))
frequencyPenalty := float32(viper.GetFloat64("frequency-penalty"))
presencePenalty := float32(viper.GetFloat64("presence-penalty"))
numGPU := int32(viper.GetInt("num-gpu-layers"))
mainGPU := int32(viper.GetInt("main-gpu"))
cfg := &models.ProviderConfig{
ModelString: viper.GetString("model"),
SystemPrompt: systemPrompt,
ProviderAPIKey: viper.GetString("provider-api-key"),
ProviderURL: viper.GetString("provider-url"),
MaxTokens: viper.GetInt("max-tokens"),
Temperature: &temperature,
TopP: &topP,
TopK: &topK,
FrequencyPenalty: &frequencyPenalty,
PresencePenalty: &presencePenalty,
StopSequences: viper.GetStringSlice("stop-sequences"),
NumGPU: &numGPU,
MainGPU: &mainGPU,
TLSSkipVerify: viper.GetBool("tls-skip-verify"),
ThinkingLevel: models.ParseThinkingLevel(viper.GetString("thinking-level")),
ModelString: viper.GetString("model"),
SystemPrompt: systemPrompt,
ProviderAPIKey: viper.GetString("provider-api-key"),
ProviderURL: viper.GetString("provider-url"),
MaxTokens: viper.GetInt("max-tokens"),
StopSequences: viper.GetStringSlice("stop-sequences"),
NumGPU: &numGPU,
MainGPU: &mainGPU,
TLSSkipVerify: viper.GetBool("tls-skip-verify"),
ThinkingLevel: models.ParseThinkingLevel(viper.GetString("thinking-level")),
}
// Only set generation parameter pointers when the user has explicitly
// provided a value. This leaves nil pointers for unset params, allowing
// per-model defaults (modelSettings / customModels params) to apply.
if viper.IsSet("temperature") {
v := float32(viper.GetFloat64("temperature"))
cfg.Temperature = &v
}
if viper.IsSet("top-p") {
v := float32(viper.GetFloat64("top-p"))
cfg.TopP = &v
}
if viper.IsSet("top-k") {
v := int32(viper.GetInt("top-k"))
cfg.TopK = &v
}
if viper.IsSet("frequency-penalty") {
v := float32(viper.GetFloat64("frequency-penalty"))
cfg.FrequencyPenalty = &v
}
if viper.IsSet("presence-penalty") {
v := float32(viper.GetFloat64("presence-penalty"))
cfg.PresencePenalty = &v
}
return cfg, systemPrompt, nil
@@ -200,6 +223,7 @@ func SetupAgent(ctx context.Context, opts AgentSetupOptions) (*AgentSetupResult,
SpinnerFunc: opts.SpinnerFunc,
DebugLogger: debugLogger,
AuthHandler: opts.AuthHandler,
TokenStoreFactory: opts.TokenStoreFactory,
CoreTools: opts.CoreTools,
DisableCoreTools: opts.DisableCoreTools,
ToolWrapper: toolWrapper,
@@ -243,7 +267,7 @@ func loadExtensions() (*extensions.Runner, extensionCreationOpts, error) {
return extensions.WrapToolsWithExtensions(tools, runner)
}
extTools := extensions.ExtensionToolsAsFantasy(runner.RegisteredTools(), runner)
extTools := extensions.ExtensionToolsAsLLMTools(runner.RegisteredTools(), runner)
return runner, extensionCreationOpts{
toolWrapper: wrapper,
-13
View File
@@ -325,12 +325,6 @@ func UnmarshalParts(data []byte) ([]ContentPart, error) {
// mixed TextPart and ToolCallPart content. Tool-role messages produce
// ToolResultPart entries.
func (m *Message) ToLLMMessages() []fantasy.Message {
return m.ToFantasyMessages()
}
// Deprecated: Use ToLLMMessages instead.
// ToFantasyMessages converts a Message to one or more LLM message values.
func (m *Message) ToFantasyMessages() []fantasy.Message {
switch m.Role {
case RoleAssistant:
var parts []fantasy.MessagePart
@@ -431,13 +425,6 @@ func (m *Message) ToFantasyMessages() []fantasy.Message {
// FromLLMMessage converts an LLM message into our Message type,
// extracting all content parts into the appropriate block types.
func FromLLMMessage(msg fantasy.Message) Message {
return FromFantasyMessage(msg)
}
// Deprecated: Use FromLLMMessage instead.
// FromFantasyMessage converts an LLM message into our Message type,
// extracting all content parts into the appropriate block types.
func FromFantasyMessage(msg fantasy.Message) Message {
m := Message{
Role: MessageRole(msg.Role),
Parts: make([]ContentPart, 0),
+234 -11
View File
@@ -2,6 +2,8 @@ package models
import (
"log"
"os"
"strings"
"github.com/spf13/viper"
)
@@ -31,7 +33,7 @@ func loadCustomModelsFromConfig() map[string]ModelInfo {
// modelConfigToModelInfo converts a CustomModelConfig to a ModelInfo.
func modelConfigToModelInfo(modelID string, cfg CustomModelConfig) ModelInfo {
return ModelInfo{
info := ModelInfo{
ID: modelID,
Name: cfg.Name,
Attachment: cfg.Attachment,
@@ -48,21 +50,242 @@ func modelConfigToModelInfo(modelID string, cfg CustomModelConfig) ModelInfo {
Output: cfg.Limit.Output,
},
}
// Convert custom model generation params if any are set.
if p := convertGenerationParams(cfg.Params); p != nil {
info.Params = p
}
return info
}
// LoadModelSettingsFromConfig loads per-model generation parameter overrides
// from the config file. Keys are "provider/model" strings. Returns nil if
// no model settings are configured.
func LoadModelSettingsFromConfig() map[string]*GenerationParams {
if !viper.IsSet("modelSettings") {
return nil
}
var settings map[string]GenerationParamsConfig
if err := viper.UnmarshalKey("modelSettings", &settings); err != nil {
log.Printf("Warning: Failed to parse modelSettings: %v", err)
return nil
}
result := make(map[string]*GenerationParams, len(settings))
for modelKey, cfg := range settings {
if p := convertGenerationParams(cfg); p != nil {
result[modelKey] = p
}
}
return result
}
// convertGenerationParams converts a GenerationParamsConfig to a GenerationParams.
// Returns nil if no parameters are set.
func convertGenerationParams(cfg GenerationParamsConfig) *GenerationParams {
p := &GenerationParams{}
any := false
if cfg.MaxTokens != nil {
p.MaxTokens = cfg.MaxTokens
any = true
}
if cfg.Temperature != nil {
p.Temperature = cfg.Temperature
any = true
}
if cfg.TopP != nil {
p.TopP = cfg.TopP
any = true
}
if cfg.TopK != nil {
p.TopK = cfg.TopK
any = true
}
if cfg.FrequencyPenalty != nil {
p.FrequencyPenalty = cfg.FrequencyPenalty
any = true
}
if cfg.PresencePenalty != nil {
p.PresencePenalty = cfg.PresencePenalty
any = true
}
if len(cfg.StopSequences) > 0 {
p.StopSequences = cfg.StopSequences
any = true
}
if cfg.ThinkingLevel != "" {
p.ThinkingLevel = ParseThinkingLevel(cfg.ThinkingLevel)
any = true
}
if cfg.SystemPrompt != "" {
p.SystemPrompt = cfg.SystemPrompt
any = true
}
if !any {
return nil
}
return p
}
// ApplyModelSettings merges per-model generation parameter defaults from the
// registry into a ProviderConfig. Model-level params are only applied for
// fields where the user has not explicitly set a value (i.e., the
// corresponding viper key is not set via CLI flag or global config).
//
// The lookup order is:
// 1. modelSettings["provider/model"] from config (highest model-level priority)
// 2. ModelInfo.Params from custom model definitions
//
// Both are overridden by explicit CLI flags / global config values.
func ApplyModelSettings(config *ProviderConfig, modelInfo *ModelInfo) {
provider, modelName, err := ParseModelString(config.ModelString)
if err != nil {
return
}
// Collect model-level params: modelSettings override > custom model params.
// modelSettings takes priority because it's the more specific/intentional config.
var params *GenerationParams
// First check modelSettings from config.
if settings := LoadModelSettingsFromConfig(); settings != nil {
modelKey := provider + "/" + modelName
if p, ok := settings[modelKey]; ok {
params = p
}
}
// Fall back to ModelInfo.Params (from custom model definitions).
if params == nil && modelInfo != nil && modelInfo.Params != nil {
params = modelInfo.Params
}
if params == nil {
return
}
// Apply each parameter only when the user hasn't explicitly set it.
// We check viper.IsSet() which returns true only when the key was
// set via CLI flag, environment variable, or config file global section.
if params.MaxTokens != nil && !isExplicitlySet("max-tokens") {
config.MaxTokens = *params.MaxTokens
}
if params.Temperature != nil && !isExplicitlySet("temperature") {
config.Temperature = params.Temperature
}
if params.TopP != nil && !isExplicitlySet("top-p") {
config.TopP = params.TopP
}
if params.TopK != nil && !isExplicitlySet("top-k") {
config.TopK = params.TopK
}
if params.FrequencyPenalty != nil && !isExplicitlySet("frequency-penalty") {
config.FrequencyPenalty = params.FrequencyPenalty
}
if params.PresencePenalty != nil && !isExplicitlySet("presence-penalty") {
config.PresencePenalty = params.PresencePenalty
}
if len(params.StopSequences) > 0 && !isExplicitlySet("stop-sequences") {
config.StopSequences = params.StopSequences
}
if params.ThinkingLevel != "" && !isExplicitlySet("thinking-level") {
config.ThinkingLevel = params.ThinkingLevel
}
if params.SystemPrompt != "" && config.SystemPrompt == "" {
// Resolve file paths: if the value points to an existing file, read it.
// We check config.SystemPrompt == "" rather than isExplicitlySet because
// viper.BindPFlag causes IsSet to return true even for unset flags.
config.SystemPrompt = LoadSystemPromptValue(params.SystemPrompt)
}
}
// LoadSystemPromptValue resolves a system prompt value that may be either
// inline text or a file path. If the value is a path to an existing file,
// its contents are read and returned. Otherwise the string is returned as-is.
// This mirrors config.LoadSystemPrompt but lives in the models package to
// avoid circular dependencies.
func LoadSystemPromptValue(input string) string {
if input == "" {
return ""
}
if info, err := os.Stat(input); err == nil && !info.IsDir() {
content, err := os.ReadFile(input)
if err != nil {
log.Printf("Warning: failed to read system prompt file %q: %v", input, err)
return input
}
return strings.TrimSpace(string(content))
}
return input
}
// isExplicitlySet returns true when the user has explicitly set a config key
// via CLI flag, environment variable, or the global section of the config file.
// Model-level defaults should not override explicitly set values.
func isExplicitlySet(key string) bool {
// viper.IsSet returns true if the key has been set in any of the
// data stores (flag, env, config file, default). We need to check
// whether the value was set at the global config level (not just
// as a default). For generation params, the global config keys use
// hyphenated names (e.g. "max-tokens", "top-p").
//
// Since viper merges all sources, IsSet returns true even for config
// file values. This means global config file values (e.g.
// temperature: 0.7 at the top level) will correctly take precedence
// over model-level defaults, which is the desired behavior.
return viper.IsSet(key)
}
// GenerationParams holds per-model generation parameter defaults.
// These are stored on ModelInfo and applied during provider creation.
// Nil pointer fields mean "no model-level default" — the global config
// or CLI flag value (if any) will be used instead.
type GenerationParams struct {
MaxTokens *int
Temperature *float32
TopP *float32
TopK *int32
FrequencyPenalty *float32
PresencePenalty *float32
StopSequences []string
ThinkingLevel ThinkingLevel
SystemPrompt string // Per-model system prompt (inline text or file path)
}
// CustomModelConfig defines a custom model configuration loaded from the config file.
// This is a duplicate here to avoid circular dependencies with internal/config.
type CustomModelConfig struct {
Name string `json:"name" yaml:"name"`
BaseURL string `json:"baseUrl,omitempty" yaml:"baseUrl,omitempty"`
APIKey string `json:"apiKey,omitempty" yaml:"apiKey,omitempty"`
Family string `json:"family,omitempty" yaml:"family,omitempty"`
Attachment bool `json:"attachment,omitempty" yaml:"attachment,omitempty"`
Reasoning bool `json:"reasoning,omitempty" yaml:"reasoning,omitempty"`
Temperature bool `json:"temperature,omitempty" yaml:"temperature,omitempty"`
Knowledge string `json:"knowledge,omitempty" yaml:"knowledge,omitempty"`
Cost CostConfig `json:"cost" yaml:"cost"`
Limit LimitConfig `json:"limit" yaml:"limit"`
Name string `json:"name" yaml:"name"`
BaseURL string `json:"baseUrl,omitempty" yaml:"baseUrl,omitempty"`
APIKey string `json:"apiKey,omitempty" yaml:"apiKey,omitempty"`
Family string `json:"family,omitempty" yaml:"family,omitempty"`
Attachment bool `json:"attachment,omitempty" yaml:"attachment,omitempty"`
Reasoning bool `json:"reasoning,omitempty" yaml:"reasoning,omitempty"`
Temperature bool `json:"temperature,omitempty" yaml:"temperature,omitempty"`
Knowledge string `json:"knowledge,omitempty" yaml:"knowledge,omitempty"`
Cost CostConfig `json:"cost" yaml:"cost"`
Limit LimitConfig `json:"limit" yaml:"limit"`
Params GenerationParamsConfig `json:"params,omitzero" yaml:"params,omitempty"`
}
// GenerationParamsConfig is the JSON/YAML-serializable form of generation
// parameter defaults. Used in both customModels[].params and modelSettings[].
type GenerationParamsConfig struct {
MaxTokens *int `json:"maxTokens,omitempty" yaml:"maxTokens,omitempty"`
Temperature *float32 `json:"temperature,omitempty" yaml:"temperature,omitempty"`
TopP *float32 `json:"topP,omitempty" yaml:"topP,omitempty"`
TopK *int32 `json:"topK,omitempty" yaml:"topK,omitempty"`
FrequencyPenalty *float32 `json:"frequencyPenalty,omitempty" yaml:"frequencyPenalty,omitempty"`
PresencePenalty *float32 `json:"presencePenalty,omitempty" yaml:"presencePenalty,omitempty"`
StopSequences []string `json:"stopSequences,omitempty" yaml:"stopSequences,omitempty"`
ThinkingLevel string `json:"thinkingLevel,omitempty" yaml:"thinkingLevel,omitempty"`
SystemPrompt string `json:"systemPrompt,omitempty" yaml:"systemPrompt,omitempty"`
}
// CostConfig defines the pricing for a custom model.
+422
View File
@@ -0,0 +1,422 @@
package models
import (
"os"
"testing"
"github.com/spf13/viper"
)
func TestConvertGenerationParams(t *testing.T) {
t.Run("empty config returns nil", func(t *testing.T) {
cfg := GenerationParamsConfig{}
p := convertGenerationParams(cfg)
if p != nil {
t.Errorf("expected nil, got %+v", p)
}
})
t.Run("temperature only", func(t *testing.T) {
temp := float32(0.7)
cfg := GenerationParamsConfig{Temperature: &temp}
p := convertGenerationParams(cfg)
if p == nil {
t.Fatal("expected non-nil")
}
if p.Temperature == nil || *p.Temperature != 0.7 {
t.Errorf("expected temperature 0.7, got %v", p.Temperature)
}
if p.TopP != nil {
t.Errorf("expected nil TopP, got %v", p.TopP)
}
})
t.Run("all params set", func(t *testing.T) {
maxTokens := 8192
temp := float32(0.5)
topP := float32(0.9)
topK := int32(50)
freqPenalty := float32(0.1)
presPenalty := float32(0.2)
cfg := GenerationParamsConfig{
MaxTokens: &maxTokens,
Temperature: &temp,
TopP: &topP,
TopK: &topK,
FrequencyPenalty: &freqPenalty,
PresencePenalty: &presPenalty,
StopSequences: []string{"STOP"},
ThinkingLevel: "high",
}
p := convertGenerationParams(cfg)
if p == nil {
t.Fatal("expected non-nil")
}
if p.MaxTokens == nil || *p.MaxTokens != 8192 {
t.Errorf("expected maxTokens 8192, got %v", p.MaxTokens)
}
if p.Temperature == nil || *p.Temperature != 0.5 {
t.Errorf("expected temperature 0.5, got %v", p.Temperature)
}
if p.TopP == nil || *p.TopP != 0.9 {
t.Errorf("expected topP 0.9, got %v", p.TopP)
}
if p.TopK == nil || *p.TopK != 50 {
t.Errorf("expected topK 50, got %v", p.TopK)
}
if p.FrequencyPenalty == nil || *p.FrequencyPenalty != 0.1 {
t.Errorf("expected frequencyPenalty 0.1, got %v", p.FrequencyPenalty)
}
if p.PresencePenalty == nil || *p.PresencePenalty != 0.2 {
t.Errorf("expected presencePenalty 0.2, got %v", p.PresencePenalty)
}
if len(p.StopSequences) != 1 || p.StopSequences[0] != "STOP" {
t.Errorf("expected stop sequences [STOP], got %v", p.StopSequences)
}
if p.ThinkingLevel != ThinkingHigh {
t.Errorf("expected thinking level high, got %v", p.ThinkingLevel)
}
})
t.Run("thinking level parsing", func(t *testing.T) {
cfg := GenerationParamsConfig{ThinkingLevel: "medium"}
p := convertGenerationParams(cfg)
if p == nil {
t.Fatal("expected non-nil")
}
if p.ThinkingLevel != ThinkingMedium {
t.Errorf("expected thinking level medium, got %v", p.ThinkingLevel)
}
})
t.Run("system prompt only", func(t *testing.T) {
cfg := GenerationParamsConfig{SystemPrompt: "You are helpful."}
p := convertGenerationParams(cfg)
if p == nil {
t.Fatal("expected non-nil")
}
if p.SystemPrompt != "You are helpful." {
t.Errorf("expected system prompt, got %q", p.SystemPrompt)
}
})
}
func TestModelConfigToModelInfoWithParams(t *testing.T) {
temp := float32(0.8)
topP := float32(0.95)
cfg := CustomModelConfig{
Name: "Test Model",
BaseURL: "http://localhost:8080/v1",
Temperature: true,
Params: GenerationParamsConfig{
Temperature: &temp,
TopP: &topP,
},
}
info := modelConfigToModelInfo("test-model", cfg)
if info.Params == nil {
t.Fatal("expected non-nil Params")
}
if info.Params.Temperature == nil || *info.Params.Temperature != 0.8 {
t.Errorf("expected temperature 0.8, got %v", info.Params.Temperature)
}
if info.Params.TopP == nil || *info.Params.TopP != 0.95 {
t.Errorf("expected topP 0.95, got %v", info.Params.TopP)
}
}
func TestModelConfigToModelInfoWithoutParams(t *testing.T) {
cfg := CustomModelConfig{
Name: "Test Model",
BaseURL: "http://localhost:8080/v1",
}
info := modelConfigToModelInfo("test-model", cfg)
if info.Params != nil {
t.Errorf("expected nil Params, got %+v", info.Params)
}
}
func TestApplyModelSettings(t *testing.T) {
// Save and restore viper state.
originalViper := viper.AllSettings()
defer func() {
viper.Reset()
for k, v := range originalViper {
viper.Set(k, v)
}
}()
t.Run("applies model params when not explicitly set", func(t *testing.T) {
viper.Reset()
temp := float32(0.8)
topK := int32(50)
maxTokens := 4096
modelInfo := &ModelInfo{
ID: "test-model",
Params: &GenerationParams{
Temperature: &temp,
TopK: &topK,
MaxTokens: &maxTokens,
},
}
config := &ProviderConfig{
ModelString: "custom/test-model",
}
ApplyModelSettings(config, modelInfo)
if config.Temperature == nil || *config.Temperature != 0.8 {
t.Errorf("expected temperature 0.8, got %v", config.Temperature)
}
if config.TopK == nil || *config.TopK != 50 {
t.Errorf("expected topK 50, got %v", config.TopK)
}
if config.MaxTokens != 4096 {
t.Errorf("expected maxTokens 4096, got %d", config.MaxTokens)
}
})
t.Run("explicit viper values take precedence", func(t *testing.T) {
viper.Reset()
viper.Set("temperature", 0.3)
temp := float32(0.8)
modelInfo := &ModelInfo{
ID: "test-model",
Params: &GenerationParams{
Temperature: &temp,
},
}
explicitTemp := float32(0.3)
config := &ProviderConfig{
ModelString: "custom/test-model",
Temperature: &explicitTemp,
}
ApplyModelSettings(config, modelInfo)
// Temperature should NOT be overridden because it's explicitly set in viper
if config.Temperature == nil || *config.Temperature != 0.3 {
t.Errorf("expected temperature 0.3 (explicit), got %v", config.Temperature)
}
})
t.Run("nil model info is safe", func(t *testing.T) {
viper.Reset()
config := &ProviderConfig{
ModelString: "custom/test-model",
}
// Should not panic
ApplyModelSettings(config, nil)
if config.Temperature != nil {
t.Errorf("expected nil temperature, got %v", config.Temperature)
}
})
t.Run("model info without params is safe", func(t *testing.T) {
viper.Reset()
modelInfo := &ModelInfo{ID: "test-model"}
config := &ProviderConfig{
ModelString: "custom/test-model",
}
ApplyModelSettings(config, modelInfo)
if config.Temperature != nil {
t.Errorf("expected nil temperature, got %v", config.Temperature)
}
})
t.Run("modelSettings from viper takes priority over ModelInfo.Params", func(t *testing.T) {
viper.Reset()
// Set up modelSettings in viper (simulating config file)
viper.Set("modelSettings", map[string]any{
"custom/test-model": map[string]any{
"temperature": 0.5,
"topK": 30,
},
})
// ModelInfo has different params
temp := float32(0.8)
topK := int32(50)
modelInfo := &ModelInfo{
ID: "test-model",
Params: &GenerationParams{
Temperature: &temp,
TopK: &topK,
},
}
config := &ProviderConfig{
ModelString: "custom/test-model",
}
ApplyModelSettings(config, modelInfo)
// modelSettings should win over ModelInfo.Params
if config.Temperature == nil || *config.Temperature != 0.5 {
t.Errorf("expected temperature 0.5 (from modelSettings), got %v", config.Temperature)
}
if config.TopK == nil || *config.TopK != 30 {
t.Errorf("expected topK 30 (from modelSettings), got %v", config.TopK)
}
})
t.Run("stop sequences applied from model params", func(t *testing.T) {
viper.Reset()
modelInfo := &ModelInfo{
ID: "test-model",
Params: &GenerationParams{
StopSequences: []string{"STOP", "END"},
},
}
config := &ProviderConfig{
ModelString: "custom/test-model",
}
ApplyModelSettings(config, modelInfo)
if len(config.StopSequences) != 2 || config.StopSequences[0] != "STOP" {
t.Errorf("expected stop sequences [STOP END], got %v", config.StopSequences)
}
})
t.Run("thinking level applied from model params", func(t *testing.T) {
viper.Reset()
modelInfo := &ModelInfo{
ID: "test-model",
Params: &GenerationParams{
ThinkingLevel: ThinkingHigh,
},
}
config := &ProviderConfig{
ModelString: "custom/test-model",
}
ApplyModelSettings(config, modelInfo)
if config.ThinkingLevel != ThinkingHigh {
t.Errorf("expected thinking level high, got %v", config.ThinkingLevel)
}
})
t.Run("system prompt applied from model params", func(t *testing.T) {
viper.Reset()
modelInfo := &ModelInfo{
ID: "test-model",
Params: &GenerationParams{
SystemPrompt: "You are a coding assistant.",
},
}
config := &ProviderConfig{
ModelString: "custom/test-model",
}
ApplyModelSettings(config, modelInfo)
if config.SystemPrompt != "You are a coding assistant." {
t.Errorf("expected system prompt to be set, got %q", config.SystemPrompt)
}
})
t.Run("explicit system prompt takes precedence", func(t *testing.T) {
viper.Reset()
modelInfo := &ModelInfo{
ID: "test-model",
Params: &GenerationParams{
SystemPrompt: "Model-specific prompt",
},
}
config := &ProviderConfig{
ModelString: "custom/test-model",
SystemPrompt: "Global prompt",
}
ApplyModelSettings(config, modelInfo)
// Global system prompt should NOT be overridden because config
// already has a non-empty SystemPrompt.
if config.SystemPrompt != "Global prompt" {
t.Errorf("expected global prompt preserved, got %q", config.SystemPrompt)
}
})
t.Run("system prompt from file path", func(t *testing.T) {
viper.Reset()
// Create a temp file with a system prompt
tmpFile, err := os.CreateTemp("", "kit-test-prompt-*.txt")
if err != nil {
t.Fatal(err)
}
defer func() { _ = os.Remove(tmpFile.Name()) }()
if _, err := tmpFile.WriteString(" Prompt from file "); err != nil {
t.Fatal(err)
}
_ = tmpFile.Close()
modelInfo := &ModelInfo{
ID: "test-model",
Params: &GenerationParams{
SystemPrompt: tmpFile.Name(),
},
}
config := &ProviderConfig{
ModelString: "custom/test-model",
}
ApplyModelSettings(config, modelInfo)
if config.SystemPrompt != "Prompt from file" {
t.Errorf("expected trimmed file content, got %q", config.SystemPrompt)
}
})
t.Run("modelSettings system prompt overrides custom model params", func(t *testing.T) {
viper.Reset()
viper.Set("modelSettings", map[string]any{
"custom/test-model": map[string]any{
"systemPrompt": "From modelSettings",
},
})
modelInfo := &ModelInfo{
ID: "test-model",
Params: &GenerationParams{
SystemPrompt: "From custom model",
},
}
config := &ProviderConfig{
ModelString: "custom/test-model",
}
ApplyModelSettings(config, modelInfo)
if config.SystemPrompt != "From modelSettings" {
t.Errorf("expected modelSettings prompt, got %q", config.SystemPrompt)
}
})
}
File diff suppressed because one or more lines are too long
+125 -17
View File
@@ -25,7 +25,6 @@ import (
openaisdk "github.com/charmbracelet/openai-go"
"github.com/mark3labs/kit/internal/auth"
"github.com/mark3labs/kit/internal/ui/progress"
)
const (
@@ -86,6 +85,7 @@ type ThinkingLevel string
const (
ThinkingOff ThinkingLevel = "off"
ThinkingNone ThinkingLevel = "none"
ThinkingMinimal ThinkingLevel = "minimal"
ThinkingLow ThinkingLevel = "low"
ThinkingMedium ThinkingLevel = "medium"
@@ -94,12 +94,14 @@ const (
// ThinkingLevels returns the ordered list of available thinking levels for cycling.
func ThinkingLevels() []ThinkingLevel {
return []ThinkingLevel{ThinkingOff, ThinkingMinimal, ThinkingLow, ThinkingMedium, ThinkingHigh}
return []ThinkingLevel{ThinkingOff, ThinkingNone, ThinkingMinimal, ThinkingLow, ThinkingMedium, ThinkingHigh}
}
// thinkingBudgetTokens returns the token budget for a thinking level, or 0 for "off".
// thinkingBudgetTokens returns the token budget for a thinking level, or 0 for "off" or "none".
func thinkingBudgetTokens(level ThinkingLevel) int64 {
switch level {
case ThinkingNone:
return 1024
case ThinkingMinimal:
return 1024
case ThinkingLow:
@@ -118,6 +120,8 @@ func ThinkingLevelDescription(level ThinkingLevel) string {
switch level {
case ThinkingOff:
return "No reasoning"
case ThinkingNone:
return "Minimal reasoning (OpenAI 'none')"
case ThinkingMinimal:
return "Very brief reasoning (~1k tokens)"
case ThinkingLow:
@@ -134,7 +138,7 @@ func ThinkingLevelDescription(level ThinkingLevel) string {
// ParseThinkingLevel converts a string to a ThinkingLevel, defaulting to ThinkingOff.
func ParseThinkingLevel(s string) ThinkingLevel {
switch ThinkingLevel(s) {
case ThinkingMinimal, ThinkingLow, ThinkingMedium, ThinkingHigh:
case ThinkingNone, ThinkingMinimal, ThinkingLow, ThinkingMedium, ThinkingHigh:
return ThinkingLevel(s)
default:
return ThinkingOff
@@ -159,6 +163,12 @@ type ProviderConfig struct {
TLSSkipVerify bool
ThinkingLevel ThinkingLevel
DisableCaching bool // Opt-out: set to true to disable automatic prompt caching
// ProgressReaderFunc, when set, wraps an io.Reader with progress display
// for long operations like Ollama model pulls. The returned io.ReadCloser
// must be closed when done. When nil, the raw reader is consumed directly
// with no progress UI.
ProgressReaderFunc func(io.Reader) io.ReadCloser
}
// ProviderResult contains the result of provider creation.
@@ -241,6 +251,16 @@ func CreateProvider(ctx context.Context, config *ProviderConfig) (*ProviderResul
validateModelConfig(config, modelInfo)
}
// Apply per-model generation parameter defaults. Model-level params are
// only applied for fields where the user hasn't explicitly set a value
// via CLI flag or global config.
ApplyModelSettings(config, modelInfo)
// Auto-raise MaxTokens toward the model's known output ceiling when the
// user hasn't explicitly set --max-tokens and no per-model override
// applied. Runs after ApplyModelSettings so explicit modelSettings win.
rightSizeMaxTokens(config, modelInfo)
// Create the base provider
var result *ProviderResult
var createErr error
@@ -285,9 +305,18 @@ func CreateProvider(ctx context.Context, config *ProviderConfig) (*ProviderResul
// Only add cache options for providers that don't already have
// options set, to avoid type conflicts (e.g., Anthropic has
// different types for regular options vs cache control options).
for k, v := range cacheOpts {
if _, exists := result.ProviderOptions[k]; !exists {
result.ProviderOptions[k] = v
//
// For OpenAI Responses API models, we skip merging entirely because
// ResponsesProviderOptions and ProviderOptions are incompatible types.
skipMerge := false
if provider == "openai" && openai.IsResponsesModel(modelName) {
skipMerge = true
}
if !skipMerge {
for k, v := range cacheOpts {
if _, exists := result.ProviderOptions[k]; !exists {
result.ProviderOptions[k] = v
}
}
}
}
@@ -479,6 +508,37 @@ func validateModelConfig(config *ProviderConfig, modelInfo *ModelInfo) {
}
}
// defaultRightSizeCap bounds auto-raised MaxTokens so that we don't silently
// allocate enormous output budgets for models with very high ceilings (e.g.
// Devstral at 262144, Mistral at 128000). Users who genuinely want more can
// pass --max-tokens explicitly or set modelSettings[...].maxTokens in config.
const defaultRightSizeCap = 32768
// rightSizeMaxTokens raises config.MaxTokens toward the model's known output
// ceiling when:
// - the user has not explicitly set --max-tokens (or the KIT_MAX_TOKENS env
// var, or the top-level max-tokens key in config.yaml), AND
// - no per-model override already bumped MaxTokens (ApplyModelSettings runs
// before this function), AND
// - modelInfo.Limit.Output is known and larger than the current MaxTokens.
//
// The raised value is capped at defaultRightSizeCap to keep accidental
// allocations reasonable on very-large-output models. This prevents the
// common "ghost" where the agent's reply is silently truncated at the 8192
// default even though the selected model supports 64k or 262k output tokens.
func rightSizeMaxTokens(config *ProviderConfig, modelInfo *ModelInfo) {
if modelInfo == nil || modelInfo.Limit.Output <= 0 {
return
}
if isExplicitlySet("max-tokens") {
return
}
target := min(modelInfo.Limit.Output, defaultRightSizeCap)
if config.MaxTokens < target {
config.MaxTokens = target
}
}
// clearConflictingAnthropicSamplingParams ensures that temperature and top_p are
// not both sent to the Anthropic API, which rejects requests containing both.
// When both are set (typically from defaults), top_p is cleared so that
@@ -525,6 +585,8 @@ func buildOpenAIProviderOptions(config *ProviderConfig, modelName string) fantas
// Returns nil for ThinkingOff (use the model's default).
func thinkingLevelToReasoningEffort(level ThinkingLevel) *openai.ReasoningEffort {
switch level {
case ThinkingNone:
return new(openai.ReasoningEffortNone)
case ThinkingMinimal:
return new(openai.ReasoningEffortMinimal)
case ThinkingLow:
@@ -538,6 +600,56 @@ func thinkingLevelToReasoningEffort(level ThinkingLevel) *openai.ReasoningEffort
}
}
// IsValidThinkingLevelForModel checks if a thinking level is valid for the given
// model. Some OpenAI models like gpt-5.4 don't support "minimal" and require
// "none" instead.
func IsValidThinkingLevelForModel(level ThinkingLevel, modelName string) bool {
if level == ThinkingOff {
return true
}
// Check if this is an OpenAI model that doesn't support "minimal"
// gpt-5.4 and newer gpt-5.x models use "none" instead of "minimal"
if level == ThinkingMinimal {
if strings.Contains(modelName, "gpt-5.4") ||
strings.Contains(modelName, "gpt-5-pro") ||
strings.Contains(modelName, "gpt-5-chat") {
return false
}
}
// Check if this is an OpenAI model that doesn't support "none"
// Older gpt-5 models only support "minimal", not "none"
if level == ThinkingNone {
if strings.Contains(modelName, "gpt-5") &&
!strings.Contains(modelName, "gpt-5.4") &&
!strings.Contains(modelName, "gpt-5-pro") &&
!strings.Contains(modelName, "gpt-5-chat") {
// Older gpt-5 models might not support "none"
// They only added "none" support in newer versions
return false
}
}
// All other levels are generally valid for reasoning models
return true
}
// SuggestThinkingLevelFallback returns a recommended fallback level when the
// requested level is not valid for the model. Returns ThinkingOff if no
// suitable fallback exists.
func SuggestThinkingLevelFallback(level ThinkingLevel, modelName string) ThinkingLevel {
if level == ThinkingMinimal && !IsValidThinkingLevelForModel(level, modelName) {
// For models that don't support "minimal", suggest "none" (~same token budget)
return ThinkingNone
}
if level == ThinkingNone && !IsValidThinkingLevelForModel(level, modelName) {
// For models that don't support "none", suggest "minimal" (~same token budget)
return ThinkingMinimal
}
return ThinkingOff
}
// buildAnthropicProviderOptions returns fantasy.ProviderOptions configured for
// Anthropic models with extended thinking. When thinking is enabled, it sets
// SendReasoning to true and configures the thinking budget. For thinking-off
@@ -1123,7 +1235,7 @@ func loadOllamaModelWithFallback(ctx context.Context, baseURL, modelName string,
// Phase 1: Check if model exists locally
if err := checkOllamaModelExists(client, baseURL, modelName); err != nil {
// Phase 2: Pull model if not found
if err := pullOllamaModel(ctx, client, baseURL, modelName); err != nil {
if err := pullOllamaModel(ctx, client, baseURL, modelName, config.ProgressReaderFunc); err != nil {
return nil, fmt.Errorf("failed to pull model %s: %v", modelName, err)
}
}
@@ -1212,11 +1324,7 @@ func checkOllamaModelExists(client *http.Client, baseURL, modelName string) erro
return nil
}
func pullOllamaModel(ctx context.Context, client *http.Client, baseURL, modelName string) error {
return pullOllamaModelWithProgress(ctx, client, baseURL, modelName, true)
}
func pullOllamaModelWithProgress(ctx context.Context, client *http.Client, baseURL, modelName string, showProgress bool) error {
func pullOllamaModel(ctx context.Context, client *http.Client, baseURL, modelName string, progressFn func(io.Reader) io.ReadCloser) error {
reqBody := map[string]string{"name": modelName}
jsonBody, _ := json.Marshal(reqBody)
@@ -1240,10 +1348,10 @@ func pullOllamaModelWithProgress(ctx context.Context, client *http.Client, baseU
return fmt.Errorf("failed to pull model (status %d): %s", resp.StatusCode, string(body))
}
if showProgress {
progressReader := progress.NewProgressReader(resp.Body)
defer func() { _ = progressReader.Close() }()
_, err = io.ReadAll(progressReader)
if progressFn != nil {
pr := progressFn(resp.Body)
defer func() { _ = pr.Close() }()
_, err = io.ReadAll(pr)
} else {
_, err = io.ReadAll(resp.Body)
}
+17 -5
View File
@@ -26,6 +26,11 @@ type ModelInfo struct {
ProviderNPM string // Model-specific provider npm override (e.g. "@ai-sdk/anthropic")
BaseURL string // Per-model base URL override (custom models only)
APIKey string // Per-model API key override (custom models only)
// Params holds per-model generation parameter defaults. These are applied
// when the user hasn't explicitly set the corresponding CLI flag or global
// config value. Nil pointer fields mean "no model-level default".
Params *GenerationParams
}
// SupportsCaching returns true if this model family supports prompt caching.
@@ -236,6 +241,18 @@ func (r *ModelsRegistry) LookupModel(provider, modelID string) *ModelInfo {
return &modelInfo
}
// LookupModelForSettings is a convenience function that parses a
// "provider/model" string and looks up the ModelInfo in the global registry.
// Returns nil when the model string is invalid or the model is unknown.
// Used by Kit.SetModel to pre-apply per-model settings before CreateProvider.
func LookupModelForSettings(modelString string) *ModelInfo {
provider, modelName, err := ParseModelString(modelString)
if err != nil {
return nil
}
return GetGlobalRegistry().LookupModel(provider, modelName)
}
// getRequiredEnvVars returns the required environment variables for a provider.
func (r *ModelsRegistry) getRequiredEnvVars(provider string) ([]string, error) {
providerInfo, exists := r.providers[provider]
@@ -362,11 +379,6 @@ func (r *ModelsRegistry) GetLLMProviders() []string {
return providers
}
// Deprecated: Use GetLLMProviders instead.
func (r *ModelsRegistry) GetFantasyProviders() []string {
return r.GetLLMProviders()
}
// isProviderLLMSupported checks if a provider can be used with the LLM layer.
func isProviderLLMSupported(providerID string, info *ProviderInfo) bool {
// Ollama and custom are always supported (model names are user-defined).
+148
View File
@@ -0,0 +1,148 @@
package models
import (
"testing"
"github.com/spf13/pflag"
"github.com/spf13/viper"
)
// bindMaxTokensFlag wires a fresh pflag-backed "max-tokens" key into viper so
// isExplicitlySet behaves the same way it does in production. Returns a
// cleanup function that removes the binding so sibling tests see a clean
// state.
func bindMaxTokensFlag(t *testing.T, args []string) func() {
t.Helper()
fs := pflag.NewFlagSet("test", pflag.ContinueOnError)
fs.Int("max-tokens", 8192, "")
if err := viper.BindPFlag("max-tokens", fs.Lookup("max-tokens")); err != nil {
t.Fatalf("BindPFlag: %v", err)
}
if err := fs.Parse(args); err != nil {
t.Fatalf("fs.Parse: %v", err)
}
return func() {
viper.Reset()
}
}
func TestRightSizeMaxTokens_RaisesWhenBelowCeiling(t *testing.T) {
cleanup := bindMaxTokensFlag(t, nil) // no args → flag.Changed = false
defer cleanup()
config := &ProviderConfig{MaxTokens: 8192}
modelInfo := &ModelInfo{
ID: "claude-sonnet-4-5",
Limit: Limit{Context: 200000, Output: 64000},
}
rightSizeMaxTokens(config, modelInfo)
if config.MaxTokens != 32768 {
t.Errorf("expected MaxTokens raised to defaultRightSizeCap (32768), got %d", config.MaxTokens)
}
}
func TestRightSizeMaxTokens_CapsAtDefaultRightSizeCap(t *testing.T) {
cleanup := bindMaxTokensFlag(t, nil)
defer cleanup()
config := &ProviderConfig{MaxTokens: 8192}
// Mistral Devstral has 262144 output — we should still cap at 32768.
modelInfo := &ModelInfo{
ID: "devstral-medium-latest",
Limit: Limit{Context: 262144, Output: 262144},
}
rightSizeMaxTokens(config, modelInfo)
if config.MaxTokens != defaultRightSizeCap {
t.Errorf("expected MaxTokens capped at %d, got %d", defaultRightSizeCap, config.MaxTokens)
}
}
func TestRightSizeMaxTokens_UsesExactOutputWhenBelowCap(t *testing.T) {
cleanup := bindMaxTokensFlag(t, nil)
defer cleanup()
config := &ProviderConfig{MaxTokens: 4096}
// Model with output limit smaller than the cap.
modelInfo := &ModelInfo{
ID: "gpt-4",
Limit: Limit{Context: 8192, Output: 8192},
}
rightSizeMaxTokens(config, modelInfo)
if config.MaxTokens != 8192 {
t.Errorf("expected MaxTokens raised to model output ceiling (8192), got %d", config.MaxTokens)
}
}
func TestRightSizeMaxTokens_DoesNotLowerCurrentValue(t *testing.T) {
cleanup := bindMaxTokensFlag(t, nil)
defer cleanup()
// User (via per-model settings, applied earlier) already bumped MaxTokens
// above the cap — we must not clobber their choice.
config := &ProviderConfig{MaxTokens: 100000}
modelInfo := &ModelInfo{
ID: "devstral-medium-latest",
Limit: Limit{Context: 262144, Output: 262144},
}
rightSizeMaxTokens(config, modelInfo)
if config.MaxTokens != 100000 {
t.Errorf("expected MaxTokens preserved at 100000, got %d", config.MaxTokens)
}
}
func TestRightSizeMaxTokens_RespectsExplicitFlag(t *testing.T) {
// Simulate `--max-tokens 4096` on the command line.
cleanup := bindMaxTokensFlag(t, []string{"--max-tokens", "4096"})
defer cleanup()
config := &ProviderConfig{MaxTokens: 4096}
modelInfo := &ModelInfo{
ID: "claude-sonnet-4-5",
Limit: Limit{Context: 200000, Output: 64000},
}
rightSizeMaxTokens(config, modelInfo)
if config.MaxTokens != 4096 {
t.Errorf("expected explicit --max-tokens to be preserved (4096), got %d", config.MaxTokens)
}
}
func TestRightSizeMaxTokens_NilModelInfo(t *testing.T) {
cleanup := bindMaxTokensFlag(t, nil)
defer cleanup()
config := &ProviderConfig{MaxTokens: 8192}
// Custom model / Ollama / unknown provider → no model info.
rightSizeMaxTokens(config, nil)
if config.MaxTokens != 8192 {
t.Errorf("expected MaxTokens unchanged with nil modelInfo, got %d", config.MaxTokens)
}
}
func TestRightSizeMaxTokens_ZeroOutputLimit(t *testing.T) {
cleanup := bindMaxTokensFlag(t, nil)
defer cleanup()
config := &ProviderConfig{MaxTokens: 8192}
// Model present in catalog but with no known output limit.
modelInfo := &ModelInfo{
ID: "unknown-model",
Limit: Limit{Context: 0, Output: 0},
}
rightSizeMaxTokens(config, modelInfo)
if config.MaxTokens != 8192 {
t.Errorf("expected MaxTokens unchanged with zero output limit, got %d", config.MaxTokens)
}
}
+59 -6
View File
@@ -7,10 +7,12 @@ import (
"regexp"
"strconv"
"strings"
"github.com/mark3labs/kit/internal/fences"
)
// PromptTemplate is a named prompt template with shell-style argument placeholders.
// It supports Pi-style $1, $2, $@, $ARGUMENTS, ${@:N}, ${@:N:L} syntax.
// It supports Pi-style $1, $2, $@, $+, $ARGUMENTS, ${@:N}, ${@:N:L} syntax.
type PromptTemplate struct {
// Name is the human-readable identifier for this template.
Name string
@@ -120,19 +122,28 @@ func ParseCommandArgs(input string) []string {
// argPlaceholder matches shell-style argument placeholders:
// - $1, $2, etc. - positional arguments
// - $@ - all arguments
// - $@ - all arguments (zero or more)
// - $+ - all arguments (one or more required)
// - $ARGUMENTS - all arguments (alias for $@)
// - ${@:N} - arguments from N onwards
// - ${@:N:L} - L arguments starting from N
var argPlaceholder = regexp.MustCompile(`\$\{(\d+)\}|\$\{(\d+):(\d+)\}|\$\{ARGUMENTS\}|\$\{@(:\d+)?(:\d+)?\}|\$(\d+)|\$@|\$ARGUMENTS`)
var argPlaceholder = regexp.MustCompile(`\$\{(\d+)\}|\$\{(\d+):(\d+)\}|\$\{ARGUMENTS\}|\$\{@(:\d+)?(:\d+)?\}|\$(\d+)|\$@|\$\+|\$ARGUMENTS`)
// SubstituteArgs replaces argument placeholders in content with values from args.
// Supported placeholders:
// - $N, ${N} - the Nth argument (1-indexed)
// - $@, $ARGUMENTS, ${ARGUMENTS} - all arguments joined with spaces
// - $@, $+, $ARGUMENTS, ${ARGUMENTS} - all arguments joined with spaces
// - ${@:N} - arguments from index N onwards (0-indexed)
// - ${@:N:L} - L arguments starting from index N (0-indexed)
func SubstituteArgs(content string, args []string) string {
return fences.ReplaceOutside(content, func(segment string) string {
return substituteArgsInSegment(segment, args)
})
}
// substituteArgsInSegment performs argument substitution on a single text
// segment that is known to be outside fenced code blocks.
func substituteArgsInSegment(content string, args []string) string {
return argPlaceholder.ReplaceAllStringFunc(content, func(match string) string {
// Check for ${N} or ${N:M} format
if strings.HasPrefix(match, "${") && strings.Contains(match, "}") {
@@ -191,8 +202,8 @@ func SubstituteArgs(content string, args []string) string {
if strings.HasPrefix(match, "$") && !strings.HasPrefix(match, "${") {
suffix := match[1:]
// $@ or $ARGUMENTS
if suffix == "@" || suffix == "ARGUMENTS" {
// $@, $+, or $ARGUMENTS
if suffix == "@" || suffix == "+" || suffix == "ARGUMENTS" {
return strings.Join(args, " ")
}
@@ -266,6 +277,48 @@ func joinArgsRange(args []string, start, length int) string {
return strings.Join(args[start:end], " ")
}
// HasArgPlaceholders reports whether the template content contains any
// argument placeholders ($1, $@, $ARGUMENTS, ${@:...}, etc.).
// Placeholders inside fenced code blocks and inline code spans are ignored.
func (t *PromptTemplate) HasArgPlaceholders() bool {
return argPlaceholder.MatchString(fences.StripCode(t.Content))
}
// RequiredArgs returns the number of positional arguments the template
// expects. This is determined by the highest $N or ${N} placeholder found
// in the content (1-indexed, so $2 means 2 args required). The $+
// placeholder (required variadic) ensures at least 1. Optional wildcards
// ($@, $ARGUMENTS) do not contribute to the count.
func (t *PromptTemplate) RequiredArgs() int {
content := fences.StripCode(t.Content)
maxN := 0
hasRequiredVariadic := strings.Contains(content, "$+")
for _, match := range argPlaceholder.FindAllStringSubmatch(content, -1) {
// Group 1: ${N} format — the N value.
if match[1] != "" {
if n, err := strconv.Atoi(match[1]); err == nil && n > maxN {
maxN = n
}
}
// Group 2: ${N:M} format — the N value (start index).
if match[2] != "" {
if n, err := strconv.Atoi(match[2]); err == nil && n > maxN {
maxN = n
}
}
// Group 6: $N format (no braces) — the N value.
if match[6] != "" {
if n, err := strconv.Atoi(match[6]); err == nil && n > maxN {
maxN = n
}
}
}
if hasRequiredVariadic && maxN < 1 {
maxN = 1
}
return maxN
}
// Expand substitutes arguments into the template content and returns the result.
// It first parses args from the input string, then substitutes them into the template.
func (t *PromptTemplate) Expand(argsInput string) string {
+117
View File
@@ -129,6 +129,48 @@ func TestSubstituteArgs(t *testing.T) {
args: []string{},
expected: "Args: ",
},
{
name: "$1 inside code block preserved",
content: "Use $1 here\n```bash\necho $1\n```\ndone",
args: []string{"foo"},
expected: "Use foo here\n```bash\necho $1\n```\ndone",
},
{
name: "$@ inside code block preserved",
content: "Run $@\n```\necho $@\n```\n",
args: []string{"a", "b"},
expected: "Run a b\n```\necho $@\n```\n",
},
{
name: "all placeholders inside code block",
content: "Prompt\n```\n$1 $2 $@\n```\n",
args: []string{"x"},
expected: "Prompt\n```\n$1 $2 $@\n```\n",
},
{
name: "$1 inside inline code preserved",
content: "Use `$1` here and $1 outside",
args: []string{"foo"},
expected: "Use `$1` here and foo outside",
},
{
name: "$+ required variadic",
content: "Args: $+",
args: []string{"a", "b", "c"},
expected: "Args: a b c",
},
{
name: "$+ with empty args",
content: "Args: $+",
args: []string{},
expected: "Args: ",
},
{
name: "all placeholders in inline code",
content: "Use `$1` and `$@` for args",
args: []string{"x"},
expected: "Use `$1` and `$@` for args",
},
}
for _, tt := range tests {
@@ -213,3 +255,78 @@ func TestPromptTemplateExpand(t *testing.T) {
})
}
}
func TestHasArgPlaceholders(t *testing.T) {
tests := []struct {
name string
content string
want bool
}{
{"no placeholders", "Just a plain prompt with no args", false},
{"$1 placeholder", "Create a $1 component", true},
{"$@ placeholder", "Run with args: $@", true},
{"$ARGUMENTS placeholder", "Features: $ARGUMENTS", true},
{"${1} placeholder", "Name: ${1}", true},
{"${ARGUMENTS} placeholder", "All: ${ARGUMENTS}", true},
{"${@:1} placeholder", "Rest: ${@:1}", true},
{"${@:1:2} placeholder", "Slice: ${@:1:2}", true},
{"dollar in text", "Cost is one hundred dollars", false},
{"empty content", "", false},
{"$1 inside code block only", "Prompt\n```\necho $1\n```\n", false},
{"$1 outside and inside code block", "Use $1 here\n```\necho $1\n```\n", true},
{"$@ inside code block only", "Prompt\n```bash\necho $@\n```\n", false},
{"$+ placeholder", "Run with args: $+", true},
{"$+ inside inline code only", "Use `$+` for required args", false},
{"$1 inside inline code only", "Use `$1` for positional args", false},
{"$1 outside and in inline code", "Create $1 (see `$1` syntax)", true},
{"$@ outside $1 in inline code", "Run $@ with `$1` syntax", true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tpl := &PromptTemplate{Content: tt.content}
if got := tpl.HasArgPlaceholders(); got != tt.want {
t.Errorf("HasArgPlaceholders() = %v, want %v", got, tt.want)
}
})
}
}
func TestRequiredArgs(t *testing.T) {
tests := []struct {
name string
content string
want int
}{
{"no placeholders", "Just a plain prompt", 0},
{"$1 only", "Create a $1 component", 1},
{"$1 and $2", "Create $1 with $2", 2},
{"$3 skipping $2", "Use $1 and $3", 3},
{"${1} braced", "Name: ${1}", 1},
{"${2} braced", "Name: ${1} Desc: ${2}", 2},
{"$@ only", "Run with: $@", 0},
{"$ARGUMENTS only", "Features: $ARGUMENTS", 0},
{"${ARGUMENTS} only", "All: ${ARGUMENTS}", 0},
{"$1 and $@", "Create $1 with extras: $@", 1},
{"${@:1} slice only", "Rest: ${@:1}", 0},
{"${@:1:2} slice only", "Slice: ${@:1:2}", 0},
{"mixed $1 $2 and $@", "Create $1 named $2: $@", 2},
{"empty content", "", 0},
{"$2 inside code block only", "Prompt\n```\n$1 $2\n```\n", 0},
{"$1 outside $2 inside code block", "Use $1\n```\n$2 inside\n```\n", 1},
{"$+ only", "Run with: $+", 1},
{"$+ and $2", "Create $2 with: $+", 2},
{"$+ inside inline code only", "Use `$+` for required args", 0},
{"$1 and $2 in inline code only", "Use `$1` and `$2` for args", 0},
{"$1 outside $2 in inline code", "Create $1 (see `$2`)", 1},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
tpl := &PromptTemplate{Content: tt.content}
if got := tpl.RequiredArgs(); got != tt.want {
t.Errorf("RequiredArgs() = %d, want %d", got, tt.want)
}
})
}
}
+66
View File
@@ -0,0 +1,66 @@
package session
import (
"testing"
"github.com/mark3labs/kit/internal/message"
)
// TestCompactionParentCycleRegression tests that after multiple compactions,
// newly appended messages always have a valid parent chain and BuildContext
// returns the correct messages.
func TestCompactionParentCycleRegression(t *testing.T) {
tm := InMemoryTreeSession("/test")
// Simulate a long conversation with multiple compactions.
msg1, _ := tm.AppendMessage(message.Message{Role: message.RoleUser, Parts: []message.ContentPart{message.TextContent{Text: "msg1"}}})
msg2, _ := tm.AppendMessage(message.Message{Role: message.RoleAssistant, Parts: []message.ContentPart{message.TextContent{Text: "msg2"}}})
// First compaction
comp1, _ := tm.AppendCompaction("Summary 1", msg1, 1000, 500, 1, []string{}, []string{})
msg3, _ := tm.AppendMessage(message.Message{Role: message.RoleUser, Parts: []message.ContentPart{message.TextContent{Text: "msg3"}}})
msg4, _ := tm.AppendMessage(message.Message{Role: message.RoleAssistant, Parts: []message.ContentPart{message.TextContent{Text: "msg4"}}})
// Second compaction
comp2, _ := tm.AppendCompaction("Summary 2", msg3, 1000, 500, 1, []string{}, []string{})
msg5, _ := tm.AppendMessage(message.Message{Role: message.RoleUser, Parts: []message.ContentPart{message.TextContent{Text: "msg5"}}})
msg6, _ := tm.AppendMessage(message.Message{Role: message.RoleAssistant, Parts: []message.ContentPart{message.TextContent{Text: "msg6"}}})
// Verify parent chain integrity
for _, id := range []string{msg1, msg2, comp1, msg3, msg4, comp2, msg5, msg6} {
entry := tm.GetEntry(id)
if entry == nil {
t.Fatalf("entry %s not found in index", id)
}
}
// Walk parent chain from msg6 — must reach root without cycles
visited := make(map[string]bool)
current := msg6
for current != "" {
if visited[current] {
t.Fatalf("cycle detected at entry %s", current)
}
visited[current] = true
entry := tm.GetEntry(current)
if entry == nil {
t.Fatalf("entry %s missing from index during parent walk", current)
}
parent := ""
switch e := entry.(type) {
case *MessageEntry:
parent = e.ParentID
case *CompactionEntry:
parent = e.ParentID
}
current = parent
}
// BuildContext should return: Summary2 + msg6 + msg5 + msg3 + msg4 = 5 messages
msgs, _, _ := tm.BuildContext()
if len(msgs) != 5 {
t.Fatalf("expected 5 messages, got %d: %+v", len(msgs), msgs)
}
}
+109
View File
@@ -0,0 +1,109 @@
package session
import (
"testing"
"github.com/mark3labs/kit/internal/message"
)
// TestDetectCycleWithCorruptedParentChain tests that cycle detection works
// when a corrupted session has circular parent references.
func TestDetectCycleWithCorruptedParentChain(t *testing.T) {
tm := InMemoryTreeSession("/test")
// Create normal chain: msg1 -> msg2 -> msg3
id1, _ := tm.AppendMessage(message.Message{Role: message.RoleUser, Parts: []message.ContentPart{message.TextContent{Text: "msg1"}}})
_, _ = tm.AppendMessage(message.Message{Role: message.RoleAssistant, Parts: []message.ContentPart{message.TextContent{Text: "msg2"}}})
id3, _ := tm.AppendMessage(message.Message{Role: message.RoleUser, Parts: []message.ContentPart{message.TextContent{Text: "msg3"}}})
// Simulate corruption: manually set msg1's parent to msg3, creating cycle
// This simulates the condition seen in the user's session
for _, entry := range tm.entries {
if e, ok := entry.(*MessageEntry); ok && e.ID == id1 {
e.ParentID = id3 // Create cycle: msg1 -> msg3 -> ... -> msg1
break
}
}
// DetectCycle should find the cycle
// The cycle is: id1 -> id3 -> id2 -> id1
// So detecting from id3 should find id1 as the repeat
cycle, entry := tm.DetectCycle(id3)
if !cycle {
t.Fatal("expected to detect cycle, but none found")
}
// The cycle entry could be id1 or id3 depending on where we start
if entry != id1 && entry != id3 {
t.Fatalf("expected cycle at %s or %s, got %s", id1, id3, entry)
}
// BuildContext should still work (it has its own cycle detection)
// but will truncate at the cycle point
msgs, _, _ := tm.BuildContext()
if len(msgs) == 0 {
t.Fatal("BuildContext returned no messages")
}
}
// TestAppendMessageRejectsInvalidParent tests that AppendMessage rejects
// appending when the current leaf has a broken parent chain.
func TestAppendMessageRejectsInvalidParent(t *testing.T) {
tm := InMemoryTreeSession("/test")
// Create normal message
id1, err := tm.AppendMessage(message.Message{Role: message.RoleUser, Parts: []message.ContentPart{message.TextContent{Text: "msg1"}}})
if err != nil {
t.Fatalf("failed to append msg1: %v", err)
}
// Simulate corruption: set leafID to a non-existent ID
tm.leafID = "non-existent-id"
// Next append should fail validation
_, err = tm.AppendMessage(message.Message{Role: message.RoleAssistant, Parts: []message.ContentPart{message.TextContent{Text: "msg2"}}})
if err == nil {
t.Fatal("expected error when appending with invalid leafID, got nil")
}
// Restore valid leafID
tm.leafID = id1
// Append should succeed now
_, err = tm.AppendMessage(message.Message{Role: message.RoleAssistant, Parts: []message.ContentPart{message.TextContent{Text: "msg3"}}})
if err != nil {
t.Fatalf("failed to append msg3 after restoring leafID: %v", err)
}
}
// TestBuildContextHandlesCycleGracefully tests that BuildContext handles
// cycles gracefully by truncating the branch.
func TestBuildContextHandlesCycleGracefully(t *testing.T) {
tm := InMemoryTreeSession("/test")
// Create messages
id1, _ := tm.AppendMessage(message.Message{Role: message.RoleUser, Parts: []message.ContentPart{message.TextContent{Text: "msg1"}}})
_, _ = tm.AppendMessage(message.Message{Role: message.RoleAssistant, Parts: []message.ContentPart{message.TextContent{Text: "msg2"}}})
id3, _ := tm.AppendMessage(message.Message{Role: message.RoleUser, Parts: []message.ContentPart{message.TextContent{Text: "msg3"}}})
// Verify normal case works
msgs, _, _ := tm.BuildContext()
if len(msgs) != 3 {
t.Fatalf("expected 3 messages, got %d", len(msgs))
}
// Simulate cycle: set msg1's parent to msg3
for _, entry := range tm.entries {
if e, ok := entry.(*MessageEntry); ok && e.ID == id1 {
e.ParentID = id3
break
}
}
// BuildContext should handle cycle gracefully (getBranchLocked has cycle detection)
msgs, _, _ = tm.BuildContext()
// Should only include messages from the cycle: msg3, msg2, msg1
// (msg3 is leaf, walks to msg2 -> msg1 -> msg3 (cycle detected, stops))
if len(msgs) != 3 {
t.Fatalf("expected 3 messages in cycle case, got %d: %+v", len(msgs), msgs)
}
}
+37 -1
View File
@@ -365,6 +365,9 @@ func OpenTreeSession(path string) (*TreeManager, error) {
tm.leafID = tm.EntryID(tm.entries[len(tm.entries)-1])
}
// Validate tree integrity and log diagnostics
tm.LogTreeDiagnostics()
// Open file for appending.
f, err := os.OpenFile(path, os.O_WRONLY|os.O_APPEND, 0644)
if err != nil {
@@ -410,6 +413,12 @@ func (tm *TreeManager) AppendMessage(msg message.Message) (string, error) {
tm.mu.Lock()
defer tm.mu.Unlock()
// Validate parent chain before appending to detect/prevent cycles
// that could be caused by external file corruption or race conditions.
if err := tm.validateParentChainLocked(tm.leafID, ""); err != nil {
return "", fmt.Errorf("parent chain validation failed: %w", err)
}
entry, err := NewMessageEntry(tm.leafID, msg)
if err != nil {
return "", err
@@ -518,6 +527,13 @@ func (tm *TreeManager) AppendCompaction(summary, firstKeptEntryID string, tokens
tm.mu.Lock()
defer tm.mu.Unlock()
// Validate that firstKeptEntryID exists if provided
if firstKeptEntryID != "" {
if _, ok := tm.index[firstKeptEntryID]; !ok {
return "", fmt.Errorf("first kept entry %q does not exist", firstKeptEntryID)
}
}
// The compaction entry has no parent, making it a new "root" for the
// post-compaction branch. This ensures old compacted messages are not
// traversed when walking from the current leaf.
@@ -1213,12 +1229,32 @@ func (tm *TreeManager) getBranchLocked(fromID string) []any {
}
// buildTreeNode recursively builds a TreeNode from an entry ID.
// It includes a depth limit to prevent infinite recursion in case of
// corrupted parent-child relationships.
func (tm *TreeManager) buildTreeNode(id string) *TreeNode {
return tm.buildTreeNodeDepth(id, 0, make(map[string]bool))
}
// buildTreeNodeDepth is the internal implementation with depth tracking.
func (tm *TreeManager) buildTreeNodeDepth(id string, depth int, visited map[string]bool) *TreeNode {
const maxDepth = 1000
if depth > maxDepth {
// Cycle or extremely deep tree detected, stop recursing
return nil
}
if visited[id] {
// Cycle detected, stop recursing
return nil
}
entry, ok := tm.index[id]
if !ok {
return nil
}
visited[id] = true
defer delete(visited, id)
node := &TreeNode{
Entry: entry,
ID: id,
@@ -1226,7 +1262,7 @@ func (tm *TreeManager) buildTreeNode(id string) *TreeNode {
}
for _, childID := range tm.childIndex[id] {
child := tm.buildTreeNode(childID)
child := tm.buildTreeNodeDepth(childID, depth+1, visited)
if child != nil {
node.Children = append(node.Children, child)
}
+143
View File
@@ -0,0 +1,143 @@
package session
import (
"fmt"
"log"
)
// ValidateParentChain checks that the parent ID points to an existing entry
// and that appending this entry would not create a cycle. This should be called
// before appending any entry to the tree.
// Returns an error if the parent is invalid or would create a cycle.
func (tm *TreeManager) ValidateParentChain(parentID string, newEntryID string) error {
if parentID == "" {
// Empty parent is valid (root entry)
return nil
}
// Check that parent exists
if _, ok := tm.index[parentID]; !ok {
return fmt.Errorf("parent entry %q does not exist in index", parentID)
}
// Check that we're not creating a cycle by walking up the parent chain
// from parentID and ensuring we don't hit newEntryID (or any node that
// has newEntryID as an ancestor, but since newEntryID is new, just check
// that parentID isn't newEntryID, which it can't be since we check existence)
visited := make(map[string]bool)
current := parentID
for current != "" {
if visited[current] {
return fmt.Errorf("existing cycle detected at entry %q", current)
}
visited[current] = true
// Safety check: if somehow we reach the new entry ID, that's a cycle
if current == newEntryID {
return fmt.Errorf("would create cycle: entry %q cannot be its own ancestor", newEntryID)
}
entry, ok := tm.index[current]
if !ok {
return fmt.Errorf("broken parent chain: entry %q not found", current)
}
current = tm.entryParentID(entry)
}
return nil
}
// DetectCycle walks the parent chain from the given entry ID and returns true
// if a cycle is detected. This is used for diagnostics.
func (tm *TreeManager) DetectCycle(fromID string) (cycleDetected bool, cycleEntry string) {
visited := make(map[string]bool)
current := fromID
for current != "" {
if visited[current] {
return true, current
}
visited[current] = true
entry, ok := tm.index[current]
if !ok {
return false, ""
}
current = tm.entryParentID(entry)
}
return false, ""
}
// LogTreeDiagnostics logs information about the tree structure for debugging.
// Call this after OpenTreeSession or when anomalies are detected.
func (tm *TreeManager) LogTreeDiagnostics() {
tm.mu.RLock()
defer tm.mu.RUnlock()
log.Printf("[TreeManager] Entry count: %d, Leaf ID: %s", len(tm.entries), tm.leafID)
// Check for cycles from leaf
if tm.leafID != "" {
if cycle, entry := tm.detectCycleLocked(tm.leafID); cycle {
log.Printf("[TreeManager] WARNING: Cycle detected in tree at entry %s", entry)
}
}
// Count entries by type
counts := make(map[EntryType]int)
for _, entry := range tm.entries {
var et EntryType
switch e := entry.(type) {
case *MessageEntry:
et = e.Type
case *ModelChangeEntry:
et = e.Type
case *BranchSummaryEntry:
et = e.Type
case *LabelEntry:
et = e.Type
case *SessionInfoEntry:
et = e.Type
case *ExtensionDataEntry:
et = e.Type
case *CompactionEntry:
et = e.Type
default:
et = "unknown"
}
counts[et]++
}
log.Printf("[TreeManager] Entry types: %+v", counts)
}
// detectCycleLocked is the internal version of DetectCycle (must hold read lock)
func (tm *TreeManager) detectCycleLocked(fromID string) (bool, string) {
visited := make(map[string]bool)
current := fromID
for current != "" {
if visited[current] {
return true, current
}
visited[current] = true
entry, ok := tm.index[current]
if !ok {
return false, ""
}
current = tm.entryParentID(entry)
}
return false, ""
}
// validateParentChainLocked is the internal version used by append methods.
// Must be called with the write lock held.
func (tm *TreeManager) validateParentChainLocked(parentID string, newEntryID string) error {
if parentID == "" {
return nil
}
if _, ok := tm.index[parentID]; !ok {
return fmt.Errorf("parent entry %q does not exist", parentID)
}
// Check for existing cycles in the parent chain
if cycle, entry := tm.detectCycleLocked(parentID); cycle {
return fmt.Errorf("existing cycle detected at entry %q in parent chain", entry)
}
return nil
}
+104 -36
View File
@@ -8,11 +8,11 @@ import (
"sync"
"time"
"charm.land/fantasy"
"github.com/mark3labs/kit/internal/config"
"github.com/mark3labs/mcp-go/client"
"github.com/mark3labs/mcp-go/client/transport"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
// ConnectionPoolConfig defines configuration parameters for the MCP connection pool.
@@ -60,35 +60,34 @@ type MCPConnection struct {
// creation, health monitoring, and cleanup. The pool runs background health checks
// to proactively identify and remove unhealthy connections.
type MCPConnectionPool struct {
connections map[string]*MCPConnection
config *ConnectionPoolConfig
mu sync.RWMutex
model fantasy.LanguageModel
ctx context.Context
cancel context.CancelFunc
debug bool
debugLogger DebugLogger
oauthFlow *OAuthFlowRunner
connections map[string]*MCPConnection
config *ConnectionPoolConfig
mu sync.RWMutex
ctx context.Context
cancel context.CancelFunc
debug bool
debugLogger DebugLogger
oauthFlow *OAuthFlowRunner
tokenStoreFactory TokenStoreFactory // custom factory for per-server token stores (nil = default FileTokenStore)
}
// NewMCPConnectionPool creates a new MCP connection pool with the specified configuration.
// If config is nil, default configuration values will be used. The pool starts a background
// goroutine for periodic health checks that runs until Close is called.
// The model parameter is used for MCP servers that require sampling support.
// Thread-safe for concurrent use immediately after creation.
func NewMCPConnectionPool(config *ConnectionPoolConfig, model fantasy.LanguageModel, debug bool, authHandler MCPAuthHandler) *MCPConnectionPool {
func NewMCPConnectionPool(config *ConnectionPoolConfig, debug bool, authHandler MCPAuthHandler, tokenStoreFactory TokenStoreFactory) *MCPConnectionPool {
if config == nil {
config = DefaultConnectionPoolConfig()
}
ctx, cancel := context.WithCancel(context.Background())
pool := &MCPConnectionPool{
connections: make(map[string]*MCPConnection),
config: config,
model: model,
ctx: ctx,
cancel: cancel,
debug: debug,
connections: make(map[string]*MCPConnection),
config: config,
ctx: ctx,
cancel: cancel,
debug: debug,
tokenStoreFactory: tokenStoreFactory,
}
if authHandler != nil {
@@ -244,10 +243,12 @@ func (p *MCPConnectionPool) performHealthCheck(ctx context.Context, conn *MCPCon
// createConnection creates a new connection
func (p *MCPConnectionPool) createConnection(ctx context.Context, serverName string, serverConfig config.MCPServerConfig) (*MCPConnection, error) {
oauthEnabled := p.oauthFlow != nil && !serverConfig.NoOAuth
mcpClient, err := p.createMCPClient(ctx, serverName, serverConfig)
if err != nil {
// SSE transport can return OAuth error during Start()
if p.oauthFlow != nil && IsOAuthError(err) {
if oauthEnabled && IsOAuthError(err) {
if flowErr := p.oauthFlow.RunAuthFlow(ctx, serverName, err); flowErr != nil {
return nil, fmt.Errorf("OAuth authorization failed: %w", flowErr)
}
@@ -263,7 +264,7 @@ func (p *MCPConnectionPool) createConnection(ctx context.Context, serverName str
if err := p.initializeClient(ctx, mcpClient); err != nil {
// Streamable HTTP transport returns OAuth error during Initialize()
if p.oauthFlow != nil && IsOAuthError(err) {
if oauthEnabled && IsOAuthError(err) {
if flowErr := p.oauthFlow.RunAuthFlow(ctx, serverName, err); flowErr != nil {
_ = mcpClient.Close()
return nil, fmt.Errorf("OAuth authorization failed: %w", flowErr)
@@ -306,6 +307,8 @@ func (p *MCPConnectionPool) createMCPClient(ctx context.Context, serverName stri
return p.createSSEClient(ctx, serverConfig)
case "streamable":
return p.createStreamableClient(ctx, serverConfig)
case "inprocess":
return p.createInProcessClient(serverConfig)
default:
return nil, fmt.Errorf("unsupported transport type '%s' for server %s", transportType, serverName)
}
@@ -362,20 +365,30 @@ func (p *MCPConnectionPool) createSSEClient(ctx context.Context, serverConfig co
}
}
// Enable OAuth for remote transports when an auth handler is configured.
// The OAuthConfig uses PKCE and the handler's redirect URI. Client ID and
// scopes are discovered automatically via dynamic client registration and
// server metadata (RFC 9728).
if p.oauthFlow != nil {
tokenStore, tsErr := NewFileTokenStore(serverConfig.URL)
// Enable OAuth for remote transports when an auth handler is configured
// and the server hasn't opted out via NoOAuth. Public MCP servers (e.g.
// PubMed) set NoOAuth to skip dynamic client registration and token
// exchange, which would otherwise fail with a 404.
if p.oauthFlow != nil && !serverConfig.NoOAuth {
tokenStore, tsErr := p.createTokenStore(serverConfig.URL)
if tsErr != nil {
return nil, fmt.Errorf("failed to create token store: %w", tsErr)
}
options = append(options, transport.WithOAuth(transport.OAuthConfig{
oauthCfg := transport.OAuthConfig{
RedirectURI: p.oauthFlow.handler.RedirectURI(),
PKCEEnabled: true,
TokenStore: tokenStore,
}))
}
if serverConfig.OAuthClientID != "" {
oauthCfg.ClientID = serverConfig.OAuthClientID
}
if serverConfig.OAuthClientSecret != "" {
oauthCfg.ClientSecret = serverConfig.OAuthClientSecret
}
if len(serverConfig.OAuthScopes) > 0 {
oauthCfg.Scopes = serverConfig.OAuthScopes
}
options = append(options, transport.WithOAuth(oauthCfg))
}
sseClient, err := client.NewSSEMCPClient(serverConfig.URL, options...)
@@ -409,20 +422,28 @@ func (p *MCPConnectionPool) createStreamableClient(ctx context.Context, serverCo
}
}
// Enable OAuth for remote transports when an auth handler is configured.
// The OAuthConfig uses PKCE and the handler's redirect URI. Client ID and
// scopes are discovered automatically via dynamic client registration and
// server metadata (RFC 9728).
if p.oauthFlow != nil {
tokenStore, tsErr := NewFileTokenStore(serverConfig.URL)
// Enable OAuth for remote transports when an auth handler is configured
// and the server hasn't opted out via NoOAuth.
if p.oauthFlow != nil && !serverConfig.NoOAuth {
tokenStore, tsErr := p.createTokenStore(serverConfig.URL)
if tsErr != nil {
return nil, fmt.Errorf("failed to create token store: %w", tsErr)
}
options = append(options, transport.WithHTTPOAuth(transport.OAuthConfig{
oauthCfg := transport.OAuthConfig{
RedirectURI: p.oauthFlow.handler.RedirectURI(),
PKCEEnabled: true,
TokenStore: tokenStore,
}))
}
if serverConfig.OAuthClientID != "" {
oauthCfg.ClientID = serverConfig.OAuthClientID
}
if serverConfig.OAuthClientSecret != "" {
oauthCfg.ClientSecret = serverConfig.OAuthClientSecret
}
if len(serverConfig.OAuthScopes) > 0 {
oauthCfg.Scopes = serverConfig.OAuthScopes
}
options = append(options, transport.WithHTTPOAuth(oauthCfg))
}
streamableClient, err := client.NewStreamableHttpClient(serverConfig.URL, options...)
@@ -437,6 +458,32 @@ func (p *MCPConnectionPool) createStreamableClient(ctx context.Context, serverCo
return streamableClient, nil
}
// createInProcessClient creates an in-process MCP client that communicates
// directly with an *server.MCPServer in the same process. No subprocess is
// spawned and no network I/O occurs — calls go through JSON marshal →
// MCPServer.HandleMessage → JSON unmarshal, all in-memory.
func (p *MCPConnectionPool) createInProcessClient(serverConfig config.MCPServerConfig) (client.MCPClient, error) {
srv, ok := serverConfig.InProcessServer.(*server.MCPServer)
if !ok {
return nil, fmt.Errorf("InProcessServer must be *server.MCPServer, got %T", serverConfig.InProcessServer)
}
inProcessClient, err := client.NewInProcessClient(srv)
if err != nil {
return nil, fmt.Errorf("failed to create in-process client: %w", err)
}
return inProcessClient, nil
}
// createTokenStore creates a token store for the given server URL.
// If a custom TokenStoreFactory is configured, it is used; otherwise the
// default file-backed token store is created.
func (p *MCPConnectionPool) createTokenStore(serverURL string) (transport.TokenStore, error) {
if p.tokenStoreFactory != nil {
return p.tokenStoreFactory(serverURL)
}
return NewFileTokenStore(serverURL)
}
// initializeClient initializes the client
func (p *MCPConnectionPool) initializeClient(ctx context.Context, client client.MCPClient) error {
initCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
@@ -583,6 +630,27 @@ func (p *MCPConnectionPool) GetClients() map[string]client.MCPClient {
return clients
}
// RemoveConnection closes and removes a single connection from the pool.
// Returns an error if the connection does not exist or if closing fails.
// Thread-safe for concurrent use.
func (p *MCPConnectionPool) RemoveConnection(serverName string) error {
p.mu.Lock()
defer p.mu.Unlock()
conn, exists := p.connections[serverName]
if !exists {
return fmt.Errorf("connection %q not found in pool", serverName)
}
err := conn.client.Close()
delete(p.connections, serverName)
if p.debugLogger != nil && p.debugLogger.IsDebugEnabled() {
p.debugLogger.LogDebug(fmt.Sprintf("[POOL] Removed connection %s", serverName))
}
return err
}
// Close gracefully shuts down the connection pool, closing all client connections
// and stopping the background health check goroutine. It attempts to close all
// connections even if some fail, logging any errors encountered.
-109
View File
@@ -1,109 +0,0 @@
package tools
import (
"context"
"encoding/json"
"fmt"
"charm.land/fantasy"
"github.com/mark3labs/mcp-go/mcp"
)
// mcpFantasyTool adapts an MCP tool to the fantasy.AgentTool interface.
// It bridges the MCP tool protocol with fantasy's agent tool system, handling
// name prefixing, schema conversion, connection pooling, and result marshaling.
type mcpFantasyTool struct {
toolInfo fantasy.ToolInfo
mapping *toolMapping
providerOptions fantasy.ProviderOptions
}
// Info returns the fantasy tool info including name, description, and parameter schema.
func (t *mcpFantasyTool) Info() fantasy.ToolInfo {
return t.toolInfo
}
// Run executes the MCP tool by routing through the connection pool.
// It maps the prefixed tool name back to the original name, retrieves a healthy
// connection, invokes the tool, and converts the MCP result to a fantasy ToolResponse.
func (t *mcpFantasyTool) Run(ctx context.Context, call fantasy.ToolCall) (fantasy.ToolResponse, error) {
// Parse and validate JSON arguments
var arguments any
input := call.Input
if input == "" || input == "{}" {
arguments = nil
} else {
var temp any
if err := json.Unmarshal([]byte(input), &temp); err != nil {
return fantasy.NewTextErrorResponse(fmt.Sprintf("invalid JSON arguments: %v", err)), nil
}
arguments = json.RawMessage(input)
}
// Get connection from pool with health check
conn, err := t.mapping.manager.connectionPool.GetConnectionWithHealthCheck(
ctx, t.mapping.serverName, t.mapping.serverConfig,
)
if err != nil {
return fantasy.ToolResponse{}, fmt.Errorf("failed to get healthy connection from pool: %w", err)
}
// Call the MCP tool using the original (unprefixed) name
result, err := conn.client.CallTool(ctx, mcp.CallToolRequest{
Request: mcp.Request{
Method: "tools/call",
},
Params: mcp.CallToolParams{
Name: t.mapping.originalName,
Arguments: arguments,
},
})
if err != nil {
// Handle OAuth re-authorization: token may have expired mid-session.
if t.mapping.manager.connectionPool.oauthFlow != nil && IsOAuthError(err) {
if flowErr := t.mapping.manager.connectionPool.oauthFlow.RunAuthFlow(ctx, t.mapping.serverName, err); flowErr != nil {
return fantasy.ToolResponse{}, fmt.Errorf("OAuth re-authorization failed for tool %s: %w", t.mapping.originalName, flowErr)
}
// Retry the tool call after successful re-auth.
result, err = conn.client.CallTool(ctx, mcp.CallToolRequest{
Request: mcp.Request{
Method: "tools/call",
},
Params: mcp.CallToolParams{
Name: t.mapping.originalName,
Arguments: arguments,
},
})
if err != nil {
t.mapping.manager.connectionPool.HandleConnectionError(t.mapping.serverName, err)
return fantasy.ToolResponse{}, fmt.Errorf("failed to call mcp tool after re-auth: %w", err)
}
} else {
// Mark connection as unhealthy for automatic recovery
t.mapping.manager.connectionPool.HandleConnectionError(t.mapping.serverName, err)
return fantasy.ToolResponse{}, fmt.Errorf("failed to call mcp tool: %w", err)
}
}
// Marshal the MCP result to JSON string
marshaledResult, err := json.Marshal(result)
if err != nil {
return fantasy.ToolResponse{}, fmt.Errorf("failed to marshal mcp tool result: %w", err)
}
// Return as text response, preserving error status from MCP
if result.IsError {
return fantasy.NewTextErrorResponse(string(marshaledResult)), nil
}
return fantasy.NewTextResponse(string(marshaledResult)), nil
}
// ProviderOptions returns provider-specific options for this tool.
func (t *mcpFantasyTool) ProviderOptions() fantasy.ProviderOptions {
return t.providerOptions
}
// SetProviderOptions sets provider-specific options for this tool.
func (t *mcpFantasyTool) SetProviderOptions(opts fantasy.ProviderOptions) {
t.providerOptions = opts
}
+244
View File
@@ -0,0 +1,244 @@
package tools
import (
"context"
"encoding/json"
"strings"
"testing"
"github.com/mark3labs/kit/internal/config"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
// newTestInProcessServer creates a simple MCP server with one tool for testing.
func newTestInProcessServer() *server.MCPServer {
srv := server.NewMCPServer("test-server", "1.0.0",
server.WithToolCapabilities(true),
)
srv.AddTool(
mcp.NewTool("greet",
mcp.WithDescription("Say hello"),
mcp.WithString("name", mcp.Required(), mcp.Description("Name to greet")),
),
func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
name, _ := req.GetArguments()["name"].(string)
return mcp.NewToolResultText("Hello, " + name + "!"), nil
},
)
return srv
}
func TestInProcessTransportType(t *testing.T) {
cfg := config.MCPServerConfig{
Type: "inprocess",
InProcessServer: newTestInProcessServer(),
}
if got := cfg.GetTransportType(); got != "inprocess" {
t.Errorf("GetTransportType() = %q, want %q", got, "inprocess")
}
}
func TestInProcessTransportTypeInferred(t *testing.T) {
// When Type is empty but InProcessServer is set, infer "inprocess".
cfg := config.MCPServerConfig{
InProcessServer: newTestInProcessServer(),
}
if got := cfg.GetTransportType(); got != "inprocess" {
t.Errorf("GetTransportType() = %q, want %q", got, "inprocess")
}
}
func TestInProcessValidation(t *testing.T) {
// Valid: InProcessServer is set.
validCfg := &config.Config{
MCPServers: map[string]config.MCPServerConfig{
"test": {
Type: "inprocess",
InProcessServer: newTestInProcessServer(),
},
},
}
if err := validCfg.Validate(); err != nil {
t.Errorf("expected valid config, got error: %v", err)
}
// Invalid: type is inprocess but InProcessServer is nil.
invalidCfg := &config.Config{
MCPServers: map[string]config.MCPServerConfig{
"test": {
Type: "inprocess",
},
},
}
if err := invalidCfg.Validate(); err == nil {
t.Error("expected validation error for nil InProcessServer, got nil")
}
}
func TestConnectionPoolInProcessClient(t *testing.T) {
pool := NewMCPConnectionPool(DefaultConnectionPoolConfig(), false, nil, nil)
defer func() { _ = pool.Close() }()
ctx := context.Background()
srv := newTestInProcessServer()
cfg := config.MCPServerConfig{
Type: "inprocess",
InProcessServer: srv,
}
conn, err := pool.GetConnection(ctx, "test-inproc", cfg)
if err != nil {
t.Fatalf("GetConnection failed: %v", err)
}
// Verify the connection is healthy and functional.
if !conn.isHealthy {
t.Error("expected connection to be healthy")
}
// List tools to verify the connection works end-to-end.
toolsResp, err := conn.client.ListTools(ctx, mcp.ListToolsRequest{})
if err != nil {
t.Fatalf("ListTools failed: %v", err)
}
if len(toolsResp.Tools) != 1 {
t.Fatalf("expected 1 tool, got %d", len(toolsResp.Tools))
}
if toolsResp.Tools[0].Name != "greet" {
t.Errorf("expected tool name 'greet', got %q", toolsResp.Tools[0].Name)
}
}
func TestConnectionPoolInProcessToolExecution(t *testing.T) {
pool := NewMCPConnectionPool(DefaultConnectionPoolConfig(), false, nil, nil)
defer func() { _ = pool.Close() }()
ctx := context.Background()
srv := newTestInProcessServer()
cfg := config.MCPServerConfig{
Type: "inprocess",
InProcessServer: srv,
}
conn, err := pool.GetConnection(ctx, "test-inproc", cfg)
if err != nil {
t.Fatalf("GetConnection failed: %v", err)
}
// Call the tool.
result, err := conn.client.CallTool(ctx, mcp.CallToolRequest{
Request: mcp.Request{Method: "tools/call"},
Params: mcp.CallToolParams{
Name: "greet",
Arguments: map[string]any{"name": "World"},
},
})
if err != nil {
t.Fatalf("CallTool failed: %v", err)
}
if result.IsError {
t.Error("expected non-error result")
}
if len(result.Content) == 0 {
t.Fatal("expected at least one content block")
}
text, ok := result.Content[0].(mcp.TextContent)
if !ok {
t.Fatalf("expected TextContent, got %T", result.Content[0])
}
if text.Text != "Hello, World!" {
t.Errorf("expected 'Hello, World!', got %q", text.Text)
}
}
func TestMCPToolManagerInProcess(t *testing.T) {
ctx := context.Background()
srv := newTestInProcessServer()
mgr := NewMCPToolManager()
cfg := config.MCPServerConfig{
Type: "inprocess",
InProcessServer: srv,
}
count, err := mgr.AddServer(ctx, "myserver", cfg)
if err != nil {
t.Fatalf("AddServer failed: %v", err)
}
if count != 1 {
t.Errorf("expected 1 tool, got %d", count)
}
tools := mgr.GetTools()
if len(tools) != 1 {
t.Fatalf("expected 1 tool, got %d", len(tools))
}
if tools[0].Name != "myserver__greet" {
t.Errorf("expected tool name 'myserver__greet', got %q", tools[0].Name)
}
// Execute the tool.
input, _ := json.Marshal(map[string]any{"name": "SDK"})
result, err := mgr.ExecuteTool(ctx, "myserver__greet", string(input))
if err != nil {
t.Fatalf("ExecuteTool failed: %v", err)
}
if result.IsError {
t.Error("expected non-error result")
}
if result.Content == "" {
t.Error("expected non-empty result content")
}
// Verify result contains our greeting.
if !strings.Contains(result.Content, "Hello, SDK!") {
t.Errorf("expected 'Hello, SDK!' in result, got %q", result.Content)
}
}
func TestConnectionPoolInProcessInvalidServer(t *testing.T) {
pool := NewMCPConnectionPool(DefaultConnectionPoolConfig(), false, nil, nil)
defer func() { _ = pool.Close() }()
ctx := context.Background()
// Pass a non-*server.MCPServer value.
cfg := config.MCPServerConfig{
Type: "inprocess",
InProcessServer: "not a server",
}
_, err := pool.GetConnection(ctx, "bad", cfg)
if err == nil {
t.Fatal("expected error for invalid InProcessServer type")
}
}
func TestConnectionPoolInProcessReuse(t *testing.T) {
pool := NewMCPConnectionPool(DefaultConnectionPoolConfig(), false, nil, nil)
defer func() { _ = pool.Close() }()
ctx := context.Background()
srv := newTestInProcessServer()
cfg := config.MCPServerConfig{
Type: "inprocess",
InProcessServer: srv,
}
// Get connection twice — should reuse.
conn1, err := pool.GetConnection(ctx, "reuse-test", cfg)
if err != nil {
t.Fatalf("first GetConnection failed: %v", err)
}
conn2, err := pool.GetConnection(ctx, "reuse-test", cfg)
if err != nil {
t.Fatalf("second GetConnection failed: %v", err)
}
if conn1 != conn2 {
t.Error("expected same connection object on reuse")
}
}
+821 -44
View File
@@ -2,6 +2,7 @@ package tools
import (
"context"
"encoding/base64"
"encoding/json"
"fmt"
"maps"
@@ -9,30 +10,149 @@ import (
"strings"
"sync"
"charm.land/fantasy"
log "github.com/charmbracelet/log"
"github.com/mark3labs/kit/internal/config"
"github.com/mark3labs/mcp-go/mcp"
)
// MCPTool represents a tool discovered from an MCP server. It contains all
// the metadata needed to present the tool to an LLM (name, description, JSON
// schema) plus the server origin information needed to execute it.
type MCPTool struct {
// Name is the prefixed tool name: "serverName__toolName".
Name string
// Description is the human-readable tool description.
Description string
// Parameters is the JSON Schema properties for the tool's input.
Parameters map[string]any
// Required lists the required parameter names.
Required []string
// ServerName is the MCP server this tool belongs to.
ServerName string
// OriginalName is the unprefixed tool name on the MCP server.
OriginalName string
}
// MCPToolResult is the result of executing an MCP tool via ExecuteTool.
type MCPToolResult struct {
// Content is the JSON-encoded result from the MCP server.
Content string
// IsError indicates the MCP server reported a tool-level error.
IsError bool
}
// MCPPrompt represents a prompt discovered from an MCP server.
type MCPPrompt struct {
// Name is the prompt name on the MCP server.
Name string
// Description is the human-readable prompt description.
Description string
// Arguments lists the prompt's expected arguments.
Arguments []MCPPromptArgument
// ServerName is the MCP server this prompt belongs to.
ServerName string
}
// MCPPromptArgument describes an argument that a prompt template can accept.
type MCPPromptArgument struct {
// Name is the argument name.
Name string
// Description is a human-readable description.
Description string
// Required indicates whether this argument must be provided.
Required bool
}
// MCPPromptMessage is a single message returned by a prompt expansion.
type MCPPromptMessage struct {
// Role is "user" or "assistant".
Role string
// Content is the text content of the message.
Content string
// FileParts contains binary attachments extracted from embedded resources,
// images, or audio content blocks. Empty for text-only messages.
FileParts []MCPFilePart
}
// MCPFilePart represents a binary file attachment extracted from an MCP prompt
// content block (ImageContent, AudioContent, or EmbeddedResource with blob data).
type MCPFilePart struct {
// Filename is a best-effort name derived from the resource URI or content type.
Filename string
// Data is the raw binary content (already base64-decoded).
Data []byte
// MediaType is the MIME type (e.g. "image/png", "audio/wav").
MediaType string
}
// MCPPromptResult is the result of expanding an MCP prompt via GetPrompt.
type MCPPromptResult struct {
// Description is an optional description returned by the server.
Description string
// Messages contains the expanded prompt messages.
Messages []MCPPromptMessage
}
// MCPResource represents a resource discovered from an MCP server.
type MCPResource struct {
// URI is the unique resource identifier (e.g. "file:///path" or custom scheme).
URI string
// Name is a human-readable name for the resource.
Name string
// Description is an optional description of the resource.
Description string
// MIMEType is the MIME type of the resource, if known.
MIMEType string
// ServerName is the MCP server this resource belongs to.
ServerName string
}
// MCPResourceContent is the result of reading an MCP resource via ReadResource.
type MCPResourceContent struct {
// URI is the resource URI that was read.
URI string
// MIMEType is the MIME type of the content.
MIMEType string
// Text is the text content (non-empty for text resources).
Text string
// BlobData is the decoded binary content (non-empty for blob resources).
BlobData []byte
// IsBlob is true when the content is binary (BlobData is set).
IsBlob bool
}
// MCPToolManager manages MCP (Model Context Protocol) tools and clients across multiple servers.
// It provides a unified interface for loading, managing, and executing tools from various MCP servers,
// including stdio, SSE, streamable HTTP, and built-in server types. The manager handles connection
// pooling, health checks, tool name prefixing to avoid conflicts, and sampling support for LLM interactions.
// pooling, health checks, tool name prefixing to avoid conflicts, and OAuth re-authorization.
// Thread-safe for concurrent tool invocations.
type MCPToolManager struct {
connectionPool *MCPConnectionPool
tools []fantasy.AgentTool
toolMap map[string]*toolMapping // maps prefixed tool names to their server and original name
mu sync.Mutex // protects tools and toolMap during parallel loading
model fantasy.LanguageModel // LLM model for sampling
authHandler MCPAuthHandler // OAuth handler for remote servers (nil = no OAuth)
config *config.Config
debug bool
debugLogger DebugLogger
connectionPool *MCPConnectionPool
tools []MCPTool
toolMap map[string]*toolMapping // maps prefixed tool names to their server and original name
prompts []MCPPrompt // prompts discovered from all servers
resources []MCPResource // resources discovered from all servers
subscriptions map[string]string // resource URI → server name for active subscriptions
mu sync.Mutex // protects tools, toolMap, prompts, resources during parallel loading
authHandler MCPAuthHandler // OAuth handler for remote servers (nil = no OAuth)
tokenStoreFactory TokenStoreFactory // factory for creating per-server token stores (nil = default FileTokenStore)
config *config.Config
debug bool
debugLogger DebugLogger
// onServerLoaded, if non-nil, is called when each server finishes loading.
// Called with server name, tool count, and error (nil on success).
onServerLoaded func(serverName string, toolCount int, err error)
// onToolsChanged, if non-nil, is called after AddServer or RemoveServer
// mutates the tool list. The agent layer uses this to trigger a rebuild
// so the LLM sees the updated tools.
onToolsChanged func()
// onResourcesChanged, if non-nil, is called when a subscribed resource
// is updated by the server.
onResourcesChanged func()
}
// toolMapping stores the mapping between prefixed tool names and their original details
@@ -40,27 +160,19 @@ type toolMapping struct {
serverName string
originalName string
serverConfig config.MCPServerConfig
manager *MCPToolManager
}
// NewMCPToolManager creates a new MCP tool manager instance.
// Returns an initialized manager with empty tool collections ready to load tools from MCP servers.
// The manager must be configured with SetModel and LoadTools before use.
// The manager must be configured with LoadTools before use.
func NewMCPToolManager() *MCPToolManager {
return &MCPToolManager{
tools: make([]fantasy.AgentTool, 0),
toolMap: make(map[string]*toolMapping),
tools: make([]MCPTool, 0),
toolMap: make(map[string]*toolMapping),
subscriptions: make(map[string]string),
}
}
// SetModel sets the LLM model for sampling support.
// The model is used when MCP servers request sampling operations, allowing them to
// leverage the host's LLM capabilities for text generation tasks.
// This method should be called before LoadTools if any MCP servers require sampling support.
func (m *MCPToolManager) SetModel(model fantasy.LanguageModel) {
m.model = model
}
// SetAuthHandler sets the OAuth handler for remote MCP server authentication.
// When set, remote transports (streamable HTTP, SSE) are configured with OAuth
// support, enabling automatic authorization flows when servers require authentication.
@@ -69,6 +181,20 @@ func (m *MCPToolManager) SetAuthHandler(handler MCPAuthHandler) {
m.authHandler = handler
}
// GetAuthHandler returns the OAuth handler for remote MCP server authentication.
// Returns nil if no handler is configured.
func (m *MCPToolManager) GetAuthHandler() MCPAuthHandler {
return m.authHandler
}
// SetTokenStoreFactory sets a custom factory for creating per-server OAuth token
// stores. When set, the factory is called for each remote MCP server instead of
// using the default file-based token store. This method should be called before
// LoadTools.
func (m *MCPToolManager) SetTokenStoreFactory(factory TokenStoreFactory) {
m.tokenStoreFactory = factory
}
// SetDebugLogger sets the debug logger for the tool manager.
// The logger will be used to output detailed debugging information about MCP connections,
// tool loading, and execution. If a connection pool exists, it will also be configured
@@ -87,6 +213,157 @@ func (m *MCPToolManager) SetOnServerLoaded(cb func(serverName string, toolCount
m.onServerLoaded = cb
}
// SetOnToolsChanged sets the callback that's invoked after AddServer or
// RemoveServer mutates the tool list. The agent layer uses this to trigger
// a rebuild so the LLM sees the updated tool set.
func (m *MCPToolManager) SetOnToolsChanged(cb func()) {
m.onToolsChanged = cb
}
// AddServer connects to a new MCP server at runtime and loads its tools.
// The server's tools are immediately available to the agent after this call.
// Returns the number of tools loaded from the server.
//
// If the connection pool has not been initialised yet (i.e. LoadTools was never
// called), AddServer creates one automatically using the manager's current
// configuration.
//
// Returns an error if a server with the same name is already loaded, or if
// the connection or tool loading fails.
func (m *MCPToolManager) AddServer(ctx context.Context, name string, cfg config.MCPServerConfig) (int, error) {
m.mu.Lock()
// Check for duplicate.
if _, exists := m.toolMap[name+"__"]; exists {
m.mu.Unlock()
return 0, fmt.Errorf("MCP server %q is already loaded", name)
}
// More thorough duplicate check: scan toolMap for any key with the server prefix.
prefix := name + "__"
for k := range m.toolMap {
if len(k) >= len(prefix) && k[:len(prefix)] == prefix {
m.mu.Unlock()
return 0, fmt.Errorf("MCP server %q is already loaded", name)
}
}
m.mu.Unlock()
// Lazily create the connection pool if LoadTools was never called.
m.ensureConnectionPool()
count, err := m.loadServerTools(ctx, name, cfg)
if err != nil {
return 0, fmt.Errorf("failed to add MCP server %q: %w", name, err)
}
// Notify listeners.
if m.onServerLoaded != nil {
m.onServerLoaded(name, count, nil)
}
if m.onToolsChanged != nil {
m.onToolsChanged()
}
return count, nil
}
// RemoveServer disconnects an MCP server and removes all its tools and prompts.
// After this call the agent will no longer see or be able to call tools from
// the named server. Returns an error if the server is not loaded.
func (m *MCPToolManager) RemoveServer(name string) error {
prefix := name + "__"
m.mu.Lock()
// Check the server actually has tools or prompts loaded.
found := false
for k := range m.toolMap {
if len(k) >= len(prefix) && k[:len(prefix)] == prefix {
found = true
break
}
}
if !found {
// Also check prompts — a server might expose only prompts.
for _, p := range m.prompts {
if p.ServerName == name {
found = true
break
}
}
}
if !found {
m.mu.Unlock()
return fmt.Errorf("MCP server %q is not loaded", name)
}
// Remove tools belonging to this server.
newTools := make([]MCPTool, 0, len(m.tools))
for _, t := range m.tools {
if len(t.Name) < len(prefix) || t.Name[:len(prefix)] != prefix {
newTools = append(newTools, t)
}
}
m.tools = newTools
// Remove tool mappings.
for k := range m.toolMap {
if len(k) >= len(prefix) && k[:len(prefix)] == prefix {
delete(m.toolMap, k)
}
}
// Remove prompts belonging to this server.
newPrompts := make([]MCPPrompt, 0, len(m.prompts))
for _, p := range m.prompts {
if p.ServerName != name {
newPrompts = append(newPrompts, p)
}
}
m.prompts = newPrompts
// Remove resources belonging to this server.
newResources := make([]MCPResource, 0, len(m.resources))
for _, r := range m.resources {
if r.ServerName != name {
newResources = append(newResources, r)
} else {
// Clean up any active subscription for this resource.
delete(m.subscriptions, r.URI)
}
}
m.resources = newResources
m.mu.Unlock()
// Close the connection in the pool (best-effort).
if m.connectionPool != nil {
_ = m.connectionPool.RemoveConnection(name)
}
if m.onToolsChanged != nil {
m.onToolsChanged()
}
return nil
}
// ensureConnectionPool lazily creates a connection pool if one does not exist.
// This allows AddServer to work even if LoadTools was never called.
func (m *MCPToolManager) ensureConnectionPool() {
if m.connectionPool != nil {
return
}
debug := false
if m.config != nil {
debug = m.config.Debug
}
if m.debugLogger == nil {
m.debugLogger = NewSimpleDebugLogger(debug)
}
m.connectionPool = NewMCPConnectionPool(DefaultConnectionPoolConfig(), debug, m.authHandler, m.tokenStoreFactory)
m.connectionPool.SetDebugLogger(m.debugLogger)
}
// LoadTools loads tools from all configured MCP servers based on the provided configuration.
// It initializes the connection pool, connects to each configured server, and loads their tools.
// Tools from different servers are prefixed with the server name to avoid naming conflicts.
@@ -99,7 +376,7 @@ func (m *MCPToolManager) LoadTools(ctx context.Context, cfg *config.Config) erro
if m.debugLogger == nil {
m.debugLogger = NewSimpleDebugLogger(cfg.Debug)
}
m.connectionPool = NewMCPConnectionPool(DefaultConnectionPoolConfig(), m.model, cfg.Debug, m.authHandler)
m.connectionPool = NewMCPConnectionPool(DefaultConnectionPoolConfig(), cfg.Debug, m.authHandler, m.tokenStoreFactory)
m.connectionPool.SetDebugLogger(m.debugLogger)
// Load all servers in parallel. Each server connection (subprocess
@@ -181,10 +458,10 @@ func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string,
}
// Build tools locally before acquiring the lock.
var localTools []fantasy.AgentTool
var localTools []MCPTool
localMap := make(map[string]*toolMapping)
// Convert MCP tools to fantasy AgentTools with prefixed names
// Convert MCP tools to MCPTool structs with prefixed names
for _, mcpTool := range listResults.Tools {
// Filter tools based on allowedTools/excludedTools
if len(serverConfig.AllowedTools) > 0 {
@@ -198,7 +475,7 @@ func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string,
continue
}
// Convert MCP InputSchema to map[string]any for fantasy ToolInfo
// Convert MCP InputSchema to map[string]any
marshaledSchema, err := json.Marshal(mcpTool.InputSchema)
if err != nil {
return -1, fmt.Errorf("conv mcp tool input schema fail(marshal): %w, tool name: %s", err, mcpTool.Name)
@@ -207,7 +484,7 @@ func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string,
// Fix for JSON Schema draft-07 vs draft-04 compatibility
marshaledSchema = convertExclusiveBoundsToBoolean(marshaledSchema)
// Parse into map[string]any for fantasy's parameters format
// Parse into map[string]any
var schemaMap map[string]any
if err := json.Unmarshal(marshaledSchema, &schemaMap); err != nil {
return -1, fmt.Errorf("conv mcp tool input schema fail(unmarshal): %w, tool name: %s", err, mcpTool.Name)
@@ -223,7 +500,7 @@ func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string,
// Fix for issue #89: Ensure object schemas have a properties field.
// When schema type is "object" with no properties, we keep the
// empty parameters map — fantasy handles this fine.
// empty parameters map.
if req, ok := schemaMap["required"].([]any); ok {
for _, r := range req {
@@ -241,22 +518,18 @@ func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string,
serverName: serverName,
originalName: mcpTool.Name,
serverConfig: serverConfig,
manager: m,
}
localMap[prefixedName] = mapping
// Create fantasy AgentTool
fantasyTool := &mcpFantasyTool{
toolInfo: fantasy.ToolInfo{
Name: prefixedName,
Description: mcpTool.Description,
Parameters: parameters,
Required: required,
},
mapping: mapping,
}
localTools = append(localTools, fantasyTool)
// Create MCPTool
localTools = append(localTools, MCPTool{
Name: prefixedName,
Description: mcpTool.Description,
Parameters: parameters,
Required: required,
ServerName: serverName,
OriginalName: mcpTool.Name,
})
}
// Merge into the manager under the lock.
@@ -265,15 +538,516 @@ func (m *MCPToolManager) loadServerTools(ctx context.Context, serverName string,
m.tools = append(m.tools, localTools...)
m.mu.Unlock()
// Also load prompts from this server (best-effort, non-blocking).
m.loadServerPrompts(ctx, serverName, conn)
// Also load resources from this server (best-effort, non-blocking).
m.loadServerResources(ctx, serverName, conn)
return len(localTools), nil
}
// GetTools returns all loaded tools as fantasy AgentTools from all configured MCP servers.
// ExecuteTool calls an MCP tool through the connection pool, handling health
// checks, OAuth re-authorization, and connection error tracking.
// The inputJSON parameter is the raw JSON arguments from the LLM.
// Returns the result content, error flag, and any execution error.
func (m *MCPToolManager) ExecuteTool(ctx context.Context, prefixedName, inputJSON string) (*MCPToolResult, error) {
m.mu.Lock()
mapping, ok := m.toolMap[prefixedName]
m.mu.Unlock()
if !ok {
return nil, fmt.Errorf("tool %q not found", prefixedName)
}
// Parse and validate JSON arguments
var arguments any
if inputJSON == "" || inputJSON == "{}" {
arguments = nil
} else {
var temp any
if err := json.Unmarshal([]byte(inputJSON), &temp); err != nil {
return &MCPToolResult{
Content: fmt.Sprintf("invalid JSON arguments: %v", err),
IsError: true,
}, nil
}
arguments = json.RawMessage(inputJSON)
}
// Get connection from pool with health check
conn, err := m.connectionPool.GetConnectionWithHealthCheck(
ctx, mapping.serverName, mapping.serverConfig,
)
if err != nil {
return nil, fmt.Errorf("failed to get healthy connection from pool: %w", err)
}
callRequest := mcp.CallToolRequest{
Request: mcp.Request{
Method: "tools/call",
},
Params: mcp.CallToolParams{
Name: mapping.originalName,
Arguments: arguments,
},
}
// Call the MCP tool using the original (unprefixed) name
result, err := conn.client.CallTool(ctx, callRequest)
if err != nil {
// Handle OAuth re-authorization: token may have expired mid-session.
if m.connectionPool.oauthFlow != nil && IsOAuthError(err) {
if flowErr := m.connectionPool.oauthFlow.RunAuthFlow(ctx, mapping.serverName, err); flowErr != nil {
return nil, fmt.Errorf("OAuth re-authorization failed for tool %s: %w", mapping.originalName, flowErr)
}
// Retry the tool call after successful re-auth.
result, err = conn.client.CallTool(ctx, callRequest)
if err != nil {
m.connectionPool.HandleConnectionError(mapping.serverName, err)
return nil, fmt.Errorf("failed to call mcp tool after re-auth: %w", err)
}
} else {
// Mark connection as unhealthy for automatic recovery
m.connectionPool.HandleConnectionError(mapping.serverName, err)
return nil, fmt.Errorf("failed to call mcp tool: %w", err)
}
}
// Marshal the MCP result to JSON string
marshaledResult, 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: result.IsError,
}, nil
}
// GetTools returns all loaded MCP tools from all configured MCP servers.
// Tools are returned with their prefixed names (serverName__toolName) to ensure uniqueness.
func (m *MCPToolManager) GetTools() []fantasy.AgentTool {
func (m *MCPToolManager) GetTools() []MCPTool {
return m.tools
}
// GetPrompts returns all prompts discovered from connected MCP servers.
func (m *MCPToolManager) GetPrompts() []MCPPrompt {
m.mu.Lock()
defer m.mu.Unlock()
result := make([]MCPPrompt, len(m.prompts))
copy(result, m.prompts)
return result
}
// GetPrompt retrieves and expands a specific prompt from an MCP server.
// The serverName identifies which server to query, promptName is the prompt's
// name on that server, and args are the template arguments to substitute.
// This call is lazy — it contacts the MCP server on each invocation.
func (m *MCPToolManager) GetPrompt(ctx context.Context, serverName, promptName string, args map[string]string) (*MCPPromptResult, error) {
if m.connectionPool == nil {
return nil, fmt.Errorf("no connection pool available")
}
clients := m.connectionPool.GetClients()
mcpClient, ok := clients[serverName]
if !ok {
return nil, fmt.Errorf("MCP server %q not found", serverName)
}
req := mcp.GetPromptRequest{}
req.Params.Name = promptName
if len(args) > 0 {
req.Params.Arguments = args
}
result, err := mcpClient.GetPrompt(ctx, req)
if err != nil {
return nil, fmt.Errorf("failed to get prompt %q from server %q: %w", promptName, serverName, err)
}
// Convert MCP messages to our types, extracting all content types.
var messages []MCPPromptMessage
for _, msg := range result.Messages {
text, fileParts := extractPromptContent(msg.Content)
if text != "" || len(fileParts) > 0 {
messages = append(messages, MCPPromptMessage{
Role: string(msg.Role),
Content: text,
FileParts: fileParts,
})
}
}
return &MCPPromptResult{
Description: result.Description,
Messages: messages,
}, nil
}
// extractPromptContent extracts text and binary attachments from an MCP Content value.
// Handles all MCP content types: TextContent, ImageContent, AudioContent,
// EmbeddedResource (text and blob), and ResourceLink.
func extractPromptContent(content mcp.Content) (string, []MCPFilePart) {
switch c := content.(type) {
case mcp.TextContent:
return c.Text, nil
case *mcp.TextContent:
if c != nil {
return c.Text, nil
}
return "", nil
case mcp.ImageContent:
return "", decodeBase64FilePart(c.Data, c.MIMEType, "image/png", "image.png")
case *mcp.ImageContent:
if c != nil {
return "", decodeBase64FilePart(c.Data, c.MIMEType, "image/png", "image.png")
}
return "", nil
case mcp.AudioContent:
return "", decodeBase64FilePart(c.Data, c.MIMEType, "audio/wav", "audio.wav")
case *mcp.AudioContent:
if c != nil {
return "", decodeBase64FilePart(c.Data, c.MIMEType, "audio/wav", "audio.wav")
}
return "", nil
case mcp.EmbeddedResource:
return extractEmbeddedResourceContent(c.Resource)
case *mcp.EmbeddedResource:
if c != nil {
return extractEmbeddedResourceContent(c.Resource)
}
return "", nil
case mcp.ResourceLink:
// ResourceLink is a reference without inline content — include as a
// text annotation so the LLM knows about it.
return fmt.Sprintf("[Referenced resource: %s (%s)]", c.URI, c.Name), nil
case *mcp.ResourceLink:
if c != nil {
return fmt.Sprintf("[Referenced resource: %s (%s)]", c.URI, c.Name), nil
}
return "", nil
default:
return "", nil
}
}
// extractEmbeddedResourceContent handles the two variants of embedded resource
// content: text resources are inlined as fenced code blocks, blob resources
// are base64-decoded into MCPFilePart attachments.
func extractEmbeddedResourceContent(res mcp.ResourceContents) (string, []MCPFilePart) {
switch r := res.(type) {
case mcp.TextResourceContents:
return fmt.Sprintf("[File: %s]\n```\n%s\n```", r.URI, r.Text), nil
case *mcp.TextResourceContents:
if r != nil {
return fmt.Sprintf("[File: %s]\n```\n%s\n```", r.URI, r.Text), nil
}
return "", nil
case mcp.BlobResourceContents:
return "", decodeBase64FilePart(r.Blob, r.MIMEType, "application/octet-stream", filenameFromURI(r.URI))
case *mcp.BlobResourceContents:
if r != nil {
return "", decodeBase64FilePart(r.Blob, r.MIMEType, "application/octet-stream", filenameFromURI(r.URI))
}
return "", nil
default:
return "", nil
}
}
// decodeBase64FilePart decodes base64-encoded data into an MCPFilePart.
// Returns nil on decode failure (logged as a warning).
func decodeBase64FilePart(data, mimeType, defaultMIME, filename string) []MCPFilePart {
decoded, err := base64.StdEncoding.DecodeString(data)
if err != nil {
log.Warn("mcp prompt: failed to decode base64 content", "filename", filename, "error", err)
return nil
}
if mimeType == "" {
mimeType = defaultMIME
}
return []MCPFilePart{{
Filename: filename,
Data: decoded,
MediaType: mimeType,
}}
}
// filenameFromURI extracts a filename from a URI (e.g. "file:///path/to/img.png" → "img.png").
func filenameFromURI(uri string) string {
uri = strings.TrimPrefix(uri, "file://")
if idx := strings.LastIndex(uri, "/"); idx >= 0 {
return uri[idx+1:]
}
if uri == "" {
return "resource"
}
return uri
}
// loadServerPrompts loads prompts from a single MCP server connection.
// Called inside loadServerTools after a successful connection is established.
// Thread-safe: acquires m.mu to merge results.
func (m *MCPToolManager) loadServerPrompts(ctx context.Context, serverName string, conn *MCPConnection) {
listResult, err := conn.client.ListPrompts(ctx, mcp.ListPromptsRequest{})
if err != nil {
// Prompts are optional — servers may not support them.
// Silently skip.
return
}
if len(listResult.Prompts) == 0 {
return
}
var localPrompts []MCPPrompt
for _, p := range listResult.Prompts {
var args []MCPPromptArgument
for _, a := range p.Arguments {
args = append(args, MCPPromptArgument{
Name: a.Name,
Description: a.Description,
Required: a.Required,
})
}
localPrompts = append(localPrompts, MCPPrompt{
Name: p.Name,
Description: p.Description,
Arguments: args,
ServerName: serverName,
})
}
m.mu.Lock()
m.prompts = append(m.prompts, localPrompts...)
m.mu.Unlock()
}
// loadServerResources loads resources from a single MCP server connection.
// Called inside loadServerTools after a successful connection is established.
// Thread-safe: acquires m.mu to merge results.
func (m *MCPToolManager) loadServerResources(ctx context.Context, serverName string, conn *MCPConnection) {
listResult, err := conn.client.ListResources(ctx, mcp.ListResourcesRequest{})
if err != nil {
// Resources are optional — servers may not support them.
return
}
if len(listResult.Resources) == 0 {
return
}
var localResources []MCPResource
for _, r := range listResult.Resources {
localResources = append(localResources, MCPResource{
URI: r.URI,
Name: r.Name,
Description: r.Description,
MIMEType: r.MIMEType,
ServerName: serverName,
})
}
m.mu.Lock()
m.resources = append(m.resources, localResources...)
m.mu.Unlock()
}
// GetResources returns all resources discovered from connected MCP servers.
func (m *MCPToolManager) GetResources() []MCPResource {
m.mu.Lock()
defer m.mu.Unlock()
result := make([]MCPResource, len(m.resources))
copy(result, m.resources)
return result
}
// SetOnResourcesChanged sets the callback invoked when a subscribed resource
// changes. Used by the UI layer to refresh autocomplete or re-read content.
func (m *MCPToolManager) SetOnResourcesChanged(cb func()) {
m.onResourcesChanged = cb
}
// ReadResource reads a specific resource from an MCP server by URI.
// Returns the resource content (text or binary blob).
func (m *MCPToolManager) ReadResource(ctx context.Context, serverName, uri string) (*MCPResourceContent, error) {
if m.connectionPool == nil {
return nil, fmt.Errorf("no connection pool available")
}
clients := m.connectionPool.GetClients()
mcpClient, ok := clients[serverName]
if !ok {
return nil, fmt.Errorf("MCP server %q not found", serverName)
}
req := mcp.ReadResourceRequest{}
req.Params.URI = uri
result, err := mcpClient.ReadResource(ctx, req)
if err != nil {
return nil, fmt.Errorf("failed to read resource %q from server %q: %w", uri, serverName, err)
}
if len(result.Contents) == 0 {
return nil, fmt.Errorf("resource %q returned no content", uri)
}
// Process the first content item (most resources return exactly one).
content := result.Contents[0]
switch c := content.(type) {
case mcp.TextResourceContents:
return &MCPResourceContent{
URI: c.URI,
MIMEType: c.MIMEType,
Text: c.Text,
IsBlob: false,
}, nil
case *mcp.TextResourceContents:
if c == nil {
return nil, fmt.Errorf("resource %q returned nil text content", uri)
}
return &MCPResourceContent{
URI: c.URI,
MIMEType: c.MIMEType,
Text: c.Text,
IsBlob: false,
}, nil
case mcp.BlobResourceContents:
decoded, err := base64.StdEncoding.DecodeString(c.Blob)
if err != nil {
return nil, fmt.Errorf("failed to decode blob resource %q: %w", uri, err)
}
return &MCPResourceContent{
URI: c.URI,
MIMEType: c.MIMEType,
BlobData: decoded,
IsBlob: true,
}, nil
case *mcp.BlobResourceContents:
if c == nil {
return nil, fmt.Errorf("resource %q returned nil blob content", uri)
}
decoded, err := base64.StdEncoding.DecodeString(c.Blob)
if err != nil {
return nil, fmt.Errorf("failed to decode blob resource %q: %w", uri, err)
}
return &MCPResourceContent{
URI: c.URI,
MIMEType: c.MIMEType,
BlobData: decoded,
IsBlob: true,
}, nil
default:
return nil, fmt.Errorf("resource %q returned unknown content type %T", uri, content)
}
}
// SubscribeResource subscribes to change notifications for a resource.
// When the resource changes on the server, onResourcesChanged is called
// and the resource list is refreshed automatically.
func (m *MCPToolManager) SubscribeResource(ctx context.Context, serverName, uri string) error {
if m.connectionPool == nil {
return fmt.Errorf("no connection pool available")
}
clients := m.connectionPool.GetClients()
mcpClient, ok := clients[serverName]
if !ok {
return fmt.Errorf("MCP server %q not found", serverName)
}
req := mcp.SubscribeRequest{}
req.Params.URI = uri
if err := mcpClient.Subscribe(ctx, req); err != nil {
return fmt.Errorf("failed to subscribe to resource %q on server %q: %w", uri, serverName, err)
}
m.mu.Lock()
m.subscriptions[uri] = serverName
m.mu.Unlock()
return nil
}
// UnsubscribeResource cancels change notifications for a resource.
func (m *MCPToolManager) UnsubscribeResource(ctx context.Context, serverName, uri string) error {
if m.connectionPool == nil {
return fmt.Errorf("no connection pool available")
}
clients := m.connectionPool.GetClients()
mcpClient, ok := clients[serverName]
if !ok {
return fmt.Errorf("MCP server %q not found", serverName)
}
req := mcp.UnsubscribeRequest{}
req.Params.URI = uri
if err := mcpClient.Unsubscribe(ctx, req); err != nil {
return fmt.Errorf("failed to unsubscribe from resource %q on server %q: %w", uri, serverName, err)
}
m.mu.Lock()
delete(m.subscriptions, uri)
m.mu.Unlock()
return nil
}
// RefreshServerResources re-fetches resources from a specific server.
// Called when a resource change notification is received.
func (m *MCPToolManager) RefreshServerResources(ctx context.Context, serverName string) {
if m.connectionPool == nil {
return
}
clients := m.connectionPool.GetClients()
mcpClient, ok := clients[serverName]
if !ok {
return
}
listResult, err := mcpClient.ListResources(ctx, mcp.ListResourcesRequest{})
if err != nil {
return
}
var newResources []MCPResource
for _, r := range listResult.Resources {
newResources = append(newResources, MCPResource{
URI: r.URI,
Name: r.Name,
Description: r.Description,
MIMEType: r.MIMEType,
ServerName: serverName,
})
}
m.mu.Lock()
// Remove old resources from this server, add new ones.
filtered := make([]MCPResource, 0, len(m.resources))
for _, r := range m.resources {
if r.ServerName != serverName {
filtered = append(filtered, r)
}
}
m.resources = append(filtered, newResources...)
m.mu.Unlock()
if m.onResourcesChanged != nil {
m.onResourcesChanged()
}
}
// GetLoadedServerNames returns the names of all successfully loaded MCP servers.
// This includes servers that are currently connected and have had their tools loaded,
// regardless of their current health status. Useful for debugging and status reporting.
@@ -290,6 +1064,9 @@ func (m *MCPToolManager) GetLoadedServerNames() []string {
// proper cleanup of stdio processes, network connections, and other resources.
// It is safe to call Close multiple times.
func (m *MCPToolManager) Close() error {
if m.connectionPool == nil {
return nil
}
return m.connectionPool.Close()
}
@@ -0,0 +1,323 @@
package tools
import (
"context"
"os"
"path/filepath"
"runtime"
"slices"
"strings"
"sync"
"testing"
"time"
"github.com/mark3labs/kit/internal/config"
)
// testdataDir returns the absolute path to the testdata directory.
func testdataDir(t *testing.T) string {
t.Helper()
_, file, _, ok := runtime.Caller(0)
if !ok {
t.Fatal("cannot determine test file path")
}
return filepath.Join(filepath.Dir(file), "testdata")
}
// echoServerConfig returns an MCPServerConfig for the test echo MCP server.
func echoServerConfig(t *testing.T) config.MCPServerConfig {
t.Helper()
script := filepath.Join(testdataDir(t), "echo_server.py")
if _, err := os.Stat(script); err != nil {
t.Skipf("echo_server.py not found: %v", err)
}
return config.MCPServerConfig{
Command: []string{"python3", script},
}
}
// TestMCPToolManager_AddServer_Integration tests adding a real MCP server
// at runtime and verifying tools are loaded.
func TestMCPToolManager_AddServer_Integration(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
manager := NewMCPToolManager()
defer func() { _ = manager.Close() }()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
// Track callbacks.
var mu sync.Mutex
var loadedServer string
var loadedCount int
toolsChangedCount := 0
manager.SetOnServerLoaded(func(name string, count int, err error) {
mu.Lock()
loadedServer = name
loadedCount = count
mu.Unlock()
})
manager.SetOnToolsChanged(func() {
mu.Lock()
toolsChangedCount++
mu.Unlock()
})
// Add the server.
count, err := manager.AddServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("AddServer failed: %v", err)
}
if count != 2 {
t.Errorf("Expected 2 tools from echo server, got %d", count)
}
// Verify callbacks fired.
mu.Lock()
if loadedServer != "echo" {
t.Errorf("Expected onServerLoaded for 'echo', got %q", loadedServer)
}
if loadedCount != 2 {
t.Errorf("Expected onServerLoaded count=2, got %d", loadedCount)
}
if toolsChangedCount != 1 {
t.Errorf("Expected onToolsChanged called once, got %d", toolsChangedCount)
}
mu.Unlock()
// Verify tools are accessible.
tools := manager.GetTools()
if len(tools) != 2 {
t.Fatalf("Expected 2 tools, got %d", len(tools))
}
// Verify tool names are prefixed.
toolNames := make(map[string]bool)
for _, tool := range tools {
toolNames[tool.Name] = true
}
if !toolNames["echo__echo"] {
t.Error("Expected tool 'echo__echo'")
}
if !toolNames["echo__greet"] {
t.Error("Expected tool 'echo__greet'")
}
// Verify server appears in loaded names.
names := manager.GetLoadedServerNames()
if !slices.Contains(names, "echo") {
t.Errorf("Expected 'echo' in loaded server names, got: %v", names)
}
}
// TestMCPToolManager_RemoveServer_Integration tests removing a real MCP server
// and verifying tools are cleaned up.
func TestMCPToolManager_RemoveServer_Integration(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
manager := NewMCPToolManager()
defer func() { _ = manager.Close() }()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
// Add the server first.
count, err := manager.AddServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("AddServer failed: %v", err)
}
if count != 2 {
t.Fatalf("Expected 2 tools, got %d", count)
}
var mu sync.Mutex
toolsChangedCount := 0
manager.SetOnToolsChanged(func() {
mu.Lock()
toolsChangedCount++
mu.Unlock()
})
// Remove the server.
err = manager.RemoveServer("echo")
if err != nil {
t.Fatalf("RemoveServer failed: %v", err)
}
// Verify tools are gone.
tools := manager.GetTools()
if len(tools) != 0 {
t.Errorf("Expected 0 tools after removal, got %d", len(tools))
}
// Verify callback fired.
mu.Lock()
if toolsChangedCount != 1 {
t.Errorf("Expected onToolsChanged called once, got %d", toolsChangedCount)
}
mu.Unlock()
// Verify server is gone from loaded names.
names := manager.GetLoadedServerNames()
for _, n := range names {
if n == "echo" {
t.Error("Server 'echo' should not appear in loaded names after removal")
}
}
// Removing again should error.
err = manager.RemoveServer("echo")
if err == nil {
t.Fatal("Expected error removing already-removed server")
}
if !strings.Contains(err.Error(), "not loaded") {
t.Errorf("Expected 'not loaded' error, got: %v", err)
}
}
// TestMCPToolManager_AddRemoveMultiple_Integration tests adding and removing
// multiple servers, verifying tool isolation.
func TestMCPToolManager_AddRemoveMultiple_Integration(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
manager := NewMCPToolManager()
defer func() { _ = manager.Close() }()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
// Add two servers with the same binary but different names.
count1, err := manager.AddServer(ctx, "server-a", cfg)
if err != nil {
t.Fatalf("AddServer server-a failed: %v", err)
}
count2, err := manager.AddServer(ctx, "server-b", cfg)
if err != nil {
t.Fatalf("AddServer server-b failed: %v", err)
}
totalTools := count1 + count2
if totalTools != 4 {
t.Fatalf("Expected 4 total tools (2+2), got %d", totalTools)
}
tools := manager.GetTools()
if len(tools) != 4 {
t.Fatalf("Expected 4 tools, got %d", len(tools))
}
// Remove server-a, verify server-b tools remain.
err = manager.RemoveServer("server-a")
if err != nil {
t.Fatalf("RemoveServer server-a failed: %v", err)
}
tools = manager.GetTools()
if len(tools) != 2 {
t.Fatalf("Expected 2 tools after removing server-a, got %d", len(tools))
}
// Remaining tools should all be from server-b.
for _, tool := range tools {
if !strings.HasPrefix(tool.Name, "server-b__") {
t.Errorf("Expected tool from server-b, got: %s", tool.Name)
}
}
// Remove server-b.
err = manager.RemoveServer("server-b")
if err != nil {
t.Fatalf("RemoveServer server-b failed: %v", err)
}
tools = manager.GetTools()
if len(tools) != 0 {
t.Errorf("Expected 0 tools after removing all servers, got %d", len(tools))
}
}
// TestMCPToolManager_AddServer_DuplicateDetection_Integration tests that
// adding a server with the same name as an already loaded server errors.
func TestMCPToolManager_AddServer_DuplicateDetection_Integration(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
manager := NewMCPToolManager()
defer func() { _ = manager.Close() }()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
// Add the server.
_, err := manager.AddServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("First AddServer failed: %v", err)
}
// Try to add again with the same name.
_, err = manager.AddServer(ctx, "echo", cfg)
if err == nil {
t.Fatal("Expected error adding duplicate server")
}
if !strings.Contains(err.Error(), "already loaded") {
t.Errorf("Expected 'already loaded' error, got: %v", err)
}
}
// TestMCPToolManager_AddAfterRemove_Integration tests that a server can be
// re-added after being removed.
func TestMCPToolManager_AddAfterRemove_Integration(t *testing.T) {
if testing.Short() {
t.Skip("skipping integration test in short mode")
}
manager := NewMCPToolManager()
defer func() { _ = manager.Close() }()
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
cfg := echoServerConfig(t)
// Add, remove, re-add.
_, err := manager.AddServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("First AddServer failed: %v", err)
}
err = manager.RemoveServer("echo")
if err != nil {
t.Fatalf("RemoveServer failed: %v", err)
}
count, err := manager.AddServer(ctx, "echo", cfg)
if err != nil {
t.Fatalf("Re-AddServer failed: %v", err)
}
if count != 2 {
t.Errorf("Expected 2 tools on re-add, got %d", count)
}
tools := manager.GetTools()
if len(tools) != 2 {
t.Errorf("Expected 2 tools after re-add, got %d", len(tools))
}
}
+155
View File
@@ -0,0 +1,155 @@
package tools
import (
"context"
"strings"
"sync"
"testing"
"time"
"github.com/mark3labs/kit/internal/config"
)
// TestMCPToolManager_AddServer_DuplicateName verifies that adding a server
// with a name that already exists returns an error.
func TestMCPToolManager_AddServer_DuplicateName(t *testing.T) {
manager := NewMCPToolManager()
cfg := config.MCPServerConfig{
Command: []string{"non-existent-command"},
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
// First add will fail (bad command), but let's test the duplicate detection
// by simulating a loaded server via LoadTools first.
loadCfg := &config.Config{
MCPServers: map[string]config.MCPServerConfig{
"test-server": cfg,
},
}
// This will fail to load but creates the connection pool.
_ = manager.LoadTools(ctx, loadCfg)
// Now try to add the same server name — the tools didn't load (bad command),
// so AddServer should not find a duplicate and should fail with connection error.
_, err := manager.AddServer(ctx, "test-server", cfg)
if err == nil {
t.Fatal("Expected error when adding server with bad command, got nil")
}
// It should be a connection error, not a duplicate error.
if strings.Contains(err.Error(), "already loaded") {
t.Fatalf("Should not report duplicate since server failed to load initially: %v", err)
}
}
// TestMCPToolManager_RemoveServer_NotLoaded verifies that removing a server
// that doesn't exist returns an appropriate error.
func TestMCPToolManager_RemoveServer_NotLoaded(t *testing.T) {
manager := NewMCPToolManager()
err := manager.RemoveServer("nonexistent")
if err == nil {
t.Fatal("Expected error when removing non-existent server, got nil")
}
if !strings.Contains(err.Error(), "not loaded") {
t.Errorf("Expected 'not loaded' error, got: %v", err)
}
}
// TestMCPToolManager_AddServer_CreatesConnectionPool verifies that AddServer
// lazily creates a connection pool when LoadTools was never called.
func TestMCPToolManager_AddServer_CreatesConnectionPool(t *testing.T) {
manager := NewMCPToolManager()
// Connection pool should be nil initially.
if manager.connectionPool != nil {
t.Fatal("Expected nil connection pool before any operation")
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
// AddServer with a bad command — should fail, but the pool should be created.
_, err := manager.AddServer(ctx, "lazy-server", config.MCPServerConfig{
Command: []string{"non-existent-command"},
})
if err == nil {
t.Fatal("Expected error for bad command")
}
// Connection pool should have been created.
if manager.connectionPool == nil {
t.Fatal("Expected connection pool to be created lazily by AddServer")
}
}
// TestMCPToolManager_OnToolsChanged_Callback verifies that the onToolsChanged
// callback fires on RemoveServer (we can't easily test AddServer with a real
// MCP server, but we can test the callback wiring).
func TestMCPToolManager_OnToolsChanged_Callback(t *testing.T) {
manager := NewMCPToolManager()
var mu sync.Mutex
callCount := 0
manager.SetOnToolsChanged(func() {
mu.Lock()
callCount++
mu.Unlock()
})
// RemoveServer on non-existent should NOT fire callback.
_ = manager.RemoveServer("nonexistent")
mu.Lock()
if callCount != 0 {
t.Errorf("Expected 0 callback calls for failed remove, got %d", callCount)
}
mu.Unlock()
}
// TestMCPToolManager_Close_NilPool verifies Close is safe when the connection
// pool was never initialized.
func TestMCPToolManager_Close_NilPool(t *testing.T) {
manager := NewMCPToolManager()
err := manager.Close()
if err != nil {
t.Fatalf("Expected nil error from Close with nil pool, got: %v", err)
}
}
// TestMCPConnectionPool_RemoveConnection_NotFound verifies that removing a
// non-existent connection returns an error.
func TestMCPConnectionPool_RemoveConnection_NotFound(t *testing.T) {
pool := NewMCPConnectionPool(DefaultConnectionPoolConfig(), false, nil, nil)
defer func() { _ = pool.Close() }()
err := pool.RemoveConnection("nonexistent")
if err == nil {
t.Fatal("Expected error for non-existent connection")
}
if !strings.Contains(err.Error(), "not found") {
t.Errorf("Expected 'not found' error, got: %v", err)
}
}
// TestMCPToolManager_EnsureConnectionPool_Idempotent verifies that
// ensureConnectionPool doesn't recreate an existing pool.
func TestMCPToolManager_EnsureConnectionPool_Idempotent(t *testing.T) {
manager := NewMCPToolManager()
// First call creates the pool.
manager.ensureConnectionPool()
pool1 := manager.connectionPool
if pool1 == nil {
t.Fatal("Expected pool to be created")
}
// Second call should be a no-op.
manager.ensureConnectionPool()
pool2 := manager.connectionPool
if pool1 != pool2 {
t.Fatal("Expected ensureConnectionPool to be idempotent")
}
}
+691
View File
@@ -0,0 +1,691 @@
package tools
import (
"context"
"encoding/base64"
"fmt"
"strings"
"testing"
mcpclient "github.com/mark3labs/mcp-go/client"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
// newTestPromptServer creates an in-process MCP server with prompt capabilities
// and the specified prompts + handlers. Returns an initialized MCPClient.
func newTestPromptServer(t *testing.T, prompts ...server.ServerPrompt) mcpclient.MCPClient {
t.Helper()
mcpServer := server.NewMCPServer(
"test-prompt-server", "1.0.0",
server.WithPromptCapabilities(true),
server.WithToolCapabilities(true),
)
if len(prompts) > 0 {
mcpServer.AddPrompts(prompts...)
}
// Add a dummy tool so loadServerTools has something to list.
mcpServer.AddTool(
mcp.NewTool("noop", mcp.WithDescription("no-op tool")),
func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
return mcp.NewToolResultText("ok"), nil
},
)
client, err := mcpclient.NewInProcessClient(mcpServer)
if err != nil {
t.Fatalf("NewInProcessClient: %v", err)
}
ctx := context.Background()
if err := client.Start(ctx); err != nil {
t.Fatalf("client.Start: %v", err)
}
initReq := mcp.InitializeRequest{}
initReq.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
initReq.Params.ClientInfo = mcp.Implementation{Name: "test", Version: "1.0"}
if _, err := client.Initialize(ctx, initReq); err != nil {
t.Fatalf("client.Initialize: %v", err)
}
t.Cleanup(func() { _ = client.Close() })
return client
}
// injectClientIntoManager sets up an MCPToolManager with a pre-connected
// in-process client, bypassing the normal connection pool flow.
func injectClientIntoManager(t *testing.T, serverName string, client mcpclient.MCPClient) *MCPToolManager {
t.Helper()
m := NewMCPToolManager()
// Create a minimal connection pool and inject our client.
pool := NewMCPConnectionPool(DefaultConnectionPoolConfig(), false, nil, nil)
pool.mu.Lock()
pool.connections[serverName] = &MCPConnection{
client: client,
serverName: serverName,
isHealthy: true,
}
pool.mu.Unlock()
m.connectionPool = pool
return m
}
func TestLoadServerPrompts_Basic(t *testing.T) {
ctx := context.Background()
client := newTestPromptServer(t,
server.ServerPrompt{
Prompt: mcp.NewPrompt("review-pr",
mcp.WithPromptDescription("Review a pull request"),
mcp.WithArgument("pr_number",
mcp.ArgumentDescription("The PR number to review"),
mcp.RequiredArgument(),
),
mcp.WithArgument("focus",
mcp.ArgumentDescription("Area to focus on"),
),
),
Handler: func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
prNum := req.Params.Arguments["pr_number"]
return &mcp.GetPromptResult{
Description: "PR review prompt",
Messages: []mcp.PromptMessage{
{
Role: mcp.RoleUser,
Content: mcp.TextContent{
Type: "text",
Text: fmt.Sprintf("Please review PR #%s", prNum),
},
},
},
}, nil
},
},
server.ServerPrompt{
Prompt: mcp.NewPrompt("explain-code",
mcp.WithPromptDescription("Explain a piece of code"),
),
Handler: func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
return &mcp.GetPromptResult{
Messages: []mcp.PromptMessage{
{
Role: mcp.RoleUser,
Content: mcp.TextContent{
Type: "text",
Text: "Please explain the following code.",
},
},
},
}, nil
},
},
)
m := injectClientIntoManager(t, "github", client)
conn := &MCPConnection{
client: client,
serverName: "github",
isHealthy: true,
}
m.loadServerPrompts(ctx, "github", conn)
prompts := m.GetPrompts()
if len(prompts) != 2 {
t.Fatalf("expected 2 prompts, got %d", len(prompts))
}
// Find review-pr prompt.
var reviewPR *MCPPrompt
for i := range prompts {
if prompts[i].Name == "review-pr" {
reviewPR = &prompts[i]
break
}
}
if reviewPR == nil {
t.Fatal("review-pr prompt not found")
}
if reviewPR.Description != "Review a pull request" {
t.Errorf("unexpected description: %q", reviewPR.Description)
}
if reviewPR.ServerName != "github" {
t.Errorf("unexpected server name: %q", reviewPR.ServerName)
}
if len(reviewPR.Arguments) != 2 {
t.Fatalf("expected 2 arguments, got %d", len(reviewPR.Arguments))
}
// Verify argument metadata.
arg0 := reviewPR.Arguments[0]
if arg0.Name != "pr_number" {
t.Errorf("expected first arg name 'pr_number', got %q", arg0.Name)
}
if !arg0.Required {
t.Error("expected first arg to be required")
}
arg1 := reviewPR.Arguments[1]
if arg1.Name != "focus" {
t.Errorf("expected second arg name 'focus', got %q", arg1.Name)
}
if arg1.Required {
t.Error("expected second arg to be optional")
}
}
func TestGetPrompt_ExpandsWithArgs(t *testing.T) {
ctx := context.Background()
client := newTestPromptServer(t,
server.ServerPrompt{
Prompt: mcp.NewPrompt("greet",
mcp.WithPromptDescription("Greet someone"),
mcp.WithArgument("name", mcp.RequiredArgument()),
),
Handler: func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
name := req.Params.Arguments["name"]
return &mcp.GetPromptResult{
Description: "Greeting",
Messages: []mcp.PromptMessage{
{
Role: mcp.RoleUser,
Content: mcp.TextContent{
Type: "text",
Text: fmt.Sprintf("Hello, %s!", name),
},
},
},
}, nil
},
},
)
m := injectClientIntoManager(t, "myserver", client)
result, err := m.GetPrompt(ctx, "myserver", "greet", map[string]string{"name": "World"})
if err != nil {
t.Fatalf("GetPrompt error: %v", err)
}
if result.Description != "Greeting" {
t.Errorf("unexpected description: %q", result.Description)
}
if len(result.Messages) != 1 {
t.Fatalf("expected 1 message, got %d", len(result.Messages))
}
if result.Messages[0].Role != "user" {
t.Errorf("unexpected role: %q", result.Messages[0].Role)
}
if result.Messages[0].Content != "Hello, World!" {
t.Errorf("unexpected content: %q", result.Messages[0].Content)
}
}
func TestGetPrompt_MultipleMessages(t *testing.T) {
ctx := context.Background()
client := newTestPromptServer(t,
server.ServerPrompt{
Prompt: mcp.NewPrompt("chat-starter"),
Handler: func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
return &mcp.GetPromptResult{
Messages: []mcp.PromptMessage{
{
Role: mcp.RoleUser,
Content: mcp.TextContent{Type: "text", Text: "What is Go?"},
},
{
Role: mcp.RoleAssistant,
Content: mcp.TextContent{Type: "text", Text: "Go is a programming language."},
},
{
Role: mcp.RoleUser,
Content: mcp.TextContent{Type: "text", Text: "Tell me more."},
},
},
}, nil
},
},
)
m := injectClientIntoManager(t, "server", client)
result, err := m.GetPrompt(ctx, "server", "chat-starter", nil)
if err != nil {
t.Fatalf("GetPrompt error: %v", err)
}
if len(result.Messages) != 3 {
t.Fatalf("expected 3 messages, got %d", len(result.Messages))
}
if result.Messages[0].Role != "user" {
t.Errorf("msg[0] role: got %q, want 'user'", result.Messages[0].Role)
}
if result.Messages[1].Role != "assistant" {
t.Errorf("msg[1] role: got %q, want 'assistant'", result.Messages[1].Role)
}
if result.Messages[2].Content != "Tell me more." {
t.Errorf("msg[2] content: got %q, want 'Tell me more.'", result.Messages[2].Content)
}
}
func TestGetPrompt_ServerNotFound(t *testing.T) {
m := NewMCPToolManager()
pool := NewMCPConnectionPool(DefaultConnectionPoolConfig(), false, nil, nil)
m.connectionPool = pool
_, err := m.GetPrompt(context.Background(), "nonexistent", "foo", nil)
if err == nil {
t.Fatal("expected error for nonexistent server")
}
}
func TestGetPrompt_NoPool(t *testing.T) {
m := NewMCPToolManager()
_, err := m.GetPrompt(context.Background(), "any", "foo", nil)
if err == nil {
t.Fatal("expected error with no pool")
}
}
func TestRemoveServer_RemovesPrompts(t *testing.T) {
ctx := context.Background()
client := newTestPromptServer(t,
server.ServerPrompt{
Prompt: mcp.NewPrompt("my-prompt",
mcp.WithPromptDescription("A test prompt"),
),
Handler: func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
return &mcp.GetPromptResult{
Messages: []mcp.PromptMessage{
{Role: mcp.RoleUser, Content: mcp.TextContent{Type: "text", Text: "hi"}},
},
}, nil
},
},
)
m := injectClientIntoManager(t, "testsvr", client)
// Manually populate tools and prompts as loadServerTools would.
conn := m.connectionPool.connections["testsvr"]
m.loadServerPrompts(ctx, "testsvr", conn)
// Also add a fake tool mapping so RemoveServer finds the server.
m.toolMap["testsvr__noop"] = &toolMapping{
serverName: "testsvr",
originalName: "noop",
}
m.tools = append(m.tools, MCPTool{
Name: "testsvr__noop",
ServerName: "testsvr",
})
// Verify prompts exist before removal.
if got := len(m.GetPrompts()); got != 1 {
t.Fatalf("expected 1 prompt before removal, got %d", got)
}
// Remove the server.
err := m.RemoveServer("testsvr")
if err != nil {
t.Fatalf("RemoveServer error: %v", err)
}
// Verify prompts are gone.
if got := len(m.GetPrompts()); got != 0 {
t.Fatalf("expected 0 prompts after removal, got %d", got)
}
}
func TestLoadServerPrompts_NoPromptCapability(t *testing.T) {
// Server without prompt capabilities — ListPrompts should fail gracefully.
mcpServer := server.NewMCPServer("no-prompts", "1.0.0",
server.WithToolCapabilities(true),
// No WithPromptCapabilities
)
mcpServer.AddTool(
mcp.NewTool("noop"),
func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
return mcp.NewToolResultText("ok"), nil
},
)
client, err := mcpclient.NewInProcessClient(mcpServer)
if err != nil {
t.Fatalf("NewInProcessClient: %v", err)
}
ctx := context.Background()
_ = client.Start(ctx)
initReq := mcp.InitializeRequest{}
initReq.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
initReq.Params.ClientInfo = mcp.Implementation{Name: "test", Version: "1.0"}
_, _ = client.Initialize(ctx, initReq)
t.Cleanup(func() { _ = client.Close() })
m := NewMCPToolManager()
conn := &MCPConnection{
client: client,
serverName: "no-prompts",
isHealthy: true,
}
// Should not panic or error — just silently skip.
m.loadServerPrompts(ctx, "no-prompts", conn)
if got := len(m.GetPrompts()); got != 0 {
t.Fatalf("expected 0 prompts from server without prompt capability, got %d", got)
}
}
func TestExtractPromptContent(t *testing.T) {
t.Run("TextContent", func(t *testing.T) {
text, parts := extractPromptContent(mcp.TextContent{Type: "text", Text: "hello world"})
if text != "hello world" {
t.Errorf("text = %q, want %q", text, "hello world")
}
if len(parts) != 0 {
t.Errorf("expected 0 file parts, got %d", len(parts))
}
})
t.Run("ImageContent", func(t *testing.T) {
// base64 of "fake image"
encoded := base64.StdEncoding.EncodeToString([]byte("fake image"))
text, parts := extractPromptContent(mcp.ImageContent{
Type: "image",
Data: encoded,
MIMEType: "image/png",
})
if text != "" {
t.Errorf("expected empty text, got %q", text)
}
if len(parts) != 1 {
t.Fatalf("expected 1 file part, got %d", len(parts))
}
if parts[0].MediaType != "image/png" {
t.Errorf("media type = %q, want %q", parts[0].MediaType, "image/png")
}
if parts[0].Filename != "image.png" {
t.Errorf("filename = %q, want %q", parts[0].Filename, "image.png")
}
if string(parts[0].Data) != "fake image" {
t.Errorf("data = %q, want %q", string(parts[0].Data), "fake image")
}
})
t.Run("ImageContent_DefaultMIME", func(t *testing.T) {
encoded := base64.StdEncoding.EncodeToString([]byte("img"))
_, parts := extractPromptContent(mcp.ImageContent{
Type: "image",
Data: encoded,
// no MIMEType → should default to image/png
})
if len(parts) != 1 {
t.Fatalf("expected 1 file part, got %d", len(parts))
}
if parts[0].MediaType != "image/png" {
t.Errorf("default MIME = %q, want %q", parts[0].MediaType, "image/png")
}
})
t.Run("AudioContent", func(t *testing.T) {
encoded := base64.StdEncoding.EncodeToString([]byte("fake audio"))
text, parts := extractPromptContent(mcp.AudioContent{
Type: "audio",
Data: encoded,
MIMEType: "audio/mp3",
})
if text != "" {
t.Errorf("expected empty text, got %q", text)
}
if len(parts) != 1 {
t.Fatalf("expected 1 file part, got %d", len(parts))
}
if parts[0].MediaType != "audio/mp3" {
t.Errorf("media type = %q, want %q", parts[0].MediaType, "audio/mp3")
}
if parts[0].Filename != "audio.wav" {
t.Errorf("filename = %q, want %q", parts[0].Filename, "audio.wav")
}
})
t.Run("EmbeddedResource_Text", func(t *testing.T) {
text, parts := extractPromptContent(mcp.EmbeddedResource{
Type: "resource",
Resource: mcp.TextResourceContents{
URI: "file:///project/main.go",
MIMEType: "text/x-go",
Text: "package main",
},
})
if text == "" {
t.Fatal("expected non-empty text for text resource")
}
if !strings.Contains(text, "package main") {
t.Errorf("text should contain resource content, got %q", text)
}
if !strings.Contains(text, "file:///project/main.go") {
t.Errorf("text should contain URI, got %q", text)
}
if len(parts) != 0 {
t.Errorf("expected 0 file parts for text resource, got %d", len(parts))
}
})
t.Run("EmbeddedResource_Blob", func(t *testing.T) {
blobData := []byte("binary content")
encoded := base64.StdEncoding.EncodeToString(blobData)
text, parts := extractPromptContent(mcp.EmbeddedResource{
Type: "resource",
Resource: mcp.BlobResourceContents{
URI: "file:///project/data.bin",
MIMEType: "application/octet-stream",
Blob: encoded,
},
})
if text != "" {
t.Errorf("expected empty text for blob resource, got %q", text)
}
if len(parts) != 1 {
t.Fatalf("expected 1 file part for blob resource, got %d", len(parts))
}
if parts[0].Filename != "data.bin" {
t.Errorf("filename = %q, want %q", parts[0].Filename, "data.bin")
}
if parts[0].MediaType != "application/octet-stream" {
t.Errorf("media type = %q, want %q", parts[0].MediaType, "application/octet-stream")
}
if string(parts[0].Data) != "binary content" {
t.Errorf("data = %q, want %q", string(parts[0].Data), "binary content")
}
})
t.Run("ResourceLink", func(t *testing.T) {
text, parts := extractPromptContent(mcp.ResourceLink{
Type: "resource_link",
URI: "file:///docs/readme.md",
Name: "readme.md",
})
if text == "" {
t.Fatal("expected non-empty text for resource link")
}
if !strings.Contains(text, "file:///docs/readme.md") {
t.Errorf("text should contain URI, got %q", text)
}
if !strings.Contains(text, "readme.md") {
t.Errorf("text should contain name, got %q", text)
}
if len(parts) != 0 {
t.Errorf("expected 0 file parts for resource link, got %d", len(parts))
}
})
t.Run("InvalidBase64", func(t *testing.T) {
_, parts := extractPromptContent(mcp.ImageContent{
Type: "image",
Data: "not-valid-base64!!!",
MIMEType: "image/png",
})
if len(parts) != 0 {
t.Errorf("expected 0 file parts for invalid base64, got %d", len(parts))
}
})
t.Run("NilContent", func(t *testing.T) {
text, parts := extractPromptContent((*mcp.TextContent)(nil))
if text != "" {
t.Errorf("expected empty text for nil, got %q", text)
}
if len(parts) != 0 {
t.Errorf("expected 0 parts for nil, got %d", len(parts))
}
})
}
func TestFilenameFromURI(t *testing.T) {
tests := []struct {
uri string
want string
}{
{"file:///path/to/image.png", "image.png"},
{"file:///single.txt", "single.txt"},
{"resource://server/data.json", "data.json"},
{"nopath", "nopath"},
{"", "resource"},
}
for _, tt := range tests {
t.Run(tt.uri, func(t *testing.T) {
got := filenameFromURI(tt.uri)
if got != tt.want {
t.Errorf("filenameFromURI(%q) = %q, want %q", tt.uri, got, tt.want)
}
})
}
}
func TestGetPrompt_EmbeddedResources(t *testing.T) {
ctx := context.Background()
imgData := base64.StdEncoding.EncodeToString([]byte("fake-png"))
blobData := base64.StdEncoding.EncodeToString([]byte("binary-blob"))
client := newTestPromptServer(t,
server.ServerPrompt{
Prompt: mcp.NewPrompt("review-with-files",
mcp.WithPromptDescription("Review with embedded resources"),
),
Handler: func(ctx context.Context, req mcp.GetPromptRequest) (*mcp.GetPromptResult, error) {
return &mcp.GetPromptResult{
Description: "Review prompt with embedded files",
Messages: []mcp.PromptMessage{
{
Role: mcp.RoleUser,
Content: mcp.TextContent{Type: "text", Text: "Please review these files:"},
},
{
Role: mcp.RoleUser,
Content: mcp.EmbeddedResource{
Type: "resource",
Resource: mcp.TextResourceContents{
URI: "file:///src/main.go",
MIMEType: "text/x-go",
Text: "package main\n\nfunc main() {}",
},
},
},
{
Role: mcp.RoleUser,
Content: mcp.ImageContent{
Type: "image",
Data: imgData,
MIMEType: "image/png",
},
},
{
Role: mcp.RoleUser,
Content: mcp.EmbeddedResource{
Type: "resource",
Resource: mcp.BlobResourceContents{
URI: "file:///data/model.bin",
MIMEType: "application/octet-stream",
Blob: blobData,
},
},
},
},
}, nil
},
},
)
m := injectClientIntoManager(t, "test", client)
result, err := m.GetPrompt(ctx, "test", "review-with-files", nil)
if err != nil {
t.Fatalf("GetPrompt error: %v", err)
}
if result.Description != "Review prompt with embedded files" {
t.Errorf("unexpected description: %q", result.Description)
}
// Should have 4 messages: text, embedded text resource, image, embedded blob
if len(result.Messages) != 4 {
t.Fatalf("expected 4 messages, got %d", len(result.Messages))
}
// Message 0: plain text
msg0 := result.Messages[0]
if msg0.Content != "Please review these files:" {
t.Errorf("msg[0] content = %q", msg0.Content)
}
if len(msg0.FileParts) != 0 {
t.Errorf("msg[0] expected 0 file parts, got %d", len(msg0.FileParts))
}
// Message 1: embedded text resource → inlined as text
msg1 := result.Messages[1]
if !strings.Contains(msg1.Content, "package main") {
t.Errorf("msg[1] should contain resource text, got %q", msg1.Content)
}
if len(msg1.FileParts) != 0 {
t.Errorf("msg[1] expected 0 file parts (text resource), got %d", len(msg1.FileParts))
}
// Message 2: image → file part
msg2 := result.Messages[2]
if msg2.Content != "" {
t.Errorf("msg[2] expected empty text for image, got %q", msg2.Content)
}
if len(msg2.FileParts) != 1 {
t.Fatalf("msg[2] expected 1 file part, got %d", len(msg2.FileParts))
}
if msg2.FileParts[0].MediaType != "image/png" {
t.Errorf("msg[2] file part MIME = %q", msg2.FileParts[0].MediaType)
}
if string(msg2.FileParts[0].Data) != "fake-png" {
t.Errorf("msg[2] file part data = %q", string(msg2.FileParts[0].Data))
}
// Message 3: embedded blob resource → file part
msg3 := result.Messages[3]
if msg3.Content != "" {
t.Errorf("msg[3] expected empty text for blob resource, got %q", msg3.Content)
}
if len(msg3.FileParts) != 1 {
t.Fatalf("msg[3] expected 1 file part, got %d", len(msg3.FileParts))
}
if msg3.FileParts[0].Filename != "model.bin" {
t.Errorf("msg[3] filename = %q, want %q", msg3.FileParts[0].Filename, "model.bin")
}
if string(msg3.FileParts[0].Data) != "binary-blob" {
t.Errorf("msg[3] file part data = %q", string(msg3.FileParts[0].Data))
}
}
+2 -4
View File
@@ -103,14 +103,12 @@ func TestMCPToolManager_EmptyConfig(t *testing.T) {
// Test that we can get tool info for each tool
for _, tool := range tools {
info := tool.Info()
// Check that the tool has a valid name
if info.Name == "" {
if tool.Name == "" {
t.Error("Tool has empty name")
}
t.Logf("Tool: %s, Description: %s", info.Name, info.Description)
t.Logf("Tool: %s, Description: %s", tool.Name, tool.Description)
}
}
+7
View File
@@ -6,6 +6,7 @@ import (
"net/url"
"github.com/mark3labs/mcp-go/client"
"github.com/mark3labs/mcp-go/client/transport"
)
// MCPAuthHandler is the internal interface for handling MCP OAuth flows.
@@ -21,6 +22,12 @@ type MCPAuthHandler interface {
HandleAuth(ctx context.Context, serverName string, authURL string) (callbackURL string, err error)
}
// TokenStoreFactory creates a transport.TokenStore for a given MCP server URL.
// When provided to the connection pool, it is called once per remote MCP server
// instead of using the default file-based token store. Implementations can
// return any transport.TokenStore — in-memory, database-backed, encrypted, etc.
type TokenStoreFactory func(serverURL string) (transport.TokenStore, error)
// OAuthFlowRunner handles the OAuth authorization flow when an MCP server
// returns an OAuthAuthorizationRequiredError. It coordinates dynamic client
// registration, PKCE generation, user authorization (via MCPAuthHandler),
+111
View File
@@ -0,0 +1,111 @@
#!/usr/bin/env python3
"""Minimal MCP server over stdio for testing. Exposes one tool: echo."""
import json
import sys
def read_message():
"""Read a JSON-RPC message from stdin."""
line = sys.stdin.readline()
if not line:
return None
return json.loads(line.strip())
def write_message(msg):
"""Write a JSON-RPC message to stdout."""
sys.stdout.write(json.dumps(msg) + "\n")
sys.stdout.flush()
def handle(msg):
method = msg.get("method", "")
mid = msg.get("id")
if method == "initialize":
write_message({
"jsonrpc": "2.0",
"id": mid,
"result": {
"protocolVersion": "2024-11-05",
"capabilities": {"tools": {}},
"serverInfo": {"name": "test-echo", "version": "1.0.0"},
},
})
elif method == "notifications/initialized":
pass # no response needed
elif method == "tools/list":
write_message({
"jsonrpc": "2.0",
"id": mid,
"result": {
"tools": [
{
"name": "echo",
"description": "Echoes the input text back.",
"inputSchema": {
"type": "object",
"properties": {
"text": {"type": "string", "description": "Text to echo"}
},
"required": ["text"],
},
},
{
"name": "greet",
"description": "Returns a greeting.",
"inputSchema": {
"type": "object",
"properties": {
"name": {"type": "string", "description": "Name to greet"}
},
"required": ["name"],
},
},
]
},
})
elif method == "tools/call":
tool_name = msg["params"]["name"]
args = msg["params"].get("arguments", {})
if tool_name == "echo":
text = args.get("text", "")
write_message({
"jsonrpc": "2.0",
"id": mid,
"result": {
"content": [{"type": "text", "text": text}]
},
})
elif tool_name == "greet":
name = args.get("name", "World")
write_message({
"jsonrpc": "2.0",
"id": mid,
"result": {
"content": [{"type": "text", "text": f"Hello, {name}!"}]
},
})
else:
write_message({
"jsonrpc": "2.0",
"id": mid,
"error": {"code": -32601, "message": f"Unknown tool: {tool_name}"},
})
elif method == "ping":
write_message({"jsonrpc": "2.0", "id": mid, "result": {}})
else:
if mid is not None:
write_message({
"jsonrpc": "2.0",
"id": mid,
"error": {"code": -32601, "message": f"Unknown method: {method}"},
})
if __name__ == "__main__":
while True:
msg = read_message()
if msg is None:
break
handle(msg)
+1 -25
View File
@@ -19,7 +19,7 @@ import (
// newTestInput creates an InputComponent with the given AppController (may be nil).
func newTestInput(ctrl AppController) *InputComponent {
return NewInputComponent(80, "test input", ctrl)
return NewInputComponent(80, ctrl)
}
// sendInputMsg calls component.Update with the given message, returns the
@@ -69,30 +69,6 @@ func TestInputComponent_SubmitEmitsSubmitMsg(t *testing.T) {
}
}
// TestInputComponent_CtrlD_SubmitEmitsSubmitMsg verifies that ctrl+d also
// submits the text.
func TestInputComponent_CtrlD_SubmitEmitsSubmitMsg(t *testing.T) {
ctrl := &stubAppController{}
c := newTestInput(ctrl)
c.textarea.SetValue("ctrl+d submit")
c.lastValue = "ctrl+d submit"
_, cmd := sendInputMsg(c, tea.KeyPressMsg{Code: 'd', Mod: tea.ModCtrl})
msg := runCmd(cmd)
if msg == nil {
t.Fatal("expected a cmd from ctrl+d on non-empty input")
}
sm, ok := msg.(core.SubmitMsg)
if !ok {
t.Fatalf("expected submitMsg from ctrl+d, got %T", msg)
}
if sm.Text != "ctrl+d submit" {
t.Fatalf("expected Text='ctrl+d submit', got %q", sm.Text)
}
}
// TestInputComponent_EmptySubmit_NoCmd verifies that submitting an empty or
// whitespace-only string produces no cmd.
func TestInputComponent_EmptySubmit_NoCmd(t *testing.T) {
+2 -1
View File
@@ -20,6 +20,7 @@ type SlashCommand struct {
Aliases []string
Category string // e.g., "Navigation", "System", "Info"
Complete func(prefix string) []string // optional argument tab-completion
HasArgs bool // true when the command expects arguments (e.g. prompt templates with placeholders)
}
// SlashCommands provides the global registry of all available slash commands
@@ -83,7 +84,7 @@ var SlashCommands = []SlashCommand{
},
{
Name: "/thinking",
Description: "Set thinking/reasoning level (off, minimal, low, medium, high)",
Description: "Set thinking/reasoning level (off, none, minimal, low, medium, high)",
Category: "System",
Aliases: []string{"/think"},
Complete: func(prefix string) []string {
+5
View File
@@ -25,6 +25,11 @@ type SubmitMsg struct {
// presses ESC a second time, the canceling state is reset to false.
type CancelTimerExpiredMsg struct{}
// CtrlCResetMsg is sent after a short delay when the user presses Ctrl+C to
// clear input. If the user doesn't press Ctrl+C again within the timeout,
// the ctrlCPressedOnce flag is reset so the next Ctrl+C will clear again.
type CtrlCResetMsg struct{}
// --- Tree session events ---
// TreeNodeSelectedMsg is sent when the user selects a node in the tree selector.
+3 -1
View File
@@ -139,7 +139,9 @@ func (h *CLIEventHandler) Handle(msg tea.Msg) {
case "block":
h.cli.DisplayExtensionBlock(e.Text, e.BorderColor, e.Subtitle)
default:
fmt.Println(e.Text)
// Route unstyled extension prints through the system message
// renderer so they get consistent formatting and timestamps.
h.cli.DisplayInfo(e.Text)
}
case app.StepCompleteEvent:
+8 -1
View File
@@ -29,9 +29,16 @@ type (
ExtensionCommand = commands.ExtensionCommand
)
// Re-export functions from fileutil package
// Re-export functions and types from fileutil package
var ProcessFileAttachments = fileutil.ProcessFileAttachments
// Re-export types from fileutil
type (
FileAttachmentResult = fileutil.FileAttachmentResult
FilePart = fileutil.FilePart
MCPResourceReader = fileutil.MCPResourceReader
)
// Re-export from prefs package
var (
LoadThemePreference = prefs.LoadThemePreference
+1 -3
View File
@@ -109,9 +109,7 @@ func SetupCLI(opts *CLISetupOptions) (*CLI, error) {
}
}
fmt.Println("")
// Display model info
// Display model info (the system message block provides its own spacing).
if provider != "unknown" && model != "unknown" {
cli.DisplayInfo(fmt.Sprintf("Model loaded: %s (%s)", provider, model))
}
+12 -3
View File
@@ -8,15 +8,24 @@ import (
"strings"
)
// FileSuggestion represents a single file or directory suggestion for the @
// autocomplete popup.
// FileSuggestion represents a single file, directory, or MCP resource
// suggestion for the @ autocomplete popup.
type FileSuggestion struct {
// RelPath is the path relative to the search base (e.g. "cmd/kit/main.go").
// RelPath is the path relative to the search base (e.g. "cmd/kit/main.go")
// or a display name for MCP resources (e.g. "mcp:server/resource-name").
RelPath string
// IsDir is true when the entry is a directory.
IsDir bool
// Score is the fuzzy match score (higher is better).
Score int
// IsMCPResource is true for MCP resource entries.
IsMCPResource bool
// MCPServerName is the MCP server name (set when IsMCPResource is true).
MCPServerName string
// MCPResourceURI is the MCP resource URI (set when IsMCPResource is true).
MCPResourceURI string
// MCPMIMEType is the MIME type hint from the MCP server.
MCPMIMEType string
}
// maxFileSuggestions is the maximum number of file suggestions returned.
+202 -10
View File
@@ -2,29 +2,85 @@ package fileutil
import (
"fmt"
"mime"
"net/http"
"os"
"path/filepath"
"regexp"
"strings"
"github.com/mark3labs/kit/internal/fences"
)
// FilePart represents a binary file attachment (image, audio, etc.) extracted
// from an @file reference. Callers convert this to kit.LLMFilePart before
// sending to the LLM. Defined here to avoid a circular dependency on pkg/kit.
type FilePart struct {
// Filename is the basename of the file (e.g. "photo.png").
Filename string
// Data is the raw file bytes.
Data []byte
// MediaType is the MIME type (e.g. "image/png", "audio/wav").
MediaType string
}
// MCPResourceReader is a callback function that reads an MCP resource by
// server name and URI. Returns text content, binary data, MIME type, and error.
// Used by ProcessFileAttachments to resolve @mcp:server:uri tokens.
type MCPResourceReader func(serverName, uri string) (text string, blobData []byte, mimeType string, isBlob bool, err error)
// FileAttachmentResult is the result of processing @file references in user
// input. Text files are inlined as XML in ProcessedText; binary files (images,
// audio, video, PDFs) are returned as FileParts for multimodal submission.
type FileAttachmentResult struct {
// ProcessedText is the user's text with @file tokens replaced:
// text files become XML-wrapped content, binary file tokens are removed.
ProcessedText string
// FileParts contains binary file attachments extracted from @file
// references. Empty when all referenced files are text.
FileParts []FilePart
}
// fileTokenPattern matches @file references in user text. Supports:
// - @"path with spaces.txt" (quoted)
// - @path/to/file.txt (unquoted, no spaces)
var fileTokenPattern = regexp.MustCompile(`@"[^"]+"|@[^\s]+`)
// ProcessFileAttachments scans the user's input text for @file references,
// reads each referenced file, and returns the text with @tokens replaced by
// XML-wrapped file content. Non-file @ tokens (like email addresses) are left
// unchanged.
// reads each referenced file, and returns a result containing the processed
// text and any binary file attachments. Text files are XML-wrapped inline;
// binary files (images, audio, etc.) are extracted as FileParts for multimodal
// submission. Non-file @ tokens (like email addresses) are left unchanged.
//
// Returns the original text unchanged if no valid @file references are found.
func ProcessFileAttachments(text string, cwd string) string {
// MCP resources are supported via @mcp:server:uri tokens. The optional
// mcpReader callback is used to resolve them; pass nil to skip MCP resources.
func ProcessFileAttachments(text string, cwd string, mcpReader ...MCPResourceReader) FileAttachmentResult {
var reader MCPResourceReader
if len(mcpReader) > 0 {
reader = mcpReader[0]
}
var allParts []FilePart
processed := fences.ReplaceOutside(text, func(segment string) string {
result, parts := processFileTokens(segment, cwd, reader)
allParts = append(allParts, parts...)
return result
})
return FileAttachmentResult{
ProcessedText: processed,
FileParts: allParts,
}
}
// processFileTokens handles @file replacement in a single text segment
// that is known to be outside fenced code blocks. Returns the processed
// text and any binary file parts extracted.
func processFileTokens(text string, cwd string, mcpReader MCPResourceReader) (string, []FilePart) {
tokens := fileTokenPattern.FindAllString(text, -1)
if len(tokens) == 0 {
return text
return text, nil
}
var parts []FilePart
result := text
for _, token := range tokens {
path := tokenToPath(token)
@@ -32,6 +88,43 @@ func ProcessFileAttachments(text string, cwd string) string {
continue
}
// Check for MCP resource reference: @mcp:server:uri
if strings.HasPrefix(path, "mcp:") {
if mcpReader == nil {
continue
}
mcpRef := path[4:] // strip "mcp:"
// Split into server:uri (first colon separates server from URI)
serverName, uri, ok := strings.Cut(mcpRef, ":")
if !ok || serverName == "" || uri == "" {
continue // invalid format
}
textContent, blobData, mimeType, isBlob, err := mcpReader(serverName, uri)
if err != nil {
continue // skip on error, leave token as-is
}
if isBlob {
// Binary MCP resource → extract as FilePart.
filename := filepath.Base(uri)
if filename == "." || filename == "/" {
filename = serverName + "_resource"
}
parts = append(parts, FilePart{
Filename: filename,
Data: blobData,
MediaType: mimeType,
})
result = strings.Replace(result, token, "", 1)
} else {
// Text MCP resource → inline as XML.
wrapped := fmt.Sprintf("<resource uri=\"%s\" server=\"%s\">\n%s\n</resource>", uri, serverName, textContent)
result = strings.Replace(result, token, wrapped, 1)
}
continue
}
absPath, err := resolvePath(path, cwd)
if err != nil {
// Not a valid file reference — leave the token as-is.
@@ -59,12 +152,28 @@ func ProcessFileAttachments(text string, cwd string) string {
continue
}
// Build the XML-wrapped replacement.
wrapped := wrapFileContent(absPath, content)
result = strings.Replace(result, token, wrapped, 1)
mediaType := detectMediaType(absPath, content)
if isBinaryMediaType(mediaType) {
// Binary file → extract as a FilePart for multimodal submission.
// Remove the @token from the text.
parts = append(parts, FilePart{
Filename: filepath.Base(absPath),
Data: content,
MediaType: mediaType,
})
result = strings.Replace(result, token, "", 1)
} else {
// Text file → inline as XML-wrapped content.
wrapped := wrapFileContent(absPath, content)
result = strings.Replace(result, token, wrapped, 1)
}
}
return result
// Clean up any extra whitespace left by removed binary tokens.
result = strings.TrimSpace(result)
return result, parts
}
// tokenToPath strips the @ prefix and optional quotes from a token,
@@ -127,3 +236,86 @@ func resolvePath(path string, cwd string) (string, error) {
func wrapFileContent(absPath string, content []byte) string {
return fmt.Sprintf("<file path=\"%s\">\n%s\n</file>", absPath, string(content))
}
// detectMediaType determines the MIME type of a file using extension-based
// lookup first (more reliable for known types), then falls back to content
// sniffing via net/http.DetectContentType.
func detectMediaType(path string, content []byte) string {
// Extension-based detection is more reliable for well-known types.
ext := strings.ToLower(filepath.Ext(path))
if mt := mime.TypeByExtension(ext); mt != "" {
// mime.TypeByExtension returns types like "image/png; charset=utf-8"
// — strip parameters.
if base, _, ok := strings.Cut(mt, ";"); ok {
return strings.TrimSpace(base)
}
return mt
}
// Known extensions that mime package may miss.
switch ext {
case ".webp":
return "image/webp"
case ".avif":
return "image/avif"
case ".heic", ".heif":
return "image/heif"
case ".opus":
return "audio/opus"
case ".flac":
return "audio/flac"
case ".m4a":
return "audio/mp4"
case ".wasm":
return "application/wasm"
}
// Content sniffing fallback.
if len(content) > 0 {
detected := http.DetectContentType(content)
if detected != "" && detected != "application/octet-stream" {
if base, _, ok := strings.Cut(detected, ";"); ok {
return strings.TrimSpace(base)
}
return detected
}
}
// Default: treat as plain text so it gets XML-wrapped.
return "text/plain"
}
// isBinaryMediaType returns true if the MIME type represents a binary file
// that should be sent as a multimodal FilePart rather than XML-wrapped text.
func isBinaryMediaType(mediaType string) bool {
// Image types — always binary.
if strings.HasPrefix(mediaType, "image/") {
return true
}
// Audio types — always binary.
if strings.HasPrefix(mediaType, "audio/") {
return true
}
// Video types — always binary.
if strings.HasPrefix(mediaType, "video/") {
return true
}
// Specific application types that are binary.
switch mediaType {
case "application/pdf",
"application/zip",
"application/gzip",
"application/x-tar",
"application/octet-stream",
"application/wasm",
"application/x-executable",
"application/vnd.ms-excel",
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
"application/vnd.ms-powerpoint",
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
"application/msword",
"application/vnd.openxmlformats-officedocument.wordprocessingml.document":
return true
}
return false
}
+209
View File
@@ -0,0 +1,209 @@
package fileutil
import (
"os"
"path/filepath"
"testing"
)
func TestProcessFileAttachments_TextFile(t *testing.T) {
// Create a temp text file
dir := t.TempDir()
textFile := filepath.Join(dir, "hello.txt")
if err := os.WriteFile(textFile, []byte("hello world"), 0644); err != nil {
t.Fatal(err)
}
text := "@" + textFile + " check this out"
result := ProcessFileAttachments(text, dir)
if len(result.FileParts) != 0 {
t.Errorf("expected 0 FileParts for text file, got %d", len(result.FileParts))
}
if result.ProcessedText == text {
t.Error("expected text file to be XML-wrapped, but got original text unchanged")
}
// Should contain XML wrapping
if !contains(result.ProcessedText, "<file path=") {
t.Error("expected XML <file> wrapping in processed text")
}
if !contains(result.ProcessedText, "hello world") {
t.Error("expected file content in processed text")
}
}
func TestProcessFileAttachments_BinaryFile(t *testing.T) {
// Create a minimal PNG file (binary)
dir := t.TempDir()
pngFile := filepath.Join(dir, "image.png")
// Minimal valid PNG header
pngData := []byte{
0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, // PNG signature
0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52, // IHDR chunk
0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, // 1x1
0x08, 0x02, 0x00, 0x00, 0x00, 0x90, 0x77, 0x53, 0xDE, // 8bit RGB
0x00, 0x00, 0x00, 0x0C, 0x49, 0x44, 0x41, 0x54, // IDAT chunk
0x08, 0xD7, 0x63, 0xF8, 0xCF, 0xC0, 0x00, 0x00,
0x00, 0x02, 0x00, 0x01, 0xE2, 0x21, 0xBC, 0x33,
0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4E, 0x44, // IEND chunk
0xAE, 0x42, 0x60, 0x82,
}
if err := os.WriteFile(pngFile, pngData, 0644); err != nil {
t.Fatal(err)
}
text := "@" + pngFile + " what is this image?"
result := ProcessFileAttachments(text, dir)
if len(result.FileParts) != 1 {
t.Fatalf("expected 1 FilePart for binary file, got %d", len(result.FileParts))
}
if result.FileParts[0].MediaType != "image/png" {
t.Errorf("expected media type image/png, got %s", result.FileParts[0].MediaType)
}
if result.FileParts[0].Filename != "image.png" {
t.Errorf("expected filename image.png, got %s", result.FileParts[0].Filename)
}
// The @token should be removed from the text
if contains(result.ProcessedText, "@") && contains(result.ProcessedText, pngFile) {
t.Error("expected @token to be removed from processed text for binary file")
}
if contains(result.ProcessedText, "what is this image?") {
// Good, the prompt text should remain
} else {
t.Error("expected prompt text to remain in processed text")
}
}
func TestProcessFileAttachments_MCPResource(t *testing.T) {
// Test @mcp:server:uri token processing with a mock reader
text := "@mcp:test-server:docs://readme tell me about this"
reader := func(serverName, uri string) (string, []byte, string, bool, error) {
if serverName != "test-server" || uri != "docs://readme" {
t.Errorf("unexpected server/uri: %s/%s", serverName, uri)
}
return "Hello from MCP resource", nil, "text/plain", false, nil
}
result := ProcessFileAttachments(text, "/tmp", reader)
if len(result.FileParts) != 0 {
t.Errorf("expected 0 FileParts for text MCP resource, got %d", len(result.FileParts))
}
if !contains(result.ProcessedText, "<resource uri=\"docs://readme\" server=\"test-server\">") {
t.Error("expected <resource> XML wrapping in processed text")
}
if !contains(result.ProcessedText, "Hello from MCP resource") {
t.Error("expected MCP resource content in processed text")
}
}
func TestProcessFileAttachments_MCPResource_Binary(t *testing.T) {
// Test @mcp:server:uri token processing for a binary resource
text := "@mcp:test-server:images://logo describe this"
reader := func(serverName, uri string) (string, []byte, string, bool, error) {
if serverName != "test-server" || uri != "images://logo" {
t.Errorf("unexpected server/uri: %s/%s", serverName, uri)
}
return "", []byte{0x89, 0x50, 0x4E, 0x47}, "image/png", true, nil
}
result := ProcessFileAttachments(text, "/tmp", reader)
if len(result.FileParts) != 1 {
t.Fatalf("expected 1 FilePart for binary MCP resource, got %d", len(result.FileParts))
}
if result.FileParts[0].MediaType != "image/png" {
t.Errorf("expected media type image/png, got %s", result.FileParts[0].MediaType)
}
if result.FileParts[0].Filename != "logo" {
t.Errorf("expected filename 'logo', got %s", result.FileParts[0].Filename)
}
// The @token should be removed from the text
if contains(result.ProcessedText, "@mcp:") {
t.Error("expected @mcp: token to be removed from processed text for binary resource")
}
}
func TestProcessFileAttachments_NoReader(t *testing.T) {
// Without an MCP reader, @mcp: tokens should be left as-is
text := "@mcp:server:resource this is a test"
result := ProcessFileAttachments(text, "/tmp")
if len(result.FileParts) != 0 {
t.Errorf("expected 0 FileParts, got %d", len(result.FileParts))
}
// The @mcp: token should remain unchanged since no reader was provided
if result.ProcessedText != text {
t.Errorf("expected text unchanged without reader, got: %s", result.ProcessedText)
}
}
func TestDetectMediaType(t *testing.T) {
tests := []struct {
ext string
content []byte
expected string
}{
// An intentionally-synthetic extension that is not registered
// in any system MIME database. Exercises the "unknown ext +
// no content" branch, which must return the text/plain default.
// Do not use real extensions (e.g. .go) here: CI images often
// ship /etc/mime.types with entries like ".go → text/x-go",
// which would make the assertion environment-dependent.
{".kitsyntheticext", nil, "text/plain"},
{".png", []byte{0x89, 0x50, 0x4E, 0x47}, "image/png"},
{".jpg", []byte{0xFF, 0xD8, 0xFF}, "image/jpeg"},
{".pdf", []byte{0x25, 0x50, 0x44, 0x46}, "application/pdf"},
{".txt", []byte("hello"), "text/plain"},
{".wav", nil, "audio/wav"},
{".webp", nil, "image/webp"},
}
for _, tt := range tests {
t.Run(tt.ext, func(t *testing.T) {
got := detectMediaType("test"+tt.ext, tt.content)
if got != tt.expected {
t.Errorf("detectMediaType(%q) = %q, want %q", tt.ext, got, tt.expected)
}
})
}
}
func TestIsBinaryMediaType(t *testing.T) {
tests := []struct {
mimeType string
expected bool
}{
{"image/png", true},
{"image/jpeg", true},
{"audio/wav", true},
{"video/mp4", true},
{"application/pdf", true},
{"text/plain", false},
{"text/go", false},
{"application/json", false},
}
for _, tt := range tests {
t.Run(tt.mimeType, func(t *testing.T) {
got := isBinaryMediaType(tt.mimeType)
if got != tt.expected {
t.Errorf("isBinaryMediaType(%q) = %v, want %v", tt.mimeType, got, tt.expected)
}
})
}
}
func contains(s, substr string) bool {
return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsStr(s, substr))
}
func containsStr(s, substr string) bool {
for i := 0; i <= len(s)-len(substr); i++ {
if s[i:i+len(substr)] == substr {
return true
}
}
return false
}
+1
View File
@@ -17,6 +17,7 @@ type Renderer interface {
RenderReasoningBlock(content string, timestamp time.Time) UIMessage
RenderToolMessage(toolName, toolArgs, toolResult string, isError bool) UIMessage
RenderSystemMessage(content string, timestamp time.Time) UIMessage
RenderCustomMessage(content, label string, timestamp time.Time) UIMessage
RenderErrorMessage(errorMsg string, timestamp time.Time) UIMessage
RenderDebugMessage(message string, timestamp time.Time) UIMessage
RenderDebugConfigMessage(config map[string]any, timestamp time.Time) UIMessage
+137 -40
View File
@@ -2,6 +2,7 @@ package ui
import (
"fmt"
"sort"
"strings"
"charm.land/bubbles/v2/key"
@@ -39,7 +40,6 @@ type InputComponent struct {
width int
lastValue string
popupHeight int
title string
submitNext bool // defer submit one tick so popup dismisses cleanly
// Argument completion state. When the user types "/cmd " followed by
@@ -61,6 +61,10 @@ type InputComponent struct {
// autocomplete suggestions. Set by the parent via SetCwd.
cwd string
// mcpResources is a callback that returns available MCP resources for
// the @ autocomplete popup. Set by the parent via SetMCPResourceProvider.
mcpResources func() []FileSuggestion
// appCtrl is used for slash commands that mutate app state.
// May be nil in tests; nil-safe.
appCtrl AppController
@@ -101,17 +105,17 @@ type clipboardImageMsg struct {
err error
}
// NewInputComponent creates a new InputComponent with the given width, title,
// and optional AppController. If appCtrl is nil the component still works but
// NewInputComponent creates a new InputComponent with the given width and
// optional AppController. If appCtrl is nil the component still works but
// /clear and /clear-queue are no-ops.
func NewInputComponent(width int, title string, appCtrl AppController) *InputComponent {
func NewInputComponent(width int, appCtrl AppController) *InputComponent {
ta := textarea.New()
ta.Placeholder = "Type your message..."
ta.ShowLineNumbers = false
ta.Prompt = ""
ta.CharLimit = 0
ta.SetWidth(width - 8) // Account for container padding, border and internal padding
ta.SetHeight(3) // Default to 3 lines like huh
ta.SetHeight(4) // 4 lines for comfortable multi-line input
ta.Focus()
// Override InsertNewline so only ctrl+j and shift+enter insert newlines.
@@ -136,8 +140,8 @@ func NewInputComponent(width int, title string, appCtrl AppController) *InputCom
commands: commands.SlashCommands,
width: width,
popupHeight: 7,
title: title,
appCtrl: appCtrl,
hideHint: true,
}
}
@@ -147,6 +151,12 @@ func (s *InputComponent) SetCwd(cwd string) {
s.cwd = cwd
}
// SetMCPResourceProvider sets a callback that returns MCP resource suggestions
// for the @ autocomplete popup. Called by the parent after construction.
func (s *InputComponent) SetMCPResourceProvider(fn func() []FileSuggestion) {
s.mcpResources = fn
}
// Init implements tea.Model. Starts the cursor blink animation.
func (s *InputComponent) Init() tea.Cmd {
return textarea.Blink
@@ -190,7 +200,7 @@ func (s *InputComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
case tea.KeyPressMsg:
if !s.showPopup {
switch msg.String() {
case "ctrl+d", "enter":
case "enter":
value := s.textarea.Value()
s.pushHistory(value)
s.textarea.SetValue("")
@@ -285,16 +295,25 @@ func (s *InputComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
s.textarea.CursorEnd()
return s, nil
}
selectedCmd := s.filtered[s.selected].Command
// Populate textarea with selected item and submit on next tick.
if s.argMode {
s.textarea.SetValue(s.argCommand + " " + s.filtered[s.selected].Command.Name)
s.textarea.SetValue(s.argCommand + " " + selectedCmd.Name)
} else {
s.textarea.SetValue(s.filtered[s.selected].Command.Name)
s.textarea.SetValue(selectedCmd.Name)
}
s.textarea.CursorEnd()
s.showPopup = false
s.selected = 0
s.submitNext = true
// If the selected command expects arguments, populate
// the input with the command + trailing space so the
// user can type args, instead of auto-submitting.
if !s.argMode && selectedCmd.HasArgs {
s.textarea.SetValue(selectedCmd.Name + " ")
s.textarea.CursorEnd()
} else {
s.submitNext = true
}
return s, nil
}
return s, nil
@@ -323,9 +342,46 @@ func (s *InputComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
// Check for @file trigger first.
cursorCol := len(line) // approximate: cursor is at end after typing
if hasAt, prefix, atIdx := ExtractAtPrefix(line, cursorCol); hasAt && s.cwd != "" {
suggestions := GetFileSuggestions(prefix, s.cwd)
if hasAt, prefix, atIdx := ExtractAtPrefix(line, cursorCol); hasAt {
var suggestions []FileSuggestion
// Local file suggestions (only if cwd is set).
if s.cwd != "" {
suggestions = GetFileSuggestions(prefix, s.cwd)
}
// MCP resource suggestions — merge with file suggestions.
if s.mcpResources != nil {
mcpSuggestions := s.mcpResources()
if prefix != "" {
// Fuzzy-filter MCP resources against the typed prefix.
queryLower := strings.ToLower(prefix)
var filtered []FileSuggestion
for _, r := range mcpSuggestions {
score := scoreFilePath(queryLower, r.RelPath)
if score <= 0 {
// Also try matching against the resource name without prefix.
score = scoreFilePath(queryLower, r.MCPServerName+"/"+r.RelPath)
}
if score > 0 {
r.Score = score
filtered = append(filtered, r)
}
}
mcpSuggestions = filtered
}
suggestions = append(suggestions, mcpSuggestions...)
}
if len(suggestions) > 0 {
// Sort by score descending, cap at maxFileSuggestions.
sort.Slice(suggestions, func(i, j int) bool {
return suggestions[i].Score > suggestions[j].Score
})
if len(suggestions) > maxFileSuggestions {
suggestions = suggestions[:maxFileSuggestions]
}
s.showPopup = true
s.fileMode = true
s.argMode = false
@@ -339,6 +395,8 @@ func (s *InputComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
desc := ""
if fs.IsDir {
desc = "directory"
} else if fs.IsMCPResource {
desc = "mcp:" + fs.MCPServerName
}
s.fileSynthCmds[i] = commands.SlashCommand{Name: name, Description: desc}
s.filtered[i] = FuzzyMatch{Command: &s.fileSynthCmds[i], Score: fs.Score}
@@ -461,19 +519,13 @@ func (s *InputComponent) resetHistoryBrowsing() {
s.savedInput = ""
}
// View implements tea.Model. Renders the title, textarea, autocomplete popup
// View implements tea.Model. Renders the textarea, autocomplete popup
// (if visible), and help text.
func (s *InputComponent) View() tea.View {
containerStyle := lipgloss.NewStyle()
theme := style.GetTheme()
// PaddingLeft(3) aligns with message content: border(1) + paddingLeft(2).
titleStyle := lipgloss.NewStyle().
Foreground(theme.Text).
MarginBottom(1).
PaddingLeft(3)
inputBoxStyle := lipgloss.NewStyle().
Border(lipgloss.ThickBorder()).
BorderLeft(true).
@@ -481,12 +533,12 @@ func (s *InputComponent) View() tea.View {
BorderTop(false).
BorderBottom(false).
BorderForeground(theme.Primary).
MarginTop(1).
MarginBottom(1).
PaddingLeft(2). // match message block paddingLeft
Width(s.width - 1) // full width minus left border
var view strings.Builder
view.WriteString(titleStyle.Render(s.title))
view.WriteString("\n")
view.WriteString(inputBoxStyle.Render(s.textarea.View()))
// Popup is now rendered as a centered overlay in AppModel.View()
@@ -521,12 +573,14 @@ func (s *InputComponent) View() tea.View {
} else {
hint = "^X s steer"
}
} else if availableHintWidth >= 80 {
hint = "enter submit • ctrl+j / shift+enter new line • ctrl+x e editor • ctrl+v paste image"
} else if availableHintWidth >= 67 {
hint = "enter submit • ctrl+j / shift+enter new line • ctrl+v paste image"
hint = "enter submit • ctrl+j new line • ctrl+x e editor • ctrl+v image"
} else if availableHintWidth >= 40 {
hint = "↵ submit • ctrl+j newline • ctrl+v image"
hint = "↵ submit • ctrl+j newline • ^X e editor"
} else if availableHintWidth >= 20 {
hint = "↵ submit • ctrl+j"
hint = "↵ submit • ^X e editor"
} else {
hint = "↵ submit"
}
@@ -647,9 +701,25 @@ func (s *InputComponent) renderPopupWithOptions(centered bool) string {
}
content = indicator + displayName
} else {
nameWidth := 15
if innerWidth < 25 {
nameWidth = max(innerWidth*2/5+1, 8)
// Compute nameWidth from the longest command name in the
// visible slice so we never truncate unnecessarily.
nameWidth := 0
for _, fm := range s.filtered {
if n := len([]rune(fm.Command.Name)); n > nameWidth {
nameWidth = n
}
}
nameWidth += 3 // account for indicator prefix (2) + gap before description (1)
// Ensure descriptions still get at least 20 chars when possible.
maxForName := innerWidth - 20
if maxForName < 8 {
maxForName = innerWidth * 2 / 3
}
if nameWidth > maxForName {
nameWidth = maxForName
}
if nameWidth < 8 {
nameWidth = 8
}
maxNameChars := nameWidth - 2
displayName := sc.Name
@@ -782,9 +852,25 @@ func (s *InputComponent) PendingImageCount() int {
return len(s.pendingImages)
}
// Clear clears the textarea content and resets related state. Returns true if
// there was content to clear, false if the input was already empty.
func (s *InputComponent) Clear() bool {
hadContent := s.textarea.Value() != ""
s.textarea.SetValue("")
s.textarea.CursorEnd()
s.lastValue = ""
s.showPopup = false
s.argMode = false
s.fileMode = false
s.browsingHistory = false
s.savedInput = ""
return hadContent
}
// applyFileCompletion replaces the @prefix in the textarea with the selected
// file suggestion. For directories, it keeps the popup open for further
// drilling. For files, it closes the popup and adds a trailing space.
// file or MCP resource suggestion. For directories, it keeps the popup open
// for further drilling. For files and resources, it closes the popup and adds
// a trailing space.
func (s *InputComponent) applyFileCompletion(idx int) {
if idx >= len(s.fileSuggestions) {
return
@@ -801,19 +887,30 @@ func (s *InputComponent) applyFileCompletion(idx int) {
// Reconstruct: everything before the @ on the last line + @<path>
beforeAt := lastLine[:s.fileAtStartIdx]
needsQuote := strings.Contains(suggestion.RelPath, " ")
var replacement string
if needsQuote {
replacement = `@"` + suggestion.RelPath + `"`
} else {
replacement = "@" + suggestion.RelPath
}
// For files, add a trailing space. For directories, don't — allow
// continued drilling into the directory.
if !suggestion.IsDir {
if suggestion.IsMCPResource {
// MCP resources use @mcp:server:uri format.
// Quote if the URI contains spaces.
ref := "mcp:" + suggestion.MCPServerName + ":" + suggestion.MCPResourceURI
if strings.Contains(ref, " ") {
replacement = `@"` + ref + `"`
} else {
replacement = "@" + ref
}
replacement += " "
} else {
needsQuote := strings.Contains(suggestion.RelPath, " ")
if needsQuote {
replacement = `@"` + suggestion.RelPath + `"`
} else {
replacement = "@" + suggestion.RelPath
}
// For files, add a trailing space. For directories, don't — allow
// continued drilling into the directory.
if !suggestion.IsDir {
replacement += " "
}
}
newLastLine := beforeAt + replacement
@@ -825,7 +922,7 @@ func (s *InputComponent) applyFileCompletion(idx int) {
s.textarea.SetValue(newValue)
s.textarea.CursorEnd()
if suggestion.IsDir {
if suggestion.IsDir && !suggestion.IsMCPResource {
// Keep popup open — trigger a refresh for the new directory.
s.lastValue = "" // force re-evaluation on next update tick
} else {
+1 -1
View File
@@ -156,7 +156,7 @@ func (s *StreamingMessageItem) Render(width int) string {
durationMs = time.Since(s.startTime).Milliseconds()
}
ty := createTypography(style.GetTheme())
rendered = render.ReasoningBlock(s.content, durationMs, ty, style.GetTheme())
rendered = render.ReasoningBlock(s.content, durationMs, width, ty, style.GetTheme())
} else {
// Render as assistant message
rendered = render.AssistantBlock(s.content, width, style.GetTheme())
+35 -4
View File
@@ -150,9 +150,26 @@ func (r *MessageRenderer) SetWidth(width int) {
r.width = width
}
// RenderUserMessage renders a user's input message using herald Tip alert
// RenderUserMessage renders a user's input message with a colored left border.
func (r *MessageRenderer) RenderUserMessage(content string, timestamp time.Time) UIMessage {
rendered := render.UserBlock(content, r.width, r.ty, style.GetTheme())
if strings.TrimSpace(content) == "" {
content = "(empty message)"
}
theme := style.GetTheme()
// Highlight @file tokens with accent color.
content = render.HighlightFileTokens(content, theme)
rendered := renderContentBlock(
content,
r.width,
WithAlign(lipgloss.Left),
WithBorderColor(theme.Success),
WithPaddingTop(0),
WithPaddingBottom(0),
WithMarginBottom(1),
)
return UIMessage{
Type: UserMessage,
@@ -178,7 +195,7 @@ func (r *MessageRenderer) RenderAssistantMessage(content string, timestamp time.
// as live streaming: muted italic text with margin. This is used when resuming
// sessions to display saved reasoning content.
func (r *MessageRenderer) RenderReasoningBlock(content string, timestamp time.Time) UIMessage {
rendered := render.ReasoningBlock(content, 0, r.ty, style.GetTheme())
rendered := render.ReasoningBlock(content, 0, r.width, r.ty, style.GetTheme())
return UIMessage{
Type: AssistantMessage,
@@ -200,6 +217,19 @@ func (r *MessageRenderer) RenderSystemMessage(content string, timestamp time.Tim
}
}
// RenderCustomMessage renders a message with a custom alert label (e.g. "Help").
// Content is rendered as markdown.
func (r *MessageRenderer) RenderCustomMessage(content, label string, timestamp time.Time) UIMessage {
rendered := render.CustomBlock(content, label, r.width, style.GetTheme())
return UIMessage{
Type: SystemMessage,
Content: rendered,
Height: lipgloss.Height(rendered),
Timestamp: timestamp,
}
}
// RenderDebugMessage renders diagnostic and debugging information
func (r *MessageRenderer) RenderDebugMessage(message string, timestamp time.Time) UIMessage {
header := r.ty.H6("🔍 Debug Output")
@@ -399,7 +429,8 @@ func createTypography(theme style.Theme) *herald.Typography {
herald.WithCodeLineNumbers(true),
// Customize alert labels
herald.WithAlertLabel(herald.AlertNote, "Info"),
herald.WithAlertLabel(herald.AlertTip, "You"),
herald.WithAlertLabel(herald.AlertTip, ""),
herald.WithAlertIcon(herald.AlertTip, ""),
herald.WithAlertLabel(herald.AlertWarning, "Working"),
herald.WithAlertLabel(herald.AlertCaution, "Error"),
)
+652 -54
View File
File diff suppressed because it is too large Load Diff
+151 -9
View File
@@ -515,12 +515,12 @@ func TestWindowResize_distributeHeight(t *testing.T) {
ctrl := &stubAppController{}
m, _, _ := newTestAppModel(ctrl)
// With height=30, scroll height = 30 - 1 (separator) - 9 (input) - 1 (statusBar) = 19
// With height=30, scroll height = 30 - 1 (separator) - 8 (input) - 1 (statusBar) = 20
m = sendMsg(m, tea.WindowSizeMsg{Width: 80, Height: 30})
_ = m
if m.scrollList.height != 19 {
t.Fatalf("expected scroll list height=19, got %d", m.scrollList.height)
if m.scrollList.height != 20 {
t.Fatalf("expected scroll list height=20, got %d", m.scrollList.height)
}
}
@@ -853,23 +853,165 @@ func TestSpinnerEvent_hideDoesNotTransitionState(t *testing.T) {
}
// --------------------------------------------------------------------------
// ctrl+c produces tea.Quit
// ctrl+c double-press to quit
// --------------------------------------------------------------------------
// TestCtrlC_producesQuit verifies that ctrl+c always returns a tea.Quit cmd.
// TestCtrlC_producesQuit verifies that double ctrl+c returns a tea.Quit cmd.
func TestCtrlC_producesQuit(t *testing.T) {
ctrl := &stubAppController{}
m, _, _ := newTestAppModel(ctrl)
// First Ctrl+C arms the quit flag.
updated, cmd := m.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl})
m = updated.(*AppModel)
if cmd == nil {
t.Fatal("expected a command after first ctrl+c, got nil")
}
// Should be a reset timer, not quit.
msg := cmd()
if _, ok := msg.(core.CtrlCResetMsg); !ok {
t.Fatalf("expected CtrlCResetMsg after first ctrl+c, got %T", msg)
}
// Second Ctrl+C should quit.
_, cmd = m.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl})
if cmd == nil {
t.Fatal("expected tea.Quit cmd on second ctrl+c, got nil")
}
msg = cmd()
if _, ok := msg.(tea.QuitMsg); !ok {
t.Fatalf("expected QuitMsg from second ctrl+c, got %T", msg)
}
}
// TestCtrlC_clearsInput_firstPress tests that Ctrl+C clears input on first
// press when there's content, and requires a second press to quit.
func TestCtrlC_clearsInput_firstPress(t *testing.T) {
// Create a real InputComponent to test the clear behavior
ctrl := &stubAppController{}
m, _, _ := newTestAppModel(ctrl)
// Replace with real InputComponent that has content
input := NewInputComponent(80, ctrl)
input.textarea.SetValue("some text content")
m.input = input
// First Ctrl+C should clear input, not quit
_, cmd := m.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl})
if cmd == nil {
t.Fatal("expected tea.Quit cmd on ctrl+c, got nil")
// Should have cleared the input
if input.textarea.Value() != "" {
t.Fatalf("expected input to be cleared, got %q", input.textarea.Value())
}
// Should have set ctrlCPressedOnce flag
if !m.ctrlCPressedOnce {
t.Fatal("expected ctrlCPressedOnce to be true after first Ctrl+C")
}
// The command should be a ctrlCResetCmd (not tea.Quit)
if cmd == nil {
t.Fatal("expected a command after first Ctrl+C, got nil")
}
// We verify it's a quit command by running it and checking the message type.
msg := cmd()
if _, ok := msg.(core.CtrlCResetMsg); !ok {
t.Fatalf("expected CtrlCResetMsg, got %T", msg)
}
// Second Ctrl+C should now quit
_, cmd = m.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl})
if cmd == nil {
t.Fatal("expected tea.Quit cmd on second Ctrl+C, got nil")
}
msg = cmd()
if _, ok := msg.(tea.QuitMsg); !ok {
t.Fatalf("expected QuitMsg from ctrl+c cmd, got %T", msg)
t.Fatalf("expected QuitMsg on second Ctrl+C, got %T", msg)
}
}
// TestCtrlC_resetAfterSubmit tests that the Ctrl+C flag is reset after
// submitting a message, so the next Ctrl+C clears input again.
func TestCtrlC_resetAfterSubmit(t *testing.T) {
// Use newTestAppModel but replace the input with a real InputComponent
ctrl := &stubAppController{}
m, _, _ := newTestAppModel(ctrl)
// Replace with real InputComponent
input := NewInputComponent(80, ctrl)
input.textarea.SetValue("content")
m.input = input
// First Ctrl+C clears input
updated, _ := m.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl})
m = updated.(*AppModel)
if input.textarea.Value() != "" {
t.Fatal("expected input to be cleared")
}
// Flag should be set
if !m.ctrlCPressedOnce {
t.Fatal("expected ctrlCPressedOnce to be true after first Ctrl+C")
}
// Simulate CtrlCResetMsg being processed (timer expired)
updated, _ = m.Update(core.CtrlCResetMsg{})
m = updated.(*AppModel)
// Flag should be reset
if m.ctrlCPressedOnce {
t.Fatal("expected ctrlCPressedOnce to be false after CtrlCResetMsg")
}
// Add new content to input
input.textarea.SetValue("new content")
// Next Ctrl+C should clear again (not quit) because flag was reset
_, cmd := m.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl})
if input.textarea.Value() != "" {
t.Fatalf("expected input to be cleared again, got %q", input.textarea.Value())
}
if cmd == nil {
t.Fatal("expected a command after Ctrl+C, got nil")
}
msg := cmd()
if _, ok := msg.(core.CtrlCResetMsg); !ok {
t.Fatalf("expected CtrlCResetMsg, got %T", msg)
}
}
// TestCtrlC_emptyInput_armsQuit tests that Ctrl+C on empty input still
// requires a second press to quit (consistent double-press behavior).
func TestCtrlC_emptyInput_armsQuit(t *testing.T) {
ctrl := &stubAppController{}
m, _, _ := newTestAppModel(ctrl)
// Replace with real InputComponent (empty by default)
input := NewInputComponent(80, ctrl)
m.input = input
// First Ctrl+C on empty input should arm the flag, not quit.
updated, cmd := m.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl})
m = updated.(*AppModel)
if !m.ctrlCPressedOnce {
t.Fatal("expected ctrlCPressedOnce to be true after first Ctrl+C")
}
if cmd == nil {
t.Fatal("expected a command (reset timer), got nil")
}
msg := cmd()
if _, ok := msg.(core.CtrlCResetMsg); !ok {
t.Fatalf("expected CtrlCResetMsg, got %T", msg)
}
// Second Ctrl+C should quit.
_, cmd = m.Update(tea.KeyPressMsg{Code: 'c', Mod: tea.ModCtrl})
if cmd == nil {
t.Fatal("expected tea.Quit cmd on second Ctrl+C, got nil")
}
msg = cmd()
if _, ok := msg.(tea.QuitMsg); !ok {
t.Fatalf("expected QuitMsg on second Ctrl+C, got %T", msg)
}
}
+6
View File
@@ -288,3 +288,9 @@ func (pr *ProgressReader) Close() error {
return nil
}
// NewProgressReadCloser is a convenience wrapper around NewProgressReader that
// returns an io.ReadCloser, suitable for use as a ProgressReaderFunc callback.
func NewProgressReadCloser(r io.Reader) io.ReadCloser {
return NewProgressReader(r)
}
+78 -9
View File
@@ -19,9 +19,10 @@ import (
type promptMode string
const (
promptModeSelect promptMode = "select"
promptModeConfirm promptMode = "confirm"
promptModeInput promptMode = "input"
promptModeSelect promptMode = "select"
promptModeConfirm promptMode = "confirm"
promptModeInput promptMode = "input"
promptModePassword promptMode = "password"
)
// promptResult carries the synchronous outcome of a prompt overlay update.
@@ -102,10 +103,38 @@ func newInputPrompt(message, placeholder, defaultValue string, width, height int
}
}
// Init returns the initial command for the prompt overlay. For input mode
// this starts the cursor blink animation.
// newPasswordPrompt creates a prompt overlay for password input (masked).
func newPasswordPrompt(message string, width, height int) *promptOverlay {
ta := textarea.New()
ta.Placeholder = "Enter password"
ta.ShowLineNumbers = false
ta.Prompt = ""
ta.CharLimit = 0
ta.SetWidth(width - 12) // account for border + padding
ta.SetHeight(1)
ta.Focus()
// Prevent Enter from inserting a newline — we intercept it for submit.
ta.KeyMap.InsertNewline = key.NewBinding(
key.WithKeys("ctrl+j", "shift+enter"),
)
// Enable password masking - the textarea will show dots instead of characters
// Note: textarea doesn't have built-in password masking, so we handle it in View()
return &promptOverlay{
mode: promptModePassword,
message: message,
inputTA: ta,
width: width,
height: height,
}
}
// Init returns the initial command for the prompt overlay. For input/password
// modes this starts the cursor blink animation.
func (p *promptOverlay) Init() tea.Cmd {
if p.mode == promptModeInput {
if p.mode == promptModeInput || p.mode == promptModePassword {
return textarea.Blink
}
return nil
@@ -113,13 +142,13 @@ func (p *promptOverlay) Init() tea.Cmd {
// Update handles messages for the prompt overlay. It returns a non-nil
// *promptResult when the user completes or cancels the prompt. The returned
// tea.Cmd is for textarea blink ticks (input mode only).
// tea.Cmd is for textarea blink ticks (input/password modes only).
func (p *promptOverlay) Update(msg tea.Msg) (*promptResult, tea.Cmd) {
switch msg := msg.(type) {
case tea.WindowSizeMsg:
p.width = msg.Width
p.height = msg.Height
if p.mode == promptModeInput {
if p.mode == promptModeInput || p.mode == promptModePassword {
p.inputTA.SetWidth(p.width - 12)
}
return nil, nil
@@ -132,11 +161,13 @@ func (p *promptOverlay) Update(msg tea.Msg) (*promptResult, tea.Cmd) {
return p.updateConfirm(msg)
case promptModeInput:
return p.updateInput(msg)
case promptModePassword:
return p.updatePassword(msg)
}
}
// Pass non-key messages to textarea for blink animation.
if p.mode == promptModeInput {
if p.mode == promptModeInput || p.mode == promptModePassword {
var cmd tea.Cmd
p.inputTA, cmd = p.inputTA.Update(msg)
return nil, cmd
@@ -202,6 +233,20 @@ func (p *promptOverlay) updateInput(msg tea.KeyPressMsg) (*promptResult, tea.Cmd
}
}
func (p *promptOverlay) updatePassword(msg tea.KeyPressMsg) (*promptResult, tea.Cmd) {
switch msg.String() {
case "enter":
return &promptResult{completed: true, value: p.inputTA.Value()}, nil
case "esc":
return &promptResult{cancelled: true}, nil
default:
// Delegate character input, backspace, cursor movement, etc.
var cmd tea.Cmd
p.inputTA, cmd = p.inputTA.Update(msg)
return nil, cmd
}
}
// Render returns the prompt as a styled string for inline composition in the
// AppModel layout. The prompt replaces the normal input area (below the
// separator and above the status bar) rather than taking over the full screen.
@@ -216,6 +261,8 @@ func (p *promptOverlay) Render() string {
content = p.viewConfirm(theme)
case promptModeInput:
content = p.viewInput(theme)
case promptModePassword:
content = p.viewPassword(theme)
}
return renderContentBlock(content, p.width,
@@ -286,3 +333,25 @@ func (p *promptOverlay) viewInput(theme style.Theme) string {
return strings.Join(lines, "\n")
}
func (p *promptOverlay) viewPassword(theme style.Theme) string {
var lines []string
// Add 🔐 icon to message for password prompt
lines = append(lines, lipgloss.NewStyle().Bold(true).Foreground(theme.Text).Render("🔐 "+p.message))
lines = append(lines, "")
// Mask the password input with dots
passwordValue := p.inputTA.Value()
masked := strings.Repeat("•", len([]rune(passwordValue)))
// Render the masked password in a style that looks like input
maskedStyle := lipgloss.NewStyle().Foreground(theme.Text)
cursor := lipgloss.NewStyle().Foreground(theme.Accent).Render("█")
lines = append(lines, maskedStyle.Render(masked)+cursor)
lines = append(lines, "")
lines = append(lines, lipgloss.NewStyle().
Foreground(theme.Muted).
Render(" Enter submit Esc cancel (input is hidden)"))
return strings.Join(lines, "\n")
}
+71 -2
View File
@@ -5,6 +5,7 @@ package render
import (
"fmt"
"regexp"
"strings"
"charm.land/lipgloss/v2"
@@ -13,8 +14,14 @@ import (
"github.com/mark3labs/kit/internal/ui/style"
)
// fileTokenPattern matches @file references in user text. Supports:
// - @"path with spaces.txt" (quoted)
// - @path/to/file.txt (unquoted, no spaces)
var fileTokenPattern = regexp.MustCompile(`@"[^"]+"|@[^\s]+`)
// UserBlock renders a user message with herald Tip styling.
// The width parameter controls line wrapping so long messages don't overflow.
// Any @file tokens in the content are highlighted with the theme accent color.
func UserBlock(content string, width int, ty *herald.Typography, theme style.Theme) string {
if strings.TrimSpace(content) == "" {
content = "(empty message)"
@@ -27,10 +34,23 @@ func UserBlock(content string, width int, ty *herald.Typography, theme style.The
content = lipgloss.Wrap(content, width-4, "")
}
// Highlight @file tokens with accent color so file references are
// visually distinct from surrounding prompt text.
content = HighlightFileTokens(content, theme)
rendered := ty.Tip(content)
return styleMarginBottom(theme, rendered)
}
// HighlightFileTokens wraps @file tokens in the given text with the theme
// accent color so they stand out visually in rendered user messages.
func HighlightFileTokens(text string, theme style.Theme) string {
accentStyle := lipgloss.NewStyle().Foreground(theme.Accent).Bold(true)
return fileTokenPattern.ReplaceAllStringFunc(text, func(token string) string {
return accentStyle.Render(token)
})
}
// AssistantBlock renders an assistant message with markdown styling.
func AssistantBlock(content string, width int, theme style.Theme) string {
if strings.TrimSpace(content) == "" {
@@ -43,14 +63,18 @@ func AssistantBlock(content string, width int, theme style.Theme) string {
// ReasoningBlock renders a reasoning/thinking block with muted italic text.
// If duration > 0, shows "Thought for Xs" label. Otherwise shows just "Thought".
func ReasoningBlock(content string, duration int64, ty *herald.Typography, theme style.Theme) string {
// The width parameter controls soft-wrapping so long reasoning lines don't get cut off.
func ReasoningBlock(content string, duration int64, width int, ty *herald.Typography, theme style.Theme) string {
if strings.TrimSpace(content) == "" {
return ""
}
// Match live streaming styling: muted italic text
// Match live streaming styling: muted italic text.
lines := strings.Split(strings.TrimRight(content, "\n"), "\n")
contentStr := strings.TrimLeft(strings.Join(lines, "\n"), " \t\n")
if width > 4 {
contentStr = wrapText(contentStr, width-4)
}
mutedStyle := lipgloss.NewStyle().Foreground(theme.Muted)
contentRendered := mutedStyle.Render(ty.Italic(contentStr))
@@ -85,6 +109,45 @@ func SystemBlock(content string, ty *herald.Typography, theme style.Theme) strin
return styleMarginBottom(theme, rendered)
}
// CustomBlock renders a message with herald Note styling and a custom label.
// Content is rendered as markdown before being wrapped in the alert. This
// creates a one-off Typography instance with the given label so callers
// can use any title (e.g. "Help", "Warning") without changing the shared
// typography's default "Info" label.
func CustomBlock(content, label string, width int, theme style.Theme) string {
if strings.TrimSpace(content) == "" {
content = "No content available"
}
// Render markdown first — subtract 4 for the alert bar prefix ("│ ").
mdWidth := max(width-4, 10)
rendered := style.ToMarkdown(content, mdWidth)
ty := herald.New(
herald.WithPalette(herald.ColorPalette{
Primary: theme.Primary,
Secondary: theme.Secondary,
Tertiary: theme.Info,
Accent: theme.Accent,
Highlight: theme.Highlight,
Muted: theme.Muted,
Text: theme.Text,
Surface: theme.Background,
Base: theme.CodeBg,
}),
herald.WithAlertPalette(herald.AlertPalette{
Note: theme.Info,
Tip: theme.Success,
Important: theme.Accent,
Warning: theme.Warning,
Caution: theme.Error,
}),
herald.WithAlertLabel(herald.AlertNote, label),
)
alertRendered := ty.Note(rendered)
return styleMarginBottom(theme, alertRendered)
}
// ErrorBlock renders an error message with herald Caution styling.
func ErrorBlock(errorMsg string, ty *herald.Typography, theme style.Theme) string {
rendered := ty.Caution(errorMsg)
@@ -133,3 +196,9 @@ func ToolBlock(displayName, params, body string, isError bool, width int, ty *he
func styleMarginBottom(theme style.Theme, content string) string {
return lipgloss.NewStyle().MarginBottom(1).Render(content)
}
// wrapText soft-wraps a string to the given width using lipgloss, which is
// ANSI-aware and preserves escape sequences across line breaks.
func wrapText(s string, width int) string {
return lipgloss.NewStyle().Width(width).Render(s)
}
+111
View File
@@ -0,0 +1,111 @@
package render
import (
"strings"
"testing"
"github.com/indaco/herald"
"github.com/mark3labs/kit/internal/ui/style"
)
// testTypography creates a herald Typography for tests.
func testTypography(theme style.Theme) *herald.Typography {
return herald.New(
herald.WithPalette(herald.ColorPalette{
Primary: theme.Primary,
Secondary: theme.Secondary,
Tertiary: theme.Info,
Accent: theme.Accent,
Highlight: theme.Highlight,
Muted: theme.Muted,
Text: theme.Text,
Surface: theme.Background,
Base: theme.CodeBg,
}),
herald.WithAlertLabel(herald.AlertTip, ""),
herald.WithAlertIcon(herald.AlertTip, ""),
)
}
func TestHighlightFileTokens(t *testing.T) {
theme := style.DefaultTheme()
tests := []struct {
name string
input string
wantHas []string // substrings that must be present in the output
wantNone []string // substrings that must NOT be present as plain text
}{
{
name: "no tokens",
input: "hello world",
wantHas: []string{"hello world"},
},
{
name: "single unquoted token",
input: "refactor @main.go please",
wantHas: []string{"@main.go", "refactor", "please"},
},
{
name: "quoted token with spaces",
input: `check @"path with spaces/file.txt" out`,
wantHas: []string{`@"path with spaces/file.txt"`, "check", "out"},
},
{
name: "multiple tokens",
input: "@main.go @utils.go refactor these",
wantHas: []string{"@main.go", "@utils.go", "refactor these"},
},
{
name: "path with directory",
input: "look at @internal/ui/render/blocks.go",
wantHas: []string{"@internal/ui/render/blocks.go", "look at"},
},
{
name: "empty string",
input: "",
wantHas: []string{""},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := HighlightFileTokens(tt.input, theme)
for _, want := range tt.wantHas {
if !strings.Contains(result, want) {
t.Errorf("HighlightFileTokens(%q) = %q, want substring %q", tt.input, result, want)
}
}
// If there were @tokens, the result should contain ANSI escape
// sequences (from lipgloss styling).
if fileTokenPattern.MatchString(tt.input) && !strings.Contains(result, "\x1b[") {
t.Errorf("HighlightFileTokens(%q) should contain ANSI escapes for @tokens but got %q", tt.input, result)
}
})
}
}
func TestUserBlockHighlightsFileTokens(t *testing.T) {
theme := style.DefaultTheme()
ty := testTypography(theme)
// A user message with @file tokens should contain ANSI escapes around the token.
content := "refactor @main.go and @utils.go"
result := UserBlock(content, 80, ty, theme)
// The rendered output should contain both file references.
if !strings.Contains(result, "@main.go") {
t.Errorf("UserBlock output should contain @main.go, got:\n%s", result)
}
if !strings.Contains(result, "@utils.go") {
t.Errorf("UserBlock output should contain @utils.go, got:\n%s", result)
}
// Verify ANSI codes are present (the tokens are styled).
if !strings.Contains(result, "\x1b[") {
t.Errorf("UserBlock output should contain ANSI escape codes for styled @file tokens")
}
}
+8 -2
View File
@@ -472,6 +472,10 @@ func (s *StreamComponent) renderReasoningBlock(reasoning string) string {
// Main content using Italic with Muted color for visual distinction.
content := strings.TrimLeft(strings.Join(lines, "\n"), " \t\n")
// Soft-wrap to the available width so long lines don't get cut off.
if s.width > 4 {
content = lipgloss.NewStyle().Width(s.width - 4).Render(content)
}
theme := GetTheme()
mutedStyle := lipgloss.NewStyle().Foreground(theme.Muted)
parts = append(parts, mutedStyle.Render(s.ty.Italic(content)))
@@ -588,8 +592,10 @@ func formatToolExecutionMessage(toolName string) string {
return toolName
}
// UpdateTheme refreshes the component's typography instance with colors from
// the current theme. This is called when the user changes themes via /theme.
// UpdateTheme refreshes the component's typography instance and spinner
// animation frames with colors from the current theme. This is called when
// the user changes themes via /theme.
func (s *StreamComponent) UpdateTheme() {
s.ty = createTypography(GetTheme())
s.spinnerFrames = knightRiderFrames()
}
-4
View File
@@ -200,10 +200,6 @@ func (ts *TreeSelectorComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
case key.Matches(msg, key.NewBinding(key.WithKeys("ctrl+l"))):
ts.filter = TreeFilterLabelOnly
ts.rebuildFlatList()
case key.Matches(msg, key.NewBinding(key.WithKeys("ctrl+a"))):
ts.filter = TreeFilterAll
ts.rebuildFlatList()
default:
// Typing search.
if msg.Text != "" && len(msg.Text) == 1 {
+19 -14
View File
@@ -134,23 +134,28 @@ func (ut *UsageTracker) EstimateAndUpdateUsage(inputText, outputText string) {
}
// SetContextTokens records the approximate current context window utilization.
// This should be set from FinalUsage.InputTokens, which already includes the
// full conversation history (system prompt + all previous messages). Do NOT
// add OutputTokens as that would double-count (output becomes input next turn).
// Use FinalResponse.Usage rather than aggregate TotalUsage, because TotalUsage
// sums across all tool-calling steps and overstates the actual window fill level.
//
// The value should include ALL token categories from the last API call:
//
// InputTokens + CacheReadTokens + CacheCreationTokens + OutputTokens
//
// With Anthropic prompt caching, InputTokens can be near-zero while
// CacheReadTokens holds the bulk of the context. All four must be summed
// to get the true context window fill level.
//
// OutputTokens is included because the assistant's output becomes part of
// the context on the next turn.
//
// Use FinalResponse.Usage (last step only) rather than aggregate TotalUsage,
// because TotalUsage sums across all tool-calling steps and overstates the
// actual window fill level.
//
// The value is set unconditionally (not max-only) so that context shrinks
// correctly after compaction.
func (ut *UsageTracker) SetContextTokens(tokens int) {
ut.mu.Lock()
defer ut.mu.Unlock()
// Track the maximum context seen so far. In multi-step tool calls,
// FinalUsage.InputTokens may reflect only the last step's input, which
// can be smaller than previous steps. We want to show the largest context
// the model has processed in this session.
if tokens > ut.contextTokens {
ut.contextTokens = tokens
}
// If tokens < current, we keep the larger value (no-op)
// This prevents the display from dropping during multi-step tool calls.
ut.contextTokens = tokens
}
// RenderUsageInfo generates a formatted string displaying current usage statistics
+89
View File
@@ -77,6 +77,11 @@ host, err := kit.New(ctx, &kit.Options{
// Compaction
AutoCompact: true, // Auto-compact near context limit
// In-process MCP servers (map name → *kit.MCPServer)
InProcessMCPServers: map[string]*kit.MCPServer{
"docs": mcpSrv,
},
})
```
@@ -112,6 +117,79 @@ response, err := host.Prompt(
)
```
### Dynamic MCP Server Management
Add, remove, and list MCP servers at runtime:
```go
// Add an MCP server at runtime
n, err := host.AddMCPServer(ctx, "github", kit.MCPServerConfig{
Command: "npx",
Args: []string{"-y", "@modelcontextprotocol/server-github"},
})
fmt.Printf("Loaded %d tools from MCP server\n", n)
// List connected MCP servers
for _, s := range host.ListMCPServers() {
fmt.Printf("%s: %d tools\n", s.Name, s.ToolCount)
}
// Disconnect a server and remove its tools
host.RemoveMCPServer("github")
```
### In-Process MCP Servers
Register mcp-go servers that run in the same process — no subprocess spawning,
no network I/O. This is ideal for custom tool servers implemented in Go:
```go
import (
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
// Create an mcp-go server with tools
mcpSrv := server.NewMCPServer("my-tools", "1.0.0",
server.WithToolCapabilities(true),
)
mcpSrv.AddTool(mcp.NewTool("search_docs",
mcp.WithDescription("Search documentation"),
mcp.WithString("query", mcp.Required()),
), searchHandler)
// Option 1: At init time via Options
host, _ := kit.New(ctx, &kit.Options{
InProcessMCPServers: map[string]*kit.MCPServer{
"docs": mcpSrv,
},
})
// Option 2: At runtime
n, err := host.AddInProcessMCPServer(ctx, "docs", mcpSrv)
fmt.Printf("Loaded %d tools from in-process server\n", n)
```
Kit does not take ownership of the server's lifecycle — the caller is responsible for any cleanup. In-process server tools are prefixed the same way as external MCP servers (e.g. `"docs__search_docs"`).
### MCP Prompts
MCP servers can expose prompt templates via the MCP prompts capability.
Kit exposes these through the SDK:
```go
// List prompts from all connected MCP servers
prompts := host.ListMCPPrompts()
for _, p := range prompts {
fmt.Printf("%s/%s: %s\n", p.Server, p.Name, p.Description)
}
// Get a specific prompt with arguments
msg, err := host.GetMCPPrompt(ctx, "server-name", "prompt-name", map[string]string{
"topic": "concurrency",
})
```
### Session Management
Maintain conversation context:
@@ -145,6 +223,16 @@ kit.LLMUsage // {InputTokens, OutputTokens, TotalTokens, ...}
kit.LLMResponse // {Content, FinishReason, Usage}
kit.LLMFilePart // {Filename, Data []byte, MediaType}
// MCP OAuth types
kit.MCPServer // *server.MCPServer for in-process MCP transport
kit.MCPServerConfig // Configuration for an MCP server (stdio, SSE, or in-process)
kit.MCPAuthHandler // Interface: handles user-facing OAuth authorization
kit.DefaultMCPAuthHandler // Port + callback-server mechanics; set OnAuthURL for presentation
kit.CLIMCPAuthHandler // CLI wrapper: opens browser, prints status
kit.MCPTokenStore // Persists OAuth tokens for a single MCP server
kit.MCPToken // OAuth token (access token, refresh token, expiry)
kit.MCPTokenStoreFactory // Creates an MCPTokenStore for a given server URL
// Conversion helpers
msgs := kit.ConvertToLLMMessages(&msg) // SDK Message → []LLMMessage
msg := kit.ConvertFromLLMMessage(lMsg) // LLMMessage → SDK Message
@@ -192,6 +280,7 @@ Key `Options` fields for SDK usage:
| `NoSession` | Ephemeral mode (no session persistence) |
| `SessionPath` | Open specific session file |
| `Continue` | Resume most recent session |
| `InProcessMCPServers` | Map of name → `*kit.MCPServer` for in-process MCP servers |
| `Debug` | Enable debug logging |
## Environment Variables
+5 -6
View File
@@ -22,13 +22,13 @@ func NewTreeManagerAdapter(tm *session.TreeManager) SessionManager {
// AppendMessage implements SessionManager.
func (a *treeManagerAdapter) AppendMessage(msg LLMMessage) (string, error) {
// LLMMessage is just an alias for fantasy.Message, so no conversion needed
// LLMMessage is a type alias, so no conversion needed.
return a.inner.AppendLLMMessage(msg)
}
// GetMessages implements SessionManager.
func (a *treeManagerAdapter) GetMessages() []LLMMessage {
// LLMMessage is just an alias for fantasy.Message
// LLMMessage is a type alias, so no conversion needed.
return a.inner.GetLLMMessages()
}
@@ -223,9 +223,8 @@ func (a *treeManagerAdapter) convertEntry(entry any) *BranchEntry {
}
}
// convertKitMessagesToFantasy converts kit LLM messages to fantasy messages.
// Since LLMMessage is an alias for fantasy.Message, this is a no-op.
func convertKitMessagesToFantasy(msgs []LLMMessage) []fantasy.Message {
// LLMMessage is just an alias for fantasy.Message, so we can type convert
// convertToLLMMessages converts kit LLM messages to the underlying provider
// message type. Since LLMMessage is a type alias, this is a no-op.
func convertToLLMMessages(msgs []LLMMessage) []fantasy.Message {
return msgs
}
+26 -3
View File
@@ -31,6 +31,11 @@ func (m *Kit) EstimateContextTokens() int {
// limit and should be compacted.
// Formula: contextTokens > contextWindow reserveTokens.
// Returns false if the model's context limit is unknown.
//
// When API-reported token counts are available (after at least one turn),
// the real count is used instead of the text-based heuristic. This is
// significantly more accurate because it includes system prompts, tool
// definitions, and other overhead that the heuristic cannot account for.
func (m *Kit) ShouldCompact() bool {
info := m.GetModelInfo()
if info == nil || info.Limit.Context <= 0 {
@@ -42,8 +47,18 @@ func (m *Kit) ShouldCompact() bool {
reserveTokens = m.compactionOpts.ReserveTokens
}
// Prefer the real API-reported token count when available.
m.lastInputTokensMu.RLock()
realTokens := m.lastInputTokens
m.lastInputTokensMu.RUnlock()
if realTokens > 0 {
return realTokens > info.Limit.Context-reserveTokens
}
// Fall back to text-based heuristic before first turn completes.
messages := m.session.GetMessages()
return compaction.ShouldCompact(convertKitMessagesToFantasy(messages), info.Limit.Context, reserveTokens)
return compaction.ShouldCompact(convertToLLMMessages(messages), info.Limit.Context, reserveTokens)
}
// GetContextStats returns current context usage statistics including
@@ -188,9 +203,9 @@ func (m *Kit) compactInternal(ctx context.Context, opts *CompactionOptions, cust
// custom summary. It still determines the cut point and persists a
// CompactionEntry.
func (m *Kit) applyCustomCompaction(summary string, messages []LLMMessage, opts *CompactionOptions) (*CompactionResult, error) {
originalTokens := compaction.EstimateMessageTokens(convertKitMessagesToFantasy(messages))
originalTokens := compaction.EstimateMessageTokens(convertToLLMMessages(messages))
cutPoint := compaction.FindCutPoint(convertKitMessagesToFantasy(messages), opts.KeepRecentTokens)
cutPoint := compaction.FindCutPoint(convertToLLMMessages(messages), opts.KeepRecentTokens)
if cutPoint == 0 {
cutPoint = len(messages) - 1
if cutPoint < 1 {
@@ -245,6 +260,14 @@ func (m *Kit) persistAndEmitCompaction(
); err != nil {
return fmt.Errorf("failed to persist compaction entry: %w", err)
}
// Reset the API-reported token count so GetContextStats() and
// ShouldCompact() don't use stale pre-compaction values. The next
// API call will set the accurate post-compaction count.
m.lastInputTokensMu.Lock()
m.lastInputTokens = 0
m.lastInputTokensMu.Unlock()
m.events.emit(CompactionEvent{
Summary: summary,
OriginalTokens: originalTokens,
+31 -10
View File
@@ -38,20 +38,37 @@ Guidelines:
- Be concise in your responses
- Show file paths clearly when working with files`
// setSDKDefaults registers the same viper defaults that the CLI sets via
// cobra flag bindings. This ensures the SDK behaves identically to the CLI
// even when cobra is not used.
// sdkDefaultMaxTokens is the last-resort ceiling applied when the SDK caller
// has not configured max-tokens via Options, env, config, or a per-model
// default. It matches the CLI's --max-tokens cobra default so SDK and CLI
// callers see the same base value before per-model right-sizing runs.
// It is intentionally applied on the *models.ProviderConfig struct
// (not via viper) so that viper.IsSet("max-tokens") remains false and the
// right-sizing + per-model-default paths continue to work.
const sdkDefaultMaxTokens = 8192
// setSDKDefaults registers viper defaults that match the CLI's cobra flag
// defaults for keys where SetDefault does not interfere with downstream
// viper.IsSet() checks.
//
// Keys that participate in "explicit vs unset" precedence downstream —
// max-tokens, temperature, top-p, top-k, frequency-penalty, presence-penalty,
// thinking-level — are deliberately NOT registered here. viper.SetDefault
// causes viper.IsSet() to return true, which would suppress per-model
// defaults (ApplyModelSettings) and automatic right-sizing (rightSizeMaxTokens)
// for every SDK-created Kit. Those defaults are instead applied:
//
// - max-tokens: as a last-resort struct-level floor (sdkDefaultMaxTokens)
// in kit.New() after BuildProviderConfig returns, when the resolved
// value is still zero.
// - thinking-level: handled implicitly by models.ParseThinkingLevel("")
// which returns models.ThinkingOff.
// - sampling params (temperature, top-p, top-k, frequency/presence-penalty):
// left as nil pointers so provider libraries apply their own defaults.
func setSDKDefaults() {
viper.SetDefault("model", "anthropic/claude-sonnet-4-5-20250929")
viper.SetDefault("system-prompt", defaultSystemPrompt)
viper.SetDefault("max-tokens", 4096)
viper.SetDefault("temperature", 0.7)
viper.SetDefault("top-p", 0.95)
viper.SetDefault("top-k", 40)
viper.SetDefault("frequency-penalty", 0.0)
viper.SetDefault("presence-penalty", 0.0)
viper.SetDefault("stream", true)
viper.SetDefault("thinking-level", "off")
viper.SetDefault("num-gpu-layers", -1)
viper.SetDefault("main-gpu", 0)
}
@@ -102,6 +119,10 @@ func InitConfig(configFile string, debug bool) error {
}
viper.SetEnvPrefix("KIT")
// Map hyphenated config keys (e.g. "max-tokens") to underscored env
// var names (e.g. KIT_MAX_TOKENS). Without this, AutomaticEnv looks
// for KIT_MAX-TOKENS and silently misses valid env overrides.
viper.SetEnvKeyReplacer(strings.NewReplacer("-", "_"))
viper.AutomaticEnv()
return nil
}
+136 -3
View File
@@ -23,6 +23,14 @@ const (
EventMessageUpdate EventType = "message_update"
// EventMessageEnd fires when the assistant message is complete.
EventMessageEnd EventType = "message_end"
// EventToolCallStart fires when the LLM begins generating tool call arguments.
// The tool name is known but arguments are still streaming.
EventToolCallStart EventType = "tool_call_start"
// EventToolCallDelta fires for each streamed fragment of tool call arguments.
EventToolCallDelta EventType = "tool_call_delta"
// EventToolCallEnd fires when tool argument streaming is complete, before
// the tool call is parsed and execution begins.
EventToolCallEnd EventType = "tool_call_end"
// EventToolCall fires when a tool call has been parsed and is about to execute.
EventToolCall EventType = "tool_call"
// EventToolExecutionStart fires when a tool begins executing.
@@ -45,6 +53,8 @@ const (
// EventToolOutput fires when a tool produces streaming output chunks.
EventToolOutput EventType = "tool_output"
EventStepUsage EventType = "step_usage"
// EventPasswordPrompt fires when a sudo command needs a password.
EventPasswordPrompt EventType = "password_prompt"
// EventSteerConsumed fires when one or more steering messages have been
// injected into the agent turn via PrepareStep.
EventSteerConsumed EventType = "steer_consumed"
@@ -108,6 +118,38 @@ func parseToolArgs(toolArgs string) map[string]any {
return nil
}
// ---------------------------------------------------------------------------
// Finish reason constants
// ---------------------------------------------------------------------------
// Finish reasons reported by the LLM provider on a completed turn. These
// mirror fantasy.FinishReason string values so comparisons against
// TurnEndEvent.StopReason / TurnResult.StopReason are stable across
// providers.
const (
// FinishReasonStop: the model produced a natural stop (e.g. stop sequence
// or end-of-turn signal).
FinishReasonStop = "stop"
// FinishReasonLength: the model hit the configured max_output_tokens
// budget. The response is truncated. Surface this to the user and
// consider raising --max-tokens / KIT_MAX_TOKENS / modelSettings[...]
// .maxTokens.
FinishReasonLength = "length"
// FinishReasonToolCalls: the model stopped to emit tool calls (normal
// mid-turn state during agentic loops).
FinishReasonToolCalls = "tool-calls"
// FinishReasonContentFilter: the provider's safety filter stopped
// generation.
FinishReasonContentFilter = "content-filter"
// FinishReasonError: the model stopped because of an error.
FinishReasonError = "error"
// FinishReasonOther: provider-specific reason that doesn't map to any of
// the above.
FinishReasonOther = "other"
// FinishReasonUnknown: the provider didn't report a finish reason.
FinishReasonUnknown = "unknown"
)
// ---------------------------------------------------------------------------
// Concrete event structs
// ---------------------------------------------------------------------------
@@ -122,9 +164,13 @@ func (e TurnStartEvent) EventType() EventType { return EventTurnStart }
// TurnEndEvent fires after the agent finishes processing.
type TurnEndEvent struct {
Response string
Error error
StopReason string // "end_turn", "max_tokens", "tool_use", "error", etc.
Response string
Error error
// StopReason is the LLM provider's finish reason for the final step of
// the turn. Compare against the FinishReason* constants — in particular,
// FinishReasonLength indicates the response was truncated because the
// agent hit its max_output_tokens budget.
StopReason string
}
// EventType implements Event.
@@ -178,6 +224,40 @@ type MessageEndEvent struct {
// EventType implements Event.
func (e MessageEndEvent) EventType() EventType { return EventMessageEnd }
// ToolCallStartEvent fires when the LLM begins generating tool call arguments.
// The tool name is known at this point but the full arguments are still being
// streamed. UIs can use this to show a "running" indicator immediately instead
// of waiting for the full argument JSON to finish streaming.
type ToolCallStartEvent struct {
ToolCallID string // Stable ID for correlating tool lifecycle events
ToolName string
ToolKind string // Tool classification: "execute", "edit", "read", "search", "agent"
}
// EventType implements Event.
func (e ToolCallStartEvent) EventType() EventType { return EventToolCallStart }
// ToolCallDeltaEvent fires for each streamed fragment of tool call arguments.
// Useful for live-previewing artifact content as it's generated, or showing a
// progress indicator with byte count.
type ToolCallDeltaEvent struct {
ToolCallID string // Stable ID for correlating tool lifecycle events
Delta string // JSON fragment of tool arguments
}
// EventType implements Event.
func (e ToolCallDeltaEvent) EventType() EventType { return EventToolCallDelta }
// ToolCallEndEvent fires when tool argument streaming is complete, before
// the tool call is parsed and execution begins. UIs can use this to
// transition from an "generating args" state to an "executing" state.
type ToolCallEndEvent struct {
ToolCallID string // Stable ID for correlating tool lifecycle events
}
// EventType implements Event.
func (e ToolCallEndEvent) EventType() EventType { return EventToolCallEnd }
// ToolCallEvent fires when a tool call has been parsed.
type ToolCallEvent struct {
ToolCallID string // Stable ID for correlating tool lifecycle events
@@ -299,6 +379,26 @@ type SteerConsumedEvent struct {
// EventType implements Event.
func (e SteerConsumedEvent) EventType() EventType { return EventSteerConsumed }
// PasswordPromptEvent fires when a sudo command needs a password.
// The TUI should display a password prompt and send the result back via ResponseCh.
type PasswordPromptEvent struct {
// Prompt is the message to display to the user.
Prompt string
// ResponseCh receives the password from the TUI.
// The TUI must send exactly one value: (password, false) for submit
// or ("", true) for cancel.
ResponseCh chan<- PasswordPromptResponse
}
// PasswordPromptResponse carries the password prompt result.
type PasswordPromptResponse struct {
Password string
Cancelled bool
}
// EventType implements Event.
func (e PasswordPromptEvent) EventType() EventType { return EventPasswordPrompt }
// ---------------------------------------------------------------------------
// EventBus
// ---------------------------------------------------------------------------
@@ -362,6 +462,39 @@ func (m *Kit) OnToolCall(handler func(ToolCallEvent)) func() {
})
}
// OnToolCallStart registers a handler that fires only for ToolCallStartEvent.
// This fires when the LLM begins generating tool call arguments — before the
// full argument JSON is available. Returns an unsubscribe function.
func (m *Kit) OnToolCallStart(handler func(ToolCallStartEvent)) func() {
return m.Subscribe(func(e Event) {
if tcs, ok := e.(ToolCallStartEvent); ok {
handler(tcs)
}
})
}
// OnToolCallDelta registers a handler that fires only for ToolCallDeltaEvent.
// Each delta contains a JSON fragment of tool call arguments as they stream in.
// Returns an unsubscribe function.
func (m *Kit) OnToolCallDelta(handler func(ToolCallDeltaEvent)) func() {
return m.Subscribe(func(e Event) {
if tcd, ok := e.(ToolCallDeltaEvent); ok {
handler(tcd)
}
})
}
// OnToolCallEnd registers a handler that fires only for ToolCallEndEvent.
// This fires when tool argument streaming is complete, before the tool call
// is parsed and execution begins. Returns an unsubscribe function.
func (m *Kit) OnToolCallEnd(handler func(ToolCallEndEvent)) func() {
return m.Subscribe(func(e Event) {
if tce, ok := e.(ToolCallEndEvent); ok {
handler(tce)
}
})
}
// OnToolResult registers a handler that fires only for ToolResultEvent.
// Returns an unsubscribe function.
func (m *Kit) OnToolResult(handler func(ToolResultEvent)) func() {
+32
View File
@@ -100,6 +100,38 @@ func (m *Kit) bridgeExtensions(runner *extensions.Runner) {
})
}
// Tool call input streaming events — fire as the LLM generates tool arguments.
if runner.HasHandlers(extensions.ToolCallInputStart) {
m.Subscribe(func(e Event) {
if ev, ok := e.(ToolCallStartEvent); ok {
_, _ = runner.Emit(extensions.ToolCallInputStartEvent{
ToolCallID: ev.ToolCallID,
ToolName: ev.ToolName,
ToolKind: ev.ToolKind,
})
}
})
}
if runner.HasHandlers(extensions.ToolCallInputDelta) {
m.Subscribe(func(e Event) {
if ev, ok := e.(ToolCallDeltaEvent); ok {
_, _ = runner.Emit(extensions.ToolCallInputDeltaEvent{
ToolCallID: ev.ToolCallID,
Delta: ev.Delta,
})
}
})
}
if runner.HasHandlers(extensions.ToolCallInputEnd) {
m.Subscribe(func(e Event) {
if ev, ok := e.(ToolCallEndEvent); ok {
_, _ = runner.Emit(extensions.ToolCallInputEndEvent{
ToolCallID: ev.ToolCallID,
})
}
})
}
if runner.HasHandlers(extensions.AgentEnd) {
m.Subscribe(func(e Event) {
if ev, ok := e.(TurnEndEvent); ok {
+819 -83
View File
File diff suppressed because it is too large Load Diff
+225
View File
@@ -5,6 +5,8 @@ import (
"os"
"testing"
"github.com/spf13/viper"
kit "github.com/mark3labs/kit/pkg/kit"
)
@@ -54,6 +56,225 @@ func TestNewWithOptions(t *testing.T) {
}
}
// TestNewWithGenerationOptions verifies that the SDK-only generation
// parameter overrides on Options propagate all the way through to the
// agent without requiring any viper.Set workarounds in caller code.
func TestNewWithGenerationOptions(t *testing.T) {
if os.Getenv("ANTHROPIC_API_KEY") == "" {
t.Skip("Skipping test: ANTHROPIC_API_KEY not set")
}
ctx := context.Background()
// MaxTokens override — keep ThinkingLevel off so Anthropic's thinking
// budget doesn't auto-bump MaxTokens above what we configured.
t.Run("MaxTokens", func(t *testing.T) {
defer resetViper()
const want = 12345
host, err := kit.New(ctx, &kit.Options{
Model: "anthropic/claude-sonnet-4-5-20250929",
Quiet: true,
MaxTokens: want,
})
if err != nil {
t.Fatalf("Failed to create Kit: %v", err)
}
defer func() { _ = host.Close() }()
if got := host.MaxTokens(); got != want {
t.Errorf("Options.MaxTokens=%d did not propagate; Kit.MaxTokens()=%d", want, got)
}
if !viper.IsSet("max-tokens") {
t.Error("viper.IsSet(\"max-tokens\") should be true after MaxTokens override")
}
})
// ThinkingLevel override — verified via the public getter, which
// reads back the configured (not provider-derived) level.
t.Run("ThinkingLevel", func(t *testing.T) {
defer resetViper()
const want = "high"
host, err := kit.New(ctx, &kit.Options{
Model: "anthropic/claude-sonnet-4-5-20250929",
Quiet: true,
ThinkingLevel: want,
})
if err != nil {
t.Fatalf("Failed to create Kit: %v", err)
}
defer func() { _ = host.Close() }()
if got := host.GetThinkingLevel(); got != want {
t.Errorf("Options.ThinkingLevel=%q did not propagate; Kit.GetThinkingLevel()=%q", want, got)
}
})
// Temperature override — pointer semantics let callers distinguish
// "explicitly 0.0" from "unset", which we assert by pushing a distinct
// value and reading it back off viper's merged state.
t.Run("Temperature", func(t *testing.T) {
defer resetViper()
want := float32(0.12345)
host, err := kit.New(ctx, &kit.Options{
Model: "anthropic/claude-sonnet-4-5-20250929",
Quiet: true,
Temperature: &want,
})
if err != nil {
t.Fatalf("Failed to create Kit: %v", err)
}
defer func() { _ = host.Close() }()
if !viper.IsSet("temperature") {
t.Fatal("viper.IsSet(\"temperature\") should be true after Temperature override")
}
if got := float32(viper.GetFloat64("temperature")); got != want {
t.Errorf("Options.Temperature=%v did not propagate; viper=%v", want, got)
}
})
}
// TestNewPreservesIsSetSemantics verifies that creating a Kit WITHOUT
// populating the generation-param Options fields does NOT mark those
// keys as explicitly set in viper. This is the precedence contract
// that per-model defaults (ApplyModelSettings) and right-sizing
// (rightSizeMaxTokens) rely on.
//
// Previously setSDKDefaults() used viper.SetDefault() for every param,
// which caused viper.IsSet() to return true for all of them — silently
// suppressing per-model defaults and pinning max-tokens at 4096 even
// on models with much larger output limits.
func TestNewPreservesIsSetSemantics(t *testing.T) {
if os.Getenv("ANTHROPIC_API_KEY") == "" {
t.Skip("Skipping test: ANTHROPIC_API_KEY not set")
}
defer resetViper()
ctx := context.Background()
host, err := kit.New(ctx, &kit.Options{
Model: "anthropic/claude-sonnet-4-5-20250929",
Quiet: true,
NoSession: true,
SkipConfig: true, // isolate from any ~/.kit.yml values
})
if err != nil {
t.Fatalf("Failed to create Kit: %v", err)
}
defer func() { _ = host.Close() }()
// These keys must remain "unset" from viper's perspective so the
// downstream isExplicitlySet() checks allow per-model defaults to
// take effect.
checkKeys := []string{
"max-tokens",
"temperature",
"top-p",
"top-k",
"frequency-penalty",
"presence-penalty",
"thinking-level",
}
// With SkipConfig: true, InitConfig() is not invoked, so viper has
// no env-var bindings registered. Any IsSet() here would come purely
// from SDK-side SetDefault/Set calls — which is exactly what this
// test is guarding against.
for _, k := range checkKeys {
if viper.IsSet(k) {
t.Errorf("viper.IsSet(%q) == true when no Options field set it "+
"(SDK defaults must not corrupt IsSet semantics)", k)
}
}
}
// TestNewWithProviderOptions verifies that programmatic provider overrides
// (API key, URL) take effect without env vars or config files, and that
// Options.ProviderAPIKey *wins* over any pre-existing viper state.
func TestNewWithProviderOptions(t *testing.T) {
if os.Getenv("ANTHROPIC_API_KEY") == "" {
t.Skip("Skipping test: ANTHROPIC_API_KEY not set")
}
ctx := context.Background()
t.Run("succeeds with API key from Options", func(t *testing.T) {
defer resetViper()
apiKey := os.Getenv("ANTHROPIC_API_KEY")
host, err := kit.New(ctx, &kit.Options{
Model: "anthropic/claude-sonnet-4-5-20250929",
Quiet: true,
NoSession: true,
ProviderAPIKey: apiKey,
})
if err != nil {
t.Fatalf("Failed to create Kit with ProviderAPIKey option: %v", err)
}
defer func() { _ = host.Close() }()
if got := viper.GetString("provider-api-key"); got != apiKey {
t.Errorf("Options.ProviderAPIKey did not propagate to viper; got %q (len=%d)", got, len(got))
}
})
// Override precedence: even when viper already holds a different
// provider-api-key value (as it would if a config file or earlier
// Set() call populated one), Options.ProviderAPIKey must win.
t.Run("Options override beats pre-existing viper state", func(t *testing.T) {
defer resetViper()
viper.Set("provider-api-key", "sk-config-file-placeholder")
want := "sk-from-options-override"
// Use an OpenAI-flavored model so the validation path accepts
// the placeholder without attempting a real Anthropic handshake.
host, err := kit.New(ctx, &kit.Options{
Model: "openai/gpt-4o-mini",
Quiet: true,
NoSession: true,
NoExtensions: true,
DisableCoreTools: true,
ProviderAPIKey: want,
})
// Creation may still fail if the model registry is strict, but
// we only care that the override reached viper before any
// provider handshake happened.
if host != nil {
defer func() { _ = host.Close() }()
}
_ = err
if got := viper.GetString("provider-api-key"); got != want {
t.Errorf("Options.ProviderAPIKey did not override pre-existing viper value; got %q, want %q", got, want)
}
})
// ProviderURL override must also reach viper.
t.Run("ProviderURL propagates", func(t *testing.T) {
defer resetViper()
const want = "https://custom.example.com/v1"
host, err := kit.New(ctx, &kit.Options{
Model: "anthropic/claude-sonnet-4-5-20250929",
Quiet: true,
NoSession: true,
ProviderURL: want,
})
if err != nil {
t.Fatalf("Failed to create Kit with ProviderURL option: %v", err)
}
defer func() { _ = host.Close() }()
if got := viper.GetString("provider-url"); got != want {
t.Errorf("Options.ProviderURL did not propagate; got %q, want %q", got, want)
}
})
}
func TestSessionManagement(t *testing.T) {
if os.Getenv("ANTHROPIC_API_KEY") == "" {
t.Skip("Skipping test: ANTHROPIC_API_KEY not set")
@@ -81,3 +302,7 @@ func TestSessionManagement(t *testing.T) {
t.Error("Expected non-empty session ID")
}
}
// resetViper wipes viper's global state so a test case doesn't leak
// viper.Set() calls into the next one. Used via defer in subtests.
func resetViper() { viper.Reset() }
+56
View File
@@ -0,0 +1,56 @@
package kit_test
import (
"testing"
kit "github.com/mark3labs/kit/pkg/kit"
)
// TestMCPServerStatus_TypeSurface verifies the MCPServerStatus type is
// accessible and has the expected fields.
func TestMCPServerStatus_TypeSurface(t *testing.T) {
s := kit.MCPServerStatus{
Name: "test-server",
ToolCount: 5,
}
if s.Name != "test-server" {
t.Errorf("Expected Name 'test-server', got %q", s.Name)
}
if s.ToolCount != 5 {
t.Errorf("Expected ToolCount 5, got %d", s.ToolCount)
}
}
// TestMCPServerConfig_ForDynamicAdd verifies that MCPServerConfig can be
// constructed with the expected fields for dynamic server management.
func TestMCPServerConfig_ForDynamicAdd(t *testing.T) {
// Stdio server config.
stdio := kit.MCPServerConfig{
Command: []string{"npx", "-y", "@modelcontextprotocol/server-github"},
Environment: map[string]string{"GITHUB_TOKEN": "test-token"},
}
if len(stdio.Command) != 3 {
t.Errorf("Expected 3 command parts, got %d", len(stdio.Command))
}
if stdio.Environment["GITHUB_TOKEN"] != "test-token" {
t.Error("Expected GITHUB_TOKEN in environment")
}
// Remote server config.
remote := kit.MCPServerConfig{
URL: "https://mcp.example.com/sse",
Headers: []string{"Authorization: Bearer test"},
}
if remote.URL != "https://mcp.example.com/sse" {
t.Errorf("Unexpected URL: %s", remote.URL)
}
// Config with tool filtering.
filtered := kit.MCPServerConfig{
Command: []string{"some-server"},
AllowedTools: []string{"read", "write"},
}
if len(filtered.AllowedTools) != 2 {
t.Errorf("Expected 2 allowed tools, got %d", len(filtered.AllowedTools))
}
}
-5
View File
@@ -22,11 +22,6 @@ func GetLLMProviders() []string {
return models.GetGlobalRegistry().GetLLMProviders()
}
// Deprecated: Use GetLLMProviders instead.
func GetFantasyProviders() []string {
return GetLLMProviders()
}
// GetModelsForProvider returns all known models for a provider.
func GetModelsForProvider(provider string) (map[string]ModelInfo, error) {
return models.GetGlobalRegistry().GetModelsForProvider(provider)
+45 -45
View File
@@ -5,18 +5,18 @@ import (
"fmt"
"net"
"net/http"
"os/exec"
"runtime"
"sync"
"time"
)
// MCPAuthHandler handles OAuth authorization for MCP servers.
// Implementations control the user experience — opening a browser, showing a
// prompt, displaying a URL, etc.
// prompt, displaying a URL, posting to a message bus, etc.
//
// The default implementation ([DefaultMCPAuthHandler]) opens the system browser
// and starts a local HTTP callback server to receive the authorization code.
// [DefaultMCPAuthHandler] provides the transport mechanics (port reservation
// and callback server) but performs no user-facing I/O on its own; consumers
// wire presentation via [DefaultMCPAuthHandler.OnAuthURL] or implement
// MCPAuthHandler from scratch.
type MCPAuthHandler interface {
// RedirectURI returns the OAuth redirect URI that the callback server
// will listen on. This is called during MCP transport setup — before any
@@ -37,23 +37,44 @@ type MCPAuthHandler interface {
HandleAuth(ctx context.Context, serverName string, authURL string) (callbackURL string, err error)
}
// DefaultMCPAuthHandler opens the system browser and starts a local HTTP
// callback server to receive the OAuth authorization code. It eagerly reserves
// a TCP port on construction so [RedirectURI] is stable for the lifetime of
// the handler.
// DefaultMCPAuthHandler provides the transport mechanics of an OAuth flow —
// reserving a local TCP port and running a one-shot HTTP callback server —
// without making any user-experience decisions. It performs no browser opens,
// no printing, no TUI calls; consumers attach presentation by setting
// [DefaultMCPAuthHandler.OnAuthURL] or by wrapping the handler.
//
// Create instances with [NewDefaultMCPAuthHandler] (random port) or
// [NewDefaultMCPAuthHandlerWithPort] (explicit port).
// The handler eagerly reserves a TCP port on construction so [RedirectURI] is
// stable for the lifetime of the handler. Create instances with
// [NewDefaultMCPAuthHandler] (random port) or [NewDefaultMCPAuthHandlerWithPort]
// (explicit port). Always call [DefaultMCPAuthHandler.Close] when done to
// release the port.
type DefaultMCPAuthHandler struct {
listener net.Listener
port int
mu sync.Mutex // guards listener lifecycle
// OnAuthURL, if set, is invoked exactly once per [HandleAuth] call with
// the authorization URL the user must visit. This is where consumers
// plug in their UX: open a browser, print to stderr, post to a TUI
// stream, render a QR code, etc. The handler performs no I/O on the
// URL itself; if OnAuthURL is nil the URL is silently dropped and the
// user has no way to complete the flow.
//
// OnAuthURL is called synchronously before the handler blocks on the
// callback. It must not block indefinitely — long-running work should
// be dispatched to a goroutine.
OnAuthURL func(serverName, authURL string)
}
// NewDefaultMCPAuthHandler creates a handler that listens on a random
// available port on localhost. The port is reserved immediately so
// [RedirectURI] returns a stable value. Call [DefaultMCPAuthHandler.Close]
// when the handler is no longer needed to release the port.
//
// The returned handler has no OnAuthURL hook configured and will therefore
// appear to hang on HandleAuth until the context deadline fires. Set
// OnAuthURL before using the handler, or use a higher-level wrapper such
// as [CLIMCPAuthHandler].
func NewDefaultMCPAuthHandler() (*DefaultMCPAuthHandler, error) {
listener, err := net.Listen("tcp", "localhost:0")
if err != nil {
@@ -88,9 +109,9 @@ func (h *DefaultMCPAuthHandler) Port() int {
return h.port
}
// HandleAuth opens the system browser to authURL and waits for the OAuth
// callback on the local server. It returns the full callback URL including
// query parameters (code, state, etc.).
// HandleAuth invokes [OnAuthURL] with the authorization URL (if configured)
// and waits for the OAuth callback on the local server. It returns the full
// callback URL including query parameters (code, state, etc.).
//
// If the context has no deadline, a default 2-minute timeout is applied.
// The callback server is started for each HandleAuth call and shut down
@@ -136,19 +157,13 @@ func (h *DefaultMCPAuthHandler) HandleAuth(ctx context.Context, serverName strin
Handler: mux,
}
// Start serving on the pre-reserved listener. We need to create a new
// listener on the same port because http.Server.Serve takes ownership
// and closes the listener when done. The original listener is kept open
// to reserve the port; we create a second listener via SO_REUSEADDR
// semantics (Go's default on most platforms) or, more reliably, we
// temporarily release and re-acquire.
//
// Strategy: use the held listener directly for Serve. After Serve
// returns (due to Shutdown), re-acquire the listener to keep the port
// reserved for future HandleAuth calls.
// Start serving on the pre-reserved listener. http.Server.Serve takes
// ownership and closes the listener when Shutdown is called, so we
// re-acquire a fresh listener on the same port in the deferred cleanup
// below to keep the port reserved for subsequent HandleAuth calls.
h.mu.Lock()
serveListener := h.listener
h.listener = nil // Serve will close it
h.listener = nil
h.mu.Unlock()
if serveListener == nil {
@@ -184,10 +199,11 @@ func (h *DefaultMCPAuthHandler) HandleAuth(ctx context.Context, serverName strin
}
}()
// Open the system browser.
if err := openBrowser(authURL); err != nil {
// Browser open is best-effort; the user can still navigate manually.
_ = err
// Surface the authorization URL to the consumer. This is the single
// presentation seam: the SDK itself does not open browsers, print,
// or otherwise touch the user's environment.
if h.OnAuthURL != nil {
h.OnAuthURL(serverName, authURL)
}
// Wait for the callback, a server error, or context cancellation.
@@ -214,22 +230,6 @@ func (h *DefaultMCPAuthHandler) Close() error {
return nil
}
// openBrowser opens the default system browser to the given URL. This is a
// best-effort operation — errors are returned but callers typically ignore
// them since the user can navigate manually.
func openBrowser(url string) error {
switch runtime.GOOS {
case "linux":
return exec.Command("xdg-open", url).Start()
case "windows":
return exec.Command("rundll32", "url.dll,FileProtocolHandler", url).Start()
case "darwin":
return exec.Command("open", url).Start()
default:
return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
}
}
// oauthSuccessHTML is the HTML page returned to the browser after a
// successful OAuth callback.
const oauthSuccessHTML = `<!DOCTYPE html>
+47 -15
View File
@@ -5,32 +5,49 @@ import (
"fmt"
"io"
"os"
"os/exec"
"runtime"
)
// CLIMCPAuthHandler wraps a [DefaultMCPAuthHandler] and prints status messages
// to a writer (typically stderr) so the user knows what's happening during
// OAuth authorization. This is the handler used by the CLI/TUI binary.
// CLIMCPAuthHandler is the MCP OAuth handler for CLI/TUI consumers. It wraps
// a [DefaultMCPAuthHandler] and layers standard CLI behavior on top of the
// underlying transport mechanics:
//
// For TUI integration, set NotifyFunc to route messages through the TUI's
// event system instead of (or in addition to) the writer.
// - Opens the authorization URL in the system browser
// - Prints status messages (or routes them to a TUI via [NotifyFunc])
//
// Non-CLI consumers (web apps, daemons, custom TUIs) should not use this
// handler; implement [MCPAuthHandler] directly or configure a
// [DefaultMCPAuthHandler] with a custom OnAuthURL instead.
type CLIMCPAuthHandler struct {
inner *DefaultMCPAuthHandler
w io.Writer
// NotifyFunc, when set, is called with status messages instead of writing
// to the writer. This allows the TUI to display system messages in the
// chat stream. If nil, messages are written to w.
// NotifyFunc, when set, is called with status messages instead of
// writing to the writer. This allows the TUI to display system
// messages in the chat stream. If nil, messages are written to w.
NotifyFunc func(serverName, message string)
}
// NewCLIMCPAuthHandler creates a CLI auth handler that prints status messages
// to stderr and delegates the actual OAuth flow to a [DefaultMCPAuthHandler].
// to stderr, opens the authorization URL in the system browser, and delegates
// the callback-server mechanics to a [DefaultMCPAuthHandler].
func NewCLIMCPAuthHandler() (*CLIMCPAuthHandler, error) {
inner, err := NewDefaultMCPAuthHandler()
if err != nil {
return nil, err
}
return &CLIMCPAuthHandler{inner: inner, w: os.Stderr}, nil
h := &CLIMCPAuthHandler{inner: inner, w: os.Stderr}
// Wire the CLI presentation policy into the inner handler's hook.
// This is the one place in the codebase where OAuth triggers a
// browser open; the SDK core remains I/O-free.
inner.OnAuthURL = func(serverName, authURL string) {
h.notify(serverName, fmt.Sprintf("🔐 MCP server %q requires authentication. Opening browser...", serverName))
h.notify(serverName, fmt.Sprintf(" If the browser doesn't open, visit:\n %s", authURL))
// Browser open is best-effort; the user can still navigate manually.
_ = openBrowser(authURL)
}
return h, nil
}
// RedirectURI returns the OAuth redirect URI from the inner handler.
@@ -38,17 +55,15 @@ func (h *CLIMCPAuthHandler) RedirectURI() string {
return h.inner.RedirectURI()
}
// HandleAuth prints status messages and delegates to the inner handler.
// HandleAuth delegates to the inner handler (which invokes OnAuthURL, runs
// the callback server, and returns the full callback URL) and emits a final
// success or failure notification.
func (h *CLIMCPAuthHandler) HandleAuth(ctx context.Context, serverName string, authURL string) (string, error) {
h.notify(serverName, fmt.Sprintf("🔐 MCP server %q requires authentication. Opening browser...", serverName))
h.notify(serverName, fmt.Sprintf(" If the browser doesn't open, visit:\n %s", authURL))
callbackURL, err := h.inner.HandleAuth(ctx, serverName, authURL)
if err != nil {
h.notify(serverName, fmt.Sprintf("✗ Authentication failed for %q: %v", serverName, err))
return "", err
}
h.notify(serverName, fmt.Sprintf("✓ Authenticated with %q", serverName))
return callbackURL, nil
}
@@ -66,3 +81,20 @@ func (h *CLIMCPAuthHandler) notify(serverName, message string) {
}
_, _ = fmt.Fprintln(h.w, message)
}
// openBrowser opens the system default browser at url. Intentionally
// unexported: browser opening is CLI policy, not SDK surface. Consumers
// that need similar behavior for their own UX should bring their own
// helper (or use a third-party package like github.com/pkg/browser).
func openBrowser(url string) error {
switch runtime.GOOS {
case "linux":
return exec.Command("xdg-open", url).Start()
case "windows":
return exec.Command("rundll32", "url.dll,FileProtocolHandler", url).Start()
case "darwin":
return exec.Command("open", url).Start()
default:
return fmt.Errorf("unsupported platform: %s", runtime.GOOS)
}
}
+37 -6
View File
@@ -11,6 +11,8 @@ import (
"github.com/mark3labs/kit/internal/message"
"github.com/mark3labs/kit/internal/models"
"github.com/mark3labs/kit/internal/session"
"github.com/mark3labs/mcp-go/client/transport"
"github.com/mark3labs/mcp-go/server"
)
// ==== Message Types (internal/message/content.go) ====
@@ -128,9 +130,9 @@ type SpinnerFunc = agent.SpinnerFunc
// ==== LLM Types ====
//
// These are type aliases for the corresponding charm.land/fantasy types,
// giving them clean LLM-prefixed names without leaking the dependency name.
// SDK consumers can use these types without importing charm.land/fantasy directly.
// These are type aliases for the underlying LLM provider types, giving them
// clean LLM-prefixed names without leaking the dependency name. SDK consumers
// can use these types without importing the provider package directly.
// LLMMessage represents a message in an LLM conversation, carrying a role
// and a slice of typed content parts (text, tool calls, reasoning, etc.).
@@ -176,7 +178,7 @@ type LLMMessageRole = fantasy.MessageRole
// LLMFinishReason indicates why the LLM stopped generating.
type LLMFinishReason = fantasy.FinishReason
// LLM role constants mirror fantasy.MessageRole* values under clean LLM-prefixed names.
// LLM role constants mirror the provider's role values under clean LLM-prefixed names.
const (
// LLMRoleUser identifies a user message.
LLMRoleUser = fantasy.MessageRoleUser
@@ -189,11 +191,11 @@ const (
)
// NewLLMUserMessage constructs a user-role LLMMessage with optional file
// attachments. It is equivalent to fantasy.NewUserMessage.
// attachments.
var NewLLMUserMessage = fantasy.NewUserMessage
// NewLLMSystemMessage constructs a system-role LLMMessage from one or more
// prompt strings. It is equivalent to fantasy.NewSystemMessage.
// prompt strings.
var NewLLMSystemMessage = fantasy.NewSystemMessage
// ==== Compaction Types (internal/compaction/) ====
@@ -204,6 +206,35 @@ type CompactionResult = compaction.CompactionResult
// CompactionOptions configures compaction behaviour.
type CompactionOptions = compaction.CompactionOptions
// ==== MCP OAuth Types ====
// MCPServer is an in-process MCP server from the mcp-go library.
// Pass an instance to [Kit.AddInProcessMCPServer] or
// [Options.InProcessMCPServers] to register tools without spawning a
// subprocess or making network calls.
type MCPServer = server.MCPServer
// MCPTokenStore persists OAuth tokens for a single MCP server. Implementations
// must be safe for concurrent use.
//
// This is a type alias for the mcp-go transport.TokenStore interface. SDK
// consumers can implement this interface to provide custom storage backends
// (database, encrypted file, in-memory, etc.).
type MCPTokenStore = transport.TokenStore
// MCPToken represents an OAuth token for an MCP server, containing access
// and refresh tokens along with expiration metadata.
type MCPToken = transport.Token
// MCPTokenStoreFactory creates an [MCPTokenStore] for a given MCP server URL.
// It is called once per remote MCP server during connection setup.
type MCPTokenStoreFactory func(serverURL string) (MCPTokenStore, error)
// ErrMCPNoToken is the sentinel error that [MCPTokenStore] implementations
// should return from GetToken when no token is stored for the server.
// Callers can check for this with errors.Is.
var ErrMCPNoToken = transport.ErrNoToken
// ==== Constructor & Helper Functions ====
// ParseModelString parses a model string in "provider/model" format.
+7 -6
View File
@@ -77,8 +77,8 @@ func TestLLMRoleConstants(t *testing.T) {
}
}
// TestLLMMessageAlias verifies LLMMessage is a type alias for fantasy.Message
// and can be used interchangeably.
// TestLLMMessageAlias verifies LLMMessage is a type alias for the underlying
// LLM provider message type and can be used interchangeably.
func TestLLMMessageAlias(t *testing.T) {
// Construct an LLMMessage using alias types.
msg := kit.LLMMessage{
@@ -132,8 +132,8 @@ func TestNewLLMSystemMessage(t *testing.T) {
}
}
// TestLLMUsageAlias verifies LLMUsage is a type alias for fantasy.Usage
// and carries the correct fields.
// TestLLMUsageAlias verifies LLMUsage is a type alias for the underlying
// LLM provider usage type and carries the correct fields.
func TestLLMUsageAlias(t *testing.T) {
u := kit.LLMUsage{
InputTokens: 100,
@@ -150,7 +150,7 @@ func TestLLMUsageAlias(t *testing.T) {
t.Errorf("LLMUsage.TotalTokens = %d, want 150", u.TotalTokens)
}
// Verify JSON marshaling uses snake_case (inherited from fantasy.Usage tags).
// Verify JSON marshaling uses snake_case (inherited from the provider's tags).
data, err := json.Marshal(u)
if err != nil {
t.Fatalf("LLMUsage.MarshalJSON: %v", err)
@@ -165,7 +165,8 @@ func TestLLMUsageAlias(t *testing.T) {
}
}
// TestLLMFilePartAlias verifies LLMFilePart is a type alias for fantasy.FilePart.
// TestLLMFilePartAlias verifies LLMFilePart is a type alias for the underlying
// LLM provider file part type.
func TestLLMFilePartAlias(t *testing.T) {
fp := kit.LLMFilePart{
Filename: "screenshot.png",
+1 -1
View File
@@ -17,7 +17,7 @@ import time
import os
KIT_BIN = os.path.join(os.path.dirname(__file__), "..", "output", "kit")
MODEL = "opencode/kimi-k2.5"
MODEL = os.environ.get("MODEL", "opencode/kimi-k2.5")
CWD = os.path.expanduser("~")
TIMEOUT = 60 # seconds to wait for the prompt to complete

Some files were not shown because too many files have changed in this diff Show More