746 lines
16 KiB
Go
746 lines
16 KiB
Go
package claude
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/base64"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"os"
|
|
"os/exec"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"g.fc.run/ai/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
|
|
SessionID string
|
|
Resume string
|
|
ResumeAt string
|
|
ForkSession bool
|
|
GitName string
|
|
GitEmail string
|
|
Env []string
|
|
Guidance string
|
|
Hooks map[string][]HookMatcher
|
|
MCPServers []*mcp.Server
|
|
StdioMCP []mcp.Stdio
|
|
Transcript io.Writer
|
|
OnBadLine func(line []byte, reason string)
|
|
}
|
|
|
|
var knownTypes = map[string]bool{
|
|
"assistant": true,
|
|
"user": true,
|
|
"result": true,
|
|
"system": true,
|
|
"stream_event": true,
|
|
"tool_progress": true,
|
|
"rate_limit_event": true,
|
|
"control_request": true,
|
|
"control_response": true,
|
|
"control_cancel_request": true,
|
|
"keep_alive": true,
|
|
}
|
|
|
|
type Event struct {
|
|
Type string
|
|
Subtype string
|
|
Raw json.RawMessage
|
|
Read time.Time
|
|
}
|
|
|
|
type Block struct {
|
|
Type string `json:"type"`
|
|
ID string `json:"id"`
|
|
ToolUseID string `json:"tool_use_id"`
|
|
Text string `json:"text"`
|
|
Thinking string `json:"thinking"`
|
|
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"`
|
|
ModelUsage map[string]ModelUsage `json:"modelUsage"`
|
|
}
|
|
|
|
type ModelUsage struct {
|
|
ContextWindow int `json:"contextWindow"`
|
|
}
|
|
|
|
type ContextUsage struct {
|
|
Model string
|
|
Tokens int
|
|
}
|
|
|
|
func (e Event) ContextUsage() (ContextUsage, error) {
|
|
msg := struct {
|
|
Message struct {
|
|
Model string `json:"model"`
|
|
Usage struct {
|
|
InputTokens int `json:"input_tokens"`
|
|
CacheCreationInputTokens int `json:"cache_creation_input_tokens"`
|
|
CacheReadInputTokens int `json:"cache_read_input_tokens"`
|
|
OutputTokens int `json:"output_tokens"`
|
|
} `json:"usage"`
|
|
} `json:"message"`
|
|
}{}
|
|
err := json.Unmarshal(e.Raw, &msg)
|
|
if err != nil {
|
|
return ContextUsage{}, err
|
|
}
|
|
u := msg.Message.Usage
|
|
return ContextUsage{
|
|
Model: msg.Message.Model,
|
|
Tokens: u.InputTokens + u.CacheCreationInputTokens + u.CacheReadInputTokens + u.OutputTokens,
|
|
}, nil
|
|
}
|
|
|
|
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) ParentToolUseID() (string, error) {
|
|
msg := struct {
|
|
ParentToolUseID string `json:"parent_tool_use_id"`
|
|
}{}
|
|
err := json.Unmarshal(e.Raw, &msg)
|
|
return msg.ParentToolUseID, err
|
|
}
|
|
|
|
func (e Event) UUID() (string, error) {
|
|
msg := struct {
|
|
UUID string `json:"uuid"`
|
|
}{}
|
|
err := json.Unmarshal(e.Raw, &msg)
|
|
return msg.UUID, err
|
|
}
|
|
|
|
func (e Event) TextDelta() (string, error) {
|
|
msg := struct {
|
|
ParentToolUseID *string `json:"parent_tool_use_id"`
|
|
Event struct {
|
|
Type string `json:"type"`
|
|
Delta struct {
|
|
Type string `json:"type"`
|
|
Text string `json:"text"`
|
|
} `json:"delta"`
|
|
} `json:"event"`
|
|
}{}
|
|
err := json.Unmarshal(e.Raw, &msg)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if msg.ParentToolUseID != nil || msg.Event.Type != "content_block_delta" || msg.Event.Delta.Type != "text_delta" {
|
|
return "", nil
|
|
}
|
|
return msg.Event.Delta.Text, 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
|
|
scrub []byte
|
|
|
|
pendingMu sync.Mutex
|
|
pending map[string]chan controlResult
|
|
|
|
hooks map[string]HookFunc
|
|
servers map[string]*mcp.Server
|
|
|
|
idMu sync.Mutex
|
|
id string
|
|
|
|
onBadLine func(line []byte, reason string)
|
|
lastLine atomic.Int64
|
|
}
|
|
|
|
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{},
|
|
id: initialID(opts),
|
|
onBadLine: opts.OnBadLine,
|
|
}
|
|
s.lastLine.Store(time.Now().UnixNano())
|
|
|
|
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 initialID(opts Options) string {
|
|
if opts.SessionID != "" {
|
|
return opts.SessionID
|
|
}
|
|
if opts.ForkSession {
|
|
return ""
|
|
}
|
|
return opts.Resume
|
|
}
|
|
|
|
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, "alwaysLoad": true}
|
|
}
|
|
for _, srv := range opts.StdioMCP {
|
|
servers[srv.Name] = map[string]any{"type": "stdio", "command": srv.Command, "args": srv.Args}
|
|
}
|
|
|
|
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",
|
|
"--include-partial-messages",
|
|
"--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.Guidance != "" {
|
|
args = append(args, "--append-system-prompt", opts.Guidance)
|
|
}
|
|
|
|
if opts.SessionID != "" {
|
|
args = append(args, "--session-id="+opts.SessionID)
|
|
}
|
|
|
|
if opts.Resume != "" {
|
|
args = append(args, "--resume="+opts.Resume)
|
|
}
|
|
|
|
if opts.ResumeAt != "" {
|
|
args = append(args, "--resume-session-at="+opts.ResumeAt)
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
env = 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,
|
|
)
|
|
return append(env, opts.Env...)
|
|
}
|
|
|
|
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) Pid() int {
|
|
return s.cmd.Process.Pid
|
|
}
|
|
|
|
func (s *Session) LastLine() time.Time {
|
|
return time.Unix(0, s.lastLine.Load())
|
|
}
|
|
|
|
func (s *Session) Exited() bool {
|
|
select {
|
|
case <-s.done:
|
|
return true
|
|
default:
|
|
return false
|
|
}
|
|
}
|
|
|
|
func (s *Session) ID() string {
|
|
s.idMu.Lock()
|
|
defer s.idMu.Unlock()
|
|
return s.id
|
|
}
|
|
|
|
type Image struct {
|
|
MediaType string
|
|
Data []byte
|
|
}
|
|
|
|
func (s *Session) Send(text string, images ...Image) error {
|
|
var content any = text
|
|
if len(images) > 0 {
|
|
blocks := []map[string]any{}
|
|
for _, img := range images {
|
|
blocks = append(blocks, map[string]any{
|
|
"type": "image",
|
|
"source": map[string]any{
|
|
"type": "base64",
|
|
"media_type": img.MediaType,
|
|
"data": base64.StdEncoding.EncodeToString(img.Data),
|
|
},
|
|
})
|
|
}
|
|
if text != "" {
|
|
blocks = append(blocks, map[string]any{"type": "text", "text": text})
|
|
}
|
|
content = blocks
|
|
}
|
|
|
|
return s.write(map[string]any{
|
|
"type": "user",
|
|
"message": map[string]any{"role": "user", "content": content},
|
|
"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) Kill() {
|
|
s.cmd.Process.Kill()
|
|
<-s.done
|
|
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
|
|
}
|
|
|
|
if len(s.scrub) > 0 {
|
|
line = bytes.ReplaceAll(line, s.scrub, []byte("[redacted]"))
|
|
}
|
|
|
|
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')
|
|
line = bytes.TrimRight(line, "\r\n")
|
|
if len(line) > 0 {
|
|
s.lastLine.Store(time.Now().UnixNano())
|
|
}
|
|
if len(line) > 0 && line[0] == '{' {
|
|
s.handleLine(line)
|
|
} else if len(line) > 0 {
|
|
s.badLine(line, "not a JSON object")
|
|
}
|
|
if err != nil {
|
|
if err != io.EOF {
|
|
slog.Error("read claude stdout", "err", err)
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *Session) handleLine(line []byte) {
|
|
read := time.Now()
|
|
env := envelope{}
|
|
err := json.Unmarshal(line, &env)
|
|
if err != nil {
|
|
s.badLine(line, err.Error())
|
|
return
|
|
}
|
|
|
|
if !knownTypes[env.Type] {
|
|
s.badLine(line, fmt.Sprintf("unrecognised type %q", env.Type))
|
|
}
|
|
|
|
s.writeMu.Lock()
|
|
if env.Type == "control_response" {
|
|
s.learnScrub(env.Response)
|
|
}
|
|
if env.Type != "stream_event" {
|
|
s.record("recv", line)
|
|
}
|
|
s.writeMu.Unlock()
|
|
|
|
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...), Read: read}
|
|
}
|
|
|
|
func (s *Session) badLine(line []byte, reason string) {
|
|
slog.Warn("bad claude line", "reason", reason, "line", string(line))
|
|
if s.onBadLine != nil {
|
|
s.onBadLine(line, reason)
|
|
}
|
|
}
|
|
|
|
func (s *Session) learnScrub(raw json.RawMessage) {
|
|
resp := struct {
|
|
Response struct {
|
|
Account struct {
|
|
Email string `json:"email"`
|
|
} `json:"account"`
|
|
} `json:"response"`
|
|
}{}
|
|
err := json.Unmarshal(raw, &resp)
|
|
if err != nil || resp.Response.Account.Email == "" {
|
|
return
|
|
}
|
|
s.scrub = []byte(resp.Response.Account.Email)
|
|
}
|
|
|
|
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)
|
|
}
|