Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions internal/repl/repl.go
Original file line number Diff line number Diff line change
Expand Up @@ -59,14 +59,23 @@ func Run(ag *agent.Agent, cliAgent agent.CLIAgent, quietMode bool, version strin
if cfg == nil || cfg.CloudSync == nil || !*cfg.CloudSync {
return
}
tracker.Start(api, projectID, ag.Cfg.Model)
tracker.StartWithHistory(api, projectID, ag.Cfg.Model, len(ag.History), ag.Usage.TotalTokens())
}
completeCloudSession := func() {
cfg := ag.AppConfig
if cfg == nil || cfg.CloudSync == nil || !*cfg.CloudSync {
return
}
tracker.Complete(ag.Cfg.Context.API, ag.Usage.TotalTokens(), session.SummaryFor(ag.History), ag.History)
tracker.CompleteCurrent(ag.Cfg.Context.API, ag.Usage.TotalTokens(), ag.History)
}
switchCloudProject := func(projectID int) bool {
cfg := ag.AppConfig
if cfg == nil || cfg.CloudSync == nil || !*cfg.CloudSync {
return false
}
return tracker.SwitchProject(
ag.Cfg.Context.API, projectID, ag.Cfg.Model, ag.Usage.TotalTokens(), ag.History,
)
}

