feat: thread forking

.github/workflows/e2e-tests.yml

@@ -118,6 +118,13 @@ jobs:
       mock: true
     secrets: inherit
 
+  e2e-test-thread-forking:
+    uses: ./.github/workflows/reusable-e2e-test.yml
+    with:
+      test-name: "thread-forking"
+      mock: true
+    secrets: inherit
+
   e2e-tests-complete:
     needs:
       [

proto/tim-api/tim/api/thread/v1alpha1/thread_service.proto

@@ -8,7 +8,6 @@ import "google/api/annotations.proto";
 import "google/api/client.proto";
 import "google/api/field_behavior.proto";
 import "google/protobuf/empty.proto";
-import "google/protobuf/timestamp.proto";
 import "tim/api/thread/v1alpha1/thread_types.proto";
 import "tim/api/tool/v1alpha1/tool_types.proto";
 
@@ -37,13 +36,13 @@ service ThreadService {
     option (google.api.method_signature) = "parent,thread,initial_message_text";
   }
 
-  // Fork an existing thread to create a new thread
+  // Fork a thread at a specific message (inclusive)
   rpc ForkThread(ForkThreadRequest) returns (ForkThreadResponse) {
     option (google.api.http) = {
-      post: "/v1alpha1/{parent=orgs/*/users/*/threads/*}:fork"
+      post: "/v1alpha1/{message=orgs/*/users/*/threads/*/messages/*}:fork"
       body: "*"
     };
-    option (google.api.method_signature) = "parent,thread";
+    option (google.api.method_signature) = "message";
   }
 
   // Update mutable fields on a thread
@@ -171,27 +170,24 @@ message CreateThreadRequest {
   ];
 }
 
