diff --git a/pkg/acp/agent.go b/pkg/acp/agent.go index 37bf8e400..354dec7e3 100644 --- a/pkg/acp/agent.go +++ b/pkg/acp/agent.go @@ -635,11 +635,11 @@ func (a *Agent) readResourceLink(ctx context.Context, sessionID string, rl *acp. } func resourceLinkName(rl *acp.ContentBlockResourceLink) string { - if rl.Name != "" { - return rl.Name + if name := chat.SanitizeDisplayName(rl.Name); name != "" { + return name } if path, ok := resourceLinkPath(rl.Uri); ok { - if base := filepath.Base(path); base != "." && base != string(filepath.Separator) { + if base := chat.SanitizeDisplayName(filepath.Base(path)); base != "" && base != "." && base != string(filepath.Separator) { return base } } diff --git a/pkg/acp/agent_test.go b/pkg/acp/agent_test.go index d174253cb..1ed626d8c 100644 --- a/pkg/acp/agent_test.go +++ b/pkg/acp/agent_test.go @@ -4,6 +4,7 @@ import ( "context" "io" "path/filepath" + "strings" "testing" acpsdk "github.com/coder/acp-go-sdk" @@ -70,6 +71,19 @@ func TestBuildUserContent_ResourceLinkFallbackDoesNotExposeAbsoluteURI(t *testin assert.NotContains(t, content, "/var/folders") } +func TestResourceLinkNameIsSafeAndBounded(t *testing.T) { + t.Parallel() + + longName := strings.Repeat("é", 100) + "\nforged" + assert.Equal(t, chat.SanitizeDisplayName(longName), resourceLinkName(&acpsdk.ContentBlockResourceLink{Name: longName})) + + got := resourceLinkName(&acpsdk.ContentBlockResourceLink{Uri: "file:///tmp/unsafe%0Aname.png"}) + assert.Equal(t, "unsafe_name.png", got) + assert.LessOrEqual(t, len(got), chat.MaxSanitizedFieldBytes) + + assert.Equal(t, "resource", resourceLinkName(&acpsdk.ContentBlockResourceLink{Uri: "https://example.com/private.png"})) +} + func TestBuildUserMessage_ImageContent(t *testing.T) { t.Parallel() diff --git a/pkg/acp/runagent_test.go b/pkg/acp/runagent_test.go index c36a32484..faba21a22 100644 --- a/pkg/acp/runagent_test.go +++ b/pkg/acp/runagent_test.go @@ -8,6 +8,8 @@ import ( "fmt" "io" "log/slog" + "os" + "path/filepath" "slices" "strings" "sync" @@ -19,11 +21,15 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + agentpkg "github.com/docker/docker-agent/pkg/agent" "github.com/docker/docker-agent/pkg/chat" "github.com/docker/docker-agent/pkg/effort" + "github.com/docker/docker-agent/pkg/model/provider/base" + "github.com/docker/docker-agent/pkg/modelsdev" "github.com/docker/docker-agent/pkg/runtime" "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/sessiontitle" + "github.com/docker/docker-agent/pkg/team" "github.com/docker/docker-agent/pkg/tools" skillstool "github.com/docker/docker-agent/pkg/tools/builtin/skills" "github.com/docker/docker-agent/pkg/tools/builtin/todo" @@ -182,6 +188,84 @@ func promptRequest(text string) acpsdk.PromptRequest { } } +type escapingMediaProvider struct { + requested string +} + +func (p *escapingMediaProvider) ID() modelsdev.ID { return modelsdev.ParseIDOrZero("test/media") } +func (p *escapingMediaProvider) BaseConfig() base.Config { return base.Config{} } +func (p *escapingMediaProvider) MaxTokens() int { return 0 } +func (p *escapingMediaProvider) CreateChatCompletionStream(context.Context, []chat.Message, []tools.Tool) (chat.MessageStream, error) { + return &escapingMediaStream{responses: []chat.MessageStreamResponse{ + {Choices: []chat.MessageStreamChoice{{Index: 0, Delta: chat.MessageDelta{ + Content: "completed", + Media: []chat.MediaDelta{{Data: []byte("png"), MimeType: "image/png", Name: "provider.png", RequestedPath: p.requested, Size: 3}}, + }}}}, + {Choices: []chat.MessageStreamChoice{{Index: 0, FinishReason: chat.FinishReasonStop}}, Usage: &chat.Usage{InputTokens: 1, OutputTokens: 1}}, + }}, nil +} + +type escapingMediaStream struct { + responses []chat.MessageStreamResponse + index int +} + +func (s *escapingMediaStream) Recv() (chat.MessageStreamResponse, error) { + if s.index == len(s.responses) { + return chat.MessageStreamResponse{}, io.EOF + } + response := s.responses[s.index] + s.index++ + return response, nil +} + +func (*escapingMediaStream) Close() {} + +func TestPrompt_EscapingGeneratedMediaCompletesWithoutElicitation(t *testing.T) { + workspace := t.TempDir() + external := filepath.Join(t.TempDir(), "cat.png") + store := session.NewInMemorySessionStore() + root := agentpkg.New("root", "You are a test agent", agentpkg.WithModel(&escapingMediaProvider{requested: external})) + rt, err := runtime.New(t.Context(), team.New(team.WithAgents(root)), + runtime.WithSessionCompaction(false), runtime.WithSessionStore(store)) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, rt.Close()) }) + rt.OnElicitationRequest(func(runtime.Event) { t.Error("generated media escape must not elicit") }) + agent, sess, peer := newPromptTestAgent(t, rt) + sess.sess.ID = testSessionID + sess.sess.WorkingDir = workspace + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + response, err := agent.Prompt(ctx, promptRequest("save the generated image")) + require.NoError(t, err) + assert.Equal(t, acpsdk.StopReasonEndTurn, response.StopReason) + assert.NoFileExists(t, external) + assert.Equal(t, []byte("png"), mustReadACPFile(t, filepath.Join(workspace, "cat.png"))) + stored, err := store.GetSession(t.Context(), testSessionID) + require.NoError(t, err) + require.Len(t, stored.GetAllMessages(), 2) + assistant := stored.GetAllMessages()[1].Message + assert.Equal(t, "completed", assistant.Content) + require.Len(t, assistant.MultiContent, 2) + document := assistant.MultiContent[1].Document + require.NotNil(t, document) + assert.Equal(t, "cat.png", document.Source.ArtifactPath) + assert.Equal(t, chat.ArtifactRootWorkspace, document.Source.ArtifactRoot) + assert.Equal(t, testSessionID, document.Source.ArtifactOwnerSessionID) + out, ok := peer.out.(*captureWriter) + require.True(t, ok) + assert.Contains(t, strings.Join(out.lines(), "\n"), "outside the workspace") + assert.Empty(t, peer.recordedRequests()) +} + +func mustReadACPFile(t *testing.T, path string) []byte { + t.Helper() + b, err := os.ReadFile(path) + require.NoError(t, err) + return b +} + func TestPromptReplacementCancelsQueuedTurnWithoutSideEffects(t *testing.T) { t.Parallel() diff --git a/pkg/chat/document.go b/pkg/chat/document.go index 9f822eec7..6b02aadaf 100644 --- a/pkg/chat/document.go +++ b/pkg/chat/document.go @@ -8,8 +8,7 @@ package chat // deprecated but remain supported for backward compatibility. const MessagePartTypeDocument MessagePartType = "document" -// ArtifactRootKind identifies which root a DocumentSource.ArtifactPath is -// relative to. +// ArtifactRootKind identifies which root a DocumentSource.ArtifactPath is relative to. type ArtifactRootKind string // ArtifactRootWorkspace means ArtifactPath is relative to the OWNING diff --git a/pkg/chat/media.go b/pkg/chat/media.go index f028f7b77..4373a3177 100644 --- a/pkg/chat/media.go +++ b/pkg/chat/media.go @@ -15,6 +15,16 @@ type MediaDelta struct { // Materialization chooses a fallback when no usable name is available. Name string `json:"name,omitempty"` + // RequestedPath is the prompt-directed target path the model asked for + // (e.g. echoed from an "as sunshine.jpg" instruction), when one exists. + // It is untrusted model input: the runtime routes it through + // workspacemedia.ClassifyRequestedPath, and a path escaping the workspace + // requires an explicit user confirmation before it is honored. Response + // marker extraction (the "[media-file: ...]" protocol) will populate it; + // until that lands, providers leave it empty and materialization falls + // back to Name. + RequestedPath string `json:"requested_path,omitempty"` + // Size is the byte length of Data, cached because Data itself is // dropped once the artifact is materialized. Size int64 `json:"size,omitempty"` diff --git a/pkg/cli/runner.go b/pkg/cli/runner.go index e730fb2df..ea15f9890 100644 --- a/pkg/cli/runner.go +++ b/pkg/cli/runner.go @@ -248,9 +248,10 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess case *runtime.ElicitationRequestEvent: serverURL, ok := e.Meta["docker-agent/server_url"].(string) if !ok || serverURL == "" { - slog.WarnContext(ctx, "Skipping elicitation: missing or invalid server_url (non-interactive session?)") + // Keep draining after declining forms so follow-up events cannot stall the turn. + slog.WarnContext(ctx, "Declining elicitation without form support in CLI mode", "message", e.Message) _ = rt.ResumeElicitation(ctx, "decline", nil, e.ElicitationID) - return nil + continue } result := out.PromptOAuthAuthorization(ctx, serverURL) diff --git a/pkg/cli/runner_test.go b/pkg/cli/runner_test.go index c393b1347..74d89500b 100644 --- a/pkg/cli/runner_test.go +++ b/pkg/cli/runner_test.go @@ -8,6 +8,7 @@ import ( "strings" "sync" "testing" + "time" "gotest.tools/v3/assert" @@ -33,6 +34,10 @@ func TestMain(m *testing.M) { // It emits pre-configured events from RunStream and records Resume calls. type mockRuntime struct { events []runtime.Event + // runStreamFn, when set, replaces the default pre-buffered RunStream — + // used to model a live runtime that only makes progress while the + // consumer keeps draining. + runStreamFn func(context.Context, *session.Session) <-chan runtime.Event mu sync.Mutex resumes []runtime.ResumeRequest @@ -145,7 +150,10 @@ func (m *mockRuntime) Resume(_ context.Context, req runtime.ResumeRequest) { m.resumes = append(m.resumes, req) } -func (m *mockRuntime) RunStream(_ context.Context, _ *session.Session) <-chan runtime.Event { +func (m *mockRuntime) RunStream(ctx context.Context, sess *session.Session) <-chan runtime.Event { + if m.runStreamFn != nil { + return m.runStreamFn(ctx, sess) + } ch := make(chan runtime.Event, len(m.events)) for _, e := range m.events { ch <- e @@ -596,3 +604,50 @@ func TestErrorEventReturnedNotPrinted(t *testing.T) { assert.Equal(t, errors.As(err, &runtimeErr), true) assert.Equal(t, strings.Contains(buf.String(), "model failed"), false) } + +// A non-OAuth MCP elicitation must be declined in CLI mode without abandoning the event +// stream: the runtime only makes progress while the consumer drains, so +// returning early would stall the follow-up events (redirect warning, +// assistant response) and lose the turn. The unbuffered stream below makes +// the test fail (bounded, not wedged) if Run stops consuming after the +// decline. +func TestNonOAuthElicitationDeclinedAndStreamDrained(t *testing.T) { + t.Parallel() + + drained := make(chan struct{}) + rt := &mockRuntime{ + runStreamFn: func(context.Context, *session.Session) <-chan runtime.Event { + ch := make(chan runtime.Event) // unbuffered: every send needs a live consumer + go func() { + defer close(ch) + defer close(drained) + ch <- &runtime.ElicitationRequestEvent{Type: "elicitation_request", Message: "Choose a deployment region"} + ch <- runtime.Warning("The deployment choice was declined", "test") + ch <- runtime.AgentChoice("test", "sess", "Continuing without deployment.") + }() + return ch + }, + } + + var buf bytes.Buffer + out := NewPrinter(&buf) + sess := session.New() + + err := Run(t.Context(), out, Config{}, rt, sess, []string{"hello"}) + assert.NilError(t, err) + + select { + case <-drained: + case <-time.After(10 * time.Second): + t.Fatal("the CLI stopped draining the stream after declining the elicitation") + } + + rt.mu.Lock() + defer rt.mu.Unlock() + assert.Equal(t, rt.elicitationDeclines, 1) + assert.Equal(t, rt.elicitationLastAction, tools.ElicitationAction("decline")) + assert.Check(t, strings.Contains(buf.String(), "deployment choice was declined"), + "the warning must be surfaced: %q", buf.String()) + assert.Check(t, strings.Contains(buf.String(), "Continuing without deployment."), + "the assistant response must still be printed: %q", buf.String()) +} diff --git a/pkg/runtime/elicitation.go b/pkg/runtime/elicitation.go index 85bb1ee0e..9970ef2df 100644 --- a/pkg/runtime/elicitation.go +++ b/pkg/runtime/elicitation.go @@ -455,21 +455,59 @@ func backgroundElicitationDeclinedNote(message string) string { ) } +// elicitationSpec carries an MCP request through the shared waiter registry. +type elicitationSpec struct { + message string + mode string + schema any + url string + // serverElicitationID is the originating MCP server's wire ID, if any. + // Informational only — never a routing key (#3584 review item 2a). + serverElicitationID string + meta map[string]any + // agentName and sessionID override the runtime-derived defaults (the + // shared current-agent slot and the ctx conversation ID) when the caller + // knows the owning agent/session more precisely. + agentName string + sessionID string +} + // elicitationHandler is the MCP-toolset-side hook that turns an inbound // elicitation request from a server into an ElicitationRequest event and // waits for the embedder's response, correlated by elicitation ID. func (r *LocalRuntime) elicitationHandler(ctx context.Context, req *mcp.ElicitParams) (tools.ElicitationResult, error) { slog.DebugContext(ctx, "Elicitation request received from MCP server", "message", req.Message) + return r.requestElicitation(ctx, elicitationSpec{ + message: req.Message, + mode: req.Mode, + schema: req.RequestedSchema, + url: req.URL, + serverElicitationID: req.ElicitationID, + meta: req.Meta, + }) +} +// requestElicitation emits spec as an ElicitationRequest event and waits for +// the embedder's response, correlated by elicitation ID. +func (r *LocalRuntime) requestElicitation(ctx context.Context, spec elicitationSpec) (tools.ElicitationResult, error) { // In non-interactive mode (e.g., MCP serve), there is no user to respond // to elicitation requests. Decline immediately instead of blocking forever. if r.nonInteractive { - slog.DebugContext(ctx, "Declining elicitation in non-interactive mode", "message", req.Message) + slog.DebugContext(ctx, "Declining elicitation in non-interactive mode", "message", spec.message) return tools.ElicitationResult{ Action: tools.ElicitationActionDecline, }, nil } + sessionID := spec.sessionID + if sessionID == "" { + sessionID = genai.ConversationIDFromContext(ctx) + } + agentName := spec.agentName + if agentName == "" { + agentName = r.currentAgentName() + } + // A background session (run_background_agent) marks its context so // toolset Start() OAuth fails fast instead of eliciting (#3200). Mid-call // elicitations reach here regardless of that marker, so extend the same @@ -479,8 +517,8 @@ func (r *LocalRuntime) elicitationHandler(ctx context.Context, req *mcp.ElicitPa // all can answer this request. Decline immediately with a model-readable // note instead of parking a goroutine forever (#3584). if !tools.InteractivePromptsAllowed(ctx) && !r.hasElicitationSink() { - slog.WarnContext(ctx, "Declining elicitation: background session has no UI to answer it", "message", req.Message) - r.elicitationDeclines.record(genai.ConversationIDFromContext(ctx), backgroundElicitationDeclinedNote(req.Message)) + slog.WarnContext(ctx, "Declining elicitation: background session has no UI to answer it", "message", spec.message) + r.elicitationDeclines.record(sessionID, backgroundElicitationDeclinedNote(spec.message)) return tools.ElicitationResult{ Action: tools.ElicitationActionDecline, }, nil @@ -494,7 +532,7 @@ func (r *LocalRuntime) elicitationHandler(ctx context.Context, req *mcp.ElicitPa // The registry key (and the ElicitationID surfaced to clients for // ResumeElicitation routing) is always a freshly generated, internal - // ID — never the MCP wire req.ElicitationID. The wire value is only + // ID — never the MCP wire elicitation ID. The wire value is only // ever set for URL-mode elicitations and is chosen by the originating // MCP server; two independent servers (e.g. two background jobs each // talking to their own MCP process) can legitimately reuse the same @@ -513,16 +551,15 @@ func (r *LocalRuntime) elicitationHandler(ctx context.Context, req *mcp.ElicitPa defer r.elicitationWaiters.abandon(correlationID, wt) slog.DebugContext(ctx, "Sending elicitation request event to client", - "message", req.Message, - "mode", req.Mode, - "requested_schema", req.RequestedSchema, - "url", req.URL, + "message", spec.message, + "mode", spec.mode, + "requested_schema", spec.schema, + "url", spec.url, "elicitation_id", correlationID, - "server_elicitation_id", req.ElicitationID) - slog.DebugContext(ctx, "Elicitation request meta", "meta", req.Meta) + "server_elicitation_id", spec.serverElicitationID) + slog.DebugContext(ctx, "Elicitation request meta", "meta", spec.meta) - sessionID := genai.ConversationIDFromContext(ctx) - ev := ElicitationRequest(req.Message, req.Mode, req.RequestedSchema, req.URL, correlationID, req.ElicitationID, sessionID, req.Meta, r.currentAgentName()) + ev := ElicitationRequest(spec.message, spec.mode, spec.schema, spec.url, correlationID, spec.serverElicitationID, sessionID, spec.meta, agentName) // Reliable delivery: invoked synchronously, unconditionally, and exactly // once, BEFORE anything that could block (#3584 review item 1). This diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index da70c0f43..fb83206ea 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -2,7 +2,6 @@ package runtime import ( "context" - "errors" "fmt" "log/slog" "path" @@ -1328,13 +1327,15 @@ func sanitizeToolCallName(name string) string { // generated-media manifest ([session.GeneratedMediaManifest]), the trust // anchor a resolver must consult before reading a workspace path back. // -// The requested filename is the sanitized provider display name when one -// exists, otherwise a generic "generated-N"; the writer owns MIME/extension +// The requested filename is the prompt-directed path when the provider +// surfaced one ([chat.MediaDelta.RequestedPath], populated by marker +// extraction once that lands), otherwise the sanitized provider display +// name, otherwise a generic "generated-N"; the writer owns MIME/extension // correction and collision suffixing, and the part persists the exact final -// relative path it returns. A corrected extension additionally surfaces a -// bounded user-visible notice naming the final path. Explicit -// prompt-directed naming (and its out-of-workspace confirmation flow) is -// intentionally not implemented here yet. +// path it returns. A corrected extension additionally surfaces a bounded +// user-visible notice naming the final path. A prompt-directed path that +// escapes the workspace (absolute, "..", or "~"-rooted) is redirected into +// the workspace root under the sanitized basename with a bounded warning. // // When no workspace root is available (no provenance anywhere in the parent // chain, or a malformed stored value) every item fails with the same @@ -1394,29 +1395,25 @@ func (r *LocalRuntime) materializeGeneratedMedia(ctx context.Context, sess *sess continue } - requested := safeName - generic := fmt.Sprintf("generated-%d", i+1) - if requested == "" { - requested = generic - } - res, err := workspacemediaWrite(root, requested, m.Data, m.MimeType) - if err != nil && requested != generic && errors.Is(err, workspacemedia.ErrPathEscape) { - // A provider display name the writer refuses even after display - // sanitization (e.g. a Windows-reserved name like "CON.png") must - // not cost the user the item; there is no user-chosen path to - // honor at this stage, so fall back to the generic name. - res, err = workspacemediaWrite(root, generic, m.Data, m.MimeType) - } + res, err := r.writeGeneratedMedia(generatedMediaItem{ + workspaceRoot: root, + requestedPath: m.RequestedPath, + providerName: safeName, + genericName: fmt.Sprintf("generated-%d", i+1), + data: m.Data, + mimeType: m.MimeType, + agentName: agentName, + index: i + 1, + total: len(media), + }, events) if err != nil { warnItemFailed(err) continue } - - if err := r.recordGeneratedFile(ctx, sess.ID, res.RelPath, safeMimeType); err != nil { - // The file is already a real workspace deliverable, so keep the - // reference; without the manifest record inline display will - // refuse to render it (fail closed), which the user should hear - // about. res.RelPath is writer-sanitized and workspace-relative. + manifestRecorded := true + if err := r.recordGeneratedFile(ctx, sess.ID, chat.ArtifactRootWorkspace, res.RelPath, safeMimeType); err != nil { + manifestRecorded = false + // Keep the saved file, but warn that missing manifest authorization prevents display. slog.DebugContext(ctx, "Failed to record generated media in the manifest; the file was written but may not display inline", "agent", agentName, "session_id", sess.ID, "rel_path", res.RelPath, "error", err) if events != nil { @@ -1425,6 +1422,17 @@ func (r *LocalRuntime) materializeGeneratedMedia(ctx context.Context, sess *sess } } + if manifestRecorded { + if err := r.recordGeneratedBlob(ctx, sess.ID, res.RelPath, m.Data); err != nil { + slog.DebugContext(ctx, "Failed to store portable generated media; keeping the workspace file", + "agent", agentName, "session_id", sess.ID, "rel_path", res.RelPath, "error", err) + if events != nil { + warning := generatedBlobWarning(res.RelPath) + events.Emit(Warning(chat.TruncateUTF8Bytes(warning, maxPlaceholderOrWarningBytes), agentName)) + } + } + } + if res.ExtensionCorrected && events != nil { notice := fmt.Sprintf("Saved generated media as %s: the requested extension %q does not match the returned %s data", res.RelPath, res.RequestedExtension, safeMimeType) @@ -1458,20 +1466,33 @@ func (r *LocalRuntime) sessionLookup() session.Lookup { } // recordGeneratedFile writes one manifest record after a successful -// workspace write — materialization is the only writer of the manifest. -func (r *LocalRuntime) recordGeneratedFile(ctx context.Context, sessionID, relPath, mimeType string) error { +// write — materialization is the only writer of the manifest. +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{ SessionID: sessionID, - RelPath: relPath, + RelPath: finalPath, + Root: root, MimeType: mimeType, CreatedAt: r.now(), }) } +func (r *LocalRuntime) recordGeneratedBlob(ctx context.Context, sessionID, finalPath string, data []byte) error { + blobs, ok := r.sessionStore.(session.GeneratedMediaBlobStore) + if !ok { + return fmt.Errorf("session store %T does not implement generated-media blob storage", r.sessionStore) + } + return blobs.AddGeneratedBlob(ctx, sessionID, finalPath, data) +} + +func generatedBlobWarning(finalPath string) string { + return fmt.Sprintf("Saved generated media %s in the workspace, but could not keep a portable copy with the session. %s", finalPath, retryWithDebugAdvice) +} + // workspacemediaWrite is [workspacemedia.Write] behind a package-level // indirection so tests can inject a deterministic failure for one item in a // batch [LocalRuntime.materializeGeneratedMedia] call. Production code must diff --git a/pkg/runtime/materialize_generated_media_test.go b/pkg/runtime/materialize_generated_media_test.go index 0aa36c7cb..ad12482b2 100644 --- a/pkg/runtime/materialize_generated_media_test.go +++ b/pkg/runtime/materialize_generated_media_test.go @@ -2,6 +2,7 @@ package runtime import ( "bytes" + "context" "errors" "fmt" "io/fs" @@ -427,7 +428,65 @@ func TestMaterializeGeneratedMedia_ClassifiedWriteFailureReasons(t *testing.T) { } } -// storeWithoutManifest hides the built-in store's GeneratedMediaManifest +func TestMaterializeGeneratedMedia_StoresPortableBlob(t *testing.T) { + r, store, _ := newMediaTestRuntime(t) + sess, root := workspaceSession(t, "sess-portable") + data := make([]byte, (20<<20)+1) + copy(data, "portable") + data[len(data)-1] = 0x7f + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, []chat.MediaDelta{ + {Data: data, MimeType: "image/png", Name: "cat.png", Size: int64(len(data))}, + }, "root", sink) + require.Len(t, parts, 1) + assert.FileExists(t, filepath.Join(root, "cat.png")) + stored, err := store.(session.GeneratedMediaBlobStore).LookupGeneratedBlob(t.Context(), sess.ID, "cat.png") + require.NoError(t, err) + require.Len(t, stored, len(data)) + assert.Equal(t, byte(0x7f), stored[len(stored)-1]) + assert.Empty(t, sink.warnings()) +} + +type failingBlobStore struct { + session.Store + + err error +} + +func (s failingBlobStore) AddGeneratedFile(ctx context.Context, file session.GeneratedFile) error { + return s.Store.(session.GeneratedMediaManifest).AddGeneratedFile(ctx, file) +} + +func (s failingBlobStore) LookupGeneratedFile(ctx context.Context, sessionID, relPath string) (*session.GeneratedFile, error) { + return s.Store.(session.GeneratedMediaManifest).LookupGeneratedFile(ctx, sessionID, relPath) +} + +func (s failingBlobStore) AddGeneratedBlob(context.Context, string, string, []byte) error { + return s.err +} + +func (s failingBlobStore) LookupGeneratedBlob(ctx context.Context, sessionID, relPath string) ([]byte, error) { + return s.Store.(session.GeneratedMediaBlobStore).LookupGeneratedBlob(ctx, sessionID, relPath) +} + +func TestMaterializeGeneratedMedia_BlobStoreFailureKeepsWorkspaceFileAndWarns(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + r.sessionStore = failingBlobStore{Store: r.sessionStore, err: errors.New("blob persistence failed")} + sess, root := workspaceSession(t, "sess-blob-failure") + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, []chat.MediaDelta{ + {Data: []byte("image"), MimeType: "image/png", Name: "cat.png", Size: 5}, + }, "root", sink) + require.Len(t, parts, 1) + assert.FileExists(t, filepath.Join(root, "cat.png")) + warnings := sink.warnings() + require.Len(t, warnings, 1) + assert.Contains(t, warnings[0].Message, "could not keep a portable copy with the session") + assert.Contains(t, warnings[0].Message, retryWithDebugAdvice) + assert.NotContains(t, warnings[0].Message, "too large") + assertBoundedSingleLineUTF8(t, warnings[0].Message) +} + // implementation: interface embedding only promotes session.Store's own // method set, so the type assertion in recordGeneratedFile fails. type storeWithoutManifest struct{ session.Store } diff --git a/pkg/runtime/media_escape.go b/pkg/runtime/media_escape.go new file mode 100644 index 000000000..307b761d1 --- /dev/null +++ b/pkg/runtime/media_escape.go @@ -0,0 +1,69 @@ +package runtime + +import ( + "errors" + "fmt" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/workspacemedia" +) + +type generatedMediaItem struct { + workspaceRoot string + // requestedPath is the prompt-directed target (untrusted model input, + // see chat.MediaDelta.RequestedPath); empty when the model named + // nothing explicitly. + requestedPath string + providerName string + genericName string + data []byte + mimeType string + agentName string + index, total int +} + +func (r *LocalRuntime) writeGeneratedMedia(item generatedMediaItem, events EventSink) (workspacemedia.Result, error) { + if item.requestedPath == "" { + return r.writeProviderNamedMedia(item) + } + + class, cleaned := workspacemedia.ClassifyRequestedPath(item.requestedPath) + if class == workspacemedia.PathWorkspaceRelative { + res, err := workspacemediaWrite(item.workspaceRoot, cleaned, item.data, item.mimeType) + if err == nil || !errors.Is(err, workspacemedia.ErrPathEscape) { + return res, err + } + } + return r.redirectEscapedMedia(item, events) +} + +func (r *LocalRuntime) writeProviderNamedMedia(item generatedMediaItem) (workspacemedia.Result, error) { + requested := item.providerName + if requested == "" { + requested = item.genericName + } + res, err := workspacemediaWrite(item.workspaceRoot, requested, item.data, item.mimeType) + if err != nil && requested != item.genericName && errors.Is(err, workspacemedia.ErrPathEscape) { + res, err = workspacemediaWrite(item.workspaceRoot, item.genericName, item.data, item.mimeType) + } + return res, err +} + +func (r *LocalRuntime) redirectEscapedMedia(item generatedMediaItem, events EventSink) (workspacemedia.Result, error) { + base := workspacemedia.RequestedBasename(item.requestedPath) + if base == "" { + base = item.genericName + } + res, err := workspacemediaWrite(item.workspaceRoot, base, item.data, item.mimeType) + if err != nil && base != item.genericName && errors.Is(err, workspacemedia.ErrPathEscape) { + res, err = workspacemediaWrite(item.workspaceRoot, item.genericName, item.data, item.mimeType) + } + if err != nil { + return workspacemedia.Result{}, err + } + if events != nil { + warning := fmt.Sprintf("Requested save location for generated media item %d/%d is outside the workspace or unusable; saved as %s in the workspace instead", item.index, item.total, res.RelPath) + events.Emit(Warning(chat.TruncateUTF8Bytes(warning, maxPlaceholderOrWarningBytes), item.agentName)) + } + return res, nil +} diff --git a/pkg/runtime/media_escape_test.go b/pkg/runtime/media_escape_test.go new file mode 100644 index 000000000..632148109 --- /dev/null +++ b/pkg/runtime/media_escape_test.go @@ -0,0 +1,165 @@ +package runtime + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/team" + "github.com/docker/docker-agent/pkg/workspacemedia" +) + +func TestMaterializeGeneratedMedia_RedirectsEscapingPaths(t *testing.T) { + tests := []struct{ name, requested, want string }{ + {"posix absolute", "/outside/cat.jpg", "cat.png"}, + {"windows absolute", `C:\outside\cat.jpg`, "cat.png"}, + {"windows UNC", `\\server\share\cat.jpg`, "cat.png"}, + {"windows drive relative", `C:outside\cat.jpg`, "cat.png"}, + {"posix traversal", "../../outside/cat.jpg", "cat.png"}, + {"windows traversal", `..\..\outside\cat.jpg`, "cat.png"}, + {"home", "~/outside/cat.jpg", "cat.png"}, + {"reserved basename", "/outside/CON.png", "generated-1.png"}, + {"control basename", "/outside/bad\nname.jpg", "bad-name.png"}, + {"invalid UTF-8 basename", "/outside/bad" + string([]byte{0xff}) + ".jpg", "generated-1.png"}, + {"overlong basename", "/outside/" + strings.Repeat("x", 300) + ".png", "generated-1.png"}, + {"unusable basename", "../..", "generated-1.png"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + r, store, _ := newMediaTestRuntime(t) + sess, root := workspaceSession(t, "redirect-"+tt.name) + sink := &collectingSink{} + parts := r.materializeGeneratedMedia(t.Context(), sess, []chat.MediaDelta{{Data: []byte("png"), MimeType: "image/png", Name: "provider.png", RequestedPath: tt.requested, Size: 3}}, "root", sink) + require.Len(t, parts, 1) + doc := parts[0].Document + require.NotNil(t, doc) + assert.Equal(t, tt.want, doc.Source.ArtifactPath) + assert.Equal(t, chat.ArtifactRootWorkspace, doc.Source.ArtifactRoot) + assert.Equal(t, []byte("png"), mustReadFile(t, filepath.Join(root, filepath.FromSlash(tt.want)))) + _, err := manifestOf(t, store).LookupGeneratedFile(t.Context(), sess.ID, tt.want) + require.NoError(t, err) + warnings := sink.warnings() + require.NotEmpty(t, warnings) + for _, warning := range warnings { + assert.NotContains(t, warning.Message, tt.requested) + assert.LessOrEqual(t, len(warning.Message), maxPlaceholderOrWarningBytes) + } + }) + } +} + +func TestMaterializeGeneratedMedia_RedirectsSymlinkEscapeAndAvoidsCollision(t *testing.T) { + requireSymlinkSupport(t) + r, _, _ := newMediaTestRuntime(t) + sess, root := workspaceSession(t, "redirect-symlink") + outside := t.TempDir() + require.NoError(t, os.Symlink(outside, filepath.Join(root, "link"))) + require.NoError(t, os.WriteFile(filepath.Join(root, "cat.png"), []byte("existing"), 0o644)) + parts := r.materializeGeneratedMedia(t.Context(), sess, []chat.MediaDelta{{Data: []byte("new"), MimeType: "image/png", RequestedPath: "link/cat.png", Size: 3}}, "root", &collectingSink{}) + require.Len(t, parts, 1) + assert.Equal(t, "cat-1.png", parts[0].Document.Source.ArtifactPath) + assert.Equal(t, []byte("existing"), mustReadFile(t, filepath.Join(root, "cat.png"))) + assert.Equal(t, []byte("new"), mustReadFile(t, filepath.Join(root, "cat-1.png"))) + entries, err := os.ReadDir(outside) + require.NoError(t, err) + assert.Empty(t, entries) +} + +func requireSymlinkSupport(t *testing.T) { + t.Helper() + root := t.TempDir() + err := os.Symlink("target", filepath.Join(root, "probe")) + if err != nil { + t.Skipf("symlinks unavailable: %v", err) + } +} + +func TestRunStream_EscapingGeneratedMediaCompletesWithoutElicitation(t *testing.T) { + for _, nonInteractive := range []bool{false, true} { + t.Run(fmt.Sprintf("nonInteractive=%t", nonInteractive), func(t *testing.T) { + workspace := t.TempDir() + external := filepath.Join(t.TempDir(), "cat.png") + store := session.NewInMemorySessionStore() + stream := &mockStream{responses: []chat.MessageStreamResponse{ + {Choices: []chat.MessageStreamChoice{{Index: 0, Delta: chat.MessageDelta{Content: "completed", Media: []chat.MediaDelta{{Data: []byte("png"), MimeType: "image/png", Name: "provider.png", RequestedPath: external, Size: 3}}}}}}, + {Choices: []chat.MessageStreamChoice{{Index: 0, FinishReason: chat.FinishReasonStop}}, Usage: &chat.Usage{InputTokens: 1, OutputTokens: 1}}, + }} + root := agent.New("root", "instructions", agent.WithModel(&mockProvider{id: "test/media", stream: stream})) + rt, err := New(t.Context(), team.New(team.WithAgents(root)), WithSessionCompaction(false), WithSessionStore(store)) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, rt.Close()) }) + rt.OnElicitationRequest(func(Event) { t.Error("generated media escape must not elicit") }) + opts := []session.Opt{session.WithUserMessage("save the generated image")} + if nonInteractive { + opts = append(opts, session.WithNonInteractive(true)) + } + sess := session.New(opts...) + sess.ID = "media-escape" + sess.WorkingDir = workspace + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + var warnings int + var stopped bool + for event := range rt.RunStream(ctx, sess) { + switch event.(type) { + case *WarningEvent: + warnings++ + case *StreamStoppedEvent: + stopped = true + } + } + require.NoError(t, ctx.Err()) + assert.True(t, stopped) + assert.Positive(t, warnings) + assert.NoFileExists(t, external) + assert.Equal(t, []byte("png"), mustReadFile(t, filepath.Join(workspace, "cat.png"))) + stored, err := store.GetSession(t.Context(), sess.ID) + require.NoError(t, err) + messages := stored.GetAllMessages() + require.Len(t, messages, 2) + assistant := messages[1].Message + assert.Equal(t, "completed", assistant.Content) + require.Len(t, assistant.MultiContent, 2) + document := assistant.MultiContent[1].Document + require.NotNil(t, document) + assert.Equal(t, "cat.png", document.Source.ArtifactPath) + }) + } +} + +func TestMaterializeGeneratedMedia_RedirectFailurePreservesSibling(t *testing.T) { + r, _, _ := newMediaTestRuntime(t) + sess, root := workspaceSession(t, "redirect-partial") + original := workspacemediaWrite + workspacemediaWrite = func(workspaceRoot, requested string, data []byte, mimeType string) (workspacemedia.Result, error) { + if requested == "bad.png" { + return workspacemedia.Result{}, os.ErrPermission + } + return original(workspaceRoot, requested, data, mimeType) + } + t.Cleanup(func() { workspacemediaWrite = original }) + parts := r.materializeGeneratedMedia(t.Context(), sess, []chat.MediaDelta{ + {Data: []byte("bad"), MimeType: "image/png", RequestedPath: "/outside/bad.png", Size: 3}, + {Data: []byte("good"), MimeType: "image/png", RequestedPath: "nested/good.png", Size: 4}, + }, "root", &collectingSink{}) + require.Len(t, parts, 1) + assert.Equal(t, "nested/good.png", parts[0].Document.Source.ArtifactPath) + assert.Equal(t, []byte("good"), mustReadFile(t, filepath.Join(root, "nested", "good.png"))) +} + +func mustReadFile(t *testing.T, path string) []byte { + t.Helper() + b, err := os.ReadFile(path) + require.NoError(t, err) + return b +} diff --git a/pkg/session/generated_media_manifest.go b/pkg/session/generated_media_manifest.go index 5887c487e..ebdf73b01 100644 --- a/pkg/session/generated_media_manifest.go +++ b/pkg/session/generated_media_manifest.go @@ -8,6 +8,9 @@ import ( "io/fs" "strings" "time" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/concurrent" ) var ( @@ -16,22 +19,33 @@ var ( // generated file" and refuse to read the workspace path. ErrGeneratedFileNotFound = errors.New("generated file not found in manifest") + ErrGeneratedBlobNotFound = errors.New("generated media blob not found") + // ErrInvalidGeneratedFilePath is returned for a path that can never be a - // workspacemedia.Write result (empty, absolute, traversal, NUL, ...). + // pkg/workspacemedia write result for its root kind (empty, traversal, + // NUL, wrong absolute/relative shape, ...). ErrInvalidGeneratedFilePath = errors.New("invalid generated file path") + + // ErrInvalidGeneratedFileRoot is returned for a root kind the manifest + // does not record. + ErrInvalidGeneratedFileRoot = errors.New("invalid generated file root kind") ) -// GeneratedFile is one generated-media manifest record: a workspace file -// written by materialization on behalf of the owning session. +// GeneratedFile is one generated-media manifest record: a file written by +// materialization on behalf of the owning session. type GeneratedFile struct { // SessionID is the OWNING session — the session active when the media // was generated, permanent across branch/fork. SessionID string - // RelPath is the workspace-relative, slash-separated path exactly as - // returned by workspacemedia.Write. + // RelPath is the exact workspace-relative, slash-separated path written. RelPath string + // Root is the artifact root kind the path is interpreted against: + // chat.ArtifactRootWorkspace (the default — an empty value is normalized + // to it). Resolution must require it to match the reference's root. + Root chat.ArtifactRootKind + // MimeType is the sanitized MIME type of the written content. MimeType string @@ -39,33 +53,66 @@ type GeneratedFile struct { CreatedAt time.Time } -// GeneratedMediaManifest records which workspace files generated-media +// GeneratedMediaManifest records which files generated-media // materialization wrote. It is the trust anchor for resolving a -// workspace-rooted artifact reference (chat.ArtifactRootWorkspace): a -// workspace path may only be read back if the (owner session, path) pair -// was recorded here by materialization itself — session JSON alone must -// never be able to select an arbitrary workspace file such as ".env" or a -// source file. Only materialization may call AddGeneratedFile. +// generated-media artifact reference: a path may only be read back if the (owner +// session, path) pair was recorded here by materialization itself — session +// JSON alone must never be able to select an arbitrary file such as ".env" +// or a source file. Only materialization may call AddGeneratedFile. // // Implemented by the built-in session stores; resolvers obtain it by type // asserting their session.Store. type GeneratedMediaManifest interface { // AddGeneratedFile records file. The path is validated against the - // workspacemedia.Write output shape and rejected with - // ErrInvalidGeneratedFilePath otherwise. + // pkg/workspacemedia output shape for file.Root and rejected with + // ErrInvalidGeneratedFilePath (or ErrInvalidGeneratedFileRoot) + // otherwise. AddGeneratedFile(ctx context.Context, file GeneratedFile) error // LookupGeneratedFile returns the record for (sessionID, relPath), or // ErrGeneratedFileNotFound when materialization never wrote that path // for that session. Invalid inputs fail with ErrInvalidGeneratedFilePath - // (or ErrEmptyID) rather than being normalized. + // (or ErrEmptyID) rather than being normalized. Callers must additionally + // require the returned Root to match their reference's ArtifactRoot. LookupGeneratedFile(ctx context.Context, sessionID, relPath string) (*GeneratedFile, error) } +// GeneratedMediaBlobStore persists portable copies of generated media. Blob +// lookup remains manifest-gated: callers must first validate the corresponding +// GeneratedFile record and its root kind. +type GeneratedMediaBlobStore interface { + AddGeneratedBlob(ctx context.Context, sessionID, relPath string, data []byte) error + LookupGeneratedBlob(ctx context.Context, sessionID, relPath string) ([]byte, error) +} + +// normalizeGeneratedFileRoot maps the zero value to the workspace root kind, +// so callers omitting the root and legacy rows keep their workspace meaning. +func normalizeGeneratedFileRoot(root chat.ArtifactRootKind) chat.ArtifactRootKind { + if root == "" { + return chat.ArtifactRootWorkspace + } + return root +} + +// validateGeneratedFileRecord vets a manifest record at the add boundary: +// fail closed on any (root, path) combination pkg/workspacemedia could never +// have produced, so neither a buggy writer nor a tampered caller can smuggle +// a mis-rooted path into the manifest. +func validateGeneratedFileRecord(sessionID string, root chat.ArtifactRootKind, relPath string) error { + if err := validateGeneratedFileKey(sessionID, relPath); err != nil { + return err + } + if normalizeGeneratedFileRoot(root) != chat.ArtifactRootWorkspace { + return fmt.Errorf("%w: %q", ErrInvalidGeneratedFileRoot, root) + } + return nil +} + // validateGeneratedFileKey vets a manifest key at the API boundary, on both -// write and lookup: fail closed on anything workspacemedia.Write could never -// have returned, so neither a buggy writer nor a tampered session JSON can -// smuggle an absolute or traversing path through the manifest. +// write and lookup. A key has the workspace-relative shape +// workspacemedia.Write guarantees; anything else — traversal, NUL, +// stray backslashes in a relative path — fails closed so a tampered session +// JSON cannot probe arbitrary files through the manifest. func validateGeneratedFileKey(sessionID, relPath string) error { if sessionID == "" { return ErrEmptyID @@ -73,7 +120,10 @@ func validateGeneratedFileKey(sessionID, relPath string) error { if relPath == "" { return fmt.Errorf("%w: empty path", ErrInvalidGeneratedFilePath) } - if strings.ContainsAny(relPath, "\x00\\") { + if strings.ContainsRune(relPath, '\x00') { + return fmt.Errorf("%w: %q", ErrInvalidGeneratedFilePath, relPath) + } + if strings.ContainsRune(relPath, '\\') { return fmt.Errorf("%w: %q", ErrInvalidGeneratedFilePath, relPath) } // fs.ValidPath rejects absolute paths, ".." segments, empty segments, @@ -93,9 +143,10 @@ func generatedFileKey(sessionID, relPath string) string { } func (s *InMemorySessionStore) AddGeneratedFile(_ context.Context, file GeneratedFile) error { - if err := validateGeneratedFileKey(file.SessionID, file.RelPath); err != nil { + if err := validateGeneratedFileRecord(file.SessionID, file.Root, file.RelPath); err != nil { return err } + file.Root = normalizeGeneratedFileRoot(file.Root) s.generatedFiles.Store(generatedFileKey(file.SessionID, file.RelPath), file) return nil } @@ -113,30 +164,88 @@ func (s *InMemorySessionStore) LookupGeneratedFile(_ context.Context, sessionID, // deleteGeneratedFiles prunes every manifest record owned by sessionID. func (s *InMemorySessionStore) deleteGeneratedFiles(sessionID string) { + deleteGeneratedKeys(s.generatedFiles, sessionID) +} + +func (s *InMemorySessionStore) deleteGeneratedBlobs(sessionID string) { + deleteGeneratedKeys(s.generatedBlobs, sessionID) +} + +func deleteGeneratedKeys[T any](entries *concurrent.Map[string, T], sessionID string) { prefix := generatedFileKey(sessionID, "") var doomed []string - s.generatedFiles.Range(func(key string, _ GeneratedFile) bool { + entries.Range(func(key string, _ T) bool { if strings.HasPrefix(key, prefix) { doomed = append(doomed, key) } return true }) for _, key := range doomed { - s.generatedFiles.Delete(key) + entries.Delete(key) + } +} + +func (s *InMemorySessionStore) AddGeneratedBlob(_ context.Context, sessionID, relPath string, data []byte) error { + if err := validateGeneratedFileKey(sessionID, relPath); err != nil { + return err + } + s.generatedBlobs.Store(generatedFileKey(sessionID, relPath), append([]byte(nil), data...)) + return nil +} + +func (s *InMemorySessionStore) LookupGeneratedBlob(_ context.Context, sessionID, relPath string) ([]byte, error) { + if err := validateGeneratedFileKey(sessionID, relPath); err != nil { + return nil, err + } + data, ok := s.generatedBlobs.Load(generatedFileKey(sessionID, relPath)) + if !ok { + return nil, fmt.Errorf("%w: %q", ErrGeneratedBlobNotFound, relPath) + } + return append([]byte(nil), data...), nil +} + +func (s *SQLiteSessionStore) AddGeneratedBlob(ctx context.Context, sessionID, relPath string, data []byte) error { + if err := validateGeneratedFileKey(sessionID, relPath); err != nil { + return err + } + _, err := s.db.ExecContext(ctx, ` + INSERT INTO generated_media_blobs (session_id, rel_path, data) + VALUES (?, ?, ?) + ON CONFLICT (session_id, rel_path) DO UPDATE SET data = excluded.data + `, sessionID, relPath, data) + return err +} + +func (s *SQLiteSessionStore) LookupGeneratedBlob(ctx context.Context, sessionID, relPath string) ([]byte, error) { + if err := validateGeneratedFileKey(sessionID, relPath); err != nil { + return nil, err + } + var data []byte + err := s.db.QueryRowContext(ctx, ` + SELECT data FROM generated_media_blobs + WHERE session_id = ? AND rel_path = ? + `, sessionID, relPath).Scan(&data) + if errors.Is(err, sql.ErrNoRows) { + return nil, fmt.Errorf("%w: %q", ErrGeneratedBlobNotFound, relPath) + } + if err != nil { + return nil, err } + return data, nil } func (s *SQLiteSessionStore) AddGeneratedFile(ctx context.Context, file GeneratedFile) error { - if err := validateGeneratedFileKey(file.SessionID, file.RelPath); err != nil { + if err := validateGeneratedFileRecord(file.SessionID, file.Root, file.RelPath); err != nil { return err } _, err := s.db.ExecContext(ctx, ` - INSERT INTO generated_media_manifest (session_id, rel_path, mime_type, created_at) - VALUES (?, ?, ?, ?) + INSERT INTO generated_media_manifest (session_id, rel_path, root_kind, mime_type, created_at) + VALUES (?, ?, ?, ?, ?) ON CONFLICT (session_id, rel_path) DO UPDATE SET + root_kind = excluded.root_kind, mime_type = excluded.mime_type, created_at = excluded.created_at - `, file.SessionID, file.RelPath, file.MimeType, file.CreatedAt.UTC().Format(time.RFC3339Nano)) + `, file.SessionID, file.RelPath, string(normalizeGeneratedFileRoot(file.Root)), file.MimeType, file.CreatedAt.UTC().Format(time.RFC3339Nano)) return err } @@ -145,17 +254,18 @@ func (s *SQLiteSessionStore) LookupGeneratedFile(ctx context.Context, sessionID, return nil, err } file := GeneratedFile{SessionID: sessionID, RelPath: relPath} - var createdAt string + var root, createdAt string err := s.db.QueryRowContext(ctx, ` - SELECT mime_type, created_at FROM generated_media_manifest + SELECT root_kind, mime_type, created_at FROM generated_media_manifest WHERE session_id = ? AND rel_path = ? - `, sessionID, relPath).Scan(&file.MimeType, &createdAt) + `, sessionID, relPath).Scan(&root, &file.MimeType, &createdAt) if errors.Is(err, sql.ErrNoRows) { return nil, fmt.Errorf("%w: %q", ErrGeneratedFileNotFound, relPath) } if err != nil { return nil, err } + file.Root = normalizeGeneratedFileRoot(chat.ArtifactRootKind(root)) file.CreatedAt = parseCreatedAt(createdAt) return &file, nil } diff --git a/pkg/session/generated_media_manifest_test.go b/pkg/session/generated_media_manifest_test.go index 9a32888ba..fe4702cb0 100644 --- a/pkg/session/generated_media_manifest_test.go +++ b/pkg/session/generated_media_manifest_test.go @@ -1,11 +1,14 @@ package session import ( + "path/filepath" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/chat" ) // manifestStores runs a subtest against both built-in Store implementations, @@ -24,6 +27,34 @@ func manifestStores(t *testing.T, run func(t *testing.T, store Store, manifest G }) } +func TestGeneratedMediaBlobStore_RoundTripAndCopy(t *testing.T) { + t.Parallel() + manifestStores(t, func(t *testing.T, _ Store, manifest GeneratedMediaManifest) { + t.Helper() + blobs, ok := manifest.(GeneratedMediaBlobStore) + require.True(t, ok) + ctx := t.Context() + input := make([]byte, (20<<20)+1) + copy(input, "portable image") + input[len(input)-1] = 0x7f + require.NoError(t, blobs.AddGeneratedBlob(ctx, "owner", "images/cat.png", input)) + input[0] = 'X' + + got, err := blobs.LookupGeneratedBlob(ctx, "owner", "images/cat.png") + require.NoError(t, err) + require.Len(t, got, len(input)) + assert.Equal(t, byte('p'), got[0], "store must retain its own copy") + assert.Equal(t, byte(0x7f), got[len(got)-1], "large payload must round-trip through storage") + got[1] = 'X' + again, err := blobs.LookupGeneratedBlob(ctx, "owner", "images/cat.png") + require.NoError(t, err) + assert.Equal(t, byte('o'), again[1], "lookup must return a copy") + + _, err = blobs.LookupGeneratedBlob(ctx, "owner", "missing.png") + require.ErrorIs(t, err, ErrGeneratedBlobNotFound) + }) +} + func TestGeneratedMediaManifest_RoundTrip(t *testing.T) { t.Parallel() manifestStores(t, func(t *testing.T, _ Store, manifest GeneratedMediaManifest) { @@ -69,15 +100,14 @@ func TestGeneratedMediaManifest_RefusesUnrecordedPaths(t *testing.T) { } // TestGeneratedMediaManifest_RejectsInvalidPaths pins the API-boundary -// validation on BOTH write and lookup: shapes workspacemedia.Write can never -// return (absolute, traversal, backslashes, NUL, empty/dot segments) fail -// with ErrInvalidGeneratedFilePath before touching storage. +// validation on BOTH write and lookup: shapes pkg/workspacemedia can never +// return (traversal, backslashes in relative paths, NUL, empty/dot +// segments, unclean absolutes) fail with ErrInvalidGeneratedFilePath before +// touching storage. func TestGeneratedMediaManifest_RejectsInvalidPaths(t *testing.T) { t.Parallel() invalid := []string{ "", - "/etc/passwd", - "/abs/cat.png", "../outside.png", "images/../../outside.png", "images/./cat.png", @@ -85,6 +115,8 @@ func TestGeneratedMediaManifest_RejectsInvalidPaths(t *testing.T) { "images/cat.png/", `images\cat.png`, "cat\x00.png", + "/abs\x00/cat.png", + "/abs/../cat.png", ".", "..", } @@ -107,6 +139,46 @@ func TestGeneratedMediaManifest_RejectsInvalidPaths(t *testing.T) { }) } +// TestGeneratedMediaManifest_WorkspaceRootNormalized verifies empty root +// kinds retain their workspace meaning on read-back. +func TestGeneratedMediaManifest_WorkspaceRootNormalized(t *testing.T) { + t.Parallel() + manifestStores(t, func(t *testing.T, _ Store, manifest GeneratedMediaManifest) { + t.Helper() + require.NoError(t, manifest.AddGeneratedFile(t.Context(), GeneratedFile{ + SessionID: "owner", RelPath: "cat.png", MimeType: "image/png", CreatedAt: time.Now(), + })) + got, err := manifest.LookupGeneratedFile(t.Context(), "owner", "cat.png") + require.NoError(t, err) + assert.Equal(t, chat.ArtifactRootWorkspace, got.Root) + }) +} + +// TestGeneratedMediaManifest_RejectsMisRootedRecords: the add boundary +// refuses any (root, path) combination pkg/workspacemedia could never have +// produced, and unrecorded absolute paths still fail closed on lookup — a +// tampered session JSON cannot probe /etc/passwd through the manifest. +func TestGeneratedMediaManifest_RejectsMisRootedRecords(t *testing.T) { + t.Parallel() + manifestStores(t, func(t *testing.T, _ Store, manifest GeneratedMediaManifest) { + t.Helper() + absolutePath := filepath.Join(t.TempDir(), "abs", "cat.png") + add := func(root chat.ArtifactRootKind, p string) error { + return manifest.AddGeneratedFile(t.Context(), GeneratedFile{ + SessionID: "owner", RelPath: p, Root: root, MimeType: "image/png", CreatedAt: time.Now(), + }) + } + require.ErrorIs(t, add("", absolutePath), ErrInvalidGeneratedFilePath, + "a workspace record must never carry an absolute path") + require.ErrorIs(t, add(chat.ArtifactRootWorkspace, absolutePath), ErrInvalidGeneratedFilePath) + require.ErrorIs(t, add("external", "cat.png"), ErrInvalidGeneratedFileRoot) + require.ErrorIs(t, add("attacker-root", "cat.png"), ErrInvalidGeneratedFileRoot) + + _, err := manifest.LookupGeneratedFile(t.Context(), "owner", absolutePath) + require.ErrorIs(t, err, ErrInvalidGeneratedFilePath) + }) +} + // TestGeneratedMediaManifest_DeleteSessionPrunesRecords: the manifest table // has no foreign key (the session row may not exist yet when materialization // records a file), so DeleteSession must prune records explicitly. @@ -168,3 +240,17 @@ func TestGeneratedMediaManifest_ReAddUpdatesRecord(t *testing.T) { assert.Equal(t, "image/webp", got.MimeType) }) } + +func TestGeneratedMediaManifest_DeleteSessionPrunesBlob(t *testing.T) { + t.Parallel() + manifestStores(t, func(t *testing.T, store Store, manifest GeneratedMediaManifest) { + t.Helper() + ctx := t.Context() + require.NoError(t, store.AddSession(ctx, New(WithID("doomed-blob")))) + blobs := manifest.(GeneratedMediaBlobStore) + require.NoError(t, blobs.AddGeneratedBlob(ctx, "doomed-blob", "cat.png", []byte("png"))) + require.NoError(t, store.DeleteSession(ctx, "doomed-blob")) + _, err := blobs.LookupGeneratedBlob(ctx, "doomed-blob", "cat.png") + require.ErrorIs(t, err, ErrGeneratedBlobNotFound) + }) +} diff --git a/pkg/session/migrations.go b/pkg/session/migrations.go index 415d0ecc0..5d35296c6 100644 --- a/pkg/session/migrations.go +++ b/pkg/session/migrations.go @@ -463,6 +463,26 @@ func getAllMigrations() []Migration { `, DownSQL: `DROP TABLE IF EXISTS generated_media_manifest`, }, + { + ID: 29, + Name: "029_add_generated_media_blob_table", + Description: "Store generated-media bytes in the session database for portable session rendering", + UpSQL: ` + CREATE TABLE IF NOT EXISTS generated_media_blobs ( + session_id TEXT NOT NULL, + rel_path TEXT NOT NULL, + data BLOB NOT NULL, + PRIMARY KEY (session_id, rel_path) + ) + `, + DownSQL: `DROP TABLE IF EXISTS generated_media_blobs`, + }, + { + ID: 30, + Name: "030_add_root_kind_to_generated_media_manifest", + Description: "Add root_kind to generated_media_manifest so user-confirmed out-of-workspace generated files stay manifest-gated alongside workspace-relative ones", + UpSQL: `ALTER TABLE generated_media_manifest ADD COLUMN root_kind TEXT NOT NULL DEFAULT 'workspace'`, + }, } } diff --git a/pkg/session/migrations_pinned_test.go b/pkg/session/migrations_pinned_test.go index f25a66f65..5e0cbaa73 100644 --- a/pkg/session/migrations_pinned_test.go +++ b/pkg/session/migrations_pinned_test.go @@ -39,7 +39,7 @@ func TestMigrationCatalogIsContentPinned(t *testing.T) { got := digestMigrationCatalog(getAllMigrations()) - const wantDigest = "18f9416ab037c50b6cfef0c9ea42787ce53dc303ed3c119974b4eab24eba55e6" + const wantDigest = "e2d764e159e9c952596cdac552c860df2df3c90d5cf00fdac0ff4d9eda5b15ba" if got != wantDigest { t.Fatalf(`migration catalogue content has changed. diff --git a/pkg/session/store.go b/pkg/session/store.go index d949497bc..18abd20c0 100644 --- a/pkg/session/store.go +++ b/pkg/session/store.go @@ -150,6 +150,7 @@ type Store interface { type InMemorySessionStore struct { sessions *concurrent.Map[string, *Session] generatedFiles *concurrent.Map[string, GeneratedFile] // keyed by generatedFileKey + generatedBlobs *concurrent.Map[string, []byte] // keyed by generatedFileKey messageID atomic.Int64 // counter for message IDs, incremented via Add(1) } @@ -157,6 +158,7 @@ func NewInMemorySessionStore() Store { return &InMemorySessionStore{ sessions: concurrent.NewMap[string, *Session](), generatedFiles: concurrent.NewMap[string, GeneratedFile](), + generatedBlobs: concurrent.NewMap[string, []byte](), } } @@ -236,6 +238,7 @@ func (s *InMemorySessionStore) DeleteSession(_ context.Context, id string) error } s.sessions.Delete(id) s.deleteGeneratedFiles(id) + s.deleteGeneratedBlobs(id) return nil } @@ -1001,14 +1004,22 @@ func (s *SQLiteSessionStore) DeleteSession(ctx context.Context, id string) error return ErrEmptyID } - result, err := s.db.ExecContext(ctx, "DELETE FROM sessions WHERE id = ?", id) + // These tables carry no foreign keys because media may be recorded before + // the lazily persisted session row exists, so prune them explicitly. + tx, err := s.db.BeginTx(ctx, nil) if err != nil { return err } + defer func() { _ = tx.Rollback() }() - // The manifest table carries no foreign key (see migration - // 028_add_generated_media_manifest_table), so prune explicitly. - if _, err := s.db.ExecContext(ctx, "DELETE FROM generated_media_manifest WHERE session_id = ?", id); err != nil { + result, err := tx.ExecContext(ctx, "DELETE FROM sessions WHERE id = ?", id) + if err != nil { + return err + } + if _, err := tx.ExecContext(ctx, "DELETE FROM generated_media_blobs WHERE session_id = ?", id); err != nil { + return err + } + if _, err := tx.ExecContext(ctx, "DELETE FROM generated_media_manifest WHERE session_id = ?", id); err != nil { return err } @@ -1021,7 +1032,7 @@ func (s *SQLiteSessionStore) DeleteSession(ctx context.Context, id string) error return ErrNotFound } - return nil + return tx.Commit() } // UpdateSession updates an existing session's metadata, or creates it if it doesn't exist (upsert). diff --git a/pkg/workspacemedia/classify.go b/pkg/workspacemedia/classify.go new file mode 100644 index 000000000..00c2db9aa --- /dev/null +++ b/pkg/workspacemedia/classify.go @@ -0,0 +1,85 @@ +package workspacemedia + +import ( + "slices" + "strings" + "unicode/utf8" +) + +// PathClass is the outcome of classifying a model-requested target path +// before any I/O. Classification is purely lexical: a PathWorkspaceRelative +// path can still be refused at write time when a symlinked parent resolves +// outside the root (surfaced as [ErrPathEscape] by [Write]). +type PathClass int + +const ( + // PathInvalid marks a path that names nothing usable: empty, only + // separators/dots, an invalid or Windows-reserved final segment. It is + // unusable — callers should fall back to a safe generated name. + PathInvalid PathClass = iota + + // PathWorkspaceRelative marks a path contained under the workspace + // root, directly writable via [Write]. + PathWorkspaceRelative + + // PathEscaping marks a path that targets a location outside the + // workspace: absolute, traversing above the root via "..", or rooted at + // a home directory via a leading "~" segment. + PathEscaping +) + +// ClassifyRequestedPath classifies requested and, for PathWorkspaceRelative, +// returns the normalized slash-separated relative path to hand to [Write] +// (both separator styles accepted; interior "." and ".." segments resolved +// lexically, so "a/../b.png" is contained rather than escaping). For every +// other class the second return is "". +func ClassifyRequestedPath(requested string) (PathClass, string) { + if isAbsolutePath(requested) { + return PathEscaping, "" + } + segments := splitPathSegments(requested) + if len(segments) > 0 && strings.HasPrefix(segments[0], "~") { + return PathEscaping, "" + } + + var stack []string + for _, seg := range segments { + switch seg { + case ".": + case "..": + if len(stack) == 0 { + return PathEscaping, "" + } + stack = stack[:len(stack)-1] + default: + stack = append(stack, seg) + } + } + cleaned := strings.Join(stack, "/") + if _, _, _, err := splitRequestedPath(cleaned); err != nil { + return PathInvalid, "" + } + return PathWorkspaceRelative, cleaned +} + +// RequestedBasename returns the final meaningful segment of a requested +// path — the name to redirect to when the full path is refused — or "" +// when none exists (empty, separators only, or only "."/".." segments). +func RequestedBasename(requested string) string { + segments := splitPathSegments(requested) + for _, seg := range slices.Backward(segments) { + if seg != "." && seg != ".." { + if !utf8.ValidString(seg) || len(seg) > 255 { + return "" + } + return seg + } + } + return "" +} + +// splitPathSegments splits on both separator styles, mirroring +// splitRequestedPath: model-provided paths may be Windows-style. +func splitPathSegments(requested string) []string { + return strings.FieldsFunc(requested, func(r rune) bool { return r == '/' || r == '\\' }) +} diff --git a/pkg/workspacemedia/classify_test.go b/pkg/workspacemedia/classify_test.go new file mode 100644 index 000000000..18bc168e2 --- /dev/null +++ b/pkg/workspacemedia/classify_test.go @@ -0,0 +1,66 @@ +package workspacemedia + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestClassifyRequestedPath(t *testing.T) { + t.Parallel() + tests := []struct { + requested string + class PathClass + cleaned string + }{ + {"cat.png", PathWorkspaceRelative, "cat.png"}, + {"images/cat.png", PathWorkspaceRelative, "images/cat.png"}, + {`images\cat.png`, PathWorkspaceRelative, "images/cat.png"}, + {"./images/cat.png", PathWorkspaceRelative, "images/cat.png"}, + {"a/../b.png", PathWorkspaceRelative, "b.png"}, + {"a/b/../../c.png", PathWorkspaceRelative, "c.png"}, + + {"/abs/cat.png", PathEscaping, ""}, + {`\abs\cat.png`, PathEscaping, ""}, + {`C:\abs\cat.png`, PathEscaping, ""}, + {"../cat.png", PathEscaping, ""}, + {"a/../../cat.png", PathEscaping, ""}, + {"~/cat.png", PathEscaping, ""}, + {"~", PathEscaping, ""}, + {"~user/cat.png", PathEscaping, ""}, + + {"//", PathEscaping, ""}, + + {"", PathInvalid, ""}, + {".", PathInvalid, ""}, + {"a/..", PathInvalid, ""}, + {"CON.png", PathInvalid, ""}, + {"...", PathInvalid, ""}, + } + for _, tt := range tests { + class, cleaned := ClassifyRequestedPath(tt.requested) + assert.Equal(t, tt.class, class, "class of %q", tt.requested) + assert.Equal(t, tt.cleaned, cleaned, "cleaned form of %q", tt.requested) + } +} + +func TestRequestedBasename(t *testing.T) { + t.Parallel() + tests := []struct { + requested string + want string + }{ + {"cat.png", "cat.png"}, + {"/abs/dir/cat.png", "cat.png"}, + {"../outside/cat.png", "cat.png"}, + {`C:\dir\cat.png`, "cat.png"}, + {"dir/name/..", "name"}, + {"", ""}, + {"/", ""}, + {"../..", ""}, + {"./.", ""}, + } + for _, tt := range tests { + assert.Equal(t, tt.want, RequestedBasename(tt.requested), "basename of %q", tt.requested) + } +} diff --git a/pkg/workspacemedia/writer.go b/pkg/workspacemedia/writer.go index 311d7f1cd..17849e3c3 100644 --- a/pkg/workspacemedia/writer.go +++ b/pkg/workspacemedia/writer.go @@ -50,8 +50,8 @@ var ErrPathEscape = errors.New("path escapes the workspace or has invalid segmen // Result describes a completed write. type Result struct { - // RelPath is the exact final path written, relative to the workspace - // root and slash-separated. Persist this verbatim. + // RelPath is the exact final path written, relative to the workspace root + // and slash-separated. Persist this verbatim. RelPath string // ExtensionCorrected reports that the requested filename's extension