diff --git a/forge-core/llm/content_parts_test.go b/forge-core/llm/content_parts_test.go new file mode 100644 index 00000000..ff826bb5 --- /dev/null +++ b/forge-core/llm/content_parts_test.go @@ -0,0 +1,90 @@ +package llm + +import ( + "encoding/base64" + "encoding/json" + "strings" + "testing" +) + +// TestChatMessage_TextOnlyWireIsUnchanged pins the back-compat invariant: a +// text-only ChatMessage (no Parts) marshals byte-identically to the pre-#255 +// shape — no "parts" key — so existing sessions, prompt-cache prefixes, and +// provider requests are untouched. +func TestChatMessage_TextOnlyWireIsUnchanged(t *testing.T) { + b, err := json.Marshal(ChatMessage{Role: RoleUser, Content: "hello"}) + if err != nil { + t.Fatal(err) + } + got := string(b) + if want := `{"role":"user","content":"hello"}`; got != want { + t.Errorf("text-only wire = %s, want %s", got, want) + } + if strings.Contains(got, "parts") { + t.Errorf("empty Parts must be omitted from the wire; got %s", got) + } +} + +// TestMediaRef_BytesNeverSerialized is the #255 history-bloat guard: MediaRef.Bytes +// is json:"-", so marshaling a message with media never writes the inline bytes — +// only the URI reference survives. This is what keeps session history from +// re-persisting (and re-sending) base64 media every turn. +func TestMediaRef_BytesNeverSerialized(t *testing.T) { + msg := ChatMessage{ + Role: RoleUser, + Content: "look at this", + Parts: []ContentPart{ + NewTextContentPart("look at this"), + NewMediaContentPart(ContentPartImage, MediaRef{ + MimeType: "image/png", + URI: ".forge/files/inbound/photo.png", + Bytes: []byte("SUPER-SECRET-RAW-IMAGE-BYTES"), + }), + }, + } + b, err := json.Marshal(msg) + if err != nil { + t.Fatal(err) + } + got := string(b) + if strings.Contains(got, "SUPER-SECRET-RAW-IMAGE-BYTES") { + t.Errorf("inline Bytes leaked into the serialized message:\n%s", got) + } + if b64 := base64.StdEncoding.EncodeToString([]byte("SUPER-SECRET-RAW-IMAGE-BYTES")); strings.Contains(got, b64) { + t.Errorf("inline Bytes leaked as base64 into the serialized message:\n%s", got) + } + if !strings.Contains(got, ".forge/files/inbound/photo.png") { + t.Errorf("URI reference should survive serialization; got %s", got) + } + + // Round-trip: unmarshal keeps the URI, drops the bytes (rehydration is a + // separate step). + var back ChatMessage + if err := json.Unmarshal(b, &back); err != nil { + t.Fatal(err) + } + if !back.HasMedia() { + t.Fatal("round-tripped message should still report HasMedia") + } + if got := back.Parts[1].Media.URI; got != ".forge/files/inbound/photo.png" { + t.Errorf("round-tripped URI = %q", got) + } + if len(back.Parts[1].Media.Bytes) != 0 { + t.Errorf("round-tripped Bytes must be empty (json:\"-\"); got %d bytes", len(back.Parts[1].Media.Bytes)) + } +} + +// TestChatMessage_HasMedia covers the media predicate. +func TestChatMessage_HasMedia(t *testing.T) { + if (ChatMessage{Content: "x"}).HasMedia() { + t.Error("text-only message must not report media") + } + textParts := ChatMessage{Parts: []ContentPart{NewTextContentPart("x")}} + if textParts.HasMedia() { + t.Error("text-only parts must not report media") + } + withImage := ChatMessage{Parts: []ContentPart{NewMediaContentPart(ContentPartImage, MediaRef{MimeType: "image/png"})}} + if !withImage.HasMedia() { + t.Error("message with an image part must report media") + } +} diff --git a/forge-core/llm/types.go b/forge-core/llm/types.go index 212258c0..6f27ef15 100644 --- a/forge-core/llm/types.go +++ b/forge-core/llm/types.go @@ -14,12 +14,71 @@ const ( ) // ChatMessage represents a single message in a chat conversation. +// +// Content is the flattened text-of-record and stays authoritative for +// compression, truncation, dedup, and the security scanners. Parts is the +// additive multimodal representation (#255): when a message carries media it +// holds text + image/document parts; it is empty (and omitted on the wire) for +// the common text-only case, so existing text-only behavior is unchanged. type ChatMessage struct { - Role string `json:"role"` - Content string `json:"content,omitempty"` - ToolCalls []ToolCall `json:"tool_calls,omitempty"` - ToolCallID string `json:"tool_call_id,omitempty"` - Name string `json:"name,omitempty"` + Role string `json:"role"` + Content string `json:"content,omitempty"` + Parts []ContentPart `json:"parts,omitempty"` + ToolCalls []ToolCall `json:"tool_calls,omitempty"` + ToolCallID string `json:"tool_call_id,omitempty"` + Name string `json:"name,omitempty"` +} + +// ContentPart Type discriminators. +const ( + ContentPartText = "text" + ContentPartImage = "image" + ContentPartDocument = "document" +) + +// ContentPart is one piece of a multimodal message's content — a text block or +// an image/document media block. Providers that support vision/documents +// serialize Parts into their native content-block format; when Parts is empty +// they fall back to the plain Content string (byte-identical to today). +type ContentPart struct { + Type string `json:"type"` // ContentPartText | ContentPartImage | ContentPartDocument + Text string `json:"text,omitempty"` // set when Type == ContentPartText + Media *MediaRef `json:"media,omitempty"` // set when Type == ContentPartImage | ContentPartDocument +} + +// MediaRef references the media carried by an image/document ContentPart. +// +// Bytes is the inline payload used to build a provider request, but it is +// deliberately NEVER serialized (json:"-"): persisted history (SessionData) +// stores only the URI reference, and the executor rehydrates Bytes from the URI +// per turn (RehydrateMedia). This is the #255 history-bloat guard — persisting +// inline base64 would rewrite the media into the session file on every turn and +// resend it to the model on every replay. +type MediaRef struct { + MimeType string `json:"mimeType,omitempty"` + URI string `json:"uri,omitempty"` // reference form persisted in history (e.g. a .forge/files path) + Bytes []byte `json:"-"` // inline payload; rehydrated from URI, never persisted +} + +// NewTextContentPart builds a text content part. +func NewTextContentPart(text string) ContentPart { + return ContentPart{Type: ContentPartText, Text: text} +} + +// NewMediaContentPart builds an image/document content part from a MediaRef. +// kind must be ContentPartImage or ContentPartDocument. +func NewMediaContentPart(kind string, media MediaRef) ContentPart { + return ContentPart{Type: kind, Media: &media} +} + +// HasMedia reports whether the message carries any image/document part. +func (m ChatMessage) HasMedia() bool { + for i := range m.Parts { + if m.Parts[i].Media != nil { + return true + } + } + return false } // ToolCall represents an LLM request to invoke a tool. diff --git a/forge-core/runtime/media_rehydrate.go b/forge-core/runtime/media_rehydrate.go new file mode 100644 index 00000000..08a84cb2 --- /dev/null +++ b/forge-core/runtime/media_rehydrate.go @@ -0,0 +1,77 @@ +package runtime + +import ( + "context" + "fmt" + "os" + "path/filepath" + "strings" + + "github.com/initializ/forge/forge-core/llm" +) + +// RehydrateMedia fills the inline Bytes of every image/document ContentPart in +// msgs from its persisted URI reference, reading each file confined to the +// context's files dir (WithFilesDir). It is the read side of the #255 +// history-reference design: history persists only URIs (llm.MediaRef.Bytes is +// json:"-" so it never bloats the session file or re-sends on replay), so the +// executor calls RehydrateMedia to reload the bytes before building a provider +// ChatRequest. +// +// Parts whose Bytes are already populated are left untouched (freshly ingested +// media not yet round-tripped through history). A URI that escapes the files +// dir, or that cannot be read, yields an error naming the offending part — +// callers decide whether that is fatal or degrades to a text reference. +func RehydrateMedia(ctx context.Context, msgs []llm.ChatMessage) error { + for mi := range msgs { + for pi := range msgs[mi].Parts { + media := msgs[mi].Parts[pi].Media + if media == nil || len(media.Bytes) > 0 || media.URI == "" { + continue + } + path, ok := confinedFilesPath(ctx, media.URI) + if !ok { + return fmt.Errorf("rehydrate media: uri %q escapes the files dir", media.URI) + } + // confinedFilesPath is lexical (Abs+Rel) — it stops ../ and + // absolute traversal but does not resolve symlinks, so a symlink + // INSIDE the files dir pointing outward would otherwise be followed + // by the read below. Resolve symlinks on both the target and the + // files dir (the dir itself may sit under a symlink, e.g. macOS + // /tmp -> /private/tmp) and re-confine before reading. This matters + // once inbound (untrusted) files are persisted here (#255 Phase 4). + realPath, err := resolveWithinFilesDir(ctx, path) + if err != nil { + return fmt.Errorf("rehydrate media %q: %w", media.URI, err) + } + b, err := os.ReadFile(realPath) + if err != nil { + return fmt.Errorf("rehydrate media %q: %w", media.URI, err) + } + media.Bytes = b + } + } + return nil +} + +// resolveWithinFilesDir resolves symlinks in path and confirms the real target +// still lives inside the (symlink-resolved) files dir. It complements the +// lexical confinedFilesPath: a symlink placed inside the files dir that points +// outward passes the lexical check but is caught here. Returns the resolved +// path to read, or an error if the target escapes the dir (or cannot be +// resolved — e.g. it does not exist). +func resolveWithinFilesDir(ctx context.Context, path string) (string, error) { + realPath, err := filepath.EvalSymlinks(path) + if err != nil { + return "", err + } + realDir, err := filepath.EvalSymlinks(FilesDirFromContext(ctx)) + if err != nil { + return "", err + } + rel, err := filepath.Rel(realDir, realPath) + if err != nil || rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + return "", fmt.Errorf("resolved path escapes the files dir") + } + return realPath, nil +} diff --git a/forge-core/runtime/media_rehydrate_test.go b/forge-core/runtime/media_rehydrate_test.go new file mode 100644 index 00000000..bde3dca0 --- /dev/null +++ b/forge-core/runtime/media_rehydrate_test.go @@ -0,0 +1,121 @@ +package runtime + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/initializ/forge/forge-core/llm" +) + +// TestRehydrateMedia_LoadsBytesFromURI is the read side of the #255 +// history-reference design: a persisted message carries only a URI (Bytes +// dropped by json:"-"); RehydrateMedia reloads the bytes from the files dir +// before the message is handed to a provider. +func TestRehydrateMedia_LoadsBytesFromURI(t *testing.T) { + dir := t.TempDir() + want := []byte("the-real-image-bytes") + uri := filepath.Join(dir, "photo.png") + if err := os.WriteFile(uri, want, 0o600); err != nil { + t.Fatal(err) + } + ctx := WithFilesDir(context.Background(), dir) + + msgs := []llm.ChatMessage{{ + Role: llm.RoleUser, + Parts: []llm.ContentPart{llm.NewMediaContentPart(llm.ContentPartImage, llm.MediaRef{MimeType: "image/png", URI: uri})}, + }} + if err := RehydrateMedia(ctx, msgs); err != nil { + t.Fatalf("RehydrateMedia: %v", err) + } + if got := msgs[0].Parts[0].Media.Bytes; string(got) != string(want) { + t.Errorf("rehydrated bytes = %q, want %q", got, want) + } +} + +// TestRehydrateMedia_AlreadyPopulatedIsUntouched: freshly ingested media that +// hasn't round-tripped through history keeps its bytes and is not re-read. +func TestRehydrateMedia_AlreadyPopulatedIsUntouched(t *testing.T) { + ctx := WithFilesDir(context.Background(), t.TempDir()) + msgs := []llm.ChatMessage{{ + Role: llm.RoleUser, + Parts: []llm.ContentPart{llm.NewMediaContentPart(llm.ContentPartImage, llm.MediaRef{MimeType: "image/png", URI: "/does/not/exist.png", Bytes: []byte("inline")})}, + }} + if err := RehydrateMedia(ctx, msgs); err != nil { + t.Fatalf("RehydrateMedia should skip already-populated bytes, got %v", err) + } + if got := string(msgs[0].Parts[0].Media.Bytes); got != "inline" { + t.Errorf("bytes = %q, want inline (unchanged)", got) + } +} + +// TestRehydrateMedia_RejectsPathEscape: a URI outside the files dir must be +// refused, not read — the same confinement fileArtifactFromToolResult uses. +func TestRehydrateMedia_RejectsPathEscape(t *testing.T) { + dir := t.TempDir() + // A secret sitting outside the files dir. + secret := filepath.Join(t.TempDir(), "secret.txt") + if err := os.WriteFile(secret, []byte("nope"), 0o600); err != nil { + t.Fatal(err) + } + ctx := WithFilesDir(context.Background(), dir) + msgs := []llm.ChatMessage{{ + Role: llm.RoleUser, + Parts: []llm.ContentPart{llm.NewMediaContentPart(llm.ContentPartImage, llm.MediaRef{URI: secret})}, + }} + if err := RehydrateMedia(ctx, msgs); err == nil { + t.Fatal("a URI escaping the files dir must be rejected") + } + if len(msgs[0].Parts[0].Media.Bytes) != 0 { + t.Error("bytes must not be populated from an out-of-confinement path") + } +} + +// TestRehydrateMedia_RejectsSymlinkEscape: a symlink placed INSIDE the files +// dir that points outward passes the lexical confinement but must be caught by +// the symlink resolution — the Phase-4 hardening from the #528 review. +func TestRehydrateMedia_RejectsSymlinkEscape(t *testing.T) { + dir := t.TempDir() + // A secret outside the files dir. + outside := filepath.Join(t.TempDir(), "secret.txt") + if err := os.WriteFile(outside, []byte("nope"), 0o600); err != nil { + t.Fatal(err) + } + // A symlink INSIDE the files dir pointing at it — lexically confined + // (the link path is under dir), but resolves outward. + link := filepath.Join(dir, "link.png") + if err := os.Symlink(outside, link); err != nil { + t.Skipf("symlinks unsupported on this platform: %v", err) + } + ctx := WithFilesDir(context.Background(), dir) + msgs := []llm.ChatMessage{{ + Role: llm.RoleUser, + Parts: []llm.ContentPart{llm.NewMediaContentPart(llm.ContentPartImage, llm.MediaRef{URI: link})}, + }} + if err := RehydrateMedia(ctx, msgs); err == nil { + t.Fatal("a symlink inside the files dir pointing outward must be rejected") + } + if len(msgs[0].Parts[0].Media.Bytes) != 0 { + t.Error("bytes must not be populated from a symlink escaping the files dir") + } +} + +// TestMemory_MediaCountsTowardBudget: media parts carry no Content chars but +// must charge the budget, so an image-heavy history still trims (#255). +func TestMemory_MediaCountsTowardBudget(t *testing.T) { + imagePart := llm.NewMediaContentPart(llm.ContentPartImage, llm.MediaRef{MimeType: "image/png", URI: "a.png"}) + + withMedia := NewMemory("", 0, "") + withMedia.Append(llm.ChatMessage{Role: llm.RoleUser, Content: "hi", Parts: []llm.ContentPart{imagePart}}) + textOnly := NewMemory("", 0, "") + textOnly.Append(llm.ChatMessage{Role: llm.RoleUser, Content: "hi"}) + + if withMedia.totalChars() <= textOnly.totalChars() { + t.Errorf("a message with a media part must charge more budget than text-only: media=%d text=%d", + withMedia.totalChars(), textOnly.totalChars()) + } + if got := withMedia.totalChars() - textOnly.totalChars(); got != mediaCharWeight { + t.Errorf("media budget delta = %d, want mediaCharWeight (%d)", got, mediaCharWeight) + } +} diff --git a/forge-core/runtime/memory.go b/forge-core/runtime/memory.go index cd2aa34d..7fbe519b 100644 --- a/forge-core/runtime/memory.go +++ b/forge-core/runtime/memory.go @@ -84,6 +84,13 @@ func NewMemory(systemPrompt string, maxChars int, model string) *Memory { // maxMessageChars is the per-message size cap (defense in depth). const maxMessageChars = 50_000 +// mediaCharWeight is the char-equivalent budget charged per image/document +// media part in totalChars (#255). Media carries no Content chars but consumes +// real model token budget; this conservative flat proxy keeps an image-heavy +// conversation from silently exceeding the true budget. It is a budget estimate +// only — billed truth is the provider's token accounting. +const mediaCharWeight = 4_000 + // Append adds a message to the conversation history and trims if over budget. // Individual messages exceeding maxMessageChars are truncated as a safety net. func (m *Memory) Append(msg llm.ChatMessage) { @@ -347,6 +354,16 @@ func (m *Memory) totalChars() int { for _, tc := range msg.ToolCalls { total += len(tc.Function.Name) + len(tc.Function.Arguments) } + // Media parts consume real token budget (an image is ~hundreds– + // thousands of tokens) but carry no Content chars — without a weight + // here a conversation of images would blow the true budget while trim + // thinks it is tiny (#255). Text parts are already reflected in + // msg.Content (the text-of-record), so only media parts are added. + for i := range msg.Parts { + if msg.Parts[i].Media != nil { + total += mediaCharWeight + } + } } return total }