From 35f3b45845bfcb1edadea958568a705e4933bdba 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 1/2] feat: render generated images through workspace manifest authorization Resolve authorized workspace files asynchronously outside the TUI update loop using the owning session, never viewing CWD. Reject unknown or historical external roots before any lookup. Require manifest identity, containment, regular-file and opened-file identity checks. Revalidate manifest authorization on every resolution so deletion revokes references; cache only workspace provenance. Route replacements by message ID, and preserve text and unavailable placeholders for forged or missing references. --- pkg/app/app.go | 30 + pkg/app/generated_file_test.go | 57 ++ pkg/runtime/generated_file.go | 239 +++++++ pkg/runtime/generated_file_test.go | 385 +++++++++++ pkg/runtime/loop.go | 17 +- pkg/runtime/runtime.go | 16 +- pkg/tui/components/message/message.go | 48 +- pkg/tui/components/message/message_test.go | 75 +++ .../messages/generated_media_test.go | 178 +++++ pkg/tui/components/messages/messages.go | 92 ++- pkg/tui/components/messages/messages_test.go | 99 ++- pkg/tui/page/chat/chat.go | 27 +- pkg/tui/page/chat/generated_media.go | 216 ++++++ pkg/tui/page/chat/generated_media_test.go | 636 ++++++++++++++++++ pkg/tui/page/chat/runtime_events.go | 4 + pkg/tui/types/types.go | 18 + 16 files changed, 2097 insertions(+), 40 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 0b891349b0..3f5b371972 100644 --- a/pkg/app/app.go +++ b/pkg/app/app.go @@ -1267,6 +1267,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..66217013e8 --- /dev/null +++ b/pkg/runtime/generated_file.go @@ -0,0 +1,239 @@ +package runtime + +import ( + "context" + "errors" + "fmt" + "io" + "io/fs" + "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, or missing 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 resolves; all other kinds are unavailable. + Root chat.ArtifactRootKind + // Path is the recorded workspace-relative slash-separated final path. + 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 keeps resolved workspace provenance per owner session. +// Manifest authorization is deliberately not cached: every resolution must +// observe current store state so deleting a session immediately revokes its +// generated-file references. +type generatedFileCache struct { + mu sync.Mutex + roots map[string]string // owner session ID → workspace root +} + +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 +} + +// 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 { + return nil, fmt.Errorf("%w: unresolvable root kind %q", ErrGeneratedFileUnavailable, ref.Root) + } + + if !fs.ValidPath(ref.Path) || strings.ContainsAny(ref.Path, "\\:\x00\r\n") || ref.Path == "." || strings.HasPrefix(ref.Path, "~") { + return nil, fmt.Errorf("%w: invalid workspace path", ErrGeneratedFileUnavailable) + } + + 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) + } + + 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 current manifest record for ref. +func (r *LocalRuntime) lookupGeneratedFile(ctx context.Context, ref GeneratedFileRef) (session.GeneratedFile, error) { + 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) + } + 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 +} + +// 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 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()) + } + 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(f) + if err != nil { + return nil, fmt.Errorf("reading recorded file: %w", err) + } + 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..153e683bf8 --- /dev/null +++ b/pkg/runtime/generated_file_test.go @@ -0,0 +1,385 @@ +package runtime + +import ( + "context" + "os" + "path/filepath" + "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 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) + } +} + +func TestResolveGeneratedFile_InvalidWorkspacePathRefusedBeforeLookup(t *testing.T) { + t.Parallel() + for _, target := range []string{"/outside/cat.png", "../cat.png", "~/cat.png", `C:\outside\cat.png`, `C:cat.png`, `\\server\share\cat.png`, "a/../cat.png", "a//cat.png", "cat\x00.png", "cat\n.png", ""} { + t.Run(target, func(t *testing.T) { + store := &countingStore{Store: session.NewInMemorySessionStore()} + r := &LocalRuntime{sessionStore: store, now: time.Now} + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef("owner", target)) + require.ErrorIs(t, err, ErrGeneratedFileUnavailable) + assert.Zero(t, store.lookups) + assert.Zero(t, store.getSessions) + }) + } +} + +func TestResolveGeneratedFile_ExternalRootRefusedBeforeLookup(t *testing.T) { + t.Parallel() + sess, _ := workspaceSession(t, "sess-external-ref") + inner := session.NewInMemorySessionStore() + require.NoError(t, inner.AddSession(t.Context(), sess)) + store := &countingStore{Store: inner} + r := &LocalRuntime{sessionStore: store, now: time.Now} + target := filepath.Join(t.TempDir(), "cat.png") + require.NoError(t, os.WriteFile(target, []byte("external-bytes"), 0o644)) + + _, err := r.ResolveGeneratedFile(t.Context(), GeneratedFileRef{ + OwnerSessionID: sess.ID, + Root: chat.ArtifactRootKind("external"), + Path: target, + }) + require.ErrorIs(t, err, ErrGeneratedFileUnavailable) + assert.Zero(t, store.lookups) + assert.Zero(t, store.getSessions) +} + +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_ForgedHistoricalExternalManifestRefused(t *testing.T) { + t.Parallel() + sess, _ := workspaceSession(t, "sess-historical-external") + inner := session.NewInMemorySessionStore() + require.NoError(t, inner.AddSession(t.Context(), sess)) + target := filepath.Join(t.TempDir(), "cat.png") + require.NoError(t, os.WriteFile(target, []byte("external-bytes"), 0o644)) + external := chat.ArtifactRootKind("external") + store := &fixedManifestStore{Store: inner, file: session.GeneratedFile{ + SessionID: sess.ID, RelPath: target, Root: external, MimeType: "image/png", CreatedAt: time.Now(), + }} + r := &LocalRuntime{sessionStore: store, now: time.Now} + + _, err := r.ResolveGeneratedFile(t.Context(), GeneratedFileRef{ + OwnerSessionID: sess.ID, Root: chat.ArtifactRootWorkspace, Path: target, + }) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable) +} + +type fixedManifestStore struct { + session.Store + + file session.GeneratedFile +} + +func (s *fixedManifestStore) LookupGeneratedFile(context.Context, string, string) (*session.GeneratedFile, error) { + return &s.file, nil +} + +func (s *fixedManifestStore) AddGeneratedFile(context.Context, session.GeneratedFile) error { + return nil +} + +func TestResolveGeneratedFile_LegacyManifestRoundTripsLargeWorkspaceFile(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-large-legacy") + r, store := resolverTestRuntime(t, sess) + recordWorkspaceFile(t, store, sess.ID, root, "cat.png", []byte("png")) + path := filepath.Join(root, "cat.png") + const size = (20 << 20) + 1 + require.NoError(t, os.Truncate(path, size)) + f, err := os.OpenFile(path, os.O_WRONLY, 0) + require.NoError(t, err) + _, err = f.WriteAt([]byte{0x7f}, size-1) + require.NoError(t, err) + require.NoError(t, f.Close()) + + resolved, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + require.NoError(t, err) + require.Len(t, resolved.Data, size) + assert.Equal(t, byte(0x7f), resolved.Data[size-1]) +} + +// countingStore wraps a session store to observe how often resolution hits +// the store, proving which resolution inputs are cached. +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_CachesRootButRevalidatesManifest(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, 3, store.lookups, "every resolution must revalidate manifest authorization") + assert.Equal(t, 1, store.getSessions, "repeat resolutions must reuse the cached workspace root") +} + +func TestResolveGeneratedFile_RecordSeedsOnlyRootCache(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.Equal(t, 1, store.lookups, "a freshly recorded file must still revalidate manifest authorization") + assert.Zero(t, store.getSessions, "a freshly recorded file must resolve without a session read") +} + +func TestResolveGeneratedFile_SessionDeletionRevokesCachedResolution(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-deleted") + 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("original")) + + resolved, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + require.NoError(t, err) + assert.Equal(t, []byte("original"), resolved.Data) + require.NoError(t, store.DeleteSession(t.Context(), sess.ID)) + require.NoError(t, os.WriteFile(filepath.Join(root, "cat.png"), []byte("replacement"), 0o644)) + + _, err = r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + require.ErrorIs(t, err, ErrGeneratedFileUnavailable) + assert.Equal(t, 2, store.lookups, "deletion must be observed before reading replaced workspace bytes") +} + +// 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 234fe40e29..8937c12519 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -1376,6 +1376,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)) @@ -1472,20 +1476,25 @@ func (r *LocalRuntime) sessionLookup() session.Lookup { return r.sessionStore.GetSession } -// recordGeneratedFile writes one manifest record after a successful -// write — materialization is the only writer of the manifest. +// recordGeneratedFile writes one manifest record after a successful write. +// Materialization is the only writer of the manifest; resolvers always read +// it back so authorization reflects current store state. 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 + } + return nil } func (r *LocalRuntime) recordGeneratedBlob(ctx context.Context, sessionID, finalPath string, data []byte) error { diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index d184998326..b899274f61 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -258,12 +258,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 d9f99a0170..1be31d2f04 100644 --- a/pkg/tui/components/message/message.go +++ b/pkg/tui/components/message/message.go @@ -359,7 +359,7 @@ func (mv *messageModel) IsToggleLine(lineIdx int) bool { func (mv *messageModel) RenderedSegments(width int) (AssistantSegments, bool) { msg := mv.message - if msg == nil || msg.Type != types.MessageTypeAssistant || msg.Content == "" || mv.selected || len(mv.markdownImages) != 0 { + if msg == nil || msg.Type != types.MessageTypeAssistant || msg.Content == "" || mv.selected || len(mv.markdownImages) != 0 || len(msg.AssistantMedia) != 0 { return AssistantSegments{}, false } messageStyle := styles.AssistantMessageStyle @@ -492,7 +492,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 } @@ -552,7 +552,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() } @@ -569,6 +569,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) { @@ -735,6 +736,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 02490455ab..81c3917fc7 100644 --- a/pkg/tui/components/message/message_test.go +++ b/pkg/tui/components/message/message_test.go @@ -550,3 +550,78 @@ func TestAssistantRenderedSegmentsRebuildHeaderOnWidthChange(t *testing.T) { require.Equal(t, linePlain(want), linePlain(got)) require.Equal(t, lineWidthsForMessage(want), lineWidthsForMessage(got)) } + +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..e4722c35ba --- /dev/null +++ b/pkg/tui/components/messages/generated_media_test.go @@ -0,0 +1,178 @@ +package messages + +import ( + "bytes" + "image" + "image/color" + "image/png" + "strings" + "testing" + + "github.com/charmbracelet/x/ansi" + "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" + tuiimage "github.com/docker/docker-agent/pkg/tui/image" + "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}, + }) +} + +func generatedMediaTestImage(t *testing.T, name string) tuiimage.Inline { + t.Helper() + img := image.NewRGBA(image.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 TestAppendAssistantMediaPreservesStreamedTextRendering(t *testing.T) { + m := newMediaTestModel(t) + m.AppendToLastMessage("root", "Here is your cat:") + _ = m.View() + + inline := generatedMediaTestImage(t, "cat.png") + unavailable := `Generated image "cat.png" is unavailable.` + fallback := `Generated image "cat.png" saved to: /tmp/artifacts/sess/cat.png` + m.AppendAssistantMedia("root", []types.AssistantMedia{{ID: 9, Fallback: unavailable}}) + + tuiimage.SetRenderingEnabled(true) + t.Cleanup(func() { tuiimage.SetRenderingEnabled(true) }) + view := m.View() + plain := ansi.Strip(view) + assert.Contains(t, plain, "Here is your cat:") + assert.Contains(t, plain, unavailable) + assert.NotContains(t, view, "cagent-image") + + m.UpdateAssistantMedia([]types.AssistantMedia{{ID: 9, Image: &inline, Fallback: fallback}}) + view = m.View() + plain = ansi.Strip(view) + assert.Contains(t, plain, "Here is your cat:") + assert.Contains(t, plain, "cat.png") + assert.Contains(t, view, "cagent-image") + assert.NotContains(t, plain, "unavailable") + assert.NotContains(t, plain, "saved to:") + textIndex := strings.Index(view, "Here is your cat:") + require.NotEqual(t, -1, textIndex) + assert.Less(t, textIndex, strings.Index(view, "cagent-image")) + + tuiimage.SetRenderingEnabled(false) + m.views[0] = m.createMessageView(m.messages[0]) + m.invalidateAllItems() + view = m.View() + plain = ansi.Strip(view) + assert.Contains(t, plain, "Here is your cat:") + assert.Contains(t, strings.ReplaceAll(plain, "\n", ""), fallback) + assert.NotContains(t, view, "cagent-image") +} + +// 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 ab700fd25c..754a836f1e 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 @@ -1322,7 +1335,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. @@ -1747,7 +1760,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) @@ -1857,9 +1870,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)) } @@ -2071,6 +2089,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 613466d82a..b1928cf126 100644 --- a/pkg/tui/components/messages/messages_test.go +++ b/pkg/tui/components/messages/messages_test.go @@ -120,7 +120,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) @@ -168,7 +168,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 @@ -210,7 +210,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) @@ -249,7 +249,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 @@ -317,7 +317,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) @@ -430,7 +430,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") @@ -496,7 +496,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) @@ -545,7 +545,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) @@ -593,7 +593,7 @@ func TestLoadFromSessionReasoningWithContentToolResultsStandalone(t *testing.T) }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) require.Len(t, m.messages, 3) @@ -659,7 +659,7 @@ func TestLoadFromSessionMultipleStandaloneToolCallsWithContentAndResults(t *test }, } - m.LoadFromSession(sess) + m.LoadFromSession(sess, nil) require.Len(t, m.messages, 3) @@ -1374,7 +1374,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) @@ -1425,7 +1425,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) @@ -1783,7 +1783,7 @@ func TestMessageCacheBoundsHistoricalRerender(t *testing.T) { } sess.Messages = append(sess.Messages, session.NewMessageItem(&session.Message{AgentName: "root", Message: chat.Message{Role: role, Content: body}})) } - _ = m.LoadFromSession(sess) + _ = m.LoadFromSession(sess, nil) _ = m.View() require.False(t, m.renderDirty) before := m.renderedItems.Len() @@ -1795,3 +1795,76 @@ func TestMessageCacheBoundsHistoricalRerender(t *testing.T) { require.False(t, m.renderDirty) require.Equal(t, before, m.renderedItems.Len(), "single append does not trigger history-wide cache growth") } + +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 fab0a08c0e..feec2d3ee3 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 } @@ -227,9 +228,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) + } } } @@ -638,6 +644,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..7896813107 --- /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 { + 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..a2d0b08280 --- /dev/null +++ b/pkg/tui/page/chat/generated_media_test.go @@ -0,0 +1,636 @@ +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_ExternalRootRefStaysUnavailable(t *testing.T) { + t.Parallel() + + const owner = "sess-external" + rt := &resolverTestRuntime{} + p, rec := newGeneratedMediaTestPage(t, rt) + + part := workspaceImagePart("cat.png", "/tmp/cat.png", owner) + part.Document.Source.ArtifactRoot = chat.ArtifactRootKind("external") + 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) + assert.Nil(t, media.Image) + assert.Equal(t, `Generated image "cat.png" is unavailable.`, media.Fallback) + assert.Nil(t, p.TakeRoutedTimers()) + 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 a6682714ca..a34a142042 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 From 38adfaba37bae502a998d460f19511720d18662c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Fri, 4 Sep 2026 20:09:06 +0200 Subject: [PATCH 2/2] feat: restore generated media from portable session blobs Authorize owning-session workspace references against the manifest before preferring portable blobs. Fall back to workspace files only when the optional blob interface is absent or reports blob-not-found; other errors fail closed. Portable bytes survive missing files and workspace provenance, with a recorded relative display label in that case. Generated-media reads have no application byte cap; ordinary attachment limits are unchanged. --- pkg/runtime/generated_file.go | 34 ++++++++++++++------ pkg/runtime/generated_file_test.go | 51 ++++++++++++++++++++++++++++++ 2 files changed, 76 insertions(+), 9 deletions(-) diff --git a/pkg/runtime/generated_file.go b/pkg/runtime/generated_file.go index 66217013e8..65854be444 100644 --- a/pkg/runtime/generated_file.go +++ b/pkg/runtime/generated_file.go @@ -23,12 +23,6 @@ import ( // 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). @@ -43,9 +37,9 @@ type GeneratedFileRef struct { 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). +// ResolvedGeneratedFile carries the resolved bytes and a display path. +// Portable blobs remain readable without workspace provenance; in that case +// Path is the recorded relative artifact path, not a canonical absolute path. type ResolvedGeneratedFile struct { Data []byte Path string @@ -106,6 +100,16 @@ func (r *LocalRuntime) ResolveGeneratedFile(ctx context.Context, ref GeneratedFi return nil, fmt.Errorf("%w: reference root %q does not match recorded root %q", ErrGeneratedFileUnavailable, ref.Root, record.Root) } + if blobs, ok := r.sessionStore.(session.GeneratedMediaBlobStore); ok { + data, err := blobs.LookupGeneratedBlob(ctx, ref.OwnerSessionID, ref.Path) + if err == nil { + return &ResolvedGeneratedFile{Data: data, Path: generatedFileDisplayPath(ctx, r, ref)}, nil + } + if !errors.Is(err, session.ErrGeneratedBlobNotFound) { + return nil, fmt.Errorf("%w: loading portable media: %w", ErrGeneratedFileUnavailable, err) + } + } + workspaceRoot, err := r.generatedFileWorkspaceRoot(ctx, ref.OwnerSessionID) if err != nil { return nil, err @@ -117,6 +121,18 @@ func (r *LocalRuntime) ResolveGeneratedFile(ctx context.Context, ref GeneratedFi return &ResolvedGeneratedFile{Data: data, Path: canonical}, nil } +func generatedFileDisplayPath(ctx context.Context, r *LocalRuntime, ref GeneratedFileRef) string { + root, err := r.generatedFileWorkspaceRoot(ctx, ref.OwnerSessionID) + if err != nil { + return ref.Path + } + canonicalRoot, err := filepath.EvalSymlinks(root) + if err != nil { + canonicalRoot = root + } + return filepath.Join(canonicalRoot, filepath.FromSlash(ref.Path)) +} + // lookupGeneratedFile returns the current manifest record for ref. func (r *LocalRuntime) lookupGeneratedFile(ctx context.Context, ref GeneratedFileRef) (session.GeneratedFile, error) { manifest, ok := r.sessionStore.(session.GeneratedMediaManifest) diff --git a/pkg/runtime/generated_file_test.go b/pkg/runtime/generated_file_test.go index 153e683bf8..d16763c6b8 100644 --- a/pkg/runtime/generated_file_test.go +++ b/pkg/runtime/generated_file_test.go @@ -213,9 +213,60 @@ func TestResolveGeneratedFile_ParentWorkingDirFallback(t *testing.T) { assert.Equal(t, []byte("png"), resolved.Data) } +func TestResolveGeneratedFile_BlobWithoutManifestIsRefused(t *testing.T) { + t.Parallel() + sess, _ := workspaceSession(t, "sess-forged-blob") + r, store := resolverTestRuntime(t, sess) + require.NoError(t, store.(session.GeneratedMediaBlobStore).AddGeneratedBlob(t.Context(), sess.ID, ".env", []byte("SECRET"))) + + _, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, ".env")) + assert.ErrorIs(t, err, ErrGeneratedFileUnavailable) +} + +func TestResolveGeneratedFile_PortableBlobWinsOverWorkspace(t *testing.T) { + t.Parallel() + sess, root := workspaceSession(t, "sess-portable") + + r, store := resolverTestRuntime(t, sess) + recordWorkspaceFile(t, store, sess.ID, root, "cat.png", []byte("workspace")) + require.NoError(t, store.(session.GeneratedMediaBlobStore).AddGeneratedBlob(t.Context(), sess.ID, "cat.png", []byte("database"))) + require.NoError(t, os.Remove(filepath.Join(root, "cat.png"))) + + resolved, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + require.NoError(t, err) + assert.Equal(t, []byte("database"), resolved.Data) +} + +func TestResolveGeneratedFile_PortableBlobWithoutWorkspaceRoot(t *testing.T) { + t.Parallel() + sess := &session.Session{ID: "sess-portable-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(), + })) + require.NoError(t, store.(session.GeneratedMediaBlobStore).AddGeneratedBlob(t.Context(), sess.ID, "cat.png", []byte("database"))) + + resolved, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + require.NoError(t, err) + assert.Equal(t, []byte("database"), resolved.Data) + assert.Equal(t, "cat.png", resolved.Path) +} + +func TestResolveGeneratedFile_LegacyManifestFallsBackToWorkspace(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("legacy")) + + resolved, err := r.ResolveGeneratedFile(t.Context(), workspaceRef(sess.ID, "cat.png")) + require.NoError(t, err) + assert.Equal(t, []byte("legacy"), 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(),