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_CRON_EXEC_TIMEOUT_MINUTES=10`
|
||||
- `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
|
||||
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.SetWorkspace(agent.Workspace)
|
||||
mcpTool.SetMaxInlineTextRunes(al.cfg.Tools.MCP.GetMaxInlineTextChars())
|
||||
|
||||
if registerAsHidden {
|
||||
agent.Tools.RegisterHidden(mcpTool)
|
||||
|
|
|
|||
|
|
@ -943,10 +943,21 @@ type MCPServerConfig struct {
|
|||
type MCPConfig struct {
|
||||
ToolConfig ` envPrefix:"PICOCLAW_TOOLS_MCP_"`
|
||||
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 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) {
|
||||
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) {
|
||||
jsonData := `{
|
||||
"agents": {
|
||||
|
|
|
|||
|
|
@ -462,6 +462,7 @@ func DefaultConfig() *Config {
|
|||
UseBM25: true,
|
||||
UseRegex: false,
|
||||
},
|
||||
Max: DefaultMCPMaxInlineTextChars,
|
||||
Servers: map[string]MCPServerConfig{},
|
||||
},
|
||||
AppendFile: ToolConfig{
|
||||
|
|
|
|||
|
|
@ -6,8 +6,10 @@ import (
|
|||
"fmt"
|
||||
"hash/fnv"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
|
|
@ -26,18 +28,21 @@ type MCPManager interface {
|
|||
|
||||
// MCPTool wraps an MCP tool to implement the Tool interface
|
||||
type MCPTool struct {
|
||||
manager MCPManager
|
||||
serverName string
|
||||
tool *mcp.Tool
|
||||
mediaStore media.MediaStore
|
||||
manager MCPManager
|
||||
serverName string
|
||||
tool *mcp.Tool
|
||||
mediaStore media.MediaStore
|
||||
workspace string
|
||||
maxInlineTextRunes int
|
||||
}
|
||||
|
||||
// NewMCPTool creates a new MCP tool wrapper
|
||||
func NewMCPTool(manager MCPManager, serverName string, tool *mcp.Tool) *MCPTool {
|
||||
return &MCPTool{
|
||||
manager: manager,
|
||||
serverName: serverName,
|
||||
tool: tool,
|
||||
manager: manager,
|
||||
serverName: serverName,
|
||||
tool: tool,
|
||||
maxInlineTextRunes: maxMCPInlineTextRunes,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -45,6 +50,18 @@ func (t *MCPTool) SetMediaStore(store media.MediaStore) {
|
|||
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
|
||||
// as part of a tool/function identifier for downstream providers.
|
||||
// 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{
|
||||
ForLLM: strings.Join(compactStrings(llmParts), "\n"),
|
||||
ForLLM: forLLM,
|
||||
Media: mediaRefs,
|
||||
}
|
||||
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) {
|
||||
if content == nil || content.Resource == nil {
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
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