diff --git a/claudetool/output_iframe.go b/claudetool/output_iframe.go index ff7d3c8f09a59d8a0dcb3576f71674b8e4929259..a4eb28458a06665398826c7d841716761c9b11ad 100644 --- a/claudetool/output_iframe.go +++ b/claudetool/output_iframe.go @@ -3,16 +3,26 @@ package claudetool import ( "context" "encoding/json" + "os" + "path/filepath" + "strings" "shelley.exe.dev/llm" ) // OutputIframeTool displays sandboxed HTML content to the user. -var OutputIframeTool = &llm.Tool{ - Name: outputIframeName, - Description: outputIframeDescription, - InputSchema: llm.MustSchema(outputIframeInputSchema), - Run: outputIframeRun, +// It requires a MutableWorkingDir to resolve relative file paths. +type OutputIframeTool struct { + WorkingDir *MutableWorkingDir +} + +func (t *OutputIframeTool) Tool() *llm.Tool { + return &llm.Tool{ + Name: outputIframeName, + Description: outputIframeDescription, + InputSchema: llm.MustSchema(outputIframeInputSchema), + Run: t.Run, + } } const ( @@ -34,50 +44,217 @@ Good uses: - SVG graphics The HTML should be self-contained. You can include inline \n") + } + + // Inject JS files as script tags + for _, f := range jsFiles { + injection.WriteString("\n") + } + + // Insert after or at the beginning + injectionStr := injection.String() + if idx := strings.Index(strings.ToLower(html), ""); idx != -1 { + return html[:idx+6] + "\n" + injectionStr + html[idx+6:] + } + if idx := strings.Index(strings.ToLower(html), ""); idx != -1 { + return html[:idx+6] + "\n\n" + injectionStr + "\n" + html[idx+6:] + } + // No head or html tag, just prepend + return injectionStr + html +} + +// escapeJSString escapes a string for use in a JavaScript string literal. +func escapeJSString(s string) string { + var b strings.Builder + for _, r := range s { + switch r { + case '\\': + b.WriteString("\\\\") + case '"': + b.WriteString("\\\"") + case '\n': + b.WriteString("\\n") + case '\r': + b.WriteString("\\r") + case '\t': + b.WriteString("\\t") + default: + b.WriteRune(r) + } + } + return b.String() } -func outputIframeRun(ctx context.Context, m json.RawMessage) llm.ToolOut { +func (t *OutputIframeTool) Run(ctx context.Context, m json.RawMessage) llm.ToolOut { var input struct { - HTML string `json:"html"` - Title string `json:"title"` + Path string `json:"path"` + Title string `json:"title"` + Files map[string]string `json:"files"` } if err := json.Unmarshal(m, &input); err != nil { return llm.ErrorToolOut(err) } - if input.HTML == "" { - return llm.ErrorfToolOut("html content is required") + if input.Path == "" { + return llm.ErrorfToolOut("path is required") } + // Resolve the path relative to working directory + path := input.Path + if !filepath.IsAbs(path) { + path = filepath.Join(t.WorkingDir.Get(), path) + } + + // Read the main HTML file + data, err := os.ReadFile(path) + if err != nil { + return llm.ErrorfToolOut("failed to read file: %v", err) + } + + // Read additional files + var embeddedFiles []EmbeddedFile + for name, filePath := range input.Files { + // Resolve relative paths + if !filepath.IsAbs(filePath) { + filePath = filepath.Join(t.WorkingDir.Get(), filePath) + } + content, err := os.ReadFile(filePath) + if err != nil { + return llm.ErrorfToolOut("failed to read file %q: %v", name, err) + } + embeddedFiles = append(embeddedFiles, EmbeddedFile{ + Name: name, + Path: input.Files[name], // Original path for download + Content: string(content), + Type: detectFileType(name), + }) + } + + // Inject files into the HTML for iframe display + html := injectFiles(string(data), embeddedFiles) + display := OutputIframeDisplay{ - Type: "output_iframe", - HTML: input.HTML, - Title: input.Title, + Type: "output_iframe", + HTML: html, + Title: input.Title, + Filename: filepath.Base(input.Path), + Files: embeddedFiles, } return llm.ToolOut{ diff --git a/claudetool/output_iframe_test.go b/claudetool/output_iframe_test.go index 7d478e1fb8c9174c4bb21d5af90e6ee12468415e..c2c2b981335d558e7301b558b455ac66409ab1b0 100644 --- a/claudetool/output_iframe_test.go +++ b/claudetool/output_iframe_test.go @@ -3,45 +3,150 @@ package claudetool import ( "context" "encoding/json" + "os" + "path/filepath" + "strings" "testing" ) func TestOutputIframeRun(t *testing.T) { + // Create a temp directory for test files + tmpDir := t.TempDir() + + // Create test HTML files + htmlFile := filepath.Join(tmpDir, "test.html") + if err := os.WriteFile(htmlFile, []byte("

Hello

"), 0o644); err != nil { + t.Fatal(err) + } + + chartFile := filepath.Join(tmpDir, "chart.html") + if err := os.WriteFile(chartFile, []byte("
Chart
"), 0o644); err != nil { + t.Fatal(err) + } + + // Create a data file + dataFile := filepath.Join(tmpDir, "data.json") + if err := os.WriteFile(dataFile, []byte(`{"values": [1, 2, 3]}`), 0o644); err != nil { + t.Fatal(err) + } + + // Create a CSS file + cssFile := filepath.Join(tmpDir, "styles.css") + if err := os.WriteFile(cssFile, []byte("body { color: red; }"), 0o644); err != nil { + t.Fatal(err) + } + + workingDir := &MutableWorkingDir{} + workingDir.Set(tmpDir) + + tool := &OutputIframeTool{WorkingDir: workingDir} + tests := []struct { - name string - input map[string]any - wantErr bool - wantTitle string + name string + input map[string]any + wantErr bool + wantTitle string + wantFilename string + wantFiles int + checkHTML func(html string) bool }{ { - name: "basic html", + name: "basic html file", input: map[string]any{ - "html": "

Hello

", + "path": "test.html", }, - wantErr: false, - wantTitle: "", + wantErr: false, + wantTitle: "", + wantFilename: "test.html", + wantFiles: 0, }, { name: "html with title", input: map[string]any{ - "html": "
Chart
", + "path": "chart.html", "title": "My Chart", }, - wantErr: false, - wantTitle: "My Chart", + wantErr: false, + wantTitle: "My Chart", + wantFilename: "chart.html", + wantFiles: 0, + }, + { + name: "html with data file", + input: map[string]any{ + "path": "chart.html", + "title": "Chart with Data", + "files": map[string]any{ + "data.json": "data.json", + }, + }, + wantErr: false, + wantTitle: "Chart with Data", + wantFilename: "chart.html", + wantFiles: 1, + checkHTML: func(html string) bool { + return strings.Contains(html, "window.__FILES__") && + strings.Contains(html, "data.json") + }, + }, + { + name: "html with multiple files", + input: map[string]any{ + "path": "chart.html", + "title": "Styled Chart", + "files": map[string]any{ + "data.json": "data.json", + "styles.css": "styles.css", + }, + }, + wantErr: false, + wantTitle: "Styled Chart", + wantFilename: "chart.html", + wantFiles: 2, + checkHTML: func(html string) bool { + return strings.Contains(html, "window.__FILES__") && + strings.Contains(html, "