Viewing files at commit c10cefe7abe5Commit detailsView latest

Browse files
oakmark/transport_cli.go 157 lines · 5.0 KB · go
1
package main
2
3
import (
4
	"context"
5
	"fmt"
6
	"os"
7
	"os/exec"
8
	"path/filepath"
9
	"time"
10
)
11
12
// cliTransport drives the real `oak` binary per operation. It is correct by
13
// construction (it IS the real client) and gives the honest real-agent number:
14
// the throughput a fleet of agents actually gets, including client-side BLAKE3
15
// chunking and one process spin-up per commit and per push. It is expected to
16
// be the floor, and to become generator-bound before the server does at high
17
// concurrency — which is precisely why the protocol transport exists.
18
//
19
// Push is two steps — `oak commit` then `oak push` — deliberately, NOT
20
// `oak commit --push`: on an open loopback `oak serve`, `oak commit --push`
21
// wrongly demands credentials while standalone `oak push` succeeds anonymously
22
// (filed as oak friction fb-64).
23
type cliTransport struct {
24
	oakBin string
25
	remote string
26
	owner  string
27
	name   string
28
	id     int
29
	branch string
30
	depth  int // clone depth for the session read; 0 = full, else --shallow
31
32
	baseDir string // per-agent temp root
33
	workDir string // branch/repo strategy: the long-lived clone
34
35
	sessDir    string // session strategy: current session's clone
36
	sessBranch string
37
	sessSeq    int // monotonic session counter (unique subdir + branch suffix)
38
}
39
40
func newCLITransport(oakBin, remote, owner, name string, id int, branch string, depth int) *cliTransport {
41
	return &cliTransport{oakBin: oakBin, remote: remote, owner: owner, name: name, id: id, branch: branch, depth: depth}
42
}
43
44
// oakEnv returns the process env with OAK_REMOTE pinned to this run's server,
45
// overriding any stored remote. OAK_API_KEY (for oak.space) is inherited from
46
// the parent environment untouched.
47
func (t *cliTransport) oakEnv() []string {
48
	return append(os.Environ(), "OAK_REMOTE="+t.remote)
49
}
50
51
// oak runs an oak subcommand in dir, returning combined output as the error
52
// text on failure so classifyOak can see conflict markers.
53
func (t *cliTransport) oak(ctx context.Context, dir string, args ...string) error {
54
	cmd := exec.CommandContext(ctx, t.oakBin, args...)
55
	cmd.Dir = dir
56
	cmd.Env = t.oakEnv()
57
	out, err := cmd.CombinedOutput()
58
	if err != nil {
59
		return fmt.Errorf("oak %v: %v: %s", args, err, string(out))
60
	}
61
	return nil
62
}
63
64
func (t *cliTransport) writeFiles(dir string, files []FileSpec) error {
65
	for _, f := range files {
66
		p := filepath.Join(dir, f.Path)
67
		if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
68
			return err
69
		}
70
		if err := os.WriteFile(p, f.Content, 0o644); err != nil {
71
			return err
72
		}
73
	}
74
	return nil
75
}
76
77
// cloneInto clones the repo into a fresh dir under baseDir and switches to
78
// branchName. shallow controls the clone depth for the initial download.
79
func (t *cliTransport) cloneInto(ctx context.Context, sub, branchName string, shallow bool) (string, error) {
80
	dir := filepath.Join(t.baseDir, sub)
81
	args := []string{"clone", t.owner + "/" + t.name, dir, "-r", t.remote}
82
	if shallow {
83
		args = append(args, "--shallow")
84
	}
85
	if err := t.oak(ctx, t.baseDir, args...); err != nil {
86
		return "", err
87
	}
88
	if err := t.oak(ctx, dir, "switch", "-c", branchName); err != nil {
89
		return "", err
90
	}
91
	return dir, nil
92
}
93
94
// commitPush writes the changeset and publishes it: `oak commit` then
95
// `oak push` (see the type doc for why it's two steps).
96
func (t *cliTransport) commitPush(ctx context.Context, dir string, files []FileSpec) error {
97
	if err := t.writeFiles(dir, files); err != nil {
98
		return err
99
	}
100
	if err := t.oak(ctx, dir, "commit", "--quiet"); err != nil {
101
		return err
102
	}
103
	return t.oak(ctx, dir, "push", "-r", t.remote)
104
}
105
106
func (t *cliTransport) Init(ctx context.Context) error {
107
	base, err := os.MkdirTemp("", fmt.Sprintf("oakmark-cli-a%d-", t.id))
108
	if err != nil {
109
		return err
110
	}
111
	t.baseDir = base
112
	// Setup clone is unmeasured; shallow keeps it fast (working tree identical).
113
	wd, err := t.cloneInto(ctx, "wd", t.branch, true)
114
	if err != nil {
115
		return err
116
	}
117
	t.workDir = wd
118
	return nil
119
}
120
121
func (t *cliTransport) Push(ctx context.Context, files []FileSpec) error {
122
	return t.commitPush(ctx, t.workDir, files)
123
}
124
125
func (t *cliTransport) CloneSession(ctx context.Context) error {
126
	// Fresh clone per session (the measured read). A previous session's dir is
127
	// left for Close() to reap so cleanup never steals from the next clone's time.
128
	t.sessSeq++
129
	shallow := t.depth != 0
130
	t.sessBranch = fmt.Sprintf("%s-s%d", t.branch, t.sessSeq)
131
	dir, err := t.cloneInto(ctx, fmt.Sprintf("s%d", t.sessSeq), t.sessBranch, shallow)
132
	if err != nil {
133
		return err
134
	}
135
	t.sessDir = dir
136
	return nil
137
}
138
139
func (t *cliTransport) PushCheckpoint(ctx context.Context, session int, files []FileSpec) error {
140
	return t.commitPush(ctx, t.sessDir, files)
141
}
142
143
func (t *cliTransport) CleanupSession(session int) {
144
	// Best-effort: close the ephemeral branch so refs don't accumulate on the
145
	// target. Runs on a fresh context so it still fires as the window closes.
146
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
147
	defer cancel()
148
	if t.sessDir != "" && t.sessBranch != "" {
149
		_ = t.oak(ctx, t.sessDir, "close", t.sessBranch)
150
	}
151
}
152
153
func (t *cliTransport) Close() {
154
	if t.baseDir != "" {
155
		_ = os.RemoveAll(t.baseDir)
156
	}
157
}