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/modelcontextprotocol/go-sdk/mcp"
|
||||||
|
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
"github.com/sipeed/picoclaw/pkg/media"
|
"github.com/sipeed/picoclaw/pkg/media"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
@ -352,13 +353,14 @@ func (t *MCPTool) persistLargeTextArtifact(text string) *ToolResult {
|
||||||
if limit <= 0 {
|
if limit <= 0 {
|
||||||
limit = maxMCPInlineTextRunes
|
limit = maxMCPInlineTextRunes
|
||||||
}
|
}
|
||||||
if text == "" || utf8.RuneCountInString(text) <= limit || t.workspace == "" {
|
size := utf8.RuneCountInString(text)
|
||||||
|
if text == "" || size <= limit || t.workspace == "" {
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
dir := filepath.Join(t.workspace, ".artifacts", "mcp")
|
dir := filepath.Join(t.workspace, ".artifacts", "mcp")
|
||||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||||
return nil
|
return t.largeTextArtifactFallback(text, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
pattern := fmt.Sprintf(
|
pattern := fmt.Sprintf(
|
||||||
|
|
@ -368,28 +370,44 @@ func (t *MCPTool) persistLargeTextArtifact(text string) *ToolResult {
|
||||||
)
|
)
|
||||||
tmpFile, err := os.CreateTemp(dir, pattern)
|
tmpFile, err := os.CreateTemp(dir, pattern)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return nil
|
return t.largeTextArtifactFallback(text, err)
|
||||||
}
|
}
|
||||||
path := tmpFile.Name()
|
path := tmpFile.Name()
|
||||||
if _, err = tmpFile.WriteString(text); err != nil {
|
if _, err = tmpFile.WriteString(text); err != nil {
|
||||||
_ = tmpFile.Close()
|
_ = tmpFile.Close()
|
||||||
_ = os.Remove(path)
|
_ = os.Remove(path)
|
||||||
return nil
|
return t.largeTextArtifactFallback(text, err)
|
||||||
}
|
}
|
||||||
if err = tmpFile.Close(); err != nil {
|
if err = tmpFile.Close(); err != nil {
|
||||||
_ = os.Remove(path)
|
_ = os.Remove(path)
|
||||||
return nil
|
return t.largeTextArtifactFallback(text, err)
|
||||||
}
|
}
|
||||||
|
|
||||||
return &ToolResult{
|
return &ToolResult{
|
||||||
ForLLM: fmt.Sprintf(
|
ForLLM: fmt.Sprintf(
|
||||||
"[MCP returned a large text result (%d chars); omitted from model context and saved as a local artifact.]",
|
"[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 + "]"},
|
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) {
|
func (t *MCPTool) storeEmbeddedResource(ctx context.Context, content *mcp.EmbeddedResource) (string, 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.]", ""
|
||||||
|
|
|
||||||
|
|
@ -749,3 +749,37 @@ func TestMCPTool_Execute_CustomInlineTextThreshold(t *testing.T) {
|
||||||
t.Fatalf("expected text to be omitted from ForLLM, got %q", result.ForLLM)
|
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