Fix intermediate state where newly gen frontmatter wasn't saved yet

pkg/document/document_test.go

@@ -88,9 +88,11 @@ First paragraph
 
 		frontmatter, err := doc.Frontmatter()
 		require.NoError(t, err)
-		marshaledFrontmatter, err := frontmatter.Marshal(testIdentityResolver.DocumentEnabled())
+		var docID bytes.Buffer
+		marshaledFrontmatter, err := frontmatter.Marshal(testIdentityResolver.DocumentEnabled(), &docID)
 		require.NoError(t, err)
 		assert.Regexp(t, `---\nkey: value\nrunme:\n  id: .*\n  version: v(?:[3-9]\d*|2\.\d+\.\d+|2\.\d+|\d+)\n---`, string(marshaledFrontmatter))
+		assert.Len(t, docID.String(), 26)
 	})
 
 	t.Run("Format", func(t *testing.T) {

pkg/document/editor/editor.go

@@ -38,15 +38,22 @@ func Deserialize(data []byte, identityResolver *identity.IdentityResolver) (*Not
 		},
 	}
 
+	var docID bytes.Buffer
+	raw, err := frontmatter.Marshal(identityResolver.DocumentEnabled(), &docID)
 	// Additionally, put raw frontmatter in notebook's metadata.
 	// TODO(adamb): handle the error.
-	if raw, err := frontmatter.Marshal(identityResolver.DocumentEnabled()); err == nil && len(raw) > 0 {
+	if err == nil && len(raw) > 0 {
 		notebook.Metadata[PrefixAttributeName(InternalAttributePrefix, FrontmatterKey)] = string(raw)
 	}
-
-	// Store internal ephemeral document ID if the document lifecycle ID is disabled.
-	if !identityResolver.DocumentEnabled() {
-		notebook.Metadata[PrefixAttributeName(InternalAttributePrefix, DocumentID)] = identityResolver.EphemeralDocumentID()
+	// Store document ID in metadata to bridge state where frontmatter with identity wasn't saved yet.
+	if err == nil && docID.Len() > 0 {
+		notebook.Metadata[PrefixAttributeName(InternalAttributePrefix, DocumentID)] = docID.String()
+	}
+	// Include new frontmatter in notebook to bridge state where a new file/doc was created but not serialized yet.
+	if notebook.Frontmatter == nil {
+		if fm, err := document.ParseFrontmatter(raw); err == nil {
+			notebook.Frontmatter = fm
+		}
 	}
 
 	return notebook, nil
@@ -83,7 +90,8 @@ func Serialize(notebook *Notebook, outputMetadata *document.RunmeMetadata) ([]by
 	if frontmatter != nil {
 		// if the deserializer didn't add the ID first, it means it's not required
 		requireIdentity := !frontmatter.Runme.IsEmpty() && frontmatter.Runme.ID != ""
-		raw, err = frontmatter.Marshal(requireIdentity)
+		var internalDocID bytes.Buffer
+		raw, err = frontmatter.Marshal(requireIdentity, &internalDocID)
 		if err != nil {
 			return nil, err
 		}

pkg/document/editor/editorservice/service_test.go

@@ -93,6 +93,7 @@ func Test_IdentityUnspecified(t *testing.T) {
 
 		dResp, err := deserialize(client, tt.content, identity)
 		assert.NoError(t, err)
+		assert.Contains(t, dResp.Notebook.Metadata, "runme.dev/id")
 
 		rawFrontmatter, ok := dResp.Notebook.Metadata["runme.dev/frontmatter"]
 		if tt.hasExtraFrontmatter {
@@ -139,7 +140,7 @@ func Test_IdentityAll(t *testing.T) {
 		rawFrontmatter, ok := dResp.Notebook.Metadata["runme.dev/frontmatter"]
 		assert.True(t, ok)
 
-		assert.Len(t, dResp.Notebook.Metadata, 2)
+		assert.Len(t, dResp.Notebook.Metadata, 3)
 
 		if tt.hasExtraFrontmatter {
 			assert.Contains(t, rawFrontmatter, "prop: value")
@@ -180,7 +181,7 @@ func Test_IdentityDocument(t *testing.T) {
 		rawFrontmatter, ok := dResp.Notebook.Metadata["runme.dev/frontmatter"]
 		assert.True(t, ok)
 
-		assert.Len(t, dResp.Notebook.Metadata, 2)
+		assert.Len(t, dResp.Notebook.Metadata, 3)
 
 		if tt.hasExtraFrontmatter {
 			assert.Contains(t, rawFrontmatter, "prop: value")
@@ -217,6 +218,7 @@ func Test_IdentityCell(t *testing.T) {
 	for _, tt := range tests {
 		dResp, err := deserialize(client, tt.content, identity)
 		assert.NoError(t, err)
+		assert.Contains(t, dResp.Notebook.Metadata, "runme.dev/id")
 
 		rawFrontmatter, ok := dResp.Notebook.Metadata["runme.dev/frontmatter"]
 
@@ -268,7 +270,7 @@ func Test_RunmelessFrontmatter(t *testing.T) {
 	rawFrontmatter, ok := dResp.Notebook.Metadata["runme.dev/frontmatter"]
 
 	assert.True(t, ok)
-	assert.Len(t, dResp.Notebook.Metadata, 3)
+	assert.Len(t, dResp.Notebook.Metadata, 2)
 	assert.Contains(t, rawFrontmatter, "prop: value\n")
 	assert.NotContains(t, rawFrontmatter, "id: \"123\"\n")
 	assert.NotRegexp(t, versionRegex, rawFrontmatter, "Wrong version")
@@ -290,6 +292,63 @@ func Test_RunmelessFrontmatter(t *testing.T) {
 	assert.Contains(t, content, "```js {\"id\":\""+testMockID+"\"}\n")
 }
 
+func Test_NewFile_EmptyDocument_WithIdentityAll(t *testing.T) {
+	doc := ""
+
+	identity := parserv1.RunmeIdentity_RUNME_IDENTITY_ALL
+
+	dResp, err := deserialize(client, doc, identity)
+	assert.NoError(t, err)
+
+	rawFrontmatter, ok := dResp.Notebook.Metadata["runme.dev/frontmatter"]
+	assert.True(t, ok)
+	assert.Len(t, dResp.Notebook.Metadata, 3)
+	assert.Regexp(t, versionRegex, rawFrontmatter, "Wrong version")
+
+	assert.NotNil(t, dResp.Notebook.Frontmatter)
+	prasedRunmeID := dResp.Notebook.Frontmatter.Runme.Id
+	assert.Contains(t, rawFrontmatter, "id: "+prasedRunmeID+"\n")
+
+	sResp, err := serializeWithIdentityPersistence(client, dResp.Notebook, identity)
+	assert.NoError(t, err)
+
+	content := string(sResp.Result)
+
+	assert.Contains(t, content, "runme:\n")
+	assert.Contains(t, content, "id: "+prasedRunmeID+"\n")
+	assert.Contains(t, content, "version: v")
+	assert.Regexp(t, "^---\n", content)
+}
+
+func Test_EphemeralIdentity(t *testing.T) {
+	doc := strings.Join([]string{
+		"# Test identity integration with extension",
+		"```sh\ngh auth --help\n```",
+	}, "\n")
+
+	identity := parserv1.RunmeIdentity_RUNME_IDENTITY_UNSPECIFIED
+
+	dResp, err := deserialize(client, doc, identity)
+	assert.NoError(t, err)
+
+	assert.Len(t, dResp.Notebook.Metadata, 2)
+	assert.Len(t, dResp.Notebook.Cells, 2)
+	assert.NotContains(t, dResp.Notebook.Cells[1].Metadata, "id")
+	assert.Contains(t, dResp.Notebook.Cells[1].Metadata, "runme.dev/id")
+
+	sResp, err := serializeWithIdentityPersistence(client, dResp.Notebook, identity)
+	assert.NoError(t, err)
+
+	content := string(sResp.Result)
+
+	assert.NotContains(t, content, "{\"id\":\"")
+	assert.NotContains(t, dResp.Notebook.Metadata, "id")
+	assert.NotContains(t, content, "runme:\n")
+	assert.NotContains(t, content, "id: ")
+	assert.NotContains(t, content, "version: v")
+	assert.NotRegexp(t, "^---\n", content)
+}
+
 func Test_parserServiceServer_Ast(t *testing.T) {
 	t.Run("Metadata", func(t *testing.T) {
 		os.Setenv("RUNME_AST_METADATA", "true")
@@ -409,7 +468,9 @@ func deserialize(client parserv1.ParserServiceClient, content string, idt parser
 }
 
 func serializeWithIdentityPersistence(client parserv1.ParserServiceClient, notebook *parserv1.Notebook, idt parserv1.RunmeIdentity) (*parserv1.SerializeResponse, error) {
-	persistIdentityLikeExtension(notebook)
+	if idt == parserv1.RunmeIdentity_RUNME_IDENTITY_ALL || idt == parserv1.RunmeIdentity_RUNME_IDENTITY_CELL {
+		persistIdentityLikeExtension(notebook)
+	}
 	return client.Serialize(
 		context.Background(),
 		&parserv1.SerializeRequest{

pkg/document/frontmatter.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	stderrors "errors"
+	"io"
 
 	"github.com/pelletier/go-toml/v2"
 	"github.com/pkg/errors"
@@ -87,16 +88,30 @@ func newFrontmatter() *Frontmatter {
 	}
 }
 
-// Marshal returns a marshaled frontmatter including triple-dashed lines.
-// If the identity is required, but Frontmatter is nil, a new one is created.
-func (f *Frontmatter) Marshal(requireIdentity bool) ([]byte, error) {
+// Marshal returns a marshaled frontmatter including triple-dashed lines and a document ID.
+// If the identity is required, but Frontmatter is nil, a new one including ID is created.
+func (f *Frontmatter) Marshal(requireIdentity bool, docID io.StringWriter) ([]byte, error) {
 	if f == nil {
 		if !requireIdentity {
-			return nil, nil
+			// Return a new document ID for ephemeral usage.
+			_, err := docID.WriteString(ulid.GenerateID())
+			return nil, err
 		}
 		f = newFrontmatter()
 	}
-	return f.marshal(requireIdentity)
+
+	b, err := f.marshal(requireIdentity)
+	if err != nil {
+		return nil, err
+	}
+
+	if docID != nil && !f.Runme.IsEmpty() {
+		if _, err := docID.WriteString(f.Runme.ID); err != nil {
+			return nil, err
+		}
+	}
+
+	return b, nil
 }
 
 func (f *Frontmatter) marshal(requireIdentity bool) ([]byte, error) {

pkg/document/identity/resolver.go

@@ -71,16 +71,10 @@ func (ir *IdentityResolver) DocumentEnabled() bool {
 	return ir.documentIdentity
 }
 
-// EphemeralDocumentID returns a new document ID which is not persisted.
-func (ir *IdentityResolver) EphemeralDocumentID() string {
-	return ulid.GenerateID()
-}
-
 // GetCellID returns a cell ID and a boolean indicating if it's new or from attributes.
 func (ir *IdentityResolver) GetCellID(obj any, attributes map[string]string) (string, bool) {
-	if !ir.cellIdentity {
-		return "", false
-	}
+	// we used to early return here when !ir.cellIdentity, but we
+	// need cell IDs for the ephemeral case too, ie 'runme.dev/id'
 
 	// todo(sebastian): are invalid ulid's valid IDs?
 	// Check for a valid 'id' in attributes;

pkg/document/identity/resolver_test.go

@@ -46,7 +46,7 @@ func TestIdentityResolver(t *testing.T) {
 		assert.True(t, resolver.DocumentEnabled())
 	})
 
-	t.Run("GetCellID", func(t *testing.T) {
+	t.Run("GetCellID_IdentityRequired", func(t *testing.T) {
 		id := "01HF53Z4RCVPRANKFBZYMS72QW"
 		ulid.MockGenerator(id)
 		resolver := NewResolver(CellLifecycleIdentity)
@@ -58,8 +58,15 @@ func TestIdentityResolver(t *testing.T) {
 		assert.NotEmpty(t, id)
 	})
 
-	t.Run("EphemeralDocumentID", func(t *testing.T) {
-		resolver := NewResolver(DefaultLifecycleIdentity)
-		assert.Len(t, resolver.EphemeralDocumentID(), 26)
+	t.Run("GetCellID_IdentityNotRequired", func(t *testing.T) {
+		id := "01J1D6BDDD767E819NV8W7YQC2"
+		ulid.MockGenerator(id)
+		resolver := NewResolver(UnspecifiedLifecycleIdentity)
+		obj := struct{}{}
+		attributes := map[string]string{"id": id}
+		id, ok := resolver.GetCellID(obj, attributes)
+
+		assert.True(t, ok)
+		assert.NotEmpty(t, id)
 	})
 }