diff --git a/README.md b/README.md index 053fdae..481b330 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,18 @@ grok models grok fix the flaky test in src/foo.test.ts ``` +**Generate images or short videos from chat:** + +```bash +grok "Generate a retro-futuristic logo for my CLI called Grok Forge" +grok "Edit ./assets/hero.png into a watercolor poster" +grok "Animate ./assets/cover.jpg into a 6 second cinematic push-in" +``` + +Image and video generation are exposed as agent tools inside normal chat sessions. +You keep using a text model for the session, and Grok saves generated media under +`.grok/generated-media/` by default unless you ask for a specific output path. + --- ## What you actually get @@ -101,6 +113,7 @@ grok fix the flaky test in src/foo.test.ts |--------|----------------| | **Grok-native** | Defaults tuned for Grok; models like **`grok-code-fast-1`**, **`grok-4-1-fast-reasoning`**, **`grok-4.20-multi-agent-0309`**, plus flagship and fast variants—run `grok models` for the full menu. | | **X + web search** | **`search_x`** and **`search_web`** tools—live posts and docs without pretending the internet stopped in 2023. | +| **Media generation** | Built-in **`generate_image`** and **`generate_video`** tools for text-to-image, image editing, text-to-video, and image-to-video flows. Generated files are saved locally so you can reuse them after the xAI URLs expire. | | **Sub-agents (default behavior)** | Foreground **`task`** delegation (e.g. explore vs general) plus background **`delegate`** for read-only deep dives—parallelize like you mean it. | | **Custom sub-agents** | Define named agents with **`subAgents`** in **`~/.grok/user-settings.json`** and manage them from the TUI with **`/agents`**. | | **Remote control** | Pair **Telegram** from the TUI (`/remote-control` → Telegram): DM your bot, **`/pair`**, approve the code in-terminal. Keep the CLI running while you ping it from your phone. | diff --git a/src/agent/agent.ts b/src/agent/agent.ts index 124a859..25195cc 100644 --- a/src/agent/agent.ts +++ b/src/agent/agent.ts @@ -126,6 +126,8 @@ TOOLS: - delegation_list: List running and completed background delegations. Do not poll it repeatedly. - search_web: Search the web for current information, documentation, APIs, tutorials, etc. - search_x: Search X/Twitter for real-time posts, discussions, opinions, and trends. +- generate_image: Generate a new image or edit an existing image. It saves image files locally and returns their paths. +- generate_video: Generate a new video or animate an existing image. It saves video files locally and returns their paths. - MCP tools: Enabled servers appear as tools named like mcp___. WORKFLOW: @@ -155,6 +157,8 @@ EXAMPLES: - "research how auth works" -> delegate to explore first - "investigate why this test fails" -> delegate to explore first, then continue with findings - "refactor this module" -> delegate a focused part to general when helpful +- "generate a logo" -> use generate_image +- "animate this still image" -> use generate_video - Recurring specialized workflows -> use the matching custom sub-agent via task IMPORTANT: @@ -1114,21 +1118,24 @@ function toToolResult(output: unknown): ToolResult { const r = output as { success: boolean; output?: string; + error?: string; diff?: ToolResult["diff"]; plan?: Plan; task?: ToolResult["task"]; delegation?: ToolResult["delegation"]; backgroundProcess?: ToolResult["backgroundProcess"]; + media?: ToolResult["media"]; }; return { success: r.success, output: r.output, - error: r.success ? undefined : r.output, + error: r.error ?? (r.success ? undefined : r.output), diff: r.diff, plan: r.plan, task: r.task, delegation: r.delegation, backgroundProcess: r.backgroundProcess, + media: r.media, }; } return { success: true, output: String(output) }; @@ -1141,6 +1148,8 @@ function formatSubagentActivity(toolName: string, args?: unknown): string { if (toolName === "edit_file") return `Edit ${parsed.path || "file"}`; if (toolName === "search_web") return `Web search "${truncate(parsed.query || "", 50)}"`; if (toolName === "search_x") return `X search "${truncate(parsed.query || "", 50)}"`; + if (toolName === "generate_image") return `Generate image "${truncate(parsed.prompt || "", 50)}"`; + if (toolName === "generate_video") return `Generate video "${truncate(parsed.prompt || "", 50)}"`; if (toolName === "bash") return truncate(parsed.command || "Run command", 70); return truncate(`${toolName}`, 70); } diff --git a/src/grok/media.test.ts b/src/grok/media.test.ts new file mode 100644 index 0000000..1de1eb9 --- /dev/null +++ b/src/grok/media.test.ts @@ -0,0 +1,138 @@ +import fs from "fs"; +import os from "os"; +import path from "path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { generateImageTool, generateVideoTool } from "./media"; + +const { generateImageMock, generateVideoMock } = vi.hoisted(() => ({ + generateImageMock: vi.fn(), + generateVideoMock: vi.fn(), +})); + +vi.mock("ai", () => ({ + generateImage: generateImageMock, + experimental_generateVideo: generateVideoMock, +})); + +describe("media tools", () => { + let tempDir: string; + const originalFetch = globalThis.fetch; + + beforeEach(() => { + tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "grok-media-")); + generateImageMock.mockReset(); + generateVideoMock.mockReset(); + globalThis.fetch = originalFetch; + }); + + afterEach(() => { + globalThis.fetch = originalFetch; + fs.rmSync(tempDir, { recursive: true, force: true }); + }); + + it("generates edited images from a local source path and saves them under .grok/generated-media", async () => { + const sourcePath = path.join(tempDir, "input.png"); + fs.writeFileSync(sourcePath, Buffer.from([1, 2, 3, 4])); + generateImageMock.mockResolvedValue({ + images: [{ uint8Array: new Uint8Array([9, 8, 7]), mediaType: "image/png" }], + warnings: [], + providerMetadata: { xai: { images: [{ url: "https://example.com/generated.png" }] } }, + responses: [{ modelId: "grok-imagine-image" }], + }); + + const provider = { + image: vi.fn((modelId: string) => ({ modelId })), + }; + + const result = await generateImageTool( + provider as never, + { + prompt: "Turn this into a watercolor portrait", + source: "input.png", + resolution: "2k", + }, + tempDir, + ); + + expect(result.success).toBe(true); + expect(generateImageMock).toHaveBeenCalledTimes(1); + expect(generateImageMock.mock.calls[0]?.[0]).toMatchObject({ + prompt: { + text: "Turn this into a watercolor portrait", + }, + providerOptions: { + xai: { + resolution: "2k", + }, + }, + }); + expect(provider.image).toHaveBeenCalledWith("grok-imagine-image"); + + const media = result.media ?? []; + expect(media).toHaveLength(1); + expect(media[0]?.sourcePath).toBe(sourcePath); + expect(media[0]?.url).toBe("https://example.com/generated.png"); + expect(media[0]?.path).toContain(path.join(".grok", "generated-media")); + expect(fs.existsSync(media[0]?.path ?? "")).toBe(true); + }); + + it("generates videos from a remote source image and respects output path and polling options", async () => { + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(Buffer.from([5, 4, 3, 2]), { + status: 200, + headers: { "content-type": "image/jpeg" }, + }), + ) as unknown as typeof fetch; + + generateVideoMock.mockResolvedValue({ + videos: [{ uint8Array: new Uint8Array([1, 3, 5, 7]), mediaType: "video/mp4" }], + warnings: [], + providerMetadata: { xai: { videoUrl: "https://example.com/generated.mp4" } }, + responses: [{ modelId: "grok-imagine-video" }], + }); + + const provider = { + video: vi.fn((modelId: string) => ({ modelId })), + }; + + const result = await generateVideoTool( + provider as never, + { + prompt: "Animate a slow camera push with blinking eyes", + source: "https://example.com/start-frame.jpg", + duration: 6, + resolution: "720p", + poll_interval_ms: 2000, + poll_timeout_ms: 900000, + output_path: "clips/teaser", + }, + tempDir, + ); + + expect(result.success).toBe(true); + expect(generateVideoMock).toHaveBeenCalledTimes(1); + expect(generateVideoMock.mock.calls[0]?.[0]).toMatchObject({ + prompt: { + text: "Animate a slow camera push with blinking eyes", + }, + duration: 6, + providerOptions: { + xai: { + resolution: "720p", + pollIntervalMs: 2000, + pollTimeoutMs: 900000, + }, + }, + }); + expect(provider.video).toHaveBeenCalledWith("grok-imagine-video"); + const prompt = generateVideoMock.mock.calls[0]?.[0]?.prompt as { image?: string }; + expect(prompt.image?.startsWith("data:image/jpeg;base64,")).toBe(true); + + const media = result.media ?? []; + expect(media).toHaveLength(1); + expect(media[0]?.sourceUrl).toBe("https://example.com/start-frame.jpg"); + expect(media[0]?.url).toBe("https://example.com/generated.mp4"); + expect(media[0]?.path).toBe(path.resolve(tempDir, "clips/teaser.mp4")); + expect(fs.existsSync(media[0]?.path ?? "")).toBe(true); + }); +}); diff --git a/src/grok/media.ts b/src/grok/media.ts new file mode 100644 index 0000000..eafb036 --- /dev/null +++ b/src/grok/media.ts @@ -0,0 +1,398 @@ +import { generateImage, experimental_generateVideo as generateVideo } from "ai"; +import { existsSync, mkdirSync, readFileSync, writeFileSync } from "fs"; +import { dirname, extname, isAbsolute, join, resolve } from "path"; +import type { MediaAsset, ToolResult } from "../types/index"; +import type { XaiProvider } from "./client"; + +const GENERATED_MEDIA_DIR = ".grok/generated-media"; +const IMAGE_MODEL_ID = "grok-imagine-image"; +const VIDEO_MODEL_ID = "grok-imagine-video"; + +export const IMAGE_ASPECT_RATIOS = [ + "1:1", + "16:9", + "9:16", + "4:3", + "3:4", + "3:2", + "2:3", + "2:1", + "1:2", + "19.5:9", + "9:19.5", + "20:9", + "9:20", + "auto", +] as const; + +export const VIDEO_ASPECT_RATIOS = ["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"] as const; +export const IMAGE_RESOLUTIONS = ["1k", "2k"] as const; +export const VIDEO_RESOLUTIONS = ["480p", "720p"] as const; + +export type ImageAspectRatio = (typeof IMAGE_ASPECT_RATIOS)[number]; +export type VideoAspectRatio = (typeof VIDEO_ASPECT_RATIOS)[number]; +export type ImageResolution = (typeof IMAGE_RESOLUTIONS)[number]; +export type VideoResolution = (typeof VIDEO_RESOLUTIONS)[number]; + +export interface GenerateImageToolInput { + prompt: string; + source?: string; + aspect_ratio?: ImageAspectRatio; + resolution?: ImageResolution; + n?: number; + output_path?: string; +} + +export interface GenerateVideoToolInput { + prompt: string; + source?: string; + duration?: number; + aspect_ratio?: VideoAspectRatio; + resolution?: VideoResolution; + output_path?: string; + poll_interval_ms?: number; + poll_timeout_ms?: number; +} + +interface ResolvedImageSource { + data: Uint8Array; + dataUrl: string; + mediaType: string; + sourcePath?: string; + sourceUrl?: string; +} + +export async function generateImageTool( + provider: XaiProvider, + input: GenerateImageToolInput, + cwd: string, + abortSignal?: AbortSignal, +): Promise { + try { + const source = input.source ? await resolveImageSource(input.source, cwd, abortSignal) : null; + const prompt = source ? { text: input.prompt, images: [source.data] } : input.prompt; + const response = await generateImage({ + model: provider.image(IMAGE_MODEL_ID), + prompt, + n: input.n, + aspectRatio: toSdkAspectRatio(input.aspect_ratio), + abortSignal, + providerOptions: { + xai: { + ...(input.resolution ? { resolution: input.resolution } : {}), + }, + }, + }); + + const modelId = getResponseModelId(response.responses, IMAGE_MODEL_ID); + const urls = extractImageUrls(response.providerMetadata, response.images.length); + const runId = createRunId(); + const media = response.images.map((image, index) => { + const path = buildOutputPath({ + cwd, + kind: "image", + mediaType: image.mediaType, + outputPath: input.output_path, + index, + total: response.images.length, + runId, + }); + writeFileSync(path, image.uint8Array); + return { + kind: "image", + path, + url: urls[index], + mediaType: image.mediaType, + prompt: input.prompt, + sourcePath: source?.sourcePath, + sourceUrl: source?.sourceUrl, + modelId, + } satisfies MediaAsset; + }); + + return { + success: true, + output: summarizeMediaResult("image", media, response.warnings), + media, + }; + } catch (err: unknown) { + return failureResult("Image generation failed", err); + } +} + +export async function generateVideoTool( + provider: XaiProvider, + input: GenerateVideoToolInput, + cwd: string, + abortSignal?: AbortSignal, +): Promise { + try { + const source = input.source ? await resolveImageSource(input.source, cwd, abortSignal) : null; + const prompt = source ? { text: input.prompt, image: source.dataUrl } : input.prompt; + const response = await generateVideo({ + model: provider.video(VIDEO_MODEL_ID), + prompt, + duration: input.duration, + aspectRatio: input.aspect_ratio, + abortSignal, + providerOptions: { + xai: { + ...(input.resolution ? { resolution: input.resolution } : {}), + ...(input.poll_interval_ms ? { pollIntervalMs: input.poll_interval_ms } : {}), + ...(input.poll_timeout_ms ? { pollTimeoutMs: input.poll_timeout_ms } : {}), + }, + }, + }); + + const modelId = getResponseModelId(response.responses, VIDEO_MODEL_ID); + const videoUrl = extractVideoUrl(response.providerMetadata); + const runId = createRunId(); + const media = response.videos.map((video, index) => { + const path = buildOutputPath({ + cwd, + kind: "video", + mediaType: video.mediaType, + outputPath: input.output_path, + index, + total: response.videos.length, + runId, + }); + writeFileSync(path, video.uint8Array); + return { + kind: "video", + path, + url: videoUrl, + mediaType: video.mediaType, + prompt: input.prompt, + sourcePath: source?.sourcePath, + sourceUrl: source?.sourceUrl, + durationSeconds: input.duration, + modelId, + } satisfies MediaAsset; + }); + + return { + success: true, + output: summarizeMediaResult("video", media, response.warnings), + media, + }; + } catch (err: unknown) { + return failureResult("Video generation failed", err); + } +} + +function summarizeMediaResult(kind: "image" | "video", media: MediaAsset[], warnings?: unknown[]): string { + const label = `${media.length} ${kind}${media.length === 1 ? "" : "s"}`; + const lines = [`Generated ${label}.`]; + for (const asset of media) { + const extra: string[] = []; + if (asset.url) extra.push(`url: ${asset.url}`); + if (asset.sourcePath) extra.push(`source: ${asset.sourcePath}`); + if (asset.sourceUrl) extra.push(`source_url: ${asset.sourceUrl}`); + lines.push(`- ${asset.path}${extra.length > 0 ? ` (${extra.join(", ")})` : ""}`); + } + const warningLines = (warnings ?? []).map(formatWarning).filter((warning): warning is string => Boolean(warning)); + if (warningLines.length > 0) { + lines.push("", "Warnings:"); + for (const warning of warningLines) { + lines.push(`- ${warning}`); + } + } + return lines.join("\n"); +} + +function failureResult(prefix: string, err: unknown): ToolResult { + const message = err instanceof Error ? err.message : String(err); + return { + success: false, + error: `${prefix}: ${message}`, + output: `${prefix}: ${message}`, + }; +} + +async function resolveImageSource( + source: string, + cwd: string, + abortSignal?: AbortSignal, +): Promise { + if (isUrl(source)) { + const response = await fetch(source, { signal: abortSignal }); + if (!response.ok) { + throw new Error(`Source image download failed: ${response.status} ${response.statusText}`); + } + const buffer = new Uint8Array(await response.arrayBuffer()); + const mediaType = response.headers.get("content-type") || guessMediaType(source, "image/png"); + return { + data: buffer, + dataUrl: toDataUrl(buffer, mediaType), + mediaType, + sourceUrl: source, + }; + } + + const fullPath = isAbsolute(source) ? source : resolve(cwd, source); + if (!existsSync(fullPath)) { + throw new Error(`Source image not found: ${source}`); + } + const buffer = readFileSync(fullPath); + const mediaType = guessMediaType(fullPath, "image/png"); + return { + data: buffer, + dataUrl: toDataUrl(buffer, mediaType), + mediaType, + sourcePath: fullPath, + }; +} + +function buildOutputPath({ + cwd, + kind, + mediaType, + outputPath, + index, + total, + runId, +}: { + cwd: string; + kind: "image" | "video"; + mediaType: string | undefined; + outputPath?: string; + index: number; + total: number; + runId: string; +}): string { + const extension = extensionForMediaType(kind, mediaType); + const requested = outputPath ? (isAbsolute(outputPath) ? outputPath : resolve(cwd, outputPath)) : null; + let fullPath = requested; + + if (!fullPath) { + fullPath = join(resolve(cwd, GENERATED_MEDIA_DIR), `${kind}-${runId}.${extension}`); + } + + const currentExt = extname(fullPath); + if (total > 1) { + const basePath = currentExt ? fullPath.slice(0, -currentExt.length) : fullPath; + fullPath = `${basePath}-${index + 1}.${extension}`; + } else if (!currentExt) { + fullPath = `${fullPath}.${extension}`; + } + + mkdirSync(dirname(fullPath), { recursive: true }); + return fullPath; +} + +function extensionForMediaType(kind: "image" | "video", mediaType?: string): string { + const normalized = (mediaType || "").toLowerCase(); + switch (normalized) { + case "image/jpeg": + return "jpg"; + case "image/png": + return "png"; + case "image/webp": + return "webp"; + case "image/gif": + return "gif"; + case "video/mp4": + return "mp4"; + case "video/webm": + return "webm"; + default: + return kind === "image" ? "png" : "mp4"; + } +} + +function guessMediaType(pathLike: string, fallback: string): string { + const extension = extname(pathLike).toLowerCase(); + switch (extension) { + case ".jpg": + case ".jpeg": + return "image/jpeg"; + case ".png": + return "image/png"; + case ".webp": + return "image/webp"; + case ".gif": + return "image/gif"; + case ".bmp": + return "image/bmp"; + case ".svg": + return "image/svg+xml"; + case ".mp4": + return "video/mp4"; + case ".webm": + return "video/webm"; + default: + return fallback; + } +} + +function toDataUrl(data: Uint8Array, mediaType: string): string { + return `data:${mediaType};base64,${Buffer.from(data).toString("base64")}`; +} + +function extractImageUrls(providerMetadata: unknown, count: number): Array { + const xai = getRecord(getRecord(providerMetadata)?.xai); + const images = Array.isArray(xai?.images) ? xai.images : []; + return Array.from({ length: count }, (_, index) => extractUrl(images[index])); +} + +function extractVideoUrl(providerMetadata: unknown): string | undefined { + const xai = getRecord(getRecord(providerMetadata)?.xai); + return readString(xai, "videoUrl") ?? extractUrl(xai?.video); +} + +function extractUrl(value: unknown): string | undefined { + if (typeof value === "string") return value; + const record = getRecord(value); + return readString(record, "url") ?? readString(record, "videoUrl"); +} + +function getResponseModelId(responses: unknown, fallback: string): string { + if (!Array.isArray(responses) || responses.length === 0) { + return fallback; + } + const first = getRecord(responses[0]); + return readString(first, "modelId") ?? fallback; +} + +function getRecord(value: unknown): Record | undefined { + return value && typeof value === "object" ? (value as Record) : undefined; +} + +function readString(record: Record | undefined, key: string): string | undefined { + const value = record?.[key]; + return typeof value === "string" && value.length > 0 ? value : undefined; +} + +function createRunId(): string { + return new Date().toISOString().replace(/[:.]/g, "-"); +} + +function toSdkAspectRatio(aspectRatio?: ImageAspectRatio): `${number}:${number}` | undefined { + return aspectRatio ? (aspectRatio as unknown as `${number}:${number}`) : undefined; +} + +function isUrl(value: string): boolean { + try { + const url = new URL(value); + return url.protocol === "http:" || url.protocol === "https:"; + } catch { + return false; + } +} + +function formatWarning(warning: unknown): string | undefined { + if (typeof warning === "string") { + return warning.trim() || undefined; + } + const record = getRecord(warning); + const message = readString(record, "message"); + if (message) return message.trim() || undefined; + const details = readString(record, "details"); + const feature = readString(record, "feature"); + const type = readString(record, "type"); + if (feature && details) return `${feature}: ${details}`; + if (feature) return feature; + if (details) return details; + return type; +} diff --git a/src/grok/tools.ts b/src/grok/tools.ts index 948a82a..94794db 100644 --- a/src/grok/tools.ts +++ b/src/grok/tools.ts @@ -5,6 +5,16 @@ import { editFile, readFile, writeFile } from "../tools/file"; import type { AgentMode, TaskRequest, ToolResult } from "../types/index"; import { type CustomSubagentConfig, loadValidSubAgents } from "../utils/settings"; import type { XaiProvider } from "./client"; +import { + type GenerateImageToolInput, + type GenerateVideoToolInput, + generateImageTool, + generateVideoTool, + IMAGE_ASPECT_RATIOS, + IMAGE_RESOLUTIONS, + VIDEO_ASPECT_RATIOS, + VIDEO_RESOLUTIONS, +} from "./media"; const RESPONSES_SEARCH_MODEL = "grok-4-1-fast-non-reasoning"; @@ -150,6 +160,68 @@ export function createTools( return runResponsesSearch(query, "x_search", abortSignal); }, }), + + generate_image: tool({ + description: + "Generate a new image or edit an existing image using Grok Imagine. Use when the user asks to create, redesign, restyle, or modify an image. Optionally pass a local file path or public URL in source to edit an existing image. Saves the generated image files locally and returns their paths.", + inputSchema: z.object({ + prompt: z.string().describe("Prompt describing the image to generate or the edit to apply"), + source: z + .string() + .optional() + .describe("Optional local image path or public image URL to use as the source for editing"), + aspect_ratio: z + .enum(IMAGE_ASPECT_RATIOS) + .optional() + .describe("Optional output aspect ratio. Use when the user requests a specific format."), + resolution: z.enum(IMAGE_RESOLUTIONS).optional().describe("Optional output resolution: 1k or 2k"), + n: z.number().int().min(1).max(10).optional().describe("Number of images to generate (default: 1)"), + output_path: z + .string() + .optional() + .describe("Optional file path for the generated image. For multiple images, numbered suffixes are added."), + }), + execute: async (input: GenerateImageToolInput, { abortSignal }) => { + return generateImageTool(provider, input, cwd(), abortSignal); + }, + }), + + generate_video: tool({ + description: + "Generate a new short video or animate an existing image using Grok Imagine Video. Use when the user asks for a clip, animation, cinematic shot, or motion from a still image. Optionally pass a local image path or public image URL in source for image-to-video generation. Saves the generated video files locally and returns their paths.", + inputSchema: z.object({ + prompt: z.string().describe("Prompt describing the video or motion to generate"), + source: z + .string() + .optional() + .describe("Optional local image path or public image URL to use as the starting frame"), + duration: z.number().int().min(1).max(15).optional().describe("Video duration in seconds (1-15)"), + aspect_ratio: z + .enum(VIDEO_ASPECT_RATIOS) + .optional() + .describe("Optional output aspect ratio for text-to-video or to override image-to-video framing"), + resolution: z.enum(VIDEO_RESOLUTIONS).optional().describe("Optional output resolution: 480p or 720p"), + output_path: z + .string() + .optional() + .describe("Optional file path for the generated video. For multiple videos, numbered suffixes are added."), + poll_interval_ms: z + .number() + .int() + .min(100) + .optional() + .describe("Optional polling interval in milliseconds while waiting for video generation"), + poll_timeout_ms: z + .number() + .int() + .min(1000) + .optional() + .describe("Optional timeout in milliseconds while waiting for video generation"), + }), + execute: async (input: GenerateVideoToolInput, { abortSignal }) => { + return generateVideoTool(provider, input, cwd(), abortSignal); + }, + }), }; const tools: ToolSet = { ...base }; diff --git a/src/headless/output.test.ts b/src/headless/output.test.ts index 2ca5b09..9d704e8 100644 --- a/src/headless/output.test.ts +++ b/src/headless/output.test.ts @@ -71,6 +71,22 @@ describe("headless output helpers", () => { }); }); + it("renders generated media paths in text mode tool results", () => { + const chunk: StreamChunk = { + type: "tool_result", + toolCall: toolCall("generate_image"), + toolResult: { + success: true, + output: "Generated 1 image.", + media: [{ kind: "image", path: "/tmp/generated.png", url: "https://example.com/generated.png" }], + }, + }; + + expect(renderHeadlessChunk(chunk)).toEqual({ + stderr: "\u001b[32m✓ generate_image\u001b[0m\n /tmp/generated.png (https://example.com/generated.png)\n", + }); + }); + it("emits semantic JSONL for a single step with text and tool (json emitter)", () => { const sessionId = "jsonl-test-session"; const tc = toolCall("bash"); diff --git a/src/headless/output.ts b/src/headless/output.ts index 5bd8072..24d0b24 100644 --- a/src/headless/output.ts +++ b/src/headless/output.ts @@ -94,7 +94,13 @@ export function renderHeadlessChunk(chunk: StreamChunk): HeadlessWrites { const icon = chunk.toolResult.success ? "✓" : "✗"; const color = chunk.toolResult.success ? "\x1b[32m" : "\x1b[31m"; const name = chunk.toolCall?.function.name || "tool"; - return { stderr: `${color}${icon} ${name}\x1b[0m\n` }; + const mediaLines = + chunk.toolResult.media?.map((asset) => { + const suffix = asset.url ? ` (${asset.url})` : ""; + return ` ${asset.path}${suffix}`; + }) ?? []; + const stderr = [`${color}${icon} ${name}\x1b[0m`, ...mediaLines].join("\n"); + return { stderr: `${stderr}\n` }; } case "error": diff --git a/src/storage/tool-results.ts b/src/storage/tool-results.ts new file mode 100644 index 0000000..54709bb --- /dev/null +++ b/src/storage/tool-results.ts @@ -0,0 +1,55 @@ +import type { ToolResult } from "../types/index"; + +export function extractToolResultFromOutput(output: unknown): ToolResult | null { + if (!output || typeof output !== "object") return null; + + if ("success" in output) { + const result = output as ToolResult; + return { + success: Boolean(result.success), + output: result.output, + error: result.error, + diff: result.diff, + plan: result.plan, + task: result.task, + delegation: result.delegation, + backgroundProcess: result.backgroundProcess, + media: result.media, + }; + } + + if ("type" in output && output.type === "json" && "value" in output) { + return extractToolResultFromOutput((output as { value: unknown }).value); + } + + if ("type" in output && output.type === "error-text" && "value" in output) { + return { + success: false, + error: String((output as { value: unknown }).value), + }; + } + + if ("type" in output && output.type === "text" && "value" in output) { + return { + success: true, + output: String((output as { value: unknown }).value), + }; + } + + return null; +} + +export function getOutputKind(output: unknown): string { + if (output && typeof output === "object" && "type" in output && typeof output.type === "string") { + return output.type; + } + return "json"; +} + +export function isOutputSuccess(output: unknown): boolean { + if (!output || typeof output !== "object") return true; + if ("type" in output) { + return !String(output.type).startsWith("error"); + } + return true; +} diff --git a/src/storage/transcript.test.ts b/src/storage/transcript.test.ts new file mode 100644 index 0000000..76195a2 --- /dev/null +++ b/src/storage/transcript.test.ts @@ -0,0 +1,24 @@ +import { describe, expect, it } from "vitest"; +import type { ToolResult } from "../types/index"; +import { extractToolResultFromOutput } from "./tool-results"; + +describe("transcript media tool results", () => { + it("preserves media metadata when stored tool output is normalized", () => { + const mediaResult: ToolResult = { + success: true, + output: "Generated 1 image.\n- /tmp/example.png", + media: [ + { + kind: "image", + path: "/tmp/example.png", + url: "https://example.com/generated.png", + sourcePath: "/tmp/source.png", + prompt: "Create a new hero image", + modelId: "grok-imagine-image", + }, + ], + }; + + expect(extractToolResultFromOutput(mediaResult)).toEqual(mediaResult); + }); +}); diff --git a/src/storage/transcript.ts b/src/storage/transcript.ts index 9d366f8..d072ea2 100644 --- a/src/storage/transcript.ts +++ b/src/storage/transcript.ts @@ -2,6 +2,7 @@ import type { ModelMessage } from "ai"; import { getCompactionSummaryText } from "../agent/compaction"; import type { ChatEntry, ToolCall, ToolResult } from "../types/index"; import { getDatabase, withTransaction } from "./db"; +import { extractToolResultFromOutput, getOutputKind, isOutputSuccess } from "./tool-results"; import { buildEffectiveTranscript, type LoadedTranscriptState, type PersistedCompaction } from "./transcript-view"; interface MessageRow { @@ -355,56 +356,3 @@ function toFallbackToolCall(toolCallId: string, toolName: string): ToolCall { }, }; } - -function extractToolResultFromOutput(output: unknown): ToolResult | null { - if (!output || typeof output !== "object") return null; - - if ("success" in output) { - const result = output as ToolResult; - return { - success: Boolean(result.success), - output: result.output, - error: result.error, - diff: result.diff, - plan: result.plan, - task: result.task, - delegation: result.delegation, - backgroundProcess: result.backgroundProcess, - }; - } - - if ("type" in output && output.type === "json" && "value" in output) { - return extractToolResultFromOutput((output as { value: unknown }).value); - } - - if ("type" in output && output.type === "error-text" && "value" in output) { - return { - success: false, - error: String((output as { value: unknown }).value), - }; - } - - if ("type" in output && output.type === "text" && "value" in output) { - return { - success: true, - output: String((output as { value: unknown }).value), - }; - } - - return null; -} - -function getOutputKind(output: unknown): string { - if (output && typeof output === "object" && "type" in output && typeof output.type === "string") { - return output.type; - } - return "json"; -} - -function isOutputSuccess(output: unknown): boolean { - if (!output || typeof output !== "object") return true; - if ("type" in output) { - return !String(output.type).startsWith("error"); - } - return true; -} diff --git a/src/types/index.ts b/src/types/index.ts index 0f8372e..e253557 100644 --- a/src/types/index.ts +++ b/src/types/index.ts @@ -66,6 +66,18 @@ export interface BackgroundProcessInfo { command: string; } +export interface MediaAsset { + kind: "image" | "video"; + path: string; + url?: string; + mediaType?: string; + prompt?: string; + sourcePath?: string; + sourceUrl?: string; + durationSeconds?: number; + modelId?: string; +} + export interface ToolResult { success: boolean; output?: string; @@ -75,6 +87,7 @@ export interface ToolResult { task?: TaskRun; delegation?: DelegationRun; backgroundProcess?: BackgroundProcessInfo; + media?: MediaAsset[]; } export interface ToolCall { diff --git a/src/ui/app.tsx b/src/ui/app.tsx index 8b7af76..3155b49 100644 --- a/src/ui/app.tsx +++ b/src/ui/app.tsx @@ -3228,6 +3228,10 @@ function MessageView({ entry, index, t, modeColor }: { entry: ChatEntry; index: return ; } + if ((entry.toolResult?.media?.length ?? 0) > 0) { + return ; + } + if (name === "write_file" || name === "edit_file") { const filePath = diff?.filePath || tryParseArg(entry.toolCall, "path") || args; const label = name === "write_file" ? `Write ${filePath}` : `Edit ${filePath}`; @@ -3648,6 +3652,35 @@ function ToolTextOutputView({ t, label, content }: { t: Theme; label: string; co ); } +function MediaToolResultView({ t, label, toolResult }: { t: Theme; label: string; toolResult: ToolResult }) { + const media = toolResult.media ?? []; + + return ( + + + {label} + + {toolResult.output ? ( + + + + ) : null} + {media.length > 0 ? ( + + {media.map((asset, index) => ( + + {asset.path} + {asset.url ? {`url: ${asset.url}`} : null} + {asset.sourcePath ? {`source: ${asset.sourcePath}`} : null} + {asset.sourceUrl ? {`source_url: ${asset.sourceUrl}`} : null} + + ))} + + ) : null} + + ); +} + /* ── Slash Menu ──────────────────────────────────────────────── */ function bottomAlignedModalTop(height: number, panelHeight: number): number { @@ -4110,6 +4143,7 @@ function toolArgs(tc?: ToolCall): string { if (tc.function.name === "bash") return a.command || ""; if (tc.function.name === "read_file" || tc.function.name === "write_file" || tc.function.name === "edit_file") return a.path || ""; + if (tc.function.name === "generate_image" || tc.function.name === "generate_video") return a.prompt || ""; if (tc.function.name === "task") return a.description || ""; if (tc.function.name === "delegate") return a.description || ""; if (tc.function.name === "delegation_read") return a.id || ""; @@ -4144,6 +4178,8 @@ function toolLabel(tc: ToolCall): string { if (tc.function.name === "edit_file") return `Edit ${trunc(args, 60)}`; if (tc.function.name === "search_web") return `Web Search "${trunc(args, 60)}"`; if (tc.function.name === "search_x") return `X Search "${trunc(args, 60)}"`; + if (tc.function.name === "generate_image") return `Generate image "${trunc(args, 60)}"`; + if (tc.function.name === "generate_video") return `Generate video "${trunc(args, 60)}"`; if (tc.function.name === "task") return `Task ${trunc(args, 60)}`; if (tc.function.name === "delegate") return `Background ${trunc(args, 60)}`; if (tc.function.name === "delegation_read") return `Read delegation ${trunc(args, 60)}`;