153 lines
3.1 KiB
Go
153 lines
3.1 KiB
Go
package agents
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"log/slog"
|
|
"sync"
|
|
"time"
|
|
|
|
"g.fc.run/misc/byre/internal/api"
|
|
"g.fc.run/misc/byre/internal/claude"
|
|
"g.fc.run/misc/byre/internal/git"
|
|
"g.fc.run/misc/byre/internal/store"
|
|
)
|
|
|
|
type snapRequest struct {
|
|
trigger api.SnapshotTrigger
|
|
seq int64
|
|
}
|
|
|
|
type snapState struct {
|
|
mu sync.Mutex
|
|
running bool
|
|
pending *snapRequest
|
|
closed bool
|
|
}
|
|
|
|
func (m *Manager) snapStateLocked(taskID string) *snapState {
|
|
st := m.snaps[taskID]
|
|
if st == nil {
|
|
st = &snapState{}
|
|
m.snaps[taskID] = st
|
|
}
|
|
return st
|
|
}
|
|
|
|
func (m *Manager) snapshotSoon(l *live, trigger api.SnapshotTrigger, seq int64) {
|
|
m.snapMu.Lock()
|
|
defer m.snapMu.Unlock()
|
|
|
|
st := m.snapStateLocked(l.taskID)
|
|
if st.closed {
|
|
return
|
|
}
|
|
req := &snapRequest{trigger: trigger, seq: seq}
|
|
if st.running {
|
|
st.pending = req
|
|
return
|
|
}
|
|
st.running = true
|
|
go m.runSnapshots(l, st, req)
|
|
}
|
|
|
|
func (m *Manager) runSnapshots(l *live, st *snapState, req *snapRequest) {
|
|
for {
|
|
st.mu.Lock()
|
|
m.snapMu.Lock()
|
|
closed := st.closed
|
|
m.snapMu.Unlock()
|
|
if !closed {
|
|
err := m.takeSnapshot(context.Background(), l, req.trigger, req.seq)
|
|
if err != nil {
|
|
slog.Error("take snapshot", "task", l.taskID, "trigger", req.trigger, "err", err)
|
|
}
|
|
}
|
|
st.mu.Unlock()
|
|
|
|
m.snapMu.Lock()
|
|
req = st.pending
|
|
st.pending = nil
|
|
if req == nil || st.closed {
|
|
st.running = false
|
|
m.snapMu.Unlock()
|
|
return
|
|
}
|
|
m.snapMu.Unlock()
|
|
}
|
|
}
|
|
|
|
func (m *Manager) takeSnapshot(ctx context.Context, l *live, trigger api.SnapshotTrigger, seq int64) error {
|
|
if seq == 0 {
|
|
var err error
|
|
seq, err = m.store.LastTaskSeq(ctx, l.taskID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
tree, err := git.Snapshot(l.dir)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
last, err := m.store.LastSnapshot(ctx, l.taskID)
|
|
if err != nil && !errors.Is(err, store.ErrNotFound) {
|
|
return err
|
|
}
|
|
if err == nil && last.TreeSha == tree {
|
|
return nil
|
|
}
|
|
|
|
n := last.N + 1
|
|
commit, err := git.CommitSnapshot(l.dir, tree, git.SnapshotRef(l.taskID, n), fmt.Sprintf("byre snapshot %d (%s)", n, trigger), []string{
|
|
"GIT_AUTHOR_NAME=" + l.gitName,
|
|
"GIT_AUTHOR_EMAIL=" + l.gitEmail,
|
|
"GIT_COMMITTER_NAME=" + l.gitName,
|
|
"GIT_COMMITTER_EMAIL=" + l.gitEmail,
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return m.store.CreateSnapshot(ctx, api.Snapshot{
|
|
Id: store.NewID(),
|
|
TaskId: l.taskID,
|
|
AgentSessionId: l.id,
|
|
N: n,
|
|
CommitSha: commit,
|
|
TreeSha: tree,
|
|
EventSeq: seq,
|
|
Trigger: trigger,
|
|
CreatedAt: time.Now().UTC(),
|
|
}, l.projectID)
|
|
}
|
|
|
|
func (m *Manager) dropSnapshots(repo, taskID string) error {
|
|
m.snapMu.Lock()
|
|
st := m.snapStateLocked(taskID)
|
|
st.closed = true
|
|
m.snapMu.Unlock()
|
|
|
|
st.mu.Lock()
|
|
defer st.mu.Unlock()
|
|
return git.DeleteSnapshotRefs(repo, taskID)
|
|
}
|
|
|
|
func (m *Manager) afterTool(l *live) claude.HookFunc {
|
|
return func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
|
|
text := m.deliverMidTurn(ctx, l)
|
|
if text == "" {
|
|
return map[string]any{}, nil
|
|
}
|
|
return map[string]any{
|
|
"hookSpecificOutput": map[string]any{
|
|
"hookEventName": "PostToolUse",
|
|
"additionalContext": text,
|
|
},
|
|
}, nil
|
|
}
|
|
}
|