Adding toggle for checkpointing and use LatestMessage to get message

Author: AdityaHegde

proto/gen/rill/runtime/v1/api.pb.go

@@ -4928,6 +4928,11 @@ definitions:
         description: |-
           Optional path to the file that the user is currently viewing/editing.
           This helps the agent understand which file the user is referring to in their request.
+      enableCheckpointCommits:
+        type: boolean
+        description: |-
+          Optional toggle to enable checkpoint commit before any files are changed.
+          Can be used to undo the changes reverting to the commit.
     description: Context for prompts handled by the developer_agent.
   v1DirEntry:
     type: object

proto/gen/rill/runtime/v1/api.pb.validate.go

@@ -1085,6 +1085,9 @@ message DeveloperAgentContext {
   // Optional path to the file that the user is currently viewing/editing.
   // This helps the agent understand which file the user is referring to in their request.
   string current_file_path = 2;
+  // Optional toggle to enable checkpoint commit before any files are changed.
+  // Can be used to undo the changes reverting to the commit.
+  bool enable_checkpoint_commits = 3;
 }
 
 // Request message for RuntimeService.ListConversations

proto/gen/rill/runtime/v1/runtime.swagger.yaml

@@ -20,9 +20,10 @@ type DeveloperAgent struct {
 var _ Tool[*DeveloperAgentArgs, *DeveloperAgentResult] = (*DeveloperAgent)(nil)
 
 type DeveloperAgentArgs struct {
-	Prompt          string `json:"prompt" jsonschema:"User prompt for the developer agent."`
-	InitProject     bool   `json:"init_project,omitempty" jsonschema:"Set to true if the prompt is provided as part of project initialization. Should only be used when the project directory contains an empty project (except for the boilerplate generated by UnpackEmpty)."`
-	CurrentFilePath string `json:"current_file_path,omitempty" jsonschema:"Optional path to the file that the user is currently viewing/editing."`
+	Prompt                  string `json:"prompt" jsonschema:"User prompt for the developer agent."`
+	InitProject             bool   `json:"init_project,omitempty" jsonschema:"Set to true if the prompt is provided as part of project initialization. Should only be used when the project directory contains an empty project (except for the boilerplate generated by UnpackEmpty)."`
+	CurrentFilePath         string `json:"current_file_path,omitempty" jsonschema:"Optional path to the file that the user is currently viewing/editing."`
+	EnableCheckpointCommits bool   `json:"enable_checkpoint_commits,omitempty" jsonschema:"If true, the developer agent will automatically commit checkpoints before any write is triggered."`
 }
 
 type DeveloperAgentResult struct {

proto/rill/runtime/v1/api.proto

@@ -10,6 +10,7 @@ import (
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 	"github.com/pmezard/go-difflib/difflib"
 	"github.com/rilldata/rill/runtime"
+	"github.com/rilldata/rill/runtime/pkg/mapstructureutil"
 )
 
 const WriteFileName = "write_file"
@@ -167,22 +168,45 @@ func (t *WriteFile) reconcileAndGetStatus(ctx context.Context, path string) (res
 
 // maybeCreateCheckpoint creates a checkpoint if this is the 1st write file message in the current message chain.
 func (t *WriteFile) maybeCreateCheckpoint(ctx context.Context, s *Session, path string) (string, error) {
-	// Find a write file message in the current message chain.
-	var msg *Message
-	for i := len(s.messages) - 1; i >= 0; i-- {
-		if s.messages[i].Tool == WriteFileName && s.messages[i].Type == MessageTypeResult {
-			msg = s.messages[i]
+	// Find the nearest developer agent call and make sure checkpointing is enabled.
+	nearestDevAgentCall, ok := s.LatestMessage(FilterByTool(DeveloperAgentName), FilterByType(MessageTypeCall))
+	if ok {
+		rawReq, err := s.UnmarshalMessageContent(nearestDevAgentCall)
+		if err != nil {
+			return "", err
 		}
-
-		if s.messages[i].Role == RoleUser {
-			break
+		var req DeveloperAgentArgs
+		err = mapstructureutil.WeakDecode(rawReq, &req)
+		if err != nil {
+			return "", err
+		}
+		if !req.EnableCheckpointCommits {
+			return "", nil
 		}
 	}
 
-	// If there is already a write file message then we dont need to create a checkpoint.
-	if msg != nil {
-		return "", nil
+	// Find a write file message in the current message chain.
+	previousWriteResult, ok := s.LatestMessage(FilterByTool(WriteFileName), FilterByType(MessageTypeResult))
+	// If there is already a write file message then make sure it has a checkpoint commit.
+	if ok {
+		nearestRoot, ok := s.LatestMessage(FilterByRoot())
+		// If the nearest root message is before the previous write file message, then there hasn't been a commit yet for this loop.
+		if !ok || nearestRoot.Time.Before(previousWriteResult.Time) {
+			rawRes, err := s.UnmarshalMessageContent(previousWriteResult)
+			if err != nil {
+				return "", err
+			}
+			var res WriteFileResult
+			err = mapstructureutil.WeakDecode(rawRes, &res)
+			if err != nil {
+				return "", err
+			}
+			if res.CheckpointCommitHash != "" {
+				return res.CheckpointCommitHash, nil
+			}
+		}
 	}
+	// Else, create a checkpoint commit.
 
 	repo, release, err := t.Runtime.Repo(ctx, s.InstanceID())
 	if err != nil {

runtime/ai/developer_agent.go

@@ -258,8 +258,9 @@ func (s *Server) Complete(ctx context.Context, req *runtimev1.CompleteRequest) (
 	var developerAgentArgs *ai.DeveloperAgentArgs
 	if req.DeveloperAgentContext != nil {
 		developerAgentArgs = &ai.DeveloperAgentArgs{
-			InitProject:     req.DeveloperAgentContext.InitProject,
-			CurrentFilePath: req.DeveloperAgentContext.CurrentFilePath,
+			InitProject:             req.DeveloperAgentContext.InitProject,
+			CurrentFilePath:         req.DeveloperAgentContext.CurrentFilePath,
+			EnableCheckpointCommits: req.DeveloperAgentContext.EnableCheckpointCommits,
 		}
 	}
 
@@ -379,8 +380,9 @@ func (s *Server) CompleteStreaming(req *runtimev1.CompleteStreamingRequest, stre
 	var developerAgentArgs *ai.DeveloperAgentArgs
 	if req.DeveloperAgentContext != nil {
 		developerAgentArgs = &ai.DeveloperAgentArgs{
-			InitProject:     req.DeveloperAgentContext.InitProject,
-			CurrentFilePath: req.DeveloperAgentContext.CurrentFilePath,
+			InitProject:             req.DeveloperAgentContext.InitProject,
+			CurrentFilePath:         req.DeveloperAgentContext.CurrentFilePath,
+			EnableCheckpointCommits: req.DeveloperAgentContext.EnableCheckpointCommits,
 		}
 	}
 

runtime/ai/file_write.go

@@ -21,11 +21,17 @@ export const developerChatConfig = {
 function getActiveFileContext(): Readable<Partial<RuntimeServiceCompleteBody>> {
   return derived(page, (pageState) => {
     const filePath = pageState.params?.file;
-    if (!filePath) return {} satisfies Partial<RuntimeServiceCompleteBody>;
+    // TODO: toggle for this checkpoints
+    if (!filePath) {
+      return {
+        developerAgentContext: { enableCheckpointCommits: true },
+      } satisfies Partial<RuntimeServiceCompleteBody>;
+    }
 
     return {
       developerAgentContext: {
         currentFilePath: filePath,
+        enableCheckpointCommits: true,
       },
     } satisfies Partial<RuntimeServiceCompleteBody>;
   });

runtime/server/chat.go

@@ -4590,6 +4590,14 @@ export class DeveloperAgentContext extends Message$1<DeveloperAgentContext> {
    */
   currentFilePath = "";
 
+  /**
+   * Optional toggle to enable checkpoint commit before any files are changed.
+   * Can be used to undo the changes reverting to the commit.
+   *
+   * @generated from field: bool enable_checkpoint_commits = 3;
+   */
+  enableCheckpointCommits = false;
+
   constructor(data?: PartialMessage<DeveloperAgentContext>) {
     super();
     proto3.util.initPartial(data, this);
@@ -4600,6 +4608,7 @@ export class DeveloperAgentContext extends Message$1<DeveloperAgentContext> {
   static readonly fields: FieldList = proto3.util.newFieldList(() => [
     { no: 1, name: "init_project", kind: "scalar", T: 8 /* ScalarType.BOOL */ },
     { no: 2, name: "current_file_path", kind: "scalar", T: 9 /* ScalarType.STRING */ },
+    { no: 3, name: "enable_checkpoint_commits", kind: "scalar", T: 8 /* ScalarType.BOOL */ },
   ]);
 
   static fromBinary(bytes: Uint8Array, options?: Partial<BinaryReadOptions>): DeveloperAgentContext {

web-common/src/features/editor/chat-context.ts

@@ -823,6 +823,9 @@ This should only be used when the project directory contains an empty project (e
   /** Optional path to the file that the user is currently viewing/editing.
 This helps the agent understand which file the user is referring to in their request. */
   currentFilePath?: string;
+  /** Optional toggle to enable checkpoint commit before any files are changed.
+Can be used to undo the changes reverting to the commit. */
+  enableCheckpointCommits?: boolean;
 }
 
 export interface V1DirEntry {