fix(mcp): avoid leaking large text when artifact persistence fails
This commit is contained in:
parent
b6aa585dc0
commit
cefec159ad
2 changed files with 58 additions and 6 deletions
|
|
@ -13,6 +13,7 @@ import (
|
|||
|
||||
"github.com/modelcontextprotocol/go-sdk/mcp"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/logger"
|
||||
"github.com/sipeed/picoclaw/pkg/media"
|
||||
)
|
||||
|
||||
|
|
@ -352,13 +353,14 @@ func (t *MCPTool) persistLargeTextArtifact(text string) *ToolResult {
|
|||
if limit <= 0 {
|
||||
limit = maxMCPInlineTextRunes
|
||||
}
|
||||
if text == "" || utf8.RuneCountInString(text) <= limit || t.workspace == "" {
|
||||
size := utf8.RuneCountInString(text)
|
||||
if text == "" || size <= limit || t.workspace == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
dir := filepath.Join(t.workspace, ".artifacts", "mcp")
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
return nil
|
||||
return t.largeTextArtifactFallback(text, err)
|
||||
}
|
||||
|
||||
pattern := fmt.Sprintf(
|
||||
|
|
@ -368,28 +370,44 @@ func (t *MCPTool) persistLargeTextArtifact(text string) *ToolResult {
|
|||
)
|
||||
tmpFile, err := os.CreateTemp(dir, pattern)
|
||||
if err != nil {
|
||||
return nil
|
||||
return t.largeTextArtifactFallback(text, err)
|
||||
}
|
||||
path := tmpFile.Name()
|
||||
if _, err = tmpFile.WriteString(text); err != nil {
|
||||
_ = tmpFile.Close()
|
||||
_ = os.Remove(path)
|
||||
return nil
|
||||
return t.largeTextArtifactFallback(text, err)
|
||||
}
|
||||
if err = tmpFile.Close(); err != nil {
|
||||
_ = os.Remove(path)
|
||||
return nil
|
||||
return t.largeTextArtifactFallback(text, err)
|
||||
}
|
||||
|
||||
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),
|
||||
size,
|
||||
),
|
||||
ArtifactTags: []string{"[file:" + path + "]"},
|
||||
}
|
||||
}
|
||||
|
||||
func (t *MCPTool) largeTextArtifactFallback(text string, err error) *ToolResult {
|
||||
size := utf8.RuneCountInString(text)
|
||||
logger.WarnCF("tool", "Failed to persist large MCP text artifact", map[string]any{
|
||||
"server": t.serverName,
|
||||
"tool": t.tool.Name,
|
||||
"chars": size,
|
||||
"error": err.Error(),
|
||||
})
|
||||
return &ToolResult{
|
||||
ForLLM: fmt.Sprintf(
|
||||
"[MCP returned a large text result (%d chars); omitted from model context because artifact persistence failed.]",
|
||||
size,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
func (t *MCPTool) storeEmbeddedResource(ctx context.Context, content *mcp.EmbeddedResource) (string, string, string) {
|
||||
if content == nil || content.Resource == nil {
|
||||
return "", "[MCP returned an embedded resource without data.]", ""
|
||||
|
|
|
|||
|
|
@ -749,3 +749,37 @@ func TestMCPTool_Execute_CustomInlineTextThreshold(t *testing.T) {
|
|||
t.Fatalf("expected text to be omitted from ForLLM, got %q", result.ForLLM)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMCPTool_Execute_LargeTextArtifactFailureStillOmitsContext(t *testing.T) {
|
||||
workspaceRoot := t.TempDir()
|
||||
workspaceFile := filepath.Join(workspaceRoot, "not-a-directory")
|
||||
if err := os.WriteFile(workspaceFile, []byte("x"), 0o600); err != nil {
|
||||
t.Fatalf("failed to create workspace file: %v", err)
|
||||
}
|
||||
|
||||
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(workspaceFile)
|
||||
|
||||
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, "artifact persistence failed") {
|
||||
t.Fatalf("expected persistence failure note, got %q", result.ForLLM)
|
||||
}
|
||||
if len(result.ArtifactTags) != 0 {
|
||||
t.Fatalf("expected no artifact tags on persistence failure, got %+v", result.ArtifactTags)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue