fix(mcp): preserve raw MCP payload in text artifacts
This commit is contained in:
parent
969ff6cfd5
commit
b6aa585dc0
2 changed files with 65 additions and 11 deletions
|
|
@ -272,14 +272,19 @@ func extractContentText(content []mcp.Content) string {
|
||||||
|
|
||||||
func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Content) *ToolResult {
|
func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Content) *ToolResult {
|
||||||
llmParts := make([]string, 0, len(content))
|
llmParts := make([]string, 0, len(content))
|
||||||
|
rawTextParts := make([]string, 0, len(content))
|
||||||
mediaRefs := make([]string, 0, len(content))
|
mediaRefs := make([]string, 0, len(content))
|
||||||
|
|
||||||
for _, c := range content {
|
for _, c := range content {
|
||||||
switch v := c.(type) {
|
switch v := c.(type) {
|
||||||
case *mcp.TextContent:
|
case *mcp.TextContent:
|
||||||
text := strings.TrimSpace(sanitizeToolLLMContent(v.Text))
|
rawText := strings.TrimSpace(v.Text)
|
||||||
if text != "" {
|
if rawText != "" {
|
||||||
llmParts = append(llmParts, text)
|
rawTextParts = append(rawTextParts, rawText)
|
||||||
|
}
|
||||||
|
safeText := strings.TrimSpace(sanitizeToolLLMContent(v.Text))
|
||||||
|
if safeText != "" {
|
||||||
|
llmParts = append(llmParts, safeText)
|
||||||
}
|
}
|
||||||
case *mcp.ImageContent:
|
case *mcp.ImageContent:
|
||||||
ref, note := t.storeBinaryContent(
|
ref, note := t.storeBinaryContent(
|
||||||
|
|
@ -312,10 +317,13 @@ func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Cont
|
||||||
case *mcp.ResourceLink:
|
case *mcp.ResourceLink:
|
||||||
llmParts = append(llmParts, summarizeResourceLink(v))
|
llmParts = append(llmParts, summarizeResourceLink(v))
|
||||||
case *mcp.EmbeddedResource:
|
case *mcp.EmbeddedResource:
|
||||||
ref, note := t.storeEmbeddedResource(ctx, v)
|
ref, note, rawText := t.storeEmbeddedResource(ctx, v)
|
||||||
if ref != "" {
|
if ref != "" {
|
||||||
mediaRefs = append(mediaRefs, ref)
|
mediaRefs = append(mediaRefs, ref)
|
||||||
}
|
}
|
||||||
|
if rawText != "" {
|
||||||
|
rawTextParts = append(rawTextParts, rawText)
|
||||||
|
}
|
||||||
if note != "" {
|
if note != "" {
|
||||||
llmParts = append(llmParts, note)
|
llmParts = append(llmParts, note)
|
||||||
}
|
}
|
||||||
|
|
@ -325,7 +333,8 @@ func (t *MCPTool) normalizeResultContent(ctx context.Context, content []mcp.Cont
|
||||||
}
|
}
|
||||||
|
|
||||||
forLLM := strings.Join(compactStrings(llmParts), "\n")
|
forLLM := strings.Join(compactStrings(llmParts), "\n")
|
||||||
if artifactResult := t.persistLargeTextArtifact(forLLM); artifactResult != nil {
|
rawText := strings.Join(compactStrings(rawTextParts), "\n")
|
||||||
|
if artifactResult := t.persistLargeTextArtifact(rawText); artifactResult != nil {
|
||||||
artifactResult.Media = mediaRefs
|
artifactResult.Media = mediaRefs
|
||||||
return artifactResult
|
return artifactResult
|
||||||
}
|
}
|
||||||
|
|
@ -381,27 +390,29 @@ func (t *MCPTool) persistLargeTextArtifact(text string) *ToolResult {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (t *MCPTool) storeEmbeddedResource(ctx context.Context, content *mcp.EmbeddedResource) (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.]", ""
|
||||||
}
|
}
|
||||||
|
|
||||||
resource := content.Resource
|
resource := content.Resource
|
||||||
if len(resource.Blob) > 0 {
|
if len(resource.Blob) > 0 {
|
||||||
return t.storeBinaryContent(
|
ref, note := t.storeBinaryContent(
|
||||||
ctx,
|
ctx,
|
||||||
"resource",
|
"resource",
|
||||||
normalizedMIMEType(resource.MIMEType),
|
normalizedMIMEType(resource.MIMEType),
|
||||||
resource.Blob,
|
resource.Blob,
|
||||||
content.Annotations,
|
content.Annotations,
|
||||||
)
|
)
|
||||||
|
return ref, note, ""
|
||||||
}
|
}
|
||||||
|
|
||||||
if strings.TrimSpace(resource.Text) != "" {
|
rawText := strings.TrimSpace(resource.Text)
|
||||||
return "", sanitizeToolLLMContent(resource.Text)
|
if rawText != "" {
|
||||||
|
return "", sanitizeToolLLMContent(resource.Text), rawText
|
||||||
}
|
}
|
||||||
|
|
||||||
return "", summarizeEmbeddedResource(content)
|
return "", summarizeEmbeddedResource(content), ""
|
||||||
}
|
}
|
||||||
|
|
||||||
func (t *MCPTool) storeBinaryContent(
|
func (t *MCPTool) storeBinaryContent(
|
||||||
|
|
|
||||||
|
|
@ -635,6 +635,49 @@ func TestMCPTool_Execute_LargeBase64TextIsOmittedFromContext(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestMCPTool_Execute_LargeBase64TextArtifactPreservesRawPayload(t *testing.T) {
|
||||||
|
workspace := t.TempDir()
|
||||||
|
largeBase64 := strings.Repeat("QUJD", 400)
|
||||||
|
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: largeBase64},
|
||||||
|
},
|
||||||
|
}, nil
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
mcpTool := NewMCPTool(manager, "test_server", &mcp.Tool{Name: "dump_payload"})
|
||||||
|
mcpTool.SetWorkspace(workspace)
|
||||||
|
mcpTool.SetMaxInlineTextRunes(32)
|
||||||
|
|
||||||
|
result := mcpTool.Execute(context.Background(), nil)
|
||||||
|
|
||||||
|
if !strings.Contains(result.ForLLM, "saved as a local artifact") {
|
||||||
|
t.Fatalf("expected artifact note, got %q", result.ForLLM)
|
||||||
|
}
|
||||||
|
if result.ForLLM == largeBase64OmittedMessage {
|
||||||
|
t.Fatalf("expected artifact note instead of sanitized base64 placeholder")
|
||||||
|
}
|
||||||
|
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), "]")
|
||||||
|
data, err := os.ReadFile(path)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("expected artifact file to be readable: %v", err)
|
||||||
|
}
|
||||||
|
if string(data) != largeBase64 {
|
||||||
|
t.Fatalf("expected artifact file contents to preserve raw MCP payload")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestMCPTool_Execute_LargeTextStoredAsArtifact(t *testing.T) {
|
func TestMCPTool_Execute_LargeTextStoredAsArtifact(t *testing.T) {
|
||||||
workspace := t.TempDir()
|
workspace := t.TempDir()
|
||||||
largeText := strings.Repeat("This is a large MCP text payload.\n", 800)
|
largeText := strings.Repeat("This is a large MCP text payload.\n", 800)
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue