Add Claude Code stream-json client, report gate, and terminal driver
This commit is contained in:
Ian Gulliver
@@ -0,0 +1,525 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"g.fc.run/misc/byre/internal/mcp"
|
||||
)
|
||||
|
||||
type HookFunc func(ctx context.Context, input json.RawMessage) (map[string]any, error)
|
||||
|
||||
type HookMatcher struct {
|
||||
Matcher string
|
||||
Func HookFunc
|
||||
}
|
||||
|
||||
type Options struct {
|
||||
Cwd string
|
||||
Model string
|
||||
Resume string
|
||||
ForkSession bool
|
||||
GitName string
|
||||
GitEmail string
|
||||
Hooks map[string][]HookMatcher
|
||||
MCPServers []*mcp.Server
|
||||
Transcript io.Writer
|
||||
}
|
||||
|
||||
type Event struct {
|
||||
Type string
|
||||
Subtype string
|
||||
Raw json.RawMessage
|
||||
}
|
||||
|
||||
type Block struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
Name string `json:"name"`
|
||||
Input json.RawMessage `json:"input"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
IsError bool `json:"is_error"`
|
||||
}
|
||||
|
||||
type Result struct {
|
||||
Subtype string `json:"subtype"`
|
||||
IsError bool `json:"is_error"`
|
||||
Result string `json:"result"`
|
||||
DurationMS int `json:"duration_ms"`
|
||||
NumTurns int `json:"num_turns"`
|
||||
TotalCostUSD float64 `json:"total_cost_usd"`
|
||||
}
|
||||
|
||||
func (e Event) Blocks() ([]Block, error) {
|
||||
msg := struct {
|
||||
Message struct {
|
||||
Content json.RawMessage `json:"content"`
|
||||
} `json:"message"`
|
||||
}{}
|
||||
err := json.Unmarshal(e.Raw, &msg)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(msg.Message.Content) > 0 && msg.Message.Content[0] == '"' {
|
||||
text := ""
|
||||
err = json.Unmarshal(msg.Message.Content, &text)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return []Block{{Type: "text", Text: text}}, nil
|
||||
}
|
||||
|
||||
blocks := []Block{}
|
||||
err = json.Unmarshal(msg.Message.Content, &blocks)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return blocks, nil
|
||||
}
|
||||
|
||||
func (e Event) Result() (Result, error) {
|
||||
res := Result{}
|
||||
err := json.Unmarshal(e.Raw, &res)
|
||||
return res, err
|
||||
}
|
||||
|
||||
type Session struct {
|
||||
cmd *exec.Cmd
|
||||
stdin io.WriteCloser
|
||||
events chan Event
|
||||
done chan struct{}
|
||||
|
||||
writeMu sync.Mutex
|
||||
transcript io.Writer
|
||||
|
||||
pendingMu sync.Mutex
|
||||
pending map[string]chan controlResult
|
||||
|
||||
hooks map[string]HookFunc
|
||||
servers map[string]*mcp.Server
|
||||
|
||||
idMu sync.Mutex
|
||||
id string
|
||||
}
|
||||
|
||||
type controlResult struct {
|
||||
response json.RawMessage
|
||||
err error
|
||||
}
|
||||
|
||||
type envelope struct {
|
||||
Type string `json:"type"`
|
||||
Subtype string `json:"subtype"`
|
||||
SessionID string `json:"session_id"`
|
||||
RequestID string `json:"request_id"`
|
||||
Request json.RawMessage `json:"request"`
|
||||
Response json.RawMessage `json:"response"`
|
||||
}
|
||||
|
||||
func Start(ctx context.Context, opts Options) (*Session, error) {
|
||||
s := &Session{
|
||||
events: make(chan Event, 256),
|
||||
done: make(chan struct{}),
|
||||
transcript: opts.Transcript,
|
||||
pending: map[string]chan controlResult{},
|
||||
hooks: map[string]HookFunc{},
|
||||
servers: map[string]*mcp.Server{},
|
||||
}
|
||||
|
||||
args, err := s.buildArgs(opts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
s.cmd = exec.Command("claude", args...)
|
||||
s.cmd.Dir = opts.Cwd
|
||||
s.cmd.Stderr = os.Stderr
|
||||
s.cmd.Env = buildEnv(opts)
|
||||
|
||||
s.stdin, err = s.cmd.StdinPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
stdout, err := s.cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
err = s.cmd.Start()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("start claude: %w", err)
|
||||
}
|
||||
|
||||
go s.readLoop(stdout)
|
||||
|
||||
_, err = s.Control(ctx, map[string]any{
|
||||
"subtype": "initialize",
|
||||
"hooks": s.hookConfig(opts.Hooks),
|
||||
})
|
||||
if err != nil {
|
||||
s.Close()
|
||||
return nil, fmt.Errorf("initialize: %w", err)
|
||||
}
|
||||
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *Session) buildArgs(opts Options) ([]string, error) {
|
||||
settings, err := json.Marshal(map[string]any{
|
||||
"attribution": map[string]any{"commit": "", "pr": ""},
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
servers := map[string]any{}
|
||||
for _, srv := range opts.MCPServers {
|
||||
s.servers[srv.Name] = srv
|
||||
servers[srv.Name] = map[string]any{"type": "sdk", "name": srv.Name}
|
||||
}
|
||||
|
||||
mcpConfig, err := json.Marshal(map[string]any{"mcpServers": servers})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
args := []string{
|
||||
"--output-format", "stream-json",
|
||||
"--input-format", "stream-json",
|
||||
"--verbose",
|
||||
"--permission-mode", "bypassPermissions",
|
||||
"--allow-dangerously-skip-permissions",
|
||||
"--settings", string(settings),
|
||||
"--mcp-config", string(mcpConfig),
|
||||
}
|
||||
|
||||
if opts.Model != "" {
|
||||
args = append(args, "--model", opts.Model)
|
||||
}
|
||||
|
||||
if opts.Resume != "" {
|
||||
args = append(args, "--resume="+opts.Resume)
|
||||
}
|
||||
|
||||
if opts.ForkSession {
|
||||
args = append(args, "--fork-session")
|
||||
}
|
||||
|
||||
return args, nil
|
||||
}
|
||||
|
||||
func buildEnv(opts Options) []string {
|
||||
env := []string{}
|
||||
for _, kv := range os.Environ() {
|
||||
if len(kv) >= 11 && kv[:11] == "CLAUDECODE=" {
|
||||
continue
|
||||
}
|
||||
env = append(env, kv)
|
||||
}
|
||||
|
||||
return append(env,
|
||||
"CLAUDE_CODE_ENTRYPOINT=sdk-go",
|
||||
"CLAUDE_CODE_DISABLE_AUTOUPDATER=1",
|
||||
"GIT_AUTHOR_NAME="+opts.GitName,
|
||||
"GIT_AUTHOR_EMAIL="+opts.GitEmail,
|
||||
"GIT_COMMITTER_NAME="+opts.GitName,
|
||||
"GIT_COMMITTER_EMAIL="+opts.GitEmail,
|
||||
)
|
||||
}
|
||||
|
||||
func (s *Session) hookConfig(hooks map[string][]HookMatcher) map[string]any {
|
||||
if len(hooks) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
cfg := map[string]any{}
|
||||
n := 0
|
||||
for event, matchers := range hooks {
|
||||
list := []map[string]any{}
|
||||
for _, m := range matchers {
|
||||
id := fmt.Sprintf("hook_%d", n)
|
||||
n++
|
||||
s.hooks[id] = m.Func
|
||||
list = append(list, map[string]any{
|
||||
"matcher": m.Matcher,
|
||||
"hookCallbackIds": []string{id},
|
||||
})
|
||||
}
|
||||
cfg[event] = list
|
||||
}
|
||||
|
||||
return cfg
|
||||
}
|
||||
|
||||
func (s *Session) Events() <-chan Event {
|
||||
return s.events
|
||||
}
|
||||
|
||||
func (s *Session) Done() <-chan struct{} {
|
||||
return s.done
|
||||
}
|
||||
|
||||
func (s *Session) ID() string {
|
||||
s.idMu.Lock()
|
||||
defer s.idMu.Unlock()
|
||||
return s.id
|
||||
}
|
||||
|
||||
func (s *Session) Send(text string) error {
|
||||
return s.write(map[string]any{
|
||||
"type": "user",
|
||||
"message": map[string]any{"role": "user", "content": text},
|
||||
"parent_tool_use_id": nil,
|
||||
"session_id": s.ID(),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Session) Interrupt(ctx context.Context) error {
|
||||
_, err := s.Control(ctx, map[string]any{"subtype": "interrupt"})
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Session) Control(ctx context.Context, req map[string]any) (json.RawMessage, error) {
|
||||
id := newRequestID()
|
||||
ch := make(chan controlResult, 1)
|
||||
|
||||
s.pendingMu.Lock()
|
||||
s.pending[id] = ch
|
||||
s.pendingMu.Unlock()
|
||||
|
||||
defer func() {
|
||||
s.pendingMu.Lock()
|
||||
delete(s.pending, id)
|
||||
s.pendingMu.Unlock()
|
||||
}()
|
||||
|
||||
err := s.write(map[string]any{
|
||||
"type": "control_request",
|
||||
"request_id": id,
|
||||
"request": req,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
select {
|
||||
case res := <-ch:
|
||||
return res.response, res.err
|
||||
case <-s.done:
|
||||
return nil, errors.New("claude exited")
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) Close() error {
|
||||
s.writeMu.Lock()
|
||||
s.stdin.Close()
|
||||
s.writeMu.Unlock()
|
||||
|
||||
select {
|
||||
case <-s.done:
|
||||
case <-time.After(5 * time.Second):
|
||||
s.cmd.Process.Kill()
|
||||
<-s.done
|
||||
}
|
||||
|
||||
return s.cmd.Wait()
|
||||
}
|
||||
|
||||
func (s *Session) write(v any) error {
|
||||
b, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
|
||||
s.record("send", b)
|
||||
|
||||
_, err = s.stdin.Write(append(b, '\n'))
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Session) record(dir string, line []byte) {
|
||||
if s.transcript == nil {
|
||||
return
|
||||
}
|
||||
|
||||
b, err := json.Marshal(map[string]any{"dir": dir, "msg": json.RawMessage(line)})
|
||||
if err != nil {
|
||||
slog.Error("record transcript", "err", err)
|
||||
return
|
||||
}
|
||||
|
||||
s.transcript.Write(append(b, '\n'))
|
||||
}
|
||||
|
||||
func (s *Session) readLoop(stdout io.Reader) {
|
||||
defer close(s.done)
|
||||
defer close(s.events)
|
||||
|
||||
r := bufio.NewReader(stdout)
|
||||
for {
|
||||
line, err := r.ReadBytes('\n')
|
||||
if len(line) > 1 && line[0] == '{' {
|
||||
s.handleLine(line[:len(line)-1])
|
||||
}
|
||||
if err != nil {
|
||||
if err != io.EOF {
|
||||
slog.Error("read claude stdout", "err", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Session) handleLine(line []byte) {
|
||||
s.writeMu.Lock()
|
||||
s.record("recv", line)
|
||||
s.writeMu.Unlock()
|
||||
|
||||
env := envelope{}
|
||||
err := json.Unmarshal(line, &env)
|
||||
if err != nil {
|
||||
slog.Error("parse claude message", "err", err, "line", string(line))
|
||||
return
|
||||
}
|
||||
|
||||
switch env.Type {
|
||||
case "control_response":
|
||||
s.handleControlResponse(env.Response)
|
||||
return
|
||||
|
||||
case "control_request":
|
||||
go s.handleControlRequest(env.RequestID, env.Request)
|
||||
return
|
||||
|
||||
case "control_cancel_request", "keep_alive":
|
||||
return
|
||||
}
|
||||
|
||||
if env.Type == "system" && env.Subtype == "init" {
|
||||
s.idMu.Lock()
|
||||
s.id = env.SessionID
|
||||
s.idMu.Unlock()
|
||||
}
|
||||
|
||||
s.events <- Event{Type: env.Type, Subtype: env.Subtype, Raw: append(json.RawMessage{}, line...)}
|
||||
}
|
||||
|
||||
func (s *Session) handleControlResponse(raw json.RawMessage) {
|
||||
resp := struct {
|
||||
Subtype string `json:"subtype"`
|
||||
RequestID string `json:"request_id"`
|
||||
Response json.RawMessage `json:"response"`
|
||||
Error string `json:"error"`
|
||||
}{}
|
||||
err := json.Unmarshal(raw, &resp)
|
||||
if err != nil {
|
||||
slog.Error("parse control response", "err", err)
|
||||
return
|
||||
}
|
||||
|
||||
s.pendingMu.Lock()
|
||||
ch := s.pending[resp.RequestID]
|
||||
s.pendingMu.Unlock()
|
||||
|
||||
if ch == nil {
|
||||
slog.Warn("control response for unknown request", "request_id", resp.RequestID)
|
||||
return
|
||||
}
|
||||
|
||||
if resp.Subtype == "error" {
|
||||
ch <- controlResult{err: errors.New(resp.Error)}
|
||||
return
|
||||
}
|
||||
|
||||
ch <- controlResult{response: resp.Response}
|
||||
}
|
||||
|
||||
func (s *Session) handleControlRequest(id string, raw json.RawMessage) {
|
||||
result, err := s.dispatchControl(raw)
|
||||
if err != nil {
|
||||
slog.Error("control request failed", "request_id", id, "err", err)
|
||||
s.write(map[string]any{
|
||||
"type": "control_response",
|
||||
"response": map[string]any{
|
||||
"subtype": "error",
|
||||
"request_id": id,
|
||||
"error": err.Error(),
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
s.write(map[string]any{
|
||||
"type": "control_response",
|
||||
"response": map[string]any{
|
||||
"subtype": "success",
|
||||
"request_id": id,
|
||||
"response": result,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Session) dispatchControl(raw json.RawMessage) (any, error) {
|
||||
req := struct {
|
||||
Subtype string `json:"subtype"`
|
||||
CallbackID string `json:"callback_id"`
|
||||
Input json.RawMessage `json:"input"`
|
||||
ServerName string `json:"server_name"`
|
||||
Message json.RawMessage `json:"message"`
|
||||
}{}
|
||||
err := json.Unmarshal(raw, &req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
switch req.Subtype {
|
||||
case "hook_callback":
|
||||
fn := s.hooks[req.CallbackID]
|
||||
if fn == nil {
|
||||
return nil, fmt.Errorf("no hook callback %s", req.CallbackID)
|
||||
}
|
||||
return fn(ctx, req.Input)
|
||||
|
||||
case "mcp_message":
|
||||
srv := s.servers[req.ServerName]
|
||||
if srv == nil {
|
||||
return nil, fmt.Errorf("no mcp server %s", req.ServerName)
|
||||
}
|
||||
resp, err := srv.Handle(ctx, req.Message)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return map[string]any{"mcp_response": resp}, nil
|
||||
|
||||
case "can_use_tool":
|
||||
return map[string]any{"behavior": "allow", "updatedInput": req.Input}, nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("unsupported control request %s", req.Subtype)
|
||||
}
|
||||
|
||||
func newRequestID() string {
|
||||
b := make([]byte, 8)
|
||||
rand.Read(b)
|
||||
return "req_" + hex.EncodeToString(b)
|
||||
}
|
||||
@@ -0,0 +1,313 @@
|
||||
package claude_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"g.fc.run/misc/byre/internal/claude"
|
||||
"g.fc.run/misc/byre/internal/mcp"
|
||||
"g.fc.run/misc/byre/internal/report"
|
||||
)
|
||||
|
||||
const (
|
||||
testModel = "sonnet"
|
||||
gitName = "Byre Test"
|
||||
gitEmail = "test@byre.invalid"
|
||||
)
|
||||
|
||||
func TestCommitIdentityHooksAndReport(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testContext(t)
|
||||
repo := scratchRepo(t)
|
||||
|
||||
reports := []report.Report{}
|
||||
mu := sync.Mutex{}
|
||||
gate := &report.Gate{OnReport: func(r report.Report) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
reports = append(reports, r)
|
||||
}}
|
||||
|
||||
tools := []string{}
|
||||
s := start(t, ctx, claude.Options{
|
||||
Cwd: repo,
|
||||
MCPServers: []*mcp.Server{{Name: "byre", Version: "0.0.0", Tools: []mcp.Tool{gate.Tool()}}},
|
||||
Hooks: map[string][]claude.HookMatcher{
|
||||
"Stop": {{Func: gate.StopHook}},
|
||||
"PostToolUse": {{Func: func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
|
||||
in := struct {
|
||||
ToolName string `json:"tool_name"`
|
||||
}{}
|
||||
err := json.Unmarshal(input, &in)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
tools = append(tools, in.ToolName)
|
||||
return map[string]any{}, nil
|
||||
}}},
|
||||
},
|
||||
})
|
||||
|
||||
send(t, s, `Create a file hello.txt containing exactly the word hi, then git add and git commit it with the message "add hello".`)
|
||||
waitResult(t, ctx, s, nil)
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
t.Logf("tools=%v reports=%d nudges=%d", tools, len(reports), gate.Nudges())
|
||||
|
||||
if len(reports) == 0 {
|
||||
t.Fatalf("report was never called")
|
||||
}
|
||||
if len(tools) == 0 {
|
||||
t.Fatalf("PostToolUse hook never fired")
|
||||
}
|
||||
|
||||
log := git(t, repo, "log", "-1", "--format=%an <%ae>%n%cn <%ce>%n%B")
|
||||
want := gitName + " <" + gitEmail + ">\n" + gitName + " <" + gitEmail + ">\nadd hello"
|
||||
if !strings.HasPrefix(log, want) {
|
||||
t.Fatalf("unexpected commit:\n%s", log)
|
||||
}
|
||||
lower := strings.ToLower(log)
|
||||
if strings.Contains(lower, "co-authored-by") || strings.Contains(lower, "claude") {
|
||||
t.Fatalf("agent attribution in commit:\n%s", log)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStopHookBlocks(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testContext(t)
|
||||
|
||||
blocked := false
|
||||
mu := sync.Mutex{}
|
||||
s := start(t, ctx, claude.Options{
|
||||
Cwd: scratchRepo(t),
|
||||
Hooks: map[string][]claude.HookMatcher{
|
||||
"Stop": {{Func: func(ctx context.Context, input json.RawMessage) (map[string]any, error) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if blocked {
|
||||
return map[string]any{}, nil
|
||||
}
|
||||
blocked = true
|
||||
return map[string]any{"decision": "block", "reason": "Before stopping, reply with the single word PINEAPPLE."}, nil
|
||||
}}},
|
||||
},
|
||||
})
|
||||
|
||||
send(t, s, "Reply with the single word ok.")
|
||||
text := waitResult(t, ctx, s, nil)
|
||||
|
||||
if !blocked {
|
||||
t.Fatalf("Stop hook never fired")
|
||||
}
|
||||
if !strings.Contains(text, "PINEAPPLE") {
|
||||
t.Fatalf("Stop hook block reason not acted on; assistant said %q", text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInterrupt(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testContext(t)
|
||||
s := start(t, ctx, claude.Options{Cwd: scratchRepo(t)})
|
||||
|
||||
send(t, s, "Run the bash command `sleep 120` in the foreground (not in the background), then say done.")
|
||||
|
||||
interruptedAt := time.Time{}
|
||||
waitResult(t, ctx, s, func(ev claude.Event) {
|
||||
if !interruptedAt.IsZero() || !hasToolUse(t, ev) {
|
||||
return
|
||||
}
|
||||
interruptedAt = time.Now()
|
||||
err := s.Interrupt(ctx)
|
||||
if err != nil {
|
||||
t.Errorf("interrupt: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
if interruptedAt.IsZero() {
|
||||
t.Fatalf("never saw a tool use to interrupt")
|
||||
}
|
||||
elapsed := time.Since(interruptedAt)
|
||||
t.Logf("turn ended %s after interrupt", elapsed)
|
||||
if elapsed > 30*time.Second {
|
||||
t.Fatalf("interrupt took %s to end the turn", elapsed)
|
||||
}
|
||||
|
||||
send(t, s, "Reply with the single word ok.")
|
||||
text := waitResult(t, ctx, s, nil)
|
||||
if !strings.Contains(strings.ToLower(text), "ok") {
|
||||
t.Fatalf("session unusable after interrupt; assistant said %q", text)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResume(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testContext(t)
|
||||
repo := scratchRepo(t)
|
||||
|
||||
first, err := claude.Start(ctx, options(t, "first", claude.Options{Cwd: repo}))
|
||||
if err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
send(t, first, "Remember the codeword AZURE-17. Reply with the single word ok.")
|
||||
waitResult(t, ctx, first, nil)
|
||||
id := first.ID()
|
||||
closeSession(t, first)
|
||||
|
||||
if id == "" {
|
||||
t.Fatalf("no session id from first session")
|
||||
}
|
||||
|
||||
resumed, err := claude.Start(ctx, options(t, "resumed", claude.Options{Cwd: repo, Resume: id}))
|
||||
if err != nil {
|
||||
t.Fatalf("resume: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { closeSession(t, resumed) })
|
||||
|
||||
send(t, resumed, "What was the codeword? Reply with just the codeword.")
|
||||
text := waitResult(t, ctx, resumed, nil)
|
||||
|
||||
if !strings.Contains(text, "AZURE-17") {
|
||||
t.Fatalf("resumed session lost context; assistant said %q", text)
|
||||
}
|
||||
if resumed.ID() != id {
|
||||
t.Fatalf("resume changed session id: want %s, got %s", id, resumed.ID())
|
||||
}
|
||||
}
|
||||
|
||||
func testContext(t *testing.T) context.Context {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
|
||||
t.Cleanup(cancel)
|
||||
return ctx
|
||||
}
|
||||
|
||||
func scratchRepo(t *testing.T) string {
|
||||
dir := t.TempDir()
|
||||
git(t, dir, "init", "--initial-branch=main")
|
||||
git(t, dir, "commit", "--allow-empty", "--message", "root")
|
||||
return dir
|
||||
}
|
||||
|
||||
func options(t *testing.T, suffix string, opts claude.Options) claude.Options {
|
||||
name := t.Name()
|
||||
if suffix != "" {
|
||||
name += "-" + suffix
|
||||
}
|
||||
|
||||
err := os.MkdirAll(filepath.Join("testdata", "transcripts"), 0o755)
|
||||
if err != nil {
|
||||
t.Fatalf("mkdir transcripts: %v", err)
|
||||
}
|
||||
|
||||
f, err := os.Create(filepath.Join("testdata", "transcripts", name+".jsonl"))
|
||||
if err != nil {
|
||||
t.Fatalf("create transcript: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { f.Close() })
|
||||
|
||||
opts.Model = testModel
|
||||
opts.GitName = gitName
|
||||
opts.GitEmail = gitEmail
|
||||
opts.Transcript = f
|
||||
return opts
|
||||
}
|
||||
|
||||
func start(t *testing.T, ctx context.Context, opts claude.Options) *claude.Session {
|
||||
s, err := claude.Start(ctx, options(t, "", opts))
|
||||
if err != nil {
|
||||
t.Fatalf("start: %v", err)
|
||||
}
|
||||
t.Cleanup(func() { closeSession(t, s) })
|
||||
return s
|
||||
}
|
||||
|
||||
func closeSession(t *testing.T, s *claude.Session) {
|
||||
err := s.Close()
|
||||
if err != nil {
|
||||
t.Logf("claude exit: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func send(t *testing.T, s *claude.Session, text string) {
|
||||
err := s.Send(text)
|
||||
if err != nil {
|
||||
t.Fatalf("send: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func waitResult(t *testing.T, ctx context.Context, s *claude.Session, onEvent func(claude.Event)) string {
|
||||
text := []string{}
|
||||
for {
|
||||
select {
|
||||
case ev, ok := <-s.Events():
|
||||
if !ok {
|
||||
t.Fatalf("claude exited before result")
|
||||
}
|
||||
if onEvent != nil {
|
||||
onEvent(ev)
|
||||
}
|
||||
if ev.Type == "assistant" {
|
||||
blocks, err := ev.Blocks()
|
||||
if err != nil {
|
||||
t.Fatalf("parse assistant: %v", err)
|
||||
}
|
||||
for _, b := range blocks {
|
||||
if b.Type == "text" {
|
||||
text = append(text, b.Text)
|
||||
}
|
||||
}
|
||||
}
|
||||
if ev.Type == "result" {
|
||||
res, err := ev.Result()
|
||||
if err != nil {
|
||||
t.Fatalf("parse result: %v", err)
|
||||
}
|
||||
t.Logf("result %s: %d turns, %dms, $%.4f", res.Subtype, res.NumTurns, res.DurationMS, res.TotalCostUSD)
|
||||
return strings.Join(text, "\n")
|
||||
}
|
||||
case <-ctx.Done():
|
||||
t.Fatalf("timed out waiting for result")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func hasToolUse(t *testing.T, ev claude.Event) bool {
|
||||
if ev.Type != "assistant" {
|
||||
return false
|
||||
}
|
||||
blocks, err := ev.Blocks()
|
||||
if err != nil {
|
||||
t.Fatalf("parse assistant: %v", err)
|
||||
}
|
||||
for _, b := range blocks {
|
||||
if b.Type == "tool_use" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func git(t *testing.T, dir string, args ...string) string {
|
||||
cmd := exec.Command("git", args...)
|
||||
cmd.Dir = dir
|
||||
cmd.Env = append(os.Environ(),
|
||||
"GIT_AUTHOR_NAME="+gitName,
|
||||
"GIT_AUTHOR_EMAIL="+gitEmail,
|
||||
"GIT_COMMITTER_NAME="+gitName,
|
||||
"GIT_COMMITTER_EMAIL="+gitEmail,
|
||||
)
|
||||
out, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
t.Fatalf("git %v: %v\n%s", args, err, out)
|
||||
}
|
||||
return string(out)
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
const protocolVersion = "2025-06-18"
|
||||
|
||||
type Tool struct {
|
||||
Name string
|
||||
Description string
|
||||
InputSchema json.RawMessage
|
||||
Handler func(ctx context.Context, args json.RawMessage) (string, error)
|
||||
}
|
||||
|
||||
type Server struct {
|
||||
Name string
|
||||
Version string
|
||||
Tools []Tool
|
||||
}
|
||||
|
||||
type request struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID json.RawMessage `json:"id,omitempty"`
|
||||
Method string `json:"method"`
|
||||
Params json.RawMessage `json:"params,omitempty"`
|
||||
}
|
||||
|
||||
type rpcError struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
type response struct {
|
||||
JSONRPC string `json:"jsonrpc"`
|
||||
ID json.RawMessage `json:"id"`
|
||||
Result any `json:"result,omitempty"`
|
||||
Error *rpcError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type toolDesc struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
InputSchema json.RawMessage `json:"inputSchema"`
|
||||
}
|
||||
|
||||
type content struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
type callResult struct {
|
||||
Content []content `json:"content"`
|
||||
IsError bool `json:"isError,omitempty"`
|
||||
}
|
||||
|
||||
func (s *Server) Handle(ctx context.Context, raw json.RawMessage) (json.RawMessage, error) {
|
||||
req := request{}
|
||||
err := json.Unmarshal(raw, &req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse jsonrpc: %w", err)
|
||||
}
|
||||
|
||||
if req.ID == nil {
|
||||
return json.Marshal(response{JSONRPC: "2.0", Result: map[string]any{}})
|
||||
}
|
||||
|
||||
result, rerr := s.dispatch(ctx, req)
|
||||
resp := response{JSONRPC: "2.0", ID: req.ID, Result: result, Error: rerr}
|
||||
return json.Marshal(resp)
|
||||
}
|
||||
|
||||
func (s *Server) dispatch(ctx context.Context, req request) (any, *rpcError) {
|
||||
switch req.Method {
|
||||
case "initialize":
|
||||
return map[string]any{
|
||||
"protocolVersion": protocolVersion,
|
||||
"capabilities": map[string]any{"tools": map[string]any{}},
|
||||
"serverInfo": map[string]any{"name": s.Name, "version": s.Version},
|
||||
}, nil
|
||||
|
||||
case "ping":
|
||||
return map[string]any{}, nil
|
||||
|
||||
case "tools/list":
|
||||
tools := []toolDesc{}
|
||||
for _, t := range s.Tools {
|
||||
tools = append(tools, toolDesc{Name: t.Name, Description: t.Description, InputSchema: t.InputSchema})
|
||||
}
|
||||
return map[string]any{"tools": tools}, nil
|
||||
|
||||
case "tools/call":
|
||||
return s.call(ctx, req.Params)
|
||||
}
|
||||
|
||||
return nil, &rpcError{Code: -32601, Message: fmt.Sprintf("method not found: %s", req.Method)}
|
||||
}
|
||||
|
||||
func (s *Server) call(ctx context.Context, params json.RawMessage) (any, *rpcError) {
|
||||
p := struct {
|
||||
Name string `json:"name"`
|
||||
Arguments json.RawMessage `json:"arguments"`
|
||||
}{}
|
||||
err := json.Unmarshal(params, &p)
|
||||
if err != nil {
|
||||
return nil, &rpcError{Code: -32602, Message: err.Error()}
|
||||
}
|
||||
|
||||
for _, t := range s.Tools {
|
||||
if t.Name != p.Name {
|
||||
continue
|
||||
}
|
||||
|
||||
text, err := t.Handler(ctx, p.Arguments)
|
||||
if err != nil {
|
||||
return callResult{Content: []content{{Type: "text", Text: err.Error()}}, IsError: true}, nil
|
||||
}
|
||||
|
||||
return callResult{Content: []content{{Type: "text", Text: text}}}, nil
|
||||
}
|
||||
|
||||
return nil, &rpcError{Code: -32602, Message: fmt.Sprintf("unknown tool: %s", p.Name)}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package report
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"g.fc.run/misc/byre/internal/mcp"
|
||||
)
|
||||
|
||||
const maxNudges = 3
|
||||
|
||||
const nudge = "You must call the report tool (mcp__byre__report) before ending your turn."
|
||||
|
||||
type Report struct {
|
||||
Done []string `json:"done"`
|
||||
OpenQuestions []string `json:"open_questions"`
|
||||
NextSteps []string `json:"next_steps"`
|
||||
Risks []string `json:"risks"`
|
||||
Deviations []string `json:"deviations"`
|
||||
}
|
||||
|
||||
type Gate struct {
|
||||
OnReport func(Report)
|
||||
|
||||
mu sync.Mutex
|
||||
reported bool
|
||||
nudges int
|
||||
}
|
||||
|
||||
func (g *Gate) Reset() {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
g.reported = false
|
||||
g.nudges = 0
|
||||
}
|
||||
|
||||
func (g *Gate) Nudges() int {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
return g.nudges
|
||||
}
|
||||
|
||||
func (g *Gate) Tool() mcp.Tool {
|
||||
return mcp.Tool{
|
||||
Name: "report",
|
||||
Description: "Report what you did this turn. You must call this before ending every turn. List every departure from the plan in deviations.",
|
||||
InputSchema: json.RawMessage(`{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"done": {"type": "array", "items": {"type": "string"}},
|
||||
"open_questions": {"type": "array", "items": {"type": "string"}},
|
||||
"next_steps": {"type": "array", "items": {"type": "string"}},
|
||||
"risks": {"type": "array", "items": {"type": "string"}},
|
||||
"deviations": {"type": "array", "items": {"type": "string"}}
|
||||
},
|
||||
"required": ["done"]
|
||||
}`),
|
||||
Handler: g.handle,
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gate) handle(ctx context.Context, args json.RawMessage) (string, error) {
|
||||
r := Report{}
|
||||
err := json.Unmarshal(args, &r)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("malformed report: %w", err)
|
||||
}
|
||||
|
||||
if len(r.Done) == 0 {
|
||||
return "", fmt.Errorf("report must list at least one item in done")
|
||||
}
|
||||
|
||||
g.mu.Lock()
|
||||
g.reported = true
|
||||
g.mu.Unlock()
|
||||
|
||||
if g.OnReport != nil {
|
||||
g.OnReport(r)
|
||||
}
|
||||
|
||||
return "report filed", nil
|
||||
}
|
||||
|
||||
func (g *Gate) StopHook(ctx context.Context, input json.RawMessage) (map[string]any, error) {
|
||||
g.mu.Lock()
|
||||
defer g.mu.Unlock()
|
||||
|
||||
if g.reported || g.nudges >= maxNudges {
|
||||
return map[string]any{}, nil
|
||||
}
|
||||
|
||||
g.nudges++
|
||||
return map[string]any{"decision": "block", "reason": nudge}, nil
|
||||
}
|
||||
Reference in New Issue
Block a user