-// CreateThreadRequest is used to create a new thread, which must be unique within
-// the parent thread if one is provided.
+// ForkThreadRequest is used to fork a thread at a specific message
 message ForkThreadRequest {
-  // The resource path of the thread to fork from
-  string parent = 1 [
+  // The resource path of the user message to fork before (non-inclusive).
+  // The forked thread will contain all messages BEFORE this message.
+  // Must be a user message and cannot be the first message (index 0).
+  string message = 1 [
     (google.api.field_behavior) = REQUIRED,
-    (aep.api.field_info).resource_reference = "tim.settlerlabs.com/thread",
-    (buf.validate.field).string.pattern = "^orgs/[a-fA-F0-9-]{36}/users/[a-fA-F0-9-]{36}/threads/[a-fA-F0-9-]{36}$"
+    (aep.api.field_info).resource_reference = "tim.settlerlabs.com/llm-message",
+    (buf.validate.field).string.pattern = "^orgs/[a-fA-F0-9-]{36}/users/[a-fA-F0-9-]{36}/threads/[a-fA-F0-9-]{36}/messages/[a-fA-F0-9-]{36}$"
   ];
 
-  // The timestamp from the parent thread to fork from, all events on the parent
-  // thread before this time will be included in the new thread.
-  google.protobuf.Timestamp before_time = 2 [
-    (buf.validate.field).required = true,
-    (google.api.field_behavior) = REQUIRED
-  ];
+  // Optional title for the forked thread. If not provided, will be auto-generated.
+  optional string title = 2 [(buf.validate.field).string.max_len = 255];
 }
 
 // ForkThreadResponse is the response for forking a thread
 message ForkThreadResponse {
-  // The newly created thread
+  // The newly created forked thread
   Thread thread = 1 [
     (buf.validate.field).required = true,
     (google.api.field_behavior) = REQUIRED
@@ -330,6 +326,12 @@ message EditThreadMessageRequest {
   // If true: restores files from checkpoint and returns file restoration data.
   // If false: saves without restoring (new checkpoint replaces old one).
   optional bool restore = 3;
+
+  // If true, creates a new forked thread instead of modifying the original thread.
+  // The forked thread will include all messages before the edited message,
+  // and the edited content will be added as a new message in the forked thread.
+  // The original thread remains unchanged.
+  optional bool create_fork = 4 [(google.api.field_behavior) = OPTIONAL];
 }
 
 // EditThreadMessageResponse contains the edited message and any file restoration data
@@ -343,6 +345,9 @@ message EditThreadMessageResponse {
   // Files to restore from checkpoint (if applicable)
   // The client should restore these files before continuing
   repeated FileRestoration file_restorations = 2;
+
+  // If create_fork was true in the request, this contains the newly created forked thread
+  optional Thread forked_thread = 3;
 }
 
 // ConfigureThreadWorkingDirectoryRequest configures the working directory for checkpoint creation

proto/tim-api/tim/api/thread/v1alpha1/thread_types.proto

@@ -45,6 +45,12 @@ message Thread {
       (google.api.field_behavior) = IMMUTABLE,
       (buf.validate.field).required = true
     ];
+    // The message at which this thread was forked (optional for backward compatibility)
+    // If set, this is the exact message that served as the fork point.
+    string fork_message_uid = 3 [
+      (google.api.field_behavior) = OUTPUT_ONLY,
+      (google.api.field_info).format = UUID4
+    ];
   }
 
   // If this thread was forked from another thread, this is the ID of that thread

shared/apiclient/convenience.go

@@ -3,15 +3,13 @@ package apiclient
 import (
 	"context"
 	"fmt"
-	"time"
 
 	"connectrpc.com/connect"
 	orgv1alpha1 "github.com/Greybox-Labs/tim/tim-proto/gen/tim/api/org/v1alpha1"
 	personav1alpha1 "github.com/Greybox-Labs/tim/tim-proto/gen/tim/api/persona/v1alpha1"
 	threadv1alpha1 "github.com/Greybox-Labs/tim/tim-proto/gen/tim/api/thread/v1alpha1"
 	todov1alpha1 "github.com/Greybox-Labs/tim/tim-proto/gen/tim/api/todo/v1alpha1"
 	userv1alpha1 "github.com/Greybox-Labs/tim/tim-proto/gen/tim/api/user/v1alpha1"
-	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 // Common pagination options
@@ -179,19 +177,6 @@ func (c *Client) CreateThreadWithPrompt(ctx context.Context, orgID, userID strin
 	return resp.Msg, nil
 }
 
-// ForkThreadAtTime forks a thread at a specific timestamp
-func (c *Client) ForkThreadAtTime(ctx context.Context, threadPath string, beforeTime time.Time) (*threadv1alpha1.Thread, error) {
-	req := &threadv1alpha1.ForkThreadRequest{
-		Parent:     threadPath,
-		BeforeTime: timestamppb.New(beforeTime),
-	}
-	resp, err := c.Thread.ForkThread(ctx, connect.NewRequest(req))
-	if err != nil {
-		return nil, c.handleError(err)
-	}
-	return resp.Msg.Thread, nil
-}
-
 // ListThreadsWithDefaults lists threads with default pagination
 func (c *Client) ListThreadsWithDefaults(ctx context.Context, orgID, userID string) (*threadv1alpha1.ListThreadsResponse, error) {
 	return c.ListThreadsWithPagination(ctx, orgID, userID, PaginationOptions{

tests/system/framework/api_client.py

@@ -233,7 +233,7 @@ def submit_user_message(self, thread_path: str, text: str) -> dict[str, Any]:
         return response.json()
 
     def edit_thread_message(
-        self, message_path: str, content: str, restore: bool = False
+        self, message_path: str, content: str, restore: bool = False, create_fork: bool = False
     ) -> dict[str, Any]:
         """Edit a thread message.
 
@@ -243,14 +243,43 @@ def edit_thread_message(
             restore: If true, restores files from checkpoint and returns file restoration data.
                     If false, saves without restoring (new checkpoint replaces old one).
                     If not set, returns error when checkpoint exists (to prompt user).
+            create_fork: If true, creates a new forked thread instead of modifying the original thread.
 
         Returns:
-            The updated message with optional file_restorations field
+            The updated message with optional file_restorations and forked_thread fields
         """
+        data: dict[str, Any] = {"content": content, "restore": restore}
+        if create_fork:
+            data["create_fork"] = create_fork
+
         response = self.client.post(
             self._url(f"{message_path}:edit"),
             headers=self._default_headers,
-            json={"content": content, "restore": restore},
+            json=data,
+        )
+        response.raise_for_status()
+        return response.json()
+
+    def fork_thread_from_message(
+        self, message_path: str, title: str | None = None
+    ) -> dict[str, Any]:
+        """Fork a thread at a specific message.
+
+        Args:
+            message_path: Path to the message to fork from (inclusive)
+            title: Optional title for the forked thread
+
+        Returns:
+            The forked thread response
+        """
+        data: dict[str, Any] = {}
+        if title:
+            data["title"] = title
+
+        response = self.client.post(
+            self._url(f"{message_path}:fork"),
+            headers=self._default_headers,
+            json=data,
         )
         response.raise_for_status()
         return response.json()

tests/system/framework/models/__init__.py

@@ -41,6 +41,7 @@
     LlmMessageContent,
     LlmMessageRole,
     MessageList,
+    ParentThreadId,
     Thinking,
     Thread,
     ThreadList,
@@ -73,6 +74,7 @@
     "LlmMessageContent",
     "LlmMessageRole",
     "MessageList",
+    "ParentThreadId",
     "Thinking",
     "Thread",
     "ThreadList",

tests/system/framework/models/thread.py

@@ -27,11 +27,22 @@ class LlmMessageRole(str, Enum):
     ASSISTANT = "LLM_MESSAGE_ROLE_ASSISTANT"
 
 
+class ParentThreadId(BaseModel):
+    """Parent thread reference for forked threads."""
+
+    path: str
+    before_time: str | None = Field(None, alias="beforeTime")
+    fork_message_uid: str | None = Field(None, alias="forkMessageUid")
+
+    model_config = ConfigDict(populate_by_alias=True)
+
+
 class Thread(BaseModel):
     """Thread resource."""
 
     path: str
     display_name: str | None = Field(None, alias="displayName")
+    parent_thread_id: ParentThreadId | None = Field(None, alias="parentThreadId")
     persona_uid: str | None = Field(None, alias="personaUid")
     environment_uid: str | None = Field(None, alias="environmentUid")
     create_time: str | None = Field(None, alias="createTime")