Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions pkg/app/app.go
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down
57 changes: 57 additions & 0 deletions pkg/app/generated_file_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
255 changes: 255 additions & 0 deletions pkg/runtime/generated_file.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,255 @@
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")

// 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 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
}

// 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)
}

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
}
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
}

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)
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
}
Loading
Loading