Remove unused raw protocol recorders from agent sessions
This commit is contained in:
@@ -41,7 +41,6 @@ type Options struct {
|
|||||||
Hooks map[string][]HookMatcher
|
Hooks map[string][]HookMatcher
|
||||||
MCPServers []*mcp.Server
|
MCPServers []*mcp.Server
|
||||||
StdioMCP []mcp.Stdio
|
StdioMCP []mcp.Stdio
|
||||||
Transcript io.Writer
|
|
||||||
OnBadLine func(line []byte, reason string)
|
OnBadLine func(line []byte, reason string)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -198,8 +197,6 @@ type Session struct {
|
|||||||
done chan struct{}
|
done chan struct{}
|
||||||
|
|
||||||
writeMu sync.Mutex
|
writeMu sync.Mutex
|
||||||
transcript io.Writer
|
|
||||||
scrub []byte
|
|
||||||
|
|
||||||
pendingMu sync.Mutex
|
pendingMu sync.Mutex
|
||||||
pending map[string]chan controlResult
|
pending map[string]chan controlResult
|
||||||
@@ -237,7 +234,6 @@ func Start(ctx context.Context, opts Options) (*Session, error) {
|
|||||||
s := &Session{
|
s := &Session{
|
||||||
events: make(chan Event, 256),
|
events: make(chan Event, 256),
|
||||||
done: make(chan struct{}),
|
done: make(chan struct{}),
|
||||||
transcript: opts.Transcript,
|
|
||||||
pending: map[string]chan controlResult{},
|
pending: map[string]chan controlResult{},
|
||||||
inflight: map[string]context.CancelFunc{},
|
inflight: map[string]context.CancelFunc{},
|
||||||
hooks: map[string]HookFunc{},
|
hooks: map[string]HookFunc{},
|
||||||
@@ -524,30 +520,10 @@ func (s *Session) write(v any) error {
|
|||||||
s.writeMu.Lock()
|
s.writeMu.Lock()
|
||||||
defer s.writeMu.Unlock()
|
defer s.writeMu.Unlock()
|
||||||
|
|
||||||
s.record("send", b)
|
|
||||||
|
|
||||||
_, err = s.stdin.Write(append(b, '\n'))
|
_, err = s.stdin.Write(append(b, '\n'))
|
||||||
return err
|
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) {
|
func (s *Session) readLoop(stdout io.Reader) {
|
||||||
defer close(s.done)
|
defer close(s.done)
|
||||||
defer close(s.events)
|
defer close(s.events)
|
||||||
@@ -587,15 +563,6 @@ func (s *Session) handleLine(line []byte) {
|
|||||||
s.badLine(line, fmt.Sprintf("unrecognized type %q", env.Type))
|
s.badLine(line, fmt.Sprintf("unrecognized 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 {
|
switch env.Type {
|
||||||
case "control_response":
|
case "control_response":
|
||||||
s.handleControlResponse(env.Response)
|
s.handleControlResponse(env.Response)
|
||||||
@@ -638,21 +605,6 @@ func (s *Session) badLine(line []byte, reason string) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
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) {
|
func (s *Session) handleControlResponse(raw json.RawMessage) {
|
||||||
resp := struct {
|
resp := struct {
|
||||||
Subtype string `json:"subtype"`
|
Subtype string `json:"subtype"`
|
||||||
|
|||||||
@@ -11,7 +11,6 @@ import (
|
|||||||
"io"
|
"io"
|
||||||
"log/slog"
|
"log/slog"
|
||||||
"os"
|
"os"
|
||||||
"regexp"
|
|
||||||
"strconv"
|
"strconv"
|
||||||
"sync"
|
"sync"
|
||||||
"sync/atomic"
|
"sync/atomic"
|
||||||
@@ -33,7 +32,6 @@ type Options struct {
|
|||||||
Reminder string
|
Reminder string
|
||||||
MCPServers []*mcp.Server
|
MCPServers []*mcp.Server
|
||||||
StdioMCP []mcp.Stdio
|
StdioMCP []mcp.Stdio
|
||||||
Transcript io.Writer
|
|
||||||
OnBadLine func(line []byte, reason string)
|
OnBadLine func(line []byte, reason string)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -89,8 +87,6 @@ type Session struct {
|
|||||||
done chan struct{}
|
done chan struct{}
|
||||||
|
|
||||||
writeLock chan struct{}
|
writeLock chan struct{}
|
||||||
recordMu sync.Mutex
|
|
||||||
transcript io.Writer
|
|
||||||
|
|
||||||
nextID atomic.Int64
|
nextID atomic.Int64
|
||||||
pendingMu sync.Mutex
|
pendingMu sync.Mutex
|
||||||
@@ -185,7 +181,6 @@ func newSession(opts Options, stdin *os.File, stdout io.Reader) *Session {
|
|||||||
events: make(chan Event, 256),
|
events: make(chan Event, 256),
|
||||||
writeLock: make(chan struct{}, 1),
|
writeLock: make(chan struct{}, 1),
|
||||||
done: make(chan struct{}),
|
done: make(chan struct{}),
|
||||||
transcript: opts.Transcript,
|
|
||||||
pending: map[int64]chan rpcResult{},
|
pending: map[int64]chan rpcResult{},
|
||||||
servers: map[string]*mcp.Server{},
|
servers: map[string]*mcp.Server{},
|
||||||
reminder: opts.Reminder,
|
reminder: opts.Reminder,
|
||||||
@@ -552,8 +547,6 @@ func (s *Session) write(m message, deadline time.Time) error {
|
|||||||
return os.ErrDeadlineExceeded
|
return os.ErrDeadlineExceeded
|
||||||
}
|
}
|
||||||
|
|
||||||
s.record("send", b)
|
|
||||||
|
|
||||||
err = s.stdin.SetWriteDeadline(deadline)
|
err = s.stdin.SetWriteDeadline(deadline)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
@@ -566,26 +559,6 @@ func (s *Session) write(m message, deadline time.Time) error {
|
|||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
var emailField = regexp.MustCompile(`"email":"[^"]*"`)
|
|
||||||
|
|
||||||
func (s *Session) record(dir string, line []byte) {
|
|
||||||
if s.transcript == nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
line = emailField.ReplaceAll(line, []byte(`"email":"[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.recordMu.Lock()
|
|
||||||
defer s.recordMu.Unlock()
|
|
||||||
s.transcript.Write(append(b, '\n'))
|
|
||||||
}
|
|
||||||
|
|
||||||
func (s *Session) readLoop(stdout io.Reader) {
|
func (s *Session) readLoop(stdout io.Reader) {
|
||||||
defer close(s.done)
|
defer close(s.done)
|
||||||
defer close(s.events)
|
defer close(s.events)
|
||||||
@@ -620,8 +593,6 @@ func (s *Session) handleLine(line []byte) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
s.record("recv", line)
|
|
||||||
|
|
||||||
if m.Method == "" {
|
if m.Method == "" {
|
||||||
s.handleResponse(line, m)
|
s.handleResponse(line, m)
|
||||||
return
|
return
|
||||||
|
|||||||
Reference in New Issue
Block a user