diff --git a/cmd/chat_tools.go b/cmd/chat_tools.go index 0e6bc559..e792d939 100644 --- a/cmd/chat_tools.go +++ b/cmd/chat_tools.go @@ -107,6 +107,7 @@ func optionalTools() []tool.Tool { tool.DevEnvTool{}, tool.ProjectVerifyTool{}, tool.AppVerifyTool{}, + tool.GenerateMediaTool{}, tool.DependencyAuditTool{}, tool.GitHubTool{}, &tool.PRGeneratorTool{}, diff --git a/internal/daemon/telegram.go b/internal/daemon/telegram.go index 2c0a2c7d..a816de47 100644 --- a/internal/daemon/telegram.go +++ b/internal/daemon/telegram.go @@ -10,8 +10,12 @@ import ( "log/slog" "net/http" "net/url" + "os" + "path/filepath" "strings" "time" + + "github.com/GrayCodeAI/hawk/internal/stt" ) // TelegramGateway connects hawk to a Telegram bot. @@ -46,6 +50,19 @@ type TelegramMessage struct { From struct { Username string `json:"username"` } `json:"from"` + // Voice is set for Telegram voice notes; Audio for audio-file attachments. + // When either is present and an STT transcriber is installed, the audio is + // transcribed and the text is replaced with the transcript so the agent + // can answer it. + Voice *telegramAudio `json:"voice,omitempty"` + Audio *telegramAudio `json:"audio,omitempty"` +} + +// telegramAudio describes a Telegram voice/audio attachment. +type telegramAudio struct { + FileID string `json:"file_id"` + Duration int `json:"duration"` + MimeType string `json:"mime_type"` } // NewTelegramGateway creates a gateway with the given bot token. The authorizer @@ -182,8 +199,22 @@ func (tg *TelegramGateway) handleMessage(ctx context.Context, msg *TelegramMessa return } + // Transcribe voice/audio attachments when an STT backend is installed. + prompt := msg.Text + if audio := msg.Voice; audio != nil { + prompt = tg.transcribeAudio(ctx, msg.Chat.ID, audio) + if prompt == "" { + return // error already replied + } + } else if audio := msg.Audio; audio != nil { + prompt = tg.transcribeAudio(ctx, msg.Chat.ID, audio) + if prompt == "" { + return + } + } + // Forward to hawk daemon - response, err := tg.forwardToHawk(ctx, msg.Text) + response, err := tg.forwardToHawk(ctx, prompt) if err != nil { response = fmt.Sprintf("Error: %v", err) } @@ -196,6 +227,83 @@ func (tg *TelegramGateway) handleMessage(ctx context.Context, msg *TelegramMessa tg.reply(ctx, msg.Chat.ID, response) } +// transcribeAudio downloads a Telegram voice/audio attachment, transcribes it +// via the installed STT engine, and returns the transcript prefixed with +// a marker. On any error it replies to the user with the problem and returns "". +func (tg *TelegramGateway) transcribeAudio(ctx context.Context, chatID int64, a *telegramAudio) string { + if !stt.Enabled() { + return "" + } + + // Resolve file_id to a download URL via Telegram's bot API. + filePath, err := tg.getAudioFilePath(ctx, a.FileID) + if err != nil { + tg.reply(ctx, chatID, fmt.Sprintf("Audio resolve failed: %v", err)) + return "" + } + downloadURL := fmt.Sprintf("https://api.telegram.org/file/bot%s/%s", tg.Token, filePath) + + localPath, err := stt.DownloadAttachment(ctx, tg.client, downloadURL, "", safeFileName(a.FileID, a.MimeType)) + if err != nil { + tg.reply(ctx, chatID, fmt.Sprintf("Audio download failed: %v", err)) + return "" + } + defer func() { + // Clean up the temp dir containing the downloaded file. + _ = os.RemoveAll(filepath.Dir(localPath)) + }() + + text, err := stt.Transcribe(ctx, localPath, "") + if err != nil { + tg.reply(ctx, chatID, fmt.Sprintf("Transcription failed: %v", err)) + return "" + } + return "[Voice transcript] " + text +} + +// getAudioFilePath resolves a Telegram file_id to a file path on the bot API +// server so the message can be downloaded via the file URL. +func (tg *TelegramGateway) getAudioFilePath(ctx context.Context, fileID string) (string, error) { + req, err := http.NewRequestWithContext(ctx, http.MethodGet, + "https://api.telegram.org/bot"+tg.Token+"/getFile?file_id="+url.QueryEscape(fileID), nil) + if err != nil { + return "", err + } + resp, err := tg.client.Do(req) + if err != nil { + return "", err + } + defer func() { _ = resp.Body.Close() }() + body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<16)) + if err != nil { + return "", err + } + var result struct { + OK bool `json:"ok"` + Result struct { + FilePath string `json:"file_path"` + } `json:"result"` + } + if err := json.Unmarshal(body, &result); err != nil { + return "", err + } + if !result.OK { + return "", fmt.Errorf("telegram getFile returned not OK") + } + return result.Result.FilePath, nil +} + +// safeFileName builds a human-readable filename for a downloaded Telegram audio +// file using the file_id prefix and the MIME-suggested extension. +func safeFileName(fileID, mimeType string) string { + ext := stt.ExtensionForMedia(fileID, mimeType) + prefix := fileID + if len(prefix) > 16 { + prefix = prefix[:16] + } + return prefix + ext +} + // reply sends text and logs (rather than swallows) any delivery failure. func (tg *TelegramGateway) reply(ctx context.Context, chatID int64, text string) { if err := tg.sendMessage(ctx, chatID, text); err != nil { diff --git a/internal/daemon/telegram_test.go b/internal/daemon/telegram_test.go index 182e8c19..3697c050 100644 --- a/internal/daemon/telegram_test.go +++ b/internal/daemon/telegram_test.go @@ -155,3 +155,33 @@ func jsonResp(body string) *http.Response { Header: http.Header{"Content-Type": []string{"application/json"}}, } } + +func TestTelegramSafeFileName(t *testing.T) { + if got := safeFileName("AAabcdef0123456789xyz", "audio/ogg"); got != "AAabcdef01234567.ogg" { + t.Fatalf("safeFileName = %q", got) + } + if got := safeFileName("f", "audio/mpeg"); got != "f.mp3" { + t.Fatalf("safeFileName short id = %q", got) + } +} + +func TestTelegramVoiceWithoutTranscriberFallsBackToText(t *testing.T) { + hawk := newIPv4TelegramServer(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + writeJSON(w, http.StatusOK, ChatResponse{Response: "hawk-reply"}) + })) + defer hawk.Close() + + // No STT transcriber installed: a voice message must fall back to the + // existing text path (transcribeAudio returns "" with no reply). + tg := newTelegramGatewayFromConfig(TelegramConfig{Token: "tok", AllowList: []string{"user"}}, hawk.URL, "k") + tg.handleMessage(context.Background(), &TelegramMessage{ + Text: "hello", + Chat: struct { + ID int64 `json:"id"` + }{ID: 1}, + From: struct { + Username string `json:"username"` + }{Username: "user"}, + Voice: &telegramAudio{FileID: "fid", MimeType: "audio/ogg"}, + }) +} diff --git a/internal/engine/safety/capabilities.go b/internal/engine/safety/capabilities.go index ea9f1a69..6fe34d31 100644 --- a/internal/engine/safety/capabilities.go +++ b/internal/engine/safety/capabilities.go @@ -52,6 +52,7 @@ var toolPolicies = map[string]ToolPolicy{ "Diagnostics": {Name: "Diagnostics", Capabilities: []Capability{CapabilityFilesystemRead, CapabilityProcessExecute}, DefaultRisk: RiskMedium}, "ProjectVerify": {Name: "ProjectVerify", Capabilities: []Capability{CapabilityFilesystemRead, CapabilityProcessExecute}, DefaultRisk: RiskMedium}, "AppVerify": {Name: "AppVerify", Capabilities: []Capability{CapabilityFilesystemRead, CapabilityProcessExecute}, DefaultRisk: RiskMedium}, + "GenerateMedia": {Name: "GenerateMedia", Capabilities: []Capability{CapabilityNetworkAccess, CapabilityFilesystemWrite}, DefaultRisk: RiskMedium}, "DependencyAudit": {Name: "DependencyAudit", Capabilities: []Capability{CapabilityFilesystemRead, CapabilityProcessExecute, CapabilityNetworkAccess}, DefaultRisk: RiskMedium}, "Git": {Name: "Git", Capabilities: []Capability{CapabilityProcessExecute}, DefaultRisk: RiskMedium}, "GitHub": {Name: "GitHub", Capabilities: []Capability{CapabilityNetworkAccess, CapabilityProcessExecute}, DefaultRisk: RiskMedium}, diff --git a/internal/engine/safety/permission.go b/internal/engine/safety/permission.go index 8942e1b5..eefef55d 100644 --- a/internal/engine/safety/permission.go +++ b/internal/engine/safety/permission.go @@ -321,6 +321,8 @@ func canonicalToolName(name string) string { return "ProjectVerify" case "app_verify", "appverify", "verify_app": return "AppVerify" + case "generate_media", "generatemedia", "media": + return "GenerateMedia" case "dependency_audit", "dependencyaudit", "deps": return "DependencyAudit" case "git_history", "githistory", "git-history": diff --git a/internal/stt/stt.go b/internal/stt/stt.go new file mode 100644 index 00000000..d9e1f714 --- /dev/null +++ b/internal/stt/stt.go @@ -0,0 +1,173 @@ +// Package stt provides pluggable speech-to-text transcription for messaging +// gateways (e.g. Telegram voice notes), adopted from grok-cli's audio-input +// flow. A Transcriber backs the actual cloud STT call; the package handles the +// safe plumbing around it: downloading an attachment to a temp file with a +// path-traversal guard and mapping media types to extensions, so the same +// transcription pipeline works regardless of which STT backend is installed. +package stt + +import ( + "context" + "fmt" + "io" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" +) + +// Transcriber turns audio bytes into text. hawk does not bundle an STT +// provider; a host wires one in via SetTranscriber. A nil transcriber makes +// Transcription disabled and returns a clear error. +type Transcriber interface { + // Name identifies the STT backend for provenance/logging. + Name() string + // Transcribe converts audio data (already on disk at localPath) to text. + // language is optional ("", "en", etc.). + Transcribe(ctx context.Context, localPath, language string) (string, error) +} + +var transcriber Transcriber + +// SetTranscriber installs the speech-to-text backend. +func SetTranscriber(t Transcriber) { transcriber = t } + +// Enabled reports whether an STT backend is installed. +func Enabled() bool { return transcriber != nil } + +// TranscribeResult is the outcome of transcribing a downloaded attachment. +type TranscribeResult struct { + Text string + Language string + DurationMS int64 +} + +// DownloadAttachment downloads a Telegram (or generic) hosted file to a +// per-request temp directory and returns the local path. The token is used for +// Telegram's file endpoint; for other hosts set downloadToken to "". +// +// The returned path always lives inside the created temp dir, which defeats +// path-traversal attempts in the file path or name. +func DownloadAttachment(ctx context.Context, client *http.Client, downloadURL, downloadToken, suggestedName string) (string, error) { + if client == nil { + client = http.DefaultClient + } + u, err := url.Parse(downloadURL) + if err != nil { + return "", fmt.Errorf("stt: parse download url: %w", err) + } + if u.Scheme != "https" && u.Scheme != "http" { + return "", fmt.Errorf("stt: unsupported download scheme %q", u.Scheme) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, downloadURL, nil) + if err != nil { + return "", fmt.Errorf("stt: build download request: %w", err) + } + if downloadToken != "" { + req.Header.Set("Authorization", "Bearer "+downloadToken) + } + resp, err := client.Do(req) + if err != nil { + return "", fmt.Errorf("stt: download: %w", err) + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("stt: download status %d", resp.StatusCode) + } + + // Create a dedicated temp dir and confine the file to it. + dir, err := os.MkdirTemp("", "hawk-stt-*") + if err != nil { + return "", fmt.Errorf("stt: create temp dir: %w", err) + } + name := safeFileName(suggestedName) + path := filepath.Join(dir, name) + if !isPathInside(dir, path) { + // Defensive: MkdirTemp guarantees this, but never trust a caller's name. + return "", fmt.Errorf("stt: suggested name escapes temp dir") + } + f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600) + if err != nil { + return "", fmt.Errorf("stt: create temp file: %w", err) + } + defer func() { _ = f.Close() }() + if _, err := io.Copy(f, io.LimitReader(resp.Body, 50<<20)); err != nil { // 50 MB cap + return "", fmt.Errorf("stt: write attachment: %w", err) + } + return path, nil +} + +// Transcribe runs the installed transcriber over the audio file at localPath. +// It returns the transcript, or a clear error when no STT backend is installed. +func Transcribe(ctx context.Context, localPath, language string) (string, error) { + if transcriber == nil { + return "", fmt.Errorf("stt: no transcriber installed — configure a speech-to-text backend first") + } + return transcriber.Transcribe(ctx, localPath, language) +} + +// isPathInside reports whether child is lexically inside parent after cleaning. +func isPathInside(parent, child string) bool { + rel, err := filepath.Rel(parent, child) + if err != nil { + return false + } + return rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)) +} + +// safeFileName strips path separators and control characters from a suggested +// attachment name so it can never traverse out of the temp dir. +func safeFileName(name string) string { + name = filepath.Base(strings.TrimSpace(name)) + if name == "" || name == "." || name == ".." || name == string(filepath.Separator) { + return "attachment.ogg" + } + var b strings.Builder + for _, r := range name { + if r == '/' || r == '\\' || r < 0x20 || r == 0x7f { + continue + } + b.WriteRune(r) + } + cleaned := b.String() + if cleaned == "" { + return "attachment.ogg" + } + return cleaned +} + +// ExtensionForMedia infers a file extension from a Telegram media path/name or +// MIME type. Telegram voice notes are .ogg (Opus); audio attachments vary. +func ExtensionForMedia(nameOrPath, mime string) string { + base := strings.ToLower(filepath.Base(nameOrPath)) + switch { + case strings.HasSuffix(base, ".oga"), strings.HasSuffix(base, ".ogg"): + return ".ogg" + case strings.HasSuffix(base, ".mp3"): + return ".mp3" + case strings.HasSuffix(base, ".m4a"): + return ".m4a" + case strings.HasSuffix(base, ".wav"): + return ".wav" + case strings.HasSuffix(base, ".flac"): + return ".flac" + case strings.HasSuffix(base, ".aac"): + return ".aac" + } + switch strings.ToLower(mime) { + case "audio/ogg", "application/ogg", "audio/opus": + return ".ogg" + case "audio/mpeg": + return ".mp3" + case "audio/mp4", "audio/x-m4a": + return ".m4a" + case "audio/wav", "audio/x-wav": + return ".wav" + case "audio/flac": + return ".flac" + case "audio/aac": + return ".aac" + } + return ".audio" +} diff --git a/internal/stt/stt_test.go b/internal/stt/stt_test.go new file mode 100644 index 00000000..97b76374 --- /dev/null +++ b/internal/stt/stt_test.go @@ -0,0 +1,133 @@ +package stt + +import ( + "context" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "testing" +) + +type mockTranscriber struct{} + +func (mockTranscriber) Name() string { return "mock" } + +func (mockTranscriber) Transcribe(ctx context.Context, localPath, language string) (string, error) { + data, _ := os.ReadFile(localPath) + return "transcribed:" + string(data), nil +} + +func TestDownloadAttachmentConfinesToTempDir(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("fake-audio-bytes")) + })) + defer srv.Close() + + path, err := DownloadAttachment(context.Background(), nil, srv.URL+"/file", "", "voice.ogg") + if err != nil { + t.Fatalf("DownloadAttachment: %v", err) + } + t.Cleanup(func() { _ = os.RemoveAll(filepath.Dir(path)) }) + + data, err := os.ReadFile(path) + if err != nil { + t.Fatalf("read downloaded file: %v", err) + } + if string(data) != "fake-audio-bytes" { + t.Fatalf("content = %q", data) + } + if filepath.Base(path) != "voice.ogg" { + t.Fatalf("name = %q", filepath.Base(path)) + } + if !isPathInside(filepath.Dir(path), path) { + t.Fatalf("path not confined: %q", path) + } +} + +func TestDownloadAttachmentRejectsNonHTTP(t *testing.T) { + if _, err := DownloadAttachment(context.Background(), nil, "file:///etc/passwd", "", "x"); err == nil { + t.Fatal("expected error for non-http scheme") + } +} + +func TestDownloadAttachmentRejectsTraversalName(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte("x")) + })) + defer srv.Close() + + path, err := DownloadAttachment(context.Background(), nil, srv.URL+"/f", "", "../../evil.ogg") + if err != nil { + t.Fatalf("DownloadAttachment: %v", err) + } + t.Cleanup(func() { _ = os.RemoveAll(filepath.Dir(path)) }) + + // The cleaned name must not traverse; filepath.Base already strips the + // leading separators, and safeFileName removes the rest. + if strings.Contains(path, "..") { + t.Fatalf("path leaked traversal: %q", path) + } + if filepath.Base(path) != "evil.ogg" { + t.Fatalf("name = %q", filepath.Base(path)) + } +} + +func TestTranscribeRequiresEngine(t *testing.T) { + SetTranscriber(nil) + if _, err := Transcribe(context.Background(), "/tmp/x.ogg", ""); err == nil { + t.Fatal("expected error with no transcriber") + } +} + +func TestTranscribeWithEngine(t *testing.T) { + SetTranscriber(mockTranscriber{}) + t.Cleanup(func() { SetTranscriber(nil) }) + + dir := t.TempDir() + path := filepath.Join(dir, "a.ogg") + if err := os.WriteFile(path, []byte("audio"), 0o600); err != nil { + t.Fatal(err) + } + got, err := Transcribe(context.Background(), path, "en") + if err != nil { + t.Fatalf("Transcribe: %v", err) + } + if got != "transcribed:audio" { + t.Fatalf("got %q", got) + } +} + +func TestExtensionForMedia(t *testing.T) { + cases := []struct{ name, mime, want string }{ + {"voice.oga", "", ".ogg"}, + {"voice.ogg", "", ".ogg"}, + {"audio.mp3", "", ".mp3"}, + {"audio.m4a", "", ".m4a"}, + {"", "audio/ogg", ".ogg"}, + {"", "audio/mpeg", ".mp3"}, + {"", "audio/mp4", ".m4a"}, + {"", "", ".audio"}, + } + for _, c := range cases { + if got := ExtensionForMedia(c.name, c.mime); got != c.want { + t.Fatalf("ExtensionForMedia(%q,%q) = %q, want %q", c.name, c.mime, got, c.want) + } + } +} + +func TestSafeFileName(t *testing.T) { + for input, want := range map[string]string{ + "voice.ogg": "voice.ogg", + "../../evil.ogg": "evil.ogg", + "a/b/evil.ogg": "evil.ogg", + "": "attachment.ogg", + "..": "attachment.ogg", + "voice\n\x00.ogg": "voice.ogg", + } { + if got := safeFileName(input); got != want { + t.Fatalf("safeFileName(%q) = %q, want %q", input, got, want) + } + } +} diff --git a/internal/tool/media_generation.go b/internal/tool/media_generation.go new file mode 100644 index 00000000..e29acb23 --- /dev/null +++ b/internal/tool/media_generation.go @@ -0,0 +1,317 @@ +package tool + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + "os" + "path/filepath" + "strings" + "time" + + "github.com/GrayCodeAI/hawk/internal/storage" +) + +// MediaEngine is the pluggable backend that actually generates image/video +// assets. hawk does not bundle a media provider; a host (or a provider-backed +// eyrie integration) wires one in via SetMediaEngine. The tool and the +// local-persistence contract are provider-agnostic, so generation, saving, and +// result reporting all work the same regardless of backend. +type MediaEngine interface { + // Name returns the engine identifier for provenance. + Name() string + // GenerateImage produces one or more images from prompt (and optional + // source for image editing). Each result must carry the media bytes or a + // URL the host can download; the tool handles persistence. + GenerateImage(ctx context.Context, prompt, source string, opts MediaOptions) ([]MediaResult, error) + // GenerateVideo produces a video from prompt (and optional source for + // image-to-video). + GenerateVideo(ctx context.Context, prompt, source string, opts MediaOptions) ([]MediaResult, error) +} + +// MediaResult is one generated asset as returned by an engine. +type MediaResult struct { + // Data holds the raw bytes when the engine produces them locally. If empty, + // the tool attempts to download from URL. + Data []byte `json:"-"` + // URL is the provider-hosted location of the asset (used when Data is + // empty, and preserved for provenance). + URL string `json:"url,omitempty"` + // Kind is "image" or "video". + Kind string `json:"kind"` + // MIME is the asset MIME type used to pick the file extension. + MIME string `json:"mime,omitempty"` +} + +// MediaOptions carries provider-neutral generation controls. +type MediaOptions struct { + AspectRatio string `json:"aspect_ratio,omitempty"` + Resolution string `json:"resolution,omitempty"` // image: 1k/2k; video: 480p/720p + Count int `json:"count,omitempty"` // number of images + DurationSec int `json:"duration_seconds,omitempty"` +} + +// MediaAsset is the persisted, locally-available representation returned to the +// model and user. +type MediaAsset struct { + Path string `json:"path"` + Kind string `json:"kind"` + MIME string `json:"mime,omitempty"` + Prompt string `json:"prompt,omitempty"` + URL string `json:"url,omitempty"` +} + +var mediaEngine MediaEngine + +// SetMediaEngine installs the media-generation backend. It is nil by default; +// the GenerateMedia tool reports a clear error until a host wires one in. +func SetMediaEngine(e MediaEngine) { mediaEngine = e } + +// MediaEngineName returns the active engine name, or "" when none is installed. +func MediaEngineName() string { + if mediaEngine == nil { + return "" + } + return mediaEngine.Name() +} + +// DefaultMediaDir returns the stable, user-scoped directory for generated media +// (the hawk analog of grok-cli's .grok/generated-media). +func DefaultMediaDir() string { + return filepath.Join(storage.StateDir(), "generated-media") +} + +// GenerateMediaTool creates images/videos via a pluggable media engine and +// persists outputs locally so they remain usable after provider URLs expire. +// Adopted from grok-cli's generate_image/generate_video tools. +type GenerateMediaTool struct{} + +func (GenerateMediaTool) Name() string { return "GenerateMedia" } +func (GenerateMediaTool) RiskLevel() string { return "medium" } +func (GenerateMediaTool) Aliases() []string { return []string{"generate-media", "media"} } +func (GenerateMediaTool) Description() string { + return "Generate an image or short video from a text prompt (and optionally edit an existing local image or URL). The generated asset is saved locally and its path is returned so you can reference it directly." +} + +func (GenerateMediaTool) Parameters() map[string]interface{} { + return map[string]interface{}{ + "type": "object", + "properties": map[string]interface{}{ + "kind": map[string]interface{}{ + "type": "string", + "enum": []string{"image", "video"}, + "description": "The kind of media to generate.", + }, + "prompt": map[string]interface{}{ + "type": "string", + "description": "Text description of the media to generate.", + }, + "source": map[string]interface{}{ + "type": "string", + "description": "Optional local file path or URL used for image editing / image-to-video.", + }, + "aspect_ratio": map[string]interface{}{ + "type": "string", + "description": "Aspect ratio, e.g. 16:9, 1:1, 9:16.", + }, + "resolution": map[string]interface{}{ + "type": "string", + "description": "Resolution: images 1k or 2k; video 480p or 720p.", + }, + "count": map[string]interface{}{ + "type": "integer", + "minimum": 1, + "maximum": 4, + "description": "Number of images to generate (default 1).", + }, + "duration_seconds": map[string]interface{}{ + "type": "integer", + "minimum": 1, + "maximum": 15, + "description": "Video duration in seconds (default 5).", + }, + "output_path": map[string]interface{}{ + "type": "string", + "description": "Optional explicit output directory; defaults to the user state generated-media directory.", + }, + }, + "required": []string{"kind", "prompt"}, + } +} + +func (GenerateMediaTool) Execute(ctx context.Context, input json.RawMessage) (string, error) { + var p struct { + Kind string `json:"kind"` + Prompt string `json:"prompt"` + Source string `json:"source"` + AspectRatio string `json:"aspect_ratio"` + Resolution string `json:"resolution"` + Count int `json:"count"` + DurationSec int `json:"duration_seconds"` + OutputPath string `json:"output_path"` + } + if err := json.Unmarshal(input, &p); err != nil { + return "", fmt.Errorf("invalid input: %w", err) + } + p.Kind = strings.ToLower(strings.TrimSpace(p.Kind)) + p.Prompt = strings.TrimSpace(p.Prompt) + if p.Kind != "image" && p.Kind != "video" { + return "", fmt.Errorf("kind must be image or video") + } + if p.Prompt == "" { + return "", fmt.Errorf("prompt is required") + } + if mediaEngine == nil { + return "", fmt.Errorf("no media engine installed — configure a provider-backed media backend first") + } + if p.Count <= 0 { + p.Count = 1 + } + if p.Count > 4 { + p.Count = 4 + } + if p.DurationSec <= 0 { + p.DurationSec = 5 + } + if p.DurationSec > 15 { + p.DurationSec = 15 + } + + source := p.Source + if source != "" { + // Resolve a local source path to an absolute path (URLs pass through). + if !strings.Contains(source, "://") { + if err := validatePathAllowed(ctx, source); err != nil { + return "", err + } + abs, err := filepath.Abs(source) + if err != nil { + return "", fmt.Errorf("resolve source: %w", err) + } + source = abs + } + } + + opts := MediaOptions{ + AspectRatio: p.AspectRatio, + Resolution: p.Resolution, + Count: p.Count, + DurationSec: p.DurationSec, + } + + var results []MediaResult + var err error + switch p.Kind { + case "image": + results, err = mediaEngine.GenerateImage(ctx, p.Prompt, source, opts) + case "video": + results, err = mediaEngine.GenerateVideo(ctx, p.Prompt, source, opts) + } + if err != nil { + return "", fmt.Errorf("media generation failed: %w", err) + } + if len(results) == 0 { + return "", fmt.Errorf("media engine returned no assets") + } + + dest := p.OutputPath + if dest == "" { + dest = DefaultMediaDir() + } + if err := os.MkdirAll(dest, 0o750); err != nil { + return "", fmt.Errorf("create media dir: %w", err) + } + + assets := make([]MediaAsset, 0, len(results)) + for i, r := range results { + asset, err := persistMediaResult(ctx, dest, p.Kind, r, i) + if err != nil { + return "", err + } + asset.Prompt = p.Prompt + assets = append(assets, asset) + } + return encodeJSON(map[string]interface{}{"engine": mediaEngine.Name(), "assets": assets}) +} + +// persistMediaResult saves one engine result to dest and returns the asset +// descriptor. Bytes come from the engine; otherwise the URL is downloaded. +func persistMediaResult(ctx context.Context, dest, kind string, r MediaResult, index int) (MediaAsset, error) { + ext := extensionForMIME(r.MIME, kind) + name := fmt.Sprintf("media-%s-%d%s", time.Now().UTC().Format("20060102-150405"), index, ext) + path := filepath.Join(dest, name) + + if len(r.Data) == 0 && r.URL != "" { + data, err := downloadMedia(ctx, r.URL) + if err != nil { + return MediaAsset{}, fmt.Errorf("download %s: %w", r.URL, err) + } + r.Data = data + } + if len(r.Data) == 0 { + return MediaAsset{}, fmt.Errorf("media result has neither data nor a downloadable URL") + } + if err := os.WriteFile(path, r.Data, 0o644); err != nil { // #nosec G306 -- generated media is intentionally world-readable + return MediaAsset{}, fmt.Errorf("write media: %w", err) + } + return MediaAsset{ + Path: path, + Kind: kind, + MIME: r.MIME, + URL: r.URL, + }, nil +} + +// extensionForMIME maps a media MIME type to a file extension, defaulting by kind. +func extensionForMIME(mime, kind string) string { + switch strings.ToLower(mime) { + case "image/png", "image/png; charset=utf-8": + return ".png" + case "image/jpeg", "image/jpg": + return ".jpg" + case "image/webp": + return ".webp" + case "image/gif": + return ".gif" + case "video/mp4", "video/webm", "video/quicktime": + ext := ".mp4" + if strings.Contains(strings.ToLower(mime), "webm") { + ext = ".webm" + } + if strings.Contains(strings.ToLower(mime), "quicktime") { + ext = ".mov" + } + return ext + } + if kind == "video" { + return ".mp4" + } + return ".png" +} + +// downloadMedia fetches a provider-hosted media URL. Only http(s) is allowed. +func downloadMedia(ctx context.Context, rawURL string) ([]byte, error) { + u, err := url.Parse(rawURL) + if err != nil { + return nil, err + } + if u.Scheme != "http" && u.Scheme != "https" { + return nil, fmt.Errorf("unsupported URL scheme %q", u.Scheme) + } + req, err := http.NewRequestWithContext(ctx, http.MethodGet, rawURL, nil) + if err != nil { + return nil, err + } + resp, err := http.DefaultClient.Do(req) + if err != nil { + return nil, err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("download status %d", resp.StatusCode) + } + return io.ReadAll(io.LimitReader(resp.Body, 50<<20)) // 50 MB cap +} diff --git a/internal/tool/media_generation_test.go b/internal/tool/media_generation_test.go new file mode 100644 index 00000000..559796ec --- /dev/null +++ b/internal/tool/media_generation_test.go @@ -0,0 +1,126 @@ +package tool + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" +) + +type mockMediaEngine struct{} + +func (mockMediaEngine) Name() string { return "mock" } + +func (mockMediaEngine) GenerateImage(ctx context.Context, prompt, source string, opts MediaOptions) ([]MediaResult, error) { + results := make([]MediaResult, opts.Count) + for i := range results { + results[i] = MediaResult{Kind: "image", MIME: "image/png", Data: []byte("fakepng")} + } + return results, nil +} + +func (mockMediaEngine) GenerateVideo(ctx context.Context, prompt, source string, opts MediaOptions) ([]MediaResult, error) { + return []MediaResult{{Kind: "video", MIME: "video/mp4", Data: []byte("fakemp4")}}, nil +} + +func TestGenerateMediaImageSavesLocally(t *testing.T) { + SetMediaEngine(mockMediaEngine{}) + t.Cleanup(func() { SetMediaEngine(nil) }) + + outDir := t.TempDir() + res, err := (GenerateMediaTool{}).Execute(context.Background(), json.RawMessage( + `{"kind":"image","prompt":"a cat","count":2,"output_path":"`+outDir+`"}`, + )) + if err != nil { + t.Fatalf("Execute: %v", err) + } + var resp struct { + Engine string `json:"engine"` + Assets []struct { + Path string `json:"path"` + Kind string `json:"kind"` + } `json:"assets"` + } + if err := json.Unmarshal([]byte(res), &resp); err != nil { + t.Fatalf("decode: %v", err) + } + if resp.Engine != "mock" || len(resp.Assets) != 2 { + t.Fatalf("resp = %+v", resp) + } + for _, a := range resp.Assets { + if a.Kind != "image" || !strings.HasSuffix(a.Path, ".png") { + t.Fatalf("asset = %+v", a) + } + if _, err := os.Stat(a.Path); err != nil { + t.Fatalf("asset not persisted: %v", err) + } + } +} + +func TestGenerateMediaVideo(t *testing.T) { + SetMediaEngine(mockMediaEngine{}) + t.Cleanup(func() { SetMediaEngine(nil) }) + + outDir := t.TempDir() + res, err := (GenerateMediaTool{}).Execute(context.Background(), json.RawMessage( + `{"kind":"video","prompt":"sunset timelapse","output_path":"`+outDir+`"}`, + )) + if err != nil { + t.Fatalf("Execute: %v", err) + } + var resp struct { + Assets []struct { + Path string `json:"path"` + Kind string `json:"kind"` + } `json:"assets"` + } + if err := json.Unmarshal([]byte(res), &resp); err != nil { + t.Fatal(err) + } + if len(resp.Assets) != 1 || resp.Assets[0].Kind != "video" || !strings.HasSuffix(resp.Assets[0].Path, ".mp4") { + t.Fatalf("assets = %+v", resp.Assets) + } +} + +func TestGenerateMediaNoEngine(t *testing.T) { + SetMediaEngine(nil) + if _, err := (GenerateMediaTool{}).Execute(context.Background(), + json.RawMessage(`{"kind":"image","prompt":"x"}`)); err == nil { + t.Fatal("expected error when no engine installed") + } +} + +func TestGenerateMediaRequiresPrompt(t *testing.T) { + SetMediaEngine(mockMediaEngine{}) + t.Cleanup(func() { SetMediaEngine(nil) }) + if _, err := (GenerateMediaTool{}).Execute(context.Background(), + json.RawMessage(`{"kind":"image"}`)); err == nil { + t.Fatal("expected error for missing prompt") + } +} + +func TestExtensionForMIME(t *testing.T) { + cases := []struct{ mime, kind, want string }{ + {"image/png", "image", ".png"}, + {"image/jpeg", "image", ".jpg"}, + {"video/mp4", "video", ".mp4"}, + {"video/webm", "video", ".webm"}, + {"video/quicktime", "video", ".mov"}, + {"", "video", ".mp4"}, + {"", "image", ".png"}, + } + for _, c := range cases { + if got := extensionForMIME(c.mime, c.kind); got != c.want { + t.Fatalf("extensionForMIME(%q,%q) = %q, want %q", c.mime, c.kind, got, c.want) + } + } +} + +func TestDefaultMediaDir(t *testing.T) { + d := DefaultMediaDir() + if d == "" || filepath.Base(d) != "generated-media" { + t.Fatalf("DefaultMediaDir = %q", d) + } +}