feat(mcp): store oversized text results as artifacts
This commit is contained in:
parent
bd56e10bb8
commit
49759e6209
7 changed files with 199 additions and 8 deletions
|
|
@ -528,6 +528,9 @@ For example:
|
||||||
- `PICOCLAW_TOOLS_EXEC_ENABLE_DENY_PATTERNS=false`
|
- `PICOCLAW_TOOLS_EXEC_ENABLE_DENY_PATTERNS=false`
|
||||||
- `PICOCLAW_TOOLS_CRON_EXEC_TIMEOUT_MINUTES=10`
|
- `PICOCLAW_TOOLS_CRON_EXEC_TIMEOUT_MINUTES=10`
|
||||||
- `PICOCLAW_TOOLS_MCP_ENABLED=true`
|
- `PICOCLAW_TOOLS_MCP_ENABLED=true`
|
||||||
|
- `PICOCLAW_TOOLS_MCP_MAX_INLINE_TEXT_CHARS=8192`
|
||||||
|
|
||||||
Note: Nested map-style config (for example `tools.mcp.servers.<name>.*`) is configured in `config.json` rather than
|
Note: Nested map-style config (for example `tools.mcp.servers.<name>.*`) is configured in `config.json` rather than
|
||||||
environment variables.
|
environment variables.
|
||||||
|
|
||||||
|
For MCP tools, `tools.mcp.max_inline_text_chars` controls how much text result is kept inline in model context. Above this threshold, PicoClaw saves the MCP text result as a local artifact in the agent workspace and gives the model a short note plus a structured `[file:...]` artifact path instead of injecting the full payload into context.
|
||||||
|
|
|
||||||
|
|
@ -126,6 +126,8 @@ func (al *AgentLoop) ensureMCPInitialized(ctx context.Context) error {
|
||||||
}
|
}
|
||||||
|
|
||||||
mcpTool := tools.NewMCPTool(mcpManager, serverName, tool)
|
mcpTool := tools.NewMCPTool(mcpManager, serverName, tool)
|
||||||
|
mcpTool.SetWorkspace(agent.Workspace)
|
||||||
|
mcpTool.SetMaxInlineTextRunes(al.cfg.Tools.MCP.GetMaxInlineTextChars())
|
||||||
|
|
||||||
if registerAsHidden {
|
if registerAsHidden {
|
||||||
agent.Tools.RegisterHidden(mcpTool)
|
agent.Tools.RegisterHidden(mcpTool)
|
||||||
|
|
|
||||||
|
|
@ -943,10 +943,21 @@ type MCPServerConfig struct {
|
||||||
type MCPConfig struct {
|
type MCPConfig struct {
|
||||||
ToolConfig ` envPrefix:"PICOCLAW_TOOLS_MCP_"`
|
ToolConfig ` envPrefix:"PICOCLAW_TOOLS_MCP_"`
|
||||||
Discovery ToolDiscoveryConfig ` json:"discovery"`
|
Discovery ToolDiscoveryConfig ` json:"discovery"`
|
||||||
|
// Max controls how much MCP text stays inline before it is saved as an artifact.
|
||||||
|
Max int `json:"max_inline_text_chars,omitempty" env:"PICOCLAW_TOOLS_MCP_MAX_INLINE_TEXT_CHARS"`
|
||||||
// Servers is a map of server name to server configuration
|
// Servers is a map of server name to server configuration
|
||||||
Servers map[string]MCPServerConfig `json:"servers,omitempty"`
|
Servers map[string]MCPServerConfig `json:"servers,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const DefaultMCPMaxInlineTextChars = 16 * 1024
|
||||||
|
|
||||||
|
func (c *MCPConfig) GetMaxInlineTextChars() int {
|
||||||
|
if c.Max > 0 {
|
||||||
|
return c.Max
|
||||||
|
}
|
||||||
|
return DefaultMCPMaxInlineTextChars
|
||||||
|
}
|
||||||
|
|
||||||
func LoadConfig(path string) (*Config, error) {
|
func LoadConfig(path string) (*Config, error) {
|
||||||
logger.Debugf("loading config from %s", path)
|
logger.Debugf("loading config from %s", path)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -198,6 +198,41 @@ func TestAgentConfig_FullParse(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestDefaultConfig_MCPMaxInlineTextChars(t *testing.T) {
|
||||||
|
cfg := DefaultConfig()
|
||||||
|
if cfg.Tools.MCP.GetMaxInlineTextChars() != DefaultMCPMaxInlineTextChars {
|
||||||
|
t.Fatalf(
|
||||||
|
"DefaultConfig().Tools.MCP.GetMaxInlineTextChars() = %d, want %d",
|
||||||
|
cfg.Tools.MCP.GetMaxInlineTextChars(),
|
||||||
|
DefaultMCPMaxInlineTextChars,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestLoadConfig_MCPMaxInlineTextChars(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
configPath := filepath.Join(dir, "config.json")
|
||||||
|
raw := `{
|
||||||
|
"tools": {
|
||||||
|
"mcp": {
|
||||||
|
"enabled": true,
|
||||||
|
"max_inline_text_chars": 2048
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}`
|
||||||
|
if err := os.WriteFile(configPath, []byte(raw), 0o644); err != nil {
|
||||||
|
t.Fatalf("WriteFile(configPath): %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
cfg, err := LoadConfig(configPath)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("LoadConfig() error: %v", err)
|
||||||
|
}
|
||||||
|
if got := cfg.Tools.MCP.GetMaxInlineTextChars(); got != 2048 {
|
||||||
|
t.Fatalf("cfg.Tools.MCP.GetMaxInlineTextChars() = %d, want 2048", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestConfig_BackwardCompat_NoAgentsList(t *testing.T) {
|
func TestConfig_BackwardCompat_NoAgentsList(t *testing.T) {
|
||||||
jsonData := `{
|
jsonData := `{
|
||||||
"agents": {
|
"agents": {
|
||||||
|
|
|
||||||
|
|
@ -462,6 +462,7 @@ func DefaultConfig() *Config {
|
||||||
UseBM25: true,
|
UseBM25: true,
|
||||||
UseRegex: false,
|
UseRegex: false,
|
||||||
},
|
},
|
||||||
|
Max: DefaultMCPMaxInlineTextChars,
|
||||||
Servers: map[string]MCPServerConfig{},
|
Servers: map[string]MCPServerConfig{},
|
||||||
},
|
},
|
||||||
AppendFile: ToolConfig{
|
AppendFile: ToolConfig{
|
||||||
|
|
|
||||||
|
|
@ -6,8 +6,10 @@ import (
|
||||||
"fmt"
|
"fmt"
|
||||||
"hash/fnv"
|
"hash/fnv"
|
||||||
"os"
|
"os"
|
||||||
|
"path/filepath"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
|
||||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
|
||||||
|
|
@ -30,6 +32,8 @@ type MCPTool struct {
|
||||||
serverName string
|
serverName string
|
||||||
tool *mcp.Tool
|
tool *mcp.Tool
|
||||||
mediaStore media.MediaStore
|
mediaStore media.MediaStore
|
||||||
|
workspace string
|
||||||
|
maxInlineTextRunes int
|
||||||
}
|
}
|
||||||
|
|
||||||
// NewMCPTool creates a new MCP tool wrapper
|
// NewMCPTool creates a new MCP tool wrapper
|
||||||
|
|
@ -38,6 +42,7 @@ func NewMCPTool(manager MCPManager, serverName string, tool *mcp.Tool) *MCPTool
|
||||||
manager: manager,
|
manager: manager,
|
||||||
serverName: serverName,
|
serverName: serverName,
|
||||||
tool: tool,
|
tool: tool,
|
||||||
|
maxInlineTextRunes: maxMCPInlineTextRunes,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -45,6 +50,18 @@ func (t *MCPTool) SetMediaStore(store media.MediaStore) {
|
||||||
t.mediaStore = store
|
t.mediaStore = store
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (t *MCPTool) SetWorkspace(workspace string) {
|
||||||
|
t.workspace = strings.TrimSpace(workspace)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (t *MCPTool) SetMaxInlineTextRunes(limit int) {
|
||||||
|
if limit > 0 {
|
||||||
|
t.maxInlineTextRunes = limit
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
const maxMCPInlineTextRunes = 16 * 1024
|
||||||
|
|
||||||
// sanitizeIdentifierComponent normalizes a string so it can be safely used
|
// sanitizeIdentifierComponent normalizes a string so it can be safely used
|
||||||
// as part of a tool/function identifier for downstream providers.
|
// as part of a tool/function identifier for downstream providers.
|
||||||
// It:
|
// It:
|
||||||
|
|
@ -307,13 +324,63 @@ func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Cont
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
forLLM := strings.Join(compactStrings(llmParts), "\n")
|
||||||
|
if artifactResult := t.persistLargeTextArtifact(forLLM); artifactResult != nil {
|
||||||
|
artifactResult.Media = mediaRefs
|
||||||
|
return artifactResult
|
||||||
|
}
|
||||||
|
|
||||||
result := &ToolResult{
|
result := &ToolResult{
|
||||||
ForLLM: strings.Join(compactStrings(llmParts), "\n"),
|
ForLLM: forLLM,
|
||||||
Media: mediaRefs,
|
Media: mediaRefs,
|
||||||
}
|
}
|
||||||
return result
|
return result
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (t *MCPTool) persistLargeTextArtifact(text string) *ToolResult {
|
||||||
|
text = strings.TrimSpace(text)
|
||||||
|
limit := t.maxInlineTextRunes
|
||||||
|
if limit <= 0 {
|
||||||
|
limit = maxMCPInlineTextRunes
|
||||||
|
}
|
||||||
|
if text == "" || utf8.RuneCountInString(text) <= limit || t.workspace == "" {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
dir := filepath.Join(t.workspace, ".artifacts", "mcp")
|
||||||
|
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
pattern := fmt.Sprintf(
|
||||||
|
"%s_%s_*.txt",
|
||||||
|
sanitizeIdentifierComponent(t.serverName),
|
||||||
|
sanitizeIdentifierComponent(t.tool.Name),
|
||||||
|
)
|
||||||
|
tmpFile, err := os.CreateTemp(dir, pattern)
|
||||||
|
if err != nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
path := tmpFile.Name()
|
||||||
|
if _, err = tmpFile.WriteString(text); err != nil {
|
||||||
|
_ = tmpFile.Close()
|
||||||
|
_ = os.Remove(path)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
if err = tmpFile.Close(); err != nil {
|
||||||
|
_ = os.Remove(path)
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
return &ToolResult{
|
||||||
|
ForLLM: fmt.Sprintf(
|
||||||
|
"[MCP returned a large text result (%d chars); omitted from model context and saved as a local artifact.]",
|
||||||
|
utf8.RuneCountInString(text),
|
||||||
|
),
|
||||||
|
ArtifactTags: []string{"[file:" + path + "]"},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (t *MCPTool) storeEmbeddedResource(ctx context.Context, content *mcp.EmbeddedResource) (string, string) {
|
func (t *MCPTool) storeEmbeddedResource(ctx context.Context, content *mcp.EmbeddedResource) (string, string) {
|
||||||
if content == nil || content.Resource == nil {
|
if content == nil || content.Resource == nil {
|
||||||
return "", "[MCP returned an embedded resource without data.]"
|
return "", "[MCP returned an embedded resource without data.]"
|
||||||
|
|
|
||||||
|
|
@ -634,3 +634,75 @@ func TestMCPTool_Execute_LargeBase64TextIsOmittedFromContext(t *testing.T) {
|
||||||
t.Fatalf("expected sanitized large base64 note, got %q", result.ForLLM)
|
t.Fatalf("expected sanitized large base64 note, got %q", result.ForLLM)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestMCPTool_Execute_LargeTextStoredAsArtifact(t *testing.T) {
|
||||||
|
workspace := t.TempDir()
|
||||||
|
largeText := strings.Repeat("This is a large MCP text payload.\n", 800)
|
||||||
|
manager := &MockMCPManager{
|
||||||
|
callToolFunc: func(ctx context.Context, serverName, toolName string, arguments map[string]any) (*mcp.CallToolResult, error) {
|
||||||
|
return &mcp.CallToolResult{
|
||||||
|
Content: []mcp.Content{
|
||||||
|
&mcp.TextContent{Text: largeText},
|
||||||
|
},
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
mcpTool := NewMCPTool(manager, "test_server", &mcp.Tool{Name: "dump_payload"})
|
||||||
|
mcpTool.SetWorkspace(workspace)
|
||||||
|
|
||||||
|
result := mcpTool.Execute(context.Background(), nil)
|
||||||
|
|
||||||
|
if strings.Contains(result.ForLLM, "This is a large MCP text payload") {
|
||||||
|
t.Fatalf("expected large MCP text to be omitted from ForLLM, got %q", result.ForLLM)
|
||||||
|
}
|
||||||
|
if !strings.Contains(result.ForLLM, "saved as a local artifact") {
|
||||||
|
t.Fatalf("expected artifact note, got %q", result.ForLLM)
|
||||||
|
}
|
||||||
|
if len(result.ArtifactTags) != 1 {
|
||||||
|
t.Fatalf("expected 1 artifact tag, got %d", len(result.ArtifactTags))
|
||||||
|
}
|
||||||
|
tag := result.ArtifactTags[0]
|
||||||
|
const prefix = "[file:"
|
||||||
|
if !strings.HasPrefix(tag, prefix) || !strings.HasSuffix(tag, "]") {
|
||||||
|
t.Fatalf("expected file artifact tag, got %q", tag)
|
||||||
|
}
|
||||||
|
path := strings.TrimSuffix(strings.TrimPrefix(tag, prefix), "]")
|
||||||
|
if !strings.HasPrefix(path, workspace) {
|
||||||
|
t.Fatalf("expected artifact inside workspace, got %q", path)
|
||||||
|
}
|
||||||
|
data, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("expected artifact file to be readable: %v", err)
|
||||||
|
}
|
||||||
|
if string(data) != strings.TrimSpace(largeText) {
|
||||||
|
t.Fatalf("expected artifact file contents to match source text")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMCPTool_Execute_CustomInlineTextThreshold(t *testing.T) {
|
||||||
|
workspace := t.TempDir()
|
||||||
|
text := strings.Repeat("small custom threshold text\n", 20)
|
||||||
|
manager := &MockMCPManager{
|
||||||
|
callToolFunc: func(ctx context.Context, serverName, toolName string, arguments map[string]any) (*mcp.CallToolResult, error) {
|
||||||
|
return &mcp.CallToolResult{
|
||||||
|
Content: []mcp.Content{
|
||||||
|
&mcp.TextContent{Text: text},
|
||||||
|
},
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
mcpTool := NewMCPTool(manager, "test_server", &mcp.Tool{Name: "dump_payload"})
|
||||||
|
mcpTool.SetWorkspace(workspace)
|
||||||
|
mcpTool.SetMaxInlineTextRunes(32)
|
||||||
|
|
||||||
|
result := mcpTool.Execute(context.Background(), nil)
|
||||||
|
|
||||||
|
if len(result.ArtifactTags) != 1 {
|
||||||
|
t.Fatalf("expected custom threshold to persist artifact, got %+v", result)
|
||||||
|
}
|
||||||
|
if strings.Contains(result.ForLLM, "small custom threshold text") {
|
||||||
|
t.Fatalf("expected text to be omitted from ForLLM, got %q", result.ForLLM)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue