From 47e701830dc5f87b48cadc561258e1b93b7c44e6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Fri, 31 Jul 2026 17:59:44 +0200 Subject: [PATCH 1/2] feat: keep generated-media saves inside the owning workspace Redirect absolute, traversal, home-relative and symlink-escaping targets to a sanitized workspace basename with a bounded warning. Generated media never elicits a path decision or writes externally; normal ACP and CLI turns complete without a form consumer. Preserve MIME correction, collision-safe atomic writes, partial success, owning-session provenance and workspace-only manifest authorization. Keep the pinned root-kind migration for database compatibility. --- pkg/acp/agent.go | 6 +- pkg/acp/agent_test.go | 14 ++ pkg/acp/runagent_test.go | 84 ++++++++++ pkg/chat/document.go | 3 +- pkg/chat/media.go | 10 ++ pkg/cli/runner.go | 5 +- pkg/cli/runner_test.go | 57 ++++++- pkg/runtime/elicitation.go | 61 +++++-- pkg/runtime/loop.go | 54 +++--- pkg/runtime/media_escape.go | 69 ++++++++ pkg/runtime/media_escape_test.go | 165 +++++++++++++++++++ pkg/session/generated_media_manifest.go | 94 ++++++++--- pkg/session/generated_media_manifest_test.go | 54 +++++- pkg/session/migrations.go | 6 + pkg/session/migrations_pinned_test.go | 2 +- pkg/workspacemedia/classify.go | 85 ++++++++++ pkg/workspacemedia/classify_test.go | 66 ++++++++ pkg/workspacemedia/writer.go | 4 +- 18 files changed, 756 insertions(+), 83 deletions(-) create mode 100644 pkg/runtime/media_escape.go create mode 100644 pkg/runtime/media_escape_test.go create mode 100644 pkg/workspacemedia/classify.go create mode 100644 pkg/workspacemedia/classify_test.go diff --git a/pkg/acp/agent.go b/pkg/acp/agent.go index 37bf8e400f..354dec7e39 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 d174253cb8..1ed626d8cc 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 c36a32484b..faba21a225 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 9f822eec70..6b02aadaf7 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 f028f7b77e..4373a31777 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 e730fb2df9..ea15f98902 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 c393b13472..74d89500b2 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 85bb1ee0e2..9970ef2df8 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 da70c0f434..c77fb37c69 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,23 @@ 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. + if err := r.recordGeneratedFile(ctx, sess.ID, chat.ArtifactRootWorkspace, res.RelPath, safeMimeType); err != nil { + // 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 { @@ -1458,15 +1453,16 @@ 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(), }) diff --git a/pkg/runtime/media_escape.go b/pkg/runtime/media_escape.go new file mode 100644 index 0000000000..307b761d19 --- /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 0000000000..6321481090 --- /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 5887c487e0..23d303b91e 100644 --- a/pkg/session/generated_media_manifest.go +++ b/pkg/session/generated_media_manifest.go @@ -8,6 +8,8 @@ import ( "io/fs" "strings" "time" + + "github.com/docker/docker-agent/pkg/chat" ) var ( @@ -17,21 +19,30 @@ var ( ErrGeneratedFileNotFound = errors.New("generated file not found in manifest") // 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 +50,58 @@ 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) } +// 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 +109,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 +132,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 } @@ -127,16 +167,17 @@ func (s *InMemorySessionStore) deleteGeneratedFiles(sessionID string) { } 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 +186,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 9a32888bae..6bff5ae4c7 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, @@ -69,15 +72,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 +87,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 +111,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. diff --git a/pkg/session/migrations.go b/pkg/session/migrations.go index 415d0ecc01..17682a651d 100644 --- a/pkg/session/migrations.go +++ b/pkg/session/migrations.go @@ -463,6 +463,12 @@ func getAllMigrations() []Migration { `, DownSQL: `DROP TABLE IF EXISTS generated_media_manifest`, }, + { + ID: 29, + Name: "029_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 f25a66f654..ee075528ac 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 = "3d7db51380ef6b4ef5958a85eee1be279a6a96f7a11bf77b1b68f3a821218d9d" if got != wantDigest { t.Fatalf(`migration catalogue content has changed. diff --git a/pkg/workspacemedia/classify.go b/pkg/workspacemedia/classify.go new file mode 100644 index 0000000000..00c2db9aa9 --- /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 0000000000..18bc168e2d --- /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 311d7f1cd3..17849e3c32 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 From 6cb0b30b276fd76a137003587fbbdaf455fb2a45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Fri, 4 Sep 2026 20:08:50 +0200 Subject: [PATCH 2/2] feat(#3996): persist portable generated media blobs After successful workspace materialization and manifest recording, store the bytes under the owning session and final artifact identity. Add the SQLite blob table as migration 029 and place the manifest root-kind migration at 030. The workspace deliverable remains available at its requested or collision-adjusted path; the session copy enables portable resolution without embedding generated base64 in session JSON. Blob writes are all-or-error without a per-item byte cap. A failed blob write preserves the workspace file and emits a safe portability warning. DeleteSession removes both manifest and blob rows. The in-memory store mirrors add, lookup, defensive copying, overwrite, and cleanup semantics. --- pkg/runtime/loop.go | 25 +++++++ .../materialize_generated_media_test.go | 61 +++++++++++++++- pkg/session/generated_media_manifest.go | 72 ++++++++++++++++++- pkg/session/generated_media_manifest_test.go | 42 +++++++++++ pkg/session/migrations.go | 16 ++++- pkg/session/migrations_pinned_test.go | 2 +- pkg/session/store.go | 21 ++++-- 7 files changed, 229 insertions(+), 10 deletions(-) diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index c77fb37c69..fb83206eab 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -1410,7 +1410,9 @@ func (r *LocalRuntime) materializeGeneratedMedia(ctx context.Context, sess *sess warnItemFailed(err) continue } + 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) @@ -1420,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) @@ -1468,6 +1481,18 @@ func (r *LocalRuntime) recordGeneratedFile(ctx context.Context, sessionID string }) } +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 0aa36c7cbf..ad12482b21 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/session/generated_media_manifest.go b/pkg/session/generated_media_manifest.go index 23d303b91e..ebdf73b011 100644 --- a/pkg/session/generated_media_manifest.go +++ b/pkg/session/generated_media_manifest.go @@ -10,6 +10,7 @@ import ( "time" "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/concurrent" ) var ( @@ -18,6 +19,8 @@ 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 // pkg/workspacemedia write result for its root kind (empty, traversal, // NUL, wrong absolute/relative shape, ...). @@ -74,6 +77,14 @@ type GeneratedMediaManifest interface { 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 { @@ -153,17 +164,74 @@ 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 { diff --git a/pkg/session/generated_media_manifest_test.go b/pkg/session/generated_media_manifest_test.go index 6bff5ae4c7..fe4702cb0c 100644 --- a/pkg/session/generated_media_manifest_test.go +++ b/pkg/session/generated_media_manifest_test.go @@ -27,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) { @@ -212,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 17682a651d..5d35296c6b 100644 --- a/pkg/session/migrations.go +++ b/pkg/session/migrations.go @@ -465,7 +465,21 @@ func getAllMigrations() []Migration { }, { ID: 29, - Name: "029_add_root_kind_to_generated_media_manifest", + 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 ee075528ac..5e0cbaa738 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 = "3d7db51380ef6b4ef5958a85eee1be279a6a96f7a11bf77b1b68f3a821218d9d" + 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 d949497bc0..18abd20c09 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).