From 7438c97cf8793449b4063125833d9bd4450daad7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Fri, 31 Jul 2026 21:30:31 +0200 Subject: [PATCH] feat(#3996): render generated images inline via manifest-gated resolution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Render model-generated images inline in the same assistant turn (with a filename-only fallback on non-graphics terminals), resolving strictly through LocalRuntime.ResolveGeneratedFile: one recorded reference resolves to bytes plus a validated canonical path, gated on the generated-media manifest, the owning session's persisted WorkingDir (bounded parent fallback, never the viewer cwd), os.Root containment, and symlink-free path components — a forged DocumentSource or a symlink swapped in after materialization selects nothing. Per-owner roots and manifest records are cached on the runtime and seeded at materialization. The chat page attaches a sanitized "unavailable" placeholder synchronously and resolves inside a routed tea.Cmd — never in Update — then swaps results in by ID, for live MessageAddedEvents and for restored session messages. Successful resolution shows the canonical workspace path in the non-graphics fallback; every failure (missing, replaced, tampered, unknown root kind, unrecorded) degrades to the filename-only wording. Runtimes without the capability (remote) render nothing. TUI regression tests pin that same-turn labels and fallbacks come from the final persisted Document name and the resolver-validated canonical workspace path (including collision-suffixed names) — the TUI never constructs paths or sees provisional MediaDelta names. --- pkg/app/app.go | 30 + pkg/app/generated_file_test.go | 57 ++ pkg/runtime/generated_file.go | 293 +++++++++ pkg/runtime/generated_file_test.go | 355 ++++++++++ pkg/runtime/loop.go | 21 +- pkg/runtime/media_escape_test.go | 81 +++ pkg/runtime/runtime.go | 16 +- pkg/tui/components/message/message.go | 46 +- pkg/tui/components/message/message_test.go | 75 +++ .../messages/generated_media_test.go | 120 ++++ pkg/tui/components/messages/messages.go | 92 ++- pkg/tui/components/messages/messages_test.go | 97 ++- pkg/tui/page/chat/chat.go | 27 +- pkg/tui/page/chat/generated_media.go | 216 ++++++ pkg/tui/page/chat/generated_media_test.go | 615 ++++++++++++++++++ pkg/tui/page/chat/runtime_events.go | 4 + pkg/tui/types/types.go | 18 + 17 files changed, 2125 insertions(+), 38 deletions(-) create mode 100644 pkg/app/generated_file_test.go create mode 100644 pkg/runtime/generated_file.go create mode 100644 pkg/runtime/generated_file_test.go create mode 100644 pkg/tui/components/messages/generated_media_test.go create mode 100644 pkg/tui/page/chat/generated_media.go create mode 100644 pkg/tui/page/chat/generated_media_test.go diff --git a/pkg/app/app.go b/pkg/app/app.go index d6e968abc8..e9019bc400 100644 --- a/pkg/app/app.go +++ b/pkg/app/app.go @@ -1146,6 +1146,36 @@ type contextBreakdownProvider interface { ContextBreakdown(ctx context.Context, sess *session.Session) (*runtime.ContextBreakdown, error) } +// generatedFileResolver is an optional runtime capability: resolving one +// recorded generated-media reference to its bytes and validated canonical +// path, gated on the generated-media manifest and the owning session's +// workspace (see [runtime.LocalRuntime.ResolveGeneratedFile]). Only the +// local runtime implements it; remote runtimes never deliver generated-file +// payloads, so UIs treat the missing capability as "nothing to resolve". +type generatedFileResolver interface { + ResolveGeneratedFile(ctx context.Context, ref runtime.GeneratedFileRef) (*runtime.ResolvedGeneratedFile, error) +} + +// CanResolveGeneratedFiles reports whether the runtime can resolve +// generated-media references at all, letting UIs skip resolution work +// entirely on runtimes without the capability. +func (a *App) CanResolveGeneratedFiles() bool { + _, ok := a.runtime.(generatedFileResolver) + return ok +} + +// ResolveGeneratedFile resolves one recorded generated-media reference. +// Returns an error wrapping [runtime.ErrUnsupported] when the runtime does +// not own local generated media (e.g. remote runtimes). Callers must treat +// any error as "unavailable" — never surface its text to the user. +func (a *App) ResolveGeneratedFile(ctx context.Context, ref runtime.GeneratedFileRef) (*runtime.ResolvedGeneratedFile, error) { + resolver, ok := a.runtime.(generatedFileResolver) + if !ok { + return nil, fmt.Errorf("generated file resolution: %w", runtime.ErrUnsupported) + } + return resolver.ResolveGeneratedFile(ctx, ref) +} + // ContextBreakdown returns the estimated context-window composition for the // current session. Returns an error wrapping [runtime.ErrUnsupported] when // the runtime cannot compute it (e.g. remote runtimes). diff --git a/pkg/app/generated_file_test.go b/pkg/app/generated_file_test.go new file mode 100644 index 0000000000..aea2d8733b --- /dev/null +++ b/pkg/app/generated_file_test.go @@ -0,0 +1,57 @@ +package app + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/runtime" + "github.com/docker/docker-agent/pkg/session" +) + +// resolvingRuntime is mockRuntime plus the generated-file resolver +// capability (see generatedFileResolver), recording the refs it was asked +// to resolve. +type resolvingRuntime struct { + mockRuntime + + refs []runtime.GeneratedFileRef + resolved *runtime.ResolvedGeneratedFile + err error +} + +func (r *resolvingRuntime) ResolveGeneratedFile(_ context.Context, ref runtime.GeneratedFileRef) (*runtime.ResolvedGeneratedFile, error) { + r.refs = append(r.refs, ref) + return r.resolved, r.err +} + +func TestApp_ResolveGeneratedFile_ForwardsToCapableRuntime(t *testing.T) { + t.Parallel() + rt := &resolvingRuntime{resolved: &runtime.ResolvedGeneratedFile{Data: []byte("png"), Path: "/ws/cat.png"}} + app := New(t.Context(), rt, session.New()) + ref := runtime.GeneratedFileRef{OwnerSessionID: "sess", Root: chat.ArtifactRootWorkspace, Path: "cat.png"} + + assert.True(t, app.CanResolveGeneratedFiles()) + resolved, err := app.ResolveGeneratedFile(t.Context(), ref) + + require.NoError(t, err) + assert.Equal(t, rt.resolved, resolved) + assert.Equal(t, []runtime.GeneratedFileRef{ref}, rt.refs) +} + +// TestApp_ResolveGeneratedFile_UnsupportedWithoutCapability pins the +// remote-runtime shape: a runtime without the resolver capability reports +// it upfront and resolution fails with runtime.ErrUnsupported. +func TestApp_ResolveGeneratedFile_UnsupportedWithoutCapability(t *testing.T) { + t.Parallel() + app := New(t.Context(), &mockRuntime{}, session.New()) + + assert.False(t, app.CanResolveGeneratedFiles()) + _, err := app.ResolveGeneratedFile(t.Context(), runtime.GeneratedFileRef{ + OwnerSessionID: "sess", Root: chat.ArtifactRootWorkspace, Path: "cat.png", + }) + assert.ErrorIs(t, err, runtime.ErrUnsupported) +} diff --git a/pkg/runtime/generated_file.go b/pkg/runtime/generated_file.go new file mode 100644 index 0000000000..32f1cc5211 --- /dev/null +++ b/pkg/runtime/generated_file.go @@ -0,0 +1,293 @@ +package runtime + +import ( + "context" + "errors" + "fmt" + "io" + "os" + "path" + "path/filepath" + "strings" + "sync" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" +) + +// ErrGeneratedFileUnavailable is the single caller-visible failure of +// [LocalRuntime.ResolveGeneratedFile]. Every refusal — unknown root kind, +// missing manifest record, root-kind mismatch, workspace escape, symlink +// replacement, missing or oversized file — collapses into it so UIs can +// only ever say "unavailable"; the wrapped cause is for debug logs. +var ErrGeneratedFileUnavailable = errors.New("generated file unavailable") + +// maxGeneratedFileBytes bounds how much a single resolution reads into +// memory. Matches the inline-rendering bound in pkg/tui/image; a recorded +// path whose content grew beyond it (i.e. was replaced) is refused rather +// than loaded. +const maxGeneratedFileBytes = 20 << 20 + +// GeneratedFileRef identifies one persisted generated-media reference, as +// carried by [chat.DocumentSource] (ArtifactPath/ArtifactRoot/ +// ArtifactOwnerSessionID). +type GeneratedFileRef struct { + // OwnerSessionID is the session the file was materialized under — the + // owning session, never the viewing one. + OwnerSessionID string + // Root is the root kind Path is interpreted against. Only + // chat.ArtifactRootWorkspace and chat.ArtifactRootExternal resolve; + // the empty (unknown) kind is always unavailable. + Root chat.ArtifactRootKind + // Path is the recorded final path: workspace-relative slash-separated + // for the workspace root, absolute for the external root. + Path string +} + +// ResolvedGeneratedFile is a successful resolution: the file bytes plus the +// validated canonical absolute path, safe to display verbatim (owner IDs, +// raw refs, and error details are never part of it). +type ResolvedGeneratedFile struct { + Data []byte + Path string +} + +// generatedFileCache is the per-owner-session state ResolveGeneratedFile +// needs: the resolved workspace root and the manifest records seen so far. +// It is seeded by materialization (so live rendering does no store reads) +// and filled lazily for restored sessions. Only positive results are +// cached: a miss may become a hit when materialization records a new file, +// while recorded state never silently disappears — this runtime's store is +// the only manifest writer, and an upsert for an existing path refreshes +// the entry through cacheGeneratedFile. +type generatedFileCache struct { + mu sync.Mutex + roots map[string]string // owner session ID → workspace root + files map[string]map[string]session.GeneratedFile // owner session ID → recorded path → record +} + +func (c *generatedFileCache) root(ownerID string) (string, bool) { + c.mu.Lock() + defer c.mu.Unlock() + root, ok := c.roots[ownerID] + return root, ok +} + +func (c *generatedFileCache) setRoot(ownerID, root string) { + c.mu.Lock() + defer c.mu.Unlock() + if c.roots == nil { + c.roots = make(map[string]string) + } + c.roots[ownerID] = root +} + +func (c *generatedFileCache) file(ownerID, relPath string) (session.GeneratedFile, bool) { + c.mu.Lock() + defer c.mu.Unlock() + record, ok := c.files[ownerID][relPath] + return record, ok +} + +func (c *generatedFileCache) setFile(record session.GeneratedFile) { + c.mu.Lock() + defer c.mu.Unlock() + if c.files == nil { + c.files = make(map[string]map[string]session.GeneratedFile) + } + if c.files[record.SessionID] == nil { + c.files[record.SessionID] = make(map[string]session.GeneratedFile) + } + c.files[record.SessionID][record.RelPath] = record +} + +// ResolveGeneratedFile resolves one recorded generated-media reference to +// its bytes and validated canonical path. It is the only supported read +// path for generated media: the (owner session, path) pair must have been +// recorded in the generated-media manifest by materialization, the root +// kind must match the record, and a workspace path must still be a plain +// regular file inside the owning session's workspace — a reference alone, +// however it was forged, never selects a file. +// +// It is safe for concurrent use and intended to be called off the UI +// update loop (e.g. inside a tea.Cmd). +func (r *LocalRuntime) ResolveGeneratedFile(ctx context.Context, ref GeneratedFileRef) (*ResolvedGeneratedFile, error) { + if ref.OwnerSessionID == "" { + return nil, fmt.Errorf("%w: reference without an owner session", ErrGeneratedFileUnavailable) + } + if ref.Root != chat.ArtifactRootWorkspace && ref.Root != chat.ArtifactRootExternal { + return nil, fmt.Errorf("%w: unresolvable root kind %q", ErrGeneratedFileUnavailable, ref.Root) + } + + record, err := r.lookupGeneratedFile(ctx, ref) + if err != nil { + return nil, err + } + if record.Root != ref.Root { + return nil, fmt.Errorf("%w: reference root %q does not match recorded root %q", ErrGeneratedFileUnavailable, ref.Root, record.Root) + } + + if ref.Root == chat.ArtifactRootExternal { + data, err := readExternalGeneratedFile(ref.Path) + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrGeneratedFileUnavailable, err) + } + return &ResolvedGeneratedFile{Data: data, Path: ref.Path}, nil + } + + workspaceRoot, err := r.generatedFileWorkspaceRoot(ctx, ref.OwnerSessionID) + if err != nil { + return nil, err + } + data, canonical, err := readWorkspaceGeneratedFile(workspaceRoot, ref.Path) + if err != nil { + return nil, fmt.Errorf("%w: %w", ErrGeneratedFileUnavailable, err) + } + return &ResolvedGeneratedFile{Data: data, Path: canonical}, nil +} + +// lookupGeneratedFile returns the manifest record for ref, from the cache +// or the session store. +func (r *LocalRuntime) lookupGeneratedFile(ctx context.Context, ref GeneratedFileRef) (session.GeneratedFile, error) { + if record, ok := r.generatedFiles.file(ref.OwnerSessionID, ref.Path); ok { + return record, nil + } + manifest, ok := r.sessionStore.(session.GeneratedMediaManifest) + if !ok { + return session.GeneratedFile{}, fmt.Errorf("%w: session store %T has no generated-media manifest", ErrGeneratedFileUnavailable, r.sessionStore) + } + record, err := manifest.LookupGeneratedFile(ctx, ref.OwnerSessionID, ref.Path) + if err != nil { + return session.GeneratedFile{}, fmt.Errorf("%w: %w", ErrGeneratedFileUnavailable, err) + } + r.generatedFiles.setFile(*record) + return *record, nil +} + +// generatedFileWorkspaceRoot returns the OWNING session's workspace root — +// persisted WorkingDir with the bounded parent-chain fallback, never the +// viewer's cwd — from the cache or the session store. +func (r *LocalRuntime) generatedFileWorkspaceRoot(ctx context.Context, ownerID string) (string, error) { + if root, ok := r.generatedFiles.root(ownerID); ok { + return root, nil + } + if r.sessionStore == nil { + return "", fmt.Errorf("%w: no session store to resolve the owner workspace", ErrGeneratedFileUnavailable) + } + owner, err := r.sessionStore.GetSession(ctx, ownerID) + if err != nil { + return "", fmt.Errorf("%w: loading owner session: %w", ErrGeneratedFileUnavailable, err) + } + root, err := session.ResolveWorkingDir(ctx, owner, r.sessionLookup()) + if err != nil { + return "", fmt.Errorf("%w: %w", ErrGeneratedFileUnavailable, err) + } + r.generatedFiles.setRoot(ownerID, root) + return root, nil +} + +// readWorkspaceGeneratedFile reads relPath under workspaceRoot with the +// same containment the writer enforced: os.Root confines every operation +// to the workspace, and no path component may be a symlink — the manifest +// recorded a regular file written by pkg/workspacemedia, so a symlink +// found now (even one pointing elsewhere INSIDE the workspace, e.g. at +// ".env") means the file was replaced and must not be followed. +func readWorkspaceGeneratedFile(workspaceRoot, relPath string) (data []byte, canonical string, err error) { + root, err := os.OpenRoot(workspaceRoot) + if err != nil { + return nil, "", fmt.Errorf("opening workspace root: %w", err) + } + defer root.Close() + + osRel := filepath.FromSlash(relPath) + if err := rejectSymlinkComponents(root, relPath); err != nil { + return nil, "", err + } + + f, err := root.Open(osRel) + if err != nil { + return nil, "", fmt.Errorf("opening recorded file: %w", err) + } + defer f.Close() + data, err = readRegularGeneratedFile(f, func() (os.FileInfo, error) { return root.Lstat(osRel) }) + if err != nil { + return nil, "", err + } + + // The workspace root itself may legitimately be reached through + // symlinks (e.g. macOS /tmp); canonicalize it so the displayed path is + // the real location. The recorded relative path below it is + // symlink-free (checked above), so a plain join stays canonical. + canonicalRoot, err := filepath.EvalSymlinks(workspaceRoot) + if err != nil { + canonicalRoot = workspaceRoot + } + return data, filepath.Join(canonicalRoot, osRel), nil +} + +// rejectSymlinkComponents fails when any component of the slash-separated +// relPath — intermediate directory or final file — is a symlink inside +// root. +func rejectSymlinkComponents(root *os.Root, relPath string) error { + components := strings.Split(relPath, "/") + for i := range components { + prefix := path.Join(components[:i+1]...) + fi, err := root.Lstat(filepath.FromSlash(prefix)) + if err != nil { + return fmt.Errorf("inspecting recorded path: %w", err) + } + if fi.Mode()&os.ModeSymlink != 0 { + return fmt.Errorf("recorded path component %q was replaced by a symlink", prefix) + } + if i < len(components)-1 && !fi.IsDir() { + return fmt.Errorf("recorded path component %q is not a directory", prefix) + } + } + return nil +} + +// readExternalGeneratedFile reads a user-confirmed absolute external +// target. Parent directories may legitimately be symlinks (the user +// confirmed this exact path, e.g. under macOS /tmp), but the leaf must +// still be the regular file the writer created — a symlink swapped in +// afterwards is refused. +func readExternalGeneratedFile(target string) ([]byte, error) { + f, err := os.Open(target) + if err != nil { + return nil, fmt.Errorf("opening recorded external file: %w", err) + } + defer f.Close() + return readRegularGeneratedFile(f, func() (os.FileInfo, error) { return os.Lstat(target) }) +} + +// readRegularGeneratedFile reads an opened generated file after verifying — +// against a fresh Lstat taken AFTER the open, closing the check/open race — +// that the path still names this exact regular, size-bounded file rather +// than a symlink swapped in since materialization. +func readRegularGeneratedFile(f *os.File, lstat func() (os.FileInfo, error)) ([]byte, error) { + st, err := f.Stat() + if err != nil { + return nil, fmt.Errorf("inspecting recorded file: %w", err) + } + if !st.Mode().IsRegular() { + return nil, fmt.Errorf("recorded file is not a regular file (%s)", st.Mode()) + } + if st.Size() > maxGeneratedFileBytes { + return nil, fmt.Errorf("recorded file is %d bytes, above the %d-byte display bound", st.Size(), maxGeneratedFileBytes) + } + lfi, err := lstat() + if err != nil { + return nil, fmt.Errorf("re-inspecting recorded path: %w", err) + } + if lfi.Mode()&os.ModeSymlink != 0 || !os.SameFile(lfi, st) { + return nil, errors.New("recorded path no longer names the opened file") + } + data, err := io.ReadAll(io.LimitReader(f, maxGeneratedFileBytes+1)) + if err != nil { + return nil, fmt.Errorf("reading recorded file: %w", err) + } + if len(data) > maxGeneratedFileBytes { + return nil, errors.New("recorded file grew above the display bound during read") + } + return data, nil +} diff --git a/pkg/runtime/generated_file_test.go b/pkg/runtime/generated_file_test.go new file mode 100644 index 0000000000..77634df590 --- /dev/null +++ b/pkg/runtime/generated_file_test.go @@ -0,0 +1,355 @@ +package runtime + +import ( + "context" + "os" + "path/filepath" + stdruntime "runtime" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" +) + +// resolverTestRuntime is a LocalRuntime over an in-memory store with one +// owning session that has a real workspace root. +func resolverTestRuntime(t *testing.T, sess *session.Session) (*LocalRuntime, session.Store) { + t.Helper() + store := session.NewInMemorySessionStore() + require.NoError(t, store.AddSession(t.Context(), sess)) + return &LocalRuntime{sessionStore: store, now: time.Now}, store +} + +// recordWorkspaceFile writes content at relPath under root and records it +// in the manifest, exactly as materialization would. +func recordWorkspaceFile(t *testing.T, store session.Store, sessID, root, relPath string, content []byte) { + t.Helper() + target := filepath.Join(root, filepath.FromSlash(relPath)) + require.NoError(t, os.MkdirAll(filepath.Dir(target), 0o755)) + require.NoError(t, os.WriteFile(target, content, 0o644)) + require.NoError(t, manifestOf(t, store).AddGeneratedFile(t.Context(), session.GeneratedFile{ + SessionID: sessID, + RelPath: relPath, + Root: chat.ArtifactRootWorkspace, + MimeType: "image/png", + CreatedAt: time.Now(), + })) +} + +func workspaceRef(owner, relPath string) GeneratedFileRef { + return GeneratedFileRef{OwnerSessionID: owner, Root: chat.ArtifactRootWorkspace, Path: relPath} +} + +func requireSymlinkSupport(t *testing.T) { + t.Helper() + if stdruntime.GOOS == "windows" { + t.Skip("symlink creation requires elevated privileges on Windows") + } +} + +func TestResolveGeneratedFile_WorkspaceSuccess(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-resolve") + r, store := resolverTestRuntime(t, sess) + recordWorkspaceFile(t, store, sess.ID, root, "images/cat.png", []byte("png-bytes")) + + resolved, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "images/cat.png")) + + require.NoError(t, err) + assert.Equal(t, []byte("png-bytes"), resolved.Data) + canonicalRoot, err := filepath.EvalSymlinks(root) + require.NoError(t, err) + assert.Equal(t, filepath.Join(canonicalRoot, "images", "cat.png"), resolved.Path) +} + +// TestResolveGeneratedFile_ForgedReferenceRefused is the manifest contract: +// a reference alone — however session JSON was tampered with — must never +// select a real workspace file such as ".env" or a source file. +func TestResolveGeneratedFile_ForgedReferenceRefused(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-forged") + r, _ := resolverTestRuntime(t, sess) + require.NoError(t, os.WriteFile(filepath.Join(root, ".env"), []byte("SECRET=1"), 0o600)) + require.NoError(t, os.WriteFile(filepath.Join(root, "main.go"), []byte("package main"), 0o644)) + + for _, path := range []string{".env", "main.go"} { + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, path)) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable, "unrecorded workspace file %q must not resolve", path) + } +} + +func TestResolveGeneratedFile_LegacyAndInvalidRefsRefused(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-legacy") + r, store := resolverTestRuntime(t, sess) + recordWorkspaceFile(t, store, sess.ID, root, "cat.png", []byte("png")) + + for name, ref := range map[string]GeneratedFileRef{ + "legacy empty root": {OwnerSessionID: sess.ID, Root: "", Path: "cat.png"}, + "unknown root kind": {OwnerSessionID: sess.ID, Root: "datadir", Path: "cat.png"}, + "missing owner": {Root: chat.ArtifactRootWorkspace, Path: "cat.png"}, + "traversal path": workspaceRef(sess.ID, "../cat.png"), + "absolute path as workspace": workspaceRef(sess.ID, filepath.Join(root, "cat.png")), + } { + _, err := r.ResolveGeneratedFile(t.Context(), ref) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable, name) + } +} + +// TestResolveGeneratedFile_RootKindMismatchRefused: a recorded workspace +// path must not resolve through an external-rooted reference (and vice +// versa) — the record's root kind is part of the trust decision. +func TestResolveGeneratedFile_RootKindMismatchRefused(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-rootkind") + r, store := resolverTestRuntime(t, sess) + recordWorkspaceFile(t, store, sess.ID, root, "cat.png", []byte("png")) + + _, err := r.ResolveGeneratedFile(t.Context(), GeneratedFileRef{ + OwnerSessionID: sess.ID, + Root: chat.ArtifactRootExternal, + Path: "cat.png", + }) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable) +} + +func TestResolveGeneratedFile_MissingFileRefused(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-missing") + r, store := resolverTestRuntime(t, sess) + recordWorkspaceFile(t, store, sess.ID, root, "cat.png", []byte("png")) + require.NoError(t, os.Remove(filepath.Join(root, "cat.png"))) + + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable) +} + +// TestResolveGeneratedFile_SymlinkReplacementRefused: a symlink swapped in +// after the manifest record must not be followed, even when its target is +// another file INSIDE the workspace. +func TestResolveGeneratedFile_SymlinkReplacementRefused(t *testing.T) { + t.Parallel() + requireSymlinkSupport(t) + sess, root := workspaceSession(t, "sess-symlink") + r, store := resolverTestRuntime(t, sess) + recordWorkspaceFile(t, store, sess.ID, root, "cat.png", []byte("png")) + require.NoError(t, os.WriteFile(filepath.Join(root, ".env"), []byte("SECRET=1"), 0o600)) + require.NoError(t, os.Remove(filepath.Join(root, "cat.png"))) + require.NoError(t, os.Symlink(filepath.Join(root, ".env"), filepath.Join(root, "cat.png"))) + + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable) +} + +// TestResolveGeneratedFile_SymlinkParentRefused: replacing a recorded +// path's parent directory with a symlink must equally refuse, even when +// the link stays inside the workspace. +func TestResolveGeneratedFile_SymlinkParentRefused(t *testing.T) { + t.Parallel() + requireSymlinkSupport(t) + sess, root := workspaceSession(t, "sess-symlink-parent") + r, store := resolverTestRuntime(t, sess) + recordWorkspaceFile(t, store, sess.ID, root, "images/cat.png", []byte("png")) + + other := filepath.Join(root, "other") + require.NoError(t, os.MkdirAll(other, 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(other, "cat.png"), []byte("SECRET"), 0o644)) + require.NoError(t, os.RemoveAll(filepath.Join(root, "images"))) + require.NoError(t, os.Symlink(other, filepath.Join(root, "images"))) + + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "images/cat.png")) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable) +} + +// TestResolveGeneratedFile_CrossWorkspaceOwner: resolution uses the OWNING +// session's persisted workspace, never any other session's (or the +// viewer's) directory — a same-named file elsewhere must not shadow it. +func TestResolveGeneratedFile_CrossWorkspaceOwner(t *testing.T) { + t.Parallel() + owner, ownerRoot := workspaceSession(t, "sess-owner") + r, store := resolverTestRuntime(t, owner) + recordWorkspaceFile(t, store, owner.ID, ownerRoot, "cat.png", []byte("owner-bytes")) + + otherRoot := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(otherRoot, "cat.png"), []byte("other-bytes"), 0o644)) + require.NoError(t, store.AddSession(t.Context(), &session.Session{ID: "sess-viewer", WorkingDir: otherRoot})) + + resolved, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(owner.ID, "cat.png")) + + require.NoError(t, err) + assert.Equal(t, []byte("owner-bytes"), resolved.Data) + canonicalRoot, err := filepath.EvalSymlinks(ownerRoot) + require.NoError(t, err) + assert.Equal(t, filepath.Join(canonicalRoot, "cat.png"), resolved.Path) +} + +// TestResolveGeneratedFile_ParentWorkingDirFallback: an old sub-session +// without its own WorkingDir resolves through its parent's, mirroring +// session.ResolveWorkingDir. +func TestResolveGeneratedFile_ParentWorkingDirFallback(t *testing.T) { + t.Parallel() + parent, root := workspaceSession(t, "sess-parent") + r, store := resolverTestRuntime(t, parent) + child := &session.Session{ID: "sess-child", ParentID: parent.ID} + require.NoError(t, store.AddSession(t.Context(), child)) + recordWorkspaceFile(t, store, child.ID, root, "cat.png", []byte("png")) + + resolved, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(child.ID, "cat.png")) + + require.NoError(t, err) + assert.Equal(t, []byte("png"), resolved.Data) +} + +func TestResolveGeneratedFile_NoWorkspaceRootRefused(t *testing.T) { + t.Parallel() + sess := &session.Session{ID: "sess-rootless"} + r, store := resolverTestRuntime(t, sess) + require.NoError(t, manifestOf(t, store).AddGeneratedFile(t.Context(), session.GeneratedFile{ + SessionID: sess.ID, RelPath: "cat.png", Root: chat.ArtifactRootWorkspace, MimeType: "image/png", CreatedAt: time.Now(), + })) + + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable) +} + +func TestResolveGeneratedFile_ExternalSuccessAndLeafSymlinkRefused(t *testing.T) { + t.Parallel() + requireSymlinkSupport(t) + sess, _ := workspaceSession(t, "sess-external") + r, store := resolverTestRuntime(t, sess) + + target := filepath.Join(t.TempDir(), "cat.png") + require.NoError(t, os.WriteFile(target, []byte("external-bytes"), 0o644)) + require.NoError(t, manifestOf(t, store).AddGeneratedFile(t.Context(), session.GeneratedFile{ + SessionID: sess.ID, RelPath: target, Root: chat.ArtifactRootExternal, MimeType: "image/png", CreatedAt: time.Now(), + })) + ref := GeneratedFileRef{OwnerSessionID: sess.ID, Root: chat.ArtifactRootExternal, Path: target} + + resolved, err := r.ResolveGeneratedFile(t.Context(), ref) + require.NoError(t, err) + assert.Equal(t, []byte("external-bytes"), resolved.Data) + assert.Equal(t, target, resolved.Path, "the external path the user confirmed is the display path") + + require.NoError(t, os.Remove(target)) + secret := filepath.Join(t.TempDir(), "secret") + require.NoError(t, os.WriteFile(secret, []byte("SECRET"), 0o600)) + require.NoError(t, os.Symlink(secret, target)) + _, err = r.ResolveGeneratedFile(t.Context(), ref) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable, "an external leaf replaced by a symlink must be refused") +} + +func TestResolveGeneratedFile_OversizedFileRefused(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-oversized") + r, store := resolverTestRuntime(t, sess) + recordWorkspaceFile(t, store, sess.ID, root, "cat.png", []byte("png")) + require.NoError(t, os.Truncate(filepath.Join(root, "cat.png"), maxGeneratedFileBytes+1)) + + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable) +} + +// countingStore wraps a session store to observe how often resolution hits +// the store, proving the per-session cache. +type countingStore struct { + session.Store + + getSessions int + lookups int +} + +func (s *countingStore) GetSession(ctx context.Context, id string) (*session.Session, error) { + s.getSessions++ + return s.Store.GetSession(ctx, id) +} + +func (s *countingStore) LookupGeneratedFile(ctx context.Context, sessionID, relPath string) (*session.GeneratedFile, error) { + s.lookups++ + return s.Store.(session.GeneratedMediaManifest).LookupGeneratedFile(ctx, sessionID, relPath) +} + +func (s *countingStore) AddGeneratedFile(ctx context.Context, file session.GeneratedFile) error { + return s.Store.(session.GeneratedMediaManifest).AddGeneratedFile(ctx, file) +} + +func TestResolveGeneratedFile_CachesRootAndManifestPerSession(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-cache") + inner := session.NewInMemorySessionStore() + require.NoError(t, inner.AddSession(t.Context(), sess)) + store := &countingStore{Store: inner} + r := &LocalRuntime{sessionStore: store, now: time.Now} + recordWorkspaceFile(t, store, sess.ID, root, "cat.png", []byte("png")) + + for range 3 { + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + require.NoError(t, err) + } + + assert.Equal(t, 1, store.lookups, "repeat resolutions must reuse the cached manifest record") + assert.Equal(t, 1, store.getSessions, "repeat resolutions must reuse the cached workspace root") +} + +// TestResolveGeneratedFile_RecordSeedsAndRefreshesCache: materialization's +// recordGeneratedFile both seeds the cache (no store reads on first +// resolve) and refreshes it when the same path is re-recorded — a stale +// entry must not outlive an upsert. +func TestResolveGeneratedFile_RecordSeedsAndRefreshesCache(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-cache-seed") + inner := session.NewInMemorySessionStore() + require.NoError(t, inner.AddSession(t.Context(), sess)) + store := &countingStore{Store: inner} + r := &LocalRuntime{sessionStore: store, now: time.Now} + + require.NoError(t, os.WriteFile(filepath.Join(root, "cat.png"), []byte("png"), 0o644)) + require.NoError(t, r.recordGeneratedFile(t.Context(), sess.ID, chat.ArtifactRootWorkspace, "cat.png", "image/png")) + r.generatedFiles.setRoot(sess.ID, root) // materializeGeneratedMedia seeds this alongside + + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + require.NoError(t, err) + assert.Zero(t, store.lookups, "a freshly recorded file must resolve without a manifest read") + assert.Zero(t, store.getSessions, "a freshly recorded file must resolve without a session read") + + // Re-record the same path as external: the cached record must follow, + // so the old workspace-rooted reference stops resolving. + external := filepath.Join(t.TempDir(), "cat.png") + require.NoError(t, os.WriteFile(external, []byte("ext"), 0o644)) + require.NoError(t, r.recordGeneratedFile(t.Context(), sess.ID, chat.ArtifactRootExternal, external, "image/png")) + resolved, err := r.ResolveGeneratedFile(t.Context(), GeneratedFileRef{ + OwnerSessionID: sess.ID, Root: chat.ArtifactRootExternal, Path: external, + }) + require.NoError(t, err) + assert.Equal(t, []byte("ext"), resolved.Data) +} + +// TestResolveGeneratedFile_MaterializedEndToEnd drives the real +// materialization path and resolves the reference it persisted — the +// exact live TUI flow. +func TestResolveGeneratedFile_MaterializedEndToEnd(t *testing.T) { + r, store, _ := newMediaTestRuntime(t) + sess, root := workspaceSession(t, "sess-e2e") + require.NoError(t, store.AddSession(t.Context(), sess)) + + parts := r.materializeGeneratedMedia(t.Context(), sess, []chat.MediaDelta{ + {Data: []byte{0xAA}, MimeType: "image/png", Name: "cat.png", Size: 1}, + }, "root", nil) + require.Len(t, parts, 1) + src := parts[0].Document.Source + + resolved, err := r.ResolveGeneratedFile(t.Context(), GeneratedFileRef{ + OwnerSessionID: src.ArtifactOwnerSessionID, + Root: src.ArtifactRoot, + Path: src.ArtifactPath, + }) + + require.NoError(t, err) + assert.Equal(t, []byte{0xAA}, resolved.Data) + canonicalRoot, err := filepath.EvalSymlinks(root) + require.NoError(t, err) + assert.Equal(t, filepath.Join(canonicalRoot, "cat.png"), resolved.Path) +} diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 4d8dc2f912..ff7e3491df 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -1370,6 +1370,10 @@ func (r *LocalRuntime) materializeGeneratedMedia(ctx context.Context, sess *sess if rootErr != nil { slog.DebugContext(ctx, "No workspace root for generated media; dropping every media item, keeping the rest of the turn", "agent", agentName, "session_id", sess.ID, "error", rootErr) + } else { + // Seed the resolver cache so live inline rendering of this turn's + // media does not have to re-resolve the root from the store. + r.generatedFiles.setRoot(sess.ID, root) } parts := make([]chat.MessagePart, 0, len(media)) @@ -1419,7 +1423,9 @@ func (r *LocalRuntime) materializeGeneratedMedia(ctx context.Context, sess *sess // The file is already a real workspace deliverable, so keep the // reference; without the manifest record inline display will // refuse to render it (fail closed), which the user should hear - // about. res.RelPath is writer-sanitized and workspace-relative. + // about. res.RelPath is the writer-sanitized final path + // (workspace-relative, or the confirmed absolute external + // target), safe to show. slog.DebugContext(ctx, "Failed to record generated media in the manifest; the file was written but may not display inline", "agent", agentName, "session_id", sess.ID, "rel_path", res.RelPath, "error", err) if events != nil { @@ -1461,19 +1467,26 @@ func (r *LocalRuntime) sessionLookup() session.Lookup { } // recordGeneratedFile writes one manifest record after a successful -// write — materialization is the only writer of the manifest. +// write — materialization is the only writer of the manifest — and mirrors +// it into the resolver cache (an upsert for an existing path refreshes the +// cached entry). func (r *LocalRuntime) recordGeneratedFile(ctx context.Context, sessionID string, root chat.ArtifactRootKind, finalPath, mimeType string) error { manifest, ok := r.sessionStore.(session.GeneratedMediaManifest) if !ok { return fmt.Errorf("session store %T does not implement the generated-media manifest", r.sessionStore) } - return manifest.AddGeneratedFile(ctx, session.GeneratedFile{ + file := session.GeneratedFile{ SessionID: sessionID, RelPath: finalPath, Root: root, MimeType: mimeType, CreatedAt: r.now(), - }) + } + if err := manifest.AddGeneratedFile(ctx, file); err != nil { + return err + } + r.generatedFiles.setFile(file) + return nil } // workspacemediaWrite is [workspacemedia.Write] behind a package-level diff --git a/pkg/runtime/media_escape_test.go b/pkg/runtime/media_escape_test.go index 75c3fba495..5cb7a8756a 100644 --- a/pkg/runtime/media_escape_test.go +++ b/pkg/runtime/media_escape_test.go @@ -1,6 +1,7 @@ package runtime import ( + "fmt" "os" "path/filepath" "testing" @@ -417,3 +418,83 @@ func TestMaterializeGeneratedMedia_EscapeDirectoryTargetGenericName(t *testing.T _, err := os.Stat(final) require.NoError(t, err) } + +// TestMaterializeGeneratedMedia_EscapeTraversalRequested: a prompt-directed +// "../" path is an escape like any absolute one — the user is asked about +// the resolved absolute target and a decline redirects into the workspace. +func TestMaterializeGeneratedMedia_EscapeTraversalRequested(t *testing.T) { + r, _ := newEscapeTestRuntime(t) + sess, root := workspaceSession(t, "sess-escape-traversal") + + parts, sink, req := materializeWithAnswer(t, r, sess, escapeMedia("../escaped-cat.png"), tools.ElicitationActionDecline, nil) + + resolvedTarget := filepath.Join(filepath.Dir(root), "escaped-cat.png") + assert.Contains(t, req.Message, resolvedTarget, "the user must see the resolved absolute target, not the raw ../ path") + _, err := os.Stat(resolvedTarget) + assert.True(t, os.IsNotExist(err), "nothing may be written outside the workspace on decline") + + require.Len(t, parts, 2) + assert.Equal(t, "escaped-cat.png", parts[0].Document.Source.ArtifactPath) + assert.Equal(t, chat.ArtifactRootWorkspace, parts[0].Document.Source.ArtifactRoot) + require.Len(t, sink.warnings(), 1) + assert.Contains(t, sink.warnings()[0].Message, "outside the workspace") + assertSiblingPreserved(t, parts, root) +} + +// TestMaterializeGeneratedMedia_EscapeTildeRequested: "~/..." expands to +// the real home directory and is treated as an escape; without a user to +// ask it redirects into the workspace and never touches home. +func TestMaterializeGeneratedMedia_EscapeTildeRequested(t *testing.T) { + r, _ := newEscapeTestRuntime(t) + r.nonInteractive = true + r.OnElicitationRequest(func(e Event) { t.Errorf("no elicitation must be emitted in non-interactive mode, got %T", e) }) + sess, root := workspaceSession(t, "sess-escape-tilde") + filename := fmt.Sprintf("cagent-test-escape-%d.png", time.Now().UnixNano()) + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, escapeMedia("~/"+filename), "root", sink) + + require.Len(t, parts, 2) + assert.Equal(t, filename, parts[0].Document.Source.ArtifactPath) + assert.Equal(t, chat.ArtifactRootWorkspace, parts[0].Document.Source.ArtifactRoot) + home, err := os.UserHomeDir() + require.NoError(t, err) + _, err = os.Stat(filepath.Join(home, filename)) + assert.True(t, os.IsNotExist(err), "a redirected tilde path must never be written under home") + require.Len(t, sink.warnings(), 1) + assert.Contains(t, sink.warnings()[0].Message, "outside the workspace") + assertSiblingPreserved(t, parts, root) +} + +// TestMaterializeGeneratedMedia_SymlinkedParentRedirects: a lexically +// workspace-relative path whose parent directory is a symlink out of the +// workspace escapes at I/O time; the writer refuses it and the bytes are +// redirected into the workspace root like any unconfirmed escape. +func TestMaterializeGeneratedMedia_SymlinkedParentRedirects(t *testing.T) { + requireSymlinkSupport(t) + r, store := newEscapeTestRuntime(t) + r.OnElicitationRequest(func(e Event) { t.Errorf("a symlinked-parent escape must redirect, not elicit, got %T", e) }) + sess, root := workspaceSession(t, "sess-escape-symlink-parent") + outside := t.TempDir() + require.NoError(t, os.Symlink(outside, filepath.Join(root, "link"))) + + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, []chat.MediaDelta{ + {Data: []byte{0xAA}, MimeType: "image/png", RequestedPath: "link/cat.png", Size: 1}, + }, "root", sink) + + require.Len(t, parts, 1) + doc := parts[0].Document + assert.Equal(t, "cat.png", doc.Source.ArtifactPath) + assert.Equal(t, chat.ArtifactRootWorkspace, doc.Source.ArtifactRoot) + entries, err := os.ReadDir(outside) + require.NoError(t, err) + assert.Empty(t, entries, "nothing may be written through the symlinked parent") + data, err := os.ReadFile(filepath.Join(root, "cat.png")) + require.NoError(t, err) + assert.Equal(t, []byte{0xAA}, data) + _, err = manifestOf(t, store).LookupGeneratedFile(t.Context(), sess.ID, "cat.png") + require.NoError(t, err) + require.Len(t, sink.warnings(), 1) + assert.Contains(t, sink.warnings()[0].Message, "escapes the workspace") +} diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index 375f2f8026..883424c5dd 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -252,12 +252,16 @@ type LocalRuntime struct { elicitationSinkMu sync.RWMutex onElicitationRequest func(Event) sessionStore session.Store - workingDir string // Working directory for hooks execution - env []string // Environment variables for hooks execution - modelSwitcherCfg *ModelSwitcherConfig - providerRegistry *provider.Registry - gatewayModels gatewayModelsCache - dmrModels dmrModelsCache + // generatedFiles caches per-owner-session workspace roots and manifest + // records for [LocalRuntime.ResolveGeneratedFile]. Seeded by + // materialization, filled lazily for restored sessions. + generatedFiles generatedFileCache + workingDir string // Working directory for hooks execution + env []string // Environment variables for hooks execution + modelSwitcherCfg *ModelSwitcherConfig + providerRegistry *provider.Registry + gatewayModels gatewayModelsCache + dmrModels dmrModelsCache // hooksRegistry is the runtime-private hooks.Registry used to build // every Executor. It carries the runtime-owned builtin hooks diff --git a/pkg/tui/components/message/message.go b/pkg/tui/components/message/message.go index 831b674310..3e8762292f 100644 --- a/pkg/tui/components/message/message.go +++ b/pkg/tui/components/message/message.go @@ -322,7 +322,7 @@ func (mv *messageModel) isSpinnerDriven() bool { case types.MessageTypeSpinner, types.MessageTypeLoading: return true case types.MessageTypeAssistant: - return mv.message.Content == "" + return mv.message.Content == "" && len(mv.message.AssistantMedia) == 0 } return false } @@ -382,7 +382,7 @@ func (mv *messageModel) render(width int) string { noTopPaddingStyle := messageStyle.PaddingTop(0) return noTopPaddingStyle.Width(width).Render(topRow + "\n" + content) case types.MessageTypeAssistant: - if msg.Content == "" { + if msg.Content == "" && len(msg.AssistantMedia) == 0 { return mv.spinner.View() } @@ -399,6 +399,7 @@ func (mv *messageModel) render(width int) string { codeBlocks = nil } rendered, codeBlocks = replaceMarkdownImagePlaceholders(rendered, codeBlocks, imagePlaceholders) + rendered = appendAssistantMediaLines(rendered, msg.AssistantMedia, innerRenderWidth) var prefix string if !mv.sameAgentAsPrevious(msg) { @@ -565,6 +566,47 @@ func replaceMarkdownImagePlaceholders(rendered string, codeBlocks []markdown.Cod return strings.Join(result, "\n"), codeBlocks } +// appendAssistantMediaLines appends generated-media blocks after the +// rendered markdown, preserving the text-then-media order of the assistant +// turn. Appending never shifts earlier lines, so code-block coordinates +// computed for the markdown remain valid. +func appendAssistantMediaLines(rendered string, media []types.AssistantMedia, width int) string { + blocks := make([]string, 0, len(media)) + for _, m := range media { + if lines := assistantMediaLines(m, width); len(lines) > 0 { + blocks = append(blocks, strings.Join(lines, "\n")) + } + } + if len(blocks) == 0 { + return rendered + } + joined := strings.Join(blocks, "\n\n") + if rendered = strings.TrimRight(rendered, "\n\r\t "); rendered == "" { + return joined + } + return rendered + "\n\n" + joined +} + +// assistantMediaLines renders one generated-media item: a muted name label +// plus kitty marker rows when the image is renderable (mirroring the +// markdown-image layout above), or the item's safe textual fallback when +// graphics are unavailable or the image never decoded. +func assistantMediaLines(media types.AssistantMedia, width int) []string { + if media.Image != nil { + if markers := tuiimage.RenderMarkers(*media.Image, width); len(markers) > 0 { + lines := make([]string, 0, len(markers)+1) + if media.Image.Name != "" { + lines = append(lines, " "+styles.MutedStyle.Render(media.Image.Name)) + } + return append(lines, markers...) + } + } + if media.Fallback == "" { + return nil + } + return []string{styles.MutedStyle.Width(width).Render(media.Fallback)} +} + // renderAssistantMarkdown renders streamed assistant content using a per-message // IncrementalRenderer. The renderer remembers the last rendered stable prefix // so each new chunk only re-parses the trailing region. The first render at a diff --git a/pkg/tui/components/message/message_test.go b/pkg/tui/components/message/message_test.go index cccb24d725..89ecfbd2b8 100644 --- a/pkg/tui/components/message/message_test.go +++ b/pkg/tui/components/message/message_test.go @@ -464,3 +464,78 @@ func TestAgentReturnRespectsNarrowWidths(t *testing.T) { } } } + +func testInlineImage(t *testing.T, name string) tuiimage.Inline { + t.Helper() + img := stdimage.NewRGBA(stdimage.Rect(0, 0, 2, 1)) + img.Set(0, 0, color.RGBA{R: 255, A: 255}) + var data bytes.Buffer + require.NoError(t, png.Encode(&data, img)) + inline, ok := tuiimage.FromBytes(name, "image/png", data.Bytes()) + require.True(t, ok) + return inline +} + +func TestAssistantMediaRendersInlineAfterText(t *testing.T) { + tuiimage.SetRenderingEnabled(true) + + inline := testInlineImage(t, "cat.png") + msg := types.Agent(types.MessageTypeAssistant, "assistant", "Here is your cat:") + msg.AssistantMedia = []types.AssistantMedia{{Image: &inline, Fallback: `Generated image "cat.png" saved to: /tmp/cat.png`}} + mv := New(animation.NewRuntime(), msg, nil) + mv.SetSize(80, 0) + + view := mv.View() + assert.Contains(t, view, "cagent-image", "generated media must emit terminal image markers") + plain := ansi.Strip(view) + assert.Contains(t, plain, "cat.png", "image name must label the rendered image") + assert.NotContains(t, plain, "saved to:", "the textual fallback must not show when the image renders inline") + assert.Less(t, strings.Index(plain, "Here is your cat:"), strings.Index(view, "cagent-image"), + "streamed text must precede the generated media in the same turn") +} + +func TestAssistantMediaOnlyReplacesSpinnerWithVisibleContent(t *testing.T) { + tuiimage.SetRenderingEnabled(true) + + inline := testInlineImage(t, "cat.png") + msg := types.Agent(types.MessageTypeAssistant, "assistant", "") + msg.AssistantMedia = []types.AssistantMedia{{Image: &inline, Fallback: `Generated image "cat.png" saved to: /tmp/cat.png`}} + mv := New(animation.NewRuntime(), msg, nil) + mv.SetSize(80, 0) + + assert.False(t, mv.isSpinnerDriven(), "a media-only assistant message is real content, not a spinner placeholder") + view := mv.View() + assert.Contains(t, view, "cagent-image", "a media-only turn must render the image, not a spinner") + assert.Contains(t, ansi.Strip(view), "cat.png") +} + +func TestAssistantMediaGraphicsDisabledShowsFallback(t *testing.T) { + tuiimage.SetRenderingEnabled(false) + defer tuiimage.SetRenderingEnabled(true) + + inline := testInlineImage(t, "cat.png") + msg := types.Agent(types.MessageTypeAssistant, "assistant", "") + msg.AssistantMedia = []types.AssistantMedia{{Image: &inline, Fallback: `Generated image "cat.png" saved to: /tmp/artifacts/sess/cat.png`}} + mv := New(animation.NewRuntime(), msg, nil) + mv.SetSize(80, 0) + + view := mv.View() + assert.NotContains(t, view, "cagent-image", "no image markers may be emitted while graphics are disabled") + plain := ansi.Strip(view) + assert.Contains(t, plain, `Generated image "cat.png" saved to:`, "the fallback must make the generated file visible") + assert.Contains(t, strings.ReplaceAll(plain, "\n", ""), "/tmp/artifacts/sess/cat.png") +} + +func TestAssistantMediaUnrenderableImageShowsFallback(t *testing.T) { + tuiimage.SetRenderingEnabled(true) + + msg := types.Agent(types.MessageTypeAssistant, "assistant", "Result:") + msg.AssistantMedia = []types.AssistantMedia{{Fallback: `Generated image "cat.png" is unavailable.`}} + mv := New(animation.NewRuntime(), msg, nil) + mv.SetSize(80, 0) + + plain := ansi.Strip(mv.View()) + assert.Contains(t, plain, `Generated image "cat.png" is unavailable.`) + assert.Less(t, strings.Index(plain, "Result:"), strings.Index(plain, "unavailable"), + "text must keep preceding the failed media item") +} diff --git a/pkg/tui/components/messages/generated_media_test.go b/pkg/tui/components/messages/generated_media_test.go new file mode 100644 index 0000000000..0d653acc49 --- /dev/null +++ b/pkg/tui/components/messages/generated_media_test.go @@ -0,0 +1,120 @@ +package messages + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/tui/animation" + "github.com/docker/docker-agent/pkg/tui/service" + "github.com/docker/docker-agent/pkg/tui/types" +) + +func newMediaTestModel(t *testing.T) *model { + t.Helper() + m := NewScrollableView(animation.NewRuntime(), 80, 24, &service.SessionState{}).(*model) + m.SetSize(80, 24) + return m +} + +func assistantSessionItem(agent, content string) session.Item { + return session.NewMessageItem(&session.Message{ + AgentName: agent, + Message: chat.Message{Role: chat.MessageRoleAssistant, Content: content}, + }) +} + +// TestLoadFromSession_AttachesGeneratedMediaAtPosition: restored media joins +// the assistant message built for its exact session position — not the +// newest message. +func TestLoadFromSession_AttachesGeneratedMediaAtPosition(t *testing.T) { + t.Parallel() + m := newMediaTestModel(t) + sess := &session.Session{ + ID: "sess-restore", + Messages: []session.Item{ + session.NewMessageItem(&session.Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "draw"}}), + assistantSessionItem("root", "Here is your cat:"), + assistantSessionItem("root", "Anything else?"), + }, + } + media := map[int][]types.AssistantMedia{ + 1: {{ID: 7, Fallback: `Generated image "cat.png" is unavailable.`}}, + } + + m.LoadFromSession(sess, media) + + require.Len(t, m.messages, 3) + require.Len(t, m.messages[1].AssistantMedia, 1, "media must join the message at its session position") + assert.EqualValues(t, 7, m.messages[1].AssistantMedia[0].ID) + assert.Empty(t, m.messages[2].AssistantMedia, "later messages must stay media-free") +} + +// TestLoadFromSession_MediaOnlyAssistantMessage: an assistant message with +// no text but restored media still becomes a visible message, mirroring +// AppendAssistantMedia's media-only turn. +func TestLoadFromSession_MediaOnlyAssistantMessage(t *testing.T) { + t.Parallel() + m := newMediaTestModel(t) + sess := &session.Session{ + ID: "sess-media-only", + Messages: []session.Item{assistantSessionItem("root", "")}, + } + + m.LoadFromSession(sess, map[int][]types.AssistantMedia{ + 0: {{ID: 3, Fallback: `Generated image "cat.png" is unavailable.`}}, + }) + + require.Len(t, m.messages, 1) + assert.Equal(t, types.MessageTypeAssistant, m.messages[0].Type) + assert.Empty(t, m.messages[0].Content) + require.Len(t, m.messages[0].AssistantMedia, 1) + + // Without media (e.g. no resolver capability) the empty turn stays + // invisible, as before. + m.LoadFromSession(sess, nil) + assert.Empty(t, m.messages) +} + +// TestUpdateAssistantMedia_ReplacesByID: a resolution result replaces its +// placeholder wherever it sits — including a non-final message — while +// zero-ID and unmatched items stay untouched. +func TestUpdateAssistantMedia_ReplacesByID(t *testing.T) { + t.Parallel() + m := newMediaTestModel(t) + m.AppendAssistantMedia("root", []types.AssistantMedia{ + {ID: 1, Fallback: "placeholder one"}, + {Fallback: "legacy, final"}, + }) + m.AddUserMessage("and another") + m.AppendAssistantMedia("root", []types.AssistantMedia{{ID: 2, Fallback: "placeholder two"}}) + + m.UpdateAssistantMedia([]types.AssistantMedia{ + {ID: 1, Fallback: "resolved one"}, + {ID: 99, Fallback: "unknown id, dropped"}, + }) + + require.Len(t, m.messages, 3) + first := m.messages[0].AssistantMedia + require.Len(t, first, 2) + assert.Equal(t, "resolved one", first[0].Fallback, "the matching placeholder must be replaced in place") + assert.Equal(t, "legacy, final", first[1].Fallback, "zero-ID items are final and untouched") + assert.Equal(t, "placeholder two", m.messages[2].AssistantMedia[0].Fallback, + "an unmatched placeholder must keep waiting for its own result") +} + +// TestUpdateAssistantMedia_StaleResultIsNoOp: results whose placeholders no +// longer exist (e.g. the list was reloaded) change nothing. +func TestUpdateAssistantMedia_StaleResultIsNoOp(t *testing.T) { + t.Parallel() + m := newMediaTestModel(t) + m.AppendAssistantMedia("root", []types.AssistantMedia{{ID: 5, Fallback: "placeholder"}}) + + cmd := m.UpdateAssistantMedia([]types.AssistantMedia{{ID: 42, Fallback: "stale"}}) + + assert.Nil(t, cmd) + assert.Equal(t, "placeholder", m.messages[0].AssistantMedia[0].Fallback) +} diff --git a/pkg/tui/components/messages/messages.go b/pkg/tui/components/messages/messages.go index 087828677a..adf3316064 100644 --- a/pkg/tui/components/messages/messages.go +++ b/pkg/tui/components/messages/messages.go @@ -87,13 +87,26 @@ type Model interface { AppendToolOutput(msg *runtime.ToolCallOutputEvent) tea.Cmd AddToolResult(msg *runtime.ToolCallResponseEvent, status types.ToolStatus) tea.Cmd AppendToLastMessage(agentName, content string) tea.Cmd + // AppendAssistantMedia attaches generated media to the agent's current + // assistant message (or starts a media-only one), so it renders in the + // same assistant turn as the streamed text. + AppendAssistantMedia(agentName string, media []types.AssistantMedia) tea.Cmd + // UpdateAssistantMedia replaces previously attached media items — + // wherever they sit in the list — with the given resolved items, matched + // by types.AssistantMedia.ID. Items with unknown or zero IDs are + // ignored, so a stale asynchronous result is harmless. + UpdateAssistantMedia(media []types.AssistantMedia) tea.Cmd AppendReasoning(agentName, content string) tea.Cmd AddShellOutputMessage(content string) tea.Cmd // AddAgentReturn appends the UI-only "child returned control to parent" // delegation transition. It is never persisted, so it does not reappear // when the session is reloaded. AddAgentReturn(fromAgent, toAgent string) tea.Cmd - LoadFromSession(sess *session.Session) tea.Cmd + // LoadFromSession rebuilds the list from a persisted session. + // generatedMedia carries the restored generated-media items to attach, + // keyed by the owning message's index in sess.Messages; nil when the + // caller cannot resolve generated media. + LoadFromSession(sess *session.Session, generatedMedia map[int][]types.AssistantMedia) tea.Cmd // StopAnimations unregisters every view from the animation coordinator. // Call it when the list is discarded or its host view goes away, so @@ -1141,7 +1154,7 @@ func (m *model) shouldCacheMessage(index int) bool { case types.MessageTypeToolResult: return true case types.MessageTypeAssistant: - return strings.Trim(msg.Content, "\r\n\t ") != "" + return strings.Trim(msg.Content, "\r\n\t ") != "" || len(msg.AssistantMedia) > 0 case types.MessageTypeAssistantReasoningBlock: // Cacheable once spinners/fades have settled. Content mutations go // through invalidateItem, which drops any stale entry. @@ -1447,7 +1460,7 @@ func (m *model) addMessage(msg *types.Message) tea.Cmd { return tea.Batch(cmds...) } -func (m *model) LoadFromSession(sess *session.Session) tea.Cmd { +func (m *model) LoadFromSession(sess *session.Session, generatedMedia map[int][]types.AssistantMedia) tea.Cmd { appendSessionMessage := func(msg *types.Message, view layout.Model) { m.messages = append(m.messages, msg) m.views = append(m.views, view) @@ -1557,9 +1570,14 @@ func (m *model) LoadFromSession(sess *session.Session) tea.Cmd { m.messages[lastIdx].Content += smsg.Message.ReasoningContent } - // Step 2: Handle assistant content - this breaks the reasoning block chain - if hasContent { + // Step 2: Handle assistant content — this breaks the reasoning + // block chain. Restored generated media joins the same message + // (or forms a media-only one), mirroring AppendAssistantMedia's + // live behavior. + restoredMedia := generatedMedia[pos] + if hasContent || len(restoredMedia) > 0 { msg := types.Agent(types.MessageTypeAssistant, smsg.AgentName, smsg.Message.Content) + msg.AssistantMedia = restoredMedia appendSessionMessage(msg, m.createMessageView(msg)) } @@ -1753,6 +1771,70 @@ func (m *model) AppendToLastMessage(agentName, content string) tea.Cmd { return m.addMessage(types.Agent(types.MessageTypeAssistant, agentName, content)) } +// AppendAssistantMedia mirrors AppendToLastMessage for generated media: it +// replaces a pending spinner and joins the agent's current assistant +// message so the media renders inside the same turn as the streamed text, +// or starts a media-only assistant message when there is none. +func (m *model) AppendAssistantMedia(agentName string, media []types.AssistantMedia) tea.Cmd { + if len(media) == 0 { + return nil + } + m.removeSpinner() + + if len(m.messages) > 0 { + lastIdx := len(m.messages) - 1 + lastMsg := m.messages[lastIdx] + if lastMsg.Type == types.MessageTypeAssistant && lastMsg.Sender == agentName { + lastMsg.AssistantMedia = append(lastMsg.AssistantMedia, media...) + cmd := m.views[lastIdx].(message.Model).SetMessage(lastMsg) + m.invalidateItem(lastIdx) + return cmd + } + } + + msg := types.Agent(types.MessageTypeAssistant, agentName, "") + msg.AssistantMedia = media + return m.addMessage(msg) +} + +// UpdateAssistantMedia replaces attached media items in place by ID. See +// Model.UpdateAssistantMedia. +func (m *model) UpdateAssistantMedia(media []types.AssistantMedia) tea.Cmd { + byID := make(map[uint64]types.AssistantMedia, len(media)) + for _, item := range media { + if item.ID != 0 { + byID[item.ID] = item + } + } + if len(byID) == 0 { + return nil + } + + var cmds []tea.Cmd + for i, msg := range m.messages { + changed := false + for j, item := range msg.AssistantMedia { + if resolved, ok := byID[item.ID]; ok { + msg.AssistantMedia[j] = resolved + changed = true + } + } + if !changed { + continue + } + if view, ok := m.views[i].(message.Model); ok { + if cmd := view.SetMessage(msg); cmd != nil { + cmds = append(cmds, cmd) + } + } + m.invalidateItem(i) + } + if len(cmds) == 0 { + return nil + } + return tea.Batch(cmds...) +} + func (m *model) AppendReasoning(agentName, content string) tea.Cmd { m.removeSpinner() diff --git a/pkg/tui/components/messages/messages_test.go b/pkg/tui/components/messages/messages_test.go index a3cc7c95cc..09041995fe 100644 --- a/pkg/tui/components/messages/messages_test.go +++ b/pkg/tui/components/messages/messages_test.go @@ -117,7 +117,7 @@ func TestLoadFromSessionIncludesReasoningContent(t *testing.T) { }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // Expect: user message + reasoning block + assistant content = 3 messages require.Len(t, m.messages, 3) @@ -165,7 +165,7 @@ func TestLoadFromSessionReasoningOrderWithToolCalls(t *testing.T) { }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // Expect: reasoning block (reasoning only) + assistant content + standalone tool call = 3 messages // The content breaks the reasoning block chain, so tool calls become standalone @@ -207,7 +207,7 @@ func TestLoadFromSessionReasoningOnlyNoContent(t *testing.T) { }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // Expect: just the reasoning block (no assistant content) require.Len(t, m.messages, 1) @@ -246,7 +246,7 @@ func TestLoadFromSessionToolCallsOnlyNoReasoning(t *testing.T) { }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // Expect: assistant content + standalone tool call = 2 messages // Tool calls without reasoning should NOT go into a reasoning block @@ -314,7 +314,7 @@ func TestLoadFromSessionWithToolResults(t *testing.T) { }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // Expect: reasoning block (reasoning + 2 tool calls with results) + assistant content = 2 messages require.Len(t, m.messages, 2) @@ -427,7 +427,7 @@ func TestLoadFromSessionCombinesConsecutiveReasoningBlocks(t *testing.T) { }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // Should have: 1 combined reasoning block + 1 assistant content = 2 messages require.Len(t, m.messages, 2, "consecutive reasoning blocks should be combined into one") @@ -493,7 +493,7 @@ func TestLoadFromSessionStandaloneToolCallsWithResults(t *testing.T) { }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // Expect: standalone tool call (not in reasoning block) require.Len(t, m.messages, 1) @@ -542,7 +542,7 @@ func TestLoadFromSessionToolCallsDuringReasoningNoContent(t *testing.T) { }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // Expect: reasoning block only (tool call inside it) require.Len(t, m.messages, 1) @@ -590,7 +590,7 @@ func TestLoadFromSessionReasoningWithContentToolResultsStandalone(t *testing.T) }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) require.Len(t, m.messages, 3) @@ -656,7 +656,7 @@ func TestLoadFromSessionMultipleStandaloneToolCallsWithContentAndResults(t *test }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) require.Len(t, m.messages, 3) @@ -1325,7 +1325,7 @@ func TestLoadFromSessionReasoningBlockAgentBadges(t *testing.T) { }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // user + (reasoning block + content) x 2 require.Len(t, m.messages, 5) @@ -1376,7 +1376,7 @@ func TestLoadFromSessionReasoningAfterTransferTaskShowsAgentBadge(t *testing.T) }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) // standalone transfer_task tool call + developer reasoning block require.Len(t, m.messages, 2) @@ -1709,3 +1709,76 @@ func TestAgentReturnIsInertInList(t *testing.T) { assert.Contains(t, out, types.AgentReturnLabel) assert.NotContains(t, out, types.MessageCopyLabel) } + +func assistantTestMedia(fallback string) []types.AssistantMedia { + return []types.AssistantMedia{{Fallback: fallback}} +} + +func TestAppendAssistantMediaJoinsSameAgentAssistantMessage(t *testing.T) { + t.Parallel() + + m := NewScrollableView(animation.NewRuntime(), 80, 24, &service.SessionState{}).(*model) + m.AddUserMessage("draw") + m.AppendToLastMessage("root", "Here it is:") + require.Equal(t, 1, m.MessageTypeCount(types.MessageTypeAssistant)) + + // The join path returns SetMessage's follow-up command, which is nil + // without markdown references — same contract as AppendToLastMessage — + // so assert on list state, not the command. + m.AppendAssistantMedia("root", assistantTestMedia("img-1")) + + require.Equal(t, 1, m.MessageTypeCount(types.MessageTypeAssistant), + "media must join the agent's current assistant message, not start a new one") + last := m.messages[len(m.messages)-1] + assert.Equal(t, "Here it is:", last.Content) + require.Len(t, last.AssistantMedia, 1) + assert.Equal(t, "img-1", last.AssistantMedia[0].Fallback) + + m.AppendAssistantMedia("root", assistantTestMedia("img-2")) + require.Len(t, last.AssistantMedia, 2, "subsequent media joins the same turn") + assert.Equal(t, []string{"img-1", "img-2"}, []string{last.AssistantMedia[0].Fallback, last.AssistantMedia[1].Fallback}, + "media order must follow append order") +} + +func TestAppendAssistantMediaStartsMediaOnlyMessageAndReplacesSpinner(t *testing.T) { + t.Parallel() + + m := NewScrollableView(animation.NewRuntime(), 80, 24, &service.SessionState{}).(*model) + m.AddAssistantMessage("root", "") + require.Equal(t, 1, m.MessageTypeCount(types.MessageTypeSpinner)) + + require.NotNil(t, m.AppendAssistantMedia("root", assistantTestMedia("img-1"))) + + assert.Zero(t, m.MessageTypeCount(types.MessageTypeSpinner), "the pending spinner must be replaced by the media message") + require.Equal(t, 1, m.MessageTypeCount(types.MessageTypeAssistant)) + last := m.messages[len(m.messages)-1] + assert.Empty(t, last.Content) + assert.Equal(t, "root", last.Sender) + require.Len(t, last.AssistantMedia, 1) +} + +func TestAppendAssistantMediaDifferentAgentStartsNewMessage(t *testing.T) { + t.Parallel() + + m := NewScrollableView(animation.NewRuntime(), 80, 24, &service.SessionState{}).(*model) + m.AddUserMessage("draw") + m.AppendToLastMessage("root", "parent text") + require.Equal(t, 1, m.MessageTypeCount(types.MessageTypeAssistant)) + + require.NotNil(t, m.AppendAssistantMedia("researcher", assistantTestMedia("img-1"))) + + require.Equal(t, 2, m.MessageTypeCount(types.MessageTypeAssistant), + "another agent's media must not be attached to the previous agent's message") + assert.Empty(t, m.messages[1].AssistantMedia) + assert.Equal(t, "researcher", m.messages[2].Sender) +} + +func TestAppendAssistantMediaEmptyIsNoOp(t *testing.T) { + t.Parallel() + + m := NewScrollableView(animation.NewRuntime(), 80, 24, &service.SessionState{}).(*model) + m.AddAssistantMessage("root", "") + + assert.Nil(t, m.AppendAssistantMedia("root", nil)) + assert.Equal(t, 1, m.MessageTypeCount(types.MessageTypeSpinner), "empty media must not disturb the pending spinner") +} diff --git a/pkg/tui/page/chat/chat.go b/pkg/tui/page/chat/chat.go index 825db8ad7a..dbe67991a0 100644 --- a/pkg/tui/page/chat/chat.go +++ b/pkg/tui/page/chat/chat.go @@ -171,11 +171,12 @@ type Page interface { // the current app.Session().ID after a session restore or in-place // replace. SetRoutingID(id string) - // TakeRoutedTimers returns and clears the routed one-shot timer commands - // armed by the most recent Update. The active page's Update already - // returns them inside its regular command; the appModel calls this for - // background pages — whose regular commands are discarded — so - // presentation deadlines keep running while a tab is hidden. + // TakeRoutedTimers returns and clears the routed one-shot commands + // (presentation timers, generated-media resolution) armed by the most + // recent Update. The active page's Update already returns them inside + // its regular command; the appModel calls this for background pages — + // whose regular commands are discarded — so those deadlines and + // resolutions keep running while a tab is hidden. TakeRoutedTimers() tea.Cmd VisualGeneration() uint64 } @@ -226,9 +227,10 @@ type chatPage struct { // addressed to; empty for standalone pages (timers then fire unrouted, // which is correct when this is the only page). routingID string - // pendingTimers holds the routed timer commands armed by the current - // Update, so they can be re-collected via TakeRoutedTimers when the - // regular command is discarded (background tabs). + // pendingTimers holds the routed one-shot commands (presentation timers, + // generated-media resolution) armed by the current Update, so they can be + // re-collected via TakeRoutedTimers when the regular command is discarded + // (background tabs). pendingTimers []tea.Cmd // Track whether we've received content from an assistant response @@ -493,7 +495,11 @@ func (p *chatPage) Init() tea.Cmd { if sess := p.app.Session(); sess != nil { p.sidebar.LoadFromSession(sess) if len(sess.Messages) > 0 { - cmds = append(cmds, p.messages.LoadFromSession(sess)) + restoredMedia, mediaRequests := p.collectRestoredGeneratedMedia(sess) + cmds = append(cmds, p.messages.LoadFromSession(sess, restoredMedia)) + if resolve := p.resolveGeneratedMediaCmd(mediaRequests); resolve != nil { + cmds = append(cmds, resolve) + } } } @@ -625,6 +631,9 @@ func (p *chatPage) update(msg tea.Msg) (layout.Model, tea.Cmd) { case msgtypes.ClearQueueMsg: return p.handleClearQueue() + case generatedMediaResolvedMsg: + return p, p.messages.UpdateAssistantMedia(msg.media) + case msgtypes.ThemeChangedMsg: // Theme changed - forward to all child components to invalidate caches var cmds []tea.Cmd diff --git a/pkg/tui/page/chat/generated_media.go b/pkg/tui/page/chat/generated_media.go new file mode 100644 index 0000000000..8f865aff58 --- /dev/null +++ b/pkg/tui/page/chat/generated_media.go @@ -0,0 +1,216 @@ +package chat + +import ( + "context" + "fmt" + "log/slog" + "strings" + "sync/atomic" + + tea "charm.land/bubbletea/v2" + + "github.com/docker/docker-agent/pkg/app" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/runtime" + "github.com/docker/docker-agent/pkg/session" + tuiimage "github.com/docker/docker-agent/pkg/tui/image" + msgtypes "github.com/docker/docker-agent/pkg/tui/messages" + "github.com/docker/docker-agent/pkg/tui/types" +) + +// Generated media rendering. +// +// The run loop materializes model-generated images into the owning +// session's workspace and persists manifest-gated references (see +// materializeGeneratedMedia in pkg/runtime/loop.go). This file renders +// those references: a sanitized "unavailable" placeholder is attached +// synchronously, and the actual bytes + validated canonical path are +// resolved through the runtime's generated-file resolver capability inside +// a tea.Cmd — never synchronously in Update — then swapped in by ID via +// messages.Model.UpdateAssistantMedia. Runtimes without the capability +// (e.g. remote) render nothing, as before. + +// generatedMediaIDs issues process-unique placeholder IDs, so a resolution +// result can never match a placeholder from another page or an earlier +// session load. +var generatedMediaIDs atomic.Uint64 + +// generatedMediaRequest pairs one pending placeholder with the reference +// the resolver needs. +type generatedMediaRequest struct { + id uint64 + ref runtime.GeneratedFileRef + name string + mimeType string +} + +// generatedMediaResolvedMsg delivers asynchronously resolved media items +// back to the page, which applies them by ID. +type generatedMediaResolvedMsg struct { + media []types.AssistantMedia +} + +// handleMessageAdded surfaces model-generated media in the turn it was +// produced: the run loop announces the persisted assistant message via +// MessageAddedEvent; without this handler the file sits in the workspace +// with nothing visible in the chat. +func (p *chatPage) handleMessageAdded(msg *runtime.MessageAddedEvent) tea.Cmd { + if msg.Message == nil { + // The payload is process-local (json:"-"): events decoded from a + // remote runtime carry only IDs. Nothing to resolve or render. + return nil + } + if p.streamCancelled || msg.Message.Message.Role != chat.MessageRoleAssistant { + return nil + } + if !p.app.CanResolveGeneratedFiles() { + return nil + } + placeholders, requests := generatedImageMedia(msg.Message.Message.MultiContent) + if len(placeholders) == 0 { + return nil + } + + p.hasReceivedAssistantContent = true + p.setPendingResponse(false) + agentName := msg.Message.AgentName + if agentName == "" { + agentName = msg.AgentName + } + return tea.Batch( + p.sidebar.SetAgentActivity(agentName), + p.messages.AppendAssistantMedia(agentName, placeholders), + p.resolveGeneratedMediaCmd(requests), + ) +} + +// collectRestoredGeneratedMedia extracts the generated media of every +// restored assistant message, keyed by its index in sess.Messages, for +// messages.Model.LoadFromSession, plus the resolution requests to run +// asynchronously. Nil when the runtime cannot resolve generated files. +func (p *chatPage) collectRestoredGeneratedMedia(sess *session.Session) (map[int][]types.AssistantMedia, []generatedMediaRequest) { + if !p.app.CanResolveGeneratedFiles() { + return nil, nil + } + var restored map[int][]types.AssistantMedia + var requests []generatedMediaRequest + for pos, item := range sess.Messages { + if !item.IsMessage() || item.Message.Implicit || item.Message.Message.Role != chat.MessageRoleAssistant { + continue + } + placeholders, reqs := generatedImageMedia(item.Message.Message.MultiContent) + if len(placeholders) == 0 { + continue + } + if restored == nil { + restored = make(map[int][]types.AssistantMedia) + } + restored[pos] = placeholders + requests = append(requests, reqs...) + } + return restored, requests +} + +// generatedImageMedia extracts the generated images from an assistant +// message's parts: document parts carrying an owner-qualified generated-file +// reference and an image MIME type. Every extracted item starts as a +// sanitized "unavailable" placeholder; items whose root kind the resolver +// supports additionally get a resolution request. References with an +// unknown (empty) root kind stay unavailable by design. User attachments +// (inline sources) and ownerless references are not extracted. +func generatedImageMedia(parts []chat.MessagePart) ([]types.AssistantMedia, []generatedMediaRequest) { + var media []types.AssistantMedia + var requests []generatedMediaRequest + for _, part := range parts { + doc := part.Document + if part.Type != chat.MessagePartTypeDocument || doc == nil { + continue + } + src := doc.Source + if src.ArtifactPath == "" || src.ArtifactOwnerSessionID == "" { + continue + } + if !chat.IsImageMimeType(doc.MimeType) { + continue + } + + name := chat.SanitizeDisplayName(doc.Name) + if name == "" { + name = "generated media" + } + item := types.AssistantMedia{Fallback: fmt.Sprintf("Generated image %q is unavailable.", name)} + if src.ArtifactRoot == chat.ArtifactRootWorkspace || src.ArtifactRoot == chat.ArtifactRootExternal { + item.ID = generatedMediaIDs.Add(1) + requests = append(requests, generatedMediaRequest{ + id: item.ID, + ref: runtime.GeneratedFileRef{ + OwnerSessionID: src.ArtifactOwnerSessionID, + Root: src.ArtifactRoot, + Path: src.ArtifactPath, + }, + name: name, + mimeType: doc.MimeType, + }) + } + media = append(media, item) + } + return media, requests +} + +// resolveGeneratedMediaCmd resolves the requested items on a background +// goroutine and routes the results back to this page (its tab may be +// hidden — or another tab active — by the time they arrive). The command +// is also recorded like a routed timer so an update on a hidden tab keeps +// the resolution armed. +func (p *chatPage) resolveGeneratedMediaCmd(requests []generatedMediaRequest) tea.Cmd { + if len(requests) == 0 { + return nil + } + application, ctx, routingID := p.app, p.ctx(), p.routingID + cmd := func() tea.Msg { + media := make([]types.AssistantMedia, 0, len(requests)) + for _, req := range requests { + media = append(media, resolveGeneratedImage(ctx, application, req)) + } + var inner tea.Msg = generatedMediaResolvedMsg{media: media} + if routingID == "" { + return inner + } + return msgtypes.RoutedMsg{SessionID: routingID, Inner: inner} + } + p.pendingTimers = append(p.pendingTimers, cmd) + return cmd +} + +// resolveGeneratedImage resolves one generated image through the runtime's +// manifest-gated resolver and prepares it for terminal rendering. Every +// returned item carries a Fallback built exclusively from safe display +// data — the sanitized name and, when resolution validated it, the +// canonical workspace path. Raw resolver errors, references, and owner +// session IDs never reach the fallback; they go to the debug log only. +func resolveGeneratedImage(ctx context.Context, application *app.App, req generatedMediaRequest) types.AssistantMedia { + fallback := fmt.Sprintf("Generated image %q is unavailable.", req.name) + resolved, err := application.ResolveGeneratedFile(ctx, req.ref) + if err != nil { + slog.DebugContext(ctx, "Generated file could not be resolved for display", "name", req.name, "error", err) + return types.AssistantMedia{ID: req.id, Fallback: fallback} + } + if resolved.Path != "" && displaySafePath(resolved.Path) { + fallback = fmt.Sprintf("Generated image %q saved to: %s", req.name, resolved.Path) + } + inline, decoded := tuiimage.FromBytes(req.name, req.mimeType, resolved.Data) + if !decoded { + slog.DebugContext(ctx, "Generated file could not be decoded for display", "name", req.name) + return types.AssistantMedia{ID: req.id, Fallback: fallback} + } + return types.AssistantMedia{ID: req.id, Image: &inline, Fallback: fallback} +} + +// displaySafePath reports whether path can be shown verbatim in the chat: +// the canonical path comes from our own validated resolver, but a filename +// containing control characters must still never reach the terminal. +func displaySafePath(path string) bool { + return !strings.ContainsFunc(path, func(r rune) bool { + return r < 0x20 || r == 0x7f + }) +} diff --git a/pkg/tui/page/chat/generated_media_test.go b/pkg/tui/page/chat/generated_media_test.go new file mode 100644 index 0000000000..ce10a376cc --- /dev/null +++ b/pkg/tui/page/chat/generated_media_test.go @@ -0,0 +1,615 @@ +package chat + +import ( + "bytes" + "context" + "image" + "image/color" + "image/png" + "sync" + "testing" + + tea "charm.land/bubbletea/v2" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/app" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/runtime" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/tui/animation" + "github.com/docker/docker-agent/pkg/tui/components/messages" + "github.com/docker/docker-agent/pkg/tui/service" + "github.com/docker/docker-agent/pkg/tui/types" +) + +// resolverResult is one canned outcome of the fake resolver runtime, keyed +// by the reference's path. +type resolverResult struct { + data []byte + path string + err error +} + +// resolverTestRuntime is queueTestRuntime plus the local-runtime-only +// generated-file resolver capability (see pkg/app's generatedFileResolver), +// with canned per-path outcomes and a record of every resolved ref. The +// real manifest/workspace security behind the capability is covered in +// pkg/runtime; these tests pin the TUI contract around it. +type resolverTestRuntime struct { + queueTestRuntime + + mu sync.Mutex + results map[string]resolverResult + refs []runtime.GeneratedFileRef +} + +func (r *resolverTestRuntime) ResolveGeneratedFile(_ context.Context, ref runtime.GeneratedFileRef) (*runtime.ResolvedGeneratedFile, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.refs = append(r.refs, ref) + res, ok := r.results[ref.Path] + if !ok || res.err != nil { + return nil, runtime.ErrGeneratedFileUnavailable + } + return &runtime.ResolvedGeneratedFile{Data: res.data, Path: res.path}, nil +} + +func (r *resolverTestRuntime) resolveCalls() int { + r.mu.Lock() + defer r.mu.Unlock() + return len(r.refs) +} + +// mediaRecordingMessages wraps the real [messages.Model], recording +// AppendAssistantMedia and UpdateAssistantMedia calls while forwarding them +// so the real list state mutates. Mirrors recordingMessages in +// image_output_guard_integration_test.go. +type mediaRecordingMessages struct { + messages.Model + + mediaAgents []string + mediaCalls [][]types.AssistantMedia + mediaUpdates [][]types.AssistantMedia +} + +func (r *mediaRecordingMessages) AppendAssistantMedia(agentName string, media []types.AssistantMedia) tea.Cmd { + r.mediaAgents = append(r.mediaAgents, agentName) + r.mediaCalls = append(r.mediaCalls, media) + return r.Model.AppendAssistantMedia(agentName, media) +} + +func (r *mediaRecordingMessages) UpdateAssistantMedia(media []types.AssistantMedia) tea.Cmd { + r.mediaUpdates = append(r.mediaUpdates, media) + return r.Model.UpdateAssistantMedia(media) +} + +func newGeneratedMediaTestPage(t *testing.T, rt runtime.Runtime) (*chatPage, *mediaRecordingMessages) { + t.Helper() + return newGeneratedMediaTestPageWithSession(t, rt, session.New()) +} + +func testPNGBytes(t *testing.T) []byte { + t.Helper() + img := image.NewRGBA(image.Rect(0, 0, 2, 1)) + img.Set(0, 0, color.RGBA{B: 255, A: 255}) + var data bytes.Buffer + require.NoError(t, png.Encode(&data, img)) + return data.Bytes() +} + +func newGeneratedMediaTestPageWithSession(t *testing.T, rt runtime.Runtime, sess *session.Session) (*chatPage, *mediaRecordingMessages) { + t.Helper() + p := New(animation.NewRuntime(), t.Context(), app.New(t.Context(), rt, sess), service.NewSessionState(sess)).(*chatPage) + rec := &mediaRecordingMessages{Model: p.messages} + p.messages = rec + return p, rec +} + +// workspaceImagePart builds the document part a MessageAddedEvent carries +// for a workspace-materialized generated image. +func workspaceImagePart(name, relPath, owner string) chat.MessagePart { + return chat.MessagePart{ + Type: chat.MessagePartTypeDocument, + Document: &chat.Document{ + Name: name, + MimeType: "image/png", + Source: chat.DocumentSource{ + ArtifactPath: relPath, + ArtifactRoot: chat.ArtifactRootWorkspace, + ArtifactOwnerSessionID: owner, + }, + }, + } +} + +// assistantMessageAdded builds the event the run loop emits after +// persisting the "root" agent's assistant message. +func assistantMessageAdded(sessionID string, parts ...chat.MessagePart) *runtime.MessageAddedEvent { + msg := &session.Message{ + AgentName: "root", + Message: chat.Message{ + Role: chat.MessageRoleAssistant, + MultiContent: parts, + }, + } + return runtime.MessageAdded(sessionID, msg, "root").(*runtime.MessageAddedEvent) +} + +// resolveArmedMedia runs the asynchronous resolution command the page armed +// (recorded like a routed timer, so it survives background-tab dispatch) +// and returns the resolved-media message it produced. +func resolveArmedMedia(t *testing.T, p *chatPage) generatedMediaResolvedMsg { + t.Helper() + cmd := p.TakeRoutedTimers() + require.NotNil(t, cmd, "an async resolution command must be armed") + for _, msg := range runTimerCmd(t, cmd) { + if resolved, ok := msg.(generatedMediaResolvedMsg); ok { + return resolved + } + } + t.Fatal("the armed commands produced no generatedMediaResolvedMsg") + return generatedMediaResolvedMsg{} +} + +func TestMessageAdded_TextAndWorkspaceImageJoinSameTurn(t *testing.T) { + t.Parallel() + + const owner = "sess-media" + rt := &resolverTestRuntime{results: map[string]resolverResult{ + "cat.png": {data: testPNGBytes(t), path: "/workspace/cat.png"}, + }} + p, rec := newGeneratedMediaTestPage(t, rt) + + rec.AddUserMessage("draw a cat") + handled, _ := p.handleRuntimeEvent(runtime.StreamStarted(owner, "root")) + require.True(t, handled) + handled, _ = p.handleRuntimeEvent(runtime.AgentChoice("root", owner, "Here is your cat:")) + require.True(t, handled) + require.Equal(t, 1, rec.MessageTypeCount(types.MessageTypeAssistant)) + + handled, _ = p.handleRuntimeEvent(assistantMessageAdded(owner, + chat.MessagePart{Type: chat.MessagePartTypeText, Text: "Here is your cat:"}, + workspaceImagePart("cat.png", "cat.png", owner), + )) + require.True(t, handled, "MessageAddedEvent must be a recognized runtime event") + + // The placeholder is attached synchronously; resolution must not have + // happened inside the update path. + require.Equal(t, []string{"root"}, rec.mediaAgents) + require.Len(t, rec.mediaCalls, 1) + require.Len(t, rec.mediaCalls[0], 1) + placeholder := rec.mediaCalls[0][0] + assert.Nil(t, placeholder.Image, "nothing is decoded before the async resolution") + assert.Equal(t, `Generated image "cat.png" is unavailable.`, placeholder.Fallback) + assert.NotZero(t, placeholder.ID, "a resolvable item must carry a replacement ID") + assert.Zero(t, rt.resolveCalls(), "the resolver must never run synchronously inside Update") + + resolved := resolveArmedMedia(t, p) + _, _ = p.update(resolved) + + assert.Equal(t, 1, rt.resolveCalls()) + require.Len(t, rec.mediaUpdates, 1) + media := rec.mediaUpdates[0][0] + assert.Equal(t, placeholder.ID, media.ID) + require.NotNil(t, media.Image, "a resolvable file must be decoded for inline rendering") + assert.Equal(t, "cat.png", media.Image.Name) + assert.NotEmpty(t, media.Image.PNGData) + assert.Equal(t, `Generated image "cat.png" saved to: /workspace/cat.png`, media.Fallback, + "the fallback must surface the resolver-validated canonical workspace path") + + assert.Equal(t, 1, rec.MessageTypeCount(types.MessageTypeAssistant), + "the media must join the streamed-text assistant message, not open a new turn") + assert.Zero(t, rec.MessageTypeCount(types.MessageTypeSpinner)) + assert.True(t, p.hasReceivedAssistantContent) +} + +// TestMessageAdded_LabelsUseFinalPersistedDocumentName is the live-repro +// pin for prompt-directed naming ("Generate an image of a red panda coding +// at a terminal as assets/red-panda-terminal.jpg"): the same-turn inline +// label and fallback must carry the FINAL persisted Document name and the +// resolver's canonical workspace path — after marker/prompt naming, MIME +// extension correction, and collision suffixing — never a provisional or +// TUI-constructed one. The TUI gets both exclusively from the trusted +// runtime: the name from the MessageAddedEvent document, the path from +// ResolveGeneratedFile. +func TestMessageAdded_LabelsUseFinalPersistedDocumentName(t *testing.T) { + t.Parallel() + + const owner = "sess-final-name" + rt := &resolverTestRuntime{results: map[string]resolverResult{ + "assets/red-panda-terminal.png": {data: testPNGBytes(t), path: "/workspace/assets/red-panda-terminal.png"}, + "assets/red-panda-terminal-1.png": {data: testPNGBytes(t), path: "/workspace/assets/red-panda-terminal-1.png"}, + }} + p, rec := newGeneratedMediaTestPage(t, rt) + + handled, _ := p.handleRuntimeEvent(assistantMessageAdded(owner, + chat.MessagePart{Type: chat.MessagePartTypeText, Text: "Here is your red panda coding at a terminal:"}, + workspaceImagePart("red-panda-terminal.png", "assets/red-panda-terminal.png", owner), + workspaceImagePart("red-panda-terminal-1.png", "assets/red-panda-terminal-1.png", owner), + )) + require.True(t, handled) + + require.Len(t, rec.mediaCalls, 1) + require.Len(t, rec.mediaCalls[0], 2) + assert.Equal(t, `Generated image "red-panda-terminal.png" is unavailable.`, rec.mediaCalls[0][0].Fallback, + "even the pre-resolution placeholder must name the final persisted file") + assert.Equal(t, `Generated image "red-panda-terminal-1.png" is unavailable.`, rec.mediaCalls[0][1].Fallback, + "a collision-suffixed final name must be shown as persisted") + + _, _ = p.update(resolveArmedMedia(t, p)) + require.Len(t, rec.mediaUpdates, 1) + require.Len(t, rec.mediaUpdates[0], 2) + + resolved := rec.mediaUpdates[0][0] + require.NotNil(t, resolved.Image) + assert.Equal(t, "red-panda-terminal.png", resolved.Image.Name, + "the inline label must be the final persisted document name, not a provisional generated-N one") + assert.Equal(t, `Generated image "red-panda-terminal.png" saved to: /workspace/assets/red-panda-terminal.png`, resolved.Fallback, + "the fallback must carry the resolver-validated canonical workspace path") + + suffixed := rec.mediaUpdates[0][1] + require.NotNil(t, suffixed.Image) + assert.Equal(t, "red-panda-terminal-1.png", suffixed.Image.Name) + assert.Equal(t, `Generated image "red-panda-terminal-1.png" saved to: /workspace/assets/red-panda-terminal-1.png`, suffixed.Fallback) +} + +func TestMessageAdded_MediaOnlyTurnReplacesSpinner(t *testing.T) { + t.Parallel() + + const owner = "sess-media-only" + rt := &resolverTestRuntime{results: map[string]resolverResult{ + "cat.png": {data: testPNGBytes(t), path: "/workspace/cat.png"}, + }} + p, rec := newGeneratedMediaTestPage(t, rt) + + rec.AddUserMessage("draw a cat") + _, _ = p.handleRuntimeEvent(runtime.StreamStarted(owner, "root")) + require.Equal(t, 1, rec.MessageTypeCount(types.MessageTypeSpinner), + "a real pending spinner must exist before the media arrives") + + handled, _ := p.handleRuntimeEvent(assistantMessageAdded(owner, workspaceImagePart("cat.png", "cat.png", owner))) + require.True(t, handled) + + assert.Zero(t, rec.MessageTypeCount(types.MessageTypeSpinner), + "a media-only turn must replace the pending spinner immediately, before resolution") + require.Equal(t, 1, rec.MessageTypeCount(types.MessageTypeAssistant), + "a media-only turn must add a visible assistant message") + require.Len(t, rec.mediaCalls, 1) + assert.True(t, p.hasReceivedAssistantContent, "media-only output counts as assistant content") + + _, _ = p.update(resolveArmedMedia(t, p)) + require.Len(t, rec.mediaUpdates, 1) + require.NotNil(t, rec.mediaUpdates[0][0].Image) +} + +func TestMessageAdded_PreservesMediaOrder(t *testing.T) { + t.Parallel() + + const owner = "sess-order" + rt := &resolverTestRuntime{results: map[string]resolverResult{ + "first.png": {data: testPNGBytes(t), path: "/workspace/first.png"}, + "second.png": {data: testPNGBytes(t), path: "/workspace/second.png"}, + }} + p, rec := newGeneratedMediaTestPage(t, rt) + + handled, _ := p.handleRuntimeEvent(assistantMessageAdded(owner, + chat.MessagePart{Type: chat.MessagePartTypeText, Text: "two images"}, + workspaceImagePart("first.png", "first.png", owner), + workspaceImagePart("second.png", "second.png", owner), + )) + require.True(t, handled) + + require.Len(t, rec.mediaCalls, 1) + require.Len(t, rec.mediaCalls[0], 2) + + _, _ = p.update(resolveArmedMedia(t, p)) + require.Len(t, rec.mediaUpdates, 1) + require.Len(t, rec.mediaUpdates[0], 2) + assert.Equal(t, "first.png", rec.mediaUpdates[0][0].Image.Name) + assert.Equal(t, "second.png", rec.mediaUpdates[0][1].Image.Name) + assert.Equal(t, rec.mediaCalls[0][0].ID, rec.mediaUpdates[0][0].ID, + "resolved items must target their placeholders in order") + assert.Equal(t, rec.mediaCalls[0][1].ID, rec.mediaUpdates[0][1].ID) +} + +// TestMessageAdded_NilMessageIsNoOp pins the remote-runtime shape: the +// Message payload is process-local (json:"-"), so a decoded remote event +// carries only IDs. It must be a defined no-op — no panic, no resolution +// attempt, no list mutation. +func TestMessageAdded_NilMessageIsNoOp(t *testing.T) { + t.Parallel() + + p, rec := newGeneratedMediaTestPage(t, &resolverTestRuntime{}) + _, _ = p.handleRuntimeEvent(runtime.StreamStarted("sess-remote", "root")) + spinners := rec.MessageTypeCount(types.MessageTypeSpinner) + + var handled bool + require.NotPanics(t, func() { + handled, _ = p.handleRuntimeEvent(runtime.MessageAdded("sess-remote", nil, "root")) + }) + assert.True(t, handled) + assert.Empty(t, rec.mediaCalls, "a payload-less event must not produce media") + assert.Equal(t, spinners, rec.MessageTypeCount(types.MessageTypeSpinner), "the pending spinner must be untouched") + assert.False(t, p.hasReceivedAssistantContent) +} + +// TestMessageAdded_NoResolverCapabilityIsNoOp covers runtimes that cannot +// resolve generated files (e.g. remote runtimes): even a fully +// workspace-backed payload must render nothing rather than guess. +func TestMessageAdded_NoResolverCapabilityIsNoOp(t *testing.T) { + t.Parallel() + + p, rec := newGeneratedMediaTestPage(t, queueTestRuntime{}) + + handled, cmd := p.handleRuntimeEvent(assistantMessageAdded("sess-no-cap", + workspaceImagePart("cat.png", "cat.png", "sess-no-cap"))) + assert.True(t, handled) + assert.Nil(t, cmd) + assert.Empty(t, rec.mediaCalls) + assert.Nil(t, p.TakeRoutedTimers(), "no resolution may be armed without the capability") +} + +func TestMessageAdded_UnresolvableFallsBackToFilenameOnly(t *testing.T) { + t.Parallel() + + const owner = "sess-missing" + rt := &resolverTestRuntime{} // every resolution fails + p, rec := newGeneratedMediaTestPage(t, rt) + + handled, _ := p.handleRuntimeEvent(assistantMessageAdded(owner, + workspaceImagePart("cat.png", "missing-file.png", owner))) + require.True(t, handled) + + _, _ = p.update(resolveArmedMedia(t, p)) + + require.Len(t, rec.mediaUpdates, 1) + media := rec.mediaUpdates[0][0] + assert.Nil(t, media.Image, "an unresolvable file has nothing to render") + assert.Equal(t, `Generated image "cat.png" is unavailable.`, media.Fallback, + "an unresolved file falls back to the display filename only — no guessed path") + assert.NotContains(t, media.Fallback, owner, "the owner session ID must never leak into the fallback") + assert.NotContains(t, media.Fallback, "missing-file.png", "the raw reference must never leak into the fallback") + assert.NotContains(t, media.Fallback, "workspace", "no root kind or path may be shown for an unresolved file") + assert.NotContains(t, media.Fallback, "unavailable:", "raw resolver errors must never leak into the fallback") +} + +func TestMessageAdded_UndecodableFallsBackToCanonicalPath(t *testing.T) { + t.Parallel() + + const owner = "sess-corrupt" + rt := &resolverTestRuntime{results: map[string]resolverResult{ + "cat.png": {data: []byte("not really a png"), path: "/workspace/cat.png"}, + }} + p, rec := newGeneratedMediaTestPage(t, rt) + + handled, _ := p.handleRuntimeEvent(assistantMessageAdded(owner, workspaceImagePart("cat.png", "cat.png", owner))) + require.True(t, handled) + _, _ = p.update(resolveArmedMedia(t, p)) + + require.Len(t, rec.mediaUpdates, 1) + media := rec.mediaUpdates[0][0] + assert.Nil(t, media.Image, "undecodable bytes must not be handed to the renderer") + assert.Equal(t, `Generated image "cat.png" saved to: /workspace/cat.png`, media.Fallback, + "a resolved-but-undecodable file must surface its validated canonical path") +} + +// TestMessageAdded_ControlCharPathStaysUnavailable: a canonical path that +// cannot be shown verbatim (control characters) must degrade to the +// unavailable wording, never reach the terminal. +func TestMessageAdded_ControlCharPathStaysUnavailable(t *testing.T) { + t.Parallel() + + const owner = "sess-hostile-path" + rt := &resolverTestRuntime{results: map[string]resolverResult{ + "cat.png": {data: []byte("not a png"), path: "/workspace/\x1b[31mcat.png"}, + }} + p, rec := newGeneratedMediaTestPage(t, rt) + + _, _ = p.handleRuntimeEvent(assistantMessageAdded(owner, workspaceImagePart("cat.png", "cat.png", owner))) + _, _ = p.update(resolveArmedMedia(t, p)) + + require.Len(t, rec.mediaUpdates, 1) + fallback := rec.mediaUpdates[0][0].Fallback + assert.Equal(t, `Generated image "cat.png" is unavailable.`, fallback) + assert.NotContains(t, fallback, "\x1b") +} + +func TestMessageAdded_SanitizesHostileDisplayName(t *testing.T) { + t.Parallel() + + const owner = "sess-hostile" + rt := &resolverTestRuntime{} // resolution fails; only the name reaches the fallback + p, rec := newGeneratedMediaTestPage(t, rt) + + _, _ = p.handleRuntimeEvent(assistantMessageAdded(owner, + workspaceImagePart("../evil/\x1b[31mname.png", "cat.png", owner))) + _, _ = p.update(resolveArmedMedia(t, p)) + + require.Len(t, rec.mediaUpdates, 1) + fallback := rec.mediaUpdates[0][0].Fallback + assert.NotContains(t, fallback, "..", "traversal-like sequences must be sanitized out of the display name") + assert.NotContains(t, fallback, "", "angle brackets must be sanitized out of the display name") + assert.NotContains(t, fallback, "\x1b", "control characters must never reach the terminal") + assert.Contains(t, fallback, "name.png") +} + +func TestMessageAdded_IgnoresNonGeneratedAndNonImageParts(t *testing.T) { + t.Parallel() + + const owner = "sess-skip" + rt := &resolverTestRuntime{} + p, rec := newGeneratedMediaTestPage(t, rt) + + pdf := workspaceImagePart("doc.pdf", "doc.pdf", owner) + pdf.Document.MimeType = "application/pdf" + + handled, cmd := p.handleRuntimeEvent(assistantMessageAdded(owner, + chat.MessagePart{Type: chat.MessagePartTypeText, Text: "just text"}, + // User-attached image: inline bytes, no generated-file reference. + chat.MessagePart{Type: chat.MessagePartTypeDocument, Document: &chat.Document{ + Name: "attached.png", + MimeType: "image/png", + Source: chat.DocumentSource{InlineData: testPNGBytes(t)}, + }}, + // Ownerless reference: never resolved against a guessed session. + chat.MessagePart{Type: chat.MessagePartTypeDocument, Document: &chat.Document{ + Name: "legacy.png", + MimeType: "image/png", + Source: chat.DocumentSource{ArtifactPath: "legacy.png"}, + }}, + pdf, + )) + assert.True(t, handled) + assert.Nil(t, cmd) + assert.Empty(t, rec.mediaCalls, "no part above is a generated workspace-backed image") + assert.Zero(t, rt.resolveCalls()) +} + +// TestMessageAdded_UnknownRootRefStaysUnavailable: references whose root +// kind is unknown (empty ArtifactRoot, owner present) show the sanitized +// unavailable fallback without ever hitting the resolver. +func TestMessageAdded_UnknownRootRefStaysUnavailable(t *testing.T) { + t.Parallel() + + const owner = "sess-legacy" + rt := &resolverTestRuntime{} + p, rec := newGeneratedMediaTestPage(t, rt) + + part := workspaceImagePart("cat.png", "cat.png", owner) + part.Document.Source.ArtifactRoot = "" + handled, _ := p.handleRuntimeEvent(assistantMessageAdded(owner, part)) + require.True(t, handled) + + require.Len(t, rec.mediaCalls, 1) + media := rec.mediaCalls[0][0] + assert.Zero(t, media.ID, "an unknown-root item is final: nothing will replace it") + assert.Nil(t, media.Image) + assert.Equal(t, `Generated image "cat.png" is unavailable.`, media.Fallback) + assert.Nil(t, p.TakeRoutedTimers(), "no resolution may be armed for an unknown-root reference") + assert.Zero(t, rt.resolveCalls()) +} + +func TestMessageAdded_NonAssistantRoleIsNoOp(t *testing.T) { + t.Parallel() + + const owner = "sess-role" + p, rec := newGeneratedMediaTestPage(t, &resolverTestRuntime{}) + + msg := &session.Message{AgentName: "root", Message: chat.Message{ + Role: chat.MessageRoleTool, + MultiContent: []chat.MessagePart{workspaceImagePart("cat.png", "cat.png", owner)}, + }} + handled, cmd := p.handleRuntimeEvent(runtime.MessageAdded(owner, msg, "root")) + assert.True(t, handled) + assert.Nil(t, cmd) + assert.Empty(t, rec.mediaCalls) +} + +func TestMessageAdded_CancelledStreamIsNoOp(t *testing.T) { + t.Parallel() + + const owner = "sess-cancelled" + p, rec := newGeneratedMediaTestPage(t, &resolverTestRuntime{}) + p.streamCancelled = true + + handled, cmd := p.handleRuntimeEvent(assistantMessageAdded(owner, workspaceImagePart("cat.png", "cat.png", owner))) + assert.True(t, handled) + assert.Nil(t, cmd) + assert.Empty(t, rec.mediaCalls, "no media may be appended after the user cancelled the stream") +} + +// restoredMediaSession builds a persisted session whose SECOND assistant +// message carries generated media, so targeting the right historical +// message (not the newest) is exercised. +func restoredMediaSession(owner string) *session.Session { + sess := session.New() + sess.ID = owner + sess.Messages = []session.Item{ + session.NewMessageItem(&session.Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "draw a cat"}}), + session.NewMessageItem(&session.Message{AgentName: "root", Message: chat.Message{ + Role: chat.MessageRoleAssistant, Content: "Working on it.", + }}), + session.NewMessageItem(&session.Message{AgentName: "root", Message: chat.Message{ + Role: chat.MessageRoleAssistant, + Content: "Here is your cat:", + MultiContent: []chat.MessagePart{ + {Type: chat.MessagePartTypeText, Text: "Here is your cat:"}, + workspaceImagePart("cat.png", "cat.png", owner), + }, + }}), + session.NewMessageItem(&session.Message{AgentName: "root", Message: chat.Message{ + Role: chat.MessageRoleAssistant, Content: "Anything else?", + }}), + } + return sess +} + +// TestInit_RestoredSessionResolvesGeneratedMedia: a restored session's +// generated media is attached at load (sanitized placeholder) and resolved +// asynchronously, exactly like the live path. +func TestInit_RestoredSessionResolvesGeneratedMedia(t *testing.T) { + t.Parallel() + + const owner = "sess-restored" + rt := &resolverTestRuntime{results: map[string]resolverResult{ + "cat.png": {data: testPNGBytes(t), path: "/workspace/cat.png"}, + }} + p, rec := newGeneratedMediaTestPageWithSession(t, rt, restoredMediaSession(owner)) + + _ = p.Init() + + require.Equal(t, 3, rec.MessageTypeCount(types.MessageTypeAssistant)) + assert.Zero(t, rt.resolveCalls(), "restoring a session must not resolve synchronously") + + _, _ = p.update(resolveArmedMedia(t, p)) + + assert.Equal(t, 1, rt.resolveCalls()) + assert.Equal(t, []runtime.GeneratedFileRef{{ + OwnerSessionID: owner, + Root: chat.ArtifactRootWorkspace, + Path: "cat.png", + }}, rt.refs, "the persisted owner reference must be resolved as-is") + require.Len(t, rec.mediaUpdates, 1) + media := rec.mediaUpdates[0][0] + require.NotNil(t, media.Image) + assert.Equal(t, `Generated image "cat.png" saved to: /workspace/cat.png`, media.Fallback) +} + +// TestCollectRestoredGeneratedMedia_TargetsOwningMessage pins the position +// mapping LoadFromSession consumes: media lands on the exact session index +// of the assistant message that carries the reference. +func TestCollectRestoredGeneratedMedia_TargetsOwningMessage(t *testing.T) { + t.Parallel() + + const owner = "sess-positions" + sess := restoredMediaSession(owner) + p, _ := newGeneratedMediaTestPageWithSession(t, &resolverTestRuntime{}, sess) + + restored, requests := p.collectRestoredGeneratedMedia(sess) + + require.Len(t, restored, 1) + require.Len(t, restored[2], 1, "the media must be keyed to the carrying message's session position") + assert.Equal(t, `Generated image "cat.png" is unavailable.`, restored[2][0].Fallback) + require.Len(t, requests, 1) + assert.Equal(t, restored[2][0].ID, requests[0].id) +} + +// TestCollectRestoredGeneratedMedia_NoCapability: without the resolver +// capability (remote runtimes) a restored session renders no media at all — +// the pre-resolver behavior. +func TestCollectRestoredGeneratedMedia_NoCapability(t *testing.T) { + t.Parallel() + + const owner = "sess-remote-restore" + sess := restoredMediaSession(owner) + p, rec := newGeneratedMediaTestPageWithSession(t, queueTestRuntime{}, sess) + + restored, requests := p.collectRestoredGeneratedMedia(sess) + assert.Nil(t, restored) + assert.Nil(t, requests) + + _ = p.Init() + require.Equal(t, 3, rec.MessageTypeCount(types.MessageTypeAssistant)) + assert.Nil(t, p.TakeRoutedTimers()) +} diff --git a/pkg/tui/page/chat/runtime_events.go b/pkg/tui/page/chat/runtime_events.go index 1b07ea65f0..d0d9f54c68 100644 --- a/pkg/tui/page/chat/runtime_events.go +++ b/pkg/tui/page/chat/runtime_events.go @@ -32,6 +32,7 @@ import ( // - AgentChoiceEvent → Append text to message // - AgentChoiceReasoningEvent → Append reasoning block // - UserMessageEvent → Replace loading with user message +// - MessageAddedEvent → Render generated media (local runs only) // // Tool Events: // - PartialToolCallEvent → Show tool call in progress @@ -92,6 +93,9 @@ func (p *chatPage) handleRuntimeEvent(msg tea.Msg) (bool, tea.Cmd) { case *runtime.AgentChoiceReasoningEvent: return true, p.handleAgentChoiceReasoning(msg) + case *runtime.MessageAddedEvent: + return true, p.handleMessageAdded(msg) + case *runtime.ShellOutputEvent: return true, p.messages.AddShellOutputMessage(msg.Output) diff --git a/pkg/tui/types/types.go b/pkg/tui/types/types.go index 68dbc6ebec..4e25403662 100644 --- a/pkg/tui/types/types.go +++ b/pkg/tui/types/types.go @@ -62,6 +62,21 @@ const ( ToolStatusError ) +// AssistantMedia is one generated-media item attached to an assistant +// message (an artifact-backed image produced by the model itself, resolved +// by the chat page). Image is non-nil only when the artifact bytes were +// resolved and decoded for terminal rendering; Fallback always carries the +// safe textual description shown when inline rendering is unavailable +// (graphics disabled, undecodable bytes, or unresolvable artifact). +type AssistantMedia struct { + // ID links an item awaiting asynchronous resolution to the result that + // replaces it (see messages.Model.UpdateAssistantMedia). Zero means + // static: the item is final and never replaced. + ID uint64 + Image *tuiimage.Inline + Fallback string +} + // Message represents a single message in the chat type Message struct { Type MessageType @@ -72,6 +87,9 @@ type Message struct { ToolStatus ToolStatus // Status for tool calls ToolResult *tools.ToolCallResult // Result of tool call (when completed) Images []tuiimage.Inline // Prepared terminal images from the result + // AssistantMedia holds generated media rendered as part of an assistant + // turn, after the message's text content. + AssistantMedia []AssistantMedia // StartedAt records when a tool call entered ToolStatusRunning. // Used to display elapsed time for long-running tool calls. StartedAt *time.Time