// Graceful interrupt handling
Expand Down Expand Up @@ -214,8 +223,12 @@ func Run(ag *agent.Agent, cliAgent agent.CLIAgent, quietMode bool, version strin
id := strings.TrimPrefix(input, "/project ")
var pid int
if _, err := fmt.Sscanf(id, "%d", &pid); err == nil {
movedCloudSession := switchCloudProject(pid)
ag.Cfg.Context.ProjectID = pid
term.PrintSystem(fmt.Sprintf("Project set to #%d", pid))
if movedCloudSession {
term.PrintSystem(fmt.Sprintf("Cloud session moved to project #%d", pid))
}
} else {
term.PrintError("Invalid project ID")
}
Expand Down
65 changes: 64 additions & 1 deletion internal/session/cloud_session.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,18 +53,81 @@ func ApplyCloudSyncChoice(cfg *api.Config, line string) bool {
// CloudSessionTracker manages the lifecycle of a cloud-tracked agent session.
// Zero value is ready to use — no initialisation needed.
type CloudSessionTracker struct {
cloudID string
cloudID string
projectID int
historyStart int
tokenStart int
}

// Start opens a cloud session the first time it is called with a valid API
// client and non-zero project ID. Subsequent calls are no-ops (idempotent).
func (t *CloudSessionTracker) Start(client *api.APIClient, projectID int, model string) {
t.StartWithHistory(client, projectID, model, 0, 0)
}

// StartWithHistory opens a cloud session and records the point in the local
// conversation at which it began. This lets a single local REPL conversation
// move between projects without copying the earlier project's messages into
// the next cloud session.
func (t *CloudSessionTracker) StartWithHistory(
client *api.APIClient, projectID int, model string, historyLength int, totalTokens int,
) {
if client == nil || projectID == 0 || t.cloudID != "" {
return
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
t.cloudID = client.CreateAgentSession(ctx, projectID, model)
if t.cloudID != "" {
t.projectID = projectID
t.historyStart = historyLength
t.tokenStart = totalTokens
}
}

// SwitchProject finalizes the current cloud session and starts another one for
// projectID. Local history remains intact, but only messages entered after the
// switch belong to the new cloud session. It returns true when a tracked
// session was actually moved to a different project.
func (t *CloudSessionTracker) SwitchProject(
client *api.APIClient, projectID int, model string, totalTokens int, messages []api.Message,
) bool {
if client == nil {
return false
}
if t.cloudID == "" {
t.StartWithHistory(client, projectID, model, len(messages), totalTokens)
return false
}
if t.projectID == projectID {
return false
}

t.CompleteCurrent(client, totalTokens, messages)
t.cloudID = ""
t.projectID = 0
t.historyStart = 0
t.tokenStart = 0
t.StartWithHistory(client, projectID, model, len(messages), totalTokens)
return true
}

// CompleteCurrent finalizes the active cloud session using only the portion
// of the local history that belongs to its current project.
func (t *CloudSessionTracker) CompleteCurrent(client *api.APIClient, totalTokens int, messages []api.Message) {
start := t.historyStart
if start < 0 {
start = 0
}
Comment on lines +117 to +121

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 LOGICGROUNDED Out-of-range historyStart uploads full replacement history to wrong cloud session

When historyStart exceeds the current len(messages) (for example, after /load replaces ag.History with a shorter saved session), CompleteCurrent resets start to 0 and uploads messages[0:]. This violates the project-scoping invariant: messages that predate the active cloud session are uploaded to the finalizing session. Cap start at len(messages) instead so the slice is empty when local history has been truncated.

Example:

historyStart = 5 after a cloud session began with 5 local messages
after /load, len(messages) = 2
actual: projectMessages = messages[0:] (uploads 2 loaded messages to the old project's cloud session)
expected: projectMessages = messages[5:] = nil (uploads nothing)

Current:

func (t *CloudSessionTracker) CompleteCurrent(client *api.APIClient, totalTokens int, messages []api.Message) {
	start := t.historyStart
	if start < 0 || start > len(messages) {
		start = 0
	}

Proposed:

start := t.historyStart
if start < 0 {
	start = 0
}
if start > len(messages) {
	start = len(messages)
}
projectMessages := messages[start:]
Suggested change
func (t *CloudSessionTracker) CompleteCurrent(client *api.APIClient, totalTokens int, messages []api.Message) {
start := t.historyStart
if start < 0 || start > len(messages) {
start = 0
}
start := t.historyStart
if start < 0 {
start = 0
}
if start > len(messages) {
start = len(messages)
}
projectMessages := messages[start:]
More Info
  • Threat model: A user who starts a cloud-tracked REPL, then /loads a shorter saved session, and later exits (or switches projects again) causes the finalizing cloud session to receive messages from the loaded session. Those messages may belong to a different project or an earlier conversation, leaking context and inflating the event count/token totals for the wrong project.
  • Specific code citations: internal/session/cloud_session.go lines 117-121: CompleteCurrent slices messages[start:] after the start > len(messages) guard resets start to 0. The retrieval snapshot for internal/repl/repl.go shows /load replacing ag.History without touching the tracker.
  • Existing protections: The projectTokens < 0 clamp prevents negative token totals, but there is no equivalent guard for the history slice. SwitchProject and StartWithHistory only set historyStart to len(messages) at session creation; neither /load nor any other visible path updates it when ag.History is replaced.
  • Proposed mitigation: Split the bounds check so start > len(messages) caps start at len(messages), producing an empty slice when local history has been truncated. Optionally reset historyStart/tokenStart when /load replaces history, but the local cap is the minimal fix.
  • Alternative mitigations considered: Resetting the tracker on /load is broader and would lose the in-progress cloud session; capping the slice preserves the existing session while avoiding cross-project data leakage.
  • Severity calibration: Score 4: the failure path requires a /load after cloud session start, but when it happens it silently uploads wrong messages to the wrong project. It is not a crash but a data-scoping/information-leak bug.
Prompt To Fix With AI
This is a comment left during a code review.
Path: internal/session/cloud_session.go
Line: 117-121

Comment:
**Out-of-range historyStart uploads full replacement history to wrong cloud session**

When `historyStart` exceeds the current `len(messages)` (for example, after `/load` replaces `ag.History` with a shorter saved session), `CompleteCurrent` resets `start` to `0` and uploads `messages[0:]`. This violates the project-scoping invariant: messages that predate the active cloud session are uploaded to the finalizing session. Cap `start` at `len(messages)` instead so the slice is empty when local history has been truncated.

Example:
historyStart = 5 after a cloud session began with 5 local messages
after /load, len(messages) = 2
actual: projectMessages = messages[0:] (uploads 2 loaded messages to the old project's cloud session)
expected: projectMessages = messages[5:] = nil (uploads nothing)

Threat model:
A user who starts a cloud-tracked REPL, then `/load`s a shorter saved session, and later exits (or switches projects again) causes the finalizing cloud session to receive messages from the loaded session. Those messages may belong to a different project or an earlier conversation, leaking context and inflating the event count/token totals for the wrong project.

Specific code citations:
`internal/session/cloud_session.go` lines 117-121: `CompleteCurrent` slices `messages[start:]` after the `start > len(messages)` guard resets `start` to `0`. The retrieval snapshot for `internal/repl/repl.go` shows `/load` replacing `ag.History` without touching the tracker.

Existing protections:
The `projectTokens < 0` clamp prevents negative token totals, but there is no equivalent guard for the history slice. `SwitchProject` and `StartWithHistory` only set `historyStart` to `len(messages)` at session creation; neither `/load` nor any other visible path updates it when `ag.History` is replaced.

Proposed mitigation:
Split the bounds check so `start > len(messages)` caps `start` at `len(messages)`, producing an empty slice when local history has been truncated. Optionally reset `historyStart`/`tokenStart` when `/load` replaces history, but the local cap is the minimal fix.

Alternative mitigations considered:
Resetting the tracker on `/load` is broader and would lose the in-progress cloud session; capping the slice preserves the existing session while avoiding cross-project data leakage.

Severity calibration:
Score 4: the failure path requires a `/load` after cloud session start, but when it happens it silently uploads wrong messages to the wrong project. It is not a crash but a data-scoping/information-leak bug.

How can I resolve this? If you propose a fix, please make it concise.

if start > len(messages) {
start = len(messages)
}
projectMessages := messages[start:]
projectTokens := totalTokens - t.tokenStart
if projectTokens < 0 {
projectTokens = 0
}
t.Complete(client, projectTokens, SummaryFor(projectMessages), projectMessages)
}

// Complete patches the cloud session as finished and uploads the full message
Expand Down
97 changes: 97 additions & 0 deletions internal/session/cloud_session_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -319,3 +319,100 @@ func TestCloudTracker_Complete_SkipsUploadWhenNoMessages(t *testing.T) {
t.Errorf("expected 1 HTTP call (PATCH only), got %d", calls)
}
}

func TestCloudTracker_CompleteCurrent_SkipsHistoryBeforeTruncatedSession(t *testing.T) {
calls := 0
client, _ := newTestClient(t, func(w http.ResponseWriter, r *http.Request) {
calls++
_, _ = w.Write([]byte(`{}`))
})

tracker := CloudSessionTracker{
cloudID: "cloud-xyz",
historyStart: 5,
}
// A /load can replace the local conversation with a shorter one. None of
// those loaded messages belong to the active cloud session.
messages := []api.Message{
{Role: "user", Content: "loaded conversation"},
{Role: "assistant", Content: "loaded response"},
}
tracker.CompleteCurrent(client, 100, messages)

// Complete may PATCH the session, but must not upload loaded messages.
if calls != 1 {
t.Errorf("expected 1 HTTP call (PATCH only), got %d", calls)
}
}

func TestCloudTracker_SwitchProject_FinalizesOldSessionAndScopesNewHistory(t *testing.T) {
var createdProjects []int
var patchedIDs []string
var uploadedEventCounts []int
var patchedTokens []int
client, _ := newTestClient(t, func(w http.ResponseWriter, r *http.Request) {
switch {
case r.Method == http.MethodPost && r.URL.Path == "/api/agent-sessions":
var body struct {
ProjectID int `json:"project_id"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode create request: %v", err)
}
createdProjects = append(createdProjects, body.ProjectID)
_, _ = w.Write([]byte("{\"session_id\":\"cloud-" + string(rune('a'+len(createdProjects)-1)) + "\"}"))
case r.Method == http.MethodPatch:
var body struct {
TotalTokens int `json:"total_tokens"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode finalize request: %v", err)
}
patchedIDs = append(patchedIDs, r.URL.Path)
patchedTokens = append(patchedTokens, body.TotalTokens)
_, _ = w.Write([]byte("{}"))
case r.Method == http.MethodPost && r.URL.Path == "/api/agent-sessions/cloud-a/events":
fallthrough
case r.Method == http.MethodPost && r.URL.Path == "/api/agent-sessions/cloud-b/events":
var body struct {
Events []json.RawMessage `json:"events"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
t.Fatalf("decode events request: %v", err)
}
uploadedEventCounts = append(uploadedEventCounts, len(body.Events))
_, _ = w.Write([]byte("{}"))
default:
t.Fatalf("unexpected request: %s %s", r.Method, r.URL.Path)
}
})

initialMessages := []api.Message{
{Role: "user", Content: "test project one"},
{Role: "assistant", Content: "project one response"},
}
var tracker CloudSessionTracker
tracker.StartWithHistory(client, 1, "claude-sonnet-4-6", 0, 0)
if !tracker.SwitchProject(client, 2, "claude-sonnet-4-6", 100, initialMessages) {
t.Fatal("expected a project switch to move the cloud session")
}

messages := append(initialMessages,
api.Message{Role: "user", Content: "test project two"},
api.Message{Role: "assistant", Content: "project two response"},
)
tracker.CompleteCurrent(client, 150, messages)

if len(createdProjects) != 2 || createdProjects[0] != 1 || createdProjects[1] != 2 {
t.Errorf("created projects: got %v, want [1 2]", createdProjects)
}
if len(patchedIDs) != 2 || patchedIDs[0] != "/api/agent-sessions/cloud-a" || patchedIDs[1] != "/api/agent-sessions/cloud-b" {
t.Errorf("patched sessions: got %v", patchedIDs)
}
if len(patchedTokens) != 2 || patchedTokens[0] != 100 || patchedTokens[1] != 50 {
t.Errorf("per-project token totals: got %v, want [100 50]", patchedTokens)
}
if len(uploadedEventCounts) != 2 || uploadedEventCounts[0] != 2 || uploadedEventCounts[1] != 2 {
t.Errorf("uploaded event counts: got %v, want [2 2]", uploadedEventCounts)
}
}
Loading