Cyrene_v2/internal/runtime/loop/checkpoint.go

package loop

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/agent"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/persist"
)

// checkpointStore is the minimal surface the loop needs to persist a durable
// snapshot and resume an interrupted run (design §1 durable execution).
// *persist.Store satisfies it in production; tests substitute a fake to assert
// the resume path without a real database.
type checkpointStore interface {
	// Save upserts a checkpoint.
	Save(ctx context.Context, cp *persist.Checkpoint) error
	// LoadLatest returns the most recently updated checkpoint for agentID, or
	// ErrNotFound when the agent has none.
	LoadLatest(ctx context.Context, agentID string) (*persist.Checkpoint, error)
}

// checkpointState is the durable snapshot persisted at every step. Carrying the
// message window and the step index lets a long-running task be resumed from the
// latest checkpoint without losing context (design §1 durable execution).
type checkpointState struct {
	Step       int             `json:"step"`
	Messages   []model.Message `json:"messages"`
	StopReason string          `json:"stop_reason,omitempty"`
}

// saveCheckpoint persists a snapshot at step. It never fails the run: a nil store
// is a no-op and a marshal/save error is logged.
func (l *Loop) saveCheckpoint(ctx context.Context, sessionID string, step int, messages []model.Message, stopReason string) {
	if l.checkpoint == nil {
		return
	}
	state := checkpointState{
		Step:       step,
		Messages:   messages,
		StopReason: stopReason,
	}
	buf, err := json.Marshal(state)
	if err != nil {
		l.log.Warn("checkpoint marshal", "step", step, "error", err)
		return
	}
	cp := &persist.Checkpoint{
		ID:      fmt.Sprintf("%s:%d", sessionID, step),
		AgentID: sessionID,
		Step:    step,
		State:   buf,
	}
	if err := l.checkpoint.Save(ctx, cp); err != nil {
		l.log.Warn("checkpoint save", "step", step, "error", err)
	}
}

// loadCheckpoint resumes an interrupted run from the session's latest durable
// checkpoint (design §1 durable execution). It returns the restored message
// window, the next loop step to run (the checkpoint step + 1) and whether a
// resumable checkpoint was found. When no checkpoint exists, or the latest one
// reflects a completed/blocked run, it returns the caller's messages with step 1
// so the loop starts fresh.
//
// Restoration is idempotent: the checkpoint captures the message window exactly as
// it stood at the end of the last completed step (assistant tool calls plus their
// results are already fused in), so resuming at the next step rebuilds the model
// context from history and never re-executes a completed step's tool call nor
// regenerates a finished answer. Loading is read-only — a resume never deletes the
// checkpoint, so every step's snapshot stays queryable.
func (l *Loop) loadCheckpoint(ctx context.Context, sessionID string, messages []model.Message) ([]model.Message, int, bool) {
	if l.checkpoint == nil {
		return messages, 1, false
	}
	cp, err := l.checkpoint.LoadLatest(ctx, sessionID)
	if err != nil {
		if !errors.Is(err, persist.ErrNotFound) {
			l.log.Warn("checkpoint load", "session", sessionID, "error", err)
		}
		return messages, 1, false
	}
	st := checkpointState{}
	if err := json.Unmarshal(cp.State, &st); err != nil {
		l.log.Warn("checkpoint unmarshal", "session", sessionID, "step", cp.Step, "error", err)
		return messages, 1, false
	}
	if !resumableStopReason(st.StopReason) {
		l.log.Debug("checkpoint not resumable", "session", sessionID, "step", cp.Step, "stop_reason", st.StopReason)
		return messages, 1, false
	}
	l.log.Info("checkpoint resumed", "session", sessionID, "step", cp.Step, "next_step", cp.Step+1, "messages", len(st.Messages))
	return st.Messages, cp.Step + 1, true
}

// resumableStopReason reports whether a checkpoint whose last step ended with the
// given stop reason represents an in-flight run worth resuming. Only a
// mid-tool-loop step is genuinely "interrupted": the loop folded the tool results
// into the window and was about to continue, so picking up at the next step
// rebuilds context without duplicating work. A checkpoint carrying a final answer,
// a budget exhaustion, a safety-net termination or a guardrail denial is a
// completed/blocked run — resuming it would regenerate an already-produced answer
// (a duplicate side effect) — so those fall through to a fresh start (a caller
// opening a genuinely new conversation uses a distinct sessionID, which misses
// LoadLatest anyway).
func resumableStopReason(stopReason string) bool {
	switch stopReason {
	case "", agent.StopReasonToolCalls:
		return true
	default:
		return false
	}
}